prometheusAIR's picture
Fix critical bug: SDE stochastic_sampling destroys the conditioning frame into colorful static
66412be verified
|
Raw
History Blame Contribute Delete
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`](https://huggingface.co/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](https://github.com/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](https://huggingface.co/nvidia/Cosmos3-Super-Image2Video-4Step);
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](https://openmdw.ai/license/1-1/) license, same
as the original model β€” see [`LICENSE`](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.
1. **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.)
2. **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
0) 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.
```python
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:
```bash
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)
```bash
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
```bash
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
```bash
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](BIAS.md), [Explainability](EXPLAINABILITY.md),
[Safety & Security](SAFETY.md), and [Privacy](PRIVACY.md) subcards (included
in this repo), and the Limitations / Ethical Considerations sections of
[NVIDIA's original model card](https://huggingface.co/nvidia/Cosmos3-Super-Image2Video-4Step).
Report security vulnerabilities [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).