Spaces:
Running
Running
Upload 13 files
Browse files- README.md +5 -5
- app.py +120 -12
- model.py +13 -0
- multit2i.py +186 -16
- requirements.txt +12 -1
- tagger/character_series_dict.csv +0 -0
- tagger/danbooru_e621.csv +0 -0
- tagger/fl2sd3longcap.py +74 -0
- tagger/output.py +16 -0
- tagger/tag_group.csv +0 -0
- tagger/tagger.py +549 -0
- tagger/utils.py +45 -0
- tagger/v2.py +260 -0
README.md
CHANGED
@@ -1,13 +1,13 @@
|
|
1 |
---
|
2 |
-
title:
|
3 |
-
emoji:
|
4 |
-
colorFrom:
|
5 |
-
colorTo:
|
6 |
sdk: gradio
|
7 |
sdk_version: 4.39.0
|
8 |
app_file: app.py
|
9 |
-
pinned: false
|
10 |
short_description: Text-to-Image
|
|
|
11 |
---
|
12 |
|
13 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
|
|
1 |
---
|
2 |
+
title: Free Multi Models Text-to-Image Heavy-Armed Demo
|
3 |
+
emoji: 🌐🌊
|
4 |
+
colorFrom: blue
|
5 |
+
colorTo: purple
|
6 |
sdk: gradio
|
7 |
sdk_version: 4.39.0
|
8 |
app_file: app.py
|
|
|
9 |
short_description: Text-to-Image
|
10 |
+
pinned: true
|
11 |
---
|
12 |
|
13 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
app.py
CHANGED
@@ -1,33 +1,105 @@
|
|
1 |
import gradio as gr
|
|
|
2 |
from multit2i import (
|
3 |
load_models,
|
4 |
-
find_model_list,
|
5 |
infer_multi,
|
|
|
6 |
save_gallery_images,
|
7 |
change_model,
|
8 |
get_model_info_md,
|
9 |
loaded_models,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
10 |
)
|
11 |
|
12 |
|
13 |
-
models
|
14 |
-
load_models(models,
|
15 |
|
16 |
|
17 |
-
css = """
|
|
|
|
|
18 |
|
19 |
with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
|
20 |
with gr.Column():
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
27 |
results = gr.Gallery(label="Gallery", interactive=False, show_download_button=True, show_share_button=False,
|
28 |
container=True, format="png", object_fit="contain")
|
29 |
image_files = gr.Files(label="Download", interactive=False)
|
30 |
-
clear_results = gr.Button("Clear Gallery
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
31 |
gr.Markdown(
|
32 |
f"""This demo was created in reference to the following demos.
|
33 |
- [Nymbo/Flood](https://huggingface.co/spaces/Nymbo/Flood).
|
@@ -40,13 +112,49 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
|
|
40 |
gr.on(
|
41 |
triggers=[run_button.click, prompt.submit],
|
42 |
fn=infer_multi,
|
43 |
-
inputs=[prompt,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
44 |
outputs=[results],
|
45 |
queue=True,
|
46 |
show_progress="full",
|
47 |
show_api=True,
|
48 |
).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
|
|
|
49 |
clear_results.click(lambda: (None, None), None, [results, image_files], queue=False, show_api=False)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
50 |
|
51 |
demo.queue()
|
52 |
demo.launch()
|
|
|
1 |
import gradio as gr
|
2 |
+
from model import models
|
3 |
from multit2i import (
|
4 |
load_models,
|
|
|
5 |
infer_multi,
|
6 |
+
infer_multi_random,
|
7 |
save_gallery_images,
|
8 |
change_model,
|
9 |
get_model_info_md,
|
10 |
loaded_models,
|
11 |
+
get_positive_prefix,
|
12 |
+
get_positive_suffix,
|
13 |
+
get_negative_prefix,
|
14 |
+
get_negative_suffix,
|
15 |
+
get_recom_prompt_type,
|
16 |
+
set_recom_prompt_preset,
|
17 |
+
get_tag_type,
|
18 |
+
)
|
19 |
+
from tagger.tagger import (
|
20 |
+
predict_tags_wd,
|
21 |
+
remove_specific_prompt,
|
22 |
+
convert_danbooru_to_e621_prompt,
|
23 |
+
insert_recom_prompt,
|
24 |
+
)
|
25 |
+
from tagger.fl2sd3longcap import predict_tags_fl2_sd3
|
26 |
+
from tagger.v2 import (
|
27 |
+
V2_ALL_MODELS,
|
28 |
+
v2_random_prompt,
|
29 |
+
)
|
30 |
+
from tagger.utils import (
|
31 |
+
V2_ASPECT_RATIO_OPTIONS,
|
32 |
+
V2_RATING_OPTIONS,
|
33 |
+
V2_LENGTH_OPTIONS,
|
34 |
+
V2_IDENTITY_OPTIONS,
|
35 |
)
|
36 |
|
37 |
|
38 |
+
load_models(models, 5)
|
39 |
+
#load_models(models, 20) # Fetching 20 models at the same time. default: 5
|
40 |
|
41 |
|
42 |
+
css = """
|
43 |
+
#model_info { text-align: center; display:block; }
|
44 |
+
"""
|
45 |
|
46 |
with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
|
47 |
with gr.Column():
|
48 |
+
with gr.Accordion("Advanced settings", open=True):
|
49 |
+
with gr.Accordion("Recommended Prompt", open=False):
|
50 |
+
recom_prompt_preset = gr.Radio(label="Set Presets", choices=get_recom_prompt_type(), value="Common")
|
51 |
+
with gr.Row():
|
52 |
+
positive_prefix = gr.CheckboxGroup(label="Use Positive Prefix", choices=get_positive_prefix(), value=[])
|
53 |
+
positive_suffix = gr.CheckboxGroup(label="Use Positive Suffix", choices=get_positive_suffix(), value=["Common"])
|
54 |
+
negative_prefix = gr.CheckboxGroup(label="Use Negative Prefix", choices=get_negative_prefix(), value=[], visible=False)
|
55 |
+
negative_suffix = gr.CheckboxGroup(label="Use Negative Suffix", choices=get_negative_suffix(), value=["Common"], visible=False)
|
56 |
+
with gr.Accordion("Prompt Transformer", open=False):
|
57 |
+
v2_rating = gr.Radio(label="Rating", choices=list(V2_RATING_OPTIONS), value="sfw")
|
58 |
+
v2_aspect_ratio = gr.Radio(label="Aspect ratio", info="The aspect ratio of the image.", choices=list(V2_ASPECT_RATIO_OPTIONS), value="square", visible=False)
|
59 |
+
v2_length = gr.Radio(label="Length", info="The total length of the tags.", choices=list(V2_LENGTH_OPTIONS), value="long")
|
60 |
+
v2_identity = gr.Radio(label="Keep identity", info="How strictly to keep the identity of the character or subject. If you specify the detail of subject in the prompt, you should choose `strict`. Otherwise, choose `none` or `lax`. `none` is very creative but sometimes ignores the input prompt.", choices=list(V2_IDENTITY_OPTIONS), value="lax")
|
61 |
+
v2_ban_tags = gr.Textbox(label="Ban tags", info="Tags to ban from the output.", placeholder="alternate costumen, ...", value="censored")
|
62 |
+
v2_tag_type = gr.Radio(label="Tag Type", info="danbooru for common, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru", visible=False)
|
63 |
+
v2_model = gr.Dropdown(label="Model", choices=list(V2_ALL_MODELS.keys()), value=list(V2_ALL_MODELS.keys())[0])
|
64 |
+
with gr.Accordion("Model", open=True):
|
65 |
+
model_name = gr.Dropdown(label="Select Model", show_label=False, choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0], allow_custom_value=True)
|
66 |
+
model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]), elem_id="model_info")
|
67 |
+
with gr.Group():
|
68 |
+
with gr.Accordion("Prompt from Image File", open=False):
|
69 |
+
tagger_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
|
70 |
+
with gr.Accordion(label="Advanced options", open=False):
|
71 |
+
tagger_general_threshold = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True)
|
72 |
+
tagger_character_threshold = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
|
73 |
+
tagger_tag_type = gr.Radio(label="Convert tags to", info="danbooru for common, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
|
74 |
+
tagger_recom_prompt = gr.Radio(label="Insert reccomended prompt", choices=["None", "Animagine", "Pony"], value="None", interactive=True)
|
75 |
+
tagger_keep_tags = gr.Radio(label="Remove tags leaving only the following", choices=["body", "dress", "all"], value="all")
|
76 |
+
tagger_algorithms = gr.CheckboxGroup(["Use WD Tagger", "Use Florence-2-SD3-Long-Captioner"], label="Algorithms", value=["Use WD Tagger"])
|
77 |
+
tagger_generate_from_image = gr.Button(value="Generate Tags from Image")
|
78 |
+
with gr.Row():
|
79 |
+
v2_character = gr.Textbox(label="Character", placeholder="hatsune miku", scale=2)
|
80 |
+
v2_series = gr.Textbox(label="Series", placeholder="vocaloid", scale=2)
|
81 |
+
random_prompt = gr.Button(value="Extend Prompt 🎲", size="sm", scale=1)
|
82 |
+
clear_prompt = gr.Button(value="Clear Prompt 🗑️", size="sm", scale=1)
|
83 |
+
prompt = gr.Text(label="Prompt", lines=1, max_lines=8, placeholder="1girl, solo, ...", show_copy_button=True)
|
84 |
+
neg_prompt = gr.Text(label="Negative Prompt", lines=1, max_lines=8, placeholder="", visible=False)
|
85 |
+
with gr.Row():
|
86 |
+
run_button = gr.Button("Generate Image", scale=6)
|
87 |
+
random_button = gr.Button("Random Model 🎲", scale=3)
|
88 |
+
image_num = gr.Number(label="Count", minimum=1, maximum=16, value=1, step=1, interactive=True, scale=1)
|
89 |
results = gr.Gallery(label="Gallery", interactive=False, show_download_button=True, show_share_button=False,
|
90 |
container=True, format="png", object_fit="contain")
|
91 |
image_files = gr.Files(label="Download", interactive=False)
|
92 |
+
clear_results = gr.Button("Clear Gallery / Download")
|
93 |
+
examples = gr.Examples(
|
94 |
+
examples = [
|
95 |
+
["souryuu asuka langley, 1girl, neon genesis evangelion, plugsuit, pilot suit, red bodysuit, sitting, crossing legs, black eye patch, cat hat, throne, symmetrical, looking down, from bottom, looking at viewer, outdoors"],
|
96 |
+
["sailor moon, magical girl transformation, sparkles and ribbons, soft pastel colors, crescent moon motif, starry night sky background, shoujo manga style"],
|
97 |
+
["kafuu chino, 1girl, solo"],
|
98 |
+
["1girl"],
|
99 |
+
["beautiful sunset"],
|
100 |
+
],
|
101 |
+
inputs=[prompt],
|
102 |
+
)
|
103 |
gr.Markdown(
|
104 |
f"""This demo was created in reference to the following demos.
|
105 |
- [Nymbo/Flood](https://huggingface.co/spaces/Nymbo/Flood).
|
|
|
112 |
gr.on(
|
113 |
triggers=[run_button.click, prompt.submit],
|
114 |
fn=infer_multi,
|
115 |
+
inputs=[prompt, neg_prompt, results, image_num, model_name,
|
116 |
+
positive_prefix, positive_suffix, negative_prefix, negative_suffix],
|
117 |
+
outputs=[results],
|
118 |
+
queue=True,
|
119 |
+
show_progress="full",
|
120 |
+
show_api=True,
|
121 |
+
).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
|
122 |
+
gr.on(
|
123 |
+
triggers=[random_button.click],
|
124 |
+
fn=infer_multi_random,
|
125 |
+
inputs=[prompt, neg_prompt, results, image_num,
|
126 |
+
positive_prefix, positive_suffix, negative_prefix, negative_suffix],
|
127 |
outputs=[results],
|
128 |
queue=True,
|
129 |
show_progress="full",
|
130 |
show_api=True,
|
131 |
).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
|
132 |
+
clear_prompt.click(lambda: (None, None, None), None, [prompt, v2_series, v2_character], queue=False, show_api=False)
|
133 |
clear_results.click(lambda: (None, None), None, [results, image_files], queue=False, show_api=False)
|
134 |
+
recom_prompt_preset.change(set_recom_prompt_preset, [recom_prompt_preset],
|
135 |
+
[positive_prefix, positive_suffix, negative_prefix, negative_suffix], queue=False, show_api=False)
|
136 |
+
random_prompt.click(
|
137 |
+
v2_random_prompt, [prompt, v2_series, v2_character, v2_rating, v2_aspect_ratio, v2_length,
|
138 |
+
v2_identity, v2_ban_tags, v2_model], [prompt, v2_series, v2_character], show_api=False,
|
139 |
+
).success(
|
140 |
+
get_tag_type, [positive_prefix, positive_suffix, negative_prefix, negative_suffix], [v2_tag_type], queue=False, show_api=False
|
141 |
+
).success(
|
142 |
+
convert_danbooru_to_e621_prompt, [prompt, v2_tag_type], [prompt], queue=False, show_api=False,
|
143 |
+
)
|
144 |
+
tagger_generate_from_image.click(
|
145 |
+
predict_tags_wd,
|
146 |
+
[tagger_image, prompt, tagger_algorithms, tagger_general_threshold, tagger_character_threshold],
|
147 |
+
[v2_series, v2_character, prompt, gr.Button(visible=False)],
|
148 |
+
show_api=False,
|
149 |
+
).success(
|
150 |
+
predict_tags_fl2_sd3, [tagger_image, prompt, tagger_algorithms], [prompt], show_api=False,
|
151 |
+
).success(
|
152 |
+
remove_specific_prompt, [prompt, tagger_keep_tags], [prompt], queue=False, show_api=False,
|
153 |
+
).success(
|
154 |
+
convert_danbooru_to_e621_prompt, [prompt, tagger_tag_type], [prompt], queue=False, show_api=False,
|
155 |
+
).success(
|
156 |
+
insert_recom_prompt, [prompt, neg_prompt, tagger_recom_prompt], [prompt, neg_prompt], queue=False, show_api=False,
|
157 |
+
)
|
158 |
|
159 |
demo.queue()
|
160 |
demo.launch()
|
model.py
ADDED
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from multit2i import find_model_list
|
2 |
+
|
3 |
+
|
4 |
+
models = find_model_list("digiplay", [], "", "likes", 100)
|
5 |
+
|
6 |
+
|
7 |
+
# Examples:
|
8 |
+
#models = ['yodayo-ai/kivotos-xl-2.0', 'yodayo-ai/holodayo-xl-2.1'] # specific models
|
9 |
+
#models = find_model_list("John6666", [], "", "last_modified", 20) # John6666's latest 20 models
|
10 |
+
#models = find_model_list("John6666", ["anime"], "", "last_modified", 20) # John6666's latest 20 models with 'anime' tag
|
11 |
+
#models = find_model_list("John6666", [], "anime", "last_modified", 20) # John6666's latest 20 models without 'anime' tag
|
12 |
+
#models = find_model_list("", [], "", "last_modified", 20) # latest 20 text-to-image models of huggingface
|
13 |
+
#models = find_model_list("", [], "", "downloads", 20) # monthly most downloaded 20 text-to-image models of huggingface
|
multit2i.py
CHANGED
@@ -1,12 +1,18 @@
|
|
1 |
import gradio as gr
|
2 |
import asyncio
|
|
|
3 |
from pathlib import Path
|
4 |
|
5 |
|
|
|
6 |
loaded_models = {}
|
7 |
model_info_dict = {}
|
8 |
|
9 |
|
|
|
|
|
|
|
|
|
10 |
def list_sub(a, b):
|
11 |
return [e for e in a if e not in b]
|
12 |
|
@@ -27,7 +33,7 @@ def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="l
|
|
27 |
if not sort: sort = "last_modified"
|
28 |
models = []
|
29 |
try:
|
30 |
-
model_infos = api.list_models(author=author,
|
31 |
tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit * 5)
|
32 |
except Exception as e:
|
33 |
print(f"Error: Failed to list models.")
|
@@ -55,7 +61,8 @@ def get_t2i_model_info_dict(repo_id: str):
|
|
55 |
if model.private or model.gated: return info
|
56 |
try:
|
57 |
tags = model.tags
|
58 |
-
except Exception:
|
|
|
59 |
return info
|
60 |
if not 'diffusers' in model.tags: return info
|
61 |
if 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
|
@@ -63,8 +70,7 @@ def get_t2i_model_info_dict(repo_id: str):
|
|
63 |
elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3"
|
64 |
else: info["ver"] = "Other"
|
65 |
info["url"] = f"https://huggingface.co/{repo_id}/"
|
66 |
-
if model.card_data and model.card_data.tags
|
67 |
-
info["tags"] = model.card_data.tags
|
68 |
info["downloads"] = model.downloads
|
69 |
info["likes"] = model.likes
|
70 |
info["last_modified"] = model.last_modified.strftime("lastmod: %Y-%m-%d")
|
@@ -101,12 +107,44 @@ def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
|
|
101 |
return gr.update(value=output_images), gr.update(value=output_paths)
|
102 |
|
103 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
104 |
def load_model(model_name: str):
|
105 |
global loaded_models
|
106 |
global model_info_dict
|
107 |
if model_name in loaded_models.keys(): return loaded_models[model_name]
|
108 |
try:
|
109 |
-
loaded_models[model_name] =
|
110 |
print(f"Loaded: {model_name}")
|
111 |
except Exception as e:
|
112 |
if model_name in loaded_models.keys(): del loaded_models[model_name]
|
@@ -115,8 +153,10 @@ def load_model(model_name: str):
|
|
115 |
return None
|
116 |
try:
|
117 |
model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
|
|
|
118 |
except Exception as e:
|
119 |
if model_name in model_info_dict.keys(): del model_info_dict[model_name]
|
|
|
120 |
print(e)
|
121 |
return loaded_models[model_name]
|
122 |
|
@@ -126,21 +166,124 @@ async def async_load_models(models: list, limit: int=5):
|
|
126 |
async def async_load_model(model: str):
|
127 |
async with sem:
|
128 |
try:
|
129 |
-
return load_model
|
130 |
except Exception as e:
|
131 |
print(e)
|
132 |
tasks = [asyncio.create_task(async_load_model(model)) for model in models]
|
133 |
-
return await asyncio.
|
134 |
|
135 |
|
136 |
def load_models(models: list, limit: int=5):
|
137 |
-
loop = asyncio.
|
138 |
try:
|
139 |
loop.run_until_complete(async_load_models(models, limit))
|
140 |
except Exception as e:
|
141 |
print(e)
|
142 |
pass
|
143 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
144 |
|
145 |
|
146 |
def get_model_info_md(model_name: str):
|
@@ -152,29 +295,56 @@ def change_model(model_name: str):
|
|
152 |
return get_model_info_md(model_name)
|
153 |
|
154 |
|
155 |
-
def infer(prompt: str,
|
156 |
from PIL import Image
|
157 |
import random
|
158 |
seed = ""
|
159 |
rand = random.randint(1, 500)
|
160 |
for i in range(rand):
|
161 |
seed += " "
|
162 |
-
rprompt = ", highly detailed, masterpiece, best quality, very aesthetic, absurdres, " if recom_prompt else ""
|
163 |
caption = model_name.split("/")[-1]
|
164 |
try:
|
165 |
model = load_model(model_name)
|
166 |
if not model: return (Image.Image(), None)
|
167 |
-
image_path = model(prompt +
|
168 |
-
image = Image.open(image_path).convert('
|
169 |
except Exception as e:
|
170 |
print(e)
|
171 |
return (Image.Image(), None)
|
172 |
return (image, caption)
|
173 |
|
174 |
|
175 |
-
def infer_multi(prompt: str,
|
|
|
|
|
176 |
image_num = int(image_num)
|
177 |
images = results if results else []
|
178 |
-
|
179 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
180 |
yield images
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import gradio as gr
|
2 |
import asyncio
|
3 |
+
from threading import RLock
|
4 |
from pathlib import Path
|
5 |
|
6 |
|
7 |
+
lock = RLock()
|
8 |
loaded_models = {}
|
9 |
model_info_dict = {}
|
10 |
|
11 |
|
12 |
+
def to_list(s):
|
13 |
+
return [x.strip() for x in s.split(",")]
|
14 |
+
|
15 |
+
|
16 |
def list_sub(a, b):
|
17 |
return [e for e in a if e not in b]
|
18 |
|
|
|
33 |
if not sort: sort = "last_modified"
|
34 |
models = []
|
35 |
try:
|
36 |
+
model_infos = api.list_models(author=author, pipeline_tag="text-to-image",
|
37 |
tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit * 5)
|
38 |
except Exception as e:
|
39 |
print(f"Error: Failed to list models.")
|
|
|
61 |
if model.private or model.gated: return info
|
62 |
try:
|
63 |
tags = model.tags
|
64 |
+
except Exception as e:
|
65 |
+
print(e)
|
66 |
return info
|
67 |
if not 'diffusers' in model.tags: return info
|
68 |
if 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
|
|
|
70 |
elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3"
|
71 |
else: info["ver"] = "Other"
|
72 |
info["url"] = f"https://huggingface.co/{repo_id}/"
|
73 |
+
info["tags"] = model.card_data.tags if model.card_data and model.card_data.tags else []
|
|
|
74 |
info["downloads"] = model.downloads
|
75 |
info["likes"] = model.likes
|
76 |
info["last_modified"] = model.last_modified.strftime("lastmod: %Y-%m-%d")
|
|
|
107 |
return gr.update(value=output_images), gr.update(value=output_paths)
|
108 |
|
109 |
|
110 |
+
def load_from_model(model_name: str, hf_token: str = None):
|
111 |
+
import httpx
|
112 |
+
import huggingface_hub
|
113 |
+
from gradio.exceptions import ModelNotFoundError
|
114 |
+
model_url = f"https://huggingface.co/{model_name}"
|
115 |
+
api_url = f"https://api-inference.huggingface.co/models/{model_name}"
|
116 |
+
print(f"Fetching model from: {model_url}")
|
117 |
+
|
118 |
+
headers = {"Authorization": f"Bearer {hf_token}"} if hf_token is not None else {}
|
119 |
+
response = httpx.request("GET", api_url, headers=headers)
|
120 |
+
if response.status_code != 200:
|
121 |
+
raise ModelNotFoundError(
|
122 |
+
f"Could not find model: {model_name}. If it is a private or gated model, please provide your Hugging Face access token (https://huggingface.co/settings/tokens) as the argument for the `hf_token` parameter."
|
123 |
+
)
|
124 |
+
headers["X-Wait-For-Model"] = "true"
|
125 |
+
client = huggingface_hub.InferenceClient(model=model_name, headers=headers, token=hf_token)
|
126 |
+
inputs = gr.components.Textbox(label="Input")
|
127 |
+
outputs = gr.components.Image(label="Output")
|
128 |
+
fn = client.text_to_image
|
129 |
+
|
130 |
+
def query_huggingface_inference_endpoints(*data):
|
131 |
+
return fn(*data)
|
132 |
+
|
133 |
+
interface_info = {
|
134 |
+
"fn": query_huggingface_inference_endpoints,
|
135 |
+
"inputs": inputs,
|
136 |
+
"outputs": outputs,
|
137 |
+
"title": model_name,
|
138 |
+
}
|
139 |
+
return gr.Interface(**interface_info)
|
140 |
+
|
141 |
+
|
142 |
def load_model(model_name: str):
|
143 |
global loaded_models
|
144 |
global model_info_dict
|
145 |
if model_name in loaded_models.keys(): return loaded_models[model_name]
|
146 |
try:
|
147 |
+
loaded_models[model_name] = load_from_model(model_name)
|
148 |
print(f"Loaded: {model_name}")
|
149 |
except Exception as e:
|
150 |
if model_name in loaded_models.keys(): del loaded_models[model_name]
|
|
|
153 |
return None
|
154 |
try:
|
155 |
model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
|
156 |
+
print(f"Assigned: {model_name}")
|
157 |
except Exception as e:
|
158 |
if model_name in model_info_dict.keys(): del model_info_dict[model_name]
|
159 |
+
print(f"Failed to assigned: {model_name}")
|
160 |
print(e)
|
161 |
return loaded_models[model_name]
|
162 |
|
|
|
166 |
async def async_load_model(model: str):
|
167 |
async with sem:
|
168 |
try:
|
169 |
+
return await asyncio.to_thread(load_model, model)
|
170 |
except Exception as e:
|
171 |
print(e)
|
172 |
tasks = [asyncio.create_task(async_load_model(model)) for model in models]
|
173 |
+
return await asyncio.gather(*tasks, return_exceptions=True)
|
174 |
|
175 |
|
176 |
def load_models(models: list, limit: int=5):
|
177 |
+
loop = asyncio.new_event_loop()
|
178 |
try:
|
179 |
loop.run_until_complete(async_load_models(models, limit))
|
180 |
except Exception as e:
|
181 |
print(e)
|
182 |
pass
|
183 |
+
finally:
|
184 |
+
loop.close()
|
185 |
+
|
186 |
+
|
187 |
+
positive_prefix = {
|
188 |
+
"Pony": to_list("score_9, score_8_up, score_7_up"),
|
189 |
+
"Pony Anime": to_list("source_anime, anime, score_9, score_8_up, score_7_up"),
|
190 |
+
}
|
191 |
+
positive_suffix = {
|
192 |
+
"Common": to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres"),
|
193 |
+
"Anime": to_list("anime artwork, anime style, studio anime, highly detailed"),
|
194 |
+
}
|
195 |
+
negative_prefix = {
|
196 |
+
"Pony": to_list("score_6, score_5, score_4"),
|
197 |
+
"Pony Anime": to_list("score_6, score_5, score_4, source_pony, source_furry, source_cartoon"),
|
198 |
+
"Pony Real": to_list("score_6, score_5, score_4, source_anime, source_pony, source_furry, source_cartoon"),
|
199 |
+
}
|
200 |
+
negative_suffix = {
|
201 |
+
"Common": to_list("lowres, (bad), bad hands, bad feet, text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"),
|
202 |
+
"Pony Anime": to_list("busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends"),
|
203 |
+
"Pony Real": to_list("ugly, airbrushed, simple background, cgi, cartoon, anime"),
|
204 |
+
}
|
205 |
+
positive_all = negative_all = []
|
206 |
+
for k, v in (positive_prefix | positive_suffix).items():
|
207 |
+
positive_all = positive_all + v + [s.replace("_", " ") for s in v]
|
208 |
+
positive_all = list_uniq(positive_all)
|
209 |
+
for k, v in (negative_prefix | negative_suffix).items():
|
210 |
+
negative_all = negative_all + v + [s.replace("_", " ") for s in v]
|
211 |
+
positive_all = list_uniq(positive_all)
|
212 |
+
|
213 |
+
|
214 |
+
def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
|
215 |
+
def flatten(src):
|
216 |
+
return [item for row in src for item in row]
|
217 |
+
prompts = to_list(prompt)
|
218 |
+
neg_prompts = to_list(neg_prompt)
|
219 |
+
prompts = list_sub(prompts, positive_all)
|
220 |
+
neg_prompts = list_sub(neg_prompts, negative_all)
|
221 |
+
last_empty_p = [""] if not prompts and type != "None" else []
|
222 |
+
last_empty_np = [""] if not neg_prompts and type != "None" else []
|
223 |
+
prefix_ps = flatten([positive_prefix.get(s, []) for s in pos_pre])
|
224 |
+
suffix_ps = flatten([positive_suffix.get(s, []) for s in pos_suf])
|
225 |
+
prefix_nps = flatten([negative_prefix.get(s, []) for s in neg_pre])
|
226 |
+
suffix_nps = flatten([negative_suffix.get(s, []) for s in neg_suf])
|
227 |
+
prompt = ", ".join(list_uniq(prefix_ps + prompts + suffix_ps) + last_empty_p)
|
228 |
+
neg_prompt = ", ".join(list_uniq(prefix_nps + neg_prompts + suffix_nps) + last_empty_np)
|
229 |
+
return prompt, neg_prompt
|
230 |
+
|
231 |
+
|
232 |
+
recom_prompt_type = {
|
233 |
+
"None": ([], [], [], []),
|
234 |
+
"Auto": ([], [], [], []),
|
235 |
+
"Common": ([], ["Common"], [], ["Common"]),
|
236 |
+
"Animagine": ([], ["Common", "Anime"], [], ["Common"]),
|
237 |
+
"Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]),
|
238 |
+
"Pony Anime": (["Pony", "Pony Anime"], ["Common", "Anime"], ["Pony", "Pony Anime"], ["Common", "Pony Anime"]),
|
239 |
+
"Pony Real": (["Pony"], ["Common"], ["Pony", "Pony Real"], ["Common", "Pony Real"]),
|
240 |
+
}
|
241 |
+
|
242 |
+
|
243 |
+
enable_auto_recom_prompt = False
|
244 |
+
def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
|
245 |
+
global enable_auto_recom_prompt
|
246 |
+
if type == "Auto": enable_auto_recom_prompt = True
|
247 |
+
else: enable_auto_recom_prompt = False
|
248 |
+
pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
|
249 |
+
return recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
|
250 |
+
|
251 |
+
|
252 |
+
def set_recom_prompt_preset(type: str = "None"):
|
253 |
+
pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
|
254 |
+
return pos_pre, pos_suf, neg_pre, neg_suf
|
255 |
+
|
256 |
+
|
257 |
+
def get_recom_prompt_type():
|
258 |
+
type = list(recom_prompt_type.keys())
|
259 |
+
type.remove("Auto")
|
260 |
+
return type
|
261 |
+
|
262 |
+
|
263 |
+
def get_positive_prefix():
|
264 |
+
return list(positive_prefix.keys())
|
265 |
+
|
266 |
+
|
267 |
+
def get_positive_suffix():
|
268 |
+
return list(positive_suffix.keys())
|
269 |
+
|
270 |
+
|
271 |
+
def get_negative_prefix():
|
272 |
+
return list(negative_prefix.keys())
|
273 |
+
|
274 |
+
|
275 |
+
def get_negative_suffix():
|
276 |
+
return list(negative_suffix.keys())
|
277 |
+
|
278 |
+
|
279 |
+
def get_tag_type(pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
|
280 |
+
tag_type = "danbooru"
|
281 |
+
words = pos_pre + pos_suf + neg_pre + neg_suf
|
282 |
+
for word in words:
|
283 |
+
if "Pony" in word:
|
284 |
+
tag_type = "e621"
|
285 |
+
break
|
286 |
+
return tag_type
|
287 |
|
288 |
|
289 |
def get_model_info_md(model_name: str):
|
|
|
295 |
return get_model_info_md(model_name)
|
296 |
|
297 |
|
298 |
+
def infer(prompt: str, neg_prompt: str, model_name: str):
|
299 |
from PIL import Image
|
300 |
import random
|
301 |
seed = ""
|
302 |
rand = random.randint(1, 500)
|
303 |
for i in range(rand):
|
304 |
seed += " "
|
|
|
305 |
caption = model_name.split("/")[-1]
|
306 |
try:
|
307 |
model = load_model(model_name)
|
308 |
if not model: return (Image.Image(), None)
|
309 |
+
image_path = model(prompt + seed)
|
310 |
+
image = Image.open(image_path).convert('RGBA')
|
311 |
except Exception as e:
|
312 |
print(e)
|
313 |
return (Image.Image(), None)
|
314 |
return (image, caption)
|
315 |
|
316 |
|
317 |
+
async def infer_multi(prompt: str, neg_prompt: str, results: list, image_num: float, model_name: str,
|
318 |
+
pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], progress=gr.Progress(track_tqdm=True)):
|
319 |
+
from tqdm.asyncio import tqdm_asyncio
|
320 |
image_num = int(image_num)
|
321 |
images = results if results else []
|
322 |
+
prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
|
323 |
+
tasks = [asyncio.to_thread(infer, prompt, neg_prompt, model_name) for i in range(image_num)]
|
324 |
+
#results = await asyncio.gather(*tasks, return_exceptions=True)
|
325 |
+
results = await tqdm_asyncio.gather(*tasks)
|
326 |
+
if not results: results = []
|
327 |
+
for result in results:
|
328 |
+
with lock:
|
329 |
+
if result and result[1]: images.append(result)
|
330 |
yield images
|
331 |
+
|
332 |
+
|
333 |
+
async def infer_multi_random(prompt: str, neg_prompt: str, results: list, image_num: float,
|
334 |
+
pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], progress=gr.Progress(track_tqdm=True)):
|
335 |
+
from tqdm.asyncio import tqdm_asyncio
|
336 |
+
import random
|
337 |
+
image_num = int(image_num)
|
338 |
+
images = results if results else []
|
339 |
+
random.seed()
|
340 |
+
model_names = random.choices(list(loaded_models.keys()), k = image_num)
|
341 |
+
prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
|
342 |
+
tasks = [asyncio.to_thread(infer, prompt, neg_prompt, model_name) for model_name in model_names]
|
343 |
+
#results = await asyncio.gather(*tasks, return_exceptions=True)
|
344 |
+
results = await tqdm_asyncio.gather(*tasks)
|
345 |
+
if not results: results = []
|
346 |
+
for result in results:
|
347 |
+
with lock:
|
348 |
+
if result and result[1]: images.append(result)
|
349 |
+
yield images
|
350 |
+
|
requirements.txt
CHANGED
@@ -1 +1,12 @@
|
|
1 |
-
huggingface_hub
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
huggingface_hub
|
2 |
+
torch
|
3 |
+
torchvision
|
4 |
+
accelerate
|
5 |
+
transformers
|
6 |
+
optimum[onnxruntime]
|
7 |
+
spaces
|
8 |
+
dartrs
|
9 |
+
httpx==0.13.3
|
10 |
+
httpcore
|
11 |
+
googletrans==4.0.0rc1
|
12 |
+
timm
|
tagger/character_series_dict.csv
ADDED
The diff for this file is too large to render.
See raw diff
|
|
tagger/danbooru_e621.csv
ADDED
The diff for this file is too large to render.
See raw diff
|
|
tagger/fl2sd3longcap.py
ADDED
@@ -0,0 +1,74 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import AutoProcessor, AutoModelForCausalLM
|
2 |
+
import spaces
|
3 |
+
import re
|
4 |
+
from PIL import Image
|
5 |
+
|
6 |
+
import subprocess
|
7 |
+
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
|
8 |
+
|
9 |
+
fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).eval()
|
10 |
+
fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
|
11 |
+
|
12 |
+
|
13 |
+
def fl_modify_caption(caption: str) -> str:
|
14 |
+
"""
|
15 |
+
Removes specific prefixes from captions if present, otherwise returns the original caption.
|
16 |
+
Args:
|
17 |
+
caption (str): A string containing a caption.
|
18 |
+
Returns:
|
19 |
+
str: The caption with the prefix removed if it was present, or the original caption.
|
20 |
+
"""
|
21 |
+
# Define the prefixes to remove
|
22 |
+
prefix_substrings = [
|
23 |
+
('captured from ', ''),
|
24 |
+
('captured at ', '')
|
25 |
+
]
|
26 |
+
|
27 |
+
# Create a regex pattern to match any of the prefixes
|
28 |
+
pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings])
|
29 |
+
replacers = {opening.lower(): replacer for opening, replacer in prefix_substrings}
|
30 |
+
|
31 |
+
# Function to replace matched prefix with its corresponding replacement
|
32 |
+
def replace_fn(match):
|
33 |
+
return replacers[match.group(0).lower()]
|
34 |
+
|
35 |
+
# Apply the regex to the caption
|
36 |
+
modified_caption = re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
|
37 |
+
|
38 |
+
# If the caption was modified, return the modified version; otherwise, return the original
|
39 |
+
return modified_caption if modified_caption != caption else caption
|
40 |
+
|
41 |
+
|
42 |
+
@spaces.GPU
|
43 |
+
def fl_run_example(image):
|
44 |
+
task_prompt = "<DESCRIPTION>"
|
45 |
+
prompt = task_prompt + "Describe this image in great detail."
|
46 |
+
|
47 |
+
# Ensure the image is in RGB mode
|
48 |
+
if image.mode != "RGB":
|
49 |
+
image = image.convert("RGB")
|
50 |
+
|
51 |
+
inputs = fl_processor(text=prompt, images=image, return_tensors="pt")
|
52 |
+
generated_ids = fl_model.generate(
|
53 |
+
input_ids=inputs["input_ids"],
|
54 |
+
pixel_values=inputs["pixel_values"],
|
55 |
+
max_new_tokens=1024,
|
56 |
+
num_beams=3
|
57 |
+
)
|
58 |
+
generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
|
59 |
+
parsed_answer = fl_processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
|
60 |
+
return fl_modify_caption(parsed_answer["<DESCRIPTION>"])
|
61 |
+
|
62 |
+
|
63 |
+
def predict_tags_fl2_sd3(image: Image.Image, input_tags: str, algo: list[str]):
|
64 |
+
def to_list(s):
|
65 |
+
return [x.strip() for x in s.split(",") if not s == ""]
|
66 |
+
|
67 |
+
def list_uniq(l):
|
68 |
+
return sorted(set(l), key=l.index)
|
69 |
+
|
70 |
+
if not "Use Florence-2-SD3-Long-Captioner" in algo:
|
71 |
+
return input_tags
|
72 |
+
tag_list = list_uniq(to_list(input_tags) + to_list(fl_run_example(image) + ", "))
|
73 |
+
tag_list.remove("")
|
74 |
+
return ", ".join(tag_list)
|
tagger/output.py
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from dataclasses import dataclass
|
2 |
+
|
3 |
+
|
4 |
+
@dataclass
|
5 |
+
class UpsamplingOutput:
|
6 |
+
upsampled_tags: str
|
7 |
+
|
8 |
+
copyright_tags: str
|
9 |
+
character_tags: str
|
10 |
+
general_tags: str
|
11 |
+
rating_tag: str
|
12 |
+
aspect_ratio_tag: str
|
13 |
+
length_tag: str
|
14 |
+
identity_tag: str
|
15 |
+
|
16 |
+
elapsed_time: float = 0.0
|
tagger/tag_group.csv
ADDED
The diff for this file is too large to render.
See raw diff
|
|
tagger/tagger.py
ADDED
@@ -0,0 +1,549 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from PIL import Image
|
2 |
+
import torch
|
3 |
+
import gradio as gr
|
4 |
+
import spaces
|
5 |
+
from transformers import (
|
6 |
+
AutoImageProcessor,
|
7 |
+
AutoModelForImageClassification,
|
8 |
+
)
|
9 |
+
from pathlib import Path
|
10 |
+
|
11 |
+
|
12 |
+
WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
|
13 |
+
WD_MODEL_NAME = WD_MODEL_NAMES[0]
|
14 |
+
|
15 |
+
wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
|
16 |
+
wd_model.to("cuda" if torch.cuda.is_available() else "cpu")
|
17 |
+
wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
|
18 |
+
|
19 |
+
|
20 |
+
def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
|
21 |
+
return (
|
22 |
+
[f"1{noun}"]
|
23 |
+
+ [f"{num}{noun}s" for num in range(minimum + 1, maximum + 1)]
|
24 |
+
+ [f"{maximum+1}+{noun}s"]
|
25 |
+
)
|
26 |
+
|
27 |
+
|
28 |
+
PEOPLE_TAGS = (
|
29 |
+
_people_tag("girl") + _people_tag("boy") + _people_tag("other") + ["no humans"]
|
30 |
+
)
|
31 |
+
|
32 |
+
|
33 |
+
RATING_MAP = {
|
34 |
+
"sfw": "safe",
|
35 |
+
"general": "safe",
|
36 |
+
"sensitive": "sensitive",
|
37 |
+
"questionable": "nsfw",
|
38 |
+
"explicit": "explicit, nsfw",
|
39 |
+
}
|
40 |
+
DANBOORU_TO_E621_RATING_MAP = {
|
41 |
+
"sfw": "rating_safe",
|
42 |
+
"general": "rating_safe",
|
43 |
+
"safe": "rating_safe",
|
44 |
+
"sensitive": "rating_safe",
|
45 |
+
"nsfw": "rating_explicit",
|
46 |
+
"explicit, nsfw": "rating_explicit",
|
47 |
+
"explicit": "rating_explicit",
|
48 |
+
"rating:safe": "rating_safe",
|
49 |
+
"rating:general": "rating_safe",
|
50 |
+
"rating:sensitive": "rating_safe",
|
51 |
+
"rating:questionable, nsfw": "rating_explicit",
|
52 |
+
"rating:explicit, nsfw": "rating_explicit",
|
53 |
+
}
|
54 |
+
|
55 |
+
|
56 |
+
# https://github.com/toriato/stable-diffusion-webui-wd14-tagger/blob/a9eacb1eff904552d3012babfa28b57e1d3e295c/tagger/ui.py#L368
|
57 |
+
kaomojis = [
|
58 |
+
"0_0",
|
59 |
+
"(o)_(o)",
|
60 |
+
"+_+",
|
61 |
+
"+_-",
|
62 |
+
"._.",
|
63 |
+
"<o>_<o>",
|
64 |
+
"<|>_<|>",
|
65 |
+
"=_=",
|
66 |
+
">_<",
|
67 |
+
"3_3",
|
68 |
+
"6_9",
|
69 |
+
">_o",
|
70 |
+
"@_@",
|
71 |
+
"^_^",
|
72 |
+
"o_o",
|
73 |
+
"u_u",
|
74 |
+
"x_x",
|
75 |
+
"|_|",
|
76 |
+
"||_||",
|
77 |
+
]
|
78 |
+
|
79 |
+
|
80 |
+
def replace_underline(x: str):
|
81 |
+
return x.strip().replace("_", " ") if x not in kaomojis else x.strip()
|
82 |
+
|
83 |
+
|
84 |
+
def to_list(s):
|
85 |
+
return [x.strip() for x in s.split(",") if not s == ""]
|
86 |
+
|
87 |
+
|
88 |
+
def list_sub(a, b):
|
89 |
+
return [e for e in a if e not in b]
|
90 |
+
|
91 |
+
|
92 |
+
def list_uniq(l):
|
93 |
+
return sorted(set(l), key=l.index)
|
94 |
+
|
95 |
+
|
96 |
+
def load_dict_from_csv(filename):
|
97 |
+
dict = {}
|
98 |
+
if not Path(filename).exists():
|
99 |
+
if Path('./tagger/', filename).exists(): filename = str(Path('./tagger/', filename))
|
100 |
+
else: return dict
|
101 |
+
try:
|
102 |
+
with open(filename, 'r', encoding="utf-8") as f:
|
103 |
+
lines = f.readlines()
|
104 |
+
except Exception:
|
105 |
+
print(f"Failed to open dictionary file: {filename}")
|
106 |
+
return dict
|
107 |
+
for line in lines:
|
108 |
+
parts = line.strip().split(',')
|
109 |
+
dict[parts[0]] = parts[1]
|
110 |
+
return dict
|
111 |
+
|
112 |
+
|
113 |
+
anime_series_dict = load_dict_from_csv('character_series_dict.csv')
|
114 |
+
|
115 |
+
|
116 |
+
def character_list_to_series_list(character_list):
|
117 |
+
output_series_tag = []
|
118 |
+
series_tag = ""
|
119 |
+
series_dict = anime_series_dict
|
120 |
+
for tag in character_list:
|
121 |
+
series_tag = series_dict.get(tag, "")
|
122 |
+
if tag.endswith(")"):
|
123 |
+
tags = tag.split("(")
|
124 |
+
character_tag = "(".join(tags[:-1])
|
125 |
+
if character_tag.endswith(" "):
|
126 |
+
character_tag = character_tag[:-1]
|
127 |
+
series_tag = tags[-1].replace(")", "")
|
128 |
+
|
129 |
+
if series_tag:
|
130 |
+
output_series_tag.append(series_tag)
|
131 |
+
|
132 |
+
return output_series_tag
|
133 |
+
|
134 |
+
|
135 |
+
def select_random_character(series: str, character: str):
|
136 |
+
from random import seed, randrange
|
137 |
+
seed()
|
138 |
+
character_list = list(anime_series_dict.keys())
|
139 |
+
character = character_list[randrange(len(character_list) - 1)]
|
140 |
+
series = anime_series_dict.get(character.split(",")[0].strip(), "")
|
141 |
+
return series, character
|
142 |
+
|
143 |
+
|
144 |
+
def danbooru_to_e621(dtag, e621_dict):
|
145 |
+
def d_to_e(match, e621_dict):
|
146 |
+
dtag = match.group(0)
|
147 |
+
etag = e621_dict.get(replace_underline(dtag), "")
|
148 |
+
if etag:
|
149 |
+
return etag
|
150 |
+
else:
|
151 |
+
return dtag
|
152 |
+
|
153 |
+
import re
|
154 |
+
tag = re.sub(r'[\w ]+', lambda wrapper: d_to_e(wrapper, e621_dict), dtag, 2)
|
155 |
+
return tag
|
156 |
+
|
157 |
+
|
158 |
+
danbooru_to_e621_dict = load_dict_from_csv('danbooru_e621.csv')
|
159 |
+
|
160 |
+
|
161 |
+
def convert_danbooru_to_e621_prompt(input_prompt: str = "", prompt_type: str = "danbooru"):
|
162 |
+
if prompt_type == "danbooru": return input_prompt
|
163 |
+
tags = input_prompt.split(",") if input_prompt else []
|
164 |
+
people_tags: list[str] = []
|
165 |
+
other_tags: list[str] = []
|
166 |
+
rating_tags: list[str] = []
|
167 |
+
|
168 |
+
e621_dict = danbooru_to_e621_dict
|
169 |
+
for tag in tags:
|
170 |
+
tag = replace_underline(tag)
|
171 |
+
tag = danbooru_to_e621(tag, e621_dict)
|
172 |
+
if tag in PEOPLE_TAGS:
|
173 |
+
people_tags.append(tag)
|
174 |
+
elif tag in DANBOORU_TO_E621_RATING_MAP.keys():
|
175 |
+
rating_tags.append(DANBOORU_TO_E621_RATING_MAP.get(tag.replace(" ",""), ""))
|
176 |
+
else:
|
177 |
+
other_tags.append(tag)
|
178 |
+
|
179 |
+
rating_tags = sorted(set(rating_tags), key=rating_tags.index)
|
180 |
+
rating_tags = [rating_tags[0]] if rating_tags else []
|
181 |
+
rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
|
182 |
+
|
183 |
+
output_prompt = ", ".join(people_tags + other_tags + rating_tags)
|
184 |
+
|
185 |
+
return output_prompt
|
186 |
+
|
187 |
+
|
188 |
+
def translate_prompt(prompt: str = ""):
|
189 |
+
def translate_to_english(prompt):
|
190 |
+
import httpcore
|
191 |
+
setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
|
192 |
+
from googletrans import Translator
|
193 |
+
translator = Translator()
|
194 |
+
try:
|
195 |
+
translated_prompt = translator.translate(prompt, src='auto', dest='en').text
|
196 |
+
return translated_prompt
|
197 |
+
except Exception as e:
|
198 |
+
print(e)
|
199 |
+
return prompt
|
200 |
+
|
201 |
+
def is_japanese(s):
|
202 |
+
import unicodedata
|
203 |
+
for ch in s:
|
204 |
+
name = unicodedata.name(ch, "")
|
205 |
+
if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name:
|
206 |
+
return True
|
207 |
+
return False
|
208 |
+
|
209 |
+
def to_list(s):
|
210 |
+
return [x.strip() for x in s.split(",")]
|
211 |
+
|
212 |
+
prompts = to_list(prompt)
|
213 |
+
outputs = []
|
214 |
+
for p in prompts:
|
215 |
+
p = translate_to_english(p) if is_japanese(p) else p
|
216 |
+
outputs.append(p)
|
217 |
+
|
218 |
+
return ", ".join(outputs)
|
219 |
+
|
220 |
+
|
221 |
+
def translate_prompt_to_ja(prompt: str = ""):
|
222 |
+
def translate_to_japanese(prompt):
|
223 |
+
import httpcore
|
224 |
+
setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
|
225 |
+
from googletrans import Translator
|
226 |
+
translator = Translator()
|
227 |
+
try:
|
228 |
+
translated_prompt = translator.translate(prompt, src='en', dest='ja').text
|
229 |
+
return translated_prompt
|
230 |
+
except Exception as e:
|
231 |
+
print(e)
|
232 |
+
return prompt
|
233 |
+
|
234 |
+
def is_japanese(s):
|
235 |
+
import unicodedata
|
236 |
+
for ch in s:
|
237 |
+
name = unicodedata.name(ch, "")
|
238 |
+
if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name:
|
239 |
+
return True
|
240 |
+
return False
|
241 |
+
|
242 |
+
def to_list(s):
|
243 |
+
return [x.strip() for x in s.split(",")]
|
244 |
+
|
245 |
+
prompts = to_list(prompt)
|
246 |
+
outputs = []
|
247 |
+
for p in prompts:
|
248 |
+
p = translate_to_japanese(p) if not is_japanese(p) else p
|
249 |
+
outputs.append(p)
|
250 |
+
|
251 |
+
return ", ".join(outputs)
|
252 |
+
|
253 |
+
|
254 |
+
def tags_to_ja(itag, dict):
|
255 |
+
def t_to_j(match, dict):
|
256 |
+
tag = match.group(0)
|
257 |
+
ja = dict.get(replace_underline(tag), "")
|
258 |
+
if ja:
|
259 |
+
return ja
|
260 |
+
else:
|
261 |
+
return tag
|
262 |
+
|
263 |
+
import re
|
264 |
+
tag = re.sub(r'[\w ]+', lambda wrapper: t_to_j(wrapper, dict), itag, 2)
|
265 |
+
|
266 |
+
return tag
|
267 |
+
|
268 |
+
|
269 |
+
def convert_tags_to_ja(input_prompt: str = ""):
|
270 |
+
tags = input_prompt.split(",") if input_prompt else []
|
271 |
+
out_tags = []
|
272 |
+
|
273 |
+
tags_to_ja_dict = load_dict_from_csv('all_tags_ja_ext.csv')
|
274 |
+
dict = tags_to_ja_dict
|
275 |
+
for tag in tags:
|
276 |
+
tag = replace_underline(tag)
|
277 |
+
tag = tags_to_ja(tag, dict)
|
278 |
+
out_tags.append(tag)
|
279 |
+
|
280 |
+
return ", ".join(out_tags)
|
281 |
+
|
282 |
+
|
283 |
+
enable_auto_recom_prompt = True
|
284 |
+
|
285 |
+
|
286 |
+
animagine_ps = to_list("masterpiece, best quality, very aesthetic, absurdres")
|
287 |
+
animagine_nps = to_list("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]")
|
288 |
+
pony_ps = to_list("score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres")
|
289 |
+
pony_nps = to_list("source_pony, score_6, score_5, score_4, busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends")
|
290 |
+
other_ps = to_list("anime artwork, anime style, studio anime, highly detailed, cinematic photo, 35mm photograph, film, bokeh, professional, 4k, highly detailed")
|
291 |
+
other_nps = to_list("photo, deformed, black and white, realism, disfigured, low contrast, drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly")
|
292 |
+
default_ps = to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres")
|
293 |
+
default_nps = to_list("score_6, score_5, score_4, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]")
|
294 |
+
def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
|
295 |
+
global enable_auto_recom_prompt
|
296 |
+
prompts = to_list(prompt)
|
297 |
+
neg_prompts = to_list(neg_prompt)
|
298 |
+
|
299 |
+
prompts = list_sub(prompts, animagine_ps + pony_ps)
|
300 |
+
neg_prompts = list_sub(neg_prompts, animagine_nps + pony_nps)
|
301 |
+
|
302 |
+
last_empty_p = [""] if not prompts and type != "None" else []
|
303 |
+
last_empty_np = [""] if not neg_prompts and type != "None" else []
|
304 |
+
|
305 |
+
if type == "Auto":
|
306 |
+
enable_auto_recom_prompt = True
|
307 |
+
else:
|
308 |
+
enable_auto_recom_prompt = False
|
309 |
+
if type == "Animagine":
|
310 |
+
prompts = prompts + animagine_ps
|
311 |
+
neg_prompts = neg_prompts + animagine_nps
|
312 |
+
elif type == "Pony":
|
313 |
+
prompts = prompts + pony_ps
|
314 |
+
neg_prompts = neg_prompts + pony_nps
|
315 |
+
|
316 |
+
prompt = ", ".join(list_uniq(prompts) + last_empty_p)
|
317 |
+
neg_prompt = ", ".join(list_uniq(neg_prompts) + last_empty_np)
|
318 |
+
|
319 |
+
return prompt, neg_prompt
|
320 |
+
|
321 |
+
|
322 |
+
def load_model_prompt_dict():
|
323 |
+
import json
|
324 |
+
dict = {}
|
325 |
+
path = 'model_dict.json' if Path('model_dict.json').exists() else './tagger/model_dict.json'
|
326 |
+
try:
|
327 |
+
with open('model_dict.json', encoding='utf-8') as f:
|
328 |
+
dict = json.load(f)
|
329 |
+
except Exception:
|
330 |
+
pass
|
331 |
+
return dict
|
332 |
+
|
333 |
+
|
334 |
+
model_prompt_dict = load_model_prompt_dict()
|
335 |
+
|
336 |
+
|
337 |
+
def insert_model_recom_prompt(prompt: str = "", neg_prompt: str = "", model_name: str = "None"):
|
338 |
+
if not model_name or not enable_auto_recom_prompt: return prompt, neg_prompt
|
339 |
+
prompts = to_list(prompt)
|
340 |
+
neg_prompts = to_list(neg_prompt)
|
341 |
+
prompts = list_sub(prompts, animagine_ps + pony_ps + other_ps)
|
342 |
+
neg_prompts = list_sub(neg_prompts, animagine_nps + pony_nps + other_nps)
|
343 |
+
last_empty_p = [""] if not prompts and type != "None" else []
|
344 |
+
last_empty_np = [""] if not neg_prompts and type != "None" else []
|
345 |
+
ps = []
|
346 |
+
nps = []
|
347 |
+
if model_name in model_prompt_dict.keys():
|
348 |
+
ps = to_list(model_prompt_dict[model_name]["prompt"])
|
349 |
+
nps = to_list(model_prompt_dict[model_name]["negative_prompt"])
|
350 |
+
else:
|
351 |
+
ps = default_ps
|
352 |
+
nps = default_nps
|
353 |
+
prompts = prompts + ps
|
354 |
+
neg_prompts = neg_prompts + nps
|
355 |
+
prompt = ", ".join(list_uniq(prompts) + last_empty_p)
|
356 |
+
neg_prompt = ", ".join(list_uniq(neg_prompts) + last_empty_np)
|
357 |
+
return prompt, neg_prompt
|
358 |
+
|
359 |
+
|
360 |
+
tag_group_dict = load_dict_from_csv('tag_group.csv')
|
361 |
+
|
362 |
+
|
363 |
+
def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"):
|
364 |
+
def is_dressed(tag):
|
365 |
+
import re
|
366 |
+
p = re.compile(r'dress|cloth|uniform|costume|vest|sweater|coat|shirt|jacket|blazer|apron|leotard|hood|sleeve|skirt|shorts|pant|loafer|ribbon|necktie|bow|collar|glove|sock|shoe|boots|wear|emblem')
|
367 |
+
return p.search(tag)
|
368 |
+
|
369 |
+
def is_background(tag):
|
370 |
+
import re
|
371 |
+
p = re.compile(r'background|outline|light|sky|build|day|screen|tree|city')
|
372 |
+
return p.search(tag)
|
373 |
+
|
374 |
+
un_tags = ['solo']
|
375 |
+
group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
|
376 |
+
keep_group_dict = {
|
377 |
+
"body": ['groups', 'body_parts'],
|
378 |
+
"dress": ['groups', 'body_parts', 'attire'],
|
379 |
+
"all": group_list,
|
380 |
+
}
|
381 |
+
|
382 |
+
def is_necessary(tag, keep_tags, group_dict):
|
383 |
+
if keep_tags == "all":
|
384 |
+
return True
|
385 |
+
elif tag in un_tags or group_dict.get(tag, "") in explicit_group:
|
386 |
+
return False
|
387 |
+
elif keep_tags == "body" and is_dressed(tag):
|
388 |
+
return False
|
389 |
+
elif is_background(tag):
|
390 |
+
return False
|
391 |
+
else:
|
392 |
+
return True
|
393 |
+
|
394 |
+
if keep_tags == "all": return input_prompt
|
395 |
+
keep_group = keep_group_dict.get(keep_tags, keep_group_dict["body"])
|
396 |
+
explicit_group = list(set(group_list) ^ set(keep_group))
|
397 |
+
|
398 |
+
tags = input_prompt.split(",") if input_prompt else []
|
399 |
+
people_tags: list[str] = []
|
400 |
+
other_tags: list[str] = []
|
401 |
+
|
402 |
+
group_dict = tag_group_dict
|
403 |
+
for tag in tags:
|
404 |
+
tag = replace_underline(tag)
|
405 |
+
if tag in PEOPLE_TAGS:
|
406 |
+
people_tags.append(tag)
|
407 |
+
elif is_necessary(tag, keep_tags, group_dict):
|
408 |
+
other_tags.append(tag)
|
409 |
+
|
410 |
+
output_prompt = ", ".join(people_tags + other_tags)
|
411 |
+
|
412 |
+
return output_prompt
|
413 |
+
|
414 |
+
|
415 |
+
def sort_taglist(tags: list[str]):
|
416 |
+
if not tags: return []
|
417 |
+
character_tags: list[str] = []
|
418 |
+
series_tags: list[str] = []
|
419 |
+
people_tags: list[str] = []
|
420 |
+
group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
|
421 |
+
group_tags = {}
|
422 |
+
other_tags: list[str] = []
|
423 |
+
rating_tags: list[str] = []
|
424 |
+
|
425 |
+
group_dict = tag_group_dict
|
426 |
+
group_set = set(group_dict.keys())
|
427 |
+
character_set = set(anime_series_dict.keys())
|
428 |
+
series_set = set(anime_series_dict.values())
|
429 |
+
rating_set = set(DANBOORU_TO_E621_RATING_MAP.keys()) | set(DANBOORU_TO_E621_RATING_MAP.values())
|
430 |
+
|
431 |
+
for tag in tags:
|
432 |
+
tag = replace_underline(tag)
|
433 |
+
if tag in PEOPLE_TAGS:
|
434 |
+
people_tags.append(tag)
|
435 |
+
elif tag in rating_set:
|
436 |
+
rating_tags.append(tag)
|
437 |
+
elif tag in group_set:
|
438 |
+
elem = group_dict[tag]
|
439 |
+
group_tags[elem] = group_tags[elem] + [tag] if elem in group_tags else [tag]
|
440 |
+
elif tag in character_set:
|
441 |
+
character_tags.append(tag)
|
442 |
+
elif tag in series_set:
|
443 |
+
series_tags.append(tag)
|
444 |
+
else:
|
445 |
+
other_tags.append(tag)
|
446 |
+
|
447 |
+
output_group_tags: list[str] = []
|
448 |
+
for k in group_list:
|
449 |
+
output_group_tags.extend(group_tags.get(k, []))
|
450 |
+
|
451 |
+
rating_tags = [rating_tags[0]] if rating_tags else []
|
452 |
+
rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
|
453 |
+
|
454 |
+
output_tags = character_tags + series_tags + people_tags + output_group_tags + other_tags + rating_tags
|
455 |
+
|
456 |
+
return output_tags
|
457 |
+
|
458 |
+
|
459 |
+
def sort_tags(tags: str):
|
460 |
+
if not tags: return ""
|
461 |
+
taglist: list[str] = []
|
462 |
+
for tag in tags.split(","):
|
463 |
+
taglist.append(tag.strip())
|
464 |
+
taglist = list(filter(lambda x: x != "", taglist))
|
465 |
+
return ", ".join(sort_taglist(taglist))
|
466 |
+
|
467 |
+
|
468 |
+
def postprocess_results(results: dict[str, float], general_threshold: float, character_threshold: float):
|
469 |
+
results = {
|
470 |
+
k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True)
|
471 |
+
}
|
472 |
+
|
473 |
+
rating = {}
|
474 |
+
character = {}
|
475 |
+
general = {}
|
476 |
+
|
477 |
+
for k, v in results.items():
|
478 |
+
if k.startswith("rating:"):
|
479 |
+
rating[k.replace("rating:", "")] = v
|
480 |
+
continue
|
481 |
+
elif k.startswith("character:"):
|
482 |
+
character[k.replace("character:", "")] = v
|
483 |
+
continue
|
484 |
+
|
485 |
+
general[k] = v
|
486 |
+
|
487 |
+
character = {k: v for k, v in character.items() if v >= character_threshold}
|
488 |
+
general = {k: v for k, v in general.items() if v >= general_threshold}
|
489 |
+
|
490 |
+
return rating, character, general
|
491 |
+
|
492 |
+
|
493 |
+
def gen_prompt(rating: list[str], character: list[str], general: list[str]):
|
494 |
+
people_tags: list[str] = []
|
495 |
+
other_tags: list[str] = []
|
496 |
+
rating_tag = RATING_MAP[rating[0]]
|
497 |
+
|
498 |
+
for tag in general:
|
499 |
+
if tag in PEOPLE_TAGS:
|
500 |
+
people_tags.append(tag)
|
501 |
+
else:
|
502 |
+
other_tags.append(tag)
|
503 |
+
|
504 |
+
all_tags = people_tags + other_tags
|
505 |
+
|
506 |
+
return ", ".join(all_tags)
|
507 |
+
|
508 |
+
|
509 |
+
@spaces.GPU()
|
510 |
+
def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8):
|
511 |
+
inputs = wd_processor.preprocess(image, return_tensors="pt")
|
512 |
+
|
513 |
+
outputs = wd_model(**inputs.to(wd_model.device, wd_model.dtype))
|
514 |
+
logits = torch.sigmoid(outputs.logits[0]) # take the first logits
|
515 |
+
|
516 |
+
# get probabilities
|
517 |
+
results = {
|
518 |
+
wd_model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
|
519 |
+
}
|
520 |
+
# rating, character, general
|
521 |
+
rating, character, general = postprocess_results(
|
522 |
+
results, general_threshold, character_threshold
|
523 |
+
)
|
524 |
+
prompt = gen_prompt(
|
525 |
+
list(rating.keys()), list(character.keys()), list(general.keys())
|
526 |
+
)
|
527 |
+
output_series_tag = ""
|
528 |
+
output_series_list = character_list_to_series_list(character.keys())
|
529 |
+
if output_series_list:
|
530 |
+
output_series_tag = output_series_list[0]
|
531 |
+
else:
|
532 |
+
output_series_tag = ""
|
533 |
+
return output_series_tag, ", ".join(character.keys()), prompt, gr.update(interactive=True)
|
534 |
+
|
535 |
+
|
536 |
+
def predict_tags_wd(image: Image.Image, input_tags: str, algo: list[str], general_threshold: float = 0.3,
|
537 |
+
character_threshold: float = 0.8, input_series: str = "", input_character: str = ""):
|
538 |
+
if not "Use WD Tagger" in algo and len(algo) != 0:
|
539 |
+
return input_series, input_character, input_tags, gr.update(interactive=True)
|
540 |
+
return predict_tags(image, general_threshold, character_threshold)
|
541 |
+
|
542 |
+
|
543 |
+
def compose_prompt_to_copy(character: str, series: str, general: str):
|
544 |
+
characters = character.split(",") if character else []
|
545 |
+
serieses = series.split(",") if series else []
|
546 |
+
generals = general.split(",") if general else []
|
547 |
+
tags = characters + serieses + generals
|
548 |
+
cprompt = ",".join(tags) if tags else ""
|
549 |
+
return cprompt
|
tagger/utils.py
ADDED
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
from dartrs.v2 import AspectRatioTag, LengthTag, RatingTag, IdentityTag
|
3 |
+
|
4 |
+
|
5 |
+
V2_ASPECT_RATIO_OPTIONS: list[AspectRatioTag] = [
|
6 |
+
"ultra_wide",
|
7 |
+
"wide",
|
8 |
+
"square",
|
9 |
+
"tall",
|
10 |
+
"ultra_tall",
|
11 |
+
]
|
12 |
+
V2_RATING_OPTIONS: list[RatingTag] = [
|
13 |
+
"sfw",
|
14 |
+
"general",
|
15 |
+
"sensitive",
|
16 |
+
"nsfw",
|
17 |
+
"questionable",
|
18 |
+
"explicit",
|
19 |
+
]
|
20 |
+
V2_LENGTH_OPTIONS: list[LengthTag] = [
|
21 |
+
"very_short",
|
22 |
+
"short",
|
23 |
+
"medium",
|
24 |
+
"long",
|
25 |
+
"very_long",
|
26 |
+
]
|
27 |
+
V2_IDENTITY_OPTIONS: list[IdentityTag] = [
|
28 |
+
"none",
|
29 |
+
"lax",
|
30 |
+
"strict",
|
31 |
+
]
|
32 |
+
|
33 |
+
|
34 |
+
# ref: https://qiita.com/tregu148/items/fccccbbc47d966dd2fc2
|
35 |
+
def gradio_copy_text(_text: None):
|
36 |
+
gr.Info("Copied!")
|
37 |
+
|
38 |
+
|
39 |
+
COPY_ACTION_JS = """\
|
40 |
+
(inputs, _outputs) => {
|
41 |
+
// inputs is the string value of the input_text
|
42 |
+
if (inputs.trim() !== "") {
|
43 |
+
navigator.clipboard.writeText(inputs);
|
44 |
+
}
|
45 |
+
}"""
|
tagger/v2.py
ADDED
@@ -0,0 +1,260 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import time
|
2 |
+
import torch
|
3 |
+
from typing import Callable
|
4 |
+
from pathlib import Path
|
5 |
+
|
6 |
+
from dartrs.v2 import (
|
7 |
+
V2Model,
|
8 |
+
MixtralModel,
|
9 |
+
MistralModel,
|
10 |
+
compose_prompt,
|
11 |
+
LengthTag,
|
12 |
+
AspectRatioTag,
|
13 |
+
RatingTag,
|
14 |
+
IdentityTag,
|
15 |
+
)
|
16 |
+
from dartrs.dartrs import DartTokenizer
|
17 |
+
from dartrs.utils import get_generation_config
|
18 |
+
|
19 |
+
|
20 |
+
import gradio as gr
|
21 |
+
from gradio.components import Component
|
22 |
+
|
23 |
+
|
24 |
+
try:
|
25 |
+
from output import UpsamplingOutput
|
26 |
+
except:
|
27 |
+
from .output import UpsamplingOutput
|
28 |
+
|
29 |
+
|
30 |
+
V2_ALL_MODELS = {
|
31 |
+
"dart-v2-moe-sft": {
|
32 |
+
"repo": "p1atdev/dart-v2-moe-sft",
|
33 |
+
"type": "sft",
|
34 |
+
"class": MixtralModel,
|
35 |
+
},
|
36 |
+
"dart-v2-sft": {
|
37 |
+
"repo": "p1atdev/dart-v2-sft",
|
38 |
+
"type": "sft",
|
39 |
+
"class": MistralModel,
|
40 |
+
},
|
41 |
+
}
|
42 |
+
|
43 |
+
|
44 |
+
def prepare_models(model_config: dict):
|
45 |
+
model_name = model_config["repo"]
|
46 |
+
tokenizer = DartTokenizer.from_pretrained(model_name)
|
47 |
+
model = model_config["class"].from_pretrained(model_name)
|
48 |
+
|
49 |
+
return {
|
50 |
+
"tokenizer": tokenizer,
|
51 |
+
"model": model,
|
52 |
+
}
|
53 |
+
|
54 |
+
|
55 |
+
def normalize_tags(tokenizer: DartTokenizer, tags: str):
|
56 |
+
"""Just remove unk tokens."""
|
57 |
+
return ", ".join([tag for tag in tokenizer.tokenize(tags) if tag != "<|unk|>"])
|
58 |
+
|
59 |
+
|
60 |
+
@torch.no_grad()
|
61 |
+
def generate_tags(
|
62 |
+
model: V2Model,
|
63 |
+
tokenizer: DartTokenizer,
|
64 |
+
prompt: str,
|
65 |
+
ban_token_ids: list[int],
|
66 |
+
):
|
67 |
+
output = model.generate(
|
68 |
+
get_generation_config(
|
69 |
+
prompt,
|
70 |
+
tokenizer=tokenizer,
|
71 |
+
temperature=1,
|
72 |
+
top_p=0.9,
|
73 |
+
top_k=100,
|
74 |
+
max_new_tokens=256,
|
75 |
+
ban_token_ids=ban_token_ids,
|
76 |
+
),
|
77 |
+
)
|
78 |
+
|
79 |
+
return output
|
80 |
+
|
81 |
+
|
82 |
+
def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
|
83 |
+
return (
|
84 |
+
[f"1{noun}"]
|
85 |
+
+ [f"{num}{noun}s" for num in range(minimum + 1, maximum + 1)]
|
86 |
+
+ [f"{maximum+1}+{noun}s"]
|
87 |
+
)
|
88 |
+
|
89 |
+
|
90 |
+
PEOPLE_TAGS = (
|
91 |
+
_people_tag("girl") + _people_tag("boy") + _people_tag("other") + ["no humans"]
|
92 |
+
)
|
93 |
+
|
94 |
+
|
95 |
+
def gen_prompt_text(output: UpsamplingOutput):
|
96 |
+
# separate people tags (e.g. 1girl)
|
97 |
+
people_tags = []
|
98 |
+
other_general_tags = []
|
99 |
+
|
100 |
+
for tag in output.general_tags.split(","):
|
101 |
+
tag = tag.strip()
|
102 |
+
if tag in PEOPLE_TAGS:
|
103 |
+
people_tags.append(tag)
|
104 |
+
else:
|
105 |
+
other_general_tags.append(tag)
|
106 |
+
|
107 |
+
return ", ".join(
|
108 |
+
[
|
109 |
+
part.strip()
|
110 |
+
for part in [
|
111 |
+
*people_tags,
|
112 |
+
output.character_tags,
|
113 |
+
output.copyright_tags,
|
114 |
+
*other_general_tags,
|
115 |
+
output.upsampled_tags,
|
116 |
+
output.rating_tag,
|
117 |
+
]
|
118 |
+
if part.strip() != ""
|
119 |
+
]
|
120 |
+
)
|
121 |
+
|
122 |
+
|
123 |
+
def elapsed_time_format(elapsed_time: float) -> str:
|
124 |
+
return f"Elapsed: {elapsed_time:.2f} seconds"
|
125 |
+
|
126 |
+
|
127 |
+
def parse_upsampling_output(
|
128 |
+
upsampler: Callable[..., UpsamplingOutput],
|
129 |
+
):
|
130 |
+
def _parse_upsampling_output(*args) -> tuple[str, str, dict]:
|
131 |
+
output = upsampler(*args)
|
132 |
+
|
133 |
+
return (
|
134 |
+
gen_prompt_text(output),
|
135 |
+
elapsed_time_format(output.elapsed_time),
|
136 |
+
gr.update(interactive=True),
|
137 |
+
gr.update(interactive=True),
|
138 |
+
)
|
139 |
+
|
140 |
+
return _parse_upsampling_output
|
141 |
+
|
142 |
+
|
143 |
+
class V2UI:
|
144 |
+
model_name: str | None = None
|
145 |
+
model: V2Model
|
146 |
+
tokenizer: DartTokenizer
|
147 |
+
|
148 |
+
input_components: list[Component] = []
|
149 |
+
generate_btn: gr.Button
|
150 |
+
|
151 |
+
def on_generate(
|
152 |
+
self,
|
153 |
+
model_name: str,
|
154 |
+
copyright_tags: str,
|
155 |
+
character_tags: str,
|
156 |
+
general_tags: str,
|
157 |
+
rating_tag: RatingTag,
|
158 |
+
aspect_ratio_tag: AspectRatioTag,
|
159 |
+
length_tag: LengthTag,
|
160 |
+
identity_tag: IdentityTag,
|
161 |
+
ban_tags: str,
|
162 |
+
*args,
|
163 |
+
) -> UpsamplingOutput:
|
164 |
+
if self.model_name is None or self.model_name != model_name:
|
165 |
+
models = prepare_models(V2_ALL_MODELS[model_name])
|
166 |
+
self.model = models["model"]
|
167 |
+
self.tokenizer = models["tokenizer"]
|
168 |
+
self.model_name = model_name
|
169 |
+
|
170 |
+
# normalize tags
|
171 |
+
# copyright_tags = normalize_tags(self.tokenizer, copyright_tags)
|
172 |
+
# character_tags = normalize_tags(self.tokenizer, character_tags)
|
173 |
+
# general_tags = normalize_tags(self.tokenizer, general_tags)
|
174 |
+
|
175 |
+
ban_token_ids = self.tokenizer.encode(ban_tags.strip())
|
176 |
+
|
177 |
+
prompt = compose_prompt(
|
178 |
+
prompt=general_tags,
|
179 |
+
copyright=copyright_tags,
|
180 |
+
character=character_tags,
|
181 |
+
rating=rating_tag,
|
182 |
+
aspect_ratio=aspect_ratio_tag,
|
183 |
+
length=length_tag,
|
184 |
+
identity=identity_tag,
|
185 |
+
)
|
186 |
+
|
187 |
+
start = time.time()
|
188 |
+
upsampled_tags = generate_tags(
|
189 |
+
self.model,
|
190 |
+
self.tokenizer,
|
191 |
+
prompt,
|
192 |
+
ban_token_ids,
|
193 |
+
)
|
194 |
+
elapsed_time = time.time() - start
|
195 |
+
|
196 |
+
return UpsamplingOutput(
|
197 |
+
upsampled_tags=upsampled_tags,
|
198 |
+
copyright_tags=copyright_tags,
|
199 |
+
character_tags=character_tags,
|
200 |
+
general_tags=general_tags,
|
201 |
+
rating_tag=rating_tag,
|
202 |
+
aspect_ratio_tag=aspect_ratio_tag,
|
203 |
+
length_tag=length_tag,
|
204 |
+
identity_tag=identity_tag,
|
205 |
+
elapsed_time=elapsed_time,
|
206 |
+
)
|
207 |
+
|
208 |
+
|
209 |
+
def parse_upsampling_output_simple(upsampler: UpsamplingOutput):
|
210 |
+
return gen_prompt_text(upsampler)
|
211 |
+
|
212 |
+
|
213 |
+
v2 = V2UI()
|
214 |
+
|
215 |
+
|
216 |
+
def v2_upsampling_prompt(model: str = "dart-v2-moe-sft", copyright: str = "", character: str = "",
|
217 |
+
general_tags: str = "", rating: str = "nsfw", aspect_ratio: str = "square",
|
218 |
+
length: str = "very_long", identity: str = "lax", ban_tags: str = "censored"):
|
219 |
+
raw_prompt = parse_upsampling_output_simple(v2.on_generate(model, copyright, character, general_tags,
|
220 |
+
rating, aspect_ratio, length, identity, ban_tags))
|
221 |
+
return raw_prompt
|
222 |
+
|
223 |
+
|
224 |
+
def load_dict_from_csv(filename):
|
225 |
+
dict = {}
|
226 |
+
if not Path(filename).exists():
|
227 |
+
if Path('./tagger/', filename).exists(): filename = str(Path('./tagger/', filename))
|
228 |
+
else: return dict
|
229 |
+
try:
|
230 |
+
with open(filename, 'r', encoding="utf-8") as f:
|
231 |
+
lines = f.readlines()
|
232 |
+
except Exception:
|
233 |
+
print(f"Failed to open dictionary file: {filename}")
|
234 |
+
return dict
|
235 |
+
for line in lines:
|
236 |
+
parts = line.strip().split(',')
|
237 |
+
dict[parts[0]] = parts[1]
|
238 |
+
return dict
|
239 |
+
|
240 |
+
|
241 |
+
anime_series_dict = load_dict_from_csv('character_series_dict.csv')
|
242 |
+
|
243 |
+
|
244 |
+
def select_random_character(series: str, character: str):
|
245 |
+
from random import seed, randrange
|
246 |
+
seed()
|
247 |
+
character_list = list(anime_series_dict.keys())
|
248 |
+
character = character_list[randrange(len(character_list) - 1)]
|
249 |
+
series = anime_series_dict.get(character.split(",")[0].strip(), "")
|
250 |
+
return series, character
|
251 |
+
|
252 |
+
|
253 |
+
def v2_random_prompt(general_tags: str = "", copyright: str = "", character: str = "", rating: str = "nsfw",
|
254 |
+
aspect_ratio: str = "square", length: str = "very_long", identity: str = "lax",
|
255 |
+
ban_tags: str = "censored", model: str = "dart-v2-moe-sft"):
|
256 |
+
if copyright == "" and character == "":
|
257 |
+
copyright, character = select_random_character("", "")
|
258 |
+
raw_prompt = v2_upsampling_prompt(model, copyright, character, general_tags, rating,
|
259 |
+
aspect_ratio, length, identity, ban_tags)
|
260 |
+
return raw_prompt, copyright, character
|