Update app.py

app.py

@@ -3,6 +3,8 @@ import requests
 import os
 from PIL import Image
 from io import BytesIO
+from tqdm import tqdm
+import time
 
 # Definindo as informações do repositório e a trigger word
 repo = "artificialguybr/StudioGhibli.Redmond-V2"
@@ -24,19 +26,35 @@ def generate_image(prompt):
         },
     }
 
-    try:
-        response = requests.post(api_url, headers=headers, json=payload)
-        response.raise_for_status()  # Verifica se a requisição foi bem-sucedida
-        image_data = response.json().get("data", [])[0]  # Obtém os dados da imagem
-        return Image.open(BytesIO(image_data))
-    except requests.exceptions.HTTPError as http_err:
-        print(f"HTTP error occurred: {http_err}")
-    except requests.exceptions.RequestException as err:
-        print(f"An error occurred: {err}")
-    except requests.exceptions.JSONDecodeError as json_err:
-        print(f"JSON decode error: {json_err}")
-    except Exception as e:
-        print(f"An unexpected error occurred: {e}")
+    error_count = 0
+    pbar = tqdm(total=None, desc="Loading model")
+    while True:
+        try:
+            response = requests.post(api_url, headers=headers, json=payload)
+            if response.status_code == 200:
+                image_data = response.json().get("data", [])[0]  # Obtém os dados da imagem
+                return Image.open(BytesIO(image_data))
+            elif response.status_code == 503:
+                # 503 é quando o modelo está em inicialização
+                time.sleep(1)
+                pbar.update(1)
+            elif response.status_code == 500 and error_count < 5:
+                # Erro interno do servidor, tenta novamente até 5 vezes
+                time.sleep(1)
+                error_count += 1
+            else:
+                raise Exception(f"API Error: {response.status_code}")
+        except requests.exceptions.HTTPError as http_err:
+            print(f"HTTP error occurred: {http_err}")
+        except requests.exceptions.RequestException as err:
+            print(f"An error occurred: {err}")
+        except requests.exceptions.JSONDecodeError as json_err:
+            print(f"JSON decode error: {json_err}")
+            # Se a resposta não for JSON, tenta carregar a imagem diretamente
+            return Image.open(BytesIO(response.content))
+        except Exception as e:
+            print(f"An unexpected error occurred: {e}")
+            break  # Sai do loop se um erro inesperado ocorrer
 
 # Criando a interface do Gradio
 iface = gr.Interface(