app.py

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,
        selected_ratings,
        n_neighbours,
        api_username,
        api_key,
    ):
        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)
        faiss.normalize_L2(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)
            faiss.normalize_L2(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,
                selected_ratings,
                api_username,
                api_key,
            )
            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():
            with gr.Column():
                positive_tags = gr.Textbox(label="Positive tags")
                negative_tags = gr.Textbox(label="Negative tags")
                n_neighbours = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=5,
                    step=1,
                    label="# of images",
                )

            with gr.Column():
                api_username = gr.Textbox(label="Danbooru API Username")
                api_key = gr.Textbox(label="Danbooru API Key")
                selected_ratings = gr.CheckboxGroup(
                    choices=["General", "Sensitive", "Questionable", "Explicit"],
                    value=["General", "Sensitive"],
                    label="Ratings",
                )
        find_btn = gr.Button("Find similar images")

        similar_images = gr.Gallery(label="Similar images", columns=[5])

        examples = gr.Examples(
            [
                [
                    "marcille_donato",
                    "",
                    ["General", "Sensitive"],
                    5,
                    "",
                    "",
                ],
                [
                    "yellow_eyes,red_horns",
                    "",
                    ["General", "Sensitive"],
                    5,
                    "",
                    "",
                ],
                [
                    "artoria_pendragon_(fate),solo",
                    "excalibur_(fate/stay_night),green_eyes,monochrome,blonde_hair",
                    ["General", "Sensitive"],
                    5,
                    "",
                    "",
                ],
            ],
            inputs=[
                positive_tags,
                negative_tags,
                selected_ratings,
                n_neighbours,
                api_username,
                api_key,
            ],
            outputs=[similar_images],
            fn=predictor.predict,
            run_on_click=True,
            cache_examples=False,
        )

        find_btn.click(
            fn=predictor.predict,
            inputs=[
                positive_tags,
                negative_tags,
                selected_ratings,
                n_neighbours,
                api_username,
                api_key,
            ],
            outputs=[similar_images],
        )

    demo.queue()
    demo.launch()


if __name__ == "__main__":
    main()