Update app.py

app.py

@@ -8,15 +8,28 @@ pipe.to("cpu")  # Use "cuda" if GPU is available
 
 unet = pipe.unet
 
-def generate_image(prompt, unet):
-    added_cond_kwargs = {"text_embeds": pipe.get_text_embedding(prompt)}
-    image = unet(prompt, **added_cond_kwargs).images[0]
+def generate_image(prompt, unet, pipe):
+    # Encode the prompt
+    text_encoding = pipe.text_encoder(prompt, return_tensors="pt").to(unet.device)
+    
+    # Generate the image
+    image = unet(text_embeddings=text_encoding.last_hidden_state).images[0]
     return image
 
 def chatbot(prompt):
     # Generate the image based on the user's input
-    image = generate_image(prompt, unet)
+    image = generate_image(prompt, unet, pipe)
     return image
+
+def get_aug_embed(self, text_embeds, image):
+    if text_embeds is None:
+        text_embeds = self.text_encoder(
+            text_embeds=text_embeds,
+            image=image,
+            height=self.unet.config.sample_size,
+            width=self.unet.config.sample_size,
+        )
+    return text_embeds
     
 # Create the Gradio interface
 interface = gr.Interface(