Update app.py

app.py

@@ -10,12 +10,18 @@ from huggingface_hub import whoami, get_space_runtime
 
 from translator import translate_texts
 
+# ------------------------------------------------------------------
+# 模型配置
+# ------------------------------------------------------------------
 MODEL_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
 MODEL_FILENAME = "model.onnx"
 LABEL_FILENAME = "selected_tags.csv"
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
+# ------------------------------------------------------------------
+# Tagger 类 (全局实例化)
+# ------------------------------------------------------------------
 class Tagger:
     def __init__(self):
         self.hf_token = HF_TOKEN
@@ -48,6 +54,8 @@ class Tagger:
             print(f"❌ 模型或标签加载失败: {e}")
             raise RuntimeError(f"模型初始化失败: {e}")
 
+
+    # ------------------------- preprocess -------------------------
     def _preprocess(self, img: Image.Image) -> np.ndarray:
         if img is None: raise ValueError("输入图像不能为空")
         if img.mode != "RGB": img = img.convert("RGB")
@@ -58,6 +66,7 @@ class Tagger:
             canvas = canvas.resize((self.input_size, self.input_size), Image.BICUBIC)
         return np.array(canvas)[:, :, ::-1].astype(np.float32)
 
+    # --------------------------- predict --------------------------
     def predict(self, img: Image.Image, gen_th: float = 0.35, char_th: float = 0.85):
         if self.model is None: raise RuntimeError("模型未成功加载，无法进行预测。")
         inp_name = self.model.get_inputs()[0].name
@@ -86,12 +95,16 @@ class Tagger:
 
         return res, tag_categories_for_translation
 
+# 全局 Tagger 实例
 try:
     tagger_instance = Tagger()
 except RuntimeError as e:
     print(f"应用启动时Tagger初始化失败: {e}")
     tagger_instance = None
 
+# ------------------------------------------------------------------
+# Gradio UI
+# ------------------------------------------------------------------
 custom_css = """
 .label-container { max-height: 300px; overflow-y: auto; border: 1px solid #ddd; padding: 10px; border-radius: 5px; background-color: #f9f9f9; }
 .tag-item { display: flex; justify-content: space-between; align-items: center; margin: 2px 0; padding: 2px 5px; border-radius: 3px; background-color: #fff; transition: background-color 0.2s; }
@@ -171,8 +184,15 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
             gr.Markdown("### 标签汇总结果")
             out_summary = gr.Textbox(label="标签汇总", lines=5, show_copy_button=True)
             
+    def get_token_from_request(request: gr.Request) -> str | None:
+        """Helper function to extract token from request headers."""
+        auth_header = request.headers.get("authorization")
+        if auth_header and auth_header.startswith("Bearer "):
+            return auth_header.split(" ")[1]
+        return None
+
     def check_user_status(request: gr.Request):
-        token = request.token
+        token = get_token_from_request(request)
         if token:
             try:
                 user_info = whoami(token=token)
@@ -225,7 +245,8 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
 
         yield gr.update(interactive=False, value="🔄 处理中..."), gr.update(visible=True, value="🔄 正在分析..."), *["<p>分析中...</p>"]*3, "分析中...", {}, {}
 
-        token, is_owner = request.token, False
+        token = get_token_from_request(request)
+        is_owner = False
         if token:
             try:
                 user_info = whoami(token=token)
@@ -300,4 +321,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
 if __name__ == "__main__":
     if tagger_instance is None:
         print("CRITICAL: Tagger 未能初始化，应用功能将受限。")
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)