[feat] cosmos3 PR2: DiT sound pathway + sound packing (bit-exact parity)
Activate the dormant sound MoT heads and pack the sound modality, completing the t2vs DiT path (audio component 2). DiT forward (`Cosmos3VFMTransformer`): - Encode sound: pack each [C,T] latent -> [T,C], `audio_proj_in` (= framework `sound2llm`) + `audio_modality_embed`, scatter timestep embeds onto noisy frames, scatter into the (full) gen split. Mirrors framework `_encode_sound`. - Decode sound: `audio_proj_out` (= `llm2sound`) on noisy sound hidden states -> `_unpack_sound` to per-sample [C,T] (framework `_decode_sound`). - Reuses the general `_scatter_timestep_embeds` (works for (T,1,1) shapes). Sound packing (`sequence_packing.py`): `Cosmos3SoundItem` + sound fields on `Cosmos3SampleInputs`/`Cosmos3PackedSequence`/`to_dit_kwargs`. Sound shares the vision "full" split (preserving the causal+full 2-split invariant), with `(T,1,1)` token shapes, a `(T,1)` condition mask, mse-loss/timestep bookkeeping per noisy frame, and 3D-MRoPE temporal positions starting at the vision temporal offset (parallel to vision; `start_frame_offset=0`, tcf=1; does not advance the offset). Mirrors framework `_pack_sound_tokens`. Parity (`test_cosmos3_sound_parity.py`, framework model+pack = oracle): - sound packing fields + `position_ids` exact vs framework `pack_input_sequence` (has_sound) across t2vs/i2vs cases. - `preds_vision` AND `preds_sound`: max abs diff = 0.0, mean abs diff = 0.0 (both native-pack and framework-pack inputs). Tiny framework builder gains a `sound_gen` flag. Full suite: 127 passed.
This commit is contained in:
SolitaryThinker
@@ -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]:
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 <worktree> && <fv-cosmos3 python> -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)
|
||||
Reference in New Issue
Block a user