Upload 2 files

app.py

@@ -48,7 +48,9 @@ load_models(models, 10)
 #load_models(models, 20) # Fetching 20 models at the same time. default: 5
 
 
-css = """"""
+css = """
+#model_info { text-align: center; display:block; }
+"""
 
 with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
     with gr.Column(): 
@@ -61,13 +63,13 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
                 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]))
+            model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]), elem_id="model_info")
         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)
         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)
+            run_button = gr.Button("Generate Image", scale=6)
+            random_button = gr.Button("Random Model 🎲", scale=3)
+            image_num = gr.Number(label="Count", 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)

multit2i.py

@@ -1,8 +1,10 @@
 import gradio as gr
 import asyncio
+from threading import RLock, Thread
 from pathlib import Path
 
 
+lock = RLock()
 loaded_models = {}
 model_info_dict = {}
 
@@ -59,7 +61,8 @@ def get_t2i_model_info_dict(repo_id: str):
     if model.private or model.gated: return info
     try:
         tags = model.tags
-    except Exception:
+    except Exception as e:
+        print(e)
         return info
     if not 'diffusers' in model.tags: return info
     if 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
@@ -74,7 +77,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
 
 
@@ -110,32 +113,35 @@ def load_model(model_name: str):
     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}')
+        with lock:
+            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]
+        with lock:
+            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)
+        with lock:
+            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]
+        with lock:
+            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):
-    from tqdm.asyncio import tqdm_asyncio
+async def async_load_models(models: list, limit: int=5, wait=10):
     sem = asyncio.Semaphore(limit)
     async def async_load_model(model: str):
         async with sem:
            try:
-               return load_model(model)
+               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 tqdm_asyncio.gather(*tasks)
+    return await asyncio.gather(*tasks, return_exceptions=True)
 
 
 def load_models(models: list, limit: int=5):
@@ -271,20 +277,23 @@ def infer(prompt: str, neg_prompt: str, 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
+    #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:
-        images.append(result)
+        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)):
-    import asyncio
+    #from tqdm.asyncio import tqdm_asyncio
     import random
     image_num = int(image_num)
     images = results if results else []
@@ -293,6 +302,10 @@ async def infer_multi_random(prompt: str, neg_prompt: str, results: list, image_
     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)
+    #await tqdm_asyncio.gather(*tasks)
+    if not results: results = []
     for result in results:
-        images.append(result)
+        with lock:
+            if result and result[1]: images.append(result)
         yield images
+