Upload 10 files

app.py

@@ -1,92 +1,94 @@
 import gradio as gr
 from model import models
-from multit2i import (
-    load_models, infer_fn, infer_rand_fn, save_gallery,
+from multit2i import (load_models, infer_fn, infer_rand_fn, save_gallery,
     change_model, warm_model, get_model_info_md, loaded_models,
     get_positive_prefix, get_positive_suffix, get_negative_prefix, get_negative_suffix,
-    get_recom_prompt_type, set_recom_prompt_preset, get_tag_type,
-)
-from tagger.tagger import (
-    predict_tags_wd, remove_specific_prompt, convert_danbooru_to_e621_prompt,
-    insert_recom_prompt, compose_prompt_to_copy,
-)
+    get_recom_prompt_type, set_recom_prompt_preset, get_tag_type)
+from tagger.tagger import (predict_tags_wd, remove_specific_prompt, convert_danbooru_to_e621_prompt,
+    insert_recom_prompt, compose_prompt_to_copy)
 from tagger.fl2sd3longcap import predict_tags_fl2_sd3
 from tagger.v2 import V2_ALL_MODELS, v2_random_prompt
-from tagger.utils import (
-    V2_ASPECT_RATIO_OPTIONS, V2_RATING_OPTIONS,
-    V2_LENGTH_OPTIONS, V2_IDENTITY_OPTIONS,
-)
+from tagger.utils import (V2_ASPECT_RATIO_OPTIONS, V2_RATING_OPTIONS,
+    V2_LENGTH_OPTIONS, V2_IDENTITY_OPTIONS)
 
-
-max_images = 8
+max_images = 6
+MAX_SEED = 2**32-1
 load_models(models)
 
 css = """
 .model_info { text-align: center; }
-.output { width=112px; height=112px; !important; }
-.gallery { width=100%; min_height=768px; !important; }
+.output { width=112px; height=112px; max_width=112px; max_height=112px; !important; }
+.gallery { min_width=512px; min_height=512px; max_height=1024px; !important; }
 """
 
 with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
-    with gr.Column(): 
-        with gr.Group():
-            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]), elem_classes="model_info")
-        with gr.Group():
-            with gr.Accordion("Prompt from Image File", open=False):
-                tagger_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
-                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 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"])
-                tagger_generate_from_image = gr.Button(value="Generate Tags from Image")
-            with gr.Row():
-                v2_character = gr.Textbox(label="Character", placeholder="hatsune miku", scale=2)
-                v2_series = gr.Textbox(label="Series", placeholder="vocaloid", scale=2)
-                random_prompt = gr.Button(value="Extend Prompt 🎲", size="sm", scale=1)
-                clear_prompt = gr.Button(value="Clear Prompt 🗑️", size="sm", scale=1)
-            prompt = gr.Text(label="Prompt", lines=2, 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.Accordion("Advanced options", open=False):
-                width = gr.Number(label="Width", info="If 0, the default value is used.", maximum=1216, step=32, value=0)
-                height = gr.Number(label="Height", info="If 0, the default value is used.", maximum=1216, step=32, value=0)
-                steps = gr.Number(label="Number of inference steps", info="If 0, the default value is used.", maximum=100, step=1, value=0)
-                cfg = gr.Number(label="Guidance scale", info="If 0, the default value is used.", maximum=30.0, step=0.1, value=0)
-            with gr.Accordion("Recommended Prompt", open=False):
-                recom_prompt_preset = gr.Radio(label="Set Presets", choices=get_recom_prompt_type(), value="Common")
+    with gr.Row():
+        with gr.Column(scale=10): 
+            with gr.Group():
+                with gr.Accordion("Prompt from Image File", open=False):
+                    tagger_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
+                    with gr.Accordion(label="Advanced options", open=False):
+                        with gr.Row():
+                            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():
+                            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():
+                        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():
+                        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])
+                    v2_copy = gr.Button(value="Copy to clipboard", variant="secondary", size="sm", interactive=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=[])
-                    negative_suffix = gr.CheckboxGroup(label="Use Negative Suffix", choices=get_negative_suffix(), value=["Common"])
-            with gr.Accordion("Prompt Transformer", open=False):
-                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")
-                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])
-                v2_copy = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
-        image_num = gr.Slider(label="Number of images", minimum=1, maximum=max_images, value=1, step=1, interactive=True, scale=1)
-        with gr.Row():
-            run_button = gr.Button("Generate Image", scale=6)
-            random_button = gr.Button("Random Model 🎲", scale=3)
-            stop_button = gr.Button('Stop', interactive=False, scale=1)
-    with gr.Column(): 
-        with gr.Group():
+                    v2_character = gr.Textbox(label="Character", placeholder="hatsune miku", scale=2)
+                    v2_series = gr.Textbox(label="Series", placeholder="vocaloid", scale=2)
+                    random_prompt = gr.Button(value="Extend Prompt 🎲", variant="secondary", size="sm", scale=1)
+                    clear_prompt = gr.Button(value="Clear Prompt 🗑️", variant="secondary", size="sm", scale=1)
+                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():
+                        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)
+                    with gr.Row():
+                        steps = gr.Slider(label="Number of inference steps", info="If 0, the default value is used.", maximum=100, step=1, value=0)
+                        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)
+                    recom_prompt_preset = gr.Radio(label="Set Presets", choices=get_recom_prompt_type(), value="Common")
+                    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=[])
+                        negative_suffix = gr.CheckboxGroup(label="Use Negative Suffix", choices=get_negative_suffix(), value=["Common"])
+
+                image_num = gr.Slider(label="Number of images", minimum=1, maximum=max_images, value=1, step=1, interactive=True, scale=1)
             with gr.Row():
