Update app.py

app.py

@@ -115,7 +115,7 @@ def custom_get_clip_prompt_embeds(self, prompt, num_images_per_prompt, device):
     return prompt_embeds, pooled_prompt_embeds
 
 # Override the high-level encode_prompt to use T5 encoding and return three outputs.
-# --- KEY CHANGE: Return token_ids as a tuple (text_token_ids, dummy_img_token_ids)
+# --- KEY CHANGE: Return token_ids as a single tensor.
 def custom_encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance=False,
                          negative_prompt=None, prompt_embeds=None, prompt_2=None, **kwargs):
     text_inputs = self.tokenizer(
@@ -142,18 +142,13 @@ def custom_encode_prompt(self, prompt, device, num_images_per_prompt, do_classif
          pooled_uncond_embeddings = uncond_embeddings.mean(dim=1)
          text_embeddings = torch.cat([uncond_embeddings, text_embeddings], dim=0)
          pooled_text_embeddings = torch.cat([pooled_uncond_embeddings, pooled_text_embeddings], dim=0)
-         token_ids_text = text_inputs.input_ids
+         token_ids = text_inputs.input_ids
     else:
-         token_ids_text = text_inputs.input_ids
+         token_ids = text_inputs.input_ids
     text_embeddings = text_embeddings.repeat_interleave(num_images_per_prompt, dim=0)
     pooled_text_embeddings = pooled_text_embeddings.repeat_interleave(num_images_per_prompt, dim=0)
-    token_ids_text = token_ids_text.repeat_interleave(num_images_per_prompt, dim=0)
-    
-    # --- Create dummy image token ids with the same shape as text token ids.
-    dummy_img_ids = torch.full_like(token_ids_text, fill_value=t5_tokenizer.pad_token_id)
-    
-    # Return a tuple of token id tensors.
-    return text_embeddings, pooled_text_embeddings, (token_ids_text, dummy_img_ids)
+    token_ids = token_ids.repeat_interleave(num_images_per_prompt, dim=0)
+    return text_embeddings, pooled_text_embeddings, token_ids
 
 pipe._get_clip_prompt_embeds = custom_get_clip_prompt_embeds.__get__(pipe)
 pipe._encode_prompt = custom_encode_prompt.__get__(pipe)
@@ -171,7 +166,6 @@ def patched_time_embed(self, timestep, guidance, pooled_projections):
     text_out = self.fixed_text_proj(pooled_projections)
     return time_out + text_out
 
-# Patch the forward method.
 pipe.transformer.time_text_embed.forward = patched_time_embed.__get__(pipe.transformer.time_text_embed)
 
 # ----- HISTORY FUNCTIONS & GRADIO INTERFACE -----