Create OmniOpenAIServeImage Class To manage edit_images and generate_image

Author: zhcn000000

vllm_omni/entrypoints/openai/api_server.py

@@ -0,0 +1,31 @@
+from vllm.entrypoints.openai.engine.protocol import (
+    ModelCard,
+    ModelList,
+    ModelPermission,
+)
+from vllm.entrypoints.openai.models.protocol import BaseModelPath
+
+
+class DiffusionServingModels:
+    """Minimal OpenAIServingModels implementation for diffusion-only servers.
+
+    vLLM's /v1/models route expects `app.state.openai_serving_models` to expose
+    `show_available_models()`. In pure diffusion mode we don't initialize the
+    full OpenAIServingModels (it depends on LLM-specific processors), so we
+    provide a lightweight fallback.
+    """
+
+    def __init__(self, base_model_paths: list[BaseModelPath]) -> None:
+        self._base_model_paths = base_model_paths
+
+    async def show_available_models(self) -> ModelList:
+        return ModelList(
+            data=[
+                ModelCard(
+                    id=base_model.name,
+                    root=base_model.model_path,
+                    permission=[ModelPermission()],
+                )
+                for base_model in self._base_model_paths
+            ]
+        )

vllm_omni/entrypoints/openai/diffusion_models.py

@@ -10,6 +10,8 @@
 
 import base64
 import io
+import json
+from typing import Any
 
 import PIL.Image
 
@@ -63,3 +65,25 @@ def encode_image_base64(image: PIL.Image.Image) -> str:
     image.save(buffer, format="PNG")
     buffer.seek(0)
     return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def apply_stage_default_sampling_params(
+    default_params_json: str | None,
+    sampling_params: Any,
+    stage_key: str,
+) -> None:
+    """
+    Update a stage's sampling parameters with vLLM-Omni defaults.
+
+    Args:
+        default_params_json: JSON string of stage-keyed default parameters
+        sampling_params: The sampling parameters object to update
+        stage_key: The stage ID/key in the pipeline
+    """
+    if default_params_json is not None:
+        default_params_dict = json.loads(default_params_json)
+        if stage_key in default_params_dict:
+            stage_defaults = default_params_dict[stage_key]
+            for param_name, param_value in stage_defaults.items():
+                if hasattr(sampling_params, param_name):
+                    setattr(sampling_params, param_name, param_value)

vllm_omni/entrypoints/openai/image_api_utils.py

@@ -4,6 +4,8 @@
 from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.images import (
     ImageData,
+    ImageEditRequest,
+    ImageEditResponse,
     ImageGenerationRequest,
     ImageGenerationResponse,
     ResponseFormat,
@@ -19,6 +21,8 @@
     "ImageData",
     "ImageGenerationRequest",
     "ImageGenerationResponse",
+    "ImageEditRequest",
+    "ImageEditResponse",
     "ResponseFormat",
     "VideoData",
     "VideoGenerationRequest",

vllm_omni/entrypoints/openai/protocol/__init__.py

@@ -7,6 +7,7 @@
 for text-to-image generation, with vllm-omni specific extensions.
 """
 
+import json
 from enum import Enum
 from typing import Any
 
@@ -127,5 +128,84 @@ class ImageGenerationResponse(BaseModel):
 
     created: int = Field(..., description="Unix timestamp of when the generation completed")
     data: list[ImageData] = Field(..., description="Array of generated images")
-    output_format: str = Field(None, description="The output format of the image generation")
-    size: str = Field(None, description="The size of the image generated")
+
+
+class ImageEditResponse(BaseModel):
+    """
+    OpenAI DALL-E compatible image generation response.
+
+    Returns generated images with metadata.
+    """
+
+    created: int = Field(..., description="Unix timestamp of when the generation completed")
+    data: list[ImageData] = Field(..., description="Array of generated images")
+    output_format: str = Field(..., description="The output format of the image generation")
+    size: str = Field(..., description="The size of the image generated")
+
+
+class ImageEditRequest(BaseModel):
+    prompt: str = Field(..., description="Text description of the desired image edit")
+    model: str | None = Field(
+        default=None,
+        description="Model to use (optional, uses server's configured model if omitted)",
+    )
+    n: int = Field(default=1, ge=1, le=10, description="Number of images to generate")
+    size: str | None = Field(
+        default=None,
+        description="Image dimensions in WIDTHxHEIGHT format (e.g., '1024x1024', uses model defaults if omitted)",
+    )
+    response_format: ResponseFormat = Field(default=ResponseFormat.B64_JSON, description="Format of the returned image")
+    user: str | None = Field(default=None, description="User identifier for tracking")
+
+    # vllm-omni extensions for diffusion control
+    negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the image")
+    num_inference_steps: int | None = Field(
+        default=None,
+        ge=1,
+        le=200,
+        description="Number of diffusion sampling steps (uses model defaults if not specified)",
+    )
+    guidance_scale: float | None = Field(
+        default=None,
+        ge=0.0,
+        le=20.0,
+        description="Classifier-free guidance scale (uses model defaults if not specified)",
+    )
+    true_cfg_scale: float | None = Field(
+        default=None,
+        ge=0.0,
+        le=20.0,
+        description="True CFG scale (model-specific parameter, may be ignored if not supported)",
+    )
+    seed: int | None = Field(default=None, description="Random seed for reproducibility")
+    generator_device: str | None = Field(
+        default=None,
+        description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.",
+    )
+    lora: dict[str, Any] | None = Field(
+        default=None,
+        description=(
+            "Optional LoRA adapter for this request. Expected shape: "
+            "{name/path/scale/int_id}. Field names are flexible "
+            "(e.g. name|lora_name|adapter, path|lora_path|local_path, "
+            "scale|lora_scale, int_id|lora_int_id)."
+        ),
+    )
+
+    @field_validator("lora")
+    @classmethod
+    def validate_lora(cls, v):
+        """Validate LoRA field - must be a dict if provided."""
+        if isinstance(v, str):
+            try:
+                v_dict = json.loads(v)
+                if isinstance(v_dict, dict):
+                    return v_dict
+                else:
+                    raise ValueError("LoRA field must be a JSON object (dict)")
+            except json.JSONDecodeError:
+                raise ValueError("LoRA field must be a valid JSON string representing a dict")
+        elif isinstance(v, dict) or v is None:
+            return v
+        else:
+            raise ValueError("LoRA field must be either a dict or a JSON string representing a dict")