app.py

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.
- Displays categorized tags in a structured format.
- Includes a separate tab for image captioning using Florence 2. This model supports CUDA, MPS or CPU if one of them is available.
- Supports various captioning tasks (e.g., Caption, Detailed Caption, Object Detection), as well it can display output text and images for tasks that generate visual outputs.

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) - The space will restart every 2 days to ensure stability and performance. It uses a background scheduler to handle the restart process."

# Using "JohnSmith9982/small_and_pretty" theme
with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", 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