app.py

import spaces
import os
import json
import time
import torch
from PIL import Image
from tqdm import tqdm
import gradio as gr

from safetensors.torch import save_file
from src.pipeline import FluxPipeline
from src.transformer_flux import FluxTransformer2DModel
from src.lora_helper import set_single_lora, set_multi_lora, unset_lora

# Initialize the image processor
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")

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
pipe.to("cuda")

def clear_cache(transformer):
    for name, attn_processor in transformer.attn_processors.items():
        attn_processor.bank_kv.clear()

# Define the Gradio interface with token verification
@spaces.GPU()
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:
        # Ensure height and width are divisible by 8
        height = int(height)
        width = int(width)
        
        if height % 8 != 0 or width % 8 != 0:
            # Adjust to nearest multiple of 8
            height = (height // 8) * 8
            width = (width // 8) * 8
            print(f"Dimensions adjusted to be divisible by 8: {height}x{width}")
        
        # 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=height,
            width=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
    except Exception as e:
        error_message = f"Error during generation: {str(e)}"
        print(error_message)
        return f"ERROR: {error_message}"

# Define the Gradio interface components
control_types = ["Ghibli"]

# Create the Gradio Blocks interface
with gr.Blocks() as demo:
    gr.Markdown("# Ghibli Studio Control Image Generation with EasyControl")
    
    # 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.）")
    
    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))")
    gr.Markdown("**NOTE**: Both height and width must be divisible by 8. Values will be automatically adjusted if needed.")

    with gr.Tab("Ghibli Condition Generation"):
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="Prompt", value="Ghibli Studio style, Charming hand-drawn anime-style illustration")
                spatial_img = gr.Image(label="Ghibli Image", type="pil")
                height = gr.Slider(minimum=256, maximum=1024, step=8, label="Height", value=768)
                width = gr.Slider(minimum=256, maximum=1024, step=8, label="Width", value=768)
                seed = gr.Number(label="Seed", value=42)
                control_type = gr.Dropdown(choices=control_types, label="Control Type", value="Ghibli")
                single_generate_btn = gr.Button("Generate Image")
            with gr.Column():
                single_output_image = gr.Image(label="Generated Image")

        # 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=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 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=inputs,
        outputs=single_output_image
    )

# Launch the Gradio app
demo.queue().launch()