Update app.py

app.py

@@ -8,9 +8,7 @@ from tqdm import tqdm
 import gradio as gr
 import base64
 import io
-from fastapi import Request, HTTPException
-import secrets
-from functools import wraps
+from fastapi import Request
 
 from safetensors.torch import save_file
 from src.pipeline import FluxPipeline
@@ -24,9 +22,6 @@ lora_base_path = "./models"
 # Environment variable for API token (set this in your Hugging Face space settings)
 API_TOKEN = os.environ.get("HF_TOKEN")
 
-# Generate a secure random session token for UI access (regenerates on restart)
-UI_SESSION_TOKEN = secrets.token_urlsafe(32)
-
 # Initialize the pipeline
 pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
 transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16)
@@ -37,60 +32,29 @@ def clear_cache(transformer):
     for name, attn_processor in transformer.attn_processors.items():
         attn_processor.bank_kv.clear()
 
-# Token verification middleware
-def verify_token(request: Request):
-    """Verify the token from the request header"""
-    try:
+# Token verification function
+def verify_token(token):
+    """Verify if the provided token matches the API token"""
+    return API_TOKEN and token == API_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 False
+            return {"error": "Unauthorized access. Invalid or missing token in Authorization header."}
         
         token = auth_header.replace("Bearer ", "")
-        
-        # Check against your API token
-        if API_TOKEN and token == API_TOKEN:
-            return True
-        
-        return False
-    except Exception as e:
-        print(f"Token verification error: {e}")
-        return False
-
-# Auth decorator for Gradio functions
-def require_auth(func):
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        # Extract the request object
-        request = kwargs.get('request')
-        
-        if not request:
-            return {"error": "Authentication required"}
-        
-        # Check for API token in header for programmatic access
-        auth_header = request.headers.get("Authorization")
-        if auth_header and auth_header.startswith("Bearer "):
-            token = auth_header.replace("Bearer ", "")
-            if API_TOKEN and token == API_TOKEN:
-                return func(*args, **kwargs)
-        
-        # Check for UI session token in cookies for UI access
-        auth_cookie = request.cookies.get("hf_ui_auth")
-        if auth_cookie == UI_SESSION_TOKEN:
-            return func(*args, **kwargs)
-            
-        # If no valid auth found, return error
-        if request.is_from_ui:
-            return "Unauthorized: Valid authentication required"
-        else:
-            return {"error": "Unauthorized access. Invalid or missing token."}
+        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"
     
-    return wrapper
-
-# Define the Gradio interface with token verification
-@spaces.GPU()
-@require_auth
-def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type, request: gr.Request = None):
-
     try:
         # Set the control type
         if control_type == "Ghibli":
@@ -134,22 +98,26 @@ control_types = ["Ghibli"]
 
 # Example data
 single_examples = [
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/00.png"), 680, 1024, 5, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/02.png"), 560, 1024, 42, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/03.png"), 568, 1024, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/04.png"), 768, 672, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/06.png"), 896, 1024, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/07.png"), 528, 800, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/08.png"), 696, 1024, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/09.png"), 896, 1024, 1, "Ghibli"],
+    ["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],
+    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/03.png"), 568, 1024, 1, "Ghibli", API_TOKEN],
+    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/04.png"), 768, 672, 1, "Ghibli", API_TOKEN],
+    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/06.png"), 896, 1024, 1, "Ghibli", API_TOKEN],
+    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/07.png"), 528, 800, 1, "Ghibli", API_TOKEN],
+    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/08.png"), 696, 1024, 1, "Ghibli", API_TOKEN],
+    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/09.png"), 896, 1024, 1, "Ghibli", API_TOKEN],
 ]
 
 # Create the Gradio Blocks interface
 with gr.Blocks() as demo:
     gr.Markdown("# Ghibli Studio Control Image Generation with EasyControl")
+    gr.Markdown("⚠️ **AUTHENTICATION REQUIRED**: You must enter a valid API token to use this interface.")
     gr.Markdown("The model is trained on **only 100 real Asian faces** paired with **GPT-4o-generated Ghibli-style counterparts**, and it preserves facial features while applying the iconic anime aesthetic.")
     gr.Markdown("Generate images using EasyControl with Ghibli control LoRAs.（Due to hardware constraints, only low-resolution images can be generated. For high-resolution (1024+), please set up your own environment.）")
     
