restyle_frames.py: --control edge uploads a canny-style map for comfyui_sd controlnet

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
type-two 2026-08-16 16:00:05 +10:00
parent ecdaeda47e
commit e0699d47b5

View File

@ -81,13 +81,20 @@ def job_output_asset(jid):
return mine[0]['id'] if mine else None
def farm_edit(Image, src_png, out_png, operator, params, tmp_dir):
def farm_edit(Image, src_png, out_png, operator, params, tmp_dir, control='flat'):
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)
if control == 'edge':
# canny-style conditioning image (thin white lines on black) for
# ControlNet graphs with no preprocessor node (comfyui_sd)
from PIL import ImageChops, ImageFilter, ImageOps
edges = ImageOps.autocontrast(im.convert('L').filter(ImageFilter.FIND_EDGES))
ImageChops.multiply(edges, alpha).convert('RGB').save(tmp_in)
else:
flat = Image.new('RGB', im.size, (128, 128, 128))
flat.paste(im, (0, 0), alpha)
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']
@ -140,6 +147,8 @@ def main():
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('--control', default='flat', choices=['flat', 'edge'],
help='what to upload as the job asset: flattened frame, or an edge map (controlnet)')
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')
@ -176,7 +185,7 @@ def main():
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)
ok = farm_edit(Image, src, out_png, a.operator, dict(params), tmp_dir, a.control)
done += ok
failed += (not ok)
if not ok: