Update app.py

app.py

@@ -6,6 +6,7 @@ import json
 import os
 import tempfile
 import logging
+import threading
 from PIL import Image
 from io import BytesIO
 
@@ -47,30 +48,50 @@ def download_and_split_image(url):
     return [save_image(i) for i in images]
 
 def niji_api(prompt):
-    try:
-        response = requests.post(fetapi, headers={'Content-Type': 'application/json'}, data=json.dumps({'msg': prompt}))
-        response.raise_for_status()  # Check for HTTP errors.
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Failed to make POST request")
-        raise ValueError("Invalid Response") 
-    data = response.json()
-    message_id = data['messageId']
-    progress = 0
-    while progress < 100:
+    message_id = None
+    images = []
+
+    def function():
+        nonlocal message_id, images
         try:
-            response = requests.get(f'{odnapi}/message/{message_id}', headers={'Authorization': auth_token})
-            response.raise_for_status()
+            response = requests.post(fetapi, headers={'Content-Type': 'application/json'}, data=json.dumps({'msg': prompt}))
+            response.raise_for_status()  # Check for HTTP errors.
         except requests.exceptions.RequestException as e:
-            logger.warning(f"Failure in getting message response")
-            continue
+            logger.error(f"Failed to make POST request")
+            raise ValueError("Invalid Response") 
         data = response.json()
-        progress = data.get('progress', 0)
-        if progress_image_url:= data.get('progressImageUrl'):
-            yield [(img, f"{progress}% done") for img in download_and_split_image(progress_image_url)]
-        time.sleep(1)
-    # Process the final image urls
-    image_url = data['response']['imageUrl']
-    yield [(img, f"image {idx+1}/4") for idx, img in enumerate(download_and_split_image(image_url))]
+        message_id = data['messageId']
+        
+        progress = 0
+        while progress < 100:
+            try:
+                response = requests.get(f'{odnapi}/message/{message_id}', headers={'Authorization': auth_token})
+                response.raise_for_status()
+            except requests.exceptions.RequestException as e:
+                logger.warning(f"Failure in getting message response")
+                continue
+            
+            data = response.json()
+            progress = data.get('progress', 0)
+            
+            if progress_image_url:= data.get('progressImageUrl'):
+                images.extend([(img, f"{progress}% done") for img in download_and_split_image(progress_image_url)])
+            time.sleep(1)
+        
+        # Process the final image urls
+        image_url = data['response']['imageUrl']
+        images.extend([(img, f"image {idx+1}/4") for idx, img in enumerate(download_and_split_image(image_url))])
+        
+    t = threading.Timer(180, function=function)
+    t.start()
+    t.join()
+
+    if t.is_alive():
+        # Thread is still executing, raise a TimeoutError
+        t.cancel()
+        raise TimeoutError("Operation timed out after 3 minutes")
+
+    return images
 
 with gr.Blocks() as demo:
     gr.HTML('''