Add quick tryon function for API

Dockerfile

@@ -27,4 +27,4 @@ ENV CUDA_VISIBLE_DEVICES=0
 EXPOSE 7860
 
 # Command to run the application
-CMD ["python", "app.py"] 
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"] 

app.py

@@ -1,10 +1,15 @@
 import spaces
+
+# For fast API
 from fastapi import FastAPI, File, UploadFile
 import pickle
 import uvicorn
 import pandas as pd
 import io
+
+# Remove gradio 
 # import gradio as gr
+
 from PIL import Image
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
@@ -30,6 +35,7 @@ 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
 
+# For use ZeroGPU 
 # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def pil_to_binary_mask(pil_image, threshold=0):
@@ -46,7 +52,7 @@ def pil_to_binary_mask(pil_image, threshold=0):
     return output_mask
 
 
-base_path = 'yisol/IDM-VTON'
+base_path = 'nami0342/GenAI_VTON_API'
 example_path = os.path.join(os.path.dirname(__file__), 'example')
 
 unet = UNet2DConditionModel.from_pretrained(
@@ -128,7 +134,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,is_automaskchecked,is_checked_crop,denoise_steps,seed):
     device = "cuda"
 
     openpose_model.preprocessor.body_estimation.model.to(device)
@@ -136,7 +142,7 @@ def start_tryon(dict,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 = dict["background"].convert("RGB")    
+    human_img_orig = dict["background"].convert("RGB")
     
     if is_checked_crop:
         width, height = human_img_orig.size
@@ -153,7 +159,7 @@ def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_ste
         human_img = human_img_orig.resize((768,1024))
 
 
-    if is_checked:
+    if is_automaskchecked:
         keypoints = openpose_model(human_img.resize((384,512)))
         model_parse, _ = parsing_model(human_img.resize((384,512)))
         mask, mask_gray = get_mask_location('hd', "upper_body", model_parse, keypoints)
@@ -288,7 +294,7 @@ async def vton_run(
     upload_human: UploadFile = File(...),
     upload_cloth: UploadFile = File(...),
     input_prompt: str = None,
-    is_checked: bool = True,
+    is_automaskchecked: bool = True,
     is_checked_crop: bool = True,
     denoise_steps: int = 30,
     seed: int = 42
@@ -296,7 +302,7 @@ async def vton_run(
     target_human = Image.open(io.BytesIO(await upload_human.read()))
     target_cloth = Image.open(io.BytesIO(await upload_cloth.read()))
 
-    results = start_tryon(target_human, target_cloth, input_prompt, is_checked, is_checked_crop, denoise_steps, seed)
+    results = start_tryon(target_human, target_cloth, input_prompt, is_automaskchecked, is_checked_crop, denoise_steps, seed)
     return results[0]
 
 

requirements.txt

@@ -20,4 +20,6 @@ av==10.0.0
 fvcore==0.1.5.post20221221
 cloudpickle==2.2.1
 omegaconf==2.3.0
-pycocotools==2.0.7
+pycocotools==2.0.7
+fastapi==0.99.1
+uvicorn