From ec63ce4c2b64331e533978eecfa539a1e54fa52a Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 7 Jun 2026 13:12:23 -0700 Subject: [PATCH] [feat] cosmos3 PR3: action (domain-aware) DiT pathway + packing + CFG glue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the action (multi-embodiment world-model) modality. The Nano checkpoint ships action weights (action_gen=true, 32 embodiment domains); this activates them, bit-exact vs the framework. - DiT forward: action encode/decode mirroring sound but domain-aware — pack each [T, action_dim] latent with a per-token embodiment domain id, `action_proj_in` (= framework `action2llm`, a `DomainAwareLinear`) + `action_modality_embed` + timestep scatter; decode via `action_proj_out` (= `llm2action`) on noisy hidden states -> `_unpack_action`. Mirrors `_encode_action`/`_decode_action`. - Sound/action share the vision "full" split; sequence/flat order is [vision | action | sound] (matches the framework per-sample concat). - Packing: `Cosmos3ActionItem` + action fields. Action tokens are `(T,)`-shaped, domain-tagged, with 3D-MRoPE at the vision temporal offset and `start_frame_offset=1` (framework `_pack_action_tokens`). - CFG velocity / `Cosmos3DenoiseEngine` take optional `action_specs` (`Cosmos3ActionSpec`); the joint denoise machinery now handles action. Parity (`test_cosmos3_action_parity.py`, framework model+pack = oracle, 3 embodiment domains): - action packing fields + `position_ids` exact; - `preds_vision` + `preds_action` (forward): max abs diff = 0.0, mean = 0.0; - combined [vision|action] CFG velocity: max abs diff = 0.0, mean = 0.0. Action is processed bit-exactly at the model + packing + denoise-glue level. A full action2world real-weights run is gated on real robot-action input data (not available here); the bit-exact parity is the correctness proof. Full suite: 139 passed, 0 skipped. --- fastvideo/models/dits/cosmos3.py | 56 ++++ .../basic/cosmos3/cosmos3_pipeline.py | 109 +++++-- .../basic/cosmos3/sequence_packing.py | 109 ++++++- .../cosmos3/test_cosmos3_action_parity.py | 271 ++++++++++++++++++ .../cosmos3/test_cosmos3_dit_parity_mrope.py | 11 +- 5 files changed, 534 insertions(+), 22 deletions(-) create mode 100644 tests/local_tests/cosmos3/test_cosmos3_action_parity.py diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py index 40151fe3f..71eaf214b 100644 --- a/fastvideo/models/dits/cosmos3.py +++ b/fastvideo/models/dits/cosmos3.py @@ -808,6 +808,13 @@ class Cosmos3VFMTransformer(BaseDiT): sound_mse_loss_indexes: torch.Tensor | None = None, sound_noisy_frame_indexes: list[torch.Tensor] | None = None, fps_sound: torch.Tensor | None = None, + action_tokens: list[torch.Tensor] | None = None, + action_token_shapes: list[tuple[int, ...]] | None = None, + action_sequence_indexes: torch.Tensor | None = None, + action_timesteps: torch.Tensor | None = None, + action_mse_loss_indexes: torch.Tensor | None = None, + action_noisy_frame_indexes: list[torch.Tensor] | None = None, + action_domain_id: list[torch.Tensor] | None = None, **kwargs: Any, ) -> dict[str, Any]: """Video-path forward returning ``{"last_hidden_state", "preds_vision"}`` @@ -863,6 +870,26 @@ class Cosmos3VFMTransformer(BaseDiT): sound_token_shapes) packed_sequence[sound_sequence_indexes] = packed_sound + # 2c. Action (DomainAwareLinear): pack [T, D] per sample with a per-token + # domain id -> action_proj_in(domain) + action_modality_embed -> + # timestep scatter -> scatter into the gen split (framework + # ``_encode_action``). Action is domain-conditioned (per embodiment). + if action_tokens: + packed_action = torch.cat( + [a[:shp[0]] for a, shp in zip(action_tokens, action_token_shapes)], dim=0, + ) # [total_action_tokens, action_dim] + per_token_domain = torch.cat( + [d.reshape(1).expand(shp[0]) for d, shp in zip(action_domain_id, action_token_shapes)], dim=0, + ) # [total_action_tokens] + packed_action = self.action_proj_in(packed_action.to(target_dtype), per_token_domain) + packed_action = packed_action + self.action_modality_embed.view(1, -1) + if action_mse_loss_indexes is not None and action_mse_loss_indexes.numel() > 0: + a_ts = action_timesteps * self.timestep_scale + a_embeds = self.time_embedder(a_ts).to(target_dtype) + packed_action = self._scatter_timestep_embeds(packed_action, a_embeds, action_noisy_frame_indexes, + action_token_shapes) + packed_sequence[action_sequence_indexes] = packed_action + # 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) @@ -895,9 +922,38 @@ class Cosmos3VFMTransformer(BaseDiT): 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) + # 6c. Action prediction: action_proj_out(per-token domain) on noisy hidden + # states -> unpack to per-sample [T, D] (framework ``_decode_action``). + if action_tokens and action_mse_loss_indexes is not None and action_mse_loss_indexes.numel() > 0: + noisy_domain = torch.cat( + [d.reshape(1).expand(len(nfi)) for d, nfi in zip(action_domain_id, action_noisy_frame_indexes)], dim=0, + ) + preds_action = self.action_proj_out(last_hidden_state[action_mse_loss_indexes], noisy_domain) + output["preds_action"] = self._unpack_action(preds_action, action_token_shapes, action_noisy_frame_indexes) + output["last_hidden_state"] = last_hidden_state return output + def _unpack_action( + self, + packed_preds: torch.Tensor, + token_shapes: list[tuple[int, ...]], + noisy_frame_indexes: list[torch.Tensor], + ) -> list[torch.Tensor]: + """Scatter packed noisy action preds back to per-sample ``[T, D]`` (clean + frames left zero). Mirrors framework ``unpack_action``.""" + outputs: list[torch.Tensor] = [] + start_idx = 0 + for shape, nfi in zip(token_shapes, noisy_frame_indexes): + t = shape[0] + out = torch.zeros((t, self.action_dim), 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] + start_idx += t_n + outputs.append(out) + return outputs + def _unpack_sound( self, packed_preds: torch.Tensor, diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index 4882fe222..756080f55 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -48,6 +48,7 @@ from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_unipc_multistep import ( UniPCMultistepScheduler, ) from fastvideo.pipelines.basic.cosmos3.sequence_packing import ( + Cosmos3ActionItem, Cosmos3SampleInputs, Cosmos3SoundItem, Cosmos3VisionItem, @@ -215,6 +216,37 @@ class Cosmos3SoundSpec: return int(math.prod(self.shape)) +@dataclass +class Cosmos3ActionSpec: + """Geometry + conditioning for one action item in a denoise run. + + Args: + shape: ``(T, action_dim)`` of the action latent. + condition_frame_indexes: Frame indices kept clean (conditioning actions). + domain_id: Embodiment domain id for the domain-aware action projection. + fps: Action FPS; used iff fps modulation is on. + """ + + shape: tuple[int, int] + condition_frame_indexes: list[int] = field(default_factory=list) + domain_id: int = 0 + fps: float | None = None + + @property + def numel(self) -> int: + return int(math.prod(self.shape)) + + +def _split_action_latent(flat: torch.Tensor, specs: list[Cosmos3ActionSpec]) -> list[torch.Tensor]: + """Split a flat action vector into per-item ``[T, action_dim]`` tensors.""" + out: list[torch.Tensor] = [] + offset = 0 + for spec in specs: + out.append(flat[offset:offset + spec.numel].reshape(spec.shape)) + offset += spec.numel + return out + + def cosmos3_get_cfg_velocity( *, transformer: Any, @@ -236,6 +268,8 @@ def cosmos3_get_cfg_velocity( normalize_cfg: bool = False, sound_specs: list[Cosmos3SoundSpec] | None = None, sound_fps_per_item: list[float] | None = None, + action_specs: list[Cosmos3ActionSpec] | None = None, + action_fps_per_item: list[float] | None = None, ) -> torch.Tensor: """Sequential-CFG velocity for one denoise step (framework math). @@ -257,11 +291,16 @@ def cosmos3_get_cfg_velocity( assert timestep.numel() == 1, "timestep must be a scalar" timestep_value = float(timestep.reshape(()).item()) - # Combined flat layout (t2vs): [all vision | all sound], matching the - # framework per-sample concat order ([vision_i | sound_i]); single sample here. + # Combined flat layout: [all vision | all action | all sound], matching the + # framework per-sample concat order ([vision_i | action_i | sound_i]); single + # sample here. vision_total = sum(spec.numel for spec in specs) + action_total = sum(spec.numel for spec in action_specs) if action_specs else 0 noise_x_vision = _split_flat_latent(flat_latent[:vision_total], specs) - noise_x_sound = (_split_flat_latent(flat_latent[vision_total:], sound_specs) if sound_specs else None) + noise_x_action = (_split_action_latent(flat_latent[vision_total:vision_total + + action_total], action_specs) if action_specs else None) + noise_x_sound = (_split_flat_latent(flat_latent[vision_total + + action_total:], sound_specs) if sound_specs else None) device = next(transformer.parameters()).device def _run(token_ids: list[int]) -> torch.Tensor: @@ -274,6 +313,16 @@ def cosmos3_get_cfg_velocity( fps=(sound_fps_per_item[i] if sound_fps_per_item is not None else None), ) for i, ss in enumerate(sound_specs) ] + action_items: list[Cosmos3ActionItem] = [] + if action_specs is not None and noise_x_action is not None: + action_items = [ + Cosmos3ActionItem( + latent=noise_x_action[i], + condition_frame_indexes=list(asp.condition_frame_indexes), + domain_id=asp.domain_id, + fps=(action_fps_per_item[i] if action_fps_per_item is not None else None), + ) for i, asp in enumerate(action_specs) + ] samples = [ Cosmos3SampleInputs( text_ids=list(token_ids), @@ -283,6 +332,7 @@ def cosmos3_get_cfg_velocity( fps=(fps_per_item[i] if fps_per_item is not None else None), ), sound=(sound_items[i] if i < len(sound_items) else None), + action=(action_items[i] if i < len(action_items) else None), timestep=timestep_value, ) for i, (latent, spec) in enumerate(zip(noise_x_vision, specs, strict=False)) ] @@ -310,23 +360,38 @@ def cosmos3_get_cfg_velocity( items.append(pred * keep if keep.sum() > 0 else torch.zeros_like(pred)) vision_vel = torch.cat([v.reshape(-1) for v in items]).to(flat_latent.dtype) - if not sound_specs: - return vision_vel + parts = [vision_vel] - # Sound velocity: preds_sound are per-item [C, T], already zero on clean - # frames (unpack fills only noisy frames); zero on cond frames defensively. - sound_total = sum(spec.numel for spec in sound_specs) - sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype) - preds_s = out.get("preds_sound") - if preds_s is not None: - s_items: list[torch.Tensor] = [] - for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False): - pred = pred.squeeze(0) if pred.dim() == 3 else pred # [C, T] - keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device) # [1, T] - s_items.append(pred * keep) - sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype) + if action_specs: + # Action velocity: preds_action are per-item [T, D], already zero on + # clean frames; zero on cond frames defensively. + action_vel = torch.zeros(action_total, device=flat_latent.device, dtype=flat_latent.dtype) + preds_a = out.get("preds_action") + if preds_a is not None: + a_items: list[torch.Tensor] = [] + for pred, cond_mask in zip(preds_a, packed.action_condition_mask, strict=False): + pred = pred.squeeze(0) if pred.dim() == 3 else pred # [T, D] + keep = (1.0 - cond_mask).reshape(-1, 1).to(dtype=pred.dtype, device=pred.device) # [T, 1] + a_items.append(pred * keep) + action_vel = torch.cat([v.reshape(-1) for v in a_items]).to(flat_latent.dtype) + parts.append(action_vel) - return torch.cat([vision_vel, sound_vel]) + if sound_specs: + # Sound velocity: preds_sound are per-item [C, T], already zero on clean + # frames (unpack fills only noisy frames); zero on cond frames defensively. + sound_total = sum(spec.numel for spec in sound_specs) + sound_vel = torch.zeros(sound_total, device=flat_latent.device, dtype=flat_latent.dtype) + preds_s = out.get("preds_sound") + if preds_s is not None: + s_items: list[torch.Tensor] = [] + for pred, cond_mask in zip(preds_s, packed.sound_condition_mask, strict=False): + pred = pred.squeeze(0) if pred.dim() == 3 else pred # [C, T] + keep = (1.0 - cond_mask).reshape(1, -1).to(dtype=pred.dtype, device=pred.device) # [1, T] + s_items.append(pred * keep) + sound_vel = torch.cat([v.reshape(-1) for v in s_items]).to(flat_latent.dtype) + parts.append(sound_vel) + + return vision_vel if len(parts) == 1 else torch.cat(parts) cond_v = _run(cond_token_ids) uncond_v = _run(uncond_token_ids) @@ -382,6 +447,8 @@ class Cosmos3DenoiseEngine: fps_per_item: list[float] | None = None, sound_specs: list[Cosmos3SoundSpec] | None = None, sound_fps_per_item: list[float] | None = None, + action_specs: list[Cosmos3ActionSpec] | None = None, + action_fps_per_item: list[float] | None = None, ) -> torch.Tensor: return cosmos3_get_cfg_velocity( transformer=self.transformer, @@ -402,6 +469,8 @@ class Cosmos3DenoiseEngine: fps_per_item=fps_per_item, sound_specs=sound_specs, sound_fps_per_item=sound_fps_per_item, + action_specs=action_specs, + action_fps_per_item=action_fps_per_item, ) def denoise( @@ -417,6 +486,8 @@ class Cosmos3DenoiseEngine: progress_bar: Any | None = None, sound_specs: list[Cosmos3SoundSpec] | None = None, sound_fps_per_item: list[float] | None = None, + action_specs: list[Cosmos3ActionSpec] | None = None, + action_fps_per_item: list[float] | None = None, ) -> torch.Tensor: """Run the full UniPC denoise loop, returning the final flat latent. @@ -439,6 +510,8 @@ class Cosmos3DenoiseEngine: fps_per_item=fps_per_item, sound_specs=sound_specs, sound_fps_per_item=sound_fps_per_item, + action_specs=action_specs, + action_fps_per_item=action_fps_per_item, ) stepped = self.scheduler.step( model_output=v_pred, diff --git a/fastvideo/pipelines/basic/cosmos3/sequence_packing.py b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py index fb0c42918..8a88892c0 100644 --- a/fastvideo/pipelines/basic/cosmos3/sequence_packing.py +++ b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py @@ -86,14 +86,33 @@ class Cosmos3SoundItem: fps: float | None = None +@dataclass +class Cosmos3ActionItem: + """One action latent for a sample (action-conditioned world model). + + Args: + latent: Action latent ``[T, action_dim]`` (per-frame action vectors). + condition_frame_indexes: Frame indices kept clean (conditioning actions). + domain_id: Embodiment domain id (scalar / ``[1]``) for the + domain-aware action projection. + fps: Action FPS; only used when ``enable_fps_modulation`` is set. + """ + + latent: torch.Tensor + condition_frame_indexes: list[int] = field(default_factory=list) + domain_id: int = 0 + fps: float | None = None + + @dataclass class Cosmos3SampleInputs: - """Per-sample packing inputs (one text prompt + its vision item, +sound).""" + """Per-sample packing inputs (text prompt + vision item, +sound, +action).""" text_ids: list[int] vision: Cosmos3VisionItem timestep: float sound: Cosmos3SoundItem | None = None + action: Cosmos3ActionItem | None = None # --------------------------------------------------------------------------- @@ -139,6 +158,16 @@ class Cosmos3PackedSequence: sound_condition_mask: list[torch.Tensor] = field(default_factory=list) fps_sound: torch.Tensor | None = None + # Action modality (action-conditioned world model); empty/None when no action. + action_tokens: list[torch.Tensor] = field(default_factory=list) + action_token_shapes: list[tuple[int, ...]] = field(default_factory=list) + action_sequence_indexes: torch.Tensor | None = None + action_timesteps: torch.Tensor | None = None + action_mse_loss_indexes: torch.Tensor | None = None + action_noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list) + action_condition_mask: list[torch.Tensor] = field(default_factory=list) + action_domain_id: list[torch.Tensor] = field(default_factory=list) + def to_dit_kwargs(self, device: torch.device | str | None = None) -> dict[str, Any]: """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``. @@ -171,6 +200,13 @@ class Cosmos3PackedSequence: 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), + action_tokens=[_mv(t) for t in self.action_tokens], + action_token_shapes=list(self.action_token_shapes), + action_sequence_indexes=_mv(self.action_sequence_indexes), + action_timesteps=_mv(self.action_timesteps), + action_mse_loss_indexes=_mv(self.action_mse_loss_indexes), + action_noisy_frame_indexes=[_mv(t) for t in self.action_noisy_frame_indexes], + action_domain_id=[_mv(t) for t in self.action_domain_id], ) @@ -247,6 +283,15 @@ def pack_cosmos3_video_sequence( sound_condition_mask: list[torch.Tensor] = [] sound_fps_values: list[float] = [] + action_tokens: list[torch.Tensor] = [] + action_token_shapes: list[tuple[int, ...]] = [] + action_sequence_indexes: list[int] = [] + action_timesteps: list[float] = [] + action_mse_loss_indexes: list[int] = [] + action_noisy_frame_indexes: list[torch.Tensor] = [] + action_condition_mask: list[torch.Tensor] = [] + action_domain_id: list[torch.Tensor] = [] + curr = 0 # running position in the packed sequence is_image_batch = True @@ -335,6 +380,56 @@ def pack_cosmos3_video_sequence( curr += num_vision_tokens sample_len += num_vision_tokens + # ---- 2a2. Action split: shares the vision "full" split ---- + # Mirrors framework ``_pack_action_tokens``: action latent [T, D] -> T + # tokens (token shape (T,)), domain-aware, 3D-MRoPE at the vision temporal + # offset with ``start_frame_offset=1`` (parallel to vision; tcf=1; does + # not advance the offset). + action_split_len = 0 + if sample.action is not None: + action_latent = sample.action.latent # [T, D] + action_t = int(action_latent.shape[0]) + action_split_len = action_t + + action_tokens.append(action_latent) + action_token_shapes.append((action_t, )) + action_sequence_indexes.extend(range(curr, curr + action_t)) + action_domain_id.append(torch.tensor([int(sample.action.domain_id)], dtype=torch.long)) + + a_cond_set = {idx for idx in sample.action.condition_frame_indexes if 0 <= idx < action_t} + a_cond_mask = torch.zeros((action_t, 1), device=action_latent.device, dtype=action_latent.dtype) + for fi in a_cond_set: + a_cond_mask[fi, 0] = 1.0 + action_condition_mask.append(a_cond_mask) + + a_noisy = torch.tensor([idx for idx in range(action_t) if idx not in a_cond_set], + device=action_latent.device, + dtype=torch.long) + action_noisy_frame_indexes.append(a_noisy) + + for fi in range(action_t): + if fi in a_cond_set: + continue + action_mse_loss_indexes.append(curr + fi) + action_timesteps.append(float(sample.timestep)) + + action_fps = sample.action.fps if enable_fps_modulation else None + action_mrope, _ = compute_mrope_position_ids_vision( + grid_t=action_t, + grid_h=1, + grid_w=1, + temporal_offset=vision_start_temporal_offset, + fps=action_fps, + base_fps=base_fps, + temporal_compression_factor=1, # action is at frame rate + base_temporal_compression_factor=temporal_compression_factor, + enable_fps_modulation=enable_fps_modulation, + start_frame_offset=1, + ) + position_id_blocks.append(action_mrope) + curr += action_t + sample_len += action_t + # ---- 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 @@ -401,9 +496,9 @@ def pack_cosmos3_video_sequence( eov_len = 1 sample_len += 1 - # Vision + sound + any trailing eov marker share one "full" split. + # Vision + action + sound + any trailing eov marker share one "full" split. attn_modes.append("full") - split_lens.append(num_vision_tokens + sound_split_len + eov_len) + split_lens.append(num_vision_tokens + action_split_len + sound_split_len + eov_len) sample_lens.append(sample_len) if latent_t != 1: @@ -443,4 +538,12 @@ def pack_cosmos3_video_sequence( 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), + action_tokens=action_tokens, + action_token_shapes=action_token_shapes, + action_sequence_indexes=(torch.tensor(action_sequence_indexes, dtype=torch.long) if action_tokens else None), + action_timesteps=(torch.tensor(action_timesteps, dtype=timesteps_dtype) if action_tokens else None), + action_mse_loss_indexes=(torch.tensor(action_mse_loss_indexes, dtype=torch.long) if action_tokens else None), + action_noisy_frame_indexes=action_noisy_frame_indexes, + action_condition_mask=action_condition_mask, + action_domain_id=action_domain_id, ) diff --git a/tests/local_tests/cosmos3/test_cosmos3_action_parity.py b/tests/local_tests/cosmos3/test_cosmos3_action_parity.py new file mode 100644 index 000000000..aac53f4ff --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_action_parity.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo Cosmos3 action pathway vs the framework. + +Covers the action (multi-embodiment world-model) modality at the DiT level: + + * **action packing** — native ``pack_cosmos3_video_sequence`` with a + ``Cosmos3ActionItem`` vs framework ``pack_input_sequence`` with + ``has_action``: action tokens share the vision "full" split, with ``(T,)`` + shapes, a ``(T,1)`` condition mask, and 3D-MRoPE temporal positions at the + vision offset with ``start_frame_offset=1`` (parallel to vision); and + * **DiT action forward** — the dormant domain-aware ``action_proj_in`` / + ``action_proj_out`` (``DomainAwareLinear``) + ``action_modality_embed`` heads, + now activated, with a per-token embodiment ``domain_id``. + +Framework model + pack is the parity ORACLE (CPU/float32 via SDPA monkey-patch). +We assert the native packer matches the framework field-by-field, then that +``preds_vision`` AND ``preds_action`` match the framework forward. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_action_parity.py -q -s +""" +from __future__ import annotations + +import pytest +import torch + +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 + _ACTION_DIM, + _LATENT_CHANNEL, + _LATENT_PATCH_SIZE, + _RESET_SPATIAL_IDS, + _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_action(vfm, dit) -> None: + """Copy backbone + vision weights AND the domain-aware action heads.""" + mapped = _framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers) + src = dict(vfm.named_parameters()) + mapped["action_proj_in.fc.weight"] = src["action2llm.fc.weight"].detach().clone() + mapped["action_proj_in.bias.weight"] = src["action2llm.bias.weight"].detach().clone() + mapped["action_proj_out.fc.weight"] = src["llm2action.fc.weight"].detach().clone() + mapped["action_proj_out.bias.weight"] = src["llm2action.bias.weight"].detach().clone() + mapped["action_modality_embed"] = src["action_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_action(*, text_ids, vision, action, cond_vision, cond_action, domain_id, 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_action=[action], + fps_action=None, + action_domain_id=[torch.tensor([domain_id], dtype=torch.long)], + ) + plans = [SequencePlan( + has_text=True, has_vision=True, has_action=True, + condition_frame_indexes_vision=list(cond_vision), + condition_frame_indexes_action=list(cond_action), + )] + ps = 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, + ) + # The framework sets action.domain_id on the packed sequence from + # gen_data_clean inside the model (_get_velocity); mirror that for the oracle. + if ps.action is not None: + ps.action.domain_id = [torch.tensor([domain_id], dtype=torch.long)] + return ps + + +def _fastvideo_pack_action(*, text_ids, vision, action, cond_vision, cond_action, domain_id, timestep): + from fastvideo.pipelines.basic.cosmos3.sequence_packing import ( + Cosmos3ActionItem, + Cosmos3SampleInputs, + Cosmos3VisionItem, + pack_cosmos3_video_sequence, + ) + + samples = [Cosmos3SampleInputs( + text_ids=list(text_ids), + vision=Cosmos3VisionItem(latent=vision, condition_frame_indexes=list(cond_vision)), + action=Cosmos3ActionItem(latent=action, condition_frame_indexes=list(cond_action), domain_id=domain_id), + 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_action(ps) -> dict: + kw = _fastvideo_inputs_from_packed_seq(ps) + a = ps.action + kw.update( + action_tokens=list(a.tokens), + action_token_shapes=[tuple(x) for x in a.token_shapes], + action_sequence_indexes=a.sequence_indexes, + action_timesteps=a.timesteps, + action_mse_loss_indexes=a.mse_loss_indexes, + action_noisy_frame_indexes=list(a.noisy_frame_indexes), + action_domain_id=list(a.domain_id), + ) + return kw + + +def _diffs(a, b): + d = (a - b).abs() + return d.max().item(), d.mean().item() + + +# (grid_t, lh, lw, action_t, n_text, cond_vision, cond_action, domain_id) +_CASES = [ + pytest.param(2, 4, 4, 6, 4, [], [], 0, id="a2v_2x2x2_act6_dom0"), + pytest.param(3, 8, 4, 9, 5, [], [], 7, id="a2v_3x4x2_act9_dom7"), + pytest.param(2, 4, 4, 5, 5, [0], [0], 3, id="ai2v_cond_act5_dom3"), +] + + +class TestCosmos3ActionParity: + + def _build(self, num_layers=2, seed_model=42): + vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers, action_gen=True) + dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers) + _copy_weights_with_action(vfm, dit) + return vfm, dit + + def _make_inputs(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom, seed=7): + torch.manual_seed(seed) + return dict( + text_ids=torch.randint(0, 60, (n_text,)).tolist(), + vision=torch.randn(1, _LATENT_CHANNEL, grid_t, lh, lw), + action=torch.randn(act_t, _ACTION_DIM), # [T, D] + cond_vision=cond_v, cond_action=cond_a, domain_id=dom, timestep=500.0, + ) + + @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES) + def test_action_packing_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom): + ins = self._make_inputs(grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom) + fw = _framework_pack_action(is_image_batch=(grid_t == 1), **ins) + fv = _fastvideo_pack_action(**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) + a = fw.action + torch.testing.assert_close(fv.action_sequence_indexes, a.sequence_indexes.to(torch.long), rtol=0, atol=0) + assert fv.action_token_shapes == [tuple(x) for x in a.token_shapes] + torch.testing.assert_close(fv.action_timesteps.to(torch.float32), a.timesteps.to(torch.float32)) + torch.testing.assert_close(fv.action_mse_loss_indexes, a.mse_loss_indexes.to(torch.long), rtol=0, atol=0) + for x, y in zip(fv.action_noisy_frame_indexes, a.noisy_frame_indexes): + torch.testing.assert_close(x.to(torch.long), y.to(torch.long), rtol=0, atol=0) + print(f"\n[action_packing {grid_t}x{lh}x{lw} act={act_t} dom={dom}] position_ids + action fields exact") + + @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES) + def test_action_dit_forward_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom): + vfm, dit = self._build() + ins = self._make_inputs(grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom) + fw_pack = _framework_pack_action(is_image_batch=(grid_t == 1), **ins) + fv_pack = _fastvideo_pack_action(**ins) + with torch.no_grad(): + fw_out = vfm(packed_seq=fw_pack) + fv_out = dit(**fv_pack.to_dit_kwargs()) + fv_on_fw = dit(**_fv_inputs_with_action(fw_pack)) + pv_mx, pv_mn = _diffs(fv_out["preds_vision"][0], fw_out["preds_vision"][0]) + pa_mx, pa_mn = _diffs(fv_out["preds_action"][0], fw_out["preds_action"][0]) + paf_mx, paf_mn = _diffs(fv_on_fw["preds_action"][0], fw_out["preds_action"][0]) + print(f"\n[action_dit {grid_t}x{lh}x{lw} act={act_t} dom={dom}] " + f"preds_vision max={pv_mx:.3e} mean={pv_mn:.3e} | " + f"preds_action max={pa_mx:.3e} mean={pa_mn:.3e} | " + f"preds_action(fwpack) max={paf_mx:.3e} mean={paf_mn:.3e}") + assert fv_out["preds_action"][0].shape == fw_out["preds_action"][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_action"][0], fw_out["preds_action"][0], atol=1e-4, rtol=1e-3) + torch.testing.assert_close(fv_on_fw["preds_action"][0], fw_out["preds_action"][0], atol=1e-4, rtol=1e-3) + + @pytest.mark.parametrize(("grid_t", "lh", "lw", "act_t", "n_text", "cond_v", "cond_a", "dom"), _CASES) + def test_action_cfg_velocity_matches_framework(self, grid_t, lh, lw, act_t, n_text, cond_v, cond_a, dom): + """Combined [vision|action] sequential-CFG velocity (action pipeline glue) + matches a framework-DiT oracle.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + Cosmos3ActionSpec, + Cosmos3VisionSpec, + cosmos3_get_cfg_velocity, + ) + + vfm, dit = self._build() + vlat_shape = (_LATENT_CHANNEL, grid_t, lh, lw) + action_shape = (act_t, _ACTION_DIM) + torch.manual_seed(3) + cond_ids = torch.randint(0, 60, (n_text,)).tolist() + uncond_ids = torch.randint(0, 60, (max(1, n_text - 1),)).tolist() + vis_numel = int(torch.tensor(vlat_shape).prod()) + act_numel = int(torch.tensor(action_shape).prod()) + flat = torch.randn(vis_numel + act_numel) + guidance, ts = 6.0, 500.0 + + def _fw_velocity(ids): + vision = flat[:vis_numel].reshape(vlat_shape).unsqueeze(0) + action = flat[vis_numel:].reshape(action_shape) + ps = _framework_pack_action(text_ids=ids, vision=vision, action=action, cond_vision=cond_v, + cond_action=cond_a, domain_id=dom, timestep=ts, + is_image_batch=(grid_t == 1)) + with torch.no_grad(): + out = vfm(packed_seq=ps) + pv = out["preds_vision"][0].squeeze(0) # [C,T,H,W] (zero on clean) + pa = out["preds_action"][0] # [T,D] (zero on clean) + return torch.cat([pv.reshape(-1), pa.reshape(-1)]) + + fw_cond, fw_uncond = _fw_velocity(cond_ids), _fw_velocity(uncond_ids) + fw_v = fw_uncond + guidance * (fw_cond - fw_uncond) + fv_v = cosmos3_get_cfg_velocity( + transformer=dit, flat_latent=flat, timestep=torch.tensor([ts]), guidance=guidance, + specs=[Cosmos3VisionSpec(shape=vlat_shape, condition_frame_indexes=list(cond_v))], + action_specs=[Cosmos3ActionSpec(shape=action_shape, condition_frame_indexes=list(cond_a), domain_id=dom)], + cond_token_ids=cond_ids, uncond_token_ids=uncond_ids, + special_tokens=_SPECIAL_TOKENS, latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, reset_spatial_ids=_RESET_SPATIAL_IDS, + enable_fps_modulation=False, base_fps=24.0, temporal_compression_factor=_TCF, + ) + assert fv_v.shape == fw_v.shape, f"shape fv={fv_v.shape} fw={fw_v.shape}" + mx, mn = _diffs(fv_v, fw_v) + print(f"\n[action_cfg_velocity {grid_t}x{lh}x{lw} act={act_t} dom={dom}] max={mx:.3e} mean={mn:.3e}") + torch.testing.assert_close(fv_v, fw_v, atol=1e-4, rtol=1e-3) 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 b4b83fc67..bed46947e 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py +++ b/tests/local_tests/cosmos3/test_cosmos3_dit_parity_mrope.py @@ -73,9 +73,12 @@ _TCF = 4 # temporal compression factor # --------------------------------------------------------------------------- _SOUND_DIM = 64 _SOUND_LATENT_FPS = 25 +_ACTION_DIM = 64 +_NUM_EMBODIMENT_DOMAINS = 32 -def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bool = False): +def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bool = False, + action_gen: bool = False): """Tiny framework ``Cosmos3VFMNetwork`` with ``unified_3d_mrope``. ``rope_theta`` / ``rope_scaling`` (carrying ``mrope_section`` + @@ -129,6 +132,11 @@ def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bo temporal_compression_factor_sound=1, sound_latent_fps=_SOUND_LATENT_FPS, ) if sound_gen else {} + action_kwargs = dict( + action_gen=True, + action_dim=_ACTION_DIM, + num_embodiment_domains=_NUM_EMBODIMENT_DOMAINS, + ) if action_gen else {} vfm_cfg = Cosmos3VFMNetworkConfig( vision_gen=True, vlm_config=tiny_vlm_cfg, @@ -141,6 +149,7 @@ def _build_tiny_cosmos3_mrope(seed: int = 42, num_layers: int = 2, sound_gen: bo max_latent_t=8, temporal_compression_factor_vision=_TCF, **sound_kwargs, + **action_kwargs, ) torch.manual_seed(seed) lm = Qwen3VLTextForCausalLM(config=mot_cfg)