first commit

thirdparty/matting_server/main.py

@@ -0,0 +1,96 @@
+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)