Add Cosmos3 Edge and Distilled checkpoints support (#5001)

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

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

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
 

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",

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"]}

tests/model_tests/diffusion/test_alignment.py

@@ -71,6 +71,7 @@ EXCLUDED_MODELS = [
     "OmniVoicePipeline",
     "OmniVoice",
     "Cosmos3OmniDiffusersPipeline",
+    "Cosmos3OmniPipeline",
     "DiffusersAdapterPipeline",
     "HiDreamImagePipeline",
     "DreamZeroPipeline",

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,
     }
 )

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",
 ]

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:

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)

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."
+            )

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",
 }

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,