Upload folder using huggingface_hub

.ipynb_checkpoints/handler-checkpoint.py

@@ -1,5 +1,5 @@
 import runpod
-from helpers import prepare_pipeline, get_result, b64_to_pil
+from helpers import get_result, b64_to_pil, pil_to_b64
 import base64
 from PIL import Image
 
@@ -10,16 +10,21 @@ def handler(job):
     garm_img_b64 = job['input']['garm_img_b64']
     garm_img = b64_to_pil(garm_img_b64)
 
+    body_part = job['input'].get('body_part')
+    
     denoise_steps = job['input'].get('denoise_steps') if job['input'].get('denoise_steps') else 30
 
     seed = job['input'].get('seed') if job['input'].get('seed') else 42
-
+    
     is_checked_crop = job['input'].get('is_checked_crop') if job['input'].get('is_checked_crop') else False
 
     garment_des = job['input'].get('garment_des') if job['input'].get('garment_des') else ""
     
-    result = get_result(PIPE, human_img, garm_img, denoise_steps, seed, is_checked_crop, garment_des)
-    return pil_to_b64(result)
+    output_image, mask_image = get_result(human_img, garm_img, body_part, denoise_steps, seed, is_checked_crop, garment_des)
+
+    output_image.save("output_image.jpg")
+    mask_image.save("mask_image.jpg")
+    
+    return {"output_image": pil_to_b64(output_image), "mask_image": pil_to_b64(mask_image)}
 
-PIPE = prepare_pipeline()
 runpod.serverless.start({"handler": handler})

.ipynb_checkpoints/helpers-checkpoint.py

@@ -26,7 +26,7 @@ from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_ori
 from torchvision.transforms.functional import to_pil_image
 
 
-def b64_to_pil():
+def b64_to_pil(base64_string):
     # Decode the base64 string
     image_data = base64.b64decode(base64_string)
 
@@ -34,6 +34,12 @@ def b64_to_pil():
     image = Image.open(BytesIO(image_data))
     return image
 
+def pil_to_b64(pil_img):
+    buffered = BytesIO()
+    pil_img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return img_str
+    
 def prepare_pipeline():
     pass
 base_path = 'yisol/IDM-VTON'

handler.py

@@ -1,5 +1,5 @@
 import runpod
-from helpers import prepare_pipeline, get_result, b64_to_pil
+from helpers import get_result, b64_to_pil, pil_to_b64
 import base64
 from PIL import Image
 
@@ -10,16 +10,21 @@ def handler(job):
     garm_img_b64 = job['input']['garm_img_b64']
     garm_img = b64_to_pil(garm_img_b64)
 
+    body_part = job['input'].get('body_part')
+    
     denoise_steps = job['input'].get('denoise_steps') if job['input'].get('denoise_steps') else 30
 
     seed = job['input'].get('seed') if job['input'].get('seed') else 42
-
+    
     is_checked_crop = job['input'].get('is_checked_crop') if job['input'].get('is_checked_crop') else False
 
     garment_des = job['input'].get('garment_des') if job['input'].get('garment_des') else ""
     
-    result = get_result(PIPE, human_img, garm_img, denoise_steps, seed, is_checked_crop, garment_des)
-    return pil_to_b64(result)
+    output_image, mask_image = get_result(human_img, garm_img, body_part, denoise_steps, seed, is_checked_crop, garment_des)
+
+    output_image.save("output_image.jpg")
+    mask_image.save("mask_image.jpg")
+    
+    return {"output_image": pil_to_b64(output_image), "mask_image": pil_to_b64(mask_image)}
 
-PIPE = prepare_pipeline()
 runpod.serverless.start({"handler": handler})

helpers.py

@@ -26,7 +26,7 @@ from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_ori
 from torchvision.transforms.functional import to_pil_image
 
 
-def b64_to_pil():
+def b64_to_pil(base64_string):
     # Decode the base64 string
     image_data = base64.b64decode(base64_string)
 
@@ -34,6 +34,12 @@ def b64_to_pil():
     image = Image.open(BytesIO(image_data))
     return image
 
+def pil_to_b64(pil_img):
+    buffered = BytesIO()
+    pil_img.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return img_str
+    
 def prepare_pipeline():
     pass
 base_path = 'yisol/IDM-VTON'