Upload app.py

app.py

@@ -15,7 +15,7 @@ def load_fn(models):
                 m = gr_Interface_load(f'models/{model}')
             except Exception as error:
                 print(error)
-                m = gr.Interface(lambda txt: None, ['text'], ['image'])
+                m = gr.Interface(lambda: None, ['text'], ['image'])
             models_load.update({model: m})
 
 
@@ -24,7 +24,7 @@ load_fn(models)
 
 num_models = 6
 default_models = models[:num_models]
-
+timeout = 60
 
 def extend_choices(choices):
     return choices[:num_models] + (num_models - len(choices[:num_models])) * ['NA']
@@ -49,9 +49,10 @@ async def infer(model_str, prompt, timeout):
     await asyncio.sleep(0)
     try:
         result = await asyncio.wait_for(task, timeout=timeout)
-    except Exception as e:
+    except (Exception, asyncio.TimeoutError) as e:
         print(e)
-        task.cancel()
+        print(f"Task timed-out: {model_str}")
+        if not task.done(): task.cancel()
         result = None
     return result
 
@@ -61,9 +62,10 @@ def gen_fn(model_str, prompt):
         return None
     try:
         loop = asyncio.new_event_loop()
-        result = loop.run_until_complete(infer(model_str, prompt, 60))
-    except Exception as e:
+        result = loop.run_until_complete(infer(model_str, prompt, timeout))
+    except (Exception, asyncio.CancelledError) as e:
         print(e)
+        print(f"Task aborted: {model_str}")
         result = None
     finally:
         loop.close()
@@ -76,10 +78,11 @@ def gen_fn_gallery(model_str, prompt, gallery):
         yield gallery
     try:
         loop = asyncio.new_event_loop()
-        result = loop.run_until_complete(infer(model_str, prompt, 60))
+        result = loop.run_until_complete(infer(model_str, prompt, timeout))
         if result: gallery.append(result)
-    except Exception as e:
+    except (Exception, asyncio.CancelledError) as e:
         print(e)
+        print(f"Task aborted: {model_str}")
     finally:
         loop.close()
     yield gallery
@@ -94,7 +97,7 @@ with gr.Blocks(css=CSS) as demo:
         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')
         stop_button = gr.Button('Stop', variant = 'secondary', interactive = False)
-        gen_button.click(lambda s: gr.update(interactive = True), None, stop_button)
+        gen_button.click(lambda: gr.update(interactive = True), None, stop_button)
         gr.HTML(
         """
             <div style="text-align: center; max-width: 1200px; margin: 0 auto;">
@@ -115,7 +118,7 @@ with gr.Blocks(css=CSS) as demo:
             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])
+                stop_button.click(lambda: 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, interactive = True)
             model_choice.change(update_imgbox, model_choice, output)