#!/usr/bin/env python
"""
Local generation server for your quantized Cosmos3-Super, built on the *validated*
diffusers path (NOT vLLM-Omni).

This reproduces, at startup, the exact in-memory model your streaming quantizer
already rendered from successfully -- build empty on meta, insert weight-only
quantizers, compress, stream the BF16 shards into compressed form -- then serves
that model behind a tiny HTTP API. It does NOT reload the export_hf_checkpoint
output (that unified format is for vLLM-Omni / TRT-LLM; diffusers round-trips a
ModelOpt model via modelopt_state + state_dict, which is what --cache uses below).

Nothing here is speculative: the model object served is the same one that produced
cosmos3_super_<fmt>_validate.png.

ENDPOINTS
---------
  GET  /health    -> readiness + which format is loaded
  POST /generate  -> text -> still image (JSON body; returns PNG)
  POST /animate   -> image -> video (multipart upload; returns MP4, or GIF if no
                     mp4 encoder is installed)

ENV / DEPS
----------
Run in the venv that has diffusers-from-git-main + modelopt + accelerate (your
quantization venv, e.g. /home/prometheus/ModelOpt/.venv). Extra installs:

    pip install fastapi uvicorn python-multipart      # python-multipart is REQUIRED for /animate
    pip install imageio imageio-ffmpeg                # optional: mp4 output (else /animate returns GIF)

USAGE
-----
    CUDA_VISIBLE_DEVICES=0 python serve_cosmos3_diffusers.py --format nvfp4
    # faster restarts after the first boot (writes/reads a ~36 GB cache):
    CUDA_VISIBLE_DEVICES=0 python serve_cosmos3_diffusers.py --format nvfp4 --cache ./cosmos3-cache

Text -> still image:
    curl -s -X POST http://localhost:8000/generate \
        -H 'Content-Type: application/json' \
        -d '{"prompt":"a robot arm on a workbench in a bright lab","num_inference_steps":50}' \
        --output out.png

Image -> video (upload the conditioning frame, so server-side paths never matter;
`@` makes curl attach the file from YOUR current directory, and a shell ~ is expanded
by the shell before curl runs):
    curl -s -X POST http://localhost:8000/animate \
        -F image=@out.png \
        -F 'prompt=The robotic arm slowly lowers its gripper toward the objects and holds. Static camera.' \
        -F num_frames=49 -F fps=24 \
        --output clip.mp4

Health:
    curl -s http://localhost:8000/health
"""

import argparse
import contextlib
import gc
import io
import os
import tempfile

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

import torch
from accelerate import init_empty_weights, load_checkpoint_in_model
from accelerate.utils import get_max_memory, infer_auto_device_map
from accelerate.utils.dataclasses import CustomDtype
from huggingface_hub import snapshot_download
from PIL import Image

import modelopt.torch.quantization as mtq
from diffusers import Cosmos3OmniTransformer
from diffusers.utils import export_to_gif, export_to_video

SRC_REPO = "nvidia/Cosmos3-Super"

# --- hard-won config from the validated quantizer (inlined so this file stands alone) ---
SPARE_SUBSTRINGS = [
    "time_embedder", "proj_in", "proj_out", "lm_head", "embed", "norm", "audio_proj",
]


def _is_spare(name: str) -> bool:
    return any(s in name for s in SPARE_SUBSTRINGS)


def build_quant_cfg(fmt: str) -> dict:
    if fmt == "fp8":
        return {
            "quant_cfg": {
                "*weight_quantizer": {"num_bits": (4, 3), "axis": None, "enable": True},
                "*input_quantizer": {"enable": False},
                "*output_quantizer": {"enable": False},
                "*softmax_quantizer": {"enable": False},
            },
            "algorithm": "max",
        }
    if fmt == "nvfp4":
        import copy
        base = getattr(mtq, "W4A16_NVFP4_CFG", None) or mtq.NVFP4_DEFAULT_CFG
        cfg = copy.deepcopy(base)
        # Bake weight-only INTO THE CONFIG: modelopt_state replays the config, not
        # imperative .disable() calls made after quantize. NVFP4_DEFAULT_CFG ships with
        # activation quantization enabled, so without this, a restored checkpoint comes
        # back with ~1806 dynamic activation quantizers active (~10x slower per step).
        # The drop-in loader re-disables as belt-and-braces, but the saved state should
        # be correct on its own. (The FP8 dict below already does this.)
        cfg.setdefault("quant_cfg", {})
        cfg["quant_cfg"]["*input_quantizer"] = {"enable": False}
        cfg["quant_cfg"]["*output_quantizer"] = {"enable": False}
        cfg["quant_cfg"]["*softmax_quantizer"] = {"enable": False}
        for s in SPARE_SUBSTRINGS:
            cfg["quant_cfg"][f"*{s}*weight_quantizer"] = {"enable": False}
        return cfg
    raise ValueError(f"Unknown format: {fmt!r}")


