Update app.py

app.py

@@ -184,6 +184,8 @@ def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width,
 
     # update_state(f"Done. Seed: {seed}")
 
+    return replace_nsfw_images(result)
+
 def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed):
 
     print(f"{datetime.datetime.now()} img_to_img, model: {model_path}")
@@ -234,6 +236,23 @@ def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance
         callback=pipe_callback)
 
     # update_state(f"Done. Seed: {seed}")
+
+    return replace_nsfw_images(result)
+
+    def replace_nsfw_images(results):
+
+    if is_colab:
+        return results.images
+
+    for i in range(len(results.images)):
+        if results.nsfw_content_detected[i]:
+            orig_img = results.images[i]
+            results.images[i] = Image.open("nsfw.png")
+            nsfw_label = f"NSFW: {results.nsfw_content_prob[i]:.2f}"
+            img_array = np.concatenate([np.array(orig_img), np.array(results.images[i])], axis=1)
+            display(Image.fromarray(img_array), caption=nsfw_label)
+
+    return results.images
         
 
 # css = """.finetuned-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.finetuned-diffusion-div div h1{font-weight:900;margin-bottom:7px}.finetuned-diffusion-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}