Update app.py

app.py

@@ -1,1103 +1,1103 @@
-import os
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-import io
-import copy
-import requests
-import numpy as np
-import spaces
-import gradio as gr
-from transformers import AutoProcessor, AutoModelForCausalLM
-from transformers import AutoModelForCausalLM, AutoProcessor
-from transformers.dynamic_module_utils import get_imports
-from PIL import Image, ImageDraw, ImageFont 
-import matplotlib.pyplot as plt
-import matplotlib.patches as patches
-from unittest.mock import patch
-
-import argparse
-import huggingface_hub
-import onnxruntime as rt
-import pandas as pd
-import traceback
-import tempfile
-import zipfile
-import re
-import ast
-import time
-from datetime import datetime, timezone
-from collections import defaultdict
-from classifyTags import classify_tags
-# Add scheduler code here
-from apscheduler.schedulers.background import BackgroundScheduler
-
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
-    """Work around for https://huggingface.co/microsoft/phi-1_5/discussions/72."""
-    if not str(filename).endswith("/modeling_florence2.py"):
-        return get_imports(filename)
-    imports = get_imports(filename)
-    if "flash_attn" in imports:
-        imports.remove("flash_attn")
-    return imports
-
-@spaces.GPU
-def get_device_type():
-    import torch
-    if torch.cuda.is_available():  
-        return "cuda"
-    else: 
-        if (torch.backends.mps.is_available() and torch.backends.mps.is_built()):
-            return "mps"
-        else:
-            return "cpu"
-            
-model_id = 'MiaoshouAI/Florence-2-base-PromptGen-v2.0'
-
-import subprocess
-device = get_device_type()
-if (device == "cuda"):
-    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-    model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
-    processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
-    model.to(device)
-else:
-    #https://huggingface.co/microsoft/Florence-2-base-ft/discussions/4
-    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
-        model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
-        processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
-        model.to(device)
-
-TITLE = "Multi-Tagger"
-DESCRIPTION = """
-Multi-Tagger is a powerful and versatile application that integrates two cutting-edge models: Waifu Diffusion and Florence 2. This app is designed to provide comprehensive image analysis and captioning capabilities, making it a valuable tool for AI artists, researchers, and enthusiasts.
-
-Features:
-- Supports batch processing of multiple images.
-- Tags images with multiple categories: general tags, character tags, and ratings.
-- Tags are categorized into groups (e.g., general, characters, ratings).
-- Displays categorized tags in a structured format.
-- Integrates Llama3 models to reorganize the tags into a readable English article.
-- Includes a separate tab for image captioning using Florence 2.
-- Florence 2 supports CUDA, MPS or CPU if one of them is available.
-- Supports various captioning tasks (e.g., Caption, Detailed Caption, Object Detection).
-- Displays output text and images for tasks that generate visual outputs.
-- The space will restart every 2 days to ensure stability and performance. It uses a background scheduler to handle the restart process.
-
-Example image by [me.](https://huggingface.co/Werli)
-"""
-colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
-            'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']
-
-# 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)
-        # New code to create categorized output string
-        categorized_output_strings = []
-        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)
-                
-                # Create a single string of all categorized tags
-                categorized_output_string = ', '.join([', '.join(tags) for tags in classified_tags.values()])
-                categorized_output_strings.append(categorized_output_string)
-
-                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.")
-        
-        # Collect all categorized output strings into a single string
-        final_categorized_output = ', '.join(categorized_output_strings)
-        
-        return download, sorted_general_strings, classified_tags, rating, character_res, general_res, unclassified_tags, tag_results, final_categorized_output
-# END
-
-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
-
-# END
-
-def fig_to_pil(fig):
-    buf = io.BytesIO()
-    fig.savefig(buf, format='png')
-    buf.seek(0)
-    return Image.open(buf)
-
-@spaces.GPU
-def run_example(task_prompt, image, text_input=None):
-    if text_input is None:
-        prompt = task_prompt
-    else:
-        prompt = task_prompt + text_input
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
-    generated_ids = model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=1024,
-        early_stopping=False,
-        do_sample=False,
-        num_beams=3,
-    )
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = processor.post_process_generation(
-        generated_text,
-        task=task_prompt,
-        image_size=(image.width, image.height)
-    )
-    return parsed_answer
-
-def plot_bbox(image, data):
-    fig, ax = plt.subplots()
-    ax.imshow(image)
-    for bbox, label in zip(data['bboxes'], data['labels']):
-        x1, y1, x2, y2 = bbox
-        rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none')
-        ax.add_patch(rect)
-        plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
-    ax.axis('off')
-    return fig
-
-def draw_polygons(image, prediction, fill_mask=False):
-    draw = ImageDraw.Draw(image)
-    scale = 1
-    for polygons, label in zip(prediction['polygons'], prediction['labels']):
-        color = random.choice(colormap)
-        fill_color = random.choice(colormap) if fill_mask else None
-        for _polygon in polygons:
-            _polygon = np.array(_polygon).reshape(-1, 2)
-            if len(_polygon) < 3:
-                print('Invalid polygon:', _polygon)
-                continue
-            _polygon = (_polygon * scale).reshape(-1).tolist()
-            if fill_mask:
-                draw.polygon(_polygon, outline=color, fill=fill_color)
-            else:
-                draw.polygon(_polygon, outline=color)
-            draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color)
-    return image
-
-def convert_to_od_format(data):
-    bboxes = data.get('bboxes', [])
-    labels = data.get('bboxes_labels', [])
-    od_results = {
-        'bboxes': bboxes,
-        'labels': labels
-    }
-    return od_results
-
-def draw_ocr_bboxes(image, prediction):
-    scale = 1
-    draw = ImageDraw.Draw(image)
-    bboxes, labels = prediction['quad_boxes'], prediction['labels']
-    for box, label in zip(bboxes, labels):
-        color = random.choice(colormap)
-        new_box = (np.array(box) * scale).tolist()
-        draw.polygon(new_box, width=3, outline=color)
-        draw.text((new_box[0]+8, new_box[1]+2),
-                  "{}".format(label),
-                  align="right",
-                  fill=color)
-    return image
-
-def convert_to_od_format(data):
-    bboxes = data.get('bboxes', [])
-    labels = data.get('bboxes_labels', [])
-    od_results = {
-        'bboxes': bboxes,
-        'labels': labels
-    }
-    return od_results
-
-def draw_ocr_bboxes(image, prediction):
-    scale = 1
-    draw = ImageDraw.Draw(image)
-    bboxes, labels = prediction['quad_boxes'], prediction['labels']
-    for box, label in zip(bboxes, labels):
-        color = random.choice(colormap)
-        new_box = (np.array(box) * scale).tolist()
-        draw.polygon(new_box, width=3, outline=color)
-        draw.text((new_box[0]+8, new_box[1]+2),
-                  "{}".format(label),
-                  align="right",
-                  fill=color)
-    return image
-def process_image(image, task_prompt, text_input=None):
-    # Test
-    if isinstance(image, str):  # If image is a file path
-        image = Image.open(image)  # Load image from file path
-    else:  # If image is a NumPy array
-        image = Image.fromarray(image)  # Convert NumPy array to PIL Image
-    if task_prompt == 'Caption':
-        task_prompt = '<CAPTION>'
-        results = run_example(task_prompt, image)
-        return results[task_prompt], None
-    elif task_prompt == 'Detailed Caption':
-        task_prompt = '<DETAILED_CAPTION>'
-        results = run_example(task_prompt, image)
-        return results[task_prompt], None
-    elif task_prompt == 'More Detailed Caption':
-        task_prompt = '<MORE_DETAILED_CAPTION>'
-        results = run_example(task_prompt, image)
-        return results[task_prompt], plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
-    elif task_prompt == 'Caption + Grounding':
-        task_prompt = '<CAPTION>'
-        results = run_example(task_prompt, image)
-        text_input = results[task_prompt]
-        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
-        results = run_example(task_prompt, image, text_input)
-        results['<CAPTION>'] = text_input
-        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'Detailed Caption + Grounding':
-        task_prompt = '<DETAILED_CAPTION>'
-        results = run_example(task_prompt, image)
-        text_input = results[task_prompt]
-        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
-        results = run_example(task_prompt, image, text_input)
-        results['<DETAILED_CAPTION>'] = text_input
-        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'More Detailed Caption + Grounding':
-        task_prompt = '<MORE_DETAILED_CAPTION>'
-        results = run_example(task_prompt, image)
-        text_input = results[task_prompt]
-        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
-        results = run_example(task_prompt, image, text_input)
-        results['<MORE_DETAILED_CAPTION>'] = text_input
-        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'Object Detection':
-        task_prompt = '<OD>'
-        results = run_example(task_prompt, image)
-        fig = plot_bbox(image, results['<OD>'])
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'Dense Region Caption':
-        task_prompt = '<DENSE_REGION_CAPTION>'
-        results = run_example(task_prompt, image)
-        fig = plot_bbox(image, results['<DENSE_REGION_CAPTION>'])
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'Region Proposal':
-        task_prompt = '<REGION_PROPOSAL>'
-        results = run_example(task_prompt, image)
-        fig = plot_bbox(image, results['<REGION_PROPOSAL>'])
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'Caption to Phrase Grounding':
-        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
-        results = run_example(task_prompt, image, text_input)
-        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'Referring Expression Segmentation':
-        task_prompt = '<REFERRING_EXPRESSION_SEGMENTATION>'
-        results = run_example(task_prompt, image, text_input)
-        output_image = copy.deepcopy(image)
-        output_image = draw_polygons(output_image, results['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True)
-        return results, output_image
-    elif task_prompt == 'Region to Segmentation':
-        task_prompt = '<REGION_TO_SEGMENTATION>'
-        results = run_example(task_prompt, image, text_input)
-        output_image = copy.deepcopy(image)
-        output_image = draw_polygons(output_image, results['<REGION_TO_SEGMENTATION>'], fill_mask=True)
-        return results, output_image
-    elif task_prompt == 'Open Vocabulary Detection':
-        task_prompt = '<OPEN_VOCABULARY_DETECTION>'
-        results = run_example(task_prompt, image, text_input)
-        bbox_results = convert_to_od_format(results['<OPEN_VOCABULARY_DETECTION>'])
-        fig = plot_bbox(image, bbox_results)
-        return results, fig_to_pil(fig)
-    elif task_prompt == 'Region to Category':
-        task_prompt = '<REGION_TO_CATEGORY>'
-        results = run_example(task_prompt, image, text_input)
-        return results, None
-    elif task_prompt == 'Region to Description':
-        task_prompt = '<REGION_TO_DESCRIPTION>'
-        results = run_example(task_prompt, image, text_input)
-        return results, None
-    elif task_prompt == 'OCR':
-        task_prompt = '<OCR>'
-        results = run_example(task_prompt, image)
-        return results, None
-    elif task_prompt == 'OCR with Region':
-        task_prompt = '<OCR_WITH_REGION>'
-        results = run_example(task_prompt, image)
-        output_image = copy.deepcopy(image)
-        output_image = draw_ocr_bboxes(output_image, results['<OCR_WITH_REGION>'])
-        return results, output_image
-    else:
-        return "", None  # Return empty string and None for unknown task prompts
-##############
-# 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;
-#output {
-    height: 500px; 
-    overflow: auto; 
-    border: 1px solid #ccc; 
-}
-"""
-single_task_list =[
-    'Caption', 'Detailed Caption', 'More Detailed Caption', 'Object Detection',
-    'Dense Region Caption', 'Region Proposal', 'Caption to Phrase Grounding',
-    'Referring Expression Segmentation', 'Region to Segmentation',
-    'Open Vocabulary Detection', 'Region to Category', 'Region to Description',
-    'OCR', 'OCR with Region'
-]
-cascaded_task_list =[
-    'Caption + Grounding', 'Detailed Caption + Grounding', 'More Detailed Caption + Grounding'
-]
-def update_task_dropdown(choice):
-    if choice == 'Cascaded task':
-        return gr.Dropdown(choices=cascaded_task_list, value='Caption + Grounding')
-    else:
-        return gr.Dropdown(choices=single_task_list, value='Caption')
-
-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,
-]
-
-# This is workaround will make the space restart every 2 days. (for test).
-def _restart_space():
-    HF_TOKEN = os.getenv("HF_TOKEN")
-    if not HF_TOKEN:
-        raise ValueError("HF_TOKEN environment variable is not set.")
-    huggingface_hub.HfApi().restart_space(repo_id="Werli/Multi-Tagger", token=HF_TOKEN, factory_reboot=False)
-scheduler = BackgroundScheduler()
-# Add a job to restart the space every 2 days (172800 seconds)
-restart_space_job = scheduler.add_job(_restart_space, "interval", seconds=172800)
-# Start the scheduler
-scheduler.start()
-next_run_time_utc = restart_space_job.next_run_time.astimezone(timezone.utc)
-NEXT_RESTART = f"Next Restart: {next_run_time_utc.strftime('%Y-%m-%d %H:%M:%S')} (UTC)"
-
-# Using "reilnuud/polite" theme
-with gr.Blocks(title=TITLE, css=css, theme="Werli/wd-tagger-images", fill_width=True) as demo:
-    gr.Markdown(value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
-    gr.Markdown(value=DESCRIPTION)
-    gr.Markdown(NEXT_RESTART)
-    with gr.Tab(label="Waifu Diffusion"):
-        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_output = gr.Textbox(label="Categorized 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,  categorized_output,],
-        )  
-        gr.Examples(
-            [["images/1girl.png", VIT_LARGE_MODEL_DSV3_REPO, 0.35, False, 0.85, False]], 
-            inputs=[
-                image_input,
-                model_repo,
-                general_thresh,
-                general_mcut_enabled,
-                character_thresh,
-                character_mcut_enabled,
-            ],
-        )
-    with gr.Tab(label="Florence 2 Image Captioning"):
-        with gr.Row():
-            with gr.Column(variant="panel"):
-                input_img = gr.Image(label="Input Picture")
-                task_type = gr.Radio(choices=['Single task', 'Cascaded task'], label='Task type selector', value='Single task')
-                task_prompt = gr.Dropdown(choices=single_task_list, label="Task Prompt", value="Caption")
-                task_type.change(fn=update_task_dropdown, inputs=task_type, outputs=task_prompt)
-                text_input = gr.Textbox(label="Text Input (optional)")
-                submit_btn = gr.Button(value="Submit")
-            with gr.Column(variant="panel"):
-            #OUTPUT
-                output_text = gr.Textbox(label="Output Text", show_label=True, show_copy_button=True, lines=8) # Here is the problem!
-                output_img = gr.Image(label="Output Image")
-        gr.Examples(
-            examples=[
-                ["images/image1.png", 'Object Detection'],
-                ["images/image2.png", 'OCR with Region']
-            ],
-            inputs=[input_img, task_prompt],
-            outputs=[output_text, output_img],
-            fn=process_image,
-            cache_examples=False,
-            label='Try examples'
-        )
-        submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
-        
-demo.queue(max_size=2)
+import os
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+import io
+import copy
+import requests
+import numpy as np
+import spaces
+import gradio as gr
+from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers.dynamic_module_utils import get_imports
+from PIL import Image, ImageDraw, ImageFont 
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from unittest.mock import patch
+
+import argparse
+import huggingface_hub
+import onnxruntime as rt
+import pandas as pd
+import traceback
+import tempfile
+import zipfile
+import re
+import ast
+import time
+from datetime import datetime, timezone
+from collections import defaultdict
+from classifyTags import classify_tags
+# Add scheduler code here
+from apscheduler.schedulers.background import BackgroundScheduler
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
+    """Work around for https://huggingface.co/microsoft/phi-1_5/discussions/72."""
+    if not str(filename).endswith("/modeling_florence2.py"):
+        return get_imports(filename)
+    imports = get_imports(filename)
+    if "flash_attn" in imports:
+        imports.remove("flash_attn")
+    return imports
+
+@spaces.GPU
+def get_device_type():
+    import torch
+    if torch.cuda.is_available():  
+        return "cuda"
+    else: 
+        if (torch.backends.mps.is_available() and torch.backends.mps.is_built()):
+            return "mps"
+        else:
+            return "cpu"
+            
+model_id = 'MiaoshouAI/Florence-2-base-PromptGen-v2.0'
+
+import subprocess
+device = get_device_type()
+if (device == "cuda"):
+    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+    model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
+    processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
+    model.to(device)
+else:
+    #https://huggingface.co/microsoft/Florence-2-base-ft/discussions/4
+    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+        model = AutoModelForCausalLM.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
+        processor = AutoProcessor.from_pretrained("MiaoshouAI/Florence-2-base-PromptGen-v2.0", trust_remote_code=True)
+        model.to(device)
+
+TITLE = "Multi-Tagger"
+DESCRIPTION = """
+Multi-Tagger is a powerful and versatile application that integrates two cutting-edge models: Waifu Diffusion and Florence 2. This app is designed to provide comprehensive image analysis and captioning capabilities, making it a valuable tool for AI artists, researchers, and enthusiasts.
+
+Features:
+- Supports batch processing of multiple images.
+- Tags images with multiple categories: general tags, character tags, and ratings.
+- Tags are categorized into groups (e.g., general, characters, ratings).
+- Displays categorized tags in a structured format.
+- Integrates Llama3 models to reorganize the tags into a readable English article.
+- Includes a separate tab for image captioning using Florence 2.
+- Florence 2 supports CUDA, MPS or CPU if one of them is available.
+- Supports various captioning tasks (e.g., Caption, Detailed Caption, Object Detection).
+- Displays output text and images for tasks that generate visual outputs.
+- The space will restart every 2 days to ensure stability and performance. It uses a background scheduler to handle the restart process.
+
+Example image by [me.](https://huggingface.co/Werli)
+"""
+colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
+            'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']
+
+# 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)
+        # New code to create categorized output string
+        categorized_output_strings = []
+        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)
+                
+                # Create a single string of all categorized tags
+                categorized_output_string = ', '.join([', '.join(tags) for tags in classified_tags.values()])
+                categorized_output_strings.append(categorized_output_string)
+
+                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.")
+        
+        # Collect all categorized output strings into a single string
+        final_categorized_output = ', '.join(categorized_output_strings)
+        
+        return download, sorted_general_strings, classified_tags, rating, character_res, general_res, unclassified_tags, tag_results, final_categorized_output
+# END
+
+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
+
+# END
+
+def fig_to_pil(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png')
+    buf.seek(0)
+    return Image.open(buf)
+
+@spaces.GPU
+def run_example(task_prompt, image, text_input=None):
+    if text_input is None:
+        prompt = task_prompt
+    else:
+        prompt = task_prompt + text_input
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        early_stopping=False,
+        do_sample=False,
+        num_beams=3,
+    )
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = processor.post_process_generation(
+        generated_text,
+        task=task_prompt,
+        image_size=(image.width, image.height)
+    )
+    return parsed_answer
+
+def plot_bbox(image, data):
+    fig, ax = plt.subplots()
+    ax.imshow(image)
+    for bbox, label in zip(data['bboxes'], data['labels']):
+        x1, y1, x2, y2 = bbox
+        rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none')
+        ax.add_patch(rect)
+        plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
+    ax.axis('off')
+    return fig
+
+def draw_polygons(image, prediction, fill_mask=False):
+    draw = ImageDraw.Draw(image)
+    scale = 1
+    for polygons, label in zip(prediction['polygons'], prediction['labels']):
+        color = random.choice(colormap)
+        fill_color = random.choice(colormap) if fill_mask else None
+        for _polygon in polygons:
+            _polygon = np.array(_polygon).reshape(-1, 2)
+            if len(_polygon) < 3:
+                print('Invalid polygon:', _polygon)
+                continue
+            _polygon = (_polygon * scale).reshape(-1).tolist()
+            if fill_mask:
+                draw.polygon(_polygon, outline=color, fill=fill_color)
+            else:
+                draw.polygon(_polygon, outline=color)
+            draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color)
+    return image
+
+def convert_to_od_format(data):
+    bboxes = data.get('bboxes', [])
+    labels = data.get('bboxes_labels', [])
+    od_results = {
+        'bboxes': bboxes,
+        'labels': labels
+    }
+    return od_results
+
+def draw_ocr_bboxes(image, prediction):
+    scale = 1
+    draw = ImageDraw.Draw(image)
+    bboxes, labels = prediction['quad_boxes'], prediction['labels']
+    for box, label in zip(bboxes, labels):
+        color = random.choice(colormap)
+        new_box = (np.array(box) * scale).tolist()
+        draw.polygon(new_box, width=3, outline=color)
+        draw.text((new_box[0]+8, new_box[1]+2),
+                  "{}".format(label),
+                  align="right",
+                  fill=color)
+    return image
+
+def convert_to_od_format(data):
+    bboxes = data.get('bboxes', [])
+    labels = data.get('bboxes_labels', [])
+    od_results = {
+        'bboxes': bboxes,
+        'labels': labels
+    }
+    return od_results
+
+def draw_ocr_bboxes(image, prediction):
+    scale = 1
+    draw = ImageDraw.Draw(image)
+    bboxes, labels = prediction['quad_boxes'], prediction['labels']
+    for box, label in zip(bboxes, labels):
+        color = random.choice(colormap)
+        new_box = (np.array(box) * scale).tolist()
+        draw.polygon(new_box, width=3, outline=color)
+        draw.text((new_box[0]+8, new_box[1]+2),
+                  "{}".format(label),
+                  align="right",
+                  fill=color)
+    return image
+def process_image(image, task_prompt, text_input=None):
+    # Test
+    if isinstance(image, str):  # If image is a file path
+        image = Image.open(image)  # Load image from file path
+    else:  # If image is a NumPy array
+        image = Image.fromarray(image)  # Convert NumPy array to PIL Image
+    if task_prompt == 'Caption':
+        task_prompt = '<CAPTION>'
+        results = run_example(task_prompt, image)
+        return results[task_prompt], None
+    elif task_prompt == 'Detailed Caption':
+        task_prompt = '<DETAILED_CAPTION>'
+        results = run_example(task_prompt, image)
+        return results[task_prompt], None
+    elif task_prompt == 'More Detailed Caption':
+        task_prompt = '<MORE_DETAILED_CAPTION>'
+        results = run_example(task_prompt, image)
+        return results[task_prompt], plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+    elif task_prompt == 'Caption + Grounding':
+        task_prompt = '<CAPTION>'
+        results = run_example(task_prompt, image)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        results['<CAPTION>'] = text_input
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Detailed Caption + Grounding':
+        task_prompt = '<DETAILED_CAPTION>'
+        results = run_example(task_prompt, image)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        results['<DETAILED_CAPTION>'] = text_input
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'More Detailed Caption + Grounding':
+        task_prompt = '<MORE_DETAILED_CAPTION>'
+        results = run_example(task_prompt, image)
+        text_input = results[task_prompt]
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        results['<MORE_DETAILED_CAPTION>'] = text_input
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Object Detection':
+        task_prompt = '<OD>'
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<OD>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Dense Region Caption':
+        task_prompt = '<DENSE_REGION_CAPTION>'
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<DENSE_REGION_CAPTION>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Region Proposal':
+        task_prompt = '<REGION_PROPOSAL>'
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<REGION_PROPOSAL>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Caption to Phrase Grounding':
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Referring Expression Segmentation':
+        task_prompt = '<REFERRING_EXPRESSION_SEGMENTATION>'
+        results = run_example(task_prompt, image, text_input)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True)
+        return results, output_image
+    elif task_prompt == 'Region to Segmentation':
+        task_prompt = '<REGION_TO_SEGMENTATION>'
+        results = run_example(task_prompt, image, text_input)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results['<REGION_TO_SEGMENTATION>'], fill_mask=True)
+        return results, output_image
+    elif task_prompt == 'Open Vocabulary Detection':
+        task_prompt = '<OPEN_VOCABULARY_DETECTION>'
+        results = run_example(task_prompt, image, text_input)
+        bbox_results = convert_to_od_format(results['<OPEN_VOCABULARY_DETECTION>'])
+        fig = plot_bbox(image, bbox_results)
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Region to Category':
+        task_prompt = '<REGION_TO_CATEGORY>'
+        results = run_example(task_prompt, image, text_input)
+        return results, None
+    elif task_prompt == 'Region to Description':
+        task_prompt = '<REGION_TO_DESCRIPTION>'
+        results = run_example(task_prompt, image, text_input)
+        return results, None
+    elif task_prompt == 'OCR':
+        task_prompt = '<OCR>'
+        results = run_example(task_prompt, image)
+        return results, None
+    elif task_prompt == 'OCR with Region':
+        task_prompt = '<OCR_WITH_REGION>'
+        results = run_example(task_prompt, image)
+        output_image = copy.deepcopy(image)
+        output_image = draw_ocr_bboxes(output_image, results['<OCR_WITH_REGION>'])
+        return results, output_image
+    else:
+        return "", None  # Return empty string and None for unknown task prompts
+##############
+# 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;
+#output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+}
+"""
+single_task_list =[
+    'Caption', 'Detailed Caption', 'More Detailed Caption', 'Object Detection',
+    'Dense Region Caption', 'Region Proposal', 'Caption to Phrase Grounding',
+    'Referring Expression Segmentation', 'Region to Segmentation',
+    'Open Vocabulary Detection', 'Region to Category', 'Region to Description',
+    'OCR', 'OCR with Region'
+]
+cascaded_task_list =[
+    'Caption + Grounding', 'Detailed Caption + Grounding', 'More Detailed Caption + Grounding'
+]
+def update_task_dropdown(choice):
+    if choice == 'Cascaded task':
+        return gr.Dropdown(choices=cascaded_task_list, value='Caption + Grounding')
+    else:
+        return gr.Dropdown(choices=single_task_list, value='Caption')
+
+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,
+]
+
+# This is workaround will make the space restart every 2 days. (for test).
+def _restart_space():
+    HF_TOKEN = os.getenv("HF_TOKEN")
+    if not HF_TOKEN:
+        raise ValueError("HF_TOKEN environment variable is not set.")
+    huggingface_hub.HfApi().restart_space(repo_id="Werli/Multi-Tagger", token=HF_TOKEN, factory_reboot=False)
+scheduler = BackgroundScheduler()
+# Add a job to restart the space every 2 days (172800 seconds)
+restart_space_job = scheduler.add_job(_restart_space, "interval", seconds=172800)
+# Start the scheduler
+scheduler.start()
+next_run_time_utc = restart_space_job.next_run_time.astimezone(timezone.utc)
+NEXT_RESTART = f"Next Restart: {next_run_time_utc.strftime('%Y-%m-%d %H:%M:%S')} (UTC)"
+
+# Using "JohnSmith9982/small_and_pretty" theme
+with gr.Blocks(title=TITLE, css=css, theme="Werli/wd-tagger-images", fill_width=True) as demo:
+    gr.Markdown(value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
+    gr.Markdown(value=DESCRIPTION)
+    gr.Markdown(NEXT_RESTART)
+    with gr.Tab(label="Waifu Diffusion"):
+        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_output = gr.Textbox(label="Categorized 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,  categorized_output,],
+        )  
+        gr.Examples(
+            [["images/1girl.png", VIT_LARGE_MODEL_DSV3_REPO, 0.35, False, 0.85, False]], 
+            inputs=[
+                image_input,
+                model_repo,
+                general_thresh,
+                general_mcut_enabled,
+                character_thresh,
+                character_mcut_enabled,
+            ],
+        )
+    with gr.Tab(label="Florence 2 Image Captioning"):
+        with gr.Row():
+            with gr.Column(variant="panel"):
+                input_img = gr.Image(label="Input Picture")
+                task_type = gr.Radio(choices=['Single task', 'Cascaded task'], label='Task type selector', value='Single task')
+                task_prompt = gr.Dropdown(choices=single_task_list, label="Task Prompt", value="Caption")
+                task_type.change(fn=update_task_dropdown, inputs=task_type, outputs=task_prompt)
+                text_input = gr.Textbox(label="Text Input (optional)")
+                submit_btn = gr.Button(value="Submit")
+            with gr.Column(variant="panel"):
+            #OUTPUT
+                output_text = gr.Textbox(label="Output Text", show_label=True, show_copy_button=True, lines=8) # Here is the problem!
+                output_img = gr.Image(label="Output Image")
+        gr.Examples(
+            examples=[
+                ["images/image1.png", 'Object Detection'],
+                ["images/image2.png", 'OCR with Region']
+            ],
+            inputs=[input_img, task_prompt],
+            outputs=[output_text, output_img],
+            fn=process_image,
+            cache_examples=False,
+            label='Try examples'
+        )
+        submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
+        
+demo.queue(max_size=2)
 demo.launch(debug=True) # test