#!/usr/bin/env python
"""
Standalone re-validation of the Cosmos3-Super-Image2Video-4Step FP8 quantization
(quantize_cosmos3_i2v4step_streaming.py). Exists because the first --smoke run's
render was invalid: it used the pipeline's default 35-step/guidance=6.0 schedule
instead of this checkpoint's trained 4-step sde schedule (a diffusers version gap --
Cosmos3OmniPipeline doesn't pass the scheduler_config.json's fixed_step_sampler_config
through to scheduler.set_timesteps; see _force_fixed_step_schedule in
serve_cosmos3_i2v4step_diffusers.py for the full rationale).

NOTE: export_hf_checkpoint() (what the quantize script wrote to ./cosmos3-i2v4step-fp8)
is NVIDIA's *deployment* format for vLLM/TRT-LLM -- it has no modelopt_state.pth, so
diffusers' from_pretrained / load_cosmos3_modelopt.py can't load it directly (same gap
repackage_for_hf.py exists to solve for the base model). Rather than write a second
~65GB copy to disk just to validate, this rebuilds the quantized transformer in memory
via the same build_quantized_transformer() used during quantization (re-streams the
121GB source once) and renders straight from that -- no extra disk footprint beyond the
output clip. If you want a real servable drop-in repo afterward, run
repackage_for_hf_i2v4step.py (that DOES need the disk headroom).

USAGE
    python validate_cosmos3_i2v4step_fp8.py --format fp8 --image out.png
"""
import argparse
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
from PIL import Image

from serve_cosmos3_i2v4step_diffusers import build_quantized_transformer, make_pipeline


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--format", choices=["fp8", "nvfp4"], default="fp8")
    ap.add_argument("--gpu-mem-fraction", type=float, default=0.85)
    ap.add_argument("--image", default="out.png")
    ap.add_argument("--num-frames", type=int, default=49)
    ap.add_argument("--height", type=int, default=512)
    ap.add_argument("--width", type=int, default=512)
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--out", default="cosmos3_i2v4step_fp8_validate.mp4")
    ap.add_argument(
        "--prompt",
        default="The camera holds static as the scene continues naturally, with subtle motion and realistic physics.",
    )
    args = ap.parse_args()

    print(f"[build] rebuilding quantized ({args.format}) transformer in memory ...")
    model = build_quantized_transformer(args.format, args.gpu_mem_fraction)
    pipe = make_pipeline(model)  # includes _force_fixed_step_schedule on pipe.scheduler

    cond_image = Image.open(args.image).convert("RGB")

    print(f"[render] {args.num_frames} frames @ {args.height}x{args.width}, seed={args.seed}")
    with torch.inference_mode():
        result = pipe(
            prompt=args.prompt,
            negative_prompt="",
            image=cond_image,
            num_frames=args.num_frames,
            height=args.height,
            width=args.width,
            num_inference_steps=4,   # no-op once patched; kept for clarity
            guidance_scale=1.0,       # CFG is distilled out of this checkpoint
            generator=torch.Generator(device="cuda").manual_seed(args.seed),
        )

    from diffusers.utils import export_to_video
    export_to_video(result.video, args.out, fps=24)
    print(f"[done] wrote {args.out}")


if __name__ == "__main__":
    main()