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

def initialize_auraflow_pipeline():
    """Initialize and return the AuraFlowPipeline."""
    pipeline = AuraFlowPipeline.from_pretrained(
        "fal/AuraFlow-v0.3",
        torch_dtype=torch.float16,
        variant="fp16",
    ).to("cuda")
    return pipeline

@spaces.GPU(duration=200)
def generate_image(pipeline, prompt, width, height, num_inference_steps, seed, guidance_scale):
    """Generate an image using the AuraFlowPipeline."""
    generator = torch.Generator().manual_seed(seed)

    image = pipeline(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=guidance_scale,
    ).images[0]

    return image

# Initialize the pipeline once
auraflow_pipeline = initialize_auraflow_pipeline()

# Gradio interface
def gradio_generate_image(prompt, width, height, num_inference_steps, seed, guidance_scale):
    return generate_image(auraflow_pipeline, prompt, width, height, num_inference_steps, seed, guidance_scale)

# Create Gradio Blocks
with gr.Blocks() as demo:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information
    </h1>
    """)
    gr.HTML(
        """
        <h3 style='text-align: center'>
        Follow me for more!
        <a href='https://twitter.com/kadirnar_ai' target='_blank'>Twitter</a> | <a href='https://github.com/kadirnar' target='_blank'>Github</a> | <a href='https://www.linkedin.com/in/kadir-nar/' target='_blank'>Linkedin</a>  | <a href='https://www.huggingface.co/kadirnar/' target='_blank'>HuggingFace</a>
        </h3>
        """)
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your image prompt here...")
            width_input = gr.Slider(minimum=512, maximum=2048, step=64, value=1536, label="Width")
            height_input = gr.Slider(minimum=512, maximum=2048, step=64, value=768, label="Height")
            steps_input = gr.Slider(minimum=10, maximum=100, step=1, value=50, label="Number of Inference Steps")
            seed_input = gr.Number(label="Seed", value=1)
            guidance_input = gr.Slider(minimum=1, maximum=10, step=0.1, value=3.5, label="Guidance Scale")
            generate_btn = gr.Button("Generate Image")
        
        with gr.Column():
            image_output = gr.Image(label="Generated Image")
    
    generate_btn.click(
        fn=gradio_generate_image,
        inputs=[prompt_input, width_input, height_input, steps_input, seed_input, guidance_input],
        outputs=image_output
    )

# Launch the Gradio interface
demo.launch()