| """2-step instruction editing with the akhaliq IDM v2 adapter (armV5 folded in). |
| |
| pip install "diffusers>=0.38" "transformers>=5.3" accelerate safetensors huggingface_hub pillow |
| # plus ai-toolkit (for image-conditioned prompt embeds, the exact conditioning used in training): |
| # git clone https://github.com/ostris/ai-toolkit && pip install -r ai-toolkit/requirements.txt |
| python inference_edit.py --image input.jpg --prompt "make both people laugh" --steps 2 --out out.png |
| |
| Stack: ai-toolkit prompt embeds (control_images) + reference latents (Kontext-style rope axis) |
| + v2 adapter (fp32 LoRA hooks, replaced time path) + data-space sampler, no CFG. |
| """ |
|
|
| import argparse |
| import os |
| import sys |
|
|
| import torch |
| from safetensors.torch import load_file |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--image", required=True) |
| ap.add_argument("--prompt", required=True) |
| ap.add_argument("--steps", type=int, default=2) |
| ap.add_argument("--seed", type=int, default=42) |
| ap.add_argument("--out", default="edit_out.png") |
| ap.add_argument("--base", default="OpenMOSS-Team/MOSS-Image-SFT") |
| ap.add_argument("--adapter", default="akhaliq/MOSS-Image-IDM-2-Steps") |
| ap.add_argument("--adapter_dir", default="adapter_v2") |
| ap.add_argument("--device", default="cuda") |
| args = ap.parse_args() |
|
|
| from diffusers import DiffusionPipeline |
| from huggingface_hub import hf_hub_download, snapshot_download |
| from torchvision.transforms.functional import to_tensor |
| from PIL import Image |
|
|
| device = torch.device(args.device) |
|
|
| |
| from extensions_built_in.diffusion_models.moss_image.moss_image import MossImageModel |
| from toolkit.config_modules import ModelConfig |
|
|
| mc = ModelConfig(name_or_path=args.base, 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() |
|
|
| img = Image.open(args.image).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([args.prompt], control_images=[[img_t]]) |
| ref = model.encode_images([img_t])[0] |
| del model |
| torch.cuda.empty_cache() |
|
|
| |
| pipe = DiffusionPipeline.from_pretrained(args.base, custom_pipeline="pipeline_moss_image", |
| trust_remote_code=True, |
| torch_dtype={"default": torch.bfloat16, "vae": torch.float32}) |
| pipe.to(device) |
|
|
| 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) |
|
|
| adapter = snapshot_download(args.adapter, allow_patterns=[f"{args.adapter_dir}/*.safetensors"]) |
| d = os.path.join(adapter, args.adapter_dir) |
| tr = pipe.transformer |
| for name in ("time_embedding", "shared_block_modulation", "video_head_modulation"): |
| getattr(tr, name).float() |
| tp = load_file(os.path.join(d, "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(d, "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(d, "dataspace.safetensors"), device=str(device)) |
| Q, codes = ds["latent_mix.Q"], ds["step_codes"] |
|
|
| emb = next(t for t in (cond[k][0] for k in cond.keys()) if torch.is_tensor(t) and t.dim() >= 2 and t.shape[-1] == 4096) |
| if emb.dim() == 2: |
| emb = emb[None] |
| emb = emb.to(device) |
| mask = torch.ones(emb.shape[:2], dtype=torch.bool, device=device) |
| ref_in = ref[None] if ref.dim() == 3 else ref |
|
|
| @torch.no_grad() |
| def generate(): |
| z = torch.randn(1, 128, h // 16, w // 16, |
| generator=torch.Generator("cpu").manual_seed(args.seed)).to(device) |
| x0 = torch.zeros_like(z) |
| for r in range(args.steps): |
| x = z + torch.einsum("oc,bchw->bohw", Q, x0) |
| x0 = z - pipe.transformer(x, codes[r], emb, mask, ref_latents=[ref_in]).sample.float() |
| 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] |
|
|
| generate().save(args.out) |
| print(args.out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
