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

    # ---- conditioning via ai-toolkit (same as training/demo) ----
    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()

    # ---- release pipeline + adapter + ref-conditioning patch ----
    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()