Fix critical bug: SDE stochastic_sampling destroys the conditioning frame into colorful static

66412be verified 6 days ago

7.85 kB

license: other
license_name: openmdw1.1-license
license_link: https://openmdw.ai/license/1-1/
library_name: cosmos
tags:
  - nvidia
  - cosmos
  - cosmos3
  - image-to-video
  - video-generation
  - fp8
  - quantized
  - modelopt

Cosmos3-Super-Image2Video-4Step — FP8 (community quantization)

This is a community FP8 (E4M3), weight-only quantization of NVIDIA's nvidia/Cosmos3-Super-Image2Video-4Step — a 64B-parameter image-to-video model, DMD2-distilled for 4-step, CFG-free generation — produced with NVIDIA TensorRT Model Optimizer (nvidia-modelopt).

Model Developer (original model): NVIDIA. This is not an official NVIDIA release. Everything about the model itself — architecture, training data, benchmarks, full input/output specs, limitations, and ethical considerations — is unchanged from NVIDIA's original and is documented on their model card; it isn't reproduced here. This card only covers what we changed (the quantization) and how to run this checkpoint.

License

Released under the OpenMDW-1.1 license, same as the original model — see LICENSE in this repo.

Quantization Notes

Recipe: per-tensor E4M3 weights on the 64 transformer blocks' attention/MLP linears. Activations, embeddings, norms, the bundled Qwen3 reasoner head, time_embedder, and proj_in/proj_out are left in BF16. Weight-only, so no calibration forward pass was needed. Per NVIDIA's own note, FP8 is not officially tested/supported for this model — treat this as best-effort.

Loading gotchas — read before using with diffusers: as of diffusers==0.39.0.dev0, plain Cosmos3OmniPipeline usage has two bugs against this checkpoint. Both are fixed automatically by the serving script below (serve_cosmos3_i2v4step_diffusers.py); if you're writing your own loading code, patch the scheduler yourself as shown.

Wrong step count / CFG on. Cosmos3OmniPipeline.__call__ does not read this checkpoint's scheduler/scheduler_config.json → fixed_step_sampler_config.t_list (the trained 4-step sde schedule). It silently falls back to the pipeline's generic defaults (num_inference_steps=35, guidance_scale=6.0), which is off-distribution for a DMD2-distilled 4-step checkpoint. (vLLM-Omni parses this file correctly; this gap is specific to the plain diffusers PyTorch path.)
Image conditioning gets destroyed into colorful static (critical). This checkpoint's scheduler ships stochastic_sampling=True (SDE sampling). Cosmos3 anchors the conditioning frame by zeroing the model's predicted velocity there — under a deterministic step that means "leave this frame unchanged," but FlowMatchEulerDiscreteScheduler's SDE branch computes x0 = sample - current_sigma * model_output (= sample, since velocity is
1. then prev_sample = (1 - next_sigma) * x0 + next_sigma * randn_tensor(...) — it re-noises by next_sigma regardless of velocity. Across this checkpoint's 4 steps that compounds to ~99.6% fresh noise in the conditioned frame: your input image comes out as colorful static while the genuinely-denoised motion frames still look like plausible video (this combination — garbage first frame, coherent-but-drifted rest — is exactly how it presents). Disabling stochastic_sampling restores the correct zero-velocity-is-a-no-op behavior. Confirmed by direct A/B render, same seed/image/prompt, only this flag changed.

def force_fixed_step_schedule(scheduler):
    t_list = scheduler.config.fixed_step_sampler_config["t_list"]
    orig = scheduler.set_timesteps
    scheduler.set_timesteps = lambda num_inference_steps=None, device=None, **_: (
        orig(sigmas=list(t_list), device=device)
    )
    if scheduler.config.stochastic_sampling:
        scheduler.register_to_config(stochastic_sampling=False)

pipe = ...  # Cosmos3OmniPipeline.from_pretrained(...)
force_fixed_step_schedule(pipe.scheduler)

result = pipe(prompt=..., image=..., guidance_scale=1.0, num_inference_steps=4, ...)
# guidance_scale=1.0 disables CFG; num_inference_steps is a no-op once patched

Verified against this exact repo: the patch produces a 4-iteration denoising loop (not 35) with CFG off and a correctly-preserved conditioning frame (not colorful static), matching the checkpoint's trained regime. The save/restore round trip (transformer/modelopt_state.pth) reloads correctly with 896 quantized weight wrappers active.

Usage: Run Inference (single GPU, `diffusers`)

This repo is a diffusers-loadable repackage (transformer/modelopt_state.pth), not NVIDIA's vLLM-Omni deployment export format. If you have a vLLM-Omni cluster, use NVIDIA's original BF16 checkpoint and card instead. For everyone else — anyone running this FP8 checkpoint on a single GPU — use the FastAPI server included in this repo, which loads the checkpoint, applies the 4-step scheduler fix above automatically, and exposes a plain HTTP endpoint.

1. Install

Requires a diffusers build with Cosmos3 support — not yet in a PyPI release as of this writing, so install from the exact commit this checkpoint was produced and validated against:

pip install "git+https://github.com/huggingface/diffusers.git@2c7efb95349296cf6bcce981ea036275a82a94df"
pip install nvidia-modelopt accelerate torch fastapi uvicorn python-multipart

2. Download this repo (weights + scripts together)

hf download prometheusAIR/Cosmos3-Super-Image2Video-4Step-FP8 \
    --local-dir Cosmos3-Super-Image2Video-4Step-FP8
cd Cosmos3-Super-Image2Video-4Step-FP8

This pulls the whole ~66GB repo — FP8 weights, VAE, tokenizer, and the .py scripts side by side, the same download either way.

3. Serve it

CUDA_VISIBLE_DEVICES=0 python serve_cosmos3_i2v4step_diffusers.py --repo .

One GPU with roughly 70GB+ free memory (RTX PRO 6000 96GB, H100/H200, A100-80GB, etc.) is enough — no multi-GPU sharding required for this FP8 checkpoint. The server listens on http://localhost:8000 once it's done loading.

4. Use it

curl -s -X POST http://localhost:8000/animate \
    -F image=@your_first_frame.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

num_frames / fps / height / width are adjustable. num_inference_steps and guidance_scale are intentionally not exposed — this checkpoint's 4-step, CFG-free schedule is fixed and applied automatically by the server.

Other scripts in this repo

quantize_cosmos3_i2v4step_streaming.py — reproduces this FP8 quantization from NVIDIA's BF16 source checkpoint.
repackage_for_hf_i2v4step.py — rebuilds this diffusers-loadable repo format from a quantized transformer.
load_cosmos3_modelopt.py — the underlying loader serve_cosmos3_i2v4step_diffusers.py uses; import load_pipe(...) directly if you want a pipe object instead of an HTTP server.
validate_cosmos3_i2v4step_fp8.py — a minimal standalone image→video smoke test against this checkpoint.

Responsible Use

See NVIDIA's Bias, Explainability, Safety & Security, and Privacy subcards (included in this repo), and the Limitations / Ethical Considerations sections of NVIDIA's original model card. Report security vulnerabilities here.