Reference implementation: 2-step + armV5 editing at CFG 2 (ai-toolkit conditioning)

33de70d verified about 7 hours ago

13.5 kB

"""Experiment: can the 2-step IDM adapter drive MOSS-Image-Edit-HQ (armV5)?
Grid (linoyts style): input | ours·2 steps | base·50 steps, for 4 edit cases.
- Base column: ai-toolkit inference stack, armV5, 50 steps CFG 3 (exact demo settings).
- Ours column: same ai-toolkit conditioning (cond embeds + ref latents), release
  transformer + armV5 merged + 2-step IDM adapter (fp32 LoRA hooks, time path),
  data-space sampler, 2 iterations, no CFG.
Uploads to akhaliq/MOSS-Image-IDM-LoRA eval/edit_2step/.
"""

import os
import subprocess
import sys
import zipfile
import urllib.request

from huggingface_hub import hf_hub_download, snapshot_download, upload_folder

OUT = "/workspace/edit_out"
DEST = "akhaliq/MOSS-Image-IDM-LoRA"
ARM_REPO = "akhaliq/MOSS-Image-Edit-HQ"
ARM_FILE = "weights/armV5/armV5.safetensors"

# input images are generated with the base model first (self-contained), then edited
CASES = [
    {
        "gen": "A pencil sketch of a lighthouse on a rocky cliff, rough pencil strokes, white paper",
        "edit": "turn this sketch into a colorful realistic photograph",
        "label": "Sketch to photo",
    },
    {
        "gen": "Street photo of a young woman with short straight black hair wearing a denim jacket, city crosswalk background, natural daylight",
        "edit": "change her hair to long wavy auburn hair, keep everything else the same",
        "label": "Hairstyle edit",
    },
    {
        "gen": "Two students in white polo shirts standing on a school sports field, neutral expressions, snapshot photo",
        "edit": "make both of them laugh with big open-mouth smiles, keep everything else the same",
        "label": "Expression edit",
    },
    {
        "gen": "An old black and white photograph of an elderly man with a pipe sitting in an armchair, 1950s studio portrait, film grain",
        "edit": "colorize and restore this old black and white photograph",
        "label": "Photo restoration",
    },
]


def sh(cmd, **kw):
    print(f"+ {' '.join(map(str, cmd))}", flush=True)
    subprocess.run(cmd, check=True, **kw)


