Manjushri commited on
Commit
98f4d1a
1 Parent(s): fe55238

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +39 -0
app.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import DiffusionPipeline
2
+ import gradio as gr
3
+ import numpy as np
4
+ import imageio
5
+ from PIL import Image
6
+ import torch
7
+
8
+ device = "cuda" if torch.cuda.is_available() else "cpu"
9
+
10
+ pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting")
11
+ pipe.to(device)
12
+
13
+ source_img = gr.Image(source="upload", type="numpy", tool="sketch", elem_id="source_container");
14
+
15
+ def resize(height,img):
16
+ baseheight = height
17
+ img = Image.open(img)
18
+ hpercent = (baseheight/float(img.size[1]))
19
+ wsize = int((float(img.size[0])*float(hpercent)))
20
+ img = img.resize((wsize,baseheight), Image.LANCZOS)
21
+ return img
22
+
23
+ def predict(source_img, prompt):
24
+ imageio.imwrite("data.png", source_img["image"])
25
+ imageio.imwrite("data_mask.png", source_img["mask"])
26
+
27
+ src = resize(512, "data.png")
28
+ src.save("src.png")
29
+ mask = resize(512, "data_mask.png")
30
+ mask.save("mask.png")
31
+
32
+ image = pipe(prompt, image=src, mask_image=mask, strength=0.75, num_inference_steps=10).images[0]
33
+
34
+ return image
35
+
36
+
37
+ title="Stable Diffusion 2.0 Inpainting CPU"
38
+ description="Inpainting with Stable Diffusion 2.0 <br />Warning: Slow process... ~5-10 min inference time.<br> <b>Please use 512*512 or 768x768 square .png image as input to avoid memory error!!!</b>"
39
+ gr.Interface(fn=predict, inputs=[source_img, "text"], outputs='image', title=title, description=description).launch(debug=True)