Update app.py

app.py

@@ -21,7 +21,7 @@ word_list = word_list_dataset["train"]['text']
 
 is_gpu_busy = False
 def infer(prompt):
-    #global is_gpu_busy
+    global is_gpu_busy
     samples = 4
     steps = 50
     scale = 7.5
@@ -30,31 +30,31 @@ def infer(prompt):
         if re.search(rf"\b{filter}\b", prompt):
             raise gr.Error("Unsafe content found. Please try again with different prompts.")
         
-    #generator = torch.Generator(device=device).manual_seed(seed)
-    #print("Is GPU busy? ", is_gpu_busy)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    print("Is GPU busy? ", is_gpu_busy)
     images = []
-    #if(not is_gpu_busy):
-    #    is_gpu_busy = True
-    #    images_list = pipe(
-    #        [prompt] * samples,
-    #        num_inference_steps=steps,
-    #        guidance_scale=scale,
+    if(not is_gpu_busy):
+        is_gpu_busy = True
+        images_list = pipe(
+            [prompt] * samples,
+            num_inference_steps=steps,
+            guidance_scale=scale,
             #generator=generator,
-    #    )
-    #    is_gpu_busy = False
-    #    safe_image = Image.open(r"unsafe.png")
-    #    for i, image in enumerate(images_list["sample"]):
-     #       if(images_list["nsfw_content_detected"][i]):
-     #           images.append(safe_image)
-     #       else:
-     #           images.append(image)
-    #else:
-    url = os.getenv('JAX_BACKEND_URL')
-    payload = {'prompt': prompt}
-    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)
+        )
+        is_gpu_busy = False
+        safe_image = Image.open(r"unsafe.png")
+        for i, image in enumerate(images_list["sample"]):
+            if(images_list["nsfw_content_detected"][i]):
+                images.append(safe_image)
+            else:
+                images.append(image)
+    else:
+        url = os.getenv('JAX_BACKEND_URL')
+        payload = {'prompt': prompt}
+        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)
     
     return images