import gradio as gr import asyncio from threading import RLock from pathlib import Path lock = RLock() 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 as e: print(e) 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}/" info["tags"] = model.card_data.tags if model.card_data and model.card_data.tags else [] 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_from_model(model_name: str, hf_token: str = None): import httpx import huggingface_hub from gradio.exceptions import ModelNotFoundError model_url = f"https://huggingface.co/{model_name}" api_url = f"https://api-inference.huggingface.co/models/{model_name}" print(f"Fetching model from: {model_url}") headers = {"Authorization": f"Bearer {hf_token}"} if hf_token is not None else {} response = httpx.request("GET", api_url, headers=headers) if response.status_code != 200: raise ModelNotFoundError( 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." ) headers["X-Wait-For-Model"] = "true" client = huggingface_hub.InferenceClient(model=model_name, headers=headers, token=hf_token) inputs = gr.components.Textbox(label="Input") outputs = gr.components.Image(label="Output") fn = client.text_to_image def query_huggingface_inference_endpoints(*data): return fn(*data) interface_info = { "fn": query_huggingface_inference_endpoints, "inputs": inputs, "outputs": outputs, "title": model_name, } return gr.Interface(**interface_info) 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] = load_from_model(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) print(f"Assigned: {model_name}") except Exception as e: if model_name in model_info_dict.keys(): del model_info_dict[model_name] print(f"Failed to assigned: {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 await asyncio.to_thread(load_model, model) except Exception as e: print(e) tasks = [asyncio.create_task(async_load_model(model)) for model in models] return await asyncio.gather(*tasks, return_exceptions=True) def load_models(models: list, limit: int=5): loop = asyncio.new_event_loop() try: loop.run_until_complete(async_load_models(models, limit)) except Exception as e: print(e) pass finally: loop.close() positive_prefix = { "Pony": to_list("score_9, score_8_up, score_7_up"), "Pony Anime": to_list("source_anime, 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), 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]"), "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_tag_type(pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []): tag_type = "danbooru" words = pos_pre + pos_suf + neg_pre + neg_suf for word in words: if "Pony" in word: tag_type = "e621" break return tag_type 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): 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('RGBA') except Exception as e: print(e) return (Image.Image(), None) return (image, caption) async 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)): from tqdm.asyncio import tqdm_asyncio 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) tasks = [asyncio.to_thread(infer, prompt, neg_prompt, model_name) for i in range(image_num)] #results = await asyncio.gather(*tasks, return_exceptions=True) results = await tqdm_asyncio.gather(*tasks) if not results: results = [] for result in results: with lock: if result and result[1]: images.append(result) yield images async def infer_multi_random(prompt: str, neg_prompt: str, results: list, image_num: float, pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], progress=gr.Progress(track_tqdm=True)): from tqdm.asyncio import tqdm_asyncio import random image_num = int(image_num) images = results if results else [] random.seed() model_names = random.choices(list(loaded_models.keys()), k = image_num) prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf) tasks = [asyncio.to_thread(infer, prompt, neg_prompt, model_name) for model_name in model_names] #results = await asyncio.gather(*tasks, return_exceptions=True) results = await tqdm_asyncio.gather(*tasks) if not results: results = [] for result in results: with lock: if result and result[1]: images.append(result) yield images