import gradio as gr
from brain import StoryGenerator

story_generator = StoryGenerator()


def generate_story_from_image(image, model_name):
    """Wrapper function to use with Gradio interface"""
    return story_generator.generate_story_from_image(image, model_name)


iface = gr.Interface(
    fn=generate_story_from_image,
    inputs=[
        gr.Image(type="pil"),
        gr.Dropdown(
            choices=list(story_generator.text_models.keys()),
            label="Choose a model",
            value="Mistral-7B"
        )
    ],
    outputs="text",
    title="Kids Story Generator",
    description="Upload an image, choose a model, and get a kids story based on it!",
    examples=[["assets/image.jpg", "Mistral-7B"]],
)

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