redux_experiment_final/scripts/run_matrix.py

#!/usr/bin/env python3
"""Queue the flux-redux test matrix (seed 777) against a running ComfyUI and collect outputs."""
import json, os, shutil, sys, time, urllib.request
from PIL import Image

HOST = "http://127.0.0.1:7865"
SEED = 777
IN_DIR = "/workspace/flux-inputs"
OUT_ROOT = "/workspace/outputs"
COMFY_OUT = "/workspace/ComfyUI/output"

R = ["r1.jpg", "r2.jpg", "r3.jpg", "r4.jpg"]
T = ["t1.jpeg", "t2.jpg", "t3.jpg", "t4.jpg", "t5.jpg", "t6.jpg"]

def latent_size(img_path, target_mp=1.0, step=16):
    w, h = Image.open(img_path).size
    scale = (target_mp * 1e6 / (w * h)) ** 0.5
    return max(step, round(w * scale / step) * step), max(step, round(h * scale / step) * step)

def post(prompt):
    req = urllib.request.Request(f"{HOST}/prompt",
        data=json.dumps({"prompt": prompt}).encode(), headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as r:
        d = json.loads(r.read())
    if "prompt_id" not in d:
        raise RuntimeError(d)
    return d["prompt_id"]

def flux_core(p, unet="flux1-dev.safetensors"):
    p["u"] = {"class_type": "UNETLoader", "inputs": {"unet_name": unet, "weight_dtype": "default"}}
    p["c"] = {"class_type": "DualCLIPLoader", "inputs": {
        "clip_name1": "t5xxl_fp16.safetensors", "clip_name2": "clip_l.safetensors",
        "type": "flux", "device": "default"}}
    p["v"] = {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}}

def redux_chain(p, style_img, strength, stype):
    p["si"] = {"class_type": "LoadImage", "inputs": {"image": style_img}}
    p["cvl"] = {"class_type": "CLIPVisionLoader", "inputs": {"clip_name": "sigclip_vision_patch14_384.safetensors"}}
    p["cve"] = {"class_type": "CLIPVisionEncode", "inputs": {"clip_vision": ["cvl", 0], "image": ["si", 0], "crop": "center"}}
    p["sml"] = {"class_type": "StyleModelLoader", "inputs": {"style_model_name": "flux1-redux-dev.safetensors"}}
    p["sma"] = {"class_type": "StyleModelApply", "inputs": {
        "conditioning": ["guid", 0], "style_model": ["sml", 0], "clip_vision_output": ["cve", 0],
        "strength": strength, "strength_type": stype}}

def save(p, src, prefix):
    p["dec"] = {"class_type": "VAEDecode", "inputs": {"samples": src, "vae": ["v", 0]}}
    p["save"] = {"class_type": "SaveImage", "inputs": {"images": ["dec", 0], "filename_prefix": prefix}}

def wf_style_composition(style_img, comp_img, tag):
    p = {}
    flux_core(p)
    p["txt"] = {"class_type": "CLIPTextEncode", "inputs": {"clip": ["c", 0], "text": ""}}
    p["guid"] = {"class_type": "FluxGuidance", "inputs": {"conditioning": ["txt", 0], "guidance": 3.5}}
    p["neg"] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["guid", 0]}}
    redux_chain(p, style_img, 0.5, "attn_bias")
    p["ci"] = {"class_type": "LoadImage", "inputs": {"image": comp_img}}
    p["depth"] = {"class_type": "DepthAnythingV2Preprocessor", "inputs": {
        "image": ["ci", 0], "ckpt_name": "depth_anything_v2_vitl.pth", "resolution": 1024}}
    p["cnl"] = {"class_type": "ControlNetLoader", "inputs": {"control_net_name": "FLUX.1-dev-ControlNet-Union-Pro-2.0.safetensors"}}
    p["cn"] = {"class_type": "ControlNetApplySD3", "inputs": {
        "positive": ["sma", 0], "negative": ["neg", 0], "control_net": ["cnl", 0],
        "vae": ["v", 0], "image": ["depth", 0],
        "strength": 0.7, "start_percent": 0.0, "end_percent": 0.8}}
    w, h = latent_size(os.path.join(IN_DIR, comp_img))
    p["msf"] = {"class_type": "ModelSamplingFlux", "inputs": {
        "model": ["u", 0], "max_shift": 1.15, "base_shift": 0.5, "width": w, "height": h}}
    p["lat"] = {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}}
    p["ks"] = {"class_type": "KSampler", "inputs": {
        "model": ["msf", 0], "positive": ["cn", 0], "negative": ["cn", 1], "latent_image": ["lat", 0],
        "seed": SEED, "steps": 28, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0}}
    save(p, ["ks", 0], f"flux-redux-style-composition/{tag}")
    return p

