138 lines
4.5 KiB
Python
138 lines
4.5 KiB
Python
"""
|
|
RunPod Serverless worker: SDXL inpainting.
|
|
|
|
Powers the app's Magic Eraser, Generative Fill AND Outpaint — all three go
|
|
through this one endpoint (RUNPOD_SD_INPAINT_URL). No app change needed; this
|
|
matches the exact contract the app already sends (rpInpaint):
|
|
|
|
{"input": {
|
|
"image": "<base64>", # required (raw base64 or data: URI)
|
|
"mask": "<base64>", # required — WHITE = regenerate, BLACK = keep
|
|
"prompt": "<text>", # optional ("" = Magic Eraser: clean fill)
|
|
"strength": 0.8, # optional
|
|
"guidance_scale": 7, # optional
|
|
"num_inference_steps": 35, # optional
|
|
"negative_prompt": "...", # optional
|
|
"seed": 123 # optional
|
|
}}
|
|
Returns: {"image": "<base64 PNG>"} (or {"error": "..."}).
|
|
|
|
CUDA 12.1 base + torch cu121 (same as the detect/upscale workers) so it runs on
|
|
RunPod's RTX 3090/4090 (driver 550) hosts.
|
|
"""
|
|
|
|
import base64
|
|
import io
|
|
import os
|
|
|
|
_V = "/runpod-volume"
|
|
if os.path.isdir(_V):
|
|
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
|
|
|
import torch # noqa: E402
|
|
import runpod # noqa: E402
|
|
from PIL import Image # noqa: E402
|
|
from diffusers import AutoPipelineForInpainting # noqa: E402
|
|
|
|
MODEL = os.environ.get("SD_INPAINT_MODEL", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1")
|
|
MAX_SIZE = int(os.environ.get("SD_INPAINT_MAX_SIZE", "1024"))
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
_pipe = None
|
|
|
|
|
|
def _get_pipe():
|
|
global _pipe
|
|
if _pipe is None:
|
|
try:
|
|
pipe = AutoPipelineForInpainting.from_pretrained(
|
|
MODEL, torch_dtype=torch.float16, variant="fp16"
|
|
)
|
|
except Exception:
|
|
# Repo may not ship an fp16 variant — fall back to default weights.
|
|
pipe = AutoPipelineForInpainting.from_pretrained(MODEL, torch_dtype=torch.float16)
|
|
total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if DEVICE == "cuda" else 0
|
|
if total_gb >= 20:
|
|
pipe = pipe.to("cuda")
|
|
else:
|
|
pipe.enable_model_cpu_offload()
|
|
pipe.set_progress_bar_config(disable=True)
|
|
_pipe = pipe
|
|
return _pipe
|
|
|
|
|
|
def _b64_to_image(data, mode="RGB"):
|
|
if "," in data:
|
|
data = data.split(",", 1)[1]
|
|
return Image.open(io.BytesIO(base64.b64decode(data))).convert(mode)
|
|
|
|
|
|
def _image_to_b64(img):
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
|
|
|
|
def _fit(img):
|
|
# SDXL wants multiples of 8; cap the long side at MAX_SIZE.
|
|
w, h = img.size
|
|
scale = min(1.0, MAX_SIZE / max(w, h))
|
|
nw = max(8, (int(w * scale) // 8) * 8)
|
|
nh = max(8, (int(h * scale) // 8) * 8)
|
|
return img.resize((nw, nh), Image.LANCZOS)
|
|
|
|
|
|
def handler(event):
|
|
inp = (event or {}).get("input") or {}
|
|
image_b64 = inp.get("image")
|
|
mask_b64 = inp.get("mask")
|
|
if not image_b64:
|
|
return {"error": "Missing 'image' (base64)."}
|
|
if not mask_b64:
|
|
return {"error": "Missing 'mask' (base64)."}
|
|
|
|
prompt = (inp.get("prompt") or "").strip()
|
|
negative = (inp.get("negative_prompt") or "").strip() or None
|
|
|
|
try:
|
|
image = _fit(_b64_to_image(image_b64, "RGB"))
|
|
mask = _b64_to_image(mask_b64, "L").resize(image.size, Image.NEAREST)
|
|
|
|
steps = int(inp.get("num_inference_steps") or 35)
|
|
guidance = float(inp.get("guidance_scale") or 7.0)
|
|
|
|
# Strength: empty prompt = Magic Eraser -> fully replace (1.0). With a
|
|
# prompt (Generative Fill / Outpaint) respect the caller but keep a floor
|
|
# so the masked region actually changes.
|
|
raw = inp.get("strength")
|
|
if not prompt:
|
|
strength = 1.0
|
|
else:
|
|
strength = max(0.7, min(1.0, float(raw))) if raw is not None else 0.9
|
|
|
|
seed = inp.get("seed")
|
|
generator = (
|
|
torch.Generator(device=DEVICE).manual_seed(int(seed)) if seed is not None else None
|
|
)
|
|
|
|
eff_prompt = prompt or "clean, seamless, photorealistic background, natural continuation"
|
|
|
|
result = _get_pipe()(
|
|
prompt=eff_prompt,
|
|
negative_prompt=negative,
|
|
image=image,
|
|
mask_image=mask,
|
|
height=image.height,
|
|
width=image.width,
|
|
num_inference_steps=steps,
|
|
guidance_scale=guidance,
|
|
strength=strength,
|
|
generator=generator,
|
|
).images[0]
|
|
|
|
return {"image": _image_to_b64(result)}
|
|
except Exception as exc:
|
|
return {"error": f"{type(exc).__name__}: {exc}"}
|
|
|
|
|
|
runpod.serverless.start({"handler": handler})
|