Update app.py

app.py

@@ -11,6 +11,7 @@ from PIL import Image
 from typing import List, Optional, Union
 import inspect
 import warnings
+import sentence_transformers
 
 # Load the Stable Diffusion model
 modelid = "CompVis/stable-diffusion-v1-4"
@@ -19,6 +20,8 @@ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(modelid, revision="fp16",
 pipe.to(device)
 
 url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+# Load the Sentence-BERT model for text embeddings
+text_embedding_model = sentence_transformers.SentenceTransformer("paraphrase-MiniLM-L6-v2")
 
 
 def generate_image(prompt):
@@ -28,8 +31,9 @@ def generate_image(prompt):
 
     generator = torch.Generator(device=device).manual_seed(1024)
     with autocast("cuda"):
-        prompt = torch.tensor(prompt, device=device).half()
-        image = pipe(prompt=prompt, init_image=init_img, strength=0.75, guidance_scale=7.5, generator=generator).images[0]
+        prompt_embedding = text_embedding_model.encode([prompt])[0]
+        prompt_tensor = torch.tensor(prompt_embedding, device=device).half()
+        image = pipe(prompt=prompt_tensor, init_image=init_img, strength=0.75, guidance_scale=7.5, generator=generator).images[0]
 
     return image