import gradio as gr
from models import load_model, generate_image

# Load the model on startup
pipe = load_model()

def text_to_image(prompt):
    return generate_image(pipe, prompt)

with gr.Blocks(title="AI Text-to-Image Generator") as demo:
    gr.Markdown("# AI Text-to-Image Generator\nBuilt with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)")
    gr.Markdown("Generate images from text prompts using FLUX.1-dev model.")
    
    prompt_input = gr.Textbox(
        label="Enter your prompt",
        placeholder="A beautiful landscape with mountains and a lake...",
        lines=3
    )
    
    generate_btn = gr.Button("Generate Image")
    
    output_image = gr.Image(label="Generated Image")
    
    generate_btn.click(
        fn=text_to_image,
        inputs=[prompt_input],
        outputs=[output_image]
    )

if __name__ == "__main__":
    demo.launch()