| import base64 |
| import io |
| import os |
|
|
| import runpod |
| from PIL import Image |
|
|
| from pipeline import load_pipeline |
|
|
|
|
| PIPELINE = load_pipeline(os.getenv("LORA_PATH", ".")) |
|
|
|
|
| def _as_bool(value, default: bool = True) -> bool: |
| if value is None: |
| return default |
| if isinstance(value, bool): |
| return value |
| if isinstance(value, str): |
| return value.lower() in {"1", "true", "yes", "on"} |
| return bool(value) |
|
|
|
|
| def _decode_image(data: str) -> Image.Image: |
| if data.startswith("data:image"): |
| data = data.split(",", 1)[1] |
| return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB") |
|
|
|
|
| def _encode_png(image: Image.Image) -> str: |
| buffer = io.BytesIO() |
| image.save(buffer, format="PNG") |
| return base64.b64encode(buffer.getvalue()).decode("utf-8") |
|
|
|
|
| def handler(event): |
| payload = event.get("input", {}) |
| image_b64 = payload.get("image") |
| if not image_b64: |
| return {"error": "Missing input.image base64 PNG/JPEG data."} |
|
|
| image = _decode_image(image_b64) |
| seed = payload.get("seed") |
| if seed is not None: |
| seed = int(seed) |
|
|
| output = PIPELINE( |
| image=image, |
| num_inference_steps=int(payload.get("num_inference_steps", 50)), |
| guidance_scale=float(payload.get("guidance_scale", 7.5)), |
| controlnet_conditioning_scale=float(payload.get("controlnet_conditioning_scale", 0.8)), |
| strength=float(payload.get("strength", 0.75)), |
| quantize=_as_bool(payload.get("quantize"), True), |
| n_colors=int(payload.get("n_colors", 32)), |
| seed=seed, |
| ) |
|
|
| return { |
| "image": _encode_png(output["image"]), |
| "rembg_ok": output["rembg_ok"], |
| } |
|
|
|
|
| runpod.serverless.start({"handler": handler}) |
|
|
