Upload 11 files

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+from model import models
 from multit2i import (
     load_models,
     infer_multi,
@@ -13,9 +14,8 @@ from multit2i import (
     get_negative_suffix,
     get_recom_prompt_type,
     set_recom_prompt_preset,
+    get_tag_type,
 )
-from model import models
-
 from tagger.tagger import (
     predict_tags_wd,
     remove_specific_prompt,
@@ -35,8 +35,8 @@ from tagger.utils import (
 )
 
 
-load_models(models, 10)
-#load_models(models, 20) # Fetching 20 models at the same time. default: 5 *This option is not working so far.
+load_models(models, 5)
+#load_models(models, 20) # Fetching 20 models at the same time. default: 5
 
 
 css = """
@@ -59,6 +59,7 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
                 v2_length = gr.Radio(label="Length", info="The total length of the tags.", choices=list(V2_LENGTH_OPTIONS), value="long")
                 v2_identity = gr.Radio(label="Keep identity", info="How strictly to keep the identity of the character or subject. If you specify the detail of subject in the prompt, you should choose `strict`. Otherwise, choose `none` or `lax`. `none` is very creative but sometimes ignores the input prompt.", choices=list(V2_IDENTITY_OPTIONS), value="lax")                    
                 v2_ban_tags = gr.Textbox(label="Ban tags", info="Tags to ban from the output.", placeholder="alternate costumen, ...", value="censored")
+                v2_tag_type = gr.Radio(label="Tag Type", info="danbooru for common, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru", visible=False)
                 v2_model = gr.Dropdown(label="Model", choices=list(V2_ALL_MODELS.keys()), value=list(V2_ALL_MODELS.keys())[0])
             with gr.Accordion("Model", open=True):
                 model_name = gr.Dropdown(label="Select Model", choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0])
@@ -69,7 +70,7 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
                 with gr.Accordion(label="Advanced options", open=False):
                     tagger_general_threshold = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True)
                     tagger_character_threshold = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
-                    tagger_tag_type = gr.Radio(label="Convert tags to", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
+                    tagger_tag_type = gr.Radio(label="Convert tags to", info="danbooru for common, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
                     tagger_recom_prompt = gr.Radio(label="Insert reccomended prompt", choices=["None", "Animagine", "Pony"], value="None", interactive=True)  
                     tagger_keep_tags = gr.Radio(label="Remove tags leaving only the following", choices=["body", "dress", "all"], value="all")
                 tagger_algorithms = gr.CheckboxGroup(["Use WD Tagger", "Use Florence-2-SD3-Long-Captioner"], label="Algorithms", value=["Use WD Tagger"])
@@ -103,8 +104,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")
@@ -134,8 +133,14 @@ This is due to the time it takes for Gradio to generate an example image to cach
     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)
-    random_prompt.click(v2_random_prompt, [prompt, v2_series, v2_character, v2_rating, v2_aspect_ratio, v2_length,
-                                           v2_identity, v2_ban_tags, v2_model], [prompt, v2_series, v2_character], queue=False, show_api=False)
+    random_prompt.click(
+        v2_random_prompt, [prompt, v2_series, v2_character, v2_rating, v2_aspect_ratio, v2_length,
+          v2_identity, v2_ban_tags, v2_model], [prompt, v2_series, v2_character], show_api=False,
+    ).success(
+        get_tag_type, [positive_prefix, positive_suffix, negative_prefix, negative_suffix], [v2_tag_type], queue=False, show_api=False
+    ).success(
+        convert_danbooru_to_e621_prompt, [prompt, v2_tag_type], [prompt], queue=False, show_api=False,
+    )
     tagger_generate_from_image.click(
         predict_tags_wd,
         [tagger_image, prompt, tagger_algorithms, tagger_general_threshold, tagger_character_threshold],

model.py

@@ -19,9 +19,6 @@ models = [
 ]
 
 
-models = ['yodayo-ai/kivotos-xl-2.0', 'Raelina/Rae-Diffusion-XL-V2']
-
-
 # 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
@@ -29,4 +26,3 @@ models = ['yodayo-ai/kivotos-xl-2.0', 'Raelina/Rae-Diffusion-XL-V2']
 #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,7 @@
 import gradio as gr
 import asyncio
-from threading import RLock, Thread
+import queue
+from threading import RLock
 from pathlib import Path
 
 
@@ -70,8 +71,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 +108,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 +277,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 +317,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 +333,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 +341,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:

tagger/tagger.py

@@ -31,12 +31,15 @@ PEOPLE_TAGS = (
 
 
 RATING_MAP = {
+    "sfw": "safe",
     "general": "safe",
     "sensitive": "sensitive",
     "questionable": "nsfw",
     "explicit": "explicit, nsfw",
 }
 DANBOORU_TO_E621_RATING_MAP = {
+    "sfw": "rating_safe",
+    "general": "rating_safe",
     "safe": "rating_safe",
     "sensitive": "rating_safe",
     "nsfw": "rating_explicit",