# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ GlmImagePipeline implementation for vLLM-Omni. This pipeline implements GLM-Image text-to-image generation with: - AR stage: GlmImageForConditionalGeneration generates prior tokens - DiT stage: GlmImageTransformer2DModel performs diffusion denoising - VAE: AutoencoderKL decodes latents to images """ from __future__ import annotations import inspect import json import logging import os import re from collections.abc import Iterable from typing import cast import numpy as np import PIL.Image import torch from diffusers.image_processor import VaeImageProcessor from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import ( ByT5Tokenizer, GlmImageForConditionalGeneration, GlmImageProcessor, T5EncoderModel, ) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, 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.glm_image.glm_image_transformer import ( GlmImageKVCache, GlmImageTransformer2DModel, ) from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) logger = logging.getLogger(__name__) def get_glm_image_pre_process_func(od_config: OmniDiffusionConfig): """Get pre-processing function for GLM-Image pipeline. Pre-processes condition images before they are sent to the pipeline. This is called by DiffusionEngine before batching requests. """ model_name = od_config.model if os.path.exists(model_name): model_path = model_name else: model_path = download_weights_from_hf_specific(model_name, None, ["*"]) vae_config_path = os.path.join(model_path, "vae/config.json") with open(vae_config_path) as f: vae_config = json.load(f) block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512]) vae_scale_factor = 2 ** (len(block_out_channels) - 1) image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) # GLM-Image uses patch_size=2 for transformer patch_size = 2 def pre_process_func(request: OmniDiffusionRequest): """Pre-process condition images for Image Edit mode.""" for i, prompt in enumerate(request.prompts): multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None if isinstance(prompt, str): prompt = OmniTextPrompt(prompt=prompt) if "additional_information" not in prompt: prompt["additional_information"] = {} if raw_image is None: # Text-to-image mode, no preprocessing needed continue if not isinstance(raw_image, list): raw_image = [raw_image] images = [ PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im) for im in raw_image ] preprocessed = [] height, width = None, None for img in images: if isinstance(img, PIL.Image.Image): img_h, img_w = img.size[::-1] # PIL is (width, height) else: img_h, img_w = img.shape[:2] # Align to multiple of vae_scale_factor * patch_size multiple_of = vae_scale_factor * patch_size img_h = (img_h // multiple_of) * multiple_of img_w = (img_w // multiple_of) * multiple_of processed = image_processor.preprocess(img, height=img_h, width=img_w) preprocessed.append(processed) # Use first image dimensions as default if height is None: height, width = img_h, img_w # Store in request if isinstance(prompt, str): prompt = OmniTextPrompt(prompt=prompt, additional_information={}) elif "additional_information" not in prompt: prompt["additional_information"] = {} prompt["additional_information"]["preprocessed_image"] = processed # type: ignore prompt["additional_information"]["prompt_image"] = images # type: ignore request.prompts[i] = prompt if request.sampling_params.height is None: request.sampling_params.height = height if request.sampling_params.width is None: request.sampling_params.width = width return request return pre_process_func def get_glm_image_post_process_func(od_config: OmniDiffusionConfig): """Get post-processing function for GLM-Image pipeline.""" model_name = od_config.model if os.path.exists(model_name): model_path = model_name else: model_path = download_weights_from_hf_specific(model_name, None, ["*"]) vae_config_path = os.path.join(model_path, "vae/config.json") with open(vae_config_path) as f: vae_config = json.load(f) block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512]) vae_scale_factor = 2 ** (len(block_out_channels) - 1) image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) def post_process_func(images: torch.Tensor) -> list[PIL.Image.Image]: return image_processor.postprocess(images, output_type="pil") return post_process_func def calculate_shift( image_seq_len: int, base_seq_len: int = 256, base_shift: float = 0.25, max_shift: float = 0.75, ) -> float: """Calculate timestep shift based on image sequence length.""" m = (image_seq_len / base_seq_len) ** 0.5 mu = m * max_shift + base_shift return mu def retrieve_timesteps( scheduler, num_inference_steps: int | None = None, device: str | torch.device | None = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, **kwargs, ) -> tuple[torch.Tensor, int]: """ Calls the scheduler's `set_timesteps` method and retrieves timesteps. Handles custom timesteps and sigmas schedules. """ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) accepts_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) if timesteps is not None and sigmas is not None: # Both provided - check if scheduler supports both if not accepts_timesteps and not accepts_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep or sigma schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif timesteps is not None: if not accepts_timesteps: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif sigmas is not None: if not accepts_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" sigma schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps def retrieve_latents( encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample", ) -> torch.Tensor: """Extract latents from VAE encoder output.""" if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": return encoder_output.latent_dist.sample(generator) elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": return encoder_output.latent_dist.mode() elif hasattr(encoder_output, "latents"): return encoder_output.latents else: raise AttributeError("Could not access latents of provided encoder_output") class GlmImagePipeline(nn.Module): """ GLM-Image Pipeline for text-to-image and image-to-image generation. This pipeline integrates: - AR model (GlmImageForConditionalGeneration): Generates prior image tokens - Text encoder (T5EncoderModel): Encodes glyph/text embeddings - DiT model (GlmImageTransformer2DModel): Diffusion transformer - VAE (AutoencoderKL): Encodes/decodes images to/from latent space The pipeline flow: 1. AR generates prior_token_ids from text prompt 2. T5 encodes glyph text for text rendering 3. DiT performs iterative denoising conditioned on prior tokens 4. VAE decodes final latents to image """ def __init__( self, *, od_config: OmniDiffusionConfig, prefix: str = "", ): super().__init__() self.od_config = od_config self.parallel_config = od_config.parallel_config self.device = get_local_device() model = od_config.model local_files_only = os.path.exists(model) if local_files_only: model_path = model else: model_path = download_weights_from_hf_specific(model, od_config.revision, ["*"]) # Load scheduler self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model_path, subfolder="scheduler", local_files_only=True ) # Load AR model (vision_language_encoder) logger.info("Loading GlmImageForConditionalGeneration (AR model)...") self.vision_language_encoder = GlmImageForConditionalGeneration.from_pretrained( model_path, subfolder="vision_language_encoder", local_files_only=True, torch_dtype=torch.bfloat16, ).to(self.device) self.vision_language_encoder.eval() # Load processor for AR model self.processor = GlmImageProcessor.from_pretrained(model_path, subfolder="processor", local_files_only=True) # Load text encoder (T5 for glyph embeddings) logger.info("Loading T5EncoderModel (glyph encoder)...") self.text_encoder = T5EncoderModel.from_pretrained( model_path, subfolder="text_encoder", local_files_only=True, torch_dtype=torch.bfloat16, ).to(self.device) self.text_encoder.eval() # Load tokenizer for glyph encoding self.tokenizer = ByT5Tokenizer.from_pretrained(model_path, subfolder="tokenizer", local_files_only=True) # Load VAE logger.info("Loading AutoencoderKL (VAE)...") self.vae = AutoencoderKL.from_pretrained( model_path, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16 ).to(self.device) self.vae.eval() # Load transformer (DiT) logger.info("Loading GlmImageTransformer2DModel (DiT)...") self.transformer = GlmImageTransformer2DModel(od_config=od_config) # Weight sources for DiT loading self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=od_config.model, subfolder="transformer", revision=od_config.revision, prefix="transformer.", fall_back_to_pt=True, ) ] # Configure scale factors self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) self.default_sample_size = 128 # Get transformer config for patch size self._patch_size = getattr(self.transformer, "patch_size", 2) # ==================== Input Validation ==================== def check_inputs( self, prompt: str | list[str] | None, height: int | None, width: int | None, prompt_embeds: torch.Tensor | None = None, ) -> None: """Validate input arguments before generation.""" # Check dimension alignment multiple_of = self.vae_scale_factor * self._patch_size if height is not None and height % multiple_of != 0: logger.warning( f"`height` should be divisible by {multiple_of} but is {height}. " "Dimensions will be adjusted accordingly." ) if width is not None and width % multiple_of != 0: logger.warning( f"`width` should be divisible by {multiple_of} but is {width}. Dimensions will be adjusted accordingly." ) # Check prompt/prompt_embeds mutual exclusivity if prompt is not None and prompt_embeds is not None: raise ValueError( f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. " "Please provide only one of the two." ) if prompt is None and prompt_embeds is None: raise ValueError("Provide either `prompt` or `prompt_embeds`. Cannot leave both undefined.") # Check prompt type if prompt is not None and not isinstance(prompt, (str, list)): raise ValueError(f"`prompt` must be of type `str` or `list` but is {type(prompt)}") # ==================== AR Stage Methods ==================== @staticmethod def _compute_generation_params( image_grid_thw: torch.Tensor, is_text_to_image: bool, ) -> tuple[int, int, int, int]: """ Compute AR generation parameters from image grid. Args: image_grid_thw: Image grid tensor of shape [N, 3] where each row is [t, h, w] is_text_to_image: Whether this is text-to-image (vs image-to-image) Returns: Tuple of (max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w) """ grid_sizes = [] grid_hw = [] for i in range(image_grid_thw.shape[0]): t, h, w = image_grid_thw[i].tolist() grid_sizes.append(int(h * w)) grid_hw.append((int(h), int(w))) if not is_text_to_image: # Image-to-image: only generate target image tokens max_new_tokens = grid_sizes[-1] + 1 large_image_start_offset = 0 target_grid_h, target_grid_w = grid_hw[-1] else: # Text-to-image: generate both small preview and large target total_tokens = sum(grid_sizes) max_new_tokens = total_tokens + 1 large_image_start_offset = sum(grid_sizes[1:]) target_grid_h, target_grid_w = grid_hw[0] return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w @staticmethod def _extract_large_image_tokens( outputs: torch.Tensor, input_length: int, large_image_start_offset: int, large_image_tokens: int ) -> torch.Tensor: """Extract large image tokens from AR output.""" generated_tokens = outputs[0][input_length:] large_image_start = large_image_start_offset large_image_end = large_image_start + large_image_tokens return generated_tokens[large_image_start:large_image_end] @staticmethod def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> torch.Tensor: """Upsample token IDs by 2x using nearest neighbor interpolation.""" token_ids = token_ids.view(1, 1, token_h, token_w) token_ids = torch.nn.functional.interpolate(token_ids.float(), scale_factor=2, mode="nearest").to( dtype=torch.long ) token_ids = token_ids.view(1, -1) return token_ids @torch.inference_mode() def generate_prior_tokens( self, prompt: str, height: int, width: int, image: list[PIL.Image.Image] | None = None, factor: int = 32, ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: """ Generate prior tokens using the AR model. Args: prompt: Text prompt for generation height: Target image height width: Target image width image: Optional condition images for image-to-image factor: Token factor (default 32) Returns: Tuple of (prior_token_ids, prior_token_image_ids) prior_token_image_ids is a list of tensors, one per condition image """ device = self.vision_language_encoder.device height = (height // factor) * factor width = (width // factor) * factor is_text_to_image = image is None or len(image) == 0 # Build message content content = [] if image is not None: for img in image: content.append({"type": "image", "image": img}) content.append({"type": "text", "text": prompt}) messages = [{"role": "user", "content": content}] # Apply chat template - processor will handle target dimensions and build grid inputs = self.processor.apply_chat_template( messages, tokenize=True, target_h=height, target_w=width, return_dict=True, return_tensors="pt", ).to(device) image_grid_thw = inputs.get("image_grid_thw") # Compute generation parameters from the full grid max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params( image_grid_thw=image_grid_thw, is_text_to_image=is_text_to_image ) # Process condition images if provided # Use image_grid_thw[:-1] to exclude the target image grid (last entry) prior_token_image_ids = None if image is not None and image_grid_thw is not None and len(image_grid_thw) > 1: # Get features only for condition images (exclude target image grid) condition_grid = image_grid_thw[:-1] prior_token_image_embed = self.vision_language_encoder.get_image_features( inputs["pixel_values"], condition_grid ).pooler_output prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0) flat_prior_token_image_ids = self.vision_language_encoder.get_image_tokens( prior_token_image_embed, condition_grid ) # Split by image grid sizes and convert to list split_sizes = (condition_grid.prod(dim=-1)).tolist() prior_token_image_ids_list = torch.split(flat_prior_token_image_ids, split_sizes, dim=0) # Convert to list with upsampling prior_token_image_ids = [] for i, token_ids in enumerate(prior_token_image_ids_list): grid_t, grid_h, grid_w = condition_grid[i].tolist() token_ids = token_ids.view(1, -1) # Upsample 2x (from d32 to d64) token_ids_upsampled = self._upsample_token_ids(token_ids, grid_h, grid_w) prior_token_image_ids.append(token_ids_upsampled) # Generate with AR model outputs = self.vision_language_encoder.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=True, ) # Extract and upsample tokens large_image_tokens = token_h * token_w prior_token_ids_d32 = self._extract_large_image_tokens( outputs, inputs["input_ids"].shape[-1], large_image_offset, large_image_tokens ) prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w) return prior_token_ids, prior_token_image_ids # ==================== Text Encoding Methods ==================== def get_glyph_texts(self, prompt: str | list[str]) -> list[str]: """Extract text within quotes for glyph rendering.""" prompt = prompt[0] if isinstance(prompt, list) else prompt ocr_texts = ( re.findall(r"'([^']*)'", prompt) + re.findall(r"“([^“”]*)”", prompt) + re.findall(r'"([^"]*)"', prompt) + re.findall(r"「([^「」]*)」", prompt) ) return ocr_texts def _get_glyph_embeds( self, prompt: str | list[str], max_sequence_length: int = 2048, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> torch.Tensor: """Get glyph embeddings from T5 encoder for text rendering.""" device = device or self.device dtype = dtype or self.text_encoder.dtype glyph_texts = self.get_glyph_texts(prompt) input_ids = self.tokenizer( glyph_texts if len(glyph_texts) > 0 else [""], max_length=max_sequence_length, truncation=True, ).input_ids # Pad to even length input_ids = [[self.tokenizer.pad_token_id] * ((len(ids) + 1) % 2) + ids for ids in input_ids] max_length = max(len(ids) for ids in input_ids) attention_mask = torch.tensor( [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids], device=device, ) input_ids = torch.tensor( [ids + [self.tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids], device=device, ) outputs = self.text_encoder(input_ids, attention_mask=attention_mask) glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) return glyph_embeds.to(device=device, dtype=dtype) def encode_prompt( self, prompt: str | list[str], do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, prompt_embeds: torch.Tensor | None = None, device: torch.device | None = None, dtype: torch.dtype | None = None, max_sequence_length: int = 2048, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Encode prompt into glyph embeddings for text rendering.""" device = device or self.device prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: prompt_embeds = self._get_glyph_embeds(prompt, max_sequence_length, device, dtype) seq_len = prompt_embeds.size(1) prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) negative_prompt_embeds = None if do_classifier_free_guidance: negative_prompt = [""] * batch_size negative_prompt_embeds = self._get_glyph_embeds(negative_prompt, max_sequence_length, device, dtype) seq_len = negative_prompt_embeds.size(1) negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) return prompt_embeds, negative_prompt_embeds # ==================== Latent Preparation ==================== def prepare_latents( self, batch_size: int, num_channels_latents: int, height: int, width: int, dtype: torch.dtype, device: torch.device, generator: torch.Generator | None, latents: torch.Tensor | None = None, ) -> torch.Tensor: """Prepare random noise latents.""" if latents is not None: return latents.to(device) shape = ( batch_size, num_channels_latents, int(height) // self.vae_scale_factor, int(width) // self.vae_scale_factor, ) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError(f"Passed {len(generator)} generators but batch size is {batch_size}.") latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) return latents def diffuse( self, latents: torch.Tensor, prior_token_id: torch.Tensor, prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None, timesteps: torch.Tensor, target_size: torch.Tensor, crop_coords: torch.Tensor, guidance_scale: float, do_classifier_free_guidance: bool, kv_caches: GlmImageKVCache | None = None, ) -> torch.Tensor: """ Denoising loop for diffusion process with CFG-Parallel support. Args: latents: Initial noise latents prior_token_id: Prior tokens generated by AR model prompt_embeds: Encoded positive prompt embeddings (glyph embeddings) negative_prompt_embeds: Encoded negative prompt embeddings timesteps: Denoising timesteps target_size: Target image size tensor [[height, width]] crop_coords: Crop coordinates tensor guidance_scale: CFG scale do_classifier_free_guidance: Whether to apply CFG kv_caches: Optional KV cache for Image Edit mode Returns: Denoised latents ready for VAE decode """ # Prepare conditional/unconditional drop flags prior_token_drop_cond = torch.full_like(prior_token_id, False, dtype=torch.bool) prior_token_drop_uncond = torch.full_like(prior_token_id, True, dtype=torch.bool) transformer_dtype = self.transformer.dtype # Enable CFG-parallel: rank0 computes positive, rank1 computes negative cfg_parallel_ready = do_classifier_free_guidance and get_classifier_free_guidance_world_size() > 1 for i, t in enumerate(timesteps): latent_model_input = latents.to(transformer_dtype) timestep = t.expand(latents.shape[0]) - 1 if cfg_parallel_ready: cfg_group = get_cfg_group() cfg_rank = get_classifier_free_guidance_rank() if cfg_rank == 0: # Rank 0: Compute positive (conditional) prediction local_pred = self.transformer( hidden_states=latent_model_input, encoder_hidden_states=prompt_embeds, prior_token_id=prior_token_id, prior_token_drop=prior_token_drop_cond, timestep=timestep, target_size=target_size, crop_coords=crop_coords, kv_cache=kv_caches, return_dict=False, )[0].float() else: # Rank 1: Compute negative (unconditional) prediction local_pred = self.transformer( hidden_states=latent_model_input, encoder_hidden_states=negative_prompt_embeds, prior_token_id=prior_token_id, prior_token_drop=prior_token_drop_uncond, timestep=timestep, target_size=target_size, crop_coords=crop_coords, kv_cache=kv_caches, return_dict=False, )[0].float() # All-gather predictions from all ranks gathered = cfg_group.all_gather(local_pred, separate_tensors=True) if cfg_rank == 0: # Rank 0: Combine predictions and apply CFG noise_pred_cond = gathered[0] noise_pred_uncond = gathered[1] noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) # Scheduler step latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] # Broadcast updated latents to all ranks cfg_group.broadcast(latents, src=0) else: # Sequential CFG (single GPU or no CFG) # Conditional forward pass noise_pred_cond = self.transformer( hidden_states=latent_model_input, encoder_hidden_states=prompt_embeds, prior_token_id=prior_token_id, prior_token_drop=prior_token_drop_cond, timestep=timestep, target_size=target_size, crop_coords=crop_coords, kv_cache=kv_caches, return_dict=False, )[0].float() if do_classifier_free_guidance: # Unconditional forward pass noise_pred_uncond = self.transformer( hidden_states=latent_model_input, encoder_hidden_states=negative_prompt_embeds, prior_token_id=prior_token_id, prior_token_drop=prior_token_drop_uncond, timestep=timestep, target_size=target_size, crop_coords=crop_coords, kv_cache=kv_caches, return_dict=False, )[0].float() noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) else: noise_pred = noise_pred_cond # Scheduler step latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] return latents # ==================== Main Forward Pass ==================== def _prepare_condition_image_kv_cache( self, condition_images: list[torch.Tensor], prior_token_image_ids: list[torch.Tensor], prompt_embeds: torch.Tensor, generator: torch.Generator | None = None, ) -> GlmImageKVCache: """ Prepare KV cache by running condition images through transformer at timestep 0. This is used for Image Edit mode where we need to cache the condition image's KV states for cross-attention during denoising. Args: condition_images: List of preprocessed condition images prior_token_image_ids: Prior token IDs for each condition image from AR model prompt_embeds: Prompt embeddings (used to get dtype) generator: Optional random generator Returns: GlmImageKVCache with cached KV states from condition images """ kv_caches = self.transformer.create_kv_cache() kv_caches.set_mode("write") # Prepare VAE normalization parameters latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.latent_channels, 1, 1) .to(device=self.device, dtype=prompt_embeds.dtype) ) latents_std = ( torch.tensor(self.vae.config.latents_std) .view(1, self.vae.config.latent_channels, 1, 1) .to(device=self.device, dtype=prompt_embeds.dtype) ) # Process each condition image through transformer to populate KV cache for condition_image, condition_prior_token_id in zip(condition_images, prior_token_image_ids): condition_image = condition_image.to(device=self.device, dtype=prompt_embeds.dtype) # Encode condition image to latent space # Use argmax (mode) for deterministic encoding of condition images condition_latent = retrieve_latents( self.vae.encode(condition_image), generator=generator, sample_mode="argmax" ) condition_latent = (condition_latent - latents_mean) / latents_std # Run forward pass at timestep 0 to cache KV states # Empty encoder_hidden_states since we only want to cache image features _ = self.transformer( hidden_states=condition_latent, encoder_hidden_states=torch.zeros_like(prompt_embeds)[:1, :0, ...], prior_token_id=condition_prior_token_id, prior_token_drop=torch.full_like(condition_prior_token_id, False, dtype=torch.bool), timestep=torch.zeros((1,), device=self.device), target_size=torch.tensor([condition_image.shape[-2:]], device=self.device, dtype=prompt_embeds.dtype), crop_coords=torch.zeros((1, 2), device=self.device, dtype=prompt_embeds.dtype), kv_cache=kv_caches, return_dict=False, ) return kv_caches @torch.inference_mode() def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: """ Main generation forward pass. Args: req: OmniDiffusionRequest with generation parameters Returns: DiffusionOutput containing generated image """ 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.""", ) first_prompt = req.prompts[0] prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") # Get pre-computed prompt embeddings if provided if isinstance(first_prompt, str): prompt_embeds = None else: prompt_embeds = first_prompt.get("prompt_embeds") if not isinstance(prompt_embeds, torch.Tensor): prompt_embeds = None # Get condition images for Image Edit mode # Use pre-processed images from pre_process_func preprocessed_images = ( None if isinstance(first_prompt, str) else [first_prompt.get("additional_information", {}).get("preprocessed_image")] ) condition_images = ( None if isinstance(first_prompt, str) else first_prompt.get("additional_information", {}).get("prompt_image") ) img_height = req.sampling_params.height img_width = req.sampling_params.width is_image_edit = preprocessed_images is not None # Use image dimensions as default if available height = req.sampling_params.height or img_height or self.default_sample_size * self.vae_scale_factor width = req.sampling_params.width or img_width or self.default_sample_size * self.vae_scale_factor num_inference_steps = req.sampling_params.num_inference_steps or 50 guidance_scale = req.sampling_params.guidance_scale or 1.5 # 0. Validate inputs self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds) batch_size = 1 do_classifier_free_guidance = guidance_scale > 1.0 # Set seed if provided generator = None if req.sampling_params.seed is not None: generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) # 1. Get prior tokens - either from external source (multistage) or generate internally # Check if prior_token_ids are provided externally (from AR stage in multistage mode) external_prior_tokens = req.sampling_params.extra_args.get("prior_token_ids") external_prior_image_ids = req.sampling_params.extra_args.get("prior_token_image_ids") if external_prior_tokens is not None: # Multistage mode: use externally provided prior tokens from vLLM AR stage logger.info("Using externally provided prior tokens from AR stage...") prior_token_id = external_prior_tokens if isinstance(prior_token_id, list): prior_token_id = torch.tensor(prior_token_id, dtype=torch.long, device=self.device) elif isinstance(prior_token_id, torch.Tensor): prior_token_id = prior_token_id.to(device=self.device, dtype=torch.long) # Ensure shape is [1, num_tokens] for batch processing if prior_token_id.dim() == 1: prior_token_id = prior_token_id.unsqueeze(0) prior_token_image_ids = external_prior_image_ids else: # Single-stage mode: generate prior tokens with internal AR model logger.info("Generating prior tokens with AR model...") prior_token_id, prior_token_image_ids = self.generate_prior_tokens( prompt=prompt, image=condition_images, height=height, width=width, ) # 2. Encode prompt for glyph embeddings logger.info("Encoding prompt...") prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt, do_classifier_free_guidance=do_classifier_free_guidance, num_images_per_prompt=1, prompt_embeds=prompt_embeds, device=self.device, dtype=self.transformer.dtype, ) # 3. Prepare KV cache for Image Edit mode kv_caches = None if is_image_edit and prior_token_image_ids is not None: logger.info("Preparing KV cache for Image Edit mode...") kv_caches = self._prepare_condition_image_kv_cache( condition_images=preprocessed_images, prior_token_image_ids=prior_token_image_ids, prompt_embeds=prompt_embeds, generator=generator, ) # Switch to read mode for denoising kv_caches.set_mode("read") # 4. Prepare latents latent_channels = self.transformer.in_channels latents = self.prepare_latents( batch_size=batch_size, num_channels_latents=latent_channels, height=height, width=width, dtype=prompt_embeds.dtype, device=self.device, generator=generator, ) # 5. Prepare timesteps image_seq_len = ((height // self.vae_scale_factor) * (width // self.vae_scale_factor)) // (self._patch_size**2) timesteps_array = np.linspace(self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1)[:-1] timesteps_array = timesteps_array.astype(np.int64).astype(np.float32) sigmas = timesteps_array / self.scheduler.config.num_train_timesteps mu = calculate_shift( image_seq_len, self.scheduler.config.get("base_image_seq_len", 256), self.scheduler.config.get("base_shift", 0.25), self.scheduler.config.get("max_shift", 0.75), ) timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, self.device, timesteps_array.tolist(), sigmas.tolist(), mu=mu ) # 6. Prepare conditioning tensors target_size = torch.tensor([[height, width]], dtype=prompt_embeds.dtype, device=self.device) crop_coords = torch.zeros((1, 2), dtype=prompt_embeds.dtype, device=self.device) # 7. Denoising loop with CFG-parallel support logger.info(f"Starting denoising loop with {num_inference_steps} steps...") latents = self.diffuse( latents=latents, prior_token_id=prior_token_id, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, timesteps=timesteps, target_size=target_size, crop_coords=crop_coords, guidance_scale=guidance_scale, do_classifier_free_guidance=do_classifier_free_guidance, kv_caches=kv_caches, ) # 8. VAE decode logger.info("Decoding latents with VAE...") latents = latents.to(self.vae.dtype) latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.latent_channels, 1, 1) .to(latents.device, latents.dtype) ) latents_std = ( torch.tensor(self.vae.config.latents_std) .view(1, self.vae.config.latent_channels, 1, 1) .to(latents.device, latents.dtype) ) latents = latents * latents_std + latents_mean image = self.vae.decode(latents, return_dict=False, generator=generator)[0] # 9. Leave post-process to vllm-omni pipeline return DiffusionOutput(output=image) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load transformer weights.""" # Filter weights for transformer only transformer_weights = ( (name.replace("transformer.", ""), weight) for name, weight in weights if name.startswith("transformer.") ) return self.transformer.load_weights(transformer_weights)