Update app.py

app.py

@@ -8,7 +8,6 @@ from tqdm import tqdm
 import gradio as gr
 import base64
 import io
-from fastapi import Request
 
 from safetensors.torch import save_file
 from src.pipeline import FluxPipeline
@@ -20,7 +19,7 @@ base_path = "black-forest-labs/FLUX.1-dev"
 lora_base_path = "./models"
 
 # Environment variable for API token (set this in your Hugging Face space settings)
-API_TOKEN = os.environ.get("HF_TOKEN")
+API_TOKEN = os.environ.get("HF_API_TOKEN")
 
 # Initialize the pipeline
 pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
@@ -39,21 +38,10 @@ def verify_token(token):
 
 # Define the Gradio interface with token verification
 @spaces.GPU()
-def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type, api_token="", request: gr.Request = None):
-    # Check authentication for API requests
-    if not request.is_from_ui:
-        # For API requests, check Authorization header
-        auth_header = request.headers.get("Authorization")
-        if not auth_header or not auth_header.startswith("Bearer "):
-            return {"error": "Unauthorized access. Invalid or missing token in Authorization header."}
-        
-        token = auth_header.replace("Bearer ", "")
-        if not verify_token(token):
-            return {"error": "Unauthorized access. Invalid token."}
-    else:
-        # For UI requests, check the token input
-        if not verify_token(api_token):
-            return "Unauthorized: Please enter a valid API token"
+def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type, api_token="", request=None):
+    # Check authentication
+    if not verify_token(api_token):
+        return "Unauthorized: Please enter a valid API token"
     
     try:
         # Set the control type
@@ -77,26 +65,45 @@ def single_condition_generate_image(prompt, spatial_img, height, width, seed, co
         ).images[0]
         clear_cache(pipe.transformer)
         
-        # Convert to base64 for API responses
-        if not request.is_from_ui:
-            buffered = io.BytesIO()
-            image.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            return {"meta": {"format": "png", "base64": img_str}}
-        
+        # We'll always return the PIL image for UI
+        # The API will extract base64 from the returned image
         return image
     
     except Exception as e:
         error_msg = f"Error during image generation: {str(e)}"
         print(error_msg)
-        if not request.is_from_ui:
-            return {"error": error_msg}
         return None
 
+# Define an API endpoint that uses the main function but returns proper JSON
+@spaces.GPU()
+def api_generate_image(prompt, spatial_img, height, width, seed, control_type, api_token=""):
+    # Verify the API token
+    if not verify_token(api_token):
+        return {"error": "Unauthorized access. Invalid token."}
+    
+    try:
+        # Use the same function as the UI
+        image = single_condition_generate_image(
+            prompt, spatial_img, height, width, seed, control_type, api_token
+        )
+        
+        if image is None or isinstance(image, str):
+            # Error occurred
+            error_msg = image if isinstance(image, str) else "Image generation failed"
+            return {"error": error_msg}
+        
+        # Return the image directly instead of converting to base64
+        return image
+    
+    except Exception as e:
+        error_msg = f"API error: {str(e)}"
+        print(error_msg)
+        return {"error": error_msg}
+
 # Define the Gradio interface components
 control_types = ["Ghibli"]
 
-# Example data
+# Example data - add the API token for convenience (assuming only you can see the examples)
 single_examples = [
     ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/00.png"), 680, 1024, 5, "Ghibli", API_TOKEN],
     ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/02.png"), 560, 1024, 42, "Ghibli", API_TOKEN],
@@ -151,5 +158,14 @@ with gr.Blocks() as demo:
         outputs=single_output_image
     )
 
+# Create an API endpoint that clients can use programmatically
+demo.queue()
+
+# Add the API endpoint
+demo.load(api_generate_image, 
+         inputs=[prompt, spatial_img, height, width, seed, control_type, api_token],
+         outputs=gr.JSON(),
+         api_name="generate")
+
 # Launch the Gradio app
-demo.queue().launch()
+demo.launch()