modelbeast/server/operators/mflux_image_edit/run.py
type-two 431aeb08ea feat(mflux_image_edit): Qwen-Image-Edit LoRA support
mflux-generate-qwen-edit already exposes --lora-paths/--lora-scales; wire them
to params.loras (['multiple-angles=1.0', ...]). Names resolve by filename stem
or fuzzy match across QWEN_LORA_DIRS (default ~/Documents/localmodels/QwenLora)
so John can drop new .safetensors in and use them by short name. Misses log
loudly instead of failing the job.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-25 00:57:06 +10:00

86 lines
3.0 KiB
Python

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path
# Qwen-Image-Edit LoRA search path (name → file). Drop .safetensors anywhere under these and they
# become usable as {"loras": ["multiple-angles=1.0", ...]}. Same UX as sd_local's LoRA list.
QWEN_LORA_DIRS = [Path(os.path.expanduser(d)) for d in os.environ.get(
"QWEN_LORA_DIRS", "~/Documents/localmodels/QwenLora").split(":")]
def find_lora(stem: str):
"""Match by exact filename stem, else the first file whose stem contains it (case-insensitive)."""
stem = stem.strip()
cands = [f for d in QWEN_LORA_DIRS if d.is_dir() for f in d.rglob("*.safetensors")]
for f in cands:
if f.stem == stem:
return f
low = stem.lower()
for f in cands:
if low in f.stem.lower() or low in f.parent.name.lower():
return f
return None
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())]
# optional LoRAs: ["multiple-angles=1.0", "gymnastics-pose=0.8"] → --lora-paths/--lora-scales
paths, scales = [], []
for spec in (p.get("loras") or []):
name, _, w = str(spec).partition("=")
f = find_lora(name)
if not f:
print(f"LORA MISS '{name}' — searched {[str(d) for d in QWEN_LORA_DIRS]}", flush=True)
continue
paths.append(str(f))
scales.append(str(float(w or 1.0)))
if paths:
cmd += ["--lora-paths", *paths, "--lora-scales", *scales]
print(f"loras: {list(zip(paths, scales))}", flush=True)
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)