[feat] v2: Cosmos-Predict2-2B-Video2World port (EDM-Karras) — reference bucket-C recipe

First net-new architecture ported via the self-contained recipe-package pattern
(card-declared adapter + new loop, no shared-dispatch edit):

* CosmosDenoiseLoop (v2/recipes/cosmos2/loop.py): EDM preconditioning folded into
  a flow-match Euler integrator. Faithful port of CosmosDenoisingStage — Karras
  sigma schedule (rho=7, sigma_max=80 -> sigma_min=0.002, terminal clamp), latent
  init randn*sigma_max, per-step c_in/c_skip/c_out (sigma_data=1) -> x0, CFG in x0
  space, x0 -> velocity (x-x0)/sigma, FLOW_MATCH_STEP. video2world frame-replace
  conditioning threaded but inert for the t2v preset.
* build_karras_sigmas helper added to v2/loop/sampler.py.
* CosmosDiT + CosmosT5Encoder adapters (v2/platform/backends/torch_cosmos.py),
  declared on the card via ComponentSpec.adapter (the ac29750b enabler) — DiT
  returns raw EDM output + builds the mandatory zero condition/padding masks + fps;
  T5 uses the raw last_hidden_state (no Wan zero-pad). Reuses the WanVAE adapter.
* card/program/registry (HF id nvidia/Cosmos-Predict2-2B-Video2World + arch
  fallback on CosmosTransformer3DModel) + COSMOS_NEG prompt + SamplingDefaults
  (35 steps, gs 7, 704x1280, 93f, 16fps).

CPU-verified: Karras schedule, card/program build, registry resolve (id + arch),
EDM loop runs end-to-end on the CPU toy backend. GPU load/run is BRINGUP.
v2 mini green (211 passed, 2 skipped); pre-commit clean.

v2/loop/sampler.py

@@ -29,8 +29,27 @@ def x0_from_velocity(x_t: np.ndarray, velocity: np.ndarray, sigma_t: float) -> n
     return x_t - sigma_t * velocity
 
 
-def flow_match_euler_step(x_t: np.ndarray, velocity: np.ndarray,
-                          sigma_t: float, sigma_next: float) -> np.ndarray:
+def build_karras_sigmas(num_steps: int,
+                        sigma_max: float = 80.0,
+                        sigma_min: float = 0.002,
+                        rho: float = 7.0) -> np.ndarray:
+    """Karras et al. (2022) ρ-interpolated σ schedule + the terminal σ, as Cosmos/EDM models use it.
+
+    Reproduces ``FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)`` the Cosmos pipeline configures
+    (``sigma_max=80, sigma_min=0.002, rho=7``): a length-``num_steps`` ramp ``σ_max→σ_min`` via
+    ``σ_i = (max_inv_rho + i/(n-1)·(min_inv_rho-max_inv_rho))^ρ``, then ONE appended terminal value. The
+    scheduler appends ``0.0`` and, with ``final_sigmas_type='sigma_min'``, the Cosmos stage overwrites it
+    with ``σ[-2]`` to avoid a divide-by-zero in the EDM coeffs / velocity — so the terminal here is
+    ``σ_min`` (the last ramp value), giving ``num_steps+1`` points the EDM loop integrates pairwise.
+    """
+    ramp = np.linspace(0.0, 1.0, num_steps, dtype=np.float64)
+    min_inv_rho = sigma_min**(1.0 / rho)
+    max_inv_rho = sigma_max**(1.0 / rho)
+    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho  # σ_max -> σ_min
+    return np.concatenate([sigmas, sigmas[-1:]]).astype(np.float64)  # + sigma_min terminal clamp
+
+
+def flow_match_euler_step(x_t: np.ndarray, velocity: np.ndarray, sigma_t: float, sigma_next: float) -> np.ndarray:
     """One deterministic Euler step along the flow ODE."""
     return x_t + (sigma_next - sigma_t) * velocity
 
@@ -40,8 +59,14 @@ def add_noise(x0: np.ndarray, noise: np.ndarray, sigma: float) -> np.ndarray:
     return (1.0 - sigma) * x0 + sigma * noise
 
 
