some debug added to t5_prompt_embeds

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -419,6 +419,8 @@ def _get_t5_prompt_embeds(
       num_videos_per_prompt: int = 1,
       max_sequence_length: int = 226,
   ):
+    jax.debug.print("Prompt Shape")
+    jax.debug.print(prompt.shape)
     prompt = [prompt] if isinstance(prompt, str) else prompt
     prompt = [prompt_clean(u) for u in prompt]
     batch_size = len(prompt)
@@ -433,6 +435,12 @@ def _get_t5_prompt_embeds(
         return_tensors="pt",
     )
     text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+    jax.debug.print("Text Input IDS")
+    jax.debug.print(text_input_ids)
+    jax.debug.print(text_input_ids.shape)
+    jax.debug.print("Mask")
+    jax.debug.print(mask)
+    jax.debug.print(mask.shape)
     seq_lens = mask.gt(0).sum(dim=1).long()
     prompt_embeds = self.text_encoder(text_input_ids, mask).last_hidden_state
     prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]