File size: 1,725 Bytes
e397040 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | 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})
|