-def flow_sde_step_with_logprob(x_t, velocity, sigma_t: float, sigma_next: float, *,
-                               noise=None, prev_sample=None, noise_scale: float = 0.7):
+def flow_sde_step_with_logprob(x_t,
+                               velocity,
+                               sigma_t: float,
+                               sigma_next: float,
+                               *,
+                               noise=None,
+                               prev_sample=None,
+                               noise_scale: float = 0.7):
     """FlowGRPO-style stochastic step + Gaussian log-prob (PromptRL/UniRL §6; FlowGRPO lineage).
 
     The RL *rollout* sampler (vs the deterministic ODE ``flow_match_euler_step`` used at serve time).
@@ -53,18 +78,17 @@ def flow_sde_step_with_logprob(x_t, velocity, sigma_t: float, sigma_next: float,
     x_t = np.asarray(x_t, dtype=np.float64)
     velocity = np.asarray(velocity, dtype=np.float64)
     s = min(float(sigma_t), 0.9999)
-    dt = float(sigma_next) - float(sigma_t)                       # negative (σ decreases)
+    dt = float(sigma_next) - float(sigma_t)  # negative (σ decreases)
     std = float(np.sqrt(s / (1.0 - s)) * noise_scale)
     denom = 2.0 * max(s, 1e-6)
-    mean = x_t * (1.0 + std ** 2 / denom * dt) + velocity * (1.0 + std ** 2 * (1.0 - s) / denom) * dt
+    mean = x_t * (1.0 + std**2 / denom * dt) + velocity * (1.0 + std**2 * (1.0 - s) / denom) * dt
     eff_std = max(std * np.sqrt(max(-dt, 1e-12)), 1e-6)
     if prev_sample is None:
         n = noise if noise is not None else np.zeros_like(x_t)
         prev_sample = mean + eff_std * np.asarray(n, dtype=np.float64)
     prev_sample = np.asarray(prev_sample, dtype=np.float64)
-    var = eff_std ** 2
-    log_prob = float(np.mean(-((prev_sample - mean) ** 2) / (2.0 * var)
-                             - np.log(eff_std) - 0.5 * np.log(2.0 * np.pi)))
+    var = eff_std**2
+    log_prob = float(np.mean(-((prev_sample - mean)**2) / (2.0 * var) - np.log(eff_std) - 0.5 * np.log(2.0 * np.pi)))
     return prev_sample.astype("float32"), log_prob, mean.astype("float32"), float(eff_std)
 
 
@@ -80,7 +104,7 @@ def flow_sde_ml_velocity(x_t, sample, sigma_t: float, sigma_next: float, *, nois
     dt = float(sigma_next) - float(sigma_t)
     std = float(np.sqrt(s / (1.0 - s)) * noise_scale)
     denom = 2.0 * max(s, 1e-6)
-    a = 1.0 + std ** 2 / denom * dt
-    b = (1.0 + std ** 2 * (1.0 - s) / denom) * dt
+    a = 1.0 + std**2 / denom * dt
+    b = (1.0 + std**2 * (1.0 - s) / denom) * dt
     b = b if abs(b) > 1e-8 else (1e-8 if b >= 0.0 else -1e-8)
     return ((sample - x_t * a) / b).astype("float32")

v2/platform/backends/torch_cosmos.py

@@ -0,0 +1,61 @@
+"""Cosmos-Predict2 torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter``
+so the Cosmos recipe is self-contained (no edit to the shared ``_make_dit``/``_make_text_encoder``
+dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box.
+
+* ``CosmosDiT`` — the EDM denoiser ``F_θ``. The loop hands the *already EDM-input-scaled* model input
+  (``x·c_in``) and the model timestep ``t = σ·1000``; this adapter returns the **raw** transformer
+  output (the EDM ``c_skip``/``c_out`` → x0 reconstruction + x0-space CFG live in ``CosmosDenoiseLoop``,
+  NOT here). It builds the mandatory zero ``condition_mask`` / ``padding_mask`` and ``fps`` the Cosmos
+  forward requires (the model concats the masks internally → 18ch patch_embed input). Faithful to
+  ``CosmosDenoisingStage._run_transformer`` (fastvideo/pipelines/stages/denoising.py).
+* ``CosmosT5Encoder`` — T5-Large (1024-dim). Cosmos uses the **raw** last_hidden_state (NaN→0, no
+  fixed-length zero-pad), unlike the Wan T5 convention that zero-pads to 512.
+"""
+from __future__ import annotations
+
+import torch
+
+from v2.platform.backends.torch_backend import T5Encoder, TorchComponent, _to_numpy
+
+
+class CosmosDiT(TorchComponent):
+    """``dit(model_input[C,T,h,w], text_embed[seq,1024], timestep) -> raw noise_pred[C,T,h,w]``.
+
+    ``model_input`` is pre-scaled by the loop (``x·c_in``); ``timestep`` is ``σ·1000`` (1D [B] float)."""
+
+    @torch.no_grad()
+    def __call__(self, model_input, text_embed, timestep, context=None, *, cond=None):
+        hs = self._t(model_input)  # [1, C, T, h, w]
+        ehs = self._t(text_embed)
+        b, _c, t, h, w = hs.shape
+        ts = torch.tensor([float(timestep)], device=self.device, dtype=self.dtype).expand(b)
+        condition_mask = torch.zeros(b, 1, t, h, w, device=self.device, dtype=self.dtype)  # t2v: zeros
+        padding_mask = torch.zeros(1, 1, h, w, device=self.device, dtype=self.dtype)  # resized→[h,w]
+        with self._ctx():
+            out = self.module(hidden_states=hs,
+                              timestep=ts,
+                              encoder_hidden_states=ehs,
+                              fps=24,
+                              condition_mask=condition_mask,
+                              padding_mask=padding_mask,
+                              return_dict=False)[0]
+        return self._n(out)  # RAW EDM output (loop reconstructs x0)
+
+
+class CosmosT5Encoder(T5Encoder):
+    """Cosmos T5-Large: raw last_hidden_state (NaN→0), NO zero-pad-to-max (the Wan convention would
+    mis-condition Cosmos). Truncates to ``max_length`` but returns only the real-token rows."""
+
+    @torch.no_grad()
+    def encode(self, text):
+        toks = self.tokenizer(text or "", return_tensors="pt", max_length=self.max_length, truncation=True)
+        ids = toks.input_ids.to(self.device)
+        mask = toks.attention_mask.to(self.device)
+        with self._ctx():
+            out = self.module(input_ids=ids, attention_mask=mask)
+        hidden = (out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]).squeeze(0)
+        hidden = torch.nan_to_num(hidden, nan=0.0)
+        n = int(mask.sum().item())  # zero any row beyond real-token length
+        if n < hidden.shape[0]:
+            hidden[n:] = 0.0
+        return _to_numpy(hidden)

v2/recipes/_prompts.py

@@ -18,6 +18,15 @@ WAN_NEG_CN = ("色调艳丽，过曝，静态，细节模糊不清，字幕，
               "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，"
               "背景人很多，倒着走")
 
+# Cosmos-Predict2 family negative prompt (verbatim from the Cosmos2 pipeline preset).
+COSMOS_NEG = ("The video captures a series of frames showing ugly scenes, static with no motion, motion"
+              " blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images,"
+              " poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out"
+              " colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding,"
+              " unnatural transitions, outdated special effects, fake elements, unconvincing visuals,"
+              " poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of"
+              " poor quality.")
+
 # LTX-2 family negative prompt (base / 2.3); the distilled few-step presets use "" instead.
 LTX2_NEG = ("blurry, out of focus, overexposed, underexposed, low contrast, "
             "washed out colors, excessive noise, grainy texture, poor lighting, "

v2/recipes/cosmos2/__init__.py

@@ -0,0 +1,14 @@
+"""Cosmos-Predict2 (Video2World) — EDM-Karras denoiser ported into the v2 substrate.
+
+Self-contained recipe package (the bucket-C pattern): card declares its torch adapters via
+``ComponentSpec.adapter`` (``CosmosDiT``/``CosmosT5Encoder`` in ``v2/platform/backends/torch_cosmos.py``)
+and a new ``CosmosDenoiseLoop`` (EDM preconditioning folded into flow-match Euler), reusing the Wan VAE
+adapter + T5 + ``stamp_wan21_checkpoints``. Registered in ``v2/registry.py``.
+"""
+from __future__ import annotations
+
+from v2.recipes.cosmos2.card import build_cosmos2_card
+from v2.recipes.cosmos2.loop import CosmosDenoiseLoop
+from v2.recipes.cosmos2.program import build_cosmos2_program
+
+__all__ = ["build_cosmos2_card", "build_cosmos2_program", "CosmosDenoiseLoop"]

v2/recipes/cosmos2/card.py

@@ -0,0 +1,122 @@
+"""Cosmos-Predict2-2B-Video2World ModelCard — text→video (the registered preset; video2world
+conditioning is a later capability the loop already threads).
+
+Architecture deltas vs Wan (all GPU-path, declared on the card so the recipe is self-contained):
+  * DiT  ``fastvideo.models.dits.cosmos:CosmosTransformer3DModel`` — an **EDM denoiser** (NOT flow-match);
+    the adapter (``ComponentSpec.adapter`` -> ``CosmosDiT``) returns the raw network output and
+    ``CosmosDenoiseLoop`` does the EDM ``c_in/c_skip/c_out`` → x0 reconstruction + x0-space CFG + the
+    Karras (ρ=7, σ_max=80→σ_min=0.002) schedule.
+  * VAE  ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — Wan-style (z=16, 8×/4×); reuses the v2
+    ``WanVAE`` adapter unchanged (sigma_data=1.0 makes the Cosmos *σ_data factor a no-op).
+  * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (T5-Large, 1024-dim) via ``CosmosT5Encoder``
+    (raw last_hidden_state, no Wan zero-pad).
+``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout).
+"""
+from __future__ import annotations
+
+from v2._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind
+from v2.card import (
+    CacheContract,
+    CapabilityMatrix,
+    ComponentSpec,
+    CostModel,
+    LoopSpec,
+    ModelCard,
+    ParallelismContract,
+    ParitySpec,
+    ParityTestSpec,
+    PrecisionContract,
+    RecipeSpec,
+    SamplingDefaults,
+)
+from v2.loop.policies import ClassicCFG, NoRouting, PrecisionPolicy
+from v2.parallel import ParallelPlan
+from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from
+from v2.recipes._prompts import COSMOS_NEG
+from v2.recipes.cosmos2.loop import CosmosDenoiseLoop
+from v2.recipes.wan21.card import stamp_wan21_checkpoints
+
+_COSMOS_DIT = "v2.platform.backends.torch_cosmos:CosmosDiT"
+_COSMOS_T5 = "v2.platform.backends.torch_cosmos:CosmosT5Encoder"
+
+
+def build_cosmos2_card(model_id: str = "cosmos-predict2-2b-video2world",
+                       *,
+                       checkpoint_root: str | None = None,
+                       sampling_defaults: SamplingDefaults | None = None) -> ModelCard:
+    seed = _seed_from(model_id)
+    cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7)
+    cfg = ClassicCFG()
+    precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True)
+    expert = NoRouting("transformer")
+
+    def loop_factory():
+        return CosmosDenoiseLoop(loop_id="diffusion_denoise",
+                                 cfg=cfg,
+                                 precision=precision,
+                                 expert=expert,
+                                 cost=cost,
+                                 sigma_max=80.0,
+                                 sigma_min=0.002,
+                                 sigma_data=1.0,
+                                 rho=7.0,
+                                 augment_sigma=0.001)
+
+    components = {
+        "text_encoder":
+        ComponentSpec(component_id="text_encoder",
+                      kind="text_encoder",
+                      load_id="fastvideo.models.encoders.t5:T5EncoderModel",
+                      adapter=_COSMOS_T5,
+                      factory=lambda inst: ToyTextEncoder(),
+                      required_for={"t2v"}),
+        "vae":
+        ComponentSpec(component_id="vae",
+                      kind="vae",
+                      load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan",
+                      factory=lambda inst: ToyVAE(),
+                      required_for={"t2v"}),
+        "transformer":
+        ComponentSpec(component_id="transformer",
+                      kind="dit",
+                      load_id="fastvideo.models.dits.cosmos:CosmosTransformer3DModel",
+                      adapter=_COSMOS_DIT,
+                      factory=lambda inst: ToyDiT(seed=seed),
+                      resident_for=["diffusion_denoise"],
+                      required_for={"t2v"}),
+    }
+    loops = {
+        "diffusion_denoise":
+        LoopSpec(loop_id="diffusion_denoise",
+                 kind=LoopKind.DIFFUSION_DENOISE,
+                 work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
+                 step_cost_model=cost,
+                 shared_weight_components=["transformer"],
+                 cache_policy=["feature"],
+                 graph_capture="breakable_cudagraph",
+                 loop_factory=loop_factory),
+    }
+    card = ModelCard(
+        model_id=model_id,
+        family="cosmos",
+        components=components,
+        loops=loops,
+        capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.VAE_DECODE, Capability.POLICY_ROLLOUT),
+        recipe=RecipeSpec(method="base",
+                          assumes_loop="diffusion_denoise",
+                          assumes_precision="float32",
+                          consistency_required=ConsistencyLevel.C1),
+        parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1],
+                          interleave_required=True,
+                          tests=[ParityTestSpec(name="denoise_trajectory", level=ConsistencyLevel.C1, tap="latents")]),
+        caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
+        precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
+        parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
+        sampling_defaults=sampling_defaults or SamplingDefaults(
+            num_steps=35, guidance_scale=7.0, height=704, width=1280, num_frames=93, fps=16,
+            negative_prompt=COSMOS_NEG),
+    )
+    card.validate()
+    if checkpoint_root:
+        stamp_wan21_checkpoints(card, checkpoint_root)
+    return card

v2/recipes/cosmos2/loop.py

@@ -0,0 +1,198 @@
+"""CosmosDenoiseLoop — EDM-Karras preconditioning folded into a flow-match Euler integrator.
+
+Cosmos-Predict2 is NOT a plain flow-match model despite the pipeline wrapping a
+``FlowMatchEulerDiscreteScheduler``. The network is an EDM denoiser ``F_θ`` whose preconditioned output
+reconstructs ``x0`` (``x0 = c_skip·x + c_out·F_θ(x·c_in)``, ``sigma_data=1``); CFG combines in **x0
+space**; then ``x0`` is converted to a flow-match velocity ``(x - x0)/σ`` for the deterministic Euler
+update ``x_next = x + (σ_next - σ)·v``. The σ schedule is Karras (ρ=7, σ_max=80 → σ_min=0.002), NOT the
+flow-shift linspace, and the latent is initialized at ``randn·σ_max``. Faithful port of
+``fastvideo/pipelines/stages/denoising.py:CosmosDenoisingStage`` (see that file for the exact math).
+
+video2world (frame_replace) conditioning is threaded (``conditioning_latents``/``cond_indicator`` from
+slots) but is ``None`` for the registered t2v preset → the gated injection branch is inert; the loop
+degrades to pure t2v exactly as the fastvideo pipeline does when no image/video is given.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from v2._enums import ExecutionProfile, WorkUnitKind
+from v2.loop.contracts import (
+    Done,
+    LoopResult,
+    LoopState,
+    ResourceRequest,
+    ShapeSignature,
+    StepContext,
+    StepResult,
+    WorkPlan,
+)
+from v2.loop.sampler import build_karras_sigmas
+from v2.platform import FLOW_MATCH_STEP
+from v2.recipes.wan21.loop import latent_shape
+
+COSMOS_LATENT_CHANNELS = 16  # config.in_channels(17) - 1 (the condition_mask channel)
+COSMOS_TEMPORAL_RATIO = 4
+COSMOS_SPATIAL_RATIO = 8
+
+
+class CosmosDenoiseLoop:
+
+    def __init__(self,
+                 *,
+                 loop_id,
+                 cfg,
+                 precision,
+                 expert,
+                 cost,
+                 sigma_max: float = 80.0,
+                 sigma_min: float = 0.002,
+                 sigma_data: float = 1.0,
+                 rho: float = 7.0,
+                 augment_sigma: float = 0.001,
+                 latent_channels: int = COSMOS_LATENT_CHANNELS,
+                 spatial_ratio: int = COSMOS_SPATIAL_RATIO,
+                 temporal_ratio: int = COSMOS_TEMPORAL_RATIO):
+        self.loop_id = loop_id
+        self.cfg = cfg  # carried for the WorkPlan op-structure key; x0-space CFG is done here
+        self.precision = precision
+        self.expert = expert
+        self.cost = cost
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.sigma_data = sigma_data
+        self.rho = rho
+        self.augment_sigma = augment_sigma
+        self.latent_channels = latent_channels
+        self.spatial_ratio = spatial_ratio
+        self.temporal_ratio = temporal_ratio
+
+    def init(self, req, model, ctx) -> LoopState:
+        seed = req.diffusion.seed if req.diffusion.seed is not None else 0
+        rng = np.random.default_rng(seed)
+        sig = build_karras_sigmas(req.diffusion.num_steps,
+                                  sigma_max=self.sigma_max,
+                                  sigma_min=self.sigma_min,
+                                  rho=self.rho)
+        shape = latent_shape(req,
+                             model,
+                             channels=self.latent_channels,
+                             spatial_ratio=self.spatial_ratio,
+                             temporal_ratio=self.temporal_ratio)
+        x = (rng.standard_normal(shape) * float(self.sigma_max)).astype("float32")  # EDM: randn·σ_max
+        st = LoopState(loop_id=self.loop_id,
+                       instance_id=model.card.model_id,
+                       request_id=req.request_id,
+                       profile=ctx.profile,
+                       rng=rng,
+                       seed=seed,
+                       latents={"video": x},
+                       sigmas=[float(s) for s in sig],
+                       timesteps=[float(s) * 1000.0 for s in sig])  # model timestep = σ·1000 (large for EDM)
+        st.cond["prompt_embeds"] = ctx.slots.get("text_embeds")
+        st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds")
+        st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale)
+        st.scratch["stream_video"] = bool(req.outputs.stream.get("video"))
+        # video2world (frame_replace): VAE-encoded conditioning latents + cond/uncond indicators. None for
+        # the t2v preset -> the gated injection below is skipped (pure t2v, matching the fastvideo pipeline).
+        st.scratch["conditioning_latents"] = ctx.slots.get("conditioning_latents")
+        st.scratch["cond_indicator"] = ctx.slots.get("cond_indicator")
+        st.scratch["uncond_indicator"] = ctx.slots.get("uncond_indicator")
+        st.plugin_state["cfg"] = {}
+        return st
+
+    def next(self, st: LoopState):
+        i = st.step_idx
+        if i >= len(st.sigmas) - 1:
+            return Done()
+        sigma_t, sigma_next = st.sigmas[i], st.sigmas[i + 1]
+        t = st.timesteps[i]  # σ·1000, the model timestep
+        expert_id = self.expert.expert_for(StepContext(i, t, sigma_t))
+        x = st.latents["video"]
+        pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"]
+        scale = st.scratch["guidance_scale"]
+        precision = self.precision
+        sd = float(self.sigma_data)
+        do_cfg = scale != 1.0 and ne is not None
+        cond_lat = st.scratch.get("conditioning_latents")
+        cond_ind = st.scratch.get("cond_indicator")
+        uncond_ind = st.scratch.get("uncond_indicator")
+        aug = float(self.augment_sigma)
+
+        def _x0(model: Any, dit: Any, latent_in: np.ndarray, text_embed: Any, ind: Any) -> np.ndarray:
+            # EDM preconditioning + (gated) frame-replace conditioning, returning the x0 prediction for one
+            # CFG branch. ``ind`` is the (un)cond indicator; None for t2v -> no frame injection.
+            s = float(sigma_t)
+            c_in = 1.0 / (s**2 + sd**2)**0.5
+            c_skip = sd**2 / (s**2 + sd**2)
+            c_out = s * sd / (s**2 + sd**2)**0.5
+            cur_ind = (ind * 0 if (ind is not None and aug >= s) else ind)
+            lat = np.array(latent_in, dtype="float32")
+            if cur_ind is not None and cond_lat is not None:  # video2world frame injection (inert for t2v)
+                c_in_aug = 1.0 / (aug**2 + sd**2)**0.5
+                cn = st.rng.standard_normal(lat.shape).astype("float32")
+                cf = (cond_lat + cn * aug) * c_in_aug / c_in
+                lat = cur_ind * cf + (1 - cur_ind) * lat
+            model_input = (lat * c_in).astype("float32")
+            np_pred = np.asarray(dit(model_input, text_embed, t), dtype="float32")
+            x0 = c_skip * x + c_out * np_pred
+            if cur_ind is not None and cond_lat is not None:
+                x0 = cur_ind * cond_lat + (1 - cur_ind) * x0
+            return x0
+
+        def run(model, override=None):
+            dit = model.component(expert_id)
+            if override is not None and "noise_pred" in override:
+                velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32"))
+            else:
+                cond_x0 = _x0(model, dit, x, pe, cond_ind)
+                if do_cfg:
+                    uncond_x0 = _x0(model, dit, x, ne, uncond_ind)
+                    final_x0 = cond_x0 + scale * (cond_x0 - uncond_x0)
+                else:
+                    final_x0 = cond_x0
+                velocity = precision.cast((x - final_x0) / max(float(sigma_t), 1e-6))
+            x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(precision.cast(x), velocity, sigma_t, sigma_next)
+            return StepResult(output={
+                "noise_pred": np.asarray(velocity, dtype="float32"),
+                "latents": x_next.astype("float32")
+            })
+
+        cond_bytes = sum(int(np.asarray(e).nbytes) for e in (pe, ne) if e is not None)
+        res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(x.shape)), 2.0 if do_cfg else 1.0),
+                              resident_bytes=int(x.nbytes) + cond_bytes,
+                              peak_activation_bytes=int(x.nbytes))
+        return WorkPlan(loop_id=self.loop_id,
+                        instance_id=st.instance_id,
+                        kind=WorkUnitKind.DIFFUSION_STEP,
+                        shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP,
+                                                 dims=tuple(x.shape),
+                                                 dtype=precision.compute_dtype,
+                                                 extra=(("cfg", type(self.cfg).__name__), ("edm", True))),
+                        resources=res,
+                        payload={
+                            "branch": "edm",
+                            "step": i
+                        },
+                        run=run,
+                        label=f"cosmos.denoise.{i}",
+                        capturable=False)  # EDM x0-space CFG + (gated) host-RNG frame injection -> eager path
+
+    def advance(self, st: LoopState, result: StepResult) -> LoopState:
+        st.latents["video"] = result.output["latents"]
+        if st.profile == ExecutionProfile.ROLLOUT:
+            st.trajectory.append({
+                "step": st.step_idx,
+                "sigma": st.sigmas[st.step_idx],
+                "velocity": np.asarray(result.output["noise_pred"]).copy(),
+                "latents": np.asarray(st.latents["video"]).copy()
+            })
+        st.step_idx += 1
+        return st
+
+    def finalize(self, st: LoopState) -> LoopResult:
+        return LoopResult(outputs={"latents": st.latents["video"]},
+                          metrics={"denoise_steps": float(st.step_idx)},
+                          behavior=st.trajectory or None)

