John6666 commited on
Commit
5f8d674
Β·
verified Β·
1 Parent(s): e23e556

Upload 6 files

Browse files
Files changed (5) hide show
  1. app.py +48 -1
  2. ctag.py +3 -3
  3. hft2i.py +551 -0
  4. model.py +44 -0
  5. requirements.txt +3 -1
app.py CHANGED
@@ -1,10 +1,17 @@
1
  import gradio as gr
2
  from ctag import MODELS, DEFAULT_DF, search_char_dict, on_select_df
 
 
 
 
3
 
4
  CSS = """
5
  .title { font-size: 3em; align-items: center; text-align: center; }
6
  .info { align-items: center; text-align: center; }
7
  img[src*="#center"] { display: block; margin: auto; }
 
 
 
8
  """
9
 
10
  with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=CSS) as app:
@@ -19,11 +26,51 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=CSS) as app:
19
  with gr.Group():
20
  with gr.Row(equal_height=True):
21
  search_tag = gr.Textbox(label="Output tag", value="", show_copy_button=True, interactive=False)
 
22
  search_md = gr.Markdown("<br><br><br>", elem_classes="info")
23
  search_output = gr.Dataframe(label="Select character", value=DEFAULT_DF, type="pandas", wrap=True, interactive=False)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
  gr.on(triggers=[search_input.change, search_model.change], fn=search_char_dict,
26
  inputs=[search_input, search_model], outputs=[search_output], trigger_mode="always_last")
27
- search_output.select(on_select_df, [search_output, search_detail], [search_tag, search_md])
 
 
 
 
 
 
 
 
 
 
 
28
 
29
  app.launch(ssr_mode=False)
 
1
  import gradio as gr
2
  from ctag import MODELS, DEFAULT_DF, search_char_dict, on_select_df
3
+ from hft2i import (gen_image, save_gallery, get_models, get_def_model, change_model, warm_model, get_model_info_md, get_recom_prompt_mode, update_prompt)
4
+
5
+ MAX_IMAGES = 6
6
+ MAX_SEED = 2**32-1
7
 
8
  CSS = """
9
  .title { font-size: 3em; align-items: center; text-align: center; }
10
  .info { align-items: center; text-align: center; }
11
  img[src*="#center"] { display: block; margin: auto; }
12
+ .model_info { text-align: center; }
13
+ .output { width=112px; height=112px; max_width=112px; max_height=112px; !important; }
14
+ .gallery { min_width=512px; min_height=512px; max_height=1024px; !important; }
15
  """
16
 
17
  with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=CSS) as app:
 
26
  with gr.Group():
27
  with gr.Row(equal_height=True):
28
  search_tag = gr.Textbox(label="Output tag", value="", show_copy_button=True, interactive=False)
29
+ search_tag_model = gr.Textbox(label="Model", value="", visible=False)
30
  search_md = gr.Markdown("<br><br><br>", elem_classes="info")
31
  search_output = gr.Dataframe(label="Select character", value=DEFAULT_DF, type="pandas", wrap=True, interactive=False)
32
+ with gr.Group():
33
+ prompt = gr.Text(label="Prompt", lines=2, max_lines=8, placeholder="1girl, solo, ...", show_copy_button=True)
34
+ with gr.Accordion("Advanced options", open=False):
35
+ nprompt = gr.Text(label="Negative Prompt", lines=1, max_lines=8, placeholder="")
36
+ with gr.Row():
37
+ width = gr.Slider(label="Width", info="If 0, default value is used.", maximum=2048, step=32, value=0)
38
+ height = gr.Slider(label="Height", info="If 0, default value is used.", maximum=2048, step=32, value=0)
39
+ steps = gr.Slider(label="Number of inference steps", info="If 0, default value is used.", maximum=100, step=1, value=0)
40
+ cfg = gr.Slider(label="Guidance scale", info="If 0, default value is used.", maximum=30.0, step=0.1, value=0)
41
+ seed = gr.Slider(label="Seed", info="If -1, use Randomized Seed.", minimum=-1, maximum=MAX_SEED, step=1, value=-1)
42
+ with gr.Row():
43
+ model_name = [None] * MAX_IMAGES
44
+ model_info = [None] * MAX_IMAGES
45
+ for i in range(MAX_IMAGES):
46
+ with gr.Column():
47
+ model_name[i] = gr.Dropdown(label=f"Select Model {int(i) + 1}", choices=get_models(), value=get_def_model(i), allow_custom_value=True)
48
+ model_info[i] = gr.Markdown(value=get_model_info_md(get_def_model(i)), elem_classes="model_info")
49
+ recom_prompt_mode = gr.Radio(label="Recommened prompt", choices=get_recom_prompt_mode(), value="Common")
50
+ gen_button = gr.Button("Generate Image", variant="primary")
51
+ with gr.Group():
52
+ with gr.Row():
53
+ output_images = [gr.Image(label='', elem_classes="output", type="filepath", format="png",
54
+ show_download_button=True, show_share_button=False,
55
+ interactive=False, min_width=80, visible=True, width=112, height=112) for _ in range(MAX_IMAGES)]
56
+ gallery = gr.Gallery(label="Gallery", elem_classes="gallery", interactive=False, show_download_button=True, show_share_button=False,
57
+ container=True, format="png", object_fit="cover", columns=1, rows=1)
58
+ image_files = gr.Files(label="Download", interactive=False)
59
+ clear_image_button = gr.Button("Clear Gallery / Download πŸ—‘οΈ", variant="secondary")
60
 
