Address more PR review comments

Author: ShreejithSG

examples/inference/basic/basic_glm_image.py

@@ -3,72 +3,56 @@
 GLM-Image inference example using FastVideo.
 
 This example demonstrates how to generate images using the GLM-Image model
-through FastVideo's pipeline infrastructure.
+through FastVideo's VideoGenerator API.
 
 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 PIL import Image
 
-from fastvideo.pipelines.basic.glm_image import GlmImagePipeline
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo import VideoGenerator
 
 OUTPUT_PATH = "image_output"
 
 
-
 def main():
-    # Load the GLM-Image pipeline
-    pipe = GlmImagePipeline.from_pretrained(
+    # FastVideo will automatically use the optimal default arguments for GLM-Image.
+    generator = VideoGenerator.from_pretrained(
         "zai-org/GLM-Image",
         num_gpus=1,
         trust_remote_code=True,
     )
 
-    # Create a batch with generation parameters
+    # Generate images - use save_video=False and return_frames=True
+    # to get raw frames that we can save as images
     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="",
+    os.makedirs(OUTPUT_PATH, exist_ok=True)
+    
+    # Generate first image
+    result = generator.generate_video(
+        prompt,
+        output_path=OUTPUT_PATH,
+        save_video=False,
+        return_frames=True,
         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,
+        guidance_scale=7.5,  
     )
 
-    # 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}")
+    # Save as PNG image (result is a list of frames, take the first/only one)
+    if result and len(result) > 0:
+        img = Image.fromarray(result[0])
+        img.save(os.path.join(OUTPUT_PATH, "landscape.png"))
+        print(f"Saved image to {OUTPUT_PATH}/landscape.png")
 
 
 if __name__ == "__main__":

fastvideo/attention/utils/flash_attn_no_pad.py

@@ -30,14 +30,7 @@ def flash_attn_no_pad(qkv,
     nheads = qkv.shape[-2]
     x = rearrange(qkv, "b s three h d -> b s (three h d)")
 
-    # 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, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
 
     x_unpad = rearrange(x_unpad,
                         "nnz (three h d) -> nnz three h d",
@@ -76,32 +69,12 @@ def flash_attn_no_pad_v3(qkv,
     batch_size, seqlen, _, nheads, head_dim = qkv.shape
     query, key, value = qkv.unbind(dim=2)
 
-    # 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)
-    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)
-    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, indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+        rearrange(query, "b s h d -> b s (h d)"), key_padding_mask)
+    key_unpad, _, cu_seqlens_k, _, _ = unpad_input(
+        rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
+    value_unpad, _, _, _, _ = unpad_input(
+        rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
 
     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,7 +51,9 @@ 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)
-            # Skip unknown fields - HF configs often have extra fields we don't need
+            else:
+                raise AttributeError(
+                    f"{type(arch_config).__name__} has no field '{key}'")
 
         if hasattr(arch_config, "__post_init__"):
             arch_config.__post_init__()

fastvideo/models/dits/glm_image.py

@@ -423,9 +423,9 @@ def forward(
             # 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)
+            # Convert to SDPA format: (B, 1, 1, L) for key-padding style mask
+            # True = attend, False = ignore (will be converted to additive mask by SDPA)
+            attention_mask_kv = (mix_attn_mask > 0).unsqueeze(1).unsqueeze(2)  # (B, 1, 1, L)
         else:
             attention_mask_kv = None
 

fastvideo/pipelines/composed_pipeline_base.py

@@ -287,6 +287,7 @@ def load_modules(
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
+        model_index.pop("_name_or_path", None)
         model_index.pop("workload_type", None)
         if "boundary_ratio" in model_index and model_index[
                 "boundary_ratio"] is not None:
@@ -330,8 +331,6 @@ def load_modules(
 
         modules = {}
         for module_name, value in model_index.items():
-            if module_name.startswith("_"):
-                continue
             transformers_or_diffusers, architecture = value
             if transformers_or_diffusers is None:
                 logger.warning(

fastvideo/pipelines/stages/decoding.py

@@ -244,10 +244,10 @@ def forward(
 
         # Apply post-decoding hook if configured
         # This allows pipelines to apply custom transformations after VAE decoding
-        if hasattr(fastvideo_args.pipeline_config, 'post_decoding') and \
-           fastvideo_args.pipeline_config.post_decoding is not None:
+        post_decoding_hook = getattr(fastvideo_args.pipeline_config, 'post_decoding', None)
+        if post_decoding_hook is not None:
             logger.debug("Applying post_decoding hook")
-            frames = fastvideo_args.pipeline_config.post_decoding(frames)
+            frames = post_decoding_hook(frames)
 
         # Convert to CPU float32 for compatibility
         frames = frames.cpu().float()