cosmos-framework/examples/integration/net_level.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""
"net-only" demo: drive `Cosmos3VFMNetwork` (= `model.net`) directly.

================================================================================
⚠  THIS IS A WIRING DEMO, NOT A PRODUCTION RECIPE
================================================================================
The code below shows HOW to call `net.forward` and where to put your loss /
sampler — it does NOT reproduce cosmos_framework's training or sampling recipe.
Concretely, this demo deliberately simplifies four things:

    1. WEIGHTS — `model.net` is RANDOM-INITIALIZED. The demo never loads the
       ~30 GB Cosmos3-Nano DCP shards, so losses and samples are meaningless;
       the point is to show the call sequence. For real weight loading see
       `cosmos_framework.inference.model.Cosmos3OmniModel.from_pretrained_dcp`.
    2. NOISE SCHEDULE — uses a fixed σ = 0.5 every iter.
       Real training samples σ from a logit-normal (image) / waver (video)
       distribution per `OmniMoTModel._get_train_noise_level_vision`.
    3. LOSS — plain MSE on velocity.
       Real training uses `cosmos_framework.model.vfm.algorithm.loss.flow_matching
       .compute_flow_matching_loss`, which adds per-sample weighting,
       condition-mask zeroing, and `loss_scale=10` (with separate image/video
       scaling).
    4. SAMPLER — plain Euler, no CFG, ~8 steps.
       Real inference uses UniPC (`cosmos_framework.model.vfm.diffusion.samplers.unipc`)
       with `guidance=1.5` and 35 steps.

If you train or sample with the demo's simplifications you will diverge from
the Cosmos3 recipe. Use this file to learn the API surface, then swap in the
real weights + loss + sampler when porting to your own framework.

================================================================================
WHAT THIS SHOWS
================================================================================
The previous demos (trainer_level_inference.py / trainer_level_training.py) call
`OmniMoTModel.generate_samples_from_batch` and `OmniMoTModel.training_step`,
which are 2000+ line orchestration methods.

This demo goes one level deeper: it extracts the *core denoiser network*

    net = model.net          # type: Cosmos3VFMNetwork (a plain nn.Module)

and calls `net.forward(packed_seq, fps_vision=...)` itself, then writes the
flow-matching loss and the sampling loop by hand. The point: `net` is the
unit you would port into another training framework — the surrounding
`OmniMoTModel` is just orchestration around it.

================================================================================
THE 3 LAYERS
================================================================================
  ┌────────────────────────────────────────────────────────────────────────────
  │ OmniMoTModel  (model)                                                    │
  │   ├── training_step(), generate_samples_from_batch()  ← orchestration    │
  │   ├── encode/decode = VAE                             ← cosmos_framework VAE       │
  │   ├── _pack_input_sequence(...)                       ← cosmos_framework packer    │
  │   │                                                                       │
  │   └── net = Cosmos3VFMNetwork  ◀──── THIS DEMO'S FOCUS                    │
  │           forward(packed_seq, fps_vision=...) -> {"preds_vision": ...}    │
  └────────────────────────────────────────────────────────────────────────────

`net`'s I/O contract:
  INPUT  : a `PackedSequence` (text tokens + noised vision latents +
           attention modes + position ids), built via
           `model._pack_input_sequence(...)`.
  OUTPUT : dict with `preds_vision` = list of velocity tensors [C, T, H, W],
           one per sample.

================================================================================
WHAT IS STILL "COSMOS" IN THIS DEMO
================================================================================
To USE `net`, you still need to BUILD its input. We use these cosmos_framework helpers
to do that — they are unavoidable unless you re-implement the packer:

    model.encode(...) / model.decode(...)        ← VAE (pixels ↔ latents)
    model._load_and_tokenize_text_data(...)      ← text tokenization
    build_sequence_plans_from_data_batch(...)    ← per-sample modality plan
    model._pack_input_sequence(...)              ← builds the PackedSequence
    model._add_noise_to_input(...)               ← rectified-flow noising
    model._replace_clean_with_noised(...)        ← splice xt into packed_seq
    model._prepare_inference_data(...)           ← inference-time data prep
    model._get_velocity(net=net, ...)            ← inference per-step helper
                                                    (just packs + calls net)

If you wanted ZERO cosmos_framework imports at runtime, you would re-vendor
`cosmos_framework/data/vfm/sequence_packing.py` and the VAE into your own framework.

