Fixed Cosmos3 regression (#4627)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -19,9 +19,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
 class StubScheduler:
-    def __init__(self, timesteps: list[int] | None = None, *, flow_shift: float = 1.0) -> None:
+    def __init__(
+        self,
+        timesteps: list[int] | None = None,
+        *,
+        flow_shift: float = 1.0,
+        use_karras_sigmas: bool = True,
+    ) -> None:
         self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
-        self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift)
+        self.config = SimpleNamespace(
+            num_train_timesteps=1000,
+            flow_shift=flow_shift,
+            use_karras_sigmas=use_karras_sigmas,
+        )
         self.set_timesteps_calls: list[tuple[int, torch.device]] = []
         self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
 
@@ -194,6 +204,8 @@ def make_cosmos3_pipeline():
         pipeline._base_scheduler_config = pipeline.scheduler.config
         pipeline._engine_init_flow_shift = 1.0
         pipeline._current_flow_shift = 1.0
+        pipeline._base_scheduler_use_karras_sigmas = True
+        pipeline._current_scheduler_use_karras_sigmas = True
         pipeline._guidance_scale = None
         pipeline._num_timesteps = None
         pipeline._cosmos3_branch_caches = None
@@ -292,8 +304,8 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
     class _StubDiffusersScheduler:
         from_config_calls: list[dict[str, Any]] = []
 
-        def __init__(self, *, flow_shift: float = 1.0) -> None:
-            self.config = SimpleNamespace(flow_shift=flow_shift)
+        def __init__(self, *, flow_shift: float = 1.0, use_karras_sigmas: bool = True) -> None:
+            self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
 
         @classmethod
         def from_pretrained(cls, *args, **kwargs):
@@ -302,7 +314,13 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
         @classmethod
         def from_config(cls, config, **kwargs):
             cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
-            return cls(flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)))
+            return cls(
+                flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
+                use_karras_sigmas=kwargs.get(
+                    "use_karras_sigmas",
+                    getattr(config, "use_karras_sigmas", True),
+                ),
+            )
 
     class _StubVideoProcessor:
         def __init__(self, *args, **kwargs) -> None:
@@ -349,10 +367,13 @@ def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_in
 
     pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
 
+    assert stub_real_pipeline_init.from_config_calls == []
     assert pipeline._sound_tokenizer is None
     assert pipeline.transformer.sound_gen is False
     assert not hasattr(pipeline.transformer, "audio_proj_in")
     assert not hasattr(pipeline.transformer, "audio_proj_out")
+    assert pipeline._base_scheduler_use_karras_sigmas is True
+    assert pipeline._current_scheduler_use_karras_sigmas is True
 
 
 def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeline_init) -> None:
@@ -366,8 +387,52 @@ def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeli
     assert stub_real_pipeline_init.from_config_calls
     call = stub_real_pipeline_init.from_config_calls[-1]
     assert getattr(call["config"], "flow_shift") == 1.0
-    assert call["kwargs"] == {"flow_shift": 2.5, "use_karras_sigmas": False}
+    assert call["kwargs"] == {"flow_shift": 2.5}
     assert pipeline._engine_init_flow_shift == 2.5
+    assert pipeline._base_scheduler_use_karras_sigmas is True
+
+
+def test_set_flow_shift_restores_checkpoint_scheduler_mode(
+    make_cosmos3_pipeline,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+    class RecordingScheduler:
+        from_config_calls: list[dict[str, Any]] = []
+
+        def __init__(self, *, flow_shift: float, use_karras_sigmas: bool) -> None:
+            self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
+
+        @classmethod
+        def from_config(cls, config, **kwargs):
+            cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
+            return cls(
+                flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
+                use_karras_sigmas=kwargs.get(
+                    "use_karras_sigmas",
+                    getattr(config, "use_karras_sigmas", True),
+                ),
+            )
+
+    monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", RecordingScheduler)
+    pipeline = make_cosmos3_pipeline()
+
+    pipeline._set_flow_shift(10.0, use_karras_sigmas=False)
+    assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {
+        "flow_shift": 10.0,
+        "use_karras_sigmas": False,
+    }
+    assert pipeline._current_flow_shift == 10.0
+    assert pipeline._current_scheduler_use_karras_sigmas is False
+
+    pipeline._set_flow_shift(10.0)
+    assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {"flow_shift": 10.0}
+    assert pipeline._current_flow_shift == 10.0
+    assert pipeline._current_scheduler_use_karras_sigmas is True
+
+    pipeline._set_flow_shift(10.0)
+    assert len(RecordingScheduler.from_config_calls) == 2
 
 
 def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
@@ -1022,7 +1087,12 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
     pipeline._encode_video_tensor = fake_encode
     pipeline._prepare_transfer_latents = fake_prepare
     pipeline.diffuse_transfer = fake_diffuse_transfer
-    pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
+
+    def fake_set_flow_shift(target, *, use_karras_sigmas=None):
+        captured.setdefault("flow_shifts", []).append(target)
+        captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
+
+    pipeline._set_flow_shift = fake_set_flow_shift
     pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta")
 
     control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8)
