"""4-step MOSS-Image-SFT with the rapid-moss adapter (data-space sampler, no CFG).

    pip install "diffusers>=0.38" "transformers>=5.3" accelerate safetensors huggingface_hub pillow
    python inference.py --prompt "an anime girl in a yellow raincoat at a rainy bus stop" --out out.png

The base stays bf16; the LoRA is applied unmerged in fp32 (its delta is smaller than bf16 rounding),
and the step embedder + modulation biases are replaced by fp32 tensors. Sampler:
    x0 = 0; for r in range(steps): x = z + Q x0; x0 = z - transformer(x, codes[r], text)
with one fixed noise z, no noise re-injection and no timestep schedule.
"""
import argparse
import os

import torch
from diffusers import DiffusionPipeline
from huggingface_hub import snapshot_download
from safetensors.torch import load_file


def load(base, adapter, device):
    custom = os.path.join(base, 'pipeline_moss_image.py') if os.path.isdir(base) else 'pipeline_moss_image'
    pipe = DiffusionPipeline.from_pretrained(base, custom_pipeline=custom, trust_remote_code=True,
                                             torch_dtype={'default': torch.bfloat16, 'vae': torch.float32})
    pipe.to(device)
    tr = pipe.transformer
    if not os.path.isdir(adapter):
        adapter = snapshot_download(adapter, allow_patterns=['*.safetensors'])
    for name in ('time_embedding', 'shared_block_modulation', 'video_head_modulation'):
        getattr(tr, name).float()
    time_path = load_file(os.path.join(adapter, 'time_embedding.safetensors'), device=str(device))
    result = tr.load_state_dict(time_path, strict=False)
    assert not result.unexpected_keys, result.unexpected_keys
    lora = load_file(os.path.join(adapter, 'lora.safetensors'), device=str(device))
    for i, layer in enumerate(tr.layers):
        for name, lin in layer.proj.items():
            key = f'diffusion_model.layers.{i}.proj.{name}'
            down, up = lora[key + '.lora_down.weight'], lora[key + '.lora_up.weight']
            scale = float(lora[key + '.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))
    return pipe, ds['latent_mix.Q'], ds['step_codes']


@torch.no_grad()
def generate(pipe, Q, codes, prompt, width, height, steps, seed):
    device = pipe._execution_device
    emb, mask = pipe.encode_prompt(prompt)
    z = torch.randn(1, 128, height // 16, width // 16, generator=torch.Generator('cpu').manual_seed(seed)).to(device)
    x0 = torch.zeros_like(z)
    for r in range(steps):
        x = z + torch.einsum('oc,bchw->bohw', Q, x0)
        x0 = z - pipe.transformer(x, codes[r], emb, mask).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, h, w = lat.shape
    lat = lat.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
    image = vae.decode(lat, return_dict=False)[0].clamp(-1, 1)
    return pipe.image_processor.postprocess(image, output_type='pil')[0]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prompt', required=True, help='natural language, tags or a MOSS JSON layout string')
    ap.add_argument('--width', type=int, default=1024)
    ap.add_argument('--height', type=int, default=1024)
    ap.add_argument('--steps', type=int, default=4, help='1-4; each prefix of the 4 steps is trained')
    ap.add_argument('--seed', type=int, default=42)
    ap.add_argument('--out', default='out.png')
    ap.add_argument('--base', default='OpenMOSS-Team/MOSS-Image-SFT')
    ap.add_argument('--adapter', default='darask0/rapid-moss')
    ap.add_argument('--device', default='cuda')
    args = ap.parse_args()
    if args.width % 16 or args.height % 16:
        raise SystemExit('width and height must be multiples of 16')
    pipe, Q, codes = load(args.base, args.adapter, torch.device(args.device))
    if not 1 <= args.steps <= len(codes):
        raise SystemExit(f'steps must be 1..{len(codes)}')
    generate(pipe, Q, codes, args.prompt, args.width, args.height, args.steps, args.seed).save(args.out)
    print(args.out)


if __name__ == '__main__':
    main()