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="Bedtime Stories Generator",
    description="""This space uses image captioning and language models to generate 
    stories. Upload an image, choose a model, and get a children's story based on it!
    Images are first converted to text using a pre-trained image captioning model, 
    and then the text is used as a prompt for a language model to generate a story.
    """,
    examples=[["assets/image.jpg", "Mistral-7B"]],
)

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