Upload 13 files

README.md

@@ -1,5 +1,5 @@
 ---
-title: Free Multi Models Text-to-Image Heavy-Armed Demo V2
+title: Free Multi Models Text-to-Image Heavy-Armed Demo V3
 emoji: 🌐🌊
 colorFrom: blue
 colorTo: purple

app.py

@@ -11,81 +11,84 @@ 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)
 
-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="")
-            with gr.Accordion("Advanced options", open=False):
-                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)
-            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():
-                    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):
+    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():
-                    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)
-        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", scale=6)
+                random_button = gr.Button("Random Model 🎲", variant="secondary", scale=3)
+                stop_button = gr.Button('Stop', variant="secondary", 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, container=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 = [

multit2i.py

@@ -35,7 +35,7 @@ def is_repo_name(s):
 
 def get_status(model_name: str):
     from huggingface_hub import InferenceClient
-    client = InferenceClient(timeout=10)
+    client = InferenceClient(token=HF_TOKEN, timeout=10)
     return client.get_model_status(model_name)
 
 
@@ -54,7 +54,7 @@ def is_loadable(model_name: str, force_gpu: bool = 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):
     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
@@ -67,7 +67,7 @@ def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="l
         print(e)
         return models
     for model in model_infos:
-        if not model.private and not model.gated:
+        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)
@@ -77,7 +77,7 @@ 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
@@ -86,14 +86,15 @@ def get_t2i_model_info_dict(repo_id: str):
         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"
@@ -109,7 +110,8 @@ def get_t2i_model_info_dict(repo_id: str):
 
 
 def rename_image(image_path: str | None, model_name: str, save_path: str | None = None):
-    from PIL import Image
+    from PIL import Image, ImageFile
+    ImageFile.LOAD_TRUNCATED_IMAGES = True
     from datetime import datetime, timezone, timedelta
     if image_path is None: return None
     dt_now = datetime.now(timezone(timedelta(hours=9)))
@@ -352,7 +354,7 @@ 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, seed: int = -1):
     png_path = "image.png"
@@ -372,7 +374,7 @@ def infer_body(client: InferenceClient | gr.Interface, prompt: str, neg_prompt:
         return str(Path(png_path).resolve())
     except Exception as e:
         print(e)
-        return None
+        raise Exception(e)
 
 
 async def infer(model_name: str, prompt: str, neg_prompt: str | None = None,
@@ -392,11 +394,17 @@ async def infer(model_name: str, prompt: str, neg_prompt: str | None = None,
     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}")
+    except Exception as e:
+        print(e)
+        if not task.done(): task.cancel()
+        result = None
+        raise Exception(e)
     if task.done() and result is not None:
         with lock:
             image = rename_image(result, model_name, save_path)
@@ -404,20 +412,25 @@ 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, 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, 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
@@ -432,14 +445,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, 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