-                output = [gr.Image(label='', elem_classes="output", type="filepath", format=".png",
-                          show_download_button=True, show_share_button=False, show_label=False,
-                          interactive=False, min_width=80, visible=True) for _ in range(max_images)]
-        with gr.Group():
-            results = gr.Gallery(label="Gallery", elem_classes="gallery", interactive=False, show_download_button=True, show_share_button=False,
-                                 container=True, format="png", object_fit="cover", columns=2, rows=2)
-            image_files = gr.Files(label="Download", interactive=False)
-            clear_results = gr.Button("Clear Gallery / Download 🗑️")
+                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)
+            with gr.Group():
+                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]), elem_classes="model_info")
+        with gr.Column(scale=10): 
+            with gr.Group():
+                with gr.Row():
+                    output = [gr.Image(label='', elem_classes="output", type="filepath", format="png",
+                            show_download_button=True, show_share_button=False, show_label=False,
+                            interactive=False, min_width=80, visible=True) for _ in range(max_images)]
+            with gr.Group():
+                results = gr.Gallery(label="Gallery", elem_classes="gallery", interactive=False, show_download_button=True, show_share_button=False,
+                                    container=True, format="png", object_fit="cover", columns=2, rows=2)
+                image_files = gr.Files(label="Download", interactive=False)
+                clear_results = gr.Button("Clear Gallery / Download 🗑️", variant="secondary")
     with gr.Column(): 
         examples = gr.Examples(
             examples = [
@@ -115,13 +117,13 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
         img_i = gr.Number(i, visible=False)
         image_num.change(lambda i, n: gr.update(visible = (i < n)), [img_i, image_num], o, show_api=False)
         gen_event = gr.on(triggers=[run_button.click, prompt.submit],
-         fn=lambda i, n, m, t1, t2, n1, n2, n3, n4, l1, l2, l3, l4: infer_fn(m, t1, t2, n1, n2, n3, n4, l1, l2, l3, l4) if (i < n) else None,
-         inputs=[img_i, image_num, model_name, prompt, neg_prompt, height, width, steps, cfg,
+         fn=lambda i, n, m, t1, t2, n1, n2, n3, n4, n5, l1, l2, l3, l4: infer_fn(m, t1, t2, n1, n2, n3, n4, n5, l1, l2, l3, l4) if (i < n) else None,
+         inputs=[img_i, image_num, model_name, prompt, neg_prompt, height, width, steps, cfg, seed,
                   positive_prefix, positive_suffix, negative_prefix, negative_suffix],
          outputs=[o], queue=True, show_api=False)
         gen_event2 = gr.on(triggers=[random_button.click],
-         fn=lambda i, n, m, t1, t2, n1, n2, n3, n4, l1, l2, l3, l4: infer_rand_fn(m, t1, t2, n1, n2, n3, n4, l1, l2, l3, l4) if (i < n) else None,
-         inputs=[img_i, image_num, model_name, prompt, neg_prompt, height, width, steps, cfg,
+         fn=lambda i, n, m, t1, t2, n1, n2, n3, n4, n5, l1, l2, l3, l4: infer_rand_fn(m, t1, t2, n1, n2, n3, n4, n5, l1, l2, l3, l4) if (i < n) else None,
+         inputs=[img_i, image_num, model_name, prompt, neg_prompt, height, width, steps, cfg, seed,
                   positive_prefix, positive_suffix, negative_prefix, negative_suffix],
          outputs=[o], queue=True, show_api=False)
         o.change(save_gallery, [o, results], [results, image_files], show_api=False)
@@ -135,29 +137,19 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
     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(
-            lambda: ("", "", ""), None, [v2_series, v2_character, prompt], queue=False, 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(lambda: ("", "", ""), None, [v2_series, v2_character, prompt], queue=False, show_api=False,
     ).success(
         predict_tags_wd,
         [tagger_image, prompt, tagger_algorithms, tagger_general_threshold, tagger_character_threshold],
         [v2_series, v2_character, prompt, v2_copy],
         show_api=False,
-    ).success(
-        predict_tags_fl2_sd3, [tagger_image, prompt, tagger_algorithms], [prompt], show_api=False,
-    ).success(
-        remove_specific_prompt, [prompt, tagger_keep_tags], [prompt], queue=False, show_api=False,
-    ).success(
-        convert_danbooru_to_e621_prompt, [prompt, tagger_tag_type], [prompt], queue=False, show_api=False,
-    ).success(
-        insert_recom_prompt, [prompt, neg_prompt, tagger_recom_prompt], [prompt, neg_prompt], queue=False, show_api=False,
-    ).success(
-        compose_prompt_to_copy, [v2_character, v2_series, prompt], [prompt], queue=False, show_api=False,
-    )
+    ).success(predict_tags_fl2_sd3, [tagger_image, prompt, tagger_algorithms], [prompt], show_api=False,
+    ).success(remove_specific_prompt, [prompt, tagger_keep_tags], [prompt], queue=False, show_api=False,
+    ).success(convert_danbooru_to_e621_prompt, [prompt, tagger_tag_type], [prompt], queue=False, show_api=False,
+    ).success(insert_recom_prompt, [prompt, neg_prompt, tagger_recom_prompt], [prompt, neg_prompt], queue=False, show_api=False,
+    ).success(compose_prompt_to_copy, [v2_character, v2_series, prompt], [prompt], queue=False, show_api=False)
 
-demo.queue()
-demo.launch()
+demo.queue(default_concurrency_limit=200, max_size=200)
+demo.launch(max_threads=400)

multit2i.py

@@ -3,8 +3,10 @@ import asyncio
 from threading import RLock
 from pathlib import Path
 from huggingface_hub import InferenceClient
+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
 
@@ -31,22 +33,43 @@ def is_repo_name(s):
     return re.fullmatch(r'^[^/]+?/[^/]+?$', s)
 
 
-def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30):
+def get_status(model_name: str):
+    from huggingface_hub import InferenceClient
+    client = InferenceClient(token=HF_TOKEN, timeout=10)
+    return client.get_model_status(model_name)
+
+
+def is_loadable(model_name: str, force_gpu: bool = False):
+    try:
+        status = get_status(model_name)
+    except Exception as e:
+        print(e)
+        print(f"Couldn't load {model_name}.")
+        return False
+    gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()
+    if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):
+        print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")
+    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):
     from huggingface_hub import HfApi
-    api = HfApi()
+    api = HfApi(token=HF_TOKEN)
     default_tags = ["diffusers"]
     if not sort: sort = "last_modified"
+    limit = limit * 20 if check_status and force_gpu else limit * 5
     models = []
     try:
-        model_infos = api.list_models(author=author, pipeline_tag="text-to-image",
-                                       tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit * 5)
+        model_infos = api.list_models(author=author, task="text-to-image",
+                                       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:
-           if not_tag and not_tag in model.tags: continue
+        if not model.private and not model.gated or HF_TOKEN is not None:
+           loadable = is_loadable(model.id, force_gpu) if check_status else True
+           if not_tag and not_tag in model.tags or not loadable: continue
            models.append(model.id)
            if len(models) == limit: break
     return models
@@ -54,23 +77,24 @@ def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="l
 
 def get_t2i_model_info_dict(repo_id: str):
     from huggingface_hub import HfApi
-    api = HfApi()
+    api = HfApi(token=HF_TOKEN)
     info = {"md": "None"}
     try:
         if not is_repo_name(repo_id) or not api.repo_exists(repo_id=repo_id): return info
-        model = api.model_info(repo_id=repo_id)
+        model = api.model_info(repo_id=repo_id, token=HF_TOKEN)
     except Exception as e:
         print(f"Error: Failed to get {repo_id}'s info.")
         print(e)
         return info
-    if model.private or model.gated: return info
+    if model.private or model.gated and HF_TOKEN is None: return info
     try:
         tags = model.tags
     except Exception as e:
         print(e)
         return info
     if not 'diffusers' in model.tags: return info
-    if 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
+    if 'diffusers:FluxPipeline' in tags: info["ver"] = "FLUX.1"
+    elif 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
     elif 'diffusers:StableDiffusionPipeline' in tags: info["ver"] = "SD1.5"
     elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3"
     else: info["ver"] = "Other"
@@ -118,20 +142,23 @@ def save_gallery(image_path: str | None, images: list[tuple] | None):
 
 # https://github.com/gradio-app/gradio/blob/main/gradio/external.py
 # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
-def load_from_model(model_name: str, hf_token: str = None):
+from typing import Literal
+def load_from_model(model_name: str, hf_token: str | Literal[False] | None = None):
     import httpx
     import huggingface_hub
-    from gradio.exceptions import ModelNotFoundError
+    from gradio.exceptions import ModelNotFoundError, TooManyRequestsError
     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 {}
+    headers = ({} if hf_token in [False, None] else {"Authorization": f"Bearer {hf_token}"})
     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."
         )
+    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)
@@ -140,7 +167,14 @@ def load_from_model(model_name: str, hf_token: str = None):
     fn = client.text_to_image
 
     def query_huggingface_inference_endpoints(*data, **kwargs):
