Upload 2 files

app.py

@@ -50,13 +50,18 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", fill_width=True, css=css) as demo:
                 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=None)
+                height = gr.Number(label="Height", info="If 0, the default value is used.", maximum=1216, step=32, value=None)
+                steps = gr.Number(label="Number of inference steps", info="If 0, the default value is used.", maximum=100, step=1, value=None)
+                cfg = gr.Number(label="Guidance scale", info="If 0, the default value is used.", maximum=30.0, step=0.1, value=None)
             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=[], visible=False)
-                    negative_suffix = gr.CheckboxGroup(label="Use Negative Suffix", choices=get_negative_suffix(), value=["Common"], visible=False)
+                    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)
@@ -109,12 +114,14 @@ 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, l1, l2, l3, l4: infer_fn(m, t1, t2, l1, l2, l3, l4) if (i < n) else None,
-         inputs=[img_i, image_num, model_name, prompt, neg_prompt, positive_prefix, positive_suffix, negative_prefix, negative_suffix],
+         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,
+                  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, l1, l2, l3, l4: infer_rand_fn(m, t1, t2, l1, l2, l3, l4) if (i < n) else None,
-         inputs=[img_i, image_num, model_name, prompt, neg_prompt, positive_prefix, positive_suffix, negative_prefix, negative_suffix],
+         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,
+                  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)
         stop_button.click(lambda: gr.update(interactive=False), None, stop_button, cancels=[gen_event, gen_event2], show_api=False)

multit2i.py

@@ -2,6 +2,11 @@ import gradio as gr
 import asyncio
 from threading import RLock
 from pathlib import Path
+from huggingface_hub import InferenceClient
+
+
+server_timeout = 600
+inference_timeout = 300
 
 
 lock = RLock()
@@ -80,7 +85,7 @@ def get_t2i_model_info_dict(repo_id: str):
     return info
 
 
-def rename_image(image_path: str | None, model_name: str):
+def rename_image(image_path: str | None, model_name: str, save_path: str | None = None):
     from PIL import Image
     from datetime import datetime, timezone, timedelta
     if image_path is None: return None
@@ -90,7 +95,10 @@ def rename_image(image_path: str | None, model_name: str):
         if Path(image_path).exists():
             png_path = "image.png"
             Image.open(image_path).convert('RGBA').save(png_path, "PNG")
-            new_path = str(Path(png_path).resolve().rename(Path(filename).resolve()))
+            if save_path is not None:
+                new_path = str(Path(png_path).resolve().rename(Path(save_path).resolve()))
+            else:
+                new_path = str(Path(png_path).resolve().rename(Path(filename).resolve()))
             return new_path
         else:
             return None
@@ -125,13 +133,14 @@ def load_from_model(model_name: str, hf_token: str = None):
             f"Could not find model: {model_name}. If it is a private or gated model, please provide your Hugging Face access token (https://huggingface.co/settings/tokens) as the argument for the `hf_token` parameter."
         )
     headers["X-Wait-For-Model"] = "true"
-    client = huggingface_hub.InferenceClient(model=model_name, headers=headers, token=hf_token, timeout=600)
+    client = huggingface_hub.InferenceClient(model=model_name, headers=headers,
+                                              token=hf_token, timeout=server_timeout)
     inputs = gr.components.Textbox(label="Input")
     outputs = gr.components.Image(label="Output")
     fn = client.text_to_image
 
-    def query_huggingface_inference_endpoints(*data):
-        return fn(*data)
+    def query_huggingface_inference_endpoints(*data, **kwargs):
+        return fn(*data, **kwargs)
 
     interface_info = {
         "fn": query_huggingface_inference_endpoints,
@@ -164,6 +173,34 @@ def load_model(model_name: str):
     return loaded_models[model_name]
 
 
+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)
+        if status is None or status.framework != "diffusers" and not status.state 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)
+            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
+    try:
+        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]
+
+
 def load_models(models: list):
     for model in models:
         load_model(model)
@@ -276,21 +313,48 @@ def get_model_info_md(model_name: str):
 
 
 def change_model(model_name: str):
-    load_model(model_name)
+    load_model_api(model_name)
     return get_model_info_md(model_name)
 
 
 def warm_model(model_name: str):
-    model = load_model(model_name)
+    model = load_model_api(model_name)
     if model:
         try:
             print(f"Warming model: {model_name}")
-            model(" ")
+            infer_body(model, " ")
         except Exception as e:
             print(e)
 
 
-async def infer(model_name: str, prompt: str, neg_prompt: str, timeout: float):
+# 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,
+               height: int | None = None, width: int | None = None,
+               steps: int | None = None, cfg: int | None = None):
+    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
+    try:
+        if isinstance(client, InferenceClient):
+            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)
+        else: return None
+        image.save(png_path)
+        return str(Path(png_path).resolve())
+    except Exception as e:
+        print(e)
+        return None
+
+
+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,
+               save_path: str | None = None, timeout: float = inference_timeout):
     import random
     noise = ""
     rand = random.randint(1, 500)
@@ -298,7 +362,8 @@ async def infer(model_name: str, prompt: str, neg_prompt: str, timeout: float):
         noise += " "
     model = load_model(model_name)
     if not model: return None
-    task = asyncio.create_task(asyncio.to_thread(model, f'{prompt} {noise}'))
+    task = asyncio.create_task(asyncio.to_thread(infer_body, model, f"{prompt} {noise}", neg_prompt,
+                                                 height, width, steps, cfg))
     await asyncio.sleep(0)
     try:
         result = await asyncio.wait_for(task, timeout=timeout)
@@ -309,20 +374,21 @@ async def infer(model_name: str, prompt: str, neg_prompt: str, timeout: float):
         result = None
     if task.done() and result is not None:
         with lock:
-            image = rename_image(result, model_name)
+            image = rename_image(result, model_name, save_path)
         return image
     return None
 
 
-infer_timeout = 300
-def infer_fn(model_name: str, prompt: str, neg_prompt: str,
-             pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
+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,
+             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.new_event_loop()
-        result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, infer_timeout))
+        result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
+                                               steps, cfg, save_path, inference_timeout))
     except (Exception, asyncio.CancelledError) as e:
         print(e)
         print(f"Task aborted: {model_name}")
@@ -332,8 +398,9 @@ def infer_fn(model_name: str, prompt: str, neg_prompt: str,
     return result
 
 
-def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str,
-             pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
+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,
+             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
@@ -342,7 +409,8 @@ def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str,
     try:
         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, infer_timeout))
+        result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
+                                               steps, cfg, save_path, inference_timeout))
     except (Exception, asyncio.CancelledError) as e:
         print(e)
         print(f"Task aborted: {model_name}")