MrDrmm commited on
Commit
86792a7
·
verified ·
1 Parent(s): 3ed10df

Update multit2i.py

Browse files
Files changed (1) hide show
  1. multit2i.py +502 -502
multit2i.py CHANGED
@@ -1,502 +1,502 @@
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 = 300
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):
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:
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'Model Info: {", ".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
- client = huggingface_hub.InferenceClient(model=model_name, headers=headers,
164
- token=hf_token, timeout=server_timeout)
165
- inputs = gr.components.Textbox(label="Input")
166
- outputs = gr.components.Image(label="Output")
167
- fn = client.text_to_image
168
-
169
- def query_huggingface_inference_endpoints(*data, **kwargs):
170
- try:
171
- data = fn(*data, **kwargs) # type: ignore
172
- except huggingface_hub.utils.HfHubHTTPError as e:
173
- if "429" in str(e):
174
- raise TooManyRequestsError() from e
175
- except Exception as e:
176
- raise Exception() from e
177
- return data
178
-
179
- interface_info = {
180
- "fn": query_huggingface_inference_endpoints,
181
- "inputs": inputs,
182
- "outputs": outputs,
183
- "title": model_name,
184
- }
185
- return gr.Interface(**interface_info)
186
-
187
-
188
- def load_model(model_name: str):
189
- global loaded_models
190
- global model_info_dict
191
- if model_name in loaded_models.keys(): return loaded_models[model_name]
192
- try:
193
- loaded_models[model_name] = load_from_model(model_name, hf_token=HF_TOKEN)
194
- print(f"Loaded: {model_name}")
195
- except Exception as e:
196
- if model_name in loaded_models.keys(): del loaded_models[model_name]
197
- print(f"Failed to load: {model_name}")
198
- print(e)
199
- return None
200
- try:
201
- model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
202
- print(f"Assigned: {model_name}")
203
- except Exception as e:
204
- if model_name in model_info_dict.keys(): del model_info_dict[model_name]
205
- print(f"Failed to assigned: {model_name}")
206
- print(e)
207
- return loaded_models[model_name]
208
-
209
-
210
- def load_model_api(model_name: str):
211
- global loaded_models
212
- global model_info_dict
213
- if model_name in loaded_models.keys(): return loaded_models[model_name]
214
- try:
215
- client = InferenceClient(timeout=5)
216
- status = client.get_model_status(model_name, token=HF_TOKEN)
217
- if status is None or status.framework != "diffusers" or status.state not in ["Loadable", "Loaded"]:
218
- print(f"Failed to load by API: {model_name}")
219
- return None
220
- else:
221
- loaded_models[model_name] = InferenceClient(model_name, token=HF_TOKEN, timeout=server_timeout)
222
- print(f"Loaded by API: {model_name}")
223
- except Exception as e:
224
- if model_name in loaded_models.keys(): del loaded_models[model_name]
225
- print(f"Failed to load by API: {model_name}")
226
- print(e)
227
- return None
228
- try:
229
- model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
230
- print(f"Assigned by API: {model_name}")
231
- except Exception as e:
232
- if model_name in model_info_dict.keys(): del model_info_dict[model_name]
233
- print(f"Failed to assigned by API: {model_name}")
234
- print(e)
235
- return loaded_models[model_name]
236
-
237
-
238
- def load_models(models: list):
239
- for model in models:
240
- load_model(model)
241
-
242
-
243
- positive_prefix = {
244
- "Pony": to_list("score_9, score_8_up, score_7_up"),
245
- "Pony Anime": to_list("source_anime, anime, score_9, score_8_up, score_7_up"),
246
- }
247
- positive_suffix = {
248
- "Common": to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres"),
249
- "Anime": to_list("anime artwork, anime style, studio anime, highly detailed"),
250
- }
251
- negative_prefix = {
252
- "Pony": to_list("score_6, score_5, score_4"),
253
- "Pony Anime": to_list("score_6, score_5, score_4, source_pony, source_furry, source_cartoon"),
254
- "Pony Real": to_list("score_6, score_5, score_4, source_anime, source_pony, source_furry, source_cartoon"),
255
- }
256
- negative_suffix = {
257
- "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]"),
258
- "Pony Anime": to_list("busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends"),
259
- "Pony Real": to_list("ugly, airbrushed, simple background, cgi, cartoon, anime"),
260
- }
261
- positive_all = negative_all = []
262
- for k, v in (positive_prefix | positive_suffix).items():
263
- positive_all = positive_all + v + [s.replace("_", " ") for s in v]
264
- positive_all = list_uniq(positive_all)
265
- for k, v in (negative_prefix | negative_suffix).items():
266
- negative_all = negative_all + v + [s.replace("_", " ") for s in v]
267
- positive_all = list_uniq(positive_all)
268
-
269
-
270
- def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
271
- def flatten(src):
272
- return [item for row in src for item in row]
273
- prompts = to_list(prompt)
274
- neg_prompts = to_list(neg_prompt)
275
- prompts = list_sub(prompts, positive_all)
276
- neg_prompts = list_sub(neg_prompts, negative_all)
277
- last_empty_p = [""] if not prompts and type != "None" else []
278
- last_empty_np = [""] if not neg_prompts and type != "None" else []
279
- prefix_ps = flatten([positive_prefix.get(s, []) for s in pos_pre])
280
- suffix_ps = flatten([positive_suffix.get(s, []) for s in pos_suf])
281
- prefix_nps = flatten([negative_prefix.get(s, []) for s in neg_pre])
282
- suffix_nps = flatten([negative_suffix.get(s, []) for s in neg_suf])
283
- prompt = ", ".join(list_uniq(prefix_ps + prompts + suffix_ps) + last_empty_p)
284
- neg_prompt = ", ".join(list_uniq(prefix_nps + neg_prompts + suffix_nps) + last_empty_np)
285
- return prompt, neg_prompt
286
-
287
-
288
- recom_prompt_type = {
289
- "None": ([], [], [], []),
290
- "Auto": ([], [], [], []),
291
- "Common": ([], ["Common"], [], ["Common"]),
292
- "Animagine": ([], ["Common", "Anime"], [], ["Common"]),
293
- "Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]),
294
- "Pony Anime": (["Pony", "Pony Anime"], ["Common", "Anime"], ["Pony", "Pony Anime"], ["Common", "Pony Anime"]),
295
- "Pony Real": (["Pony"], ["Common"], ["Pony", "Pony Real"], ["Common", "Pony Real"]),
296
- }
297
-
298
-
299
- enable_auto_recom_prompt = False
300
- def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
301
- global enable_auto_recom_prompt
302
- if type == "Auto": enable_auto_recom_prompt = True
303
- else: enable_auto_recom_prompt = False
304
- pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
305
- return recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
306
-
307
-
308
- def set_recom_prompt_preset(type: str = "None"):
309
- pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
310
- return pos_pre, pos_suf, neg_pre, neg_suf
311
-
312
-
313
- def get_recom_prompt_type():
314
- type = list(recom_prompt_type.keys())
315
- type.remove("Auto")
316
- return type
317
-
318
-
319
- def get_positive_prefix():
320
- return list(positive_prefix.keys())
321
-
322
-
323
- def get_positive_suffix():
324
- return list(positive_suffix.keys())
325
-
326
-
327
- def get_negative_prefix():
328
- return list(negative_prefix.keys())
329
-
330
-
331
- def get_negative_suffix():
332
- return list(negative_suffix.keys())
333
-
334
-
335
- def get_tag_type(pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
336
- tag_type = "danbooru"
337
- words = pos_pre + pos_suf + neg_pre + neg_suf
338
- for word in words:
339
- if "Pony" in word:
340
- tag_type = "e621"
341
- break
342
- return tag_type
343
-
344
-
345
- def get_model_info_md(model_name: str):
346
- if model_name in model_info_dict.keys(): return model_info_dict[model_name].get("md", "")
347
-
348
-
349
- def change_model(model_name: str):
350
- load_model_api(model_name)
351
- return get_model_info_md(model_name)
352
-
353
-
354
- def warm_model(model_name: str):
355
- model = load_model_api(model_name)
356
- if model:
357
- try:
358
- print(f"Warming model: {model_name}")
359
- infer_body(model, " ")
360
- except Exception as e:
361
- print(e)
362
-
363
-
364
- # https://huggingface.co/docs/api-inference/detailed_parameters
365
- # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
366
- def infer_body(client: InferenceClient | gr.Interface | object, model_str: str, prompt: str, neg_prompt: str = "",
367
- height: int = 0, width: int = 0, steps: int = 0, cfg: int = 0, seed: int = -1):
368
- png_path = "image.png"
369
- kwargs = {}
370
- if height > 0: kwargs["height"] = height
371
- if width > 0: kwargs["width"] = width
372
- if steps > 0: kwargs["num_inference_steps"] = steps
373
- if cfg > 0: cfg = kwargs["guidance_scale"] = cfg
374
- if seed == -1: kwargs["seed"] = randomize_seed()
375
- else: kwargs["seed"] = seed
376
- try:
377
- if isinstance(client, InferenceClient):
378
- image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
379
- elif isinstance(client, gr.Interface):
380
- image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
381
- else: return None
382
- if isinstance(image, tuple): return None
383
- return save_image(image, png_path, model_str, prompt, neg_prompt, height, width, steps, cfg, seed)
384
- except Exception as e:
385
- print(e)
386
- raise Exception() from e
387
-
388
-
389
- async def infer(model_name: str, prompt: str, neg_prompt: str ="", height: int = 0, width: int = 0,
390
- steps: int = 0, cfg: int = 0, seed: int = -1,
391
- save_path: str | None = None, timeout: float = inference_timeout):
392
- model = load_model(model_name)
393
- if not model: return None
394
- task = asyncio.create_task(asyncio.to_thread(infer_body, model, model_name, prompt, neg_prompt,
395
- height, width, steps, cfg, seed))
396
- await asyncio.sleep(0)
397
- try:
398
- result = await asyncio.wait_for(task, timeout=timeout)
399
- except asyncio.TimeoutError as e:
400
- print(e)
401
- print(f"Task timed out: {model_name}")
402
- if not task.done(): task.cancel()
403
- result = None
404
- raise Exception(f"Task timed out: {model_name}") from e
405
- except Exception as e:
406
- print(e)
407
- if not task.done(): task.cancel()
408
- result = None
409
- raise Exception() from e
410
- if task.done() and result is not None:
411
- with lock:
412
- image = rename_image(result, model_name, save_path)
413
- return image
414
- return None
415
-
416
-
417
- # https://github.com/aio-libs/pytest-aiohttp/issues/8 # also AsyncInferenceClient is buggy.
418
- def infer_fn(model_name: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
419
- steps: int = 0, cfg: int = 0, seed: int = -1,
420
- pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
421
- if model_name == 'NA':
422
- return None
423
- try:
424
- loop = asyncio.get_running_loop()
425
- except Exception:
426
- loop = asyncio.new_event_loop()
427
- try:
428
- prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
429
- result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
430
- steps, cfg, seed, save_path, inference_timeout))
431
- except (Exception, asyncio.CancelledError) as e:
432
- print(e)
433
- print(f"Task aborted: {model_name}, Error: {e}")
434
- result = None
435
- raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
436
- finally:
437
- loop.close()
438
- return result
439
-
440
-
441
- def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
442
- steps: int = 0, cfg: int = 0, seed: int = -1,
443
- pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
444
- import random
445
- if model_name_dummy == 'NA':
446
- return None
447
- random.seed()
448
- model_name = random.choice(list(loaded_models.keys()))
449
- try:
450
- loop = asyncio.get_running_loop()
451
- except Exception:
452
- loop = asyncio.new_event_loop()
453
- try:
454
- prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
455
- result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
456
- steps, cfg, seed, save_path, inference_timeout))
457
- except (Exception, asyncio.CancelledError) as e:
458
- print(e)
459
- print(f"Task aborted: {model_name}, Error: {e}")
460
- result = None
461
- raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
462
- finally:
463
- loop.close()
464
- return result
465
-
466
-
467
- def save_image(image, savefile, modelname, prompt, nprompt, height=0, width=0, steps=0, cfg=0, seed=-1):
468
- from PIL import Image, PngImagePlugin
469
- import json
470
- try:
471
- metadata = {"prompt": prompt, "negative_prompt": nprompt, "Model": {"Model": modelname.split("/")[-1]}}
472
- if steps > 0: metadata["num_inference_steps"] = steps
473
- if cfg > 0: metadata["guidance_scale"] = cfg
474
- if seed != -1: metadata["seed"] = seed
475
- if width > 0 and height > 0: metadata["resolution"] = f"{width} x {height}"
476
- metadata_str = json.dumps(metadata)
477
- info = PngImagePlugin.PngInfo()
478
- info.add_text("metadata", metadata_str)
479
- image.save(savefile, "PNG", pnginfo=info)
480
- return str(Path(savefile).resolve())
481
- except Exception as e:
482
- print(f"Failed to save image file: {e}")
483
- raise Exception(f"Failed to save image file:") from e
484
-
485
-
486
- def randomize_seed():
487
- from random import seed, randint
488
- MAX_SEED = 2**32-1
489
- seed()
490
- rseed = randint(0, MAX_SEED)
491
- return rseed
492
-
493
-
494
- from translatepy import Translator
495
- translator = Translator()
496
- def translate_to_en(input: str):
497
- try:
498
- output = str(translator.translate(input, 'English'))
499
- except Exception as e:
500
- output = input
501
- print(e)
502
- return output
 
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 = 300
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=50)
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):
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:
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'Model Info: {", ".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
+ client = huggingface_hub.InferenceClient(model=model_name, headers=headers,
164
+ token=hf_token, timeout=server_timeout)
165
+ inputs = gr.components.Textbox(label="Input")
166
+ outputs = gr.components.Image(label="Output")
167
+ fn = client.text_to_image
168
+
169
+ def query_huggingface_inference_endpoints(*data, **kwargs):
170
+ try:
171
+ data = fn(*data, **kwargs) # type: ignore
172
+ except huggingface_hub.utils.HfHubHTTPError as e:
173
+ if "429" in str(e):
174
+ raise TooManyRequestsError() from e
175
+ except Exception as e:
176
+ raise Exception() from e
177
+ return data
178
+
179
+ interface_info = {
180
+ "fn": query_huggingface_inference_endpoints,
181
+ "inputs": inputs,
182
+ "outputs": outputs,
183
+ "title": model_name,
184
+ }
185
+ return gr.Interface(**interface_info)
186
+
187
+
188
+ def load_model(model_name: str):
189
+ global loaded_models
190
+ global model_info_dict
191
+ if model_name in loaded_models.keys(): return loaded_models[model_name]
192
+ try:
193
+ loaded_models[model_name] = load_from_model(model_name, hf_token=HF_TOKEN)
194
+ print(f"Loaded: {model_name}")
195
+ except Exception as e:
196
+ if model_name in loaded_models.keys(): del loaded_models[model_name]
197
+ print(f"Failed to load: {model_name}")
198
+ print(e)
199
+ return None
200
+ try:
201
+ model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
202
+ print(f"Assigned: {model_name}")
203
+ except Exception as e:
204
+ if model_name in model_info_dict.keys(): del model_info_dict[model_name]
205
+ print(f"Failed to assigned: {model_name}")
206
+ print(e)
207
+ return loaded_models[model_name]
208
+
209
+
210
+ def load_model_api(model_name: str):
211
+ global loaded_models
212
+ global model_info_dict
213
+ if model_name in loaded_models.keys(): return loaded_models[model_name]
214
+ try:
215
+ client = InferenceClient(timeout=5)
216
+ status = client.get_model_status(model_name, token=HF_TOKEN)
217
+ if status is None or status.framework != "diffusers" or status.state not in ["Loadable", "Loaded"]:
218
+ print(f"Failed to load by API: {model_name}")
219
+ return None
220
+ else:
221
+ loaded_models[model_name] = InferenceClient(model_name, token=HF_TOKEN, timeout=server_timeout)
222
+ print(f"Loaded by API: {model_name}")
223
+ except Exception as e:
224
+ if model_name in loaded_models.keys(): del loaded_models[model_name]
225
+ print(f"Failed to load by API: {model_name}")
226
+ print(e)
227
+ return None
228
+ try:
229
+ model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
230
+ print(f"Assigned by API: {model_name}")
231
+ except Exception as e:
232
+ if model_name in model_info_dict.keys(): del model_info_dict[model_name]
233
+ print(f"Failed to assigned by API: {model_name}")
234
+ print(e)
235
+ return loaded_models[model_name]
236
+
237
+
238
+ def load_models(models: list):
239
+ for model in models:
240
+ load_model(model)
241
+
242
+
243
+ positive_prefix = {
244
+ "Pony": to_list("score_9, score_8_up, score_7_up"),
245
+ "Pony Anime": to_list("source_anime, anime, score_9, score_8_up, score_7_up"),
246
+ }
247
+ positive_suffix = {
248
+ "Common": to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres"),
249
+ "Anime": to_list("anime artwork, anime style, studio anime, highly detailed"),
250
+ }
251
+ negative_prefix = {
252
+ "Pony": to_list("score_6, score_5, score_4"),
253
+ "Pony Anime": to_list("score_6, score_5, score_4, source_pony, source_furry, source_cartoon"),
254
+ "Pony Real": to_list("score_6, score_5, score_4, source_anime, source_pony, source_furry, source_cartoon"),
255
+ }
256
+ negative_suffix = {
257
+ "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]"),
258
+ "Pony Anime": to_list("busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends"),
259
+ "Pony Real": to_list("ugly, airbrushed, simple background, cgi, cartoon, anime"),
260
+ }
261
+ positive_all = negative_all = []
262
+ for k, v in (positive_prefix | positive_suffix).items():
263
+ positive_all = positive_all + v + [s.replace("_", " ") for s in v]
264
+ positive_all = list_uniq(positive_all)
265
+ for k, v in (negative_prefix | negative_suffix).items():
266
+ negative_all = negative_all + v + [s.replace("_", " ") for s in v]
267
+ positive_all = list_uniq(positive_all)
268
+
269
+
270
+ def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
271
+ def flatten(src):
272
+ return [item for row in src for item in row]
273
+ prompts = to_list(prompt)
274
+ neg_prompts = to_list(neg_prompt)
275
+ prompts = list_sub(prompts, positive_all)
276
+ neg_prompts = list_sub(neg_prompts, negative_all)
277
+ last_empty_p = [""] if not prompts and type != "None" else []
278
+ last_empty_np = [""] if not neg_prompts and type != "None" else []
279
+ prefix_ps = flatten([positive_prefix.get(s, []) for s in pos_pre])
280
+ suffix_ps = flatten([positive_suffix.get(s, []) for s in pos_suf])
281
+ prefix_nps = flatten([negative_prefix.get(s, []) for s in neg_pre])
282
+ suffix_nps = flatten([negative_suffix.get(s, []) for s in neg_suf])
283
+ prompt = ", ".join(list_uniq(prefix_ps + prompts + suffix_ps) + last_empty_p)
284
+ neg_prompt = ", ".join(list_uniq(prefix_nps + neg_prompts + suffix_nps) + last_empty_np)
285
+ return prompt, neg_prompt
286
+
287
+
288
+ recom_prompt_type = {
289
+ "None": ([], [], [], []),
290
+ "Auto": ([], [], [], []),
291
+ "Common": ([], ["Common"], [], ["Common"]),
292
+ "Animagine": ([], ["Common", "Anime"], [], ["Common"]),
293
+ "Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]),
294
+ "Pony Anime": (["Pony", "Pony Anime"], ["Common", "Anime"], ["Pony", "Pony Anime"], ["Common", "Pony Anime"]),
295
+ "Pony Real": (["Pony"], ["Common"], ["Pony", "Pony Real"], ["Common", "Pony Real"]),
296
+ }
297
+
298
+
299
+ enable_auto_recom_prompt = False
300
+ def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
301
+ global enable_auto_recom_prompt
302
+ if type == "Auto": enable_auto_recom_prompt = True
303
+ else: enable_auto_recom_prompt = False
304
+ pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
305
+ return recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
306
+
307
+
308
+ def set_recom_prompt_preset(type: str = "None"):
309
+ pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
310
+ return pos_pre, pos_suf, neg_pre, neg_suf
311
+
312
+
313
+ def get_recom_prompt_type():
314
+ type = list(recom_prompt_type.keys())
315
+ type.remove("Auto")
316
+ return type
317
+
318
+
319
+ def get_positive_prefix():
320
+ return list(positive_prefix.keys())
321
+
322
+
323
+ def get_positive_suffix():
324
+ return list(positive_suffix.keys())
325
+
326
+
327
+ def get_negative_prefix():
328
+ return list(negative_prefix.keys())
329
+
330
+
331
+ def get_negative_suffix():
332
+ return list(negative_suffix.keys())
333
+
334
+
335
+ def get_tag_type(pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
336
+ tag_type = "danbooru"
337
+ words = pos_pre + pos_suf + neg_pre + neg_suf
338
+ for word in words:
339
+ if "Pony" in word:
340
+ tag_type = "e621"
341
+ break
342
+ return tag_type
343
+
344
+
345
+ def get_model_info_md(model_name: str):
346
+ if model_name in model_info_dict.keys(): return model_info_dict[model_name].get("md", "")
347
+
348
+
349
+ def change_model(model_name: str):
350
+ load_model_api(model_name)
351
+ return get_model_info_md(model_name)
352
+
353
+
354
+ def warm_model(model_name: str):
355
+ model = load_model_api(model_name)
356
+ if model:
357
+ try:
358
+ print(f"Warming model: {model_name}")
359
+ infer_body(model, " ")
360
+ except Exception as e:
361
+ print(e)
362
+
363
+
364
+ # https://huggingface.co/docs/api-inference/detailed_parameters
365
+ # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
366
+ def infer_body(client: InferenceClient | gr.Interface | object, model_str: str, prompt: str, neg_prompt: str = "",
367
+ height: int = 0, width: int = 0, steps: int = 0, cfg: int = 0, seed: int = -1):
368
+ png_path = "image.png"
369
+ kwargs = {}
370
+ if height > 0: kwargs["height"] = height
371
+ if width > 0: kwargs["width"] = width
372
+ if steps > 0: kwargs["num_inference_steps"] = steps
373
+ if cfg > 0: cfg = kwargs["guidance_scale"] = cfg
374
+ if seed == -1: kwargs["seed"] = randomize_seed()
375
+ else: kwargs["seed"] = seed
376
+ try:
377
+ if isinstance(client, InferenceClient):
378
+ image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
379
+ elif isinstance(client, gr.Interface):
380
+ image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
381
+ else: return None
382
+ if isinstance(image, tuple): return None
383
+ return save_image(image, png_path, model_str, prompt, neg_prompt, height, width, steps, cfg, seed)
384
+ except Exception as e:
385
+ print(e)
386
+ raise Exception() from e
387
+
388
+
389
+ async def infer(model_name: str, prompt: str, neg_prompt: str ="", height: int = 0, width: int = 0,
390
+ steps: int = 0, cfg: int = 0, seed: int = -1,
391
+ save_path: str | None = None, timeout: float = inference_timeout):
392
+ model = load_model(model_name)
393
+ if not model: return None
394
+ task = asyncio.create_task(asyncio.to_thread(infer_body, model, model_name, prompt, neg_prompt,
395
+ height, width, steps, cfg, seed))
396
+ await asyncio.sleep(0)
397
+ try:
398
+ result = await asyncio.wait_for(task, timeout=timeout)
399
+ except asyncio.TimeoutError as e:
400
+ print(e)
401
+ print(f"Task timed out: {model_name}")
402
+ if not task.done(): task.cancel()
403
+ result = None
404
+ raise Exception(f"Task timed out: {model_name}") from e
405
+ except Exception as e:
406
+ print(e)
407
+ if not task.done(): task.cancel()
408
+ result = None
409
+ raise Exception() from e
410
+ if task.done() and result is not None:
411
+ with lock:
412
+ image = rename_image(result, model_name, save_path)
413
+ return image
414
+ return None
415
+
416
+
417
+ # https://github.com/aio-libs/pytest-aiohttp/issues/8 # also AsyncInferenceClient is buggy.
418
+ def infer_fn(model_name: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
419
+ steps: int = 0, cfg: int = 0, seed: int = -1,
420
+ pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
421
+ if model_name == 'NA':
422
+ return None
423
+ try:
424
+ loop = asyncio.get_running_loop()
425
+ except Exception:
426
+ loop = asyncio.new_event_loop()
427
+ try:
428
+ prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
429
+ result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
430
+ steps, cfg, seed, save_path, inference_timeout))
431
+ except (Exception, asyncio.CancelledError) as e:
432
+ print(e)
433
+ print(f"Task aborted: {model_name}, Error: {e}")
434
+ result = None
435
+ raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
436
+ finally:
437
+ loop.close()
438
+ return result
439
+
440
+
441
+ def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
442
+ steps: int = 0, cfg: int = 0, seed: int = -1,
443
+ pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
444
+ import random
445
+ if model_name_dummy == 'NA':
446
+ return None
447
+ random.seed()
448
+ model_name = random.choice(list(loaded_models.keys()))
449
+ try:
450
+ loop = asyncio.get_running_loop()
451
+ except Exception:
452
+ loop = asyncio.new_event_loop()
453
+ try:
454
+ prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
455
+ result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
456
+ steps, cfg, seed, save_path, inference_timeout))
457
+ except (Exception, asyncio.CancelledError) as e:
458
+ print(e)
459
+ print(f"Task aborted: {model_name}, Error: {e}")
460
+ result = None
461
+ raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
462
+ finally:
463
+ loop.close()
464
+ return result
465
+
466
+
467
+ def save_image(image, savefile, modelname, prompt, nprompt, height=0, width=0, steps=0, cfg=0, seed=-1):
468
+ from PIL import Image, PngImagePlugin
469
+ import json
470
+ try:
471
+ metadata = {"prompt": prompt, "negative_prompt": nprompt, "Model": {"Model": modelname.split("/")[-1]}}
472
+ if steps > 0: metadata["num_inference_steps"] = steps
473
+ if cfg > 0: metadata["guidance_scale"] = cfg
474
+ if seed != -1: metadata["seed"] = seed
475
+ if width > 0 and height > 0: metadata["resolution"] = f"{width} x {height}"
476
+ metadata_str = json.dumps(metadata)
477
+ info = PngImagePlugin.PngInfo()
478
+ info.add_text("metadata", metadata_str)
479
+ image.save(savefile, "PNG", pnginfo=info)
480
+ return str(Path(savefile).resolve())
481
+ except Exception as e:
482
+ print(f"Failed to save image file: {e}")
483
+ raise Exception(f"Failed to save image file:") from e
484
+
485
+
486
+ def randomize_seed():
487
+ from random import seed, randint
488
+ MAX_SEED = 2**32-1
489
+ seed()
490
+ rseed = randint(0, MAX_SEED)
491
+ return rseed
492
+
493
+
494
+ from translatepy import Translator
495
+ translator = Translator()
496
+ def translate_to_en(input: str):
497
+ try:
498
+ output = str(translator.translate(input, 'English'))
499
+ except Exception as e:
500
+ output = input
501
+ print(e)
502
+ return output