app.py

import argparse
import os

import gradio as gr
import huggingface_hub
import numpy as np
import onnxruntime as rt
import pandas as pd
from PIL import Image
import traceback
import tempfile
import zipfile
import re
import ast
import time
from datetime import datetime
from collections import defaultdict
from classifyTags import classify_tags

TITLE = "WaifuDiffusion Tagger multiple images"
DESCRIPTION = """
Demo for the WaifuDiffusion tagger models  
Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)

Features of This Modified Version:
- Supports batch processing of multiple images
- Displays tag results in categorized groups: the generated tags will now be analyzed and categorized into corresponding groups.
"""

# Dataset v3 series of models:
SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"

# Dataset v2 series of models:
MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
SWIN_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
CONV_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"

# IdolSankaku series of models:
EVA02_LARGE_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-eva02-large-tagger-v1"
SWINV2_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-swinv2-tagger-v1"

# Files to download from the repos
MODEL_FILENAME = "model.onnx"
LABEL_FILENAME = "selected_tags.csv"

# LLAMA model
META_LLAMA_3_3B_REPO = "jncraton/Llama-3.2-3B-Instruct-ct2-int8"
META_LLAMA_3_8B_REPO = "avans06/Meta-Llama-3.2-8B-Instruct-ct2-int8_float16"

