Update app.py

app.py

@@ -13,7 +13,6 @@ from transformers import (
 from diffusers import DDPMScheduler,AutoencoderKL
 from typing import List
 
-
 import torch
 import os
 from transformers import AutoTokenizer
@@ -27,7 +26,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
 
-
 def pil_to_binary_mask(pil_image, threshold=0):
     np_image = np.array(pil_image)
     grayscale_image = Image.fromarray(np_image).convert("L")
@@ -41,7 +39,6 @@ def pil_to_binary_mask(pil_image, threshold=0):
     output_mask = Image.fromarray(mask)
     return output_mask
 
-
 base_path = 'yisol/IDM-VTON'
 example_path = os.path.join(os.path.dirname(__file__), 'example')
 
@@ -85,7 +82,6 @@ vae = AutoencoderKL.from_pretrained(base_path,
                                     torch_dtype=torch.float16,
 )
 
-# "stabilityai/stable-diffusion-xl-base-1.0",
 UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
     base_path,
     subfolder="unet_encoder",
@@ -124,14 +120,14 @@ 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,category):
+def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, category):
     device = "cuda"
 
     openpose_model.preprocessor.body_estimation.model.to(device)
     pipe.to(device)
     pipe.unet_encoder.to(device)
 
-    garm_img= garm_img.convert("RGB").resize((768,1024))
+    garm_img = garm_img.convert("RGB").resize((768,1024))
     human_img_orig = dict["background"].convert("RGB")
 
     if is_checked_crop:
@@ -148,33 +144,31 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
     else:
         human_img = human_img_orig.resize((768,1024))
 
-
+    status_message = ""
     if is_checked:
-        keypoints = openpose_model(human_img.resize((384,512)))
-        model_parse, _ = parsing_model(human_img.resize((384,512)))
-        mask, mask_gray = get_mask_location('hd', category, model_parse, keypoints)
-        mask = mask.resize((768,1024))
+        try:
+            keypoints = openpose_model(human_img.resize((384,512)))
+            model_parse, _ = parsing_model(human_img.resize((384,512)))
+            mask, mask_gray = get_mask_location('hd', category, model_parse, keypoints)
+            mask = mask.resize((768,1024))
+        except NotImplementedError:
+            status_message = f"카테고리 {category}에 대한 자동 마스크 생성이 구현되지 않았습니다. 사용자가 그린 마스크를 사용합니다."
+            mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
     else:
         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)
     mask_gray = to_pil_image((mask_gray+1.0)/2.0)
 
-
     human_img_arg = _apply_exif_orientation(human_img.resize((384,512)))
     human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR")
 
-
-
     args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'))
-    # verbosity = getattr(args, "verbosity", None)
     pose_img = args.func(args,human_img_arg)
     pose_img = pose_img[:,:,::-1]
     pose_img = Image.fromarray(pose_img).resize((768,1024))
 
     with torch.no_grad():
-        # Extract the images
         with torch.cuda.amp.autocast():
             with torch.no_grad():
                 prompt = "((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), the model is wearing " + garment_des
@@ -211,8 +205,6 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
                             negative_prompt=negative_prompt,
                         )
 
-
-
                     pose_img =  tensor_transfrom(pose_img).unsqueeze(0).to(device,torch.float16)
                     garm_tensor =  tensor_transfrom(garm_img).unsqueeze(0).to(device,torch.float16)
                     generator = torch.Generator(device).manual_seed(seed) if seed is not None else None
@@ -238,10 +230,9 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
     if is_checked_crop:
         out_img = images[0].resize(crop_size)
         human_img_orig.paste(out_img, (int(left), int(top)))
-        return human_img_orig, mask_gray
+        return human_img_orig, mask_gray, status_message
     else:
-        return images[0], mask_gray
-    # return images[0], mask_gray
+        return images[0], mask_gray, status_message
 
 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]
@@ -257,10 +248,8 @@ for ex_human in human_list_path:
     ex_dict['composite'] = None
     human_ex_list.append(ex_dict)
 
-
 image_blocks = gr.Blocks(theme="Nymbo/Nymbo_Theme").queue(max_size=12)
 with image_blocks as demo:
-
     with gr.Column():
         try_button = gr.Button(value="가상 피팅 시작")
         with gr.Accordion(label="고급 설정", open=False):
@@ -301,7 +290,13 @@ with image_blocks as demo:
             masked_img = gr.Image(label="마스크 적용 이미지", elem_id="masked-img",show_share_button=False)
         with gr.Column():
             image_out = gr.Image(label="결과", elem_id="output-img",show_share_button=False)
+        
+    with gr.Column():
+        status_message = gr.Textbox(label="상태", interactive=False)
 
-    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed, category], 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, category], 
+                     outputs=[image_out, masked_img, status_message], 
+                     api_name='tryon')
 
 image_blocks.launch(auth=("gini","pick"))