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"


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
@spaces.GPU()
def single_condition_generate_image(prompt, spatial_img, height=768, width=768, seed=42, control_type="Ghibli", use_zero_init=False, zero_steps=1):
    # 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,
        use_zero_init=use_zero_init,
        zero_steps=int(zero_steps)
    ).images[0]
    clear_cache(pipe.transformer)
    return image

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

# Example data
single_examples = [
    # Your examples can be added here
]

# Custom CSS for modern look and feel
custom_css = """
:root {
    --primary-color: #6366f1;
    --secondary-color: #8b5cf6;
    --background-color: #f8fafc;
    --card-background: #ffffff;
    --text-color: #1e293b;
    --border-radius: 12px;
    --shadow: 0 4px 6px -1px rgba(0,0,0,0.1), 0 2px 4px -1px rgba(0,0,0,0.06);
}

body {
    background-color: var(--background-color);
    color: var(--text-color);
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
}

.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto;
}

.main-container {
    padding: 2rem;
}

.app-header {
    margin-bottom: 2rem;
    text-align: center;
}

.app-title {
    font-size: 2.5rem;
    font-weight: 700;
    background: linear-gradient(90deg, var(--primary-color), var(--secondary-color));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 0.5rem;
}

.app-subtitle {
    font-size: 1.1rem;
    opacity: 0.8;
    max-width: 600px;
    margin: 0 auto;
}

.input-panel, .output-panel {
    background-color: var(--card-background);
    border-radius: var(--border-radius);
    box-shadow: var(--shadow);
    padding: 1.5rem;
    height: 100%;
}

.input-panel {
    border-top: 4px solid var(--primary-color);
}

.output-panel {
    border-top: 4px solid var(--secondary-color);
    display: flex;
    flex-direction: column;
}

.panel-title {
    font-weight: 600;
    font-size: 1.25rem;
    margin-bottom: 1.5rem;
    color: var(--primary-color);
}

.generate-btn {
    background: linear-gradient(90deg, var(--primary-color), var(--secondary-color)) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    border-radius: var(--border-radius) !important;
    padding: 0.75rem 1.5rem !important;
    margin-top: 1rem !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 6px -1px rgba(99, 102, 241, 0.4) !important;
}

.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 15px -3px rgba(99, 102, 241, 0.4) !important;
}

/* Input field styling */
.gr-input, .gr-textarea {
    border-radius: 8px !important;
    border: 1px solid #e2e8f0 !important;
    padding: 0.75rem !important;
}

.gr-input:focus, .gr-textarea:focus {
    border-color: var(--primary-color) !important;
    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.2) !important;
}

.gr-form {
    background: transparent !important;
    border: none !important;
}

.gr-slider {
    margin-top: 1rem !important;
}

.gr-check-radio {
    accent-color: var(--primary-color) !important;
}

.gr-panel {
    border: none !important;
    background: transparent !important;
}

.gr-tab-panel {
    border: none !important;
}

.gr-tab {
    border-radius: var(--border-radius) var(--border-radius) 0 0 !important;
}

.gr-tab.selected {
    background: linear-gradient(90deg, var(--primary-color), var(--secondary-color)) !important;
    color: white !important;
}

.gr-image {
    border-radius: var(--border-radius) !important;
    overflow: hidden !important;
}

.footer {
    text-align: center;
    margin-top: 2rem;
    opacity: 0.7;
    font-size: 0.9rem;
}

.parameter-row {
    display: flex;
    gap: 1rem;
    margin-bottom: 1rem;
}

.parameter-col {
    flex: 1;
}

@media (max-width: 768px) {
    .parameter-row {
        flex-direction: column;
    }
}
"""

