Update app.py

app.py

@@ -5,35 +5,66 @@ from torch import autocast
 from diffusers import StableDiffusionPipeline
 from datasets import load_dataset
 from PIL import Image  
+from io import BytesIO
+import base64
 import re
+import os
+import requests
+
+from share_btn import community_icon_html, loading_icon_html, share_js
 
 model_id = "CompVis/stable-diffusion-v1-4"
 device = "cuda"
 
 #If you are running this code locally, you need to either do a 'huggingface-cli login` or paste your User Access Token from here https://huggingface.co/settings/tokens into the use_auth_token field below. 
-pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token="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", revision="fp16", torch_dtype=torch.float16)
+pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True, revision="fp16", torch_dtype=torch.float16)
 pipe = pipe.to(device)
 torch.backends.cudnn.benchmark = True
 
+#When running locally, you won`t have access to this, so you can remove this part
+word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
+word_list = word_list_dataset["train"]['text']
 
-def infer(prompt, samples, steps, scale, seed):
+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 gr.Error("Unsafe content found. Please try again with different prompts.")
         
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    images_list = pipe(
-        [prompt] * samples,
-        num_inference_steps=steps,
-        guidance_scale=scale,
-        generator=generator,
-    )
+    #generator = torch.Generator(device=device).manual_seed(seed)
+    print("Is GPU busy? ", is_gpu_busy)
     images = []
-    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)
-    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_decoded = Image.open(BytesIO(base64.b64decode(image)))
+            images.append(image_decoded)
+    
+    
+    return images, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+    
     
 css = """
         .gradio-container {
@@ -118,6 +149,38 @@ css = """
             font-weight: bold;
             font-size: 115%;
         }
+        #container-advanced-btns{
+            display: flex;
+            flex-wrap: wrap;
+            justify-content: space-between;
+            align-items: center;
+        }
+        .animate-spin {
+            animation: spin 1s linear infinite;
+        }
+        @keyframes spin {
+            from {
+                transform: rotate(0deg);
+            }
+            to {
+                transform: rotate(360deg);
+            }
+        }
+        #share-btn-container {
+            display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
+        }
+        #share-btn {
+            all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;
+        }
+        #share-btn * {
+            all: unset;
+        }
+        .gr-form{
+            flex: 1 1 50%; border-top-right-radius: 0; border-bottom-right-radius: 0;
+        }
+        #prompt-container{
+            gap: 0;
+        }
 """
 
 block = gr.Blocks(css=css)
@@ -160,6 +223,7 @@ examples = [
     ],
 ]
 
+
 with block:
     gr.HTML(
         """
@@ -225,12 +289,13 @@ with block:
     )
     with gr.Group():
         with gr.Box():
-            with gr.Row().style(mobile_collapse=False, equal_height=True):
+            with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
                 text = gr.Textbox(
                     label="Enter your prompt",
                     show_label=False,
                     max_lines=1,
                     placeholder="Enter your prompt",
+                    elem_id="prompt-text-input",
                 ).style(
                     border=(True, False, True, True),
                     rounded=(True, False, False, True),
@@ -239,15 +304,22 @@ with block:
                 btn = gr.Button("Generate image").style(
                     margin=False,
                     rounded=(False, True, True, False),
+                    full_width=False,
                 )
 
         gallery = gr.Gallery(
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
 
-        advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
+        with gr.Group(elem_id="container-advanced-btns"):
+            advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
+            with gr.Group(elem_id="share-btn-container"):
+                community_icon = gr.HTML(community_icon_html, visible=False)
+                loading_icon = gr.HTML(loading_icon_html, visible=False)
+                share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)
 
         with gr.Row(elem_id="advanced-options"):
+            gr.Markdown("Advanced settings are temporarily unavailable")
             samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)
             steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
             scale = gr.Slider(
@@ -261,12 +333,14 @@ with block:
                 randomize=True,
             )
 
-        ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, steps, scale, seed], outputs=gallery, cache_examples=True)
+        ex = gr.Examples(examples=examples, fn=infer, inputs=text, outputs=[gallery, community_icon, loading_icon, share_button], cache_examples=True)
         ex.dataset.headers = [""]
 
         
-        text.submit(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)
-        btn.click(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)
+        text.submit(infer, inputs=text, outputs=[gallery, community_icon, loading_icon, share_button])
+        
+        btn.click(infer, inputs=text, outputs=[gallery, community_icon, loading_icon, share_button])
+        
         advanced_button.click(
             None,
             [],
@@ -277,6 +351,12 @@ with block:
                 options.style.display = ["none", ""].includes(options.style.display) ? "flex" : "none";
             }""",
         )
+        share_button.click(
+            None,
+            [],
+            [],
+            _js=share_js,
+        )
         gr.HTML(
             """
                 <div class="footer">
@@ -292,4 +372,4 @@ Despite how impressive being able to turn text into image is, beware to the fact
            """
         )
 
-block.queue(max_size=25).launch()
+block.queue(max_size=25, concurrency_count=2).launch()