Update run/run_ootd.py

run/run_ootd.py

@@ -11,9 +11,6 @@ from preprocess.humanparsing.run_parsing import Parsing
 from ootd.inference_ootd_hd import OOTDiffusionHD
 from ootd.inference_ootd_dc import OOTDiffusionDC
 
-# Create the outputs directory if it doesn't exist
-output_dir = Path("./outputs")
-output_dir.mkdir(exist_ok=True)
 
 import argparse
 parser = argparse.ArgumentParser(description='run ootd')
@@ -28,14 +25,16 @@ parser.add_argument('--sample', type=int, default=4, required=False)
 parser.add_argument('--seed', type=int, default=-1, required=False)
 args = parser.parse_args()
 
+
 openpose_model = OpenPose(args.gpu_id)
 parsing_model = Parsing(args.gpu_id)
 
+
 category_dict = ['upperbody', 'lowerbody', 'dress']
 category_dict_utils = ['upper_body', 'lower_body', 'dresses']
 
-model_type = args.model_type  # "hd" or "dc"
-category = args.category  # 0:upperbody; 1:lowerbody; 2:dress
+model_type = args.model_type # "hd" or "dc"
+category = args.category # 0:upperbody; 1:lowerbody; 2:dress
 cloth_path = args.cloth_path
 model_path = args.model_path
 
@@ -49,12 +48,13 @@ if model_type == "hd":
 elif model_type == "dc":
     model = OOTDiffusionDC(args.gpu_id)
 else:
-    raise ValueError("model_type must be 'hd' or 'dc'!")
+    raise ValueError("model_type must be \'hd\' or \'dc\'!")
+
 
 if __name__ == '__main__':
 
     if model_type == 'hd' and category != 0:
-        raise ValueError("model_type 'hd' requires category == 0 (upperbody)!")
+        raise ValueError("model_type \'hd\' requires category == 0 (upperbody)!")
 
     cloth_img = Image.open(cloth_path).resize((768, 1024))
     model_img = Image.open(model_path).resize((768, 1024))
@@ -66,8 +66,7 @@ if __name__ == '__main__':
     mask_gray = mask_gray.resize((768, 1024), Image.NEAREST)
     
     masked_vton_img = Image.composite(mask_gray, model_img, mask)
-    # Save the mask image to outputs directory
-    masked_vton_img.save(output_dir / 'mask.jpg')
+    masked_vton_img.save('./images_output/mask.jpg')
 
     images = model(
         model_type=model_type,
@@ -84,12 +83,5 @@ if __name__ == '__main__':
 
     image_idx = 0
     for image in images:
-        output_path = output_dir / f'out_{model_type}_{image_idx}.png'
-        image.save(output_path)  # Save images explicitly to outputs
+        image.save('./images_output/out_' + model_type + '_' + str(image_idx) + '.png')
         image_idx += 1
-        print(f"Image saved to {output_path}")
-
-    # Construct URLs for images in the outputs directory
-    base_url = "https://omyteq-ootdiffusion.hf.space/file="  # Replace with your actual space URL
-    output_urls = [f"{base_url}{output_path.name}" for output_path in output_dir.glob('*.png')]
-    print(f"Image URLs: {output_urls}")