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):
@@ -38,30 +39,50 @@ 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]
 
-    
-def gen_fn(model_str, prompt):
-    if model_str == 'NA':
-        return None
+
+async def infer(model_str, prompt, timeout):
     noise = ""
     rand = randint(1, 500)
     for i in range(rand):
         noise += " "
-    return models_load[model_str](f'{prompt} {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
+
+
+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':
-        return gallery
-    noise = ""
-    rand = randint(1, 500)
-    for i in range(rand):
-        noise += " "
+        yield gallery
     try:
-        gallery.append(models_load[model_str](f'{prompt} {noise}'))
+        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)
-    return gallery
+    finally:
+        loop.close()
+    yield gallery
 
 
 CSS="""

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