def wf_bfl_lora(style_img, comp_img, tag):
    p = {}
    flux_core(p)
    p["lora"] = {"class_type": "LoraLoaderModelOnly", "inputs": {
        "model": ["u", 0], "lora_name": "flux1-depth-dev-lora.safetensors", "strength_model": 1.0}}
    p["txt"] = {"class_type": "CLIPTextEncode", "inputs": {"clip": ["c", 0], "text": ""}}
    p["guid"] = {"class_type": "FluxGuidance", "inputs": {"conditioning": ["txt", 0], "guidance": 10.0}}
    p["neg"] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["guid", 0]}}
    redux_chain(p, style_img, 0.5, "attn_bias")
    p["ci"] = {"class_type": "LoadImage", "inputs": {"image": comp_img}}
    p["depth"] = {"class_type": "DepthAnythingV2Preprocessor", "inputs": {
        "image": ["ci", 0], "ckpt_name": "depth_anything_v2_vitl.pth", "resolution": 1024}}
    p["ip2p"] = {"class_type": "InstructPixToPixConditioning", "inputs": {
        "positive": ["sma", 0], "negative": ["neg", 0], "vae": ["v", 0], "pixels": ["depth", 0]}}
    w, h = latent_size(os.path.join(IN_DIR, comp_img))
    p["msf"] = {"class_type": "ModelSamplingFlux", "inputs": {
        "model": ["lora", 0], "max_shift": 1.15, "base_shift": 0.5, "width": w, "height": h}}
    p["ks"] = {"class_type": "KSampler", "inputs": {
        "model": ["msf", 0], "positive": ["ip2p", 0], "negative": ["ip2p", 1], "latent_image": ["ip2p", 2],
        "seed": SEED, "steps": 28, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0}}
    save(p, ["ks", 0], f"flux-redux-style-composition-bfl-lora/{tag}")
    return p

def wf_single(style_img, tag, *, unet="flux1-dev.safetensors", prompt_text="", strength=1.0,
              stype="multiply", steps=28, w=768, h=1024, guidance=3.5, prefix="flux-redux-fal-dev",
              use_guidance=True, use_msf=True):
    p = {}
    flux_core(p, unet=unet)
    p["txt"] = {"class_type": "CLIPTextEncode", "inputs": {"clip": ["c", 0], "text": prompt_text}}
    if use_guidance:
        p["guid"] = {"class_type": "FluxGuidance", "inputs": {"conditioning": ["txt", 0], "guidance": guidance}}
    else:
        p["guid"] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["txt", 0]}}  # placeholder chain
    redux_chain(p, style_img, strength, stype)
    if not use_guidance:  # schnell: wire StyleModelApply straight to text encode
        p["sma"]["inputs"]["conditioning"] = ["txt", 0]
        del p["guid"]
    model_src = ["u", 0]
    if use_msf:
        p["msf"] = {"class_type": "ModelSamplingFlux", "inputs": {
            "model": ["u", 0], "max_shift": 1.15, "base_shift": 0.5, "width": w, "height": h}}
        model_src = ["msf", 0]
    p["noise"] = {"class_type": "RandomNoise", "inputs": {"noise_seed": SEED}}
    p["guider"] = {"class_type": "BasicGuider", "inputs": {"model": model_src, "conditioning": ["sma", 0]}}
    p["samp"] = {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler"}}
    p["sched"] = {"class_type": "BasicScheduler", "inputs": {
        "model": model_src, "scheduler": "simple", "steps": steps, "denoise": 1.0}}
    p["lat"] = {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}}
    p["sca"] = {"class_type": "SamplerCustomAdvanced", "inputs": {
        "noise": ["noise", 0], "guider": ["guider", 0], "sampler": ["samp", 0],
        "sigmas": ["sched", 0], "latent_image": ["lat", 0]}}
    save(p, ["sca", 0], f"{prefix}/{tag}")
    return p

