Upload 2 files

app.py

@@ -10,7 +10,6 @@ lock = RLock()
 def load_fn(models):
     global models_load
     models_load = {}
-    
     for model in models:
         if model not in models_load.keys():
             try:
@@ -26,7 +25,7 @@ load_fn(models)
 
 num_models = 6
 max_images = 6
-timeout = 60
+inference_timeout = 300
 default_models = models[:num_models]
 
 
@@ -39,13 +38,19 @@ def update_imgbox(choices):
     return [gr.Image(None, label = m, visible = (m != 'NA')) for m in choices_plus]
 
 
-async def infer(model_str, prompt, timeout):
-    from PIL import Image
+async def infer(model_str, prompt, nprompt="", height=None, width=None, steps=None, cfg=None, timeout=inference_timeout):
+    from pathlib import Path
+    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
     noise = ""
     rand = randint(1, 500)
     for i in range(rand):
         noise += " "
-    task = asyncio.create_task(asyncio.to_thread(models_load[model_str], f'{prompt} {noise}'))
+    task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn,
+                               prompt=f'{prompt} {noise}', negative_prompt=nprompt, **kwargs))
     await asyncio.sleep(0)
     try:
         result = await asyncio.wait_for(task, timeout=timeout)
@@ -56,17 +61,20 @@ async def infer(model_str, prompt, timeout):
         result = None
     if task.done() and result is not None:
         with lock:
-            image = Image.open(result).convert('RGBA')
+            png_path = "image.png"
+            result.save(png_path)
+            image = str(Path(png_path).resolve())
         return image
     return None
 
 
-def gen_fn(model_str, prompt):
+def gen_fn(model_str, prompt, nprompt="", height=None, width=None, steps=None, cfg=None):
     if model_str == 'NA':
         return None
     try:
         loop = asyncio.new_event_loop()
-        result = loop.run_until_complete(infer(model_str, prompt, timeout))
+        result = loop.run_until_complete(infer(model_str, prompt, nprompt,
+                                         height, width, steps, cfg, inference_timeout))
     except (Exception, asyncio.CancelledError) as e:
         print(e)
         print(f"Task aborted: {model_str}")
@@ -93,7 +101,14 @@ CSS="""
 with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=CSS) as demo:
     with gr.Tab('The Dream'):
         with gr.Column(scale=2):
-            txt_input = gr.Textbox(label='Your prompt:', lines=4)
+            with gr.Group():
+                txt_input = gr.Textbox(label='Your prompt:', lines=4)
+                neg_input = gr.Textbox(label='Negative prompt:', lines=1)
+                with gr.Accordion("Advanced", open=False, visible=True):
+                    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.Row():
                 gen_button = gr.Button(f'Generate up to {int(num_models)} images in up to 3 minutes total', scale=2)
                 stop_button = gr.Button('Stop', variant='secondary', interactive=False, scale=1)
@@ -114,7 +129,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=CSS) as demo:
                                 preview=True, object_fit="cover", columns=2, rows=2) 
 
         for m, o in zip(current_models, output):
-            gen_event = gr.on(triggers=[gen_button.click, txt_input.submit], fn=gen_fn, inputs=[m, txt_input], outputs=[o])
+            gen_event = gr.on(triggers=[gen_button.click, txt_input.submit], fn=gen_fn,
+                              inputs=[m, txt_input, neg_input, height, width, steps, cfg], outputs=[o])
             o.change(add_gallery, [o, m, gallery], [gallery])
             stop_button.click(lambda: gr.update(interactive=False), None, stop_button, cancels=[gen_event])
 
@@ -127,12 +143,19 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=CSS) as demo:
     with gr.Tab('Single model'):
         with gr.Column(scale=2):
             model_choice2 = gr.Dropdown(models, label='Choose model', value=models[0])
-            txt_input2 = gr.Textbox(label='Your prompt:', lines=4)
+            with gr.Group():
+                txt_input2 = gr.Textbox(label='Your prompt:', lines=4)
+                neg_input2 = gr.Textbox(label='Negative prompt:', lines=1)
+                with gr.Accordion("Advanced", open=False, visible=True):
+                    width2 = gr.Number(label="Width", info="If 0, the default value is used.", maximum=1216, step=32, value=None)
+                    height2 = gr.Number(label="Height", info="If 0, the default value is used.", maximum=1216, step=32, value=None)
+                    steps2 = gr.Number(label="Number of inference steps", info="If 0, the default value is used.", maximum=100, step=1, value=None)
+                    cfg2 = gr.Number(label="Guidance scale", info="If 0, the default value is used.", maximum=30.0, step=0.1, value=None)
             num_images = gr.Slider(1, max_images, value=max_images, step=1, label='Number of images')
             with gr.Row():
                 gen_button2 = gr.Button('Generate', scale=2)
                 stop_button2 = gr.Button('Stop', variant='secondary', interactive=False, scale=1)
-                gen_button2.click(lambda: gr.update(interactive = True), None, stop_button2)
+                gen_button2.click(lambda: gr.update(interactive=True), None, stop_button2)
 
         with gr.Column(scale=1):
             with gr.Group():
@@ -150,8 +173,9 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=CSS) as demo:
             img_i = gr.Number(i, visible = False)
             num_images.change(lambda i, n: gr.update(visible = (i < n)), [img_i, num_images], o)
             gen_event2 = gr.on(triggers=[gen_button2.click, txt_input2.submit],
-                               fn=lambda i, n, m, t: gen_fn(m, t) if (i < n) else None,
-                               inputs=[img_i, num_images, model_choice2, txt_input2], outputs=[o])
+                               fn=lambda i, n, m, t1, t2, n1, n2, n3, n4: gen_fn(m, t1, t2, n1, n2, n3, n4) if (i < n) else None,
+                               inputs=[img_i, num_images, model_choice2, txt_input2, neg_input2,
+                                       height2, width2, steps2, cfg2], outputs=[o])
             o.change(add_gallery, [o, model_choice2, gallery2], [gallery2])
             stop_button2.click(lambda: gr.update(interactive=False), None, stop_button2, cancels=[gen_event2])
 

externalmod.py

@@ -33,6 +33,9 @@ if TYPE_CHECKING:
     from gradio.interface import Interface
 
 
+server_timeout = 600
+
+
 @document()
 def load(
     name: str,
@@ -115,7 +118,7 @@ def from_model(model_name: str, hf_token: str | None, alias: str | None, **kwarg
 
     headers["X-Wait-For-Model"] = "true"
     client = huggingface_hub.InferenceClient(
-        model=model_name, headers=headers, token=hf_token, timeout=120,
+        model=model_name, headers=headers, token=hf_token, timeout=server_timeout,
     )
 
     # For tasks that are not yet supported by the InferenceClient
@@ -365,10 +368,10 @@ def from_model(model_name: str, hf_token: str | None, alias: str | None, **kwarg
     else:
         raise ValueError(f"Unsupported pipeline type: {p}")
 
-    def query_huggingface_inference_endpoints(*data):
+    def query_huggingface_inference_endpoints(*data, **kwargs):
         if preprocess is not None:
             data = preprocess(*data)
-        data = fn(*data)  # type: ignore
+        data = fn(*data, **kwargs)  # type: ignore
         if postprocess is not None:
             data = postprocess(data)  # type: ignore
         return data