File size: 13,744 Bytes
5fbe98e
 
cc6e31a
 
5fbe98e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cc6e31a
5fbe98e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cc6e31a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5fbe98e
 
 
 
 
cc6e31a
5fbe98e
 
cc6e31a
5fbe98e
 
 
 
cc6e31a
 
5fbe98e
cc6e31a
 
5fbe98e
 
 
 
cc6e31a
5fbe98e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cc6e31a
 
 
 
 
 
 
 
 
 
5fbe98e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cc6e31a
5fbe98e
 
 
 
cc6e31a
 
5fbe98e
 
 
 
 
 
 
 
 
cc6e31a
5fbe98e
 
 
 
 
 
 
cc6e31a
 
5fbe98e
 
 
 
 
 
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
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
import gradio as gr
import asyncio
import queue
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