61
  gr.on(triggers=[search_input.change, search_model.change], fn=search_char_dict,
62
  inputs=[search_input, search_model], outputs=[search_output], trigger_mode="always_last")
63
+ search_output.select(on_select_df, [search_output, search_detail], [search_tag, search_tag_model, search_md], queue=False, show_api=False)
64
+ search_tag.change(update_prompt, [search_tag, search_tag_model], [prompt, model_name[0], recom_prompt_mode], queue=False, show_api=False)
65
+ for i, o in enumerate(output_images):
66
+ img_i = gr.Number(i, visible=False)
67
+ model_name[i].change(change_model, [model_name[i]], [model_info[i]], queue=False, show_api=False)\
68
+ .success(warm_model, [model_name[i]], None, queue=False, show_api=False)\
69
+ .then(lambda x: gr.update(label=x.split("/")[-1]), [model_name[i]], [o], queue=False, show_api=False)
70
+ gen_event = gr.on(triggers=[gen_button.click, prompt.submit],
71
+ fn=lambda m, t1, t2, n1, n2, n3, n4, n5, t3: gen_image(m, t1, t2, n1, n2, n3, n4, n5, t3),
72
+ inputs=[model_name[i], prompt, nprompt, height, width, steps, cfg, seed, recom_prompt_mode],
73
+ outputs=[o], queue=False, show_api=False)
74
+ o.change(save_gallery, [o, gallery], [gallery, image_files], show_api=False)
75
 
76
  app.launch(ssr_mode=False)