def enforce_weight_only_and_spare(model) -> tuple[int, int]:
    n_spare = n_act = 0
    for name, module in model.named_modules():
        if not (name.endswith("_quantizer") and hasattr(module, "disable")):
            continue
        if name.endswith("weight_quantizer"):
            if _is_spare(name.rsplit(".", 1)[0]):
                module.disable()
                n_spare += 1
        else:
            module.disable()
            n_act += 1
    return n_spare, n_act


def compressed_device_map(model, gpu_mem_fraction: float = 0.85) -> dict:
    max_memory = {k: v * gpu_mem_fraction for k, v in get_max_memory().items()}
    no_split = set()
    for name, module in model.named_modules():
        if name.endswith((".layers.0", ".blocks.0", ".transformer_blocks.0")):
            no_split.add(module.__class__.__name__)
    special_dtypes = {}
    for name, module in model.named_modules():
        if (
            hasattr(module, "weight")
            and hasattr(module, "weight_quantizer")
            and getattr(module.weight_quantizer, "is_enabled", True)
            and not getattr(module.weight_quantizer, "fake_quant", True)
        ):
            nb = module.weight_quantizer.num_bits
            if isinstance(nb, tuple):
                nb = nb[0] + nb[1] + 1
            special_dtypes[name + ".weight"] = CustomDtype.FP8 if nb == 8 else CustomDtype.INT4
    return infer_auto_device_map(
        model, max_memory=max_memory,
        no_split_module_classes=list(no_split), special_dtypes=special_dtypes,
    )


def _materialize_residual_meta(model) -> int:
    """Fill any leftover meta tensors (disabled-quantizer scratch) with zeros on GPU."""
    n = 0
    for _, module in model.named_modules():
        for bn, buf in list(module._buffers.items()):
            if buf is not None and getattr(buf, "is_meta", False):
                module._buffers[bn] = torch.zeros(buf.shape, dtype=buf.dtype, device="cuda")
                n += 1
        for pn, par in list(module._parameters.items()):
            if par is not None and getattr(par, "is_meta", False):
                module._parameters[pn] = torch.nn.Parameter(
                    torch.zeros(par.shape, dtype=par.dtype, device="cuda"), requires_grad=False
                )
                n += 1
    return n


def _transformer_dir() -> str:
    local_root = snapshot_download(SRC_REPO, allow_patterns=["transformer/*"])
    return os.path.join(local_root, "transformer")


def build_quantized_transformer(fmt: str, gpu_mem_fraction: float = 0.85):
    """The proven path: empty-on-meta -> quantize -> compress -> stream shards in."""
    transformer_dir = _transformer_dir()
    print(f"[build] empty transformer on meta from {transformer_dir}")
    config = Cosmos3OmniTransformer.load_config(transformer_dir)
    with init_empty_weights(include_buffers=False):
        model = Cosmos3OmniTransformer.from_config(config)

    print(f"[build] inserting weight-only {fmt} quantizers")
    mtq.quantize(model, build_quant_cfg(fmt))
    n_spare, n_act = enforce_weight_only_and_spare(model)
    print(f"[build] weight-only: disabled {n_act} activation quantizers; {n_spare} spare weight layers")

    print("[build] setting up compressed parameter shapes")
    try:
        mtq.compress(model, config=mtq.CompressConfig(quant_gemm=False))
    except (AttributeError, TypeError):
        mtq.compress(model)

    print("[build] streaming BF16 shards into compressed form (slow step)")
    load_checkpoint_in_model(
        model, checkpoint=transformer_dir,
        device_map=compressed_device_map(model, gpu_mem_fraction), dtype=torch.bfloat16,
    )
    fixed = _materialize_residual_meta(model)
    if fixed:
        print(f"[build] materialized {fixed} residual meta tensors")
    return model


