import base64
import io
import torch
from diffusers.pipelines.glm_image import GlmImagePipeline

class EndpointHandler:
    def __init__(self, path=""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.bfloat16
        self.pipe = GlmImagePipeline.from_pretrained(
            "zai-org/GLM-Image",
            torch_dtype=self.dtype,
            device_map="cuda",
            enable_model_cpu_offload=True,
        )

    def __call__(self, data):
        prompt = data.pop("inputs", "")
        params = data.pop("parameters", {})

        width = params.get("width", 1024)
        height = params.get("height", 1024)
        num_inference_steps = params.get("num_inference_steps", 50)
        guidance_scale = params.get("guidance_scale", 1.5)

        # GLM-Image requires dimensions divisible by 32
        width = (width // 32) * 32
        height = (height // 32) * 32

        image = self.pipe(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
        ).images[0]

        buf = io.BytesIO()
        image.save(buf, format="PNG")
        img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        return {"image": img_b64, "format": "png"}