Update app.py

app.py

@@ -40,9 +40,19 @@ def save_image(img, suffix='.png'):
     
 async def niji_api(prompt, progress=gr.Progress(), max_retries=5, backoff_factor=0.1):
     iters = 1
-    progress(iters/24, desc="Sending request to MidJourney Server")
+    progress(iters/32, desc="Sending request to MidJourney Server")
+    session = requests.Session()  # Using Session to reuse the underlying TCP connection
+    retries = Retry(total=max_retries, backoff_factor=backoff_factor, status_forcelist=[429, 500, 502, 503, 504])
+    adapter = HTTPAdapter(max_retries=retries)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
     try:
-        response = requests.post(fetapi, headers={'Content-Type': 'application/json'}, data=json.dumps({'msg': prompt}))
+        response = session.post(
+            fetapi,
+            headers={'Content-Type': 'application/json'},
+            data=json.dumps({'msg': prompt}),
+            timeout=5.0  # Here, the timeout duration is set to 5 seconds
+        )
         response.raise_for_status()  # Check for HTTP errors.
     except requests.exceptions.RequestException as e:
         logger.error(f"Failed to make POST request")
@@ -55,7 +65,7 @@ async def niji_api(prompt, progress=gr.Progress(), max_retries=5, backoff_factor
 
     def fetch_image(url):
         try:
-            response = session.get(url)
+            response = session.get(url, timeout=5.0)
             return Image.open(BytesIO(response.content))
         except requests.exceptions.RequestException as e:
             logger.error(f"Failed to fetch image")
@@ -68,16 +78,17 @@ async def niji_api(prompt, progress=gr.Progress(), max_retries=5, backoff_factor
 
     while prog < 100:
         try:
-            response = requests.get(f'{odnapi}/message/{message_id}?expireMins=2', headers={'Authorization': auth_token})
+            response = session.get(
+                f'{odnapi}/message/{message_id}?expireMins=2', 
+                headers={'Authorization': auth_token},
+                timeout=5.0
+            )
             response.raise_for_status()
         except requests.exceptions.RequestException as e:
             logger.warning(f"Failure in getting message response")
             continue
         data = response.json()
         prog = data.get('progress', 0)
-        if isinstance(prog, str) and not prog.isdigit():
-            prog = 0
-            
         if progress_image_url := data.get('progressImageUrl'):
             iters = -100
             yield [(img, f"{prog}% done") for img in download_and_split_image(progress_image_url)]