Update app.py

app.py

@@ -5,6 +5,8 @@ from externalmod import gr_Interface_load
 import asyncio
 import os
 from threading import RLock
+from flask import send_file
+
 
 lock = RLock()
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -29,7 +31,7 @@ MAX_SEED = 3999999999
 default_models = models[:num_models]
 inference_timeout = 600
 
-async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
+#async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
     kwargs = {"seed": seed}
     task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn, prompt=prompt, **kwargs, token=HF_TOKEN))
     await asyncio.sleep(0)
@@ -48,6 +50,31 @@ async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
         return png_path
     return None
 
+
+
+async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
+    kwargs = {"seed": seed}
+    task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn, prompt=prompt, **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()
+        return None  # Return None if the task fails or times out
+    
+    if task.done() and result is not None:
+        with lock:
+            png_path = "image.png"
+            result.save(png_path)  # Save the image to a file
+        return send_file(png_path, mimetype='image/png')  # Directly use send_file here
+    
+    return None  # Return None if no result is generated
+
+
 # Expose Gradio API
 def generate_api(model_str, prompt, seed=1):
     result = asyncio.run(infer(model_str, prompt, seed))