Spaces:
Running
Running
File size: 6,403 Bytes
b384075 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
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
|