File size: 1,067 Bytes
005a384 01914b8 3d578ab 6d978fb 01914b8 6d978fb c723b2a 01914b8 6d978fb c723b2a 01914b8 c723b2a 01914b8 c723b2a 01914b8 c723b2a 3fd067c 3d578ab 3fd067c 3bf49e0 3fd067c |
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 |
from diffusers import DiffusionPipeline
import torch
import base64
from io import BytesIO
class EndpointHandler:
def __init__(self, path=""):
print("Loading Juggernaut XL…")
self.pipe = DiffusionPipeline.from_pretrained(
path,
torch_dtype=torch.float16
).to("cuda")
def __call__(self, data):
prompt = data.get("inputs", "")
params = data.get("parameters", {})
steps = params.get("num_inference_steps", 28)
cfg = params.get("guidance_scale", 4.5)
result = self.pipe(
prompt,
num_inference_steps=steps,
guidance_scale=cfg
)
pil = result.images[0]
# Convert to base64
buffer = BytesIO()
pil.save(buffer, format="PNG")
base64_img = base64.b64encode(buffer.getvalue()).decode("utf-8")
# HuggingFace CUSTOM PIPELINE REQUIRED FORMAT
return {
"outputs": [
{
"images": [base64_img]
}
]
} |