[Bugfix] t2i online fix cache configuration failure and errors in /v1/images/generation endpoint. (#499)

Signed-off-by: David Chen <530634352@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/pipeline.yml

@@ -24,18 +24,6 @@ steps:
     agents:
       queue: "cpu_queue_premerge"
 
-  - label: "Image Generation API Test"
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/entrypoints/openai_api/
-    agents:
-      queue: "cpu_queue_premerge"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-
   - label: "Diffusion Model Test"
     timeout_in_minutes: 20
     depends_on: image-build

tests/entrypoints/openai_api/test_image_server.py

@@ -106,6 +106,20 @@ class MockGenerationResult:
         self.images = images
 
 
+class FakeAsyncOmni:
+    """Fake AsyncOmni that yields a single diffusion output."""
+
+    def __init__(self):
+        self.stage_list = ["llm", "diffusion"]
+        self.default_sampling_params_list = [{"temperature": 0.1}, {"top_p": 0.9}]
+        self.captured_sampling_params_list = None
+
+    async def generate(self, prompt, request_id, sampling_params_list):
+        self.captured_sampling_params_list = sampling_params_list
+        images = [Image.new("RGB", (64, 64), color="green")]
+        yield MockGenerationResult(images)
+
+
 @pytest.fixture
 def mock_async_diffusion():
     """Mock AsyncOmniDiffusion instance that returns fake images"""
@@ -139,6 +153,22 @@ def test_client(mock_async_diffusion):
     return TestClient(app)
 
 
+@pytest.fixture
+def async_omni_test_client():
+    """Create test client with mocked AsyncOmni engine."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+
+    app.state.engine_client = FakeAsyncOmni()
+    app.state.stage_configs = [{"stage_type": "llm"}, {"stage_type": "diffusion"}]
+
+    return TestClient(app)
+
+
 @pytest.mark.skip(reason="Async API server uses different health check mechanism")
 def test_health_endpoint(test_client):
     """Test health check endpoint - skipped for async server"""
@@ -171,6 +201,29 @@ def test_generate_single_image(test_client):
     assert img.size == (64, 64)  # Our mock returns 64x64 images
 
 
+def test_generate_images_async_omni_sampling_params(async_omni_test_client):
+    """Test AsyncOmni path uses per-stage sampling params."""
+    response = async_omni_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a cat",
+            "n": 2,
+            "size": "256x256",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200
+    engine = async_omni_test_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    assert captured[0] == {"temperature": 0.1}
+    assert captured[1]["num_outputs_per_prompt"] == 2
+    assert captured[1]["height"] == 256
+    assert captured[1]["width"] == 256
+    assert captured[1]["seed"] == 7
+
+
 def test_generate_multiple_images(test_client):
     """Test generating multiple images"""
     response = test_client.post(

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.entrypoints import omni as omni_module
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+
+def test_default_stage_config_includes_cache_backend(monkeypatch):
+    """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
+    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
+
+    omni = AsyncOmni(
+        model="dummy-model",
+        cache_backend="cache_dit",
+        cache_config='{"Fn_compute_blocks": 2}',
+        vae_use_slicing=True,
+        ulysses_degree=2,
+    )
+
+    stage_cfg = omni.stage_configs[0]
+    engine_args = stage_cfg.engine_args
+
+    assert engine_args.get("cache_backend") == "cache_dit"
+    cache_config = engine_args.get("cache_config")
+    assert cache_config["Fn_compute_blocks"] == 2
+    assert engine_args.get("vae_use_slicing") is True
+    parallel_config = engine_args.get("parallel_config")
+    if hasattr(parallel_config, "get"):
+        ulysses_degree = parallel_config.get("ulysses_degree")
+    else:
+        ulysses_degree = getattr(parallel_config, "ulysses_degree", None)
+    assert ulysses_degree == 2
+
+
+def test_default_cache_config_used_when_missing(monkeypatch):
+    """Ensure default cache_config is applied when cache_backend is set."""
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
+    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
+
+    omni = AsyncOmni(
+        model="dummy-model",
+        cache_backend="cache_dit",
+    )
+
+    engine_args = omni.stage_configs[0].engine_args
+    cache_config = engine_args.get("cache_config")
+    assert cache_config is not None
+    assert cache_config["Fn_compute_blocks"] == 1
+
+
+def test_default_stage_devices_from_sequence_parallel(monkeypatch):
+    """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
+    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
+
+    omni = AsyncOmni(
+        model="dummy-model",
+        ulysses_degree=2,
+        ring_degree=2,
+    )
+
+    stage_cfg = omni.stage_configs[0]
+    runtime = stage_cfg.runtime
+    if hasattr(runtime, "get"):
+        devices = runtime.get("devices")
+    else:
+        devices = getattr(runtime, "devices", None)
+    assert devices == "0,1,2,3"

tests/entrypoints/test_omni_stage_diffusion_config.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.entrypoints.omni_stage import _build_od_config
+
+
+def test_build_od_config_includes_diffusion_fields():
+    engine_args = {
+        "cache_backend": "cache_dit",
+        "cache_config": {"Fn_compute_blocks": 2},
+        "vae_use_slicing": True,
+    }
+    od_config = _build_od_config(engine_args, model="dummy-model")
+
+    assert od_config["model"] == "dummy-model"
+    assert od_config["cache_backend"] == "cache_dit"
+    assert od_config["cache_config"]["Fn_compute_blocks"] == 2
+    assert od_config["vae_use_slicing"] is True
+
+
+def test_build_od_config_respects_explicit_config():
+    engine_args = {
+        "od_config": {"cache_backend": "tea_cache"},
+        "cache_backend": "cache_dit",
+    }
+    od_config = _build_od_config(engine_args, model="dummy-model")
+    assert od_config == {"cache_backend": "tea_cache"}

vllm_omni/diffusion/worker/gpu_worker.py

@@ -44,6 +44,7 @@ class GPUWorker:
         self.rank = rank
         self.od_config = od_config
         self.pipeline = None
+        self.device = None
 
         self.init_device_and_model()
 
@@ -58,8 +59,8 @@ class GPUWorker:
         os.environ["RANK"] = str(rank)
         os.environ["WORLD_SIZE"] = str(world_size)
 
-        device = torch.device(f"cuda:{rank}")
-        torch.cuda.set_device(device)
+        self.device = torch.device(f"cuda:{rank}")
+        torch.cuda.set_device(self.device)
 
         # hack
         vllm_config = VllmConfig()
@@ -89,7 +90,7 @@ class GPUWorker:
             with DeviceMemoryProfiler() as m:
                 self.pipeline = model_loader.load_model(
                     od_config=self.od_config,
-                    load_device=f"cuda:{rank}",
+                    load_device=str(self.device),
                 )
             time_after_load = time.perf_counter()
 
@@ -129,6 +130,9 @@ class GPUWorker:
         # TODO: dealing with first req for now
         req = reqs[0]
 
+        if req.generator is None and req.seed is not None:
+            req.generator = torch.Generator(device=self.device).manual_seed(req.seed)
+
         # Refresh cache context if needed
         if self.cache_backend is not None and self.cache_backend.is_enabled():
             self.cache_backend.refresh(self.pipeline, req.num_inference_steps)

vllm_omni/entrypoints/async_omni.py

@@ -85,6 +85,9 @@ class AsyncOmni(OmniBase):
     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"]
@@ -92,13 +95,21 @@ class AsyncOmni(OmniBase):
             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")
+            if sequence_parallel_size is None:
+                sequence_parallel_size = ulysses_degree * ring_degree
+            num_devices = sequence_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=1,
-                sequence_parallel_size=1,
-                ulysses_degree=1,
-                ring_degree=1,
+                sequence_parallel_size=sequence_parallel_size,
+                ulysses_degree=ulysses_degree,
+                ring_degree=ring_degree,
                 cfg_parallel_size=1,
             )
         default_stage_cfg = [
@@ -114,8 +125,8 @@ class AsyncOmni(OmniBase):
                     "parallel_config": parallel_config,
                     "vae_use_slicing": kwargs.get("vae_use_slicing", False),
                     "vae_use_tiling": kwargs.get("vae_use_tiling", False),
-                    "cache_backend": kwargs.get("cache_backend", "none"),
-                    "cache_config": kwargs.get("cache_config", None),
+                    "cache_backend": cache_backend,
+                    "cache_config": cache_config,
                 },
                 "final_output": True,
                 "final_output_type": "image",

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -119,7 +119,7 @@ class AsyncOmniDiffusion:
         prompt: str,
         request_id: str | None = None,
         num_inference_steps: int = 50,
-        guidance_scale: float = 7.5,
+        guidance_scale: float | None = None,
         height: int | None = None,
         width: int | None = None,
         negative_prompt: str | None = None,
@@ -133,7 +133,7 @@ class AsyncOmniDiffusion:
             prompt: Text prompt describing the desired image
             request_id: Optional unique identifier for tracking the request
             num_inference_steps: Number of denoising steps (default: 50)
-            guidance_scale: Classifier-free guidance scale (default: 7.5)
+            guidance_scale: Classifier-free guidance scale (optional, uses model defaults if omitted)
             height: Optional image height in pixels
             width: Optional image width in pixels
             negative_prompt: Optional negative prompt for guidance
@@ -151,18 +151,21 @@ class AsyncOmniDiffusion:
             request_id = f"diff-{uuid.uuid4().hex[:16]}"
 
         # Prepare request
-        request = self._prepare_request(
-            prompt=prompt,
-            request_id=request_id,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            height=height,
-            width=width,
-            negative_prompt=negative_prompt,
-            num_outputs_per_prompt=num_outputs_per_prompt,
-            seed=seed,
+        request_kwargs = {
+            "prompt": prompt,
+            "request_id": request_id,
+            "num_inference_steps": num_inference_steps,
+            "height": height,
+            "width": width,
+            "negative_prompt": negative_prompt,
+            "num_outputs_per_prompt": num_outputs_per_prompt,
+            "seed": seed,
             **kwargs,
-        )
+        }
+        if guidance_scale is not None:
+            request_kwargs["guidance_scale"] = guidance_scale
+
+        request = self._prepare_request(**request_kwargs)
 
         logger.debug("Starting generation for request %s", request_id)
 
@@ -201,7 +204,7 @@ class AsyncOmniDiffusion:
             prompt=prompt,
             metrics={
                 "num_inference_steps": num_inference_steps,
-                "guidance_scale": guidance_scale,
+                "guidance_scale": request.guidance_scale,
             },
         )
 

vllm_omni/entrypoints/omni.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
 import multiprocessing as mp
 import os
 import time
@@ -101,6 +102,36 @@ class OmniBase:
         logger.info(f"Initializing stages for model: {model}")
         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
+
+    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
+
     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.
@@ -108,6 +139,8 @@ class OmniBase:
         # 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:
@@ -123,7 +156,13 @@ class OmniBase:
                     "devices": devices,
                     "max_batch_size": 1,
                 },
-                "engine_args": OmegaConf.create(kwargs),
+                "engine_args": OmegaConf.create(
+                    {
+                        **kwargs,
+                        "cache_backend": cache_backend,
+                        "cache_config": cache_config,
+                    }
+                ),
                 "final_output": True,
                 "final_output_type": "image",
             }

vllm_omni/entrypoints/omni_stage.py

@@ -14,6 +14,7 @@ import os
 import queue
 import sys
 import traceback
+from dataclasses import fields
 from typing import Any
 
 from vllm.inputs import TextPrompt
@@ -26,6 +27,7 @@ from vllm.v1.engine import EngineCoreOutput
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.llm_engine import LLMEngine
 
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.distributed.omni_connectors import build_stage_connectors
 from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
 from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor
@@ -46,6 +48,18 @@ from vllm_omni.utils import detect_device_type
 logger = init_logger(__name__)
 
 
+def _build_od_config(engine_args: dict[str, Any], model: str) -> dict[str, Any]:
+    """Build OmniDiffusionConfig kwargs from engine args."""
+    od_config = engine_args.get("od_config", {})
+    if not od_config:
+        od_config = {"model": model}
+        od_field_names = {f.name for f in fields(OmniDiffusionConfig)}
+        for key, value in engine_args.items():
+            if key in od_field_names:
+                od_config[key] = value
+    return od_config
+
+
 def prepare_sampling_params(sampling_params: Any, stage_type: str) -> Any:
     """Prepare sampling parameters for the given stage type.
 
@@ -1019,17 +1033,12 @@ async def _stage_worker_async(
     try:
         if stage_type == "diffusion":
             # For diffusion, we need to extract diffusion-specific config
-            od_config = engine_args.get("od_config", {})
-            if not od_config:
-                # Create default config from engine_args
-                od_config = {"model": model}
-                # Copy relevant diffusion args
-                for key in ["model", "device", "dtype", "enable_cpu_offload"]:
-                    if key in engine_args:
-                        od_config[key] = engine_args[key]
+            od_config = _build_od_config(engine_args, model)
             logger.debug(f"[Stage-%s] Initializing diffusion engine with config: {od_config}", stage_id)
             stage_engine = AsyncOmniDiffusion(
-                model=model, od_config=od_config, **{k: v for k, v in engine_args.items() if k != "od_config"}
+                model=model,
+                od_config=od_config,
+                **{k: v for k, v in engine_args.items() if k not in {"od_config", "model"}},
             )
             vllm_config = None  # Diffusion doesn't use vllm_config
         else:

vllm_omni/entrypoints/openai/api_server.py

@@ -573,8 +573,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             detail="Stage configs not found. Start server with a multi-stage omni model.",
         )
 
-    # Check for diffusion stage
+    # Check for diffusion stage and collect stage types
     has_diffusion_stage = False
+    stage_types: list[str] = []
     for stage in stage_configs:
         # Handle both dict and OmegaConf objects
         stage_type = None
@@ -593,7 +594,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
 
         if stage_type == "diffusion":
             has_diffusion_stage = True
-            break
+        stage_types.append(stage_type)
 
     if not has_diffusion_stage:
         raise HTTPException(
@@ -647,7 +648,33 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         logger.info(f"Generating {request.n} image(s) {size_str}")
 
         # Generate images using AsyncOmni (multi-stage mode)
-        result = await engine_client.generate(**gen_params)
+        result = None
+        stage_list = getattr(engine_client, "stage_list", None)
+        if isinstance(stage_list, list):
+            default_params_list = getattr(engine_client, "default_sampling_params_list", None)
+            if not isinstance(default_params_list, list):
+                default_params_list = [{} for _ in stage_types]
+            else:
+                default_params_list = list(default_params_list)
+            if len(default_params_list) != len(stage_types):
+                default_params_list = (default_params_list + [{} for _ in stage_types])[: len(stage_types)]
+
+            sampling_params_list: list[dict[str, Any]] = []
+            for idx, stage_type in enumerate(stage_types):
+                if stage_type == "diffusion":
+                    sampling_params_list.append(gen_params)
+                else:
+                    base_params = default_params_list[idx]
+                    sampling_params_list.append(dict(base_params) if isinstance(base_params, dict) else base_params)
+
+            async for output in engine_client.generate(
+                prompt=gen_params["prompt"],
+                request_id=gen_params["request_id"],
+                sampling_params_list=sampling_params_list,
+            ):
+                result = output
+        else:
+            result = await engine_client.generate(**gen_params)
 
         if result is None:
             raise HTTPException(
@@ -656,7 +683,15 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             )
 
         # Extract images from result
-        images = result.images if hasattr(result, "images") and result.images else []
+        images = []
+        if hasattr(result, "images") and result.images:
+            images = result.images
+        elif hasattr(result, "request_output"):
+            request_output = result.request_output
+            if isinstance(request_output, dict) and request_output.get("images"):
+                images = request_output["images"]
+            elif hasattr(request_output, "images") and request_output.images:
+                images = request_output.images
 
         logger.info(f"Successfully generated {len(images)} image(s)")
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -1062,7 +1062,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # Get request parameters from extra_body
             # Text-to-image parameters (ref: text_to_image.py)
             num_inference_steps = extra_body.get("num_inference_steps", 50)
-            guidance_scale = extra_body.get("guidance_scale", 7.5)
+            guidance_scale = extra_body.get("guidance_scale")
             true_cfg_scale = extra_body.get("true_cfg_scale")  # Qwen-Image specific
             seed = extra_body.get("seed")
             negative_prompt = extra_body.get("negative_prompt")
@@ -1094,7 +1094,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 "prompt": prompt,
                 "request_id": request_id,
                 "num_inference_steps": num_inference_steps,
-                "guidance_scale": guidance_scale,
                 "height": height,
                 "width": width,
                 "negative_prompt": negative_prompt,
@@ -1102,6 +1101,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 "seed": seed,
             }
 
+            if guidance_scale is not None:
+                gen_kwargs["guidance_scale"] = guidance_scale
+
             # Add Qwen-Image specific parameter
             if true_cfg_scale is not None:
                 gen_kwargs["true_cfg_scale"] = true_cfg_scale