Update app.py

app.py

@@ -122,7 +122,7 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 @spaces.GPU
-def start_tryon(imgs,garm_img,garment_des,is_checked,is_checked_crop,denoise_steps,seed, category):
+def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_steps,seed, category):
     device = "cuda"
     category = int(category) 
     if category==0:
@@ -137,7 +137,7 @@ def start_tryon(imgs,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
     pipe.unet_encoder.to(device)
 
     garm_img= garm_img.convert("RGB").resize((768,1024))
-    human_img_orig = imgs.convert("RGB")    
+    human_img_orig = dict["background"].convert("RGB")    
     
     if is_checked_crop:
         width, height = human_img_orig.size
@@ -165,7 +165,7 @@ def start_tryon(imgs,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
         mask, mask_gray = get_mask_location('hd', category, model_parse, keypoints)
         mask = mask.resize((768,1024))
     else:
-        mask = pil_to_binary_mask(imgs.convert("RGB").resize((768, 1024)))
+        mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
         # mask = transforms.ToTensor()(mask)
         # mask = mask.unsqueeze(0)
     mask_gray = (1-transforms.ToTensor()(mask)) * tensor_transfrom(human_img)
@@ -276,11 +276,11 @@ for ex_human in human_list_path:
 
 image_blocks = gr.Blocks().queue()
 with image_blocks as demo:
-    gr.Markdown("## VirtualFit")
-    gr.Markdown("VirtualFIT Demo")
+    gr.Markdown("## DressFit")
+    gr.Markdown("DressFit Demo")
     with gr.Row():
         with gr.Column():
-            imgs = gr.Image(label="Person Image", sources='upload', type="pil")
+            imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
             with gr.Row():
                 is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
             with gr.Row():
@@ -293,6 +293,7 @@ with image_blocks as demo:
                 examples_per_page=10,
                 examples=human_ex_list
             )
+        
 
         with gr.Column():
             garm_img = gr.Image(label="Garment", sources='upload', type="pil")