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