# Create the Gradio Blocks interface
with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_classes="main-container"):
        # Header
        with gr.Column(elem_classes="app-header"):
            gr.HTML("""
                <div class="app-title">FLUX Studio: Ghibli Art Creator</div>
                <div class="app-subtitle">Transform your imagination into enchanting Ghibli-style illustrations</div>
                <div class="badges">
                    <span class="badge">AI-Powered</span>
                    <span class="badge">Ghibli Style</span>
                    <span class="badge">Engine</span>
                </div>
            """)
        
        with gr.Tab("Ghibli Style Generator"):
            with gr.Row(equal_height=True):
                with gr.Column(elem_classes="input-panel"):
                    gr.HTML('<div class="panel-title">Create Your Ghibli Art</div>')
                    
                    # Information card
                    gr.HTML("""
                    <div class="info-card">
                        <h3>🎨 Ghibli Style AI Generator</h3>
                        <p>Transform your ideas into beautiful Ghibli-inspired artwork with our AI generator.</p>
                        
                        <div class="features">
                            <div class="feature">
                                <div class="feature-icon">✨</div>
                                <div class="feature-text">
                                    <h4>Magical Transformations</h4>
                                    <p>Turn ordinary photos into extraordinary Ghibli-style illustrations</p>
                                </div>
                            </div>
                            
                            <div class="feature">
                                <div class="feature-icon">🖌️</div>
                                <div class="feature-text">
                                    <h4>Artistic Control</h4>
                                    <p>Craft detailed prompts to guide the artistic direction</p>
                                </div>
                            </div>
                            
                            <div class="feature">
                                <div class="feature-icon">🏞️</div>
                                <div class="feature-text">
                                    <h4>Studio Quality</h4>
                                    <p>Generate images inspired by the legendary Ghibli animation style</p>
                                </div>
                            </div>
                        </div>
                    </div>
                    """)
                    
                    prompt = gr.Textbox(
                        label="Describe your Ghibli artwork", 
                        placeholder="E.g., A magical forest with floating lanterns, Ghibli style with vibrant colors",
                        value="Ghibli Studio style, Charming hand-drawn anime-style illustration",
                        elem_classes="prompt-input"
                    )
                    
                    spatial_img = gr.Image(
                        label="Upload a reference image (optional)", 
                        type="pil",
                        elem_classes="image-upload"
                    )
                    
                    # Hidden parameters - not displayed but needed for the function
                    height = gr.Number(label="Height", value=768, visible=False)
                    width = gr.Number(label="Width", value=768, visible=False)
                    seed = gr.Number(label="Seed", value=42, visible=False)
                    control_type = gr.Dropdown(choices=control_types, value="Ghibli", visible=False)
                    use_zero_init = gr.Checkbox(label="Use CFG zero start", value=False, visible=False)
                    zero_steps = gr.Number(label="Zero Init Steps", value=1, visible=False)
                    
                    single_generate_btn = gr.Button("✨ Create Magic", elem_classes="generate-btn")
                
                with gr.Column(elem_classes="output-panel"):
                    gr.HTML('<div class="panel-title">Generated Masterpiece</div>')
                    single_output_image = gr.Image(label="", show_label=False)
                    
                    # Added gallery of examples


            # Add examples for Single Condition Generation
            with gr.Accordion("Examples", open=False):
                gr.Examples(
                    examples=single_examples,
                    inputs=[prompt, spatial_img, height, width, seed, control_type, use_zero_init, zero_steps],
                    outputs=single_output_image,
                    fn=single_condition_generate_image,
                    cache_examples=False,
                    label="Example Prompts"
                )
        
        # Footer
        gr.HTML("""
            <div class="footer">
                <div class="footer-content">
                    <div class="footer-section">
                        <h4>About our Studio</h4>
                        <p>This is an innovative AI image generation platform that specializes in creating Ghibli-style artwork from your descriptions and reference images.</p>
                    </div>
                    <div class="footer-section">
                        <h4>How It Works</h4>
                        <p>Our AI has been trained on thousands of Ghibli-style images to capture the essence of this beloved animation style.</p>
                    </div>
                    <div class="footer-section">
                        <h4>Credits</h4>
                        <p>Powered by Infomerica .</p>
                    </div>
                </div>
                <div class="copyright">
                    © All rights reserved
                </div>
            </div>
        """)

    # Link the button to the generation function
    single_generate_btn.click(
        single_condition_generate_image,
        inputs=[prompt, spatial_img, height, width, seed, control_type, use_zero_init, zero_steps],
        outputs=single_output_image
    )

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