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.

Author: ShreejithSG

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()

fastvideo/attention/layer.py

@@ -1,8 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import TYPE_CHECKING
+
 import torch
 import torch.nn as nn
 
+from fastvideo.attention.backends.sdpa import SDPAMetadata
 from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
 from fastvideo.distributed.communication_op import (
     sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
@@ -301,6 +304,7 @@ def forward(
         k: torch.Tensor,
         v: torch.Tensor,
         freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """
         Apply local attention between query, key and value tensors.
@@ -324,6 +328,12 @@ def forward(
             cos, sin = freqs_cis
             q = _apply_rotary_emb(q, cos, sin, is_neox_style=False)
             k = _apply_rotary_emb(k, cos, sin, is_neox_style=False)
+        
+        if attention_mask is not None:
+            if ctx_attn_metadata is None:
+                ctx_attn_metadata = SDPAMetadata(current_timestep=0, attn_mask=attention_mask)
+            else:
+                ctx_attn_metadata.attn_mask = attention_mask
 
         output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
         return output

fastvideo/attention/utils/flash_attn_no_pad.py

@@ -29,8 +29,15 @@ def flash_attn_no_pad(qkv,
     seqlen = qkv.shape[1]
     nheads = qkv.shape[-2]
     x = rearrange(qkv, "b s three h d -> b s (three h d)")
-    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(
-        x, key_padding_mask)
+    
+    # Handle version differences in flash_attn unpad_input
+    unpad_output = unpad_input(x, key_padding_mask)
+    if len(unpad_output) == 5:
+        x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_output
+    else:
+        # Older versions return 4 values or different order? 
+        # Typically: x_unpad, indices, cu_seqlens, max_s
+        x_unpad, indices, cu_seqlens, max_s = unpad_output[:4]
 
     x_unpad = rearrange(x_unpad,
                         "nnz (three h d) -> nnz three h d",
@@ -69,12 +76,29 @@ def flash_attn_no_pad_v3(qkv,
     batch_size, seqlen, _, nheads, head_dim = qkv.shape
     query, key, value = qkv.unbind(dim=2)
 
-    query_unpad, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+    # Handle version differences in flash_attn unpad_input for query
+    unpad_output_q = unpad_input(
         rearrange(query, "b s h d -> b s (h d)"), key_padding_mask)
-    key_unpad, _, cu_seqlens_k, _, _ = unpad_input(
+    if len(unpad_output_q) >= 4:
+        query_unpad, indices, cu_seqlens_q, max_seqlen_q = unpad_output_q[:4]
+    else:
+        raise ValueError(f"Unexpected unpad_input output length: {len(unpad_output_q)}")
+
+    # Handle version differences for key
+    unpad_output_k = unpad_input(
         rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
-    value_unpad, _, _, _, _ = unpad_input(
+    if len(unpad_output_k) >= 3:
+         key_unpad, _, cu_seqlens_k = unpad_output_k[:3]
+    else:
+         raise ValueError(f"Unexpected unpad_input output length: {len(unpad_output_k)}")
+
+    # Handle version differences for value
+    unpad_output_v = unpad_input(
         rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
+    if len(unpad_output_v) >= 1:
+        value_unpad = unpad_output_v[0]
+    else:
+        raise ValueError(f"Unexpected unpad_input output length: {len(unpad_output_v)}")
 
     query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=nheads)
     key_unpad = rearrange(key_unpad, "nnz (h d) -> nnz h d", h=nheads)

fastvideo/configs/models/base.py

@@ -51,9 +51,7 @@ def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
         for key, value in source_model_dict.items():
             if key in valid_fields:
                 setattr(arch_config, key, value)
-            else:
-                raise AttributeError(
-                    f"{type(arch_config).__name__} has no field '{key}'")
+            # Skip unknown fields - HF configs often have extra fields we don't need
 
         if hasattr(arch_config, "__post_init__"):
             arch_config.__post_init__()

fastvideo/configs/models/dits/__init__.py

@@ -1,5 +1,6 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
+from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
@@ -9,5 +10,5 @@
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
     "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
-    "LongCatVideoConfig"
+    "LongCatVideoConfig", "GlmImageDiTConfig"
 ]

fastvideo/configs/models/dits/glm_image.py

@@ -0,0 +1,95 @@
+# 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"

fastvideo/configs/models/vaes/__init__.py

@@ -1,4 +1,5 @@
 from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
+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.stepvideovae import StepVideoVAEConfig
@@ -10,4 +11,5 @@
     "StepVideoVAEConfig",
     "CosmosVAEConfig",
     "Hunyuan15VAEConfig",
+    "GlmImageVAEConfig",
 ]

fastvideo/configs/models/vaes/glm_image.py

@@ -0,0 +1,53 @@
+# 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

fastvideo/configs/pipelines/base.py

@@ -73,6 +73,10 @@ class PipelineConfig:
                                   ...] = field(default_factory=lambda:
                                                (postprocess_text, ))
 
+    # Post-decoding hook for custom processing after VAE decoding
+    # This allows pipelines to apply custom transformations to decoded images/videos
+    post_decoding: Callable[[torch.Tensor], torch.Tensor] | None = None
+
     # StepVideo specific parameters
     pos_magic: str | None = None
     neg_magic: str | None = None