diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 4a5fd43d4..ddca947cd 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -288,12 +288,20 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch): return self class _StubDiffusersScheduler: - config = SimpleNamespace(flow_shift=1.0) + from_config_calls: list[dict[str, Any]] = [] + + def __init__(self, *, flow_shift: float = 1.0) -> None: + self.config = SimpleNamespace(flow_shift=flow_shift) @classmethod def from_pretrained(cls, *args, **kwargs): return cls() + @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))) + class _StubVideoProcessor: def __init__(self, *args, **kwargs) -> None: pass @@ -303,6 +311,7 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", _StubDiffusersScheduler) monkeypatch.setattr(pipeline_cosmos3, "VideoProcessor", _StubVideoProcessor) monkeypatch.setattr(pipeline_cosmos3, "get_local_device", lambda: torch.device("cpu")) + return _StubDiffusersScheduler def _make_od_config(*, sound_gen: bool) -> SimpleNamespace: @@ -344,6 +353,21 @@ def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_in assert not hasattr(pipeline.transformer, "audio_proj_out") +def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeline_init) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + od_config = _make_od_config(sound_gen=False) + od_config.flow_shift = 2.5 + + pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config) + + 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 pipeline._engine_init_flow_shift == 2.5 + + def test_pipeline_init_passes_tokenizer_attrs_into_transformer( stub_real_pipeline_init, monkeypatch: pytest.MonkeyPatch, @@ -1107,7 +1131,18 @@ class TestForwardRouting: def _install_forward_stubs(self, pipeline): captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []} - def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): + def fake_format( + prompt, + negative_prompt, + num_frames, + frame_rate, + height, + width, + max_sequence_length, + sp, + use_system_prompt=False, + is_t2i=False, + ): captured["format"] = { "prompt": prompt, "negative_prompt": negative_prompt, @@ -1115,7 +1150,8 @@ class TestForwardRouting: "frame_rate": frame_rate, "height": height, "width": width, - "is_t2i": kwargs["is_t2i"], + "use_system_prompt": use_system_prompt, + "is_t2i": is_t2i, } return _ids(2), _mask(), _ids(1), _mask() @@ -1145,12 +1181,26 @@ class TestForwardRouting: ( {"prompt": "A painted robot", "modalities": ["image"]}, make_sampling_params(num_outputs_per_prompt=2), - {"key": "image", "is_t2i": True, "flow": [3.0], "steps": [50, 50], "frames": 1}, + { + "key": "image", + "is_t2i": True, + "use_system_prompt": True, + "flow": [3.0], + "steps": [50, 50], + "frames": 1, + }, ), ( "A warehouse robot", make_sampling_params(), - {"key": "video", "is_t2i": False, "flow": [1.0], "steps": [35], "frames": 189}, + { + "key": "video", + "is_t2i": False, + "use_system_prompt": True, + "flow": [1.0], + "steps": [35], + "frames": 189, + }, ), ], ) @@ -1168,6 +1218,7 @@ 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 [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"] @@ -1196,6 +1247,7 @@ 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) @@ -1222,6 +1274,7 @@ class TestForwardRouting: ) assert captured["flow_shifts"][-1] == 10.0 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 @@ -1237,6 +1290,7 @@ 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) @@ -1262,6 +1316,7 @@ 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 @@ -1337,6 +1392,7 @@ class TestForwardRouting: "frame_rate": 15.0, "height": 16, "width": 16, + "use_system_prompt": False, "is_t2i": False, } assert "flow_shifts" not in captured diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 78df8f70f..b3dace893 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -129,8 +129,9 @@ COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolutio COSMOS3_INVERSE_DURATION_TEMPLATE = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." COSMOS3_INVERSE_RESOLUTION_TEMPLATE = "This video is not of {height}x{width} resolution." COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{width} resolution." -COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a given prompt." -COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a given prompt." +# NOTE: Intentional typo in "give" instead of "given" to match training setup. +COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a give prompt." +COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a give prompt." COSMOS3_T2V_DEFAULT_HEIGHT = 720 COSMOS3_T2V_DEFAULT_WIDTH = 1280 @@ -732,8 +733,16 @@ class Cosmos3OmniDiffusersPipeline( subfolder="scheduler", local_files_only=local_files_only, ) - if od_config.flow_shift is not None: - self.scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config, flow_shift=od_config.flow_shift) + 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, + ) self._cpu_scheduler_state() # --- Video processor for post-decode --- @@ -1410,7 +1419,9 @@ class Cosmos3OmniDiffusersPipeline( target = float(target_shift) if target == float(self._current_flow_shift): return - self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, flow_shift=target) + self.scheduler = UniPCMultistepScheduler.from_config( + self._base_scheduler_config, flow_shift=target, use_karras_sigmas=False + ) self._cpu_scheduler_state() self._current_flow_shift = target @@ -2909,7 +2920,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", False)) + use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", not action_enabled)) if action_enabled and action_video_tensor is None: extra_action_video = self._get_sp_param(sp, "action_video", None)