Update app.py

app.py

@@ -1,14 +1,14 @@
 import sys
+import torch
 import os
-
+import spaces
+import gradio as gr
+from PIL import Image
 sys.path.append('./')
 os.system("pip install gradio accelerate==0.25.0 torchmetrics==1.2.1 tqdm==4.66.1 fastapi==0.111.0 transformers==4.36.2 diffusers==0.25 einops==0.7.0 bitsandbytes scipy==1.11.1 opencv-python gradio==4.24.0 fvcore cloudpickle omegaconf pycocotools basicsr av onnxruntime==1.16.2 peft==0.11.1 huggingface_hub==0.24.7 --no-deps")
-import spaces
 from fastapi import FastAPI
 app = FastAPI()
 
-from PIL import Image
-import gradio as gr
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
 from src.unet_hacked_tryon import UNet2DConditionModel
@@ -20,15 +20,13 @@ from transformers import (
 )
 from diffusers import DDPMScheduler,AutoencoderKL
 from typing import List
-
-import torch
-import os
 from transformers import AutoTokenizer
 import numpy as np
+from utils_mask import get_mask_location
 from torchvision import transforms
-
-
-device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+from preprocess.humanparsing.run_parsing import Parsing
+from preprocess.openpose.run_openpose import OpenPose
+from torchvision.transforms.functional import to_pil_image
 
 def pil_to_binary_mask(pil_image, threshold=0):
     np_image = np.array(pil_image)
@@ -93,7 +91,8 @@ UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
     torch_dtype=torch.float16,
 )
 
-
+parsing_model = Parsing(0)
+openpose_model = OpenPose(0)
 
 UNet_Encoder.requires_grad_(False)
 image_encoder.requires_grad_(False)
@@ -124,15 +123,23 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 @spaces.GPU
-def start_tryon(person_img, mask_img, cloth_img, garment_des, denoise_steps=10, seed=42):
+def start_tryon(person_img,cloth_img, garment_des, denoise_steps=10, seed=42):
     # Assuming device is set up (e.g., "cuda" or "cpu")
+    device="cuda"
+    openpose_model.preprocessor.body_estimation.model.to(device)
     pipe.to(device)
     pipe.unet_encoder.to(device)
 
     # Resize and prepare images
     garm_img = cloth_img.convert("RGB").resize((768, 1024))
     human_img = person_img.convert("RGB").resize((768, 1024))
-    mask = pil_to_binary_mask(mask_img.convert("RGB").resize((768, 1024)))
+    is_checked=True;
+    if is_checked:
+        keypoints = openpose_model(human_img.resize((384,512)))
+        model_parse, _ = parsing_model(human_img.resize((384,512)))
+        mask= get_mask_location('hd', "upper_body", model_parse, keypoints)
+
+    mask = pil_to_binary_mask(mask.convert("RGB").resize((768, 1024)))
     pose_img=Image.open("00006_00.jpg")
 
     # Prepare pose image (already uploaded)
@@ -215,8 +222,6 @@ with image_blocks as demo:
     with gr.Row():
         with gr.Column():
             person_img = gr.Image(label='Person Image', sources='upload', type="pil")
-            mask_img = gr.Image(label='Mask Image', sources='upload', type="pil")
-
         with gr.Column():
             cloth_img = gr.Image(label='Garment Image', sources='upload', type="pil")
             garment_des = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", label="Garment Description")