Upload 2 files

app.py

@@ -3,6 +3,7 @@ from multit2i import (
     load_models,
     find_model_list,
     infer_multi,
+    infer_multi_random,
     save_gallery_images,
     change_model,
     get_model_info_md,
@@ -51,19 +52,22 @@ css = """"""
 
 with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
     with gr.Column(): 
-        model_name = gr.Dropdown(label="Select Model", choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0], allow_custom_value=True)
-        model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]))
         with gr.Accordion("Advanced settings", open=False):
-            image_num = gr.Slider(label="Number of Images", minimum=1, maximum=8, value=1, step=1)
             with gr.Accordion("Recommended Prompt"):
                 recom_prompt_preset = gr.Radio(label="Set Presets", choices=get_recom_prompt_type(), value="Common")
                 positive_prefix = gr.CheckboxGroup(label="Use Positive Prefix", choices=get_positive_prefix(), value=[])
                 positive_suffix = gr.CheckboxGroup(label="Use Positive Suffix", choices=get_positive_suffix(), value=["Common"])
                 negative_prefix = gr.CheckboxGroup(label="Use Negative Prefix", choices=get_negative_prefix(), value=[], visible=False)
                 negative_suffix = gr.CheckboxGroup(label="Use Negative Suffix", choices=get_negative_suffix(), value=["Common"], visible=False)
+        with gr.Group():
+            model_name = gr.Dropdown(label="Select Model", choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0])
+            model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]))
         prompt = gr.Text(label="Prompt", lines=1, max_lines=8, placeholder="1girl, solo, ...")
         neg_prompt = gr.Text(label="Negative Prompt", lines=1, max_lines=8, placeholder="", visible=False)
-        run_button = gr.Button("Generate Image")
+        with gr.Row():
+            run_button = gr.Button("Generate Image", scale=2)
+            random_button = gr.Button("Random Model 🎲", scale=1)
+            image_num = gr.Number(label="Number of images", minimum=1, maximum=16, value=1, step=1, interactive=True, scale=1)
         results = gr.Gallery(label="Gallery", interactive=False, show_download_button=True, show_share_button=False,
                               container=True, format="png", object_fit="contain")
         image_files = gr.Files(label="Download", interactive=False)
@@ -99,6 +103,16 @@ This is due to the time it takes for Gradio to generate an example image to cach
         show_progress="full",
         show_api=True,
     ).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
+    gr.on(
+        triggers=[random_button.click],
+        fn=infer_multi_random,
+        inputs=[prompt, neg_prompt, results, image_num,
+                 positive_prefix, positive_suffix, negative_prefix, negative_suffix],
+        outputs=[results],
+        queue=True,
+        show_progress="full",
+        show_api=True,
+    ).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
     clear_results.click(lambda: (None, None), None, [results, image_files], queue=False, show_api=False)
     recom_prompt_preset.change(set_recom_prompt_preset, [recom_prompt_preset],
      [positive_prefix, positive_suffix, negative_prefix, negative_suffix], queue=False, show_api=False)

multit2i.py

@@ -74,7 +74,7 @@ def get_t2i_model_info_dict(repo_id: str):
     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"]})'
+    info["md"] = f'    Model Info: {", ".join(descs)} [Model Repo]({info["url"]})'
     return info
 
 
@@ -126,6 +126,7 @@ def load_model(model_name: str):
 
 
 async def async_load_models(models: list, limit: int=5):
+    from tqdm.asyncio import tqdm_asyncio
     sem = asyncio.Semaphore(limit)
     async def async_load_model(model: str):
         async with sem:
@@ -134,22 +135,23 @@ async def async_load_models(models: list, limit: int=5):
            except Exception as e:
                print(e)
     tasks = [asyncio.create_task(async_load_model(model)) for model in models]
-    return await asyncio.wait(tasks)
+    return await tqdm_asyncio.gather(*tasks)
 
 
 def load_models(models: list, limit: int=5):
-    loop = asyncio.get_event_loop()
+    loop = asyncio.new_event_loop()
     try:
         loop.run_until_complete(async_load_models(models, limit))
     except Exception as e:
         print(e)
         pass
-    loop.close()
+    finally:
+        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"),
+    "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"),
@@ -161,7 +163,7 @@ negative_prefix = {
     "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]"),
+    "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"),
 }
@@ -248,7 +250,7 @@ def change_model(model_name: str):
     return get_model_info_md(model_name)
 
 
-def infer(prompt: str, neg_prompt: str, model_name: str, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt: str, neg_prompt: str, model_name: str):
     from PIL import Image
     import random
     seed = ""
@@ -260,18 +262,37 @@ def infer(prompt: str, neg_prompt: str, model_name: str, progress=gr.Progress(tr
         model = load_model(model_name)
         if not model: return (Image.Image(), None)
         image_path = model(prompt + seed)
-        image = Image.open(image_path).convert('RGB')
+        image = Image.open(image_path).convert('RGBA')
     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,
+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)):
+    import 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)
-    for i in range(image_num):
-        images.append(infer(prompt, neg_prompt, model_name, recom_prompt))
+    tasks = [asyncio.to_thread(infer, prompt, neg_prompt, model_name) for i in range(image_num)]
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+    for result in results:
+        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)):
+    import 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)
+    for result in results:
+        images.append(result)
         yield images