Update app.py

app.py

@@ -8,6 +8,7 @@ from huggingface_hub import whoami
 from PIL import Image
 from io import BytesIO
 import tempfile
+import time
 
 # --- Google Gemini API Configuration ---
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
@@ -70,10 +71,17 @@ def unified_image_generator(
         # Always add the prompt to the contents
         contents.append(prompt)
 
-        response = client.models.generate_content( 
-            model=GEMINI_MODEL_NAME,
-            contents=contents,
-        )
+        for attempt in range(1, 3 + 1):
+            try:
+                response = client.models.generate_content(
+                    model=GEMINI_MODEL_NAME,
+                    contents=contents,
+                )
+                break
+            except Exception as e:
+                if attempt == MAX_RETRIES:
+                    raise gr.Error(f"The Gemini API returned an error: {e}")
+                time.sleep(1)
         
         image_data = _extract_image_data_from_response(response)