diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index ba482e793..1085e3878 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -36,20 +36,36 @@ class StubScheduler: 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, - use_karras_sigmas=use_karras_sigmas, ) - self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + self.set_timesteps_calls: list[dict[str, Any]] = [] self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] - def set_timesteps(self, num_steps: int, device: torch.device) -> None: - self.set_timesteps_calls.append((num_steps, device)) - self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device) + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + *, + shift: float | None = None, + sigmas: list[float] | None = None, + ) -> None: + self.set_timesteps_calls.append( + { + "num_inference_steps": num_inference_steps, + "device": device, + "shift": shift, + "sigmas": sigmas, + } + ) + if sigmas is not None: + self.timesteps = torch.tensor(sigmas, device=device) + else: + assert num_inference_steps is not None + self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device) def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs): del kwargs @@ -200,6 +216,7 @@ def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch): def make_cosmos3_pipeline(): def _make(): from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT, Cosmos3OmniDiffusersPipeline, ) @@ -213,11 +230,10 @@ def make_cosmos3_pipeline(): pipeline.vae_scale_factor_temporal = 4 pipeline.vae_scale_factor_spatial = 8 pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0) - 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._engine_init_flow_shift = COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT + pipeline._current_flow_shift = COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT + pipeline.is_distilled_model = False + pipeline.is_edge_model = False pipeline._guidance_scale = None pipeline._num_timesteps = None pipeline._cosmos3_branch_caches = None @@ -243,6 +259,7 @@ def make_sampling_params(**overrides: Any) -> SimpleNamespace: "num_frames": None, "num_inference_steps": None, "guidance_scale": None, + "guidance_scale_provided": False, "generator": None, "seed": 123, "num_outputs_per_prompt": 1, @@ -288,6 +305,165 @@ def _mask() -> torch.Tensor: return torch.ones(1, 1, dtype=torch.long) +def _capture_tokenize_calls(pipeline: Any) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + + def _tokenize( + text: str, + max_sequence_length: int, + use_system_prompt: bool = False, + system_prompt: str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + index = len(calls) + 1 + calls.append( + { + "text": text, + "max_sequence_length": max_sequence_length, + "use_system_prompt": use_system_prompt, + "system_prompt": system_prompt, + } + ) + return _ids(index), torch.full((1, 1), index, dtype=torch.long) + + pipeline._tokenize_prompt = _tokenize + return calls + + +@pytest.mark.parametrize( + ("provided", "value", "default", "is_distilled", "expected"), + [ + (False, 1.0, 7.0, False, 7.0), + (True, 1.0, 7.0, False, 1.0), + (True, 4.5, 7.0, False, 4.5), + (False, 1.0, 7.0, True, 1.0), + (True, 4.5, 7.0, True, 1.0), + ], +) +def test_resolve_guidance_scale( + make_cosmos3_pipeline, + provided: bool, + value: float, + default: float, + is_distilled: bool, + expected: float, +) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = is_distilled + sp = make_sampling_params( + guidance_scale=value, + guidance_scale_provided=provided, + ) + + assert pipeline._resolve_guidance_scale(sp, default) == expected + + +def test_distilled_generation_accepts_t2i_and_i2v(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + common = { + "action_enabled": False, + "transfer_config": None, + "is_v2v": False, + "sound_enabled": False, + } + + pipeline._validate_distilled_generation_mode( + is_t2i=True, + image_tensor=None, + **common, + ) + pipeline._validate_distilled_generation_mode( + is_t2i=False, + image_tensor=torch.zeros(1), + **common, + ) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"action_enabled": True}, "action requests are unsupported"), + ({"transfer_config": SimpleNamespace()}, "transfer requests are unsupported"), + ({"is_v2v": True}, "video-to-video requests are unsupported"), + ({"sound_enabled": True}, "sound generation is unsupported"), + ( + {"is_t2i": False, "image_tensor": None}, + "text-to-video requests are unsupported", + ), + ( + {"is_t2i": True, "image_tensor": torch.zeros(1)}, + "image-conditioned image generation is unsupported", + ), + ], +) +def test_distilled_generation_rejects_unsupported_modes( + make_cosmos3_pipeline, + overrides: dict[str, Any], + message: str, +) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + kwargs = { + "is_t2i": True, + "image_tensor": None, + "action_enabled": False, + "transfer_config": None, + "is_v2v": False, + "sound_enabled": False, + } + kwargs.update(overrides) + + with pytest.raises(ValueError, match=message): + pipeline._validate_distilled_generation_mode(**kwargs) + + +def test_distilled_generation_rejects_robolab_policy(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + + with pytest.raises(ValueError, match="RoboLab/action policy requests are unsupported"): + pipeline._forward_robolab_policy(make_sampling_params(), None, 0.0) + + +@pytest.mark.parametrize( + ("prompt", "sampling_params", "message"), + [ + ( + {"prompt": "x", "modalities": ["video"], "generate_sound": True}, + make_sampling_params(), + "do not support sound generation", + ), + ( + {"prompt": "x", "modalities": ["video"]}, + make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}), + "do not support transfer inference", + ), + ( + { + "prompt": "x", + "modalities": ["video"], + "additional_information": { + "preprocessed_video": torch.zeros(1, 3, 5, 16, 16), + }, + }, + make_sampling_params(height=16, width=16, num_frames=5), + "do not support video-to-video generation", + ), + ], +) +def test_edge_forward_rejects_unsupported_modes( + make_cosmos3_pipeline, + prompt: dict[str, Any], + sampling_params: SimpleNamespace, + message: str, +) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_edge_model = True + + with pytest.raises(ValueError, match=message): + pipeline.forward(SimpleNamespace(prompts=[prompt], sampling_params=sampling_params)) + + def test_pipeline_registered_and_exported() -> None: from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS from vllm_omni.diffusion.models import cosmos3 @@ -303,17 +479,45 @@ def test_pipeline_registered_and_exported() -> None: assert issubclass(Cosmos3OmniDiffusersPipeline, nn.Module) assert issubclass(Cosmos3OmniDiffusersPipeline, ProgressBarMixin) assert Cosmos3OmniDiffusersPipeline.support_image_input is True - assert _DIFFUSION_MODELS["Cosmos3OmniDiffusersPipeline"] == ( - "cosmos3", - "pipeline_cosmos3", - "Cosmos3OmniDiffusersPipeline", - ) - assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func" - assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func" - assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func" - assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__ + for pipeline_name in ("Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"): + assert _DIFFUSION_MODELS[pipeline_name] == ( + "cosmos3", + "pipeline_cosmos3", + "Cosmos3OmniDiffusersPipeline", + ) + assert _DIFFUSION_PRE_PROCESS_FUNCS[pipeline_name] == "get_cosmos3_pre_process_func" + assert _DIFFUSION_POST_PROCESS_FUNCS[pipeline_name] == "get_cosmos3_post_process_func" + assert _DIFFUSION_IR_OP_PRIORITY_FUNCS[pipeline_name] == "get_cosmos3_ir_op_priority_func" + assert pipeline_name in CUSTOM_DIT_ENABLERS + + +@pytest.mark.parametrize( + "pipeline_name", + ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"], +) +def test_cosmos3_model_index_resolves_pipeline( + tmp_path, + pipeline_name: str, +) -> None: + import json + + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": pipeline_name})) + + config = OmniDiffusionConfig(model=str(tmp_path)) + config.enrich_config() + + assert config.model_class_name == pipeline_name + resolved_pipeline_cls = DiffusionModelRegistry._try_load_model_cls(config.model_class_name) + assert resolved_pipeline_cls is Cosmos3OmniDiffusersPipeline + @pytest.fixture def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch): @@ -335,25 +539,24 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch): return self class _StubDiffusersScheduler: + load_config_calls: list[dict[str, Any]] = [] from_config_calls: list[dict[str, Any]] = [] - 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) + 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() + def load_config(cls, *args, **kwargs): + cls.load_config_calls.append({"args": args, "kwargs": dict(kwargs)}) + return { + "_class_name": "UniPCMultistepScheduler", + "flow_shift": 1.0, + } @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), - ), - ) + return cls(flow_shift=kwargs.get("shift", config.get("flow_shift", 1.0))) class _StubVideoProcessor: def __init__(self, *args, **kwargs) -> None: @@ -361,13 +564,18 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(pipeline_cosmos3, "AutoTokenizer", _StubAutoTokenizer) monkeypatch.setattr(pipeline_cosmos3, "DistributedAutoencoderKLWan", _StubDiffusersVAE) - monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", _StubDiffusersScheduler) + monkeypatch.setattr(pipeline_cosmos3, "FlowUniPCMultistepScheduler", _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: +def _make_od_config( + *, + sound_gen: bool, + tf_model_config_overrides: dict[str, Any] | None = None, + model_config: dict[str, Any] | None = None, +) -> SimpleNamespace: tf_model_config = { "hidden_size": 8, "num_hidden_layers": 0, @@ -382,6 +590,8 @@ def _make_od_config(*, sound_gen: bool) -> SimpleNamespace: } if sound_gen: tf_model_config["sound_gen"] = True + if tf_model_config_overrides: + tf_model_config.update(tf_model_config_overrides) return SimpleNamespace( enable_cpu_offload=False, enable_diffusion_pipeline_profiler=False, @@ -390,7 +600,7 @@ def _make_od_config(*, sound_gen: bool) -> SimpleNamespace: flow_shift=None, quantization_config=None, custom_pipeline_args={}, - model_config={}, + model_config=model_config or {}, tf_model_config=tf_model_config, ) @@ -400,16 +610,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: +def test_pipeline_init_uses_flow_unipc_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) @@ -417,55 +624,302 @@ def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeli 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} + assert len(stub_real_pipeline_init.load_config_calls) == 1 + assert len(stub_real_pipeline_init.from_config_calls) == 1 + call = stub_real_pipeline_init.from_config_calls[0] + assert call["config"]["_class_name"] == "UniPCMultistepScheduler" + assert call["kwargs"] == { + "shift": 1.0, + "use_dynamic_shifting": False, + "prediction_type": "flow_prediction", + } + assert pipeline.is_distilled_model is False 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, +@pytest.mark.parametrize( + ("scheduler_class_name", "expected_distilled"), + [ + ("FlowMatchEulerDiscreteScheduler", True), + ("UniPCMultistepScheduler", False), + ], +) +def test_pipeline_resolves_scheduler_class_from_checkpoint_file( + tmp_path, + stub_real_pipeline_init, monkeypatch: pytest.MonkeyPatch, + scheduler_class_name: str, + expected_distilled: bool, ) -> None: + import json + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler as RealFlowUniPCMultistepScheduler, + ) - class RecordingScheduler: - from_config_calls: list[dict[str, Any]] = [] + t_list = [1.0, 0.75, 0.5, 0.25] + scheduler_dir = tmp_path / "scheduler" + scheduler_dir.mkdir() + scheduler_config = {"_class_name": scheduler_class_name} + if expected_distilled: + scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list} + (scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config)) - def __init__(self, *, flow_shift: float, use_karras_sigmas: bool) -> None: - self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas) + class StubFlowUniPCScheduler(StubScheduler): + from_config_calls: list[tuple[Any, dict[str, Any]]] = [] + + @classmethod + def load_config(cls, *args, **kwargs): + return RealFlowUniPCMultistepScheduler.load_config(*args, **kwargs) @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), - ), - ) + cls.from_config_calls.append((config, dict(kwargs))) + return cls() - monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", RecordingScheduler) - pipeline = make_cosmos3_pipeline() + class StubFlowMatchScheduler(StubScheduler): + from_config_calls: list[tuple[Any, dict[str, Any]]] = [] - 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, + @classmethod + def from_config(cls, config, **kwargs): + cls.from_config_calls.append((config, dict(kwargs))) + scheduler = cls() + scheduler.config.fixed_step_sampler_config = config["fixed_step_sampler_config"] + return scheduler + + monkeypatch.setattr(pipeline_cosmos3, "FlowUniPCMultistepScheduler", StubFlowUniPCScheduler) + monkeypatch.setattr( + pipeline_cosmos3, + "FlowMatchEulerDiscreteScheduler", + StubFlowMatchScheduler, + ) + + od_config = _make_od_config(sound_gen=False) + od_config.model = str(tmp_path) + pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config) + + assert pipeline.is_distilled_model is expected_distilled + if expected_distilled: + assert len(StubFlowMatchScheduler.from_config_calls) == 1 + assert StubFlowMatchScheduler.from_config_calls[0][1] == {"stochastic_sampling": True} + assert StubFlowUniPCScheduler.from_config_calls == [] + assert pipeline._scheduler_init_t_list == t_list + else: + assert len(StubFlowUniPCScheduler.from_config_calls) == 1 + assert StubFlowMatchScheduler.from_config_calls == [] + assert not hasattr(pipeline, "_scheduler_init_t_list") + + +def test_distilled_pipeline_initializes_sde_scheduler( + stub_real_pipeline_init, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from diffusers import FlowMatchEulerDiscreteScheduler + + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + + t_list = [1.0, 0.9375, 0.8333333333333334, 0.625] + scheduler_config = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "shift": 1.0, + "stochastic_sampling": False, + "fixed_step_sampler_config": { + "sample_type": "sde", + "t_list": t_list, + }, } - assert pipeline._current_flow_shift == 10.0 - assert pipeline._current_scheduler_use_karras_sigmas is False + monkeypatch.setattr( + stub_real_pipeline_init, + "load_config", + classmethod(lambda cls, *args, **kwargs: scheduler_config), + ) - 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 = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False)) + assert pipeline.is_distilled_model is True + assert isinstance(pipeline.scheduler, FlowMatchEulerDiscreteScheduler) + assert pipeline.scheduler.config.stochastic_sampling is True + assert pipeline._scheduler_init_t_list == t_list - pipeline._set_flow_shift(10.0) - assert len(RecordingScheduler.from_config_calls) == 2 + +@pytest.mark.parametrize( + ("fixed_step_config", "error_pattern"), + [ + pytest.param( + {"t_list": [1.0, 0.5]}, + r"fixed_step_sampler_config\.sample_type=sde", + id="missing-sample-type", + ), + pytest.param( + {"sample_type": "ode", "t_list": [1.0, 0.5]}, + r"fixed_step_sampler_config\.sample_type=sde", + id="unsupported-sample-type", + ), + pytest.param( + {"sample_type": "sde"}, + r"non-empty fixed_step_sampler_config\.t_list", + id="missing-t-list", + ), + pytest.param( + {"sample_type": "sde", "t_list": []}, + r"non-empty fixed_step_sampler_config\.t_list", + id="empty-t-list", + ), + ], +) +def test_distilled_scheduler_validates_fixed_step_config( + stub_real_pipeline_init, + monkeypatch: pytest.MonkeyPatch, + fixed_step_config: dict[str, Any], + error_pattern: str, +) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + + scheduler_config = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "fixed_step_sampler_config": fixed_step_config, + } + monkeypatch.setattr( + stub_real_pipeline_init, + "load_config", + classmethod(lambda cls, *args, **kwargs: scheduler_config), + ) + + with pytest.raises(ValueError, match=error_pattern): + Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False)) + + +def test_pipeline_init_selects_edge_transformer_from_backbone_type(stub_real_pipeline_init) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT, + Cosmos3OmniDiffusersPipeline, + ) + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + COSMOS3_EDGE_BACKBONE_TYPE, + Cosmos3EdgeVFMTransformer, + ) + + pipeline = Cosmos3OmniDiffusersPipeline( + od_config=_make_od_config( + sound_gen=False, + tf_model_config_overrides={ + "backbone_type": COSMOS3_EDGE_BACKBONE_TYPE, + "qk_norm_for_text": False, + "latent_channel": 48, + "latent_patch_size": 2, + "temporal_compression_factor": 4, + "layer_norm_epsilon": 1e-5, + }, + ) + ) + + assert isinstance(pipeline.transformer, Cosmos3EdgeVFMTransformer) + assert pipeline.is_edge_model is True + assert pipeline._engine_init_flow_shift == COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT + + +def test_pipeline_init_does_not_select_edge_from_model_index_class_name(stub_real_pipeline_init) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer + + pipeline = Cosmos3OmniDiffusersPipeline( + od_config=_make_od_config( + sound_gen=False, + model_config={"_class_name": "Cosmos3EdgeOmniDiffusersPipeline"}, + ) + ) + + assert isinstance(pipeline.transformer, Cosmos3VFMTransformer) + assert not isinstance(pipeline.transformer, Cosmos3EdgeVFMTransformer) + + +def test_flow_unipc_reuses_scheduler_and_forwards_each_request_shift(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + scheduler = pipeline.scheduler + + assert pipeline._engine_init_flow_shift == 10.0 + assert pipeline._current_flow_shift == 10.0 + assert scheduler.config.flow_shift == 1.0 + + for shift in (3.0, 10.0): + pipeline._set_flow_shift(shift) + pipeline._set_timesteps(4, torch.device("cpu"), shift=pipeline._current_flow_shift) + + assert pipeline.scheduler is scheduler + assert pipeline._current_flow_shift == 10.0 + assert [call["shift"] for call in scheduler.set_timesteps_calls] == [3.0, 10.0] + assert all(call["num_inference_steps"] == 4 for call in scheduler.set_timesteps_calls) + assert all(call["sigmas"] is None for call in scheduler.set_timesteps_calls) + + +def test_flow_unipc_reproducible_with_same_seed(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, + ) + + pipeline = make_cosmos3_pipeline() + pipeline.scheduler = FlowUniPCMultistepScheduler( + shift=1.0, + use_dynamic_shifting=False, + prediction_type="flow_prediction", + ) + pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: ( + _ids(1), + _mask(), + _ids(0), + _mask(), + ) + + def run(seed: int) -> torch.Tensor: + request = make_request_batch( + {"prompt": "A test video.", "modalities": ["video"]}, + make_sampling_params( + seed=seed, + guidance_scale=1.0, + guidance_scale_provided=True, + num_inference_steps=4, + num_frames=5, + height=16, + width=16, + extra_args={"flow_shift": 3.0}, + ), + ) + output = pipeline.forward(request) + return output.output["video"] + + first = run(123) + second = run(123) + different_seed = run(456) + + torch.testing.assert_close(first, second) + assert not torch.equal(first, different_seed) + + +def test_distilled_set_timesteps_uses_fixed_sigma_schedule(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + pipeline._scheduler_init_t_list = [1.0, 0.75, 0.5, 0.25] + + pipeline._set_timesteps( + num_inference_steps=99, + device=torch.device("cpu"), + shift=7.0, + ) + + assert pipeline.scheduler.set_timesteps_calls == [ + { + "num_inference_steps": None, + "device": torch.device("cpu"), + "shift": None, + "sigmas": pipeline._scheduler_init_t_list, + } + ] def test_pipeline_init_passes_tokenizer_attrs_into_transformer( @@ -806,12 +1260,35 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke assert merged.custom_op == ["platform_kernel", "native"] -def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None: - from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline +def test_format_and_tokenize_prompts_leaves_plain_prompts_unchanged_by_default(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_SYSTEM_PROMPT pipeline = make_cosmos3_pipeline() - captured: list[str] = [] - pipeline._tokenize_prompt = lambda text, *args, **kwargs: (captured.append(text) or _ids(len(captured)), _mask()) + calls = _capture_tokenize_calls(pipeline) + + result = 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_system_prompt=False, + is_t2i=False, + ) + + assert [call["text"] for call in calls] == ["A robot.", "bad."] + assert all(call["max_sequence_length"] == 32 for call in calls) + assert all(call["use_system_prompt"] is False for call in calls) + assert all(call["system_prompt"] == COSMOS3_SYSTEM_PROMPT for call in calls) + assert [tensor.item() for tensor in result] == [1, 1, 2, 2] + + +def test_format_and_tokenize_prompts_applies_video_templates_and_system_override(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) pipeline._format_and_tokenize_prompts( "A robot", @@ -821,21 +1298,167 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No height=720, width=1280, max_sequence_length=32, + sp=SimpleNamespace( + system_prompt="direct system prompt", + extra_args={ + "use_duration_template": True, + "use_resolution_template": True, + "system_prompt": "API system prompt", + }, + ), + use_system_prompt=True, + is_t2i=False, + ) + + assert calls == [ + { + "text": ("A robot. The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution."), + "max_sequence_length": 32, + "use_system_prompt": True, + "system_prompt": "API system prompt", + }, + { + "text": ( + "bad. The video is not 2.0 seconds long and is not of 24 FPS. This video is not of 720x1280 resolution." + ), + "max_sequence_length": 32, + "use_system_prompt": True, + "system_prompt": "API system prompt", + }, + ] + + +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 + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + "A robot", + "bad", + num_frames=1, + frame_rate=24, + height=1024, + width=768, + max_sequence_length=64, sp=SimpleNamespace( extra_args={ - "negative_metadata_mode": "inverse", - # Duration/resolution metadata templates are off by default - # (commit "Change resolution and duration templates to off by - # default"); enable them explicitly to exercise the formatting. "use_duration_template": True, "use_resolution_template": True, } ), use_system_prompt=True, + is_t2i=True, + ) + + assert [call["text"] for call in calls] == [ + "A robot. This image is of 1024x768 resolution.", + "bad. This image is not of 1024x768 resolution.", + ] + assert all("seconds" not in call["text"] and "FPS" not in call["text"] for call in calls) + assert all(call["system_prompt"] == COSMOS3_T2I_SYSTEM_PROMPT for call in calls) + + +def test_format_and_tokenize_prompts_rewrites_json_object_metadata(make_cosmos3_pipeline) -> None: + import json + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + ('{"caption": "A robot", "duration": "old", "fps": 1, "resolution": {"H": 1, "W": 2}, "aspect_ratio": "1:1"}'), + "bad", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace( + extra_args={ + "aspect_ratio": "16:9", + "use_duration_template": True, + "use_resolution_template": True, + } + ), is_t2i=False, ) - assert "The video is 2.0 seconds long" in captured[0] - assert "The video is not 2.0 seconds long" in captured[1] + + assert json.loads(calls[0]["text"]) == { + "caption": "A robot", + "duration": "2s", + "fps": 24.0, + "resolution": {"H": 720, "W": 1280}, + "aspect_ratio": "16:9", + } + assert "The video is 2.0 seconds long" not in calls[0]["text"] + assert calls[1]["text"] == ( + "bad. The video is not 2.0 seconds long and is not of 24 FPS. This video is not of 720x1280 resolution." + ) + + +def test_format_and_tokenize_prompts_removes_video_metadata_from_t2i_json(make_cosmos3_pipeline) -> None: + import json + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + '{"caption": "A robot", "duration": "2s", "fps": 24}', + "", + num_frames=1, + frame_rate=24, + height=1024, + width=768, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + is_t2i=True, + ) + + assert json.loads(calls[0]["text"]) == { + "caption": "A robot", + "resolution": {"H": 1024, "W": 768}, + } + + +@pytest.mark.parametrize( + "prompt", + [ + "{malformed", + '["A robot"]', + ], +) +def test_format_and_tokenize_prompts_falls_back_for_non_object_json( + make_cosmos3_pipeline, + prompt: str, +) -> None: + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + prompt, + "", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace( + extra_args={ + "use_duration_template": True, + "use_resolution_template": True, + } + ), + is_t2i=False, + ) + + assert calls[0]["text"] == ( + f"{prompt}. The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution." + ) + + +def test_checkpoint_key_remap() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline remaps = { "embed_tokens.weight": "transformer.language_model.embed_tokens.weight", @@ -847,9 +1470,14 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No "layers.3.self_attn.to_q.weight": "transformer.language_model.layers.3.self_attn.to_q.weight", "layers.3.self_attn.to_out.weight": "transformer.language_model.layers.3.self_attn.to_out.weight", "layers.3.self_attn.norm_q.weight": "transformer.language_model.layers.3.self_attn.norm_q.weight", + "layers.3.self_attn.k_norm_und_for_gen.weight": ( + "transformer.language_model.layers.3.self_attn.k_norm_und_for_gen.weight" + ), "layers.3.self_attn.add_q_proj.weight": "transformer.gen_layers.3.cross_attention.to_q.weight", "layers.3.self_attn.to_add_out.weight": "transformer.gen_layers.3.cross_attention.to_out.weight", "layers.3.self_attn.norm_added_q.weight": "transformer.gen_layers.3.cross_attention.norm_q.weight", + "layers.3.mlp_moe_gen.up_proj.weight": "transformer.gen_layers.3.mlp.up_proj.weight", + "layers.3.mlp_moe_gen.down_proj.weight": "transformer.gen_layers.3.mlp.down_proj.weight", "transformer.model.layers.3.self_attn.add_k_proj.weight": ( "transformer.gen_layers.3.cross_attention.to_k.weight" ), @@ -1193,9 +1821,9 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint pipeline._prepare_transfer_latents = fake_prepare pipeline.diffuse_transfer = fake_diffuse_transfer - def fake_set_flow_shift(target, *, use_karras_sigmas=None): + def fake_set_flow_shift(target): captured.setdefault("flow_shifts", []).append(target) - captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas) + pipeline._current_flow_shift = float(target) pipeline._set_flow_shift = fake_set_flow_shift pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta") @@ -1230,8 +1858,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] + # Transfer applies the V2V flow shift when building its timestep schedule. + assert [call["shift"] for call in pipeline.scheduler.set_timesteps_calls] == [10.0] assert output.output["metadata"]["video"]["fps"] == expected_fps assert output.output["payload"]["video"].device.type == "meta" @@ -1391,9 +2019,9 @@ class TestForwardRouting: pipeline._format_and_tokenize_prompts = fake_format pipeline._prepare_latents = fake_prepare - def fake_set_flow_shift(target, *, use_karras_sigmas=None): + def fake_set_flow_shift(target): captured.setdefault("flow_shifts", []).append(target) - captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas) + pipeline._current_flow_shift = float(target) pipeline._set_flow_shift = fake_set_flow_shift pipeline.diffuse = fake_diffuse @@ -1420,7 +2048,7 @@ class TestForwardRouting: { "key": "video", "is_t2i": False, - "flow": [1.0], + "flow": [10.0], "steps": [35], "frames": 189, }, @@ -1443,8 +2071,8 @@ class TestForwardRouting: assert captured["format"]["is_t2i"] is expected["is_t2i"] 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"] + assert [call["num_inference_steps"] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"] + assert all(call["shift"] == expected["flow"][0] for call in pipeline.scheduler.set_timesteps_calls) def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None: pipeline = make_cosmos3_pipeline() @@ -1491,7 +2119,6 @@ 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["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 13d1c399b..cb59dffe5 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -15,6 +15,25 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.fixture(autouse=True) +def _single_rank_tensor_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + """Provide the TP metadata required by vLLM parallel linear layers.""" + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + + from vllm_omni.diffusion.models.cosmos3 import ( + transformer_cosmos3, + transformer_cosmos3_edge, + ) + + monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(transformer_cosmos3_edge, "get_tensor_model_parallel_world_size", lambda: 1) + + @pytest.fixture def accelerator_device() -> torch.device: """Provide an accelerator device, skipping when none is available.""" @@ -40,6 +59,22 @@ def _tiny_cosmos3_config(**overrides): return config +def _tiny_cosmos3_edge_config(**overrides): + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + COSMOS3_EDGE_BACKBONE_TYPE, + ) + + config = _tiny_cosmos3_config( + latent_patch_size=2, + latent_channel=48, + temporal_compression_factor=4, + backbone_type=COSMOS3_EDGE_BACKBONE_TYPE, + qk_norm_for_text=False, + ) + config.update(overrides) + return config + + def test_mrope_position_ids_cover_text_video_sound_and_action() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( compute_mrope_position_ids_action, @@ -78,6 +113,16 @@ def test_mrope_position_ids_cover_text_video_sound_and_action() -> None: assert action_offset == 14 +def test_timestep_embedder_stores_frequencies_in_fp32() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + TimestepEmbedder, + ) + + embedder = TimestepEmbedder(hidden_size=8, frequency_embedding_size=16) + + assert embedder.freqs.dtype == torch.float32 + + @pytest.mark.parametrize( ("key", "value"), [ @@ -97,30 +142,245 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) -> Cosmos3VFMTransformer._validate_supported_config(None) +def test_edge_config_resolves_nemotron_defaults() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + layer_norm_epsilon=1e-5, + rope_scaling={}, + rope_parameters={"mrope_section": [24, 20, 20]}, + ), + dtype=torch.float32, + ) + ) + + assert model.rms_norm_eps == 1e-5 + assert model.rope_theta == 100_000_000 + assert model.mrope_section == [24, 20, 20] + assert model.latent_channel_size == 48 + assert model.latent_patch_size == 2 + assert model.temporal_compression_factor == 4 + + +def test_edge_config_requires_backbone_type() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + config = _tiny_cosmos3_edge_config() + config.pop("backbone_type") + + with pytest.raises(ValueError, match="must declare backbone_type"): + Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=config, + dtype=torch.float32, + ) + ) + + +def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + Cosmos3Relu2MLP, + ) + + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + num_hidden_layers=1, + use_k_norm_und_for_gen=False, + ), + dtype=torch.float32, + ) + ) + + layer = model.language_model.layers[0] + gen_layer = model.gen_layers[0] + assert isinstance(layer.mlp, Cosmos3Relu2MLP) + assert not hasattr(layer.mlp, "gate_proj") + assert isinstance(gen_layer.mlp, Cosmos3Relu2MLP) + assert not hasattr(gen_layer.mlp, "gate_proj") + assert not hasattr(layer.self_attn, "norm_q") + assert not hasattr(layer.self_attn, "norm_k") + assert layer.self_attn.k_norm_und_for_gen is None + + +def test_edge_creates_gen_k_norm_when_enabled() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + num_hidden_layers=1, + qk_norm_for_diffusion=True, + use_k_norm_und_for_gen=True, + ), + dtype=torch.float32, + ) + ) + + k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen + assert model.use_k_norm_und_for_gen is True + assert isinstance(k_norm, RMSNorm) + assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm) + assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm) + + +def test_edge_supports_no_gen_qk_norm_variant() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + num_hidden_layers=1, + qk_norm_for_diffusion=False, + use_k_norm_und_for_gen=False, + ), + dtype=torch.float32, + ) + ) + + assert model.qk_norm_for_diffusion is False + assert model.language_model.layers[0].self_attn.k_norm_und_for_gen is None + assert model.gen_layers[0].cross_attention.qk_norm is False + assert not hasattr(model.gen_layers[0].cross_attention, "norm_q") + assert not hasattr(model.gen_layers[0].cross_attention, "norm_k") + + +def test_edge_validates_required_relu2_weights() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + model = object.__new__(Cosmos3EdgeVFMTransformer) + nn.Module.__init__(model) + model.num_hidden_layers = 1 + model.qk_norm_for_diffusion = True + model.use_k_norm_und_for_gen = True + + complete = { + "transformer.language_model.layers.0.mlp.up_proj.weight", + "transformer.language_model.layers.0.mlp.down_proj.weight", + "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight", + "transformer.gen_layers.0.mlp.up_proj.weight", + "transformer.gen_layers.0.mlp.down_proj.weight", + } + model.validate_loaded_weights(complete) + + mlp_weights = complete - { + "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight", + } + for missing_weight in mlp_weights: + with pytest.raises(ValueError, match="missing required weights"): + model.validate_loaded_weights(complete - {missing_weight}) + + missing_k_norm = complete - {"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight"} + with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"): + model.validate_loaded_weights(missing_k_norm) + + model.use_k_norm_und_for_gen = False + model.validate_loaded_weights(missing_k_norm) + + +def test_edge_gen_cached_k_is_normalized_but_reasoner_uses_raw_k() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeCausalAttention, + ) + + class ScaleProjection(nn.Module): + def __init__(self, scale: float) -> None: + super().__init__() + self.scale = scale + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * self.scale + + class IdentityOutput(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + class RawKAttention(nn.Module): + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + del v + return q + k + + attn = object.__new__(Cosmos3EdgeCausalAttention) + nn.Module.__init__(attn) + attn.hidden_size = 4 + attn.num_heads = 2 + attn.num_kv_heads = 2 + attn.head_dim = 2 + attn.num_heads_local = 2 + attn.num_kv_heads_local = 2 + attn.to_q = ScaleProjection(1.0) + attn.to_k = ScaleProjection(1.0) + attn.to_v = ScaleProjection(1.0) + attn.to_out = IdentityOutput() + attn.k_norm_und_for_gen = RMSNorm(2, eps=1e-5) + attn.attn = RawKAttention() + + hidden_states = torch.tensor( + [[[1.0, 2.0, 3.0, 4.0], [2.0, 1.0, 4.0, 3.0]]], + dtype=torch.float32, + ) + freqs_cos = torch.ones(1, 2, 1, 2) + freqs_sin = torch.zeros(1, 2, 1, 2) + + reasoner_out, gen_k, gen_v = attn(hidden_states, freqs_cos, freqs_sin) + attn.to_k.scale = 3.0 + scaled_reasoner_out, scaled_gen_k, scaled_gen_v = attn(hidden_states, freqs_cos, freqs_sin) + + torch.testing.assert_close(gen_k, scaled_gen_k) + torch.testing.assert_close(gen_v, scaled_gen_v) + assert not torch.allclose(reasoner_out, scaled_reasoner_out) + + def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeLanguageModel, + Cosmos3EdgeVFMTransformer, + ) - model = object.__new__(Cosmos3VFMTransformer) - nn.Module.__init__(model) - model.language_model = nn.Module() - model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)]) - model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) - model.norm_moe_gen = nn.LayerNorm(2) + for transformer_cls, language_model_cls in ( + (Cosmos3VFMTransformer, Cosmos3LanguageModel), + (Cosmos3EdgeVFMTransformer, Cosmos3EdgeLanguageModel), + ): + model = object.__new__(transformer_cls) + nn.Module.__init__(model) + model.language_model = nn.Module() + model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)]) + model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) + model.norm_moe_gen = nn.LayerNorm(2) - matched = [ - name - for name, module in model.named_modules() - if any(condition(name, module) for condition in model._hsdp_shard_conditions) - ] - assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"] - assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"] - assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"] - assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"] + matched = [ + name + for name, module in model.named_modules() + if any(condition(name, module) for condition in model._hsdp_shard_conditions) + ] + assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"] + assert transformer_cls._layerwise_offload_blocks_attrs == ["gen_layers"] + assert language_model_cls._layerwise_offload_blocks_attrs == ["layers"] + assert transformer_cls._repeated_blocks == ["Cosmos3GenDecoderLayer"] - model.latent_patch_size = 2 - model.latent_channel_size = 3 - latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5) - torch.testing.assert_close(model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents) + model.latent_patch_size = 2 + model.latent_channel_size = 3 + latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5) + torch.testing.assert_close( + model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), + latents, + ) def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None: @@ -142,6 +402,17 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No assert tuple(output.shape) == (1, 2, 1, 2, 2) +def test_qwen_gen_mlp_remains_gated() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer + + model = Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1), dtype=torch.float32) + ) + + assert isinstance(model.gen_layers[0].mlp, Cosmos3GatedMLP) + assert hasattr(model.gen_layers[0].mlp, "gate_proj") + + def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 diff --git a/tests/e2e/accuracy/test_cosmos3_similarity.py b/tests/e2e/accuracy/test_cosmos3_similarity.py index ff2350096..efa095fb1 100644 --- a/tests/e2e/accuracy/test_cosmos3_similarity.py +++ b/tests/e2e/accuracy/test_cosmos3_similarity.py @@ -42,8 +42,6 @@ def _server_args() -> list[str]: return [ "--num-gpus", "1", - "--model-class-name", - "Cosmos3OmniDiffusersPipeline", "--stage-init-timeout", "900", "--init-timeout", diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index f2db9ded0..76ef84e7e 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -60,8 +60,9 @@ def test_sensenova_extra_registry_declares_request_and_response_params() -> None @pytest.mark.core_model @pytest.mark.cpu -def test_cosmos3_extra_registry_declares_request_and_response_params() -> None: - assert get_extra_body_params("Cosmos3OmniDiffusersPipeline") == frozenset( +@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) +def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_name: str) -> None: + assert get_extra_body_params(pipeline_name) == frozenset( { "flow_shift", "max_sequence_length", @@ -98,7 +99,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params() -> None: "session_id", } ) - assert get_extra_output_params("Cosmos3OmniDiffusersPipeline") == frozenset( + assert get_extra_output_params(pipeline_name) == frozenset( { "action", "raw_action_dim", @@ -106,7 +107,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params() -> None: "action_mode", } ) - assert should_init_extra_args_for_non_diffusion_stages("Cosmos3OmniDiffusersPipeline") is False + assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False @pytest.mark.core_model @@ -224,9 +225,10 @@ def test_ming_flash_omni_image_to_image_prompt_builder() -> None: @pytest.mark.core_model @pytest.mark.cpu -def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None: +@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) +def test_cosmos3_text_to_image_prompt_builder_selects_image_modality(pipeline_name: str) -> None: assert build_text_to_image_prompt( - "Cosmos3OmniDiffusersPipeline", + pipeline_name, prompt="a red sports car at golden hour", negative_prompt="blurry, distorted", height=1024, @@ -237,7 +239,7 @@ def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None: "negative_prompt": "blurry, distorted", } assert build_text_to_image_prompt( - "Cosmos3OmniDiffusersPipeline", + pipeline_name, prompt="a red sports car", negative_prompt=None, ) == {"prompt": "a red sports car", "modalities": ["image"]} diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index ae01334f6..c95644143 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -71,6 +71,7 @@ EXCLUDED_MODELS = [ "OmniVoicePipeline", "OmniVoice", "Cosmos3OmniDiffusersPipeline", + "Cosmos3OmniPipeline", "DiffusersAdapterPipeline", "HiDreamImagePipeline", "DreamZeroPipeline", diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 48e64de0a..5d63ce3c0 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -1016,6 +1016,7 @@ CUSTOM_DIT_ENABLERS.update( "Wan22VACEPipeline": enable_cache_for_wan22, "Wan22S2VPipeline": enable_cache_for_wan22_s2v, "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3, + "Cosmos3OmniPipeline": enable_cache_for_cosmos3, "Krea2Pipeline": enable_cache_for_krea2, } ) diff --git a/vllm_omni/diffusion/models/cosmos3/__init__.py b/vllm_omni/diffusion/models/cosmos3/__init__.py index 6df062b5c..af4eeaacf 100644 --- a/vllm_omni/diffusion/models/cosmos3/__init__.py +++ b/vllm_omni/diffusion/models/cosmos3/__init__.py @@ -7,10 +7,12 @@ from .pipeline_cosmos3 import ( get_cosmos3_pre_process_func, ) from .transformer_cosmos3 import Cosmos3VFMTransformer +from .transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer __all__ = [ "Cosmos3OmniDiffusersPipeline", "get_cosmos3_post_process_func", "get_cosmos3_pre_process_func", "Cosmos3VFMTransformer", + "Cosmos3EdgeVFMTransformer", ] diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 5d0e24789..58e74ccf4 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -29,6 +29,7 @@ is produced from sound latents rather than from ``multi_modal_data["audio"]``. from __future__ import annotations +import json import math import os import time @@ -39,7 +40,7 @@ from typing import Any, ClassVar import numpy as np import PIL.Image import torch -from diffusers import UniPCMultistepScheduler +from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor from torch import nn @@ -61,6 +62,9 @@ from vllm_omni.diffusion.models.interface import ( SupportsComponentDiscovery, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero +from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch @@ -94,7 +98,8 @@ from .transfer import ( transfer_max_frames_from_extra_args, uint8_cthw_to_normalized_5d, ) -from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen +from .transformer_cosmos3 import Cosmos3VFMTransformer, _tf_config_get, resolve_sound_gen +from .transformer_cosmos3_edge import COSMOS3_EDGE_BACKBONE_TYPE, Cosmos3EdgeVFMTransformer from .utils import ( COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, COSMOS3_VAE_TEMPORAL_COMPRESSION, @@ -154,6 +159,7 @@ COSMOS3_T2V_DEFAULT_WIDTH = 1280 COSMOS3_T2V_DEFAULT_NUM_FRAMES = 189 COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS = 35 COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE = 6.0 +COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT = 10.0 COSMOS3_T2I_DEFAULT_HEIGHT = 1024 COSMOS3_T2I_DEFAULT_WIDTH = 1024 @@ -162,12 +168,73 @@ COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE = 7.0 COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0 COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0) +COSMOS3_EDGE_T2V_DEFAULT_HEIGHT = 480 +COSMOS3_EDGE_T2V_DEFAULT_WIDTH = 832 +COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE = 5.0 +COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT = 3.0 +COSMOS3_EDGE_T2I_DEFAULT_HEIGHT = 640 +COSMOS3_EDGE_T2I_DEFAULT_WIDTH = 640 + +COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS = "FlowMatchEulerDiscreteScheduler" + # Truncation cap on the prompt token count (shared by T2I and T2V). Prompts # are tokenized to their natural length (no padding); this only bounds the # UND pathway / GEN cross-attention cost for pathologically long prompts. COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096 +def _ceil_video_num_frames(num_frames: int, temporal_compression_factor: int) -> int: + """Round a video length up to the causal VAE's ``factor * k + 1`` grid.""" + if temporal_compression_factor <= 0: + raise ValueError(f"Cosmos3 temporal_compression_factor must be positive, got {temporal_compression_factor}.") + if num_frames <= 1: + return num_frames + latent_frames = math.ceil((num_frames - 1) / temporal_compression_factor) + return latent_frames * temporal_compression_factor + 1 + + +def _format_json_object_prompt( + prompt: str, + *, + num_frames: int, + frame_rate: float, + height: int, + width: int, + aspect_ratio: str | None, +) -> str | None: + """Match Imaginaire4 metadata injection for JSON-object prompts.""" + try: + prompt_obj = json.loads(prompt) + except (json.JSONDecodeError, TypeError, ValueError): + return None + if not isinstance(prompt_obj, dict): + return None + + metadata: dict[str, Any] = {} + if num_frames > 1: + duration_seconds = int(num_frames / frame_rate) if frame_rate > 0 else 0 + metadata.update({"duration": f"{duration_seconds}s", "fps": float(frame_rate)}) + else: + prompt_obj.pop("duration", None) + prompt_obj.pop("fps", None) + metadata["resolution"] = {"H": int(height), "W": int(width)} + if aspect_ratio is not None: + metadata["aspect_ratio"] = aspect_ratio + + prompt_obj.update(metadata) + return json.dumps(prompt_obj) + + +def resolve_cosmos3_transformer_cls(model_config: Any) -> type[Cosmos3VFMTransformer]: + """Select the Cosmos3 transformer implementation from transformer/config.json.""" + backbone_type = _tf_config_get(model_config, "backbone_type", None) + if backbone_type is None: + return Cosmos3VFMTransformer + if backbone_type == COSMOS3_EDGE_BACKBONE_TYPE: + return Cosmos3EdgeVFMTransformer + raise ValueError(f"Unsupported Cosmos3 transformer backbone_type={backbone_type!r}.") + + # --------------------------------------------------------------------------- # Post-process function (registered in registry.py) # --------------------------------------------------------------------------- @@ -190,6 +257,9 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled video_processor = VideoProcessor(vae_scale_factor=16) + is_edge_model = ( + _tf_config_get(getattr(od_config, "tf_model_config", None), "backbone_type", None) == COSMOS3_EDGE_BACKBONE_TYPE + ) # Eager-load guardrail models at pipeline build time when the server-level # gate is on. Per-request overrides only decide whether the loaded models # are *invoked* — they cannot turn checks on without a server-side preload. @@ -410,12 +480,17 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): extra = _extra_args(request) transfer_requested = action_mode is None and has_transfer_hints(extra) - # Auto-calculate H/W from aspect ratio (720p max area) + # Resolve missing H/W. if transfer_requested: _set_transfer_size_from_image(request, image) elif request.sampling_params.height is None or request.sampling_params.width is None: if action_mode is not None: _set_action_size_from_image(request, image) + elif is_edge_model: + if request.sampling_params.height is None: + request.sampling_params.height = COSMOS3_EDGE_T2V_DEFAULT_HEIGHT + if request.sampling_params.width is None: + request.sampling_params.width = COSMOS3_EDGE_T2V_DEFAULT_WIDTH else: max_area = 720 * 1280 aspect_ratio = image.height / image.width @@ -823,27 +898,52 @@ class Cosmos3OmniDiffusersPipeline( sound_latent_fps = self._sound_tokenizer.latent_fps # --- Transformer (weights loaded later via weights_sources) --- - self.transformer = Cosmos3VFMTransformer( + transformer_cls = resolve_cosmos3_transformer_cls(od_config.tf_model_config) + self.transformer = transformer_cls( od_config=od_config, temporal_compression_factor=self.vae_scale_factor_temporal, sound_gen=sound_gen, sound_dim=sound_dim, sound_latent_fps=sound_latent_fps, ) + self.is_edge_model = transformer_cls is Cosmos3EdgeVFMTransformer # --- Scheduler --- - # Load from checkpoint to preserve solver_order, timestep_spacing, - # beta_schedule, sigma bounds, flow_shift, etc. Only override - # flow_shift when explicitly requested by the user. - self.scheduler = UniPCMultistepScheduler.from_pretrained( + # Distilled model differs from regular one only by scheduler, + # distilled one uses FlowMatchEulerDiscreteScheduler, while + # regular should use FlowUniPCMultistepScheduler + + scheduler_config = FlowUniPCMultistepScheduler.load_config( model_path, 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, + + scheduler_class_name = scheduler_config.get("_class_name") + + self.is_distilled_model = False + if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS: + fixed_step_config = scheduler_config.get("fixed_step_sampler_config") + if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde": + raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.") + t_list = fixed_step_config.get("t_list") + if not isinstance(t_list, list) or not t_list: + raise ValueError("Cosmos3 distilled scheduler requires a non-empty fixed_step_sampler_config.t_list.") + self.scheduler = FlowMatchEulerDiscreteScheduler.from_config( + scheduler_config, + stochastic_sampling=True, + ) + self._scheduler_init_t_list = list(t_list) + self.is_distilled_model = True + else: + # Preserve compatible solver settings from the checkpoint, but keep + # the base shift neutral. The concrete request shift is applied when + # FlowUniPC builds its timesteps. + self.scheduler = FlowUniPCMultistepScheduler.from_config( + scheduler_config, + shift=1.0, + use_dynamic_shifting=False, + prediction_type="flow_prediction", ) self._cpu_scheduler_state() @@ -862,14 +962,16 @@ class Cosmos3OmniDiffusersPipeline( ), ] - # Snapshot the loaded scheduler config so we can rebuild the - # scheduler at request time when a per-request flow_shift override - # is supplied (T2I uses shift=3.0; T2V/I2V use the engine default). - self._base_scheduler_config = self.scheduler.config - self._engine_init_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0) + # An engine-level override becomes the video default; otherwise use + # the matching regular/Edge default. Per-request values still take + # precedence. + default_video_flow_shift = ( + COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT if self.is_edge_model else COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT + ) + self._engine_init_flow_shift = float( + od_config.flow_shift if od_config.flow_shift is not None else default_video_flow_shift + ) 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 @@ -989,6 +1091,7 @@ class Cosmos3OmniDiffusersPipeline( "self_attn.to_out.": f"{und_lp}.self_attn.to_out.", "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.", "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.", + "self_attn.k_norm_und_for_gen.": f"{und_lp}.self_attn.k_norm_und_for_gen.", # GEN attention "self_attn.add_q_proj.": f"{gen_lp}.cross_attention.to_q.", "self_attn.add_k_proj.": f"{gen_lp}.cross_attention.to_k.", @@ -1054,6 +1157,7 @@ class Cosmos3OmniDiffusersPipeline( loaded = loader.load_weights(_remapped_weights()) self.transformer.post_load_weights() self.transformer.eval() + self.transformer.validate_loaded_weights(loaded) if getattr(self.transformer, "sound_gen", False): sound_markers = ("audio_proj_in.", "audio_proj_out.", "audio_modality_embed") missing = [marker.rstrip(".") for marker in sound_markers if not any(marker in name for name in loaded)] @@ -1359,6 +1463,8 @@ class Cosmos3OmniDiffusersPipeline( inputs: RoboLabPolicyInputs, pipeline_start: float, ) -> DiffusionOutput: + if self.is_distilled_model: + raise self._distilled_unsupported_error("RoboLab/action policy requests are unsupported.") if not getattr(self.transformer, "action_gen", False): raise ValueError( "Cosmos3 RoboLab policy serving was requested, but the transformer " @@ -1469,6 +1575,7 @@ class Cosmos3OmniDiffusersPipeline( guidance_interval=None, raw_action_dim=raw_action_dim, scheduler=scheduler, + generator=generator, ) if _is_rank_zero(): @@ -1560,38 +1667,34 @@ class Cosmos3OmniDiffusersPipeline( raise ValueError(f"Incorrect modality value in {modalities}, expected one of {accepted_modalities}.") return "image" in modalities + def _set_timesteps(self, num_inference_steps: int, device: str | torch.device, shift: float) -> None: + if self.is_distilled_model: + self.scheduler.set_timesteps(sigmas=self._scheduler_init_t_list, device=device) + else: + self.scheduler.set_timesteps( + num_inference_steps, + device=device, + shift=shift, + ) + + def _set_flow_shift(self, target_shift: float) -> None: + """Select the shift applied by FlowUniPC when timesteps are built.""" + self._current_flow_shift = float(target_shift) + @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 _resolve_seed(sp: OmniDiffusionSamplingParams, generator: torch.Generator | None, default: int = 42) -> int: + if sp.seed is not None: + return int(sp.seed) + if isinstance(generator, torch.Generator): + return int(generator.initial_seed()) + return int(default) - 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 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) - 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 _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float: + if self.is_distilled_model: + return 1.0 + if sp.guidance_scale_provided: + return float(sp.guidance_scale) + return float(default) def _cpu_scheduler_state(self) -> None: # We need to move scheduler tensors to CPU, as unipc from diffusers assumes they are on CPU. @@ -1612,6 +1715,53 @@ class Cosmos3OmniDiffusersPipeline( def num_timesteps(self): return self._num_timesteps + @staticmethod + def _distilled_unsupported_error(detail: str) -> ValueError: + return ValueError( + f"Cosmos3 distilled checkpoints support only text-to-image and image-to-video generation; {detail}" + ) + + def _validate_distilled_generation_mode( + self, + *, + is_t2i: bool, + image_tensor: torch.Tensor | None, + action_enabled: bool, + transfer_config: Cosmos3TransferConfig | None, + is_v2v: bool, + sound_enabled: bool, + ) -> None: + if not self.is_distilled_model: + return + if action_enabled: + raise self._distilled_unsupported_error("action requests are unsupported.") + if transfer_config is not None: + raise self._distilled_unsupported_error("transfer requests are unsupported.") + if is_v2v: + raise self._distilled_unsupported_error("video-to-video requests are unsupported.") + if sound_enabled: + raise self._distilled_unsupported_error("sound generation is unsupported.") + if not is_t2i and image_tensor is None: + raise self._distilled_unsupported_error("text-to-video requests are unsupported.") + if is_t2i and image_tensor is not None: + raise self._distilled_unsupported_error("image-conditioned image generation is unsupported.") + + def _validate_edge_generation_mode( + self, + *, + transfer_config: Cosmos3TransferConfig | None, + is_v2v: bool, + sound_enabled: bool, + ) -> None: + if not self.is_edge_model: + return + if sound_enabled: + raise ValueError("Cosmos3 Edge checkpoints do not support sound generation.") + if transfer_config is not None: + raise ValueError("Cosmos3 Edge checkpoints do not support transfer inference.") + if is_v2v: + raise ValueError("Cosmos3 Edge checkpoints do not support video-to-video generation.") + # -- Prompt formatting ----------------------------------------------------- @staticmethod @@ -1628,6 +1778,10 @@ class Cosmos3OmniDiffusersPipeline( """ Append duration and resolution metadata to a prompt. """ + prompt = prompt.strip() + if duration_template is None and resolution_template is None: + return prompt + parts: list[str] = [] head = prompt.rstrip(".").strip() if head: @@ -1903,15 +2057,26 @@ class Cosmos3OmniDiffusersPipeline( res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE else: res_tmpl = None - prompt = self._apply_metadata_templates( + json_prompt = _format_json_object_prompt( prompt, - num_frames, - frame_rate, - height, - width, - duration_template=dur_tmpl, - resolution_template=res_tmpl, + num_frames=num_frames, + frame_rate=frame_rate, + height=height, + width=width, + aspect_ratio=self._get_sp_param(sp, "aspect_ratio", None), ) + if json_prompt is not None: + prompt = json_prompt + else: + prompt = self._apply_metadata_templates( + prompt, + num_frames, + frame_rate, + height, + width, + duration_template=dur_tmpl, + resolution_template=res_tmpl, + ) if _is_rank_zero(): logger.info("Final prompt: '%s'", prompt) @@ -2272,6 +2437,7 @@ class Cosmos3OmniDiffusersPipeline( guidance_interval: tuple[float, float] | None = None, raw_action_dim: int | None = None, scheduler: Any | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Denoising loop with 3-mode CFG support (parallel, sequential, none). @@ -2383,11 +2549,23 @@ class Cosmos3OmniDiffusersPipeline( if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]: action_pred[..., raw_action_dim:] = 0 if action_latents is None and sound_latents is None: - latents = step_scheduler.step(video_pred, t, latents, return_dict=False)[0] + latents = step_scheduler.step( + video_pred, + t, + latents, + generator=generator, + return_dict=False, + )[0] else: packed_noise, shapes, numels = _pack_joint(video_pred, action_pred, sound_pred) packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents) - packed_next = step_scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0] + packed_next = step_scheduler.step( + packed_noise, + t, + packed_latents, + generator=generator, + return_dict=False, + )[0] unpacked = _unpack_joint(packed_next, shapes, numels) latents = unpacked[0] idx = 1 @@ -2613,6 +2791,7 @@ class Cosmos3OmniDiffusersPipeline( velocity_mask: torch.Tensor, condition_latents: torch.Tensor, guidance_interval: tuple[float, float] | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor: def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool: if interval is None: @@ -2724,7 +2903,13 @@ class Cosmos3OmniDiffusersPipeline( if isinstance(noise_pred, tuple): raise ValueError("Cosmos3 transfer diffusion expects video-only tensor predictions.") noise_pred = noise_pred * velocity_mask - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + latents = self.scheduler.step( + noise_pred, + t, + latents, + generator=generator, + return_dict=False, + )[0] latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents finally: self._cosmos3_branch_caches = None @@ -2741,6 +2926,8 @@ class Cosmos3OmniDiffusersPipeline( transfer_video_tensor: torch.Tensor | None, transfer_input_fps: float | None, ) -> DiffusionOutput: + if self.is_distilled_model: + raise self._distilled_unsupported_error("transfer requests are unsupported.") input_frames = None if transfer_video_tensor is not None: input_frames = normalized_video_to_uint8_cthw(transfer_video_tensor) @@ -2806,7 +2993,7 @@ class Cosmos3OmniDiffusersPipeline( guidance_scale = ( float(transfer_config.guidance_scale) if transfer_config.guidance_scale is not None - else float(sp.guidance_scale or COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE) + else self._resolve_guidance_scale(sp, COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE) ) flow_shift_target = float( transfer_config.flow_shift @@ -2821,11 +3008,11 @@ class Cosmos3OmniDiffusersPipeline( self._guidance_scale = guidance_scale self._num_timesteps = num_inference_steps - self._set_flow_shift(flow_shift_target, use_karras_sigmas=False) + self._set_flow_shift(flow_shift_target) generator = sp.generator + seed = self._resolve_seed(sp, generator) if generator is None: - seed = sp.seed if sp.seed is not None else 42 generator = torch.Generator(device=self.device).manual_seed(seed) cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts( @@ -2915,7 +3102,11 @@ class Cosmos3OmniDiffusersPipeline( transfer_share_vision_temporal_positions=transfer_config.share_vision_temporal_positions, ) - self.scheduler.set_timesteps(num_inference_steps, device=self.device) + self._set_timesteps( + num_inference_steps, + device=self.device, + shift=self._current_flow_shift, + ) latents = self.diffuse_transfer( latents=latents, timesteps=self.scheduler.timesteps, @@ -2930,6 +3121,7 @@ class Cosmos3OmniDiffusersPipeline( shared_kwargs=shared_kwargs, velocity_mask=velocity_mask, condition_latents=condition_latents, + generator=generator, ) output_video = self._decode_latents(latents).clamp(-1, 1) previous_output = output_video @@ -3009,6 +3201,20 @@ class Cosmos3OmniDiffusersPipeline( action_enabled = action_mode is not None transfer_config = resolve_transfer_config(sp, prompt_data) action_video_tensor = video_tensor if action_enabled else None + is_v2v = video_tensor is not None and not is_t2i and not action_enabled + self._validate_distilled_generation_mode( + is_t2i=is_t2i, + image_tensor=image_tensor, + action_enabled=action_enabled, + transfer_config=transfer_config, + is_v2v=is_v2v, + sound_enabled=sound_enabled, + ) + self._validate_edge_generation_mode( + transfer_config=transfer_config, + is_v2v=is_v2v, + sound_enabled=sound_enabled, + ) if transfer_config is not None: if is_t2i: raise ValueError("Cosmos3 transfer inference is supported only for video outputs.") @@ -3056,31 +3262,36 @@ class Cosmos3OmniDiffusersPipeline( raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.") if video_tensor is not None and is_t2i: raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.") - is_v2v = video_tensor is not None and not is_t2i and not action_enabled - # T2I and T2V share the same model + forward path; only defaults - # differ: - # T2I: 1024x1024, 50 steps, shift=3.0, guidance_interval=[400, 1000] - # T2V: 720x1280, 35 steps, shift=engine-init, no interval + # T2I and T2V share the same model and forward path; their defaults are: + # T2I: regular 1024x1024, Edge 640x640; 50 steps, shift=3.0, + # guidance_interval=[400, 1000] + # T2V: 189 frames and 35 steps; regular guidance=6, shift=10; + # Edge guidance=5, shift=3; + # no guidance interval if is_t2i: - height = sp.height or COSMOS3_T2I_DEFAULT_HEIGHT - width = sp.width or COSMOS3_T2I_DEFAULT_WIDTH + height = sp.height or ( + COSMOS3_EDGE_T2I_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2I_DEFAULT_HEIGHT + ) + width = sp.width or (COSMOS3_EDGE_T2I_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2I_DEFAULT_WIDTH) num_frames = 1 num_inference_steps = sp.num_inference_steps or COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS - guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE + guidance_scale = self._resolve_guidance_scale(sp, COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE) default_flow_shift = COSMOS3_T2I_DEFAULT_FLOW_SHIFT default_guidance_interval: tuple[float, float] | None = COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL batch_size = max(1, int(sp.num_outputs_per_prompt or 1)) else: - height = sp.height or COSMOS3_T2V_DEFAULT_HEIGHT - width = sp.width or COSMOS3_T2V_DEFAULT_WIDTH + height = sp.height or ( + COSMOS3_EDGE_T2V_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2V_DEFAULT_HEIGHT + ) + width = sp.width or (COSMOS3_EDGE_T2V_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2V_DEFAULT_WIDTH) + default_guidance_scale = ( + COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE if self.is_edge_model else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE + ) num_frames = sp.num_frames or COSMOS3_T2V_DEFAULT_NUM_FRAMES num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS - guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE - # Fall back to the engine-init shift, NOT None: passing None - # to ``_set_flow_shift`` would leak a prior T2I rebuild - # (shift=3.0) into a subsequent video request. - default_flow_shift = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift + guidance_scale = self._resolve_guidance_scale(sp, default_guidance_scale) + default_flow_shift = self._engine_init_flow_shift default_guidance_interval = None batch_size = 1 # Existing video pipeline assumes B=1. @@ -3103,9 +3314,23 @@ class Cosmos3OmniDiffusersPipeline( f"or action_chunk_size + 1; got num_frames={num_frames}, action_chunk_size={action_chunk_size}." ) num_inference_steps = sp.num_inference_steps or 30 - guidance_scale = sp.guidance_scale if sp.guidance_scale is not None else 1.0 + guidance_scale = self._resolve_guidance_scale(sp, 1.0) default_flow_shift = 5.0 + if not is_t2i and not action_enabled: + requested_num_frames = int(num_frames) + num_frames = _ceil_video_num_frames( + requested_num_frames, + self.vae_scale_factor_temporal, + ) + if num_frames != requested_num_frames and _is_rank_zero(): + logger.info( + "Rounded Cosmos3 num_frames from %d to %d for temporal compression factor %d.", + requested_num_frames, + num_frames, + self.vae_scale_factor_temporal, + ) + domain_id = None if action_enabled: domain_id = resolve_domain_id( @@ -3146,13 +3371,12 @@ class Cosmos3OmniDiffusersPipeline( self._guidance_scale = guidance_scale self._num_timesteps = num_inference_steps - # Always resolve to a concrete target shift for this request, then - # update the shared Diffusers scheduler. - self._set_flow_shift(flow_shift_target, use_karras_sigmas=False if is_v2v else None) + # Always resolve to a concrete target shift for this request. + self._set_flow_shift(flow_shift_target) generator = sp.generator + seed = self._resolve_seed(sp, generator) if generator is None: - seed = sp.seed if sp.seed is not None else 42 generator = torch.Generator(device=self.device).manual_seed(seed) # --- Format prompts & tokenize (B=1; reused across loop iterations @@ -3293,7 +3517,11 @@ class Cosmos3OmniDiffusersPipeline( ) def _run_diffusion(start_latents): - self.scheduler.set_timesteps(num_inference_steps, device=self.device) + self._set_timesteps( + num_inference_steps, + device=self.device, + shift=self._current_flow_shift, + ) scheduler = self.scheduler return self.diffuse( latents=start_latents, @@ -3314,6 +3542,7 @@ class Cosmos3OmniDiffusersPipeline( guidance_interval=guidance_interval, raw_action_dim=raw_action_dim, scheduler=scheduler, + generator=generator, ) if is_t2i and batch_size > 1: diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 21e723ced..de9767d6f 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -396,7 +396,6 @@ class TimestepEmbedder(nn.Module): hidden_size: int, frequency_embedding_size: int = 256, max_period: int = 10000, - target_dtype: torch.dtype = torch.bfloat16, ) -> None: super().__init__() # Following diffusers naming pattern here for checkpoint compatibility. @@ -407,7 +406,7 @@ class TimestepEmbedder(nn.Module): self.hidden_size = hidden_size half = frequency_embedding_size // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half) + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half) self.register_buffer("freqs", freqs, persistent=False) def forward(self, t: torch.Tensor) -> torch.Tensor: @@ -587,6 +586,7 @@ class Cosmos3CrossAttention(nn.Module): head_dim: int, rms_norm_eps: float, quant_config: QuantizationConfig | None = None, + qk_norm: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -636,8 +636,10 @@ class Cosmos3CrossAttention(nn.Module): prefix=f"{prefix}.to_out", ) - self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps) - self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.qk_norm = qk_norm + if self.qk_norm: + self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps) self.attn = FrameworkAttention( num_heads=self.num_heads, @@ -721,8 +723,9 @@ class Cosmos3CrossAttention(nn.Module): v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim) # Per-head QK norm - q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon) - k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon) + if self.qk_norm: + q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon) + k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon) # Qwen3-style RoPE q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) @@ -806,6 +809,8 @@ class Cosmos3GenDecoderLayer(nn.Module): head_dim: int, rms_norm_eps: float, quant_config: QuantizationConfig | None = None, + mlp_cls: type[nn.Module] = Cosmos3GatedMLP, + qk_norm: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -817,11 +822,12 @@ class Cosmos3GenDecoderLayer(nn.Module): head_dim=head_dim, rms_norm_eps=rms_norm_eps, quant_config=quant_config, + qk_norm=qk_norm, prefix=f"{prefix}.cross_attention", ) self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.mlp = Cosmos3GatedMLP( + self.mlp = mlp_cls( hidden_size=hidden_size, intermediate_size=intermediate_size, quant_config=quant_config, @@ -996,6 +1002,9 @@ class Cosmos3VFMTransformer(nn.Module): packed_modules_mapping = {} + _language_model_cls = Cosmos3LanguageModel + _gen_mlp_cls = Cosmos3GatedMLP + @staticmethod def _is_transformer_block(name: str, module) -> bool: return ("gen_layers" in name or "language_model.layers" in name) and name.split(".")[-1].isdigit() @@ -1032,6 +1041,25 @@ class Cosmos3VFMTransformer(nn.Module): if actual != expected: raise ValueError(f"Unsupported Cosmos3 transformer config: {key}={actual!r}; expected {expected!r}.") + @classmethod + def _resolve_rms_norm_eps(cls, model_config: Any) -> float: + return float(_tf_config_get(model_config, "rms_norm_eps", 1e-6)) + + @classmethod + def _resolve_rope_theta(cls, model_config: Any) -> float: + return float(_tf_config_get(model_config, "rope_theta", 5_000_000)) + + @classmethod + def _resolve_mrope_section(cls, model_config: Any) -> list[int]: + rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {} + return list(rope_scaling.get("mrope_section", [24, 20, 20])) + + def _language_model_kwargs(self) -> dict[str, Any]: + return {} + + def validate_loaded_weights(self, loaded: set[str]) -> None: + del loaded + def __init__( self, od_config: OmniDiffusionConfig, @@ -1044,7 +1072,6 @@ class Cosmos3VFMTransformer(nn.Module): super().__init__() model_config = od_config.tf_model_config self._validate_supported_config(model_config) - rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {} self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096)) self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36)) @@ -1053,9 +1080,10 @@ class Cosmos3VFMTransformer(nn.Module): self.head_dim = int(_tf_config_get(model_config, "head_dim", 128)) self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288)) self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936)) - self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6)) - self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000)) - self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20])) + self.rms_norm_eps = self._resolve_rms_norm_eps(model_config) + self.rope_theta = self._resolve_rope_theta(model_config) + self.mrope_section = self._resolve_mrope_section(model_config) + self.qk_norm_for_diffusion = bool(_tf_config_get(model_config, "qk_norm_for_diffusion", True)) self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2)) self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48)) self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001)) @@ -1095,10 +1123,12 @@ class Cosmos3VFMTransformer(nn.Module): ) self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size + self.use_k_norm_und_for_gen = _tf_config_get(model_config, "use_k_norm_und_for_gen", None) + dtype = od_config.dtype quant_config = getattr(od_config, "quantization_config", None) if od_config else None - self.language_model = Cosmos3LanguageModel( + self.language_model = self._language_model_cls( hidden_size=self.hidden_size, intermediate_size=self.intermediate_size, num_hidden_layers=self.num_hidden_layers, @@ -1111,12 +1141,13 @@ class Cosmos3VFMTransformer(nn.Module): mrope_section=self.mrope_section, quant_config=quant_config, prefix="language_model", + **self._language_model_kwargs(), ) # Video projection layers are small; not worth quantizing. self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size) self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim) - self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype) + self.time_embedder = TimestepEmbedder(self.hidden_size) if self.action_gen: self.action_proj_in = DomainAwareLinear( self.action_dim, @@ -1147,6 +1178,8 @@ class Cosmos3VFMTransformer(nn.Module): head_dim=self.head_dim, rms_norm_eps=self.rms_norm_eps, quant_config=quant_config, + mlp_cls=self._gen_mlp_cls, + qk_norm=self.qk_norm_for_diffusion, prefix=f"gen_layers.{i}", ) for i in range(self.num_hidden_layers) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py new file mode 100644 index 000000000..6a1b4313b --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py @@ -0,0 +1,396 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 Edge transformer variant with a Nemotron dense UND backbone.""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention + +from .transformer_cosmos3 import ( + Cosmos3VFMTransformer, + Qwen3VLTextRotaryEmbedding, + RMSNorm, + _apply_rotary_pos_emb, + _nested_get, + _tf_config_get, +) + +COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" + + +def _tf_config_nested_get(config: Any, key: str, default: Any = None) -> Any: + if config is None: + return default + if isinstance(config, dict): + value = _nested_get(config, key) + elif hasattr(config, "to_dict"): + value = _nested_get(config.to_dict(), key) + elif hasattr(config, "params"): + value = _nested_get(config.params, key) + else: + value = getattr(config, key, None) + return default if value is None else value + + +class Cosmos3Relu2MLP(nn.Module): + """Nemotron dense MLP: down_proj(relu(up_proj(x)) ** 2).""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + hidden = F.relu(self.up_proj(x)) + return self.down_proj(hidden * hidden) + + +class Cosmos3EdgeCausalAttention(nn.Module): + """Edge UND causal attention. + + Reasoner self-attention uses raw UND Q/K with RoPE because + ``qk_norm_for_text=false``. The GEN-facing UND key cache is a separate + normalized-then-RoPE'd view of the raw K tensor when GEN QK norm is enabled; + no-GEN-QK-norm checkpoints reuse the raw RoPE'd UND K. + """ + + def __init__( + self, + *, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + use_k_norm_und_for_gen: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + + tp_size = get_tensor_model_parallel_world_size() + self.num_heads_local = self.num_heads // tp_size + self.num_kv_heads_local = self.num_kv_heads // tp_size + + self.to_q = ColumnParallelLinear( + hidden_size, + self.num_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.to_out = RowParallelLinear( + self.num_heads * self.head_dim, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out", + ) + + self.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None + self.attn = FrameworkAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=True, + softmax_scale=1.0 / (self.head_dim**0.5), + num_kv_heads=self.num_kv_heads, + skip_sequence_parallel=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, S, _ = hidden_states.shape + + q = self.to_q(hidden_states).view(B, S, self.num_heads_local, self.head_dim) + k = self.to_k(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim) + v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim) + + q_reasoner, k_reasoner = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) + + if self.k_norm_und_for_gen is not None: + k_gen = F.rms_norm( + k, + (self.head_dim,), + self.k_norm_und_for_gen.weight, + eps=self.k_norm_und_for_gen.variance_epsilon, + ) + _, k_gen = _apply_rotary_pos_emb(q, k_gen, freqs_cos, freqs_sin) + else: + k_gen = k_reasoner + + out = self.attn(q_reasoner, k_reasoner, v).reshape(B, S, -1) + return self.to_out(out), k_gen, v + + +class Cosmos3EdgeUndDecoderLayer(nn.Module): + """Edge UND decoder layer: causal self-attention + ReLU2 MLP.""" + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + use_k_norm_und_for_gen: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = Cosmos3EdgeCausalAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + use_k_norm_und_for_gen=use_k_norm_und_for_gen, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.mlp = Cosmos3Relu2MLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + cos, sin = freqs + attn_out, k, v = self.self_attn(hidden_states, cos, sin) + hidden_states = residual + attn_out + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + self.mlp(hidden_states) + + return hidden_states, k, v + + +class Cosmos3EdgeLanguageModel(nn.Module): + """Nemotron dense UND tower that returns GEN-facing K plus raw V per layer.""" + + _layerwise_offload_blocks_attrs = ["layers"] + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + num_hidden_layers: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + vocab_size: int, + rms_norm_eps: float, + rope_theta: float, + mrope_section: list[int], + use_k_norm_und_for_gen: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, hidden_size) + self.rotary_emb = Qwen3VLTextRotaryEmbedding( + head_dim=head_dim, + rope_theta=rope_theta, + mrope_section=mrope_section, + ) + self.layers = nn.ModuleList( + [ + Cosmos3EdgeUndDecoderLayer( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + use_k_norm_und_for_gen=use_k_norm_und_for_gen, + quant_config=quant_config, + prefix=f"{prefix}.layers.{i}", + ) + for i in range(num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) + + def forward( + self, + text_ids: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + hidden = self.embed_tokens(text_ids) + + cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = [] + for layer in self.layers: + hidden, k, v = layer(hidden, freqs) + cached_kv.append((k, v)) + + return cached_kv + + +class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer): + """Cosmos3 Edge variant with Nemotron dense UND and shared GEN diffusion.""" + + _language_model_cls = Cosmos3EdgeLanguageModel + _gen_mlp_cls = Cosmos3Relu2MLP + + @staticmethod + def _validate_supported_config(model_config: Any) -> None: + expected_values = { + "qk_norm_for_text": False, + "position_embedding_type": "unified_3d_mrope", + "unified_3d_mrope_reset_spatial_ids": True, + "joint_attn_implementation": "two_way", + } + for key, expected in expected_values.items(): + actual = _tf_config_get(model_config, key, expected) + if actual != expected: + raise ValueError( + f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}." + ) + + backbone_type = _tf_config_get(model_config, "backbone_type", None) + if backbone_type is None: + raise ValueError( + f"Cosmos3 Edge transformer config must declare backbone_type={COSMOS3_EDGE_BACKBONE_TYPE!r}." + ) + if backbone_type != COSMOS3_EDGE_BACKBONE_TYPE: + raise ValueError( + "Unsupported Cosmos3 Edge transformer config: " + f"backbone_type={backbone_type!r}; expected {COSMOS3_EDGE_BACKBONE_TYPE!r}." + ) + + required_values = { + "latent_channel": 48, + "latent_patch_size": 2, + "temporal_compression_factor": 4, + } + for key, expected in required_values.items(): + actual = _tf_config_get(model_config, key, None) + if actual is None: + raise ValueError(f"Cosmos3 Edge transformer config must declare {key}={expected}.") + if int(actual) != expected: + raise ValueError( + f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}." + ) + + @classmethod + def _resolve_rms_norm_eps(cls, model_config: Any) -> float: + return float( + _tf_config_get( + model_config, + "rms_norm_eps", + _tf_config_get(model_config, "layer_norm_epsilon", 1e-5), + ) + ) + + @classmethod + def _resolve_rope_theta(cls, model_config: Any) -> float: + return float(_tf_config_get(model_config, "rope_theta", 100_000_000)) + + @classmethod + def _resolve_mrope_section(cls, model_config: Any) -> list[int]: + mrope_section = _tf_config_get(model_config, "mrope_section", None) + if mrope_section is None: + mrope_section = _tf_config_nested_get(model_config, "mrope_section", [24, 20, 20]) + return list(mrope_section) + + def _language_model_kwargs(self) -> dict[str, Any]: + return {"use_k_norm_und_for_gen": self.use_k_norm_und_for_gen} + + def validate_loaded_weights(self, loaded: set[str]) -> None: + missing: list[str] = [] + for layer_idx in range(self.num_hidden_layers): + required_markers = ( + f"language_model.layers.{layer_idx}.mlp.up_proj.", + f"language_model.layers.{layer_idx}.mlp.down_proj.", + f"gen_layers.{layer_idx}.mlp.up_proj.", + f"gen_layers.{layer_idx}.mlp.down_proj.", + ) + if self.use_k_norm_und_for_gen: + required_markers = ( + *required_markers, + f"language_model.layers.{layer_idx}.self_attn.k_norm_und_for_gen.", + ) + missing.extend( + marker.rstrip(".") for marker in required_markers if not any(marker in name for name in loaded) + ) + if missing: + preview = ", ".join(missing[:10]) + suffix = " ..." if len(missing) > 10 else "" + raise ValueError( + "Cosmos3 Edge transformer checkpoint is missing required weights for " + f"{preview}{suffix}. Use a converted Edge checkpoint with ReLU2 UND/GEN MLP weights." + ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 1454a37b1..5b8b2fba1 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -276,6 +276,11 @@ _DIFFUSION_MODELS = { "pipeline_cosmos3", "Cosmos3OmniDiffusersPipeline", ), + "Cosmos3OmniPipeline": ( + "cosmos3", + "pipeline_cosmos3", + "Cosmos3OmniDiffusersPipeline", + ), "SoulXSingerPipeline": ( "soulx_singer", "pipeline_soulx_singer_svs", @@ -536,6 +541,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func", + "Cosmos3OmniPipeline": "get_cosmos3_post_process_func", "HiDreamImagePipeline": "get_hidream_image_post_process_func", "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func", "Krea2Pipeline": "get_krea2_post_process_func", @@ -546,6 +552,7 @@ _DIFFUSION_IR_OP_PRIORITY_FUNCS = { # `ir_op_priority_func` function must be placed in {mod_folder}/{mod_relname}.py, # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key. "Cosmos3OmniDiffusersPipeline": "get_cosmos3_ir_op_priority_func", + "Cosmos3OmniPipeline": "get_cosmos3_ir_op_priority_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { @@ -570,6 +577,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func", "MagiHumanPipeline": "get_magi_human_pre_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", + "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func", "SoulXSingerPipeline": "get_soulxsinger_pre_process_func", "SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func", } diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 7103991c7..79f8b4360 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -153,6 +153,11 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS, "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt, }, + "Cosmos3OmniPipeline": { + "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS, + "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS, + "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt, + }, "MagiHumanPipeline": { "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS, "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,