From c54e0737d06c914ed338c1eef2eca7b8b6da2731 Mon Sep 17 00:00:00 2001 From: ShreejithSG Date: Fri, 16 Jan 2026 03:46:14 +0000 Subject: [PATCH] Address PR comments: GLM-Image Port to FastVideo- Replaced diffusers FeedForward with FastVideo MLP in DiT.- Moved SDPAMetadata import to top of attention layer for cleaner code.- Removed unnecessary __call__ override from pipeline to use standard composition.- Consolidated CFG into a single batched pass with concatenated embeddings for performance.- Fixed attention mask compatibility issues with optimized FLASH_ATTN backend.- Robustly patched flash_attn unpad utility for variable version return values. --- examples/inference/basic/basic_glm_image.py | 75 ++ fastvideo/configs/models/dits/__init__.py | 4 +- fastvideo/configs/models/dits/glm_image.py | 106 +++ fastvideo/configs/models/vaes/__init__.py | 2 + fastvideo/configs/models/vaes/glm_image.py | 54 ++ fastvideo/configs/pipelines/glm_image.py | 122 +++ fastvideo/models/dits/glm_image.py | 842 ++++++++++++++++++ fastvideo/models/loader/component_loader.py | 35 + fastvideo/models/registry.py | 7 + .../pipelines/basic/glm_image/__init__.py | 7 + .../basic/glm_image/glm_image_pipeline.py | 83 ++ fastvideo/pipelines/pipeline_registry.py | 1 + fastvideo/pipelines/stages/__init__.py | 5 + .../stages/glm_image_before_denoising.py | 232 +++++ .../pipelines/stages/glm_image_decoding.py | 73 ++ .../pipelines/stages/glm_image_denoising.py | 139 +++ 16 files changed, 1786 insertions(+), 1 deletion(-) create mode 100644 examples/inference/basic/basic_glm_image.py create mode 100644 fastvideo/configs/models/dits/glm_image.py create mode 100644 fastvideo/configs/models/vaes/glm_image.py create mode 100644 fastvideo/configs/pipelines/glm_image.py create mode 100644 fastvideo/models/dits/glm_image.py create mode 100644 fastvideo/pipelines/basic/glm_image/__init__.py create mode 100644 fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py create mode 100644 fastvideo/pipelines/stages/glm_image_before_denoising.py create mode 100644 fastvideo/pipelines/stages/glm_image_decoding.py create mode 100644 fastvideo/pipelines/stages/glm_image_denoising.py diff --git a/examples/inference/basic/basic_glm_image.py b/examples/inference/basic/basic_glm_image.py new file mode 100644 index 000000000..16c978715 --- /dev/null +++ b/examples/inference/basic/basic_glm_image.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GLM-Image inference example using FastVideo. + +This example demonstrates how to generate images using the GLM-Image model +through FastVideo's pipeline infrastructure. + +Usage: + python examples/inference/basic/basic_glm_image.py +""" + +import os + +# Set environment variables for single-GPU distributed setup +os.environ.setdefault("MASTER_ADDR", "127.0.0.1") +os.environ.setdefault("MASTER_PORT", "29501") +os.environ.setdefault("RANK", "0") +os.environ.setdefault("WORLD_SIZE", "1") + +from fastvideo.pipelines.basic.glm_image import GlmImagePipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +OUTPUT_PATH = "image_output" + + + +def main(): + # Load the GLM-Image pipeline + pipe = GlmImagePipeline.from_pretrained( + "zai-org/GLM-Image", + num_gpus=1, + trust_remote_code=True, + ) + + # Create a batch with generation parameters + prompt = ( + "A beautiful landscape photography with rolling hills, " + "a winding river, and a vibrant sunset in the background. " + "Warm golden light, photorealistic style." + ) + + batch = ForwardBatch( + data_type="image", + prompt=prompt, + negative_prompt="", + height=1024, + width=1024, + num_inference_steps=50, + guidance_scale=7.5, + guidance_rescale=0.7, + do_classifier_free_guidance=True, + num_frames=1, + seed=42, + ) + + # Generate the image + result = pipe.forward(batch, pipe.fastvideo_args) + + # Output is in result.output as a tensor [B, C, T, H, W] + print(f"Generated image tensor shape: {result.output.shape}") + + # Save the output image + import torch + from torchvision.utils import save_image + + os.makedirs(OUTPUT_PATH, exist_ok=True) + # result.output is [B, C, 1, H, W] for images + image_tensor = result.output.squeeze(2) # [B, C, H, W] + save_path = os.path.join(OUTPUT_PATH, "output.png") + save_image(image_tensor, save_path) + print(f"Image saved to {save_path}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 5ae49b488..701234257 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -2,6 +2,7 @@ from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig from fastvideo.configs.models.dits.flux_2 import Flux2Config +from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config @@ -16,5 +17,6 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", - "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config" + "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config", + "GlmImageDiTConfig" ] diff --git a/fastvideo/configs/models/dits/glm_image.py b/fastvideo/configs/models/dits/glm_image.py new file mode 100644 index 000000000..03f149f9d --- /dev/null +++ b/fastvideo/configs/models/dits/glm_image.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GLM-Image DiT configuration. + +GLM-Image uses a 7B diffusion transformer (DiT) decoder that expands tokens +from the autoregressive vision-language encoder into high-resolution images. +""" +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_blocks(n: str, m) -> bool: + return "transformer_blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class GlmImageDiTArchConfig(DiTArchConfig): + """Architecture config for GlmImageTransformer2DModel.""" + + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) + + # GLM-Image DiT settings (7B model) + # hidden_size = num_attention_heads * attention_head_dim = 64 * 64 = 4096 + hidden_size: int = 4096 + num_attention_heads: int = 32 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + num_layers: int = 30 + + # Text and condition dims + text_embed_dim: int = 1472 + time_embed_dim: int = 512 + condition_dim: int = 256 + + # VQ settings for AR tokens + prior_vq_quantizer_codebook_size: int = 16384 + + # Patch embedding + patch_size: int = 2 + + # Positional embedding max resolution + max_height: int = 2048 + max_width: int = 2048 + + # QK normalization + qk_norm: str = "layer_norm" + eps: float = 1e-5 + + # LoRA exclusions + exclude_lora_layers: list[str] = field( + default_factory=lambda: + ["image_projector", "glyph_projector", "prior_token_embedding"]) + + # Param name mappings for weight loading (HF -> custom) + param_names_mapping: dict = field( + default_factory=lambda: { + # Projectors (mapped from FeedForward net.0.proj -> fc_in, net.2 -> fc_out) + r"^image_projector\.net\.0\.proj\.(.*)$": + r"image_projector.fc_in.\1", + r"^image_projector\.net\.2\.(.*)$": r"image_projector.fc_out.\1", + r"^glyph_projector\.net\.0\.proj\.(.*)$": + r"glyph_projector.fc_in.\1", + r"^glyph_projector\.net\.2\.(.*)$": r"glyph_projector.fc_out.\1", + r"^prior_projector\.net\.0\.proj\.(.*)$": + r"prior_projector.fc_in.\1", + r"^prior_projector\.net\.2\.(.*)$": r"prior_projector.fc_out.\1", + r"^prior_token_embedding\.(.*)$": r"prior_token_embedding.\1", + + # Transformer blocks + r"^transformer_blocks\.(\d+)\.norm1\.(.*)$": + r"transformer_blocks.\1.norm1.\2", + r"^transformer_blocks\.(\d+)\.attn1\.to_q\.(.*)$": + r"transformer_blocks.\1.attn1.to_q.\2", + r"^transformer_blocks\.(\d+)\.attn1\.to_k\.(.*)$": + r"transformer_blocks.\1.attn1.to_k.\2", + r"^transformer_blocks\.(\d+)\.attn1\.to_v\.(.*)$": + r"transformer_blocks.\1.attn1.to_v.\2", + r"^transformer_blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": + r"transformer_blocks.\1.attn1.to_out.0.\2", + + # FeedForward in blocks (net.0.proj -> fc_in, net.2 -> fc_out) + r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": + r"transformer_blocks.\1.ff.fc_in.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": + r"transformer_blocks.\1.ff.fc_out.\2", + + # Output + r"^norm_out\.(.*)$": r"norm_out.\1", + r"^proj_out\.(.*)$": r"proj_out.\1", + }) + + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + def __post_init__(self): + super().__post_init__() + self.num_channels_latents = self.out_channels + + +@dataclass +class GlmImageDiTConfig(DiTConfig): + """Configuration for GLM-Image DiT model.""" + + arch_config: DiTArchConfig = field(default_factory=GlmImageDiTArchConfig) + prefix: str = "GlmImage" diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index d0bd1dd83..6f3ca9768 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -2,6 +2,7 @@ from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig +from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig @@ -21,4 +22,5 @@ __all__ = [ "OobleckVAEArchConfig", "OobleckVAEConfig", "Flux2VAEConfig", + "GlmImageVAEConfig", ] diff --git a/fastvideo/configs/models/vaes/glm_image.py b/fastvideo/configs/models/vaes/glm_image.py new file mode 100644 index 000000000..f1f9f3f94 --- /dev/null +++ b/fastvideo/configs/models/vaes/glm_image.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GLM-Image VAE configuration. + +GLM-Image uses an AutoencoderKL for encoding/decoding images to/from latent space. +This is an image-only VAE (no temporal compression). +""" +from dataclasses import dataclass, field + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class GlmImageVAEArchConfig(VAEArchConfig): + """Architecture config for GLM-Image VAE (AutoencoderKL).""" + + # Standard KL-VAE parameters + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 16 + + # Encoder/decoder architecture + block_out_channels: tuple[int, ...] = (128, 256, 512, 512) + layers_per_block: int = 2 + norm_num_groups: int = 32 + + # Scaling factor for latents (standard for SD-style VAEs) + scaling_factor: float = 0.18215 + + # Image-only VAE: no temporal compression + temporal_compression_ratio: int = 1 + spatial_compression_ratio: int = 8 + + +@dataclass +class GlmImageVAEConfig(VAEConfig): + """Configuration for GLM-Image VAE.""" + + arch_config: GlmImageVAEArchConfig = field( + default_factory=GlmImageVAEArchConfig) + + # Tiling settings for high-resolution images + use_tiling: bool = True + use_temporal_tiling: bool = False # Image model, no temporal dimension + use_parallel_tiling: bool = False + + # Tile dimensions for memory efficiency + tile_sample_min_height: int = 512 + tile_sample_min_width: int = 512 + tile_sample_stride_height: int = 384 + tile_sample_stride_width: int = 384 + + # For image models, we need both encoder and decoder + load_encoder: bool = True + load_decoder: bool = True diff --git a/fastvideo/configs/pipelines/glm_image.py b/fastvideo/configs/pipelines/glm_image.py new file mode 100644 index 000000000..910685d39 --- /dev/null +++ b/fastvideo/configs/pipelines/glm_image.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GLM-Image pipeline configuration. + +GLM-Image is a hybrid model combining: +- 9B autoregressive (AR) vision-language encoder for semantic/layout token generation +- 7B diffusion transformer (DiT) decoder for high-fidelity image generation + +The pipeline supports: +- Text-to-image (T2I) generation +- Image-to-image (I2I) editing/style transfer +""" +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config +from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +def glm_image_t5_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor: + """Postprocess T5 encoder outputs for GLM-Image. + + GLM-Image uses T5 as a glyph/text encoder for text rendering capabilities. + """ + mask: torch.Tensor = outputs.attention_mask + hidden_state: torch.Tensor = outputs.last_hidden_state + seq_lens = mask.gt(0).sum(dim=1).long() + + assert torch.isnan(hidden_state).sum() == 0, "T5 hidden states contain NaN" + + # Pack embeddings with proper padding + prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] + max_len = 512 # GLM-Image uses 512 max sequence length + prompt_embeds_tensor: torch.Tensor = torch.stack([ + torch.cat([u, u.new_zeros(max_len - u.size(0), u.size(1))]) + for u in prompt_embeds + ], + dim=0) + + return prompt_embeds_tensor + + +@dataclass +class GlmImageConfig(PipelineConfig): + """Configuration for GLM-Image T2I/I2I pipeline. + + This configuration sets up the complete GLM-Image pipeline including: + - T5 text encoder (glyph encoder for text rendering) + - Vision-language encoder (AR model for semantic tokens) + - DiT diffusion decoder + - VAE for image encoding/decoding + """ + + # Image generation flag (not video) + is_video_pipeline: bool = False + + # DiT configuration + dit_config: DiTConfig = field(default_factory=GlmImageDiTConfig) + dit_precision: str = "bf16" + + # VAE configuration + vae_config: VAEConfig = field(default_factory=GlmImageVAEConfig) + vae_precision: str = "fp32" + vae_tiling: bool = True + vae_sp: bool = False # No spatial parallelism for image VAE + + # Text encoder (T5 for glyph/text rendering) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (T5Config(), )) + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp32", )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], + ...] = field(default_factory=lambda: + (glm_image_t5_postprocess, )) + + # Flow matching scheduler settings + flow_shift: float | None = 1.0 + embedded_cfg_scale: float = 7.5 + + # AR model sampling parameters + ar_temperature: float = 0.95 + ar_top_p: float = 0.75 + ar_do_sample: bool = True + ar_max_new_tokens: int = 1024 + + # Image resolution constraints (must be divisible by 32) + default_height: int = 1024 + default_width: int = 1024 + resolution_multiple: int = 32 + + # Post-decoding hook for custom processing + post_decoding: Callable | None = None + + def __post_init__(self): + # GLM-Image needs both encoder and decoder + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + def validate_resolution(self, height: int, width: int) -> tuple[int, int]: + """Ensure resolution is divisible by 32 as required by GLM-Image.""" + h = (height // self.resolution_multiple) * self.resolution_multiple + w = (width // self.resolution_multiple) * self.resolution_multiple + return h, w + + +@dataclass +class GlmImageI2IConfig(GlmImageConfig): + """Configuration for GLM-Image image-to-image pipeline. + + Inherits from GlmImageConfig with settings optimized for I2I tasks + like style transfer, editing, and inpainting. + """ + + # Lower CFG for I2I to preserve input structure + embedded_cfg_scale: float = 5.0 + + # Strength for I2I (how much of the original to preserve) + strength: float = 0.8 diff --git a/fastvideo/models/dits/glm_image.py b/fastvideo/models/dits/glm_image.py new file mode 100644 index 000000000..23bc7f473 --- /dev/null +++ b/fastvideo/models/dits/glm_image.py @@ -0,0 +1,842 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +"""GLM-Image Transformer 2D Model. + +This module implements the diffusion transformer (DiT) decoder for GLM-Image. +The model takes tokens from the autoregressive vision-language encoder and +denoises them into high-fidelity images. +""" +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from fastvideo.layers.mlp import MLP + +from fastvideo.attention import LocalAttention +from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig +from fastvideo.layers.layernorm import ScaleResidualLayerNormScaleShift +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.rotary_embedding import _apply_rotary_emb +from fastvideo.layers.visual_embedding import Timesteps +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + +logger = init_logger(__name__) + + +# ============================================================================= +# KV Cache (for future use with caching) +# ============================================================================= +class GlmImageLayerKVCache: + """KV cache for GlmImage model per layer.""" + + def __init__(self): + self.k_cache = None + self.v_cache = None + self.mode: Optional[str] = None # "write", "read", "skip" + + def store(self, k: torch.Tensor, v: torch.Tensor): + if self.k_cache is None: + self.k_cache = k + self.v_cache = v + else: + self.k_cache = torch.cat([self.k_cache, k], dim=1) + self.v_cache = torch.cat([self.v_cache, v], dim=1) + + def get(self): + return self.k_cache, self.v_cache + + def clear(self): + self.k_cache = None + self.v_cache = None + self.mode = None + + +class GlmImageKVCache: + """Container for all layers' KV caches.""" + + def __init__(self, num_layers: int): + self.num_layers = num_layers + self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)] + + def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache: + return self.caches[layer_idx] + + def set_mode(self, mode: Optional[str]): + if mode is not None and mode not in ["write", "read", "skip"]: + raise ValueError( + f"Invalid mode: {mode}, must be one of 'write', 'read', 'skip'" + ) + for cache in self.caches: + cache.mode = mode + + def clear(self): + for cache in self.caches: + cache.clear() + + +# ============================================================================= +# Timestep and Text Projection +# ============================================================================= +class GlmImageTimestepEmbedding(nn.Module): + """ + Replacement for diffusers TimestepEmbedding using ReplicatedLinear. + Structure: linear_1 -> act(silu) -> linear_2 + """ + + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + ): + super().__init__() + if out_dim is None: + out_dim = time_embed_dim + self.linear_1 = ReplicatedLinear(in_channels, time_embed_dim, bias=True) + if act_fn == "silu": + self.act = nn.SiLU() + elif act_fn == "gelu": + self.act = nn.GELU(approximate="tanh") + else: + self.act = nn.SiLU() + self.linear_2 = ReplicatedLinear(time_embed_dim, out_dim, bias=True) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample, _ = self.linear_1(sample) + sample = self.act(sample) + sample, _ = self.linear_2(sample) + return sample + + +class GlmImageTextProjection(nn.Module): + """ + Replacement for diffusers PixArtAlphaTextProjection using ReplicatedLinear. + Structure: linear_1 -> act_1 -> linear_2 + """ + + def __init__( + self, + in_features: int, + hidden_size: int, + out_features: int = None, + act_fn: str = "silu", + ): + super().__init__() + if out_features is None: + out_features = hidden_size + self.linear_1 = ReplicatedLinear(in_features, hidden_size, bias=True) + if act_fn == "silu": + self.act_1 = nn.SiLU() + elif act_fn == "gelu_tanh": + self.act_1 = nn.GELU(approximate="tanh") + else: + self.act_1 = nn.SiLU() + self.linear_2 = ReplicatedLinear(hidden_size, out_features, bias=True) + + 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) + return hidden_states + + +class GlmImageCombinedTimestepSizeEmbeddings(nn.Module): + """Combined timestep and size embeddings using Timesteps class.""" + + def __init__( + self, + embedding_dim: int, + condition_dim: int, + pooled_projection_dim: int, + timesteps_dim: int = 256, + ): + 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 + ) + self.condition_embedder = GlmImageTextProjection( + pooled_projection_dim, embedding_dim, act_fn="silu" + ) + + 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) + ) # (B, embedding_dim) + condition_emb = self.condition_embedder( + condition_proj.to(dtype=hidden_dtype) + ) # (B, embedding_dim) + + conditioning = timesteps_emb + condition_emb + return conditioning + + +# ============================================================================= +# Image Projector +# ============================================================================= +class GlmImageImageProjector(nn.Module): + """Projects latent patches to hidden dimension.""" + + def __init__( + self, + in_channels: int = 16, + hidden_size: int = 2560, + patch_size: int = 2, + ): + super().__init__() + self.patch_size = patch_size + self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) + + 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 + + +# ============================================================================= +# AdaLayerNorm +# ============================================================================= +class GlmImageAdaLayerNormZero(nn.Module): + """Ada Layer Norm with 12-way modulation using ReplicatedLinear.""" + + def __init__(self, embedding_dim: int, dim: int) -> None: + super().__init__() + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.linear = ReplicatedLinear(embedding_dim, 12 * dim, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: 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 + ) + + 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, + ) + + +# ============================================================================= +# Attention +# ============================================================================= +class GlmImageAttention(nn.Module): + """GLM-Image attention using LocalAttention backend with ReplicatedLinear.""" + + def __init__( + self, + query_dim: int, + heads: int, + dim_head: int, + out_dim: int, + bias: bool = True, + qk_norm: str = "layer_norm", + elementwise_affine: bool = False, + eps: float = 1e-5, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + prefix: str = "", + ): + super().__init__() + + self.heads = out_dim // dim_head if out_dim is not None else heads + self.num_kv_heads = self.heads + self.dim_head = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.inner_kv_dim = self.inner_dim + self.out_dim = out_dim if out_dim is not None else query_dim + + self.to_q = ReplicatedLinear(query_dim, self.inner_dim, bias=bias) + self.to_k = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias) + self.to_v = ReplicatedLinear(query_dim, self.inner_kv_dim, bias=bias) + + self.to_out = nn.ModuleList( + [ReplicatedLinear(self.inner_dim, self.out_dim, bias=True)] + ) + + if qk_norm is None: + self.norm_q = None + self.norm_k = None + elif qk_norm == "layer_norm": + self.norm_q = nn.LayerNorm( + dim_head, eps=eps, elementwise_affine=elementwise_affine + ) + self.norm_k = nn.LayerNorm( + dim_head, eps=eps, elementwise_affine=elementwise_affine + ) + else: + raise ValueError(f"unknown qk_norm: {qk_norm}") + + self.attn = LocalAttention( + num_heads=self.heads, + head_size=dim_head, + num_kv_heads=self.heads, + softmax_scale=None, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + kv_cache: Optional[GlmImageLayerKVCache] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + dtype = encoder_hidden_states.dtype + + 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) + + # 1. QKV projections + query, _ = self.to_q(hidden_states) + key, _ = self.to_k(hidden_states) + value, _ = self.to_v(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.norm_q is not None: + query = self.norm_q(query).to(dtype=dtype) + if self.norm_k is not None: + key = self.norm_k(key).to(dtype=dtype) + + # 3. Rotational positional embeddings applied to latent stream + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + + query[:, text_seq_length:, :, :] = _apply_rotary_emb( + query[:, text_seq_length:, :, :], cos, sin, is_neox_style=True + ) + key[:, text_seq_length:, :, :] = _apply_rotary_emb( + key[:, text_seq_length:, :, :], cos, sin, is_neox_style=True + ) + + # 4. KV Cache handling + if kv_cache is not None: + if kv_cache.mode == "write": + kv_cache.store(key, value) + elif kv_cache.mode == "read": + k_cache, v_cache = kv_cache.get() + key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key + value = ( + torch.cat([v_cache, value], dim=1) if v_cache is not None else value + ) + elif kv_cache.mode == "skip": + pass + + # 5. Attention + if attention_mask is not None: + # Generate the combined mask for [text, image] + # text_attn_mask shape: (batch_size, text_seq_length) + text_attn_mask = attention_mask + device = query.device + dtype = query.dtype + + mix_attn_mask = torch.ones( + (batch_size, text_seq_length + image_seq_length), device=device + ) + # Match SGLang: text part uses provided mask, image part is always 1 + mix_attn_mask[:, :text_seq_length] = text_attn_mask.float().to(device) + + # Use 2D mask (B, L) for simplified backend handling. + # flash_attn/unpad_input and standard SDPA can both handle this. + attention_mask_kv = (mix_attn_mask > 0) # (B, L) + else: + attention_mask_kv = None + + hidden_states = self.attn(query, key, value, attention_mask=attention_mask_kv) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + # 6. Output projection + hidden_states, _ = self.to_out[0](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 + + +# ============================================================================= +# Transformer Block +# ============================================================================= +class GlmImageTransformerBlock(nn.Module): + """GLM-Image Transformer block with fused residual+norm operations.""" + + def __init__( + self, + dim: int = 2560, + num_attention_heads: int = 64, + attention_head_dim: int = 40, + time_embed_dim: int = 512, + supported_attention_backends: set[AttentionBackendEnum] | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + # 1. Attention + self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim) + + self.attn1 = GlmImageAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + out_dim=dim, + bias=True, + qk_norm="layer_norm", + elementwise_affine=False, + eps=1e-5, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1", + ) + + # 2. Feedforward with fused ScaleResidualLayerNorm + self.norm2 = ScaleResidualLayerNormScaleShift( + dim, norm_type="layer", eps=1e-5, elementwise_affine=False + ) + self.norm2_context = ScaleResidualLayerNormScaleShift( + dim, norm_type="layer", eps=1e-5, elementwise_affine=False + ) + self.ff = MLP(input_dim=dim, mlp_hidden_dim=dim * 4, output_dim=dim, act_type="gelu_pytorch_tanh") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: 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, + kv_cache: Optional[GlmImageLayerKVCache] = 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 + if attention_kwargs is None: + attention_kwargs = {} + + 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, + kv_cache=kv_cache, + **attention_kwargs, + ) + + # 3. Feedforward (fused residual + norm + scale/shift) + # Fix: hidden_states is the "residual" (base stream), attn_hidden_states is the "x" (modulated by gate) + norm_hidden_states, hidden_states = self.norm2( + hidden_states, + attn_hidden_states, + gate_msa.unsqueeze(1), + shift_mlp.unsqueeze(1), + scale_mlp.unsqueeze(1), + ) + norm_encoder_hidden_states, encoder_hidden_states = self.norm2_context( + encoder_hidden_states, + attn_encoder_hidden_states, + c_gate_msa.unsqueeze(1), + c_shift_mlp.unsqueeze(1), + c_scale_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 + + +# ============================================================================= +# Rotary Positional Embedding +# ============================================================================= +class GlmImageRotaryPosEmbed(nn.Module): + """2D rotary positional embeddings for image patches.""" + + def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None: + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.theta = theta + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size, num_channels, height, width = hidden_states.shape + height, width = height // self.patch_size, width // self.patch_size + device = hidden_states.device + + dim_h, dim_w = self.dim // 2, self.dim // 2 + h_inv_freq = 1.0 / ( + self.theta + ** ( + torch.arange(0, dim_h, 2, dtype=torch.float32, device=device)[ + : (dim_h // 2) + ].float() + / dim_h + ) + ) + w_inv_freq = 1.0 / ( + self.theta + ** ( + torch.arange(0, dim_w, 2, dtype=torch.float32, device=device)[ + : (dim_w // 2) + ].float() + / dim_w + ) + ) + h_seq = torch.arange(height, device=device) + w_seq = torch.arange(width, device=device) + freqs_h = torch.outer(h_seq, h_inv_freq) + freqs_w = torch.outer(w_seq, w_inv_freq) + + # Create position matrices for height and width + freqs_h = freqs_h.unsqueeze(1) + freqs_w = freqs_w.unsqueeze(0) + # Broadcast freqs_h and freqs_w to [height, width, dim//4] + freqs_h = freqs_h.expand(height, width, -1) + freqs_w = freqs_w.expand(height, width, -1) + + # Concatenate along last dimension to get [height, width, dim//2] + freqs = torch.cat([freqs_h, freqs_w], dim=-1) + freqs = freqs.reshape(height * width, -1) # [height * width, dim//2] + return (freqs.cos(), freqs.sin()) + + +# ============================================================================= +# Final AdaLayerNorm +# ============================================================================= +class GlmImageAdaLayerNormContinuous(nn.Module): + """ + GlmImage-only final AdaLN: LN(x) -> Linear(cond) -> chunk -> affine. + Matches Megatron: **no activation** before the Linear on conditioning embedding. + """ + + 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", + ): + super().__init__() + self.linear = nn.Linear( + conditioning_embedding_dim, embedding_dim * 2, bias=bias + ) + if norm_type == "layer_norm": + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = nn.RMSNorm(embedding_dim, eps, elementwise_affine) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + def forward( + self, x: torch.Tensor, conditioning_embedding: torch.Tensor + ) -> torch.Tensor: + # *** NO SiLU here *** + 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 + + +# ============================================================================= +# Main Model +# ============================================================================= +class GlmImageTransformer2DModel(BaseDiT): + """ + GLM-Image Transformer 2D Model. + + Aligned with SGLang implementation for numerical consistency. + """ + + _fsdp_shard_conditions = GlmImageDiTConfig().arch_config._fsdp_shard_conditions + _compile_conditions = GlmImageDiTConfig().arch_config._fsdp_shard_conditions + _supported_attention_backends = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + + param_names_mapping = GlmImageDiTConfig().arch_config.param_names_mapping + reverse_param_names_mapping = {} + lora_param_names_mapping = {} + + def __init__( + self, + config: GlmImageDiTConfig, + hf_config: dict[str, Any], + ): + super().__init__(config=config, hf_config=hf_config) + + self.config_data = config # Store config + arch_config = config.arch_config + + self.in_channels = arch_config.in_channels + self.out_channels = arch_config.out_channels + self.patch_size = arch_config.patch_size + self.num_layers = arch_config.num_layers + self.attention_head_dim = arch_config.attention_head_dim + self.num_attention_heads = arch_config.num_attention_heads + self.text_embed_dim = arch_config.text_embed_dim + self.time_embed_dim = arch_config.time_embed_dim + + # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords + # Each of these are sincos embeddings of shape 2 * condition_dim + pooled_projection_dim = 2 * 2 * arch_config.condition_dim + inner_dim = arch_config.num_attention_heads * arch_config.attention_head_dim + + # Required by BaseDiT + self.hidden_size = inner_dim + self.num_channels_latents = arch_config.out_channels + + # 1. RoPE + self.rotary_emb = GlmImageRotaryPosEmbed( + arch_config.attention_head_dim, arch_config.patch_size, theta=10000.0 + ) + + # 2. Patch & Text-timestep embedding + self.image_projector = GlmImageImageProjector( + arch_config.in_channels, inner_dim, arch_config.patch_size + ) + self.glyph_projector = MLP( + input_dim=arch_config.text_embed_dim, + mlp_hidden_dim=inner_dim, + output_dim=inner_dim, + act_type="gelu_pytorch_tanh", + ) + self.prior_token_embedding = nn.Embedding( + arch_config.prior_vq_quantizer_codebook_size, inner_dim + ) + self.prior_projector = MLP( + input_dim=inner_dim, + mlp_hidden_dim=inner_dim, + output_dim=inner_dim, + act_type="silu", + ) + + self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( + embedding_dim=arch_config.time_embed_dim, + condition_dim=arch_config.condition_dim, + pooled_projection_dim=pooled_projection_dim, + timesteps_dim=arch_config.time_embed_dim, + ) + + # 3. Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + GlmImageTransformerBlock( + inner_dim, + arch_config.num_attention_heads, + arch_config.attention_head_dim, + arch_config.time_embed_dim, + supported_attention_backends=self._supported_attention_backends, + prefix=f"transformer_blocks.{i}", + ) + for i in range(arch_config.num_layers) + ] + ) + + # 4. Output projection + self.norm_out = GlmImageAdaLayerNormContinuous( + inner_dim, arch_config.time_embed_dim, elementwise_affine=False + ) + self.proj_out = nn.Linear( + inner_dim, + arch_config.patch_size * arch_config.patch_size * arch_config.out_channels, + bias=True, + ) + + self.gradient_checkpointing = False + + self.__post_init__() + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + prior_token_id: torch.Tensor, + prior_token_drop: torch.Tensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + attention_kwargs: Optional[Dict[str, Any]] = None, + attention_mask: Optional[torch.Tensor] = None, + kv_caches: Optional[GlmImageKVCache] = None, + kv_caches_mode: Optional[str] = None, + freqs_cis: Optional[ + Union[ + Tuple[torch.Tensor, torch.Tensor], + List[Tuple[torch.Tensor, torch.Tensor]], + ] + ] = None, + guidance: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + if kv_caches is not None: + kv_caches.set_mode(kv_caches_mode) + + batch_size, num_channels, height, width = hidden_states.shape + + # SGLang does timestep -= 1.0 + timestep = timestep - 1.0 + + if isinstance(encoder_hidden_states, list): + encoder_hidden_states = encoder_hidden_states[0] + + # 1. RoPE + image_rotary_emb = freqs_cis + if image_rotary_emb is None: + image_rotary_emb = self.rotary_emb(hidden_states) + + # 2. Patch & Timestep embeddings + p = self.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) + # Using boolean indexing for sample-level dropping + if prior_token_drop.any(): + # Create a mask for the batch dimension + 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 + ) + temb = F.silu(temb) + + # 3. Transformer blocks + for idx, block in enumerate(self.transformer_blocks): + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + attention_mask, + attention_kwargs, + kv_cache=kv_caches[idx] if kv_caches is not None else None, + ) + + # 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 + ) + output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + + return output.float() + + +EntryClass = GlmImageTransformer2DModel diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index ffe5fd07f..6b097c317 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -95,6 +95,8 @@ class ComponentLoader(ABC): "image_processor": (ImageProcessorLoader, "transformers"), "feature_extractor": (ImageProcessorLoader, "transformers"), "image_encoder": (ImageEncoderLoader, "transformers"), + "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"), + "processor": (ProcessorLoader, "transformers"), "upsampler": (UpsamplerLoader, "diffusers"), "upsampler_2": (UpsamplerLoader, "diffusers"), # Stable Audio's `StableAudioMultiConditioner` bundles T5 + @@ -523,6 +525,39 @@ class ImageEncoderLoader(TextEncoderLoader): ) +class VisionLanguageEncoderLoader(ComponentLoader): + """Loader for vision-language autoregressive encoders.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + from fastvideo.distributed.parallel_state import get_local_torch_device + from fastvideo.models.encoders.glm_image_ar_loader import ( + GlmImageARLoader) + + logger.info("Loading vision-language encoder from %s", model_path) + target_device = get_local_torch_device() + loader = GlmImageARLoader( + model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=fastvideo_args.trust_remote_code, + ).to(target_device).eval() + return loader + + +class ProcessorLoader(ComponentLoader): + """Loader for HF processors that pair with vision-language encoders.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + from transformers import AutoProcessor + + logger.info("Loading processor from %s", model_path) + processor = AutoProcessor.from_pretrained( + model_path, + trust_remote_code=fastvideo_args.trust_remote_code, + ) + logger.info("Loaded processor: %s", processor.__class__.__name__) + return processor + + class ImageProcessorLoader(ComponentLoader): """Loader for image processor.""" diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 69b9d6d0a..78ca7ed42 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -61,6 +61,11 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"), } +# Text-to-image DiT models (2D image generation) +_TEXT_TO_IMAGE_DIT_MODELS = { + "GlmImageTransformer2DModel": ("dits", "glm_image", "GlmImageTransformer2DModel"), +} + _TEXT_ENCODER_MODELS = { "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"), "CLIPTextModelWithProjection": @@ -87,6 +92,7 @@ _IMAGE_ENCODER_MODELS: dict[str, tuple] = { } _VAE_MODELS = { + "AutoencoderKL": ("vaes", "autoencoderkl", "AutoencoderKL"), # Generic 2D VAE "AutoencoderKLHunyuanVideo": ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), "AutoencoderKLCausal3D": ("vaes", "gamecraftvae", "GameCraftVAE"), @@ -134,6 +140,7 @@ _UPSAMPLERS = { _LEGACY_FAST_VIDEO_MODELS = { **_TEXT_TO_VIDEO_DIT_MODELS, **_IMAGE_TO_VIDEO_DIT_MODELS, + **_TEXT_TO_IMAGE_DIT_MODELS, **_TEXT_ENCODER_MODELS, **_IMAGE_ENCODER_MODELS, **_VAE_MODELS, diff --git a/fastvideo/pipelines/basic/glm_image/__init__.py b/fastvideo/pipelines/basic/glm_image/__init__.py new file mode 100644 index 000000000..f85bca3cc --- /dev/null +++ b/fastvideo/pipelines/basic/glm_image/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GLM-Image pipeline package.""" + +from fastvideo.pipelines.basic.glm_image.glm_image_pipeline import ( + GlmImagePipeline, ) + +__all__ = ["GlmImagePipeline"] diff --git a/fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py b/fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py new file mode 100644 index 000000000..241eca6c5 --- /dev/null +++ b/fastvideo/pipelines/basic/glm_image/glm_image_pipeline.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GLM-Image diffusion pipeline implementation. +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, ) +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline +from fastvideo.pipelines.stages import InputValidationStage +from fastvideo.pipelines.stages.glm_image_before_denoising import ( + GlmImageBeforeDenoisingStage, ) +from fastvideo.pipelines.stages.glm_image_denoising import GlmImageDenoisingStage +from fastvideo.pipelines.stages.glm_image_decoding import GlmImageDecodingStage + +logger = init_logger(__name__) + + +class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "GlmImagePipeline" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + "vision_language_encoder", # AR model for prior tokens - critical for quality + "processor", # Processor for AR model + ] + + _optional_config_modules = [] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + # Initial shift, will be updated dynamically in BeforeDenoisingStage + self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=1.0) + + # Validate configuration if provided + if fastvideo_args.pipeline_config is not None: + # Basic validation or usage if needed, otherwise just pass + pass + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self.add_stage( + stage_name="input_validation_stage", + stage=InputValidationStage(), + ) + + self.add_stage( + stage_name="glm_image_before_denoising_stage", + stage=GlmImageBeforeDenoisingStage( + vae=self.get_module("vae"), + text_encoder=self.get_module("text_encoder"), + tokenizer=self.get_module("tokenizer"), + processor=self.get_module("processor"), + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vision_language_encoder=self.get_module( + "vision_language_encoder"), + ), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=GlmImageDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + pipeline=self, + ), + ) + + self.add_stage( + stage_name="decoding_stage", + stage=GlmImageDecodingStage( + vae=self.get_module("vae"), + pipeline=self, + ), + ) + + +EntryClass = GlmImagePipeline diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index dc00d35b4..def100fa2 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -16,6 +16,7 @@ from fastvideo.pipelines.lora_pipeline import LoRAPipeline logger = init_logger(__name__) + _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = { WorkloadType.I2V: "PreprocessPipelineI2V", WorkloadType.T2V: "PreprocessPipelineT2V", diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index ae4b701f6..d03512d4d 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -46,6 +46,9 @@ from fastvideo.pipelines.stages.longcat_video_vae_encoding import LongCatVideoVA from fastvideo.pipelines.stages.longcat_kv_cache_init import LongCatKVCacheInitStage from fastvideo.pipelines.stages.longcat_vc_denoising import LongCatVCDenoisingStage +# GLM-Image stages +from fastvideo.pipelines.stages.glm_image_before_denoising import GlmImageBeforeDenoisingStage + __all__ = [ "PipelineStage", "InputValidationStage", @@ -101,4 +104,6 @@ __all__ = [ "LongCatVideoVAEEncodingStage", "LongCatKVCacheInitStage", "LongCatVCDenoisingStage", + # GLM-Image stages + "GlmImageBeforeDenoisingStage", ] diff --git a/fastvideo/pipelines/stages/glm_image_before_denoising.py b/fastvideo/pipelines/stages/glm_image_before_denoising.py new file mode 100644 index 000000000..ffff9adb1 --- /dev/null +++ b/fastvideo/pipelines/stages/glm_image_before_denoising.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GLM-Image before denoising stage. + +This stage handles the preprocessing for GLM-Image generation, including: +- Prompt parsing and tokenization +- AR token generation via the vision-language encoder +- Latent preparation for the denoising process +""" + +import re +from math import sqrt + +import torch +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +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 + return m * max_shift + base_shift + + +class GlmImageBeforeDenoisingStage(PipelineStage): + + def __init__(self, + vae, + text_encoder, + tokenizer, + processor, + transformer, + scheduler, + vision_language_encoder=None): + super().__init__() + self.vae = vae + self.text_encoder = text_encoder + self.tokenizer = tokenizer + self.processor = processor + self.transformer = transformer + self.scheduler = scheduler + + if isinstance(vision_language_encoder, tuple): + self.vision_language_encoder = vision_language_encoder[0] + self.vl_processor = vision_language_encoder[1] or processor + else: + self.vision_language_encoder = vision_language_encoder + self.vl_processor = processor + + def _format_prompt(self, prompt: str, h: int, w: int) -> str: + """Add H W to prompt if missing. H, W are pixel sizes / 32.""" + th, tw = h // 32, w // 32 + tag = f"{th} {tw}" + if tag not in prompt: + # Check if there's an existing tag to replace + prompt = re.sub(r"\d+\s+\d+", "", prompt) + prompt = f"{prompt} {tag}" + return prompt + + def _parse_shape(self, prompt: str) -> tuple[int, int, int, int]: + match = re.search(r"(\d+)\s+(\d+)", prompt) + if not match: + raise ValueError( + f"Prompt must contain H W, got: {prompt}") + th, tw = int(match.group(1)), int(match.group(2)) + ratio = th / tw + pth = int(sqrt(ratio) * 16) + ptw = int(sqrt(1 / ratio) * 16) + return th, tw, pth, ptw + + def _upsample_tokens(self, tokens: torch.Tensor, th: int, + tw: int) -> torch.Tensor: + """2x nearest neighbor upsample from d32 to d16 tokens.""" + # tokens: [N] + tokens = tokens.view(1, 1, th, tw) + tokens = torch.nn.functional.interpolate(tokens.float(), + scale_factor=2, + mode="nearest").long() + return tokens.view(1, -1) + + @torch.no_grad() + def forward(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> ForwardBatch: + device = get_local_torch_device() + dtype = torch.bfloat16 + + # 1. Format and parse prompt + batch.prompt = self._format_prompt(batch.prompt, batch.height, + batch.width) + th, tw, pth, ptw = self._parse_shape(batch.prompt) + + # 2. Generate AR tokens + if self.vision_language_encoder is not None: + # Expand prompt for AR model (SGLang style) + expanded_prompt = batch.prompt.replace( + f"{th} {tw}", + f"{th} {tw}{pth} {ptw}") + + content = [{"type": "text", "text": expanded_prompt}] + if hasattr(batch, 'image') and batch.image is not None: + content.insert(0, {"type": "image", "image": batch.image}) + + messages = [{"role": "user", "content": content}] + inputs = self.vl_processor.apply_chat_template( + messages, + tokenize=True, + target_h=batch.height, + target_w=batch.width, + return_dict=True, + return_tensors="pt").to(device) + + # Total tokens: small_image + large_image + max_new = (th * tw) + (pth * ptw) + 1 + outputs = self.vision_language_encoder.generate( + **inputs, max_new_tokens=max_new, do_sample=False) + + # Extract large image tokens + gen_tokens = outputs[0][inputs.input_ids.shape[-1]:] + expected_start = pth * ptw + expected_count = th * tw + + # Handle variable-length output: take what we can, pad if needed + if len(gen_tokens) >= expected_start + expected_count: + large_tokens = gen_tokens[expected_start:expected_start + + expected_count] + else: + # Fallback: pad with zeros if not enough tokens + available = gen_tokens[expected_start:] if len( + gen_tokens) > expected_start else gen_tokens + large_tokens = torch.zeros(expected_count, + dtype=gen_tokens.dtype, + device=gen_tokens.device) + large_tokens[:min(len(available), expected_count + )] = available[:expected_count] + logger.warning( + "AR generated %d tokens, expected %d. Padding with zeros.", + len(gen_tokens), expected_start + expected_count) + + # Upsample to d16 + batch.prior_token_id = self._upsample_tokens(large_tokens, th, tw) + batch.prior_token_drop = torch.zeros(batch.prior_token_id.shape, + dtype=torch.bool, + device=device) + + else: + # No AR model - use random prior tokens (unconditional generation fallback) + # Number of patches after upsampling: (2*th) * (2*tw) + num_prior_tokens = 4 * th * tw # Upsampled from th*tw + logger.warning( + "No vision_language_encoder provided. Using random prior tokens for unconditional generation." + ) + batch.prior_token_id = torch.randint(0, + 16384, (1, num_prior_tokens), + device=device) + batch.prior_token_drop = torch.ones( + batch.prior_token_id.shape, dtype=torch.bool, + device=device) # Drop all priors + + # 3. Encode prompts with T5 + text_inputs = self.tokenizer(batch.prompt, + padding="max_length", + max_length=512, + truncation=True, + return_tensors="pt").to(device) + + prompt_embeds = self.text_encoder( + input_ids=text_inputs.input_ids, + attention_mask=text_inputs.attention_mask).last_hidden_state.to( + dtype) + + # Prepare for CFG - concatenate [neg, pos] embeddings for batched processing + if batch.do_classifier_free_guidance: + neg_text_inputs = self.tokenizer(batch.negative_prompt or "", + padding="max_length", + max_length=512, + truncation=True, + return_tensors="pt").to(device) + + neg_prompt_embeds = self.text_encoder( + input_ids=neg_text_inputs.input_ids, + attention_mask=neg_text_inputs.attention_mask + ).last_hidden_state.to(dtype) + + # Concatenate for CFG batching: [uncond, cond] - shape (2, L, D) + prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds]) + attention_mask = torch.cat( + [neg_text_inputs.attention_mask, text_inputs.attention_mask]) + else: + attention_mask = text_inputs.attention_mask + + batch.prompt_embeds = [prompt_embeds] + batch.attention_mask = attention_mask + + # 4. Prepare Latents + # GLM-Image uses 16 channels, VAE compression factor is 8 + # Add extra dimension for video compatibility: [B, C, T, H, W] with T=1 + latents = torch.randn((1, 16, 1, batch.height // 8, batch.width // 8), + device=device, + dtype=dtype) + # Flow matching doesn't use init_noise_sigma, just pure noise + batch.latents = latents + + # 5. Dynamic Flow Shift + num_patches = (batch.height // 16) * (batch.width // 16) + shift = calculate_shift(num_patches) + if hasattr(self.scheduler, 'set_shift'): + self.scheduler.set_shift(shift) + self.scheduler.set_timesteps(batch.num_inference_steps, device=device) + batch.timesteps = self.scheduler.timesteps + + return batch + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("prompt", batch.prompt, V.string_not_empty) + return result + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("prior_token_id", batch.prior_token_id, V.is_tensor) + return result diff --git a/fastvideo/pipelines/stages/glm_image_decoding.py b/fastvideo/pipelines/stages/glm_image_decoding.py new file mode 100644 index 000000000..10cbbcf3b --- /dev/null +++ b/fastvideo/pipelines/stages/glm_image_decoding.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GLM-Image decoding stage. + +Simple 2D image decoding using diffusers AutoencoderKL. +""" + +import torch +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import VerificationResult + + +class GlmImageDecodingStage(PipelineStage): + """Simple decoding stage for GLM-Image (2D image generation).""" + + def __init__(self, vae, pipeline=None): + super().__init__() + self.vae = vae + self.pipeline = pipeline + + @torch.no_grad() + def forward(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> ForwardBatch: + latents = batch.latents + device = latents.device + + # Squeeze temporal dimension if present [B, C, T, H, W] -> [B, C, H, W] + if latents.dim() == 5: + latents = latents.squeeze(2) + + # GLM-Image uses latent normalization like CogVideoX + # SGLang implementation indicates we should use latents_std and latents_mean + # and ignore the scaling_factor (0.18215) when these are present. + if hasattr(self.vae.config, + 'latents_mean') and self.vae.config.latents_mean is not None: + latents_mean = torch.tensor(self.vae.config.latents_mean, + device=device, + dtype=latents.dtype).view(1, 16, 1, 1) + latents_std = torch.tensor(self.vae.config.latents_std, + device=device, + dtype=latents.dtype).view(1, 16, 1, 1) + # Un-normalize: latents * std + mean + latents = latents * latents_std + latents_mean + elif hasattr(self.vae.config, 'scaling_factor'): + latents = latents / self.vae.config.scaling_factor + + # Ensure VAE is on the same device as latents + self.vae.to(device) + + # Decode latents to image + decoded = self.vae.decode(latents.to(self.vae.dtype)) + + # Handle different return types + # Handle different return types + image = decoded.sample if hasattr(decoded, 'sample') else decoded + + # Normalize from [-1, 1] to [0, 1] + image = (image / 2 + 0.5).clamp(0, 1) + + # Store as output (add fake temporal dim for compatibility) + batch.output = image.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] + + return batch + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + return VerificationResult() # Minimal verification + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + return VerificationResult() # Minimal verification diff --git a/fastvideo/pipelines/stages/glm_image_denoising.py b/fastvideo/pipelines/stages/glm_image_denoising.py new file mode 100644 index 000000000..7a8013d51 --- /dev/null +++ b/fastvideo/pipelines/stages/glm_image_denoising.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GLM-Image denoising stage. + +This stage performs the denoising loop for GLM-Image generation. +""" + +import torch +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.distributed import get_local_torch_device +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + + +class GlmImageDenoisingStage(DenoisingStage): + """Denoising stage for GLM-Image pipeline.""" + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + """Verify denoising stage inputs.""" + result = VerificationResult() + # Basic checks + result.add_check("timesteps", batch.timesteps, + [V.is_tensor, V.min_dims(1)]) + # Use latents or latent + latents = getattr(batch, 'latent', getattr(batch, 'latents', None)) + result.add_check("latents", latents, [V.is_tensor, V.with_dims(5)]) + + result.add_check("num_inference_steps", batch.num_inference_steps, + V.positive_int) + + # We handle prompt_embeds as concatenated [neg, pos] for CFG + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + + # Explicitly skip negative_prompt_embeds logic from base class since we pre-concat + + return result + + def forward(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> ForwardBatch: + device = get_local_torch_device() + dtype = torch.bfloat16 + + # Parameters + num_inference_steps = batch.num_inference_steps + guidance_scale = getattr(batch, 'guidance_scale', 1.5) + + # Get latents from batch + latents = getattr(batch, 'latent', getattr(batch, 'latents', None)) + if latents is None: + raise ValueError("No latents found in batch.") + + # GLM-Image is 2D, squeeze temporal dimension if present [B, C, T, H, W] -> [B, C, H, W] + if latents.dim() == 5: + latents = latents.squeeze(2) + + timesteps = batch.timesteps + prompt_embeds = batch.prompt_embeds[0] # (2, L, D) for CFG + attention_mask = getattr(batch, 'attention_mask', None) + + # Prepare target_size and crop_coords + bs = 2 if guidance_scale > 1.0 else 1 + target_size = torch.tensor([[batch.height, batch.width]], + device=device, + dtype=torch.long).repeat(bs, 1) + crop_coords = torch.zeros((bs, 2), device=device, dtype=torch.long) + + # Prepare prior tokens for CFG + prior_token_id = batch.prior_token_id + if guidance_scale > 1.0 and prior_token_id.shape[0] == 1: + prior_token_id = prior_token_id.repeat(2, 1) + # Create drop mask: [False] for cond, [True] for uncond (batch dimension indexing) + prior_token_drop = torch.tensor([False, True], device=device) + else: + prior_token_drop = getattr(batch, 'prior_token_drop', + torch.tensor([False], device=device)) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # Prepare model input (CFG pass) + if guidance_scale > 1.0: + latent_model_input = torch.cat([latents] * 2) + else: + latent_model_input = latents + + # Scale model input for flow matching / diffusion + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t) + latent_model_input = latent_model_input.to(dtype) + t_expand = t.expand(latent_model_input.shape[0]) + + with torch.no_grad(), set_forward_context(current_timestep=i, + attn_metadata=None, + forward_batch=batch): + # Predict noise + noise_pred = self.transformer(latent_model_input, + prompt_embeds, + t_expand, + prior_token_id, + prior_token_drop, + target_size, + crop_coords, + attention_mask=attention_mask) + + # Apply Classifier-Free Guidance + if guidance_scale > 1.0: + noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * ( + noise_pred_cond - noise_pred_uncond) + + # Apply guidance rescale if requested + guidance_rescale = getattr(batch, 'guidance_rescale', + 0.0) + if guidance_rescale > 0.0: + std_text = noise_pred_cond.std(dim=list( + range(1, noise_pred_cond.ndim)), + keepdim=True) + std_cfg = noise_pred.std(dim=list( + range(1, noise_pred.ndim)), + keepdim=True) + noise_pred_rescaled = noise_pred * (std_text / + std_cfg) + noise_pred = guidance_rescale * noise_pred_rescaled + ( + 1 - guidance_rescale) * noise_pred + + # Step scheduler + latents = self.scheduler.step(noise_pred, + t, + latents, + return_dict=False)[0] + + progress_bar.update() + + # Store result in batch (add temporal dim back for video-compatible decoding) + batch.latents = latents.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] + return batch