#!/usr/bin/env python3
"""Run full MOSS-Image-SFT quantized WEIGHTS (zero LoRA).

Usage: use downloaded exact original SFT components plus a quantized
transformer directory produced by create_full_sft_checkpoints.py.
MOSS teacher/default is 50 denoising steps and CFG=3; no 4-step student.

Hardware requirements depend on the quantization format, library and device kernels.
"""
import argparse
import gc
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import sys
import time
import torch

def sha256(path):
    h=hashlib.sha256()
    with open(path,"rb") as f:
        for chunk in iter(lambda:f.read(4*1024*1024),b""):
            h.update(chunk)
    return h.hexdigest()

def pinned_loader(root):
    spec=importlib.util.spec_from_file_location("moss_sft_native_model",root/"pipeline_moss_image.py")
    module=importlib.util.module_from_spec(spec)
    sys.modules[spec.name]=module
    spec.loader.exec_module(module)
    return module

def check_artifact(variant):
    manifest=json.loads((variant/"QUANTIZATION_RECEIPT.json").read_text())
    done=json.loads((variant/"COMPLETE.json").read_text())
    verified = done.get("verified_all_disk_weights", done.get("reload_verified", False))
    if not verified or done["variant"]!=manifest["variant"]:
        raise RuntimeError("Incomplete or unverified quantized weight export")
    if manifest["adapter"] is not None or manifest["lora"] is not None or manifest["student"] is not None:
        raise RuntimeError("An adapter is present, expected full SFT only")
    for fn, digest in done["checkpoint_sha256"].items():
        f=variant/"transformer"/fn
        if not f.is_file() or sha256(f)!=digest:
            raise RuntimeError(f"Missing/modified full SFT checkpoint shard: {fn}")
    return manifest


def verify_pinned_source(source, expected_revision):
    """Verify original model identity for a verified local snapshot or an immutable HF cache snapshot.

    A bare directory with unverified model files must not be accepted. The
    official HF snapshot is identified by the exact commit in its cache path,
    not a mutable branch/ref name. A local model_receipt.json is also valid.
    """
    source=Path(source).resolve()
    receipt=source/"model_receipt.json"
    if receipt.is_file():
        data=json.loads(receipt.read_text())
        if data.get("revision")!=expected_revision or not data.get("complete"):
            raise RuntimeError("Original SFT local receipt revision/completeness mismatch")
    elif not (source.name==expected_revision and source.parent.name=="snapshots"):
        raise RuntimeError("Unverified original components. Use HF snapshot_download with exact revision "
                           +expected_revision+" or original complete local model_receipt.json")
    required=("pipeline_moss_image.py","text_encoder/config.json",
              "tokenizer/tokenizer_config.json","vae/config.json","scheduler/scheduler_config.json")
    missing=[part for part in required if not (source/part).is_file()]
    if missing:
        raise RuntimeError("Original SFT source components missing: "+", ".join(missing))
    return source


class StreamedQwenText:
    """BF16 Qwen text encoder streamed via Accelerate; no LoRA."""
    def __init__(self,root,device="cuda",cache_size=16,encoder_root=None):
        from transformers import AutoModel,AutoTokenizer
        from accelerate import cpu_offload
        from compact_components import inspect_qwen
        self.device=torch.device(device)
        self.tokenizer=AutoTokenizer.from_pretrained(root/"tokenizer",local_files_only=True)
        if encoder_root is None:
            # Identical original unquantized tokenizer/encoder for BF16 reference.
            enc=AutoModel.from_pretrained(root/"text_encoder",dtype=torch.bfloat16,
                                          local_files_only=True).eval()
            self.quantization="original_bf16"
        else:
            # This is real saved FP8WO tensor-subclass storage, not a BF16
            # encoder dynamically quantized when loading or a cached BF16 bank.
            encoder_root=Path(encoder_root).resolve()
            info=inspect_qwen(encoder_root,verify_sha=True)
            import torchao
            import torchao.quantization
            enc=AutoModel.from_pretrained(encoder_root/"text_encoder",
                dtype=torch.bfloat16,use_safetensors=False,local_files_only=True).eval()
            real=sum(type(m.weight).__name__=="AffineQuantizedTensor"
                     for n,m in enc.named_modules()
                     if n.startswith("language_model.") and isinstance(m,torch.nn.Linear))
            if real!=info["converted_linear_weights"]:
                raise RuntimeError(f"Packed Qwen was not actually rehydrated: {real}")
            self.quantization="original_qwen3.5_fp8_weight_only_248_linear"
        if hasattr(enc,"visual"):
            # Image input is not used by the MOSS-SFT text-only caption path.
            del enc.visual
            gc.collect()
        # Hold original CPU storage directly, not 15+ GiB of cloned duplicate
        # Qwen weights. Accelerate replaces module slots with meta after the
        # state_dict tensors have been captured by its CPU weights_map.
        weights={k:v.detach() for k,v in enc.state_dict().items()}
        cpu_offload(enc,execution_device=self.device,offload_buffers=True,
                    state_dict=weights)
        del weights
        self.enc=enc
        self.cache={}
        self.cache_size=cache_size

    @torch.no_grad()
    def encode(self,prompt):
        if prompt in self.cache:
            return self.cache[prompt]
        formatted=self.tokenizer.apply_chat_template(
            [{"role":"user","content":prompt}],tokenize=False,
            add_generation_prompt=True,enable_thinking=True)
        inp=self.tokenizer(formatted,return_tensors="pt",truncation=True,
                           max_length=4096).to(self.device)
        result=self.enc(**inp,use_cache=False,return_dict=True)
        embedding=result.last_hidden_state[:1]
        data=(embedding,torch.ones(embedding.shape[:2],
                                   device=embedding.device,dtype=torch.bool))
        if len(self.cache)>=self.cache_size:
            self.cache.clear()
        self.cache[prompt]=data
        return data


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--variant",type=Path,required=True,help="folder with transformer/*.bin and COMPLETE.json")
    ap.add_argument("--source",type=Path,required=True,
                    help="original pinned MOSS-Image-SFT, for tokenizer/text encoder/VAE/scheduler")
    ap.add_argument("--prompt",required=True)
    ap.add_argument("--out",type=Path,required=True)
    ap.add_argument("--seed",type=int,default=123)
    ap.add_argument("--steps",type=int,default=50)
    ap.add_argument("--cfg",type=float,default=3.0)
    ap.add_argument("--width",type=int,default=1024)
    ap.add_argument("--height",type=int,default=1024)
    ap.add_argument("--encoder-root",type=Path,default=None,
                    help="Optional separately packaged original Qwen3.5 FP8WO encoder; validates all four shards")
    ap.add_argument("--vae-root",type=Path,default=None,
                    help="Optional locally verified original Flux2 VAE FP16 or BF16 component")
    ap.add_argument("--experimental-high-res",action="store_true")
    ap.add_argument("--allow-unsupported-device",action="store_true",
                    help="Experimental override for unsupported hardware; does NOT guarantee execution")
    args=ap.parse_args()
    if args.steps<1 or not 0<=args.cfg<=10:
        raise ValueError("invalid inference steps/guidance")
    if args.width%16 or args.height%16 or min(args.width,args.height)<512:
        raise ValueError("width and height must be multiples of 16, at least 512")
    if (args.height,args.width)!=(1024,1024) and not args.experimental_high_res:
        raise ValueError("non-1024 sizes require --experimental-high-res (not quality-certified)")
    if not torch.cuda.is_available():
        raise RuntimeError("This inference CLI needs a CUDA GPU")
    gpu=torch.cuda.get_device_name(0)
    cc=torch.cuda.get_device_capability(0)

    args.source=args.source.resolve()
    args.variant=args.variant.resolve()
    receipt=check_artifact(args.variant)
    method=receipt.get("quantization")
    required_cc=((8,9) if method in ("fp8dynamic-row","fp8dynamic-tensor")
                 else (10,0) if method=="nvfp4" else (0,0))
    if cc < required_cc and not args.allow_unsupported_device:
        raise RuntimeError(f"GPU {gpu} compute capability {cc} below minimum for {method} ({required_cc}); "
                           "use a compatible accelerator/runtime or explicit experimental override.")
    verify_pinned_source(args.source, receipt["source_revision"])
    import torchao # noqa: F401; register tensor subclasses for torch.load
    import torchao.quantization # noqa: F401; activate tensor-subclass safe globals
    import torchao.prototype.mx_formats # noqa: F401; NVFP4Tensor registration
    from diffusers import AutoencoderKLFlux2
    module=pinned_loader(args.source)
    start=time.perf_counter()
    transformer=module.MOSSImageTransformer.from_pretrained(
        args.variant/"transformer",torch_dtype=torch.bfloat16,
        use_safetensors=False,local_files_only=True).eval()
    selected=receipt.get("selected_quantized_linears",
                         receipt.get("conversion",{}).get("selected_linear_count"))
    if selected is None or selected < 200:
        raise RuntimeError("Quantized weight count missing or implausibly small")
    converted=sum(type(m.weight).__name__!="Parameter" for m in transformer.modules()
                  if isinstance(m,torch.nn.Linear))
    if converted!=selected:
        raise RuntimeError(f"Rehydration mismatch: {converted} != {selected}; will not fallback")
    transformer=transformer.to("cuda")
    print(json.dumps({"loaded_gpu":gpu, "sm":cc, "variant":receipt["variant"],
       "selected":converted, "load_seconds":round(time.perf_counter()-start,2)}),flush=True)
    # Qwen3.5 text encoder stays in host RAM, streams to GPU module by module.
    # This is slow per new prompt, but prevents a resident >13 GiB encoder.
    text=StreamedQwenText(args.source,encoder_root=args.encoder_root)
    positive=text.encode(args.prompt)
    negative=text.encode("") if args.cfg!=1.0 else None
    if args.vae_root is None:
        vae=AutoencoderKLFlux2.from_pretrained(args.source/"vae",
                torch_dtype=torch.float32,local_files_only=True).eval().to("cuda")
        vae_kind="original_fp32"
    else:
        from compact_components import inspect_vae
        vae_info,dtype=inspect_vae(args.vae_root,verify_sha=True)
        vae=AutoencoderKLFlux2.from_pretrained(Path(args.vae_root)/"vae",
                torch_dtype=dtype,local_files_only=True).eval().to("cuda")
        vae_kind=vae_info["dtype"]
    scheduler=module.MOSSImageScheduler.from_pretrained(args.source/"scheduler",
            local_files_only=True)
    pipe=module.MOSSImagePipeline(transformer,scheduler,None,None,vae)
    generator=torch.Generator("cpu").manual_seed(args.seed)
    kw={}
    if negative is not None:
        kw["negative_prompt_embeds"]=negative[0]
        kw["negative_prompt_attention_mask"]=negative[1]
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()
    t0=time.perf_counter()
    with torch.inference_mode():
        im=pipe(prompt_embeds=positive[0],prompt_attention_mask=positive[1],
                width=args.width,height=args.height,
                num_inference_steps=args.steps,guidance_scale=args.cfg,
                uncond_skip_layers=[22],generator=generator,**kw).images[0]
    torch.cuda.synchronize()
    args.out.parent.mkdir(parents=True,exist_ok=True)
    im.save(args.out)
    print(json.dumps({"image":str(args.out),"variant":receipt["variant"],"lora":False,
       "encoder_precision":text.quantization,"vae_precision":vae_kind,
       "seconds_image":round(time.perf_counter()-t0,2),
       "peak_cuda_alloc_gib":round(torch.cuda.max_memory_allocated()/2**30,3),
       "peak_cuda_reserved_gib":round(torch.cuda.max_memory_reserved()/2**30,3)}),flush=True)

if __name__=="__main__":
    main()