Auto-sync from GitHub Actions

api_fastapi.py

@@ -26,7 +26,6 @@ FastAPI 适配层（可选）。
 
 from __future__ import annotations
 
-import asyncio
 from fastapi import FastAPI
 from pydantic import BaseModel, Field
 
@@ -130,8 +129,7 @@ async def related(body: RelatedIn) -> list[RelatedTagOut]:
     - show_nsfw：是否包含 NSFW 标签，默认 True
     """
     tagger = await DanbooruTagger.get_instance()
-    results = await asyncio.to_thread(
-        tagger.get_related,
+    results = await tagger.get_related_async(
         body.tags,
         set(body.tags),   # exclude 已选标签自身
         body.limit,

core/engine.py

@@ -163,9 +163,9 @@ class DanbooruTagger:
 
     _instance: Optional['DanbooruTagger'] = None
     _lock: Optional[asyncio.Lock] = None
-    # 进程级搜索并发信号量：串行化 search()，避免多个 model.encode()
-    # 并发抢占 CPU 而拖垮事件循环。
-    _search_sem: Optional[asyncio.Semaphore] = None
+    # 进程级 CPU 并发闸门：串行化所有 CPU 密集型操作（search / get_related /
+    # get_group_candidates），避免并发抢占 CPU 而拖垮 asyncio 事件循环。
+    _cpu_sem: Optional[asyncio.Semaphore] = None
 
     @classmethod
     def is_ready(cls) -> bool:
@@ -346,10 +346,15 @@ class DanbooruTagger:
 
     # ── 搜索 ──────────────────────────────────────────────────────────────
 
-    def _encode_queries(self, queries: list[str]) -> torch.Tensor:
-        """批量编码查询词，命中 embedding 缓存的跳过 model.encode。"""
+    def _encode_queries(self, queries: list[str]) -> tuple[torch.Tensor, list[bool]]:
+        """批量编码查询词，命中 embedding 缓存的跳过 model.encode。
+
+        Returns:
+            (q_emb, hit_mask): 编码后的张量 (Q, D) 以及每个 query 是否命中缓存。
+        """
         cached_vecs: list[Optional[torch.Tensor]] = [self._emb_cache.get(q) for q in queries]
         uncached_idx = [i for i, v in enumerate(cached_vecs) if v is None]
+        hit_mask = [v is not None for v in cached_vecs]
 
         if uncached_idx:
             uncached_texts = [queries[i] for i in uncached_idx]
@@ -362,7 +367,7 @@ class DanbooruTagger:
                 self._emb_cache.put(queries[i], emb)
                 cached_vecs[i] = emb
 
-        return torch.stack(cached_vecs)  # type: ignore[arg-type]
+        return torch.stack(cached_vecs), hit_mask  # type: ignore[arg-type]
 
     def search(self, request: SearchRequest) -> SearchResponse:
         if not self.is_loaded:
@@ -396,7 +401,7 @@ class DanbooruTagger:
             extra_segments = []
             queries  = [request.query]
 
-        q_emb = self._encode_queries(queries)
+        q_emb, hit_mask = self._encode_queries(queries)
 
         tl    = request.target_layers
         k     = request.top_k
@@ -473,17 +478,23 @@ class DanbooruTagger:
         # 对每个候选标签，计算其与完整原始查询（而非分词片段）的语义相似度，
         # 将相似度作为软因子乘入 final_score，使仅由分词碎片匹配到的噪声
         # 标签自然下沉，同时不硬过滤任何结果。
+        # 批量矩阵乘法替代逐条 torch.dot，O(R*L) 降为 O(L) + O(R)。
         full_q = q_emb[0]          # queries[0] 始终为完整原始查询
-        alpha  = 0.3 if SearchRequest.use_segmentation else 0   # 一致性调节强度（0=不调节, 1=完全按一致性重排），仅在启用分词时有意义
-        for r in final.values():
-            idx = self._name_to_idx[r.tag]
-            max_co = 0.0
+        alpha  = 0.3 if request.use_segmentation else 0   # 一致性调节强度（0=不调节, 1=完全按一致性重排），仅在启用分词时有意义
+        if alpha > 0 and final:
+            tag_list = list(final.keys())
+            tag_indices = [self._name_to_idx[t] for t in tag_list]
+            idx_tensor = torch.tensor(tag_indices, dtype=torch.long, device=full_q.device)
+            max_co = torch.zeros(len(tag_indices), device=full_q.device)
             for ln, attr, _ in _LAYER_SPEC:
                 if ln not in tl:
                     continue
-                max_co = max(max_co, float(torch.dot(full_q, getattr(self, attr)[idx])))
-            # coherence=0 时最多扣 alpha=15%；coherence=1 时不扣分
-            r.final_score = round(r.final_score * (1.0 - alpha + alpha * max_co), 4)
+                emb_selected = getattr(self, attr)[idx_tensor]       # (R, D)
+                co = (full_q.unsqueeze(0) @ emb_selected.T).squeeze(0)  # (R,)
+                max_co = torch.maximum(max_co, co)
+            for i, tag in enumerate(tag_list):
+                r = final[tag]
+                r.final_score = round(r.final_score * (1.0 - alpha + alpha * float(max_co[i])), 4)
 
         # Group expand 处理（在 guaranteed_tags 之前，因为会改分数）
         if request.group_mode == "expand" and self._tag_to_groups:
@@ -537,27 +548,59 @@ class DanbooruTagger:
 
         tags_all = ', '.join(r.tag for r in valid)
         tags_sfw = ', '.join(r.tag for r in valid if r.nsfw != '1')
+        cached_queries = [q for q, hit in zip(queries, hit_mask) if hit]
         response = SearchResponse(
             tags_all=tags_all, tags_sfw=tags_sfw,
             results=valid, keywords=keywords, segments=extra_segments,
+            cached_queries=cached_queries,
         )
         self._search_cache.put(cache_key, response)
         return response
 
+    # ── CPU 并发闸门（类级信号量，所有 CPU 密集型操作共享）───────────────
+
     @classmethod
-    def _get_search_sem(cls) -> asyncio.Semaphore:
-        if cls._search_sem is None:
-            cls._search_sem = asyncio.Semaphore(1)
-        return cls._search_sem
+    def _get_cpu_sem(cls) -> asyncio.Semaphore:
+        if cls._cpu_sem is None:
+            cls._cpu_sem = asyncio.Semaphore(1)
+        return cls._cpu_sem
 
     async def search_async(self, request: SearchRequest) -> SearchResponse:
-        """search() 的并发安全异步封装：信号量串行化 + 线程池执行。
+        """search() 的并发安全异步封装：共享闸门串行化 + 线程池执行。
 
         所有异步入口（MCP / API / UI）都应改用本方法，而非各自
-        asyncio.to_thread(self.search)，以共享同一个并发闸门。
+        asyncio.to_thread(self.search)，以共享同一个 CPU 并发闸门。
+        包含 60 秒超时，防止异常卡死导致信号量永久泄漏。
         """
