Update app.py

app.py

@@ -6,8 +6,6 @@ import torch
 from PIL import Image
 from tqdm import tqdm
 import gradio as gr
-import base64
-import io
 
 from safetensors.torch import save_file
 from src.pipeline import FluxPipeline
@@ -19,9 +17,8 @@ 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_API_TOKEN")
+API_TOKEN = os.environ.get("HF_TOKEN")
 
-# 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)
 pipe.transformer = transformer
@@ -31,100 +28,52 @@ def clear_cache(transformer):
     for name, attn_processor in transformer.attn_processors.items():
         attn_processor.bank_kv.clear()
 
-# 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=None):
-    # Check authentication
-    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=""):
+    # Check if API token is required and valid
+    if API_TOKEN and api_token != API_TOKEN:
+        return "ERROR: Invalid API token. Please provide a valid token to generate images."
     
-    try:
-        # Set the control type
-        if control_type == "Ghibli":
-            lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
-        set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
-        
-        # Process the image
-        spatial_imgs = [spatial_img] if spatial_img else []
-        image = pipe(
-            prompt,
-            height=int(height),
-            width=int(width),
-            guidance_scale=3.5,
-            num_inference_steps=25,
-            max_sequence_length=512,
-            generator=torch.Generator("cpu").manual_seed(seed), 
-            subject_images=[],
-            spatial_images=spatial_imgs,
-            cond_size=512,
-        ).images[0]
-        clear_cache(pipe.transformer)
-        
-        # We'll always return the PIL image for UI
-        # The API will extract base64 from the returned image
-        return image
+    # Set the control type
+    if control_type == "Ghibli":
+        lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
+    set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
     
-    except Exception as e:
-        error_msg = f"Error during image generation: {str(e)}"
-        print(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}
+    # Process the image
+    spatial_imgs = [spatial_img] if spatial_img else []
+    image = pipe(
+        prompt,
+        height=int(height),
+        width=int(width),
+        guidance_scale=3.5,
+        num_inference_steps=25,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(seed), 
+        subject_images=[],
+        spatial_images=spatial_imgs,
+        cond_size=512,
+    ).images[0]
+    clear_cache(pipe.transformer)
+    return image
 
 # Define the Gradio interface components
 control_types = ["Ghibli"]
 
-# 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],
-    ["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.")
+    
+    # Only show token field if API token is required
+    if API_TOKEN:
+        gr.Markdown("⚠️ **AUTHENTICATION REQUIRED**: Please enter your API token to use this service.")
+        api_token = gr.Textbox(label="API Token", type="password", value="")
+    else:
+        api_token = gr.Textbox(visible=False, value="")  # Hidden field with empty value
+    
     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))")
 
@@ -141,31 +90,53 @@ with gr.Blocks() as demo:
             with gr.Column():
                 single_output_image = gr.Image(label="Generated Image")
 
-        # Add examples for Single Condition Generation (including the token)
+        # Set up examples (with token automatically added if present)
+        example_inputs = [prompt, spatial_img, height, width, seed, control_type]
+        if API_TOKEN:
+            # Add token to examples for convenience
+            example_data = [
+                ["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],
+            ]
+            example_inputs.append(api_token)
+        else:
+            # Use examples without token
+            example_data = [
+                ["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"],
+            ]
+
         gr.Examples(
-            examples=single_examples,
-            inputs=[prompt, spatial_img, height, width, seed, control_type, api_token],
+            examples=example_data,
+            inputs=example_inputs,
             outputs=single_output_image,
             fn=single_condition_generate_image,
             cache_examples=False,
             label="Single Condition Examples"
         )
 
-    # Link the buttons to the functions, including the token
+    # Link the buttons to the functions with API token included
+    inputs = [prompt, spatial_img, height, width, seed, control_type]
+    if API_TOKEN:
+        inputs.append(api_token)
+    
     single_generate_btn.click(
         single_condition_generate_image,
-        inputs=[prompt, spatial_img, height, width, seed, control_type, api_token],
+        inputs=inputs,
         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.launch()
+demo.queue().launch()