def main():
    os.makedirs("/workspace", exist_ok=True)
    os.makedirs(OUT, exist_ok=True)
    os.chdir("/workspace")
    sh([sys.executable, "-m", "pip", "install", "-q",
        "torch==2.15.0", "torchvision", "torchaudio", "diffusers>=0.38", "transformers>=5.3",
        "accelerate", "safetensors", "pillow", "numpy<2"])

    # ai-toolkit @ the demo's pinned commit
    url = "https://github.com/ostris/ai-toolkit/archive/14f8f95.zip"
    urllib.request.urlretrieve(url, "/workspace/aitk.zip")
    with zipfile.ZipFile("/workspace/aitk.zip") as zf:
        zf.extractall("/workspace")
    aitk = next(p for p in os.listdir("/workspace") if p.startswith("ai-toolkit-"))
    os.rename(os.path.join("/workspace", aitk), "/workspace/ai-toolkit")
    sh([sys.executable, "-m", "pip", "install", "-q", "--no-cache-dir", "-r", "requirements.txt"],
       cwd="/workspace/ai-toolkit")
    sys.path.insert(0, "/workspace/ai-toolkit")

    import torch
    from safetensors.torch import load_file
    from torchvision.transforms.functional import to_tensor
    from PIL import Image

    device = torch.device("cuda")

    # ---------- 0. input images via base T2I (release pipeline) ----------
    from diffusers import DiffusionPipeline
    pipe = DiffusionPipeline.from_pretrained(
        "OpenMOSS-Team/MOSS-Image-SFT", custom_pipeline="pipeline_moss_image",
        trust_remote_code=True, torch_dtype={"default": torch.bfloat16, "vae": torch.float32})
    pipe.to(device)
    inputs = []
    for i, case in enumerate(CASES):
        img = pipe(case["gen"], height=1024, width=1024, num_inference_steps=50,
                   guidance_scale=3.0, generator=torch.Generator("cpu").manual_seed(100 + i)).images[0]
        img.save(os.path.join(OUT, f"{i}_input.jpg"), quality=92)
        inputs.append(img)
        print(f"input {i}: {case['label']}", flush=True)
    clean_transformer_state = {k: v.detach().clone() for k, v in pipe.transformer.state_dict().items()}

    # ---------- 1. ai-toolkit stack: conditioning + base edits ----------
    from extensions_built_in.diffusion_models.moss_image.moss_image import MossImageModel
    from extensions_built_in.diffusion_models.moss_image.src.pipeline import MossImagePipeline
    from toolkit.config_modules import ModelConfig

    mc = ModelConfig(name_or_path="OpenMOSS-Team/MOSS-Image-SFT", arch="moss_image",
                     quantize=False, quantize_te=False, model_kwargs={"edit": True, "kv_cache": True})
    model = MossImageModel(device="cuda", model_config=mc)
    model.torch_dtype = torch.bfloat16
    model.load_model()

    def merge_lora(transformer, sd):
        """Merge an ai-toolkit/kohya LoRA (alpha==rank assumed, per the demo's merge)."""
        deltas = {}
        for k, v in sd.items():
            if ".lora_" not in k:
                continue
            mod, kind = k.replace("diffusion_model.", "").replace("transformer.", "").rsplit(".lora_", 1)
            kind = {"down.weight": "A.weight", "up.weight": "B.weight"}.get(kind, kind)
            deltas.setdefault(mod, {})[kind] = v.float()
        fused = {}
        for mod, pair in deltas.items():
            base, _, sub = mod.rpartition(".proj.")
            if sub in ("q", "k", "v"):
                fused.setdefault(base + ".proj.qkv", {})[sub] = pair
            else:
                fused[mod] = pair
        state = transformer.state_dict()
        n = 0
        for mod, pair in fused.items():
            if mod.endswith(".proj.qkv") and mod + ".weight" not in state:
                # release transformer: separate q/k/v linears
                base2 = mod[: -len(".qkv")]
                for sub, p in pair.items():
                    wname = f"{base2}.{sub}.weight"
                    assert wname in state, wname
                    delta = p["B.weight"] @ p["A.weight"]
                    state[wname] = (state[wname].float() + delta.to(state[wname].device)).to(state[wname].dtype)
                    n += 1
                continue
            wname = mod + ".weight"
            assert wname in state, wname
            if set(pair) == {"q", "k", "v"}:
                down = torch.cat([pair[p]["A.weight"] for p in ("q", "k", "v")], dim=0)
                up = torch.block_diag(*[pair[p]["B.weight"] for p in ("q", "k", "v")])
                delta = up @ down
            else:
                delta = pair["B.weight"] @ pair["A.weight"]
            state[wname] = (state[wname].float() + delta.to(state[wname].device)).to(state[wname].dtype)
            n += 1
        transformer.load_state_dict(state)
        print(f"merged {n} lora modules", flush=True)

    arm_sd = load_file(hf_hub_download(ARM_REPO, ARM_FILE))
    merge_lora(model.transformer, arm_sd)
    model.pipeline = MossImagePipeline(model)

    conds = []
    for i, case in enumerate(CASES):
        img = inputs[i].convert("RGB")
        w, h = img.size
        scale = 1024.0 / max(w, h)
        if scale < 1.0:
            w, h = int(w * scale), int(h * scale)
        w, h = max(16, w - w % 16), max(16, h - h % 16)
        if (w, h) != img.size:
            img = img.resize((w, h))
        img_t = (to_tensor(img) * 2.0 - 1.0).to(torch.bfloat16)
        cond = model.get_prompt_embeds([case["edit"]], control_images=[[img_t]])
        uncond = model.get_prompt_embeds([""], control_images=[[img_t]])
        ref = model.encode_images([img_t])[0]
        # base column: exact demo settings
        images = model.pipeline(cond, uncond, height=h, width=w, num_inference_steps=50,
                                guidance_scale=3.0, ref_latents=[[ref]], isolate_refs=True,
                                generator=torch.Generator("cpu").manual_seed(42))
        images[0].save(os.path.join(OUT, f"{i}_base50.jpg"), quality=92)
        conds.append({"cond": cond, "uncond": uncond, "ref": ref, "w": w, "h": h})
        print(f"base50 edit {i}: {case['label']}", flush=True)

    del model
    torch.cuda.empty_cache()

    # ---------- 2. ours: release transformer + armV5 + 2-step IDM adapter ----------
    # reload armV5 onto the (still loaded) release transformer
    merge_lora(pipe.transformer, arm_sd)

    # edit conditioning patch (vendored mechanism, downloaded from the demo Space)
    me_path = hf_hub_download("akhaliq/moss-image-edit-hq-demo", "moss_edit.py", repo_type="space")
    sys.path.insert(0, os.path.dirname(me_path))
    import moss_edit
    moss_edit.patch(pipe.transformer)

    # 2-step adapter: time path + fp32 LoRA hooks + Q/codes (from repo root)
    adapter = snapshot_download(DEST, allow_patterns=["lora.safetensors", "time_embedding.safetensors", "dataspace.safetensors"])
    tr = pipe.transformer
    for name in ("time_embedding", "shared_block_modulation", "video_head_modulation"):
        getattr(tr, name).float()
    tp = load_file(os.path.join(adapter, "time_embedding.safetensors"), device=str(device))
    res = tr.load_state_dict(tp, strict=False)
    assert not res.unexpected_keys, res.unexpected_keys
    lora = load_file(os.path.join(adapter, "lora.safetensors"), device=str(device))
    for li, layer in enumerate(tr.layers):
        for name, lin in layer.proj.items():
            k = f"diffusion_model.layers.{li}.proj.{name}"
            down, up = lora[k + ".lora_down.weight"], lora[k + ".lora_up.weight"]
            scale = float(lora[k + ".alpha"]) / down.shape[0]

            def hook(module, inputs, output, down=down, up=up, scale=scale):
                return output + (scale * ((inputs[0].float() @ down.t()) @ up.t())).to(output.dtype)

            lin.register_forward_hook(hook)
    ds = load_file(os.path.join(adapter, "dataspace.safetensors"), device=str(device))
    Q, codes = ds["latent_mix.Q"], ds["step_codes"]

    def unpack(cond_emb):
        emb = cond_emb
        if not torch.is_tensor(emb):
            cands = [emb[k][0] for k in emb.keys()]
            emb = next(t for t in cands if torch.is_tensor(t) and t.dim() >= 2 and t.shape[-1] == 4096)
        if emb.dim() == 2:
            emb = emb[None]
        return emb.to(device)

    @torch.no_grad()
    def rapid_edit(cond_emb, uncond_emb, ref, w, h, cfg=1.0, seed=42):
        emb = unpack(cond_emb)
        mask = torch.ones(emb.shape[:2], dtype=torch.bool, device=device)
        uncond = unpack(uncond_emb) if uncond_emb is not None else None
        umask = torch.ones(uncond.shape[:2], dtype=torch.bool, device=device) if uncond is not None else None
        z = torch.randn(1, 128, h // 16, w // 16, generator=torch.Generator("cpu").manual_seed(seed)).to(device)
        ref_in = ref[None] if ref.dim() == 3 else ref
        x0 = torch.zeros_like(z)
        for r in range(2):
            x = z + torch.einsum("oc,bchw->bohw", Q, x0)
            v_pos = pipe.transformer(x, codes[r], emb, mask, ref_latents=[ref_in]).sample.float()
            if cfg != 1.0 and uncond is not None:
                v_neg = pipe.transformer(x, codes[r], uncond, umask, ref_latents=[ref_in],
                                         skip_layers=(22,)).sample.float()
                v_pos = v_neg + cfg * (v_pos - v_neg)
            x0 = z - v_pos
        vae = pipe.vae
        lat = x0.to(vae.dtype)
        bn, eps = vae.bn, vae.config.batch_norm_eps
        lat = lat * (bn.running_var + eps).sqrt().view(1, -1, 1, 1) + bn.running_mean.view(1, -1, 1, 1)
        b, c, hh, ww = lat.shape
        lat = lat.reshape(b, c // 4, 2, 2, hh, ww).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * hh, 2 * ww)
        image = vae.decode(lat, return_dict=False)[0].clamp(-1, 1)
        return pipe.image_processor.postprocess(image, output_type="pil")[0]

    for i, case in enumerate(CASES):
        c = conds[i]
        img = rapid_edit(c["cond"], None, c["ref"], c["w"], c["h"], cfg=1.0)
        img.save(os.path.join(OUT, f"{i}_ours2.jpg"), quality=92)
        img = rapid_edit(c["cond"], c["uncond"], c["ref"], c["w"], c["h"], cfg=2.0)
        img.save(os.path.join(OUT, f"{i}_ours2cfg2.jpg"), quality=92)
        img = rapid_edit(c["cond"], c["uncond"], c["ref"], c["w"], c["h"], cfg=3.0)
        img.save(os.path.join(OUT, f"{i}_ours2cfg.jpg"), quality=92)
        print(f"ours 2-step edits (cfg1/2/3) {i}: {case['label']}", flush=True)

    # ---------- 3. grid: input | ours·2 steps | base·50 steps ----------
    from PIL import ImageDraw
    cell = 384
    label_h, head_h = 26, 26
    n = len(CASES)
    grid = Image.new("RGB", (cell * 5, (cell + label_h) * n + head_h), "white")
    draw = ImageDraw.Draw(grid)
    for c, title in enumerate(["input", "2 steps", "2 steps CFG 2", "2 steps CFG 3", "base · 50 steps"]):
        draw.text((c * cell + cell // 2 - 40, 6), title, fill="black")
    for r, case in enumerate(CASES):
        draw.text((6, head_h + r * (cell + label_h) + cell + 6), case["label"], fill="black")
        for c, tag in enumerate(["input", "ours2", "ours2cfg2", "ours2cfg", "base50"]):
            img = Image.open(os.path.join(OUT, f"{r}_{tag}.jpg"))
            img.thumbnail((cell, cell))
            grid.paste(img, (c * cell + (cell - img.width) // 2, head_h + r * (cell + label_h)))
    grid.save(os.path.join(OUT, "grid_edit_2step.jpg"), quality=90)

    upload_folder(repo_id=DEST, folder_path=OUT, path_in_repo="eval/edit_2step",
                  commit_message="Edit experiment: armV5 + 2-step IDM adapter vs armV5 @ 50 steps")
    print("EDIT_EVAL_UPLOADED", flush=True)


if __name__ == "__main__":
    main()

	"""Experiment: can the 2-step IDM adapter drive MOSS-Image-Edit-HQ (armV5)?

	Grid (linoyts style): input \| ours·2 steps \| base·50 steps, for 4 edit cases.
	- Base column: ai-toolkit inference stack, armV5, 50 steps CFG 3 (exact demo settings).
	- Ours column: same ai-toolkit conditioning (cond embeds + ref latents), release
	transformer + armV5 merged + 2-step IDM adapter (fp32 LoRA hooks, time path),
	data-space sampler, 2 iterations, no CFG.
	Uploads to akhaliq/MOSS-Image-IDM-LoRA eval/edit_2step/.
	"""

	import os
	import subprocess
	import sys
	import zipfile
	import urllib.request

	from huggingface_hub import hf_hub_download, snapshot_download, upload_folder

	OUT = "/workspace/edit_out"
	DEST = "akhaliq/MOSS-Image-IDM-LoRA"
	ARM_REPO = "akhaliq/MOSS-Image-Edit-HQ"
	ARM_FILE = "weights/armV5/armV5.safetensors"

	# input images are generated with the base model first (self-contained), then edited
	CASES = [
	{
	"gen": "A pencil sketch of a lighthouse on a rocky cliff, rough pencil strokes, white paper",
	"edit": "turn this sketch into a colorful realistic photograph",
	"label": "Sketch to photo",
	},
	{
	"gen": "Street photo of a young woman with short straight black hair wearing a denim jacket, city crosswalk background, natural daylight",
	"edit": "change her hair to long wavy auburn hair, keep everything else the same",
	"label": "Hairstyle edit",
	},
	{
	"gen": "Two students in white polo shirts standing on a school sports field, neutral expressions, snapshot photo",
	"edit": "make both of them laugh with big open-mouth smiles, keep everything else the same",
	"label": "Expression edit",
	},
	{
	"gen": "An old black and white photograph of an elderly man with a pipe sitting in an armchair, 1950s studio portrait, film grain",
	"edit": "colorize and restore this old black and white photograph",
	"label": "Photo restoration",
	},
	]


	def sh(cmd, **kw):
	print(f"+ {' '.join(map(str, cmd))}", flush=True)
	subprocess.run(cmd, check=True, **kw)


	def main():
	os.makedirs("/workspace", exist_ok=True)
	os.makedirs(OUT, exist_ok=True)
	os.chdir("/workspace")
	sh([sys.executable, "-m", "pip", "install", "-q",
	"torch==2.15.0", "torchvision", "torchaudio", "diffusers>=0.38", "transformers>=5.3",
	"accelerate", "safetensors", "pillow", "numpy<2"])

	# ai-toolkit @ the demo's pinned commit
	url = "https://github.com/ostris/ai-toolkit/archive/14f8f95.zip"
	urllib.request.urlretrieve(url, "/workspace/aitk.zip")
	with zipfile.ZipFile("/workspace/aitk.zip") as zf:
	zf.extractall("/workspace")
	aitk = next(p for p in os.listdir("/workspace") if p.startswith("ai-toolkit-"))
	os.rename(os.path.join("/workspace", aitk), "/workspace/ai-toolkit")
	sh([sys.executable, "-m", "pip", "install", "-q", "--no-cache-dir", "-r", "requirements.txt"],
	cwd="/workspace/ai-toolkit")
	sys.path.insert(0, "/workspace/ai-toolkit")

	import torch
	from safetensors.torch import load_file
	from torchvision.transforms.functional import to_tensor
	from PIL import Image

	device = torch.device("cuda")

	# ---------- 0. input images via base T2I (release pipeline) ----------
	from diffusers import DiffusionPipeline
	pipe = DiffusionPipeline.from_pretrained(
	"OpenMOSS-Team/MOSS-Image-SFT", custom_pipeline="pipeline_moss_image",
	trust_remote_code=True, torch_dtype={"default": torch.bfloat16, "vae": torch.float32})
	pipe.to(device)
	inputs = []
	for i, case in enumerate(CASES):
	img = pipe(case["gen"], height=1024, width=1024, num_inference_steps=50,
	guidance_scale=3.0, generator=torch.Generator("cpu").manual_seed(100 + i)).images[0]
	img.save(os.path.join(OUT, f"{i}_input.jpg"), quality=92)
	inputs.append(img)
	print(f"input {i}: {case['label']}", flush=True)
	clean_transformer_state = {k: v.detach().clone() for k, v in pipe.transformer.state_dict().items()}

	# ---------- 1. ai-toolkit stack: conditioning + base edits ----------
	from extensions_built_in.diffusion_models.moss_image.moss_image import MossImageModel
	from extensions_built_in.diffusion_models.moss_image.src.pipeline import MossImagePipeline
	from toolkit.config_modules import ModelConfig

	mc = ModelConfig(name_or_path="OpenMOSS-Team/MOSS-Image-SFT", arch="moss_image",
	quantize=False, quantize_te=False, model_kwargs={"edit": True, "kv_cache": True})
	model = MossImageModel(device="cuda", model_config=mc)
	model.torch_dtype = torch.bfloat16
	model.load_model()

	def merge_lora(transformer, sd):
	"""Merge an ai-toolkit/kohya LoRA (alpha==rank assumed, per the demo's merge)."""
	deltas = {}
	for k, v in sd.items():
	if ".lora_" not in k:
	continue
	mod, kind = k.replace("diffusion_model.", "").replace("transformer.", "").rsplit(".lora_", 1)
	kind = {"down.weight": "A.weight", "up.weight": "B.weight"}.get(kind, kind)
	deltas.setdefault(mod, {})[kind] = v.float()
	fused = {}
	for mod, pair in deltas.items():
	base, _, sub = mod.rpartition(".proj.")
	if sub in ("q", "k", "v"):
	fused.setdefault(base + ".proj.qkv", {})[sub] = pair
	else:
	fused[mod] = pair
	state = transformer.state_dict()
	n = 0
	for mod, pair in fused.items():
	if mod.endswith(".proj.qkv") and mod + ".weight" not in state:
	# release transformer: separate q/k/v linears
	base2 = mod[: -len(".qkv")]
	for sub, p in pair.items():
	wname = f"{base2}.{sub}.weight"
	assert wname in state, wname
	delta = p["B.weight"] @ p["A.weight"]
	state[wname] = (state[wname].float() + delta.to(state[wname].device)).to(state[wname].dtype)
	n += 1
	continue
	wname = mod + ".weight"
	assert wname in state, wname
	if set(pair) == {"q", "k", "v"}:
	down = torch.cat([pair[p]["A.weight"] for p in ("q", "k", "v")], dim=0)
	up = torch.block_diag(*[pair[p]["B.weight"] for p in ("q", "k", "v")])
	delta = up @ down
	else:
	delta = pair["B.weight"] @ pair["A.weight"]
	state[wname] = (state[wname].float() + delta.to(state[wname].device)).to(state[wname].dtype)
	n += 1
	transformer.load_state_dict(state)
	print(f"merged {n} lora modules", flush=True)

	arm_sd = load_file(hf_hub_download(ARM_REPO, ARM_FILE))
	merge_lora(model.transformer, arm_sd)
	model.pipeline = MossImagePipeline(model)

	conds = []
	for i, case in enumerate(CASES):
	img = inputs[i].convert("RGB")
	w, h = img.size
	scale = 1024.0 / max(w, h)
	if scale < 1.0:
	w, h = int(w * scale), int(h * scale)
	w, h = max(16, w - w % 16), max(16, h - h % 16)
	if (w, h) != img.size:
	img = img.resize((w, h))
	img_t = (to_tensor(img) * 2.0 - 1.0).to(torch.bfloat16)
	cond = model.get_prompt_embeds([case["edit"]], control_images=[[img_t]])
	uncond = model.get_prompt_embeds([""], control_images=[[img_t]])
	ref = model.encode_images([img_t])[0]
	# base column: exact demo settings
	images = model.pipeline(cond, uncond, height=h, width=w, num_inference_steps=50,
	guidance_scale=3.0, ref_latents=[[ref]], isolate_refs=True,
	generator=torch.Generator("cpu").manual_seed(42))
	images[0].save(os.path.join(OUT, f"{i}_base50.jpg"), quality=92)
	conds.append({"cond": cond, "uncond": uncond, "ref": ref, "w": w, "h": h})
	print(f"base50 edit {i}: {case['label']}", flush=True)

	del model
	torch.cuda.empty_cache()

	# ---------- 2. ours: release transformer + armV5 + 2-step IDM adapter ----------
	# reload armV5 onto the (still loaded) release transformer
	merge_lora(pipe.transformer, arm_sd)

	# edit conditioning patch (vendored mechanism, downloaded from the demo Space)
	me_path = hf_hub_download("akhaliq/moss-image-edit-hq-demo", "moss_edit.py", repo_type="space")
	sys.path.insert(0, os.path.dirname(me_path))
	import moss_edit
	moss_edit.patch(pipe.transformer)

	# 2-step adapter: time path + fp32 LoRA hooks + Q/codes (from repo root)
	adapter = snapshot_download(DEST, allow_patterns=["lora.safetensors", "time_embedding.safetensors", "dataspace.safetensors"])
	tr = pipe.transformer
	for name in ("time_embedding", "shared_block_modulation", "video_head_modulation"):
	getattr(tr, name).float()
	tp = load_file(os.path.join(adapter, "time_embedding.safetensors"), device=str(device))
	res = tr.load_state_dict(tp, strict=False)
	assert not res.unexpected_keys, res.unexpected_keys
	lora = load_file(os.path.join(adapter, "lora.safetensors"), device=str(device))
	for li, layer in enumerate(tr.layers):
	for name, lin in layer.proj.items():
	k = f"diffusion_model.layers.{li}.proj.{name}"
	down, up = lora[k + ".lora_down.weight"], lora[k + ".lora_up.weight"]
	scale = float(lora[k + ".alpha"]) / down.shape[0]

	def hook(module, inputs, output, down=down, up=up, scale=scale):
	return output + (scale * ((inputs[0].float() @ down.t()) @ up.t())).to(output.dtype)

	lin.register_forward_hook(hook)
	ds = load_file(os.path.join(adapter, "dataspace.safetensors"), device=str(device))
	Q, codes = ds["latent_mix.Q"], ds["step_codes"]

	def unpack(cond_emb):
	emb = cond_emb
	if not torch.is_tensor(emb):
	cands = [emb[k][0] for k in emb.keys()]
	emb = next(t for t in cands if torch.is_tensor(t) and t.dim() >= 2 and t.shape[-1] == 4096)
	if emb.dim() == 2:
	emb = emb[None]
	return emb.to(device)

	@torch.no_grad()
	def rapid_edit(cond_emb, uncond_emb, ref, w, h, cfg=1.0, seed=42):
	emb = unpack(cond_emb)
	mask = torch.ones(emb.shape[:2], dtype=torch.bool, device=device)
	uncond = unpack(uncond_emb) if uncond_emb is not None else None
	umask = torch.ones(uncond.shape[:2], dtype=torch.bool, device=device) if uncond is not None else None
	z = torch.randn(1, 128, h // 16, w // 16, generator=torch.Generator("cpu").manual_seed(seed)).to(device)
	ref_in = ref[None] if ref.dim() == 3 else ref
	x0 = torch.zeros_like(z)
	for r in range(2):
	x = z + torch.einsum("oc,bchw->bohw", Q, x0)
	v_pos = pipe.transformer(x, codes[r], emb, mask, ref_latents=[ref_in]).sample.float()
	if cfg != 1.0 and uncond is not None:
	v_neg = pipe.transformer(x, codes[r], uncond, umask, ref_latents=[ref_in],
	skip_layers=(22,)).sample.float()
	v_pos = v_neg + cfg * (v_pos - v_neg)
	x0 = z - v_pos
	vae = pipe.vae
	lat = x0.to(vae.dtype)
	bn, eps = vae.bn, vae.config.batch_norm_eps
	lat = lat * (bn.running_var + eps).sqrt().view(1, -1, 1, 1) + bn.running_mean.view(1, -1, 1, 1)
	b, c, hh, ww = lat.shape
	lat = lat.reshape(b, c // 4, 2, 2, hh, ww).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * hh, 2 * ww)
	image = vae.decode(lat, return_dict=False)[0].clamp(-1, 1)
	return pipe.image_processor.postprocess(image, output_type="pil")[0]

	for i, case in enumerate(CASES):
	c = conds[i]
	img = rapid_edit(c["cond"], None, c["ref"], c["w"], c["h"], cfg=1.0)
	img.save(os.path.join(OUT, f"{i}_ours2.jpg"), quality=92)
	img = rapid_edit(c["cond"], c["uncond"], c["ref"], c["w"], c["h"], cfg=2.0)
	img.save(os.path.join(OUT, f"{i}_ours2cfg2.jpg"), quality=92)
	img = rapid_edit(c["cond"], c["uncond"], c["ref"], c["w"], c["h"], cfg=3.0)
	img.save(os.path.join(OUT, f"{i}_ours2cfg.jpg"), quality=92)
	print(f"ours 2-step edits (cfg1/2/3) {i}: {case['label']}", flush=True)

	# ---------- 3. grid: input \| ours·2 steps \| base·50 steps ----------
	from PIL import ImageDraw
	cell = 384
	label_h, head_h = 26, 26
	n = len(CASES)
	grid = Image.new("RGB", (cell * 5, (cell + label_h) * n + head_h), "white")
	draw = ImageDraw.Draw(grid)
	for c, title in enumerate(["input", "2 steps", "2 steps CFG 2", "2 steps CFG 3", "base · 50 steps"]):
	draw.text((c * cell + cell // 2 - 40, 6), title, fill="black")
	for r, case in enumerate(CASES):
	draw.text((6, head_h + r * (cell + label_h) + cell + 6), case["label"], fill="black")
	for c, tag in enumerate(["input", "ours2", "ours2cfg2", "ours2cfg", "base50"]):
	img = Image.open(os.path.join(OUT, f"{r}_{tag}.jpg"))
	img.thumbnail((cell, cell))
	grid.paste(img, (c * cell + (cell - img.width) // 2, head_h + r * (cell + label_h)))
	grid.save(os.path.join(OUT, "grid_edit_2step.jpg"), quality=90)

	upload_folder(repo_id=DEST, folder_path=OUT, path_in_repo="eval/edit_2step",
	commit_message="Edit experiment: armV5 + 2-step IDM adapter vs armV5 @ 50 steps")
	print("EDIT_EVAL_UPLOADED", flush=True)


	if __name__ == "__main__":
	main()