New features!

Big changes to the code!

Add a functionality to save and include some outputs to the zip file:
- Modified the `predict` method in the `Predictor` class to create and save:
- A `.txt` file containing "Output (string)" as well the "Categorized Output (string)".
- A `categorized_tag.json` file containing "Categorized (tags)".
- A copy of the uploaded image(s) in PNG format.
- Updated the `create_file` method to handle both text and JSON files.
- Ensured that all the created files are included in the zip file during the download process.

Also:
- Fixed the images grid position by creating a custom CSS, so it shows all the images correctly.
- Used a minify to clean some parts of the code without changing its functionality.
- Removed unnecessary comments and duplicate code.

app.py

@@ -1,56 +1,30 @@
 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
+import io,copy,requests,numpy as np,spaces,gradio as gr
+from transformers import AutoProcessor,AutoModelForCausalLM,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 PIL import Image,ImageDraw,ImageFont
+import matplotlib.pyplot as plt,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
+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 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
-
+import json
+os.environ['PYTORCH_ENABLE_MPS_FALLBACK']='1'
+
+def fixed_get_imports(filename:str|os.PathLike)->list[str]:
+	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"
-            
+	import torch
+	if torch.cuda.is_available():return'cuda'
+	elif 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
@@ -61,7 +35,6 @@ if (device == "cuda"):
     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)
@@ -75,13 +48,12 @@ 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 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.
+- Includes a separate tab for image captioning using Florence 2. Supports CUDA, MPS or CPU if one of them is available.
+- Supports various captioning tasks (e.g., Caption, Detailed Caption, Object Detection), 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']
+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"
@@ -89,277 +61,87 @@ 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
+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()  # 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
+	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 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 __init__(self,repoId:str,device:str=None,loadModel:bool=False):
+		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/1073741824
+				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,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
-            
+	import ctranslate2,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
-
+	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=.7,sampling_topp=.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):
@@ -412,7 +194,7 @@ class Predictor:
 
         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(
@@ -421,18 +203,21 @@ class Predictor:
             )
         # 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)
+    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.name
+        return file_path
 
     def predict(
         self,
@@ -465,14 +250,13 @@ class Predictor:
             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 = ""
-            # New code to create categorized output string
+            # Create categorized output string
             categorized_output_strings = []
             rating = None
             character_res = None
@@ -567,6 +351,22 @@ class Predictor:
                     # Collect all categorized output strings into a single string
                     final_categorized_output = ', '.join(categorized_output_strings)
                     
+                    # 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")
@@ -602,15 +402,16 @@ class Predictor:
                     print(traceback.format_exc())
                     print("Error predict: " + str(e))
             # Result
+            # Zip creation logic:
             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")
+                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 llama3_reorganize_model_repo:
                 llama3_reorganize.release_vram()
                 del llama3_reorganize
@@ -620,11 +421,9 @@ class Predictor:
             print("Predict is complete.")
             
             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": "", 
@@ -636,272 +435,80 @@ def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state:
     }
     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"]
-
-def append_gallery(gallery: list, image: str):
-    if gallery is None:
-        gallery = []
-    if not image:
-        return gallery, None
-    
-    gallery.append(image)
-
-    return gallery, None
-
-
-def extend_gallery(gallery: list, images):
-    if gallery is None:
-        gallery = []
-    if not images:
-        return gallery
-    
-    # Combine the new images with the existing gallery images
-    gallery.extend(images)
-
-    return gallery
-
-def remove_image_from_gallery(gallery: list, selected_image: str):
-    if not gallery or not selected_image:
-        return gallery
-
-    selected_image = ast.literal_eval(selected_image) # Use ast.literal_eval to parse text into a tuple.
-    # Remove the selected image from the gallery
-    if selected_image in gallery:
-        gallery.remove(selected_image)
-    return gallery
-
+def 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
 # END
 
-def fig_to_pil(fig):
-    buf = io.BytesIO()
-    fig.savefig(buf, format='png')
-    buf.seek(0)
-    return Image.open(buf)
-
+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, None
-    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
+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=.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):
+	if isinstance(image,str):image=Image.open(image)
+	else:image=Image.fromarray(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,None
+	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
+
+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']
 
-# 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')
@@ -909,7 +516,6 @@ def update_task_dropdown(choice):
         return gr.Dropdown(choices=single_task_list, value='Caption')
 
 args = parse_args()
-
 predictor = Predictor()
 
 dropdown_list = [
@@ -933,19 +539,23 @@ llama_list = [
     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()
+	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."
+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 = """
+div.progress-level div.progress-level-inner {text-align: left !important; width: 55.5% !important;}
+#output {height: 500px; overflow: auto; border: 1px solid #ccc;}
+label.float.svelte-i3tvor {position: relative !important;}
+.reduced-height.svelte-11chud3 {height: calc(80% - var(--size-10));}
+"""
 
 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>")
@@ -1030,7 +640,7 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                         size="lg",
                     )
             with gr.Column(variant="panel"):
-                download_file = gr.File(label="Output (Download)") # 0
+                download_file = gr.File(label="Output (Download All)") # 0
                 character_res = gr.Label(label="Output (characters)") # 1
                 sorted_general_strings = gr.Textbox(label="Output (string)", show_label=True, show_copy_button=True) # 2
                 final_categorized_output = gr.Textbox(label="Categorized Output (string)", show_label=True, show_copy_button=True) # 3
@@ -1112,5 +722,4 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
             label='Try examples'
         )
         submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
-        
 demo.queue(max_size=2).launch()