diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 2a1eaaeda..02e728b65 100644 --- a/.buildkite/pipeline.yml +++ b/.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 diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 183fa0519..45f4b33ad 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/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( diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py new file mode 100644 index 000000000..1a4582681 --- /dev/null +++ b/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" diff --git a/tests/entrypoints/test_omni_stage_diffusion_config.py b/tests/entrypoints/test_omni_stage_diffusion_config.py new file mode 100644 index 000000000..5fe04cbbd --- /dev/null +++ b/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"} diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py index 13780dff0..987ee1782 100644 --- a/vllm_omni/diffusion/worker/gpu_worker.py +++ b/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) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 8356adac0..bd0ebd393 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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", diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 14f3bc0f6..0a4976a37 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/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, }, ) diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index fc6646984..f4abb6c96 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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", } diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 3f45a97f2..814e26009 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 3e0042549..49958eb66 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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)") diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 72012d926..1f225e5b9 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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