Add addt logging

Author: kevin314

fastvideo/pipelines/stages/denoising.py

@@ -660,19 +660,31 @@ def forward(
                 pipeline.add_module("transformer", self.transformer)
             fastvideo_args.model_loaded["transformer"] = True
 
-        # Setup precision and autocast settings
-        target_dtype = torch.bfloat16
+        # Setup precision to match diffusers exactly
+        # Diffusers uses transformer.dtype (bfloat16) and converts inputs before transformer calls
+        # For FSDP wrapped models, we need to access the underlying module
+        if hasattr(self.transformer, 'module'):
+            transformer_dtype = next(self.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.transformer.parameters()).dtype
+        target_dtype = transformer_dtype
         autocast_enabled = (target_dtype != torch.float32
                             ) and not fastvideo_args.disable_autocast
 
         # Get latents and setup
         latents = batch.latents
         num_inference_steps = batch.num_inference_steps
         guidance_scale = batch.guidance_scale
+
+        sum_value = latents.float().sum().item()
+        # Write to output file
+        with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+            f.write(f"Denoising init: latents sum = {sum_value:.6f}, shape = {latents.shape}\n")
 
 
-        # Setup scheduler timesteps like Diffusers does
-        # Diffusers uses set_timesteps without custom sigmas, letting the scheduler generate them
+        # Setup scheduler timesteps - use default scheduler sigma generation
+        # The torch.linspace(0, 1, num_inference_steps) approach was incorrect for FlowMatchEulerDiscreteScheduler
+        # Let the scheduler generate its own sigmas using the configured sigma_max, sigma_min, etc.
         self.scheduler.set_timesteps(num_inference_steps, device=latents.device)
         timesteps = self.scheduler.timesteps
 

fastvideo/pipelines/stages/latent_preparation.py

@@ -306,6 +306,18 @@ def forward(
             if self.vae is not None:
                 # Move VAE to correct device before encoding
                 self.vae = self.vae.to(device)
+                
+                # Log VAE info and input video stats
+                print(f"[FASTVIDEO VAE DEBUG] VAE model: {type(self.vae).__name__}")
+                print(f"[FASTVIDEO VAE DEBUG] VAE config z_dim: {self.vae.config.z_dim}")
+                print(f"[FASTVIDEO VAE DEBUG] Input video shape: {video.shape}, dtype: {video.dtype}, device: {video.device}")
+                print(f"[FASTVIDEO VAE DEBUG] Input video sum: {video.float().sum().item():.6f}")
+                with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+                    f.write(f"FastVideo VAE: model_type = {type(self.vae).__name__}\n")
+                    f.write(f"FastVideo VAE: z_dim = {self.vae.config.z_dim}\n")
+                    f.write(f"FastVideo VAE: input_video_shape = {video.shape}\n")
+                    f.write(f"FastVideo VAE: input_video_sum = {video.float().sum().item():.6f}\n")
+                
                 if isinstance(generator, list):
                     init_latents = []
                     for i in range(batch_size):
@@ -361,9 +373,17 @@ def forward(
                     latents_std = torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype)
                     print(f"[FASTVIDEO CONDITIONING DEBUG] latents_mean = {self.vae.config.latents_mean}, latents_std = {self.vae.config.latents_std}")
                     print(f"[FASTVIDEO CONDITIONING DEBUG] scheduler.sigma_data = {self.scheduler.sigma_data}")
+                    with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+                        f.write(f"FastVideo Conditioning: scheduler.sigma_data = {self.scheduler.sigma_data}\n")
+                        f.write(f"FastVideo Conditioning: latents_mean = {self.vae.config.latents_mean}\n")
+                        f.write(f"FastVideo Conditioning: latents_std = {self.vae.config.latents_std}\n")
                     print(f"[FASTVIDEO CONDITIONING DEBUG] Before normalization sum = {init_latents.float().sum().item()}")
+                    with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+                        f.write(f"FastVideo Conditioning: before_normalization_sum = {init_latents.float().sum().item():.6f}\n")
                     init_latents = (init_latents - latents_mean) / latents_std * self.scheduler.sigma_data
                     print(f"[FASTVIDEO CONDITIONING DEBUG] After normalization sum = {init_latents.float().sum().item()}")
+                    with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+                        f.write(f"FastVideo Conditioning: after_normalization_sum = {init_latents.float().sum().item():.6f}\n")
 
                 conditioning_latents = init_latents
                 print(f"[FASTVIDEO CONDITIONING DEBUG] Final conditioning_latents sum = {conditioning_latents.float().sum().item()}")
@@ -441,6 +461,38 @@ def forward(
         if conditioning_latents is not None:
             logger.info(f"CosmosLatentPreparationStage - conditioning_latents shape: {conditioning_latents.shape}")
 
+        # Log tensor sums to fastvideo_hidden_states.log
+        sum_value = latents.float().sum().item()
+        print(f"FastVideo LatentPreparation: latents sum = {sum_value:.6f}")
+        with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+            f.write(f"FastVideo LatentPreparation: latents sum = {sum_value:.6f}\n")
+        
+        if conditioning_latents is not None:
+            sum_value = conditioning_latents.float().sum().item()
+            print(f"FastVideo LatentPreparation: conditioning_latents sum = {sum_value:.6f}")
+            with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+                f.write(f"FastVideo LatentPreparation: conditioning_latents sum = {sum_value:.6f}\n")
+        
+        sum_value = cond_indicator.float().sum().item()
+        print(f"FastVideo LatentPreparation: cond_indicator sum = {sum_value:.6f}")
+        with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+            f.write(f"FastVideo LatentPreparation: cond_indicator sum = {sum_value:.6f}\n")
+        
+        sum_value = uncond_indicator.float().sum().item()
+        print(f"FastVideo LatentPreparation: uncond_indicator sum = {sum_value:.6f}")
+        with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+            f.write(f"FastVideo LatentPreparation: uncond_indicator sum = {sum_value:.6f}\n")
+        
+        sum_value = cond_mask.float().sum().item()
+        print(f"FastVideo LatentPreparation: cond_mask sum = {sum_value:.6f}")
+        with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+            f.write(f"FastVideo LatentPreparation: cond_mask sum = {sum_value:.6f}\n")
+        
+        sum_value = uncond_mask.float().sum().item()
+        print(f"FastVideo LatentPreparation: uncond_mask sum = {sum_value:.6f}")
+        with open("/workspace/FastVideo/fastvideo_hidden_states.log", "a") as f:
+            f.write(f"FastVideo LatentPreparation: uncond_mask sum = {sum_value:.6f}\n")
+
         return batch
 
     def verify_input(self, batch: ForwardBatch,