# --- optional fast-restart cache (modelopt_state + weights, per ModelOpt docs) ----------
def _cache_paths(cache_dir: str, fmt: str):
    return (os.path.join(cache_dir, f"modelopt_state_{fmt}.pt"),
            os.path.join(cache_dir, f"weights_{fmt}.pt"))


def save_quantized(model, fmt: str, cache_dir: str) -> None:
    try:
        from modelopt.torch.opt import modelopt_state
    except ImportError:
        from modelopt.torch.opt.conversion import modelopt_state
    os.makedirs(cache_dir, exist_ok=True)
    state_path, weights_path = _cache_paths(cache_dir, fmt)
    print(f"[cache] writing {state_path} + {weights_path} (large; one time)")
    torch.save(modelopt_state(model), state_path)
    torch.save(model.state_dict(), weights_path)


def try_restore_quantized(fmt: str, cache_dir: str):
    """Restore the compressed model from cache. Returns model or None (caller rebuilds)."""
    state_path, weights_path = _cache_paths(cache_dir, fmt)
    if not (os.path.isfile(state_path) and os.path.isfile(weights_path)):
        return None
    try:
        try:
            from modelopt.torch.opt import restore_from_modelopt_state
        except ImportError:
            from modelopt.torch.opt.conversion import restore_from_modelopt_state
        print(f"[cache] restoring from {state_path}")
        config = Cosmos3OmniTransformer.load_config(_transformer_dir())
        with init_empty_weights(include_buffers=False):
            model = Cosmos3OmniTransformer.from_config(config)
        state = torch.load(state_path, map_location="cpu", weights_only=False)
        restore_from_modelopt_state(model, state)  # replays quantize + compress structure
        weights = torch.load(weights_path, map_location="cpu", weights_only=False)
        model.load_state_dict(weights, strict=False, assign=True)
        _materialize_residual_meta(model)
        print("[cache] restore OK")
        return model
    except Exception as e:
        import traceback
        print(f"[cache] restore failed ({type(e).__name__}: {e}); falling back to full rebuild")
        traceback.print_exc()
        return None


# --- pipeline assembly (mirrors the validated render_from_memory) -----------------------
def make_pipeline(model, flow_shift: float = 3.0):
    from diffusers import Cosmos3OmniPipeline
    from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler

    model.to("cuda")
    # dtype consistency: keep FP8/NVFP4 weights, but bring stray fp32 buffers to bf16 and
    # cast time-embedder inputs at the fp32->bf16 boundary (the validated nudges).
    for m in model.modules():
        for bn, buf in list(m._buffers.items()):
            if buf is not None and buf.dtype == torch.float32:
                m._buffers[bn] = buf.to(torch.bfloat16)

    def _cast_bf16(_m, args):
        return tuple(
            a.to(torch.bfloat16)
            if torch.is_tensor(a) and a.is_floating_point() and a.dtype != torch.bfloat16 else a
            for a in args
        )

    for name, m in model.named_modules():
        if "time_embedder" in name and hasattr(m, "linear_1"):
            m.register_forward_pre_hook(_cast_bf16)

    pipe = Cosmos3OmniPipeline.from_pretrained(
        SRC_REPO, transformer=model, torch_dtype=torch.bfloat16,
        enable_safety_checker=False,  # local single-user server; revisit if exposing it
    )
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
    for name, comp in pipe.components.items():
        if name != "transformer" and isinstance(comp, torch.nn.Module):
            comp.to("cuda")
    return pipe


# --- HTTP server ------------------------------------------------------------------------
import asyncio
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import Response
from pydantic import BaseModel

STATE: dict = {}
_gen_lock = asyncio.Lock()  # one generation at a time on a single GPU


# ---- text -> still image -------------------------------------------------------
class GenRequest(BaseModel):
    prompt: str
    negative_prompt: str = ""
    num_inference_steps: int = 50
    guidance_scale: float = 4.0
    height: int = 1024
    width: int = 1024
    num_frames: int = 1            # 1 = still image; >1 = video frames (heavier)
    seed: int | None = 1234        # null -> random each call


def _run_generation(req: GenRequest) -> bytes:
    pipe = STATE["pipe"]
    gen = torch.Generator(device="cuda").manual_seed(int(req.seed)) if req.seed is not None else None
    with torch.inference_mode():
        result = pipe(
            prompt=req.prompt,
            negative_prompt=req.negative_prompt,
            num_frames=req.num_frames,
            height=req.height,
            width=req.width,
            num_inference_steps=req.num_inference_steps,
            guidance_scale=req.guidance_scale,
            generator=gen,
        )
    img = result.video[0]  # PIL image for the first (or only) frame
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    del result
    gc.collect()
    torch.cuda.empty_cache()
    return buf.getvalue()


