serve_cosmos3_i2v4step_diffusers.py

#!/usr/bin/env python
"""
Local generation server for your quantized Cosmos3-Super-Image2Video-4Step, built on
the *validated* diffusers path (NOT vLLM-Omni). Sibling of serve_cosmos3_diffusers.py
(the nvidia/Cosmos3-Super base-model server) -- same build recipe, adapted for three
real differences in this checkpoint:

  1. NO /generate (text -> still) ENDPOINT. This checkpoint is Image2Video-specific
     (unlike the T2I/T2V/I2V-omni base model); NVIDIA's docs never demonstrate a
     text-only call for it, and there's no local way to verify one works without
     the checkpoint in hand. Rather than ship a silently-untested endpoint, this
     server exposes /animate only. If a text-only path turns out to work when you
     test it, port /generate over from serve_cosmos3_diffusers.py at that point.
  2. NO SCHEDULER OBJECT SWAP, BUT A PATCHED set_timesteps. serve_cosmos3_diffusers.py
     replaces the shipped scheduler with UniPCMultistepScheduler(flow_shift=...); this
     checkpoint instead keeps FlowMatchEulerDiscreteScheduler as shipped (do not swap
     it) but needs one targeted patch: Cosmos3OmniPipeline.__call__ always calls
     scheduler.set_timesteps(num_inference_steps, device=device) with no passthrough
     for this checkpoint's scheduler_config.json fixed_step_sampler_config (verified by
     running it -- an unpatched call silently executed the pipeline's 35-step default
     instead of the checkpoint's trained 4-step sde schedule). make_pipeline() below
     applies _force_fixed_step_schedule() to fix this.
  3. NO num_inference_steps IN THE REQUEST; guidance_scale IS FIXED AT 1.0. Per
     NVIDIA's model card these are fixed by the distilled checkpoint -- CFG is
     distilled out (do_classifier_free_guidance is guidance_scale != 1.0 in the
     pipeline source, so the 6.0 pipeline default would silently re-enable it) and
     the step count comes from the scheduler patch above, not from the caller.

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

USAGE
-----
    CUDA_VISIBLE_DEVICES=0 python serve_cosmos3_i2v4step_diffusers.py --repo prometheusAIR/Cosmos3-I2V4Step-fp8
    CUDA_VISIBLE_DEVICES=0 python serve_cosmos3_i2v4step_diffusers.py --repo ./cosmos3-i2v4step-fp8-hf

    # or rebuild from the bf16 source (the original streaming path):
    CUDA_VISIBLE_DEVICES=0 python serve_cosmos3_i2v4step_diffusers.py --format fp8
    CUDA_VISIBLE_DEVICES=0 python serve_cosmos3_i2v4step_diffusers.py --format fp8 --cache ./cosmos3-i2v4step-cache

Image -> video:
    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 asyncio
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 fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import Response
from huggingface_hub import snapshot_download
from PIL import Image
from pydantic import BaseModel

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-Image2Video-4Step"

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
        # Return the base preset UNMODIFIED -- on this installed modelopt version,
        # W4A16_NVFP4_CFG's quant_cfg is a LIST (newer format), not the dict this
        # function used to assume (confirmed empirically: "TypeError: list indices
        # must be integers or slices, not str" trying dict-style key assignment
        # here). enforce_weight_only_and_spare() below disables activations/spares
        # on the model's actual inserted quantizer modules AFTER mtq.quantize(),
        # which is config-form agnostic -- no need to pre-bake into the config dict.
        # load_cosmos3_modelopt.py's loader re-applies the same disabling on every
        # restore anyway (modelopt_state replays the config, not our imperative
        # .disable() calls), so skipping the pre-bake here loses nothing.
        base = getattr(mtq, "W4A16_NVFP4_CFG", None) or mtq.NVFP4_DEFAULT_CFG
        return copy.deepcopy(base)
    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:
    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):
    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)
        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


def _force_fixed_step_schedule(scheduler) -> bool:
    """See module docstring point 2. Patches THIS scheduler instance's set_timesteps
    to always use the checkpoint's own fixed_step_sampler_config.t_list, regardless of
    whatever num_inference_steps Cosmos3OmniPipeline.__call__ passes in internally.

    ALSO disables stochastic_sampling (forces the deterministic ODE branch) -- CONFIRMED
    by direct A/B render (same seed/image/prompt), not speculative. This checkpoint's
    scheduler ships stochastic_sampling=True (SDE). Cosmos3's image-conditioning anchors
    frame 0 by zeroing the model's predicted velocity there; that only means "leave this
    position unchanged" under the deterministic step (prev_sample = sample + dt*0). The
    SDE branch instead computes x0 = sample - current_sigma*model_output (= sample when
    velocity is 0), then prev_sample = (1-next_sigma)*x0 + next_sigma*randn_tensor(...) --
    it re-noises by next_sigma regardless of velocity, with no zero-velocity no-op. Over
    this checkpoint's 4 steps that compounds to ~99.6% fresh noise in the conditioned
    frame by the end: the input image comes out as colorful static while the genuinely
    denoised motion frames still look like a plausible video. Disabling stochastic_
    sampling restores the correct zero-velocity-is-a-no-op behavior."""
    cfg = getattr(scheduler.config, "fixed_step_sampler_config", None)
    t_list = cfg.get("t_list") if isinstance(cfg, dict) else None
    if not t_list:
        print("[warn] no fixed_step_sampler_config.t_list on this scheduler; leaving set_timesteps unpatched")
        return False
    _orig = scheduler.set_timesteps

    def _patched(num_inference_steps=None, device=None, sigmas=None, mu=None, timesteps=None):
        return _orig(sigmas=list(t_list), device=device)

    scheduler.set_timesteps = _patched
    if getattr(scheduler.config, "stochastic_sampling", False):
        scheduler.register_to_config(stochastic_sampling=False)
        print("[build] disabled stochastic_sampling (SDE re-noising corrupts the "
              "image-conditioned frame -- see docstring)")
    print(f"[build] forced fixed {len(t_list)}-step sde schedule: {t_list}")
    return True


# --- pipeline assembly --------------------------------------------------------------
def make_pipeline(model):
    from diffusers import Cosmos3OmniPipeline

    model.to("cuda")
    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
    )
    # Keep the shipped FlowMatchEulerDiscreteScheduler object, but patch its
    # set_timesteps to actually honor the checkpoint's fixed 4-step sde schedule.
    _force_fixed_step_schedule(pipe.scheduler)
    for name, comp in pipe.components.items():
        if name != "transformer" and isinstance(comp, torch.nn.Module):
            comp.to("cuda")
    return pipe


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


def _run_i2v(pil_image, prompt, negative_prompt, num_frames, fps, height, width, 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():
        # num_inference_steps is a no-op once _force_fixed_step_schedule has patched
        # the scheduler (it always substitutes the checkpoint's own t_list).
        # guidance_scale=1.0 disables CFG -- see module docstring point 3.
        result = pipe(
            prompt=prompt, negative_prompt=negative_prompt,
            image=image, num_frames=num_frames, fps=fps,
            height=height, width=width,
            num_inference_steps=4, guidance_scale=1.0,
            generator=gen, output_type="pil",
        )
    frames = result.video
    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"]
    repo = STATE.get("repo")
    if repo:
        from load_cosmos3_modelopt import load_pipe

        path = repo
        if not os.path.isdir(path):
            from huggingface_hub import snapshot_download as _snap
            print(f"[serve] fetching full snapshot of {repo} from the Hub ...")
            path = _snap(repo)
        print(f"[serve] loading drop-in checkpoint: {path}")
        pipe = load_pipe(path)
        # load_cosmos3_modelopt.load_pipe() is shared with the base model and doesn't
        # know about this checkpoint's fixed-step schedule -- patch it here instead.
        _force_fixed_step_schedule(pipe.scheduler)
        STATE["pipe"] = pipe
        print(f"[serve] ready (drop-in: {repo}, {fmt.upper()} assumed from repo contents)")
        yield
        STATE.clear()
        return
    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)
    print(f"[ready] serving {fmt.upper()} Cosmos3-Super-Image2Video-4Step 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("/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),
    height: int = Form(1024),
    width: int = Form(1024),
    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, 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="fp8")
    ap.add_argument("--repo", default=None,
                    help="Serve a published drop-in checkpoint: an HF repo id or a local "
                         "repackaged dir. Skips the bf16 rebuild entirely; requires "
                         "load_cosmos3_modelopt.py next to this file.")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=8000)
    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, gpu_mem_fraction=args.gpu_mem_fraction, cache_dir=args.cache, repo=args.repo,
    )

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