handling both shein images and viton standard images

app.py

@@ -1,4 +1,5 @@
 import logging
+import math
 import gradio as gr
 from PIL import Image
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
@@ -135,16 +136,25 @@ CATEGORY = "upper_body" # "lower_body"
 
 @spaces.GPU
 def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked_crop,denoise_steps,seed):
-    device = "cuda"
-    # device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+    #device = "cuda"
+    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
     
     openpose_model.preprocessor.body_estimation.model.to(device)
     pipe.to(device)
     pipe.unet_encoder.to(device)
+    
+    #human_img_orig = dict["background"].convert("RGB")   # ImageEditor
+    human_img_orig = dict.convert("RGB")     # Image
+
+    # Derive HEIGHT & WIDTH such that width is not more than 1000. This will cater to both Shein images (4160x6240) of 3:4 AR and model standard images ( 768x1024 ) of 2:3 AR
+    WIDTH, HEIGHT = human_img_orig.size
+    division_factor = math.ceil(WIDTH/1000)
+    WIDTH = int(WIDTH/division_factor)
+    HEIGHT = int(HEIGHT/division_factor)
+    POSE_WIDTH = int(WIDTH/2)
+    POSE_HEIGHT = int(HEIGHT/2)
 
     garm_img= garm_img.convert("RGB").resize((WIDTH,HEIGHT))
-    # human_img_orig = dict["background"].convert("RGB")    
-    human_img_orig = dict.convert("RGB")     
     if is_checked_crop:
         width, height = human_img_orig.size
         target_width = int(min(width, height * (3 / 4)))
@@ -272,8 +282,8 @@ human_list = os.listdir(os.path.join(example_path,"human"))
 human_list_path = [os.path.join(example_path,"human",human) for human in human_list]
 
 human_ex_list = []
-human_ex_list = human_list_path
-""" if using ImageEditor instead of Image while taking input, use this
+human_ex_list = human_list_path # Image
+""" if using ImageEditor instead of Image while taking input, use this - ImageEditor
 for ex_human in human_list_path:
     ex_dict= {}
     ex_dict['background'] = ex_human
@@ -293,7 +303,7 @@ with image_blocks as demo:
         with gr.Column():
             # changing from ImageEditor to Image to allow easy passing of data through API
             # instead of passing {"dictionary": <>} ( which is failing ), we can directly pass the image
-            # imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
+            #imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
             imgs = gr.Image(sources='upload', type='pil',label='Human. Mask with pen or use auto-masking')
             with gr.Row():
                 is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)