Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import requests
 import os
+import time
 
 # Environment variables for API details
 API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")  # Fetching the API token from environment variable
@@ -12,19 +13,25 @@ def query_huggingface_api(api_url, prompt):
     response = requests.post(api_url, headers=headers, json=data)
     
     if response.status_code == 200:
-        # Assuming the API returns binary image data
-        return response.content
+        return response.content, None  # Return the image and no error
     else:
-        return None, f"Error {response.status_code}: {response.text}"
+        return None, f"Error {response.status_code}: {response.text}"  # Return None and the error message
 
 # Gradio function for generating the image
 def generate_image(api_url, prompt):
-    result, error = query_huggingface_api(f"https://api-inference.huggingface.co/models/{api_url}", prompt)
-    
-    if result:
-        return result, None
-    else:
-        return None, error
+    # Attempt to query the API with retry logic for loading models
+    for attempt in range(5):  # Try up to 5 times
+        result, error = query_huggingface_api(f"https://api-inference.huggingface.co/models/{api_url}", prompt)
+        
+        if result:
+            return result, None
+        elif "Model is currently loading" in error:
+            estimated_time = float(error.split("estimated_time\":")[1].split("}")[0])  # Extract estimated time from error message
+            time.sleep(estimated_time + 5)  # Wait for the model to load, with an additional buffer time
+        else:
+            return None, error  # Return the error if it's not a loading issue
+
+    return None, "Model is still loading after multiple attempts. Please try again later."  # Final error if all attempts fail
 
 # Create Gradio Blocks Interface
 with gr.Blocks() as demo:
@@ -39,12 +46,13 @@ with gr.Blocks() as demo:
         with gr.Column():
             text_input = gr.Textbox(
                 label="Enter your prompt", 
-                placeholder="Type something here..."
+                placeholder="Type something here...",
+                value="psg stocking, blue eyes, blue hair, colored inner hair, hair bow, long hair, multicolored hair, pink hair, two-tone hair, bangs, blunt bangs,"
             )
             model_input = gr.Textbox(
                 label="Model URL", 
                 placeholder="Enter the model URL...",
-                value="user/sdwarm"
+                value=""
             )
             generate_btn = gr.Button("Generate Image")