| import gradio as gr | |
| import asyncio | |
| from pathlib import Path | |
| loaded_models = {} | |
| model_info_dict = {} | |
| def to_list(s): | |
| return [x.strip() for x in s.split(",")] | |
| def list_sub(a, b): | |
| return [e for e in a if e not in b] | |
| def list_uniq(l): | |
| return sorted(set(l), key=l.index) | |
| def is_repo_name(s): | |
| import re | |
| return re.fullmatch(r'^[^/]+?/[^/]+?$', s) | |
| def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30): | |
| from huggingface_hub import HfApi | |
| api = HfApi() | |
| default_tags = ["diffusers"] | |
| if not sort: sort = "last_modified" | |
| models = [] | |
| try: | |
| model_infos = api.list_models(author=author, pipeline_tag="text-to-image", | |
| tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit * 5) | |
| except Exception as e: | |
| print(f"Error: Failed to list models.") | |
| print(e) | |
| return models | |
| for model in model_infos: | |
| if not model.private and not model.gated: | |
| if not_tag and not_tag in model.tags: continue | |
| models.append(model.id) | |
| if len(models) == limit: break | |
| return models | |
| def get_t2i_model_info_dict(repo_id: str): | |
| from huggingface_hub import HfApi | |
| api = HfApi() | |
| info = {"md": "None"} | |
| try: | |
| if not is_repo_name(repo_id) or not api.repo_exists(repo_id=repo_id): return info | |
| model = api.model_info(repo_id=repo_id) | |
| except Exception as e: | |
| print(f"Error: Failed to get {repo_id}'s info.") | |
| print(e) | |
| return info | |
| if model.private or model.gated: return info | |
| try: | |
| tags = model.tags | |
| except Exception: | |
| return info | |
| if not 'diffusers' in model.tags: return info | |
| if 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL" | |
| elif 'diffusers:StableDiffusionPipeline' in tags: info["ver"] = "SD1.5" | |
| elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3" | |
| else: info["ver"] = "Other" | |
| info["url"] = f"https://huggingface.co./{repo_id}/" | |
| if model.card_data and model.card_data.tags: | |
| info["tags"] = model.card_data.tags | |
| info["downloads"] = model.downloads | |
| info["likes"] = model.likes | |
| info["last_modified"] = model.last_modified.strftime("lastmod: %Y-%m-%d") | |
| un_tags = ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl'] | |
| descs = [info["ver"]] + list_sub(info["tags"], un_tags) + [f'DLs: {info["downloads"]}'] + [f'❤: {info["likes"]}'] + [info["last_modified"]] | |
| info["md"] = f'Model Info: {", ".join(descs)} [Model Repo]({info["url"]})' | |
| return info | |
| def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)): | |
| from datetime import datetime, timezone, timedelta | |
| progress(0, desc="Updating gallery...") | |
| dt_now = datetime.now(timezone(timedelta(hours=9))) | |
| basename = dt_now.strftime('%Y%m%d_%H%M%S_') | |
| i = 1 | |
| if not images: return images | |
| output_images = [] | |
| output_paths = [] | |
| for image in images: | |
| filename = f'{image[1]}_{basename}{str(i)}.png' | |
| i += 1 | |
| oldpath = Path(image[0]) | |
| newpath = oldpath | |
| try: | |
| if oldpath.stem == "image" and oldpath.exists(): | |
| newpath = oldpath.resolve().rename(Path(filename).resolve()) | |
| except Exception as e: | |
| print(e) | |
| pass | |
| finally: | |
| output_paths.append(str(newpath)) | |
| output_images.append((str(newpath), str(filename))) | |
| progress(1, desc="Gallery updated.") | |
| return gr.update(value=output_images), gr.update(value=output_paths) | |
| def load_model(model_name: str): | |
| global loaded_models | |
| global model_info_dict | |
| if model_name in loaded_models.keys(): return loaded_models[model_name] | |
| try: | |
| loaded_models[model_name] = gr.load(f'models/{model_name}') | |
| print(f"Loaded: {model_name}") | |
| except Exception as e: | |
| if model_name in loaded_models.keys(): del loaded_models[model_name] | |
| print(f"Failed to load: {model_name}") | |
| print(e) | |
| return None | |
| try: | |
| model_info_dict[model_name] = get_t2i_model_info_dict(model_name) | |
| except Exception as e: | |
| if model_name in model_info_dict.keys(): del model_info_dict[model_name] | |
| print(e) | |
| return loaded_models[model_name] | |
| async def async_load_models(models: list, limit: int=5): | |
| sem = asyncio.Semaphore(limit) | |
| async def async_load_model(model: str): | |
| async with sem: | |
| try: | |
| return load_model(model) | |
| except Exception as e: | |
| print(e) | |
| tasks = [asyncio.create_task(async_load_model(model)) for model in models] | |
| return await asyncio.wait(tasks) | |
| def load_models(models: list, limit: int=5): | |
| loop = asyncio.get_event_loop() | |
| try: | |
| loop.run_until_complete(async_load_models(models, limit)) | |
| except Exception as e: | |
| print(e) | |
| pass | |
| loop.close() | |
| positive_prefix = { | |
| "Pony": to_list("score_9, score_8_up, score_7_up"), | |
| "Pony Anime": to_list("source_anime, score_9, score_8_up, score_7_up"), | |
| } | |
| positive_suffix = { | |
| "Common": to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres"), | |
| "Anime": to_list("anime artwork, anime style, studio anime, highly detailed"), | |
| } | |
| negative_prefix = { | |
| "Pony": to_list("score_6, score_5, score_4"), | |
| "Pony Anime": to_list("score_6, score_5, score_4, source_pony, source_furry, source_cartoon"), | |
| "Pony Real": to_list("score_6, score_5, score_4, source_anime, source_pony, source_furry, source_cartoon"), | |
| } | |
| negative_suffix = { | |
| "Common": 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]"), | |
| "Pony Anime": to_list("busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends"), | |
| "Pony Real": to_list("ugly, airbrushed, simple background, cgi, cartoon, anime"), | |
| } | |
| positive_all = negative_all = [] | |
| for k, v in (positive_prefix | positive_suffix).items(): | |
| positive_all = positive_all + v + [s.replace("_", " ") for s in v] | |
| positive_all = list_uniq(positive_all) | |
| for k, v in (negative_prefix | negative_suffix).items(): | |
| negative_all = negative_all + v + [s.replace("_", " ") for s in v] | |
| positive_all = list_uniq(positive_all) | |
| def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []): | |
| def flatten(src): | |
| return [item for row in src for item in row] | |
| prompts = to_list(prompt) | |
| neg_prompts = to_list(neg_prompt) | |
| prompts = list_sub(prompts, positive_all) | |
| neg_prompts = list_sub(neg_prompts, negative_all) | |
| last_empty_p = [""] if not prompts and type != "None" else [] | |
| last_empty_np = [""] if not neg_prompts and type != "None" else [] | |
| prefix_ps = flatten([positive_prefix.get(s, []) for s in pos_pre]) | |
| suffix_ps = flatten([positive_suffix.get(s, []) for s in pos_suf]) | |
| prefix_nps = flatten([negative_prefix.get(s, []) for s in neg_pre]) | |
| suffix_nps = flatten([negative_suffix.get(s, []) for s in neg_suf]) | |
| prompt = ", ".join(list_uniq(prefix_ps + prompts + suffix_ps) + last_empty_p) | |
| neg_prompt = ", ".join(list_uniq(prefix_nps + neg_prompts + suffix_nps) + last_empty_np) | |
| return prompt, neg_prompt | |
| recom_prompt_type = { | |
| "None": ([], [], [], []), | |
| "Auto": ([], [], [], []), | |
| "Common": ([], ["Common"], [], ["Common"]), | |
| "Animagine": ([], ["Common", "Anime"], [], ["Common"]), | |
| "Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]), | |
| "Pony Anime": (["Pony", "Pony Anime"], ["Common", "Anime"], ["Pony", "Pony Anime"], ["Common", "Pony Anime"]), | |
| "Pony Real": (["Pony"], ["Common"], ["Pony", "Pony Real"], ["Common", "Pony Real"]), | |
| } | |
| enable_auto_recom_prompt = False | |
| def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"): | |
| global enable_auto_recom_prompt | |
| if type == "Auto": enable_auto_recom_prompt = True | |
| else: enable_auto_recom_prompt = False | |
| pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], [])) | |
| return recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf) | |
| def set_recom_prompt_preset(type: str = "None"): | |
| pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], [])) | |
| return pos_pre, pos_suf, neg_pre, neg_suf | |
| def get_recom_prompt_type(): | |
| type = list(recom_prompt_type.keys()) | |
| type.remove("Auto") | |
| return type | |
| def get_positive_prefix(): | |
| return list(positive_prefix.keys()) | |
| def get_positive_suffix(): | |
| return list(positive_suffix.keys()) | |
| def get_negative_prefix(): | |
| return list(negative_prefix.keys()) | |
| def get_negative_suffix(): | |
| return list(negative_suffix.keys()) | |
| def get_model_info_md(model_name: str): | |
| if model_name in model_info_dict.keys(): return model_info_dict[model_name].get("md", "") | |
| def change_model(model_name: str): | |
| load_model(model_name) | |
| return get_model_info_md(model_name) | |
| def infer(prompt: str, neg_prompt: str, model_name: str, progress=gr.Progress(track_tqdm=True)): | |
| from PIL import Image | |
| import random | |
| seed = "" | |
| rand = random.randint(1, 500) | |
| for i in range(rand): | |
| seed += " " | |
| caption = model_name.split("/")[-1] | |
| try: | |
| model = load_model(model_name) | |
| if not model: return (Image.Image(), None) | |
| image_path = model(prompt + seed) | |
| image = Image.open(image_path).convert('RGB') | |
| except Exception as e: | |
| print(e) | |
| return (Image.Image(), None) | |
| return (image, caption) | |
| def infer_multi(prompt: str, neg_prompt: str, results: list, image_num: float, model_name: str, | |
| pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], progress=gr.Progress(track_tqdm=True)): | |
| image_num = int(image_num) | |
| images = results if results else [] | |
| prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf) | |
| for i in range(image_num): | |
| images.append(infer(prompt, neg_prompt, model_name, recom_prompt)) | |
| yield images | |