-        return fn(*data, **kwargs)
+        try:
+            data = fn(*data, **kwargs)  # type: ignore
+        except huggingface_hub.utils.HfHubHTTPError as e:
+            if "429" in str(e):
+                raise TooManyRequestsError() from e
+        except Exception as e:
+            raise Exception() from e
+        return data
 
     interface_info = {
         "fn": query_huggingface_inference_endpoints,
@@ -156,7 +190,7 @@ 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] = load_from_model(model_name)
+        loaded_models[model_name] = load_from_model(model_name, hf_token=HF_TOKEN)
         print(f"Loaded: {model_name}")
     except Exception as e:
         if model_name in loaded_models.keys(): del loaded_models[model_name]
@@ -179,12 +213,12 @@ def load_model_api(model_name: str):
     if model_name in loaded_models.keys(): return loaded_models[model_name]
     try:
         client = InferenceClient(timeout=5)
-        status = client.get_model_status(model_name)
+        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"]:
             print(f"Failed to load by API: {model_name}")
             return None
         else:
-            loaded_models[model_name] = InferenceClient(model_name, timeout=server_timeout)
+            loaded_models[model_name] = InferenceClient(model_name, token=HF_TOKEN, timeout=server_timeout)
             print(f"Loaded by API: {model_name}")
     except Exception as e:
         if model_name in loaded_models.keys(): del loaded_models[model_name]
