Refactor GLM-Image decoding stage to inherit from our DecodingStage

fastvideo/pipelines/stages/glm_image_decoding.py

@@ -1,73 +1,51 @@
 # SPDX-License-Identifier: Apache-2.0
 """
 GLM-Image decoding stage.
-
-Simple 2D image decoding using diffusers AutoencoderKL.
 """
 
 import torch
+from fastvideo.distributed import get_local_torch_device
 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
+from fastvideo.pipelines.stages.decoding import DecodingStage
+from fastvideo.utils import PRECISION_TO_TYPE
 
 
-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
+class GlmImageDecodingStage(DecodingStage):
+    """Decoding stage for GLM-Image (2D image generation)."""
 
     @torch.no_grad()
-    def forward(self, batch: ForwardBatch,
-                fastvideo_args: FastVideoArgs) -> ForwardBatch:
-        latents = batch.latents
-        device = latents.device
+    def decode(self, latents: torch.Tensor,
+               fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        """
+        Decode latent representations into pixel space using VAE.
+        
+        Args:
+            latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)
+            fastvideo_args: Configuration containing VAE precision and tiling settings
+            
+        Returns:
+            Decoded image tensor with shape (batch, channels, frames, height, width),
+            normalized to [0, 1] range
+        """
+        self.vae = self.vae.to(get_local_torch_device())
+        latents = latents.to(get_local_torch_device())
+
+        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        latents = self._denormalize_latents(latents)
 
-        # 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
+        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            if not vae_autocast_enabled:
+                latents = latents.to(vae_dtype)
+            decoded = self.vae.decode(latents)
 
-        # 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
+        image = image.unsqueeze(2)
+        return image

fastvideo/pipelines/stages/glm_image_denoising.py

@@ -15,7 +15,6 @@ from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.attention.backends.sdpa import SDPAMetadata
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.attention.selector import backend_name_to_enum
 
 
 class GlmImageDenoisingStage(DenoisingStage):
@@ -109,9 +108,9 @@ class GlmImageDenoisingStage(DenoisingStage):
                 attn_metadata = None
                 if text_attention_mask is not None:
                     first_block = self.transformer.transformer_blocks[0]
-                    attn_backend_name = first_block.attn1.attn.backend.name if hasattr(first_block.attn1.attn, 'backend') else None
+                    attn_backend = getattr(first_block.attn1.attn, 'backend', None)
                     
-                    if attn_backend_name == "SDPA" or backend_name_to_enum(attn_backend_name) == AttentionBackendEnum.TORCH_SDPA:
+                    if attn_backend is not None and attn_backend == AttentionBackendEnum.TORCH_SDPA:
                         batch_size = latent_model_input.shape[0]
                         # Create combined mask: text part uses provided mask, image part is always 1
                         mix_attn_mask = torch.ones(