# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ BagelPipeline implementation for vLLM-Omni. """ from __future__ import annotations import json import os from collections.abc import Iterable from copy import deepcopy from dataclasses import dataclass from math import isqrt from typing import 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.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.worker.request_batch import DiffusionRequestBatch 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__) @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 @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", ] 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.""", ) # 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. first_prompt = req.prompts[0] prompt = first_prompt if isinstance(req.prompts[0], str) else (req.prompts[0].get("prompt") or "") max_hw = int(self.bagel.max_latent_size * self.bagel.latent_downsample) if req.sampling_params.height is None and req.sampling_params.width is None: height = width = max_hw else: height = int(req.sampling_params.height) if req.sampling_params.height is not None else max_hw width = int(req.sampling_params.width) if req.sampling_params.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) extra_args = getattr(req.sampling_params, "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(req.sampling_params.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 = req.sampling_params.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 req.sampling_params.kv_metadata and "ropes" in req.sampling_params.kv_metadata: gen_context["ropes"] = req.sampling_params.kv_metadata["ropes"] else: gen_context["ropes"] = [seq_len] if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata: image_shape = tuple(req.sampling_params.kv_metadata["image_shape"]) branch_kvs = getattr(req.sampling_params, "cfg_branch_past_key_values", None) or {} branch_metadata = getattr(req.sampling_params, "cfg_branch_kv_metadata", None) or {} active_branch = getattr(req.sampling_params, "cfg_active_branch", None) branch_roles = getattr(req.sampling_params, "cfg_branch_roles", None) or list(branch_kvs.keys()) cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) or branch_kvs.get("cfg_text") cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) or branch_metadata.get( "cfg_text" ) cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or branch_kvs.get("cfg_img") cfg_img_metadata = getattr(req.sampling_params, "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 req.sampling_params.kv_metadata and "ropes" in req.sampling_params.kv_metadata: cfg_img_context["ropes"] = req.sampling_params.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 image_shape = (resized_h, resized_w) logger.info(f"img2img: resized image to {resized_w}x{resized_h}") 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 req.sampling_params.seed is not None: torch.manual_seed(req.sampling_params.seed) if self.device.type == "cuda": torch.cuda.manual_seed(req.sampling_params.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, req.sampling_params.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) 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=req.sampling_params.return_trajectory_latents, scheduler=self.scheduler, scheduler_kwargs=self.scheduler_kwargs, ) img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], 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 req.sampling_params.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 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())