@@ -329,49 +363,58 @@ def warm_model(model_name: str):
 
 # 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, prompt: str, neg_prompt: str | None = None,
+def infer_body(client: InferenceClient | gr.Interface | object, prompt: str, neg_prompt: str | None = None,
                height: int | None = None, width: int | None = None,
-               steps: int | None = None, cfg: int | None = None):
+               steps: int | None = None, cfg: int | None = None, seed: int = -1):
     png_path = "image.png"
     kwargs = {}
     if height is not None and height >= 256: kwargs["height"] = height
     if width is not None and width >= 256: kwargs["width"] = width
     if steps is not None and steps >= 1: kwargs["num_inference_steps"] = steps
     if cfg is not None and cfg > 0: cfg = kwargs["guidance_scale"] = cfg
+    if seed >= 0: kwargs["seed"] = seed
     try:
         if isinstance(client, InferenceClient):
-            image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
+            image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
         elif isinstance(client, gr.Interface):
-            image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
+            image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
         else: return None
+        if isinstance(image, tuple): return None
         image.save(png_path)
         return str(Path(png_path).resolve())
     except Exception as e:
         print(e)
-        return None
+        raise Exception() from e
 
 
 async def infer(model_name: str, prompt: str, neg_prompt: str | None = None,
                height: int | None = None, width: int | None = None,
-               steps: int | None = None, cfg: int | None = None,
+               steps: int | None = None, cfg: int | None = None, seed: int = -1,
                save_path: str | None = None, timeout: float = inference_timeout):
     import random
     noise = ""
