#!/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)