import runpod
from helpers import get_result, b64_to_pil, pil_to_b64
import base64
from PIL import Image

def handler(job):
    human_img_b64 = job['input']['human_img_b64']
    human_img = b64_to_pil(human_img_b64)

    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 ""
    
    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)}

runpod.serverless.start({"handler": handler})