Update app.py

app.py

@@ -9,11 +9,11 @@ pipe.to("cpu")  # Use "cuda" if GPU is available
 unet = pipe.unet
 
 def generate_image(prompt, unet, pipe):
-    # Encode the prompt
-    text_encoding = pipe.text_encoder(prompt, return_tensors="pt").to(unet.device)
+    # Tokenize the prompt
+    tokens = pipe.tokenizer(prompt, padding=True, truncation=True, max_length=50, return_tensors="pt").to(unet.device)
     
     # Generate the image
-    image = unet(text_embeddings=text_encoding.last_hidden_state).images[0]
+    image = unet(text_embeddings=tokens.input_ids).images[0]
     return image
 
 def chatbot(prompt):