Update app.py

app.py

@@ -20,9 +20,6 @@ from preprocess.openpose.run_openpose import OpenPose
 from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation
 from torchvision.transforms.functional import to_pil_image
 
-topwears = ["shirt", "t-shirt", "top", "blouse", "sweatshirt"]
-bottomwears = ["short", "shorts", "trousers", "leggings", "sweatshirt", "jeans", "skirts", "joggers", "pants", "dhoti", "lungi", "capris", "palazzos"]
-
 
 def pil_to_binary_mask(pil_image, threshold=0):
     np_image = np.array(pil_image)
@@ -38,7 +35,7 @@ def pil_to_binary_mask(pil_image, threshold=0):
     return output_mask
 
 
-base_path = './IDM-VTON'
+base_path = 'yisol/IDM-VTON'
 example_path = os.path.join(os.path.dirname(__file__), 'example')
 
 unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet", torch_dtype=torch.float16,)
@@ -64,7 +61,12 @@ vae.requires_grad_(False)
 unet.requires_grad_(False)
 text_encoder_one.requires_grad_(False)
 text_encoder_two.requires_grad_(False)
-tensor_transfrom = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5]),])
+tensor_transfrom = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+    )
 
 pipe = TryonPipeline.from_pretrained(
         base_path,
@@ -82,7 +84,7 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 @spaces.GPU
-def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_steps,seed):
+def start_tryon(dict,garm_img,garment_des,cloth_type,is_checked,is_checked_crop,denoise_steps,seed):
     device = "cuda"
     
     openpose_model.preprocessor.body_estimation.model.to(device)
@@ -110,15 +112,7 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
     if is_checked:
         keypoints = openpose_model(human_img.resize((384,512)))
         model_parse, _ = parsing_model(human_img.resize((384,512)))
-        
-        # using lambda functions to check if the description contains any -
-        contains_word = lambda s, l: any(map(lambda x: x in s, l))
-        # topwears
-        if contains_word(desc,topwears):
-            mask, mask_gray = get_mask_location('hd', "upper_body", model_parse, keypoints)
-        # bottomwears
-        if contains_word(desc,bottomwears):
-            mask, mask_gray = get_mask_location('hd', "lower_body", model_parse, keypoints)
+        mask, mask_gray = get_mask_location('hd', cloth_type, model_parse, keypoints)
         mask = mask.resize((768,1024))
     else:
         mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
@@ -165,17 +159,7 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
                     if not isinstance(negative_prompt, List):
                         negative_prompt = [negative_prompt] * 1
                     with torch.inference_mode():
-                        (
-                            prompt_embeds_c,
-                            _,
-                            _,
-                            _,
-                        ) = pipe.encode_prompt(
-                            prompt,
-                            num_images_per_prompt=1,
-                            do_classifier_free_guidance=False,
-                            negative_prompt=negative_prompt,
-                        )
+                        (prompt_embeds_c,_,_,_,) = pipe.encode_prompt(prompt,num_images_per_prompt=1,do_classifier_free_guidance=False,negative_prompt=negative_prompt,)
 
 
 
@@ -209,6 +193,19 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
         return images[0], mask_gray
     # return images[0], mask_gray
 
+
+def main_(imgs, topwear_img, topwear_des, bottomwear_img, bottomwear_des, dress_img, dress_des prompt, is_checked,is_checked_crop, denoise_steps, seed):
+    if dress_img!=None:
+        return start_tryon(imgs,dress_img,dress_des,"dresses",is_checked,is_checked_crop,denoise_steps,seed)
+    elif topwear_img!=None and bottomwear_img==None:
+        return start_tryon(imgs,topwear_img,topwear_des,"upper_body",is_checked,is_checked_crop,denoise_steps,seed)
+    elif topwear_img==None and bottomwear_img!=None:
+        return start_tryon(imgs,bottomwear_img,bottomwear_des,"lower_body",is_checked,is_checked_crop,denoise_steps,seed)
+    elif topwear_img!=None and bottomwear_img!=None:
+        half_img, half_mask = start_tryon(imgs,topwear_img,topwear_des,"upper_body",is_checked,is_checked_crop,denoise_steps,seed)
+        return start_tryon(imgs,half_img,bottomwear_des,"lower_body",is_checked,is_checked_crop,denoise_steps,seed)
+    
+        
 garm_list = os.listdir(os.path.join(example_path,"cloth"))
 garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list]
 