ctag.py CHANGED
@@ -136,12 +136,12 @@ def search_char_dict(q: str, models: list[str], progress=gr.Progress(track_tqdm=
136
  def on_select_df(df: pd.DataFrame, is_detail: bool, evt: gr.SelectData):
137
  d = tag_dict.get(evt.value, None)
138
  if d is None: return ""
139
- #print(d)
140
  if is_detail: info = find_char_info(d["name"])
141
  else: info = None
142
- #print(info)
143
  if info is not None:
144
  md = f'## [{info["name"]}]({info["wiki"]}) / [{info["series"]}]({info["wiki"]}) / {info["gender"]}\n![{info["name"]}]({info["image"]}#center)\n[{info["desc"]}]({info["wiki"]})'
145
  else: md = f'## {d["name"]} / {d["series"]}' if d["series"] else f'## {d["name"]}'
146
  md += f'\n<br>Tag is for {model_dict[evt.value]}.'
147
- return d["tag"], md
 
136
  def on_select_df(df: pd.DataFrame, is_detail: bool, evt: gr.SelectData):
137
  d = tag_dict.get(evt.value, None)
138
  if d is None: return ""
139
+ #print(d) #
140
  if is_detail: info = find_char_info(d["name"])
141
  else: info = None
142
+ #print(info) #
143
  if info is not None:
144
  md = f'## [{info["name"]}]({info["wiki"]}) / [{info["series"]}]({info["wiki"]}) / {info["gender"]}\n![{info["name"]}]({info["image"]}#center)\n[{info["desc"]}]({info["wiki"]})'
145
  else: md = f'## {d["name"]} / {d["series"]}' if d["series"] else f'## {d["name"]}'
146
  md += f'\n<br>Tag is for {model_dict[evt.value]}.'
147
+ return d["tag"], model_dict[evt.value], md
hft2i.py ADDED
@@ -0,0 +1,551 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import asyncio
3
+ from threading import RLock
4
+ from pathlib import Path
5
+ from huggingface_hub import InferenceClient
6
+ import os
7
+
8
+
9
+ HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None # If private or gated models aren't used, ENV setting is unnecessary.
10
+ server_timeout = 600
11
+ inference_timeout = 600
12
+
13
+
14
+ lock = RLock()
15
+ loaded_models = {}
16
+ model_info_dict = {}
17
+
18
+
19
+ def to_list(s):
20
+ return [x.strip() for x in s.split(",")]
21
+
22
+
23
+ def list_sub(a, b):
24
+ return [e for e in a if e not in b]
25
+
26
+
27
+ def list_uniq(l):
28
+ return sorted(set(l), key=l.index)
29
+
30
+
31
+ def is_repo_name(s):
32
+ import re
33
+ return re.fullmatch(r'^[^/]+?/[^/]+?$', s)
34
+
35
+
36
+ def get_status(model_name: str):
37
+ from huggingface_hub import InferenceClient
38
+ client = InferenceClient(token=HF_TOKEN, timeout=10)
39
+ return client.get_model_status(model_name)
40
+
41
+
42
+ def is_loadable(model_name: str, force_gpu: bool = False):
43
+ try:
44
+ status = get_status(model_name)
45
+ except Exception as e:
46
+ print(e)
47
+ print(f"Couldn't load {model_name}.")
48
+ return False
49
+ gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()
50
+ if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):
51
+ print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")
52
+ return status is not None and status.state in ["Loadable", "Loaded"] and (not force_gpu or gpu_state)
53
+
54
+
55
+ def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False, public=False):
56
+ from huggingface_hub import HfApi
57
+ api = HfApi(token=HF_TOKEN)
58
+ default_tags = ["diffusers"]
59
+ if not sort: sort = "last_modified"
60
+ limit = limit * 20 if check_status and force_gpu else limit * 5
61
+ models = []
62
+ try:
63
+ model_infos = api.list_models(author=author, #task="text-to-image",
64
+ tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)
65
+ except Exception as e:
66
+ print(f"Error: Failed to list models.")
67
+ print(e)
68
+ return models
69
+ for model in model_infos:
70
+ if not model.private and not model.gated or (HF_TOKEN is not None and not public):
71
+ loadable = is_loadable(model.id, force_gpu) if check_status else True
72
+ if not_tag and not_tag in model.tags or not loadable or "not-for-all-audiences" in model.tags: continue
73
+ models.append(model.id)
74
+ if len(models) == limit: break
75
+ return models
76
+
77
+
78
+ def get_t2i_model_info_dict(repo_id: str):
79
+ from huggingface_hub import HfApi
80
+ api = HfApi(token=HF_TOKEN)
81
+ info = {"md": "None"}
82
+ try:
83
+ if not is_repo_name(repo_id) or not api.repo_exists(repo_id=repo_id): return info
84
+ model = api.model_info(repo_id=repo_id, token=HF_TOKEN)
85
+ except Exception as e:
86
+ print(f"Error: Failed to get {repo_id}'s info.")
87
+ print(e)
88
+ return info
89
+ if model.private or model.gated and HF_TOKEN is None: return info
90
+ try:
91
+ tags = model.tags
92
+ except Exception as e:
93
+ print(e)
94
+ return info
95
+ if not 'diffusers' in model.tags: return info
96
+ if 'diffusers:FluxPipeline' in tags: info["ver"] = "FLUX.1"
97
+ elif 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
98
+ elif 'diffusers:StableDiffusionPipeline' in tags: info["ver"] = "SD1.5"
99
+ elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3"
100
+ else: info["ver"] = "Other"
101
+ info["url"] = f"https://huggingface.co/{repo_id}/"
102
+ info["tags"] = model.card_data.tags if model.card_data and model.card_data.tags else []
103
+ info["downloads"] = model.downloads
104
+ info["likes"] = model.likes
105
+ info["last_modified"] = model.last_modified.strftime("lastmod: %Y-%m-%d")
106
+ un_tags = ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl']
107
+ descs = [info["ver"]] + list_sub(info["tags"], un_tags) + [f'DLs: {info["downloads"]}'] + [f'πŸ’•: {info["likes"]}'] + [info["last_modified"]]
108
+ info["md"] = f'{", ".join(descs)} [Model Repo]({info["url"]})'
109
+ return info
110
+
111
+
112
+ def rename_image(image_path: str | None, model_name: str, save_path: str | None = None):
113
+ import shutil
114
+ from datetime import datetime, timezone, timedelta
115
+ if image_path is None: return None
116
+ dt_now = datetime.now(timezone(timedelta(hours=9)))
117
+ filename = f"{model_name.split('/')[-1]}_{dt_now.strftime('%Y%m%d_%H%M%S')}.png"
118
+ try:
119
+ if Path(image_path).exists():
120
+ png_path = "image.png"
121
+ if str(Path(image_path).resolve()) != str(Path(png_path).resolve()): shutil.copy(image_path, png_path)
122
+ if save_path is not None:
123
+ new_path = str(Path(png_path).resolve().rename(Path(save_path).resolve()))
124
+ else:
125
+ new_path = str(Path(png_path).resolve().rename(Path(filename).resolve()))
126
+ return new_path
127
+ else:
128
+ return None
129
+ except Exception as e:
130
+ print(e)
131
+ return None
132
+
133
+
134
+ def save_gallery(image_path: str | None, images: list[tuple] | None):
135
+ if images is None: images = []
136
+ files = [i[0] for i in images]
137
+ if image_path is None: return images, files
138
+ files.insert(0, str(image_path))
139
+ images.insert(0, (str(image_path), Path(image_path).stem))
140
+ return images, files
141
+
142
+
143
+ # https://github.com/gradio-app/gradio/blob/main/gradio/external.py
144
+ # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
145
+ from typing import Literal
146
+ def load_from_model(model_name: str, hf_token: str | Literal[False] | None = None):
147
+ import httpx
148
+ import huggingface_hub
149
+ from gradio.exceptions import ModelNotFoundError, TooManyRequestsError
150
+ model_url = f"https://huggingface.co/{model_name}"
151
+ api_url = f"https://api-inference.huggingface.co/models/{model_name}"
152
+ print(f"Fetching model from: {model_url}")
153
+
154
+ headers = ({} if hf_token in [False, None] else {"Authorization": f"Bearer {hf_token}"})
155
+ response = httpx.request("GET", api_url, headers=headers)
156
+ if response.status_code != 200:
157
+ raise ModelNotFoundError(
158
+ 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."
159
+ )
160
+ p = response.json().get("pipeline_tag")
161
+ if p != "text-to-image": raise ModelNotFoundError(f"This model isn't for text-to-image or unsupported: {model_name}.")
162
+ headers["X-Wait-For-Model"] = "true"
163
+ kwargs = {}
164
+ if hf_token is not None: kwargs["token"] = hf_token
165
+ client = huggingface_hub.InferenceClient(model=model_name, headers=headers, timeout=server_timeout, **kwargs)
166
+ inputs = gr.components.Textbox(label="Input")
167
+ outputs = gr.components.Image(label="Output")
168
+ fn = client.text_to_image
169
+
170
+ def query_huggingface_inference_endpoints(*data, **kwargs):
171
+ try:
172
+ data = fn(*data, **kwargs) # type: ignore
173
+ except huggingface_hub.utils.HfHubHTTPError as e:
174
+ print(e)
175
+ if "429" in str(e): raise TooManyRequestsError() from e
176
+ except Exception as e:
177
+ print(e)
178
+ raise Exception() from e
179
+ return data
180
+
181
+ interface_info = {
182
+ "fn": query_huggingface_inference_endpoints,
183
+ "inputs": inputs,
184
+ "outputs": outputs,
185
+ "title": model_name,
186
+ }
187
+ return gr.Interface(**interface_info)
188
+
189
+
190
+ def load_model(model_name: str):
191
+ global loaded_models
192
+ global model_info_dict
193
+ if model_name in loaded_models.keys(): return loaded_models[model_name]
194
+ try:
195
+ loaded_models[model_name] = load_from_model(model_name, hf_token=HF_TOKEN)
196
+ print(f"Loaded: {model_name}")
197
+ except Exception as e:
198
+ if model_name in loaded_models.keys(): del loaded_models[model_name]
199
+ print(f"Failed to load: {model_name}")
200
+ print(e)
201
+ return None
202
+ try:
203
+ model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
204
+ print(f"Assigned: {model_name}")
205
+ except Exception as e:
206
+ if model_name in model_info_dict.keys(): del model_info_dict[model_name]
207
+ print(f"Failed to assigned: {model_name}")
208
+ print(e)
209
+ return loaded_models[model_name]
210
+
211
+
212
+ def load_model_api(model_name: str):
213
+ global loaded_models
214
+ global model_info_dict
215
+ if model_name in loaded_models.keys(): return loaded_models[model_name]
216
+ try:
217
+ client = InferenceClient(timeout=5)
218
+ status = client.get_model_status(model_name, token=HF_TOKEN)
219
+ if status is None or status.framework != "diffusers" or status.state not in ["Loadable", "Loaded"]:
220
+ print(f"Failed to load by API: {model_name}")
221
+ return None
222
+ else:
223
+ kwargs = {}
224
+ if HF_TOKEN is not None: kwargs["token"] = HF_TOKEN
225
+ loaded_models[model_name] = InferenceClient(model_name, timeout=server_timeout, **kwargs)
226
+ print(f"Loaded by API: {model_name}")
227
+ except Exception as e:
228
+ if model_name in loaded_models.keys(): del loaded_models[model_name]
229
+ print(f"Failed to load by API: {model_name}")
230
+ print(e)
231
+ return None
232
+ try:
233
+ model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
234
+ print(f"Assigned by API: {model_name}")
235
+ except Exception as e:
236
+ if model_name in model_info_dict.keys(): del model_info_dict[model_name]
237
+ print(f"Failed to assigned by API: {model_name}")
238
+ print(e)
239
+ return loaded_models[model_name]
240
+
241
+
242
+ def load_models(models: list):
243
+ for model in models:
244
+ load_model(model)
245
+
246
+
247
+ positive_prefix = {
248
+ "Pony": to_list("score_9, score_8_up, score_7_up"),
249
+ "Pony Anime": to_list("source_anime, anime, score_9, score_8_up, score_7_up"),
250
+ }
251
+ positive_suffix = {
252
+ "Common": to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres"),
253
+ "Anime": to_list("anime artwork, anime style, studio anime, highly detailed"),
254
+ }
255
+ negative_prefix = {
256
+ "Pony": to_list("score_6, score_5, score_4"),
257
+ "Pony Anime": to_list("score_6, score_5, score_4, source_pony, source_furry, source_cartoon"),
258
+ "Pony Real": to_list("score_6, score_5, score_4, source_anime, source_pony, source_furry, source_cartoon"),
259
+ }
260
+ negative_suffix = {
261
+ "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]"),
262
+ "Pony Anime": to_list("busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends"),
263
+ "Pony Real": to_list("ugly, airbrushed, simple background, cgi, cartoon, anime"),
264
+ }
265
+ positive_all = negative_all = []
266
+ for k, v in (positive_prefix | positive_suffix).items():
267
+ positive_all = positive_all + v + [s.replace("_", " ") for s in v]
268
+ positive_all = list_uniq(positive_all)
269
+ for k, v in (negative_prefix | negative_suffix).items():
270
+ negative_all = negative_all + v + [s.replace("_", " ") for s in v]
271
+ positive_all = list_uniq(positive_all)
272
+
273
+
274
+ def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
275
+ def flatten(src):
276
+ return [item for row in src for item in row]
277
+ prompts = to_list(prompt) if prompt else []
278
+ neg_prompts = to_list(neg_prompt) if neg_prompt else []
279
+ prompts = list_sub(prompts, positive_all)
280
+ neg_prompts = list_sub(neg_prompts, negative_all)
281
+ last_empty_p = [""] if not prompts and type != "None" else []
282
+ last_empty_np = [""] if not neg_prompts and type != "None" else []
283
+ prefix_ps = flatten([positive_prefix.get(s, []) for s in pos_pre])
284
+ suffix_ps = flatten([positive_suffix.get(s, []) for s in pos_suf])
285
+ prefix_nps = flatten([negative_prefix.get(s, []) for s in neg_pre])
286
+ suffix_nps = flatten([negative_suffix.get(s, []) for s in neg_suf])
287
+ prompt = ", ".join(list_uniq(prefix_ps + prompts + suffix_ps) + last_empty_p)
288
+ neg_prompt = ", ".join(list_uniq(prefix_nps + neg_prompts + suffix_nps) + last_empty_np)
289
+ return prompt, neg_prompt
290
+
291
+
292
+ recom_prompt_type = {
293
+ "None": ([], [], [], []),
294
+ "Common": ([], ["Common"], [], ["Common"]),
295
+ "Animagine": ([], ["Common", "Anime"], [], ["Common"]),
296
+ "Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]),
297
+ "Pony Anime": (["Pony", "Pony Anime"], ["Common", "Anime"], ["Pony", "Pony Anime"], ["Common", "Pony Anime"]),
298
+ "Pony Real": (["Pony"], ["Common"], ["Pony", "Pony Real"], ["Common", "Pony Real"]),
299
+ }
300
+
301
+
302
+ def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
303
+ pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
304
+ return recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
305
+
306
+
307
+ def set_recom_prompt_preset(type: str = "None"):
308
+ pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
309
+ return pos_pre, pos_suf, neg_pre, neg_suf
310
+
311
+
312
+ def get_recom_prompt_type():
313
+ return list(recom_prompt_type.keys())
314
+
315
+
316
+ def get_positive_prefix():
317
+ return list(positive_prefix.keys())
318
+
319
+
320
+ def get_positive_suffix():
321
+ return list(positive_suffix.keys())
322
+
323
+
324
+ def get_negative_prefix():
325
+ return list(negative_prefix.keys())
326
+
327
+
328
+ def get_negative_suffix():
329
+ return list(negative_suffix.keys())
330
+
331
+
332
+ def get_tag_type(pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
333
+ tag_type = "danbooru"
334
+ words = pos_pre + pos_suf + neg_pre + neg_suf
335
+ for word in words:
336
+ if "Pony" in word:
337
+ tag_type = "e621"
338
+ break
339
+ return tag_type
340
+
341
+
342
+ def get_model_info_md(model_name: str):
343
+ if model_name in model_info_dict.keys(): return model_info_dict[model_name].get("md", "")
344
+
345
+
346
+ def change_model(model_name: str):
347
+ load_model_api(model_name)
348
+ return get_model_info_md(model_name)
349
+
350
+
351
+ def warm_model(model_name: str):
352
+ model = load_model_api(model_name)
353
+ if model:
354
+ try:
355
+ print(f"Warming model: {model_name}")
356
+ infer_body(model, model_name, " ")
357
+ except Exception as e:
358
+ print(e)
359
+
360
+
361
+ # https://huggingface.co/docs/api-inference/detailed_parameters
362
+ # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
363
+ def infer_body(client: InferenceClient | gr.Interface | object, model_str: str, prompt: str, neg_prompt: str = "",
364
+ height: int = 0, width: int = 0, steps: int = 0, cfg: int = 0, seed: int = -1):
365
+ png_path = "image.png"
366
+ kwargs = {}
367
+ if height > 0: kwargs["height"] = height
368
+ if width > 0: kwargs["width"] = width
369
+ if steps > 0: kwargs["num_inference_steps"] = steps
370
+ if cfg > 0: cfg = kwargs["guidance_scale"] = cfg
371
+ if seed == -1: kwargs["seed"] = randomize_seed()
372
+ else: kwargs["seed"] = seed
373
+ if HF_TOKEN is not None: kwargs["token"] = HF_TOKEN
374
+ try:
375
+ if isinstance(client, InferenceClient):
376
+ image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
377
+ elif isinstance(client, gr.Interface):
378
+ image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
379
+ else: return None
380
+ if isinstance(image, tuple): return None
381
+ return save_image(image, png_path, model_str, prompt, neg_prompt, height, width, steps, cfg, seed)
382
+ except Exception as e:
383
+ print(e)
384
+ raise Exception() from e
385
+
386
+
387
+ async def infer(model_name: str, prompt: str, neg_prompt: str ="", height: int = 0, width: int = 0,
388
+ steps: int = 0, cfg: int = 0, seed: int = -1,
389
+ save_path: str | None = None, timeout: float = inference_timeout):
390
+ model = load_model(model_name)
391
+ if not model: return None
392
+ task = asyncio.create_task(asyncio.to_thread(infer_body, model, model_name, prompt, neg_prompt,
393
+ height, width, steps, cfg, seed))
394
+ await asyncio.sleep(0)
395
+ try:
396
+ result = await asyncio.wait_for(task, timeout=timeout)
397
+ except asyncio.TimeoutError as e:
398
+ print(e)
399
+ print(f"Task timed out: {model_name}")
400
+ if not task.done(): task.cancel()
401
+ result = None
402
+ raise Exception(f"Task timed out: {model_name}") from e
403
+ except Exception as e:
404
+ print(e)
405
+ if not task.done(): task.cancel()
406
+ result = None
407
+ raise Exception() from e
408
+ if task.done() and result is not None:
409
+ with lock:
410
+ image = rename_image(result, model_name, save_path)
411
+ return image
412
+ return None
413
+
414
+
415
+ # https://github.com/aio-libs/pytest-aiohttp/issues/8 # also AsyncInferenceClient is buggy.
416
+ def infer_fn(model_name: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
417
+ steps: int = 0, cfg: int = 0, seed: int = -1,
418
+ pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
419
+ if model_name in ["NA", ""]: return gr.update()
420
+ try:
421
+ loop = asyncio.get_running_loop()
422
+ except Exception:
423
+ loop = asyncio.new_event_loop()
424
+ try:
425
+ prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
426
+ result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
427
+ steps, cfg, seed, save_path, inference_timeout))
428
+ except (Exception, asyncio.CancelledError) as e:
429
+ print(e)
430
+ print(f"Task aborted: {model_name}, Error: {e}")
431
+ result = None
432
+ raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
433
+ finally:
434
+ loop.close()
435
+ return result
436
+
437
+
438
+ def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
439
+ steps: int = 0, cfg: int = 0, seed: int = -1,
440
+ pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
441
+ import random
442
+ if model_name_dummy in ["NA", ""]: return gr.update()
443
+ random.seed()
444
+ model_name = random.choice(list(loaded_models.keys()))
445
+ try:
446
+ loop = asyncio.get_running_loop()
447
+ except Exception:
448
+ loop = asyncio.new_event_loop()
449
+ try:
450
+ prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
451
+ result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
452
+ steps, cfg, seed, save_path, inference_timeout))
453
+ except (Exception, asyncio.CancelledError) as e:
454
+ print(e)
455
+ print(f"Task aborted: {model_name}, Error: {e}")
456
+ result = None
457
+ raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
458
+ finally:
459
+ loop.close()
460
+ return result
461
+
462
+
463
+ def save_image(image, savefile, modelname, prompt, nprompt, height=0, width=0, steps=0, cfg=0, seed=-1):
464
+ from PIL import Image, PngImagePlugin
465
+ import json
466
+ try:
467
+ metadata = {"prompt": prompt, "negative_prompt": nprompt, "Model": {"Model": modelname.split("/")[-1]}}
468
+ if steps > 0: metadata["num_inference_steps"] = steps
469
+ if cfg > 0: metadata["guidance_scale"] = cfg
470
+ if seed != -1: metadata["seed"] = seed
471
+ if width > 0 and height > 0: metadata["resolution"] = f"{width} x {height}"
472
+ metadata_str = json.dumps(metadata)
473
+ info = PngImagePlugin.PngInfo()
474
+ info.add_text("metadata", metadata_str)
475
+ image.save(savefile, "PNG", pnginfo=info)
476
+ return str(Path(savefile).resolve())
477
+ except Exception as e:
478
+ print(f"Failed to save image file: {e}")
479
+ raise Exception(f"Failed to save image file:") from e
480
+
481
+
482
+ def randomize_seed():
483
+ from random import seed, randint
484
+ MAX_SEED = 2**32-1
485
+ seed()
486
+ rseed = randint(0, MAX_SEED)
487
+ return rseed
488
+
489
+
490
+ from translatepy import Translator
491
+ translator = Translator()
492
+ def translate_to_en(input: str):
493
+ try:
494
+ output = str(translator.translate(input, 'English'))
495
+ except Exception as e:
496
+ output = input
497
+ print(e)
498
+ return output
499
+
500
+
501
+ def get_recom_prompt_mode():
502
+ return list(recom_prompt_type.keys())
503
+
504
+
505
+ def get_models():
506
+ return [""] + list(loaded_models.keys())
507
+
508
+
509
+ def get_def_model(i: int):
510
+ if i == 0: return list(loaded_models.keys())[0]
511
+ else: return ""
512
+
513
+
514
+ def warm_models(models: list[str]):
515
+ for model in models:
516
+ asyncio.new_event_loop().run_in_executor(None, warm_model, model)
517
+
518
+
519
+ def gen_image(model_name: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
520
+ steps: int = 0, cfg: int = 0, seed: int = -1, recom_mode = "None"):
521
+ if model_name in ["NA", ""]: return gr.update()
522
+ try:
523
+ loop = asyncio.get_running_loop()
524
+ except Exception:
525
+ loop = asyncio.new_event_loop()
526
+ try:
527
+ prompt, neg_prompt = insert_recom_prompt(prompt, neg_prompt, recom_mode)
528
+ result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
529
+ steps, cfg, seed, None, inference_timeout))
530
+ except (Exception, asyncio.CancelledError) as e:
531
+ print(e)
532
+ print(f"Task aborted: {model_name}, Error: {e}")
533
+ result = None
534
+ raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
535
+ finally:
536
+ loop.close()
537
+ return result
538
+
539
+
540
+ from model import models, models_animagine, models_noob, models_ill
541
+ all_models = list_uniq(models_animagine + models_noob + models_ill + models)
542
+ load_models(all_models)
543
+ warm_models(models_animagine + models_noob + models_ill)
544
+
545
+
546
+ def update_prompt(tag: str, model_type: str):
547
+ if model_type == "Animagine 3.1": model = models_animagine[0]
548
+ elif model_type == "Illustrious": model = models_ill[0]
549
+ else: model = models_noob[0]
550
+ recom_mode = "Animagine"
551
+ return tag, model, recom_mode
model.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from hft2i import find_model_list
2
+
3
+
4
+ models_animagine = ["cagliostrolab/animagine-xl-3.1", "yodayo-ai/kivotos-xl-2.0"]
5
+ models_noob = ["John6666/noobai-xl-nai-xl-epsilonpred11version-sdxl"]
6
+ models_ill = ["Raelina/Raehoshi-illust-XL-3", "John6666/illustrious-xl-early-release-v0-sdxl"]
7
+
8
+
9
+ models = [
10
+ "yodayo-ai/clandestine-xl-1.0",
11
+ "yodayo-ai/kivotos-xl-2.0",
12
+ "yodayo-ai/holodayo-xl-2.1",
13
+ "cagliostrolab/animagine-xl-3.1",
14
+ "votepurchase/ponyDiffusionV6XL",
15
+ "eienmojiki/Anything-XL",
16
+ "eienmojiki/Starry-XL-v5.2",
17
+ "digiplay/MilkyWonderland_v1",
18
+ "digiplay/majicMIX_sombre_v2",
19
+ "digiplay/majicMIX_realistic_v7",
20
+ "votepurchase/counterfeitV30_v30",
21
+ "Meina/MeinaMix_V11",
22
+ "KBlueLeaf/Kohaku-XL-Epsilon-rev3",
23
+ "KBlueLeaf/Kohaku-XL-Zeta",
24
+ "kayfahaarukku/UrangDiffusion-1.4",
25
+ "Eugeoter/artiwaifu-diffusion-2.0",
26
+ "Raelina/Rae-Diffusion-XL-V2",
27
+ "Raelina/Raemu-XL-V4",
28
+ ]
29
+
30
+
31
+ #models += find_model_list("John6666", ["anime", "illustrious"], "", "last_modified", 50, public=True)
32
+ #models += find_model_list("John6666", ["anime", "animagine"], "", "last_modified", 50, public=True)
33
+
34
+
35
+ #models = find_model_list("Disty0", [], "", "last_modified", 100)
36
+
37
+
38
+ # Examples:
39
+ #models = ["yodayo-ai/kivotos-xl-2.0", "yodayo-ai/holodayo-xl-2.1"] # specific models
40
+ #models = find_model_list("John6666", [], "", "last_modified", 20) # John6666's latest 20 models
41
+ #models = find_model_list("John6666", ["anime"], "", "last_modified", 20) # John6666's latest 20 models with "anime" tag
42
+ #models = find_model_list("John6666", [], "anime", "last_modified", 20) # John6666's latest 20 models without "anime" tag
43
+ #models = find_model_list("", [], "", "last_modified", 20) # latest 20 text-to-image models of huggingface
44
+ #models = find_model_list("", [], "", "downloads", 20) # monthly most downloaded 20 text-to-image models of huggingface
requirements.txt CHANGED
@@ -1,2 +1,4 @@
 
1
  pandas
2
- pykakasi
 
 
1
+ huggingface-hub
2
  pandas
3
+ pykakasi
4
+ translatepy