[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>
This commit is contained in:
Chenguang Zhengandyinpe authored and GitHub committed
1 parent
8722c6374c
commit
acf66ee2cc
19 files changed
+2506
-20
No files matched your search
@@ -0,0 +1,197 @@
|
||||
# Qwen-Image-Edit Online Serving
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen_image_edit>.
|
||||
|
||||
|
||||
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"
|
||||
``````
|
||||
@@ -0,0 +1,178 @@
|
||||
# Qwen-Image/Z-Image Online Serving
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen_image>.
|
||||
|
||||
|
||||
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"
|
||||
``````
|
||||
@@ -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 |
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
+12
@@ -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"
|
||||
@@ -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 |
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
+12
@@ -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"
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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=<ConfigGroup>` 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
|
||||
|
||||
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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())
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
+127
-13
@@ -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
|
||||
Reference in new issue
Block a user