v2/recipes/cosmos2/program.py

@@ -0,0 +1,35 @@
+"""Cosmos-Predict2 t2v program: text_encode → diffusion_denoise (EDM) → vae_decode.
+
+Same inline shape as the Wan t2v program — the EDM specifics live entirely in ``CosmosDenoiseLoop`` and
+the ``CosmosDiT`` adapter, so the node graph is unchanged. (video2world conditioning would add a
+VAE-encode node writing ``conditioning_latents``/``cond_indicator`` into slots; the loop already reads
+them.)
+"""
+from __future__ import annotations
+
+from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind
+from v2.recipes.common import text_encode_node_fn as _text_encode
+
+
+def _vae_decode(instance, slots, request, ctx) -> None:
+    slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"])
+
+
+def build_cosmos2_program() -> Program:
+    return Program(
+        program_id="cosmos2.t2v.inline",
+        kind=ProgramKind.INLINE,
+        nodes=[
+            ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")),
+            ModelLoopNode("denoise",
+                          loop_id="diffusion_denoise",
+                          output_slot="denoise_out",
+                          reads=("text_embeds", "neg_text_embeds"),
+                          writes=("denoise_out", )),
+            ComponentNode("vae_decode", fn=_vae_decode, reads=("denoise_out", ), writes=("video", )),
+        ],
+        output_artifacts={
+            "video": "video",
+            "latents": "denoise_out"
+        },
+    ).validate()

