From c5862363126f1a3c0ee60c3d68d6ae14d55b8c0d Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Tue, 7 Jul 2026 00:02:18 +0000 Subject: [PATCH 01/11] [None][feat] Add GlmImage text-to-image pipeline Signed-off-by: Joseph Loftin --- docs/source/models/visual-generation.md | 4 + examples/visual_gen/README.md | 2 + .../configs/glm-image-fp8-1gpu.yaml | 29 + examples/visual_gen/models/glm_image.py | 82 ++ .../visual_gen/serve/configs/glm_image.yml | 5 + .../_torch/visual_gen/models/__init__.py | 2 + .../visual_gen/models/glm_image/__init__.py | 19 + .../models/glm_image/pipeline_glm_image.py | 1101 +++++++++++++++++ .../models/glm_image/transformer_glm_image.py | 859 +++++++++++++ .../_torch/visual_gen/pipeline_registry.py | 1 + .../test_lists/test-db/l0_b200.yml | 2 + .../visual_gen/test_glm_image_pipeline.py | 682 ++++++++++ .../visual_gen/test_glm_image_transformer.py | 260 ++++ 13 files changed, 3048 insertions(+) create mode 100644 examples/visual_gen/configs/glm-image-fp8-1gpu.yaml create mode 100644 examples/visual_gen/models/glm_image.py create mode 100644 examples/visual_gen/serve/configs/glm_image.yml create mode 100644 tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py create mode 100644 tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py create mode 100644 tests/unittest/_torch/visual_gen/test_glm_image_transformer.py diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index dcfd46dc1375..20342fdf73d8 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -47,6 +47,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode | `nvidia/Cosmos3-Edge` | Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native) | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | Text-to-Video | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | Text-to-Video | +| `zai-org/GLM-Image` | Text-to-Image | Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via `model_index.json`; LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The `AutoPipeline` registry selects the appropriate pipeline class automatically. @@ -67,6 +68,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **Qwen-Image-Edit-2511** | Yes | Yes | No | No | Yes | No | No | Yes | Yes | No | No | No | No | No | | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | | **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | Yes | No | No | No | No | +| **GlmImage** [^9] | Yes | Yes | No | No | No | No | No | No | No | Yes | No | No | No | No | [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable. @@ -80,6 +82,8 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^7]: `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` — a distilled version of Wan2.2-TI2V-5B with 3 denoising steps. CFG parallelism, TeaCache, and Cache-DiT are not applicable. +[^9]: GlmImage currently supports single-GPU text-to-image with BF16 parity vs `diffusers` (cosine >= 0.99 on the full transformer). FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints. Image-to-image conditioning, sequence/CFG parallelism, parallel VAE, and caching (TeaCache / Cache-DiT) are not yet supported. + ## Quick Start Here is a simple example to generate a video with Wan 2.1: diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index 9792f0e4d3a6..a6f755e2cab5 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -25,6 +25,7 @@ python models/cosmos3_ti2v.py --prompt "A robot arm picks fruit in a grocery sto python models/qwen_image.py python models/qwen_image_layered.py --image /path/to/image.png python models/qwen_image_edit.py --image /path/to/source.png --prompt "Make the image look like a watercolor painting" +python models/glm_image.py python models/hunyuan_t2v.py # With engine config (quant, parallelism, etc.) @@ -37,6 +38,7 @@ python models/cosmos3_ti2v.py --visual_gen_args configs/cosmos3-nano-1gpu.yaml - python models/qwen_image.py --visual_gen_args configs/qwen-image-fp8-1gpu.yaml python models/qwen_image_layered.py --visual_gen_args configs/qwen-image-layered-1gpu.yaml --image /path/to/image.png python models/qwen_image_edit.py --visual_gen_args configs/qwen-image-edit-2511-fp4-1gpu.yaml --image /path/to/source.png --prompt "Make the image look like a watercolor painting" +python models/glm_image.py --visual_gen_args configs/glm-image-fp8-1gpu.yaml python models/hunyuan_t2v.py --visual_gen_args configs/hunyuan-t2v-fp8-1gpu.yaml ``` diff --git a/examples/visual_gen/configs/glm-image-fp8-1gpu.yaml b/examples/visual_gen/configs/glm-image-fp8-1gpu.yaml new file mode 100644 index 000000000000..611039a23591 --- /dev/null +++ b/examples/visual_gen/configs/glm-image-fp8-1gpu.yaml @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# 1-GPU GlmImage text-to-image with FP8 dynamic quantization. +# Model: zai-org/GLM-Image +# Shared by offline examples (--visual_gen_args) and trtllm-serve. +# +# GlmImage constraints: single-GPU text-to-image only. Image-to-image +# conditioning, sequence/CFG parallelism, and caching are not yet supported. +quant_config: + quant_algo: FP8 + dynamic: true +parallel_config: + cfg_size: 1 + ulysses_size: 1 +cuda_graph_config: + enable: false diff --git a/examples/visual_gen/models/glm_image.py b/examples/visual_gen/models/glm_image.py new file mode 100644 index 000000000000..093548e24bfb --- /dev/null +++ b/examples/visual_gen/models/glm_image.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GlmImage text-to-image generation. + +Usage: + python glm_image.py + python glm_image.py --visual_gen_args ../configs/glm-image-fp8-1gpu.yaml +""" + +import argparse +from pathlib import Path + +from tensorrt_llm import VisualGen, VisualGenArgs + + +def _output_paths(output_path: str, num_images: int) -> str | list[str]: + if num_images == 1: + return output_path + + path = Path(output_path) + return [str(path.with_name(f"{path.stem}_{index}{path.suffix}")) for index in range(num_images)] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + default="zai-org/GLM-Image", + help="Hugging Face model id or local checkpoint path.", + ) + parser.add_argument( + "--visual_gen_args", + "--extra_visual_gen_options", + dest="visual_gen_args", + help="Optional VisualGenArgs YAML file.", + ) + parser.add_argument( + "--prompt", + default="A serene mountain lake at sunrise, watercolor style, highly detailed", + help="Text prompt for image generation.", + ) + parser.add_argument( + "--num_images_per_prompt", + type=int, + default=1, + help="Number of images to generate for the prompt.", + ) + parser.add_argument( + "--output_path", + default="glm_image_output.png", + help="Image output path. Multiple images append an index before the suffix.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.num_images_per_prompt < 1: + raise ValueError("--num_images_per_prompt must be >= 1") + extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None + visual_gen = VisualGen(model=args.model, args=extra_args) + params = visual_gen.default_params + params.num_images_per_prompt = args.num_images_per_prompt + output = visual_gen.generate(inputs=args.prompt, params=params) + saved = output.save(_output_paths(args.output_path, args.num_images_per_prompt)) + print(f"Saved image(s) to {saved}") + + +if __name__ == "__main__": + main() diff --git a/examples/visual_gen/serve/configs/glm_image.yml b/examples/visual_gen/serve/configs/glm_image.yml new file mode 100644 index 000000000000..138fd010c5a4 --- /dev/null +++ b/examples/visual_gen/serve/configs/glm_image.yml @@ -0,0 +1,5 @@ +# GlmImage text-to-image (single GPU). +# Image-to-image conditioning, sequence/CFG parallelism, and caching are not yet supported. +parallel_config: + cfg_size: 1 + ulysses_size: 1 diff --git a/tensorrt_llm/_torch/visual_gen/models/__init__.py b/tensorrt_llm/_torch/visual_gen/models/__init__.py index 4f5eb8b4ffa5..f23fe3d815fa 100644 --- a/tensorrt_llm/_torch/visual_gen/models/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/__init__.py @@ -35,6 +35,7 @@ from ..pipeline_registry import AutoPipeline, register_pipeline from .cosmos3 import Cosmos3OmniMoTPipeline from .flux import Flux2Pipeline, FluxPipeline +from .glm_image import GlmImagePipeline from .hunyuan_video1_5 import HunyuanVideo15Pipeline from .ltx2 import LTX2Pipeline # noqa: F401 from .qwen_image import QwenImageEditPlusPipeline, QwenImagePipeline @@ -46,6 +47,7 @@ "BasePipeline", "FluxPipeline", "Flux2Pipeline", + "GlmImagePipeline", "QwenImageEditPlusPipeline", "QwenImageLayeredPipeline", "QwenImagePipeline", diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py new file mode 100644 index 000000000000..16b659411d12 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .pipeline_glm_image import GlmImagePipeline +from .transformer_glm_image import GlmImageAttention, GlmImageTransformer2DModel + +__all__ = ["GlmImagePipeline", "GlmImageTransformer2DModel", "GlmImageAttention"] diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py new file mode 100644 index 000000000000..7f76dd93a9a6 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py @@ -0,0 +1,1101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import re +import time +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import PIL +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.glm_image.pipeline_glm_image import ( + AutoencoderKL, + FlowMatchEulerDiscreteScheduler, + T5EncoderModel, +) +from diffusers.utils.torch_utils import randn_tensor +from transformers import ByT5Tokenizer, GlmImageForConditionalGeneration, GlmImageProcessor + +from tensorrt_llm import logger +from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig +from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline + +from .transformer_glm_image import GlmImageTransformer2DModel + +# ------------------------------------------------------------------ +# HF Port +# ------------------------------------------------------------------ +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents + + +def retrieve_latents( + encoder_output: torch.Tensor, + generator: Optional[torch.Generator] = None, + sample_mode: str = "sample", +): + 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") + + +# Copied from diffusers.pipelines.cogview4.pipeline_cogview4.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + 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: + 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 and sigmas is 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 timesteps is None and 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" sigmas 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 calculate_shift( + image_seq_len, + base_seq_len: int = 256, + base_shift: float = 0.25, + max_shift: float = 0.75, +) -> float: + m = (image_seq_len / base_seq_len) ** 0.5 + mu = m * max_shift + base_shift + return mu + + +@register_pipeline( + "GlmImagePipeline", + hf_ids=["zai-org/GLM-Image"], + doc="GlmImage family (text-to-image).", +) +class GlmImagePipeline(BasePipeline): + # ------------------------------------------------------------------ + # HF Port + # ------------------------------------------------------------------ + @staticmethod + def _validate_and_normalize_images( + image: Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]], + batch_size: int, + ) -> List[List[PIL.Image.Image]]: + """ + Validate and normalize image inputs to List[List[PIL.Image]]. + + Rules: + - batch_size > 1: Only accepts List[List[PIL.Image]], each sublist must have equal length + - batch_size == 1: Accepts List[PIL.Image] for legacy compatibility (converted to [[img1, img2, ...]]) + - Other formats raise ValueError + + Args: + image: Input images in various formats + batch_size: Number of prompts in the batch + + Returns: + Normalized images as List[List[PIL.Image]], or None if no images provided + """ + if image is None or len(image) == 0: + return None + + first_element = image[0] + + if batch_size == 1: + # Legacy format: List[PIL.Image] -> [[img1, img2, ...]] + if not isinstance(first_element, (list, tuple)): + return [list(image)] + # Already in List[List[PIL.Image]] format + if len(image) != 1: + raise ValueError( + f"For batch_size=1 with List[List[PIL.Image]] format, expected 1 image list, got {len(image)}." + ) + return [list(image[0])] + + # batch_size > 1: must be List[List[PIL.Image]] + if not isinstance(first_element, (list, tuple)): + raise ValueError( + f"For batch_size > 1, images must be List[List[PIL.Image]] format. " + f"Got List[{type(first_element).__name__}] instead. " + f"Each prompt requires its own list of condition images." + ) + + if len(image) != batch_size: + raise ValueError( + f"Number of image lists ({len(image)}) must match batch size ({batch_size})." + ) + + # Validate homogeneous: all sublists must have same length + num_input_images_per_prompt = len(image[0]) + for idx, imgs in enumerate(image): + if len(imgs) != num_input_images_per_prompt: + raise ValueError( + f"All prompts must have the same number of condition images. " + f"Prompt 0 has {num_input_images_per_prompt} images, but prompt {idx} has {len(imgs)} images." + ) + + return [list(imgs) for imgs in image] + + def generate_prior_tokens( + self, + prompt: Union[str, List[str]], + height: int, + width: int, + image: Optional[List[List[PIL.Image.Image]]] = None, + device: Optional[torch.device] = None, + generator: Optional[torch.Generator] = None, + ): + """ + Generate prior tokens for the DiT model using the AR model. + + Args: + prompt: Single prompt or list of prompts + height: Target image height + width: Target image width + image: Normalized image input as List[List[PIL.Image]]. Should be pre-validated + using _validate_and_normalize_images() before calling this method. + device: Target device + generator: Random generator for reproducibility + + Returns: + Tuple of: + - prior_token_ids: Tensor of shape (batch_size, num_tokens) with upsampled prior tokens + - prior_token_image_ids_per_sample: List of tensors, one per sample. Each tensor contains + the upsampled prior token ids for all condition images in that sample. None for t2i. + - source_image_grid_thw_per_sample: List of tensors, one per sample. Each tensor has shape + (num_condition_images, 3) with upsampled grid info. None for t2i. + """ + device = device or self._execution_device + + # Normalize prompt to list format + prompt_list = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt_list) + + # Image is already normalized by _validate_and_normalize_images(): None or List[List[PIL.Image]] + is_text_to_image = image is None + # Build messages for each sample in the batch + all_messages = [] + for idx, p in enumerate(prompt_list): + content = [] + if not is_text_to_image: + for img in image[idx]: + content.append({"type": "image", "image": img}) + content.append({"type": "text", "text": p}) + all_messages.append([{"role": "user", "content": content}]) + # Process with the processor (supports batch with left padding) + inputs = self.processor.apply_chat_template( + all_messages, + tokenize=True, + padding=True if batch_size > 1 else False, + target_h=height, + target_w=width, + return_dict=True, + return_tensors="pt", + ).to(device) + + image_grid_thw = inputs.get("image_grid_thw") + images_per_sample = inputs.get("images_per_sample") + + # Determine number of condition images and grids per sample + num_condition_images = 0 if is_text_to_image else len(image[0]) + if images_per_sample is not None: + num_grids_per_sample = images_per_sample[0].item() + else: + # Fallback for batch_size=1: total grids is for single sample + num_grids_per_sample = image_grid_thw.shape[0] + + # Compute generation params (same for all samples in homogeneous batch) + first_sample_grids = image_grid_thw[:num_grids_per_sample] + max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params( + image_grid_thw=first_sample_grids, is_text_to_image=is_text_to_image + ) + + # Generate source image tokens (prior_token_image_ids) for i2i mode + prior_token_image_ids = None + source_image_grid_thw = None + if not is_text_to_image: + # Extract source grids by selecting condition image indices (skip target grids) + # Grid order from processor: [s0_cond1, s0_cond2, ..., s0_target, s1_cond1, s1_cond2, ..., s1_target, ...] + # We need indices: [0, 1, ..., num_condition_images-1, num_grids_per_sample, num_grids_per_sample+1, ...] + source_indices = [] + for sample_idx in range(batch_size): + base = sample_idx * num_grids_per_sample + source_indices.extend(range(base, base + num_condition_images)) + source_grids = image_grid_thw[source_indices] + + if len(source_grids) > 0: + prior_token_image_embed = self.vision_language_encoder.get_image_features( + inputs["pixel_values"], source_grids + ).pooler_output + prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0) + prior_token_image_ids_d32 = self.vision_language_encoder.get_image_tokens( + prior_token_image_embed, source_grids + ) + # Upsample each source image's prior tokens to match VAE/DiT resolution + split_sizes = source_grids.prod(dim=-1).tolist() + prior_ids_per_source = torch.split(prior_token_image_ids_d32, split_sizes) + upsampled_prior_ids = [] + for i, prior_ids in enumerate(prior_ids_per_source): + t, h, w = source_grids[i].tolist() + upsampled = self._upsample_token_ids(prior_ids, int(h), int(w)) + upsampled_prior_ids.append(upsampled.squeeze(0)) + prior_token_image_ids = torch.cat(upsampled_prior_ids, dim=0) + # Upsample grid dimensions for later splitting + upsampled_grids = source_grids.clone() + upsampled_grids[:, 1] = upsampled_grids[:, 1] * 2 + upsampled_grids[:, 2] = upsampled_grids[:, 2] * 2 + source_image_grid_thw = upsampled_grids + + # Generate with AR model + # Set torch random seed from generator for reproducibility + # (transformers generate() doesn't accept generator parameter) + if generator is not None: + seed = generator.initial_seed() + torch.manual_seed(seed) + if device is not None and device.type == "cuda": + torch.cuda.manual_seed(seed) + outputs = self.vision_language_encoder.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=True, + ) + + # Extract and upsample prior tokens for each sample + # For left-padded inputs, generated tokens start after the padded input sequence + all_prior_token_ids = [] + max_input_length = inputs["input_ids"].shape[-1] + for idx in range(batch_size): + # For left-padded sequences, generated tokens start at max_input_length + # (padding is on the left, so all sequences end at the same position) + prior_token_ids_d32 = self._extract_large_image_tokens( + outputs[idx : idx + 1], max_input_length, large_image_offset, token_h * token_w + ) + prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w) + all_prior_token_ids.append(prior_token_ids) + prior_token_ids = torch.cat(all_prior_token_ids, dim=0) + + # Split prior_token_image_ids and source_image_grid_thw into per-sample lists for easier consumption + prior_token_image_ids_per_sample = None + source_image_grid_thw_per_sample = None + if prior_token_image_ids is not None and source_image_grid_thw is not None: + # Split grids: each sample has num_condition_images grids + source_image_grid_thw_per_sample = list( + torch.split(source_image_grid_thw, num_condition_images) + ) + # Split prior_token_image_ids: tokens per sample may vary due to different image sizes + tokens_per_image = source_image_grid_thw.prod(dim=-1).tolist() + tokens_per_sample = [] + for i in range(batch_size): + start_idx = i * num_condition_images + end_idx = start_idx + num_condition_images + tokens_per_sample.append(sum(tokens_per_image[start_idx:end_idx])) + prior_token_image_ids_per_sample = list( + torch.split(prior_token_image_ids, tokens_per_sample) + ) + + return prior_token_ids, prior_token_image_ids_per_sample, source_image_grid_thw_per_sample + + def encode_prompt( + self, + prompt: Union[str, List[str]], + do_classifier_free_guidance: bool = True, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + max_sequence_length: int = 2048, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + Whether to use classifier free guidance or not. + num_images_per_prompt (`int`, *optional*, defaults to 1): + Number of images that should be generated per prompt. torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + max_sequence_length (`int`, defaults to `2048`): + Maximum sequence length in encoded prompt. Can be set to other values but may lead to poorer results. + """ + device = device or self._execution_device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._get_glyph_embeds(prompt, max_sequence_length, device, dtype) + + # Repeat embeddings for num_images_per_prompt + if num_images_per_prompt > 1: + prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + + # For GLM-Image, negative_prompt must be "" instead of None + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = "" + negative_prompt = ( + batch_size * [negative_prompt] + if isinstance(negative_prompt, str) + else negative_prompt + ) + negative_prompt_embeds = self._get_glyph_embeds( + negative_prompt, max_sequence_length, device, dtype + ) + + if num_images_per_prompt > 1: + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave( + num_images_per_prompt, dim=0 + ) + + return prompt_embeds, negative_prompt_embeds + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + 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"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + @staticmethod + def _compute_generation_params( + image_grid_thw, + is_text_to_image: bool, + ): + 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: + max_new_tokens = grid_sizes[-1] + 1 + large_image_start_offset = 0 + target_grid_h, target_grid_w = grid_hw[-1] + else: + 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: + 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: + 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 + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + def _get_glyph_embeds( + self, + prompt: Union[str, List[str]] = None, + max_sequence_length: int = 2048, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + """Get glyph embeddings for each prompt in the batch.""" + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + # get_glyph_texts now returns a list of lists (one per prompt) + all_glyph_texts = self.get_glyph_texts(prompt) + + all_glyph_embeds = [] + for glyph_texts in all_glyph_texts: + if len(glyph_texts) == 0: + glyph_texts = [""] + input_ids = self.tokenizer( + glyph_texts, + max_length=max_sequence_length, + truncation=True, + ).input_ids + input_ids = [ + [self.tokenizer.pad_token_id] * ((len(input_ids) + 1) % 2) + input_ids_ + for input_ids_ in input_ids + ] + max_length = max(len(input_ids_) for input_ids_ in input_ids) + attention_mask = torch.tensor( + [ + [1] * len(input_ids_) + [0] * (max_length - len(input_ids_)) + for input_ids_ in input_ids + ], + device=device, + ) + input_ids = torch.tensor( + [ + input_ids_ + [self.tokenizer.pad_token_id] * (max_length - len(input_ids_)) + for input_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) + all_glyph_embeds.append(glyph_embeds) + + # Pad to same sequence length and stack (use left padding to match transformers) + max_seq_len = max(emb.size(1) for emb in all_glyph_embeds) + padded_embeds = [] + for emb in all_glyph_embeds: + if emb.size(1) < max_seq_len: + pad = torch.zeros( + emb.size(0), + max_seq_len - emb.size(1), + emb.size(2), + device=device, + dtype=emb.dtype, + ) + emb = torch.cat([pad, emb], dim=1) # left padding + padded_embeds.append(emb) + + glyph_embeds = torch.cat(padded_embeds, dim=0) + return glyph_embeds.to(device=device, dtype=dtype) + + def get_glyph_texts(self, prompt): + """Extract glyph texts from prompt(s). Returns a list of lists for batch processing.""" + if isinstance(prompt, str): + prompt = [prompt] + all_ocr_texts = [] + for p in prompt: + ocr_texts = ( + re.findall(r"'([^']*)'", p) + + re.findall(r"\u201c([^\u201c\u201d]*)\u201d", p) + + re.findall(r'"([^"]*)"', p) + + re.findall(r"「([^「」]*)」", p) + ) + all_ocr_texts.append(ocr_texts) + return all_ocr_texts + + @property + def guidance_scale(self): + return self._guidance_scale + + # ------------------------------------------------------------------ + # TRT-LLM + # ------------------------------------------------------------------ + def __init__(self, pipeline_config: DiffusionPipelineConfig): + super().__init__(pipeline_config) + + def load_standard_components( + self, + checkpoint_dir: str, + device: torch.device, + skip_components: Optional[list] = None, + **kwargs, + ) -> None: + skip_components = skip_components or [] + + # Tokenizer (ByT5Tokenizer) + if PipelineComponent.TOKENIZER not in skip_components: + logger.info("Loading tokenizer (ByT5Tokenizer)...") + tokenizer_path = os.path.join(checkpoint_dir, PipelineComponent.TOKENIZER) + self.tokenizer = ByT5Tokenizer.from_pretrained( + tokenizer_path, torch_dtype=self.pipeline_config.torch_dtype + ) + + # Text Encoder (T5EncoderModel) + if PipelineComponent.TEXT_ENCODER not in skip_components: + logger.info("Loading text encoder (T5EncoderModel)...") + text_encoder_path = os.path.join(checkpoint_dir, PipelineComponent.TEXT_ENCODER) + self.text_encoder = T5EncoderModel.from_pretrained( + text_encoder_path, torch_dtype=self.pipeline_config.torch_dtype + ).to(device) + + # VAE (AutoencoderKL) + if PipelineComponent.VAE not in skip_components: + logger.info("Loading VAE (AutoencoderKL)...") + vae_path = os.path.join(checkpoint_dir, PipelineComponent.VAE) + self.vae = AutoencoderKL.from_pretrained( + vae_path, torch_dtype=self.pipeline_config.torch_dtype + ).to(device) + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + # Scheduler (FlowMatchEulerDiscreteScheduler) + if PipelineComponent.SCHEDULER not in skip_components: + logger.info("Loading Scheduler (FlowMatchEulerDiscreteScheduler)...") + scheduler_path = os.path.join(checkpoint_dir, PipelineComponent.SCHEDULER) + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + scheduler_path, torch_dtype=self.pipeline_config.torch_dtype + ) + + # Vision Language Encoder (GlmImageForConditionalGeneration) + if PipelineComponent.VISION_LANGUAGE_ENCODER not in skip_components: + logger.info("Loading processor (GlmImageProcessor)...") + processor_path = os.path.join(checkpoint_dir, "processor") + self.processor = GlmImageProcessor.from_pretrained(processor_path) + + logger.info("Loading vision language encoder (GlmImageForConditionalGeneration)...") + vision_language_encoder_path = os.path.join( + checkpoint_dir, PipelineComponent.VISION_LANGUAGE_ENCODER + ) + self.vision_language_encoder = GlmImageForConditionalGeneration.from_pretrained( + vision_language_encoder_path, torch_dtype=self.pipeline_config.torch_dtype + ).to(device) + + @property + def device(self): + if self.transformer is not None: + return next(self.transformer.parameters()).device + return torch.device("cuda:0") + + @property + def dtype(self): + return self.pipeline_config.torch_dtype + + @property + def default_generation_params(self) -> dict: + """Model-specific defaults for None fields in VisualGenParams.""" + return { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "guidance_scale": 1.5, + "max_sequence_length": 2048, + } + + @property + def default_warmup_resolutions(self) -> List[Tuple[int, int]]: + return [(1024, 1024)] + + @property + def default_warmup_num_frames(self) -> List[int]: + # Image model: a single "frame" per sample. + return [1] + + @property + def resolution_multiple_of(self) -> Tuple[int, int]: + patch_size = self.transformer.config.patch_size if self.transformer is not None else 2 + multiple = getattr(self, "vae_scale_factor", 16) * patch_size + return (multiple, multiple) + + def infer(self, req): + """Run inference from DiffusionRequest.""" + params = req.params + if getattr(params, "image", None) is not None: + raise NotImplementedError( + "image-to-image conditioning is not yet supported by the " + "TensorRT-LLM GlmImage pipeline; coming in a follow-up MR" + ) + generator = None + if params.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(params.seed) + return self.forward( + prompt=req.prompt, + height=params.height, + width=params.width, + num_inference_steps=params.num_inference_steps, + guidance_scale=params.guidance_scale, + generator=generator, + num_images_per_prompt=params.num_images_per_prompt, + max_sequence_length=params.max_sequence_length, + ) + + def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: + with torch.no_grad(): + self.forward( + prompt="warmup", + height=height, + width=width, + num_inference_steps=steps, + generator=torch.Generator(device=self.device).manual_seed(42), + ) + + @torch.inference_mode() + def forward( + self, + prompt: Optional[Union[str, List[str]]] = None, + image: Optional[ + Union[ + torch.Tensor, + PIL.Image.Image, + np.ndarray, + List[torch.Tensor], + List[PIL.Image.Image], + List[np.ndarray], + ] + ] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + guidance_scale: float = 1.5, + num_images_per_prompt: int = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + prior_token_ids: Optional[torch.Tensor] = None, + prior_token_image_ids: Optional[List[torch.Tensor]] = None, + source_image_grid_thw: Optional[List[torch.Tensor]] = None, + crops_coords_top_left: Tuple[int, int] = (0, 0), + attention_kwargs: Optional[Dict[str, Any]] = None, + max_sequence_length: int = 2048, + ) -> PipelineOutput: + """ + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. Must contain shape info in the format 'H + W' where H and W are token dimensions (d32). Example: "A beautiful sunset36 24" + generates a 1152x768 image. + image: Optional condition images for image-to-image generation. + height (`int`, *optional*): + The height in pixels. If not provided, derived from prompt shape info. + width (`int`, *optional*): + The width in pixels. If not provided, derived from prompt shape info. + num_inference_steps (`int`, *optional*, defaults to `50`): + The number of denoising steps for DiT. + guidance_scale (`float`, *optional*, defaults to `1.5`): + Guidance scale for classifier-free guidance. + num_images_per_prompt (`int`, *optional*, defaults to `1`): + The number of images to generate per prompt. + generator (`torch.Generator`, *optional*): + Random generator for reproducibility. + + Returns: + PipelineOutput with image tensor ``(B, H, W, C)`` dtype uint8. + """ + if image is not None: + raise NotImplementedError( + "image-to-image conditioning is not yet supported by the " + "TensorRT-LLM GlmImage pipeline; coming in a follow-up MR" + ) + + pipeline_start = time.time() + timer = CudaPhaseTimer() + timer.mark_pre_start() + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self.device + + # 1. Validate and normalize image format + normalized_image = self._validate_and_normalize_images(image, batch_size) + + # 2. Generate prior tokens (batch mode) + # Get a single generator for AR model (use first if list provided) + logger.info("Generating prior tokens...") + ar_generator = generator[0] if isinstance(generator, list) else generator + if prior_token_ids is None: + prior_token_ids, prior_token_image_ids_per_sample, source_image_grid_thw_per_sample = ( + self.generate_prior_tokens( + prompt=prompt, + image=normalized_image, + height=height, + width=width, + device=device, + generator=ar_generator, + ) + ) + else: + # User provided prior_token_ids directly (from generate_prior_tokens) + prior_token_image_ids_per_sample = prior_token_image_ids + source_image_grid_thw_per_sample = source_image_grid_thw + + # 3. Preprocess images for VAE encoding + preprocessed_images = None + if normalized_image is not None: + preprocessed_images = [] + for prompt_images in normalized_image: + prompt_preprocessed = [] + for img in prompt_images: + image_height, image_width = ( + img.size[::-1] if isinstance(img, PIL.Image.Image) else img.shape[:2] + ) + multiple_of = self.vae_scale_factor * self.transformer.config.patch_size + image_height = (image_height // multiple_of) * multiple_of + image_width = (image_width // multiple_of) * multiple_of + img = self.image_processor.preprocess( + img, height=image_height, width=image_width + ) + prompt_preprocessed.append(img) + height = height or image_height + width = width or image_width + preprocessed_images.append(prompt_preprocessed) + + # 4. Encode input prompt + logger.info("Encoding prompt...") + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + self.do_classifier_free_guidance, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + device=device, + dtype=self.dtype, + ) + + # 5. Prepare latents + latent_channels = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size=batch_size * num_images_per_prompt, + num_channels_latents=latent_channels, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + + if normalized_image is not None: + latents_mean = torch.tensor(self.vae.config.latents_mean).view( + 1, self.vae.config.latent_channels, 1, 1 + ) + latents_std = torch.tensor(self.vae.config.latents_std).view( + 1, self.vae.config.latent_channels, 1, 1 + ) + + latents_mean = latents_mean.to(device=device, dtype=prompt_embeds.dtype) + latents_std = latents_std.to(device=device, dtype=prompt_embeds.dtype) + + # Process each sample's condition images + for prompt_idx in range(batch_size): + prompt_images = preprocessed_images[prompt_idx] + prompt_prior_ids = prior_token_image_ids_per_sample[prompt_idx] + prompt_grid_thw = source_image_grid_thw_per_sample[prompt_idx] + + # Split this sample's prior_token_image_ids by each image's token count + split_sizes = prompt_grid_thw.prod(dim=-1).tolist() + prior_ids_per_image = torch.split(prompt_prior_ids, split_sizes) + # Process each condition image for this sample + for condition_image, condition_image_prior_token_id in zip( + prompt_images, prior_ids_per_image + ): + condition_image = condition_image.to(device=device, dtype=prompt_embeds.dtype) + condition_latent = retrieve_latents( + self.vae.encode(condition_image), generator=generator, sample_mode="argmax" + ) + condition_latent = (condition_latent - latents_mean) / latents_std + + _ = self.transformer( + hidden_states=condition_latent, + encoder_hidden_states=torch.zeros_like(prompt_embeds)[:1, :0, ...], + prior_token_id=condition_image_prior_token_id, + prior_token_drop=torch.full_like( + condition_image_prior_token_id, False, dtype=torch.bool + ), + timestep=torch.zeros((1,), device=device), + target_size=torch.tensor([condition_image.shape[-2:]], device=device), + crop_coords=torch.zeros((1, 2), device=device), + attention_kwargs=attention_kwargs, + ) + + # 6. Prepare additional timestep conditions + target_size = (height, width) + target_size = torch.tensor([target_size], dtype=prompt_embeds.dtype, device=device) + crops_coords_top_left = torch.tensor( + [crops_coords_top_left], dtype=prompt_embeds.dtype, device=device + ) + + target_size = target_size.repeat(batch_size * num_images_per_prompt, 1) + crops_coords_top_left = crops_coords_top_left.repeat(batch_size * num_images_per_prompt, 1) + + # Prepare timesteps + image_seq_len = ((height // self.vae_scale_factor) * (width // self.vae_scale_factor)) // ( + self.transformer.config.patch_size**2 + ) + timesteps = ( + np.linspace(self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1)[ + :-1 + ] + if timesteps is None + else np.array(timesteps) + ) + timesteps = timesteps.astype(np.int64).astype(np.float32) + sigmas = timesteps / self.scheduler.config.num_train_timesteps if sigmas is None else sigmas + 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, device, timesteps, sigmas, mu=mu + ) + self._num_timesteps = len(timesteps) + + # 7. Denoising loop + transformer_dtype = self.dtype + + # Repeat prior_token_ids for num_images_per_prompt + if num_images_per_prompt > 1: + prior_token_ids = prior_token_ids.repeat_interleave(num_images_per_prompt, dim=0) + prior_token_drop_cond = torch.full_like(prior_token_ids, False, dtype=torch.bool) + prior_token_drop_uncond = torch.full_like(prior_token_ids, True, dtype=torch.bool) + + # Batched CFG concatenates the positive/negative streams, so align their + # lengths and carry a text attention mask (None when no padding is needed). + text_attention_mask = neg_text_attention_mask = None + if self.do_classifier_free_guidance: + ( + prompt_embeds, + negative_prompt_embeds, + text_attention_mask, + neg_text_attention_mask, + ) = self._align_cfg_embeds(prompt_embeds, negative_prompt_embeds) + + def forward_fn( + latents, + extra_stream_latents, + step_index, + timestep, + encoder_hidden_states, + extra_tensors, + ): + """Forward function for GlmImage transformer.""" + return self.transformer( + hidden_states=latents.to(transformer_dtype), + encoder_hidden_states=encoder_hidden_states, + prior_token_id=extra_tensors["prior_token_id"], + prior_token_drop=extra_tensors["prior_token_drop"], + timestep=timestep - 1, + target_size=extra_tensors["target_size"], + crop_coords=extra_tensors["crop_coords"], + attention_mask=extra_tensors.get("attention_mask"), + attention_kwargs=attention_kwargs, + return_dict=False, + )[0].float() + + extra_cfg_tensors = { + "prior_token_id": (prior_token_ids, prior_token_ids), + "prior_token_drop": (prior_token_drop_cond, prior_token_drop_uncond), + "target_size": (target_size, target_size), + "crop_coords": (crops_coords_top_left, crops_coords_top_left), + } + if text_attention_mask is not None: + extra_cfg_tensors["attention_mask"] = (text_attention_mask, neg_text_attention_mask) + + timer.mark_denoise_start() + latents = self.denoise( + latents=latents, + scheduler=self.scheduler, + prompt_embeds=prompt_embeds, + guidance_scale=self.guidance_scale, + forward_fn=forward_fn, + timesteps=timesteps, + neg_prompt_embeds=( + negative_prompt_embeds if self.do_classifier_free_guidance else None + ), + extra_cfg_tensors=extra_cfg_tensors, + ) + timer.mark_post_start() + + # Decode + logger.info("Decoding...") + image = self.decode_latents(latents, lambda lat: self._decode_latents(lat, generator)) + + if self.rank == 0: + logger.info("Pipeline total: %.2fs", time.time() - pipeline_start) + + timer.mark_end() + return timer.fill(PipelineOutput(image=image)) + + def _decode_latents( + self, latents: torch.Tensor, generator: Optional[torch.Generator] = None + ) -> torch.Tensor: + 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] + + image = (image / 2 + 0.5).clamp(0, 1) + image = image.permute(0, 2, 3, 1) # (B, C, H, W) -> (B, H, W, C) + image = (image * 255).round().to(torch.uint8) + return image + + @staticmethod + def _align_cfg_embeds( + prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Left-pad positive/negative encoder streams to a common length for batched CFG. + + Returns the (padded) embeds and text attention masks, or ``None`` masks when both + streams already match and no padding is needed. + """ + pos_len = prompt_embeds.shape[1] + neg_len = negative_prompt_embeds.shape[1] + if pos_len == neg_len: + return prompt_embeds, negative_prompt_embeds, None, None + + max_len = max(pos_len, neg_len) + + def _left_pad(embeds: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + batch, seq_len, dim = embeds.shape + mask = torch.ones(batch, seq_len, device=embeds.device, dtype=torch.long) + if seq_len < max_len: + pad_len = max_len - seq_len + embeds = torch.cat([embeds.new_zeros(batch, pad_len, dim), embeds], dim=1) + mask = torch.cat([mask.new_zeros(batch, pad_len), mask], dim=1) + return embeds, mask + + prompt_embeds, pos_mask = _left_pad(prompt_embeds) + negative_prompt_embeds, neg_mask = _left_pad(negative_prompt_embeds) + return prompt_embeds, negative_prompt_embeds, pos_mask, neg_mask + + def load_weights(self, weights: dict) -> None: + """Load transformer weights.""" + if self.transformer is not None and hasattr(self.transformer, "load_weights"): + logger.info("Loading transformer weights...") + transformer_weights = weights.get("transformer", weights) + self.transformer.load_weights(transformer_weights) + logger.info("Transformer weights loaded successfully.") + + self._target_dtype = self.pipeline_config.torch_dtype + + if self.transformer is not None: + self.transformer.eval() + + def _init_transformer(self) -> None: + """Initialize GlmImage transformer with quantization support.""" + logger.info("Creating HunyuanVideo1.5 transformer with quantization support...") + self.transformer = GlmImageTransformer2DModel( + model_config=self.pipeline_config.model_configs["transformer"] + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py new file mode 100644 index 000000000000..eb996b1fa22d --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py @@ -0,0 +1,859 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Optional, Tuple + +import torch +import torch.nn.functional as F +from diffusers.models.embeddings import Timesteps +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.transformers.transformer_glm_image import GlmImageRotaryPosEmbed +from tqdm import tqdm + +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm._torch.visual_gen.quantization import DynamicLinearWeightLoader +from tensorrt_llm._torch.visual_gen.utils import SequenceSharder +from tensorrt_llm.models.modeling_utils import QuantConfig + + +def _aux_linear( + in_features: int, + out_features: int, + bias: bool = True, + model_config: Optional[DiffusionModelConfig] = None, +) -> Linear: + """Build a quant-aware Linear for an auxiliary (non-block) projection.""" + return Linear( + in_features, + out_features, + bias=bias, + dtype=model_config.torch_dtype if model_config else None, + quant_config=model_config.quant_config if model_config else None, + skip_create_weights_in_init=( + model_config.skip_create_weights_in_init if model_config else False + ), + force_dynamic_quantization=( + model_config.force_dynamic_quantization if model_config else False + ), + ) + + +class GlmImageGELU(torch.nn.Module): + def __init__( + self, + dim_in: int, + dim_out: int, + approximate: str = "none", + bias: bool = True, + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + self.proj = _aux_linear(dim_in, dim_out, bias=bias, model_config=model_config) + self.approximate = approximate + + def gelu(self, gate: torch.Tensor) -> torch.Tensor: + return F.gelu(gate, approximate=self.approximate) + + def forward(self, hidden_states): + hidden_states = self.proj(hidden_states) + hidden_states = self.gelu(hidden_states) + return hidden_states + + +class GlmImageLinearActivation(torch.nn.Module): + def __init__( + self, + dim_in: int, + dim_out: int, + bias: bool = True, + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + self.proj = _aux_linear(dim_in, dim_out, bias=bias, model_config=model_config) + self.activation = F.silu + + def forward(self, hidden_states): + hidden_states = self.proj(hidden_states) + return self.activation(hidden_states) + + +class GlmImageFeedForward(torch.nn.Module): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "geglu", + inner_dim=None, + bias: bool = True, + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + if inner_dim is None: + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + if activation_fn == "gelu": + act_fn = GlmImageGELU(dim, inner_dim, bias=bias, model_config=model_config) + elif activation_fn == "gelu-approximate": + act_fn = GlmImageGELU( + dim, inner_dim, approximate="tanh", bias=bias, model_config=model_config + ) + elif activation_fn == "linear-silu": + act_fn = GlmImageLinearActivation(dim, inner_dim, bias=bias, model_config=model_config) + else: + raise ValueError(f"Unsupported activation_fn={activation_fn} for GlmImageFeedForward") + + self.net = torch.nn.ModuleList([]) + self.net.append(act_fn) + self.net.append(torch.nn.Dropout(dropout)) + self.net.append(_aux_linear(inner_dim, dim_out, bias=bias, model_config=model_config)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +class GlmImageTimestepEmbedding(torch.nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + out_dim: int = None, + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + self.linear_1 = _aux_linear(in_channels, time_embed_dim, model_config=model_config) + self.act = torch.nn.SiLU() + time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim + self.linear_2 = _aux_linear(time_embed_dim, time_embed_dim_out, model_config=model_config) + + def forward(self, sample): + sample = self.linear_1(sample) + sample = self.act(sample) + sample = self.linear_2(sample) + return sample + + +class GlmImagePixArtAlphaTextProjection(torch.nn.Module): + def __init__( + self, + in_features, + hidden_size, + out_features=None, + act_fn="gelu_tanh", + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + if out_features is None: + out_features = hidden_size + self.linear_1 = _aux_linear(in_features, hidden_size, bias=True, model_config=model_config) + if act_fn == "gelu_tanh": + self.act_1 = torch.nn.GELU(approximate="tanh") + elif act_fn == "silu": + self.act_1 = torch.nn.SiLU() + else: + raise ValueError(f"Unknown activation function: {act_fn}") + self.linear_2 = _aux_linear(hidden_size, out_features, bias=True, model_config=model_config) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +class GlmImageCombinedTimestepSizeEmbeddings(torch.nn.Module): + def __init__( + self, + embedding_dim: int, + condition_dim: int, + pooled_projection_dim: int, + timesteps_dim: int = 256, + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + + self.time_proj = Timesteps( + num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.condition_proj = Timesteps( + num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0 + ) + self.timestep_embedder = GlmImageTimestepEmbedding( + in_channels=timesteps_dim, time_embed_dim=embedding_dim, model_config=model_config + ) + self.condition_embedder = GlmImagePixArtAlphaTextProjection( + pooled_projection_dim, embedding_dim, act_fn="silu", model_config=model_config + ) + + def forward( + self, + timestep: torch.Tensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + hidden_dtype: torch.dtype, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + + crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(crop_coords.size(0), -1) + target_size_proj = self.condition_proj(target_size.flatten()).view(target_size.size(0), -1) + + # (B, 2 * condition_dim) + condition_proj = torch.cat([crop_coords_proj, target_size_proj], dim=1) + + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + condition_emb = self.condition_embedder(condition_proj.to(dtype=hidden_dtype)) + + conditioning = timesteps_emb + condition_emb + conditioning = F.silu(conditioning) + + return conditioning + + +class GlmImageImageProjector(torch.nn.Module): + def __init__( + self, + in_channels: int = 16, + hidden_size: int = 2560, + patch_size: int = 2, + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + self.patch_size = patch_size + + self.proj = _aux_linear(in_channels * patch_size**2, hidden_size, model_config=model_config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, channel, height, width = hidden_states.shape + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size + + hidden_states = hidden_states.reshape( + batch_size, + channel, + post_patch_height, + self.patch_size, + post_patch_width, + self.patch_size, + ) + hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + hidden_states = self.proj(hidden_states) + + return hidden_states + + +class GlmImageAdaLayerNormZero(torch.nn.Module): + def __init__( + self, + embedding_dim: int, + dim: int, + model_config: Optional[DiffusionModelConfig] = None, + ) -> None: + super().__init__() + + self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm_context = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.linear = _aux_linear(embedding_dim, 12 * dim, bias=True, model_config=model_config) + + def forward( + self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + dtype = hidden_states.dtype + norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) + norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) + + emb = self.linear(temb) + ( + shift_msa, + c_shift_msa, + scale_msa, + c_scale_msa, + gate_msa, + c_gate_msa, + shift_mlp, + c_shift_mlp, + scale_mlp, + c_scale_mlp, + gate_mlp, + c_gate_mlp, + ) = emb.chunk(12, dim=1) + + hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + encoder_hidden_states = norm_encoder_hidden_states * ( + 1 + c_scale_msa.unsqueeze(1) + ) + c_shift_msa.unsqueeze(1) + + return ( + hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) + + +class GlmImageAdaLayerNormContinuous(torch.nn.Module): + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + bias: bool = True, + norm_type: str = "layer_norm", + model_config: Optional[DiffusionModelConfig] = None, + ): + super().__init__() + self.linear = _aux_linear( + conditioning_embedding_dim, embedding_dim * 2, bias=bias, model_config=model_config + ) + if norm_type == "layer_norm": + self.norm = torch.nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = RMSNorm(hidden_size=embedding_dim, eps=eps, has_weights=elementwise_affine) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: + emb = self.linear(conditioning_embedding.to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class GlmImageAttention(Attention): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + eps: float = 1e-6, + dtype: Optional[torch.dtype] = None, + config: Optional[DiffusionModelConfig] = None, + layer_idx: int = 0, + ): + config = config or DiffusionModelConfig() + super().__init__( + num_attention_heads=num_attention_heads, + head_dim=attention_head_dim, + hidden_size=dim, + config=config, + qk_norm_mode="per_head", + qkv_mode=QKVMode.FUSE_QKV, + qk_norm=True, + ) + + self.heads = num_attention_heads + self.head_dim = attention_head_dim + + self.add_q_proj = Linear( + dim, + dim, + bias=True, + dtype=dtype, + mapping=self.mapping, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + ) + self.add_k_proj = Linear( + dim, + dim, + bias=True, + dtype=dtype, + mapping=self.mapping, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + ) + self.add_v_proj = Linear( + dim, + dim, + bias=True, + dtype=dtype, + mapping=self.mapping, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + ) + + # QK-norms, applied per-head on the head_dim. + self.norm_added_q = RMSNorm( + hidden_size=attention_head_dim, eps=eps, dtype=dtype, has_weights=True + ) + self.norm_added_k = RMSNorm( + hidden_size=attention_head_dim, eps=eps, dtype=dtype, has_weights=True + ) + + self.to_out = torch.nn.ModuleList( + [ + Linear( + dim, + dim, + bias=True, + dtype=dtype, + mapping=self.mapping, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + ), + torch.nn.Dropout(0.0), + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ): + batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape + batch_size, image_seq_length, embed_dim = hidden_states.shape + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + # QKV Proj + query, key, value = self.get_qkv(hidden_states) + query = query.unflatten(2, (self.heads, -1)) + key = key.unflatten(2, (self.heads, -1)) + value = value.unflatten(2, (self.heads, -1)) + + # 2. QK normalization + if self.qk_norm: + query, key = self.apply_qk_norm(query, key) + + # 3. Rotational positional embeddings applied to latent stream + if image_rotary_emb is not None: + from diffusers.models.embeddings import apply_rotary_emb + + query[:, text_seq_length:, :, :] = apply_rotary_emb( + query[:, text_seq_length:, :, :], + image_rotary_emb, + sequence_dim=1, + use_real_unbind_dim=-2, + ) + key[:, text_seq_length:, :, :] = apply_rotary_emb( + key[:, text_seq_length:, :, :], + image_rotary_emb, + sequence_dim=1, + use_real_unbind_dim=-2, + ) + + # 4. Attention + if attention_mask is not None: + text_attn_mask = attention_mask + assert text_attn_mask.dim() == 2, ( + "the shape of text_attn_mask should be (batch_size, text_seq_length)" + ) + text_attn_mask = text_attn_mask.float().to(query.device) + mix_attn_mask = torch.ones( + (batch_size, text_seq_length + image_seq_length), device=query.device + ) + mix_attn_mask[:, :text_seq_length] = text_attn_mask + mix_attn_mask = mix_attn_mask.unsqueeze(2) + attn_mask_matrix = mix_attn_mask @ mix_attn_mask.transpose(1, 2) + attention_mask = (attn_mask_matrix > 0).unsqueeze(1).to(query.dtype) + + hidden_states = self._attn_impl(query, key, value, key_padding_mask=attention_mask) + + hidden_states = hidden_states.to(query.dtype) + + # 5. Output projection + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + + encoder_hidden_states, hidden_states = hidden_states.split( + [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1 + ) + + return hidden_states, encoder_hidden_states + + +class GlmImageTransformerBlock(torch.nn.Module): + def __init__( + self, + dim: int = 2560, + num_attention_heads: int = 64, + attention_head_dim: int = 40, + time_embed_dim: int = 512, + eps: float = 1e-6, + dtype: Optional[torch.dtype] = None, + config: Optional[DiffusionModelConfig] = None, + layer_idx=0, + ): + super().__init__() + + # 1. Attention + self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim, model_config=config) + self.attn1 = GlmImageAttention( + dim, num_attention_heads, attention_head_dim, eps, dtype, config, layer_idx + ) + + # 2. Feedforward + self.norm2 = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm2_context = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.ff = GlmImageFeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate", model_config=config + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] + | list[tuple[torch.Tensor, torch.Tensor]] + | None = None, + attention_mask: dict[str, torch.Tensor] | None = None, + attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # 1. Timestep conditioning + ( + norm_hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + norm_encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) = self.norm1(hidden_states, encoder_hidden_states, temb) + + # 2. Attention + attention_kwargs = attention_kwargs or {} + + attn_hidden_states, attn_encoder_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + **attention_kwargs, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa.unsqueeze(1) + encoder_hidden_states = ( + encoder_hidden_states + attn_encoder_hidden_states * c_gate_msa.unsqueeze(1) + ) + + # 3. Feedforward + norm_hidden_states = self.norm2(hidden_states) * ( + 1 + scale_mlp.unsqueeze(1) + ) + shift_mlp.unsqueeze(1) + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) * ( + 1 + c_scale_mlp.unsqueeze(1) + ) + c_shift_mlp.unsqueeze(1) + + ff_output = self.ff(norm_hidden_states) + ff_output_context = self.ff(norm_encoder_hidden_states) + hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) + + return hidden_states, encoder_hidden_states + + +class GlmImageTransformer2DModel(BaseDiffusionModel): + def __init__(self, model_config: DiffusionModelConfig): + super().__init__(model_config) + + vgm = model_config.visual_gen_mapping + num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 32) + self.sharder = SequenceSharder.from_vgm(vgm, num_attention_heads=num_heads) + + pretrained_config = model_config.pretrained_config + + dtype = model_config.torch_dtype + quant_config = model_config.quant_config + skip_create_weights = model_config.skip_create_weights_in_init + force_dynamic_quant = model_config.force_dynamic_quantization + + attention_head_dim = getattr(pretrained_config, "attention_head_dim", 128) + condition_dim = getattr(pretrained_config, "condition_dim", 256) + in_channels = getattr(pretrained_config, "in_channels", 16) + num_attention_heads = getattr(pretrained_config, "num_attention_heads", 32) + num_layers = getattr(pretrained_config, "num_layers", 30) + out_channels = getattr(pretrained_config, "out_channels", 16) + patch_size = getattr(pretrained_config, "patch_size", 2) + prior_vq_quantizer_codebook_size = getattr( + pretrained_config, "prior_vq_quantizer_codebook_size", 16384 + ) + text_embed_dim = getattr(pretrained_config, "text_embed_dim", 1472) + time_embed_dim = getattr(pretrained_config, "time_embed_dim", 512) + pooled_projection_dim = 2 * 2 * condition_dim + inner_dim = num_attention_heads * attention_head_dim + + self.config = type( + "Config", + (), + { + "attention_head_dim": attention_head_dim, + "condition_dim": condition_dim, + "in_channels": in_channels, + "num_attention_heads": num_attention_heads, + "num_layers": num_layers, + "out_channels": out_channels, + "patch_size": patch_size, + "prior_vq_quantizer_codebook_size": prior_vq_quantizer_codebook_size, + "text_embed_dim": text_embed_dim, + "time_embed_dim": time_embed_dim, + "pooled_projection_dim": pooled_projection_dim, + "inner_dim": inner_dim, + "dtype": dtype, + "quant_config": quant_config, + "skip_create_weights": skip_create_weights, + "force_dynamic_quant": force_dynamic_quant, + }, + ) + + # 1. RoPE + self.rope = GlmImageRotaryPosEmbed(attention_head_dim, patch_size, theta=10000.0) + + # 2. Patch & Text-timestep embedding + self.image_projector = GlmImageImageProjector( + in_channels, inner_dim, patch_size, model_config=model_config + ) + self.glyph_projector = GlmImageFeedForward( + text_embed_dim, + inner_dim, + inner_dim=inner_dim, + activation_fn="gelu", + model_config=model_config, + ) + self.prior_token_embedding = Embedding(prior_vq_quantizer_codebook_size, inner_dim) + self.prior_projector = GlmImageFeedForward( + inner_dim, + inner_dim, + inner_dim=inner_dim, + activation_fn="linear-silu", + model_config=model_config, + ) + + self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( + embedding_dim=time_embed_dim, + condition_dim=condition_dim, + pooled_projection_dim=pooled_projection_dim, + timesteps_dim=time_embed_dim, + model_config=model_config, + ) + + # 3. Transformer blocks + self.transformer_blocks = torch.nn.ModuleList( + [ + GlmImageTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + time_embed_dim, + config=model_config, + dtype=dtype, + layer_idx=i, + ) + for i in range(num_layers) + ] + ) + + # 4. Output projection + self.norm_out = GlmImageAdaLayerNormContinuous( + inner_dim, time_embed_dim, elementwise_affine=False, model_config=model_config + ) + self.proj_out = _aux_linear( + inner_dim, patch_size * patch_size * out_channels, model_config=model_config + ) + + self.gradient_checkpointing = False + + self.apply_quant_config_exclude_modules() + self.__post_init__() + + def apply_quant_config_exclude_modules(self) -> None: + """Opt excluded Linears out of quantization (mirrors the Wan transformer).""" + quant_config = self.model_config.quant_config + if quant_config is None or quant_config.exclude_modules is None: + return + no_quant_config = QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) + for name, module in self.named_modules(): + if ( + isinstance(module, Linear) + and getattr(module, "quant_config", None) is not None + and quant_config.is_module_excluded_from_quantization(name) + ): + module.quant_config = no_quant_config + module._weights_created = False + module.create_weights() + + def __post_init__(self) -> None: + for _, module in self.named_modules(): + if callable(getattr(module, "create_weights", None)): + module.create_weights() + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + prior_token_id: torch.Tensor, + prior_token_drop: torch.Tensor, + timestep: torch.LongTensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] + | list[tuple[torch.Tensor, torch.Tensor]] + | None = None, + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: + batch_size, num_channels, height, width = hidden_states.shape + + # 1. RoPE + if image_rotary_emb is None: + image_rotary_emb = self.rope(hidden_states) + + # 2. Patch & Timestep embeddings + p = self.config.patch_size + post_patch_height = height // p + post_patch_width = width // p + + hidden_states = self.image_projector(hidden_states) + encoder_hidden_states = self.glyph_projector(encoder_hidden_states) + prior_embedding = self.prior_token_embedding(prior_token_id) + prior_embedding[prior_token_drop] *= 0.0 + prior_hidden_states = self.prior_projector(prior_embedding) + + hidden_states = hidden_states + prior_hidden_states + + temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) + + # 3. Transformer blocks + for idx, block in enumerate(self.transformer_blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + attention_mask, + attention_kwargs, + ) + else: + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + attention_mask, + attention_kwargs, + ) + + # 4. Output norm & projection + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, post_patch_height, post_patch_width, -1, p, p + ) + + # Rearrange tensor from (B, H_p, W_p, C, p, p) to (B, C, H_p * p, W_p * p) + output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + def load_weights(self, weights: dict) -> None: + """Load weights into the transformer. + + Args: + weights: Dictionary of parameter name -> tensor + """ + + # Map fused QKV layer names to original HF checkpoint names + # HF checkpoint has separate to_q, to_k, to_v / add_q_proj, add_k_proj, add_v_proj + # We fuse them into qkv_proj / add_qkv_proj for better performance + params_map = { + "add_qkv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], + "qkv_proj": ["to_q", "to_k", "to_v"], + } + + loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map) + + # Track prefixes of wrapper projectors whose sub-Linears are loaded + # by the parent's load_weights — the generic Linear loader must skip + # them (their FUSED weight modes would look for nonexistent checkpoint + # keys via params_map and error). + managed_prefixes = set() + + for name, module in tqdm(self.named_modules(), desc="Loading weights"): + if any(name.startswith(p) for p in managed_prefixes): + continue + + # Create weights for modules with skip_create_weights_in_init=True + # This must be done before loading weights (following Wan pattern) + if callable(getattr(module, "create_weights", None)): + module.create_weights() + + if len(module._parameters) == 0: + continue + + if isinstance(module, Embedding): + # Embedding subclasses Linear but must never be weight-quantized + weight_dicts = loader.get_linear_weights(module, name, weights) + if weight_dicts: + module.load_weights(weight_dicts) + elif isinstance(module, Linear): + weight_dicts = loader.get_linear_weights(module, name, weights) + + if weight_dicts: + loader.load_linear_weights(module, name, weight_dicts) + else: + module_weights = loader.filter_weights(name, weights) + for param_name, param in module._parameters.items(): + if param is not None and param_name in module_weights: + param.data.copy_( + module_weights[param_name].to(self.model_config.torch_dtype) + ) + + def post_load_weights(self) -> None: + """Call post_load_weights on all Linear modules and normalize dtypes.""" + compute_dtype = self.model_config.torch_dtype + quantized_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) + for _, module in self.named_modules(): + if isinstance(module, Linear): + module.post_load_weights() + weight = getattr(module, "weight", None) + if ( + weight is not None + and weight.is_floating_point() + and weight.dtype not in quantized_dtypes + ): + module.to(compute_dtype) + continue + for param in module._parameters.values(): + if param is not None and param.is_floating_point(): + param.data = param.data.to(compute_dtype) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index fb497d86dd10..751f7aec8da0 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -58,6 +58,7 @@ class PipelineComponent(str, Enum): IMAGE_PROCESSOR = "image_processor" SOUND_TOKENIZER = "sound_tokenizer" GUIDER = "guider" + VISION_LANGUAGE_ENCODER = "vision_language_encoder" @dataclass diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 0ca143964ea6..c0e6548c4bda 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -273,6 +273,8 @@ l0_b200: - examples/visual_gen/test_visual_gen_hunyuan.py::test_hunyuan_t2v_example TIMEOUT (30) - unittest/_torch/visual_gen/test_cosmos3_edge.py - unittest/_torch/visual_gen/test_cosmos3_example_prompts.py + - unittest/_torch/visual_gen/test_glm_image_transformer.py + - unittest/_torch/visual_gen/test_glm_image_pipeline.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py new file mode 100644 index 000000000000..0afc2b1f199a --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py @@ -0,0 +1,682 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import os +from pathlib import Path +from typing import Optional + +from tensorrt_llm._torch.modules.linear import Linear + +os.environ["TLLM_DISABLE_MPI"] = "1" + +import numpy as np +import pytest +import torch +import torch.nn.functional as F +from diffusers import DiffusionPipeline + +from tensorrt_llm._torch.visual_gen.models.glm_image import GlmImagePipeline +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader +from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + yield + os.environ.pop("TLLM_DISABLE_MPI", None) + + +# ============================================================================ +# Test constants +# ============================================================================ + +# Canonical HuggingFace Hub ID for GlmImage. +GLM_IMAGE_HF_ID = "zai-org/GLM-Image" + + +def _llm_models_root() -> Optional[str]: + """Return the LLM_MODELS_ROOT path if it resolves to an existing directory.""" + root = Path("/home/scratch.trt_llm_data_ci/llm-models/") + if "LLM_MODELS_ROOT" in os.environ: + root = Path(os.environ["LLM_MODELS_ROOT"]) + if not root.exists(): + root = Path("/scratch.trt_llm_data/llm-models/") + return str(root) if root.exists() else None + + +def _resolve_glm_checkpoint() -> str: + """Resolve the GlmImage checkpoint reference. + + Resolution order: + 1. ``GLM_IMAGE_MODEL_PATH`` env var (explicit local path or HF Hub ID). + 2. ``/GLM-Image`` staged checkpoint, when present locally. + 3. The canonical HF Hub ID — the pipeline loader downloads it on demand. + """ + explicit = os.environ.get("GLM_IMAGE_MODEL_PATH") + if explicit: + return explicit + root = _llm_models_root() + if root is not None: + staged = os.path.join(root, "GLM-Image") + if os.path.isdir(staged): + return staged + return GLM_IMAGE_HF_ID + + +# Plain HF Hub ID by default; the pipeline loader downloads it on first use. +# Point GLM_IMAGE_MODEL_PATH / LLM_MODELS_ROOT at a local checkpoint to avoid +# re-downloading. +GLM_IMAGE_PATH = _resolve_glm_checkpoint() + +# GlmImage takes a plain text prompt and derives the latent grid from the +# explicit height/width (each must be divisible by 32). 256 == 8 * 32. +PROMPT = "A dinosaur walking through the jungle" +HEIGHT = 256 +WIDTH = 256 +NUM_STEPS = 30 +SEED = 42 +GUIDANCE_SCALE = 1.5 +COS_SIM_THRESHOLD = 0.99 + + +# ============================================================================ +# Helpers +# ============================================================================ + + +def _load_trtllm_pipeline(checkpoint_path: str, skip_components=None, **kwargs): + """Load TRTLLM GlmImage pipeline without torch.compile or warmup.""" + + args = VisualGenArgs( + model=checkpoint_path, torch_compile_config=TorchCompileConfig(enable=False), **kwargs + ) + return PipelineLoader(args).load(skip_warmup=True, skip_components=skip_components) + + +def _load_hf_pipeline(checkpoint_path: str): + """Load HuggingFace diffusers pipeline (auto-detects class from model_index.json).""" + hf_pipe = DiffusionPipeline.from_pretrained( + checkpoint_path, + torch_dtype=torch.bfloat16, + ) + hf_pipe = hf_pipe.to("cuda") + hf_pipe.set_progress_bar_config(disable=True) + return hf_pipe + + +def _teardown_pipeline(pipe) -> None: + """Release a pipeline reference and reclaim GPU memory.""" + del pipe + gc.collect() + torch.cuda.empty_cache() + torch._dynamo.reset() + + +def _capture_trtllm_image( + pipeline, + prompt: str, + height: int, + width: int, + num_inference_steps: int, + guidance_scale: float, + seed: int, +): + """Run full TRTLLM pipeline including VAE decode; return the raw pipeline output.""" + with torch.no_grad(): + return pipeline.forward( + prompt=prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=torch.Generator(device="cuda").manual_seed(seed), + ) + + +def _capture_hf_image( + hf_pipe, + prompt: str, + height: int, + width: int, + num_inference_steps: int, + guidance_scale: float, + seed: int, +): + """Run HF pipeline with output_type='np'; return the raw pipeline output.""" + generator = torch.Generator(device="cuda").manual_seed(seed) + return hf_pipe( + prompt=prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + output_type="np", + ) + + +def _to_image_tensor(images) -> torch.Tensor: + """Convert (B, H, W, C) images to a float (H, W, C) tensor in [0, 1].""" + + if isinstance(images, np.ndarray): + images = torch.from_numpy(images) + images = images[0] + if images.dtype == torch.uint8: + return images.float() / 255.0 + return images.float() + + +def _cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> float: + """Cosine similarity between two tensors (flattened to 1D, cast to float32 on CPU).""" + a_flat = a.float().cpu().reshape(-1) + b_flat = b.float().cpu().reshape(-1) + return F.cosine_similarity(a_flat.unsqueeze(0), b_flat.unsqueeze(0)).clamp(-1.0, 1.0).item() + + +def _assert_pipeline_matches_hf( + checkpoint_path: str, + height: int, + width: int, + model_label: str, +) -> None: + """Run TRTLLM and HF pipelines sequentially, compare decoded image output.""" + # --- TRTLLM --- + trtllm_pipe = None + try: + trtllm_pipe = _load_trtllm_pipeline(checkpoint_path) + trtllm_output = _capture_trtllm_image( + trtllm_pipe, + prompt=PROMPT, + height=height, + width=width, + num_inference_steps=NUM_STEPS, + guidance_scale=GUIDANCE_SCALE, + seed=SEED, + ) + trtllm_image = _to_image_tensor(trtllm_output.image) + finally: + _teardown_pipeline(trtllm_pipe) + + # --- HF reference --- + hf_pipe = None + try: + hf_pipe = _load_hf_pipeline(checkpoint_path) + hf_output = _capture_hf_image( + hf_pipe, + prompt=PROMPT, + height=height, + width=width, + num_inference_steps=NUM_STEPS, + guidance_scale=GUIDANCE_SCALE, + seed=SEED, + ) + hf_image = _to_image_tensor(hf_output.images) + finally: + _teardown_pipeline(hf_pipe) + + # --- Compare --- + assert trtllm_image.numel() == hf_image.numel(), ( + f"{model_label}: element count mismatch — " + f"TRTLLM {trtllm_image.shape} ({trtllm_image.numel()}) vs " + f"HF {hf_image.shape} ({hf_image.numel()})" + ) + + cos_sim = _cosine_similarity(trtllm_image, hf_image) + print(f"\n {model_label} cosine similarity: {cos_sim:.6f}") + assert cos_sim >= COS_SIM_THRESHOLD, ( + f"{model_label}: cosine similarity {cos_sim:.6f} < {COS_SIM_THRESHOLD}. " + f"TRTLLM pipeline output diverges from the HuggingFace reference. " + f"Image shapes — TRTLLM: {trtllm_image.shape}, HF: {hf_image.shape}." + ) + + +# ============================================================================ +# Tests +# ============================================================================ + + +@pytest.mark.integration +class TestGlmImagePipelineCorrectness: + """GlmImage T2I correctness vs HuggingFace reference (256x256).""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_cosine_similarity(self): + _assert_pipeline_matches_hf( + checkpoint_path=GLM_IMAGE_PATH, + height=HEIGHT, + width=WIDTH, + model_label="GlmImage-T2I", + ) + + +@pytest.mark.integration +class TestGlmImageGeneration: + """Generation shape tests for the GlmImage pipeline. + + Validates that single-prompt generation returns a (B, H, W, C) image + batch, matching the current pipeline contract. + """ + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_single_prompt(self): + """Single prompt returns a (B, H, W, C) image batch.""" + + pipe = None + try: + pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH) + result = _capture_trtllm_image( + pipe, + prompt=PROMPT, + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_STEPS, + guidance_scale=GUIDANCE_SCALE, + seed=SEED, + ) + images = _to_image_tensor(result.image) + # _to_image_tensor drops the batch dim, so the remaining tensor is (H, W, C). + assert images.ndim == 3, f"Expected 3D (H,W,C), got {images.ndim}D" + H, W, C = images.shape + assert H == HEIGHT and W == WIDTH and C == 3 + finally: + _teardown_pipeline(pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_batch_prompts(self): + """A list of prompts returns one image per prompt in a single batched forward.""" + + prompts = [PROMPT, "A neon city skyline reflected in a river at night"] + pipe = None + try: + pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH) + with torch.no_grad(): + result = pipe.forward( + prompt=prompts, + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_STEPS, + guidance_scale=GUIDANCE_SCALE, + generator=torch.Generator(device="cuda").manual_seed(SEED), + ) + image = result.image + assert image.shape[0] == len(prompts), ( + f"Expected batch dim {len(prompts)}, got {image.shape[0]}" + ) + assert tuple(image.shape[1:]) == (HEIGHT, WIDTH, 3), ( + f"Expected (H,W,C)=({HEIGHT},{WIDTH},3), got {tuple(image.shape[1:])}" + ) + finally: + _teardown_pipeline(pipe) + + +def test_image_conditioning_not_supported(): + """Passing a condition image raises NotImplementedError (I2I lands in a follow-up MR).""" + # __new__ skips __init__; the guard fires before any self/model access, so no GPU needed. + pipe = GlmImagePipeline.__new__(GlmImagePipeline) + with pytest.raises(NotImplementedError, match="image-to-image"): + pipe.forward(prompt=PROMPT, image=torch.zeros(1)) + + +# ============================================================================= +# Quantization Optimization Tests +# ============================================================================= + + +@pytest.mark.integration +class TestGlmImageQuantizationOptimizations: + """FP8 + TRTLLM attention on GlmImage.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_fp8_trtllm_attention(self): + pipe_args = { + "quant_config": {"quant_algo": "FP8", "dynamic": True}, + "attention_config": AttentionConfig(backend="TRTLLM"), + } + + pipe = None + try: + pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH, **pipe_args) + result = _capture_trtllm_image( + pipe, + prompt=PROMPT, + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_STEPS, + guidance_scale=GUIDANCE_SCALE, + seed=SEED, + ) + + images = _to_image_tensor(result.image) + assert images.ndim == 3, f"Expected 3D (H,W,C), got {images.ndim}D" + H, W, C = images.shape + assert H == HEIGHT and W == WIDTH and C == 3 + finally: + _teardown_pipeline(pipe) + + +# ============================================================================= +# Quantization / dtype feature tests (transformer only) +# ============================================================================= + +_SKIP_AUX = [ + PipelineComponent.TEXT_ENCODER, + PipelineComponent.TOKENIZER, + PipelineComponent.VAE, + PipelineComponent.SCHEDULER, + PipelineComponent.VISION_LANGUAGE_ENCODER, +] + + +def _assert_fp8_blocks_quantized(pipe) -> None: + """Assert at least one transformer-block Linear is FP8 (float8_e4m3fn) with a weight_scale.""" + for name, module in pipe.transformer.named_modules(): + if isinstance(module, Linear) and "transformer_blocks." in name: + assert module.weight.dtype == torch.float8_e4m3fn, ( + f"{name}: expected float8_e4m3fn, got {module.weight.dtype}" + ) + assert hasattr(module, "weight_scale"), f"{name}: missing weight_scale" + return + pytest.fail("No FP8 Linear found in transformer blocks") + + +def _assert_nvfp4_blocks_quantized(pipe) -> None: + """Assert at least one transformer-block Linear is NVFP4 (packed FP4) with a two-level scale.""" + from tensorrt_llm.quantization.utils import fp4_utils + + for name, module in pipe.transformer.named_modules(): + if isinstance(module, Linear) and "transformer_blocks." in name: + assert module.weight.dtype == fp4_utils.float4_e2m1x2, ( + f"{name}: expected float4_e2m1x2, got {module.weight.dtype}" + ) + assert hasattr(module, "weight_scale"), f"{name}: missing weight_scale" + assert hasattr(module, "weight_scale_2"), f"{name}: missing weight_scale_2" + return + pytest.fail("No NVFP4 Linear found in transformer blocks") + + +def _skip_if_no_fp8_ops() -> None: + try: + if not hasattr(torch.ops, "tensorrt_llm"): + pytest.skip("tensorrt_llm torch ops not available") + _ = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor + _ = torch.ops.tensorrt_llm.quantize_e4m3_activation + except (AttributeError, RuntimeError) as e: + pytest.skip(f"FP8 quantization ops not available: {e}") + + +def _skip_if_no_nvfp4_ops() -> None: + if torch.cuda.get_device_capability(0) < (10, 0): + pytest.skip("NVFP4 requires SM>=10.0 (Blackwell+)") + try: + _ = torch.ops.trtllm.fp4_quantize + except (AttributeError, RuntimeError) as e: + pytest.skip(f"fp4_quantize op not available: {e}") + + +def _transformer_inputs(transformer, device: str = "cuda", dtype=torch.bfloat16, seed: int = 42): + """Build a minimal valid GlmImage transformer forward batch.""" + config = transformer.config + in_channels = config.in_channels + text_embed_dim = config.text_embed_dim + codebook_size = config.prior_vq_quantizer_codebook_size + + g = torch.Generator(device=device).manual_seed(seed) + batch_size = 1 + height = width = 32 + seq_len = 6 + return { + "hidden_states": torch.randn( + batch_size, in_channels, height, width, device=device, dtype=dtype, generator=g + ), + "encoder_hidden_states": torch.randn( + batch_size, seq_len, text_embed_dim, device=device, dtype=dtype, generator=g + ), + "prior_token_id": torch.randint( + 0, min(codebook_size, 64), size=(batch_size,), device=device, generator=g + ), + "prior_token_drop": torch.zeros(batch_size, dtype=torch.bool, device=device), + "timestep": torch.randint(0, 1000, size=(batch_size,), device=device, generator=g), + "target_size": torch.tensor( + [[height, width]] * batch_size, dtype=torch.float32, device=device + ), + "crop_coords": torch.tensor([[0, 0]] * batch_size, dtype=torch.float32, device=device), + } + + +def _run_transformer(transformer, inputs: dict) -> torch.Tensor: + """Run the transformer on cloned inputs; return the float32 sample output.""" + cloned = {k: (v.clone() if torch.is_tensor(v) else v) for k, v in inputs.items()} + with torch.no_grad(): + out = transformer(**cloned, return_dict=False)[0] + return out.float() + + +@pytest.mark.integration +class TestGlmImagePipelineFeatures: + """Quantization loading, dtype layout, numerical accuracy, and memory for GlmImage.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_parameter_dtypes(self): + """BF16 pipeline: every transformer param is on CUDA; all non-scale params are BF16.""" + pipe = None + try: + pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH, skip_components=_SKIP_AUX) + bf16_count = 0 + for name, param in pipe.transformer.named_parameters(): + assert param.device.type == "cuda", f"{name} not on CUDA" + if "scale" not in name.lower(): + assert param.dtype == torch.bfloat16, ( + f"{name}: expected bfloat16, got {param.dtype}" + ) + bf16_count += 1 + assert bf16_count > 0, "No BF16 parameters found" + finally: + _teardown_pipeline(pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_fp8_weights_loaded(self): + """FP8 transformer blocks have float8_e4m3fn weights and weight_scale.""" + _skip_if_no_fp8_ops() + pipe = None + try: + pipe = _load_trtllm_pipeline( + GLM_IMAGE_PATH, + skip_components=_SKIP_AUX, + quant_config={"quant_algo": "FP8", "dynamic": True}, + ) + _assert_fp8_blocks_quantized(pipe) + finally: + _teardown_pipeline(pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_fp8_block_scales_weights_loaded(self): + """FP8_BLOCK_SCALES transformer blocks have float8_e4m3fn weights and weight_scale.""" + _skip_if_no_fp8_ops() + pipe = None + try: + pipe = _load_trtllm_pipeline( + GLM_IMAGE_PATH, + skip_components=_SKIP_AUX, + quant_config={"quant_algo": "FP8_BLOCK_SCALES", "dynamic": True}, + ) + _assert_fp8_blocks_quantized(pipe) + finally: + _teardown_pipeline(pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_nvfp4_weights_loaded(self): + """NVFP4 transformer blocks have packed FP4 weights with a two-level scale.""" + _skip_if_no_nvfp4_ops() + pipe = None + try: + pipe = _load_trtllm_pipeline( + GLM_IMAGE_PATH, + skip_components=_SKIP_AUX, + quant_config={"quant_algo": "NVFP4", "dynamic": True}, + ) + _assert_nvfp4_blocks_quantized(pipe) + finally: + _teardown_pipeline(pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_fp8_single_layer_accuracy(self): + """FP8 qkv_proj output matches the BF16 F.linear reference (cos_sim > 0.99).""" + _skip_if_no_fp8_ops() + bf16_pipe = None + fp8_pipe = None + try: + bf16_pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH, skip_components=_SKIP_AUX) + fp8_pipe = _load_trtllm_pipeline( + GLM_IMAGE_PATH, + skip_components=_SKIP_AUX, + quant_config={"quant_algo": "FP8", "dynamic": True}, + ) + linear_bf16 = bf16_pipe.transformer.transformer_blocks[0].attn1.qkv_proj + linear_fp8 = fp8_pipe.transformer.transformer_blocks[0].attn1.qkv_proj + + weight = linear_bf16.weight.data.clone() + bias = linear_bf16.bias.data.clone() if linear_bf16.bias is not None else None + x = torch.randn( + 1024, + linear_bf16.in_features, + dtype=torch.bfloat16, + device="cuda", + generator=torch.Generator("cuda").manual_seed(42), + ) + + with torch.no_grad(): + ref = F.linear(x, weight, bias) + fp8_out = linear_fp8(x) + + cos_sim = F.cosine_similarity( + fp8_out.flatten().float(), ref.flatten().float(), dim=0 + ).item() + mse = F.mse_loss(fp8_out.float(), ref.float()).item() + print(f"\n FP8 qkv_proj: cos_sim={cos_sim:.6f}, mse={mse:.6f}") + assert cos_sim > 0.99, f"cos_sim too low: {cos_sim:.6f}" + assert mse < 1.0, f"MSE too high: {mse:.6f}" + finally: + if fp8_pipe is not None: + _teardown_pipeline(fp8_pipe) + if bf16_pipe is not None: + _teardown_pipeline(bf16_pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_fp8_memory_savings(self): + """FP8 transformer uses less parameter memory than BF16.""" + _skip_if_no_fp8_ops() + + def _mem_gb(pipe): + return ( + sum(p.numel() * p.element_size() for p in pipe.transformer.parameters()) / 1024**3 + ) + + bf16_pipe = None + fp8_pipe = None + try: + bf16_pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH, skip_components=_SKIP_AUX) + fp8_pipe = _load_trtllm_pipeline( + GLM_IMAGE_PATH, + skip_components=_SKIP_AUX, + quant_config={"quant_algo": "FP8", "dynamic": True}, + ) + bf16_gb = _mem_gb(bf16_pipe) + fp8_gb = _mem_gb(fp8_pipe) + ratio = bf16_gb / fp8_gb + print(f"\n BF16={bf16_gb:.3f} GB, FP8={fp8_gb:.3f} GB, ratio={ratio:.2f}x") + assert ratio > 1.9, f"Expected FP8 memory reduction, got {ratio:.2f}x" + finally: + if fp8_pipe is not None: + _teardown_pipeline(fp8_pipe) + if bf16_pipe is not None: + _teardown_pipeline(bf16_pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"]) + def test_fp8_e2e_accuracy(self, quant_algo): + """FP8 / FP8_BLOCK_SCALES full-transformer output close to BF16 (cos_sim > 0.99).""" + _skip_if_no_fp8_ops() + bf16_pipe = None + quant_pipe = None + try: + bf16_pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH, skip_components=_SKIP_AUX) + quant_pipe = _load_trtllm_pipeline( + GLM_IMAGE_PATH, + skip_components=_SKIP_AUX, + quant_config={"quant_algo": quant_algo, "dynamic": True}, + ) + inputs = _transformer_inputs(bf16_pipe.transformer) + out_bf16 = _run_transformer(bf16_pipe.transformer, inputs) + out_quant = _run_transformer(quant_pipe.transformer, inputs) + + assert not torch.isnan(out_bf16).any(), "BF16 output contains NaN" + assert not torch.isinf(out_bf16).any(), "BF16 output contains Inf" + assert not torch.isnan(out_quant).any(), f"{quant_algo} output contains NaN" + assert not torch.isinf(out_quant).any(), f"{quant_algo} output contains Inf" + + cos_sim = F.cosine_similarity(out_quant.flatten(), out_bf16.flatten(), dim=0).item() + mse = F.mse_loss(out_quant, out_bf16).item() + print( + f"\n {quant_algo} E2E ({len(bf16_pipe.transformer.transformer_blocks)} layers): " + f"cos_sim={cos_sim:.6f}, mse={mse:.6f}" + ) + # Block-scales is a coarser (per-block) scheme than per-tensor FP8, so it + # tolerates a slightly lower cosine-similarity floor. + min_cos_sim = 0.975 if quant_algo == "FP8_BLOCK_SCALES" else 0.99 + assert cos_sim > min_cos_sim, f"{quant_algo} cos_sim too low: {cos_sim:.6f}" + finally: + if quant_pipe is not None: + _teardown_pipeline(quant_pipe) + if bf16_pipe is not None: + _teardown_pipeline(bf16_pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_nvfp4_e2e_accuracy(self): + """NVFP4 full-transformer output close to BF16 (cos_sim > 0.90).""" + _skip_if_no_nvfp4_ops() + bf16_pipe = None + nvfp4_pipe = None + try: + bf16_pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH, skip_components=_SKIP_AUX) + nvfp4_pipe = _load_trtllm_pipeline( + GLM_IMAGE_PATH, + skip_components=_SKIP_AUX, + quant_config={"quant_algo": "NVFP4", "dynamic": True}, + ) + inputs = _transformer_inputs(bf16_pipe.transformer) + out_bf16 = _run_transformer(bf16_pipe.transformer, inputs) + out_nvfp4 = _run_transformer(nvfp4_pipe.transformer, inputs) + + assert not torch.isnan(out_nvfp4).any(), "NVFP4 output contains NaN" + assert not torch.isinf(out_nvfp4).any(), "NVFP4 output contains Inf" + + cos_sim = F.cosine_similarity(out_nvfp4.flatten(), out_bf16.flatten(), dim=0).item() + mse = F.mse_loss(out_nvfp4, out_bf16).item() + print( + f"\n NVFP4 E2E ({len(bf16_pipe.transformer.transformer_blocks)} layers): " + f"cos_sim={cos_sim:.6f}, mse={mse:.6f}" + ) + # NOTE: AdaLayerNorm modulation linears hurt accuracy at nvfp4, consider adding omissions later + assert cos_sim > 0.90, f"NVFP4 cos_sim too low: {cos_sim:.6f}" + finally: + if nvfp4_pipe is not None: + _teardown_pipeline(nvfp4_pipe) + if bf16_pipe is not None: + _teardown_pipeline(bf16_pipe) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py b/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py new file mode 100644 index 000000000000..5c90b25ad009 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.visual_gen import DiffusionModelConfig +from tensorrt_llm.llmapi import QuantConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.visual_gen import AttentionConfig + +CONFIG = { + "text_embed_dim": 1472, + "in_channels": 16, + "attention_head_dim": 128, + "num_attention_heads": 32, +} + + +def _create_model_config(config_dict: dict, backend: str = "VANILLA") -> DiffusionModelConfig: + """Create DiffusionModelConfig from config dict.""" + pretrained_config = SimpleNamespace(**config_dict) + return DiffusionModelConfig( + pretrained_config=pretrained_config, + quant_config=QuantConfig(), + mapping=Mapping(), + attention=AttentionConfig(backend=backend), + skip_create_weights_in_init=False, + ) + + +def _make_inputs(config, dtype, device): + batch_size = 1 + height, width = 32, 32 + + sequence_length = 6 + text_embed_dim = config["text_embed_dim"] + num_channels = config["in_channels"] + + generator = torch.Generator(device=device).manual_seed(42) + + hidden_states = torch.randn( + batch_size, + num_channels, + height, + width, + device=device, + dtype=dtype, + generator=generator, + ) + + encoder_hidden_states = torch.randn( + batch_size, sequence_length, text_embed_dim, device=device, dtype=dtype, generator=generator + ) + + prior_token_id = torch.randint(0, 64, size=(batch_size,), generator=generator, device=device) + prior_token_drop = torch.zeros(batch_size, dtype=torch.bool, device=device) + + timestep = torch.randint(0, 1000, size=(batch_size,), generator=generator, device=device) + target_size = torch.tensor([[height, width]] * batch_size, dtype=torch.float32, device=device) + crop_coords = torch.tensor([[0, 0]] * batch_size, dtype=torch.float32, device=device) + + return ( + hidden_states, + encoder_hidden_states, + prior_token_id, + prior_token_drop, + target_size, + crop_coords, + timestep, + ) + + +class TestGlmImageTransformer(unittest.TestCase): + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_glm_image_structure(self): + from tensorrt_llm._torch.visual_gen.models.glm_image import GlmImageTransformer2DModel + + config = deepcopy(CONFIG) + config["num_layers"] = 1 + model_config = _create_model_config(config) + model = GlmImageTransformer2DModel(model_config) + + # Check components are present in model + components = [ + "rope", + "image_projector", + "glyph_projector", + "prior_token_embedding", + "prior_projector", + "time_condition_embed", + "transformer_blocks", + "norm_out", + "proj_out", + ] + + for component in components: + self.assertTrue(hasattr(model, component)) + + self.assertEqual(len(model.transformer_blocks), 1) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_glm_image_forward_sanity(self): + from tensorrt_llm._torch.visual_gen.models.glm_image import GlmImageTransformer2DModel + + config = deepcopy(CONFIG) + config["num_layers"] = 1 + model_config = _create_model_config(config) + dtype = torch.bfloat16 + + model = GlmImageTransformer2DModel(model_config).to(self.DEVICE, dtype=dtype).eval() + + ( + hidden_states, + encoder_hidden_states, + prior_token_id, + prior_token_drop, + target_size, + crop_coords, + timestep, + ) = _make_inputs(config, dtype, self.DEVICE) + + with torch.no_grad(): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + prior_token_id=prior_token_id, + prior_token_drop=prior_token_drop, + target_size=target_size, + crop_coords=crop_coords, + ) + + # Model returns {"sample": tensor} + if isinstance(output, dict): + output = output["sample"] + + # Note: With random weights, NaN can occur. For unit tests, we only check shape. + # Full numerical correctness is tested in TestFluxHuggingFaceComparison. + self.assertEqual(output.shape, hidden_states.shape) + + +class TestGlmImageHuggingFaceComparison(unittest.TestCase): + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_glm_image_allclose_to_hf(self): + try: + from diffusers import GlmImageTransformer2DModel as HFGlmImageTransformer2DModel + except ImportError: + self.skipTest("diffusers not installed") + + from tensorrt_llm._torch.visual_gen.models.glm_image import GlmImageTransformer2DModel + + torch.manual_seed(42) + + # Create TRT-LLM Model + config = deepcopy(CONFIG) + config["num_layers"] = 1 + model_config = _create_model_config(config) + dtype = torch.bfloat16 + trtllm_model = GlmImageTransformer2DModel(model_config).to(self.DEVICE, dtype=dtype).eval() + + hf_model = ( + HFGlmImageTransformer2DModel( + in_channels=config["in_channels"], + num_layers=config["num_layers"], + attention_head_dim=config["attention_head_dim"], + num_attention_heads=config["num_attention_heads"], + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + # Copy weights from HF to TRT-LLM + hf_state_dict = hf_model.state_dict() + for name, _ in hf_state_dict.items(): + print(f"{name}") + + print("\n") + for name, _ in trtllm_model.named_modules(): + print(f"{name}") + + trtllm_model.load_weights(hf_state_dict) + + # Create inputs + ( + hidden_states, + encoder_hidden_states, + prior_token_id, + prior_token_drop, + target_size, + crop_coords, + timestep, + ) = _make_inputs(config, dtype, self.DEVICE) + + with torch.no_grad(): + hf_output = hf_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + prior_token_id=prior_token_id, + prior_token_drop=prior_token_drop, + target_size=target_size, + crop_coords=crop_coords, + return_dict=False, + )[0] + + trtllm_output = trtllm_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + prior_token_id=prior_token_id, + prior_token_drop=prior_token_drop, + target_size=target_size, + crop_coords=crop_coords, + ) + + # Model returns {"sample": tensor} + if isinstance(trtllm_output, dict): + trtllm_output = trtllm_output["sample"] + + # Compare outputs + hf_output = hf_output.float() + trtllm_output = trtllm_output.float() + + cos_sim = F.cosine_similarity( + hf_output.flatten().unsqueeze(0), trtllm_output.flatten().unsqueeze(0) + ).item() + + max_diff = (hf_output - trtllm_output).abs().max().item() + + print("\n[GlmImage HF Comparison]") + print(f" Cosine similarity: {cos_sim:.6f}") + print(f" Max diff: {max_diff:.6f}") + + self.assertGreater(cos_sim, 0.99, f"Cosine similarity too low: {cos_sim}") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 8ecf63d3aa77559d16cb8cec055110280747ea36 Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Tue, 7 Jul 2026 21:53:41 +0000 Subject: [PATCH 02/11] Feedback Signed-off-by: Joseph Loftin --- .../visual_gen/models/glm_image/__init__.py | 2 +- .../models/glm_image/pipeline_glm_image.py | 4 +- .../models/glm_image/transformer_glm_image.py | 44 +++++++++++-------- .../visual_gen/test_glm_image_pipeline.py | 12 +---- 4 files changed, 30 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py index 16b659411d12..4d52909b942a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/__init__.py @@ -16,4 +16,4 @@ from .pipeline_glm_image import GlmImagePipeline from .transformer_glm_image import GlmImageAttention, GlmImageTransformer2DModel -__all__ = ["GlmImagePipeline", "GlmImageTransformer2DModel", "GlmImageAttention"] +__all__ = ["GlmImageAttention", "GlmImagePipeline", "GlmImageTransformer2DModel"] diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py index 7f76dd93a9a6..629b6dcb5831 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py @@ -513,7 +513,7 @@ def do_classifier_free_guidance(self): def _get_glyph_embeds( self, - prompt: Union[str, List[str]] = None, + prompt: Optional[Union[str, List[str]]] = None, max_sequence_length: int = 2048, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, @@ -1095,7 +1095,7 @@ def load_weights(self, weights: dict) -> None: def _init_transformer(self) -> None: """Initialize GlmImage transformer with quantization support.""" - logger.info("Creating HunyuanVideo1.5 transformer with quantization support...") + logger.info("Creating GlmImage transformer with quantization support...") self.transformer = GlmImageTransformer2DModel( model_config=self.pipeline_config.model_configs["transformer"] ) diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py index eb996b1fa22d..c42c44222df0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -138,7 +138,7 @@ def __init__( self, in_channels: int, time_embed_dim: int, - out_dim: int = None, + out_dim: Optional[int] = None, model_config: Optional[DiffusionModelConfig] = None, ): super().__init__() @@ -277,7 +277,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> Tuple[torch.Tensor, torch.Tensor]: dtype = hidden_states.dtype norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) @@ -433,8 +433,8 @@ def forward( attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ): - batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape - batch_size, image_seq_length, embed_dim = hidden_states.shape + batch_size, text_seq_length, _ = encoder_hidden_states.shape + batch_size, image_seq_length, _ = hidden_states.shape hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) # QKV Proj @@ -525,13 +525,16 @@ def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - temb: torch.Tensor | None = None, - image_rotary_emb: tuple[torch.Tensor, torch.Tensor] - | list[tuple[torch.Tensor, torch.Tensor]] - | None = None, - attention_mask: dict[str, torch.Tensor] | None = None, - attention_kwargs: dict[str, Any] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + temb: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[ + Union[ + Tuple[torch.Tensor, torch.Tensor], + List[Tuple[torch.Tensor, torch.Tensor]], + ] + ] = None, + attention_mask: Optional[Dict[str, torch.Tensor]] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: # 1. Timestep conditioning ( norm_hidden_states, @@ -720,13 +723,16 @@ def forward( timestep: torch.LongTensor, target_size: torch.Tensor, crop_coords: torch.Tensor, - attention_kwargs: dict[str, Any] | None = None, + attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = True, - attention_mask: torch.Tensor | None = None, - image_rotary_emb: tuple[torch.Tensor, torch.Tensor] - | list[tuple[torch.Tensor, torch.Tensor]] - | None = None, - ) -> tuple[torch.Tensor] | Transformer2DModelOutput: + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[ + Union[ + Tuple[torch.Tensor, torch.Tensor], + List[Tuple[torch.Tensor, torch.Tensor]], + ] + ] = None, + ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: batch_size, num_channels, height, width = hidden_states.shape # 1. RoPE @@ -749,7 +755,7 @@ def forward( temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) # 3. Transformer blocks - for idx, block in enumerate(self.transformer_blocks): + for block in self.transformer_blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( block, diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py index 0afc2b1f199a..ccd949d71c95 100644 --- a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py @@ -18,9 +18,7 @@ from pathlib import Path from typing import Optional -from tensorrt_llm._torch.modules.linear import Linear - -os.environ["TLLM_DISABLE_MPI"] = "1" +os.environ.setdefault("TLLM_DISABLE_MPI", "1") import numpy as np import pytest @@ -28,17 +26,11 @@ import torch.nn.functional as F from diffusers import DiffusionPipeline +from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.visual_gen.models.glm_image import GlmImagePipeline from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs - -@pytest.fixture(autouse=True, scope="module") -def _cleanup_mpi_env(): - yield - os.environ.pop("TLLM_DISABLE_MPI", None) - - # ============================================================================ # Test constants # ============================================================================ From a8f33db8e9f777e4c08589727484f484f88922ac Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Mon, 20 Jul 2026 22:30:55 +0000 Subject: [PATCH 03/11] Normalized timesteps; Remove config Signed-off-by: Joseph Loftin --- examples/visual_gen/README.md | 1 - .../configs/glm-image-fp8-1gpu.yaml | 29 ------------------- examples/visual_gen/models/glm_image.py | 1 - .../models/glm_image/pipeline_glm_image.py | 2 +- .../models/glm_image/transformer_glm_image.py | 10 ++++++- .../visual_gen/test_glm_image_pipeline.py | 2 +- .../visual_gen/test_glm_image_transformer.py | 8 +++-- 7 files changed, 17 insertions(+), 36 deletions(-) delete mode 100644 examples/visual_gen/configs/glm-image-fp8-1gpu.yaml diff --git a/examples/visual_gen/README.md b/examples/visual_gen/README.md index a6f755e2cab5..eab1b15cbf33 100644 --- a/examples/visual_gen/README.md +++ b/examples/visual_gen/README.md @@ -38,7 +38,6 @@ python models/cosmos3_ti2v.py --visual_gen_args configs/cosmos3-nano-1gpu.yaml - python models/qwen_image.py --visual_gen_args configs/qwen-image-fp8-1gpu.yaml python models/qwen_image_layered.py --visual_gen_args configs/qwen-image-layered-1gpu.yaml --image /path/to/image.png python models/qwen_image_edit.py --visual_gen_args configs/qwen-image-edit-2511-fp4-1gpu.yaml --image /path/to/source.png --prompt "Make the image look like a watercolor painting" -python models/glm_image.py --visual_gen_args configs/glm-image-fp8-1gpu.yaml python models/hunyuan_t2v.py --visual_gen_args configs/hunyuan-t2v-fp8-1gpu.yaml ``` diff --git a/examples/visual_gen/configs/glm-image-fp8-1gpu.yaml b/examples/visual_gen/configs/glm-image-fp8-1gpu.yaml deleted file mode 100644 index 611039a23591..000000000000 --- a/examples/visual_gen/configs/glm-image-fp8-1gpu.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# 1-GPU GlmImage text-to-image with FP8 dynamic quantization. -# Model: zai-org/GLM-Image -# Shared by offline examples (--visual_gen_args) and trtllm-serve. -# -# GlmImage constraints: single-GPU text-to-image only. Image-to-image -# conditioning, sequence/CFG parallelism, and caching are not yet supported. -quant_config: - quant_algo: FP8 - dynamic: true -parallel_config: - cfg_size: 1 - ulysses_size: 1 -cuda_graph_config: - enable: false diff --git a/examples/visual_gen/models/glm_image.py b/examples/visual_gen/models/glm_image.py index 093548e24bfb..f47a19469a6a 100644 --- a/examples/visual_gen/models/glm_image.py +++ b/examples/visual_gen/models/glm_image.py @@ -16,7 +16,6 @@ Usage: python glm_image.py - python glm_image.py --visual_gen_args ../configs/glm-image-fp8-1gpu.yaml """ import argparse diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py index 629b6dcb5831..6c6863e574f1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py @@ -987,7 +987,7 @@ def forward_fn( encoder_hidden_states=encoder_hidden_states, prior_token_id=extra_tensors["prior_token_id"], prior_token_drop=extra_tensors["prior_token_drop"], - timestep=timestep - 1, + timestep=(timestep - 1) / self.scheduler.config.num_train_timesteps, target_size=extra_tensors["target_size"], crop_coords=extra_tensors["crop_coords"], attention_mask=extra_tensors.get("attention_mask"), diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py index c42c44222df0..1ca35857b612 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py @@ -607,6 +607,7 @@ def __init__(self, model_config: DiffusionModelConfig): ) text_embed_dim = getattr(pretrained_config, "text_embed_dim", 1472) time_embed_dim = getattr(pretrained_config, "time_embed_dim", 512) + num_train_timesteps = getattr(pretrained_config, "num_train_timesteps", 1000) pooled_projection_dim = 2 * 2 * condition_dim inner_dim = num_attention_heads * attention_head_dim @@ -624,6 +625,7 @@ def __init__(self, model_config: DiffusionModelConfig): "prior_vq_quantizer_codebook_size": prior_vq_quantizer_codebook_size, "text_embed_dim": text_embed_dim, "time_embed_dim": time_embed_dim, + "num_train_timesteps": num_train_timesteps, "pooled_projection_dim": pooled_projection_dim, "inner_dim": inner_dim, "dtype": dtype, @@ -720,7 +722,7 @@ def forward( encoder_hidden_states: torch.Tensor, prior_token_id: torch.Tensor, prior_token_drop: torch.Tensor, - timestep: torch.LongTensor, + timestep: torch.Tensor, target_size: torch.Tensor, crop_coords: torch.Tensor, attention_kwargs: Optional[Dict[str, Any]] = None, @@ -733,6 +735,10 @@ def forward( ] ] = None, ) -> Union[Tuple[torch.Tensor], Transformer2DModelOutput]: + """ + Args: + timestep: Normalized scheduler timestep tensor in [0, 1]. + """ batch_size, num_channels, height, width = hidden_states.shape # 1. RoPE @@ -752,6 +758,8 @@ def forward( hidden_states = hidden_states + prior_hidden_states + # GlmImage timestep embeddings use the scheduler's 1000-step scale internally. + timestep = timestep * self.config.num_train_timesteps temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) # 3. Transformer blocks diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py index ccd949d71c95..9b4730c3ba3b 100644 --- a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py @@ -440,7 +440,7 @@ def _transformer_inputs(transformer, device: str = "cuda", dtype=torch.bfloat16, 0, min(codebook_size, 64), size=(batch_size,), device=device, generator=g ), "prior_token_drop": torch.zeros(batch_size, dtype=torch.bool, device=device), - "timestep": torch.randint(0, 1000, size=(batch_size,), device=device, generator=g), + "timestep": torch.tensor([0.5] * batch_size, device=device, dtype=dtype), "target_size": torch.tensor( [[height, width]] * batch_size, dtype=torch.float32, device=device ), diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py b/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py index 5c90b25ad009..af9c7efd5746 100644 --- a/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py @@ -33,6 +33,8 @@ "num_attention_heads": 32, } +NUM_TRAIN_TIMESTEPS = 1000 + def _create_model_config(config_dict: dict, backend: str = "VANILLA") -> DiffusionModelConfig: """Create DiffusionModelConfig from config dict.""" @@ -73,7 +75,8 @@ def _make_inputs(config, dtype, device): prior_token_id = torch.randint(0, 64, size=(batch_size,), generator=generator, device=device) prior_token_drop = torch.zeros(batch_size, dtype=torch.bool, device=device) - timestep = torch.randint(0, 1000, size=(batch_size,), generator=generator, device=device) + # Normalized timestep in [0, 1]; 0.5 is exact in bf16 so it round-trips cleanly. + timestep = torch.tensor([0.5] * batch_size, device=device, dtype=dtype) target_size = torch.tensor([[height, width]] * batch_size, dtype=torch.float32, device=device) crop_coords = torch.tensor([[0, 0]] * batch_size, dtype=torch.float32, device=device) @@ -214,10 +217,11 @@ def test_glm_image_allclose_to_hf(self): ) = _make_inputs(config, dtype, self.DEVICE) with torch.no_grad(): + # HF consumes the raw timestep; TRT-LLM consumes the normalized one. hf_output = hf_model( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, - timestep=timestep, + timestep=timestep * NUM_TRAIN_TIMESTEPS, prior_token_id=prior_token_id, prior_token_drop=prior_token_drop, target_size=target_size, From c247840c06fca293e7490a76d614e97a62ecad8a Mon Sep 17 00:00:00 2001 From: jloftin Date: Thu, 13 Aug 2026 22:38:39 +0000 Subject: [PATCH 04/11] Review, E2E, golden Signed-off-by: jloftin --- .../visual_gen/serve/configs/glm_image.yml | 15 +- .../models/glm_image/pipeline_glm_image.py | 40 +-- .../models/glm_image/transformer_glm_image.py | 91 +++---- .../glm_image_fp8_blockwise_lpips_golden.json | 25 ++ .../glm_image_lpips_golden.json | 21 ++ .../glm_image_nvfp4_lpips_golden.json | 25 ++ .../visual_gen_lpips_golden_media.zip | 4 +- .../visual_gen/test_visual_gen_glm.py | 239 ++++++++++++++++++ .../test_lists/test-db/l0_b200.yml | 3 + .../visual_gen/test_glm_image_pipeline.py | 60 ++++- .../visual_gen/test_glm_image_transformer.py | 2 +- 11 files changed, 441 insertions(+), 84 deletions(-) create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_fp8_blockwise_lpips_golden.json create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_lpips_golden.json create mode 100644 tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_nvfp4_lpips_golden.json create mode 100644 tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py diff --git a/examples/visual_gen/serve/configs/glm_image.yml b/examples/visual_gen/serve/configs/glm_image.yml index 138fd010c5a4..c91cbe130a55 100644 --- a/examples/visual_gen/serve/configs/glm_image.yml +++ b/examples/visual_gen/serve/configs/glm_image.yml @@ -1,5 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. # GlmImage text-to-image (single GPU). -# Image-to-image conditioning, sequence/CFG parallelism, and caching are not yet supported. parallel_config: cfg_size: 1 ulysses_size: 1 diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py index 6c6863e574f1..e392b7a6f565 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py @@ -22,14 +22,15 @@ import numpy as np import PIL import torch +from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler from diffusers.image_processor import VaeImageProcessor -from diffusers.pipelines.glm_image.pipeline_glm_image import ( - AutoencoderKL, - FlowMatchEulerDiscreteScheduler, +from diffusers.utils.torch_utils import randn_tensor +from transformers import ( + ByT5Tokenizer, + GlmImageForConditionalGeneration, + GlmImageProcessor, T5EncoderModel, ) -from diffusers.utils.torch_utils import randn_tensor -from transformers import ByT5Tokenizer, GlmImageForConditionalGeneration, GlmImageProcessor from tensorrt_llm import logger from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig @@ -323,16 +324,19 @@ def generate_prior_tokens( # Generate with AR model # Set torch random seed from generator for reproducibility # (transformers generate() doesn't accept generator parameter) - if generator is not None: - seed = generator.initial_seed() - torch.manual_seed(seed) - if device is not None and device.type == "cuda": - torch.cuda.manual_seed(seed) - outputs = self.vision_language_encoder.generate( - **inputs, - max_new_tokens=max_new_tokens, - do_sample=True, - ) + # fork_rng keeps the reseed out of the process-global RNG + fork_devices = [device] if device is not None and device.type == "cuda" else [] + with torch.random.fork_rng(devices=fork_devices, enabled=generator is not None): + if generator is not None: + seed = generator.initial_seed() + torch.manual_seed(seed) + if device is not None and device.type == "cuda": + torch.cuda.manual_seed(seed) + outputs = self.vision_language_encoder.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=True, + ) # Extract and upsample prior tokens for each sample # For left-padded inputs, generated tokens start after the padded input sequence @@ -797,6 +801,12 @@ def forward( self._attention_kwargs = attention_kwargs self._current_timestep = None + # forward() can be called directly, bypassing infer()'s VisualGenParams defaults + if height is None: + height = self.default_generation_params["height"] + if width is None: + width = self.default_generation_params["width"] + if prompt is not None and isinstance(prompt, str): batch_size = 1 elif prompt is not None and isinstance(prompt, list): diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py index 1ca35857b612..2519d6574d2b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py @@ -352,7 +352,6 @@ def __init__( dim: int, num_attention_heads: int, attention_head_dim: int, - eps: float = 1e-6, dtype: Optional[torch.dtype] = None, config: Optional[DiffusionModelConfig] = None, layer_idx: int = 0, @@ -365,50 +364,16 @@ def __init__( config=config, qk_norm_mode="per_head", qkv_mode=QKVMode.FUSE_QKV, - qk_norm=True, + qk_norm=False, ) self.heads = num_attention_heads self.head_dim = attention_head_dim - self.add_q_proj = Linear( - dim, - dim, - bias=True, - dtype=dtype, - mapping=self.mapping, - quant_config=self.quant_config, - skip_create_weights_in_init=self.skip_create_weights_in_init, - force_dynamic_quantization=self.force_dynamic_quantization, - ) - self.add_k_proj = Linear( - dim, - dim, - bias=True, - dtype=dtype, - mapping=self.mapping, - quant_config=self.quant_config, - skip_create_weights_in_init=self.skip_create_weights_in_init, - force_dynamic_quantization=self.force_dynamic_quantization, - ) - self.add_v_proj = Linear( - dim, - dim, - bias=True, - dtype=dtype, - mapping=self.mapping, - quant_config=self.quant_config, - skip_create_weights_in_init=self.skip_create_weights_in_init, - force_dynamic_quantization=self.force_dynamic_quantization, - ) - - # QK-norms, applied per-head on the head_dim. - self.norm_added_q = RMSNorm( - hidden_size=attention_head_dim, eps=eps, dtype=dtype, has_weights=True - ) - self.norm_added_k = RMSNorm( - hidden_size=attention_head_dim, eps=eps, dtype=dtype, has_weights=True - ) + # GLM-Image uses parameter-free per-head LayerNorm for Q/K, not the base + # class's learned RMSNorm (which has no checkpoint counterpart). + self.norm_q = torch.nn.LayerNorm(attention_head_dim, eps=1e-5, elementwise_affine=False) + self.norm_k = torch.nn.LayerNorm(attention_head_dim, eps=1e-5, elementwise_affine=False) self.to_out = torch.nn.ModuleList( [ @@ -426,6 +391,10 @@ def __init__( ] ) + def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Parameter-free per-head LayerNorm on 4D tensors [B, S, H, D].""" + return self.norm_q(q), self.norm_k(k) + def forward( self, hidden_states: torch.Tensor, @@ -444,8 +413,7 @@ def forward( value = value.unflatten(2, (self.heads, -1)) # 2. QK normalization - if self.qk_norm: - query, key = self.apply_qk_norm(query, key) + query, key = self.apply_qk_norm(query, key) # 3. Rotational positional embeddings applied to latent stream if image_rotary_emb is not None: @@ -470,14 +438,15 @@ def forward( assert text_attn_mask.dim() == 2, ( "the shape of text_attn_mask should be (batch_size, text_seq_length)" ) - text_attn_mask = text_attn_mask.float().to(query.device) + # Backends take a [B, S_kv] bool key-padding mask (True = valid); masking + # the padded key columns matches the reference [B, 1, S, S] score mask. mix_attn_mask = torch.ones( - (batch_size, text_seq_length + image_seq_length), device=query.device + (batch_size, text_seq_length + image_seq_length), + dtype=torch.bool, + device=query.device, ) - mix_attn_mask[:, :text_seq_length] = text_attn_mask - mix_attn_mask = mix_attn_mask.unsqueeze(2) - attn_mask_matrix = mix_attn_mask @ mix_attn_mask.transpose(1, 2) - attention_mask = (attn_mask_matrix > 0).unsqueeze(1).to(query.dtype) + mix_attn_mask[:, :text_seq_length] = text_attn_mask.bool().to(query.device) + attention_mask = mix_attn_mask hidden_states = self._attn_impl(query, key, value, key_padding_mask=attention_mask) @@ -511,7 +480,7 @@ def __init__( # 1. Attention self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim, model_config=config) self.attn1 = GlmImageAttention( - dim, num_attention_heads, attention_head_dim, eps, dtype, config, layer_idx + dim, num_attention_heads, attention_head_dim, dtype, config, layer_idx ) # 2. Feedforward @@ -639,6 +608,15 @@ def __init__(self, model_config: DiffusionModelConfig): self.rope = GlmImageRotaryPosEmbed(attention_head_dim, patch_size, theta=10000.0) # 2. Patch & Text-timestep embedding + # NOTE: image_projector quantization is excluded when in_channels * patch_size**2 < 128. + # GLM-Image has 64, which is below the 128-block size required by + # fp8_block_scaling_gemm (causes NVRTC compilation failure). This layer runs + # once per forward pass (not in the block loop), so the perf impact is negligible. + if in_channels * patch_size**2 < 128 and quant_config is not None: + if quant_config.exclude_modules is None: + quant_config.exclude_modules = [] + if "*image_projector*" not in quant_config.exclude_modules: + quant_config.exclude_modules.append("*image_projector*") self.image_projector = GlmImageImageProjector( in_channels, inner_dim, patch_size, model_config=model_config ) @@ -807,26 +785,15 @@ def load_weights(self, weights: dict) -> None: weights: Dictionary of parameter name -> tensor """ - # Map fused QKV layer names to original HF checkpoint names - # HF checkpoint has separate to_q, to_k, to_v / add_q_proj, add_k_proj, add_v_proj - # We fuse them into qkv_proj / add_qkv_proj for better performance + # Map fused QKV layer name to original HF checkpoint names + # We fuse to_q, to_k, to_v into qkv_proj for better performance params_map = { - "add_qkv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], "qkv_proj": ["to_q", "to_k", "to_v"], } loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map) - # Track prefixes of wrapper projectors whose sub-Linears are loaded - # by the parent's load_weights — the generic Linear loader must skip - # them (their FUSED weight modes would look for nonexistent checkpoint - # keys via params_map and error). - managed_prefixes = set() - for name, module in tqdm(self.named_modules(), desc="Loading weights"): - if any(name.startswith(p) for p in managed_prefixes): - continue - # Create weights for modules with skip_create_weights_in_init=True # This must be done before loading weights (following Wan pattern) if callable(getattr(module, "create_weights", None)): diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_fp8_blockwise_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_fp8_blockwise_lpips_golden.json new file mode 100644 index 000000000000..2583616c4294 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_fp8_blockwise_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "glm_image_fp8_blockwise_lpips_golden.png", + "model": "GLM-Image", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 1024, + "width": 1024, + "num_inference_steps": 30, + "guidance_scale": 1.5, + "seed": 42, + "feature_config": { + "quantization": "FP8_BLOCK_SCALES", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc25", + "tensorrt_llm_commit": "14ad9b524dd98dd514dd265601ac9c993b2d093b", + "container_image": "artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607311529-16970", + "sha256": "04ed621df875251e1bab8e5685b42c94e95cccef1152ae601f2945e4f1a0c1c5" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_lpips_golden.json new file mode 100644 index 000000000000..2eb5d410a99a --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_lpips_golden.json @@ -0,0 +1,21 @@ +{ + "image": "glm_image_lpips_golden.png", + "model": "GLM-Image", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 1024, + "width": 1024, + "num_inference_steps": 30, + "guidance_scale": 1.5, + "seed": 42, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc25", + "tensorrt_llm_commit": "14ad9b524dd98dd514dd265601ac9c993b2d093b", + "container_image": "artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607311529-16970", + "sha256": "c14e7900e45c35d8b123bf6654f5110adb0862372db0afae421b7734490a93e2" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_nvfp4_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_nvfp4_lpips_golden.json new file mode 100644 index 000000000000..8445e2e18324 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/glm_image_nvfp4_lpips_golden.json @@ -0,0 +1,25 @@ +{ + "image": "glm_image_nvfp4_lpips_golden.png", + "model": "GLM-Image", + "source": "TensorRT-LLM VisualGen", + "prompt": "a tiny astronaut hatching from an egg on the moon", + "height": 1024, + "width": 1024, + "num_inference_steps": 30, + "guidance_scale": 1.5, + "seed": 42, + "feature_config": { + "quantization": "NVFP4", + "cuda_graph": false + }, + "torch_compile": false, + "deterministic_algorithms": true, + "lpips_net": "alex", + "lpips_threshold": 0.05, + "diffusers_version": "0.39.0", + "torch_version": "2.12.0a0+5aff3928d8.nv26.05", + "tensorrt_llm_version": "1.3.0rc25", + "tensorrt_llm_commit": "14ad9b524dd98dd514dd265601ac9c993b2d093b", + "container_image": "artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607311529-16970", + "sha256": "8489b4a994a3ae93c109d556e3ba3c932f19f150de3f53808d0c3e1b20688fcc" +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip index acb2348552ae..8b553103413c 100644 --- a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:69011916707974699428b03ecfad96b96bf76750926452ae2779d62ada23e5b0 -size 34534138 +oid sha256:7043356a539371abc61e23c275df3e8388ff4d80894bf9d7fb7845eaf764741b +size 38528317 diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py new file mode 100644 index 000000000000..87d758dff00e --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_glm.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Single-device GLM-Image visual-quality regression tests.""" + +from dataclasses import dataclass + +import pytest +import torch +from defs.examples.visual_gen.visual_gen_test_utils import ( + FeatureConfigState, + _assert_feature_quantization_installed, + _assert_lpips_below_threshold, + _assert_resolved_single_device_feature_config, + _assert_single_device_feature_executed, + _build_single_device_feature_args, + _cleanup_cuda, + _cleanup_single_device_feature_pipeline, + _disable_inductor_compile_worker_quiesce, + _fixed_nvfp4_quantization_backend, + _golden_media_path, + _lpips_deterministic_algorithms, + _lpips_model_path, + _preserve_lpips_candidate_on_failure, + _run_lpips_eval, + _run_reusable_image_lpips_eval, + _run_single_device_feature_generator, + _skip_if_missing, + _validate_single_feature_config, +) + +GLM_IMAGE_CHECKPOINT_SUBDIR = "GLM-Image" +GLM_IMAGE_LPIPS_PROMPT = "a tiny astronaut hatching from an egg on the moon" +# GLM's native resolution; the AR prior stage degrades badly below it. +GLM_IMAGE_LPIPS_HEIGHT = 1024 +GLM_IMAGE_LPIPS_WIDTH = 1024 +GLM_IMAGE_LPIPS_NUM_INFERENCE_STEPS = 30 +GLM_IMAGE_LPIPS_GUIDANCE_SCALE = 1.5 +GLM_IMAGE_LPIPS_SEED = 42 +GLM_IMAGE_LPIPS_THRESHOLD = 0.05 +GLM_IMAGE_FEATURE_LPIPS_THRESHOLD = 0.05 +GLM_IMAGE_SUPPORTED_FEATURES = frozenset({"fp8-blockwise", "nvfp4"}) + + +@dataclass(frozen=True) +class GlmImageFeatureProfile: + id: str + features: FeatureConfigState + + +@dataclass(frozen=True) +class GlmImageAccuracyCase: + id: str + checkpoint_subdir: str + golden_file: str + features: FeatureConfigState + lpips_threshold: float + + +GLM_IMAGE_FEATURE_PROFILES = ( + GlmImageFeatureProfile( + id="fp8-blockwise", + features=FeatureConfigState(quantization="FP8_BLOCK_SCALES"), + ), + GlmImageFeatureProfile( + id="nvfp4", + features=FeatureConfigState(quantization="NVFP4"), + ), +) + + +def _build_glm_image_accuracy_cases(): + cases = [] + for profile in GLM_IMAGE_FEATURE_PROFILES: + _validate_single_feature_config( + profile.features, + GLM_IMAGE_SUPPORTED_FEATURES, + "GLM-Image", + ) + case_id = profile.id + cases.append( + pytest.param( + GlmImageAccuracyCase( + id=case_id, + checkpoint_subdir=GLM_IMAGE_CHECKPOINT_SUBDIR, + golden_file=f"glm_image_{profile.id.replace('-', '_')}_lpips_golden.png", + features=profile.features, + lpips_threshold=GLM_IMAGE_FEATURE_LPIPS_THRESHOLD, + ), + id=case_id, + ) + ) + return cases + + +GLM_IMAGE_ACCURACY_CASES = _build_glm_image_accuracy_cases() + + +def _glm_image_generator(device): + # GlmImagePipeline.forward takes a generator rather than a seed + return torch.Generator(device=device).manual_seed(GLM_IMAGE_LPIPS_SEED) + + +def _generate_glm_image_lpips_image(model_path, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.media.encoding import save_image + from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs + + _skip_if_missing(model_path, "GLM-Image checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + with _lpips_deterministic_algorithms(): + args = VisualGenArgs( + model=model_path, + torch_compile_config=TorchCompileConfig(enable=False), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + result = pipeline.forward( + prompt=GLM_IMAGE_LPIPS_PROMPT, + height=GLM_IMAGE_LPIPS_HEIGHT, + width=GLM_IMAGE_LPIPS_WIDTH, + num_inference_steps=GLM_IMAGE_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=GLM_IMAGE_LPIPS_GUIDANCE_SCALE, + generator=_glm_image_generator(pipeline.device), + ) + generated_image = result.image[0].detach().cpu() + finally: + del pipeline + _cleanup_cuda() + + save_image(generated_image, output_path) + + +def _generate_glm_image_feature_image(case, output_path): + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.media.encoding import save_image + + model_path = _lpips_model_path(case.checkpoint_subdir) + _skip_if_missing(model_path, f"{case.checkpoint_subdir} checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + pipeline = None + with _lpips_deterministic_algorithms(), _fixed_nvfp4_quantization_backend(case.features): + args = _build_single_device_feature_args( + model_path, + case.features, + resolution=(GLM_IMAGE_LPIPS_HEIGHT, GLM_IMAGE_LPIPS_WIDTH), + num_frames=1, + ) + try: + pipeline = PipelineLoader(args).load(skip_warmup=False) + _assert_resolved_single_device_feature_config( + pipeline, + case.features, + resolution=(GLM_IMAGE_LPIPS_HEIGHT, GLM_IMAGE_LPIPS_WIDTH), + num_frames=1, + ) + _assert_feature_quantization_installed(pipeline, case.features) + result = pipeline.forward( + prompt=GLM_IMAGE_LPIPS_PROMPT, + height=GLM_IMAGE_LPIPS_HEIGHT, + width=GLM_IMAGE_LPIPS_WIDTH, + num_inference_steps=GLM_IMAGE_LPIPS_NUM_INFERENCE_STEPS, + guidance_scale=GLM_IMAGE_LPIPS_GUIDANCE_SCALE, + generator=_glm_image_generator(pipeline.device), + ) + _assert_single_device_feature_executed(pipeline, case.features) + generated_image = result.image[0].detach().cpu() + finally: + try: + if pipeline is not None: + _cleanup_single_device_feature_pipeline(pipeline) + del pipeline + finally: + _cleanup_cuda() + + save_image(generated_image, output_path) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_glm_image_lpips_against_golden(tmp_path): + generated_path = tmp_path / "glm_image_generated.png" + golden_path = _golden_media_path( + tmp_path, "glm_image_lpips_golden.png", "GLM-Image LPIPS golden image" + ) + _generate_glm_image_lpips_image(_lpips_model_path(GLM_IMAGE_CHECKPOINT_SUBDIR), generated_path) + score = _run_lpips_eval( + tmp_path, + "glm_image", + "image", + GLM_IMAGE_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _assert_lpips_below_threshold(score, GLM_IMAGE_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("case", GLM_IMAGE_ACCURACY_CASES) +def test_glm_image_feature_accuracy_against_golden( + request, + tmp_path, + case, + _visual_gen_lpips_scorer, +): + generated_path = tmp_path / f"glm_image_{case.id}_generated.png" + reference_path = _golden_media_path( + tmp_path, + case.golden_file, + f"{case.checkpoint_subdir} {case.id} LPIPS golden image", + ) + + _run_single_device_feature_generator( + case.features, _generate_glm_image_feature_image, case, generated_path + ) + score = _run_reusable_image_lpips_eval( + f"glm_image_{case.id}", + reference_path, + generated_path, + _visual_gen_lpips_scorer, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + case.lpips_threshold, + generated_path, + f"glm_image_{case.id}_generated.png", + ) + _assert_lpips_below_threshold(score, case.lpips_threshold) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index c0e6548c4bda..a5440bccbe82 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -367,6 +367,7 @@ l0_b200: - examples/visual_gen/test_visual_gen_wan.py::test_visual_gen_api_walkthrough - examples/visual_gen/test_visual_gen_flux.py::test_flux1_lpips_against_golden - examples/visual_gen/test_visual_gen_flux.py::test_flux2_lpips_against_golden + - examples/visual_gen/test_visual_gen_glm.py::test_glm_image_lpips_against_golden - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-fp8-blockwise] - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-nvfp4] - examples/visual_gen/test_visual_gen_flux.py::test_flux_accuracy_against_golden[flux1-cuda-graph] @@ -378,6 +379,8 @@ l0_b200: - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwenimage_feature_accuracy_against_golden[cuda-graph] - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[fp8-blockwise] - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[nvfp4] + - examples/visual_gen/test_visual_gen_glm.py::test_glm_image_feature_accuracy_against_golden[fp8-blockwise] + - examples/visual_gen/test_visual_gen_glm.py::test_glm_image_feature_accuracy_against_golden[nvfp4] - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[fp8-blockwise] - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[nvfp4] - examples/visual_gen/test_visual_gen_ltx2.py::test_ltx2_feature_accuracy_against_golden[cuda-graph] diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py index 9b4730c3ba3b..3131f09e6f41 100644 --- a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py @@ -75,13 +75,16 @@ def _resolve_glm_checkpoint() -> str: # GlmImage takes a plain text prompt and derives the latent grid from the # explicit height/width (each must be divisible by 32). 256 == 8 * 32. -PROMPT = "A dinosaur walking through the jungle" +PROMPT = "a tiny astronaut hatching from an egg on the moon" +BATCH_PROMPTS = ["a sunset over mountains", "a cat on a roof"] +# Quoted text routes through the glyph encoder, which produces the text attention mask. +GLYPH_PROMPT = "A sign that says 'OPEN'" HEIGHT = 256 WIDTH = 256 NUM_STEPS = 30 SEED = 42 GUIDANCE_SCALE = 1.5 -COS_SIM_THRESHOLD = 0.99 +COS_SIM_THRESHOLD = 0.999 # ============================================================================ @@ -290,7 +293,7 @@ def test_single_prompt(self): def test_batch_prompts(self): """A list of prompts returns one image per prompt in a single batched forward.""" - prompts = [PROMPT, "A neon city skyline reflected in a river at night"] + prompts = BATCH_PROMPTS pipe = None try: pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH) @@ -313,6 +316,57 @@ def test_batch_prompts(self): finally: _teardown_pipeline(pipe) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_quoted_prompt(self): + """A quoted prompt exercises the glyph path and its text attention mask.""" + + pipe = None + try: + pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH) + result = _capture_trtllm_image( + pipe, + prompt=GLYPH_PROMPT, + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_STEPS, + guidance_scale=GUIDANCE_SCALE, + seed=SEED, + ) + images = _to_image_tensor(result.image) + assert images.shape == (HEIGHT, WIDTH, 3), ( + f"Expected (H,W,C)=({HEIGHT},{WIDTH},3), got {tuple(images.shape)}" + ) + finally: + _teardown_pipeline(pipe) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_num_images_per_prompt(self): + """num_images_per_prompt returns that many images for a single prompt.""" + + num_images_per_prompt = 2 + pipe = None + try: + pipe = _load_trtllm_pipeline(GLM_IMAGE_PATH) + with torch.no_grad(): + result = pipe.forward( + prompt=PROMPT, + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_STEPS, + guidance_scale=GUIDANCE_SCALE, + num_images_per_prompt=num_images_per_prompt, + generator=torch.Generator(device="cuda").manual_seed(SEED), + ) + image = result.image + assert image.shape[0] == num_images_per_prompt, ( + f"Expected batch dim {num_images_per_prompt}, got {image.shape[0]}" + ) + assert tuple(image.shape[1:]) == (HEIGHT, WIDTH, 3), ( + f"Expected (H,W,C)=({HEIGHT},{WIDTH},3), got {tuple(image.shape[1:])}" + ) + finally: + _teardown_pipeline(pipe) + def test_image_conditioning_not_supported(): """Passing a condition image raises NotImplementedError (I2I lands in a follow-up MR).""" diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py b/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py index af9c7efd5746..95dc2ceba8db 100644 --- a/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_glm_image_transformer.py @@ -257,7 +257,7 @@ def test_glm_image_allclose_to_hf(self): print(f" Cosine similarity: {cos_sim:.6f}") print(f" Max diff: {max_diff:.6f}") - self.assertGreater(cos_sim, 0.99, f"Cosine similarity too low: {cos_sim}") + self.assertGreater(cos_sim, 0.999, f"Cosine similarity too low: {cos_sim}") if __name__ == "__main__": From 934087a1669cfce1f6fca0bd748c7ee753dad215 Mon Sep 17 00:00:00 2001 From: jloftin Date: Thu, 13 Aug 2026 22:50:51 +0000 Subject: [PATCH 05/11] fix threshold markdown Signed-off-by: jloftin --- docs/source/models/visual-generation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 20342fdf73d8..784b995b3c04 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -82,7 +82,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^7]: `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` — a distilled version of Wan2.2-TI2V-5B with 3 denoising steps. CFG parallelism, TeaCache, and Cache-DiT are not applicable. -[^9]: GlmImage currently supports single-GPU text-to-image with BF16 parity vs `diffusers` (cosine >= 0.99 on the full transformer). FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints. Image-to-image conditioning, sequence/CFG parallelism, parallel VAE, and caching (TeaCache / Cache-DiT) are not yet supported. +[^9]: GlmImage currently supports single-GPU text-to-image with BF16 parity vs `diffusers` (cosine >= 0.999 on the full transformer). FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints. Image-to-image conditioning, sequence/CFG parallelism, parallel VAE, and caching (TeaCache / Cache-DiT) are not yet supported. ## Quick Start From 0ac489fe97bc40140e0412303b7905d186ba87ef Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Fri, 14 Aug 2026 11:48:48 -0700 Subject: [PATCH 06/11] Update docs/source/models/visual-generation.md Co-authored-by: Zhenhua Wang <4936589+zhenhuaw-me@users.noreply.github.com> Signed-off-by: Joseph Loftin --- docs/source/models/visual-generation.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 784b995b3c04..321bf58f055a 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -80,10 +80,6 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^6]: Qwen-Image-Layered supports baseline BF16 image-conditioned layer decomposition and returns the generated RGBA layer stack as a saveable image grid. FP8 blockwise, NVFP4, cache acceleration, attention-parallel/Sage/VSA backends, Tensor Parallelism, and `trtllm-serve` image-edit routing are not enabled for this pipeline yet. -[^7]: `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` — a distilled version of Wan2.2-TI2V-5B with 3 denoising steps. CFG parallelism, TeaCache, and Cache-DiT are not applicable. - -[^9]: GlmImage currently supports single-GPU text-to-image with BF16 parity vs `diffusers` (cosine >= 0.999 on the full transformer). FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints. Image-to-image conditioning, sequence/CFG parallelism, parallel VAE, and caching (TeaCache / Cache-DiT) are not yet supported. - ## Quick Start Here is a simple example to generate a video with Wan 2.1: From 9f6762878fe4499a9ed42d87d6687e2cc30c83e9 Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Fri, 14 Aug 2026 11:48:58 -0700 Subject: [PATCH 07/11] Update docs/source/models/visual-generation.md Co-authored-by: Zhenhua Wang <4936589+zhenhuaw-me@users.noreply.github.com> Signed-off-by: Joseph Loftin --- docs/source/models/visual-generation.md | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 321bf58f055a..9a01a830dfae 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -68,7 +68,6 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **Qwen-Image-Edit-2511** | Yes | Yes | No | No | Yes | No | No | Yes | Yes | No | No | No | No | No | | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | | **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | Yes | No | No | No | No | -| **GlmImage** [^9] | Yes | Yes | No | No | No | No | No | No | No | Yes | No | No | No | No | [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable. From 34f2201d630cfceed01f214ede0e859d86573087 Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Fri, 14 Aug 2026 11:56:55 -0700 Subject: [PATCH 08/11] Update docs/source/models/visual-generation.md Co-authored-by: Zhenhua Wang <4936589+zhenhuaw-me@users.noreply.github.com> Signed-off-by: Joseph Loftin --- docs/source/models/visual-generation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 9a01a830dfae..ed9ab93666c4 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -61,7 +61,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | | **Wan 2.1 VSA** [^2] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes | | **Wan 2.2** | Yes | Yes | Yes [^3] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | -| **FastWan 2.2** | Yes | Yes | No | No | No [^7] | No | No | Yes | Yes | Yes | No | No | No | No | +| **FastWan 2.2** | Yes | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | | **LTX-2** | Yes | Yes | Yes [^4] | Yes | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image-Layered** [^6] | No | No | No | No | No | No | No | Yes | Yes | No | No | No | No | No | From 988a243a6fab6c6aea112b701c7f50fa1cadf44d Mon Sep 17 00:00:00 2001 From: jloftin Date: Fri, 14 Aug 2026 23:02:08 +0000 Subject: [PATCH 09/11] Review 2 Signed-off-by: jloftin --- .../visual_gen/serve/configs/glm_image.yml | 18 - .../models/glm_image/pipeline_glm_image.py | 361 +++--------------- .../models/glm_image/transformer_glm_image.py | 61 +-- .../visual_gen/test_glm_image_pipeline.py | 23 +- 4 files changed, 119 insertions(+), 344 deletions(-) delete mode 100644 examples/visual_gen/serve/configs/glm_image.yml diff --git a/examples/visual_gen/serve/configs/glm_image.yml b/examples/visual_gen/serve/configs/glm_image.yml deleted file mode 100644 index c91cbe130a55..000000000000 --- a/examples/visual_gen/serve/configs/glm_image.yml +++ /dev/null @@ -1,18 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# GlmImage text-to-image (single GPU). -parallel_config: - cfg_size: 1 - ulysses_size: 1 diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py index e392b7a6f565..c4e26f2f209d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py @@ -22,8 +22,7 @@ import numpy as np import PIL import torch -from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler -from diffusers.image_processor import VaeImageProcessor +from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SchedulerMixin from diffusers.utils.torch_utils import randn_tensor from transformers import ( ByT5Tokenizer, @@ -40,36 +39,19 @@ from .transformer_glm_image import GlmImageTransformer2DModel + # ------------------------------------------------------------------ # HF Port # ------------------------------------------------------------------ -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents - - -def retrieve_latents( - encoder_output: torch.Tensor, - generator: Optional[torch.Generator] = None, - sample_mode: str = "sample", -): - 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") - - # Copied from diffusers.pipelines.cogview4.pipeline_cogview4.retrieve_timesteps def retrieve_timesteps( - scheduler, + scheduler: SchedulerMixin, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, - **kwargs, -): + **kwargs: Any, +) -> Tuple[torch.Tensor, int]: r""" Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. @@ -132,7 +114,7 @@ def retrieve_timesteps( def calculate_shift( - image_seq_len, + image_seq_len: int, base_seq_len: int = 256, base_shift: float = 0.25, max_shift: float = 0.75, @@ -151,75 +133,14 @@ class GlmImagePipeline(BasePipeline): # ------------------------------------------------------------------ # HF Port # ------------------------------------------------------------------ - @staticmethod - def _validate_and_normalize_images( - image: Union[List[PIL.Image.Image], List[List[PIL.Image.Image]]], - batch_size: int, - ) -> List[List[PIL.Image.Image]]: - """ - Validate and normalize image inputs to List[List[PIL.Image]]. - - Rules: - - batch_size > 1: Only accepts List[List[PIL.Image]], each sublist must have equal length - - batch_size == 1: Accepts List[PIL.Image] for legacy compatibility (converted to [[img1, img2, ...]]) - - Other formats raise ValueError - - Args: - image: Input images in various formats - batch_size: Number of prompts in the batch - - Returns: - Normalized images as List[List[PIL.Image]], or None if no images provided - """ - if image is None or len(image) == 0: - return None - - first_element = image[0] - - if batch_size == 1: - # Legacy format: List[PIL.Image] -> [[img1, img2, ...]] - if not isinstance(first_element, (list, tuple)): - return [list(image)] - # Already in List[List[PIL.Image]] format - if len(image) != 1: - raise ValueError( - f"For batch_size=1 with List[List[PIL.Image]] format, expected 1 image list, got {len(image)}." - ) - return [list(image[0])] - - # batch_size > 1: must be List[List[PIL.Image]] - if not isinstance(first_element, (list, tuple)): - raise ValueError( - f"For batch_size > 1, images must be List[List[PIL.Image]] format. " - f"Got List[{type(first_element).__name__}] instead. " - f"Each prompt requires its own list of condition images." - ) - - if len(image) != batch_size: - raise ValueError( - f"Number of image lists ({len(image)}) must match batch size ({batch_size})." - ) - - # Validate homogeneous: all sublists must have same length - num_input_images_per_prompt = len(image[0]) - for idx, imgs in enumerate(image): - if len(imgs) != num_input_images_per_prompt: - raise ValueError( - f"All prompts must have the same number of condition images. " - f"Prompt 0 has {num_input_images_per_prompt} images, but prompt {idx} has {len(imgs)} images." - ) - - return [list(imgs) for imgs in image] - def generate_prior_tokens( self, prompt: Union[str, List[str]], height: int, width: int, - image: Optional[List[List[PIL.Image.Image]]] = None, device: Optional[torch.device] = None, generator: Optional[torch.Generator] = None, - ): + ) -> torch.Tensor: """ Generate prior tokens for the DiT model using the AR model. @@ -227,18 +148,11 @@ def generate_prior_tokens( prompt: Single prompt or list of prompts height: Target image height width: Target image width - image: Normalized image input as List[List[PIL.Image]]. Should be pre-validated - using _validate_and_normalize_images() before calling this method. device: Target device generator: Random generator for reproducibility Returns: - Tuple of: - - prior_token_ids: Tensor of shape (batch_size, num_tokens) with upsampled prior tokens - - prior_token_image_ids_per_sample: List of tensors, one per sample. Each tensor contains - the upsampled prior token ids for all condition images in that sample. None for t2i. - - source_image_grid_thw_per_sample: List of tensors, one per sample. Each tensor has shape - (num_condition_images, 3) with upsampled grid info. None for t2i. + prior_token_ids: Tensor of shape (batch_size, num_tokens) with upsampled prior tokens """ device = device or self._execution_device @@ -246,17 +160,10 @@ def generate_prior_tokens( prompt_list = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt_list) - # Image is already normalized by _validate_and_normalize_images(): None or List[List[PIL.Image]] - is_text_to_image = image is None # Build messages for each sample in the batch - all_messages = [] - for idx, p in enumerate(prompt_list): - content = [] - if not is_text_to_image: - for img in image[idx]: - content.append({"type": "image", "image": img}) - content.append({"type": "text", "text": p}) - all_messages.append([{"role": "user", "content": content}]) + all_messages = [ + [{"role": "user", "content": [{"type": "text", "text": p}]}] for p in prompt_list + ] # Process with the processor (supports batch with left padding) inputs = self.processor.apply_chat_template( all_messages, @@ -271,8 +178,6 @@ def generate_prior_tokens( image_grid_thw = inputs.get("image_grid_thw") images_per_sample = inputs.get("images_per_sample") - # Determine number of condition images and grids per sample - num_condition_images = 0 if is_text_to_image else len(image[0]) if images_per_sample is not None: num_grids_per_sample = images_per_sample[0].item() else: @@ -282,45 +187,9 @@ def generate_prior_tokens( # Compute generation params (same for all samples in homogeneous batch) first_sample_grids = image_grid_thw[:num_grids_per_sample] max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params( - image_grid_thw=first_sample_grids, is_text_to_image=is_text_to_image + image_grid_thw=first_sample_grids ) - # Generate source image tokens (prior_token_image_ids) for i2i mode - prior_token_image_ids = None - source_image_grid_thw = None - if not is_text_to_image: - # Extract source grids by selecting condition image indices (skip target grids) - # Grid order from processor: [s0_cond1, s0_cond2, ..., s0_target, s1_cond1, s1_cond2, ..., s1_target, ...] - # We need indices: [0, 1, ..., num_condition_images-1, num_grids_per_sample, num_grids_per_sample+1, ...] - source_indices = [] - for sample_idx in range(batch_size): - base = sample_idx * num_grids_per_sample - source_indices.extend(range(base, base + num_condition_images)) - source_grids = image_grid_thw[source_indices] - - if len(source_grids) > 0: - prior_token_image_embed = self.vision_language_encoder.get_image_features( - inputs["pixel_values"], source_grids - ).pooler_output - prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0) - prior_token_image_ids_d32 = self.vision_language_encoder.get_image_tokens( - prior_token_image_embed, source_grids - ) - # Upsample each source image's prior tokens to match VAE/DiT resolution - split_sizes = source_grids.prod(dim=-1).tolist() - prior_ids_per_source = torch.split(prior_token_image_ids_d32, split_sizes) - upsampled_prior_ids = [] - for i, prior_ids in enumerate(prior_ids_per_source): - t, h, w = source_grids[i].tolist() - upsampled = self._upsample_token_ids(prior_ids, int(h), int(w)) - upsampled_prior_ids.append(upsampled.squeeze(0)) - prior_token_image_ids = torch.cat(upsampled_prior_ids, dim=0) - # Upsample grid dimensions for later splitting - upsampled_grids = source_grids.clone() - upsampled_grids[:, 1] = upsampled_grids[:, 1] * 2 - upsampled_grids[:, 2] = upsampled_grids[:, 2] * 2 - source_image_grid_thw = upsampled_grids - # Generate with AR model # Set torch random seed from generator for reproducibility # (transformers generate() doesn't accept generator parameter) @@ -350,28 +219,7 @@ def generate_prior_tokens( ) prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w) all_prior_token_ids.append(prior_token_ids) - prior_token_ids = torch.cat(all_prior_token_ids, dim=0) - - # Split prior_token_image_ids and source_image_grid_thw into per-sample lists for easier consumption - prior_token_image_ids_per_sample = None - source_image_grid_thw_per_sample = None - if prior_token_image_ids is not None and source_image_grid_thw is not None: - # Split grids: each sample has num_condition_images grids - source_image_grid_thw_per_sample = list( - torch.split(source_image_grid_thw, num_condition_images) - ) - # Split prior_token_image_ids: tokens per sample may vary due to different image sizes - tokens_per_image = source_image_grid_thw.prod(dim=-1).tolist() - tokens_per_sample = [] - for i in range(batch_size): - start_idx = i * num_condition_images - end_idx = start_idx + num_condition_images - tokens_per_sample.append(sum(tokens_per_image[start_idx:end_idx])) - prior_token_image_ids_per_sample = list( - torch.split(prior_token_image_ids, tokens_per_sample) - ) - - return prior_token_ids, prior_token_image_ids_per_sample, source_image_grid_thw_per_sample + return torch.cat(all_prior_token_ids, dim=0) def encode_prompt( self, @@ -383,7 +231,7 @@ def encode_prompt( device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, max_sequence_length: int = 2048, - ): + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: r""" Encodes the prompt into text encoder hidden states. @@ -440,15 +288,15 @@ def encode_prompt( def prepare_latents( self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - ): + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]], + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: if latents is not None: return latents.to(device) @@ -467,10 +315,7 @@ def prepare_latents( return latents @staticmethod - def _compute_generation_params( - image_grid_thw, - is_text_to_image: bool, - ): + def _compute_generation_params(image_grid_thw: torch.Tensor) -> Tuple[int, int, int, int]: grid_sizes = [] grid_hw = [] @@ -479,15 +324,9 @@ def _compute_generation_params( grid_sizes.append(int(h * w)) grid_hw.append((int(h), int(w))) - if not is_text_to_image: - max_new_tokens = grid_sizes[-1] + 1 - large_image_start_offset = 0 - target_grid_h, target_grid_w = grid_hw[-1] - else: - 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] + max_new_tokens = sum(grid_sizes) + 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 @@ -512,7 +351,7 @@ def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> return token_ids @property - def do_classifier_free_guidance(self): + def do_classifier_free_guidance(self) -> bool: return self._guidance_scale > 1 def _get_glyph_embeds( @@ -521,7 +360,7 @@ def _get_glyph_embeds( max_sequence_length: int = 2048, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, - ): + ) -> torch.Tensor: """Get glyph embeddings for each prompt in the batch.""" device = device or self._execution_device dtype = dtype or self.text_encoder.dtype @@ -579,7 +418,7 @@ def _get_glyph_embeds( glyph_embeds = torch.cat(padded_embeds, dim=0) return glyph_embeds.to(device=device, dtype=dtype) - def get_glyph_texts(self, prompt): + def get_glyph_texts(self, prompt: Union[str, List[str]]) -> List[List[str]]: """Extract glyph texts from prompt(s). Returns a list of lists for batch processing.""" if isinstance(prompt, str): prompt = [prompt] @@ -595,13 +434,13 @@ def get_glyph_texts(self, prompt): return all_ocr_texts @property - def guidance_scale(self): + def guidance_scale(self) -> float: return self._guidance_scale # ------------------------------------------------------------------ # TRT-LLM # ------------------------------------------------------------------ - def __init__(self, pipeline_config: DiffusionPipelineConfig): + def __init__(self, pipeline_config: DiffusionPipelineConfig) -> None: super().__init__(pipeline_config) def load_standard_components( @@ -609,7 +448,7 @@ def load_standard_components( checkpoint_dir: str, device: torch.device, skip_components: Optional[list] = None, - **kwargs, + **kwargs: Any, ) -> None: skip_components = skip_components or [] @@ -638,7 +477,6 @@ def load_standard_components( ).to(device) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) # Scheduler (FlowMatchEulerDiscreteScheduler) if PipelineComponent.SCHEDULER not in skip_components: @@ -663,13 +501,13 @@ def load_standard_components( ).to(device) @property - def device(self): + def device(self) -> torch.device: if self.transformer is not None: return next(self.transformer.parameters()).device return torch.device("cuda:0") @property - def dtype(self): + def dtype(self) -> torch.dtype: return self.pipeline_config.torch_dtype @property @@ -694,17 +532,19 @@ def default_warmup_num_frames(self) -> List[int]: @property def resolution_multiple_of(self) -> Tuple[int, int]: + # vae_scale_factor * patch_size is the 16px latent pitch; the AR prior grid is + # generated at d32 and upsampled 2x, so the resolution must clear 32. patch_size = self.transformer.config.patch_size if self.transformer is not None else 2 - multiple = getattr(self, "vae_scale_factor", 16) * patch_size + multiple = getattr(self, "vae_scale_factor", 8) * patch_size * 2 return (multiple, multiple) - def infer(self, req): + def infer(self, req: Any) -> PipelineOutput: """Run inference from DiffusionRequest.""" params = req.params if getattr(params, "image", None) is not None: raise NotImplementedError( - "image-to-image conditioning is not yet supported by the " - "TensorRT-LLM GlmImage pipeline; coming in a follow-up MR" + "The TensorRT-LLM GlmImage pipeline is text-to-image only; " + "image-to-image conditioning is not available." ) generator = None if params.seed is not None: @@ -756,8 +596,6 @@ def forward( prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, prior_token_ids: Optional[torch.Tensor] = None, - prior_token_image_ids: Optional[List[torch.Tensor]] = None, - source_image_grid_thw: Optional[List[torch.Tensor]] = None, crops_coords_top_left: Tuple[int, int] = (0, 0), attention_kwargs: Optional[Dict[str, Any]] = None, max_sequence_length: int = 2048, @@ -789,8 +627,8 @@ def forward( """ if image is not None: raise NotImplementedError( - "image-to-image conditioning is not yet supported by the " - "TensorRT-LLM GlmImage pipeline; coming in a follow-up MR" + "The TensorRT-LLM GlmImage pipeline is text-to-image only; " + "image-to-image conditioning is not available." ) pipeline_start = time.time() @@ -806,6 +644,7 @@ def forward( height = self.default_generation_params["height"] if width is None: width = self.default_generation_params["width"] + self.validate_resolution(height, width, 1) if prompt is not None and isinstance(prompt, str): batch_size = 1 @@ -816,51 +655,20 @@ def forward( device = self.device - # 1. Validate and normalize image format - normalized_image = self._validate_and_normalize_images(image, batch_size) - - # 2. Generate prior tokens (batch mode) + # 1. Generate prior tokens (batch mode) # Get a single generator for AR model (use first if list provided) logger.info("Generating prior tokens...") ar_generator = generator[0] if isinstance(generator, list) else generator if prior_token_ids is None: - prior_token_ids, prior_token_image_ids_per_sample, source_image_grid_thw_per_sample = ( - self.generate_prior_tokens( - prompt=prompt, - image=normalized_image, - height=height, - width=width, - device=device, - generator=ar_generator, - ) + prior_token_ids = self.generate_prior_tokens( + prompt=prompt, + height=height, + width=width, + device=device, + generator=ar_generator, ) - else: - # User provided prior_token_ids directly (from generate_prior_tokens) - prior_token_image_ids_per_sample = prior_token_image_ids - source_image_grid_thw_per_sample = source_image_grid_thw - - # 3. Preprocess images for VAE encoding - preprocessed_images = None - if normalized_image is not None: - preprocessed_images = [] - for prompt_images in normalized_image: - prompt_preprocessed = [] - for img in prompt_images: - image_height, image_width = ( - img.size[::-1] if isinstance(img, PIL.Image.Image) else img.shape[:2] - ) - multiple_of = self.vae_scale_factor * self.transformer.config.patch_size - image_height = (image_height // multiple_of) * multiple_of - image_width = (image_width // multiple_of) * multiple_of - img = self.image_processor.preprocess( - img, height=image_height, width=image_width - ) - prompt_preprocessed.append(img) - height = height or image_height - width = width or image_width - preprocessed_images.append(prompt_preprocessed) - - # 4. Encode input prompt + + # 2. Encode input prompt logger.info("Encoding prompt...") prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt, @@ -873,7 +681,7 @@ def forward( dtype=self.dtype, ) - # 5. Prepare latents + # 3. Prepare latents latent_channels = self.transformer.config.in_channels latents = self.prepare_latents( batch_size=batch_size * num_images_per_prompt, @@ -886,50 +694,7 @@ def forward( latents=latents, ) - if normalized_image is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view( - 1, self.vae.config.latent_channels, 1, 1 - ) - latents_std = torch.tensor(self.vae.config.latents_std).view( - 1, self.vae.config.latent_channels, 1, 1 - ) - - latents_mean = latents_mean.to(device=device, dtype=prompt_embeds.dtype) - latents_std = latents_std.to(device=device, dtype=prompt_embeds.dtype) - - # Process each sample's condition images - for prompt_idx in range(batch_size): - prompt_images = preprocessed_images[prompt_idx] - prompt_prior_ids = prior_token_image_ids_per_sample[prompt_idx] - prompt_grid_thw = source_image_grid_thw_per_sample[prompt_idx] - - # Split this sample's prior_token_image_ids by each image's token count - split_sizes = prompt_grid_thw.prod(dim=-1).tolist() - prior_ids_per_image = torch.split(prompt_prior_ids, split_sizes) - # Process each condition image for this sample - for condition_image, condition_image_prior_token_id in zip( - prompt_images, prior_ids_per_image - ): - condition_image = condition_image.to(device=device, dtype=prompt_embeds.dtype) - condition_latent = retrieve_latents( - self.vae.encode(condition_image), generator=generator, sample_mode="argmax" - ) - condition_latent = (condition_latent - latents_mean) / latents_std - - _ = self.transformer( - hidden_states=condition_latent, - encoder_hidden_states=torch.zeros_like(prompt_embeds)[:1, :0, ...], - prior_token_id=condition_image_prior_token_id, - prior_token_drop=torch.full_like( - condition_image_prior_token_id, False, dtype=torch.bool - ), - timestep=torch.zeros((1,), device=device), - target_size=torch.tensor([condition_image.shape[-2:]], device=device), - crop_coords=torch.zeros((1, 2), device=device), - attention_kwargs=attention_kwargs, - ) - - # 6. Prepare additional timestep conditions + # 4. Prepare additional timestep conditions target_size = (height, width) target_size = torch.tensor([target_size], dtype=prompt_embeds.dtype, device=device) crops_coords_top_left = torch.tensor( @@ -963,7 +728,7 @@ def forward( ) self._num_timesteps = len(timesteps) - # 7. Denoising loop + # 5. Denoising loop transformer_dtype = self.dtype # Repeat prior_token_ids for num_images_per_prompt @@ -984,13 +749,13 @@ def forward( ) = self._align_cfg_embeds(prompt_embeds, negative_prompt_embeds) def forward_fn( - latents, - extra_stream_latents, - step_index, - timestep, - encoder_hidden_states, - extra_tensors, - ): + latents: torch.Tensor, + extra_stream_latents: Optional[torch.Tensor], + step_index: int, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + extra_tensors: Dict[str, Any], + ) -> torch.Tensor: """Forward function for GlmImage transformer.""" return self.transformer( hidden_states=latents.to(transformer_dtype), diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py index 2519d6574d2b..fdbbd131e8fd 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py @@ -29,7 +29,6 @@ from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization import DynamicLinearWeightLoader -from tensorrt_llm._torch.visual_gen.utils import SequenceSharder from tensorrt_llm.models.modeling_utils import QuantConfig @@ -63,7 +62,7 @@ def __init__( approximate: str = "none", bias: bool = True, model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() self.proj = _aux_linear(dim_in, dim_out, bias=bias, model_config=model_config) self.approximate = approximate @@ -71,7 +70,7 @@ def __init__( def gelu(self, gate: torch.Tensor) -> torch.Tensor: return F.gelu(gate, approximate=self.approximate) - def forward(self, hidden_states): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.proj(hidden_states) hidden_states = self.gelu(hidden_states) return hidden_states @@ -84,12 +83,12 @@ def __init__( dim_out: int, bias: bool = True, model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() self.proj = _aux_linear(dim_in, dim_out, bias=bias, model_config=model_config) self.activation = F.silu - def forward(self, hidden_states): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.proj(hidden_states) return self.activation(hidden_states) @@ -102,10 +101,10 @@ def __init__( mult: int = 4, dropout: float = 0.0, activation_fn: str = "geglu", - inner_dim=None, + inner_dim: Optional[int] = None, bias: bool = True, model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() if inner_dim is None: inner_dim = int(dim * mult) @@ -140,14 +139,14 @@ def __init__( time_embed_dim: int, out_dim: Optional[int] = None, model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() self.linear_1 = _aux_linear(in_channels, time_embed_dim, model_config=model_config) self.act = torch.nn.SiLU() time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim self.linear_2 = _aux_linear(time_embed_dim, time_embed_dim_out, model_config=model_config) - def forward(self, sample): + def forward(self, sample: torch.Tensor) -> torch.Tensor: sample = self.linear_1(sample) sample = self.act(sample) sample = self.linear_2(sample) @@ -157,12 +156,12 @@ def forward(self, sample): class GlmImagePixArtAlphaTextProjection(torch.nn.Module): def __init__( self, - in_features, - hidden_size, - out_features=None, - act_fn="gelu_tanh", + in_features: int, + hidden_size: int, + out_features: Optional[int] = None, + act_fn: str = "gelu_tanh", model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() if out_features is None: out_features = hidden_size @@ -175,7 +174,7 @@ def __init__( raise ValueError(f"Unknown activation function: {act_fn}") self.linear_2 = _aux_linear(hidden_size, out_features, bias=True, model_config=model_config) - def forward(self, caption): + def forward(self, caption: torch.Tensor) -> torch.Tensor: hidden_states = self.linear_1(caption) hidden_states = self.act_1(hidden_states) hidden_states = self.linear_2(hidden_states) @@ -190,7 +189,7 @@ def __init__( pooled_projection_dim: int, timesteps_dim: int = 256, model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() self.time_proj = Timesteps( @@ -237,7 +236,7 @@ def __init__( hidden_size: int = 2560, patch_size: int = 2, model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() self.patch_size = patch_size @@ -277,7 +276,7 @@ def __init__( def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> Tuple[torch.Tensor, ...]: dtype = hidden_states.dtype norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) @@ -327,7 +326,7 @@ def __init__( bias: bool = True, norm_type: str = "layer_norm", model_config: Optional[DiffusionModelConfig] = None, - ): + ) -> None: super().__init__() self.linear = _aux_linear( conditioning_embedding_dim, embedding_dim * 2, bias=bias, model_config=model_config @@ -355,7 +354,7 @@ def __init__( dtype: Optional[torch.dtype] = None, config: Optional[DiffusionModelConfig] = None, layer_idx: int = 0, - ): + ) -> None: config = config or DiffusionModelConfig() super().__init__( num_attention_heads=num_attention_heads, @@ -401,7 +400,15 @@ def forward( encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - ): + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Only VANILLA honors key_padding_mask; TRTLLM and CUTEDSL drop it, and FA4 + # reads it as a valid prefix while the pipeline left-pads. + if attention_mask is not None and self.attn_backend != "VANILLA": + raise NotImplementedError( + "Padded GlmImage prompts require the VANILLA attention backend, " + f"got {self.attn_backend}." + ) + batch_size, text_seq_length, _ = encoder_hidden_states.shape batch_size, image_seq_length, _ = hidden_states.shape hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -473,8 +480,8 @@ def __init__( eps: float = 1e-6, dtype: Optional[torch.dtype] = None, config: Optional[DiffusionModelConfig] = None, - layer_idx=0, - ): + layer_idx: int = 0, + ) -> None: super().__init__() # 1. Attention @@ -501,7 +508,7 @@ def forward( List[Tuple[torch.Tensor, torch.Tensor]], ] ] = None, - attention_mask: Optional[Dict[str, torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, attention_kwargs: Optional[Dict[str, Any]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # 1. Timestep conditioning @@ -550,12 +557,12 @@ def forward( class GlmImageTransformer2DModel(BaseDiffusionModel): - def __init__(self, model_config: DiffusionModelConfig): + def __init__(self, model_config: DiffusionModelConfig) -> None: super().__init__(model_config) vgm = model_config.visual_gen_mapping - num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 32) - self.sharder = SequenceSharder.from_vgm(vgm, num_attention_heads=num_heads) + if vgm is not None and vgm.ulysses_size > 1: + raise NotImplementedError(f"GlmImage requires ulysses_size=1, got {vgm.ulysses_size}.") pretrained_config = model_config.pretrained_config diff --git a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py index 3131f09e6f41..b95c40c0afed 100644 --- a/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_glm_image_pipeline.py @@ -27,7 +27,7 @@ from diffusers import DiffusionPipeline from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.visual_gen.models.glm_image import GlmImagePipeline +from tensorrt_llm._torch.visual_gen.models.glm_image import GlmImageAttention, GlmImagePipeline from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader from tensorrt_llm.visual_gen.args import AttentionConfig, TorchCompileConfig, VisualGenArgs @@ -376,6 +376,27 @@ def test_image_conditioning_not_supported(): pipe.forward(prompt=PROMPT, image=torch.zeros(1)) +def test_resolution_must_be_multiple_of_32(): + """1008 clears the 16px latent pitch but not the 32px prior-token grid.""" + pipe = GlmImagePipeline.__new__(GlmImagePipeline) + pipe.transformer = None + assert pipe.resolution_multiple_of == (32, 32) + with pytest.raises(ValueError, match="must be multiples of"): + pipe.forward(prompt=PROMPT, height=1008, width=1008) + + +def test_padded_text_mask_requires_vanilla_backend(): + """Backends that drop or misread key_padding_mask are rejected, not silently wrong.""" + attn = GlmImageAttention.__new__(GlmImageAttention) + attn.attn_backend = "TRTLLM" + with pytest.raises(NotImplementedError, match="VANILLA"): + attn.forward( + torch.zeros(1, 4, 8), + encoder_hidden_states=torch.zeros(1, 2, 8), + attention_mask=torch.ones(1, 2, dtype=torch.long), + ) + + # ============================================================================= # Quantization Optimization Tests # ============================================================================= From d23a50ec8cb10ee42f9fe798def83f3de81c21e2 Mon Sep 17 00:00:00 2001 From: jloftin Date: Fri, 14 Aug 2026 23:48:50 +0000 Subject: [PATCH 10/11] Review 3 Signed-off-by: jloftin --- .../visual_gen/models/glm_image/pipeline_glm_image.py | 10 ++++------ .../models/glm_image/transformer_glm_image.py | 8 ++++++-- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py index c4e26f2f209d..0b4d161ef3d5 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py @@ -605,14 +605,12 @@ def forward( Args: prompt (`str` or `list[str]`, *optional*): - The prompt or prompts to guide the image generation. Must contain shape info in the format 'H - W' where H and W are token dimensions (d32). Example: "A beautiful sunset36 24" - generates a 1152x768 image. - image: Optional condition images for image-to-image generation. + The prompt or prompts to guide the image generation. + image: Must be `None`; this pipeline is text-to-image. height (`int`, *optional*): - The height in pixels. If not provided, derived from prompt shape info. + The height in pixels, a multiple of 32. Defaults to 1024. width (`int`, *optional*): - The width in pixels. If not provided, derived from prompt shape info. + The width in pixels, a multiple of 32. Defaults to 1024. num_inference_steps (`int`, *optional*, defaults to `50`): The number of denoising steps for DiT. guidance_scale (`float`, *optional*, defaults to `1.5`): diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py index fdbbd131e8fd..32c183229748 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/transformer_glm_image.py @@ -561,8 +561,12 @@ def __init__(self, model_config: DiffusionModelConfig) -> None: super().__init__(model_config) vgm = model_config.visual_gen_mapping - if vgm is not None and vgm.ulysses_size > 1: - raise NotImplementedError(f"GlmImage requires ulysses_size=1, got {vgm.ulysses_size}.") + if vgm is not None and (vgm.ulysses_size > 1 or vgm.cp_size > 1): + raise NotImplementedError( + "GlmImage requires ulysses_size=1, ring_size=1 and attn2d_size=(1, 1); got " + f"ulysses_size={vgm.ulysses_size}, ring_size={vgm.ring_size}, " + f"attn2d_size=({vgm.attn2d_row_size}, {vgm.attn2d_col_size})." + ) pretrained_config = model_config.pretrained_config From abec434bfdb0b2a596b8f770eb22bcbc4e099d19 Mon Sep 17 00:00:00 2001 From: Joseph Loftin Date: Wed, 19 Aug 2026 20:38:41 +0000 Subject: [PATCH 11/11] restore wan Signed-off-by: Joseph Loftin --- docs/source/models/visual-generation.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index ed9ab93666c4..25f4b4b26e78 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -61,13 +61,14 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | | **Wan 2.1 VSA** [^2] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes | | **Wan 2.2** | Yes | Yes | Yes [^3] | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | -| **FastWan 2.2** | Yes | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | +| **FastWan 2.2** | Yes | Yes | No | No | No [^7] | No | No | Yes | Yes | Yes | No | No | No | No | | **LTX-2** | Yes | Yes | Yes [^4] | Yes | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image-Layered** [^6] | No | No | No | No | No | No | No | Yes | Yes | No | No | No | No | No | | **Qwen-Image-Edit-2511** | Yes | Yes | No | No | Yes | No | No | Yes | Yes | No | No | No | No | No | | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | | **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | Yes | No | No | No | No | +| **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | Yes | No | No | No | No | [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable. @@ -79,6 +80,8 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^6]: Qwen-Image-Layered supports baseline BF16 image-conditioned layer decomposition and returns the generated RGBA layer stack as a saveable image grid. FP8 blockwise, NVFP4, cache acceleration, attention-parallel/Sage/VSA backends, Tensor Parallelism, and `trtllm-serve` image-edit routing are not enabled for this pipeline yet. +[^7]: `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` — a distilled version of Wan2.2-TI2V-5B with 3 denoising steps. CFG parallelism, TeaCache, and Cache-DiT are not applicable. + ## Quick Start Here is a simple example to generate a video with Wan 2.1: