Upgrade to PixAI Tagger v1.0 and prefer Index-Translate

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+.venv/
+assets/
+assets-v1.0/
+.cache/

README.md

@@ -14,6 +14,47 @@ hf_oauth: true
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
+## 模型与翻译
+
+默认使用 [PixAI Tagger v1.0](https://huggingface.co/pixai-labs/pixai-tagger-v1.0)，
+模型权重、配置与官方推理代码固定在 `9fe10addf9326e292da8a85a98ea74cd91b41771`。
+采用官方 1008 × 1008 等比缩放和补边预处理，透明背景合成为白色。
+
+- 页面及任务 API 的默认阈值：通用 `0.17`、角色 `0.27`；请求中显式设置的阈值仍会生效。
+- IP 标签使用模型 `copyright` 类别及其真实置信度，默认阈值 `0.24`。
+- 原有三个展示类别及任务 API 六个顶层返回字段保持兼容。
+- 六类原始标签及分数保存在页面“推理元数据”与任务结果的 `metadata.tags` 中：
+  `general`、`character`、`copyright`、`style`、`meta`、`rating`。
+  后三类使用官方推荐阈值 `0.15`、`0.17`、`0.41`。
+
+翻译默认调用 [Index-Translate 官方公开 API](https://github.com/bilibili/Index-Translate)：
+`https://index-translate.bilibili.com/v1/chat/completions`，模型 `Index-Translate-35B-A3B`，默认无需密钥。
+标签分批翻译，使用编号核对译文，保持标签与翻译一一对应。Index 请求失败或译文校验失败时，整组标签按顺序回退到原腾讯云账号池、百度账号池，全部失败时保留原文。
+
+可选翻译配置：
+
+| Variable / Secret | 默认值 / 用途 |
+| --- | --- |
+| `INDEX_TRANSLATE_URL` | 上述完整接口地址 |
+| `INDEX_TRANSLATE_MODEL` | `Index-Translate-35B-A3B` |
+| `INDEX_TRANSLATE_TIMEOUT` | `30` 秒 |
+| `INDEX_TRANSLATE_API_KEY`（Secret） | 仅接入需要鉴权的兼容服务时设置 |
+
+原 `TENCENT_CREDENTIALS_JSON`、`TENCENT_SECRET_ID`、`TENCENT_SECRET_KEY`、
+`BAIDU_CREDENTIALS_JSON` 等凭据配置继续用于回退，应保存在 Space Secrets 中。
+HF 模型下载令牌不会发送给翻译服务。
+
+升级部署检查：
+
+1. 安装更新后的 `requirements.txt`，保留现有 Job API Secrets 和图片域名白名单。
+2. 如设置过 `ASSETS_REPO_ID` / `ASSETS_REVISION`，更新为上述 v1.0 仓库和提交；旧版资源不能用于新 handler。
+   `MODEL_DIR` 现在作为 Hub 缓存目录，默认 `./assets-v1.0`；运行用户必须可写。
+3. 首次启动需下载约 1.95 GB 权重。CPU 可运行，页面和 Job API 共用一个模型实例及推理锁；请预留下载空间和启动时间。
+4. 等 Space 构建完成，确认启动日志显示 v1.0 加载成功，再测试页面图片反推和 `/api/jobs` 创建/轮询。
+5. 如需回滚，将 Space 代码恢复到升级前提交 `57942b0f23e74896acf5f1588e8a7d4810bc284f`，并恢复原模型变量，再等待重新构建。
+
+本地关键测试：`python -m unittest discover -s tests -v`。
+
 ## 自定义任务 API
 
 该 Space 在保留原 Gradio 页面功能的同时提供受共享密钥保护的异步标签分析接口：
@@ -27,8 +68,8 @@ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-
 ```json
 {
   "input_image_url": "https://example.com/input.png",
-  "general_threshold": 0.35,
-  "character_threshold": 0.85,
+  "general_threshold": 0.17,
+  "character_threshold": 0.27,
   "show_confidence": true,
   "show_general": true,
   "show_character": true,

app.py

@@ -5,7 +5,6 @@ os.environ.setdefault("HF_HUB_DISABLE_IMPLICIT_TOKEN", "1")
 
 import json
 import time
-import shutil
 import threading
 import warnings
 from html import escape
@@ -121,10 +120,8 @@ def validate_and_open_image(image_path: str) -> Image.Image:
     try:
         with Image.open(image_path) as img:
             img.load()
-            if img.mode != "RGB":
-                img = img.convert("RGB")
-            else:
-                img = img.copy()
+            # 保留透明通道，让模型官方预处理统一合成白底。
+            img = img.copy()
     except Image.DecompressionBombWarning:
         raise ImageValidationError("图片在解码阶段触发像素炸弹保护，已拒绝处理。")
     except Exception as e:
@@ -134,57 +131,53 @@ def validate_and_open_image(image_path: str) -> Image.Image:
 
 
 # ------------------------------------------------------------------
-# 新版 PixAI Tagger v0.9 模型配置
+# PixAI Tagger v1.0 模型配置
 # ------------------------------------------------------------------
-ASSETS_REPO_ID = os.environ.get("ASSETS_REPO_ID", "pixai-labs/pixai-tagger-v0.9")
-ASSETS_REVISION = os.environ.get("ASSETS_REVISION")
+ASSETS_REPO_ID = os.environ.get("ASSETS_REPO_ID", "pixai-labs/pixai-tagger-v1.0")
+# 固定模型权重和自定义推理代码，避免上游 main 更新后混用不同版本。
+ASSETS_REVISION = os.environ.get("ASSETS_REVISION", "9fe10addf9326e292da8a85a98ea74cd91b41771")
 ASSETS_HF_TOKEN = os.environ.get("HF_TOKEN")
-MODEL_DIR = os.environ.get("MODEL_DIR", "./assets")
+MODEL_DIR = os.environ.get("MODEL_DIR", "./assets-v1.0")
 
 REQUIRED_FILES = [
-    "model_v0.9.pth",
-    "tags_v0.9_13k.json",
-    "char_ip_map.json",
+    "config.json",
+    "preprocessor_config.json",
+    "tagger_pipeline.py",
+    "model.safetensors",
 ]
 
 
-def ensure_assets(repo_id: str, revision: Optional[str], target_dir: str) -> None:
-    """下载受限标签模型资源，并复制到 handler 期望的本地目录。
+def ensure_assets(repo_id: str, revision: Optional[str], target_dir: str) -> str:
+    """下载同一版本的模型、配置和官方推理代码，返回本地快照路径。
 
     Args:
         repo_id: Hugging Face 模型仓库 ID。
         revision: 可选模型版本；为空时使用仓库默认版本。
-        target_dir: handler 读取模型资源的本地目录。
+        target_dir: Hugging Face 模型缓存目录。
 
     Returns:
-        None: 文件已存在或下载复制完成后不返回数据。
+        已完整下载的模型快照目录。
     """
     target = Path(target_dir)
     target.mkdir(parents=True, exist_ok=True)
 
-    missing = [fname for fname in REQUIRED_FILES if not (target / fname).exists()]
-    if not missing:
-        return
-
     snapshot_path = snapshot_download(
         repo_id=repo_id,
         revision=revision,
         allow_patterns=REQUIRED_FILES,
-        # 该受限仓库必须鉴权；token 只发往 Hub 下载端点，不参与任何 Job API 请求。
+        cache_dir=str(target),
+        # token 仅用于 Hub 下载；模型加载与翻译请求均不传递此凭据。
         token=ASSETS_HF_TOKEN or False,
     )
 
     for fname in REQUIRED_FILES:
         src = Path(snapshot_path) / fname
-        dst = target / fname
-
         if not src.exists():
             raise FileNotFoundError(
                 f"模型资源缺失：'{fname}' 未在 {repo_id} @ {revision or 'default'} 中找到。"
             )
 
-        if src.resolve() != dst.resolve():
-            shutil.copyfile(src, dst)
+    return snapshot_path
 
 
 # ------------------------------------------------------------------
@@ -198,48 +191,28 @@ class Tagger:
 
     def _load_model_and_labels(self) -> None:
         try:
-            ensure_assets(ASSETS_REPO_ID, ASSETS_REVISION, MODEL_DIR)
-            self.handler = EndpointHandler(MODEL_DIR)
+            snapshot_path = ensure_assets(ASSETS_REPO_ID, ASSETS_REVISION, MODEL_DIR)
+            self.handler = EndpointHandler(snapshot_path)
             self.device = getattr(self.handler, "device", "unknown")
-            print(f"✅ PixAI Tagger v0.9 加载成功，设备：{str(self.device).upper()}")
+            print(f"✅ PixAI Tagger v1.0 加载成功，设备：{str(self.device).upper()}")
         except Exception as e:
-            print(f"❌ PixAI Tagger v0.9 加载失败: {e}")
+            print(f"❌ PixAI Tagger v1.0 加载失败: {e}")
             raise RuntimeError(f"模型初始化失败: {e}") from e
 
     @staticmethod
     def _display_tag(tag: str) -> str:
         return str(tag).replace("_", " ")
 
-    @staticmethod
-    def _get_score(scores: dict, tag: str) -> float:
-        """
-        handler 通常以原始 tag 作为分数字典 key。
-        这里额外兼容空格/下划线两种写法，避免 key 不一致时取不到分数。
-        """
-        if not isinstance(scores, dict):
-            return 0.0
-
-        candidates = [
-            tag,
-            str(tag).replace("_", " "),
-            str(tag).replace(" ", "_"),
-        ]
-
-        for key in candidates:
-            if key in scores:
-                try:
-                    return float(scores[key])
-                except Exception:
-                    return 0.0
+    def predict(self, img: Image.Image, gen_th: float = 0.17, char_th: float = 0.27):
+        """运行 v1.0 反推并生成页面和任务 API 共用的标签结果。
 
-        return 0.0
+        Args:
+            img: 已校验的 PIL 图片，透明背景由官方预处理处理。
+            gen_th: 通用标签的置信度阈值。
+            char_th: 角色标签的置信度阈值。
 
-    def predict(self, img: Image.Image, gen_th: float = 0.30, char_th: float = 0.85):
-        """
-        返回结构保持原 app.py 的 UI 处理习惯：
-        - general：通用/特征标签，带置信度
-        - characters：角色标签，带置信度
-        - ips：IP 标签，新模型不返回评分标签，因此原 ratings 改为 ips，且 IP 不展示伪造置信度
+        Returns:
+            三类展示标签、对应翻译顺序，以及包含六类原始标签的元数据。
         """
         if self.handler is None:
             raise RuntimeError("模型未成功加载，无法进行预测。")
@@ -265,48 +238,37 @@ class Tagger:
         out = self.handler(data)
         latency = round(time.time() - started, 4)
 
-        feature_tags = out.get("feature", []) or []
-        character_tags = out.get("character", []) or []
-        ip_tags = out.get("ip", []) or []
-
-        feature_scores = out.get("feature_scores", {}) or {}
-        character_scores = out.get("character_scores", {}) or {}
-
-        general = {
-            self._display_tag(tag): self._get_score(feature_scores, tag)
-            for tag in feature_tags
-        }
-        characters = {
-            self._display_tag(tag): self._get_score(character_scores, tag)
-            for tag in character_tags
-        }
-
-        # IP 标签没有评分，使用 None 表示“不显示置信度”
-        ips = {
-            self._display_tag(tag): None
-            for tag in ip_tags
-        }
-
-        general = dict(sorted(general.items(), key=lambda kv: kv[1], reverse=True))
-        characters = dict(sorted(characters.items(), key=lambda kv: kv[1], reverse=True))
-
-        res = {
-            "general": general,
-            "characters": characters,
-            "ips": ips,
-        }
-
+        # 保持页面和既有任务 API 的三个展示类别，IP 使用 v1.0 的真实置信度。
+        res = {}
+        for display_category, model_category in (
+            ("general", "feature"), ("characters", "character"), ("ips", "ip")
+        ):
+            res[display_category] = {
+                self._display_tag(tag): score
+                for tag, score in sorted(
+                    out[f"{model_category}_scores"].items(),
+                    key=lambda item: item[1], reverse=True,
+                )
+            }
         tag_categories_for_translation = {
-            "general": list(general.keys()),
-            "characters": list(characters.keys()),
-            "ips": list(ips.keys()),
+            category: list(tags) for category, tags in res.items()
         }
 
         raw_meta = {
+            "model": ASSETS_REPO_ID,
+            "revision": ASSETS_REVISION,
             "device": str(self.device),
             "latency_s_total": latency,
-            "_params": out.get("_params", params),
-            "_timings": out.get("_timings", {}),
+            "_params": out["_params"],
+            "_timings": out["_timings"],
+            "tags": {
+                "general": out["feature_scores"],
+                "character": out["character_scores"],
+                "copyright": out["ip_scores"],
+                "style": out["style_scores"],
+                "meta": out["meta_scores"],
+                "rating": out["rating_scores"],
+            },
         }
 
         return res, tag_categories_for_translation, raw_meta
@@ -337,7 +299,7 @@ def predict_tags_with_shared_slot(
     """在共享单推理门闩内调用现有标签模型。
 
     Args:
-        image: 已完成安全校验与 RGB 转换的输入图片。
+        image: 已完成安全校验的输入图片。
         general_threshold: 通用标签阈值。
         character_threshold: 角色标签阈值。
 
@@ -470,8 +432,8 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
     gr.Markdown("# 🖼️ AI 图像标签分析器")
     gr.Markdown(
         "上传图片自动识别标签，支持中英文显示和一键复制。"
-        f"**当前模型：pixai-labs/pixai-tagger-v0.9**　|　**{DEVICE_LABEL}**\n\n"
-        "说明：新版模型不再返回评分标签，本页面已将原“评分标签”区域改为“IP 标签”。"
+        f"**当前模型：{ASSETS_REPO_ID}**　|　**{DEVICE_LABEL}**\n\n"
+        "通用、角色和 IP 标签支持中英文展示；风格、元数据和评级标签可在下方“推理元数据”的 tags 中查看。"
     )
 
     state_res = gr.State({})
@@ -480,7 +442,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
 
     with gr.Row():
         with gr.Column(scale=1):
-            img_in = gr.Image(type="filepath", label="上传图片", height=300)
+            img_in = gr.Image(type="filepath", image_mode=None, label="上传图片", height=300)
 
             btn = gr.Button("🚀 开始分析", variant="primary", elem_classes=["btn-analyze-container"])
 
@@ -488,7 +450,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
                 gen_slider = gr.Slider(
                     0,
                     1,
-                    value=0.30,
+                    value=0.17,
                     step=0.01,
                     label="通用标签阈值",
                     info="越高 → 标签更少更准",
@@ -496,15 +458,15 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
                 char_slider = gr.Slider(
                     0,
                     1,
-                    value=0.85,
+                    value=0.27,
                     step=0.01,
                     label="角色标签阈值",
-                    info="推荐保持较高阈值",
+                    info="v1.0 推荐 0.27；提高阈值可减少误识别",
                 )
                 show_tag_scores = gr.Checkbox(
                     True,
                     label="在列表中显示标签置信度",
-                    info="IP 标签不返回置信度，因此不会显示分数。",
+                    info="通用、角色和 IP 标签均显示模型置信度。",
                 )
 
             with gr.Accordion("📊 标签汇总设置", open=True):
@@ -523,10 +485,10 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
                 with gr.TabItem("🏷️ 通用标签"):
                     out_general = gr.HTML(label="General Tags")
                 with gr.TabItem("👤 角色标签"):
-                    gr.Markdown("<p style='color:gray; font-size:small;'>提示：角色标签由模型推断，建议保持较高阈值。</p>")
+                    gr.Markdown("<p style='color:gray; font-size:small;'>提示：角色标签由模型推断，默认使用 v1.0 推荐阈值 0.27。</p>")
                     out_char = gr.HTML(label="Character Tags")
                 with gr.TabItem("🌐 IP 标签"):
-                    gr.Markdown("<p style='color:gray; font-size:small;'>提示：新版模型输出 IP 标签，但不返回评分标签/评分置信度。</p>")
+                    gr.Markdown("<p style='color:gray; font-size:small;'>提示：IP 标签来自 v1.0 的版权/作品类别，支持真实置信度。</p>")
                     out_ip = gr.HTML(label="IP Tags")
 
             gr.Markdown("### 标签汇总结果")
@@ -858,7 +820,7 @@ def execute_image_inversion_job(
 
     Args:
         payload: 已通过公开 Schema 校验的标签分析参数。
-        input_image: 已由任务 API 安全抓取并转换为 RGB 的图片。
+        input_image: 已由任务 API 安全抓取并解码的图片。
         job_id: 当前异步任务标识，仅用于隔离任务上下文。
 
     Returns:

handler.py

@@ -1,6 +1,5 @@
 import base64
 import io
-import json
 import logging
 import os
 import time
@@ -8,272 +7,130 @@ from pathlib import Path
 from typing import Any
 
 import requests
-import timm
 import torch
-import torchvision.transforms as transforms
 from PIL import Image
-
-
-class TaggingHead(torch.nn.Module):
-    def __init__(self, input_dim, num_classes):
-        super().__init__()
-        self.input_dim = input_dim
-        self.num_classes = num_classes
-        self.head = torch.nn.Sequential(torch.nn.Linear(input_dim, num_classes))
-
-    def forward(self, x):
-        logits = self.head(x)
-        probs = torch.nn.functional.sigmoid(logits)
-        return probs
-
-
-def get_tags(tags_file: Path) -> tuple[dict[str, int], int, int]:
-    with tags_file.open("r", encoding="utf-8") as f:
-        tag_info = json.load(f)
-    tag_map = tag_info["tag_map"]
-    tag_split = tag_info["tag_split"]
-    gen_tag_count = tag_split["gen_tag_count"]
-    character_tag_count = tag_split["character_tag_count"]
-    return tag_map, gen_tag_count, character_tag_count
-
-
-def get_character_ip_mapping(mapping_file: Path):
-    with mapping_file.open("r", encoding="utf-8") as f:
-        mapping = json.load(f)
-    return mapping
-
-
-def get_encoder():
-    base_model_repo = "hf_hub:SmilingWolf/wd-eva02-large-tagger-v3"
-    encoder = timm.create_model(base_model_repo, pretrained=False)
-    encoder.reset_classifier(0)
-    return encoder
-
-
-def get_decoder():
-    decoder = TaggingHead(1024, 13461)
-    return decoder
-
-
-def get_model():
-    encoder = get_encoder()
-    decoder = get_decoder()
-    model = torch.nn.Sequential(encoder, decoder)
-    return model
-
-
-def load_model(weights_file, device):
-    model = get_model()
-    states_dict = torch.load(weights_file, map_location=device, weights_only=True)
-    model.load_state_dict(states_dict)
-    model.to(device)
-    model.eval()
-    return model
-
-
-def pure_pil_alpha_to_color_v2(
-    image: Image.Image, color: tuple[int, int, int] = (255, 255, 255)
-) -> Image.Image:
-    """
-    Convert a PIL image with an alpha channel to a RGB image.
-    This is a workaround for the fact that the model expects a RGB image, but the image may have an alpha channel.
-    This function will convert the image to a RGB image, and fill the alpha channel with the given color.
-    The alpha channel is the 4th channel of the image.
-    """
-    image.load()  # needed for split()
-    background = Image.new("RGB", image.size, color)
-    background.paste(image, mask=image.split()[3])  # 3 is the alpha channel
-    return background
-
-
-def pil_to_rgb(image: Image.Image) -> Image.Image:
-    if image.mode == "RGBA":
-        image = pure_pil_alpha_to_color_v2(image)
-    elif image.mode == "P":
-        image = pure_pil_alpha_to_color_v2(image.convert("RGBA"))
-    else:
-        image = image.convert("RGB")
-    return image
+from transformers import pipeline
 
 
 class EndpointHandler:
-    def __init__(self, path: str):
-        repo_path = Path(path)
-        assert repo_path.is_dir(), f"Model directory not found: {repo_path}"
-        weights_file = repo_path / "model_v0.9.pth"
-        tags_file = repo_path / "tags_v0.9_13k.json"
-        mapping_file = repo_path / "char_ip_map.json"
-        if not weights_file.exists():
-            raise FileNotFoundError(f"Model file not found: {weights_file}")
-        if not tags_file.exists():
-            raise FileNotFoundError(f"Tags file not found: {tags_file}")
-        if not mapping_file.exists():
-            raise FileNotFoundError(f"Mapping file not found: {mapping_file}")
+    """加载 PixAI Tagger v1.0，并兼容应用现有的标签返回字段。"""
 
-        # Robust device selection: prefer CPU unless CUDA is truly usable
-        force_cpu = os.environ.get("FORCE_CPU", "0") in {"1", "true", "TRUE", "yes", "on"}
+    def __init__(self, path: str):
+        """从已经下载完整的本地模型目录初始化官方推理流程。
+
+        Args:
+            path: 包含配置、预处理器、官方模型代码和 safetensors 权重的目录。
+        """
+        repo_path = Path(path).resolve()
+        for filename in (
+            "config.json",
+            "preprocessor_config.json",
+            "tagger_pipeline.py",
+            "model.safetensors",
+        ):
+            if not (repo_path / filename).is_file():
+                raise FileNotFoundError(f"PixAI Tagger v1.0 文件不存在：{repo_path / filename}")
+
+        force_cpu = os.environ.get("FORCE_CPU", "0").lower() in {"1", "true", "yes", "on"}
+        self.device = "cpu"
         if not force_cpu and torch.cuda.is_available():
             try:
-                # Probe that CUDA can actually be used (driver present)
                 torch.zeros(1).to("cuda")
                 self.device = "cuda"
-            except Exception:
-                self.device = "cpu"
-        else:
-            self.device = "cpu"
-        self.model = load_model(str(weights_file), self.device)
-        self.transform = transforms.Compose(
-            [
-                transforms.Resize((448, 448)),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-            ]
+            except (RuntimeError, AssertionError):
+                logging.warning("CUDA 初始化失败，改用 CPU 加载 PixAI Tagger v1.0")
+
+        # v1.0 的等比缩放、补边和透明背景处理必须使用随模型发布的处理器。
+        # 代码和权重均从同一个本地快照读取，避免初始化时混入其他版本。
+        self.tagger = pipeline(
+            task="tagger",
+            model=str(repo_path),
+            image_processor=str(repo_path),
+            trust_remote_code=True,
+            token=False,
+            model_kwargs={"local_files_only": True},
+            device=self.device,
         )
         self.fetch_image_timeout = 5.0
-        self.default_general_threshold = 0.3
-        self.default_character_threshold = 0.85
-
-        tag_map, self.gen_tag_count, self.character_tag_count = get_tags(tags_file)
-
-        # Invert the tag_map for efficient index-to-tag lookups
-        self.index_to_tag_map = {v: k for k, v in tag_map.items()}
-
-        self.character_ip_mapping = get_character_ip_mapping(mapping_file)
+        self.default_general_threshold = 0.17
+        self.default_character_threshold = 0.27
 
     def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
-        inputs = data.pop("inputs", data)
-
-        fetch_start_time = time.time()
+        """识别图片并返回六类标签、可选置信度和推理耗时。
+
+        Args:
+            data: inputs 接收 PIL 图片、url 或 base64 image；parameters 控制
+                通用与角色阈值、threshold/topk 模式、标签数量和置信度输出。
+
+        Returns:
+            feature 对应 general，ip 对应 copyright；character、style、meta、
+            rating 保留官方分类名称，另含 _params 和 _timings 元数据。
+        """
+        inputs = data.get("inputs", data)
+        fetch_start_time = time.perf_counter()
         if isinstance(inputs, Image.Image):
             image = inputs
-        elif image_url := inputs.pop("url", None):
-            with requests.get(
-                image_url, stream=True, timeout=self.fetch_image_timeout
-            ) as res:
+        elif image_url := inputs.get("url"):
+            with requests.get(image_url, stream=True, timeout=self.fetch_image_timeout) as res:
                 res.raise_for_status()
                 image = Image.open(res.raw)
-        elif image_base64_encoded := inputs.pop("image", None):
+                image.load()
+        elif image_base64_encoded := inputs.get("image"):
             image = Image.open(io.BytesIO(base64.b64decode(image_base64_encoded)))
+            image.load()
         else:
-            raise ValueError(f"No image or url provided: {data}")
-        # remove alpha channel if it exists
-        image = pil_to_rgb(image)
-        fetch_time = time.time() - fetch_start_time
-
-        parameters = data.pop("parameters", {})
-        general_threshold = parameters.pop(
-            "general_threshold", self.default_general_threshold
-        )
-        character_threshold = parameters.pop(
-            "character_threshold", self.default_character_threshold
-        )
-        # Optional behavior controls
-        mode = parameters.pop("mode", "threshold")  # "threshold" | "topk"
-        include_scores = bool(parameters.pop("include_scores", False))
-        topk_general = int(parameters.pop("topk_general", 25))
-        topk_character = int(parameters.pop("topk_character", 10))
-
-        inference_start_time = time.time()
-        with torch.inference_mode():
-            # Preprocess image on CPU
-            image_tensor = self.transform(image).unsqueeze(0)
-            # Pin memory and use non_blocking transfer only when using CUDA
-            if self.device == "cuda":
-                image_tensor = image_tensor.pin_memory().to(self.device, non_blocking=True)
-            else:
-                image_tensor = image_tensor.to(self.device)
-
-            # Run model on GPU
-            probs = self.model(image_tensor)[0]  # Get probs for the single image
-
-            if mode == "topk":
-                # Select top-k by category, independent of thresholds
-                gen_slice = probs[: self.gen_tag_count]
-                char_slice = probs[self.gen_tag_count :]
-                k_gen = max(0, min(int(topk_general), self.gen_tag_count))
-                k_char = max(0, min(int(topk_character), self.character_tag_count))
-                gen_scores, gen_idx = (torch.tensor([]), torch.tensor([], dtype=torch.long))
-                char_scores, char_idx = (torch.tensor([]), torch.tensor([], dtype=torch.long))
-                if k_gen > 0:
-                    gen_scores, gen_idx = torch.topk(gen_slice, k_gen)
-                if k_char > 0:
-                    char_scores, char_idx = torch.topk(char_slice, k_char)
-                    char_idx = char_idx + self.gen_tag_count
-
-                # Merge for unified post-processing
-                combined_indices = torch.cat((gen_idx, char_idx)).cpu()
-                combined_scores = torch.cat((gen_scores, char_scores)).cpu()
-            else:
-                # Perform thresholding directly on the GPU
-                general_mask = probs[: self.gen_tag_count] > general_threshold
-                character_mask = probs[self.gen_tag_count :] > character_threshold
-
-                # Get the indices of positive tags on the GPU
-                general_indices = general_mask.nonzero(as_tuple=True)[0]
-                character_indices = (
-                    character_mask.nonzero(as_tuple=True)[0] + self.gen_tag_count
-                )
-
-                # Combine indices and move the small result tensor to the CPU
-                combined_indices = torch.cat((general_indices, character_indices)).cpu()
-                combined_scores = probs[combined_indices].detach().float().cpu()
-
-        inference_time = time.time() - inference_start_time
-
-        post_process_start_time = time.time()
-
-        cur_gen_tags = []
-        cur_char_tags = []
-        gen_scores_out: dict[str, float] = {}
-        char_scores_out: dict[str, float] = {}
-
-        # Use the efficient pre-computed map for lookups
-        for pos, i in enumerate(combined_indices):
-            idx = int(i.item())
-            tag = self.index_to_tag_map[idx]
-            if idx < self.gen_tag_count:
-                cur_gen_tags.append(tag)
-                if include_scores:
-                    score = float(combined_scores[pos].item())
-                    gen_scores_out[tag] = score
-            else:
-                cur_char_tags.append(tag)
-                if include_scores:
-                    score = float(combined_scores[pos].item())
-                    char_scores_out[tag] = score
-
-        ip_tags = []
-        for tag in cur_char_tags:
-            if tag in self.character_ip_mapping:
-                ip_tags.extend(self.character_ip_mapping[tag])
-        ip_tags = sorted(set(ip_tags))
-        post_process_time = time.time() - post_process_start_time
-
-        logging.info(
-            f"Timing - Fetch: {fetch_time:.3f}s, Inference: {inference_time:.3f}s, Post-process: {post_process_time:.3f}s, Total: {fetch_time + inference_time + post_process_time:.3f}s"
-        )
-
-        out: dict[str, Any] = {
-            "feature": cur_gen_tags,
-            "character": cur_char_tags,
-            "ip": ip_tags,
-            "_timings": {
-                "fetch_s": round(fetch_time, 4),
-                "inference_s": round(inference_time, 4),
-                "post_process_s": round(post_process_time, 4),
-                "total_s": round(fetch_time + inference_time + post_process_time, 4),
-            },
-            "_params": {
-                "mode": mode,
-                "general_threshold": general_threshold,
-                "character_threshold": character_threshold,
-                "topk_general": topk_general,
-                "topk_character": topk_character,
-            },
+            raise ValueError("No image or url provided")
+        fetch_time = time.perf_counter() - fetch_start_time
+
+        parameters = data.get("parameters", {})
+        general_threshold = float(parameters.get("general_threshold", self.default_general_threshold))
+        character_threshold = float(parameters.get("character_threshold", self.default_character_threshold))
+        mode = parameters.get("mode", "threshold")
+        if mode not in {"threshold", "topk"}:
+            raise ValueError("mode 必须为 threshold 或 topk")
+        include_scores = bool(parameters.get("include_scores", False))
+        topk_general = max(0, int(parameters.get("topk_general", 25)))
+        topk_character = max(0, int(parameters.get("topk_character", 10)))
+        thresholds = {"general": general_threshold, "character": character_threshold}
+        if mode == "topk":
+            # 官方流程只有阈值筛选；先取齐这两类概率，再按各自数量截取。
+            thresholds = {"general": -1.0, "character": -1.0}
+
+        inference_start_time = time.perf_counter()
+        results = self.tagger(image, threshold=thresholds)["results"]
+        inference_time = time.perf_counter() - inference_start_time
+
+        post_process_start_time = time.perf_counter()
+        out: dict[str, Any] = {}
+        for source_category, output_category in (
+            ("general", "feature"),
+            ("character", "character"),
+            ("copyright", "ip"),
+            ("style", "style"),
+            ("meta", "meta"),
+            ("rating", "rating"),
+        ):
+            scores = sorted(results[source_category].items(), key=lambda item: item[1], reverse=True)
+            if mode == "topk" and source_category in {"general", "character"}:
+                limit = topk_general if source_category == "general" else topk_character
+                scores = scores[:limit]
+            out[output_category] = [tag for tag, _ in scores]
+            if include_scores:
+                out[f"{output_category}_scores"] = {tag: float(score) for tag, score in scores}
+        post_process_time = time.perf_counter() - post_process_start_time
+
+        out["_timings"] = {
+            "fetch_s": round(fetch_time, 4),
+            "inference_s": round(inference_time, 4),
+            "post_process_s": round(post_process_time, 4),
+            "total_s": round(fetch_time + inference_time + post_process_time, 4),
         }
-
-        if include_scores:
-            out["feature_scores"] = gen_scores_out
-            out["character_scores"] = char_scores_out
-
+        out["_params"] = {
+            "mode": mode,
+            "general_threshold": general_threshold,
+            "character_threshold": character_threshold,
+            "topk_general": topk_general,
+            "topk_character": topk_character,
+        }
+        logging.info("PixAI Tagger v1.0 耗时：%s", out["_timings"])
         return out

image_inversion_job_api.py

@@ -215,8 +215,8 @@ class ImageInversionJobRequest(BaseModel):
     )
 
     input_image_url: str = Field(min_length=1, max_length=4096)
-    general_threshold: float = Field(default=0.35, ge=0.0, le=1.0)
-    character_threshold: float = Field(default=0.85, ge=0.0, le=1.0)
+    general_threshold: float = Field(default=0.17, ge=0.0, le=1.0)
+    character_threshold: float = Field(default=0.27, ge=0.0, le=1.0)
     show_confidence: bool = True
     show_general: bool = True
     show_character: bool = True
@@ -453,7 +453,7 @@ def decode_image_bytes(
     content_type: str,
     settings: ImageInversionJobSettings,
 ) -> Image.Image:
-    """校验远程内容并解码为脱离底层流的 RGB 图片。
+    """校验远程内容并解码为脱离底层流的图片，保留透明通道。
 
     Args:
         image_bytes: 在字节上限内完整读取的图片内容。
@@ -461,7 +461,7 @@ def decode_image_bytes(
         settings: 图片字节、像素及单边尺寸限制。
 
     Returns:
-        完成 EXIF 方向修正并复制到内存的 RGB PIL 图片。
+        完成 EXIF 方向修正并复制到内存的 PIL 图片。
     """
     if not image_bytes or len(image_bytes) > settings.max_image_bytes:
         raise ImageSourceError(
@@ -496,7 +496,8 @@ def decode_image_bytes(
                         "The remote image exceeds the configured dimension limit.",
                     )
                 source_image.load()
-                return ImageOps.exif_transpose(source_image).convert("RGB").copy()
+                # 保留透明通道，交给模型官方预处理合成白底。
+                return ImageOps.exif_transpose(source_image).copy()
     except ImageSourceError:
         raise
     except (
@@ -547,7 +548,7 @@ def _fetch_remote_image_inner(
         settings: 图片白名单、超时、体积、像素与尺寸边界。
 
     Returns:
-        可直接交给标签模型的 RGB PIL 图片。
+        可直接交给标签模型且保留透明通道的 PIL 图片。
     """
     started_at = time.monotonic()
     validated_url = validate_image_url(url, settings)
@@ -679,7 +680,7 @@ def fetch_remote_image(
         settings: 白名单、硬超时、字节、像素与尺寸边界。
 
     Returns:
-        在总期限内完成验证、EXIF 转正和 RGB 转换的 PIL 图片。
+        在总期限内完成验证、EXIF 转正并保留透明通道的 PIL 图片。
     """
     if not IMAGE_FETCH_SLOTS.acquire(blocking=False):
         raise ImageSourceError(

requirements.txt

@@ -1,13 +1,14 @@
 gradio==5.29.0
 fastapi>=0.115.2,<1
-huggingface_hub>=0.24.0
+huggingface_hub>=0.34,<1.0
 numpy
 pandas
-Pillow
+Pillow>=10
 onnxruntime
-torch
-torchvision
-timm
+torch>=2.6,<3
+torchvision>=0.21,<1
+timm==1.0.27
+transformers==4.57.6
 requests
 httpx==0.28.1
 pydantic>=2,<2.12

tests/test_app_v1.py

@@ -0,0 +1,160 @@
+import ast
+import io
+import json
+import os
+import tempfile
+import time
+import unittest
+from html import escape
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+from PIL import Image
+
+from image_inversion_job_api import (
+    ImageInversionJobRequest,
+    ImageInversionJobSettings,
+    decode_image_bytes,
+)
+
+
+class AppV1IntegrationTests(unittest.TestCase):
+    def setUp(self):
+        """隔离加载应用定义，为集成测试准备六类模型结果。"""
+        # 只加载本次集成涉及的实际业务定义，避免测试时启动 Gradio 或下载权重。
+        app_path = Path(__file__).resolve().parents[1] / "app.py"
+        self.tree = ast.parse(app_path.read_text(encoding="utf-8"))
+        definition_names = {
+            "Tagger", "ImageValidationError", "_format_size", "validate_and_open_image",
+            "format_tags_html", "generate_summary_text_content", "execute_image_inversion_job",
+        }
+        constant_names = {
+            "MAX_UPLOAD_BYTES", "MAX_IMAGE_SIDE", "MAX_IMAGE_PIXELS",
+            "MAX_DECOMPRESSED_BYTES", "ALLOWED_IMAGE_FORMATS", "ASSETS_REPO_ID",
+            "ASSETS_REVISION",
+        }
+        definitions = [
+            node for node in ast.walk(self.tree)
+            if (
+                isinstance(node, (ast.ClassDef, ast.FunctionDef))
+                and node.name in definition_names
+            ) or (
+                isinstance(node, ast.Assign)
+                and any(isinstance(target, ast.Name) and target.id in constant_names for target in node.targets)
+            )
+        ]
+        self.scope = {
+            "Image": Image, "ImageInversionJobRequest": ImageInversionJobRequest,
+            "os": os, "time": time, "json": json, "escape": escape,
+        }
+        exec(compile(ast.Module(body=definitions, type_ignores=[]), str(app_path), "exec"), self.scope)
+        self.tagger = self.scope["Tagger"].__new__(self.scope["Tagger"])
+        self.tagger.device = "cpu"
+        self.handler_output = {
+            "feature_scores": {"blue_eyes": 0.71, "long_hair": 0.98},
+            "character_scores": {"test_character": 0.84},
+            "ip_scores": {"test_series": 0.76},
+            "style_scores": {"anime": 0.89},
+            "meta_scores": {"highres": 0.91},
+            "rating_scores": {"general": 0.99},
+            "_params": {"general_threshold": 0.17, "character_threshold": 0.27},
+            "_timings": {"inference_s": 0.125},
+        }
+        self.tagger.handler = Mock(return_value=self.handler_output)
+        self.scope["tagger_instance"] = self.tagger
+
+    def test_predict_preserves_all_categories_and_real_ip_confidence(self):
+        """验证完整类别元数据和真实 IP 置信度。"""
+        image = Image.new("RGBA", (32, 32), (10, 20, 30, 0))
+        result, translation_order, metadata = self.tagger.predict(image)
+
+        self.assertEqual(list(result), ["general", "characters", "ips"])
+        self.assertEqual(result["general"], {"long hair": 0.98, "blue eyes": 0.71})
+        self.assertEqual(translation_order["general"], ["long hair", "blue eyes"])
+        self.assertEqual(result["ips"], {"test series": 0.76})
+        self.assertEqual(metadata["tags"], {
+            "general": self.handler_output["feature_scores"],
+            "character": self.handler_output["character_scores"],
+            "copyright": self.handler_output["ip_scores"],
+            "style": self.handler_output["style_scores"],
+            "meta": self.handler_output["meta_scores"],
+            "rating": self.handler_output["rating_scores"],
+        })
+        call_data = self.tagger.handler.call_args.args[0]
+        self.assertIs(call_data["inputs"], image)
+        self.assertTrue(call_data["parameters"]["include_scores"])
+        self.assertEqual(call_data["parameters"]["general_threshold"], 0.17)
+        self.assertEqual(call_data["parameters"]["character_threshold"], 0.27)
+
+    def test_job_keeps_six_fields_and_translation_alignment(self):
+        """验证任务接口保持六个字段且翻译不串位。"""
+        self.scope["translate_texts"] = Mock(return_value=["长发", "蓝眼睛", "测试角色", "测试作品"])
+        payload = ImageInversionJobRequest(input_image_url="https://example.com/image.png")
+        result = self.scope["execute_image_inversion_job"](payload, Image.new("RGB", (32, 32)), "job-id")
+
+        self.assertEqual(set(result), {
+            "status_markdown", "general_tags_html", "character_tags_html",
+            "ip_tags_html", "summary_text", "metadata",
+        })
+        self.scope["translate_texts"].assert_called_once_with(
+            ["long hair", "blue eyes", "test character", "test series"],
+            src_lang="auto", tgt_lang="zh",
+        )
+        self.assertIn("长发", result["general_tags_html"])
+        self.assertIn("测试角色", result["character_tags_html"])
+        self.assertIn("测试作品", result["ip_tags_html"])
+        self.assertIn("0.760", result["ip_tags_html"])
+        self.assertIn("test series/*测试作品*/", result["summary_text"])
+        self.assertEqual(len(result["metadata"]["tags"]), 6)
+
+        # 整个翻译调用异常也不能影响标签和新增元数据的交付。
+        self.scope["translate_texts"].side_effect = RuntimeError("translation unavailable")
+        untranslated = self.scope["execute_image_inversion_job"](payload, Image.new("RGB", (32, 32)), "job-id")
+        self.assertEqual(set(untranslated), set(result))
+        self.assertIn("test series", untranslated["summary_text"])
+        self.assertNotIn("测试作品", untranslated["ip_tags_html"])
+        self.assertEqual(untranslated["metadata"]["tags"], result["metadata"]["tags"])
+
+    def test_ui_and_api_defaults_match_v1_thresholds(self):
+        """验证页面和任务接口的默认阈值与 v1.0 推荐值一致。"""
+        slider_nodes = [
+            node for node in ast.walk(self.tree)
+            if isinstance(node, ast.Assign)
+            and any(isinstance(target, ast.Name) and target.id in {"gen_slider", "char_slider"} for target in node.targets)
+        ]
+        self.scope["gr"] = SimpleNamespace(Slider=lambda *args, **kwargs: kwargs)
+        exec(compile(ast.Module(body=slider_nodes, type_ignores=[]), "app-sliders", "exec"), self.scope)
+        payload = ImageInversionJobRequest(input_image_url="https://example.com/image.png")
+
+        self.assertEqual(self.scope["gen_slider"]["value"], 0.17)
+        self.assertEqual(self.scope["char_slider"]["value"], 0.27)
+        self.assertEqual(payload.general_threshold, 0.17)
+        self.assertEqual(payload.character_threshold, 0.27)
+
+    def test_upload_and_job_decode_preserve_transparency(self):
+        """验证上传和远程解码均将透明背景留给官方处理器。"""
+        settings = ImageInversionJobSettings(api_key="x" * 32, allowed_hosts=frozenset({"example.com"}))
+        for mode in ("RGBA", "P"):
+            with self.subTest(mode=mode):
+                image = Image.new("RGBA", (32, 32), (10, 20, 30, 0))
+                if mode == "P":
+                    image = image.convert("P")
+                buffer = io.BytesIO()
+                image.save(buffer, format="PNG")
+                decoded = decode_image_bytes(buffer.getvalue(), "image/png", settings)
+                self.assertEqual(decoded.convert("RGBA").getpixel((0, 0))[3], 0)
+
+                # 仅创建并显式删除当前测试的单个文件，不创建需递归清理的目录。
+                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as source:
+                    source.write(buffer.getvalue())
+                    source_path = Path(source.name).resolve()
+                try:
+                    uploaded = self.scope["validate_and_open_image"](str(source_path))
+                    self.assertEqual(uploaded.convert("RGBA").getpixel((0, 0))[3], 0)
+                finally:
+                    source_path.unlink()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_handler.py

@@ -0,0 +1,158 @@
+"""不下载模型权重，验证 v1.0 推理适配层与官方 pipeline 的调用契约。"""
+
+import base64
+import importlib.util
+import io
+import os
+import sys
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock, patch
+
+from PIL import Image
+
+
+class EndpointHandlerTest(unittest.TestCase):
+    """覆盖标签映射、真实分数、预处理委托和本地加载边界。"""
+
+    @classmethod
+    def setUpClass(cls):
+        """隔离导入推理依赖，让适配层测试可在没有 GPU 和权重时运行。"""
+        spec = importlib.util.spec_from_file_location(
+            "pixai_handler_test", Path(__file__).resolve().parents[1] / "handler.py"
+        )
+        cls.module = importlib.util.module_from_spec(spec)
+        cls.torch = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=Mock(return_value=False)),
+            zeros=Mock(),
+        )
+        with patch.dict(
+            sys.modules,
+            {"torch": cls.torch, "transformers": SimpleNamespace(pipeline=Mock())},
+        ):
+            spec.loader.exec_module(cls.module)
+
+    def setUp(self):
+        """构造各类别独立的分数，避免将版权或风格误当成通用标签。"""
+        self.results = {
+            "general": {"blue_hair": 0.7, "1girl": 0.99, "hat": 0.2},
+            "character": {"character_b": 0.3, "character_a": 0.9},
+            "copyright": {"original": 0.64},
+            "style": {"style_a": 0.6},
+            "meta": {"highres": 0.8},
+            "rating": {"rating:g": 0.88},
+        }
+        self.pipe = Mock(return_value={"results": self.results})
+        self.factory = Mock(return_value=self.pipe)
+        self.torch.cuda.is_available.return_value = False
+        self.torch.zeros.reset_mock()
+        self.torch.zeros.return_value.to.side_effect = None
+        self.model_path = Path(__file__).resolve().parents[1]
+        with patch.object(Path, "is_file", return_value=True), patch.object(
+            self.module, "pipeline", self.factory
+        ):
+            self.handler = self.module.EndpointHandler(str(self.model_path))
+
+    def test_local_snapshot_and_official_processor(self):
+        """本地四文件快照与官方自定义任务共同决定模型和预处理版本。"""
+        self.factory.assert_called_once_with(
+            task="tagger",
+            model=str(self.model_path),
+            image_processor=str(self.model_path),
+            trust_remote_code=True,
+            token=False,
+            model_kwargs={"local_files_only": True},
+            device="cpu",
+        )
+        image = Image.new("RGBA", (80, 40), (255, 0, 0, 0))
+        data = {"inputs": image, "parameters": {"include_scores": True}}
+        result = self.handler(data)
+        self.pipe.assert_called_once_with(
+            image, threshold={"general": 0.17, "character": 0.27}
+        )
+        self.assertEqual(image.mode, "RGBA")
+        self.assertEqual(image.size, (80, 40))
+        self.assertIs(data["inputs"], image)
+        self.assertEqual(data["parameters"], {"include_scores": True})
+        self.assertEqual(result["feature"], ["1girl", "blue_hair", "hat"])
+        self.assertEqual(result["feature_scores"], self.results["general"])
+        for category in ("character", "style", "meta", "rating"):
+            self.assertEqual(result[f"{category}_scores"], self.results[category])
+        self.assertEqual(result["ip"], ["original"])
+        self.assertEqual(result["ip_scores"], {"original": 0.64})
+        self.assertEqual(result["_params"]["general_threshold"], 0.17)
+        self.assertEqual(set(result["_timings"]), {"fetch_s", "inference_s", "post_process_s", "total_s"})
+
+    def test_custom_thresholds_and_optional_scores(self):
+        """阈值交由官方流程筛选，关闭分数时不泄漏分数字段。"""
+        image = Image.new("RGB", (32, 32))
+        result = self.handler(
+            {"inputs": image, "parameters": {"general_threshold": 0.4, "character_threshold": 0.6}}
+        )
+        self.pipe.assert_called_once_with(image, threshold={"general": 0.4, "character": 0.6})
+        self.assertFalse(any(name.endswith("_scores") for name in result))
+
+    def test_topk_ranks_categories_independently(self):
+        """Top-K 不先经过默认阈值，且保留版权等其他分类的原始分数。"""
+        image = Image.new("RGB", (32, 32))
+        result = self.handler(
+            {
+                "inputs": image,
+                "parameters": {"mode": "topk", "topk_general": 2, "topk_character": 1, "include_scores": True},
+            }
+        )
+        self.pipe.assert_called_once_with(image, threshold={"general": -1.0, "character": -1.0})
+        self.assertEqual(result["feature"], ["1girl", "blue_hair"])
+        self.assertEqual(result["character"], ["character_a"])
+        self.assertEqual(result["character_scores"], {"character_a": 0.9})
+        self.assertEqual(result["ip_scores"], {"original": 0.64})
+        zero_result = self.handler(
+            {"inputs": image, "parameters": {"mode": "topk", "topk_general": 0, "topk_character": 0}}
+        )
+        self.assertEqual(zero_result["feature"], [])
+        self.assertEqual(zero_result["character"], [])
+
+    def test_base64_and_remote_images_are_decoded_before_inference(self):
+        """URL 图片在响应关闭前解码，base64 和透明图片也交由官方处理器。"""
+        image_bytes = io.BytesIO()
+        Image.new("RGBA", (80, 40), (0, 255, 0, 0)).save(image_bytes, format="PNG")
+        image_content = image_bytes.getvalue()
+        self.handler({"inputs": {"image": base64.b64encode(image_content).decode()}})
+        decoded_image = self.pipe.call_args.args[0]
+        self.assertEqual((decoded_image.mode, decoded_image.size), ("RGBA", (80, 40)))
+
+        response = MagicMock()
+        response.raw = io.BytesIO(image_content)
+        response.__enter__.return_value = response
+        response.__exit__.side_effect = lambda *_: response.raw.close()
+        with patch.object(self.module.requests, "get", return_value=response) as get:
+            self.handler({"inputs": {"url": "https://example.org/image.png"}})
+        get.assert_called_once_with("https://example.org/image.png", stream=True, timeout=5.0)
+        response.raise_for_status.assert_called_once_with()
+        self.assertTrue(response.raw.closed)
+        self.assertEqual(self.pipe.call_args.args[0].getpixel((0, 0)), (0, 255, 0, 0))
+
+    def test_incomplete_snapshot_fails_before_model_loading(self):
+        """缺失新版本预处理器时直接报错，避免加载旧版或远程补齐混用。"""
+        with patch.object(
+            Path, "is_file", autospec=True, side_effect=lambda path: path.name != "preprocessor_config.json"
+        ), patch.object(self.module, "pipeline") as factory:
+            with self.assertRaisesRegex(FileNotFoundError, "preprocessor_config.json"):
+                self.module.EndpointHandler(str(self.model_path))
+        factory.assert_not_called()
+
+    def test_cuda_probe_failure_uses_cpu(self):
+        """存在 CUDA 接口但驱动不能工作时，仍能使用 CPU 加载。"""
+        self.torch.cuda.is_available.return_value = True
+        self.torch.zeros.return_value.to.side_effect = RuntimeError("driver unavailable")
+        with patch.dict(os.environ, {"FORCE_CPU": "0"}), patch.object(
+            Path, "is_file", return_value=True
+        ), patch.object(self.module, "pipeline", self.factory):
+            handler = self.module.EndpointHandler(str(self.model_path))
+        self.assertEqual(handler.device, "cpu")
+        self.assertEqual(self.factory.call_args.kwargs["device"], "cpu")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_translator.py

@@ -0,0 +1,145 @@
+import json
+import unittest
+from unittest.mock import Mock, patch
+
+import requests
+
+import translator
+
+
+class TranslatorTests(unittest.TestCase):
+    def test_index_preserves_order_and_does_not_call_legacy_providers(self):
+        response = Mock()
+        response.json.return_value = {
+            "choices": [{
+                "finish_reason": "stop",
+                "message": {"content": '{"2":"蓝眼睛","0":"长发","1":"长发"}'},
+            }]
+        }
+        with (
+            patch.object(translator.requests, "post", return_value=response) as post,
+            patch.object(translator, "INDEX_TRANSLATE_API_KEY", ""),
+            patch.object(translator, "_translate_with_tencent_pool") as tencent,
+            patch.object(translator, "_translate_with_baidu_pool") as baidu,
+        ):
+            result = translator.translate_texts(["long_hair", "long_hair", "blue_eyes"])
+
+        self.assertEqual(result, ["长发", "长发", "蓝眼睛"])
+        tencent.assert_not_called()
+        baidu.assert_not_called()
+        self.assertEqual(post.call_args.args, (translator.INDEX_TRANSLATE_URL,))
+        payload = post.call_args.kwargs["json"]
+        self.assertEqual(payload["model"], translator.INDEX_TRANSLATE_MODEL)
+        self.assertEqual(payload["chat_template_kwargs"], {"enable_thinking": False})
+        self.assertEqual(payload["temperature"], 0)
+        self.assertEqual(
+            post.call_args.kwargs["headers"],
+            {"User-Agent": "Index-Translate-Client/1.0"},
+        )
+
+    def test_index_batches_by_item_count_and_text_length(self):
+        for texts, batch_sizes in [
+            ([f"tag_{index}" for index in range(33)], [32, 1]),
+            (["a" * 1000, "b" * 1000, "c" * 1000], [2, 1]),
+        ]:
+            with self.subTest(batch_sizes=batch_sizes):
+                responses = []
+                expected = []
+                for batch_number, size in enumerate(batch_sizes):
+                    translations = {str(index): f"译文{batch_number}-{index}" for index in range(size)}
+                    expected.extend(translations.values())
+                    response = Mock()
+                    response.json.return_value = {
+                        "choices": [{
+                            "finish_reason": "stop",
+                            "message": {"content": json.dumps(translations)},
+                        }]
+                    }
+                    responses.append(response)
+                with (
+                    patch.object(translator.requests, "post", side_effect=responses) as post,
+                    patch.object(translator, "INDEX_TRANSLATE_API_KEY", "test-key"),
+                ):
+                    self.assertEqual(translator.translate_texts(texts), expected)
+                self.assertEqual(post.call_count, 2)
+                self.assertEqual(post.call_args.kwargs["headers"]["Authorization"], "Bearer test-key")
+                sources = [
+                    json.loads(call.kwargs["json"]["messages"][0]["content"].split("〖源文〗\n")[1].split("\n\n")[0])
+                    for call in post.call_args_list
+                ]
+                self.assertEqual([len(source) for source in sources], batch_sizes)
+                self.assertEqual([value for source in sources for value in source.values()], texts)
+
+    def test_invalid_index_responses_fall_back_without_misaligning_tags(self):
+        cases = [
+            ("not json", "stop"),
+            ('{"0":"长发"}', "stop"),
+            ('{"0":"长发","0":"蓝眼睛"}', "stop"),
+            ('{"0":"长发","1":"蓝眼睛","1":"眼睛"}', "stop"),
+            ('{"0":"长发","2":"蓝眼睛"}', "stop"),
+            ('{"0":"长发","1":""}', "stop"),
+            ('{"0":"长发","1":null}', "stop"),
+            ('{"0":"长发","1":"蓝眼睛"}', "length"),
+            (None, "stop"),
+        ]
+        for content, finish_reason in cases:
+            with self.subTest(content=content, finish_reason=finish_reason):
+                response = Mock()
+                response.json.return_value = {
+                    "choices": [{"finish_reason": finish_reason, "message": {"content": content}}]
+                }
+                with (
+                    patch.object(translator.requests, "post", return_value=response),
+                    patch.object(translator, "_translate_with_tencent_pool", return_value=["腾讯长发", "腾讯蓝眼睛"]) as tencent,
+                    patch.object(translator, "_translate_with_baidu_pool") as baidu,
+                ):
+                    result = translator.translate_texts(["long_hair", "blue_eyes"], "en", "zh")
+                self.assertEqual(result, ["腾讯长发", "腾讯蓝眼睛"])
+                tencent.assert_called_once_with(["long_hair", "blue_eyes"], "en", "zh")
+                baidu.assert_not_called()
+
+    def test_index_http_and_transport_errors_use_existing_fallback_chain(self):
+        for error in [requests.Timeout(), requests.HTTPError("429")]:
+            with self.subTest(error=type(error).__name__):
+                with (
+                    patch.object(translator.requests, "post", side_effect=error),
+                    patch.object(translator, "_translate_with_tencent_pool", return_value=None) as tencent,
+                    patch.object(translator, "_translate_with_baidu_pool", return_value=["百度长发"]) as baidu,
+                ):
+                    self.assertEqual(translator.translate_texts(["long_hair"]), ["百度长发"])
+                tencent.assert_called_once()
+                baidu.assert_called_once()
+
+    def test_later_index_batch_failure_retranslates_all_inputs_in_order(self):
+        texts = [f"tag_{index}" for index in range(33)]
+        response = Mock()
+        response.json.return_value = {
+            "choices": [{
+                "finish_reason": "stop",
+                "message": {"content": json.dumps({str(index): "Index译文" for index in range(32)})},
+            }]
+        }
+        expected = [f"腾讯译文{index}" for index in range(33)]
+        with (
+            patch.object(translator.requests, "post", side_effect=[response, requests.Timeout()]),
+            patch.object(translator, "_translate_with_tencent_pool", return_value=expected) as tencent,
+        ):
+            self.assertEqual(translator.translate_texts(texts), expected)
+        tencent.assert_called_once_with(texts, "auto", "zh")
+
+    def test_incomplete_legacy_results_return_originals_and_empty_inputs_do_not_call_api(self):
+        for result in (None, [], ["缺少一项"]):
+            with self.subTest(result=result):
+                with (
+                    patch.object(translator.requests, "post", side_effect=requests.Timeout()),
+                    patch.object(translator, "_translate_with_tencent_pool", return_value=result),
+                    patch.object(translator, "_translate_with_baidu_pool", return_value=result),
+                ):
+                    self.assertEqual(translator.translate_texts(["long_hair", "blue_eyes"]), ["long_hair", "blue_eyes"])
+        with patch.object(translator.requests, "post") as post:
+            self.assertEqual(translator.translate_texts([]), [])
+        post.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

translator.py

@@ -1,7 +1,12 @@
 """
 translator.py
-腾讯云 (批量接口+动态分包+多账号轮询) + 百度翻译 API 封装
-⚠️ 需在 HF 空间的 “Variables” 页设置以下环境变量
+Index-Translate 优先，失败后依次回退腾讯云和百度翻译。
+Index 官方公开接口无需密钥；可在 HF 空间设置以下可选环境变量：
+INDEX_TRANSLATE_URL      完整的 OpenAI Chat Completions 接口地址
+INDEX_TRANSLATE_MODEL    默认 Index-Translate-35B-A3B
+INDEX_TRANSLATE_API_KEY  自部署接口需要鉴权时配置
+INDEX_TRANSLATE_TIMEOUT  单次请求超时秒数，默认 30
+旧翻译接口继续使用以下环境变量：
 ------------------------------------------------------------------
 TENCENT_CREDENTIALS_JSON  形如:
 [
@@ -23,6 +28,13 @@ from typing import List, Sequence, Optional, Dict, Any
 
 import requests
 
+INDEX_TRANSLATE_URL = os.environ.get(
+    "INDEX_TRANSLATE_URL", "https://index-translate.bilibili.com/v1/chat/completions"
+)
+INDEX_TRANSLATE_MODEL = os.environ.get("INDEX_TRANSLATE_MODEL", "Index-Translate-35B-A3B")
+INDEX_TRANSLATE_API_KEY = os.environ.get("INDEX_TRANSLATE_API_KEY", "")
+INDEX_TRANSLATE_TIMEOUT = float(os.environ.get("INDEX_TRANSLATE_TIMEOUT", "30"))
+
 _tencent_creds_list = json.loads(os.environ.get("TENCENT_CREDENTIALS_JSON", "[]"))
 
 # 兼容旧的单账号配置
@@ -44,6 +56,89 @@ _baidu_creds_list = json.loads(os.environ.get("BAIDU_CREDENTIALS_JSON", "[]"))
 _baidu_idx: int = 0
 
 
+def _translate_with_index(texts: Sequence[str], src="auto", tgt="zh") -> Optional[List[str]]:
+    """通过 Index-Translate 分批翻译，响应不完整时交给旧接口处理。
+
+    Args:
+        texts: 按顺序排列的待翻译标签。
+        src: 源语言代码，auto 表示自动识别。
+        tgt: 目标语言代码或语言名称。
+
+    Returns:
+        与输入逐项对应的译文；请求或响应校验失败时返回 None。
+    """
+    languages = {"zh": "中文", "en": "英语", "ja": "日语", "ko": "韩语"}
+    source_language = "" if src == "auto" else languages.get(src, src)
+    target_language = languages.get(tgt, tgt)
+    headers = {"User-Agent": "Index-Translate-Client/1.0"}
+    if INDEX_TRANSLATE_API_KEY:
+        headers["Authorization"] = f"Bearer {INDEX_TRANSLATE_API_KEY}"
+
+    # 控制生成长度，避免大量标签挤满输出预算而丢失后半段译文。
+    chunks = []
+    chunk = []
+    char_count = 0
+    for text in texts:
+        if chunk and (len(chunk) >= 32 or char_count + len(text) > 2000):
+            chunks.append(chunk)
+            chunk = []
+            char_count = 0
+        chunk.append(text)
+        char_count += len(text)
+    if chunk:
+        chunks.append(chunk)
+
+    results = []
+    try:
+        for chunk in chunks:
+            # 用编号保持标签身份，模型调整 JSON 键顺序也不会造成译文错位。
+            source = {str(index): text for index, text in enumerate(chunk)}
+            prompt = (
+                f"请将以下{source_language}图像标签翻译成{target_language}，并且严格遵循所有约束要求。\n\n"
+                f"〖源文〗\n{json.dumps(source, ensure_ascii=False)}\n\n"
+                "〖约束要求〗\n"
+                "1. 〖硬性要求〗 保留 JSON 对象的所有 key，不得增删、合并或调整 key 对应的条目，仅翻译各 value。\n"
+                "2. 〖注意〗 每个 value 是独立的图像标签，按动漫图像标签语境翻译。\n\n"
+                "请以相同的 JSON 格式输出翻译结果，key 保持不变，value 为对应译文。只输出 JSON，不要有任何额外说明。"
+            )
+            response = requests.post(
+                INDEX_TRANSLATE_URL,
+                headers=headers,
+                json={
+                    "model": INDEX_TRANSLATE_MODEL,
+                    "messages": [{"role": "user", "content": prompt}],
+                    "temperature": 0,
+                    "max_tokens": 2048,
+                    "stream": False,
+                    "chat_template_kwargs": {"enable_thinking": False},
+                },
+                timeout=INDEX_TRANSLATE_TIMEOUT,
+            )
+            response.raise_for_status()
+            choice = response.json()["choices"][0]
+            if choice["finish_reason"] != "stop":
+                raise ValueError("Index-Translate 未完整生成译文")
+            content = choice["message"]["content"]
+            if not isinstance(content, str) or not content.lstrip().startswith("{"):
+                raise ValueError("Index-Translate 未返回 JSON 对象")
+
+            # 保留原始键值对才能发现重复编号，防止 json.loads 静默覆盖错误结果。
+            pairs = json.loads(content, object_pairs_hook=list)
+            translated = dict(pairs)
+            if len(pairs) != len(source) or set(translated) != set(source):
+                raise ValueError("Index-Translate 返回的标签编号不完整或重复")
+            for key, text in source.items():
+                value = translated[key]
+                if not isinstance(value, str) or (text.strip() and not value.strip()):
+                    raise ValueError("Index-Translate 返回了空译文或非文本译文")
+                results.append(value.strip())
+    except (requests.RequestException, ValueError, KeyError, IndexError, TypeError) as error:
+        # 不输出请求正文或鉴权信息；失败时沿用已有翻译服务。
+        print(f"[translator] Index-Translate failed ({type(error).__name__}); falling back to Tencent/Baidu.")
+        return None
+    return results
+
+
 # 腾讯云翻译逻辑 (批量接口 TextTranslateBatch)
 def _sign(key: bytes, msg: str) -> bytes:
     return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()
@@ -227,20 +322,27 @@ def _translate_with_baidu_pool(texts: Sequence[str], src="auto", tgt="zh") -> Op
 def translate_texts(texts: Sequence[str],
                     src_lang: str = "auto",
                     tgt_lang: str = "zh") -> List[str]:
-    """
-    逻辑：
-    1. 尝试腾讯云 (批量接口 + 多账号轮询)
-    2. 失败降级到百度云
-    3. 还失败返回原文
+    """优先使用 Index-Translate，失败后依次尝试腾讯云、百度和原文。
+
+    Args:
+        texts: 按顺序排列的待翻译标签。
+        src_lang: 源语言代码，auto 表示自动识别。
+        tgt_lang: 目标语言代码。
+
+    Returns:
+        与输入等长且逐项对应的译文，所有翻译接口失败时返回原文。
     """
     if not texts:
         return []
 
-    # 1. 优先尝试腾讯云
-    out = _translate_with_tencent_pool(texts, src_lang, tgt_lang)
+    out = _translate_with_index(texts, src_lang, tgt_lang)
+
+    if out is None:
+        out = _translate_with_tencent_pool(texts, src_lang, tgt_lang)
     
     # 2. 失败降级到百度
-    if out is None:
+    if out is None or len(out) != len(texts):
         out = _translate_with_baidu_pool(texts, src_lang, tgt_lang)
-        
-    return out or list(texts)
+
+    # 旧接口也不能把缺项结果传给上层，否则标签与译文会逐项错位。
+    return out if out is not None and len(out) == len(texts) else list(texts)