Update app.py

app.py

@@ -6,32 +6,16 @@ model_id = "SG161222/RealVisXL_V4.0"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 pipe.to("cpu")  # Use "cuda" if GPU is available
 
-unet = pipe.unet
-
-def generate_image(prompt, unet, pipe):
-    # Tokenize the prompt
-    tokens = pipe.tokenizer(prompt, padding=True, truncation=True, max_length=50, return_tensors="pt").to(unet.device)
-    
-    # Generate the image
-    text_embeds = pipe.text_encoder(tokens.input_ids)
-    image = unet(text_embeds=text_embeds).images[0]
+def generate_image(prompt, pipe):
+    # Generate the image using the pipeline
+    image = pipe(prompt).images[0]
     return image
 
 def chatbot(prompt):
     # Generate the image based on the user's input
-    image = generate_image(prompt, unet, pipe)
+    image = generate_image(prompt, 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(
     fn=chatbot,
@@ -42,4 +26,4 @@ interface = gr.Interface(
 )
 
 # Launch the interface
-interface.launch()
+interface.launch()