Update app.py

app.py

@@ -11,17 +11,23 @@ API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you've set this environment vari
 API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"
 
 # Function to call Hugging Face API and get the generated image
-def generate_image(prompt):
+def generate_image(prompt, progress=gr.Progress()):
     headers = {"Authorization": f"Bearer {API_TOKEN}"}
     data = {"inputs": prompt}
     
+    # Initialize progress
+    progress(0, desc="Starting image generation...")
+    
     response = requests.post(API_URL, headers=headers, json=data)
     
     if response.status_code == 200:
+        progress(50, desc="Processing image...")
         image_bytes = BytesIO(response.content)
         image = Image.open(image_bytes)
+        progress(100, desc="Completed")
         return image
     else:
+        progress(100, desc="Error occurred")
         return f"Error: {response.status_code}, {response.text}"
 
 # Create Gradio interface
@@ -35,8 +41,8 @@ def create_ui():
         with gr.Row():
             output_image = gr.Image(label="Generated Image")
         
-        # Link the button to the function
-        generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
+        # Link the button to the function with progress
+        generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image, show_progress=True)
     
     return ui