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>

docs/.nav.yml

@@ -12,7 +12,6 @@ nav:
     - Offline Inference:
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - LTX-2: user_guide/examples/offline_inference/ltx2.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md

docs/models/supported_models.md

@@ -34,7 +34,6 @@ th {
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
-|`LTX2Pipeline` | LTX-2 | `Lightricks/LTX-2` |
 
 
 ## List of Supported Models for NPU

docs/user_guide/examples/offline_inference/ltx2.md

@@ -1,42 +0,0 @@
-# LTX-2 Text-To-Video
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/ltx2>.
-
-The `Lightricks/LTX-2` pipeline generates high-quality videos from text prompts using a 19B parameter DiT architecture. It supports up to 4K video generation with synchronized audio.
-
-## Local CLI Usage
-
-```bash
-python text_to_video.py \
-  --prompt "A panda riding a bicycle through a forest, cinematic lighting" \
-  --height 512 \
-  --width 768 \
-  --num_frames 121 \
-  --num_inference_steps 40 \
-  --guidance_scale 4.0 \
-  --output ltx2_output.mp4
-```
-
-Key arguments:
-
-- `--prompt`: Text description of the video to generate (string).
-- `--height/--width`: Output resolution (defaults 512x768). Dimensions must be divisible by 32.
-- `--num_frames`: Number of frames (should be 8*n+1, e.g., 25, 81, 121). Default is 121.
-- `--guidance_scale`: Classifier-free guidance scale (default: 4.0). Range 3.0-5.0 recommended.
-- `--negative_prompt`: Optional text describing what to avoid in the video.
-- `--num_inference_steps`: Number of denoising steps (default: 40).
-- `--seed`: Random seed for reproducibility (default: 42).
-- `--fps`: Frames per second for the saved MP4 (default: 24).
-- `--output`: Path to save the generated video.
-
-## Example materials
-
-??? abstract "text_to_video.py"
-    ``````py
-    --8<-- "examples/offline_inference/ltx2/text_to_video.py"
-    ``````
-
-??? abstract "text_to_video.md"
-    ``````md
-    --8<-- "examples/offline_inference/ltx2/text_to_video.md"
-    ``````

examples/offline_inference/ltx2/text_to_video.md

@@ -1,77 +0,0 @@
-# LTX-2 Text-to-Video Generation
-
-This example demonstrates how to use Lightricks LTX-2 model for text-to-video generation with vLLM-Omni.
-
-## Model Information
-
-LTX-2 is a state-of-the-art text-to-video model by Lightricks that can generate high-quality videos from text prompts. It supports:
-- Text-to-video generation
-- Image-to-video generation (conditioned on an input image)
-- High-resolution video generation (up to 4K)
-- Synchronized audio-video generation
-
-For more information, see:
-- [Lightricks/LTX-2 on HuggingFace](https://huggingface.co/Lightricks/LTX-2)
-- [LTX-2 GitHub Repository](https://github.com/Lightricks/LTX-2)
-
-## Requirements
-
-- Python >= 3.10
-- diffusers >= 0.36.0 (already included in vllm-omni dependencies)
-- GPU with sufficient VRAM (recommended: >= 24GB for high-quality generation)
-
-## Usage
-
-### Basic Text-to-Video Generation
-
-```bash
-python text_to_video.py \
-    --model Lightricks/LTX-2 \
-    --prompt "A panda riding a bicycle through a forest, cinematic lighting" \
-    --height 512 \
-    --width 768 \
-    --num_frames 121 \
-    --num_inference_steps 40 \
-    --guidance_scale 4.0 \
-    --output output.mp4
-```
-
-### Parameters
-
-- `--model`: Model ID or local path (default: `Lightricks/LTX-2`)
-- `--prompt`: Text description of the video to generate
-- `--negative_prompt`: Text describing what to avoid in the video
-- `--height`: Video height in pixels, must be divisible by 32 (default: 512)
-- `--width`: Video width in pixels, must be divisible by 32 (default: 768)
-- `--num_frames`: Number of frames, should be 8*n+1 (e.g., 25, 81, 121) (default: 121)
-- `--num_inference_steps`: Number of denoising steps (default: 40)
-- `--guidance_scale`: Classifier-free guidance scale (default: 4.0)
-- `--seed`: Random seed for reproducibility (default: 42)
-- `--fps`: Output video frames per second (default: 24)
-- `--output`: Output video file path (default: ltx2_output.mp4)
-
-### Example Prompts
-
-```bash
-# Cinematic scene
-python text_to_video.py --prompt "A serene lakeside sunrise with mist over the water, cinematic"
-
-# Action scene
-python text_to_video.py --prompt "A futuristic cityscape with flying cars, neon lights, cyberpunk style"
-
-# Nature scene
-python text_to_video.py --prompt "A waterfall in a lush rainforest, birds flying overhead"
-```
-
-## Performance Tips
-
-1. **Resolution**: Start with lower resolutions (512x768) for faster generation. Higher resolutions require more VRAM.
-2. **Frames**: More frames take longer to generate. Start with 25 or 81 frames for testing.
-3. **Inference Steps**: 40 steps provide good quality. Fewer steps are faster but may reduce quality.
-4. **Guidance Scale**: Values between 3.0-5.0 typically work well. Higher values follow the prompt more strictly.
-
-## Notes
-
-- The first run will download the model from HuggingFace (approximately 20-30GB).
-- Generation time depends on your GPU and the parameters chosen.
-- For best results, use descriptive prompts with style keywords (e.g., "cinematic", "high quality", "detailed").

examples/offline_inference/ltx2/text_to_video.py

@@ -1,124 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-from pathlib import Path
-
-import numpy as np
-import torch
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video with LTX-2 T2V.")
-    parser.add_argument(
-        "--model",
-        default="Lightricks/LTX-2",
-        help="LTX-2 model ID or local path.",
-    )
-    parser.add_argument(
-        "--prompt",
-        default="A serene lakeside sunrise with mist over the water.",
-        help="Text prompt for video generation.",
-    )
-    parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--guidance_scale", type=float, default=4.0, help="CFG scale.")
-    parser.add_argument("--height", type=int, default=512, help="Video height (must be divisible by 32).")
-    parser.add_argument("--width", type=int, default=768, help="Video width (must be divisible by 32).")
-    parser.add_argument(
-        "--num_frames", type=int, default=121, help="Number of frames (should be 8*n+1, e.g., 25, 81, 121)."
-    )
-    parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
-    parser.add_argument("--output", type=str, default="ltx2_output.mp4", help="Path to save the video (mp4).")
-    parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
-
-    omni = Omni(
-        model=args.model,
-    )
-
-    frames = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
-    )
-
-    # Extract video frames from OmniRequestOutput
-    if isinstance(frames, list) and len(frames) > 0:
-        first_item = frames[0]
-
-        # Check if it's an OmniRequestOutput
-        if hasattr(first_item, "final_output_type"):
-            if first_item.final_output_type != "image":
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', "
-                    f"expected 'image' for video frames output."
-                )
-
-            # Pipeline mode: extract from nested request_output
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                        frames = inner_output.images[0] if inner_output.images else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
-            # Diffusion mode: use direct images field
-            elif hasattr(first_item, "images") and first_item.images:
-                frames = first_item.images
-            else:
-                raise ValueError("No video frames found in OmniRequestOutput.")
-
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    try:
-        from diffusers.utils import export_to_video
-    except ImportError:
-        raise ImportError("diffusers is required for export_to_video.")
-
-    # frames may be np.ndarray (preferred) or torch.Tensor
-    # export_to_video expects a list of frames with values in [0, 1]
-    if isinstance(frames, torch.Tensor):
-        video_tensor = frames.detach().cpu()
-        if video_tensor.dim() == 5:
-            # [B, C, F, H, W] or [B, F, H, W, C]
-            if video_tensor.shape[1] in (3, 4):
-                video_tensor = video_tensor[0].permute(1, 2, 3, 0)
-            else:
-                video_tensor = video_tensor[0]
-        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
-            video_tensor = video_tensor.permute(1, 2, 3, 0)
-        # If float, assume [-1,1] and normalize to [0,1]
-        if video_tensor.is_floating_point():
-            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
-        video_array = video_tensor.float().numpy()
-    else:
-        video_array = frames
-        if hasattr(video_array, "shape") and video_array.ndim == 5:
-            video_array = video_array[0]
-
-    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
-    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
-        video_array = list(video_array)
-
-    export_to_video(video_array, str(output_path), fps=args.fps)
-    print(f"Saved generated video to {output_path}")
-
-
-if __name__ == "__main__":
-    main()

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/diffusion/models/ltx2/__init__.py

@@ -1,13 +0,0 @@
-"""LTX-2 model components for text-to-video and image-to-video generation."""
-
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
-    LTX2Pipeline,
-    get_ltx2_post_process_func,
-    get_ltx2_pre_process_func,
-)
-
-__all__ = [
-    "LTX2Pipeline",
-    "get_ltx2_post_process_func",
-    "get_ltx2_pre_process_func",
-]

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -1,186 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import os
-from typing import Any
-
-import torch
-from diffusers import LTXVideoPipeline
-from torch import nn
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-
-
-def get_ltx2_post_process_func(
-    od_config: OmniDiffusionConfig,
-):
-    """Post-process function for LTX-2: convert video output to numpy array."""
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)
-
-    def post_process_func(
-        video: torch.Tensor,
-        output_type: str = "np",
-    ):
-        if output_type == "latent":
-            return video
-        return video_processor.postprocess_video(video, output_type=output_type)
-
-    return post_process_func
-
-
-def get_ltx2_pre_process_func(
-    od_config: OmniDiffusionConfig,
-):
-    """Pre-process function for LTX-2: validate and prepare image inputs for I2V mode."""
-    import PIL.Image
-
-    def pre_process_func(requests: list[OmniDiffusionRequest]) -> list[OmniDiffusionRequest]:
-        for req in requests:
-            # Load image if path is provided
-            if req.image_path is not None and req.pil_image is None:
-                req.pil_image = PIL.Image.open(req.image_path).convert("RGB")
-
-        return requests
-
-    return pre_process_func
-
-
-class LTX2Pipeline(nn.Module):
-    """Wrapper for Lightricks LTX-2 text-to-video pipeline.
-
-    This model uses the diffusers LTXVideoPipeline for efficient video generation.
-    Supports text-to-video and image-to-video generation.
-    """
-
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.od_config = od_config
-        self.device = get_local_device()
-
-        model = od_config.model
-        local_files_only = os.path.exists(model)
-
-        # Load the LTX-2 pipeline from diffusers
-        dtype = getattr(od_config, "dtype", torch.bfloat16)
-
-        self.pipe = LTXVideoPipeline.from_pretrained(
-            model,
-            torch_dtype=dtype,
-            local_files_only=local_files_only,
-        ).to(self.device)
-
-        # Store configuration
-        self.output_type = self.od_config.output_type
-
-        # No weights_sources needed since we use the diffusers pipeline directly
-        self.weights_sources = []
-
-    @torch.no_grad()
-    def forward(
-        self,
-        req: OmniDiffusionRequest,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int | None = None,
-        guidance_scale: float = 4.0,
-        frame_num: int | None = None,
-        output_type: str | None = "np",
-        generator: torch.Generator | None = None,
-        **kwargs: Any,
-    ) -> DiffusionOutput:
-        """Generate video from text prompt.
-
-        Args:
-            req: OmniDiffusionRequest containing generation parameters
-            prompt: Text prompt for video generation
-            negative_prompt: Negative text prompt
-            height: Video height (must be divisible by 32)
-            width: Video width (must be divisible by 32)
-            num_inference_steps: Number of diffusion steps
-            guidance_scale: Guidance scale for classifier-free guidance
-            frame_num: Number of frames to generate
-            output_type: Output format ("np", "pt", or "pil")
-            generator: Random generator for reproducibility
-            **kwargs: Additional arguments passed to the pipeline
-
-        Returns:
-            DiffusionOutput containing the generated video frames
-        """
-        # Extract parameters from request
-        prompt = req.prompt if req.prompt is not None else prompt
-        negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
-
-        if prompt is None:
-            raise ValueError("Prompt is required for LTX-2 video generation.")
-
-        # Set default values for video generation
-        height = req.height or height or 512
-        width = req.width or width or 768
-        num_frames = req.num_frames if req.num_frames else frame_num or 121
-        num_steps = req.num_inference_steps or num_inference_steps or 40
-
-        # Ensure dimensions are divisible by 32 (LTX-2 requirement)
-        height = (height // 32) * 32
-        width = (width // 32) * 32
-
-        # Ensure num_frames follows LTX-2 convention (divisible by 8 + 1)
-        # e.g., 25, 41, 81, 121, etc.
-        if (num_frames - 1) % 8 != 0:
-            num_frames = ((num_frames - 1) // 8) * 8 + 1
-        num_frames = max(num_frames, 25)  # Minimum 25 frames
-
-        # Prepare generator
-        if generator is None:
-            seed = req.seed if hasattr(req, "seed") and req.seed is not None else 42
-            generator = torch.Generator(device=self.device).manual_seed(seed)
-
-        # Prepare image input for I2V mode if available
-        image = None
-        if hasattr(req, "pil_image") and req.pil_image is not None:
-            image = req.pil_image
-
-        # Call the diffusers pipeline
-        pipeline_kwargs = {
-            "prompt": prompt,
-            "negative_prompt": negative_prompt,
-            "height": height,
-            "width": width,
-            "num_frames": num_frames,
-            "num_inference_steps": num_steps,
-            "guidance_scale": guidance_scale,
-            "generator": generator,
-            "output_type": output_type,
-        }
-
-        # Add image if provided (for I2V mode)
-        if image is not None:
-            pipeline_kwargs["image"] = image
-
-        # Add any extra kwargs
-        pipeline_kwargs.update(kwargs)
-
-        # Generate video
-        output = self.pipe(**pipeline_kwargs)
-
-        # Extract frames from the pipeline output
-        if hasattr(output, "frames"):
-            frames = output.frames
-        elif hasattr(output, "videos"):
-            frames = output.videos
-        else:
-            # Fallback: assume the output itself is the frames
-            frames = output
-
-        # Return as DiffusionOutput
-        return DiffusionOutput(output=frames)

vllm_omni/diffusion/registry.py

@@ -79,11 +79,6 @@ _DIFFUSION_MODELS = {
         "pipeline_flux2_klein",
         "Flux2KleinPipeline",
     ),
-    "LTX2Pipeline": (
-        "ltx2",
-        "pipeline_ltx2",
-        "LTX2Pipeline",
-    ),
 }
 
 
@@ -105,19 +100,10 @@ def initialize_model(
     if model_class is not None:
         model = model_class(od_config=od_config)
         # Configure VAE memory optimization settings from config
-        # For pipelines that wrap diffusers pipelines (like LTX2Pipeline),
-        # VAE is accessed via model.pipe.vae
-        vae = None
-        if hasattr(model, "vae"):
-            vae = model.vae
-        elif hasattr(model, "pipe") and hasattr(model.pipe, "vae"):
-            vae = model.pipe.vae
-
-        if vae is not None:
-            if hasattr(vae, "use_slicing"):
-                vae.use_slicing = od_config.vae_use_slicing
-            if hasattr(vae, "use_tiling"):
-                vae.use_tiling = od_config.vae_use_tiling
+        if hasattr(model.vae, "use_slicing"):
+            model.vae.use_slicing = od_config.vae_use_slicing
+        if hasattr(model.vae, "use_tiling"):
+            model.vae.use_tiling = od_config.vae_use_tiling
 
         return model
     else:
@@ -141,7 +127,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
     "Flux2KleinPipeline": "get_flux2_klein_post_process_func",
-    "LTX2Pipeline": "get_ltx2_post_process_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {
@@ -154,7 +139,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",
     "WanPipeline": "get_wan22_pre_process_func",
     "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func",
-    "LTX2Pipeline": "get_ltx2_pre_process_func",
 }
 
 

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]:
-        """Create default diffusion stage configuration."""
-        # TODO: here is different from the Omni class. We should merge the two in the future.
-        cache_backend = kwargs.get("cache_backend", "none")
-        cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None))
-
-        devices = "0"
-        if "parallel_config" in kwargs:
-            parallel_config = kwargs["parallel_config"]
-            num_devices = kwargs["parallel_config"].world_size
-            for i in range(1, num_devices):
-                devices += f",{i}"
-        else:
-            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 _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> list[Any]:
+        """Create default diffusion stage configuration using StageConfigFactory.
+        
+        This method uses the structured factory pattern to build a robust,
+        type-safe stage configuration for async mode.
+        
+        Args:
+            kwargs: Dictionary of configuration parameters
+            
+        Returns:
+            List containing a single OmegaConf DictConfig for the diffusion stage
+        """
+        return StageConfigFactory.create_default_diffusion_async(kwargs)
 
     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":
-            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
+        """Delegate to StageConfigFactory for cache config defaults."""
+        return StageConfigFactory._get_default_cache_config(cache_backend)
 
     def _normalize_cache_config(self, cache_backend: str | None, cache_config: Any | None) -> Any | 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 = self._get_default_cache_config(cache_backend)
-        return cache_config
+        """Delegate to StageConfigFactory for cache config normalization."""
+        return StageConfigFactory._normalize_cache_config(cache_backend, cache_config)
 
-    def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
-        """Create default diffusion stage configuration."""
-        # 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"])
-        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.
-        devices = "0"
-        if "parallel_config" in kwargs:
-            num_devices = kwargs["parallel_config"].world_size
-            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 _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> list[Any]:
+        """Create default diffusion stage configuration using StageConfigFactory.
+        
+        This method uses the structured factory pattern to build a robust,
+        type-safe stage configuration instead of manually constructing dictionaries.
+        
+        Args:
+            kwargs: Dictionary of configuration parameters
+            
+        Returns:
+            List containing a single OmegaConf DictConfig for the diffusion stage
+        """
+        return StageConfigFactory.create_default_diffusion(kwargs)
 
     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()]