[Model] Support Boogu/Boogu-Image-0.1-Base and Boogu/Boogu-Image-0.1-Edit (#4995)
Signed-off-by: zzehli <jaeli_ottawa@outlook.com> Signed-off-by: jl <jaeli_ottawa@outlook.com> Signed-off-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
JaeandCursor authored and GitHub committed
1 parent
d34b91722f
commit
ee33954dff
26 files changed
+4973
No files matched your search
@@ -649,6 +649,10 @@ steps:
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit.py
|
||||
tests/e2e/online_serving/test_boogu_image_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
tests/e2e/online_serving/test_krea2_expansion.py
|
||||
tests/e2e/offline_inference/test_krea2_expansion.py
|
||||
|
||||
@@ -54,6 +54,7 @@ th {
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | |
|
||||
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
# Boogu-Image
|
||||
|
||||
> Text-to-image and image-editing online serving (Boogu-Image-0.1-Base / -Edit)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Boogu
|
||||
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing)
|
||||
- Task: Text-to-image generation and text-guided image editing (TI2I)
|
||||
- Mode: Online serving with the OpenAI-compatible API
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a known-good starting point for serving
|
||||
`Boogu/Boogu-Image-0.1-Base` with vLLM-Omni's native pipeline (no
|
||||
`--diffusion-load-format diffusers`, and the upstream `boogu` package is not
|
||||
required).
|
||||
|
||||
Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
|
||||
family. This recipe covers the Base text-to-image checkpoint, which pairs a
|
||||
Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match
|
||||
Euler scheduler with time-shift. It handles photorealistic generation and
|
||||
Chinese/English text rendering.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Base>
|
||||
- Project page: <https://boogu.org>
|
||||
- GitHub: <https://github.com/boogu-project/Boogu-Image>
|
||||
- Related example: [`examples/online_serving/text_to_image/`](../../examples/online_serving/text_to_image/README.md)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents tested configurations for CUDA GPU serving. The native
|
||||
pipeline runs single-GPU; multi-GPU parallelism, CPU offload, and cache
|
||||
acceleration are not yet supported for this model (see Notes).
|
||||
|
||||
## GPU
|
||||
|
||||
### 1 x A100/H100 (Single GPU, 40GB+ VRAM)
|
||||
|
||||
The model footprint is roughly 34.6 GiB on GPU, so a 40GB+ card is recommended.
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Base --omni --port 8091
|
||||
```
|
||||
|
||||
!!! note
|
||||
If you hit Out-of-Memory (OOM) on a smaller card, enable VAE slicing and
|
||||
tiling to reduce peak memory: `--vae-use-slicing --vae-use-tiling`.
|
||||
|
||||
#### Verification
|
||||
|
||||
After the server is ready, test with a simple request:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Boogu/Boogu-Image-0.1-Base",
|
||||
"prompt": "A mountain lake at sunset, photorealistic, cinematic lighting",
|
||||
"size": "1024x1024",
|
||||
"num_inference_steps": 28,
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > output.png
|
||||
```
|
||||
|
||||
Or via the chat-completions endpoint (parameters go in `extra_body`):
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A mountain lake at sunset, photorealistic, cinematic lighting"}
|
||||
],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 28,
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Memory usage:** ~34.6 GiB on GPU. Use `--vae-use-slicing --vae-use-tiling`
|
||||
to trim peak VRAM if needed.
|
||||
- **Key flags:**
|
||||
- `--omni` — enables vLLM-Omni diffusion serving.
|
||||
- **Guidance:** Boogu-Image uses `guidance_scale` (mapped to the upstream
|
||||
`text_guidance_scale`); the default is `4.0`. Classifier-free guidance is
|
||||
active whenever `guidance_scale > 1.0`.
|
||||
- **Recommended settings:** `num_inference_steps=28`-`50`, `guidance_scale=4.0`.
|
||||
The model's maximum native resolution is 2K.
|
||||
- **Known limitations (not yet supported):** CPU offload
|
||||
(`--enable-cpu-offload` / `--enable-layerwise-offload`), Cache-DiT
|
||||
(`--cache-backend cache_dit`), and multi-GPU parallelism (TP / SP / CFG /
|
||||
HSDP) are planned follow-ups and are not validated for this model yet.
|
||||
|
||||
## Image editing (Boogu-Image-0.1-Edit)
|
||||
|
||||
The Edit checkpoint is served by the same native pipeline (`BooguImagePipeline`);
|
||||
the image-editing (TI2I) path activates automatically when a request carries a
|
||||
reference image. The Base text-to-image path is unaffected (no reference image
|
||||
is sent).
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
Edit an image with `/v1/images/edits` (the model-card example — change a photo
|
||||
to a colored-pencil drawing). Diffusion parameters are plain multipart form
|
||||
fields; add `guidance_scale_2` to enable image guidance (double-guidance path):
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/images/edits \
|
||||
-F model="Boogu/Boogu-Image-0.1-Edit" \
|
||||
-F image="@input.png" \
|
||||
-F prompt="Change the style to a colored pencil drawing." \
|
||||
-F num_inference_steps=28 \
|
||||
-F guidance_scale=5.0 \
|
||||
-F guidance_scale_2=2.0 \
|
||||
-F seed=42 \
|
||||
| jq -r '.data[0].b64_json' | base64 -d > edited.png
|
||||
```
|
||||
|
||||
Or via chat completions (attach the image as a data URL; parameters go in
|
||||
`extra_body`):
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
|
||||
{"type": "text", "text": "Change the style to a colored pencil drawing."}
|
||||
]}
|
||||
],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 28,
|
||||
"guidance_scale": 5.0,
|
||||
"guidance_scale_2": 1.0,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited.png
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Single reference image:** only one input image is supported for now (the
|
||||
upstream "Only support 1 reference image for now" limit).
|
||||
- **Guidance semantics:**
|
||||
- `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
|
||||
`4.0`); `> 1.0` enables text CFG. Editing typically uses `5.0`.
|
||||
- `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
|
||||
default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
|
||||
(3 model predictions per step), steering more strongly toward the reference
|
||||
image.
|
||||
- **Output resolution:** the output size follows the reference image (upstream
|
||||
`align_res`, on by default for a single reference), so `height`/`width` are
|
||||
derived from the input and requested sizes are not applied for edits.
|
||||
- **Same limitations** as the Base checkpoint apply (no CPU offload, Cache-DiT,
|
||||
or multi-GPU parallelism yet).
|
||||
- **Offline editing:** the shared offline example
|
||||
[`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_to_image.md)
|
||||
supports Boogu-Image-Edit directly
|
||||
(`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional
|
||||
`--guidance-scale-2`).
|
||||
@@ -37,6 +37,7 @@ recipes/
|
||||
| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
|
||||
| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
|
||||
| [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
|
||||
| [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ |
|
||||
| [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB |
|
||||
| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
|
||||
| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_boogu_image_edit_single_device",
|
||||
"description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. Text-only edit path (2 predictions/step); a double-guidance (guidance_scale_2) case is a follow-up since the random dataset cannot inject extra_body params. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Edit",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps28_i2i",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 28,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": [1, 2, 4],
|
||||
"warmup-requests": 4,
|
||||
"warmup-concurrency": 4,
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
|
||||
"latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
|
||||
"peak_memory_mb_max": 40500,
|
||||
"peak_memory_mb_mean": 40500
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,35 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_boogu_image_single_device",
|
||||
"description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0611/0.0609/0.0611, latency_mean 16.37/31.22/55.77 s, peak mem 36754 MB) with ~10% margin.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Base",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps28",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 28,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": [1, 2, 4],
|
||||
"warmup-requests": 4,
|
||||
"warmup-concurrency": 4,
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
|
||||
"latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
|
||||
"peak_memory_mb_max": 40500,
|
||||
"peak_memory_mb_mean": 40500
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
Whitespace-only changes.
@@ -0,0 +1,161 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""L1 unit tests for the Boogu time-shifting flow-match Euler scheduler.
|
||||
|
||||
Golden values were generated from the upstream boogu scheduler
|
||||
(scripts/test_ported_scheduler_parity.py) with the released checkpoint's
|
||||
scheduler_config.json, so CI does not need the `boogu` package installed.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.boogu_image import FlowMatchEulerDiscreteScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# Ground truth from the released checkpoint's scheduler/scheduler_config.json.
|
||||
CHECKPOINT_CONFIG = dict(
|
||||
num_train_timesteps=1000,
|
||||
do_shift=True,
|
||||
dynamic_time_shift=False,
|
||||
time_shift_version="v1",
|
||||
seq_len=4096,
|
||||
base_shift=0.5,
|
||||
max_shift=1.15,
|
||||
)
|
||||
|
||||
NUM_TOKENS = 4096 * 4 # 1024x1024 resolution
|
||||
|
||||
# Generated from the upstream boogu scheduler with CHECKPOINT_CONFIG.
|
||||
GOLDEN_TIMESTEPS = {
|
||||
4: [0.00000000, 0.09546924, 0.24048907, 0.48715585],
|
||||
10: [
|
||||
0.00000000,
|
||||
0.03398621,
|
||||
0.07335263,
|
||||
0.11948693,
|
||||
0.17429835,
|
||||
0.24048907,
|
||||
0.32201332,
|
||||
0.42489707,
|
||||
0.55880034,
|
||||
0.74024153,
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def make_scheduler(**overrides) -> FlowMatchEulerDiscreteScheduler:
|
||||
return FlowMatchEulerDiscreteScheduler(**{**CHECKPOINT_CONFIG, **overrides})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_steps", sorted(GOLDEN_TIMESTEPS))
|
||||
def test_golden_timesteps_checkpoint_config(n_steps):
|
||||
scheduler = make_scheduler()
|
||||
scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
|
||||
np.testing.assert_allclose(
|
||||
scheduler.timesteps.numpy(),
|
||||
np.asarray(GOLDEN_TIMESTEPS[n_steps], dtype=np.float32),
|
||||
atol=1e-7,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_steps", [4, 10, 30])
|
||||
def test_timestep_grid_structure(n_steps):
|
||||
scheduler = make_scheduler()
|
||||
scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
|
||||
|
||||
assert scheduler.timesteps.shape == (n_steps,)
|
||||
assert scheduler._timesteps.shape == (n_steps + 1,)
|
||||
# Ascending 0 -> 1 convention with the terminal 1.0 appended.
|
||||
assert scheduler.timesteps[0].item() == 0.0
|
||||
assert torch.all(scheduler._timesteps[1:] > scheduler._timesteps[:-1])
|
||||
assert scheduler._timesteps[-1].item() == 1.0
|
||||
assert torch.equal(scheduler._timesteps[:-1], scheduler.timesteps)
|
||||
|
||||
|
||||
def test_no_shift_is_uniform_grid():
|
||||
n_steps = 8
|
||||
scheduler = make_scheduler(do_shift=False)
|
||||
scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
|
||||
expected = np.linspace(0, 1, n_steps + 1, dtype=np.float32)[:-1]
|
||||
np.testing.assert_allclose(scheduler.timesteps.numpy(), expected, atol=1e-7)
|
||||
|
||||
|
||||
def test_static_shift_ignores_num_tokens():
|
||||
# dynamic_time_shift=False: the shift comes from configured seq_len, so
|
||||
# per-call num_tokens must not change the schedule.
|
||||
scheduler_a = make_scheduler()
|
||||
scheduler_a.set_timesteps(num_inference_steps=10, num_tokens=NUM_TOKENS)
|
||||
scheduler_b = make_scheduler()
|
||||
scheduler_b.set_timesteps(num_inference_steps=10, num_tokens=320 * 320)
|
||||
assert torch.equal(scheduler_a.timesteps, scheduler_b.timesteps)
|
||||
|
||||
|
||||
def test_step_arithmetic_and_index_advancement():
|
||||
n_steps = 4
|
||||
scheduler = make_scheduler()
|
||||
scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
|
||||
|
||||
sample = torch.zeros(1, 4, 8, 8)
|
||||
model_output = torch.ones(1, 4, 8, 8)
|
||||
|
||||
assert scheduler.step_index is None
|
||||
out = scheduler.step(model_output, scheduler.timesteps[0], sample)
|
||||
assert scheduler.step_index == 1
|
||||
|
||||
# prev_sample = sample + (t_next - t) * model_output
|
||||
expected_dt = scheduler._timesteps[1] - scheduler._timesteps[0]
|
||||
assert torch.allclose(out.prev_sample, expected_dt * model_output)
|
||||
|
||||
# Full loop reaches the end of the grid.
|
||||
sample = out.prev_sample
|
||||
for t in scheduler.timesteps[1:]:
|
||||
sample = scheduler.step(model_output, t, sample, return_dict=False)[0]
|
||||
assert scheduler.step_index == n_steps
|
||||
# Integrating a constant velocity of 1 over t in [0, 1] gives exactly 1.
|
||||
assert torch.allclose(sample, torch.ones_like(sample), atol=1e-6)
|
||||
|
||||
|
||||
def test_step_rejects_integer_timestep():
|
||||
scheduler = make_scheduler()
|
||||
scheduler.set_timesteps(num_inference_steps=4, num_tokens=NUM_TOKENS)
|
||||
with pytest.raises(ValueError, match="integer indices"):
|
||||
scheduler.step(torch.zeros(1), 0, torch.zeros(1))
|
||||
|
||||
|
||||
def test_set_begin_index():
|
||||
scheduler = make_scheduler()
|
||||
scheduler.set_timesteps(num_inference_steps=4, num_tokens=NUM_TOKENS)
|
||||
scheduler.set_begin_index(2)
|
||||
assert scheduler.begin_index == 2
|
||||
|
||||
scheduler.step(torch.zeros(1), scheduler.timesteps[2], torch.zeros(1))
|
||||
assert scheduler.step_index == 3
|
||||
|
||||
|
||||
def test_len_and_config_roundtrip():
|
||||
scheduler = make_scheduler()
|
||||
assert len(scheduler) == CHECKPOINT_CONFIG["num_train_timesteps"]
|
||||
# ConfigMixin must retain the checkpoint config (needed for from_pretrained).
|
||||
assert scheduler.config.do_shift is True
|
||||
assert scheduler.config.dynamic_time_shift is False
|
||||
assert scheduler.config.time_shift_version == "v1"
|
||||
assert scheduler.config.seq_len == 4096
|
||||
|
||||
|
||||
def test_matches_upstream_boogu_scheduler():
|
||||
"""Differential check against the upstream boogu package, when installed."""
|
||||
boogu_schedulers = pytest.importorskip("boogu.schedulers.scheduling_flow_match_euler_discrete_time_shifting")
|
||||
UpstreamScheduler = boogu_schedulers.FlowMatchEulerDiscreteScheduler
|
||||
|
||||
for dynamic in (False, True):
|
||||
for version in ("v1", "v2"):
|
||||
for n_steps in (4, 10, 30):
|
||||
config = {**CHECKPOINT_CONFIG, "dynamic_time_shift": dynamic, "time_shift_version": version}
|
||||
upstream = UpstreamScheduler(**config)
|
||||
ported = make_scheduler(**config)
|
||||
upstream.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
|
||||
ported.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
|
||||
assert torch.equal(upstream.timesteps, ported.timesteps), (dynamic, version, n_steps)
|
||||
assert torch.equal(upstream._timesteps, ported._timesteps), (dynamic, version, n_steps)
|
||||
@@ -0,0 +1,433 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
# Small mock dimensions (head_dim = HIDDEN_SIZE / NUM_HEADS = 16 must equal
|
||||
# sum(AXES_DIM_ROPE)).
|
||||
HIDDEN_SIZE = 64
|
||||
NUM_HEADS = 4
|
||||
NUM_KV_HEADS = 2
|
||||
HEAD_DIM = HIDDEN_SIZE // NUM_HEADS
|
||||
AXES_DIM_ROPE = (8, 4, 4)
|
||||
AXES_LENS = (32, 16, 16)
|
||||
MULTIPLE_OF = 32
|
||||
NORM_EPS = 1e-5
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _init_distributed():
|
||||
"""Initialize the minimal single-rank distributed environment required by
|
||||
the vLLM parallel linear layers (tensor-parallel group must exist)."""
|
||||
from vllm.distributed.parallel_state import (
|
||||
cleanup_dist_env_and_memory,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
|
||||
os.environ.setdefault("MASTER_ADDR", "localhost")
|
||||
os.environ.setdefault("MASTER_PORT", "29512")
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
local_rank=0,
|
||||
distributed_init_method="env://",
|
||||
)
|
||||
initialize_model_parallel()
|
||||
yield
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _force_default_gemm(monkeypatch):
|
||||
"""Force CPU-compatible GEMM dispatch for tests using CPU tensors.
|
||||
|
||||
vLLM's dispatch_unquantized_gemm() selects the backend by platform, not by
|
||||
tensor device; CPU test tensors can crash on non-default backends."""
|
||||
from vllm.model_executor.layers.utils import default_unquantized_gemm
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
|
||||
lambda: default_unquantized_gemm,
|
||||
)
|
||||
|
||||
|
||||
def _randomize_parameters(module: torch.nn.Module) -> None:
|
||||
"""Fill parameters with small random values.
|
||||
|
||||
vLLM parallel linears allocate weights with `torch.empty` (real weights
|
||||
arrive via `load_weights`), so uninitialized memory must be overwritten
|
||||
before a forward pass."""
|
||||
with torch.no_grad():
|
||||
for param in module.parameters():
|
||||
param.uniform_(-0.02, 0.02)
|
||||
|
||||
|
||||
def _identity_rotary_emb(batch_size: int, seq_len: int) -> torch.Tensor:
|
||||
"""Complex rotary frequencies encoding a zero rotation."""
|
||||
return torch.polar(
|
||||
torch.ones(batch_size, seq_len, HEAD_DIM // 2),
|
||||
torch.zeros(batch_size, seq_len, HEAD_DIM // 2),
|
||||
)
|
||||
|
||||
|
||||
def _tiny_tf_model_config(**overrides):
|
||||
config = {
|
||||
"patch_size": 2,
|
||||
"in_channels": 4,
|
||||
"hidden_size": HIDDEN_SIZE,
|
||||
"num_layers": 4,
|
||||
"num_double_stream_layers": 2,
|
||||
"num_refiner_layers": 2,
|
||||
"num_attention_heads": NUM_HEADS,
|
||||
"num_kv_heads": NUM_KV_HEADS,
|
||||
"multiple_of": MULTIPLE_OF,
|
||||
"norm_eps": NORM_EPS,
|
||||
"axes_dim_rope": list(AXES_DIM_ROPE),
|
||||
"axes_lens": list(AXES_LENS),
|
||||
"instruction_feature_configs": {
|
||||
"instruction_feat_dim": 32,
|
||||
"reduce_type": "mean",
|
||||
"num_instruction_feature_layers": 1,
|
||||
},
|
||||
"prompt_tuning_configs": {"use_prompt_tuning": False},
|
||||
"timestep_scale": 1.0,
|
||||
}
|
||||
config.update(overrides)
|
||||
return config
|
||||
|
||||
|
||||
def _tiny_od_config(**overrides):
|
||||
from vllm_omni.diffusion.data import TransformerConfig
|
||||
|
||||
return SimpleNamespace(
|
||||
tf_model_config=TransformerConfig.from_dict(_tiny_tf_model_config(**overrides)),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def test_boogu_image_transformer_import():
|
||||
from vllm_omni.diffusion.models.boogu_image import BooguImageTransformer2DModel
|
||||
|
||||
assert BooguImageTransformer2DModel is not None
|
||||
|
||||
|
||||
def test_single_stream_block_shape():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformerBlock,
|
||||
)
|
||||
|
||||
block = BooguImageTransformerBlock(
|
||||
dim=HIDDEN_SIZE,
|
||||
num_attention_heads=NUM_HEADS,
|
||||
num_kv_heads=NUM_KV_HEADS,
|
||||
multiple_of=MULTIPLE_OF,
|
||||
ffn_dim_multiplier=None,
|
||||
norm_eps=NORM_EPS,
|
||||
modulation=True,
|
||||
)
|
||||
_randomize_parameters(block)
|
||||
|
||||
batch_size, seq_len = 1, 16
|
||||
hidden_states = torch.randn(batch_size, seq_len, HIDDEN_SIZE)
|
||||
attention_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
|
||||
rotary_emb = _identity_rotary_emb(batch_size, seq_len)
|
||||
temb = torch.randn(batch_size, min(HIDDEN_SIZE, 1024))
|
||||
|
||||
out = block(hidden_states, attention_mask, rotary_emb, temb)
|
||||
assert out.shape == hidden_states.shape
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
with pytest.raises(ValueError, match="temb"):
|
||||
block(hidden_states, attention_mask, rotary_emb, None)
|
||||
|
||||
|
||||
def test_single_stream_block_no_modulation():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformerBlock,
|
||||
)
|
||||
|
||||
block = BooguImageTransformerBlock(
|
||||
dim=HIDDEN_SIZE,
|
||||
num_attention_heads=NUM_HEADS,
|
||||
num_kv_heads=NUM_KV_HEADS,
|
||||
multiple_of=MULTIPLE_OF,
|
||||
ffn_dim_multiplier=None,
|
||||
norm_eps=NORM_EPS,
|
||||
modulation=False,
|
||||
)
|
||||
_randomize_parameters(block)
|
||||
|
||||
batch_size, seq_len = 2, 8
|
||||
hidden_states = torch.randn(batch_size, seq_len, HIDDEN_SIZE)
|
||||
out = block(hidden_states, None, _identity_rotary_emb(batch_size, seq_len))
|
||||
assert out.shape == hidden_states.shape
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
|
||||
def test_double_stream_block_shape():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageDoubleStreamTransformerBlock,
|
||||
)
|
||||
|
||||
block = BooguImageDoubleStreamTransformerBlock(
|
||||
dim=HIDDEN_SIZE,
|
||||
num_attention_heads=NUM_HEADS,
|
||||
num_kv_heads=NUM_KV_HEADS,
|
||||
multiple_of=MULTIPLE_OF,
|
||||
ffn_dim_multiplier=None,
|
||||
norm_eps=NORM_EPS,
|
||||
modulation=True,
|
||||
)
|
||||
_randomize_parameters(block)
|
||||
|
||||
batch_size = 1
|
||||
img_len, instruct_len = 16, 8
|
||||
total_len = img_len + instruct_len
|
||||
|
||||
img_hidden_states = torch.randn(batch_size, img_len, HIDDEN_SIZE)
|
||||
instruct_hidden_states = torch.randn(batch_size, instruct_len, HIDDEN_SIZE)
|
||||
img_attention_mask = torch.ones(batch_size, img_len, dtype=torch.bool)
|
||||
joint_attention_mask = torch.ones(batch_size, total_len, dtype=torch.bool)
|
||||
image_rotary_emb = _identity_rotary_emb(batch_size, img_len)
|
||||
rotary_emb = _identity_rotary_emb(batch_size, total_len)
|
||||
temb = torch.randn(batch_size, min(HIDDEN_SIZE, 1024))
|
||||
|
||||
img_out, instruct_out = block(
|
||||
img_hidden_states,
|
||||
instruct_hidden_states,
|
||||
img_attention_mask,
|
||||
joint_attention_mask,
|
||||
image_rotary_emb,
|
||||
rotary_emb,
|
||||
temb=temb,
|
||||
encoder_seq_lengths=[instruct_len] * batch_size,
|
||||
seq_lengths=[total_len] * batch_size,
|
||||
)
|
||||
|
||||
assert img_out.shape == img_hidden_states.shape
|
||||
assert instruct_out.shape == instruct_hidden_states.shape
|
||||
assert torch.isfinite(img_out).all()
|
||||
assert torch.isfinite(instruct_out).all()
|
||||
|
||||
|
||||
def test_transformer_instantiates():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
|
||||
assert BooguImageTransformer2DModel._repeated_blocks == [
|
||||
"BooguImageTransformerBlock",
|
||||
"BooguImageContextRefinerTransformerBlock",
|
||||
"BooguImageSingleStreamTransformerBlock",
|
||||
]
|
||||
assert BooguImageTransformer2DModel._layerwise_offload_blocks_attrs == [
|
||||
"single_stream_layers",
|
||||
"double_stream_layers",
|
||||
]
|
||||
|
||||
assert len(model.noise_refiner) == 2
|
||||
assert len(model.ref_image_refiner) == 2
|
||||
assert len(model.context_refiner) == 2
|
||||
assert len(model.double_stream_layers) == 2
|
||||
assert len(model.single_stream_layers) == 2 # num_layers - num_double_stream_layers
|
||||
assert model.image_index_embedding.shape == (5, HIDDEN_SIZE)
|
||||
# patch_size^2 * in_channels -> hidden_size
|
||||
assert model.x_embedder.in_features == 2 * 2 * 4
|
||||
assert model.x_embedder.out_features == HIDDEN_SIZE
|
||||
|
||||
|
||||
def test_transformer_preprocesses_multiple_instruction_feature_layers():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
instruction_feature_configs = {
|
||||
"instruction_feat_dim": 32,
|
||||
"reduce_type": "concat",
|
||||
"num_instruction_feature_layers": 2,
|
||||
}
|
||||
model = BooguImageTransformer2DModel(
|
||||
od_config=_tiny_od_config(instruction_feature_configs=instruction_feature_configs)
|
||||
)
|
||||
hidden_states = [torch.randn(1, 8, 32), torch.randn(1, 8, 32)]
|
||||
|
||||
processed = model.preprocess_instruction_hidden_states(hidden_states)
|
||||
|
||||
assert model.preprocessed_instruction_feat_dim == 64
|
||||
assert processed.shape == (1, 8, 64)
|
||||
assert torch.equal(processed, torch.cat(hidden_states, dim=-1))
|
||||
|
||||
|
||||
def test_transformer_validates_rope_dims():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="axes_dim_rope"):
|
||||
BooguImageTransformer2DModel(od_config=_tiny_od_config(axes_dim_rope=[8, 8, 8]))
|
||||
|
||||
|
||||
def test_transformer_rejects_prompt_tuning():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="[Pp]rompt tuning"):
|
||||
BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
|
||||
|
||||
|
||||
def _native_to_checkpoint_name(name: str) -> str:
|
||||
"""Inverse of ``load_weights`` remapping: native param name -> diffusers name.
|
||||
|
||||
- ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
|
||||
- promoted joint-attention projections move back under ``.processor.``.
|
||||
"""
|
||||
if ".to_out." in name:
|
||||
name = name.replace(".to_out.", ".to_out.0.")
|
||||
for proj in (
|
||||
"img_to_q",
|
||||
"img_to_k",
|
||||
"img_to_v",
|
||||
"instruct_to_q",
|
||||
"instruct_to_k",
|
||||
"instruct_to_v",
|
||||
"instruct_out",
|
||||
"img_out",
|
||||
):
|
||||
token = f".img_instruct_attn.{proj}."
|
||||
if token in name:
|
||||
name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
|
||||
break
|
||||
return name
|
||||
|
||||
|
||||
def test_transformer_load_weights_round_trip():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
native_params = dict(model.named_parameters())
|
||||
|
||||
# Build synthetic diffusers-named weights (one per native parameter).
|
||||
checkpoint_weights = {}
|
||||
for native_name, param in native_params.items():
|
||||
checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
|
||||
|
||||
# The remapping must be a bijection over the parameter set.
|
||||
assert len(checkpoint_weights) == len(native_params)
|
||||
|
||||
loaded = model.load_weights(list(checkpoint_weights.items()))
|
||||
|
||||
# No missing / unexpected parameters.
|
||||
assert loaded == set(native_params.keys())
|
||||
|
||||
# Values landed on the right parameters (TP=1: weight_loader copies verbatim).
|
||||
reloaded = dict(model.named_parameters())
|
||||
for native_name in native_params:
|
||||
expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
|
||||
assert torch.allclose(reloaded[native_name], expected)
|
||||
|
||||
|
||||
def test_transformer_load_weights_warns_for_unexpected_and_unloaded(caplog):
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
native_params = dict(model.named_parameters())
|
||||
loaded_name = next(iter(native_params))
|
||||
checkpoint_name = _native_to_checkpoint_name(loaded_name)
|
||||
|
||||
loaded = model.load_weights(
|
||||
[
|
||||
(checkpoint_name, torch.randn_like(native_params[loaded_name])),
|
||||
("unexpected.weight", torch.ones(1)),
|
||||
]
|
||||
)
|
||||
|
||||
assert loaded == {loaded_name}
|
||||
assert "Skipping unexpected checkpoint weight unexpected.weight" in caplog.text
|
||||
assert "Model parameters not loaded from checkpoint" in caplog.text
|
||||
assert next(name for name in native_params if name != loaded_name) in caplog.text
|
||||
|
||||
|
||||
def test_transformer_forward_t2i_shape():
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageDoubleStreamRotaryPosEmbed,
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
_randomize_parameters(model)
|
||||
model.eval()
|
||||
|
||||
batch_size = 1
|
||||
in_channels = 4
|
||||
latent_h = latent_w = 8 # multiples of patch_size (2)
|
||||
instruct_len = 8
|
||||
instruction_feat_dim = 32 # matches _tiny_tf_model_config
|
||||
|
||||
latents = torch.randn(batch_size, in_channels, latent_h, latent_w)
|
||||
timestep = torch.full((batch_size,), 0.5)
|
||||
instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
|
||||
instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
|
||||
freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
|
||||
|
||||
with torch.no_grad():
|
||||
out = model(latents, timestep, instruction_hidden_states, freqs_cis, instruction_attention_mask)
|
||||
|
||||
assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
|
||||
def test_transformer_forward_ti2i_shape():
|
||||
"""Editing path: a non-empty ``ref_image_hidden_states`` exercises the
|
||||
reference-image patch embedder + refiner and must not change the output
|
||||
shape (the output tracks the noise-latent dimensions)."""
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageDoubleStreamRotaryPosEmbed,
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
_randomize_parameters(model)
|
||||
model.eval()
|
||||
|
||||
batch_size = 1
|
||||
in_channels = 4
|
||||
latent_h = latent_w = 8
|
||||
ref_h, ref_w = 6, 10 # a differently-sized reference latent
|
||||
instruct_len = 8
|
||||
instruction_feat_dim = 32
|
||||
|
||||
latents = torch.randn(batch_size, in_channels, latent_h, latent_w)
|
||||
timestep = torch.full((batch_size,), 0.5)
|
||||
instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
|
||||
instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
|
||||
freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
|
||||
|
||||
# One sample, one reference image (Boogu editing supports a single ref).
|
||||
ref_image_hidden_states = [[torch.randn(in_channels, ref_h, ref_w)]]
|
||||
|
||||
with torch.no_grad():
|
||||
out = model(
|
||||
latents,
|
||||
timestep,
|
||||
instruction_hidden_states,
|
||||
freqs_cis,
|
||||
instruction_attention_mask,
|
||||
ref_image_hidden_states=ref_image_hidden_states,
|
||||
)
|
||||
|
||||
assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
|
||||
assert torch.isfinite(out).all()
|
||||
@@ -0,0 +1,900 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""L1 unit tests for the native Boogu-Image pipeline.
|
||||
|
||||
Two groups:
|
||||
|
||||
1. Constructor tests: all ``from_pretrained`` calls are mocked (Ovis pattern)
|
||||
so the pipeline ``__init__`` wiring is exercised without downloading
|
||||
weights or building the real transformer.
|
||||
2. Prompt-encoding tests: a pipeline shell (``object.__new__``) is wired with
|
||||
small deterministic fakes so the ported chat templating, processor kwargs,
|
||||
CFG handling, and reshape logic can be verified numerically on CPU.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig, TransformerConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
_MODULE = "vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image"
|
||||
|
||||
_EMBED_DIM = 8
|
||||
_SEQ_LEN = 6
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constructor tests (mocked from_pretrained)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_dependencies(mocker, monkeypatch):
|
||||
"""Mock external components so ``__init__`` runs without real weights."""
|
||||
# The full VLM wrapper has an ``lm_head``; the pipeline must strip it and
|
||||
# keep the inner ``.model`` as the encoder.
|
||||
inner_encoder = mocker.MagicMock(name="inner_qwen3vl_model")
|
||||
inner_encoder.dtype = torch.float32
|
||||
mllm_wrapper = mocker.MagicMock(name="qwen3vl_wrapper")
|
||||
mllm_wrapper.model.to.return_value = inner_encoder
|
||||
|
||||
mock_processor = mocker.MagicMock(name="processor")
|
||||
|
||||
mock_vae = mocker.MagicMock(name="vae")
|
||||
mock_vae.config.block_out_channels = [128, 256, 512, 512] # scale factor 8
|
||||
mock_vae.to.return_value = mock_vae
|
||||
|
||||
mock_scheduler = mocker.MagicMock(name="scheduler")
|
||||
|
||||
monkeypatch.setattr(
|
||||
f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
|
||||
lambda *a, **k: mock_scheduler,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
|
||||
lambda *a, **k: mllm_wrapper,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
|
||||
lambda *a, **k: mock_processor,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
f"{_MODULE}.AutoencoderKL.from_pretrained",
|
||||
lambda *a, **k: mock_vae,
|
||||
)
|
||||
|
||||
mock_transformer_cls = mocker.MagicMock(name="transformer_cls")
|
||||
mock_transformer_instance = mocker.MagicMock(name="transformer")
|
||||
mock_transformer_cls.return_value = mock_transformer_instance
|
||||
monkeypatch.setattr(f"{_MODULE}.BooguImageTransformer2DModel", mock_transformer_cls)
|
||||
|
||||
# Treat the dummy model id as a local path: skips hub prefetch.
|
||||
mocker.patch("os.path.exists", return_value=True)
|
||||
|
||||
return {
|
||||
"inner_encoder": inner_encoder,
|
||||
"mllm_wrapper": mllm_wrapper,
|
||||
"processor": mock_processor,
|
||||
"vae": mock_vae,
|
||||
"scheduler": mock_scheduler,
|
||||
"transformer": mock_transformer_instance,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def boogu_pipeline(mock_dependencies):
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
BooguImagePipeline,
|
||||
)
|
||||
|
||||
od_config = OmniDiffusionConfig(
|
||||
model="dummy-boogu",
|
||||
tf_model_config=TransformerConfig(params={}),
|
||||
dtype=torch.float32,
|
||||
num_gpus=1,
|
||||
)
|
||||
return BooguImagePipeline(od_config=od_config)
|
||||
|
||||
|
||||
def test_boogu_image_pipeline_import():
|
||||
from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline
|
||||
|
||||
assert BooguImagePipeline is not None
|
||||
|
||||
|
||||
def test_component_discovery_declarations():
|
||||
from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline
|
||||
|
||||
# CPU offload / HSDP discovery must find ``mllm`` (there is no
|
||||
# ``text_encoder`` attribute on this pipeline).
|
||||
assert BooguImagePipeline._dit_modules == ["transformer"]
|
||||
assert BooguImagePipeline._encoder_modules == ["mllm"]
|
||||
assert BooguImagePipeline._vae_modules == ["vae"]
|
||||
|
||||
|
||||
def test_constructor_wires_components(boogu_pipeline, mock_dependencies):
|
||||
assert boogu_pipeline.scheduler is mock_dependencies["scheduler"]
|
||||
assert boogu_pipeline.processor is mock_dependencies["processor"]
|
||||
assert boogu_pipeline.vae is mock_dependencies["vae"]
|
||||
assert boogu_pipeline.transformer is mock_dependencies["transformer"]
|
||||
assert boogu_pipeline.vae_scale_factor == 8
|
||||
assert boogu_pipeline.default_sample_size == 128
|
||||
assert hasattr(boogu_pipeline, "load_weights")
|
||||
|
||||
|
||||
def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
|
||||
# Upstream encodes with the inner Qwen3VLModel, not the generation wrapper.
|
||||
assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"]
|
||||
|
||||
|
||||
def test_constructor_weights_sources(boogu_pipeline):
|
||||
(source,) = boogu_pipeline.weights_sources
|
||||
assert source.model_or_path == "dummy-boogu"
|
||||
assert source.subfolder == "transformer"
|
||||
assert source.prefix == "transformer."
|
||||
assert source.fall_back_to_pt is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("parallel_config", "cache_backend", "message"),
|
||||
[
|
||||
(DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
|
||||
(DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"),
|
||||
(DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"),
|
||||
(DiffusionParallelConfig(cfg_parallel_size=2), "none", "CFG parallelism"),
|
||||
(
|
||||
DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2),
|
||||
"none",
|
||||
"HSDP",
|
||||
),
|
||||
(DiffusionParallelConfig(), "cache_dit", "Cache backend 'cache_dit'"),
|
||||
(DiffusionParallelConfig(), "tea_cache", "Cache backend 'tea_cache'"),
|
||||
],
|
||||
)
|
||||
def test_constructor_rejects_unsupported_execution_modes(
|
||||
mock_dependencies,
|
||||
parallel_config,
|
||||
cache_backend,
|
||||
message,
|
||||
):
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
BooguImagePipeline,
|
||||
)
|
||||
|
||||
od_config = OmniDiffusionConfig(
|
||||
model="dummy-boogu",
|
||||
tf_model_config=TransformerConfig(params={}),
|
||||
dtype=torch.float32,
|
||||
parallel_config=parallel_config,
|
||||
cache_backend=cache_backend,
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
BooguImagePipeline(od_config=od_config)
|
||||
|
||||
# Validation happens before any checkpoint component is constructed.
|
||||
mock_dependencies["mllm_wrapper"].model.to.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompt-encoding tests (deterministic fakes, no constructor)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _RecordingProcessor:
|
||||
"""Fake Qwen3VLProcessor: deterministic token ids derived from the text."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def apply_chat_template(self, prompts, **kwargs):
|
||||
self.calls.append({"prompts": prompts, "kwargs": kwargs})
|
||||
batch = len(prompts)
|
||||
input_ids = torch.zeros(batch, _SEQ_LEN, dtype=torch.long)
|
||||
for i, messages in enumerate(prompts):
|
||||
system_text = messages[0]["content"][0]["text"]
|
||||
user_text = messages[1]["content"][0]["text"]
|
||||
input_ids[i, 0] = len(system_text) % 997
|
||||
input_ids[i, 1] = len(user_text) % 997
|
||||
input_ids[i, 2:] = torch.arange(2, _SEQ_LEN) + i
|
||||
attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long)
|
||||
attention_mask[:, -1] = 0 # fake right-padding
|
||||
return {"input_ids": input_ids, "attention_mask": attention_mask}
|
||||
|
||||
|
||||
class _FakeMLLM:
|
||||
"""Fake Qwen3VLModel: hidden states are a deterministic function of ids."""
|
||||
|
||||
dtype = torch.bfloat16
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, input_ids=None, attention_mask=None, output_hidden_states=False, **kwargs):
|
||||
self.calls.append({"output_hidden_states": output_hidden_states, **kwargs})
|
||||
base = input_ids.to(torch.float32).unsqueeze(-1).repeat(1, 1, _EMBED_DIM)
|
||||
hidden = base + torch.arange(_EMBED_DIM, dtype=torch.float32)
|
||||
return SimpleNamespace(last_hidden_state=hidden, hidden_states=(hidden - 1, hidden))
|
||||
|
||||
|
||||
def _make_encode_pipeline(num_instruction_feature_layers: int = 1):
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
SYSTEM_PROMPT_4_T2I_UNIFIED,
|
||||
SYSTEM_PROMPT_4_TI2I_UNIFIED,
|
||||
BooguImagePipeline,
|
||||
)
|
||||
|
||||
pipeline = object.__new__(BooguImagePipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline._execution_device = torch.device("cpu")
|
||||
pipeline.processor = _RecordingProcessor()
|
||||
pipeline.mllm = _FakeMLLM()
|
||||
pipeline.transformer = SimpleNamespace(
|
||||
instruction_feature_configs={
|
||||
"instruction_feat_dim": _EMBED_DIM,
|
||||
"num_instruction_feature_layers": num_instruction_feature_layers,
|
||||
"reduce_type": "mean",
|
||||
},
|
||||
dtype=torch.float32,
|
||||
)
|
||||
pipeline.SYSTEM_PROMPT_4_T2I = SYSTEM_PROMPT_4_T2I_UNIFIED
|
||||
pipeline.SYSTEM_PROMPT_DROP = SYSTEM_PROMPT_4_TI2I_UNIFIED
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_apply_chat_template_system_prompt_selection():
|
||||
pipeline = _make_encode_pipeline()
|
||||
|
||||
messages = pipeline._apply_chat_template("a cat on a mat")
|
||||
assert messages[0]["role"] == "system"
|
||||
assert messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_T2I
|
||||
assert messages[1]["role"] == "user"
|
||||
assert messages[1]["content"][0]["text"] == "a cat on a mat"
|
||||
|
||||
# Empty and whitespace-only instructions select the DROP prompt (this is
|
||||
# the path the default negative prompt "" takes).
|
||||
for empty in ("", " ", None):
|
||||
messages = pipeline._apply_chat_template(empty)
|
||||
assert messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_DROP
|
||||
|
||||
|
||||
def test_processor_called_with_upstream_kwargs():
|
||||
pipeline = _make_encode_pipeline()
|
||||
pipeline.encode_prompt("a dog", do_classifier_free_guidance=False)
|
||||
|
||||
(call,) = pipeline.processor.calls
|
||||
kwargs = call["kwargs"]
|
||||
assert kwargs["padding"] == "longest"
|
||||
assert kwargs["padding_side"] == "right"
|
||||
assert kwargs["truncation"] is False
|
||||
assert kwargs["max_length"] == 1280 # upstream __call__ default
|
||||
assert kwargs["tokenize"] is True
|
||||
assert kwargs["return_dict"] is True
|
||||
assert kwargs["return_tensors"] == "pt"
|
||||
|
||||
|
||||
def test_encode_prompt_shapes_and_dtype():
|
||||
pipeline = _make_encode_pipeline()
|
||||
embeds, mask, neg_embeds, neg_mask = pipeline.encode_prompt(["a dog", "a cat"])
|
||||
|
||||
assert embeds.shape == (2, _SEQ_LEN, _EMBED_DIM)
|
||||
assert mask.shape == (2, _SEQ_LEN)
|
||||
assert neg_embeds.shape == (2, _SEQ_LEN, _EMBED_DIM)
|
||||
assert neg_mask.shape == (2, _SEQ_LEN)
|
||||
# Cast to the MLLM dtype, mask passed through from the processor.
|
||||
assert embeds.dtype == torch.bfloat16
|
||||
assert neg_embeds.dtype == torch.bfloat16
|
||||
assert torch.equal(mask[:, -1], torch.zeros(2, dtype=torch.long))
|
||||
|
||||
|
||||
def test_single_layer_encoding_requests_hidden_states_once():
|
||||
pipeline = _make_encode_pipeline()
|
||||
|
||||
pipeline._get_instruction_feature_embeds("a dog")
|
||||
|
||||
assert pipeline.mllm.calls == [{"output_hidden_states": True, "return_dict": True}]
|
||||
|
||||
|
||||
def test_single_layer_encoding_propagates_mllm_failure():
|
||||
pipeline = _make_encode_pipeline()
|
||||
failure = RuntimeError("mllm failed")
|
||||
|
||||
class _FailingMLLM:
|
||||
dtype = torch.bfloat16
|
||||
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
self.calls += 1
|
||||
raise failure
|
||||
|
||||
pipeline.mllm = _FailingMLLM()
|
||||
|
||||
with pytest.raises(RuntimeError, match="mllm failed") as exc_info:
|
||||
pipeline._get_instruction_feature_embeds("a dog")
|
||||
|
||||
assert exc_info.value is failure
|
||||
assert pipeline.mllm.calls == 1
|
||||
|
||||
|
||||
def test_encode_prompt_cfg_negative_default_is_empty_string():
|
||||
pipeline = _make_encode_pipeline()
|
||||
pipeline.encode_prompt("a dog")
|
||||
|
||||
assert len(pipeline.processor.calls) == 2
|
||||
negative_messages = pipeline.processor.calls[1]["prompts"]
|
||||
assert len(negative_messages) == 1
|
||||
# The default "" negative prompt goes through the DROP system prompt.
|
||||
assert negative_messages[0][0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_DROP
|
||||
assert negative_messages[0][1]["content"][0]["text"] == ""
|
||||
|
||||
|
||||
def test_encode_prompt_without_cfg_skips_negative():
|
||||
pipeline = _make_encode_pipeline()
|
||||
embeds, mask, neg_embeds, neg_mask = pipeline.encode_prompt("a dog", do_classifier_free_guidance=False)
|
||||
|
||||
assert len(pipeline.processor.calls) == 1
|
||||
assert embeds.shape == (1, _SEQ_LEN, _EMBED_DIM)
|
||||
assert neg_embeds is None
|
||||
assert neg_mask is None
|
||||
|
||||
|
||||
def test_encode_prompt_explicit_negative_prompt():
|
||||
pipeline = _make_encode_pipeline()
|
||||
pipeline.encode_prompt("a dog", negative_prompt="blurry, low quality")
|
||||
|
||||
negative_messages = pipeline.processor.calls[1]["prompts"]
|
||||
assert negative_messages[0][0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_T2I
|
||||
assert negative_messages[0][1]["content"][0]["text"] == "blurry, low quality"
|
||||
|
||||
|
||||
def test_encode_prompt_negative_batch_mismatch_raises():
|
||||
pipeline = _make_encode_pipeline()
|
||||
with pytest.raises(ValueError, match="batch size"):
|
||||
pipeline.encode_prompt(["a dog", "a cat"], negative_prompt=["only one"])
|
||||
|
||||
|
||||
def test_encode_prompt_num_images_per_prompt_repeats_embeds():
|
||||
pipeline = _make_encode_pipeline()
|
||||
embeds, mask, neg_embeds, neg_mask = pipeline.encode_prompt("a dog", num_images_per_prompt=3)
|
||||
|
||||
assert embeds.shape == (3, _SEQ_LEN, _EMBED_DIM)
|
||||
assert mask.shape == (3, _SEQ_LEN)
|
||||
assert torch.equal(embeds[0], embeds[1])
|
||||
assert torch.equal(embeds[0], embeds[2])
|
||||
assert torch.equal(mask[0], mask[1])
|
||||
assert neg_embeds.shape == (3, _SEQ_LEN, _EMBED_DIM)
|
||||
|
||||
|
||||
def test_encode_prompt_precomputed_embeds_bypass_encoder():
|
||||
pipeline = _make_encode_pipeline()
|
||||
precomputed = torch.randn(1, _SEQ_LEN, _EMBED_DIM)
|
||||
precomputed_mask = torch.ones(1, _SEQ_LEN, dtype=torch.long)
|
||||
neg_precomputed = torch.randn(1, _SEQ_LEN, _EMBED_DIM)
|
||||
|
||||
embeds, mask, neg_embeds, _ = pipeline.encode_prompt(
|
||||
"ignored",
|
||||
prompt_embeds=precomputed,
|
||||
prompt_attention_mask=precomputed_mask,
|
||||
negative_prompt_embeds=neg_precomputed,
|
||||
)
|
||||
|
||||
assert len(pipeline.processor.calls) == 0
|
||||
assert torch.equal(embeds, precomputed)
|
||||
assert torch.equal(neg_embeds, neg_precomputed)
|
||||
|
||||
|
||||
def test_reshape_embeds_and_mask_list_branch():
|
||||
# Multi-layer configs return a list of per-layer tensors; the reshape
|
||||
# helper must handle both forms.
|
||||
pipeline = _make_encode_pipeline()
|
||||
layer = torch.arange(2 * _SEQ_LEN * _EMBED_DIM, dtype=torch.float32).view(2, _SEQ_LEN, _EMBED_DIM)
|
||||
mask = torch.ones(2, _SEQ_LEN, dtype=torch.long)
|
||||
|
||||
batch_size, seq_len, reshaped, reshaped_mask = pipeline._reshape_embeds_and_mask([layer, layer + 1], mask, 2)
|
||||
|
||||
assert batch_size == 2
|
||||
assert seq_len == _SEQ_LEN
|
||||
assert isinstance(reshaped, list) and len(reshaped) == 2
|
||||
assert reshaped[0].shape == (4, _SEQ_LEN, _EMBED_DIM)
|
||||
# repeat(1, n, 1).view(b*n, ...) interleaves per-sample: rows [s0, s0, s1, s1]
|
||||
assert torch.equal(reshaped[0][0], reshaped[0][1])
|
||||
assert torch.equal(reshaped[0][2], reshaped[0][3])
|
||||
|
||||
|
||||
class _FakeTransformer:
|
||||
"""Fake denoiser: velocity 0 (identity flow) with the config the loop reads."""
|
||||
|
||||
in_channels = 4
|
||||
axes_dim_rope = (8, 4, 4)
|
||||
axes_lens = (32, 16, 16)
|
||||
dtype = torch.float32
|
||||
instruction_feature_configs = {
|
||||
"instruction_feat_dim": _EMBED_DIM,
|
||||
"num_instruction_feature_layers": 1,
|
||||
"reduce_type": "mean",
|
||||
}
|
||||
|
||||
def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
|
||||
return torch.zeros_like(latents)
|
||||
|
||||
|
||||
class _FakeScheduler:
|
||||
def set_timesteps(self, num_inference_steps, device=None, num_tokens=None):
|
||||
self.timesteps = torch.linspace(0, 1, num_inference_steps + 1)[:-1]
|
||||
|
||||
def step(self, model_output, t, latents, return_dict=False):
|
||||
return (latents,)
|
||||
|
||||
|
||||
class _FakeDecodeVAE:
|
||||
dtype = torch.float32
|
||||
|
||||
def __init__(self):
|
||||
self.config = SimpleNamespace(scaling_factor=1.0, shift_factor=0.0, block_out_channels=[128, 256, 512, 512])
|
||||
|
||||
def decode(self, latents, return_dict=False):
|
||||
batch = latents.shape[0]
|
||||
return (torch.zeros(batch, 3, 16, 16),)
|
||||
|
||||
|
||||
def _make_forward_pipeline():
|
||||
pipeline = _make_encode_pipeline()
|
||||
pipeline.transformer = _FakeTransformer()
|
||||
pipeline.scheduler = _FakeScheduler()
|
||||
pipeline.vae = _FakeDecodeVAE()
|
||||
pipeline.vae_scale_factor = 8
|
||||
pipeline.default_sample_size = 128
|
||||
return pipeline
|
||||
|
||||
|
||||
def _make_request_batch(prompt, **sampling_overrides):
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
sampling = OmniDiffusionSamplingParams(**sampling_overrides)
|
||||
return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
|
||||
|
||||
|
||||
def test_forward_returns_diffusion_output():
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
|
||||
pipeline = _make_forward_pipeline()
|
||||
req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2)
|
||||
|
||||
out = pipeline.forward(req)
|
||||
|
||||
assert isinstance(out, DiffusionOutput)
|
||||
assert isinstance(out.output, torch.Tensor)
|
||||
assert out.output.shape[0] == 1
|
||||
assert torch.isfinite(out.output).all()
|
||||
# CFG default is on (text guidance 4.0), so the encoder ran twice (pos + neg).
|
||||
assert len(pipeline.processor.calls) == 2
|
||||
|
||||
|
||||
def test_forward_cfg_off_when_guidance_one():
|
||||
pipeline = _make_forward_pipeline()
|
||||
req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2, guidance_scale=1.0)
|
||||
# guidance_scale=1.0 is falsy-adjacent but explicitly disables CFG; the
|
||||
# request layer would set guidance_scale_provided, so emulate that here.
|
||||
req.sampling_params.guidance_scale_provided = True
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
# Only the positive prompt is encoded when CFG is off.
|
||||
assert len(pipeline.processor.calls) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Editing / TI2I: pre-process function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_edit_od_config(tmp_path, block_out_channels=(128, 256, 512, 512)):
|
||||
import json
|
||||
|
||||
vae_dir = tmp_path / "vae"
|
||||
vae_dir.mkdir(parents=True, exist_ok=True)
|
||||
(vae_dir / "config.json").write_text(json.dumps({"block_out_channels": list(block_out_channels)}))
|
||||
return OmniDiffusionConfig(
|
||||
model=str(tmp_path),
|
||||
tf_model_config=TransformerConfig(params={}),
|
||||
dtype=torch.float32,
|
||||
num_gpus=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"factory_name",
|
||||
["get_boogu_image_pre_process_func", "get_boogu_image_post_process_func"],
|
||||
)
|
||||
@pytest.mark.parametrize("config_contents", [None, "{not valid json"])
|
||||
def test_process_factory_reports_invalid_vae_config(tmp_path, factory_name, config_contents):
|
||||
import importlib
|
||||
|
||||
vae_dir = tmp_path / "vae"
|
||||
vae_dir.mkdir()
|
||||
if config_contents is not None:
|
||||
(vae_dir / "config.json").write_text(config_contents)
|
||||
|
||||
od_config = OmniDiffusionConfig(
|
||||
model=str(tmp_path),
|
||||
tf_model_config=TransformerConfig(params={}),
|
||||
dtype=torch.float32,
|
||||
num_gpus=1,
|
||||
)
|
||||
factory = getattr(importlib.import_module(_MODULE), factory_name)
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"Failed to load Boogu VAE config from .*vae/config\.json"):
|
||||
factory(od_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("height", "width"), [(0, 16), (16, 0), (0, 0)])
|
||||
def test_image_processor_rejects_non_positive_dimensions(height, width):
|
||||
from vllm_omni.diffusion.models.boogu_image.image_processor import BooguImageProcessor
|
||||
|
||||
image_processor = BooguImageProcessor()
|
||||
image = torch.zeros(1, 3, 16, 16)
|
||||
|
||||
with pytest.raises(ValueError, match=rf"height={height}, width={width}"):
|
||||
image_processor.get_new_height_width(image, height=height, width=width)
|
||||
|
||||
|
||||
def _make_diffusion_request(prompt, **sampling_overrides):
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
sampling = OmniDiffusionSamplingParams(**sampling_overrides)
|
||||
return OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="req-0")
|
||||
|
||||
|
||||
def test_pre_process_no_image_is_noop(tmp_path):
|
||||
import PIL.Image # noqa: F401 (import guard: PIL must be available)
|
||||
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
get_boogu_image_pre_process_func,
|
||||
)
|
||||
|
||||
pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
|
||||
|
||||
# Text-to-image request: no multimodal image -> returned unchanged.
|
||||
req = _make_diffusion_request({"prompt": "a cat"}, height=123, width=456)
|
||||
out = pre(req)
|
||||
assert "additional_information" not in out.prompt
|
||||
assert out.sampling_params.height == 123
|
||||
assert out.sampling_params.width == 456
|
||||
|
||||
# A plain-string prompt cannot carry an image either.
|
||||
str_req = _make_diffusion_request("a cat")
|
||||
assert pre(str_req).prompt == "a cat"
|
||||
|
||||
|
||||
def test_pre_process_populates_reference_and_align_res(tmp_path):
|
||||
import PIL.Image
|
||||
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
get_boogu_image_pre_process_func,
|
||||
)
|
||||
|
||||
pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
|
||||
|
||||
image = PIL.Image.new("RGB", (1000, 500)) # (width, height)
|
||||
req = _make_diffusion_request({"prompt": "make it winter", "multi_modal_data": {"image": image}})
|
||||
out = pre(req)
|
||||
|
||||
ai = out.prompt["additional_information"]
|
||||
assert "prompt_image" in ai and "preprocessed_image" in ai
|
||||
|
||||
# VLM copy is a PIL image, downscaled, never upscaled, aligned to 16.
|
||||
prompt_image = ai["prompt_image"]
|
||||
assert isinstance(prompt_image, PIL.Image.Image)
|
||||
assert prompt_image.width % 16 == 0 and prompt_image.height % 16 == 0
|
||||
assert prompt_image.width <= 1000 and prompt_image.height <= 500
|
||||
assert max(prompt_image.width, prompt_image.height) <= 768
|
||||
|
||||
# VAE copy is a normalized [1, C, H, W] tensor aligned to 16.
|
||||
vae = ai["preprocessed_image"]
|
||||
assert isinstance(vae, torch.Tensor) and vae.ndim == 4 and vae.shape[0] == 1
|
||||
assert vae.shape[-1] % 16 == 0 and vae.shape[-2] % 16 == 0
|
||||
assert -1.0 <= float(vae.min()) and float(vae.max()) <= 1.0
|
||||
|
||||
# align_res: the request resolution follows the VAE-encoded reference dims.
|
||||
assert out.sampling_params.height == vae.shape[-2]
|
||||
assert out.sampling_params.width == vae.shape[-1]
|
||||
|
||||
|
||||
def test_pre_process_rejects_multiple_images(tmp_path):
|
||||
import PIL.Image
|
||||
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
get_boogu_image_pre_process_func,
|
||||
)
|
||||
|
||||
pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
|
||||
imgs = [PIL.Image.new("RGB", (64, 64)), PIL.Image.new("RGB", (64, 64))]
|
||||
req = _make_diffusion_request({"prompt": "combine", "multi_modal_data": {"image": imgs}})
|
||||
|
||||
with pytest.raises(ValueError, match="single reference image"):
|
||||
pre(req)
|
||||
|
||||
|
||||
def test_pre_process_single_image_in_list_is_accepted(tmp_path):
|
||||
import PIL.Image
|
||||
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
get_boogu_image_pre_process_func,
|
||||
)
|
||||
|
||||
pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
|
||||
req = _make_diffusion_request({"prompt": "edit", "multi_modal_data": {"image": [PIL.Image.new("RGB", (128, 128))]}})
|
||||
out = pre(req)
|
||||
assert "preprocessed_image" in out.prompt["additional_information"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Editing / TI2I: chat template + image-aware encoding
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_edit_encode_pipeline():
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
||||
SYSTEM_PROMPT_4_TI2I_UNIFIED,
|
||||
)
|
||||
|
||||
pipeline = _make_encode_pipeline()
|
||||
pipeline.SYSTEM_PROMPT_4_TI2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
|
||||
pipeline.SYSTEM_PROMPT_4_I2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_apply_chat_template_ti2i_places_image_before_text():
|
||||
import PIL.Image
|
||||
|
||||
pipeline = _make_edit_encode_pipeline()
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
|
||||
messages = pipeline._apply_chat_template("turn day into night", [image])
|
||||
assert messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_TI2I
|
||||
user_content = messages[1]["content"]
|
||||
# Image content comes first, then the instruction text.
|
||||
assert user_content[0]["type"] == "image"
|
||||
assert user_content[0]["image"] is image
|
||||
assert user_content[-1] == {"type": "text", "text": "turn day into night"}
|
||||
|
||||
# Empty instruction with an image selects the I2I system prompt.
|
||||
empty_messages = pipeline._apply_chat_template("", [image])
|
||||
assert empty_messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_I2I
|
||||
|
||||
|
||||
class _ImageAwareRecordingProcessor:
|
||||
"""Records whether reference images reached the processor."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def apply_chat_template(self, prompts, **kwargs):
|
||||
has_image = []
|
||||
for messages in prompts:
|
||||
user_content = messages[1]["content"]
|
||||
has_image.append(any(c.get("type") == "image" for c in user_content))
|
||||
self.calls.append({"prompts": prompts, "kwargs": kwargs, "has_image": has_image})
|
||||
batch = len(prompts)
|
||||
input_ids = torch.arange(batch * _SEQ_LEN, dtype=torch.long).view(batch, _SEQ_LEN)
|
||||
attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long)
|
||||
return {"input_ids": input_ids, "attention_mask": attention_mask}
|
||||
|
||||
|
||||
def test_encode_prompt_attaches_images_to_positive_only():
|
||||
import PIL.Image
|
||||
|
||||
pipeline = _make_edit_encode_pipeline()
|
||||
pipeline.processor = _ImageAwareRecordingProcessor()
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
|
||||
pipeline.encode_prompt(
|
||||
"add a rainbow",
|
||||
do_classifier_free_guidance=True,
|
||||
input_images=[[image]],
|
||||
)
|
||||
|
||||
# Positive call carries the image; the negative (CFG) call is image-free.
|
||||
assert pipeline.processor.calls[0]["has_image"] == [True]
|
||||
assert pipeline.processor.calls[1]["has_image"] == [False]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Editing / TI2I: reference-latent VAE encode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeRefVAE:
|
||||
"""Fake VAE whose ``encode`` yields a known latent for shape/scaling checks."""
|
||||
|
||||
dtype = torch.float32
|
||||
|
||||
def __init__(self, scaling_factor=2.0, shift_factor=0.5):
|
||||
self.config = SimpleNamespace(scaling_factor=scaling_factor, shift_factor=shift_factor)
|
||||
self._latent = torch.ones(1, 4, 3, 5)
|
||||
|
||||
def encode(self, img):
|
||||
dist = SimpleNamespace(sample=lambda generator=None: self._latent.clone())
|
||||
return SimpleNamespace(latent_dist=dist)
|
||||
|
||||
|
||||
def test_build_ref_latents_shape_and_scaling():
|
||||
pipeline = _make_encode_pipeline()
|
||||
pipeline.vae = _FakeRefVAE(scaling_factor=2.0, shift_factor=0.5)
|
||||
|
||||
preprocessed = torch.zeros(1, 3, 48, 80) # normalized image tensor
|
||||
ref_latents = pipeline._build_ref_latents([preprocessed], num_images_per_prompt=1, device=torch.device("cpu"))
|
||||
|
||||
assert len(ref_latents) == 1
|
||||
(sample_latents,) = ref_latents
|
||||
assert isinstance(sample_latents, list) and len(sample_latents) == 1
|
||||
latent = sample_latents[0]
|
||||
# squeeze(0) -> [C, H, W]; (1 - shift) * scaling = (1 - 0.5) * 2 = 1.0
|
||||
assert latent.shape == (4, 3, 5)
|
||||
assert torch.allclose(latent, torch.ones(4, 3, 5))
|
||||
|
||||
|
||||
def test_build_ref_latents_expands_per_output_and_handles_none():
|
||||
pipeline = _make_encode_pipeline()
|
||||
pipeline.vae = _FakeRefVAE()
|
||||
|
||||
preprocessed = torch.zeros(1, 3, 48, 80)
|
||||
ref_latents = pipeline._build_ref_latents([preprocessed, None], num_images_per_prompt=2, device=torch.device("cpu"))
|
||||
|
||||
# Two samples x 2 outputs each = 4 entries; the None sample stays None.
|
||||
assert len(ref_latents) == 4
|
||||
assert ref_latents[0] is ref_latents[1] # same sample repeated
|
||||
assert ref_latents[2] is None and ref_latents[3] is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Editing / TI2I: forward CFG branch selection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _RecordingRefTransformer(_FakeTransformer):
|
||||
"""Counts predictions per step and records the reference-latent argument."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
|
||||
self.calls.append(kwargs.get("ref_image_hidden_states"))
|
||||
return torch.zeros_like(latents)
|
||||
|
||||
|
||||
class _EditForwardVAE(_FakeDecodeVAE):
|
||||
"""Adds a fake ``encode`` so the editing forward path can build ref latents."""
|
||||
|
||||
def encode(self, img):
|
||||
dist = SimpleNamespace(sample=lambda generator=None: torch.zeros(1, 4, 8, 8))
|
||||
return SimpleNamespace(latent_dist=dist)
|
||||
|
||||
|
||||
def _make_edit_forward_pipeline():
|
||||
pipeline = _make_edit_encode_pipeline()
|
||||
# Image-aware processor: reference images appear as ``{"type": "image"}``
|
||||
# content entries, which the default text-first fake cannot parse.
|
||||
pipeline.processor = _ImageAwareRecordingProcessor()
|
||||
pipeline.transformer = _RecordingRefTransformer()
|
||||
pipeline.scheduler = _FakeScheduler()
|
||||
pipeline.vae = _EditForwardVAE()
|
||||
pipeline.vae_scale_factor = 8
|
||||
pipeline.default_sample_size = 128
|
||||
return pipeline
|
||||
|
||||
|
||||
def _make_edit_request(**sampling_overrides):
|
||||
import PIL.Image
|
||||
|
||||
image = PIL.Image.new("RGB", (64, 64))
|
||||
prompt = {
|
||||
"prompt": "make it winter",
|
||||
"additional_information": {
|
||||
"prompt_image": image,
|
||||
"preprocessed_image": torch.zeros(1, 3, 64, 64),
|
||||
},
|
||||
}
|
||||
sampling = _sampling(**sampling_overrides)
|
||||
return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
|
||||
|
||||
|
||||
def _sampling(**overrides):
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
return OmniDiffusionSamplingParams(**overrides)
|
||||
|
||||
|
||||
def _count_per_step(calls, num_steps):
|
||||
assert len(calls) % num_steps == 0
|
||||
return len(calls) // num_steps
|
||||
|
||||
|
||||
def test_forward_ti2i_text_only_two_predictions_with_ref():
|
||||
pipeline = _make_edit_forward_pipeline()
|
||||
num_steps = 2
|
||||
req = _make_edit_request(height=64, width=64, num_inference_steps=num_steps, guidance_scale=5.0)
|
||||
req.sampling_params.guidance_scale_provided = True
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
# Text-only ti2i: cond+ref and neg+ref -> 2 predictions/step, both carry ref.
|
||||
assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
|
||||
assert all(ref is not None for ref in pipeline.transformer.calls)
|
||||
|
||||
|
||||
def test_forward_ti2i_double_guidance_three_predictions():
|
||||
pipeline = _make_edit_forward_pipeline()
|
||||
num_steps = 2
|
||||
req = _make_edit_request(
|
||||
height=64, width=64, num_inference_steps=num_steps, guidance_scale=5.0, guidance_scale_2=2.0
|
||||
)
|
||||
req.sampling_params.guidance_scale_provided = True
|
||||
req.sampling_params.guidance_scale_2_provided = True
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
# Double guidance: cond+ref, neg+ref, neg+no-ref -> 3 predictions/step.
|
||||
assert _count_per_step(pipeline.transformer.calls, num_steps) == 3
|
||||
# Exactly one of the three per step drops the reference (neg+no-ref).
|
||||
per_step = [pipeline.transformer.calls[i : i + 3] for i in range(0, len(pipeline.transformer.calls), 3)]
|
||||
for step_calls in per_step:
|
||||
assert sum(ref is None for ref in step_calls) == 1
|
||||
|
||||
|
||||
def test_forward_ti2i_image_only_two_predictions_drop_ref():
|
||||
pipeline = _make_edit_forward_pipeline()
|
||||
num_steps = 2
|
||||
# text guidance 1.0 (off, provided) + image guidance 2.0 (provided).
|
||||
req = _make_edit_request(
|
||||
height=64, width=64, num_inference_steps=num_steps, guidance_scale=1.0, guidance_scale_2=2.0
|
||||
)
|
||||
req.sampling_params.guidance_scale_provided = True
|
||||
req.sampling_params.guidance_scale_2_provided = True
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
# Image-only ti2i: cond+ref and cond+no-ref -> 2 predictions/step.
|
||||
assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
|
||||
per_step = [pipeline.transformer.calls[i : i + 2] for i in range(0, len(pipeline.transformer.calls), 2)]
|
||||
for step_calls in per_step:
|
||||
assert sum(ref is None for ref in step_calls) == 1
|
||||
|
||||
|
||||
def test_forward_ti2i_no_guidance_single_prediction_with_ref():
|
||||
pipeline = _make_edit_forward_pipeline()
|
||||
num_steps = 2
|
||||
# Both guidances off/unprovided -> image guidance forced to 1.0, no CFG.
|
||||
req = _make_edit_request(height=64, width=64, num_inference_steps=num_steps, guidance_scale=1.0)
|
||||
req.sampling_params.guidance_scale_provided = True
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
assert _count_per_step(pipeline.transformer.calls, num_steps) == 1
|
||||
assert all(ref is not None for ref in pipeline.transformer.calls)
|
||||
|
||||
|
||||
def test_forward_image_guidance_ignored_without_reference():
|
||||
# guidance_scale_2 is set but there is no reference image (t2i request);
|
||||
# image guidance must be forced off so this stays plain t2i CFG.
|
||||
pipeline = _make_edit_forward_pipeline()
|
||||
num_steps = 2
|
||||
sampling = _sampling(height=64, width=64, num_inference_steps=num_steps, guidance_scale=5.0, guidance_scale_2=2.0)
|
||||
sampling.guidance_scale_provided = True
|
||||
sampling.guidance_scale_2_provided = True
|
||||
req = SimpleNamespace(prompts=[{"prompt": "a cat"}], sampling_params=sampling)
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
# t2i text CFG: cond + uncond -> 2 predictions/step, no ref anywhere.
|
||||
assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
|
||||
assert all(ref is None for ref in pipeline.transformer.calls)
|
||||
@@ -0,0 +1,94 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Online serving tests for ``Boogu/Boogu-Image-0.1-Base`` (text-to-image).
|
||||
|
||||
- ``test_text_to_image_001``: single chat request, default server (``_get_default_case``).
|
||||
- ``test_batch_001``: concurrent prompts via ``send_diffusion_request([cfg0, cfg1, ...])`` — one
|
||||
dict per prompt; each entry carries its own ``messages`` / ``negative_prompt`` (see
|
||||
``TEST_PROMPTS``).
|
||||
|
||||
Boogu-Image's pipeline reads ``sp.guidance_scale`` (not ``true_cfg_scale``), so the
|
||||
``extra_body`` uses ``guidance_scale``. These are smoke tests (minimal steps); numeric
|
||||
parity vs. the Diffusers baseline is covered separately (see the support plan, step 14).
|
||||
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image.py \
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Boogu/Boogu-Image-0.1-Base"
|
||||
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
|
||||
NEGATIVE_PROMPT = ""
|
||||
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
TEST_PROMPTS: list[dict[str, str]] = [
|
||||
{"prompt": "a cup of coffee on a table", "negative_prompt": ""},
|
||||
{"prompt": "a toy dinosaur on a sandy beach", "negative_prompt": ""},
|
||||
{"prompt": "a futuristic city skyline at sunset", "negative_prompt": ""},
|
||||
]
|
||||
|
||||
|
||||
def _get_default_case(model: str):
|
||||
"""Return a single default ``OmniServerParams`` row (no extra ``server_args``)."""
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(model=model),
|
||||
id="default",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)
|
||||
def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""Default Boogu-Image T2I smoke (single ``default`` server config)."""
|
||||
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)
|
||||
def test_batch_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""Concurrent T2I: one ``request_config`` dict per prompt (``send_diffusion_request`` list mode)."""
|
||||
request_config = [
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"messages": dummy_messages_from_mix_data(content_text=prompt["prompt"]),
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": prompt["negative_prompt"],
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
for prompt in TEST_PROMPTS
|
||||
]
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
@@ -0,0 +1,156 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Online serving tests for Boogu-Image-0.1-Edit (image-to-image via chat completions).
|
||||
|
||||
The Edit checkpoint shares the ``BooguImagePipeline`` class with the Base
|
||||
text-to-image checkpoint; the editing (TI2I) path activates when the request
|
||||
carries a reference image. Text guidance rides on ``guidance_scale`` (upstream
|
||||
default 4.0), and image guidance rides on ``guidance_scale_2`` (upstream default
|
||||
1.0 = off; a value > 1 enables the double-guidance path). Only a single
|
||||
reference image is supported.
|
||||
|
||||
- ``test_single_image_to_image_001``: one reference image, text guidance only.
|
||||
- ``test_double_guidance_001``: one reference image, text + image guidance.
|
||||
- ``test_images_edits_endpoint_001``: same text-guided edit via ``/v1/images/edits``.
|
||||
- ``test_images_edits_endpoint_double_guidance_001``: double guidance via ``/v1/images/edits``.
|
||||
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image_edit.py \
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from tests.helpers.runtime import (
|
||||
OmniServer,
|
||||
OmniServerParams,
|
||||
OpenAIClientHandler,
|
||||
dummy_messages_from_mix_data,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
EDIT_PROMPT = "Change the style to a colored pencil drawing."
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
"""Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
),
|
||||
id="default",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Single-reference text-guided edit smoke for ``Boogu/Boogu-Image-0.1-Edit``."""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"num_inference_steps": 2,
|
||||
# Boogu text guidance (upstream default 4.0); > 1 enables CFG.
|
||||
"guidance_scale": 5.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Single-reference double-guidance edit (text + image guidance) smoke."""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 5.0,
|
||||
# Image guidance > 1 activates the 3-prediction double-guidance path.
|
||||
"guidance_scale_2": 2.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
def _edits_request_config(model: str, *, guidance_scale_2: float | None = None) -> dict:
|
||||
"""Multipart ``/v1/images/edits`` config: one synthetic reference + smoke-depth params."""
|
||||
image_bytes = base64.b64decode(generate_synthetic_image(512, 512)["base64"])
|
||||
data = {
|
||||
"model": model,
|
||||
"prompt": EDIT_PROMPT,
|
||||
"num_inference_steps": 2,
|
||||
# Boogu text guidance (upstream default 4.0); > 1 enables CFG.
|
||||
"guidance_scale": 5.0,
|
||||
"seed": 42,
|
||||
}
|
||||
if guidance_scale_2 is not None:
|
||||
data["guidance_scale_2"] = guidance_scale_2
|
||||
return {
|
||||
"data": data,
|
||||
"files": {"image": ("ref.jpg", image_bytes, "image/jpeg")},
|
||||
"timeout": 300,
|
||||
}
|
||||
|
||||
|
||||
def _assert_edits_response_has_image(responses) -> None:
|
||||
assert responses and responses[0].success, "Expected a successful /v1/images/edits response"
|
||||
body = responses[0].json_body
|
||||
assert body and body.get("data"), "No image data returned by /v1/images/edits"
|
||||
assert body["data"][0].get("b64_json"), "Missing b64_json in /v1/images/edits response"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_images_edits_endpoint_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Single-reference text-guided edit via the OpenAI-compatible ``/v1/images/edits`` endpoint."""
|
||||
responses = openai_client.send_images_edits_http_request(_edits_request_config(omni_server.model))
|
||||
_assert_edits_response_has_image(responses)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_images_edits_endpoint_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Double-guidance edit (text + image guidance) via ``/v1/images/edits``."""
|
||||
responses = openai_client.send_images_edits_http_request(
|
||||
_edits_request_config(omni_server.model, guidance_scale_2=2.0)
|
||||
)
|
||||
_assert_edits_response_has_image(responses)
|
||||
@@ -0,0 +1,107 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Edit`` (image editing / TI2I).
|
||||
|
||||
The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image
|
||||
text-guided editing and the text+image double-guidance path. This expansion
|
||||
module adds the two edit paths smoke does not exercise:
|
||||
|
||||
- ``test_image_only_guidance`` — ``guidance_scale=1.0`` (text CFG off) with
|
||||
``guidance_scale_2=2.0`` (image CFG on): the image-only branch of upstream's
|
||||
CFG-branch priority (double > text-only > image-only > t2i), which no existing
|
||||
e2e test reaches.
|
||||
- ``test_align_res_output_size`` — a non-square reference with no requested
|
||||
``height`` / ``width``: asserts the output resolution follows the reference
|
||||
(upstream ``align_res``, on by default for a single reference).
|
||||
|
||||
The Edit checkpoint shares ``BooguImagePipeline`` with the Base checkpoint; the
|
||||
TI2I path activates when the request carries a reference image. Only a single
|
||||
reference image is supported. Cases stay smoke-depth (``num_inference_steps=2``).
|
||||
CPU offload, Cache-DiT, and multi-GPU parallelism are not yet supported for this
|
||||
model and are intentionally omitted.
|
||||
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image_edit_expansion.py \
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Boogu/Boogu-Image-0.1-Edit"
|
||||
EDIT_PROMPT = "Change the style to a colored pencil drawing."
|
||||
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
"""Return the single-card server configuration supported for Boogu-Image-Edit today."""
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(model=model),
|
||||
id="default",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""Image-only guidance: text CFG off (1.0), image CFG on (2.0) -> image-only CFG branch."""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"num_inference_steps": 2,
|
||||
# Text guidance off, image guidance on -> upstream's image-only branch.
|
||||
"guidance_scale": 1.0,
|
||||
"guidance_scale_2": 2.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""Output resolution follows the reference image (upstream ``align_res``).
|
||||
|
||||
No ``height`` / ``width`` is requested. The reference is 512x384 — both
|
||||
already multiples of 32 (>= the ``vae_scale_factor * 2`` alignment) and well
|
||||
under the VAE ``max_pixels`` / ``max_side`` limits — so the align-res target
|
||||
equals the reference dims regardless of the exact VAE scale factor.
|
||||
"""
|
||||
ref_width, ref_height = 512, 384
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(ref_width, ref_height)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 5.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
responses = openai_client.send_diffusion_request(request_config)
|
||||
|
||||
assert responses and responses[0].images, "No image returned for the edit request"
|
||||
out_image = responses[0].images[0]
|
||||
assert (out_image.width, out_image.height) == (ref_width, ref_height), (
|
||||
f"align_res should size the output to the reference dims {(ref_width, ref_height)}; "
|
||||
f"got {(out_image.width, out_image.height)}"
|
||||
)
|
||||
@@ -0,0 +1,127 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Base`` (text-to-image).
|
||||
|
||||
The nightly smoke module (``test_boogu_image.py``) already covers the default
|
||||
512x512 single request and a concurrent batch. This expansion module adds the
|
||||
request-level paths that smoke does not exercise, run against every currently
|
||||
supported single-GPU server configuration:
|
||||
|
||||
Server rows (single card):
|
||||
- ``default`` — no extra ``server_args``.
|
||||
- ``vae_slicing_tiling`` — ``--vae-use-slicing --vae-use-tiling`` (the recipe's
|
||||
OOM mitigation; applied generically in ``registry.initialize_model``).
|
||||
|
||||
Cases (one per test, each parametrized over both server rows):
|
||||
- ``test_non_square_resolution`` — 768x1024, exercises the patchify / RoPE path
|
||||
at a non-square aspect ratio.
|
||||
- ``test_cfg_off`` — ``guidance_scale=1.0``, exercises the no-CFG denoise branch
|
||||
(single prediction per step).
|
||||
- ``test_high_resolution`` — 1024x1024, the recipe-recommended resolution.
|
||||
|
||||
Boogu-Image reads ``sp.guidance_scale`` (not ``true_cfg_scale``). These stay
|
||||
smoke-depth (``num_inference_steps=2``); numeric parity vs. Diffusers is covered
|
||||
by the local parity harness (support plan, step 14). CPU offload, Cache-DiT, and
|
||||
multi-GPU parallelism are not yet supported for this model and are intentionally
|
||||
omitted (add rows here when those features land).
|
||||
|
||||
From ``tests/``::
|
||||
|
||||
pytest -s -v e2e/online_serving/test_boogu_image_expansion.py \
|
||||
-m "full_model and diffusion and H100" --run-level=full_model
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "Boogu/Boogu-Image-0.1-Base"
|
||||
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
|
||||
NEGATIVE_PROMPT = ""
|
||||
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
"""Return the single-card server configurations supported for Boogu-Image today."""
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(model=model),
|
||||
id="default",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=["--vae-use-slicing", "--vae-use-tiling"],
|
||||
),
|
||||
id="vae_slicing_tiling",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_non_square_resolution(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""Non-square (768x1024) T2I exercises the patchify / RoPE path off the square grid."""
|
||||
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 768,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_cfg_off(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""``guidance_scale=1.0`` exercises the no-CFG denoise branch (one prediction per step)."""
|
||||
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"guidance_scale": 1.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_high_resolution(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
"""1024x1024 T2I at the recipe-recommended resolution."""
|
||||
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 2,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
@@ -1562,6 +1562,7 @@ def test_image_edit_parameter_pass(async_omni_test_client):
|
||||
"output_format": "jpeg",
|
||||
"num_inference_steps": 20,
|
||||
"guidance_scale": 8.0,
|
||||
"guidance_scale_2": 2.0,
|
||||
"seed": 1234,
|
||||
"negative_prompt": "negative",
|
||||
"n": 2,
|
||||
@@ -1576,6 +1577,7 @@ def test_image_edit_parameter_pass(async_omni_test_client):
|
||||
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.guidance_scale_2 == 2.0
|
||||
assert captured_sampling_params.seed == 1234
|
||||
assert captured_sampling_params.num_outputs_per_prompt == 2
|
||||
assert captured_sampling_params.width == 16
|
||||
|
||||
@@ -45,6 +45,7 @@ EXCLUDED_MODELS = [
|
||||
"WanI2VDMD2Pipeline",
|
||||
"LongCatImagePipeline",
|
||||
"BagelPipeline",
|
||||
"BooguImagePipeline",
|
||||
"LancePipeline",
|
||||
"MingImagePipeline",
|
||||
"InternVLAA1Pipeline",
|
||||
|
||||
@@ -15,6 +15,8 @@ class DiffusionModelMetadata:
|
||||
QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
|
||||
# Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3.
|
||||
HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3
|
||||
# Boogu-Image editing (TI2I) supports a single reference image for now.
|
||||
BOOGU_IMAGE_MAX_INPUT_IMAGES = 1
|
||||
|
||||
|
||||
_DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
|
||||
@@ -26,6 +28,13 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
|
||||
supports_multimodal_inputs=True,
|
||||
max_multimodal_image_inputs=HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
|
||||
),
|
||||
# Shared by the Base (text-to-image) and Edit (TI2I) checkpoints, which use
|
||||
# the same ``BooguImagePipeline`` class. Text-to-image requests simply carry
|
||||
# no reference image.
|
||||
"BooguImagePipeline": DiffusionModelMetadata(
|
||||
supports_multimodal_inputs=True,
|
||||
max_multimodal_image_inputs=BOOGU_IMAGE_MAX_INPUT_IMAGES,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
from .boogu_image_transformer import BooguImageTransformer2DModel
|
||||
from .image_processor import BooguImageProcessor
|
||||
from .pipeline_boogu_image import BooguImagePipeline, get_boogu_image_pre_process_func
|
||||
from .scheduling_flow_match_euler_discrete_time_shifting import FlowMatchEulerDiscreteScheduler
|
||||
|
||||
__all__ = [
|
||||
"BooguImagePipeline",
|
||||
"BooguImageProcessor",
|
||||
"BooguImageTransformer2DModel",
|
||||
"FlowMatchEulerDiscreteScheduler",
|
||||
"get_boogu_image_pre_process_func",
|
||||
]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,213 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Native port of the upstream ``BooguImageProcessor``.
|
||||
|
||||
Ported from the external ``boogu`` package
|
||||
(``boogu/pipelines/image_processor.py``) so the native vLLM-Omni pipeline can
|
||||
reproduce Boogu-Image's reference-image and VLM preprocessing without requiring
|
||||
``pip install boogu-image``.
|
||||
|
||||
Only the two Boogu-specific methods are ported: ``get_new_height_width`` (the
|
||||
``max_pixels`` / ``max_side_length`` aware downscale-only resize) and
|
||||
``preprocess`` (which routes through it). Everything else is inherited from the
|
||||
stock diffusers ``VaeImageProcessor``.
|
||||
"""
|
||||
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
from diffusers.image_processor import (
|
||||
PipelineImageInput,
|
||||
VaeImageProcessor,
|
||||
is_valid_image_imagelist,
|
||||
)
|
||||
|
||||
|
||||
class BooguImageProcessor(VaeImageProcessor):
|
||||
"""VaeImageProcessor variant with Boogu-Image pixel/side-length constraints.
|
||||
|
||||
Resizing never upscales (the ratio is clamped to ``<= 1``) and always aligns
|
||||
the target height/width to multiples of ``vae_scale_factor``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
do_resize: bool = True,
|
||||
vae_scale_factor: int = 16,
|
||||
resample: str = "lanczos",
|
||||
max_pixels: int | None = None,
|
||||
max_side_length: int | None = None,
|
||||
do_normalize: bool = True,
|
||||
do_binarize: bool = False,
|
||||
do_convert_rgb: bool = False,
|
||||
do_convert_grayscale: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
do_resize=do_resize,
|
||||
vae_scale_factor=vae_scale_factor,
|
||||
resample=resample,
|
||||
do_normalize=do_normalize,
|
||||
do_binarize=do_binarize,
|
||||
do_convert_rgb=do_convert_rgb,
|
||||
do_convert_grayscale=do_convert_grayscale,
|
||||
)
|
||||
self.max_pixels = max_pixels
|
||||
self.max_side_length = max_side_length
|
||||
|
||||
def get_new_height_width(
|
||||
self,
|
||||
image: PIL.Image.Image | np.ndarray | torch.Tensor,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
max_pixels: int | None = None,
|
||||
max_side_length: int | None = None,
|
||||
) -> tuple[int, int]:
|
||||
"""Return target ``(height, width)`` after downscale + alignment.
|
||||
|
||||
Faithful port of upstream ``BooguImageProcessor.get_new_height_width``.
|
||||
"""
|
||||
if height is None:
|
||||
if isinstance(image, PIL.Image.Image):
|
||||
height = image.height
|
||||
elif isinstance(image, torch.Tensor):
|
||||
height = image.shape[2]
|
||||
else:
|
||||
height = image.shape[1]
|
||||
|
||||
if width is None:
|
||||
if isinstance(image, PIL.Image.Image):
|
||||
width = image.width
|
||||
elif isinstance(image, torch.Tensor):
|
||||
width = image.shape[3]
|
||||
else:
|
||||
width = image.shape[2]
|
||||
|
||||
if max_side_length is None:
|
||||
max_side_length = self.max_side_length
|
||||
|
||||
if max_pixels is None:
|
||||
max_pixels = self.max_pixels
|
||||
|
||||
if height <= 0 or width <= 0:
|
||||
raise ValueError(f"Image height and width must be positive, got height={height}, width={width}")
|
||||
|
||||
max_side_length_ratio = 1.0
|
||||
if max_side_length is not None:
|
||||
if height > width:
|
||||
max_side_length_ratio = max_side_length / height
|
||||
else:
|
||||
max_side_length_ratio = max_side_length / width
|
||||
|
||||
cur_pixels = height * width
|
||||
max_pixels_ratio = (max_pixels / cur_pixels) ** 0.5 if max_pixels is not None else 1.0
|
||||
# Clamp ratio to <=1 to avoid upscaling input images in preprocessing.
|
||||
ratio = min(max_pixels_ratio, max_side_length_ratio, 1.0)
|
||||
|
||||
new_height, new_width = (
|
||||
int(height * ratio) // self.config.vae_scale_factor * self.config.vae_scale_factor,
|
||||
int(width * ratio) // self.config.vae_scale_factor * self.config.vae_scale_factor,
|
||||
)
|
||||
return new_height, new_width
|
||||
|
||||
def preprocess(
|
||||
self,
|
||||
image: PipelineImageInput,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
max_pixels: int | None = None,
|
||||
max_side_length: int | None = None,
|
||||
resize_mode: str = "default",
|
||||
crops_coords: tuple[int, int, int, int] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Preprocess an image into a normalized ``[B, C, H, W]`` tensor.
|
||||
|
||||
Faithful port of upstream ``BooguImageProcessor.preprocess`` (PixArt-style
|
||||
downscale). Only the PIL branch is exercised by the native pipeline, but
|
||||
the numpy/tensor branches are kept for parity.
|
||||
"""
|
||||
supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
|
||||
|
||||
if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3:
|
||||
if isinstance(image, torch.Tensor):
|
||||
image = image.unsqueeze(1)
|
||||
else:
|
||||
if image.shape[-1] == 1:
|
||||
image = np.expand_dims(image, axis=0)
|
||||
else:
|
||||
image = np.expand_dims(image, axis=-1)
|
||||
|
||||
if isinstance(image, list) and isinstance(image[0], np.ndarray) and image[0].ndim == 4:
|
||||
warnings.warn(
|
||||
"Passing `image` as a list of 4d np.ndarray is deprecated."
|
||||
"Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray",
|
||||
FutureWarning,
|
||||
)
|
||||
image = np.concatenate(image, axis=0)
|
||||
if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4:
|
||||
warnings.warn(
|
||||
"Passing `image` as a list of 4d torch.Tensor is deprecated."
|
||||
"Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor",
|
||||
FutureWarning,
|
||||
)
|
||||
image = torch.cat(image, axis=0)
|
||||
|
||||
if not is_valid_image_imagelist(image):
|
||||
supported_formats_str = ", ".join(str(x) for x in supported_formats)
|
||||
raise ValueError(f"Input is in incorrect format. Currently, we only support {supported_formats_str}")
|
||||
|
||||
if not isinstance(image, list):
|
||||
image = [image]
|
||||
|
||||
if isinstance(image[0], PIL.Image.Image):
|
||||
if crops_coords is not None:
|
||||
image = [i.crop(crops_coords) for i in image]
|
||||
if self.config.do_resize:
|
||||
height, width = self.get_new_height_width(image[0], height, width, max_pixels, max_side_length)
|
||||
image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image]
|
||||
if self.config.do_convert_rgb:
|
||||
image = [self.convert_to_rgb(i) for i in image]
|
||||
elif self.config.do_convert_grayscale:
|
||||
image = [self.convert_to_grayscale(i) for i in image]
|
||||
image = self.pil_to_numpy(image)
|
||||
image = self.numpy_to_pt(image)
|
||||
|
||||
elif isinstance(image[0], np.ndarray):
|
||||
image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)
|
||||
image = self.numpy_to_pt(image)
|
||||
height, width = self.get_new_height_width(image, height, width, max_pixels, max_side_length)
|
||||
if self.config.do_resize:
|
||||
image = self.resize(image, height, width)
|
||||
|
||||
elif isinstance(image[0], torch.Tensor):
|
||||
image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
|
||||
|
||||
if self.config.do_convert_grayscale and image.ndim == 3:
|
||||
image = image.unsqueeze(1)
|
||||
|
||||
channel = image.shape[1]
|
||||
if channel == self.config.vae_latent_channels:
|
||||
return image
|
||||
|
||||
height, width = self.get_new_height_width(image, height, width, max_pixels, max_side_length)
|
||||
if self.config.do_resize:
|
||||
image = self.resize(image, height, width)
|
||||
|
||||
do_normalize = self.config.do_normalize
|
||||
if do_normalize and image.min() < 0:
|
||||
warnings.warn(
|
||||
"Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range "
|
||||
f"for image tensor is [0,1] when passing as pytorch tensor or numpy Array. You passed `image` with "
|
||||
f"value range [{image.min()},{image.max()}]",
|
||||
FutureWarning,
|
||||
)
|
||||
do_normalize = False
|
||||
if do_normalize:
|
||||
image = self.normalize(image)
|
||||
|
||||
if self.config.do_binarize:
|
||||
image = self.binarize(image)
|
||||
|
||||
return image
|
||||
@@ -0,0 +1,793 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Native vLLM-Omni pipeline for Boogu-Image-0.1.
|
||||
|
||||
Ported from the upstream ``boogu`` package
|
||||
(``boogu/pipelines/boogu/pipeline_boogu.py``) with the following changes:
|
||||
|
||||
- Diffusers ``DiffusionPipeline``/``register_modules`` machinery replaced by a
|
||||
plain ``nn.Module`` constructed from ``OmniDiffusionConfig`` (components are
|
||||
loaded from the checkpoint subfolders; transformer weights arrive later via
|
||||
``weights_sources`` + ``load_weights``).
|
||||
- Upstream ``encode_instruction`` is exposed as ``encode_prompt`` (the
|
||||
vLLM-Omni convention, also hooked by the prompt-embed cache); only the
|
||||
text-to-image path is ported. Instruction rewriting, prompt tuning,
|
||||
reference-image (TI2I/I2I) encoding, double guidance, and vision-token
|
||||
stripping are not ported.
|
||||
- ``forward()`` (denoise loop + VAE decode) is intentionally not implemented
|
||||
yet; it lands together with the post-process registration.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from typing import ClassVar, cast
|
||||
|
||||
import PIL.Image
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageDoubleStreamRotaryPosEmbed,
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
from vllm_omni.diffusion.models.boogu_image.image_processor import BooguImageProcessor
|
||||
from vllm_omni.diffusion.models.boogu_image.scheduling_flow_match_euler_discrete_time_shifting import (
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
)
|
||||
from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Reference-image preprocessing limits (upstream ``BooguImagePipeline.__call__``
|
||||
# defaults). The VLM copy is aggressively downscaled for the Qwen3VL encoder;
|
||||
# the VAE copy keeps near-native resolution for the reference latents.
|
||||
_MAX_VLM_INPUT_PIL_PIXELS = 384 * 384
|
||||
_MAX_VLM_INPUT_PIL_SIDE_LENGTH = 384 * 2
|
||||
_MAX_INPUT_IMAGE_PIXELS = 2048 * 2048
|
||||
_MAX_INPUT_IMAGE_SIDE_LENGTH = 2048 * 2
|
||||
|
||||
|
||||
def _load_vae_scale_factor(model_path: str) -> int:
|
||||
vae_config_path = os.path.join(model_path, "vae/config.json")
|
||||
try:
|
||||
with open(vae_config_path) as f:
|
||||
vae_config = json.load(f)
|
||||
except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
raise RuntimeError(f"Failed to load Boogu VAE config from {vae_config_path}: {exc}") from exc
|
||||
|
||||
if "block_out_channels" not in vae_config:
|
||||
return 8
|
||||
return 2 ** (len(vae_config["block_out_channels"]) - 1)
|
||||
|
||||
|
||||
def get_boogu_image_post_process_func(od_config: OmniDiffusionConfig):
|
||||
"""Build the post-process callable that converts decoded tensors to images.
|
||||
|
||||
Upstream ``BooguImageProcessor`` only customizes *pre*-processing; the
|
||||
``postprocess`` path is inherited from the stock diffusers
|
||||
``VaeImageProcessor``, so we reuse it directly here.
|
||||
"""
|
||||
model_name = od_config.model
|
||||
if os.path.exists(model_name):
|
||||
model_path = model_name
|
||||
else:
|
||||
model_path = download_weights_from_hf_specific(model_name, None, ["*"])
|
||||
|
||||
vae_scale_factor = _load_vae_scale_factor(model_path)
|
||||
|
||||
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
|
||||
|
||||
def post_process_func(images: torch.Tensor):
|
||||
return image_processor.postprocess(images)
|
||||
|
||||
return post_process_func
|
||||
|
||||
|
||||
def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
"""Build the pre-process callable for Boogu-Image reference (edit) input.
|
||||
|
||||
Text-to-image requests carry no image and are passed through unchanged (the
|
||||
Base checkpoint shares this pipeline class). Edit (TI2I) requests carry a
|
||||
single reference PIL image on ``prompt["multi_modal_data"]["image"]``; it is
|
||||
resized twice — once for the Qwen3VL encoder (``prompt_image``) and once for
|
||||
the VAE reference latents (``preprocessed_image``) — and stashed in
|
||||
``additional_information`` for ``forward`` to consume. Mirrors upstream
|
||||
``preprocess_vlm_input_pil_images`` + ``prepare_image``.
|
||||
|
||||
For a single reference image, upstream ``align_res`` (default ``True``)
|
||||
derives the output resolution from the VAE-encoded reference dimensions, so
|
||||
the request height/width are overwritten accordingly.
|
||||
"""
|
||||
model_name = od_config.model
|
||||
if os.path.exists(model_name):
|
||||
model_path = model_name
|
||||
else:
|
||||
model_path = download_weights_from_hf_specific(model_name, None, ["*"])
|
||||
|
||||
vae_scale_factor = _load_vae_scale_factor(model_path)
|
||||
|
||||
# Upstream builds ``BooguImageProcessor(vae_scale_factor=vae_scale_factor*2)``
|
||||
# so all resize targets align to multiples of ``vae_scale_factor * 2``.
|
||||
image_processor = BooguImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_resize=True)
|
||||
|
||||
def pre_process_func(request: OmniDiffusionRequest):
|
||||
prompt = request.prompt
|
||||
if isinstance(prompt, str):
|
||||
# Plain-text prompt cannot carry an image -> text-to-image, no-op.
|
||||
return request
|
||||
|
||||
multi_modal_data = prompt.get("multi_modal_data") or {}
|
||||
raw_image = multi_modal_data.get("image")
|
||||
if not raw_image:
|
||||
# No reference image -> text-to-image, no-op (Base checkpoint).
|
||||
return request
|
||||
|
||||
if isinstance(raw_image, list):
|
||||
if len(raw_image) > 1:
|
||||
raise ValueError(f"Boogu-Image editing supports a single reference image; received {len(raw_image)}.")
|
||||
raw_image = raw_image[0]
|
||||
|
||||
if isinstance(raw_image, str):
|
||||
image = PIL.Image.open(raw_image)
|
||||
else:
|
||||
image = cast(PIL.Image.Image, raw_image)
|
||||
image = image.convert("RGB")
|
||||
|
||||
if "additional_information" not in prompt:
|
||||
prompt["additional_information"] = {}
|
||||
|
||||
# VLM-resized copy (PIL) for the Qwen3VL instruction encoder.
|
||||
vlm_height, vlm_width = image_processor.get_new_height_width(
|
||||
image, None, None, _MAX_VLM_INPUT_PIL_PIXELS, _MAX_VLM_INPUT_PIL_SIDE_LENGTH
|
||||
)
|
||||
prompt_image = image_processor.resize(image, vlm_height, vlm_width)
|
||||
|
||||
# VAE-ready copy (normalized [1, C, H, W] tensor) for reference latents.
|
||||
preprocessed_image = image_processor.preprocess(
|
||||
image, max_pixels=_MAX_INPUT_IMAGE_PIXELS, max_side_length=_MAX_INPUT_IMAGE_SIDE_LENGTH
|
||||
)
|
||||
|
||||
# align_res: single-image output resolution follows the reference dims.
|
||||
request.sampling_params.height = int(preprocessed_image.shape[-2])
|
||||
request.sampling_params.width = int(preprocessed_image.shape[-1])
|
||||
|
||||
prompt["additional_information"]["preprocessed_image"] = preprocessed_image
|
||||
prompt["additional_information"]["prompt_image"] = prompt_image
|
||||
request.prompt = prompt
|
||||
return request
|
||||
|
||||
return pre_process_func
|
||||
|
||||
|
||||
# System prompts matching upstream dataset logic (ported verbatim from
|
||||
# ``BooguImagePipeline.__init__``).
|
||||
SYSTEM_PROMPT_4_TI2I_UNIFIED = (
|
||||
"Describe the key features of the input image (color, shape, size, texture, objects, background), "
|
||||
"then explain how the user's text instruction should alter or modify the image. Generate a new image "
|
||||
"that meets the user's requirements while maintaining consistency with the original input where appropriate."
|
||||
)
|
||||
SYSTEM_PROMPT_4_T2I_UNIFIED = (
|
||||
"You are a helpful assistant that generates high-quality images based on user instructions. "
|
||||
"The instructions are as follows."
|
||||
)
|
||||
|
||||
|
||||
class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, SupportImageInput):
|
||||
"""Boogu-Image text-to-image and image-editing (TI2I) pipeline.
|
||||
|
||||
Native vLLM-Omni implementation. A request with a reference image (edit /
|
||||
TI2I) is served by the same class as text-to-image; the reference latents
|
||||
and Qwen3VL image tokens are threaded through ``forward`` and the ported
|
||||
transformer's reference-image refiner path.
|
||||
"""
|
||||
|
||||
supports_request_batch = False
|
||||
|
||||
support_image_input: ClassVar[bool] = True
|
||||
color_format: ClassVar[str] = "RGB"
|
||||
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["mllm"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.od_config = od_config
|
||||
self._raise_unsupported_features()
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=od_config.model,
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
fall_back_to_pt=True,
|
||||
)
|
||||
]
|
||||
|
||||
self._execution_device = get_local_device()
|
||||
model = od_config.model
|
||||
local_files_only = os.path.exists(model)
|
||||
|
||||
# See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
|
||||
# race; prefetch the whole component set before any from_pretrained.
|
||||
boogu_subfolders = ["scheduler", "vae", "mllm", "processor"]
|
||||
prefetch_subfolders(model, boogu_subfolders, local_files_only=local_files_only)
|
||||
|
||||
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
||||
model, subfolder="scheduler", local_files_only=local_files_only
|
||||
)
|
||||
|
||||
mllm = from_pretrained_with_prefetch(
|
||||
Qwen3VLForConditionalGeneration.from_pretrained,
|
||||
model,
|
||||
subfolder="mllm",
|
||||
prefetch_list=boogu_subfolders,
|
||||
local_files_only=local_files_only,
|
||||
torch_dtype=od_config.dtype,
|
||||
)
|
||||
# Upstream reuses the full VLM as an optional instruction rewriter and
|
||||
# encodes with its inner model (no ``lm_head``); the rewriter is not
|
||||
# ported, so keep only the inner ``Qwen3VLModel`` as the encoder.
|
||||
if hasattr(mllm, "lm_head"):
|
||||
mllm = mllm.model
|
||||
self.mllm = mllm.to(self._execution_device)
|
||||
|
||||
self.processor = Qwen3VLProcessor.from_pretrained(
|
||||
model, subfolder="processor", local_files_only=local_files_only
|
||||
)
|
||||
|
||||
self.vae = from_pretrained_with_prefetch(
|
||||
AutoencoderKL.from_pretrained,
|
||||
model,
|
||||
subfolder="vae",
|
||||
prefetch_list=boogu_subfolders,
|
||||
local_files_only=local_files_only,
|
||||
).to(self._execution_device)
|
||||
|
||||
self.transformer = BooguImageTransformer2DModel(od_config=od_config)
|
||||
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
|
||||
self.default_sample_size = 128
|
||||
|
||||
self.SYSTEM_PROMPT_4_T2I = SYSTEM_PROMPT_4_T2I_UNIFIED
|
||||
# Upstream uses the TI2I prompt for empty instructions (the default
|
||||
# negative prompt "" hits this path).
|
||||
self.SYSTEM_PROMPT_DROP = SYSTEM_PROMPT_4_TI2I_UNIFIED
|
||||
# Edit (TI2I / I2I) system prompts (image present in the chat template).
|
||||
self.SYSTEM_PROMPT_4_TI2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
|
||||
self.SYSTEM_PROMPT_4_I2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
|
||||
|
||||
def _raise_unsupported_features(self) -> None:
|
||||
"""Reject execution modes that do not have Boogu-specific support."""
|
||||
parallel_config = self.od_config.parallel_config
|
||||
if parallel_config.tensor_parallel_size > 1:
|
||||
raise NotImplementedError("Tensor parallelism is not supported by BooguImagePipeline.")
|
||||
if (parallel_config.sequence_parallel_size or 1) > 1:
|
||||
raise NotImplementedError("Sequence parallelism is not supported by BooguImagePipeline.")
|
||||
if parallel_config.cfg_parallel_size > 1:
|
||||
raise NotImplementedError("CFG parallelism is not supported by BooguImagePipeline.")
|
||||
if parallel_config.use_hsdp:
|
||||
raise NotImplementedError("HSDP is not supported by BooguImagePipeline.")
|
||||
if self.od_config.cache_backend not in (None, "", "none"):
|
||||
raise NotImplementedError(
|
||||
f"Cache backend '{self.od_config.cache_backend}' is not supported by BooguImagePipeline."
|
||||
)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Prompt encoding (upstream ``encode_instruction``, t2i path)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _apply_chat_template(
|
||||
self,
|
||||
instruction: str,
|
||||
input_pil_images: list[PIL.Image.Image] | None = None,
|
||||
) -> list[dict]:
|
||||
"""Build the chat messages for one instruction (text-to-image or edit).
|
||||
|
||||
Mirrors upstream ``_apply_chat_template`` (``system_prompt_follows_task_type``
|
||||
is always ``False`` here): the system prompt is picked by whether images
|
||||
are present and whether the instruction is empty, and reference images
|
||||
are placed *before* the instruction text in the user turn.
|
||||
"""
|
||||
user_text_content = [{"type": "text", "text": instruction}]
|
||||
|
||||
has_images = input_pil_images is not None and len(input_pil_images) > 0
|
||||
instruction_empty = instruction is None or len(instruction.strip()) == 0
|
||||
|
||||
if not has_images:
|
||||
system_prompt = self.SYSTEM_PROMPT_DROP if instruction_empty else self.SYSTEM_PROMPT_4_T2I
|
||||
else:
|
||||
system_prompt = self.SYSTEM_PROMPT_4_I2I if instruction_empty else self.SYSTEM_PROMPT_4_TI2I
|
||||
|
||||
system_role = {"role": "system", "content": [{"type": "text", "text": system_prompt}]}
|
||||
if not has_images:
|
||||
return [system_role, {"role": "user", "content": user_text_content}]
|
||||
|
||||
images_content = [{"type": "image", "image": pil_img} for pil_img in input_pil_images]
|
||||
return [system_role, {"role": "user", "content": images_content + user_text_content}]
|
||||
|
||||
def _get_instruction_feature_embeds(
|
||||
self,
|
||||
instruction: str | list[str],
|
||||
input_pil_images: list[list[PIL.Image.Image] | None] | None = None,
|
||||
device: torch.device | None = None,
|
||||
max_sequence_length: int = 256,
|
||||
truncate_instruction_sequence: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Encode instructions (and optional reference images) with Qwen3VL.
|
||||
|
||||
``input_pil_images`` is a per-sample list (outer length == batch size);
|
||||
each entry is the sample's already-VLM-resized reference images or
|
||||
``None``. Returns the last hidden state (or the last-N layers as a list
|
||||
when the transformer config asks for more than one) and the attention
|
||||
mask.
|
||||
"""
|
||||
device = device or self._execution_device
|
||||
instruction = [instruction] if isinstance(instruction, str) else instruction
|
||||
|
||||
if input_pil_images is None:
|
||||
per_sample_images: list[list[PIL.Image.Image] | None] = [None] * len(instruction)
|
||||
else:
|
||||
assert len(input_pil_images) == len(instruction), (
|
||||
"`input_pil_images` outer length must match the instruction batch size."
|
||||
)
|
||||
per_sample_images = input_pil_images
|
||||
|
||||
prompts = [self._apply_chat_template(text, per_sample_images[i]) for i, text in enumerate(instruction)]
|
||||
|
||||
vlm_inputs = self.processor.apply_chat_template(
|
||||
prompts,
|
||||
padding="longest",
|
||||
max_length=max_sequence_length,
|
||||
truncation=truncate_instruction_sequence,
|
||||
padding_side="right",
|
||||
return_tensors="pt",
|
||||
tokenize=True,
|
||||
return_dict=True,
|
||||
)
|
||||
for k in vlm_inputs.keys():
|
||||
if isinstance(vlm_inputs[k], torch.Tensor):
|
||||
vlm_inputs[k] = vlm_inputs[k].to(device)
|
||||
|
||||
final_instruction_mask = vlm_inputs["attention_mask"]
|
||||
|
||||
num_instruction_feature_layers = self.transformer.instruction_feature_configs.get(
|
||||
"num_instruction_feature_layers", 1
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
text_encoder_outputs = self.mllm(**vlm_inputs, output_hidden_states=True, return_dict=True)
|
||||
if num_instruction_feature_layers > 1:
|
||||
instruction_feats = list(text_encoder_outputs.hidden_states)[-num_instruction_feature_layers:]
|
||||
else:
|
||||
instruction_feats = text_encoder_outputs.hidden_states[-1]
|
||||
|
||||
dtype = self.mllm.dtype if self.mllm is not None else self.transformer.dtype
|
||||
|
||||
if isinstance(instruction_feats, (list, tuple)):
|
||||
final_instruction_feats = [feat.to(dtype=dtype, device=device) for feat in instruction_feats]
|
||||
else:
|
||||
final_instruction_feats = instruction_feats.to(dtype=dtype, device=device)
|
||||
final_instruction_mask = final_instruction_mask.to(device=device)
|
||||
|
||||
return final_instruction_feats, final_instruction_mask
|
||||
|
||||
def _reshape_embeds_and_mask(self, embeds, mask, num_images_per_prompt: int):
|
||||
"""Duplicate embeddings/mask for each generation per prompt (mps-friendly)."""
|
||||
if isinstance(embeds, (list, tuple)):
|
||||
batch_size, seq_len, _ = embeds[0].shape
|
||||
reshaped_embeds = []
|
||||
for embed in embeds:
|
||||
embed = embed.repeat(1, num_images_per_prompt, 1)
|
||||
reshaped_embeds.append(embed.view(batch_size * num_images_per_prompt, seq_len, -1))
|
||||
else:
|
||||
batch_size, seq_len, _ = embeds.shape
|
||||
embeds = embeds.repeat(1, num_images_per_prompt, 1)
|
||||
reshaped_embeds = embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
|
||||
|
||||
mask = mask.repeat(num_images_per_prompt, 1)
|
||||
reshaped_mask = mask.view(batch_size * num_images_per_prompt, -1)
|
||||
|
||||
return batch_size, seq_len, reshaped_embeds, reshaped_mask
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
do_classifier_free_guidance: bool = True,
|
||||
negative_prompt: str | list[str] | None = None,
|
||||
num_images_per_prompt: int = 1,
|
||||
device: torch.device | None = None,
|
||||
prompt_embeds: torch.Tensor | None = None,
|
||||
negative_prompt_embeds: torch.Tensor | None = None,
|
||||
prompt_attention_mask: torch.Tensor | None = None,
|
||||
negative_prompt_attention_mask: torch.Tensor | None = None,
|
||||
max_sequence_length: int = 1280,
|
||||
truncate_instruction_sequence: bool = False,
|
||||
input_images: list[list[PIL.Image.Image] | None] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
|
||||
"""Encode prompt (and negative prompt for CFG) into Qwen3VL hidden states.
|
||||
|
||||
Port of upstream ``encode_instruction`` for text-to-image and the
|
||||
text-guided image-editing (TI2I) path. Reference images are attached to
|
||||
the *positive* instruction only (upstream default
|
||||
``use_input_images_4_neg_instruct=False``). Instruction rewriting,
|
||||
prompt tuning, and double-guidance empty instructions are not ported.
|
||||
The default ``max_sequence_length`` matches the upstream ``__call__``
|
||||
default (1280), not the upstream ``encode_instruction`` default (256).
|
||||
|
||||
Args:
|
||||
input_images: Per-sample list (outer length == batch size) of
|
||||
already-VLM-resized reference images, or ``None`` for pure
|
||||
text-to-image.
|
||||
|
||||
Returns:
|
||||
``(prompt_embeds, prompt_attention_mask, negative_prompt_embeds,
|
||||
negative_prompt_attention_mask)`` where each embeds tensor has shape
|
||||
``[batch_size * num_images_per_prompt, seq_len, dim]``. The negative
|
||||
pair is ``None`` when ``do_classifier_free_guidance`` is off and no
|
||||
precomputed negative embeddings were passed.
|
||||
"""
|
||||
device = device or self._execution_device
|
||||
|
||||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||||
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds, prompt_attention_mask = self._get_instruction_feature_embeds(
|
||||
instruction=prompt,
|
||||
input_pil_images=input_images,
|
||||
device=device,
|
||||
max_sequence_length=max_sequence_length,
|
||||
truncate_instruction_sequence=truncate_instruction_sequence,
|
||||
)
|
||||
|
||||
batch_size, _, prompt_embeds, prompt_attention_mask = self._reshape_embeds_and_mask(
|
||||
prompt_embeds, prompt_attention_mask, num_images_per_prompt
|
||||
)
|
||||
|
||||
if do_classifier_free_guidance and negative_prompt_embeds is None:
|
||||
negative_prompt = negative_prompt if negative_prompt is not None else ""
|
||||
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
|
||||
|
||||
if batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has"
|
||||
f" batch size {batch_size}. Please make sure that passed `negative_prompt`"
|
||||
" matches the batch size of `prompt`."
|
||||
)
|
||||
|
||||
negative_prompt_embeds, negative_prompt_attention_mask = self._get_instruction_feature_embeds(
|
||||
instruction=negative_prompt,
|
||||
device=device,
|
||||
max_sequence_length=max_sequence_length,
|
||||
truncate_instruction_sequence=truncate_instruction_sequence,
|
||||
)
|
||||
|
||||
_, _, negative_prompt_embeds, negative_prompt_attention_mask = self._reshape_embeds_and_mask(
|
||||
negative_prompt_embeds, negative_prompt_attention_mask, num_images_per_prompt
|
||||
)
|
||||
|
||||
return (
|
||||
prompt_embeds,
|
||||
prompt_attention_mask,
|
||||
negative_prompt_embeds,
|
||||
negative_prompt_attention_mask,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Denoise loop + VAE decode (upstream ``__call__`` / ``processing``, t2i)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
|
||||
"""Sample initial noise latents (upstream ``prepare_latents``)."""
|
||||
height = int(height) // self.vae_scale_factor
|
||||
width = int(width) // self.vae_scale_factor
|
||||
shape = (batch_size, num_channels_latents, height, width)
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
return latents
|
||||
|
||||
def _resolve_output_size(self, height, width):
|
||||
"""t2i branch of upstream ``_resolve_output_and_original_size``.
|
||||
|
||||
Clamps the working resolution to ``max_input_image_pixels`` (2048**2),
|
||||
rounding down to a multiple of ``vae_scale_factor * 2``; the requested
|
||||
size is remembered so the decoded image can be resized back.
|
||||
"""
|
||||
img_scale_num = self.vae_scale_factor * 2
|
||||
ori_height, ori_width = height, width
|
||||
max_pixels = 2048 * 2048
|
||||
cur_pixels = height * width
|
||||
ratio = min((max_pixels / cur_pixels) ** 0.5, 1.0)
|
||||
height = int(height * ratio) // img_scale_num * img_scale_num
|
||||
width = int(width * ratio) // img_scale_num * img_scale_num
|
||||
return height, width, ori_height, ori_width
|
||||
|
||||
def predict(
|
||||
self, t, latents, instruction_embeds, freqs_cis, instruction_attention_mask, ref_image_hidden_states=None
|
||||
):
|
||||
"""One transformer velocity prediction (upstream ``predict``).
|
||||
|
||||
``ref_image_hidden_states`` is ``None`` for text-to-image, or the
|
||||
per-sample reference latents (``list[list[Tensor[C, H, W]]]``) for the
|
||||
image-editing path.
|
||||
"""
|
||||
timestep = t.expand(latents.shape[0]).to(latents.dtype)
|
||||
return self.transformer(
|
||||
latents,
|
||||
timestep,
|
||||
instruction_embeds,
|
||||
freqs_cis,
|
||||
instruction_attention_mask,
|
||||
ref_image_hidden_states=ref_image_hidden_states,
|
||||
)
|
||||
|
||||
def _encode_vae_image(self, img: torch.Tensor, generator=None) -> torch.Tensor:
|
||||
"""Encode an image tensor into the VAE latent space (upstream ``encode_vae``).
|
||||
|
||||
Upstream leaves ``latent_dist.sample()`` unseeded; the native path
|
||||
threads the request generator through so a fixed seed gives a
|
||||
reproducible reference latent.
|
||||
"""
|
||||
z0 = self.vae.encode(img.to(dtype=self.vae.dtype)).latent_dist.sample(generator=generator)
|
||||
if self.vae.config.shift_factor is not None:
|
||||
z0 = z0 - self.vae.config.shift_factor
|
||||
if self.vae.config.scaling_factor is not None:
|
||||
z0 = z0 * self.vae.config.scaling_factor
|
||||
return z0.to(dtype=self.vae.dtype)
|
||||
|
||||
def _build_ref_latents(
|
||||
self,
|
||||
preprocessed_images: list[torch.Tensor | None],
|
||||
num_images_per_prompt: int,
|
||||
device: torch.device,
|
||||
generator=None,
|
||||
) -> list[list[torch.Tensor] | None]:
|
||||
"""VAE-encode per-sample reference images into the transformer's format.
|
||||
|
||||
Mirrors upstream ``prepare_image``: returns a list of length
|
||||
``batch_size * num_images_per_prompt`` where each entry is either
|
||||
``None`` (no reference / text-to-image) or a list of ``[C, H, W]``
|
||||
reference latents (one per reference image). Boogu editing uses a single
|
||||
reference image, so each non-empty entry is a one-element list.
|
||||
"""
|
||||
# ``latent_dist.sample`` accepts only a single generator; a per-output
|
||||
# generator list (num_outputs_per_prompt > 1) falls back to unseeded.
|
||||
vae_generator = generator if isinstance(generator, torch.Generator) else None
|
||||
|
||||
ref_latents: list[list[torch.Tensor] | None] = []
|
||||
for image in preprocessed_images:
|
||||
if image is None:
|
||||
sample_latents: list[torch.Tensor] | None = None
|
||||
else:
|
||||
latent = self._encode_vae_image(image.to(device=device), generator=vae_generator).squeeze(0)
|
||||
sample_latents = [latent]
|
||||
for _ in range(num_images_per_prompt):
|
||||
ref_latents.append(sample_latents)
|
||||
return ref_latents
|
||||
|
||||
@staticmethod
|
||||
def _extract_reference_images(
|
||||
prompts: list,
|
||||
) -> tuple[list[PIL.Image.Image | None], list[torch.Tensor | None]]:
|
||||
"""Pull per-sample reference images out of ``additional_information``.
|
||||
|
||||
Returns ``(prompt_images, preprocessed_images)`` where entries are
|
||||
``None`` for pure text-to-image samples. Populated by
|
||||
:func:`get_boogu_image_pre_process_func`.
|
||||
"""
|
||||
prompt_images: list[PIL.Image.Image | None] = []
|
||||
preprocessed_images: list[torch.Tensor | None] = []
|
||||
for p in prompts:
|
||||
if isinstance(p, str):
|
||||
prompt_images.append(None)
|
||||
preprocessed_images.append(None)
|
||||
continue
|
||||
ai = p.get("additional_information") or {}
|
||||
prompt_images.append(ai.get("prompt_image"))
|
||||
preprocessed_images.append(ai.get("preprocessed_image"))
|
||||
return prompt_images, preprocessed_images
|
||||
|
||||
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
|
||||
# Prompt / negative-prompt extraction (mirrors the Ovis pattern; the
|
||||
# online API sometimes passes ``{"negative_prompt": None}``).
|
||||
prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
|
||||
if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
|
||||
negative_prompt = None
|
||||
else:
|
||||
negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
|
||||
|
||||
# Reference (edit / TI2I) images, if any.
|
||||
prompt_images, preprocessed_images = self._extract_reference_images(req.prompts)
|
||||
has_reference = any(img is not None for img in preprocessed_images)
|
||||
task_type = "ti2i" if has_reference else "t2i"
|
||||
|
||||
sp = req.sampling_params
|
||||
device = self._execution_device
|
||||
|
||||
height = sp.height or self.default_sample_size * self.vae_scale_factor
|
||||
width = sp.width or self.default_sample_size * self.vae_scale_factor
|
||||
num_inference_steps = sp.num_inference_steps or 50
|
||||
# Upstream default text guidance is 4.0; the engine coerces an unset
|
||||
# guidance_scale to 1.0, so only honor a caller-provided value.
|
||||
text_guidance_scale = sp.guidance_scale if sp.guidance_scale_provided else 4.0
|
||||
# Image guidance rides on ``guidance_scale_2`` (upstream default 1.0 =
|
||||
# off); only a caller-provided value enables the double-guidance path.
|
||||
image_guidance_scale = sp.guidance_scale_2 if sp.guidance_scale_2_provided else 1.0
|
||||
if not has_reference:
|
||||
image_guidance_scale = 1.0
|
||||
num_images_per_prompt = sp.num_outputs_per_prompt if sp.num_outputs_per_prompt > 0 else 1
|
||||
generator = sp.generator
|
||||
max_sequence_length = sp.max_sequence_length or 1280
|
||||
output_type = sp.output_type or "pil"
|
||||
cfg_range = (0.0, 1.0)
|
||||
|
||||
# Negative instruction embeddings are needed whenever text guidance is
|
||||
# active (t2i text CFG, ti2i text-only, and ti2i double guidance).
|
||||
do_classifier_free_guidance = text_guidance_scale > 1.0
|
||||
|
||||
batch_size = len(prompt)
|
||||
|
||||
# Per-sample VLM reference images for the positive instruction only
|
||||
# (upstream default ``use_input_images_4_neg_instruct=False``).
|
||||
input_images = None
|
||||
if has_reference:
|
||||
input_images = [[img] if img is not None else None for img in prompt_images]
|
||||
|
||||
# 1. Encode prompts.
|
||||
(
|
||||
instruction_embeds,
|
||||
instruction_attention_mask,
|
||||
negative_instruction_embeds,
|
||||
negative_instruction_attention_mask,
|
||||
) = self.encode_prompt(
|
||||
prompt,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
negative_prompt=negative_prompt,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
device=device,
|
||||
max_sequence_length=max_sequence_length,
|
||||
input_images=input_images,
|
||||
)
|
||||
|
||||
# 2. Resolve working / output resolution.
|
||||
height, width, ori_height, ori_width = self._resolve_output_size(height, width)
|
||||
|
||||
# 3. Reference latents (edit path) and initial noise latents.
|
||||
dtype = self.vae.dtype
|
||||
latent_channels = self.transformer.in_channels
|
||||
ref_latents = None
|
||||
if has_reference:
|
||||
ref_latents = self._build_ref_latents(preprocessed_images, num_images_per_prompt, device, generator)
|
||||
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
latent_channels,
|
||||
height,
|
||||
width,
|
||||
instruction_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
)
|
||||
|
||||
freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(
|
||||
self.transformer.axes_dim_rope,
|
||||
self.transformer.axes_lens,
|
||||
theta=10000,
|
||||
)
|
||||
|
||||
# 4. Timesteps (the ported scheduler consumes ``num_tokens``).
|
||||
num_tokens = latents.shape[-2] * latents.shape[-1]
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device, num_tokens=num_tokens)
|
||||
timesteps = self.scheduler.timesteps
|
||||
num_timesteps = len(timesteps)
|
||||
|
||||
# 5. Denoise loop with (sequential) classifier-free guidance. Reproduces
|
||||
# the branch priority of upstream ``processing`` (double > text-only >
|
||||
# image-only > t2i text). Reference latents are kept in the conditional
|
||||
# (and, for text-only ti2i, the unconditional) predictions.
|
||||
with self.progress_bar(total=num_timesteps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
in_cfg_range = cfg_range[0] <= i / num_timesteps <= cfg_range[1]
|
||||
text_gs = text_guidance_scale if in_cfg_range else 1.0
|
||||
image_gs = image_guidance_scale if in_cfg_range else 1.0
|
||||
|
||||
model_pred = self.predict(
|
||||
t, latents, instruction_embeds, freqs_cis, instruction_attention_mask, ref_latents
|
||||
)
|
||||
|
||||
if task_type == "ti2i" and text_gs > 1.0 and image_gs > 1.0:
|
||||
# Double guidance: 3 predictions (cond+ref, neg+ref, neg+no-ref).
|
||||
model_pred_drop_text = self.predict(
|
||||
t,
|
||||
latents,
|
||||
negative_instruction_embeds,
|
||||
freqs_cis,
|
||||
negative_instruction_attention_mask,
|
||||
ref_latents,
|
||||
)
|
||||
model_pred_drop_all = self.predict(
|
||||
t,
|
||||
latents,
|
||||
negative_instruction_embeds,
|
||||
freqs_cis,
|
||||
negative_instruction_attention_mask,
|
||||
None,
|
||||
)
|
||||
delta_text = model_pred - model_pred_drop_text
|
||||
delta_image = model_pred_drop_text - model_pred_drop_all
|
||||
model_pred = model_pred + (text_gs - 1) * delta_text + (image_gs - 1) * delta_image
|
||||
elif task_type == "ti2i" and text_gs > 1.0:
|
||||
# Text-only ti2i guidance: reference kept in the uncond pred.
|
||||
model_pred_drop_text = self.predict(
|
||||
t,
|
||||
latents,
|
||||
negative_instruction_embeds,
|
||||
freqs_cis,
|
||||
negative_instruction_attention_mask,
|
||||
ref_latents,
|
||||
)
|
||||
model_pred = model_pred + (text_gs - 1) * (model_pred - model_pred_drop_text)
|
||||
elif task_type == "ti2i" and image_gs > 1.0:
|
||||
# Image-only ti2i guidance: drop the reference in the uncond pred.
|
||||
model_pred_drop_image = self.predict(
|
||||
t, latents, instruction_embeds, freqs_cis, instruction_attention_mask, None
|
||||
)
|
||||
model_pred = model_pred + (image_gs - 1) * (model_pred - model_pred_drop_image)
|
||||
elif text_gs > 1.0:
|
||||
# Text-to-image classifier-free guidance.
|
||||
model_pred_drop_all = self.predict(
|
||||
t,
|
||||
latents,
|
||||
negative_instruction_embeds,
|
||||
freqs_cis,
|
||||
negative_instruction_attention_mask,
|
||||
None,
|
||||
)
|
||||
model_pred = model_pred + (text_gs - 1) * (model_pred - model_pred_drop_all)
|
||||
|
||||
latents = self.scheduler.step(model_pred, t, latents, return_dict=False)[0]
|
||||
latents = latents.to(dtype=instruction_embeds.dtype)
|
||||
progress_bar.update()
|
||||
|
||||
# 6. Decode.
|
||||
if output_type == "latent":
|
||||
image = latents
|
||||
else:
|
||||
latents = latents.to(dtype=dtype)
|
||||
if self.vae.config.scaling_factor is not None:
|
||||
latents = latents / self.vae.config.scaling_factor
|
||||
if self.vae.config.shift_factor is not None:
|
||||
latents = latents + self.vae.config.shift_factor
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
if (ori_height, ori_width) != (height, width):
|
||||
image = F.interpolate(image, size=(ori_height, ori_width), mode="bilinear")
|
||||
|
||||
return DiffusionOutput(output=image)
|
||||
+318
@@ -0,0 +1,318 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Ported for vLLM-Omni native Boogu-Image support from
|
||||
# boogu/schedulers/scheduling_flow_match_euler_discrete_time_shifting.py
|
||||
# (https://github.com/boogu-project/Boogu-Image). The scheduling logic is
|
||||
# kept verbatim; only debug comments were removed and typing modernized.
|
||||
#
|
||||
# --- Upstream header preserved below ---
|
||||
#
|
||||
# Copyright (C) 2026 Boogu Team.
|
||||
#
|
||||
# This file is adapted by Boogu Team from prior open-source scheduler work.
|
||||
# Boogu-specific modifications include static/dynamic time-shift handling used
|
||||
# by the released Boogu pipeline.
|
||||
#
|
||||
# Original work:
|
||||
# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
||||
from diffusers.utils import BaseOutput, logging
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
|
||||
"""
|
||||
Output class for the scheduler's `step` function output.
|
||||
|
||||
Args:
|
||||
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
||||
denoising loop.
|
||||
"""
|
||||
|
||||
prev_sample: torch.FloatTensor
|
||||
|
||||
|
||||
class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
|
||||
"""
|
||||
Euler scheduler with Boogu's training-consistent time shifting.
|
||||
|
||||
Timesteps run ascending 0 -> 1 (not descending sigmas as in stock
|
||||
diffusers), `set_timesteps` accepts a `num_tokens` argument for the
|
||||
dynamic-shift variants, and `step()` integrates with `t_next - t`
|
||||
directly.
|
||||
|
||||
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
|
||||
methods the library implements for all schedulers such as loading and saving.
|
||||
|
||||
Args:
|
||||
num_train_timesteps (`int`, defaults to 1000):
|
||||
The number of diffusion steps to train the model.
|
||||
do_shift (`bool`, defaults to `True`):
|
||||
Whether to apply the training-consistent time shift in `set_timesteps`.
|
||||
dynamic_time_shift (`bool`, defaults to `True`):
|
||||
If `True`, the shift depends on the per-sample `num_tokens`; if `False`, it uses the configured `seq_len`.
|
||||
time_shift_version (`str`, defaults to `"v2"`):
|
||||
`"v1"` (logistic transform with a linear mu mapping) or `"v2"` (rational transform).
|
||||
seq_len (`int`, *optional*):
|
||||
Token count used to compute the static shift when `dynamic_time_shift=False` (mirrors training).
|
||||
base_shift (`float`, defaults to 0.5):
|
||||
v1 linear mapping lower bound (matches training defaults).
|
||||
max_shift (`float`, defaults to 1.15):
|
||||
v1 linear mapping upper bound (matches training defaults).
|
||||
"""
|
||||
|
||||
_compatibles = []
|
||||
order = 1
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_train_timesteps: int = 1000,
|
||||
do_shift: bool = True,
|
||||
dynamic_time_shift: bool = True,
|
||||
time_shift_version: str = "v2",
|
||||
# seq_len is used to mirror training-side static time shift (when dynamic_time_shift=False)
|
||||
# In training, seq_len is the token count used to compute shift.
|
||||
seq_len: int | None = None,
|
||||
# v1 linear mapping range (matches training defaults)
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
time_shift_v2_half_scaling_factor: float = 60.0,
|
||||
):
|
||||
timesteps = torch.linspace(0, 1, num_train_timesteps + 1, dtype=torch.float32)[:-1]
|
||||
|
||||
self.timesteps = timesteps
|
||||
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
self.time_shift_v2_scaling_factor = time_shift_v2_half_scaling_factor * 2
|
||||
|
||||
@property
|
||||
def step_index(self):
|
||||
"""
|
||||
The index counter for current timestep. It will increase 1 after each scheduler step.
|
||||
"""
|
||||
return self._step_index
|
||||
|
||||
@property
|
||||
def begin_index(self):
|
||||
"""
|
||||
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
|
||||
"""
|
||||
return self._begin_index
|
||||
|
||||
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
|
||||
def set_begin_index(self, begin_index: int = 0):
|
||||
"""
|
||||
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
|
||||
|
||||
Args:
|
||||
begin_index (`int`):
|
||||
The begin index for the scheduler.
|
||||
"""
|
||||
self._begin_index = begin_index
|
||||
|
||||
def index_for_timestep(self, timestep, schedule_timesteps=None):
|
||||
if schedule_timesteps is None:
|
||||
schedule_timesteps = self._timesteps
|
||||
|
||||
indices = (schedule_timesteps == timestep).nonzero()
|
||||
|
||||
# The sigma index that is taken for the **very** first `step`
|
||||
# is always the second index (or the last index if there is only 1)
|
||||
# This way we can ensure we don't accidentally skip a sigma in
|
||||
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
|
||||
pos = 1 if len(indices) > 1 else 0
|
||||
|
||||
return indices[pos].item()
|
||||
|
||||
# --- Helpers to mirror training-side shift logic ---
|
||||
@staticmethod
|
||||
def _get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15):
|
||||
m = (y2 - y1) / (x2 - x1)
|
||||
b = y1 - m * x1
|
||||
return lambda x: m * x + b
|
||||
|
||||
@staticmethod
|
||||
def _time_shift_v1(t_np: np.ndarray, mu: float, sigma: float = 1.0) -> np.ndarray:
|
||||
# Matches training: t <- 1 - t; logistic transform; then t <- 1 - t
|
||||
eps = 1e-8
|
||||
t1 = 1.0 - t_np
|
||||
t1 = np.clip(t1, eps, 1.0 - eps)
|
||||
num = math.exp(mu)
|
||||
denom = num + np.power(1.0 / t1 - 1.0, sigma)
|
||||
y = num / denom
|
||||
out = 1.0 - y
|
||||
return out.astype(np.float32)
|
||||
|
||||
@staticmethod
|
||||
def _time_shift_v2(t_np: np.ndarray, m: float) -> np.ndarray:
|
||||
# Matches training: t' = t / (m - m t + t)
|
||||
return (t_np / (m - m * t_np + t_np)).astype(np.float32)
|
||||
|
||||
def set_timesteps(
|
||||
self,
|
||||
num_inference_steps: int | None = None,
|
||||
device: str | torch.device | None = None,
|
||||
timesteps: list[float] | None = None,
|
||||
num_tokens: int | None = None,
|
||||
):
|
||||
"""
|
||||
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
|
||||
|
||||
Args:
|
||||
num_inference_steps (`int`):
|
||||
The number of diffusion steps used when generating samples with a pre-trained model.
|
||||
device (`str` or `torch.device`, *optional*):
|
||||
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
||||
timesteps (`list[float]`, *optional*):
|
||||
Custom timesteps to use. If provided, `num_inference_steps` is ignored.
|
||||
num_tokens (`int`, *optional*):
|
||||
Per-sample token count, used by the dynamic time-shift variants.
|
||||
"""
|
||||
|
||||
if timesteps is None:
|
||||
self.num_inference_steps = num_inference_steps
|
||||
t_arr = np.linspace(0, 1, num_inference_steps + 1, dtype=np.float32)[:-1]
|
||||
|
||||
# Apply training-consistent time shift only when requested
|
||||
if self.config.do_shift:
|
||||
# dynamic or static
|
||||
if self.config.dynamic_time_shift:
|
||||
# dynamic: depend on per-sample token count
|
||||
if self.config.time_shift_version == "v1":
|
||||
# In training dynamic v1: mu is computed from tokens' linear map where
|
||||
# tokens are approximately (H_lat//2)*(W_lat//2). We approximate this with num_tokens//4.
|
||||
if num_tokens is not None and num_tokens > 0:
|
||||
tokens_reduced = max(1, int(num_tokens) // 4)
|
||||
lin = self._get_lin_function(y1=self.config.base_shift, y2=self.config.max_shift)
|
||||
mu = lin(tokens_reduced) ## 4096 for 1024x1024 resolution
|
||||
|
||||
t_arr = self._time_shift_v1(t_arr, mu, sigma=1.0)
|
||||
# else: no-op if we lack num_tokens
|
||||
elif self.config.time_shift_version == "v2":
|
||||
# MUST remain identical to current behavior when v2 + dynamic=True
|
||||
# m = sqrt(num_tokens) / 40; t' = t / (m - m t + t)
|
||||
# When input resolution is 320 * 320, m = 1, when input resolution is 512 * 512,
|
||||
# m = 1.6, when input resolution is 1024 * 1024, m = 3.2
|
||||
if num_tokens is not None and num_tokens > 0:
|
||||
m = float(np.sqrt(num_tokens)) / self.time_shift_v2_scaling_factor
|
||||
t_arr = self._time_shift_v2(t_arr, m)
|
||||
# else: no-op
|
||||
else:
|
||||
# static: depend on seq_len configured at scheduler init
|
||||
if self.config.time_shift_version == "v1":
|
||||
if self.config.seq_len is not None and self.config.seq_len > 0:
|
||||
lin = self._get_lin_function(y1=self.config.base_shift, y2=self.config.max_shift)
|
||||
mu = lin(int(self.config.seq_len))
|
||||
t_arr = self._time_shift_v1(t_arr, mu, sigma=1.0)
|
||||
elif self.config.time_shift_version == "v2":
|
||||
if self.config.seq_len is not None and self.config.seq_len > 0:
|
||||
# training static v2 uses m = sqrt(seq_len) / 40
|
||||
m = float(np.sqrt(self.config.seq_len)) / self.time_shift_v2_scaling_factor
|
||||
t_arr = self._time_shift_v2(t_arr, m)
|
||||
|
||||
timesteps = t_arr
|
||||
|
||||
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)
|
||||
_timesteps = torch.cat([timesteps, torch.ones(1, device=timesteps.device)])
|
||||
|
||||
self.timesteps = timesteps
|
||||
self._timesteps = _timesteps
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
|
||||
def _init_step_index(self, timestep):
|
||||
if self.begin_index is None:
|
||||
if isinstance(timestep, torch.Tensor):
|
||||
timestep = timestep.to(self.timesteps.device)
|
||||
self._step_index = self.index_for_timestep(timestep)
|
||||
else:
|
||||
self._step_index = self._begin_index
|
||||
|
||||
def step(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: float | torch.FloatTensor,
|
||||
sample: torch.FloatTensor,
|
||||
generator: torch.Generator | None = None, # noqa: ARG002 - kept for API compatibility
|
||||
return_dict: bool = True,
|
||||
) -> FlowMatchEulerDiscreteSchedulerOutput | tuple:
|
||||
"""
|
||||
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
||||
process from the learned model outputs (most often the predicted noise).
|
||||
|
||||
Args:
|
||||
model_output (`torch.FloatTensor`):
|
||||
The direct output from learned diffusion model.
|
||||
timestep (`float`):
|
||||
The current discrete timestep in the diffusion chain.
|
||||
sample (`torch.FloatTensor`):
|
||||
A current instance of a sample created by the diffusion process.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
A random number generator.
|
||||
return_dict (`bool`):
|
||||
Whether or not to return a [`FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
|
||||
|
||||
Returns:
|
||||
[`FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
|
||||
If return_dict is `True`, [`FlowMatchEulerDiscreteSchedulerOutput`] is returned, otherwise a tuple is
|
||||
returned where the first element is the sample tensor.
|
||||
"""
|
||||
|
||||
if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor):
|
||||
raise ValueError(
|
||||
(
|
||||
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
|
||||
" `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
|
||||
" one of the `scheduler.timesteps` as a timestep."
|
||||
),
|
||||
)
|
||||
|
||||
if self.step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
# Upcast to avoid precision issues when computing prev_sample
|
||||
sample = sample.to(torch.float32)
|
||||
t = self._timesteps[self.step_index]
|
||||
t_next = self._timesteps[self.step_index + 1]
|
||||
|
||||
prev_sample = sample + (t_next - t) * model_output
|
||||
|
||||
# Cast sample back to model compatible dtype
|
||||
prev_sample = prev_sample.to(model_output.dtype)
|
||||
|
||||
# upon completion increase step index by one
|
||||
self._step_index += 1
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample,)
|
||||
|
||||
return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
|
||||
|
||||
def __len__(self):
|
||||
return self.config.num_train_timesteps
|
||||
@@ -141,6 +141,11 @@ _DIFFUSION_MODELS = {
|
||||
"pipeline_bagel",
|
||||
"BagelPipeline",
|
||||
),
|
||||
"BooguImagePipeline": (
|
||||
"boogu_image",
|
||||
"pipeline_boogu_image",
|
||||
"BooguImagePipeline",
|
||||
),
|
||||
"LancePipeline": (
|
||||
"lance",
|
||||
"pipeline_lance",
|
||||
@@ -503,6 +508,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"GlmImagePipeline": "get_glm_image_post_process_func",
|
||||
"ZImagePipeline": "get_post_process_func",
|
||||
"OvisImagePipeline": "get_ovis_image_post_process_func",
|
||||
"BooguImagePipeline": "get_boogu_image_post_process_func",
|
||||
"WanPipeline": "get_wan22_post_process_func",
|
||||
"WanVACEPipeline": "get_wan22_vace_post_process_func",
|
||||
"LTX2Pipeline": "get_ltx2_post_process_func",
|
||||
@@ -566,6 +572,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
|
||||
# `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
|
||||
# where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
|
||||
"GlmImagePipeline": "get_glm_image_pre_process_func",
|
||||
"BooguImagePipeline": "get_boogu_image_pre_process_func",
|
||||
"QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
|
||||
"QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func",
|
||||
"LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func",
|
||||
|
||||
@@ -56,6 +56,13 @@ class OmniDiffusionRequest:
|
||||
):
|
||||
self.sampling_params.do_classifier_free_guidance = True
|
||||
|
||||
# Detect whether the caller explicitly provided guidance_scale_2
|
||||
# BEFORE auto-filling it. Pipelines with a model-specific second
|
||||
# guidance default (e.g. Boogu image_guidance_scale=1.0) must be able
|
||||
# to tell an explicit value apart from the guidance_scale fallback.
|
||||
if self.sampling_params.guidance_scale_2 is not None:
|
||||
self.sampling_params.guidance_scale_2_provided = True
|
||||
|
||||
# Auto-fill guidance_scale_2 from the (now-resolved) guidance_scale
|
||||
# so downstream code always has a valid value.
|
||||
if self.sampling_params.guidance_scale_2 is None:
|
||||
|
||||
@@ -1955,6 +1955,7 @@ async def edit_images(
|
||||
negative_prompt: str | None = Form(None),
|
||||
num_inference_steps: int | None = Form(None),
|
||||
guidance_scale: float | None = Form(None),
|
||||
guidance_scale_2: float | None = Form(None),
|
||||
strength: float | None = Form(None),
|
||||
true_cfg_scale: float | None = Form(None),
|
||||
seed: int | None = Form(None),
|
||||
@@ -2120,6 +2121,7 @@ async def edit_images(
|
||||
# 3.4 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, "guidance_scale_2", guidance_scale_2)
|
||||
_update_if_not_none(gen_params, "strength", strength)
|
||||
_update_if_not_none(gen_params, "true_cfg_scale", true_cfg_scale)
|
||||
# If seed is not provided, generate a random one to ensure
|
||||
@@ -2185,6 +2187,8 @@ async def edit_images(
|
||||
extra_body["num_inference_steps"] = num_inference_steps
|
||||
if guidance_scale is not None:
|
||||
extra_body["guidance_scale"] = guidance_scale
|
||||
if guidance_scale_2 is not None:
|
||||
extra_body["guidance_scale_2"] = guidance_scale_2
|
||||
if strength is not None:
|
||||
extra_body["strength"] = strength
|
||||
if true_cfg_scale is not None:
|
||||
|
||||
@@ -248,6 +248,7 @@ class OmniDiffusionSamplingParams:
|
||||
guidance_scale: float = 0.0
|
||||
guidance_scale_provided: bool = False
|
||||
guidance_scale_2: float | None = None
|
||||
guidance_scale_2_provided: bool = False
|
||||
guidance_rescale: float = 0.0
|
||||
strength: float | None = None # I2I: Z-Image specific now, uses to control denoising start timestep
|
||||
decode_timestep: float | list[float] | None = None
|
||||
|
||||
Reference in new issue
Block a user