Add text support

README.md

@@ -1,13 +1,13 @@
 ---
-title: Search Anime Image By Image
+title: Search Anime Image By Image or Text
 emoji: 👁
 colorFrom: blue
 colorTo: pink
 sdk: gradio
-sdk_version: 5.6.0
+sdk_version: 5.16.0
 app_file: app.py
 pinned: true
 license: openrail
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -12,10 +12,11 @@ from cheesechaser.datapool import YandeWebpDataPool, ZerochanWebpDataPool, Gelbo
     KonachanWebpDataPool, AnimePicturesWebpDataPool, DanbooruNewestWebpDataPool, Rule34WebpDataPool
 from hfutils.operate import get_hf_fs, get_hf_client
 from hfutils.utils import TemporaryDirectory
-from imgutils.tagging import wd14
+from realutils.metrics import siglip
 
 from pools import quick_webp_pool
 
+siglip._REPO_ID = "SmilingWolf/swinv2_siglip_beta"
 _REPO_ID = 'deepghs/anime_sites_indices'
 
 hf_fs = get_hf_fs()
@@ -95,15 +96,24 @@ def _get_index_info(repo_id: str, model_name: str):
     return image_ids, knn_index
 
 
-def search(model_name: str, img_input, n_neighbours: int):
+def search(model_name: str, img_input, str_input: str, n_neighbours: int):
     images_ids, knn_index = _get_index_info(_REPO_ID, model_name)
-    embeddings = wd14.get_wd14_tags(
-        img_input,
-        model_name="SwinV2_v3",
-        fmt="embedding",
-    )
-    embeddings = np.expand_dims(embeddings, 0)
-    faiss.normalize_L2(embeddings)
+
+    if str_input == "":
+        embeddings = siglip.get_siglip_image_embedding(
+            img_input,
+            model_name="smilingwolf/siglip_swinv2_base_2025_02_08_13h25m57s",
+            fmt="embeddings",
+        )
+    else:
+        embeddings = siglip.get_siglip_text_embedding(
+            str_input,
+            model_name="smilingwolf/siglip_swinv2_base_2025_02_08_13h25m57s",
+            fmt="embeddings",
+        )
+
+    # In the model, the "embeddings" output node is already normalized.
+    # Ask for the "encodings" output if you want the raw logits
 
     dists, indexes = knn_index.search(embeddings, k=n_neighbours)
     neighbours_ids = images_ids[indexes][0]
@@ -123,8 +133,8 @@ if __name__ == "__main__":
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
-                img_input = gr.Image(type="pil", label="Input")
-
+                img_input = gr.Image(type="pil", label="Image input")
+                str_input = gr.Textbox(label="Text input (leave empty to use image input)")
             with gr.Column():
                 with gr.Row():
                     n_model = gr.Dropdown(
@@ -150,6 +160,7 @@ if __name__ == "__main__":
             inputs=[
                 n_model,
                 img_input,
+                str_input,
                 n_neighbours,
             ],
             outputs=[similar_images],

requirements.txt

@@ -2,5 +2,6 @@ pillow>=9.0.0
 faiss-cpu
 dghs-imgutils
 onnxruntime
-gradio==5.5.0
+gradio==5.16.0
 cheesechaser>=0.1.6
+dghs-realutils