Update app.py

app.py

@@ -49,13 +49,50 @@ def gen_fn(model_str, prompt):
     noise = str('') #str(randint(0, 99999999999))
     return models_load[model_str](f'{prompt} {noise}')
 
-def gen_fnseed(model_str, prompt, useseed, **kwargs):
+#def gen_fnseed(model_str, prompt, useseed, **kwargs):
+#    kwargs = {}
+#    if model_str == 'NA':
+#        return None
+#    noise = str('') #str(randint(0, 99999999999))
+#    kwargs["seed"] = useseed
+#    return models_load[model_str](f'{prompt} {noise}', **kwargs)
+
+async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
+    from pathlib import Path
     kwargs = {}
+    noise = ""
+    kwargs["seed"] = seed
+    task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn,
+                               prompt=f'{prompt} {noise}', **kwargs, token=HF_TOKEN))
+    await asyncio.sleep(0)
+    try:
+        result = await asyncio.wait_for(task, timeout=timeout)
+    except (Exception, asyncio.TimeoutError) as e:
+        print(e)
+        print(f"Task timed out: {model_str}")
+        if not task.done(): task.cancel()
+        result = None
+    if task.done() and result is not None:
+        with lock:
+            png_path = "image.png"
+            result.save(png_path)
+            image = str(Path(png_path).resolve())
+        return image
+    return None
+
+def gen_fnseed(model_str, prompt, seed=1):
     if model_str == 'NA':
         return None
-    noise = str('') #str(randint(0, 99999999999))
-    kwargs["seed"] = useseed
-    return models_load[model_str](f'{prompt} {noise}', **kwargs)
+    try:
+        loop = asyncio.new_event_loop()
+        result = loop.run_until_complete(infer(model_str, prompt, seed, inference_timeout))
+    except (Exception, asyncio.CancelledError) as e:
+        print(e)
+        print(f"Task aborted: {model_str}")
+        result = None
+    finally:
+        loop.close()
+    return result
 
 def gen_fnsix(model_str, prompt):
     if model_str == 'NA':
@@ -107,7 +144,7 @@ with gr.Blocks() as demo:
     with gr.Tab('Seed it!'):
         model_choiceseed = gr.Dropdown(models, label = f'Choose a model from the {len(models)} available! Try clearing the box and typing on it to filter them!', value = models[0], filterable = True)
         txt_inputseed = gr.Textbox(label = 'Your prompt:')
-        useseed = gr.Slider(label="Use a seed to replicate the same image later", info="Max 3999999999", minimum=0, maximum=MAX_SEED, step=1, value=1)
+        seed = gr.Slider(label="Use a seed to replicate the same image later", info="Max 3999999999", minimum=0, maximum=MAX_SEED, step=1, value=1)
         
         max_imagesseed = 1
         num_imagesseed = gr.Slider(1, max_imagesone, value = max_imagesone, step = 1, label = 'One, because more would make it produce identical images with the seed', visible = False)
@@ -126,7 +163,7 @@ with gr.Blocks() as demo:
 
             gen_eventseed = gr.on(triggers=[gen_buttonseed.click, txt_inputseed.submit],
                                fn=lambda i, n, m, t, n1: gen_fnseed(m, t, n1) if (i < n) else None,
-                               inputs=[img_is, num_imagesseed, model_choiceseed, txt_inputseed, useseed], outputs=[o],
+                               inputs=[img_is, num_imagesseed, model_choiceseed, txt_inputseed, seed], outputs=[o],
                                        concurrency_limit=None, queue=False) # Be sure to delete ", queue=False" when activating the stop button
                         
             #stop_button.click(lambda s: gr.update(interactive = False), None, stop_button, cancels = [gen_event])