app.py

import os, io, copy, json, requests, spaces, gradio as gr, numpy as np
import argparse, huggingface_hub, onnxruntime as rt, pandas as pd, traceback, tempfile, zipfile, re, ast, time
from datetime import datetime, timezone
from collections import defaultdict
from PIL import Image, ImageOps
from modules.booru import booru_gradio, on_select
from apscheduler.schedulers.background import BackgroundScheduler
from modules.classifyTags import classify_tags, process_tags
from modules.beautify_model import beautify_list, beautify_class
from modules.tag_enhancer import prompt_summarizer

os.environ['PYTORCH_ENABLE_MPS_FALLBACK']='1'
os.environ['OMP_NUM_THREADS'] = '8'  # Optimize CPU utilization? Test...
#os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

TITLE = "Multi-Tagger v1.3"
DESCRIPTION = """
Multi-Tagger is a versatile application for advanced image analysis and captioning. Supports <b>CUDA</b> and <b>CPU</b>.
"""

# 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"

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=.05);parser.add_argument('--score-general-threshold', type=float, default=.35);parser.add_argument('--score-character-threshold', type=float, default=.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):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();self.checkpoints=[('Start', self.start_time)]
	def checkpoint(self, label='Checkpoint'):now=time.perf_counter();self.checkpoints.append((label, now))
	def report(self, is_clear_checkpoints=True):
		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()
	def report_all(self):
		print('\n> Execution Time Report:');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();self.checkpoints=[('Start', self.start_time)]
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, content: str, directory: str, fileName: str) -> str:
        # Write the content to a file
        file_path = os.path.join(directory, fileName)
        if fileName.endswith('.json'):
            with open(file_path, 'w', encoding="utf-8") as file:
                file.write(content)
        else:
            with open(file_path, 'w+', encoding="utf-8") as file:
                file.write(content)
        return file_path
    def predict(
        self,
        gallery,
        model_repo,
        model_repo_2,
        general_thresh,
        general_mcut_enabled,
        character_thresh,
        character_mcut_enabled,
        characters_merge_enabled,
        beautify_model_repo,
        additional_tags_prepend,
        additional_tags_append,
        tag_results,
        progress=gr.Progress()
        ):
        # Clear tag_results before starting a new prediction
        tag_results.clear()
        gallery_len = len(gallery)
        print(f"Predict load model: {model_repo}, gallery length: {gallery_len}")
        timer = Timer()  # Create a timer
        progressRatio = 0.5 if beautify_model_repo else 1
        progressTotal = gallery_len + 1
        current_progress = 0
        # Initialize variables that need to be accessible throughout the function
        final_categorized_output = ""
        categorized_output_strings = []
        txt_infos = []
        output_dir = tempfile.mkdtemp()
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        self.load_model(model_repo)
        current_progress += progressRatio/progressTotal;
        progress(current_progress, desc="Initialize wd model finished")
        timer.checkpoint(f"Initialize wd model")
        if beautify_model_repo:
            print(f"Load model {beautify_model_repo}")
            beautify = beautify_class(beautify_model_repo, loadModel=True)
            current_progress += progressRatio/progressTotal;
            progress(current_progress, desc="Initialize beautify model finished")
            timer.checkpoint(f"Initialize beautify model")
        timer.report()
        # Dictionary to track counters for each filename
        name_counters = defaultdict(int)
        for idx, value in enumerate(gallery):
            try:
                image_path = value[0]
                image_name = os.path.splitext(os.path.basename(image_path))[0]
                # Increment the counter for the current name
                name_counters[image_name] += 1
                if name_counters[image_name] > 1:
                    image_name = f"{image_name}_{name_counters[image_name]:02d}"
                image = self.prepare_image(image_path)
                # Run first model
                print(f"Gallery {idx:02d}: Starting run first model ({model_repo})...")
                self.load_model(model_repo)
                input_name = self.model.get_inputs()[0].name
                label_name = self.model.get_outputs()[0].name
                preds = self.model.run([label_name], {input_name: image})[0]
                labels = list(zip(self.tag_names, preds[0].astype(float)))
                # Process first model results
                ratings_names = [labels[i] for i in self.rating_indexes]
                rating = dict(ratings_names)

                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_temp = mcut_threshold(general_probs)
                else:
                    general_thresh_temp = general_thresh
                general_res = [x for x in general_names if x[1] > general_thresh_temp]
                general_res = dict(general_res)
                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_temp = mcut_threshold(character_probs)
                    character_thresh_temp = max(0.15, character_thresh_temp)
                else:
                    character_thresh_temp = character_thresh
                character_res = [x for x in character_names if x[1] > character_thresh_temp]
                character_res = dict(character_res)
                # Collect tags from first model
                character_list_1 = list(character_res.keys())
                sorted_general_list_1 = sorted(general_res.items(), key=lambda x: x[1], reverse=True)
                sorted_general_list_1 = [x[0] for x in sorted_general_list_1]
                # Run second model if selected and different from first
                if model_repo_2 and model_repo_2 != model_repo:
                    print(f"Gallery {idx:02d}: Starting run second model ({model_repo_2})...")
                    self.load_model(model_repo_2)
                    preds_2 = self.model.run([label_name], {input_name: image})[0]
                    labels_2 = list(zip(self.tag_names, preds_2[0].astype(float)))
                    # Process second model results
                    general_names_2 = [labels_2[i] for i in self.general_indexes]
                    if general_mcut_enabled:
                        general_probs_2 = np.array([x[1] for x in general_names_2])
                        general_thresh_temp_2 = mcut_threshold(general_probs_2)
                    else:
                        general_thresh_temp_2 = general_thresh
                    general_res_2 = [x for x in general_names_2 if x[1] > general_thresh_temp_2]
                    general_res_2 = dict(general_res_2)
                    character_names_2 = [labels_2[i] for i in self.character_indexes]
                    if character_mcut_enabled:
                        character_probs_2 = np.array([x[1] for x in character_names_2])
                        character_thresh_temp_2 = mcut_threshold(character_probs_2)
                        character_thresh_temp_2 = max(0.15, character_thresh_temp_2)
                    else:
                        character_thresh_temp_2 = character_thresh
                    character_res_2 = [x for x in character_names_2 if x[1] > character_thresh_temp_2]
                    character_res_2 = dict(character_res_2)
                    # Collect tags from second model
                    character_list_2 = list(character_res_2.keys())
                    sorted_general_list_2 = sorted(general_res_2.items(), key=lambda x: x[1], reverse=True)
                    sorted_general_list_2 = [x[0] for x in sorted_general_list_2]
                    # Combine results from both models (+ remove duplicates)
                    combined_character_list = list(set(character_list_1 + character_list_2))
                    combined_general_list = list(set(sorted_general_list_1 + sorted_general_list_2))
                else:
                    # Only first model was used
                    combined_character_list = character_list_1
                    combined_general_list = sorted_general_list_1
                # Remove values from combined_character_list that already exist in combined_general_list
                combined_character_list = [item for item in combined_character_list if item not in combined_general_list]
                # Handle prepend/append tags
                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]
                # Remove values from combined_general_list that already exist in prepend_list or append_list
                if prepend_list:
                    combined_general_list = [item for item in combined_general_list if item not in prepend_list]
                if append_list:
                    combined_general_list = [item for item in combined_general_list if item not in append_list]
                combined_general_list = prepend_list + combined_general_list + append_list
                sorted_general_strings = ", ".join((combined_character_list if characters_merge_enabled else []) + combined_general_list).replace("(", "\\(").replace(")", "\\)")
                classified_tags, unclassified_tags = classify_tags(combined_general_list)
                # Create a single string of ALL categorized tags for the current image
                categorized_output_string = ', '.join([', '.join(tags) for tags in classified_tags.values()])
                categorized_output_strings.append(categorized_output_string)
                # Collect all categorized output strings into a single string
                final_categorized_output = ', '.join(categorized_output_strings).replace("(", "\\(").replace(")", "\\)")
                # Create a .txt file for "Output (string)" and "Categorized Output (string)"
                txt_content = f"Output (string): {sorted_general_strings}\nCategorized Output (string): {final_categorized_output}"
                txt_file = self.create_file(txt_content, output_dir, f"{image_name}_output.txt")
                txt_infos.append({"path": txt_file, "name": f"{image_name}_output.txt"})
                # Create a .json file for "Categorized (tags)"
                json_content = json.dumps(classified_tags, indent=4)
                json_file = self.create_file(json_content, output_dir, f"{image_name}_categorized_tags.json")
                txt_infos.append({"path": json_file, "name": f"{image_name}_categorized_tags.json"})
                # Save a copy of the uploaded image in PNG format
                image_path = value[0]
                image = Image.open(image_path)
                image.save(os.path.join(output_dir, f"{image_name}.png"), format="PNG")
                txt_infos.append({"path": os.path.join(output_dir, f"{image_name}.png"), "name": f"{image_name}.png"})
                current_progress += progressRatio/progressTotal;
                progress(current_progress, desc=f"image{idx:02d}, predict finished")
                timer.checkpoint(f"image{idx:02d}, predict finished")
                if beautify_model_repo:
                    print(f"Starting beautify...")
                    beautify_strings = beautify.beautify(sorted_general_strings)
                    # Handle potential None returns from beautify
                    if beautify_strings is None:
                        beautify_strings = "Beautify failed - see console logs"
                    else:
                        beautify_strings = re.sub(r"Title:", "", beautify_strings)
                        beautify_strings = re.sub(r"\n+", ",", beautify_strings)
                        beautify_strings = re.sub(r",,+", ",", beautify_strings)
                    sorted_general_strings += ",\n\n" + beautify_strings
                    current_progress += progressRatio/progressTotal;
                    progress(current_progress, desc=f"image{idx:02d}, beautify finished!")
                    timer.checkpoint(f"image{idx:02d}, beautify finished!")
                txt_file = self.create_file(sorted_general_strings, output_dir, image_name + ".txt")
                txt_infos.append({"path":txt_file, "name": image_name + ".txt"})

                # Store the result in tag_results using image_path as the key
                tag_results[image_path] = {
                    "strings": sorted_general_strings,
                    "strings2": categorized_output_string,  # Store the categorized output string here
                    "classified_tags": classified_tags,
                    "rating": rating,
                    "character_res": character_res,
                    "general_res": general_res,
                    "unclassified_tags": unclassified_tags,
                    "summarize_tags": ""  # Initialize as empty string
                }
                timer.report()
            except Exception as e:
                print(traceback.format_exc())
                print("Error predict: " + str(e))
        # Zip creation logic:
        download = []
        if txt_infos is not None and len(txt_infos) > 0:
            downloadZipPath = os.path.join(output_dir, "Multi-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)
        # End zip creation logic

        if beautify_model_repo:
            beautify.release_vram()
            del beautify
        progress(1, desc=f"Predict completed")
        timer.report_all()  # Print all recorded times
        print("Predict is complete.")
        # Make sure all required variables are returned with proper defaults
        if 'sorted_general_strings' not in locals():
            sorted_general_strings = ""
        if 'final_categorized_output' not in locals():
            final_categorized_output = ""
        if 'classified_tags' not in locals():
            classified_tags = {}
        if 'rating' not in locals():
            rating = {}
        if 'character_res' not in locals():
            character_res = {}
        if 'general_res' not in locals():
            general_res = {}
        if 'unclassified_tags' not in locals():
            unclassified_tags = []

        return download, sorted_general_strings, final_categorized_output, classified_tags, rating, character_res, general_res, unclassified_tags, tag_results
def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state: gr.SelectData):
    if not selected_state:
        return selected_state
    tag_result = { 
        "strings": "", 
        "strings2": "", 
        "classified_tags": "{}", 
        "rating": "", 
        "character_res": "", 
        "general_res": "", 
        "unclassified_tags": "{}", 
        "summarize_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["strings2"], tag_result["classified_tags"], tag_result["rating"], tag_result["character_res"], tag_result["general_res"], tag_result["unclassified_tags"], tag_result["summarize_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
	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)
	if selected_image in gallery:gallery.remove(selected_image)
	return gallery
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, 
]
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)
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."

