File size: 3,147 Bytes
5af9786 1d16b0c 5af9786 137b45f 5af9786 137b45f 5af9786 137b45f 5af9786 1d16b0c 137b45f 1d16b0c 5af9786 137b45f 5af9786 1d16b0c 5af9786 137b45f 1d16b0c 5af9786 1d16b0c 5af9786 137b45f 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 82 83 84 85 86 | import torch
import base64
import io
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline, DPMSolverMultistepScheduler
import os
class EndpointHandler():
def __init__(self, path=""):
print("Loading ARX Elite Pipeline (Juggernaut XL)...")
# 1. Point directly to the NEW model
model_path = os.path.join(path, "juggernaut.safetensors")
# 2. Load the core pipeline
self.pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
model_path,
torch_dtype=torch.float16,
use_safetensors=True,
safety_checker=None
)
# 3. Upgrade the Scheduler to DPM++ 2M Karras
self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
self.pipe.scheduler.config,
use_karras_sigmas=True,
algorithm_type="sde-dpmsolver++"
)
# 4. Load IP-Adapter
self.pipe.load_ip_adapter(
"h94/IP-Adapter",
subfolder="sdxl_models",
weight_name="ip-adapter_sdxl.bin"
)
self.pipe.to("cuda")
print("ARX Elite 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):
inputs = data.pop("inputs", data)
prompt = inputs.get("prompt", "masterpiece, best quality, highly detailed, photorealistic")
# Added realism enforcers to the negative prompt
negative_prompt = inputs.get("negative_prompt", "blurry, lowres, bad anatomy, worst quality, ugly, deformed eyes, cartoon, illustration, painting, 3d render")
strength = float(inputs.get("strength", 0.55))
guidance_scale = float(inputs.get("guidance_scale", 7.5))
num_inference_steps = int(inputs.get("steps", 30))
ip_adapter_scale = float(inputs.get("ip_adapter_scale", 0.75))
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": "Missing image inputs."}
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)} |