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]
                }
            ]
        }