Spaces:
Running on L40S
Running on L40S
File size: 17,695 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 | # Cosmos3 Integration Demos
Minimal worked examples for **taking Cosmos3 into your own training / inference
framework**. Each demo is self-contained (one Python file) and runs end-to-end
on a single 80 GB GPU.
> **These demos use RANDOM main-transformer weights.** They do not load the
> ~30 GB Cosmos3-Nano DCP shards — only `config.json` is fetched. Losses,
> pixels, and samples are therefore *not meaningful*; the point is to show the
> API call sequence and tensor shapes so you can wire OmniMoTModel into your
> own code. For real weight loading see
> [`cosmos_framework.inference.model.Cosmos3OmniModel.from_pretrained_dcp`](../../cosmos_framework/inference/model.py)
> and the production CLIs in `cosmos_framework.scripts.{inference,train}`.
>
> This directory is **integration docs by example**, not a model zoo. It does
> not introduce any new training recipe — every file shows how to call code
> that already exists in `cosmos_framework/` from a *plain* PyTorch loop.
## Modality coverage
All three demos cover all four generation modes that Cosmos3-Nano supports:
| | T2I (image) | T2V (video) | ACTION_FDM | T2VS (sound+video) |
| ---------------------------- | :---------: | :---------: | :--------: | :----------------: |
| `trainer_level_inference.py` | ✅ | ✅ | ✅¹ | ✅¹ |
| `trainer_level_training.py` | ✅ | ✅ | ✅ | ✅ |
| `net_level.py` train | ✅ | ✅ | ✅ | ✅ |
| `net_level.py` sample | ✅ | ✅ | ✅¹ | ✅¹ |
¹ For ACTION_FDM and T2VS, the demos feed the model **random** conditioning
(video / actions / audio waveforms). The call sequence runs end-to-end —
loss + backward in training, sampler + decode in inference — but the
*output* is visual / audio noise. The wiring is what's being demonstrated.
For meaningful samples, swap in real conditioning data via your loader.
---
## 1. Pick the right demo
Two integration levels, four cases:
| | **Trainer-level** | **Net-level** |
| --------------- | --------------------------------------------------------- | ----------------------------------- |
| Module used | `OmniMoTModel` | `model.net` (= `Cosmos3VFMNetwork`) |
| Entry call(s) | `training_step` / `generate_samples_from_batch` | `net.forward(packed_seq, ...)` |
| Loss + sampler | written by cosmos_framework (rectified-flow, UniPC) | written by **you** in the demo |
| Effort to adopt | Lowest | Higher (you control loss & sampler) |
| File | `trainer_level_inference.py`, `trainer_level_training.py` | `net_level.py` |
```
┌────────────────────────────────────────────────────────────────────────────┐
│ OmniMoTModel ◀── Cases 1 & 2 plug in here (high-level integration) │
│ ├── training_step(batch, iter) → (aux, loss) │
│ ├── generate_samples_from_batch(batch) → {"vision": [...]} │
│ ├── encode / decode (VAE) │
│ ├── _pack_input_sequence(...) (PackedSequence builder) │
│ │ │
│ └── net = Cosmos3VFMNetwork ◀── Cases 3 & 4 plug in here (low-level) │
│ forward(packed_seq, fps_vision=...) → {"preds_vision": [...]} │
└────────────────────────────────────────────────────────────────────────────┘
```
### Decision matrix
| If you want to… | Use |
| ---------------------------------------------------------------------------- | ---------------------------- |
| Drop Cosmos3 into your training framework with minimum work | `trainer_level_training.py` |
| Drop Cosmos3 into your serving / batch-inference framework with minimum work | `trainer_level_inference.py` |
| Write a custom diffusion loss / curriculum / RL objective | `net_level.py` (train) |
| Write a custom sampler / guidance / consistency scheme | `net_level.py` (infer) |
---
## 2. The four cases
### Case 1 — `trainer_level_inference.py` (trainer-level inference)
What you replace from cosmos_framework: the `OmniInference` pipeline, Ray serving,
the CLI entry-point in `cosmos_framework.scripts.inference`. You keep `OmniMoTModel`
and its built-in CFG + UniPC/EDM sampler.
Has a `--mode {t2i,t2v,action_fdm,t2vs}` flag. T2I/T2V batches come from
cosmos_framework's `get_sample_data` helper; action_fdm and t2vs are hand-built with
random conditioning. The model call is identical for all modes:
```python
model = Cosmos3OmniModel.from_pretrained_dcp(ckpt_dir).model # OmniMoTModel
batch = build_t2iv_batch(model, ..., num_frames) # or build_action_fdm_batch / build_t2vs_batch
out = model.generate_samples_from_batch(batch, seed=[0]) # ← THE call
pixels = model.decode(out["vision"][0]) # VAE decode
# T2VS only — sound output:
# waveform = model.decode_sound(out["sound"][0])
```
### Case 2 — `trainer_level_training.py` (trainer-level training)
What you replace: `cosmos_framework.scripts.train`, the `Trainer` class, callbacks,
FSDP wiring, dataloaders. You keep `model.training_step`, which packages
flow-matching loss + sampling + packing.
```python
model = Cosmos3OmniModel.from_pretrained_dcp(ckpt_dir).model
opt = torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=1e-5)
for it, batch in enumerate(my_loader): # ← your dataloader
aux, loss = model.training_step(batch, iteration=it)
loss.backward()
opt.step(); opt.zero_grad()
```
The demo round-robins through 4 batch builders so you can read the exact
`data_batch` shape `training_step` expects for every modality:
| Helper | Modality | Key fields |
| -------------------------- | ---------- | --------------------------------------------------------------------------------------- |
| `make_text_to_image_batch` | T2I | `images`, `text_token_ids`, `image_size`, `fps` |
| `make_text_to_video_batch` | T2V | `video`, `text_token_ids`, `image_size`, `fps`, `num_frames` |
| `make_action_fdm_batch` | Action FDM | + `action`, `domain_id`, `raw_action_dim`, `mode`, `sequence_plan` |
| `make_sound_video_batch` | T2VS | + `sound` (stereo @ 48 kHz, multiple of AVAE hop=1920), `sequence_plan(has_sound=True)` |
> **⚠ Gotcha — video shape differs between training and inference batches.**
> Training (`training_step`, `is_preprocessed=True`) expects a **flat** list:
> `batch[model.input_video_key] = [video]` → `[1, C, T, H, W]`.
> Inference (`cosmos_framework.inference.action.build_action_batch`) uses **nested**:
> `batch[model.input_video_key] = [[video]]` (one extra `[]`).
> Copying an inference batch into a training loop fails inside
> `_normalize_video_databatch_inplace` with an opaque error — use the flat
> convention when calling `training_step`.
### Case 3 — `net_level.py` (net-level inference)
What you replace: everything in case 1 *plus* the cosmos_framework sampler (UniPC/EDM).
You write the sampling loop by hand and call `net.forward` per step.
`sample(model, net, batch)` is generic across modalities — it splits the
final flat trajectory back into vision/action/sound chunks using the same
offset layout as `_get_velocity`, and decodes each:
```python
net = model.net # Cosmos3VFMNetwork
seq_plans, gen_clean, cond_tokens, _, xt = model._prepare_inference_data(batch, seed=[0])
for step in range(num_steps): # ← Your sampling loop
t = 1.0 - step / num_steps
v = model._get_velocity(net=net, noise_x=xt, timestep=..., text_tokens=cond_tokens, ...)
xt = [x + dt * v_i for x, v_i in zip(xt, v)]
# Per-modality reshape + decode (offsets mirror _get_velocity's split)
vision_latent = xt[0][:vision_dim].reshape(gen_clean.x0_tokens_vision[0].shape)
pixels = model.decode(vision_latent) # always
# action: xt[0][vision_dim:vision_dim+action_dim].reshape(...) # if has_action
# sound : model.decode_sound(xt[0][...sound_slice].reshape(...)) # if has_sound
```
`sample()` returns `{"pixels", "action"?, "sound_waveform"?}`. Plain Euler,
no CFG — production cosmos_framework uses UniPC + CFG; only the integrator differs.
### Case 4 — `net_level.py` (net-level training)
What you replace: everything in case 2 *plus* the flow-matching loss and
the noise schedule. You write the loss explicitly. Same per-modality batch
builders as Case 2 (T2I / T2V / ACTION_FDM / T2VS) round-robin into one
`train_one_step` that calls `net.forward` directly.
```python
net = model.net
# Build the input contract using cosmos_framework helpers
gen_clean = model.get_data_and_condition(batch, iteration=it)
text_indexes = model._load_and_tokenize_text_data(batch, iteration=it)
seq_plans = build_sequence_plans_from_data_batch(batch, model.input_video_key, model.input_image_key)
sigmas = sample_my_sigmas(gen_clean.batch_size) # ← your noise schedule
packed_seq = model._pack_input_sequence(seq_plans, text_indexes, gen_clean, (sigmas*1000).cpu())
gen_noised = model._add_noise_to_input(gen_clean, packed_seq, sigmas, iteration=it)
model._replace_clean_with_noised(packed_seq, gen_noised); packed_seq.to_cuda()
# The bare-net forward — this is the one line that survives a port
out = net(packed_seq, fps_vision=gen_clean.fps_vision) # ← Your forward call
# Your loss — here flow-matching MSE, but it can be anything
v_pred, v_target = out["preds_vision"], gen_noised.vt_target_vision
loss = sum(F.mse_loss(p.float(), t.float()) for p, t in zip(v_pred, v_target))
loss.backward() # ← Your code
```
---
## 3. What you "extract" at each level
A pure level-A extraction (zero `import cosmos_framework`) is **not feasible without
re-vendoring** — `Cosmos3VFMNetwork.forward` takes a `PackedSequence`, which
~2400 lines of `cosmos_framework/data/vfm/sequence_packing.py` build. These demos show
the realistic options:
| Cosmos surface you keep | Trainer-level | Net-level |
| ----------------------------------------------------------- | :-----------: | :---------------: |
| `Cosmos3OmniModel.from_pretrained_dcp` (loader) | ✅ | ✅ |
| VAE (`model.encode` / `model.decode`) | ✅ | ✅ |
| Text tokenizer (`model.vlm_tokenizer` + `tokenize_caption`) | ✅ | ✅ |
| Sequence packer (`model._pack_input_sequence`) | ✅ | ✅ |
| Noise scheduler (`model._add_noise_to_input`) | ✅ | ❌ (your sigma) |
| Flow-matching loss (`model._compute_losses`) | ✅ | ❌ (your loss) |
| Sampler (`UniPC` / `EDM` in `model.sampler`) | ✅ | ❌ (your sampler) |
| Trainer / callbacks / FSDP / dataloader | ❌ | ❌ |
The "❌" cells are exactly what you replace in net-level integration.
> **Note on underscore-prefixed methods.** Net-level integration depends on
> several `_method` names on `OmniMoTModel` — `_pack_input_sequence`,
> `_load_and_tokenize_text_data`, `_add_noise_to_input`,
> `_replace_clean_with_noised`, `_prepare_inference_data`, `_get_velocity`.
> The underscore is Python convention for "internal," but **these are the
> intended net-level integration surface today** and are exercised by the
> demos in CI. Treat them as stable for integration purposes; if cosmos_framework
> ever promotes them to public names, the demos will be updated.
---
## 4. Running the demos
### Prerequisites
1. **Install cosmos_framework** as a library (`pip install -e .` from the repo root,
or activate the project's `.venv`).
2. **A single ≥ 80 GB GPU.** For training, the demos use SGD (zero optimizer
state); switching to AdamW for the full 8 B model OOMs on one 80 GB GPU.
3. **HF cache access** for the auxiliary sub-models — Qwen3-VL tokenizer,
Wan2.2 VAE, AVAE — and the Cosmos3-Nano `config.json` (single ~5 KB file).
The main ~30 GB transformer DCP is **not** downloaded; the demos run with
random main-transformer weights.
### Common flags
```bash
PYTHONPATH=. python examples/integration/<demo>.py # fetches config.json
PYTHONPATH=. python examples/integration/<demo>.py --config-dir /path/with/config.json # local config
```
### Verified runs (single H100 80 GB)
All four modalities run end-to-end in every demo. Output shapes are
deterministic (driven by the config + input shape), but **pixel / sound /
loss values are not meaningful** because the main transformer is random:
| Demo / mode | Output shape (verified) |
| ---------------------------------------------- | ------------------------------------------------- |
| `trainer_level_inference.py --mode t2i` | `pixels [3, 1, 128, 128]` |
| `trainer_level_inference.py --mode t2v` | `pixels [3, 33, 128, 128]` |
| `trainer_level_inference.py --mode action_fdm` | `pixels [3, 5, 128, 128]` |
| `trainer_level_inference.py --mode t2vs` | `pixels [3, 5, 128, 128]` + `sound [2, 15360]` |
| `trainer_level_training.py --num-iters 4` | 4 iters round-robin T2I / T2V / ACTION_FDM / T2VS |
| `net_level.py --sample-mode t2i` | `pixels [3, 1, 128, 128]` |
| `net_level.py --sample-mode t2v` | `pixels [3, 17, 128, 128]` |
| `net_level.py --sample-mode action_fdm` | `pixels [3, 5, 128, 128]` + `action [4, 64]` |
| `net_level.py --sample-mode t2vs` | `pixels [3, 5, 128, 128]` + `sound [2, 15360]` |
> **Why t2v differs:** `trainer_level_inference.py` defaults to `--num-frames 33`
> (matches cosmos_framework's default sample args), while `net_level.py` defaults to
> 17 frames inside `make_text_to_video_batch` to keep the net-level demo
> fast. Same model, same code path — only the batch's `num_frames` differs.
```bash
# Point HF_HOME at a writable cache (any path); aux sub-models + the
# Cosmos3-Nano config.json auto-download into $HF_HOME/hub/... on first use.
export HF_HOME=$HOME/cosmos_assets/hf_cache
# Case 1 — trainer-level inference (default: t2i)
PYTHONPATH=. .venv/bin/python examples/integration/trainer_level_inference.py
# Other modes:
# --mode t2v --num-frames 33
# --mode action_fdm
# --mode t2vs
# Case 2 — trainer-level training, round-robins through all 4 modalities
PYTHONPATH=. .venv/bin/python examples/integration/trainer_level_training.py \
--num-iters 4
# Cases 3 + 4 — net-level training + Euler sampling for a chosen mode
PYTHONPATH=. .venv/bin/python examples/integration/net_level.py \
--num-train-iters 4 --num-sample-steps 8 \
--sample-mode t2i # or t2v / action_fdm / t2vs
```
To run against a non-default config (e.g. Cosmos3-Super) point `--config-dir`
at a directory containing that model's `config.json`.
---
## 5. Where to look next in the cosmos_framework source
| Topic | File |
| ------------------------------- | ------------------------------------------------------------------------------ |
| OmniMoTModel definition | `cosmos_framework/model/vfm/omni_mot_model.py` |
| Cosmos3VFMNetwork (`model.net`) | `cosmos_framework/model/vfm/mot/cosmos3_vfm_network.py` |
| PackedSequence + packer | `cosmos_framework/data/vfm/sequence_packing.py` |
| Rectified-flow loss | `cosmos_framework/model/vfm/algorithm/loss/flow_matching.py` |
| UniPC / EDM samplers | `cosmos_framework/model/vfm/diffusion/samplers/` |
| Checkpoint loader | `cosmos_framework/inference/model.py` (`Cosmos3OmniModel.from_pretrained_dcp`) |
| Default sample args | `cosmos_framework/inference/defaults/<mode>/sample_args.json` |
| FSDP / parallelism wrapping | `cosmos_framework/utils/vfm/parallelism.py` (`ParallelDims`) |
| Production trainer (skipped) | `cosmos_framework/scripts/train.py`, `examples/toml/*.toml` |
|