cleanup

app.py

@@ -1,5 +1,6 @@
 import pathlib
 import os
+from PIL import Image
 
 import gradio as gr
 import torch
@@ -8,6 +9,7 @@ from diffusers import StableDiffusionPipeline
 import utils
 
 use_auth_token = os.environ["HF_AUTH_TOKEN"]
+NSFW_IMAGE = Image.open("nsfw.png")
 
 # Instantiate the pipeline.
 device, revision, torch_dtype = (
@@ -51,19 +53,25 @@ def replace_concept_tokens(text: str):
         text = text.replace(concept_token, dummy_tokens)
     return text
 
-
+all_generated_images = []
 def inference(
     prompt: str, guidance_scale: int, num_inference_steps: int, seed: int
 ):
     prompt = replace_concept_tokens(prompt)
     generator = torch.Generator(device=device).manual_seed(seed)
-    img_list = pipeline(
+    output = pipeline(
         prompt=[prompt] * 2,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
         generator=generator,
-    ).images
-    return img_list
+    )
+    img_list, nsfw_list = output.images, output.nsfw_content_detected
+    for img, nsfw in zip(img_list, nsfw_list):
+        if nsfw:
+            all_generated_images.append(NSFW_IMAGE)
+        else:
+            all_generated_images.append(img)
+    return all_generated_images
 
 DEFAULT_PROMPT = (
     "A watercolor painting on textured paper of a <det-logo> using soft strokes,"
@@ -77,11 +85,11 @@ with gr.Blocks() as demo:
         interactive=True,
     )
     guidance_scale = gr.Slider(
-        minimum=1.0, maximum=10.0, value=3.0, label="Guidance Scale", interactive=True
+        minimum=1.0, maximum=20.0, value=3.0, label="Guidance Scale", interactive=True
     )
     num_inference_steps = gr.Slider(
         minimum=25,
-        maximum=60,
+        maximum=75,
         value=40,
         label="Num Inference Steps",
         interactive=True,
@@ -90,18 +98,12 @@ with gr.Blocks() as demo:
     seed = gr.Slider(
         minimum=2147483147,
         maximum=2147483647,
-        value=2147483397,
         label="Seed",
         interactive=True,
+        randomize=True
     )
-    # output = gr.Textbox(
-    #     label="output", placeholder=use_auth_token[:5], interactive=False
-    # )
-    # gr.Button("test").click(
-    #     lambda s: replace_concept_tokens(s), inputs=[prompt], outputs=output
-    # )
 
-    generate_btn = gr.Button(label="Generate")
+    generate_btn = gr.Button(value="Generate")
     gallery = gr.Gallery(
         label="Generated Images",
         value=[],