anal

app.py

@@ -2,30 +2,90 @@ from click import command
 import streamlit as st 
 import os
 from authtoken import auth_token
+from datasets import load_dataset
+from PIL import Image  
+import re
+import os
+import requests
+
+model_id = "CompVis/stable-diffusion-v1-4"
+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']
+
+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
+    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.")
+        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 
+    
+    prompt=st.text_input('Enter your prompt')
+    true=st.button('Generate')
+    if true==True:
+        infer(prompt)
+
+    
+
+
+
+
+
+
+
+
+
 
 
-from diffusers import StableDiffusionPipeline
-st.session_state.ac=True 
-if st.session_state.ac==True:
-    os.system('pip3 install torch torchvision torchaudio')
-    import torch
-    from torch import autocast  
-    st.session_state.ac=False
 
-st.title('Stable Diffusion (AI that can imagine)')
+
+
+
+
+
+
+
+
+
+
+
+
+
+# from diffusers import StableDiffusionPipeline
+# st.session_state.ac=True 
+# if st.session_state.ac==True:
+#     os.system('pip3 install torch torchvision torchaudio')
+#     import torch
+#     from torch import autocast  
+#     st.session_state.ac=False
+
+# st.title('Stable Diffusion (AI that can imagine)')
 
     
 
-prompt=st.text_input('Enter the prompt')
-modelid = "CompVis/stable-diffusion-v1-4"
-device = "cuda"
-pipe = StableDiffusionPipeline.from_pretrained(modelid, revision="fp16", torch_dtype=torch.float16, use_auth_token=auth_token) 
-pipe.to(device) 
-def generate(): 
-    with autocast(device): 
-        image = pipe(prompt.get(), guidance_scale=8.5)["sample"][0]
+# prompt=st.text_input('Enter the prompt')
+# modelid = "CompVis/stable-diffusion-v1-4"
+# device = "cuda"
+# pipe = StableDiffusionPipeline.from_pretrained(modelid, revision="fp16", torch_dtype=torch.float16, use_auth_token=auth_token) 
+# pipe.to(device) 
+# def generate(): 
+#     with autocast(device): 
+#         image = pipe(prompt.get(), guidance_scale=8.5)["sample"][0]
     
-    image.save(prompt+'.png')
-    img = st.image(prompt+'.png')
+#     image.save(prompt+'.png')
+#     img = st.image(prompt+'.png')
 
-st.button('Generate',command=generate)
+# st.button('Generate',command=generate)