From 32370968d00e34d704688c0b56792d279aac8aae Mon Sep 17 00:00:00 2001 From: prettygirlisnotme <53518583+prettygirlisnotme@users.noreply.github.com> Date: Sun, 6 Sep 2026 08:44:15 +0800 Subject: [PATCH] [Model] Add Boogu-Image Turbo native DMD path (#6699) Signed-off-by: Dragonborn <1078757180@qq.com> --- .../boogu_image/test_pipeline_boogu_image.py | 284 +++++++++++++++++- .../test_pipeline_boogu_image_cuda.py | 52 ++++ .../diffusion/models/boogu_image/__init__.py | 6 +- .../boogu_image/pipeline_boogu_image.py | 212 +++++++++++-- 4 files changed, 531 insertions(+), 23 deletions(-) create mode 100644 tests/diffusion/models/boogu_image/test_pipeline_boogu_image_cuda.py 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 fc1e0298e..d59544c3b 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py @@ -103,9 +103,10 @@ def boogu_pipeline(mock_dependencies): def test_boogu_image_pipeline_import(): - from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline + from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline, BooguImageTurboPipeline assert BooguImagePipeline is not None + assert issubclass(BooguImageTurboPipeline, BooguImagePipeline) def test_component_discovery_declarations(): @@ -492,10 +493,16 @@ class _FakeTransformer: class _FakeScheduler: + def __init__(self): + self.set_calls = 0 + self.step_calls = 0 + def set_timesteps(self, num_inference_steps, device=None, num_tokens=None): + self.set_calls += 1 self.timesteps = torch.linspace(0, 1, num_inference_steps + 1)[:-1] def step(self, model_output, t, latents, return_dict=False): + self.step_calls += 1 return (latents,) @@ -510,8 +517,15 @@ class _FakeDecodeVAE: return (torch.zeros(batch, 3, 16, 16),) -def _make_forward_pipeline(): - pipeline = _make_encode_pipeline() +def _make_forward_pipeline(pipeline_cls=None): + if pipeline_cls is None: + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + + pipeline_cls = BooguImagePipeline + base_pipeline = _make_encode_pipeline() + pipeline = object.__new__(pipeline_cls) + nn.Module.__init__(pipeline) + pipeline.__dict__.update(base_pipeline.__dict__) pipeline.transformer = _FakeTransformer() pipeline.scheduler = _FakeScheduler() pipeline.vae = _FakeDecodeVAE() @@ -558,6 +572,8 @@ def test_forward_returns_diffusion_output(): 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 + assert pipeline.scheduler.set_calls == 1 + assert pipeline.scheduler.step_calls == 2 def test_forward_cfg_off_when_guidance_one(): @@ -605,6 +621,218 @@ def test_double_guidance_combine_matches_legacy_formula(): torch.testing.assert_close(actual, legacy, rtol=0, atol=1e-5) +class _RecordingDMDTransformer(_FakeTransformer): + def __init__(self, velocity=0.0): + self.timesteps = [] + self.refs = [] + self.freqs_real_calls = [] + self.velocity = velocity + + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + self.timesteps.append(float(timestep[0])) + self.refs.append(kwargs.get("ref_image_hidden_states")) + self.freqs_real_calls.append(freqs_real) + return torch.full_like(latents, self.velocity) + + +def _make_turbo_forward_pipeline(): + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImageTurboPipeline + + pipeline = _make_forward_pipeline(BooguImageTurboPipeline) + pipeline.transformer = _RecordingDMDTransformer(velocity=0.125) + return pipeline + + +def test_turbo_dmd_uses_four_steps_with_real_rope_without_regular_scheduler(): + pipeline = _make_turbo_forward_pipeline() + renoise_calls = [] + original_renoise = pipeline._renoise_dmd_latents + + def recording_renoise(latents, sigma, generator=None): + renoise_calls.append(sigma) + return original_renoise(latents, sigma, generator) + + pipeline._renoise_dmd_latents = recording_renoise + req = _make_request_batch( + "a cat", + height=64, + width=64, + output_type="latent", + generator=torch.Generator(device="cpu").manual_seed(7), + ) + + out = pipeline.forward(req)[0] + + assert out.output.shape[0] == 1 + assert len(pipeline.transformer.timesteps) == 4 + assert len(pipeline.transformer.freqs_real_calls) == 4 + assert all( + not freq.is_complex() + for freqs_real in pipeline.transformer.freqs_real_calls + for pair in freqs_real + for freq in pair + ) + expected_sigmas = torch.linspace(0.001, 1.0, 5, dtype=torch.bfloat16)[:-1].float().tolist() + assert pipeline.transformer.timesteps == pytest.approx(expected_sigmas) + assert pipeline.scheduler.set_calls == 0 + assert pipeline.scheduler.step_calls == 0 + assert len(renoise_calls) == 3 # The final x0 must not be re-noised. + assert len(pipeline.processor.calls) == 1 # DMD has no negative/CFG encoding. + + +def test_turbo_dmd_decode_uses_vae_attention_context(): + pipeline = _make_turbo_forward_pipeline() + events = [] + + class _RecordingContext: + def __init__(self, device): + self.device = device + + def __enter__(self): + events.append(("enter", self.device.type)) + + def __exit__(self, exc_type, exc_value, traceback): + events.append(("exit", self.device.type)) + + pipeline._vae_attention_context = lambda device: _RecordingContext(device) + original_decode = pipeline.vae.decode + + def recording_decode(latents, return_dict=False): + events.append(("decode", latents.device.type)) + return original_decode(latents, return_dict=return_dict) + + pipeline.vae.decode = recording_decode + req = _make_request_batch( + "a cat", + height=64, + width=64, + output_type="pt", + generator=torch.Generator(device="cpu").manual_seed(7), + ) + + pipeline.forward(req) + + assert events == [("enter", "cpu"), ("decode", "cpu"), ("exit", "cpu")] + + +def test_turbo_dmd_step_and_renoise_match_upstream_equations(): + from diffusers.utils.torch_utils import randn_tensor + + pipeline = _make_turbo_forward_pipeline() + latents = torch.full((2, 4, 3, 3), 2.0) + embeds = torch.zeros(2, _SEQ_LEN, _EMBED_DIM) + mask = torch.ones(2, _SEQ_LEN, dtype=torch.long) + + x0 = pipeline._predict_dmd_student_step(latents, torch.tensor(0.25), embeds, None, mask) + assert torch.allclose(x0, latents + 0.75 * 0.125) + + expected_generator = torch.Generator(device="cpu").manual_seed(13) + expected_noise = randn_tensor(x0.shape, generator=expected_generator, device=x0.device, dtype=x0.dtype) + expected = 0.4 * expected_noise + 0.6 * x0 + + actual_generator = torch.Generator(device="cpu").manual_seed(13) + actual = pipeline._renoise_dmd_latents(x0, torch.tensor(0.6), actual_generator) + assert torch.allclose(actual, expected, rtol=1e-6, atol=1e-7) + + +def test_turbo_dmd_seeded_renoise_is_deterministic(): + def run_once(): + pipeline = _make_turbo_forward_pipeline() + req = _make_request_batch( + "a cat", + height=64, + width=64, + output_type="latent", + generator=torch.Generator(device="cpu").manual_seed(11), + ) + return pipeline.forward(req)[0].output + + assert torch.equal(run_once(), run_once()) + + +def test_turbo_dmd_custom_timesteps_and_validation(): + pipeline = _make_turbo_forward_pipeline() + sigmas = pipeline._build_dmd_student_sigmas(99, torch.device("cpu"), torch.float32, 0.001, [1000, 500, 0]) + assert torch.equal(sigmas, torch.tensor([1.0, 0.5, 0.0])) + + with pytest.raises(ValueError, match="non-empty 1D"): + pipeline._build_dmd_student_sigmas(4, torch.device("cpu"), torch.float32, 0.001, []) + with pytest.raises(ValueError, match="non-empty 1D"): + pipeline._build_dmd_student_sigmas(4, torch.device("cpu"), torch.float32, 0.001, [[0.1]]) + with pytest.raises(ValueError, match="finite values"): + pipeline._build_dmd_student_sigmas(4, torch.device("cpu"), torch.float32, 0.001, [0.0, 1001.0]) + + +def test_turbo_dmd_conditioning_sigma_override_reaches_forward_loop(): + pipeline = _make_turbo_forward_pipeline() + req = _make_request_batch( + "a cat", + height=64, + width=64, + output_type="latent", + extra_args={"dmd_conditioning_sigma": 0.125}, + ) + + pipeline.forward(req) + + assert pipeline.transformer.timesteps[0] == pytest.approx(0.125) + + +def test_turbo_dmd_rejects_timesteps_and_sigmas_together(): + pipeline = _make_turbo_forward_pipeline() + req = _make_request_batch( + "a cat", + height=64, + width=64, + timesteps=torch.tensor([0.0, 0.5]), + sigmas=[0.0, 0.5], + ) + with pytest.raises(ValueError, match="only one of timesteps or sigmas"): + pipeline.forward(req) + + +def test_turbo_dmd_rejects_zero_steps(): + pipeline = _make_turbo_forward_pipeline() + req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=0) + with pytest.raises(ValueError, match="num_inference_steps must be >= 1"): + pipeline.forward(req) + + +def test_turbo_dmd_normalizes_internal_dummy_guidance_only(): + pipeline = _make_turbo_forward_pipeline() + req = _make_request_batch( + "dummy run", + height=64, + width=64, + num_inference_steps=1, + guidance_scale=0.0, + guidance_scale_2=0.0, + output_type="latent", + ) + req.is_dummy_run = lambda: True + + pipeline.forward(req) + + assert len(pipeline.transformer.timesteps) == 1 + assert len(pipeline.processor.calls) == 1 + + +@pytest.mark.parametrize( + "sampling_overrides", + [ + {"guidance_scale": 0.0}, + {"guidance_scale": 2.0}, + {"guidance_scale_2": 2.0}, + {"extra_args": {"empty_instruction_guidance_scale": 1.0}}, + ], +) +def test_turbo_dmd_rejects_non_unit_guidance(sampling_overrides): + pipeline = _make_turbo_forward_pipeline() + req = _make_request_batch("a cat", height=64, width=64, **sampling_overrides) + with pytest.raises(ValueError, match="requires guidance_scale=1.0"): + pipeline.forward(req) + + # --------------------------------------------------------------------------- # Editing / TI2I: pre-process function # --------------------------------------------------------------------------- @@ -1316,3 +1544,53 @@ def test_supports_request_batch_enabled(): from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline assert BooguImagePipeline.supports_request_batch is True + + +def test_turbo_request_batching_disabled_and_fails_closed(): + from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline + + assert BooguImageTurboPipeline.supports_request_batch is False + + pipeline = _make_turbo_forward_pipeline() + sampling = dict(height=64, width=64, num_inference_steps=4, guidance_scale=1.0, output_type="latent") + req = _wrap_request_batch( + [ + ("a cat", _sampling(**sampling)), + ("a dog", _sampling(**sampling)), + ] + ) + + with pytest.raises(RuntimeError, match="does not support request batching"): + pipeline.forward(req) + + +def _make_turbo_edit_forward_pipeline(): + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImageTurboPipeline + + base_pipeline = _make_edit_forward_pipeline() + pipeline = object.__new__(BooguImageTurboPipeline) + nn.Module.__init__(pipeline) + pipeline.__dict__.update(base_pipeline.__dict__) + pipeline.transformer = _RecordingDMDTransformer() + return pipeline + + +def test_turbo_dmd_ti2i_keeps_reference_latents_and_uses_zero_conditioning_sigma(): + pipeline = _make_turbo_edit_forward_pipeline() + req = _make_edit_request( + height=64, + width=64, + num_inference_steps=4, + guidance_scale=1.0, + output_type="latent", + generator=torch.Generator(device="cpu").manual_seed(3), + ) + req.sampling_params.guidance_scale_provided = True + + pipeline.forward(req) + + assert len(pipeline.transformer.timesteps) == 4 + assert pipeline.transformer.timesteps[0] == pytest.approx(0.0) + assert all(ref is not None for ref in pipeline.transformer.refs) + assert pipeline.scheduler.set_calls == 0 + assert pipeline.scheduler.step_calls == 0 diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image_cuda.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image_cuda.py new file mode 100644 index 000000000..fe44392fe --- /dev/null +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image_cuda.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Focused CUDA correctness coverage for Boogu-Image Turbo DMD helpers.""" + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline + +pytestmark = [pytest.mark.advanced_model, pytest.mark.cuda, pytest.mark.diffusion] + + +class _CudaDMDPipeline(BooguImageTurboPipeline): + def predict( + self, + t, + latents, + instruction_embeds, + freqs_cis, + instruction_attention_mask, + ref_image_hidden_states=None, + ): + return torch.full_like(latents, 0.125) + + +def test_boogu_turbo_dmd_helpers_cuda_bf16(): + if not torch.cuda.is_available(): + pytest.skip("Boogu-Image Turbo DMD CUDA correctness requires a CUDA device") + + pipeline = object.__new__(_CudaDMDPipeline) + nn.Module.__init__(pipeline) + device = torch.device("cuda") + latents = torch.full((2, 4, 8, 8), 2.0, device=device, dtype=torch.bfloat16) + embeds = torch.zeros(2, 4, 8, device=device, dtype=torch.bfloat16) + mask = torch.ones(2, 4, device=device, dtype=torch.long) + + sigmas = pipeline._build_dmd_student_sigmas(4, device, latents.dtype, 0.001) + assert sigmas.device.type == "cuda" + assert sigmas.dtype == torch.bfloat16 + + x0 = pipeline._predict_dmd_student_step(latents, sigmas[0], embeds, None, mask) + expected_x0 = latents + (1.0 - sigmas[0]) * torch.full_like(latents, 0.125) + torch.testing.assert_close(x0, expected_x0) + + first_generator = torch.Generator(device=device).manual_seed(13) + second_generator = torch.Generator(device=device).manual_seed(13) + first = pipeline._renoise_dmd_latents(x0, sigmas[1], first_generator) + second = pipeline._renoise_dmd_latents(x0, sigmas[1], second_generator) + assert first.device.type == "cuda" + assert torch.equal(first, second) diff --git a/vllm_omni/diffusion/models/boogu_image/__init__.py b/vllm_omni/diffusion/models/boogu_image/__init__.py index 7ec9c167c..1ccd3447c 100644 --- a/vllm_omni/diffusion/models/boogu_image/__init__.py +++ b/vllm_omni/diffusion/models/boogu_image/__init__.py @@ -1,10 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from .boogu_image_transformer import BooguImageTransformer2DModel from .image_processor import BooguImageProcessor -from .pipeline_boogu_image import BooguImagePipeline, get_boogu_image_pre_process_func +from .pipeline_boogu_image import BooguImagePipeline, BooguImageTurboPipeline, get_boogu_image_pre_process_func from .scheduling_flow_match_euler_discrete_time_shifting import FlowMatchEulerDiscreteScheduler __all__ = [ "BooguImagePipeline", + "BooguImageTurboPipeline", "BooguImageProcessor", "BooguImageTransformer2DModel", "FlowMatchEulerDiscreteScheduler", 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 cc03b75df..257e4e21a 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -16,6 +16,8 @@ Ported from the upstream ``boogu`` package CFG branches use vLLM-Omni's shared two-branch/N-branch parallel helpers. - Instruction rewriting, prompt tuning, and vision-token stripping are not ported. +- ``BooguImagePipeline`` preserves the regular scheduler/CFG path, while + ``BooguImageTurboPipeline`` selects the upstream few-step DMD student path. """ import json @@ -44,6 +46,7 @@ from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_p from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( BooguImageDoubleStreamRotaryPosEmbed, BooguImageTransformer2DModel, + RotaryFrequencyTables, ) 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 ( @@ -218,6 +221,9 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports """ supports_request_batch = True + # Turbo subclasses select the upstream few-step DMD student loop. Base and + # Edit requests keep the regular scheduler and CFG path. + _is_turbo: ClassVar[bool] = False support_image_input: ClassVar[bool] = True color_format: ClassVar[str] = "RGB" @@ -655,6 +661,108 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports combined = self.cfg_normalize_function(positive_with_reference, combined) return combined + def _build_dmd_student_sigmas( + self, + num_inference_steps: int, + device: torch.device, + dtype: torch.dtype, + conditioning_sigma: float, + timesteps: list[float] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Build the ascending sigma schedule used by Boogu's DMD student. + + The upstream Turbo pipeline uses ``linspace(conditioning_sigma, 1, + steps + 1)[:-1]``. Explicit schedules are accepted in either the + normalized [0, 1] form or the 0..1000 training-timestep form. + """ + if timesteps is not None: + # Validate request metadata on CPU once. Keeping scalar extraction + # out of the denoise loop avoids a CUDA synchronization per step. + dmd_sigmas = torch.as_tensor(timesteps, device="cpu", dtype=torch.float32) + if dmd_sigmas.ndim != 1 or dmd_sigmas.numel() == 0: + raise ValueError("DMD sigmas must be a non-empty 1D sequence.") + if dmd_sigmas.max().item() > 1.0: + dmd_sigmas = dmd_sigmas / 1000.0 + if ( + not torch.isfinite(dmd_sigmas).all().item() + or (dmd_sigmas < 0).any().item() + or (dmd_sigmas > 1).any().item() + ): + raise ValueError("DMD sigmas must be finite values in [0, 1] (or 0..1000 timesteps).") + return dmd_sigmas.to(device=device, dtype=dtype) + else: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1 for DMD student inference.") + if not 0.0 <= conditioning_sigma <= 1.0: + raise ValueError(f"DMD conditioning sigma must be in [0, 1], got {conditioning_sigma}.") + dmd_sigmas = torch.linspace( + conditioning_sigma, + 1.0, + num_inference_steps + 1, + device=device, + dtype=dtype, + )[:-1] + return dmd_sigmas + + def _predict_dmd_student_step( + self, + latents: torch.Tensor, + sigma: torch.Tensor | float, + instruction_embeds: torch.Tensor, + freqs_real: RotaryFrequencyTables, + instruction_attention_mask: torch.Tensor, + ref_latents: list[list[torch.Tensor] | None] | None = None, + ) -> torch.Tensor: + """Predict x0 and apply the upstream DMD ``x + (1-sigma)*velocity``.""" + sigma_tensor = torch.as_tensor(sigma, device=latents.device, dtype=latents.dtype).reshape(()) + model_pred = self.predict( + sigma_tensor, + latents, + instruction_embeds, + freqs_real, + instruction_attention_mask, + ref_image_hidden_states=ref_latents, + ) + sigma_expanded = sigma_tensor.reshape((1,) + (1,) * (latents.ndim - 1)) + return latents + (1.0 - sigma_expanded) * model_pred + + def _renoise_dmd_latents( + self, + latents: torch.Tensor, + sigma: torch.Tensor | float, + generator: torch.Generator | list[torch.Generator] | None = None, + ) -> torch.Tensor: + """Renoise an intermediate DMD x0 with the next sigma and seeded RNG.""" + sigma_tensor = torch.as_tensor(sigma, device=latents.device, dtype=latents.dtype).reshape(()) + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + sigma_expanded = sigma_tensor.reshape((1,) + (1,) * (latents.ndim - 1)) + return (1.0 - sigma_expanded) * noise + sigma_expanded * latents + + def _decode_output( + self, + latents: torch.Tensor, + output_type: str, + dtype: torch.dtype, + height: int, + width: int, + ori_height: int, + ori_width: int, + ) -> DiffusionOutput: + 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 + 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") + + return DiffusionOutput(output=image) + def _encode_vae_image(self, img: torch.Tensor, generator=None) -> torch.Tensor: """Encode an image tensor into the VAE latent space (upstream ``encode_vae``). @@ -758,6 +866,11 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports has_reference = any(img is not None for img in preprocessed_images) task_type = "ti2i" if has_reference else "t2i" + # Turbo-specific schedules and extra arguments are not all represented + # in the request-batch compatibility key, so Turbo stays at batch=1. + if self._is_turbo and req.num_reqs > 1: + raise RuntimeError("BooguImageTurboPipeline does not support request batching.") + # Fail-closed: a batched ti2i must never reach here (it is gated to batch=1). if has_reference and req.num_reqs > 1: raise RuntimeError( @@ -775,15 +888,25 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports 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 + num_inference_steps = ( + (sp.num_inference_steps if sp.num_inference_steps is not None else 4) + if self._is_turbo + else (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 + text_guidance_scale = ( + (sp.guidance_scale if sp.guidance_scale is not None else 1.0) + if self._is_turbo + else (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 + image_guidance_scale = ( + (sp.guidance_scale_2 if sp.guidance_scale_2 is not None else 1.0) + if self._is_turbo + else (sp.guidance_scale_2 if sp.guidance_scale_2_provided else 1.0) + ) num_images_per_prompt = sp.num_outputs_per_prompt if sp.num_outputs_per_prompt > 0 else 1 # Per-request noise generators, collated into one flat list of length # batch_size * num_images_per_prompt (request-major, output-minor). @@ -798,6 +921,26 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports output_type = sp.output_type or "pil" cfg_range = (0.0, 1.0) + extra_args = getattr(sp, "extra_args", None) or {} + empty_instruction_guidance_scale = float(extra_args.get("empty_instruction_guidance_scale", 0.0)) + is_dummy_run = callable(getattr(req, "is_dummy_run", None)) and req.is_dummy_run() + if self._is_turbo and is_dummy_run: + # The engine's generic startup request uses guidance_scale=0.0. + # Turbo's equivalent no-CFG value is 1.0, so normalize only this + # internal warmup while keeping real request validation strict. + text_guidance_scale = 1.0 + image_guidance_scale = 1.0 + empty_instruction_guidance_scale = 0.0 + if self._is_turbo and ( + text_guidance_scale != 1.0 or image_guidance_scale != 1.0 or empty_instruction_guidance_scale != 0.0 + ): + raise ValueError( + "Boogu-Image Turbo DMD inference requires guidance_scale=1.0, " + "guidance_scale_2=1.0, and empty_instruction_guidance_scale=0.0." + ) + if not has_reference: + image_guidance_scale = 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 @@ -852,6 +995,42 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports theta=10000, ) + # Turbo uses the standalone upstream DMD loop and deliberately bypasses + # the regular scheduler/CFG path. Base/Edit continue below unchanged. + if self._is_turbo: + conditioning_sigma = float(extra_args.get("dmd_conditioning_sigma", 0.0 if has_reference else 0.001)) + if sp.timesteps is not None and sp.sigmas is not None: + raise ValueError("Boogu-Image Turbo accepts only one of timesteps or sigmas.") + # ``timesteps`` mirrors the upstream Turbo API. ``sigmas`` is kept + # as the vLLM-native spelling for the same normalized DMD schedule. + dmd_timesteps = sp.timesteps if sp.timesteps is not None else sp.sigmas + dmd_sigmas = self._build_dmd_student_sigmas( + num_inference_steps, + device=device, + dtype=latents.dtype, + conditioning_sigma=conditioning_sigma, + timesteps=dmd_timesteps, + ) + with self.progress_bar(total=len(dmd_sigmas)) as progress_bar: + for index, sigma in enumerate(dmd_sigmas): + latents = self._predict_dmd_student_step( + latents, + sigma, + instruction_embeds, + freqs_real, + instruction_attention_mask, + ref_latents, + ).to(dtype=dtype) + if index + 1 < len(dmd_sigmas): + latents = self._renoise_dmd_latents(latents, dmd_sigmas[index + 1], generator).to(dtype=dtype) + progress_bar.update() + output = self._decode_output(latents, output_type, dtype, height, width, ori_height, ori_width) + return split_diffusion_output_by_request( + output, + req, + num_outputs_per_prompt=num_images_per_prompt, + ) + # 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) @@ -963,21 +1142,16 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports 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 - 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") - + output = self._decode_output(latents, output_type, dtype, height, width, ori_height, ori_width) return split_diffusion_output_by_request( - DiffusionOutput(output=image), + output, req, num_outputs_per_prompt=num_images_per_prompt, ) + + +class BooguImageTurboPipeline(BooguImagePipeline): + """Boogu-Image Turbo pipeline using the upstream few-step DMD semantics.""" + + supports_request_batch = False + _is_turbo: ClassVar[bool] = True