import io
import logging
import os

# ！！！必须在 import transformers 之前设置，否则 huggingface_hub 读不到这个环境变量！！！
if "HF_ENDPOINT" not in os.environ:
    os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import torch
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import Response
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="BiRefNet Matting Service")

# 模块加载时就把模型读进内存/显存，常驻复用（BiRefNet_lite = small 变体）
# 模型缓存走 HuggingFace 默认目录（~/.cache/huggingface）：本地存在则离线加载，不存在则自动下载
# 可通过环境变量 BIREFNET_REPO_ID 覆盖模型来源（本地目录或 HuggingFace repo id），
# 默认使用 BiRefNet_lite（small 变体）。
_REPO_ID = os.environ.get("BIREFNET_REPO_ID", "ZhengPeng7/BiRefNet_lite")

try:
    model = AutoModelForImageSegmentation.from_pretrained(
        _REPO_ID,
        trust_remote_code=True,
        local_files_only=True,
    )
    logger.info("BiRefNet 命中本地缓存，离线加载完成")
except (OSError, ValueError) as e:
    logger.info(f"本地未找到 BiRefNet 模型缓存，自动从 HuggingFace 下载 ... ({e})")
    model = AutoModelForImageSegmentation.from_pretrained(
        _REPO_ID,
        trust_remote_code=True,
    )
    logger.info("BiRefNet 模型下载完成")
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
    model.to(device)
    model.half()  # GPU 半精度省显存
else:
    # 模型权重可能为 fp16（如 ModelScope 镜像），CPU 推理统一转 fp32 避免 dtype 不匹配
    model.to(device).float()

transform = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)


@app.post("/matting")
async def matting(file: UploadFile = File(...)):
    raw = await file.read()
    try:
        image = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="invalid image")

    w, h = image.size
    inp = transform(image).unsqueeze(0).to(device)
    if device == "cuda":
        inp = inp.half()

    with torch.no_grad():
        preds = model(inp)[-1].sigmoid().cpu()

    # 单通道概率图 → resize 回原尺寸 → 作为 alpha 合成透明底 PNG
    mask = transforms.ToPILImage()(preds[0].squeeze()).resize((w, h))
    result = image.convert("RGBA")
    result.putalpha(mask)

    buf = io.BytesIO()
    result.save(buf, format="PNG")
    return Response(content=buf.getvalue(), media_type="image/png")


@app.get("/healthz")
def healthz():
    # 能响应即代表模型已加载完成（模块级加载在 import 时完成）
    return {"status": "ok"}


if __name__ == "__main__":
    # 模型已在模块 import 时于模块级加载，这里直接传 app 对象启动，
    # 不要用 "main:app" 字符串 + reload 方式，否则会重复加载模型。
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8001)