Upload 3 files

README.md

@@ -4,7 +4,7 @@ emoji: 🖼️📅
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 5.9.1
+sdk_version: 5.13.1
 app_file: app.py
 pinned: false
 short_description: Text-to-Image

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from model import models
 from multit2i import (load_models, infer_fn, infer_rand_fn, save_gallery,
-    change_model, warm_model, get_model_info_md, loaded_models,
+    change_model, warm_model, get_model_info_md, loaded_models, warm_models,
     get_positive_prefix, get_positive_suffix, get_negative_prefix, get_negative_suffix,
     get_recom_prompt_type, set_recom_prompt_preset, get_tag_type, randomize_seed, translate_to_en)
 from tagger.tagger import (predict_tags_wd, remove_specific_prompt, convert_danbooru_to_e621_prompt,
@@ -14,8 +14,11 @@ from tagger.utils import (V2_ASPECT_RATIO_OPTIONS, V2_RATING_OPTIONS,
 max_images = 6
 MAX_SEED = 2**32-1
 load_models(models)
+warm_models(models[0:max_images])
 
 css = """
+.title { font-size: 3em; align-items: center; text-align: center; }
+.info { align-items: center; text-align: center; }
 .model_info { text-align: center; }
 .output { width=112px; height=112px; max_width=112px; max_height=112px; !important; }
 .gallery { min_width=512px; min_height=512px; max_height=1024px; !important; }
@@ -29,24 +32,24 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
                     with gr.Accordion("Prompt from Image File", open=False):
                         tagger_image = gr.Image(label="Input image", type="pil", format="png", sources=["upload", "clipboard"], height=256)
                         with gr.Accordion(label="Advanced options", open=False):
-                            with gr.Row():
+                            with gr.Row(equal_height=True):
                                 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 common, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
-                            with gr.Row():
+                            with gr.Row(equal_height=True):
                                 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"])
                         tagger_generate_from_image = gr.Button(value="Generate Tags from Image", variant="secondary")
                     with gr.Accordion("Prompt Transformer", open=False):
-                        with gr.Row():
+                        with gr.Row(equal_height=True):
                             v2_character = gr.Textbox(label="Character", placeholder="hatsune miku", scale=2)
                             v2_series = gr.Textbox(label="Series", placeholder="vocaloid", scale=2)
-                        with gr.Row():
+                        with gr.Row(equal_height=True):
                             v2_rating = gr.Radio(label="Rating", choices=list(V2_RATING_OPTIONS), value="sfw")
                             v2_aspect_ratio = gr.Radio(label="Aspect ratio", info="The aspect ratio of the image.", choices=list(V2_ASPECT_RATIO_OPTIONS), value="square", visible=False)
                             v2_length = gr.Radio(label="Length", info="The total length of the tags.", choices=list(V2_LENGTH_OPTIONS), value="long")
-                        with gr.Row():
+                        with gr.Row(equal_height=True):
                             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)
@@ -56,26 +59,26 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
                     prompt = gr.Text(label="Prompt", lines=2, max_lines=8, placeholder="1girl, solo, ...", show_copy_button=True)
                     with gr.Accordion("Advanced options", open=False):
                         neg_prompt = gr.Text(label="Negative Prompt", lines=1, max_lines=8, placeholder="")
-                        with gr.Row():
+                        with gr.Row(equal_height=True):
                             width = gr.Slider(label="Width", info="If 0, the default value is used.", maximum=1216, step=32, value=0)
                             height = gr.Slider(label="Height", info="If 0, the default value is used.", maximum=1216, step=32, value=0)
                             steps = gr.Slider(label="Number of inference steps", info="If 0, the default value is used.", maximum=100, step=1, value=0)
-                        with gr.Row():
+                        with gr.Row(equal_height=True):
                             cfg = gr.Slider(label="Guidance scale", info="If 0, the default value is used.", maximum=30.0, step=0.1, value=0)
                             seed = gr.Slider(label="Seed", info="Randomize Seed if -1.", minimum=-1, maximum=MAX_SEED, step=1, value=-1)
                             seed_rand = gr.Button("Randomize Seed 🎲", size="sm", variant="secondary")
                         recom_prompt_preset = gr.Radio(label="Set Presets", choices=get_recom_prompt_type(), value="Common")
-                        with gr.Row():
+                        with gr.Row(equal_height=True):
                             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=[])
                             negative_suffix = gr.CheckboxGroup(label="Use Negative Suffix", choices=get_negative_suffix(), value=["Common"])
-                    with gr.Row():
+                    with gr.Row(equal_height=True):
                         image_num = gr.Slider(label="Number of images", minimum=1, maximum=max_images, value=1, step=1, interactive=True, scale=2)
                         trans_prompt = gr.Button(value="Translate 📝", variant="secondary", size="sm", scale=2)
                         clear_prompt = gr.Button(value="Clear 🗑️", variant="secondary", size="sm", scale=1)
                     
-                with gr.Row():
+                with gr.Row(equal_height=True):
                     run_button = gr.Button("Generate Image", variant="primary", scale=6)
                     random_button = gr.Button("Random Model 🎲", variant="secondary", scale=3)
                     #stop_button = gr.Button('Stop', variant="stop", interactive=False, scale=1)
@@ -121,7 +124,6 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
                 image_metadata = gr.Image(label="Image with metadata", type="pil", sources=["upload"])
             with gr.Column():
                 result_metadata = gr.Textbox(label="Metadata", show_label=True, show_copy_button=True, interactive=False, container=True, max_lines=99)
-
                 image_metadata.change(
                     fn=extract_exif_data,
                     inputs=[image_metadata],
@@ -132,10 +134,9 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
 [Nymbo/Flood](https://huggingface.co/spaces/Nymbo/Flood), 
 [Yntec/ToyWorldXL](https://huggingface.co/spaces/Yntec/ToyWorldXL), 
 [Yntec/Diffusion80XX](https://huggingface.co/spaces/Yntec/Diffusion80XX).
-            """
-    )
+        """, elem_classes="info")
     gr.DuplicateButton(value="Duplicate Space")
-    gr.Markdown(f"Just a few edits to *model.py* are all it takes to complete your own collection.")
+    gr.Markdown(f"Just a few edits to *model.py* are all it takes to complete your own collection.", elem_classes="info")
 
     #gr.on(triggers=[run_button.click, prompt.submit, random_button.click], fn=lambda: gr.update(interactive=True), inputs=None, outputs=stop_button, show_api=False)
     model_name.change(change_model, [model_name], [model_info], queue=False, show_api=False)\

multit2i.py

@@ -8,7 +8,7 @@ import os
 
 HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None # If private or gated models aren't used, ENV setting is unnecessary.
 server_timeout = 600
-inference_timeout = 300
+inference_timeout = 600
 
 
 lock = RLock()
@@ -52,7 +52,7 @@ def is_loadable(model_name: str, force_gpu: bool = False):
     return status is not None and status.state in ["Loadable", "Loaded"] and (not force_gpu or gpu_state)
 
 
-def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False):
+def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False, public=False):
     from huggingface_hub import HfApi
     api = HfApi(token=HF_TOKEN)
     default_tags = ["diffusers"]
@@ -61,13 +61,13 @@ def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="l
     models = []
     try:
         model_infos = api.list_models(author=author, #task="text-to-image",
-                                       tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)
+                                      tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)
     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 or HF_TOKEN is not None:
+        if not model.private and not model.gated or (HF_TOKEN is not None and not public):
            loadable = is_loadable(model.id, force_gpu) if check_status else True
            if not_tag and not_tag in model.tags or not loadable or "not-for-all-audiences" in model.tags: continue
            models.append(model.id)
@@ -104,8 +104,8 @@ def get_t2i_model_info_dict(repo_id: str):
     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"]})'
+    descs = [info["ver"]] + list_sub(info["tags"], un_tags) + [f'DLs: {info["downloads"]}'] + [f'💕: {info["likes"]}'] + [info["last_modified"]]
+    info["md"] = f'{", ".join(descs)} [Model Repo]({info["url"]})'
     return info
 
 
@@ -160,8 +160,9 @@ def load_from_model(model_name: str, hf_token: str | Literal[False] | None = Non
     p = response.json().get("pipeline_tag")
     if p != "text-to-image": raise ModelNotFoundError(f"This model isn't for text-to-image or unsupported: {model_name}.")
     headers["X-Wait-For-Model"] = "true"
-    client = huggingface_hub.InferenceClient(model=model_name, headers=headers,
-                                              token=hf_token, timeout=server_timeout)
+    kwargs = {}
+    if hf_token is not None: kwargs["token"] = hf_token
+    client = huggingface_hub.InferenceClient(model=model_name, headers=headers, timeout=server_timeout, **kwargs)
     inputs = gr.components.Textbox(label="Input")
     outputs = gr.components.Image(label="Output")
     fn = client.text_to_image
@@ -170,9 +171,10 @@ def load_from_model(model_name: str, hf_token: str | Literal[False] | None = Non
         try:
             data = fn(*data, **kwargs)  # type: ignore
         except huggingface_hub.utils.HfHubHTTPError as e:
-            if "429" in str(e):
-                raise TooManyRequestsError() from e
+            print(e)
+            if "429" in str(e): raise TooManyRequestsError() from e
         except Exception as e:
+            print(e)
             raise Exception() from e
         return data
 
@@ -210,29 +212,29 @@ def load_model(model_name: str):
 def load_model_api(model_name: str):
     global loaded_models
     global model_info_dict
-    if model_name in loaded_models.keys(): return loaded_models[model_name]
     try:
-        client = InferenceClient(timeout=5)
-        status = client.get_model_status(model_name, token=HF_TOKEN)
-        if status is None or status.framework != "diffusers" or status.state not in ["Loadable", "Loaded"]:
+        loaded = False
+        client = InferenceClient(timeout=5, token=HF_TOKEN)
+        status = client.get_model_status(model_name)
+        if status is None or status.framework != "diffusers" or not status.loaded or status.state not in ["Loadable", "Loaded"]:
             print(f"Failed to load by API: {model_name}")
-            return None
         else:
-            loaded_models[model_name] = InferenceClient(model_name, token=HF_TOKEN, timeout=server_timeout)
+            loaded_models[model_name] = InferenceClient(model_name, timeout=server_timeout)
+            loaded = True
             print(f"Loaded by API: {model_name}")
     except Exception as e:
-        if model_name in loaded_models.keys(): del loaded_models[model_name]
         print(f"Failed to load by API: {model_name}")
         print(e)
-        return None
+        loaded = False
     try:
-        model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
-        print(f"Assigned by API: {model_name}")
+        if loaded:
+            model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
+            print(f"Assigned by API: {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 by API: {model_name}")
         print(e)
-    return loaded_models[model_name]
+    return loaded_models[model_name] if model_name in loaded_models.keys() else None
 
 
 def load_models(models: list):
@@ -270,8 +272,8 @@ 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 = to_list(prompt) if prompt else []
+    neg_prompts = to_list(neg_prompt) if neg_prompt else []
     prompts = list_sub(prompts, positive_all)
     neg_prompts = list_sub(neg_prompts, negative_all)
     last_empty_p = [""] if not prompts and type != "None" else []
@@ -287,7 +289,6 @@ def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos
 
 recom_prompt_type = {
     "None": ([], [], [], []),
-    "Auto": ([], [], [], []),
     "Common": ([], ["Common"], [], ["Common"]),
     "Animagine": ([], ["Common", "Anime"], [], ["Common"]),
     "Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]),
@@ -296,11 +297,7 @@ recom_prompt_type = {
 }
 
 
-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)
 
@@ -311,9 +308,7 @@ def set_recom_prompt_preset(type: str = "None"):
 
 
 def get_recom_prompt_type():
-    type = list(recom_prompt_type.keys())
-    type.remove("Auto")
-    return type
+    return list(recom_prompt_type.keys())
 
 
 def get_positive_prefix():
@@ -356,11 +351,16 @@ def warm_model(model_name: str):
     if model:
         try:
             print(f"Warming model: {model_name}")
-            infer_body(model, " ")
+            infer_body(model, model_name, " ")
         except Exception as e:
             print(e)
 
 
+def warm_models(models: list[str]):
+    for model in models:
+        asyncio.new_event_loop().run_in_executor(None, warm_model, model)
+
+
 # https://huggingface.co/docs/api-inference/detailed_parameters
 # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
 def infer_body(client: InferenceClient | gr.Interface | object, model_str: str, prompt: str, neg_prompt: str = "",
@@ -375,21 +375,22 @@ def infer_body(client: InferenceClient | gr.Interface | object, model_str: str,
     else: kwargs["seed"] = seed
     try:
         if isinstance(client, InferenceClient):
-            image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
+            image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
         elif isinstance(client, gr.Interface):
-            image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
+            if HF_TOKEN is not None: kwargs["token"] = HF_TOKEN
+            image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
         else: return None
         if isinstance(image, tuple): return None
         return save_image(image, png_path, model_str, prompt, neg_prompt, height, width, steps, cfg, seed)
     except Exception as e:
         print(e)
-        raise Exception() from e
+        raise Exception(e) from e
 
 
 async def infer(model_name: str, prompt: str, neg_prompt: str ="", height: int = 0, width: int = 0,
                steps: int = 0, cfg: int = 0, seed: int = -1,
                save_path: str | None = None, timeout: float = inference_timeout):
-    model = load_model(model_name)
+    model = load_model_api(model_name)
     if not model: return None
     task = asyncio.create_task(asyncio.to_thread(infer_body, model, model_name, prompt, neg_prompt,
                                                  height, width, steps, cfg, seed))
@@ -406,7 +407,7 @@ async def infer(model_name: str, prompt: str, neg_prompt: str ="", height: int =
         print(e)
         if not task.done(): task.cancel()
         result = None
-        raise Exception() from e
+        raise Exception(e) from e
     if task.done() and result is not None:
         with lock:
             image = rename_image(result, model_name, save_path)
@@ -418,8 +419,7 @@ async def infer(model_name: str, prompt: str, neg_prompt: str ="", height: int =
 def infer_fn(model_name: str, prompt: str, neg_prompt: str = "", height: int = 0, width: int = 0,
              steps: int = 0, cfg: int = 0, seed: int = -1,
              pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
-    if model_name == 'NA':
-        return None
+    if model_name in ["NA", ""]: return gr.update()
     try:
         loop = asyncio.get_running_loop()
     except Exception:
@@ -442,8 +442,7 @@ def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str = "", heig
              steps: int = 0, cfg: int = 0, seed: int = -1,
              pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
     import random
-    if model_name_dummy == 'NA':
-        return None
+    if model_name_dummy in ["NA", ""]: return gr.update()
     random.seed()
     model_name = random.choice(list(loaded_models.keys()))
     try: