From 5c68193009c551dba503eff3f7dfa5affd04bbb2 Mon Sep 17 00:00:00 2001 From: ShreejithSG Date: Wed, 28 Jan 2026 05:57:34 +0000 Subject: [PATCH] Refactor GLM-Image decoding stage to inherit from our DecodingStage --- .../pipelines/stages/glm_image_decoding.py | 88 +++++++------------ .../pipelines/stages/glm_image_denoising.py | 5 +- 2 files changed, 35 insertions(+), 58 deletions(-) diff --git a/fastvideo/pipelines/stages/glm_image_decoding.py b/fastvideo/pipelines/stages/glm_image_decoding.py index 10cbbcf3b..da62908a3 100644 --- a/fastvideo/pipelines/stages/glm_image_decoding.py +++ b/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 diff --git a/fastvideo/pipelines/stages/glm_image_denoising.py b/fastvideo/pipelines/stages/glm_image_denoising.py index bfddcf4b7..344626569 100644 --- a/fastvideo/pipelines/stages/glm_image_denoising.py +++ b/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(