Address code review feedback

- Remove unused `field` import
- Move `json` import to top of file
- Add comment explaining local DiffusionParallelConfig import

Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>

tests/entrypoints/test_stage_config.py

@@ -0,0 +1,396 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for structured stage configuration builder."""
+
+import json
+from unittest.mock import Mock
+
+import pytest
+from omegaconf import OmegaConf
+
+from vllm_omni.entrypoints.stage_config import (
+    StageConfig,
+    StageConfigFactory,
+    StageRuntimeConfig,
+)
+
+
+class TestStageRuntimeConfig:
+    """Tests for StageRuntimeConfig dataclass."""
+    
+    def test_runtime_config_creation(self):
+        """Test basic runtime config creation."""
+        config = StageRuntimeConfig(
+            process=True,
+            devices="0,1,2",
+            max_batch_size=4
+        )
+        
+        assert config.process is True
+        assert config.devices == "0,1,2"
+        assert config.max_batch_size == 4
+
+
+class TestStageConfig:
+    """Tests for StageConfig dataclass."""
+    
+    def test_stage_config_creation(self):
+        """Test basic stage config creation."""
+        runtime = StageRuntimeConfig(
+            process=True,
+            devices="0",
+            max_batch_size=1
+        )
+        
+        config = StageConfig(
+            stage_id=0,
+            stage_type="diffusion",
+            runtime=runtime,
+            engine_args={"model": "test-model"},
+            final_output=True,
+            final_output_type="image"
+        )
+        
+        assert config.stage_id == 0
+        assert config.stage_type == "diffusion"
+        assert config.runtime.devices == "0"
+        assert config.engine_args["model"] == "test-model"
+        assert config.final_output is True
+        assert config.final_output_type == "image"
+    
+    def test_to_omegaconf(self):
+        """Test conversion to OmegaConf DictConfig."""
+        runtime = StageRuntimeConfig(
+            process=True,
+            devices="0,1",
+            max_batch_size=2
+        )
+        
+        config = StageConfig(
+            stage_id=0,
+            stage_type="diffusion",
+            runtime=runtime,
+            engine_args={"cache_backend": "cache_dit"},
+            final_output=True,
+            final_output_type="image"
+        )
+        
+        omega_config = config.to_omegaconf()
+        
+        assert omega_config.stage_id == 0
+        assert omega_config.stage_type == "diffusion"
+        assert omega_config.runtime.process is True
+        assert omega_config.runtime.devices == "0,1"
+        assert omega_config.runtime.max_batch_size == 2
+        assert omega_config.engine_args.cache_backend == "cache_dit"
+        assert omega_config.final_output is True
+        assert omega_config.final_output_type == "image"
+
+
+class TestDeviceStringGeneration:
+    """Tests for device string generation helper."""
+    
+    def test_device_string_from_parallel_config(self):
+        """Test device string generation from parallel_config."""
+        # Mock parallel_config with world_size
+        parallel_config = Mock()
+        parallel_config.world_size = 4
+        
+        devices = StageConfigFactory._get_device_string(parallel_config)
+        
+        assert devices == "0,1,2,3"
+    
+    def test_device_string_from_num_devices(self):
+        """Test device string generation from num_devices parameter."""
+        devices = StageConfigFactory._get_device_string(None, num_devices=3)
+        
+        assert devices == "0,1,2"
+    
+    def test_device_string_default_single_device(self):
+        """Test default device string is single device."""
+        devices = StageConfigFactory._get_device_string(None, None)
+        
+        assert devices == "0"
+    
+    def test_device_string_single_device(self):
+        """Test device string with single device."""
+        parallel_config = Mock()
+        parallel_config.world_size = 1
+        
+        devices = StageConfigFactory._get_device_string(parallel_config)
+        
+        assert devices == "0"
+
+
+class TestCacheConfigNormalization:
+    """Tests for cache configuration normalization."""
+    
+    def test_cache_dit_defaults(self):
+        """Test default cache_dit configuration."""
+        config = StageConfigFactory._get_default_cache_config("cache_dit")
+        
+        assert config is not None
+        assert config["Fn_compute_blocks"] == 1
+        assert config["Bn_compute_blocks"] == 0
+        assert config["max_warmup_steps"] == 4
+        assert config["residual_diff_threshold"] == 0.24
+        assert config["max_continuous_cached_steps"] == 3
+        assert config["enable_taylorseer"] is False
+        assert config["taylorseer_order"] == 1
+        assert config["scm_steps_mask_policy"] is None
+        assert config["scm_steps_policy"] == "dynamic"
+    
+    def test_tea_cache_defaults(self):
+        """Test default tea_cache configuration."""
+        config = StageConfigFactory._get_default_cache_config("tea_cache")
+        
+        assert config is not None
+        assert config["rel_l1_thresh"] == 0.2
+    
+    def test_no_cache_defaults(self):
+        """Test no defaults for unknown cache backend."""
+        config = StageConfigFactory._get_default_cache_config("unknown")
+        
+        assert config is None
+    
+    def test_normalize_json_string_cache_config(self):
+        """Test normalization of JSON string cache_config."""
+        cache_config_str = '{"Fn_compute_blocks": 2, "max_warmup_steps": 6}'
+        
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            cache_config_str
+        )
+        
+        assert normalized is not None
+        assert normalized["Fn_compute_blocks"] == 2
+        assert normalized["max_warmup_steps"] == 6
+    
+    def test_normalize_invalid_json_uses_defaults(self):
+        """Test that invalid JSON falls back to defaults."""
+        cache_config_str = '{invalid json}'
+        
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            cache_config_str
+        )
+        
+        # Should fall back to defaults for cache_dit
+        assert normalized is not None
+        assert normalized["Fn_compute_blocks"] == 1
+    
+    def test_normalize_dict_cache_config(self):
+        """Test normalization of dict cache_config."""
+        cache_config_dict = {"Fn_compute_blocks": 3}
+        
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            cache_config_dict
+        )
+        
+        assert normalized == cache_config_dict
+    
+    def test_normalize_none_with_backend_uses_defaults(self):
+        """Test that None cache_config with backend uses defaults."""
+        normalized = StageConfigFactory._normalize_cache_config(
+            "cache_dit", 
+            None
+        )
+        
+        assert normalized is not None
+        assert normalized["Fn_compute_blocks"] == 1
+    
+    def test_normalize_none_without_backend_returns_none(self):
+        """Test that None cache_config without backend returns None."""
+        normalized = StageConfigFactory._normalize_cache_config(
+            "none", 
+            None
+        )
+        
+        assert normalized is None
+
+
+class TestDefaultDiffusionConfigBuilder:
+    """Tests for default diffusion stage configuration builder."""
+    
+    def test_default_diffusion_config_structure(self):
+        """Test that the builder produces correct schema structure."""
+        kwargs = {
+            "model": "test-model",
+            "dtype": "float16",
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        
+        assert len(configs) == 1
+        config = configs[0]
+        
+        # Verify structure
+        assert config.stage_id == 0
+        assert config.stage_type == "diffusion"
+        assert config.runtime.process is True
+        assert config.runtime.devices == "0"
+        assert config.runtime.max_batch_size == 1
+        assert config.final_output is True
+        assert config.final_output_type == "image"
+        
+        # Verify engine_args
+        assert config.engine_args.model == "test-model"
+        assert config.engine_args.dtype == "float16"  # Should be stringified
+        assert config.engine_args.model_stage == "diffusion"
+        assert config.engine_args.cache_backend == "none"
+    
+    def test_default_diffusion_with_parallel_config(self):
+        """Test builder with parallel configuration."""
+        parallel_config = Mock()
+        parallel_config.world_size = 4
+        
+        kwargs = {
+            "model": "test-model",
+            "parallel_config": parallel_config,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should generate device string for 4 devices
+        assert config.runtime.devices == "0,1,2,3"
+    
+    def test_default_diffusion_with_cache_backend(self):
+        """Test builder with cache backend."""
+        kwargs = {
+            "model": "test-model",
+            "cache_backend": "cache_dit",
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should have cache_dit defaults
+        assert config.engine_args.cache_backend == "cache_dit"
+        assert config.engine_args.cache_config is not None
+        assert config.engine_args.cache_config["Fn_compute_blocks"] == 1
+    
+    def test_default_diffusion_preserves_kwargs(self):
+        """Test that builder preserves additional kwargs in engine_args."""
+        kwargs = {
+            "model": "test-model",
+            "custom_param": "custom_value",
+            "another_param": 42,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should preserve custom parameters
+        assert config.engine_args.custom_param == "custom_value"
+        assert config.engine_args.another_param == 42
+    
+    def test_dtype_stringification(self):
+        """Test that dtype is converted to string."""
+        # Use a mock object for dtype
+        class DType:
+            def __str__(self):
+                return "torch.float16"
+        
+        kwargs = {
+            "model": "test-model",
+            "dtype": DType(),
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        assert config.engine_args.dtype == "torch.float16"
+
+
+class TestAsyncDiffusionConfigBuilder:
+    """Tests for async diffusion stage configuration builder."""
+    
+    def test_async_diffusion_with_parallel_config(self):
+        """Test async builder with parallel_config."""
+        parallel_config = Mock()
+        parallel_config.world_size = 2
+        
+        kwargs = {
+            "parallel_config": parallel_config,
+            "vae_use_slicing": True,
+            "vae_use_tiling": False,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion_async(kwargs)
+        config = configs[0]
+        
+        assert config.runtime.devices == "0,1"
+        assert config.engine_args.parallel_config == parallel_config
+        assert config.engine_args.vae_use_slicing is True
+        assert config.engine_args.vae_use_tiling is False
+    
+    def test_async_diffusion_builds_parallel_config(self):
+        """Test async builder builds parallel_config from parameters."""
+        kwargs = {
+            "tensor_parallel_size": 2,
+            "ulysses_degree": 2,
+            "ring_degree": 1,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion_async(kwargs)
+        config = configs[0]
+        
+        # Should build parallel_config and device string
+        # tensor_parallel_size=2 * sequence_parallel_size=2 = 4 devices
+        assert config.runtime.devices == "0,1,2,3"
+        assert config.engine_args.parallel_config is not None
+        assert config.engine_args.parallel_config.tensor_parallel_size == 2
+        assert config.engine_args.parallel_config.ulysses_degree == 2
+    
+    def test_async_diffusion_includes_async_specific_args(self):
+        """Test that async builder includes async-specific engine args."""
+        kwargs = {
+            "enable_cpu_offload": True,
+            "enforce_eager": True,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion_async(kwargs)
+        config = configs[0]
+        
+        assert config.engine_args.enable_cpu_offload is True
+        assert config.engine_args.enforce_eager is True
+        assert config.engine_args.model_stage == "diffusion"
+
+
+class TestCacheConfigIntegration:
+    """Integration tests for cache config normalization."""
+    
+    def test_cache_backend_integration(self):
+        """Test cache_backend arguments are correctly nested in config."""
+        kwargs = {
+            "model": "test-model",
+            "cache_backend": "cache_dit",
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Verify cache_dit defaults are in engine_args
+        assert config.engine_args.cache_backend == "cache_dit"
+        assert config.engine_args.cache_config is not None
+        assert "Fn_compute_blocks" in config.engine_args.cache_config
+        assert "max_warmup_steps" in config.engine_args.cache_config
+    
+    def test_custom_cache_config_override(self):
+        """Test custom cache_config overrides defaults."""
+        custom_config = {"Fn_compute_blocks": 5, "custom_key": "custom_value"}
+        
+        kwargs = {
+            "model": "test-model",
+            "cache_backend": "cache_dit",
+            "cache_config": custom_config,
+        }
+        
+        configs = StageConfigFactory.create_default_diffusion(kwargs)
+        config = configs[0]
+        
+        # Should use custom config
+        assert config.engine_args.cache_config["Fn_compute_blocks"] == 5
+        assert config.engine_args.cache_config["custom_key"] == "custom_value"

vllm_omni/entrypoints/async_omni.py

@@ -29,6 +29,7 @@ from vllm_omni.entrypoints.log_utils import (
 )
 from vllm_omni.entrypoints.omni import OmniBase
 from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.stage_config import StageConfigFactory
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
@@ -115,61 +116,19 @@ class AsyncOmni(OmniBase):
             self.output_handler,
         )
 
-    def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+    def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> list[Any]:
-        """Create default diffusion stage configuration."""
+        """Create default diffusion stage configuration using StageConfigFactory.
-        # TODO: here is different from the Omni class. We should merge the two in the future.
+        
-        cache_backend = kwargs.get("cache_backend", "none")
+        This method uses the structured factory pattern to build a robust,
-        cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None))
+        type-safe stage configuration for async mode.
-
+        
-        devices = "0"
+        Args:
-        if "parallel_config" in kwargs:
+            kwargs: Dictionary of configuration parameters
-            parallel_config = kwargs["parallel_config"]
+            
-            num_devices = kwargs["parallel_config"].world_size
+        Returns:
-            for i in range(1, num_devices):
+            List containing a single OmegaConf DictConfig for the diffusion stage
-                devices += f",{i}"
+        """
-        else:
+        return StageConfigFactory.create_default_diffusion_async(kwargs)
-            ulysses_degree = kwargs.get("ulysses_degree") or 1
-            ring_degree = kwargs.get("ring_degree") or 1
-            sequence_parallel_size = kwargs.get("sequence_parallel_size")
-            tensor_parallel_size = kwargs.get("tensor_parallel_size") or 1
-            if sequence_parallel_size is None:
-                sequence_parallel_size = ulysses_degree * ring_degree
-            num_devices = sequence_parallel_size * tensor_parallel_size
-            for i in range(1, num_devices):
-                devices += f",{i}"
-            parallel_config = DiffusionParallelConfig(
-                pipeline_parallel_size=1,
-                data_parallel_size=1,
-                tensor_parallel_size=tensor_parallel_size,
-                sequence_parallel_size=sequence_parallel_size,
-                ulysses_degree=ulysses_degree,
-                ring_degree=ring_degree,
-                cfg_parallel_size=1,
-            )
-        default_stage_cfg = [
-            {
-                "stage_id": 0,
-                "stage_type": "diffusion",
-                "runtime": {
-                    "process": True,
-                    "devices": devices,
-                    "max_batch_size": 1,
-                },
-                "engine_args": {
-                    "parallel_config": parallel_config,
-                    "vae_use_slicing": kwargs.get("vae_use_slicing", False),
-                    "vae_use_tiling": kwargs.get("vae_use_tiling", False),
-                    "cache_backend": cache_backend,
-                    "cache_config": cache_config,
-                    "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
-                    "enforce_eager": kwargs.get("enforce_eager", False),
-                },
-                "final_output": True,
-                "final_output_type": "image",
-            }
-        ]
-        default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
-        return default_stage_cfg
 
     def _process_stage_ready(self, stage: OmniStage, stage_id: int, result: dict[str, Any]) -> None:
         # Store vllm_config received from worker process (may be None for diffusion stages)

vllm_omni/entrypoints/omni.py

@@ -30,6 +30,7 @@ from vllm_omni.distributed.ray_utils.utils import (
 )
 from vllm_omni.entrypoints.log_utils import OrchestratorMetrics
 from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.stage_config import StageConfigFactory
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
@@ -123,72 +124,26 @@ class OmniBase:
         self._initialize_stages(model, kwargs)
 
     def _get_default_cache_config(self, cache_backend: str | None) -> dict[str, Any] | None:
-        if cache_backend == "cache_dit":
+        """Delegate to StageConfigFactory for cache config defaults."""
-            return {
+        return StageConfigFactory._get_default_cache_config(cache_backend)
-                "Fn_compute_blocks": 1,
-                "Bn_compute_blocks": 0,
-                "max_warmup_steps": 4,
-                "residual_diff_threshold": 0.24,
-                "max_continuous_cached_steps": 3,
-                "enable_taylorseer": False,
-                "taylorseer_order": 1,
-                "scm_steps_mask_policy": None,
-                "scm_steps_policy": "dynamic",
-            }
-        if cache_backend == "tea_cache":
-            return {
-                "rel_l1_thresh": 0.2,
-            }
-        return None
 
     def _normalize_cache_config(self, cache_backend: str | None, cache_config: Any | None) -> Any | None:
-        if isinstance(cache_config, str):
+        """Delegate to StageConfigFactory for cache config normalization."""
-            try:
+        return StageConfigFactory._normalize_cache_config(cache_backend, cache_config)
-                cache_config = json.loads(cache_config)
-            except json.JSONDecodeError:
-                logger.warning("Invalid cache_config JSON, using defaults.")
-                cache_config = None
-        if cache_config is None and cache_backend not in (None, "", "none"):
-            cache_config = self._get_default_cache_config(cache_backend)
-        return cache_config
 
-    def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+    def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> list[Any]:
-        """Create default diffusion stage configuration."""
+        """Create default diffusion stage configuration using StageConfigFactory.
-        # We temporally create a default config for diffusion stage.
+        
-        # In the future, we should merge the default config with the user-provided config.
+        This method uses the structured factory pattern to build a robust,
-        # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
+        type-safe stage configuration instead of manually constructing dictionaries.
-        if "dtype" in kwargs:
+        
-            kwargs["dtype"] = str(kwargs["dtype"])
+        Args:
-        cache_backend = kwargs.get("cache_backend", "none")
+            kwargs: Dictionary of configuration parameters
-        cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None))
+            
-        # TODO: hack, calculate devices based on parallel config.
+        Returns:
-        devices = "0"
+            List containing a single OmegaConf DictConfig for the diffusion stage
-        if "parallel_config" in kwargs:
+        """
-            num_devices = kwargs["parallel_config"].world_size
+        return StageConfigFactory.create_default_diffusion(kwargs)
-            for i in range(1, num_devices):
-                devices += f",{i}"
-        default_stage_cfg = [
-            {
-                "stage_id": 0,
-                "stage_type": "diffusion",
-                "runtime": {
-                    "process": True,
-                    "devices": devices,
-                    "max_batch_size": 1,
-                },
-                "engine_args": OmegaConf.create(
-                    {
-                        **kwargs,
-                        "cache_backend": cache_backend,
-                        "cache_config": cache_config,
-                    }
-                ),
-                "final_output": True,
-                "final_output_type": "image",
-            }
-        ]
-        default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
-        return default_stage_cfg
 
     def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
         """Initialize stage list management."""

vllm_omni/entrypoints/stage_config.py

@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Structured dataclass definitions and factory for stage configurations.
+
+This module provides a robust, type-safe approach to building stage configurations
+programmatically, replacing ad-hoc dictionary construction.
+"""
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from omegaconf import DictConfig, OmegaConf
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class StageRuntimeConfig:
+    """Runtime configuration for a stage.
+    
+    Attributes:
+        process: Whether to run this stage in a separate process
+        devices: Comma-separated string of device IDs (e.g., "0,1,2")
+        max_batch_size: Maximum batch size for this stage
+    """
+    process: bool
+    devices: str
+    max_batch_size: int
+
+
+@dataclass
+class StageConfig:
+    """Complete configuration for a single stage in the pipeline.
+    
+    Attributes:
+        stage_id: Unique identifier for this stage
+        stage_type: Type of stage (e.g., "diffusion", "llm")
+        runtime: Runtime configuration for this stage
+        engine_args: Dictionary of engine-specific arguments
+        final_output: Whether this stage produces final output
+        final_output_type: Type of final output (e.g., "image", "text", "audio")
+    """
+    stage_id: int
+    stage_type: str
+    runtime: StageRuntimeConfig
+    engine_args: dict[str, Any]
+    final_output: bool
+    final_output_type: str
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Converts the dataclass to a plain dictionary.
+        
+        Returns:
+            Dictionary with stage configuration
+        """
+        return {
+            "stage_id": self.stage_id,
+            "stage_type": self.stage_type,
+            "runtime": {
+                "process": self.runtime.process,
+                "devices": self.runtime.devices,
+                "max_batch_size": self.runtime.max_batch_size,
+            },
+            "engine_args": self.engine_args,  # Pass through as-is
+            "final_output": self.final_output,
+            "final_output_type": self.final_output_type,
+        }
+    
+    def to_omegaconf(self) -> DictConfig:
+        """Converts the dataclass to an OmegaConf DictConfig.
+        
+        Returns:
+            OmegaConf DictConfig compatible with OmniStage initialization
+        """
+        # Convert to dict first, then wrap in OmegaConf
+        # This allows engine_args to remain as whatever type it already is
+        return OmegaConf.create(self.to_dict())
+
+
+class StageConfigFactory:
+    """Factory for creating stage configurations with validation and normalization.
+    
+    This factory provides methods to create structured StageConfig objects
+    with proper validation, type checking, and normalization of parameters.
+    """
+    
+    @staticmethod
+    def _get_device_string(parallel_config: Any | None, num_devices: int | None = None) -> str:
+        """Generate device string from parallel configuration or device count.
+        
+        Args:
+            parallel_config: Parallel configuration object with world_size attribute
+            num_devices: Number of devices (used if parallel_config is None)
+            
+        Returns:
+            Comma-separated device string (e.g., "0,1,2,3")
+        """
+        if parallel_config is not None:
+            device_count = parallel_config.world_size
+        elif num_devices is not None:
+            device_count = num_devices
+        else:
+            device_count = 1
+        
+        return ",".join(str(i) for i in range(device_count))
+    
+    @staticmethod
+    def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | None:
+        """Get default cache configuration for a given backend.
+        
+        Args:
+            cache_backend: Name of the cache backend (e.g., "cache_dit", "tea_cache")
+            
+        Returns:
+            Dictionary with default cache configuration, or None if no defaults exist
+        """
+        if cache_backend == "cache_dit":
+            return {
+                "Fn_compute_blocks": 1,
+                "Bn_compute_blocks": 0,
+                "max_warmup_steps": 4,
+                "residual_diff_threshold": 0.24,
+                "max_continuous_cached_steps": 3,
+                "enable_taylorseer": False,
+                "taylorseer_order": 1,
+                "scm_steps_mask_policy": None,
+                "scm_steps_policy": "dynamic",
+            }
+        if cache_backend == "tea_cache":
+            return {
+                "rel_l1_thresh": 0.2,
+            }
+        return None
+    
+    @staticmethod
+    def _normalize_cache_config(cache_backend: str | None, cache_config: Any | None) -> Any | None:
+        """Normalize cache configuration, applying defaults if needed.
+        
+        Args:
+            cache_backend: Name of the cache backend
+            cache_config: User-provided cache configuration (may be JSON string or dict)
+            
+        Returns:
+            Normalized cache configuration dictionary or None
+        """
+        if isinstance(cache_config, str):
+            try:
+                cache_config = json.loads(cache_config)
+            except json.JSONDecodeError:
+                logger.warning("Invalid cache_config JSON, using defaults.")
+                cache_config = None
+        
+        if cache_config is None and cache_backend not in (None, "", "none"):
+            cache_config = StageConfigFactory._get_default_cache_config(cache_backend)
+        
+        return cache_config
+    
+    @classmethod
+    def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[DictConfig]:
+        """Create default diffusion stage configuration from kwargs.
+        
+        This method builds a structured StageConfig for a diffusion model
+        with proper validation and normalization of all parameters.
+        
+        Args:
+            kwargs: Dictionary of configuration parameters including:
+                - dtype: Data type for the model (converted to string)
+                - cache_backend: Cache backend name (default: "none")
+                - cache_config: Cache configuration (optional)
+                - parallel_config: Parallel configuration with world_size attribute
+                - Additional engine arguments to pass through
+                
+        Returns:
+            List containing a single OmegaConf DictConfig for the diffusion stage
+        """
+        # Convert dtype to string to avoid OmegaConf serialization issues
+        if "dtype" in kwargs:
+            kwargs["dtype"] = str(kwargs["dtype"])
+        
+        # Normalize cache configuration
+        cache_backend = kwargs.get("cache_backend", "none")
+        cache_config = cls._normalize_cache_config(
+            cache_backend, 
+            kwargs.get("cache_config", None)
+        )
+        
+        # Generate device string from parallel configuration
+        devices = cls._get_device_string(kwargs.get("parallel_config"))
+        
+        # Build runtime configuration
+        runtime_config = StageRuntimeConfig(
+            process=True,
+            devices=devices,
+            max_batch_size=1,
+        )
+        
+        # Build engine arguments with cache config
+        # Wrap in OmegaConf.create to match original behavior
+        engine_args = OmegaConf.create({
+            **kwargs,
+            "cache_backend": cache_backend,
+            "cache_config": cache_config,
+        })
+        # Set model_stage after OmegaConf creation to match original behavior
+        engine_args["model_stage"] = "diffusion"
+        
+        # Create structured stage configuration
+        stage_config = StageConfig(
+            stage_id=0,
+            stage_type="diffusion",
+            runtime=runtime_config,
+            engine_args=engine_args,
+            final_output=True,
+            final_output_type="image",
+        )
+        
+        # Convert to OmegaConf for compatibility with existing code
+        return [stage_config.to_omegaconf()]
+    
+    @classmethod
+    def create_default_diffusion_async(cls, kwargs: dict[str, Any]) -> list[dict]:
+        """Create default diffusion stage configuration for async mode.
+        
+        This is similar to create_default_diffusion but handles async-specific
+        parallel configuration logic where parallel_config may be built from
+        individual parameters if not provided.
+        
+        Args:
+            kwargs: Dictionary of configuration parameters
+            
+        Returns:
+            List containing a single dict for the diffusion stage (not wrapped in OmegaConf)
+        """
+        # Local import to avoid potential circular dependencies with diffusion module
+        from vllm_omni.diffusion.data import DiffusionParallelConfig
+        
+        # Normalize cache configuration
+        cache_backend = kwargs.get("cache_backend", "none")
+        cache_config = cls._normalize_cache_config(
+            cache_backend, 
+            kwargs.get("cache_config", None)
+        )
+        
+        # Build or extract parallel_config
+        if "parallel_config" in kwargs:
+            parallel_config = kwargs["parallel_config"]
+            devices = cls._get_device_string(parallel_config)
+        else:
+            # Build parallel_config from individual parameters
+            ulysses_degree = kwargs.get("ulysses_degree") or 1
+            ring_degree = kwargs.get("ring_degree") or 1
+            sequence_parallel_size = kwargs.get("sequence_parallel_size")
+            tensor_parallel_size = kwargs.get("tensor_parallel_size") or 1
+            if sequence_parallel_size is None:
+                sequence_parallel_size = ulysses_degree * ring_degree
+            num_devices = sequence_parallel_size * tensor_parallel_size
+            devices = cls._get_device_string(None, num_devices)
+            parallel_config = DiffusionParallelConfig(
+                pipeline_parallel_size=1,
+                data_parallel_size=1,
+                tensor_parallel_size=tensor_parallel_size,
+                sequence_parallel_size=sequence_parallel_size,
+                ulysses_degree=ulysses_degree,
+                ring_degree=ring_degree,
+                cfg_parallel_size=1,
+            )
+        
+        # Build runtime configuration
+        runtime_config = StageRuntimeConfig(
+            process=True,
+            devices=devices,
+            max_batch_size=1,
+        )
+        
+        # Build engine arguments specific to async mode
+        # Note: We don't wrap in OmegaConf.create here because parallel_config
+        # may not be serializable. The original code used a plain dict.
+        engine_args = {
+            "parallel_config": parallel_config,
+            "vae_use_slicing": kwargs.get("vae_use_slicing", False),
+            "vae_use_tiling": kwargs.get("vae_use_tiling", False),
+            "cache_backend": cache_backend,
+            "cache_config": cache_config,
+            "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
+            "enforce_eager": kwargs.get("enforce_eager", False),
+            "model_stage": "diffusion",
+        }
+        
+        # Create structured stage configuration
+        stage_config = StageConfig(
+            stage_id=0,
+            stage_type="diffusion",
+            runtime=runtime_config,
+            engine_args=engine_args,
+            final_output=True,
+            final_output_type="image",
+        )
+        
+        # Return as plain dict for async mode (engine_args contains non-serializable parallel_config)
+        return [stage_config.to_dict()]