import spaces
import gradio as gr

from models.story_generator import generate_story

@spaces.GPU
def run(prompt, length):
    story = generate_story(prompt, length)
    return story


demo = gr.Interface(
    fn=run,
    inputs=[
        gr.Textbox(
            lines=3,
            label="Story Prompt",
            placeholder="Tell me a story about Akbar and Birbal..."
        ),
        gr.Dropdown(
            ["short", "medium", "long"],
            value="short",
            label="Story Length"
        )
    ],
    outputs=gr.Textbox(
        lines=20,
        label="Generated Story"
    ),
    title="Smart Cultural Storyteller"
)

demo.launch()