@@ -1055,6 +1125,8 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
     assert captured["format_frame_rate"] == expected_fps
     assert captured["shared_kwargs"]["fps"] == expected_fps
     assert captured["flow_shifts"] == [10.0]
+    # Transfer uses the V2V flow-sigma schedule (karras off) so flow_shift applies.
+    assert captured["scheduler_use_karras_sigmas"] == [False]
     assert output.custom_output["fps"] == expected_fps
     assert output.output["video"].device.type == "meta"
 
@@ -1068,7 +1140,7 @@ def test_forward_transfer_runs_multichunk_overlap_path(
 
     pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16)
     pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
-    pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
+    pipeline._set_flow_shift = lambda target, **_kwargs: captured.setdefault("flow_shifts", []).append(target)
 
     original_prepare = pipeline._prepare_transfer_latents
 
@@ -1213,7 +1285,12 @@ class TestForwardRouting:
 
         pipeline._format_and_tokenize_prompts = fake_format
         pipeline._prepare_latents = fake_prepare
-        pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
+
+        def fake_set_flow_shift(target, *, use_karras_sigmas=None):
+            captured.setdefault("flow_shifts", []).append(target)
+            captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
+
+        pipeline._set_flow_shift = fake_set_flow_shift
         pipeline.diffuse = fake_diffuse
         pipeline._decode_latents = lambda latents: latents
         return captured
@@ -1227,7 +1304,6 @@ class TestForwardRouting:
                 {
                     "key": "image",
                     "is_t2i": True,
-                    "use_system_prompt": True,
                     "flow": [3.0],
                     "steps": [50, 50],
                     "frames": 1,
@@ -1239,7 +1315,6 @@ class TestForwardRouting:
                 {
                     "key": "video",
                     "is_t2i": False,
-                    "use_system_prompt": True,
                     "flow": [1.0],
                     "steps": [35],
                     "frames": 189,
@@ -1261,9 +1336,9 @@ class TestForwardRouting:
 
         assert expected["key"] in output.output
         assert captured["format"]["is_t2i"] is expected["is_t2i"]
-        assert captured["format"]["use_system_prompt"] is expected["use_system_prompt"]
         assert captured["format"]["num_frames"] == expected["frames"]
         assert captured["flow_shifts"] == expected["flow"]
+        assert captured["scheduler_use_karras_sigmas"] == [None]
         assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
 
     def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
@@ -1290,7 +1365,6 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
-        assert captured["format"]["use_system_prompt"] is True
 
         video_tensor = torch.zeros(1, 3, 5, 16, 16)
         v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
@@ -1316,8 +1390,8 @@ class TestForwardRouting:
             )
         )
         assert captured["flow_shifts"][-1] == 10.0
+        assert captured["scheduler_use_karras_sigmas"][-1] is False
         assert captured["format"]["negative_prompt"] == ""
-        assert captured["format"]["use_system_prompt"] is True
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
         assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
 
@@ -1333,7 +1407,6 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
-        assert captured["format"]["use_system_prompt"] is True
         assert output.output["audio_sample_rate"] == 10
 
         pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
@@ -1359,7 +1432,6 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
-        assert captured["format"]["use_system_prompt"] is False
         assert output.custom_output["action"].shape == (1, 2, 2)
         assert "action_only_output" not in output.custom_output
 

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -733,16 +733,11 @@ class Cosmos3OmniDiffusersPipeline(
             subfolder="scheduler",
             local_files_only=local_files_only,
         )
