Spaces:
Running
Running
| import gradio as gr | |
| import html | |
| from src.i18n import t | |
| from src.prompt_parser import parse_prompt | |
| from src.tag_warehouse import TagWarehouse, STYLE_ICONS | |
| from src.variation_engine import generate_variations, generate_negative_prompt | |
| from src.ensemble_tagger import get_ensemble_tagger | |
| from src.image_tagger import _TAGGER_DEPS_OK | |
| from src.presets import ( | |
| PRESETS, PRESET_GROUPS, get_preset_categories, | |
| get_user_preset_names, save_user_preset, delete_user_preset | |
| ) | |
| from src.prompt_analyzer import analyze_prompt, format_analysis_html | |
| from src.session_history import HistoryEntry, add_to_history, get_history | |
| from src.synonym_data import reload_synonym_groups | |
| from src.tag_format import to_booru_tag, to_prompt_tag | |
| MAX_OUTPUTS = 10 | |
| warehouse = TagWarehouse() | |
| ALL_CATEGORIES = [ | |
| "quality", "lighting", "composition", "effects", "atmosphere", | |
| "framing", "style", "colors", "background", "pose", "expression", | |
| "clothing", "special_fx", "year_meta", "nsfw", | |
| "animal", "furry", "object", "food", | |
| "vehicle", "weapon", "architecture", | |
| "demon", "angelic", | |
| "hair", "body", "accessory", "season", | |
| ] | |
| CATEGORY_ICONS = { | |
| "quality": "⭐", "lighting": "💡", "composition": "📐", | |
| "effects": "🌀", "atmosphere": "🌙", "framing": "🖼️", | |
| "style": "🎨", "colors": "🌈", "background": "🏞️", | |
| "pose": "🧍", "expression": "😊", "clothing": "👗", | |
| "special_fx": "✨", "year_meta": "📅", "nsfw": "🔞", | |
| "animal": "🐾", "furry": "🐺", "object": "🏺", "food": "🍕", | |
| "vehicle": "🚗", "weapon": "⚔️", "architecture": "🏛️", | |
| "demon": "👿", "angelic": "👼", | |
| "hair": "💇", "body": "💪", "accessory": "💍", "season": "🌸", | |
| } | |
| DEFAULT_CHECKED = {"quality", "lighting", "expression"} | |
| def _build_artist_choice(a): | |
| icon = STYLE_ICONS.get(a.get('style', ''), '') | |
| style = a.get('style', '') | |
| return f"{icon} {a['tag']} [{style}]" if style else f"{icon} {a['tag']}" | |
| artist_styles = [""] + warehouse.get_artist_styles() | |
| all_artists = warehouse.get_all_artists() | |
| all_tandems = warehouse.get_all_tandems() | |
| STYLECHOICES = [(s.title() if s else t("artist_all", "en"), s) for s in artist_styles] | |
| ARTISTCHOICES = [ | |
| (_build_artist_choice(a), a["tag"]) | |
| for a in sorted(all_artists, key=lambda a: a.get("popularity", 0), reverse=True)[:50] | |
| ] | |
| TANDEMCHOICES = [(t("tandem_placeholder", "en"), None)] + [ | |
| ( | |
| f"🤝 {tand['artists'][0]} × {tand['artists'][1]} [{tand.get('style_blend','')}] ({tand.get('compatibility',0):.0%})", | |
| i, | |
| ) | |
| for i, tand in enumerate(all_tandems) | |
| ] | |
| COPY_JS = """(text) => { | |
| if (!text || !text.trim()) return text; | |
| try { | |
| navigator.clipboard.writeText(text); | |
| } catch(e) { | |
| try { | |
| const ta = document.createElement('textarea'); | |
| ta.value = text; | |
| ta.style.position = 'fixed'; | |
| ta.style.opacity = '0'; | |
| document.body.appendChild(ta); | |
| ta.focus(); | |
| ta.select(); | |
| document.execCommand('copy'); | |
| document.body.removeChild(ta); | |
| } catch(e2) {} | |
| } | |
| const btn = document.activeElement; | |
| if (btn && btn.classList.contains('whyx-copy-btn')) { | |
| const origHTML = btn.innerHTML; | |
| btn.innerHTML = '<span style="color:#34D399;font-weight:700;font-size:16px;">✓</span>'; | |
| btn.style.borderColor = '#34D399'; | |
| btn.style.background = 'rgba(52,211,153,0.15)'; | |
| btn.style.transform = 'scale(1.1)'; | |
| setTimeout(() => { | |
| btn.innerHTML = origHTML; | |
| btn.style.borderColor = ''; | |
| btn.style.background = ''; | |
| btn.style.transform = ''; | |
| }, 1400); | |
| } | |
| return text; | |
| }""" | |
| def on_generate(prompt, model, rating, num_variations, creativity, weight_mode, mode, lang, artist_style, selected_artists, use_tandems, selected_tandem_idx, web_enrich, seed, current_preset_state, fx_count, blacklist_text, mirror_blacklist, strip_quality, strip_artist, strip_lora, strip_meta, min_tags, user_preset_state, output_format, *checks) -> tuple[str, ...]: | |
| try: | |
| if not prompt or not prompt.strip(): | |
| return tuple([""] * (MAX_OUTPUTS * 2)) | |
| parsed = parse_prompt(prompt) | |
| if parsed is None: | |
| msg = t("error_parse", lang) | |
| return tuple([msg] + [""] * (MAX_OUTPUTS - 1) + [""] * MAX_OUTPUTS) | |
| active = [c for c, on in zip(ALL_CATEGORIES, checks) if on] | |
| if not active: | |
| active = ["quality", "lighting", "expression"] | |
| if rating in ("pg", "pg13") and "nsfw" in active: | |
| active.remove("nsfw") | |
| selected_tandem = None | |
| if use_tandems and selected_tandem_idx is not None: | |
| try: | |
| idx = int(selected_tandem_idx) | |
| if 0 <= idx < len(all_tandems): | |
| selected_tandem = all_tandems[idx] | |
| except (ValueError, TypeError): | |
| pass | |
| selected_presets = [p.strip() for p in (current_preset_state or "").split(",") if p.strip()] | |
| if user_preset_state: | |
| selected_presets.extend( | |
| p.strip() for p in str(user_preset_state).split(",") if p.strip() | |
| ) | |
| fx_count = {"off": 0, "light": 2, "rich": 4}.get(fx_count, 0) | |
| # Blacklist / Tags-to-remove (A): comma- or newline-separated. | |
| exclude_tags = [ | |
| t.strip().lower() | |
| for t in str(blacklist_text or "").replace("\n", ",").split(",") | |
| if t.strip() | |
| ] | |
| mirror = bool(mirror_blacklist) | |
| results = generate_variations( | |
| parsed=parsed, | |
| selected_categories=active, | |
| num_variations=int(num_variations), | |
| creativity=creativity, | |
| model=model.lower(), | |
| rating=rating, | |
| warehouse=warehouse, | |
| artist_style=artist_style or "", | |
| selected_artists=selected_artists or [], | |
| use_tandems=bool(use_tandems), | |
| selected_tandem=selected_tandem, | |
| weight_mode=weight_mode, | |
| mode=mode, | |
| web_enrich=web_enrich == "1", | |
| selected_presets=selected_presets, | |
| fx_count=fx_count, | |
| seed=int(seed) if seed is not None and str(seed).strip() else None, | |
| exclude_tags=exclude_tags, | |
| strip_quality=bool(strip_quality), | |
| strip_artist=bool(strip_artist), | |
| strip_lora=bool(strip_lora), | |
| strip_meta=bool(strip_meta), | |
| min_tags=int(min_tags or 0), | |
| output_format=output_format or "prompt", | |
| ) | |
| out = list(results) | |
| while len(out) < MAX_OUTPUTS: | |
| out.append("") | |
| neg_results = generate_negative_prompt( | |
| parsed=parsed, | |
| selected_categories=active, | |
| num_variations=int(num_variations), | |
| rating=rating, | |
| warehouse=warehouse, | |
| model=model.lower(), | |
| positive_tags=out[:MAX_OUTPUTS], | |
| extra_negative=exclude_tags if mirror else None, | |
| output_format=output_format or "prompt", | |
| ) | |
| neg_out = list(neg_results) | |
| while len(neg_out) < MAX_OUTPUTS: | |
| neg_out.append("") | |
| add_to_history( | |
| prompt, out[:MAX_OUTPUTS], model=model, rating=rating, | |
| num_variations=int(num_variations), creativity=creativity, | |
| weight_mode=weight_mode, | |
| ) | |
| return tuple(out[:MAX_OUTPUTS]) + tuple(neg_out[:MAX_OUTPUTS]) | |
| except Exception as exc: | |
| import traceback | |
| traceback.print_exc() | |
| err_msg = t("generation_error", lang).format(exc=html.escape(str(exc))) | |
| return tuple([err_msg] + [""] * (MAX_OUTPUTS - 1) + [""] * MAX_OUTPUTS) | |
| def _preset_status(key: str, lang: str, **kwargs) -> str: | |
| return f'<div class="whyx-info-text">{t(key, lang).format(**kwargs) if kwargs else t(key, lang)}</div>' | |
| def on_user_preset_save(name: str, tags: str, lang: str): | |
| name = (name or "").strip() | |
| tags_list = [t.strip() for t in str(tags or "").replace("\n", ",").split(",") if t.strip()] | |
| if not name: | |
| return gr.update(choices=get_user_preset_names()), "", _preset_status("user_preset_need_name", lang) | |
| save_user_preset(name, tags_list) | |
| return gr.update(choices=get_user_preset_names(), value=name), "", _preset_status("user_preset_saved", lang, name=name) | |
| def on_user_preset_apply(name: str, current_state: str, lang: str): | |
| if not name: | |
| return current_state, _preset_status("user_preset_pick", lang) | |
| state = (current_state or "").strip() | |
| keys = [k for k in state.split(",") if k.strip()] if state else [] | |
| if name not in keys: | |
| keys.append(name) | |
| return ",".join(keys), _preset_status("user_preset_applied", lang, name=name) | |
| def on_user_preset_delete(name: str, lang: str): | |
| if not name: | |
| return gr.update(choices=get_user_preset_names()), _preset_status("user_preset_pick", lang) | |
| delete_user_preset(name) | |
| return gr.update(choices=get_user_preset_names(), value=None), _preset_status("user_preset_deleted", lang, name=name) | |
| def on_web_search(query: str, lang: str) -> str: | |
| from src.tag_searcher import search_tags | |
| lc = "ru" if lang == "RU" else "en" | |
| if not query or not query.strip(): | |
| return f'<div class="whyx-info-text" style="color:rgba(148,163,184,0.6);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;">🔍 {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:rgba(56,189,248,0.1);border:1px solid rgba(56,189,248,0.2);border-radius:4px;padding:1px 6px;margin:2px;font-size:12px;">{safe_tag} <span style="color:rgba(148,163,184,0.6);font-size:10px;">[{safe_cat}]</span></span>') | |
| if results["web"]: | |
| parts.append(f'<div style="margin-top:8px;font-size:13px;font-weight:600;">🌐 {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:rgba(99,102,241,0.1);border:1px solid rgba(99,102,241,0.2);border-radius:4px;padding:1px 6px;margin:2px;font-size:12px;">{safe_tag} <span style="color:rgba(148,163,184,0.6);font-size:10px;">({count} posts)</span></span>') | |
| if not parts: | |
| return f'<div class="whyx-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 _format_history_html(lang: str) -> str: | |
| lc = "ru" if lang == "RU" else "en" | |
| history = get_history() | |
| lines = [] | |
| lines.append("<div style='background:rgba(15,23,42,0.5);border:1px solid rgba(56,189,248,0.12);border-radius:10px;padding:10px 14px;margin-top:4px;'>") | |
| entries = history.get_all()[:5] | |
| if not entries: | |
| lines.append(f"<div style='color:var(--text-dim);font-size:12px;'>{t('history_empty', lc)}</div>") | |
| else: | |
| for entry in entries: | |
| liked = entry.get_liked_results() | |
| lines.append("<div style='border-bottom:1px solid rgba(56,189,248,0.06);padding:6px 0;'>") | |
| lines.append("<div style='font-size:11px;color:var(--text-dim);'>") | |
| lines.append(f"<span style='color:var(--text);'>{t('history_prompt', lc)}:</span> <span style='color:var(--text-faint);'>") | |
| p = entry.prompt[:80] + "..." if len(entry.prompt) > 80 else entry.prompt | |
| lines.append(f"{html.escape(p)}</span></div>") | |
| if liked: | |
| lines.append(f"<div style='margin-top:3px;font-size:11px;color:#F43F5E;'>♥ {len(liked)} {t('history_liked', lc).lower()}</div>") | |
| lines.append("</div>") | |
| favs = history.get_favorites() | |
| if favs: | |
| lines.append("<div style='border-top:1px solid rgba(56,189,248,0.08);margin-top:8px;padding-top:8px;'>") | |
| lines.append(f"<div style='font-size:12px;color:var(--text);font-weight:600;margin-bottom:4px;'>♥ {t('favorites_label', lc)}</div>") | |
| for entry, idx, text in favs[:5]: | |
| t_text = text[:60] + "..." if len(text) > 60 else text | |
| lines.append(f"<div style='font-size:11px;color:var(--text-dim);padding:2px 0;'><span style='color:#F43F5E;'>♥</span> {html.escape(t_text)}</div>") | |
| lines.append("</div>") | |
| else: | |
| lines.append("<div style='border-top:1px solid rgba(56,189,248,0.08);margin-top:8px;padding-top:8px;'>") | |
| lines.append(f"<div style='font-size:11px;color:var(--text-faint);'>{t('history_no_likes', lc)}</div>") | |
| lines.append("</div>") | |
| lines.append("</div>") | |
| return "\n".join(lines) | |
| def _heart_click(idx: int, lang, *results) -> tuple: | |
| lc = "ru" if lang == "RU" else "en" | |
| history = get_history() | |
| entries = history.get_all() | |
| heart_updates = [] | |
| for i in range(MAX_OUTPUTS): | |
| if entries and i < len(entries[0].results): | |
| heart_updates.append("♥" if entries[0].is_liked(i) else "♡") | |
| else: | |
| heart_updates.append("♡") | |
| if entries and idx < len(results) and results[idx]: | |
| entries[0].toggle_like(idx) | |
| heart_updates[idx] = "♥" if entries[0].is_liked(idx) else "♡" | |
| return (gr.update(value=_format_history_html(lc)),) + tuple(gr.update(value=h) for h in heart_updates) | |
| def on_preset_change(prompt, *group_and_state): | |
| all_groups = list(group_and_state[:-1]) | |
| current_active = group_and_state[-1] | |
| active = None | |
| for vals in all_groups: | |
| if vals and isinstance(vals, list) and len(vals) > 0: | |
| active = vals[0] | |
| break | |
| if active == current_active: | |
| resets = [gr.update(value=[]) for _ in PRESET_GROUPS] | |
| return (*[c in DEFAULT_CHECKED for c in ALL_CATEGORIES], prompt, "", *resets) | |
| if active: | |
| cats = get_preset_categories(active) | |
| resets = [] | |
| for group in PRESET_GROUPS: | |
| if active in group["presets"]: | |
| resets.append(gr.update(value=[active])) | |
| else: | |
| resets.append(gr.update(value=[])) | |
| # Presets are now layered inside the engine (protected additions), | |
| # so the prompt text itself is left untouched. | |
| return (*[c in cats for c in ALL_CATEGORIES], prompt, active, *resets) | |
| else: | |
| resets = [gr.update(value=[]) for _ in PRESET_GROUPS] | |
| return (*[c in DEFAULT_CHECKED for c in ALL_CATEGORIES], prompt, "", *resets) | |
| def _clear_preset_state() -> str: | |
| return "" | |
| def on_rating_change(rating: str, *checks) -> list[bool]: | |
| checks = list(checks) | |
| if rating in ("pg", "pg13"): | |
| idx = ALL_CATEGORIES.index("nsfw") | |
| if idx < len(checks): | |
| checks[idx] = False | |
| return (*checks,) | |
| def on_artist_filter_change(style: str, query: str): | |
| 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", [])) | |
| ] | |
| # Lazy-load: limit to top 50 by popularity unless searching | |
| 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="whyx-artist-preview"> | |
| <div class="whyx-artist-preview-title">{title}</div> | |
| <div class="whyx-artist-preview-empty">{empty}</div> | |
| </div> | |
| """ | |
| lines = [] | |
| for name in artist_names: | |
| safe_name = html.escape(name) | |
| 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='whyx-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(56,189,248,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 | |
| if len(artist_names) == 1: | |
| 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:#38BDF8;'>✨ {label}: {names}</div>") | |
| return "<div class='whyx-artist-preview'><div class='whyx-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_prompt_analyze(prompt: str, lang: str): | |
| lc = "ru" if lang == "RU" else "en" | |
| data = analyze_prompt(prompt, warehouse) | |
| return gr.update(value=format_analysis_html(data, lc)) | |
| def _reload_data(lang): | |
| warehouse.reload_pools() | |
| warehouse.reload_artists() | |
| reload_synonym_groups() | |
| from src.prompt_rewriter import reload_map | |
| from src.tag_searcher import reload_cooccurrence | |
| from src.dedup_engine import reload_groups | |
| reload_map() | |
| reload_cooccurrence() | |
| reload_groups() | |
| 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="whyx-artist-preview"> | |
| <div class="whyx-artist-preview-title">{title}</div> | |
| <div class="whyx-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="whyx-artist-preview"> | |
| <div class="whyx-artist-preview-title">{title}</div> | |
| <div class="whyx-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='whyx-artist-preview'><div class='whyx-artist-preview-title'>{title}</div><div style='font-size:12px;color:#38BDF8;'>✨ {label}: {names}</div></div>") | |
| def _build_tagger_html(result: dict, lc: str) -> str: | |
| lines = [] | |
| lines.append("<div class='whyx-tagger-panel'>") | |
| # Natural-language caption (Qwen-VL) — shown first when the ensemble emits it. | |
| nl_caption = result.get("nl_caption") | |
| if nl_caption: | |
| lines.append( | |
| "<div class='whyx-tagger-panel-title'>🗣️ " + t("tagger_caption_label", lc) + "</div>" | |
| f"<div class='whyx-tagger-nl-caption'>{html.escape(nl_caption)}</div>" | |
| ) | |
| # Pose summary from the pose tagger, when available. | |
| pose_tags = result.get("pose_tags") or [] | |
| ppl = result.get("people_count", 0) | |
| if pose_tags or ppl: | |
| safe = [html.escape(p) for p in pose_tags] | |
| line = f"<div class='whyx-tagger-panel-title'>{t('tagger_pose_label', lc)}</div>" | |
| label = t("tagger_pose_single" if ppl == 1 else "tagger_pose_many", lc) | |
| if ppl: | |
| line += f"<div class='whyx-tagger-pose-people'>{label.format(n=ppl)}</div>" | |
| if safe: | |
| line += f"<div class='whyx-tagger-pose-tags'>{', '.join(safe)}</div>" | |
| lines.append(line) | |
| lines.append(f"<div class='whyx-tagger-panel-title'>{t('tagger_ratings', lc)}</div>") | |
| rating_colors = {"general": "#34D399", "sensitive": "#FBBF24", "questionable": "#F97316", "explicit": "#EF4444"} | |
| for name, score in result["ratings"].items(): | |
| color = rating_colors.get(name, "var(--text-dim)") | |
| pct = int(score * 100) | |
| bar = "▰" * (pct // 10) + "▱" * (10 - pct // 10) | |
| lines.append(f"<div class='whyx-tagger-rating-row'><span class='whyx-tagger-rating-name' style='color:{color};'>{name}</span> <span class='whyx-tagger-rating-bar' style='color:{color};'>{bar}</span> <span class='whyx-tagger-rating-pct'>({score:.1%})</span></div>") | |
| if result["characters"]: | |
| lines.append(f"<div class='whyx-tagger-panel-title' style='margin-top:10px;'>{t('tagger_characters', lc)}</div>") | |
| lines.append("<div class='whyx-tagger-badges'>") | |
| for name, score in list(result["characters"].items())[:6]: | |
| safe = html.escape(name.replace("_", " ").replace("(", "\\(").replace(")", "\\)")) | |
| lines.append(f"<span class='whyx-tagger-char-badge'>{safe} <span class='whyx-tagger-badge-pct'>{score:.0%}</span></span>") | |
| lines.append("</div>") | |
| lines.append(f"<div class='whyx-tagger-panel-title' style='margin-top:10px;'>{t('tagger_top_general', lc)}</div>") | |
| lines.append("<div class='whyx-tagger-gen-list'>") | |
| for name, score in list(result["general"].items())[:14]: | |
| safe = html.escape(name.replace("_", " ").replace("(", "\\(").replace(")", "\\)")) | |
| pct = int(score * 100) | |
| bar = "▰" * (pct // 10) + "▱" * (10 - pct // 10) | |
| lines.append(f"<div class='whyx-tagger-gen-row'><span class='whyx-tagger-gen-bar'>{bar}</span><span class='whyx-tagger-gen-name'>{safe}</span><span class='whyx-tagger-badge-pct'>{score:.0%}</span></div>") | |
| lines.append("</div>") | |
| lines.append("</div>") | |
| return "\n".join(lines) | |
| def on_tag_image(image, gen_threshold, char_threshold, fmt, lang, progress=gr.Progress()): | |
| lc = "ru" if lang == "RU" else "en" | |
| empty_chips = gr.update(choices=[], value=[], visible=False) | |
| if image is None: | |
| return ( | |
| gr.update(value=f'<div style="color:#F87171;font-size:12px;">{t("tagger_no_image", lc)}</div>'), | |
| gr.update(visible=False), | |
| empty_chips, empty_chips, | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| "", "", | |
| ) | |
| try: | |
| progress(0, desc=t("tagger_processing", lc)) | |
| from src.ensemble_tagger import get_ensemble_tagger | |
| ens = get_ensemble_tagger() | |
| if not _TAGGER_DEPS_OK: | |
| msg = f'<div style="color:#FBBF24;font-size:13px;line-height:1.5;">{t("tagger_unavailable", lc)}</div>' | |
| return ( | |
| gr.update(value=msg), | |
| gr.update(visible=False), | |
| empty_chips, empty_chips, | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| "", "", | |
| ) | |
| progress(0.35, desc=t("tagger_ratings_label", lc)) | |
| result = ens.tag_image(image, gen_threshold, char_threshold, mode="ensemble") | |
| progress(0.9, desc=t("tagger_results", lc)) | |
| if not result["general"] and not result["characters"]: | |
| html_out = f'<div style="color:#FBBF24;font-size:12px;">{t("tagger_no_tags", lc)}</div>' | |
| return ( | |
| gr.update(value=html_out), | |
| gr.update(visible=False), | |
| empty_chips, empty_chips, | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| "", "", | |
| ) | |
| html_out = _build_tagger_html(result, lc) | |
| general_names = list(result["general"].keys())[:60] | |
| char_names = list(result["characters"].keys()) | |
| general_choices = [(n.replace("_", " "), n) for n in general_names] | |
| char_choices = [(n.replace("_", " "), n) for n in char_names] | |
| raw_caption = result["caption"] | |
| esc_taglist = result["taglist"] | |
| tags_box = esc_taglist if fmt == "prompt" else raw_caption | |
| return ( | |
| gr.update(value=html_out), | |
| gr.update(visible=True), | |
| gr.update(choices=general_choices, value=general_names, visible=True), | |
| gr.update(choices=char_choices, value=char_names, visible=True), | |
| gr.update(value=tags_box, visible=True), | |
| gr.update(visible=True), | |
| raw_caption, | |
| esc_taglist, | |
| ) | |
| except Exception as exc: | |
| import traceback | |
| traceback.print_exc() | |
| err = t("tagger_error", lc).format(exc=html.escape(str(exc))) | |
| return ( | |
| gr.update(value=f'<div style="color:#F87171;font-size:12px;">{err}</div>'), | |
| gr.update(visible=False), | |
| empty_chips, empty_chips, | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| "", "", | |
| ) | |
| def _join_selected(chips, chars, fmt): | |
| selected = (list(chips or []) + list(chars or [])) | |
| if fmt == "raw": | |
| return ", ".join(to_booru_tag(t) for t in selected) | |
| return ", ".join(to_prompt_tag(t) for t in selected) | |
| def on_tagger_selection_change(chips, chars, fmt, raw_state, esc_state): | |
| return _join_selected(chips, chars, fmt) | |
| def on_tagger_format_change(fmt, chips, chars): | |
| return _join_selected(chips, chars, fmt) | |
| def on_tagger_apply_tags(current_prompt, tags): | |
| if tags and isinstance(tags, str) and tags.strip(): | |
| if current_prompt and current_prompt.strip(): | |
| return f"{current_prompt.strip()}, {tags.strip()}" | |
| return tags.strip() | |
| return current_prompt if current_prompt else "" | |
| __all__ = [ | |
| "MAX_OUTPUTS", "ALL_CATEGORIES", "CATEGORY_ICONS", "DEFAULT_CHECKED", | |
| "warehouse", "artist_styles", "all_artists", "all_tandems", | |
| "STYLECHOICES", "ARTISTCHOICES", "TANDEMCHOICES", "COPY_JS", | |
| "on_generate", "_preset_status", "on_user_preset_save", "on_user_preset_apply", | |
| "on_user_preset_delete", "on_web_search", "_format_history_html", "_heart_click", | |
| "on_preset_change", "_clear_preset_state", "on_rating_change", | |
| "on_artist_filter_change", "_format_artist_info", "on_artist_selection_change", | |
| "on_prompt_analyze", "_reload_data", "on_artist_style_recommendations", | |
| "_build_tagger_html", "on_tag_image", "_join_selected", | |
| "on_tagger_selection_change", "on_tagger_format_change", "on_tagger_apply_tags", | |
| ] | |