shootstuff/LUSTIFY-v2.0

1import base642import io3from typing import Any, Dict4 5import torch6from PIL import Image7from diffusers import (8    StableDiffusionXLPipeline,9    StableDiffusionXLImg2ImgPipeline,10    DPMSolverMultistepScheduler,11)12 13DEVICE = "cuda" if torch.cuda.is_available() else "cpu"14DTYPE = torch.float16 if DEVICE == "cuda" else torch.float3215 16 17def _decode_image(b64: str) -> Image.Image:18    """Decode a base64 string (optionally a data: URL) into a PIL RGB image."""19    if b64.strip().startswith("data:") and "," in b64:20        b64 = b64.split(",", 1)[1]21    raw = base64.b64decode(b64)22    return Image.open(io.BytesIO(raw)).convert("RGB")23 24 25def _encode_image(img: Image.Image) -> str:26    """Encode a PIL image as a base64 PNG string."""27    buf = io.BytesIO()28    img.save(buf, format="PNG")29    return base64.b64encode(buf.getvalue()).decode("utf-8")30 31 32class EndpointHandler:33    """34    Dual-mode SDXL handler for LUSTIFY-v2.0.35 36    Request shape (HF Inference Endpoints):37        {38          "inputs": "<prompt>",39          "parameters": {40            "negative_prompt": "...",        # optional41            "num_inference_steps": 30,        # optional42            "guidance_scale": 5.0,            # optional (author recommends 4-7)43            "width": 1024, "height": 1024,    # txt2img only44            "seed": 12345,                    # optional, for reproducibility45            "image": "<base64>",              # PRESENCE switches to img2img46            "strength": 0.6                   # img2img only (0-1)47          }48        }49 50    Response: {"image": "<base64 png>", "mode": "txt2img"|"img2img", "parameters": {...}}51    """52 53    def __init__(self, path: str = ""):54        # Base text-to-image pipeline. add_watermarker=False avoids the optional55        # invisible-watermark dependency.56        self.txt2img = StableDiffusionXLPipeline.from_pretrained(57            path,58            torch_dtype=DTYPE,59            use_safetensors=True,60            add_watermarker=False,61        )62        # DPM++ 2M SDE Karras — the checkpoint author's recommended sampler.63        self.txt2img.scheduler = DPMSolverMultistepScheduler.from_config(64            self.txt2img.scheduler.config,65            algorithm_type="sde-dpmsolver++",66            use_karras_sigmas=True,67        )68        self.txt2img.to(DEVICE)69 70        # img2img reuses the exact same weights/components — no extra VRAM cost.71        self.img2img = StableDiffusionXLImg2ImgPipeline(**self.txt2img.components)72        self.img2img.to(DEVICE)73 74        if DEVICE == "cuda":75            self.txt2img.enable_vae_slicing()76            try:77                self.txt2img.enable_xformers_memory_efficient_attention()78                self.img2img.enable_xformers_memory_efficient_attention()79            except Exception:80                # xformers is optional; the pipelines run fine without it.81                pass82 83    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:84        prompt = data.get("inputs") or data.get("prompt")85        params = data.get("parameters") or {}86 87        if not prompt:88            return {"error": "No prompt provided. Send {'inputs': '<prompt>'}."}89 90        negative_prompt = params.get("negative_prompt")91        num_inference_steps = int(params.get("num_inference_steps", 30))92        guidance_scale = float(params.get("guidance_scale", 5.0))93        width = int(params.get("width", 1024))94        height = int(params.get("height", 1024))95 96        seed = params.get("seed")97        generator = None98        if seed is not None:99            generator = torch.Generator(device=DEVICE).manual_seed(int(seed))100 101        init_b64 = params.get("image")102        strength = float(params.get("strength", 0.6))103 104        try:105            if init_b64:106                init_image = _decode_image(init_b64)107                result = self.img2img(108                    prompt=prompt,109                    negative_prompt=negative_prompt,110                    image=init_image,111                    strength=strength,112                    num_inference_steps=num_inference_steps,113                    guidance_scale=guidance_scale,114                    generator=generator,115                )116                mode = "img2img"117            else:118                result = self.txt2img(119                    prompt=prompt,120                    negative_prompt=negative_prompt,121                    width=width,122                    height=height,123                    num_inference_steps=num_inference_steps,124                    guidance_scale=guidance_scale,125                    generator=generator,126                )127                mode = "txt2img"128        except Exception as e:129            return {130                "error": f"{type(e).__name__}: {e}",131                "mode": "img2img" if init_b64 else "txt2img",132            }133 134        image = result.images[0]135        return {136            "image": _encode_image(image),137            "mode": mode,138            "parameters": {139                "num_inference_steps": num_inference_steps,140                "guidance_scale": guidance_scale,141                "strength": strength if init_b64 else None,142                "width": width,143                "height": height,144                "seed": int(seed) if seed is not None else None,145            },146        }147