Update app.py

app.py

@@ -48,50 +48,30 @@ def download_and_split_image(url):
     return [save_image(i) for i in images]
 
 def niji_api(prompt):
-    message_id = None
-    images = []
-
-    def function():
-        nonlocal message_id, images
+    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:
         try:
-            response = requests.post(fetapi, headers={'Content-Type': 'application/json'}, data=json.dumps({'msg': prompt}))
-            response.raise_for_status()  # Check for HTTP errors.
+            response = requests.get(f'{odnapi}/message/{message_id}', headers={'Authorization': auth_token})
+            response.raise_for_status()
         except requests.exceptions.RequestException as e:
-            logger.error(f"Failed to make POST request")
-            raise ValueError("Invalid Response") 
+            logger.warning(f"Failure in getting message response")
+            continue
         data = response.json()
-        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
+        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))]
 
 with gr.Blocks() as demo:
     gr.HTML('''