from typing import Dict, Any, List
import torch
from diffusers import FluxControlNetModel, FluxControlNetPipeline
from PIL import Image
import requests
from io import BytesIO
import base64
import os
from huggingface_hub import login

class EndpointHandler:
    def __init__(self, path: str = ""):
        login(token=os.getenv("HF_TOKEN"))
        # 加载 ControlNet 模型
        self.controlnet = FluxControlNetModel.from_pretrained(
            path, torch_dtype=torch.bfloat16
        )
        # 加载主流水线（基础模型来自 black-forest-labs/FLUX.1-dev）
        self.pipe = FluxControlNetPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            controlnet=self.controlnet,
            torch_dtype=torch.bfloat16
        )
        self.pipe.to("cuda")

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        # 兼容包装与非包装请求体
        payload = data.get("inputs", data)  # 支持 UI 及直接 POST 两种格式 

        # 读取图像字节
        img_bytes = None
        url = payload.get("control_image_url")
        b64 = payload.get("control_image_base64")
        if b64:
            img_bytes = base64.b64decode(b64)
        elif url:
            resp = requests.get(url)
            resp.raise_for_status()  # 捕获 4xx/5xx 错误 :contentReference[oaicite:3]{index=3}
            img_bytes = resp.content
        else:
            raise ValueError("请在 inputs 中提供 control_image_url 或 control_image_base64")

        # 用 PIL 解析并转换
        try:
            control_image = Image.open(BytesIO(img_bytes)).convert("RGB")
        except Exception as e:
            # 可能是数据损坏或格式不符
            raise ValueError(f"无法识别图像文件: {str(e)}")

        # 可选：调整尺寸（默认 4× 放大）
        w, h = control_image.size
        factor = data.get("upscale_factor", 4)
        control_image = control_image.resize((w * factor, h * factor))

        # 推理参数
        steps = data.get("num_inference_steps", 28)
        scale = data.get("controlnet_conditioning_scale", 0.6)
        guidance = data.get("guidance_scale", 3.5)

        # 执行推理
        output = self.pipe(
            prompt="",
            control_image=control_image,
            num_inference_steps=steps,
            controlnet_conditioning_scale=scale,
            guidance_scale=guidance,
            height=control_image.height,
            width=control_image.width
        )

        # 将 PIL 图像转换为 Base64
        results = []
        for img in output.images:
            buf = BytesIO()
            img.save(buf, format="PNG")
            img_b64 = base64.b64encode(buf.getvalue()).decode()
            results.append({"image_base64": img_b64})

        return results