css = """
#custom-gallery {--row-height: 180px;display: grid;grid-auto-rows: min-content;gap: 10px;}
#custom-gallery .thumbnail-item {height: var(--row-height);width: 100%;position: relative;overflow: hidden;border-radius: 8px;box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);transition: transform 0.2s ease,  box-shadow 0.2s ease;}
#custom-gallery .thumbnail-item:hover {transform: translateY(-3px);box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);}
#custom-gallery .thumbnail-item img {width: auto;height: 100%;max-width: 100%;max-height: var(--row-height);object-fit: contain;margin: 0 auto;display: block;}
#custom-gallery .thumbnail-item img.portrait {max-width: 100%;}
#custom-gallery .thumbnail-item img.landscape {max-height: 100%;}
.gallery-container {max-height: 500px;overflow-y: auto;padding-right: 0px;--size-80: 500px;}
.thumbnails {display: flex;position: absolute;bottom: 0;width: 120px;overflow-x: scroll;padding-top: 320px;padding-bottom: 280px;padding-left: 4px;flex-wrap: wrap;}
#custom-gallery .thumbnail-item img {width: auto;height: 100%;max-width: 100%;max-height: var(--row-height);object-fit: initial;width: fit-content;margin: 0px auto;display: block;}
"""
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(value=f"<p style='text-align: center;'>{DESCRIPTION}</p>")
    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=2,
                        show_share_button=False, 
                        interactive=True, 
                        height="auto",
                        label="Grid of images", 
                        preview=False, 
                        elem_id="custom-gallery"  # Added for custom styling
                    )
                with gr.Column(variant="panel"):
                    model_repo = gr.Dropdown(dropdown_list, value=EVA02_LARGE_MODEL_DSV3_REPO, label="1st Model", )
                    PLUS = "+?"
                    gr.Markdown(value=f"<p style='text-align: center;'>{PLUS}</p>")
                    model_repo_2 = gr.Dropdown([None] + dropdown_list, value=None, label="2nd Model (Optional)", info="Select another model for diversified results.", )
                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():
                    beautify_model_repo = gr.Dropdown([None] + beautify_list, value=None, label="Beautify Model", info="Use a model to describe or 'beautify' a single image into a readable English article.", )
                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, beautify_model_repo, additional_tags_prepend, additional_tags_append, ], variant="secondary", size="lg", )
                with gr.Row():
                    rating = gr.Label(label="Rating")
            with gr.Column(variant="panel"):
                download_file = gr.File(label="Download")  # 0
                character_res = gr.Label(label="Output (characters)")  # 1
                sorted_general_strings = gr.Textbox(label="Output", show_label=True, show_copy_button=True, lines=5)  # 2
                final_categorized_output = gr.Textbox(label="Categorized", info="If tagging multiple images and got long tags,  please select an image to display tags correctly.", show_label=True, show_copy_button=True, lines=5)  # 3
                pe_generate_btn = gr.Button(value="SUMMARIZE TAGS", size="lg", variant="primary") # 4
                summarize_tags = gr.Textbox(label="Summarized Tags", show_label=True, show_copy_button=True, lines=6)  # 5
                prompt_summarizer_model = gr.Radio(["Medium", "Long", "Flux"], label="Model Choice", value="Medium", info="Summarize your prompts with medium or long answers. It's recommended for Flux.") # 6
                categorized = gr.JSON(label="Categorized (tags) - JSON")  # 7
                general_res = gr.Label(label="Output (tags)")  # 8
                unclassified = gr.JSON(label="Unclassified (tags)")  # 9
                clear.add([download_file, sorted_general_strings, final_categorized_output, categorized, rating, character_res, general_res, unclassified, prompt_summarizer_model, summarize_tags, ])
            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, final_categorized_output, categorized, rating, character_res, general_res, unclassified, summarize_tags])
            # Event to remove a selected image from the gallery
            remove_button.click(remove_image_from_gallery, inputs=[gallery, selected_image], outputs=gallery)
            # Event to for the Prompt Beautify Button
            pe_generate_btn.click(lambda tags, model:prompt_summarizer('', '', tags, model)[0], inputs=[final_categorized_output, prompt_summarizer_model], outputs=[summarize_tags])
        submit.click(predictor.predict, inputs=[gallery, model_repo, model_repo_2, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, characters_merge_enabled, beautify_model_repo, additional_tags_prepend, additional_tags_append, tag_results, ], outputs=[download_file, sorted_general_strings, final_categorized_output, categorized, rating, character_res, general_res, unclassified, tag_results, ], )
        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, ],)
        gr.Markdown(NEXT_RESTART)
    with gr.Tab("Booru Image Fetcher"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("### ⚙️ Search Parameters")
                site = gr.Dropdown(label="Select Source", choices=["Gelbooru (Not working)", "Rule34", "Xbooru"], value="Xbooru")
                Tags = gr.Textbox(label="Tags (comma-separated)", placeholder="e.g. solo,  1girl,  1boy,  artist name,  character,  black hair,  granblue fantasy,  ...", lines=3)
                exclude_tags = gr.Textbox(label="Exclude Tags (comma-separated)", placeholder="e.g. animated,  watermark,  username,  ...", lines=3)
                score = gr.Number(label="Minimum Score", value=0)
                count = gr.Slider(label="Number of Images", minimum=1, maximum=20, step=1, value=1)
                Safe = gr.Checkbox(label="Include Safe", value=True)
                Questionable = gr.Checkbox(label="Include Questionable", value=True)
                Explicit = gr.Checkbox(label="Include Explicit (18+)", value=False)
                submit_btn = gr.Button("Fetch Images", variant="primary")
            with gr.Column():
                gr.Markdown("### 📄 Results")
                images_output = gr.Gallery(label="Images", columns=3, rows=2, object_fit="contain", height=500)
                tags_output = gr.Textbox(label="Tags", placeholder="Select an image to display tags", lines=6, show_copy_button=True)
                post_url_output = gr.Textbox(label="Post URL", lines=1, show_copy_button=True)
                image_url_output = gr.Textbox(label="Image URL", lines=1, show_copy_button=True)
        # State to store tags, URLs
        tags_state = gr.State([])
        post_url_state = gr.State([])
        image_url_state = gr.State([])
        submit_btn.click(fn=booru_gradio, inputs=[Tags, exclude_tags, score, count, Safe, Questionable, Explicit, site], outputs=[images_output, tags_state, post_url_state, image_url_state], )
        images_output.select(fn=on_select, inputs=[tags_state, post_url_state, image_url_state], outputs=[tags_output, post_url_output, image_url_output], )
    with gr.Tab(label="Misc"): 
        with gr.Row():
            with gr.Column(variant="panel"):
                input_tags = gr.Textbox(label="Input Tags", placeholder="1girl,  cat,  horns,  blue hair,  ...\nor\n? 1girl 1234567? cat 1234567? horns 1234567? blue hair 1234567? ...", lines=4)
                submit_button = gr.Button(value="SUBMIT", variant="primary", size="lg")
            with gr.Column(variant="panel"):
                categorized_string = gr.Textbox(label="Categorized (string)", show_label=True, show_copy_button=True, lines=8)
                categorized_json = gr.JSON(label="Categorized (tags) - JSON")
                submit_button.click(process_tags, inputs=[input_tags], outputs=[categorized_string, categorized_json])
            with gr.Column(variant="panel"):
                pe_generate_btn = gr.Button(value="SUMMARIZE TAGS", size="lg", variant="primary")
                summarize_tags = gr.Textbox(label="Summarized Tags", show_label=True, show_copy_button=True, lines=5)
                prompt_summarizer_model = gr.Radio(["Medium", "Long", "Flux"], label="Model Choice", value="Medium", info="Summarize your prompts with medium or long answers. It's recommended for Flux.")
                pe_generate_btn.click(lambda tags, model:prompt_summarizer('', '', tags, model)[0], inputs=[categorized_string, prompt_summarizer_model], outputs=[summarize_tags])
demo.queue(max_size=10).launch(show_error=True)