First commit

README.md

@@ -9,5 +9,3 @@ app_file: app.py
 pinned: false
 license: apache-2.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,138 @@
+import json
+
+import faiss
+import flax
+import gradio as gr
+import jax
+import numpy as np
+import pandas as pd
+import requests
+
+from Models.CLIP import CLIP
+
+
+def danbooru_id_to_url(image_id, selected_ratings, api_username="", api_key=""):
+    headers = {"User-Agent": "image_similarity_tool"}
+    ratings_to_letters = {
+        "General": "g",
+        "Sensitive": "s",
+        "Questionable": "q",
+        "Explicit": "e",
+    }
+
+    acceptable_ratings = [ratings_to_letters[x] for x in selected_ratings]
+
+    image_url = f"https://danbooru.donmai.us/posts/{image_id}.json"
+    if api_username != "" and api_key != "":
+        image_url = f"{image_url}?api_key={api_key}&login={api_username}"
+
+    r = requests.get(image_url, headers=headers)
+    if r.status_code != 200:
+        return None
+
+    content = json.loads(r.text)
+    image_url = content["large_file_url"] if "large_file_url" in content else None
+    image_url = image_url if content["rating"] in acceptable_ratings else None
+    return image_url
+
+
+class Predictor:
+    def __init__(self):
+        self.base_model = "wd-v1-4-convnext-tagger-v2"
+
+        with open(f"data/{self.base_model}/clip.msgpack", "rb") as f:
+            data = f.read()
+
+        self.params = flax.serialization.msgpack_restore(data)["model"]
+        self.model = CLIP()
+
+        self.tags_df = pd.read_csv("data/selected_tags.csv")
+
+        self.images_ids = np.load("index/cosine_ids.npy")
+
+        self.knn_index = faiss.read_index("index/cosine_knn.index")
+
+        config = json.loads(open("index/cosine_infos.json").read())["index_param"]
+        faiss.ParameterSpace().set_index_parameters(self.knn_index, config)
+
+    def predict(self, positive_tags, negative_tags, n_neighbours=5):
+        tags_df = self.tags_df
+        model = self.model
+
+        num_classes = len(tags_df)
+
+        positive_tags = positive_tags.split(",")
+        negative_tags = negative_tags.split(",")
+
+        positive_tags_idxs = tags_df.index[tags_df["name"].isin(positive_tags)].tolist()
+        negative_tags_idxs = tags_df.index[tags_df["name"].isin(negative_tags)].tolist()
+
+        tags = np.zeros((1, num_classes), dtype=np.float32)
+        tags[0][positive_tags_idxs] = 1
+        emb_from_logits = model.apply(
+            {"params": self.params},
+            tags,
+            method=model.encode_text,
+        )
+        emb_from_logits = jax.device_get(emb_from_logits)
+
+        if len(negative_tags_idxs) > 0:
+            tags = np.zeros((1, num_classes), dtype=np.float32)
+            tags[0][negative_tags_idxs] = 1
+
+            neg_emb_from_logits = model.apply(
+                {"params": self.params},
+                tags,
+                method=model.encode_text,
+            )
+            neg_emb_from_logits = jax.device_get(neg_emb_from_logits)
+            emb_from_logits = emb_from_logits - neg_emb_from_logits
+
+        faiss.normalize_L2(emb_from_logits)
+
+        dists, indexes = self.knn_index.search(emb_from_logits, k=n_neighbours)
+        neighbours_ids = self.images_ids[indexes][0]
+        neighbours_ids = [int(x) for x in neighbours_ids]
+
+        captions = []
+        image_urls = []
+        for image_id, dist in zip(neighbours_ids, dists[0]):
+            current_url = danbooru_id_to_url(
+                image_id,
+                [
+                    "General",
+                    "Sensitive",
+                    "Questionable",
+                    "Explicit",
+                ],
+            )
+            if current_url is not None:
+                image_urls.append(current_url)
+                captions.append(f"{image_id}/{dist:.2f}")
+        return list(zip(image_urls, captions))
+
+
+def main():
+    predictor = Predictor()
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            positive_tags = gr.Textbox(label="Positive tags")
+            negative_tags = gr.Textbox(label="Negative tags")
+
+        find_btn = gr.Button("Find similar images")
+
+        similar_images = gr.Gallery(label="Similar images", columns=[5])
+
+        find_btn.click(
+            fn=predictor.predict,
+            inputs=[positive_tags, negative_tags],
+            outputs=[similar_images],
+        )
+
+    demo.queue()
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()

data/wd-v1-4-convnext-tagger-v2/clip.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3be3b97824313f01d9f1d74c43e441199b7ea485f5698d2008739f34c3e41200
+size 48689306

index/cosine_ids.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df724519c8c1981e49d80e2430261deb4fb6edf6d9c04e134427879710747394
+size 21830676

index/cosine_infos.json

@@ -0,0 +1 @@
+{"index_key": "OPQ256_1280,IVF16384_HNSW32,PQ256x8", "index_param": "nprobe=16,efSearch=32,ht=2048", "index_path": "/home/SmilingWolf/eval/index/ConvNextBV1_01_14_2023_08h37m46s_cosine_knn.index", "size in bytes": 1535843672, "avg_search_speed_ms": 10.164478485783887, "99p_search_speed_ms": 12.419190758373587, "reconstruction error %": 22.007358074188232, "nb vectors": 5457637, "vectors dimension": 1024, "compression ratio": 14.555180035276402}

index/cosine_knn.index

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a718ab8370df8b9d84002c55f945ef241e4cc3450d306c2ecd97661f51022ad
+size 1535843672

requirements.txt

@@ -0,0 +1,3 @@
+faiss
+jax[cpu]
+flax