John6666 commited on
Commit
0d44125
1 Parent(s): b384075

Upload 13 files

Browse files
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 = find_model_list("Yntec", [], "", "likes", 30)
14
- load_models(models, 10)
15
 
16
 
17
- css = """"""
 
 
18
 
19
  with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
20
  with gr.Column():
21
- model_name = gr.Dropdown(label="Select Model", choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0], allow_custom_value=True)
22
- model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]))
23
- image_num = gr.Slider(label="Number of Images", minimum=1, maximum=8, value=1, step=1)
24
- recom_prompt = gr.Checkbox(label="Recommended Prompt", value=True)
25
- prompt = gr.Text(label="Prompt", lines=1, max_lines=8, placeholder="1girl, solo, ...")
26
- run_button = gr.Button("Generate Image")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 and Download")
 
 
 
 
 
 
 
 
 
 
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, model_name, recom_prompt, image_num, results],
 
 
 
 
 
 
 
 
 
 
 
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("Yntec", [], "", "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, task="text-to-image", pipeline_tag="text-to-image",
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] = gr.load(f'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(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.wait(tasks)
134
 
135
 
136
  def load_models(models: list, limit: int=5):
137
- loop = asyncio.get_event_loop()
138
  try:
139
  loop.run_until_complete(async_load_models(models, limit))
140
  except Exception as e:
141
  print(e)
142
  pass
143
- loop.close()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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, model_name: str, recom_prompt: bool, progress=gr.Progress(track_tqdm=True)):
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 + rprompt + seed)
168
- image = Image.open(image_path).convert('RGB')
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, model_name: str, recom_prompt: bool, image_num: float, results: list, progress=gr.Progress(track_tqdm=True)):
 
 
176
  image_num = int(image_num)
177
  images = results if results else []
178
- for i in range(image_num):
179
- images.append(infer(prompt, model_name, recom_prompt))
 
 
 
 
 
 
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