Update app.py

app.py

@@ -6,26 +6,15 @@ import numpy as np
 import onnxruntime as rt
 import pandas as pd
 from PIL import Image
-from huggingface_hub import login
-from translator import translate_texts, translate_texts_with_dynamic_keys
+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"
-OWNER_USERNAME = os.environ.get("OWNER_USERNAME", "")
 
-HF_TOKEN = os.environ.get("HF_TOKEN", "")
-if HF_TOKEN:
-    try:
-        login(token=HF_TOKEN)
-        print("✅ HF_TOKEN 登录成功")
-    except Exception as e:
-        print(f"⚠️ HF_TOKEN 登录失败: {e}")
-else:
-    print("⚠️ 未检测到 HF_TOKEN，私有模型可能下载失败")
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
 class Tagger:
     def __init__(self):
@@ -59,12 +48,9 @@ class Tagger:
             print(f"❌ 模型或标签加载失败: {e}")
             raise RuntimeError(f"模型初始化失败: {e}")
 
-
     def _preprocess(self, img: Image.Image) -> np.ndarray:
-        if img is None:
-            raise ValueError("输入图像不能为空")
-        if img.mode != "RGB":
-            img = img.convert("RGB")
+        if img is None: raise ValueError("输入图像不能为空")
+        if img.mode != "RGB": img = img.convert("RGB")
         size = max(img.size)
         canvas = Image.new("RGB", (size, size), (255, 255, 255))
         canvas.paste(img, ((size - img.width) // 2, (size - img.height) // 2))
@@ -73,43 +59,32 @@ class Tagger:
         return np.array(canvas)[:, :, ::-1].astype(np.float32)
 
     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("模型未成功加载，无法进行预测。")
         inp_name = self.model.get_inputs()[0].name
         outputs = self.model.run(None, {inp_name: self._preprocess(img)[None, ...]})[0][0]
 
         res = {"ratings": {}, "general": {}, "characters": {}}
         tag_categories_for_translation = {"ratings": [], "general": [], "characters": []}
 
-        for idx in self.categories["rating"]:
-            tag_name = self.tag_names[idx].replace("_", " ")
-            res["ratings"][tag_name] = float(outputs[idx])
-            tag_categories_for_translation["ratings"].append(tag_name)
-
-        for idx in self.categories["general"]:
-            if outputs[idx] > gen_th:
-                tag_name = self.tag_names[idx].replace("_", " ")
-                res["general"][tag_name] = float(outputs[idx])
-                tag_categories_for_translation["general"].append(tag_name)
-
-        for idx in self.categories["character"]:
-            if outputs[idx] > char_th:
-                tag_name = self.tag_names[idx].replace("_", " ")
-                res["characters"][tag_name] = float(outputs[idx])
-                tag_categories_for_translation["characters"].append(tag_name)
-
-
-        res["general"] = dict(sorted(res["general"].items(), key=lambda kv: kv[1], reverse=True))
-        res["characters"] = dict(sorted(res["characters"].items(), key=lambda kv: kv[1], reverse=True))
-        res["ratings"] = dict(sorted(res["ratings"].items(), key=lambda kv: kv[1], reverse=True))
-
-
-        tag_categories_for_translation["general"] = list(res["general"].keys())
-        tag_categories_for_translation["characters"] = list(res["characters"].keys())
-        tag_categories_for_translation["ratings"] = list(res["ratings"].keys())
+        for cat_key, cat_indices in self.categories.items():
+            sub_res = {}
+            if cat_key == "rating":
+                for idx in cat_indices:
+                    tag_name = self.tag_names[idx].replace("_", " ")
+                    sub_res[tag_name] = float(outputs[idx])
+            else:
+                threshold = char_th if cat_key == "character" else gen_th
+                for idx in cat_indices:
+                    if outputs[idx] > threshold:
+                        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
-
 try:
     tagger_instance = Tagger()
 except RuntimeError as e:
@@ -134,10 +109,13 @@ function copyToClipboard(text) {
     }
     navigator.clipboard.writeText(text).then(() => {
         const feedback = document.createElement('div');
-        let displayText = String(text);
-        displayText = displayText.substring(0, 30) + (displayText.length > 30 ? '...' : '');
+        let displayText = String(text).substring(0, 30) + (String(text).length > 30 ? '...' : '');
         feedback.textContent = '已复制: ' + displayText;
-        feedback.style.cssText = 'position: fixed; bottom: 20px; left: 50%; transform: translateX(-50%); background-color: #4CAF50; color: white; padding: 10px 20px; border-radius: 5px; z-index: 10000; transition: opacity 0.5s ease-out;';
+        Object.assign(feedback.style, {
+            position: 'fixed', bottom: '20px', left: '50%', transform: 'translateX(-50%)',
+            backgroundColor: '#4CAF50', color: 'white', padding: '10px 20px',
+            borderRadius: '5px', zIndex: '10000', transition: 'opacity 0.5s ease-out'
+        });
         document.body.appendChild(feedback);
         setTimeout(() => {
             feedback.style.opacity = '0';
@@ -152,11 +130,14 @@ 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({})
-    state_tag_categories_for_translation = gr.State({})
-
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -164,164 +145,158 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI 图像标签分析器", css=cus
             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="通用标签阈值", info="越高 → 标签更少更准")
-                char_slider = gr.Slider(0, 1, value=0.85, step=0.01, label="角色标签阈值", info="推荐保持较高阈值")
+                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="角色标签阈值")
                 show_tag_scores = gr.Checkbox(True, label="在列表中显示标签置信度")
-                
-                with gr.Group(visible=False) as guest_api_group:
-                    gr.Markdown("### 访客翻译API配置\n由于您不是本空间所有者，需要提供自己的翻译API密钥才能使用翻译功能。")
-                    guest_tencent_id = gr.Textbox(label="腾讯云 Secret ID", type="password")
-                    guest_tencent_key = gr.Textbox(label="腾讯云 Secret Key", type="password")
-                    guest_baidu_json = gr.TextArea(
-                        label="百度翻译凭证 JSON",
-                        placeholder='[{"app_id": "...", "secret_key": "..."}, ...]',
-                        lines=3
-                    )
-            
+
+            with gr.Accordion("🔑 自定义翻译密钥 (可选)", open=False, visible=False) as api_key_accordion:
+                gr.Markdown("如果你不是空间所有者，需要在这里提供自己的API密钥才能使用翻译功能。")
+                tencent_id_in = gr.Textbox(label="腾讯云 Secret ID", lines=1)
+                tencent_key_in = gr.Textbox(label="腾讯云 Secret Key", lines=1, type="password")
+                baidu_json_in = gr.Textbox(label="百度翻译凭证 (JSON 格式)", lines=3, placeholder='[{"app_id": "...", "secret_key": "..."}]')
+
             with gr.Accordion("📊 标签汇总设置", open=True):
-                gr.Markdown("选择要包含在下方汇总文本框中的标签类别：")
-                with gr.Row():
-                    sum_general = gr.Checkbox(True, label="通用标签", min_width=50)
-                    sum_char = gr.Checkbox(True, label="角色标签", min_width=50)
-                    sum_rating = gr.Checkbox(False, label="评分标签", min_width=50)
-                sum_sep = gr.Dropdown(["逗号", "换行", "空格"], value="逗号", label="标签之间的分隔符")
+                gr.CheckboxGroup(["通用标签", "角色标签", "评分标签"], value=["通用标签", "角色标签"], label="汇总类别")
+                sum_sep = gr.Dropdown(["逗号", "换行", "空格"], value="逗号", label="标签分隔符")
                 sum_show_zh = gr.Checkbox(False, label="在汇总中显示中文翻译")
 
             processing_info = gr.Markdown("", visible=False)
 
         with gr.Column(scale=2):
             with gr.Tabs():
-                with gr.TabItem("🏷️ 通用标签"):
-                    out_general = gr.HTML(label="General Tags")
-                with gr.TabItem("👤 角色标签"):
-                    gr.Markdown("<p style='color:gray; font-size:small;'>提示：角色标签推测基于截至2024年2月的数据。</p>")
-                    out_char = gr.HTML(label="Character Tags")
-                with gr.TabItem("⭐ 评分标签"):
-                    out_rating = gr.HTML(label="Rating Tags")
-            
+                with gr.TabItem("🏷️ 通用标签"): out_general = gr.HTML(label="General Tags")
+                with gr.TabItem("👤 角色标签"): out_char = gr.HTML(label="Character Tags")
+                with gr.TabItem("⭐ 评分标签"): out_rating = gr.HTML(label="Rating Tags")
             gr.Markdown("### 标签汇总结果")
-            out_summary = gr.Textbox(
-                label="标签汇总",
-                placeholder="分析完成后，此处将显示汇总的英文标签...",
-                lines=5,
-                show_copy_button=True
-            )
-
-    def format_tags_html(tags_dict, translations_list, category_name, show_scores=True, show_translation_in_list=True):
+            out_summary = gr.Textbox(label="标签汇总", lines=5, show_copy_button=True)
+            
+    def check_user_status(request: gr.Request):
+        token = request.token
+        if token:
+            try:
+                user_info = whoami(token=token)
+                space_runtime = get_space_runtime()
+                if space_runtime and user_info and user_info["name"] == space_runtime.owner:
+                    return f"✅ 以所有者 **{user_info['fullname']}** 身份登录，将使用空间配置的密钥。", 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}")
+                return "⚠️ 无法验证您的登录状态。请提供 API 密钥。", gr.update(visible=True, open=True)
+        return "ℹ️ **访客模式**。如需使用翻译功能，请<a href='/login?redirect=/'>登录</a>或提供 API 密钥。", gr.update(visible=True, open=True)
+        
+    def format_tags_html(tags_dict, translations_list, show_scores):
         if not tags_dict: return "<p>暂无标签</p>"
         html = '<div class="label-container">'
-        tag_keys = list(tags_dict.keys())
-        for i, tag in enumerate(tag_keys):
-            score = tags_dict[tag]
+        for i, (tag, score) in enumerate(tags_dict.items()):
             escaped_tag = tag.replace("'", "\\'")
             html += '<div class="tag-item">'
             tag_display_html = f'<span class="tag-en" onclick="copyToClipboard(\'{escaped_tag}\')">{tag}</span>'
-            if show_translation_in_list and i < len(translations_list) and translations_list[i]:
+            if i < len(translations_list) and translations_list[i]:
                 tag_display_html += f'<span class="tag-zh">({translations_list[i]})</span>'
             html += f'<div>{tag_display_html}</div>'
             if show_scores: html += f'<span class="tag-score">{score:.3f}</span>'
             html += '</div>'
-        html += '</div>'
-        return html
-
-    def generate_summary_text_content(current_res, current_translations_dict, s_gen, s_char, s_rat, s_sep_type, s_show_zh):
-        if not current_res: return "请先分析图像或选择要汇总的标签类别。"
-        summary_parts, separator = [], {"逗号": ", ", "换行": "\n", "空格": " "}.get(s_sep_type, ", ")
-        categories_to_summarize = []
-        if s_gen: categories_to_summarize.append("general")
-        if s_char: categories_to_summarize.append("characters")
-        if s_rat: categories_to_summarize.append("ratings")
-        if not categories_to_summarize: return "请至少选择一个标签类别进行汇总。"
-
-        for cat_key in categories_to_summarize:
+        return html + '</div>'
+
+    def generate_summary_text_content(current_res, translations, sum_cats, sep_type, show_zh):
+        if not current_res: return "请先分析图像。"
+        parts, sep = [], {"逗号": ", ", "换行": "\n", "空格": " "}.get(sep_type, ", ")
+        cat_map = {"通用标签": "general", "角色标签": "characters", "评分标签": "ratings"}
+        for cat_name in sum_cats:
+            cat_key = cat_map[cat_name]
             if current_res.get(cat_key):
-                tags_to_join = []
-                cat_tags_en = list(current_res[cat_key].keys())
-                cat_translations = current_translations_dict.get(cat_key, [])
-                for i, en_tag in enumerate(cat_tags_en):
-                    if s_show_zh and i < len(cat_translations) and cat_translations[i]:
-                        tags_to_join.append(f"{en_tag}({cat_translations[i]})")
-                    else:
-                        tags_to_join.append(en_tag)
-                if tags_to_join: summary_parts.append(separator.join(tags_to_join))
-        joiner = "\n\n" if separator != "\n" and len(summary_parts) > 1 else separator
-        final_summary = joiner.join(summary_parts)
-        return final_summary if final_summary else "选定的类别中没有找到标签。"
-
-    def process_image_and_generate_outputs(img, g_th, c_th, s_scores, s_gen, s_char, s_rat, s_sep, s_zh_in_sum, guest_tc_id, guest_tc_key, guest_bd_json, request: gr.Request):
+                tags_en, trans = list(current_res[cat_key].keys()), translations.get(cat_key, [])
+                tags_to_join = [f"{en}({zh})" if show_zh and i < len(trans) and trans[i] else en for i, en in enumerate(tags_en)]
+                if tags_to_join: parts.append(sep.join(tags_to_join))
+        return "\n".join(parts) if parts else "选定的类别中没有找到标签。"
+
+    def process_image_and_generate_outputs(
+        img, g_th, c_th, s_scores,
+        user_tencent_id, user_tencent_key, user_baidu_json,
+        sum_cats, s_sep, s_zh_in_sum,
+        request: gr.Request
+    ):
         if img is None:
-            yield (gr.update(interactive=True, value="🚀 开始分析"), gr.update(visible=True, value="❌ 请先上传图片。"), "", "", "", "", gr.update(placeholder="请先上传图片并开始分析..."), {}, {}, {})
-            return
+            raise gr.Error("请先上传图片。")
         if tagger_instance is None:
-            yield (gr.update(interactive=True, value="🚀 开始分析"), gr.update(visible=True, value="❌ 分析器未成功初始化，请检查控制台错误。"), "", "", "", "", gr.update(placeholder="分析器初始化失败..."), {}, {}, {})
-            return
-
-        yield (gr.update(interactive=False, value="🔄 处理中..."), gr.update(visible=True, value="🔄 正在分析图像..."), gr.HTML(value="<p>分析中...</p>"), gr.HTML(value="<p>分析中...</p>"), gr.HTML(value="<p>分析中...</p>"), gr.update(value="分析中..."), {}, {}, {})
+            raise gr.Error("分析器未成功初始化，请检查后台错误。")
+
+        yield gr.update(interactive=False, value="🔄 处理中..."), gr.update(visible=True, value="🔄 正在分析..."), *["<p>分析中...</p>"]*3, "分析中...", {}, {}
+
+        token, is_owner = request.token, False
+        if token:
+            try:
+                user_info = whoami(token=token)
+                space_runtime = get_space_runtime()
+                if space_runtime and user_info and user_info["name"] == space_runtime.owner:
+                    is_owner = True
+            except Exception: pass
         
+        final_tencent_id, final_tencent_key, baidu_json_str = (
+            (os.environ.get("TENCENT_SECRET_ID"), os.environ.get("TENCENT_SECRET_KEY"), os.environ.get("BAIDU_CREDENTIALS_JSON", "[]"))
+            if is_owner else (user_tencent_id, user_tencent_key, user_baidu_json)
+        )
+        
+        final_baidu_creds_list = []
+        if baidu_json_str and baidu_json_str.strip():
+            try:
+                parsed_data = json.loads(baidu_json_str)
+                if isinstance(parsed_data, list): final_baidu_creds_list = parsed_data
+            except json.JSONDecodeError: print("提供的百度凭证JSON无效。")
+
         try:
-            res, tag_categories_original_order = tagger_instance.predict(img, g_th, c_th)
-            all_tags_to_translate = [tag for cat_key in ["general", "characters", "ratings"] for tag in tag_categories_original_order.get(cat_key, [])]
+            res, tag_cats_original = tagger_instance.predict(img, g_th, c_th)
+            all_tags = [tag for cat in tag_cats_original.values() for tag in cat]
             
-            all_translations_flat = []
-            if all_tags_to_translate:
-                is_owner = request.username and request.username == OWNER_USERNAME
-                if is_owner:
-                    print("- [Auth] 所有者身份，使用预置密钥进行翻译。")
-                    all_translations_flat = translate_texts(all_tags_to_translate)
-                else:
-                    print("- [Auth] 访客身份，使用用户提供的密钥进行翻译。")
-                    if not guest_tc_id and not guest_bd_json:
-                        print("  - [Warning] 访客未提供任何API密钥，将跳过翻译。")
-                        all_translations_flat = all_tags_to_translate
-                    else:
-                        all_translations_flat = translate_texts_with_dynamic_keys(all_tags_to_translate, guest_tc_id, guest_tc_key, guest_bd_json)
-
-            current_translations_dict = {}
-            offset = 0
-            for cat_key in ["general", "characters", "ratings"]:
-                num_tags = len(tag_categories_original_order.get(cat_key, []))
-                current_translations_dict[cat_key] = all_translations_flat[offset : offset + num_tags]
-                offset += num_tags
+            translations_flat = translate_texts(
+                all_tags,
+                tencent_secret_id=final_tencent_id,
+                tencent_secret_key=final_tencent_key,
+                baidu_credentials_list=final_baidu_creds_list
+            ) if all_tags else []
+            
+            translations, offset = {}, 0
+            for cat_key, tags in tag_cats_original.items():
+                translations[cat_key] = translations_flat[offset : offset + len(tags)]
+                offset += len(tags)
             
-            general_html = format_tags_html(res.get("general", {}), current_translations_dict.get("general", []), "general", s_scores)
-            char_html = format_tags_html(res.get("characters", {}), current_translations_dict.get("characters", []), "characters", s_scores)
-            rating_html = format_tags_html(res.get("ratings", {}), current_translations_dict.get("ratings", []), "ratings", s_scores)
-            summary_text = generate_summary_text_content(res, current_translations_dict, s_gen, s_char, s_rat, s_sep, s_zh_in_sum)
+            outputs_html = {k: format_tags_html(res.get(k, {}), translations.get(k, []), s_scores) for k in ["general", "characters", "ratings"]}
+            summary = generate_summary_text_content(res, translations, sum_cats, s_sep, s_zh_in_sum)
 
-            yield (gr.update(interactive=True, value="🚀 开始分析"), gr.update(visible=True, value="✅ 分析完成!"), general_html, char_html, rating_html, gr.update(value=summary_text), res, current_translations_dict, tag_categories_original_order)
+            yield gr.update(interactive=True, value="🚀 开始分析"), gr.update(visible=True, value="✅ 分析完成!"), outputs_html["general"], outputs_html["characters"], outputs_html["ratings"], summary, res, translations
             
         except Exception as e:
             import traceback
-            tb_str = traceback.format_exc()
-            print(f"处理时发生错误: {e}\n{tb_str}")
-            yield (gr.update(interactive=True, value="🚀 开始分析"), gr.update(visible=True, value=f"❌ 处理失败: {str(e)}"), "<p>处理出错</p>", "<p>处理出错</p>", "<p>处理出错</p>", gr.update(value=f"错误: {str(e)}", placeholder="分析失败..."), {}, {}, {})
-
-    def update_summary_display(s_gen, s_char, s_rat, s_sep, s_zh_in_sum, current_res, current_translations):
-        if not current_res: return gr.update(placeholder="请先完成一次图像分析以生成汇总。", value="")
-        new_summary_text = generate_summary_text_content(current_res, current_translations, s_gen, s_char, s_rat, s_sep, s_zh_in_sum)
-        return gr.update(value=new_summary_text)
-
-    def check_user_auth(request: gr.Request):
-        if not OWNER_USERNAME: print("⚠️ 警告: 未设置 OWNER_USERNAME 环境变量。所有用户都将被视为访客。")
-        if not request.username or request.username != OWNER_USERNAME:
-            print(f"- [Auth] 访客 '{request.username}' 已连接，显示 API Key 输入框。")
-            return gr.update(visible=True)
-        print(f"- [Auth] 所有者 '{request.username}' 已连接。")
-        return gr.update(visible=False)
-
-    demo.load(fn=check_user_auth, inputs=None, outputs=[guest_api_group])
+            traceback.print_exc()
+            raise gr.Error(f"处理时发生错误: {e}")
 
+    demo.load(fn=check_user_status, outputs=[user_status_md, api_key_accordion], queue=False)
+    
     btn.click(
         process_image_and_generate_outputs,
-        inputs=[img_in, gen_slider, char_slider, show_tag_scores, sum_general, sum_char, sum_rating, sum_sep, sum_show_zh, guest_tencent_id, guest_tencent_key, guest_baidu_json],
-        outputs=[btn, processing_info, out_general, out_char, out_rating, out_summary, state_res, state_translations_dict, state_tag_categories_for_translation]
+        inputs=[
+            img_in, gen_slider, char_slider, show_tag_scores,
+            tencent_id_in, tencent_key_in, baidu_json_in,
+            sum_general, sum_sep, sum_show_zh
+        ],
+        outputs=[
+            btn, processing_info,
+            out_general, out_char, out_rating,
+            out_summary,
+            state_res, state_translations_dict
+        ],
     )
 
-    summary_controls = [sum_general, sum_char, sum_rating, sum_sep, sum_show_zh]
+    summary_controls = [sum_general, sum_sep, sum_show_zh]
     for ctrl in summary_controls:
-        ctrl.change(fn=update_summary_display, inputs=summary_controls + [state_res, state_translations_dict], outputs=[out_summary])
+        ctrl.change(
+            fn=lambda r, t, c, s, z: generate_summary_text_content(r, t, c, s, z),
+            inputs=[state_res, state_translations_dict] + summary_controls,
+            outputs=[out_summary],
+        )
     
 if __name__ == "__main__":
     if tagger_instance is None:
-        print("CRITICAL: Tagger 未能初始化，应用功能将受限。请检查之前的错误信息。")
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+        print("CRITICAL: Tagger 未能初始化，应用功能将受限。")
+    demo.launch(server_name="0.0.0.0", server_port=7860)