From bfbd90ea3ac39f987f3658ed761374bacff6e97f Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Thu, 18 Jun 2026 15:38:46 +0000 Subject: [PATCH] =?UTF-8?q?[feat]=20v2:=20Cosmos-Predict2-2B-Video2World?= =?UTF-8?q?=20port=20(EDM-Karras)=20=E2=80=94=20reference=20bucket-C=20rec?= =?UTF-8?q?ipe?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 46 +++++-- v2/platform/backends/torch_cosmos.py | 61 +++++++++ v2/recipes/_prompts.py | 9 ++ v2/recipes/cosmos2/__init__.py | 14 ++ v2/recipes/cosmos2/card.py | 122 +++++++++++++++++ v2/recipes/cosmos2/loop.py | 198 +++++++++++++++++++++++++++ v2/recipes/cosmos2/program.py | 35 +++++ v2/registry.py | 74 +++++----- 8 files changed, 513 insertions(+), 46 deletions(-) create mode 100644 v2/platform/backends/torch_cosmos.py create mode 100644 v2/recipes/cosmos2/__init__.py create mode 100644 v2/recipes/cosmos2/card.py create mode 100644 v2/recipes/cosmos2/loop.py create mode 100644 v2/recipes/cosmos2/program.py diff --git a/v2/loop/sampler.py b/v2/loop/sampler.py index 4cd39ee1c..d1b1b5d17 100644 --- a/v2/loop/sampler.py +++ b/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") diff --git a/v2/platform/backends/torch_cosmos.py b/v2/platform/backends/torch_cosmos.py new file mode 100644 index 000000000..5a92d8c3a --- /dev/null +++ b/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) diff --git a/v2/recipes/_prompts.py b/v2/recipes/_prompts.py index 373cca1f6..ec1b360ad 100644 --- a/v2/recipes/_prompts.py +++ b/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, " diff --git a/v2/recipes/cosmos2/__init__.py b/v2/recipes/cosmos2/__init__.py new file mode 100644 index 000000000..846303504 --- /dev/null +++ b/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"] diff --git a/v2/recipes/cosmos2/card.py b/v2/recipes/cosmos2/card.py new file mode 100644 index 000000000..1b4dc6ade --- /dev/null +++ b/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 diff --git a/v2/recipes/cosmos2/loop.py b/v2/recipes/cosmos2/loop.py new file mode 100644 index 000000000..1f9551d76 --- /dev/null +++ b/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) diff --git a/v2/recipes/cosmos2/program.py b/v2/recipes/cosmos2/program.py new file mode 100644 index 000000000..139390efa --- /dev/null +++ b/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() diff --git a/v2/registry.py b/v2/registry.py index 8e230b055..6d20d990d 100644 --- a/v2/registry.py +++ b/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))