diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index d0f866b5a..9568f60c4 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index dda7c5f2d..ca3c38a91 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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` | ✅︎ | ✅︎ | | ✅︎ | diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md new file mode 100644 index 000000000..b9bff789f --- /dev/null +++ b/recipes/Boogu/Boogu-Image.md @@ -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: +- Project page: +- GitHub: +- 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,"}}, + {"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`). diff --git a/recipes/README.md b/recipes/README.md index be4c7cc08..7129918e1 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 | diff --git a/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json new file mode 100644 index 000000000..028889af1 --- /dev/null +++ b/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json @@ -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 + } + } + ] + } +] diff --git a/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json b/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json new file mode 100644 index 000000000..4ee6ca577 --- /dev/null +++ b/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json @@ -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 + } + } + ] + } +] diff --git a/tests/diffusion/models/boogu_image/__init__.py b/tests/diffusion/models/boogu_image/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_scheduler.py b/tests/diffusion/models/boogu_image/test_boogu_image_scheduler.py new file mode 100644 index 000000000..a92d9e979 --- /dev/null +++ b/tests/diffusion/models/boogu_image/test_boogu_image_scheduler.py @@ -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) diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py new file mode 100644 index 000000000..5551ba191 --- /dev/null +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -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.`` -> ``.to_out.0.`` (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() diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py new file mode 100644 index 000000000..d38be97fa --- /dev/null +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py @@ -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) diff --git a/tests/e2e/online_serving/test_boogu_image.py b/tests/e2e/online_serving/test_boogu_image.py new file mode 100644 index 000000000..5a5494d7d --- /dev/null +++ b/tests/e2e/online_serving/test_boogu_image.py @@ -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) diff --git a/tests/e2e/online_serving/test_boogu_image_edit.py b/tests/e2e/online_serving/test_boogu_image_edit.py new file mode 100644 index 000000000..45f999d26 --- /dev/null +++ b/tests/e2e/online_serving/test_boogu_image_edit.py @@ -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) diff --git a/tests/e2e/online_serving/test_boogu_image_edit_expansion.py b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py new file mode 100644 index 000000000..efc968eb6 --- /dev/null +++ b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py @@ -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)}" + ) diff --git a/tests/e2e/online_serving/test_boogu_image_expansion.py b/tests/e2e/online_serving/test_boogu_image_expansion.py new file mode 100644 index 000000000..ef6d3a62e --- /dev/null +++ b/tests/e2e/online_serving/test_boogu_image_expansion.py @@ -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) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 5a7ef7bf0..34ae5ba20 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -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 diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 4b51220c7..4bc740b01 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -45,6 +45,7 @@ EXCLUDED_MODELS = [ "WanI2VDMD2Pipeline", "LongCatImagePipeline", "BagelPipeline", + "BooguImagePipeline", "LancePipeline", "MingImagePipeline", "InternVLAA1Pipeline", diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index f33463384..7e02cf5d8 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -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, + ), } diff --git a/vllm_omni/diffusion/models/boogu_image/__init__.py b/vllm_omni/diffusion/models/boogu_image/__init__.py new file mode 100644 index 000000000..7ec9c167c --- /dev/null +++ b/vllm_omni/diffusion/models/boogu_image/__init__.py @@ -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", +] diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py new file mode 100644 index 000000000..e7489c868 --- /dev/null +++ b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py @@ -0,0 +1,1372 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Native vLLM-Omni port of the Boogu-Image transformer. +# +# Ported from the upstream `boogu` package +# (boogu/models/transformers/transformer_boogu.py and friends) with the +# following changes: +# - Diffusers mixins (ModelMixin/ConfigMixin/PeftAdapterMixin) removed; +# configuration is read from `od_config.tf_model_config`. +# - diffusers Attention + custom processors replaced by the vLLM-Omni +# `Attention` layer (backend/SP/KV-quant handled by the framework). +# - Q/K/V and FFN projections use vLLM parallel linears for TP support. +# - Training-time features (prompt tuning, gradient checkpointing) and +# inference caches (TeaCache/TaylorSeer) are not ported. + +import itertools +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import Timesteps, get_1d_rotary_pos_embed +from einops import rearrange, repeat +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + """Apply rotary embeddings in complex form (upstream `use_real=False` path). + + Args: + x: Query or key tensor of shape [B, S, H, D]. + freqs_cis: Complex frequency tensor of shape [B, S, D // 2]. + """ + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], x.shape[-1] // 2, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + return x_out.type_as(x) + + +def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return F.silu(x.float(), inplace=False).to(x.dtype) * y + + +class TimestepEmbedding(nn.Module): + def __init__(self, in_channels: int, time_embed_dim: int) -> None: + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=True) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim, bias=True) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + return self.linear_2(self.act(self.linear_1(sample))) + + +class LuminaRMSNormZero(nn.Module): + """AdaRMS modulation: projects `temb` into scale/gate terms.""" + + def __init__(self, embedding_dim: int, norm_eps: float) -> None: + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(min(embedding_dim, 1024), 4 * embedding_dim, bias=True) + self.norm = RMSNorm(embedding_dim, eps=norm_eps) + + def forward( + self, x: torch.Tensor, emb: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = self.linear(self.silu(emb)) + scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + return x, gate_msa, scale_mlp, gate_mlp + + +class LuminaLayerNormContinuous(nn.Module): + """Final AdaLN + optional output projection (upstream `norm_out`).""" + + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + bias: bool = True, + out_dim: int | None = None, + ) -> None: + super().__init__() + self.silu = nn.SiLU() + self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias) if out_dim is not None else None + + def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: + scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) + x = self.norm(x) * (1 + scale)[:, None, :] + if self.linear_2 is not None: + x = self.linear_2(x) + return x + + +class LuminaFeedForward(nn.Module): + """SwiGLU feed-forward with tensor-parallel projections.""" + + def __init__( + self, + dim: int, + inner_dim: int, + multiple_of: int = 256, + ffn_dim_multiplier: float | None = None, + ) -> None: + super().__init__() + if ffn_dim_multiplier is not None: + inner_dim = int(ffn_dim_multiplier * inner_dim) + inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) + + self.linear_1 = ColumnParallelLinear(dim, inner_dim, bias=False) # gate + self.linear_3 = ColumnParallelLinear(dim, inner_dim, bias=False) # input + self.linear_2 = RowParallelLinear(inner_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h1, _ = self.linear_1(x) + h2, _ = self.linear_3(x) + out, _ = self.linear_2(swiglu(h1, h2)) + return out + + +class Lumina2CombinedTimestepCaptionEmbedding(nn.Module): + def __init__( + self, + hidden_size: int = 4096, + instruction_feat_dim: int = 2048, + frequency_embedding_size: int = 256, + norm_eps: float = 1e-5, + timestep_scale: float = 1.0, + ) -> None: + super().__init__() + self.time_proj = Timesteps( + num_channels=frequency_embedding_size, + flip_sin_to_cos=True, + downscale_freq_shift=0.0, + scale=timestep_scale, + ) + self.timestep_embedder = TimestepEmbedding( + in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024) + ) + self.caption_embedder = nn.Sequential( + RMSNorm(instruction_feat_dim, eps=norm_eps), + nn.Linear(instruction_feat_dim, hidden_size, bias=True), + ) + + def forward( + self, + timestep: torch.Tensor, + instruction_hidden_states: torch.Tensor, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + timestep_proj = self.time_proj(timestep).to(dtype=dtype) + time_embed = self.timestep_embedder(timestep_proj) + caption_embed = self.caption_embedder(instruction_hidden_states) + return time_embed, caption_embed + + +class BooguImageDoubleStreamRotaryPosEmbed(nn.Module): + """3-axis RoPE producing per-segment frequency tensors. + + Returns separate frequency tensors for instruction tokens, reference-image + tokens, noise-image tokens, the full joint sequence, and the combined + (ref + noise) image sequence used by the double-stream blocks. + """ + + def __init__( + self, + theta: int, + axes_dim: tuple[int, int, int], + axes_lens: tuple[int, int, int] = (300, 512, 512), + patch_size: int = 2, + ) -> None: + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + self.axes_lens = axes_lens + self.patch_size = patch_size + + @staticmethod + def get_freqs_cis( + axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int + ) -> list[torch.Tensor]: + freqs_cis = [] + freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + for d, e in zip(axes_dim, axes_lens): + emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype) + freqs_cis.append(emb) + return freqs_cis + + def _get_freqs_cis(self, freqs_cis: list[torch.Tensor], ids: torch.Tensor) -> torch.Tensor: + device = ids.device + if ids.device.type == "mps": + ids = ids.to("cpu") + + result = [] + for i in range(len(self.axes_dim)): + freqs = freqs_cis[i].to(ids.device) + index = ids[:, :, i : i + 1].repeat(1, 1, freqs.shape[-1]).to(torch.int64) + result.append(torch.gather(freqs.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index)) + return torch.cat(result, dim=-1).to(device) + + def forward( + self, + freqs_cis: list[torch.Tensor], + attention_mask: torch.Tensor, + l_effective_ref_img_len: list[list[int]], + l_effective_img_len: list[int], + ref_img_sizes: list[list[tuple[int, int]] | None], + img_sizes: list[tuple[int, int]], + device: torch.device, + ): + batch_size = len(attention_mask) + p = self.patch_size + + encoder_seq_len = attention_mask.shape[1] + l_effective_cap_len = attention_mask.sum(dim=1).tolist() + + seq_lengths = [ + cap_len + sum(ref_img_len) + img_len + for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len) + ] + + max_seq_len = max(seq_lengths) + max_ref_img_len = max(sum(ref_img_len) for ref_img_len in l_effective_ref_img_len) + max_img_len = max(l_effective_img_len) + + position_ids = torch.zeros(batch_size, max_seq_len, 3, dtype=torch.int32, device=device) + + for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)): + position_ids[i, :cap_seq_len] = repeat( + torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3" + ) + + pe_shift = cap_seq_len + pe_shift_len = cap_seq_len + + if ref_img_sizes[i] is not None: + for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]): + H, W = ref_img_size + ref_H_tokens, ref_W_tokens = H // p, W // p + assert ref_H_tokens * ref_W_tokens == ref_img_len + + row_ids = repeat( + torch.arange(ref_H_tokens, dtype=torch.int32, device=device), + "h -> h w", + w=ref_W_tokens, + ).flatten() + col_ids = repeat( + torch.arange(ref_W_tokens, dtype=torch.int32, device=device), + "w -> h w", + h=ref_H_tokens, + ).flatten() + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = pe_shift + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = row_ids + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = col_ids + + pe_shift += max(ref_H_tokens, ref_W_tokens) + pe_shift_len += ref_img_len + + H, W = img_sizes[i] + H_tokens, W_tokens = H // p, W // p + assert H_tokens * W_tokens == l_effective_img_len[i] + + row_ids = repeat(torch.arange(H_tokens, dtype=torch.int32, device=device), "h -> h w", w=W_tokens).flatten() + col_ids = repeat(torch.arange(W_tokens, dtype=torch.int32, device=device), "w -> h w", h=H_tokens).flatten() + + assert pe_shift_len + l_effective_img_len[i] == seq_len + position_ids[i, pe_shift_len:seq_len, 0] = pe_shift + position_ids[i, pe_shift_len:seq_len, 1] = row_ids + position_ids[i, pe_shift_len:seq_len, 2] = col_ids + + freqs_cis = self._get_freqs_cis(freqs_cis, position_ids) + + cap_freqs_cis = torch.zeros( + batch_size, encoder_seq_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype + ) + ref_img_freqs_cis = torch.zeros( + batch_size, max_ref_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype + ) + img_freqs_cis = torch.zeros(batch_size, max_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype) + + combined_img_seq_lengths = [ + sum(ref_img_len) + img_len for ref_img_len, img_len in zip(l_effective_ref_img_len, l_effective_img_len) + ] + max_combined_img_len = max(combined_img_seq_lengths) + + combined_img_freqs_cis = torch.zeros( + batch_size, max_combined_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype + ) + + for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate( + zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths) + ): + cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len] + ref_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[i, cap_seq_len : cap_seq_len + sum(ref_img_len)] + img_freqs_cis[i, :img_len] = freqs_cis[ + i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len + ] + + combined_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[i, cap_seq_len : cap_seq_len + sum(ref_img_len)] + combined_img_freqs_cis[i, sum(ref_img_len) : sum(ref_img_len) + img_len] = freqs_cis[ + i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len + ] + + return ( + cap_freqs_cis, + ref_img_freqs_cis, + img_freqs_cis, + freqs_cis, + l_effective_cap_len, + seq_lengths, + combined_img_freqs_cis, + combined_img_seq_lengths, + ) + + +def _concat_instruction_image_features( + img_tensors: list[torch.Tensor], + instruct_tensors: list[torch.Tensor], + encoder_seq_lengths: list[int], + seq_lengths: list[int], +) -> list[torch.Tensor]: + """Pack per-stream tensors into joint sequences (instruction first, then image).""" + assert len(img_tensors) == len(instruct_tensors) + + batch_size = img_tensors[0].shape[0] + max_seq_len = max(seq_lengths) + + concatenated_list = [] + for img_tensor, instruct_tensor in zip(img_tensors, instruct_tensors): + feature_dim = img_tensor.shape[-1] + concatenated = img_tensor.new_zeros(batch_size, max_seq_len, feature_dim) + for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): + concatenated[i, :encoder_seq_len] = instruct_tensor[i, :encoder_seq_len] + concatenated[i, encoder_seq_len:seq_len] = img_tensor[i, : seq_len - encoder_seq_len] + concatenated_list.append(concatenated) + + return concatenated_list + + +def _split_instruction_image_features( + hidden_states: torch.Tensor, + encoder_seq_lengths: list[int], + seq_lengths: list[int], +) -> tuple[torch.Tensor, torch.Tensor]: + """Unpack a joint sequence back into (instruction, image) streams.""" + batch_size = hidden_states.shape[0] + feature_dim = hidden_states.shape[-1] + + max_instruct_len = max(encoder_seq_lengths) + max_img_len = max(seq_len - encoder_seq_len for seq_len, encoder_seq_len in zip(seq_lengths, encoder_seq_lengths)) + + instruct_hidden_states = hidden_states.new_zeros(batch_size, max_instruct_len, feature_dim) + img_hidden_states = hidden_states.new_zeros(batch_size, max_img_len, feature_dim) + + for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): + img_len = seq_len - encoder_seq_len + instruct_hidden_states[i, :encoder_seq_len] = hidden_states[i, :encoder_seq_len] + img_hidden_states[i, :img_len] = hidden_states[i, encoder_seq_len:seq_len] + + return instruct_hidden_states, img_hidden_states + + +class BooguImageSelfAttention(nn.Module): + """GQA self-attention with QK RMSNorm and complex RoPE. + + Replaces the upstream diffusers `Attention` + Boogu attention processors; + the vLLM-Omni `Attention` layer picks the kernel backend and handles + SP/KV-cache concerns. + """ + + def __init__(self, dim: int, num_attention_heads: int, num_kv_heads: int) -> None: + super().__init__() + self.head_dim = dim // num_attention_heads + kv_dim = self.head_dim * num_kv_heads + + self.to_q = ColumnParallelLinear(dim, dim, bias=False) + self.to_k = ColumnParallelLinear(dim, kv_dim, bias=False) + self.to_v = ColumnParallelLinear(dim, kv_dim, bias=False) + self.norm_q = RMSNorm(self.head_dim, eps=1e-5) + self.norm_k = RMSNorm(self.head_dim, eps=1e-5) + self.to_out = RowParallelLinear(dim, dim, bias=False) + + self.num_local_heads = self.to_q.output_size_per_partition // self.head_dim + self.num_local_kv_heads = self.to_k.output_size_per_partition // self.head_dim + + self.attn = Attention( + num_heads=self.num_local_heads, + head_size=self.head_dim, + softmax_scale=self.head_dim**-0.5, + causal=False, + num_kv_heads=self.num_local_kv_heads, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + rotary_emb: torch.Tensor | None, + ) -> torch.Tensor: + dtype = hidden_states.dtype + + query, _ = self.to_q(hidden_states) + key, _ = self.to_k(hidden_states) + value, _ = self.to_v(hidden_states) + + query = query.unflatten(-1, (self.num_local_heads, self.head_dim)) + key = key.unflatten(-1, (self.num_local_kv_heads, self.head_dim)) + value = value.unflatten(-1, (self.num_local_kv_heads, self.head_dim)) + + query = self.norm_q(query) + key = self.norm_k(key) + + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + query, key = query.to(dtype), key.to(dtype) + + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + attn_output = self.attn(query, key, value, attn_metadata) + attn_output = attn_output.flatten(2, 3).type_as(hidden_states) + + attn_output, _ = self.to_out(attn_output) + return attn_output + + +class BooguImageJointAttention(nn.Module): + """Joint instruction+image attention for the double-stream blocks. + + The upstream implementation keeps the per-stream Q/K/V and output + projections on the attention *processor* + (`BooguImageDoubleStreamSelfAttnProcessor*`); here they are promoted onto + this module. The final joint output projection corresponds to the upstream + `img_instruct_attn.to_out[0]`. + """ + + def __init__(self, dim: int, num_attention_heads: int, num_kv_heads: int) -> None: + super().__init__() + self.head_dim = dim // num_attention_heads + kv_dim = self.head_dim * num_kv_heads + + self.img_to_q = ColumnParallelLinear(dim, dim, bias=False) + self.img_to_k = ColumnParallelLinear(dim, kv_dim, bias=False) + self.img_to_v = ColumnParallelLinear(dim, kv_dim, bias=False) + + self.instruct_to_q = ColumnParallelLinear(dim, dim, bias=False) + self.instruct_to_k = ColumnParallelLinear(dim, kv_dim, bias=False) + self.instruct_to_v = ColumnParallelLinear(dim, kv_dim, bias=False) + + self.norm_q = RMSNorm(self.head_dim, eps=1e-5) + self.norm_k = RMSNorm(self.head_dim, eps=1e-5) + + # Per-stream output projections (attention output is head-sharded + # under TP, hence row-parallel). + self.instruct_out = RowParallelLinear(dim, dim, bias=False) + self.img_out = RowParallelLinear(dim, dim, bias=False) + # Final joint projection applied to the merged full-dim sequence. + self.to_out = ReplicatedLinear(dim, dim, bias=False) + + self.num_local_heads = self.img_to_q.output_size_per_partition // self.head_dim + self.num_local_kv_heads = self.img_to_k.output_size_per_partition // self.head_dim + + self.attn = Attention( + num_heads=self.num_local_heads, + head_size=self.head_dim, + softmax_scale=self.head_dim**-0.5, + causal=False, + num_kv_heads=self.num_local_kv_heads, + ) + + def forward( + self, + img_hidden_states: torch.Tensor, + instruct_hidden_states: torch.Tensor, + joint_attention_mask: torch.Tensor | None, + rotary_emb: torch.Tensor | None, + encoder_seq_lengths: list[int], + seq_lengths: list[int], + ) -> torch.Tensor: + dtype = img_hidden_states.dtype + batch_size = img_hidden_states.shape[0] + + img_query, _ = self.img_to_q(img_hidden_states) + img_key, _ = self.img_to_k(img_hidden_states) + img_value, _ = self.img_to_v(img_hidden_states) + + instruct_query, _ = self.instruct_to_q(instruct_hidden_states) + instruct_key, _ = self.instruct_to_k(instruct_hidden_states) + instruct_value, _ = self.instruct_to_v(instruct_hidden_states) + + query, key, value = _concat_instruction_image_features( + [img_query, img_key, img_value], + [instruct_query, instruct_key, instruct_value], + encoder_seq_lengths, + seq_lengths, + ) + + query = query.view(batch_size, -1, self.num_local_heads, self.head_dim) + key = key.view(batch_size, -1, self.num_local_kv_heads, self.head_dim) + value = value.view(batch_size, -1, self.num_local_kv_heads, self.head_dim) + + query = self.norm_q(query) + key = self.norm_k(key) + + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + query, key = query.to(dtype), key.to(dtype) + + attn_metadata = AttentionMetadata(attn_mask=joint_attention_mask) if joint_attention_mask is not None else None + attn_output = self.attn(query, key, value, attn_metadata) + attn_output = attn_output.flatten(2, 3).to(dtype) + + instruct_attn_out, img_attn_out = _split_instruction_image_features( + attn_output, encoder_seq_lengths, seq_lengths + ) + instruct_projected, _ = self.instruct_out(instruct_attn_out) + img_projected, _ = self.img_out(img_attn_out) + + merged = _concat_instruction_image_features( + [img_projected], [instruct_projected], encoder_seq_lengths, seq_lengths + )[0] + merged, _ = self.to_out(merged) + return merged + + +class BooguImageTransformerBlock(nn.Module): + """Single-stream block: GQA attention + SwiGLU FFN with optional AdaRMS modulation.""" + + def __init__( + self, + dim: int, + num_attention_heads: int, + num_kv_heads: int, + multiple_of: int, + ffn_dim_multiplier: float | None, + norm_eps: float, + modulation: bool = True, + ) -> None: + super().__init__() + self.head_dim = dim // num_attention_heads + self.modulation = modulation + + self.attn = BooguImageSelfAttention(dim, num_attention_heads, num_kv_heads) + self.feed_forward = LuminaFeedForward( + dim=dim, + inner_dim=4 * dim, + multiple_of=multiple_of, + ffn_dim_multiplier=ffn_dim_multiplier, + ) + + if modulation: + self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + else: + self.norm1 = RMSNorm(dim, eps=norm_eps) + + self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.norm2 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + image_rotary_emb: torch.Tensor | None, + temb: torch.Tensor | None = None, + ) -> torch.Tensor: + if self.modulation: + if temb is None: + raise ValueError("temb must be provided when modulation is enabled") + norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) + attn_output = self.attn(norm_hidden_states, attention_mask, image_rotary_emb) + hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output) + mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1))) + hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output) + else: + norm_hidden_states = self.norm1(hidden_states) + attn_output = self.attn(norm_hidden_states, attention_mask, image_rotary_emb) + hidden_states = hidden_states + self.norm2(attn_output) + mlp_output = self.feed_forward(self.ffn_norm1(hidden_states)) + hidden_states = hidden_states + self.ffn_norm2(mlp_output) + + return hidden_states + + +class BooguImageNoiseRefinerTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageRefImgRefinerTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageContextRefinerTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageSingleStreamTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageDoubleStreamTransformerBlock(nn.Module): + """Double-stream block: instruction and image tokens are processed in + parallel streams coupled by a joint attention, plus an extra image + self-attention.""" + + def __init__( + self, + dim: int, + num_attention_heads: int, + num_kv_heads: int, + multiple_of: int, + ffn_dim_multiplier: float | None, + norm_eps: float, + modulation: bool = True, + ) -> None: + super().__init__() + self.head_dim = dim // num_attention_heads + self.num_attention_heads = num_attention_heads + self.modulation = modulation + self.hidden_size = dim + + self.img_instruct_attn = BooguImageJointAttention(dim, num_attention_heads, num_kv_heads) + self.img_self_attn = BooguImageSelfAttention(dim, num_attention_heads, num_kv_heads) + + self.img_feed_forward = LuminaFeedForward( + dim=dim, + inner_dim=4 * dim, + multiple_of=multiple_of, + ffn_dim_multiplier=ffn_dim_multiplier, + ) + + if modulation: + # Image modulation terms: cross-attn, MLP, self-attn. + self.img_norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + self.img_norm2 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + self.img_norm3 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + else: + self.img_norm1 = RMSNorm(dim, eps=norm_eps) + self.img_norm2 = RMSNorm(dim, eps=norm_eps) + self.img_norm3 = RMSNorm(dim, eps=norm_eps) + + self.img_ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.img_attn_norm = RMSNorm(dim, eps=norm_eps) + self.img_self_attn_norm = RMSNorm(dim, eps=norm_eps) + self.img_ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + self.instruct_feed_forward = LuminaFeedForward( + dim=dim, + inner_dim=4 * dim, + multiple_of=multiple_of, + ffn_dim_multiplier=ffn_dim_multiplier, + ) + + if modulation: + # Instruction modulation terms: cross-attn, MLP. + self.instruct_norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + self.instruct_norm2 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + else: + self.instruct_norm1 = RMSNorm(dim, eps=norm_eps) + self.instruct_norm2 = RMSNorm(dim, eps=norm_eps) + + self.instruct_ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.instruct_attn_norm = RMSNorm(dim, eps=norm_eps) + self.instruct_ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + def forward( + self, + img_hidden_states: torch.Tensor, # [B, L_img, D] (ref_img + noise_img) + instruct_hidden_states: torch.Tensor, # [B, L_instruct, D] + img_attention_mask: torch.Tensor | None, # [B, L_img] + joint_attention_mask: torch.Tensor | None, # [B, L_total] + image_rotary_emb: torch.Tensor | None, # [B, L_img, head_dim // 2] + rotary_emb: torch.Tensor | None, # [B, L_total, head_dim // 2] + temb: torch.Tensor | None = None, + encoder_seq_lengths: list[int] | None = None, + seq_lengths: list[int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.modulation and temb is None: + raise ValueError("temb must be provided when modulation is enabled") + + batch_size = img_hidden_states.shape[0] + L_instruct = instruct_hidden_states.shape[1] + L_img = img_hidden_states.shape[1] + + if self.modulation: + img_norm1_out, img_gate_msa, img_scale_mlp, img_gate_mlp = self.img_norm1(img_hidden_states, temb) + img_norm2_out, img_shift_mlp, _, _ = self.img_norm2(img_hidden_states, temb) + img_norm3_out, img_gate_self, _, _ = self.img_norm3(img_hidden_states, temb) + + ( + instruct_norm1_out, + instruct_gate_msa, + instruct_scale_mlp, + instruct_gate_mlp, + ) = self.instruct_norm1(instruct_hidden_states, temb) + instruct_norm2_out, instruct_shift_mlp, _, _ = self.instruct_norm2(instruct_hidden_states, temb) + + joint_attn_out = self.img_instruct_attn( + img_norm1_out, + instruct_norm1_out, + joint_attention_mask, + rotary_emb, + encoder_seq_lengths, + seq_lengths, + ) + + # Split the joint output back into instruction/image segments. + instruct_attn_out = instruct_hidden_states.new_zeros(batch_size, L_instruct, self.hidden_size) + img_attn_out = img_hidden_states.new_zeros(batch_size, L_img, self.hidden_size) + for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): + instruct_attn_out[i, :encoder_seq_len] = joint_attn_out[i, :encoder_seq_len] + img_attn_out[i, : seq_len - encoder_seq_len] = joint_attn_out[i, encoder_seq_len:seq_len] + + img_self_attn_out = self.img_self_attn(img_norm3_out, img_attention_mask, image_rotary_emb) + + img_hidden_states = img_hidden_states + img_gate_msa.unsqueeze(1).tanh() * self.img_attn_norm(img_attn_out) + img_hidden_states = img_hidden_states + img_gate_self.unsqueeze(1).tanh() * self.img_self_attn_norm( + img_self_attn_out + ) + + img_mlp_input = (1 + img_scale_mlp.unsqueeze(1)) * img_norm2_out + img_shift_mlp.unsqueeze(1) + img_mlp_out = self.img_feed_forward(self.img_ffn_norm1(img_mlp_input)) + img_hidden_states = img_hidden_states + img_gate_mlp.unsqueeze(1).tanh() * self.img_ffn_norm2(img_mlp_out) + + instruct_hidden_states = instruct_hidden_states + instruct_gate_msa.unsqueeze( + 1 + ).tanh() * self.instruct_attn_norm(instruct_attn_out) + + instruct_mlp_input = ( + 1 + instruct_scale_mlp.unsqueeze(1) + ) * instruct_norm2_out + instruct_shift_mlp.unsqueeze(1) + instruct_mlp_out = self.instruct_feed_forward(self.instruct_ffn_norm1(instruct_mlp_input)) + instruct_hidden_states = instruct_hidden_states + instruct_gate_mlp.unsqueeze( + 1 + ).tanh() * self.instruct_ffn_norm2(instruct_mlp_out) + + else: + img_norm1_out = self.img_norm1(img_hidden_states) + img_norm3_out = self.img_norm3(img_hidden_states) + instruct_norm1_out = self.instruct_norm1(instruct_hidden_states) + + joint_attn_out = self.img_instruct_attn( + img_norm1_out, + instruct_norm1_out, + joint_attention_mask, + rotary_emb, + encoder_seq_lengths, + seq_lengths, + ) + + instruct_attn_out = instruct_hidden_states.new_zeros(batch_size, L_instruct, self.hidden_size) + img_attn_out = img_hidden_states.new_zeros(batch_size, L_img, self.hidden_size) + for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): + instruct_attn_out[i, :encoder_seq_len] = joint_attn_out[i, :encoder_seq_len] + img_attn_out[i, : seq_len - encoder_seq_len] = joint_attn_out[i, encoder_seq_len:seq_len] + + img_self_attn_out = self.img_self_attn(img_norm3_out, img_attention_mask, image_rotary_emb) + + img_hidden_states = img_hidden_states + self.img_attn_norm(img_attn_out) + img_hidden_states = img_hidden_states + self.img_self_attn_norm(img_self_attn_out) + img_norm2_out = self.img_norm2(img_hidden_states) + img_mlp_out = self.img_feed_forward(self.img_ffn_norm1(img_norm2_out)) + img_hidden_states = img_hidden_states + self.img_ffn_norm2(img_mlp_out) + + instruct_hidden_states = instruct_hidden_states + self.instruct_attn_norm(instruct_attn_out) + instruct_norm2_out = self.instruct_norm2(instruct_hidden_states) + instruct_mlp_out = self.instruct_feed_forward(self.instruct_ffn_norm1(instruct_norm2_out)) + instruct_hidden_states = instruct_hidden_states + self.instruct_ffn_norm2(instruct_mlp_out) + + return img_hidden_states, instruct_hidden_states + + +def _cal_preprocessed_instruction_feat_dim(instruction_feature_configs: dict) -> int: + num_instruction_feature_layers = max(instruction_feature_configs.get("num_instruction_feature_layers", 1), 1) + instruction_feat_dim = instruction_feature_configs.get("instruction_feat_dim", 4096) + reduce_type = instruction_feature_configs.get("reduce_type", "concat") + if "cat" in reduce_type.lower(): + return num_instruction_feature_layers * instruction_feat_dim + elif "mean" in reduce_type.lower(): + return instruction_feat_dim + else: + raise ValueError(f"Invalid reduce_type: {reduce_type}") + + +class BooguImageTransformer2DModel(nn.Module): + """Boogu-Image transformer with mixed stream topology. + + Early layers use double-stream (dual-stream) processing, then switch to + single-stream joint processing. Context/noise/reference-image refiner + blocks run before the main stack. + """ + + # Noise refiner, reference image refiner, and double stream layers are + # kept out of regional torch.compile for numerical stability (matches the + # upstream `_repeated_blocks`). + _repeated_blocks = [ + "BooguImageTransformerBlock", + "BooguImageContextRefinerTransformerBlock", + "BooguImageSingleStreamTransformerBlock", + ] + _layerwise_offload_blocks_attrs = ["single_stream_layers", "double_stream_layers"] + + def __init__(self, od_config: OmniDiffusionConfig) -> None: + super().__init__() + self.od_config = od_config + cfg = od_config.tf_model_config + + patch_size = cfg.get("patch_size", 2) + in_channels = cfg.get("in_channels", 16) + out_channels = cfg.get("out_channels", None) + hidden_size = cfg.get("hidden_size", 2304) + num_layers = cfg.get("num_layers", 26) + num_double_stream_layers = cfg.get("num_double_stream_layers", 2) + num_refiner_layers = cfg.get("num_refiner_layers", 2) + num_attention_heads = cfg.get("num_attention_heads", 24) + num_kv_heads = cfg.get("num_kv_heads", 8) + multiple_of = cfg.get("multiple_of", 256) + ffn_dim_multiplier = cfg.get("ffn_dim_multiplier", None) + norm_eps = cfg.get("norm_eps", 1e-5) + axes_dim_rope = tuple(cfg.get("axes_dim_rope", (40, 40, 40))) + axes_lens = tuple(cfg.get("axes_lens", (2048, 1664, 1664))) + instruction_feature_configs = cfg.get( + "instruction_feature_configs", + {"instruction_feat_dim": 1024, "reduce_type": "mean", "num_instruction_feature_layers": 1}, + ) + prompt_tuning_configs = cfg.get("prompt_tuning_configs", {"use_prompt_tuning": False}) + timestep_scale = cfg.get("timestep_scale", 1.0) + + if (hidden_size // num_attention_heads) != sum(axes_dim_rope): + raise ValueError( + f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) " + f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})" + ) + if num_double_stream_layers > num_layers: + raise ValueError( + f"num_double_stream_layers ({num_double_stream_layers}) cannot be greater than " + f"num_layers ({num_layers})" + ) + if prompt_tuning_configs.get("use_prompt_tuning", False): + raise NotImplementedError( + "Prompt tuning is a training-time feature and is not supported by the " + "native vLLM-Omni BooguImageTransformer2DModel." + ) + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = out_channels or in_channels + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.num_double_stream_layers = num_double_stream_layers + self.num_single_stream_layers = num_layers - num_double_stream_layers + self.axes_dim_rope = axes_dim_rope + self.axes_lens = axes_lens + self.instruction_feature_configs = instruction_feature_configs + self.preprocessed_instruction_feat_dim = _cal_preprocessed_instruction_feat_dim(instruction_feature_configs) + + self.rope_embedder = BooguImageDoubleStreamRotaryPosEmbed( + theta=10000, + axes_dim=axes_dim_rope, + axes_lens=axes_lens, + patch_size=patch_size, + ) + + self.x_embedder = nn.Linear( + in_features=patch_size * patch_size * in_channels, + out_features=hidden_size, + ) + self.ref_image_patch_embedder = nn.Linear( + in_features=patch_size * patch_size * in_channels, + out_features=hidden_size, + ) + + self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding( + hidden_size=hidden_size, + instruction_feat_dim=self.preprocessed_instruction_feat_dim, + norm_eps=norm_eps, + timestep_scale=timestep_scale, + ) + + self.noise_refiner = nn.ModuleList( + [ + BooguImageNoiseRefinerTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_refiner_layers) + ] + ) + + self.ref_image_refiner = nn.ModuleList( + [ + BooguImageRefImgRefinerTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_refiner_layers) + ] + ) + + self.context_refiner = nn.ModuleList( + [ + BooguImageContextRefinerTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=False, + ) + for _ in range(num_refiner_layers) + ] + ) + + # Mixed architecture: dual-stream first, then single-stream. + self.double_stream_layers = nn.ModuleList( + [ + BooguImageDoubleStreamTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_double_stream_layers) + ] + ) + + self.single_stream_layers = nn.ModuleList( + [ + BooguImageSingleStreamTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(self.num_single_stream_layers) + ] + ) + + self.norm_out = LuminaLayerNormContinuous( + embedding_dim=hidden_size, + conditioning_embedding_dim=min(hidden_size, 1024), + elementwise_affine=False, + eps=1e-6, + bias=True, + out_dim=patch_size * patch_size * self.out_channels, + ) + + # Distinguish multiple reference images (max 5). + self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) + + def preprocess_instruction_hidden_states(self, raw_instruction_hidden_states): + """Reduce the raw MLLM hidden states to the transformer feature dim. + + Mirrors upstream ``preprocess_instruction_hidden_states``: a single + tensor passes through unchanged; a list of per-layer states is combined + by ``concat`` or ``mean`` according to ``instruction_feature_configs``. + """ + cfg = self.instruction_feature_configs + num_instruction_feature_layers = max(cfg.get("num_instruction_feature_layers", 1), 1) + reduce_type = cfg.get("reduce_type", "concat") + + if isinstance(raw_instruction_hidden_states, torch.Tensor): + instruction_hidden_states = raw_instruction_hidden_states + elif isinstance(raw_instruction_hidden_states, (list, tuple)): + assert len(raw_instruction_hidden_states) == num_instruction_feature_layers + if "cat" in reduce_type.lower(): + instruction_hidden_states = torch.cat(raw_instruction_hidden_states, dim=-1) + elif "mean" in reduce_type.lower(): + instruction_hidden_states = torch.mean(torch.stack(raw_instruction_hidden_states), dim=0) + else: + raise ValueError(f"Invalid reduce_type: {reduce_type}") + else: + raise ValueError( + "Invalid type of raw_instruction_hidden_states, expected torch.Tensor or list, " + f"but got {type(raw_instruction_hidden_states)}" + ) + + assert self.preprocessed_instruction_feat_dim == instruction_hidden_states.shape[-1] + return instruction_hidden_states + + def flat_and_pad_to_seq(self, hidden_states, ref_image_hidden_states): + """Flatten patch tokens and pad to batched sequences. + + Ported from upstream; for text-to-image ``ref_image_hidden_states`` is + ``None`` and the reference-image branch collapses to zero-length. + """ + batch_size = len(hidden_states) + p = self.patch_size + device = hidden_states[0].device + + img_sizes = [(img.size(1), img.size(2)) for img in hidden_states] + l_effective_img_len = [(H // p) * (W // p) for (H, W) in img_sizes] + + if ref_image_hidden_states is not None: + ref_img_sizes = [ + [(img.size(1), img.size(2)) for img in imgs] if imgs is not None else None + for imgs in ref_image_hidden_states + ] + l_effective_ref_img_len = [ + [(ref_img_size[0] // p) * (ref_img_size[1] // p) for ref_img_size in _ref_img_sizes] + if _ref_img_sizes is not None + else [0] + for _ref_img_sizes in ref_img_sizes + ] + else: + ref_img_sizes = [None for _ in range(batch_size)] + l_effective_ref_img_len = [[0] for _ in range(batch_size)] + + max_ref_img_len = max(sum(ref_img_len) for ref_img_len in l_effective_ref_img_len) + max_img_len = max(l_effective_img_len) + + # Reference-image patch embeddings. + flat_ref_img_hidden_states = [] + for i in range(batch_size): + if ref_img_sizes[i] is not None: + imgs = [] + for ref_img in ref_image_hidden_states[i]: + C, H, W = ref_img.size() + ref_img = rearrange(ref_img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p) + imgs.append(ref_img) + flat_ref_img_hidden_states.append(torch.cat(imgs, dim=0)) + else: + flat_ref_img_hidden_states.append(None) + + # Noise-image patch embeddings. + flat_hidden_states = [] + for i in range(batch_size): + img = hidden_states[i] + C, H, W = img.size() + img = rearrange(img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p) + flat_hidden_states.append(img) + + padded_ref_img_hidden_states = torch.zeros( + batch_size, + max_ref_img_len, + flat_hidden_states[0].shape[-1], + device=device, + dtype=flat_hidden_states[0].dtype, + ) + padded_ref_img_mask = torch.zeros(batch_size, max_ref_img_len, dtype=torch.bool, device=device) + for i in range(batch_size): + if ref_img_sizes[i] is not None: + padded_ref_img_hidden_states[i, : sum(l_effective_ref_img_len[i])] = flat_ref_img_hidden_states[i] + padded_ref_img_mask[i, : sum(l_effective_ref_img_len[i])] = True + + padded_hidden_states = torch.zeros( + batch_size, + max_img_len, + flat_hidden_states[0].shape[-1], + device=device, + dtype=flat_hidden_states[0].dtype, + ) + padded_img_mask = torch.zeros(batch_size, max_img_len, dtype=torch.bool, device=device) + for i in range(batch_size): + padded_hidden_states[i, : l_effective_img_len[i]] = flat_hidden_states[i] + padded_img_mask[i, : l_effective_img_len[i]] = True + + return ( + padded_hidden_states, + padded_ref_img_hidden_states, + padded_img_mask, + padded_ref_img_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + ) + + def img_patch_embed_and_refine( + self, + hidden_states, + ref_image_hidden_states, + padded_img_mask, + padded_ref_img_mask, + noise_rotary_emb, + ref_img_rotary_emb, + l_effective_ref_img_len, + l_effective_img_len, + temb, + ): + """Embed image patches and run the refiner blocks. + + The reference-image refiner is skipped when there are no reference-image + tokens (text-to-image), which is numerically identical to upstream (the + combined sequence only reads ``[:sum(ref_img_len)]`` = empty) while + avoiding a degenerate zero-length attention. + """ + batch_size = len(hidden_states) + max_combined_img_len = max( + img_len + sum(ref_img_len) for img_len, ref_img_len in zip(l_effective_img_len, l_effective_ref_img_len) + ) + + hidden_states = self.x_embedder(hidden_states) + ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states) + + for i in range(batch_size): + shift = 0 + for j, ref_img_len in enumerate(l_effective_ref_img_len[i]): + ref_image_hidden_states[i, shift : shift + ref_img_len, :] = ( + ref_image_hidden_states[i, shift : shift + ref_img_len, :] + self.image_index_embedding[j] + ) + shift += ref_img_len + + for layer in self.noise_refiner: + hidden_states = layer(hidden_states, padded_img_mask, noise_rotary_emb, temb) + + flat_l_effective_ref_img_len = list(itertools.chain(*l_effective_ref_img_len)) + num_ref_images = len(flat_l_effective_ref_img_len) + max_ref_img_len = max(flat_l_effective_ref_img_len) + + if max_ref_img_len > 0: + batch_ref_img_mask = ref_image_hidden_states.new_zeros(num_ref_images, max_ref_img_len, dtype=torch.bool) + batch_ref_image_hidden_states = ref_image_hidden_states.new_zeros( + num_ref_images, max_ref_img_len, self.hidden_size + ) + batch_ref_img_rotary_emb = hidden_states.new_zeros( + num_ref_images, max_ref_img_len, ref_img_rotary_emb.shape[-1], dtype=ref_img_rotary_emb.dtype + ) + batch_temb = temb.new_zeros(num_ref_images, *temb.shape[1:], dtype=temb.dtype) + + # Flatten reference images into a temporary batch. + idx = 0 + for i in range(batch_size): + shift = 0 + for ref_img_len in l_effective_ref_img_len[i]: + batch_ref_img_mask[idx, :ref_img_len] = True + batch_ref_image_hidden_states[idx, :ref_img_len] = ref_image_hidden_states[ + i, shift : shift + ref_img_len + ] + batch_ref_img_rotary_emb[idx, :ref_img_len] = ref_img_rotary_emb[i, shift : shift + ref_img_len] + batch_temb[idx] = temb[i] + shift += ref_img_len + idx += 1 + + for layer in self.ref_image_refiner: + batch_ref_image_hidden_states = layer( + batch_ref_image_hidden_states, batch_ref_img_mask, batch_ref_img_rotary_emb, batch_temb + ) + + # Restore reference-image sequence layout. + idx = 0 + for i in range(batch_size): + shift = 0 + for ref_img_len in l_effective_ref_img_len[i]: + ref_image_hidden_states[i, shift : shift + ref_img_len] = batch_ref_image_hidden_states[ + idx, :ref_img_len + ] + shift += ref_img_len + idx += 1 + + combined_img_hidden_states = hidden_states.new_zeros(batch_size, max_combined_img_len, self.hidden_size) + for i, (ref_img_len, img_len) in enumerate(zip(l_effective_ref_img_len, l_effective_img_len)): + combined_img_hidden_states[i, : sum(ref_img_len)] = ref_image_hidden_states[i, : sum(ref_img_len)] + combined_img_hidden_states[i, sum(ref_img_len) : sum(ref_img_len) + img_len] = hidden_states[i, :img_len] + + return combined_img_hidden_states + + def forward( + self, + hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.Tensor, + instruction_hidden_states: torch.Tensor, + freqs_cis: torch.Tensor, + instruction_attention_mask: torch.Tensor, + ref_image_hidden_states: list[list[torch.Tensor]] | None = None, + ) -> torch.Tensor: + """Denoise one step: refiner -> double-stream -> fuse -> single-stream -> unpatchify. + + Ported from upstream ``BooguImageTransformer2DModel.forward`` with the + TeaCache/TaylorSeer/PEFT/gradient-checkpointing branches removed. Returns + the velocity prediction as a ``[B, C, H_lat, W_lat]`` tensor. + """ + instruction_hidden_states = self.preprocess_instruction_hidden_states(instruction_hidden_states) + + batch_size = len(hidden_states) + is_hidden_states_tensor = isinstance(hidden_states, torch.Tensor) + if is_hidden_states_tensor: + assert hidden_states.ndim == 4 + hidden_states = [_hidden_states for _hidden_states in hidden_states] + + device = hidden_states[0].device + + # Timestep and instruction embedding. + temb, instruction_hidden_states = self.time_caption_embed( + timestep, instruction_hidden_states, hidden_states[0].dtype + ) + + # Flatten and pad token sequences. + ( + hidden_states, + ref_image_hidden_states, + img_mask, + ref_img_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + ) = self.flat_and_pad_to_seq(hidden_states, ref_image_hidden_states) + + # Build rotary embeddings and sequence lengths. + ( + context_rotary_emb, + ref_img_rotary_emb, + noise_rotary_emb, + rotary_emb, + encoder_seq_lengths, + seq_lengths, + combined_img_rotary_emb, + combined_img_seq_lengths, + ) = self.rope_embedder( + freqs_cis, + instruction_attention_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + device, + ) + + # Context refinement. + for layer in self.context_refiner: + instruction_hidden_states = layer(instruction_hidden_states, instruction_attention_mask, context_rotary_emb) + + # Image patch embedding and refinement. + combined_img_hidden_states = self.img_patch_embed_and_refine( + hidden_states, + ref_image_hidden_states, + img_mask, + ref_img_mask, + noise_rotary_emb, + ref_img_rotary_emb, + l_effective_ref_img_len, + l_effective_img_len, + temb, + ) + + instruct_hidden_states = instruction_hidden_states + img_hidden_states = combined_img_hidden_states + + # Joint mask for [instruct + image]. + max_seq_len = max(seq_lengths) + joint_attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool) + for i, seq_len in enumerate(seq_lengths): + joint_attention_mask[i, :seq_len] = True + + # Dual-stream (double-stream) stage. + if self.num_double_stream_layers > 0: + max_img_len = max(combined_img_seq_lengths) + img_attention_mask = hidden_states.new_zeros(batch_size, max_img_len, dtype=torch.bool) + for i, img_seq_len in enumerate(combined_img_seq_lengths): + img_attention_mask[i, :img_seq_len] = True + + for layer in self.double_stream_layers: + img_hidden_states, instruct_hidden_states = layer( + img_hidden_states, + instruct_hidden_states, + img_attention_mask, + joint_attention_mask, + combined_img_rotary_emb, + rotary_emb, + temb, + encoder_seq_lengths, + seq_lengths, + ) + + # Fuse streams to joint sequence. + joint_hidden_states = hidden_states.new_zeros(batch_size, max(seq_lengths), self.hidden_size) + for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): + joint_hidden_states[i, :encoder_seq_len] = instruct_hidden_states[i, :encoder_seq_len] + joint_hidden_states[i, encoder_seq_len:seq_len] = img_hidden_states[i, : seq_len - encoder_seq_len] + + # Single-stream stage. + hidden_states = joint_hidden_states + for layer in self.single_stream_layers: + hidden_states = layer(hidden_states, joint_attention_mask, rotary_emb, temb) + + # Output projection. + hidden_states = self.norm_out(hidden_states, temb) + + # Reshape back to image format. + p = self.patch_size + output = [] + for i, (img_size, img_len, seq_len) in enumerate(zip(img_sizes, l_effective_img_len, seq_lengths)): + height, width = img_size + img_tokens = hidden_states[i][seq_len - img_len : seq_len] + img_output = rearrange( + img_tokens, + "(h w) (p1 p2 c) -> c (h p1) (w p2)", + h=height // p, + w=width // p, + p1=p, + p2=p, + ) + output.append(img_output) + + if is_hidden_states_tensor: + output = torch.stack(output, dim=0) + + return output + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load diffusers-named checkpoint weights into the native module. + + Two name promotions relative to upstream (see step 8/10 findings): + + - ``*.img_instruct_attn.processor.{img,instruct}_{to_q,to_k,to_v}`` / + ``{instruct,img}_out`` -> drop ``.processor`` (upstream keeps the + joint-attention projections on the attention processor; the native + module hosts them directly). + - ``*.to_out.0.weight`` -> ``*.to_out.weight`` (diffusers wraps the + output projection in a ``ModuleList``; the native module uses a plain + linear). + """ + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + original_name = name + if ".img_instruct_attn.processor." in name: + name = name.replace(".img_instruct_attn.processor.", ".img_instruct_attn.") + if ".to_out.0." in name: + name = name.replace(".to_out.0.", ".to_out.") + + if name not in params_dict: + logger.warning("Skipping unexpected checkpoint weight %s", original_name) + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + unloaded_params = sorted(params_dict.keys() - loaded_params) + if unloaded_params: + logger.warning("Model parameters not loaded from checkpoint: %s", unloaded_params) + + return loaded_params diff --git a/vllm_omni/diffusion/models/boogu_image/image_processor.py b/vllm_omni/diffusion/models/boogu_image/image_processor.py new file mode 100644 index 000000000..3edf67bcc --- /dev/null +++ b/vllm_omni/diffusion/models/boogu_image/image_processor.py @@ -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 diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py new file mode 100644 index 000000000..7ef6877c8 --- /dev/null +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -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) diff --git a/vllm_omni/diffusion/models/boogu_image/scheduling_flow_match_euler_discrete_time_shifting.py b/vllm_omni/diffusion/models/boogu_image/scheduling_flow_match_euler_discrete_time_shifting.py new file mode 100644 index 000000000..656594221 --- /dev/null +++ b/vllm_omni/diffusion/models/boogu_image/scheduling_flow_match_euler_discrete_time_shifting.py @@ -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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index aea257dd6..41412424d 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -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", diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 42181419c..3ad3f1541 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -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: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index a9251a8d8..c9ddab2e9 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -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: diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 73ac80642..f576b9724 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -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