910 lines
38 KiB
Python
910 lines
38 KiB
Python
# 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 (vLLM): GLM-Image AR stage 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,
|
|
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.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
|
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, DiffusionPipelineProfilerMixin):
|
|
"""
|
|
GLM-Image Pipeline for text-to-image and image-to-image generation.
|
|
|
|
This pipeline integrates:
|
|
- AR stage (vLLM): 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 stage provides prior_token_ids (and optionally prior_token_image_ids)
|
|
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 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, quant_config=od_config.quantization_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)
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
|
)
|
|
|
|
# ==================== 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)}")
|
|
|
|
# ==================== 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)
|
|
# Move prior token ids to device (may come from CPU due to serialization)
|
|
if isinstance(condition_prior_token_id, torch.Tensor):
|
|
condition_prior_token_id = condition_prior_token_id.to(device=self.device)
|
|
|
|
# Encode condition image to latent space
|
|
# Use argmax (mode) for deterministic encoding of condition images
|
|
condition_latent = retrieve_latents(
|
|
self.vae.encode(condition_image.unsqueeze(0)), 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 "")
|
|
|
|
# NOTE: DiffusionEngine does an internal warmup "dummy run" during
|
|
# initialization. That request carries the fixed dummy request_id and does not carry
|
|
# Stage-0 (AR) outputs via req.extra. For GLM-Image, we allow that
|
|
# specific warmup request to proceed by synthesizing minimal prior
|
|
# tokens, while still raising a clear error for real requests.
|
|
is_dummy_warmup = req.is_dummy_run()
|
|
|
|
# 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 not isinstance(first_prompt, str):
|
|
img = first_prompt.get("additional_information", {}).get("preprocessed_image")
|
|
if img is not None:
|
|
preprocessed_images = [img]
|
|
|
|
# Priority: prompt dict (from ar2diffusion) > sampling_params
|
|
# ar2diffusion returns adjusted height/width that matches prior_token_ids
|
|
if not isinstance(first_prompt, str):
|
|
ar_height = first_prompt.get("height")
|
|
ar_width = first_prompt.get("width")
|
|
else:
|
|
ar_height = ar_width = None
|
|
|
|
img_height = req.sampling_params.height
|
|
img_width = req.sampling_params.width
|
|
|
|
# Warmup may include a dummy image for image-input-capable models.
|
|
# Treat that as t2i warmup to avoid requiring i2i-only KV-cache inputs.
|
|
is_image_edit = (preprocessed_images is not None) and (not is_dummy_warmup)
|
|
|
|
# Use prompt dict dimensions (from ar2diffusion) as priority, then sampling_params
|
|
height = (
|
|
ar_height or req.sampling_params.height or img_height or self.default_sample_size * self.vae_scale_factor
|
|
)
|
|
width = ar_width or 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
|
|
|
|
# Ensure dimensions are multiples of vae_scale_factor * patch_size
|
|
multiple_of = self.vae_scale_factor * self._patch_size
|
|
height = height // multiple_of * multiple_of
|
|
width = width // multiple_of * multiple_of
|
|
|
|
self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds)
|
|
|
|
batch_size = 1
|
|
do_classifier_free_guidance = guidance_scale > 1.0
|
|
|
|
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)
|
|
# First try sampling_params.extra_args (for direct API usage)
|
|
# Then try prompt["extra"] (for multistage ar2diffusion output)
|
|
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 None and isinstance(first_prompt, dict):
|
|
prompt_extra = first_prompt.get("extra", {})
|
|
external_prior_tokens = prompt_extra.get("prior_token_ids")
|
|
external_prior_image_ids = prompt_extra.get("prior_token_image_ids")
|
|
|
|
if external_prior_tokens is not None:
|
|
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)
|
|
if prior_token_id.dim() == 1:
|
|
prior_token_id = prior_token_id.unsqueeze(0)
|
|
|
|
# Validate that prior_token_id seq_len matches dimensions
|
|
prior_seq_len = prior_token_id.shape[1]
|
|
expected_seq_len = (height // self.vae_scale_factor // self._patch_size) * (
|
|
width // self.vae_scale_factor // self._patch_size
|
|
)
|
|
if prior_seq_len != expected_seq_len:
|
|
raise ValueError(
|
|
f"prior_token_ids seq_len ({prior_seq_len}) doesn't match dimensions "
|
|
f"({height}x{width}, expected seq_len={expected_seq_len}). "
|
|
f"This indicates a mismatch between AR output and Diffusion input. "
|
|
f"Please ensure ar2diffusion returns correct height/width."
|
|
)
|
|
|
|
prior_token_image_ids = None
|
|
if external_prior_image_ids is not None:
|
|
if isinstance(external_prior_image_ids, torch.Tensor):
|
|
prior_token_image_ids = [external_prior_image_ids]
|
|
else:
|
|
prior_token_image_ids = list(external_prior_image_ids)
|
|
|
|
normalized: list[torch.Tensor] = []
|
|
for item in prior_token_image_ids:
|
|
if isinstance(item, torch.Tensor):
|
|
tensor_item = item.to(device=self.device, dtype=torch.long)
|
|
else:
|
|
tensor_item = torch.tensor(item, dtype=torch.long, device=self.device)
|
|
if tensor_item.dim() == 1:
|
|
tensor_item = tensor_item.unsqueeze(0)
|
|
normalized.append(tensor_item)
|
|
prior_token_image_ids = normalized
|
|
elif is_dummy_warmup:
|
|
# Synthesize a minimal, shape-correct prior token sequence.
|
|
# The diffusion transformer expects prior_token_id shape:
|
|
# [B, (H_lat/p)*(W_lat/p)] where H_lat=H/vae_scale_factor.
|
|
h_lat = height // self.vae_scale_factor
|
|
w_lat = width // self.vae_scale_factor
|
|
seq_len = (h_lat * w_lat) // (self._patch_size**2)
|
|
prior_token_id = torch.zeros((1, seq_len), dtype=torch.long, device=self.device)
|
|
prior_token_image_ids = None
|
|
else:
|
|
raise ValueError(
|
|
"GLM-Image diffusion pipeline expects prior tokens from the vLLM AR stage. "
|
|
"Pass them via req.extra['prior_token_ids'] (and for i2i, req.extra['prior_token_image_ids'])."
|
|
)
|
|
|
|
# 2. Encode prompt for glyph embeddings
|
|
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 None:
|
|
raise ValueError(
|
|
"Image edit (i2i) requires req.extra['prior_token_image_ids'] from the AR stage to build KV cache."
|
|
)
|
|
if is_image_edit and prior_token_image_ids is not None:
|
|
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,
|
|
)
|
|
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
|
|
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
|
|
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]
|
|
|
|
return DiffusionOutput(
|
|
output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
"""Load transformer weights."""
|
|
transformer_weights = (
|
|
(name.replace("transformer.", "", 1), weight) for name, weight in weights if name.startswith("transformer.")
|
|
)
|
|
loaded = self.transformer.load_weights(transformer_weights)
|
|
return {f"transformer.{name}" for name in loaded}
|