Update app.py

app.py

@@ -9,6 +9,8 @@ import logging
 import threading
 from PIL import Image
 from io import BytesIO
+from requests.adapters import HTTPAdapter
+from urllib3.util import Retry
 
 odnapi = os.getenv("odnapi_url")
 fetapi = os.getenv("fetapi_url")
@@ -18,14 +20,6 @@ auth_token = os.getenv("auth_token")
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def fetch_image(url):
-    try:
-        response = requests.get(url)
-        return Image.open(BytesIO(response.content))
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Failed to fetch image")
-        raise
-
 def split_image(img):
     width, height = img.size
     width_cut = width // 2
@@ -41,34 +35,67 @@ def save_image(img, suffix='.png'):
     with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
         img.save(tmp, 'PNG')
         return tmp.name
-
-def download_and_split_image(url):
-    img = fetch_image(url)
-    images = split_image(img)
-    return [save_image(i) for i in images]
-
-def niji_api(prompt):
+    
+async def niji_api(prompt, progress=gr.Progress(), max_retries=5, backoff_factor=0.1):
+    iters = 1
+    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}))
         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:
+    prog = 0
+    iters += 5
+    progress(iters/48, desc="Waiting in the generate queue")
+
+    def fetch_image(url):
         try:
-            response = requests.get(f'{odnapi}/message/{message_id}?expireMins=2', headers={'Authorization': auth_token})
+            response = session.get(url)
+            return Image.open(BytesIO(response.content))
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Failed to fetch image")
+            raise
+
+    def download_and_split_image(url):
+        img = fetch_image(url)
+        images = split_image(img)
+        return [save_image(i) for i in images]
+
+    while prog < 100:
+        try:
+            response = session.get(f'{odnapi}/message/{message_id}?expireMins=2', 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'):
-            yield [(img, f"{progress}% done") for img in download_and_split_image(progress_image_url)]
-        time.sleep(1)
+        prog = data.get('progress', 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)]
+
+        wait_time = random.uniform(1, 2) 
+        await asyncio.sleep(wait_time)
+
+        r = iters/48
+        if r < 0.4:
+          desc = "Waiting in the generate queue"
+        elif r < 0.6:
+          desc = "Still queueing"
+        elif r < 0.8:
+          desc = "Almost done"
+        if iters > 0:
+            progress(r, desc=desc)
+            iters += random.uniform(1, 2)
+
     # 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))]