diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py index 2a0094dd2..40151fe3f 100644 --- a/fastvideo/models/dits/cosmos3.py +++ b/fastvideo/models/dits/cosmos3.py @@ -801,9 +801,17 @@ class Cosmos3VFMTransformer(BaseDiT): vision_mse_loss_indexes: torch.Tensor, vision_noisy_frame_indexes: list[torch.Tensor], fps_vision: torch.Tensor | None = None, + sound_tokens: list[torch.Tensor] | None = None, + sound_token_shapes: list[tuple[int, int, int]] | None = None, + sound_sequence_indexes: torch.Tensor | None = None, + sound_timesteps: torch.Tensor | None = None, + sound_mse_loss_indexes: torch.Tensor | None = None, + sound_noisy_frame_indexes: list[torch.Tensor] | None = None, + fps_sound: torch.Tensor | None = None, **kwargs: Any, ) -> dict[str, Any]: - """Video-path forward returning ``{"last_hidden_state", "preds_vision"}``. + """Video-path forward returning ``{"last_hidden_state", "preds_vision"}`` + (plus ``"preds_sound"`` when sound tokens are provided for t2vs). Inputs mirror the framework ``PackedSequence`` fields for a single text(causal)+vision(full) sample. This is the surface the parity test @@ -838,6 +846,23 @@ class Cosmos3VFMTransformer(BaseDiT): packed_sequence[vision_sequence_indexes] = packed_vision + # 2b. Sound (t2vs): pack [C, T] -> audio_proj_in -> + modality embed -> + # timestep scatter -> scatter into the (full) gen split. Mirrors the + # framework ``_encode_sound``; sound shares the vision "full" split. + if sound_tokens: + packed_sound = torch.cat( + [s[:, :shp[0]].permute(1, 0) for s, shp in zip(sound_tokens, sound_token_shapes)], + dim=0, + ) # [total_sound_tokens, sound_dim] + packed_sound, _ = self.audio_proj_in(packed_sound.to(target_dtype)) + packed_sound = packed_sound + self.audio_modality_embed + if sound_mse_loss_indexes is not None and sound_mse_loss_indexes.numel() > 0: + s_ts = sound_timesteps * self.timestep_scale + s_embeds = self.time_embedder(s_ts).to(target_dtype) + packed_sound = self._scatter_timestep_embeds(packed_sound, s_embeds, sound_noisy_frame_indexes, + sound_token_shapes) + packed_sequence[sound_sequence_indexes] = packed_sound + # 3. RoPE for the full packed sequence (cos/sin per token); split und/gen cos, sin = self.rotary_emb(position_ids, device=device, dtype=target_dtype) # [N, head_dim] und_idx, gen_idx = self._mode_indices(split_lens, attn_modes, device) @@ -864,9 +889,35 @@ class Cosmos3VFMTransformer(BaseDiT): output["preds_vision"] = self._unpatchify_and_unpack(preds, vision_token_shapes, vision_noisy_frame_indexes, original_shapes) + # 6b. Sound prediction (t2vs): audio_proj_out on noisy sound hidden + # states -> unpack to per-sample [C, T] (framework ``_decode_sound``). + if sound_tokens and sound_mse_loss_indexes is not None and sound_mse_loss_indexes.numel() > 0: + preds_sound, _ = self.audio_proj_out(last_hidden_state[sound_mse_loss_indexes]) + output["preds_sound"] = self._unpack_sound(preds_sound, sound_token_shapes, sound_noisy_frame_indexes) + output["last_hidden_state"] = last_hidden_state return output + def _unpack_sound( + self, + packed_preds: torch.Tensor, + token_shapes: list[tuple[int, int, int]], + noisy_frame_indexes: list[torch.Tensor], + ) -> list[torch.Tensor]: + """Scatter packed noisy sound preds back to per-sample ``[C, T]`` (clean + frames left zero). Mirrors framework ``unpack_sound_latents``.""" + outputs: list[torch.Tensor] = [] + start_idx = 0 + for shape, nfi in zip(token_shapes, noisy_frame_indexes): + t = shape[0] + out = torch.zeros((self.sound_dim, t), device=packed_preds.device, dtype=packed_preds.dtype) + t_n = len(nfi) + if t_n > 0: + out[:, nfi] = packed_preds[start_idx:start_idx + t_n].T + start_idx += t_n + outputs.append(out) + return outputs + @staticmethod def _mode_indices(split_lens: list[int], attn_modes: list[str], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/fastvideo/pipelines/basic/cosmos3/sequence_packing.py b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py index 9be22f7c8..fb0c42918 100644 --- a/fastvideo/pipelines/basic/cosmos3/sequence_packing.py +++ b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py @@ -69,13 +69,31 @@ class Cosmos3VisionItem: fps: float | None = None +@dataclass +class Cosmos3SoundItem: + """One sound latent for a sample (t2vs). + + Args: + latent: AVAE sound latent ``[C, T]`` (channels, temporal frames). + condition_frame_indexes: Latent-frame indices that are *conditioned* + (clean). ``[]`` for t2vs (all frames generated). + fps: Sound latent FPS (``sound_latent_fps``, e.g. 25); 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).""" + """Per-sample packing inputs (one text prompt + its vision item, +sound).""" text_ids: list[int] vision: Cosmos3VisionItem timestep: float + sound: Cosmos3SoundItem | None = None # --------------------------------------------------------------------------- @@ -111,6 +129,16 @@ class Cosmos3PackedSequence: vision_condition_mask: list[torch.Tensor] fps_vision: torch.Tensor | None = None + # Sound modality (t2vs); empty/None when no sound. + sound_tokens: list[torch.Tensor] = field(default_factory=list) + sound_token_shapes: list[tuple[int, int, int]] = field(default_factory=list) + sound_sequence_indexes: torch.Tensor | None = None + sound_timesteps: torch.Tensor | None = None + sound_mse_loss_indexes: torch.Tensor | None = None + sound_noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list) + sound_condition_mask: list[torch.Tensor] = field(default_factory=list) + fps_sound: torch.Tensor | None = None + def to_dit_kwargs(self, device: torch.device | str | None = None) -> dict[str, Any]: """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``. @@ -136,6 +164,13 @@ class Cosmos3PackedSequence: vision_mse_loss_indexes=_mv(self.vision_mse_loss_indexes), vision_noisy_frame_indexes=[_mv(t) for t in self.vision_noisy_frame_indexes], fps_vision=self.fps_vision, + sound_tokens=[_mv(t) for t in self.sound_tokens], + sound_token_shapes=list(self.sound_token_shapes), + sound_sequence_indexes=_mv(self.sound_sequence_indexes), + sound_timesteps=_mv(self.sound_timesteps), + sound_mse_loss_indexes=_mv(self.sound_mse_loss_indexes), + sound_noisy_frame_indexes=[_mv(t) for t in self.sound_noisy_frame_indexes], + fps_sound=_mv(self.fps_sound), ) @@ -203,6 +238,15 @@ def pack_cosmos3_video_sequence( vision_condition_mask: list[torch.Tensor] = [] fps_values: list[float] = [] + sound_tokens: list[torch.Tensor] = [] + sound_token_shapes: list[tuple[int, int, int]] = [] + sound_sequence_indexes: list[int] = [] + sound_timesteps: list[float] = [] + sound_mse_loss_indexes: list[int] = [] + sound_noisy_frame_indexes: list[torch.Tensor] = [] + sound_condition_mask: list[torch.Tensor] = [] + sound_fps_values: list[float] = [] + curr = 0 # running position in the packed sequence is_image_batch = True @@ -236,6 +280,8 @@ def pack_cosmos3_video_sequence( # End of text modality: bump temporal offset before vision. temporal_offset += temporal_modality_margin + # Sound shares the vision temporal start (parallel temporal positions). + vision_start_temporal_offset = temporal_offset # ---- 2. Vision split (full) ---- latent = sample.vision.latent @@ -289,6 +335,57 @@ def pack_cosmos3_video_sequence( curr += num_vision_tokens sample_len += num_vision_tokens + # ---- 2b. Sound split (t2vs): shares the vision "full" split ---- + # Mirrors framework ``_pack_sound_tokens``: sound latent [C, T] -> T + # tokens (token shape (T,1,1)), packed right after vision, with 3D-MRoPE + # temporal positions starting at the vision temporal offset (parallel to + # vision, start_frame_offset=0, tcf=1) and NOT advancing it. + sound_split_len = 0 + if sample.sound is not None: + sound_latent = sample.sound.latent + sound_latent = sound_latent.squeeze(0) if sound_latent.dim() == 3 else sound_latent # [C, T] + _sc, sound_t = sound_latent.shape + sound_split_len = sound_t + + sound_tokens.append(sound_latent) + sound_token_shapes.append((sound_t, 1, 1)) + sound_sequence_indexes.extend(range(curr, curr + sound_t)) + + s_cond_set = {idx for idx in sample.sound.condition_frame_indexes if 0 <= idx < sound_t} + s_cond_mask = torch.zeros((sound_t, 1), device=sound_latent.device, dtype=sound_latent.dtype) + for fi in s_cond_set: + s_cond_mask[fi, 0] = 1.0 + sound_condition_mask.append(s_cond_mask) + + s_noisy = torch.tensor([idx for idx in range(sound_t) if idx not in s_cond_set], + device=sound_latent.device, + dtype=torch.long) + sound_noisy_frame_indexes.append(s_noisy) + + for fi in range(sound_t): + if fi in s_cond_set: + continue + sound_mse_loss_indexes.append(curr + fi) # 1 token per sound frame + sound_timesteps.append(float(sample.timestep)) + + sound_fps = sample.sound.fps if enable_fps_modulation else None + if sound_fps is not None: + sound_fps_values.append(float(sound_fps)) + sound_mrope, _ = compute_mrope_position_ids_vision( + grid_t=sound_t, + grid_h=1, + grid_w=1, + temporal_offset=vision_start_temporal_offset, + fps=sound_fps, + base_fps=base_fps, + temporal_compression_factor=1, # sound latent already at sound_latent_fps + enable_fps_modulation=enable_fps_modulation, + start_frame_offset=0, + ) + position_id_blocks.append(sound_mrope) + curr += sound_t + sample_len += sound_t + # ---- 3. Optional end-of-generation marker ---- eov_len = 0 if include_end_of_generation_token: @@ -304,9 +401,9 @@ def pack_cosmos3_video_sequence( eov_len = 1 sample_len += 1 - # The vision split and any trailing eov marker share the "full" split. + # Vision + sound + any trailing eov marker share one "full" split. attn_modes.append("full") - split_lens.append(num_vision_tokens + eov_len) + split_lens.append(num_vision_tokens + sound_split_len + eov_len) sample_lens.append(sample_len) if latent_t != 1: @@ -338,4 +435,12 @@ def pack_cosmos3_video_sequence( 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), + sound_tokens=sound_tokens, + sound_token_shapes=sound_token_shapes, + sound_sequence_indexes=(torch.tensor(sound_sequence_indexes, dtype=torch.long) if sound_tokens else None), + sound_timesteps=(torch.tensor(sound_timesteps, dtype=timesteps_dtype) if sound_tokens else None), + sound_mse_loss_indexes=(torch.tensor(sound_mse_loss_indexes, dtype=torch.long) if sound_tokens else None), + sound_noisy_frame_indexes=sound_noisy_frame_indexes, + sound_condition_mask=sound_condition_mask, + fps_sound=(torch.tensor(sound_fps_values, dtype=torch.float32) if sound_fps_values else None), ) diff --git a/tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py b/tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py index c05a09af5..b4b83fc67 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py +++ b/tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py @@ -71,13 +71,20 @@ _TCF = 4 # temporal compression factor # --------------------------------------------------------------------------- # Tiny model builders (framework + FastVideo) with unified_3d_mrope. # --------------------------------------------------------------------------- -def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2): +_SOUND_DIM = 64 +_SOUND_LATENT_FPS = 25 + + +def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bool = False): """Tiny framework ``Cosmos3VFMNetwork`` with ``unified_3d_mrope``. ``rope_theta`` / ``rope_scaling`` (carrying ``mrope_section`` + ``mrope_interleaved``) are threaded through the materialized text config; ``position_embedding_type="unified_3d_mrope"`` leaves ``latent_pos_embed`` as ``None`` so positions ride solely on the 3D rotary embedding. + + ``sound_gen=True`` additionally builds the sound MoT heads (``sound2llm`` / + ``llm2sound`` / ``sound_modality_embed``) for the t2vs parity test. """ from cosmos_framework.model.vfm.mot.cosmos3_vfm_network import ( Cosmos3VFMNetwork, @@ -116,6 +123,12 @@ def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2): include_visual=False, ) tiny_vlm_cfg = Qwen3VLConfig(text_config=tiny_text_dict) + sound_kwargs = dict( + sound_gen=True, + sound_dim=_SOUND_DIM, + temporal_compression_factor_sound=1, + sound_latent_fps=_SOUND_LATENT_FPS, + ) if sound_gen else {} vfm_cfg = Cosmos3VFMNetworkConfig( vision_gen=True, vlm_config=tiny_vlm_cfg, @@ -127,6 +140,7 @@ def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2): max_latent_w=16, max_latent_t=8, temporal_compression_factor_vision=_TCF, + **sound_kwargs, ) torch.manual_seed(seed) lm = Qwen3VLTextForCausalLM(config=mot_cfg) diff --git a/tests/local_tests/cosmos3/test_cosmos3_sound_parity.py b/tests/local_tests/cosmos3/test_cosmos3_sound_parity.py new file mode 100644 index 000000000..98ae95bec --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_sound_parity.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo Cosmos3 sound (t2vs) pathway vs the framework. + +Covers the audio generation pathway end-to-end at the DiT level: + + * **sound packing** — FastVideo's native packer + (``pack_cosmos3_video_sequence`` with a ``Cosmos3SoundItem``) vs the + framework ``pack_input_sequence`` with ``has_sound``: sound tokens share the + vision "full" split, with ``(T,1,1)`` shapes, a ``(T,1)`` condition mask, + and 3D-MRoPE temporal positions starting at the vision temporal offset + (parallel to vision); and + * **DiT sound forward** — the dormant ``audio_proj_in`` / ``audio_proj_out`` / + ``audio_modality_embed`` heads, now activated (framework ``sound2llm`` / + ``llm2sound`` / ``sound_modality_embed``). + +Both tiny models are built sound-enabled from the SAME config, framework weights +(incl. the sound heads) are copied into the FastVideo DiT, and the FRAMEWORK +model + framework pack is the parity ORACLE (CPU/float32 via the SDPA +monkey-patch). We assert the native packer matches the framework field-by-field, +then that ``preds_vision`` AND ``preds_sound`` match the framework forward. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_sound_parity.py -q -s +""" +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.", +) + +from .test_cosmos3_dit_parity import ( # noqa: E402 + _fastvideo_inputs_from_packed_seq, + _framework_to_fastvideo_state_dict, +) +from .test_cosmos3_dit_parity_mrope import ( # noqa: E402 + _LATENT_CHANNEL, + _LATENT_PATCH_SIZE, + _RESET_SPATIAL_IDS, + _SOUND_DIM, + _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] +_apply_sdpa_patches() + +_SPECIAL_TOKENS = {"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62} + + +def _copy_weights_with_sound(vfm, dit) -> None: + """Copy backbone + vision weights AND the sound MoT heads into the DiT.""" + mapped = _framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers) + src = dict(vfm.named_parameters()) + mapped["audio_proj_in.weight"] = src["sound2llm.weight"].detach().clone() + mapped["audio_proj_in.bias"] = src["sound2llm.bias"].detach().clone() + mapped["audio_proj_out.weight"] = src["llm2sound.weight"].detach().clone() + mapped["audio_proj_out.bias"] = src["llm2sound.bias"].detach().clone() + mapped["audio_modality_embed"] = src["sound_modality_embed"].detach().clone() + dst = dict(dit.named_parameters()) + with torch.no_grad(): + for name, tensor in mapped.items(): + assert name in dst, f"DiT missing param {name!r}" + assert dst[name].shape == tensor.shape, f"shape mismatch {name}" + dst[name].copy_(tensor.to(dst[name].dtype)) + + +def _framework_pack_sound(*, text_ids, vision, sound, cond_vision, cond_sound, timestep, is_image_batch): + from cosmos_framework.data.vfm.sequence_packing import ( + GenerationDataClean, + SequencePlan, + pack_input_sequence, + ) + + gen = GenerationDataClean( + batch_size=1, + is_image_batch=is_image_batch, + x0_tokens_vision=[vision], + fps_vision=None, + num_vision_items_per_sample=[1], + x0_tokens_sound=[sound], + fps_sound=None, + ) + plans = [SequencePlan( + has_text=True, has_vision=True, has_sound=True, + condition_frame_indexes_vision=list(cond_vision), + condition_frame_indexes_sound=list(cond_sound), + )] + return pack_input_sequence( + sequence_plans=plans, + input_text_indexes=[list(text_ids)], + gen_data_clean=gen, + input_timesteps=torch.tensor([timestep], 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_sound(*, text_ids, vision, sound, cond_vision, cond_sound, timestep): + from fastvideo.pipelines.basic.cosmos3.sequence_packing import ( + Cosmos3SampleInputs, + Cosmos3SoundItem, + Cosmos3VisionItem, + pack_cosmos3_video_sequence, + ) + + samples = [Cosmos3SampleInputs( + text_ids=list(text_ids), + vision=Cosmos3VisionItem(latent=vision, condition_frame_indexes=list(cond_vision)), + sound=Cosmos3SoundItem(latent=sound, condition_frame_indexes=list(cond_sound)), + timestep=float(timestep), + )] + 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, + ) + + +def _fv_inputs_with_sound(ps) -> dict: + """Framework PackedSequence (with sound) -> native DiT forward kwargs.""" + kw = _fastvideo_inputs_from_packed_seq(ps) + s = ps.sound + kw.update( + sound_tokens=list(s.tokens), + sound_token_shapes=[tuple(x) for x in s.token_shapes], + sound_sequence_indexes=s.sequence_indexes, + sound_timesteps=s.timesteps, + sound_mse_loss_indexes=s.mse_loss_indexes, + sound_noisy_frame_indexes=list(s.noisy_frame_indexes), + fps_sound=None, + ) + return kw + + +def _diffs(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float]: + d = (a - b).abs() + return d.max().item(), d.mean().item() + + +# (grid_t, latent_h, latent_w, sound_t, n_text, cond_vision, cond_sound) +_CASES = [ + pytest.param(2, 4, 4, 5, 4, [], [], id="t2vs_2x2x2_snd5"), + pytest.param(3, 8, 4, 8, 5, [], [], id="t2vs_3x4x2_snd8"), + pytest.param(2, 4, 4, 6, 5, [0], [], id="i2vs_cond_snd6"), +] + + +class TestCosmos3SoundParity: + + def _build(self, num_layers=2, seed_model=42): + vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers, sound_gen=True) + dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers) + _copy_weights_with_sound(vfm, dit) + return vfm, dit + + def _make_inputs(self, grid_t, latent_h, latent_w, sound_t, n_text, cond_v, cond_s, seed=7): + torch.manual_seed(seed) + vision = torch.randn(1, _LATENT_CHANNEL, grid_t, latent_h, latent_w) + sound = torch.randn(_SOUND_DIM, sound_t) # [C, T] + text_ids = torch.randint(0, 60, (n_text,)).tolist() + return dict(text_ids=text_ids, vision=vision, sound=sound, + cond_vision=cond_v, cond_sound=cond_s, timestep=500.0) + + @pytest.mark.parametrize(("grid_t", "lh", "lw", "snd_t", "n_text", "cond_v", "cond_s"), _CASES) + def test_sound_packing_matches_framework(self, grid_t, lh, lw, snd_t, n_text, cond_v, cond_s): + ins = self._make_inputs(grid_t, lh, lw, snd_t, n_text, cond_v, cond_s) + fw = _framework_pack_sound(is_image_batch=(grid_t == 1), **ins) + fv = _fastvideo_pack_sound(**ins) + + 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) + assert int(fv.sequence_length) == int(fw.sequence_length) + torch.testing.assert_close(fv.position_ids, fw.position_ids, rtol=0, atol=0) # [3, seq], exact + # Sound fields. + s = fw.sound + torch.testing.assert_close(fv.sound_sequence_indexes, s.sequence_indexes.to(torch.long), rtol=0, atol=0) + assert fv.sound_token_shapes == [tuple(x) for x in s.token_shapes] + torch.testing.assert_close(fv.sound_timesteps.to(torch.float32), s.timesteps.to(torch.float32)) + torch.testing.assert_close(fv.sound_mse_loss_indexes, s.mse_loss_indexes.to(torch.long), rtol=0, atol=0) + for a, b in zip(fv.sound_noisy_frame_indexes, s.noisy_frame_indexes): + torch.testing.assert_close(a.to(torch.long), b.to(torch.long), rtol=0, atol=0) + for a, b in zip(fv.sound_condition_mask, s.condition_mask): + torch.testing.assert_close(a.flatten().to(torch.float32), b.flatten().to(torch.float32)) + print(f"\n[sound_packing {grid_t}x{lh}x{lw} snd={snd_t}] position_ids + sound fields exact") + + @pytest.mark.parametrize(("grid_t", "lh", "lw", "snd_t", "n_text", "cond_v", "cond_s"), _CASES) + def test_sound_dit_forward_matches_framework(self, grid_t, lh, lw, snd_t, n_text, cond_v, cond_s): + vfm, dit = self._build() + ins = self._make_inputs(grid_t, lh, lw, snd_t, n_text, cond_v, cond_s) + fw_pack = _framework_pack_sound(is_image_batch=(grid_t == 1), **ins) + fv_pack = _fastvideo_pack_sound(**ins) + + with torch.no_grad(): + fw_out = vfm(packed_seq=fw_pack) # framework model + framework pack (oracle) + fv_out = dit(**fv_pack.to_dit_kwargs()) # native model + native pack + # Also run the native DiT on the framework pack to isolate the forward. + fv_on_fw = dit(**_fv_inputs_with_sound(fw_pack)) + + # preds_vision parity. + pv_max, pv_mean = _diffs(fv_out["preds_vision"][0], fw_out["preds_vision"][0]) + # preds_sound parity. + ps_max, ps_mean = _diffs(fv_out["preds_sound"][0], fw_out["preds_sound"][0]) + # native-DiT-on-framework-pack (forward only) parity. + psf_max, psf_mean = _diffs(fv_on_fw["preds_sound"][0], fw_out["preds_sound"][0]) + print(f"\n[sound_dit {grid_t}x{lh}x{lw} snd={snd_t}] " + f"preds_vision max={pv_max:.3e} mean={pv_mean:.3e} | " + f"preds_sound max={ps_max:.3e} mean={ps_mean:.3e} | " + f"preds_sound(fwpack) max={psf_max:.3e} mean={psf_mean:.3e}") + + assert fv_out["preds_sound"][0].shape == fw_out["preds_sound"][0].shape + torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3) + torch.testing.assert_close(fv_out["preds_sound"][0], fw_out["preds_sound"][0], atol=1e-4, rtol=1e-3) + torch.testing.assert_close(fv_on_fw["preds_sound"][0], fw_out["preds_sound"][0], atol=1e-4, rtol=1e-3)