"""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()