-    rand = random.randint(1, 500)
-    for i in range(rand):
-        noise += " "
+    if seed < 0:
+        rand = random.randint(1, 500)
+        for i in range(rand):
+            noise += " "
     model = load_model(model_name)
     if not model: return None
     task = asyncio.create_task(asyncio.to_thread(infer_body, model, f"{prompt} {noise}", neg_prompt,
-                                                 height, width, steps, cfg))
+                                                 height, width, steps, cfg, seed))
     await asyncio.sleep(0)
     try:
         result = await asyncio.wait_for(task, timeout=timeout)
-    except (Exception, asyncio.TimeoutError) as e:
+    except asyncio.TimeoutError as e:
         print(e)
         print(f"Task timed out: {model_name}")
         if not task.done(): task.cancel()
         result = None
+        raise Exception(f"Task timed out: {model_name}") from e
+    except Exception as e:
+        print(e)
+        if not task.done(): task.cancel()
+        result = None
+        raise Exception() from e
     if task.done() and result is not None:
         with lock:
             image = rename_image(result, model_name, save_path)
@@ -379,27 +422,32 @@ async def infer(model_name: str, prompt: str, neg_prompt: str | None = None,
     return None
 
 
+# https://github.com/aio-libs/pytest-aiohttp/issues/8 # also AsyncInferenceClient is buggy.
 def infer_fn(model_name: str, prompt: str, neg_prompt: str | None = None, height: int | None = None,
-             width: int | None = None, steps: int | None = None, cfg: int | None = None,
+             width: int | None = None, steps: int | None = None, cfg: int | None = None, 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
     try:
-        prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
+        loop = asyncio.get_running_loop()
+    except Exception:
         loop = asyncio.new_event_loop()
+    try:
+        prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
         result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
-                                               steps, cfg, save_path, inference_timeout))
+                                               steps, cfg, seed, save_path, inference_timeout))
     except (Exception, asyncio.CancelledError) as e:
         print(e)
