import gradio as gr
import torch
from diffusers import AuraFlowPipeline
import spaces
import numpy as np

pipeline = AuraFlowPipeline.from_pretrained(
    "fal/AuraFlow-v0.3",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
).to("cuda") 

STYLE_PRESETS = {
    "None": "",
    "Comic": ", in comic book style, bold outlines, vibrant colors, dynamic shading",
    "Watercolor": ", in watercolor style, soft edges, translucent colors, delicate brushstrokes",
    "Oil Painting": ", in oil painting style, rich textures, bold brushstrokes, deep colors",
    "Cyberpunk": ", in cyberpunk style, neon lights, dark atmosphere, futuristic elements",
    "Photorealistic": ", in photorealistic style, highly detailed, lifelike textures, realistic lighting"
}

examples = [
    {"prompt": "A rustic village nestled in a golden autumn valley, with rolling hills and a winding river bathed in warm light", "style": "Oil Painting"},
    {"prompt": "A majestic dragon soaring high above a range of snow-capped mountains under a golden sunset sky", "style": "Comic"},
    {"prompt": "A shiba inu on a rocky cliff overlooking a vibrant sunset ocean view", "style": "Photorealistic"},
    {"prompt": "A futuristic city skyline glowing with neon lights, towering skyscrapers, and flying cars under a stormy night", "style": "Cyberpunk"},
    
]

@spaces.GPU(duration=120)
def generate_images(
    prompt,
    negative_prompt,
    style,
    width=1024,
    height=1024,
    steps=20, 
    guidance=5.0,
    seed=1,
    num_images=1,
):
    generator = torch.Generator(device="cuda").manual_seed(seed)
    styled_prompt = f"{prompt}{STYLE_PRESETS[style]}"
    gallery = []

    for i in range(num_images):
        image = pipeline(
            prompt=styled_prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            num_inference_steps=steps,
            guidance_scale=guidance,
            generator=generator,
            output_type="pil",
        ).images[0]
        gallery.append((image, ""))

    torch.cuda.empty_cache()
    return gallery

def interface_fn(
    prompt,
    negative_prompt,
    style,
    width,
    height,
    steps,
    guidance,
    seed,
    num_images,
    randomize_seed,
    history,
    progress: gr.Progress = gr.Progress(track_tqdm=True)
):
    if not prompt:
        raise gr.Error("Please enter a prompt!")
    if randomize_seed:
        seed = np.random.randint(0, 1000000)
    
    gallery = generate_images(
        prompt=prompt,
        negative_prompt=negative_prompt,
        style=style,
        width=width,
        height=height,
        steps=steps,
        guidance=guidance,
        seed=seed,
        num_images=num_images
    )

    updated_history = update_history(gallery, history)
    return gallery, seed, updated_history

def update_history(new_images, history):
    if history is None:
        history = []
    for img in reversed(new_images):
        history.insert(0, img[0])
    return history

def clear_result():
    return gr.update(value=[]), gr.update(value=None)
    
#my custom css for layout
custom_css = """
    .gr-button {margin: 5px;}
    .output-image {border-radius: 8px;}
    #advanced_options {margin-top: 20px;}
    .style-dropdown {width: 100%; max-width: 800px;}
    .gr-textbox {width: 100%;}
    .example-row {margin-top: 20px;}
    .example-button {white-space: normal; height: auto; min-height: 60px;}
    /* Center the Generated Images gallery */
    #output-gallery {
        display: block;
        width: 100%;
        text-align: center;
    }
    #output-gallery .gallery {
        display: inline-flex;
        justify-content: center;
        align-items: center;
        flex-wrap: wrap;
        margin: 0 auto;
    }
    #output-gallery .gallery > div {
        display: flex;
        justify-content: center;
        align-items: center;
        margin: 5px;
    }
    #output-gallery img {
        display: block;
        margin: 0 auto;
    }
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as interface:
    gr.Markdown("# AuraFlow v0.3 Image Generator")
    gr.Markdown("Enter a prompt and select a style to generate images. Use the advanced settings for more control or try an example below.")
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Enter your creative prompt here",
                lines=3
            )
            neg_prompt_input = gr.Textbox(
                label="Negative Prompt",
                placeholder="What you don’t want in the image",
                lines=2
            )
            style_input = gr.Dropdown(
                choices=list(STYLE_PRESETS.keys()),
                value="None",
                label="Art Style",
                elem_classes=["style-dropdown"]
            )
            with gr.Accordion("Advanced Settings", open=False, elem_id="advanced_options"):
                width_input = gr.Slider(256, 1536, step=256, value=1024, label="Width")
                height_input = gr.Slider(256, 1536, step=256, value=1024, label="Height")
                steps_input = gr.Slider(1, 50, step=1, value=20, label="Inference Steps") 
                guidance_input = gr.Slider(0, 10, step=0.5, value=5.0, label="Guidance Scale")
                with gr.Row():
                    seed_input = gr.Number(value=1, label="Seed", visible=False)
                    randomize_seed_input = gr.Checkbox(value=True, label="Randomize Seed")
                num_images_input = gr.Slider(1, 4, step=1, value=1, label="Number of Images")
        
        with gr.Column(scale=2):
            image_output = gr.Gallery(
                label="Generated Images",
                show_label=True,
                preview=True,
                elem_id="output-gallery"
            )
            with gr.Row():
                clear_btn = gr.Button("Clear", variant="secondary")
                generate_btn = gr.Button("Generate", variant="primary")
            
            history_gallery = gr.Gallery(
                label="History",
                columns=6,
                object_fit="contain",
                interactive=False
            )

    with gr.Row(equal_height=True, elem_classes=["example-row"]):
        gr.Markdown("### Try these examples:")
    with gr.Row(equal_height=True, elem_classes=["example-row"]):
        for ex in examples:
            with gr.Column(scale=1, min_width=200):
                btn = gr.Button(ex["prompt"], variant="secondary", elem_classes=["example-button"])
                btn.click(
                    fn=lambda p=ex["prompt"], s=ex["style"]: (gr.update(value=p), gr.update(value=s)),
                    inputs=[],
                    outputs=[prompt_input, style_input]
                ).then(
                    fn=interface_fn,
                    inputs=[prompt_input, neg_prompt_input, style_input, width_input, height_input,
                            steps_input, guidance_input, seed_input, num_images_input, randomize_seed_input, history_gallery],
                    outputs=[image_output, seed_input, history_gallery]
                )
    generate_btn.click(
        fn=lambda: clear_result(),
        inputs=[],
        outputs=[image_output, seed_input]
    ).then(
        fn=interface_fn,
        inputs=[prompt_input, neg_prompt_input, style_input, width_input, height_input,
                steps_input, guidance_input, seed_input, num_images_input, randomize_seed_input, history_gallery],
        outputs=[image_output, seed_input, history_gallery]
    )

    clear_btn.click(
        fn=clear_result,
        inputs=[],
        outputs=[image_output, seed_input]
    ).then(
        fn=lambda x: (gr.update(value=[]), gr.update(value=None), x),
        inputs=[history_gallery],
        outputs=[image_output, seed_input, history_gallery]
    )

    randomize_seed_input.change(
        fn=lambda randomize: gr.update(visible=not randomize),
        inputs=randomize_seed_input,
        outputs=seed_input
    )

interface.launch(
    show_error=True,
    server_name="0.0.0.0",
    server_port=7860,
    share=True
)