1238 lines
57 KiB
Python
1238 lines
57 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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).
|
||
- 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.
|
||
- ``BooguImagePipeline`` preserves the regular scheduler/CFG path, while
|
||
``BooguImageTurboPipeline`` selects the upstream few-step DMD student path.
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
from collections.abc import Iterable
|
||
from contextlib import nullcontext
|
||
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 torch.nn.attention import SDPBackend, sdpa_kernel
|
||
from transformers import AutoModel, Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen3VLModel, Qwen3VLProcessor
|
||
from vllm.config.quantization import QuantizationConfigArgs
|
||
from vllm.logger import init_logger
|
||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
||
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
|
||
|
||
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
|
||
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 (
|
||
FlowMatchEulerDiscreteScheduler,
|
||
)
|
||
from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
|
||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||
from vllm_omni.diffusion.models.utils import create_transformers_model_with_vllm_linears
|
||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
|
||
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
||
from vllm_omni.platforms import current_omni_platform
|
||
from vllm_omni.quantization.component_config import resolve_component_quant_config
|
||
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
||
|
||
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 _boogu_batch_compatibility_key(has_reference: bool) -> tuple:
|
||
"""Separate T2I and TI2I denoise paths.
|
||
|
||
Shared shape/guidance fields, including ``guidance_scale_2_provided``,
|
||
are checked by ``RequestBatchSamplingParamsKey``.
|
||
"""
|
||
if not has_reference:
|
||
return ("boogu_image", "t2i")
|
||
return ("boogu_image", "ti2i")
|
||
|
||
|
||
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.
|
||
request.batch_compatibility_key = _boogu_batch_compatibility_key(False)
|
||
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 (Base checkpoint).
|
||
request.batch_compatibility_key = _boogu_batch_compatibility_key(False)
|
||
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
|
||
request.batch_compatibility_key = _boogu_batch_compatibility_key(True)
|
||
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(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 /
|
||
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 = 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"
|
||
|
||
_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()
|
||
transformer_quant_config = resolve_component_quant_config(od_config.quantization_config, "transformer")
|
||
mllm_quant_config = resolve_component_quant_config(od_config.quantization_config, "mllm")
|
||
self.weights_sources = [
|
||
DiffusersPipelineLoader.ComponentSource(
|
||
model_or_path=od_config.model,
|
||
subfolder="transformer",
|
||
revision=od_config.revision,
|
||
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,
|
||
revision=od_config.revision,
|
||
)
|
||
|
||
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
||
model,
|
||
subfolder="scheduler",
|
||
local_files_only=local_files_only,
|
||
revision=od_config.revision,
|
||
)
|
||
self.mllm = self._load_mllm(model, local_files_only, mllm_quant_config, boogu_subfolders)
|
||
|
||
self.processor = Qwen3VLProcessor.from_pretrained(
|
||
model,
|
||
subfolder="processor",
|
||
local_files_only=local_files_only,
|
||
revision=od_config.revision,
|
||
)
|
||
|
||
self.vae = from_pretrained_with_prefetch(
|
||
AutoencoderKL.from_pretrained,
|
||
model,
|
||
subfolder="vae",
|
||
prefetch_list=boogu_subfolders,
|
||
local_files_only=local_files_only,
|
||
revision=od_config.revision,
|
||
).to(self._execution_device)
|
||
|
||
self.transformer = BooguImageTransformer2DModel(
|
||
od_config=od_config,
|
||
quant_config=transformer_quant_config,
|
||
prefix="transformer",
|
||
)
|
||
|
||
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 _load_mllm(
|
||
self,
|
||
model_path: str,
|
||
local_files_only: bool,
|
||
quant_config: QuantizationConfig | None,
|
||
prefetch_list: list[str],
|
||
) -> Qwen3VLModel:
|
||
"""Load MLLM through HF, or prepare it for deferred online FP8 loading."""
|
||
config = None
|
||
if quant_config is not None:
|
||
if not isinstance(quant_config, (Fp8Config, DiffusionFp8Config)):
|
||
raise ValueError(
|
||
"Boogu MLLM only supports FP8 quantization. Set mllm to null to disable online quantization."
|
||
)
|
||
config = Qwen3VLConfig.from_pretrained(
|
||
model_path, subfolder="mllm", local_files_only=local_files_only, revision=self.od_config.revision
|
||
)
|
||
|
||
# With no MLLM quantization override, or with checkpoint quantization
|
||
# metadata, preserve HF model initialization and weight loading.
|
||
if config is None or getattr(config, "quantization_config", None):
|
||
mllm = from_pretrained_with_prefetch(
|
||
Qwen3VLForConditionalGeneration.from_pretrained,
|
||
model_path,
|
||
subfolder="mllm",
|
||
prefetch_list=prefetch_list,
|
||
local_files_only=local_files_only,
|
||
torch_dtype=self.od_config.dtype,
|
||
revision=self.od_config.revision,
|
||
)
|
||
# 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
|
||
return mllm.to(self._execution_device)
|
||
|
||
online_quant_config = OnlineQuantizationConfig(
|
||
QuantizationConfigArgs(linear="fp8_per_block", ignore=quant_config.ignored_layers)
|
||
)
|
||
|
||
mllm = create_transformers_model_with_vllm_linears(
|
||
AutoModel,
|
||
config,
|
||
online_quant_config,
|
||
dtype=self.od_config.dtype,
|
||
device=self._execution_device,
|
||
prefix="mllm",
|
||
skip_modules=("mllm.visual",),
|
||
)
|
||
|
||
self.weights_sources.append(
|
||
DiffusersPipelineLoader.ComponentSource(
|
||
model_or_path=model_path,
|
||
subfolder="mllm",
|
||
revision=self.od_config.revision,
|
||
prefix="mllm.",
|
||
)
|
||
)
|
||
|
||
return mllm.requires_grad_(False).eval()
|
||
|
||
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:
|
||
if parallel_config.cfg_parallel_size > 1:
|
||
# SP × CFG-Parallel is unvalidated: the combined path runs
|
||
# without crashing but drifts far from the SP=1 baseline
|
||
# (~19 dB PSNR at 4 steps), so guidance semantics differ.
|
||
# Remove this guard only alongside a parity test.
|
||
raise NotImplementedError(
|
||
"BooguImagePipeline: sequence parallelism combined with CFG parallelism is not validated."
|
||
)
|
||
if parallel_config.ring_degree > 1 or getattr(parallel_config, "allgather_degree", 1) > 1:
|
||
raise NotImplementedError(
|
||
"BooguImagePipeline currently supports sequence parallelism through Ulysses only."
|
||
)
|
||
if not current_omni_platform.is_cuda():
|
||
raise NotImplementedError("BooguImagePipeline sequence parallelism currently requires CUDA.")
|
||
if parallel_config.ulysses_mode == "strict":
|
||
# Subscript rather than get(): a missing key means the config
|
||
# schema has drifted and we would silently validate against a
|
||
# stale default, not the actual head count.
|
||
tf_params = self.od_config.tf_model_config.params
|
||
num_heads = tf_params["num_attention_heads"]
|
||
num_kv_heads = tf_params["num_kv_heads"]
|
||
if (
|
||
num_heads % parallel_config.ulysses_degree != 0
|
||
or num_kv_heads % parallel_config.ulysses_degree != 0
|
||
):
|
||
raise ValueError(
|
||
"BooguImagePipeline GQA heads are not divisible by the "
|
||
"configured Ulysses degree; set "
|
||
"parallel_config.ulysses_mode='advanced_uaa'."
|
||
)
|
||
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)
|
||
mapper = WeightsMapper(orig_to_new_prefix={"mllm.lm_head.": None, "mllm.model.": "mllm."})
|
||
return loader.load_weights(weights, mapper=mapper)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 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)
|
||
|
||
# repeat_interleave (not repeat/tile) so mask rows stay request-major
|
||
# [p0, p0, p1, p1] to match the reshaped embeds above; a plain repeat
|
||
# tiles to [p0, p1, p0, p1] and mismatches when batch_size > 1 and
|
||
# num_images_per_prompt > 1.
|
||
reshaped_mask = mask.repeat_interleave(num_images_per_prompt, dim=0)
|
||
|
||
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_real, 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_real,
|
||
instruction_attention_mask,
|
||
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 _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``).
|
||
|
||
Upstream leaves ``latent_dist.sample()`` unseeded; the native path
|
||
threads the request generator through so a fixed seed gives a
|
||
reproducible reference latent.
|
||
"""
|
||
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],
|
||
num_images_per_prompt: int,
|
||
device: torch.device,
|
||
generators: list[torch.Generator | None] | None = 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.
|
||
|
||
``generators`` holds one generator (or ``None``) per request, so each
|
||
reference latent is sampled with its own request's seed and batched
|
||
edits stay reproducible against the single-request path.
|
||
"""
|
||
ref_latents: list[list[torch.Tensor] | None] = []
|
||
for idx, image in enumerate(preprocessed_images):
|
||
if image is None:
|
||
sample_latents: list[torch.Tensor] | None = None
|
||
else:
|
||
generator = generators[idx] if generators is not None else None
|
||
latent = self._encode_vae_image(image.to(device=device), generator=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) -> list[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"
|
||
|
||
# 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.")
|
||
|
||
sampling_params_list = req.sampling_params_list
|
||
# Shared shape/step/guidance fields are guaranteed identical across the
|
||
# batch by RequestBatchSamplingParamsKey; request-local values (seeds,
|
||
# reference generators) are read per request below.
|
||
sp = sampling_params_list[0]
|
||
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 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 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 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).
|
||
generator = req.collate_request_generators(num_images_per_prompt, None)
|
||
# One generator per request for reference-latent sampling; a per-output
|
||
# generator list is not usable by ``latent_dist.sample`` and falls back
|
||
# to unseeded, matching the single-request path.
|
||
ref_generators = [
|
||
s.generator if isinstance(s.generator, torch.Generator) else None for s in sampling_params_list
|
||
]
|
||
max_sequence_length = sp.max_sequence_length or 1280
|
||
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
|
||
|
||
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, ref_generators)
|
||
|
||
latents = self.prepare_latents(
|
||
batch_size * num_images_per_prompt,
|
||
latent_channels,
|
||
height,
|
||
width,
|
||
instruction_embeds.dtype,
|
||
device,
|
||
generator,
|
||
)
|
||
|
||
freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
|
||
self.transformer.axes_dim_rope,
|
||
self.transformer.axes_lens,
|
||
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)
|
||
timesteps = self.scheduler.timesteps
|
||
num_timesteps = len(timesteps)
|
||
|
||
# 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
|
||
|
||
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).
|
||
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,
|
||
)
|
||
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,
|
||
)
|
||
elif task_type == "ti2i" and text_gs > 1.0:
|
||
# Text-only ti2i guidance: reference kept in the uncond pred.
|
||
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,
|
||
)
|
||
elif task_type == "ti2i" and image_gs > 1.0:
|
||
# Image-only ti2i guidance: drop the reference in the uncond pred.
|
||
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,
|
||
)
|
||
elif text_gs > 1.0:
|
||
# Text-to-image classifier-free guidance.
|
||
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,
|
||
)
|
||
|
||
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()
|
||
|
||
# 6. Decode.
|
||
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,
|
||
)
|
||
|
||
|
||
class BooguImageTurboPipeline(BooguImagePipeline):
|
||
"""Boogu-Image Turbo pipeline using the upstream few-step DMD semantics."""
|
||
|
||
supports_request_batch = False
|
||
_is_turbo: ClassVar[bool] = True
|