# ---- image -> video (i2v) ------------------------------------------------------
def _run_i2v(pil_image, prompt, negative_prompt, num_frames, fps,
             height, width, steps, guidance, seed) -> tuple[bytes, str]:
    pipe = STATE["pipe"]
    image = pil_image.convert("RGB")  # the pipeline resizes this to (height, width)
    gen = torch.Generator(device="cuda").manual_seed(int(seed)) if seed >= 0 else None
    with torch.inference_mode():
        result = pipe(
            prompt=prompt, negative_prompt=negative_prompt,
            image=image, num_frames=num_frames, fps=fps,
            height=height, width=width,
            num_inference_steps=steps, guidance_scale=guidance,
            enable_safety_check=False, generator=gen, output_type="pil",
        )
    frames = result.video  # list of PIL frames
    try:
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tf:
            path = tf.name
        export_to_video(frames, path, fps=int(round(fps)))
        media = "video/mp4"
    except Exception:  # no mp4 backend installed -> GIF (PIL-only, always works)
        with tempfile.NamedTemporaryFile(suffix=".gif", delete=False) as tf:
            path = tf.name
        export_to_gif(frames, path)
        media = "image/gif"
    data = open(path, "rb").read()
    os.remove(path)
    del result
    gc.collect()
    torch.cuda.empty_cache()
    return data, media


@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
    fmt = STATE["fmt"]
    cache_dir = STATE.get("cache_dir")
    model = None
    if cache_dir:
        model = try_restore_quantized(fmt, cache_dir)
    if model is None:
        model = build_quantized_transformer(fmt, STATE["gpu_mem_fraction"])
        if cache_dir:
            try:
                save_quantized(model, fmt, cache_dir)
            except Exception as e:
                print(f"[cache] save failed ({type(e).__name__}: {e}); continuing without cache")
    STATE["pipe"] = make_pipeline(model, STATE["flow_shift"])
    print(f"[ready] serving {fmt.upper()} Cosmos3-Super on diffusers")
    yield
    STATE.clear()


app = FastAPI(lifespan=lifespan)


@app.get("/health")
async def health():
    return {"status": "ok" if "pipe" in STATE else "loading", "format": STATE.get("fmt")}


@app.post("/generate")
async def generate(req: GenRequest):
    async with _gen_lock:
        loop = asyncio.get_running_loop()
        png = await loop.run_in_executor(None, _run_generation, req)
    return Response(content=png, media_type="image/png")


@app.post("/animate")
async def animate(
    image: UploadFile = File(...),
    prompt: str = Form(...),
    negative_prompt: str = Form(""),
    num_frames: int = Form(49),         # ~2.04s @ 24fps; 4n+1 maps cleanly to the VAE's 4x temporal compression
    fps: float = Form(24.0),            # native framerate; it conditions duration + audio length, so keep 24
    height: int = Form(1024),
    width: int = Form(1024),
    num_inference_steps: int = Form(35),  # video default (the still path uses 50)
    guidance_scale: float = Form(6.0),    # video default (the still path uses 4.0)
    seed: int = Form(1234),               # pass -1 for a random clip each call
):
    pil = Image.open(io.BytesIO(await image.read()))
    async with _gen_lock:
        loop = asyncio.get_running_loop()
        data, media = await loop.run_in_executor(
            None, _run_i2v, pil, prompt, negative_prompt, num_frames, fps,
            height, width, num_inference_steps, guidance_scale, seed,
        )
    return Response(content=data, media_type=media)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--format", choices=["fp8", "nvfp4"], default="nvfp4")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=8000)
    ap.add_argument("--flow-shift", type=float, default=3.0)
    ap.add_argument("--gpu-mem-fraction", type=float, default=0.85)
    ap.add_argument("--cache", default=None,
                    help="Dir for a fast-restart cache. First boot rebuilds + writes it; "
                         "later boots restore from it. Any restore error -> full rebuild.")
    args = ap.parse_args()

    STATE.update(
        fmt=args.format, flow_shift=args.flow_shift,
        gpu_mem_fraction=args.gpu_mem_fraction, cache_dir=args.cache,
    )

    import uvicorn
    uvicorn.run(app, host=args.host, port=args.port)