python app.py -h to learn how to use.'
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'
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'
EVA02_LARGE_MODEL_IS_DSV1_REPO = 'deepghs/idolsankaku-eva02-large-tagger-v1'
SWINV2_MODEL_IS_DSV1_REPO = 'deepghs/idolsankaku-swinv2-tagger-v1'
# Global variables for model components (for memory management)
CURRENT_MODEL = None
CURRENT_MODEL_NAME = None
CURRENT_TAGS_DF = None
CURRENT_TAG_NAMES = None
CURRENT_RATING_INDEXES = None
CURRENT_GENERAL_INDEXES = None
CURRENT_CHARACTER_INDEXES = None
CURRENT_MODEL_TARGET_SIZE = None
# Custom CSS for gallery styling
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;}
code {background: rgba(192, 132, 252, 0.15); padding: 0.15rem 0.5rem; border-radius: 6px; font-size: 0.9rem;color: #f0d9ff;}
"""
MODEL_FILENAME = 'model.onnx'
LABEL_FILENAME = 'selected_tags.csv'
class Timer:
"""Utility class for measuring execution time of different operations"""
def __init__(self):
self.start_time = time.perf_counter()
self.checkpoints = [('Start', self.start_time)]
def checkpoint(self, label='Checkpoint'):
"""Add a checkpoint with a label"""
now = time.perf_counter()
self.checkpoints.append((label, now))
def report(self, is_clear_checkpoints=True):
"""Report time elapsed since last checkpoint"""
max_label_length = max(len(label) for (label, _) in self.checkpoints) if self.checkpoints else 0
prev_time = self.checkpoints[0][1] if self.checkpoints else 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
if is_clear_checkpoints:
self.checkpoints.clear()
self.checkpoint()
def report_all(self):
"""Report all checkpoint times including total execution time"""
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 if self.checkpoints else 0
print(f"{'Total Execution Time'.ljust(max_label_length)}: {total_time:.3f} seconds\n")
self.checkpoints.clear()
def restart(self):
"""Restart the timer"""
self.start_time = time.perf_counter()
self.checkpoints = [('Start', self.start_time)]
def parse_args() -> argparse.Namespace:
"""Parse command line arguments"""
parser = argparse.ArgumentParser(description='Multi-Tagger - Image analysis and captioning')
# Gradio arguments
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')
# CLI analysis arguments (BETA)
parser.add_argument('--analyze', '-a', nargs='+', metavar='PATH_OR_URL',
help='Analyze images from paths or URLs. Skips Gradio launch.')
parser.add_argument('--output', '-o', type=str, default='output',
help='Output directory for tag files and downloaded images (default: output)')
parser.add_argument('--silent', '-s', action='store_true',
help='Run silently without printing results.')
parser.add_argument('--ct', action='store_true',
help='Print only character tags')
parser.add_argument('--cttx', action='store_true',
help='Print only categorized list')
parser.add_argument('--model', type=str, default=EVA02_LARGE_MODEL_DSV3_REPO,
help='WD model repository to use for analysis')
parser.add_argument('--general-threshold', type=float, default=0.35,
help='General tags threshold (default: 0.35)')
parser.add_argument('--character-threshold', type=float, default=0.85,
help='Character tags threshold (default: 0.85)')
return parser.parse_args()
def load_labels(dataframe) -> tuple:
"""Load tag names and their category indexes from the dataframe"""
name_series = dataframe['name']
tag_names = name_series.tolist()
# Find indexes for different tag categories
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):
"""Calculate threshold using Maximum Change in second derivative (MCut) method"""
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
def _download_model_files(model_repo):
"""Download model files from HuggingFace Hub"""
import huggingface_hub
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 create_optimized_ort_session(model_path):
"""Create an optimized ONNX Runtime session with GPU support (PixAI optimized version)"""
# Session options for better performance
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess_options.intra_op_num_threads = 0 # Use all available cores
sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
sess_options.enable_mem_pattern = True
sess_options.enable_cpu_mem_arena = True
# Check available providers
available_providers = ort.get_available_providers()
print(f"Available ONNX Runtime providers: {available_providers}")
# Use appropriate execution providers (in order of preference)
providers = []
# Use CUDA if available with optimized settings
if 'CUDAExecutionProvider' in available_providers:
cuda_provider = ('CUDAExecutionProvider', {
'device_id': 0,
'arena_extend_strategy': 'kNextPowerOfTwo',
'gpu_mem_limit': 4 * 1024 * 1024 * 1024, # 4GB VRAM
'cudnn_conv_algo_search': 'EXHAUSTIVE',
'do_copy_in_default_stream': True,
})
providers.append(cuda_provider)
print("Using CUDA provider for ONNX inference")
else:
print("CUDA provider not available, falling back to CPU")
# Always include CPU as fallback
providers.append('CPUExecutionProvider')
try:
session = ort.InferenceSession(model_path, sess_options, providers=providers)
print(f"Model loaded with providers: {session.get_providers()}")
return session
except Exception as e:
print(f"Failed to create ONNX session: {e}")
raise
def _load_model_components_optimized(model_repo):
"""Load and optimize model components"""
global CURRENT_MODEL, CURRENT_MODEL_NAME, CURRENT_TAGS_DF, CURRENT_TAG_NAMES
global CURRENT_RATING_INDEXES, CURRENT_GENERAL_INDEXES, CURRENT_CHARACTER_INDEXES, CURRENT_MODEL_TARGET_SIZE
# Only reload if model changed
if model_repo == CURRENT_MODEL_NAME and CURRENT_MODEL is not None:
return
# Download files
csv_path, model_path = _download_model_files(model_repo)
# Load optimized ONNX model
CURRENT_MODEL = create_optimized_ort_session(model_path)
# Load tags
tags_df = pd.read_csv(csv_path)
tag_names, rating_indexes, general_indexes, character_indexes = load_labels(tags_df)
# Store in global variables
CURRENT_TAGS_DF = tags_df
CURRENT_TAG_NAMES = tag_names
CURRENT_RATING_INDEXES = rating_indexes
CURRENT_GENERAL_INDEXES = general_indexes
CURRENT_CHARACTER_INDEXES = character_indexes
# Get model input size
_, height, width, _ = CURRENT_MODEL.get_inputs()[0].shape
CURRENT_MODEL_TARGET_SIZE = height
CURRENT_MODEL_NAME = model_repo
def _raw_predict(image_array, model_session):
"""Run raw prediction using the model session"""
input_name = model_session.get_inputs()[0].name
label_name = model_session.get_outputs()[0].name
preds = model_session.run([label_name], {input_name: image_array})[0]
return preds[0].astype(float)
def unload_model():
"""Explicitly unload the current model from memory"""
global CURRENT_MODEL, CURRENT_MODEL_NAME
global CURRENT_RATING_INDEXES, CURRENT_GENERAL_INDEXES, CURRENT_CHARACTER_INDEXES, CURRENT_MODEL_TARGET_SIZE
# Delete the model session
if CURRENT_MODEL is not None:
del CURRENT_MODEL
CURRENT_MODEL = None
# Clear other large objects
CURRENT_TAGS_DF = None
CURRENT_TAG_NAMES = None
CURRENT_RATING_INDEXES = None
CURRENT_GENERAL_INDEXES = None
CURRENT_CHARACTER_INDEXES = None
CURRENT_MODEL_TARGET_SIZE = None
CURRENT_MODEL_NAME = None
# Force garbage collection
import gc
gc.collect()
# Clear CUDA cache if using GPU
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
print('VRAM is cleared!')
except ImportError:
pass
def prepare_image(path):
print('Preparing image for model input')
image = Image.open(path)
image = image.convert('RGBA')
target_size = CURRENT_MODEL_TARGET_SIZE
# Create white background and composite
canvas = Image.new('RGBA', image.size, (255, 255, 255))
canvas.alpha_composite(image)
image = canvas.convert('RGB')
# Pad 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 needed
if max_dim != target_size:
padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC)
# Convert to array and preprocess
image_array = np.asarray(padded_image, dtype=np.float32)
image_array = image_array[:, :, ::-1] # BGR to RGB
return np.expand_dims(image_array, axis=0)
def create_file(content: str, directory: str, fileName: str) -> str:
"""Creating a file with the given content"""
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 download_image_from_url(url, output_dir):
"""Download an image from a URL and return the local path"""
try:
response = requests.get(url, timeout=30, stream=True)
response.raise_for_status()
# Get filename from URL or generate one
filename = os.path.basename(url.split('?')[0])
if not filename or '.' not in filename:
from datetime import datetime
filename = f"downloaded_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
filepath = os.path.join(output_dir, filename)
with open(filepath, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
print(f"Downloaded: {url} -> {filepath}")
return filepath
except Exception as e:
print(f"Failed to download {url}: {e}")
return None
def run_cli_analysis(args):
"""Run CLI-based image analysis"""
if not args.analyze:
return
# Ensure output directory exists
output_dir = args.output if args.output else './output'
os.makedirs(output_dir, exist_ok=True)
# Process each input - could be a URL or local path
image_paths = []
for item in args.analyze:
# Check if it's a URL
if item.startswith(('http://', 'https://')):
path = download_image_from_url(item, output_dir)
if path:
image_paths.append(path)
else:
# Local file path - handle glob patterns
if os.path.isfile(item):
image_paths.append(item)
elif os.path.isdir(item):
# If it's a directory, find all images
for root, dirs, files in os.walk(item):
for f in files:
if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif')):
image_paths.append(os.path.join(root, f))
else:
print(f"Warning: {item} is not a valid file or URL")
if not image_paths:
print("No valid images found to analyze.")
return
print(f"Analyzing {len(image_paths)} image(s)...")
try:
timer = Timer()
timer.checkpoint("Start analysis")
# Prepare gallery format
gallery = [(p, os.path.basename(p)) for p in image_paths]
# Load model
_load_model_components_optimized(args.model)
timer.checkpoint("Model loaded")
name_counters = defaultdict(int)
results = []
for idx, (image_path, _) in enumerate(gallery):
try:
image_name = os.path.splitext(os.path.basename(image_path))[0]
name_counters[image_name] += 1
if name_counters[image_name] > 1:
image_name = f"{image_name}_{name_counters[image_name]:02d}"
# Prepare and predict
image = prepare_image(image_path)
preds = _raw_predict(image, CURRENT_MODEL)
labels = list(zip(CURRENT_TAG_NAMES, preds))
# Process ratings
ratings_names = [labels[i] for i in CURRENT_RATING_INDEXES]
rating = dict(ratings_names)
# Process general tags
general_names = [labels[i] for i in CURRENT_GENERAL_INDEXES]
general_thresh = args.general_threshold
general_res = [x for x in general_names if x[1] > general_thresh]
general_res = dict(general_res)
# Process character tags
character_names = [labels[i] for i in CURRENT_CHARACTER_INDEXES]
character_thresh = args.character_threshold
character_res = [x for x in character_names if x[1] > character_thresh]
character_res = dict(character_res)
# Sort
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]
character_list = list(character_res.keys())
# Format output
combined_list = character_list + sorted_general_list
sorted_strings = ', '.join(combined_list).replace('(', '\\(').replace(')', '\\)').replace('_', ' ')
# Categorized output
categorized_strings = categorize_tags_output(sorted_strings, character_res).replace('(', '\\(').replace(')', '\\)')
categorized_j = generate_tags_json(sorted_strings, character_res)
categorized_text = "\n ".join(f"{k}: {', '.join(v)}" for k, v in categorized_j.items())
result = {
'image': image_name,
'rating': rating,
'character_tags': character_list,
'general_tags': sorted_general_list,
'tags_string': sorted_strings,
'categorized_strings': categorized_strings,
'categorized_text': categorized_text,
'categorized_json': categorized_j
}
results.append(result)
# Save files if output dir specified
create_file(sorted_strings, output_dir, f"{image_name}.txt")
create_file(categorized_strings, output_dir, f"{image_name}_categorized.txt")
# Save JSON
json_content = json.dumps(result, indent=2, ensure_ascii=False)
create_file(json_content, output_dir, f"{image_name}_result.json")
timer.checkpoint(f"Image {idx+1}: {image_name}")
except Exception as e:
print(f"Error processing {image_path}: {e}")
import traceback
traceback.print_exc()
continue
timer.report_all()
# Print results unless silent
if not args.silent:
for r in results:
print(f"\n{'='*60}")
print(f"Image: {r['image']}")
print(f"Rating: {r['rating']}")
if args.ct:
print(f"Character Tags: {', '.join(r['character_tags'])}")
elif args.cttx:
print(f"Categorized Output:")
for cat, tags in r['categorized_json'].items():
print(f" {cat}: {', '.join(tags)}")
else:
print(f"Character Tags: {', '.join(r['character_tags'])}")
print(f"General Tags: {', '.join(r['general_tags'])}")
print(f"Tags String: {r['tags_string']}")
print(f"Categorized: {r['categorized_strings']}")
print(f"Categorized list:\n{r['categorized_text']}")
print(f"{'='*60}\n")
print(f"\nResults saved to: {os.path.abspath(output_dir)}")
finally:
# Clean VRAM
unload_model()
import gc
gc.collect()
def predict(gallery, model_repo, model_repo_2, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, characters_merge_enabled, additional_tags_prepend, additional_tags_append, tag_results, progress=gr.Progress()):
"""Main prediction function for processing images"""
tag_results.clear()
gallery_len = len(gallery)
print(f"Predict load model: {model_repo}, gallery length: {gallery_len}")
timer = Timer()
progressRatio = 1
progressTotal = gallery_len + 1
current_progress = 0
txt_infos = []
output_dir = tempfile.mkdtemp()
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# Load initial model
_load_model_components_optimized(model_repo)
current_progress += progressRatio / progressTotal
progress(current_progress, desc='Initialize WD model finished')
timer.checkpoint("Initialize WD model")
timer.report()
name_counters = defaultdict(int)
for (idx, value) in enumerate(gallery):
try:
# Handle duplicate filenames
image_path = value[0]
image_name = os.path.splitext(os.path.basename(image_path))[0]
name_counters[image_name] += 1
if name_counters[image_name] > 1:
image_name = f"{image_name}_{name_counters[image_name]:02d}"
## Prepare image
image = prepare_image(image_path)
print(f"Gallery {idx:02d}: Starting run first model ({model_repo})...")
## Load and run first model
_load_model_components_optimized(model_repo)
preds = _raw_predict(image, CURRENT_MODEL)
labels = list(zip(CURRENT_TAG_NAMES, preds))
# Process ratings
ratings_names = [labels[i] for i in CURRENT_RATING_INDEXES]
rating = dict(ratings_names)
# Process general tags
general_names = [labels[i] for i in CURRENT_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)
# Process character tags
character_names = [labels[i] for i in CURRENT_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)
character_list_1 = list(character_res.keys())
# Sort general tags by confidence
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]
# Handle second model if provided
if model_repo_2 and model_repo_2 != model_repo:
print(f"Gallery {idx:02d}: Starting run second model ({model_repo_2})...")
_load_model_components_optimized(model_repo_2)
preds_2 = _raw_predict(image, CURRENT_MODEL)
labels_2 = list(zip(CURRENT_TAG_NAMES, preds_2))
# Process general tags from second model
general_names_2 = [labels_2[i] for i in CURRENT_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)
# Process character tags from second model
character_names_2 = [labels_2[i] for i in CURRENT_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)
character_list_2 = list(character_res_2.keys())
# Sort general tags from second model
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
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:
combined_character_list = character_list_1
combined_general_list = sorted_general_list_1
if not characters_merge_enabled:
combined_character_list = [item for item in combined_character_list if item not in combined_general_list]
# Handle additional 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()]
# Avoid duplicates in prepend/append lists
if prepend_list and append_list:
append_list = [item for item in append_list if item not in prepend_list]
# Remove prepended tags from main list
if prepend_list:
combined_general_list = [item for item in combined_general_list if item not in prepend_list]
# Remove appended tags from main list
if append_list:
combined_general_list = [item for item in combined_general_list if item not in append_list]
# Combine all tags
combined_general_list = prepend_list + combined_general_list + append_list
# Format output string
sorted_general_strings = ', '.join((combined_character_list if characters_merge_enabled else []) + combined_general_list).replace('(', '\\(').replace(')', '\\)').replace('_', ' ')
# Generate categorized output
categorized_strings = categorize_tags_output(sorted_general_strings, character_res).replace('(', '\\(').replace(')', '\\)')
categorized_j = generate_tags_json(sorted_general_strings, character_res)
categorized_text = "\n ".join(f"{k}: {', '.join(v)}" for k, v in categorized_j.items())
# Create output files
txt_content = f"Output (string): {sorted_general_strings}\n\nCategorized Output: {categorized_strings}"
txt_file = create_file(txt_content, output_dir, f"{image_name}_output.txt")
txt_infos.append({'path': txt_file, 'name': f"{image_name}_output.txt"})
# Save image copy
img = Image.open(image_path)
img.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"})
# Create tags text file
txt_file = create_file(sorted_general_strings, output_dir, image_name + '.txt')
# Create categorized list file
categorized_file = create_file(categorized_strings, output_dir, f"{image_name}_categorized.txt")
txt_infos.append({'path': categorized_file, 'name': f"{image_name}_categorized.txt"})
txt_infos.append({'path': txt_file, 'name': image_name + '.txt'})
# Store results
tag_results[image_path] = {'strings': sorted_general_strings, 'categorized_strings': categorized_strings, 'categorized_text': categorized_text, 'rating': rating, 'character_res': character_res}
# Update progress
current_progress += progressRatio / progressTotal
progress(current_progress, desc=f"Image {idx+1}/{gallery_len} processed")
timer.checkpoint(f"Image {idx+1}: {image_name} processed")
timer.report()
except Exception as e:
print(traceback.format_exc())
print('Error predict: ' + str(e))
continue
# Create download zip after all images are processed
download = None
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:
taggers_zip.write(info['path'], arcname=info['name'])
download = downloadZipPath
# Return first image results as default
first_image_results = '', {}, '', '', ''
if gallery and len(gallery) > 0:
first_image_path = gallery[0][0]
if first_image_path in tag_results:
first_result = tag_results[first_image_path]
character_tags_formatted = ", ".join([name.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
for name in first_result['character_res'].keys()])
first_image_results = (first_result['strings'], first_result['rating'], character_tags_formatted, first_result.get('categorized_strings', ''), first_result.get('categorized_text', ''))
progress(1.0, desc=f"Predict completed for {gallery_len} image(s)")
timer.report_all()
print('Predict is complete.')
# Clean up
unload_model()
return (download, first_image_results[0], first_image_results[1], first_image_results[2], first_image_results[3], first_image_results[4], tag_results)
def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state: gr.SelectData):
"""Return first image results if no selection"""
if not selected_state and gallery and len(gallery) > 0:
first_image_path = gallery[0][0]
if first_image_path in tag_results:
first_result = tag_results[first_image_path]
character_tags_formatted = ", ".join([name.replace("(", "\\(").replace(")", "\\)").replace("_", " ") for name in first_result['character_res'].keys()])
return (first_result['strings'], first_result['rating'], character_tags_formatted, first_result.get('categorized_strings', ''), first_result.get('categorized_text', ''))
if not selected_state:
return '', {}, '', '', ''
# Get selected image path
selected_value = selected_state.value
image_path = None
if isinstance(selected_value, dict) and 'image' in selected_value:
image_path = selected_value['image']['path']
elif isinstance(selected_value, (list, tuple)) and len(selected_value) > 0:
image_path = selected_value[0]
else:
image_path = str(selected_value)
# Return stored results
if image_path in tag_results:
result = tag_results[image_path]
character_tags_formatted = ", ".join([name.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
for name in result['character_res'].keys()])
return (result['strings'], result['rating'], character_tags_formatted, result.get('categorized_strings', ''), result.get('categorized_text', ''))
return '', {}, '', '', ''
def append_gallery(gallery: list, image: str):
"""Add a single media file (image or video) to the gallery"""
return handle_single_media_upload(image, gallery)
def extend_gallery(gallery: list, images):
"""Add multiple media files (images or videos) to the gallery"""
return handle_multiple_media_uploads(images, gallery)
def auto_convert_tags_to_prompt(categorized_tags_text, system_prompt_key, llama_model_name):
"""Auto-convert categorized list to prompt if enabled"""
if not categorized_tags_text or not str(categorized_tags_text).strip():
return ""
# Get the system prompt text from the key
prompt_text = SYSTEM_PROMPTS.get(system_prompt_key, SYSTEM_PROMPTS["Creative"])
# Generate response using the selected model and system prompt
result = ""
for chunk in generate_response(categorized_tags_text, system_prompt=prompt_text, model_name=llama_model_name):
result = chunk
return result
# Parse arguments
args = parse_args()
# If --analyze is provided, run CLI mode and exit
if args.analyze:
run_cli_analysis(args)
exit(0)
dropdown_list = [
EVA02_LARGE_MODEL_DSV3_REPO, VIT_LARGE_MODEL_DSV3_REPO, SWINV2_MODEL_DSV3_REPO,
CONV_MODEL_DSV3_REPO, VIT_MODEL_DSV3_REPO, MOAT_MODEL_DSV2_REPO,
SWIN_MODEL_DSV2_REPO, CONV_MODEL_DSV2_REPO, CONV2_MODEL_DSV2_REPO,
VIT_MODEL_DSV2_REPO, EVA02_LARGE_MODEL_IS_DSV1_REPO, SWINV2_MODEL_IS_DSV1_REPO
]
# Llama model choices
LLAMA_MODEL_CHOICES = ["gemma-3-1B"]
DEFAULT_LLAMA_MODEL = "gemma-3-1B"
with gr.Blocks(title=TITLE, css=css, theme="Werli/Purple-Crimson-Gradio-Theme", fill_width=True) as demo:
gr.Markdown(value=f"{DESCRIPTION}
") with gr.Tab(label='WD'): with gr.Row(): with gr.Column(): with gr.Column(variant='panel'): image_input = gr.Image(label='Upload an Image (or paste from clipboard)', type='filepath', sources=['upload', 'clipboard'], height=150) with gr.Row(): upload_button = gr.UploadButton('UPLOAD VIDEOS', file_types=['image', 'video'], file_count='multiple', size='md') gallery = gr.Gallery(columns=2, show_share_button=False, interactive=True, height='auto', label='Grid of images', preview=False, elem_id='custom-gallery') submit = gr.Button(value='START', variant='primary', size='lg') clear = gr.ClearButton(components=[gallery], value='CLEAR GALLERY', variant='secondary', size='sm') 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"{PLUS}
") model_repo_2 = gr.Dropdown([None] + dropdown_list, value=None, label='2nd Model (Optional)', info='Select another model for diversified results.') with gr.Accordion("OTHER", open=False): general_thresh = gr.Slider(0, 1, step=args.score_slider_step, value=args.score_general_threshold, label='General Tags Threshold', scale=3) character_thresh = gr.Slider(0, 1, step=args.score_slider_step, value=args.score_character_threshold, label='Character Tags Threshold', scale=3) general_mcut_enabled = gr.Checkbox(value=False, label='Use MCut threshold', scale=1) character_mcut_enabled = gr.Checkbox(value=False, label='Use MCut threshold', scale=1) characters_merge_enabled = gr.Checkbox(value=False, label='Merge characters into the string output', scale=1) additional_tags_prepend = gr.Text(label='Prepend Additional tags (comma split)') additional_tags_append = gr.Text(label='Append Additional tags (comma split)') rating = gr.Label(label='RATING') with gr.Row(): clear = gr.ClearButton(components=[gallery, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, characters_merge_enabled, additional_tags_prepend, additional_tags_append], value="CLEAR EVERYTHING", variant='secondary', size='lg' ) gr.Markdown('[Based on SmilingWolf/wd-tagger](https://huggingface.co/spaces/SmilingWolf/wd-tagger)') with gr.Column(variant='panel'): download_file = gr.File(label='DOWNLOAD') character_res = gr.Textbox(label='CHARACTER TAGS',show_copy_button=True,lines=2) sorted_general_strings = gr.Textbox(label='OUTPUT',show_label=True,show_copy_button=True,lines=5,max_lines=25) categorized_strings = gr.Textbox(label='CATEGORIZED TAGS',show_label=True,show_copy_button=True,lines=5,max_lines=25) tags_string = gr.Textbox(label='CATEGORIZED LIST (EDITABLE)',show_copy_button=True,lines=5,max_lines=25) convert_tags_output = gr.Textbox(label='PROMPT',show_label=True,lines=5, show_copy_button=True,max_lines=20) with gr.Row(): submit_prompt_to_convert = gr.Button(value='CREATE A PROMPT (RANDOM)', variant='primary', size='md') # System prompt and model selection for conversion with gr.Accordion("Prompt Settings", open=False): system_prompt_dropdown = gr.Dropdown( choices=list(SYSTEM_PROMPTS.keys()), value="Creative", label="System Prompt Style", info="Choose how the tags should be converted into a prompt" ) llama_model_dropdown = gr.Dropdown( choices=LLAMA_MODEL_CHOICES, value=DEFAULT_LLAMA_MODEL, label="LLM Model", info="Select the LLM model for tag-to-prompt conversion" ) auto_convert_checkbox = gr.Checkbox( value=False, label="Auto-convert CATEGORIZED LIST to prompt after analysis", info="When enabled, the prompt will be automatically generated after analysis completes" ) # State to store results tag_results = gr.State({}) # Event handlers image_input.change(append_gallery,inputs=[gallery, image_input], outputs=[gallery, image_input]) upload_button.upload(extend_gallery,inputs=[gallery, upload_button], outputs=gallery) gallery.select(get_selection_from_gallery,inputs=[gallery, tag_results], outputs=[sorted_general_strings, rating, character_res, categorized_strings, tags_string]) # Auto-convert trigger after predict completes def handle_after_predict(download, output_str, rating_val, char_tags, cat_str, cat_tags_str, tag_results_dict, auto_convert, sys_prompt, llm_model): """Handle post-predict: if auto-convert is enabled, run conversion""" prompt_output = "" if auto_convert and cat_tags_str and str(cat_tags_str).strip(): prompt_output = auto_convert_tags_to_prompt(cat_tags_str, sys_prompt, llm_model) return download, output_str, rating_val, char_tags, cat_str, cat_tags_str, prompt_output submit.click( predict, inputs=[gallery, model_repo, model_repo_2, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, characters_merge_enabled, additional_tags_prepend, additional_tags_append, tag_results], outputs=[download_file, sorted_general_strings, rating, character_res, categorized_strings, tags_string, tag_results] ).then( handle_after_predict, inputs=[download_file, sorted_general_strings, rating, character_res, categorized_strings, tags_string, tag_results, auto_convert_checkbox, system_prompt_dropdown, llama_model_dropdown], outputs=[download_file, sorted_general_strings, rating, character_res, categorized_strings, tags_string, convert_tags_output] ) # Convert tags button submit_prompt_to_convert.click( generate_response, inputs=[tags_string, system_prompt_dropdown, llama_model_dropdown], outputs=[convert_tags_output] ) with gr.Tab("PixAI"): pixai_components = create_pixai_interface( system_prompt_options=list(SYSTEM_PROMPTS.keys()), default_system_prompt="Creative", llama_model_choices=LLAMA_MODEL_CHOICES, default_llama_model=DEFAULT_LLAMA_MODEL ) with gr.Tab("ComfyUI Metadata Extractor"): comfy_interface = create_multi_comfy() demo.queue().launch()