Update app.py

app.py

@@ -1,31 +1,21 @@
 import gradio as gr
-from datasets import load_dataset
 from PIL import Image  
 
 import re
 import os
 import requests
 
+from muse import PipelineMuse
+
 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']
+device = "cuda" if torch.cuda.is_available()
+pipe = PipelineMuse.from_pretrained("openMUSE/muse-cc12m-200k").to(device)
+if device == "cuda":
+    pipe.transformer.enable_xformers_memory_efficient_attention()
 
-is_gpu_busy = False
-def infer(prompt, negative, scale):
-    global is_gpu_busy
-    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 = []
-    url = os.getenv('JAX_BACKEND_URL')
-    payload = {'prompt': prompt, 'negative_prompt': negative, '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)
-    
+def infer(prompt, negative, scale):    
+    images = pipe(prompt, timesteps=12, guidance_scale=scale, num_images_per_prompt=4, use_fp16=device == "cuda")
     return images