Upload 3 files

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
+from model import models
 from multit2i import (
     load_models,
-    find_model_list,
     infer_multi,
     infer_multi_random,
     save_gallery_images,
@@ -14,37 +14,11 @@ from multit2i import (
     get_negative_suffix,
     get_recom_prompt_type,
     set_recom_prompt_preset,
+    get_tag_type,
 )
 
 
-models = [
-    'yodayo-ai/kivotos-xl-2.0',
-    'yodayo-ai/holodayo-xl-2.1',
-    'cagliostrolab/animagine-xl-3.1',
-    'votepurchase/ponyDiffusionV6XL',
-    'eienmojiki/Anything-XL',
-    'eienmojiki/Starry-XL-v5.2',
-    'digiplay/majicMIX_sombre_v2',
-    'digiplay/majicMIX_realistic_v7',
-    'votepurchase/counterfeitV30_v30',
-    'Meina/MeinaMix_V11',
-    'KBlueLeaf/Kohaku-XL-Epsilon-rev3',
-    'kayfahaarukku/UrangDiffusion-1.1',
-    'Raelina/Rae-Diffusion-XL-V2',
-    'Raelina/Raemu-XL-V4',
-]
-
-
-# Examples:
-#models = ['yodayo-ai/kivotos-xl-2.0', 'yodayo-ai/holodayo-xl-2.1'] # specific models
-#models = find_model_list("John6666", [], "", "last_modified", 20) # John6666's latest 20 models
-#models = find_model_list("John6666", ["anime"], "", "last_modified", 20) # John6666's latest 20 models with 'anime' tag
-#models = find_model_list("John6666", [], "anime", "last_modified", 20) # John6666's latest 20 models without 'anime' tag
-#models = find_model_list("", [], "", "last_modified", 20) # latest 20 text-to-image models of huggingface
-#models = find_model_list("", [], "", "downloads", 20) # monthly most downloaded 20 text-to-image models of huggingface
-
-
-load_models(models, 10)
+load_models(models, 5)
 #load_models(models, 20) # Fetching 20 models at the same time. default: 5
 
 
@@ -54,18 +28,21 @@ css = """
 
 with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
     with gr.Column(): 
-        with gr.Accordion("Advanced settings", open=False):
-            with gr.Accordion("Recommended Prompt"):
+        with gr.Accordion("Advanced settings", open=True):
+            with gr.Accordion("Recommended Prompt", open=False):
                 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.Row():
+                    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.Accordion("Model", open=True):
+                model_name = gr.Dropdown(label="Select Model", show_label=False, 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]), elem_id="model_info")
         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]), 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)
+            clear_prompt = gr.Button(value="Clear Prompt 🗑️", size="sm", scale=1)
+            prompt = gr.Text(label="Prompt", lines=1, max_lines=8, placeholder="1girl, solo, ...", show_copy_button=True)
+            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=6)
             random_button = gr.Button("Random Model 🎲", scale=3)
@@ -88,8 +65,6 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
         f"""This demo was created in reference to the following demos.
 - [Nymbo/Flood](https://huggingface.co/spaces/Nymbo/Flood).
 - [Yntec/ToyWorldXL](https://huggingface.co/spaces/Yntec/ToyWorldXL).
-<br>The first startup takes a mind-boggling amount of time, but not so much after the second.
-This is due to the time it takes for Gradio to generate an example image to cache.
             """
     )
     gr.DuplicateButton(value="Duplicate Space")
@@ -115,6 +90,7 @@ 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)
+    clear_prompt.click(lambda: None, None, [prompt], 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)

model.py

@@ -0,0 +1,28 @@
+from multit2i import find_model_list
+
+
+models = [
+    'yodayo-ai/kivotos-xl-2.0',
+    'yodayo-ai/holodayo-xl-2.1',
+    'cagliostrolab/animagine-xl-3.1',
+    'votepurchase/ponyDiffusionV6XL',
+    'eienmojiki/Anything-XL',
+    'eienmojiki/Starry-XL-v5.2',
+    'digiplay/majicMIX_sombre_v2',
+    'digiplay/majicMIX_realistic_v7',
+    'votepurchase/counterfeitV30_v30',
+    'Meina/MeinaMix_V11',
+    'KBlueLeaf/Kohaku-XL-Epsilon-rev3',
+    'kayfahaarukku/UrangDiffusion-1.1',
+    'Raelina/Rae-Diffusion-XL-V2',
+    'Raelina/Raemu-XL-V4',
+]
+
+
+# Examples:
+#models = ['yodayo-ai/kivotos-xl-2.0', 'yodayo-ai/holodayo-xl-2.1'] # specific models
+#models = find_model_list("John6666", [], "", "last_modified", 20) # John6666's latest 20 models
+#models = find_model_list("John6666", ["anime"], "", "last_modified", 20) # John6666's latest 20 models with 'anime' tag
+#models = find_model_list("John6666", [], "anime", "last_modified", 20) # John6666's latest 20 models without 'anime' tag
+#models = find_model_list("", [], "", "last_modified", 20) # latest 20 text-to-image models of huggingface
+#models = find_model_list("", [], "", "downloads", 20) # monthly most downloaded 20 text-to-image models of huggingface

multit2i.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import asyncio
-from threading import RLock, Thread
+from threading import RLock
 from pathlib import Path
 
 
@@ -70,8 +70,7 @@ def get_t2i_model_info_dict(repo_id: str):
     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["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")
@@ -108,31 +107,61 @@ def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
     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:
-        with lock:
-            loaded_models[model_name] = gr.load(f'models/{model_name}')
+        loaded_models[model_name] = load_from_model(model_name)
         print(f"Loaded: {model_name}")
     except Exception as e:
-        with lock:
-            if model_name in loaded_models.keys(): del loaded_models[model_name]
+        if model_name in loaded_models.keys(): del loaded_models[model_name]
         print(f"Failed to load: {model_name}")
         print(e)
         return None
     try:
-        with lock:
-            model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
+        model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
+        print(f"Assigned: {model_name}")
     except Exception as e:
-        with lock:
-            if model_name in model_info_dict.keys(): del model_info_dict[model_name]
+        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, wait=10):
+async def async_load_models(models: list, limit: int=5):
     sem = asyncio.Semaphore(limit)
     async def async_load_model(model: str):
         async with sem:
@@ -247,6 +276,16 @@ 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", "")
 
@@ -277,13 +316,13 @@ 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)):
-    #from tqdm.asyncio import tqdm_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)
+    #results = await asyncio.gather(*tasks, return_exceptions=True)
+    results = await tqdm_asyncio.gather(*tasks)
     if not results: results = []
     for result in results:
         with lock:
@@ -293,7 +332,7 @@ async def infer_multi(prompt: str, neg_prompt: str, results: list, image_num: fl
 
 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
+    from tqdm.asyncio import tqdm_asyncio
     import random
     image_num = int(image_num)
     images = results if results else []
@@ -301,8 +340,8 @@ async def infer_multi_random(prompt: str, neg_prompt: str, results: list, image_
     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)
-    #await tqdm_asyncio.gather(*tasks)
+    #results = await asyncio.gather(*tasks, return_exceptions=True)
+    results = await tqdm_asyncio.gather(*tasks)
     if not results: results = []
     for result in results:
         with lock: