Update app.py

app.py

@@ -63,11 +63,45 @@ class TextProjection(torch.nn.Module):
     def forward(self, x):
         return self.proj(x.to(dtype))
 
+# Add this override to your existing pipeline setup
+class T5FluxPipeline(FluxPipeline):
+    def _get_clip_prompt_embeds(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance):
+        """Modified to work with T5 outputs"""
+        # Get T5 embeddings
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=512,
+            truncation=True,
+            return_tensors="pt",
+        ).to(device)
+        
+        text_outputs = self.text_encoder(**text_inputs)
+        prompt_embeds = text_outputs.last_hidden_state
+        
+        # Use mean pooling instead of CLIP's pooler_output
+        pooled_prompt_embeds = prompt_embeds.mean(dim=1)
+        
+        # Expand for batch
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+
+        # Handle guidance
+        if do_classifier_free_guidance:
+            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
+            negative_pooled = torch.zeros_like(pooled_prompt_embeds)
+            
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+            pooled_prompt_embeds = torch.cat([negative_pooled, pooled_prompt_embeds])
+
+        return prompt_embeds, pooled_prompt_embeds
+
+
 # Initialize pipeline components
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
 # Custom pipeline with T5 support
-pipe = DiffusionPipeline.from_pretrained(
+pipe = T5FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     text_encoder=t5_text_encoder,
     tokenizer=t5_tokenizer,