-        async with self._get_search_sem():
-            return await asyncio.to_thread(self.search, request)
+        async with self._get_cpu_sem():
+            return await asyncio.wait_for(
+                asyncio.to_thread(self.search, request),
+                timeout=60.0,
+            )
+
+    async def get_related_async(
+        self,
+        seed_tags: list[str],
+        exclude: set[str] | None = None,
+        limit: int = 20,
+        show_nsfw: bool = True,
+    ) -> list:
+        """get_related() 的并发安全异步封装，共享同一个 CPU 闸门。"""
+        async with self._get_cpu_sem():
+            return await asyncio.to_thread(
+                self.get_related, seed_tags, exclude, limit, show_nsfw,
+            )
+
+    async def get_group_candidates_async(
+        self,
+        selected_tags: list[str],
+        show_nsfw: bool = True,
+    ) -> list[dict]:
+        """get_group_candidates() 的并发安全异步封装，共享同一个 CPU 闸门。"""
+        async with self._get_cpu_sem():
+            return await asyncio.to_thread(
+                self.get_group_candidates, selected_tags, show_nsfw,
+            )
 
     def _apply_group_expand(self, final: dict[str, TagResult]) -> None:
         """expand 模式：提升同 group 标签的分数。"""

core/models.py

@@ -60,4 +60,5 @@ class SearchResponse:
     tags_sfw: str
     results: list[TagResult]
     keywords: list[str]
