Update app.py

app.py

@@ -1,10 +1,26 @@
+from diffusers import StableDiffusionInpaintPipeline
 import gradio as gr
 import numpy as np
 import imageio
 from PIL import Image
+from io import BytesIO
+import os
+
+MY_SECRET_TOKEN=os.environ.get('HF_TOKEN_SD')
+
+
+print("hello sylvain")
+
+YOUR_TOKEN=MY_SECRET_TOKEN
+
+device="cpu"
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
+pipe.to(device)
 
 source_img = gr.Image(source="upload", type="numpy", tool="sketch", elem_id="source_container");
-outputs = [gr.outputs.Image(type="file",label="output"),gr.outputs.Image(type="file",label="Mask")]
+gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
+
 def resize(height,img):
   baseheight = height
   img = Image.open(img)
@@ -14,12 +30,6 @@ def resize(height,img):
   return img
 
 def predict(source_img):
-    
-    #print(sketch)
-    #print(sketch.mode)
-    #sketch_png = resize(512,source_img)
-    #sketch_png.save('source.png')
-    #print(sketch_png)
     imageio.imwrite("data.png", source_img["image"])
     imageio.imwrite("data_mask.png", source_img["mask"]) 
     
@@ -27,8 +37,17 @@ def predict(source_img):
     src.save("src.png")
     mask = resize(512, "data_mask.png")  
     mask.save("mask.png")
-    return src, mask
+    
+    images_list = img_pipe([prompt] * 1, init_image=src, mask_image=mask, strength=0.75)
+    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
 
 custom_css="style.css"
 
-gr.Interface(fn=predict, inputs=source_img, outputs=outputs, css=custom_css).launch(enable_queue=True)
+gr.Interface(fn=predict, inputs=source_img, outputs=gallery, css=custom_css).launch(enable_queue=True)