app.py

import spaces
import os
import json
import time
import torch #a
from PIL import Image
from tqdm import tqdm
import gradio as gr
import uuid
from datetime import datetime
from typing import List, Dict, Optional

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"

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

# Gallery storage
GALLERY_DIR = "gallery"
os.makedirs(GALLERY_DIR, exist_ok=True)
GALLERY_DB = os.path.join(GALLERY_DIR, "gallery_db.json")

# Initialize gallery database
if not os.path.exists(GALLERY_DB):
    with open(GALLERY_DB, "w") as f:
        json.dump({"images": []}, f)

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

def add_to_gallery(image: Image.Image, prompt: str, control_type: str) -> str:
    """Save image to gallery and return its path"""
    image_id = str(uuid.uuid4())
    filename = f"{image_id}.png"
    filepath = os.path.join(GALLERY_DIR, filename)
    
    image.save(filepath)
    
    # Update gallery database
    with open(GALLERY_DB, "r") as f:
        db = json.load(f)
    
    db["images"].append({
        "id": image_id,
        "filename": filename,
        "prompt": prompt,
        "control_type": control_type,
        "created_at": datetime.now().isoformat()
    })
    
    with open(GALLERY_DB, "w") as f:
        json.dump(db, f, indent=2)
    
    return filepath

def get_gallery_images() -> List[Dict]:
    """Get all gallery images from database"""
    try:
        with open(GALLERY_DB, "r") as f:
            db = json.load(f)
        return db["images"]
    except:
        return []

def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type, progress=gr.Progress()):
    # 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, device="cpu")
    
    # Process the image
    spatial_imgs = [spatial_img] if spatial_img else []
    
    progress(0, desc="Starting generation...")
    image = pipe(
        prompt,
        height=int(height),
        width=int(width),
        guidance_scale=3.5,
        num_inference_steps=15,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(seed), 
        subject_images=[],
        spatial_images=spatial_imgs,
        cond_size=512,
    ).images[0]
    
    # Save to gallery
    image_path = add_to_gallery(image, prompt, control_type)
    clear_cache(pipe.transformer)
    
    return image

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

# Example data
single_examples = [
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/00.png"), 512, 512, 5, "Ghibli"],
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/02.png"), 512, 512, 42, "Ghibli"],
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/03.png"), 512, 512, 1, "Ghibli"],
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/04.png"), 512, 512, 1, "Ghibli"],
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/06.png"), 512, 512, 1, "Ghibli"],
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/07.png"), 512, 512, 1, "Ghibli"],
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/08.png"), 512, 512, 1, "Ghibli"],
    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/09.png"), 512, 512, 1, "Ghibli"],
]

# Create the Gradio Blocks interface
with gr.Blocks() as demo:
    gr.Markdown("# Ghibli Studio Control Image Generation with EasyControl")
    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.（Running on CPU due to free tier limitations; expect slower performance and lower resolution.）")
    
    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))")

    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=512, step=64, label="Height", value=512)
                width = gr.Slider(minimum=256, maximum=512, step=64, label="Width", value=512)
                seed = gr.Number(label="Seed", value=42)
                control_type = gr.Dropdown(choices=control_types, label="Control Type")
                single_generate_btn = gr.Button("Generate Image")
            with gr.Column():
                single_output_image = gr.Image(label="Generated Image")

        gr.Examples(
            examples=single_examples,
            inputs=[prompt, spatial_img, height, width, seed, control_type],
            outputs=single_output_image,
            fn=single_condition_generate_image,
            cache_examples=False,
            label="Single Condition Examples"
        )

    with gr.Tab("Gallery"):
        gallery = gr.Gallery(
            label="Generated Images",
            show_label=True,
            elem_id="gallery"
        )
        
        refresh_btn = gr.Button("Refresh Gallery")
        
        def load_gallery():
            images = get_gallery_images()
            return [os.path.join(GALLERY_DIR, img["filename"]) for img in images]
        
        refresh_btn.click(
            fn=load_gallery,
            outputs=gallery
        )
        
        # Load gallery on page load
        demo.load(
            fn=load_gallery,
            outputs=gallery
        )

    single_generate_btn.click(
        single_condition_generate_image,
        inputs=[prompt, spatial_img, height, width, seed, control_type],
        outputs=single_output_image,
        concurrency_limit=1  # Process one at a time
    )

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