handler.py

import os
import glob
import io
import re
import base64
import logging
import time
from typing import Dict, List, Any, Union, Optional
import torch
from PIL import Image
from diffusers import (
    AutoPipelineForText2Image,
    StableDiffusionPipeline,
    StableDiffusionXLPipeline,
    UNet2DConditionModel,
)
from transformers import (
    CLIPTextModel,
    CLIPTokenizer,
    CLIPTextModelWithProjection,
    AutoTokenizer,
)
from safetensors import safe_open

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("hf-endpoint-handler")

class EndpointHandler:
    """
    Custom handler for Hugging Face Dedicated Inference Endpoints.
    
    Handles Text-to-Image generation with automatic single-file detection,
    Flux/SDXL/SD1.5 architecture auto-detection, step-checkpoint filtering,
    missing UNet/TextEncoder base fallback, and LoRA loading.
    """

    def __init__(self, path: str = ""):
        logger.info(f"Initializing EndpointHandler from model directory: '{path}'")

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32

        # 1. Try standard diffusers folder load first
        try:
            logger.info("Attempting standard diffusers repository load...")
            self.pipeline = AutoPipelineForText2Image.from_pretrained(
                path,
                torch_dtype=self.dtype,
                use_safetensors=True,
            )
            logger.info("Successfully loaded standard diffusers folder pipeline.")
        except Exception as e:
            logger.warning(f"Standard load failed ({e}). Running smart single-file checkpoint loader...")
            self.pipeline = self._load_single_file_checkpoint(path)

        # 2. Apply GPU optimizations
        if self.device == "cuda":
            self.pipeline.to("cuda")

            if hasattr(self.pipeline, "enable_vae_slicing"):
                self.pipeline.enable_vae_slicing()
            if hasattr(self.pipeline, "enable_vae_tiling"):
                self.pipeline.enable_vae_tiling()

        logger.info(f"Pipeline successfully ready on {self.device} with class: {type(self.pipeline).__name__}")

    def _select_best_checkpoint(self, path: str) -> str:
        """
        Filters out training step checkpoints (_000002250, optimizer.pt, etc.)
        and prioritizes merged or main model weight files.
        """
        candidate_files = []
        if os.path.isdir(path):
            for ext in ("*.safetensors", "*.ckpt", "*.pt", "*.bin"):
                candidate_files.extend(glob.glob(os.path.join(path, ext)))
                candidate_files.extend(glob.glob(os.path.join(path, "**", ext), recursive=True))
        elif os.path.isfile(path):
            return path

        if not candidate_files:
            raise FileNotFoundError(f"No model checkpoint files found in directory '{path}'")

        filtered_files = []
        for f in candidate_files:
            fname = os.path.basename(f).lower()
            # Ignore optimizer files and subfolder weights
            if "optimizer" in fname or "text_encoder" in f or "unet" in f or "vae" in f:
                continue
            # Ignore intermediate training step checkpoints (e.g. _000002250.safetensors)
            if re.search(r"_\d{5,}", fname) or re.search(r"step_?\d+", fname):
                logger.info(f"Filtering out intermediate step checkpoint file: {os.path.basename(f)}")
                continue
            filtered_files.append(f)

        if not filtered_files:
            logger.warning("All files matched step pattern, falling back to full list...")
            filtered_files = candidate_files

        # Priority 1: Merged files
        merged = [f for f in filtered_files if "merged" in os.path.basename(f).lower()]
        if merged:
            logger.info(f"Selected merged checkpoint: {os.path.basename(merged[0])}")
            return merged[0]

        # Priority 2: Flux files
        flux = [f for f in filtered_files if "flux" in os.path.basename(f).lower()]
        if flux:
            logger.info(f"Selected Flux model checkpoint: {os.path.basename(flux[0])}")
            return flux[0]

        logger.info(f"Selected primary checkpoint file: {os.path.basename(filtered_files[0])}")
        return filtered_files[0]

    def _inspect_keys(self, target_file: str) -> dict:
        """
        Inspects safetensors header keys to determine architecture and component presence.
        """
        is_lora = False
        is_flux = False
        is_sdxl = False
        has_unet = False

        if target_file.endswith(".safetensors"):
            try:
                with safe_open(target_file, framework="pt") as f:
                    keys = f.keys()
                    for k in keys:
                        if "lora_" in k or ".lora_down" in k or ".lora_up" in k:
                            is_lora = True
                        if "double_blocks" in k or "single_blocks" in k or "guidance_in" in k:
                            is_flux = True
                        if "conditioner.embedders" in k or "text_encoders.top" in k:
                            is_sdxl = True
                        if "model.diffusion_model" in k or "unet" in k:
                            has_unet = True
            except Exception as e:
                logger.warning(f"Failed to inspect safetensors keys: {e}")

        filename_lower = os.path.basename(target_file).lower()
        if "flux" in filename_lower:
            is_flux = True

        return {
            "is_lora": is_lora,
            "is_flux": is_flux,
            "is_sdxl": is_sdxl,
            "has_unet": has_unet,
        }

    def _load_single_file_checkpoint(self, path: str):
        target_file = self._select_best_checkpoint(path)
        info = self._inspect_keys(target_file)
        logger.info(f"Checkpoint inspection result for {os.path.basename(target_file)}: {info}")

        # --- CASE A: FLUX Architecture ---
        if info["is_flux"]:
            logger.info("Flux architecture detected. Initializing Flux pipeline...")
            try:
                from diffusers import FluxPipeline
                if info["is_lora"]:
                    logger.info("Loading Flux base model and attaching LoRA weights...")
                    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=self.dtype)
                    pipe.load_lora_weights(target_file)
                    return pipe
                else:
                    try:
                        return FluxPipeline.from_single_file(target_file, torch_dtype=self.dtype)
                    except Exception as err:
                        logger.warning(f"Flux single-file load failed ({err}), loading base FLUX.1-schnell...")
                        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=self.dtype)
                        if info["is_lora"] or "merged" in target_file.lower():
                            try:
                                pipe.load_lora_weights(target_file)
                            except Exception:
                                pass
                        return pipe
            except Exception as flux_err:
                logger.warning(f"Flux pipeline load failed ({flux_err}). Falling back to SDXL...")

        # --- CASE B: SDXL Architecture ---
        if info["is_sdxl"] or "xl" in os.path.basename(target_file).lower():
            logger.info("SDXL architecture detected. Attempting SDXL single file load...")
            if info["is_lora"]:
                pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=self.dtype)
                pipe.load_lora_weights(target_file)
                return pipe

            try:
                return StableDiffusionXLPipeline.from_single_file(target_file, torch_dtype=self.dtype)
            except Exception as err:
                logger.warning(f"SDXL single-file missing components ({err}). Supplying SDXL base components...")
                unet = UNet2DConditionModel.from_pretrained(
                    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", torch_dtype=self.dtype
                )
                text_encoder_1 = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=self.dtype)
                text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
                    "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", torch_dtype=self.dtype
                )
                tokenizer_1 = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
                tokenizer_2 = AutoTokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")

                try:
                    return StableDiffusionXLPipeline.from_single_file(
                        target_file,
                        unet=unet,
                        text_encoder=text_encoder_1,
                        text_encoder_2=text_encoder_2,
                        tokenizer=tokenizer_1,
                        tokenizer_2=tokenizer_2,
                        torch_dtype=self.dtype,
                    )
                except Exception:
                    pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=self.dtype)
                    try:
                        pipe.load_lora_weights(target_file)
                    except Exception:
                        pass
                    return pipe

        # --- CASE C: Standard SD 1.5 Architecture ---
        logger.info("Attempting SD 1.5 single file pipeline load...")
        if info["is_lora"]:
            pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=self.dtype)
            pipe.load_lora_weights(target_file)
            return pipe

        try:
            return StableDiffusionPipeline.from_single_file(target_file, torch_dtype=self.dtype)
        except Exception as err:
            logger.warning(f"SD 1.5 single file missing components ({err}). Supplying base UNet and CLIP...")
            unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", torch_dtype=self.dtype)
            text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=self.dtype)
            tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")

            try:
                return StableDiffusionPipeline.from_single_file(
                    target_file,
                    unet=unet,
                    text_encoder=text_encoder,
                    tokenizer=tokenizer,
                    torch_dtype=self.dtype,
                )
            except Exception:
                pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=self.dtype)
                try:
                    pipe.load_lora_weights(target_file)
                except Exception:
                    pass
                return pipe

    def _encode_image_to_base64(self, image: Image.Image, image_format: str = "PNG", quality: int = 95) -> str:
        buffer = io.BytesIO()
        if image_format.upper() in ["JPG", "JPEG"]:
            image.save(buffer, format="JPEG", quality=quality)
        else:
            image.save(buffer, format="PNG")
        buffer.seek(0)
        return base64.b64encode(buffer.getvalue()).decode("utf-8")

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

        if isinstance(inputs, dict):
            prompt = inputs.get("prompt", "")
            negative_prompt = inputs.get("negative_prompt", parameters.get("negative_prompt", None))
        else:
            prompt = str(inputs) if inputs else ""
            negative_prompt = parameters.get("negative_prompt", None)

        if not prompt.strip():
            raise ValueError("Parameter 'prompt' (or 'inputs') must be a non-empty string.")

        pipe_name = type(self.pipeline).__name__.lower()
        default_dim = 1024 if ("xl" in pipe_name or "flux" in pipe_name) else 512

        height = (int(parameters.get("height", default_dim)) // 8) * 8
        width = (int(parameters.get("width", default_dim)) // 8) * 8

        num_inference_steps = int(parameters.get("num_inference_steps", 28 if "flux" in pipe_name else 30))
        guidance_scale = float(parameters.get("guidance_scale", 3.5 if "flux" in pipe_name else 7.5))
        seed = parameters.get("seed", None)
        num_images_per_prompt = min(max(int(parameters.get("num_images_per_prompt", 1)), 1), 4)

        output_format = str(parameters.get("output_format", "pil")).lower()
        image_format = str(parameters.get("image_format", "PNG")).upper()

        return {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "height": height,
            "width": width,
            "num_inference_steps": num_inference_steps,
            "guidance_scale": guidance_scale,
            "seed": int(seed) if seed is not None else None,
            "num_images_per_prompt": num_images_per_prompt,
            "output_format": output_format,
            "image_format": image_format,
        }

    def __call__(self, data: Dict[str, Any]) -> Union[List[Dict[str, Any]], Image.Image]:
        start_time = time.time()

        try:
            params = self._extract_parameters(data)
            logger.info(f"Processing prompt: '{params['prompt'][:60]}...'")

            generator = None
            if params["seed"] is not None:
                generator = torch.Generator(device=self.device).manual_seed(params["seed"])

            generation_args = {
                "prompt": params["prompt"],
                "negative_prompt": params["negative_prompt"],
                "height": params["height"],
                "width": params["width"],
                "num_inference_steps": params["num_inference_steps"],
                "guidance_scale": params["guidance_scale"],
                "num_images_per_prompt": params["num_images_per_prompt"],
                "generator": generator,
            }

            # Flux pipelines do not use negative_prompt
            if "flux" in type(self.pipeline).__name__.lower():
                generation_args.pop("negative_prompt", None)

            # Filter out None values
            generation_args = {k: v for k, v in generation_args.items() if v is not None}

            with torch.inference_mode():
                output = self.pipeline(**generation_args)
                images = output.images

            elapsed_seconds = round(time.time() - start_time, 3)
            logger.info(f"Generated {len(images)} image(s) in {elapsed_seconds}s")

            if params["output_format"] == "pil" and len(images) == 1:
                return images[0]

            response_payload = []
            for index, img in enumerate(images):
                b64_image = self._encode_image_to_base64(img, image_format=params["image_format"])
                response_payload.append({
                    "image": b64_image,
                    "format": params["image_format"],
                    "width": img.width,
                    "height": img.height,
                    "index": index,
                    "execution_time": elapsed_seconds,
                })

            return response_payload

        except Exception as err:
            logger.error(f"Inference execution failed: {str(err)}", exc_info=True)
            return [{"error": str(err), "status": "failed"}]