1218 lines
56 KiB
Python
1218 lines
56 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.quantization.component_config import resolve_component_quant_config
|
|
|
|
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, request_id: str) -> tuple:
|
|
"""Request-batch isolation key. ``forward`` reads shared guidance/shape fields
|
|
from the batch's first request, so t2i and ti2i must not share a key.
|
|
|
|
ti2i is held at batch=1 (request-unique key) because
|
|
``guidance_scale_2_provided`` is absent from ``RequestBatchSamplingParamsKey``
|
|
while ``forward`` reads it from the first request, so mixed-``_provided`` edits
|
|
with equal numeric ``guidance_scale_2`` would co-batch into the wrong mode.
|
|
"""
|
|
if not has_reference:
|
|
return ("boogu_image", "t2i")
|
|
return ("boogu_image", "ti2i", request_id)
|
|
|
|
|
|
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, request.request_id)
|
|
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, request.request_id)
|
|
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, request.request_id)
|
|
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):
|
|
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:
|
|
raise NotImplementedError("Sequence parallelism is not supported by BooguImagePipeline.")
|
|
if parallel_config.use_hsdp:
|
|
raise NotImplementedError("HSDP is not supported by BooguImagePipeline.")
|
|
if self.od_config.cache_backend not in (None, "", "none"):
|
|
raise NotImplementedError(
|
|
f"Cache backend '{self.od_config.cache_backend}' is not supported by BooguImagePipeline."
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
loader = AutoWeightsLoader(self)
|
|
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.")
|
|
|
|
# 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(
|
|
f"BooguImagePipeline received a batched TI2I (edit) request "
|
|
f"(num_reqs={req.num_reqs}); TI2I batching is gated to batch=1 "
|
|
"pending guidance-mode / compatibility-key validation."
|
|
)
|
|
|
|
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
|