levihsu commited on
Commit
909b6ba
β€’
1 Parent(s): 1385c1b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +261 -0
app.py ADDED
@@ -0,0 +1,261 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
+ from pathlib import Path
4
+ import sys
5
+ import torch
6
+ from PIL import Image, ImageOps
7
+
8
+ from utils_ootd import get_mask_location
9
+
10
+ PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
11
+ sys.path.insert(0, str(PROJECT_ROOT))
12
+
13
+ from preprocess.openpose.run_openpose import OpenPose
14
+ from preprocess.humanparsing.run_parsing import Parsing
15
+ from ootd.inference_ootd_hd import OOTDiffusionHD
16
+ from ootd.inference_ootd_dc import OOTDiffusionDC
17
+
18
+ import spaces
19
+
20
+
21
+ openpose_model_hd = OpenPose(0)
22
+ parsing_model_hd = Parsing(0)
23
+ ootd_model_hd = OOTDiffusionHD(0)
24
+
25
+ openpose_model_dc = OpenPose(1)
26
+ parsing_model_dc = Parsing(1)
27
+ ootd_model_dc = OOTDiffusionDC(1)
28
+
29
+
30
+ category_dict = ['upperbody', 'lowerbody', 'dress']
31
+ category_dict_utils = ['upper_body', 'lower_body', 'dresses']
32
+
33
+ example_path = os.path.join(os.path.dirname(__file__), 'run/examples')
34
+ model_hd = os.path.join(example_path, 'model/model_1.png')
35
+ garment_hd = os.path.join(example_path, 'garment/03244_00.jpg')
36
+ model_dc = os.path.join(example_path, 'model/model_8.png')
37
+ garment_dc = os.path.join(example_path, 'garment/048554_1.jpg')
38
+
39
+ @spaces.GPU
40
+ def process_hd(vton_img, garm_img, n_samples, n_steps, image_scale, seed):
41
+ model_type = 'hd'
42
+ category = 0 # 0:upperbody; 1:lowerbody; 2:dress
43
+
44
+ with torch.no_grad():
45
+ garm_img = Image.open(garm_img).resize((768, 1024))
46
+ vton_img = Image.open(vton_img).resize((768, 1024))
47
+ keypoints = openpose_model_hd(vton_img.resize((384, 512)))
48
+ model_parse, _ = parsing_model_hd(vton_img.resize((384, 512)))
49
+
50
+ mask, mask_gray = get_mask_location(model_type, category_dict_utils[category], model_parse, keypoints)
51
+ mask = mask.resize((768, 1024), Image.NEAREST)
52
+ mask_gray = mask_gray.resize((768, 1024), Image.NEAREST)
53
+
54
+ masked_vton_img = Image.composite(mask_gray, vton_img, mask)
55
+
56
+ images = ootd_model_hd(
57
+ model_type=model_type,
58
+ category=category_dict[category],
59
+ image_garm=garm_img,
60
+ image_vton=masked_vton_img,
61
+ mask=mask,
62
+ image_ori=vton_img,
63
+ num_samples=n_samples,
64
+ num_steps=n_steps,
65
+ image_scale=image_scale,
66
+ seed=seed,
67
+ )
68
+
69
+ return images
70
+
71
+ def process_dc(vton_img, garm_img, category, n_samples, n_steps, image_scale, seed):
72
+ model_type = 'dc'
73
+ if category == 'Upper-body':
74
+ category = 0
75
+ elif category == 'Lower-body':
76
+ category = 1
77
+ else:
78
+ category =2
79
+
80
+ with torch.no_grad():
81
+ garm_img = Image.open(garm_img).resize((768, 1024))
82
+ vton_img = Image.open(vton_img).resize((768, 1024))
83
+ keypoints = openpose_model_dc(vton_img.resize((384, 512)))
84
+ model_parse, _ = parsing_model_dc(vton_img.resize((384, 512)))
85
+
86
+ mask, mask_gray = get_mask_location(model_type, category_dict_utils[category], model_parse, keypoints)
87
+ mask = mask.resize((768, 1024), Image.NEAREST)
88
+ mask_gray = mask_gray.resize((768, 1024), Image.NEAREST)
89
+
90
+ masked_vton_img = Image.composite(mask_gray, vton_img, mask)
91
+
92
+ images = ootd_model_dc(
93
+ model_type=model_type,
94
+ category=category_dict[category],
95
+ image_garm=garm_img,
96
+ image_vton=masked_vton_img,
97
+ mask=mask,
98
+ image_ori=vton_img,
99
+ num_samples=n_samples,
100
+ num_steps=n_steps,
101
+ image_scale=image_scale,
102
+ seed=seed,
103
+ )
104
+
105
+ return images
106
+
107
+
108
+ block = gr.Blocks().queue()
109
+ with block:
110
+ with gr.Row():
111
+ gr.Markdown("# OOTDiffusion Demo")
112
+ with gr.Row():
113
+ gr.Markdown("## Half-body")
114
+ with gr.Row():
115
+ gr.Markdown("***Support upper-body garments***")
116
+ with gr.Row():
117
+ with gr.Column():
118
+ vton_img = gr.Image(label="Model", sources='upload', type="filepath", height=384, value=model_hd)
119
+ example = gr.Examples(
120
+ inputs=vton_img,
121
+ examples_per_page=14,
122
+ examples=[
123
+ os.path.join(example_path, 'model/model_1.png'),
124
+ os.path.join(example_path, 'model/model_2.png'),
125
+ os.path.join(example_path, 'model/model_3.png'),
126
+ os.path.join(example_path, 'model/model_4.png'),
127
+ os.path.join(example_path, 'model/model_5.png'),
128
+ os.path.join(example_path, 'model/model_6.png'),
129
+ os.path.join(example_path, 'model/model_7.png'),
130
+ os.path.join(example_path, 'model/01008_00.jpg'),
131
+ os.path.join(example_path, 'model/07966_00.jpg'),
132
+ os.path.join(example_path, 'model/05997_00.jpg'),
133
+ os.path.join(example_path, 'model/02849_00.jpg'),
134
+ os.path.join(example_path, 'model/14627_00.jpg'),
135
+ os.path.join(example_path, 'model/09597_00.jpg'),
136
+ os.path.join(example_path, 'model/01861_00.jpg'),
137
+ ])
138
+ with gr.Column():
139
+ garm_img = gr.Image(label="Garment", sources='upload', type="filepath", height=384, value=garment_hd)
140
+ example = gr.Examples(
141
+ inputs=garm_img,
142
+ examples_per_page=14,
143
+ examples=[
144
+ os.path.join(example_path, 'garment/03244_00.jpg'),
145
+ os.path.join(example_path, 'garment/00126_00.jpg'),
146
+ os.path.join(example_path, 'garment/03032_00.jpg'),
147
+ os.path.join(example_path, 'garment/06123_00.jpg'),
148
+ os.path.join(example_path, 'garment/02305_00.jpg'),
149
+ os.path.join(example_path, 'garment/00055_00.jpg'),
150
+ os.path.join(example_path, 'garment/00470_00.jpg'),
151
+ os.path.join(example_path, 'garment/02015_00.jpg'),
152
+ os.path.join(example_path, 'garment/10297_00.jpg'),
153
+ os.path.join(example_path, 'garment/07382_00.jpg'),
154
+ os.path.join(example_path, 'garment/07764_00.jpg'),
155
+ os.path.join(example_path, 'garment/00151_00.jpg'),
156
+ os.path.join(example_path, 'garment/12562_00.jpg'),
157
+ os.path.join(example_path, 'garment/04825_00.jpg'),
158
+ ])
159
+ with gr.Column():
160
+ result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True, scale=1)
161
+ with gr.Column():
162
+ run_button = gr.Button(value="Run")
163
+ n_samples = gr.Slider(label="Images", minimum=1, maximum=4, value=1, step=1)
164
+ n_steps = gr.Slider(label="Steps", minimum=20, maximum=40, value=20, step=1)
165
+ # scale = gr.Slider(label="Scale", minimum=1.0, maximum=12.0, value=5.0, step=0.1)
166
+ image_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=5.0, value=2.0, step=0.1)
167
+ seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=-1)
168
+
169
+ ips = [vton_img, garm_img, n_samples, n_steps, image_scale, seed]
170
+ run_button.click(fn=process_hd, inputs=ips, outputs=[result_gallery])
171
+
172
+
173
+ with gr.Row():
174
+ gr.Markdown("## Full-body")
175
+ with gr.Row():
176
+ gr.Markdown("***Support upper-body/lower-body/dresses; garment category must be paired!!!***")
177
+ with gr.Row():
178
+ with gr.Column():
179
+ vton_img_dc = gr.Image(label="Model", sources='upload', type="filepath", height=384, value=model_dc)
180
+ example = gr.Examples(
181
+ label="Examples (upper-body/lower-body)",
182
+ inputs=vton_img_dc,
183
+ examples_per_page=7,
184
+ examples=[
185
+ os.path.join(example_path, 'model/model_8.png'),
186
+ os.path.join(example_path, 'model/049447_0.jpg'),
187
+ os.path.join(example_path, 'model/049713_0.jpg'),
188
+ os.path.join(example_path, 'model/051482_0.jpg'),
189
+ os.path.join(example_path, 'model/051918_0.jpg'),
190
+ os.path.join(example_path, 'model/051962_0.jpg'),
191
+ os.path.join(example_path, 'model/049205_0.jpg'),
192
+ ])
193
+ example = gr.Examples(
194
+ label="Examples (dress)",
195
+ inputs=vton_img_dc,
196
+ examples_per_page=7,
197
+ examples=[
198
+ os.path.join(example_path, 'model/model_9.png'),
199
+ os.path.join(example_path, 'model/052767_0.jpg'),
200
+ os.path.join(example_path, 'model/052472_0.jpg'),
201
+ os.path.join(example_path, 'model/053514_0.jpg'),
202
+ os.path.join(example_path, 'model/053228_0.jpg'),
203
+ os.path.join(example_path, 'model/052964_0.jpg'),
204
+ os.path.join(example_path, 'model/053700_0.jpg'),
205
+ ])
206
+ with gr.Column():
207
+ garm_img_dc = gr.Image(label="Garment", sources='upload', type="filepath", height=384, value=garment_dc)
208
+ category_dc = gr.Dropdown(label="Garment category (important option!!!)", choices=["Upper-body", "Lower-body", "Dress"], value="Upper-body")
209
+ example = gr.Examples(
210
+ label="Examples (upper-body)",
211
+ inputs=garm_img_dc,
212
+ examples_per_page=7,
213
+ examples=[
214
+ os.path.join(example_path, 'garment/048554_1.jpg'),
215
+ os.path.join(example_path, 'garment/049920_1.jpg'),
216
+ os.path.join(example_path, 'garment/049965_1.jpg'),
217
+ os.path.join(example_path, 'garment/049949_1.jpg'),
218
+ os.path.join(example_path, 'garment/050181_1.jpg'),
219
+ os.path.join(example_path, 'garment/049805_1.jpg'),
220
+ os.path.join(example_path, 'garment/050105_1.jpg'),
221
+ ])
222
+ example = gr.Examples(
223
+ label="Examples (lower-body)",
224
+ inputs=garm_img_dc,
225
+ examples_per_page=7,
226
+ examples=[
227
+ os.path.join(example_path, 'garment/051827_1.jpg'),
228
+ os.path.join(example_path, 'garment/051946_1.jpg'),
229
+ os.path.join(example_path, 'garment/051473_1.jpg'),
230
+ os.path.join(example_path, 'garment/051515_1.jpg'),
231
+ os.path.join(example_path, 'garment/051517_1.jpg'),
232
+ os.path.join(example_path, 'garment/051988_1.jpg'),
233
+ os.path.join(example_path, 'garment/051412_1.jpg'),
234
+ ])
235
+ example = gr.Examples(
236
+ label="Examples (dress)",
237
+ inputs=garm_img_dc,
238
+ examples_per_page=7,
239
+ examples=[
240
+ os.path.join(example_path, 'garment/053290_1.jpg'),
241
+ os.path.join(example_path, 'garment/053744_1.jpg'),
242
+ os.path.join(example_path, 'garment/053742_1.jpg'),
243
+ os.path.join(example_path, 'garment/053786_1.jpg'),
244
+ os.path.join(example_path, 'garment/053790_1.jpg'),
245
+ os.path.join(example_path, 'garment/053319_1.jpg'),
246
+ os.path.join(example_path, 'garment/052234_1.jpg'),
247
+ ])
248
+ with gr.Column():
249
+ result_gallery_dc = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True, scale=1)
250
+ with gr.Column():
251
+ run_button_dc = gr.Button(value="Run")
252
+ n_samples_dc = gr.Slider(label="Images", minimum=1, maximum=4, value=1, step=1)
253
+ n_steps_dc = gr.Slider(label="Steps", minimum=20, maximum=40, value=20, step=1)
254
+ # scale_dc = gr.Slider(label="Scale", minimum=1.0, maximum=12.0, value=5.0, step=0.1)
255
+ image_scale_dc = gr.Slider(label="Guidance scale", minimum=1.0, maximum=5.0, value=2.0, step=0.1)
256
+ seed_dc = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=-1)
257
+
258
+ ips_dc = [vton_img_dc, garm_img_dc, category_dc, n_samples_dc, n_steps_dc, image_scale_dc, seed_dc]
259
+ run_button_dc.click(fn=process_dc, inputs=ips_dc, outputs=[result_gallery_dc])
260
+
261
+ block.launch()