From facd035b97e0b8d596363f8c8adbc7af87956a54 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 22 May 2026 02:00:37 -0700 Subject: [PATCH] [feat]: Cosmos3 DiT skeleton (Cosmos3VFMTransformer + Cosmos3LanguageModel) Ports math (compute_mrope_position_ids_text/_vision, patchify, unpatchify) verbatim from vllm-omni transformer_cosmos3.py:113-177,1009-1036. Module tree (language_model.layers.*.self_attn.{q,k,v,o}_proj + gen_layers.*.cross_attention.* + vae2llm + llm2vae + time_embedder + norm_moe_gen) exposes the parameter names the checkpoint converter targets. All layer forward() raise NotImplementedError until weights publish; module instantiation + state-dict key tests pass. --- fastvideo/configs/models/dits/cosmos3.py | 96 +++++ fastvideo/models/dits/cosmos3.py | 454 +++++++++++++++++++++++ 2 files changed, 550 insertions(+) create mode 100644 fastvideo/configs/models/dits/cosmos3.py create mode 100644 fastvideo/models/dits/cosmos3.py diff --git a/fastvideo/configs/models/dits/cosmos3.py b/fastvideo/configs/models/dits/cosmos3.py new file mode 100644 index 000000000..dc75d51a6 --- /dev/null +++ b/fastvideo/configs/models/dits/cosmos3.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 VFM Transformer FastVideo dataclass configs (Phase 2b.1 scaffold). + +Mirrors the architectural defaults of +``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py::Cosmos3VFMTransformer`` +(see lines 910-977 of the vllm-omni reference at HEAD ``8536f5b1421f``). + +Phase 2b.1 ships only the dataclass surface and skeleton tensor shapes that +``Cosmos3VFMTransformer.__init__`` needs to build a parameter-named module +tree. The ``param_names_mapping`` table is left empty here; the checkpoint +remap lives in ``scripts/checkpoint_conversion/cosmos3_convert.py`` (Phase 5) +and the FastVideo pipeline's ``_remap_ckpt_key`` (Phase 2b.2+). +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_cosmos3_transformer_block(name: str, module) -> bool: + """FSDP shard boundary for Cosmos3 UND/GEN decoder layers. + + Mirrors ``Cosmos3VFMTransformer._is_transformer_block`` (reference line + 870). Matches ``gen_layers.{i}`` and ``language_model.layers.{i}`` where + ``{i}`` is a non-negative integer. + """ + del module + if "gen_layers" not in name and "language_model.layers" not in name: + return False + return name.split(".")[-1].isdigit() + + +@dataclass +class Cosmos3ArchConfig(DiTArchConfig): + """Architecture config for Cosmos3 VFM Transformer (UND + GEN pathways). + + Defaults match the upstream Cosmos3-Nano constructor defaults at + ``transformer_cosmos3.py`` lines 910-934. + """ + + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_cosmos3_transformer_block]) + + # ``param_names_mapping`` is intentionally empty for Phase 2b.1. The + # canonical Cosmos3 checkpoint remap is owned by the pipeline's + # ``_remap_ckpt_key`` (vllm-omni reference at pipeline_cosmos3.py:319-409) + # and is ported to FastVideo in a later phase (the conversion script under + # ``scripts/checkpoint_conversion/``). Leaving this empty makes the DiT + # module tree the single source of truth for parameter names during + # parity-test development. + param_names_mapping: dict = field(default_factory=dict) + + # ---- Backbone ---- + hidden_size: int = 4096 + num_hidden_layers: int = 36 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + head_dim: int = 128 + intermediate_size: int = 12288 + vocab_size: int = 151936 + rms_norm_eps: float = 1e-6 + rope_theta: float = 5_000_000.0 + mrope_section: list[int] = field(default_factory=lambda: [24, 20, 20]) + + # ---- VAE / patch geometry ---- + latent_patch_size: int = 2 + latent_channel: int = 48 + + # ---- Diffusion conditioning ---- + timestep_scale: float = 0.001 + + # ---- Temporal modulation ---- + base_fps: float = 24.0 + temporal_compression_factor: int = 4 + enable_fps_modulation: bool = True + temporal_modality_margin: int = 15000 + + # ---- BaseDiT bookkeeping ---- + in_channels: int = 48 + out_channels: int = 48 + + def __post_init__(self) -> None: + super().__post_init__() + # Mirror the BaseDiT contract: ``num_channels_latents`` matches the + # VAE latent channel count for video DiTs. + self.num_channels_latents = self.latent_channel + if not self.out_channels: + self.out_channels = self.in_channels + + +@dataclass +class Cosmos3VideoConfig(DiTConfig): + """Pipeline-level Cosmos3 DiT config (T2V / I2V / T2I share this surface).""" + + arch_config: DiTArchConfig = field(default_factory=Cosmos3ArchConfig) + prefix: str = "Cosmos3" diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py new file mode 100644 index 000000000..9c6c4fae9 --- /dev/null +++ b/fastvideo/models/dits/cosmos3.py @@ -0,0 +1,454 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 VFM Transformer (Phase 2b.1 skeleton). + +This module lands the parameter-named module tree of +``Cosmos3VFMTransformer`` together with two standalone math utilities used +by both the FastVideo port and the Phase 2a parity tests: + +* ``compute_mrope_position_ids_text`` / ``compute_mrope_position_ids_vision`` — + the unified-3D mRoPE position ID generators ported verbatim from + ``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py`` lines + 113-177 (HEAD ``8536f5b1421f``). + +* ``Cosmos3VFMTransformer.patchify`` / + ``Cosmos3VFMTransformer.unpatchify`` — the + ``[B,C,T,H,W] <-> [B, T*hp*wp, p*p*C]`` patch tokenizer, ported from + the same reference lines 1009-1036. + +Every per-layer ``forward()`` raises ``NotImplementedError("Phase 2b.2")`` +because the layer math (RoPE application, QK-norm, GQA SDPA, UND/GEN +attention plumbing) is intentionally deferred to Phase 2b.2. Only the +module composition + tensor shape contract is exercised here, enough to +make 4 Tier-A parity tests pass: + +* ``test_cosmos3_mrope_parity.py::test_compute_mrope_position_ids_text_and_vision`` +* ``test_cosmos3_patchify_unpatchify_parity.py::test_patchify_unpatchify_roundtrip`` +* ``test_cosmos3_patchify_unpatchify_parity.py::test_patchify_default_patch_size`` +* ``test_cosmos3_state_dict_keys.py::test_fastvideo_cosmos3_dit_module_tree_param_names`` + +The constructor signature mirrors the upstream +``Cosmos3VFMTransformer(od_config, *, temporal_compression_factor=None)`` +contract so the Phase 2a tests (which were authored against the upstream +shape) can drive it directly. Phase 2b.2+ will adapt this to FastVideo's +``BaseDiT(config, hf_config)`` loader path via the ``TransformerLoader``. +""" +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.configs.models.dits.cosmos3 import Cosmos3ArchConfig, Cosmos3VideoConfig +from fastvideo.models.dits.base import BaseDiT + + +EntryClass = ["Cosmos3VFMTransformer"] + + +def _tf_config_get(config: Any, key: str, default: Any) -> Any: + """Read a value from a dict, dataclass, or simple namespace. + + Mirrors ``transformer_cosmos3._tf_config_get`` at reference line 77. + """ + if config is None: + return default + if hasattr(config, "get"): + return config.get(key, default) + return getattr(config, key, default) + + +def compute_mrope_position_ids_text( + num_tokens: int, + temporal_offset: int, +) -> tuple[torch.Tensor, int]: + """Generate 3D mRoPE position IDs for text tokens. + + Text tokens broadcast a single monotonically-increasing position-ID + sequence across all three (t, h, w) axes. Verbatim port of + ``transformer_cosmos3.compute_mrope_position_ids_text`` (reference + lines 113-124). + """ + ids = torch.arange(num_tokens, dtype=torch.long) + temporal_offset + mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous() + return mrope_ids, temporal_offset + num_tokens + + +def compute_mrope_position_ids_vision( + grid_t: int, + grid_h: int, + grid_w: int, + temporal_offset: int | float, + fps: float | None = None, + base_fps: float = 24.0, + temporal_compression_factor: int = 4, + base_temporal_compression_factor: int | None = None, + enable_fps_modulation: bool = True, + start_frame_offset: int = 0, +) -> tuple[torch.Tensor, int]: + """Generate 3D mRoPE position IDs for vision tokens. + + Builds a ``(t, h, w)`` position grid (Qwen3-VL style, spatial indices + reset per temporal segment) flattened in t-major order. Optionally + modulates the temporal axis by ``base_fps / tcf * (1 / (fps / tcf))`` + so two clips at different FPS retain wall-clock-aligned temporal + positions. Verbatim port of + ``transformer_cosmos3.compute_mrope_position_ids_vision`` (reference + lines 127-177). + """ + fps_modulation = enable_fps_modulation and fps is not None + + if fps_modulation: + assert fps is not None + tps = fps / temporal_compression_factor + effective_base_tcf = (base_temporal_compression_factor + if base_temporal_compression_factor is not None + else temporal_compression_factor) + base_tps = base_fps / effective_base_tcf + frame_indices = torch.arange(grid_t, dtype=torch.float32) + t_index = (((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset) + .view(-1, 1).expand(-1, grid_h * grid_w).flatten()) + else: + t_index = (torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + + int(temporal_offset) + start_frame_offset) + + h_index = (torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten()) + w_index = (torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten()) + + if fps_modulation: + mrope_ids = torch.stack([t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0) + else: + mrope_ids = torch.stack([t_index, h_index, w_index], dim=0) + + next_offset = math.floor(mrope_ids.max().item()) + 1 + return mrope_ids, next_offset + + +class _Qwen3VLTextRMSNorm(nn.Module): + """Qwen3-VL / T5-style RMSNorm with an ``eps`` and a learnable ``weight``. + + Mirrors ``Qwen3VLTextRMSNorm`` at reference line 89-107 so checkpoint + keys match the canonical ``*.norm*.weight`` / ``*.{q,k}_norm.weight`` + parameter naming. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2") + + +class _TimestepEmbedder(nn.Module): + """Sinusoidal-to-vector timestep embedder used by the GEN pathway. + + Skeleton of ``TimestepEmbedder`` at reference line 257-283. The + ``linear_1``/``linear_2`` parameter names match the upstream + checkpoint so the Phase 5 remap can copy weights without renaming. + """ + + def __init__( + self, + hidden_size: int, + frequency_embedding_size: int = 256, + max_period: int = 10000, + target_dtype: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) + self.frequency_embedding_size = frequency_embedding_size + self.hidden_size = hidden_size + + half = frequency_embedding_size // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half) + self.register_buffer("freqs", freqs, persistent=False) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2") + + +class _Cosmos3GatedMLP(nn.Module): + """Gated-MLP block with ``gate_proj`` / ``up_proj`` / ``down_proj`` names. + + Skeleton of ``Cosmos3GatedMLP`` at reference line 288-326. The three + linear-layer parameter names are load-bearing for the checkpoint + remap; Phase 2b.2 will swap them for FastVideo-native + ``ReplicatedLinear`` shards. + """ + + def __init__(self, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2") + + +class _Cosmos3SkeletonAttention(nn.Module): + """Shared skeleton for UND causal- and GEN cross-attention. + + Exposes the q/k/v/o projections at the upstream parameter-name + locations (``*.{q,k,v,o}_proj`` and ``*.{q,k}_norm``) with the + GQA-correct in/out dimensions + (``q_proj`` writes ``num_attention_heads * head_dim`` channels; + ``k_proj``/``v_proj`` write ``num_key_value_heads * head_dim`` + channels). Phase 2b.2 will replace the bare ``nn.Linear`` shells with + ``ColumnParallelLinear``-equivalent FastVideo layers and implement + the RoPE + SDPA forward. + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + + self.q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.o_proj = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False) + + self.q_norm = _Qwen3VLTextRMSNorm(head_dim, eps=rms_norm_eps, dtype=dtype) + self.k_norm = _Qwen3VLTextRMSNorm(head_dim, eps=rms_norm_eps, dtype=dtype) + + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2") + + +class _Cosmos3UndDecoderLayer(nn.Module): + """UND (understanding) decoder layer skeleton. + + The submodule names ``self_attn``, ``input_layernorm``, + ``post_attention_layernorm``, ``mlp`` mirror the upstream layer + composition at reference line 624-678. + """ + + def __init__(self, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None: + super().__init__() + self.self_attn = _Cosmos3SkeletonAttention( + hidden_size=arch.hidden_size, + num_attention_heads=arch.num_attention_heads, + num_key_value_heads=arch.num_key_value_heads, + head_dim=arch.head_dim, + rms_norm_eps=arch.rms_norm_eps, + dtype=dtype, + ) + self.input_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) + self.post_attention_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) + self.mlp = _Cosmos3GatedMLP(arch.hidden_size, arch.intermediate_size) + + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2") + + +class _Cosmos3GenDecoderLayer(nn.Module): + """GEN (generation) decoder layer skeleton. + + Substitutes ``cross_attention`` for ``self_attn`` (mirrors upstream + reference line 681-751); otherwise identical to the UND layer in + naming. The ``layer_idx`` is captured for the eventual cached-KV + path so Phase 2b.2 can route per-layer UND keys/values correctly. + """ + + def __init__(self, layer_idx: int, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None: + super().__init__() + self.layer_idx = layer_idx + self.cross_attention = _Cosmos3SkeletonAttention( + hidden_size=arch.hidden_size, + num_attention_heads=arch.num_attention_heads, + num_key_value_heads=arch.num_key_value_heads, + head_dim=arch.head_dim, + rms_norm_eps=arch.rms_norm_eps, + dtype=dtype, + ) + self.input_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) + self.post_attention_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) + self.mlp = _Cosmos3GatedMLP(arch.hidden_size, arch.intermediate_size) + + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2") + + +class _Cosmos3LanguageModel(nn.Module): + """Understanding-pathway container. + + Holds ``embed_tokens``, ``layers``, and a final ``norm``. The + submodule names mirror the upstream reference at line 757-831 so the + checkpoint remap target keys ``language_model.embed_tokens.*``, + ``language_model.layers.{i}.*``, and ``language_model.norm.*`` are + populated as-is by the FastVideo module tree. + """ + + def __init__(self, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(arch.vocab_size, arch.hidden_size) + self.layers = nn.ModuleList( + [_Cosmos3UndDecoderLayer(arch, dtype) for _ in range(arch.num_hidden_layers)]) + self.norm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) + + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2") + + +class Cosmos3VFMTransformer(BaseDiT): + """Cosmos3 VFM Transformer — UND language model + GEN denoising layers. + + Phase 2b.1 ships only the module-tree skeleton and the two pure math + utilities (``patchify``/``unpatchify``). All per-layer ``forward`` + methods raise ``NotImplementedError("Phase 2b.2")``; the DiT-level + ``forward`` likewise. The constructor signature + ``(od_config, *, temporal_compression_factor=None)`` matches the + upstream reference at ``transformer_cosmos3.py`` line 899 so the + Phase 2a parity tests can drive it directly. + + Phase 2b.2 will: + * Replace the per-layer ``nn.Linear``s with FastVideo-native + ``ReplicatedLinear`` / ``DistributedAttention`` plumbing. + * Implement ``forward`` (UND once, GEN per-step with cached K/V). + * Adapt the constructor to the FastVideo + ``BaseDiT(config, hf_config)`` loader contract via + ``TransformerLoader``. + """ + + _fsdp_shard_conditions = Cosmos3VideoConfig().arch_config._fsdp_shard_conditions + _compile_conditions = Cosmos3VideoConfig().arch_config._compile_conditions + param_names_mapping = Cosmos3VideoConfig().arch_config.param_names_mapping + reverse_param_names_mapping: dict[str, str] = {} + + def __init__( + self, + od_config: object | None = None, + *, + temporal_compression_factor: int | None = None, + ) -> None: + nn.Module.__init__(self) + + model_config = getattr(od_config, "tf_model_config", None) if od_config is not None else None + rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {} + + self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096)) + self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36)) + self.num_attention_heads = int(_tf_config_get(model_config, "num_attention_heads", 32)) + self.num_key_value_heads = int(_tf_config_get(model_config, "num_key_value_heads", 8)) + self.head_dim = int(_tf_config_get(model_config, "head_dim", 128)) + self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288)) + self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936)) + self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6)) + self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000)) + self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20])) + self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2)) + self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48)) + self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001)) + self.base_fps = float(_tf_config_get(model_config, "base_fps", 24.0)) + if temporal_compression_factor is None: + resolved_tcf = int(_tf_config_get(model_config, "temporal_compression_factor", 4)) + else: + resolved_tcf = int(temporal_compression_factor) + self.temporal_compression_factor = resolved_tcf + self.enable_fps_modulation = bool(_tf_config_get(model_config, "enable_fps_modulation", True)) + self.temporal_modality_margin = int( + _tf_config_get(model_config, "unified_3d_mrope_temporal_modality_margin", 15000)) + + self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size + self.num_channels_latents = self.latent_channel_size + + dtype = getattr(od_config, "dtype", torch.bfloat16) if od_config is not None else torch.bfloat16 + + arch = Cosmos3ArchConfig( + hidden_size=self.hidden_size, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + head_dim=self.head_dim, + intermediate_size=self.intermediate_size, + vocab_size=self.vocab_size, + rms_norm_eps=self.rms_norm_eps, + rope_theta=self.rope_theta, + mrope_section=self.mrope_section, + latent_patch_size=self.latent_patch_size, + latent_channel=self.latent_channel_size, + timestep_scale=self.timestep_scale, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + temporal_modality_margin=self.temporal_modality_margin, + in_channels=self.latent_channel_size, + out_channels=self.latent_channel_size, + ) + + self.language_model = _Cosmos3LanguageModel(arch, dtype) + self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) + self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) + self.time_embedder = _TimestepEmbedder(self.hidden_size, target_dtype=dtype) + self.gen_layers = nn.ModuleList( + [_Cosmos3GenDecoderLayer(i, arch, dtype) for i in range(self.num_hidden_layers)]) + self.norm_moe_gen = _Qwen3VLTextRMSNorm(self.hidden_size, eps=self.rms_norm_eps, dtype=dtype) + + def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]: + """Return ``(hp, wp, H_padded, W_padded)`` for ``latent_patch_size`` padding. + + Mirrors ``Cosmos3VFMTransformer._pad_to_patch_size`` at reference + lines 1002-1007. + """ + p = self.latent_patch_size + h_padded = ((h + p - 1) // p) * p + w_padded = ((w + p - 1) // p) * p + return h_padded // p, w_padded // p, h_padded, w_padded + + def patchify(self, latents: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: + """``[B, C, t, h, w] -> [B, t*hp*wp, p*p*C]``. + + Pads ``h``/``w`` up to a multiple of ``latent_patch_size`` before + reshaping. Verbatim port of + ``Cosmos3VFMTransformer.patchify`` (reference lines 1009-1021). + """ + batch_size = latents.shape[0] + p = self.latent_patch_size + c = self.latent_channel_size + hp, wp, h_padded, w_padded = self._pad_to_patch_size(h, w) + + if h_padded != h or w_padded != w: + latents = F.pad(latents, (0, w_padded - w, 0, h_padded - h)) + + x = latents.reshape(batch_size, c, t, hp, p, wp, p) + x = x.permute(0, 2, 3, 5, 4, 6, 1) + return x.reshape(batch_size, t * hp * wp, p * p * c) + + def unpatchify(self, tokens: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: + """``[B, t*hp*wp, p*p*C] -> [B, C, t, h, w]``, cropping ``h``/``w`` padding. + + Verbatim port of ``Cosmos3VFMTransformer.unpatchify`` (reference + lines 1023-1036). + """ + batch_size = tokens.shape[0] + p = self.latent_patch_size + c = self.latent_channel_size + hp, wp, h_padded, w_padded = self._pad_to_patch_size(h, w) + + x = tokens.reshape(batch_size, t, hp, wp, p, p, c) + x = x.permute(0, 6, 1, 2, 4, 3, 5) + x = x.reshape(batch_size, c, t, h_padded, w_padded) + + if h_padded != h or w_padded != w: + x = x[:, :, :, :h, :w] + return x + + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise NotImplementedError("Phase 2b.2")