-    segments: list[str] = field(default_factory=list)  # 分隔符切分后的原始从句级片段
+    segments: list[str] = field(default_factory=list)  # 分隔符切分后的原始从句级片段
+    cached_queries: list[str] = field(default_factory=list)  # 命中 emb 缓存的查询文本

mcp_server.py

@@ -64,8 +64,9 @@ Only supported for general, copyright, and character tag searches; **artists and
 - search_mode: Preset strategy. Pick the one that matches your intent.
     "full_scene"       — Full scene → prompt (e.g. "一个穿着白色水手服的少女在雨中奔跑")
     "concept_explore"  — Vague concept exploration, broad recall (e.g. "赛博朋克服装", "兔耳朵", "中国风汉服")
-    "subject_describe" — Describe a subject to find matching tags (e.g. "EVA中蓝发的驾驶员", "两侧有开口，前方有拉绳的运动短裤")
+    "subject_describe" — Describe **one** subject to find matching tags (e.g. "EVA中蓝发的驾驶员", "两侧有开口，前方有拉绳的运动短裤")
     "precise_lookup"   — Precise lookup / spell fix (e.g. "selafuku", "thighhigh")
+- HINT: In subject_describe mode, the tokenizer is disabled, and you can only describe one thing at a time. To search for multiple things at once, use concept_explore or full_scene.
 - category: Filter to a specific tag category. Default "all".
     "all"       — All (通用 + 版权 + 人物)
     "general"   — Visual attributes, clothing, pose, background, etc.
@@ -258,8 +259,7 @@ JSON array sorted by aggregated NPMI score (descending). Each result:
         elif t in corrections:
             corrected_tags.append(corrections[t])
 
