Spaces:
Running
Running
import gradio as gr | |
import asyncio | |
from pathlib import Path | |
loaded_models = {} | |
model_info_dict = {} | |
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, task="text-to-image", 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() | |
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, model_name: str, recom_prompt: bool, progress=gr.Progress(track_tqdm=True)): | |
from PIL import Image | |
import random | |
seed = "" | |
rand = random.randint(1, 500) | |
for i in range(rand): | |
seed += " " | |
rprompt = ", highly detailed, masterpiece, best quality, very aesthetic, absurdres, " if recom_prompt else "" | |
caption = model_name.split("/")[-1] | |
try: | |
model = load_model(model_name) | |
if not model: return (Image.Image(), None) | |
image_path = model(prompt + rprompt + 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, model_name: str, recom_prompt: bool, image_num: float, results: list, progress=gr.Progress(track_tqdm=True)): | |
image_num = int(image_num) | |
images = results if results else [] | |
for i in range(image_num): | |
images.append(infer(prompt, model_name, recom_prompt)) | |
yield images | |