not-lain commited on
Commit
ba5420f
·
1 Parent(s): 74c1094

add inpaint

Browse files
Files changed (1) hide show
  1. app.py +46 -3
app.py CHANGED
@@ -50,7 +50,7 @@ def prepare_image_and_mask(
50
  return background, mask
51
 
52
 
53
- def inpaint(
54
  image,
55
  padding_top=0,
56
  padding_bottom=0,
@@ -79,6 +79,31 @@ def inpaint(
79
  return result
80
 
81
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82
  def rmbg(image=None, url=None):
83
  if image is None:
84
  image = url
@@ -102,6 +127,8 @@ def main(*args, progress=gr.Progress(track_tqdm=True)):
102
  if api_num == 1:
103
  return rmbg(*args)
104
  elif api_num == 2:
 
 
105
  return inpaint(*args)
106
 
107
 
@@ -131,9 +158,25 @@ outpaint_tab = gr.Interface(
131
  api_name="outpainting",
132
  )
133
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
134
  demo = gr.TabbedInterface(
135
- [rmbg_tab, outpaint_tab],
136
- ["remove background", "outpainting"],
137
  title="Utilities that require GPU",
138
  )
139
 
 
50
  return background, mask
51
 
52
 
53
+ def outpaint(
54
  image,
55
  padding_top=0,
56
  padding_bottom=0,
 
79
  return result
80
 
81
 
82
+ def inpaint(
83
+ image,
84
+ mask,
85
+ prompt="",
86
+ num_inference_steps=28,
87
+ guidance_scale=50,
88
+ ):
89
+ background = image.convert("RGB")
90
+ mask = mask.convert("L")
91
+
92
+ result = pipe(
93
+ prompt=prompt,
94
+ height=background.height,
95
+ width=background.width,
96
+ image=background,
97
+ mask_image=mask,
98
+ num_inference_steps=num_inference_steps,
99
+ guidance_scale=guidance_scale,
100
+ ).images[0]
101
+
102
+ result = result.convert("RGBA")
103
+
104
+ return result
105
+
106
+
107
  def rmbg(image=None, url=None):
108
  if image is None:
109
  image = url
 
127
  if api_num == 1:
128
  return rmbg(*args)
129
  elif api_num == 2:
130
+ return outpaint(*args)
131
+ elif api_num == 3:
132
  return inpaint(*args)
133
 
134
 
 
158
  api_name="outpainting",
159
  )
160
 
161
+
162
+ inpaint_tab = gr.Interface(
163
+ fn=main,
164
+ inputs=[
165
+ gr.Number(3, visible=False),
166
+ gr.Image(label="image"),
167
+ gr.Image(label="mask"),
168
+ gr.Text(label="prompt"),
169
+ gr.Number(value=50, label="num_inference_steps"),
170
+ gr.Number(value=28, label="guidance_scale"),
171
+ ],
172
+ outputs=["image"],
173
+ api_name="inpaint",
174
+ examples=[[3, "./assets/rocket.png", "./assets/Inpainting mask.png"]],
175
+ )
176
+
177
  demo = gr.TabbedInterface(
178
+ [rmbg_tab, outpaint_tab, inpaint_tab],
179
+ ["remove background", "outpainting", "inpainting"],
180
  title="Utilities that require GPU",
181
  )
182