def pairs_2img():
    combos = [(r, t) for r in R for t in T]                       # 24 r x t
    combos += [(a, b) for a in R for b in R if a != b]            # 12 ordered r x r
    return combos

def main():
    queue = []  # (workflow_dir, tag, prompt)
    for r, t in pairs_2img():
        tag = f"{r.split('.')[0]}x{t.split('.')[0]}"
        queue.append(("flux-redux-style-composition", tag, wf_style_composition(r, t, tag)))
    for r, t in pairs_2img():
        tag = f"{r.split('.')[0]}x{t.split('.')[0]}"
        queue.append(("flux-redux-style-composition-bfl-lora", tag, wf_bfl_lora(r, t, tag)))
    for r in R:
        tag = r.split(".")[0]
        queue.append(("flux-redux-fal-dev", tag, wf_single(r, tag)))
        queue.append(("flux-redux-prompt", tag, wf_single(
            r, tag, prompt_text="a person walking through a rainy city street at night, cinematic lighting",
            strength=0.5, stype="attn_bias", w=1024, h=1024, prefix="flux-redux-prompt")))
        queue.append(("flux-redux-schnell", tag, wf_single(
            r, tag, unet="flux1-schnell.safetensors", steps=4, w=1024, h=1024,
            prefix="flux-redux-schnell", use_guidance=False, use_msf=False)))

    ids = {}
    for wf, tag, prompt in queue:
        pid = post(prompt)
        ids[pid] = (wf, tag)
        print(f"queued {wf}/{tag} -> {pid}", flush=True)

    # poll history until all done
    pending = set(ids)
    errors = []
    while pending:
        time.sleep(10)
        for pid in list(pending):
            try:
                with urllib.request.urlopen(f"{HOST}/history/{pid}") as r:
                    h = json.loads(r.read())
            except Exception:
                continue
            if pid not in h:
                continue
            st = h[pid].get("status", {})
            if st.get("completed"):
                pending.discard(pid)
                print(f"done {ids[pid][0]}/{ids[pid][1]} ({len(ids)-len(pending)}/{len(ids)})", flush=True)
            elif st.get("status_str") == "error":
                pending.discard(pid)
                errors.append(ids[pid])
                msgs = [m for m in st.get("messages", []) if m[0] == "execution_error"]
                detail = msgs[-1][1].get("exception_message", "?") if msgs else "?"
                print(f"ERROR {ids[pid][0]}/{ids[pid][1]}: {detail[:300]}", flush=True)

    # collect outputs
    for wf in set(w for w, _, _ in queue):
        os.makedirs(os.path.join(OUT_ROOT, wf), exist_ok=True)
        src_dir = os.path.join(COMFY_OUT, wf)
        if not os.path.isdir(src_dir):
            continue
        for f in sorted(os.listdir(src_dir)):
            tag = f.split("_")[0]
            shutil.copy(os.path.join(src_dir, f), os.path.join(OUT_ROOT, wf, f"{tag}.png"))
    print("COLLECTED OUTPUTS", flush=True)
    for wf in sorted(set(w for w, _, _ in queue)):
        n = len(os.listdir(os.path.join(OUT_ROOT, wf))) if os.path.isdir(os.path.join(OUT_ROOT, wf)) else 0
        print(f"  {wf}: {n} images", flush=True)
    if errors:
        print("RUN ERRORS:", errors, flush=True)
        sys.exit(1)
    print("MATRIX COMPLETE", flush=True)

if __name__ == "__main__":
    main()