# https://github.com/toriato/stable-diffusion-webui-wd14-tagger/blob/a9eacb1eff904552d3012babfa28b57e1d3e295c/tagger/ui.py#L368
kaomojis = [
    "0_0",
    "(o)_(o)",
    "+_+",
    "+_-",
    "._.",
    "<o>_<o>",
    "<|>_<|>",
    "=_=",
    ">_<",
    "3_3",
    "6_9",
    ">_o",
    "@_@",
    "^_^",
    "o_o",
    "u_u",
    "x_x",
    "|_|",
    "||_||",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--score-slider-step", type=float, default=0.05)
    parser.add_argument("--score-general-threshold", type=float, default=0.35)
    parser.add_argument("--score-character-threshold", type=float, default=0.85)
    parser.add_argument("--share", action="store_true")
    return parser.parse_args()


def load_labels(dataframe) -> list[str]:
    name_series = dataframe["name"]
    name_series = name_series.map(
        lambda x: x.replace("_", " ") if x not in kaomojis else x
    )
    tag_names = name_series.tolist()

    rating_indexes = list(np.where(dataframe["category"] == 9)[0])
    general_indexes = list(np.where(dataframe["category"] == 0)[0])
    character_indexes = list(np.where(dataframe["category"] == 4)[0])
    return tag_names, rating_indexes, general_indexes, character_indexes


def mcut_threshold(probs):
    """
    Maximum Cut Thresholding (MCut)
    Largeron, C., Moulin, C., & Gery, M. (2012). MCut: A Thresholding Strategy
     for Multi-label Classification. In 11th International Symposium, IDA 2012
     (pp. 172-183).
    """
    sorted_probs = probs[probs.argsort()[::-1]]
    difs = sorted_probs[:-1] - sorted_probs[1:]
    t = difs.argmax()
    thresh = (sorted_probs[t] + sorted_probs[t + 1]) / 2
    return thresh

class Timer:
    def __init__(self):
        self.start_time  = time.perf_counter()  # Record the start time
        self.checkpoints = [("Start", self.start_time)]  # Store checkpoints

    def checkpoint(self, label="Checkpoint"):
        """Record a checkpoint with a given label."""
        now = time.perf_counter()
        self.checkpoints.append((label, now))

    def report(self, is_clear_checkpoints = True):
        # Determine the max label width for alignment
        max_label_length = max(len(label) for label, _ in self.checkpoints)

        prev_time = self.checkpoints[0][1]
        for label, curr_time in self.checkpoints[1:]:
            elapsed = curr_time - prev_time
            print(f"{label.ljust(max_label_length)}: {elapsed:.3f} seconds")
            prev_time = curr_time
        
        if is_clear_checkpoints:
            self.checkpoints.clear()
            self.checkpoint()  # Store checkpoints

    def report_all(self):
        """Print all recorded checkpoints and total execution time with aligned formatting."""
        print("\n> Execution Time Report:")

        # Determine the max label width for alignment
        max_label_length = max(len(label) for label, _ in self.checkpoints) if len(self.checkpoints) > 0 else 0

        prev_time = self.start_time
        for label, curr_time in self.checkpoints[1:]:
            elapsed = curr_time - prev_time
            print(f"{label.ljust(max_label_length)}: {elapsed:.3f} seconds")
            prev_time = curr_time
        
        total_time = self.checkpoints[-1][1] - self.start_time
        print(f"{'Total Execution Time'.ljust(max_label_length)}: {total_time:.3f} seconds\n")

        self.checkpoints.clear()

    def restart(self):
        self.start_time  = time.perf_counter()  # Record the start time
        self.checkpoints = [("Start", self.start_time)]  # Store checkpoints

class Llama3Reorganize:
    def __init__(
        self,
        repoId: str,
        device: str = None,
        loadModel: bool = False,
    ):
        """Initializes the Llama model.

        Args:
          repoId: LLAMA model repo.
          device: Device to use for computation (cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, 
            ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia).
          localFilesOnly:  If True, avoid downloading the file and return the path to the
            local cached file if it exists.
        """
        self.modelPath = self.download_model(repoId)

        if device is None:
            import torch
            self.totalVram = 0
            if torch.cuda.is_available():
                try:
                    deviceId = torch.cuda.current_device()
                    self.totalVram = torch.cuda.get_device_properties(deviceId).total_memory/(1024*1024*1024)
                except Exception as e:
                    print(traceback.format_exc())
                    print("Error detect vram: " + str(e))
                device = "cuda" if self.totalVram > (8 if "8B" in repoId else 4) else "cpu"
            else:
                device = "cpu"

        self.device = device
        self.system_prompt = "Please reorganize the following English labels describing a single image into a readable English article of no more than 300 words. If there are name labels, include them in the article. Since the target audience is not general, explicit content can be written in full without hesitation. No introduction is needed; directly reply with the English article:"

        if loadModel:
            self.load_model()

    def download_model(self, repoId):
        import warnings
        import requests
        allowPatterns = [
            "config.json",
            "generation_config.json",
            "model.bin",
            "pytorch_model.bin",
            "pytorch_model.bin.index.json",
            "pytorch_model-*.bin",
            "sentencepiece.bpe.model",
            "tokenizer.json",
            "tokenizer_config.json",
            "shared_vocabulary.txt",
            "shared_vocabulary.json",
            "special_tokens_map.json",
            "spiece.model",
            "vocab.json",
            "model.safetensors",
            "model-*.safetensors",
            "model.safetensors.index.json",
            "quantize_config.json",
            "tokenizer.model",
            "vocabulary.json",
            "preprocessor_config.json",
            "added_tokens.json"
        ]

        kwargs = {"allow_patterns": allowPatterns,}

        try:
            return huggingface_hub.snapshot_download(repoId, **kwargs)
        except (
            huggingface_hub.utils.HfHubHTTPError,
            requests.exceptions.ConnectionError,
        ) as exception:
            warnings.warn(
                "An error occured while synchronizing the model %s from the Hugging Face Hub:\n%s",
                repoId,
                exception,
            )
            warnings.warn(
                "Trying to load the model directly from the local cache, if it exists."
            )

            kwargs["local_files_only"] = True
            return huggingface_hub.snapshot_download(repoId, **kwargs)


    def load_model(self):
        import ctranslate2
        import transformers
        try:
            print('\n\nLoading model: %s\n\n' % self.modelPath)
            kwargsTokenizer = {"pretrained_model_name_or_path": self.modelPath}
            kwargsModel = {"device": self.device, "model_path": self.modelPath, "compute_type": "auto"}
            self.roleSystem = {"role": "system", "content": self.system_prompt}
            self.Model = ctranslate2.Generator(**kwargsModel)

            self.Tokenizer = transformers.AutoTokenizer.from_pretrained(**kwargsTokenizer)
            self.terminators = [self.Tokenizer.eos_token_id, self.Tokenizer.convert_tokens_to_ids("<|eot_id|>")]

        except Exception as e:
            self.release_vram()
            raise e
            

    def release_vram(self):
        try:
            import torch
            if torch.cuda.is_available():
                if getattr(self, "Model", None) is not None and getattr(self.Model, "unload_model", None) is not None:
                    self.Model.unload_model()
                    
                if getattr(self, "Tokenizer", None) is not None:
                    del self.Tokenizer
                if getattr(self, "Model", None) is not None:
                    del self.Model
                import gc
                gc.collect()
                try:
                    torch.cuda.empty_cache()
                except Exception as e:
                    print(traceback.format_exc())
                    print("\tcuda empty cache, error: " + str(e))
                print("release vram end.")
        except Exception as e:
            print(traceback.format_exc())
            print("Error release vram: " + str(e))

    def reorganize(self, text: str, max_length: int = 400):
        output = None
        result = None
        try:
            input_ids = self.Tokenizer.apply_chat_template([self.roleSystem, {"role": "user", "content": text + "\n\nHere's the reorganized English article:"}], tokenize=False, add_generation_prompt=True)
            source = self.Tokenizer.convert_ids_to_tokens(self.Tokenizer.encode(input_ids))
            output = self.Model.generate_batch([source], max_length=max_length, max_batch_size=2, no_repeat_ngram_size=3, beam_size=2, sampling_temperature=0.7, sampling_topp=0.9, include_prompt_in_result=False, end_token=self.terminators)
            target = output[0]
            result = self.Tokenizer.decode(target.sequences_ids[0])

            if len(result) > 2:
                if result[0] == "\"" and result[len(result) - 1] == "\"":
                    result = result[1:-1]
                elif result[0] == "'" and result[len(result) - 1] == "'":
                    result = result[1:-1]
                elif result[0] == "「" and result[len(result) - 1] == "」":
                    result = result[1:-1]
                elif result[0] == "『" and result[len(result) - 1] == "』":
                    result = result[1:-1]
        except Exception as e:
            print(traceback.format_exc())
            print("Error reorganize text: " + str(e))

        return result


class Predictor:
    def __init__(self):
        self.model_target_size = None
        self.last_loaded_repo = None

    def download_model(self, model_repo):
        csv_path = huggingface_hub.hf_hub_download(
            model_repo,
            LABEL_FILENAME,
        )
        model_path = huggingface_hub.hf_hub_download(
            model_repo,
            MODEL_FILENAME,
        )
        return csv_path, model_path

    def load_model(self, model_repo):
        if model_repo == self.last_loaded_repo:
            return

        csv_path, model_path = self.download_model(model_repo)

        tags_df = pd.read_csv(csv_path)
        sep_tags = load_labels(tags_df)

        self.tag_names = sep_tags[0]
        self.rating_indexes = sep_tags[1]
        self.general_indexes = sep_tags[2]
        self.character_indexes = sep_tags[3]

        model = rt.InferenceSession(model_path)
        _, height, width, _ = model.get_inputs()[0].shape
        self.model_target_size = height

        self.last_loaded_repo = model_repo
        self.model = model

    def prepare_image(self, path):
        image = Image.open(path)
        image = image.convert("RGBA")
        target_size = self.model_target_size

        canvas = Image.new("RGBA", image.size, (255, 255, 255))
        canvas.alpha_composite(image)
        image = canvas.convert("RGB")

        # Pad image to square
        image_shape = image.size
        max_dim = max(image_shape)
        pad_left = (max_dim - image_shape[0]) // 2
        pad_top = (max_dim - image_shape[1]) // 2

        padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
        padded_image.paste(image, (pad_left, pad_top))

        # Resize
        if max_dim != target_size:
            padded_image = padded_image.resize(
                (target_size, target_size),
                Image.BICUBIC,
            )

        # Convert to numpy array
        image_array = np.asarray(padded_image, dtype=np.float32)

        # Convert PIL-native RGB to BGR
        image_array = image_array[:, :, ::-1]

        return np.expand_dims(image_array, axis=0)

    def create_file(self, text: str, directory: str, fileName: str) -> str:
        # Write the text to a file
        with open(os.path.join(directory, fileName), 'w+', encoding="utf-8") as file:
            file.write(text)

        return file.name

    def predict(
        self,
        gallery,
        model_repo,
        general_thresh,
        general_mcut_enabled,
        character_thresh,
        character_mcut_enabled,
        characters_merge_enabled,
        llama3_reorganize_model_repo,
        additional_tags_prepend,
        additional_tags_append,
        tag_results,
        progress=gr.Progress()
    ):
        gallery_len = len(gallery)
        print(f"Predict load model: {model_repo}, gallery length: {gallery_len}")

        timer = Timer()  # Create a timer
        progressRatio = 0.5 if llama3_reorganize_model_repo else 1
        progressTotal = gallery_len + 1
        current_progress = 0

        self.load_model(model_repo)
        current_progress += progressRatio/progressTotal;
        progress(current_progress, desc="Initialize wd model finished")
        timer.checkpoint(f"Initialize wd model")

        # Result
        txt_infos = []
        output_dir = tempfile.mkdtemp()
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        sorted_general_strings = ""
        rating = None
        character_res = None
        general_res = None

        if llama3_reorganize_model_repo:
            print(f"Llama3 reorganize load model {llama3_reorganize_model_repo}")
            llama3_reorganize = Llama3Reorganize(llama3_reorganize_model_repo, loadModel=True)
            current_progress += progressRatio/progressTotal;
            progress(current_progress, desc="Initialize llama3 model finished")
            timer.checkpoint(f"Initialize llama3 model")
            
        timer.report()

        prepend_list = [tag.strip() for tag in additional_tags_prepend.split(",") if tag.strip()]
        append_list = [tag.strip() for tag in additional_tags_append.split(",") if tag.strip()]
        if prepend_list and append_list:
            append_list = [item for item in append_list if item not in prepend_list]
            
        # Dictionary to track counters for each filename
        name_counters = defaultdict(int)
        for idx, value in enumerate(gallery):
            try:
                image_path = value[0]
                image_name = os.path.splitext(os.path.basename(image_path))[0]

                # Increment the counter for the current name
                name_counters[image_name] += 1
                if name_counters[image_name] > 1:
                    image_name = f"{image_name}_{name_counters[image_name]:02d}"

                image = self.prepare_image(image_path)

                input_name = self.model.get_inputs()[0].name
                label_name = self.model.get_outputs()[0].name
                print(f"Gallery {idx:02d}: Starting run wd model...")
                preds = self.model.run([label_name], {input_name: image})[0]

                labels = list(zip(self.tag_names, preds[0].astype(float)))

                # First 4 labels are actually ratings: pick one with argmax
                ratings_names = [labels[i] for i in self.rating_indexes]
                rating = dict(ratings_names)

                # Then we have general tags: pick any where prediction confidence > threshold
                general_names = [labels[i] for i in self.general_indexes]

                if general_mcut_enabled:
                    general_probs = np.array([x[1] for x in general_names])
                    general_thresh = mcut_threshold(general_probs)

                general_res = [x for x in general_names if x[1] > general_thresh]
                general_res = dict(general_res)

                # Everything else is characters: pick any where prediction confidence > threshold
                character_names = [labels[i] for i in self.character_indexes]

                if character_mcut_enabled:
                    character_probs = np.array([x[1] for x in character_names])
                    character_thresh = mcut_threshold(character_probs)
                    character_thresh = max(0.15, character_thresh)

                character_res = [x for x in character_names if x[1] > character_thresh]
                character_res = dict(character_res)
                character_list = list(character_res.keys())

                sorted_general_list = sorted(
                    general_res.items(),
                    key=lambda x: x[1],
                    reverse=True,
                )
                sorted_general_list = [x[0] for x in sorted_general_list]
                #Remove values from character_list that already exist in sorted_general_list
                character_list = [item for item in character_list if item not in sorted_general_list]
                #Remove values from sorted_general_list that already exist in prepend_list or append_list
                if prepend_list:
                    sorted_general_list = [item for item in sorted_general_list if item not in prepend_list]
                if append_list:
                    sorted_general_list = [item for item in sorted_general_list if item not in append_list]

                sorted_general_list = prepend_list + sorted_general_list + append_list

                sorted_general_strings = ", ".join((character_list if characters_merge_enabled else []) + sorted_general_list).replace("(", "\(").replace(")", "\)")

                classified_tags, unclassified_tags = classify_tags(sorted_general_list)

                current_progress += progressRatio/progressTotal;
                progress(current_progress, desc=f"image{idx:02d}, predict finished")
                timer.checkpoint(f"image{idx:02d}, predict finished")
                
                if llama3_reorganize_model_repo:
                    print(f"Starting reorganize with llama3...")
                    reorganize_strings = llama3_reorganize.reorganize(sorted_general_strings)
                    reorganize_strings = re.sub(r" *Title: *", "", reorganize_strings)
                    reorganize_strings = re.sub(r"\n+", ",", reorganize_strings)
                    reorganize_strings = re.sub(r",,+", ",", reorganize_strings)
                    sorted_general_strings += "," + reorganize_strings

                    current_progress += progressRatio/progressTotal;
                    progress(current_progress, desc=f"image{idx:02d}, llama3 reorganize finished")
                    timer.checkpoint(f"image{idx:02d}, llama3 reorganize finished")

                txt_file = self.create_file(sorted_general_strings, output_dir, image_name + ".txt")
                txt_infos.append({"path":txt_file, "name": image_name + ".txt"})

                tag_results[image_path] = { "strings": sorted_general_strings, "classified_tags": classified_tags, "rating": rating, "character_res": character_res, "general_res": general_res, "unclassified_tags": unclassified_tags }
                timer.report()
            except Exception as e:
                print(traceback.format_exc())
                print("Error predict: " + str(e))

        # Result
        download = []
        if txt_infos is not None and len(txt_infos) > 0:
            downloadZipPath = os.path.join(output_dir, "images-tagger-" + datetime.now().strftime("%Y%m%d-%H%M%S") + ".zip")
            with zipfile.ZipFile(downloadZipPath, 'w', zipfile.ZIP_DEFLATED) as taggers_zip:
                for info in txt_infos:
                    # Get file name from lookup
                    taggers_zip.write(info["path"], arcname=info["name"])
            download.append(downloadZipPath)
            
        if llama3_reorganize_model_repo:
            llama3_reorganize.release_vram()
            del llama3_reorganize
            
        progress(1, desc=f"Predict completed")
        timer.report_all()  # Print all recorded times
        print("Predict is complete.")

        return download, sorted_general_strings, classified_tags, rating, character_res, general_res, unclassified_tags, tag_results
    
def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state: gr.SelectData):
    if not selected_state:
        return selected_state

    tag_result = { "strings": "", "classified_tags": "{}", "rating": "", "character_res": "", "general_res": "", "unclassified_tags": "{}" }
    if selected_state.value["image"]["path"] in tag_results:
        tag_result = tag_results[selected_state.value["image"]["path"]]

    return (selected_state.value["image"]["path"], selected_state.value["caption"]), tag_result["strings"], tag_result["classified_tags"], tag_result["rating"], tag_result["character_res"], tag_result["general_res"], tag_result["unclassified_tags"]

