Merge branch 'main' of https://github.com/vllm-project/vllm-omni
This commit is contained in:
Taichang Zhou committed
commit
354777d5e7
29 files changed
+2031
-454
No files matched your search
+72
-37
@@ -233,43 +233,43 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Test with H100"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_async_omni.py
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
# - label: "Omni Model Test with H100"
|
||||
# timeout_in_minutes: 30
|
||||
# depends_on: image-build
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
# - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
|
||||
# - pytest -s -v tests/e2e/online_serving/test_async_omni.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 2
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
timeout_in_minutes: 20
|
||||
@@ -305,3 +305,38 @@ steps:
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# timeout_in_minutes: 30
|
||||
# depends_on: image-build
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
@@ -15,6 +15,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
---
|
||||
|
||||
*Latest News* 🔥
|
||||
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability.
|
||||
- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
|
||||
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ nav:
|
||||
- Serving:
|
||||
- OpenAI-Compatible API:
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Examples:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
|
||||
@@ -204,7 +204,151 @@ Key point for writing the example:
|
||||
|
||||
+ Save or display the generated results so users can validate the integration.
|
||||
|
||||
## Step 5: Open a Pull Request
|
||||
## Step 5: TeaCache Coefficient Estimation (Optional)
|
||||
|
||||
If your model supports TeaCache acceleration, you need to estimate the polynomial coefficients for optimal caching performance.
|
||||
|
||||
### 5.1 Add Extractor Function
|
||||
|
||||
First, implement an extractor function in `vllm_omni/diffusion/cache/teacache/extractors.py`. The extractor extracts the modulated input and defines how to run transformer blocks:
|
||||
|
||||
```python
|
||||
def extract_your_model_context(
|
||||
module: nn.Module,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
**kwargs: Any,
|
||||
) -> CacheContext:
|
||||
# 1. Preprocessing
|
||||
temb = module.time_embed(timestep)
|
||||
|
||||
# 2. Extract modulated input (for cache decision)
|
||||
modulated_input = module.transformer_blocks[0].norm1(hidden_states, temb)
|
||||
|
||||
# 3. Define transformer execution
|
||||
def run_transformer_blocks():
|
||||
h = hidden_states
|
||||
for block in module.transformer_blocks:
|
||||
h = block(h, temb=temb)
|
||||
return (h,)
|
||||
|
||||
# 4. Define postprocessing
|
||||
def postprocess(h):
|
||||
return module.proj_out(module.norm_out(h, temb))
|
||||
|
||||
return CacheContext(
|
||||
modulated_input=modulated_input,
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
temb=temb,
|
||||
run_transformer_blocks=run_transformer_blocks,
|
||||
postprocess=postprocess,
|
||||
)
|
||||
```
|
||||
|
||||
Register it in `EXTRACTOR_REGISTRY`:
|
||||
```python
|
||||
EXTRACTOR_REGISTRY = {
|
||||
...
|
||||
"YourTransformer2DModel": extract_your_model_context,
|
||||
}
|
||||
```
|
||||
|
||||
### 5.2 Add Adapter for Coefficient Estimation
|
||||
|
||||
Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`:
|
||||
|
||||
```python
|
||||
class YourModelAdapter:
|
||||
@staticmethod
|
||||
def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
|
||||
# Load your pipeline
|
||||
...
|
||||
|
||||
@staticmethod
|
||||
def get_transformer(pipeline: Any) -> tuple[Any, str]:
|
||||
return pipeline.transformer, "YourTransformer2DModel"
|
||||
|
||||
@staticmethod
|
||||
def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
|
||||
registry = HookRegistry.get_or_create(transformer)
|
||||
registry.register_hook(hook._HOOK_NAME, hook)
|
||||
|
||||
_MODEL_ADAPTERS["YourModel"] = YourModelAdapter
|
||||
```
|
||||
|
||||
### 5.3 Run Coefficient Estimation
|
||||
|
||||
Use the provided script to estimate coefficients:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.cache.teacache.coefficient_estimator import (
|
||||
TeaCacheCoefficientEstimator,
|
||||
)
|
||||
from datasets import load_dataset
|
||||
from tqdm import tqdm
|
||||
|
||||
# Load model
|
||||
estimator = TeaCacheCoefficientEstimator(
|
||||
model_path="/path/to/model",
|
||||
model_type="Bagel", # Your model type
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
# Load prompts (paper suggests ~70 prompts)
|
||||
dataset = load_dataset("nateraw/parti-prompts", split="train")
|
||||
prompts = dataset["Prompt"][:70]
|
||||
|
||||
# Collect data
|
||||
for prompt in tqdm(prompts):
|
||||
estimator.collect_from_prompt(prompt, num_inference_steps=50)
|
||||
|
||||
# Estimate coefficients
|
||||
coeffs = estimator.estimate(poly_order=4)
|
||||
print(f"Coefficients: {coeffs}")
|
||||
```
|
||||
|
||||
### 5.4 Interpreting Coefficient Estimation Results
|
||||
|
||||
The estimator outputs statistics and polynomial coefficients. Here's how to interpret them:
|
||||
|
||||
**Example Output:**
|
||||
```
|
||||
Data statistics:
|
||||
Count: 48
|
||||
Input Diffs (x): min=1.1089e-02, max=5.2555e-02, mean=2.8435e-02
|
||||
Output Diffs (y): min=2.8242e-02, max=2.9792e-01, mean=7.0312e-02
|
||||
Coefficients: [1333131.29, -168644.23, 7950.51, -163.75, 1.26]
|
||||
```
|
||||
|
||||
**What to Check:**
|
||||
- **Count**: Number of timestep pairs analyzed. Should be at least 30-50 for reliable estimation. Low count suggests insufficient prompts or inference steps.
|
||||
- **Input/Output Ranges**: Verify output differences correlate with input differences. If ranges seem unusual, check your prompt diversity.
|
||||
- **Coefficient Magnitude**: Extremely large values (>1e8) may indicate numerical instability - try collecting more diverse data.
|
||||
|
||||
**Troubleshooting:**
|
||||
- If results seem unreliable, try:
|
||||
- Increasing number of prompts (100+ recommended)
|
||||
- Using more diverse prompts from multiple datasets
|
||||
- Adjusting `num_inference_steps` (try 20, 50, 100)
|
||||
|
||||
### 5.5 Add Coefficients to Config
|
||||
|
||||
Add the estimated coefficients to `vllm_omni/diffusion/cache/teacache/config.py`:
|
||||
|
||||
```python
|
||||
_MODEL_COEFFICIENTS = {
|
||||
...
|
||||
"YourTransformer2DModel": [
|
||||
1.04730573e+06, # a4
|
||||
-1.34150749e+05, # a3
|
||||
6.51517806e+03, # a2
|
||||
-1.41209108e+02, # a1
|
||||
1.17241808e+00, # a0
|
||||
],
|
||||
}
|
||||
```
|
||||
## Step 6: Open a Pull Request
|
||||
|
||||
When submitting a pull request to add support for a new model, please include the following information in the PR description:
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
|
||||
# --8<-- [start:pre-built-wheels]
|
||||
|
||||
#### Installation of vLLM
|
||||
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
@@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.14.0rc1 \
|
||||
vllm/vllm-omni:v0.14.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
|
||||
```
|
||||
|
||||
|
||||
@@ -80,7 +80,7 @@ docker run --rm \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.14.0rc1 \
|
||||
vllm/vllm-omni-rocm:v0.14.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -99,7 +99,7 @@ docker run --rm -it \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--entrypoint bash \
|
||||
vllm/vllm-omni-rocm:v0.14.0rc1
|
||||
vllm/vllm-omni-rocm:v0.14.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -0,0 +1,205 @@
|
||||
# Image Edit API
|
||||
|
||||
vLLM-Omni provides an OpenAI DALL-E compatible API for image edit using diffusion models.
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Start the Server
|
||||
|
||||
For example...
|
||||
|
||||
```bash
|
||||
# Qwen-Image
|
||||
vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
|
||||
|
||||
|
||||
### Generate Images
|
||||
|
||||
**Using curl:**
|
||||
|
||||
```bash
|
||||
curl -s -D >(grep -i x-request-id >&2) \
|
||||
-o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
|
||||
-X POST "http://localhost:8000/v1/images/edits" \
|
||||
-F "model=xxx" \
|
||||
-F "image=@./xx.png" \
|
||||
-F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
|
||||
-F "size=1024x1024" \
|
||||
-F "output_format=png"
|
||||
```
|
||||
|
||||
|
||||
**Using OpenAI SDK:**
|
||||
|
||||
```python
|
||||
import base64
|
||||
from openai import OpenAI
|
||||
from pathlib import Path
|
||||
client = OpenAI(
|
||||
api_key="None",
|
||||
base_url="http://localhost:8000/v1"
|
||||
)
|
||||
|
||||
input_image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
|
||||
|
||||
result = client.images.edit(
|
||||
image=[],
|
||||
model="Qwen-Image-Edit-2511",
|
||||
prompt="Change the bears in the two input images into walking together.",
|
||||
size='512x512',
|
||||
stream=False,
|
||||
output_format='jpeg',
|
||||
# url格式
|
||||
extra_body={
|
||||
"url": [input_image_url1,input_image_url],
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1,
|
||||
"seed": 777,
|
||||
}
|
||||
)
|
||||
|
||||
image_base64 = result.data[0].b64_json
|
||||
image_bytes = base64.b64decode(image_base64)
|
||||
|
||||
# Save the image to a file
|
||||
with open("edit_out_http.jpeg", "wb") as f:
|
||||
f.write(image_bytes)
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
POST /v1/images/edits
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
### Request Parameters
|
||||
|
||||
#### OpenAI Standard Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `prompt` | string | **required** | A text description of the desired image |
|
||||
| `model` | string | server's model | Model to use (optional, should match server if specified) |
|
||||
| `image` | string or array | **required** | The image(s) to edit. |
|
||||
| `n` | integer | 1 | Number of images to generate (1-10) |
|
||||
| `size` | string | "auto" | Image dimensions in WxH format (e.g., "1024x1024", "512x512"), when set to auto, it decide size from first input image. |
|
||||
| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
|
||||
| `user` | string | null | User identifier for tracking |
|
||||
| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
|
||||
| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
|
||||
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
|
||||
|
||||
#### vllm-omni Extension Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `url` | string or array | None | The image(s) to edit. |
|
||||
| `negative_prompt` | string | null | Text describing what to avoid in the image |
|
||||
| `num_inference_steps` | integer | model defaults | Number of diffusion steps |
|
||||
| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
|
||||
| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
|
||||
| `seed` | integer | null | Random seed for reproducibility |
|
||||
|
||||
### Response Format
|
||||
|
||||
```json
|
||||
{
|
||||
"created": 1701234567,
|
||||
"data": [
|
||||
{
|
||||
"b64_json": "<base64-encoded PNG>",
|
||||
"url": null,
|
||||
"revised_prompt": null
|
||||
}
|
||||
],
|
||||
"output_format": null,
|
||||
"size": null,
|
||||
}
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
### Multiple Images input
|
||||
|
||||
```bash
|
||||
curl -s -D >(grep -i x-request-id >&2) \
|
||||
-o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
|
||||
-X POST "http://localhost:8000/v1/images/edits" \
|
||||
-F "model=xxx" \
|
||||
-F "image=@xx.png" \
|
||||
-F "image=@xx.png"
|
||||
-F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
|
||||
-F "size=1024x1024" \
|
||||
-F "output_format=png"
|
||||
```
|
||||
|
||||
|
||||
## Parameter Handling
|
||||
|
||||
The API passes parameters directly to the diffusion pipeline without model-specific transformation:
|
||||
|
||||
- **Default values**: When parameters are not specified, the underlying model uses its own defaults
|
||||
- **Pass-through design**: User-provided values are forwarded directly to the diffusion engine
|
||||
- **Minimal validation**: Only basic type checking and range validation at the API level
|
||||
|
||||
### Parameter Compatibility
|
||||
|
||||
The API passes parameters directly to the diffusion pipeline without model-specific validation.
|
||||
|
||||
- Unsupported parameters may be silently ignored by the model
|
||||
- Incompatible values will result in errors from the underlying pipeline
|
||||
- Recommended values vary by model - consult model documentation
|
||||
|
||||
**Best Practice:** Start with the model's recommended parameters, then adjust based on your needs.
|
||||
|
||||
## Error Responses
|
||||
|
||||
### 400 Bad Request
|
||||
|
||||
Invalid parameters (e.g., model mismatch):
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": "Invalid size format: '1024x'. Expected format: 'WIDTHxHEIGHT' (e.g., '1024x1024')."
|
||||
}
|
||||
```
|
||||
|
||||
### 422 Unprocessable Entity
|
||||
|
||||
Validation errors (missing required fields):
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": "Field 'image' or 'url' is required"
|
||||
}
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Server Not Running
|
||||
|
||||
```bash
|
||||
# Check if server is responding
|
||||
curl -X http://localhost:8000/v1/images/edit \
|
||||
-F "prompt='test'"
|
||||
```
|
||||
|
||||
### Out of Memory
|
||||
|
||||
If you encounter OOM errors:
|
||||
1. Reduce image size: `"size": "512x512"`
|
||||
2. Reduce inference steps: `"num_inference_steps": 25`
|
||||
|
||||
## Development
|
||||
|
||||
Enable debug logging to see prompts and generation details:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image-Edit-2511 --omni \
|
||||
--uvicorn-log-level debug
|
||||
```
|
||||
@@ -18,20 +18,19 @@ The following table shows which models are currently supported by parallelism me
|
||||
|
||||
### ImageGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
|
||||
|-------|------------------|------------|---------|--------------|--------------------------|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
|
||||
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
|
||||
|
||||
!!! note "TP Limitations for Diffusion Models"
|
||||
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
|
||||
|
||||
+3
-2
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "vllm-omni"
|
||||
version = "0.14.0rc1"
|
||||
version = "0.14.0"
|
||||
description = "A framework for efficient model inference with omni-modality models"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10,<3.14"
|
||||
@@ -59,7 +59,8 @@ dev = [
|
||||
"psutil>=7.2.0",
|
||||
"soundfile>=0.13.1",
|
||||
"imageio[ffmpeg]>=0.6.0",
|
||||
"opencv-python>=4.12.0.88"
|
||||
"opencv-python>=4.12.0.88",
|
||||
"mooncake-transfer-engine==0.3.8.post1"
|
||||
]
|
||||
|
||||
docs = [
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
# stage config for running BAGEL with Mooncake connector for CI e2e tests.
|
||||
# This config is optimized for single GPU tests with Mooncake inter-stage communication.
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: BagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.35
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
distributed_executor_backend: mp
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
tensor_parallel_size: 1
|
||||
omni_kv_config:
|
||||
need_send_cache: true
|
||||
kv_transfer_criteria:
|
||||
type: prefill_finished
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 52
|
||||
detokenize: true
|
||||
repetition_penalty: 1.05
|
||||
output_connectors:
|
||||
to_stage_1: mooncake_connector
|
||||
- stage_id: 1
|
||||
stage_type: diffusion
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.55
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
distributed_executor_backend: mp
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
tensor_parallel_size: 1
|
||||
omni_kv_config:
|
||||
need_recv_cache: true
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: image
|
||||
is_comprehension: false
|
||||
default_sampling_params:
|
||||
seed: 52
|
||||
input_connectors:
|
||||
from_stage_0: mooncake_connector
|
||||
|
||||
# Top-level runtime config with Mooncake connector
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
connectors:
|
||||
mooncake_connector:
|
||||
name: MooncakeConnector
|
||||
extra:
|
||||
host: "${MOONCAKE_HOST}"
|
||||
metadata_server: "http://${MOONCAKE_HOST}:${MOONCAKE_HTTP_PORT}/metadata"
|
||||
master: "${MOONCAKE_HOST}:${MOONCAKE_RPC_PORT}"
|
||||
segment: 64000000
|
||||
localbuf: 64000000
|
||||
proto: tcp
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,83 @@
|
||||
# stage config for running BAGEL with SharedMemory connector for CI e2e tests.
|
||||
# This config is optimized for single GPU tests with SharedMemory inter-stage communication.
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: BagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.35
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
tensor_parallel_size: 1
|
||||
omni_kv_config:
|
||||
need_send_cache: true
|
||||
kv_transfer_criteria:
|
||||
type: prefill_finished #or special token generated
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 52
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: diffusion
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.55
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
tensor_parallel_size: 1
|
||||
omni_kv_config:
|
||||
need_recv_cache: true
|
||||
engine_input_source: [0]
|
||||
|
||||
final_output: true
|
||||
final_output_type: image
|
||||
is_comprehension: false
|
||||
default_sampling_params:
|
||||
seed: 52
|
||||
|
||||
# Runtime edges
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
|
||||
# Distributed connectors configuration (optional)
|
||||
# More connectors will be supported in the future.
|
||||
connectors:
|
||||
shared_memory_connector:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536 # 64KB threshold
|
||||
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,308 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
End-to-end test for Bagel text2img generation.
|
||||
|
||||
This test validates that the Bagel model generates images that match
|
||||
expected reference pixel values within a ±5 tolerance.
|
||||
|
||||
Equivalent to running:
|
||||
python3 examples/offline_inference/bagel/end2end.py \
|
||||
--prompts "A futuristic city skyline at twilight, cyberpunk style" \
|
||||
--modality text2img --step 15
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
# Reference pixel data extracted from the known-good output image
|
||||
# Each entry contains (x, y) position and expected (R, G, B) values
|
||||
# "Generated with seed=52, num_inference_steps=15,
|
||||
# prompt='A futuristic city skyline at twilight, cyberpunk style'"
|
||||
REFERENCE_PIXELS = [
|
||||
{"position": (100, 100), "rgb": (68, 107, 134)},
|
||||
{"position": (400, 50), "rgb": (95, 139, 166)},
|
||||
{"position": (700, 100), "rgb": (99, 122, 151)},
|
||||
{"position": (150, 400), "rgb": (111, 125, 153)},
|
||||
{"position": (512, 512), "rgb": (97, 107, 131)},
|
||||
{"position": (700, 400), "rgb": (48, 64, 98)},
|
||||
{"position": (100, 700), "rgb": (79, 63, 84)},
|
||||
{"position": (400, 700), "rgb": (40, 58, 79)},
|
||||
{"position": (700, 700), "rgb": (60, 75, 103)},
|
||||
{"position": (256, 256), "rgb": (97, 128, 156)},
|
||||
]
|
||||
|
||||
# Maximum allowed difference per color channel
|
||||
PIXEL_TOLERANCE = 5
|
||||
|
||||
# Default test prompt
|
||||
DEFAULT_PROMPT = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|im_end|>"
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
"""Find and return a free ephemeral port by binding to port 0."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
s.listen(1)
|
||||
port = s.getsockname()[1]
|
||||
return port
|
||||
|
||||
|
||||
def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
|
||||
"""Configure sampling parameters for Bagel text2img generation.
|
||||
|
||||
Args:
|
||||
omni: The Omni instance to get default params from.
|
||||
max_tokens: Maximum tokens for the first stage.
|
||||
num_inference_steps: Number of inference steps for the diffusion stage.
|
||||
|
||||
Returns:
|
||||
Configured sampling params list.
|
||||
"""
|
||||
params_list = omni.default_sampling_params_list
|
||||
params_list[0].max_tokens = max_tokens # type: ignore
|
||||
if len(params_list) > 1:
|
||||
params_list[1].num_inference_steps = num_inference_steps # type: ignore
|
||||
return params_list
|
||||
|
||||
|
||||
def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
|
||||
"""Extract the generated image from Omni outputs.
|
||||
|
||||
Args:
|
||||
omni_outputs: List of outputs from omni.generate().
|
||||
|
||||
Returns:
|
||||
The first generated PIL Image, or None if no image found.
|
||||
"""
|
||||
for req_output in omni_outputs:
|
||||
if images := getattr(req_output, "images", None):
|
||||
return images[0]
|
||||
if hasattr(req_output, "request_output") and req_output.request_output:
|
||||
for stage_out in req_output.request_output:
|
||||
if hasattr(stage_out, "images") and stage_out.images:
|
||||
return stage_out.images[0]
|
||||
return None
|
||||
|
||||
|
||||
def _validate_pixels(
|
||||
image: Image.Image,
|
||||
reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS,
|
||||
tolerance: int = PIXEL_TOLERANCE,
|
||||
) -> None:
|
||||
"""Validate that image pixels match expected reference values.
|
||||
|
||||
Args:
|
||||
image: The PIL Image to validate.
|
||||
reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B).
|
||||
tolerance: Maximum allowed difference per color channel.
|
||||
|
||||
Raises:
|
||||
AssertionError: If any pixel differs beyond tolerance.
|
||||
"""
|
||||
for ref in reference_pixels:
|
||||
x, y = ref["position"]
|
||||
expected = ref["rgb"]
|
||||
actual = image.getpixel((x, y))[:3]
|
||||
assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), (
|
||||
f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}"
|
||||
)
|
||||
|
||||
|
||||
def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Image:
|
||||
"""Generate an image using Bagel model with configured parameters.
|
||||
|
||||
Args:
|
||||
omni: The Omni instance to use for generation.
|
||||
prompt: The text prompt for image generation.
|
||||
|
||||
Returns:
|
||||
The generated PIL Image.
|
||||
|
||||
Raises:
|
||||
AssertionError: If no image is generated or size is incorrect.
|
||||
"""
|
||||
params_list = _configure_sampling_params(omni)
|
||||
|
||||
omni_outputs = list(
|
||||
omni.generate(
|
||||
prompts=[{"prompt": prompt, "modalities": ["image"]}],
|
||||
sampling_params_list=params_list,
|
||||
)
|
||||
)
|
||||
|
||||
generated_image = _extract_generated_image(omni_outputs)
|
||||
assert generated_image is not None, "No images generated"
|
||||
assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
|
||||
|
||||
return generated_image
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_bagel_text2img_shared_memory_connector():
|
||||
"""Test Bagel text2img with shared memory connector."""
|
||||
config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
|
||||
omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
|
||||
|
||||
try:
|
||||
generated_image = _generate_bagel_image(omni)
|
||||
_validate_pixels(generated_image)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
|
||||
def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool:
|
||||
"""Wait for a port to become available.
|
||||
|
||||
Args:
|
||||
host: The host address.
|
||||
port: The port number.
|
||||
timeout: Maximum seconds to wait.
|
||||
|
||||
Returns:
|
||||
True if port becomes available, False otherwise.
|
||||
"""
|
||||
for _ in range(timeout):
|
||||
try:
|
||||
with socket.create_connection((host, port), timeout=1):
|
||||
return True
|
||||
except (TimeoutError, ConnectionRefusedError):
|
||||
time.sleep(1)
|
||||
return False
|
||||
|
||||
|
||||
def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None:
|
||||
"""Clean up any existing mooncake_master processes.
|
||||
|
||||
Args:
|
||||
timeout_secs: Maximum seconds to wait for graceful termination.
|
||||
"""
|
||||
subprocess.run(
|
||||
["pkill", "-f", "mooncake_master"],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < timeout_secs:
|
||||
result = subprocess.run(
|
||||
["pgrep", "-f", "mooncake_master"],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
break
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
subprocess.run(
|
||||
["pkill", "-9", "-f", "mooncake_master"],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
time.sleep(1)
|
||||
|
||||
|
||||
def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
|
||||
"""Load Mooncake config from YAML and substitute placeholders.
|
||||
|
||||
Args:
|
||||
host: Mooncake host address.
|
||||
rpc_port: RPC port for Mooncake master.
|
||||
http_port: HTTP metadata server port.
|
||||
|
||||
Returns:
|
||||
Path to the temporary config file with substituted values.
|
||||
"""
|
||||
config_path = str(Path(__file__).parent / "stage_configs" / "bagel_mooncake_ci.yaml")
|
||||
with open(config_path) as f:
|
||||
config_content = f.read()
|
||||
|
||||
# Substitute placeholders
|
||||
config_content = config_content.replace("${MOONCAKE_HOST}", host)
|
||||
config_content = config_content.replace("${MOONCAKE_RPC_PORT}", str(rpc_port))
|
||||
config_content = config_content.replace("${MOONCAKE_HTTP_PORT}", str(http_port))
|
||||
|
||||
# Write to temp file
|
||||
temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
|
||||
temp_file.write(config_content)
|
||||
temp_file.close()
|
||||
return temp_file.name
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_bagel_text2img_mooncake_connector():
|
||||
"""Test Bagel text2img with Mooncake connector for inter-stage communication."""
|
||||
MOONCAKE_HOST = "127.0.0.1"
|
||||
MOONCAKE_RPC_PORT = _find_free_port()
|
||||
MOONCAKE_HTTP_PORT = _find_free_port()
|
||||
MOONCAKE_METRICS_PORT = _find_free_port()
|
||||
|
||||
mooncake_master_proc = None
|
||||
temp_config_file = None
|
||||
omni = None
|
||||
|
||||
try:
|
||||
_cleanup_mooncake_processes()
|
||||
|
||||
# Start mooncake_master
|
||||
mooncake_master_proc = subprocess.Popen(
|
||||
[
|
||||
"mooncake_master",
|
||||
f"--rpc_port={MOONCAKE_RPC_PORT}",
|
||||
"--enable_http_metadata_server=true",
|
||||
"--http_metadata_server_host=0.0.0.0",
|
||||
f"--http_metadata_server_port={MOONCAKE_HTTP_PORT}",
|
||||
f"--metrics_port={MOONCAKE_METRICS_PORT}",
|
||||
],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
preexec_fn=os.setsid,
|
||||
)
|
||||
|
||||
assert _wait_for_port(MOONCAKE_HOST, MOONCAKE_RPC_PORT), "mooncake_master failed to start"
|
||||
|
||||
# Create temp config and initialize Omni
|
||||
temp_config_file = _load_mooncake_config(
|
||||
host=MOONCAKE_HOST,
|
||||
rpc_port=MOONCAKE_RPC_PORT,
|
||||
http_port=MOONCAKE_HTTP_PORT,
|
||||
)
|
||||
|
||||
omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
|
||||
|
||||
generated_image = _generate_bagel_image(omni)
|
||||
_validate_pixels(generated_image)
|
||||
|
||||
finally:
|
||||
if omni:
|
||||
omni.close()
|
||||
if temp_config_file:
|
||||
try:
|
||||
os.unlink(temp_config_file)
|
||||
except OSError:
|
||||
pass
|
||||
if mooncake_master_proc:
|
||||
try:
|
||||
os.killpg(os.getpgid(mooncake_master_proc.pid), signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
@@ -9,6 +9,7 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
|
||||
|
||||
import base64
|
||||
import io
|
||||
from argparse import Namespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
@@ -115,9 +116,11 @@ class FakeAsyncOmni:
|
||||
self.stage_list = ["llm", "diffusion"]
|
||||
self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
|
||||
self.captured_sampling_params_list = None
|
||||
self.captured_prompt = None
|
||||
|
||||
async def generate(self, prompt, request_id, sampling_params_list):
|
||||
self.captured_sampling_params_list = sampling_params_list
|
||||
self.captured_prompt = prompt
|
||||
images = [Image.new("RGB", (64, 64), color="green")]
|
||||
yield MockGenerationResult(images)
|
||||
|
||||
@@ -133,6 +136,8 @@ def mock_async_diffusion():
|
||||
# Return n PIL images wrapped in result object
|
||||
print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
|
||||
n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
|
||||
mock.captured_sampling_params_list = kwargs["sampling_params_list"]
|
||||
mock.captured_prompt = kwargs["prompt"]
|
||||
images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
|
||||
return MockGenerationResult(images)
|
||||
|
||||
@@ -155,6 +160,10 @@ def test_client(mock_async_diffusion):
|
||||
app.state.diffusion_engine = mock_async_diffusion # Also set for health endpoint
|
||||
app.state.stage_configs = [{"stage_type": "diffusion"}]
|
||||
app.state.diffusion_model_name = "Qwen/Qwen-Image" # For models endpoint
|
||||
app.state.args = Namespace(
|
||||
default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
|
||||
max_generated_image_size=4096, # 64*64
|
||||
)
|
||||
|
||||
return TestClient(app)
|
||||
|
||||
@@ -171,7 +180,10 @@ def async_omni_test_client():
|
||||
|
||||
app.state.engine_client = FakeAsyncOmni()
|
||||
app.state.stage_configs = [{"stage_type": "llm"}, {"stage_type": "diffusion"}]
|
||||
|
||||
app.state.args = Namespace(
|
||||
default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
|
||||
max_generated_image_size=4096, # 64*64
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@@ -526,3 +538,279 @@ def test_model_field_omitted_works(test_client):
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
def make_test_image_bytes(size=(64, 64)) -> bytes:
|
||||
img = Image.new(
|
||||
"RGB",
|
||||
size,
|
||||
)
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_image_edit_images_processing(async_omni_test_client):
|
||||
img_bytes_1 = make_test_image_bytes((16, 16))
|
||||
img_bytes_2 = make_test_image_bytes((32, 32))
|
||||
|
||||
# uploadfile with image key
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[
|
||||
("image", img_bytes_1),
|
||||
("image", img_bytes_2),
|
||||
],
|
||||
data={"prompt": "hello world."},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
engine = async_omni_test_client.app.state.engine_client
|
||||
captured_prompt = engine.captured_prompt
|
||||
processed_images = captured_prompt["multi_modal_data"]["image"]
|
||||
assert len(processed_images) == 2
|
||||
assert isinstance(processed_images[0], Image.Image)
|
||||
assert isinstance(processed_images[1], Image.Image)
|
||||
assert processed_images[0].size == (16, 16)
|
||||
assert processed_images[1].size == (32, 32)
|
||||
|
||||
# uploadfile with image[] key
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[
|
||||
("image[]", img_bytes_2),
|
||||
("image[]", img_bytes_1),
|
||||
],
|
||||
data={"prompt": "hello world."},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
engine = async_omni_test_client.app.state.engine_client
|
||||
captured_prompt = engine.captured_prompt
|
||||
processed_images = captured_prompt["multi_modal_data"]["image"]
|
||||
assert len(processed_images) == 2
|
||||
assert isinstance(processed_images[0], Image.Image)
|
||||
assert isinstance(processed_images[1], Image.Image)
|
||||
assert processed_images[0].size == (32, 32)
|
||||
assert processed_images[1].size == (16, 16)
|
||||
|
||||
# base64 url
|
||||
buf1 = io.BytesIO()
|
||||
img1 = Image.new("RGB", (16, 16))
|
||||
img1.save(buf1, format="PNG")
|
||||
b64_1 = "data:image/png;base64," + base64.b64encode(buf1.getvalue()).decode()
|
||||
|
||||
buf2 = io.BytesIO()
|
||||
img2 = Image.new("RGB", (24, 24))
|
||||
img2.save(buf2, format="PNG")
|
||||
b64_2 = "data:image/png;base64," + base64.b64encode(buf2.getvalue()).decode()
|
||||
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
data={
|
||||
"prompt": "hello from base64",
|
||||
"url": [b64_1, b64_2],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
processed_images = engine.captured_prompt["multi_modal_data"]["image"]
|
||||
assert len(processed_images) == 2
|
||||
assert isinstance(processed_images[0], Image.Image)
|
||||
assert isinstance(processed_images[1], Image.Image)
|
||||
assert processed_images[0].size == (16, 16)
|
||||
assert processed_images[1].size == (24, 24)
|
||||
|
||||
|
||||
def test_image_edit_parameter_pass(async_omni_test_client):
|
||||
img_bytes_1 = make_test_image_bytes((16, 16))
|
||||
|
||||
# uploadfile with image key
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"size": "16x24",
|
||||
"output_format": "jpeg",
|
||||
"num_inference_steps": 20,
|
||||
"guidance_scale": 8.0,
|
||||
"seed": 1234,
|
||||
"negative_prompt": "negative",
|
||||
"n": 2,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
engine = async_omni_test_client.app.state.engine_client
|
||||
captured_prompt = engine.captured_prompt
|
||||
captured_sampling_params = engine.captured_sampling_params_list[-1]
|
||||
|
||||
assert captured_prompt["prompt"] == "hello world."
|
||||
assert captured_prompt["negative_prompt"] == "negative"
|
||||
assert captured_sampling_params.num_inference_steps == 20
|
||||
assert captured_sampling_params.guidance_scale == 8.0
|
||||
assert captured_sampling_params.seed == 1234
|
||||
assert captured_sampling_params.num_outputs_per_prompt == 2
|
||||
assert captured_sampling_params.width == 16
|
||||
assert captured_sampling_params.height == 24
|
||||
|
||||
data = response.json()
|
||||
# All images should be valid
|
||||
for img_data in data["data"]:
|
||||
assert "b64_json" in img_data
|
||||
img_bytes = base64.b64decode(img_data["b64_json"])
|
||||
img = Image.open(io.BytesIO(img_bytes))
|
||||
assert img.format.lower() == "jpeg"
|
||||
assert data["output_format"] == "jpeg"
|
||||
assert data["size"] == "16x24"
|
||||
|
||||
|
||||
def test_image_edit_parameter_default(async_omni_test_client):
|
||||
img_bytes_1 = make_test_image_bytes((24, 16))
|
||||
|
||||
# uploadfile with image key
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"size": "auto",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
engine = async_omni_test_client.app.state.engine_client
|
||||
captured_sampling_params = engine.captured_sampling_params_list[-1]
|
||||
|
||||
assert captured_sampling_params.width == 24
|
||||
assert captured_sampling_params.height == 16
|
||||
assert captured_sampling_params.num_outputs_per_prompt == 1
|
||||
assert captured_sampling_params.num_inference_steps == 4
|
||||
assert captured_sampling_params.guidance_scale == 7.5
|
||||
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"size": "96x96",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
def test_image_edit_parameter_default_single_stage(test_client):
|
||||
img_bytes_1 = make_test_image_bytes((24, 16))
|
||||
|
||||
# uploadfile with image key
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
engine = test_client.app.state.engine_client
|
||||
captured_sampling_params = engine.captured_sampling_params_list[0]
|
||||
|
||||
assert captured_sampling_params.width == 24
|
||||
assert captured_sampling_params.height == 16
|
||||
assert captured_sampling_params.num_outputs_per_prompt == 1
|
||||
assert captured_sampling_params.num_inference_steps == 4
|
||||
assert captured_sampling_params.guidance_scale == 7.5
|
||||
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"size": "96x96",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
def test_image_edit_compression_jpeg(test_client):
|
||||
img_bytes_1 = make_test_image_bytes((16, 16))
|
||||
# uploadfile with image key
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={"prompt": "hello world.", "output_format": "jpeg", "output_compression": 100},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
|
||||
img = Image.open(io.BytesIO(img_bytes_100))
|
||||
assert img.format.lower() == "jpeg"
|
||||
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"output_format": "jpeg",
|
||||
"output_compression": 50,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
|
||||
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"output_format": "jpeg",
|
||||
"output_compression": 10,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
|
||||
|
||||
assert len(img_bytes_10) < len(img_bytes_50)
|
||||
assert len(img_bytes_50) < len(img_bytes_100)
|
||||
|
||||
|
||||
def test_image_edit_compression_png(async_omni_test_client):
|
||||
img_bytes_1 = make_test_image_bytes((16, 16))
|
||||
# uploadfile with image key
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={"prompt": "hello world.", "output_format": "PNG", "output_compression": 100},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
|
||||
img = Image.open(io.BytesIO(img_bytes_100))
|
||||
assert img.format.lower() == "png"
|
||||
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"output_format": "PNG",
|
||||
"output_compression": 50,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
|
||||
|
||||
response = async_omni_test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes_1)],
|
||||
data={
|
||||
"prompt": "hello world.",
|
||||
"output_format": "PNG",
|
||||
"output_compression": 10,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
|
||||
|
||||
assert len(img_bytes_10) < len(img_bytes_50)
|
||||
assert len(img_bytes_50) < len(img_bytes_100)
|
||||
@@ -10,26 +10,8 @@ from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
AttentionMetadata,
|
||||
)
|
||||
|
||||
# Import flash attention functions with fallback chain from utils/fa.py
|
||||
# FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
|
||||
from vllm_omni.diffusion.attention.backends.utils.fa import (
|
||||
HAS_FLASH_ATTN,
|
||||
_pad_input,
|
||||
_unpad_input,
|
||||
_upad_input,
|
||||
flash_attn_func,
|
||||
flash_attn_varlen_func,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
if not HAS_FLASH_ATTN:
|
||||
raise ImportError(
|
||||
"FlashAttentionBackend requires Flash Attention. "
|
||||
"Please install one of: fa3-fwd, flash-attention, or flash-attn. "
|
||||
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
|
||||
)
|
||||
|
||||
|
||||
class FlashAttentionBackend(AttentionBackend):
|
||||
accept_output_buffer: bool = True
|
||||
@@ -74,6 +56,24 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
attn_metadata: AttentionMetadata = None,
|
||||
) -> torch.Tensor:
|
||||
"""CUDA/ROCm flash attention implementation."""
|
||||
# Import flash attention functions with fallback chain from utils/fa.py
|
||||
# FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
|
||||
from vllm_omni.diffusion.attention.backends.utils.fa import (
|
||||
HAS_FLASH_ATTN,
|
||||
_pad_input,
|
||||
_unpad_input,
|
||||
_upad_input,
|
||||
flash_attn_func,
|
||||
flash_attn_varlen_func,
|
||||
)
|
||||
|
||||
if not HAS_FLASH_ATTN:
|
||||
raise ImportError(
|
||||
"FlashAttentionBackend requires Flash Attention. "
|
||||
"Please install one of: fa3-fwd, flash-attention, or flash-attn. "
|
||||
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
|
||||
)
|
||||
|
||||
query_length = query.size(1)
|
||||
attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
|
||||
# Contains at least one padding token in the sequence
|
||||
@@ -122,7 +122,15 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
attn_metadata: AttentionMetadata = None,
|
||||
) -> torch.Tensor:
|
||||
"""NPU attention implementation using mindiesd."""
|
||||
from mindiesd import attention_forward
|
||||
try:
|
||||
from mindiesd import attention_forward
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"FlashAttentionBackend NPU implementation requires MindIE-SD. "
|
||||
"Please install MindIE-SD to enable NPU attention support. "
|
||||
"For installation details, see https://gitcode.com/Ascend/MindIE-SD"
|
||||
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
|
||||
)
|
||||
|
||||
attention_mask = attn_metadata.attn_mask if attn_metadata else None
|
||||
output = attention_forward(
|
||||
|
||||
@@ -13,6 +13,29 @@ from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _maybe_reshape_attn_mask(query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None = None):
|
||||
"""
|
||||
Reshape Attention Mask
|
||||
[batch_size, seq_len_k] -> [batch_size, 1, seq_len_q, seq_len_k]
|
||||
"""
|
||||
# Skip Attention Mask if all values are 1, `None` mask can speedup the computation
|
||||
if attn_mask is not None and torch.all(attn_mask != 0):
|
||||
attn_mask = None
|
||||
|
||||
# Reshape Attention Mask
|
||||
# [batch_size, seq_len_k] -> [batch_size, 1, seq_len_q, seq_len_k]
|
||||
if (
|
||||
attn_mask is not None
|
||||
and attn_mask.ndim == 2
|
||||
and attn_mask.shape[0] == query.shape[0]
|
||||
and attn_mask.shape[1] == key.shape[1]
|
||||
):
|
||||
B, Sq, Skv = attn_mask.shape[0], query.shape[1], key.shape[1]
|
||||
attn_mask = attn_mask.to(torch.bool)
|
||||
attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous()
|
||||
return attn_mask
|
||||
|
||||
|
||||
class SDPABackend(AttentionBackend):
|
||||
accept_output_buffer: bool = True
|
||||
|
||||
@@ -47,16 +70,15 @@ class SDPAImpl(AttentionImpl):
|
||||
self.causal = causal
|
||||
self.softmax_scale = softmax_scale
|
||||
|
||||
def forward(
|
||||
def forward_cuda(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata = None,
|
||||
attn_metadata: AttentionMetadata | None = None,
|
||||
) -> torch.Tensor:
|
||||
query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
|
||||
attention_mask = attn_metadata.attn_mask if attn_metadata else None
|
||||
|
||||
output = torch.nn.functional.scaled_dot_product_attention(
|
||||
query,
|
||||
key,
|
||||
@@ -68,3 +90,33 @@ class SDPAImpl(AttentionImpl):
|
||||
)
|
||||
out = output.permute(0, 2, 1, 3)
|
||||
return out
|
||||
|
||||
def forward_xpu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.forward_cuda(query, key, value, attn_metadata)
|
||||
|
||||
def forward_hip(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.forward_cuda(query, key, value, attn_metadata)
|
||||
|
||||
def forward_npu(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_metadata: AttentionMetadata | None = None,
|
||||
) -> torch.Tensor:
|
||||
if attn_metadata:
|
||||
attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask)
|
||||
setattr(attn_metadata, "attn_mask", attention_mask)
|
||||
return self.forward_cuda(query, key, value, attn_metadata)
|
||||
@@ -0,0 +1,197 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from vllm.config import LoadConfig
|
||||
|
||||
from vllm_omni.diffusion.cache.teacache.extractors import get_extractor
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
|
||||
|
||||
|
||||
class DataCollectionHook(ModelHook):
|
||||
"""Hook to collect modulated inputs and model outputs for TeaCache coefficient estimation."""
|
||||
|
||||
_HOOK_NAME = "teacache_collector"
|
||||
|
||||
def __init__(self, transformer_type: str):
|
||||
super().__init__()
|
||||
self.transformer_type = transformer_type
|
||||
self.extractor_fn = None
|
||||
self.current_trajectory: list[tuple[np.ndarray, np.ndarray]] = []
|
||||
|
||||
def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
|
||||
self.extractor_fn = get_extractor(self.transformer_type)
|
||||
return module
|
||||
|
||||
def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
|
||||
ctx = self.extractor_fn(module, *args, **kwargs)
|
||||
modulated_input_cpu = ctx.modulated_input.detach().cpu().numpy()
|
||||
|
||||
outputs = ctx.run_transformer_blocks()
|
||||
ctx.hidden_states = outputs[0]
|
||||
if len(outputs) > 1 and ctx.encoder_hidden_states is not None:
|
||||
ctx.encoder_hidden_states = outputs[1]
|
||||
|
||||
model_output_cpu = ctx.hidden_states.detach().cpu().numpy()
|
||||
self.current_trajectory.append((modulated_input_cpu, model_output_cpu))
|
||||
|
||||
return ctx.postprocess(ctx.hidden_states)
|
||||
|
||||
def start_collection(self):
|
||||
self.current_trajectory = []
|
||||
|
||||
def stop_collection(self) -> list[tuple[np.ndarray, np.ndarray]]:
|
||||
return list(self.current_trajectory)
|
||||
|
||||
|
||||
class BagelAdapter:
|
||||
"""Adapter for Bagel model."""
|
||||
|
||||
@staticmethod
|
||||
def load_pipeline(model_path: str, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> BagelPipeline:
|
||||
od_config = OmniDiffusionConfig.from_kwargs(model=model_path, dtype=dtype)
|
||||
od_config.model_class_name = "BagelPipeline"
|
||||
|
||||
pipeline = BagelPipeline(od_config=od_config)
|
||||
loader = DiffusersPipelineLoader(LoadConfig())
|
||||
loader.load_weights(pipeline)
|
||||
pipeline.to(device)
|
||||
return pipeline
|
||||
|
||||
@staticmethod
|
||||
def get_transformer(pipeline: Any) -> tuple[Any, str]:
|
||||
return pipeline.bagel, "Bagel"
|
||||
|
||||
@staticmethod
|
||||
def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
|
||||
original_forward_flow = transformer._forward_flow
|
||||
|
||||
def forward_alias(self, *args, **kwargs):
|
||||
return original_forward_flow(*args, **kwargs)
|
||||
|
||||
transformer.forward = types.MethodType(forward_alias, transformer)
|
||||
registry = HookRegistry.get_or_create(transformer)
|
||||
registry.register_hook(hook._HOOK_NAME, hook)
|
||||
transformer._forward_flow = transformer.forward
|
||||
|
||||
|
||||
class DefaultAdapter:
|
||||
"""Default adapter for standard diffusers pipelines."""
|
||||
|
||||
@staticmethod
|
||||
def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
|
||||
raise NotImplementedError("DefaultAdapter.load_pipeline not implemented")
|
||||
|
||||
@staticmethod
|
||||
def get_transformer(pipeline: Any) -> tuple[Any, str]:
|
||||
return pipeline.transformer, pipeline.transformer.__class__.__name__
|
||||
|
||||
@staticmethod
|
||||
def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
|
||||
registry = HookRegistry.get_or_create(transformer)
|
||||
registry.register_hook(hook._HOOK_NAME, hook)
|
||||
|
||||
|
||||
_MODEL_ADAPTERS: dict[str, type] = {
|
||||
"Bagel": BagelAdapter,
|
||||
}
|
||||
|
||||
_EPSILON = 1e-6
|
||||
|
||||
|
||||
def calculate_relative_l1(tensor_current: np.ndarray, tensor_next: np.ndarray) -> float:
|
||||
"""Calculate relative L1 distance (Eq. 4 from TeaCache paper)."""
|
||||
diff = np.abs(tensor_current - tensor_next).sum()
|
||||
norm = np.abs(tensor_current).sum() + _EPSILON
|
||||
return diff / norm
|
||||
|
||||
|
||||
def estimate_teacache_coefficients(
|
||||
collected_data: list[list[tuple[np.ndarray, np.ndarray]]], poly_order: int = 4
|
||||
) -> list[float]:
|
||||
"""Estimate polynomial coefficients for TeaCache using np.polyfit."""
|
||||
input_diffs, output_diffs = [], []
|
||||
|
||||
for sample in collected_data:
|
||||
for t in range(len(sample) - 1):
|
||||
feat_in_curr, feat_out_curr = sample[t]
|
||||
feat_in_next, feat_out_next = sample[t + 1]
|
||||
input_diffs.append(calculate_relative_l1(feat_in_curr, feat_in_next))
|
||||
output_diffs.append(calculate_relative_l1(feat_out_curr, feat_out_next))
|
||||
|
||||
x = np.array(input_diffs, dtype=np.float64)
|
||||
y = np.array(output_diffs, dtype=np.float64)
|
||||
|
||||
print("Data statistics:")
|
||||
print(f" Count: {len(x)}")
|
||||
print(f" Input Diffs (x): min={x.min():.4e}, max={x.max():.4e}, mean={x.mean():.4e}")
|
||||
print(f" Output Diffs (y): min={y.min():.4e}, max={y.max():.4e}, mean={y.mean():.4e}")
|
||||
|
||||
return np.polyfit(x, y, poly_order).tolist()
|
||||
|
||||
|
||||
class TeaCacheCoefficientEstimator:
|
||||
"""Model-agnostic helper class to collect data and estimate TeaCache coefficients."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
model_type: str = "Bagel",
|
||||
device: str = "cuda",
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
):
|
||||
# Add validation here ⬇️
|
||||
if model_type not in _MODEL_ADAPTERS:
|
||||
available_types = list(_MODEL_ADAPTERS.keys())
|
||||
raise ValueError(
|
||||
f"Unsupported model_type: '{model_type}'. "
|
||||
f"Available types: {available_types}. "
|
||||
f"To add support for a new model, add an entry to _MODEL_ADAPTERS."
|
||||
)
|
||||
|
||||
adapter = _MODEL_ADAPTERS.get(model_type, DefaultAdapter)
|
||||
self.pipeline = adapter.load_pipeline(model_path, device, dtype)
|
||||
self.transformer, self.transformer_type = adapter.get_transformer(self.pipeline)
|
||||
self.hook = DataCollectionHook(self.transformer_type)
|
||||
self.collected_data: list[list[tuple[np.ndarray, np.ndarray]]] = []
|
||||
adapter.install_hook(self.transformer, self.hook)
|
||||
|
||||
def collect_from_prompt(self, prompt: str, **generate_kwargs):
|
||||
self.hook.start_collection()
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
|
||||
req = OmniDiffusionRequest(
|
||||
prompt=prompt,
|
||||
num_inference_steps=generate_kwargs.get("num_inference_steps", 20),
|
||||
seed=generate_kwargs.get("seed", 42),
|
||||
)
|
||||
self.pipeline.forward(req)
|
||||
trajectory = self.hook.stop_collection()
|
||||
if trajectory:
|
||||
self.collected_data.append(trajectory)
|
||||
|
||||
def estimate(self, poly_order: int = 4) -> list[float]:
|
||||
"""Estimate polynomial coefficients from collected data.
|
||||
|
||||
Args:
|
||||
poly_order: Order of polynomial fit (default: 4)
|
||||
|
||||
Returns:
|
||||
List of polynomial coefficients [a_n, a_{n-1}, ..., a_1, a_0]
|
||||
|
||||
Raises:
|
||||
RuntimeError: If no data has been collected
|
||||
"""
|
||||
if not self.collected_data:
|
||||
raise RuntimeError(
|
||||
"No data collected for coefficient estimation. "
|
||||
"Call collect_from_prompt() at least once before calling estimate()."
|
||||
)
|
||||
return estimate_teacache_coefficients(self.collected_data, poly_order)
|
||||
@@ -9,10 +9,11 @@ import torch.nn as nn
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.activation import get_act_fn
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
@@ -35,9 +36,9 @@ class FeedForward(nn.Module):
|
||||
inner_dim = int(dim * mult)
|
||||
dim_out = dim_out if dim_out is not None else dim
|
||||
|
||||
self.w_in = ReplicatedLinear(dim, inner_dim, bias=bias, return_bias=False)
|
||||
self.w_in = ColumnParallelLinear(dim, inner_dim, bias=bias, return_bias=False)
|
||||
self.act = get_act_fn("gelu_pytorch_tanh")
|
||||
self.w_out = ReplicatedLinear(inner_dim, dim_out, bias=bias, return_bias=False)
|
||||
self.w_out = RowParallelLinear(inner_dim, dim_out, bias=bias, return_bias=False)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.w_in(hidden_states)
|
||||
@@ -83,12 +84,11 @@ class LongCatImageAttention(nn.Module):
|
||||
hidden_size=query_dim,
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=self.heads,
|
||||
disable_tp=True,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
if not self.pre_only:
|
||||
self.to_out = torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)
|
||||
self.to_out = RowParallelLinear(self.inner_dim, self.out_dim, bias=out_bias)
|
||||
|
||||
if self.added_kv_proj_dim is not None:
|
||||
self.norm_added_q = RMSNorm(dim_head, eps=eps)
|
||||
@@ -98,11 +98,10 @@ class LongCatImageAttention(nn.Module):
|
||||
hidden_size=self.added_kv_proj_dim,
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=self.heads,
|
||||
disable_tp=True,
|
||||
bias=added_proj_bias,
|
||||
)
|
||||
|
||||
self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)
|
||||
self.to_add_out = RowParallelLinear(self.inner_dim, query_dim, bias=out_bias)
|
||||
|
||||
self.attn = Attention(
|
||||
num_heads=heads,
|
||||
@@ -185,22 +184,26 @@ class LongCatImageAttention(nn.Module):
|
||||
"""
|
||||
qkv, _ = self.to_qkv(hidden_states)
|
||||
|
||||
query, key, value = qkv.chunk(3, dim=-1)
|
||||
q_size = self.to_qkv.num_heads * self.head_dim
|
||||
kv_size = self.to_qkv.num_kv_heads * self.head_dim
|
||||
query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
query = query.unflatten(-1, (self.heads, -1))
|
||||
key = key.unflatten(-1, (self.heads, -1))
|
||||
value = value.unflatten(-1, (self.heads, -1))
|
||||
query = query.unflatten(-1, (self.to_qkv.num_heads, -1))
|
||||
key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
|
||||
value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
|
||||
|
||||
query = self.norm_q(query)
|
||||
key = self.norm_k(key)
|
||||
|
||||
if self.added_kv_proj_dim is not None:
|
||||
encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
|
||||
encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1)
|
||||
q_size = self.add_kv_proj.num_heads * self.head_dim
|
||||
kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
|
||||
encoder_query, encoder_key, encoder_value = encoder_qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
|
||||
encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
|
||||
encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
|
||||
encoder_query = encoder_query.unflatten(-1, (self.add_kv_proj.num_heads, -1))
|
||||
encoder_key = encoder_key.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
|
||||
encoder_value = encoder_value.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
|
||||
|
||||
# Apply RMSNorm to text Q/K
|
||||
encoder_query = self.norm_added_q(encoder_query)
|
||||
@@ -284,11 +287,14 @@ class LongCatImageAttention(nn.Module):
|
||||
encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
|
||||
[encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
|
||||
)
|
||||
hidden_states = self.to_out(hidden_states)
|
||||
hidden_states, _ = self.to_out(hidden_states)
|
||||
encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)
|
||||
|
||||
return hidden_states, encoder_hidden_states
|
||||
else:
|
||||
# For single-stream blocks, there's no to_out (RowParallelLinear) to handle the reduction
|
||||
if get_tensor_model_parallel_world_size() > 1:
|
||||
hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1)
|
||||
return hidden_states
|
||||
|
||||
|
||||
|
||||
@@ -232,6 +232,22 @@ class OmniServeCommand(CLISubcommand):
|
||||
help="Number of devices for CFG parallel computation for diffusion models. "
|
||||
"Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.",
|
||||
)
|
||||
|
||||
# Default sampling parameters
|
||||
omni_config_group.add_argument(
|
||||
"--default-sampling-params",
|
||||
type=str,
|
||||
help="Json str for Default sampling parameters, \n"
|
||||
'Structure: {"<stage_id>": {<sampling_param>: value, ...}, ...}\n'
|
||||
'e.g., \'{"0": {"num_inference_steps":50, "guidance_scale":1}}\'. '
|
||||
"Currently only supports diffusion models.",
|
||||
)
|
||||
# Diffusion model mixed precision
|
||||
omni_config_group.add_argument(
|
||||
"--max-generated-image-size",
|
||||
type=float,
|
||||
help="The max size of generate image (height * width).",
|
||||
)
|
||||
return serve_parser
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import multiprocessing
|
||||
import multiprocessing.forkserver as forkserver
|
||||
import os
|
||||
@@ -12,11 +15,13 @@ from argparse import Namespace
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from http import HTTPStatus
|
||||
from typing import Any, cast
|
||||
from typing import Annotated, Any, cast
|
||||
|
||||
import httpx
|
||||
import vllm.envs as envs
|
||||
from fastapi import Depends, HTTPException, Request
|
||||
from fastapi import Depends, File, Form, HTTPException, Request, UploadFile
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from PIL import Image
|
||||
from starlette.datastructures import State
|
||||
from starlette.routing import Route
|
||||
from vllm import SamplingParams
|
||||
@@ -890,6 +895,251 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
|
||||
Raises:
|
||||
HTTPException: For validation errors, missing engine, or generation failures
|
||||
"""
|
||||
# Get engine client (AsyncOmni) from app state
|
||||
engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
|
||||
|
||||
# Validate model field (warn if mismatch, don't error)
|
||||
if request.model is not None and request.model != model_name:
|
||||
logger.warning(
|
||||
f"Model mismatch: request specifies '{request.model}' but "
|
||||
f"server is running '{model_name}'. Using server model."
|
||||
)
|
||||
|
||||
try:
|
||||
# Build params - pass through user values directly
|
||||
prompt: OmniTextPrompt = {"prompt": request.prompt}
|
||||
if request.negative_prompt is not None:
|
||||
prompt["negative_prompt"] = request.negative_prompt
|
||||
gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
|
||||
|
||||
# Parse per-request LoRA (compatible with chat's extra_body.lora shape).
|
||||
lora_request, lora_scale = _parse_lora_request(request.lora)
|
||||
_update_if_not_none(gen_params, "lora_request", lora_request)
|
||||
_update_if_not_none(gen_params, "lora_scale", lora_scale)
|
||||
|
||||
# Parse and add size if provided
|
||||
width, height = None, None
|
||||
if request.size:
|
||||
width, height = parse_size(request.size)
|
||||
size_str = f"{width}x{height}"
|
||||
else:
|
||||
size_str = "model default"
|
||||
_update_if_not_none(gen_params, "width", width)
|
||||
_update_if_not_none(gen_params, "height", height)
|
||||
|
||||
# 3.3 Add optional parameters ONLY if provided
|
||||
_update_if_not_none(gen_params, "num_inference_steps", request.num_inference_steps)
|
||||
_update_if_not_none(gen_params, "guidance_scale", request.guidance_scale)
|
||||
_update_if_not_none(gen_params, "true_cfg_scale", request.true_cfg_scale)
|
||||
# If seed is not provided, generate a random one to ensure
|
||||
# a proper generator is initialized in the backend.
|
||||
# This fixes issues where using the default global generator
|
||||
# might produce blurry images in some environments.
|
||||
_update_if_not_none(gen_params, "seed", random.randint(0, 2**32 - 1) if request.seed is None else request.seed)
|
||||
|
||||
request_id = f"img_gen_{uuid.uuid4().hex}"
|
||||
|
||||
logger.info(f"Generating {request.n} image(s) {size_str}")
|
||||
|
||||
# Generate images using AsyncOmni (multi-stage mode)
|
||||
result = await _generate_with_async_omni(
|
||||
engine_client=engine_client,
|
||||
gen_params=gen_params,
|
||||
stage_types=stage_types,
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
if result is None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
|
||||
detail="No output generated from multi-stage pipeline.",
|
||||
)
|
||||
|
||||
# Extract images from result
|
||||
images = _extract_images_from_result(result)
|
||||
|
||||
logger.info(f"Successfully generated {len(images)} image(s)")
|
||||
|
||||
# Encode images to base64
|
||||
image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
|
||||
|
||||
return ImageGenerationResponse(
|
||||
created=int(time.time()),
|
||||
data=image_data,
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
# Re-raise HTTPExceptions as-is
|
||||
raise
|
||||
except ValueError as e:
|
||||
logger.error(f"Validation error: {e}")
|
||||
raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Image generation failed: {e}")
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image generation failed: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/images/edits",
|
||||
responses={
|
||||
HTTPStatus.OK.value: {"model": ImageGenerationResponse},
|
||||
HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
|
||||
HTTPStatus.SERVICE_UNAVAILABLE.value: {"model": ErrorResponse},
|
||||
HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
|
||||
},
|
||||
)
|
||||
async def edit_images(
|
||||
raw_request: Request,
|
||||
image: list[UploadFile] | None = File(None),
|
||||
image_array: list[UploadFile] | None = File(None, alias="image[]"),
|
||||
url: list[str] | None = Form(None),
|
||||
url_array: list[str] | None = Form(None, alias="url[]"),
|
||||
prompt: str = Form(...),
|
||||
model: str = Form(None),
|
||||
n: int = Form(1),
|
||||
size: str = Form("auto"),
|
||||
response_format: str = Form("b64_json"),
|
||||
output_format: str | None = Form("png"),
|
||||
background: str | None = Form("auto"),
|
||||
output_compression: Annotated[int, Form(ge=0, le=100)] = 100,
|
||||
user: str | None = Form(None), # unused now
|
||||
# vllm-omni extensions for diffusion control
|
||||
negative_prompt: str | None = Form(None),
|
||||
num_inference_steps: int | None = Form(None),
|
||||
guidance_scale: float | None = Form(None),
|
||||
true_cfg_scale: float | None = Form(None),
|
||||
seed: int | None = Form(None),
|
||||
# vllm-omni extension for per-request LoRA.
|
||||
lora: str | None = Form(None), # Json string
|
||||
) -> ImageGenerationResponse:
|
||||
"""
|
||||
OpenAI-compatible image edit endpoint.
|
||||
"""
|
||||
# 1. get engine and model
|
||||
engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
|
||||
if model is not None and model != model_name:
|
||||
logger.warning(
|
||||
f"Model mismatch: request specifies '{model}' but server is running '{model_name}'. Using server model."
|
||||
)
|
||||
# 2. get output format & compression
|
||||
output_format = _choose_output_format(output_format, background)
|
||||
if response_format != "b64_json":
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Only response_format 'b64_json' is supported now.",
|
||||
)
|
||||
try:
|
||||
# 2. Build prompt & images params
|
||||
prompt: OmniTextPrompt = {"prompt": prompt}
|
||||
if negative_prompt is not None:
|
||||
prompt["negative_prompt"] = negative_prompt
|
||||
input_images_list = []
|
||||
images = image or image_array
|
||||
urls = url or url_array
|
||||
if images:
|
||||
input_images_list.extend(images)
|
||||
if urls:
|
||||
input_images_list.extend(urls)
|
||||
if not input_images_list:
|
||||
raise HTTPException(status_code=422, detail="Field 'image' or 'url' is required")
|
||||
pil_images = await _load_input_images(input_images_list)
|
||||
prompt["multi_modal_data"] = {}
|
||||
prompt["multi_modal_data"]["image"] = pil_images
|
||||
|
||||
# 3 Build sample params
|
||||
gen_params = OmniDiffusionSamplingParams()
|
||||
# 3.0 Init with system default values
|
||||
app_state_args = getattr(raw_request.app.state, "args", None)
|
||||
default_sample_param = getattr(app_state_args, "default_sampling_params", None)
|
||||
# Currently only have one diffusion stage
|
||||
diffusion_stage_id = [i for i, t in enumerate(stage_types) if t == "diffusion"][0]
|
||||
apply_stage_default_sampling_params(
|
||||
default_sample_param,
|
||||
gen_params,
|
||||
str(diffusion_stage_id),
|
||||
)
|
||||
_update_if_not_none(gen_params, "num_outputs_per_prompt", n)
|
||||
# 3.1 Parse per-request LoRA (compatible with chat's extra_body.lora shape).
|
||||
lora_dict = _get_lora_from_json_str(lora)
|
||||
lora_request, lora_scale = _parse_lora_request(lora_dict)
|
||||
_update_if_not_none(gen_params, "lora_request", lora_request)
|
||||
_update_if_not_none(gen_params, "lora_scale", lora_scale)
|
||||
# 3.2 Parse and add size if provided
|
||||
max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None)
|
||||
width, height = None, None
|
||||
if size.lower() == "auto":
|
||||
width, height = pil_images[0].size # Use first image size
|
||||
else:
|
||||
width, height = parse_size(size)
|
||||
if max_generated_image_size is not None and (width * height > max_generated_image_size):
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail=f"Requested image size {width}x{height} exceeds the maximum allowed "
|
||||
f"size of {max_generated_image_size} pixels.",
|
||||
)
|
||||
|
||||
size_str = f"{width}x{height}"
|
||||
_update_if_not_none(gen_params, "width", width)
|
||||
_update_if_not_none(gen_params, "height", height)
|
||||
|
||||
# 3.3 Add optional parameters ONLY if provided
|
||||
_update_if_not_none(gen_params, "num_inference_steps", num_inference_steps)
|
||||
_update_if_not_none(gen_params, "guidance_scale", guidance_scale)
|
||||
_update_if_not_none(gen_params, "true_cfg_scale", true_cfg_scale)
|
||||
# If seed is not provided, generate a random one to ensure
|
||||
# a proper generator is initialized in the backend.
|
||||
# This fixes issues where using the default global generator
|
||||
# might produce blurry images in some environments.
|
||||
_update_if_not_none(gen_params, "seed", seed or random.randint(0, 2**32 - 1))
|
||||
|
||||
# 4. Generate images using AsyncOmni (multi-stage mode)
|
||||
request_id = f"img_edit_{int(time.time())}"
|
||||
logger.info(f"Generating {n} image(s) {size_str}")
|
||||
result = await _generate_with_async_omni(
|
||||
engine_client=engine_client,
|
||||
gen_params=gen_params,
|
||||
stage_types=stage_types,
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
# 5. Extract images from result
|
||||
images = _extract_images_from_result(result)
|
||||
logger.info(f"Successfully generated {len(images)} image(s)")
|
||||
|
||||
# Encode images to base64
|
||||
image_data = [
|
||||
ImageData(
|
||||
b64_json=_encode_image_base64_with_compression(
|
||||
img, format=output_format, output_compression=output_compression
|
||||
),
|
||||
revised_prompt=None,
|
||||
)
|
||||
for img in images
|
||||
]
|
||||
|
||||
return ImageGenerationResponse(
|
||||
created=int(time.time()),
|
||||
data=image_data,
|
||||
output_format=output_format,
|
||||
size=size_str,
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
# Re-raise HTTPExceptions as-is
|
||||
raise
|
||||
except ValueError as e:
|
||||
logger.error(f"Validation error: {e}")
|
||||
raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Image edit failed: {e}")
|
||||
raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image edit failed: {str(e)}")
|
||||
|
||||
|
||||
def _get_engine_and_model(raw_request: Request):
|
||||
# Get engine client (AsyncOmni) from app state
|
||||
engine_client: EngineClient | AsyncOmni | None = getattr(raw_request.app.state, "engine_client", None)
|
||||
if engine_client is None or not hasattr(engine_client, "stage_list"):
|
||||
@@ -942,157 +1192,229 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
|
||||
else:
|
||||
model_name = "unknown"
|
||||
|
||||
# Validate model field (warn if mismatch, don't error)
|
||||
if request.model is not None and request.model != model_name:
|
||||
logger.warning(
|
||||
f"Model mismatch: request specifies '{request.model}' but "
|
||||
f"server is running '{model_name}'. Using server model."
|
||||
)
|
||||
return engine_client, model_name, stage_types
|
||||
|
||||
|
||||
def _get_lora_from_json_str(lora_body):
|
||||
if lora_body is None:
|
||||
return None
|
||||
try:
|
||||
# Build params - pass through user values directly
|
||||
prompt: OmniTextPrompt = {"prompt": request.prompt}
|
||||
gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
|
||||
lora_dict = json.loads(lora_body)
|
||||
except json.JSONDecodeError:
|
||||
raise HTTPException(status_code=400, detail="Invalid LoRA JSON string")
|
||||
|
||||
# Parse per-request LoRA (compatible with chat's extra_body.lora shape).
|
||||
if request.lora is not None:
|
||||
if not isinstance(request.lora, dict):
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid lora field: expected an object.",
|
||||
)
|
||||
lora_body = request.lora
|
||||
lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
|
||||
lora_path = (
|
||||
lora_body.get("local_path")
|
||||
or lora_body.get("path")
|
||||
or lora_body.get("lora_path")
|
||||
or lora_body.get("lora_local_path")
|
||||
)
|
||||
lora_scale = lora_body.get("scale")
|
||||
if lora_scale is None:
|
||||
lora_scale = lora_body.get("lora_scale")
|
||||
lora_int_id = lora_body.get("int_id")
|
||||
if lora_int_id is None:
|
||||
lora_int_id = lora_body.get("lora_int_id")
|
||||
if lora_int_id is None and lora_path:
|
||||
lora_int_id = stable_lora_int_id(str(lora_path))
|
||||
if not isinstance(lora_dict, dict):
|
||||
raise HTTPException(status_code=400, detail="LoRA must be a JSON object")
|
||||
|
||||
if not lora_name or not lora_path:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid lora object: both name and path are required.",
|
||||
)
|
||||
return lora_dict
|
||||
|
||||
gen_params.lora_request = LoRARequest(str(lora_name), int(lora_int_id), str(lora_path))
|
||||
if lora_scale is not None:
|
||||
gen_params.lora_scale = float(lora_scale)
|
||||
|
||||
# Parse and add size if provided
|
||||
if request.size:
|
||||
width, height = parse_size(request.size)
|
||||
gen_params.height = height
|
||||
gen_params.width = width
|
||||
size_str = f"{width}x{height}"
|
||||
else:
|
||||
size_str = "model default"
|
||||
|
||||
# Add optional parameters ONLY if provided
|
||||
if request.num_inference_steps is not None:
|
||||
gen_params.num_inference_steps = request.num_inference_steps
|
||||
if request.negative_prompt is not None:
|
||||
prompt["negative_prompt"] = request.negative_prompt
|
||||
if request.guidance_scale is not None:
|
||||
gen_params.guidance_scale = request.guidance_scale
|
||||
if request.true_cfg_scale is not None:
|
||||
gen_params.true_cfg_scale = request.true_cfg_scale
|
||||
if request.seed is not None:
|
||||
gen_params.seed = request.seed
|
||||
else:
|
||||
# If seed is not provided, generate a random one to ensure
|
||||
# a proper generator is initialized in the backend.
|
||||
# This fixes issues where using the default global generator
|
||||
# might produce blurry images in some environments.
|
||||
gen_params.seed = random.randint(0, 2**32 - 1)
|
||||
|
||||
request_id = f"img_gen_{uuid.uuid4().hex}"
|
||||
|
||||
logger.info(f"Generating {request.n} image(s) {size_str}")
|
||||
|
||||
# Generate images using AsyncOmni (multi-stage mode)
|
||||
engine_client = cast(AsyncOmni, engine_client)
|
||||
result = None
|
||||
stage_list = getattr(engine_client, "stage_list", None)
|
||||
if isinstance(stage_list, list):
|
||||
default_params_list: list[OmniSamplingParams] | None = getattr(
|
||||
engine_client, "default_sampling_params_list", None
|
||||
)
|
||||
if not isinstance(default_params_list, list):
|
||||
default_params_list = [
|
||||
OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
|
||||
]
|
||||
else:
|
||||
default_params_list = list(default_params_list)
|
||||
if len(default_params_list) != len(stage_types):
|
||||
default_params_list = (
|
||||
default_params_list
|
||||
+ [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
|
||||
)[: len(stage_types)]
|
||||
|
||||
sampling_params_list: list[OmniSamplingParams] = []
|
||||
for idx, stage_type in enumerate(stage_types):
|
||||
if stage_type == "diffusion":
|
||||
sampling_params_list.append(gen_params)
|
||||
else:
|
||||
base_params = default_params_list[idx]
|
||||
sampling_params_list.append(base_params)
|
||||
|
||||
async for output in engine_client.generate(
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
sampling_params_list=sampling_params_list,
|
||||
):
|
||||
result = output
|
||||
else:
|
||||
result = await engine_client.generate(
|
||||
prompt=prompt, request_id=request_id, sampling_params_list=[gen_params]
|
||||
)
|
||||
|
||||
if result is None:
|
||||
def _parse_lora_request(lora_body: dict[str, Any]):
|
||||
if lora_body is not None:
|
||||
if not isinstance(lora_body, dict):
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
|
||||
detail="No output generated from multi-stage pipeline.",
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid lora field: expected an object.",
|
||||
)
|
||||
lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
|
||||
lora_path = (
|
||||
lora_body.get("local_path")
|
||||
or lora_body.get("path")
|
||||
or lora_body.get("lora_path")
|
||||
or lora_body.get("lora_local_path")
|
||||
)
|
||||
lora_scale = lora_body.get("scale")
|
||||
if lora_scale is None:
|
||||
lora_scale = lora_body.get("lora_scale")
|
||||
lora_int_id = lora_body.get("int_id")
|
||||
if lora_int_id is None:
|
||||
lora_int_id = lora_body.get("lora_int_id")
|
||||
if lora_int_id is None and lora_path:
|
||||
lora_int_id = stable_lora_int_id(str(lora_path))
|
||||
|
||||
if not lora_name or not lora_path:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid lora object: both name and path are required.",
|
||||
)
|
||||
|
||||
# Extract images from result
|
||||
images = []
|
||||
if hasattr(result, "images") and result.images:
|
||||
images = result.images
|
||||
elif hasattr(result, "request_output"):
|
||||
request_output = result.request_output
|
||||
if isinstance(request_output, dict) and request_output.get("images"):
|
||||
images = request_output["images"]
|
||||
elif hasattr(request_output, "images") and request_output.images:
|
||||
images = request_output.images
|
||||
return LoRARequest(str(lora_name), int(lora_int_id), str(lora_path)), lora_scale
|
||||
return None, None
|
||||
|
||||
logger.info(f"Successfully generated {len(images)} image(s)")
|
||||
|
||||
# Encode images to base64
|
||||
image_data = [ImageData(b64_json=encode_image_base64(img), revised_prompt=None) for img in images]
|
||||
async def _generate_with_async_omni(
|
||||
engine_client: AsyncOmni | Any,
|
||||
gen_params: Any,
|
||||
stage_types: list[str],
|
||||
**kwargs,
|
||||
):
|
||||
engine_client = cast(AsyncOmni, engine_client)
|
||||
result = None
|
||||
stage_list = getattr(engine_client, "stage_list", None)
|
||||
if isinstance(stage_list, list):
|
||||
default_params_list: list[OmniSamplingParams] | None = getattr(
|
||||
engine_client, "default_sampling_params_list", None
|
||||
)
|
||||
if not isinstance(default_params_list, list):
|
||||
default_params_list = [
|
||||
OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
|
||||
]
|
||||
else:
|
||||
default_params_list = list(default_params_list)
|
||||
if len(default_params_list) != len(stage_types):
|
||||
default_params_list = (
|
||||
default_params_list
|
||||
+ [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
|
||||
)[: len(stage_types)]
|
||||
|
||||
return ImageGenerationResponse(
|
||||
created=int(time.time()),
|
||||
data=image_data,
|
||||
sampling_params_list: list[OmniSamplingParams] = []
|
||||
for idx, stage_type in enumerate(stage_types):
|
||||
if stage_type == "diffusion":
|
||||
sampling_params_list.append(gen_params)
|
||||
else:
|
||||
base_params = default_params_list[idx]
|
||||
sampling_params_list.append(base_params)
|
||||
|
||||
async for output in engine_client.generate(
|
||||
sampling_params_list=sampling_params_list,
|
||||
**kwargs,
|
||||
):
|
||||
result = output
|
||||
else:
|
||||
result = await engine_client.generate(
|
||||
sampling_params_list=[gen_params],
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
# Re-raise HTTPExceptions as-is
|
||||
raise
|
||||
except ValueError as e:
|
||||
logger.error(f"Validation error: {e}")
|
||||
raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.exception(f"Image generation failed: {e}")
|
||||
if result is None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image generation failed: {str(e)}"
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
|
||||
detail="No output generated from multi-stage pipeline.",
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _update_if_not_none(object: any, key: str, val: any) -> None:
|
||||
if val is not None:
|
||||
setattr(object, key, val)
|
||||
|
||||
|
||||
def _extract_images_from_result(result: Any) -> list[Any]:
|
||||
images = []
|
||||
if hasattr(result, "images") and result.images:
|
||||
images = result.images
|
||||
elif hasattr(result, "request_output"):
|
||||
request_output = result.request_output
|
||||
if isinstance(request_output, dict) and request_output.get("images"):
|
||||
images = request_output["images"]
|
||||
elif hasattr(request_output, "images") and request_output.images:
|
||||
images = request_output.images
|
||||
return images
|
||||
|
||||
|
||||
async def _load_input_images(
|
||||
inputs: list[str],
|
||||
) -> list[Image.Image]:
|
||||
"""
|
||||
convert to PIL.Image.Image list
|
||||
"""
|
||||
if isinstance(inputs, str):
|
||||
inputs = [inputs]
|
||||
|
||||
images: list[Image.Image] = []
|
||||
|
||||
for inp in inputs:
|
||||
# 1. URL + base64
|
||||
if isinstance(inp, str) and inp.startswith("data:image"):
|
||||
try:
|
||||
_, b64_data = inp.split(",", 1)
|
||||
image_bytes = base64.b64decode(b64_data)
|
||||
img = Image.open(io.BytesIO(image_bytes))
|
||||
images.append(img)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid base64 image: {e}")
|
||||
|
||||
# 2. URL
|
||||
elif isinstance(inp, str) and inp.startswith("http"):
|
||||
async with httpx.AsyncClient(timeout=60) as client:
|
||||
try:
|
||||
resp = await client.get(inp)
|
||||
resp.raise_for_status()
|
||||
img = Image.open(io.BytesIO(resp.content))
|
||||
images.append(img)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to download image from URL {inp}: {e}")
|
||||
|
||||
# 3. UploadFile
|
||||
elif hasattr(inp, "file"):
|
||||
try:
|
||||
img_data = await inp.read()
|
||||
img = Image.open(io.BytesIO(img_data))
|
||||
images.append(img)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to open uploaded file: {e}")
|
||||
else:
|
||||
raise ValueError(f"Unsupported input: {inp}")
|
||||
|
||||
if not images:
|
||||
raise ValueError("No valid input images found")
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def _choose_output_format(output_format: str | None, background: str | None) -> str:
|
||||
# Normalize and choose extension
|
||||
fmt = (output_format or "").lower()
|
||||
if fmt in {"jpg", "png", "webp", "jpeg"}:
|
||||
return fmt
|
||||
# If transparency requested, prefer png
|
||||
if (background or "auto").lower() == "transparent":
|
||||
return "png"
|
||||
# Default
|
||||
return "jpeg"
|
||||
|
||||
|
||||
def _encode_image_base64_with_compression(
|
||||
image: Image.Image, format: str = "png", output_compression: int = 100
|
||||
) -> str:
|
||||
"""Encode PIL Image to base64 PNG string.
|
||||
|
||||
Args:
|
||||
image: PIL Image object
|
||||
format: Output image format (e.g., "PNG", "JPEG", "WEBP")
|
||||
output_compression: Compression level (0-100%), 100 for best quality
|
||||
Returns:
|
||||
Base64-encoded image as string
|
||||
"""
|
||||
buffer = io.BytesIO()
|
||||
save_kwargs = {}
|
||||
if format in ("jpg", "jpeg", "webp"):
|
||||
save_kwargs["quality"] = output_compression
|
||||
elif format == "png":
|
||||
save_kwargs["compress_level"] = max(0, min(9, 9 - output_compression // 11)) # Map 0-100 to 9-0
|
||||
|
||||
image.save(buffer, format=format, **save_kwargs)
|
||||
buffer.seek(0)
|
||||
return base64.b64encode(buffer.read()).decode("utf-8")
|
||||
|
||||
|
||||
def apply_stage_default_sampling_params(
|
||||
default_params_json: str | None,
|
||||
sampling_params: any,
|
||||
stage_key: str,
|
||||
) -> None:
|
||||
"""
|
||||
Update a stage's sampling parameters with vLLM-Omni defaults.
|
||||
|
||||
Args:
|
||||
default_params_json: JSON string of stage-keyed default parameters
|
||||
sampling_params: The sampling parameters object to update
|
||||
stage_key: The stage ID/key in the pipeline
|
||||
"""
|
||||
if default_params_json is not None:
|
||||
default_params_dict = json.loads(default_params_json)
|
||||
if stage_key in default_params_dict:
|
||||
stage_defaults = default_params_dict[stage_key]
|
||||
for param_name, param_value in stage_defaults.items():
|
||||
if hasattr(sampling_params, param_name):
|
||||
setattr(sampling_params, param_name, param_value)
|
||||
@@ -123,3 +123,5 @@ class ImageGenerationResponse(BaseModel):
|
||||
|
||||
created: int = Field(..., description="Unix timestamp of when the generation completed")
|
||||
data: list[ImageData] = Field(..., description="Array of generated images")
|
||||
output_format: str = Field(None, description="The output format of the image generation")
|
||||
size: str = Field(None, description="The size of the image generated")
|
||||
@@ -287,8 +287,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
|
||||
generators.append(generator)
|
||||
except ValueError as e:
|
||||
# TODO: Use a vllm-specific Validation Error
|
||||
return self.create_error_response(str(e))
|
||||
return self.create_error_response(e)
|
||||
|
||||
assert len(generators) == 1
|
||||
(result_generator,) = generators
|
||||
@@ -316,8 +315,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
request_metadata,
|
||||
)
|
||||
except ValueError as e:
|
||||
# TODO: Use a vllm-specific Validation Error
|
||||
return self.create_error_response(str(e))
|
||||
return self.create_error_response(e)
|
||||
|
||||
async def _preprocess_chat(
|
||||
self,
|
||||
@@ -634,7 +632,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
)
|
||||
except RuntimeError as e:
|
||||
logger.exception("Error in reasoning parser creation.")
|
||||
data = self.create_streaming_error_response(str(e))
|
||||
data = self.create_streaming_error_response(e)
|
||||
yield f"data: {data}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
return
|
||||
@@ -646,7 +644,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
tool_parsers = [None] * num_choices
|
||||
except Exception as e:
|
||||
logger.exception("Error in tool parser creation.")
|
||||
data = self.create_streaming_error_response(str(e))
|
||||
data = self.create_streaming_error_response(e)
|
||||
yield f"data: {data}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
return
|
||||
@@ -1286,9 +1284,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# TODO: Use a vllm-specific Validation Error
|
||||
logger.exception("Error in chat completion stream generator.")
|
||||
data = self.create_streaming_error_response(str(e))
|
||||
data = self.create_streaming_error_response(e)
|
||||
yield f"data: {data}\n\n"
|
||||
# Send the final done message after all response.n are finished
|
||||
yield "data: [DONE]\n\n"
|
||||
@@ -1313,8 +1310,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
except asyncio.CancelledError:
|
||||
return self.create_error_response("Client disconnected")
|
||||
except ValueError as e:
|
||||
# TODO: Use a vllm-specific Validation Error
|
||||
return self.create_error_response(str(e))
|
||||
return self.create_error_response(e)
|
||||
|
||||
assert final_outputs is not None
|
||||
|
||||
@@ -1487,7 +1483,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
)
|
||||
except RuntimeError as e:
|
||||
logger.exception("Error in reasoning parser creation.")
|
||||
return self.create_error_response(str(e))
|
||||
return self.create_error_response(e)
|
||||
# If the reasoning parser is enabled,
|
||||
# tool calls are extracted exclusively from the content.
|
||||
reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request)
|
||||
@@ -1572,7 +1568,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
tool_parser = self.tool_parser(tokenizer)
|
||||
except RuntimeError as e:
|
||||
logger.exception("Error in tool parser creation.")
|
||||
return self.create_error_response(str(e))
|
||||
return self.create_error_response(e)
|
||||
|
||||
tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request)
|
||||
# In the OpenAI API the finish_reason is "tools_called"
|
||||
|
||||
@@ -296,7 +296,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
except asyncio.CancelledError:
|
||||
return self.create_error_response("Client disconnected")
|
||||
except ValueError as e:
|
||||
return self.create_error_response(str(e))
|
||||
return self.create_error_response(e)
|
||||
except Exception as e:
|
||||
logger.exception("Speech generation failed: %s", e)
|
||||
return self.create_error_response(f"Speech generation failed: {e}")
|
||||
@@ -11,7 +11,7 @@ stage_args:
|
||||
model_arch: BagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.4
|
||||
gpu_memory_utilization: 0.35
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
@@ -42,7 +42,7 @@ stage_args:
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.4
|
||||
gpu_memory_utilization: 0.55
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
|
||||
@@ -11,7 +11,7 @@ stage_args:
|
||||
model_arch: BagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.4
|
||||
gpu_memory_utilization: 0.35
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
@@ -41,11 +41,11 @@ stage_args:
|
||||
- stage_id: 1
|
||||
stage_type: diffusion
|
||||
runtime:
|
||||
devices: "1"
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
gpu_memory_utilization: 0.4
|
||||
gpu_memory_utilization: 0.55
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: image
|
||||
|
||||
@@ -39,7 +39,7 @@ stage_args:
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: false
|
||||
enforce_eager: true # haven't supported talker ACL graph on NPU
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
|
||||
@@ -44,7 +44,7 @@ stage_args:
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: false
|
||||
enforce_eager: true # haven't supported talker ACL graph on NPU
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent # Output codec codes for code2wav
|
||||
# tensor_parallel_size: 2
|
||||
|
||||
@@ -35,7 +35,7 @@ from vllm_ascend.compilation.acl_graph import update_full_graph_params
|
||||
from vllm_ascend.ops.rotary_embedding import update_cos_sin
|
||||
from vllm_ascend.utils import ProfileExecuteDuration
|
||||
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import KVCacheTransferData
|
||||
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
|
||||
|
||||
@@ -65,7 +65,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
|
||||
# each model stage has their own hidden size
|
||||
self.hidden_size = self.model_config.hf_text_config.hidden_size
|
||||
self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
|
||||
self.omni_connector = None
|
||||
# Initialize KV cache manager (preserve vllm_config fallback behavior)
|
||||
self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
|
||||
|
||||
def _make_buffer(self, *size, dtype, numpy=True):
|
||||
# Prevent ray from pinning the buffer due to large size
|
||||
@@ -92,7 +93,13 @@ class NPUARModelRunner(OmniNPUModelRunner):
|
||||
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
# [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
|
||||
self.kv_extracted_req_ids = self._handle_finished_requests_kv_transfer(scheduler_output)
|
||||
self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
|
||||
finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
|
||||
kv_caches=self.kv_caches,
|
||||
block_size=self.cache_config.block_size,
|
||||
cache_dtype=str(self.cache_config.cache_dtype),
|
||||
request_id_resolver=self._resolve_global_request_id,
|
||||
)
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
|
||||
with ProfileExecuteDuration().capture_async("prepare input"):
|
||||
@@ -499,161 +506,6 @@ class NPUARModelRunner(OmniNPUModelRunner):
|
||||
return hidden_states if self.pcp_size == 1 else self.pcp_manager.get_restore_hidden_states(
|
||||
hidden_states)
|
||||
|
||||
def _handle_finished_requests_kv_transfer(self, scheduler_output: SchedulerOutput) -> list[str]:
|
||||
"""Handle KV cache transfer for finished requests.
|
||||
|
||||
Returns list of request IDs that were processed (for Scheduler to free blocks).
|
||||
"""
|
||||
finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
|
||||
if not finished_reqs:
|
||||
return []
|
||||
|
||||
logger.debug(f"Processing KV transfer for {len(finished_reqs)} requests")
|
||||
|
||||
extracted_ids = []
|
||||
for req_id, data in finished_reqs.items():
|
||||
try:
|
||||
seq_len = data.get("seq_len", 0)
|
||||
block_ids = data.get("block_ids", [])
|
||||
if not block_ids:
|
||||
logger.warning(f"Request {req_id} has no block IDs, skipping")
|
||||
continue
|
||||
|
||||
# Extract KV cache from GPU blocks -> CPU tensors
|
||||
kv_data = self._extract_kv_cache(req_id, block_ids, seq_len)
|
||||
if kv_data:
|
||||
# Transfer to downstream stage via connector
|
||||
self._transfer_kv_cache(kv_data)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed KV transfer for {req_id}: {e}")
|
||||
finally:
|
||||
extracted_ids.append(req_id)
|
||||
|
||||
return extracted_ids
|
||||
|
||||
def _extract_kv_cache(self, req_id: str, block_ids: list[int], seq_len: int) -> KVCacheTransferData | None:
|
||||
"""Extract KV cache from GPU blocks for a single request."""
|
||||
num_layers = len(self.kv_caches)
|
||||
key_cache = [None] * num_layers
|
||||
value_cache = [None] * num_layers
|
||||
|
||||
for layer_idx, kv_tensor in enumerate(self.kv_caches):
|
||||
# Validate block IDs
|
||||
max_block = kv_tensor.shape[1] - 1
|
||||
valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block]
|
||||
if not valid_ids:
|
||||
continue
|
||||
|
||||
# Extract and reshape: [2, n_blocks, block_size, n_heads, head_dim]
|
||||
# -> [2, seq_len, n_heads, head_dim]
|
||||
selected = kv_tensor[:, valid_ids] # [2, n_valid, block_size, n_heads, head_dim]
|
||||
n_kv, n_blks, blk_sz, n_heads, d_head = selected.shape
|
||||
flat = selected.reshape(n_kv, n_blks * blk_sz, n_heads, d_head)
|
||||
if seq_len < flat.shape[1]:
|
||||
flat = flat[:, :seq_len]
|
||||
|
||||
# Move to CPU
|
||||
flat_cpu = flat.detach().cpu().contiguous()
|
||||
key_cache[layer_idx] = flat_cpu[0]
|
||||
value_cache[layer_idx] = flat_cpu[1]
|
||||
|
||||
if not any(k is not None for k in key_cache):
|
||||
return None
|
||||
|
||||
return KVCacheTransferData(
|
||||
request_id=req_id,
|
||||
layer_blocks={"key_cache": key_cache, "value_cache": value_cache},
|
||||
block_ids=block_ids,
|
||||
metadata={
|
||||
"block_size": self.cache_config.block_size,
|
||||
"num_layers": num_layers,
|
||||
"dtype": str(self.cache_config.cache_dtype),
|
||||
"seq_len": seq_len,
|
||||
},
|
||||
)
|
||||
|
||||
def _transfer_kv_cache(self, kv_data: KVCacheTransferData) -> None:
|
||||
"""Transfer KV cache data to downstream stage via OmniConnector."""
|
||||
connector = self._get_or_create_connector()
|
||||
if not connector:
|
||||
return
|
||||
|
||||
# Resolve global request ID if available
|
||||
transfer_req_id = self._resolve_global_request_id(kv_data.request_id)
|
||||
from_stage, to_stage = self._detect_transfer_stages()
|
||||
|
||||
# Prepare data and transfer with retry
|
||||
data_dict = kv_data.to_dict()
|
||||
data_dict["request_id"] = transfer_req_id
|
||||
|
||||
success, size, _ = self._transfer_with_retry(
|
||||
connector, from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict
|
||||
)
|
||||
|
||||
if success:
|
||||
logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes")
|
||||
else:
|
||||
logger.error(f"KV transfer FAILED: {transfer_req_id}")
|
||||
|
||||
def _get_or_create_connector(self) -> Any | None:
|
||||
"""Get existing connector or create one from config."""
|
||||
if self.omni_connector:
|
||||
return self.omni_connector
|
||||
|
||||
from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
|
||||
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
|
||||
|
||||
config = self._get_omni_connector_config()
|
||||
if not config or not isinstance(config, dict):
|
||||
logger.warning("No valid OmniConnector config found")
|
||||
return None
|
||||
|
||||
c_type = config.get("type")
|
||||
if not c_type:
|
||||
logger.error("OmniConnector config missing 'type' field")
|
||||
return None
|
||||
|
||||
c_extra = {k: v for k, v in config.items() if k != "type"}
|
||||
self.omni_connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra))
|
||||
return self.omni_connector
|
||||
|
||||
def _get_omni_connector_config(self) -> dict[str, Any] | None:
|
||||
"""Get OmniConnector configuration from model config."""
|
||||
# Primary: omni_kv_config from YAML
|
||||
omni_kv = getattr(self.model_config, "omni_kv_config", None)
|
||||
if isinstance(omni_kv, dict):
|
||||
cfg = omni_kv.get("connector_config")
|
||||
if isinstance(cfg, dict) and cfg:
|
||||
return cfg
|
||||
|
||||
# Fallback: kv_transfer_config
|
||||
kv_cfg = getattr(self.vllm_config, "kv_transfer_config", None)
|
||||
if kv_cfg:
|
||||
direct = getattr(kv_cfg, "omni_connector_config", None)
|
||||
if isinstance(direct, dict) and direct:
|
||||
return direct
|
||||
extra = getattr(kv_cfg, "kv_connector_extra_config", None)
|
||||
if isinstance(extra, dict):
|
||||
omni = extra.get("omni_connector_config")
|
||||
if isinstance(omni, dict) and omni:
|
||||
return omni
|
||||
|
||||
return None
|
||||
|
||||
def _detect_transfer_stages(self) -> tuple[str, str]:
|
||||
"""Detect source and target stages for KV transfer."""
|
||||
omni_kv = getattr(self.model_config, "omni_kv_config", None)
|
||||
if isinstance(omni_kv, dict):
|
||||
from_s = omni_kv.get("omni_from_stage")
|
||||
to_s = omni_kv.get("omni_to_stage")
|
||||
if from_s and to_s:
|
||||
return str(from_s), str(to_s)
|
||||
|
||||
raise ValueError(
|
||||
"KV transfer stages not configured. Please set 'omni_from_stage' and 'omni_to_stage' in omni_kv_config."
|
||||
)
|
||||
|
||||
def _resolve_global_request_id(self, req_id: str) -> str:
|
||||
"""Resolve global request ID from request state."""
|
||||
req_state = self.requests.get(req_id)
|
||||
@@ -669,32 +521,3 @@ class NPUARModelRunner(OmniNPUModelRunner):
|
||||
return global_id.decode("utf-8")
|
||||
return str(global_id)
|
||||
return req_id
|
||||
|
||||
def _transfer_with_retry(
|
||||
self,
|
||||
connector: Any,
|
||||
from_stage: str,
|
||||
to_stage: str,
|
||||
request_id: str,
|
||||
data: dict[str, Any],
|
||||
max_retries: int = 3,
|
||||
) -> tuple[bool, int, dict[str, Any] | None]:
|
||||
"""Transfer data with retry and exponential backoff."""
|
||||
import time
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
put_key = f"omni_{from_stage}_to_{to_stage}_{request_id}"
|
||||
success, size, metadata = connector.put(
|
||||
from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data
|
||||
)
|
||||
if success:
|
||||
return success, size, metadata
|
||||
logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Transfer attempt {attempt + 1} exception: {e}")
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(0.1 * (2**attempt))
|
||||
|
||||
return False, 0, None
|
||||
@@ -54,12 +54,13 @@ class OmniNPUModelRunner(NPUModelRunner):
|
||||
self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
|
||||
)
|
||||
hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size
|
||||
self.talker_mtp_input_ids = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
|
||||
max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
|
||||
self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
|
||||
self.talker_mtp_inputs_embeds = self._make_buffer(
|
||||
self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False
|
||||
max_batch_size, hidden_size, dtype=self.dtype, numpy=False
|
||||
)
|
||||
self.last_talker_hidden = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
|
||||
self.text_step = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
|
||||
self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
|
||||
self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
|
||||
|
||||
def _init_mrope_positions(self, req_state: CachedRequestState):
|
||||
image_grid_thw = []
|
||||
@@ -590,12 +591,16 @@ class OmniNPUModelRunner(NPUModelRunner):
|
||||
model_instance=self.model,
|
||||
):
|
||||
if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
|
||||
num_tokens_padded_talker_mtp = num_tokens_padded
|
||||
if num_tokens_padded_talker_mtp == self.max_num_tokens:
|
||||
num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
|
||||
hidden_states = self.talker_mtp(
|
||||
self.talker_mtp_input_ids.gpu[:num_tokens_padded],
|
||||
self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded],
|
||||
self.last_talker_hidden.gpu[:num_tokens_padded],
|
||||
self.text_step.gpu[:num_tokens_padded],
|
||||
self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
|
||||
self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
|
||||
self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp],
|
||||
self.text_step.gpu[:num_tokens_padded_talker_mtp],
|
||||
)
|
||||
self.compilation_config.cache_dir = None
|
||||
hidden_states = self._generate_dummy_run_hidden_states(
|
||||
input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds
|
||||
)
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
__version__ = "0.14.0rc1"
|
||||
__version_tuple__ = (0, 14, "0rc1")
|
||||
__version__ = "0.14.0"
|
||||
__version_tuple__ = (0, 14, 0)
|
||||
# TODO: add auto version generation
|
||||
Reference in new issue
Block a user