perf: 二次元识别换用轻量版 wd-vit-tagger-v3（340MB），推理速度提升约 4 倍

internal/hentai/detector.go

@@ -39,7 +39,7 @@ var serverScript []byte
 const detectorStartTimeout = 600 * time.Second
 
 // requestTimeout 单张图片检测的超时时间。
-const requestTimeout = 30 * time.Second
+const requestTimeout = 120 * time.Second
 
 // Detector 封装一个本地 python 服务的客户端。
 type Detector struct {

internal/hentai/server.py

@@ -47,9 +47,9 @@ MODEL_DIR_NAME = "wd-tagger"
 MODEL_FILENAME = "model.onnx"
 TAGS_FILENAME = "selected_tags.csv"
 
-# HuggingFace 下载地址（wd-vit-large-tagger-v3）
+# HuggingFace 下载地址（wd-vit-tagger-v3，轻量版，约 340MB）
-HF_MODEL_URL = "https://huggingface.co/SmilingWolf/wd-vit-large-tagger-v3/resolve/main/model.onnx"
+HF_MODEL_URL = "https://huggingface.co/SmilingWolf/wd-vit-tagger-v3/resolve/main/model.onnx"
-HF_TAGS_URL = "https://huggingface.co/SmilingWolf/wd-vit-large-tagger-v3/resolve/main/selected_tags.csv"
+HF_TAGS_URL = "https://huggingface.co/SmilingWolf/wd-vit-tagger-v3/resolve/main/selected_tags.csv"
 
 # 模型输入尺寸
 INPUT_SIZE = 448
@@ -68,8 +68,8 @@ RATING_TAGS = {
 # 模型目录由 Go 侧作为第二个参数传入
 MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else ""
 
-# 推理并发数
+# 推理并发数（ViT-Large 模型较大，默认 2 避免 CPU 争抢）
-INFER_WORKERS = int(os.environ.get("KM_HENTAI_WORKERS", "4") or "4")
+INFER_WORKERS = int(os.environ.get("KM_HENTAI_WORKERS", "2") or "2")
 if INFER_WORKERS < 1:
     INFER_WORKERS = 1
 
@@ -199,7 +199,7 @@ def _preprocess(img_bytes):
     步骤：
     1. 解码为 PIL Image
     2. 转为 RGB（如果有 alpha 通道）
-    3. 缩放到 INPUT_SIZE x INPUT_SIZE（保持宽高比，用 padding 补齐）
+    3. 直接 resize 到 INPUT_SIZE x INPUT_SIZE（官方推荐方式）
     4. 归一化（ImageNet 参数）
     5. 转为 NHWC 格式的 float32 数组（模型期望 [batch, H, W, C]）
     """
@@ -207,21 +207,11 @@ def _preprocess(img_bytes):
     if img.mode != "RGB":
         img = img.convert("RGB")
 
-    # 保持宽高比缩放，短边补白色
+    # 直接 resize 到目标尺寸（wd-tagger 官方推荐）
-    w, h = img.size
+    img = img.resize((INPUT_SIZE, INPUT_SIZE), Image.LANCZOS)
-    scale = INPUT_SIZE / max(w, h)
-    new_w = int(w * scale)
-    new_h = int(h * scale)
-    img = img.resize((new_w, new_h), Image.LANCZOS)
-
-    # 创建白色背景，把缩放后的图贴到中间
-    canvas = Image.new("RGB", (INPUT_SIZE, INPUT_SIZE), (255, 255, 255))
-    paste_x = (INPUT_SIZE - new_w) // 2
-    paste_y = (INPUT_SIZE - new_h) // 2
-    canvas.paste(img, (paste_x, paste_y))
 
     # 转 numpy 数组并归一化
-    arr = np.array(canvas, dtype=np.float32) / 255.0
+    arr = np.array(img, dtype=np.float32) / 255.0
     arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
 
     # HWC -> NHWC (模型期望 [batch, height, width, channels])