================================================================================
RUN
================================================================================
    PYTHONPATH=. python examples/integration/net_level.py
    PYTHONPATH=. python examples/integration/net_level.py --config-dir /path/to/dir/with/config.json
"""

from cosmos_framework.inference.common.init import init_script

init_script(training=True)

import argparse
import json
from pathlib import Path

import attrs
import torch
import torch.nn.functional as F

from cosmos_framework.configs.base.defaults.compile import CompileConfig
from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
from cosmos_framework.data.vfm.action.domain_utils import get_domain_id
from cosmos_framework.data.vfm.action.transforms import build_sequence_plan_from_mode
from cosmos_framework.data.vfm.sequence_packing import SequencePlan, build_sequence_plans_from_data_batch
from cosmos_framework.inference.args import DEFAULT_CHECKPOINT
from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel
from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption


def _load_omni_model(*, config_dir_arg: str | None):
    """Build OmniMoTModel with RANDOM main-transformer weights — wiring demo only.

    This helper exists so the demo can run without downloading the ~30 GB transformer
    DCP. Only ``config.json`` is fetched (single ~5 KB file) and the main net is
    instantiated via ``hydra.utils.instantiate`` with random parameters. Auxiliary
    sub-models (Qwen3-VL tokenizer, Wan2.2 VAE, AVAE) still load from the HF cache
    during ``Cosmos3OmniModel.__init__`` — they are not stubbed out.

    For REAL weight loading, see
    :func:`cosmos_framework.inference.model.Cosmos3OmniModel.from_pretrained_dcp`.
    """
    if config_dir_arg is None:
        from huggingface_hub import hf_hub_download
        config_dir = Path(hf_hub_download(
            repo_id=DEFAULT_CHECKPOINT.hf.repository,
            filename="config.json",
            revision=DEFAULT_CHECKPOINT.hf.revision,
        )).parent
    else:
        config_dir = Path(config_dir_arg)
    # Shipped DCPs nest config.json one level deeper under model/.
    if not (config_dir / "config.json").exists() and (config_dir / "model" / "config.json").exists():
        config_dir = config_dir / "model"
    print(f"Loading config from: {config_dir / 'config.json'}")

    # Shipped configs carry stale `cosmos3._src.*` dotted module strings in `_type` / `_target_`
    # fields. cosmos_framework's CONFIG_REPLACEMENTS_INVERSE only rewrites the slash-form
    # paths, so we rewrite the dotted form here before constructing the config.
    config_text = (config_dir / "config.json").read_text()
    for _old, _new in [
        ("cosmos3._src.vfm.configs.base.", "cosmos_framework.configs.base."),
        ("cosmos3._src.vfm.models.", "cosmos_framework.model.vfm."),
        ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.vfm.tokenizers."),
        ("cosmos3._src.imaginaire.", "cosmos_framework."),
    ]:
        config_text = config_text.replace(_old, _new)
    config = Cosmos3OmniConfig(model=json.loads(config_text)["model"])
    config.parallelism = attrs.asdict(ParallelismConfig())
    config.compile = attrs.asdict(CompileConfig(enabled=False))
    return Cosmos3OmniModel(config).model


# ────────────────────────────────────────────────────────────────────────────
# Batch builders — same shapes as trainer_level_training.py uses.
# ────────────────────────────────────────────────────────────────────────────
def _tokenize(model, caption: str, device) -> torch.Tensor:
    ids = tokenize_caption(
        caption, model.vlm_tokenizer,
        is_video=False, use_system_prompt=model.vlm_config.use_system_prompt,
    )
    return torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)  # [1, N_tok]


def make_text_to_image_batch(model, *, caption: str, h: int = 128, w: int = 128, device="cuda") -> dict:
    image = (torch.randn(1, 3, 1, h, w, device=device) * 0.3).clamp(-1, 1)
    return {
        model.input_image_key:   [image],
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([1], device=device),
        "is_preprocessed":  True,
    }


def make_text_to_video_batch(model, *, caption: str, num_frames: int = 17,
                     h: int = 128, w: int = 128, device="cuda") -> dict:
    video = (torch.randn(1, 3, num_frames, h, w, device=device) * 0.3).clamp(-1, 1)
    return {
        model.input_video_key:   [video],
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([num_frames], device=device),
        "is_preprocessed":  True,
    }


def make_sound_video_batch(model, *, caption: str, num_video_frames: int = 5,
                           audio_hop_count: int = 8, h: int = 128, w: int = 128,
                           device="cuda") -> dict:
    """Joint text→video+sound (t2vs). See trainer_level_training.py for full contract."""
    waveform = (torch.randn(2, audio_hop_count * 1920, device=device) * 0.1).clamp(-1, 1)
    video = (torch.randn(1, 3, num_video_frames, h, w, device=device) * 0.3).clamp(-1, 1)
    sequence_plan = SequencePlan(has_text=True, has_vision=True, has_sound=True)
    return {
        model.input_video_key:   [video],
        "sound":                 [waveform],
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([num_video_frames], device=device),
        "sequence_plan":    [sequence_plan],
        "is_preprocessed":  True,
    }


def make_action_fdm_batch(model, *, caption: str, num_video_frames: int = 5,
                          action_chunk: int = 4, raw_action_dim: int = 7,
                          h: int = 128, w: int = 128,
                          domain_name: str = "bridge_orig_lerobot", device="cuda") -> dict:
    """Forward-dynamics action batch. See trainer_level_training.py for the contract."""
    video = (torch.randn(1, 3, num_video_frames, h, w, device=device) * 0.3).clamp(-1, 1)
    action = torch.zeros(action_chunk, model.config.max_action_dim, device=device)
    action[:, :raw_action_dim] = torch.randn(action_chunk, raw_action_dim, device=device) * 0.1
    sequence_plan = build_sequence_plan_from_mode(
        mode="forward_dynamics",
        video_length=num_video_frames,
        action_length=action_chunk,
        has_text=True,
    )
    return {
        model.input_video_key:   [video],
        "action":                [action],
        "raw_action_dim":        [torch.tensor(raw_action_dim, dtype=torch.long, device=device)],
        "mode":                  ["forward_dynamics"],
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([num_video_frames], device=device),
        "domain_id":        [torch.tensor(get_domain_id(domain_name), dtype=torch.long, device=device)],
        "sequence_plan":    [sequence_plan],
        "is_preprocessed":  True,
    }


# ════════════════════════════════════════════════════════════════════════════
# TRAINING: forward + backward through `net` only, custom flow-matching loss.
# ════════════════════════════════════════════════════════════════════════════
def train_one_step(model, net, batch, *, iteration: int) -> torch.Tensor:
    """One rectified-flow training step using only `net.forward`.

    Equivalent to a single `model.training_step(batch, iteration)` call, but
    with the network forward, the loss, and the backward all written here so
    you can see — and replace — each piece.
    """
    # ── 1. Build the input contract for `net`. These calls are all cosmos_framework
    #       preprocessing — they go away if you re-implement the packer.
    input_text_indexes = model._load_and_tokenize_text_data(batch, iteration)
    sequence_plans = build_sequence_plans_from_data_batch(
        data_batch=batch,
        input_video_key=model.input_video_key,
        input_image_key=model.input_image_key,
    )
    gen_data_clean = model.get_data_and_condition(batch, iteration=iteration)

    # Pick a mid-range noise level for the demo (real training samples sigma
    # from a per-modality distribution; see cosmos_framework.model.vfm.omni_mot_model
    # `_get_train_noise_level_vision`).
    B = gen_data_clean.batch_size
    # tensor_kwargs_fp32 = {"dtype": float32, "device": ...} — keeps demo
    # tensors on the same device / dtype the model expects.
    sigmas = torch.full((B, 1), 0.5, **model.tensor_kwargs_fp32)        # [B, 1]
    timesteps = (sigmas * 1000.0).cpu()                                 # [B, 1] on cpu

    packed_seq = model._pack_input_sequence(
        sequence_plans, input_text_indexes, gen_data_clean, timesteps,
    )
    gen_data_noised = model._add_noise_to_input(
        gen_data_clean, packed_seq, sigmas, iteration=iteration,
    )
    model._replace_clean_with_noised(packed_seq, gen_data_noised)
    packed_seq.to_cuda()

    # ── 2. THE bare-net forward pass. This is the single line that you would
    #       call from your own training loop after porting `net` into it.
    out = net(packed_seq, fps_vision=gen_data_clean.fps_vision)         # type: dict
    v_pred = out["preds_vision"]                                        # list of [C, T, H, W]

    # ── 3. Custom flow-matching loss (MSE on velocity). This is what
    #       `cosmos_framework.model.vfm.algorithm.loss.flow_matching.compute_flow_matching_loss`
    #       computes, minus the per-sample weighting & condition masking.
    v_target = gen_data_noised.vt_target_vision                         # list of [C, T, H, W]
    loss = sum(F.mse_loss(p.float(), t.float()) for p, t in zip(v_pred, v_target))

    # ── 4. Backward. Your code, not cosmos_framework's.
    loss.backward()
    return loss


# ════════════════════════════════════════════════════════════════════════════
# INFERENCE: hand-written Euler sampling loop, each step a `net.forward`.
# Generic across modalities: returns whatever the batch's sequence plans say
# is present (vision always; action and/or sound when configured).
# ════════════════════════════════════════════════════════════════════════════
@torch.no_grad()
def sample(model, net, batch, *, num_steps: int = 12) -> dict:
    """N-step Euler integration of dx/dt = v(x,t) — no cosmos_framework sampler involved.

    Production cosmos_framework uses UniPC/EDM (cosmos_framework.model.vfm.diffusion.samplers.*).
    Plain Euler keeps the loop on one screen and surfaces where `net` is called.

    Returns a dict:
        "pixels":         Tensor[3, T, H, W] in [0,1] — always present
        "action":         Tensor[T_action, action_dim] — only if has_action
        "sound_waveform": Tensor[C_audio, N_samples]    — only if has_sound

    For ACTION_FDM and T2VS the demo's batch uses random conditioning, so
    these outputs are noise. The wiring is what's being demonstrated.
    """
    # `_prepare_inference_data` tokenizes cond+uncond captions, builds the per-
    # sample `sequence_plans`, and constructs the initial noise tensor. The
    # noise layout per sample (flat [D]) is concatenation of:
    #     [vision_flat | action_flat (if has_action) | sound_flat (if has_sound)]
    # — same layout `_get_velocity` consumes.
    sequence_plans, gen_data_clean, cond_tokens, _, initial_noise = model._prepare_inference_data(
        batch, seed=[0], has_negative_prompt=False,
    )

    xt = initial_noise                                                   # list[B=1] of flat [D]
    dt = -1.0 / num_steps  # integrate from t=1 (noise) → t=0 (clean)

    for step in range(num_steps):
        t_now = 1.0 - step / num_steps
        timestep = torch.tensor([[t_now * 1000.0]], **model.tensor_kwargs_fp32)  # [1, 1]

        # `_get_velocity` (a) reshapes `xt` back to per-modality tokens,
        # (b) packs them into a PackedSequence with the current timestep,
        # (c) calls `net(packed_seq, ...)`, (d) returns flat velocity.
        v = model._get_velocity(
            net=net,
            noise_x=xt,
            timestep=timestep,
            text_tokens=cond_tokens,
            sequence_plans=sequence_plans,
            gen_data_clean=gen_data_clean,
        )
        xt = [x + dt * v_i for x, v_i in zip(xt, v)]

    # Split the final flat trajectory back into per-modality tensors.
    # Offsets exactly mirror `_get_velocity`'s split.
    has_action = model.config.action_gen and any(p.has_action for p in sequence_plans)
    has_sound  = model.config.sound_gen and any(p.has_sound for p in sequence_plans)

    flat = xt[0]
    offset = 0

    vision_shape = gen_data_clean.x0_tokens_vision[0].shape               # [1, C, T, H, W]
    vision_dim = int(torch.tensor(vision_shape).prod())
    vision_latent = flat[offset : offset + vision_dim].reshape(vision_shape)
    offset += vision_dim

    out: dict = {}
    pixels = model.decode(vision_latent)                                  # [1, 3, T, H, W] in [-1, 1]
    out["pixels"] = (pixels[0].clamp(-1, 1) + 1.0) / 2.0                  # [3, T, H, W] in [0, 1]

    if has_action and gen_data_clean.x0_tokens_action is not None:
        action_shape = gen_data_clean.x0_tokens_action[0].shape           # [T_action, action_dim]
        action_dim = int(torch.tensor(action_shape).prod())
        out["action"] = flat[offset : offset + action_dim].reshape(action_shape)
        offset += action_dim

    if has_sound and gen_data_clean.x0_tokens_sound is not None and sequence_plans[0].has_sound:
        sound_shape = gen_data_clean.x0_tokens_sound[0].shape             # [C_sound, T_sound]
        sound_dim = int(torch.tensor(sound_shape).prod())
        sound_latent = flat[offset : offset + sound_dim].reshape(sound_shape)
        offset += sound_dim
        out["sound_waveform"] = model.decode_sound(sound_latent)          # [C_audio, N_samples]

    return out


# ════════════════════════════════════════════════════════════════════════════
# main
# ════════════════════════════════════════════════════════════════════════════
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config-dir", type=str, default=None,
                        help="Local directory containing config.json (architecture only — weights are "
                             "randomly initialized). If omitted, fetches Cosmos3-Nano's config.json from HF.")
    parser.add_argument("--num-train-iters", type=int, default=2)
    parser.add_argument("--num-sample-steps", type=int, default=12)
    parser.add_argument("--sample-mode", type=str, default="t2i",
                        choices=["t2i", "t2v", "action_fdm", "t2vs"],
                        help="Which modality to sample. action_fdm/t2vs use random conditioning → noise output.")
    parser.add_argument("--skip-sample", action="store_true",
                        help="Skip the inference section (saves ~1 min).")
    args = parser.parse_args()

    output_dir = Path("outputs/net_level").absolute()
    output_dir.mkdir(parents=True, exist_ok=True)

    # 1) Build OmniMoTModel (random weights — see module docstring) + grab the bare net
    model = _load_omni_model(config_dir_arg=args.config_dir)

    net = model.net  # ← Cosmos3VFMNetwork; THIS is the unit you'd port
    print(f"net type:   {type(net).__name__}")
    print(f"net params: {sum(p.numel() for p in net.parameters()) / 1e9:.2f} B")

    # 2) TRAINING — forward + backward through `net` ------------------------
    net.train()
    optimizer = torch.optim.SGD([p for p in net.parameters() if p.requires_grad], lr=1e-5)

    caption_img = "A neon city street at night, rain reflecting the signs."
    caption_vid = "A camera dollies through a forest of giant glowing mushrooms."
    caption_act = "A robot arm picks up a red block from the table."
    caption_snd = "Wind howling through pine trees, distant thunder."

    def next_batch(it: int):
        kind = ["T2I", "T2V", "ACTION_FDM", "T2VS"][it % 4]
        if kind == "T2I":
            return (kind, make_text_to_image_batch(model, caption=caption_img))
        if kind == "T2V":
            return (kind, make_text_to_video_batch(model, caption=caption_vid))
        if kind == "ACTION_FDM":
            return (kind, make_action_fdm_batch(model, caption=caption_act))
        return (kind, make_sound_video_batch(model, caption=caption_snd))

    for it in range(args.num_train_iters):
        kind, batch = next_batch(it)
        loss = train_one_step(model, net, batch, iteration=it)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        print(f"[train] iter {it:>3d}  [{kind}]  loss={loss.item():.4f}")

    # 3) INFERENCE — sampling loop where each step is a `net.forward` -------
    if not args.skip_sample:
        net.eval()

        sample_caption = {
            "t2i":        "A robot standing on a rooftop at sunset.",
            "t2v":        "A camera dollies through a forest of giant glowing mushrooms.",
            "action_fdm": "A robot arm picks up a red block from the table.",
            "t2vs":       "Wind howling through pine trees, distant thunder.",
        }[args.sample_mode]
        sample_builder = {
            "t2i":        lambda: make_text_to_image_batch(model, caption=sample_caption),
            "t2v":        lambda: make_text_to_video_batch(model, caption=sample_caption),
            "action_fdm": lambda: make_action_fdm_batch(model, caption=sample_caption),
            "t2vs":       lambda: make_sound_video_batch(model, caption=sample_caption),
        }[args.sample_mode]

        out = sample(model, net, sample_builder(), num_steps=args.num_sample_steps)

        from cosmos_framework.tools.visualize.video import save_img_or_video
        sample_dir = output_dir / f"sample_{args.sample_mode}"
        sample_dir.mkdir(parents=True, exist_ok=True)
        save_img_or_video(out["pixels"], str(sample_dir / "output"), fps=16.0)
        print(f"[infer] {args.sample_mode}: pixels saved to {sample_dir / 'output'}")

        if "sound_waveform" in out:
            torch.save(out["sound_waveform"].cpu(), sample_dir / "sound.pt")
            print(f"[infer] {args.sample_mode}: sound shape={tuple(out['sound_waveform'].shape)} → sound.pt")
        if "action" in out:
            torch.save(out["action"].cpu(), sample_dir / "action.pt")
            print(f"[infer] {args.sample_mode}: action shape={tuple(out['action'].shape)} → action.pt")


if __name__ == "__main__":
    main()