Upload runpod_handler.py with huggingface_hub

runpod_handler.py

@@ -0,0 +1,64 @@
+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})