Upload 2 files

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from random import randint
 from all_models import models
 from externalmod import gr_Interface_load
+import asyncio
 
 
 def load_fn(models):
@@ -13,6 +14,7 @@ def load_fn(models):
             try:
                 m = gr_Interface_load(f'models/{model}')
             except Exception as error:
+                print(error)
                 m = gr.Interface(lambda txt: None, ['text'], ['image'])
             models_load.update({model: m})
 
@@ -24,25 +26,70 @@ num_models = 6
 default_models = models[:num_models]
 
 
-
 def extend_choices(choices):
-    return choices[0:num_models] + (num_models - len(choices[0:num_models])) * ['NA']
+    return choices[:num_models] + (num_models - len(choices[:num_models])) * ['NA']
 
 
 def update_imgbox(choices):
-    choices_plus = extend_choices(choices[0:num_models])
+    choices_plus = extend_choices(choices[:num_models])
     return [gr.Image(None, label = m, visible = (m != 'NA')) for m in choices_plus]
 
-    
-def gen_fn(model_str, prompt):
-    if model_str == 'NA':
-        return None
-    noise = str('') #str(randint(0, 99999999999))
-    return models_load[model_str](f'{prompt} {noise}')
+
+def update_imgbox_gallery(choices):
+    choices_plus = extend_choices(choices[:num_models])
+    return [gr.Gallery(None, label = m, visible = (m != 'NA')) for m in choices_plus]
 
 
+async def infer(model_str, prompt, timeout):
+    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}'))
+    await asyncio.sleep(0)
+    try:
+        result = await asyncio.wait_for(task, timeout=timeout)
+    except Exception as e:
+        print(e)
+        task.cancel()
+        result = None
+    return result
 
-with gr.Blocks() as demo:
+
+def gen_fn(model_str, prompt):
+    if model_str == 'NA':
+        return None
+    try:
+        loop = asyncio.new_event_loop()
+        result = loop.run_until_complete(infer(model_str, prompt, 60))
+    except Exception as e:
+        print(e)
+        result = None
+    finally:
+        loop.close()
+    return result
+
+
+def gen_fn_gallery(model_str, prompt, gallery):
+    if gallery is None: gallery = []
+    if model_str == 'NA':
+        yield gallery
+    try:
+        loop = asyncio.new_event_loop()
+        result = loop.run_until_complete(infer(model_str, prompt, 60))
+        if result: gallery.append(result)
+    except Exception as e:
+        print(e)
+    finally:
+        loop.close()
+    yield gallery
+
+
+CSS="""
+.output { width=480px; height=480px !important; }
+"""
+
+with gr.Blocks(css=CSS) as demo:
     with gr.Tab('The Dream'): 
         txt_input = gr.Textbox(label = 'Your prompt:', lines=4) #.style(container=False,min_width=1200)
         gen_button = gr.Button('Generate up to 6 images in up to 3 minutes total')
@@ -61,19 +108,18 @@ with gr.Blocks() as demo:
         """
                )
         with gr.Row():
-            output = [gr.Image(label = m, min_width=480) for m in default_models]
+            output = [gr.Image(label = m, show_download_button=True, elem_classes="output", show_share_button=True) for m in default_models]
+            #output = [gr.Gallery(label = m, show_download_button=True, elem_classes="output", interactive=False, show_share_button=True, container=True, format="png", object_fit="contain") for m in default_models]
             current_models = [gr.Textbox(m, visible = False) for m in default_models]
                         
             for m, o in zip(current_models, output):
                 gen_event = gen_button.click(gen_fn, [m, txt_input], o)
+                #gen_event = gen_button.click(gen_fn_gallery, [m, txt_input, o], o)
                 stop_button.click(lambda s: gr.update(interactive = False), None, stop_button, cancels = [gen_event])
         with gr.Accordion('Model selection'):
-            #model_choice = gr.CheckboxGroup(models, label = f'Choose up to {num_models} different models from the 866 available!', value = default_models, multiselect = True, max_choices = num_models, interactive = True, filterable = False)
-            #model_choice.change(update_imgbox, model_choice, output)
-            #model_choice.change(extend_choices, model_choice, current_models)
             model_choice = gr.CheckboxGroup(models, label = f'Choose up to {num_models} different models from the 866 available!', value = default_models, interactive = True)
-            #model_choice.change(lambda x: x[0:num_models], [model_choice], [model_choice])
             model_choice.change(update_imgbox, model_choice, output)
+            #model_choice.change(update_imgbox_gallery, model_choice, output)
             model_choice.change(extend_choices, model_choice, current_models)
         with gr.Row():
             gr.HTML(
@@ -84,5 +130,5 @@ with gr.Blocks() as demo:
     """
 )
 
-demo.queue()                        
-demo.launch()
+demo.queue()
+demo.launch()

externalmod.py

@@ -115,7 +115,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
+        model=model_name, headers=headers, token=hf_token, timeout=120,
     )
 
     # For tasks that are not yet supported by the InferenceClient