debugging i2v

examples/inference/basic/basic.py

@@ -14,9 +14,9 @@ def main():
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
         use_fsdp_inference=True,
-        dit_cpu_offload=False,
+        dit_cpu_offload=True,
         vae_cpu_offload=False,
-        text_encoder_cpu_offload=False,
+        text_encoder_cpu_offload=True,
         # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
         pin_cpu_memory=True,
         # image_encoder_cpu_offload=False,

fastvideo/models/dits/wanvideo.py

@@ -303,6 +303,8 @@ class WanTransformerBlock(nn.Module):
         temb: torch.Tensor,
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
     ) -> torch.Tensor:
+        cuda_memory_before = torch.cuda.memory_allocated()
+        logger.info(f"cuda memory 6: {cuda_memory_before / 1024 / 1024 / 1024} GB")
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
         bs, seq_length, _ = hidden_states.shape
@@ -312,10 +314,15 @@ class WanTransformerBlock(nn.Module):
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
             6, dim=1)
         assert shift_msa.dtype == torch.float32
+        cuda_memory_before = torch.cuda.memory_allocated()
+        logger.info(f"cuda memory 7: {cuda_memory_before / 1024 / 1024 / 1024} GB")
 
         # 1. Self-attention
+        logger.info(f"hidden_states shape: {hidden_states.shape}")
         norm_hidden_states = (self.norm1(hidden_states.float()) *
                               (1 + scale_msa) + shift_msa).to(orig_dtype)
+        cuda_memory_before = torch.cuda.memory_allocated()
+        logger.info(f"cuda memory 8: {cuda_memory_before / 1024 / 1024 / 1024} GB")
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
@@ -603,6 +610,8 @@ class WanTransformer3DModel(CachableDiT):
                 **kwargs) -> torch.Tensor:
         forward_batch = get_forward_context().forward_batch
         enable_teacache = forward_batch is not None and forward_batch.enable_teacache
+        cuda_memory_before = torch.cuda.memory_allocated()
+        logger.info(f"cuda memory 1: {cuda_memory_before / 1024 / 1024 / 1024} GB")
 
         orig_dtype = hidden_states.dtype
         if not isinstance(encoder_hidden_states, torch.Tensor):
@@ -638,9 +647,15 @@ class WanTransformer3DModel(CachableDiT):
         hidden_states = self.patch_embedding(hidden_states)
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
+        cuda_memory_before = torch.cuda.memory_allocated()
+        logger.info(f"cuda memory 2: {cuda_memory_before / 1024 / 1024 / 1024} GB")
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
             timestep, encoder_hidden_states, encoder_hidden_states_image)
+        cuda_memory_before = torch.cuda.memory_allocated()
+        logger.info(f"cuda memory 3: {cuda_memory_before / 1024 / 1024 / 1024} GB")
         timestep_proj = timestep_proj.unflatten(1, (6, -1))
+        cuda_memory_before = torch.cuda.memory_allocated()
+        logger.info(f"cuda memory 4: {cuda_memory_before / 1024 / 1024 / 1024} GB")
 
         if encoder_hidden_states_image is not None:
             encoder_hidden_states = torch.concat(
@@ -671,6 +686,8 @@ class WanTransformer3DModel(CachableDiT):
                         timestep_proj, freqs_cis)
             else:
                 for block in self.blocks:
+                    cuda_memory_before = torch.cuda.memory_allocated()
+                    logger.info(f"cuda memory 5: {cuda_memory_before / 1024 / 1024 / 1024} GB")
                     hidden_states = block(hidden_states, encoder_hidden_states,
                                           timestep_proj, freqs_cis)
 

fastvideo/pipelines/stages/denoising.py

@@ -245,7 +245,9 @@ class DenoisingStage(PipelineStage):
                     logger.info(
                         f"mask[0][0][:, ::2, ::2] shape: {mask2[0][0][:, ::2, ::2].shape}"
                     )
-                    temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
+                    temp_ts = (mask2[0][0][:, ::2, ::2] * timestep)
+                    logger.info(f"temp_ts shape before flatten: {temp_ts.shape}")
+                    temp_ts = temp_ts.flatten()
                     logger.info(f"temp_ts: {temp_ts}")
                     logger.info(f"temp_ts shape: {temp_ts.shape}")
                     temp_ts = torch.cat([
@@ -320,6 +322,8 @@ class DenoisingStage(PipelineStage):
                             # fastvideo_args=fastvideo_args
                     ):
                         # Run transformer
+                        cuda_memory_before = torch.cuda.memory_allocated()
+                        logger.info(f"cuda memory before transformer: {cuda_memory_before / 1024 / 1024 / 1024} GB")
                         noise_pred = self.transformer(
                             latent_model_input,
                             prompt_embeds,