v2/registry.py

@@ -21,7 +21,8 @@ from __future__ import annotations
 import json
 import os
 from dataclasses import dataclass
-from typing import Any, Callable
+from typing import Any
+from collections.abc import Callable
 
 
 @dataclass(frozen=True)
@@ -52,31 +53,33 @@ def _entries() -> list[ModelEntry]:
         build_wan_t2v_program,
     )
     from v2.recipes.wan21.i2v import build_wan21_i2v_card, build_wan21_i2v_program, build_wan22_i2v_a14b_card
+    from v2.recipes.cosmos2 import build_cosmos2_card, build_cosmos2_program
     return [
-        ModelEntry(("Wan-AI/Wan2.1-T2V-1.3B-Diffusers",), build_wan21_card, build_wan_t2v_program),
-        ModelEntry(("Wan-AI/Wan2.1-T2V-14B-Diffusers",), build_wan_t2v_14b_card, build_wan_t2v_program),
+        ModelEntry(("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", ), build_wan21_card, build_wan_t2v_program),
+        ModelEntry(("Wan-AI/Wan2.1-T2V-14B-Diffusers", ), build_wan_t2v_14b_card, build_wan_t2v_program),
         # Wan2.1 i2v cluster: CLIP image encoder + first-frame VAE conditioning ([mask|cond] -> 36ch DiT).
         # Fun-1.3B-InP is GPU-verified; the 14B variants reuse the same i2v card/path (weights GPU-pending).
-        ModelEntry(("weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers",),
-                   build_wan21_i2v_card, build_wan21_i2v_program),
-        ModelEntry(("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",),
+        ModelEntry(("weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", ), build_wan21_i2v_card, build_wan21_i2v_program),
+        ModelEntry(("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", ),
                    lambda: build_wan21_i2v_card("wan2.1-i2v-14b-480p", flow_shift=3.0, height=480, width=832),
                    build_wan21_i2v_program),
-        ModelEntry(("Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",),
+        ModelEntry(("Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", ),
                    lambda: build_wan21_i2v_card("wan2.1-i2v-14b-720p", flow_shift=5.0, height=720, width=1280),
                    build_wan21_i2v_program),
         # Wan2.2-I2V-A14B: MoE (2 experts + boundary) + i2v conditioning. Reuses Wan adapter + i2v + MoE.
-        ModelEntry(("Wan-AI/Wan2.2-I2V-A14B-Diffusers",),
-                   build_wan22_i2v_a14b_card, build_wan21_i2v_program),
-        ModelEntry(("Wan-AI/Wan2.2-TI2V-5B-Diffusers",), build_wan22_ti2v_card, build_wan_t2v_program),
-        ModelEntry(("Wan-AI/Wan2.2-T2V-A14B-Diffusers",), build_wan22_a14b_card, build_wan_t2v_program),
-        ModelEntry(("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",), build_wan_causal_card, build_wan_causal_program),
+        ModelEntry(("Wan-AI/Wan2.2-I2V-A14B-Diffusers", ), build_wan22_i2v_a14b_card, build_wan21_i2v_program),
+        ModelEntry(("Wan-AI/Wan2.2-TI2V-5B-Diffusers", ), build_wan22_ti2v_card, build_wan_t2v_program),
+        ModelEntry(("Wan-AI/Wan2.2-T2V-A14B-Diffusers", ), build_wan22_a14b_card, build_wan_t2v_program),
+        ModelEntry(("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", ), build_wan_causal_card, build_wan_causal_program),
         # LTX-2 two-stage distilled (ships a spatial_upsampler: base -> upsample -> refine):
-        ModelEntry(("FastVideo/LTX2-Distilled-Diffusers",), build_ltx2_card, build_ltx2_program),
+        ModelEntry(("FastVideo/LTX2-Distilled-Diffusers", ), build_ltx2_card, build_ltx2_program),
         # LTX-2 single-stage, non-distilled base (request-driven many-step, video-only):
-        ModelEntry(("Davids048/LTX2-Base-Diffusers",), build_ltx2_base_card, build_ltx2_base_program),
+        ModelEntry(("Davids048/LTX2-Base-Diffusers", ), build_ltx2_base_card, build_ltx2_base_program),
         # LTX-2.3 distilled — single-stage, JOINT text->video+audio (separate connectors + gated attn):
-        ModelEntry(("FastVideo/LTX-2.3-Distilled-Diffusers",), build_ltx2_3_card, build_ltx2_3_program),
+        ModelEntry(("FastVideo/LTX-2.3-Distilled-Diffusers", ), build_ltx2_3_card, build_ltx2_3_program),
+        # Cosmos-Predict2-2B-Video2World — EDM-Karras denoiser (new CosmosDenoiseLoop + CosmosDiT adapter),
+        # reusing the Wan VAE adapter + T5. Registered for t2v (video2world conditioning threads later).
+        ModelEntry(("nvidia/Cosmos-Predict2-2B-Video2World", ), build_cosmos2_card, build_cosmos2_program),
     ]
 
 
@@ -89,6 +92,7 @@ def read_arch_signature(root: str) -> dict:
     keys): pipeline / transformer / VAE class names + the fields that split same-class variants
     (VAE ``z_dim``, a second transformer / ``boundary_ratio`` for MoE, a ``spatial_upsampler``).
     ``root`` is a local diffusers dir or a snapshot of just the ``*.json`` configs."""
+
     def _load(*parts: str) -> dict:
         p = os.path.join(root, *parts)
         if os.path.exists(p):
@@ -114,6 +118,9 @@ def select_by_architecture(sig: dict):
     """FALLBACK: map an architecture signature -> (build_card, build_program) by class names, so a local
     path / renamed repo / new distilled variant of a known arch still resolves with no registry entry."""
     tr, pipe = sig.get("transformer_cls"), sig.get("pipeline")
+    if tr == "CosmosTransformer3DModel":
+        from v2.recipes.cosmos2 import build_cosmos2_card, build_cosmos2_program
+        return build_cosmos2_card, build_cosmos2_program
     if tr == "LTX2Transformer3DModel":
         from v2.recipes.ltx2 import (
             build_ltx2_base_card,
@@ -122,8 +129,8 @@ def select_by_architecture(sig: dict):
             build_ltx2_program,
         )
         if sig.get("has_spatial_upsampler"):
-            return build_ltx2_card, build_ltx2_program          # two-stage (base -> upsample -> refine)
-        return build_ltx2_base_card, build_ltx2_base_program    # single-stage
+            return build_ltx2_card, build_ltx2_program  # two-stage (base -> upsample -> refine)
+        return build_ltx2_base_card, build_ltx2_base_program  # single-stage
     if tr == "CausalWanTransformer3DModel":
         from v2.recipes.wan_causal import build_wan_causal_card, build_wan_causal_program
         return build_wan_causal_card, build_wan_causal_program
@@ -134,20 +141,18 @@ def select_by_architecture(sig: dict):
             build_wan22_ti2v_card,
             build_wan_t2v_program,
         )
-        if pipe == "WanDMDPipeline":   # FastWan: detected by pipeline class -> precise, not a load crash
-            raise ValueError(
-                "v2 registry: WanDMD/FastWan is not supported via the generic Wan path — its checkpoint's "
-                "to_gate_compress param mapping differs from the generic WanTransformer3DModel load. See "
-                "examples/inference/basic/V2_PORTING_STATUS.md.")
+        if pipe == "WanDMDPipeline":  # FastWan: detected by pipeline class -> precise, not a load crash
+            raise ValueError("v2 registry: WanDMD/FastWan is not supported via the generic Wan path — its checkpoint's "
+                             "to_gate_compress param mapping differs from the generic WanTransformer3DModel load. See "
+                             "examples/inference/basic/V2_PORTING_STATUS.md.")
         if sig.get("has_transformer_2") or sig.get("boundary_ratio"):
-            return build_wan22_a14b_card, build_wan_t2v_program     # Wan2.2 MoE (two experts)
+            return build_wan22_a14b_card, build_wan_t2v_program  # Wan2.2 MoE (two experts)
         if sig.get("vae_z_dim") == 48:
-            return build_wan22_ti2v_card, build_wan_t2v_program     # Wan2.2-TI2V-5B (z_dim=48 VAE)
-        return build_wan21_card, build_wan_t2v_program              # Wan2.1
-    raise ValueError(
-        f"v2 registry: unsupported architecture (transformer={tr!r}, pipeline={pipe!r}). Supported "
-        f"transformers: WanTransformer3DModel / CausalWanTransformer3DModel / LTX2Transformer3DModel. "
-        f"See examples/inference/basic/V2_PORTING_STATUS.md.")
+            return build_wan22_ti2v_card, build_wan_t2v_program  # Wan2.2-TI2V-5B (z_dim=48 VAE)
+        return build_wan21_card, build_wan_t2v_program  # Wan2.1
+    raise ValueError(f"v2 registry: unsupported architecture (transformer={tr!r}, pipeline={pipe!r}). Supported "
+                     f"transformers: WanTransformer3DModel / CausalWanTransformer3DModel / LTX2Transformer3DModel. "
+                     f"See examples/inference/basic/V2_PORTING_STATUS.md.")
 
 
 def resolve(model_path: str, root: str | None = None):
@@ -158,15 +163,14 @@ def resolve(model_path: str, root: str | None = None):
     ``*.json`` config snapshot). Raises ``ValueError`` if unregistered with ``root=None`` (so the caller
     knows to fetch the configs), or if the architecture is unsupported. Shared by every v2 entrypoint."""
     entries = _entries()
-    for e in entries:                                    # 1. exact HF id
+    for e in entries:  # 1. exact HF id
         if model_path in e.hf_ids:
             return e.build_card, e.build_program
-    short = _short(model_path)                            # 2. short repo name (renamed / forked copy)
+    short = _short(model_path)  # 2. short repo name (renamed / forked copy)
     for e in entries:
         if any(_short(h) == short for h in e.hf_ids):
             return e.build_card, e.build_program
-    if root is None:                                      # 3. architecture inference (needs the configs)
-        raise ValueError(
-            f"v2 registry: {model_path!r} is not registered; pass the checkpoint root for architecture "
-            f"inference, or add a ModelEntry in v2/registry.py.")
+    if root is None:  # 3. architecture inference (needs the configs)
+        raise ValueError(f"v2 registry: {model_path!r} is not registered; pass the checkpoint root for architecture "
+                         f"inference, or add a ModelEntry in v2/registry.py.")
     return select_by_architecture(read_arch_signature(root))