diff --git a/fastvideo/pipelines/basic/cosmos3/sequence_packing.py b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py new file mode 100644 index 000000000..16cd730f9 --- /dev/null +++ b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo-native Cosmos3 sequence packing (video subset). + +Numerical-parity port of the official ``cosmos_framework`` data packer +(``cosmos_framework.data.vfm.sequence_packing.pack_input_sequence``) restricted +to the VIDEO generation path that the FastVideo Cosmos3 DiT consumes (T2V / I2V +/ T2I). It builds, per sample, two splits: + +* a ``causal`` text split (prompt token ids, plus the trailing ``eos`` and + ``start_of_generation`` markers the framework appends when a generation + modality follows), and +* a ``full`` vision split (VAE latent patch tokens). + +The 3D-MRoPE position ids ``[3, seq]`` are produced exactly like the framework: +text tokens broadcast a single monotone id across the (t, h, w) axes, the +temporal offset is bumped by ``temporal_modality_margin`` at the text->vision +boundary, and vision tokens lay out a (T, H, W) grid with spatial ids reset per +segment. Condition frames (I2V cond frame 0, T2I single conditioned frame, ...) +are kept in the packed sequence and rope grid but excluded from the MSE-loss / +timestep bookkeeping, mirroring the framework. + +The output ``Cosmos3PackedSequence`` maps 1:1 onto the +``Cosmos3VFMTransformer.forward`` kwargs via :meth:`to_dit_kwargs`. This module +is pure torch/python; it imports no diffusers/transformers model classes. + +Reference of record: ``cosmos_framework`` (NVIDIA), the parity oracle used by +``tests/local_tests/cosmos3/test_cosmos3_packing_parity.py``. +""" +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Any + +import torch + +from fastvideo.models.dits.cosmos3 import ( + compute_mrope_position_ids_text, + compute_mrope_position_ids_vision, +) + +__all__ = [ + "Cosmos3VisionItem", + "Cosmos3SampleInputs", + "Cosmos3PackedSequence", + "pack_cosmos3_video_sequence", +] + + +# --------------------------------------------------------------------------- +# Inputs +# --------------------------------------------------------------------------- +@dataclass +class Cosmos3VisionItem: + """One vision latent for a sample. + + Args: + latent: VAE latent ``[C, T, H, W]`` (a leading batch axis of size 1 is + accepted and squeezed). + condition_frame_indexes: Latent-frame indices that are *conditioned* + (clean) rather than noisy. ``[]`` for T2V, ``[0]`` for I2V, and the + single conditioned frame for T2I. + fps: Frames-per-second for this clip; only used when + ``enable_fps_modulation`` is set. + """ + + latent: torch.Tensor + condition_frame_indexes: list[int] = field(default_factory=list) + fps: float | None = None + + +@dataclass +class Cosmos3SampleInputs: + """Per-sample packing inputs (one text prompt + its vision item).""" + + text_ids: list[int] + vision: Cosmos3VisionItem + timestep: float + + +# --------------------------------------------------------------------------- +# Output +# --------------------------------------------------------------------------- +@dataclass +class Cosmos3PackedSequence: + """Packed-sequence inputs consumed by ``Cosmos3VFMTransformer.forward``. + + Field names mirror the framework ``PackedSequence`` (+ its ``vision`` + ``ModalityData``) so the parity test can compare field-by-field. + """ + + # Sequence structure. + sample_lens: list[int] + split_lens: list[int] + attn_modes: list[str] + sequence_length: int + is_image_batch: bool + + # Text modality. + text_ids: torch.Tensor + text_indexes: torch.Tensor + position_ids: torch.Tensor # [3, sequence_length] + + # Vision modality. + vision_tokens: list[torch.Tensor] + vision_token_shapes: list[tuple[int, int, int]] + vision_sequence_indexes: torch.Tensor + vision_timesteps: torch.Tensor + vision_mse_loss_indexes: torch.Tensor + vision_noisy_frame_indexes: list[torch.Tensor] + vision_condition_mask: list[torch.Tensor] + fps_vision: torch.Tensor | None = None + + def to_dit_kwargs(self) -> dict[str, Any]: + """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``.""" + return dict( + text_ids=self.text_ids, + text_indexes=self.text_indexes, + position_ids=self.position_ids, + sequence_length=int(self.sequence_length), + split_lens=list(self.split_lens), + attn_modes=list(self.attn_modes), + vision_tokens=list(self.vision_tokens), + vision_token_shapes=list(self.vision_token_shapes), + vision_sequence_indexes=self.vision_sequence_indexes, + vision_timesteps=self.vision_timesteps, + vision_mse_loss_indexes=self.vision_mse_loss_indexes, + vision_noisy_frame_indexes=list(self.vision_noisy_frame_indexes), + fps_vision=self.fps_vision, + ) + + +# --------------------------------------------------------------------------- +# Packing +# --------------------------------------------------------------------------- +def pack_cosmos3_video_sequence( + samples: list[Cosmos3SampleInputs], + special_tokens: dict[str, int], + *, + latent_patch_size: int = 2, + include_end_of_generation_token: bool = False, + temporal_modality_margin: int = 15_000, + reset_spatial_ids: bool = True, + enable_fps_modulation: bool = False, + base_fps: float = 24.0, + temporal_compression_factor: int = 4, + initial_mrope_temporal_offset: int | float = 0, +) -> Cosmos3PackedSequence: + """Pack prompts + vision latents into the Cosmos3 DiT packed-sequence inputs. + + Video subset of ``cosmos_framework`` ``pack_input_sequence`` under + ``unified_3d_mrope``: each sample is ``[causal text, full vision]``. + + Args: + samples: Per-sample text prompt token ids + vision item + timestep. + special_tokens: Must contain ``eos_token_id`` and + ``start_of_generation`` (and ``end_of_generation`` if + ``include_end_of_generation_token``). ``bos_token_id`` is honored if + present (prepended) to match the framework. + latent_patch_size: Latent patch size used by the DiT. + include_end_of_generation_token: Append the framework's end-of-generation + marker after the vision split. + temporal_modality_margin: Temporal-offset bump applied at the + text->vision boundary (``unified_3d_mrope_temporal_modality_margin``). + reset_spatial_ids: Reset vision spatial ids to 0 per segment. + enable_fps_modulation: Use float, fps-scaled temporal positions. + base_fps: Base FPS used when ``enable_fps_modulation``. + temporal_compression_factor: VAE temporal compression factor. + initial_mrope_temporal_offset: Per-sample starting temporal offset. + + Returns: + A :class:`Cosmos3PackedSequence`. + """ + assert "eos_token_id" in special_tokens, "special_tokens must contain eos_token_id" + assert "start_of_generation" in special_tokens, "special_tokens must contain start_of_generation" + if latent_patch_size < 1: + raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}") + + # Build-time accumulators (concatenated across samples). + sample_lens: list[int] = [] + split_lens: list[int] = [] + attn_modes: list[str] = [] + + text_ids: list[int] = [] + text_indexes: list[int] = [] + position_id_blocks: list[torch.Tensor] = [] # each [3, n] + + vision_tokens: list[torch.Tensor] = [] + vision_token_shapes: list[tuple[int, int, int]] = [] + vision_sequence_indexes: list[int] = [] + vision_timesteps: list[float] = [] + vision_mse_loss_indexes: list[int] = [] + vision_noisy_frame_indexes: list[torch.Tensor] = [] + vision_condition_mask: list[torch.Tensor] = [] + fps_values: list[float] = [] + + curr = 0 # running position in the packed sequence + is_image_batch = True + + for sample in samples: + temporal_offset: int | float = initial_mrope_temporal_offset + sample_len = 0 + + # ---- 1. Text split (causal) ---- + if "bos_token_id" in special_tokens: + shifted_text_ids = [special_tokens["bos_token_id"], *sample.text_ids] + else: + shifted_text_ids = list(sample.text_ids) + # The video path always has a following generation modality, so the + # framework appends eos + start_of_generation. + shifted_text_ids = [*shifted_text_ids, special_tokens["eos_token_id"], special_tokens["start_of_generation"]] + text_split_len = len(shifted_text_ids) + + text_ids.extend(shifted_text_ids) + text_indexes.extend(range(curr, curr + text_split_len)) + + text_mrope, temporal_offset = compute_mrope_position_ids_text( + num_tokens=text_split_len, + temporal_offset=int(temporal_offset), + ) + position_id_blocks.append(text_mrope) + + attn_modes.append("causal") + split_lens.append(text_split_len) + curr += text_split_len + sample_len += text_split_len + + # End of text modality: bump temporal offset before vision. + temporal_offset += temporal_modality_margin + + # ---- 2. Vision split (full) ---- + latent = sample.vision.latent + latent = latent.squeeze(0) if latent.dim() == 5 else latent # [C, T, H, W] + _c, latent_t, latent_h, latent_w = latent.shape + patch_h = math.ceil(latent_h / latent_patch_size) + patch_w = math.ceil(latent_w / latent_patch_size) + num_vision_tokens = latent_t * patch_h * patch_w + + vision_tokens.append(sample.vision.latent) + vision_token_shapes.append((latent_t, patch_h, patch_w)) + vision_sequence_indexes.extend(range(curr, curr + num_vision_tokens)) + + condition_set = {idx for idx in sample.vision.condition_frame_indexes if 0 <= idx < latent_t} + cond_mask = torch.zeros((latent_t, 1, 1), device=latent.device, dtype=latent.dtype) + for frame_idx in condition_set: + cond_mask[frame_idx, 0, 0] = 1.0 + vision_condition_mask.append(cond_mask) + + noisy_frames = torch.tensor( + [idx for idx in range(latent_t) if idx not in condition_set], + device=latent.device, + dtype=torch.long, + ) + vision_noisy_frame_indexes.append(noisy_frames) + + # MSE-loss indices + per-token timesteps cover only the noisy frames. + frame_token_stride = patch_h * patch_w + for frame_idx in range(latent_t): + if frame_idx in condition_set: + continue + frame_start = curr + frame_idx * frame_token_stride + vision_mse_loss_indexes.extend(range(frame_start, frame_start + frame_token_stride)) + vision_timesteps.extend([float(sample.timestep)] * frame_token_stride) + + vision_fps = sample.vision.fps if enable_fps_modulation else None + if vision_fps is not None: + fps_values.append(float(vision_fps)) + vision_mrope, temporal_offset = compute_mrope_position_ids_vision( + grid_t=latent_t, + grid_h=patch_h, + grid_w=patch_w, + temporal_offset=temporal_offset, + fps=vision_fps, + base_fps=base_fps, + temporal_compression_factor=temporal_compression_factor, + enable_fps_modulation=enable_fps_modulation, + ) + position_id_blocks.append(vision_mrope) + + curr += num_vision_tokens + sample_len += num_vision_tokens + + # ---- 3. Optional end-of-generation marker ---- + eov_len = 0 + if include_end_of_generation_token: + assert "end_of_generation" in special_tokens, ("special_tokens must contain end_of_generation when " + "include_end_of_generation_token=True") + text_ids.append(special_tokens["end_of_generation"]) + text_indexes.append(curr) + eov_dtype = torch.float32 if enable_fps_modulation else torch.long + eov_ids = torch.full((3, 1), temporal_offset, dtype=eov_dtype) + position_id_blocks.append(eov_ids) + temporal_offset += 1 + curr += 1 + eov_len = 1 + sample_len += 1 + + # The vision split and any trailing eov marker share the "full" split. + attn_modes.append("full") + split_lens.append(num_vision_tokens + eov_len) + sample_lens.append(sample_len) + + if latent_t != 1: + is_image_batch = False + + sequence_length = sum(sample_lens) + + # position_ids: float iff any block is float (fps modulation path). + any_float = any(b.dtype.is_floating_point for b in position_id_blocks) + if any_float: + position_id_blocks = [b.to(torch.float32) for b in position_id_blocks] + position_ids = torch.cat(position_id_blocks, dim=1) # [3, sequence_length] + + timesteps_dtype = torch.float32 + return Cosmos3PackedSequence( + sample_lens=sample_lens, + split_lens=split_lens, + attn_modes=attn_modes, + sequence_length=sequence_length, + is_image_batch=is_image_batch, + text_ids=torch.tensor(text_ids, dtype=torch.long), + text_indexes=torch.tensor(text_indexes, dtype=torch.long), + position_ids=position_ids, + vision_tokens=vision_tokens, + vision_token_shapes=vision_token_shapes, + vision_sequence_indexes=torch.tensor(vision_sequence_indexes, dtype=torch.long), + vision_timesteps=torch.tensor(vision_timesteps, dtype=timesteps_dtype), + vision_mse_loss_indexes=torch.tensor(vision_mse_loss_indexes, dtype=torch.long), + vision_noisy_frame_indexes=vision_noisy_frame_indexes, + vision_condition_mask=vision_condition_mask, + fps_vision=(torch.tensor(fps_values, dtype=torch.float32) if fps_values else None), + ) diff --git a/tests/local_tests/cosmos3/test_cosmos3_packing_parity.py b/tests/local_tests/cosmos3/test_cosmos3_packing_parity.py new file mode 100644 index 000000000..37fd3dc07 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_packing_parity.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo Cosmos3 sequence packing vs the OFFICIAL framework. + +FastVideo's native packer +(``fastvideo.pipelines.basic.cosmos3.sequence_packing.pack_cosmos3_video_sequence``) +builds the packed-sequence inputs the ``Cosmos3VFMTransformer`` consumes. This +test asserts, for the SAME logical inputs (prompt token ids, vision latents, +condition-frame indices, diffusion timestep, fps), that FastVideo's packing +matches the official ``cosmos_framework.data.vfm.sequence_packing.pack_input_sequence`` +oracle field-by-field: + + * ``position_ids`` (exact, ``[3, seq]``), + * ``text_ids`` / ``text_indexes``, + * ``split_lens`` / ``attn_modes`` / ``sample_lens`` / ``sequence_length``, + * vision ``sequence_indexes`` / ``token_shapes`` / ``timesteps`` / + ``mse_loss_indexes`` / ``noisy_frame_indexes`` / ``condition_mask``. + +Coverage spans T2V (no condition frames), I2V (condition frame 0), and T2I +(single conditioned frame), across multiple grids, plus a multi-sample batch. + +Then BOTH the framework-packed and FastVideo-packed inputs are fed through the +SAME tiny FastVideo DiT (framework weights copied in as in the existing DiT +parity tests). Asserting bit-identical DiT output confirms FastVideo's own +packing drives the DiT to the same result as the framework's packing. + +The official framework is the parity ORACLE; it runs on CPU / float32 via the +SDPA monkey-patch from ``test_cosmos3_reference_forward``. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_packing_parity.py -q +""" +from __future__ import annotations + +import pytest +import torch + +# The official framework provides the parity oracle. +cosmos_framework = pytest.importorskip( + "cosmos_framework", + reason="cosmos_framework not installed; run in fv-cosmos3 env.", +) + +# Reuse the DiT parity helpers (weight copy + framework->DiT kwarg builder) and +# the mRoPE tiny-model builders (real-checkpoint rope constants). +from .test_cosmos3_dit_parity import ( # noqa: E402 + _copy_weights, + _fastvideo_inputs_from_packed_seq, +) +from .test_cosmos3_dit_parity_mrope import ( # noqa: E402 + _LATENT_CHANNEL, + _LATENT_PATCH_SIZE, + _MROPE_SECTION, + _RESET_SPATIAL_IDS, + _ROPE_THETA, + _TCF, + _TEMPORAL_MODALITY_MARGIN, + _build_tiny_cosmos3_mrope, + _build_tiny_fastvideo_dit_mrope, +) +from .test_cosmos3_reference_forward import _apply_sdpa_patches # noqa: E402 + +pytestmark = [pytest.mark.local] + +# Ensure the CPU/float32 SDPA patches are installed (idempotent). +_apply_sdpa_patches() + +# Tiny special-token ids (kept < tiny vocab_size=64). The video path appends +# eos + start_of_generation after the prompt tokens. +_SPECIAL_TOKENS = { + "start_of_generation": 60, + "end_of_generation": 61, + "eos_token_id": 62, +} + + +# --------------------------------------------------------------------------- +# Builders for the two packers from the SAME logical sample inputs. +# --------------------------------------------------------------------------- +def _make_vision(grid_t: int, latent_h: int, latent_w: int, seed: int) -> torch.Tensor: + """Deterministic VAE latent ``[1, C, T, H, W]``.""" + torch.manual_seed(seed) + return torch.randn(1, _LATENT_CHANNEL, grid_t, latent_h, latent_w) + + +def _framework_pack( + *, + text_ids_per_sample: list[list[int]], + visions: list[torch.Tensor], + cond_frames_per_sample: list[list[int]], + timesteps: list[float], + is_image_batch: bool, +): + from cosmos_framework.data.vfm.sequence_packing import ( + GenerationDataClean, + SequencePlan, + pack_input_sequence, + ) + + gen_data_clean = GenerationDataClean( + batch_size=len(visions), + is_image_batch=is_image_batch, + x0_tokens_vision=list(visions), + fps_vision=None, + num_vision_items_per_sample=[1] * len(visions), + ) + plans = [ + SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=list(cf)) + for cf in cond_frames_per_sample + ] + return pack_input_sequence( + sequence_plans=plans, + input_text_indexes=[list(t) for t in text_ids_per_sample], + gen_data_clean=gen_data_clean, + input_timesteps=torch.tensor(timesteps, dtype=torch.float32), + special_tokens=_SPECIAL_TOKENS, + latent_patch_size=_LATENT_PATCH_SIZE, + include_end_of_generation_token=False, + position_embedding_type="unified_3d_mrope", + unified_3d_mrope_reset_spatial_ids=_RESET_SPATIAL_IDS, + unified_3d_mrope_temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TCF, + ) + + +def _fastvideo_pack( + *, + text_ids_per_sample: list[list[int]], + visions: list[torch.Tensor], + cond_frames_per_sample: list[list[int]], + timesteps: list[float], +): + from fastvideo.pipelines.basic.cosmos3.sequence_packing import ( + Cosmos3SampleInputs, + Cosmos3VisionItem, + pack_cosmos3_video_sequence, + ) + + samples = [ + Cosmos3SampleInputs( + text_ids=list(t), + vision=Cosmos3VisionItem(latent=v, condition_frame_indexes=list(cf)), + timestep=float(ts), + ) + for t, v, cf, ts in zip(text_ids_per_sample, visions, cond_frames_per_sample, timesteps) + ] + return pack_cosmos3_video_sequence( + samples, + _SPECIAL_TOKENS, + latent_patch_size=_LATENT_PATCH_SIZE, + include_end_of_generation_token=False, + temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, + reset_spatial_ids=_RESET_SPATIAL_IDS, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TCF, + ) + + +# --------------------------------------------------------------------------- +# Field-by-field comparison. +# --------------------------------------------------------------------------- +def _assert_packs_match(fw, fv) -> None: + """Assert the framework PackedSequence and FastVideo pack agree field-by-field.""" + # Structure. + assert fv.split_lens == list(fw.split_lens), f"split_lens: fv={fv.split_lens} fw={list(fw.split_lens)}" + assert fv.attn_modes == list(fw.attn_modes), f"attn_modes: fv={fv.attn_modes} fw={list(fw.attn_modes)}" + assert fv.sample_lens == list(fw.sample_lens), f"sample_lens: fv={fv.sample_lens} fw={list(fw.sample_lens)}" + assert int(fv.sequence_length) == int(fw.sequence_length) + + # Text. + torch.testing.assert_close(fv.text_ids, fw.text_ids.to(torch.long), rtol=0, atol=0) + torch.testing.assert_close(fv.text_indexes, fw.text_indexes.to(torch.long), rtol=0, atol=0) + + # position_ids: exact, [3, seq], same dtype. + assert fv.position_ids.shape == fw.position_ids.shape, ( + f"position_ids shape: fv={tuple(fv.position_ids.shape)} fw={tuple(fw.position_ids.shape)}") + assert fv.position_ids.dtype == fw.position_ids.dtype, ( + f"position_ids dtype: fv={fv.position_ids.dtype} fw={fw.position_ids.dtype}") + torch.testing.assert_close(fv.position_ids, fw.position_ids, rtol=0, atol=0) + + # Vision. + fwv = fw.vision + torch.testing.assert_close(fv.vision_sequence_indexes, fwv.sequence_indexes.to(torch.long), rtol=0, atol=0) + assert fv.vision_token_shapes == [tuple(s) for s in fwv.token_shapes], ( + f"token_shapes: fv={fv.vision_token_shapes} fw={[tuple(s) for s in fwv.token_shapes]}") + torch.testing.assert_close(fv.vision_timesteps.to(torch.float32), fwv.timesteps.to(torch.float32)) + torch.testing.assert_close(fv.vision_mse_loss_indexes, fwv.mse_loss_indexes.to(torch.long), rtol=0, atol=0) + assert len(fv.vision_noisy_frame_indexes) == len(fwv.noisy_frame_indexes) + for a, b in zip(fv.vision_noisy_frame_indexes, fwv.noisy_frame_indexes): + torch.testing.assert_close(a.to(torch.long), b.to(torch.long), rtol=0, atol=0) + assert len(fv.vision_condition_mask) == len(fwv.condition_mask) + for a, b in zip(fv.vision_condition_mask, fwv.condition_mask): + torch.testing.assert_close(a.flatten().to(torch.float32), b.flatten().to(torch.float32)) + + +# (grid_t, latent_h, latent_w, n_text, cond_frames, id) — single-sample cases. +_CASES = [ + pytest.param(1, 8, 8, 4, [], id="t2i_1x4x4"), + pytest.param(1, 4, 4, 5, [0], id="t2i_cond_1x2x2"), + pytest.param(2, 4, 4, 4, [], id="t2v_2x2x2"), + pytest.param(3, 8, 4, 6, [], id="t2v_3x4x2"), + pytest.param(2, 4, 4, 5, [0], id="i2v_2x2x2"), + pytest.param(3, 4, 4, 4, [0], id="i2v_3x2x2"), +] + + +class TestCosmos3PackingParity: + + # -- Field-by-field packing parity ------------------------------------- + @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w", "n_text", "cond"), _CASES) + def test_packing_fields_match_framework(self, grid_t, latent_h, latent_w, n_text, cond): + torch.manual_seed(0) + text_ids = torch.randint(0, 60, (n_text,)).tolist() + vision = _make_vision(grid_t, latent_h, latent_w, seed=123) + timestep = 500.0 + + fw = _framework_pack( + text_ids_per_sample=[text_ids], + visions=[vision], + cond_frames_per_sample=[cond], + timesteps=[timestep], + is_image_batch=(grid_t == 1), + ) + fv = _fastvideo_pack( + text_ids_per_sample=[text_ids], + visions=[vision], + cond_frames_per_sample=[cond], + timesteps=[timestep], + ) + _assert_packs_match(fw, fv) + + def test_packing_fields_match_framework_multi_sample(self): + """A batch of two samples (T2V + I2V) packs identically to the framework.""" + torch.manual_seed(1) + t0 = torch.randint(0, 60, (3,)).tolist() + t1 = torch.randint(0, 60, (5,)).tolist() + v0 = _make_vision(2, 4, 4, seed=11) + v1 = _make_vision(2, 4, 4, seed=22) + kwargs = dict( + text_ids_per_sample=[t0, t1], + visions=[v0, v1], + cond_frames_per_sample=[[], [0]], + timesteps=[500.0, 250.0], + ) + fw = _framework_pack(is_image_batch=False, **kwargs) + fv = _fastvideo_pack(**kwargs) + _assert_packs_match(fw, fv) + + # -- End-to-end: FastVideo packing drives the DiT identically ---------- + @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w", "n_text", "cond"), _CASES) + def test_fastvideo_packing_drives_dit_like_framework(self, grid_t, latent_h, latent_w, n_text, cond): + """Feed BOTH the framework-packed and FastVideo-packed inputs through the + SAME FastVideo DiT (framework weights copied in); assert identical output. + """ + num_layers = 2 + torch.manual_seed(0) + text_ids = torch.randint(0, 60, (n_text,)).tolist() + vision = _make_vision(grid_t, latent_h, latent_w, seed=123) + timestep = 500.0 + + fw_pack = _framework_pack( + text_ids_per_sample=[text_ids], + visions=[vision], + cond_frames_per_sample=[cond], + timesteps=[timestep], + is_image_batch=(grid_t == 1), + ) + fv_pack = _fastvideo_pack( + text_ids_per_sample=[text_ids], + visions=[vision], + cond_frames_per_sample=[cond], + timesteps=[timestep], + ) + # Guard: the two packs must agree before we trust the DiT comparison. + _assert_packs_match(fw_pack, fv_pack) + + # One DiT instance, framework weights copied in (parity oracle weights). + vfm = _build_tiny_cosmos3_mrope(seed=42, num_layers=num_layers) + dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers) + _copy_weights(vfm, dit) + + with torch.no_grad(): + out_fw = dit(**_fastvideo_inputs_from_packed_seq(fw_pack)) + out_fv = dit(**fv_pack.to_dit_kwargs()) + + # last_hidden_state must be bit-identical. + lhs_fw = out_fw["last_hidden_state"] + lhs_fv = out_fv["last_hidden_state"] + assert lhs_fw.shape == lhs_fv.shape + max_abs_lhs = (lhs_fw - lhs_fv).abs().max().item() + print(f"\n[packing->dit {grid_t}x{latent_h}x{latent_w} cond={cond}] " + f"last_hidden_state max abs diff = {max_abs_lhs:.3e}") + torch.testing.assert_close(lhs_fv, lhs_fw, rtol=0, atol=0) + + # preds_vision must be bit-identical when there are noisy frames to + # predict. (A fully-conditioned clip has no noisy patches, so the DiT + # emits no "preds_vision" — both packs agree the mse-loss set is empty, + # already asserted by the field-parity guard above.) + has_preds = fv_pack.vision_mse_loss_indexes.numel() > 0 + assert ("preds_vision" in out_fw) == has_preds + assert ("preds_vision" in out_fv) == has_preds + if has_preds: + pv_fw = out_fw["preds_vision"][0] + pv_fv = out_fv["preds_vision"][0] + assert pv_fw.shape == pv_fv.shape + max_abs_pv = (pv_fw - pv_fv).abs().max().item() + print(f"[packing->dit {grid_t}x{latent_h}x{latent_w} cond={cond}] " + f"preds_vision max abs diff = {max_abs_pv:.3e}") + torch.testing.assert_close(pv_fv, pv_fw, rtol=0, atol=0)