Spaces:
Running
Running
| """Tools tab: prompt analyzer, artist recommender and web tag search.""" | |
| import html | |
| import time | |
| import gradio as gr | |
| from src.i18n import t | |
| from src.prompt_analyzer import analyze_prompt, format_analysis_html, format_score_badge | |
| from src.ui.common import ( | |
| MAX_OUTPUTS, warehouse, STYLECHOICES, ARTISTCHOICES, | |
| all_artists, build_artist_choice, tree_header, | |
| ) | |
| # Debounce for prompt analyzer (300ms) | |
| _ANALYZER_LAST = {"time": 0.0, "prompt": "", "html": "", "badge": ""} | |
| def web_search_empty_html(lc: str) -> str: | |
| """The empty-result strip. | |
| Kept in one place because it is written three times over: the initial | |
| value, the clear button's reset and update_lang. Duplicating it is how the | |
| strip ended up still English after a switch to RU/ZH -- two of the three | |
| copies were never re-rendered. The colour is the token, not rgba(): the | |
| hard-coded slate ignored every theme but the default one. | |
| """ | |
| return ( | |
| '<div class="art-info-text" style="color:var(--text-faint);font-size:12px;">' | |
| + t("web_search_empty", lc) + "</div>" | |
| ) | |
| def on_web_search(query: str, lang: str) -> str: | |
| from src.tag_searcher import search_tags | |
| lc = lang.lower() if lang in ("EN", "RU", "ZH") else "en" | |
| if not query or not query.strip(): | |
| return f'<div class="art-info-text" style="color:var(--text-faint);font-size:12px;">{t("web_search_empty", lc)}</div>' | |
| try: | |
| results = search_tags(query, warehouse) | |
| parts = [] | |
| if results["local"]: | |
| parts.append(f'<div style="margin-top:6px;font-size:13px;font-weight:600;color:var(--text);">🔍 {t("web_search_offline", lc)} ({len(results["local"])})</div>') | |
| for r in results["local"][:15]: | |
| safe_tag = html.escape(r["tag"]) | |
| safe_cat = html.escape(r["category"]) | |
| parts.append(f'<span style="display:inline-block;background:var(--accent-glow);border:1px solid var(--accent-border);border-radius:4px;padding:1px 6px;margin:2px;font-size:12px;color:var(--text);">{safe_tag} <span style="color:var(--text-faint);font-size:10px;">[{safe_cat}]</span></span>') | |
| if results["web"]: | |
| parts.append(f'<div style="margin-top:8px;font-size:13px;font-weight:600;color:var(--text);">🌐 {t("web_search_online", lc)} ({len(results["web"])})</div>') | |
| for r in results["web"][:MAX_OUTPUTS]: | |
| safe_tag = html.escape(r.get("tag", "")) | |
| count = r.get("post_count", 0) | |
| parts.append(f'<span style="display:inline-block;background:var(--accent-glow);border:1px solid var(--accent-border);border-radius:4px;padding:1px 6px;margin:2px;font-size:12px;color:var(--text);">{safe_tag} <span style="color:var(--text-faint);font-size:10px;">({count} posts)</span></span>') | |
| if not parts: | |
| return f'<div class="art-info-text" style="font-size:12px;">{t("web_enrich_no_results", lc)}</div>' | |
| return '<div style="line-height:1.8;">' + "".join(parts) + '</div>' | |
| except Exception as exc: | |
| return f'<div style="color:#ef4444;font-size:12px;">{t("web_search_error", lc).format(exc=html.escape(str(exc)))}</div>' | |
| def on_artist_filter_change(style: str, query: str): | |
| try: | |
| from src.artist_recommender import recommend_artists | |
| filtered = recommend_artists(warehouse, query=query or "", style=style or "", limit=50) | |
| except Exception: | |
| q = (query or "").lower().strip() | |
| filtered = all_artists | |
| if style: | |
| filtered = [a for a in filtered if a.get("style") == style] | |
| if q: | |
| filtered = [ | |
| a for a in filtered | |
| if q in a["tag"].lower() | |
| or q in a.get("style", "").lower() | |
| or any(q in t.lower() for t in a.get("signature_tags", [])) | |
| ] | |
| if not q and len(filtered) > 50: | |
| filtered = sorted(filtered, key=lambda a: a.get("popularity", 0), reverse=True)[:50] | |
| choices = [(build_artist_choice(a), a["tag"]) for a in filtered] | |
| return gr.update(choices=choices, value=[]) | |
| def format_artist_info(artist_names: list[str], lang: str = "en") -> str: | |
| title = t("artist_preview_title", lang) | |
| if not artist_names: | |
| empty = t("artist_preview_empty", lang) | |
| return f""" | |
| <div class="art-artist-preview"> | |
| <div class="art-artist-preview-title">{title}</div> | |
| <div class="art-artist-preview-empty">{empty}</div> | |
| </div> | |
| """ | |
| lines = [] | |
| for name in artist_names: | |
| safe_name = html.escape(name) | |
| # Use vivid recommender description (local, no network) with fallback to warehouse | |
| try: | |
| from src.artist_recommender import get_vivid_description as _vivid | |
| desc = _vivid(warehouse, name) | |
| # If vivid is generic, fallback to warehouse description | |
| if "not in local pool" in desc: | |
| desc = warehouse.get_artist_description(name, lang) or desc | |
| except Exception: | |
| desc = warehouse.get_artist_description(name, lang) | |
| url = warehouse.get_artist_danbooru_url(name) | |
| if desc: | |
| safe_url = html.escape(url) if url else "" | |
| safe_desc = html.escape(desc) | |
| link = f"<a class='art-artist-link' href='{safe_url}' target='_blank'>{t('danbooru_link', lang)}</a>" if url else "" | |
| lines.append(f""" | |
| <div style='margin:8px 0;padding:8px 0;border-bottom:1px solid rgba(79,70,229,0.08);'> | |
| <div style='display:flex;justify-content:space-between;align-items:center;gap:8px;flex-wrap:wrap;'> | |
| <strong style='color:var(--text);font-size:14px;'>{safe_name}</strong> | |
| {link} | |
| </div> | |
| <div style='color:var(--text-dim);font-size:12px;margin-top:4px;line-height:1.4;'>{safe_desc}</div> | |
| </div> | |
| """) | |
| # Similar artists from the first selected one (vivid recommender) | |
| if len(artist_names) == 1: | |
| try: | |
| from src.artist_recommender import get_similar_artists as _similar | |
| similar = _similar(warehouse, artist_names[0], limit=3) | |
| except Exception: | |
| similar = warehouse.get_similar_artists(artist_names[0], limit=3) | |
| if similar: | |
| names = ", ".join(html.escape(a["tag"]) for a in similar) | |
| label = t("similar_artists", lang) | |
| lines.append(f"<div style='margin-top:8px;font-size:12px;color:#7C3AED;'>✨ {label}: {names}</div>") | |
| return "<div class='art-artist-preview'><div class='art-artist-preview-title'>" + title + "</div>" + "".join(lines) + "</div>" | |
| def on_artist_selection_change(selected_artists, lang): | |
| return gr.update(value=format_artist_info(selected_artists or [], lang)) | |
| def on_artist_style_recommendations(style: str, lang: str): | |
| title = t("artist_preview_title", lang) | |
| if not style: | |
| empty = t("artist_preview_empty", lang) | |
| return gr.update(value=f""" | |
| <div class="art-artist-preview"> | |
| <div class="art-artist-preview-title">{title}</div> | |
| <div class="art-artist-preview-empty">{empty}</div> | |
| </div> | |
| """) | |
| artists = warehouse.get_artists_by_style(style) | |
| artists = sorted(artists, key=lambda a: a.get("popularity", 0), reverse=True)[:3] | |
| if not artists: | |
| empty = t("artist_preview_empty", lang) | |
| return gr.update(value=f""" | |
| <div class="art-artist-preview"> | |
| <div class="art-artist-preview-title">{title}</div> | |
| <div class="art-artist-preview-empty">{empty}</div> | |
| </div> | |
| """) | |
| names = ", ".join(html.escape(a["tag"]) for a in artists) | |
| label = t("top_artists_in_style", lang) | |
| return gr.update(value=f"<div class='art-artist-preview'><div class='art-artist-preview-title'>{title}</div><div style='font-size:12px;color:#7C3AED;'>✨ {label}: {names}</div></div>") | |
| def on_prompt_analyze(prompt: str, lang: str): | |
| # Debounce 300ms — avoid parsing on every keystroke | |
| now = time.time() | |
| key = (prompt or "").strip() | |
| if key == _ANALYZER_LAST["prompt"] and (now - _ANALYZER_LAST["time"]) < 0.30: | |
| return gr.update(value=_ANALYZER_LAST["html"]), gr.update(value=_ANALYZER_LAST["badge"]) | |
| lc = lang.lower() if lang in ("EN", "RU", "ZH") else "en" | |
| data = analyze_prompt(prompt, warehouse) | |
| html_out = format_analysis_html(data, lc) | |
| badge_out = format_score_badge(data, lc) | |
| _ANALYZER_LAST["time"] = now | |
| _ANALYZER_LAST["prompt"] = key | |
| _ANALYZER_LAST["html"] = html_out | |
| _ANALYZER_LAST["badge"] = badge_out | |
| return gr.update(value=html_out), gr.update(value=badge_out) | |
| def build(lang_state) -> dict: | |
| """Build the Tools tab inside the current Blocks context.""" | |
| c: dict = {} | |
| with gr.TabItem(f"🔧 {t('tab_tools', 'en')}") as tab_tools: | |
| with gr.Column(elem_classes=["art-tree"]): | |
| # Tree root: Analyzer | |
| with gr.Column(elem_classes=["art-tree-node"]): | |
| analyzer_header_html = gr.HTML(tree_header("📊", "analyzer_section", "en")) | |
| analyzer_md = gr.Markdown( | |
| value=f'<div class="art-info-text">{t("analyzer_desc", "en")}</div>', | |
| visible=True, | |
| ) | |
| # Tree root: Artist (slide-down) | |
| with gr.Column(elem_classes=["art-tree-node"]): | |
| artist_header_html = gr.HTML(tree_header("🖌️", "artist_section", "en")) | |
| with gr.Row(): | |
| artist_style_trigger = gr.Button( | |
| f'▸ {t("artist_style_placeholder", "en")}', | |
| elem_classes=["art-slide-trigger"], elem_id="art-trig-artist-style", | |
| ) | |
| artist_search_trigger = gr.Button( | |
| f'▸ {t("artist_search_label", "en")}', | |
| elem_classes=["art-slide-trigger"], elem_id="art-trig-artist-search", | |
| ) | |
| artist_multiselect_trigger = gr.Button( | |
| f'▸ {t("artist_select", "en")}', | |
| elem_classes=["art-slide-trigger"], elem_id="art-trig-artist-multiselect", | |
| ) | |
| artist_style_dropdown = gr.Radio( | |
| choices=STYLECHOICES, value="", | |
| label="", show_label=False, | |
| elem_id="art-panel-artist-style", | |
| elem_classes=["art-slide-panel", "art-slide-closed", "art-chip-group"], | |
| ) | |
| artist_search = gr.Textbox( | |
| label=t("artist_search_label", "en"), | |
| placeholder=t("artist_search_placeholder", "en"), | |
| info=t("artist_search_desc", "en"), | |
| elem_id="art-panel-artist-search", | |
| elem_classes=["art-slide-panel", "art-slide-closed"], | |
| ) | |
| artist_multiselect = gr.CheckboxGroup( | |
| choices=ARTISTCHOICES, value=[], | |
| label="", show_label=False, | |
| elem_id="art-panel-artist-multiselect", | |
| elem_classes=["art-slide-panel", "art-slide-closed", "art-chip-group"], | |
| ) | |
| artist_style_dropdown.change( | |
| fn=on_artist_filter_change, | |
| inputs=[artist_style_dropdown, artist_search], | |
| outputs=[artist_multiselect], | |
| ) | |
| artist_search.change( | |
| fn=on_artist_filter_change, | |
| inputs=[artist_style_dropdown, artist_search], | |
| outputs=[artist_multiselect], | |
| ) | |
| artist_info_md = gr.HTML( | |
| value=format_artist_info([], "en"), | |
| visible=True, | |
| ) | |
| artist_multiselect.change( | |
| fn=on_artist_selection_change, | |
| inputs=[artist_multiselect, lang_state], | |
| outputs=[artist_info_md], | |
| ) | |
| artist_style_dropdown.change( | |
| fn=on_artist_style_recommendations, | |
| inputs=[artist_style_dropdown, lang_state], | |
| outputs=[artist_info_md], | |
| ) | |
| # Tree root: Web Search | |
| with gr.Column(elem_classes=["art-tree-node"]): | |
| web_search_panel = gr.Accordion(t("web_search_label", "en"), open=False, elem_classes=["art-tree-branch"]) | |
| with web_search_panel: | |
| web_search_input = gr.Textbox( | |
| label="", show_label=False, | |
| placeholder=t("web_search_placeholder", "en"), | |
| scale=1, | |
| ) | |
| with gr.Row(): | |
| web_search_btn = gr.Button(t("web_search_btn", "en"), scale=0, variant="secondary") | |
| web_search_clear = gr.ClearButton(value="✕", scale=0, variant="secondary") | |
| web_search_output = gr.HTML(value=web_search_empty_html("en")) | |
| c.update( | |
| tab_tools=tab_tools, | |
| analyzer_header_html=analyzer_header_html, | |
| analyzer_md=analyzer_md, | |
| artist_header_html=artist_header_html, | |
| artist_style_trigger=artist_style_trigger, | |
| artist_style_dropdown=artist_style_dropdown, | |
| artist_search_trigger=artist_search_trigger, | |
| artist_search=artist_search, | |
| artist_multiselect_trigger=artist_multiselect_trigger, | |
| artist_multiselect=artist_multiselect, | |
| artist_info_md=artist_info_md, | |
| web_search_panel=web_search_panel, | |
| web_search_input=web_search_input, | |
| web_search_btn=web_search_btn, | |
| web_search_clear=web_search_clear, | |
| web_search_output=web_search_output, | |
| ) | |
| web_search_btn.click( | |
| fn=on_web_search, | |
| inputs=[web_search_input, lang_state], | |
| outputs=[web_search_output], | |
| ) | |
| web_search_input.submit( | |
| fn=on_web_search, | |
| inputs=[web_search_input, lang_state], | |
| outputs=[web_search_output], | |
| ) | |
| web_search_clear.click( | |
| fn=lambda lang: web_search_empty_html(lang.lower() if lang in ("EN", "RU", "ZH") else "en"), | |
| inputs=[lang_state], | |
| outputs=[web_search_output], | |
| ) | |
| return c | |