diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index f2527bd99..105b012bb 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -1150,6 +1150,72 @@ def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None: assert cfg.fps == sp.resolved_frame_rate == 12.0 +def test_transfer_vae_executor_requires_distributed_vae() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + executor = object() + pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: True) + assert pipeline._transfer_vae_executor() is executor + + pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: False) + assert pipeline._transfer_vae_executor() is None + + +def test_sync_transfer_overlap_slices_output_rank_and_broadcasts() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + class RecordingExecutor: + rank = 0 + + def __init__(self) -> None: + self.value = None + + def broadcast_tensor(self, value): + self.value = value + return value + + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + output = torch.arange(1 * 3 * 5 * 2 * 2).reshape(1, 3, 5, 2, 2) + executor = RecordingExecutor() + + overlap = pipeline._sync_transfer_overlap( + output, + overlap_frames=2, + reference_video=output, + vae_executor=executor, + ) + + assert overlap is not None + torch.testing.assert_close(overlap, output[:, :, -2:]) + assert executor.value is overlap + + +def test_sync_transfer_overlap_allocates_receive_buffer_on_non_output_rank() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + class ReceivingExecutor: + rank = 1 + + def broadcast_tensor(self, value): + assert value.shape == (1, 3, 2, 4, 5) + return torch.ones_like(value) + + pipeline = SimpleNamespace(device=torch.device("cpu"), vae=SimpleNamespace(dtype=torch.float32)) + reference = torch.zeros(1, 3, 6, 4, 5) + + overlap = Cosmos3OmniDiffusersPipeline._sync_transfer_overlap( + pipeline, + torch.empty(0), + overlap_frames=2, + reference_video=reference, + vae_executor=ReceivingExecutor(), + ) + + assert overlap is not None + assert torch.equal(overlap, torch.ones_like(overlap)) + + def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transfer diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 744c12bcc..2220c88f9 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -3215,6 +3215,46 @@ class Cosmos3OmniDiffusersPipeline( self.transformer.reset_cache() return latents + def _transfer_vae_executor(self) -> Any | None: + """Return the active distributed VAE executor, if any.""" + executor = getattr(self.vae, "distributed_executor", None) + is_distributed_enabled = getattr(self.vae, "is_distributed_enabled", None) + if executor is None or not callable(is_distributed_enabled) or not is_distributed_enabled(): + return None + return executor + + def _sync_transfer_overlap( + self, + output_video: torch.Tensor, + *, + overlap_frames: int, + reference_video: torch.Tensor, + vae_executor: Any | None, + ) -> torch.Tensor | None: + """Broadcast only the decoded frames needed to condition the next chunk.""" + if overlap_frames <= 0: + return None + + is_output_rank = vae_executor is None or vae_executor.rank == 0 + if is_output_rank: + overlap = output_video[:, :, -overlap_frames:].contiguous() + else: + overlap = torch.empty( + ( + reference_video.shape[0], + reference_video.shape[1], + overlap_frames, + reference_video.shape[3], + reference_video.shape[4], + ), + device=self.device, + dtype=self.vae.dtype, + ) + + if vae_executor is not None: + overlap = vae_executor.broadcast_tensor(overlap) + return overlap + def _forward_transfer( self, *, @@ -3354,6 +3394,8 @@ class Cosmos3OmniDiffusersPipeline( output_chunks: list[torch.Tensor] = [] control_chunks_per_hint: dict[str, list[torch.Tensor]] = {key: [] for key in per_hint_frames} previous_output: torch.Tensor | None = None + vae_executor = self._transfer_vae_executor() + is_output_rank = vae_executor is None or vae_executor.rank == 0 for chunk_id in range(num_chunks): start_frame = chunk_id * stride @@ -3448,16 +3490,26 @@ class Cosmos3OmniDiffusersPipeline( generator=generator, ) output_video = self._decode_latents(latents).clamp(-1, 1) - previous_output = output_video + if chunk_id + 1 < num_chunks: + previous_output = self._sync_transfer_overlap( + output_video, + overlap_frames=min(transfer_config.num_conditional_frames, chunk_frames), + reference_video=target_norm, + vae_executor=vae_executor, + ) - if chunk_id == 0: - output_chunks.append(output_video) - for key, control in control_norms.items(): - control_chunks_per_hint[key].append(control) - else: - output_chunks.append(output_video[:, :, current_conditional_frames:]) - for key, control in control_norms.items(): - control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:]) + if is_output_rank: + if chunk_id == 0: + output_chunks.append(output_video) + for key, control in control_norms.items(): + control_chunks_per_hint[key].append(control) + else: + output_chunks.append(output_video[:, :, current_conditional_frames:]) + for key, control in control_norms.items(): + control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:]) + + if not is_output_rank: + return DiffusionOutput(output={"video": output_video}, custom_output={"fps": frame_rate}) full_output = torch.cat(output_chunks, dim=2)[:, :, :total_frames] full_controls = {