File size: 1,416 Bytes
8204928 d66b293 06a49e4 8204928 d66b293 06a49e4 8204928 d66b293 06a49e4 d66b293 32b6c5f d66b293 06a49e4 8204928 06a49e4 |
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 |
import base64
import io
from PIL import Image
import torch
from diffusers import StableDiffusionImg2ImgPipeline
# Global pipeline instance
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = None
class EndpointHandler:
def __init__(self, model_dir: str):
# model_dir is ignored; HF clones your repo here
pass
def init(self):
global pipe
if pipe is None:
# Load your InstantID img2img model
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"karthikAI/InstantID-i2i",
revision="main",
torch_dtype=torch.float16
).to(torch_device)
def inference(self, model_inputs: dict) -> dict:
# 1) decode base64 image
b64 = model_inputs.get("inputs")
if b64 is None:
return {"error": "No 'inputs' key with base64 image provided."}
img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
# 2) extract prompt
prompt = model_inputs.get("parameters", {}).get("prompt", "")
# 3) minimal call: prompt + image only
out = pipe(prompt=prompt, image=img)
result_img = out.images[0]
# 4) encode output
buf = io.BytesIO()
result_img.save(buf, format="PNG")
b64_out = base64.b64encode(buf.getvalue()).decode()
return {"generated_image_base64": b64_out} |