Update app.py

app.py

@@ -10,8 +10,10 @@ import json
 
 # Project by Nymbo
 
-API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
+API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
 timeout = 100
+max_retries = 3
+retry_delay = 5
 
 # Function to query the API and return the generated image
 def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
@@ -42,24 +44,38 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
         }
     }
 
-    # Send the request to the API and handle the response
-    response = requests.post(API_URL, json=payload, timeout=timeout)
-    if response.status_code != 200:
-        print(f"Error: Failed to get image. Response status: {response.status_code}")
-        print(f"Response content: {response.text}")
-        if response.status_code == 503:
-            raise gr.Error(f"{response.status_code} : The model is being loaded")
-        raise gr.Error(f"{response.status_code}")
+    # Add retry logic for model loading
+    for attempt in range(max_retries):
+        try:
+            # Send the request to the API and handle the response
+            response = requests.post(API_URL, json=payload, timeout=timeout)
+            if response.status_code == 200:
+                try:
+                    # Convert the response content into an image
+                    image_bytes = response.content
+                    image = Image.open(io.BytesIO(image_bytes))
+                    print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
+                    return image
+                except Exception as e:
+                    print(f"Error when trying to open the image: {e}")
+                    return None
+            elif response.status_code == 503 and attempt < max_retries - 1:
+                print(f"Model is loading, attempt {attempt + 1}/{max_retries}. Waiting {retry_delay} seconds...")
+                time.sleep(retry_delay)
+                continue
+            else:
+                print(f"Error: Failed to get image. Response status: {response.status_code}")
+                print(f"Response content: {response.text}")
+                if response.status_code == 503:
+                    raise gr.Error(f"Model is still loading after {max_retries} attempts. Please try again in a few minutes.")
+                raise gr.Error(f"{response.status_code}")
+        except requests.exceptions.Timeout:
+            if attempt < max_retries - 1:
+                print(f"Request timed out, attempt {attempt + 1}/{max_retries}. Retrying...")
+                continue
+            raise gr.Error("Request timed out. Please try again.")
     
-    try:
-        # Convert the response content into an image
-        image_bytes = response.content
-        image = Image.open(io.BytesIO(image_bytes))
-        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
-        return image
-    except Exception as e:
-        print(f"Error when trying to open the image: {e}")
-        return None
+    return None
 
 # CSS to style the app
 css = """