anal

app.py

@@ -1,5 +1,8 @@
 from click import command
 import streamlit as st 
+import torch
+from torch import autocast
+from diffusers import StableDiffusionPipeline
 import os
 from authtoken import auth_token
 from datasets import load_dataset
@@ -13,30 +16,57 @@ device = "cuda"
 # 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(model_id, use_auth_token=auth_token, revision="fp16", torch_dtype=torch.float16)
+pipe = pipe.to(device)
+torch.backends.cudnn.benchmark = True
+
 is_gpu_busy = False
+
 def infer(prompt):
     global is_gpu_busy
     samples = 4
     steps = 50
     scale = 7.5
-    # #When running locally you can also remove this filter
+    #When running locally you can also remove this filter
     # for filter in word_list:
     #     if re.search(rf"\b{filter}\b", prompt):
-    #         raise st.Error("Unsafe content found. Please try again with different prompts.")
+    #         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)
     images = []
-    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)
-    st.image(images)
-    return images 
+    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)
+    
+    return images
     
 prompt=st.text_input('Enter your prompt')
 true=st.button('Generate')
 if true==True:
-    infer(prompt)
+    a=infer(prompt)
+    for i in a:
+        st.image(a)