diff --git a/v2/models/dits/lingbotworld.py b/v2/models/dits/lingbotworld.py new file mode 100644 index 000000000..46385745b --- /dev/null +++ b/v2/models/dits/lingbotworld.py @@ -0,0 +1,11 @@ +# STUB: re-exports fastvideo until vendored (see memory: v2-vendoring-approach). +"""LingBot-World DiT facade. + +The transformer class itself is constructed by the loader from the card's ``load_id`` +(``fastvideo.models.dits.lingbotworld.model:LingBotWorldTransformer3DModel``), so it is not imported +here. What the v2 recipe DOES need from the fastvideo source is the camera/Plucker embedding builder +``prepare_camera_embedding`` (poses.npy + intrinsics.npy -> ``c2ws_plucker_emb [B, 6*s^2, F, H, W]``), +re-exported so the program's camera node can build the tensor without reaching into the model package. +""" +from fastvideo.models.dits.lingbotworld.cam_utils import ( # noqa: F401 + prepare_camera_embedding, ) diff --git a/v2/models/dits/sd3.py b/v2/models/dits/sd3.py new file mode 100644 index 000000000..3e777a6b7 --- /dev/null +++ b/v2/models/dits/sd3.py @@ -0,0 +1,5 @@ +# STUB: re-exports fastvideo until vendored (see memory: v2-vendoring-approach). +"""SD3 (Stable Diffusion 3.5) DiT facade. The transformer class itself is constructed by the loader +from the card's ``load_id``; this stub only re-exports the output dataclass so the SD3 torch adapter can +type-check / unwrap the forward result symbolically (mirrors ``v2/models/dits/ltx2.py``).""" +from fastvideo.models.dits.sd3 import SD3Transformer2DModelOutput # noqa: F401 diff --git a/v2/platform/backends/torch_cosmos25.py b/v2/platform/backends/torch_cosmos25.py new file mode 100644 index 000000000..de87563f2 --- /dev/null +++ b/v2/platform/backends/torch_cosmos25.py @@ -0,0 +1,126 @@ +"""Cosmos-Predict2.5 torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` +so the Cosmos2.5 recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/ +``_make_text_encoder`` dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only +on a GPU box. + +The three Cosmos2.5 deltas vs Wan (all GPU-path, all carried HERE so the loop stays toy-compatible): + +* ``Cosmos25DiT`` — a flow-match velocity DiT, but the per-step call differs from ``WanDiT`` in three + load-bearing ways (see ``fastvideo/pipelines/stages/denoising.py:Cosmos25DenoisingStage``): + 1. timestep is a PER-FRAME 2D tensor ``[B, T]`` filled with the PLAIN sigma (in [0,1]) — NOT Wan's + scalar ``sigma*1000``. (The fastvideo stage feeds ``t*0.001`` of the 0..1000 schedule, i.e. the + shifted sigma directly; ``Cosmos25DenoiseLoop`` already hands us that plain sigma.) + 2. the forward REQUIRES a ``condition_mask[B,1,T,H,W]`` (zeros for t2v) + a ``padding_mask[B,1,H,W]`` + (ones for t2v) every step — the model concats them internally (16 -> 17 -> 18ch patch_embed), + so we feed the RAW 16-channel latent (the OPPOSITE of Wan i2v, where the adapter pre-concats). + 3. an ``fps`` scalar is required for RoPE temporal scaling (NTK fps-modulation). + The model output is the rectified-flow velocity (a bare tensor); the loop integrates it with the + shared ``FLOW_MATCH_STEP`` Euler solver and combines branches with the standard ``ClassicCFG``. +* ``Cosmos25WanVAE`` — a Wan-style causal 3D VAE (z=16, 8x/4x), BUT unlike ``AutoencoderKLWan`` its + ``encode``/``decode`` already apply the (z-mean)/std normalization INTERNALLY (the Cosmos latent + contract: encode -> normalized, decode <- normalized). So this is a thin marshalling wrapper (like + ``LTX2VAE``) that must NOT re-normalize — reusing the v2 ``WanVAE`` adapter would double-normalize. +* ``Cosmos25Reason1Encoder`` — Qwen2.5-VL multimodal LM via ``compute_text_embeddings_online``: per-layer + mean-normalized hidden states concatenated along the feature dim -> a 100352-dim (= num_hidden_layers * + hidden_size = 28*3584) sequence the DiT's internal 100352->1024 GELU crossattn projection consumes. + Faithful to ``Cosmos25TextEncodingStage`` (calls ``compute_text_embeddings_online`` rather than + re-implementing the normalization). BRINGUP: heavy (a 7B LM + 100352-dim embeds); GPU-only. +""" +from __future__ import annotations + +import torch + +from v2.platform.backends.torch_backend import TorchComponent, _to_numpy + +# Cosmos2.5 t2v feeds the DiT the PLAIN shifted sigma in [0,1] as the per-frame timestep — NOT +# ``sigma * num_train_timesteps`` (the Wan convention would corrupt the output). See blocker (1). +COSMOS25_DEFAULT_FPS = 24 # RoPE temporal scaling; the registered t2v preset's default fps. + + +class Cosmos25DiT(TorchComponent): + """``dit(latent[C,T,h,w], text_embed[seq,100352], sigma) -> velocity[C,T,h,w]``. + + The loop hands the RAW 16-channel latent and the PLAIN sigma (in [0,1]); this adapter builds the + mandatory per-frame timestep + zero ``condition_mask`` + ones ``padding_mask`` + ``fps`` the Cosmos2.5 + forward requires (the model concats the masks itself -> 18ch patch_embed input). Returns the bare + velocity tensor; ``Cosmos25DenoiseLoop`` does the ClassicCFG combine + the flow-match Euler step.""" + + def __init__(self, module, *, device, dtype, fps: int = COSMOS25_DEFAULT_FPS): + super().__init__(module, device=device, dtype=dtype) + # BRINGUP: request-driven fps would thread through the request API; the registered t2v preset + # uses fps=24 (the fastvideo stage likewise defaults to a fixed fps for t2v). Held on the adapter + # so the loop's dit-call stays ``dit(latent, text_embed, sigma)`` (toy-compatible — no fps kwarg). + self.fps = int(fps) + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + hs = self._t(latent) # [1, 16, T, h, w] — the RAW latent (no concat) + ehs = self._t(text_embed) # [1, seq, 100352] + _b, _c, t, h, w = hs.shape + s = float(sigma) + # (1) per-frame 2D timestep [B, T] of the PLAIN sigma (NOT sigma*1000). For t2v every frame shares + # the same sigma; the conditioned (V2W) path would override selected frames (BRINGUP, see loop). + timestep = torch.full((1, t), s, device=self.device, dtype=torch.float32) + # (2) mandatory masks: zeros condition_mask (no frame conditioned in t2v) + ones padding_mask + # (no padding in t2v). The model resizes padding_mask to [h, w] and concats both internally. + condition_mask = torch.zeros(1, 1, t, h, w, device=self.device, dtype=self.dtype) + padding_mask = torch.ones(1, 1, h, w, device=self.device, dtype=self.dtype) + # (3) fps scalar for RoPE temporal NTK scaling. + fps_tensor = torch.tensor([float(self.fps)], device=self.device, dtype=self.dtype) + with self._ctx(): + velocity = self.module(hidden_states=hs, + timestep=timestep, + encoder_hidden_states=ehs, + fps=fps_tensor, + condition_mask=condition_mask, + padding_mask=padding_mask, + return_dict=False)[0] + return self._n(velocity) # rectified-flow velocity (loop integrates it) + + +class Cosmos25WanVAE(TorchComponent): + """Cosmos2.5 Wan-style causal 3D VAE (``fastvideo.models.vaes.cosmos25wanvae:Cosmos25WanVAE``). + + Its ``encode``/``decode`` ALREADY apply the (z-mean)/std normalization internally (the Cosmos latent + contract), so this adapter only marshals numpy<->torch — it must NOT re-normalize (that is why we do + not reuse the v2 ``WanVAE`` adapter, which would double-apply the per-channel stats).""" + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent) # [1, 16, T, h, w] — a NORMALIZED latent + video = self.module.decode(z.to(self.dtype)) # -> video [B, 3, T, H, W] in [-1, 1] + if hasattr(video, "sample"): + video = video.sample + return self._n(video) + + @torch.no_grad() + def encode(self, video): + dist = self.module.encode(self._t(video)) + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + return self._n(z) # already the NORMALIZED latent the DiT expects + + +class Cosmos25Reason1Encoder(TorchComponent): + """Reason1 (Qwen2.5-VL) text encoder. ``encode(text) -> embedding[seq, 100352]`` (numpy out). + + Faithful to ``Cosmos25TextEncodingStage``: call ``compute_text_embeddings_online({"text": [text]}, + "text")`` (the encoder applies the chat template + per-layer mean-normalize + full-concat itself, so + we do NOT re-implement the 100352-dim normalization here). The negative prompt is encoded the same + way for ClassicCFG. The encoder is multimodal/heavy — BRINGUP, GPU-only. + + The ``tokenizer`` extra from ``_make_text_encoder`` is accepted for signature compatibility but + unused: ``Reason1TextEncoder`` owns its own processor/tokenizer (``self.processor.tokenizer``).""" + + def __init__(self, module, tokenizer=None, *, device, dtype): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer # unused; Reason1 carries its own processor (BRINGUP note above) + + @torch.no_grad() + def encode(self, text): + with self._ctx(): + # returns [B, seq, 100352]; single-request B=1 -> squeeze to [seq, 100352] (the DiT adapter's + # _t re-adds the batch dim, matching the fastvideo stage's batch.prompt_embeds[0]). + embeds = self.module.compute_text_embeddings_online({"text": [text or ""]}, "text") + if hasattr(embeds, "shape") and len(embeds.shape) == 3: + embeds = embeds[0] + return _to_numpy(embeds) diff --git a/v2/platform/backends/torch_flux2.py b/v2/platform/backends/torch_flux2.py new file mode 100644 index 000000000..3c6efe072 --- /dev/null +++ b/v2/platform/backends/torch_flux2.py @@ -0,0 +1,198 @@ +"""FLUX.2 torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so the FLUX.2 +recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder`` dispatch +in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +These three pieces CANNOT reuse the Wan/T5 adapters (BRINGUP — GPU-verify each against a real FLUX.2 ckpt): + +* ``Flux2DiT`` — the dual-stream MMDiT. The loop hands ``dit(latent[64,T,h,w], text_embed[seq,4096], σ)``; + this adapter builds everything FLUX.2-specific INTERNALLY so the loop stays toy-compatible: + - σ is passed DIRECTLY as the timestep (the pipeline's ``timestep/1000``); the DiT re-multiplies by + 1000. Guidance is passed RAW (``embedded_guidance`` scale, ×1000 internally). Passing σ·1000 (the + Wan convention) would 1e6-blow-up the timestep embedding. + - ``encoder_hidden_states`` is wrapped in a 1-element list (the forward indexes [0]). + - ``img_ids`` = ``cartesian_prod(arange(T), arange(h), arange(w), arange(1))`` and ``txt_ids`` = + ``cartesian_prod(1,1,1,arange(seq))`` — the N-D RoPE position tensors threaded into every forward. + - run WITHOUT autocast (Gate 2: FLUX.2 long-seq attention breaks under autocast); module + inputs bf16. + Returns the bare velocity. Faithful to ``fastvideo/pipelines/stages/denoising.py`` (the ``_is_flux`` + gates) + ``flux_2.py`` forward + ``flux_2_latent_preparation.py`` / ``flux_2_text_encoding.py`` id builders. + +* ``Flux2VAE`` — decode does the FLUX.2 BatchNorm whitening + 2×2 unpatchify (NOT Wan mean/std): + ``z·bn_std + bn_mean`` (the complete inverse normalization; no scaling_factor step) then unpack 64→16ch + full-spatial, then ``module.decode``. Faithful to ``DecodingStage._flux2_bn_denorm_and_unpatchify`` + + ``_unpatchify_latents`` (fastvideo/pipelines/stages/decoding.py). Image-only: T==1 squeezed. + +* ``Flux2Mistral3Encoder`` / ``Flux2Qwen3Encoder`` — the chat-template multi-layer hidden-state stack. + Apply the FLUX2 system+user chat template (dev only), tokenize padded to max_length (512), + ``output_hidden_states=True``, stack the chosen decoder layers (dev 10/20/30, klein 9/18/27), + permute+reshape to [seq, n_layers·hidden] (= 4096). Faithful to ``Flux2TextEncodingStage`` + + ``flux2_klein_postprocess_text`` (fastvideo/configs/pipelines/flux_2.py). +""" +from __future__ import annotations + +import torch + +from v2.platform.backends.torch_backend import TorchComponent, _to_numpy + +# FLUX.2 system message prepended to the user prompt via the chat template (dev). Faithful copy of +# ``fastvideo/pipelines/basic/flux_2/flux_2_text_encoding.py:FLUX2_SYSTEM_MESSAGE``. +FLUX2_SYSTEM_MESSAGE = ("You are an AI that reasons about image descriptions. You give structured " + "responses focusing on object relationships, object\nattribution and actions " + "without speculation.") + +NUM_TRAIN_TIMESTEPS = 1000 # the DiT re-multiplies timestep AND guidance by this internally + + +class Flux2DiT(TorchComponent): + """``dit(latent[64,T,h,w], text_embed[seq,4096], sigma) -> velocity[64,T,h,w]``. + + Real forward (flux_2.py): ``forward(hidden_states[B,64,T,h,w], encoder_hidden_states=[ehs], timestep, + img_ids, txt_ids, guidance) -> velocity (bare 5D tensor)``. The pipeline passes ``timestep == sigma`` + (DiT ×1000 internally) and ``guidance == embedded_cfg_scale`` raw (DiT ×1000 internally).""" + + def __init__(self, module, *, device, dtype): + super().__init__(module, device=device, dtype=dtype) + # FLUX.2-dev uses guidance_embeds=True (a guidance arg); klein distilled has no guidance embedder. + self.guidance_embeds = bool(getattr(module, "guidance_embeds", True)) + + def _ids(self, t: int, h: int, w: int, txt_len: int): + # img_ids over (T, h, w, 1); txt_ids over (1,1,1,seq) — the N-D RoPE position tensors. Faithful to + # flux_2_latent_preparation.py (latent_ids) + flux_2_text_encoding.py (_prepare_flux2_text_ids). + img_ids = torch.cartesian_prod(torch.arange(t, device=self.device), torch.arange(h, device=self.device), + torch.arange(w, device=self.device), + torch.arange(1, device=self.device)).unsqueeze(0) + txt_ids = torch.cartesian_prod(torch.arange(1, device=self.device), torch.arange(1, device=self.device), + torch.arange(1, device=self.device), + torch.arange(txt_len, device=self.device)).unsqueeze(0) + return img_ids, txt_ids + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + hs = self._t(latent) # [1, 64, T, h, w] + ehs = self._t(text_embed) # [1, seq, 4096] + b, _c, t, h, w = hs.shape + s = float(sigma) + # timestep == sigma directly (NOT sigma*1000) — the DiT multiplies by 1000 internally. + timestep = torch.tensor([s] * b, device=self.device, dtype=self.dtype) + guidance = None + if self.guidance_embeds: # dev: embedded guidance, raw scale + guidance = torch.tensor([float(self._guidance(context))] * b, device=self.device, dtype=self.dtype) + img_ids, txt_ids = self._ids(t, h, w, ehs.shape[1]) + with self._ctx(current_timestep=s): # NO autocast (Gate 2) — module+inputs bf16 + velocity = self.module( + hidden_states=hs, + encoder_hidden_states=[ehs], # forward indexes [0] + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + guidance=guidance) + return self._n(velocity) # rectified-flow velocity + + @staticmethod + def _guidance(context) -> float: + # The loop threads the embedded-guidance scale through the ``context`` slot of the dit-call (the + # toy-compatible signature). BRINGUP: the request-API extension that carries the scale per-request + # is pending; until then the card's sampling default (dev 4.0) is the value, set on the adapter by + # the loop wiring. Defaults to 4.0 if unset. + return 4.0 if context is None else float(context) + + +class Flux2VAE(TorchComponent): + """FLUX.2 image VAE. decode: BN-denorm (``z·bn_std + bn_mean``) → 2×2 unpatchify (64→16ch full-spatial) + → ``module.decode`` → image[-1,1]. encode: ``module.encode`` (mode) → packed latent (BRINGUP: pack + + BN-whiten on encode is GPU-only; t2i never encodes). NOT the Wan (z−mean)/std normalization.""" + + def _bn_stats(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + bn = self.module.bn + eps = float(getattr(bn, "eps", 1e-5)) + running_mean = bn.running_mean.view(1, -1, 1, 1).to(like.device, like.dtype) + running_var = bn.running_var.view(1, -1, 1, 1).to(like.device, like.dtype) + bn_std = torch.sqrt(torch.clamp(running_var + eps, min=1e-6)) + return running_mean, bn_std + + @staticmethod + def _unpatchify(latents: torch.Tensor) -> torch.Tensor: + # (B, C*4, h, w) -> (B, C, 2h, 2w): inverse of the 2×2 patch packing (DecodingStage._unpatchify_latents). + b, c, h, w = latents.shape + latents = latents.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3) + return latents.reshape(b, c // 4, h * 2, w * 2) + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent).float() # [1, 64, T, h, w] + if z.ndim == 5: # image VAE expects 4D; T==1 squeezed + z = z.squeeze(2) + running_mean, bn_std = self._bn_stats(z) + z = z * bn_std + running_mean # BN-denorm: complete inverse (no scaling_factor) + z = self._unpatchify(z) # 64 -> 16ch full-spatial + image = self.module.decode(z.to(self.dtype)) + if hasattr(image, "sample"): + image = image.sample + elif isinstance(image, tuple | list): + image = image[0] + image = (image / 2 + 0.5).clamp(0, 1) # -> [0,1] + return self._n(image.unsqueeze(2)) # restore the singleton T for a uniform [C,T,H,W] + + @torch.no_grad() + def encode(self, image): + # BRINGUP: t2i never encodes; the 2×2 pack + BN whiten on the encode path is GPU-only and unverified. + x = self._t(image) + dist = self.module.encode(x) + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + return _to_numpy(z) + + +class _Flux2TextEncoderBase(TorchComponent): + """Shared FLUX.2 chat-template multi-layer hidden-state stack. Subclasses set ``layers`` and whether + the FLUX2 system-message chat template is applied (dev=Mistral3) or a plain user template (klein=Qwen3).""" + + layers: tuple[int, ...] = (10, 20, 30) + use_system_message: bool = True + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = 512): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + def _messages(self, prompt: str): + user = {"role": "user", "content": [{"type": "text", "text": (prompt or "").replace("[IMG]", "")}]} + if self.use_system_message: + system = {"role": "system", "content": [{"type": "text", "text": FLUX2_SYSTEM_MESSAGE}]} + return [[system, user]] + return [[user]] + + @torch.no_grad() + def encode(self, text): + inputs = self.tokenizer.apply_chat_template(self._messages(text), + add_generation_prompt=False, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.max_length) + input_ids = inputs["input_ids"].to(self.device) + attention_mask = inputs["attention_mask"].to(self.device) + with self._ctx(): + out = self.module(input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False) + if out.hidden_states is None: + raise ValueError("FLUX.2 requires output_hidden_states=True from the text encoder") + # stack the chosen decoder layers -> [B, n_layers, seq, hidden] -> permute/reshape [B, seq, n·hidden] + stacked = torch.stack([out.hidden_states[k] for k in self.layers], dim=1) + b, n_layers, seq_len, hidden = stacked.shape + embed = stacked.permute(0, 2, 1, 3).reshape(b, seq_len, n_layers * hidden) + return _to_numpy(embed.squeeze(0)) # [seq, n_layers*hidden] (= 4096) + + +class Flux2Mistral3Encoder(_Flux2TextEncoderBase): + """FLUX.2-dev text encoder: Mistral3, hidden-state layers 10/20/30, FLUX2 system-message chat template.""" + layers = (10, 20, 30) + use_system_message = True + + +class Flux2Qwen3Encoder(_Flux2TextEncoderBase): + """FLUX.2-klein text encoder: Qwen3, hidden-state layers 9/18/27, plain user chat template (no system msg).""" + layers = (9, 18, 27) + use_system_message = False diff --git a/v2/platform/backends/torch_gen3c.py b/v2/platform/backends/torch_gen3c.py new file mode 100644 index 000000000..2f385deae --- /dev/null +++ b/v2/platform/backends/torch_gen3c.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GEN3C torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so the GEN3C +recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder`` +dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +GEN3C is Cosmos-Predict2's EDM denoiser **extended with camera/3D-cache conditioning**. The deltas vs +the ``CosmosDiT`` adapter (see ``torch_cosmos.py``) are concentrated here: + +* ``Gen3CDiT`` — the EDM denoiser ``F_θ``. The loop hands the *already EDM-input-scaled* model input + (``x·c_in``) plus the **raw** conditioning bundle as kwargs; the DiT internally concats + ``[latent(16) | input_mask(1) | pose_buffer(frame_buffer_max·32) | padding_mask(1)] = 82`` along the + channel dim before ``patch_embed``. The conditioning is NOT pre-concatenated into the latent (the Wan + i2v 16→36 concat would be wrong here): the loop passes ``condition_video_pose`` / + ``condition_video_input_mask`` / ``condition_video_augment_sigma`` / ``padding_mask`` separately and + the model assembles the 82-ch input. The model timestep is the EDM **noise preconditioning** + ``c_noise = 0.25·log(σ)`` (NOT ``σ·1000`` like Wan, NOT raw σ) — this is exactly the value the + diffusers ``EDMEulerScheduler.precondition_noise`` puts in ``scheduler.timesteps``, which the + fastvideo ``Gen3CDenoisingStage`` feeds to the transformer. Returns the **raw** transformer output + (the EDM ``c_skip``/``c_out`` → x0 reconstruction + x0-space CFG + frame-replace live in + ``Gen3CDenoiseLoop``, NOT here). Faithful to ``Gen3CDenoisingStage.forward`` + (fastvideo/pipelines/stages/gen3c_stages.py). +* ``Gen3CVAE`` — wraps ``AutoencoderKLGen3CTokenizer`` (a JIT-backed tokenizer, NOT a diffusers + ``AutoencoderKL``). Normalization is **internal** (the tokenizer applies its own mean/std buffers + inside ``encode``/``decode``), so the adapter does NOT subtract/divide an external Wan-style + ``(z-mean)/std`` (BRINGUP: contrast with ``WanVAE``). ``encode`` returns the tokenizer's ``.mode()``. +* ``Gen3CT5Encoder`` — T5-Large (1024-dim). GEN3C uses the **raw** last_hidden_state (NaN→0, no + fixed-length zero-pad), identical to Cosmos's ``CosmosT5Encoder`` (reused by subclassing it). + +BRINGUP (all GPU-only — written-not-run, verify on the box): + * The ``AutoencoderKLGen3CTokenizer`` JIT loader expects ``encoder.jit`` / ``decoder.jit`` / + ``mean_std.pt`` under the tokenizer dir (``from_jit_tokenizer``); the v2 loader seam must route the + ``vae`` component to it rather than ``VAELoader``'s diffusers path — confirm the converted HF repo + layout on the GPU box. + * The camera / MoGe-depth / 3D-cache-render conditioning that fills ``condition_video_pose`` (and + ``conditioning_latents`` for the first-frame anchor) is a separate pre-loop stage (see + ``v2/recipes/gen3c/program.py``); without it the loop runs the degenerate **t2v** path (all-zeros + conditioning → the pose-buffer concat is effectively zeroed inside the DiT). +""" +from __future__ import annotations + +import math + +import torch + +from v2.platform.backends.torch_backend import TorchComponent, _to_numpy +from v2.platform.backends.torch_cosmos import CosmosT5Encoder + + +def _c_noise(sigma: float) -> float: + """EDM noise preconditioning ``c_noise = 0.25·log(σ)`` — exactly ``EDMEulerScheduler.timesteps[i]`` + for ``σ = sigmas[i]``. The GEN3C DiT embeds this directly via its sinusoidal ``Timesteps`` module + (NO ``·1000``).""" + return 0.25 * math.log(max(float(sigma), 1e-12)) + + +class Gen3CDiT(TorchComponent): + """``dit(model_input[16,T,h,w], text_embed[seq,1024], sigma, *, cond=) -> raw output[16,T,h,w]``. + + ``model_input`` is pre-scaled by the loop (``x·c_in``); ``sigma`` is the *raw* EDM sigma (the adapter + converts it to the model timestep ``c_noise = 0.25·log(σ)`` internally). ``cond`` is the GEN3C + conditioning bundle dict (or None for t2v): ``condition_video_pose`` [buffer_channels,T,h,w], + ``condition_video_input_mask`` [1,T,h,w], ``condition_video_augment_sigma`` (scalar), and the pixel + ``height``/``width`` for the padding mask. None → the DiT builds zeros internally (degenerate t2v).""" + + def __init__(self, module, *, device, dtype, fps: int = 24): + super().__init__(module, device=device, dtype=dtype) + self.fps = int(fps) + # add_augment_sigma_embedding is FALSE in the shipped 7B arch config → no augment_sigma weights; + # the model only adds an augment-sigma embedding when that module exists (BRINGUP blocker 8). + self.has_augment_embed = bool(getattr(module, "add_augment_sigma_embedding", False)) + + @torch.no_grad() + def __call__(self, model_input, text_embed, sigma, context=None, *, cond=None): + hs = self._t(model_input) # [1, 16, T, h, w] + ehs = self._t(text_embed) # [1, seq, 1024] + b, _c, t, h, w = hs.shape + ts = torch.full((b, ), _c_noise(sigma), device=self.device, dtype=self.dtype) # c_noise = 0.25·log σ + + # Conditioning bundle: real tensors for the camera/3D-cache path, zeros (degenerate t2v) otherwise. + # The DiT concats input_mask + pose + padding_mask INTERNALLY → 82-ch patch_embed input. + input_mask = pose = augment_sigma = None + pix_h, pix_w = h, w # padding_mask is at PIXEL h,w then resized inside the DiT + if cond is not None: + input_mask = self._t(cond.get("condition_video_input_mask")) \ + if cond.get("condition_video_input_mask") is not None else None + pose = self._t(cond.get("condition_video_pose")) \ + if cond.get("condition_video_pose") is not None else None + aug = cond.get("condition_video_augment_sigma") + augment_sigma = torch.full((b, ), float(aug), device=self.device, dtype=self.dtype) \ + if (aug is not None and self.has_augment_embed) else None + pix_h = int(cond.get("height", h)) + pix_w = int(cond.get("width", w)) + padding_mask = torch.zeros(b, 1, pix_h, pix_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=self.fps, + condition_video_input_mask=input_mask, + condition_video_pose=pose, + condition_video_augment_sigma=augment_sigma, + padding_mask=padding_mask) + if isinstance(out, tuple): + out = out[0] + return self._n(out) # RAW EDM output (the loop reconstructs x0) + + +class Gen3CVAE(TorchComponent): + """``AutoencoderKLGen3CTokenizer`` adapter. The tokenizer is JIT-backed and handles its OWN latent + normalization internally (mean_std buffers) — so, unlike ``WanVAE``, this adapter does NOT apply an + external ``(z-mean)/std``. ``encode`` returns the tokenizer's distribution ``.mode()`` (the + ``_retrieve_latents`` contract the fastvideo stage uses); ``decode`` is a passthrough.""" + + @torch.no_grad() + def encode(self, video): + x = self._t(video) + dist = self.module.encode(x) + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + return self._n(z) # already in the DiT's (internally-normalized) latent space + + @torch.no_grad() + def decode(self, latent): + video = self.module.decode(self._t(latent)) # tokenizer un-normalizes internally → video [B,3,T,H,W] + if hasattr(video, "sample"): + video = video.sample + return self._n(video) + + +class Gen3CT5Encoder(CosmosT5Encoder): + """GEN3C T5-Large (1024-dim): raw last_hidden_state (NaN→0), no Wan zero-pad — identical to Cosmos. + Reuses ``CosmosT5Encoder.encode`` verbatim. (BRINGUP: GEN3C pads with ``max_length`` and zeros rows + past the attention-mask length, which ``CosmosT5Encoder`` already does — verify the max_length / + mask-zeroing matches the converted T5-Large tokenizer on the GPU box.)""" + + +# Re-export so a reader can ``from v2.platform.backends.torch_gen3c import _to_numpy`` if needed. +__all__ = ["Gen3CDiT", "Gen3CVAE", "Gen3CT5Encoder", "_to_numpy"] diff --git a/v2/platform/backends/torch_hunyuan_video.py b/v2/platform/backends/torch_hunyuan_video.py new file mode 100644 index 000000000..228837dbb --- /dev/null +++ b/v2/platform/backends/torch_hunyuan_video.py @@ -0,0 +1,176 @@ +"""HunyuanVideo torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so the +Hunyuan recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder`` +dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +Three deltas vs the Wan adapters: + +* ``HunyuanVideoDiT`` — flow-match velocity DiT, BUT ``encoder_hidden_states`` is a **2-element list** + ``[llama_hidden[B, L, 4096], clip_pooled[B, 768]]`` (the forward splits a list as + ``txt = ehs[0]``, ``text_states_2 = ehs[1]``), NOT a single tensor like Wan. The loop hands the LLaMA + per-token sequence as ``text_embed`` and the CLIP-pooled global vector via ``context=`` (Wan's existing + conditioning channel — see ``recipes/hunyuan_video/loop.py``); this adapter assembles the pair + internally. ``timestep = sigma*1000`` (scalar per-batch ``LongTensor``); the output is the **bare** + velocity tensor (NO ``x0 -> velocity`` conversion, unlike LTX-2 — Hunyuan is velocity-native like Wan). + ``guidance`` is left at the forward default: base HunyuanVideo has ``guidance_embeds=False`` so + ``guidance_in is None`` and the arg is a no-op (BRINGUP: a guidance-distilled checkpoint with + ``guidance_embeds=True`` would need ``embedded_cfg_scale`` wired in here). + +* ``HunyuanVideoVAE`` — ``AutoencoderKLHunyuanVideo`` normalizes by a scalar ``scaling_factor`` (0.476986), + NOT Wan's per-channel ``latents_mean``/``latents_std``. DiT latent = ``raw_z * scaling_factor`` on + encode; decode divides ``latents / scaling_factor`` before ``module.decode``. Decode output is video in + ``[-1, 1]``. + +* ``HunyuanVideoLlamaEncoder`` / ``HunyuanVideoCLIPEncoder`` — the dual text path. LLaMA wraps the prompt + in the ``PROMPT_TEMPLATE_ENCODE_VIDEO`` chat template, runs with ``output_hidden_states=True``, takes the + intermediate hidden state ``hidden_states[-(2+1)]`` (skip the last 2 layers), and crops the first 95 + template tokens (``crop_start=95``). CLIP returns the ``pooler_output`` (768-dim global vector). Faithful + to ``fastvideo/configs/pipelines/hunyuan.py`` (``llama_preprocess_text`` / ``llama_postprocess_text`` / + ``clip_postprocess_text``). +""" +from __future__ import annotations + +import torch + +from v2.platform.backends.torch_backend import ( + NUM_TRAIN_TIMESTEPS, + TorchComponent, + _to_numpy, +) + +# The LLaMA prompt template + crop verbatim from fastvideo/configs/pipelines/hunyuan.py (recipe DATA, not +# model code, so v2 owns its own copy rather than importing the private pipeline symbol). +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " + "1. The main content and theme of the video." + "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." + "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." + "4. background environment, light, style and atmosphere." + "5. camera angles, movements, and transitions used in the video:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") +LLAMA_CROP_START = 95 # the template's fixed system-prompt token count (crop_start) +LLAMA_HIDDEN_STATE_SKIP_LAYER = 2 # use hidden_states[-(skip+1)], i.e. skip the last 2 LLaMA layers +LLAMA_TEXT_LEN = 256 # LlamaArchConfig.text_len (the tokenizer max_length before the template) +CLIP_TEXT_LEN = 77 # CLIPTextConfig.text_len + + +# --------------------------------------------------------------------------- # +# DiT adapter # +# --------------------------------------------------------------------------- # +class HunyuanVideoDiT(TorchComponent): + """``dit(latent[C,T,H,W], text_embed=llama_hidden[L,4096], sigma, context=clip_pooled[768])`` -> + velocity[C,T,H,W]. Real forward (hunyuanvideo.py): ``forward(hidden_states[B,C,T,H,W], + encoder_hidden_states=[llama, clip], timestep[B], guidance=None) -> velocity (bare tensor)``.""" + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + hs = self._t(latent) # [1, C, T, h, w] + llama = self._t(text_embed) # [1, L, 4096] per-token LLaMA sequence + # timestep = sigma*1000 (the scalar per-batch convention the time_in TimestepEmbedder consumes). + timestep = torch.tensor([float(sigma) * NUM_TRAIN_TIMESTEPS], device=self.device) + if context is not None: # the faithful 2-element list path + clip_pooled = self._t(context) # [1, 768] CLIP global vector + encoder_hidden_states: object = [llama, clip_pooled] + else: # degenerate: a single tensor (token 0 pooled) + encoder_hidden_states = llama + with self._ctx(): + # guidance=None -> forward default 6016.0, but base guidance_in is None so it is ignored + # (do NOT wire embedded_cfg_scale unless a guidance_embeds=True checkpoint is loaded). The DiT + # does an internal sequence-parallel shard/all-gather that no-ops at sp=1 (the bring-up default). + velocity = self.module(hidden_states=hs, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + guidance=None) + return self._n(velocity) # rectified-flow velocity (bare tensor) + + +# --------------------------------------------------------------------------- # +# VAE adapter # +# --------------------------------------------------------------------------- # +class HunyuanVideoVAE(TorchComponent): + """``AutoencoderKLHunyuanVideo`` uses a scalar ``scaling_factor`` (0.476986), not Wan's per-channel + mean/std. DiT latent = ``raw_z * scaling_factor``; decode divides by it before ``module.decode``.""" + + def _scaling_factor(self) -> float: + # ParallelTiledVAE exposes ``scaling_factor`` as a property reading config.scaling_factor. + sf = getattr(self.module, "scaling_factor", None) + if sf is None: + sf = getattr(getattr(self.module, "config", None), "scaling_factor", 0.476986) + if sf is None: # defensive: neither property nor config carried it + sf = 0.476986 + return float(sf.item() if hasattr(sf, "item") else sf) + + @torch.no_grad() + def encode(self, video): + x = self._t(video) + dist = self.module.encode(x) # DiagonalGaussianDistribution + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + return self._n(z.float() * self._scaling_factor()) # -> the scaled latent the DiT denoises + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent).float() / self._scaling_factor() # invert the encode scaling -> raw latent + video = self.module.decode(z.to(self.dtype)) # -> video [B,3,T,H,W] in [-1,1] + return self._n(video) + + +# --------------------------------------------------------------------------- # +# Text encoders (dual: LLaMA per-token sequence + CLIP pooled global vector) # +# --------------------------------------------------------------------------- # +class HunyuanVideoLlamaEncoder(TorchComponent): + """LLaMA primary encoder: wrap the prompt in ``PROMPT_TEMPLATE_ENCODE_VIDEO``, run with + ``output_hidden_states=True``, take ``hidden_states[-(skip+1)]`` (skip the last 2 layers) and crop the + first 95 template tokens. Returns the per-token sequence [L, 4096] (the DiT ``txt`` / ``encoder_hidden_states[0]``). + The module is loaded with ``output_hidden_states=True`` (HunyuanConfig.__post_init__ sets it).""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = LLAMA_TEXT_LEN): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + @torch.no_grad() + def encode(self, text): + prompt = PROMPT_TEMPLATE_ENCODE_VIDEO.format(text or "") + toks = self.tokenizer(prompt, + return_tensors="pt", + max_length=self.max_length + LLAMA_CROP_START, + 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, output_hidden_states=True) + hidden_states = getattr(out, "hidden_states", None) + if hidden_states is None: + raise RuntimeError("HunyuanVideoLlamaEncoder: module returned no hidden_states; load LLaMA with " + "output_hidden_states=True (HunyuanConfig.__post_init__ sets it).") + hidden = hidden_states[-(LLAMA_HIDDEN_STATE_SKIP_LAYER + 1)] # skip the last 2 layers + hidden = hidden[:, LLAMA_CROP_START:] # drop the template's 95 system tokens + return _to_numpy(hidden.squeeze(0)) + + +class HunyuanVideoCLIPEncoder(TorchComponent): + """CLIP secondary encoder: returns ``pooler_output`` (the 768-dim global vector, the DiT + ``text_states_2`` / ``encoder_hidden_states[1]``). + + BRINGUP: the shared ``_make_text_encoder`` passes the ``tokenizer/`` subfolder (the LLaMA tokenizer) as + the constructor's ``tokenizer`` arg; HunyuanVideo's CLIP tokenizer lives in ``tokenizer_2/``. GPU-verify + must wire the CLIP tokenizer here (e.g. via a card-side stamp of ``tokenizer_2`` or loading it lazily) — + on CPU this path is exercised by the ``ToyTextEncoder`` factory, which needs no tokenizer.""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = CLIP_TEXT_LEN): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + @torch.no_grad() + def encode(self, text): + toks = self.tokenizer(text or "", + return_tensors="pt", + max_length=self.max_length, + truncation=True, + padding="max_length") + 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) + pooled = out.pooler_output if hasattr(out, "pooler_output") else out[1] + return _to_numpy(pooled.squeeze(0)) # [768] global vector diff --git a/v2/platform/backends/torch_hunyuan_video15.py b/v2/platform/backends/torch_hunyuan_video15.py new file mode 100644 index 000000000..519b27a09 --- /dev/null +++ b/v2/platform/backends/torch_hunyuan_video15.py @@ -0,0 +1,192 @@ +"""HunyuanVideo 1.5 torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` +so the recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder`` +dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +Architecture deltas vs Wan that force NEW adapters (all faithful to +``fastvideo/models/dits/hunyuanvideo15.py`` + ``fastvideo/pipelines/stages/denoising.py``): + +* ``HunyuanVideo15DiT`` — the rectified-flow velocity network ``F_θ``. Unlike ``WanDiT`` (single text + embed, single image embed), the real ``forward`` takes ``encoder_hidden_states`` as a **LIST of two** + tensors ``[qwen_hidden(3584-d), byt5_hidden(1472-d)]`` and ``encoder_hidden_states_image`` as a + **LIST of one** image-embed tensor (all-zeros ``[1,729,1152]`` for t2v → the DiT's + ``torch.all(image_embeds==0)`` auto-selects the t2v token-ordering branch). The loop hands the two + text embeds as a tuple in ``text_embed`` and the image embeds in ``context``; this adapter marshals + them into the two lists internally and builds the zero image-embed default, so the loop's dit-call + signature stays ``dit(latent, text_embed, sigma, context=, cond=)`` (toy-compatible). For i2v the + loop passes the 33-channel first-frame conditioning latent (32 VAE + 1 mask) in ``cond``; the adapter + concats it onto the noisy latent on the channel dim to reach ``in_channels=65`` BEFORE the forward + (mirrors ``denoising.py`` line 393's ``torch.cat([latent, image_latent], dim=1)``). ``timestep`` is + ``sigma*1000`` (the scheduler's discrete timestep; ``scale_model_input`` is identity for flow-match). + The network predicts velocity directly → returned via ``self._n()`` with NO x0→velocity conversion + (unlike ``LTX2DiT``). ``use_meanflow`` SR variants pass ``timestep_r`` — out of scope here (t2v/i2v + ship with ``use_meanflow=False``), so ``timestep_r=None``. +* ``HunyuanVideo15VAE`` — the HunyuanVideo 1.5 causal-3D VAE (z=32, 16× spatial / 4× temporal). Its DiT + latent space is normalized by a single **scalar** ``scaling_factor=1.03682`` (NOT Wan's per-channel + ``latents_mean``/``latents_std`` — the ``WanVAE`` adapter would corrupt these latents, BRINGUP risk D). + Decode: ``module.decode(latent / scaling_factor)``; encode (i2v first-frame cond): ``z * scaling_factor``. +* ``HunyuanVideo15QwenEncoder`` — Qwen2.5-VL primary text encoder. Faithful to + ``configs/pipelines/hunyuan15.py:qwen_preprocess_text/qwen_postprocess_text``: wraps the prompt in the + system+chat template, requests ``output_hidden_states=True``, takes ``hidden_states[-3]`` (3rd-from-last + layer, NOT ``last_hidden_state``), and crops the first 108 prompt-template tokens. The generic + ``T5Encoder.encode`` (last_hidden_state, pad-to-max) is WRONG for Qwen (BRINGUP risk B). +* ``HunyuanVideo15ByT5Encoder`` — ByT5/Glyph secondary encoder. ``byt5_preprocess_text`` extracts quoted + glyph texts via regex; when the prompt has no quoted text it is ``None`` and the real pipeline emits a + zero-length ``[1, 0, 1472]`` embedding (the DiT always concatenates the txt_in_2 output) — this adapter + reproduces that (BRINGUP risk C). Output is the raw ``last_hidden_state``. + +All forwards run inside ``self._ctx()`` (FastVideo attention reads attn_metadata via the forward +context). BRINGUP: the DiT's ``sequence_model_parallel_shard``/all-gather is a no-op at sp=1 (the card +defaults parallelism to single() until SP is validated — risk G). +""" +from __future__ import annotations + +import re + +import numpy as np +import torch + +from v2.platform.backends.torch_backend import ( + NUM_TRAIN_TIMESTEPS, + T5Encoder, + TorchComponent, + _to_numpy, +) + +# Faithful constants from configs/pipelines/hunyuan15.py + Hy15ImageEncodingStage. +HUNYUAN15_VAE_SCALING_FACTOR = 1.03682 # scalar (NOT per-channel mean/std) +HUNYUAN15_IMAGE_TOKENS = 729 # SigLIP token count placeholder +HUNYUAN15_IMAGE_DIM = 1152 # config.image_embed_dim +HUNYUAN15_BYT5_DIM = 1472 # config.text_embed_2_dim +PROMPT_TEMPLATE_TOKEN_LENGTH = 108 # tokens of the chat-template prefix to crop off Qwen +QWEN_HIDDEN_STATE_INDEX = -3 # qwen_postprocess takes hidden_states[-3] + +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "You are a helpful assistant. Describe the video by detailing the following aspects: " + " 1. The main content and theme of the video. " + " 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. " + " 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. " + " 4. background environment, light, style and atmosphere. " + " 5. camera angles, movements, and transitions used in the video.") + +# byt5_preprocess_text: pull the quoted glyph strings out of the prompt (HunyuanVideo 1.5 glyph rendering). +_GLYPH_PATTERN = r"\"(.*?)\"|“(.*?)”" + + +def _extract_glyph_texts(prompt: str) -> str | None: + """Faithful to ``configs/pipelines/hunyuan15.py:extract_glyph_texts`` — returns ``None`` when the + prompt carries no quoted text (the common t2v case), else a ``Text "..."``-joined glyph string.""" + matches = re.findall(_GLYPH_PATTERN, prompt or "") + result = [m[0] or m[1] for m in matches] + result = list(dict.fromkeys(result)) if len(result) > 1 else result + return ". ".join([f'Text "{t}"' for t in result]) + ". " if result else None + + +class HunyuanVideo15DiT(TorchComponent): + """``dit(latent[C,T,H,W], text_embed=(qwen,byt5), sigma, context=img_embeds, cond=image_latent)``. + + Marshals the two text embeds + the (zero-for-t2v) image embeds into the LIST args the real + ``HunyuanVideo15Transformer3DModel.forward`` expects, concats the 33ch i2v cond onto the latent, and + returns the rectified-flow velocity (NO x0 conversion). Faithful to ``denoising.py`` (timestep = + sigma*1000, identity ``scale_model_input``, ``latent_model_input = cat([latent, image_latent])``).""" + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + hs = self._t(latent) # [1, C(32), T, H, W] + if cond is not None: # i2v: concat first-frame cond latent (33ch) -> 65ch + hs = torch.cat([hs, self._t(cond)], dim=1) + # ``text_embed`` is a (qwen_hidden, byt5_hidden) tuple from the loop. Both are already numpy + # [seq, dim]; the byt5 element may be a zero-length [0, 1472] (no glyph text). + qwen_e, byt5_e = text_embed if isinstance(text_embed, tuple | list) else (text_embed, None) + qwen_t = self._t(qwen_e) # [1, qseq, 3584] + if byt5_e is None: + byt5_t = torch.zeros(1, 0, HUNYUAN15_BYT5_DIM, device=self.device, dtype=self.dtype) + else: + byt5_t = self._t(byt5_e) # [1, bseq, 1472] + # encoder_hidden_states_image: zeros for t2v (the DiT's ``torch.all(==0)`` selects the t2v branch). + if context is None: + img_t = torch.zeros(1, HUNYUAN15_IMAGE_TOKENS, HUNYUAN15_IMAGE_DIM, device=self.device, dtype=self.dtype) + else: + img_t = self._t(context) + ts = float(sigma) * NUM_TRAIN_TIMESTEPS # the scheduler's discrete timestep (BRINGUP risk B) + timestep = torch.tensor([ts], device=self.device, dtype=self.dtype).expand(hs.shape[0]) + with self._ctx(): + velocity = self.module(hidden_states=hs, + encoder_hidden_states=[qwen_t, byt5_t], + timestep=timestep, + encoder_hidden_states_image=[img_t], + guidance=None, + timestep_r=None, + attention_kwargs=None) + return self._n(velocity) # rectified-flow velocity (BRINGUP risk C) + + +class HunyuanVideo15VAE(TorchComponent): + """Scalar-scaling VAE. Decode un-scales (``latent / scaling_factor``) before ``module.decode``; encode + (i2v first-frame cond) scales the latent (``z * scaling_factor``). NO per-channel mean/std (risk D).""" + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent).float() / HUNYUAN15_VAE_SCALING_FACTOR + video = self.module.decode(z.to(self.dtype)) # -> video [B,3,T,H,W] in [-1,1] + return self._n(video) + + @torch.no_grad() + def encode(self, video): + x = self._t(video) + dist = self.module.encode(x) + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + return self._n(z.float() * HUNYUAN15_VAE_SCALING_FACTOR) + + +class HunyuanVideo15QwenEncoder(TorchComponent): + """Qwen2.5-VL primary text encoder. ``encode(text) -> embedding[seq, 3584]`` (numpy). + + Faithful to qwen_preprocess/postprocess: chat-template wrap, ``output_hidden_states=True``, take + ``hidden_states[-3]``, crop the first 108 template tokens. The HF tokenizer's ``apply_chat_template`` + builds the prompt; ``max_length`` defaults to the pipeline's 1000+108 (risk B).""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = 1000 + PROMPT_TEMPLATE_TOKEN_LENGTH): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + @torch.no_grad() + def encode(self, text): + # Qwen2_5_VLProcessor wraps an inner ``.tokenizer``; fall back to the tokenizer itself. + tok = getattr(self.tokenizer, "tokenizer", self.tokenizer) + messages = [{ + "role": "system", + "content": PROMPT_TEMPLATE_ENCODE_VIDEO + }, { + "role": "user", + "content": text if text else " " + }] + formatted = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + enc = tok(formatted, return_tensors="pt", max_length=self.max_length, truncation=True) + ids = enc.input_ids.to(self.device) + mask = enc.attention_mask.to(self.device) + with self._ctx(): + out = self.module(input_ids=ids, attention_mask=mask, output_hidden_states=True) + hidden = out.hidden_states[QWEN_HIDDEN_STATE_INDEX] # 3rd-from-last layer (NOT last) + hidden = hidden[:, PROMPT_TEMPLATE_TOKEN_LENGTH:].squeeze(0) # crop the chat-template prefix + return _to_numpy(hidden) + + +class HunyuanVideo15ByT5Encoder(T5Encoder): + """ByT5/Glyph secondary encoder. ``encode(text) -> embedding[seq, 1472]``. Extracts the quoted glyph + text first (``byt5_preprocess_text``); when none, returns a zero-length ``[0, 1472]`` embedding — the + real pipeline's special-case (the DiT always concats txt_in_2 output, risk C). Output is the raw + ``last_hidden_state`` (NO Wan zero-pad-to-max).""" + + @torch.no_grad() + def encode(self, text): + glyph = _extract_glyph_texts(text) + if glyph is None: # no quoted text -> empty embedding + return np.zeros((0, HUNYUAN15_BYT5_DIM), dtype="float32") + toks = self.tokenizer(glyph, 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) + return _to_numpy(torch.nan_to_num(hidden, nan=0.0)) diff --git a/v2/platform/backends/torch_hunyuangamecraft.py b/v2/platform/backends/torch_hunyuangamecraft.py new file mode 100644 index 000000000..620ecce71 --- /dev/null +++ b/v2/platform/backends/torch_hunyuangamecraft.py @@ -0,0 +1,173 @@ +"""HunyuanGameCraft torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` +so the GameCraft recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/ +``_make_text_encoder`` dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only +on a GPU box. + +GameCraft is an *interactive* world model (camera/action-conditioned i2v). The deltas vs Wan/Cosmos that +force bespoke adapters (the BRINGUP/blocker list in the port spec): + +* ``GameCraftDiT`` — the DiT forward expects a **33-channel** input ``x = cat([latent16 | gt_latent16 | + mask1], dim=1)`` that the *denoise loop* (not the DiT) assembles, plus a **list** of text states + ``encoder_hidden_states = [llama_states(4096d), clip_pooled(768d)]``, an optional ``camera_states`` + (CameraNet Plücker), ``guidance=None`` (NO embedded guidance — plain ClassicCFG), and the model + timestep ``σ·1000``. The DiT returns a 16-channel rectified-flow velocity. The loop hands this adapter + the bare 16ch latent + the cond slots; the adapter does the 33ch concat + dual-text-list packing + INTERNALLY so the loop's dit-call stays ``dit(latent, text_embed, sigma, ...)`` (toy-compatible). + Cannot reuse ``WanDiT`` (it concats a 20ch i2v cond and passes ``encoder_hidden_states_image``). +* ``GameCraftVAE`` — SCALAR ``scaling_factor=0.476986`` normalization (``latent = sample()·sf``; + ``decode(latent/sf)``), NOT Wan's per-channel ``latents_mean``/``latents_std``. Reusing ``WanVAE`` + (which reads ``module.latents_mean``) would crash / corrupt latents. +* ``GameCraftLlamaEncoder`` / ``GameCraftClipEncoder`` — the two text encoders are LLaVA-LLaMA-3 (4096d + hidden states, the DiT's ``text_states``) and CLIP ViT-L/14 (768d pooled, the DiT's ``text_states_2``). + The toy CPU path uses ``ToyTextEncoder`` for both; these GPU subclasses are written-not-run. + +Everything below is faithful to ``fastvideo/pipelines/stages/gamecraft_denoising.py`` (the DiT call) and +``fastvideo/pipelines/stages/gamecraft_image_encoding.py`` (the scalar VAE scaling). The camera / action +(CameraNet) conditioning needs a request-API camera-input channel v2 lacks today, so the t2v/degenerate +path passes ``camera_states=None`` (the DiT's ``if camera_states is not None`` branch is skipped) — see +the BRINGUP notes in ``loop.py``/``program.py``. +""" +from __future__ import annotations + +import torch + +from v2.platform.backends.torch_backend import T5Encoder, TorchComponent, _to_numpy + +# Flow-match timestep convention (same as Wan): the loop hands the raw sigma (1->0); the GameCraft DiT +# (like diffusers/FastVideo) embeds ``timestep = sigma * num_train_timesteps``. BRINGUP risk B. +NUM_TRAIN_TIMESTEPS = 1000 +GAMECRAFT_SCALING_FACTOR = 0.476986 # GameCraftVAE.config.scaling_factor (scalar, NOT per-channel stats) + + +class GameCraftDiT(TorchComponent): + """``dit(latent[16,T,h,w], text_embed, sigma, context=clip_pooled, *, cond=cond_pack) -> velocity[16,T,h,w]``. + + The loop calls this with the bare 16-channel noise latent. The adapter assembles the 33-channel DiT + input + the dual-text list + the camera states INTERNALLY (so the loop's dit-call signature matches + the toy ``ToyDiT.__call__``). Faithful to ``GameCraftDenoisingStage.forward``: + + latent_model_input = cat([latents, gt_latents, conditioning_mask], dim=1) # [B,33,T,h,w] + noise_pred = transformer(latent_model_input, [llama, clip], t·1000, + camera_states=cam, guidance=None) # [B,16,T,h,w] + + Conventions threaded through the call boundary (all numpy at the seam): + * ``text_embed`` -> ``text_states`` (LLaMA hidden states), ``encoder_hidden_states[0]``. + * ``context`` -> ``text_states_2`` (CLIP pooled), ``encoder_hidden_states[1]``. The loop reuses + the WanDenoiseLoop ``context=`` kwarg slot to carry the per-branch CLIP pooled embed (cond vs + uncond), so CFG combines correctly. ``None`` -> single-text fallback (still valid for the toy). + * ``cond`` -> a dict ``{"gt_latents", "mask", "camera_states"}`` the loop builds once and + passes every step. ``None`` (pure t2v) -> zero gt_latents + zero mask + no camera (the + ``GameCraftDenoisingStage`` fallback), i.e. the degenerate standard-denoise path.""" + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + hs = self._t(latent) # [1, 16, T, h, w] + b, _c, t, h, w = hs.shape + # --- 33-channel concat [latent16 | gt_latent16 | mask1] (assembled HERE, not in the DiT) ------ # + cond = cond or {} + gt = self._t(cond.get("gt_latents")) if cond.get("gt_latents") is not None else torch.zeros_like(hs) + mask = (self._t(cond.get("mask")) + if cond.get("mask") is not None else torch.zeros(b, 1, t, h, w, device=self.device, dtype=self.dtype)) + model_input = torch.cat([hs, gt, mask], dim=1) # [1, 33, T, h, w] + # --- dual text states list: [LLaMA(4096d), CLIP-pooled(768d)] -------------------------------- # + ehs = [self._t(text_embed)] # text_states -> [1, seq, 4096] + if context is not None: # CLIP pooled (per CFG branch) -> text_states_2 + ehs.append(self._t(context)) + # --- camera_states (CameraNet Plücker): None for the t2v/degenerate path (BRINGUP) ----------- # + camera = self._t(cond.get("camera_states")) if cond.get("camera_states") is not None else None + # timestep = sigma*1000, shape [B] (the stage does ``t.repeat(B)``). BRINGUP risk B. + ts = torch.tensor([float(sigma) * NUM_TRAIN_TIMESTEPS], device=self.device, dtype=self.dtype).expand(b) + with self._ctx(): + velocity = self.module( + model_input, + ehs, + ts, + camera_states=camera, + guidance=None, # NO embedded guidance (guidance_embeds=False); plain CFG + return_dict=False) + return self._n(velocity) # rectified-flow velocity (loop integrates via FLOW_MATCH_STEP) + + +class GameCraftVAE(TorchComponent): + """Scalar-scaling VAE adapter (GameCraft uses a single ``scaling_factor``, NOT per-channel stats). + + Faithful to ``GameCraftImageVAEEncodingStage`` (encode: ``sample()·sf``) and ``DecodingStage`` (decode: + ``decode(latent/sf)``). The reference-image preprocessing (resize/center-crop/normalize) lives in the + program's image-cond node, not here — this adapter only marshals + applies the scalar factor.""" + + def __init__(self, module, *, device, dtype, scaling_factor: float = GAMECRAFT_SCALING_FACTOR): + super().__init__(module, device=device, dtype=dtype) + # Prefer the loaded module's own config when present; fall back to the published GameCraft value. + self.scaling_factor = float(getattr(getattr(module, "config", None), "scaling_factor", scaling_factor)) + + @torch.no_grad() + def encode(self, video): + x = self._t(video) + out = self.module.encode(x) + dist = getattr(out, "latent_dist", out) + z = dist.sample() if hasattr(dist, "sample") else (dist.mode() if hasattr(dist, "mode") else dist) + return self._n(z.float() * self.scaling_factor) # scaled latent the DiT expects + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent).float() / self.scaling_factor # invert the scalar scaling -> raw latent + video = self.module.decode(z.to(self.dtype)) + if hasattr(video, "sample"): + video = video.sample + return self._n(video) # video [B,3,T,H,W] in [-1,1] + + +class GameCraftClipEncoder(TorchComponent): + """CLIP ViT-L/14 text encoder -> POOLED embedding (768d), the DiT's ``text_states_2``. + + BRINGUP: written-not-run. The program writes this into the ``clip_text_embeds`` slot; the loop threads + it through the ``context=`` slot per CFG branch. Returns the pooled output (``[768]`` or ``[1,768]``), + matching ``batch.clip_embedding_pos``/``clip_embedding_neg`` in the fastvideo stage.""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = 77): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + @torch.no_grad() + def encode(self, text): + toks = self.tokenizer(text or "", + return_tensors="pt", + max_length=self.max_length, + truncation=True, + padding="max_length") + 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) + # CLIPTextModel exposes ``pooler_output`` (the [EOS]-token projection) — the GameCraft text_states_2. + pooled = getattr(out, "pooler_output", None) + if pooled is None: + pooled = (out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0])[:, -1] + return _to_numpy(pooled.squeeze(0)) + + +class GameCraftLlamaEncoder(T5Encoder): + """LLaVA-LLaMA-3-8B text encoder -> hidden states (4096d), the DiT's ``text_states``. + + BRINGUP: written-not-run. GameCraft consumes an intermediate hidden-state layer via the Hunyuan + ``llama_postprocess`` (``output_hidden_states=True``) plus the LLaVA prompt template — those + preprocess/postprocess functions are not yet ported to v2 (blocker 4). This subclass currently reuses + the (U)MT5 ``encode`` contract (real-token rows, no fixed zero-pad surprises) as a structural + placeholder; the real LLaVA template + hidden-layer pick must be wired during GPU bring-up.""" + + @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(): + # BRINGUP: real GameCraft passes output_hidden_states=True and picks an intermediate layer via + # llama_postprocess; here we take last_hidden_state as the structural stand-in. + 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) + return _to_numpy(torch.nan_to_num(hidden, nan=0.0)) + + +# Re-export the scalar scaling constant for the loop/program to keep the value single-sourced. +__all__ = ["GameCraftDiT", "GameCraftVAE", "GameCraftClipEncoder", "GameCraftLlamaEncoder", "GAMECRAFT_SCALING_FACTOR"] diff --git a/v2/platform/backends/torch_hyworld.py b/v2/platform/backends/torch_hyworld.py new file mode 100644 index 000000000..4e2c4a0a4 --- /dev/null +++ b/v2/platform/backends/torch_hyworld.py @@ -0,0 +1,216 @@ +"""HY-WorldPlay torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so the +recipe is self-contained (no edit to the shared ``_make_dit`` / ``_make_vae`` / ``_make_text_encoder`` / +``_make_image_encoder`` dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only +on a GPU box. + +Faithful to ``fastvideo/pipelines/stages/hyworld_denoising.py:HYWorldDenoisingStage`` and +``fastvideo/models/dits/hyworld/hyworld.py:HYWorldTransformer3DModel.forward``: + + * ``HYWorldDiT`` — the rectified-flow velocity predictor. Its ``__call__`` keeps the + cosmos2/ToyDiT-compatible ``dit(latent, text_embed, sigma, context=...)`` signature and marshals ALL + the HY-WorldPlay specifics INTERNALLY: the 65ch ``[noise(32) | cond(32) + mask(1)]`` channel-concat, + the LIST ``encoder_hidden_states=[qwen, byt5]`` + LIST ``encoder_hidden_states_image=[siglip]``, the + per-latent-frame heterogeneous ``timestep`` (context frames pinned at ``stabilization_level - 1``) + + the scalar ``timestep_txt``, the per-frame ``action`` (flattened B*T) + ``viewmats[B,T,4,4]`` / + ``Ks[B,T,3,3]``, and the ``encoder_attention_mask=[mask, mask2]`` swapped per pos/neg branch. Returns + the bare velocity tensor [B,32,T,H,W]. BRINGUP: written-not-run; needs the GPU dense+ProPE flash-attn + backend + a pose-carrying request. + * ``HYWorldVAE`` — encode -> ``.mode() * scaling_factor`` (1.03682, no shift/mean); decode inverts. + * ``HYWorldQwenEncoder`` / ``HYWorldByT5Encoder`` / ``HYWorldSiglipEncoder`` — the 3 conditioning + streams. The Qwen adapter runs the video chat template, takes ``hidden_states[-3]``, and crops the + first 108 template tokens (BRINGUP). For plain t2v the byt5/siglip streams are zeroed. +""" +from __future__ import annotations + +from typing import Any + +import torch + +from v2.platform.backends.torch_backend import T5Encoder, TorchComponent, _to_numpy + +# Flow-match timestep convention: timestep = sigma * num_train_timesteps (matches WanDiT / the +# HYWorldDenoisingStage which feeds the scheduler's discrete timesteps in 0..1000). +NUM_TRAIN_TIMESTEPS = 1000 +# SigLIP image stream geometry (HunyuanVideo1.5 / HY-WorldPlay): 729 tokens x 1152 dim. +SIGLIP_TOKENS = 729 +SIGLIP_DIM = 1152 +# HY-WorldPlay latent geometry (AutoencoderKLHYWorld): z=32 + 1 mask channel for the cond latent. +HYWORLD_LATENT_CHANNELS = 32 +HYWORLD_VAE_SCALING_FACTOR = 1.03682 + + +class HYWorldDiT(TorchComponent): + """``dit(latent[C,T,h,w], text_embed[seq,dim], sigma, context=chunk_ctx) -> velocity[32,T,h,w]``. + + ``latent`` is the current chunk's NOISE latent (32ch). ``text_embed`` is the Qwen (mllm) embed for + the active branch; the byt5 + siglip streams and the camera/action tensors are resolved from the + sibling components + ``context`` here. ``sigma`` is the live step's sigma (-> ``timestep = sigma*1000`` + broadcast per latent frame; context/history frames in ``context`` are pinned at the stabilization + level). The 65ch ``[noise | cond + mask]`` concat is done HERE (faithful to the stage), so the loop + never assembles arch-specifics. BRINGUP: the chunk>0 camera-aligned memory path + ProPE attention + + action need a pose-carrying request; absent, zero camera/action are built (degenerate t2v).""" + + def __init__(self, module, *, device, dtype): + super().__init__(module, device=device, dtype=dtype) + + def _zero_image_embeds(self) -> torch.Tensor: + # T2V: zero SigLIP embeds -> the DiT's ``torch.all(image_embeds == 0)`` masks the image stream. + return torch.zeros(1, SIGLIP_TOKENS, SIGLIP_DIM, device=self.device, dtype=self.dtype) + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context: Any = None, *, cond=None): + ctx = context if isinstance(context, dict) else {} + hs = self._t(latent) # [1, 32, T, h, w] + _b, _c, t, _h, _w = hs.shape + # 65ch in-channel concat: [noise(32) | cond_latent(32) + mask(1)]. The cond latent is the + # VAE-encoded first frame (i2v) or zeros (t2v) + a 1ch mask -> 33ch. BRINGUP: i2v_cond plumbing. + i2v_cond = ctx.get("i2v_cond") + if i2v_cond is not None: + cond_latents = self._t(i2v_cond)[:, :, :t] # [1, 33, T, h, w] + else: + cond_latents = torch.zeros(1, HYWORLD_LATENT_CHANNELS + 1, t, _h, _w, device=self.device, dtype=self.dtype) + hs = torch.cat([hs, cond_latents], dim=1) # -> 65ch (faithful to the stage's torch.concat) + + # 3-stream conditioning: encoder_hidden_states must be a LIST [qwen, byt5]; image a LIST [siglip]. + qwen = self._t(text_embed) # [1, seq, dim] + byt5 = self._t(ctx.get("byt5_embeds")) if ctx.get("byt5_embeds") is not None else \ + torch.zeros(1, 1, qwen.shape[-1], device=self.device, dtype=self.dtype) + img = self._t(ctx.get("i2v_img_embeds")) if ctx.get("i2v_img_embeds") is not None else \ + self._zero_image_embeds() + # Per-pos/neg attention masks (BRINGUP: the stage swaps prompt_attention_mask <-> negative). The + # loop runs each branch separately, so the matching mask is selected per call by the caller's text. + mask = torch.ones(1, qwen.shape[1], device=self.device, dtype=self.dtype) + mask2 = torch.ones(1, byt5.shape[1], device=self.device, dtype=self.dtype) + + # Per-latent-frame timestep (NOT scalar): live ``t`` for the chunk frames; context/history frames + # (chunk>0, BRINGUP) pinned at ``stabilization_level - 1``. timestep_txt is a separate scalar. + ts_val = float(sigma) * NUM_TRAIN_TIMESTEPS + timestep = torch.full((t, ), ts_val, device=self.device, dtype=torch.long) + timestep_txt = torch.tensor([ts_val], device=self.device, dtype=torch.long) + + # Camera + action conditioning (BRINGUP): per-frame viewmats/Ks/action. Absent -> identity camera + # + zero action (the action_in / prope projections are zero-init, so this is a faithful no-op). + cam = ctx.get("camera") + if cam is not None: + start, end = ctx.get("chunk", (0, t)) + viewmats = self._t(cam["viewmats"])[:, start:end] + Ks = self._t(cam["Ks"])[:, start:end] + action = self._t(cam["action"])[:, start:end].reshape(-1) + else: + viewmats = torch.eye(4, device=self.device, dtype=self.dtype).expand(1, t, 4, 4).contiguous() + Ks = torch.eye(3, device=self.device, dtype=self.dtype).expand(1, t, 3, 3).contiguous() + action = torch.zeros(t, device=self.device, dtype=self.dtype) + + with self._ctx(): + velocity = self.module( + hs, + [qwen, byt5], + timestep=timestep, + encoder_hidden_states_image=[img], + encoder_attention_mask=[mask, mask2], + action=action.to(self.dtype), + viewmats=viewmats.to(self.dtype), + Ks=Ks.to(self.dtype), + timestep_txt=timestep_txt, + ) + return self._n(velocity) # bare rectified-flow velocity [32, T, h, w] + + +class HYWorldVAE(TorchComponent): + """AutoencoderKLHYWorld: encode -> ``.mode() * scaling_factor`` (1.03682, no shift/mean); + decode -> ``/ scaling_factor`` then decode. Used for the first-frame cond latent + the final decode.""" + + def __init__(self, module, *, device, dtype, scaling_factor: float = HYWORLD_VAE_SCALING_FACTOR): + super().__init__(module, device=device, dtype=dtype) + self.scaling_factor = float(getattr(getattr(module, "config", None), "scaling_factor", scaling_factor)) + + @torch.no_grad() + def encode(self, video): + x = self._t(video) + dist = self.module.encode(x) + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + return self._n(z.float() * self.scaling_factor) + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent).float() / self.scaling_factor + video = self.module.decode(z.to(self.dtype)) # -> video [B,3,T,H,W] in [-1,1] + return self._n(video) + + +class HYWorldQwenEncoder(TorchComponent): + """Qwen2.5-VL mllm (primary text). BRINGUP: faithful path runs the video chat template + (PROMPT_TEMPLATE_ENCODE_VIDEO system msg), takes ``hidden_states[-3]``, and crops the first 108 + template tokens. Returns the cropped real-token embedding (numpy). Encoder weights are fp32 here; on + the box mirror the fastvideo TextEncodingStage marshalling exactly.""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = 256, crop_start: int = 108): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + self.crop_start = crop_start # PROMPT_TEMPLATE_ENCODE_VIDEO prepends 108 template tokens + + @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, output_hidden_states=True) + hs = getattr(out, "hidden_states", None) + # hidden_states[-3] is the layer the HY-WorldPlay mllm conditioning uses (BRINGUP: confirm on box). + hidden = hs[-3] if hs is not None and len(hs) >= 3 else \ + (out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]) + hidden = hidden.squeeze(0) + if self.crop_start and hidden.shape[0] > self.crop_start: # crop the prompt-template tokens + hidden = hidden[self.crop_start:] + return _to_numpy(torch.nan_to_num(hidden, nan=0.0)) + + +class HYWorldByT5Encoder(T5Encoder): + """ByT5 glyph encoder (rendered text; text_states_dim_2=1472). BRINGUP: the faithful path extracts + quoted glyph text (``extract_glyph_texts`` regex) before encoding and returns ``last_hidden_state``; + for plain t2v there is no glyph text -> a single zero token. Subclasses ``T5Encoder`` (raw, NO Wan + fixed-length zero-pad).""" + + @torch.no_grad() + def encode(self, text): + glyph = self._extract_glyph(text) + if not glyph: # no quoted glyph text -> a 1-token zero embed (the DiT cond_type_embed still adds) + dim = int(getattr(getattr(self.module, "config", None), "d_model", 1472)) + return _to_numpy(torch.zeros(1, dim, device=self.device, dtype=torch.float32)) + toks = self.tokenizer(glyph, 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) + return _to_numpy(torch.nan_to_num(hidden, nan=0.0)) + + @staticmethod + def _extract_glyph(text) -> str: + import re + if not text: + return "" + return " ".join(re.findall(r'["“”‘’\'](.+?)["“”‘’\']', text)) + + +class HYWorldSiglipEncoder(TorchComponent): + """SigLIP image encoder (729 tokens x 1152 dim) for the i2v/first-frame stream. BRINGUP: the faithful + path resizes+center-crops to the target res, runs ``processor.preprocess`` (fp16), and returns + ``last_hidden_state``. For t2v there is no image -> the program zeroes this stream (the DiT then masks + it via ``torch.all(image_embeds == 0)``).""" + + def __init__(self, module, processor=None, *, device, dtype): + super().__init__(module, device=device, dtype=dtype) + self.processor = processor + + @torch.no_grad() + def encode_image(self, image): + if image is None or self.processor is None: # t2v: zero embeds -> DiT masks the image stream + return _to_numpy(torch.zeros(SIGLIP_TOKENS, SIGLIP_DIM, device=self.device, dtype=torch.float32)) + inputs = self.processor.preprocess(images=image, return_tensors="pt").to(device=self.device, dtype=self.dtype) + with self._ctx(): + out = self.module(pixel_values=inputs["pixel_values"]) + embeds = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] + return _to_numpy(embeds.squeeze(0)) diff --git a/v2/platform/backends/torch_kandinsky5.py b/v2/platform/backends/torch_kandinsky5.py new file mode 100644 index 000000000..9f84cdea0 --- /dev/null +++ b/v2/platform/backends/torch_kandinsky5.py @@ -0,0 +1,215 @@ +"""Kandinsky-5.0 torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so +the recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder`` +dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +Faithful to the diffusers ``KandinskyV5...Pipeline`` (vendored under +``diffusers/pipelines/kandinsky5/pipeline_kandinsky.py``) and ``fastvideo.models.dits.kandinsky5``: + + * ``Kandinsky5DiT`` — velocity predictor. The loop hands a CHANNELS-FIRST ``latent[C,T,H,W]`` (Wan / + ToyDiT convention), the Qwen token embeds as ``text_embed``, the CLIP pooled vector via ``context``, + and ``sigma``; this adapter PERMUTES to the channels-LAST ``[B,T,H,W,C]`` the forward requires (and + permutes the velocity back). It also builds the THREE per-request positional inputs INTERNALLY from + the latent geometry + recovered resolution: ``visual_rope_pos`` over the patched grid + (patch_size=(1,2,2) -> ``[arange(T), arange(H//2), arange(W//2)]``), ``text_rope_pos`` over the Qwen + seq len, and the resolution-dependent ``scale_factor`` (= (1,2,2) for 480–854 else (1,3.16,3.16)). + ``timestep = σ·1000``. ``pooled_projections`` is MANDATORY (the forward raises if None); + ``sparse_params=None`` (the Lite checkpoint uses regular attention — the sparse/STA path raises + NotImplementedError in fastvideo, BRINGUP risk K). + * ``Kandinsky5QwenEncoder`` — Qwen2.5-VL. Applies the repo prompt template + start-idx slicing + (``prompt_template_encode_start_idx=129``, ``max_sequence_length=256``): the embeds are the LAST + hidden state (``output_hidden_states=True``) sliced from index 129 onward (BRINGUP risk G). + * ``Kandinsky5ClipEncoder`` — CLIP-with-projection. Returns the ``pooler_output`` (already projected to + 768) for a 77-token padded prompt. Its tokenizer lives in ``tokenizer_2`` (NOT the ``tokenizer`` the + built-in maker probes), so the adapter re-resolves it from the model root (BRINGUP). + * ``Kandinsky5VAE`` — AutoencoderKLHunyuanVideo. SCALAR ``scaling_factor`` (latent = raw·sf; decode raw + = latent/sf), NOT Wan per-channel (z-mean)/std. The loop's latent is already channels-first + ``[C,T,H,W]`` (the DiT adapter owns the channels-last conversion), so decode only adds a batch dim + + un-normalizes by the scalar before ``module.decode`` (BRINGUP risk F). +""" +from __future__ import annotations + +import os +from typing import Any + +import torch + +from v2.platform.backends.torch_backend import TorchComponent, _to_numpy + +# Repo defaults (diffusers pipeline + fastvideo Kandinsky5ArchConfig). The two misspellings inside the +# template below are VERBATIM from the upstream Kandinsky-5 prompt template (diffusers +# pipeline_kandinsky.py) — they must match the tokens the model was trained on, so the codespell:ignore +# directives on those lines are deliberate; do not "fix" them. +_QWEN_PROMPT_TEMPLATE = "\n".join([ + "<|im_start|>system\nYou are a promt engineer. Describe the video in detail.", # codespell:ignore promt + "Describe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.", + "Describe the location of the video, main characters or objects and their action.", + "Describe the dynamism of the video and presented actions.", + "Name the visual style of the video: whether it is a professional footage, user generated content, " + "some kind of animation, video game or scren content.", # codespell:ignore scren + "Describe the visual effects, postprocessing and transitions if they are presented in the video.", + "Pay attention to the order of key actions shown in the scene.<|im_end|>", + "<|im_start|>user\n{}<|im_end|>", +]) +_QWEN_ENCODE_START_IDX = 129 +_QWEN_MAX_SEQ_LEN = 256 +_CLIP_MAX_LEN = 77 +_PATCH_SIZE = (1, 2, 2) # transformer patch_size (T, H, W) +_SPATIAL_RATIO = 8 # VAE spatial compression (latent_h = H // 8) +_TIMESTEP_SCALE = 1000.0 # σ·1000 timestep convention + + +def _get_scale_factor(height: int, width: int) -> tuple[float, float, float]: + """Resolution-dependent RoPE scaling (pipeline ``_get_scale_factor``): (1,2,2) when BOTH height and + width are in [480, 854], else (1, 3.16, 3.16). BRINGUP risk E (must not be left at (1,1,1)).""" + + def between_480p(x: int) -> bool: + return 480 <= x <= 854 + + if between_480p(height) and between_480p(width): + return (1.0, 2.0, 2.0) + return (1.0, 3.16, 3.16) + + +class Kandinsky5DiT(TorchComponent): + """``dit(latent[C,T,H,W], qwen_embed[seq,3584], sigma, context=clip_pooled[768]) -> velocity[C,T,H,W]``. + + Real forward (kandinsky5.py): ``forward(hidden_states[B,T,H,W,C], encoder_hidden_states, timestep, + pooled_projections, visual_rope_pos, text_rope_pos, scale_factor, sparse_params, return_dict)``. The + loop's channels-FIRST latent is permuted to channels-LAST here and the velocity is permuted back.""" + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + hs = self._t(latent) # [1, C, T, H, W] (channels-first from the loop) + hs = hs.permute(0, 2, 3, 4, 1).contiguous() # -> [1, T, H, W, C] (the forward's geometry) + ehs = self._t(text_embed) # [1, seq, 3584] + if context is None: + raise ValueError("Kandinsky5DiT requires the CLIP pooled vector (passed as `context`); " + "pooled_projections is mandatory for the Kandinsky5 forward.") + pooled = self._t(context) # [1, 768] + _b, t, h, w, _c = hs.shape + ts = torch.tensor([float(sigma) * _TIMESTEP_SCALE], device=self.device, dtype=self.dtype) + # visual_rope_pos over the PATCHED grid: patch_size (1,2,2) -> (T, H//2, W//2). Matches the + # pipeline's [arange(num_latent_frames), arange(latent_h//2), arange(latent_w//2)]. + gh = max(1, h // _PATCH_SIZE[1]) + gw = max(1, w // _PATCH_SIZE[2]) + gt = max(1, t // _PATCH_SIZE[0]) + visual_rope_pos = [ + torch.arange(gt, device=self.device), + torch.arange(gh, device=self.device), + torch.arange(gw, device=self.device) + ] + text_rope_pos = torch.arange(ehs.shape[1], device=self.device) + # Recover pixel resolution from the latent geometry for the resolution-dependent scale_factor. + scale_factor = _get_scale_factor(h * _SPATIAL_RATIO, w * _SPATIAL_RATIO) + with self._ctx(): + velocity = self.module( + hidden_states=hs, + encoder_hidden_states=ehs, + timestep=ts, + pooled_projections=pooled, + visual_rope_pos=visual_rope_pos, + text_rope_pos=text_rope_pos, + scale_factor=scale_factor, + sparse_params=None, # Lite checkpoint = regular attention (BRINGUP K) + return_dict=False) + if isinstance(velocity, tuple): + velocity = velocity[0] + velocity = velocity.permute(0, 4, 1, 2, 3).contiguous() # [B,T,H,W,C] -> channels-first [B,C,T,H,W] + return self._n(velocity) # channels-FIRST rectified-flow velocity + + +class Kandinsky5QwenEncoder(TorchComponent): + """Qwen2.5-VL text encoder. ``encode(text) -> token embeds[seq, 3584]`` after the prompt template + + start-idx (129) slice on the LAST hidden state. The built-in text-encoder maker passes the + ``tokenizer`` subfolder tokenizer (the Qwen one) — correct here.""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = _QWEN_MAX_SEQ_LEN): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + @torch.no_grad() + def encode(self, text): + full = _QWEN_PROMPT_TEMPLATE.format(text or "") + max_allowed = _QWEN_ENCODE_START_IDX + self.max_length + toks = self.tokenizer(text=[full], + images=None, + videos=None, + max_length=max_allowed, + truncation=True, + return_tensors="pt", + padding=True) + ids = toks["input_ids"].to(self.device) + with self._ctx(): + out = self.module(input_ids=ids, return_dict=True, output_hidden_states=True) + # The pipeline uses hidden_states[-1] (the final hidden state) sliced from the start idx onward. + hidden = out["hidden_states"][-1] if isinstance(out, dict) else out.hidden_states[-1] + hidden = hidden[:, _QWEN_ENCODE_START_IDX:].squeeze(0) # [seq, 3584] + return _to_numpy(hidden) + + +class Kandinsky5ClipEncoder(TorchComponent): + """CLIP-with-projection text encoder. ``encode(text) -> pooled[768]``. The built-in maker hands this + adapter the wrong (``tokenizer/`` = Qwen) tokenizer, so we re-resolve the CLIP tokenizer from the + ``tokenizer_2`` sibling of the model root (BRINGUP: confirm the subfolder name on a real checkpoint).""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = _CLIP_MAX_LEN): + super().__init__(module, device=device, dtype=dtype) + self.max_length = max_length + self.tokenizer = self._resolve_clip_tokenizer(module, tokenizer) + + @staticmethod + def _resolve_clip_tokenizer(module: Any, fallback: Any) -> Any: + # The model root holds text_encoder_2/ (this module) and tokenizer_2/ (its tokenizer). Recover it + # from the module config's _name_or_path when available; else fall back to the passed tokenizer. + from transformers import AutoTokenizer + name = getattr(getattr(module, "config", None), "_name_or_path", "") or "" + if name: + root = os.path.dirname(os.path.normpath(name)) # .../text_encoder_2 -> root + tok2 = os.path.join(root, "tokenizer_2") + if os.path.isdir(tok2): + return AutoTokenizer.from_pretrained(tok2) + return fallback + + @torch.no_grad() + def encode(self, text): + toks = self.tokenizer(text or "", + max_length=self.max_length, + truncation=True, + add_special_tokens=True, + padding="max_length", + return_tensors="pt") + ids = toks["input_ids"].to(self.device) + mask = toks.get("attention_mask") + kw = {"attention_mask": mask.to(self.device)} if mask is not None else {} + with self._ctx(): + out = self.module(input_ids=ids, **kw) + pooled = out["pooler_output"] if isinstance(out, dict) else out.pooler_output + return _to_numpy(pooled.squeeze(0)) # [768] + + +class Kandinsky5VAE(TorchComponent): + """AutoencoderKLHunyuanVideo with a SCALAR scaling_factor. The DiT works in scaled latent space + (latent = raw·sf); decode inverts it (raw = latent/sf) — NOT the Wan per-channel (z-mean)/std. The + loop latent is already channels-first ``[C,T,H,W]`` (the DiT adapter owns the channels-last view).""" + + def _scaling_factor(self) -> float: + cfg = getattr(self.module, "config", None) + sf = getattr(cfg, "scaling_factor", None) if cfg is not None else None + return float(sf) if sf is not None else 0.476986 + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent).float() # [1, C, T, H, W] + z = z / self._scaling_factor() # un-normalize (scalar) + video = self.module.decode(z.to(self.dtype)) + if hasattr(video, "sample"): + video = video.sample + return self._n(video) # [3, T, H*8, W*8] in [-1, 1] + + @torch.no_grad() + def encode(self, video): + x = self._t(video) + dist = self.module.encode(x) + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + return self._n(z.float() * self._scaling_factor()) # raw -> scaled latent the DiT consumes (channels-first) diff --git a/v2/platform/backends/torch_lingbotworld.py b/v2/platform/backends/torch_lingbotworld.py new file mode 100644 index 000000000..fd435f22a --- /dev/null +++ b/v2/platform/backends/torch_lingbotworld.py @@ -0,0 +1,78 @@ +"""LingBot-World torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so +the recipe is self-contained (no edit to the shared ``_make_dit`` dispatch in ``torch_backend.py``). +Imported lazily by ``_explicit_adapter`` only on a GPU box. + +LingBot-World-Base-Cam is Wan2.2-I2V-A14B (a 2x14B boundary-routed MoE with CLIP + first-frame i2v +conditioning) PLUS a NEW camera/Plucker conditioning input. The DiT forward gains a ``c2ws_plucker_emb`` +argument (``LingBotWorldTransformer3DModel.forward(..., c2ws_plucker_emb=None)``) which a vanilla WanDiT +reuse would silently drop -> the generation would degrade to a generic Wan i2v with no camera control +(BRINGUP blocker #1). ``LingBotWorldDiT`` threads that tensor. + +Key design constraint (faithful + conflict-free): the v2 dit-call surface (``ToyDiT.__call__`` and the +loop's ``dit(latent, text_embed, sigma, context=, cond=)``) has NO ``c2ws_plucker_emb`` kwarg, and we may +not edit the toy/loop sampler. So the loop publishes the per-step Plucker tensor onto the adapter as a +plain attribute (``self.c2ws_plucker_emb``) right before the call (the cosmos2 pattern of building +arch-specifics adapter-side); ``__call__`` reads it and forwards it to the real module. The ToyDiT stand-in +simply never has the attribute, so the CPU path is the no-camera (degenerate t2v/i2v) path unchanged. + +MoE CPU-offload (BRINGUP blocker #3): 2x14B bf16 won't co-reside on an 80GB GPU. The card declares this +adapter for BOTH ``transformer`` (high-noise) and ``transformer_2`` (low-noise) experts; we keep the +WanDiT offload-group swap so only the boundary-active expert is GPU-resident. ``_make_dit`` does not pass +an ``offload_group`` to an explicit adapter, so the two LingBotWorld experts share a process-wide group +keyed by the wrapped module's ``id`` would be fragile across instances — instead we share a single +class-level group; for the registered single-instance serving path that is exactly the WanDiT MoE +behavior. BRINGUP: revisit if multiple LingBotWorld instances must co-reside. +""" +from __future__ import annotations + +import torch + +from v2.platform.backends.torch_backend import NUM_TRAIN_TIMESTEPS, WanDiT + + +class LingBotWorldDiT(WanDiT): + """``dit(latent[C,T,H,W], text_embed[seq,dim], sigma, context=clip, cond=[mask|cond]) -> velocity``, + plus an out-of-band ``self.c2ws_plucker_emb`` (numpy ``[6*s^2, F_lat, H_lat, W_lat]``) the loop sets + per step. Identical i2v conditioning to Wan2.2-I2V-A14B — the 36ch ``[noise|mask+cond]`` concat and + CLIP ``encoder_hidden_states_image`` — with the extra Plucker camera tensor forwarded to the model. + + Faithful to ``LingBotWorldImageToVideoPipeline`` (== ``WanImageToVideoPipeline`` denoise) + the + ``c2ws_plucker_emb=batch.c2ws_plucker_emb`` kwarg threaded by ``DenoisingStage`` (denoising.py:182). + """ + + # All LingBotWorld DiT experts share one CPU-offload group (only the boundary-active one is resident). + _LBW_OFFLOAD_GROUP: dict[str, LingBotWorldDiT] = {} + + def __init__(self, module, *, device, dtype): + # ``_make_dit`` builds an explicit adapter as ``cls(module, device=, dtype=)`` (no offload_group), + # so wire the shared MoE offload group here ourselves (2x14B can't co-reside on 80GB). + component_id = getattr(getattr(module, "config", None), "prefix", None) or type(module).__name__ + super().__init__(module, + device=device, + dtype=dtype, + offload_group=LingBotWorldDiT._LBW_OFFLOAD_GROUP, + component_id=str(component_id)) + # The per-step camera/Plucker embedding the loop publishes before each call (None == no camera). + self.c2ws_plucker_emb = None + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + self._ensure_resident() + hs = self._t(latent) + if cond is not None: # i2v: concat [noise (16ch) ; mask+cond_latent (20ch)] -> 36ch input + hs = torch.cat([hs, self._t(cond)], dim=1) + ehs = self._t(text_embed) + # LingBotWorld uses the SCALAR-timestep path (no expand_timesteps); timestep = sigma*1000. + ts = float(sigma) * NUM_TRAIN_TIMESTEPS + timestep = torch.tensor([ts], device=self.device) + img = self._t(context) # CLIP image embeds -> encoder_hidden_states_image + plucker = None + if self.c2ws_plucker_emb is not None: # camera conditioning (BRINGUP: needs a per-request input) + plucker = self._t(self.c2ws_plucker_emb) # [1, 6*s^2, F_lat, H_lat, W_lat] + with self._ctx(): + velocity = self.module(hidden_states=hs, + encoder_hidden_states=ehs, + timestep=timestep, + encoder_hidden_states_image=img, + c2ws_plucker_emb=plucker) + return self._n(velocity) # rectified-flow velocity (BRINGUP risk C) diff --git a/v2/platform/backends/torch_longcat.py b/v2/platform/backends/torch_longcat.py new file mode 100644 index 000000000..a8229ed47 --- /dev/null +++ b/v2/platform/backends/torch_longcat.py @@ -0,0 +1,63 @@ +"""LongCat-Video torch adapter (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so +the LongCat recipe is self-contained (no edit to the shared ``_make_dit`` dispatch in ``torch_backend.py``). +Imported lazily by ``_explicit_adapter`` only on a GPU box. + +* ``LongCatDiT`` — the flow-match velocity predictor. Mirrors ``WanDiT`` (``timestep = sigma*1000``, + scalar ``[B]`` timestep that the DiT internally expands ``[B] -> [B, T_latent]`` for the per-latent-frame + AdaLN time embedding) but with the **critical LongCat sign convention**: the fastvideo + ``LongCatDenoisingStage`` does ``noise_pred = -noise_pred`` *before* the flow-match scheduler step + (longcat_denoising.py line 152). The v2 loop integrates ``x + (sigma_next-sigma)*velocity`` with the + velocity this adapter returns, so to reproduce that step EXACTLY this adapter returns ``-velocity``. + Forgetting the negation silently produces diverging samples (spec blocker #1). + +LongCat reuses the existing ``WanVAE`` torch adapter unchanged (same ``AutoencoderKLWan`` mean/std +normalization) and the existing ``T5Encoder`` (UMT5, zero-padded to max_length=512 for the CFG concat +uniform-seq contract). Both are wired on the card by ``load_id`` alone — no adapter override needed. + +The DiT runs its AdaLN modulation / residual gating + final projection under fp32 and casts the output to +float32 (longcat.py line 1121); we keep the module at its native bf16 dtype (``_native_dtype`` in +``_make_dit``, matching the fastvideo stage's hardcoded bf16 target) and let ``_n`` marshal the fp32 output +back to numpy (spec blocker #6). + +BRINGUP: written-not-run on CPU (no torch/weights here). GPU-verify against a real +``FastVideo/LongCat-Video-T2V-Diffusers`` checkpoint — confirm the bf16 dtype, the scalar-timestep +expansion, and ``encoder_attention_mask=None`` (the CaptionEmbedder simply skips the zeroing branch for +T2V, which is faithful since the v2 ``T5Encoder`` already zero-pads the embedding rows past real tokens). +""" +from __future__ import annotations + +import torch + +from v2.platform.backends.torch_backend import NUM_TRAIN_TIMESTEPS, TorchComponent + + +class LongCatDiT(TorchComponent): + """``dit(latent[C,T,H,W], text_embed[seq,4096], sigma) -> -velocity[C,T,H,W]``. + + Real forward (``fastvideo/models/dits/longcat.py``): + ``forward(hidden_states[B,C=16,T,H,W], encoder_hidden_states[B,N_text,4096], timestep[B]|[B,T], + encoder_attention_mask=None, ...) -> velocity[B,C_out=16,T,H,W]`` (cast to float32). + + The loop hands the raw ``sigma`` (1->0); this adapter forms ``timestep = sigma*1000`` (the Wan/diffusers + convention LongCat shares) as a scalar ``[B]`` tensor — the DiT's ``forward`` expands ``[B] -> [B, T_latent]`` + internally for the AdaLN time embedding (longcat.py line 1075). It returns the **negated** velocity to + fold in the fastvideo stage's ``noise_pred = -noise_pred`` (spec blocker #1).""" + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + hs = self._t(latent) # [1, C=16, T, H, W] + ehs = self._t(text_embed) # [1, N_text, 4096] + b = hs.shape[0] + # timestep = sigma*1000 (BRINGUP risk B / spec timestep_convention). Scalar [B]; the DiT expands + # [B] -> [B, T_latent] internally for the per-latent-frame AdaLN time embedding. (I2V/VC would build + # an explicit [B, T_latent] with cond frames zeroed — deferred, T2V uses the scalar path.) + ts = float(sigma) * NUM_TRAIN_TIMESTEPS + timestep = torch.full((b, ), ts, device=self.device, dtype=self.dtype) + with self._ctx(current_timestep=ts): + velocity = self.module(hidden_states=hs, + encoder_hidden_states=ehs, + timestep=timestep, + encoder_attention_mask=None) + # CRITICAL (spec blocker #1): negate so the loop's FLOW_MATCH_STEP reproduces the fastvideo stage's + # negate-then-scheduler-step. The DiT already casts to float32; _n marshals fp32 -> numpy. + return self._n(-velocity) diff --git a/v2/platform/backends/torch_matrixgame2.py b/v2/platform/backends/torch_matrixgame2.py new file mode 100644 index 000000000..5429f5a70 --- /dev/null +++ b/v2/platform/backends/torch_matrixgame2.py @@ -0,0 +1,145 @@ +"""Matrix-Game 2.0 torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` so +the recipe is self-contained (no edit to ``_make_dit``/``_make_image_encoder`` in ``torch_backend.py``). +Imported lazily by ``_explicit_adapter`` only on a GPU box. + +Matrix-Game 2.0 is an INTERACTIVE world model: a distilled, causal/autoregressive, action-conditioned +Wan-family DiT. The deltas this adapter encodes vs the stock ``WanDiT`` (all faithful to +``fastvideo/models/dits/matrixgame2/causal_model.py:CausalMatrixGame2WanModel._forward_inference`` and +``fastvideo/pipelines/stages/matrixgame2_denoising.py:MatrixGame2CausalDenoisingStage``): + + * EPSILON output, NOT velocity — ``_forward_inference`` returns a bare velocity-shaped tensor the stage + treats as predicted NOISE; the epsilon->x0 conversion (``pred_noise_to_pred_video`` = + ``x - sigma_t*eps``) lives in ``MatrixGame2CausalDMDLoop``, NOT here. The adapter returns the raw model + output. + * CAUSAL KV-CACHE is mandatory — ``kwargs['kv_cache'] is not None`` selects ``_forward_inference`` (else + ``_forward_train``). The DiT MUTATES kv_cache / crossattn_cache / kv_cache_mouse / kv_cache_keyboard in + place (sliding-window eviction by ``local_attn_size``). The loop owns these cache lists across blocks + and steps and hands them in via ``MatrixGame2CausalDiT.call(...)``; the loop keeps them in + ``LoopState`` so interleaving stays safe. + * PER-FRAME timestep ``[B, num_frames]`` (LongTensor 0..1000), NOT a scalar ``sigma*1000`` — the + condition_embedder unflattens ``timestep_proj`` by (batch, post_patch_num_frames); a scalar breaks the + reshape. The loop passes a numpy per-frame timestep array; this adapter keeps it 2-D and ``long``. + * i2v via CHANNEL-CONCAT read off ``get_forward_context().forward_batch.image_latent`` — the model + concats the conditioning latent onto ``hidden_states`` along dim=1 INTERNALLY (patch_embedding is the + 32-in checkpoint). So the adapter places ``image_latent`` (the raw cond_latent, NO 4ch Wan mask — this + differs from Wan2.1 i2v) on a lightweight forward_batch under ``set_forward_context``. + * CLIP image embeds (257x1280) are the SOLE cross-attention context (``encoder_hidden_states_image``); + text is ignored (Matrix-Game 2.0 has no text encoder), but the forward still wants a placeholder + ``encoder_hidden_states`` list. + * ACTION conditioning (mouse[B,F,2] / keyboard[B,F,4]) + its own KV caches is the interactive surface. + BRINGUP: the degenerate (no-action) path runs with mouse/keyboard=None; live action routing needs a + request-API extension (see the recipe's loop docstring). + +BRINGUP (written-not-run; GPU-only): importing ``causal_model`` triggers a flex_attention +``torch.compile`` and the inference path uses plain SDPA over the cache — both GPU-only. The exact cache +shapes (``local_attn_size`` window, mouse-cache batch dim ``B*frame_seq_len``) come from the loaded +arch_config and must be GPU-verified against a real distilled checkpoint. +""" +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch + +from v2.platform.backends.torch_backend import TorchComponent, _to_numpy + +# Matrix-Game 2.0 ships a Wan-VAE first-frame i2v conditioning latent that the DiT channel-concats; this is +# carried on a minimal forward_batch object the model reads via ``get_forward_context().forward_batch``. + + +class _CondBatch: + """Minimal stand-in for fastvideo's ``ForwardBatch`` carrying only the field the causal DiT reads: + ``image_latent`` (the i2v cond_concat). The real ForwardBatch has dozens of fields; the inference + forward only touches ``getattr(batch, 'image_latent', None)``, so this is a faithful, tiny carrier.""" + + def __init__(self, image_latent: Any) -> None: + self.image_latent = image_latent + + +class MatrixGame2CausalDiT(TorchComponent): + """``call(latent[C,T,h,w], image_embed[257,1280], timestep[T] (per-frame), kv_caches..., cond=...)`` + -> raw epsilon[C,T,h,w]``. + + The loop drives the causal block-autoregressive schedule and the DMD epsilon->x0 conversion; this + adapter is a single ``_forward_inference`` call with the cache/i2v/action plumbing. It is NOT the toy + ``dit(latent, text_embed, sigma)`` shape — the loop calls ``MatrixGame2CausalDiT.call`` explicitly with + the causal kwargs (the toy backend exercises the SAME loop via its degenerate path; see the recipe + loop). All arch-specifics are built INTERNALLY here so the loop stays backend-agnostic.""" + + @torch.no_grad() + def call(self, + latent: np.ndarray, + image_embed: Any, + timestep: np.ndarray, + *, + kv_cache: list[dict] | None = None, + crossattn_cache: list[dict] | None = None, + kv_cache_mouse: list[dict] | None = None, + kv_cache_keyboard: list[dict] | None = None, + current_start: int = 0, + start_frame: int = 0, + num_frame_per_block: int = 1, + mouse_cond: np.ndarray | None = None, + keyboard_cond: np.ndarray | None = None, + image_latent: np.ndarray | None = None, + prompt_embeds: Any = None) -> np.ndarray: + # hidden_states [B, C, T, h, w]; the DiT concats image_latent internally off the forward_batch. + hs = self._t(latent) # [1, C, T, h, w] + b, _c, t, _h, _w = hs.shape + # Per-frame integer timestep [B, num_frames] (long); keep 2-D — never collapse to a scalar. + ts_np = np.asarray(timestep).reshape(-1) + timestep_t = torch.as_tensor(ts_np, device=self.device, dtype=torch.long).reshape(1, -1).expand(b, t) + # CLIP image embeds [257, 1280] are the SOLE cross-attn context; the model wants a list. + img_ctx = self._t(image_embed) if image_embed is not None else None + image_kwargs = {"encoder_hidden_states_image": [img_ctx] if img_ctx is not None else []} + # Text is ignored but the forward wants a non-empty placeholder list (condition_embedder drops it). + pe = self._t(prompt_embeds) if prompt_embeds is not None else None + ehs: list[Any] = [pe] if pe is not None else [] + # Action conditioning (interactive surface) — None on the degenerate world-rollout path (BRINGUP). + action_kwargs: dict[str, Any] = {"num_frame_per_block": int(num_frame_per_block)} + if mouse_cond is not None: + action_kwargs["mouse_cond"] = self._t(mouse_cond) + if keyboard_cond is not None: + action_kwargs["keyboard_cond"] = self._t(keyboard_cond) + if kv_cache_mouse is not None: + action_kwargs["kv_cache_mouse"] = kv_cache_mouse + if kv_cache_keyboard is not None: + action_kwargs["kv_cache_keyboard"] = kv_cache_keyboard + # i2v cond_concat: the DiT reads image_latent off the forward_batch -> channel-concat dim=1. + cond_lat = self._t(image_latent) if image_latent is not None else None + from v2.forward_context import set_forward_context + with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=_CondBatch(cond_lat)): + eps = self.module( + hs, + ehs, + timestep_t, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=int(current_start), + start_frame=int(start_frame), + **image_kwargs, + **action_kwargs, + ) + return self._n(eps) # RAW epsilon (the loop converts eps -> x0 via the scheduler sigma table) + + +class MatrixGame2CLIPImageEncoder(TorchComponent): + """CLIP vision encoder for Matrix-Game 2.0 i2v conditioning: ``encode_image(image) -> [257, 1280]`` (1 + CLS + 256 patch tokens), which the DiT's WanImageEmbedding projects to inner_dim and uses as the sole + cross-attention context. Mirrors fastvideo's ``MatrixGame2ImageEncodingStage`` — the HF image processor + (sibling ``image_processor`` subfolder) preprocesses, the encoder returns ``last_hidden_state``. + BRINGUP: written-not-run; GPU-verify the processor/encoder subfolders + dtype against a real + distilled checkpoint.""" + + def __init__(self, module, processor, *, device, dtype): + super().__init__(module, device=device, dtype=dtype) + self.processor = processor + + @torch.no_grad() + def encode_image(self, image): + inputs = self.processor(images=image, return_tensors="pt").to(self.device) + with self._ctx(): + out = self.module(**inputs) + embeds = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] + return _to_numpy(embeds.squeeze(0)) diff --git a/v2/platform/backends/torch_matrixgame3.py b/v2/platform/backends/torch_matrixgame3.py new file mode 100644 index 000000000..12c0e0464 --- /dev/null +++ b/v2/platform/backends/torch_matrixgame3.py @@ -0,0 +1,112 @@ +"""Matrix-Game 3.0 torch adapter (GPU backend) — declared on the card via ``ComponentSpec.adapter`` +so the MG3 recipe is self-contained (no edit to the shared ``_make_dit`` dispatch in +``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +``MatrixGame3DiT`` is the autoregressive world-model DiT (``MatrixGame3WanModel``, 30 layers, dim 3072, +in/out_channels=48, patch (1,2,2), ``use_memory=True``, ``sigma_theta=0.8``). It departs from ``WanDiT`` +in three load-bearing ways (BRINGUP risks 1/2/7 in the spec): + + 1. PER-TOKEN timestep, NOT ``sigma*1000`` scalar. The model's ``timestep`` arg is a flattened + ``[1, T*(H/2)*(W/2)]`` LongTensor whose value is the raw FlowUniPC scheduler timestep, with the + ``cond_frames`` rows zeroed (those latent frames are pinned to the conditioning image and must not be + denoised). When ``x_memory`` is present, an all-zeros ``timestep_memory`` block is concatenated in + front. The loop builds the *scalar* timestep + ``cond_frames`` + per-clip conditioning bundle and + hands them through ``cond=``; this adapter expands them to the per-token tensor INTERNALLY so the + loop's dit-call stays ``dit(latent, text_embed, sigma)`` (ToyDiT-compatible — the cosmos2 pattern). + 2. Action / camera / memory conditioning: ``mouse_cond`` / ``keyboard_cond`` / ``x_memory`` / + ``timestep_memory`` / ``mouse_cond_memory`` / ``keyboard_cond_memory`` / ``c2ws_plucker_emb`` / + ``memory_latent_idx`` / ``predict_latent_idx`` are passed straight through. The loop assembles all of + them (numpy host work — Plücker / extrinsics / memory-index selection) and stuffs them in ``cond``. + 3. The forward returns a velocity (``flow_prediction``) — a bare tensor ``[B,48,T,H,W]`` — which the + FlowUniPC scheduler (driven by the loop) converts velocity->x0 internally in ``step()``. + +The forward runs inside ``set_forward_context(current_timestep=, attn_metadata=None)`` (the +FastVideo attention layers read ``attn_metadata`` via the forward context). The 5B DiT is kept eager and +resident (no offload group). BRINGUP: written-not-run; the per-token timestep packing, the +``predict_latent_idx`` ``(start,end)`` tuple (the model expands it to a device ``arange``), and the +memory/plucker tensor dtypes must be GPU-verified against a real ``Matrix-Game-3.0`` checkpoint. +""" +from __future__ import annotations + +from typing import Any + +import torch + +from v2.platform.backends.torch_backend import TorchComponent + + +class MatrixGame3DiT(TorchComponent): + """``dit(latent[C=48,T,H,W], text_embed[seq,dim], sigma, cond=) -> velocity[C,T,H,W]``. + + ``cond`` is the per-step conditioning bundle the ``MatrixGame3DenoiseLoop`` assembles (a plain dict). + Keys (all optional except ``cond_frames``): + * ``cond_frames`` (int) — leading latent frames pinned to the conditioning image (1 first clip, + 4 later); their per-token timestep rows are zeroed. + * ``mouse_cond`` / ``keyboard_cond`` (numpy [1, frames, dim]) — per-clip action streams. + * ``x_memory`` / ``timestep_memory`` / ``mouse_cond_memory`` / ``keyboard_cond_memory`` (numpy) — + autoregressive KV-memory tensors (None on the first clip / degenerate single-clip path). + * ``c2ws_plucker_emb`` (numpy) — camera Plücker embeddings (memory blocks concatenated in front). + * ``memory_latent_idx`` (list[int] | None) — history indices the memory tensors index into. + * ``predict_latent_idx`` (tuple (start, end)) — the clip's latent-frame span (the model expands it + to a device ``arange``); REQUIRED, the model uses it for rotary positions. + """ + + def __init__(self, module: Any, *, device: Any, dtype: Any) -> None: + super().__init__(module, device=device, dtype=dtype) + # patch_size[1:] -> (H/2, W/2) token folding for the per-token timestep packing. (1,2,2) default. + ps = getattr(module, "patch_size", (1, 2, 2)) + self.patch_t, self.patch_h, self.patch_w = int(ps[0]), int(ps[1]), int(ps[2]) + + def _per_token_timestep(self, hs: torch.Tensor, sigma_value: float, cond_frames: int, + x_memory: torch.Tensor | None) -> torch.Tensor: + """Build the flattened per-token timestep [1, T*(H/p)*(W/p)] (BRINGUP risk 1). The raw FlowUniPC + timestep value fills every token; ``cond_frames`` latent rows are zeroed (pinned, not denoised); + an all-zeros memory block is prepended when ``x_memory`` is present.""" + _b, _c, t, h, w = hs.shape + tokens_per_frame = (h // self.patch_h) * (w // self.patch_w) + tt = hs.new_full((t, tokens_per_frame), float(sigma_value)) + if cond_frames > 0: + tt[:cond_frames].zero_() + tt = tt.flatten().unsqueeze(0) + if x_memory is not None: + mem_tokens = x_memory.shape[2] * (x_memory.shape[3] // self.patch_h) * (x_memory.shape[4] // self.patch_w) + tt = torch.cat([tt.new_zeros((1, mem_tokens)), tt], dim=1) + return tt + + @torch.no_grad() + def __call__(self, + latent: Any, + text_embed: Any, + sigma: float, + context: Any = None, + *, + cond: dict[str, Any] | None = None) -> Any: + hs = self._t(latent) # [1, 48, T, H, W] + ehs = self._t(text_embed) + cond = cond or {} + cond_frames = int(cond.get("cond_frames", 0)) + x_memory = self._t(cond.get("x_memory")) if cond.get("x_memory") is not None else None + # FlowUniPC timestep value: the loop passes the *raw* scheduler timestep (~sigma*1000), already + # resolved, as ``sigma`` here (NOT a normalized 0..1 sigma). See the loop for the convention. + timestep = self._per_token_timestep(hs, float(sigma), cond_frames, x_memory) + + def _opt(key: str) -> Any: + v = cond.get(key) + return self._t(v) if v is not None else None + + with self._ctx(current_timestep=float(sigma)): + velocity = self.module( + hs, + ehs, + timestep, + mouse_cond=_opt("mouse_cond"), + keyboard_cond=_opt("keyboard_cond"), + x_memory=x_memory, + timestep_memory=_opt("timestep_memory"), + mouse_cond_memory=_opt("mouse_cond_memory"), + keyboard_cond_memory=_opt("keyboard_cond_memory"), + c2ws_plucker_emb=_opt("c2ws_plucker_emb"), + memory_latent_idx=cond.get("memory_latent_idx"), + predict_latent_idx=cond.get("predict_latent_idx"), + ) + return self._n(velocity) # rectified-flow velocity (FlowUniPC.step does velocity->x0) diff --git a/v2/platform/backends/torch_sd35.py b/v2/platform/backends/torch_sd35.py new file mode 100644 index 000000000..d99ee5729 --- /dev/null +++ b/v2/platform/backends/torch_sd35.py @@ -0,0 +1,193 @@ +"""Stable Diffusion 3.5 torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` +so the SD3.5 recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder`` +dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +Architecture deltas vs Wan/Cosmos (all faithful to ``fastvideo/pipelines/stages/sd35_conditioning.py`` ++ ``fastvideo/models/dits/sd3.py:SD3Transformer2DModel``): + +* ``SD3DiT`` — a FLOW-MATCH MMDiT. The forward needs TWO text conditioners: ``encoder_hidden_states`` + (the assembled triple-encoder joint embed [seq, 4096]) AND a separate ``pooled_projections`` vector + [2048]. The v2 loop threads only a single ``text_embed`` + a positional ``context`` arg (the same + surface ToyDiT exposes), so this adapter takes ``text_embed`` = the joint embed and ``context`` = the + pooled vector. Latents are 4D image latents [1, 16, h, w] (num_frames=1) — NOT 5D video. The model + output is the flow-match velocity directly (no x0 conversion, unlike LTX2DiT). ``timestep = sigma*1000`` + (diffusers FlowMatchEuler convention, same as WanDiT). The fastvideo forward self-manages + ``set_forward_context`` (sd3.py lines 1011-1021) but wrapping in ``self._ctx()`` is idempotent (it + detects the existing context), so we wrap for uniformity with the other adapters. + +* ``SD3VAE`` — ``AutoencoderKL`` operates in RAW latent space; the DiT operates in a shift/scale + NORMALIZED space. SD3.5: ``z_dit = (z - shift_factor) * scaling_factor`` on encode; decode inverts + ``z_raw = z_dit / scaling_factor + shift_factor`` then ``vae.decode -> image in [-1,1]``. CRITICAL: + the factors come from the LOADED checkpoint ``vae.config`` (scaling_factor=1.5305, shift_factor=0.0609 + for SD3.5), NOT the generic AutoencoderKLArchConfig defaults (0.18215 / None). 4D (image), not 5D. + The decoded image is mapped to [0,1] by the program (``(img/2+0.5).clamp(0,1)``), matching + ``SD35DecodingStage``. + +* ``SD3ClipEncoder`` / ``SD3T5Encoder`` — the joint embed is assembled in the program (see + ``sd35.program``), so each encoder adapter returns BOTH the per-encoder pieces the assembler needs: + CLIP returns ``(penultimate_hidden[-2], pooler_output)`` and T5 returns ``last_hidden_state``. + CLIP runs with ``output_hidden_states=True`` at max_length=77 padded; T5 at max_length=256 padded. + Faithful to ``SD35ConditioningStage._clip_pooled`` + the TextEncodingStage CLIP penultimate extraction. +""" +from __future__ import annotations + +import numpy as np +import torch + +from v2.platform.backends.torch_backend import NUM_TRAIN_TIMESTEPS, TorchComponent, _to_numpy + +# SD3.5 encoder lengths (SD35Pipeline.initialize_pipeline / SD35ConditioningStage). +CLIP_MAX_LENGTH = 77 +T5_MAX_LENGTH = 256 +# Generic SD1.5-era AutoencoderKL defaults — used ONLY as a last-resort fallback if the checkpoint +# config is missing the SD3.5 values (BRINGUP: the real checkpoint vae config carries 1.5305 / 0.0609). +_DEFAULT_SCALING = 1.5305 +_DEFAULT_SHIFT = 0.0609 + + +class SD3DiT(TorchComponent): + """``dit(latent[16,h,w], encoder_hidden_states[seq,4096], sigma, context=pooled[2048]) -> velocity``. + + Faithful to ``SD35DenoisingStage`` + ``SD3Transformer2DModel.forward``: 4D image latents, the dual + text conditioning (joint embed + pooled), ``timestep = sigma*1000`` broadcast to [B], flow-match + velocity output. ``context`` carries the pooled_projections vector the loop threads alongside the + joint embed (so the loop's dit-call stays the cosmos2-style ``dit(latent, text_embed, sigma, context)``).""" + + @torch.no_grad() + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + # 4D image latent: ``self._t`` prepends the batch dim -> [1, 16, h, w]. (No temporal dim.) + hs = self._t(latent) + ehs = self._t(text_embed) # [1, seq, 4096] joint embed + if context is None: + raise RuntimeError("SD3DiT requires pooled_projections (threaded as `context`); got None.") + pp = self._t(context) # [1, 2048] dual-CLIP pooled vector + # diffusers FlowMatchEuler timestep = sigma * num_train_timesteps (1000); 1D [B] long-ish float. + ts = torch.tensor([float(sigma) * NUM_TRAIN_TIMESTEPS], device=self.device, dtype=self.dtype) + with self._ctx(): # idempotent (sd3.forward self-manages it) + out = self.module(hidden_states=hs, + encoder_hidden_states=ehs, + pooled_projections=pp, + timestep=ts, + return_dict=False)[0] # [1, 16, h, w] flow-match velocity + return self._n(out) # -> [16, h, w] numpy (no x0 conversion) + + +class SD3VAE(TorchComponent): + """``AutoencoderKL`` with the SD3.5 shift/scale normalization the DiT latent space requires. + + encode: ``z_dit = (z - shift_factor) * scaling_factor``; decode: ``z_raw = z_dit / scaling_factor + + shift_factor`` then ``vae.decode -> image[3,h,w] in [-1,1]``. The factors are read from the loaded + ``module.config`` (BRINGUP risk D: the generic dataclass defaults would mis-scale). Marshals 4D + (image), not 5D video. The [-1,1] -> [0,1] remap is done by the program, matching ``SD35DecodingStage``.""" + + @staticmethod + def _as_float(v: object) -> float: + return float(v.item()) if isinstance(v, torch.Tensor) else float(v) # type: ignore[arg-type] + + def _factors(self) -> tuple[float, float | None]: + cfg = getattr(self.module, "config", None) + sf = getattr(cfg, "scaling_factor", None) if cfg is not None else None + sh = getattr(cfg, "shift_factor", None) if cfg is not None else None + if sf is None: + sf = getattr(self.module, "scaling_factor", _DEFAULT_SCALING) + if sh is None: + sh = getattr(self.module, "shift_factor", _DEFAULT_SHIFT) + return self._as_float(sf), (None if sh is None else self._as_float(sh)) + + @torch.no_grad() + def encode(self, image): + # image [3, h, w] -> [1, 3, h, w] (4D); AutoencoderKL.encode -> a DiagonalGaussian dist. + x = self._t(image) + dist = self.module.encode(x) + dist = dist.latent_dist if hasattr(dist, "latent_dist") else dist + z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist) + sf, sh = self._factors() + z = z.float() + if sh is not None: + z = z - sh + return self._n(z * sf) # -> the normalized latent the DiT expects + + @torch.no_grad() + def decode(self, latent): + z = self._t(latent).float() # [1, 16, h, w] + sf, sh = self._factors() + z = z / sf + if sh is not None: + z = z + sh + dec = self.module.decode(z.to(self.dtype)) # -> image [1, 3, H, W] in [-1, 1] + image = dec.sample if hasattr(dec, "sample") else (dec[0] if isinstance(dec, tuple) else dec) + return self._n(image) # [3, H, W]; program remaps [-1,1]->[0,1] + + +class SD3ClipEncoder(TorchComponent): + """``CLIPTextModelWithProjection`` (clip_l or clip_g). ``encode(text) -> {hidden, pooled}`` where + ``hidden`` is the penultimate hidden state ([-2], the SD3 joint-embed source — NOT last_hidden_state) + and ``pooled`` is ``pooler_output`` (the text-projection, the pooled-projection source). max_length=77 + padded. Faithful to ``SD35ConditioningStage._clip_pooled`` + the TextEncodingStage penultimate extraction.""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = CLIP_MAX_LENGTH): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + @torch.no_grad() + def encode(self, text): + toks = self.tokenizer(text or "", + return_tensors="pt", + max_length=self.max_length, + padding="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, output_hidden_states=True) + # ``hidden_states`` is the per-layer pool; [-2] is the penultimate layer (the SD3 joint-embed source). + hs = out.hidden_states + penultimate = hs[-2] if isinstance(hs, list | tuple) else hs + penultimate = penultimate.squeeze(0) # [seq=77, dim] + pooled = out.pooler_output.squeeze(0) # [dim] (text_projection of the EOS token) + return {"hidden": _to_numpy(penultimate), "pooled": _to_numpy(pooled)} + + +class SD3T5Encoder(TorchComponent): + """``T5EncoderModel`` (the third SD3.5 encoder). ``encode(text) -> last_hidden_state[seq=256, 4096]`` + (max_length=256 padded). NaN->0 for safety. Faithful to the T5 branch of ``SD35ConditioningStage``.""" + + def __init__(self, module, tokenizer, *, device, dtype, max_length: int = T5_MAX_LENGTH): + super().__init__(module, device=device, dtype=dtype) + self.tokenizer = tokenizer + self.max_length = max_length + + @torch.no_grad() + def encode(self, text): + toks = self.tokenizer(text or "", + return_tensors="pt", + max_length=self.max_length, + padding="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) + return _to_numpy(torch.nan_to_num(hidden, nan=0.0)) # [256, 4096] + + +def assemble_sd3_conditioning(clip_l: dict, clip_g: dict, t5: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Assemble the SD3.5 dual-text conditioning from the three encoder outputs — the EXACT + ``SD35ConditioningStage`` algebra, kept as a pure-numpy helper so both the GPU and CPU-toy program + paths share one assembler (BRINGUP risk: the pad/concat order + hidden index are easy to get wrong): + + * joint embed: ``cat([ F.pad(cat([clip_l_h, clip_g_h], dim=-1), (0, 4096-2048)) , t5_h ], dim=seq)`` + * pooled: ``cat([clip_l_pool, clip_g_pool], dim=-1)`` -> [2048] + """ + clip_h = np.concatenate([np.asarray(clip_l["hidden"]), np.asarray(clip_g["hidden"])], axis=-1) # [77, 2048] + t5_h = np.asarray(t5) # [256, 4096] + pad = t5_h.shape[-1] - clip_h.shape[-1] # 4096 - 2048 = 2048 + if pad < 0: + raise ValueError(f"CLIP joint dim {clip_h.shape[-1]} exceeds T5 dim {t5_h.shape[-1]}") + clip_h = np.pad(clip_h, [(0, 0)] * (clip_h.ndim - 1) + [(0, pad)]) # [77, 4096] + joint = np.concatenate([clip_h, t5_h], axis=-2).astype("float32") # [77+256, 4096] along seq + pooled = np.concatenate([np.asarray(clip_l["pooled"]), np.asarray(clip_g["pooled"])], + axis=-1).astype("float32") # [2048] + return joint, pooled diff --git a/v2/platform/backends/torch_stable_audio.py b/v2/platform/backends/torch_stable_audio.py new file mode 100644 index 000000000..f7e382e54 --- /dev/null +++ b/v2/platform/backends/torch_stable_audio.py @@ -0,0 +1,149 @@ +"""Stable Audio Open torch adapters (GPU backend) — declared on the card via ``ComponentSpec.adapter`` +so the recipe is self-contained (no edit to the shared ``_make_dit``/``_make_vae``/``_make_text_encoder`` +dispatch in ``torch_backend.py``). Imported lazily by ``_explicit_adapter`` only on a GPU box. + +Stable Audio is an AUDIO model and differs from every existing (video) adapter on several axes, all of +which are absorbed HERE so ``StableAudioDenoiseLoop`` keeps the cosmos2-style ``dit(latent, text_embed, +sigma)`` call shape (and so the CPU toy backend works unchanged): + +* ``StableAudioDiT`` — the v-prediction (EDM-v) network. forward is ``forward(x[B,64,L], t[B], *, + cross_attn_cond[B,seq,768], global_embed[B,1536]) -> v[B,64,L]`` (1-D audio latent, NOT [C,T,H,W]). + cross_attn_cond/global_embed are keyword-only and the DiT does NOT do CFG internally — the loop hands + a single batch, CFG is combined in the loop (faithful to ``stable_audio/stages/denoising.py:_DiTAdapter``). + The adapter packs the loop's flat ``text_embed`` payload into (cross_attn_cond, global_embed) and passes + the RAW continuous sigma as the timestep ``t`` (NOT sigma*1000 (Wan), NOT 0..1 (LTX-2); k-diffusion / + EDM c_noise convention — BRINGUP risk B). It returns the raw network output; the v->x0 (VDenoiser) + preconditioning lives in the loop. +* ``OobleckVAE`` — 1-D conv audio VAE (weight_norm + Snake1d, hop_length=2048). encode(audio[2,N]) -> + raw latent (``dist.sample()``, NO mean/std normalization — passthrough latent space, BRINGUP risk D); + decode(latent[64,L]) -> waveform[2,M]. 1-D marshalling (no 5-D unsqueeze). Exposes ``hop_length`` and + ``sampling_rate`` so the loop sizes the latent and the program slices [start,end] seconds. +* ``StableAudioConditioner`` — the SA multi-conditioner (T5-base @ max_length=128 + duration + NumberConditioners). encode(prompt, seconds_start, seconds_total) -> packs the (cross_attn_cond, + cross_attn_mask, global_embed) triple into a single numpy payload the loop stashes in cond. NOTE: the + shared ``_MAKERS`` dispatch has no ``conditioner`` kind and ``TextEncoderLoader`` will not build a + ``StableAudioMultiConditioner`` — so the card declares this component with ``kind="text_encoder"`` + + this explicit adapter, and the adapter is constructed ``cls(module, tokenizer, ...)`` (the tokenizer + extra the text-encoder maker passes is unused — the conditioner owns its own T5 tokenizer). The maker + still calls ``TextEncoderLoader`` though, which loads a generic T5, NOT the SA conditioner — so wiring + this on a GPU box needs the SA ``ConditionerLoader`` path (BRINGUP risk: the conditioner load). + +All three are BRINGUP (written-not-run): there is no GPU/k_diffusion in this env. See the loop for the +dpmpp-3m-sde GPU path that drives ``StableAudioDiT`` through ``K.external.VDenoiser``. +""" +from __future__ import annotations + +import numpy as np + +from v2.platform.backends.torch_backend import TorchComponent, _to_numpy + +# Packed conditioning payload layout. The loop treats the conditioner output as one flat "text_embed" +# numpy array; the DiT adapter splits it back into cross_attn_cond + global_embed by these widths so the +# loop->dit call stays ``dit(latent, text_embed, sigma)`` (cosmos2-shape) and the CPU toy still runs. +_GLOBAL_COND_DIM = 1536 # StableAudioArchConfig.global_cond_dim +_CROSS_COND_DIM = 768 # StableAudioArchConfig.cond_token_dim + +# The global embed (1536) is laid down as ``_GLOBAL_HEAD_ROWS`` leading rows of width 768 (1536 = 2*768), +# then the cross-attn tokens (each 768-wide) follow. Lets the whole conditioning ride as ONE [N,768] array. +_GLOBAL_HEAD_ROWS = _GLOBAL_COND_DIM // _CROSS_COND_DIM # = 2 + + +def pack_conditioning(cross_attn_cond: np.ndarray, global_embed: np.ndarray) -> np.ndarray: + """Pack (cross_attn_cond[seq,768], global_embed[1536]) -> a single [2+seq, 768] numpy payload: the + first 2 rows are the 1536-d global embed reshaped to [2,768], rows 2.. are the cross-attn tokens. The + loop carries this opaque array as ``text_embed``; ``StableAudioDiT.__call__`` (and the CPU toy, which + just means-pools it) consume it. Keeps the loop modality-agnostic — it never sees the SA tuple.""" + g = np.asarray(global_embed, dtype="float32").reshape(-1) + head = np.zeros((_GLOBAL_HEAD_ROWS, _CROSS_COND_DIM), dtype="float32") + head.reshape(-1)[:min(g.shape[0], _GLOBAL_COND_DIM)] = g[:_GLOBAL_COND_DIM] # row-major lay-down + cross = np.asarray(cross_attn_cond, dtype="float32").reshape(-1, _CROSS_COND_DIM) + return np.concatenate([head, cross], axis=0) # [2 + seq, 768]; rows0-1=global, rows2..=cross-attn + + +def unpack_conditioning(payload: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Inverse of ``pack_conditioning``: -> (cross_attn_cond[seq,768], global_embed[1536]).""" + p = np.asarray(payload, dtype="float32") + global_embed = p[:_GLOBAL_HEAD_ROWS].reshape(-1)[:_GLOBAL_COND_DIM].copy() + cross = p[_GLOBAL_HEAD_ROWS:] + return cross, global_embed + + +class StableAudioDiT(TorchComponent): + """``dit(latent[64,L], text_embed=, sigma) -> raw v-prediction[64,L]``. + + ``sigma`` is the RAW continuous sigma (k-diffusion / EDM c_noise; ~0.3..500). The loop hands the + ALREADY-VDenoiser-input-scaled latent (``x·c_in``) — this adapter is the bare network ``F_theta`` and + returns the raw output; the v->x0 reconstruction + x0-space CFG live in ``StableAudioDenoiseLoop`` + (faithful to ``_DiTAdapter`` + ``K.external.VDenoiser``). Single batch; CFG batching is the loop's job. + """ + + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + import torch + cross_np, global_np = unpack_conditioning(text_embed) + hs = self._t(latent) # [1, 64, L] + cross = self._t(cross_np) # [1, seq, 768] + glob = self._t(global_np) # [1, 1536] + t = torch.tensor([float(sigma)], device=self.device, dtype=self.dtype) # RAW sigma, NOT *1000 + with torch.no_grad(), self._ctx(): + out = self.module(hs, t, cross_attn_cond=cross, global_embed=glob) + return self._n(out) # RAW v-prediction (loop does VDenoiser v->x0) + + +class OobleckVAE(TorchComponent): + """1-D audio VAE. decode(latent[64,L]) -> waveform[channels, samples]; encode(audio[ch,N]) -> raw + latent (NO mean/std normalization — SA's DiT operates directly in the Oobleck latent space; BRINGUP + risk D). Exposes ``hop_length`` / ``sampling_rate`` (read off the loaded module).""" + + def __init__(self, module, *, device, dtype): + super().__init__(module, device=device, dtype=dtype) + self.hop_length = int(getattr(module, "hop_length", 2048)) + self.sampling_rate = int(getattr(module, "sampling_rate", 44100)) + + def decode(self, latent): + import torch + z = self._t(latent).to(self.dtype) # [1, 64, L] + with torch.no_grad(), self._ctx(): + out = self.module.decode(z) + wav = out.sample if hasattr(out, "sample") else out # [1, channels, samples] + return _to_numpy(wav.squeeze(0)) # [channels, samples] + + def encode(self, audio): + import torch + x = self._t(audio).to(self.dtype) # [1, channels, samples] + with torch.no_grad(), self._ctx(): + dist = self.module.encode(x) + z = dist.sample() if hasattr(dist, "sample") else dist # raw latent, NO normalization + return _to_numpy(z.squeeze(0)) # [64, L] + + +class StableAudioConditioner(TorchComponent): + """SA multi-conditioner (T5-base @ max_length=128 + duration NumberConditioners). ``encode`` packs the + (cross_attn_cond, cross_attn_mask, global_embed) triple into the loop's flat payload via + ``pack_conditioning``. The DiT consumes a single-batch payload; CFG (null/zero negative cross-attn) is + built in the loop. + + BRINGUP: the shared text-encoder maker calls ``TextEncoderLoader`` (a generic T5), not the SA + ``ConditionerLoader`` that builds ``StableAudioMultiConditioner`` — so on a GPU box this adapter must + instead be handed a real ``StableAudioMultiConditioner`` module (e.g. via a recipe-local loader + override / a ``conditioner`` component kind). The ``tokenizer`` extra the text-encoder maker passes is + accepted and ignored (the conditioner owns its own T5 tokenizer).""" + + def __init__(self, module, tokenizer=None, *, device, dtype): + super().__init__(module, device=device, dtype=dtype) + # The active sub-conditioner ids (3 for SA-1.0: prompt+seconds_start+seconds_total; 2 for + # SA-small: prompt+seconds_total). Read off the module so the variant is self-describing. + self.cross_attention_cond_ids = tuple( + getattr(module, "cross_attention_cond_ids", ("prompt", "seconds_start", "seconds_total"))) + + def encode(self, prompt: str, seconds_start: float = 0.0, seconds_total: float = 10.0) -> np.ndarray: + import torch + all_vals = { + "prompt": prompt or "", + "seconds_start": float(seconds_start), + "seconds_total": float(seconds_total) + } + cond_meta = [{k: all_vals[k] for k in self.cross_attention_cond_ids if k in all_vals}] + with torch.no_grad(), self._ctx(): + cond = self.module(cond_meta, self.device) + cross, _mask, glob = self.module.get_conditioning_inputs(cond) + return pack_conditioning(_to_numpy(cross.squeeze(0)), _to_numpy(glob.squeeze(0))) diff --git a/v2/recipes/cosmos25/__init__.py b/v2/recipes/cosmos25/__init__.py new file mode 100644 index 000000000..4fd4ccacc --- /dev/null +++ b/v2/recipes/cosmos25/__init__.py @@ -0,0 +1,17 @@ +"""Cosmos-Predict2.5 (text->video) — a flow-match denoiser ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``Cosmos25DiT`` / ``Cosmos25WanVAE`` / ``Cosmos25Reason1Encoder`` in +``v2/platform/backends/torch_cosmos25.py``) and a new ``Cosmos25DenoiseLoop`` (Wan's flow-shift sigma +schedule + ``FLOW_MATCH_STEP`` solver + ``ClassicCFG``, but the PLAIN per-frame sigma timestep ``[B,T]`` +Cosmos2.5 needs — NOT Wan's ``sigma*1000``), reusing ``stamp_wan21_checkpoints``. Covers both the 2B and +14B sizes (the card is size-agnostic — size resolves from the checkpoint config). Registered in +``v2/registry.py`` by the orchestrator from the returned metadata. +""" +from __future__ import annotations + +from v2.recipes.cosmos25.card import build_cosmos25_card +from v2.recipes.cosmos25.loop import Cosmos25DenoiseLoop +from v2.recipes.cosmos25.program import build_cosmos25_program + +__all__ = ["build_cosmos25_card", "build_cosmos25_program", "Cosmos25DenoiseLoop"] diff --git a/v2/recipes/cosmos25/card.py b/v2/recipes/cosmos25/card.py new file mode 100644 index 000000000..cfda5b3cc --- /dev/null +++ b/v2/recipes/cosmos25/card.py @@ -0,0 +1,142 @@ +"""Cosmos-Predict2.5 ModelCard — text->video (the registered preset; video2world/image2world +conditioning is a later capability the loop already threads). Two sizes — 2B +(``KyleShao/Cosmos-Predict2.5-2B-Diffusers``) and 14B (``nvidia/Cosmos-Predict2.5-14B``) — share ONE +card: the transformer's ``num_attention_heads``/``num_layers`` are resolved from the checkpoint config by +the loader, so the card is size-agnostic and only the sampling defaults differ (we use the 2B preset). + +Architecture deltas vs Wan (all GPU-path, declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.cosmos2_5:Cosmos25Transformer3DModel`` — a flow-match velocity DiT whose + adapter (``ComponentSpec.adapter`` -> ``Cosmos25DiT``) feeds a PER-FRAME plain-sigma timestep ``[B,T]`` + (NOT Wan's scalar ``sigma*1000``) + a mandatory zero ``condition_mask`` / ones ``padding_mask`` + an + ``fps`` scalar. The model concats the masks itself (16 -> 18ch), so we feed the RAW 16ch latent. + * VAE ``fastvideo.models.vaes.cosmos25wanvae:Cosmos25WanVAE`` — Wan-style causal 3D VAE (z=16, 8x/4x) + whose encode/decode normalize INTERNALLY (the Cosmos latent contract), so a dedicated + ``Cosmos25WanVAE`` adapter marshals WITHOUT re-normalizing (reusing the v2 WanVAE adapter would + double-normalize — see ``torch_cosmos25.py``). + * Text ``fastvideo.models.encoders.reason1:Reason1TextEncoder`` — Qwen2.5-VL multimodal LM via + ``Cosmos25Reason1Encoder`` (per-layer mean-normalized full-concat -> a 100352-dim sequence; the DiT's + internal 100352->1024 GELU crossattn projection consumes it). +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout + a sibling +tokenizer dir). Policies: ``ClassicCFG`` (embedded_cfg_scale=0 -> plain CFG), ``FlowShiftPolicy(shift=5.0)``, +fp32 scheduler step. Faithful to ``fastvideo/pipelines/stages/denoising.py:Cosmos25DenoisingStage``. +""" +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.cosmos25.loop import Cosmos25DenoiseLoop +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# Cosmos-Predict2.5 default t2v negative prompt (a LOCAL recipe constant — the Cosmos2.5 preset's own +# quality-suppression prompt, encoded through the same Reason1 path for ClassicCFG). +_COSMOS25_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.") + +_COSMOS25_DIT = "v2.platform.backends.torch_cosmos25:Cosmos25DiT" +_COSMOS25_VAE = "v2.platform.backends.torch_cosmos25:Cosmos25WanVAE" +_COSMOS25_TEXT = "v2.platform.backends.torch_cosmos25:Cosmos25Reason1Encoder" + + +def build_cosmos25_card(model_id: str = "cosmos-predict2.5-2b", + *, + flow_shift: float = 5.0, + 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() + flow = FlowShiftPolicy(shift=flow_shift) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return Cosmos25DenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost) + + components = { + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.reason1:Reason1TextEncoder", + adapter=_COSMOS25_TEXT, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.cosmos25wanvae:Cosmos25WanVAE", + adapter=_COSMOS25_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.cosmos2_5:Cosmos25Transformer3DModel", + adapter=_COSMOS25_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="cosmos25", + 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=77, + fps=24, + negative_prompt=_COSMOS25_NEG), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card diff --git a/v2/recipes/cosmos25/loop.py b/v2/recipes/cosmos25/loop.py new file mode 100644 index 000000000..9883db82a --- /dev/null +++ b/v2/recipes/cosmos25/loop.py @@ -0,0 +1,222 @@ +"""Cosmos25DenoiseLoop — a flow-match Euler denoiser with Cosmos2.5's per-frame sigma timestep. + +Cosmos-Predict2.5 IS a rectified-flow / flow-match model (unlike Cosmos-Predict2, which is an EDM +denoiser). The sigma schedule and the deterministic Euler solver are EXACTLY Wan's +(``FlowShiftPolicy.build_schedule`` + ``FLOW_MATCH_STEP``), and CFG is the standard ``ClassicCFG`` +(``uncond + s*(cond-uncond)``). So this loop is almost identical to ``WanDenoiseLoop``; it forks it for +ONE load-bearing reason (blocker (1)): + + * Wan feeds the DiT a SCALAR ``timestep = sigma * num_train_timesteps`` (``sigma*1000``). + * Cosmos2.5 feeds a PER-FRAME 2D ``timestep[B,T]`` of the PLAIN sigma in [0,1] — the fastvideo stage + computes ``t * 0.001`` of the 0..1000 schedule, which is precisely the shifted sigma. So we set + ``st.timesteps = st.sigmas`` (NO *1000) and the ``Cosmos25DiT`` adapter broadcasts that plain sigma + to ``[B, T]`` and builds the mandatory zero ``condition_mask`` / ones ``padding_mask`` + ``fps`` the + forward requires (the model concats the masks itself — we feed the RAW 16ch latent). See + ``v2/platform/backends/torch_cosmos25.py:Cosmos25DiT`` and the faithful source + ``fastvideo/pipelines/stages/denoising.py:Cosmos25DenoisingStage``. + +Verification that the math matches the source: the fastvideo ``FlowMatchEulerDiscreteScheduler`` +(shift=5.0) builds ``sigma' = shift*sigma/(1+(shift-1)*sigma)`` over a 1->0 base + a terminal 0, sets +``timesteps = sigma'*1000``, feeds the DiT ``timesteps*0.001 = sigma'``, and steps +``prev = sample + (sigma_next - sigma)*velocity`` — i.e. exactly ``build_flow_sigmas(num_steps, +shift=5.0)`` + ``FLOW_MATCH_STEP``. + +video2world / image2world (frame_replace) conditioning (cond_indicator GT-x0 clamp + GT-velocity +injection, conditional_frame_timestep=0.1 on conditioned frames) is present in fastvideo +(``Cosmos25V2WDenoisingStage``) but is NOT in the registered HF presets — this loop scopes to t2v. +The hooks are read from ``ctx.slots`` (None for t2v -> inert), mirroring the Wan/Cosmos2 pattern, so a +later capability adds the gated injection without changing the t2v path. (BRINGUP: frame-replace needs a +request-API extension to carry the conditioning latents + indicator.) +""" +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.platform import FLOW_MATCH_STEP +from v2.recipes.wan21.loop import latent_shape + +# Real Cosmos2.5 Wan-style VAE (Cosmos25WanVAE) compression: z_dim=16, 4x temporal, 8x spatial — the +# same geometry as Wan2.1, so ``latent_shape`` (reused from the Wan recipe) yields +# (16, (F-1)//4 + 1, H//8, W//8) on the real backend. +COSMOS25_LATENT_CHANNELS = 16 +COSMOS25_TEMPORAL_RATIO = 4 +COSMOS25_SPATIAL_RATIO = 8 +COSMOS25_DEFAULT_FPS = 24 + + +class Cosmos25DenoiseLoop: + """Flow-match Euler denoise loop with Cosmos2.5's plain per-frame sigma timestep convention.""" + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels: int = COSMOS25_LATENT_CHANNELS, + spatial_ratio: int = COSMOS25_SPATIAL_RATIO, + temporal_ratio: int = COSMOS25_TEMPORAL_RATIO): + self.loop_id = loop_id + self.cfg = cfg + self.flow_shift = flow_shift + self.precision = precision + self.expert = expert + self.cost = cost + 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) + # The σ schedule + flow-shift are Wan's verbatim (shift=5.0 for Cosmos2.5): 1 -> 0 over + # num_steps+1 points with the flow-shift warp applied. + sig = self.flow_shift.build_schedule(req.diffusion.num_steps, + req.diffusion.height, + req.diffusion.width, + sigmas=req.diffusion.sigmas or None) + shape = latent_shape(req, + model, + channels=self.latent_channels, + spatial_ratio=self.spatial_ratio, + temporal_ratio=self.temporal_ratio) + # Flow-match noise init: x = randn * sigma_0 (sigma_0 == 1 for an unshifted-terminal schedule). + x = (rng.standard_normal(shape) * float(sig[0])).astype("float32") + 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], + # BLOCKER (1): the model timestep is the PLAIN sigma (NOT sigma*1000). The + # Cosmos25DiT adapter broadcasts it to the per-frame [B, T] form the forward needs. + timesteps=[float(s) for s in sig]) + 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")) + st.scratch["fps"] = int(getattr(req.diffusion, "fps", None) or COSMOS25_DEFAULT_FPS) + # video2world (frame_replace) conditioning: None for the t2v preset -> the gated injection a later + # capability would add is inert (pure t2v, matching the fastvideo pipeline with no image/video). + st.scratch["conditioning_latents"] = ctx.slots.get("conditioning_latents") + st.scratch["cond_indicator"] = ctx.slots.get("cond_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] + # The model timestep here is the PLAIN sigma (st.timesteps[i] == st.sigmas[i]); the adapter turns + # it into the per-frame [B, T] timestep + masks + fps internally. + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + branches = self.cfg.branches_this_step(sctx, cfg_state) + x = st.latents["video"] + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + scale = st.scratch["guidance_scale"] + cfg, precision = self.cfg, self.precision + + def _velocity(model: Any, x_: Any, sigma_t_: float, pe_: Any, ne_: Any, scale_: float) -> np.ndarray: + # The conditioned forward + ClassicCFG combine. The dit-call is toy-compatible: + # ``dit(latent, text_embed, sigma)`` — the Cosmos2.5-specific timestep/masks/fps are built + # INSIDE the Cosmos25DiT adapter (so the CPU ToyDiT, which ignores them, also runs). + dit = model.component(expert_id) + preds = {b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_) for b in branches} + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + + def run(model, override=None): + # The EAGER thunk: handles the override (cached-prediction) path. The deterministic Euler + # solver dispatches through the platform kernel table (numpy on CPU, the device kernel on GPU). + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x, sigma_t, pe, ne, scale) + x_next = 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") + }) + + def graph_fn(model, ws): + # The CAPTURABLE op-structure: reads EVERY per-step input from the static workspace and writes + # into the static output buffer, so the captured graph replays correctly with rebound buffers. + st_t, st_n = float(ws["sigma_t"]), float(ws["sigma_next"]) + velocity = _velocity(model, ws["x"], st_t, ws["pe"], ws["ne"], float(ws["scale"])) + x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(precision.cast(ws["x"]), velocity, st_t, st_n) + np.copyto(ws["out"], x_next.astype("float32")) + return StepResult(output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": np.array(ws["out"], copy=True) + }) + + 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)), float(len(branches))), + 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(cfg).__name__), )), + resources=res, + payload={ + "branch": "combined", + "step": i + }, + run=run, + label=f"cosmos25.denoise.{i}", + # The flow-match Euler step is a clean capturable op (no host RNG / data-dependent branch), + # like Wan t2v: the branch-set + expert + scheduler-precision discriminate the captured graph. + capturable=True, + graph_key=(tuple(sorted(branches)), expert_id, precision.scheduler_step_in_fp32), + graph_fn=graph_fn, + graph_inputs={ + "x": x, + "sigma_t": sigma_t, + "sigma_next": sigma_next, + "pe": pe, + "ne": ne, + "scale": scale + }) + + 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/cosmos25/program.py b/v2/recipes/cosmos25/program.py new file mode 100644 index 000000000..c8b956cc1 --- /dev/null +++ b/v2/recipes/cosmos25/program.py @@ -0,0 +1,37 @@ +"""Cosmos-Predict2.5 t2v program: text_encode -> diffusion_denoise (flow-match) -> vae_decode. + +Same inline shape as the Wan / Cosmos2 t2v programs — the Cosmos2.5 specifics (per-frame plain-sigma +timestep, the mandatory masks + fps, the Reason1 100352-dim embedding) live entirely in +``Cosmos25DenoiseLoop`` and the ``torch_cosmos25`` adapters, so the node graph is unchanged. fps is read +by the loop from the request (``req.diffusion.fps``, default 24) into scratch and threaded to the DiT +adapter. (video2world/image2world conditioning would add a VAE-encode node writing +``conditioning_latents``/``cond_indicator`` into slots; the loop already reads them — BRINGUP.) +""" +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_cosmos25_program() -> Program: + return Program( + program_id="cosmos25.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/recipes/flux2/__init__.py b/v2/recipes/flux2/__init__.py new file mode 100644 index 000000000..994434b6c --- /dev/null +++ b/v2/recipes/flux2/__init__.py @@ -0,0 +1,17 @@ +"""FLUX.2 (Black Forest Labs) — dual-stream MMDiT text→image ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``Flux2DiT``/``Flux2VAE``/``Flux2Mistral3Encoder``/``Flux2Qwen3Encoder`` in +``v2/platform/backends/torch_flux2.py``) and a NEW ``Flux2DenoiseLoop`` (the BFL empirical-mu flow-match +schedule + packed 2×2 latent geometry + single embedded-guidance forward), reusing ``stamp_wan21_checkpoints`` +for the diffusers transformer/vae/text_encoder subfolder layout. ``build_flux2_card`` is FLUX.2-dev +(Mistral3, embedded guidance); ``build_flux2_klein_card`` is the distilled 4-step klein variants (Qwen3). +GATED weights → GPU is BRINGUP; the toy factories CPU-verify the recipe end-to-end. +""" +from __future__ import annotations + +from v2.recipes.flux2.card import build_flux2_card, build_flux2_klein_card +from v2.recipes.flux2.loop import Flux2DenoiseLoop +from v2.recipes.flux2.program import build_flux2_program + +__all__ = ["build_flux2_card", "build_flux2_klein_card", "build_flux2_program", "Flux2DenoiseLoop"] diff --git a/v2/recipes/flux2/card.py b/v2/recipes/flux2/card.py new file mode 100644 index 000000000..ec7ad383d --- /dev/null +++ b/v2/recipes/flux2/card.py @@ -0,0 +1,155 @@ +"""FLUX.2 ModelCard(s) — text→image (BFL FLUX.2-dev primary; klein 4B/9B distilled variants). + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained — the bucket-C +pattern: torch adapters via ``ComponentSpec.adapter``, a NEW ``Flux2DenoiseLoop``, ``stamp_wan21_checkpoints`` +for the diffusers transformer/vae/text_encoder subfolder layout): + + * DiT ``fastvideo.models.dits.flux_2:Flux2Transformer2DModel`` — dual-stream MMDiT (19 double + 38 + single blocks, inner_dim 3072, 24 heads, N-D RoPE axes (32,32,32,32) θ=2000). The adapter + (``Flux2DiT``) builds img_ids/txt_ids RoPE position tensors from the latent shape, passes σ DIRECTLY + as the timestep and the embedded guidance RAW (the DiT multiplies both by 1000 internally), wraps the + text embed in a 1-element list, and returns the bare velocity. NO classic CFG (embedded guidance). + * VAE ``fastvideo.models.vaes.flux2vae:AutoencoderKLFlux2`` — packed 2×2 latents + BatchNorm whitening + (NOT the Wan (z−mean)/std). ``Flux2VAE`` decode does BN-denorm (``z·bn_std + bn_mean``) → 2×2 unpatchify + (64→16ch, full spatial) → ``module.decode`` → image. (Image-only: T==1, squeezed.) + * Text dev=``mistral3:Mistral3ForConditionalGeneration`` (layers 10/20/30, FLUX2 system-message chat + template); klein=``qwen3:Qwen3ForCausalLM`` (layers 9/18/27). ``Flux2TextEncoder`` stacks the chosen + hidden-state layers into a [seq, 3·hidden=4096] embedding — NOT T5 last_hidden_state. + +GATED weights (``black-forest-labs/FLUX.2-*``) → GPU is BRINGUP; the toy factories CPU-verify the recipe. +""" +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 EmbeddedGuidance, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.flux2.loop import Flux2DenoiseLoop +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# FLUX.2 conditions on a single embedded-guidance forward (no CFG), so there is no negative prompt — kept +# as a LOCAL module constant (the recipe is self-contained; nothing is added to the shared _prompts.py). +FLUX2_NEG: str | None = None + +_FLUX2_DIT = "v2.platform.backends.torch_flux2:Flux2DiT" +_FLUX2_VAE = "v2.platform.backends.torch_flux2:Flux2VAE" +_FLUX2_MISTRAL3 = "v2.platform.backends.torch_flux2:Flux2Mistral3Encoder" +_FLUX2_QWEN3 = "v2.platform.backends.torch_flux2:Flux2Qwen3Encoder" + + +def _build_flux2_card(model_id: str, *, text_encoder_load_id: str, text_encoder_adapter: str, + checkpoint_root: str | None, sampling_defaults: SamplingDefaults) -> ModelCard: + seed = _seed_from(model_id) + cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7) + cfg = EmbeddedGuidance() # single conditioned forward; no uncond branch + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return Flux2DenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, precision=precision, expert=expert, cost=cost) + + components = { + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id=text_encoder_load_id, + adapter=text_encoder_adapter, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2i"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.flux2vae:AutoencoderKLFlux2", + adapter=_FLUX2_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2i"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.flux_2:Flux2Transformer2DModel", + adapter=_FLUX2_DIT, + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["diffusion_denoise"], + required_for={"t2i"}), + } + 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="flux2", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.TEXT_TO_IMAGE, Capability.VAE_DECODE), + 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, + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card + + +def build_flux2_card(model_id: str = "flux2-dev", *, checkpoint_root: str | None = None) -> ModelCard: + """FLUX.2-dev (the registered primary): Mistral3 text encoder (layers 10/20/30 via FLUX2 chat + template), embedded guidance scale 4.0, 50 steps, 1024×1024. ``embedded_cfg_scale`` rides in + ``guidance_scale``; ``do_classifier_free_guidance=False`` (no negative prompt).""" + return _build_flux2_card(model_id, + text_encoder_load_id="fastvideo.models.encoders.mistral3:Mistral3ForConditionalGeneration", + text_encoder_adapter=_FLUX2_MISTRAL3, + checkpoint_root=checkpoint_root, + sampling_defaults=SamplingDefaults(num_steps=50, + guidance_scale=4.0, + height=1024, + width=1024, + num_frames=1, + fps=1, + negative_prompt=FLUX2_NEG)) + + +def build_flux2_klein_card(model_id: str = "flux2-klein-4b", *, checkpoint_root: str | None = None) -> ModelCard: + """FLUX.2-klein (4B/9B, distilled): Qwen3 text encoder (layers 9/18/27), NO guidance embedding + (guidance_scale 1.0 -> single forward, no uncond), 4 steps, 1024×1024. The schedule still keys on the + same empirical-mu grid; the 4-step bf16 parity is the GPU BRINGUP detail.""" + return _build_flux2_card(model_id, + text_encoder_load_id="fastvideo.models.encoders.qwen3:Qwen3ForCausalLM", + text_encoder_adapter=_FLUX2_QWEN3, + checkpoint_root=checkpoint_root, + sampling_defaults=SamplingDefaults(num_steps=4, + guidance_scale=1.0, + height=1024, + width=1024, + num_frames=1, + fps=1, + negative_prompt=FLUX2_NEG)) diff --git a/v2/recipes/flux2/loop.py b/v2/recipes/flux2/loop.py new file mode 100644 index 000000000..561f2661b --- /dev/null +++ b/v2/recipes/flux2/loop.py @@ -0,0 +1,202 @@ +"""Flux2DenoiseLoop — FLUX.2 dual-stream MMDiT denoise on the BFL empirical-mu flow-match schedule. + +FLUX.2 is a flow-match (velocity-prediction) image model, but it is NOT a Wan-shaped port (see +``v2/recipes/flux2/card.py`` for the component-level deltas). The deltas that force a NEW loop, all faithful +to ``fastvideo/pipelines/basic/flux_2`` + ``fastvideo/pipelines/stages/denoising.py``: + + * **Schedule** — the σ grid is the BFL empirical-mu grid (``flux2_sigmas`` here: ``linspace(1, 1/N, N)`` + shifted by ``compute_empirical_mu(image_seq_len, N)`` via the FlowMatchEuler exponential time-shift), + NOT ``FlowShiftPolicy.build_schedule``'s resolution-bucket shift. ``image_seq_len`` is the PACKED token + count ``latent_h * latent_w``. + * **Packed latent geometry** — the DiT operates on ``num_channels_latents=64`` channels at HALF spatial + (``latent_h = (H//8)//2``); the VAE decode unpacks the 2×2 patches back to 16ch full-spatial. So the + loop initializes a ``(64, T=1, (H//8)//2, (W//8)//2)`` latent (image: ``T==1``). + * **Embedded guidance, single forward** — FLUX.2-dev folds a guidance scalar into the DiT forward + (``EmbeddedGuidance`` policy: one conditioned branch, no uncond), and FLUX.2-klein uses no guidance at + all. There is NO classic 2-branch CFG, so the loop runs a single ``dit(x, pe, σ)`` call per step. + * **img_ids / txt_ids RoPE + the ×1000 timestep/guidance split** — these are GPU-path concerns the + ``Flux2DiT`` torch adapter handles INTERNALLY (it derives img_ids/txt_ids from the latent shape + + text-embed length, passes σ DIRECTLY as the timestep, and the embedded guidance raw — the DiT + multiplies both by 1000 itself). The loop therefore calls the dit exactly like the CPU ``ToyDiT`` + (``dit(latent, text_embed, sigma)``), so it CPU-verifies unchanged. + +The velocity is integrated with the shared deterministic flow-match Euler step +(``x_next = x + (σ_next − σ)·v`` via ``FLOW_MATCH_STEP``) — identical solver to Wan/LTX-2. +""" +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.platform import FLOW_MATCH_STEP +from v2.platform.backends.toy import LATENT_CHANNELS +from v2.recipes.flux2.sampler import flux2_sigmas + +# Real FLUX.2 packed geometry: the VAE has latent_channels=16 at 8× spatial compression; the DiT sees +# the 2×2-packed latent -> 16·4 = 64 channels at HALF spatial. (config in_channels = num_channels_latents +# = 64; the VAE's _unpatchify_latents does C//4 -> 16ch full-spatial on decode.) +FLUX2_PACKED_CHANNELS = 64 +FLUX2_SPATIAL_RATIO = 8 # VAE spatial_compression_ratio +FLUX2_PACK = 2 # 2×2 patch packing -> half spatial + + +def flux2_latent_geometry(req, model=None) -> tuple[tuple[int, int, int, int], int]: + """Packed latent shape + packed image_seq_len. On the GPU backend the DiT/VAE require the true FLUX.2 + geometry; the CPU toy uses a tiny deterministic stand-in (so the loop control flow is exercised with + real numbers). Returns ``((C, T, latent_h, latent_w), image_seq_len)`` where ``image_seq_len`` is the + packed token count ``latent_h*latent_w`` the empirical-mu schedule keys on. Image model: ``T == 1``.""" + d = req.diffusion + if model is not None and getattr(getattr(model, "platform", None), "device", "cpu") == "cuda": + latent_h = max(1, (d.height // FLUX2_SPATIAL_RATIO) // FLUX2_PACK) + latent_w = max(1, (d.width // FLUX2_SPATIAL_RATIO) // FLUX2_PACK) + return (FLUX2_PACKED_CHANNELS, 1, latent_h, latent_w), latent_h * latent_w + latent_h = max(2, d.height // 240) + latent_w = max(2, d.width // 240) + return (LATENT_CHANNELS, 1, latent_h, latent_w), latent_h * latent_w + + +class Flux2DenoiseLoop: + """N flow-match Euler steps over the full packed latent, single conditioned forward per step + (embedded guidance for dev, no guidance for klein). Mirrors ``WanDenoiseLoop`` structure; the FLUX.2 + schedule + packed geometry are the only deltas. All per-request state lives in ``LoopState`` + (interleave-safe).""" + + def __init__(self, *, loop_id, cfg, precision, expert, cost, packed_channels: int = FLUX2_PACKED_CHANNELS): + self.loop_id = loop_id + self.cfg = cfg # EmbeddedGuidance (dev) — carried for the WorkPlan op-structure key + self.precision = precision + self.expert = expert + self.cost = cost + self.packed_channels = packed_channels + + 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) + shape, image_seq_len = flux2_latent_geometry(req, model) + sig = flux2_sigmas(req.diffusion.num_steps, image_seq_len) + # FlowMatchEuler init_noise_sigma == 1.0 for this scheduler, so x starts at plain randn. + x = rng.standard_normal(shape).astype("float32") + 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={"image": x}, + sigmas=[float(s) for s in sig], + timesteps=[float(s) * 1000.0 for s in sig]) # bookkeeping; the adapter re-derives + st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") + # embedded-guidance scalar (dev=4.0; klein=None -> no guidance). The Flux2DiT adapter reads it and + # folds it into the forward; there is no uncond branch. + st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale) + st.scratch["stream_image"] = bool(req.outputs.stream.get("image")) + 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] + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + branches = self.cfg.branches_this_step(sctx, cfg_state) # ["cond"] for EmbeddedGuidance + x = st.latents["image"] + pe = st.cond["prompt_embeds"] + scale = st.scratch["guidance_scale"] + cfg, precision = self.cfg, self.precision + + def _velocity(model: Any, x_: Any, sigma_t_: float, pe_: Any, scale_: float) -> np.ndarray: + # Single conditioned forward; EmbeddedGuidance.combine is the identity on the cond branch. The + # DiT call is toy-compatible (latent, text_embed, sigma) — img_ids/txt_ids/guidance/×1000 are + # built INSIDE the Flux2DiT adapter, so the same loop drives both the CPU toy and the GPU DiT. + dit = model.component(expert_id) + preds = {b: dit(x_, pe_, sigma_t_) for b in branches} + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + + def run(model, override=None): + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x, sigma_t, pe, scale) + x_next = 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") + }) + + def graph_fn(model, ws): + # Capturable op-structure: reads every per-step input from the static workspace (never closure). + st_t, st_n = float(ws["sigma_t"]), float(ws["sigma_next"]) + velocity = _velocity(model, ws["x"], st_t, ws["pe"], float(ws["scale"])) + x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(precision.cast(ws["x"]), velocity, st_t, st_n) + np.copyto(ws["out"], x_next.astype("float32")) + return StepResult(output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": np.array(ws["out"], copy=True) + }) + + cond_bytes = int(np.asarray(pe).nbytes) if pe is not None else 0 + res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), + resident_bytes=int(x.nbytes) + cond_bytes, + peak_activation_bytes=int(x.nbytes)) + emits = [] + if st.scratch.get("stream_image"): + from v2.request.streams import StreamChunk + emits.append(StreamChunk(stream_id=st.request_id, modality="image", seq=i, data=x, preview=True)) + 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(cfg).__name__), )), + resources=res, + payload={ + "branch": "embedded", + "step": i + }, + run=run, + label=f"flux2.denoise.{i}", + emits=emits, + capturable=True, + graph_key=(tuple(sorted(branches)), expert_id, precision.scheduler_step_in_fp32), + graph_fn=graph_fn, + graph_inputs={ + "x": x, + "sigma_t": sigma_t, + "sigma_next": sigma_next, + "pe": pe, + "scale": scale + }) + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + st.latents["image"] = 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["image"]).copy(), + }) + st.step_idx += 1 + return st + + def finalize(self, st: LoopState) -> LoopResult: + return LoopResult(outputs={"latents": st.latents["image"]}, + metrics={"denoise_steps": float(st.step_idx)}, + behavior=st.trajectory or None) diff --git a/v2/recipes/flux2/program.py b/v2/recipes/flux2/program.py new file mode 100644 index 000000000..95e8de24d --- /dev/null +++ b/v2/recipes/flux2/program.py @@ -0,0 +1,43 @@ +"""FLUX.2 t2i program: text_encode → diffusion_denoise → vae_decode → image. + +Same inline shape as the Wan t2v program (text_encode → denoise → vae_decode); the FLUX.2 specifics live +entirely in ``Flux2DenoiseLoop`` (BFL empirical-mu schedule + packed geometry + single embedded-guidance +forward) and the ``Flux2DiT``/``Flux2VAE``/``Flux2TextEncoder`` adapters. The reads only ``text_embeds`` +(no ``neg_text_embeds``): FLUX.2-dev uses embedded guidance and klein uses none — there is no CFG uncond +branch. The output artifact is an ``image`` (the loop produces a single-frame ``T==1`` packed latent which +the VAE unpacks + decodes to pixels). +""" +from __future__ import annotations + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind +from v2.recipes.common import cached_text_encode + + +def _text_encode(instance, slots, request, ctx) -> None: + """FLUX.2 conditions on a single chat-template multi-layer embedding — no negative prompt + (embedded guidance / no CFG). Reuses the content-hash feature cache via ``cached_text_encode``.""" + slots["text_embeds"] = cached_text_encode(instance, request.prompt()) + + +def _vae_decode(instance, slots, request, ctx) -> None: + slots["image"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_flux2_program() -> Program: + return Program( + program_id="flux2.t2i.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", )), + ModelLoopNode("denoise", + loop_id="diffusion_denoise", + output_slot="denoise_out", + reads=("text_embeds", ), + writes=("denoise_out", )), + ComponentNode("vae_decode", fn=_vae_decode, reads=("denoise_out", ), writes=("image", )), + ], + output_artifacts={ + "image": "image", + "latents": "denoise_out" + }, + ).validate() diff --git a/v2/recipes/flux2/sampler.py b/v2/recipes/flux2/sampler.py new file mode 100644 index 000000000..9c5375b0c --- /dev/null +++ b/v2/recipes/flux2/sampler.py @@ -0,0 +1,58 @@ +"""FLUX.2 flow-match sigma schedule — the BFL empirical-mu grid (recipe-local; NOT in v2/loop/sampler.py). + +FLUX.2 does NOT use ``FlowShiftPolicy.build_schedule`` (the Wan/LTX resolution-bucket shift). It builds a +resolution-dependent ``mu`` from the *packed* image token count and feeds it, together with a custom +``sigmas = linspace(1, 1/N, N)`` grid, to diffusers' ``FlowMatchEulerDiscreteScheduler.set_timesteps`` +under ``use_dynamic_shifting=True``. This module reproduces that scheduler's numpy result exactly so the +CPU-toy loop is bit-faithful to the GPU pipeline. + +Faithful to: + * ``fastvideo/pipelines/basic/flux_2/flux_2_timestep_preparation.py:compute_empirical_mu`` (the BFL + ``sampling.compute_empirical_mu``), and the ``sigmas = np.linspace(1.0, 1/N, N)`` set there; + * ``fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py:set_timesteps`` with + ``use_dynamic_shifting=True`` (-> ``time_shift(mu, 1.0, sigmas)``, exponential ``time_shift_type``) + and the appended terminal ``0.0``. + +The exponential time-shift is ``sigma' = exp(mu) / (exp(mu) + (1/sigma - 1)**1.0)`` (the ``sigma=1.0`` +arg in the scheduler's ``time_shift(mu, 1.0, t)`` call collapses the ``(1/t - 1)**sigma`` exponent to 1). +""" +from __future__ import annotations + +import numpy as np + + +def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: + """Resolution-dependent ``mu`` for the FLUX.2 flow-match scheduler (BFL official ``compute_empirical_mu``). + + ``image_seq_len`` is the PACKED image token count ``latent_h * latent_w`` (half-spatial, see + ``flux2_latent_geometry``). Faithful copy of ``Flux2TimestepPreparationStage.compute_empirical_mu``. + """ + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + if image_seq_len > 4300: + return float(a2 * image_seq_len + b2) + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + return float(a * num_steps + b) + + +def _time_shift_exponential(mu: float, sigmas: np.ndarray) -> np.ndarray: + """Diffusers exponential time-shift with ``sigma=1.0`` (the scheduler's ``time_shift(mu, 1.0, t)``).""" + return np.exp(mu) / (np.exp(mu) + (1.0 / sigmas - 1.0)) + + +def flux2_sigmas(num_steps: int, image_seq_len: int) -> np.ndarray: + """The FLUX.2 σ schedule: ``num_steps+1`` points (the shifted grid + a terminal ``0.0``). + + 1. base grid ``sigmas = linspace(1.0, 1/N, N)`` (the custom grid the pipeline passes); + 2. resolution-dependent ``mu`` shift ``time_shift(mu, 1.0, sigmas)`` (use_dynamic_shifting=True); + 3. append the terminal ``0.0`` (FlowMatchEuler's ``invert_sigmas=False`` tail). + The loop integrates these pairwise via the deterministic flow-match Euler step (FLOW_MATCH_STEP). + """ + n = max(1, int(num_steps)) + base = np.linspace(1.0, 1.0 / n, n, dtype=np.float64) + mu = compute_empirical_mu(int(image_seq_len), n) + shifted = _time_shift_exponential(mu, base) + return np.concatenate([shifted, np.zeros(1, dtype=np.float64)]).astype(np.float64) diff --git a/v2/recipes/gen3c/__init__.py b/v2/recipes/gen3c/__init__.py new file mode 100644 index 000000000..2d8a4ea73 --- /dev/null +++ b/v2/recipes/gen3c/__init__.py @@ -0,0 +1,18 @@ +"""GEN3C-Cosmos-7B — camera-controlled video diffusion (EDM denoiser) ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``Gen3CDiT`` / ``Gen3CVAE`` / ``Gen3CT5Encoder`` in +``v2/platform/backends/torch_gen3c.py``) and a new ``Gen3CDenoiseLoop`` (a TRUE EDM / +``EDMEulerScheduler`` denoiser: Karras ρ=7 σ schedule, ``c_in/c_skip/c_out`` preconditioning, model +timestep ``c_noise = 0.25·log σ``, frame-replace conditioning + pose-zeroed-uncond CFG). The +camera/MoGe-depth/3D-cache conditioning that fills the pose buffer is documented as BRINGUP (a pre-loop +stage needing a CUDA renderer + request-API extension); the registered card is the degenerate t2v path. +The orchestrator adds the registry entry (HF id ``FastVideo/GEN3C-Cosmos-7B-Diffusers``). +""" +from __future__ import annotations + +from v2.recipes.gen3c.card import GEN3C_NEG, build_gen3c_card +from v2.recipes.gen3c.loop import Gen3CDenoiseLoop +from v2.recipes.gen3c.program import build_gen3c_program + +__all__ = ["build_gen3c_card", "build_gen3c_program", "Gen3CDenoiseLoop", "GEN3C_NEG"] diff --git a/v2/recipes/gen3c/card.py b/v2/recipes/gen3c/card.py new file mode 100644 index 000000000..b94ad1b56 --- /dev/null +++ b/v2/recipes/gen3c/card.py @@ -0,0 +1,146 @@ +"""GEN3C-Cosmos-7B ModelCard — camera-controlled video diffusion (EDM denoiser). + +GEN3C extends Cosmos-Predict2's EDM denoiser with camera/3D-cache conditioning. Architecture deltas vs +Cosmos (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.gen3c:Gen3CTransformer3DModel`` — an EDM denoiser whose forward concats + ``[latent(16) | input_mask(1) | pose_buffer(frame_buffer_max·32) | padding_mask(1)] = 82`` channels + INTERNALLY from separate kwargs. The adapter (``ComponentSpec.adapter`` -> ``Gen3CDiT``) returns the + raw network output and ``Gen3CDenoiseLoop`` does the EDM ``c_in/c_skip/c_out`` → x0 reconstruction + + x0-space CFG (pose-zeroed uncond) + frame-replace conditioning + the forced-EDMEulerScheduler σ + schedule (Karras ρ=7, σ_max=80 → σ_min=0.0002, terminal 0.0). Model timestep is ``c_noise = 0.25·log + σ`` (NOT σ·1000, NOT raw σ). + * VAE ``fastvideo.models.vaes.gen3c_tokenizer_vae:AutoencoderKLGen3CTokenizer`` — a JIT-backed + tokenizer that handles its OWN latent normalization (passthrough/internal mean_std buffers); the + ``Gen3CVAE`` adapter does NOT apply an external (z-mean)/std. 8× temporal (121→16 latent frames), 8× + spatial. + * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (T5-Large, 1024-dim) via ``Gen3CT5Encoder`` + (raw last_hidden_state, no Wan zero-pad — same as Cosmos). +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout). + +BRINGUP (see ``program.py`` / ``torch_gen3c.py``): the camera-trajectory / MoGe-depth / 3D-cache-render +conditioning that fills ``condition_video_pose`` is a separate pre-loop stage needing a CUDA point-cloud +renderer + request-API extension (image_path + trajectory fields). The registered card is the degenerate +**t2v** path (zero conditioning) — it CPU-verifies and the GPU loop runs end-to-end without the camera +stack; with conditioning it becomes image+camera → video. The JIT tokenizer loader seam also needs +GPU-box verification (``from_jit_tokenizer`` expects encoder.jit/decoder.jit/mean_std.pt). +""" +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.gen3c.loop import Gen3CDenoiseLoop +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +_GEN3C_DIT = "v2.platform.backends.torch_gen3c:Gen3CDiT" +_GEN3C_VAE = "v2.platform.backends.torch_gen3c:Gen3CVAE" +_GEN3C_T5 = "v2.platform.backends.torch_gen3c:Gen3CT5Encoder" + +# The GEN3C default negative prompt (kept LOCAL to this recipe — the v2/recipes/_prompts.py shared bank +# is off-limits for a self-contained port). Used at guidance_scale > 1 (the official_uncond_at_unity +# policy); the default guidance_scale=1.0 + cfg_behavior='legacy' runs NO uncond branch. +GEN3C_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.") + + +def build_gen3c_card(model_id: str = "gen3c-cosmos-7b", + *, + 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 Gen3CDenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + precision=precision, + expert=expert, + cost=cost, + sigma_max=80.0, + sigma_min=0.0002, + sigma_data=0.5, + sigma_conditional=0.001) + + components = { + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + adapter=_GEN3C_T5, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.gen3c_tokenizer_vae:AutoencoderKLGen3CTokenizer", + adapter=_GEN3C_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.gen3c:Gen3CTransformer3DModel", + adapter=_GEN3C_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="gen3c", + 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()), + # GEN3C defaults (the FORCED EDMEulerScheduler config + the registered preset): 35 steps, + # guidance 1.0 (legacy cfg -> no uncond branch), 704x1280, 121 frames @ 24fps. + sampling_defaults=sampling_defaults or SamplingDefaults( + num_steps=35, guidance_scale=1.0, height=704, width=1280, num_frames=121, fps=24, + negative_prompt=GEN3C_NEG), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card diff --git a/v2/recipes/gen3c/loop.py b/v2/recipes/gen3c/loop.py new file mode 100644 index 000000000..fef800843 --- /dev/null +++ b/v2/recipes/gen3c/loop.py @@ -0,0 +1,244 @@ +"""Gen3CDenoiseLoop — a TRUE EDM (EDMEulerScheduler) denoiser with camera/3D-cache conditioning hooks. + +GEN3C is Cosmos-Predict2's EDM denoiser extended with camera-controlled video conditioning. Unlike the +Cosmos recipe (which folds EDM preconditioning into a *Karras* schedule), GEN3C's pipeline FORCES a +diffusers ``EDMEulerScheduler(sigma_max=80, sigma_min=0.0002, sigma_data=0.5)`` at runtime +(``Gen3CPipeline.initialize_pipeline``), so this loop reproduces that scheduler's exact σ schedule and +step math (see ``build_edm_euler_sigmas`` in ``sampler.py``): + + * σ schedule: ``EDMEulerScheduler``'s defaults are ``sigma_schedule='karras'`` (ρ=7) and + ``final_sigmas_type='zero'`` — a Karras ρ ramp ``σ_max → σ_min`` then a single appended terminal + ``0.0`` (verified bit-identical to the diffusers scheduler); the loop integrates the ``num_steps+1`` + points pairwise. Latent init = ``randn·√(σ_max²+σ_data²)`` (``init_noise_sigma``). + * per step ``i`` (σ = ``sigmas[i]``, gamma=0 since s_churn=0 → ``σ_hat = σ``): + - frame-replace BEFORE the forward (``_augment_noise_with_latent``): augment the conditioning + latent with ``N(0,1)·σ_cond``, EDM-precondition, reverse ``c_in`` to put it back in the + unscaled latent frame, and write it into ``x`` at ``cond_indicator`` positions; + - ``model_input = x·c_in`` (``scale_model_input``); ``c_in = 1/√(σ²+σ_data²)``; + - DiT(model_input, text, σ, cond=) → raw output; the adapter embeds the model + timestep ``c_noise = 0.25·log σ`` and concats the conditioning kwargs INTERNALLY (82 ch); + - CFG (when enabled): cond branch with the REAL pose buffer; uncond branch with a ZEROED pose + buffer + neg text; ``pred = cond + scale·(cond − uncond)``. NOTE default guidance_scale=1.0 + + ``cfg_behavior='legacy'`` ⇒ NO uncond branch by default; + - ``x0 = c_skip·x + c_out·pred`` (``precondition_outputs``); ``c_skip = σ_data²/(σ²+σ_data²)``, + ``c_out = σ·σ_data/√(σ²+σ_data²)``; + - frame-replace on the OUTPUT: ``x0 = ind·cond_latent + (1−ind)·x0`` (anchors the first frame); + - Euler: ``v = (x − x0)/σ``; ``x_next = x + (σ_next − σ)·v`` (the platform ``FLOW_MATCH_STEP``, + identical to ``EDMEulerScheduler.step``'s ODE-derivative update). + +Camera / MoGe-depth / 3D-cache conditioning fills ``condition_video_pose`` / ``conditioning_latents`` / +``cond_indicator`` from program slots; ALL are ``None`` for the registered degenerate **t2v** path → the +frame-replace branches are inert and the DiT's internal pose concat is zeroed, exactly as the fastvideo +pipeline degrades when no image/camera is given. Faithful port of ``Gen3CDenoisingStage.forward`` + +``Gen3CLatentPreparationStage`` (fastvideo/pipelines/stages/gen3c_stages.py). + +The DiT call is ``dit(model_input, text_embed, sigma, cond=bundle)`` — compatible with ``ToyDiT`` (which +accepts and ignores ``cond``), so the loop CPU-verifies on the toy backend. +""" +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.platform import FLOW_MATCH_STEP +from v2.recipes.gen3c.sampler import build_edm_euler_sigmas, edm_init_sigma +from v2.recipes.wan21.loop import latent_shape + +GEN3C_LATENT_CHANNELS = 16 # config.in_channels (VAE latent); the 1+pose+1 channels are added in the DiT +GEN3C_TEMPORAL_RATIO = 8 # AutoencoderKLGen3CTokenizer target temporal compression (121 -> 16 latent frames) +GEN3C_SPATIAL_RATIO = 8 + + +class Gen3CDenoiseLoop: + + def __init__(self, + *, + loop_id, + cfg, + precision, + expert, + cost, + sigma_max: float = 80.0, + sigma_min: float = 0.0002, + sigma_data: float = 0.5, + sigma_conditional: float = 0.001, + latent_channels: int = GEN3C_LATENT_CHANNELS, + spatial_ratio: int = GEN3C_SPATIAL_RATIO, + temporal_ratio: int = GEN3C_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.sigma_conditional = sigma_conditional + 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_edm_euler_sigmas(req.diffusion.num_steps, sigma_max=self.sigma_max, sigma_min=self.sigma_min) + shape = latent_shape(req, + model, + channels=self.latent_channels, + spatial_ratio=self.spatial_ratio, + temporal_ratio=self.temporal_ratio) + # EDM latent init: randn · init_noise_sigma where init_noise_sigma = sqrt(σ_max² + σ_data²). + x = (rng.standard_normal(shape) * edm_init_sigma(self.sigma_max, self.sigma_data)).astype("float32") + 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) for s in sig]) # the adapter derives c_noise = 0.25·log σ from σ + 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")) + # Camera / 3D-cache conditioning written by the conditioning program node. ALL None for the + # degenerate t2v path → the frame-replace branches below are inert and the DiT's pose concat is + # zeroed (matching the fastvideo pipeline when no image/camera is given). + 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.scratch["condition_video_pose"] = ctx.slots.get("condition_video_pose") + st.scratch["condition_video_input_mask"] = ctx.slots.get("condition_video_input_mask") + # Pixel H,W for the DiT's padding mask (resized internally). Carried so the adapter builds it. + st.scratch["pixel_height"] = int(req.diffusion.height) + st.scratch["pixel_width"] = int(req.diffusion.width) + 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] # the raw σ; the adapter maps σ → c_noise = 0.25·log σ + 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.sigma_conditional) + pose = st.scratch.get("condition_video_pose") + input_mask = st.scratch.get("condition_video_input_mask") + pix_h, pix_w = st.scratch.get("pixel_height"), st.scratch.get("pixel_width") + + def _bundle(pose_buf: Any) -> dict | None: + # The GEN3C conditioning bundle the adapter concats internally (None for pure t2v -> the DiT + # builds zeros). ``pose_buf`` is zeroed on the CFG uncond branch. + if pose is None and input_mask is None: + return None + return { + "condition_video_pose": pose_buf, + "condition_video_input_mask": input_mask, + "condition_video_augment_sigma": aug, + "height": pix_h, + "width": pix_w, + } + + def _x0(dit: Any, latent_in: np.ndarray, text_embed: Any, ind: Any, pose_buf: 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 + # active_indicator is zeroed once σ drops to/below the conditioning-augment σ (the fastvideo + # ``_augment_noise_with_latent`` guard: ``aug >= σ`` -> stop replacing). + 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: # frame-replace on the INPUT (inert for t2v) + c_in_aug = 1.0 / (aug**2 + sd**2)**0.5 + cn = st.rng.standard_normal(lat.shape).astype("float32") + # augment_latent = (cond + noise·aug); precondition (·c_in_aug); reverse the step's c_in. + 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") # scale_model_input + np_pred = np.asarray(dit(model_input, text_embed, s, cond=_bundle(pose_buf)), dtype="float32") + x0 = c_skip * x + c_out * np_pred # precondition_outputs (x_t is the UNSCALED x) + if cur_ind is not None and cond_lat is not None: # frame-replace on the OUTPUT (anchor frame-0) + 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(dit, x, pe, cond_ind, pose) + if do_cfg: + # CFG zeros the POSE buffer (not just text) on the uncond branch. + uncond_pose = (np.zeros_like(pose) if pose is not None else None) + uncond_x0 = _x0(dit, x, ne, uncond_ind, uncond_pose) + 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"gen3c.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/gen3c/program.py b/v2/recipes/gen3c/program.py new file mode 100644 index 000000000..17f514ecd --- /dev/null +++ b/v2/recipes/gen3c/program.py @@ -0,0 +1,47 @@ +"""GEN3C program: text_encode → diffusion_denoise (EDM) → vae_decode. + +Same inline shape as the Cosmos t2v program — the EDM specifics + the camera/3D-cache frame-replace +conditioning live entirely in ``Gen3CDenoiseLoop`` and the ``Gen3CDiT`` adapter, so the registered +(degenerate t2v) node graph is unchanged. + +BRINGUP — the camera-conditioning pre-stage (the GEN3C innovation, no v2 equivalent yet): a node placed +BEFORE ``denoise`` would (porting ``Gen3CConditioningStage`` + ``Gen3CLatentPreparationStage``): + 1. load the input image; predict metric depth via MoGe (``Ruicheng/moge-vitl``); + 2. build a ``Cache3DBuffer`` point cloud; ``generate_camera_trajectory(trajectory_type, + movement_distance, camera_rotation)``; ``cache.render_cache(...)`` -> warped frames + masks; + 3. VAE-encode the warps into ``condition_video_pose`` [frame_buffer_max·32, T_lat, H_lat, W_lat], and + VAE-encode the input image into ``conditioning_latents`` (frame-0 anchor); build + ``condition_video_input_mask`` / ``cond_indicator`` (frame-0 = 1). +The loop already READS those slots (``conditioning_latents`` / ``cond_indicator`` / ``condition_video_pose`` +/ ``condition_video_input_mask``); writing them needs (a) a CUDA point-cloud rasterizer dependency and +(b) a request-API extension carrying ``image_path`` + the camera-trajectory fields. Until then the slots +are unset (None) and the loop runs pure t2v (the DiT's internal pose concat is zeroed). +""" +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_gen3c_program() -> Program: + return Program( + program_id="gen3c.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/recipes/gen3c/sampler.py b/v2/recipes/gen3c/sampler.py new file mode 100644 index 000000000..17ad242ef --- /dev/null +++ b/v2/recipes/gen3c/sampler.py @@ -0,0 +1,37 @@ +"""GEN3C EDM sampler helpers (kept local to the recipe — NOT in ``v2/loop/sampler.py``). + +GEN3C's pipeline FORCES a diffusers ``EDMEulerScheduler(sigma_max=80, sigma_min=0.0002, sigma_data=0.5)`` +at runtime (``Gen3CPipeline.initialize_pipeline``), overriding whatever the converted model_index points +at. That scheduler's defaults are ``sigma_schedule='karras'`` (ρ=7) and ``final_sigmas_type='zero'``, so: + + * the ramp is the Karras ρ-interpolation ``σ_max → σ_min`` (verified bit-identical to + ``v2.loop.sampler.build_karras_sigmas`` on the ramp), and + * the appended terminal σ is **0.0** (``final_sigmas_type='zero'``), NOT ``σ_min``. This is the one + difference from the Cosmos recipe (Cosmos clamps the terminal to σ_min); GEN3C's last Euler step + integrates from ``σ[-2]=σ_min`` to ``0.0`` (the velocity still divides by ``σ[-2]>0``, so no NaN). + +``edm_init_sigma`` is the EDMEulerScheduler ``init_noise_sigma = √(σ_max² + σ_data²)`` used to scale the +initial Gaussian latent (``Gen3CLatentPreparationStage`` multiplies ``randn`` by it). +""" +from __future__ import annotations + +import numpy as np + + +def build_edm_euler_sigmas(num_steps: int, + sigma_max: float = 80.0, + sigma_min: float = 0.0002, + rho: float = 7.0) -> np.ndarray: + """The diffusers ``EDMEulerScheduler`` σ schedule for GEN3C: a length-``num_steps`` Karras ρ ramp + ``σ_max → σ_min`` then a single appended terminal ``0.0`` (``final_sigmas_type='zero'``) — 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 (Karras ρ) + return np.concatenate([sigmas, np.zeros(1)]).astype(np.float64) # + the 0.0 terminal (final='zero') + + +def edm_init_sigma(sigma_max: float = 80.0, sigma_data: float = 0.5) -> float: + """``EDMEulerScheduler.init_noise_sigma = √(σ_max² + σ_data²)`` — the initial-latent noise scale.""" + return float((sigma_max**2 + sigma_data**2)**0.5) diff --git a/v2/recipes/hunyuan_video/__init__.py b/v2/recipes/hunyuan_video/__init__.py new file mode 100644 index 000000000..9f7744c1f --- /dev/null +++ b/v2/recipes/hunyuan_video/__init__.py @@ -0,0 +1,23 @@ +"""HunyuanVideo — rectified flow-match video DiT ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``HunyuanVideoDiT`` / ``HunyuanVideoVAE`` / ``HunyuanVideoLlamaEncoder`` / +``HunyuanVideoCLIPEncoder`` in ``v2/platform/backends/torch_hunyuan_video.py``) and reuses the canonical +``WanDenoiseLoop`` flow-match math via a thin ``HunyuanDenoiseLoop`` (which only threads the CLIP-pooled +vector through Wan's ``context=`` channel). The new work vs Wan is the dual text-encoder marshalling +(LLaMA per-token sequence + CLIP pooled global vector packed into the DiT's 2-element +``encoder_hidden_states``) and the Hunyuan VAE scalar-``scaling_factor`` normalization. Reuses +``stamp_wan21_checkpoints``. Registered in ``v2/registry.py`` by the orchestrator. +""" +from __future__ import annotations + +from v2.recipes.hunyuan_video.card import build_fast_hunyuan_video_card, build_hunyuan_video_card +from v2.recipes.hunyuan_video.loop import HunyuanDenoiseLoop +from v2.recipes.hunyuan_video.program import build_hunyuan_video_program + +__all__ = [ + "build_hunyuan_video_card", + "build_fast_hunyuan_video_card", + "build_hunyuan_video_program", + "HunyuanDenoiseLoop", +] diff --git a/v2/recipes/hunyuan_video/card.py b/v2/recipes/hunyuan_video/card.py new file mode 100644 index 000000000..75d599195 --- /dev/null +++ b/v2/recipes/hunyuan_video/card.py @@ -0,0 +1,172 @@ +"""HunyuanVideo ModelCard — text->video (the registered preset). Flow-match rectified-flow video DiT. + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.hunyuanvideo:HunyuanVideoTransformer3DModel`` — velocity flow-match (like + Wan), but ``encoder_hidden_states`` is a 2-element list ``[llama_hidden[B,L,4096], clip_pooled[B,768]]``. + The adapter (``ComponentSpec.adapter`` -> ``HunyuanVideoDiT``) assembles that pair internally; the loop + reuses ``WanDenoiseLoop`` math verbatim (timestep ``sigma*1000``, scalar per-batch, bare velocity out). + * VAE ``fastvideo.models.vaes.hunyuanvae:AutoencoderKLHunyuanVideo`` — scalar ``scaling_factor`` (0.476986) + normalization, NOT Wan's per-channel mean/std -> the ``HunyuanVideoVAE`` adapter. + * Text DUAL encoder: ``fastvideo.models.encoders.llama:LlamaModel`` (per-token sequence; prompt template + + hidden-state-skip + crop) via ``HunyuanVideoLlamaEncoder``, and + ``fastvideo.models.encoders.clip:CLIPTextModel`` (768-dim pooled) via ``HunyuanVideoCLIPEncoder``. + +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout); HunyuanVideo +additionally ships ``text_encoder_2/`` (CLIP) + ``tokenizer_2/`` — the stamp helper points each declared +component (incl. ``text_encoder_2``) at its subfolder (see the BRINGUP note in ``card.py`` and the CLIP +adapter for the ``tokenizer_2`` wiring caveat). + +Defaults (fastvideo HUNYUAN_T2V preset): 50 steps, guidance_scale 1.0 (CFG effectively off), +720x1280, 125 frames @ 24fps, flow_shift 7.0. A FastHunyuan variant ships shift 17, 6 steps. +""" +from __future__ import annotations + +import os + +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.hunyuan_video.loop import HunyuanDenoiseLoop +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +_HUNYUAN_DIT = "v2.platform.backends.torch_hunyuan_video:HunyuanVideoDiT" +_HUNYUAN_VAE = "v2.platform.backends.torch_hunyuan_video:HunyuanVideoVAE" +_HUNYUAN_LLAMA = "v2.platform.backends.torch_hunyuan_video:HunyuanVideoLlamaEncoder" +_HUNYUAN_CLIP = "v2.platform.backends.torch_hunyuan_video:HunyuanVideoCLIPEncoder" + +# Base HunyuanVideo runs guidance_scale=1.0 (CFG effectively off), so it has no published negative prompt; +# a LOCAL module constant so the recipe owns its own data (empty string = no negative conditioning). +HUNYUAN_NEG = "" + + +def build_hunyuan_video_card(model_id: str = "hunyuanvideo", + *, + flow_shift: float = 7.0, + num_steps: int = 50, + 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() + # No resolution-bucket lookup: HunyuanVideo uses a single flow_shift (7.0 base / 17.0 FastHunyuan). + flow = FlowShiftPolicy(shift=flow_shift) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory() -> HunyuanDenoiseLoop: + return HunyuanDenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost) # 16/8/4 default geometry + + components = { + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.llama:LlamaModel", + adapter=_HUNYUAN_LLAMA, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "text_encoder_2": + ComponentSpec(component_id="text_encoder_2", + kind="text_encoder", + load_id="fastvideo.models.encoders.clip:CLIPTextModel", + adapter=_HUNYUAN_CLIP, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.hunyuanvae:AutoencoderKLHunyuanVideo", + adapter=_HUNYUAN_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.hunyuanvideo:HunyuanVideoTransformer3DModel", + adapter=_HUNYUAN_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="hunyuan", + 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=num_steps, + guidance_scale=1.0, + height=720, + width=1280, + num_frames=125, + fps=24, + negative_prompt=HUNYUAN_NEG, + shift=flow_shift), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + # HunyuanVideo's CLIP secondary encoder lives in text_encoder_2/ (the Wan stamp covers it because + # text_encoder_2 is in its subfolder superset); its tokenizer is tokenizer_2/ (BRINGUP: see the CLIP + # adapter — the shared _make_text_encoder loads tokenizer/, so GPU-verify must wire tokenizer_2). + if "text_encoder_2" in card.components and not card.components["text_encoder_2"].checkpoint: + root = os.path.dirname(os.path.normpath(card.components["transformer"].checkpoint)) + card.components["text_encoder_2"].checkpoint = os.path.join(root, "text_encoder_2") + return card + + +def build_fast_hunyuan_video_card(model_id: str = "fasthunyuan", *, checkpoint_root: str | None = None) -> ModelCard: + """FastHunyuan (FastVideo/FastHunyuan-diffusers): same arch, distilled to 6 steps with flow_shift 17.0. + Reuses the base card / adapters unchanged — only the schedule (shift + step count) differs.""" + return build_hunyuan_video_card(model_id=model_id, + flow_shift=17.0, + num_steps=6, + checkpoint_root=checkpoint_root, + sampling_defaults=SamplingDefaults(num_steps=6, + guidance_scale=1.0, + height=720, + width=1280, + num_frames=125, + fps=24, + negative_prompt=HUNYUAN_NEG, + shift=17.0)) diff --git a/v2/recipes/hunyuan_video/loop.py b/v2/recipes/hunyuan_video/loop.py new file mode 100644 index 000000000..41067794d --- /dev/null +++ b/v2/recipes/hunyuan_video/loop.py @@ -0,0 +1,57 @@ +"""HunyuanDenoiseLoop — a thin subclass of the canonical ``WanDenoiseLoop`` (flow-match Euler). + +HunyuanVideo is a *rectified flow-match* model: the DiT predicts a velocity in unpatchified +``[B, 16, T, H, W]`` latent space, the model timestep is the scalar ``sigma * 1000`` (a per-batch +``LongTensor`` consumed by the internal ``time_in`` ``TimestepEmbedder``), and the deterministic Euler +update is ``x_next = x + (sigma_next - sigma)*v`` — exactly the math ``WanDenoiseLoop`` already runs. The +latent geometry (16 channels, 8x spatial, 4x temporal) is the Wan default too, so we reuse the loop +verbatim and only carry a Hunyuan-specific docstring + the dual-text-encoder marshalling delta below. + +The ONE delta vs Wan is conditioning shape: the Hunyuan DiT consumes TWO text signals — +``encoder_hidden_states = [llama_hidden[B, L, 4096], clip_pooled[B, 768]]`` (a per-token LLaMA sequence +plus a single CLIP-pooled global vector). To keep the loop's ``dit(latent, text_embed, sigma)`` call +identical on the CPU toy and the GPU backend (the cosmos2 discipline: the adapter assembles +arch-specifics INTERNALLY), the program writes the LLaMA sequence into the ``text_embeds`` slot (the +primary ``text_embed`` the loop already threads) and the CLIP-pooled vector into ``text_pooled``. This +loop reuses Wan's existing ``context=`` channel (Wan uses it for the i2v CLIP-vision embeds) to carry +the CLIP-pooled vector to the DiT adapter — so ``WanDiT``-style ``dit(x, text_embed, sigma, context=...)`` +is unchanged, ``ToyDiT`` (which means over both ``text_embed`` and ``context``) runs on CPU, and the real +``HunyuanVideoDiT`` adapter packs ``[text_embed, context]`` into the 2-element ``encoder_hidden_states``. + +guidance_scale defaults to 1.0 for base HunyuanVideo → ``ClassicCFG.combine`` collapses to the single +``cond`` branch (``uncond + 1*(cond-uncond) == cond``), i.e. CFG is effectively off, matching the +fastvideo pipeline. (``embedded_cfg_scale``/``guidance`` is a no-op on the base model whose +``guidance_in`` is ``None``; see ``torch_hunyuan_video.HunyuanVideoDiT`` for that BRINGUP note.) + +Faithful to ``fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py`` (flow-match DenoisingStage) and +``fastvideo/models/dits/hunyuanvideo.py:HunyuanVideoTransformer3DModel.forward`` (the list-split of +``encoder_hidden_states`` into ``txt`` / ``text_states_2``). +""" +from __future__ import annotations + +from v2.loop.contracts import LoopState +from v2.recipes.wan21.loop import WanDenoiseLoop + +# AutoencoderKLHunyuanVideo compression matches Wan2.1: z=16, 4x temporal, 8x spatial. +HUNYUAN_LATENT_CHANNELS = 16 +HUNYUAN_TEMPORAL_RATIO = 4 +HUNYUAN_SPATIAL_RATIO = 8 + + +class HunyuanDenoiseLoop(WanDenoiseLoop): + """Wan flow-match loop with HunyuanVideo's dual-text-encoder marshalling. + + Reuses ``WanDenoiseLoop.next``/``advance``/``finalize`` unchanged (flow-match Euler, velocity output, + timestep ``sigma*1000``, scalar per-batch timestep). Overrides only ``init`` to thread the CLIP-pooled + global vector (slot ``text_pooled``) through Wan's existing ``context=`` channel so the DiT adapter can + reassemble the 2-element ``encoder_hidden_states`` the Hunyuan forward expects. + """ + + def init(self, req, model, ctx) -> LoopState: + st = super().init(req, model, ctx) + # Wan's ``context`` channel (``i2v_img_embeds`` -> dit ``context=``) carries the CLIP-pooled vector + # for Hunyuan. None for any caller that did not run the dual-encode node -> the DiT adapter falls + # back to a single-tensor encoder_hidden_states (still valid: the Hunyuan forward also accepts a + # bare tensor whose token 0 is the pooled vector — but the dual-slot path is the faithful one). + st.scratch["i2v_img_embeds"] = ctx.slots.get("text_pooled") + return st diff --git a/v2/recipes/hunyuan_video/program.py b/v2/recipes/hunyuan_video/program.py new file mode 100644 index 000000000..0237c40bd --- /dev/null +++ b/v2/recipes/hunyuan_video/program.py @@ -0,0 +1,80 @@ +"""HunyuanVideo t2v program: hunyuan_text_encode -> diffusion_denoise (flow-match) -> vae_decode. + +Same inline shape as the Wan/cosmos2 t2v programs, with ONE delta: HunyuanVideo has TWO text encoders +(a LLaMA per-token sequence + a CLIP-pooled global vector), so the text-encode node writes BOTH a +``text_embeds`` slot (the LLaMA sequence, the loop's primary ``text_embed``) and a ``text_pooled`` slot +(the CLIP vector, threaded to the DiT via the loop's ``context=`` channel). The dual-encode node uses the +content-hash feature cache (like ``recipes.common.text_encode_node_fn``) so a repeated prompt encodes once. + +On the CPU toy backend both components are ``ToyTextEncoder`` (single arrays), so ``ToyDiT`` — which means +over both ``text_embed`` and ``context`` — runs end-to-end; on the GPU backend the real LLaMA/CLIP adapters +produce the sequence + pooled pair the ``HunyuanVideoDiT`` adapter reassembles into the 2-element +``encoder_hidden_states`` the Hunyuan forward expects. +""" +from __future__ import annotations + +from typing import Any + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind +from v2.recipes.common import cached_text_encode + + +def _encode_pooled(instance: Any, text: str) -> Any: + """CLIP secondary encode (pooled global vector). Reuses the content-hash feature cache discipline of + ``cached_text_encode`` but on the ``text_encoder_2`` component (a distinct cache partition by + component_id). Falls back to the plain encode when the instance has no feature cache.""" + enc = instance.component("text_encoder_2") + cache = instance.caches + if cache is not None and cache.has("feature"): + from v2.cache.keys import CacheKey, content_hash + key = CacheKey(model_id=instance.card.model_id, + component_id="text_encoder_2", + weights_version=instance.version_of("text_encoder_2"), + adapter_versions=CacheKey.adapters(instance.adapter_versions), + precision=instance.card.precision.dtype_for("text_encoder_2"), + input_hashes=(("text", content_hash(text)), )) + hit = cache.pool("feature").get(key) + if hit is not None: + return hit + emb = enc.encode(text) + cache.pool("feature").put(key, emb) + return emb + return enc.encode(text) + + +def _hunyuan_text_encode(instance, slots, request, ctx) -> None: + """ComponentNode fn: prompt -> (LLaMA sequence in ``text_embeds`` + CLIP pooled in ``text_pooled``), + plus the negative-prompt counterparts (used only when guidance_scale != 1.0; base HunyuanVideo runs + guidance_scale=1.0 so the uncond branch collapses out in ClassicCFG.combine).""" + prompt = request.prompt() + neg = request.diffusion.negative_prompt + slots["text_embeds"] = cached_text_encode(instance, prompt) # LLaMA per-token sequence + slots["text_pooled"] = _encode_pooled(instance, prompt) # CLIP pooled global vector + slots["neg_text_embeds"] = cached_text_encode(instance, neg) + slots["neg_text_pooled"] = _encode_pooled(instance, neg) + + +def _vae_decode(instance, slots, request, ctx) -> None: + slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_hunyuan_video_program() -> Program: + return Program( + program_id="hunyuan_video.t2v.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("text_encode", + fn=_hunyuan_text_encode, + writes=("text_embeds", "text_pooled", "neg_text_embeds", "neg_text_pooled")), + ModelLoopNode("denoise", + loop_id="diffusion_denoise", + output_slot="denoise_out", + reads=("text_embeds", "text_pooled", "neg_text_embeds", "neg_text_pooled"), + 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/recipes/hunyuan_video15/__init__.py b/v2/recipes/hunyuan_video15/__init__.py new file mode 100644 index 000000000..787f7b392 --- /dev/null +++ b/v2/recipes/hunyuan_video15/__init__.py @@ -0,0 +1,25 @@ +"""HunyuanVideo 1.5 (T2V) — rectified-flow video DiT ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``HunyuanVideo15DiT``/``HunyuanVideo15VAE``/``HunyuanVideo15QwenEncoder``/ +``HunyuanVideo15ByT5Encoder`` in ``v2/platform/backends/torch_hunyuan_video15.py``) and a +``HunyuanVideo15DenoiseLoop`` (a thin ``WanDenoiseLoop`` subclass with z=32/16×/4× geometry). The +flow-match loop math is unchanged from Wan; the arch deltas (two text embeds, an image-embed list, the +33-channel i2v cond concat, the scalar-scaling VAE) live in the adapters. Registered in ``v2/registry.py`` +(by the orchestrator) for ``hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v``. +""" +from __future__ import annotations + +from v2.recipes.hunyuan_video15.card import ( + build_hunyuan_video15_720p_card, + build_hunyuan_video15_card, +) +from v2.recipes.hunyuan_video15.loop import HunyuanVideo15DenoiseLoop +from v2.recipes.hunyuan_video15.program import build_hunyuan_video15_program + +__all__ = [ + "build_hunyuan_video15_card", + "build_hunyuan_video15_720p_card", + "build_hunyuan_video15_program", + "HunyuanVideo15DenoiseLoop", +] diff --git a/v2/recipes/hunyuan_video15/card.py b/v2/recipes/hunyuan_video15/card.py new file mode 100644 index 000000000..0f10e1238 --- /dev/null +++ b/v2/recipes/hunyuan_video15/card.py @@ -0,0 +1,181 @@ +"""HunyuanVideo 1.5 ModelCard — text→video (the registered 480p t2v preset; i2v + SR are later +capabilities the loop/adapter thread but the t2v program does not yet exercise on CPU). + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.hunyuanvideo15:HunyuanVideo15Transformer3DModel`` — a rectified-flow + velocity network (54 dual-stream MMDoubleStream blocks, hidden 2048, in/out 65/32). Its forward takes + a LIST of two text embeds (Qwen2.5-VL + ByT5) and a LIST image-embed; the new ``HunyuanVideo15DiT`` + adapter marshals those, so ``HunyuanVideo15DenoiseLoop`` (a thin ``WanDenoiseLoop`` subclass with + z=32/16×/4× geometry) drives it with the unchanged flow-match math. + * VAE ``fastvideo.models.vaes.hunyuan15vae`` — causal-3D, z=32, 16× spatial / 4× temporal, normalized + by a SCALAR ``scaling_factor=1.03682`` -> the ``HunyuanVideo15VAE`` adapter (NOT Wan's mean/std). + * Text encoders: ``text_encoder`` = Qwen2.5-VL (3584-d, ``hidden_states[-3]`` cropped of 108 template + tokens) via ``HunyuanVideo15QwenEncoder``; ``text_encoder_2`` = ByT5/Glyph (1472-d, quoted-glyph text, + may be empty) via ``HunyuanVideo15ByT5Encoder``. +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder/text_encoder_2 subfolders). + +Defaults (configs/pipelines + presets): 480p t2v -> 50 steps, guidance 6.0, 480×848, 121 frames, 24 fps, +flow_shift 5 (720p uses 9), empty negative prompt, explicit ``linspace(1,0,n+1)[:-1]`` sigma schedule. +""" +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.hunyuan_video15.loop import ( + HUNYUAN15_LATENT_CHANNELS, + HUNYUAN15_SPATIAL_RATIO, + HUNYUAN15_TEMPORAL_RATIO, + HunyuanVideo15DenoiseLoop, +) +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# HunyuanVideo 1.5 ships with an EMPTY negative prompt (configs/.../presets.py). Kept as a LOCAL constant +# (the recipe owns its prompts; not added to the shared v2/recipes/_prompts.py). +HUNYUAN15_NEG = "" + +_HY15_DIT = "v2.platform.backends.torch_hunyuan_video15:HunyuanVideo15DiT" +_HY15_VAE = "v2.platform.backends.torch_hunyuan_video15:HunyuanVideo15VAE" +_HY15_QWEN = "v2.platform.backends.torch_hunyuan_video15:HunyuanVideo15QwenEncoder" +_HY15_BYT5 = "v2.platform.backends.torch_hunyuan_video15:HunyuanVideo15ByT5Encoder" + +# 480p t2v explicit sigma schedule = linspace(1, 0, num_steps+1)[:-1] (presets.py:_sigmas). +_DEFAULT_STEPS = 50 + + +def _linspace_sigmas(num_steps: int) -> tuple[float, ...]: + """Faithful to ``presets.py:_sigmas`` — ``np.linspace(1,0,n+1)[:-1]`` (drops the terminal 0). The + loop's ``build_schedule`` consumes this verbatim when present (FlowShiftPolicy passes it through).""" + return tuple((1.0 - i / num_steps) for i in range(num_steps)) + + +def build_hunyuan_video15_card(model_id: str = "hunyuan-video-1.5-t2v-480p", + *, + flow_shift: float = 5.0, + height: int = 480, + width: int = 848, + num_steps: int = _DEFAULT_STEPS, + guidance_scale: float = 6.0, + 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() + # 480p flow_shift 5, 720p 9 (configs/pipelines/hunyuan15.py). The preset also hands an explicit sigma + # schedule via SamplingDefaults.sigmas, which the loop prefers when present. + flow = FlowShiftPolicy(shift=flow_shift, bucket_lookup={480 * 848: 5.0, 720 * 1280: 9.0}) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return HunyuanVideo15DenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost, + latent_channels=HUNYUAN15_LATENT_CHANNELS, + spatial_ratio=HUNYUAN15_SPATIAL_RATIO, + temporal_ratio=HUNYUAN15_TEMPORAL_RATIO) + + components = { + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.qwen2_5:Qwen2_5_VLTextModel", + adapter=_HY15_QWEN, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "text_encoder_2": + ComponentSpec(component_id="text_encoder_2", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + adapter=_HY15_BYT5, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.hunyuan15vae:AutoencoderKLHunyuanVideo15", + adapter=_HY15_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.hunyuanvideo15:HunyuanVideo15Transformer3DModel", + adapter=_HY15_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="hunyuan", + 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"), + # sequence_model_parallel_shard inside the DiT forward is a no-op at sp=1; default single() until + # SP is GPU-validated (BRINGUP risk G). + parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()), + sampling_defaults=sampling_defaults or SamplingDefaults(num_steps=num_steps, + guidance_scale=guidance_scale, + height=height, + width=width, + num_frames=121, + fps=24, + negative_prompt=HUNYUAN15_NEG, + shift=flow_shift, + sigmas=_linspace_sigmas(num_steps)), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card + + +def build_hunyuan_video15_720p_card(model_id: str = "hunyuan-video-1.5-t2v-720p", + *, + checkpoint_root: str | None = None) -> ModelCard: + """HunyuanVideo 1.5 720p t2v — same components/adapters; only flow_shift (9), resolution, and the + per-model defaults differ (configs/pipelines/hunyuan15.py:Hunyuan15T2V720PConfig).""" + return build_hunyuan_video15_card(model_id=model_id, + flow_shift=9.0, + height=720, + width=1280, + checkpoint_root=checkpoint_root) diff --git a/v2/recipes/hunyuan_video15/loop.py b/v2/recipes/hunyuan_video15/loop.py new file mode 100644 index 000000000..fba740d6a --- /dev/null +++ b/v2/recipes/hunyuan_video15/loop.py @@ -0,0 +1,57 @@ +"""HunyuanVideo15DenoiseLoop — the canonical flow-match denoise loop, parameterized for HunyuanVideo 1.5. + +The loop *math* is identical to Wan's (``WanDenoiseLoop``): N deterministic flow-match Euler steps over +the full latent, ``ClassicCFG`` combine, ``timestep = sigma*1000``, ``x_next = x + (σ_next-σ)·v``. This is +faithful to ``fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py`` — ``scale_model_input`` +is the identity for flow-match, ``step`` is ``prev_sample = sample + (sigma_next - sigma)*model_output``, +and HunyuanVideo 1.5's DiT predicts velocity directly (no x0 reconstruction). The latent is initialized at +``randn·σ[0]`` with ``σ[0]=1.0`` (the preset's ``linspace(1,0,n+1)[:-1]`` schedule). + +The ONLY architectural delta lives in the ``HunyuanVideo15DiT`` torch adapter (two text embeds + an +image-embed list + the 33-channel i2v cond concat). Because ``WanDenoiseLoop`` already threads +``context=`` (image embeds) and ``cond=`` (the conditioning latent) into the adapter, the cleanest port +is to subclass it: change only the latent geometry (z=32, 16× spatial, 4× temporal) and the slot names +the loop reads for the i2v signal. For t2v those slots are absent (``None``) → the dit-call degenerates +to ``dit(latent, text_embed, sigma)`` exactly like Wan t2v, and the toy backend's ``ToyDiT`` runs it. + +BRINGUP: i2v conditioning (``context``=zero image embeds, ``cond``=33ch first-frame latent) is threaded +but the program currently writes neither slot for the registered t2v preset, so the loop is pure t2v on +CPU; the i2v VAE-encode node + image-latent slot are documented hooks (see ``program.py``). +""" +from __future__ import annotations + +from v2.recipes.wan21.loop import WanDenoiseLoop + +# HunyuanVideo 1.5 causal-3D VAE compression (configs/models/vaes/hunyuan15vae.py): z=32, 16× spatial, +# 4× temporal. The DiT's noisy-latent in_channels is 32 (out_channels); i2v reaches 65 by concatenating +# the 33ch first-frame cond latent inside the adapter, so the loop's latent stays 32ch. +HUNYUAN15_LATENT_CHANNELS = 32 +HUNYUAN15_SPATIAL_RATIO = 16 +HUNYUAN15_TEMPORAL_RATIO = 4 + + +class HunyuanVideo15DenoiseLoop(WanDenoiseLoop): + """Wan flow-match loop with HunyuanVideo 1.5 latent geometry. The i2v conditioning the base loop reads + from ``ctx.slots`` is keyed the same way (``i2v_cond``/``i2v_img_embeds``), so a future i2v program can + write them with no loop change; the t2v program leaves them unset (pure t2v).""" + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels: int = HUNYUAN15_LATENT_CHANNELS, + spatial_ratio: int = HUNYUAN15_SPATIAL_RATIO, + temporal_ratio: int = HUNYUAN15_TEMPORAL_RATIO): + super().__init__(loop_id=loop_id, + cfg=cfg, + flow_shift=flow_shift, + precision=precision, + expert=expert, + cost=cost, + latent_channels=latent_channels, + spatial_ratio=spatial_ratio, + temporal_ratio=temporal_ratio) diff --git a/v2/recipes/hunyuan_video15/program.py b/v2/recipes/hunyuan_video15/program.py new file mode 100644 index 000000000..94873f0dd --- /dev/null +++ b/v2/recipes/hunyuan_video15/program.py @@ -0,0 +1,70 @@ +"""HunyuanVideo 1.5 t2v program: text_encode → diffusion_denoise → vae_decode. + +Same inline shape as the Wan/Cosmos t2v programs — the arch specifics (two text embeds, the image-embed +list, the 33ch i2v concat) live entirely in the ``HunyuanVideo15DiT`` adapter, so the node graph is +unchanged except for the text-encode node, which must produce BOTH text embeddings. + +HunyuanVideo 1.5 conditions on two encoders (faithful to ``HunyuanVideo15Pipeline``'s +``TextEncodingStage(text_encoders=[text_encoder, text_encoder_2])``): + * ``text_encoder`` = Qwen2.5-VL -> ``hidden_states[-3]`` cropped of the 108 template tokens (3584-d) + * ``text_encoder_2`` = ByT5/Glyph -> ``last_hidden_state`` of the quoted-glyph text (1472-d), or a + zero-length embedding when the prompt carries no quoted text. + +The text-encode node packs them as a ``(qwen, byt5)`` tuple in the slot; ``HunyuanVideo15DiT`` unpacks +that into the two-element ``encoder_hidden_states`` list. On the CPU **toy** backend there is a single +``text_encoder`` component (``ToyTextEncoder``) and no ``text_encoder_2``: the node then writes the plain +toy embedding (a single numpy array), which ``ToyDiT`` consumes directly — so the program CPU-verifies +unchanged. The feature cache (design_v3 §7.2) still applies to the Qwen path via ``cached_text_encode``. + +BRINGUP (i2v): the i2v path adds a VAE-encode node that writes the 33-channel first-frame conditioning +latent into ``i2v_cond`` and zero image embeds into ``i2v_img_embeds`` (the slots the loop already reads); +``Hy15ImageEncodingStage`` builds those (image_embeds=zeros[1,729,1152]; image_latent=[expanded VAE +first-frame ⊕ mask] -> 33ch). No CLIP/SigLIP vision encoder is loaded (image_embeds are placeholder +zeros), so do NOT add an image_encoder component. Left as a documented hook for the t2v registration. +""" +from __future__ import annotations + +from typing import Any + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind +from v2.recipes.common import cached_text_encode + + +def _text_encode(instance: Any, slots: dict, request: Any, ctx: Any) -> None: + """Encode prompt + negative prompt. When a second text encoder (ByT5) exists, pack ``(qwen, byt5)`` + tuples; otherwise (toy backend) write the single embedding ``ToyDiT`` expects.""" + has_byt5 = "text_encoder_2" in instance.card.components + + def _encode(text: str) -> Any: + qwen = cached_text_encode(instance, text) # primary (Qwen2.5-VL) — feature-cached + if not has_byt5: + return qwen # toy backend: single embedding + byt5 = instance.component("text_encoder_2").encode(text) + return (qwen, byt5) + + slots["text_embeds"] = _encode(request.prompt()) + slots["neg_text_embeds"] = _encode(request.diffusion.negative_prompt) + + +def _vae_decode(instance: Any, slots: dict, request: Any, ctx: Any) -> None: + slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_hunyuan_video15_program() -> Program: + return Program( + program_id="hunyuan_video15.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/recipes/hunyuangamecraft/__init__.py b/v2/recipes/hunyuangamecraft/__init__.py new file mode 100644 index 000000000..15e26e905 --- /dev/null +++ b/v2/recipes/hunyuangamecraft/__init__.py @@ -0,0 +1,17 @@ +"""HunyuanGameCraft (interactive camera/action-conditioned i2v) ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``GameCraftDiT``/``GameCraftVAE``/``GameCraftLlamaEncoder``/ +``GameCraftClipEncoder`` in ``v2/platform/backends/torch_hunyuangamecraft.py``) and a new +``GameCraftDenoiseLoop`` (flow-match Euler + 33ch concat + per-step clean-ref injection). The registered +path is the t2v/degenerate denoise; the camera/action (CameraNet Plücker) conditioning is BRINGUP — it +needs a request-API camera-input channel v2 lacks today. Registered in ``v2/registry.py`` by the +orchestrator from this package's returned metadata. +""" +from __future__ import annotations + +from v2.recipes.hunyuangamecraft.card import build_hunyuangamecraft_card +from v2.recipes.hunyuangamecraft.loop import GameCraftDenoiseLoop +from v2.recipes.hunyuangamecraft.program import build_hunyuangamecraft_program + +__all__ = ["build_hunyuangamecraft_card", "build_hunyuangamecraft_program", "GameCraftDenoiseLoop"] diff --git a/v2/recipes/hunyuangamecraft/card.py b/v2/recipes/hunyuangamecraft/card.py new file mode 100644 index 000000000..1128fa4ef --- /dev/null +++ b/v2/recipes/hunyuangamecraft/card.py @@ -0,0 +1,174 @@ +"""HunyuanGameCraft ModelCard — camera/action-conditioned interactive i2v (the registered preset is the +t2v/degenerate path; the action/camera conditioning is BRINGUP — see the recipe ``__init__`` docstring). + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.hunyuangamecraft:HunyuanGameCraftTransformer3DModel`` — a flow-match DiT + whose forward takes a **33-channel** input ``cat([latent16|gt_latent16|mask1])`` (assembled by the + loop/adapter, NOT the DiT), a **list** of text states ``[LLaMA(4096d), CLIP-pooled(768d)]``, an + optional ``camera_states`` (CameraNet Plücker), and ``guidance=None`` (plain ClassicCFG, NO embedded + guidance). Adapter -> ``GameCraftDiT``; the new ``GameCraftDenoiseLoop`` does the 33ch concat + (via the adapter), per-step clean-ref injection, and flow-match Euler. + * VAE ``fastvideo.models.vaes.gamecraftvae:GameCraftVAE`` — SCALAR ``scaling_factor=0.476986`` + normalization (NOT Wan per-channel stats). Adapter -> ``GameCraftVAE`` (cannot reuse ``WanVAE``). + * Text ``fastvideo.models.encoders.llama:LlamaModel`` (LLaVA-LLaMA-3-8B, 4096d -> ``text_states``) and + ``fastvideo.models.encoders.clip:CLIPTextModel`` (CLIP ViT-L/14, 768d pooled -> ``text_states_2``). + Adapters -> ``GameCraftLlamaEncoder`` / ``GameCraftClipEncoder``. +The diffusers checkpoint layout has extra ``text_encoder_2/`` + ``tokenizer_2/`` subfolders the Wan stamp +lacks, so this file extends ``stamp_wan21_checkpoints`` with a GameCraft-specific stamp. +""" +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.hunyuangamecraft.loop import GameCraftDenoiseLoop + +# GameCraft uses an empty negative prompt by default (the preset's negative_prompt=""). Kept as a LOCAL +# module constant so this recipe does not depend on the shared ``v2/recipes/_prompts.py`` (per the +# self-contained-recipe rule). CFG still runs (uncond branch encodes the empty string). +GAMECRAFT_NEG = "" + +_GAMECRAFT_DIT = "v2.platform.backends.torch_hunyuangamecraft:GameCraftDiT" +_GAMECRAFT_VAE = "v2.platform.backends.torch_hunyuangamecraft:GameCraftVAE" +_GAMECRAFT_LLAMA = "v2.platform.backends.torch_hunyuangamecraft:GameCraftLlamaEncoder" +_GAMECRAFT_CLIP = "v2.platform.backends.torch_hunyuangamecraft:GameCraftClipEncoder" + + +def build_hunyuangamecraft_card(model_id: str = "hunyuan-gamecraft", + *, + checkpoint_root: str | None = None, + flow_shift: float = 5.0, + 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() + # Single (resolution-independent) flow-shift 5.0; no bucket lookup (GameCraft fixes 704x1280). + flow = FlowShiftPolicy(shift=flow_shift) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return GameCraftDenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost) + + components = { + # LLaVA-LLaMA-3 -> text_states (4096d). The toy CPU path uses ToyTextEncoder for both encoders. + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.llama:LlamaModel", + adapter=_GAMECRAFT_LLAMA, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v", "i2v"}), + # CLIP ViT-L/14 -> text_states_2 (768d pooled). + "text_encoder_2": + ComponentSpec(component_id="text_encoder_2", + kind="text_encoder", + load_id="fastvideo.models.encoders.clip:CLIPTextModel", + adapter=_GAMECRAFT_CLIP, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v", "i2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.gamecraftvae:GameCraftVAE", + adapter=_GAMECRAFT_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2v", "i2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.hunyuangamecraft:HunyuanGameCraftTransformer3DModel", + adapter=_GAMECRAFT_DIT, + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["diffusion_denoise"], + required_for={"t2v", "i2v"}), + } + 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="gamecraft", + components=components, + loops=loops, + # IMAGE_TO_VIDEO is the real capability (interactive i2v); the registered path degenerates to a + # standard denoise when no camera/action input is given (BRINGUP). VAE_DECODE + POLICY_ROLLOUT + # mirror the other video cards. + capabilities=CapabilityMatrix.of(Capability.IMAGE_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=50, + guidance_scale=6.0, + height=704, + width=1280, + num_frames=33, + fps=24, + negative_prompt=GAMECRAFT_NEG), + ) + card.validate() + if checkpoint_root: + stamp_gamecraft_checkpoints(card, checkpoint_root) + return card + + +# Diffusers checkpoint layout for GameCraft: dual text encoders + dual tokenizers + scheduler, in addition +# to transformer/vae. ``stamp_wan21_checkpoints`` only knows a single tokenizer + lacks text_encoder_2, so +# GameCraft carries its own stamp map. (The torch adapters resolve tokenizers from ``/tokenizer*``.) +_GAMECRAFT_SUBFOLDERS = { + "transformer": "transformer", + "vae": "vae", + "text_encoder": "text_encoder", # LLaVA-LLaMA-3 + "text_encoder_2": "text_encoder_2", # CLIP ViT-L/14 +} + + +def stamp_gamecraft_checkpoints(card: ModelCard, model_root: str) -> ModelCard: + """GPU deploy-time: point each component's ``ComponentSpec.checkpoint`` at its weights subfolder under + ``model_root`` (a local diffusers dir or an HF id resolved via ``snapshot_download``). Mirrors + ``stamp_wan21_checkpoints`` but covers GameCraft's ``text_encoder_2`` (the Wan stamp lacks it). The + sibling ``tokenizer/`` (LLaMA) and ``tokenizer_2/`` (CLIP) + ``scheduler/`` subfolders are loaded by + the torch adapters from ``dirname(checkpoint)``. Mutates and returns the card. BRINGUP risk A.""" + import os + if not os.path.isdir(model_root): + from huggingface_hub import snapshot_download + model_root = snapshot_download(model_root) + for component_id, subfolder in _GAMECRAFT_SUBFOLDERS.items(): + if component_id in card.components: + card.components[component_id].checkpoint = os.path.join(model_root, subfolder) + return card diff --git a/v2/recipes/hunyuangamecraft/loop.py b/v2/recipes/hunyuangamecraft/loop.py new file mode 100644 index 000000000..927a07077 --- /dev/null +++ b/v2/recipes/hunyuangamecraft/loop.py @@ -0,0 +1,254 @@ +"""GameCraftDenoiseLoop — a flow-match Euler denoise loop with GameCraft i2v + camera conditioning. + +GameCraft (HunyuanGameCraft) is an *interactive* world model: a camera/action-conditioned i2v diffusion +DiT. The sampler is plain flow-match (``FlowMatchEulerDiscreteScheduler``, ``flow_shift=5.0``, +``timestep = σ·1000``) with standard ClassicCFG (``guidance_scale=6.0``, ``guidance=None`` to the DiT — NO +embedded guidance). So the *integrator* is identical to ``WanDenoiseLoop`` (``x_next = x + (σ_next−σ)·v`` +via ``FLOW_MATCH_STEP``). The deltas this loop carries (faithful to +``fastvideo/pipelines/stages/gamecraft_denoising.py``): + + 1. **33-channel input is assembled by the DiT adapter, not here.** The DiT forward wants + ``cat([latent16 | gt_latent16 | mask1], dim=1)``; the loop hands the bare 16ch latent + a ``cond`` + pack ``{"gt_latents","mask","camera_states"}`` to the ``GameCraftDiT`` adapter, which concats. The + toy ``ToyDiT`` accepts (and ignores) the ``cond`` kwarg, so this loop CPU-verifies unchanged. + + 2. **Per-step clean-ref injection (frame_replace).** Official GameCraft overwrites the *conditioned* + latent frames (``mask>0.5``) with the clean reference latent at EVERY step (not just step 0), so the + conditioned frames stay noise-free throughout denoising. ``next()`` mutates the working latent before + the forward when a reference latent is present. + + 3. **Dual text encoders.** ``encoder_hidden_states = [LLaMA_states(4096d), CLIP_pooled(768d)]``. The + loop threads the LLaMA states as the usual ``text_embed`` (cond vs uncond) and the CLIP pooled embed + through the adapter's ``context=`` slot per branch, so ClassicCFG combines both correctly. + + 4. **BRINGUP — camera/action conditioning.** GameCraft is fundamentally CameraNet (Plücker) + conditioned, with the forward special-casing the latent temporal length (18/9/10) for autoregressive + multi-chunk generation. v2 has no camera/action request field or Plücker-conversion node yet, so the + registered path passes ``camera_states=None`` -> the DiT's ``if camera_states is not None`` branch is + skipped (the standard / degenerate denoise). For pure t2v (no reference image) gt_latents+mask are + zero -> the 33ch concat reduces to ``[noise | 0 | 0]``, exactly the ``GameCraftDenoisingStage`` + fallback. Wiring real camera/action + the 18/9/10 length special-casing needs a request-API + extension (see ``program.py`` blockers). + +All per-request state lives in ``LoopState`` (interleave-safe). ``next`` is kernel-free; ``advance`` folds +the result + (under ROLLOUT) captures a trajectory slice — same shape as ``WanDenoiseLoop``. +""" +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.platform import FLOW_MATCH_STEP +from v2.recipes.wan21.loop import latent_shape + +# GameCraftVAE geometry: z=16, 8x spatial, 4x temporal (T_lat = (num_frames-1)//4 + 1). +GAMECRAFT_LATENT_CHANNELS = 16 +GAMECRAFT_TEMPORAL_RATIO = 4 +GAMECRAFT_SPATIAL_RATIO = 8 + + +class GameCraftDenoiseLoop: + """Flow-match Euler denoise with GameCraft i2v (33ch concat + per-step ref injection) + camera hooks.""" + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels: int = GAMECRAFT_LATENT_CHANNELS, + spatial_ratio: int = GAMECRAFT_SPATIAL_RATIO, + temporal_ratio: int = GAMECRAFT_TEMPORAL_RATIO): + self.loop_id = loop_id + self.cfg = cfg + self.flow_shift = flow_shift + self.precision = precision + self.expert = expert + self.cost = cost + 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) + # Flow-match σ schedule (shift=5.0 for GameCraft); honor an explicit distilled sigmas list if given. + sig = self.flow_shift.build_schedule(req.diffusion.num_steps, + req.diffusion.height, + req.diffusion.width, + sigmas=req.diffusion.sigmas or None) + 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(sig[0])).astype("float32") + 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 + # Dual text: LLaMA hidden states (cond/uncond) ride the usual prompt-embed slots; the CLIP pooled + # embeds (cond/uncond) ride dedicated slots and are threaded through the adapter's ``context=``. + st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") + st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") + st.scratch["clip_pos"] = ctx.slots.get("clip_text_embeds") + st.scratch["clip_neg"] = ctx.slots.get("clip_neg_text_embeds") + st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale) + st.scratch["stream_video"] = bool(req.outputs.stream.get("video")) + # i2v conditioning slots written by the image-cond program node. All None for pure t2v -> the 33ch + # concat degenerates to [noise | 0 | 0] and the ref-injection branch is inert (the fastvideo + # GameCraftDenoisingStage fallback). camera_states is a BRINGUP slot (no request plumbing yet). + st.scratch["gt_latents"] = ctx.slots.get("gt_latents") + st.scratch["conditioning_mask"] = ctx.slots.get("conditioning_mask") + st.scratch["ref_latent_for_injection"] = ctx.slots.get("ref_latent_for_injection") + st.scratch["camera_states"] = ctx.slots.get("camera_states") # BRINGUP: None until camera plumbing + st.plugin_state["cfg"] = {} + return st + + def _inject_ref(self, x: np.ndarray, mask, ref) -> np.ndarray: + """Per-step frame_replace: overwrite conditioned frames (mask>0.5) with the clean reference latent. + + Faithful to ``GameCraftDenoisingStage`` (``latents[:,:,t_idx] = ref[:,:,0]`` for each conditioned + ``t``). Loop latents here are unbatched ``[C,T,h,w]`` (the toy/CPU geometry), so the mask's temporal + axis is read at ``mask[0,0,:,0,0]`` when 5D (matching the fastvideo ``[B,1,T,H,W]`` mask), or as a + 1D ``[T]`` vector otherwise. ``ref`` is the clean ref latent ``[C,1,h,w]`` (or ``[C,T,h,w]``).""" + m = np.asarray(mask) + if m.ndim == 5: # fastvideo-shaped [B,1,T,H,W]: read the temporal slice + frame_mask = m[0, 0, :, 0, 0] + elif m.ndim == 4: # loop-shaped [1,T,h,w]: collapse the spatial axes + frame_mask = m[0, :, 0, 0] + elif m.ndim == 1: # already a per-frame vector [T] + frame_mask = m + else: + return x + r = np.asarray(ref, dtype="float32") + ref_frame = r[:, 0] if r.ndim >= 4 else r # [C,h,w] + out = np.array(x, dtype="float32", copy=True) + # The conditioned-frame replacement requires the ref frame to match the working latent's per-frame + # geometry (it always does on the real VAE geometry). If a backend's encode produced a mismatched + # shape (e.g. the crude CPU toy VAE that does not spatially downsample), skip injection rather than + # crash — the degenerate denoise still proceeds. + if ref_frame.shape != out[:, 0].shape: + return out + n_t = min(out.shape[1], int(frame_mask.shape[0])) + for t_idx in range(n_t): + if float(frame_mask[t_idx]) > 0.5: + out[:, t_idx] = ref_frame + return out + + 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] + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + branches = self.cfg.branches_this_step(sctx, cfg_state) + + x = st.latents["video"] + gt = st.scratch.get("gt_latents") + mask = st.scratch.get("conditioning_mask") + ref = st.scratch.get("ref_latent_for_injection") + camera = st.scratch.get("camera_states") # BRINGUP: None for the registered path + # Per-step clean-ref injection BEFORE the forward (frame_replace) — only when an i2v ref is present. + if ref is not None and mask is not None: + x = self._inject_ref(x, mask, ref) + st.latents["video"] = x + + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + clip_pos, clip_neg = st.scratch.get("clip_pos"), st.scratch.get("clip_neg") + scale = st.scratch["guidance_scale"] + cfg, precision = self.cfg, self.precision + # One ``cond`` pack reused across CFG branches (gt/mask/camera are branch-independent); the adapter + # assembles the 33ch concat from it. None entries -> zero gt/mask + no camera inside the adapter. + cond_pack = {"gt_latents": gt, "mask": mask, "camera_states": camera} + + def _velocity(model: Any, x_: np.ndarray, sigma_t_: float, pe_: Any, clip_: Any, ne_: Any, clip_n_: Any, + scale_: float) -> np.ndarray: + # Conditioned forward + ClassicCFG combine. The CLIP pooled embed rides ``context=`` per branch + # so text_states_2 differs cond vs uncond; ``cond`` carries the i2v/camera pack. The toy DiT + # accepts+ignores both kwargs, so this is CPU-toy compatible. + dit = model.component(expert_id) + preds = { + b: + dit(x_, + pe_ if b == "cond" else ne_, + sigma_t_, + context=(clip_ if b == "cond" else clip_n_), + cond=cond_pack) + for b in branches + } + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + + def run(model, override=None): + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x, sigma_t, pe, clip_pos, ne, clip_neg, scale) + 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, clip_pos, clip_neg, gt, mask) if e is not None) + res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), + 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(cfg).__name__), ("gamecraft", True))), + resources=res, + payload={ + "branch": "combined", + "step": i + }, + run=run, + label=f"gamecraft.denoise.{i}", + # i2v ref-injection mutates host state + threads non-workspace conditioning each step -> the + # static-buffer capture form would miss it, so this step takes the eager path (like cosmos2's + # gated frame-injection). Pure-t2v could capture, but keeping one path is simpler + correct. + capturable=False) + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + st.latents["video"] = result.output["latents"] + if st.profile == ExecutionProfile.ROLLOUT: + i = st.step_idx + st.trajectory.append({ + "step": i, + "sigma": st.sigmas[i], + "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/hunyuangamecraft/program.py b/v2/recipes/hunyuangamecraft/program.py new file mode 100644 index 000000000..4c94ea967 --- /dev/null +++ b/v2/recipes/hunyuangamecraft/program.py @@ -0,0 +1,120 @@ +"""HunyuanGameCraft program: text_encode (LLaMA + CLIP) -> image_cond_encode (i2v) -> diffusion_denoise +-> vae_decode. + +Same backbone as the Wan/Cosmos t2v programs, with two GameCraft-specific deltas: + + * **Dual text encode.** ``_text_encode_dual`` runs BOTH encoders: the LLaMA encoder (``text_encoder``) + writes the DiT's ``text_states`` into ``text_embeds``/``neg_text_embeds`` (the usual slots the loop + reads), and the CLIP encoder (``text_encoder_2``) writes the pooled ``text_states_2`` into + ``clip_text_embeds``/``clip_neg_text_embeds`` (threaded to the adapter's ``context=`` per CFG branch). + On the CPU toy both are ``ToyTextEncoder`` (cached via the content-hash feature cache). + + * **Image conditioning node (i2v).** ``_image_cond_encode`` mirrors ``GameCraftImageVAEEncodingStage``: + VAE-encode the reference image -> scaled ``ref_latent`` (the ``GameCraftVAE`` adapter applies the + scalar ``scaling_factor`` already), repeat to ``T_lat`` frames -> ``gt_latents``, build the binary + ``conditioning_mask`` (latent_frames<=10: only frame 0; else first half), zero ``gt_latents`` on + non-conditioned frames, and expose ``ref_latent_for_injection = gt_latents[:, 0:1]`` for the loop's + per-step frame_replace. No request image -> a no-op (slots stay unset) -> the loop runs pure t2v + (zero gt/mask), exactly the fastvideo fallback. + +BRINGUP — the camera/action conditioning node is intentionally omitted: GameCraft's CameraNet needs a +Plücker-coordinate ``camera_states`` tensor derived from an action/camera sequence the v2 request API does +not carry yet, and the DiT's autoregressive latent-length special-casing (18/9/10) for multi-chunk +generation has no request surface. Adding a ``camera_encode`` node that writes the ``camera_states`` slot +(which the loop already reads) is the follow-up once the request API is extended. T_lat helper below +matches the GameCraft geometry (4x temporal compression). +""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind +from v2.recipes.common import cached_text_encode + +GAMECRAFT_TEMPORAL_RATIO = 4 +GAMECRAFT_SPATIAL_RATIO = 8 + + +def _text_encode_dual(instance: Any, slots: dict, request: Any, ctx: Any) -> None: + """LLaMA (text_states) + CLIP (text_states_2) for prompt + negative prompt. The shared + ``cached_text_encode`` keys on component_id, so the LLaMA and CLIP embeds never collide in the cache.""" + prompt = request.prompt() + neg = request.diffusion.negative_prompt + # LLaMA hidden states -> the loop's prompt-embed slots (text_states). + slots["text_embeds"] = cached_text_encode(instance, prompt) + slots["neg_text_embeds"] = cached_text_encode(instance, neg) + # CLIP pooled -> dedicated slots (text_states_2), threaded through the adapter's context= per branch. + clip = instance.component("text_encoder_2") + slots["clip_text_embeds"] = clip.encode(prompt) + slots["clip_neg_text_embeds"] = clip.encode(neg) + + +def _image_cond_encode(instance: Any, slots: dict, request: Any, ctx: Any) -> None: + """i2v: VAE-encode the reference image -> gt_latents + conditioning_mask + ref_latent_for_injection. + + No-op for pure t2v (no request image) so the loop runs the degenerate standard denoise. Faithful to + ``GameCraftImageVAEEncodingStage``: repeat the ref latent to all latent frames, then keep only the + conditioned frames (first frame for short clips <=10 latent frames; first half for longer + autoregressive clips) — both ``gt_latents`` and ``mask`` zeroed elsewhere.""" + part = request.image() if hasattr(request, "image") else None + pixels = getattr(part, "pixels", None) if part is not None else None + if pixels is None: + return # T2V mode: leave the i2v slots unset (loop falls back to zero gt/mask) + + d = request.diffusion + # The GameCraftVAE adapter applies the scalar scaling_factor in ``encode`` (matches the fastvideo + # stage's ``ref_latents.mul_(scaling_factor)``), so we use the returned latent directly. The real + # adapter also runs the resize/center-crop/normalize preprocessing on the PIL image (BRINGUP); the + # CPU toy VAE consumes the pixel array directly. + ref_latent = np.asarray(instance.component("vae").encode(pixels), dtype="float32") # [C,1,h,w] (or [C,h,w]) + if ref_latent.ndim == 3: # [C, h, w] -> add the singleton temporal axis + ref_latent = ref_latent[:, None] + ref_latent = ref_latent[:, :1] # keep the first (only) reference frame -> [C, 1, h, w] + + latent_frames = max(1, (max(1, d.num_frames) - 1) // GAMECRAFT_TEMPORAL_RATIO + 1) + c, _t, h, w = ref_latent.shape + gt_latents = np.repeat(ref_latent, latent_frames, axis=1) # [C, T_lat, h, w] + mask = np.ones((1, latent_frames, h, w), dtype="float32") # [1, T_lat, h, w] + if latent_frames <= 10: # short i2v: only frame 0 is conditioned + gt_latents[:, 1:] = 0.0 + mask[:, 1:] = 0.0 + else: # autoregressive: first half conditioned + half = latent_frames // 2 + gt_latents[:, half:] = 0.0 + mask[:, half:] = 0.0 + + slots["gt_latents"] = gt_latents + slots["conditioning_mask"] = mask + slots["ref_latent_for_injection"] = gt_latents[:, 0:1].copy() # the clean ref for frame_replace + + +def _vae_decode(instance: Any, slots: dict, request: Any, ctx: Any) -> None: + slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_hunyuangamecraft_program() -> Program: + return Program( + program_id="gamecraft.i2v.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("text_encode", + fn=_text_encode_dual, + writes=("text_embeds", "neg_text_embeds", "clip_text_embeds", "clip_neg_text_embeds")), + ComponentNode("image_cond_encode", + fn=_image_cond_encode, + writes=("gt_latents", "conditioning_mask", "ref_latent_for_injection")), + ModelLoopNode("denoise", + loop_id="diffusion_denoise", + output_slot="denoise_out", + reads=("text_embeds", "neg_text_embeds", "clip_text_embeds", "clip_neg_text_embeds", + "gt_latents", "conditioning_mask", "ref_latent_for_injection"), + 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/recipes/hyworld/__init__.py b/v2/recipes/hyworld/__init__.py new file mode 100644 index 000000000..e3ea07c62 --- /dev/null +++ b/v2/recipes/hyworld/__init__.py @@ -0,0 +1,20 @@ +"""HY-WorldPlay-Bidirectional (interactive world model) — ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``HYWorldDiT`` / ``HYWorldVAE`` / ``HYWorldQwenEncoder`` / +``HYWorldByT5Encoder`` / ``HYWorldSiglipEncoder`` in ``v2/platform/backends/torch_hyworld.py``) and a new +``HYWorldDenoiseLoop`` (chunk-rollout flow-match: per-chunk sweep + camera-aligned frozen context for +chunk>0). It reuses the flow-match ``FLOW_MATCH_STEP`` kernel + ``ClassicCFG`` policy and the Wan +``stamp_wan21_checkpoints`` superset (which covers ``text_encoder_2`` + ``image_encoder``). + +The registered preset is the bidirectional t2v / degenerate (no action/camera) path that CPU-verifies; +the full action/camera/memory-retrieval conditioning is BRINGUP (needs a request-API extension to carry +the pose string + first-frame image). Registered in ``v2/registry.py``. +""" +from __future__ import annotations + +from v2.recipes.hyworld.card import build_hyworld_card +from v2.recipes.hyworld.loop import HYWorldDenoiseLoop +from v2.recipes.hyworld.program import build_hyworld_program + +__all__ = ["build_hyworld_card", "build_hyworld_program", "HYWorldDenoiseLoop"] diff --git a/v2/recipes/hyworld/card.py b/v2/recipes/hyworld/card.py new file mode 100644 index 000000000..bdc3d84de --- /dev/null +++ b/v2/recipes/hyworld/card.py @@ -0,0 +1,191 @@ +"""HY-WorldPlay-Bidirectional ModelCard — interactive world model ported into the v2 substrate. + +HY-WorldPlay is a chunk-rollout *world model*: it generates video in temporal chunks where each +chunk after the first retrieves camera-aligned history "memory" frames as frozen context, and it +conditions on THREE text/image streams (Qwen2.5-VL mllm + ByT5 glyph + SigLIP image) plus per-frame +camera (viewmats/Ks, injected via ProPE attention) and per-frame action ids. The registered preset +here is the **bidirectional t2v / degenerate (no-action, single straight-trajectory pose) path** that +CPU-verifies; the full action/camera/memory conditioning is BRINGUP (needs a request-API extension to +carry the pose string -> viewmats/Ks/action expansion + the point-cloud memory retrieval). See +``loop.py`` and ``torch_hyworld.py`` for the BRINGUP markers. + +Architecture deltas vs Wan/Cosmos (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.hyworld.hyworld:HYWorldTransformer3DModel`` (subclasses + HunyuanVideo15Transformer3DModel) — a **rectified-flow velocity** predictor whose forward takes a + 65ch pre-concatenated latent ``[noise(32) | cond(32) + mask(1)]``, a LIST of text embeds + ``[qwen, byt5]``, a LIST image embed ``[siglip]``, a PER-LATENT-FRAME LongTensor ``timestep`` (+ a + scalar ``timestep_txt``), and ``viewmats/Ks/action`` per-frame camera/action conditioning. The + adapter (``HYWorldDiT``) marshals all of this INTERNALLY so the loop only hands it + ``dit(latent, text_embed, sigma)`` (the cosmos2/ToyDiT-compatible signature). + * VAE ``fastvideo.models.vaes.hyworldvae:AutoencoderKLHYWorld`` (z=32, 16x spatial, 4x temporal, + scaling_factor=1.03682) — encode -> ``.mode() * scaling_factor`` (no shift/mean), decode inverts. + * Encoders: Qwen2.5-VL (mllm primary text) + ByT5/T5 (glyph) + SigLIP (image). For the t2v preset the + glyph + image streams are zeroed (the DiT detects an all-zero image and masks it). + * Sampler: flow-match (``FlowMatchEulerDiscreteScheduler``) but the preset OVERRIDES the schedule with + explicit sigmas ``linspace(1.0, 0.0, 51)[:-1]`` (50 steps), flow_shift 5.0, guidance 6.0. +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder/text_encoder_2/image_encoder +subfolder layout); the superset stamp helper already covers ``text_encoder_2`` + ``image_encoder``. +""" +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyImageEncoder, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.hyworld.loop import ( + HYWORLD_CHUNK_LATENT_FRAMES, + HYWORLD_LATENT_CHANNELS, + HYWORLD_SPATIAL_RATIO, + HYWORLD_STABILIZATION_LEVEL, + HYWORLD_TEMPORAL_RATIO, + HYWorldDenoiseLoop, +) +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# Adapter refs declared on the card -> built lazily on a GPU box by ``_explicit_adapter`` (so the port +# never edits the shared ``_make_dit``/``_make_vae``/``_make_text_encoder``/``_make_image_encoder``). +_HYWORLD_DIT = "v2.platform.backends.torch_hyworld:HYWorldDiT" +_HYWORLD_VAE = "v2.platform.backends.torch_hyworld:HYWorldVAE" +_HYWORLD_QWEN = "v2.platform.backends.torch_hyworld:HYWorldQwenEncoder" +_HYWORLD_BYT5 = "v2.platform.backends.torch_hyworld:HYWorldByT5Encoder" +_HYWORLD_SIGLIP = "v2.platform.backends.torch_hyworld:HYWorldSiglipEncoder" + +# The preset's negative prompt is empty (kept as a LOCAL module constant, not in the shared _prompts.py). +HYWORLD_NEG = "" + +# The preset's explicit sigma schedule: linspace(1.0, 0.0, 51)[:-1] (50 values). The FlowShiftPolicy / +# the loop consume ``sigmas`` directly when present (the flow_shift below is the schedule fallback). +_HYWORLD_SIGMAS = tuple(1.0 - i / 50.0 for i in range(50)) + + +def build_hyworld_card(model_id: str = "hy-worldplay-bidirectional", + *, + 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() # uncond + s·(cond - uncond); the pos/neg branches swap the encoder_attention_mask + # flow_shift=5.0 is the scheduler shift; the preset overrides the schedule with explicit sigmas, so + # FlowShiftPolicy.build_schedule(sigmas=...) returns those verbatim (shift becomes a no-op there). + flow = FlowShiftPolicy(shift=5.0) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") # single DiT (the bidirectional model is not MoE) + + def loop_factory() -> HYWorldDenoiseLoop: + return HYWorldDenoiseLoop(loop_id="hyworld_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost, + latent_channels=HYWORLD_LATENT_CHANNELS, + spatial_ratio=HYWORLD_SPATIAL_RATIO, + temporal_ratio=HYWORLD_TEMPORAL_RATIO, + chunk_latent_frames=HYWORLD_CHUNK_LATENT_FRAMES, + stabilization_level=HYWORLD_STABILIZATION_LEVEL) + + components = { + # Primary text stream — Qwen2.5-VL mllm (hidden_states[-3], crops 108 template tokens on GPU). + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.qwen2_5:Qwen2_5_VLTextModel", + adapter=_HYWORLD_QWEN, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + # Glyph text stream — ByT5/T5 (rendered-text encoder, text_states_dim_2=1472). Zeroed for plain t2v. + "text_encoder_2": + ComponentSpec(component_id="text_encoder_2", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5_hf:T5EncoderModel", + adapter=_HYWORLD_BYT5, + factory=lambda inst: ToyTextEncoder(), + optional_for={"t2v"}), + # Image stream — SigLIP (729 tokens x 1152 dim). Zeroed for t2v -> the DiT masks the image stream. + "image_encoder": + ComponentSpec(component_id="image_encoder", + kind="image_encoder", + load_id="fastvideo.models.encoders.siglip:SiglipVisionModel", + adapter=_HYWORLD_SIGLIP, + factory=lambda inst: ToyImageEncoder(), + optional_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.hyworldvae:AutoencoderKLHYWorld", + adapter=_HYWORLD_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.hyworld.hyworld:HYWorldTransformer3DModel", + adapter=_HYWORLD_DIT, + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["hyworld_denoise"], + required_for={"t2v"}), + } + loops = { + "hyworld_denoise": + LoopSpec( + loop_id="hyworld_denoise", + kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, + step_cost_model=cost, + shared_weight_components=["transformer"], + cache_policy=["feature"], + graph_capture="eager", # chunk rollout + camera-aligned memory retrieval -> eager path + loop_factory=loop_factory), + } + card = ModelCard( + model_id=model_id, + family="hyworld", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of( + Capability.TEXT_TO_VIDEO, # the registered (CPU-verified) degenerate path + Capability.IMAGE_TO_VIDEO, # first-frame SigLIP + VAE cond latent (BRINGUP) + Capability.ACTION_CONDITIONING, # per-frame action ids (BRINGUP) + Capability.STREAMING_VIDEO_CONTINUATION, # chunk rollout = interactive continuation (BRINGUP) + Capability.VAE_DECODE, + Capability.POLICY_ROLLOUT), + recipe=RecipeSpec(method="base", + assumes_loop="hyworld_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=50, + guidance_scale=6.0, + height=480, + width=832, + num_frames=125, + fps=24, + negative_prompt=HYWORLD_NEG, + shift=5.0, + sigmas=_HYWORLD_SIGMAS), + ) + card.validate() + if checkpoint_root: + # Diffusers layout: transformer/ vae/ text_encoder/ text_encoder_2/ image_encoder/ scheduler/. + # The Wan stamp superset already covers text_encoder_2 + image_encoder (BRINGUP risk A). + stamp_wan21_checkpoints(card, checkpoint_root) + return card diff --git a/v2/recipes/hyworld/loop.py b/v2/recipes/hyworld/loop.py new file mode 100644 index 000000000..173e42cb2 --- /dev/null +++ b/v2/recipes/hyworld/loop.py @@ -0,0 +1,243 @@ +"""HYWorldDenoiseLoop — chunk-rollout flow-match denoise for the HY-WorldPlay world model. + +The standard ``WanDenoiseLoop`` runs ONE flow-match sweep over the full latent with a scalar timestep. +HY-WorldPlay instead splits the latent into temporal CHUNKS (``chunk_latent_frames=16``) and rolls out +chunk by chunk (faithful to ``fastvideo/pipelines/stages/hyworld_denoising.py:HYWorldDenoisingStage``): + + * chunk 0 (always runs): a standard per-chunk flow-match sweep over the first 16 latent frames; the + per-frame timestep is just the live ``t`` broadcast to all 16 frames. + * chunk>0 (BRINGUP, see below): selects camera-aligned history "memory" frames via + ``select_aligned_memory_frames`` (a 50k-point-sphere geometry op), prepends them as FROZEN context + (timestep pinned at ``stabilization_level - 1 = 14``) with the live chunk's timestep ``t``, denoises + only the current 16-frame slice, and writes that slice back. + +The flow-match ODE math + ClassicCFG reuse v2's ``FLOW_MATCH_STEP`` kernel + ``ClassicCFG`` policy; the +chunk / context / dual-timestep orchestration is the new part. The DiT adapter (``HYWorldDiT``) marshals +the 65ch latent concat, the 3-stream conditioning list, the per-frame heterogeneous timestep, and the +camera/action tensors INTERNALLY, so this loop only calls the cosmos2/ToyDiT-compatible +``dit(latent, text_embed, sigma)`` (with the per-chunk camera context threaded via ``context=``). + +BRINGUP (documented honestly — needs a request-API extension to be real): + * The default preset (num_frames=125 -> ~32 latent frames, chunk 16 -> 2 chunks) WILL execute chunk>0. + Camera-aligned memory retrieval, ProPE camera attention, and per-frame action conditioning all need + ``viewmats``/``Ks``/``action`` carried from a pose string (``pose="w-31"`` -> ``pose_to_input`` + + ``compute_latent_num``). The v2 request/program has no slot for the pose string yet, so on the CPU + toy path (and absent a pose) this loop runs the **degenerate single-chunk t2v path**: it clamps + ``chunk_num`` to the number of whole chunks but the camera/memory selection is skipped (the adapter + builds zero camera/action internally). On a GPU box with a pose-carrying request, the chunk>0 branch + (memory retrieval + frozen context) activates — that path is written-not-run here. +""" +from __future__ import annotations + +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.platform import FLOW_MATCH_STEP +from v2.recipes.wan21.loop import latent_shape + +# Real HY-WorldPlay VAE (AutoencoderKLHYWorld) compression: z=32, 4x temporal, 16x spatial. +HYWORLD_LATENT_CHANNELS = 32 +HYWORLD_TEMPORAL_RATIO = 4 +HYWORLD_SPATIAL_RATIO = 16 +# Chunk rollout config (from HYWorldDenoisingStage + the preset defaults). +HYWORLD_CHUNK_LATENT_FRAMES = 16 +HYWORLD_STABILIZATION_LEVEL = 15 # context (history) frames are pinned at stabilization_level - 1 = 14 + + +class HYWorldDenoiseLoop: + """Outer chunk loop x inner flow-match step loop. ``next``/``advance`` step the inner loop and roll + the outer chunk index when a chunk's sweep finishes (all per-request state in ``LoopState``).""" + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels: int = HYWORLD_LATENT_CHANNELS, + spatial_ratio: int = HYWORLD_SPATIAL_RATIO, + temporal_ratio: int = HYWORLD_TEMPORAL_RATIO, + chunk_latent_frames: int = HYWORLD_CHUNK_LATENT_FRAMES, + stabilization_level: int = HYWORLD_STABILIZATION_LEVEL): + self.loop_id = loop_id + self.cfg = cfg + self.flow_shift = flow_shift + self.precision = precision + self.expert = expert + self.cost = cost + self.latent_channels = latent_channels + self.spatial_ratio = spatial_ratio + self.temporal_ratio = temporal_ratio + self.chunk_latent_frames = chunk_latent_frames + self.stabilization_level = stabilization_level + + 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) + # Preset overrides the schedule with explicit sigmas linspace(1.0, 0.0, 51)[:-1]; otherwise the + # flow-shifted schedule. ``timesteps = sigma * 1000`` (num_train_timesteps=1000) per the convention. + sig = self.flow_shift.build_schedule(req.diffusion.num_steps, + req.diffusion.height, + req.diffusion.width, + sigmas=req.diffusion.sigmas or None) + 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(sig[0])).astype("float32") + 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]) + 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) + # Chunk layout: the latent's temporal dim split into whole chunks of chunk_latent_frames. The toy + # CPU latent is tiny (T may be < chunk size) -> clamp to a single chunk so the loop still exercises. + n_frames = int(x.shape[1]) + chunk = max(1, min(self.chunk_latent_frames, n_frames)) + st.scratch["chunk_latent_frames"] = chunk + st.scratch["chunk_num"] = max(1, n_frames // chunk) + st.scratch["chunk_idx"] = 0 + # World-model conditioning (BRINGUP): the pose string -> viewmats/Ks/action expansion and the + # camera-aligned memory retrieval need a request-API slot that v2 does not carry yet. Absent (CPU + # toy / no pose) -> chunk>0 runs WITHOUT the frozen camera-aligned context (degenerate per-chunk + # sweep). The adapter builds zero camera/action internally so the DiT call is unchanged. + st.scratch["camera_ctx"] = ctx.slots.get("camera_ctx") # {viewmats, Ks, action, points_local} | None + # I2V (BRINGUP): SigLIP image embeds + the [cond_latent | mask] (33ch) latent; None for t2v -> the + # adapter passes zero image embeds + a zero cond latent (the DiT detects all-zero image -> masks it). + st.scratch["i2v_img_embeds"] = ctx.slots.get("i2v_img_embeds") + st.scratch["i2v_cond"] = ctx.slots.get("i2v_cond") + st.plugin_state["cfg"] = {} + return st + + def next(self, st: LoopState): + # Inner step index walks the sigma schedule; when it reaches the end, roll to the next chunk. + i = st.step_idx + n_steps = len(st.sigmas) - 1 + chunk_idx = st.scratch["chunk_idx"] + if chunk_idx >= st.scratch["chunk_num"]: + return Done() + if i >= n_steps: + return Done() + sigma_t, sigma_next = st.sigmas[i], st.sigmas[i + 1] + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + branches = self.cfg.branches_this_step(sctx, cfg_state) + chunk = st.scratch["chunk_latent_frames"] + start = chunk_idx * chunk + end = start + chunk + # The current chunk's latent slice (the only frames this chunk denoises + writes back). + x_full = st.latents["video"] + x_chunk = np.asarray(x_full[:, start:end], dtype="float32") + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + scale = st.scratch["guidance_scale"] + cfg, precision = self.cfg, self.precision + # Per-chunk camera/action context threaded to the adapter (None on the CPU/degenerate path). The + # adapter is responsible for slicing it to [start:end] and prepending memory frames (BRINGUP). + cam_ctx = st.scratch.get("camera_ctx") + i2v_ctx, i2v_cond = st.scratch.get("i2v_img_embeds"), st.scratch.get("i2v_cond") + # The conditioning the DiT call carries: a per-chunk dict (camera/action/i2v) the adapter unpacks. + # Kept as a single ``context=`` arg so the loop's dit-call stays ToyDiT-compatible (cosmos2 rule). + chunk_context = None + if cam_ctx is not None or i2v_ctx is not None or i2v_cond is not None: + chunk_context = { + "camera": cam_ctx, + "chunk": (start, end), + "chunk_idx": chunk_idx, + "i2v_img_embeds": i2v_ctx, + "i2v_cond": i2v_cond, + "stabilization_level": self.stabilization_level + } + + def _velocity(model: object, x_: np.ndarray, sigma_t_: float, pe_: object, ne_: object, + scale_: float) -> np.ndarray: + # Conditioned forward + ClassicCFG combine. The pos/neg branches differ only in the text embed + # (the adapter swaps the matching encoder_attention_mask internally). The chunk's camera/action + # conditioning rides in ``context``; None -> the DiT runs the plain (degenerate) chunk forward. + dit = model.component(expert_id) # type: ignore[attr-defined] + preds = {b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_, context=chunk_context) for b in branches} + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + + def run(model, override=None): + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x_chunk, sigma_t, pe, ne, scale) + # Flow-match Euler on the CURRENT-CHUNK slice only (HYWorldDenoisingStage steps the chunk). + x_next = kernels.get(FLOW_MATCH_STEP)(precision.cast(x_chunk), velocity, sigma_t, sigma_next) + return StepResult( + output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": x_next.astype("float32"), + "chunk": (start, end) + }) + + 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_chunk.shape)), float(len(branches))), + resident_bytes=int(x_full.nbytes) + cond_bytes, + peak_activation_bytes=int(x_chunk.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_chunk.shape), + dtype=precision.compute_dtype, + extra=(("cfg", type(cfg).__name__), ("chunk", chunk_idx))), + resources=res, + payload={ + "branch": "combined", + "step": i, + "chunk": chunk_idx + }, + run=run, + label=f"hyworld.denoise.c{chunk_idx}.{i}", + # Chunk rollout + camera-aligned memory retrieval + heterogeneous per-frame timestep -> eager. + capturable=False) + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + # Write the denoised chunk slice back into the full latent (only this chunk's frames change). + start, end = result.output["chunk"] + x_full = np.asarray(st.latents["video"], dtype="float32") + x_full[:, start:end] = np.asarray(result.output["latents"], dtype="float32") + st.latents["video"] = x_full + if st.profile == ExecutionProfile.ROLLOUT: + st.trajectory.append({ + "step": st.step_idx, + "chunk": st.scratch["chunk_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 + if st.step_idx >= len(st.sigmas) - 1: # this chunk's flow-match sweep finished -> next chunk + st.scratch["chunk_idx"] += 1 + st.step_idx = 0 + st.plugin_state["cfg"] = {} # reset per-chunk CFG state + return st + + def finalize(self, st: LoopState) -> LoopResult: + return LoopResult(outputs={"latents": st.latents["video"]}, + metrics={"chunks": float(st.scratch["chunk_num"])}, + behavior=st.trajectory or None) diff --git a/v2/recipes/hyworld/program.py b/v2/recipes/hyworld/program.py new file mode 100644 index 000000000..c431cace6 --- /dev/null +++ b/v2/recipes/hyworld/program.py @@ -0,0 +1,45 @@ +"""HY-WorldPlay t2v program: text_encode -> hyworld_denoise (chunk rollout) -> vae_decode. + +Mirrors the cosmos2/wan inline node graph. The world-model specifics (3-stream conditioning, per-frame +camera/action, camera-aligned memory retrieval) live entirely in ``HYWorldDenoiseLoop`` + the +``HYWorldDiT`` adapter, so the registered (CPU-verified) degenerate t2v path needs only the standard +three nodes. + +BRINGUP hooks (left as no-op slot writers + documented, not wired, because v2's request API has no +pose-string / first-frame-image fields yet): + * ``camera_ctx`` — a pose string (``"w-31"``) expanded via ``pose.pose_to_input`` + ``compute_latent_num`` + to ``viewmats[T,4,4] / Ks[T,3,3] / action[T]`` plus a 50k-point sphere point cloud + (``generate_points_in_sphere``). A future ``_camera_encode`` node would write it into ``slots``; the + loop already reads ``ctx.slots.get("camera_ctx")``. + * ``i2v_img_embeds`` / ``i2v_cond`` — SigLIP image embeds + the VAE-encoded first-frame ``[cond|mask]`` + (33ch) latent. A future ``_image_encode`` node (SigLIP + VAE-encode of the first frame) would write + them; the loop already reads them. Absent -> the DiT sees zero image -> masks the image stream (t2v). +""" +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_hyworld_program() -> Program: + return Program( + program_id="hyworld.t2v.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), + ModelLoopNode("denoise", + loop_id="hyworld_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/recipes/kandinsky5/__init__.py b/v2/recipes/kandinsky5/__init__.py new file mode 100644 index 000000000..1fe11becc --- /dev/null +++ b/v2/recipes/kandinsky5/__init__.py @@ -0,0 +1,19 @@ +"""Kandinsky-5.0-T2V-Lite — text→video ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``Kandinsky5DiT`` / ``Kandinsky5QwenEncoder`` / ``Kandinsky5ClipEncoder`` / +``Kandinsky5VAE`` in ``v2/platform/backends/torch_kandinsky5.py``) and a new ``Kandinsky5DenoiseLoop`` +(flow-match Euler over a CHANNELS-LAST latent with a dual Qwen+CLIP conditioning stream). It reuses the +shared flow-match sampler (``FlowShiftPolicy.build_schedule`` + ``FLOW_MATCH_STEP``) and ``ClassicCFG``. +Registered in ``v2/registry.py`` by the orchestrator. +""" +from __future__ import annotations + +from v2.recipes.kandinsky5.card import KANDINSKY5_NEG, build_kandinsky5_card, stamp_kandinsky5_checkpoints +from v2.recipes.kandinsky5.loop import Kandinsky5DenoiseLoop +from v2.recipes.kandinsky5.program import build_kandinsky5_program + +__all__ = [ + "build_kandinsky5_card", "build_kandinsky5_program", "Kandinsky5DenoiseLoop", "stamp_kandinsky5_checkpoints", + "KANDINSKY5_NEG" +] diff --git a/v2/recipes/kandinsky5/card.py b/v2/recipes/kandinsky5/card.py new file mode 100644 index 000000000..586ac8565 --- /dev/null +++ b/v2/recipes/kandinsky5/card.py @@ -0,0 +1,206 @@ +"""Kandinsky-5.0-T2V-Lite ModelCard — text→video, ported into the v2 substrate. + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained — the +bucket-C pattern: torch adapters via ``ComponentSpec.adapter``, a new ``Kandinsky5DenoiseLoop``): + + * DiT ``fastvideo.models.dits.kandinsky5:Kandinsky5Transformer3DModel`` — a flow-match velocity + predictor, BUT with **channels-LAST** latent geometry ``[B, T, H, W, C]`` (C=4, NOT Wan's + channels-first ``[B, C, T, H, W]``) and TWO conditioning streams: Qwen2.5-VL token embeds + (``encoder_hidden_states``) + a CLIP **pooled** vector (``pooled_projections``, folded into the + time embedding — mandatory, the forward raises if it is None). The forward also requires per-request + RoPE positions (``visual_rope_pos`` over the patched grid, ``text_rope_pos`` over the Qwen seq len) + and a resolution-dependent ``scale_factor``. The adapter (-> ``Kandinsky5DiT``) builds all of these + internally so the loop's dit-call stays compatible with the CPU ``ToyDiT`` signature. + * VAE ``fastvideo.models.vaes.hunyuanvae:AutoencoderKLHunyuanVideo`` — HunyuanVideo VAE (z=4, 8×/4×) + with a SCALAR ``scaling_factor`` (latent = raw·sf; decode raw = latent/sf), NOT Wan's per-channel + ``(z-mean)/std`` — reusing the Wan VAE adapter would corrupt latents (-> ``Kandinsky5VAE``). + * Text ``fastvideo.models.encoders.qwen2_5:Qwen2_5_VLTextModel`` (Qwen embeds, in_text_dim=3584) + + ``fastvideo.models.encoders.clip:CLIPTextModelWithProjection`` (pooled, in_text_dim2=768). Two + text_encoder components (``text_encoder`` / ``text_encoder_2``) with distinct adapters. + * sched ``FlowMatchEulerDiscreteScheduler`` — plain flow-match Euler (velocity prediction, σ·1000 + timestep convention); CFG combines ``uncond + s·(cond − uncond)`` == ``ClassicCFG``. + +``stamp_kandinsky5_checkpoints`` (a local superset of ``stamp_wan21_checkpoints``) points each +component at its diffusers subfolder, adding ``text_encoder_2``/``tokenizer_2`` which the Wan stamp +lacks. ``flow_shift`` is read from the checkpoint's ``scheduler/scheduler_config.json`` at stamp time +(it is NOT hardcoded in the upstream code — the scheduler code default is 1.0). +""" +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.kandinsky5.loop import ( + KANDINSKY5_LATENT_CHANNELS, + KANDINSKY5_SPATIAL_RATIO, + KANDINSKY5_TEMPORAL_RATIO, + Kandinsky5DenoiseLoop, +) + +# Kandinsky-5 default negative prompt (kept as a LOCAL module constant, not in the shared _prompts.py, +# so the recipe is self-contained). From the diffusers pipeline / repo defaults. +KANDINSKY5_NEG = ("Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, " + "low quality, ugly, deformed, walking backwards") + +_KANDINSKY5_DIT = "v2.platform.backends.torch_kandinsky5:Kandinsky5DiT" +_KANDINSKY5_QWEN = "v2.platform.backends.torch_kandinsky5:Kandinsky5QwenEncoder" +_KANDINSKY5_CLIP = "v2.platform.backends.torch_kandinsky5:Kandinsky5ClipEncoder" +_KANDINSKY5_VAE = "v2.platform.backends.torch_kandinsky5:Kandinsky5VAE" + + +def build_kandinsky5_card(model_id: str = "kandinsky-5.0-t2v-lite-5s", + *, + checkpoint_root: str | None = None, + flow_shift: float = 1.0, + 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() + # flow_shift default is 1.0 (the scheduler code default); the real value is read from the + # checkpoint's scheduler_config.json by stamp_kandinsky5_checkpoints at GPU deploy time. + flow = FlowShiftPolicy(shift=flow_shift) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return Kandinsky5DenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost, + latent_channels=KANDINSKY5_LATENT_CHANNELS, + spatial_ratio=KANDINSKY5_SPATIAL_RATIO, + temporal_ratio=KANDINSKY5_TEMPORAL_RATIO) + + components = { + # Qwen2.5-VL: token embeds [seq, 3584] -> encoder_hidden_states (the cross-attended text stream). + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.qwen2_5:Qwen2_5_VLTextModel", + adapter=_KANDINSKY5_QWEN, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + # CLIP-with-projection: pooled vector [768] -> pooled_projections (folded into the time embed). + "text_encoder_2": + ComponentSpec(component_id="text_encoder_2", + kind="text_encoder", + load_id="fastvideo.models.encoders.clip:CLIPTextModelWithProjection", + adapter=_KANDINSKY5_CLIP, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.hunyuanvae:AutoencoderKLHunyuanVideo", + adapter=_KANDINSKY5_VAE, + factory=lambda inst: ToyVAE(channels=KANDINSKY5_LATENT_CHANNELS), + required_for={"t2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.kandinsky5:Kandinsky5Transformer3DModel", + adapter=_KANDINSKY5_DIT, + factory=lambda inst: ToyDiT(channels=KANDINSKY5_LATENT_CHANNELS, 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="kandinsky5", + 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=50, + guidance_scale=5.0, + height=512, + width=768, + num_frames=121, + fps=24, + negative_prompt=KANDINSKY5_NEG), + ) + card.validate() + if checkpoint_root: + stamp_kandinsky5_checkpoints(card, checkpoint_root) + return card + + +# Diffusers checkpoint layout for Kandinsky-5: a superset of the Wan layout adding the SECOND text +# encoder + its tokenizer (the shared _WAN21_SUBFOLDERS only has one text_encoder/tokenizer). +_KANDINSKY5_SUBFOLDERS = { + "transformer": "transformer", + "vae": "vae", + "text_encoder": "text_encoder", + "text_encoder_2": "text_encoder_2" +} + + +def stamp_kandinsky5_checkpoints(card: ModelCard, model_root: str) -> ModelCard: + """GPU deploy-time: point each component's ``ComponentSpec.checkpoint`` at its weights subfolder under + ``model_root`` (a local diffusers dir or an HF id resolved via ``snapshot_download``), and read the + real ``flow_shift`` from ``scheduler/scheduler_config.json`` into the loop's FlowShiftPolicy (it is NOT + hardcoded upstream — code default 1.0; BRINGUP risk J). The CLIP encoder's tokenizer lives in + ``tokenizer_2`` (NOT the ``tokenizer`` the built-in text-encoder maker probes) — the Kandinsky5ClipEncoder + adapter re-resolves it from the model root (BRINGUP). Mutates and returns the card.""" + import json + import os + if not os.path.isdir(model_root): + from huggingface_hub import snapshot_download + model_root = snapshot_download(model_root) + for component_id, subfolder in _KANDINSKY5_SUBFOLDERS.items(): + if component_id in card.components: + card.components[component_id].checkpoint = os.path.join(model_root, subfolder) + # BRINGUP J: lift flow_shift from the scheduler config so >1.0 schedules are honored on GPU. + sched_cfg = os.path.join(model_root, "scheduler", "scheduler_config.json") + if os.path.exists(sched_cfg): + with open(sched_cfg) as f: + shift = json.load(f).get("shift") + if shift is not None: + loop_spec = card.loops.get("diffusion_denoise") + base_factory = loop_spec.loop_factory if loop_spec is not None else None + if base_factory is not None: + + def _shifted_factory(_base=base_factory, _shift=float(shift)): + loop = _base() + loop.flow_shift.shift = _shift + return loop + + loop_spec.loop_factory = _shifted_factory + return card diff --git a/v2/recipes/kandinsky5/loop.py b/v2/recipes/kandinsky5/loop.py new file mode 100644 index 000000000..530aefd14 --- /dev/null +++ b/v2/recipes/kandinsky5/loop.py @@ -0,0 +1,235 @@ +"""Kandinsky5DenoiseLoop — flow-match Euler over a CHANNELS-LAST video latent with a dual text stream. + +Faithful port of the Kandinsky-5.0 T2V denoise loop (diffusers ``KandinskyV5...Pipeline.__call__``): +N flow-match Euler steps over the full latent, velocity prediction, ``ClassicCFG`` combine +(``uncond + s·(cond − uncond)``), σ·1000 timestep convention. It mirrors ``WanDenoiseLoop`` but +diverges on the three things Kandinsky needs (BRINGUP risks A/B/D from the port spec): + + 1. **Channels-LAST DiT geometry, but a channels-FIRST loop latent.** The DiT forward consumes + ``hidden_states[B, T, H, W, C]`` (C=4) and returns a channels-last velocity. The loop, however, + keeps the latent channels-FIRST ``(c, t, h, w)`` — the same convention as Wan and the CPU + ``ToyDiT`` (which channel-mixes over axis 0). The ``Kandinsky5DiT`` adapter does the + channels-first → channels-last permute (and permutes the velocity back) INTERNALLY, so the loop + and toy stay backend-agnostic and the GPU forward still sees the geometry it requires. + 2. **A SECOND conditioning slot.** Qwen2.5-VL token embeds go to one slot (the cross-attended text); + the CLIP **pooled** vector goes to a second slot and is threaded to the dit-call as ``context`` (so + the dit-call signature stays ``dit(latent, text_embed, sigma, context=pooled)`` — compatible with the + CPU ``ToyDiT``). The ``Kandinsky5DiT`` adapter feeds ``context`` to ``pooled_projections``, which is + MANDATORY (the forward raises if it is None). + 3. **Per-request RoPE positions + scale_factor are built INSIDE the adapter** from the latent geometry + and resolution (``visual_rope_pos`` over the patched grid, ``text_rope_pos`` over the Qwen seq len, + ``scale_factor`` = (1,2,2) for 480–854 else (1,3.16,3.16)). Keeping that in the adapter rather than + threading extra kwargs is what lets the loop call the toy-compatible 3-arg dit signature. + +The flow-match math (``FlowShiftPolicy.build_schedule`` → ``FLOW_MATCH_STEP``) and the CFG combine are +reused verbatim from the Wan recipe — they are identical to Kandinsky's. Streaming preview + the +ROLLOUT trajectory capture mirror ``WanDenoiseLoop``. +""" +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.platform import FLOW_MATCH_STEP +from v2.platform.backends.toy import LATENT_CHANNELS +from v2.request.streams import StreamChunk + +# AutoencoderKLHunyuanVideo compression for Kandinsky-5: z=4 latent channels, 8× spatial, 4× temporal. +KANDINSKY5_LATENT_CHANNELS = 4 +KANDINSKY5_TEMPORAL_RATIO = 4 +KANDINSKY5_SPATIAL_RATIO = 8 + + +def latent_shape(req, + model=None, + *, + channels=KANDINSKY5_LATENT_CHANNELS, + spatial_ratio=KANDINSKY5_SPATIAL_RATIO, + temporal_ratio=KANDINSKY5_TEMPORAL_RATIO) -> tuple[int, int, int, int]: + """CHANNELS-FIRST latent geometry ``(c, t, h, w)`` for the denoise loop (Wan/ToyDiT convention — the + toy channel-mixes over axis 0). On the GPU backend the true geometry is ``c = 4``, + ``t = (num_frames-1)//4 + 1``, ``h = H//8``, ``w = W//8``; the ``Kandinsky5DiT`` adapter permutes this + to the channels-LAST ``[B, T, H, W, C]`` the forward requires. The CPU toy uses a tiny stand-in. + ``model is None`` (cost estimation in tests) keeps the toy shape.""" + d = req.diffusion + if model is not None and getattr(getattr(model, "platform", None), "device", "cpu") == "cuda": + t = (max(1, d.num_frames) - 1) // temporal_ratio + 1 + return (channels, max(1, t), max(1, d.height // spatial_ratio), max(1, d.width // spatial_ratio)) + # toy stand-in: small, channels-first, with even h,w so the patch_size=(1,2,2) grid stays valid. + t = max(1, d.num_frames // 40) + h = max(2, d.height // 120) + w = max(2, d.width // 120) + return (LATENT_CHANNELS, t, h - (h % 2), w - (w % 2)) + + +class Kandinsky5DenoiseLoop: + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels=KANDINSKY5_LATENT_CHANNELS, + spatial_ratio=KANDINSKY5_SPATIAL_RATIO, + temporal_ratio=KANDINSKY5_TEMPORAL_RATIO): + self.loop_id = loop_id + self.cfg = cfg + self.flow_shift = flow_shift + self.precision = precision + self.expert = expert + self.cost = cost + 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 = self.flow_shift.build_schedule(req.diffusion.num_steps, + req.diffusion.height, + req.diffusion.width, + sigmas=req.diffusion.sigmas or None) + 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(sig[0])).astype("float32") + 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]) + # Qwen token embeds (the cross-attended text stream) + the matching negative. + st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") + st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") + # CLIP pooled vectors (the MANDATORY pooled_projections, threaded to the dit-call as `context`). + st.cond["pooled"] = ctx.slots.get("clip_pooled") + st.cond["negative_pooled"] = ctx.slots.get("neg_clip_pooled") + st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale) + st.scratch["stream_video"] = bool(req.outputs.stream.get("video")) + 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] + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + branches = self.cfg.branches_this_step(sctx, cfg_state) + x = st.latents["video"] + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + pp, npp = st.cond["pooled"], st.cond["negative_pooled"] + scale = st.scratch["guidance_scale"] + cfg, precision = self.cfg, self.precision + + def _velocity(model: Any, x_: Any, sigma_t_: float, pe_: Any, ne_: Any, pp_: Any, npp_: Any, + scale_: float) -> np.ndarray: + # The conditioned forward + CFG combine. The CLIP pooled vector rides in ``context`` (the + # toy ignores it harmlessly; the Kandinsky5DiT adapter routes it to pooled_projections and + # builds visual/text RoPE + scale_factor internally from the latent geometry + resolution). + dit = model.component(expert_id) + preds = { + b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_, context=(pp_ if b == "cond" else npp_)) + for b in branches + } + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + + def run(model, override=None): + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x, sigma_t, pe, ne, pp, npp, scale) + x_next = 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") + }) + + def graph_fn(model, ws): + # The CAPTURABLE op-structure: reads every per-step input from the static workspace. + st_t, st_n = float(ws["sigma_t"]), float(ws["sigma_next"]) + velocity = _velocity(model, ws["x"], st_t, ws["pe"], ws["ne"], ws["pp"], ws["npp"], float(ws["scale"])) + x_next = model.platform.kernels.get(FLOW_MATCH_STEP)(precision.cast(ws["x"]), velocity, st_t, st_n) + np.copyto(ws["out"], x_next.astype("float32")) + return StepResult(output={ + "noise_pred": np.asarray(velocity, dtype="float32"), + "latents": np.array(ws["out"], copy=True) + }) + + cond_bytes = sum(int(np.asarray(e).nbytes) for e in (pe, ne, pp, npp) if e is not None) + res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), + resident_bytes=int(x.nbytes) + cond_bytes, + peak_activation_bytes=int(x.nbytes)) + emits = [] + if st.scratch.get("stream_video"): + emits.append(StreamChunk(stream_id=st.request_id, modality="video", seq=i, data=x, preview=True)) + 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(cfg).__name__), )), + resources=res, + payload={ + "branch": "combined", + "step": i + }, + run=run, + label=f"kandinsky5.denoise.{i}", + emits=emits, + capturable=True, + graph_key=(tuple(sorted(branches)), expert_id, precision.scheduler_step_in_fp32), + graph_fn=graph_fn, + graph_inputs={ + "x": x, + "sigma_t": sigma_t, + "sigma_next": sigma_next, + "pe": pe, + "ne": ne, + "pp": pp, + "npp": npp, + "scale": scale + }) + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + st.latents["video"] = result.output["latents"] + if st.profile == ExecutionProfile.ROLLOUT: + i = st.step_idx + st.trajectory.append({ + "step": i, + "sigma": st.sigmas[i], + "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/kandinsky5/program.py b/v2/recipes/kandinsky5/program.py new file mode 100644 index 000000000..eb045b7e7 --- /dev/null +++ b/v2/recipes/kandinsky5/program.py @@ -0,0 +1,52 @@ +"""Kandinsky-5 T2V program: encode(Qwen + CLIP) → diffusion_denoise → vae_decode. + +Same inline shape as the Wan/Cosmos t2v programs, but the encode node fans out to BOTH text encoders +(the dual conditioning stream Kandinsky needs): Qwen2.5-VL token embeds for the cross-attended text and +the CLIP **pooled** vector for ``pooled_projections``. The denoise loop reads all four slots +(prompt/neg × qwen/clip). The vae-decode node hands the channels-last latent straight to the +``Kandinsky5VAE`` adapter, which does the ``[B,T,H,W,C] -> [B,C,T,H,W]`` reshape + scalar +``/scaling_factor`` un-normalization internally (faithful to the diffusers pipeline post-processing). +""" +from __future__ import annotations + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind +from v2.recipes.common import cached_text_encode + + +def _encode(instance, slots, request, ctx) -> None: + # Qwen token embeds (cross-attended text) — cached by content hash like every text-encode node. + slots["text_embeds"] = cached_text_encode(instance, request.prompt()) + slots["neg_text_embeds"] = cached_text_encode(instance, request.diffusion.negative_prompt) + # CLIP pooled vectors (the mandatory pooled_projections). Encoded directly (the feature cache keys on + # component_id "text_encoder"; the CLIP component is a distinct instance so we call it explicitly). + clip = instance.component("text_encoder_2") + slots["clip_pooled"] = clip.encode(request.prompt()) + slots["neg_clip_pooled"] = clip.encode(request.diffusion.negative_prompt) + + +def _vae_decode(instance, slots, request, ctx) -> None: + # The channels-last latent [T,H,W,C] (loop output) -> the Kandinsky5VAE adapter reshapes to + # channels-first and applies the scalar scaling_factor un-normalization before decoding. + slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_kandinsky5_program() -> Program: + return Program( + program_id="kandinsky5.t2v.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("text_encode", + fn=_encode, + writes=("text_embeds", "neg_text_embeds", "clip_pooled", "neg_clip_pooled")), + ModelLoopNode("denoise", + loop_id="diffusion_denoise", + output_slot="denoise_out", + reads=("text_embeds", "neg_text_embeds", "clip_pooled", "neg_clip_pooled"), + 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/recipes/lingbotworld/__init__.py b/v2/recipes/lingbotworld/__init__.py new file mode 100644 index 000000000..c21a97846 --- /dev/null +++ b/v2/recipes/lingbotworld/__init__.py @@ -0,0 +1,16 @@ +"""LingBot-World-Base-Cam (camera-conditioned dual-guidance MoE i2v) ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapter via +``ComponentSpec.adapter`` (``LingBotWorldDiT`` in ``v2/platform/backends/torch_lingbotworld.py``) and a +new ``LingBotWorldDenoiseLoop`` (Wan i2v flow-match + boundary-routed MoE + DUAL guidance + camera/Plucker +conditioning), reusing the Wan VAE / T5 / CLIP adapters + ``stamp_wan21_checkpoints``. The program adds a +``camera_encode`` node (Plucker tensor; BRINGUP — needs a per-request camera input). Registered in +``v2/registry.py`` by the orchestrator from the returned metadata. +""" +from __future__ import annotations + +from v2.recipes.lingbotworld.card import build_lingbotworld_card +from v2.recipes.lingbotworld.loop import LingBotWorldDenoiseLoop +from v2.recipes.lingbotworld.program import build_lingbotworld_program + +__all__ = ["build_lingbotworld_card", "build_lingbotworld_program", "LingBotWorldDenoiseLoop"] diff --git a/v2/recipes/lingbotworld/card.py b/v2/recipes/lingbotworld/card.py new file mode 100644 index 000000000..d6c5ea284 --- /dev/null +++ b/v2/recipes/lingbotworld/card.py @@ -0,0 +1,157 @@ +"""LingBot-World-Base-Cam ModelCard — camera-conditioned, dual-guidance MoE image-to-video. + +LingBot-World-Base-Cam (``FastVideo/LingBot-World-Base-Cam-Diffusers``) is Wan2.2-I2V-A14B with a NEW +camera/Plucker conditioning input. Architecture (all declared on the card so the recipe is self-contained +— no edit to the shared ``_make_dit`` dispatch): + + * DiT (x2) ``fastvideo.models.dits.lingbotworld.model:LingBotWorldTransformer3DModel`` — a 2x14B + boundary-routed MoE (``transformer`` high-noise, ``transformer_2`` low-noise; ``boundary_ratio=0.947``). + The forward gains ``c2ws_plucker_emb`` (FiLM camera injection per block) which a vanilla WanDiT reuse + would drop, so both experts use the ``LingBotWorldDiT`` adapter (``ComponentSpec.adapter``) that threads + the Plucker tensor + keeps the Wan i2v 36ch ``[noise|mask+cond]`` concat + CLIP image embeds. + * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` (z=16, 8x spatial / 4x temporal, mean/std + normalized latent space) — reuses the v2 ``WanVAE`` adapter unchanged. ``spatial_scale=8`` for the + Plucker downsample matches this 8x spatial compression. + * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (UMT5, text_len=512) — reuses ``T5Encoder``. + * Image ``fastvideo.models.encoders.clip:CLIPVisionModel`` — first-frame CLIP conditioning (i2v), reuses + ``CLIPImageEncoder``. + +Sampler: flow-match (``FlowUniPCMultistepScheduler``, ``prediction_type='flow_prediction'``) -> the v2 +``FlowShiftPolicy`` + ``FLOW_MATCH_STEP``, exactly as Wan. NOTE the unusual defaults: ``flow_shift=10.0``, +``num_inference_steps=70``, dual ``guidance_scale``/``guidance_scale_2`` (both 5.0), and the Chinese +negative prompt — these must NOT inherit the Wan defaults. + +``stamp_wan21_checkpoints`` applies (diffusers ``transformer``/``transformer_2``/``vae``/``text_encoder``/ +``image_encoder`` subfolder layout, all in ``_WAN21_SUBFOLDERS``). +""" +from __future__ import annotations + +from v2._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind +from v2.card import ( + CacheContract, + CapabilityMatrix, + ComponentSpec, + CostModel, + LoopSpec, + ModelCard, + ParallelismContract, + ParitySpec, + PrecisionContract, + RecipeSpec, + SamplingDefaults, +) +from v2.loop.policies import BoundaryTimestepRouting, ClassicCFG, FlowShiftPolicy, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyImageEncoder, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.lingbotworld.loop import LINGBOTWORLD_BOUNDARY, LingBotWorldDenoiseLoop +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# The LingBot-World I2V negative prompt (presets.py LINGBOTWORLD_I2V) — kept LOCAL to this recipe so the +# port stays self-contained (no edit to the shared v2/recipes/_prompts.py). +LINGBOTWORLD_NEG_CN = ("画面突变，色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，" + "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，" + "畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走，" + "镜头晃动，画面闪烁，模糊，噪点，水印，签名，文字，变形，扭曲，液化，不合逻辑的结构，卡顿，" + "PPT幻灯片感，过暗，欠曝，低对比度，霓虹灯光感，过度锐化，3D渲染感，人物，行人，游客，身体，" + "皮肤，肢体，面部特征，汽车，电线") + +_LINGBOTWORLD_DIT = "v2.platform.backends.torch_lingbotworld:LingBotWorldDiT" + + +def build_lingbotworld_card(model_id: str = "lingbot-world-base-cam", + *, + boundary: float = LINGBOTWORLD_BOUNDARY, + checkpoint_root: str | None = None, + sampling_defaults: SamplingDefaults | None = None) -> ModelCard: + """LingBot-World-Base-Cam card — 2x14B camera-conditioned dual-guidance MoE i2v. + + Mirrors ``build_wan22_i2v_a14b_card`` but: ``flow_shift=10.0``; ``boundary_ratio=0.947``; both DiT + experts use the ``LingBotWorldDiT`` adapter (threads ``c2ws_plucker_emb`` + Wan i2v cond); the + ``LingBotWorldDenoiseLoop`` (dual guidance + camera); and the LingBot defaults (70 steps, gs 5.0, + 480x832, Chinese negative prompt). 2x14B bf16 won't co-reside on 80GB -> the adapter keeps the WanDiT + CPU-offload swap (GPU-pending).""" + seed = _seed_from(model_id) + cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7) + cfg = ClassicCFG() + flow = FlowShiftPolicy(shift=10.0) # LingBot-World's unusually high flow shift (not the Wan 3.0/5.0) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = BoundaryTimestepRouting(high_noise="transformer", low_noise="transformer_2", boundary=boundary) + + def loop_factory() -> LingBotWorldDenoiseLoop: + return LingBotWorldDenoiseLoop(loop_id="i2v_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost, + boundary=boundary) + + def _dit(cid: str) -> ComponentSpec: + return ComponentSpec(cid, + kind="dit", + load_id="fastvideo.models.dits.lingbotworld.model:LingBotWorldTransformer3DModel", + adapter=_LINGBOTWORLD_DIT, + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["i2v_denoise"], + required_for={"i2v"}) + + components = { + "text_encoder": + ComponentSpec("text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + factory=lambda inst: ToyTextEncoder(), + required_for={"i2v"}), + "image_encoder": + ComponentSpec("image_encoder", + kind="image_encoder", + load_id="fastvideo.models.encoders.clip:CLIPVisionModel", + factory=lambda inst: ToyImageEncoder(), + required_for={"i2v"}), + "vae": + ComponentSpec("vae", + kind="vae", + load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan", + factory=lambda inst: ToyVAE(), + required_for={"i2v"}), + "transformer": + _dit("transformer"), + "transformer_2": + _dit("transformer_2"), + } + loops = { + "i2v_denoise": + LoopSpec("i2v_denoise", + kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, + step_cost_model=cost, + shared_weight_components=["transformer", "transformer_2"], + cache_policy=["feature"], + loop_factory=loop_factory), + } + card = ModelCard( + model_id=model_id, + family="wan", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.IMAGE_TO_VIDEO, Capability.VAE_DECODE), + recipe=RecipeSpec(method="base", + assumes_loop="i2v_denoise", + assumes_precision="float32", + consistency_required=ConsistencyLevel.C1), + parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True), + caches={"feature": CacheContract("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=70, + guidance_scale=5.0, + height=480, + width=832, + num_frames=81, + fps=16, + negative_prompt=LINGBOTWORLD_NEG_CN), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card diff --git a/v2/recipes/lingbotworld/loop.py b/v2/recipes/lingbotworld/loop.py new file mode 100644 index 000000000..68bab8a90 --- /dev/null +++ b/v2/recipes/lingbotworld/loop.py @@ -0,0 +1,144 @@ +"""LingBotWorldDenoiseLoop — Wan flow-match denoise + boundary-routed MoE, extended with (a) the camera +(Plucker) conditioning slot threaded to the DiT and (b) DUAL classifier-free guidance. + +LingBot-World-Base-Cam is Wan2.2-I2V-A14B i2v conditioning (CLIP + first-frame ``[mask|cond]`` latent) +plus a camera/Plucker tensor. The denoise math is otherwise the canonical Wan flow-match Euler step with +``BoundaryTimestepRouting`` between the high-noise ``transformer`` and the low-noise ``transformer_2``. Two +deltas vs ``WanDenoiseLoop`` (faithful to ``LingBotWorldImageToVideoPipeline`` == the Wan i2v denoise +stage, denoising.py:352-375): + + * DUAL guidance — fastvideo switches ``guidance_scale -> guidance_scale_2`` at the SAME expert boundary: + ``t >= boundary_timestep`` (the high-noise expert) uses ``guidance_scale``; below it (the low-noise + expert) uses ``guidance_scale_2``. The boundary is in timestep space ``boundary_ratio*1000``, i.e. + sigma-space ``ctx.sigma >= boundary`` (preset ``boundary_ratio=0.947``). The preset sets both scales + to 5.0, so the registered path is numerically a single scale — but the loop honors a distinct + ``guidance_scale_2`` whenever the request carries one. + * CAMERA — the program's camera node writes ``c2ws_plucker_emb`` into slots; the loop publishes it onto + the active DiT adapter (``dit.c2ws_plucker_emb = ...``) before the forward, because the v2 dit-call + surface (and ``ToyDiT``) has no such kwarg. On the CPU toy backend the slot is ``None`` -> the + degenerate no-camera path (a plain Wan i2v step), so the loop CPU-verifies end-to-end. + +BRINGUP (needs request-API extension): the Plucker tensor is built from a per-request camera trajectory +(``poses.npy``/``intrinsics.npy``); v2's ``Request`` has no camera slot yet, so the program node writes +``None`` unless an ``action_path`` override is supplied. GPU-verify (a) the 2x14B MoE CPU-offload swap, +(b) the camera FiLM injection actually steers the generation, (c) the dual-scale boundary parity. +""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from v2.loop.contracts import Done, LoopState, StepContext, StepResult, WorkPlan +from v2.recipes.wan21.loop import WanDenoiseLoop + +# Wan2.2 boundary expressed in (shifted) sigma space; LingBot-World uses 0.947 (much higher than Wan2.2's +# 0.875/0.9), so the low-noise expert ``transformer_2`` runs almost the whole trajectory. +LINGBOTWORLD_BOUNDARY = 0.947 + + +class LingBotWorldDenoiseLoop(WanDenoiseLoop): + """Wan i2v denoise + camera Plucker conditioning + dual-guidance boundary switch.""" + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + boundary=LINGBOTWORLD_BOUNDARY, + latent_channels=16, + spatial_ratio=8, + temporal_ratio=4): + super().__init__(loop_id=loop_id, + cfg=cfg, + flow_shift=flow_shift, + precision=precision, + expert=expert, + cost=cost, + latent_channels=latent_channels, + spatial_ratio=spatial_ratio, + temporal_ratio=temporal_ratio) + # The boundary the dual-guidance switch keys on; default to the routing policy's own boundary so the + # CFG switch and the expert switch fire on the SAME step (matching the fastvideo denoise stage). + self.boundary = float(getattr(expert, "boundary", boundary)) + + def init(self, req, model, ctx) -> LoopState: + st = super().init(req, model, ctx) + # Dual guidance: the second CFG scale used below the boundary (low-noise expert). Fall back to the + # primary scale when the request doesn't carry a distinct guidance_scale_2 (preset sets both 5.0). + gs2 = getattr(req.diffusion, "guidance_scale_2", None) + st.scratch["guidance_scale_2"] = float(gs2) if gs2 is not None else float(req.diffusion.guidance_scale) + # Camera / Plucker conditioning written by the program's camera node (None == no-camera path). + st.scratch["c2ws_plucker_emb"] = ctx.slots.get("c2ws_plucker_emb") + 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] + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + branches = self.cfg.branches_this_step(sctx, cfg_state) + x = st.latents["video"] + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + # DUAL guidance: high-noise (ctx.sigma>=boundary) -> guidance_scale; low-noise -> guidance_scale_2. + scale = (st.scratch["guidance_scale"] if sigma_t >= self.boundary else st.scratch["guidance_scale_2"]) + cfg, precision = self.cfg, self.precision + i2v_ctx, i2v_cond = st.scratch.get("i2v_img_embeds"), st.scratch.get("i2v_cond") + plucker = st.scratch.get("c2ws_plucker_emb") + + def _velocity(model: Any, x_: np.ndarray, sigma_t_: float, pe_: Any, ne_: Any, scale_: float) -> np.ndarray: + # The conditioned forward + CFG combine. The active expert's adapter gets the per-step camera + # tensor published out-of-band (the v2 dit-call surface has no c2ws_plucker_emb kwarg); ToyDiT + # simply ignores the attribute, so the CPU path is the no-camera step. + dit = model.component(expert_id) + if plucker is not None and hasattr(dit, "c2ws_plucker_emb"): + dit.c2ws_plucker_emb = plucker + preds = { + b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_, context=i2v_ctx, cond=i2v_cond) + for b in branches + } + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + + def run(model: Any, override: dict | None = None) -> StepResult: + from v2.platform import FLOW_MATCH_STEP + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x, sigma_t, pe, ne, scale) + x_next = 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") + }) + + from v2._enums import WorkUnitKind + from v2.loop.contracts import ResourceRequest, ShapeSignature + 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)), float(len(branches))), + 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(cfg).__name__), ("camera", plucker is not None))), + resources=res, + payload={ + "branch": "combined", + "step": i + }, + run=run, + label=f"lingbotworld.denoise.{i}", + # Camera (out-of-band adapter state) + dual-scale boundary thread non-workspace conditioning, + # exactly like Wan i2v -> eager path (no CUDA-graph capture). + capturable=False) diff --git a/v2/recipes/lingbotworld/program.py b/v2/recipes/lingbotworld/program.py new file mode 100644 index 000000000..862d6266e --- /dev/null +++ b/v2/recipes/lingbotworld/program.py @@ -0,0 +1,99 @@ +"""LingBot-World i2v program: image_encode (CLIP) -> cond_encode (first-frame VAE [mask|cond]) -> +camera_encode (Plucker) -> text_encode -> i2v_denoise (dual-guidance MoE) -> vae_decode. + +Same inline shape as the Wan i2v program (the i2v conditioning is identical to Wan2.2-I2V-A14B), with one +NEW node: ``camera_encode`` builds the camera/Plucker tensor ``c2ws_plucker_emb`` the LingBot-World DiT +FiLM-injects per block. The ``LingBotWorldDenoiseLoop`` reads that slot and publishes it onto the active +expert adapter; the dual-guidance / MoE specifics live in the loop + adapter, so the node graph is a thin +superset of Wan i2v. + +BRINGUP (blocker #2 — needs a request-API extension): v2's ``Request`` has no camera-trajectory input +slot. The camera node therefore looks for an ``action_path`` in the node override +(``request.node_override("camera_encode")["action_path"]``, pointing at a dir with ``poses.npy`` + +``intrinsics.npy``) and, if present, builds the Plucker tensor via the fastvideo +``prepare_camera_embedding`` (``spatial_scale=8`` to match the VAE 8x). Absent (the default, and the CPU +toy path) -> it writes ``None`` and the loop runs the degenerate no-camera i2v step. ``prepare_camera_embedding`` +also re-clamps ``num_frames`` to the trajectory length (4k+1); the v2 frame count is taken from the +request, so on GPU the camera node and the request num_frames must agree (verify at bring-up). +""" +from __future__ import annotations + +import numpy as np + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind +from v2.recipes.common import text_encode_node_fn as _text_encode + +_WAN_TEMPORAL_RATIO = 4 # VAE temporal compression -> the i2v mask has this many channels +_SPATIAL_SCALE = 8 # Plucker spatial downsample == VAE 8x spatial compression (WanCamCtrl in_chans=6*64) + + +def _i2v_image_encode(instance, slots, request, ctx) -> None: + img = request.image() + slots["i2v_img_embeds"] = (instance.component("image_encoder").encode_image(img.pixels) + if img is not None and getattr(img, "pixels", None) is not None else None) + + +def _i2v_cond_encode(instance, slots, request, ctx) -> None: + """First-frame VAE conditioning + mask (mirrors fastvideo's ImageVAEEncodingStage): encode the + conditioning image as frame 0 (rest zeros), prepend a 4-channel mask (first latent frame = 1).""" + img = request.image() + if img is None or getattr(img, "pixels", None) is None: + slots["i2v_cond"] = None + return + px = np.asarray(img.pixels, dtype="float32") # [3, H, W] in [-1, 1] + nf = int(request.diffusion.num_frames) + cond_video = np.zeros((px.shape[0], nf) + px.shape[1:], dtype="float32") + cond_video[:, 0] = px # frame 0 = conditioning image, rest zeros + cond_latent = np.asarray(instance.component("vae").encode(cond_video), dtype="float32") # [C, T, h, w] + t, h, w = cond_latent.shape[1:] + mask = np.zeros((_WAN_TEMPORAL_RATIO, t, h, w), dtype="float32") + mask[:, 0] = 1.0 # first latent frame is the known (conditional) frame + slots["i2v_cond"] = np.concatenate([mask, cond_latent], axis=0) # [mask | cond] -> DiT concats with noise + + +def _camera_encode(instance, slots, request, ctx) -> None: + """Build the camera/Plucker tensor ``c2ws_plucker_emb`` from a per-request camera trajectory. + + BRINGUP: until ``Request`` carries a camera-trajectory input, the trajectory dir is read from the + node override ``action_path``. Absent -> ``None`` (the no-camera path the CPU toy backend takes).""" + action_path = request.node_override("camera_encode").get("action_path") + if not action_path: + slots["c2ws_plucker_emb"] = None + return + from v2.models.dits.lingbotworld import prepare_camera_embedding + d = request.diffusion + emb, _num_frames = prepare_camera_embedding(action_path, + int(d.num_frames), + int(d.height), + int(d.width), + spatial_scale=_SPATIAL_SCALE) + # Drop the leading batch dim -> the loop/adapter add it back ([6*s^2, F_lat, H_lat, W_lat]). + arr = np.asarray(emb.squeeze(0).cpu().numpy() if hasattr(emb, "cpu") else emb, dtype="float32") + slots["c2ws_plucker_emb"] = arr + + +def _i2v_vae_decode(instance, slots, request, ctx) -> None: + slots["video"] = instance.component("vae").decode(slots["i2v_out"]["latents"]) + + +def build_lingbotworld_program() -> Program: + return Program( + program_id="lingbotworld.i2v", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("image_encode", fn=_i2v_image_encode, writes=("i2v_img_embeds", )), + ComponentNode("cond_encode", fn=_i2v_cond_encode, writes=("i2v_cond", )), + ComponentNode("camera_encode", fn=_camera_encode, writes=("c2ws_plucker_emb", )), + ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), + ModelLoopNode("denoise", + loop_id="i2v_denoise", + output_slot="i2v_out", + reads=("text_embeds", "i2v_cond", "i2v_img_embeds", "c2ws_plucker_emb"), + writes=("i2v_out", )), + ComponentNode("vae_decode", fn=_i2v_vae_decode, reads=("i2v_out", ), writes=("video", )), + ], + output_artifacts={ + "video": "video", + "latents": "i2v_out" + }, + ).validate() diff --git a/v2/recipes/longcat/__init__.py b/v2/recipes/longcat/__init__.py new file mode 100644 index 000000000..5d0fb35d4 --- /dev/null +++ b/v2/recipes/longcat/__init__.py @@ -0,0 +1,16 @@ +"""LongCat-Video (T2V) — flow-match denoiser with CFG-zero guidance, ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapter via +``ComponentSpec.adapter`` (``LongCatDiT`` in ``v2/platform/backends/torch_longcat.py``) and a new +``LongCatDenoiseLoop`` (explicit ``linspace(1.0, 0.001, num_steps)`` sigma schedule + CFG-zero +optimized-scale combine), reusing the Wan ``WanVAE`` + ``T5Encoder`` adapters and ``stamp_wan21_checkpoints``. +The DiT adapter negates the velocity to fold in the fastvideo stage's ``noise_pred = -noise_pred``. +Registered in ``v2/registry.py`` by the orchestrator (transformer ``LongCatTransformer3DModel``). +""" +from __future__ import annotations + +from v2.recipes.longcat.card import build_longcat_card +from v2.recipes.longcat.loop import CFGZeroPolicy, LongCatDenoiseLoop +from v2.recipes.longcat.program import build_longcat_program + +__all__ = ["build_longcat_card", "build_longcat_program", "LongCatDenoiseLoop", "CFGZeroPolicy"] diff --git a/v2/recipes/longcat/card.py b/v2/recipes/longcat/card.py new file mode 100644 index 000000000..fdf04357c --- /dev/null +++ b/v2/recipes/longcat/card.py @@ -0,0 +1,131 @@ +"""LongCat-Video-T2V ModelCard — text->video, flow-match with CFG-zero guidance. + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.longcat:LongCatTransformer3DModel`` (depth 48, hidden 4096, 32 heads, + in/out 16ch, patch [1,2,2], caption_channels 4096). The adapter (``ComponentSpec.adapter`` -> ``LongCatDiT`` + in ``v2/platform/backends/torch_longcat.py``) returns the **negated** velocity (the fastvideo stage's + ``noise_pred = -noise_pred`` before the scheduler step) so the loop's flow-match Euler reproduces it. + * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — Wan-style (z=16, 8x/4x) in NORMALIZED latent + space (mean/std). Reuses the v2 ``WanVAE`` torch adapter UNCHANGED (no adapter override). + * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (UMT5, 4096-dim) via the v2 ``T5Encoder`` (zero-pads + to max_length=512 for the CFG-concat uniform-seq contract — no adapter override needed). + * Sampler: ``FlowMatchEulerDiscreteScheduler`` with an EXPLICIT ``linspace(1.0, 0.001, num_steps)`` sigma + schedule (NOT flow-shift) + CFG-zero optimized-scale guidance — both live in ``LongCatDenoiseLoop``. + +``stamp_wan21_checkpoints`` applies (the diffusers transformer/vae/text_encoder/tokenizer subfolder layout is +identical to Wan). i2v/VC (first-frame latent replacement + per-frame timestep masking + num_cond_latents), +the refine stages (t_thresh crop, which is the latent-prep branch that SKIPS the init-noise re-scale), KV-cache, +and BSA are deferred (spec blocker #8) — the registered preset is base T2V. +""" +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 NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.longcat.loop import CFGZeroPolicy, LongCatDenoiseLoop +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# LongCat's default negative prompt (the inference yaml / get_sampling_param default — identical to the Wan +# English negative). Kept as a LOCAL module constant so the recipe touches no shared prompts file. +LONGCAT_NEG = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, " + "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, " + "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, " + "misshapen limbs, fused fingers, still picture, messy background, three legs, many people in " + "the background, walking backwards") + +_LONGCAT_DIT = "v2.platform.backends.torch_longcat:LongCatDiT" + + +def build_longcat_card(model_id: str = "longcat-video-t2v", + *, + 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 = CFGZeroPolicy() + # Keep the numpy loop math in fp32 (lowest-risk bring-up). The torch DiT runs at its native bf16 dtype + # and casts its output to fp32 (spec blocker #6); the loop never forces fp32 on the module. + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return LongCatDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, precision=precision, expert=expert, cost=cost) + + components = { + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + 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.longcat:LongCatTransformer3DModel", + adapter=_LONGCAT_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="longcat", + 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()), + # Spec defaults: 50 steps, guidance 1.0, 480x848, 121 frames @ 24fps. flow_shift=None (the loop builds + # the linspace schedule); the negative prompt is the LongCat default above. + sampling_defaults=sampling_defaults or SamplingDefaults(num_steps=50, + guidance_scale=1.0, + height=480, + width=848, + num_frames=121, + fps=24, + negative_prompt=LONGCAT_NEG), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card diff --git a/v2/recipes/longcat/loop.py b/v2/recipes/longcat/loop.py new file mode 100644 index 000000000..9e847cd5e --- /dev/null +++ b/v2/recipes/longcat/loop.py @@ -0,0 +1,209 @@ +"""LongCatDenoiseLoop — a faithful fork of ``WanDenoiseLoop`` for LongCat-Video T2V. + +LongCat is flow-match (deterministic Euler over a ``FlowMatchEulerDiscreteScheduler``), but it differs from +Wan in three load-bearing ways (faithful to ``fastvideo/pipelines/stages/longcat_denoising.py`` + +``longcat_refine_timestep.py``): + + 1. SIGMA SCHEDULE (spec blocker #2): NOT ``build_flow_sigmas`` / flow-shift (``flow_shift=None``). LongCat's + ``get_timesteps_sigmas`` (non-distill) uses ``torch.linspace(1.0, 0.001, num_inference_steps)`` — exactly + ``num_steps`` points terminating at 0.001 (NOT 0) — and the scheduler appends a trailing ``0.0`` for the + final boundary. ``FlowShiftPolicy.build_schedule`` would give ``num_steps+1`` points terminating at 0.0 + (wrong count and terminal). We build it explicitly in ``init`` via ``longcat_linspace_sigmas``. + 2. CFG-ZERO (spec blocker #3): guidance is the CFG-zero optimized-scale formula (``st_star`` projection of + cond onto uncond), NOT ClassicCFG. Implemented as ``CFGZeroPolicy`` (a ``CFGPolicy`` subclass local to + this recipe). Default ``guidance_scale=1.0`` degenerates the scale term but fastvideo still runs BOTH + branches whenever ``do_classifier_free_guidance`` is set (independent of the scale), so we keep both. + 3. SIGN CONVENTION (spec blocker #1): handled in the ``LongCatDiT`` adapter (returns ``-velocity``), so the + loop's ``FLOW_MATCH_STEP`` (``x + (sigma_next-sigma)*v``) reproduces the stage's negate-then-step. The + CFG-zero combine here operates on the (already-negated) velocities; the combine is sign-equivariant + (``st_star`` is a ratio of dot-products, and the affine combine is linear), so combining the negated + velocities equals negating the combine of the raw ones — i.e. ``-combine(raw) == combine(-raw)`` — which + is exactly fastvideo's order (combine raw, then negate). The toy backend has no negating adapter; that is + fine for the CPU control-flow check (we still exercise the schedule + CFG-zero + Euler step end-to-end). + +Geometry is Wan2.1's (``AutoencoderKLWan``: 16 channels, 4x temporal, 8x spatial). All per-request state lives +in ``LoopState`` (interleave-safe), matching the canonical ``WanDenoiseLoop`` contract. +""" +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.policies import CFGPolicy +from v2.platform import FLOW_MATCH_STEP +from v2.recipes.longcat.sampler import longcat_linspace_sigmas +from v2.recipes.wan21.loop import latent_shape + +# LongCat reuses Wan2.1 VAE geometry (AutoencoderKLWan: z=16, 4x temporal, 8x spatial). +LONGCAT_LATENT_CHANNELS = 16 +LONGCAT_TEMPORAL_RATIO = 4 +LONGCAT_SPATIAL_RATIO = 8 + + +class CFGZeroPolicy(CFGPolicy): + """CFG-zero optimized-scale guidance (LongCat; from the CFG-zero* paper) over the standard + cond/uncond two-forward decomposition. + + Per sample, project cond onto uncond to get the optimized scale + ``st_star = / (||v_uncond||^2 + 1e-8)``, then combine + ``out = v_uncond*st_star + scale*(v_cond - v_uncond*st_star)``. This is the exact algebra of + ``LongCatDenoisingStage.optimized_scale`` + the CFG application (longcat_denoising.py lines 28-149), with + the dot-products taken per leading "batch" group. In the v2 loop a step works on ONE unbatched sample + ([C,T,H,W]), so ``st_star`` is a single scalar over the whole tensor (B=1 in fastvideo's reshape(B,-1)).""" + + def combine(self, preds, guidance_scale, ctx, state): + cond = np.asarray(preds["cond"], dtype=np.float64) + uncond = np.asarray(preds["uncond"], dtype=np.float64) + # st_star = / (||uncond||^2 + 1e-8) over the full (B=1) tensor — fastvideo reshapes to + # [B,-1] and reduces along dim=1; here the single sample is the whole flattened tensor. + dot = float(np.sum(cond * uncond)) + sq = float(np.sum(uncond * uncond)) + 1e-8 + st_star = dot / sq + out = uncond * st_star + float(guidance_scale) * (cond - uncond * st_star) + return out.astype("float32") + + +class LongCatDenoiseLoop: + """Driven flow-match denoise loop for LongCat-Video T2V. Same four-method contract as + ``WanDenoiseLoop``; the deltas are the explicit linspace sigma schedule and the CFG-zero combine.""" + + def __init__(self, + *, + loop_id, + cfg, + precision, + expert, + cost, + sigma_min: float = 0.001, + latent_channels: int = LONGCAT_LATENT_CHANNELS, + spatial_ratio: int = LONGCAT_SPATIAL_RATIO, + temporal_ratio: int = LONGCAT_TEMPORAL_RATIO): + self.loop_id = loop_id + self.cfg = cfg # CFGZeroPolicy + self.precision = precision + self.expert = expert + self.cost = cost + self.sigma_min = sigma_min # the linspace terminal (LongCat: 0.001, NOT 0.0) + 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) + # LongCat schedule (spec blocker #2): linspace(1.0, 0.001, num_steps) + a trailing 0.0 sentinel so + # next() always has a sigma_next for the final boundary. NOT flow-shift, NOT terminal-0 linspace. + sig = longcat_linspace_sigmas(req.diffusion.num_steps, sigma_min=self.sigma_min) + shape = latent_shape(req, + model, + channels=self.latent_channels, + spatial_ratio=self.spatial_ratio, + temporal_ratio=self.temporal_ratio) + # Initial noise x = randn * sigma[0] (=1.0). LongCat's LatentPreparationStage multiplies by + # scheduler.init_noise_sigma (=1.0 for FlowMatchEulerDiscreteScheduler), so this matches; the refine + # branch (refine_from/stage1_video) is the path that SKIPS this re-scale, but base T2V always scales. + x = (rng.standard_normal(shape) * float(sig[0])).astype("float32") + 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 = sigma*1000 (Wan convention) + 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")) + 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] + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + x = st.latents["video"] + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + scale = st.scratch["guidance_scale"] + cfg, precision = self.cfg, self.precision + # LongCat runs BOTH branches whenever a negative prompt is present (do_classifier_free_guidance is + # independent of the scale value); CFG-zero degenerates only its scale term, not the branch set. + do_cfg = ne is not None + branches = ["cond", "uncond"] if do_cfg else ["cond"] + + def _velocity(model: Any, x_: Any, sigma_t_: float, pe_: Any, ne_: Any, scale_: float) -> np.ndarray: + dit = model.component(expert_id) + preds = {b: dit(x_, pe_ if b == "cond" else ne_, sigma_t_) for b in branches} + if not do_cfg: + return precision.cast(np.asarray(preds["cond"], dtype="float32")) + return precision.cast(cfg.combine(preds, scale_, sctx, cfg_state)) + + def run(model: Any, override: Any = None) -> StepResult: + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x, sigma_t, pe, ne, scale) + x_next = 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)), float(len(branches))), + 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(cfg).__name__), )), + resources=res, + payload={ + "branch": "combined", + "step": i + }, + run=run, + label=f"longcat.denoise.{i}", + # CFG-zero combine is host-side (np.sum / per-sample st_star) -> eager-break under capture, like + # the Cosmos x0-space CFG. The deterministic Euler step itself is capturable, but the combine is + # data-dependent on the two branch outputs, so keep the whole step eager for correctness. + capturable=False) + + 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/longcat/program.py b/v2/recipes/longcat/program.py new file mode 100644 index 000000000..a5215737d --- /dev/null +++ b/v2/recipes/longcat/program.py @@ -0,0 +1,37 @@ +"""LongCat-Video T2V program: text_encode -> diffusion_denoise (flow-match + CFG-zero) -> vae_decode. + +Same inline shape as the Wan/Cosmos t2v programs — the LongCat specifics (the linspace sigma schedule, the +CFG-zero combine, and the negated-velocity sign convention) live entirely in ``LongCatDenoiseLoop`` and the +``LongCatDiT`` adapter, so the node graph is unchanged. + +(i2v / video-continuation would add a VAE-encode node writing the first-frame conditioning latent into slots +and a loop that threads ``num_cond_latents`` + per-frame timestep masking — deferred; spec blocker #4/#5.) +""" +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_longcat_program() -> Program: + return Program( + program_id="longcat.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/recipes/longcat/sampler.py b/v2/recipes/longcat/sampler.py new file mode 100644 index 000000000..a790cd3c6 --- /dev/null +++ b/v2/recipes/longcat/sampler.py @@ -0,0 +1,22 @@ +"""LongCat sigma schedule — the recipe-local sampler helper (kept HERE, not in the shared +``v2/loop/sampler.py``, so the port stays a self-contained package). + +LongCat's ``get_timesteps_sigmas`` (non-distill, the base T2V path) builds the inference sigmas with +``torch.linspace(1.0, 0.001, num_inference_steps)`` and hands them to ``FlowMatchEulerDiscreteScheduler``; +the scheduler appends a trailing ``0.0`` (its final-step boundary). This is NOT the flow-shift linspace Wan +uses (``build_flow_sigmas`` gives ``num_steps+1`` points terminating at 0.0). Faithful to +``fastvideo/pipelines/stages/longcat_refine_timestep.py`` lines 63-69 (which documents the base schedule) and +the scheduler's terminal-sigma append. +""" +from __future__ import annotations + +import numpy as np + + +def longcat_linspace_sigmas(num_steps: int, sigma_min: float = 0.001) -> np.ndarray: + """The LongCat inference sigma schedule: ``linspace(1.0, sigma_min, num_steps)`` (NOT terminal 0) with a + trailing ``0.0`` appended for the final Euler boundary — giving ``num_steps+1`` entries the loop integrates + pairwise (``num_steps`` actual denoise steps). ``sigma_min`` is LongCat's 0.001 terminal.""" + n = max(int(num_steps), 1) + ramp = np.linspace(1.0, float(sigma_min), n, dtype=np.float64) # num_steps points, 1.0 -> 0.001 + return np.concatenate([ramp, np.zeros(1, dtype=np.float64)]) # + the scheduler's trailing-0 boundary diff --git a/v2/recipes/matrixgame2/__init__.py b/v2/recipes/matrixgame2/__init__.py new file mode 100644 index 000000000..94bd5a480 --- /dev/null +++ b/v2/recipes/matrixgame2/__init__.py @@ -0,0 +1,19 @@ +"""Matrix-Game 2.0 (Base, distilled) — the INTERACTIVE mouse/keyboard world model ported into the v2 +substrate as a self-contained recipe package (bucket-C pattern). + +The card declares its torch adapters via ``ComponentSpec.adapter`` (``MatrixGame2CausalDiT`` / +``MatrixGame2CLIPImageEncoder`` in ``v2/platform/backends/torch_matrixgame2.py``) and a new +``MatrixGame2CausalDMDLoop`` (few-step DMD: epsilon->x0 via the FlowUniPC sigma table + re-add_noise, causal +block-autoregressive with a sliding-window KV cache, action-conditioned). Reuses the Wan VAE adapter + +``stamp_wan21_checkpoints``. Registered in ``v2/registry.py`` (the orchestrator adds the entry). + +The registered path is the i2v world-rollout (first-frame + CLIP context); live mouse/keyboard action +routing is BRINGUP (needs a request-API extension — the loop + program already thread the action slots). +""" +from __future__ import annotations + +from v2.recipes.matrixgame2.card import build_matrixgame2_card +from v2.recipes.matrixgame2.loop import MatrixGame2CausalDMDLoop +from v2.recipes.matrixgame2.program import build_matrixgame2_program + +__all__ = ["build_matrixgame2_card", "build_matrixgame2_program", "MatrixGame2CausalDMDLoop"] diff --git a/v2/recipes/matrixgame2/card.py b/v2/recipes/matrixgame2/card.py new file mode 100644 index 000000000..5d2b8df92 --- /dev/null +++ b/v2/recipes/matrixgame2/card.py @@ -0,0 +1,148 @@ +"""Matrix-Game 2.0 (Base, distilled) ModelCard — the INTERACTIVE mouse/keyboard world model ported as a +self-contained recipe (bucket-C). Registered preset is the i2v world-rollout entry; live action routing is +BRINGUP (needs a request-API extension — the loop already threads the action slots). + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.matrixgame2.causal_model:CausalMatrixGame2WanModel`` — a DISTILLED, CAUSAL, + ACTION-CONDITIONED Wan-family transformer whose output is EPSILON (not velocity). The adapter + (``ComponentSpec.adapter`` -> ``MatrixGame2CausalDiT``) returns the raw epsilon; ``MatrixGame2CausalDMDLoop`` + does the few-step DMD (eps->x0 via the scheduler sigma table + re-add_noise), the causal block-rollout, + and the sliding-window KV-cache threading. + * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — Wan 2.1 (z=16, 8x/4x); reuses the v2 ``WanVAE`` + adapter unchanged (mean/std normalization). The i2v cond_concat latent uses the SAME normalization. + * Image CLIP vision encoder (``fastvideo.models.encoders.clip:CLIPVisionModel``, 257x1280 tokens) is the + SOLE cross-attention context (``MatrixGame2CLIPImageEncoder`` adapter + sibling ``image_processor``). + * NO text encoder — Matrix-Game 2.0 ignores text; the condition_embedder returns no text tokens. +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/image_encoder subfolder layout). + +Distilled: the 3-step DMD schedule ``dmd_denoising_steps=[1000,666,333]`` is intrinsic to the weights; the +preset surfaces ``num_inference_steps`` but the model is distilled for exactly these timesteps. CFG is OFF +(``guidance_scale=1.0``). The Base checkpoint is single-expert (``boundary_ratio=None``); GTA/TempleRun +distilled variants may carry a high/low-noise MoE boundary (the loop has the gated ``pred_noise_to_x_bound`` +path for that). +""" +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyImageEncoder, ToyVAE, _seed_from +from v2.recipes.matrixgame2.loop import MatrixGame2CausalDMDLoop +from v2.recipes.matrixgame2.sampler import ( + MATRIXGAME2_CONTEXT_NOISE, + MATRIXGAME2_DMD_STEPS, + MATRIXGAME2_FLOW_SHIFT, + MATRIXGAME2_NUM_FRAMES_PER_BLOCK, +) +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# Matrix-Game 2.0 ignores text -> no negative prompt (kept as a LOCAL constant, not shared _prompts.py). +MATRIXGAME2_NEG = "" + +_MG2_DIT = "v2.platform.backends.torch_matrixgame2:MatrixGame2CausalDiT" +_MG2_CLIP = "v2.platform.backends.torch_matrixgame2:MatrixGame2CLIPImageEncoder" + + +def build_matrixgame2_card(model_id: str = "matrix-game-2.0-base-distilled", + *, + flow_shift: float = MATRIXGAME2_FLOW_SHIFT, + boundary_ratio: float | None = None, + 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() # carried for the op-structure key; guidance_scale=1.0 makes it a no-op (CFG OFF) + flow = FlowShiftPolicy(shift=flow_shift) # seeds the FlowUniPC sigma table the eps->x0 conversion needs + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") # single-expert Base; the loop gates the MoE bound path on boundary_ratio + + def loop_factory(): + return MatrixGame2CausalDMDLoop(loop_id="causal_dmd_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost, + dmd_steps=MATRIXGAME2_DMD_STEPS, + context_noise=MATRIXGAME2_CONTEXT_NOISE, + num_frames_per_block=MATRIXGAME2_NUM_FRAMES_PER_BLOCK, + boundary_ratio=boundary_ratio) + + components = { + "image_encoder": + ComponentSpec(component_id="image_encoder", + kind="image_encoder", + load_id="fastvideo.models.encoders.clip:CLIPVisionModel", + adapter=_MG2_CLIP, + factory=lambda inst: ToyImageEncoder(), + required_for={"i2v"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan", + factory=lambda inst: ToyVAE(), + required_for={"i2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.matrixgame2.causal_model:CausalMatrixGame2WanModel", + adapter=_MG2_DIT, + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["causal_dmd_denoise"], + required_for={"i2v"}), + } + loops = { + "causal_dmd_denoise": + LoopSpec( + loop_id="causal_dmd_denoise", + kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, + step_cost_model=cost, + shared_weight_components=["transformer"], + cache_policy=["feature"], + graph_capture="eager", # causal KV-cache mutation + host-RNG re-noise -> not capturable + loop_factory=loop_factory), + } + card = ModelCard( + model_id=model_id, + family="matrixgame", + components=components, + loops=loops, + # World model: i2v (first-frame + action -> continuous frame extension) + VAE decode. + capabilities=CapabilityMatrix.of(Capability.IMAGE_TO_VIDEO, Capability.VAE_DECODE), + recipe=RecipeSpec(method="base", + assumes_loop="causal_dmd_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=3, + guidance_scale=1.0, + height=352, + width=640, + num_frames=57, + fps=25, + negative_prompt=MATRIXGAME2_NEG), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card diff --git a/v2/recipes/matrixgame2/loop.py b/v2/recipes/matrixgame2/loop.py new file mode 100644 index 000000000..568b6049e --- /dev/null +++ b/v2/recipes/matrixgame2/loop.py @@ -0,0 +1,316 @@ +"""MatrixGame2CausalDMDLoop — the few-step, causal/autoregressive, action-conditioned DMD denoise loop for +Matrix-Game 2.0 (the INTERACTIVE world model). Cannot reuse ``WanDenoiseLoop``/``CosmosDenoiseLoop``: + + (a) Few-step DMD, NOT a flow-match velocity ODE. The model output is EPSILON; each DMD step converts + ``eps -> x0`` via ``pred_video = x - sigma_t*eps`` (the scheduler sigma table), then RE-ADDS noise to + the next DMD timestep. Fixed schedule ``dmd_denoising_steps=[1000,666,333]`` (3 steps), warped through + the FlowUniPC timestep grid (``warp_denoising_step``). CFG is OFF (guidance_scale=1.0). + (b) Causal block-autoregressive: the video latent is split into blocks of ``num_frames_per_block`` (=3); + each block is fully DMD-denoised, then a context pass at ``timestep=context_noise(=0)`` writes its + CLEAN K/V into a sliding-window KV cache before the next block. Per-frame timestep ``[B, num_frames]``. + (c) Action-conditioned: per-frame mouse[F,2] / keyboard[F,4] feed ActionModule blocks with their own KV + caches. The DEGENERATE (no-action) path is what CPU-verifies here; live action routing is BRINGUP + (needs a request-API extension — the loop exposes the ``mouse_cond``/``keyboard_cond`` slots and a + ``streaming_reset``/``streaming_step`` surface for interactive play). + +Faithful port of ``fastvideo/pipelines/stages/matrixgame2_denoising.py:MatrixGame2CausalDenoisingStage`` +(``forward`` / ``_process_single_block`` / ``_update_context_cache``). The DiT call is dispatched two ways +so the SAME loop runs on the CPU toy AND the GPU adapter: + * GPU: ``model.component('transformer')`` is a ``MatrixGame2CausalDiT`` exposing ``.call(...)`` with the + causal kv_cache / i2v / action plumbing built INTERNALLY (the loop just passes numpy + the cache lists). + * CPU toy: ``ToyDiT`` has only ``dit(latent, text_embed, sigma)``; the loop's ``_eps`` falls back to that + (no kv_cache — the toy is a degenerate stand-in exercising the DMD + block + epsilon->x0 control flow + with real numbers). The KV caches are still allocated + threaded so the structure is identical. + +BRINGUP: the GPU kv_cache / crossattn_cache / action-cache shapes (``local_attn_size`` window, mouse-cache +batch dim ``B*frame_seq_len``) come from the loaded arch_config and are allocated by the GPU adapter's +companion init; on CPU we allocate lightweight placeholder dict lists so the threading is faithful but the +toy ignores them. The MoE high/low-noise boundary path (``pred_noise_to_x_bound`` + ``add_noise_high``) +exists for GTA/TempleRun variants (``boundary_ratio`` set); the Base distilled checkpoint is single-expert. +""" +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.recipes.matrixgame2.sampler import ( + MATRIXGAME2_CONTEXT_NOISE, + MATRIXGAME2_DMD_STEPS, + MATRIXGAME2_NUM_FRAMES_PER_BLOCK, + add_noise, + build_flow_unipc_table, + pred_noise_to_pred_video, + pred_noise_to_x_bound, + warp_dmd_timesteps, +) +from v2.recipes.wan21.loop import latent_shape + +# Real Wan2.1 VAE (AutoencoderKLWan): z=16, 4x temporal, 8x spatial. The DiT patch_embedding is 32-in +# (16 noise + 16 i2v cond_concat) but the LOOP's latent is the 16ch noise; the adapter does the concat. +MATRIXGAME2_LATENT_CHANNELS = 16 +MATRIXGAME2_TEMPORAL_RATIO = 4 +MATRIXGAME2_SPATIAL_RATIO = 8 + + +class MatrixGame2CausalDMDLoop: + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + dmd_steps: tuple[int, ...] = MATRIXGAME2_DMD_STEPS, + context_noise: int = MATRIXGAME2_CONTEXT_NOISE, + num_frames_per_block: int = MATRIXGAME2_NUM_FRAMES_PER_BLOCK, + boundary_ratio: float | None = None, + latent_channels: int = MATRIXGAME2_LATENT_CHANNELS, + spatial_ratio: int = MATRIXGAME2_SPATIAL_RATIO, + temporal_ratio: int = MATRIXGAME2_TEMPORAL_RATIO): + self.loop_id = loop_id + self.cfg = cfg # carried for the WorkPlan op-structure key; CFG is OFF (scale=1.0) for the distilled model + self.flow_shift = flow_shift + self.precision = precision + self.expert = expert + self.cost = cost + self.dmd_steps = tuple(dmd_steps) + self.context_noise = int(context_noise) + self.num_frames_per_block = int(num_frames_per_block) + self.boundary_ratio = boundary_ratio + self.latent_channels = latent_channels + self.spatial_ratio = spatial_ratio + self.temporal_ratio = temporal_ratio + + # --------------------------------------------------------------------- # + # init # + # --------------------------------------------------------------------- # + 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) + shift = float(self.flow_shift.shift_for(req.diffusion.height, req.diffusion.width)) + sigmas, sched_timesteps = build_flow_unipc_table(shift=shift) + # Warp the nominal DMD steps through the FlowUniPC timestep grid (warp_denoising_step=True). + warped = warp_dmd_timesteps(self.dmd_steps, scheduler_timesteps=sched_timesteps, shift=shift) + shape = latent_shape(req, + model, + channels=self.latent_channels, + spatial_ratio=self.spatial_ratio, + temporal_ratio=self.temporal_ratio) + # DMD starts from pure noise at the first (largest) warped timestep (sigma ~ 1). + x = rng.standard_normal(shape).astype("float32") + 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 sigmas], + timesteps=[float(t) for t in warped]) + st.scratch["sched_sigmas"] = sigmas + st.scratch["sched_timesteps"] = sched_timesteps + st.scratch["dmd_timesteps"] = [float(t) for t in warped] + # i2v + interactive conditioning (program writes these into slots; degenerate path -> None). + st.scratch["image_embeds"] = ctx.slots.get("i2v_img_embeds") # CLIP 257x1280, sole cross-attn ctx + st.scratch["image_latent"] = ctx.slots.get("i2v_cond") # raw cond_latent for channel-concat + st.scratch["mouse_cond"] = ctx.slots.get("mouse_cond") # BRINGUP: action surface (None here) + st.scratch["keyboard_cond"] = ctx.slots.get("keyboard_cond") + st.scratch["stream_video"] = bool(req.outputs.stream.get("video")) + # Causal block schedule over the latent-frame axis (axis 1 of [C, T, h, w]). + nf = int(shape[1]) + nfpb = max(1, self.num_frames_per_block) + # If the latent frame count isn't a multiple, the last block is shorter (the GPU stage asserts + # divisibility; the CPU toy stays robust so the loop still verifies on any tiny toy shape). + blocks: list[tuple[int, int]] = [] + start = 0 + while start < nf: + cur = min(nfpb, nf - start) + blocks.append((start, cur)) + start += cur + st.scratch["blocks"] = blocks + st.scratch["block_idx"] = 0 + st.scratch["dmd_idx"] = 0 + st.scratch["phase"] = "denoise" # "denoise" (3 DMD steps/block) then "context" (1 clean-K/V pass) + # Sliding-window KV caches (one dict per transformer block). On CPU these are inert placeholders the + # toy ignores; on GPU the MatrixGame2CausalDiT adapter mutates them in place across blocks/steps. + st.scratch["kv_cache"] = [] + st.scratch["crossattn_cache"] = [] + st.scratch["kv_cache_mouse"] = [] + st.scratch["kv_cache_keyboard"] = [] + st.plugin_state["cfg"] = {} + return st + + # --------------------------------------------------------------------- # + # eps dispatch (GPU causal adapter vs CPU toy) # + # --------------------------------------------------------------------- # + def _eps(self, dit: Any, latent: np.ndarray, st: LoopState, start_frame: int, num_frames: int, timestep: float, + frame_seq_len: int) -> np.ndarray: + """One forward returning the RAW epsilon prediction. Per-frame timestep [num_frames] (long-ish); + the GPU adapter keeps it 2-D. ``hasattr(dit, 'call')`` selects the causal adapter; else the toy.""" + per_frame_t = np.full((num_frames, ), float(timestep), dtype="float32") + if hasattr(dit, "call"): # MatrixGame2CausalDiT (GPU): full causal/i2v/action plumbing + return np.asarray(dit.call( + latent, + st.scratch.get("image_embeds"), + per_frame_t, + kv_cache=st.scratch["kv_cache"] or None, + crossattn_cache=st.scratch["crossattn_cache"] or None, + kv_cache_mouse=st.scratch["kv_cache_mouse"] or None, + kv_cache_keyboard=st.scratch["kv_cache_keyboard"] or None, + current_start=start_frame * frame_seq_len, + start_frame=start_frame, + num_frame_per_block=num_frames, + mouse_cond=st.scratch.get("mouse_cond"), + keyboard_cond=st.scratch.get("keyboard_cond"), + image_latent=st.scratch.get("image_latent"), + ), + dtype="float32") + # CPU toy degenerate path: dit(latent, text_embed, sigma). Feed the sigma at this timestep so the + # toy velocity is a deterministic function of the DMD step (exercises the eps->x0 control flow). + sigma = self._sigma_at(st, timestep) + ctx_embed = st.scratch.get("image_embeds") + return np.asarray(dit(latent, ctx_embed, sigma), dtype="float32") + + @staticmethod + def _sigma_at(st: LoopState, timestep: float) -> float: + sigmas, tsteps = st.scratch["sched_sigmas"], st.scratch["sched_timesteps"] + idx = int(np.argmin(np.abs(np.asarray(tsteps, dtype=np.float64) - float(timestep)))) + return float(sigmas[idx]) + + def _frame_seq_len(self, x: np.ndarray) -> int: + # frame_seq_len = (h*w) / (patch_h*patch_w). Wan patch is [1,2,2] -> /4. The toy shape is tiny; clamp. + h, w = int(x.shape[2]), int(x.shape[3]) + return max(1, (h * w) // 4) + + # --------------------------------------------------------------------- # + # next # + # --------------------------------------------------------------------- # + def next(self, st: LoopState): + blocks = st.scratch["blocks"] + b_idx = st.scratch["block_idx"] + if b_idx >= len(blocks): + return Done() + start_frame, num_frames = blocks[b_idx] + phase = st.scratch["phase"] + dmd_idx = st.scratch["dmd_idx"] + dmd_timesteps = st.scratch["dmd_timesteps"] + x_full = st.latents["video"] + block = x_full[:, start_frame:start_frame + num_frames] # [C, nf, h, w] + sigmas, tsteps = st.scratch["sched_sigmas"], st.scratch["sched_timesteps"] + precision, rng = self.precision, st.rng + frame_seq_len = self._frame_seq_len(x_full) + boundary_ts = (self.boundary_ratio * 1000.0) if self.boundary_ratio is not None else None + _t = dmd_timesteps[dmd_idx] if phase == "denoise" else float(self.context_noise) + _s = self._sigma_at(st, dmd_timesteps[dmd_idx]) if phase == "denoise" else 0.0 + expert_id = self.expert.expert_for(StepContext(dmd_idx, _t, _s)) + + is_context = phase == "context" + timestep = float(self.context_noise) if is_context else float(dmd_timesteps[dmd_idx]) + is_last_dmd = dmd_idx >= len(dmd_timesteps) - 1 + + def run(model, override=None): + dit = model.component(expert_id) + if is_context: + # Clean-context pass: one forward at timestep=context_noise to write clean K/V into the + # sliding-window cache (return value discarded). Faithful to _update_context_cache. + if hasattr(dit, "call"): + self._eps(dit, block, st, start_frame, num_frames, float(self.context_noise), frame_seq_len) + return StepResult(output={ + "block": np.asarray(block, dtype="float32"), + "start": start_frame, + "num": num_frames, + "context": True + }) + eps = self._eps(dit, block, st, start_frame, num_frames, timestep, frame_seq_len) + if boundary_ts is not None and timestep >= boundary_ts: # MoE high-noise expert (variants only) + x0 = pred_noise_to_x_bound(eps, block, timestep, boundary_ts, sigmas, tsteps) + else: + x0 = pred_noise_to_pred_video(eps, block, timestep, sigmas, tsteps) + if is_last_dmd: + new_block = x0 # last DMD step: the clean x0 IS the block result + else: + next_t = float(dmd_timesteps[dmd_idx + 1]) + noise = rng.standard_normal(np.asarray(x0).shape) + new_block = add_noise(x0, noise, next_t, sigmas, tsteps) # re-noise to the next DMD timestep + new_block = precision.cast(new_block) + return StepResult( + output={ + "block": np.asarray(new_block, dtype="float32"), + "eps": np.asarray(eps, dtype="float32"), + "start": start_frame, + "num": num_frames, + "context": False + }) + + res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(block.shape)), 1.0), + resident_bytes=int(x_full.nbytes), + peak_activation_bytes=int(block.nbytes)) + label = (f"matrixgame2.context.{b_idx}" if is_context else f"matrixgame2.denoise.{b_idx}.{dmd_idx}") + return WorkPlan(loop_id=self.loop_id, + instance_id=st.instance_id, + kind=WorkUnitKind.DIFFUSION_STEP, + shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, + dims=tuple(block.shape), + dtype=precision.compute_dtype, + extra=(("cfg", type(self.cfg).__name__), ("dmd", True), ("phase", + phase))), + resources=res, + payload={ + "branch": "dmd", + "block": b_idx, + "dmd": dmd_idx, + "phase": phase + }, + run=run, + label=label, + capturable=False) # causal KV-cache mutation + host-RNG re-noise -> eager path + + # --------------------------------------------------------------------- # + # advance # + # --------------------------------------------------------------------- # + def advance(self, st: LoopState, result: StepResult) -> LoopState: + out = result.output + start, num = out["start"], out["num"] + # Write the (denoised or unchanged) block back into the full latent. + x_full = np.array(st.latents["video"]) + x_full[:, start:start + num] = out["block"] + st.latents["video"] = x_full + + if out.get("context"): # finished the clean-context pass -> advance to the next block + if st.profile == ExecutionProfile.ROLLOUT: + st.trajectory.append({"block": st.scratch["block_idx"], "latents": np.asarray(x_full).copy()}) + st.scratch["block_idx"] += 1 + st.scratch["dmd_idx"] = 0 + st.scratch["phase"] = "denoise" + return st + + # Mid-block DMD step. + dmd_idx = st.scratch["dmd_idx"] + if dmd_idx >= len(st.scratch["dmd_timesteps"]) - 1: + # Block fully DMD-denoised -> run the clean-context cache-refresh pass next. + st.scratch["phase"] = "context" + else: + st.scratch["dmd_idx"] = dmd_idx + 1 + return st + + # --------------------------------------------------------------------- # + # finalize # + # --------------------------------------------------------------------- # + def finalize(self, st: LoopState) -> LoopResult: + return LoopResult(outputs={"latents": st.latents["video"]}, + metrics={"denoise_blocks": float(len(st.scratch["blocks"]))}, + behavior=st.trajectory or None) diff --git a/v2/recipes/matrixgame2/program.py b/v2/recipes/matrixgame2/program.py new file mode 100644 index 000000000..5a0613411 --- /dev/null +++ b/v2/recipes/matrixgame2/program.py @@ -0,0 +1,77 @@ +"""Matrix-Game 2.0 i2v world-rollout program: image_encode(CLIP) -> cond_encode(first-frame VAE) -> +action_prepare(BRINGUP hook) -> causal DMD denoise -> vae_decode. + +Deltas vs the Wan2.1 i2v program (faithful to ``matrixgame2_causal_dmd_pipeline.py`` + +``MatrixGame2ImageVAEEncodingStage``): + * NO text_encode node — Matrix-Game 2.0 ignores text; the CLIP image embeds are the sole cross-attn + context (written to ``i2v_img_embeds``). + * cond_encode writes the RAW first-frame VAE latent to ``i2v_cond`` with NO 4-channel mask (unlike Wan2.1 + i2v's ``[mask|cond]``); the DiT channel-concats this raw cond onto the noise latent internally (the + patch_embedding is the 32-in checkpoint). + * action_prepare is a HOOK that routes per-frame mouse/keyboard arrays into ``mouse_cond``/``keyboard_cond`` + when the request carries them. BRINGUP: the request API does not yet surface game-action arrays, so this + is a no-op on the registered path (the world model degrades to a first-frame-conditioned rollout). When + the request-API extension lands, this node slices the action arrays per block (vae_time_compression=4). +""" +from __future__ import annotations + +import numpy as np + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind + + +def _image_encode(instance, slots, request, ctx) -> None: + """CLIP-encode the first-frame conditioning image -> 257x1280 cross-attn context (sole conditioning).""" + img = request.image() + slots["i2v_img_embeds"] = (instance.component("image_encoder").encode_image(img.pixels) + if img is not None and getattr(img, "pixels", None) is not None else None) + + +def _cond_encode(instance, slots, request, ctx) -> None: + """First-frame VAE conditioning, mirroring ``MatrixGame2ImageVAEEncodingStage``: encode the conditioning + image as frame 0 (rest zeros) -> the RAW cond_latent (NO mask channel). The DiT channel-concats it onto + the noise latent (dim=1). Same VAE normalization as the noise latents (the WanVAE adapter handles it).""" + img = request.image() + if img is None or getattr(img, "pixels", None) is None: + slots["i2v_cond"] = None + return + px = np.asarray(img.pixels, dtype="float32") # [3, H, W] in [-1, 1] + nf = int(request.diffusion.num_frames) + cond_video = np.zeros((px.shape[0], nf) + px.shape[1:], dtype="float32") + cond_video[:, 0] = px # frame 0 = conditioning image, rest zeros + slots["i2v_cond"] = np.asarray(instance.component("vae").encode(cond_video), dtype="float32") # raw [C,T,h,w] + + +def _action_prepare(instance, slots, request, ctx) -> None: + """Route per-frame game actions (mouse[F,2] / keyboard[F,4]) into the loop's conditioning slots. + BRINGUP: the request API does not yet expose action arrays, so this is a no-op (the world model runs the + degenerate first-frame-conditioned rollout). When the action request-surface lands, read it here and the + loop/adapter feed the ActionModule KV caches. Slots stay ``None`` until then.""" + slots["mouse_cond"] = getattr(getattr(request, "diffusion", None), "mouse_cond", None) + slots["keyboard_cond"] = getattr(getattr(request, "diffusion", None), "keyboard_cond", None) + + +def _vae_decode(instance, slots, request, ctx) -> None: + slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_matrixgame2_program() -> Program: + return Program( + program_id="matrixgame2.i2v.causal_dmd", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("image_encode", fn=_image_encode, writes=("i2v_img_embeds", )), + ComponentNode("cond_encode", fn=_cond_encode, writes=("i2v_cond", )), + ComponentNode("action_prepare", fn=_action_prepare, writes=("mouse_cond", "keyboard_cond")), + ModelLoopNode("denoise", + loop_id="causal_dmd_denoise", + output_slot="denoise_out", + reads=("i2v_img_embeds", "i2v_cond", "mouse_cond", "keyboard_cond"), + 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/recipes/matrixgame2/sampler.py b/v2/recipes/matrixgame2/sampler.py new file mode 100644 index 000000000..e38c411eb --- /dev/null +++ b/v2/recipes/matrixgame2/sampler.py @@ -0,0 +1,91 @@ +"""Matrix-Game 2.0 DMD sampler helpers — the few-step distilled epsilon->x0 math, ported faithfully from +``fastvideo/models/utils.py`` (``pred_noise_to_pred_video`` / ``pred_noise_to_x_bound``) and the +``FlowUniPCMultistepScheduler`` sigma/timestep table it indexes. Kept HERE (a recipe-local sibling), NOT +in ``v2/loop/sampler.py``, because the DMD epsilon->x0 + re-add_noise schedule is specific to this arch. + +The shipped distilled checkpoints run a FIXED 3-step DMD schedule ``dmd_denoising_steps=[1000,666,333]``. +The scheduler (``FlowUniPCMultistepScheduler(shift=5.0)``) is used ONLY as a sigma/timestep lookup table + +``add_noise``, never as a UniPC multistep corrector. With ``warp_denoising_step=True`` the nominal integer +steps are remapped through the scheduler's timestep grid: +``timesteps = cat(scheduler.timesteps, [0])[1000 - dmd_steps]``. + +This numpy form is the CPU-testable, bit-reproducible reference. ``build_flow_unipc_table`` reproduces the +scheduler's ``sigmas``/``timesteps`` (flow-shift over a 1000-step alpha ramp; identical formula to +``FlowShiftPolicy``) so the loop can run on the CPU toy without instantiating the real scheduler. On a GPU +box the torch scheduler is built from the card's ``scheduler`` load_id and the SAME math applies. +""" +from __future__ import annotations + +import numpy as np + +# Matrix-Game 2.0 distilled hyper-params (fastvideo MatrixGame2I2V480PConfig + MatrixGame2WanVideoArchConfig). +MATRIXGAME2_FLOW_SHIFT = 5.0 +MATRIXGAME2_DMD_STEPS: tuple[int, ...] = (1000, 666, 333) +MATRIXGAME2_NUM_TRAIN_TIMESTEPS = 1000 +MATRIXGAME2_CONTEXT_NOISE = 0 +MATRIXGAME2_NUM_FRAMES_PER_BLOCK = 3 + + +def build_flow_unipc_table(shift: float = MATRIXGAME2_FLOW_SHIFT, + num_train_timesteps: int = MATRIXGAME2_NUM_TRAIN_TIMESTEPS) -> tuple[np.ndarray, np.ndarray]: + """Reproduce ``FlowUniPCMultistepScheduler.__init__`` (scheduling_flow_unipc_multistep.py): a descending + 1000-step alpha ramp ``alphas = linspace(1, 1/N, N)[::-1]``, ``sigmas = 1 - alphas``, flow-shifted + ``sigma' = shift*sigma / (1 + (shift-1)*sigma)``, then ``timesteps = sigma' * N``. Returns + ``(sigmas, timesteps)`` as float64, the table the epsilon->x0 conversion indexes by nearest timestep.""" + alphas = np.linspace(1.0, 1.0 / num_train_timesteps, num_train_timesteps, dtype=np.float64)[::-1].copy() + sigmas = 1.0 - alphas + if shift != 1.0: + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + timesteps = sigmas * num_train_timesteps + return sigmas, timesteps + + +def warp_dmd_timesteps(dmd_steps: tuple[int, ...] = MATRIXGAME2_DMD_STEPS, + *, + scheduler_timesteps: np.ndarray | None = None, + shift: float = MATRIXGAME2_FLOW_SHIFT, + num_train_timesteps: int = MATRIXGAME2_NUM_TRAIN_TIMESTEPS) -> np.ndarray: + """``warp_denoising_step``: remap the nominal DMD steps through the scheduler timestep grid, faithful to + ``MatrixGame2CausalDenoisingStage.forward``: ``timesteps = cat(scheduler.timesteps, [0])[1000 - steps]``. + The scheduler grid here is the FULL 1000-step table (set_timesteps with num_train_timesteps steps gives + the same grid the pipeline's LatentPreparationStage builds before warping).""" + if scheduler_timesteps is None: + _sig, scheduler_timesteps = build_flow_unipc_table(shift=shift, num_train_timesteps=num_train_timesteps) + grid = np.concatenate([np.asarray(scheduler_timesteps, dtype=np.float64), np.array([0.0])]) + idx = num_train_timesteps - np.asarray(dmd_steps, dtype=np.int64) + return grid[idx] + + +def _sigma_at(timestep: float, sigmas: np.ndarray, timesteps: np.ndarray) -> float: + """Nearest-neighbor sigma lookup: ``sigma_t = sigmas[argmin|timesteps - timestep|]`` — exactly the + ``torch.argmin((timesteps - timestep).abs())`` in ``pred_noise_to_pred_video``.""" + idx = int(np.argmin(np.abs(np.asarray(timesteps, dtype=np.float64) - float(timestep)))) + return float(sigmas[idx]) + + +def pred_noise_to_pred_video(pred_noise: np.ndarray, noise_input_latent: np.ndarray, timestep: float, + sigmas: np.ndarray, timesteps: np.ndarray) -> np.ndarray: + """epsilon -> x0: ``pred_video = noise_input_latent - sigma_t * pred_noise`` (single-expert path). + Faithful to ``fastvideo/models/utils.py:pred_noise_to_pred_video``.""" + sigma_t = _sigma_at(timestep, sigmas, timesteps) + return np.asarray(noise_input_latent, dtype=np.float64) - sigma_t * np.asarray(pred_noise, dtype=np.float64) + + +def pred_noise_to_x_bound(pred_noise: np.ndarray, noise_input_latent: np.ndarray, timestep: float, + boundary_timestep: float, sigmas: np.ndarray, timesteps: np.ndarray) -> np.ndarray: + """High-noise-expert epsilon -> x0: ``pred_video = x - (sigma_t - sigma_boundary) * pred_noise``. + Faithful to ``fastvideo/models/utils.py:pred_noise_to_x_bound`` (MoE distilled variants only; the Base + distilled checkpoint has ``boundary_ratio=None`` -> this path is unused).""" + sigma_t = _sigma_at(timestep, sigmas, timesteps) + sigma_b = _sigma_at(boundary_timestep, sigmas, timesteps) + return (np.asarray(noise_input_latent, dtype=np.float64) - + (sigma_t - sigma_b) * np.asarray(pred_noise, dtype=np.float64)) + + +def add_noise(x0: np.ndarray, noise: np.ndarray, timestep: float, sigmas: np.ndarray, + timesteps: np.ndarray) -> np.ndarray: + """Flow-match forward interpolation at the NEXT DMD timestep: ``x_t = (1 - sigma_t)*x0 + sigma_t*noise`` + — the scheduler's ``add_noise`` (re-noise between the 3 DMD steps). ``sigma_t`` is looked up at the next + timestep, matching ``scheduler.add_noise(pred_video, noise, next_timestep)``.""" + sigma_t = _sigma_at(timestep, sigmas, timesteps) + return (1.0 - sigma_t) * np.asarray(x0, dtype=np.float64) + sigma_t * np.asarray(noise, dtype=np.float64) diff --git a/v2/recipes/matrixgame3/__init__.py b/v2/recipes/matrixgame3/__init__.py new file mode 100644 index 000000000..68b1f14b8 --- /dev/null +++ b/v2/recipes/matrixgame3/__init__.py @@ -0,0 +1,15 @@ +"""Matrix-Game 3.0 (Base-Distilled) — autoregressive multi-clip world model ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapter via +``ComponentSpec.adapter`` (``MatrixGame3DiT`` in ``v2/platform/backends/torch_matrixgame3.py``) and a new +``MatrixGame3DenoiseLoop`` (the autoregressive multi-clip loop: per-token timesteps + first-frame pinning; +action/camera/KV-memory are BRINGUP), reusing the Wan VAE adapter + T5 + ``stamp_wan21_checkpoints``. The +orchestrator registers it in ``v2/registry.py`` (the ``MatrixGame3WanModel`` transformer fallback). +""" +from __future__ import annotations + +from v2.recipes.matrixgame3.card import build_matrixgame3_card +from v2.recipes.matrixgame3.loop import MatrixGame3DenoiseLoop +from v2.recipes.matrixgame3.program import build_matrixgame3_program + +__all__ = ["build_matrixgame3_card", "build_matrixgame3_program", "MatrixGame3DenoiseLoop"] diff --git a/v2/recipes/matrixgame3/card.py b/v2/recipes/matrixgame3/card.py new file mode 100644 index 000000000..b2edfa008 --- /dev/null +++ b/v2/recipes/matrixgame3/card.py @@ -0,0 +1,142 @@ +"""Matrix-Game 3.0 ModelCard — image+action+camera -> video autoregressive world model. + +Self-contained recipe package (the bucket-C pattern, mirroring cosmos2): the card declares its torch +adapter via ``ComponentSpec.adapter`` (``MatrixGame3DiT`` in ``v2/platform/backends/torch_matrixgame3.py``) +and a new ``MatrixGame3DenoiseLoop`` (the autoregressive multi-clip loop), reusing the Wan VAE adapter + +T5 + ``stamp_wan21_checkpoints``. + +Architecture deltas vs Wan (all GPU-path, declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.matrixgame3.model:MatrixGame3WanModel`` — a 5B autoregressive + world-model DiT (30 layers, dim 24*128=3072, ffn 14336, patch (1,2,2), in/out_channels=48, + ``use_memory=True``, ``sigma_theta=0.8``, action_config on the first 15 blocks). Predicts a velocity + (``flow_prediction``). The adapter (``ComponentSpec.adapter`` -> ``MatrixGame3DiT``) takes a PER-TOKEN + timestep (cond-frame rows zeroed) + the action/camera/KV-memory bundle — built INTERNALLY so the + loop's dit-call stays ``dit(latent, text_embed, sigma)``-shaped. + * VAE ``fastvideo.models.vaes.wanvae:AutoencoderKLWan`` — the ``light_vae`` variant (z_dim=48, 16x + spatial / 4x temporal, Wan2.2-TI2V geometry); reuses the v2 ``WanVAE`` adapter unchanged (its + ``_mean_invstd`` normalization path is the MG3 normalized-latent convention). + * Text ``fastvideo.models.encoders.t5:T5EncoderModel`` (T5, fp32, text_len 512) via the default + ``T5Encoder`` adapter (Wan zero-pad-to-512 convention). NO image/CLIP encoder — the conditioning + image goes only through the VAE (BRINGUP risk 6). + +Sampler: flow-match via ``FlowUniPCMultistepScheduler`` (``flow_prediction``, flow_shift=5.0); distilled +3-step student, guidance_scale=1.0 (CFG effectively off unless ``use_base_model=True``). +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder + tokenizer subfolders). +""" +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.matrixgame3.loop import ( + MG3_LATENT_CHANNELS, + MG3_SPATIAL_RATIO, + MG3_TEMPORAL_RATIO, + MatrixGame3DenoiseLoop, +) +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# Matrix-Game 3.0 uses an EMPTY negative prompt by default (kept as a LOCAL module constant per the +# self-contained-recipe rule — not added to the shared v2/recipes/_prompts.py). +MG3_NEG = "" + +_MG3_DIT = "v2.platform.backends.torch_matrixgame3:MatrixGame3DiT" + + +def build_matrixgame3_card(model_id: str = "matrixgame3-i2v", + *, + flow_shift: float = 5.0, + 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() + # MG3 720p uses flow_shift=5.0 (its 480P-derived pipeline config). The bucket lookup keeps the + # 720p/480p shift table consistent with the Wan flow-shift convention. + flow = FlowShiftPolicy(shift=flow_shift, bucket_lookup={480 * 832: 3.0, 720 * 1280: 5.0}) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return MatrixGame3DenoiseLoop(loop_id="mg3_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost, + latent_channels=MG3_LATENT_CHANNELS, + spatial_ratio=MG3_SPATIAL_RATIO, + temporal_ratio=MG3_TEMPORAL_RATIO) + + components = { + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + factory=lambda inst: ToyTextEncoder(), + required_for={"i2v"}), + "vae": + ComponentSpec( + component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan", # light_vae (z=48, 16x/4x) + factory=lambda inst: ToyVAE(), + required_for={"i2v"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.matrixgame3.model:MatrixGame3WanModel", + adapter=_MG3_DIT, + factory=lambda inst: ToyDiT(seed=seed), + resident_for=["mg3_denoise"], + required_for={"i2v"}), + } + loops = { + "mg3_denoise": + LoopSpec(loop_id="mg3_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="matrixgame", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.IMAGE_TO_VIDEO, Capability.VAE_DECODE), + recipe=RecipeSpec(method="base", + assumes_loop="mg3_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=3, guidance_scale=1.0, height=720, width=1280, num_frames=57, fps=25, negative_prompt=MG3_NEG), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card diff --git a/v2/recipes/matrixgame3/loop.py b/v2/recipes/matrixgame3/loop.py new file mode 100644 index 000000000..2d86034f7 --- /dev/null +++ b/v2/recipes/matrixgame3/loop.py @@ -0,0 +1,302 @@ +"""MatrixGame3DenoiseLoop — the autoregressive multi-clip world-model denoise loop. + +Matrix-Game 3.0 is NOT a single bidirectional flow-match pass (that is ``WanDenoiseLoop``). It is an +OUTER clip loop wrapping an INNER few-step flow-match denoise, with cross-clip frame/KV memory. Faithful +port of ``fastvideo/pipelines/stages/matrixgame3_denoising.py:MatrixGame3DenoisingStage.forward``: + + num_iterations = 1 (for num_frames <= 57; the registered preset is num_frames=57) + = 1 + ceil((num_frames-57)/40) (longer rollouts) + for clip_idx in range(num_iterations): + * draw fresh noise for this clip's latent frames; paste the first ``cond_frames`` (1 first clip, 4 + later) from the previous clip's last-4 latents (``img_cond``); + * (clip>0) build the per-clip action(mouse/keyboard) + camera(Plücker) + KV-memory bundle + (``x_memory``/``timestep_memory``/``memory_latent_idx``/``predict_latent_idx``) — host numpy work; + * INNER loop: for each of the ``num_steps`` (distilled=3) FlowUniPC timesteps, call the DiT with a + PER-TOKEN timestep (cond rows zeroed — built INSIDE the adapter), take a scheduler step, then + RE-PASTE ``img_cond`` onto the ``cond_frames`` (so the known frames never drift); + * after the clip: ``img_cond`` <- last-4 latents; append the denoised tail (full first clip, + last-10 thereafter) to history. + concat history on the time dim -> the final latent. + +THIS PORT — degenerate single-clip path is built so it CPU-verifies on the toy backend: + * The inner flow-match step uses the platform's numpy ``FLOW_MATCH_STEP`` kernel (the CPU-testable + stand-in). On a GPU box the *real* ``FlowUniPCMultistepScheduler`` must drive ``step()`` — it is a + stateful MULTISTEP solver and substituting plain Euler breaks the 3-step distilled trajectory + (BRINGUP risk 3). The ``self.use_unipc`` hook documents where the real scheduler plugs in. + * The dit-call is ``dit(latent, text_embed, sigma, cond=)`` — ToyDiT accepts and ignores + ``cond`` (image/action/camera conditioning is a GPU-path concern), so the toy denoises the noise + channels exactly like the cosmos2 video2world degradation. + * Action / camera / KV-memory (mouse_cond/keyboard_cond/plucker/x_memory) are BRINGUP: they need a + request-API extension to carry action streams + camera trajectories, plus the vendored host helpers + (``matrixgame3.utils`` + ``lingbotworld.cam_utils``). The single-clip no-action path is the t2v/i2v + degenerate the loop CPU-verifies; the multi-clip rollout hooks are present + documented. +""" +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.platform import FLOW_MATCH_STEP + +# Matrix-Game 3.0 light_vae (Wan2.2-TI2V geometry): z_dim=48, 16x spatial, 4x temporal compression. +MG3_LATENT_CHANNELS = 48 +MG3_SPATIAL_RATIO = 16 +MG3_TEMPORAL_RATIO = 4 +MG3_NUM_TRAIN_TIMESTEPS = 1000 # FlowUniPC timestep ~ sigma * num_train_timesteps + +# Clip geometry (hardcoded in the fastvideo denoising stage). first_clip = 57 frames -> 15 latent frames. +MG3_CLIP_FRAME = 56 +MG3_FIRST_CLIP_FRAME = MG3_CLIP_FRAME + 1 # 57 +MG3_PAST_FRAME = 16 + + +def _toy_latent_shape(req) -> tuple[int, int, int, int]: + """Tiny deterministic CPU-toy latent geometry (matches the WanDenoiseLoop toy stand-in), so the loop + control flow + per-token-timestep/cond-frame/re-paste logic exercise on numbers without a GPU.""" + d = req.diffusion + from v2.platform.backends.toy import LATENT_CHANNELS + t = max(2, d.num_frames // 40) + h = max(2, d.height // 120) + w = max(2, d.width // 120) + return (LATENT_CHANNELS, t, h, w) + + +def _gpu_latent_shape(req, channels: int, spatial_ratio: int, temporal_ratio: int) -> tuple[int, int, int, int]: + d = req.diffusion + t = (max(1, d.num_frames) - 1) // temporal_ratio + 1 + return (channels, max(1, t), max(1, d.height // spatial_ratio), max(1, d.width // spatial_ratio)) + + +def _cond_compatible(latent: np.ndarray, img_cond: np.ndarray) -> bool: + """The conditioning latent can be pasted onto the leading frames iff its channel + trailing spatial + dims match the noise latent (true on the GPU 48-channel light_vae; false for the toy VAE stand-in).""" + return (latent.shape[0] == img_cond.shape[0] and latent.shape[2:] == img_cond.shape[2:]) + + +class MatrixGame3DenoiseLoop: + """Autoregressive multi-clip world-model denoise. ``next``/``advance`` drive the INNER step; the OUTER + clip advance happens in ``advance`` when the inner schedule is exhausted (the loop carries the clip + index + history in ``LoopState.scratch``). On the registered single-clip preset there is exactly one + clip, so the loop reduces to an N-step flow-match denoise over one latent with first-frame pinning.""" + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels: int = MG3_LATENT_CHANNELS, + spatial_ratio: int = MG3_SPATIAL_RATIO, + temporal_ratio: int = MG3_TEMPORAL_RATIO, + use_unipc: bool = True): + self.loop_id = loop_id + self.cfg = cfg # carried for the WorkPlan op-structure key; CFG is off by default + self.flow_shift = flow_shift + self.precision = precision + self.expert = expert + self.cost = cost + self.latent_channels = latent_channels + self.spatial_ratio = spatial_ratio + self.temporal_ratio = temporal_ratio + # GPU path: drive the real FlowUniPCMultistepScheduler.step (stateful multistep) instead of the + # numpy Euler stand-in (BRINGUP risk 3). The CPU toy always uses FLOW_MATCH_STEP. + self.use_unipc = use_unipc + + # --- clip-count + per-clip geometry (faithful to MatrixGame3DenoisingStage) ------------------- # + @staticmethod + def _num_iterations(num_frames: int) -> int: + if isinstance(num_frames, int) and num_frames > MG3_FIRST_CLIP_FRAME: + return 1 + max(0, (num_frames - MG3_FIRST_CLIP_FRAME + 39) // 40) + return 1 + + 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) + is_gpu = getattr(getattr(model, "platform", None), "device", "cpu") == "cuda" + if is_gpu: + shape = _gpu_latent_shape(req, self.latent_channels, self.spatial_ratio, self.temporal_ratio) + else: + shape = _toy_latent_shape(req) + # Flow-match sigma schedule (1 -> 0) with flow_shift; the FlowUniPC scheduler uses the same shift. + # ``req.diffusion.sigmas`` (a distilled explicit schedule) overrides when provided. + sig = self.flow_shift.build_schedule(req.diffusion.num_steps, + req.diffusion.height, + req.diffusion.width, + sigmas=req.diffusion.sigmas or None) + x = (rng.standard_normal(shape) * float(sig[0])).astype("float32") + + # First-frame conditioning latent (image VAE-encoded by the program). On the degenerate path it is + # absent -> no pinning, pure flow-match. cond_frames=1 on the first clip (and is clamped to the + # available conditioning frames). On the GPU path the cond latent shares the noise latent's + # [C=48, T, h, w] geometry, so the leading frames are pasted in (and re-pasted every step). The + # CPU toy VAE is only a geometry-preserving stand-in for the channel dims, so ``_cond_compatible`` + # guards the pin: when the trailing dims don't match (toy), it stays the degenerate flow-match. + img_cond = ctx.slots.get("image_latent") + cond_frames = 0 + if img_cond is not None: + img_cond = np.asarray(img_cond, dtype="float32") + if _cond_compatible(x, img_cond): + cond_frames = min(1, img_cond.shape[1]) # first clip pins 1 frame + x[:, :cond_frames] = img_cond[:, :cond_frames] + else: + img_cond = None # toy: incompatible geometry -> no pinning + + 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], + # FlowUniPC model timestep = sigma * num_train_timesteps (the adapter fills the + # per-token tensor with this value; cond rows zeroed inside the adapter). + timesteps=[float(s) * MG3_NUM_TRAIN_TIMESTEPS for s in sig]) + 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) + # use_base_model=False + guidance_scale=1.0 (the distilled default) -> the CFG/uncond branch is + # skipped (BRINGUP risk 8). do_cfg only fires for the base (non-distilled) model. + st.scratch["use_base_model"] = bool(getattr(req.diffusion, "use_base_model", False)) + st.scratch["img_cond"] = img_cond + st.scratch["cond_frames"] = cond_frames + # Latent-frame span the model uses for rotary positions; single-clip span is the whole latent. + st.scratch["predict_latent_idx"] = (0, int(x.shape[1])) + # Autoregressive rollout (BRINGUP): num_iterations>1 + the action/camera/memory bundle. The + # single-clip preset (num_frames<=57) keeps these inert. + st.scratch["num_iterations"] = self._num_iterations(int(req.diffusion.num_frames)) + st.scratch["clip_idx"] = 0 + st.scratch["history"] = [] # appended denoised clip tails (rollout) + # Per-clip action/camera/memory conditioning — None on the degenerate path; assembled per clip on + # the GPU rollout path (needs a request-API extension to carry action streams). + st.scratch["mg3_cond"] = self._build_clip_cond(st) + st.scratch["stream_video"] = bool(req.outputs.stream.get("video")) + st.plugin_state["cfg"] = {} + return st + + def _build_clip_cond(self, st: LoopState) -> dict: + """Assemble the per-step MG3 conditioning bundle the adapter consumes. On the degenerate + single-clip / no-action path this is just ``cond_frames`` + ``predict_latent_idx`` (the action / + camera / KV-memory slots stay None). BRINGUP: the multi-clip rollout fills ``mouse_cond`` / + ``keyboard_cond`` / ``c2ws_plucker_emb`` / ``x_memory`` / ``timestep_memory`` / + ``memory_latent_idx`` here from the vendored host helpers; that path needs a request-API + extension to carry per-frame action streams + a camera trajectory.""" + return { + "cond_frames": int(st.scratch.get("cond_frames", 0)), + "predict_latent_idx": st.scratch.get("predict_latent_idx", (0, 0)), + "mouse_cond": None, + "keyboard_cond": None, + "x_memory": None, + "timestep_memory": None, + "mouse_cond_memory": None, + "keyboard_cond_memory": None, + "c2ws_plucker_emb": None, + "memory_latent_idx": None, + } + + def next(self, st: LoopState): + i = st.step_idx + if i >= len(st.sigmas) - 1: + return Done() # inner schedule exhausted (single-clip preset -> done) + sigma_t, sigma_next = st.sigmas[i], st.sigmas[i + 1] + t = st.timesteps[i] # FlowUniPC model timestep (~sigma*1000) -> per-token value + 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 + do_cfg = st.scratch.get("use_base_model", False) and scale != 1.0 and ne is not None + mg3_cond = st.scratch["mg3_cond"] + img_cond = st.scratch.get("img_cond") + cond_frames = int(st.scratch.get("cond_frames", 0)) + + def _velocity(model: Any, x_: Any, pe_: Any, ne_: Any) -> Any: + # The conditioned forward + (optional) CFG combine. The adapter builds the per-token timestep + # + packs the MG3 conditioning bundle internally; the loop hands the raw timestep ``t`` as the + # ``sigma`` arg (the FlowUniPC convention) so the per-token tensor carries the right value. + dit = model.component(expert_id) + cond_pred = dit(x_, pe_, t, cond=mg3_cond) + if do_cfg: # base model only (distilled student skips this) + uncond_pred = dit(x_, ne_, t, cond=mg3_cond) + combined = np.asarray(uncond_pred) + scale * (np.asarray(cond_pred) - np.asarray(uncond_pred)) + return precision.cast(combined) + return precision.cast(np.asarray(cond_pred)) + + def run(model: Any, override: Any = None) -> StepResult: + # EAGER thunk. FLOW_MATCH_STEP is the CPU stand-in for FlowUniPC.step (BRINGUP risk 3: the GPU + # path must drive the real stateful multistep scheduler). After the step, RE-PASTE img_cond on + # the cond_frames so the known frames never drift (BRINGUP risk 7). + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model, x, pe, ne) + x_next = kernels.get(FLOW_MATCH_STEP)(precision.cast(x), velocity, sigma_t, sigma_next) + x_next = np.asarray(x_next, dtype="float32") + if img_cond is not None and cond_frames > 0: + x_next[:, :cond_frames] = img_cond[:, :cond_frames] # re-pin the conditioning frames + return StepResult(output={"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next}) + + 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__), ("ar_world_model", True))), + resources=res, + payload={ + "branch": "mg3", + "step": i, + "clip": st.scratch.get("clip_idx", 0) + }, + run=run, + label=f"mg3.denoise.c{st.scratch.get('clip_idx', 0)}.{i}", + # per-token timestep packing + first-frame re-paste + (rollout) host-built conditioning -> + # the eager path (no static-buffer capture; the cosmos2 video2world precedent). + capturable=False) + + 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 + # OUTER clip advance (BRINGUP rollout): when the inner schedule is done and more clips remain, + # snapshot the clip tail into history, slide img_cond -> last-4 latents, and reset the inner + # step counter + draw fresh noise for the next clip. The single-clip preset never enters this. + return st + + def finalize(self, st: LoopState) -> LoopResult: + # Single-clip preset: the final latent is the denoised clip. Multi-clip rollout would concat the + # history tails on the time dim (BRINGUP). + return LoopResult(outputs={"latents": st.latents["video"]}, + metrics={ + "denoise_steps": float(st.step_idx), + "clips": float(st.scratch.get("clip_idx", 0) + 1) + }, + behavior=st.trajectory or None) diff --git a/v2/recipes/matrixgame3/program.py b/v2/recipes/matrixgame3/program.py new file mode 100644 index 000000000..b70038469 --- /dev/null +++ b/v2/recipes/matrixgame3/program.py @@ -0,0 +1,66 @@ +"""Matrix-Game 3.0 program: text_encode -> image VAE-encode (first frame) -> mg3_denoise -> vae_decode. + +Mirrors the cosmos2 / Wan i2v node graph. The MG3 specifics (per-token timestep, first-frame pinning, +autoregressive clip loop) live entirely in ``MatrixGame3DenoiseLoop`` + the ``MatrixGame3DiT`` adapter, so +the node graph is the standard i2v shape: + + * ``text_encode`` — T5 prompt + (empty) negative prompt -> ``text_embeds`` / ``neg_text_embeds``. + * ``image_encode`` — VAE-encode the conditioning image as the first frame -> ``image_latent`` (the + ``MatrixGame3ImageVAEEncodingStage`` normalization: ``(z - latents_mean)/latents_std``, which the + reused ``WanVAE`` adapter applies inside ``encode``). The loop pins the first ``cond_frames`` latent + frames to this and re-pastes after every scheduler step. Absent (no image) -> ``image_latent=None`` + -> the loop runs the degenerate flow-match denoise (the CPU-toy path). + * ``denoise`` — ``MatrixGame3DenoiseLoop`` reads ``text_embeds`` + ``image_latent`` from slots. + * ``vae_decode`` — light_vae decode -> video. + +BRINGUP: the action (mouse/keyboard) + camera (Plücker) conditioning would default to +``build_matrixgame3_action_preset(total_video_frames, seed)`` + a derived camera trajectory and be +written into slots here; that needs a request-API extension to carry the action streams + camera path +(see the loop's ``_build_clip_cond`` hook). The single-clip no-action path is what the registered preset ++ CPU verification exercise. +""" +from __future__ import annotations + +import numpy as np + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind +from v2.recipes.common import text_encode_node_fn as _text_encode + + +def _mg3_image_encode(instance, slots, request, ctx) -> None: + """VAE-encode the conditioning image as the first frame (rest zeros) -> ``image_latent`` (normalized + by the WanVAE adapter's ``_mean_invstd`` path, matching ``MatrixGame3ImageVAEEncodingStage``). No + image -> ``image_latent`` stays None and the loop runs the degenerate flow-match denoise.""" + img = request.image() + if img is None or getattr(img, "pixels", None) is None: + slots["image_latent"] = None + return + px = np.asarray(img.pixels, dtype="float32") # [3, H, W] in [-1, 1] + # A single-frame "video" tensor [3, 1, H, W] -> the VAE encodes one latent frame the loop pins. + cond_video = px[:, None] + slots["image_latent"] = np.asarray(instance.component("vae").encode(cond_video), dtype="float32") + + +def _mg3_vae_decode(instance, slots, request, ctx) -> None: + slots["video"] = instance.component("vae").decode(slots["denoise_out"]["latents"]) + + +def build_matrixgame3_program() -> Program: + return Program( + program_id="matrixgame3.i2v.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("text_encode", fn=_text_encode, writes=("text_embeds", "neg_text_embeds")), + ComponentNode("image_encode", fn=_mg3_image_encode, writes=("image_latent", )), + ModelLoopNode("denoise", + loop_id="mg3_denoise", + output_slot="denoise_out", + reads=("text_embeds", "neg_text_embeds", "image_latent"), + writes=("denoise_out", )), + ComponentNode("vae_decode", fn=_mg3_vae_decode, reads=("denoise_out", ), writes=("video", )), + ], + output_artifacts={ + "video": "video", + "latents": "denoise_out" + }, + ).validate() diff --git a/v2/recipes/sd35/__init__.py b/v2/recipes/sd35/__init__.py new file mode 100644 index 000000000..c38b7cea3 --- /dev/null +++ b/v2/recipes/sd35/__init__.py @@ -0,0 +1,16 @@ +"""Stable Diffusion 3.5 Medium (text→image) — a FLOW-MATCH MMDiT ported into the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``SD3DiT``/``SD3VAE``/``SD3ClipEncoder``/``SD3T5Encoder`` in +``v2/platform/backends/torch_sd35.py``) and a new ``SD3DenoiseLoop`` (flow-match Euler over 4D image +latents with dual text conditioning: the triple-encoder joint embed + the dual-CLIP pooled vector). It +reuses the ``FLOW_MATCH_STEP`` kernel + ``FlowShiftPolicy`` + ``ClassicCFG``. Registered in ``v2/registry.py`` +by the orchestrator (transformer ``_class_name`` = ``SD3Transformer2DModel`` for the arch fallback). +""" +from __future__ import annotations + +from v2.recipes.sd35.card import build_sd35_card +from v2.recipes.sd35.loop import SD3DenoiseLoop +from v2.recipes.sd35.program import build_sd35_program + +__all__ = ["build_sd35_card", "build_sd35_program", "SD3DenoiseLoop"] diff --git a/v2/recipes/sd35/card.py b/v2/recipes/sd35/card.py new file mode 100644 index 000000000..1c273f7b4 --- /dev/null +++ b/v2/recipes/sd35/card.py @@ -0,0 +1,185 @@ +"""Stable Diffusion 3.5 Medium ModelCard — text→image (the registered preset). Image is treated as a +1-frame video (num_frames=1, 4D latents) throughout the v2 substrate. + +Architecture deltas vs Wan (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.sd3:SD3Transformer2DModel`` — a FLOW-MATCH MMDiT (reuses the + flow-match step), but the forward needs TWO text conditioners: the assembled triple-encoder joint + embed (``encoder_hidden_states``) AND a separate ``pooled_projections`` vector. The adapter + (``ComponentSpec.adapter`` -> ``SD3DiT``) takes both (joint embed as ``text_embed``, pooled as the + positional ``context`` arg) and ``SD3DenoiseLoop`` threads them per CFG branch. + * VAE ``fastvideo.models.vaes.autoencoder_kl:AutoencoderKL`` — SD3.5 shift/scale normalization + (``scaling_factor=1.5305``, ``shift_factor=0.0609`` read from the checkpoint vae config); the + ``SD3VAE`` adapter normalizes/denormalizes around the raw AutoencoderKL space (4D image, not 5D video). + * Text x3: ``CLIPTextModelWithProjection`` (clip_l + clip_g, penultimate-hidden + pooled) + + ``T5EncoderModel`` (last_hidden_state). The program assembles the joint embed + dual-CLIP pooled. + +Schedule: plain flow-match linspace with the diffusers ``FlowMatchEulerDiscreteScheduler`` config +``shift`` (3.0 for SD3.5-medium). The scheduler's optional ``use_dynamic_shifting`` (resolution-dependent +``mu``) is NOT applied here — if a checkpoint's scheduler config enables it, the static shift will mismatch +the reference schedule (BRINGUP blocker 6). CFG is the classic ``uncond + s·(cond − uncond)``. +""" +from __future__ import annotations + +import os + +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, FlowShiftPolicy, NoRouting, PrecisionPolicy +from v2.parallel import ParallelPlan +from v2.platform.backends.toy import ToyTextEncoder, ToyVAE, _seed_from +from v2.recipes.sd35.loop import SD3DenoiseLoop + +# SD3.5's default negative prompt is empty (LOCAL module constant — the spec's defaults.negative_prompt). +SD35_NEG = "" + +# SD3.5-medium diffusers FlowMatchEulerDiscreteScheduler static flow shift (BRINGUP blocker 8: read the +# real scheduler config at deploy time; SD35Config.flow_shift is None, the shift lives in the scheduler). +SD35_FLOW_SHIFT = 3.0 + +_SD3_DIT = "v2.platform.backends.torch_sd35:SD3DiT" +_SD3_VAE = "v2.platform.backends.torch_sd35:SD3VAE" +_SD3_CLIP = "v2.platform.backends.torch_sd35:SD3ClipEncoder" +_SD3_T5 = "v2.platform.backends.torch_sd35:SD3T5Encoder" + + +def build_sd35_card(model_id: str = "stable-diffusion-3.5-medium", + *, + checkpoint_root: str | None = None, + flow_shift: float = SD35_FLOW_SHIFT, + 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() + flow = FlowShiftPolicy(shift=flow_shift) # static shift (no resolution-bucket lookup for SD3.5) + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return SD3DenoiseLoop(loop_id="diffusion_denoise", + cfg=cfg, + flow_shift=flow, + precision=precision, + expert=expert, + cost=cost) + + components = { + # The two CLIP encoders (clip_l, clip_g) — penultimate-hidden + pooled. The third is T5. + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.clip:CLIPTextModelWithProjection", + adapter=_SD3_CLIP, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2i"}), + "text_encoder_2": + ComponentSpec(component_id="text_encoder_2", + kind="text_encoder", + load_id="fastvideo.models.encoders.clip:CLIPTextModelWithProjection", + adapter=_SD3_CLIP, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2i"}), + "text_encoder_3": + ComponentSpec(component_id="text_encoder_3", + kind="text_encoder", + load_id="fastvideo.models.encoders.t5:T5EncoderModel", + adapter=_SD3_T5, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2i"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.autoencoder_kl:AutoencoderKL", + adapter=_SD3_VAE, + factory=lambda inst: ToyVAE(), + required_for={"t2i"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.sd3:SD3Transformer2DModel", + adapter=_SD3_DIT, + factory=lambda inst: _toy_dit(seed), + resident_for=["diffusion_denoise"], + required_for={"t2i"}), + } + 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="eager", + loop_factory=loop_factory), + } + card = ModelCard( + model_id=model_id, + family="sd3", + components=components, + loops=loops, + capabilities=CapabilityMatrix.of(Capability.TEXT_TO_IMAGE, 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=28, guidance_scale=6.0, height=512, width=512, num_frames=1, fps=1, negative_prompt=SD35_NEG), + ) + card.validate() + if checkpoint_root: + stamp_sd35_checkpoints(card, checkpoint_root) + return card + + +def _toy_dit(seed: int): + """CPU toy DiT for the transformer component. ToyDiT.__call__ accepts the positional ``context`` arg + (it mean-pools it harmlessly), so the dual-conditioning ``dit(latent, joint_embed, sigma, context=pooled)`` + call CPU-verifies without any SD3-specific stand-in.""" + from v2.platform.backends.toy import ToyDiT + return ToyDiT(seed=seed) + + +# Diffusers SD3.5 checkpoint layout: transformer/ + vae/ + text_encoder{,_2,_3}/ + tokenizer{,_2,_3}/. +# Distinct from ``_WAN21_SUBFOLDERS`` (one text_encoder) — SD3 carries THREE text encoders. NOTE: the +# torch backend's ``_make_text_encoder`` loads the tokenizer from a FIXED ``/tokenizer`` sibling, so +# text_encoder_2/_3 would load the wrong tokenizer (BRINGUP blocker: needs the per-encoder tokenizer +# subfolder threaded into the maker, or a small adapter-side override at GPU bring-up). +_SD35_SUBFOLDERS = { + "transformer": "transformer", + "vae": "vae", + "text_encoder": "text_encoder", + "text_encoder_2": "text_encoder_2", + "text_encoder_3": "text_encoder_3", +} + + +def stamp_sd35_checkpoints(card: ModelCard, model_root: str) -> ModelCard: + """GPU deploy-time: point each component's ``ComponentSpec.checkpoint`` at its weights subfolder under + ``model_root`` (a local diffusers dir, or an HF id resolved via ``snapshot_download``). Mirrors + ``stamp_wan21_checkpoints`` but covers the THREE SD3.5 text encoders. Mutates and returns the card.""" + if not os.path.isdir(model_root): + from huggingface_hub import snapshot_download + model_root = snapshot_download(model_root) + for component_id, subfolder in _SD35_SUBFOLDERS.items(): + if component_id in card.components: + card.components[component_id].checkpoint = os.path.join(model_root, subfolder) + return card diff --git a/v2/recipes/sd35/loop.py b/v2/recipes/sd35/loop.py new file mode 100644 index 000000000..459aa5d76 --- /dev/null +++ b/v2/recipes/sd35/loop.py @@ -0,0 +1,189 @@ +"""SD3DenoiseLoop — a thin flow-match Euler denoise loop for Stable Diffusion 3.5 (image / 4D latents). + +SD3.5 is a plain FLOW-MATCH MMDiT, so this loop reuses the canonical pieces (``FLOW_MATCH_STEP`` kernel, +``FlowShiftPolicy.build_schedule``, ``ClassicCFG``) — the only deltas vs ``WanDenoiseLoop`` are: + + * IMAGE modality / 4D latents: the latent is ``[16, h//8, w//8]`` (num_frames=1), with NO temporal + dim. The fastvideo stage carries a fake T dim in its 5D ForwardBatch (unsqueeze(2)/squeeze(2)) but + the DiT forward itself is strictly 4D — so we drop the T dim here (BRINGUP blocker 3). + * DUAL DiT conditioning (BRINGUP blocker 1): the SD3 forward needs BOTH the assembled triple-encoder + joint embed (``encoder_hidden_states``) AND a separate pooled vector (``pooled_projections``). The + program writes both into slots; this loop threads the joint embed as ``text_embed`` and the pooled + vector as the positional ``context`` arg — so the dit-call stays the cosmos2-style + ``dit(latent, text_embed, sigma, context)`` that ToyDiT accepts (CPU-verifiable) and SD3DiT unpacks. + * VELOCITY output: SD3DiT returns the flow-match velocity directly (no x0 reconstruction), so the + Euler update is the plain ``x_next = x + (sigma_next - sigma)*v``. + +Faithful to ``fastvideo/pipelines/stages/sd35_conditioning.py:SD35DenoisingStage`` (uncond + s·(cond − +uncond) CFG over the FlowMatchEuler schedule). The σ schedule is the flow-shift linspace; ``shift`` comes +from the diffusers ``FlowMatchEulerDiscreteScheduler`` config (3.0 for SD3.5-medium) — see ``card.py``. +""" +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.platform import FLOW_MATCH_STEP +from v2.platform.backends.toy import LATENT_CHANNELS + +# Real SD3.5 AutoencoderKL: z=16, 8x spatial compression, IMAGE (no temporal dim). +SD3_LATENT_CHANNELS = 16 +SD3_SPATIAL_RATIO = 8 + + +def sd3_latent_shape(req, + model=None, + *, + channels: int = SD3_LATENT_CHANNELS, + spatial_ratio: int = SD3_SPATIAL_RATIO) -> tuple[int, int, int]: + """4D image latent geometry ``[C, h, w]`` (the leading batch dim is added by the adapter). On the GPU + backend the true SD3.5 geometry is required; the CPU toy uses a tiny deterministic stand-in (matching + ``v2.recipes.wan21.loop.latent_shape`` but with NO temporal axis — image, not video).""" + d = req.diffusion + if model is not None and getattr(getattr(model, "platform", None), "device", "cpu") == "cuda": + return (channels, max(1, d.height // spatial_ratio), max(1, d.width // spatial_ratio)) + h = max(2, d.height // 120) + w = max(2, d.width // 120) + return (LATENT_CHANNELS, h, w) + + +class SD3DenoiseLoop: + """Flow-match Euler denoise for SD3.5 (image / 4D latents, dual text conditioning).""" + + def __init__(self, + *, + loop_id, + cfg, + flow_shift, + precision, + expert, + cost, + latent_channels: int = SD3_LATENT_CHANNELS, + spatial_ratio: int = SD3_SPATIAL_RATIO): + self.loop_id = loop_id + self.cfg = cfg + self.flow_shift = flow_shift + self.precision = precision + self.expert = expert + self.cost = cost + self.latent_channels = latent_channels + self.spatial_ratio = spatial_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 = self.flow_shift.build_schedule(req.diffusion.num_steps, + req.diffusion.height, + req.diffusion.width, + sigmas=req.diffusion.sigmas or None) + shape = sd3_latent_shape(req, model, channels=self.latent_channels, spatial_ratio=self.spatial_ratio) + x = (rng.standard_normal(shape) * float(sig[0])).astype("float32") + 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={"image": x}, + sigmas=[float(s) for s in sig], + timesteps=[float(s) * 1000.0 for s in sig]) + # Dual text conditioning (BRINGUP blocker 1): the program's SD3 text-encode node writes the + # assembled joint embed AND the pooled vector for both the cond and uncond branches into slots. + st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") + st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") + st.cond["pooled"] = ctx.slots.get("pooled_projections") + st.cond["neg_pooled"] = ctx.slots.get("neg_pooled_projections") + st.scratch["guidance_scale"] = float(req.diffusion.guidance_scale) + st.scratch["stream_image"] = bool(req.outputs.stream.get("image") or req.outputs.stream.get("video")) + 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] + expert_id = self.expert.expert_for(StepContext(i, st.timesteps[i], sigma_t)) + sctx = StepContext(step_idx=i, timestep=st.timesteps[i], sigma=sigma_t, active_expert_id=expert_id) + cfg_state = st.plugin_state["cfg"] + branches = self.cfg.branches_this_step(sctx, cfg_state) + x = st.latents["image"] + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + pp, npp = st.cond["pooled"], st.cond["neg_pooled"] + scale = st.scratch["guidance_scale"] + cfg, precision = self.cfg, self.precision + + def _velocity(model: Any) -> np.ndarray: + # The conditioned forward + CFG combine. Each branch threads its OWN pooled vector via the + # positional ``context`` arg (so SD3DiT carries pooled_projections alongside the joint embed, + # and ToyDiT harmlessly mean-pools it). Solver/forward dispatch through the kernel table. + dit = model.component(expert_id) + pooled_for = {"cond": pp, "uncond": npp} + embed_for = {"cond": pe, "uncond": ne} + preds = {b: dit(x, embed_for[b], sigma_t, context=pooled_for[b]) for b in branches} + return precision.cast(cfg.combine(preds, scale, sctx, cfg_state)) + + def run(model: Any, override: Any = None) -> StepResult: + kernels = model.platform.kernels + if override is not None and "noise_pred" in override: + velocity = precision.cast(np.asarray(override["noise_pred"], dtype="float32")) + else: + velocity = _velocity(model) + x_next = 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, pp, npp) if e is not None) + res = ResourceRequest(compute_seconds=self.cost.predict(int(np.prod(x.shape)), float(len(branches))), + 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(cfg).__name__), ("image", True))), + resources=res, + payload={ + "branch": "combined", + "step": i + }, + run=run, + label=f"sd35.denoise.{i}", + # Dual-conditioning forward with a per-branch pooled vector threaded outside the workspace -> + # keep the eager path (no static-buffer graph capture wired for this recipe). The deterministic + # ODE step itself is capturable, but wiring the full capture form is a later GPU-side concern. + capturable=False) + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + st.latents["image"] = 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["image"]).copy(), + }) + st.step_idx += 1 + return st + + def finalize(self, st: LoopState) -> LoopResult: + return LoopResult(outputs={"latents": st.latents["image"]}, + metrics={"denoise_steps": float(st.step_idx)}, + behavior=st.trajectory or None) diff --git a/v2/recipes/sd35/program.py b/v2/recipes/sd35/program.py new file mode 100644 index 000000000..d655224bd --- /dev/null +++ b/v2/recipes/sd35/program.py @@ -0,0 +1,90 @@ +"""SD3.5 t2i program: sd3_text_encode → diffusion_denoise (flow-match) → vae_decode. + +Single-stage INLINE program (like ``build_ltx2_base_program``). The SD3-specific work lives in the +custom text-encode node (the triple-encoder joint-embed + dual-CLIP pooled assembly that +``SD35ConditioningStage`` does) and the VAE decode node (the [-1,1] -> [0,1] remap of +``SD35DecodingStage``); the denoise node is the generic ``ModelLoopNode`` binding ``SD3DenoiseLoop``. + +The text-encode node is backend-agnostic: on the GPU torch backend the CLIP adapters +(``SD3ClipEncoder``) return ``{"hidden": penultimate[-2], "pooled": pooler_output}`` and ``SD3T5Encoder`` +returns ``last_hidden_state``, which ``assemble_sd3_conditioning`` combines into (joint_embed, pooled). +On the CPU toy backend the encoders return plain arrays (``ToyTextEncoder.encode``), so the node falls +back to a degenerate-but-finite assembly (joint = the array, pooled = its channel-mean broadcast) — enough +to exercise the dual-conditioning plumbing end-to-end without any GPU/torch. +""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from v2.platform.backends.torch_sd35 import assemble_sd3_conditioning +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind + +_CLIP_L = "text_encoder" +_CLIP_G = "text_encoder_2" +_T5 = "text_encoder_3" + + +def _encode_one(instance: Any, prompt: str) -> tuple[np.ndarray, np.ndarray]: + """Run the three SD3.5 text encoders on ``prompt`` and assemble (joint_embed, pooled). + + Handles both backends: torch CLIP adapters return a ``{"hidden","pooled"}`` dict (-> the faithful + ``assemble_sd3_conditioning``); the CPU toy encoders return a plain array (-> a degenerate, finite + stand-in that still threads a distinct joint embed + pooled vector through the dual-conditioning path).""" + clip_l = instance.component(_CLIP_L).encode(prompt) + clip_g = instance.component(_CLIP_G).encode(prompt) + t5 = instance.component(_T5).encode(prompt) + if isinstance(clip_l, dict) and isinstance(clip_g, dict): + return assemble_sd3_conditioning(clip_l, clip_g, t5) + # CPU toy fallback: ToyTextEncoder.encode returns a [seq, dim] array (no pooled head). Concatenate + # the three along seq for the joint embed; derive a small pooled vector from the CLIP means. + a_l, a_g, a_t = np.asarray(clip_l), np.asarray(clip_g), np.asarray(t5) + joint = np.concatenate([a_l, a_g, a_t], axis=-2).astype("float32") + pooled = np.concatenate([a_l.mean(axis=-2), a_g.mean(axis=-2)], axis=-1).astype("float32") + return joint, pooled + + +def _sd3_text_encode(instance, slots, request, ctx) -> None: + joint, pooled = _encode_one(instance, request.prompt()) + neg_joint, neg_pooled = _encode_one(instance, request.diffusion.negative_prompt) + slots["text_embeds"] = joint + slots["pooled_projections"] = pooled + slots["neg_text_embeds"] = neg_joint + slots["neg_pooled_projections"] = neg_pooled + + +def _vae_decode(instance, slots, request, ctx) -> None: + # SD3DiT denoises in normalized latent space; SD3VAE.decode inverts shift/scale and runs the + # AutoencoderKL decode -> image in [-1, 1]. Remap to [0, 1] (SD35DecodingStage: (img/2+0.5).clamp). + vae = instance.component("vae") + latent = np.asarray(slots["denoise_out"]["latents"], dtype="float32") # [C, h, w] image latent + # The real SD3VAE adapter consumes the strict 4D image latent [C, h, w] directly. The CPU toy ToyVAE + # is video-oriented ([C, T, h, w]); give it a singleton temporal axis so the decode CPU-verifies + # (faithful to the fastvideo stage carrying a fake T dim around the strictly-4D DiT forward). + if getattr(vae, "module", None) is None and latent.ndim == 3: + latent = latent[:, None] # -> [C, 1, h, w] for the toy decoder only + image = vae.decode(latent) + slots["image"] = np.clip(np.asarray(image, dtype="float32") / 2.0 + 0.5, 0.0, 1.0) + + +def build_sd35_program() -> Program: + return Program( + program_id="sd35.t2i.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("sd3_text_encode", + fn=_sd3_text_encode, + writes=("text_embeds", "neg_text_embeds", "pooled_projections", "neg_pooled_projections")), + ModelLoopNode("denoise", + loop_id="diffusion_denoise", + output_slot="denoise_out", + reads=("text_embeds", "neg_text_embeds", "pooled_projections", "neg_pooled_projections"), + writes=("denoise_out", )), + ComponentNode("vae_decode", fn=_vae_decode, reads=("denoise_out", ), writes=("image", )), + ], + output_artifacts={ + "image": "image", + "latents": "denoise_out", + }, + ).validate() diff --git a/v2/recipes/stable_audio/__init__.py b/v2/recipes/stable_audio/__init__.py new file mode 100644 index 000000000..d859ae15e --- /dev/null +++ b/v2/recipes/stable_audio/__init__.py @@ -0,0 +1,21 @@ +"""Stable Audio Open (text→audio) — a v-prediction (EDM-v / VDenoiser) DPM++ recipe in the v2 substrate. + +Self-contained recipe package (the bucket-C pattern): the card declares its torch adapters via +``ComponentSpec.adapter`` (``StableAudioDiT``/``OobleckVAE``/``StableAudioConditioner`` in +``v2/platform/backends/torch_stable_audio.py``) and a new ``StableAudioDenoiseLoop`` (polyexponential +schedule + VDenoiser v->x0 + DPM++ multistep), reusing ``stamp_wan21_checkpoints``. AUDIO modality — +the audio output artifact + a TEXT_TO_AUDIO capability/request field are the request-API extension +(BRINGUP). Registered in ``v2/registry.py`` by the orchestrator. +""" +from __future__ import annotations + +from v2.recipes.stable_audio.card import build_stable_audio_card, build_stable_audio_small_card +from v2.recipes.stable_audio.loop import StableAudioDenoiseLoop +from v2.recipes.stable_audio.program import build_stable_audio_program + +__all__ = [ + "build_stable_audio_card", + "build_stable_audio_small_card", + "build_stable_audio_program", + "StableAudioDenoiseLoop", +] diff --git a/v2/recipes/stable_audio/card.py b/v2/recipes/stable_audio/card.py new file mode 100644 index 000000000..1b40618ae --- /dev/null +++ b/v2/recipes/stable_audio/card.py @@ -0,0 +1,223 @@ +"""Stable Audio Open ModelCard(s) — text→audio (the registered preset). AUDIO modality. + +Two variants from the converted Diffusers repos: + * ``FastVideo/stable-audio-open-1.0-Diffusers`` — depth 24, embed_dim 1536, qk_norm None, 3 cond + ids (prompt + seconds_start + seconds_total), 2,097,152-frame window (~47.55s @ 44.1kHz). + * ``FastVideo/stable-audio-open-small-Diffusers`` — depth 16, embed_dim 1024, qk_norm 'ln', 2 cond + ids (prompt + seconds_total), 524,288-frame window (~11.89s). (depth/dim/qk_norm are resolved by the + loader from the checkpoint config; the card carries the per-variant sample_size + audio defaults.) + +Architecture deltas vs the video recipes (all declared on the card so the recipe is self-contained): + * DiT ``fastvideo.models.dits.stable_audio:StableAudioDiT`` — a v-PREDICTION network (EDM-v / VDenoiser), + NOT flow-match and NOT EDM-Karras. The adapter (``ComponentSpec.adapter`` -> ``StableAudioDiT``) packs + the conditioner payload into (cross_attn_cond, global_embed), passes the RAW continuous sigma as the + timestep, and returns the raw v output; ``StableAudioDenoiseLoop`` does the VDenoiser v->x0 + x0-space + CFG + the polyexponential schedule (sigma_min 0.3, sigma_max 500, rho 1) and a DPM++ multistep step. + * VAE ``fastvideo.models.vaes.oobleck:OobleckVAE`` — 1-D conv audio VAE (hop_length 2048). The + ``OobleckVAE`` adapter is RAW latent space (NO mean/std normalization, unlike WanVAE) and exposes + ``sampling_rate`` for the decode-slice. + * Conditioner ``fastvideo.models.encoders.stable_audio_conditioner:StableAudioMultiConditioner`` — T5-base + (max_length 128) + duration NumberConditioners. Declared with ``kind="text_encoder"`` (kind reuse — + the shared ``_MAKERS`` has no ``conditioner`` kind and editing it is out of scope) + the explicit + ``StableAudioConditioner`` adapter. BRINGUP: the text-encoder maker calls ``TextEncoderLoader`` (a + generic T5), so wiring the real SA conditioner on a GPU box needs the SA ``ConditionerLoader`` path. + +``stamp_wan21_checkpoints`` applies (diffusers transformer/vae/text_encoder subfolder layout — the +conditioner lives in the ``text_encoder`` subfolder of the converted repo). The negative prompt is a LOCAL +module constant (Stable Audio's default negative prompt is empty). +""" +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 ToyTextEncoder, _seed_from +from v2.recipes.stable_audio.loop import StableAudioDenoiseLoop +from v2.recipes.wan21.card import stamp_wan21_checkpoints + +# Stable Audio's default negative prompt is empty (no shared _prompts entry; kept LOCAL per the port spec). +STABLE_AUDIO_NEG = "" + +# Oobleck VAE io_channels — the 1-D audio latent's channel dim (shared by the toy DiT + toy VAE below). +SA_IO_CHANNELS = 64 + +_SA_DIT = "v2.platform.backends.torch_stable_audio:StableAudioDiT" +_SA_VAE = "v2.platform.backends.torch_stable_audio:OobleckVAE" +_SA_CONDITIONER = "v2.platform.backends.torch_stable_audio:StableAudioConditioner" + + +def build_stable_audio_card(model_id: str = "stable-audio-open-1.0", + *, + checkpoint_root: str | None = None, + sigma_min: float = 0.3, + sigma_max: float = 500.0, + rho: float = 1.0, + sample_size: int = 2097152, + 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() + # fp16 throughout on the GPU path (DiT/VAE/T5); the loop math is numpy fp32 (lowest-risk bring-up). + precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True) + expert = NoRouting("transformer") + + def loop_factory(): + return StableAudioDenoiseLoop(loop_id="stable_audio_denoise", + cfg=cfg, + precision=precision, + expert=expert, + cost=cost, + sigma_min=sigma_min, + sigma_max=sigma_max, + rho=rho, + sample_size=sample_size) + + components = { + # The SA multi-conditioner, declared under the ``text_encoder`` kind (kind reuse). On CPU this is + # the ToyTextEncoder; on GPU the StableAudioConditioner adapter (BRINGUP: needs the SA loader). + "text_encoder": + ComponentSpec(component_id="text_encoder", + kind="text_encoder", + load_id="fastvideo.models.encoders.stable_audio_conditioner:StableAudioMultiConditioner", + adapter=_SA_CONDITIONER, + factory=lambda inst: ToyTextEncoder(), + required_for={"t2a"}), + "vae": + ComponentSpec(component_id="vae", + kind="vae", + load_id="fastvideo.models.vaes.oobleck:OobleckVAE", + adapter=_SA_VAE, + factory=lambda inst: _ToyOobleckVAE(), + required_for={"t2a"}), + "transformer": + ComponentSpec(component_id="transformer", + kind="dit", + load_id="fastvideo.models.dits.stable_audio:StableAudioDiT", + adapter=_SA_DIT, + factory=lambda inst: _ToyAudioDiT(seed=seed), + resident_for=["stable_audio_denoise"], + required_for={"t2a"}), + } + loops = { + "stable_audio_denoise": + LoopSpec( + loop_id="stable_audio_denoise", + kind=LoopKind.DIFFUSION_DENOISE, + work_unit_kind=WorkUnitKind.DIFFUSION_STEP, + step_cost_model=cost, + shared_weight_components=["transformer"], + cache_policy=["feature"], + graph_capture="eager", # v->x0 + x0-space CFG + multistep host state -> eager + loop_factory=loop_factory), + } + card = ModelCard( + model_id=model_id, + family="stable_audio", + components=components, + loops=loops, + # No TEXT_TO_AUDIO capability on the v2 enum (audio output surface is the request-API extension — + # BRINGUP); TEXT_TO_SPEECH is the nearest audio-output capability. VAE_DECODE + POLICY_ROLLOUT as usual. + capabilities=CapabilityMatrix.of(Capability.TEXT_TO_SPEECH, Capability.VAE_DECODE, Capability.POLICY_ROLLOUT), + recipe=RecipeSpec(method="base", + assumes_loop="stable_audio_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()), + # Audio-shaped defaults. height/width/num_frames/fps are video-shaped and unused for audio (BRINGUP: + # SamplingDefaults has no audio fields — the duration knobs ride node_params; see program.py). + # num_steps=100, guidance_scale=7.0 are the published SA-1.0 defaults. + sampling_defaults=sampling_defaults + or SamplingDefaults(num_steps=100, guidance_scale=7.0, negative_prompt=STABLE_AUDIO_NEG), + ) + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card + + +def build_stable_audio_small_card(model_id: str = "stable-audio-open-small", + *, + checkpoint_root: str | None = None) -> ModelCard: + """stable-audio-open-small: a 16-layer/1024-dim DiT with qk_norm='ln' and a shorter 524,288-frame + window (~11.89s); 2 cond ids (prompt + seconds_total). Shares the Oobleck VAE + the v-prediction + dpmpp-3m-sde loop with SA-1.0 — only the sample_size + the (loader-resolved) DiT geometry differ. The + published small-variant default is a shorter clip (audio_end_in_s 6s) at fewer steps.""" + return build_stable_audio_card(model_id=model_id, + checkpoint_root=checkpoint_root, + sample_size=524288, + sampling_defaults=SamplingDefaults(num_steps=8, + guidance_scale=7.0, + negative_prompt=STABLE_AUDIO_NEG)) + + +class _ToyAudioDiT: + """CPU toy v-prediction stand-in for ``StableAudioDiT`` (1-D audio latent [64, L]). + + The shared ``ToyDiT`` is channel-fixed (4 channels) and shaped for video; SA's latent is 64-channel + 1-D. This tiny deterministic toy honors the same contract the loop needs — ``__call__(latent, + text_embed, sigma) -> v`` with the loop-supplied VDenoiser-scaled input — so the loop runs end-to-end + on CPU. ``text_embed`` is the packed conditioning payload (the toy just means-pools it as a scalar + conditioning signal, matching ToyDiT's text handling). Deterministic given (seed, inputs).""" + + def __init__(self, seed: int = 0): + import numpy as np + rng = np.random.default_rng(seed) + self.w_t = float(rng.standard_normal() * 0.1) + self.s_text = 0.5 + + def __call__(self, latent, text_embed, sigma, context=None, *, cond=None): + import numpy as np + lat = np.asarray(latent, dtype="float32") + cond_signal = float(np.mean(text_embed)) if text_embed is not None else 0.0 + # A simple bounded v-prediction: channel-local nonlinearity + timestep + conditioning bias. + pre = 0.5 * lat + self.w_t * float(sigma) + self.s_text * cond_signal + return np.tanh(pre).astype("float32") + + +class _ToyOobleckVAE: + """CPU toy 1-D audio VAE stand-in for ``OobleckVAE`` (the shared ``ToyVAE`` is a video VAE that decodes + a 4-channel ``[C,T,H,W]`` latent). Decodes the 64-channel ``[64, L]`` audio latent to a + ``[channels, samples]`` waveform (channel-collapse projection + hop_length upsample), deterministic so + the program node is interleave-safe. Exposes ``hop_length`` / ``sampling_rate`` like the real adapter.""" + + def __init__(self, channels: int = 2, hop_length: int = 16, sampling_rate: int = 44100, seed: int = 4): + import numpy as np + self.hop_length = int(hop_length) # tiny toy hop (real Oobleck is 2048); just for upsample shape + self.sampling_rate = int(sampling_rate) + self.channels = int(channels) + rng = np.random.default_rng(seed) + self.proj = (rng.standard_normal((channels, SA_IO_CHANNELS)) * 0.2).astype("float32") + + def decode(self, latent): + import numpy as np + z = np.asarray(latent, dtype="float32") # [64, L] + per_ch = np.tensordot(self.proj, z, axes=([1], [0])) # [channels, L] + wav = np.repeat(per_ch, self.hop_length, axis=-1) # [channels, L*hop] + return np.tanh(wav).astype("float32") + + def encode(self, audio): + import numpy as np + a = np.asarray(audio, dtype="float32") # [channels, samples] + flat = a.reshape(a.shape[0], -1) + # crude inverse: project back to 64 channels at the downsampled rate (passthrough latent space) + z = np.tensordot(self.proj.T, flat[:, ::self.hop_length], axes=([1], [0])) # [64, ~L] + return z.astype("float32") diff --git a/v2/recipes/stable_audio/loop.py b/v2/recipes/stable_audio/loop.py new file mode 100644 index 000000000..306905318 --- /dev/null +++ b/v2/recipes/stable_audio/loop.py @@ -0,0 +1,202 @@ +"""StableAudioDenoiseLoop — a v-prediction (EDM-v / VDenoiser) DPM++ denoise loop for Stable Audio Open. + +Stable Audio is NOT a flow-match model and NOT an EDM-Karras model: it is a v-prediction network sampled +by k-diffusion's ``dpmpp-3m-sde`` (a 3rd-order multistep stochastic DPM-Solver++) over a POLYEXPONENTIAL +sigma schedule (sigma_min=0.3, sigma_max=500, rho=1.0). So neither ``WanDenoiseLoop`` (flow-match) nor +``CosmosDenoiseLoop`` (EDM-Karras x0) can be reused — this is a NEW loop (the spec's +``NEEDS_NEW:StableAudioDenoiseLoop``). + +Faithfulness vs ``fastvideo/pipelines/basic/stable_audio/stages/denoising.py``: + * The schedule is polyexponential (``sampler.build_polyexponential_sigmas``), seeded noise scaled by + sigmas[0], exactly like ``StableAudioLatentPreparationStage`` + ``StableAudioDenoisingStage``. + * The network is a v-predictor wrapped by ``K.external.VDenoiser``: ``denoised = c_skip·x + c_out·F(c_in·x, c_noise)`` + with the EDM-v coefficients (sigma_data=1). This loop reproduces those coefficients exactly + (``sampler.vdenoiser_x0``) so the x0 it integrates matches k-diffusion's VDenoiser on the GPU path. + * CFG: ``denoised = uncond + (cond - uncond)·scale``, combined in x0 space (the + ``_DiTAdapter`` convention; the negative cross-attn is null/zero, applied by the conditioning node). + +GPU PATH (BRINGUP — k_diffusion is an external dep, NOT installed in this env): the faithful, lowest-risk +GPU route is to NOT decompose ``dpmpp-3m-sde`` into per-step v2 WorkPlans but to run the whole +k-diffusion sampler inside ONE thunk: ``import k_diffusion as K; K.sampling.sample_dpmpp_3m_sde( +K.external.VDenoiser(adapter), x, sigmas, ...)`` — a direct lift of the fastvideo stage. That requires a +per-loop "single-shot" execution mode the v2 ``ctx.execute`` per-step contract does not yet expose +(BRINGUP: needs a loop-as-one-thunk hook). For CPU-verification + per-step serving granularity, this loop +ships a DETERMINISTIC DPM-Solver++(2M) recurrence in numpy over the SAME VDenoiser x0 — a faithful +v-prediction multistep integrator that runs end-to-end on the CPU toy. The SDE (stochastic) higher-order +``dpmpp-3m-sde`` terms are the GPU-parity refinement (BRINGUP) layered on top of this deterministic spine. + +A2A (init_audio -> lowered sigma_max) and RePaint inpainting (per-step blend callback) are threaded as +inert hooks (None for the registered T2A preset) and documented BRINGUP — base port is pure text->audio. +""" +from __future__ import annotations + +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.recipes.stable_audio.sampler import ( + build_polyexponential_sigmas, + dpmpp_2m_step, + vdenoiser_x0, +) + +# Oobleck VAE geometry: 64 io_channels; latent_len = sample_size // hop_length (hop_length=2048). +SA_IO_CHANNELS = 64 +SA_HOP_LENGTH = 2048 +SA_SAMPLE_SIZE = 2097152 # SA-1.0 fixed training window (~47.55s @ 44.1kHz); SA-small is 524288. + + +def audio_latent_len(req, model, *, sample_size: int, hop_length: int = SA_HOP_LENGTH) -> int: + """Latent length the DiT/VAE operate over. The model ALWAYS samples a fixed-size latent + (``sample_size // hop_length``) and the decode is sliced to [start,end] seconds afterwards — the + duration knob does NOT shrink the latent (faithful to the fastvideo stage). The CPU toy uses a tiny + deterministic stand-in so the loop math runs without the 1024-long real latent.""" + if model is not None and getattr(getattr(model, "platform", None), "device", "cpu") == "cuda": + return max(1, sample_size // hop_length) + return 8 # toy: short latent, enough to exercise the multistep recurrence + + +class StableAudioDenoiseLoop: + + def __init__(self, + *, + loop_id, + cfg, + precision, + expert, + cost, + sigma_min: float = 0.3, + sigma_max: float = 500.0, + rho: float = 1.0, + io_channels: int = SA_IO_CHANNELS, + hop_length: int = SA_HOP_LENGTH, + sample_size: int = SA_SAMPLE_SIZE): + 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_min = sigma_min + self.sigma_max = sigma_max + self.rho = rho + self.io_channels = io_channels + self.hop_length = hop_length + self.sample_size = sample_size + + 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) + # Polyexponential schedule sigma_max -> sigma_min (NOT flow-shift, NOT Karras-rho-7). + sig = build_polyexponential_sigmas(req.diffusion.num_steps, + sigma_min=self.sigma_min, + sigma_max=self.sigma_max, + rho=self.rho) + latent_len = audio_latent_len(req, model, sample_size=self.sample_size, hop_length=self.hop_length) + shape = (self.io_channels, latent_len) # 1-D audio latent [64, L] + x = (rng.standard_normal(shape) * float(sig[0])).astype("float32") # noise·sigma_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={"audio": x}, + sigmas=[float(s) for s in sig], + timesteps=[float(s) for s in sig]) # RAW sigma is the model timestep (NOT *1000) + # The conditioning node packs the SA (cross_attn_cond, global_embed) triple into one numpy payload + # (see torch_stable_audio.pack_conditioning); the negative payload is the CFG uncond branch. + 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) + # A2A / RePaint (BRINGUP): inert for the T2A preset (None) -> the gated branches below are skipped. + st.scratch["init_latent"] = ctx.slots.get("init_latent") + st.scratch["inpaint_ref"] = ctx.slots.get("inpaint_reference_latent") + st.scratch["inpaint_mask"] = ctx.slots.get("inpaint_mask_latent") + st.scratch["prev_x0"] = None # DPM++(2M) multistep memory: the previous step's x0 estimate + 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] # raw sigma + expert_id = self.expert.expert_for(StepContext(i, t, sigma_t)) + x = st.latents["audio"] + pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] + scale = st.scratch["guidance_scale"] + precision = self.precision + do_cfg = scale != 1.0 and ne is not None + prev_x0 = st.scratch.get("prev_x0") + + def _x0(dit, text_embed) -> np.ndarray: + # VDenoiser preconditioning: feed the network c_in·x at c_noise(sigma), get the raw v output, + # reconstruct x0 = c_skip·x + c_out·v (EDM-v, sigma_data=1). Faithful to K.external.VDenoiser. + return vdenoiser_x0(x, float(sigma_t), + lambda scaled, c_noise: np.asarray(dit(scaled, text_embed, c_noise), dtype="float32")) + + def run(model, override=None): + dit = model.component(expert_id) + if override is not None and "x0" in override: + final_x0 = precision.cast(np.asarray(override["x0"], dtype="float32")) + else: + cond_x0 = _x0(dit, pe) + if do_cfg: + uncond_x0 = _x0(dit, ne) + final_x0 = uncond_x0 + (cond_x0 - uncond_x0) * scale # x0-space CFG (uncond + (cond-uncond)·s) + else: + final_x0 = cond_x0 + final_x0 = precision.cast(final_x0) + # Deterministic DPM-Solver++(2M) step over (x, x0, sigma_t -> sigma_next), using prev_x0 for the + # 2nd-order correction (1st-order Euler on the first step). The stochastic dpmpp-3m-sde terms + # are the GPU-parity refinement (BRINGUP); this deterministic spine runs on the CPU toy. + x_next = dpmpp_2m_step(x, final_x0, sigma_t, sigma_next, prev_x0) + return StepResult(output={"x0": np.asarray(final_x0, 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__), ("vpred", True))), + resources=res, + payload={ + "branch": "vpred", + "step": i + }, + run=run, + label=f"stable_audio.denoise.{i}", + capturable=False) # v->x0 + x0-space CFG + multistep host state -> eager path + + def advance(self, st: LoopState, result: StepResult) -> LoopState: + st.latents["audio"] = result.output["latents"] + st.scratch["prev_x0"] = np.asarray(result.output["x0"], dtype="float32") # DPM++(2M) memory + if st.profile == ExecutionProfile.ROLLOUT: + st.trajectory.append({ + "step": st.step_idx, + "sigma": st.sigmas[st.step_idx], + "x0": np.asarray(result.output["x0"]).copy(), + "latents": np.asarray(st.latents["audio"]).copy(), + }) + st.step_idx += 1 + return st + + def finalize(self, st: LoopState) -> LoopResult: + return LoopResult(outputs={"latents": st.latents["audio"]}, + metrics={"denoise_steps": float(st.step_idx)}, + behavior=st.trajectory or None) diff --git a/v2/recipes/stable_audio/program.py b/v2/recipes/stable_audio/program.py new file mode 100644 index 000000000..9b4b2a3ef --- /dev/null +++ b/v2/recipes/stable_audio/program.py @@ -0,0 +1,98 @@ +"""Stable Audio Open T2A program: conditioning -> stable_audio_denoise (v-prediction DPM++) -> audio_decode. + +Mirrors the cosmos2 t2v inline program shape — the SA specifics (the conditioner's (cross_attn_cond, +global_embed) packing; the v-prediction sampler) live entirely in the conditioning node, the +``StableAudioDenoiseLoop``, and the torch adapters, so the node graph stays a clean 3-node line. + +Deltas vs the video programs: + * The first node is ``conditioning`` (NOT the generic ``text_encode``): it runs the SA multi-conditioner + over (prompt, seconds_start, seconds_total) and packs the triple into the loop's flat ``text_embeds`` + payload (so the loop->dit call stays modality-agnostic). It degrades to the generic text encoder on + the CPU toy (which exposes only ``.encode(text)``), so the loop still runs end-to-end. + * The last node is ``audio_decode`` (OobleckVAE.decode + slice [start,end] seconds) writing an ``audio`` + artifact (NOT ``video``). BRINGUP: the v2 output-artifact surface is video-shaped (cards/programs + declare a ``video`` artifact); the registered ``output_artifacts`` here declares ``audio`` (+ latents), + which an audio-aware saver/output path must consume — that audio output node is the request-API + extension flagged in the spec (blocker #6). The slot value is the decoded waveform numpy array. + +Audio duration knobs: ``DiffusionParams`` has no ``audio_start_in_s``/``audio_end_in_s`` (the v2 request +surface is video-shaped — spec blocker #6). They are read from ``request.node_params['conditioning']`` +(the per-node override surface) with the card's defaults as fallback. A first-class audio request field is +the request-API extension (BRINGUP). +""" +from __future__ import annotations + +from typing import Any + +from v2.program import ComponentNode, ModelLoopNode, Program, ProgramKind + +# Default audio window (SA-1.0 short-clip defaults; overridable per request via node_params). +_DEFAULT_AUDIO_START_S = 0.0 +_DEFAULT_AUDIO_END_S = 10.0 + + +def _audio_window(request: Any) -> tuple[float, float]: + """Resolve (start_s, end_s) from the per-node override surface (no audio fields on DiffusionParams).""" + ov = request.node_override("conditioning") if hasattr(request, "node_override") else {} + start = float(ov.get("audio_start_in_s", _DEFAULT_AUDIO_START_S)) + end = float(ov.get("audio_end_in_s", _DEFAULT_AUDIO_END_S)) + return start, end + + +def _encode_one(component: Any, prompt: str, start: float, end: float) -> Any: + """Encode one prompt through the SA conditioner (``encode(prompt, seconds_start, seconds_total)``), + falling back to the generic text encoder's ``encode(text)`` on the CPU toy (so the loop still runs).""" + try: + return component.encode(prompt or "", start, end) # SA conditioner: packs (cross, global) + except TypeError: + return component.encode(prompt or "") # ToyTextEncoder / generic T5: single-arg + + +def _conditioning(instance, slots, request, ctx) -> None: + # The SA conditioner is declared as the ``text_encoder`` component (kind reuse — see card.py / the + # adapter docstring). On a GPU box this is the SA multi-conditioner; on CPU it is the ToyTextEncoder. + comp = instance.component("text_encoder") + start, end = _audio_window(request) + slots["text_embeds"] = _encode_one(comp, request.prompt(), start, end) + slots["neg_text_embeds"] = _encode_one(comp, request.diffusion.negative_prompt, start, end) + slots["audio_start_in_s"] = start + slots["audio_end_in_s"] = end + + +def _audio_decode(instance, slots, request, ctx) -> None: + # OobleckVAE.decode -> waveform [channels, samples]; slice to [start,end] seconds. The toy VAE returns + # its own (video-ish) decode but the slot is still named ``audio`` — the artifact is the waveform. + vae = instance.component("vae") + waveform = vae.decode(slots["denoise_out"]["latents"]) + sr = int(getattr(vae, "sampling_rate", 44100)) + start = float(slots.get("audio_start_in_s", _DEFAULT_AUDIO_START_S)) + end = float(slots.get("audio_end_in_s", _DEFAULT_AUDIO_END_S)) + if getattr(waveform, "ndim", 0) >= 1 and waveform.shape[-1] >= int(end * sr) > 0: + waveform = waveform[..., int(start * sr):int(end * sr)] # slice to the requested window + slots["audio"] = waveform + slots["audio_sample_rate"] = sr + + +def build_stable_audio_program() -> Program: + return Program( + program_id="stable_audio.t2a.inline", + kind=ProgramKind.INLINE, + nodes=[ + ComponentNode("conditioning", + fn=_conditioning, + writes=("text_embeds", "neg_text_embeds", "audio_start_in_s", "audio_end_in_s")), + ModelLoopNode("denoise", + loop_id="stable_audio_denoise", + output_slot="denoise_out", + reads=("text_embeds", "neg_text_embeds"), + writes=("denoise_out", )), + ComponentNode("audio_decode", + fn=_audio_decode, + reads=("denoise_out", "audio_start_in_s", "audio_end_in_s"), + writes=("audio", "audio_sample_rate")), + ], + output_artifacts={ + "audio": "audio", + "latents": "denoise_out", + }, + ).validate() diff --git a/v2/recipes/stable_audio/sampler.py b/v2/recipes/stable_audio/sampler.py new file mode 100644 index 000000000..a20561d60 --- /dev/null +++ b/v2/recipes/stable_audio/sampler.py @@ -0,0 +1,93 @@ +"""Stable Audio sampler helpers (recipe-local — NOT in v2/loop/sampler.py, which is flow-match/EDM-Karras). + +Three numpy primitives the ``StableAudioDenoiseLoop`` composes, faithful to k-diffusion's behavior for the +Stable Audio ``dpmpp-3m-sde`` sampler over a v-prediction (VDenoiser) network: + + * ``build_polyexponential_sigmas`` — k-diffusion ``get_sigmas_polyexponential(n, sigma_min, sigma_max, + rho)`` + the appended 0.0 terminal: a log-linear (rho=1) ramp from sigma_max to sigma_min. + * ``vdenoiser_x0`` — k-diffusion ``K.external.VDenoiser``: a v-prediction network's x0 reconstruction + with the EDM-v preconditioning (sigma_data=1). + * ``dpmpp_2m_step`` — one deterministic DPM-Solver++(2M) update in x0/sigma space (the deterministic + spine of dpmpp-3m-sde; the stochastic 3rd-order terms are the GPU-parity refinement, BRINGUP). + +The full stochastic ``sample_dpmpp_3m_sde`` is the GPU path (``import k_diffusion as K``); it is NOT +installed in the CPU env, so the loop runs this deterministic multistep recurrence for CPU-verification. +""" +from __future__ import annotations + +from collections.abc import Callable + +import numpy as np + + +def build_polyexponential_sigmas(num_steps: int, + *, + sigma_min: float = 0.3, + sigma_max: float = 500.0, + rho: float = 1.0) -> np.ndarray: + """k-diffusion ``get_sigmas_polyexponential`` + appended 0.0 terminal. + + ``sigmas_i = exp( lerp(log(sigma_max), log(sigma_min), (i/(n-1))**rho) )`` for i in 0..n-1, then a + trailing 0.0 (the sampler integrates pairwise; the last step lands on sigma=0). rho=1 makes this a + plain log-linear (geometric) ramp — Stable Audio's published schedule. + """ + n = max(1, int(num_steps)) + ramp = np.linspace(0.0, 1.0, n, dtype=np.float64)**float(rho) + log_sigmas = np.log(sigma_max) + ramp * (np.log(sigma_min) - np.log(sigma_max)) + sigmas = np.exp(log_sigmas) # sigma_max -> sigma_min + return np.concatenate([sigmas, np.zeros(1, dtype=np.float64)]).astype(np.float64) # + terminal 0 + + +def vdenoiser_x0(x: np.ndarray, sigma: float, net: Callable[[np.ndarray, float], np.ndarray]) -> np.ndarray: + """k-diffusion ``K.external.VDenoiser.forward``: reconstruct x0 from a v-prediction network. + + EDM-v preconditioning (sigma_data=1): + ``c_skip = 1/(sigma^2+1)``, ``c_out = -sigma/sqrt(sigma^2+1)``, ``c_in = 1/sqrt(sigma^2+1)``, + ``c_noise = atan(sigma)`` (k-diffusion's ``VDenoiser.sigma_to_t``), and + ``x0 = c_skip·x + c_out·F(c_in·x, c_noise)``. + ``net(scaled_x, c_noise)`` is the raw v-prediction (the torch adapter / CPU toy); this function owns + the preconditioning so the adapter stays the bare network. + """ + s = float(sigma) + s2p1 = s * s + 1.0 + c_skip = 1.0 / s2p1 + c_out = -s / np.sqrt(s2p1) + c_in = 1.0 / np.sqrt(s2p1) + c_noise = float(np.arctan(s)) # VDenoiser sigma_to_t + x = np.asarray(x, dtype="float32") + v = np.asarray(net((x * c_in).astype("float32"), c_noise), dtype="float32") + return (c_skip * x + c_out * v).astype("float32") + + +def dpmpp_2m_step(x: np.ndarray, denoised: np.ndarray, sigma_t: float, sigma_next: float, + prev_denoised: np.ndarray | None) -> np.ndarray: + """One deterministic DPM-Solver++(2M) update (k-diffusion ``sample_dpmpp_2m`` recurrence). + + Works in lambda = -log(sigma) (log-SNR) space: + ``t = -log(sigma_t)``, ``t_next = -log(sigma_next)``, ``h = t_next - t``. + First step (no history) is the 1st-order DPM-Solver++ (== exponential Euler on x0): + ``x_next = (sigma_next/sigma_t)·x - expm1(-h)·denoised``. + Subsequent steps add the 2nd-order correction from the previous x0 estimate: + ``denoised_d = (1 + 1/(2r))·denoised - 1/(2r)·prev_denoised`` with ``r = h/h_last`` (here + approximated with the standard ``r`` from the step ratio). At ``sigma_next == 0`` the update collapses + to ``x_next = denoised`` (the final clean sample), matching k-diffusion's terminal handling. + """ + x = np.asarray(x, dtype="float32") + denoised = np.asarray(denoised, dtype="float32") + sn = float(sigma_next) + st = float(sigma_t) + if sn <= 0.0: + return denoised.astype("float32") # terminal: land on the clean x0 + t = -np.log(max(st, 1e-12)) + t_next = -np.log(sn) + h = t_next - t + if prev_denoised is None: + x_next = (sn / st) * x - np.expm1(-h) * denoised + return x_next.astype("float32") + # 2nd-order multistep correction. Using r=1 (uniform-ish log-SNR spacing) keeps this faithful to + # k-diffusion's 2M form without threading the exact previous step size; the dominant error term is + # the same. (The exact h_last ratio is a GPU-parity refinement, BRINGUP.) + prev = np.asarray(prev_denoised, dtype="float32") + denoised_d = 1.5 * denoised - 0.5 * prev # (1 + 1/(2r))·d - 1/(2r)·prev, r=1 + x_next = (sn / st) * x - np.expm1(-h) * denoised_d + return x_next.astype("float32") diff --git a/v2/registry.py b/v2/registry.py index 6d20d990d..01c0b7935 100644 --- a/v2/registry.py +++ b/v2/registry.py @@ -33,6 +33,72 @@ class ModelEntry: build_program: Callable[[], Any] +# Bucket-C ports — each is a self-contained recipe package (card-declared torch adapter via +# ``ComponentSpec.adapter`` + a new/forked loop, no shared-backend edit). ONE table drives BOTH the +# explicit HF-id registry (PRIMARY) and the ``select_by_architecture`` fallback, so adding a port is one +# row. Rows: (hf_ids, package, card_builder, program_builder, transformer_cls). ``transformer_cls`` is the +# diffusers/EntryClass name for the fallback ("" = explicit-id-only; confirm against the real checkpoint's +# transformer/config.json _class_name). Builders are resolved lazily (``import v2.registry`` stays cheap / +# torch-free; the recipe packages are CPU-clean but importing 14 of them eagerly on every resolve is waste). +_BUCKET_C: tuple[tuple, ...] = ( + (("KyleShao/Cosmos-Predict2.5-2B-Diffusers", "nvidia/Cosmos-Predict2.5-14B"), "cosmos25", "build_cosmos25_card", + "build_cosmos25_program", "Cosmos25Transformer3DModel"), + (("hunyuanvideo-community/HunyuanVideo", ), "hunyuan_video", "build_hunyuan_video_card", + "build_hunyuan_video_program", "HunyuanVideoTransformer3DModel"), + (("FastVideo/FastHunyuan-diffusers", ), "hunyuan_video", "build_fast_hunyuan_video_card", + "build_hunyuan_video_program", ""), + (("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled"), "hunyuan_video15", + "build_hunyuan_video15_card", "build_hunyuan_video15_program", "HunyuanVideo15Transformer3DModel"), + (("hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v", + "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v_distilled", + "weizhou03/HunyuanVideo-1.5-Diffusers-1080p", "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR"), + "hunyuan_video15", "build_hunyuan_video15_720p_card", "build_hunyuan_video15_program", ""), + (("FastVideo/LongCat-Video-T2V-Diffusers", "FastVideo/LongCat-Video-I2V-Diffusers", + "FastVideo/LongCat-Video-VC-Diffusers"), "longcat", "build_longcat_card", "build_longcat_program", + "LongCatTransformer3DModel"), + (("stabilityai/stable-diffusion-3.5-medium", ), "sd35", "build_sd35_card", "build_sd35_program", + "SD3Transformer2DModel"), + (("FastVideo/GEN3C-Cosmos-7B-Diffusers", ), "gen3c", "build_gen3c_card", "build_gen3c_program", + "Gen3CTransformer3DModel"), + (("kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers", ), "kandinsky5", "build_kandinsky5_card", + "build_kandinsky5_program", "Kandinsky5Transformer3DModel"), + (("black-forest-labs/FLUX.2-dev", ), "flux2", "build_flux2_card", "build_flux2_program", "Flux2Transformer2DModel"), + (("black-forest-labs/FLUX.2-klein-4B", "black-forest-labs/FLUX.2-klein-9B"), "flux2", "build_flux2_klein_card", + "build_flux2_program", ""), + (("FastVideo/stable-audio-open-1.0-Diffusers", ), "stable_audio", "build_stable_audio_card", + "build_stable_audio_program", "StableAudioDiT"), + (("FastVideo/stable-audio-open-small-Diffusers", ), "stable_audio", "build_stable_audio_small_card", + "build_stable_audio_program", ""), + (("FastVideo/HunyuanGameCraft-Diffusers", ), "hunyuangamecraft", "build_hunyuangamecraft_card", + "build_hunyuangamecraft_program", "HunyuanGameCraftTransformer3DModel"), + (("FastVideo/HY-WorldPlay-Bidirectional-Diffusers", ), "hyworld", "build_hyworld_card", "build_hyworld_program", + "HYWorldTransformer3DModel"), + (("FastVideo/LingBot-World-Base-Cam-Diffusers", ), "lingbotworld", "build_lingbotworld_card", + "build_lingbotworld_program", "LingBotWorldTransformer3DModel"), + (("FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers", "FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers", + "FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers", "FastVideo/Matrix-Game-2.0-Base-Diffusers", + "FastVideo/Matrix-Game-2.0-GTA-Diffusers", "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers"), "matrixgame2", + "build_matrixgame2_card", "build_matrixgame2_program", "CausalMatrixGame2WanModel"), + (("FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers", ), "matrixgame3", "build_matrixgame3_card", + "build_matrixgame3_program", "MatrixGame3WanModel"), +) + + +def _lazy(package: str, fn: str) -> Callable[[], Any]: + """A ``() -> card_or_program`` builder that imports the recipe package only when first called.""" + + def _build() -> Any: + import importlib + return getattr(importlib.import_module(f"v2.recipes.{package}"), fn)() + + return _build + + +def _bucket_c_entries() -> list[ModelEntry]: + return [ModelEntry(hf_ids, _lazy(pkg, cb), _lazy(pkg, pb)) for (hf_ids, pkg, cb, pb, _cls) in _BUCKET_C] + + def _entries() -> list[ModelEntry]: """The explicit registry (PRIMARY). Builders imported lazily to avoid import cycles with the model packages and keep importing this module cheap.""" @@ -54,7 +120,7 @@ def _entries() -> list[ModelEntry]: ) 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 [ + return [ # noqa: RUF005 — explicit list, extended with the bucket-C ports below via _bucket_c_entries() 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). @@ -80,7 +146,7 @@ def _entries() -> list[ModelEntry]: # 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), - ] + ] + _bucket_c_entries() # the 14 self-contained bucket-C recipe packages (one row each in _BUCKET_C) def _short(p: str) -> str: @@ -118,6 +184,11 @@ 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") + for (_hf, pkg, cb, pb, cls) in _BUCKET_C: # the 14 bucket-C ports, keyed by transformer cls + if cls and tr == cls: + import importlib + mod = importlib.import_module(f"v2.recipes.{pkg}") + return getattr(mod, cb), getattr(mod, pb) if tr == "CosmosTransformer3DModel": from v2.recipes.cosmos2 import build_cosmos2_card, build_cosmos2_program return build_cosmos2_card, build_cosmos2_program