[Bugfix] Fix DataType Handling in Default Diffusion Config (#2530)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/entrypoints/test_utils.py

@@ -5,14 +5,17 @@ from collections import Counter
 from dataclasses import dataclass
 
 import pytest
+import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.utils import (
     _convert_dataclasses_to_dict,
     _filter_dict_like_object,
     filter_dataclass_kwargs,
+    load_and_resolve_stage_configs,
     resolve_model_config_path,
 )
 
@@ -304,3 +307,18 @@ class TestResolveModelConfigPath:
 
         assert result is not None
         assert "glm_image.yaml" in result
+
+
+class TestLoadAndResolveStageConfigs:
+    def test_load_and_resolve_with_kwargs(self):
+        """Ensure that dtype survives default stage creation."""
+        kwargs = {"dtype": torch.float32}
+        config_path, stage_configs = load_and_resolve_stage_configs(
+            model="black-forest-labs/FLUX.2-klein-4B",
+            stage_configs_path=None,
+            kwargs=kwargs,
+            default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
+        )
+        assert config_path is None
+        assert len(stage_configs) == 1
+        assert "dtype" in stage_configs[0]["engine_args"]