handler.py

import base64
import io
from typing import Any, Dict

import torch
from PIL import Image
from diffusers import (
    StableDiffusionXLPipeline,
    StableDiffusionXLImg2ImgPipeline,
    DPMSolverMultistepScheduler,
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32


def _decode_image(b64: str) -> Image.Image:
    """Decode a base64 string (optionally a data: URL) into a PIL RGB image."""
    if b64.strip().startswith("data:") and "," in b64:
        b64 = b64.split(",", 1)[1]
    raw = base64.b64decode(b64)
    return Image.open(io.BytesIO(raw)).convert("RGB")


def _encode_image(img: Image.Image) -> str:
    """Encode a PIL image as a base64 PNG string."""
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


class EndpointHandler:
    """
    Dual-mode SDXL handler for LUSTIFY-v2.0.

    Request shape (HF Inference Endpoints):
        {
          "inputs": "<prompt>",
          "parameters": {
            "negative_prompt": "...",        # optional
            "num_inference_steps": 30,        # optional
            "guidance_scale": 5.0,            # optional (author recommends 4-7)
            "width": 1024, "height": 1024,    # txt2img only
            "seed": 12345,                    # optional, for reproducibility
            "image": "<base64>",              # PRESENCE switches to img2img
            "strength": 0.6                   # img2img only (0-1)
          }
        }

    Response: {"image": "<base64 png>", "mode": "txt2img"|"img2img", "parameters": {...}}
    """

    def __init__(self, path: str = ""):
        # Base text-to-image pipeline. add_watermarker=False avoids the optional
        # invisible-watermark dependency.
        self.txt2img = StableDiffusionXLPipeline.from_pretrained(
            path,
            torch_dtype=DTYPE,
            use_safetensors=True,
            add_watermarker=False,
        )
        # DPM++ 2M SDE Karras — the checkpoint author's recommended sampler.
        self.txt2img.scheduler = DPMSolverMultistepScheduler.from_config(
            self.txt2img.scheduler.config,
            algorithm_type="sde-dpmsolver++",
            use_karras_sigmas=True,
        )
        self.txt2img.to(DEVICE)

        # img2img reuses the exact same weights/components — no extra VRAM cost.
        self.img2img = StableDiffusionXLImg2ImgPipeline(**self.txt2img.components)
        self.img2img.to(DEVICE)

        if DEVICE == "cuda":
            self.txt2img.enable_vae_slicing()
            try:
                self.txt2img.enable_xformers_memory_efficient_attention()
                self.img2img.enable_xformers_memory_efficient_attention()
            except Exception:
                # xformers is optional; the pipelines run fine without it.
                pass

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        prompt = data.get("inputs") or data.get("prompt")
        params = data.get("parameters") or {}

        if not prompt:
            return {"error": "No prompt provided. Send {'inputs': '<prompt>'}."}

        negative_prompt = params.get("negative_prompt")
        num_inference_steps = int(params.get("num_inference_steps", 30))
        guidance_scale = float(params.get("guidance_scale", 5.0))
        width = int(params.get("width", 1024))
        height = int(params.get("height", 1024))

        seed = params.get("seed")
        generator = None
        if seed is not None:
            generator = torch.Generator(device=DEVICE).manual_seed(int(seed))

        init_b64 = params.get("image")
        strength = float(params.get("strength", 0.6))

        try:
            if init_b64:
                init_image = _decode_image(init_b64)
                result = self.img2img(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    image=init_image,
                    strength=strength,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=guidance_scale,
                    generator=generator,
                )
                mode = "img2img"
            else:
                result = self.txt2img(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=width,
                    height=height,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=guidance_scale,
                    generator=generator,
                )
                mode = "txt2img"
        except Exception as e:
            return {
                "error": f"{type(e).__name__}: {e}",
                "mode": "img2img" if init_b64 else "txt2img",
            }

        image = result.images[0]
        return {
            "image": _encode_image(image),
            "mode": mode,
            "parameters": {
                "num_inference_steps": num_inference_steps,
                "guidance_scale": guidance_scale,
                "strength": strength if init_b64 else None,
                "width": width,
                "height": height,
                "seed": int(seed) if seed is not None else None,
            },
        }