97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
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)
|