84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
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 the fleet-standard stash, which is
|
|
# present on BOTH mflux nodes (m3ultra + ultra) — a job can dispatch to either,
|
|
# so a node-local dir would resolve on one and fail on the other.
|
|
# mflux natively takes --lora-paths/--lora-scales; a LoRA for the wrong base
|
|
# model is a SILENT no-op (no error, the edit just ignores it).
|
|
LORA_DIRS = [Path(d).expanduser() for d in
|
|
(os.environ.get("QWEN_LORA_DIRS") or "").split(os.pathsep) if d.strip()] \
|
|
or [Path.home() / "Documents" / "localmodels" / "QwenLora"]
|
|
if p.get("loras"):
|
|
paths, scales = [], []
|
|
avail = {}
|
|
for d in LORA_DIRS:
|
|
for f in d.rglob("*.safetensors"):
|
|
avail.setdefault(f.name.lower(), f)
|
|
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)
|