Update app.py

app.py

@@ -6,22 +6,25 @@ import numpy as np
 import onnxruntime as rt
 import pandas as pd
 from PIL import Image
-from huggingface_hub import whoami, get_space_runtime, HfApi
-from huggingface_hub.hf_api import SpaceRuntime
-
+from huggingface_hub import whoami, HfApi
 from translator import translate_texts
 
 # ------------------------------------------------------------------
-# 模型配置
+# Model Configuration
 # ------------------------------------------------------------------
 MODEL_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
 MODEL_FILENAME = "model.onnx"
 LABEL_FILENAME = "selected_tags.csv"
 
+# It's recommended to manage the token within the HF Spaces secrets
 HF_TOKEN = os.environ.get("HF_TOKEN")
+# A more robust way to get the space owner
+SPACE_ID = os.environ.get("SPACE_ID")
+SPACE_OWNER = SPACE_ID.split('/')[0] if SPACE_ID else None
+
 
 # ------------------------------------------------------------------
-# Tagger 类 (全局实例化)
+# Tagger Class (Global Instance)
 # ------------------------------------------------------------------
 class Tagger:
     def __init__(self):
@@ -50,15 +53,14 @@ class Tagger:
             }
             self.model = rt.InferenceSession(model_path)
             self.input_size = self.model.get_inputs()[0].shape[1]
-            print("✅ 模型和标签加载成功")
+            print("✅ Model and labels loaded successfully.")
         except Exception as e:
-            print(f"❌ 模型或标签加载失败: {e}")
-            raise RuntimeError(f"模型初始化失败: {e}")
-
+            print(f"❌ Failed to load model or labels: {e}")
+            raise RuntimeError(f"Model initialization failed: {e}")
 
     # ------------------------- preprocess -------------------------
     def _preprocess(self, img: Image.Image) -> np.ndarray:
-        if img is None: raise ValueError("输入图像不能为空")
+        if img is None: raise ValueError("Input image cannot be None.")
         if img.mode != "RGB": img = img.convert("RGB")
         size = max(img.size)
         canvas = Image.new("RGB", (size, size), (255, 255, 255))
@@ -69,7 +71,7 @@ class Tagger:
 
     # --------------------------- predict --------------------------
     def predict(self, img: Image.Image, gen_th: float = 0.35, char_th: float = 0.85):
-        if self.model is None: raise RuntimeError("模型未成功加载，无法进行预测。")
+        if self.model is None: raise RuntimeError("Model not loaded, cannot predict.")
         inp_name = self.model.get_inputs()[0].name
         outputs = self.model.run(None, {inp_name: self._preprocess(img)[None, ...]})[0][0]
 
@@ -89,18 +91,17 @@ class Tagger:
                         tag_name = self.tag_names[idx].replace("_", " ")
                         sub_res[tag_name] = float(outputs[idx])
 
-            # Use the correct key for 'character'
             res_key = "characters" if cat_key == "character" else cat_key
             res[res_key] = dict(sorted(sub_res.items(), key=lambda kv: kv[1], reverse=True))
             tag_categories_for_translation[res_key] = list(res[res_key].keys())
 
         return res, tag_categories_for_translation
 
-# 全局 Tagger 实例
+# Global Tagger instance
 try:
     tagger_instance = Tagger()
 except RuntimeError as e:
-    print(f"应用启动时Tagger初始化失败: {e}")
+    print(f"Tagger initialization failed on app startup: {e}")
     tagger_instance = None
 
 # ------------------------------------------------------------------
@@ -145,12 +146,12 @@ function copyToClipboard(text) {
 with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=custom_css, js=_js_functions) as demo:
     gr.Markdown("# 🖼️ AI 图像标签分析器")
     gr.Markdown("上传图片自动识别标签，支持中英文显示和一键复制。[NovelAI在线绘画](https://nai.idlecloud.cc/)")
-    
+
     with gr.Row():
         with gr.Column(scale=1):
             login_button = gr.LoginButton(value="🤗 通过 Hugging Face 登录")
             user_status_md = gr.Markdown("ℹ️ 正在检查登录状态...")
-    
+
     state_res = gr.State({})
     state_translations_dict = gr.State({})
 
@@ -158,7 +159,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
         with gr.Column(scale=1):
             img_in = gr.Image(type="pil", label="上传图片", height=300)
             btn = gr.Button("🚀 开始分析", variant="primary", elem_classes=["btn-analyze-container"])
-            
+
             with gr.Accordion("⚙️ 高级设置", open=False):
                 gen_slider = gr.Slider(0, 1, value=0.35, step=0.01, label="通用标签阈值")
                 char_slider = gr.Slider(0, 1, value=0.85, step=0.01, label="角色标签阈值")
@@ -191,30 +192,35 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
             return auth_header.split(" ")[1]
         return None
 
-    def check_user_is_owner(user_info: dict | None, space_runtime: SpaceRuntime | None) -> bool:
-        if not user_info or not space_runtime:
+    def is_user_space_owner(user_info: dict | None) -> bool:
+        """
+        Robustly checks if the user is the owner of the space by parsing SPACE_ID.
+        """
+        if not user_info or not SPACE_OWNER:
+            if not SPACE_OWNER:
+                print("⚠️ Warning: SPACE_ID environment variable not found.")
             return False
-        if user_info.get("name") == space_runtime.owner:
-            return True
-        user_orgs = user_info.get("orgs", [])
-        if any(org.get("name") == space_runtime.owner for org in user_orgs):
-            return True
-            
-        return False
+        
+        user_name = user_info.get("name")
+        user_orgs = [org.get("name") for org in user_info.get("orgs", [])]
+
+        print(f"ℹ️ [Auth Check] Space Owner: '{SPACE_OWNER}', User: '{user_name}', User Orgs: {user_orgs}")
+
+        is_owner = (user_name == SPACE_OWNER) or (SPACE_OWNER in user_orgs)
+        return is_owner
 
     def check_user_status(request: gr.Request):
         token = get_token_from_request(request)
         if token:
             try:
                 user_info = whoami(token=token)
-                space_runtime = get_space_runtime()
                 
-                if check_user_is_owner(user_info, space_runtime):
+                if is_user_space_owner(user_info):
                     return f"✅ 以所有者 **{user_info.get('fullname', user_info.get('name'))}** 身份登录，将使用空间配置的密钥。", gr.update(visible=False)
                 else:
                     return f"👋 你好, **{user_info.get('fullname', '用户')}**！请在下方提供你自己的翻译 API 密钥。", gr.update(visible=True, open=True)
             except Exception as e:
-                print(f"获取用户信息时出错: {e}")
+                print(f"Error getting user info: {e}")
                 return "⚠️ 无法验证您的登录状态。请提供 API 密钥。", gr.update(visible=True, open=True)
         return "ℹ️ **访客模式**。如需使用翻译功能，请<a href='/login?redirect=/'>登录</a>或提供 API 密钥。", gr.update(visible=True, open=True)
         
@@ -262,8 +268,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
         if token:
             try:
                 user_info = whoami(token=token)
-                space_runtime = get_space_runtime()
-                if check_user_is_owner(user_info, space_runtime):
+                if is_user_space_owner(user_info):
                     is_owner = True
             except Exception: pass
         
@@ -332,5 +337,5 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
     
 if __name__ == "__main__":
     if tagger_instance is None:
-        print("CRITICAL: Tagger 未能初始化，应用功能将受限。")
+        print("CRITICAL: Tagger failed to initialize. App functionality will be limited.")
     demo.launch(server_name="0.0.0.0", server_port=7860)