-        print(f"Task aborted: {model_name}")
+        print(f"Task aborted: {model_name}, Error: {e}")
         result = None
+        raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
     finally:
         loop.close()
     return result
 
 
 def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str | None = None, height: int | None = None,
-             width: int | None = None, steps: int | None = None, cfg: int | None = None,
+             width: int | None = None, steps: int | None = None, cfg: int | None = None, 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':
@@ -407,14 +455,18 @@ def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str | None = N
     random.seed()
     model_name = random.choice(list(loaded_models.keys()))
     try:
-        prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
+        loop = asyncio.get_running_loop()
+    except Exception:
         loop = asyncio.new_event_loop()
+    try:
+        prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
         result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
-                                               steps, cfg, save_path, inference_timeout))
+                                               steps, cfg, seed, save_path, inference_timeout))
     except (Exception, asyncio.CancelledError) as e:
         print(e)
-        print(f"Task aborted: {model_name}")
+        print(f"Task aborted: {model_name}, Error: {e}")
         result = None
+        raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
     finally:
         loop.close()
     return result

tagger/fl2sd3longcap.py

@@ -1,5 +1,5 @@
-from transformers import AutoProcessor, AutoModelForCausalLM
 import spaces
+from transformers import AutoProcessor, AutoModelForCausalLM
 import re
 from PIL import Image 
 import torch
@@ -8,9 +8,13 @@ import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to(device).eval()
-fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
 
+try:
+    fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to("cpu").eval()
+    fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
+except Exception as e:
+    print(e)
+    fl_model = fl_processor = None
 
 def fl_modify_caption(caption: str) -> str:
     """
@@ -41,7 +45,7 @@ def fl_modify_caption(caption: str) -> str:
     return modified_caption if modified_caption != caption else caption
 
 
-@spaces.GPU
+@spaces.GPU(duration=30)
 def fl_run_example(image):
     task_prompt = "<DESCRIPTION>"
     prompt = task_prompt + "Describe this image in great detail."
@@ -50,6 +54,7 @@ def fl_run_example(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
 
+    fl_model.to(device)
     inputs = fl_processor(text=prompt, images=image, return_tensors="pt").to(device)
     generated_ids = fl_model.generate(
         input_ids=inputs["input_ids"],
@@ -57,6 +62,7 @@ def fl_run_example(image):
         max_new_tokens=1024,
         num_beams=3
     )
+    fl_model.to("cpu")
     generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
     parsed_answer = fl_processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
     return fl_modify_caption(parsed_answer["<DESCRIPTION>"])

tagger/tagger.py

@@ -1,7 +1,7 @@
+import spaces
 from PIL import Image
 import torch
 import gradio as gr
-import spaces
 from transformers import (
     AutoImageProcessor,
     AutoModelForImageClassification,
@@ -12,10 +12,15 @@ from pathlib import Path
 WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
 WD_MODEL_NAME = WD_MODEL_NAMES[0]
 
-wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
-wd_model.to("cuda" if torch.cuda.is_available() else "cpu")
-wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+default_device = device
 
+try:
+    wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True).to(default_device).eval()
+    wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+except Exception as e:
+    print(e)
+    wd_model = wd_processor = None
 
 def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
     return (
@@ -506,7 +511,7 @@ def gen_prompt(rating: list[str], character: list[str], general: list[str]):
     return ", ".join(all_tags)
 
 
-@spaces.GPU()
+@spaces.GPU(duration=30)
 def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8):
     inputs = wd_processor.preprocess(image, return_tensors="pt")
 
@@ -514,9 +519,11 @@ def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_t
     logits = torch.sigmoid(outputs.logits[0])  # take the first logits
 
     # get probabilities
+    if device != default_device: wd_model.to(device=device)
     results = {
         wd_model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
     }
+    if device != default_device: wd_model.to(device=default_device)
     # rating, character, general
     rating, character, general = postprocess_results(
         results, general_threshold, character_threshold