Spaces:
Sleeping
Sleeping
Upload 9 files
Browse files- app.py +673 -370
- modules/__pycache__/booru.cpython-311.pyc +0 -0
- modules/__pycache__/classifyTags.cpython-311.pyc +0 -0
- modules/__pycache__/pixai.cpython-311.pyc +0 -0
- modules/booru.py +51 -1
- modules/classifyTags.py +337 -156
- modules/pixai.py +810 -0
- requirements.txt +19 -20
app.py
CHANGED
|
@@ -1,520 +1,823 @@
|
|
| 1 |
-
import os, io,
|
| 2 |
-
import
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
from datetime import datetime, timezone
|
| 4 |
from collections import defaultdict
|
| 5 |
from PIL import Image, ImageOps
|
| 6 |
-
from modules.booru import booru_gradio, on_select
|
| 7 |
from apscheduler.schedulers.background import BackgroundScheduler
|
| 8 |
-
from modules.classifyTags import
|
| 9 |
-
from modules.
|
| 10 |
-
from modules.
|
| 11 |
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
|
|
|
| 15 |
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
|
|
|
|
|
|
| 19 |
"""
|
| 20 |
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
EVA02_LARGE_MODEL_IS_DSV1_REPO =
|
| 35 |
-
SWINV2_MODEL_IS_DSV1_REPO =
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 44 |
class Timer:
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
class Predictor:
|
|
|
|
|
|
|
| 57 |
def __init__(self):
|
| 58 |
-
self.
|
| 59 |
self.last_loaded_repo = None
|
| 60 |
-
|
| 61 |
-
csv_path = huggingface_hub.hf_hub_download(model_repo, LABEL_FILENAME, )
|
| 62 |
-
model_path = huggingface_hub.hf_hub_download(model_repo, MODEL_FILENAME, )
|
| 63 |
-
return csv_path, model_path
|
| 64 |
def load_model(self, model_repo):
|
| 65 |
-
if
|
|
|
|
| 66 |
return
|
| 67 |
-
|
| 68 |
-
tags_df = pd.read_csv(csv_path)
|
| 69 |
-
sep_tags = load_labels(tags_df)
|
| 70 |
-
self.tag_names = sep_tags[0]
|
| 71 |
-
self.rating_indexes = sep_tags[1]
|
| 72 |
-
self.general_indexes = sep_tags[2]
|
| 73 |
-
self.character_indexes = sep_tags[3]
|
| 74 |
-
model = rt.InferenceSession(model_path)
|
| 75 |
-
_, height, width, _ = model.get_inputs()[0].shape
|
| 76 |
-
self.model_target_size = height
|
| 77 |
self.last_loaded_repo = model_repo
|
| 78 |
-
|
| 79 |
def prepare_image(self, path):
|
|
|
|
| 80 |
image = Image.open(path)
|
| 81 |
-
image = image.convert(
|
| 82 |
-
target_size =
|
| 83 |
-
|
|
|
|
|
|
|
| 84 |
canvas.alpha_composite(image)
|
| 85 |
-
image = canvas.convert(
|
| 86 |
-
|
|
|
|
| 87 |
image_shape = image.size
|
| 88 |
max_dim = max(image_shape)
|
| 89 |
pad_left = (max_dim - image_shape[0]) // 2
|
| 90 |
pad_top = (max_dim - image_shape[1]) // 2
|
| 91 |
-
padded_image = Image.new(
|
| 92 |
padded_image.paste(image, (pad_left, pad_top))
|
| 93 |
-
|
|
|
|
| 94 |
if max_dim != target_size:
|
| 95 |
-
padded_image = padded_image.resize(
|
| 96 |
-
|
| 97 |
-
|
| 98 |
-
)
|
| 99 |
-
# Convert to numpy array
|
| 100 |
image_array = np.asarray(padded_image, dtype=np.float32)
|
| 101 |
-
|
| 102 |
-
image_array = image_array[:, :, ::-1]
|
| 103 |
return np.expand_dims(image_array, axis=0)
|
| 104 |
|
| 105 |
def create_file(self, content: str, directory: str, fileName: str) -> str:
|
| 106 |
-
|
| 107 |
file_path = os.path.join(directory, fileName)
|
| 108 |
if fileName.endswith('.json'):
|
| 109 |
-
with open(file_path, 'w', encoding=
|
| 110 |
file.write(content)
|
| 111 |
else:
|
| 112 |
-
with open(file_path, 'w+', encoding=
|
| 113 |
file.write(content)
|
| 114 |
return file_path
|
| 115 |
-
|
| 116 |
-
|
| 117 |
-
|
| 118 |
-
|
| 119 |
-
|
| 120 |
-
general_thresh,
|
| 121 |
-
general_mcut_enabled,
|
| 122 |
-
character_thresh,
|
| 123 |
-
character_mcut_enabled,
|
| 124 |
-
characters_merge_enabled,
|
| 125 |
-
beautify_model_repo,
|
| 126 |
-
additional_tags_prepend,
|
| 127 |
-
additional_tags_append,
|
| 128 |
-
tag_results,
|
| 129 |
-
progress=gr.Progress()
|
| 130 |
-
):
|
| 131 |
-
# Clear tag_results before starting a new prediction
|
| 132 |
tag_results.clear()
|
| 133 |
gallery_len = len(gallery)
|
| 134 |
print(f"Predict load model: {model_repo}, gallery length: {gallery_len}")
|
| 135 |
-
|
| 136 |
-
|
|
|
|
| 137 |
progressTotal = gallery_len + 1
|
| 138 |
current_progress = 0
|
| 139 |
-
# Initialize variables that need to be accessible throughout the function
|
| 140 |
-
final_categorized_output = ""
|
| 141 |
-
categorized_output_strings = []
|
| 142 |
txt_infos = []
|
| 143 |
output_dir = tempfile.mkdtemp()
|
|
|
|
| 144 |
if not os.path.exists(output_dir):
|
| 145 |
os.makedirs(output_dir)
|
|
|
|
|
|
|
| 146 |
self.load_model(model_repo)
|
| 147 |
-
current_progress += progressRatio/progressTotal
|
| 148 |
-
progress(current_progress, desc=
|
| 149 |
-
timer.checkpoint(
|
| 150 |
-
if beautify_model_repo:
|
| 151 |
-
print(f"Load model {beautify_model_repo}")
|
| 152 |
-
beautify = beautify_class(beautify_model_repo, loadModel=True)
|
| 153 |
-
current_progress += progressRatio/progressTotal;
|
| 154 |
-
progress(current_progress, desc="Initialize beautify model finished")
|
| 155 |
-
timer.checkpoint(f"Initialize beautify model")
|
| 156 |
timer.report()
|
| 157 |
-
|
| 158 |
name_counters = defaultdict(int)
|
| 159 |
-
|
|
|
|
| 160 |
try:
|
|
|
|
| 161 |
image_path = value[0]
|
| 162 |
image_name = os.path.splitext(os.path.basename(image_path))[0]
|
| 163 |
-
# Increment the counter for the current name
|
| 164 |
name_counters[image_name] += 1
|
| 165 |
if name_counters[image_name] > 1:
|
| 166 |
image_name = f"{image_name}_{name_counters[image_name]:02d}"
|
|
|
|
|
|
|
| 167 |
image = self.prepare_image(image_path)
|
| 168 |
-
# Run first model
|
| 169 |
print(f"Gallery {idx:02d}: Starting run first model ({model_repo})...")
|
|
|
|
|
|
|
| 170 |
self.load_model(model_repo)
|
| 171 |
-
|
| 172 |
-
|
| 173 |
-
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
ratings_names = [labels[i] for i in self.rating_indexes]
|
| 177 |
rating = dict(ratings_names)
|
| 178 |
|
| 179 |
-
|
|
|
|
| 180 |
if general_mcut_enabled:
|
| 181 |
general_probs = np.array([x[1] for x in general_names])
|
| 182 |
general_thresh_temp = mcut_threshold(general_probs)
|
| 183 |
else:
|
| 184 |
general_thresh_temp = general_thresh
|
|
|
|
| 185 |
general_res = [x for x in general_names if x[1] > general_thresh_temp]
|
| 186 |
general_res = dict(general_res)
|
| 187 |
-
|
|
|
|
|
|
|
| 188 |
if character_mcut_enabled:
|
| 189 |
character_probs = np.array([x[1] for x in character_names])
|
| 190 |
character_thresh_temp = mcut_threshold(character_probs)
|
| 191 |
character_thresh_temp = max(0.15, character_thresh_temp)
|
| 192 |
else:
|
| 193 |
character_thresh_temp = character_thresh
|
|
|
|
| 194 |
character_res = [x for x in character_names if x[1] > character_thresh_temp]
|
| 195 |
character_res = dict(character_res)
|
| 196 |
-
# Collect tags from first model
|
| 197 |
character_list_1 = list(character_res.keys())
|
|
|
|
|
|
|
| 198 |
sorted_general_list_1 = sorted(general_res.items(), key=lambda x: x[1], reverse=True)
|
| 199 |
sorted_general_list_1 = [x[0] for x in sorted_general_list_1]
|
| 200 |
-
|
|
|
|
| 201 |
if model_repo_2 and model_repo_2 != model_repo:
|
| 202 |
print(f"Gallery {idx:02d}: Starting run second model ({model_repo_2})...")
|
| 203 |
self.load_model(model_repo_2)
|
| 204 |
-
preds_2 =
|
| 205 |
-
labels_2 = list(zip(
|
| 206 |
-
|
| 207 |
-
|
|
|
|
| 208 |
if general_mcut_enabled:
|
| 209 |
general_probs_2 = np.array([x[1] for x in general_names_2])
|
| 210 |
general_thresh_temp_2 = mcut_threshold(general_probs_2)
|
| 211 |
else:
|
| 212 |
general_thresh_temp_2 = general_thresh
|
|
|
|
| 213 |
general_res_2 = [x for x in general_names_2 if x[1] > general_thresh_temp_2]
|
| 214 |
general_res_2 = dict(general_res_2)
|
| 215 |
-
|
|
|
|
|
|
|
| 216 |
if character_mcut_enabled:
|
| 217 |
character_probs_2 = np.array([x[1] for x in character_names_2])
|
| 218 |
character_thresh_temp_2 = mcut_threshold(character_probs_2)
|
| 219 |
character_thresh_temp_2 = max(0.15, character_thresh_temp_2)
|
| 220 |
else:
|
| 221 |
character_thresh_temp_2 = character_thresh
|
|
|
|
| 222 |
character_res_2 = [x for x in character_names_2 if x[1] > character_thresh_temp_2]
|
| 223 |
character_res_2 = dict(character_res_2)
|
| 224 |
-
# Collect tags from second model
|
| 225 |
character_list_2 = list(character_res_2.keys())
|
|
|
|
|
|
|
| 226 |
sorted_general_list_2 = sorted(general_res_2.items(), key=lambda x: x[1], reverse=True)
|
| 227 |
sorted_general_list_2 = [x[0] for x in sorted_general_list_2]
|
| 228 |
-
|
|
|
|
| 229 |
combined_character_list = list(set(character_list_1 + character_list_2))
|
| 230 |
combined_general_list = list(set(sorted_general_list_1 + sorted_general_list_2))
|
| 231 |
else:
|
| 232 |
-
# Only first model was used
|
| 233 |
combined_character_list = character_list_1
|
| 234 |
combined_general_list = sorted_general_list_1
|
| 235 |
-
|
| 236 |
-
|
| 237 |
-
|
| 238 |
-
|
| 239 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 240 |
if prepend_list and append_list:
|
| 241 |
append_list = [item for item in append_list if item not in prepend_list]
|
| 242 |
-
|
|
|
|
| 243 |
if prepend_list:
|
| 244 |
combined_general_list = [item for item in combined_general_list if item not in prepend_list]
|
|
|
|
|
|
|
| 245 |
if append_list:
|
| 246 |
combined_general_list = [item for item in combined_general_list if item not in append_list]
|
|
|
|
|
|
|
| 247 |
combined_general_list = prepend_list + combined_general_list + append_list
|
| 248 |
-
|
| 249 |
-
|
| 250 |
-
|
| 251 |
-
|
| 252 |
-
|
| 253 |
-
|
| 254 |
-
|
| 255 |
-
#
|
| 256 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 257 |
txt_file = self.create_file(txt_content, output_dir, f"{image_name}_output.txt")
|
| 258 |
-
txt_infos.append({
|
| 259 |
-
|
| 260 |
-
|
| 261 |
-
json_file = self.create_file(json_content, output_dir, f"{image_name}_categorized_tags.json")
|
| 262 |
-
txt_infos.append({"path": json_file, "name": f"{image_name}_categorized_tags.json"})
|
| 263 |
-
# Save a copy of the uploaded image in PNG format
|
| 264 |
image_path = value[0]
|
| 265 |
image = Image.open(image_path)
|
| 266 |
-
image.save(os.path.join(output_dir, f"{image_name}.png"), format=
|
| 267 |
-
txt_infos.append({
|
| 268 |
-
|
| 269 |
-
|
| 270 |
-
|
| 271 |
-
|
| 272 |
-
|
| 273 |
-
|
| 274 |
-
|
| 275 |
-
|
| 276 |
-
|
| 277 |
-
|
| 278 |
-
|
| 279 |
-
|
| 280 |
-
|
| 281 |
-
|
| 282 |
-
current_progress += progressRatio/progressTotal;
|
| 283 |
-
progress(current_progress, desc=f"image{idx:02d}, beautify finished!")
|
| 284 |
-
timer.checkpoint(f"image{idx:02d}, beautify finished!")
|
| 285 |
-
txt_file = self.create_file(sorted_general_strings, output_dir, image_name + ".txt")
|
| 286 |
-
txt_infos.append({"path":txt_file, "name": image_name + ".txt"})
|
| 287 |
-
|
| 288 |
-
# Store the result in tag_results using image_path as the key
|
| 289 |
tag_results[image_path] = {
|
| 290 |
-
|
| 291 |
-
|
| 292 |
-
|
| 293 |
-
|
| 294 |
-
|
| 295 |
-
|
| 296 |
-
"unclassified_tags": unclassified_tags,
|
| 297 |
-
"summarize_tags": "" # Initialize as empty string
|
| 298 |
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 299 |
timer.report()
|
|
|
|
| 300 |
except Exception as e:
|
| 301 |
print(traceback.format_exc())
|
| 302 |
-
print(
|
| 303 |
-
|
|
|
|
| 304 |
download = []
|
| 305 |
if txt_infos is not None and len(txt_infos) > 0:
|
| 306 |
-
downloadZipPath = os.path.join(
|
|
|
|
|
|
|
|
|
|
| 307 |
with zipfile.ZipFile(downloadZipPath, 'w', zipfile.ZIP_DEFLATED) as taggers_zip:
|
| 308 |
for info in txt_infos:
|
| 309 |
-
|
| 310 |
-
|
|
|
|
| 311 |
download.append(downloadZipPath)
|
| 312 |
-
# End zip creation logic
|
| 313 |
|
| 314 |
-
if beautify_model_repo:
|
| 315 |
-
beautify.release_vram()
|
| 316 |
-
del beautify
|
| 317 |
progress(1, desc=f"Predict completed")
|
| 318 |
-
timer.report_all()
|
| 319 |
-
print(
|
| 320 |
-
|
| 321 |
-
|
| 322 |
-
|
| 323 |
-
if
|
| 324 |
-
|
| 325 |
-
|
| 326 |
-
|
| 327 |
-
|
| 328 |
-
|
| 329 |
-
|
| 330 |
-
|
| 331 |
-
|
| 332 |
-
|
| 333 |
-
|
| 334 |
-
|
| 335 |
-
|
| 336 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 337 |
def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state: gr.SelectData):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 338 |
if not selected_state:
|
| 339 |
-
return
|
| 340 |
-
|
| 341 |
-
|
| 342 |
-
|
| 343 |
-
|
| 344 |
-
|
| 345 |
-
|
| 346 |
-
|
| 347 |
-
|
| 348 |
-
|
| 349 |
-
|
| 350 |
-
|
| 351 |
-
|
| 352 |
-
|
| 353 |
-
|
| 354 |
-
|
| 355 |
-
|
| 356 |
-
|
| 357 |
-
|
| 358 |
-
|
| 359 |
-
|
| 360 |
-
|
| 361 |
-
|
| 362 |
-
|
| 363 |
-
|
| 364 |
-
|
| 365 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 366 |
args = parse_args()
|
| 367 |
predictor = Predictor()
|
| 368 |
dropdown_list = [
|
| 369 |
-
EVA02_LARGE_MODEL_DSV3_REPO,
|
| 370 |
-
|
| 371 |
-
|
| 372 |
-
|
| 373 |
-
VIT_MODEL_DSV3_REPO,
|
| 374 |
-
# ---
|
| 375 |
-
MOAT_MODEL_DSV2_REPO,
|
| 376 |
-
SWIN_MODEL_DSV2_REPO,
|
| 377 |
-
CONV_MODEL_DSV2_REPO,
|
| 378 |
-
CONV2_MODEL_DSV2_REPO,
|
| 379 |
-
VIT_MODEL_DSV2_REPO,
|
| 380 |
-
# ---
|
| 381 |
-
EVA02_LARGE_MODEL_IS_DSV1_REPO,
|
| 382 |
-
SWINV2_MODEL_IS_DSV1_REPO,
|
| 383 |
]
|
|
|
|
| 384 |
def _restart_space():
|
| 385 |
-
|
| 386 |
-
|
| 387 |
-
|
| 388 |
-
|
| 389 |
-
|
| 390 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 391 |
scheduler.start()
|
| 392 |
-
next_run_time_utc=restart_space_job.next_run_time.astimezone(timezone.utc)
|
| 393 |
-
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."
|
| 394 |
|
| 395 |
-
css =
|
| 396 |
-
#custom-gallery {--row-height: 180px;display: grid;grid-auto-rows: min-content;gap: 10px;}
|
| 397 |
-
#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;}
|
| 398 |
-
#custom-gallery .thumbnail-item:hover {transform: translateY(-3px);box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);}
|
| 399 |
-
#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;}
|
| 400 |
-
#custom-gallery .thumbnail-item img.portrait {max-width: 100%;}
|
| 401 |
-
#custom-gallery .thumbnail-item img.landscape {max-height: 100%;}
|
| 402 |
-
.gallery-container {max-height: 500px;overflow-y: auto;padding-right: 0px;--size-80: 500px;}
|
| 403 |
-
.thumbnails {display: flex;position: absolute;bottom: 0;width: 120px;overflow-x: scroll;padding-top: 320px;padding-bottom: 280px;padding-left: 4px;flex-wrap: wrap;}
|
| 404 |
-
#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;}
|
| 405 |
-
"""
|
| 406 |
-
with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True) as demo:
|
| 407 |
gr.Markdown(value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
|
| 408 |
-
#gr.Markdown(value=DESCRIPTION)
|
| 409 |
gr.Markdown(value=f"<p style='text-align: center;'>{DESCRIPTION}</p>")
|
| 410 |
-
|
|
|
|
| 411 |
with gr.Row():
|
| 412 |
with gr.Column():
|
| 413 |
-
|
| 414 |
-
with gr.Column(variant=
|
| 415 |
-
|
| 416 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 417 |
with gr.Row():
|
| 418 |
-
upload_button = gr.UploadButton(
|
| 419 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 420 |
gallery = gr.Gallery(
|
| 421 |
columns=2,
|
| 422 |
-
show_share_button=False,
|
| 423 |
-
interactive=True,
|
| 424 |
-
height=
|
| 425 |
-
label=
|
| 426 |
-
preview=False,
|
| 427 |
-
elem_id=
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 428 |
)
|
| 429 |
-
|
| 430 |
-
model_repo = gr.Dropdown(dropdown_list, value=EVA02_LARGE_MODEL_DSV3_REPO, label="1st Model", )
|
| 431 |
-
PLUS = "+?"
|
| 432 |
gr.Markdown(value=f"<p style='text-align: center;'>{PLUS}</p>")
|
| 433 |
-
model_repo_2 = gr.Dropdown(
|
| 434 |
-
|
| 435 |
-
|
| 436 |
-
|
|
|
|
|
|
|
|
|
|
| 437 |
with gr.Row():
|
| 438 |
-
|
| 439 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 440 |
with gr.Row():
|
| 441 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 442 |
with gr.Row():
|
| 443 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 444 |
with gr.Row():
|
| 445 |
-
additional_tags_prepend = gr.Text(
|
| 446 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 447 |
with gr.Row():
|
| 448 |
clear = gr.ClearButton(
|
| 449 |
-
components=[
|
| 450 |
-
|
| 451 |
-
|
| 452 |
-
|
| 453 |
-
|
| 454 |
-
|
| 455 |
-
|
| 456 |
-
|
| 457 |
-
|
| 458 |
-
|
| 459 |
-
|
| 460 |
-
|
| 461 |
-
|
| 462 |
-
|
| 463 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 464 |
tag_results = gr.State({})
|
| 465 |
-
|
| 466 |
-
|
| 467 |
-
|
| 468 |
-
|
| 469 |
-
|
| 470 |
-
|
| 471 |
-
|
| 472 |
-
|
| 473 |
-
|
| 474 |
-
|
| 475 |
-
|
| 476 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 477 |
gr.Examples(
|
| 478 |
-
[[
|
| 479 |
-
inputs=[image_input, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled
|
| 480 |
-
|
|
|
|
|
|
|
|
|
|
| 481 |
with gr.Tab("Booru Image Fetcher"):
|
| 482 |
-
|
| 483 |
-
|
| 484 |
-
|
| 485 |
-
|
| 486 |
-
|
| 487 |
-
exclude_tags = gr.Textbox(label="Exclude Tags (comma-separated)", placeholder="e.g. animated, watermark, username, ...", lines=3)
|
| 488 |
-
score = gr.Number(label="Minimum Score", value=0)
|
| 489 |
-
count = gr.Slider(label="Number of Images", minimum=1, maximum=20, step=1, value=1)
|
| 490 |
-
Safe = gr.Checkbox(label="Include Safe", value=True)
|
| 491 |
-
Questionable = gr.Checkbox(label="Include Questionable", value=True)
|
| 492 |
-
Explicit = gr.Checkbox(label="Include Explicit (18+)", value=False)
|
| 493 |
-
submit_btn = gr.Button("Fetch Images", variant="primary")
|
| 494 |
-
with gr.Column():
|
| 495 |
-
gr.Markdown("### 📄 Results")
|
| 496 |
-
images_output = gr.Gallery(label="Images", columns=3, rows=2, object_fit="contain", height=500)
|
| 497 |
-
tags_output = gr.Textbox(label="Tags", placeholder="Select an image to display tags", lines=6, show_copy_button=True)
|
| 498 |
-
post_url_output = gr.Textbox(label="Post URL", lines=1, show_copy_button=True)
|
| 499 |
-
image_url_output = gr.Textbox(label="Image URL", lines=1, show_copy_button=True)
|
| 500 |
-
# State to store tags, URLs
|
| 501 |
-
tags_state = gr.State([])
|
| 502 |
-
post_url_state = gr.State([])
|
| 503 |
-
image_url_state = gr.State([])
|
| 504 |
-
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], )
|
| 505 |
-
images_output.select(fn=on_select, inputs=[tags_state, post_url_state, image_url_state], outputs=[tags_output, post_url_output, image_url_output], )
|
| 506 |
-
with gr.Tab(label="Misc"):
|
| 507 |
-
with gr.Row():
|
| 508 |
-
with gr.Column(variant="panel"):
|
| 509 |
-
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)
|
| 510 |
-
submit_button = gr.Button(value="START", variant="primary", size="lg")
|
| 511 |
-
with gr.Column(variant="panel"):
|
| 512 |
-
categorized_string = gr.Textbox(label="Categorized (string)", show_label=True, show_copy_button=True, lines=8)
|
| 513 |
-
categorized_json = gr.JSON(label="Categorized (tags) - JSON")
|
| 514 |
-
submit_button.click(process_tags, inputs=[input_tags], outputs=[categorized_string, categorized_json])
|
| 515 |
-
with gr.Column(variant="panel"):
|
| 516 |
-
pe_generate_btn = gr.Button(value="SUMMARIZE TAGS", size="lg", variant="primary")
|
| 517 |
-
summarize_tags = gr.Textbox(label="Summarized Tags", show_label=True, show_copy_button=True, lines=5)
|
| 518 |
-
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.")
|
| 519 |
-
pe_generate_btn.click(lambda tags, model:prompt_summarizer('', '', tags, model)[0], inputs=[categorized_string, prompt_summarizer_model], outputs=[summarize_tags])
|
| 520 |
-
demo.queue(max_size=10).launch(show_error=True)
|
|
|
|
| 1 |
+
import os, io, json, requests, spaces, argparse, traceback, tempfile, zipfile, re, ast, time
|
| 2 |
+
import gradio as gr
|
| 3 |
+
import numpy as np
|
| 4 |
+
import huggingface_hub
|
| 5 |
+
import onnxruntime as ort
|
| 6 |
+
import pandas as pd
|
| 7 |
from datetime import datetime, timezone
|
| 8 |
from collections import defaultdict
|
| 9 |
from PIL import Image, ImageOps
|
|
|
|
| 10 |
from apscheduler.schedulers.background import BackgroundScheduler
|
| 11 |
+
from modules.classifyTags import categorize_tags_output, generate_tags_json
|
| 12 |
+
from modules.pixai import create_pixai_interface
|
| 13 |
+
from modules.booru import create_booru_interface
|
| 14 |
|
| 15 |
+
""" For GPU install all the requirements.txt and the following:
|
| 16 |
+
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
|
| 17 |
+
pip install onnxruntime-gpu
|
| 18 |
+
"""
|
| 19 |
|
| 20 |
+
""" It's recommended to create a venv:
|
| 21 |
+
python -m venv venv
|
| 22 |
+
venv\Scripts\activate
|
| 23 |
+
pip install ...
|
| 24 |
+
python app.py
|
| 25 |
"""
|
| 26 |
|
| 27 |
+
TITLE = 'Multi-Tagger v1.3'
|
| 28 |
+
DESCRIPTION = '\nMulti-Tagger is a versatile application for advanced image analysis and captioning. Supports <b>CUDA</b> and <b>CPU</b>.\n'
|
| 29 |
+
|
| 30 |
+
SWINV2_MODEL_DSV3_REPO = 'SmilingWolf/wd-swinv2-tagger-v3'
|
| 31 |
+
CONV_MODEL_DSV3_REPO = 'SmilingWolf/wd-convnext-tagger-v3'
|
| 32 |
+
VIT_MODEL_DSV3_REPO = 'SmilingWolf/wd-vit-tagger-v3'
|
| 33 |
+
VIT_LARGE_MODEL_DSV3_REPO = 'SmilingWolf/wd-vit-large-tagger-v3'
|
| 34 |
+
EVA02_LARGE_MODEL_DSV3_REPO = 'SmilingWolf/wd-eva02-large-tagger-v3'
|
| 35 |
+
MOAT_MODEL_DSV2_REPO = 'SmilingWolf/wd-v1-4-moat-tagger-v2'
|
| 36 |
+
SWIN_MODEL_DSV2_REPO = 'SmilingWolf/wd-v1-4-swinv2-tagger-v2'
|
| 37 |
+
CONV_MODEL_DSV2_REPO = 'SmilingWolf/wd-v1-4-convnext-tagger-v2'
|
| 38 |
+
CONV2_MODEL_DSV2_REPO = 'SmilingWolf/wd-v1-4-convnextv2-tagger-v2'
|
| 39 |
+
VIT_MODEL_DSV2_REPO = 'SmilingWolf/wd-v1-4-vit-tagger-v2'
|
| 40 |
+
EVA02_LARGE_MODEL_IS_DSV1_REPO = 'deepghs/idolsankaku-eva02-large-tagger-v1'
|
| 41 |
+
SWINV2_MODEL_IS_DSV1_REPO = 'deepghs/idolsankaku-swinv2-tagger-v1'
|
| 42 |
+
|
| 43 |
+
# Global variables for model components (for memory management)
|
| 44 |
+
CURRENT_MODEL = None
|
| 45 |
+
CURRENT_MODEL_NAME = None
|
| 46 |
+
CURRENT_TAGS_DF = None
|
| 47 |
+
CURRENT_TAG_NAMES = None
|
| 48 |
+
CURRENT_RATING_INDEXES = None
|
| 49 |
+
CURRENT_GENERAL_INDEXES = None
|
| 50 |
+
CURRENT_CHARACTER_INDEXES = None
|
| 51 |
+
CURRENT_MODEL_TARGET_SIZE = None
|
| 52 |
+
|
| 53 |
+
# Custom CSS for gallery styling
|
| 54 |
+
css = """
|
| 55 |
+
#custom-gallery {--row-height: 180px;display: grid;grid-auto-rows: min-content;gap: 10px;}
|
| 56 |
+
#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;}
|
| 57 |
+
#custom-gallery .thumbnail-item:hover {transform: translateY(-3px);box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);}
|
| 58 |
+
#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;}
|
| 59 |
+
#custom-gallery .thumbnail-item img.portrait {max-width: 100%;}
|
| 60 |
+
#custom-gallery .thumbnail-item img.landscape {max-height: 100%;}
|
| 61 |
+
.gallery-container {max-height: 500px;overflow-y: auto;padding-right: 0px;--size-80: 500px;}
|
| 62 |
+
.thumbnails {display: flex;position: absolute;bottom: 0;width: 120px;overflow-x: scroll;padding-top: 320px;padding-bottom: 280px;padding-left: 4px;flex-wrap: wrap;}
|
| 63 |
+
#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;}
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
MODEL_FILENAME = 'model.onnx'
|
| 67 |
+
LABEL_FILENAME = 'selected_tags.csv'
|
| 68 |
+
|
| 69 |
class Timer:
|
| 70 |
+
"""Utility class for measuring execution time of different operations"""
|
| 71 |
+
|
| 72 |
+
def __init__(self):
|
| 73 |
+
self.start_time = time.perf_counter()
|
| 74 |
+
self.checkpoints = [('Start', self.start_time)]
|
| 75 |
+
|
| 76 |
+
def checkpoint(self, label='Checkpoint'):
|
| 77 |
+
"""Add a checkpoint with a label"""
|
| 78 |
+
now = time.perf_counter()
|
| 79 |
+
self.checkpoints.append((label, now))
|
| 80 |
+
|
| 81 |
+
def report(self, is_clear_checkpoints=True):
|
| 82 |
+
"""Report time elapsed since last checkpoint"""
|
| 83 |
+
max_label_length = max(len(label) for (label, _) in self.checkpoints) if self.checkpoints else 0
|
| 84 |
+
prev_time = self.checkpoints[0][1] if self.checkpoints else self.start_time
|
| 85 |
+
|
| 86 |
+
for (label, curr_time) in self.checkpoints[1:]:
|
| 87 |
+
elapsed = curr_time - prev_time
|
| 88 |
+
print(f"{label.ljust(max_label_length)}: {elapsed:.3f} seconds")
|
| 89 |
+
prev_time = curr_time
|
| 90 |
+
|
| 91 |
+
if is_clear_checkpoints:
|
| 92 |
+
self.checkpoints.clear()
|
| 93 |
+
self.checkpoint()
|
| 94 |
+
|
| 95 |
+
def report_all(self):
|
| 96 |
+
"""Report all checkpoint times including total execution time"""
|
| 97 |
+
print('\n> Execution Time Report:')
|
| 98 |
+
max_label_length = max(len(label) for (label, _) in self.checkpoints) if len(self.checkpoints) > 0 else 0
|
| 99 |
+
prev_time = self.start_time
|
| 100 |
+
|
| 101 |
+
for (label, curr_time) in self.checkpoints[1:]:
|
| 102 |
+
elapsed = curr_time - prev_time
|
| 103 |
+
print(f"{label.ljust(max_label_length)}: {elapsed:.3f} seconds")
|
| 104 |
+
prev_time = curr_time
|
| 105 |
+
|
| 106 |
+
total_time = self.checkpoints[-1][1] - self.start_time if self.checkpoints else 0
|
| 107 |
+
print(f"{'Total Execution Time'.ljust(max_label_length)}: {total_time:.3f} seconds\n")
|
| 108 |
+
self.checkpoints.clear()
|
| 109 |
+
|
| 110 |
+
def restart(self):
|
| 111 |
+
"""Restart the timer"""
|
| 112 |
+
self.start_time = time.perf_counter()
|
| 113 |
+
self.checkpoints = [('Start', self.start_time)]
|
| 114 |
+
|
| 115 |
+
def parse_args() -> argparse.Namespace:
|
| 116 |
+
"""Parse command line arguments"""
|
| 117 |
+
parser = argparse.ArgumentParser()
|
| 118 |
+
parser.add_argument('--score-slider-step', type=float, default=0.05)
|
| 119 |
+
parser.add_argument('--score-general-threshold', type=float, default=0.35)
|
| 120 |
+
parser.add_argument('--score-character-threshold', type=float, default=0.85)
|
| 121 |
+
parser.add_argument('--share', action='store_true')
|
| 122 |
+
return parser.parse_args()
|
| 123 |
+
|
| 124 |
+
def load_labels(dataframe) -> tuple:
|
| 125 |
+
"""Load tag names and their category indexes from the dataframe"""
|
| 126 |
+
name_series = dataframe['name']
|
| 127 |
+
tag_names = name_series.tolist()
|
| 128 |
+
|
| 129 |
+
# Find indexes for different tag categories
|
| 130 |
+
rating_indexes = list(np.where(dataframe['category'] == 9)[0])
|
| 131 |
+
general_indexes = list(np.where(dataframe['category'] == 0)[0])
|
| 132 |
+
character_indexes = list(np.where(dataframe['category'] == 4)[0])
|
| 133 |
+
|
| 134 |
+
return tag_names, rating_indexes, general_indexes, character_indexes
|
| 135 |
+
|
| 136 |
+
def mcut_threshold(probs):
|
| 137 |
+
"""Calculate threshold using Maximum Change in second derivative (MCut) method"""
|
| 138 |
+
sorted_probs = probs[probs.argsort()[::-1]]
|
| 139 |
+
difs = sorted_probs[:-1] - sorted_probs[1:]
|
| 140 |
+
t = difs.argmax()
|
| 141 |
+
thresh = (sorted_probs[t] + sorted_probs[t + 1]) / 2
|
| 142 |
+
return thresh
|
| 143 |
+
|
| 144 |
+
def _download_model_files(model_repo):
|
| 145 |
+
"""Download model files from HuggingFace Hub"""
|
| 146 |
+
csv_path = huggingface_hub.hf_hub_download(model_repo, LABEL_FILENAME)
|
| 147 |
+
model_path = huggingface_hub.hf_hub_download(model_repo, MODEL_FILENAME)
|
| 148 |
+
return csv_path, model_path
|
| 149 |
+
|
| 150 |
+
def create_optimized_ort_session(model_path):
|
| 151 |
+
"""Create an optimized ONNX Runtime session with GPU support"""
|
| 152 |
+
# Configure session options for better performance
|
| 153 |
+
sess_options = ort.SessionOptions()
|
| 154 |
+
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
| 155 |
+
sess_options.intra_op_num_threads = 0 # Use all available cores
|
| 156 |
+
sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
|
| 157 |
+
sess_options.enable_mem_pattern = True
|
| 158 |
+
sess_options.enable_cpu_mem_arena = True
|
| 159 |
+
|
| 160 |
+
# Check available providers
|
| 161 |
+
available_providers = ort.get_available_providers()
|
| 162 |
+
print(f"Available ONNX Runtime providers: {available_providers}")
|
| 163 |
+
|
| 164 |
+
# Configure execution providers (prefer CUDA if available)
|
| 165 |
+
providers = []
|
| 166 |
+
|
| 167 |
+
# Use CUDA if available
|
| 168 |
+
if 'CUDAExecutionProvider' in available_providers:
|
| 169 |
+
providers.append('CUDAExecutionProvider')
|
| 170 |
+
print("Using CUDA provider for ONNX inference")
|
| 171 |
+
else:
|
| 172 |
+
print("CUDA provider not available, falling back to CPU")
|
| 173 |
+
|
| 174 |
+
# Always include CPU as fallback
|
| 175 |
+
providers.append('CPUExecutionProvider')
|
| 176 |
+
|
| 177 |
+
try:
|
| 178 |
+
session = ort.InferenceSession(model_path, sess_options, providers=providers)
|
| 179 |
+
print(f"Model loaded with providers: {session.get_providers()}")
|
| 180 |
+
return session
|
| 181 |
+
except Exception as e:
|
| 182 |
+
print(f"Failed to create ONNX session: {e}")
|
| 183 |
+
raise
|
| 184 |
+
|
| 185 |
+
def _load_model_components_optimized(model_repo):
|
| 186 |
+
"""Load and optimize model components"""
|
| 187 |
+
global CURRENT_MODEL, CURRENT_MODEL_NAME, CURRENT_TAGS_DF, CURRENT_TAG_NAMES
|
| 188 |
+
global CURRENT_RATING_INDEXES, CURRENT_GENERAL_INDEXES, CURRENT_CHARACTER_INDEXES, CURRENT_MODEL_TARGET_SIZE
|
| 189 |
+
|
| 190 |
+
# Only reload if model changed
|
| 191 |
+
if model_repo == CURRENT_MODEL_NAME and CURRENT_MODEL is not None:
|
| 192 |
+
return
|
| 193 |
+
|
| 194 |
+
# Download files
|
| 195 |
+
csv_path, model_path = _download_model_files(model_repo)
|
| 196 |
+
|
| 197 |
+
# Load optimized ONNX model
|
| 198 |
+
CURRENT_MODEL = create_optimized_ort_session(model_path)
|
| 199 |
+
|
| 200 |
+
# Load tags
|
| 201 |
+
tags_df = pd.read_csv(csv_path)
|
| 202 |
+
tag_names, rating_indexes, general_indexes, character_indexes = load_labels(tags_df)
|
| 203 |
+
|
| 204 |
+
# Store in global variables
|
| 205 |
+
CURRENT_TAGS_DF = tags_df
|
| 206 |
+
CURRENT_TAG_NAMES = tag_names
|
| 207 |
+
CURRENT_RATING_INDEXES = rating_indexes
|
| 208 |
+
CURRENT_GENERAL_INDEXES = general_indexes
|
| 209 |
+
CURRENT_CHARACTER_INDEXES = character_indexes
|
| 210 |
+
|
| 211 |
+
# Get model input size
|
| 212 |
+
_, height, width, _ = CURRENT_MODEL.get_inputs()[0].shape
|
| 213 |
+
CURRENT_MODEL_TARGET_SIZE = height
|
| 214 |
+
CURRENT_MODEL_NAME = model_repo
|
| 215 |
+
|
| 216 |
+
def _raw_predict(image_array, model_session):
|
| 217 |
+
"""Run raw prediction using the model session"""
|
| 218 |
+
input_name = model_session.get_inputs()[0].name
|
| 219 |
+
label_name = model_session.get_outputs()[0].name
|
| 220 |
+
preds = model_session.run([label_name], {input_name: image_array})[0]
|
| 221 |
+
return preds[0].astype(float)
|
| 222 |
+
|
| 223 |
+
def unload_model():
|
| 224 |
+
"""Explicitly unload the current model from memory"""
|
| 225 |
+
global CURRENT_MODEL, CURRENT_MODEL_NAME, CURRENT_TAGS_DF, CURRENT_TAG_NAMES
|
| 226 |
+
global CURRENT_RATING_INDEXES, CURRENT_GENERAL_INDEXES, CURRENT_CHARACTER_INDEXES, CURRENT_MODEL_TARGET_SIZE
|
| 227 |
+
|
| 228 |
+
# Delete the model session
|
| 229 |
+
if CURRENT_MODEL is not None:
|
| 230 |
+
del CURRENT_MODEL
|
| 231 |
+
CURRENT_MODEL = None
|
| 232 |
+
|
| 233 |
+
# Clear other large objects
|
| 234 |
+
CURRENT_TAGS_DF = None
|
| 235 |
+
CURRENT_TAG_NAMES = None
|
| 236 |
+
CURRENT_RATING_INDEXES = None
|
| 237 |
+
CURRENT_GENERAL_INDEXES = None
|
| 238 |
+
CURRENT_CHARACTER_INDEXES = None
|
| 239 |
+
CURRENT_MODEL_TARGET_SIZE = None
|
| 240 |
+
CURRENT_MODEL_NAME = None
|
| 241 |
+
|
| 242 |
+
# Force garbage collection
|
| 243 |
+
import gc
|
| 244 |
+
gc.collect()
|
| 245 |
+
|
| 246 |
+
# Clear CUDA cache if using GPU
|
| 247 |
+
try:
|
| 248 |
+
import torch
|
| 249 |
+
if torch.cuda.is_available():
|
| 250 |
+
torch.cuda.empty_cache()
|
| 251 |
+
except ImportError:
|
| 252 |
+
pass
|
| 253 |
+
|
| 254 |
+
def cleanup_after_processing():
|
| 255 |
+
"""Cleanup resources after processing"""
|
| 256 |
+
unload_model()
|
| 257 |
+
|
| 258 |
class Predictor:
|
| 259 |
+
"""Main predictor class for handling image tagging"""
|
| 260 |
+
|
| 261 |
def __init__(self):
|
| 262 |
+
self.model_components = None
|
| 263 |
self.last_loaded_repo = None
|
| 264 |
+
|
|
|
|
|
|
|
|
|
|
| 265 |
def load_model(self, model_repo):
|
| 266 |
+
"""Load model if not already loaded"""
|
| 267 |
+
if model_repo == self.last_loaded_repo and self.model_components is not None:
|
| 268 |
return
|
| 269 |
+
_load_model_components_optimized(model_repo)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 270 |
self.last_loaded_repo = model_repo
|
| 271 |
+
|
| 272 |
def prepare_image(self, path):
|
| 273 |
+
"""Prepare image for model input"""
|
| 274 |
image = Image.open(path)
|
| 275 |
+
image = image.convert('RGBA')
|
| 276 |
+
target_size = CURRENT_MODEL_TARGET_SIZE
|
| 277 |
+
|
| 278 |
+
# Create white background and composite
|
| 279 |
+
canvas = Image.new('RGBA', image.size, (255, 255, 255))
|
| 280 |
canvas.alpha_composite(image)
|
| 281 |
+
image = canvas.convert('RGB')
|
| 282 |
+
|
| 283 |
+
# Pad to square
|
| 284 |
image_shape = image.size
|
| 285 |
max_dim = max(image_shape)
|
| 286 |
pad_left = (max_dim - image_shape[0]) // 2
|
| 287 |
pad_top = (max_dim - image_shape[1]) // 2
|
| 288 |
+
padded_image = Image.new('RGB', (max_dim, max_dim), (255, 255, 255))
|
| 289 |
padded_image.paste(image, (pad_left, pad_top))
|
| 290 |
+
|
| 291 |
+
# Resize if needed
|
| 292 |
if max_dim != target_size:
|
| 293 |
+
padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC)
|
| 294 |
+
|
| 295 |
+
# Convert to array and preprocess
|
|
|
|
|
|
|
| 296 |
image_array = np.asarray(padded_image, dtype=np.float32)
|
| 297 |
+
image_array = image_array[:, :, ::-1] # BGR to RGB
|
|
|
|
| 298 |
return np.expand_dims(image_array, axis=0)
|
| 299 |
|
| 300 |
def create_file(self, content: str, directory: str, fileName: str) -> str:
|
| 301 |
+
"""Create a file with the given content"""
|
| 302 |
file_path = os.path.join(directory, fileName)
|
| 303 |
if fileName.endswith('.json'):
|
| 304 |
+
with open(file_path, 'w', encoding='utf-8') as file:
|
| 305 |
file.write(content)
|
| 306 |
else:
|
| 307 |
+
with open(file_path, 'w+', encoding='utf-8') as file:
|
| 308 |
file.write(content)
|
| 309 |
return file_path
|
| 310 |
+
|
| 311 |
+
def predict(self, gallery, model_repo, model_repo_2, general_thresh, general_mcut_enabled,
|
| 312 |
+
character_thresh, character_mcut_enabled, characters_merge_enabled,
|
| 313 |
+
additional_tags_prepend, additional_tags_append, tag_results, progress=gr.Progress()):
|
| 314 |
+
"""Main prediction function for processing images"""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 315 |
tag_results.clear()
|
| 316 |
gallery_len = len(gallery)
|
| 317 |
print(f"Predict load model: {model_repo}, gallery length: {gallery_len}")
|
| 318 |
+
|
| 319 |
+
timer = Timer()
|
| 320 |
+
progressRatio = 1
|
| 321 |
progressTotal = gallery_len + 1
|
| 322 |
current_progress = 0
|
|
|
|
|
|
|
|
|
|
| 323 |
txt_infos = []
|
| 324 |
output_dir = tempfile.mkdtemp()
|
| 325 |
+
|
| 326 |
if not os.path.exists(output_dir):
|
| 327 |
os.makedirs(output_dir)
|
| 328 |
+
|
| 329 |
+
# Load initial model
|
| 330 |
self.load_model(model_repo)
|
| 331 |
+
current_progress += progressRatio / progressTotal
|
| 332 |
+
progress(current_progress, desc='Initialize wd model finished')
|
| 333 |
+
timer.checkpoint("Initialize wd model")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 334 |
timer.report()
|
| 335 |
+
|
| 336 |
name_counters = defaultdict(int)
|
| 337 |
+
|
| 338 |
+
for (idx, value) in enumerate(gallery):
|
| 339 |
try:
|
| 340 |
+
# Handle duplicate filenames
|
| 341 |
image_path = value[0]
|
| 342 |
image_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
|
|
| 343 |
name_counters[image_name] += 1
|
| 344 |
if name_counters[image_name] > 1:
|
| 345 |
image_name = f"{image_name}_{name_counters[image_name]:02d}"
|
| 346 |
+
|
| 347 |
+
# Prepare image
|
| 348 |
image = self.prepare_image(image_path)
|
|
|
|
| 349 |
print(f"Gallery {idx:02d}: Starting run first model ({model_repo})...")
|
| 350 |
+
|
| 351 |
+
# Load and run first model
|
| 352 |
self.load_model(model_repo)
|
| 353 |
+
preds = _raw_predict(image, CURRENT_MODEL)
|
| 354 |
+
labels = list(zip(CURRENT_TAG_NAMES, preds))
|
| 355 |
+
|
| 356 |
+
# Process ratings
|
| 357 |
+
ratings_names = [labels[i] for i in CURRENT_RATING_INDEXES]
|
|
|
|
| 358 |
rating = dict(ratings_names)
|
| 359 |
|
| 360 |
+
# Process general tags
|
| 361 |
+
general_names = [labels[i] for i in CURRENT_GENERAL_INDEXES]
|
| 362 |
if general_mcut_enabled:
|
| 363 |
general_probs = np.array([x[1] for x in general_names])
|
| 364 |
general_thresh_temp = mcut_threshold(general_probs)
|
| 365 |
else:
|
| 366 |
general_thresh_temp = general_thresh
|
| 367 |
+
|
| 368 |
general_res = [x for x in general_names if x[1] > general_thresh_temp]
|
| 369 |
general_res = dict(general_res)
|
| 370 |
+
|
| 371 |
+
# Process character tags
|
| 372 |
+
character_names = [labels[i] for i in CURRENT_CHARACTER_INDEXES]
|
| 373 |
if character_mcut_enabled:
|
| 374 |
character_probs = np.array([x[1] for x in character_names])
|
| 375 |
character_thresh_temp = mcut_threshold(character_probs)
|
| 376 |
character_thresh_temp = max(0.15, character_thresh_temp)
|
| 377 |
else:
|
| 378 |
character_thresh_temp = character_thresh
|
| 379 |
+
|
| 380 |
character_res = [x for x in character_names if x[1] > character_thresh_temp]
|
| 381 |
character_res = dict(character_res)
|
|
|
|
| 382 |
character_list_1 = list(character_res.keys())
|
| 383 |
+
|
| 384 |
+
# Sort general tags by confidence
|
| 385 |
sorted_general_list_1 = sorted(general_res.items(), key=lambda x: x[1], reverse=True)
|
| 386 |
sorted_general_list_1 = [x[0] for x in sorted_general_list_1]
|
| 387 |
+
|
| 388 |
+
# Handle second model if provided
|
| 389 |
if model_repo_2 and model_repo_2 != model_repo:
|
| 390 |
print(f"Gallery {idx:02d}: Starting run second model ({model_repo_2})...")
|
| 391 |
self.load_model(model_repo_2)
|
| 392 |
+
preds_2 = _raw_predict(image, CURRENT_MODEL)
|
| 393 |
+
labels_2 = list(zip(CURRENT_TAG_NAMES, preds_2))
|
| 394 |
+
|
| 395 |
+
# Process general tags from second model
|
| 396 |
+
general_names_2 = [labels_2[i] for i in CURRENT_GENERAL_INDEXES]
|
| 397 |
if general_mcut_enabled:
|
| 398 |
general_probs_2 = np.array([x[1] for x in general_names_2])
|
| 399 |
general_thresh_temp_2 = mcut_threshold(general_probs_2)
|
| 400 |
else:
|
| 401 |
general_thresh_temp_2 = general_thresh
|
| 402 |
+
|
| 403 |
general_res_2 = [x for x in general_names_2 if x[1] > general_thresh_temp_2]
|
| 404 |
general_res_2 = dict(general_res_2)
|
| 405 |
+
|
| 406 |
+
# Process character tags from second model
|
| 407 |
+
character_names_2 = [labels_2[i] for i in CURRENT_CHARACTER_INDEXES]
|
| 408 |
if character_mcut_enabled:
|
| 409 |
character_probs_2 = np.array([x[1] for x in character_names_2])
|
| 410 |
character_thresh_temp_2 = mcut_threshold(character_probs_2)
|
| 411 |
character_thresh_temp_2 = max(0.15, character_thresh_temp_2)
|
| 412 |
else:
|
| 413 |
character_thresh_temp_2 = character_thresh
|
| 414 |
+
|
| 415 |
character_res_2 = [x for x in character_names_2 if x[1] > character_thresh_temp_2]
|
| 416 |
character_res_2 = dict(character_res_2)
|
|
|
|
| 417 |
character_list_2 = list(character_res_2.keys())
|
| 418 |
+
|
| 419 |
+
# Sort general tags from second model
|
| 420 |
sorted_general_list_2 = sorted(general_res_2.items(), key=lambda x: x[1], reverse=True)
|
| 421 |
sorted_general_list_2 = [x[0] for x in sorted_general_list_2]
|
| 422 |
+
|
| 423 |
+
# Combine results from both models
|
| 424 |
combined_character_list = list(set(character_list_1 + character_list_2))
|
| 425 |
combined_general_list = list(set(sorted_general_list_1 + sorted_general_list_2))
|
| 426 |
else:
|
|
|
|
| 427 |
combined_character_list = character_list_1
|
| 428 |
combined_general_list = sorted_general_list_1
|
| 429 |
+
|
| 430 |
+
# Remove characters from general tags if merging is disabled
|
| 431 |
+
if not characters_merge_enabled:
|
| 432 |
+
combined_character_list = [item for item in combined_character_list
|
| 433 |
+
if item not in combined_general_list]
|
| 434 |
+
|
| 435 |
+
# Handle additional tags
|
| 436 |
+
prepend_list = [tag.strip() for tag in additional_tags_prepend.split(',') if tag.strip()]
|
| 437 |
+
append_list = [tag.strip() for tag in additional_tags_append.split(',') if tag.strip()]
|
| 438 |
+
|
| 439 |
+
# Avoid duplicates in prepend/append lists
|
| 440 |
if prepend_list and append_list:
|
| 441 |
append_list = [item for item in append_list if item not in prepend_list]
|
| 442 |
+
|
| 443 |
+
# Remove prepended tags from main list
|
| 444 |
if prepend_list:
|
| 445 |
combined_general_list = [item for item in combined_general_list if item not in prepend_list]
|
| 446 |
+
|
| 447 |
+
# Remove appended tags from main list
|
| 448 |
if append_list:
|
| 449 |
combined_general_list = [item for item in combined_general_list if item not in append_list]
|
| 450 |
+
|
| 451 |
+
# Combine all tags
|
| 452 |
combined_general_list = prepend_list + combined_general_list + append_list
|
| 453 |
+
|
| 454 |
+
# Format output string
|
| 455 |
+
sorted_general_strings = ', '.join(
|
| 456 |
+
(combined_character_list if characters_merge_enabled else []) +
|
| 457 |
+
combined_general_list
|
| 458 |
+
).replace('(', '\\(').replace(')', '\\)').replace('_', ' ')
|
| 459 |
+
|
| 460 |
+
# Generate categorized output
|
| 461 |
+
categorized_strings = categorize_tags_output(sorted_general_strings, character_res).replace('(', '\\(').replace(')', '\\)')
|
| 462 |
+
categorized_json = generate_tags_json(sorted_general_strings, character_res)
|
| 463 |
+
|
| 464 |
+
# Create output files
|
| 465 |
+
txt_content = f"Output (string): {sorted_general_strings}\n\nCategorized Output: {categorized_strings}"
|
| 466 |
txt_file = self.create_file(txt_content, output_dir, f"{image_name}_output.txt")
|
| 467 |
+
txt_infos.append({'path': txt_file, 'name': f"{image_name}_output.txt"})
|
| 468 |
+
|
| 469 |
+
# Save image copy
|
|
|
|
|
|
|
|
|
|
| 470 |
image_path = value[0]
|
| 471 |
image = Image.open(image_path)
|
| 472 |
+
image.save(os.path.join(output_dir, f"{image_name}.png"), format='PNG')
|
| 473 |
+
txt_infos.append({'path': os.path.join(output_dir, f"{image_name}.png"), 'name': f"{image_name}.png"})
|
| 474 |
+
|
| 475 |
+
# Create tags text file
|
| 476 |
+
txt_file = self.create_file(sorted_general_strings, output_dir, image_name + '.txt')
|
| 477 |
+
# Create categorized tags file
|
| 478 |
+
categorized_file = self.create_file(categorized_strings, output_dir, f"{image_name}_categorized.txt")
|
| 479 |
+
txt_infos.append({'path': categorized_file, 'name': f"{image_name}_categorized.txt"})
|
| 480 |
+
txt_infos.append({'path': txt_file, 'name': image_name + '.txt'})
|
| 481 |
+
|
| 482 |
+
# Create JSON file
|
| 483 |
+
json_content = json.dumps(categorized_json, indent=2, ensure_ascii=False)
|
| 484 |
+
json_file = self.create_file(json_content, output_dir, f"{image_name}_categorized.json")
|
| 485 |
+
txt_infos.append({'path': json_file, 'name': f"{image_name}_categorized.json"})
|
| 486 |
+
|
| 487 |
+
# Store results
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 488 |
tag_results[image_path] = {
|
| 489 |
+
'strings': sorted_general_strings,
|
| 490 |
+
'categorized_strings': categorized_strings,
|
| 491 |
+
'categorized_json': categorized_json,
|
| 492 |
+
'rating': rating,
|
| 493 |
+
'character_res': character_res,
|
| 494 |
+
'general_res': general_res
|
|
|
|
|
|
|
| 495 |
}
|
| 496 |
+
|
| 497 |
+
# Update progress
|
| 498 |
+
current_progress += progressRatio / progressTotal
|
| 499 |
+
progress(current_progress, desc=f"image{idx:02d}, predict finished")
|
| 500 |
+
timer.checkpoint(f"image{idx:02d}, predict finished")
|
| 501 |
timer.report()
|
| 502 |
+
|
| 503 |
except Exception as e:
|
| 504 |
print(traceback.format_exc())
|
| 505 |
+
print('Error predict: ' + str(e))
|
| 506 |
+
|
| 507 |
+
# Create download zip
|
| 508 |
download = []
|
| 509 |
if txt_infos is not None and len(txt_infos) > 0:
|
| 510 |
+
downloadZipPath = os.path.join(
|
| 511 |
+
output_dir,
|
| 512 |
+
'Multi-Tagger-' + datetime.now().strftime('%Y%m%d-%H%M%S') + '.zip'
|
| 513 |
+
)
|
| 514 |
with zipfile.ZipFile(downloadZipPath, 'w', zipfile.ZIP_DEFLATED) as taggers_zip:
|
| 515 |
for info in txt_infos:
|
| 516 |
+
taggers_zip.write(info['path'], arcname=info['name'])
|
| 517 |
+
# If using GPU, model will auto unload after zip file creation
|
| 518 |
+
cleanup_after_processing() # Comment here to turn off this behavior
|
| 519 |
download.append(downloadZipPath)
|
|
|
|
| 520 |
|
|
|
|
|
|
|
|
|
|
| 521 |
progress(1, desc=f"Predict completed")
|
| 522 |
+
timer.report_all()
|
| 523 |
+
print('Predict is complete.')
|
| 524 |
+
|
| 525 |
+
# Return first image results as default
|
| 526 |
+
first_image_results = '', {}, {}, {}, '', {}
|
| 527 |
+
if gallery and len(gallery) > 0:
|
| 528 |
+
first_image_path = gallery[0][0]
|
| 529 |
+
if first_image_path in tag_results:
|
| 530 |
+
first_result = tag_results[first_image_path]
|
| 531 |
+
character_tags_formatted = ", ".join([name.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
|
| 532 |
+
for name in first_result['character_res'].keys()])
|
| 533 |
+
first_image_results = (
|
| 534 |
+
first_result['strings'],
|
| 535 |
+
first_result['rating'],
|
| 536 |
+
character_tags_formatted,
|
| 537 |
+
first_result['general_res'],
|
| 538 |
+
first_result.get('categorized_strings', ''),
|
| 539 |
+
first_result.get('categorized_json', {})
|
| 540 |
+
)
|
| 541 |
+
|
| 542 |
+
|
| 543 |
+
return (
|
| 544 |
+
download,
|
| 545 |
+
first_image_results[0],
|
| 546 |
+
first_image_results[1],
|
| 547 |
+
first_image_results[2],
|
| 548 |
+
first_image_results[3],
|
| 549 |
+
first_image_results[4],
|
| 550 |
+
first_image_results[5],
|
| 551 |
+
tag_results
|
| 552 |
+
)
|
| 553 |
+
|
| 554 |
def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state: gr.SelectData):
|
| 555 |
+
# Return first image results if no selection
|
| 556 |
+
if not selected_state and gallery and len(gallery) > 0:
|
| 557 |
+
first_image_path = gallery[0][0]
|
| 558 |
+
if first_image_path in tag_results:
|
| 559 |
+
first_result = tag_results[first_image_path]
|
| 560 |
+
character_tags_formatted = ", ".join([name.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
|
| 561 |
+
for name in first_result['character_res'].keys()])
|
| 562 |
+
return (
|
| 563 |
+
first_result['strings'],
|
| 564 |
+
first_result['rating'],
|
| 565 |
+
character_tags_formatted,
|
| 566 |
+
first_result['general_res'],
|
| 567 |
+
first_result.get('categorized_strings', ''),
|
| 568 |
+
first_result.get('categorized_json', {})
|
| 569 |
+
)
|
| 570 |
+
|
| 571 |
if not selected_state:
|
| 572 |
+
return '', {}, '', {}, '', {}
|
| 573 |
+
|
| 574 |
+
# Get selected image path
|
| 575 |
+
selected_value = selected_state.value
|
| 576 |
+
image_path = None
|
| 577 |
+
|
| 578 |
+
if isinstance(selected_value, dict) and 'image' in selected_value:
|
| 579 |
+
image_path = selected_value['image']['path']
|
| 580 |
+
elif isinstance(selected_value, (list, tuple)) and len(selected_value) > 0:
|
| 581 |
+
image_path = selected_value[0]
|
| 582 |
+
else:
|
| 583 |
+
image_path = str(selected_value)
|
| 584 |
+
|
| 585 |
+
# Return stored results
|
| 586 |
+
if image_path in tag_results:
|
| 587 |
+
result = tag_results[image_path]
|
| 588 |
+
|
| 589 |
+
character_tags_formatted = ", ".join([name.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
|
| 590 |
+
for name in result['character_res'].keys()])
|
| 591 |
+
return (
|
| 592 |
+
result['strings'],
|
| 593 |
+
result['rating'],
|
| 594 |
+
character_tags_formatted,
|
| 595 |
+
result['general_res'],
|
| 596 |
+
result.get('categorized_strings', ''),
|
| 597 |
+
result.get('categorized_json', {})
|
| 598 |
+
)
|
| 599 |
+
|
| 600 |
+
return '', {}, '', {}, '', {}
|
| 601 |
+
|
| 602 |
+
def append_gallery(gallery: list, image: str):
|
| 603 |
+
"""Add a single image to the gallery"""
|
| 604 |
+
if gallery is None:
|
| 605 |
+
gallery = []
|
| 606 |
+
if not image:
|
| 607 |
+
return gallery, None
|
| 608 |
+
gallery.append(image)
|
| 609 |
+
return gallery, None
|
| 610 |
+
|
| 611 |
+
def extend_gallery(gallery: list, images):
|
| 612 |
+
"""Add multiple images to the gallery"""
|
| 613 |
+
if gallery is None:
|
| 614 |
+
gallery = []
|
| 615 |
+
if not images:
|
| 616 |
+
return gallery
|
| 617 |
+
gallery.extend(images)
|
| 618 |
+
return gallery
|
| 619 |
+
|
| 620 |
+
# Parse arguments and initialize predictor
|
| 621 |
args = parse_args()
|
| 622 |
predictor = Predictor()
|
| 623 |
dropdown_list = [
|
| 624 |
+
EVA02_LARGE_MODEL_DSV3_REPO, VIT_LARGE_MODEL_DSV3_REPO, SWINV2_MODEL_DSV3_REPO,
|
| 625 |
+
CONV_MODEL_DSV3_REPO, VIT_MODEL_DSV3_REPO, MOAT_MODEL_DSV2_REPO,
|
| 626 |
+
SWIN_MODEL_DSV2_REPO, CONV_MODEL_DSV2_REPO, CONV2_MODEL_DSV2_REPO,
|
| 627 |
+
VIT_MODEL_DSV2_REPO, EVA02_LARGE_MODEL_IS_DSV1_REPO, SWINV2_MODEL_IS_DSV1_REPO
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 628 |
]
|
| 629 |
+
|
| 630 |
def _restart_space():
|
| 631 |
+
"""Restart the HuggingFace Space periodically for stability"""
|
| 632 |
+
HF_TOKEN = os.getenv('HF_TOKEN')
|
| 633 |
+
if not HF_TOKEN:
|
| 634 |
+
raise ValueError('HF_TOKEN environment variable is not set.')
|
| 635 |
+
huggingface_hub.HfApi().restart_space(
|
| 636 |
+
repo_id='Werli/Multi-Tagger',
|
| 637 |
+
token=HF_TOKEN,
|
| 638 |
+
factory_reboot=False
|
| 639 |
+
)
|
| 640 |
+
|
| 641 |
+
# Setup scheduler for periodic restarts
|
| 642 |
+
scheduler = BackgroundScheduler()
|
| 643 |
+
restart_space_job = scheduler.add_job(_restart_space, 'interval', seconds=172800)
|
| 644 |
scheduler.start()
|
| 645 |
+
next_run_time_utc = restart_space_job.next_run_time.astimezone(timezone.utc)
|
| 646 |
+
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."
|
| 647 |
|
| 648 |
+
with gr.Blocks(title=TITLE, css=css, theme='Werli/Purple-Crimson-Gradio-Theme', fill_width=True) as demo:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 649 |
gr.Markdown(value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
|
|
|
|
| 650 |
gr.Markdown(value=f"<p style='text-align: center;'>{DESCRIPTION}</p>")
|
| 651 |
+
|
| 652 |
+
with gr.Tab(label='Waifu Diffusion'):
|
| 653 |
with gr.Row():
|
| 654 |
with gr.Column():
|
| 655 |
+
|
| 656 |
+
with gr.Column(variant='panel'):
|
| 657 |
+
image_input = gr.Image(
|
| 658 |
+
label='Upload an Image or clicking paste from clipboard button',
|
| 659 |
+
type='filepath',
|
| 660 |
+
sources=['upload', 'clipboard'],
|
| 661 |
+
height=150
|
| 662 |
+
)
|
| 663 |
with gr.Row():
|
| 664 |
+
upload_button = gr.UploadButton(
|
| 665 |
+
'Upload multiple images',
|
| 666 |
+
file_types=['image'],
|
| 667 |
+
file_count='multiple',
|
| 668 |
+
size='sm'
|
| 669 |
+
)
|
| 670 |
gallery = gr.Gallery(
|
| 671 |
columns=2,
|
| 672 |
+
show_share_button=False,
|
| 673 |
+
interactive=True,
|
| 674 |
+
height='auto',
|
| 675 |
+
label='Grid of images',
|
| 676 |
+
preview=False,
|
| 677 |
+
elem_id='custom-gallery'
|
| 678 |
+
)
|
| 679 |
+
submit = gr.Button(value='Analyze Images', variant='primary', size='lg')
|
| 680 |
+
with gr.Column(variant='panel'):
|
| 681 |
+
model_repo = gr.Dropdown(
|
| 682 |
+
dropdown_list,
|
| 683 |
+
value=EVA02_LARGE_MODEL_DSV3_REPO,
|
| 684 |
+
label='1st Model'
|
| 685 |
)
|
| 686 |
+
PLUS = '+?'
|
|
|
|
|
|
|
| 687 |
gr.Markdown(value=f"<p style='text-align: center;'>{PLUS}</p>")
|
| 688 |
+
model_repo_2 = gr.Dropdown(
|
| 689 |
+
[None] + dropdown_list,
|
| 690 |
+
value=None,
|
| 691 |
+
label='2nd Model (Optional)',
|
| 692 |
+
info='Select another model for diversified results.'
|
| 693 |
+
)
|
| 694 |
+
|
| 695 |
with gr.Row():
|
| 696 |
+
general_thresh = gr.Slider(
|
| 697 |
+
0, 1,
|
| 698 |
+
step=args.score_slider_step,
|
| 699 |
+
value=args.score_general_threshold,
|
| 700 |
+
label='General Tags Threshold',
|
| 701 |
+
scale=3
|
| 702 |
+
)
|
| 703 |
+
general_mcut_enabled = gr.Checkbox(
|
| 704 |
+
value=False,
|
| 705 |
+
label='Use MCut threshold',
|
| 706 |
+
scale=1
|
| 707 |
+
)
|
| 708 |
+
|
| 709 |
with gr.Row():
|
| 710 |
+
character_thresh = gr.Slider(
|
| 711 |
+
0, 1,
|
| 712 |
+
step=args.score_slider_step,
|
| 713 |
+
value=args.score_character_threshold,
|
| 714 |
+
label='Character Tags Threshold',
|
| 715 |
+
scale=3
|
| 716 |
+
)
|
| 717 |
+
character_mcut_enabled = gr.Checkbox(
|
| 718 |
+
value=False,
|
| 719 |
+
label='Use MCut threshold',
|
| 720 |
+
scale=1
|
| 721 |
+
)
|
| 722 |
+
|
| 723 |
with gr.Row():
|
| 724 |
+
characters_merge_enabled = gr.Checkbox(
|
| 725 |
+
value=False,
|
| 726 |
+
label='Merge characters into the string output',
|
| 727 |
+
scale=1
|
| 728 |
+
)
|
| 729 |
+
|
| 730 |
with gr.Row():
|
| 731 |
+
additional_tags_prepend = gr.Text(
|
| 732 |
+
label='Prepend Additional tags (comma split)'
|
| 733 |
+
)
|
| 734 |
+
additional_tags_append = gr.Text(
|
| 735 |
+
label='Append Additional tags (comma split)'
|
| 736 |
+
)
|
| 737 |
+
|
| 738 |
with gr.Row():
|
| 739 |
clear = gr.ClearButton(
|
| 740 |
+
components=[
|
| 741 |
+
gallery, model_repo, general_thresh, general_mcut_enabled,
|
| 742 |
+
character_thresh, character_mcut_enabled, characters_merge_enabled,
|
| 743 |
+
additional_tags_prepend, additional_tags_append
|
| 744 |
+
],
|
| 745 |
+
variant='secondary',
|
| 746 |
+
size='lg'
|
| 747 |
+
)
|
| 748 |
+
|
| 749 |
+
with gr.Column(variant='panel'):
|
| 750 |
+
download_file = gr.File(label='Download')
|
| 751 |
+
character_res = gr.Textbox(
|
| 752 |
+
label="Character tags",
|
| 753 |
+
show_copy_button=True,
|
| 754 |
+
lines=3
|
| 755 |
+
)
|
| 756 |
+
sorted_general_strings = gr.Textbox(
|
| 757 |
+
label='Output',
|
| 758 |
+
show_label=True,
|
| 759 |
+
show_copy_button=True,
|
| 760 |
+
lines=5
|
| 761 |
+
)
|
| 762 |
+
categorized_strings = gr.Textbox(
|
| 763 |
+
label='Categorized',
|
| 764 |
+
show_label=True,
|
| 765 |
+
show_copy_button=True,
|
| 766 |
+
lines=5
|
| 767 |
+
)
|
| 768 |
+
tags_json = gr.JSON(
|
| 769 |
+
label='Categorized Tags (JSON)',
|
| 770 |
+
visible=True
|
| 771 |
+
)
|
| 772 |
+
rating = gr.Label(label='Rating')
|
| 773 |
+
general_res = gr.Textbox(
|
| 774 |
+
label="General tags",
|
| 775 |
+
show_copy_button=True,
|
| 776 |
+
lines=3,
|
| 777 |
+
visible=False # Temp
|
| 778 |
+
)
|
| 779 |
+
# State to store results
|
| 780 |
tag_results = gr.State({})
|
| 781 |
+
|
| 782 |
+
# Event handlers
|
| 783 |
+
image_input.change(
|
| 784 |
+
append_gallery,
|
| 785 |
+
inputs=[gallery, image_input],
|
| 786 |
+
outputs=[gallery, image_input]
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
upload_button.upload(
|
| 790 |
+
extend_gallery,
|
| 791 |
+
inputs=[gallery, upload_button],
|
| 792 |
+
outputs=gallery
|
| 793 |
+
)
|
| 794 |
+
|
| 795 |
+
gallery.select(
|
| 796 |
+
get_selection_from_gallery,
|
| 797 |
+
inputs=[gallery, tag_results],
|
| 798 |
+
outputs=[sorted_general_strings, rating, character_res, general_res, categorized_strings, tags_json]
|
| 799 |
+
)
|
| 800 |
+
|
| 801 |
+
submit.click(
|
| 802 |
+
predictor.predict,
|
| 803 |
+
inputs=[
|
| 804 |
+
gallery, model_repo, model_repo_2, general_thresh, general_mcut_enabled,
|
| 805 |
+
character_thresh, character_mcut_enabled, characters_merge_enabled,
|
| 806 |
+
additional_tags_prepend, additional_tags_append, tag_results
|
| 807 |
+
],
|
| 808 |
+
outputs=[download_file, sorted_general_strings, rating, character_res, general_res, categorized_strings, tags_json, tag_results]
|
| 809 |
+
)
|
| 810 |
+
|
| 811 |
gr.Examples(
|
| 812 |
+
[['images/1girl.png', EVA02_LARGE_MODEL_DSV3_REPO, 0.35, False, 0.85, False]],
|
| 813 |
+
inputs=[image_input, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled]
|
| 814 |
+
)
|
| 815 |
+
gr.Markdown('[Based on SmilingWolf/wd-tagger](https://huggingface.co/spaces/SmilingWolf/wd-tagger) <p style="text-align:right"><a href="https://huggingface.co/spaces/John6666/danbooru-tags-transformer-v2-with-wd-tagger-b">Prompt Enhancer</a></p>')
|
| 816 |
+
with gr.Tab("PixAI"):
|
| 817 |
+
pixai_interface = create_pixai_interface()
|
| 818 |
with gr.Tab("Booru Image Fetcher"):
|
| 819 |
+
booru_interface = create_booru_interface()
|
| 820 |
+
|
| 821 |
+
gr.Markdown(NEXT_RESTART)
|
| 822 |
+
|
| 823 |
+
demo.queue(max_size=5).launch(show_error=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
modules/__pycache__/booru.cpython-311.pyc
ADDED
|
Binary file (13 kB). View file
|
|
|
modules/__pycache__/classifyTags.cpython-311.pyc
ADDED
|
Binary file (38.3 kB). View file
|
|
|
modules/__pycache__/pixai.cpython-311.pyc
ADDED
|
Binary file (39.5 kB). View file
|
|
|
modules/booru.py
CHANGED
|
@@ -2,6 +2,19 @@ import requests,re,base64,io,numpy as np
|
|
| 2 |
from PIL import Image,ImageOps
|
| 3 |
import torch,gradio as gr
|
| 4 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 5 |
# Helper to load image from URL
|
| 6 |
def loadImageFromUrl(url):
|
| 7 |
response = requests.get(url, timeout=10)
|
|
@@ -108,4 +121,41 @@ def on_select(evt: gr.SelectData, tags_list, post_url_list, image_url_list):
|
|
| 108 |
idx = evt.index
|
| 109 |
if idx < len(tags_list):
|
| 110 |
return tags_list[idx], post_url_list[idx], image_url_list[idx]
|
| 111 |
-
return "No tags", "", ""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
from PIL import Image,ImageOps
|
| 3 |
import torch,gradio as gr
|
| 4 |
|
| 5 |
+
# Custom CSS for gallery styling
|
| 6 |
+
css = """
|
| 7 |
+
#custom-gallery {--row-height: 180px;display: grid;grid-auto-rows: min-content;gap: 10px;}
|
| 8 |
+
#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;}
|
| 9 |
+
#custom-gallery .thumbnail-item:hover {transform: translateY(-3px);box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);}
|
| 10 |
+
#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;}
|
| 11 |
+
#custom-gallery .thumbnail-item img.portrait {max-width: 100%;}
|
| 12 |
+
#custom-gallery .thumbnail-item img.landscape {max-height: 100%;}
|
| 13 |
+
.gallery-container {max-height: 500px;overflow-y: auto;padding-right: 0px;--size-80: 500px;}
|
| 14 |
+
.thumbnails {display: flex;position: absolute;bottom: 0;width: 120px;overflow-x: scroll;padding-top: 320px;padding-bottom: 280px;padding-left: 4px;flex-wrap: wrap;}
|
| 15 |
+
#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;}
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
# Helper to load image from URL
|
| 19 |
def loadImageFromUrl(url):
|
| 20 |
response = requests.get(url, timeout=10)
|
|
|
|
| 121 |
idx = evt.index
|
| 122 |
if idx < len(tags_list):
|
| 123 |
return tags_list[idx], post_url_list[idx], image_url_list[idx]
|
| 124 |
+
return "No tags", "", ""
|
| 125 |
+
|
| 126 |
+
def create_booru_interface():
|
| 127 |
+
with gr.Blocks(css=css, fill_width=True) as demo:
|
| 128 |
+
with gr.Row():
|
| 129 |
+
with gr.Column():
|
| 130 |
+
gr.Markdown("### ⚙️ Search Parameters")
|
| 131 |
+
site = gr.Dropdown(label="Select Source", choices=["Gelbooru (Not working)", "Rule34", "Xbooru"], value="Xbooru")
|
| 132 |
+
Tags = gr.Textbox(label="Tags (comma-separated)", placeholder="e.g. solo, 1girl, 1boy, artist name, character, black hair, granblue fantasy, ...", lines=3)
|
| 133 |
+
exclude_tags = gr.Textbox(label="Exclude Tags (comma-separated)", placeholder="e.g. animated, watermark, username, ...", lines=3)
|
| 134 |
+
score = gr.Number(label="Minimum Score", value=0)
|
| 135 |
+
count = gr.Slider(label="Number of Images", minimum=1, maximum=20, step=1, value=1)
|
| 136 |
+
Safe = gr.Checkbox(label="Include Safe", value=True)
|
| 137 |
+
Questionable = gr.Checkbox(label="Include Questionable", value=True)
|
| 138 |
+
Explicit = gr.Checkbox(label="Include Explicit (18+)", value=False)
|
| 139 |
+
submit_btn = gr.Button("Fetch Images", variant="primary")
|
| 140 |
+
with gr.Column():
|
| 141 |
+
gr.Markdown("### 📄 Results")
|
| 142 |
+
images_output = gr.Gallery(
|
| 143 |
+
columns=2,
|
| 144 |
+
show_share_button=False,
|
| 145 |
+
interactive=True,
|
| 146 |
+
height='auto',
|
| 147 |
+
label='Grid of images',
|
| 148 |
+
preview=False,
|
| 149 |
+
elem_id='custom-gallery'
|
| 150 |
+
)
|
| 151 |
+
tags_output = gr.Textbox(label="Tags", placeholder="Select an image to display tags", lines=6, show_copy_button=True)
|
| 152 |
+
post_url_output = gr.Textbox(label="Post URL", lines=2, show_copy_button=True)
|
| 153 |
+
image_url_output = gr.Textbox(label="Image URL", lines=2, show_copy_button=True)
|
| 154 |
+
# State to store tags, URLs
|
| 155 |
+
tags_state = gr.State([])
|
| 156 |
+
post_url_state = gr.State([])
|
| 157 |
+
image_url_state = gr.State([])
|
| 158 |
+
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], )
|
| 159 |
+
images_output.select(fn=on_select, inputs=[tags_state, post_url_state, image_url_state], outputs=[tags_output, post_url_output, image_url_output], )
|
| 160 |
+
|
| 161 |
+
return demo
|
modules/classifyTags.py
CHANGED
|
@@ -1,156 +1,337 @@
|
|
| 1 |
-
|
| 2 |
-
import
|
| 3 |
-
|
| 4 |
-
#
|
| 5 |
-
|
| 6 |
-
'
|
| 7 |
-
'
|
| 8 |
-
'
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
'
|
| 13 |
-
'
|
| 14 |
-
'
|
| 15 |
-
'
|
| 16 |
-
'
|
| 17 |
-
'
|
| 18 |
-
'
|
| 19 |
-
'
|
| 20 |
-
'
|
| 21 |
-
'
|
| 22 |
-
'
|
| 23 |
-
'
|
| 24 |
-
'
|
| 25 |
-
'
|
| 26 |
-
'
|
| 27 |
-
'
|
| 28 |
-
'
|
| 29 |
-
'
|
| 30 |
-
'
|
| 31 |
-
'
|
| 32 |
-
'
|
| 33 |
-
'
|
| 34 |
-
'
|
| 35 |
-
'
|
| 36 |
-
'
|
| 37 |
-
'
|
| 38 |
-
'
|
| 39 |
-
'
|
| 40 |
-
'
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
node.
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
|
| 78 |
-
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
|
| 82 |
-
|
| 83 |
-
|
| 84 |
-
|
| 85 |
-
|
| 86 |
-
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
|
| 90 |
-
|
| 91 |
-
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
|
| 95 |
-
|
| 96 |
-
|
| 97 |
-
|
| 98 |
-
|
| 99 |
-
|
| 100 |
-
|
| 101 |
-
|
| 102 |
-
|
| 103 |
-
|
| 104 |
-
|
| 105 |
-
|
| 106 |
-
|
| 107 |
-
|
| 108 |
-
|
| 109 |
-
|
| 110 |
-
|
| 111 |
-
|
| 112 |
-
|
| 113 |
-
|
| 114 |
-
|
| 115 |
-
|
| 116 |
-
|
| 117 |
-
|
| 118 |
-
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
|
| 123 |
-
|
| 124 |
-
|
| 125 |
-
|
| 126 |
-
|
| 127 |
-
|
| 128 |
-
|
| 129 |
-
|
| 130 |
-
|
| 131 |
-
|
| 132 |
-
|
| 133 |
-
|
| 134 |
-
|
| 135 |
-
|
| 136 |
-
|
| 137 |
-
|
| 138 |
-
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
|
| 142 |
-
|
| 143 |
-
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
|
| 147 |
-
|
| 148 |
-
|
| 149 |
-
|
| 150 |
-
|
| 151 |
-
|
| 152 |
-
|
| 153 |
-
|
| 154 |
-
|
| 155 |
-
|
| 156 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import re
|
| 2 |
+
from collections import defaultdict
|
| 3 |
+
|
| 4 |
+
# Test: Define priority tags that should always come first
|
| 5 |
+
PRIORITY_TAGS = [
|
| 6 |
+
'1girl', '2girls', '3girls', '4girls', '5girls', '6+girls', 'multiple_girls',
|
| 7 |
+
'1boy', '2boys', '3boys', '4boys', '5boys', '6+boys', 'multiple_boys',
|
| 8 |
+
'male_focus', 'female_focus', 'other_focus'
|
| 9 |
+
]
|
| 10 |
+
|
| 11 |
+
categories = {
|
| 12 |
+
'Explicit':['sex', '69', 'paizuri', 'cum', 'precum', 'areola_slip', 'hetero', 'erection', 'oral', 'fellatio', 'yaoi', 'ejaculation', 'ejaculating', 'masturbation', 'handjob', 'bulge', 'rape', '_rape', 'doggystyle', 'threesome', 'missionary', 'object_insertion', 'nipple', 'nipples', 'pussy', 'anus', 'penis', 'groin', 'testicles', 'testicle', 'anal', 'cameltoe', 'areolae', 'dildo', 'clitoris', 'top-down_bottom-up', 'gag', 'groping', 'gagged', 'gangbang', 'orgasm', 'femdom', 'incest', 'bukkake', 'breast_out', 'vaginal', 'vagina', 'public_indecency', 'breast_sucking', 'folded', 'cunnilingus', '_cunnilingus', 'foreskin', 'bestiality', 'footjob', 'uterus', 'womb', 'flaccid', 'defloration', 'butt_plug', 'cowgirl_position', 'reverse_cowgirl_position', 'squatting_cowgirl_position', 'reverse_upright_straddle', 'irrumatio', 'deepthroat', 'pokephilia', 'gaping', 'orgy', 'cleft_of_venus', 'futanari', 'futasub', 'futa', 'cumdrip', 'fingering', 'vibrator', 'partially_visible_vulva', 'penetration', 'penetrated', 'cumshot', 'exhibitionism', 'breast_milk', 'grinding', 'clitoral', 'urethra', 'phimosis', 'cervix', 'impregnation', 'tribadism', 'molestation', 'pubic_hair', 'clothed_female_nude_male', 'clothed_male_nude_female', 'clothed_female_nude_female', 'clothed_male_nude_male', 'sex_machine', 'milking_machine', 'ovum', 'chikan', 'pussy_juice_drip_through_clothes', 'ejaculating_while_penetrated', 'suspended_congress', 'reverse_suspended_congress', 'spread_pussy_under_clothes', 'anilingus', 'reach-around', 'humping', 'consensual_tentacles', 'tentacle_pit', 'cum_in_'],
|
| 13 |
+
'Appearance Status':['backless', 'bandaged_neck', 'bleeding', 'blood', '_blood', 'blush', 'body_writing', 'bodypaint', 'bottomless', 'breath', 'bruise', 'butt_crack', 'cold', 'covered_mouth', 'crack', 'cross-section', 'crotchless', 'crying', 'curvy', 'cuts', 'dirty', 'dripping', 'drunk', 'from_mouth', 'glowing', 'hairy', 'halterneck', 'hot', 'injury', 'latex', 'leather', 'levitation', 'lipstick_mark', '_markings', 'makeup', 'mole', 'moles', 'no_bra', 'nosebleed', 'nude', 'outfit', 'pantylines', 'peeing', 'piercing', '_piercing', 'piercings', 'pregnant', 'public_nudity', 'reverse', '_skin', '_submerged', 'saliva', 'scar', 'scratches', 'see-through', 'shadow', 'shibari', 'sideless', 'skindentation', 'sleeping', 'tan', 'soap_bubbles', 'steam', 'steaming_body', 'stitches', 'sweat', 'sweatdrop', 'sweaty', 'tanlines', 'tattoo', 'tattoo', 'tears', 'topless', 'transparent', 'trefoil', 'trembling', 'veins', 'visible_air', 'wardrobe_malfunction', 'wet', 'x-ray', 'unconscious', 'handprint'],
|
| 14 |
+
'Action Pose':['afloat', 'afterimage', 'against_fourth_wall', 'against_wall', 'aiming', 'all_fours',"another's_mouth",'arm_', 'arm_support', 'arms_', 'arms_behind_back', 'asphyxiation', 'attack', 'back', 'ballet', 'bara', 'bathing', 'battle', 'bdsm', 'beckoning', 'bent_over', 'bite_mark', 'biting', 'bondage', 'breast_suppress', 'breathing', 'burning', 'bust_cup', 'carry', 'carrying', 'caught', 'chained', 'cheek_squash', 'chewing', 'cigarette', 'clapping', 'closed_eye', 'come_hither', 'cooking', 'covering', 'cuddling', 'dancing', '_docking', 'destruction', 'dorsiflexion', 'dreaming', 'dressing', 'drinking', 'driving', 'dropping', 'eating', 'exercise', 'expansion', 'exposure', 'facing', 'failure', 'fallen_down', 'falling', 'feeding', 'fetal_position', 'fighting', 'finger_on_trigger', 'finger_to_cheek', 'finger_to_mouth', 'firing', 'fishing', 'flashing', 'fleeing', 'flexible', 'flexing', 'floating', 'flying', 'fourth_wall', 'freediving', 'frogtie', '_grab', 'girl_on_top', 'giving', 'grabbing', 'grabbing_', 'gymnastics', '_hold', 'hadanugi_dousa', 'hairdressing', 'hand_', 'hand_on', 'hand_on_wall', 'hands_', 'headpat', 'hiding', 'holding', 'hug', 'hugging', 'imagining', 'in_container', 'in_mouth', 'in_palm', 'jealous', 'jumping', 'kabedon', 'kicking', 'kiss', 'kissing', 'kneeling', '_lift', 'lactation', 'laundry', 'licking', 'lifted_by_self', 'looking', 'lowleg', 'lying', 'melting', 'midair', 'moaning', '_open', 'on_back', 'on_bed', 'on_ground', 'on_lap', 'on_one_knee', 'one_eye_closed', 'open_', 'over_mouth', 'own_mouth', '_peek', '_pose', '_press', '_pull', 'padding', 'paint', 'painting_(action)', 'palms_together', 'pee', 'peeking', 'pervert', 'petting', 'pigeon-toed', 'piggyback', 'pinching', 'pinky_out', 'pinned', 'plantar_flexion', 'planted', 'playing', 'pocky', 'pointing', 'poke', 'poking', 'pouring', 'pov', 'praying', 'presenting', 'profanity', 'pulled_by_self', 'pulling', 'pump_action', 'punching', '_rest', 'raised', 'reaching', 'reading', 'reclining', 'reverse_grip', 'riding', 'running', '_slip', 'salute', 'screaming', 'seiza', 'selfie', 'sewing', 'shaking', 'shoe_dangle', 'shopping', 'shouting', 'showering', 'shushing', 'singing', 'sitting', 'slapping', 'smell', 'smelling', 'smoking', 'smother', 'solo', 'spanked', 'spill', 'spilling', 'spinning', 'splashing', 'split', 'squatting', 'squeezed', 'breasts_squeezed_together', 'standing', 'standing_on_', 'staring', 'straddling', 'strangling', 'stretching', 'surfing', 'suspension', 'swimming', 'talking', 'teardrop', 'tearing_clothes', 'throwing', 'tied_up', 'tiptoes', 'toe_scrunch', 'toothbrush', 'trigger_discipline', 'tripping', 'tsundere', 'turning_head', 'twitching', 'two-handed', 'tying', '_up', 'unbuttoned', 'undressed', 'undressing', 'unsheathed', 'unsheathing', 'unzipped', 'unzipping', 'upright_straddle', 'v', 'V', 'vore', '_wielding', 'wading', 'walk-in', 'walking', 'wariza', 'waving', 'wedgie', 'wrestling', 'writing', 'yawning', 'yokozuwari', '_conscious', 'massage', 'struggling', 'shrugging', 'drugged', 'tentacles_under_clothes', 'restrained_by_tentacles', 'tentacles_around_arms', 'tentacles_around_legs', 'restrained_legs', 'restrained_tail', 'restrained_arms', 'tentacles_on_female', 'archery', 'cleaning', 'tempura', 'facepalm', 'sadism'],
|
| 15 |
+
'Headwear':['antennae', 'antlers', 'aura', 'bandaged_head', 'bandana', 'bandeau', 'beanie', 'beanie', 'beret', 'bespectacled', 'blindfold', 'bonnet', '_cap', 'circlet', 'crown', '_drill', '_drills', 'diadem', '_eyewear', 'ear_covers', 'ear_ornament', 'ear_tag', 'earbuds', 'earclip', 'earmuffs', 'earphones', 'earpiece', 'earring', 'earrings', 'eyeliner', 'eyepatch', 'eyewear_on_head', 'facial', 'fedora', 'glasses', 'goggles', '_headwear', 'hachimaki', 'hair_', 'hair_bobbles', 'hair_ornament', 'hair_rings', 'hair_tie', 'hairband', 'hairclip', 'hairpin', 'hairpods', 'halo', 'hat', 'head-mounted_display', 'head_wreath', 'headband', 'headdress', 'headgear', 'headphones', 'headpiece', 'headset', 'helm', 'helmet', 'hood', 'kabuto_(helmet)', 'kanzashi', '_mask', 'maid_headdress', 'mask', 'mask', 'mechanical_ears', 'mechanical_eye', 'mechanical_horns', 'mob_cap', 'monocle', 'neck_ruff', 'nightcap', 'on_head', 'pince-nez', 'qingdai_guanmao', 'scarf_over_mouth', 'scrunchie', 'sunglasses',"tam_o'_shanter",'tate_eboshi', 'tiara', 'topknot', 'turban', 'veil', 'visor', 'wig', 'mitre', 'tricorne', 'bicorne'],
|
| 16 |
+
'Handwear':['arm_warmers', 'armband', 'armlet', 'bandaged_arm', 'bandaged_fingers', 'bandaged_hand', 'bandaged_wrist', 'bangle', 'bracelet', 'bracelets', 'bracer', 'cuffs', 'elbow_pads', '_gauntlets', '_glove', '_gloves', 'gauntlets', 'gloves', 'kote', 'kurokote', 'mechanical_arm', 'mechanical_arms', 'mechanical_hands', 'mittens', 'mitts', 'nail_polish', 'prosthetic_arm', 'wrist_cuffs', 'wrist_guards', 'wristband', 'yugake'],
|
| 17 |
+
'One-Piece Outfit':['bodystocking', 'bodysuit', 'dress', 'furisode', 'gown', 'hanfu', 'jumpsuit', 'kimono', 'leotard', 'microdress', 'one-piece', 'overalls', 'robe', 'spacesuit', 'sundress', 'yukata'],
|
| 18 |
+
'Upper Body Clothing':['aiguillette', 'apron', '_apron', 'armor', '_armor', 'ascot', 'babydoll', 'bikini', '_bikini', 'blazer', '_blazer', 'blouse', '_blouse', 'bowtie', '_bowtie', 'bra', '_bra', 'breast_curtain', 'breast_curtains', 'breast_pocket', 'breastplate', 'bustier', 'camisole', 'cape', 'capelet', 'cardigan', 'center_opening', 'chemise', 'chest_jewel', 'choker', 'cloak', 'coat', 'coattails', 'collar', '_collar', 'corset', 'criss-cross_halter', 'crop_top', 'dougi', 'feather_boa', 'gakuran', 'hagoromo', 'hanten_(clothes)', 'haori', 'harem_pants', 'harness', 'hoodie', 'jacket', '_jacket', 'japanese_clothes', 'kappougi', 'kariginu', 'lapels', 'lingerie', '_lingerie', 'maid', 'mechanical_wings', 'mizu_happi', 'muneate', 'neckerchief', 'necktie', 'negligee', 'nightgown', 'pajamas', '_pajamas', 'pauldron', 'pauldrons', 'plunging_neckline', 'raincoat', 'rei_no_himo', 'sailor_collar', 'sarashi', 'scarf', 'serafuku', 'shawl', 'shirt', 'shoulder_', 'sleepwear', 'sleeve', 'sleeveless', 'sleeves', '_sleeves', 'sode', 'spaghetti_strap', 'sportswear', 'strapless', 'suit', 'sundress', 'suspenders', 'sweater', 'swimsuit', '_top', '_torso', 't-shirt', 'tabard', 'tailcoat', 'tank_top', 'tasuki', 'tie_clip', 'tunic', 'turtleneck', 'tuxedo', '_uniform', 'undershirt', 'uniform', 'v-neck', 'vambraces', 'vest', 'waistcoat'],
|
| 19 |
+
'Lower Body Clothing':['bare_hips', 'bloomers', 'briefs', 'buruma', 'crotch_seam', 'cutoffs', 'denim', 'faulds', 'fundoshi', 'g-string', 'garter_straps', 'hakama', 'hip_vent', 'jeans', 'knee_pads', 'loincloth', 'mechanical_tail', 'microskirt', 'miniskirt', 'overskirt', 'panties', 'pants', 'pantsu', 'panty_straps', 'pelvic_curtain', 'petticoat', 'sarong', 'shorts', 'side_slit', 'skirt', 'sweatpants', 'swim_trunks', 'thong', 'underwear', 'waist_cape'],
|
| 20 |
+
'Foot & Legwear':['anklet', 'bandaged_leg', 'boot', 'boots', '_footwear', 'flats', 'flip-flops', 'geta', 'greaves', '_heels', 'kneehigh', 'kneehighs', '_legwear', 'leg_warmers', 'leggings', 'loafers', 'mary_janes', 'mechanical_legs', 'okobo', 'over-kneehighs', 'pantyhose', 'prosthetic_leg', 'pumps', '_shoe', '_sock', 'sandals', 'shoes', 'skates', 'slippers', 'sneakers', 'socks', 'spikes', 'tabi', 'tengu-geta', 'thigh_strap', 'thighhighs', 'uwabaki', 'zouri', 'legband', 'ankleband'],
|
| 21 |
+
'Other Accessories':['alternate_', 'anklet', 'badge', 'beads', 'belt', 'belts', 'bow', '_bow', 'brooch', 'buckle', 'button', 'buttons', '_clothes', '_costume', '_cutout', 'casual', 'charm', 'clothes_writing', 'clothing_aside', 'costume', 'cow_print', 'cross', 'd-pad', 'double-breasted', 'drawstring', 'epaulettes', 'fabric', 'fishnets', 'floral_print', 'formal', 'frills', '_garter', 'gem', 'holster', 'jewelry', '_knot', 'lace', 'lanyard', 'leash', 'magatama', 'mechanical_parts', 'medal', 'medallion', 'naked_bandage', 'necklace', '_ornament', '(ornament)', 'o-ring', 'obi', 'obiage', 'obijime', '_pin', '_print', 'padlock', 'patterned_clothing', 'pendant', 'piercing', 'plaid', 'pocket', 'polka_dot', 'pom_pom_(clothes)', 'pom_pom_(clothes)', 'pouch', 'ribbon', '_ribbon', '_stripe', '_stripes', 'sash', 'shackles', 'shimenawa', 'shrug_(clothing)', 'skin_tight', 'spandex', 'strap', 'sweatband', '_trim', 'tassel', 'zettai_ryouiki', 'zipper'],
|
| 22 |
+
'Facial Expression':['ahegao', 'anger_vein', 'angry', 'annoyed', 'confused', 'drooling', 'embarrassed', 'expressionless', 'eye_contact', '_face', 'frown', 'fucked_silly', 'furrowed_brow', 'glaring', 'gloom_(expression)', 'grimace', 'grin', 'happy', 'jitome', 'laughing', '_mouth', 'nervous', 'notice_lines', 'o_o', 'parted_lips', 'pout', 'puff_of_air', 'restrained', 'sad', 'sanpaku', 'scared', 'scowl', 'serious', 'shaded_face', 'shy', 'sigh', 'sleepy', 'smile', 'smirk', 'smug', 'snot', 'spoken_ellipsis', 'spoken_exclamation_mark', 'spoken_interrobang', 'spoken_question_mark', 'squiggle', 'surprised', 'tareme', 'tearing_up', 'thinking', 'tongue', 'tongue_out', 'torogao', 'tsurime', 'turn_pale', 'wide-eyed', 'wince', 'worried', 'heartbeat'],
|
| 23 |
+
'Facial Emoji':['!!', '!', '!?', '+++', '+_+', '...', '...?', '._.', '03:00', '0_0', ':/', ':3', ':<', ':>', ':>=', ':d', ':i', ':o', ':p', ':q', ':t', ':x', ':|', ';(', ';)', ';3', ';d', ';o', ';p', ';q', '=_=', '>:(', '>:)', '>_<', '>_o', '>o<', '?', '??', '@_@', '\\m/', '\n/', '\\o/', '\\||/', '^^^', '^_^', 'c:', 'd:', 'o_o', 'o3o', 'u_u', 'w', 'x', 'x_x', 'xd', 'zzz', '|_|'],
|
| 24 |
+
'Head':['afro', 'ahoge', 'animal_ear_fluff', '_bangs', '_bun', 'bald', 'beard', 'blunt_bangs', 'blunt_ends', 'bob_cut', 'bowl_cut', 'braid', 'braids', 'buzz_cut', 'circle_cut', 'colored_tips', 'cowlick', 'dot_nose', 'dreadlocks', '_ear', '_ears', '_eye', '_eyes', 'enpera', 'eyeball', 'eyebrow', 'eyebrow_cut', 'eyebrows', 'eyelashes', 'eyeshadow', 'faceless', 'facepaint', 'facial_mark', 'fang', 'forehead', 'freckles', 'goatee', '_hair', 'very_long_hair', 'hair_bun', 'hair_flaps', 'hair_intakes', 'hair_tubes', 'tentacle_hair', '_horn', '_horns', 'half_updo', 'head_tilt', 'heterochromia', 'hime_cut', 'hime_cut', 'horns', '_in_eye', 'inverted_bob', 'kemonomimi_mode', 'lips', 'mascara', 'mohawk', 'mouth_', 'mustache', 'nose', 'one-eyed', 'one_eye', 'one_side_up', '_pupils', 'parted_bangs', 'pompadour', 'ponytail', 'ringlets', '_sclera', 'sideburns', 'sidecut', 'sidelock', 'sidelocks', 'skull', 'snout', 'stubble', 'swept_bangs', 'tails', 'teeth', 'third_eye', 'twintails', 'two_side_up', 'undercut', 'updo', 'v-shaped_eyebrows', 'whiskers'],
|
| 25 |
+
'Hands':['_arm', '_arms', 'claws', '_finger', '_fingers', 'fingernails', '_hand', '_nail', '_nails', 'palms', 'rings', 'thumbs_up'],
|
| 26 |
+
'Upper Body':['abs', 'armpit', 'armpits', 'backboob', 'belly', 'biceps', 'breast_rest', 'breasts', 'button_gap', 'cleavage', 'collarbone', 'dimples_of_venus', 'downblouse', 'flat_chest', 'linea_alba', 'median_furrow', 'midriff', 'nape', 'navel', 'pectorals', 'ribs', '_shoulder', '_shoulders', 'shoulder_blades', 'sideboob', 'sidetail', 'spine', 'stomach', 'strap_gap', 'toned', 'underboob', 'underbust'],
|
| 27 |
+
'Lower Body':['ankles', 'ass', 'barefoot', 'crotch', 'feet', 'highleg', 'hip_bones', 'hooves', 'kneepits', 'knees', 'legs', 'soles', 'tail', 'thigh_gap', 'thighlet', 'thighs', 'toenail', 'toenails', 'toes', 'wide_hips'],
|
| 28 |
+
'Creature':['(animal)', 'anglerfish', 'animal', 'bear', 'bee', 'bird', 'bug', 'butterfly', 'cat', 'chick', 'chicken', 'chinese_zodiac', 'clownfish', 'coral', 'crab', 'creature', 'crow', 'dog', 'dove', 'dragon', 'duck', 'eagle', 'fish', 'fish', 'fox', 'fox', 'frog', 'frog', 'goldfish', 'hamster', 'horse', 'jellyfish', 'ladybug', 'lion', 'mouse', 'octopus', 'owl', 'panda', 'penguin', 'pig', 'pigeon', 'rabbit', 'rooster', 'seagull', 'shark', 'sheep', 'shrimp', 'snail', 'snake', 'squid', 'starfish', 'tanuki', 'tentacles', 'goo_tentacles', 'plant_tentacles', 'crotch_tentacles', 'mechanical_tentacles', 'squidward_tentacles', 'suction_tentacles', 'penis_tentacles', 'translucent_tentacles', 'back_tentacles', 'red_tentacles', 'green_tentacles', 'blue_tentacles', 'black_tentacles', 'pink_tentacles', 'purple_tentacles', 'face_tentacles', 'tentacles_everywhere', 'milking_tentacles', 'tiger', 'turtle', 'weasel', 'whale', 'wolf', 'parrot', 'sparrow', 'unicorn'],
|
| 29 |
+
'Plant':['bamboo', 'bouquet', 'branch', 'bush', 'cherry_blossoms', 'clover', 'daisy', '(flower)', 'flower', 'flower', 'gourd', 'hibiscus', 'holly', 'hydrangea', 'leaf', 'lily_pad', 'lotus', 'moss', 'palm_leaf', 'palm_tree', 'petals', 'plant', 'plum_blossoms', 'rose', 'spider_lily', 'sunflower', 'thorns', 'tree', 'tulip', 'vines', 'wisteria', 'acorn'],
|
| 30 |
+
'Food':['apple', 'baguette', 'banana', 'baozi', 'beans', 'bento', 'berry', 'blueberry', 'bread', 'broccoli', 'burger', 'cabbage', 'cake', 'candy', 'carrot', 'cheese', 'cherry', 'chili_pepper', 'chocolate', 'coconut', 'cookie', 'corn', 'cream', 'crepe', 'cucumber', 'cucumber', 'cupcake', 'curry', 'dango', 'dessert', 'doughnut', 'egg', 'eggplant', '_(food)', '_(fruit)', 'food', 'french_fries', 'fruit', 'grapes', 'ice_cream', 'icing', 'lemon', 'lettuce', 'lollipop', 'macaron', 'mandarin_orange', 'meat', 'melon', 'mochi', 'mushroom', 'noodles', 'omelet', 'omurice', 'onigiri', 'onion', 'pancake', 'parfait', 'pasties', 'pastry', 'peach', 'pineapple', 'pizza', 'popsicle', 'potato', 'pudding', 'pumpkin', 'radish', 'ramen', 'raspberry', 'rice', 'roasted_sweet_potato', 'sandwich', 'sausage', 'seaweed', 'skewer', 'spitroast', 'spring_onion', 'strawberry', 'sushi', 'sweet_potato', 'sweets', 'taiyaki', 'takoyaki', 'tamagoyaki', 'tempurakanbea', 'toast', 'tomato', 'vegetable', 'wagashi', 'wagashi', 'watermelon', 'jam', 'popcorn'],
|
| 31 |
+
'Beverage':['alcohol', 'beer', 'coffee', 'cola', 'drink', 'juice', 'juice_box', 'milk', 'sake', 'soda', 'tea', '_tea', 'whiskey', 'wine', 'cocktail'],
|
| 32 |
+
'Music':['band', 'baton_(conducting)', 'beamed', 'cello', 'concert', 'drum', 'drumsticks', 'eighth_note', 'flute', 'guitar', 'harp', 'horn', '(instrument)', 'idol', 'instrument', 'k-pop', 'lyre', '(music)', 'megaphone', 'microphone', 'music', 'musical_note', 'phonograph', 'piano', 'plectrum', 'quarter_note', 'recorder', 'sixteenth_note', 'sound_effects', 'trumpet', 'utaite', 'violin', 'whistle'],
|
| 33 |
+
'Weapons & Equipment':['ammunition', 'arrow_(projectile)', 'axe', 'bandolier', 'baseball_bat', 'beretta_92', 'bolt_action', 'bomb', 'bullet', 'bullpup', 'cannon', 'chainsaw', 'crossbow', 'dagger', 'energy_sword', 'explosive', 'fighter_jet', 'gohei', 'grenade', 'gun', 'hammer', 'handgun', 'holstered', 'jet', 'katana', 'knife', 'kunai', 'lance', 'mallet', 'nata_(tool)', 'polearm', 'quiver', 'rapier', 'revolver', 'rifle', 'rocket_launcher', 'scabbard', 'scope', 'scythe', 'sheath', 'sheathed', 'shield', 'shotgun', 'shuriken', 'spear', 'staff', 'suppressor', 'sword', 'tank', 'tantou', 'torpedo', 'trident', '(weapon)', 'wand', 'weapon', 'whip', 'yumi_(bow)', 'h&k_hk416', 'rocket_launcher', 'heckler_&_koch', '_weapon'],
|
| 34 |
+
'Vehicles':['aircraft', 'airplane', 'bicycle', 'boat', 'car', 'caterpillar_tracks', 'flight_deck', 'helicopter', 'motor_vehicle', 'motorcycle', 'ship', 'spacecraft', 'spoiler_(automobile)', 'train', 'truck', 'watercraft', 'wheel', 'wheelbarrow', 'wheelchair', 'inflatable_raft'],
|
| 35 |
+
'Buildings':['apartment', 'aquarium', 'architecture', 'balcony', 'building', 'cafe', 'castle', 'church', 'gym', 'hallway', 'hospital', 'house', 'library', '(place)', 'porch', 'restaurant', 'restroom', 'rooftop', 'shop', 'skyscraper', 'stadium', 'stage', 'temple', 'toilet', 'tower', 'train_station', 'veranda'],
|
| 36 |
+
'Indoor':['bath', 'bathroom', 'bathtub', 'bed', 'bed_sheet', 'bedroom', 'blanket', 'bookshelf', 'carpet', 'ceiling', 'chair', 'chalkboard', 'classroom', 'counter', 'cupboard', 'curtains', 'cushion', 'dakimakura', 'desk', 'door', 'doorway', 'drawer', '_floor', 'floor', 'futon', 'indoors', 'interior', 'kitchen', 'kotatsu', 'locker', 'mirror', 'pillow', 'room', 'rug', 'school_desk', 'shelf', 'shouji', 'sink', 'sliding_doors', 'stairs', 'stool', 'storeroom', 'table', 'tatami', 'throne', 'window', 'windowsill', 'bathhouse', 'chest_of_drawers'],
|
| 37 |
+
'Outdoor':['alley', 'arch', 'beach', 'bridge', 'bus_stop', 'bush', 'cave', '(city)', 'city', 'cliff', 'crescent', 'crosswalk', 'day', 'desert', 'fence', 'ferris_wheel', 'field', 'forest', 'grass', 'graveyard', 'hill', 'lake', 'lamppost', 'moon', 'mountain', 'night', 'ocean', 'onsen', 'outdoors', 'path', 'pool', 'poolside', 'railing', 'railroad', 'river', 'road', 'rock', 'sand', 'shore', 'sky', 'smokestack', 'snow', 'snowball', 'snowman', 'street', 'sun', 'sunlight', 'sunset', 'tent', 'torii', 'town', 'tree', 'turret', 'utility_pole', 'valley', 'village', 'waterfall'],
|
| 38 |
+
'Objects':['anchor', 'android', 'armchair', '(bottle)', 'backpack', 'bag', 'ball', 'balloon', 'bandages', 'bandaid', 'bandaids', 'banknote', 'banner', 'barcode', 'barrel', 'baseball', 'basket', 'basketball', 'beachball', 'bell', 'bench', 'binoculars', 'board_game', 'bone', 'book', 'bottle', 'bowl', 'box', 'box_art', 'briefcase', 'broom', 'bucket', '(chess)', '(computer)', '(computing)', '(container)', 'cage', 'calligraphy_brush', 'camera', 'can', 'candle', 'candlestand', 'cane', 'card', 'cartridge', 'cellphone', 'chain', 'chandelier', 'chess', 'chess_piece', 'choko_(cup)', 'chopsticks', 'cigar', 'clipboard', 'clock', 'clothesline', 'coin', 'comb', 'computer', 'condom', 'controller', 'cosmetics', 'couch', 'cowbell', 'crazy_straw', 'cup', 'cutting_board', 'dice', 'digital_media_player', 'doll', 'drawing_tablet', 'drinking_straw', 'easel', 'electric_fan', 'emblem', 'envelope', 'eraser', 'feathers', 'figure', 'fire', 'fishing_rod', 'flag', 'flask', 'folding_fan', 'fork', 'frying_pan', '(gemstone)', 'game_console', 'gears', 'gemstone', 'gift', 'glass', 'glowstick', 'gold', 'handbag', 'handcuffs', 'handheld_game_console', 'hose', 'id_card', 'innertube', 'iphone',"jack-o'-lantern",'jar', 'joystick', 'key', 'keychain', 'kiseru', 'ladder', 'ladle', 'lamp', 'lantern', 'laptop', 'letter', 'letterboxed', 'lifebuoy', 'lipstick', 'liquid', 'lock', 'lotion', '_machine', 'map', 'marker', 'model_kit', 'money', 'monitor', 'mop', 'mug', 'needle', 'newspaper', 'nintendo', 'nintendo_switch', 'notebook', '(object)', 'ofuda', 'orb', 'origami', '(playing_card)', 'pack', 'paddle', 'paintbrush', 'pan', 'paper', 'parasol', 'patch', 'pc', 'pen', 'pencil', 'pencil', 'pendant_watch', 'phone', 'pill', 'pinwheel', 'plate', 'playstation', 'pocket_watch', 'pointer', 'poke_ball', 'pole', 'quill', 'racket', 'randoseru', 'remote_control', 'ring', 'rope', 'sack', 'saddle', 'sakazuki', 'satchel', 'saucer', 'scissors', 'scroll', 'seashell', 'seatbelt', 'shell', 'shide', 'shopping_cart', 'shovel', 'shower_head', 'silk', 'sketchbook', 'smartphone', 'soap', 'sparkler', 'spatula', 'speaker', 'spoon', 'statue', 'stethoscope', 'stick', 'sticker', 'stopwatch', 'string', 'stuffed_', 'stylus', 'suction_cups', 'suitcase', 'surfboard', 'syringe', 'talisman', 'tanzaku', 'tape', 'teacup', 'teapot', 'teddy_bear', 'television', 'test_tube', 'tiles', 'tokkuri', 'tombstone', 'torch', 'towel', 'toy', 'traffic_cone', 'tray', 'treasure_chest', 'uchiwa', 'umbrella', 'vase', 'vial', 'video_game', 'viewfinder', 'volleyball', 'wallet', 'watch', 'watch', 'whisk', 'whiteboard', 'wreath', 'wrench', 'wristwatch', 'yunomi', 'ace_of_hearts', 'inkwell', 'compass', 'ipod', 'sunscreen', 'rocket', 'cobblestone'],
|
| 39 |
+
'Character Design':['+boys', '+girls', '1other', '39', '_boys', '_challenge', '_connection', '_female', '_fur', '_girls', '_interface', '_male', '_man', '_person', 'abyssal_ship', 'age_difference', 'aged_down', 'aged_up', 'albino', 'alien', 'alternate_muscle_size', 'ambiguous_gender', 'amputee', 'androgynous', 'angel', 'animalization', 'ass-to-ass', 'assault_visor', 'au_ra', 'baby', 'bartender', 'beak', 'bishounen', 'borrowed_character', 'boxers', 'boy', 'breast_envy', 'breathing_fire', 'bride', 'broken', 'brother_and_sister', 'brothers', 'camouflage', 'cheating_(relationship)', 'cheerleader', 'chibi', 'child', 'clone', 'command_spell', 'comparison', 'contemporary', 'corpse', 'corruption', 'cosplay', 'couple', 'creature_and_personification', 'crossdressing', 'crossover', 'cyberpunk', 'cyborg', 'cyclops', 'damaged', 'dancer', 'danmaku', 'darkness', 'death', 'defeat', 'demon', 'disembodied_', 'draph', 'drone', 'duel', 'dwarf', 'egyptian', 'electricity', 'elezen', 'elf', 'enmaided', 'erune', 'everyone', 'evolutionary_line', 'expressions', 'fairy', 'family', 'fangs', 'fantasy', 'fashion', 'fat', 'father_and_daughter', 'father_and_son', 'fewer_digits', 'fins', 'flashback', 'fluffy', 'fumo_(doll)', 'furry', 'fusion', 'fuuin_no_tsue', 'gameplay_mechanics', 'genderswap', 'ghost', 'giant', 'giantess', 'gibson_les_paul', 'girl', 'goblin', 'groom', 'guro', 'gyaru', 'habit', 'harem', 'harpy', 'harvin', 'heads_together', 'health_bar', 'height_difference', 'hitodama', 'horror_(theme)', 'humanization', 'husband_and_wife', 'hydrokinesis', 'hypnosis', 'hyur', 'idol', 'insignia', 'instant_loss', 'interracial', 'interspecies', 'japari_bun', 'jeweled_branch_of_hourai', 'jiangshi', 'jirai_kei', 'joints', 'karakasa_obake', 'keyhole', 'kitsune', 'knight', 'kodona', 'kogal', 'kyuubi', 'lamia', 'left-handed', 'loli', 'lolita', 'look-alike', 'machinery', 'magic', 'male_focus', 'manly', 'matching_outfits', 'mature_female', 'mecha', 'mermaid', 'meta', 'miko', 'milestone_celebration', 'military', 'mind_control', 'miniboy', 'minigirl',"miqo'te",'monster', 'monsterification', 'mother_and_daughter', 'mother_and_son', 'multiple_others', 'muscular', 'nanodesu_(phrase)', 'narrow_waist', 'nekomata', 'netorare', 'ninja', 'no_humans', 'nontraditional', 'nun', 'nurse', 'object_namesake', 'obliques', 'office_lady', 'old', 'on_body', 'onee-shota', 'oni', 'orc', 'others', 'otoko_no_ko', 'oversized_object', 'paint_splatter', 'pantyshot', 'pawpads', 'persona', 'personality', 'personification', 'pet_play', 'petite', 'pirate', 'playboy_bunny', 'player_2', 'plugsuit', 'plump', 'poi', 'pokemon', 'police', 'policewoman', 'pom_pom_(cheerleading)', 'princess', 'prosthesis', 'pun', 'puppet', 'race_queen', 'radio_antenna', 'real_life_insert', 'redesign', 'reverse_trap', 'rigging', 'robot', 'rod_of_remorse', 'sailor', 'salaryman', 'samurai', 'sangvis_ferri', 'scales', 'scene_reference', 'school', 'sheikah', 'shota', 'shrine', 'siblings', 'side-by-side', 'sidesaddle', 'sisters', 'size_difference', 'skeleton', 'skinny', 'slave', 'slime_(substance)', 'soldier', 'spiked_shell', 'spokencharacter', 'steampunk', 'streetwear', 'striker_unit', 'strongman', 'submerged', 'suggestive', 'super_saiyan', 'superhero', 'surreal', 'take_your_pick', 'tall', 'talons', 'taur', 'teacher', 'team_rocket', 'three-dimensional_maneuver_gear', 'time_paradox', 'tomboy', 'traditional_youkai', 'transformation', 'trick_or_treat', 'tusks', 'twins', 'ufo', 'under_covers', 'v-fin', 'v-fin', 'vampire', 'virtual_youtuber', 'waitress', 'watching_television', 'wedding', 'what', 'when_you_see_it', 'wife_and_wife', 'wing', 'wings', 'witch', 'world_war_ii', 'yandere', 'year_of', 'yes', 'yin_yang', 'yordle',"you're_doing_it_wrong",'you_gonna_get_raped', 'yukkuri_shiteitte_ne', 'yuri', 'zombie', '(alice_in_wonderland)', '(arknights)', '(blue_archive)', '(cosplay)', '(creature)', '(emblem)', '(evangelion)', '(fate)', '(fate/stay_night)', '(ff11)', '(fire_emblem)', '(genshin_impact)', '(grimm)', '(houseki_no_kuni)', '(hyouka)', '(idolmaster)', '(jojo)', '(kancolle)', '(kantai_collection)', '(kill_la_kill)', '(league_of_legends)', '(legends)', '(lyomsnpmp)', '(machimazo)', '(madoka_magica)', '(mecha)', '(meme)', '(nier:automata)', '(organ)', '(overwatch)', '(pokemon)', '(project_moon)', '(project_sekai)', '(sao)', '(senran_kagura)', '(splatoon)', '(touhou)', '(tsukumo_sana)', '(youkai_watch)', '(yu-gi-oh!_gx)', '(zelda)', 'sextuplets', 'imperial_japanese_army', 'extra_faces', '_miku'],
|
| 40 |
+
'Composition':['abstract', 'anime_coloring', 'animification', 'back-to-back', 'bad_anatomy', 'blurry', 'border', 'bound', 'cameo', 'cheek-to-cheek', 'chromatic_aberration', 'close-up', 'collage', 'color_guide', 'colorful', 'comic', 'contrapposto', 'cover', 'cowboy_shot', 'crosshatching', 'depth_of_field', 'dominatrix', 'dutch_angle', '_focus', 'face-to-face', 'fake_screenshot', 'film_grain', 'fisheye', 'flat_color', 'foreshortening', 'from_above', 'from_behind', 'from_below', 'from_side', 'full_body', 'glitch', 'greyscale', 'halftone', 'head_only', 'heads-up_display', 'high_contrast', 'horizon', '_inset', 'inset', 'jaggy_lines', '1koma', '2koma', '3koma', '4koma', '5koma', 'leaning', 'leaning_forward', 'leaning_to_the_side', 'left-to-right_manga', 'lens_flare', 'limited_palette', 'lineart', 'lineup', 'lower_body', '(medium)', 'marker_(medium)', 'meme', 'mixed_media', 'monochrome', 'multiple_views', 'muted_color', 'oekaki', 'on_side', 'out_of_frame', 'outline', 'painting', 'parody', 'partially_colored', 'partially_underwater_shot', 'perspective', 'photorealistic', 'picture_frame', 'pillarboxed', 'portrait', 'poster_(object)', 'product_placement', 'profile', 'realistic', 'recording', 'retro_artstyle', '(style)', '_style', 'sandwiched', 'science_fiction', 'sepia', 'shikishi', 'side-by-side', 'sideways', 'sideways_glance', 'silhouette', 'sketch', 'spot_color', 'still_life', 'straight-on', 'symmetry', '(texture)', 'tachi-e', 'taking_picture', 'tegaki', 'too_many', 'traditional_media', 'turnaround', 'underwater', 'upper_body', 'upside-down', 'upskirt', 'variations', 'wide_shot', '_design', 'symbolism', 'rounded_corners', 'surrounded'],
|
| 41 |
+
'Season':['akeome', 'anniversary', 'autumn', 'birthday', 'christmas', '_day', 'festival', 'halloween', 'kotoyoro', 'nengajou', 'new_year', 'spring_(season)', 'summer', 'tanabata', 'valentine', 'winter'],
|
| 42 |
+
'Background':['simple_background', '_background', 'backlighting', 'bloom', 'bokeh', 'brick_wall', 'bubble', 'cable', 'caustics', 'cityscape', 'cloud', 'confetti', 'constellation', 'contrail', 'crowd', 'crystal', 'dark', 'debris', 'dusk', 'dust', 'egasumi', 'embers', 'emphasis_lines', 'energy', 'evening', 'explosion', 'fireworks', 'fog', 'footprints', 'glint', 'graffiti', 'ice', 'industrial_pipe', 'landscape', 'light', 'light_particles', 'light_rays', 'lightning', 'lights', 'moonlight', 'motion_blur', 'motion_lines', 'mountainous_horizon', 'nature', '(planet)', 'pagoda', 'people', 'pillar', 'planet', 'power_lines', 'puddle', 'rain', 'rainbow', 'reflection', 'ripples', 'rubble', 'ruins', 'scenery', 'shade', 'shooting_star', 'sidelighting', 'smoke', 'snowflakes', 'snowing', 'space', 'sparkle', 'sparks', 'speed_lines', 'spider_web', 'spotlight', 'star_(sky)', 'stone_wall', 'sunbeam', 'sunburst', 'sunrise', '_theme', 'tile_wall', 'twilight', 'wall_clock', 'wall_of_text', 'water', 'waves', 'wind', 'wire', 'wooden_wall', 'lighthouse'],
|
| 43 |
+
'Patterns':['arrow', 'bass_clef', 'blank_censor', 'circle', 'cube', 'heart', 'hexagon', 'hexagram', 'light_censor', '(pattern)', 'pattern', 'pentagram', 'roman_numeral', '(shape)', '(symbol)', 'shape', 'sign', 'symbol', 'tally', 'treble_clef', 'triangle', 'tube', 'yagasuri'],
|
| 44 |
+
'Censorship':['blur_censor', '_censor', '_censoring', 'censored', 'character_censor', 'convenient', 'hair_censor', 'heart_censor', 'identity_censor', 'maebari', 'novelty_censor', 'soap_censor', 'steam_censor', 'tail_censor', 'uncensored'],
|
| 45 |
+
'Others':['2007', '2008', '2009', '2010', '2011', '2012', '2013', '2014', '2015', '2016', '2017', '2018', '2019', '2020', '2021', '2022', '2023', '2024', 'artist', 'artist_name', 'artistic_error', 'asian', '(company)', 'character_name', 'content_rating', 'copyright', 'cover_page', 'dated', 'english_text', 'japan', 'layer', 'logo', 'name', 'numbered', 'page_number', 'pixiv_id', 'language', 'reference_sheet', 'signature', 'speech_bubble', 'subtitled', 'text', 'thank_you', 'typo', 'username', 'wallpaper', 'watermark', 'web_address', 'screwdriver', 'translated'],
|
| 46 |
+
'Quality Tags':['masterpiece', '_quality', 'highres', 'absurdres', 'ultra-detailed', 'lowres']}
|
| 47 |
+
|
| 48 |
+
# Build a trie for efficient prefix matching
|
| 49 |
+
class TrieNode:
|
| 50 |
+
def __init__(self):
|
| 51 |
+
self.children = {}
|
| 52 |
+
self.category = None
|
| 53 |
+
|
| 54 |
+
class TagTrie:
|
| 55 |
+
def __init__(self):
|
| 56 |
+
self.root = TrieNode()
|
| 57 |
+
self._build_trie()
|
| 58 |
+
|
| 59 |
+
def _build_trie(self):
|
| 60 |
+
for category, tags in categories.items():
|
| 61 |
+
for tag in tags:
|
| 62 |
+
node = self.root
|
| 63 |
+
for char in tag:
|
| 64 |
+
if char not in node.children:
|
| 65 |
+
node.children[char] = TrieNode()
|
| 66 |
+
node = node.children[char]
|
| 67 |
+
node.category = category
|
| 68 |
+
|
| 69 |
+
def find_category(self, tag):
|
| 70 |
+
node = self.root
|
| 71 |
+
matched_category = None
|
| 72 |
+
|
| 73 |
+
# Try exact match first
|
| 74 |
+
for char in tag:
|
| 75 |
+
if char in node.children:
|
| 76 |
+
node = node.children[char]
|
| 77 |
+
if node.category:
|
| 78 |
+
matched_category = node.category
|
| 79 |
+
else:
|
| 80 |
+
break
|
| 81 |
+
|
| 82 |
+
# If exact match found, return it
|
| 83 |
+
if matched_category and node.children == {}:
|
| 84 |
+
return matched_category
|
| 85 |
+
|
| 86 |
+
# If partial match found, check if it's a valid prefix
|
| 87 |
+
if matched_category:
|
| 88 |
+
return matched_category
|
| 89 |
+
|
| 90 |
+
# Try substring matching for longer than 3 characters
|
| 91 |
+
for i in range(len(tag)):
|
| 92 |
+
for j in range(i+4, len(tag)+1): # Only check substrings longer than 3 chars
|
| 93 |
+
substring = tag[i:j]
|
| 94 |
+
node = self.root
|
| 95 |
+
valid = True
|
| 96 |
+
for char in substring:
|
| 97 |
+
if char in node.children:
|
| 98 |
+
node = node.children[char]
|
| 99 |
+
else:
|
| 100 |
+
valid = False
|
| 101 |
+
break
|
| 102 |
+
if valid and node.category:
|
| 103 |
+
return node.category
|
| 104 |
+
|
| 105 |
+
return None
|
| 106 |
+
|
| 107 |
+
tag_trie = TagTrie()
|
| 108 |
+
|
| 109 |
+
def normalize_tag(tag):
|
| 110 |
+
"""Normalize tag by converting spaces/hyphens to underscores"""
|
| 111 |
+
return re.sub(r'[-\s]+', '_', tag.strip())
|
| 112 |
+
|
| 113 |
+
def classify_single_tag(tag):
|
| 114 |
+
"""Classify a single tag into its category"""
|
| 115 |
+
normalized_tag = normalize_tag(tag)
|
| 116 |
+
|
| 117 |
+
# Try exact match through Trie lookup first
|
| 118 |
+
category = tag_trie.find_category(normalized_tag)
|
| 119 |
+
|
| 120 |
+
# If no match and has underscores, try parts
|
| 121 |
+
if not category and '_' in normalized_tag:
|
| 122 |
+
parts = normalized_tag.split('_')
|
| 123 |
+
for part in parts:
|
| 124 |
+
if len(part) > 3: # Only check parts longer than 3 characters
|
| 125 |
+
category = tag_trie.find_category(part)
|
| 126 |
+
if category:
|
| 127 |
+
break
|
| 128 |
+
|
| 129 |
+
# Special handling for escaped parentheses
|
| 130 |
+
if not category and ('\\(' in normalized_tag or '\\)' in normalized_tag):
|
| 131 |
+
unescaped = normalized_tag.replace('\\(', '(').replace('\\)', ')')
|
| 132 |
+
category = tag_trie.find_category(unescaped)
|
| 133 |
+
|
| 134 |
+
if not category and '_' in unescaped:
|
| 135 |
+
parts = unescaped.split('_')
|
| 136 |
+
for part in parts:
|
| 137 |
+
if len(part) > 3:
|
| 138 |
+
category = tag_trie.find_category(part)
|
| 139 |
+
if category:
|
| 140 |
+
break
|
| 141 |
+
|
| 142 |
+
return category if category else 'Uncategorized'
|
| 143 |
+
|
| 144 |
+
def extract_priority_and_character_tags(tags_list, character_tags):
|
| 145 |
+
"""
|
| 146 |
+
Extract priority tags and character tags from the tags list
|
| 147 |
+
|
| 148 |
+
Args:
|
| 149 |
+
tags_list (list): List of all tags
|
| 150 |
+
character_tags (dict): Dictionary of character tags with confidence scores
|
| 151 |
+
|
| 152 |
+
Returns:
|
| 153 |
+
tuple: (priority_tags, character_tag_names, remaining_tags)
|
| 154 |
+
"""
|
| 155 |
+
priority_tags_found = []
|
| 156 |
+
character_tag_names = list(character_tags.keys()) if character_tags else []
|
| 157 |
+
remaining_tags = []
|
| 158 |
+
|
| 159 |
+
# Convert priority tags to set for faster lookup
|
| 160 |
+
priority_set = set(PRIORITY_TAGS)
|
| 161 |
+
|
| 162 |
+
for tag in tags_list:
|
| 163 |
+
if tag in priority_set:
|
| 164 |
+
priority_tags_found.append(tag)
|
| 165 |
+
elif tag in character_tag_names:
|
| 166 |
+
# Character tags are already handled separately
|
| 167 |
+
remaining_tags.append(tag)
|
| 168 |
+
else:
|
| 169 |
+
remaining_tags.append(tag)
|
| 170 |
+
|
| 171 |
+
return priority_tags_found, character_tag_names, remaining_tags
|
| 172 |
+
|
| 173 |
+
def classify_tags_for_display(tag_string, character_tags=None):
|
| 174 |
+
"""
|
| 175 |
+
Classify a string of tags and organize them by categories with priority ordering for display
|
| 176 |
+
|
| 177 |
+
Args:
|
| 178 |
+
tag_string (str): Comma-separated tags string
|
| 179 |
+
character_tags (dict): Dictionary of character tags with confidence scores
|
| 180 |
+
|
| 181 |
+
Returns:
|
| 182 |
+
str: Categorized and organized tags as a comma-separated string
|
| 183 |
+
"""
|
| 184 |
+
if not tag_string:
|
| 185 |
+
return ""
|
| 186 |
+
|
| 187 |
+
# Split tags by common delimiters
|
| 188 |
+
delimiters = r'[,\n\r\.!?]+'
|
| 189 |
+
raw_tags = re.split(delimiters, tag_string)
|
| 190 |
+
|
| 191 |
+
# Clean and normalize tags
|
| 192 |
+
cleaned_tags = []
|
| 193 |
+
for tag in raw_tags:
|
| 194 |
+
tag = tag.strip()
|
| 195 |
+
if tag:
|
| 196 |
+
cleaned_tags.append(tag)
|
| 197 |
+
|
| 198 |
+
# Extract priority and character tags
|
| 199 |
+
priority_tags_found, character_tag_names, remaining_tags = extract_priority_and_character_tags(cleaned_tags, character_tags)
|
| 200 |
+
|
| 201 |
+
# Classify remaining tags
|
| 202 |
+
categorized = defaultdict(list)
|
| 203 |
+
uncategorized = []
|
| 204 |
+
|
| 205 |
+
for tag in remaining_tags:
|
| 206 |
+
# Skip character tags as they're already in their own list
|
| 207 |
+
if tag in character_tag_names:
|
| 208 |
+
continue
|
| 209 |
+
|
| 210 |
+
category = classify_single_tag(tag)
|
| 211 |
+
if category == 'Uncategorized':
|
| 212 |
+
uncategorized.append(tag)
|
| 213 |
+
else:
|
| 214 |
+
categorized[category].append(tag)
|
| 215 |
+
|
| 216 |
+
# Build result string with priority ordering
|
| 217 |
+
result_parts = []
|
| 218 |
+
|
| 219 |
+
# 1. Add priority subject tags first
|
| 220 |
+
result_parts.extend(priority_tags_found)
|
| 221 |
+
|
| 222 |
+
# 2. Add character tags next
|
| 223 |
+
result_parts.extend(character_tag_names)
|
| 224 |
+
|
| 225 |
+
# 3. Add categorized tags in category order
|
| 226 |
+
for category in categories.keys():
|
| 227 |
+
if category in categorized and categorized[category]:
|
| 228 |
+
result_parts.extend(categorized[category])
|
| 229 |
+
|
| 230 |
+
# 4. Add uncategorized tags at the end
|
| 231 |
+
result_parts.extend(uncategorized)
|
| 232 |
+
|
| 233 |
+
# Process tags: replace underscores with spaces and handle escaped characters
|
| 234 |
+
processed_tags = []
|
| 235 |
+
for tag in result_parts:
|
| 236 |
+
processed_tag = tag.replace('_', ' ').replace('\\(', '(').replace('\\)', ')')
|
| 237 |
+
processed_tags.append(processed_tag)
|
| 238 |
+
|
| 239 |
+
return ', '.join(processed_tags)
|
| 240 |
+
|
| 241 |
+
def generate_categorized_json(tag_string, character_tags=None):
|
| 242 |
+
"""
|
| 243 |
+
Generate JSON object organizing tags by categories
|
| 244 |
+
|
| 245 |
+
Args:
|
| 246 |
+
tag_string (str): Comma-separated tags string
|
| 247 |
+
character_tags (dict): Dictionary of character tags with confidence scores
|
| 248 |
+
|
| 249 |
+
Returns:
|
| 250 |
+
dict: JSON-compatible dictionary with categories as keys and tag lists as values
|
| 251 |
+
"""
|
| 252 |
+
if not tag_string:
|
| 253 |
+
return {}
|
| 254 |
+
|
| 255 |
+
# Split tags by common delimiters
|
| 256 |
+
delimiters = r'[,\n\r\.!?]+'
|
| 257 |
+
raw_tags = re.split(delimiters, tag_string)
|
| 258 |
+
|
| 259 |
+
# Clean and normalize tags
|
| 260 |
+
cleaned_tags = []
|
| 261 |
+
for tag in raw_tags:
|
| 262 |
+
tag = tag.strip()
|
| 263 |
+
if tag:
|
| 264 |
+
cleaned_tags.append(tag)
|
| 265 |
+
|
| 266 |
+
# Extract priority and character tags
|
| 267 |
+
priority_tags_found, character_tag_names, remaining_tags = extract_priority_and_character_tags(cleaned_tags, character_tags)
|
| 268 |
+
|
| 269 |
+
# Classify remaining tags
|
| 270 |
+
categorized = defaultdict(list)
|
| 271 |
+
uncategorized = []
|
| 272 |
+
|
| 273 |
+
for tag in remaining_tags:
|
| 274 |
+
# Skip character tags as they're already in their own list
|
| 275 |
+
if tag in character_tag_names:
|
| 276 |
+
continue
|
| 277 |
+
|
| 278 |
+
category = classify_single_tag(tag)
|
| 279 |
+
if category == 'Uncategorized':
|
| 280 |
+
uncategorized.append(tag)
|
| 281 |
+
else:
|
| 282 |
+
# Store the original tag (with underscores) for JSON
|
| 283 |
+
categorized[category].append(tag)
|
| 284 |
+
|
| 285 |
+
# Build JSON result
|
| 286 |
+
json_result = {}
|
| 287 |
+
|
| 288 |
+
# Add special categories if they have content
|
| 289 |
+
if priority_tags_found:
|
| 290 |
+
# Process priority tags for display (replace underscores with spaces) # Replacement is not 100% necessary, but will do anyway
|
| 291 |
+
processed_priority = [tag.replace('_', ' ').replace('\\(', '(').replace('\\)', ')') for tag in priority_tags_found]
|
| 292 |
+
json_result['Subject'] = processed_priority
|
| 293 |
+
|
| 294 |
+
if character_tag_names:
|
| 295 |
+
# Process character tags for display
|
| 296 |
+
processed_characters = [tag.replace('_', ' ').replace('\\(', '(').replace('\\)', ')') for tag in character_tag_names]
|
| 297 |
+
json_result['Characters'] = processed_characters
|
| 298 |
+
|
| 299 |
+
# Add categorized tags (process for display)
|
| 300 |
+
for category, tags in categorized.items():
|
| 301 |
+
if tags:
|
| 302 |
+
processed_tags = [tag.replace('_', ' ').replace('\\(', '(').replace('\\)', ')') for tag in tags]
|
| 303 |
+
json_result[category] = processed_tags
|
| 304 |
+
|
| 305 |
+
# Add uncategorized tags if any
|
| 306 |
+
if uncategorized:
|
| 307 |
+
processed_uncategorized = [tag.replace('_', ' ').replace('\\(', '(').replace('\\)', ')') for tag in uncategorized]
|
| 308 |
+
json_result['Uncategorized'] = processed_uncategorized
|
| 309 |
+
|
| 310 |
+
return json_result
|
| 311 |
+
|
| 312 |
+
|
| 313 |
+
def categorize_tags_output(tag_string, character_tags=None):
|
| 314 |
+
"""
|
| 315 |
+
Main function to categorize tags output for display
|
| 316 |
+
|
| 317 |
+
Args:
|
| 318 |
+
tag_string (str): Raw tags string from the model
|
| 319 |
+
character_tags (dict): Dictionary of character tags with confidence scores
|
| 320 |
+
|
| 321 |
+
Returns:
|
| 322 |
+
str: Organized, categorized tags string
|
| 323 |
+
"""
|
| 324 |
+
return classify_tags_for_display(tag_string, character_tags)
|
| 325 |
+
|
| 326 |
+
def generate_tags_json(tag_string, character_tags=None):
|
| 327 |
+
"""
|
| 328 |
+
Main function to generate categorized JSON
|
| 329 |
+
|
| 330 |
+
Args:
|
| 331 |
+
tag_string (str): Raw tags string from the model
|
| 332 |
+
character_tags (dict): Dictionary of character tags with confidence scores
|
| 333 |
+
|
| 334 |
+
Returns:
|
| 335 |
+
dict: JSON object with categorized tags
|
| 336 |
+
"""
|
| 337 |
+
return generate_categorized_json(tag_string, character_tags)
|
modules/pixai.py
ADDED
|
@@ -0,0 +1,810 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os, json, zipfile, tempfile, time, traceback
|
| 2 |
+
import gradio as gr
|
| 3 |
+
import pandas as pd
|
| 4 |
+
import numpy as np
|
| 5 |
+
import onnxruntime as ort
|
| 6 |
+
from collections import defaultdict
|
| 7 |
+
from typing import Union, Dict, Any, Tuple, List
|
| 8 |
+
from PIL import Image
|
| 9 |
+
from huggingface_hub import hf_hub_download
|
| 10 |
+
from huggingface_hub.errors import EntryNotFoundError
|
| 11 |
+
from datetime import datetime
|
| 12 |
+
|
| 13 |
+
# Global variables for model components (for memory management)
|
| 14 |
+
CURRENT_MODEL = None
|
| 15 |
+
CURRENT_MODEL_NAME = None
|
| 16 |
+
CURRENT_TAGS_DF = None
|
| 17 |
+
CURRENT_D_IPS = None
|
| 18 |
+
CURRENT_PREPROCESS_FUNC = None
|
| 19 |
+
CURRENT_THRESHOLDS = None
|
| 20 |
+
CURRENT_CATEGORY_NAMES = None
|
| 21 |
+
|
| 22 |
+
css = """
|
| 23 |
+
#custom-gallery {--row-height: 180px;display: grid;grid-auto-rows: min-content;gap: 10px;}
|
| 24 |
+
#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;}
|
| 25 |
+
#custom-gallery .thumbnail-item:hover {transform: translateY(-3px);box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);}
|
| 26 |
+
#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;}
|
| 27 |
+
#custom-gallery .thumbnail-item img.portrait {max-width: 100%;}
|
| 28 |
+
#custom-gallery .thumbnail-item img.landscape {max-height: 100%;}
|
| 29 |
+
.gallery-container {max-height: 500px;overflow-y: auto;padding-right: 0px;--size-80: 500px;}
|
| 30 |
+
.thumbnails {display: flex;position: absolute;bottom: 0;width: 120px;overflow-x: scroll;padding-top: 320px;padding-bottom: 280px;padding-left: 4px;flex-wrap: wrap;}
|
| 31 |
+
#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;}
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
def preprocess_on_gpu(img, device='cuda'):
|
| 35 |
+
"""Preprocess image on GPU using PyTorch"""
|
| 36 |
+
import torch
|
| 37 |
+
import torchvision.transforms as transforms
|
| 38 |
+
# Convert PIL to tensor and move to GPU
|
| 39 |
+
transform = transforms.Compose([transforms.Resize((448, 448)), transforms.ToTensor(), transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])])
|
| 40 |
+
# Move to GPU if available
|
| 41 |
+
tensor_img = transform(img).unsqueeze(0)
|
| 42 |
+
if torch.cuda.is_available():
|
| 43 |
+
tensor_img = tensor_img.to(device)
|
| 44 |
+
return tensor_img.cpu().numpy()
|
| 45 |
+
|
| 46 |
+
class Timer: # Report the execution time & process
|
| 47 |
+
def __init__(self):
|
| 48 |
+
self.start_time = time.perf_counter()
|
| 49 |
+
self.checkpoints = [('Start', self.start_time)]
|
| 50 |
+
|
| 51 |
+
def checkpoint(self, label='Checkpoint'):
|
| 52 |
+
now = time.perf_counter()
|
| 53 |
+
self.checkpoints.append((label, now))
|
| 54 |
+
|
| 55 |
+
def report(self, is_clear_checkpoints=True):
|
| 56 |
+
max_label_length = max(len(label) for (label, _) in self.checkpoints) if self.checkpoints else 0
|
| 57 |
+
prev_time = self.checkpoints[0][1] if self.checkpoints else self.start_time
|
| 58 |
+
|
| 59 |
+
for (label, curr_time) in self.checkpoints[1:]:
|
| 60 |
+
elapsed = curr_time - prev_time
|
| 61 |
+
print(f"{label.ljust(max_label_length)}: {elapsed:.3f} seconds")
|
| 62 |
+
prev_time = curr_time
|
| 63 |
+
|
| 64 |
+
if is_clear_checkpoints:
|
| 65 |
+
self.checkpoints.clear()
|
| 66 |
+
self.checkpoint()
|
| 67 |
+
|
| 68 |
+
def report_all(self):
|
| 69 |
+
print('\n> Execution Time Report:')
|
| 70 |
+
max_label_length = max(len(label) for (label, _) in self.checkpoints) if len(self.checkpoints) > 0 else 0
|
| 71 |
+
prev_time = self.start_time
|
| 72 |
+
|
| 73 |
+
for (label, curr_time) in self.checkpoints[1:]:
|
| 74 |
+
elapsed = curr_time - prev_time
|
| 75 |
+
print(f"{label.ljust(max_label_length)}: {elapsed:.3f} seconds")
|
| 76 |
+
prev_time = curr_time
|
| 77 |
+
|
| 78 |
+
total_time = self.checkpoints[-1][1] - self.start_time if self.checkpoints else 0
|
| 79 |
+
print(f"{'Total Execution Time'.ljust(max_label_length)}: {total_time:.3f} seconds\n") # Performance tests
|
| 80 |
+
self.checkpoints.clear()
|
| 81 |
+
|
| 82 |
+
def restart(self):
|
| 83 |
+
self.start_time = time.perf_counter()
|
| 84 |
+
self.checkpoints = [('Start', self.start_time)]
|
| 85 |
+
|
| 86 |
+
def _get_repo_id(model_name: str) -> str:
|
| 87 |
+
"""Get the repository ID for the specified model name."""
|
| 88 |
+
if '/' in model_name:
|
| 89 |
+
return model_name
|
| 90 |
+
else:
|
| 91 |
+
return f'deepghs/pixai-tagger-{model_name}-onnx'
|
| 92 |
+
|
| 93 |
+
def _download_model_files(model_name: str):
|
| 94 |
+
"""Download all required model files."""
|
| 95 |
+
repo_id = _get_repo_id(model_name)
|
| 96 |
+
|
| 97 |
+
# Download the necessary files using hf_hub_download instead of local cache...
|
| 98 |
+
model_path = hf_hub_download(
|
| 99 |
+
repo_id=repo_id,
|
| 100 |
+
filename='model.onnx',
|
| 101 |
+
library_name="pixai-tagger"
|
| 102 |
+
)
|
| 103 |
+
tags_path = hf_hub_download(
|
| 104 |
+
repo_id=repo_id,
|
| 105 |
+
filename='selected_tags.csv',
|
| 106 |
+
library_name="pixai-tagger"
|
| 107 |
+
)
|
| 108 |
+
preprocess_path = hf_hub_download(
|
| 109 |
+
repo_id=repo_id,
|
| 110 |
+
filename='preprocess.json',
|
| 111 |
+
library_name="pixai-tagger"
|
| 112 |
+
)
|
| 113 |
+
try:
|
| 114 |
+
thresholds_path = hf_hub_download(
|
| 115 |
+
repo_id=repo_id,
|
| 116 |
+
filename='thresholds.csv',
|
| 117 |
+
library_name="pixai-tagger"
|
| 118 |
+
)
|
| 119 |
+
except EntryNotFoundError:
|
| 120 |
+
thresholds_path = None
|
| 121 |
+
|
| 122 |
+
return model_path, tags_path, preprocess_path, thresholds_path
|
| 123 |
+
|
| 124 |
+
def create_optimized_ort_session(model_path):
|
| 125 |
+
"""Create an optimized ONNX Runtime session with GPU support"""
|
| 126 |
+
# Test: Session options for better performance
|
| 127 |
+
sess_options = ort.SessionOptions()
|
| 128 |
+
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
| 129 |
+
sess_options.intra_op_num_threads = 0 # Use all available cores
|
| 130 |
+
sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
|
| 131 |
+
sess_options.enable_mem_pattern = True
|
| 132 |
+
sess_options.enable_cpu_mem_arena = True
|
| 133 |
+
|
| 134 |
+
# Check available providers
|
| 135 |
+
available_providers = ort.get_available_providers()
|
| 136 |
+
print(f"Available ONNX Runtime providers: {available_providers}")
|
| 137 |
+
|
| 138 |
+
# Use appropriate execution providers (in order of preference)
|
| 139 |
+
providers = []
|
| 140 |
+
|
| 141 |
+
# Use CUDA if available
|
| 142 |
+
if 'CUDAExecutionProvider' in available_providers:
|
| 143 |
+
cuda_provider = ('CUDAExecutionProvider', {
|
| 144 |
+
'device_id': 0,
|
| 145 |
+
'arena_extend_strategy': 'kNextPowerOfTwo',
|
| 146 |
+
'gpu_mem_limit': 4 * 1024 * 1024 * 1024, # 4GB VRAM
|
| 147 |
+
'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
| 148 |
+
'do_copy_in_default_stream': True,
|
| 149 |
+
})
|
| 150 |
+
providers.append(cuda_provider)
|
| 151 |
+
print("Using CUDA provider for ONNX inference")
|
| 152 |
+
else:
|
| 153 |
+
print("CUDA provider not available, falling back to CPU")
|
| 154 |
+
|
| 155 |
+
# Always include CPU as fallback (FOR HF)
|
| 156 |
+
providers.append('CPUExecutionProvider')
|
| 157 |
+
|
| 158 |
+
try:
|
| 159 |
+
session = ort.InferenceSession(model_path, sess_options, providers=providers)
|
| 160 |
+
print(f"Model loaded with providers: {session.get_providers()}")
|
| 161 |
+
return session
|
| 162 |
+
except Exception as e:
|
| 163 |
+
print(f"Failed to create ONNX session: {e}")
|
| 164 |
+
raise
|
| 165 |
+
|
| 166 |
+
def _load_model_components_optimized(model_name: str):
|
| 167 |
+
global CURRENT_MODEL, CURRENT_MODEL_NAME, CURRENT_TAGS_DF, CURRENT_D_IPS
|
| 168 |
+
global CURRENT_PREPROCESS_FUNC, CURRENT_THRESHOLDS, CURRENT_CATEGORY_NAMES
|
| 169 |
+
|
| 170 |
+
# Only reload if model changed
|
| 171 |
+
if CURRENT_MODEL_NAME != model_name:
|
| 172 |
+
# Download files
|
| 173 |
+
model_path, tags_path, preprocess_path, thresholds_path = _download_model_files(model_name)
|
| 174 |
+
|
| 175 |
+
# Load optimized ONNX model
|
| 176 |
+
CURRENT_MODEL = create_optimized_ort_session(model_path)
|
| 177 |
+
|
| 178 |
+
# Load tags
|
| 179 |
+
CURRENT_TAGS_DF = pd.read_csv(tags_path)
|
| 180 |
+
CURRENT_D_IPS = {}
|
| 181 |
+
|
| 182 |
+
if 'ips' in CURRENT_TAGS_DF.columns:
|
| 183 |
+
CURRENT_TAGS_DF['ips'] = CURRENT_TAGS_DF['ips'].fillna('{}').map(json.loads)
|
| 184 |
+
for name, ips in zip(CURRENT_TAGS_DF['name'], CURRENT_TAGS_DF['ips']):
|
| 185 |
+
if ips:
|
| 186 |
+
CURRENT_D_IPS[name] = ips
|
| 187 |
+
|
| 188 |
+
# Load preprocessing
|
| 189 |
+
with open(preprocess_path, 'r') as f:
|
| 190 |
+
data_ = json.load(f)
|
| 191 |
+
# Simple preprocessing function
|
| 192 |
+
def transform(img):
|
| 193 |
+
# Ensure image is in RGB mode
|
| 194 |
+
if img.mode != 'RGB':
|
| 195 |
+
img = img.convert('RGB')
|
| 196 |
+
|
| 197 |
+
# Resize to 448x448 <- Very important.
|
| 198 |
+
img = img.resize((448, 448), Image.Resampling.LANCZOS)
|
| 199 |
+
|
| 200 |
+
# Convert to numpy array and normalize
|
| 201 |
+
img_array = np.array(img).astype(np.float32)
|
| 202 |
+
|
| 203 |
+
# Normalize pixel values to [0, 1]
|
| 204 |
+
img_array = img_array / 255.0
|
| 205 |
+
|
| 206 |
+
# Normalize with ImageNet mean and std
|
| 207 |
+
mean = np.array([0.48145466, 0.4578275, 0.40821073]).astype(np.float32)
|
| 208 |
+
std = np.array([0.26862954, 0.26130258, 0.27577711]).astype(np.float32)
|
| 209 |
+
img_array = (img_array - mean) / std
|
| 210 |
+
|
| 211 |
+
# Transpose to (C, H, W)
|
| 212 |
+
img_array = np.transpose(img_array, (2, 0, 1))
|
| 213 |
+
return img_array
|
| 214 |
+
|
| 215 |
+
CURRENT_PREPROCESS_FUNC = transform
|
| 216 |
+
|
| 217 |
+
# Load thresholds
|
| 218 |
+
CURRENT_THRESHOLDS = {}
|
| 219 |
+
CURRENT_CATEGORY_NAMES = {}
|
| 220 |
+
|
| 221 |
+
if thresholds_path and os.path.exists(thresholds_path):
|
| 222 |
+
df_category_thresholds = pd.read_csv(thresholds_path, keep_default_na=False)
|
| 223 |
+
for item in df_category_thresholds.to_dict('records'):
|
| 224 |
+
if item['category'] not in CURRENT_THRESHOLDS:
|
| 225 |
+
CURRENT_THRESHOLDS[item['category']] = item['threshold']
|
| 226 |
+
CURRENT_CATEGORY_NAMES[item['category']] = item['name']
|
| 227 |
+
else:
|
| 228 |
+
# Default thresholds if file doesn't exist
|
| 229 |
+
CURRENT_THRESHOLDS = {0: 0.3, 4: 0.85, 9: 0.85}
|
| 230 |
+
CURRENT_CATEGORY_NAMES = {0: 'general', 4: 'character', 9: 'rating'}
|
| 231 |
+
|
| 232 |
+
CURRENT_MODEL_NAME = model_name
|
| 233 |
+
|
| 234 |
+
return (CURRENT_MODEL, CURRENT_TAGS_DF, CURRENT_D_IPS, CURRENT_PREPROCESS_FUNC,
|
| 235 |
+
CURRENT_THRESHOLDS, CURRENT_CATEGORY_NAMES)
|
| 236 |
+
|
| 237 |
+
def _raw_predict(image: Image.Image, model_name: str):
|
| 238 |
+
"""Make a raw prediction with the PixAI tagger model."""
|
| 239 |
+
try:
|
| 240 |
+
# Ensure we have a PIL Image
|
| 241 |
+
if not isinstance(image, Image.Image):
|
| 242 |
+
raise ValueError("Input must be a PIL Image") # <-
|
| 243 |
+
|
| 244 |
+
# Load model components
|
| 245 |
+
model, _, _, preprocess_func, _, _ = _load_model_components_optimized(model_name)
|
| 246 |
+
|
| 247 |
+
# Preprocess image
|
| 248 |
+
input_tensor = preprocess_func(image)
|
| 249 |
+
|
| 250 |
+
# Add batch dimension
|
| 251 |
+
if len(input_tensor.shape) == 3:
|
| 252 |
+
input_tensor = np.expand_dims(input_tensor, axis=0)
|
| 253 |
+
|
| 254 |
+
# Run inference
|
| 255 |
+
output_names = [output.name for output in model.get_outputs()]
|
| 256 |
+
output_values = model.run(output_names, {'input': input_tensor.astype(np.float32)})
|
| 257 |
+
|
| 258 |
+
return {name: value[0] for name, value in zip(output_names, output_values)}
|
| 259 |
+
|
| 260 |
+
except Exception as e:
|
| 261 |
+
raise RuntimeError(f"Error processing image: {str(e)}")
|
| 262 |
+
|
| 263 |
+
def get_pixai_tags(
|
| 264 |
+
image: Union[str, Image.Image],
|
| 265 |
+
model_name: str = 'deepghs/pixai-tagger-v0.9-onnx',
|
| 266 |
+
thresholds: Union[float, Dict[Any, float]] = None,
|
| 267 |
+
fmt='all'
|
| 268 |
+
):
|
| 269 |
+
try:
|
| 270 |
+
# Load image if it's a path
|
| 271 |
+
if isinstance(image, str):
|
| 272 |
+
pil_image = Image.open(image)
|
| 273 |
+
elif isinstance(image, Image.Image):
|
| 274 |
+
pil_image = image
|
| 275 |
+
else:
|
| 276 |
+
raise ValueError("Image must be a file path or PIL Image")
|
| 277 |
+
|
| 278 |
+
# Load model components
|
| 279 |
+
_, df_tags, d_ips, _, default_thresholds, category_names = _load_model_components_optimized(model_name)
|
| 280 |
+
|
| 281 |
+
values = _raw_predict(pil_image, model_name)
|
| 282 |
+
prediction = values.get('prediction', np.array([]))
|
| 283 |
+
|
| 284 |
+
if prediction.size == 0:
|
| 285 |
+
raise RuntimeError("Model did not return valid predictions")
|
| 286 |
+
|
| 287 |
+
tags = {}
|
| 288 |
+
|
| 289 |
+
# Process tags by category
|
| 290 |
+
for category in sorted(set(df_tags['category'].tolist())):
|
| 291 |
+
mask = df_tags['category'] == category
|
| 292 |
+
tag_names = df_tags.loc[mask, 'name']
|
| 293 |
+
category_pred = prediction[mask]
|
| 294 |
+
|
| 295 |
+
# Determine threshold for this category
|
| 296 |
+
if isinstance(thresholds, float):
|
| 297 |
+
category_threshold = thresholds
|
| 298 |
+
elif isinstance(thresholds, dict) and \
|
| 299 |
+
(category in thresholds or category_names.get(category, '') in thresholds):
|
| 300 |
+
if category in thresholds:
|
| 301 |
+
category_threshold = thresholds[category]
|
| 302 |
+
elif category_names.get(category, '') in thresholds:
|
| 303 |
+
category_threshold = thresholds[category_names[category]]
|
| 304 |
+
else:
|
| 305 |
+
category_threshold = 0.85
|
| 306 |
+
else:
|
| 307 |
+
category_threshold = default_thresholds.get(category, 0.85)
|
| 308 |
+
|
| 309 |
+
# Apply threshold
|
| 310 |
+
pred_mask = category_pred >= category_threshold
|
| 311 |
+
filtered_tag_names = tag_names[pred_mask].tolist()
|
| 312 |
+
filtered_predictions = category_pred[pred_mask].tolist()
|
| 313 |
+
|
| 314 |
+
# Sort by confidence
|
| 315 |
+
cate_tags = dict(sorted(
|
| 316 |
+
zip(filtered_tag_names, filtered_predictions),
|
| 317 |
+
key=lambda x: (-x[1], x[0])
|
| 318 |
+
))
|
| 319 |
+
|
| 320 |
+
category_name = category_names.get(category, f"category_{category}")
|
| 321 |
+
values[category_name] = cate_tags
|
| 322 |
+
tags.update(cate_tags)
|
| 323 |
+
|
| 324 |
+
values['tag'] = tags
|
| 325 |
+
|
| 326 |
+
# Handle IPs if available
|
| 327 |
+
if 'ips' in df_tags.columns:
|
| 328 |
+
ips_mapping, ips_counts = {}, defaultdict(int)
|
| 329 |
+
for tag, _ in tags.items():
|
| 330 |
+
if tag in d_ips:
|
| 331 |
+
ips_mapping[tag] = d_ips[tag]
|
| 332 |
+
for ip_name in d_ips[tag]:
|
| 333 |
+
ips_counts[ip_name] += 1
|
| 334 |
+
values['ips_mapping'] = ips_mapping
|
| 335 |
+
values['ips_count'] = dict(ips_counts)
|
| 336 |
+
values['ips'] = [x for x, _ in sorted(ips_counts.items(), key=lambda x: (-x[1], x[0]))]
|
| 337 |
+
|
| 338 |
+
# Return based on format
|
| 339 |
+
if fmt == 'all':
|
| 340 |
+
# Return all available categories
|
| 341 |
+
available_categories = [category_names.get(cat, f"category_{cat}")
|
| 342 |
+
for cat in sorted(set(df_tags['category'].tolist()))]
|
| 343 |
+
return tuple(values.get(cat, {}) for cat in available_categories)
|
| 344 |
+
elif fmt in values:
|
| 345 |
+
return values[fmt]
|
| 346 |
+
else:
|
| 347 |
+
return values
|
| 348 |
+
|
| 349 |
+
except Exception as e:
|
| 350 |
+
raise RuntimeError(f"Error processing image: {str(e)}")
|
| 351 |
+
|
| 352 |
+
def format_ips_output(ips_result, ips_mapping):
|
| 353 |
+
"""Format IP detection output as a single string with proper escaping."""
|
| 354 |
+
if not ips_result and not ips_mapping:
|
| 355 |
+
return ""
|
| 356 |
+
|
| 357 |
+
# Format detected IPs
|
| 358 |
+
ips_list = []
|
| 359 |
+
if ips_result:
|
| 360 |
+
ips_list = [ip.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
|
| 361 |
+
for ip in ips_result]
|
| 362 |
+
|
| 363 |
+
# Format character-to-IP mapping
|
| 364 |
+
mapping_list = []
|
| 365 |
+
if ips_mapping:
|
| 366 |
+
for char, ips in ips_mapping.items():
|
| 367 |
+
formatted_char = char.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
|
| 368 |
+
formatted_ips = [ip.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
|
| 369 |
+
for ip in ips]
|
| 370 |
+
mapping_list.append(f"{formatted_char}: {', '.join(formatted_ips)}")
|
| 371 |
+
|
| 372 |
+
# Combine all into a single string
|
| 373 |
+
result_parts = []
|
| 374 |
+
if ips_list:
|
| 375 |
+
result_parts.append(", ".join(ips_list))
|
| 376 |
+
if mapping_list:
|
| 377 |
+
result_parts.extend(mapping_list)
|
| 378 |
+
|
| 379 |
+
return ", ".join(result_parts)
|
| 380 |
+
|
| 381 |
+
def process_single_image(
|
| 382 |
+
image_path,
|
| 383 |
+
model_name="deepghs/pixai-tagger-v0.9-onnx", ###
|
| 384 |
+
general_threshold=0.3,
|
| 385 |
+
character_threshold=0.85,
|
| 386 |
+
progress=None,
|
| 387 |
+
idx=0,
|
| 388 |
+
total_images=1
|
| 389 |
+
):
|
| 390 |
+
"""Process a single image and return all formatted outputs."""
|
| 391 |
+
try:
|
| 392 |
+
if image_path is None:
|
| 393 |
+
return "", "", "", "", {}, {}
|
| 394 |
+
|
| 395 |
+
if progress:
|
| 396 |
+
progress((idx)/total_images, desc=f"Processing image {idx+1}/{total_images}")
|
| 397 |
+
|
| 398 |
+
# Load image from path
|
| 399 |
+
pil_image = Image.open(image_path)
|
| 400 |
+
|
| 401 |
+
# Set thresholds
|
| 402 |
+
thresholds = {
|
| 403 |
+
'general': general_threshold,
|
| 404 |
+
'character': character_threshold
|
| 405 |
+
}
|
| 406 |
+
|
| 407 |
+
# Get all tag categories
|
| 408 |
+
all_categories = get_pixai_tags(
|
| 409 |
+
pil_image, model_name, thresholds, fmt='all'
|
| 410 |
+
)
|
| 411 |
+
|
| 412 |
+
# Ensure we have at least 3 categories (general, character, rating)
|
| 413 |
+
while len(all_categories) < 3:
|
| 414 |
+
all_categories += ({},)
|
| 415 |
+
|
| 416 |
+
general_tags = all_categories[0] if len(all_categories) > 0 else {}
|
| 417 |
+
character_tags = all_categories[1] if len(all_categories) > 1 else {}
|
| 418 |
+
rating_tags = all_categories[2] if len(all_categories) > 2 else {}
|
| 419 |
+
|
| 420 |
+
# Get IP detection data
|
| 421 |
+
ips_result = get_pixai_tags(pil_image, model_name, thresholds, fmt='ips') or []
|
| 422 |
+
ips_mapping = get_pixai_tags(pil_image, model_name, thresholds, fmt='ips_mapping') or {}
|
| 423 |
+
|
| 424 |
+
# Format character tags (names only)
|
| 425 |
+
character_names = [name.replace("(", "\\(").replace(")", "\\)").replace("_", " ") # Replacement shouldn't be necessary here, but I'll do anyway
|
| 426 |
+
for name in character_tags.keys()]
|
| 427 |
+
character_output = ", ".join(character_names)
|
| 428 |
+
|
| 429 |
+
# Format general tags (names only)
|
| 430 |
+
general_names = [name.replace("(", "\\(").replace(")", "\\)").replace("_", " ")
|
| 431 |
+
for name in general_tags.keys()]
|
| 432 |
+
general_output = ", ".join(general_names)
|
| 433 |
+
|
| 434 |
+
# Format IP detection output
|
| 435 |
+
ips_output = format_ips_output(ips_result, ips_mapping)
|
| 436 |
+
|
| 437 |
+
# Format combined tags (Character tags first, then General tags, then IP tags)
|
| 438 |
+
combined_parts = []
|
| 439 |
+
if character_names:
|
| 440 |
+
combined_parts.append(", ".join(character_names))
|
| 441 |
+
if general_names:
|
| 442 |
+
combined_parts.append(", ".join(general_names))
|
| 443 |
+
if ips_output:
|
| 444 |
+
combined_parts.append(ips_output)
|
| 445 |
+
|
| 446 |
+
combined_output = ", ".join(combined_parts)
|
| 447 |
+
|
| 448 |
+
# Get detailed JSON data
|
| 449 |
+
json_data = {
|
| 450 |
+
"character_tags": character_tags,
|
| 451 |
+
"general_tags": general_tags,
|
| 452 |
+
"rating_tags": rating_tags,
|
| 453 |
+
"ips_result": ips_result,
|
| 454 |
+
"ips_mapping": ips_mapping
|
| 455 |
+
}
|
| 456 |
+
|
| 457 |
+
# Format rating as label-compatible dict
|
| 458 |
+
rating_output = {k.replace("(", "\\(").replace(")", "\\)").replace("_", " "): v
|
| 459 |
+
for k, v in rating_tags.items()}
|
| 460 |
+
|
| 461 |
+
return (
|
| 462 |
+
character_output, # Character tags
|
| 463 |
+
general_output, # General tags
|
| 464 |
+
ips_output, # IP Detection
|
| 465 |
+
combined_output, # Combined tags
|
| 466 |
+
json_data, # Detailed JSON
|
| 467 |
+
rating_output # Rating <- Not working atm
|
| 468 |
+
)
|
| 469 |
+
except Exception as e:
|
| 470 |
+
error_msg = f"Error: {str(e)}"
|
| 471 |
+
# Return error message for all 6 outputs
|
| 472 |
+
return error_msg, error_msg, error_msg, error_msg, {}, {} # 6
|
| 473 |
+
|
| 474 |
+
"""GPU"""
|
| 475 |
+
def unload_model():
|
| 476 |
+
"""Explicitly unload the current model from memory"""
|
| 477 |
+
global CURRENT_MODEL, CURRENT_MODEL_NAME, CURRENT_TAGS_DF, CURRENT_D_IPS
|
| 478 |
+
global CURRENT_PREPROCESS_FUNC, CURRENT_THRESHOLDS, CURRENT_CATEGORY_NAMES
|
| 479 |
+
# Delete the model session
|
| 480 |
+
if CURRENT_MODEL is not None:
|
| 481 |
+
del CURRENT_MODEL
|
| 482 |
+
CURRENT_MODEL = None
|
| 483 |
+
# Clear other large objects
|
| 484 |
+
CURRENT_TAGS_DF = None
|
| 485 |
+
CURRENT_D_IPS = None
|
| 486 |
+
CURRENT_PREPROCESS_FUNC = None
|
| 487 |
+
CURRENT_THRESHOLDS = None
|
| 488 |
+
CURRENT_CATEGORY_NAMES = None
|
| 489 |
+
CURRENT_MODEL_NAME = None
|
| 490 |
+
# Force garbage collection
|
| 491 |
+
import gc
|
| 492 |
+
gc.collect()
|
| 493 |
+
# Clear CUDA cache if using GPU
|
| 494 |
+
try:
|
| 495 |
+
import torch
|
| 496 |
+
if torch.cuda.is_available():
|
| 497 |
+
torch.cuda.empty_cache()
|
| 498 |
+
except ImportError:
|
| 499 |
+
pass
|
| 500 |
+
# print("Model unloaded and memory cleared")
|
| 501 |
+
def cleanup_after_processing():
|
| 502 |
+
unload_model()
|
| 503 |
+
|
| 504 |
+
def process_gallery_images(
|
| 505 |
+
gallery,
|
| 506 |
+
model_name,
|
| 507 |
+
general_threshold,
|
| 508 |
+
character_threshold,
|
| 509 |
+
progress=gr.Progress()
|
| 510 |
+
):
|
| 511 |
+
"""Process all images in the gallery and return results with download file."""
|
| 512 |
+
if not gallery:
|
| 513 |
+
return [], "", "", "", {}, {}, {}, None
|
| 514 |
+
|
| 515 |
+
tag_results = {}
|
| 516 |
+
txt_infos = []
|
| 517 |
+
output_dir = tempfile.mkdtemp()
|
| 518 |
+
|
| 519 |
+
if not os.path.exists(output_dir):
|
| 520 |
+
os.makedirs(output_dir)
|
| 521 |
+
|
| 522 |
+
total_images = len(gallery)
|
| 523 |
+
timer = Timer()
|
| 524 |
+
|
| 525 |
+
try:
|
| 526 |
+
for idx, image_data in enumerate(gallery):
|
| 527 |
+
try:
|
| 528 |
+
image_path = image_data[0] if isinstance(image_data, (list, tuple)) else image_data
|
| 529 |
+
|
| 530 |
+
# Process image
|
| 531 |
+
results = process_single_image(
|
| 532 |
+
image_path, model_name, general_threshold, character_threshold,
|
| 533 |
+
progress, idx, total_images
|
| 534 |
+
)
|
| 535 |
+
|
| 536 |
+
# Store results
|
| 537 |
+
tag_results[image_path] = {
|
| 538 |
+
'character_tags': results[0],
|
| 539 |
+
'general_tags': results[1],
|
| 540 |
+
'ips_detection': results[2],
|
| 541 |
+
'combined_tags': results[3],
|
| 542 |
+
'json_data': results[4],
|
| 543 |
+
'rating': results[5]
|
| 544 |
+
}
|
| 545 |
+
|
| 546 |
+
# Create output files with descriptive names
|
| 547 |
+
image_name = os.path.splitext(os.path.basename(image_path))[0]
|
| 548 |
+
|
| 549 |
+
# Save all output files with descriptive prefixes
|
| 550 |
+
files_to_create = [
|
| 551 |
+
(f"character_tags-{image_name}.txt", results[0]),
|
| 552 |
+
(f"general_tags-{image_name}.txt", results[1]),
|
| 553 |
+
(f"ips_detection-{image_name}.txt", results[2]),
|
| 554 |
+
(f"combined_tags-{image_name}.txt", results[3]),
|
| 555 |
+
(f"detailed_json-{image_name}.json", json.dumps(results[4], indent=4, ensure_ascii=False))
|
| 556 |
+
]
|
| 557 |
+
|
| 558 |
+
for file_name, content in files_to_create:
|
| 559 |
+
file_path = os.path.join(output_dir, file_name)
|
| 560 |
+
with open(file_path, 'w', encoding='utf-8') as f:
|
| 561 |
+
f.write(content if isinstance(content, str) else content)
|
| 562 |
+
txt_infos.append({'path': file_path, 'name': file_name})
|
| 563 |
+
|
| 564 |
+
# Copy original image
|
| 565 |
+
original_image = Image.open(image_path)
|
| 566 |
+
image_copy_path = os.path.join(output_dir, f"{image_name}{os.path.splitext(image_path)[1]}")
|
| 567 |
+
original_image.save(image_copy_path)
|
| 568 |
+
txt_infos.append({'path': image_copy_path, 'name': f"{image_name}{os.path.splitext(image_path)[1]}"})
|
| 569 |
+
|
| 570 |
+
timer.checkpoint(f"image{idx:02d}, processed")
|
| 571 |
+
|
| 572 |
+
except Exception as e:
|
| 573 |
+
print(f"Error processing image {image_path}: {str(e)}")
|
| 574 |
+
print(traceback.format_exc())
|
| 575 |
+
continue
|
| 576 |
+
|
| 577 |
+
# Create zip file
|
| 578 |
+
download_zip_path = os.path.join(output_dir, f"Multi-Tagger-{datetime.now().strftime('%Y%m%d-%H%M%S')}.zip")
|
| 579 |
+
with zipfile.ZipFile(download_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
|
| 580 |
+
for info in txt_infos:
|
| 581 |
+
zipf.write(info['path'], arcname=info['name'])
|
| 582 |
+
# If using GPU, model will auto unload after zip file creation
|
| 583 |
+
cleanup_after_processing() # Comment here to turn off this behavior
|
| 584 |
+
|
| 585 |
+
progress(1.0, desc="Processing complete")
|
| 586 |
+
timer.report_all()
|
| 587 |
+
print('Processing is complete.')
|
| 588 |
+
|
| 589 |
+
# Return first image results as default if available even if we are tagging 1000+ images.
|
| 590 |
+
first_image_results = ("", "", "", {}, {}, "") # 6
|
| 591 |
+
if gallery and len(gallery) > 0:
|
| 592 |
+
first_image_path = gallery[0][0] if isinstance(gallery[0], (list, tuple)) else gallery[0]
|
| 593 |
+
if first_image_path in tag_results:
|
| 594 |
+
result = tag_results[first_image_path]
|
| 595 |
+
first_image_results = (
|
| 596 |
+
result['character_tags'],
|
| 597 |
+
result['general_tags'],
|
| 598 |
+
result['combined_tags'],
|
| 599 |
+
result['json_data'],
|
| 600 |
+
result['rating'],
|
| 601 |
+
result['ips_detection']
|
| 602 |
+
)
|
| 603 |
+
|
| 604 |
+
return tag_results, first_image_results[0], first_image_results[1], first_image_results[2], first_image_results[3], first_image_results[4], first_image_results[5], download_zip_path
|
| 605 |
+
|
| 606 |
+
except Exception as e:
|
| 607 |
+
print(f"Error in process_gallery_images: {str(e)}")
|
| 608 |
+
print(traceback.format_exc())
|
| 609 |
+
progress(1.0, desc="Processing failed")
|
| 610 |
+
return {}, "", "", "", {}, {}, "", None
|
| 611 |
+
|
| 612 |
+
def get_selection_from_gallery(gallery, tag_results, selected_state: gr.SelectData):
|
| 613 |
+
"""Handle gallery image selection and update UI with stored results."""
|
| 614 |
+
if not selected_state or not tag_results:
|
| 615 |
+
return "", "", "", {}, {}, ""
|
| 616 |
+
|
| 617 |
+
# Get selected image path
|
| 618 |
+
selected_value = selected_state.value
|
| 619 |
+
if isinstance(selected_value, dict) and 'image' in selected_value:
|
| 620 |
+
image_path = selected_value['image']['path']
|
| 621 |
+
elif isinstance(selected_value, (list, tuple)) and len(selected_value) > 0:
|
| 622 |
+
image_path = selected_value[0]
|
| 623 |
+
else:
|
| 624 |
+
image_path = str(selected_value)
|
| 625 |
+
|
| 626 |
+
# Retrieve stored results
|
| 627 |
+
if image_path in tag_results:
|
| 628 |
+
result = tag_results[image_path]
|
| 629 |
+
return (
|
| 630 |
+
result['character_tags'],
|
| 631 |
+
result['general_tags'],
|
| 632 |
+
result['combined_tags'],
|
| 633 |
+
result['json_data'],
|
| 634 |
+
result['rating'],
|
| 635 |
+
result['ips_detection']
|
| 636 |
+
)
|
| 637 |
+
|
| 638 |
+
# Return empty if not found
|
| 639 |
+
return "", "", "", {}, {}, ""
|
| 640 |
+
|
| 641 |
+
def append_gallery(gallery, image):
|
| 642 |
+
"""Add a single image to the gallery."""
|
| 643 |
+
if gallery is None:
|
| 644 |
+
gallery = []
|
| 645 |
+
if not image:
|
| 646 |
+
return gallery, None
|
| 647 |
+
gallery.append(image)
|
| 648 |
+
return gallery, None
|
| 649 |
+
|
| 650 |
+
def extend_gallery(gallery, images):
|
| 651 |
+
"""Add multiple images to the gallery."""
|
| 652 |
+
if gallery is None:
|
| 653 |
+
gallery = []
|
| 654 |
+
if not images:
|
| 655 |
+
return gallery
|
| 656 |
+
gallery.extend(images)
|
| 657 |
+
return gallery
|
| 658 |
+
|
| 659 |
+
def create_pixai_interface():
|
| 660 |
+
"""Create the PixAI Gradio interface"""
|
| 661 |
+
with gr.Blocks(css=css, fill_width=True) as demo:
|
| 662 |
+
# gr.Markdown("Upload anime-style images to extract tags using PixAI")
|
| 663 |
+
# State to store results
|
| 664 |
+
tag_results = gr.State({})
|
| 665 |
+
selected_image = gr.Textbox(label='Selected Image', visible=False)
|
| 666 |
+
|
| 667 |
+
with gr.Row():
|
| 668 |
+
with gr.Column():
|
| 669 |
+
# Image upload section
|
| 670 |
+
with gr.Column(variant='panel'):
|
| 671 |
+
image_input = gr.Image(
|
| 672 |
+
label='Upload an Image or clicking paste from clipboard button',
|
| 673 |
+
type='filepath',
|
| 674 |
+
sources=['upload', 'clipboard'],
|
| 675 |
+
height=150
|
| 676 |
+
)
|
| 677 |
+
with gr.Row():
|
| 678 |
+
upload_button = gr.UploadButton(
|
| 679 |
+
'Upload multiple images',
|
| 680 |
+
file_types=['image'],
|
| 681 |
+
file_count='multiple',
|
| 682 |
+
size='sm'
|
| 683 |
+
)
|
| 684 |
+
gallery = gr.Gallery(
|
| 685 |
+
columns=2,
|
| 686 |
+
show_share_button=False,
|
| 687 |
+
interactive=True,
|
| 688 |
+
height='auto',
|
| 689 |
+
label='Grid of images',
|
| 690 |
+
preview=False,
|
| 691 |
+
elem_id='custom-gallery'
|
| 692 |
+
)
|
| 693 |
+
run_button = gr.Button("Analyze Images", variant="primary", size='lg')
|
| 694 |
+
model_dropdown = gr.Dropdown(
|
| 695 |
+
choices=["deepghs/pixai-tagger-v0.9-onnx"],
|
| 696 |
+
value="deepghs/pixai-tagger-v0.9-onnx",
|
| 697 |
+
label="Model"
|
| 698 |
+
)
|
| 699 |
+
# Threshold controls
|
| 700 |
+
with gr.Row():
|
| 701 |
+
general_threshold = gr.Slider(
|
| 702 |
+
minimum=0.0, maximum=1.0, value=0.30, step=0.05,
|
| 703 |
+
label="General Tags Threshold", scale=3
|
| 704 |
+
)
|
| 705 |
+
character_threshold = gr.Slider(
|
| 706 |
+
minimum=0.0, maximum=1.0, value=0.85, step=0.05,
|
| 707 |
+
label="Character Tags Threshold", scale=3
|
| 708 |
+
)
|
| 709 |
+
|
| 710 |
+
with gr.Row():
|
| 711 |
+
clear = gr.ClearButton(
|
| 712 |
+
components=[gallery, model_dropdown, general_threshold, character_threshold],
|
| 713 |
+
variant='secondary',
|
| 714 |
+
size='lg'
|
| 715 |
+
)
|
| 716 |
+
clear.add([tag_results])
|
| 717 |
+
detailed_json_output = gr.JSON(label="Detailed JSON")
|
| 718 |
+
|
| 719 |
+
with gr.Column(variant='panel'):
|
| 720 |
+
|
| 721 |
+
download_file = gr.File(label="Download")
|
| 722 |
+
|
| 723 |
+
# Output blocks
|
| 724 |
+
character_tags_output = gr.Textbox(
|
| 725 |
+
label="Character tags",
|
| 726 |
+
show_copy_button=True,
|
| 727 |
+
lines=3
|
| 728 |
+
)
|
| 729 |
+
general_tags_output = gr.Textbox(
|
| 730 |
+
label="General tags",
|
| 731 |
+
show_copy_button=True,
|
| 732 |
+
lines=3
|
| 733 |
+
)
|
| 734 |
+
ips_detection_output = gr.Textbox(
|
| 735 |
+
label="IPs Detection",
|
| 736 |
+
show_copy_button=True,
|
| 737 |
+
lines=5
|
| 738 |
+
)
|
| 739 |
+
combined_tags_output = gr.Textbox(
|
| 740 |
+
label="Combined tags",
|
| 741 |
+
show_copy_button=True,
|
| 742 |
+
lines=6
|
| 743 |
+
)
|
| 744 |
+
rating_output = gr.Label(label="Rating")
|
| 745 |
+
|
| 746 |
+
# Clear button targets
|
| 747 |
+
clear.add([
|
| 748 |
+
download_file,
|
| 749 |
+
character_tags_output,
|
| 750 |
+
general_tags_output,
|
| 751 |
+
ips_detection_output,
|
| 752 |
+
combined_tags_output,
|
| 753 |
+
rating_output,
|
| 754 |
+
detailed_json_output
|
| 755 |
+
])
|
| 756 |
+
|
| 757 |
+
# Event handlers
|
| 758 |
+
image_input.change(
|
| 759 |
+
append_gallery,
|
| 760 |
+
inputs=[gallery, image_input],
|
| 761 |
+
outputs=[gallery, image_input]
|
| 762 |
+
)
|
| 763 |
+
|
| 764 |
+
upload_button.upload(
|
| 765 |
+
extend_gallery,
|
| 766 |
+
inputs=[gallery, upload_button],
|
| 767 |
+
outputs=gallery
|
| 768 |
+
)
|
| 769 |
+
|
| 770 |
+
gallery.select(
|
| 771 |
+
get_selection_from_gallery,
|
| 772 |
+
inputs=[gallery, tag_results],
|
| 773 |
+
outputs=[
|
| 774 |
+
character_tags_output,
|
| 775 |
+
general_tags_output,
|
| 776 |
+
combined_tags_output,
|
| 777 |
+
detailed_json_output,
|
| 778 |
+
rating_output,
|
| 779 |
+
ips_detection_output
|
| 780 |
+
]
|
| 781 |
+
)
|
| 782 |
+
|
| 783 |
+
run_button.click(
|
| 784 |
+
process_gallery_images,
|
| 785 |
+
inputs=[gallery, model_dropdown, general_threshold, character_threshold],
|
| 786 |
+
outputs=[
|
| 787 |
+
tag_results,
|
| 788 |
+
character_tags_output,
|
| 789 |
+
general_tags_output,
|
| 790 |
+
combined_tags_output,
|
| 791 |
+
detailed_json_output,
|
| 792 |
+
rating_output,
|
| 793 |
+
ips_detection_output,
|
| 794 |
+
download_file
|
| 795 |
+
]
|
| 796 |
+
)
|
| 797 |
+
|
| 798 |
+
gr.Markdown('[Based on Source code for imgutils.tagging.pixai](https://dghs-imgutils.deepghs.org/main/_modules/imgutils/tagging/pixai.html) & [pixai-labs/pixai-tagger-demo](https://huggingface.co/spaces/pixai-labs/pixai-tagger-demo)')
|
| 799 |
+
|
| 800 |
+
return demo
|
| 801 |
+
|
| 802 |
+
# Export public API
|
| 803 |
+
__all__ = [
|
| 804 |
+
'get_pixai_tags',
|
| 805 |
+
'process_single_image',
|
| 806 |
+
'process_gallery_images',
|
| 807 |
+
'create_pixai_interface',
|
| 808 |
+
'unload_model',
|
| 809 |
+
'cleanup_after_processing'
|
| 810 |
+
]
|
requirements.txt
CHANGED
|
@@ -1,20 +1,19 @@
|
|
| 1 |
-
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
https://github.com/kingbri1/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu128torch2.7.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
|
|
|
|
| 1 |
+
torch
|
| 2 |
+
gradio
|
| 3 |
+
safetensors
|
| 4 |
+
sentencepiece
|
| 5 |
+
pillow
|
| 6 |
+
requests
|
| 7 |
+
numpy
|
| 8 |
+
timm
|
| 9 |
+
einops
|
| 10 |
+
optimum
|
| 11 |
+
accelerate
|
| 12 |
+
opencv-python
|
| 13 |
+
onnxruntime>=1.12.0
|
| 14 |
+
matplotlib
|
| 15 |
+
apscheduler
|
| 16 |
+
spaces
|
| 17 |
+
pandas==2.1.2
|
| 18 |
+
huggingface-hub
|
| 19 |
+
transformers
|
|
|