@@ -229,7 +226,7 @@ for ex_human in human_list_path:
 image_blocks = gr.Blocks(theme="Nymbo/Alyx_Theme").queue()
 with image_blocks as demo:
     gr.HTML("<center><h1>Virtual Try-On</h1></center>")
-    gr.HTML("<center><p>Upload an image of a person and an image of a garment ✨</p></center>")
+    gr.HTML("<center><p>Upload an image of a person and images of the clothes✨</p></center>")
     with gr.Row():
         with gr.Column():
             imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
@@ -237,39 +234,33 @@ with image_blocks as demo:
                 is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
             with gr.Row():
                 is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
-
-            example = gr.Examples(
-                inputs=imgs,
-                examples_per_page=10,
-                examples=human_ex_list
-            )
+            example = gr.Examples(inputs=imgs, examples_per_page=10, examples=human_ex_list)
 
         with gr.Column():
-            garm_img = gr.Image(label="Topwear", sources='upload', type="pil")
+            topwear_image = gr.Image(label="Topwear", sources='upload', type="pil")
             with gr.Row(elem_id="prompt-container"):
                 with gr.Row():
-                    prompt = gr.Textbox(placeholder="Description of topwear ex) Short Sleeve Black Round Neck T-shirts", show_label=False, elem_id="prompt")
-            example = gr.Examples(
-                inputs=garm_img,
-                examples_per_page=8,
-                examples=garm_list_path)
-
+                    topwear_desc = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt")
+            example = gr.Examples(inputs=topwear_img, examples_per_page=8,examples=garm_list_path)
         with gr.Column():
-            garm_img = gr.Image(label="Bottomwear", sources='upload', type="pil")
+            bottomwear_image = gr.Image(label="Bottomwear", sources='upload', type="pil")
             with gr.Row(elem_id="prompt-container"):
                 with gr.Row():
-                    prompt = gr.Textbox(placeholder="Description of bottomwear ex) Olive Cargo Pants", show_label=False, elem_id="prompt")
-            example = gr.Examples(
-                inputs=garm_img,
-                examples_per_page=8,
-                examples=garm_list_path)
-            
+                    bottomwear_desc = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt")
+            example = gr.Examples(inputs=bottomwear_img, examples_per_page=8,examples=garm_list_path)
         with gr.Column():
-            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
-            masked_img = gr.Image(label="Masked image output", elem_id="masked-img",show_share_button=False)
+            dress_image = gr.Image(label="Dress", sources='upload', type="pil")
+            with gr.Row(elem_id="prompt-container"):
+                with gr.Row():
+                    dress_desc = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt")
+            example = gr.Examples(inputs=dress_img, examples_per_page=8,examples=garm_list_path)
+
         with gr.Column():
             # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
             image_out = gr.Image(label="Output", elem_id="output-img",show_share_button=False)
+            with gr.Accordion("Debug Info", open=False):
+                masked_img = gr.Image(label="Masked image output", elem_id="masked-img",show_share_button=False)
+
 
 
 
@@ -281,8 +272,10 @@ with image_blocks as demo:
                 seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)
 
 
+    try_button.click(fn=main_, inputs=[imgs,topwear_image,topwear_desc,bottomwear_image,bottomwear_desc,dress_image,dress_desc,prompt,is_checked,is_checked_crop,denoise_steps,seed],
+                     outputs=[image_out, masked_img], api_name='tryon')
+
+            
 
-    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
 
-    
 image_blocks.launch()