Spaces:
Running
on
Zero
Running
on
Zero
omyteq agency
commited on
Commit
•
b3767b7
1
Parent(s):
5d6a57f
Update run/run_ootd.py
Browse files- run/run_ootd.py +9 -17
run/run_ootd.py
CHANGED
@@ -11,9 +11,6 @@ from preprocess.humanparsing.run_parsing import Parsing
|
|
11 |
from ootd.inference_ootd_hd import OOTDiffusionHD
|
12 |
from ootd.inference_ootd_dc import OOTDiffusionDC
|
13 |
|
14 |
-
# Create the outputs directory if it doesn't exist
|
15 |
-
output_dir = Path("./outputs")
|
16 |
-
output_dir.mkdir(exist_ok=True)
|
17 |
|
18 |
import argparse
|
19 |
parser = argparse.ArgumentParser(description='run ootd')
|
@@ -28,14 +25,16 @@ parser.add_argument('--sample', type=int, default=4, required=False)
|
|
28 |
parser.add_argument('--seed', type=int, default=-1, required=False)
|
29 |
args = parser.parse_args()
|
30 |
|
|
|
31 |
openpose_model = OpenPose(args.gpu_id)
|
32 |
parsing_model = Parsing(args.gpu_id)
|
33 |
|
|
|
34 |
category_dict = ['upperbody', 'lowerbody', 'dress']
|
35 |
category_dict_utils = ['upper_body', 'lower_body', 'dresses']
|
36 |
|
37 |
-
model_type = args.model_type
|
38 |
-
category = args.category
|
39 |
cloth_path = args.cloth_path
|
40 |
model_path = args.model_path
|
41 |
|
@@ -49,12 +48,13 @@ if model_type == "hd":
|
|
49 |
elif model_type == "dc":
|
50 |
model = OOTDiffusionDC(args.gpu_id)
|
51 |
else:
|
52 |
-
raise ValueError("model_type must be 'hd' or 'dc'!")
|
|
|
53 |
|
54 |
if __name__ == '__main__':
|
55 |
|
56 |
if model_type == 'hd' and category != 0:
|
57 |
-
raise ValueError("model_type 'hd' requires category == 0 (upperbody)!")
|
58 |
|
59 |
cloth_img = Image.open(cloth_path).resize((768, 1024))
|
60 |
model_img = Image.open(model_path).resize((768, 1024))
|
@@ -66,8 +66,7 @@ if __name__ == '__main__':
|
|
66 |
mask_gray = mask_gray.resize((768, 1024), Image.NEAREST)
|
67 |
|
68 |
masked_vton_img = Image.composite(mask_gray, model_img, mask)
|
69 |
-
|
70 |
-
masked_vton_img.save(output_dir / 'mask.jpg')
|
71 |
|
72 |
images = model(
|
73 |
model_type=model_type,
|
@@ -84,12 +83,5 @@ if __name__ == '__main__':
|
|
84 |
|
85 |
image_idx = 0
|
86 |
for image in images:
|
87 |
-
|
88 |
-
image.save(output_path) # Save images explicitly to outputs
|
89 |
image_idx += 1
|
90 |
-
print(f"Image saved to {output_path}")
|
91 |
-
|
92 |
-
# Construct URLs for images in the outputs directory
|
93 |
-
base_url = "https://omyteq-ootdiffusion.hf.space/file=" # Replace with your actual space URL
|
94 |
-
output_urls = [f"{base_url}{output_path.name}" for output_path in output_dir.glob('*.png')]
|
95 |
-
print(f"Image URLs: {output_urls}")
|
|
|
11 |
from ootd.inference_ootd_hd import OOTDiffusionHD
|
12 |
from ootd.inference_ootd_dc import OOTDiffusionDC
|
13 |
|
|
|
|
|
|
|
14 |
|
15 |
import argparse
|
16 |
parser = argparse.ArgumentParser(description='run ootd')
|
|
|
25 |
parser.add_argument('--seed', type=int, default=-1, required=False)
|
26 |
args = parser.parse_args()
|
27 |
|
28 |
+
|
29 |
openpose_model = OpenPose(args.gpu_id)
|
30 |
parsing_model = Parsing(args.gpu_id)
|
31 |
|
32 |
+
|
33 |
category_dict = ['upperbody', 'lowerbody', 'dress']
|
34 |
category_dict_utils = ['upper_body', 'lower_body', 'dresses']
|
35 |
|
36 |
+
model_type = args.model_type # "hd" or "dc"
|
37 |
+
category = args.category # 0:upperbody; 1:lowerbody; 2:dress
|
38 |
cloth_path = args.cloth_path
|
39 |
model_path = args.model_path
|
40 |
|
|
|
48 |
elif model_type == "dc":
|
49 |
model = OOTDiffusionDC(args.gpu_id)
|
50 |
else:
|
51 |
+
raise ValueError("model_type must be \'hd\' or \'dc\'!")
|
52 |
+
|
53 |
|
54 |
if __name__ == '__main__':
|
55 |
|
56 |
if model_type == 'hd' and category != 0:
|
57 |
+
raise ValueError("model_type \'hd\' requires category == 0 (upperbody)!")
|
58 |
|
59 |
cloth_img = Image.open(cloth_path).resize((768, 1024))
|
60 |
model_img = Image.open(model_path).resize((768, 1024))
|
|
|
66 |
mask_gray = mask_gray.resize((768, 1024), Image.NEAREST)
|
67 |
|
68 |
masked_vton_img = Image.composite(mask_gray, model_img, mask)
|
69 |
+
masked_vton_img.save('./images_output/mask.jpg')
|
|
|
70 |
|
71 |
images = model(
|
72 |
model_type=model_type,
|
|
|
83 |
|
84 |
image_idx = 0
|
85 |
for image in images:
|
86 |
+
image.save('./images_output/out_' + model_type + '_' + str(image_idx) + '.png')
|
|
|
87 |
image_idx += 1
|
|
|
|
|
|
|
|
|
|
|
|