Update app.py

app.py

@@ -28,6 +28,9 @@ from torchvision import transforms
 from preprocess.humanparsing.run_parsing import Parsing
 from preprocess.openpose.run_openpose import OpenPose
 from torchvision.transforms.functional import to_pil_image
+import apply_net
+from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation
+
 
 def pil_to_binary_mask(pil_image, threshold=0):
     np_image = np.array(pil_image)
@@ -124,7 +127,7 @@ pipe = TryonPipeline.from_pretrained(
 pipe.unet_encoder = UNet_Encoder
 
 @spaces.GPU
-def start_tryon(person_img,cloth_img, garment_des, denoise_steps=10, seed=42):
+def start_tryon(boy,girl,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)
@@ -133,7 +136,7 @@ def start_tryon(person_img,cloth_img, garment_des, denoise_steps=10, seed=42):
 
     # Resize and prepare images
     garm_img = cloth_img.convert("RGB").resize((768, 1024))
-    human_img = person_img.convert("RGB")
+    human_img = person_img.convert("RGB").resize((768,1024))
 
     is_checked=True;
     if is_checked:
@@ -142,19 +145,25 @@ def start_tryon(person_img,cloth_img, garment_des, denoise_steps=10, seed=42):
         mask, mask_gray= get_mask_location('hd', "upper_body", model_parse, keypoints)
         mask = mask.resize((768,1024))
 
-    #mask = pil_to_binary_mask(mask.convert("RGB").resize((768, 1024)))
-    pose_img=Image.open("00006_00.jpg")
+    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))
 
-    # Prepare pose image (already uploaded)
-    pose_img = pose_img.resize((768, 1024))
+    if boy:
+        prompt = "A boy is wearing"+garment_des
+    if girl:
+        prompt= "A girl is wearing"+garment_des
 
-    
-    
     # Embedding generation for prompts
     with torch.no_grad():
         with torch.cuda.amp.autocast():
             # Generate text embeddings for garment description
-            prompt = f"model is wearing {garment_des}"
+            prompt = prompt
             negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
             with torch.inference_mode():
              (
@@ -168,7 +177,7 @@ def start_tryon(person_img,cloth_img, garment_des, denoise_steps=10, seed=42):
                 do_classifier_free_guidance=True,
                 negative_prompt=negative_prompt,
              )
-            prompt = "a photo of " + garment_des
+            prompt = "A photo of " + garment_des
             negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
             if not isinstance(prompt, List):
                         prompt = [prompt] * 1
@@ -203,6 +212,7 @@ def start_tryon(person_img,cloth_img, garment_des, denoise_steps=10, seed=42):
                 num_inference_steps=denoise_steps,
                 generator=generator,
                 strength=1.0,
+                pose_img=pose_img_tensor.to(device, torch.float16),
                 text_embeds_cloth=prompt_embeds_cloth.to(device, torch.float16),
                 cloth=garm_tensor.to(device, torch.float16),
                 pose_img = None,
@@ -225,6 +235,9 @@ with image_blocks as demo:
     with gr.Row():
         with gr.Column():
             person_img = gr.Image(label='Person Image', sources='upload', type="pil")
+            boy = gr.Checkbox(label="Yes", info="Boy",value=True)
+            girl = gr.Checkbox(label="Yes", info="Girl",value=False)
+
         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")
@@ -234,6 +247,6 @@ with image_blocks as demo:
             image_out = gr.Image(label="Output Image", elem_id="output-img", show_share_button=False)
 
     try_button = gr.Button(value="Try-on")
-    try_button.click(fn=start_tryon, inputs=[person_img, cloth_img, garment_des], outputs=[image_out], api_name='tryon')
+    try_button.click(fn=start_tryon, inputs=[boy,girl,person_img, cloth_img, garment_des], outputs=[image_out], api_name='tryon')
 
 image_blocks.launch()