Cosmos3 transfer fix (#5614)

Signed-off-by: Maciej Bala <mbala@nvidia.com>

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1019,11 +1019,19 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest
     cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": True}))
     assert cfg is not None
     assert list(cfg.hints) == ["edge"]
+    assert cfg.hints["edge"].control_weight == 1.0
+    assert cfg.normalized_control_weights == [1.0]
     assert cfg.guidance_scale == 3.0
     assert cfg.control_guidance == 1.5
     assert cfg.flow_shift == 10.0
     assert cfg.num_video_frames_per_chunk == 93
     assert cfg.share_vision_temporal_positions is True
+    assert cfg.emphasize_control_in_prompt is True
+    prompt_ablation_cfg = transfer.resolve_transfer_config(
+        make_sampling_params(extra_args={"edge": True, "emphasize_control_in_prompt": False})
+    )
+    assert prompt_ablation_cfg is not None
+    assert prompt_ablation_cfg.emphasize_control_in_prompt is False
     # fps omitted (no fps/frame_rate on the sampling params) -> wsm preset default (10) applies.
     defaulted_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"wsm": True}))
     assert defaulted_fps_cfg is not None
@@ -1092,6 +1100,40 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest
     assert "preprocessed_video" not in additional
 
 
+def test_transfer_control_weight_validation_and_normalization() -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transfer
+
+    single = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": 0.4}}))
+    assert single is not None
+    assert single.hints["edge"].control_weight == 0.4
+    assert single.normalized_control_weights == [1.0]
+
+    multiple = transfer.resolve_transfer_config(
+        make_sampling_params(
+            extra_args={
+                "edge": {"control_weight": 1.0},
+                "depth": {"control_weight": 3.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)},
+            }
+        )
+    )
+    assert multiple is not None
+    assert multiple.normalized_control_weights == [0.25, 0.75]
+
+    with pytest.raises(ValueError, match="Unsupported.*weight"):
+        transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"weight": 0.5}}))
+    with pytest.raises(ValueError, match="finite and non-negative"):
+        transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": -0.5}}))
+    with pytest.raises(ValueError, match="positive sum"):
+        transfer.resolve_transfer_config(
+            make_sampling_params(
+                extra_args={
+                    "edge": {"control_weight": 0.0},
+                    "depth": {"control_weight": 0.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)},
+                }
+            )
+        )
+
+
 def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
     """When fps and frame_rate differ, transfer must pick frame_rate -- the same
     precedence as OmniDiffusionSamplingParams.resolved_frame_rate. Uses the real
@@ -1352,6 +1394,63 @@ def test_format_and_tokenize_prompts_applies_video_templates_and_system_override
     ]
 
 
+def test_format_and_tokenize_prompts_applies_transfer_prompt_contract(make_cosmos3_pipeline) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE,
+        COSMOS3_TRANSFER_SYSTEM_PROMPT,
+    )
+
+    pipeline = make_cosmos3_pipeline()
+    calls = _capture_tokenize_calls(pipeline)
+    directive = COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names="edge")
+
+    pipeline._format_and_tokenize_prompts(
+        "A robot",
+        "bad",
+        num_frames=48,
+        frame_rate=24,
+        height=720,
+        width=1280,
+        max_sequence_length=32,
+        sp=SimpleNamespace(
+            extra_args={
+                "use_resolution_template": True,
+                "system_prompt": "request-level system prompt",
+            }
+        ),
+        use_system_prompt=True,
+        is_t2i=False,
+        system_prompt=COSMOS3_TRANSFER_SYSTEM_PROMPT,
+        prompt_suffix=directive,
+        use_duration_template=True,
+        use_resolution_template=True,
+        negative_metadata_mode="same",
+    )
+
+    metadata = "The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution."
+    assert calls[0]["text"] == f"A robot. {metadata} {directive}"
+    assert calls[1]["text"] == f"bad. {metadata}"
+    assert all(call["use_system_prompt"] is True for call in calls)
+    assert all(call["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT for call in calls)
+
+
+def test_format_and_tokenize_prompts_rejects_unknown_negative_metadata_mode(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+
+    with pytest.raises(ValueError, match="negative_metadata_mode"):
+        pipeline._format_and_tokenize_prompts(
+            "A robot",
+            "bad",
+            num_frames=48,
+            frame_rate=24,
+            height=720,
+            width=1280,
+            max_sequence_length=32,
+            sp=SimpleNamespace(extra_args={}),
+            negative_metadata_mode="unexpected",
+        )
+
+
 def test_format_and_tokenize_prompts_uses_image_templates_for_t2i(make_cosmos3_pipeline) -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_T2I_SYSTEM_PROMPT
 
@@ -1740,6 +1839,7 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
         control_guidance=1.5,
         control_guidance_interval=None,
         control_latents=[torch.zeros_like(latents)],
+        control_weights=[1.0],
         shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
         velocity_mask=velocity_mask,
         condition_latents=torch.zeros_like(latents),
@@ -1750,6 +1850,9 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
         (2, False),
         (1, True),
     ]
+    assert pipeline.transformer.calls[0]["kwargs"]["control_weights"] == [1.0]
+    assert "control_weights" not in pipeline.transformer.calls[1]["kwargs"]
+    assert pipeline.transformer.calls[2]["kwargs"]["control_weights"] == [1.0]
     torch.testing.assert_close(result, torch.full_like(latents, 254.0))
 
 
@@ -1857,14 +1960,35 @@ def test_diffuse_transfer_interval_switches_branch_counts(make_cosmos3_pipeline,
     torch.testing.assert_close(result, torch.full_like(latents, 508.0))
 
 
-@pytest.mark.parametrize(("hint_key", "expected_fps"), [("edge", 8.0), ("wsm", 10.0)])
-def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint_key: str, expected_fps: float) -> None:
+@pytest.mark.parametrize(
+    ("hint_key", "emphasize_control", "expected_fps", "negative_prompt"),
+    [
+        ("edge", None, 8.0, None),
+        ("wsm", None, 10.0, ""),
+        ("edge", False, 8.0, "custom negative"),
+    ],
+)
+def test_forward_transfer_uses_transfer_prompt_contract_and_source_fps_except_wsm(
+    make_cosmos3_pipeline,
+    hint_key: str,
+    emphasize_control: bool | None,
+    expected_fps: float,
+    negative_prompt: str | None,
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE,
+        COSMOS3_TRANSFER_SYSTEM_PROMPT,
+    )
+
     pipeline = make_cosmos3_pipeline()
     captured: dict[str, Any] = {}
 
     def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs):
-        del prompt, negative_prompt, num_frames, height, width, args, kwargs
+        del num_frames, height, width, args
+        captured["prompt"] = prompt
+        captured["negative_prompt"] = negative_prompt
         captured["format_frame_rate"] = frame_rate
+        captured["format_kwargs"] = kwargs
         return _ids(2), _mask(), _ids(1), _mask()
 
     def fake_encode(video: torch.Tensor) -> torch.Tensor:
@@ -1901,33 +2025,48 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
     pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta")
 
     control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8)
+    extra_args = {
+        hint_key: {"control": control},
+        "max_frames": 5,
+        "num_video_frames_per_chunk": 5,
+        "show_control_condition": True,
+    }
+    if emphasize_control is not None:
+        extra_args["emphasize_control_in_prompt"] = emphasize_control
+    prompt_data = {
+        "prompt": "transfer",
+        "modalities": ["video"],
+        "additional_information": {
+            "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16),
+            "transfer_input_fps": 8.0,
+        },
+    }
+    if negative_prompt is not None:
+        prompt_data["negative_prompt"] = negative_prompt
     request = SimpleNamespace(
-        prompts=[
-            {
-                "prompt": "transfer",
-                "modalities": ["video"],
-                "additional_information": {
-                    "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16),
-                    "transfer_input_fps": 8.0,
-                },
-            }
-        ],
+        prompts=[prompt_data],
         sampling_params=make_sampling_params(
             height=16,
             width=16,
             # fps omitted (no frame_rate) -> non-wsm uses the source video fps (8), wsm uses its preset (10).
-            extra_args={
-                hint_key: {"control": control},
-                "max_frames": 5,
-                "num_video_frames_per_chunk": 5,
-                "show_control_condition": True,
-            },
+            extra_args=extra_args,
         ),
     )
 
     output = pipeline.forward(request)
 
     assert captured["format_frame_rate"] == expected_fps
+    expected_negative_prompt = "" if negative_prompt is None else negative_prompt
+    assert captured["negative_prompt"] == expected_negative_prompt
+    assert captured["format_kwargs"]["use_system_prompt"] is True
+    assert captured["format_kwargs"]["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT
+    assert captured["format_kwargs"]["use_duration_template"] is True
+    assert captured["format_kwargs"]["use_resolution_template"] is True
+    assert captured["format_kwargs"]["negative_metadata_mode"] == "same"
+    expected_suffix = (
+        None if emphasize_control is False else COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names=hint_key)
+    )
+    assert captured["format_kwargs"]["prompt_suffix"] == expected_suffix
     assert captured["shared_kwargs"]["fps"] == expected_fps
     assert captured["flow_shifts"] == [10.0]
     # Transfer applies the V2V flow shift when building its timestep schedule.

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -504,6 +504,35 @@ def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device(
     assert generator_param.device == accelerator_device
 
 
+def test_multi_control_attention_weights_target_outputs() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention
+
+    class ControlValueAttention:
+        def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+            del k
+            # Each per-control K/V layout is [text, control_i, target].
+            return v[:, 1:2].expand(q.shape[0], q.shape[1], -1, -1)
+
+    attention = SimpleNamespace(multi_control_attn=ControlValueAttention())
+    q = torch.zeros(1, 3, 1, 1)
+    k = torch.zeros_like(q)
+    v = torch.tensor([[[[10.0]], [[20.0]], [[30.0]]]])
+    text_kv = torch.zeros(1, 1, 1, 1)
+
+    output = Cosmos3CrossAttention._forward_multi_control(
+        attention,
+        q,
+        k,
+        v,
+        text_kv,
+        text_kv,
+        control_token_sizes=(1, 1),
+        control_weights=(0.25, 0.75),
+    )
+
+    torch.testing.assert_close(output, torch.tensor([[[10.0], [20.0], [17.5]]]))
+
+
 def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
@@ -521,9 +550,21 @@ def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatc
         video_shape=(1, 2, 2),
         fps=24.0,
         control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)],
+        control_weights=[1.0, 3.0],
     )
 
     assert tuple(output.shape) == tuple(hidden_states.shape)
+    with pytest.raises(ValueError, match="control_weights length"):
+        model(
+            hidden_states=hidden_states,
+            timestep=torch.tensor([1.0]),
+            text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+            text_mask=torch.ones(1, 2, dtype=torch.long),
+            video_shape=(1, 2, 2),
+            fps=24.0,
+            control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)],
+            control_weights=[1.0],
+        )
     with pytest.raises(ValueError, match="control latent shape"):
         model(
             hidden_states=hidden_states,