updated

app.py

@@ -2,9 +2,87 @@ import os
 import gradio as gr
 from PIL import Image
 import io
-from utils import query_hf_api
+import requests
+from typing import Optional, Tuple
 
-def generate_image(prompt: str) -> Image.Image:
+def load_environment():
+    """
+    Attempt to load environment variables with error handling.
+    
+    Returns:
+        Optional[str]: Hugging Face Token or None
+    """
+    try:
+        from dotenv import load_dotenv
+        load_dotenv()
+    except ImportError:
+        print("python-dotenv not installed. Ensure HF_TOKEN is set in environment.")
+    
+    return os.getenv("HF_TOKEN")
+
+def query_hf_api(
+    prompt: str, 
+    model_url: str = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0",
+    max_retries: int = 3
+) -> Optional[bytes]:
+    """
+    Query the Hugging Face Inference API with robust error handling and retry mechanism.
+    
+    Args:
+        prompt (str): Text prompt for image generation
+        model_url (str): URL of the Hugging Face model
+        max_retries (int): Maximum number of retry attempts
+    
+    Returns:
+        Optional[bytes]: Generated image bytes or None
+    """
+    # Validate inputs
+    if not prompt or not prompt.strip():
+        raise ValueError("Prompt cannot be empty")
+    
+    # Load token
+    HF_TOKEN = load_environment()
+    if not HF_TOKEN:
+        raise ValueError("Hugging Face token not found. Set HF_TOKEN in .env or environment variables.")
+    
+    # Prepare headers
+    headers = {
+        "Authorization": f"Bearer {HF_TOKEN}",
+        "Content-Type": "application/json"
+    }
+    
+    # Payload with additional configuration
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "negative_prompt": "low quality, bad anatomy, blurry",
+            "num_inference_steps": 50,
+        }
+    }
+    
+    # Retry mechanism
+    for attempt in range(max_retries):
+        try:
+            response = requests.post(
+                model_url, 
+                headers=headers, 
+                json=payload,
+                timeout=120  # 2-minute timeout
+            )
+            
+            response.raise_for_status()  # Raise exception for bad status codes
+            
+            return response.content
+        
+        except requests.exceptions.RequestException as e:
+            print(f"Request error (Attempt {attempt + 1}/{max_retries}): {e}")
+            
+            if attempt == max_retries - 1:
+                raise RuntimeError(f"Failed to generate image after {max_retries} attempts: {e}")
+
+    raise RuntimeError("Unexpected error in image generation")
+
+def generate_image(prompt: str) -> Tuple[Optional[Image.Image], str]:
     """
     Generate an image from a text prompt.
     
@@ -12,20 +90,24 @@ def generate_image(prompt: str) -> Image.Image:
         prompt (str): Text description for image generation
     
     Returns:
-        Image.Image: Generated PIL Image
+        Tuple[Optional[Image.Image], str]: Generated PIL Image and status message
     """
     try:
+        # Validate prompt
+        if not prompt or not prompt.strip():
+            return None, "Error: Prompt cannot be empty"
+        
         # Generate image bytes
         image_bytes = query_hf_api(prompt)
         
         # Convert to PIL Image
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         
-        return image
+        return image, "Image generated successfully!"
     
     except Exception as e:
         print(f"Image generation error: {e}")
-        return None
+        return None, f"Error: {str(e)}"
 
 def create_gradio_interface():
     """
@@ -52,23 +134,6 @@ def create_gradio_interface():
                     lines=3
                 )
                 
-                # Advanced Options
-                with gr.Accordion("Advanced Options", open=False):
-                    steps_slider = gr.Slider(
-                        minimum=10, 
-                        maximum=100, 
-                        value=50, 
-                        step=1, 
-                        label="Inference Steps"
-                    )
-                    guidance_slider = gr.Slider(
-                        minimum=1, 
-                        maximum=20, 
-                        value=7.5, 
-                        step=0.5, 
-                        label="Guidance Scale"
-                    )
-                
                 # Generate Button
                 generate_button = gr.Button("✨ Generate Image", variant="primary")
             
@@ -80,14 +145,14 @@ def create_gradio_interface():
                     interactive=False
                 )
         
-        # Error Handling Output
-        error_output = gr.Textbox(label="Status", visible=False)
+        # Status Output
+        status_output = gr.Textbox(label="Status")
         
         # Event Handlers
         generate_button.click(
             fn=generate_image,
             inputs=[text_input],
-            outputs=[output_image, error_output],
+            outputs=[output_image, status_output],
             api_name="generate"
         )