Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 from datasets import load_dataset
 from PIL import Image  
+import torch
+from diffusers import StableDiffusionPipeline
 
 import re
 import os
@@ -8,28 +10,17 @@ import requests
 
 from share_btn import community_icon_html, loading_icon_html, share_js
 
-word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
-word_list = word_list_dataset["train"]['text']
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16) 
 
 is_gpu_busy = False
 def infer(person1, person2, scale):
     global is_gpu_busy
     #FIXME: did somebody say f"prompt" injections? :P
     prompt = f"Cross between {person1} and {person2}"
-    for filter in word_list:
-        if re.search(rf"\b{filter}\b", prompt):
-            raise gr.Error("Unsafe content found. Please try again with different prompts.")
 
-    images = []
-    person_1 = ''
-    url = os.getenv('JAX_BACKEND_URL')
-    payload = {'prompt': prompt, 'guidance_scale': scale}
-    images_request = requests.post(url, json = payload)
-    for image in images_request.json()["images"]:
-        image_b64 = (f"data:image/jpeg;base64,{image}")
-        images.append(image_b64)
+    image = pipe(prompt).images[0] 
     
-    return images
+    return image
     
     
 css = """
@@ -300,11 +291,11 @@ with block:
         #        randomize=True,
         #    )
 
-        ex = gr.Examples(examples=examples, fn=infer, inputs=[person1, person2, guidance_scale], outputs=[gallery, community_icon, loading_icon, share_button], cache_examples=False)
+        ex = gr.Examples(examples=examples, fn=infer, inputs=[person1, person2], outputs=[gallery, community_icon, loading_icon, share_button], cache_examples=False)
         ex.dataset.headers = [""]
-        negative.submit(infer, inputs=[person1, person2, guidance_scale], outputs=[gallery], postprocess=False)
-        text.submit(infer, inputs=[person1, person2, guidance_scale], outputs=[gallery], postprocess=False)
-        btn.click(infer, inputs=[person1, person2, guidance_scale], outputs=[gallery], postprocess=False)
+        negative.submit(infer, inputs=[person1, person2], outputs=[gallery], postprocess=False)
+        text.submit(infer, inputs=[person1, person2], outputs=[gallery], postprocess=False)
+        btn.click(infer, inputs=[person1, person2], outputs=[gallery], postprocess=False)
         
         #advanced_button.click(
         #    None,