modelbeast/server/operators/mflux_image_edit/run.py
m3ultra 4c05a9daa8 docs: wan_video + ardy_motion measured RAM/speed (two-machine verification)
wan_video 5B q8 verified clean on M3 Ultra (148.7s) and M2 Max (503.7s),
identical 9.2GB max RSS on both; 81GB 'footprint' on m2max identified as
MLX opportunistic cache, not demand. Frame-identical output at same seed
across machines. HARDWARE.md gains wan_video and ardy_motion rows.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-08-05 16:44:08 +10:00

76 lines
2.9 KiB
Python

import argparse
import json
import subprocess
import sys
from pathlib import Path
ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)
if not a.input:
print("ERROR: no input image")
sys.exit(1)
if not p.get("prompt"):
print("ERROR: prompt is required — describe the edit you want")
sys.exit(1)
cli = Path(sys.executable).parent / "mflux-generate-qwen-edit"
if not cli.exists():
print(f"ERROR: {cli} missing. Run scripts/install_mflux.sh")
sys.exit(1)
src = Path(a.input[0])
outdir = Path(a.outdir)
out = outdir / f"{src.stem}_edited.png"
# ALL inputs go to Qwen (it composes up to ~3 images: "put the shirt from image 2 on the person
# in image 1"). Single-image jobs behave exactly as before.
cmd = [str(cli), "--image-paths", *[str(Path(x).resolve()) for x in a.input],
"--prompt", p["prompt"],
"--steps", str(p.get("steps", 25)),
"--seed", str(p.get("seed", 42)),
"--guidance", str(p.get("guidance", 4.0)),
"--output", str(out.resolve())]
# Qwen-Edit LoRAs: "name" or "name:scale", comma-separated. Names resolve
# (case-insensitive prefix match) against ~/Documents/localmodels/qwen-loras,
# which mirrors ultra's civit/_parked-qwen stash. mflux natively takes
# --lora-paths/--lora-scales; a LoRA for the wrong base is a silent no-op.
LORA_DIR = Path.home() / "Documents" / "localmodels" / "qwen-loras"
if p.get("loras"):
paths, scales = [], []
avail = {f.name.lower(): f for f in LORA_DIR.rglob("*.safetensors")}
for item in str(p["loras"]).split(","):
item = item.strip()
if not item:
continue
name, _, scale = item.partition(":")
name = name.strip().lower()
hit = avail.get(name) or avail.get(name + ".safetensors") or next(
(f for k, f in sorted(avail.items()) if k.startswith(name)), None)
if not hit:
print(f"ERROR: no LoRA matching '{name}' under {LORA_DIR}")
print("available:", ", ".join(sorted(f.stem for f in avail.values())[:40]))
sys.exit(1)
paths.append(str(hit))
scales.append(str(float(scale) if scale.strip() else 1.0))
if paths:
cmd += ["--lora-paths", *paths, "--lora-scales", *scales]
print("+", " ".join(cmd), flush=True)
print("(first run downloads Qwen-Image-Edit weights)", flush=True)
res = subprocess.run(cmd, cwd=str(outdir), stdout=sys.stdout, stderr=subprocess.STDOUT)
if res.returncode != 0:
sys.exit(res.returncode)
pngs = sorted(outdir.rglob("*.png"))
if not pngs:
print("ERROR: no edited image produced")
sys.exit(1)
img = out if out.exists() else pngs[-1]
(outdir / "result.json").write_text(json.dumps(
{"outputs": [{"path": img.name, "meta": {"tool": "qwen-image-edit"}}]}))
print(f"done: {img.name}", flush=True)