Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import onnxruntime as rt
 import pandas as pd
 from PIL import Image
+from huggingface_hub import login
 
 # 模型配置
 MODEL_REPO = "SmilingWolf/wd-swinv2-tagger-v3"  # 默认模型
@@ -12,6 +13,11 @@ MODEL_FILENAME = "model.onnx"
 LABEL_FILENAME = "selected_tags.csv"
 HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
+if not os.environ.get("HF_TOKEN"):
+    print("⚠️ 警告：未检测到HF_TOKEN，部分模型可能需要认证")
+else:
+    login(token=os.environ.get("HF_TOKEN"))
+
 # 标签处理配置
 kaomojis = [
     "0_0",
@@ -40,34 +46,45 @@ class Tagger:
         self.model = None
         self.tag_names = []
         self.model_size = None
+        self.hf_token = os.environ.get("HF_TOKEN", "")  # 从环境变量获取
         self._init_model()
 
     def _init_model(self):
         """初始化模型和标签"""
-        # 下载模型文件
-        label_path = huggingface_hub.hf_hub_download(
-            MODEL_REPO, 
-            LABEL_FILENAME,
-            token=HF_TOKEN
-        )
-        model_path = huggingface_hub.hf_hub_download(
-            MODEL_REPO,
-            MODEL_FILENAME,
-            token=HF_TOKEN
-        )
-        
-        # 加载标签
-        tags_df = pd.read_csv(label_path)
-        self.tag_names = tags_df["name"].tolist()
-        self.categories = {
-            "rating": np.where(tags_df["category"] == 9)[0],
-            "general": np.where(tags_df["category"] == 0)[0],
-            "character": np.where(tags_df["category"] == 4)[0]
-        }
-        
-        # 加载ONNX模型
-        self.model = rt.InferenceSession(model_path)
-        self.model_size = self.model.get_inputs()[0].shape[1]
+        try:
+            label_path = huggingface_hub.hf_hub_download(
+                MODEL_REPO, 
+                LABEL_FILENAME,
+                token=self.hf_token
+            )
+            model_path = huggingface_hub.hf_hub_download(
+                MODEL_REPO,
+                MODEL_FILENAME,
+                token=self.hf_token
+            )
+            
+            # 加载标签
+            tags_df = pd.read_csv(label_path)
+            self.tag_names = tags_df["name"].tolist()
+            self.categories = {
+                "rating": np.where(tags_df["category"] == 9)[0],
+                "general": np.where(tags_df["category"] == 0)[0],
+                "character": np.where(tags_df["category"] == 4)[0]
+            }
+            
+            # 加载ONNX模型
+            self.model = rt.InferenceSession(model_path)
+            self.model_size = self.model.get_inputs()[0].shape[1]
+        except huggingface_hub.utils.HfHubHTTPError as e:
+            if "401" in str(e):
+                raise RuntimeError(
+                    "模型下载认证失败，请：\n"
+                    "1. 访问https://huggingface.co/SmilingWolf/wd-swinv2-tagger-v3\n"
+                    "2. 点击Agree and continue\n"
+                    "3. 确保HF_TOKEN已正确设置"
+                )
+            else:
+                raise
 
     def _preprocess(self, img):
         """图像预处理"""