File size: 3,011 Bytes
5af9786 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | import torch
import base64
import io
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline
import os
class EndpointHandler():
def __init__(self, path=""):
# 'path' is the folder where HF automatically loaded your repo files
print("Loading ARX Pipeline...")
# 1. Point directly to the model you uploaded to the repo
model_path = os.path.join(path, "biglust.safetensors")
# 2. Load the pipeline
self.pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
model_path,
torch_dtype=torch.float16,
use_safetensors=True,
safety_checker=None
)
# 3. Load IP-Adapter (HF will download this from the public hub automatically)
self.pipe.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter_sdxl.bin"
)
self.pipe.to("cuda")
print("ARX Pipeline Ready.")
def decode_base64_image(self, image_string):
if "," in image_string:
image_string = image_string.split(",")[1]
image_bytes = base64.b64decode(image_string)
return Image.open(io.BytesIO(image_bytes)).convert("RGB")
def encode_image_base64(self, image):
buffered = io.BytesIO()
image.save(buffered, format="PNG")
return base64.b64encode(buffered.getvalue()).decode('utf-8')
def __call__(self, data):
"""
data param format: {"inputs": { "prompt": "...", "init_image": "..." }}
"""
# HF wraps payloads in an "inputs" key automatically
inputs = data.pop("inputs", data)
prompt = inputs.get("prompt", "masterpiece, best quality")
negative_prompt = inputs.get("negative_prompt", "lowres, bad anatomy, worst quality, ugly")
strength = float(inputs.get("strength", 0.65))
guidance_scale = float(inputs.get("guidance_scale", 7.0))
num_inference_steps = int(inputs.get("steps", 25))
ip_adapter_scale = float(inputs.get("ip_adapter_scale", 0.50))
init_image_b64 = inputs.get("init_image")
ip_adapter_image_b64 = inputs.get("ip_adapter_image")
if not init_image_b64 or not ip_adapter_image_b64:
return {"error": "Both init_image and ip_adapter_image must be provided."}
init_image = self.decode_base64_image(init_image_b64).resize((1024, 1024))
ip_image = self.decode_base64_image(ip_adapter_image_b64).resize((1024, 1024))
self.pipe.set_ip_adapter_scale(ip_adapter_scale)
# Generate!
result = self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
image=init_image,
ip_adapter_image=ip_image,
strength=strength,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps
).images[0]
return {"image": self.encode_image_base64(result)} |