-    results = await asyncio.to_thread(
-        tagger.get_related,
+    results = await tagger.get_related_async(
         corrected_tags,
         set(corrected_tags),
         limit,

ui_nicegui.py

@@ -56,6 +56,18 @@ class _SuppressMCPNoise(logging.Filter):
 
 logging.getLogger("uvicorn.error").addFilter(_SuppressMCPNoise())
 
+# suppress MCP OAuth discovery 404 noise (clients probing .well-known/oauth-authorization-server)
+class _SuppressOAuthNoise(logging.Filter):
+    _MARKER = ".well-known/oauth-authorization-server"
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        if self._MARKER in record.getMessage():
+            return False
+        return True
+
+logging.getLogger("uvicorn.access").addFilter(_SuppressOAuthNoise())
+logging.getLogger("nicegui").addFilter(_SuppressOAuthNoise())
+
 # ── 表格列定义 ─────────────────────────────────────────────────────────────────
 
 TABLE_COLUMNS = [
@@ -150,6 +162,9 @@ class DanbooruSearchUI:
         self.selected_chips_container = None  # 已选标签 chip 容器
         self.current_related: list = []
         self.chip_extra_selected: set = set()
+        # 去抖任务句柄（取消旧任务避免 CPU 洪峰）
+        self._debounce_related_task = None  # type: asyncio.Task | None
+        self._debounce_group_task = None    # type: asyncio.Task | None
 
         # tag -> prompt 权重，范围 [0.1, 1.9]，默认 1.0
         self.tag_weights: dict[str, float] = {}
@@ -1233,24 +1248,31 @@ class DanbooruSearchUI:
             # 分词筛选 chips
             self.current_filter_keyword = 'ALL'   # 新搜索默认选中"全部"
             self.keywords_container.clear()
+            cached_set = set(response.cached_queries) if response.cached_queries else set()
             with self.keywords_container:
                 ui.label('分词筛选:').classes('text-sm text-gray-500 font-bold mr-2')
                 ui.chip('全部', on_click=lambda: self._filter_by_source('ALL')) \
                     .props('color=primary text-color=white clickable')
                 use_seg = self.input_segment.value if self.input_segment else True
                 if use_seg:
-                    ui.chip('整句',
-                            on_click=lambda: self._filter_by_source(self.current_query_str)) \
-                        .props('color=grey-4 text-color=black clickable')
+                    whole = ui.chip('整句',
+                            on_click=lambda: self._filter_by_source(self.current_query_str))
+                    whole.props('color=grey-4 text-color=black clickable')
+                    if self.current_query_str in cached_set:
+                        whole.style('outline: 1px dashed rgba(128,128,128,0.3); outline-offset: 1px;')
                     # 从句级原始片段（分隔符切分后未 jieba 的长片段，区别于关键词）
                     for seg in response.segments:
-                        ui.chip(seg,
-                                on_click=lambda s=seg: self._filter_by_source(s)) \
-                            .props('color=blue-1 text-color=blue-8 clickable')
+                        sc = ui.chip(seg,
+                                on_click=lambda s=seg: self._filter_by_source(s))
+                        sc.props('color=blue-1 text-color=blue-8 clickable')
+                        if seg in cached_set:
+                            sc.style('outline: 1px dashed rgba(128,128,128,0.3); outline-offset: 1px;')
                     for kw in response.keywords:
-                        ui.chip(kw,
-                                on_click=lambda k=kw: self._filter_by_source(k)) \
-                            .props('color=grey-4 text-color=black clickable')
+                        kc = ui.chip(kw,
+                                on_click=lambda k=kw: self._filter_by_source(k))
+                        kc.props('color=grey-4 text-color=black clickable')
+                        if kw in cached_set:
+                            kc.style('outline: 1px dashed rgba(128,128,128,0.3); outline-offset: 1px;')
                 else:
                     ui.label('(分词已关闭)').classes('text-xs text-gray-400')
 
@@ -1423,25 +1445,31 @@ class DanbooruSearchUI:
             self._render_related_list(merged, show_nsfw)
 
     def _refresh_related_from_selection(self, selected_tags: list[str], show_nsfw: bool):
-        """仅刷新关联推荐列表。"""
+        """仅刷新关联推荐列表（300ms 去抖，避免快速勾选产生 CPU 洪峰）。"""
+        # 取消上次未执行的刷新
+        if self._debounce_related_task and not self._debounce_related_task.done():
+            self._debounce_related_task.cancel()
         async def _do():
+            await asyncio.sleep(0.3)
             if not selected_tags:
                 self._refresh_related([], show_nsfw)
                 return
             tagger = await DanbooruTagger.get_instance()
-            related = await run.io_bound(
-                tagger.get_related,
+            related = await tagger.get_related_async(
                 selected_tags,
                 set(selected_tags),
                 50,
                 show_nsfw,
             )
             self._refresh_related(related, show_nsfw)
-        asyncio.ensure_future(_do())
+        self._debounce_related_task = asyncio.ensure_future(_do())
 
     def _refresh_group_from_selection(self, selected_tags: list[str], show_nsfw: bool):
-        """仅刷新同类扩展区域。"""
+        """仅刷新同类扩展区域（300ms 去抖，避免快速勾选产生 CPU 洪峰）。"""
+        if self._debounce_group_task and not self._debounce_group_task.done():
+            self._debounce_group_task.cancel()
         async def _do():
+            await asyncio.sleep(0.3)
             if not selected_tags:
                 if self.group_expansion_container is not None:
                     self.group_expansion_container.clear()
@@ -1449,13 +1477,12 @@ class DanbooruSearchUI:
                         ui.label('请先搜索并勾选标签…').classes('text-sm text-gray-400 italic p-4')
                 return
             tagger = await DanbooruTagger.get_instance()
-            group_data = await run.io_bound(
-                tagger.get_group_candidates,
+            group_data = await tagger.get_group_candidates_async(
                 selected_tags,
                 show_nsfw,
             )
             self._render_group_expansion(group_data, selected_tags, show_nsfw)
-        asyncio.ensure_future(_do())
+        self._debounce_group_task = asyncio.ensure_future(_do())
 
     def _render_group_expansion(self, group_data: list, selected_tags: list[str], show_nsfw: bool):
         """渲染 Group 同类扩展区域。"""
@@ -1753,6 +1780,7 @@ if __name__ in {'__main__', '__mp_main__'}:
     async def head_root():
         return PlainTextResponse('')
 
+
     ui.run(
         host=host,
         port=port,
@@ -1761,3 +1789,4 @@ if __name__ in {'__main__', '__mp_main__'}:
         show=not is_cloud(),
         reconnect_timeout=120,
     )
+