From 11a5842bd12255ee5ef4ab89340bceab6affc9b0 Mon Sep 17 00:00:00 2001
From: Alex Brooks <albrooks@redhat.com>
Date: Wed, 4 Mar 2026 07:49:28 -0700
Subject: [PATCH] [Bugfix] Fix Dtype Parsing (#1391)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
---
 tests/entrypoints/test_omni_diffusion.py | 89 ++++++++++++++++++++++++
 vllm_omni/entrypoints/omni.py            |  8 ++-
 2 files changed, 95 insertions(+), 2 deletions(-)

diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py
index 29e6dcf6a..90b0fd053 100644
--- a/tests/entrypoints/test_omni_diffusion.py
+++ b/tests/entrypoints/test_omni_diffusion.py
@@ -4,6 +4,7 @@ from queue import Empty, Queue
 from typing import Any
 
 import pytest
+import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
@@ -1340,3 +1341,91 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
 
     # Verify stop_stage_worker was called (process should be set)
     assert omni.stage_list[0]._proc is not None
+
+
+# Tests below are for diffusion dtype normalization fix from the following:
+# https://github.com/vllm-project/vllm-omni/pull/1391
+# In the future we should ensure dtypes are parsed in less hacky way and make
+# these tests more atomic.
+@pytest.mark.parametrize("dtype", ["float16", torch.float16])
+def test_dtype_normalization_valid_types(
+    monkeypatch, dtype: str | torch.dtype, mocker: MockerFixture, fake_stage_config
+):
+    """Ensure Diffusion Config builder coerces valid types correctly."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        # Return not stage configs to fall back to the diffusion cfg builder
+        return []
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+
+    omni = Omni(model="any", init_timeout=1, dtype=dtype)
+
+    # Dtype parsing being checked is on the diffusion path
+    assert len(omni.stage_configs) == 1
+    assert omni.stage_configs[0]["stage_type"] == "diffusion"
+
+    # Regardless of whether a str / dtype is passed, it resolves correctly
+    engine_args = omni.stage_configs[0]["engine_args"]
+    assert "dtype" in engine_args
+    assert isinstance(engine_args["dtype"], str)
+    assert engine_args["dtype"] == "float16"
+
+
+def test_dtype_normalization_invalid_types(monkeypatch, mocker: MockerFixture, fake_stage_config):
+    """Ensure Diffusion Config builder correctly handles bad dtype overrides."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        # Return not stage configs to fall back to the diffusion cfg builder
+        return []
+
+    class NotATorchDtype:
+        pass
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+
+    # Raise TypeError if we get an unrecognized type
+    with pytest.raises(TypeError):
+        Omni(model="any", init_timeout=1, dtype=NotATorchDtype)
diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py
index 58dd88eb0..9ecc5f762 100644
--- a/vllm_omni/entrypoints/omni.py
+++ b/vllm_omni/entrypoints/omni.py
@@ -13,6 +13,7 @@ from typing import Any, Literal, overload
 
 import huggingface_hub
 import msgspec.msgpack
+import torch
 import zmq
 from omegaconf import OmegaConf
 from tqdm.auto import tqdm
@@ -217,8 +218,11 @@ class OmniBase:
         # We temporally create a default config for diffusion stage.
         # In the future, we should merge the default config with the user-provided config.
         # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
-        if "dtype" in kwargs:
-            kwargs["dtype"] = str(kwargs["dtype"])
+        if "dtype" in kwargs and not isinstance(kwargs["dtype"], str):
+            if not isinstance(kwargs["dtype"], torch.dtype):
+                raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(kwargs['dtype']).__name__}")
+            kwargs["dtype"] = str(kwargs["dtype"]).removeprefix("torch.")
+
         cache_backend = kwargs.get("cache_backend", "none")
         cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None))
         # TODO: hack, calculate devices based on parallel config.