From 3b335814cc5aa9cd836a95e6b97edb5ddf9c3193 Mon Sep 17 00:00:00 2001
From: Rahul Steiger <rsteiger@nvidia.com>
Date: Tue, 8 Sep 2026 01:34:08 +0200
Subject: [PATCH] [Bugfix][Cosmos3] Fix distributed Transfer output envelope
 (#7205)

Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Co-authored-by: Nils Hoffmann <nhoffmann@nvidia.com>
---
 .../models/cosmos3/test_cosmos3_pipeline.py   | 36 +++++++++++++++++++
 .../models/cosmos3/pipeline_cosmos3.py        |  7 +++-
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index 0b2ed3a57..72e4ab962 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
@@ -2401,6 +2401,42 @@ def test_forward_transfer_runs_multichunk_overlap_path(
     torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2))
 
 
+def test_forward_transfer_non_output_rank_uses_canonical_envelope(
+    make_cosmos3_pipeline,
+    sequential_cfg_parallel,
+) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.vae.distributed_executor = SimpleNamespace(rank=1)
+    pipeline.vae.is_distributed_enabled = lambda: True
+    pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1")
+    pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
+    pipeline._set_flow_shift = lambda *_args, **_kwargs: None
+    decoded = torch.zeros(1, 3, 1, 16, 16)
+    pipeline._decode_latents = lambda latents: decoded
+
+    request = SimpleNamespace(
+        prompts=[{"prompt": "transfer", "modalities": ["video"]}],
+        sampling_params=make_sampling_params(
+            height=16,
+            width=16,
+            num_inference_steps=1,
+            guidance_scale=1.0,
+            extra_args={
+                "edge": {"control": torch.zeros(3, 1, 16, 16, dtype=torch.uint8)},
+                "max_frames": 1,
+                "num_video_frames_per_chunk": 1,
+            },
+        ),
+    )
+
+    output = pipeline.forward(request)
+
+    assert set(output.output) == {"payload", "metadata"}
+    assert set(output.output["payload"]) == {"video"}
+    torch.testing.assert_close(output.output["payload"]["video"], decoded)
+    assert output.output["metadata"] == {"video": {"fps": 24.0}}
+
+
 def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
     """With cache-dit active the uncond pass is kept even outside the guidance
     interval (so cache-dit's has_separate_cfg parity stays in phase), and the
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index a171e2d3b..49a4e4027 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -3522,7 +3522,12 @@ class Cosmos3OmniDiffusersPipeline(
                         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})
+            return DiffusionOutput(
+                output={
+                    "payload": {"video": output_video},
+                    "metadata": {"video": {"fps": frame_rate}},
+                },
+            )
 
         full_output = torch.cat(output_chunks, dim=2)[:, :, :total_frames]
         full_controls = {