removed debug from i2v 2.1

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -236,25 +236,14 @@ def __call__(
           scheduler_state=scheduler_state,
           rng=inference_rng,
       )
-      max_logging.log(f"[DEBUG CALL] latents shape after loop: {latents.shape}")
-      max_logging.log(f"[DEBUG CALL] NaNs in latents AFTER loop: {jnp.isnan(latents).any()}, Infs: {jnp.isinf(latents).any()}")
       if self.config.expand_timesteps:
          latents = (1 - first_frame_mask) * condition + first_frame_mask * latents
-      max_logging.log(f"[DEBUG CALL] NaNs in latents AFTER frame mask: {jnp.isnan(latents).any()}")
       latents_bcthw = jnp.transpose(latents, (0, 4, 1, 2, 3))
-      max_logging.log(f"[DEBUG CALL] NaNs in latents BEFORE denorm: {jnp.isnan(latents_bcthw).any()}")
       latents_denorm_bcthw = self._denormalize_latents(latents_bcthw)
-      max_logging.log(f"[DEBUG CALL] NaNs in latents AFTER denorm: {jnp.isnan(latents_denorm_bcthw).any()}")
-
 
     if output_type == "latent":
       return jnp.transpose(latents_denorm_bcthw, (0, 2, 3, 4, 1))
-    max_logging.log(f"[DEBUG CALL] NaNs in latents BEFORE decode: {jnp.isnan(latents_denorm_bcthw).any()}")
     decoded_video = self._decode_latents_to_video(latents_denorm_bcthw)
-    if isinstance(decoded_video, np.ndarray):
-      max_logging.log(f"[DEBUG CALL] NaNs in video AFTER decode: {np.isnan(decoded_video).any()}")
-    else:
-      max_logging.log(f"[DEBUG CALL] Decoded video type: {type(decoded_video)}")
     return decoded_video