Fix GLM-Image seed handling and add SSIM test

fastvideo/pipelines/stages/glm_image_before_denoising.py

@@ -203,6 +203,13 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         # 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
+        
+        # Set seed for deterministic latent generation
+        if batch.seed is not None:
+            torch.manual_seed(batch.seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(batch.seed)
+        
         latents = torch.randn((1, 16, 1, batch.height // 8, batch.width // 8),
                               device=device,
                               dtype=dtype)

fastvideo/tests/ssim/test_glm_image_similarity.py

@@ -1,15 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-"""
-SSIM-based similarity tests for GLM-Image generation.
-
-Tests GLM-Image text-to-image generation with various prompts and backends.
-"""
+"""SSIM-based regression tests for GLM-Image generation."""
 
 import os
 
 import pytest
 import torch
-import hashlib
+from PIL import Image
+import numpy as np
 
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
@@ -18,26 +15,24 @@ logger = init_logger(__name__)
 
 # Device-specific reference folder
 device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = "_reference_videos"
-
 if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
+    device_reference_folder = "A40_reference_videos"
 elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
+    device_reference_folder = "L40S_reference_videos"
 elif "H100" in device_name:
-    device_reference_folder = "H100" + device_reference_folder_suffix
+    device_reference_folder = "H100_reference_videos"
 else:
-    logger.warning(f"Unsupported device for ssim tests: {device_name}")
+    logger.warning(f"Unsupported device: {device_name}, using L40S references")
+    device_reference_folder = "L40S_reference_videos"
 
 # Test prompts
 TEST_PROMPTS = [
     "A beautiful landscape photography with rolling hills, a winding river, and a vibrant sunset in the background. Warm golden light, photorealistic style.",
-    "A majestic mountain landscape at golden hour with a calm lake reflecting snow-capped peaks",
 ]
 
 # Fixed reference filenames (avoid hashing prompts)
 REFERENCE_FILENAMES = {
-    TEST_PROMPTS[0]: "landscape_ref.png"
+    TEST_PROMPTS[0]: "landscape_ref.png",
 }
 
 # GLM-Image parameters
@@ -114,7 +109,7 @@ def _ssim_torch(x: torch.Tensor, y: torch.Tensor, window_size: int = 11, sigma:
 
 
 @pytest.mark.parametrize("prompt", TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"])
+@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
 def test_glm_image_similarity(prompt, ATTENTION_BACKEND):
     """
     Test GLM-Image generation similarity against reference images.

generate_diffusers_reference.py

@@ -1,83 +0,0 @@
-#!/usr/bin/env python3
-"""
-Generate reference image using official diffusers GlmImagePipeline.
-
-This script generates a golden reference image that will be committed to the repo
-for SSIM regression testing.
-"""
-
-import os
-import hashlib
-import torch
-from diffusers.pipelines.glm_image import GlmImagePipeline
-from PIL import Image
-
-# Test prompts (matching test_glm_image_similarity.py)
-TEST_PROMPTS = [
-    "A beautiful landscape photography with rolling hills, a winding river, and a vibrant sunset in the background. Warm golden light, photorealistic style.",
-    "A majestic mountain landscape at golden hour with a calm lake reflecting snow-capped peaks",
-]
-
-# Fixed reference filenames matching the SSIM test.
-REFERENCE_FILENAMES = {
-    TEST_PROMPTS[0]: "landscape_ref.png",
-    TEST_PROMPTS[1]: "mountain_ref.png",
-}
-
-# GLM-Image parameters (matching test_glm_image_similarity.py)
-GLM_IMAGE_PARAMS = {
-    "height": 1024,
-    "width": 1024,
-    "num_inference_steps": 50,
-    "guidance_scale": 7.5,
-    "seed": 1024,
-}
-
-# Determine device folder
-device_name = torch.cuda.get_device_name()
-if "A40" in device_name:
-    device_folder = "A40_reference_videos"
-elif "L40S" in device_name:
-    device_folder = "L40S_reference_videos"
-elif "H100" in device_name:
-    device_folder = "H100_reference_videos"
-else:
-    print(f"Warning: Unsupported device {device_name}, using H100_reference_videos")
-    device_folder = "H100_reference_videos"
-
-# Setup paths
-base_dir = os.path.dirname(os.path.abspath(__file__))
-reference_dir = os.path.join(base_dir, "fastvideo", "tests", "ssim", device_folder, "glm_image")
-os.makedirs(reference_dir, exist_ok=True)
-
-print(f"Loading GLM-Image pipeline from zai-org/GLM-Image...")
-pipe = GlmImagePipeline.from_pretrained(
-    "zai-org/GLM-Image",
-    torch_dtype=torch.bfloat16,
-    device_map="cuda",
-)
-
-print(f"Generating reference images...")
-for prompt in TEST_PROMPTS:
-    image_filename = REFERENCE_FILENAMES[prompt]
-    reference_path = os.path.join(reference_dir, image_filename)
-    
-    print(f"\nGenerating: {prompt[:50]}...")
-    print(f"  Output: {reference_path}")
-    
-    generator = torch.Generator(device="cuda").manual_seed(GLM_IMAGE_PARAMS["seed"])
-    
-    result = pipe(
-        prompt=prompt,
-        height=GLM_IMAGE_PARAMS["height"],
-        width=GLM_IMAGE_PARAMS["width"],
-        num_inference_steps=GLM_IMAGE_PARAMS["num_inference_steps"],
-        guidance_scale=GLM_IMAGE_PARAMS["guidance_scale"],
-        generator=generator,
-    )
-    
-    img = result.images[0]
-    img.save(reference_path)
-    print(f"Saved to {reference_path}")
-
-print(f"\n✓ All reference images generated in {reference_dir}")