From 62f0866a464b62dcc2b66ff4e5a8e8eb4a5b3e5d Mon Sep 17 00:00:00 2001 From: Mihir Jagtap Date: Thu, 5 Feb 2026 19:32:10 -0800 Subject: [PATCH] Create gamecraft_denoising.py --- .../pipelines/stages/gamecraft_denoising.py | 346 ++++++++++++++++++ 1 file changed, 346 insertions(+) create mode 100644 fastvideo/pipelines/stages/gamecraft_denoising.py diff --git a/fastvideo/pipelines/stages/gamecraft_denoising.py b/fastvideo/pipelines/stages/gamecraft_denoising.py new file mode 100644 index 000000000..a15f8712e --- /dev/null +++ b/fastvideo/pipelines/stages/gamecraft_denoising.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GameCraft-specific denoising stage for hybrid history conditioning. + +HunyuanGameCraft uses a 33-channel input format: +- 16 channels: noisy latents +- 16 channels: conditioning/history latents +- 1 channel: mask indicating history vs predicted frames + +The model outputs 16-channel noise predictions which are used to update +only the non-history portion of the latents. +""" + +import inspect +import weakref + +import torch +from tqdm.auto import tqdm + +from fastvideo.attention import get_attn_backend +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.platforms import AttentionBackendEnum + +logger = init_logger(__name__) + + +class GameCraftDenoisingStage(PipelineStage): + """ + Denoising stage for HunyuanGameCraft with hybrid history conditioning. + + This stage handles the special 33-channel input format required by GameCraft: + - Concatenates latents (16ch) + conditioning (16ch) + mask (1ch) + - Passes 33-channel tensor to transformer + - Gets 16-channel noise prediction + - Updates latents via scheduler step + + For text-to-video (no history), conditioning is zeros and mask is ones. + For video continuation, conditioning contains history frames and mask + indicates which frames are history (0) vs predicted (1). + """ + + def __init__( + self, + transformer, + scheduler, + pipeline=None, + vae=None, + ) -> None: + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + self.vae = vae + self.pipeline = weakref.ref(pipeline) if pipeline else None + + # Get attention backend + attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads + self.attn_backend = get_attn_backend( + head_size=attn_head_size, + dtype=torch.float16, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + ) + + def _prepare_conditioning_latents( + self, + latents: torch.Tensor, + history_latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Prepare conditioning latents and mask for GameCraft. + + Args: + latents: Current noisy latents [B, 16, T, H, W] + history_latents: Optional history latents for video continuation + + Returns: + conditioning_latents: [B, 16, T, H, W] - zeros for t2v, history for continuation + mask: [B, 1, T, H, W] - ones for t2v, mixed for continuation + """ + batch_size, channels, num_frames, height, width = latents.shape + device = latents.device + dtype = latents.dtype + + if history_latents is None: + # Text-to-video mode: no history conditioning + # The model was trained with first half as "history" (mask=1) and second half as "generate" (mask=0) + # For T2V, we use zeros for conditioning but still use the expected mask pattern + conditioning_latents = torch.zeros( + batch_size, channels, num_frames, height, width, + device=device, dtype=dtype + ) + # Create mask: first half = 1 (history placeholder), second half = 0 (generate) + # This matches the training format even for T2V + mask = torch.zeros( + batch_size, 1, num_frames, height, width, + device=device, dtype=dtype + ) + num_history = num_frames // 2 + mask[:, :, :num_history, :, :] = 1.0 # First half marked as "history" + else: + # Video continuation mode: use history as conditioning + conditioning_latents = history_latents.clone() + + # Create mask: 0 for history frames, 1 for frames to generate + # Typically first half are history, second half are generated + mask = torch.ones( + batch_size, 1, num_frames, height, width, + device=device, dtype=dtype + ) + + # For continuation, first frame(s) come from history + if num_frames > 1: + # First half is history (mask=0), second half is predicted (mask=1) + num_history = num_frames // 2 + mask[:, :, :num_history, :, :] = 0.0 + conditioning_latents[:, :, num_history:, :, :] = 0.0 + else: + # Single frame: use first frame from history + mask[:, :, 0, :, :] = 0.0 + + return conditioning_latents, mask + + def prepare_extra_func_kwargs(self, func, kwargs): + """Filter kwargs to only include parameters the function accepts.""" + # Get the function's signature + sig = inspect.signature(func) + valid_keys = set(sig.parameters.keys()) + + # Filter to only valid kwargs + return {k: v for k, v in kwargs.items() if k in valid_keys and v is not None} + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """ + Run the GameCraft denoising loop with hybrid conditioning. + + Args: + batch: The current batch information + fastvideo_args: The inference arguments + + Returns: + batch: Updated batch with denoised latents + """ + device = get_local_torch_device() + target_dtype = torch.bfloat16 + + # Get batch data + latents = batch.latents # [B, 16, T, H, W] + timesteps = batch.timesteps + + # Get prompt embeddings from batch (it's a list for multiple encoders) + # Convert to target dtype + prompt_embeds = [ + emb.to(target_dtype) if emb is not None else None + for emb in batch.prompt_embeds + ] + neg_prompt_embeds = [ + emb.to(target_dtype) if emb is not None else None + for emb in (batch.negative_prompt_embeds or []) + ] + + # Get CLIP pooled embeddings if available + clip_embedding_pos = getattr(batch, 'clip_embedding_pos', None) + if clip_embedding_pos is not None: + clip_embedding_pos = clip_embedding_pos.to(target_dtype) + clip_embedding_neg = getattr(batch, 'clip_embedding_neg', None) + if clip_embedding_neg is not None: + clip_embedding_neg = clip_embedding_neg.to(target_dtype) + + # Get history latents if available (for video continuation) + history_latents = getattr(batch, 'history_latents', None) + + # Prepare conditioning latents and mask + conditioning_latents, mask = self._prepare_conditioning_latents( + latents, history_latents + ) + + # Keep a copy of conditioning for each step + gt_latents_concat = conditioning_latents.clone() + mask_concat = mask + + # Prepare extra step kwargs for scheduler + extra_step_kwargs = {} + if "generator" in set( + inspect.signature(self.scheduler.step).parameters.keys() + ): + extra_step_kwargs["generator"] = batch.generator + + # Get camera states for CameraNet conditioning + # Convert camera_states to target dtype and device + camera_states = batch.camera_states + if camera_states is not None: + camera_states = camera_states.to(device=device, dtype=target_dtype) + camera_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + {"camera_states": camera_states}, + ) + + # Prepare positional conditioning kwargs + pos_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_2": clip_embedding_pos, + "encoder_attention_mask": batch.prompt_attention_mask, + }, + ) + + neg_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_2": clip_embedding_neg, + "encoder_attention_mask": batch.negative_attention_mask, + }, + ) + + # Denoising loop + num_inference_steps = len(timesteps) + num_warmup_steps = max( + len(timesteps) - num_inference_steps * self.scheduler.order, 0 + ) + + with tqdm(total=num_inference_steps, desc="Denoising") as progress_bar: + for i, t in enumerate(timesteps): + # Debug: print timestep and latent stats on first step + if i == 0: + logger.info(f"[DEBUG] Timestep 0: t={t.item() if hasattr(t, 'item') else t}") + logger.info(f"[DEBUG] Latents shape: {latents.shape}, mean: {latents.mean().item():.4f}, std: {latents.std().item():.4f}") + logger.info(f"[DEBUG] Mask shape: {mask_concat.shape}, unique values: {mask_concat.unique().tolist()}") + + # Concatenate latents with conditioning for 33-channel input + # [B, 16, T, H, W] + [B, 16, T, H, W] + [B, 1, T, H, W] = [B, 33, T, H, W] + latents_concat = torch.cat( + [latents, gt_latents_concat, mask_concat], dim=1 + ) + + latent_model_input = latents_concat.to(target_dtype) + + # Scale model input + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t + ) + + # Expand timestep + t_expand = t.unsqueeze(0) if t.dim() == 0 else t + + # Get primary text embeddings (LLaMA) + current_prompt_embeds = prompt_embeds[0] if len(prompt_embeds) > 0 else None + + # Run transformer forward pass with positive conditioning + batch.is_cfg_negative = False + with set_forward_context( + current_timestep=i, + attn_metadata=None, + forward_batch=batch, + ): + noise_pred = self.transformer( + latent_model_input, + current_prompt_embeds, + t_expand, + guidance=None, + **pos_cond_kwargs, + **camera_kwargs, + ) + + # Debug: print noise prediction stats on first step + if i == 0: + logger.info(f"[DEBUG] Noise pred shape: {noise_pred.shape}, mean: {noise_pred.mean().item():.4f}, std: {noise_pred.std().item():.4f}") + + # Apply classifier-free guidance if enabled + if batch.do_classifier_free_guidance: + # Get negative embeddings + current_neg_embeds = neg_prompt_embeds[0] if neg_prompt_embeds and len(neg_prompt_embeds) > 0 else None + + batch.is_cfg_negative = True + with set_forward_context( + current_timestep=i, + attn_metadata=None, + forward_batch=batch, + ): + noise_pred_uncond = self.transformer( + latent_model_input, + current_neg_embeds, + t_expand, + guidance=None, + **neg_cond_kwargs, + **camera_kwargs, + ) + + noise_pred = noise_pred_uncond + batch.guidance_scale * ( + noise_pred - noise_pred_uncond + ) + + # Debug: print noise prediction after CFG on first step + if i == 0: + logger.info(f"[DEBUG] After CFG - Noise pred mean: {noise_pred.mean().item():.4f}, std: {noise_pred.std().item():.4f}") + + # Scheduler step - only update non-history frames (second half) + # The model was trained to always keep first half fixed and generate second half + # noise_pred is [B, 16, T, H, W], latents is [B, 16, T, H, W] + num_frames = latents.shape[2] + num_history = num_frames // 2 + + if num_frames > 1 and num_history > 0: + # Only update second half of frames (the generated portion) + latents[:, :, num_history:, :, :] = self.scheduler.step( + noise_pred[:, :, num_history:, :, :], + t, + latents[:, :, num_history:, :, :], + **extra_step_kwargs, + return_dict=False + )[0] + else: + # Single frame or edge case: update all + latents = self.scheduler.step( + noise_pred, + t, + latents, + **extra_step_kwargs, + return_dict=False + )[0] + + # Debug: print latent stats after scheduler step + if i == 0 or i == len(timesteps) - 1: + logger.info(f"[DEBUG] Step {i} - Latents after step: mean: {latents.mean().item():.4f}, std: {latents.std().item():.4f}") + + # Update progress bar + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and + (i + 1) % self.scheduler.order == 0 + ): + progress_bar.update() + + # Update batch with final latents + batch.latents = latents + + return batch