From 8f284b343af4904225088ff8a9ffab4d43768568 Mon Sep 17 00:00:00 2001 From: hurukawa <61525444+nagisa-kunhah@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:37:21 +0800 Subject: [PATCH] [Feature][Diffusion]: CFG parallelism support for Boogu-Image (#6786) Signed-off-by: nagisa <1434936049@qq.com> Signed-off-by: nagisa-kun <1434936049@qq.com> --- .buildkite/cuda/test-nightly.yml | 11 + .../diffusion/parallelism/cfg_parallel.md | 2 +- docs/user_guide/diffusion_features.md | 1 + recipes/Boogu/Boogu-Image.md | 94 ++++++- .../test_boogu_image_edit_vllm_omni.json | 181 ++++++++++++- .../tests/test_boogu_image_vllm_omni.json | 44 ++++ .../distributed/test_cfg_parallel.py | 69 ++++- .../boogu_image/test_pipeline_boogu_image.py | 100 +++++++- .../test_boogu_image_edit_expansion.py | 84 +++++- .../test_boogu_image_expansion.py | 18 +- .../diffusion/distributed/cfg_parallel.py | 21 +- .../diffusion/distributed/parallel_state.py | 5 + .../boogu_image/pipeline_boogu_image.py | 240 ++++++++++++++---- 13 files changed, 792 insertions(+), 78 deletions(-) diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 1e6bc8f4a..716bbc6f5 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -229,6 +229,17 @@ steps: --run-level "full_model" mirror_hardwares: h100_2 + - label: ":full_moon: Diffusion X2I(&A&T) · Boogu Edit CFG · 3-GPU" + timeout_in_minutes: 120 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_boogu_image_edit_expansion.py + -k "test_double_guidance_cfg_parallel and cfg_parallel_3_full" + -m "slow and diffusion and cuda and H100 and cards_3" + --run-level "full_model" + mirror_hardwares: h100_3 + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · 4-GPU" timeout_in_minutes: 120 commands: diff --git a/docs/user_guide/diffusion/parallelism/cfg_parallel.md b/docs/user_guide/diffusion/parallelism/cfg_parallel.md index 554110668..544754403 100644 --- a/docs/user_guide/diffusion/parallelism/cfg_parallel.md +++ b/docs/user_guide/diffusion/parallelism/cfg_parallel.md @@ -94,7 +94,7 @@ In `DiffusionParallelConfig` !!! info - Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism. + Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** and **Boogu-Image-Edit double guidance** also support `cfg_parallel_size=3` for full three-way branch execution. Boogu double guidance additionally supports size 2 via round-robin assignment of its three branches. --- diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 3f6a33878..9e867dd47 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -117,6 +117,7 @@ The following tables show which models support each feature: | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | |--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| | **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| **Boogu-Image** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index b1edf0b4e..741ea606b 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -39,9 +39,10 @@ rendering. The Edit and Edit-Turbo checkpoints use the same native ## 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). +This recipe documents supported configurations for CUDA GPU serving. The native +pipeline supports single-GPU inference and classifier-free-guidance (CFG) +parallelism. CPU offload, cache acceleration, and other multi-GPU dimensions +remain unsupported (see Notes). ## GPU @@ -108,8 +109,24 @@ curl -s http://localhost:8091/v1/chat/completions \ 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. + (`--cache-backend cache_dit`), and TP / SP / HSDP multi-GPU parallelism. + +### 2 x H100 (CFG parallel, Base T2I) + +Base T2I has two predictions per guided denoising step, so two CFG ranks map +one-to-one to the positive and negative-text branches: + +```bash +vllm serve Boogu/Boogu-Image-0.1-Base \ + --omni \ + --cfg-parallel-size 2 \ + --port 8091 +``` + +Use `guidance_scale > 1.0` to activate the parallel two-branch path. Requests +with `guidance_scale=1.0` remain valid on the same server: every rank evaluates +only the positive branch, no negative embeddings are built, and no guidance is +applied. ## Image editing (Boogu-Image-0.1-Edit) @@ -124,6 +141,33 @@ is sent). vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091 ``` +For text-only or image-only guidance, use two CFG ranks: + +```bash +vllm serve Boogu/Boogu-Image-0.1-Edit \ + --omni \ + --cfg-parallel-size 2 \ + --port 8091 +``` + +Double guidance has three predictions per denoising step. It supports either +two ranks (round-robin assignment: rank 0 runs branches 0 and 2) or three ranks +(one branch per rank): + +```bash +# Practical two-GPU configuration +vllm serve Boogu/Boogu-Image-0.1-Edit \ + --omni \ + --cfg-parallel-size 2 \ + --port 8091 + +# Maximum branch parallelism +vllm serve Boogu/Boogu-Image-0.1-Edit \ + --omni \ + --cfg-parallel-size 3 \ + --port 8091 +``` + ### Verification Edit an image with `/v1/images/edits` (the model-card example — change a photo @@ -175,11 +219,29 @@ curl -s http://localhost:8091/v1/chat/completions \ 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. +- **CFG-size recommendations:** + + | Request mode | Scales | Predictions per step | Recommended CFG size | + | --- | --- | ---: | ---: | + | Base T2I text guidance | `guidance_scale > 1`, no reference | 2 | 2 | + | Edit text-only guidance | `guidance_scale > 1`, `guidance_scale_2 = 1` | 2 | 2 | + | Edit image-only guidance | `guidance_scale = 1`, `guidance_scale_2 > 1` | 2 | 2 | + | Edit double guidance | both scales `> 1` | 3 | 2 (practical) or 3 (maximum parallelism) | + | CFG off | both scales `= 1` | 1 | Reuse a size-2 server safely; no speedup expected | + + Double guidance preserves Boogu's original combination: + + ```text + pred = positive_with_reference + + (text_scale - 1) * (positive_with_reference - negative_with_reference) + + (image_scale - 1) * (negative_with_reference - uncond) + ``` - **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). +- **Same limitations** as the Base checkpoint apply except that CFG parallelism + is supported as described above; CPU offload, Cache-DiT, TP, SP, and HSDP are + still unsupported. - **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 @@ -251,3 +313,21 @@ curl -s http://localhost:8091/v1/chat/completions \ the regular Edit checkpoint applies. - **Known limitations:** the same single-GPU limitations as Base and Edit apply; CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated. + +## Performance validation + +The checked-in single-device A40 measurements remain the current Boogu +baseline: Base T2I is about 16.37 s request latency at concurrency 1 and Edit +text-only guidance is about 29.64 s for 512x512, 28-step requests, with roughly +36.8 GB peak memory. CFG-parallel A/B numbers must be collected on the same +hardware and software revision before reporting a speedup; do not compare a +multi-GPU result against these A40 numbers as if it were a controlled A/B. + +Run the dedicated perf rows in +`tests/dfx/perf/tests/test_boogu_image_vllm_omni.json` and +`tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json` on 2x/3x H100, then +record latency, throughput, and per-GPU peak memory beside the single-device +run. The expected shape is close to the general CFG-parallel range (roughly +1.5x-1.8x for balanced two-branch workloads); the three-branch Edit result will +depend on whether two-rank round-robin or full three-rank execution is used and +must be reported from measurement rather than assumed. 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 index e2f96f626..70e0728a0 100644 --- a/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json +++ b/tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json @@ -1,7 +1,7 @@ [ { "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.", + "description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. This row measures text-only edit guidance (2 predictions/step); dedicated rows below use extra-body to measure double guidance. 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", @@ -33,5 +33,184 @@ } } ] + }, + { + "test_name": "test_boogu_image_edit_double_guidance_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device sequential three-branch double-guidance control row for direct A/B comparison with cfg_parallel_size=2 and 3 on identical H100 hardware.", + "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_double_cfg1", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 28, + "num-prompts": 10, + "max-concurrency": 1, + "warmup-requests": 4, + "warmup-concurrency": 1, + "warmup-num-inference-steps": 28, + "enable-negative-prompt": true, + "extra-body": { + "guidance_scale": 5.0, + "guidance_scale_2": 2.0, + "seed": 42 + } + } + ] + }, + { + "test_name": "test_boogu_image_edit_text_cfg_parallel_2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "Two-device CFG-parallel Edit text-only guidance measurement row. Compare with the single-device row on identical H100 hardware and software.", + "server_type": "vllm-omni", + "server_params": { + "model": "Boogu/Boogu-Image-0.1-Edit", + "serve_args": { + "cfg-parallel-size": 2, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps28_i2i_text_cfg2", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 28, + "num-prompts": 10, + "max-concurrency": 1, + "warmup-requests": 4, + "warmup-concurrency": 1, + "warmup-num-inference-steps": 28, + "enable-negative-prompt": true, + "extra-body": { + "guidance_scale": 5.0, + "guidance_scale_2": 1.0, + "seed": 42 + } + } + ] + }, + { + "test_name": "test_boogu_image_edit_double_cfg_parallel_2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "Two-device round-robin execution of the three Edit double-guidance branches. Populate A/B latency, throughput, and per-GPU memory after GPU validation.", + "server_type": "vllm-omni", + "server_params": { + "model": "Boogu/Boogu-Image-0.1-Edit", + "serve_args": { + "cfg-parallel-size": 2, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps28_i2i_double_cfg2", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 28, + "num-prompts": 10, + "max-concurrency": 1, + "warmup-requests": 4, + "warmup-concurrency": 1, + "warmup-num-inference-steps": 28, + "enable-negative-prompt": true, + "extra-body": { + "guidance_scale": 5.0, + "guidance_scale_2": 2.0, + "seed": 42 + } + } + ] + }, + { + "test_name": "test_boogu_image_edit_double_cfg_parallel_3", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 3 + } + }, + "full_model", + "diffusion" + ], + "description": "Full three-device execution of the three Edit double-guidance branches. Compare directly with sequential and cfg_parallel_size=2 on identical H100 hardware.", + "server_type": "vllm-omni", + "server_params": { + "model": "Boogu/Boogu-Image-0.1-Edit", + "serve_args": { + "cfg-parallel-size": 3, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps28_i2i_double_cfg3", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 28, + "num-prompts": 10, + "max-concurrency": 1, + "warmup-requests": 4, + "warmup-concurrency": 1, + "warmup-num-inference-steps": 28, + "enable-negative-prompt": true, + "extra-body": { + "guidance_scale": 5.0, + "guidance_scale_2": 2.0, + "seed": 42 + } + } + ] } ] diff --git a/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json b/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json index 22218268e..b0f35bd4a 100644 --- a/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_boogu_image_vllm_omni.json @@ -33,5 +33,49 @@ } } ] + }, + { + "test_name": "test_boogu_image_cfg_parallel_2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "Two-device CFG-parallel Base T2I measurement row. Run beside test_boogu_image_single_device on the same H100 revision and populate latency, throughput, speedup, and per-GPU peak-memory baselines from the resulting report.", + "server_type": "vllm-omni", + "server_params": { + "model": "Boogu/Boogu-Image-0.1-Base", + "serve_args": { + "cfg-parallel-size": 2, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps28_cfg2", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 28, + "num-prompts": 10, + "max-concurrency": 1, + "warmup-requests": 4, + "warmup-concurrency": 1, + "warmup-num-inference-steps": 28, + "enable-negative-prompt": true, + "extra-body": { + "guidance_scale": 4.0, + "seed": 42 + } + } + ] } ] diff --git a/tests/diffusion/distributed/test_cfg_parallel.py b/tests/diffusion/distributed/test_cfg_parallel.py index 8991ddec3..2ade13549 100644 --- a/tests/diffusion/distributed/test_cfg_parallel.py +++ b/tests/diffusion/distributed/test_cfg_parallel.py @@ -17,7 +17,12 @@ import torch from tests.helpers.mark import hardware_marks, hardware_test from vllm_omni.diffusion.distributed import parallel_state -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap +from vllm_omni.diffusion.distributed.cfg_parallel import ( + CFGParallelMixin, + _dispatch_branches, + _get_cfg_world_size_or_one, + _wrap, +) from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, get_classifier_free_guidance_rank, @@ -393,6 +398,68 @@ def _build_pipeline_on_device( # --------------------------------------------------------------------------- +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_cfg_world_size_defaults_only_when_group_is_uninitialized(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(parallel_state, "_CFG", None) + + assert not parallel_state.is_cfg_group_initialized() + assert _get_cfg_world_size_or_one() == 1 + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_cfg_world_size_does_not_swallow_initialized_group_errors(monkeypatch: pytest.MonkeyPatch): + class BrokenCfgGroup: + @property + def world_size(self) -> int: + raise AssertionError("CFG group is broken") + + monkeypatch.setattr(parallel_state, "_CFG", BrokenCfgGroup()) + + assert parallel_state.is_cfg_group_initialized() + with pytest.raises(AssertionError, match="CFG group is broken"): + _get_cfg_world_size_or_one() + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cpu +def test_cfg_defaults_to_sequential_without_initialized_group(): + pipeline = _build_pipeline_on_device( + TestCFGPipeline, + hidden_dim=32, + model_seed=42, + dtype=torch.float32, + device=torch.device("cpu"), + ) + positive_kwargs, negative_kwargs = _make_two_branch_inputs( + batch_size=1, + channels=4, + height=4, + width=4, + dtype=torch.float32, + device=torch.device("cpu"), + input_seed=123, + ) + + with torch.no_grad(): + positive = pipeline.predict_noise(**positive_kwargs) + negative = pipeline.predict_noise(**negative_kwargs) + actual = pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=5.0, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + + expected = negative + 5.0 * (positive - negative) + torch.testing.assert_close(actual, expected) + + @pytest.mark.core_model @pytest.mark.diffusion @pytest.mark.cpu diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index 84bac1b2b..fc2b5d362 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py @@ -192,7 +192,6 @@ def test_constructor_weights_sources(boogu_pipeline): (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", @@ -227,6 +226,26 @@ def test_constructor_rejects_unsupported_execution_modes( mock_dependencies["mllm_wrapper"].model.to.assert_not_called() +@pytest.mark.parametrize("cfg_parallel_size", [2, 3]) +def test_constructor_accepts_cfg_parallel(mock_dependencies, cfg_parallel_size): + 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=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size), + ) + + pipeline = BooguImagePipeline(od_config=od_config) + + assert pipeline.od_config.parallel_config.cfg_parallel_size == cfg_parallel_size + assert hasattr(pipeline, "predict_noise_maybe_with_cfg") + assert hasattr(pipeline, "predict_noise_with_multi_branch_cfg") + + # --------------------------------------------------------------------------- # Prompt-encoding tests (deterministic fakes, no constructor) # --------------------------------------------------------------------------- @@ -526,6 +545,14 @@ def test_forward_returns_diffusion_output(): def test_forward_cfg_off_when_guidance_one(): pipeline = _make_forward_pipeline() + cfg_calls = [] + original_predict_noise_maybe_with_cfg = pipeline.predict_noise_maybe_with_cfg + + def record_cfg_call(**kwargs): + cfg_calls.append(kwargs) + return original_predict_noise_maybe_with_cfg(**kwargs) + + pipeline.predict_noise_maybe_with_cfg = record_cfg_call 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. @@ -535,6 +562,30 @@ def test_forward_cfg_off_when_guidance_one(): # Only the positive prompt is encoded when CFG is off. assert len(pipeline.processor.calls) == 1 + assert len(cfg_calls) == 2 + assert all(call["do_true_cfg"] is False for call in cfg_calls) + assert all(call["negative_kwargs"] is None for call in cfg_calls) + + +def test_double_guidance_combine_matches_legacy_formula(): + pipeline = _make_forward_pipeline() + positive_with_reference = torch.randn(2, 4, 8, 8) + negative_with_reference = torch.randn(2, 4, 8, 8) + uncond = torch.randn(2, 4, 8, 8) + text_scale = 5.0 + image_scale = 2.0 + + actual = pipeline.combine_multi_branch_cfg_noise( + [positive_with_reference, negative_with_reference, uncond], + {"text": text_scale, "image": image_scale}, + ) + legacy = ( + positive_with_reference + + (text_scale - 1.0) * (positive_with_reference - negative_with_reference) + + (image_scale - 1.0) * (negative_with_reference - uncond) + ) + + torch.testing.assert_close(actual, legacy, rtol=0, atol=1e-5) # --------------------------------------------------------------------------- @@ -802,6 +853,39 @@ def test_build_ref_latents_expands_per_output_and_handles_none(): assert ref_latents[2] is None and ref_latents[3] is None +def test_vae_attention_context_is_noop_on_cpu(): + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import ( + BooguImagePipeline, + ) + + with BooguImagePipeline._vae_attention_context(torch.device("cpu")): + pass + + +def test_vae_attention_context_prioritizes_efficient_cuda(monkeypatch): + from vllm_omni.diffusion.models.boogu_image import pipeline_boogu_image + + recorded = {} + + def fake_sdpa_kernel(backends, *, set_priority): + recorded["backends"] = backends + recorded["set_priority"] = set_priority + return pipeline_boogu_image.nullcontext() + + monkeypatch.setattr(pipeline_boogu_image, "sdpa_kernel", fake_sdpa_kernel) + + with pipeline_boogu_image.BooguImagePipeline._vae_attention_context(torch.device("cuda")): + pass + + assert recorded == { + "backends": [ + pipeline_boogu_image.SDPBackend.EFFICIENT_ATTENTION, + pipeline_boogu_image.SDPBackend.MATH, + ], + "set_priority": True, + } + + # --------------------------------------------------------------------------- # Editing / TI2I: forward CFG branch selection # --------------------------------------------------------------------------- @@ -812,9 +896,11 @@ class _RecordingRefTransformer(_FakeTransformer): def __init__(self): self.calls = [] + self.instruction_embed_calls = [] def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): self.calls.append(kwargs.get("ref_image_hidden_states")) + self.instruction_embed_calls.append(instruction_embeds) return torch.zeros_like(latents) @@ -876,6 +962,9 @@ def test_forward_ti2i_text_only_two_predictions_with_ref(): # 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) + for i in range(0, len(pipeline.transformer.instruction_embed_calls), 2): + positive, negative = pipeline.transformer.instruction_embed_calls[i : i + 2] + assert positive is not negative def test_forward_ti2i_double_guidance_three_predictions(): @@ -895,6 +984,11 @@ def test_forward_ti2i_double_guidance_three_predictions(): 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 + embed_calls = pipeline.transformer.instruction_embed_calls + for i in range(0, len(embed_calls), 3): + positive, negative_with_reference, uncond = embed_calls[i : i + 3] + assert positive is not negative_with_reference + assert negative_with_reference is uncond def test_forward_ti2i_image_only_two_predictions_drop_ref(): @@ -914,6 +1008,10 @@ def test_forward_ti2i_image_only_two_predictions_drop_ref(): 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 + embed_calls = pipeline.transformer.instruction_embed_calls + for i in range(0, len(embed_calls), 2): + positive_with_reference, positive_without_reference = embed_calls[i : i + 2] + assert positive_with_reference is positive_without_reference def test_forward_ti2i_no_guidance_single_prediction_with_ref(): diff --git a/tests/e2e/online_serving/test_boogu_image_edit_expansion.py b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py index df7650b65..d5ded8802 100644 --- a/tests/e2e/online_serving/test_boogu_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py @@ -6,7 +6,8 @@ L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (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: +module adds the edit paths and CFG-parallel configurations 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 @@ -15,13 +16,19 @@ module adds the two edit paths smoke does not exercise: - ``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). +- ``test_double_guidance_cfg_parallel`` — three-branch double guidance with + CFG sizes 2 and 3. +- ``test_cfg_off_on_cfg_parallel_server`` — both scales at 1.0 on a size-2 CFG + server. Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint; the TI2I path activates when the request carries a reference image. Only a single reference image is supported. The regular Edit cases stay smoke-depth (``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step -configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet -supported for these models and are intentionally omitted. +configuration. Regular Edit CFG-parallel coverage includes two-card +text-only/image-only/CFG-off paths, two-card round-robin double guidance, and +full three-card double guidance. Edit-Turbo remains single-card-only here; +CPU offload and Cache-DiT are intentionally omitted for both checkpoints. From ``tests/``:: @@ -47,16 +54,49 @@ EDIT_TURBO_REVISION = "hotfix-1k-20260708" EDIT_PROMPT = "Change the style to a colored pencil drawing." SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) +PARALLEL_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +PARALLEL_3_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=3) def _get_diffusion_feature_cases(model: str): - """Return the single-card server configuration supported for Boogu-Image-Edit today.""" + """Return single-card and two-card CFG configurations for two-branch paths.""" return [ pytest.param( OmniServerParams(model=model), id="default", marks=SINGLE_CARD_FEATURE_MARKS, ), + pytest.param( + OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]), + id="cfg_parallel_2", + marks=PARALLEL_2_FEATURE_MARKS, + ), + ] + + +def _get_double_guidance_cases(model: str): + """Return two-rank round-robin and full three-rank double-guidance rows.""" + return [ + pytest.param( + OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]), + id="cfg_parallel_2_round_robin", + marks=PARALLEL_2_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams(model=model, server_args=["--cfg-parallel-size", "3"]), + id="cfg_parallel_3_full", + marks=PARALLEL_3_FEATURE_MARKS, + ), + ] + + +def _get_cfg_off_cases(model: str): + return [ + pytest.param( + OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]), + id="cfg_parallel_2_cfg_off", + marks=PARALLEL_2_FEATURE_MARKS, + ) ] @@ -106,6 +146,42 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien openai_client.send_diffusion_request(request_config) +@pytest.mark.parametrize("omni_server", _get_double_guidance_cases(MODEL), indirect=True) +def test_double_guidance_cfg_parallel(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + """Double guidance exercises three branches on two and three CFG ranks.""" + 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, + "guidance_scale_2": 2.0, + "seed": 42, + }, + } + openai_client.send_diffusion_request(request_config) + + +@pytest.mark.parametrize("omni_server", _get_cfg_off_cases(MODEL), indirect=True) +def test_cfg_off_on_cfg_parallel_server(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + """A CFG process group must not force negative embeddings or guidance.""" + 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": 1.0, + "guidance_scale_2": 1.0, + "seed": 42, + }, + } + openai_client.send_diffusion_request(request_config) + + @pytest.mark.parametrize( ("omni_server", "sampling_params"), _get_align_res_cases(), diff --git a/tests/e2e/online_serving/test_boogu_image_expansion.py b/tests/e2e/online_serving/test_boogu_image_expansion.py index a56e5a338..df845f2e4 100644 --- a/tests/e2e/online_serving/test_boogu_image_expansion.py +++ b/tests/e2e/online_serving/test_boogu_image_expansion.py @@ -7,14 +7,15 @@ 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: +supported server configuration: -Server rows (single card): +Server rows: - ``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``). +- ``cfg_parallel_2`` — two-card CFG branch parallelism. -Cases (one per test, each parametrized over both server rows): +Cases (one per test, each parametrized over all 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 @@ -24,8 +25,7 @@ Cases (one per test, each parametrized over both server rows): 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). +non-CFG multi-GPU parallelism remain unsupported. From ``tests/``:: @@ -49,10 +49,11 @@ T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting" NEGATIVE_PROMPT = "" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) +PARALLEL_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) def _get_diffusion_feature_cases(model: str): - """Return the single-card server configurations supported for Boogu-Image today.""" + """Return the validated single-card and CFG-parallel configurations.""" return [ pytest.param( OmniServerParams(model=model), @@ -67,6 +68,11 @@ def _get_diffusion_feature_cases(model: str): id="vae_slicing_tiling", marks=SINGLE_CARD_FEATURE_MARKS, ), + pytest.param( + OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]), + id="cfg_parallel_2", + marks=PARALLEL_2_FEATURE_MARKS, + ), ] diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 4c20f1587..1e09c9d43 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/vllm_omni/diffusion/distributed/cfg_parallel.py @@ -16,11 +16,26 @@ from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, get_classifier_free_guidance_world_size, + is_cfg_group_initialized, ) logger = init_logger(__name__) +def _get_cfg_world_size_or_one() -> int: + """Return the CFG world size, defaulting only when no group exists. + + Diffusion workers initialize the CFG process group even when its world + size is one. Pipelines are also invoked directly by unit tests and some + offline integrations, though, where no distributed groups exist. Treat + that explicit case like a one-rank CFG group. Errors from an existing CFG + coordinator must propagate instead of silently selecting sequential CFG. + """ + if not is_cfg_group_initialized(): + return 1 + return get_classifier_free_guidance_world_size() + + def _wrap(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: """Normalize prediction to tuple form.""" return pred if isinstance(pred, tuple) else (pred,) @@ -106,7 +121,7 @@ class CFGParallelMixin(metaclass=ABCMeta): """ if do_true_cfg: # Automatically detect CFG parallel configuration - cfg_parallel_ready = get_classifier_free_guidance_world_size() > 1 + cfg_parallel_ready = _get_cfg_world_size_or_one() > 1 if cfg_parallel_ready: cfg_group = get_cfg_group() @@ -253,7 +268,7 @@ class CFGParallelMixin(metaclass=ABCMeta): """ if do_true_cfg: n_branches = len(branches_kwargs) - cfg_world_size = get_classifier_free_guidance_world_size() + cfg_world_size = _get_cfg_world_size_or_one() cfg_parallel_ready = cfg_world_size > 1 if cfg_parallel_ready: @@ -474,7 +489,7 @@ class CFGParallelMixin(metaclass=ABCMeta): Returns: True if CFG parallel configuration is valid, False otherwise """ - if get_classifier_free_guidance_world_size() == 1: + if _get_cfg_world_size_or_one() == 1: return True if true_cfg_scale <= 1: diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 85ba9cd6c..335418254 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -322,6 +322,11 @@ def is_pipeline_last_stage(): # CFG +def is_cfg_group_initialized() -> bool: + """Return whether the classifier-free-guidance group has been created.""" + return _CFG is not None + + def get_cfg_group() -> GroupCoordinator: assert _CFG is not None, "classifier_free_guidance parallel group is not initialized" return _CFG diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index 02b333f4b..8b6fc3a45 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -11,17 +11,17 @@ Ported from the upstream ``boogu`` package 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. + vLLM-Omni convention, also hooked by the prompt-embed cache). +- Text-to-image and single-reference TI2I inference share one native pipeline; + CFG branches use vLLM-Omni's shared two-branch/N-branch parallel helpers. +- Instruction rewriting, prompt tuning, and vision-token stripping are not + ported. """ import json import os from collections.abc import Iterable +from contextlib import nullcontext from typing import ClassVar, cast import PIL.Image @@ -31,11 +31,13 @@ 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 torch.nn.attention import SDPBackend, sdpa_kernel 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.cfg_parallel import CFGParallelMixin 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 @@ -189,7 +191,7 @@ SYSTEM_PROMPT_4_T2I_UNIFIED = ( ) -class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, SupportImageInput): +class BooguImagePipeline(CFGParallelMixin, 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 / @@ -299,8 +301,6 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery 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"): @@ -562,6 +562,79 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery ref_image_hidden_states=ref_image_hidden_states, ) + def predict_noise(self, **kwargs) -> torch.Tensor: + """Run one Boogu CFG branch through the native transformer.""" + return self.predict(**kwargs) + + def combine_cfg_noise( + self, + positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], + negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], + true_cfg_scale: float, + cfg_normalize: bool = False, + kwargs: dict | None = None, + ) -> torch.Tensor: + """Preserve Boogu's sequential two-branch CFG operation order.""" + positive_items = positive_noise_pred if isinstance(positive_noise_pred, tuple) else (positive_noise_pred,) + negative_items = negative_noise_pred if isinstance(negative_noise_pred, tuple) else (negative_noise_pred,) + if len(positive_items) != 1 or len(negative_items) != 1: + raise ValueError("Boogu CFG expects exactly one prediction tensor per branch.") + + positive = positive_items[0] + negative = negative_items[0] + combined = positive + (true_cfg_scale - 1) * (positive - negative) + if cfg_normalize: + combined = self.cfg_normalize_function(positive, combined) + return combined + + def combine_multi_branch_cfg_noise( + self, + predictions: list[torch.Tensor], + true_cfg_scale: float | dict[str, float], + cfg_normalize: bool = False, + ) -> torch.Tensor: + """Combine Boogu CFG branches using the original operation order. + + Although the usual two- and three-branch CFG formulas can be rewritten + algebraically, changing their floating-point operation order causes + small per-step differences that accumulate across the denoise loop. + Keep the sequential Boogu implementation's order so parallel branch + combination does not introduce an additional source of numeric drift. + + The two-branch formula is:: + + positive + (scale - 1) * (positive - negative) + + The three-branch order is ``[positive_with_reference, + negative_with_reference, negative_without_reference]`` and combines as:: + + positive_with_reference + + (text_scale - 1) * (positive_with_reference - negative_with_reference) + + (image_scale - 1) * (negative_with_reference - uncond) + """ + if len(predictions) == 2: + if not isinstance(true_cfg_scale, float): + raise TypeError("Boogu two-branch CFG requires a scalar guidance scale.") + positive, negative = predictions + combined = positive + (true_cfg_scale - 1) * (positive - negative) + if cfg_normalize: + combined = self.cfg_normalize_function(positive, combined) + return combined + if len(predictions) != 3: + return super().combine_multi_branch_cfg_noise(predictions, true_cfg_scale, cfg_normalize) + if not isinstance(true_cfg_scale, dict): + raise TypeError("Boogu three-branch CFG requires text and image guidance scales.") + + positive_with_reference, negative_with_reference, uncond = predictions + combined = ( + positive_with_reference + + (true_cfg_scale["text"] - 1) * (positive_with_reference - negative_with_reference) + + (true_cfg_scale["image"] - 1) * (negative_with_reference - uncond) + ) + if cfg_normalize: + combined = self.cfg_normalize_function(positive_with_reference, combined) + return combined + def _encode_vae_image(self, img: torch.Tensor, generator=None) -> torch.Tensor: """Encode an image tensor into the VAE latent space (upstream ``encode_vae``). @@ -569,13 +642,35 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery 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) + with self._vae_attention_context(img.device): + 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) + @staticmethod + def _vae_attention_context(device: torch.device): + """Pin the VAE SDPA backend so worker topology cannot change results. + + Boogu's float32 VAE attention can otherwise choose a different SDPA + backend in the in-process CFG=1 executor and the multi-process CFG>1 + executor. The resulting small reference-latent differences accumulate + over the denoising loop. Prefer the default fast CUDA backend used by + Boogu and keep the math implementation as a CUDA fallback. + + This pin is intentionally CUDA-only. CPU and vendor-specific accelerator + backends keep their platform defaults, so topology-independent VAE + parity is not guaranteed for those devices by this context. + """ + if device.type != "cuda": + return nullcontext() + return sdpa_kernel( + [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH], + set_priority=True, + ) + def _build_ref_latents( self, preprocessed_images: list[torch.Tensor | None], @@ -722,71 +817,107 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery 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. + # 5. Denoise loop with shared sequential/parallel CFG execution. + # Reproduces the branch priority of upstream ``processing`` (double > + # text-only > image-only > t2i text). Reference latents stay attached + # to the same branches as in the original sequential implementation. 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_real, instruction_attention_mask, ref_latents + positive_kwargs = dict( + t=t, + latents=latents, + instruction_embeds=instruction_embeds, + freqs_real=freqs_real, + instruction_attention_mask=instruction_attention_mask, + ref_image_hidden_states=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_real, - negative_instruction_attention_mask, - ref_latents, + negative_with_reference_kwargs = dict( + t=t, + latents=latents, + instruction_embeds=negative_instruction_embeds, + freqs_real=freqs_real, + instruction_attention_mask=negative_instruction_attention_mask, + ref_image_hidden_states=ref_latents, ) - model_pred_drop_all = self.predict( - t, - latents, - negative_instruction_embeds, - freqs_real, - negative_instruction_attention_mask, - None, + uncond_kwargs = dict( + t=t, + latents=latents, + instruction_embeds=negative_instruction_embeds, + freqs_real=freqs_real, + instruction_attention_mask=negative_instruction_attention_mask, + ref_image_hidden_states=None, + ) + model_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={"text": text_gs, "image": image_gs}, + branches_kwargs=[positive_kwargs, negative_with_reference_kwargs, uncond_kwargs], + cfg_normalize=False, ) - 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_real, - negative_instruction_attention_mask, - ref_latents, + negative_kwargs = dict( + t=t, + latents=latents, + instruction_embeds=negative_instruction_embeds, + freqs_real=freqs_real, + instruction_attention_mask=negative_instruction_attention_mask, + ref_image_hidden_states=ref_latents, + ) + model_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=text_gs, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, ) - 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_real, instruction_attention_mask, None + negative_kwargs = dict(positive_kwargs, ref_image_hidden_states=None) + model_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=image_gs, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, ) - 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_real, - negative_instruction_attention_mask, - None, + negative_kwargs = dict( + t=t, + latents=latents, + instruction_embeds=negative_instruction_embeds, + freqs_real=freqs_real, + instruction_attention_mask=negative_instruction_attention_mask, + ref_image_hidden_states=None, + ) + model_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=text_gs, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + else: + # CFG-off requests remain valid even if the server owns a + # CFG process group: every rank evaluates only the positive + # branch and no negative embeddings are required. + model_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=False, + true_cfg_scale=1.0, + positive_kwargs=positive_kwargs, + negative_kwargs=None, + cfg_normalize=False, ) - 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] + do_true_cfg = text_gs > 1.0 or (task_type == "ti2i" and image_gs > 1.0) + latents = self.scheduler_step_maybe_with_cfg(model_pred, t, latents, do_true_cfg=do_true_cfg) latents = latents.to(dtype=instruction_embeds.dtype) progress_bar.update() @@ -799,7 +930,8 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery 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] + with self._vae_attention_context(latents.device): + 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")