From acf66ee2cc3799f22e01ec5eeddfe86f39a14388 Mon Sep 17 00:00:00 2001 From: Chenguang Zheng <645327136@qq.com> Date: Wed, 17 Dec 2025 17:40:45 +0800 Subject: [PATCH] [Entrypoints] Support Online Serving for Diffusion-only Models (#259) Signed-off-by: Chenguang ZHENG <645327136@qq.com> Signed-off-by: yinpe <11810305@mail.sustech.edu.cn> Co-authored-by: yinpe <11810305@mail.sustech.edu.cn> --- .../examples/online_serving/image_to_image.md | 197 +++++++++ .../examples/online_serving/text_to_image.md | 178 ++++++++ .../online_serving/image_to_image/README.md | 179 ++++++++ .../image_to_image/gradio_demo.py | 182 +++++++++ .../image_to_image/openai_chat_client.py | 155 +++++++ .../image_to_image/run_server.sh | 12 + .../online_serving/text_to_image/README.md | 156 +++++++ .../text_to_image/gradio_demo.py | 172 ++++++++ .../text_to_image/openai_chat_client.py | 133 ++++++ .../text_to_image/run_curl_text_to_image.sh | 34 ++ .../text_to_image/run_server.sh | 12 + vllm_omni/diffusion/data.py | 20 + vllm_omni/entrypoints/__init__.py | 21 + vllm_omni/entrypoints/async_diffusion.py | 269 ++++++++++++ vllm_omni/entrypoints/cli/serve.py | 75 +++- vllm_omni/entrypoints/openai/__init__.py | 33 ++ vllm_omni/entrypoints/openai/api_server.py | 172 +++++++- vllm_omni/entrypoints/openai/serving_chat.py | 386 +++++++++++++++++- vllm_omni/outputs.py | 140 ++++++- 19 files changed, 2506 insertions(+), 20 deletions(-) create mode 100644 docs/user_guide/examples/online_serving/image_to_image.md create mode 100644 docs/user_guide/examples/online_serving/text_to_image.md create mode 100644 examples/online_serving/image_to_image/README.md create mode 100644 examples/online_serving/image_to_image/gradio_demo.py create mode 100644 examples/online_serving/image_to_image/openai_chat_client.py create mode 100755 examples/online_serving/image_to_image/run_server.sh create mode 100644 examples/online_serving/text_to_image/README.md create mode 100644 examples/online_serving/text_to_image/gradio_demo.py create mode 100644 examples/online_serving/text_to_image/openai_chat_client.py create mode 100755 examples/online_serving/text_to_image/run_curl_text_to_image.sh create mode 100755 examples/online_serving/text_to_image/run_server.sh create mode 100644 vllm_omni/entrypoints/async_diffusion.py diff --git a/docs/user_guide/examples/online_serving/image_to_image.md b/docs/user_guide/examples/online_serving/image_to_image.md new file mode 100644 index 000000000..096488b92 --- /dev/null +++ b/docs/user_guide/examples/online_serving/image_to_image.md @@ -0,0 +1,197 @@ +# Qwen-Image-Edit Online Serving + +Source . + + +This example demonstrates how to deploy Qwen-Image-Edit model for online image editing service using vLLM-Omni. + +## Start Server + +### Basic Start + +```bash +vllm serve Qwen/Qwen-Image-Edit --omni --port 8092 +``` + +### Start with Parameters + +Or use the startup script: + +```bash +bash run_server.sh +``` + +## API Calls + +### Method 1: Using curl (Image Editing) + +```bash +# Image editing +bash run_curl_image_edit.sh input.png "Convert this image to watercolor style" + +# Or execute directly +IMG_B64=$(base64 -w0 input.png) +curl -s http://localhost:8092/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d "{ + \"messages\": [{ + \"role\": \"user\", + \"content\": [ + {\"type\": \"text\", \"text\": \"Convert this image to watercolor style\"}, + {\"type\": \"image_url\", \"image_url\": {\"url\": \"data:image/png;base64,\$IMG_B64\"}} + ] + }], + \"extra_body\": { + \"height\": 1024, + \"width\": 1024, + \"num_inference_steps\": 50, + \"guidance_scale\": 7.5, + \"seed\": 42 + } + }" | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png +``` + +### Method 2: Using Python Client + +```bash +python openai_chat_client.py --input input.png --prompt "Convert to oil painting style" --output output.png +``` + +### Method 3: Using Gradio Demo + +```bash +python gradio_demo.py +# Visit http://localhost:7861 +``` + +## Request Format + +### Image Editing (Using image_url Format) + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Convert this image to watercolor style"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}} + ] + } + ] +} +``` + +### Image Editing (Using Simplified image Format) + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"text": "Convert this image to watercolor style"}, + {"image": "BASE64_IMAGE_DATA"} + ] + } + ] +} +``` + +### Image Editing with Parameters + +Use `extra_body` to pass generation parameters: + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Convert to ink wash painting style"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}} + ] + } + ], + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "guidance_scale": 7.5, + "seed": 42 + } +} +``` + +## Generation Parameters (extra_body) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `num_inference_steps` | int | 50 | Number of denoising steps | +| `guidance_scale` | float | 7.5 | CFG guidance scale | +| `seed` | int | None | Random seed (reproducible) | +| `negative_prompt` | str | None | Negative prompt | +| `num_outputs_per_prompt` | int | 1 | Number of images to generate | + +## Response Format + +```json +{ + "id": "chatcmpl-xxx", + "created": 1234567890, + "model": "Qwen/Qwen-Image-Edit", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": [{ + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,..." + } + }] + }, + "finish_reason": "stop" + }], + "usage": {...} +} +``` + +## Common Editing Instructions Examples + +| Instruction | Description | +|-------------|-------------| +| `Convert this image to watercolor style` | Style transfer | +| `Convert the image to black and white` | Desaturation | +| `Enhance the color saturation` | Color adjustment | +| `Convert to cartoon style` | Cartoonization | +| `Add vintage filter effect` | Filter effect | +| `Convert daytime scene to nighttime` | Scene conversion | + +## File Description + +| File | Description | +|------|-------------| +| `run_server.sh` | Server startup script | +| `run_curl_image_edit.sh` | curl image editing example | +| `openai_chat_client.py` | Python client | +| `gradio_demo.py` | Gradio interactive interface | + +## Example materials + +??? abstract "gradio_demo.py" + ``````py + --8<-- "examples/online_serving/qwen_image_edit/gradio_demo.py" + `````` +??? abstract "openai_chat_client.py" + ``````py + --8<-- "examples/online_serving/qwen_image_edit/openai_chat_client.py" + `````` +??? abstract "run_curl_image_edit.sh" + ``````sh + --8<-- "examples/online_serving/qwen_image_edit/run_curl_image_edit.sh" + `````` +??? abstract "run_server.sh" + ``````sh + --8<-- "examples/online_serving/qwen_image_edit/run_server.sh" + `````` diff --git a/docs/user_guide/examples/online_serving/text_to_image.md b/docs/user_guide/examples/online_serving/text_to_image.md new file mode 100644 index 000000000..12d4de636 --- /dev/null +++ b/docs/user_guide/examples/online_serving/text_to_image.md @@ -0,0 +1,178 @@ +# Qwen-Image/Z-Image Online Serving + +Source . + + +This example demonstrates how to deploy Qwen-Image model for online image generation service using vLLM-Omni. + +## Start Server + +### Basic Start + +```bash +vllm serve Qwen/Qwen-Image --omni --port 8091 +``` + +### Start with Parameters + + +Or use the startup script: + +```bash +bash run_server.sh +``` + +## API Calls + +### Method 1: Using curl + +```bash +# Basic text-to-image generation +bash run_curl_text_to_image.sh + +# Or execute directly +curl -s http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + {"role": "user", "content": "A beautiful landscape painting"} + ], + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42 + } + }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png +``` + +### Method 2: Using Python Client + +```bash +python openai_chat_client.py --prompt "A beautiful landscape painting" --output output.png +``` + +### Method 3: Using Gradio Demo + +```bash +python gradio_demo.py +# Visit http://localhost:7860 +``` + +## Request Format + +### Simple Text-to-Image Generation + +```json +{ + "messages": [ + {"role": "user", "content": "A beautiful landscape painting"} + ] +} +``` + +### Generation with Parameters + +Use `extra_body` to pass generation parameters: + +```json +{ + "messages": [ + {"role": "user", "content": "A beautiful landscape painting"} + ], + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42 + } +} +``` + +### Multimodal Input (Text + Structured Content) + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "A beautiful landscape painting"} + ] + } + ] +} +``` + +## Generation Parameters (extra_body) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `height` | int | None | Image height in pixels | +| `width` | int | None | Image width in pixels | +| `num_inference_steps` | int | 50 | Number of denoising steps | +| `true_cfg_scale` | float | 4.0 | Qwen-Image CFG scale | +| `seed` | int | None | Random seed (reproducible) | +| `negative_prompt` | str | None | Negative prompt | +| `num_outputs_per_prompt` | int | 1 | Number of images to generate | + +## Response Format + +```json +{ + "id": "chatcmpl-xxx", + "created": 1234567890, + "model": "Qwen/Qwen-Image", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": [{ + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,..." + } + }] + }, + "finish_reason": "stop" + }], + "usage": {...} +} +``` + +## Extract Image + +```bash +# Extract base64 from response and decode to image +cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png +``` + +## File Description + +| File | Description | +|------|-------------| +| `run_server.sh` | Server startup script | +| `run_curl_text_to_image.sh` | curl example | +| `openai_chat_client.py` | Python client | +| `gradio_demo.py` | Gradio interactive interface | + +## Example materials + +??? abstract "gradio_demo.py" + ``````py + --8<-- "examples/online_serving/qwen_image/gradio_demo.py" + `````` +??? abstract "openai_chat_client.py" + ``````py + --8<-- "examples/online_serving/qwen_image/openai_chat_client.py" + `````` +??? abstract "run_curl_text_to_image.sh" + ``````sh + --8<-- "examples/online_serving/qwen_image/run_curl_text_to_image.sh" + `````` +??? abstract "run_server.sh" + ``````sh + --8<-- "examples/online_serving/qwen_image/run_server.sh" + `````` diff --git a/examples/online_serving/image_to_image/README.md b/examples/online_serving/image_to_image/README.md new file mode 100644 index 000000000..d9c3f11ff --- /dev/null +++ b/examples/online_serving/image_to_image/README.md @@ -0,0 +1,179 @@ +# Qwen-Image-Edit Online Serving + +This example demonstrates how to deploy Qwen-Image-Edit model for online image editing service using vLLM-Omni. + +## Start Server + +### Basic Start + +```bash +vllm serve Qwen/Qwen-Image-Edit --omni --port 8092 +``` + +### Start with Parameters + + +Or use the startup script: + +```bash +bash run_server.sh +``` + +## API Calls + +### Method 1: Using curl (Image Editing) + +```bash +# Image editing +bash run_curl_image_edit.sh input.png "Convert this image to watercolor style" + +# Or execute directly +IMG_B64=$(base64 -w0 input.png) +curl -s http://localhost:8092/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d "{ + \"messages\": [{ + \"role\": \"user\", + \"content\": [ + {\"type\": \"text\", \"text\": \"Convert this image to watercolor style\"}, + {\"type\": \"image_url\", \"image_url\": {\"url\": \"data:image/png;base64,\$IMG_B64\"}} + ] + }], + \"extra_body\": { + \"height\": 1024, + \"width\": 1024, + \"num_inference_steps\": 50, + \"guidance_scale\": 7.5, + \"seed\": 42 + } + }" | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png +``` + +### Method 2: Using Python Client + +```bash +python openai_chat_client.py --input input.png --prompt "Convert to oil painting style" --output output.png +``` + +### Method 3: Using Gradio Demo + +```bash +python gradio_demo.py +# Visit http://localhost:7861 +``` + +## Request Format + +### Image Editing (Using image_url Format) + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Convert this image to watercolor style"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}} + ] + } + ] +} +``` + +### Image Editing (Using Simplified image Format) + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"text": "Convert this image to watercolor style"}, + {"image": "BASE64_IMAGE_DATA"} + ] + } + ] +} +``` + +### Image Editing with Parameters + +Use `extra_body` to pass generation parameters: + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Convert to ink wash painting style"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}} + ] + } + ], + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "guidance_scale": 7.5, + "seed": 42 + } +} +``` + +## Generation Parameters (extra_body) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `height` | int | None | Output image height in pixels | +| `width` | int | None | Output image width in pixels | +| `size` | str | None | Output image size (e.g., "1024x1024") | +| `num_inference_steps` | int | 50 | Number of denoising steps | +| `guidance_scale` | float | 7.5 | CFG guidance scale | +| `seed` | int | None | Random seed (reproducible) | +| `negative_prompt` | str | None | Negative prompt | +| `num_outputs_per_prompt` | int | 1 | Number of images to generate | + +## Response Format + +```json +{ + "id": "chatcmpl-xxx", + "created": 1234567890, + "model": "Qwen/Qwen-Image-Edit", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": [{ + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,..." + } + }] + }, + "finish_reason": "stop" + }], + "usage": {...} +} +``` + +## Common Editing Instructions Examples + +| Instruction | Description | +|-------------|-------------| +| `Convert this image to watercolor style` | Style transfer | +| `Convert the image to black and white` | Desaturation | +| `Enhance the color saturation` | Color adjustment | +| `Convert to cartoon style` | Cartoonization | +| `Add vintage filter effect` | Filter effect | +| `Convert daytime scene to nighttime` | Scene conversion | + +## File Description + +| File | Description | +|------|-------------| +| `run_server.sh` | Server startup script | +| `run_curl_image_edit.sh` | curl image editing example | +| `openai_chat_client.py` | Python client | +| `gradio_demo.py` | Gradio interactive interface | diff --git a/examples/online_serving/image_to_image/gradio_demo.py b/examples/online_serving/image_to_image/gradio_demo.py new file mode 100644 index 000000000..db3198c79 --- /dev/null +++ b/examples/online_serving/image_to_image/gradio_demo.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +""" +Qwen-Image-Edit Gradio Demo for online serving. + +Usage: + python gradio_demo.py [--server http://localhost:8092] [--port 7861] +""" + +import argparse +import base64 +from io import BytesIO + +import gradio as gr +import requests +from PIL import Image + + +def edit_image( + input_image: Image.Image, + prompt: str, + steps: int, + guidance_scale: float, + seed: int | None, + negative_prompt: str, + server_url: str, +) -> Image.Image | None: + """Edit an image using the chat completions API.""" + if input_image is None: + raise gr.Error("Please upload an image first") + + # Convert image to base64 + buffer = BytesIO() + input_image.save(buffer, format="PNG") + image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8") + + # Build user message with text and image + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}, + ], + } + ] + + # Build extra_body with generation parameters + extra_body = { + "num_inference_steps": steps, + "guidance_scale": guidance_scale, + } + if seed is not None and seed >= 0: + extra_body["seed"] = seed + if negative_prompt: + extra_body["negative_prompt"] = negative_prompt + + # Build request payload + payload = {"messages": messages, "extra_body": extra_body} + + try: + response = requests.post( + f"{server_url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=300, + ) + response.raise_for_status() + data = response.json() + + content = data["choices"][0]["message"]["content"] + if isinstance(content, list) and len(content) > 0: + image_url = content[0].get("image_url", {}).get("url", "") + if image_url.startswith("data:image"): + _, b64_data = image_url.split(",", 1) + image_bytes = base64.b64decode(b64_data) + return Image.open(BytesIO(image_bytes)) + + return None + + except Exception as e: + print(f"Error: {e}") + raise gr.Error(f"Edit failed: {e}") + + +def create_demo(server_url: str): + """Create Gradio demo interface.""" + + with gr.Blocks(title="Qwen-Image-Edit Demo") as demo: + gr.Markdown("# Qwen-Image-Edit Online Editing") + gr.Markdown("Upload an image and describe the editing effect you want") + + with gr.Row(): + with gr.Column(scale=1): + input_image = gr.Image( + label="Input Image", + type="pil", + ) + prompt = gr.Textbox( + label="Edit Instruction", + placeholder="Describe the editing effect you want...", + lines=2, + ) + negative_prompt = gr.Textbox( + label="Negative Prompt", + placeholder="Describe what you don't want...", + lines=2, + ) + + with gr.Row(): + steps = gr.Slider( + label="Inference Steps", + minimum=10, + maximum=100, + value=50, + step=5, + ) + guidance_scale = gr.Slider( + label="Guidance Scale (CFG)", + minimum=1.0, + maximum=20.0, + value=7.5, + step=0.5, + ) + + with gr.Row(): + seed = gr.Number( + label="Random Seed (-1 for random)", + value=-1, + precision=0, + ) + + edit_btn = gr.Button("Edit Image", variant="primary") + + with gr.Column(scale=1): + output_image = gr.Image( + label="Edited Image", + type="pil", + ) + + # Examples + gr.Examples( + examples=[ + ["Convert this image to watercolor style"], + ["Convert the image to black and white"], + ["Enhance the color saturation"], + ["Convert to cartoon style"], + ["Add vintage filter effect"], + ["Convert daytime to nighttime"], + ["Convert to oil painting style"], + ["Add dreamy blur effect"], + ], + inputs=[prompt], + ) + + def process_edit(img, p, st, g, se, n): + actual_seed = se if se >= 0 else None + return edit_image(img, p, st, g, actual_seed, n, server_url) + + edit_btn.click( + fn=process_edit, + inputs=[input_image, prompt, steps, guidance_scale, seed, negative_prompt], + outputs=[output_image], + ) + + return demo + + +def main(): + parser = argparse.ArgumentParser(description="Qwen-Image-Edit Gradio Demo") + parser.add_argument("--server", default="http://localhost:8092", help="Server URL") + parser.add_argument("--port", type=int, default=7861, help="Gradio port") + parser.add_argument("--share", action="store_true", help="Create public link") + + args = parser.parse_args() + + print(f"Connecting to server: {args.server}") + demo = create_demo(args.server) + demo.launch(server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/image_to_image/openai_chat_client.py b/examples/online_serving/image_to_image/openai_chat_client.py new file mode 100644 index 000000000..01ec3cb1d --- /dev/null +++ b/examples/online_serving/image_to_image/openai_chat_client.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +""" +Qwen-Image-Edit OpenAI-compatible chat client for image editing. + +Usage: + python openai_chat_client.py --input qwen_image_output.png --prompt "Convert to watercolor style" --output output.png + python openai_chat_client.py --input input.png --prompt "Convert to oil painting" --seed 42 +""" + +import argparse +import base64 +from io import BytesIO +from pathlib import Path + +import requests +from PIL import Image + + +def edit_image( + input_image: str | Path, + prompt: str, + server_url: str = "http://localhost:8092", + height: int | None = None, + width: int | None = None, + steps: int | None = None, + guidance_scale: float | None = None, + seed: int | None = None, + negative_prompt: str | None = None, +) -> bytes | None: + """Edit an image using the chat completions API. + + Args: + input_image: Path to input image + prompt: Text description of the edit + server_url: Server URL + height: Output image height in pixels + width: Output image width in pixels + steps: Number of inference steps + guidance_scale: CFG guidance scale + seed: Random seed + negative_prompt: Negative prompt + + Returns: + Edited image bytes or None if failed + """ + # Read and encode input image + input_path = Path(input_image) + if not input_path.exists(): + print(f"Error: Input image not found: {input_path}") + return None + + with open(input_path, "rb") as f: + image_bytes = f.read() + image_b64 = base64.b64encode(image_bytes).decode("utf-8") + + # Detect image type + img = Image.open(BytesIO(image_bytes)) + mime_type = f"image/{img.format.lower()}" if img.format else "image/png" + + # Build user message with text and image + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{image_b64}"}}, + ], + } + ] + + # Build extra_body with generation parameters + extra_body = {} + if steps is not None: + extra_body["num_inference_steps"] = steps + if guidance_scale is not None: + extra_body["guidance_scale"] = guidance_scale + if seed is not None: + extra_body["seed"] = seed + if negative_prompt: + extra_body["negative_prompt"] = negative_prompt + + # Build request payload + payload = {"messages": messages} + if extra_body: + payload["extra_body"] = extra_body + + # Send request + try: + response = requests.post( + f"{server_url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=300, + ) + response.raise_for_status() + data = response.json() + + # Extract image from response + content = data["choices"][0]["message"]["content"] + if isinstance(content, list) and len(content) > 0: + image_url = content[0].get("image_url", {}).get("url", "") + if image_url.startswith("data:image"): + _, b64_data = image_url.split(",", 1) + return base64.b64decode(b64_data) + + print(f"Unexpected response format: {content}") + return None + + except Exception as e: + print(f"Error: {e}") + return None + + +def main(): + parser = argparse.ArgumentParser(description="Qwen-Image-Edit chat client") + parser.add_argument("--input", "-i", required=True, help="Input image path") + parser.add_argument("--prompt", "-p", required=True, help="Edit prompt") + parser.add_argument("--output", "-o", default="output.png", help="Output file") + parser.add_argument("--server", "-s", default="http://localhost:8092", help="Server URL") + parser.add_argument("--height", type=int, default=1024, help="Output image height") + parser.add_argument("--width", type=int, default=1024, help="Output image width") + parser.add_argument("--steps", type=int, default=50, help="Inference steps") + parser.add_argument("--guidance", type=float, default=7.5, help="Guidance scale") + parser.add_argument("--seed", type=int, help="Random seed") + parser.add_argument("--negative", help="Negative prompt") + + args = parser.parse_args() + + print(f"Input: {args.input}") + print(f"Prompt: {args.prompt}") + + image_bytes = edit_image( + input_image=args.input, + prompt=args.prompt, + server_url=args.server, + height=args.height, + width=args.width, + steps=args.steps, + guidance_scale=args.guidance, + seed=args.seed, + negative_prompt=args.negative, + ) + + if image_bytes: + output_path = Path(args.output) + output_path.write_bytes(image_bytes) + print(f"Image saved to: {output_path}") + print(f"Size: {len(image_bytes) / 1024:.1f} KB") + else: + print("Failed to edit image") + exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/image_to_image/run_server.sh b/examples/online_serving/image_to_image/run_server.sh new file mode 100755 index 000000000..6b8e081d9 --- /dev/null +++ b/examples/online_serving/image_to_image/run_server.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# Qwen-Image-Edit online serving startup script + +MODEL="${MODEL:-Qwen/Qwen-Image-Edit}" +PORT="${PORT:-8092}" + +echo "Starting Qwen-Image-Edit server..." +echo "Model: $MODEL" +echo "Port: $PORT" + +vllm serve "$MODEL" --omni \ + --port "$PORT" diff --git a/examples/online_serving/text_to_image/README.md b/examples/online_serving/text_to_image/README.md new file mode 100644 index 000000000..60494b443 --- /dev/null +++ b/examples/online_serving/text_to_image/README.md @@ -0,0 +1,156 @@ +# Qwen-Image Online Serving + +This example demonstrates how to deploy Qwen-Image model for online image generation service using vLLM-Omni. + +## Start Server + +### Basic Start + +```bash +vllm serve Qwen/Qwen-Image --omni --port 8091 +``` + +### Start with Parameters + +Or use the startup script: + +```bash +bash run_server.sh +``` + +## API Calls + +### Method 1: Using curl + +```bash +# Basic text-to-image generation +bash run_curl_text_to_image.sh + +# Or execute directly +curl -s http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + {"role": "user", "content": "A beautiful landscape painting"} + ], + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42 + } + }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png +``` + +### Method 2: Using Python Client + +```bash +python openai_chat_client.py --prompt "A beautiful landscape painting" --output output.png +``` + +### Method 3: Using Gradio Demo + +```bash +python gradio_demo.py +# Visit http://localhost:7860 +``` + +## Request Format + +### Simple Text Generation + +```json +{ + "messages": [ + {"role": "user", "content": "A beautiful landscape painting"} + ] +} +``` + +### Generation with Parameters + +Use `extra_body` to pass generation parameters: + +```json +{ + "messages": [ + {"role": "user", "content": "A beautiful landscape painting"} + ], + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "true_cfg_scale": 4.0, + "seed": 42 + } +} +``` + +### Multimodal Input (Text + Structured Content) + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "A beautiful landscape painting"} + ] + } + ] +} +``` + +## Generation Parameters (extra_body) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `height` | int | None | Image height in pixels | +| `width` | int | None | Image width in pixels | +| `size` | str | None | Image size (e.g., "1024x1024") | +| `num_inference_steps` | int | 50 | Number of denoising steps | +| `true_cfg_scale` | float | 4.0 | Qwen-Image CFG scale | +| `seed` | int | None | Random seed (reproducible) | +| `negative_prompt` | str | None | Negative prompt | +| `num_outputs_per_prompt` | int | 1 | Number of images to generate | + +## Response Format + +```json +{ + "id": "chatcmpl-xxx", + "created": 1234567890, + "model": "Qwen/Qwen-Image", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": [{ + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,..." + } + }] + }, + "finish_reason": "stop" + }], + "usage": {...} +} +``` + +## Extract Image + +```bash +# Extract base64 from response and decode to image +cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png +``` + +## File Description + +| File | Description | +|------|-------------| +| `run_server.sh` | Server startup script | +| `run_curl_text_to_image.sh` | curl example | +| `openai_chat_client.py` | Python client | +| `gradio_demo.py` | Gradio interactive interface | diff --git a/examples/online_serving/text_to_image/gradio_demo.py b/examples/online_serving/text_to_image/gradio_demo.py new file mode 100644 index 000000000..20d59fea7 --- /dev/null +++ b/examples/online_serving/text_to_image/gradio_demo.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +""" +Qwen-Image Gradio Demo for online serving. + +Usage: + python gradio_demo.py [--server http://localhost:8091] [--port 7860] +""" + +import argparse +import base64 +from io import BytesIO + +import gradio as gr +import requests +from PIL import Image + + +def generate_image( + prompt: str, + height: int, + width: int, + steps: int, + cfg_scale: float, + seed: int | None, + negative_prompt: str, + server_url: str, +) -> Image.Image | None: + """Generate an image using the chat completions API.""" + messages = [{"role": "user", "content": prompt}] + + # Build extra_body with generation parameters + extra_body = { + "height": height, + "width": width, + "num_inference_steps": steps, + "true_cfg_scale": cfg_scale, + } + if seed is not None and seed >= 0: + extra_body["seed"] = seed + if negative_prompt: + extra_body["negative_prompt"] = negative_prompt + + # Build request payload + payload = {"messages": messages, "extra_body": extra_body} + + try: + response = requests.post( + f"{server_url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=300, + ) + response.raise_for_status() + data = response.json() + + content = data["choices"][0]["message"]["content"] + if isinstance(content, list) and len(content) > 0: + image_url = content[0].get("image_url", {}).get("url", "") + if image_url.startswith("data:image"): + _, b64_data = image_url.split(",", 1) + image_bytes = base64.b64decode(b64_data) + return Image.open(BytesIO(image_bytes)) + + return None + + except Exception as e: + print(f"Error: {e}") + raise gr.Error(f"Generation failed: {e}") + + +def create_demo(server_url: str): + """Create Gradio demo interface.""" + + with gr.Blocks(title="Qwen-Image Demo") as demo: + gr.Markdown("# Qwen-Image Online Generation") + gr.Markdown("Generate images using Qwen-Image model") + + with gr.Row(): + with gr.Column(scale=1): + prompt = gr.Textbox( + label="Prompt", + placeholder="Describe the image you want to generate...", + lines=3, + ) + negative_prompt = gr.Textbox( + label="Negative Prompt", + placeholder="Describe what you don't want...", + lines=2, + ) + + with gr.Row(): + height = gr.Slider( + label="Height", + minimum=256, + maximum=2048, + value=1024, + step=64, + ) + width = gr.Slider( + label="Width", + minimum=256, + maximum=2048, + value=1024, + step=64, + ) + + with gr.Row(): + steps = gr.Slider( + label="Inference Steps", + minimum=10, + maximum=100, + value=50, + step=5, + ) + cfg_scale = gr.Slider( + label="True CFG Scale", + minimum=1.0, + maximum=20.0, + value=4.0, + step=0.5, + ) + + with gr.Row(): + seed = gr.Number( + label="Random Seed (-1 for random)", + value=-1, + precision=0, + ) + + generate_btn = gr.Button("Generate Image", variant="primary") + + with gr.Column(scale=1): + output_image = gr.Image( + label="Generated Image", + type="pil", + ) + + # Examples + gr.Examples( + examples=[ + ["A beautiful landscape painting with misty mountains", "", 1024, 1024, 50, 4.0, 42], + ["A cute cat sitting on a windowsill with sunlight", "", 1024, 1024, 50, 4.0, 123], + ["Cyberpunk style futuristic city with neon lights", "blurry, low quality", 1024, 768, 50, 4.0, 456], + ["Chinese ink painting of bamboo forest with a house", "", 768, 1024, 50, 4.0, 789], + ], + inputs=[prompt, negative_prompt, height, width, steps, cfg_scale, seed], + ) + + generate_btn.click( + fn=lambda p, h, w, st, c, se, n: generate_image(p, h, w, st, c, se if se >= 0 else None, n, server_url), + inputs=[prompt, height, width, steps, cfg_scale, seed, negative_prompt], + outputs=[output_image], + ) + + return demo + + +def main(): + parser = argparse.ArgumentParser(description="Qwen-Image Gradio Demo") + parser.add_argument("--server", default="http://localhost:8091", help="Server URL") + parser.add_argument("--port", type=int, default=7860, help="Gradio port") + parser.add_argument("--share", action="store_true", help="Create public link") + + args = parser.parse_args() + + print(f"Connecting to server: {args.server}") + demo = create_demo(args.server) + demo.launch(server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_image/openai_chat_client.py b/examples/online_serving/text_to_image/openai_chat_client.py new file mode 100644 index 000000000..cc5e00139 --- /dev/null +++ b/examples/online_serving/text_to_image/openai_chat_client.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +Qwen-Image OpenAI-compatible chat client for image generation. + +Usage: + python openai_chat_client.py --prompt "A beautiful landscape" --output output.png + python openai_chat_client.py --prompt "A sunset" --height 1024 --width 1024 --steps 50 --seed 42 +""" + +import argparse +import base64 +from pathlib import Path + +import requests + + +def generate_image( + prompt: str, + server_url: str = "http://localhost:8091", + height: int | None = None, + width: int | None = None, + steps: int | None = None, + true_cfg_scale: float | None = None, + seed: int | None = None, + negative_prompt: str | None = None, + num_outputs_per_prompt: int = 1, +) -> bytes | None: + """Generate an image using the chat completions API. + + Args: + prompt: Text description of the image + server_url: Server URL + height: Image height in pixels + width: Image width in pixels + steps: Number of inference steps + true_cfg_scale: Qwen-Image CFG scale + seed: Random seed + negative_prompt: Negative prompt + num_outputs_per_prompt: Number of images to generate + + Returns: + Image bytes or None if failed + """ + messages = [{"role": "user", "content": prompt}] + + # Build extra_body with generation parameters + extra_body = {} + if height is not None: + extra_body["height"] = height + if width is not None: + extra_body["width"] = width + if steps is not None: + extra_body["num_inference_steps"] = steps + if true_cfg_scale is not None: + extra_body["true_cfg_scale"] = true_cfg_scale + if seed is not None: + extra_body["seed"] = seed + if negative_prompt: + extra_body["negative_prompt"] = negative_prompt + if num_outputs_per_prompt > 1: + extra_body["num_outputs_per_prompt"] = num_outputs_per_prompt + + # Build request payload + payload = {"messages": messages} + if extra_body: + payload["extra_body"] = extra_body + + # Send request + try: + response = requests.post( + f"{server_url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=300, + ) + response.raise_for_status() + data = response.json() + + # Extract image from response + content = data["choices"][0]["message"]["content"] + if isinstance(content, list) and len(content) > 0: + image_url = content[0].get("image_url", {}).get("url", "") + if image_url.startswith("data:image"): + _, b64_data = image_url.split(",", 1) + return base64.b64decode(b64_data) + + print(f"Unexpected response format: {content}") + return None + + except Exception as e: + print(f"Error: {e}") + return None + + +def main(): + parser = argparse.ArgumentParser(description="Qwen-Image chat client") + parser.add_argument("--prompt", "-p", default="a cup of coffee on the table", help="Text prompt") + parser.add_argument("--output", "-o", default="qwen_image_output.png", help="Output file") + parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL") + parser.add_argument("--height", type=int, default=1024, help="Image height") + parser.add_argument("--width", type=int, default=1024, help="Image width") + parser.add_argument("--steps", type=int, default=50, help="Inference steps") + parser.add_argument("--cfg-scale", type=float, default=4.0, help="True CFG scale") + parser.add_argument("--seed", type=int, default=42, help="Random seed") + parser.add_argument("--negative", help="Negative prompt") + + args = parser.parse_args() + + print(f"Generating image for: {args.prompt}") + + image_bytes = generate_image( + prompt=args.prompt, + server_url=args.server, + height=args.height, + width=args.width, + steps=args.steps, + true_cfg_scale=args.cfg_scale, + seed=args.seed, + negative_prompt=args.negative, + ) + + if image_bytes: + output_path = Path(args.output) + output_path.write_bytes(image_bytes) + print(f"Image saved to: {output_path}") + print(f"Size: {len(image_bytes) / 1024:.1f} KB") + else: + print("Failed to generate image") + exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_image/run_curl_text_to_image.sh b/examples/online_serving/text_to_image/run_curl_text_to_image.sh new file mode 100755 index 000000000..d90d11c5a --- /dev/null +++ b/examples/online_serving/text_to_image/run_curl_text_to_image.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# Qwen-Image text-to-image curl example + +SERVER="${SERVER:-http://localhost:8091}" +PROMPT="${PROMPT:-a cup of coffee on the table}" +OUTPUT="${OUTPUT:-qwen_image_output.png}" + +echo "Generating image..." +echo "Prompt: $PROMPT" +echo "Output: $OUTPUT" + +curl -s "$SERVER/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "{ + \"messages\": [ + {\"role\": \"user\", \"content\": \"$PROMPT\"} + ], + \"extra_body\": { + \"height\": 1024, + \"width\": 1024, + \"num_inference_steps\": 50, + \"true_cfg_scale\": 4.0, + \"seed\": 42, + \"num_outputs_per_prompt\": 1 + } + }" | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > "$OUTPUT" + +if [ -f "$OUTPUT" ]; then + echo "Image saved to: $OUTPUT" + echo "Size: $(du -h "$OUTPUT" | cut -f1)" +else + echo "Failed to generate image" + exit 1 +fi diff --git a/examples/online_serving/text_to_image/run_server.sh b/examples/online_serving/text_to_image/run_server.sh new file mode 100755 index 000000000..b25337a0e --- /dev/null +++ b/examples/online_serving/text_to_image/run_server.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# Qwen-Image online serving startup script + +MODEL="${MODEL:-Qwen/Qwen-Image}" +PORT="${PORT:-8091}" + +echo "Starting Qwen-Image server..." +echo "Model: $MODEL" +echo "Port: $PORT" + +vllm serve "$MODEL" --omni \ + --port "$PORT" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index cd886f3ae..34692f5c4 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -388,6 +388,26 @@ class OmniDiffusionConfig: f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})" ) + # Convert string dtype to torch.dtype if needed + if isinstance(self.dtype, str): + dtype_map = { + "auto": torch.bfloat16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float16": torch.float16, + "fp16": torch.float16, + "half": torch.float16, + "float32": torch.float32, + "fp32": torch.float32, + "float": torch.float32, + } + dtype_lower = self.dtype.lower() + if dtype_lower in dtype_map: + self.dtype = dtype_map[dtype_lower] + else: + logger.warning(f"Unknown dtype string '{self.dtype}', defaulting to bfloat16") + self.dtype = torch.bfloat16 + # Convert cache_config dict to DiffusionCacheConfig if needed if isinstance(self.cache_config, dict): self.cache_config = DiffusionCacheConfig.from_dict(self.cache_config) diff --git a/vllm_omni/entrypoints/__init__.py b/vllm_omni/entrypoints/__init__.py index e69de29bb..1670d8681 100644 --- a/vllm_omni/entrypoints/__init__.py +++ b/vllm_omni/entrypoints/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +""" +vLLM-Omni entrypoints module. + +Provides high-level interfaces for running omni models including: +- AsyncOmni: Async orchestrator for multi-stage LLM pipelines +- AsyncOmniDiffusion: Async interface for diffusion model inference +- Omni: Unified entrypoint that auto-selects between LLM and Diffusion +""" + +from vllm_omni.entrypoints.async_diffusion import AsyncOmniDiffusion +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.omni import Omni + +__all__ = [ + "AsyncOmni", + "AsyncOmniDiffusion", + "Omni", +] diff --git a/vllm_omni/entrypoints/async_diffusion.py b/vllm_omni/entrypoints/async_diffusion.py new file mode 100644 index 000000000..b83f624ad --- /dev/null +++ b/vllm_omni/entrypoints/async_diffusion.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Async entrypoint for vLLM-Omni diffusion model inference. + +Provides an asynchronous interface for running diffusion models, +enabling concurrent request handling and streaming generation. +""" + +import asyncio +import uuid +from collections.abc import AsyncGenerator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import fields +from typing import Any + +from PIL import Image +from vllm.logger import init_logger +from vllm.transformers_utils.config import get_hf_file_to_dict + +from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) + + +class AsyncOmniDiffusion: + """Async entry point for vLLM-Omni diffusion model inference. + + This class provides an asynchronous interface for running diffusion models, + enabling concurrent request handling. It wraps the DiffusionEngine and + provides async methods for image generation. + + Args: + model: Model name or path to load + od_config: Optional OmniDiffusionConfig. If not provided, it will be + created from kwargs + **kwargs: Additional keyword arguments passed to OmniDiffusionConfig + + Example: + >>> async_diffusion = AsyncOmniDiffusion(model="Qwen/Qwen-Image") + >>> result = await async_diffusion.generate( + ... prompt="A beautiful sunset over the ocean", + ... request_id="req-1", + ... ) + >>> print(result.images) + """ + + def __init__( + self, + model: str, + od_config: OmniDiffusionConfig | None = None, + **kwargs: Any, + ): + self.model = model + + # Build config + if od_config is None: + od_config = OmniDiffusionConfig.from_kwargs(model=model, **kwargs) + elif isinstance(od_config, dict): + od_config = OmniDiffusionConfig.from_kwargs(**od_config) + + self.od_config = od_config + + # Load model class name and transformer config + config_dict = get_hf_file_to_dict("model_index.json", od_config.model) + od_config.model_class_name = config_dict.get("_class_name", None) + + tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model) + od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict) + + # Initialize engine + self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config) + + # Thread pool for running sync engine in async context + self._executor = ThreadPoolExecutor(max_workers=1) + self._closed = False + + logger.info("AsyncOmniDiffusion initialized with model: %s", model) + + def _prepare_request( + self, + prompt: str, + request_id: str | None = None, + **kwargs: Any, + ) -> OmniDiffusionRequest: + """Prepare a diffusion request from prompt and parameters. + + Args: + prompt: Text prompt for image generation + request_id: Optional unique identifier for the request + **kwargs: Additional generation parameters + + Returns: + OmniDiffusionRequest ready for processing + """ + if request_id is None: + request_id = f"diff-{uuid.uuid4().hex[:16]}" + + field_names = {f.name for f in fields(OmniDiffusionRequest)} + + init_kwargs = { + "prompt": prompt, + "request_id": request_id, + } + + for key, value in kwargs.items(): + if key in field_names: + init_kwargs[key] = value + + return OmniDiffusionRequest(**init_kwargs) + + async def generate( + self, + prompt: str, + request_id: str | None = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + height: int | None = None, + width: int | None = None, + negative_prompt: str | None = None, + num_outputs_per_prompt: int = 1, + seed: int | None = None, + **kwargs: Any, + ) -> OmniRequestOutput: + """Generate images asynchronously from a text prompt. + + Args: + 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) + height: Optional image height in pixels + width: Optional image width in pixels + negative_prompt: Optional negative prompt for guidance + num_outputs_per_prompt: Number of images to generate (default: 1) + seed: Optional random seed for reproducibility + **kwargs: Additional generation parameters + + Returns: + OmniRequestOutput containing generated images + + Raises: + RuntimeError: If generation fails + """ + if request_id is None: + 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, + **kwargs, + ) + + logger.debug("Starting generation for request %s", request_id) + + # Run engine in thread pool + loop = asyncio.get_event_loop() + try: + result = await loop.run_in_executor( + self._executor, + self.engine.step, + [request], + ) + except Exception as e: + logger.error("Generation failed for request %s: %s", request_id, e) + raise RuntimeError(f"Diffusion generation failed: {e}") from e + + # Process results + images: list[Image.Image] = [] + if result is not None: + if isinstance(result, list): + for item in result: + if isinstance(item, Image.Image): + images.append(item) + elif isinstance(result, Image.Image): + images.append(result) + + logger.debug( + "Generation completed for request %s, produced %d images", + request_id, + len(images), + ) + + return OmniRequestOutput.from_diffusion( + request_id=request_id, + images=images, + prompt=prompt, + metrics={ + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + }, + ) + + async def generate_stream( + self, + prompt: str, + request_id: str | None = None, + **kwargs: Any, + ) -> AsyncGenerator[OmniRequestOutput, None]: + """Generate images with streaming progress updates. + + Currently, diffusion models don't support true streaming, so this + yields a single result after generation completes. Future implementations + may support step-by-step progress updates. + + Args: + prompt: Text prompt describing the desired image + request_id: Optional unique identifier for tracking the request + **kwargs: Additional generation parameters + + Yields: + OmniRequestOutput with generation progress/results + """ + result = await self.generate(prompt=prompt, request_id=request_id, **kwargs) + yield result + + def close(self) -> None: + """Close the engine and release resources. + + Should be called when done using the AsyncOmniDiffusion instance. + """ + if self._closed: + return + self._closed = True + + try: + self.engine.close() + except Exception as e: + logger.warning("Error closing diffusion engine: %s", e) + + try: + self._executor.shutdown(wait=False) + except Exception as e: + logger.warning("Error shutting down executor: %s", e) + + logger.info("AsyncOmniDiffusion closed") + + def shutdown(self) -> None: + """Alias for close() method.""" + self.close() + + def __del__(self) -> None: + """Best-effort cleanup on deletion.""" + try: + self.close() + except Exception: + pass + + @property + def is_running(self) -> bool: + """Check if the engine is running.""" + return not self._closed + + @property + def is_stopped(self) -> bool: + """Check if the engine is stopped.""" + return self._closed diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 7a07e94f0..4680083fe 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -1,5 +1,8 @@ """ Omni serve command for vLLM-Omni. + +Supports both multi-stage LLM models (e.g., Qwen2.5-Omni) and +diffusion models (e.g., Qwen-Image) through the same CLI interface. """ import argparse @@ -15,8 +18,19 @@ from vllm_omni.entrypoints.openai.api_server import omni_run_server logger = init_logger(__name__) -DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni LLM -completions via HTTP. Defaults to Qwen/Qwen2.5-Omni-7B if no model is specified. +DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni models +via HTTP. Supports both multi-stage LLM models and diffusion models. + +The server automatically detects the model type: +- LLM models: Served via /v1/chat/completions endpoint +- Diffusion models: Served via /v1/images/generations endpoint + +Examples: + # Start an Omni LLM server + vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 + + # Start a diffusion model server + vllm serve Qwen/Qwen-Image --omni --port 8091 Search by using: `--help=` to explore options by section (e.g., --help=ModelConfig, --help=Frontend) @@ -38,6 +52,13 @@ class OmniServeCommand(CLISubcommand): uvloop.run(omni_run_server(args)) def validate(self, args: argparse.Namespace) -> None: + # Skip validation for diffusion models as they have different requirements + from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model + + model = getattr(args, "model_tag", None) or getattr(args, "model", None) + if model and is_diffusion_model(model): + logger.info("Detected diffusion model: %s", model) + return validate_parsed_serve_args(args) def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser: @@ -52,7 +73,7 @@ class OmniServeCommand(CLISubcommand): serve_parser.add_argument( "--omni", action="store_true", - help="Enable vLLM-Omni mode", + help="Enable vLLM-Omni mode for multi-modal and diffusion models", ) serve_parser.add_argument( "--stage-configs-path", @@ -108,6 +129,54 @@ class OmniServeCommand(CLISubcommand): default=None, help="The address of the Ray cluster to connect to.", ) + + # Diffusion model specific arguments + serve_parser.add_argument( + "--num-gpus", + type=int, + default=1, + help="Number of GPUs to use for diffusion model inference.", + ) + + # Cache optimization parameters + serve_parser.add_argument( + "--cache-backend", + type=str, + default="none", + help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit'", + ) + serve_parser.add_argument( + "--cache-config", + type=str, + default=None, + help="JSON string of cache configuration (e.g., '{\"rel_l1_thresh\": 0.2}').", + ) + + # VAE memory optimization parameters + serve_parser.add_argument( + "--vae-use-slicing", + action="store_true", + help="Enable VAE slicing for memory optimization (useful for mitigating OOM issues).", + ) + serve_parser.add_argument( + "--vae-use-tiling", + action="store_true", + help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).", + ) + + # Video model parameters (e.g., Wan2.2) - engine-level + serve_parser.add_argument( + "--boundary-ratio", + type=float, + default=None, + help="Boundary split ratio for low/high DiT in video models (e.g., 0.875 for Wan2.2).", + ) + serve_parser.add_argument( + "--flow-shift", + type=float, + default=None, + help="Scheduler flow_shift for video models (e.g., 5.0 for 720p, 12.0 for 480p).", + ) return serve_parser diff --git a/vllm_omni/entrypoints/openai/__init__.py b/vllm_omni/entrypoints/openai/__init__.py index e69de29bb..58a96cc0e 100644 --- a/vllm_omni/entrypoints/openai/__init__.py +++ b/vllm_omni/entrypoints/openai/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +OpenAI-compatible API entrypoints for vLLM-Omni. + +Provides: +- omni_run_server: Main server entry point (auto-detects model type) +- omni_run_diffusion_server: Server for diffusion models +- OmniOpenAIServingChat: Unified chat completion handler for both LLM and diffusion models +""" + +from vllm_omni.entrypoints.openai.api_server import ( + build_async_diffusion, + build_async_omni, + omni_diffusion_init_app_state, + omni_init_app_state, + omni_run_diffusion_server, + omni_run_server, +) +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + +__all__ = [ + # Server functions + "omni_run_server", + "omni_run_diffusion_server", + "build_async_omni", + "build_async_diffusion", + "omni_init_app_state", + "omni_diffusion_init_app_state", + # Serving classes + "OmniOpenAIServingChat", +] diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index b40a04e28..1f4798af4 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1,9 +1,11 @@ +import json import multiprocessing import multiprocessing.forkserver as forkserver import os from argparse import Namespace from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from dataclasses import fields from http import HTTPStatus from typing import Any @@ -38,6 +40,9 @@ from vllm.logger import init_logger from vllm.transformers_utils.tokenizer import MistralTokenizer from vllm.utils import decorate_logs +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model +from vllm_omni.entrypoints.async_diffusion import AsyncOmniDiffusion from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat @@ -45,13 +50,73 @@ logger = init_logger(__name__) async def omni_run_server(args, **uvicorn_kwargs) -> None: - """Run a single-worker API server.""" + """Run a single-worker API server. + + Automatically detects if the model is a diffusion model and routes + to the appropriate server implementation. + """ # Add process-specific prefix to stdout and stderr. decorate_logs("APIServer") listen_address, sock = setup_server(args) - await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) + + # Check if model is a diffusion model + if is_diffusion_model(args.model): + logger.info("Detected diffusion model, starting diffusion API server") + await omni_run_diffusion_server_worker(listen_address, sock, args, **uvicorn_kwargs) + else: + await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) + + +async def omni_run_diffusion_server(args, **uvicorn_kwargs) -> None: + """Run a diffusion model API server.""" + + # Add process-specific prefix to stdout and stderr. + decorate_logs("DiffusionAPIServer") + + listen_address, sock = setup_server(args) + await omni_run_diffusion_server_worker(listen_address, sock, args, **uvicorn_kwargs) + + +async def omni_run_diffusion_server_worker(listen_address, sock, args, **uvicorn_kwargs) -> None: + """Run a diffusion model API server worker.""" + + # Load logging config for uvicorn if specified + log_config = load_log_config(args.log_config_file) + if log_config is not None: + uvicorn_kwargs["log_config"] = log_config + + async with build_async_diffusion(args) as diffusion_engine: + app = build_app(args) + + await omni_diffusion_init_app_state(diffusion_engine, app.state, args) + + logger.info("Starting vLLM Diffusion API server on %s", listen_address) + + shutdown_task = await serve_http( + app, + sock=sock, + enable_ssl_refresh=getattr(args, "enable_ssl_refresh", False), + host=args.host, + port=args.port, + log_level=args.uvicorn_log_level, + access_log=not getattr(args, "disable_uvicorn_access_log", False), + timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE, + ssl_keyfile=getattr(args, "ssl_keyfile", None), + ssl_certfile=getattr(args, "ssl_certfile", None), + ssl_ca_certs=getattr(args, "ssl_ca_certs", None), + ssl_cert_reqs=getattr(args, "ssl_cert_reqs", 0), + h11_max_incomplete_event_size=getattr(args, "h11_max_incomplete_event_size", None), + h11_max_header_count=getattr(args, "h11_max_header_count", None), + **uvicorn_kwargs, + ) + + # NB: Await server shutdown only after the backend context is exited + try: + await shutdown_task + finally: + sock.close() async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None: @@ -147,6 +212,60 @@ async def build_async_omni( yield async_omni +@asynccontextmanager +async def build_async_diffusion( + args: Namespace, + **kwargs: Any, +) -> AsyncIterator[AsyncOmniDiffusion]: + """Build an AsyncOmniDiffusion instance from command-line arguments. + + Creates an async context manager that yields an AsyncOmniDiffusion + instance configured from the provided arguments. + + Args: + args: Parsed command-line arguments containing model and configuration + **kwargs: Additional keyword arguments passed to AsyncOmniDiffusion + + Yields: + AsyncOmniDiffusion instance ready for use + """ + diffusion_engine: AsyncOmniDiffusion | None = None + + try: + # Build diffusion kwargs by extracting matching OmniDiffusionConfig fields from args + config_field_names = {f.name for f in fields(OmniDiffusionConfig)} + diffusion_kwargs: dict[str, Any] = {"model": args.model} + + for field_name in config_field_names: + if not hasattr(args, field_name): + continue + value = getattr(args, field_name) + if value is None: + continue + # Special handling for cache_config JSON string + if field_name == "cache_config" and isinstance(value, str): + try: + value = json.loads(value) + except json.JSONDecodeError as e: + logger.warning(f"Failed to parse cache_config JSON: {e}") + continue + diffusion_kwargs[field_name] = value + + diffusion_kwargs.update(kwargs) + logger.info(f"diffusion_kwargs: {diffusion_kwargs}") + logger.info( + "Building AsyncOmniDiffusion with model=%s, num_gpus=%s", + args.model, + diffusion_kwargs.get("num_gpus", 1), + ) + diffusion_engine = AsyncOmniDiffusion(**diffusion_kwargs) + + yield diffusion_engine + finally: + if diffusion_engine: + diffusion_engine.shutdown() + + @asynccontextmanager async def build_async_omni_from_stage_config( args: Namespace, @@ -318,6 +437,49 @@ def Omnichat(request: Request) -> OmniOpenAIServingChat | None: return request.app.state.openai_serving_chat +async def omni_diffusion_init_app_state( + diffusion_engine: AsyncOmniDiffusion, + state: State, + args: Namespace, +) -> None: + """Initialize the FastAPI application state for diffusion model API server. + + Sets up the application state with diffusion model information and + chat completion handler for image generation via /v1/chat/completions. + + Args: + diffusion_engine: AsyncOmniDiffusion engine instance + state: FastAPI application state object to initialize + args: Parsed command-line arguments + """ + if args.served_model_name is not None: + served_model_names = args.served_model_name + else: + served_model_names = [args.model] + + model_name = served_model_names[0] if served_model_names else args.model + + state.diffusion_engine = diffusion_engine + state.log_stats = not getattr(args, "disable_log_stats", False) + + # Initialize chat handler with diffusion engine (uses /v1/chat/completions endpoint) + # Note: Request-level parameters (num_inference_steps, guidance_scale, seed, height, width, etc.) + # are passed per-request via the API, not as server defaults + state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion( + diffusion_engine=diffusion_engine, + model_name=model_name, + ) + + # Set other handlers to None for diffusion-only mode + state.engine_client = None + state.vllm_config = None + + state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) + state.server_load_metrics = 0 + + logger.info("Diffusion API server initialized for model: %s", model_name) + + @router.post( "/v1/chat/completions", dependencies=[Depends(validate_json_request)], @@ -337,9 +499,13 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re try: generator = await handler.create_chat_completion(request, raw_request) except Exception as e: + logger.exception("Chat completion failed: %s", e) raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e + if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), status_code=generator.error.code) + return JSONResponse( + content=generator.model_dump(), status_code=generator.code if hasattr(generator, "code") else 400 + ) elif isinstance(generator, ChatCompletionResponse): return JSONResponse(content=generator.model_dump()) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 94ab41906..5d42ae3d2 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -6,10 +6,11 @@ import uuid from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence from datetime import datetime, timedelta, timezone from io import BytesIO -from typing import Any +from typing import TYPE_CHECKING, Any, Optional import jinja2 from fastapi import Request +from PIL import Image from pydantic import TypeAdapter try: @@ -71,10 +72,52 @@ from vllm.utils import as_list from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures from vllm_omni.outputs import OmniRequestOutput +if TYPE_CHECKING: + from vllm_omni.entrypoints.async_diffusion import AsyncOmniDiffusion + logger = init_logger(__name__) class OmniOpenAIServingChat(OpenAIServingChat): + """OpenAI-compatible chat serving for both LLM and Diffusion models. + + This class extends OpenAIServingChat to support: + - Standard LLM chat completions + - Diffusion model image generation via chat interface + + For diffusion mode, use the `for_diffusion` class method to create an instance. + """ + + # Diffusion mode attributes + _diffusion_mode: bool = False + _diffusion_engine: Optional["AsyncOmniDiffusion"] = None + _diffusion_model_name: str = "" + + @classmethod + def for_diffusion( + cls, + diffusion_engine: "AsyncOmniDiffusion", + model_name: str, + ) -> "OmniOpenAIServingChat": + """Create a chat serving instance for diffusion models. + + Args: + diffusion_engine: The async diffusion engine + model_name: Name of the model being served + + Returns: + OmniOpenAIServingChat instance configured for diffusion mode + + Note: + Request-level parameters (num_inference_steps, guidance_scale, seed, + height, width, num_frames, fps, etc.) are passed per-request via the API. + """ + instance = cls.__new__(cls) + instance._diffusion_mode = True + instance._diffusion_engine = diffusion_engine + instance._diffusion_model_name = model_name + return instance + async def create_chat_completion( self, request: ChatCompletionRequest, @@ -86,7 +129,14 @@ class OmniOpenAIServingChat(OpenAIServingChat): See https://platform.openai.com/docs/api-reference/chat/create for the API specification. This API mimics the OpenAI Chat Completion API. + + For diffusion models, this generates images and returns them + in a chat completion response format. """ + # Handle diffusion mode + if self._diffusion_mode: + return await self._create_diffusion_chat_completion(request, raw_request) + error_check_ret = await self._check_model(request) if error_check_ret is not None: logger.error("Error with model %s", error_check_ret) @@ -438,6 +488,8 @@ class OmniOpenAIServingChat(OpenAIServingChat): ) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role) elif omni_outputs.final_output_type == "audio": choices_data = self._create_audio_choice(omni_outputs, role) + elif omni_outputs.final_output_type == "image": + choices_data = self._create_image_choice(omni_outputs, role) else: logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}") continue @@ -780,3 +832,335 @@ class OmniOpenAIServingChat(OpenAIServingChat): ) choices.append(choice_data) return choices + + def _create_image_choice(self, omni_outputs: OmniRequestOutput, role: str): + """Create chat completion response choices for image output. + + Converts image tensor or PIL Image output from diffusion models + into base64-encoded image data for API response. + + Args: + omni_outputs: Output containing image data from diffusion stage + role: The role for the response message (e.g., "assistant") + + Returns: + List of ChatCompletionResponseChoice with image content + """ + from PIL import Image + + choices: list[ChatCompletionResponseChoice] = [] + final_res = omni_outputs.request_output + + # Handle different image output formats + images = [] + + # First check omni_outputs.images directly (for diffusion mode via from_diffusion) + if omni_outputs.images: + images = omni_outputs.images + # Fall back to request_output for pipeline mode + elif final_res is not None: + if hasattr(final_res, "multimodal_output") and final_res.multimodal_output: + image_data = final_res.multimodal_output.get("image") + if image_data is not None: + if isinstance(image_data, Image.Image): + images.append(image_data) + elif hasattr(image_data, "cpu"): # Tensor + import numpy as np + + # Convert tensor to PIL Image + img_array = image_data.float().detach().cpu().numpy() + # Handle different tensor formats (CHW -> HWC) + if img_array.ndim == 3 and img_array.shape[0] in [1, 3, 4]: + img_array = np.transpose(img_array, (1, 2, 0)) + # Normalize to 0-255 + if img_array.max() <= 1.0: + img_array = (img_array * 255).astype(np.uint8) + else: + img_array = img_array.astype(np.uint8) + # Handle grayscale + if img_array.ndim == 2: + images.append(Image.fromarray(img_array, mode="L")) + elif img_array.shape[-1] == 1: + images.append(Image.fromarray(img_array.squeeze(-1), mode="L")) + elif img_array.shape[-1] == 3: + images.append(Image.fromarray(img_array, mode="RGB")) + elif img_array.shape[-1] == 4: + images.append(Image.fromarray(img_array, mode="RGBA")) + elif hasattr(final_res, "images") and final_res.images: + images = final_res.images + + # Convert images to base64 + image_contents = [] + for img in images: + with BytesIO() as buffer: + img.save(buffer, format="PNG") + img_bytes = buffer.getvalue() + img_base64 = base64.b64encode(img_bytes).decode("utf-8") + image_contents.append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{img_base64}", + }, + } + ) + + # Create message content + if len(image_contents) == 1: + content = image_contents + elif len(image_contents) > 1: + content = image_contents + else: + content = [{"type": "text", "text": "Image generation completed but no images were produced."}] + + # Create response choice + choice_data = ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role=role, content=content), + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + choices.append(choice_data) + + return choices + + # ==================== Diffusion Mode Methods ==================== + + async def _create_diffusion_chat_completion( + self, + request: ChatCompletionRequest, + raw_request: Request | None = None, + ) -> ChatCompletionResponse | ErrorResponse: + """Generate images via chat completion interface for diffusion models. + + Args: + request: Chat completion request + raw_request: Raw FastAPI request object + + Returns: + ChatCompletionResponse with generated images or ErrorResponse + """ + try: + request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}" + created_time = int(time.time()) + + # Convert messages to dict format + messages = [] + for msg in request.messages: + if hasattr(msg, "model_dump"): + messages.append(msg.model_dump()) + elif isinstance(msg, dict): + messages.append(msg) + else: + messages.append({"role": getattr(msg, "role", "user"), "content": getattr(msg, "content", "")}) + + # Extract prompt and images from messages + prompt, reference_images = self._extract_diffusion_prompt_and_images(messages) + + if not prompt: + return self._create_error_response("No text prompt found in messages") + + # Extract generation parameters from extra_body (preferred) + # Reference: text_to_image.py and text_to_video.py for supported parameters + extra_body = getattr(request, "extra_body", None) or {} + + # Parse size if provided (supports "1024x1024" format) + height = extra_body.get("height") + width = extra_body.get("width") + if "size" in extra_body: + try: + size_str = extra_body["size"] + if isinstance(size_str, str) and "x" in size_str.lower(): + w, h = size_str.lower().split("x") + width, height = int(w), int(h) + except ValueError: + logger.warning("Invalid size format: %s", extra_body.get("size")) + + # 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) + true_cfg_scale = extra_body.get("true_cfg_scale") # Qwen-Image specific + seed = extra_body.get("seed") + negative_prompt = extra_body.get("negative_prompt") + num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1) + + # Text-to-video parameters (ref: text_to_video.py) + num_frames = extra_body.get("num_frames") + guidance_scale_2 = extra_body.get("guidance_scale_2") # For video high-noise CFG + + logger.info( + "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s", + request_id, + prompt[:50] + "..." if len(prompt) > 50 else prompt, + len(reference_images), + {k: v for k, v in extra_body.items() if v is not None}, + ) + + # Decode reference images if provided + pil_images: list[Image.Image] = [] + for img_b64 in reference_images: + try: + img_bytes = base64.b64decode(img_b64) + pil_images.append(Image.open(BytesIO(img_bytes))) + except Exception as e: + logger.warning("Failed to decode reference image: %s", e) + + # Build generation kwargs + gen_kwargs: dict[str, Any] = { + "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, + } + + # Add Qwen-Image specific parameter + if true_cfg_scale is not None: + gen_kwargs["true_cfg_scale"] = true_cfg_scale + + # Add video generation parameters if set + if num_frames is not None: + gen_kwargs["num_frames"] = num_frames + if guidance_scale_2 is not None: + gen_kwargs["guidance_scale_2"] = guidance_scale_2 + + # Add reference image if provided + if pil_images: + gen_kwargs["pil_image"] = pil_images[0] + + # Generate image + result = await self._diffusion_engine.generate(**gen_kwargs) + + # Convert images to base64 content + image_contents: list[dict[str, Any]] = [] + for img in result.images: + with BytesIO() as buffer: + img.save(buffer, format="PNG") + img_bytes = buffer.getvalue() + img_base64 = base64.b64encode(img_bytes).decode("utf-8") + image_contents.append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{img_base64}", + }, + } + ) + + # Build response + if not image_contents: + content = "Image generation completed but no images were produced." + else: + content = image_contents + + # Use model_construct to bypass validation for multimodal content + # (ChatMessage.content only accepts str, but we need list for images) + message = ChatMessage.model_construct(role="assistant", content=content) + choice = ChatCompletionResponseChoice.model_construct( + index=0, + message=message, + finish_reason="stop", + logprobs=None, + stop_reason=None, + ) + + response = ChatCompletionResponse( + id=request_id, + created=created_time, + model=self._diffusion_model_name, + choices=[choice], + usage=UsageInfo( + prompt_tokens=len(prompt.split()), + completion_tokens=1, + total_tokens=len(prompt.split()) + 1, + ), + ) + + logger.info( + "Diffusion chat completed for request %s: %d images", + request_id, + len(result.images), + ) + + return response + + except Exception as e: + logger.exception("Diffusion chat completion failed: %s", e) + return self._create_error_response( + f"Image generation failed: {str(e)}", + status_code=500, + ) + + def _extract_diffusion_prompt_and_images( + self, + messages: list[dict[str, Any]], + ) -> tuple[str, list[str]]: + """Extract text prompt and base64 images from chat messages. + + Args: + messages: List of chat messages + + Returns: + Tuple of (prompt_text, list_of_base64_images) + """ + prompt_parts: list[str] = [] + images: list[str] = [] + + for message in messages: + role = message.get("role", "") + if role != "user": + continue + + content = message.get("content", "") + + # String content + if isinstance(content, str): + prompt_parts.append(content) + continue + + # List of content items + if isinstance(content, list): + for item in content: + if isinstance(item, str): + prompt_parts.append(item) + elif isinstance(item, dict): + # Handle {"type": "text", "text": "..."} format + if item.get("type") == "text": + prompt_parts.append(item.get("text", "")) + # Handle {"text": "..."} format + elif "text" in item and "type" not in item: + prompt_parts.append(item["text"]) + # Handle {"type": "image_url", "image_url": {"url": "..."}} + elif item.get("type") == "image_url": + url = item.get("image_url", {}).get("url", "") + if url.startswith("data:image"): + try: + _, b64_data = url.split(",", 1) + images.append(b64_data) + except ValueError: + logger.warning("Invalid data URL format") + # Handle {"image": "base64..."} format + elif "image" in item: + images.append(item["image"]) + + prompt = " ".join(prompt_parts).strip() + return prompt, images + + def _create_error_response( + self, + message: str, + err_type: str = "BadRequestError", + status_code: int = 400, + ) -> ErrorResponse: + """Create an error response following OpenAI error format.""" + return ErrorResponse( + message=message, + type=err_type, + code=status_code, + ) diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 47b58f7f3..1fe3112d7 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -1,6 +1,8 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field +from typing import Any import torch +from PIL import Image from vllm.outputs import RequestOutput from vllm.v1.outputs import ModelRunnerOutput @@ -20,20 +22,132 @@ class OmniModelRunnerOutput(ModelRunnerOutput): @dataclass -class OmniRequestOutput(RequestOutput): - """Request output for omni pipeline stages. +class OmniRequestOutput: + """Unified request output for both pipeline stages and diffusion models. - Wraps a standard RequestOutput with stage-specific metadata, - indicating which stage produced the output and what type of - final output it represents. + This class handles outputs from: + 1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output) + 2. Diffusion models (with images, prompt, metrics) Attributes: - stage_id: Identifier of the stage that produced this output - final_output_type: Type of final output (e.g., "text", "image", - "audio", "latents") - request_output: The underlying RequestOutput from the stage + request_id: Unique identifier for this request + finished: Whether generation is complete + stage_id: Identifier of the stage that produced this output (pipeline mode) + final_output_type: Type of output ("text", "image", "audio", "latents") + request_output: The underlying RequestOutput from the stage (pipeline mode) + images: List of generated PIL images (diffusion mode) + prompt: The prompt used for generation (diffusion mode) + latents: Optional tensor of latent representations (diffusion mode) + metrics: Optional dictionary of generation metrics """ - stage_id: int - final_output_type: str - request_output: RequestOutput + request_id: str = "" + finished: bool = True + + # Pipeline stage fields + stage_id: int | None = None + final_output_type: str = "text" + request_output: RequestOutput | None = None + + # Diffusion model fields + images: list[Image.Image] = field(default_factory=list) + prompt: str | None = None + latents: torch.Tensor | None = None + metrics: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_pipeline( + cls, + stage_id: int, + final_output_type: str, + request_output: RequestOutput, + ) -> "OmniRequestOutput": + """Create output from pipeline stage. + + Args: + stage_id: Stage identifier + final_output_type: Type of output + request_output: The stage's output + + Returns: + OmniRequestOutput configured for pipeline mode + """ + return cls( + request_id=getattr(request_output, "request_id", ""), + stage_id=stage_id, + final_output_type=final_output_type, + request_output=request_output, + finished=True, + ) + + @classmethod + def from_diffusion( + cls, + request_id: str, + images: list[Image.Image], + prompt: str | None = None, + metrics: dict[str, Any] | None = None, + latents: torch.Tensor | None = None, + ) -> "OmniRequestOutput": + """Create output from diffusion model. + + Args: + request_id: Request identifier + images: Generated images + prompt: The prompt used + metrics: Generation metrics + latents: Optional latent tensors + + Returns: + OmniRequestOutput configured for diffusion mode + """ + return cls( + request_id=request_id, + final_output_type="image", + images=images, + prompt=prompt, + latents=latents, + metrics=metrics or {}, + finished=True, + ) + + @property + def num_images(self) -> int: + """Return the number of generated images.""" + return len(self.images) + + @property + def is_diffusion_output(self) -> bool: + """Check if this is a diffusion model output.""" + return len(self.images) > 0 or self.final_output_type == "image" + + @property + def is_pipeline_output(self) -> bool: + """Check if this is a pipeline stage output.""" + return self.stage_id is not None and self.request_output is not None + + def to_dict(self) -> dict[str, Any]: + """Convert to dictionary for JSON serialization.""" + result = { + "request_id": self.request_id, + "finished": self.finished, + "final_output_type": self.final_output_type, + } + + if self.is_diffusion_output: + result.update( + { + "num_images": self.num_images, + "prompt": self.prompt, + "metrics": self.metrics, + } + ) + + if self.is_pipeline_output: + result.update( + { + "stage_id": self.stage_id, + } + ) + + return result