1016 lines
41 KiB
Python
1016 lines
41 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: 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)
|