Ding ZuhaoandCursor
154a1f25f5
Signed-off-by: Ding Zuhao <e1583181@u.nus.edu> Co-authored-by: Cursor <cursoragent@cursor.com>
1593 lines
72 KiB
Python
1593 lines
72 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""
|
|
BagelPipeline implementation for vLLM-Omni.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from collections.abc import Iterable, Sequence
|
|
from copy import copy, deepcopy
|
|
from dataclasses import dataclass
|
|
from math import isqrt
|
|
from typing import TYPE_CHECKING, Any, ClassVar
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from torch import nn
|
|
from transformers import AutoTokenizer, SiglipImageProcessor, SiglipVisionConfig, SiglipVisionModel
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader
|
|
from vllm.transformers_utils.configs.bagel import BagelConfig
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
get_cfg_group,
|
|
get_classifier_free_guidance_world_size,
|
|
)
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
|
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.errors import OmniClientError
|
|
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
|
|
|
|
from .autoencoder import AutoEncoder, AutoEncoderParams, DistributedAutoEncoder
|
|
from .bagel_transformer import Bagel, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
|
from vllm_omni.diffusion.worker.utils import StepRequestState
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
|
|
|
|
|
|
@dataclass
|
|
class BagelGenParams:
|
|
num_timesteps: int = 50
|
|
timestep_shift: float = 3.0
|
|
cfg_text_scale: float = 4.0
|
|
cfg_img_scale: float = 1.5
|
|
cfg_interval: tuple = (0.4, 1.0)
|
|
cfg_renorm_min: float = 0.0
|
|
cfg_renorm_type: str = "global"
|
|
|
|
|
|
def add_special_tokens(tokenizer):
|
|
all_special_tokens = []
|
|
for k, v in tokenizer.special_tokens_map.items():
|
|
if isinstance(v, str):
|
|
all_special_tokens.append(v)
|
|
elif isinstance(v, list):
|
|
all_special_tokens += v
|
|
|
|
new_tokens = []
|
|
|
|
if "<|im_start|>" not in all_special_tokens:
|
|
new_tokens.append("<|im_start|>")
|
|
|
|
if "<|im_end|>" not in all_special_tokens:
|
|
new_tokens.append("<|im_end|>")
|
|
|
|
if "<|vision_start|>" not in all_special_tokens:
|
|
new_tokens.append("<|vision_start|>")
|
|
|
|
if "<|vision_end|>" not in all_special_tokens:
|
|
new_tokens.append("<|vision_end|>")
|
|
|
|
num_new_tokens = tokenizer.add_tokens(new_tokens)
|
|
bos_token_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
|
eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
|
|
start_of_image = tokenizer.convert_tokens_to_ids("<|vision_start|>")
|
|
end_of_image = tokenizer.convert_tokens_to_ids("<|vision_end|>")
|
|
|
|
new_token_ids = dict(
|
|
bos_token_id=bos_token_id,
|
|
eos_token_id=eos_token_id,
|
|
start_of_image=start_of_image,
|
|
end_of_image=end_of_image,
|
|
)
|
|
|
|
return tokenizer, new_token_ids, num_new_tokens
|
|
|
|
|
|
def get_bagel_post_process_func(od_config: OmniDiffusionConfig):
|
|
# BagelPipeline returns PIL.Image.Image directly.
|
|
def post_process_func(x):
|
|
return x
|
|
|
|
return post_process_func
|
|
|
|
|
|
def _resolve_bagel_image_geometry(od_config: OmniDiffusionConfig) -> tuple[int, int]:
|
|
"""Return the effective ``(latent_downsample, max_latent_size)``.
|
|
|
|
Some BAGEL checkpoints advertise a stale ``max_latent_size`` in
|
|
``config.json``. Model loading already corrects it from the positional
|
|
embedding weight, so admission preprocessing must do the same or it can
|
|
compute a different img2img shape from the Worker.
|
|
"""
|
|
|
|
model = getattr(od_config, "model", None)
|
|
if not model:
|
|
raise ValueError("BAGEL img2img preprocessing requires od_config.model.")
|
|
if os.path.exists(model):
|
|
model_path = model
|
|
else:
|
|
model_path = download_weights_from_hf_specific(
|
|
model,
|
|
None,
|
|
["*"],
|
|
revision=getattr(od_config, "revision", None),
|
|
)
|
|
|
|
config_path = os.path.join(model_path, "config.json")
|
|
with open(config_path, encoding="utf-8") as f:
|
|
bagel_cfg = json.load(f)
|
|
|
|
vae_config = bagel_cfg.get("vae_config") or {}
|
|
latent_downsample = int(vae_config.get("downsample", 8)) * int(bagel_cfg.get("latent_patch_size", 2))
|
|
max_latent_size = int(bagel_cfg.get("max_latent_size", 32))
|
|
|
|
index_path = os.path.join(model_path, "model.safetensors.index.json")
|
|
if os.path.exists(index_path):
|
|
with open(index_path, encoding="utf-8") as f:
|
|
weight_map = json.load(f).get("weight_map") or {}
|
|
position_key = next(
|
|
(key for key in ("latent_pos_embed.pos_embed", "bagel.latent_pos_embed.pos_embed") if key in weight_map),
|
|
None,
|
|
)
|
|
if position_key is not None:
|
|
from safetensors import safe_open
|
|
|
|
shard_path = os.path.join(model_path, weight_map[position_key])
|
|
with safe_open(shard_path, framework="pt") as f:
|
|
num_positions = int(f.get_slice(position_key).get_shape()[0])
|
|
inferred_size = isqrt(num_positions)
|
|
if inferred_size * inferred_size != num_positions:
|
|
raise ValueError(f"BAGEL latent position embedding length must be a square, got {num_positions}.")
|
|
max_latent_size = inferred_size
|
|
|
|
return latent_downsample, max_latent_size
|
|
|
|
|
|
def _bagel_effective_image_size(
|
|
image_size: tuple[int, int],
|
|
*,
|
|
latent_downsample: int,
|
|
max_latent_size: int,
|
|
) -> tuple[int, int]:
|
|
"""Match BAGEL's img2img resize and return ``(height, width)``."""
|
|
|
|
width, height = image_size
|
|
if width <= 0 or height <= 0:
|
|
raise ValueError(f"BAGEL img2img input must have positive dimensions, got {width}x{height}.")
|
|
|
|
max_image_size = int(max_latent_size * latent_downsample)
|
|
scale = min(max_image_size / max(width, height), 1.0)
|
|
min_image_size = min(256, max_image_size)
|
|
scale = max(scale, min_image_size / min(width, height))
|
|
resized_width = max(
|
|
latent_downsample,
|
|
int(round(width * scale / latent_downsample) * latent_downsample),
|
|
)
|
|
resized_height = max(
|
|
latent_downsample,
|
|
int(round(height * scale / latent_downsample) * latent_downsample),
|
|
)
|
|
return min(resized_height, max_image_size), min(resized_width, max_image_size)
|
|
|
|
|
|
def _bagel_canvas_requested(sampling: OmniDiffusionSamplingParams) -> bool:
|
|
return (sampling.height is not None and not sampling.height_not_provided) or (
|
|
sampling.width is not None and not sampling.width_not_provided
|
|
)
|
|
|
|
|
|
def _bagel_align_requested_canvas(
|
|
sampling: OmniDiffusionSamplingParams,
|
|
*,
|
|
latent_downsample: int,
|
|
max_latent_size: int,
|
|
) -> None:
|
|
"""Floor an explicitly requested canvas to the latent stride; reject sides over the checkpoint limit."""
|
|
|
|
max_image_size = int(max_latent_size * latent_downsample)
|
|
resized: list[str] = []
|
|
for name in ("width", "height"):
|
|
value = getattr(sampling, name)
|
|
if value is None or getattr(sampling, f"{name}_not_provided"):
|
|
continue
|
|
value = int(value)
|
|
if value > max_image_size:
|
|
raise OmniClientError(
|
|
f"Requested {name}={value} exceeds the BAGEL checkpoint limit of {max_image_size} "
|
|
f"(max_latent_size={max_latent_size}, latent_downsample={latent_downsample})."
|
|
)
|
|
aligned = max(latent_downsample, (value // latent_downsample) * latent_downsample)
|
|
if aligned != value:
|
|
setattr(sampling, name, aligned)
|
|
resized.append(f"{name} {value}->{aligned}")
|
|
if resized:
|
|
logger.warning(
|
|
"BAGEL: requested canvas is not a multiple of the latent stride %d and cannot be "
|
|
"generated as-is; resizing %s.",
|
|
latent_downsample,
|
|
", ".join(resized),
|
|
)
|
|
|
|
|
|
def get_bagel_pre_process_func(od_config: OmniDiffusionConfig):
|
|
"""Resolve BAGEL execution mode and step-batch compatibility."""
|
|
|
|
step_execution = bool(getattr(od_config, "step_execution", False))
|
|
image_geometry: tuple[int, int] | None = None
|
|
|
|
def pre_process_func(request: OmniDiffusionRequest):
|
|
nonlocal image_geometry
|
|
|
|
sampling = request.sampling_params
|
|
modalities = (request.prompt.get("modalities") or []) if isinstance(request.prompt, dict) else []
|
|
kv_metadata = getattr(sampling, "kv_metadata", None) or {}
|
|
image_shape = kv_metadata.get("image_shape")
|
|
if _bagel_canvas_requested(sampling) and "text" not in modalities:
|
|
if image_geometry is None:
|
|
image_geometry = _resolve_bagel_image_geometry(od_config)
|
|
latent_downsample, max_latent_size = image_geometry
|
|
_bagel_align_requested_canvas(
|
|
sampling,
|
|
latent_downsample=latent_downsample,
|
|
max_latent_size=max_latent_size,
|
|
)
|
|
elif image_shape is not None:
|
|
sampling.height, sampling.width = (int(value) for value in image_shape)
|
|
sampling.height_not_provided = False
|
|
sampling.width_not_provided = False
|
|
elif isinstance(request.prompt, dict):
|
|
multi_modal_data = request.prompt.get("multi_modal_data") or {}
|
|
image_input = multi_modal_data.get("img2img")
|
|
if image_input is None and "text" not in modalities:
|
|
image_input = multi_modal_data.get("image")
|
|
if isinstance(image_input, list):
|
|
image_input = image_input[0] if image_input else None
|
|
if image_input is not None:
|
|
if isinstance(image_input, str):
|
|
with Image.open(image_input) as image:
|
|
input_size = image.size
|
|
else:
|
|
input_size = image_input.size
|
|
if image_geometry is None:
|
|
image_geometry = _resolve_bagel_image_geometry(od_config)
|
|
latent_downsample, max_latent_size = image_geometry
|
|
sampling.height, sampling.width = _bagel_effective_image_size(
|
|
input_size,
|
|
latent_downsample=latent_downsample,
|
|
max_latent_size=max_latent_size,
|
|
)
|
|
sampling.height_not_provided = True
|
|
sampling.width_not_provided = True
|
|
|
|
extra_args = request.sampling_params.extra_args or {}
|
|
cfg_interval = tuple(extra_args.get("cfg_interval", (0.4, 1.0)))
|
|
request.batch_compatibility_key = (
|
|
"bagel_cfg",
|
|
float(extra_args.get("cfg_text_scale", 4.0)),
|
|
float(extra_args.get("cfg_img_scale", 1.5)),
|
|
cfg_interval,
|
|
extra_args.get("cfg_renorm_type", "global"),
|
|
float(extra_args.get("cfg_renorm_min", 0.0)),
|
|
)
|
|
|
|
# Preserve legacy plain-string behavior: only an explicit modality
|
|
# request selects Bagel.generate_text() instead of stepwise denoising.
|
|
if step_execution and isinstance(request.prompt, dict):
|
|
modalities = request.prompt.get("modalities") or []
|
|
if "text" in modalities:
|
|
request.use_step_execution = False
|
|
return request
|
|
|
|
return pre_process_func
|
|
|
|
|
|
@dataclass
|
|
class _VaeCfg:
|
|
z_channels: int = 16
|
|
downsample: int = 8
|
|
|
|
|
|
@dataclass
|
|
class _VitCfg:
|
|
patch_size: int = 14
|
|
hidden_size: int = 1152
|
|
|
|
|
|
def default_ae_params() -> AutoEncoderParams:
|
|
return AutoEncoderParams(
|
|
resolution=256,
|
|
in_channels=3,
|
|
downsample=8,
|
|
ch=128,
|
|
out_ch=3,
|
|
ch_mult=[1, 2, 4, 4],
|
|
num_res_blocks=2,
|
|
z_channels=16,
|
|
scale_factor=0.3611,
|
|
shift_factor=0.1159,
|
|
)
|
|
|
|
|
|
class SiglipNaViTWrapper(nn.Module):
|
|
def __init__(self, vision_model):
|
|
super().__init__()
|
|
# If input is SiglipVisionModel, unwrap it to get SiglipVisionTransformer
|
|
if hasattr(vision_model, "vision_model"):
|
|
self.vision_model = vision_model.vision_model
|
|
else:
|
|
self.vision_model = vision_model
|
|
|
|
def forward(self, packed_pixel_values, packed_flattened_position_ids, cu_seqlens, max_seqlen):
|
|
patch_embed = self.vision_model.embeddings.patch_embedding
|
|
w = patch_embed.weight.view(patch_embed.weight.shape[0], -1)
|
|
x = F.linear(packed_pixel_values, w, patch_embed.bias)
|
|
pos = self.vision_model.embeddings.position_embedding(packed_flattened_position_ids)
|
|
x = x + pos
|
|
hidden_states = x.unsqueeze(0)
|
|
seq_len = x.shape[0]
|
|
mask = torch.full((1, 1, seq_len, seq_len), torch.finfo(x.dtype).min, device=x.device, dtype=x.dtype)
|
|
cu_seqlens_list = cu_seqlens.tolist()
|
|
for i in range(len(cu_seqlens_list) - 1):
|
|
start = cu_seqlens_list[i]
|
|
end = cu_seqlens_list[i + 1]
|
|
mask[..., start:end, start:end] = 0.0
|
|
|
|
outputs = self.vision_model.encoder(inputs_embeds=hidden_states, attention_mask=mask)
|
|
return outputs.last_hidden_state.squeeze(0)
|
|
|
|
|
|
class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin):
|
|
"""Bagel generation pipeline (MoT) packaged for vllm-omni diffusion engine.
|
|
|
|
This pipeline is self-contained and uses the ported Bagel core files.
|
|
"""
|
|
|
|
_dit_modules: ClassVar[list[str]] = ["language_model.model"]
|
|
_encoder_modules: ClassVar[list[str]] = []
|
|
_vae_modules: ClassVar[list[str]] = ["vae"]
|
|
_resident_modules: ClassVar[list[str]] = [
|
|
"bagel.time_embedder",
|
|
"bagel.vae2llm",
|
|
"bagel.llm2vae",
|
|
"bagel.latent_pos_embed",
|
|
"bagel.vit_model",
|
|
"bagel.connector",
|
|
"bagel.vit_pos_embed",
|
|
]
|
|
# LoRA scanning target. ``_dit_modules`` holds the dotted path
|
|
# "language_model.model", which DiffusionLoRAManager cannot resolve as a
|
|
# pipeline attribute, so the top-level ``bagel`` module (LLM plus ViT,
|
|
# connector, and VAE projections) is declared explicitly here.
|
|
_lora_components: ClassVar[list[str]] = ["bagel"]
|
|
supports_step_execution: ClassVar[bool] = True
|
|
|
|
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
|
super().__init__()
|
|
self.od_config = od_config
|
|
self.device = get_local_device()
|
|
|
|
self.scheduler: object | None = None
|
|
self.scheduler_kwargs: dict = {}
|
|
|
|
model = od_config.model
|
|
local_files_only = os.path.exists(model)
|
|
if local_files_only:
|
|
model_path = model
|
|
else:
|
|
# Download everything required (ema.safetensors, ae.safetensors, tokenizer files, configs).
|
|
model_path = download_weights_from_hf_specific(model, od_config.revision, ["*"])
|
|
|
|
# Load Bagel top-level config for VAE settings.
|
|
cfg_path = os.path.join(model_path, "config.json")
|
|
with open(cfg_path, encoding="utf-8") as f:
|
|
bagel_cfg = json.load(f)
|
|
|
|
vae_cfg_dict = bagel_cfg.get("vae_config") or {}
|
|
vae_cfg = _VaeCfg(
|
|
z_channels=int(vae_cfg_dict.get("z_channels", 16)),
|
|
downsample=int(vae_cfg_dict.get("downsample", 8)),
|
|
)
|
|
|
|
# LLM config: Bagel MoT requires explicitly setting layer_module
|
|
llm_cfg_path = os.path.join(model_path, "llm_config.json")
|
|
llm_config = Qwen2MoTConfig.from_json_file(llm_cfg_path)
|
|
llm_config.qk_norm = True
|
|
llm_config.tie_word_embeddings = False
|
|
# Allow overriding from vllm-omni config if user wants MoE/vanilla.
|
|
llm_config.layer_module = od_config.override_transformer_cls_name or "Qwen2MoTDecoderLayer"
|
|
|
|
# Tokenizer and special tokens.
|
|
# Bagel uses a Qwen2 tokenizer variant; prefer trust_remote_code to get the
|
|
# correct tokenizer implementation from the checkpoint repo when available.
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
model_path,
|
|
local_files_only=True,
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
# Try finding vision_config or interpolate from top-level config
|
|
vit_cfg_dict = bagel_cfg.get("vit_config") or {}
|
|
vit_cfg = _VitCfg(
|
|
patch_size=int(vit_cfg_dict.get("patch_size", 14)),
|
|
hidden_size=int(vit_cfg_dict.get("hidden_size", 1152)),
|
|
)
|
|
vit_config_path = os.path.join(model_path, "vit_config.json")
|
|
vit_conf = SiglipVisionConfig.from_json_file(vit_config_path)
|
|
if vit_conf.num_hidden_layers == 27:
|
|
vit_conf.num_hidden_layers = 26
|
|
vit_conf.vision_use_head = False
|
|
self.vit_model = SiglipVisionModel(vit_conf)
|
|
self.image_processor = SiglipImageProcessor.from_pretrained(model_path, local_files_only=True)
|
|
|
|
if self.vit_model:
|
|
self.vit_model = SiglipNaViTWrapper(self.vit_model)
|
|
vit_cfg.hidden_size = self.vit_model.vision_model.config.hidden_size
|
|
vit_cfg.patch_size = self.vit_model.vision_model.config.patch_size
|
|
|
|
self.tokenizer, self.new_token_ids, _ = add_special_tokens(self.tokenizer)
|
|
|
|
tok_len = len(self.tokenizer)
|
|
required_max_id = max(int(v) for v in self.new_token_ids.values())
|
|
llm_config.vocab_size = max(
|
|
int(getattr(llm_config, "vocab_size", tok_len)),
|
|
int(tok_len),
|
|
int(required_max_id + 1),
|
|
)
|
|
|
|
parallel_config = od_config.parallel_config if od_config else None
|
|
quant_config = od_config.quantization_config
|
|
# Bagel uses explicit prefixes ("bagel.language_model", "bagel") because
|
|
# its model structure nests components under a top-level "bagel" module,
|
|
# unlike other pipelines where the transformer is the root module.
|
|
# This ensures ComponentQuantizationConfig prefix matching works correctly.
|
|
self.language_model = Qwen2MoTForCausalLM(
|
|
llm_config, parallel_config=parallel_config, quant_config=quant_config, prefix="bagel.language_model"
|
|
)
|
|
self.transformer = self.language_model.model
|
|
ae_params: AutoEncoderParams = default_ae_params()
|
|
self.vae = DistributedAutoEncoder(ae_params)
|
|
|
|
self.bagel = Bagel(
|
|
language_model=self.language_model,
|
|
vit_model=self.vit_model,
|
|
parallel_config=parallel_config,
|
|
quant_config=quant_config,
|
|
prefix="bagel",
|
|
config=BagelConfig(
|
|
llm_config=llm_config,
|
|
vae_config=vae_cfg,
|
|
vit_config=vit_cfg,
|
|
vit_max_num_patch_per_side=int(bagel_cfg.get("vit_max_num_patch_per_side", 70)),
|
|
connector_act=str(bagel_cfg.get("connector_act", "gelu_pytorch_tanh")),
|
|
interpolate_pos=bool(bagel_cfg.get("interpolate_pos", False)),
|
|
latent_patch_size=int(bagel_cfg.get("latent_patch_size", 2)),
|
|
max_latent_size=int(bagel_cfg.get("max_latent_size", 32)),
|
|
timestep_shift=float(bagel_cfg.get("timestep_shift", 1.0)),
|
|
),
|
|
)
|
|
|
|
# Let vLLM loader download and stream all *.safetensors under model root.
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder=None,
|
|
revision=od_config.revision,
|
|
prefix="",
|
|
fall_back_to_pt=False,
|
|
)
|
|
]
|
|
|
|
# Defer device placement to the weight-loading/offload path in three cases:
|
|
# 1. Quantization: When quantization is enabled, vLLM linear layers live on meta
|
|
# device until the weight loader materializes them. Calling .to(device) would fail on those meta tensors,
|
|
# so we skip it entirely and let the weight loader handle device placement.
|
|
# 2. Layerwise offload: modules should be initialized on CPU first, then
|
|
# selectively materialized/moved by the offloader.
|
|
# 3. HSDP: weights should be loaded on CPU first and sharded afterwards,
|
|
# rather than eagerly placing the full model on one GPU.
|
|
if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp):
|
|
self.to(self.device)
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
|
)
|
|
|
|
@staticmethod
|
|
def _decode_image_from_latent(
|
|
bagel: Bagel, vae: AutoEncoder, latent: torch.Tensor, image_shape: tuple[int, int]
|
|
) -> Image.Image:
|
|
H, W = image_shape
|
|
h, w = H // bagel.latent_downsample, W // bagel.latent_downsample
|
|
p = bagel.latent_patch_size
|
|
c = bagel.latent_channel
|
|
latent = latent.reshape(1, h, w, p, p, c)
|
|
latent = torch.einsum("nhwpqc->nchpwq", latent)
|
|
latent = latent.reshape(1, c, h * p, w * p)
|
|
|
|
# Cast to VAE dtype (e.g. bfloat16) as latents might remain float32 from generation loop
|
|
vae_dtype = next(vae.parameters()).dtype
|
|
latent = latent.to(vae_dtype)
|
|
|
|
image = vae.decode(latent)
|
|
image = (image * 0.5 + 0.5).clamp(0, 1)[0].permute(1, 2, 0) * 255
|
|
return Image.fromarray(image.to(torch.uint8).cpu().numpy())
|
|
|
|
def _regen_init_noise_on_device(self, gen_input: dict, seed: int | None) -> None:
|
|
"""Resample ``gen_input["packed_init_noises"]`` on-device with a fresh
|
|
per-call ``torch.Generator``.
|
|
|
|
``Bagel.prepare_input`` (and the Lance video equivalent) call
|
|
``torch.randn`` with no device or generator, falling back to CPU+fp32
|
|
via the global RNG. Upstream Lance samples directly on CUDA+bf16 via
|
|
``torch.Generator(device=cuda).manual_seed(seed)`` (lance.py:1536),
|
|
so for the same seed the two sides land on different noise streams.
|
|
Mutates ``gen_input`` in place; no-op if seed is unset or device is CPU.
|
|
"""
|
|
if seed is None or self.device.type != "cuda":
|
|
return
|
|
ref = gen_input["packed_init_noises"]
|
|
gen_input["packed_init_noises"] = torch.randn(
|
|
ref.shape,
|
|
generator=torch.Generator(device=self.device).manual_seed(int(seed)),
|
|
device=self.device,
|
|
dtype=self.od_config.dtype,
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
|
|
if len(req.prompts) > 1:
|
|
logger.warning(
|
|
"""This model only supports a single prompt, not a batched request.""",
|
|
"""Taking only the first image for now.""",
|
|
)
|
|
return self._forward_single(req.prompts[0], req.sampling_params)
|
|
|
|
@torch.inference_mode()
|
|
def _forward_single(
|
|
self,
|
|
first_prompt: OmniPromptType,
|
|
sampling: OmniDiffusionSamplingParams,
|
|
*,
|
|
prepare_only: bool = False,
|
|
) -> DiffusionOutput | dict[str, Any]:
|
|
# TODO: In online mode, sometimes it receives [{"prompts": None}, {...}], so cannot use .get("...", "")
|
|
# TODO: May be some data formatting operations on the API side. Hack for now.
|
|
prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
|
|
|
|
max_hw = int(self.bagel.max_latent_size * self.bagel.latent_downsample)
|
|
if sampling.height is None and sampling.width is None:
|
|
height = width = max_hw
|
|
else:
|
|
height = int(sampling.height) if sampling.height is not None else max_hw
|
|
width = int(sampling.width) if sampling.width is not None else max_hw
|
|
if height > max_hw or width > max_hw:
|
|
raise ValueError(
|
|
f"Requested resolution {height}x{width} exceeds Bagel checkpoint limit "
|
|
f"{max_hw}x{max_hw} (max_latent_size={self.bagel.max_latent_size}, "
|
|
f"latent_downsample={self.bagel.latent_downsample})."
|
|
)
|
|
image_shape = (height, width)
|
|
image_shape_requested = _bagel_canvas_requested(sampling)
|
|
|
|
extra_args = getattr(sampling, "extra_args", {}) or {}
|
|
cfg_text_scale = extra_args.get("cfg_text_scale", 4.0)
|
|
cfg_img_scale = extra_args.get("cfg_img_scale", 1.5)
|
|
|
|
cfg_interval = extra_args.get("cfg_interval", (0.4, 1.0))
|
|
cfg_renorm_type = extra_args.get("cfg_renorm_type", "global")
|
|
cfg_renorm_min = extra_args.get("cfg_renorm_min", 0.0)
|
|
|
|
gen_params = BagelGenParams(
|
|
num_timesteps=int(sampling.num_inference_steps or 50),
|
|
timestep_shift=float(extra_args.get("timestep_shift", 3.0)),
|
|
cfg_text_scale=cfg_text_scale,
|
|
cfg_img_scale=cfg_img_scale,
|
|
cfg_interval=cfg_interval,
|
|
cfg_renorm_type=cfg_renorm_type,
|
|
cfg_renorm_min=cfg_renorm_min,
|
|
)
|
|
|
|
gen_context = {
|
|
"kv_lens": [0],
|
|
"ropes": [0],
|
|
"past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers),
|
|
}
|
|
cfg_text_context = deepcopy(gen_context)
|
|
cfg_img_context = deepcopy(gen_context)
|
|
|
|
injected_kv = sampling.past_key_values
|
|
if injected_kv is not None:
|
|
logger.info("Using injected KV Cache (direct)")
|
|
injected_kv = NaiveCache.from_object(injected_kv)
|
|
gen_context["past_key_values"] = injected_kv
|
|
seq_len = injected_kv.key_cache[0].shape[0]
|
|
gen_context["kv_lens"] = [seq_len]
|
|
if sampling.kv_metadata and "ropes" in sampling.kv_metadata:
|
|
gen_context["ropes"] = sampling.kv_metadata["ropes"]
|
|
else:
|
|
gen_context["ropes"] = [seq_len]
|
|
|
|
if sampling.kv_metadata and "image_shape" in sampling.kv_metadata:
|
|
if not image_shape_requested:
|
|
image_shape = tuple(sampling.kv_metadata["image_shape"])
|
|
elif tuple(sampling.kv_metadata["image_shape"]) != image_shape:
|
|
logger.info(
|
|
"Generating at requested %dx%d instead of KV image shape %s",
|
|
image_shape[1],
|
|
image_shape[0],
|
|
sampling.kv_metadata["image_shape"],
|
|
)
|
|
|
|
branch_kvs = getattr(sampling, "cfg_branch_past_key_values", None) or {}
|
|
branch_metadata = getattr(sampling, "cfg_branch_kv_metadata", None) or {}
|
|
active_branch = getattr(sampling, "cfg_active_branch", None)
|
|
branch_roles = getattr(sampling, "cfg_branch_roles", None) or list(branch_kvs.keys())
|
|
|
|
cfg_text_kv = getattr(sampling, "cfg_text_past_key_values", None) or branch_kvs.get("cfg_text")
|
|
cfg_text_metadata = getattr(sampling, "cfg_text_kv_metadata", None) or branch_metadata.get("cfg_text")
|
|
cfg_img_kv = getattr(sampling, "cfg_img_past_key_values", None) or branch_kvs.get("cfg_img")
|
|
cfg_img_metadata = getattr(sampling, "cfg_img_kv_metadata", None) or branch_metadata.get("cfg_img")
|
|
|
|
cfg_parallel_contract = (
|
|
active_branch is not None or bool(branch_roles) or cfg_text_kv is not None or cfg_img_kv is not None
|
|
)
|
|
if cfg_parallel_contract:
|
|
logger.info(
|
|
"CFG enabled with injected branch KV context roles=%s active=%s",
|
|
branch_roles,
|
|
active_branch,
|
|
)
|
|
|
|
if cfg_text_kv is not None:
|
|
cfg_text_kv = NaiveCache.from_object(cfg_text_kv)
|
|
cfg_text_seq_len = cfg_text_kv.key_cache[0].shape[0]
|
|
cfg_text_context["past_key_values"] = cfg_text_kv
|
|
cfg_text_context["kv_lens"] = [cfg_text_seq_len]
|
|
if cfg_text_metadata and "ropes" in cfg_text_metadata:
|
|
cfg_text_context["ropes"] = cfg_text_metadata["ropes"]
|
|
else:
|
|
cfg_text_context["ropes"] = [cfg_text_seq_len]
|
|
else:
|
|
# No cfg_text companion received. For text2img this is the
|
|
# expected path: original BAGEL uses an empty KV cache (0
|
|
# tokens) as the text-unconditional branch. Keep the default
|
|
# empty NaiveCache in cfg_text_context and preserve the
|
|
# original cfg_text_scale so CFG still applies.
|
|
pass
|
|
|
|
if cfg_img_kv is None:
|
|
# text2img multi-stage: cfg_img reuses gen KV (positive prompt,
|
|
# no image), mirroring forward_cache_update_text on cfg_img_context
|
|
# in the single-stage path.
|
|
cfg_img_seq_len = injected_kv.key_cache[0].shape[0]
|
|
cfg_img_context["past_key_values"] = injected_kv
|
|
cfg_img_context["kv_lens"] = [cfg_img_seq_len]
|
|
if sampling.kv_metadata and "ropes" in sampling.kv_metadata:
|
|
cfg_img_context["ropes"] = sampling.kv_metadata["ropes"]
|
|
else:
|
|
cfg_img_context["ropes"] = [cfg_img_seq_len]
|
|
else:
|
|
cfg_img_kv = NaiveCache.from_object(cfg_img_kv)
|
|
cfg_img_seq_len = cfg_img_kv.key_cache[0].shape[0]
|
|
cfg_img_context["past_key_values"] = cfg_img_kv
|
|
cfg_img_context["kv_lens"] = [cfg_img_seq_len]
|
|
if cfg_img_metadata and "ropes" in cfg_img_metadata:
|
|
cfg_img_context["ropes"] = cfg_img_metadata["ropes"]
|
|
else:
|
|
cfg_img_context["ropes"] = [cfg_img_seq_len]
|
|
|
|
else:
|
|
image_input = (
|
|
None
|
|
if isinstance(first_prompt, str)
|
|
else (
|
|
(first_prompt.get("multi_modal_data") or {}).get("image")
|
|
or (first_prompt.get("multi_modal_data") or {}).get("img2img")
|
|
)
|
|
)
|
|
if image_input and not isinstance(image_input, list):
|
|
image_input = [image_input]
|
|
if image_input:
|
|
image_input = [Image.open(image) if isinstance(image, str) else image for image in image_input]
|
|
|
|
if image_input:
|
|
# If we have an image, we prefill with it
|
|
if self.image_processor and self.vae:
|
|
|
|
def vit_transforms(img):
|
|
return self.image_processor(images=img, return_tensors="pt").pixel_values[0]
|
|
|
|
stride = self.bagel.latent_downsample
|
|
max_img_size = int(self.bagel.max_latent_size * stride)
|
|
|
|
def _resize_to_stride(img):
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
w, h = img.size
|
|
# Scale down if longest edge exceeds max
|
|
scale = min(max_img_size / max(w, h), 1.0)
|
|
# Scale up if shortest edge is too small (min 256)
|
|
min_img_size = min(256, max_img_size)
|
|
scale = max(scale, min_img_size / min(w, h))
|
|
new_w = max(stride, int(round(w * scale / stride) * stride))
|
|
new_h = max(stride, int(round(h * scale / stride) * stride))
|
|
# Clamp to max
|
|
new_w = min(new_w, max_img_size)
|
|
new_h = min(new_h, max_img_size)
|
|
if new_w != w or new_h != h:
|
|
img = img.resize((new_w, new_h), Image.BICUBIC)
|
|
return img
|
|
|
|
image_input = [_resize_to_stride(img) for img in image_input]
|
|
|
|
resized_w, resized_h = image_input[0].size
|
|
if not image_shape_requested:
|
|
image_shape = (resized_h, resized_w)
|
|
logger.info(
|
|
"img2img: resized source to %dx%d, generating at %dx%d",
|
|
resized_w,
|
|
resized_h,
|
|
image_shape[1],
|
|
image_shape[0],
|
|
)
|
|
|
|
def vae_transforms(img):
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
# Convert to [-1, 1] tensor (H, W, C) -> (C, H, W)
|
|
arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0
|
|
return arr.permute(2, 0, 1)
|
|
|
|
# Update gen_context with image (VAE + ViT)
|
|
gen_input_vae, newlens_vae, new_rope_vae = self.bagel.prepare_vae_images(
|
|
curr_kvlens=gen_context["kv_lens"],
|
|
curr_rope=gen_context["ropes"],
|
|
images=image_input,
|
|
transforms=vae_transforms,
|
|
new_token_ids=self.new_token_ids,
|
|
)
|
|
for k, v in gen_input_vae.items():
|
|
if torch.is_tensor(v):
|
|
gen_input_vae[k] = v.to(self.device)
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
gen_context["past_key_values"] = self.bagel.forward_cache_update_vae(
|
|
self.vae, gen_context["past_key_values"], **gen_input_vae
|
|
)
|
|
gen_context["kv_lens"] = newlens_vae
|
|
gen_context["ropes"] = new_rope_vae
|
|
|
|
gen_input_img, newlens_img, new_rope_img = self.bagel.prepare_vit_images(
|
|
curr_kvlens=gen_context["kv_lens"],
|
|
curr_rope=gen_context["ropes"],
|
|
images=image_input,
|
|
transforms=vit_transforms,
|
|
new_token_ids=self.new_token_ids,
|
|
)
|
|
for k, v in gen_input_img.items():
|
|
if torch.is_tensor(v):
|
|
gen_input_img[k] = v.to(self.device)
|
|
for k in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"):
|
|
gen_input_img.pop(k, None)
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
|
|
gen_context["past_key_values"], **gen_input_img
|
|
)
|
|
gen_context["kv_lens"] = newlens_img
|
|
gen_context["ropes"] = new_rope_img
|
|
|
|
cfg_text_context = deepcopy(gen_context)
|
|
|
|
# Strip <|im_start|>/<|im_end|> wrappers that end2end.py may have
|
|
# already added, so prepare_prompts doesn't double-add bos/eos.
|
|
clean_prompt = prompt.removeprefix("<|im_start|>").removesuffix("<|im_end|>")
|
|
|
|
# Update gen_context with text prompt
|
|
generation_input, newlens, new_rope = self.bagel.prepare_prompts(
|
|
curr_kvlens=gen_context["kv_lens"],
|
|
curr_rope=gen_context["ropes"],
|
|
prompts=[clean_prompt],
|
|
tokenizer=self.tokenizer,
|
|
new_token_ids=self.new_token_ids,
|
|
)
|
|
# Fail fast with a clear error instead of CUDA gather OOB.
|
|
max_tid = int(generation_input["packed_text_ids"].max().item())
|
|
emb_n = int(self.language_model.vocab_size)
|
|
if max_tid >= emb_n:
|
|
raise ValueError(
|
|
"Tokenizer/model vocab mismatch: max token id "
|
|
f"{max_tid} >= embed_tokens size {emb_n}. "
|
|
"This usually means you're not using the tokenizer shipped with the Bagel checkpoint, "
|
|
"or llm_config.vocab_size is smaller than the tokenizer vocab."
|
|
)
|
|
for k, v in generation_input.items():
|
|
if torch.is_tensor(v):
|
|
generation_input[k] = v.to(self.device)
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
gen_context["past_key_values"] = self.bagel.forward_cache_update_text(
|
|
gen_context["past_key_values"], **generation_input
|
|
)
|
|
gen_context["kv_lens"] = newlens
|
|
gen_context["ropes"] = new_rope
|
|
|
|
# cfg_text_context: update with negative prompt (no text condition).
|
|
# When empty, keep cfg_text_context as-is (kv_lens=0) to match
|
|
# original BAGEL.
|
|
prompt_negative = first_prompt.get("negative_prompt") if isinstance(first_prompt, dict) else None
|
|
neg_prompt = prompt_negative if prompt_negative is not None else extra_args.get("negative_prompt", "")
|
|
if neg_prompt:
|
|
neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts(
|
|
curr_kvlens=cfg_text_context["kv_lens"],
|
|
curr_rope=cfg_text_context["ropes"],
|
|
prompts=[neg_prompt],
|
|
tokenizer=self.tokenizer,
|
|
new_token_ids=self.new_token_ids,
|
|
)
|
|
for k, v in neg_input.items():
|
|
if torch.is_tensor(v):
|
|
neg_input[k] = v.to(self.device)
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
cfg_text_context["past_key_values"] = self.bagel.forward_cache_update_text(
|
|
cfg_text_context["past_key_values"], **neg_input
|
|
)
|
|
cfg_text_context["kv_lens"] = neg_newlens
|
|
cfg_text_context["ropes"] = neg_rope
|
|
|
|
# cfg_img_context: update with text prompt (no image condition)
|
|
cfg_img_generation_input, cfg_img_newlens, cfg_img_new_rope = self.bagel.prepare_prompts(
|
|
curr_kvlens=cfg_img_context["kv_lens"],
|
|
curr_rope=cfg_img_context["ropes"],
|
|
prompts=[clean_prompt],
|
|
tokenizer=self.tokenizer,
|
|
new_token_ids=self.new_token_ids,
|
|
)
|
|
for k, v in cfg_img_generation_input.items():
|
|
if torch.is_tensor(v):
|
|
cfg_img_generation_input[k] = v.to(self.device)
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
cfg_img_context["past_key_values"] = self.bagel.forward_cache_update_text(
|
|
cfg_img_context["past_key_values"], **cfg_img_generation_input
|
|
)
|
|
cfg_img_context["kv_lens"] = cfg_img_newlens
|
|
cfg_img_context["ropes"] = cfg_img_new_rope
|
|
|
|
# ---- Detect output modality and think mode ----
|
|
modalities = first_prompt.get("modalities", []) if isinstance(first_prompt, dict) else []
|
|
is_text_output = "text" in modalities
|
|
think_enabled = extra_args.get("think", False)
|
|
think_text = None
|
|
|
|
if think_enabled and injected_kv is None:
|
|
max_think_tokens = int(extra_args.get("max_think_tokens", 1000))
|
|
do_sample = bool(extra_args.get("do_sample", False))
|
|
text_temperature = float(extra_args.get("text_temperature", 0.3))
|
|
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
start_input = self.bagel.prepare_start_tokens(
|
|
gen_context["kv_lens"], gen_context["ropes"], self.new_token_ids
|
|
)
|
|
for k, v in start_input.items():
|
|
if torch.is_tensor(v):
|
|
start_input[k] = v.to(self.device)
|
|
|
|
gen_ctx_copy = deepcopy(gen_context)
|
|
token_ids = self.bagel.generate_text(
|
|
past_key_values=gen_ctx_copy["past_key_values"],
|
|
max_length=max_think_tokens,
|
|
do_sample=do_sample,
|
|
temperature=text_temperature,
|
|
end_token_id=self.new_token_ids["eos_token_id"],
|
|
**start_input,
|
|
)
|
|
# token_ids shape: (seq_len, batch=1)
|
|
decoded = self.tokenizer.decode(token_ids[:, 0].tolist())
|
|
# Strip chat markers to get clean text
|
|
think_text = decoded.split("<|im_end|>")[0]
|
|
if "<|im_start|>" in think_text:
|
|
think_text = think_text.split("<|im_start|>")[-1]
|
|
logger.info("Think mode generated %d tokens", token_ids.shape[0])
|
|
|
|
if not is_text_output:
|
|
# Use the autoregressive KV cache from think generation
|
|
# directly, instead of decode→re-encode which adds extra
|
|
# bos/eos and may alter tokenization.
|
|
num_think_tokens = token_ids.shape[0]
|
|
gen_context["past_key_values"] = gen_ctx_copy["past_key_values"]
|
|
gen_context["kv_lens"] = [kl + num_think_tokens for kl in gen_context["kv_lens"]]
|
|
gen_context["ropes"] = [r + num_think_tokens for r in gen_context["ropes"]]
|
|
|
|
# ---- Text-only output (text2text / img2text) ----
|
|
if is_text_output and injected_kv is None:
|
|
if think_text is not None:
|
|
# Think mode already generated the text (including reasoning)
|
|
text_output = think_text
|
|
else:
|
|
max_text_tokens = int(extra_args.get("max_think_tokens", 500))
|
|
do_sample = bool(extra_args.get("do_sample", False))
|
|
text_temperature = float(extra_args.get("text_temperature", 0.3))
|
|
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
start_input = self.bagel.prepare_start_tokens(
|
|
gen_context["kv_lens"], gen_context["ropes"], self.new_token_ids
|
|
)
|
|
for k, v in start_input.items():
|
|
if torch.is_tensor(v):
|
|
start_input[k] = v.to(self.device)
|
|
token_ids = self.bagel.generate_text(
|
|
past_key_values=gen_context["past_key_values"],
|
|
max_length=max_text_tokens,
|
|
do_sample=do_sample,
|
|
temperature=text_temperature,
|
|
end_token_id=self.new_token_ids["eos_token_id"],
|
|
**start_input,
|
|
)
|
|
decoded = self.tokenizer.decode(token_ids[:, 0].tolist())
|
|
text_output = decoded.split("<|im_end|>")[0]
|
|
if "<|im_start|>" in text_output:
|
|
text_output = text_output.split("<|im_start|>")[-1]
|
|
|
|
return DiffusionOutput(
|
|
output={
|
|
"payload": {"text": text_output},
|
|
"metadata": {"text": {"text_output": text_output}},
|
|
},
|
|
stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
|
|
)
|
|
|
|
# ---- Image generation (text2img / img2img) ----
|
|
if sampling.seed is not None:
|
|
torch.manual_seed(sampling.seed)
|
|
if self.device.type == "cuda":
|
|
torch.cuda.manual_seed(sampling.seed)
|
|
|
|
generation_input = self.bagel.prepare_vae_latent(
|
|
curr_kvlens=gen_context["kv_lens"],
|
|
curr_rope=gen_context["ropes"],
|
|
image_sizes=[image_shape],
|
|
new_token_ids=self.new_token_ids,
|
|
)
|
|
# Fail fast for special tokens used by the image path as well.
|
|
max_tid_img = int(generation_input["packed_text_ids"].max().item())
|
|
emb_n = int(self.language_model.vocab_size)
|
|
if max_tid_img >= emb_n:
|
|
raise ValueError(
|
|
"Tokenizer/model vocab mismatch (image path): max token id "
|
|
f"{max_tid_img} >= embed_tokens size {emb_n}. "
|
|
"This indicates the tokenizer token IDs do not match the checkpoint embeddings."
|
|
)
|
|
# Position ids must be non-negative; negative ids can trigger CUDA gather OOB inside RoPE.
|
|
min_pid = int(generation_input["packed_position_ids"].min().item())
|
|
if min_pid < 0:
|
|
raise ValueError(f"Invalid packed_position_ids: min={min_pid} (must be >= 0)")
|
|
# Latent position embedding bounds check: ids must be < max_latent_size^2.
|
|
max_lat_pid = int(generation_input["packed_vae_position_ids"].max().item())
|
|
max_lat_pid_allowed = int(self.bagel.max_latent_size * self.bagel.max_latent_size) - 1
|
|
if max_lat_pid > max_lat_pid_allowed:
|
|
raise ValueError(
|
|
"Invalid packed_vae_position_ids (latent position embedding OOB): "
|
|
f"max={max_lat_pid} > allowed_max={max_lat_pid_allowed}. "
|
|
f"Requested image_shape={image_shape}, max_latent_size={self.bagel.max_latent_size}."
|
|
)
|
|
for k, v in generation_input.items():
|
|
if torch.is_tensor(v):
|
|
generation_input[k] = v.to(self.device)
|
|
|
|
# NOTE: For now we disable device specific noise regeneration so that e2e tests can run
|
|
# on both CUDA and ROCm. Context: https://github.com/vllm-project/vllm-omni/pull/4081
|
|
# self._regen_init_noise_on_device(generation_input, sampling.seed)
|
|
|
|
# text cfg
|
|
generation_input_cfg_text = self.bagel.prepare_vae_latent_cfg(
|
|
curr_kvlens=cfg_text_context["kv_lens"],
|
|
curr_rope=cfg_text_context["ropes"],
|
|
image_sizes=[image_shape],
|
|
)
|
|
# img cfg
|
|
generation_input_cfg_img = self.bagel.prepare_vae_latent_cfg(
|
|
curr_kvlens=cfg_img_context["kv_lens"],
|
|
curr_rope=cfg_img_context["ropes"],
|
|
image_sizes=[image_shape],
|
|
)
|
|
for k, v in generation_input_cfg_text.items():
|
|
if torch.is_tensor(v):
|
|
generation_input_cfg_text[k] = v.to(self.device)
|
|
for k, v in generation_input_cfg_img.items():
|
|
if torch.is_tensor(v):
|
|
generation_input_cfg_img[k] = v.to(self.device)
|
|
|
|
if prepare_only:
|
|
return {
|
|
"generation_input": generation_input,
|
|
"cfg_text_packed_position_ids": generation_input_cfg_text["cfg_packed_position_ids"],
|
|
"cfg_img_packed_position_ids": generation_input_cfg_img["cfg_packed_position_ids"],
|
|
"gen_context": gen_context,
|
|
"cfg_text_context": cfg_text_context,
|
|
"cfg_img_context": cfg_img_context,
|
|
"gen_params": gen_params,
|
|
"image_shape": image_shape,
|
|
"think_text": think_text,
|
|
}
|
|
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
latents, trajectory_latents, trajectory_timesteps, trajectory_log_probs = self.bagel.generate_image(
|
|
past_key_values=gen_context["past_key_values"],
|
|
cfg_text_past_key_values=cfg_text_context["past_key_values"],
|
|
cfg_img_past_key_values=cfg_img_context["past_key_values"],
|
|
num_timesteps=gen_params.num_timesteps,
|
|
timestep_shift=gen_params.timestep_shift,
|
|
cfg_text_scale=gen_params.cfg_text_scale,
|
|
cfg_img_scale=gen_params.cfg_img_scale,
|
|
cfg_interval=gen_params.cfg_interval,
|
|
cfg_renorm_min=gen_params.cfg_renorm_min,
|
|
cfg_renorm_type=gen_params.cfg_renorm_type,
|
|
**generation_input,
|
|
cfg_text_packed_position_ids=generation_input_cfg_text["cfg_packed_position_ids"],
|
|
cfg_img_packed_position_ids=generation_input_cfg_img["cfg_packed_position_ids"],
|
|
return_trajectory_latents=sampling.return_trajectory_latents,
|
|
scheduler=self.scheduler,
|
|
scheduler_kwargs=self.scheduler_kwargs,
|
|
)
|
|
|
|
return self._build_image_output(
|
|
latents[0],
|
|
image_shape,
|
|
trajectory_latents=trajectory_latents,
|
|
trajectory_timesteps=trajectory_timesteps,
|
|
trajectory_log_probs=trajectory_log_probs,
|
|
return_trajectory_decoded=sampling.return_trajectory_decoded,
|
|
think_text=think_text,
|
|
)
|
|
|
|
def _build_image_output(
|
|
self,
|
|
latent: torch.Tensor,
|
|
image_shape: tuple[int, int],
|
|
*,
|
|
trajectory_latents: list[torch.Tensor] | None = None,
|
|
trajectory_timesteps: list[torch.Tensor] | None = None,
|
|
trajectory_log_probs: list[torch.Tensor] | None = None,
|
|
return_trajectory_decoded: bool = False,
|
|
think_text: str | None = None,
|
|
) -> DiffusionOutput:
|
|
img = self._decode_image_from_latent(self.bagel, self.vae, latent, image_shape)
|
|
|
|
# Build trajectory output when requested
|
|
trajectory_latents_stacked: torch.Tensor | None = None
|
|
trajectory_timesteps_stacked: torch.Tensor | None = None
|
|
trajectory_decoded: list[Image.Image] | None = None
|
|
if trajectory_latents:
|
|
trajectory_latents_stacked = torch.stack(trajectory_latents)
|
|
trajectory_timesteps_stacked = torch.stack(trajectory_timesteps)
|
|
if return_trajectory_decoded:
|
|
trajectory_decoded = [
|
|
self._decode_image_from_latent(self.bagel, self.vae, lat, image_shape) for lat in trajectory_latents
|
|
]
|
|
|
|
trajectory_log_probs_stacked: torch.Tensor | None = None
|
|
if trajectory_log_probs:
|
|
trajectory_log_probs_stacked = torch.stack(trajectory_log_probs)
|
|
|
|
payload = {"image": img}
|
|
metadata = {}
|
|
if think_text is not None:
|
|
metadata["text"] = {"think_text": think_text}
|
|
trajectory_payload = {}
|
|
if trajectory_latents_stacked is not None:
|
|
trajectory_payload["latents"] = trajectory_latents_stacked
|
|
if trajectory_timesteps_stacked is not None:
|
|
trajectory_payload["timesteps"] = trajectory_timesteps_stacked
|
|
if trajectory_log_probs_stacked is not None:
|
|
trajectory_payload["log_probs"] = trajectory_log_probs_stacked
|
|
if trajectory_decoded is not None:
|
|
trajectory_payload["decoded"] = trajectory_decoded
|
|
if trajectory_payload:
|
|
payload["trajectory"] = trajectory_payload
|
|
metadata["trajectory"] = {"type": "denoising"}
|
|
|
|
return DiffusionOutput(
|
|
output={
|
|
"payload": payload,
|
|
"metadata": metadata,
|
|
},
|
|
stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
|
|
)
|
|
|
|
def prepare_encode(
|
|
self,
|
|
state: StepRequestState,
|
|
**kwargs: object,
|
|
) -> StepRequestState:
|
|
"""Populate *state* with BAGEL inputs, latents, timesteps, and CFG config."""
|
|
del kwargs
|
|
if self.bagel._sp_size > 1:
|
|
raise NotImplementedError("BAGEL step execution does not currently support sequence parallelism.")
|
|
|
|
sampling = state.sampling
|
|
prompt = state.prompt if state.prompt is not None else ""
|
|
if isinstance(prompt, dict):
|
|
if "text" in (prompt.get("modalities") or []):
|
|
raise NotImplementedError("BAGEL text output is not supported by step execution.")
|
|
|
|
ctx = self._forward_single(prompt, sampling, prepare_only=True)
|
|
if not isinstance(ctx, dict):
|
|
raise RuntimeError("BAGEL step preparation did not produce an image-generation context.")
|
|
|
|
generation_input = ctx["generation_input"]
|
|
latents = generation_input["packed_init_noises"]
|
|
gen_params = ctx["gen_params"]
|
|
timesteps, dts = self.bagel.prepare_denoise_schedule(
|
|
latents,
|
|
gen_params.num_timesteps,
|
|
gen_params.timestep_shift,
|
|
)
|
|
|
|
# Keep request-local scheduler progress while sharing immutable config
|
|
# and tensor buffers with the pipeline scheduler.
|
|
req_scheduler = copy(self.scheduler)
|
|
|
|
# Populate state from generation context.
|
|
state.latents = latents
|
|
state.timesteps = timesteps
|
|
state.step_index = 0
|
|
state.scheduler = req_scheduler
|
|
state.do_true_cfg = gen_params.cfg_text_scale > 1.0
|
|
state.img_shapes = [ctx["image_shape"]]
|
|
state.extra.update(
|
|
{
|
|
"bagel_generation_input": {
|
|
key: value for key, value in generation_input.items() if key != "packed_init_noises"
|
|
},
|
|
"bagel_cfg_text_packed_position_ids": ctx["cfg_text_packed_position_ids"],
|
|
"bagel_cfg_img_packed_position_ids": ctx["cfg_img_packed_position_ids"],
|
|
"bagel_gen_context": ctx["gen_context"],
|
|
"bagel_cfg_text_context": ctx["cfg_text_context"],
|
|
"bagel_cfg_img_context": ctx["cfg_img_context"],
|
|
"bagel_gen_params": gen_params,
|
|
"bagel_dts": dts,
|
|
"bagel_image_shape": ctx["image_shape"],
|
|
"bagel_scheduler_kwargs": dict(self.scheduler_kwargs),
|
|
"bagel_think_text": ctx["think_text"],
|
|
"bagel_trajectory_latents": (
|
|
# Trajectory output needs immutable snapshots because each
|
|
# later scheduler update replaces or mutates the latents.
|
|
[latents.clone()] if sampling.return_trajectory_latents and len(timesteps) > 0 else []
|
|
),
|
|
"bagel_trajectory_timesteps": [],
|
|
"bagel_trajectory_log_probs": [],
|
|
}
|
|
)
|
|
return state
|
|
|
|
@staticmethod
|
|
def _pack_step_generation_inputs(
|
|
states: Sequence[StepRequestState],
|
|
) -> dict[str, torch.Tensor]:
|
|
generation_inputs = [state.extra["bagel_generation_input"] for state in states]
|
|
packed_seqlens = [generation_input["packed_seqlens"] for generation_input in generation_inputs]
|
|
if any(seqlens.numel() != 1 for seqlens in packed_seqlens):
|
|
raise ValueError("BAGEL step execution expects one packed sequence per request.")
|
|
seq_lengths = [int(seqlens.item()) for seqlens in packed_seqlens]
|
|
if len(set(seq_lengths)) != 1:
|
|
raise ValueError("BAGEL step batching requires matching packed sequence lengths.")
|
|
|
|
query_offset = 0
|
|
text_indexes = []
|
|
vae_indexes = []
|
|
for generation_input, seq_len in zip(generation_inputs, seq_lengths, strict=True):
|
|
text_indexes.append(generation_input["packed_text_indexes"] + query_offset)
|
|
vae_indexes.append(generation_input["packed_vae_token_indexes"] + query_offset)
|
|
query_offset += seq_len
|
|
|
|
return {
|
|
"packed_text_ids": torch.cat([item["packed_text_ids"] for item in generation_inputs]),
|
|
"packed_text_indexes": torch.cat(text_indexes),
|
|
"packed_vae_position_ids": torch.cat([item["packed_vae_position_ids"] for item in generation_inputs]),
|
|
"packed_vae_token_indexes": torch.cat(vae_indexes),
|
|
"packed_seqlens": torch.cat(packed_seqlens),
|
|
"packed_position_ids": torch.cat([item["packed_position_ids"] for item in generation_inputs]),
|
|
}
|
|
|
|
def _build_denoise_kwargs(
|
|
self,
|
|
input_batch: InputBatch,
|
|
states: Sequence[StepRequestState],
|
|
) -> dict[str, Any]:
|
|
parallel_config = getattr(self.bagel, "parallel_config", None)
|
|
cfg_parallel_size = getattr(parallel_config, "cfg_parallel_size", 1)
|
|
|
|
packed = self._pack_step_generation_inputs(states)
|
|
vae_lengths = [int(state.latents.shape[0]) for state in states]
|
|
if sum(vae_lengths) != int(input_batch.latents.shape[0]):
|
|
raise ValueError("BAGEL packed latent rows do not match InputBatch.latents.")
|
|
|
|
cfg_settings = {
|
|
(
|
|
state.extra["bagel_gen_params"].cfg_text_scale,
|
|
state.extra["bagel_gen_params"].cfg_img_scale,
|
|
tuple(state.extra["bagel_gen_params"].cfg_interval),
|
|
state.extra["bagel_gen_params"].cfg_renorm_type,
|
|
state.extra["bagel_gen_params"].cfg_renorm_min,
|
|
)
|
|
for state in states
|
|
}
|
|
if len(cfg_settings) != 1:
|
|
raise ValueError("Mixed BAGEL CFG settings cannot share one step batch.")
|
|
cfg_text_scale, cfg_img_scale, _cfg_interval, cfg_renorm_type, cfg_renorm_min = next(iter(cfg_settings))
|
|
|
|
cfg_text_scales = []
|
|
cfg_img_scales = []
|
|
for state in states:
|
|
gen_params = state.extra["bagel_gen_params"]
|
|
timestep = state.current_timestep
|
|
if timestep is None:
|
|
raise ValueError(f"BAGEL request {state.request_id} has no current timestep.")
|
|
t_value = float(timestep.item())
|
|
in_cfg_window = t_value > gen_params.cfg_interval[0] and t_value <= gen_params.cfg_interval[1]
|
|
text_scale = gen_params.cfg_text_scale if in_cfg_window else 1.0
|
|
cfg_text_scales.append(text_scale)
|
|
cfg_img_scales.append(gen_params.cfg_img_scale if in_cfg_window and text_scale > 1.0 else 1.0)
|
|
|
|
use_cfg_text = any(scale > 1.0 for scale in cfg_text_scales)
|
|
use_cfg_img = use_cfg_text and any(scale > 1.0 for scale in cfg_img_scales)
|
|
configured_cfg_text = any(state.extra["bagel_gen_params"].cfg_text_scale > 1.0 for state in states)
|
|
configured_cfg_img = any(
|
|
state.extra["bagel_gen_params"].cfg_text_scale > 1.0 and state.extra["bagel_gen_params"].cfg_img_scale > 1.0
|
|
for state in states
|
|
)
|
|
build_cfg_text = use_cfg_text or (cfg_parallel_size > 1 and configured_cfg_text)
|
|
build_cfg_img = use_cfg_img or (cfg_parallel_size > 1 and configured_cfg_img)
|
|
gen_cache = NaiveCache.merge([state.extra["bagel_gen_context"]["past_key_values"] for state in states])
|
|
cfg_branch_pids = None
|
|
cfg_branch_caches = None
|
|
if build_cfg_text:
|
|
cfg_branch_pids = [
|
|
packed["packed_position_ids"],
|
|
torch.cat([state.extra["bagel_cfg_text_packed_position_ids"] for state in states]),
|
|
]
|
|
cfg_branch_caches = [
|
|
gen_cache,
|
|
NaiveCache.merge([state.extra["bagel_cfg_text_context"]["past_key_values"] for state in states]),
|
|
]
|
|
if build_cfg_img:
|
|
cfg_branch_pids.append(
|
|
torch.cat([state.extra["bagel_cfg_img_packed_position_ids"] for state in states])
|
|
)
|
|
cfg_branch_caches.append(
|
|
NaiveCache.merge([state.extra["bagel_cfg_img_context"]["past_key_values"] for state in states])
|
|
)
|
|
|
|
return {
|
|
"x_t": input_batch.latents,
|
|
"timestep": input_batch.timesteps,
|
|
"past_key_values": gen_cache,
|
|
"cfg_renorm_min": cfg_renorm_min,
|
|
"cfg_renorm_type": cfg_renorm_type,
|
|
"cfg_text_scale": cfg_text_scale,
|
|
"cfg_img_scale": cfg_img_scale,
|
|
"cfg_branch_pids": cfg_branch_pids,
|
|
"cfg_branch_caches": cfg_branch_caches,
|
|
"cfg_vae_lengths": vae_lengths if use_cfg_text else None,
|
|
"cfg_text_scales": cfg_text_scales if use_cfg_text else None,
|
|
"cfg_img_scales": cfg_img_scales if use_cfg_text else None,
|
|
**packed,
|
|
}
|
|
|
|
def _denoise_step_cfg_parallel(self, denoise_kwargs: dict[str, Any]) -> torch.Tensor:
|
|
cfg_branch_pids = denoise_kwargs["cfg_branch_pids"]
|
|
cfg_branch_caches = denoise_kwargs["cfg_branch_caches"]
|
|
if cfg_branch_pids is None or cfg_branch_caches is None:
|
|
return self.bagel.forward(**denoise_kwargs)
|
|
|
|
cfg_world_size = get_classifier_free_guidance_world_size()
|
|
num_branches = len(cfg_branch_pids)
|
|
if cfg_world_size == 2 and num_branches == 3:
|
|
raise ValueError(
|
|
f"Image CFG (cfg_img_scale={denoise_kwargs['cfg_img_scale']}) requires cfg_parallel_size=3, "
|
|
"but got cfg_parallel_size=2. Use cfg_parallel_size=3 to enable image CFG in parallel mode."
|
|
)
|
|
|
|
common_keys = (
|
|
"x_t",
|
|
"timestep",
|
|
"packed_vae_token_indexes",
|
|
"packed_vae_position_ids",
|
|
"packed_text_ids",
|
|
"packed_text_indexes",
|
|
"packed_seqlens",
|
|
)
|
|
common = {key: denoise_kwargs[key] for key in common_keys}
|
|
branches_kwargs = [
|
|
{**common, "packed_position_ids": position_ids, "past_key_values": cache}
|
|
for position_ids, cache in zip(cfg_branch_pids, cfg_branch_caches, strict=True)
|
|
]
|
|
result = self.bagel.predict_noise_with_multi_branch_cfg(
|
|
do_true_cfg=denoise_kwargs["cfg_text_scales"] is not None,
|
|
true_cfg_scale={
|
|
"cfg_text_scale": denoise_kwargs["cfg_text_scale"],
|
|
"cfg_img_scale": denoise_kwargs["cfg_img_scale"],
|
|
"cfg_renorm_type": denoise_kwargs["cfg_renorm_type"],
|
|
"cfg_renorm_min": denoise_kwargs["cfg_renorm_min"],
|
|
"cfg_vae_lengths": denoise_kwargs["cfg_vae_lengths"],
|
|
"cfg_text_scales": denoise_kwargs["cfg_text_scales"],
|
|
"cfg_img_scales": denoise_kwargs["cfg_img_scales"],
|
|
},
|
|
branches_kwargs=branches_kwargs,
|
|
)
|
|
if not isinstance(result, torch.Tensor):
|
|
raise TypeError("BAGEL CFG parallel step must return one velocity tensor.")
|
|
return result
|
|
|
|
def denoise_step(
|
|
self,
|
|
input_batch: InputBatch,
|
|
*,
|
|
states: Sequence[StepRequestState] | None = None,
|
|
**kwargs: object,
|
|
) -> torch.Tensor:
|
|
"""One denoise step: pack request-local BAGEL state and run the model."""
|
|
del kwargs
|
|
states = tuple(states or input_batch.states)
|
|
if not states:
|
|
raise ValueError("BAGEL denoise_step requires at least one request state.")
|
|
denoise_kwargs = self._build_denoise_kwargs(input_batch, states)
|
|
parallel_config = getattr(self.bagel, "parallel_config", None)
|
|
use_cfg_parallel = (
|
|
getattr(parallel_config, "cfg_parallel_size", 1) > 1 and denoise_kwargs["cfg_branch_pids"] is not None
|
|
)
|
|
|
|
with torch.autocast(
|
|
device_type=self.device.type,
|
|
enabled=self.device.type != "cpu",
|
|
dtype=self.od_config.dtype,
|
|
):
|
|
if use_cfg_parallel:
|
|
if any(state.step_index == 0 for state in states):
|
|
get_cfg_group().broadcast(input_batch.latents, src=0)
|
|
return self._denoise_step_cfg_parallel(denoise_kwargs)
|
|
return self.bagel.forward(**denoise_kwargs)
|
|
|
|
def step_scheduler(
|
|
self,
|
|
state: StepRequestState,
|
|
noise_pred: torch.Tensor,
|
|
**kwargs: object,
|
|
) -> None:
|
|
"""One scheduler step: update ``state.latents`` and advance ``step_index``."""
|
|
del kwargs
|
|
t = state.current_timestep
|
|
if t is None or state.latents is None:
|
|
raise ValueError(f"BAGEL request {state.request_id} is not ready for a scheduler step.")
|
|
dt = state.extra["bagel_dts"][state.step_index]
|
|
log_prob = None
|
|
if state.scheduler is not None:
|
|
output = state.scheduler.step(
|
|
noise_pred.to(state.latents.device),
|
|
t,
|
|
state.latents,
|
|
dt,
|
|
**state.extra["bagel_scheduler_kwargs"],
|
|
)
|
|
state.latents = output.prev_sample
|
|
log_prob = getattr(output, "log_prob", None)
|
|
else:
|
|
state.latents = state.latents - noise_pred.to(state.latents.device) * dt
|
|
|
|
if state.sampling.return_trajectory_latents:
|
|
# Preserve this opt-in step snapshot before the next update.
|
|
state.extra["bagel_trajectory_latents"].append(state.latents.clone())
|
|
state.extra["bagel_trajectory_timesteps"].append(t)
|
|
if log_prob is not None:
|
|
state.extra["bagel_trajectory_log_probs"].append(log_prob)
|
|
state.step_index += 1
|
|
|
|
def post_decode(
|
|
self,
|
|
state: StepRequestState,
|
|
**kwargs: object,
|
|
) -> DiffusionOutput:
|
|
"""Decode final latents from *state*."""
|
|
del kwargs
|
|
if state.latents is None:
|
|
raise ValueError(f"BAGEL request {state.request_id} has no latents to decode.")
|
|
image_shape = state.extra["bagel_image_shape"]
|
|
return_trajectory_decoded = state.sampling.return_trajectory_decoded
|
|
return self._build_image_output(
|
|
state.latents,
|
|
image_shape,
|
|
trajectory_latents=state.extra["bagel_trajectory_latents"],
|
|
trajectory_timesteps=state.extra["bagel_trajectory_timesteps"],
|
|
trajectory_log_probs=state.extra["bagel_trajectory_log_probs"],
|
|
return_trajectory_decoded=return_trajectory_decoded,
|
|
think_text=state.extra["bagel_think_text"],
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
state = self.state_dict()
|
|
allowed = set(state.keys())
|
|
shapes = {k: tuple(v.shape) for k, v in state.items()}
|
|
|
|
tp_aware_params = {name for name, p in self.named_parameters() if hasattr(p, "weight_loader")}
|
|
|
|
# Expand allowed/tp_aware_params with stacked param source names.
|
|
# The model fuses several checkpoint projections into merged layers:
|
|
# QKV: q/k/v_proj → qkv_proj, q/k/v_proj_moe_gen → qkv_proj.gen_exp
|
|
# and remaps non-stacked weights like:
|
|
# {norm}_moe_gen.weight → {norm}.gen_weight (MoTRMSNorm layers)
|
|
# We expand allowed names so _filtered_weights does not drop them.
|
|
_stacked_expansions = [
|
|
# text QKV
|
|
(".qkv_proj", ".q_proj"),
|
|
(".qkv_proj", ".k_proj"),
|
|
(".qkv_proj", ".v_proj"),
|
|
# gen QKV
|
|
(".qkv_proj.gen_exp", ".q_proj_moe_gen"),
|
|
(".qkv_proj.gen_exp", ".k_proj_moe_gen"),
|
|
(".qkv_proj.gen_exp", ".v_proj_moe_gen"),
|
|
# gen o_proj (non-stacked, but still remapped)
|
|
(".o_proj.gen_exp", ".o_proj_moe_gen"),
|
|
# text FFN gate+up
|
|
(".mlp.gate_up_proj", ".mlp.gate_proj"),
|
|
(".mlp.gate_up_proj", ".mlp.up_proj"),
|
|
# gen FFN gate+up
|
|
(".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.gate_proj"),
|
|
(".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.up_proj"),
|
|
# MoTRMSNorm gen_weight ← checkpoint _moe_gen.weight
|
|
(".input_layernorm.gen_", ".input_layernorm_moe_gen."),
|
|
(".post_attention_layernorm.gen_", ".post_attention_layernorm_moe_gen."),
|
|
(".q_norm.gen_", ".q_norm_moe_gen."),
|
|
(".k_norm.gen_", ".k_norm_moe_gen."),
|
|
(".norm.gen_", ".norm_moe_gen."),
|
|
]
|
|
stacked_source_names: set[str] = set()
|
|
for name in list(allowed):
|
|
for target_suffix, source_suffix in _stacked_expansions:
|
|
if target_suffix in name:
|
|
stacked_source_names.add(name.replace(target_suffix, source_suffix))
|
|
allowed.update(stacked_source_names)
|
|
tp_aware_params.update(stacked_source_names)
|
|
|
|
def _normalize_name(name: str) -> str:
|
|
# Common wrappers/prefixes in checkpoints.
|
|
for pfx in ("module.", "model."):
|
|
if name.startswith(pfx):
|
|
name = name[len(pfx) :]
|
|
# Common component renames across repos.
|
|
if name.startswith("vae_model."):
|
|
name = "vae." + name[len("vae_model.") :]
|
|
# Bagel `ae.safetensors` commonly stores AE weights without a top-level prefix.
|
|
# Map them into this pipeline's `vae.*` namespace.
|
|
if name.startswith("encoder.") or name.startswith("decoder."):
|
|
name = "vae." + name
|
|
return name
|
|
|
|
def _iter_candidate_names(name: str) -> Iterable[str]:
|
|
"""Yield candidate parameter names in this pipeline for a checkpoint key.
|
|
|
|
The upstream Bagel repo typically stores Bagel-core layers (time_embedder,
|
|
latent_pos_embed, vae2llm, llm2vae, etc.) at the top-level of the model,
|
|
while this vllm-omni integration nests them under `self.bagel`.
|
|
"""
|
|
n = _normalize_name(name)
|
|
yield n
|
|
|
|
# Map Bagel core layers from top-level -> `bagel.*` namespace.
|
|
for pfx in ("time_embedder.", "latent_pos_embed.", "vae2llm.", "llm2vae."):
|
|
if n.startswith(pfx):
|
|
yield "bagel." + n
|
|
break
|
|
|
|
# Map connector and vit_pos_embed to `bagel.*`
|
|
for pfx in ("connector.", "vit_pos_embed."):
|
|
if n.startswith(pfx):
|
|
yield "bagel." + n
|
|
break
|
|
|
|
if n.startswith("vit_model."):
|
|
yield "bagel." + n # matches self.bagel.vit_model
|
|
elif n.startswith("vision_model."):
|
|
yield "bagel.vit_model." + n
|
|
elif n.startswith("model.vision_model."):
|
|
yield "bagel.vit_model." + n[len("model.") :]
|
|
|
|
def _filtered_weights():
|
|
total = 0
|
|
kept = 0
|
|
shape_mismatch = 0
|
|
for name, tensor in weights:
|
|
total += 1
|
|
picked = None
|
|
for cand in _iter_candidate_names(name):
|
|
if cand in allowed:
|
|
# Only accept if tensor shape matches target param/buffer shape.
|
|
if tuple(tensor.shape) == shapes.get(cand) or cand in tp_aware_params:
|
|
picked = cand
|
|
break
|
|
else:
|
|
if cand.endswith("bagel.latent_pos_embed.pos_embed") and tensor.ndim == 2:
|
|
npos, hdim = tensor.shape
|
|
side = isqrt(int(npos))
|
|
if side * side == int(npos) and hdim == int(self.bagel.hidden_size):
|
|
param = self.bagel.latent_pos_embed.pos_embed
|
|
# Resize in-place to keep the same Parameter object.
|
|
param.data = param.data.new_empty((npos, hdim))
|
|
# Update model bookkeeping so position-id generation matches.
|
|
self.bagel.max_latent_size = int(side)
|
|
if hasattr(self.bagel, "config"):
|
|
setattr(self.bagel.config, "max_latent_size", int(side))
|
|
if hasattr(self.bagel.latent_pos_embed, "max_num_patch_per_side"):
|
|
self.bagel.latent_pos_embed.max_num_patch_per_side = int(side)
|
|
shapes[cand] = (npos, hdim)
|
|
picked = cand
|
|
break
|
|
# Handle flattened patch embedding for SigLIP
|
|
if cand.endswith("embeddings.patch_embedding.weight") and tensor.ndim == 2:
|
|
# Checkpoint has (Hidden, C*P*P), model expects (Hidden, C, P, P)
|
|
if shapes.get(cand) is not None:
|
|
target_shape = shapes[cand]
|
|
if tensor.numel() == torch.prod(torch.tensor(target_shape)):
|
|
# Reshape tensor to match target
|
|
tensor = tensor.view(target_shape)
|
|
picked = cand
|
|
break
|
|
|
|
shape_mismatch += 1
|
|
# Keep this quiet; shape mismatches are expected for ignored modules.
|
|
if picked is not None:
|
|
kept += 1
|
|
yield picked, tensor
|
|
# else: ignore extra weights (e.g. connector/vision/und)
|
|
logger.info_once(
|
|
"BagelPipeline weight filter kept %d/%d tensors (shape mismatches seen: %d)",
|
|
kept,
|
|
total,
|
|
shape_mismatch,
|
|
)
|
|
|
|
loader = AutoWeightsLoader(self)
|
|
return loader.load_weights(_filtered_weights())
|