foitin/pipeline/restyle_frames.py
2026-08-16 15:43:29 +10:00

192 lines
8.0 KiB
Python

#!/usr/bin/env python3
"""restyle_frames.py — generate frames@<style>/ variant dirs for a packed
character by restyling its existing base frames. Two engines:
pillow-xray deterministic local pass (no AI): dark silhouette fill +
glowing edge lines. Instant, zero flicker — the "digital
guts" tier. Needs Pillow only.
farm any MODELBEAST image->image operator (mflux_image_edit,
comfyui_sd, ...). Frame is flattened onto neutral gray,
submitted, and the result is re-masked with the original
alpha so the sprite silhouette is untouched. Fixed seed
per style keeps frame-to-frame drift down (some shimmer
is period-correct for digitized fighters). SLOW operators
(mflux_image_edit ≈ 4 min/frame on m3ultra) are hero-tier
only — check timing on one frame before a big batch.
Usage:
python3 pipeline/restyle_frames.py characters/vesper --style xray --engine pillow-xray
python3 pipeline/restyle_frames.py characters/vesper --style sketch --engine farm \
--operator mflux_image_edit --prompt "Redraw as a pencil sketch ..." \
--op-params '{"steps": 20}' --moves idle,hit --limit 10
Resume-safe: frames whose output already exists are skipped. Variant dirs
must keep the SAME frame count as base (engine falls back per-frame past
the variant's end, which pops mid-move) — this script always writes 1:1.
Heartbeat: --hb writes progress lines to ~/.jobs/foitin-restyle.hb.
"""
import argparse, json, mimetypes, os, sys, time, urllib.request, uuid
HOST = os.environ.get('MB_HOST', 'http://100.89.131.57:8777')
HB = os.path.expanduser('~/.jobs/foitin-restyle.hb')
def hb(msg, enabled):
if not enabled:
return
os.makedirs(os.path.dirname(HB), exist_ok=True)
with open(HB, 'a') as f:
f.write(time.strftime('%H:%M:%S ') + msg + '\n')
# ---------------------------------------------------------------- farm client
def token():
t = os.environ.get('MB_TOKEN')
if t:
return t
for line in open(os.path.expanduser('~/Documents/backnforth/.env')):
if line.startswith('MB_TOKEN='):
return line.split('=', 1)[1].strip()
sys.exit('no MB_TOKEN')
def req(path, data=None, headers=None, raw=False):
h = {'Authorization': f'Bearer {token()}'}
h.update(headers or {})
r = urllib.request.Request(HOST + path, data=data, headers=h)
body = urllib.request.urlopen(r, timeout=300).read()
if raw:
return body
s = body.decode('utf-8', 'replace')
return json.loads(''.join(c if c >= ' ' or c in '\t' else ' ' for c in s))
def upload(path):
boundary = uuid.uuid4().hex
ctype = mimetypes.guess_type(path)[0] or 'application/octet-stream'
body = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; '
f'filename="{os.path.basename(path)}"\r\nContent-Type: {ctype}\r\n\r\n').encode() \
+ open(path, 'rb').read() + f'\r\n--{boundary}--\r\n'.encode()
a = req('/api/assets', data=body, headers={'Content-Type': f'multipart/form-data; boundary={boundary}'})
return a.get('id') or (a.get('items') or [a])[0].get('id')
def job_output_asset(jid):
a = req('/api/assets?limit=100')
items = a if isinstance(a, list) else a.get('items', [])
imgs = [x for x in items if str(x.get('filename', x.get('name', ''))).lower()
.endswith(('.png', '.jpg', '.webp'))]
mine = [x for x in imgs if x.get('parent_job') == jid]
return mine[0]['id'] if mine else None
def farm_edit(Image, src_png, out_png, operator, params, tmp_dir):
im = Image.open(src_png).convert('RGBA')
alpha = im.split()[3]
flat = Image.new('RGB', im.size, (128, 128, 128))
flat.paste(im, (0, 0), alpha)
tmp_in = os.path.join(tmp_dir, 'mb_in.png')
flat.save(tmp_in)
jid = req('/api/jobs', data=json.dumps(
{'operator': operator, 'asset_id': upload(tmp_in), 'params': params}).encode(),
headers={'Content-Type': 'application/json'})['id']
while True:
time.sleep(5)
st = req(f'/api/jobs/{jid}').get('status')
if st in ('done', 'error', 'cancelled'):
break
if st != 'done':
return False
gid = job_output_asset(jid)
if not gid:
return False
tmp_out = os.path.join(tmp_dir, 'mb_out.png')
open(tmp_out, 'wb').write(req(f'/api/assets/{gid}/file', raw=True))
edited = Image.open(tmp_out).convert('RGB')
if edited.size != im.size:
edited = edited.resize(im.size, Image.LANCZOS)
edited = edited.convert('RGBA')
edited.putalpha(alpha)
edited.save(out_png)
return True
# ------------------------------------------------------------- pillow engine
def pillow_xray(Image, src_png, out_png, fill=(10, 22, 34), line=(87, 230, 255)):
from PIL import ImageChops, ImageFilter, ImageOps
im = Image.open(src_png).convert('RGBA')
alpha = im.split()[3]
gray = im.convert('L')
edges = ImageOps.autocontrast(gray.filter(ImageFilter.FIND_EDGES))
edges = edges.filter(ImageFilter.MaxFilter(3))
edges = ImageChops.multiply(edges, alpha) # clip glow to the silhouette
out = Image.new('RGBA', im.size, (0, 0, 0, 0))
out.paste(Image.new('RGBA', im.size, fill + (255,)), (0, 0), alpha)
out.paste(Image.new('RGBA', im.size, line + (255,)), (0, 0), edges)
out.putalpha(alpha)
out.save(out_png)
return True
# ----------------------------------------------------------------------- main
def main():
ap = argparse.ArgumentParser()
ap.add_argument('char_dir', help='packed character folder (characters/<id>)')
ap.add_argument('--style', required=True, help='variant name -> frames@<style>/')
ap.add_argument('--engine', default='pillow-xray', choices=['pillow-xray', 'farm'])
ap.add_argument('--operator', default='mflux_image_edit')
ap.add_argument('--prompt', default='')
ap.add_argument('--op-params', default='{}', help='extra operator params as JSON')
ap.add_argument('--moves', default='', help='comma list; default = all moves')
ap.add_argument('--limit', type=int, default=0, help='max frames this run (0 = all)')
ap.add_argument('--hb', action='store_true', help='heartbeat to ~/.jobs/foitin-restyle.hb')
a = ap.parse_args()
from PIL import Image
params = json.loads(a.op_params)
if a.prompt:
params['prompt'] = a.prompt
params.setdefault('seed', 42)
moves_root = os.path.join(a.char_dir, 'moves')
moves = sorted(a.moves.split(',') if a.moves else os.listdir(moves_root))
tmp_dir = f'/tmp/restyle_{os.getpid()}'
os.makedirs(tmp_dir, exist_ok=True)
done = failed = 0
for mv in moves:
frames_dir = os.path.join(moves_root, mv, 'frames')
if not os.path.isdir(frames_dir):
continue
out_dir = os.path.join(moves_root, mv, 'frames@' + a.style)
os.makedirs(out_dir, exist_ok=True)
pngs = sorted(f for f in os.listdir(frames_dir) if f.endswith('.png'))
for i, f in enumerate(pngs):
out_png = os.path.join(out_dir, f)
if os.path.exists(out_png):
continue
if a.limit and done >= a.limit:
hb(f'limit {a.limit} reached', a.hb)
print(f'[restyle] limit reached ({done} frames)')
return
src = os.path.join(frames_dir, f)
if a.engine == 'pillow-xray':
ok = pillow_xray(Image, src, out_png)
else:
ok = farm_edit(Image, src, out_png, a.operator, dict(params), tmp_dir)
done += ok
failed += (not ok)
if not ok:
print(f'[restyle] FAILED {mv}/{f}', flush=True)
if done % 20 == 0 or a.engine == 'farm':
hb(f'{mv} {i + 1}/{len(pngs)} (total {done} done {failed} failed)', a.hb)
print(f'[restyle] {mv} {i + 1}/{len(pngs)} (total {done})', flush=True)
hb(f'ALL DONE ({done} frames, {failed} failed)', a.hb)
print(f'[restyle] ALL DONE ({done} frames, {failed} failed)')
main()