Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Roger Wang
|
86d9280274 | ||
|
Roger Wang
|
ce3e067854 | ||
|
|
1deb4fdf92 | ||
|
Lancer
|
9aae75c3ee | ||
|
汪志鹏
|
284575ab82 | ||
|
Lancer
|
2472c360aa |
No files matched your search
+114
-36
@@ -261,39 +261,117 @@ steps:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - |
|
||||
# timeout 30m bash -c '
|
||||
# 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: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
|
||||
- label: "Bagel Text2Img Model Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
|
||||
|
||||
- label: "Bagel Img2Img Model Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
|
||||
|
||||
- label: "Bagel Online Serving Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_bagel_online.py
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
|
||||
@@ -17,7 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
|
||||
*Latest News* 🔥
|
||||
|
||||
- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
|
||||
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
|
||||
- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
|
||||
- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
|
||||
- [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. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
|
||||
|
||||
@@ -100,6 +100,7 @@ nav:
|
||||
- api/vllm_omni
|
||||
- CLI Reference: cli
|
||||
- Community:
|
||||
- Governance: community/governance.md
|
||||
- community/*
|
||||
- Slack: https://slack.vllm.ai
|
||||
- Blog: https://blog.vllm.ai
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
# Governance
|
||||
|
||||
vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
|
||||
|
||||
## Values
|
||||
|
||||
vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
|
||||
|
||||
### Design Values
|
||||
|
||||
1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
|
||||
2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
|
||||
3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
|
||||
4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
|
||||
5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
|
||||
|
||||
### Collaboration Values
|
||||
|
||||
1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
|
||||
2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
|
||||
|
||||
## Project Maintainers
|
||||
|
||||
### Lead Maintainers
|
||||
|
||||
Lead maintainers are responsible for the overall direction and strategy of the project:
|
||||
|
||||
- [@Gaohan123](https://github.com/Gaohan123)
|
||||
- [@hsliuustc0106](https://github.com/hsliuustc0106)
|
||||
- [@ywang96](https://github.com/ywang96)
|
||||
|
||||
### Active Committers
|
||||
|
||||
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
|
||||
|
||||
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
|
||||
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
|
||||
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
|
||||
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
|
||||
- [@wtomin](https://github.com/wtomin):
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
|
||||
|
||||
## Meetings
|
||||
|
||||
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
|
||||
|
||||
## Committer Nomination Process
|
||||
|
||||
Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
|
||||
@@ -24,22 +24,22 @@ The following table shows which models are currently supported by parallelism me
|
||||
|
||||
### ImageGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
|
||||
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
|
||||
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ |
|
||||
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A |
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
|
||||
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A | ❌ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | ❌ |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ✅ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | ✅ |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ |
|
||||
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ |
|
||||
|
||||
!!! 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.
|
||||
|
||||
@@ -10,7 +10,7 @@ stage_args:
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: BagelForConditionalGeneration
|
||||
model_arch: OmniBagelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.45
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
End-to-end test for Bagel img2img generation.
|
||||
|
||||
This test validates that the Bagel model generates images from an input image
|
||||
and text prompt that match expected reference pixel values within a ±5 tolerance.
|
||||
|
||||
Equivalent to running:
|
||||
python3 examples/offline_inference/bagel/end2end.py \
|
||||
--prompts "Change the grass color to red" \
|
||||
--modality img2img --step 15 \
|
||||
--image-path 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
|
||||
"""
|
||||
|
||||
import socket
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from vllm.assets.image import ImageAsset
|
||||
|
||||
from tests.utils import hardware_test
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
# Reference pixel data extracted from the known-good output image
|
||||
# Generated with seed=52, num_inference_steps=15,
|
||||
# prompt='Change the grass color to red',
|
||||
# input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
|
||||
REFERENCE_PIXELS = [
|
||||
{"position": (100, 100), "rgb": (157, 172, 217)},
|
||||
{"position": (400, 50), "rgb": (105, 144, 218)},
|
||||
{"position": (700, 100), "rgb": (118, 159, 233)},
|
||||
{"position": (150, 400), "rgb": (195, 34, 60)},
|
||||
{"position": (512, 336), "rgb": (222, 214, 193)},
|
||||
{"position": (700, 400), "rgb": (197, 15, 43)},
|
||||
{"position": (100, 600), "rgb": (105, 13, 18)},
|
||||
{"position": (400, 600), "rgb": (169, 33, 44)},
|
||||
{"position": (700, 600), "rgb": (101, 86, 93)},
|
||||
{"position": (256, 256), "rgb": (181, 202, 222)},
|
||||
]
|
||||
|
||||
PIXEL_TOLERANCE = 5
|
||||
|
||||
DEFAULT_PROMPT = "<|fim_middle|><|im_start|>Change the grass color to red<|im_end|>"
|
||||
|
||||
EXPECTED_OUTPUT_SIZE = (1024, 672)
|
||||
|
||||
|
||||
def _load_input_image() -> Image.Image:
|
||||
"""Load the test input image via vllm's ImageAsset."""
|
||||
return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
|
||||
|
||||
|
||||
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 img2img 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
|
||||
params_list[1].extra_args = { # type: ignore
|
||||
"cfg_text_scale": 4.0,
|
||||
"cfg_img_scale": 1.5,
|
||||
}
|
||||
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_img2img(
|
||||
omni: Omni,
|
||||
input_image: Image.Image,
|
||||
prompt: str = DEFAULT_PROMPT,
|
||||
) -> Image.Image:
|
||||
"""Generate an image using Bagel model with img2img pipeline.
|
||||
|
||||
Args:
|
||||
omni: The Omni instance to use for generation.
|
||||
input_image: The input PIL Image for img2img.
|
||||
prompt: The text prompt for image editing.
|
||||
|
||||
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,
|
||||
"multi_modal_data": {"img2img": input_image},
|
||||
"modalities": ["img2img"],
|
||||
}
|
||||
],
|
||||
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 == EXPECTED_OUTPUT_SIZE, f"Expected {EXPECTED_OUTPUT_SIZE}, got {generated_image.size}"
|
||||
|
||||
return generated_image
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_bagel_img2img_shared_memory_connector():
|
||||
"""Test Bagel img2img with shared memory connector."""
|
||||
input_image = _load_input_image()
|
||||
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_img2img(omni, input_image)
|
||||
_validate_pixels(generated_image)
|
||||
finally:
|
||||
omni.close()
|
||||
@@ -14,10 +14,6 @@ Equivalent to running:
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
@@ -37,23 +33,23 @@ from vllm_omni.entrypoints.omni import Omni
|
||||
# "Generated with seed=52, num_inference_steps=15,
|
||||
# prompt='A futuristic city skyline at twilight, cyberpunk style'"
|
||||
REFERENCE_PIXELS = [
|
||||
{"position": (100, 100), "rgb": (49, 96, 134)},
|
||||
{"position": (400, 50), "rgb": (63, 127, 167)},
|
||||
{"position": (700, 100), "rgb": (70, 101, 141)},
|
||||
{"position": (150, 400), "rgb": (115, 90, 150)},
|
||||
{"position": (512, 512), "rgb": (98, 86, 119)},
|
||||
{"position": (700, 400), "rgb": (29, 42, 91)},
|
||||
{"position": (100, 700), "rgb": (47, 50, 88)},
|
||||
{"position": (400, 700), "rgb": (36, 52, 91)},
|
||||
{"position": (700, 700), "rgb": (45, 58, 99)},
|
||||
{"position": (256, 256), "rgb": (62, 94, 135)},
|
||||
{"position": (100, 100), "rgb": (121, 118, 100)},
|
||||
{"position": (400, 50), "rgb": (163, 162, 143)},
|
||||
{"position": (700, 100), "rgb": (170, 156, 127)},
|
||||
{"position": (150, 400), "rgb": (129, 127, 112)},
|
||||
{"position": (512, 512), "rgb": (135, 61, 59)},
|
||||
{"position": (700, 400), "rgb": (205, 107, 43)},
|
||||
{"position": (100, 700), "rgb": (197, 177, 157)},
|
||||
{"position": (400, 700), "rgb": (139, 107, 86)},
|
||||
{"position": (700, 700), "rgb": (247, 205, 146)},
|
||||
{"position": (256, 256), "rgb": (171, 160, 153)},
|
||||
]
|
||||
|
||||
# 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|>"
|
||||
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
End-to-end online serving test for Bagel text2img and img2img generation.
|
||||
|
||||
This test validates that the Bagel model can serve image generation requests
|
||||
via the OpenAI-compatible chat completions API.
|
||||
|
||||
Equivalent to running:
|
||||
vllm-omni serve "ByteDance-Seed/BAGEL-7B-MoT" --omni --port 8091
|
||||
|
||||
# text2img
|
||||
python3 examples/online_serving/bagel/openai_chat_client.py \\
|
||||
--prompt "A cute cat" --modality text2img
|
||||
|
||||
# img2img
|
||||
python3 examples/online_serving/bagel/openai_chat_client.py \\
|
||||
--prompt "Let the woman wear a blue dress" --modality img2img \\
|
||||
--image-url women.jpg
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from PIL import Image
|
||||
from vllm.assets.image import ImageAsset
|
||||
|
||||
from tests.utils import hardware_test
|
||||
|
||||
MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
|
||||
STAGE_CONFIGS_PATH = str(
|
||||
Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml"
|
||||
)
|
||||
|
||||
TEXT2IMG_PROMPT = "A cute cat"
|
||||
IMG2IMG_PROMPT = "Change the grass color to red"
|
||||
|
||||
|
||||
class BagelOmniServer:
|
||||
"""Context manager to start/stop a vLLM-Omni server for Bagel model tests."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = MODEL,
|
||||
stage_configs_path: str = STAGE_CONFIGS_PATH,
|
||||
env_dict: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self.model = model
|
||||
self.stage_configs_path = stage_configs_path
|
||||
self.env_dict = env_dict
|
||||
self.proc: subprocess.Popen | None = None
|
||||
self.host = "127.0.0.1"
|
||||
self.port = _find_free_port()
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return f"http://{self.host}:{self.port}"
|
||||
|
||||
def _start_server(self) -> None:
|
||||
env = os.environ.copy()
|
||||
if self.env_dict is not None:
|
||||
env.update(self.env_dict)
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"vllm_omni.entrypoints.cli.main",
|
||||
"serve",
|
||||
self.model,
|
||||
"--omni",
|
||||
"--host",
|
||||
self.host,
|
||||
"--port",
|
||||
str(self.port),
|
||||
"--stage-configs-path",
|
||||
self.stage_configs_path,
|
||||
"--stage-init-timeout",
|
||||
"300",
|
||||
]
|
||||
|
||||
self.proc = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
try:
|
||||
if not _wait_for_port(self.host, self.port, timeout=600, proc=self.proc):
|
||||
self.terminate()
|
||||
raise RuntimeError(f"Server failed to start within 600 seconds on {self.host}:{self.port}")
|
||||
except Exception:
|
||||
self.terminate()
|
||||
raise
|
||||
|
||||
def __enter__(self):
|
||||
self._start_server()
|
||||
return self
|
||||
|
||||
def terminate(self) -> None:
|
||||
if self.proc:
|
||||
try:
|
||||
os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
try:
|
||||
self.proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
try:
|
||||
os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
self.proc.wait()
|
||||
self.proc = None
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.terminate()
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
s.listen(1)
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def _wait_for_port(host: str, port: int, timeout: int = 600, proc: subprocess.Popen | None = None) -> bool:
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
if proc is not None and proc.poll() is not None:
|
||||
# Server process exited early
|
||||
return False
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
sock.settimeout(1)
|
||||
if sock.connect_ex((host, port)) == 0:
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(2)
|
||||
return False
|
||||
|
||||
|
||||
def _send_chat_request(
|
||||
server_url: str,
|
||||
prompt: str,
|
||||
*,
|
||||
modality: str = "text2img",
|
||||
image: Image.Image | None = None,
|
||||
timeout: int = 300,
|
||||
) -> dict[str, Any]:
|
||||
"""Send a chat completion request matching the openai_chat_client.py format."""
|
||||
content: list[dict[str, Any]] = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
|
||||
|
||||
if image is not None:
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format="JPEG")
|
||||
b64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/jpeg;base64,{b64_data}"},
|
||||
}
|
||||
)
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
}
|
||||
|
||||
if modality in ("text2img", "img2img"):
|
||||
payload["modalities"] = ["image"]
|
||||
|
||||
resp = requests.post(
|
||||
f"{server_url}/v1/chat/completions",
|
||||
headers={"Content-Type": "application/json"},
|
||||
json=payload,
|
||||
timeout=timeout,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _extract_image_from_response(data: dict[str, Any]) -> Image.Image | None:
|
||||
"""Extract the generated PIL Image from a chat completion response."""
|
||||
for choice in data.get("choices", []):
|
||||
content = choice.get("message", {}).get("content")
|
||||
if isinstance(content, list) and content:
|
||||
first_item = content[0]
|
||||
if isinstance(first_item, dict) and "image_url" in first_item:
|
||||
url = first_item["image_url"].get("url", "")
|
||||
if url.startswith("data:image"):
|
||||
_, b64 = url.split(",", 1)
|
||||
return Image.open(BytesIO(base64.b64decode(b64)))
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_bagel_text2img_online():
|
||||
"""Test Bagel text2img via OpenAI-compatible chat completions API."""
|
||||
with BagelOmniServer() as server:
|
||||
response_data = _send_chat_request(
|
||||
server.base_url,
|
||||
TEXT2IMG_PROMPT,
|
||||
modality="text2img",
|
||||
)
|
||||
|
||||
image = _extract_image_from_response(response_data)
|
||||
assert image is not None, f"No image in response: {response_data}"
|
||||
image.load()
|
||||
|
||||
w, h = image.size
|
||||
assert w > 0 and h > 0, f"Invalid image size: {image.size}"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_bagel_img2img_online():
|
||||
"""Test Bagel img2img via OpenAI-compatible chat completions API."""
|
||||
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
|
||||
|
||||
with BagelOmniServer() as server:
|
||||
response_data = _send_chat_request(
|
||||
server.base_url,
|
||||
IMG2IMG_PROMPT,
|
||||
modality="img2img",
|
||||
image=input_image,
|
||||
)
|
||||
|
||||
image = _extract_image_from_response(response_data)
|
||||
assert image is not None, f"No image in response: {response_data}"
|
||||
image.load()
|
||||
|
||||
w, h = image.size
|
||||
assert w > 0 and h > 0, f"Invalid image size: {image.size}"
|
||||
@@ -157,3 +157,29 @@ class RotaryEmbedding(CustomOp):
|
||||
sin,
|
||||
interleaved=self.interleaved,
|
||||
)
|
||||
|
||||
|
||||
def apply_rope_to_qk(
|
||||
rope: RotaryEmbedding,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Apply rotary positional embeddings to query and key tensors.
|
||||
|
||||
Args:
|
||||
rope: RotaryEmbedding instance for applying position embeddings
|
||||
query: Query tensor [B, S, H, D]
|
||||
key: Key tensor [B, S, H, D]
|
||||
image_rotary_emb: Tuple of (cos, sin) tensors or None
|
||||
|
||||
Returns:
|
||||
Tuple of (query, key) with RoPE applied if rotary embeddings provided
|
||||
"""
|
||||
if image_rotary_emb is not None:
|
||||
cos, sin = image_rotary_emb
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
query = rope(query, cos, sin)
|
||||
key = rope(key, cos, sin)
|
||||
return query, key
|
||||
@@ -30,7 +30,7 @@ if TYPE_CHECKING:
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -224,12 +224,7 @@ class FluxAttention(torch.nn.Module):
|
||||
key = torch.cat([encoder_key, key], dim=1)
|
||||
value = torch.cat([encoder_value, value], dim=1)
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
cos, sin = image_rotary_emb # [S, D/2]
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
query = self.rope(query, cos, sin)
|
||||
key = self.rope(key, cos, sin)
|
||||
query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb) # [S, D/2]
|
||||
|
||||
hidden_states = self.attn(
|
||||
query,
|
||||
@@ -475,6 +470,12 @@ class FluxTransformer2DModel(nn.Module):
|
||||
# -- typically a transformer layer
|
||||
# used for torch compile optimizations
|
||||
_repeated_blocks = ["FluxTransformerBlock"]
|
||||
|
||||
@staticmethod
|
||||
def _is_transformer_block(name: str, module) -> bool:
|
||||
return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_transformer_block]
|
||||
packed_modules_mapping = {
|
||||
"to_qkv": ["to_q", "to_k", "to_v"],
|
||||
"add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],
|
||||
|
||||
@@ -160,10 +160,10 @@ class FluxPipeline(nn.Module, CFGParallelMixin):
|
||||
)
|
||||
self.text_encoder = CLIPTextModel.from_pretrained(
|
||||
model, subfolder="text_encoder", local_files_only=local_files_only
|
||||
)
|
||||
).to(self.device)
|
||||
self.text_encoder_2 = T5EncoderModel.from_pretrained(
|
||||
model, subfolder="text_encoder_2", local_files_only=local_files_only
|
||||
)
|
||||
).to(self.device)
|
||||
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
|
||||
self.device
|
||||
)
|
||||
|
||||
@@ -553,6 +553,12 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
|
||||
_repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
|
||||
|
||||
@staticmethod
|
||||
def _is_transformer_block(name: str, module) -> bool:
|
||||
return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_transformer_block]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: int = 1,
|
||||
|
||||
@@ -366,7 +366,7 @@ class Flux2Pipeline(nn.Module, SupportImageInput):
|
||||
)
|
||||
self.text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
|
||||
model, subfolder="text_encoder", local_files_only=local_files_only
|
||||
)
|
||||
).to(self._execution_device)
|
||||
self.tokenizer = PixtralProcessor.from_pretrained(
|
||||
model, subfolder="tokenizer", local_files_only=local_files_only
|
||||
)
|
||||
|
||||
@@ -28,6 +28,7 @@ from diffusers.models.embeddings import (
|
||||
)
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.normalization import AdaLayerNormContinuous
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
@@ -39,8 +40,15 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.sp_plan import (
|
||||
SequenceParallelInput,
|
||||
SequenceParallelOutput,
|
||||
)
|
||||
from vllm_omni.diffusion.forward_context import get_forward_context
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
|
||||
|
||||
logger = init_logger(__name__)
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
@@ -98,6 +106,7 @@ class Flux2FeedForward(nn.Module):
|
||||
class Flux2Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
parallel_config: DiffusionParallelConfig,
|
||||
query_dim: int,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
@@ -112,6 +121,7 @@ class Flux2Attention(nn.Module):
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.parallel_config = parallel_config
|
||||
self.head_dim = dim_head
|
||||
self.inner_dim = out_dim if out_dim is not None else dim_head * heads
|
||||
self.query_dim = query_dim
|
||||
@@ -208,33 +218,87 @@ class Flux2Attention(nn.Module):
|
||||
encoder_query = self.norm_added_q(encoder_query)
|
||||
encoder_key = self.norm_added_k(encoder_key)
|
||||
|
||||
query = torch.cat([encoder_query, query], dim=1)
|
||||
key = torch.cat([encoder_key, key], dim=1)
|
||||
value = torch.cat([encoder_value, value], dim=1)
|
||||
sp_size = self.parallel_config.sequence_parallel_size
|
||||
forward_ctx = get_forward_context()
|
||||
use_sp_joint_attention = sp_size is not None and sp_size > 1 and not forward_ctx.split_text_embed_in_sp
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
cos, sin = image_rotary_emb
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
query = self.rope(query, cos, sin)
|
||||
key = self.rope(key, cos, sin)
|
||||
if use_sp_joint_attention and image_rotary_emb is not None:
|
||||
cos, sin = image_rotary_emb
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
txt_len = encoder_query.shape[1]
|
||||
txt_cos, img_cos = cos[:txt_len], cos[txt_len:]
|
||||
txt_sin, img_sin = sin[:txt_len], sin[txt_len:]
|
||||
query = self.rope(query, img_cos, img_sin)
|
||||
key = self.rope(key, img_cos, img_sin)
|
||||
encoder_query = self.rope(encoder_query, txt_cos, txt_sin)
|
||||
encoder_key = self.rope(encoder_key, txt_cos, txt_sin)
|
||||
|
||||
attn_metadata = None
|
||||
if attention_mask is not None:
|
||||
if attention_mask.dim() == 3:
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
attn_metadata = AttentionMetadata(attn_mask=attention_mask)
|
||||
attn_metadata = AttentionMetadata(
|
||||
joint_query=encoder_query,
|
||||
joint_key=encoder_key,
|
||||
joint_value=encoder_value,
|
||||
joint_strategy="front",
|
||||
)
|
||||
hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None)
|
||||
encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None)
|
||||
if hidden_states_mask is not None:
|
||||
attn_metadata.attn_mask = hidden_states_mask
|
||||
if encoder_hidden_states_mask is not None:
|
||||
attn_metadata.joint_attn_mask = encoder_hidden_states_mask
|
||||
|
||||
hidden_states = self.attn(query, key, value, attn_metadata)
|
||||
hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
|
||||
hidden_states = self.attn(query, key, value, attn_metadata)
|
||||
hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
|
||||
|
||||
if encoder_hidden_states is not None:
|
||||
context_len = encoder_hidden_states.shape[1]
|
||||
encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
|
||||
[context_len, hidden_states.shape[1] - context_len],
|
||||
dim=1,
|
||||
)
|
||||
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
|
||||
txt_len = encoder_hidden_states.shape[1]
|
||||
encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
|
||||
[txt_len, hidden_states.shape[1] - txt_len],
|
||||
dim=1,
|
||||
)
|
||||
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
|
||||
else:
|
||||
query = torch.cat([encoder_query, query], dim=1)
|
||||
key = torch.cat([encoder_key, key], dim=1)
|
||||
value = torch.cat([encoder_value, value], dim=1)
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
cos, sin = image_rotary_emb
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
query = self.rope(query, cos, sin)
|
||||
key = self.rope(key, cos, sin)
|
||||
|
||||
attn_metadata = None
|
||||
if attention_mask is not None:
|
||||
if attention_mask.dim() == 3:
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
attn_metadata = AttentionMetadata(attn_mask=attention_mask)
|
||||
|
||||
hidden_states = self.attn(query, key, value, attn_metadata)
|
||||
hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
|
||||
|
||||
context_len = encoder_hidden_states.shape[1]
|
||||
encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
|
||||
[context_len, hidden_states.shape[1] - context_len],
|
||||
dim=1,
|
||||
)
|
||||
encoder_hidden_states = self.to_add_out(encoder_hidden_states)
|
||||
else:
|
||||
if image_rotary_emb is not None:
|
||||
cos, sin = image_rotary_emb
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
query = self.rope(query, cos, sin)
|
||||
key = self.rope(key, cos, sin)
|
||||
|
||||
attn_metadata = None
|
||||
if attention_mask is not None:
|
||||
if attention_mask.dim() == 3:
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
attn_metadata = AttentionMetadata(attn_mask=attention_mask)
|
||||
|
||||
hidden_states = self.attn(query, key, value, attn_metadata)
|
||||
hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
|
||||
|
||||
hidden_states = self.to_out[0](hidden_states)
|
||||
hidden_states = self.to_out[1](hidden_states)
|
||||
@@ -251,6 +315,7 @@ class Flux2ParallelSelfAttention(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
parallel_config: DiffusionParallelConfig,
|
||||
query_dim: int,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
@@ -265,6 +330,7 @@ class Flux2ParallelSelfAttention(nn.Module):
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.parallel_config = parallel_config
|
||||
self.head_dim = dim_head
|
||||
self.inner_dim = out_dim if out_dim is not None else dim_head * heads
|
||||
self.query_dim = query_dim
|
||||
@@ -325,20 +391,57 @@ class Flux2ParallelSelfAttention(nn.Module):
|
||||
query = self.norm_q(query)
|
||||
key = self.norm_k(key)
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
sp_size = self.parallel_config.sequence_parallel_size
|
||||
forward_ctx = get_forward_context()
|
||||
text_seq_len = kwargs.get("text_seq_len", None)
|
||||
use_sp_single_stream = (
|
||||
sp_size is not None and sp_size > 1 and not forward_ctx.split_text_embed_in_sp and text_seq_len is not None
|
||||
)
|
||||
|
||||
if use_sp_single_stream and image_rotary_emb is not None:
|
||||
cos, sin = image_rotary_emb
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
query = self.rope(query, cos, sin)
|
||||
key = self.rope(key, cos, sin)
|
||||
txt_cos, img_cos = cos[:text_seq_len], cos[text_seq_len:]
|
||||
txt_sin, img_sin = sin[:text_seq_len], sin[text_seq_len:]
|
||||
|
||||
attn_metadata = None
|
||||
if attention_mask is not None:
|
||||
if attention_mask.dim() == 3:
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
attn_metadata = AttentionMetadata(attn_mask=attention_mask)
|
||||
img_query = query[:, text_seq_len:]
|
||||
img_key = key[:, text_seq_len:]
|
||||
img_value = value[:, text_seq_len:]
|
||||
text_query = query[:, :text_seq_len]
|
||||
text_key = key[:, :text_seq_len]
|
||||
text_value = value[:, :text_seq_len]
|
||||
|
||||
img_query = self.rope(img_query, img_cos, img_sin)
|
||||
img_key = self.rope(img_key, img_cos, img_sin)
|
||||
text_query = self.rope(text_query, txt_cos, txt_sin)
|
||||
text_key = self.rope(text_key, txt_cos, txt_sin)
|
||||
|
||||
attn_metadata = AttentionMetadata(
|
||||
joint_query=text_query,
|
||||
joint_key=text_key,
|
||||
joint_value=text_value,
|
||||
joint_strategy="front",
|
||||
)
|
||||
hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None)
|
||||
encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None)
|
||||
if hidden_states_mask is not None:
|
||||
attn_metadata.attn_mask = hidden_states_mask
|
||||
if encoder_hidden_states_mask is not None:
|
||||
attn_metadata.joint_attn_mask = encoder_hidden_states_mask
|
||||
|
||||
attn_output = self.attn(img_query, img_key, img_value, attn_metadata)
|
||||
else:
|
||||
query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb)
|
||||
|
||||
attn_metadata = None
|
||||
if attention_mask is not None:
|
||||
if attention_mask.dim() == 3:
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
attn_metadata = AttentionMetadata(attn_mask=attention_mask)
|
||||
|
||||
attn_output = self.attn(query, key, value, attn_metadata)
|
||||
|
||||
attn_output = self.attn(query, key, value, attn_metadata)
|
||||
attn_output = attn_output.flatten(2, 3).to(query.dtype)
|
||||
|
||||
mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)
|
||||
@@ -350,6 +453,7 @@ class Flux2ParallelSelfAttention(nn.Module):
|
||||
class Flux2SingleTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
parallel_config: DiffusionParallelConfig,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
@@ -361,6 +465,7 @@ class Flux2SingleTransformerBlock(nn.Module):
|
||||
super().__init__()
|
||||
self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
||||
self.attn = Flux2ParallelSelfAttention(
|
||||
parallel_config=parallel_config,
|
||||
query_dim=dim,
|
||||
dim_head=attention_head_dim,
|
||||
heads=num_attention_heads,
|
||||
@@ -383,6 +488,13 @@ class Flux2SingleTransformerBlock(nn.Module):
|
||||
split_hidden_states: bool = False,
|
||||
text_seq_len: int | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Forward pass for Flux2SingleTransformerBlock with SP support.
|
||||
|
||||
In SP mode: image hidden_states is chunked (B, img_len/SP, D),
|
||||
text encoder_hidden_states is full (B, txt_len, D).
|
||||
The block concatenates them for joint attention.
|
||||
"""
|
||||
if encoder_hidden_states is not None:
|
||||
text_seq_len = encoder_hidden_states.shape[1]
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
|
||||
@@ -412,6 +524,7 @@ class Flux2SingleTransformerBlock(nn.Module):
|
||||
class Flux2TransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
parallel_config: DiffusionParallelConfig,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
@@ -425,6 +538,7 @@ class Flux2TransformerBlock(nn.Module):
|
||||
self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
|
||||
|
||||
self.attn = Flux2Attention(
|
||||
parallel_config=parallel_config,
|
||||
query_dim=dim,
|
||||
added_kv_proj_dim=dim,
|
||||
dim_head=attention_head_dim,
|
||||
@@ -525,6 +639,47 @@ class Flux2PosEmbed(nn.Module):
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
class Flux2RopePrepare(nn.Module):
|
||||
"""Prepares RoPE embeddings for sequence parallel.
|
||||
|
||||
This module encapsulates the RoPE computation for Flux.2-klein.
|
||||
For dual-stream attention, text components (outputs 0, 1) are replicated
|
||||
across SP ranks, while image components (outputs 2, 3) are sharded.
|
||||
|
||||
NOTE: The hidden_states projection is handled separately in forward()
|
||||
so that _sp_plan can shard it at the root level.
|
||||
"""
|
||||
|
||||
def __init__(self, pos_embed: Flux2PosEmbed):
|
||||
super().__init__()
|
||||
self.pos_embed = pos_embed
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img_ids: torch.Tensor,
|
||||
txt_ids: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Compute RoPE embeddings for text and image sequences.
|
||||
|
||||
Args:
|
||||
img_ids: Image position IDs (img_seq_len, n_axes)
|
||||
txt_ids: Text position IDs (txt_seq_len, n_axes)
|
||||
|
||||
Returns:
|
||||
Tuple of cosine / sine components for text & image
|
||||
in the order: (txt_cos, txt_sin, img_cos, img_sin)
|
||||
|
||||
NOTE: careful about output orders if this is refactored in the
|
||||
future; we need to match the _sp_plan indices, since text
|
||||
components (0 & 1) need to be replicated across SP ranks,
|
||||
while image components (2 & 3) must be sharded.
|
||||
"""
|
||||
img_freqs_cos, img_freqs_sin = self.pos_embed(img_ids)
|
||||
txt_freqs_cos, txt_freqs_sin = self.pos_embed(txt_ids)
|
||||
return txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin
|
||||
|
||||
|
||||
class Flux2TimestepGuidanceEmbeddings(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -580,10 +735,29 @@ class Flux2Modulation(nn.Module):
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
"""
|
||||
The Transformer model introduced in Flux 2.
|
||||
|
||||
Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
|
||||
"""
|
||||
|
||||
_repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
|
||||
|
||||
@staticmethod
|
||||
def _is_transformer_block(name: str, module) -> bool:
|
||||
return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_transformer_block]
|
||||
_sp_plan = {
|
||||
"": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True),
|
||||
},
|
||||
"rope_prepare": {
|
||||
2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True),
|
||||
3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True),
|
||||
},
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
"""SP plan: shard hidden_states at root level, shard img_freqs at rope_prepare, gather output at proj_out."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: int = 1,
|
||||
@@ -600,6 +774,7 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
rope_theta: int = 2000,
|
||||
eps: float = 1e-6,
|
||||
guidance_embeds: bool = True,
|
||||
od_config: OmniDiffusionConfig = None,
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
):
|
||||
super().__init__()
|
||||
@@ -622,6 +797,13 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
guidance_embeds=guidance_embeds,
|
||||
)
|
||||
|
||||
if od_config is not None:
|
||||
self.parallel_config = od_config.parallel_config
|
||||
else:
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
self.parallel_config = DiffusionParallelConfig()
|
||||
|
||||
self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope))
|
||||
self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
|
||||
in_channels=timestep_guidance_channels,
|
||||
@@ -637,9 +819,12 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False)
|
||||
self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
|
||||
|
||||
self.rope_prepare = Flux2RopePrepare(self.pos_embed)
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
Flux2TransformerBlock(
|
||||
parallel_config=self.parallel_config,
|
||||
dim=self.inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
@@ -655,6 +840,7 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
self.single_transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
Flux2SingleTransformerBlock(
|
||||
parallel_config=self.parallel_config,
|
||||
dim=self.inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
@@ -695,6 +881,10 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
|
||||
num_txt_tokens = encoder_hidden_states.shape[1]
|
||||
|
||||
sp_size = self.parallel_config.sequence_parallel_size
|
||||
if sp_size is not None and sp_size > 1:
|
||||
get_forward_context().split_text_embed_in_sp = False
|
||||
|
||||
timestep = timestep.to(hidden_states.dtype) * 1000
|
||||
if guidance is not None:
|
||||
guidance = guidance.to(hidden_states.dtype) * 1000
|
||||
@@ -705,21 +895,44 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
double_stream_mod_txt = self.double_stream_modulation_txt(temb)
|
||||
single_stream_mod = self.single_stream_modulation(temb)[0]
|
||||
|
||||
hidden_states = self.x_embedder(hidden_states)
|
||||
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
|
||||
|
||||
if img_ids.ndim == 3:
|
||||
img_ids = img_ids[0]
|
||||
if txt_ids.ndim == 3:
|
||||
txt_ids = txt_ids[0]
|
||||
|
||||
image_rotary_emb = self.pos_embed(img_ids)
|
||||
text_rotary_emb = self.pos_embed(txt_ids)
|
||||
hidden_states = self.x_embedder(hidden_states)
|
||||
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
|
||||
|
||||
txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin = self.rope_prepare(img_ids, txt_ids)
|
||||
|
||||
concat_rotary_emb = (
|
||||
torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
|
||||
torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
|
||||
torch.cat([txt_freqs_cos, img_freqs_cos], dim=0),
|
||||
torch.cat([txt_freqs_sin, img_freqs_sin], dim=0),
|
||||
)
|
||||
|
||||
# Create separate masks for image and text portions for Ulysses SP joint attention
|
||||
hidden_states_mask = None
|
||||
encoder_hidden_states_mask = None
|
||||
ctx = get_forward_context()
|
||||
if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
|
||||
batch_size = hidden_states.shape[0]
|
||||
img_padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
|
||||
|
||||
hidden_states_mask = torch.ones(
|
||||
batch_size,
|
||||
img_padded_seq_len,
|
||||
dtype=torch.bool,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
hidden_states_mask[:, ctx.sp_original_seq_len :] = False
|
||||
if hidden_states_mask.all():
|
||||
hidden_states_mask = None
|
||||
|
||||
if hidden_states_mask is not None:
|
||||
joint_attention_kwargs["hidden_states_mask"] = hidden_states_mask
|
||||
if encoder_hidden_states_mask is not None:
|
||||
joint_attention_kwargs["encoder_hidden_states_mask"] = encoder_hidden_states_mask
|
||||
|
||||
for block in self.transformer_blocks:
|
||||
encoder_hidden_states, hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
@@ -739,6 +952,7 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
temb_mod_params=single_stream_mod,
|
||||
image_rotary_emb=concat_rotary_emb,
|
||||
joint_attention_kwargs=joint_attention_kwargs,
|
||||
text_seq_len=num_txt_tokens,
|
||||
)
|
||||
|
||||
hidden_states = hidden_states[:, num_txt_tokens:, ...]
|
||||
|
||||
@@ -218,7 +218,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput):
|
||||
model,
|
||||
subfolder="text_encoder",
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
).to(self._execution_device)
|
||||
self.tokenizer = Qwen2TokenizerFast.from_pretrained(
|
||||
model,
|
||||
subfolder="tokenizer",
|
||||
|
||||
@@ -49,7 +49,7 @@ from vllm_omni.diffusion.forward_context import (
|
||||
get_forward_context,
|
||||
is_forward_context_available,
|
||||
)
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
|
||||
|
||||
ADALN_EMBED_DIM = 256
|
||||
SEQ_MULTI_OF = 32
|
||||
@@ -329,10 +329,7 @@ class ZImageAttention(nn.Module):
|
||||
query = self.norm_q(query)
|
||||
key = self.norm_k(key)
|
||||
|
||||
cos = cos.to(query.dtype)
|
||||
sin = sin.to(query.dtype)
|
||||
query = self.rope(query, cos, sin)
|
||||
key = self.rope(key, cos, sin)
|
||||
query, key = apply_rope_to_qk(self.rope, query, key, (cos, sin))
|
||||
# Cast to correct dtype
|
||||
dtype = query.dtype
|
||||
query, key = query.to(dtype), key.to(dtype)
|
||||
|
||||
Reference in new issue
Block a user