def append_gallery(gallery: list, image: str):
    if gallery is None:
        gallery = []
    if not image:
        return gallery, None
    
    gallery.append(image)

    return gallery, None


def extend_gallery(gallery: list, images):
    if gallery is None:
        gallery = []
    if not images:
        return gallery
    
    # Combine the new images with the existing gallery images
    gallery.extend(images)

    return gallery

def remove_image_from_gallery(gallery: list, selected_image: str):
    if not gallery or not selected_image:
        return gallery

    selected_image = ast.literal_eval(selected_image) #Use ast.literal_eval to parse text into a tuple.
    # Remove the selected image from the gallery
    if selected_image in gallery:
        gallery.remove(selected_image)
    return gallery


def main():
    # Custom CSS to set the height of the gr.Dropdown menu
    css = """
    div.progress-level div.progress-level-inner {
        text-align: left !important;
        width: 55.5% !important;
    }
    """

    args = parse_args()

    predictor = Predictor()

    dropdown_list = [
        EVA02_LARGE_MODEL_DSV3_REPO,
        SWINV2_MODEL_DSV3_REPO,
        CONV_MODEL_DSV3_REPO,
        VIT_MODEL_DSV3_REPO,
        VIT_LARGE_MODEL_DSV3_REPO,
        # ---
        MOAT_MODEL_DSV2_REPO,
        SWIN_MODEL_DSV2_REPO,
        CONV_MODEL_DSV2_REPO,
        CONV2_MODEL_DSV2_REPO,
        VIT_MODEL_DSV2_REPO,
        # ---
        SWINV2_MODEL_IS_DSV1_REPO,
        EVA02_LARGE_MODEL_IS_DSV1_REPO,
    ]

    llama_list = [
        META_LLAMA_3_3B_REPO,
        META_LLAMA_3_8B_REPO,
    ]

    with gr.Blocks(title=TITLE, css = css) as demo:
        gr.Markdown(
            value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>"
        )
        gr.Markdown(value=DESCRIPTION)
        with gr.Row():
            with gr.Column():
                submit = gr.Button(value="Submit", variant="primary", size="lg")
                with gr.Column(variant="panel"):
                    # Create an Image component for uploading images
                    image_input = gr.Image(label="Upload an Image or clicking paste from clipboard button", type="filepath", sources=["upload", "clipboard"], height=150)
                    with gr.Row():
                        upload_button = gr.UploadButton("Upload multiple images", file_types=["image"], file_count="multiple", size="sm")
                        remove_button = gr.Button("Remove Selected Image", size="sm")
                    gallery = gr.Gallery(columns=5, rows=5, show_share_button=False, interactive=True, height="500px", label="Gallery that displaying a grid of images")

                model_repo = gr.Dropdown(
                    dropdown_list,
                    value=EVA02_LARGE_MODEL_DSV3_REPO,
                    label="Model",
                )
                with gr.Row():
                    general_thresh = gr.Slider(
                        0,
                        1,
                        step=args.score_slider_step,
                        value=args.score_general_threshold,
                        label="General Tags Threshold",
                        scale=3,
                    )
                    general_mcut_enabled = gr.Checkbox(
                        value=False,
                        label="Use MCut threshold",
                        scale=1,
                    )
                with gr.Row():
                    character_thresh = gr.Slider(
                        0,
                        1,
                        step=args.score_slider_step,
                        value=args.score_character_threshold,
                        label="Character Tags Threshold",
                        scale=3,
                    )
                    character_mcut_enabled = gr.Checkbox(
                        value=False,
                        label="Use MCut threshold",
                        scale=1,
                    )
                with gr.Row():
                    characters_merge_enabled = gr.Checkbox(
                        value=True,
                        label="Merge characters into the string output",
                        scale=1,
                    )
                with gr.Row():
                    llama3_reorganize_model_repo = gr.Dropdown(
                        [None] + llama_list,
                        value=None,
                        label="Llama3 Model",
                        info="Use the Llama3 model to reorganize the article (Note: very slow)",
                    )
                with gr.Row():
                    additional_tags_prepend = gr.Text(label="Prepend Additional tags (comma split)")
                    additional_tags_append  = gr.Text(label="Append Additional tags (comma split)")
                with gr.Row():
                    clear = gr.ClearButton(
                        components=[
                            gallery,
                            model_repo,
                            general_thresh,
                            general_mcut_enabled,
                            character_thresh,
                            character_mcut_enabled,
                            characters_merge_enabled,
                            llama3_reorganize_model_repo,
                            additional_tags_prepend,
                            additional_tags_append,
                        ],
                        variant="secondary",
                        size="lg",
                    )
            with gr.Column(variant="panel"):
                download_file = gr.File(label="Output (Download)")
                sorted_general_strings = gr.Textbox(label="Output (string)", show_label=True, show_copy_button=True)
                categorized = gr.JSON(label="Categorized (tags)")
                rating = gr.Label(label="Rating")
                character_res = gr.Label(label="Output (characters)")
                general_res = gr.Label(label="Output (tags)")
                unclassified = gr.JSON(label="Unclassified (tags)")
                clear.add(
                    [
                        download_file,
                        sorted_general_strings,
                        categorized,
                        rating,
                        character_res,
                        general_res,
                        unclassified,
                    ]
                )
                
            tag_results = gr.State({})
            # Define the event listener to add the uploaded image to the gallery
            image_input.change(append_gallery, inputs=[gallery, image_input], outputs=[gallery, image_input])
            # When the upload button is clicked, add the new images to the gallery
            upload_button.upload(extend_gallery, inputs=[gallery, upload_button], outputs=gallery)
            # Event to update the selected image when an image is clicked in the gallery
            selected_image = gr.Textbox(label="Selected Image", visible=False)
            gallery.select(get_selection_from_gallery, inputs=[gallery, tag_results], outputs=[selected_image, sorted_general_strings, categorized, rating, character_res, general_res, unclassified])
            # Event to remove a selected image from the gallery
            remove_button.click(remove_image_from_gallery, inputs=[gallery, selected_image], outputs=gallery)

        submit.click(
            predictor.predict,
            inputs=[
                gallery,
                model_repo,
                general_thresh,
                general_mcut_enabled,
                character_thresh,
                character_mcut_enabled,
                characters_merge_enabled,
                llama3_reorganize_model_repo,
                additional_tags_prepend,
                additional_tags_append,
                tag_results,
            ],
            outputs=[download_file, sorted_general_strings, categorized, rating, character_res, general_res, unclassified, tag_results,],
        )
        
        gr.Examples(
            [["power.jpg", SWINV2_MODEL_DSV3_REPO, 0.35, False, 0.85, False]], 
            inputs=[
                image_input,
                model_repo,
                general_thresh,
                general_mcut_enabled,
                character_thresh,
                character_mcut_enabled,
            ],
        )

    demo.queue(max_size=2)
    demo.launch(inbrowser=True)


if __name__ == "__main__":
    main()