+    # Authentication input - visible at the top of the interface
+    api_token = gr.Textbox(label="API Token (Required)", type="password", value="")
+    
     gr.Markdown("**[Attention!!]**：The recommended prompts for using Ghibli Control LoRA should include the trigger words: `Ghibli Studio style, Charming hand-drawn anime-style illustration`")
     gr.Markdown("😊😊If you like this demo, please give us a star (github: [EasyControl](https://github.com/Xiaojiu-z/EasyControl))")
 
@@ -166,95 +134,22 @@ with gr.Blocks() as demo:
             with gr.Column():
                 single_output_image = gr.Image(label="Generated Image")
 
-        # Add examples for Single Condition Generation
+        # Add examples for Single Condition Generation (including the token)
         gr.Examples(
             examples=single_examples,
-            inputs=[prompt, spatial_img, height, width, seed, control_type],
+            inputs=[prompt, spatial_img, height, width, seed, control_type, api_token],
             outputs=single_output_image,
             fn=single_condition_generate_image,
             cache_examples=False,
             label="Single Condition Examples"
         )
 
-    # Link the buttons to the functions
+    # Link the buttons to the functions, including the token
     single_generate_btn.click(
         single_condition_generate_image,
-        inputs=[prompt, spatial_img, height, width, seed, control_type],
+        inputs=[prompt, spatial_img, height, width, seed, control_type, api_token],
         outputs=single_output_image
     )
 
-# Create a custom authentication interface
-def auth_interface(demo):
-    def check_auth(username, password):
-        # You could use a more sophisticated verification here
-        if API_TOKEN and password == API_TOKEN:
-            # Set the UI session token as a cookie
-            gr.Info("Authentication successful")
-            return {"hf_ui_auth": UI_SESSION_TOKEN}, True
-        else:
-            gr.Warning("Authentication failed")
-            return {}, False
-    
-    with gr.Blocks() as auth_block:
-        gr.Markdown("# Authentication Required")
-        gr.Markdown("This Hugging Face space requires authentication.")
-        
-        username = gr.Textbox(label="Username (any value)")
-        password = gr.Textbox(label="API Token", type="password")
-        submit_btn = gr.Button("Login")
-        
-        submit_btn.click(
-            fn=check_auth,
-            inputs=[username, password],
-            outputs=[gr.JSON(visible=False), gr.Checkbox(visible=False)],
-            _js="""
-            function(auth_result, success) {
-                if (success) {
-                    // Set the auth cookie
-                    const cookies = auth_result;
-                    for (const [key, value] of Object.entries(cookies)) {
-                        document.cookie = `${key}=${value}; path=/; max-age=86400; SameSite=Strict`;
-                    }
-                    // Refresh the page to load the actual UI
-                    window.location.reload();
-                }
-                return [auth_result, success];
-            }
-            """
-        )
-    
-    return auth_block
-
-# Launch the Gradio app with custom auth
-if API_TOKEN:
-    # Define middleware to check cookie-based auth before showing UI
-    def auth_middleware(app):
-        @app.middleware("http")
-        async def auth_middleware(request, call_next):
-            # Check for UI auth cookie
-            cookies = request.cookies
-            has_valid_auth = cookies.get("hf_ui_auth") == UI_SESSION_TOKEN
-            
-            # Check for API token in header
-            auth_header = request.headers.get("Authorization")
-            if auth_header and auth_header.startswith("Bearer "):
-                token = auth_header.replace("Bearer ", "")
-                has_valid_auth = has_valid_auth or (token == API_TOKEN)
-            
-            # If accessing the main UI without auth, redirect to auth UI
-            if not has_valid_auth and request.url.path == "/" and "text/html" in request.headers.get("accept", ""):
-                # Show auth UI instead of main UI
-                auth_ui = auth_interface(demo)
-                auth_response = await auth_ui.process_api(request)
-                return auth_response
-            
-            # Otherwise proceed normally (API validation happens in the endpoint function)
-            return await call_next(request)
-        
-        return app
-    
-    # Launch with auth middleware
-    demo.queue().launch(middleware=auth_middleware)
-else:
-    print("WARNING: No API_TOKEN set. Running without authentication!")
-    demo.queue().launch()
+# Launch the Gradio app
+demo.queue().launch()