-        init_flow_shift = (
-            od_config.flow_shift
-            if od_config.flow_shift is not None
-            else getattr(self.scheduler.config, "flow_shift", 1.0)
-        )
-        self.scheduler = UniPCMultistepScheduler.from_config(
-            self.scheduler.config,
-            flow_shift=init_flow_shift,
-            use_karras_sigmas=False,
-        )
+        if od_config.flow_shift is not None:
+            self.scheduler = UniPCMultistepScheduler.from_config(
+                self.scheduler.config,
+                flow_shift=od_config.flow_shift,
+            )
         self._cpu_scheduler_state()
 
         # --- Video processor for post-decode ---
@@ -766,6 +761,8 @@ class Cosmos3OmniDiffusersPipeline(
         self._base_scheduler_config = self.scheduler.config
         self._engine_init_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0)
         self._current_flow_shift = self._engine_init_flow_shift
+        self._base_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas(self.scheduler.config)
+        self._current_scheduler_use_karras_sigmas = self._base_scheduler_use_karras_sigmas
 
         self._guidance_scale = None
         self._num_timesteps = None
@@ -1406,24 +1403,38 @@ class Cosmos3OmniDiffusersPipeline(
             raise ValueError(f"Incorrect modality value in {modalities}, expected one of {accepted_modalities}.")
         return "image" in modalities
 
-    def _set_flow_shift(self, target_shift: float) -> None:
-        """Set the UniPC ``flow_shift`` to a concrete target value.
+    @staticmethod
+    def _scheduler_use_karras_sigmas(config: Any) -> bool | None:
+        value = getattr(config, "use_karras_sigmas", None)
+        return None if value is None else bool(value)
+
+    def _set_flow_shift(self, target_shift: float, *, use_karras_sigmas: bool | None = None) -> None:
+        """Set UniPC scheduler mode for a concrete request.
 
         The scheduler is rebuilt from the saved base config if
-        the target differs from the current shift.  Tracking
-        ``self._current_flow_shift`` explicitly is required because the
-        previous mode may have rebuilt the scheduler - we cannot rely on
-        ``self.scheduler.config.flow_shift`` reflecting the last requested
-        target if a rebuild was skipped via the equality check.
+        the target differs from the current shift or Karras-sigma mode.
+        Tracking explicit scheduler state is required because the previous
+        mode may have rebuilt the scheduler - we cannot rely on
+        ``self.scheduler.config`` reflecting the last requested target if a
+        rebuild was skipped via the equality check.
         """
         target = float(target_shift)
-        if target == float(self._current_flow_shift):
-            return
-        self.scheduler = UniPCMultistepScheduler.from_config(
-            self._base_scheduler_config, flow_shift=target, use_karras_sigmas=False
+        target_use_karras_sigmas = (
+            self._base_scheduler_use_karras_sigmas if use_karras_sigmas is None else bool(use_karras_sigmas)
         )
+        if (
+            target == float(self._current_flow_shift)
+            and target_use_karras_sigmas == self._current_scheduler_use_karras_sigmas
+        ):
+            return
+
+        scheduler_kwargs: dict[str, Any] = {"flow_shift": target}
+        if use_karras_sigmas is not None:
+            scheduler_kwargs["use_karras_sigmas"] = bool(use_karras_sigmas)
+        self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, **scheduler_kwargs)
         self._cpu_scheduler_state()
         self._current_flow_shift = target
+        self._current_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas(self.scheduler.config)
 
     def _cpu_scheduler_state(self) -> None:
         # We need to move scheduler tensors to CPU, as unipc from diffusers assumes they are on CPU.
@@ -2619,7 +2630,7 @@ class Cosmos3OmniDiffusersPipeline(
 
         self._guidance_scale = guidance_scale
         self._num_timesteps = num_inference_steps
-        self._set_flow_shift(flow_shift_target)
+        self._set_flow_shift(flow_shift_target, use_karras_sigmas=False)
 
         generator = sp.generator
         if generator is None:
@@ -2918,7 +2929,7 @@ class Cosmos3OmniDiffusersPipeline(
             self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
             or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH
         )
-        use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", not action_enabled))
+        use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", is_v2v))
 
         if action_enabled and action_video_tensor is None:
             extra_action_video = self._get_sp_param(sp, "action_video", None)
@@ -2942,7 +2953,7 @@ class Cosmos3OmniDiffusersPipeline(
 
         # Always resolve to a concrete target shift for this request, then
         # update the shared Diffusers scheduler.
-        self._set_flow_shift(flow_shift_target)
+        self._set_flow_shift(flow_shift_target, use_karras_sigmas=False if is_v2v else None)
 
         generator = sp.generator
         if generator is None: