debug statements added

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -403,14 +403,11 @@ def encode_image(self, image: PipelineImageInput, num_videos_per_prompt: int = 1
           image = [image]
       image_inputs = self.image_processor(images=image, return_tensors="np")
       pixel_values = jnp.array(image_inputs.pixel_values)
-      max_logging.log(f"[DEBUG ENC] pixel_values shape: {pixel_values.shape}")
 
       image_encoder_output = self.image_encoder(pixel_values, output_hidden_states=True)
       image_embeds = image_encoder_output.hidden_states[-2]
-      max_logging.log(f"[DEBUG ENC] Shape of image_embeds from image_encoder: {image_embeds.shape}")
-      
+
       image_embeds = jnp.repeat(image_embeds, num_videos_per_prompt, axis=0)
-      max_logging.log(f"[DEBUG ENC] Shape of image_embeds after repeat: {image_embeds.shape}")
       return image_embeds
 
 

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -203,7 +203,12 @@ def __call__(
       t0 = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[0]
       dummy_noise = jnp.zeros_like(latents)
       # This call initializes the internal state arrays
-      _, scheduler_state = self.scheduler.step(scheduler_state, dummy_noise, t0, latents)
+      step_output = self.scheduler.step(scheduler_state, dummy_noise, t0, latents)
+      max_logging.log(f"[DEBUG] scheduler.step output type: {type(step_output)}")
+      max_logging.log(f"[DEBUG] scheduler.step output value: {step_output}")
+      _, scheduler_state = step_output
+      max_logging.log(f"[DEBUG] After prime step: scheduler_state type: {type(scheduler_state)}")
+      max_logging.log(f"[DEBUG] After prime step: scheduler_state value: {scheduler_state}")
       max_logging.log(f"[DEBUG] Scheduler state primed: step_index={scheduler_state.step_index is not None}, last_sample={scheduler_state.last_sample is not None}")
 
     graphdef, state, rest_of_state = nnx.split(self.transformer, nnx.Param, ...)