Upload 3 files

README.md

@@ -4,7 +4,7 @@ emoji: 🌐🌊
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 4.40.0
+sdk_version: 4.42.0
 app_file: app.py
 short_description: Text-to-Image
 license: mit

app.py

@@ -1,24 +1,18 @@
 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_SEED = 2**32-1
 load_models(models)
 
 css = """
@@ -51,10 +45,13 @@ 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)
             neg_prompt = gr.Text(label="Negative Prompt", lines=1, max_lines=8, placeholder="")
             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.Row():
+                    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)
+                with gr.Row():
+                    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)
+                    seed = gr.Slider(label="Seed", info="Randomize Seed if -1.", minimum=-1, maximum=MAX_SEED, step=1, value=-1)
             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():
@@ -63,12 +60,14 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
                     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)
+                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", 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)
@@ -115,13 +114,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 +134,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()

multit2i.py

@@ -6,7 +6,7 @@ from huggingface_hub import InferenceClient
 import os
 
 
-HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None
+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
 
@@ -33,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(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()
     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", token=HF_TOKEN,
-                                       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 and HF_TOKEN is None:
-           if not_tag and not_tag in model.tags: continue
+        if not model.private and not model.gated:
+           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
@@ -333,13 +354,14 @@ def warm_model(model_name: str):
 # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
 def infer_body(client: InferenceClient | gr.Interface, 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, token=HF_TOKEN)
@@ -355,17 +377,18 @@ def infer_body(client: InferenceClient | gr.Interface, prompt: str, neg_prompt:
 
 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)
@@ -382,7 +405,7 @@ async def infer(model_name: str, prompt: str, neg_prompt: str | None = None,
 
 
 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
@@ -390,7 +413,7 @@ def infer_fn(model_name: str, prompt: str, neg_prompt: str | None = None, height
         prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
         loop = asyncio.new_event_loop()
         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}")
@@ -401,7 +424,7 @@ def infer_fn(model_name: str, prompt: str, neg_prompt: str | None = None, height
 
 
 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':
@@ -412,7 +435,7 @@ def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str | None = N
         prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
         loop = asyncio.new_event_loop()
         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}")