Update app.py

app.py

@@ -1,28 +1,28 @@
 import gradio as gr
-from brain import img2txt, generate_story
-import os
+from brain import StoryGenerator
 
+story_generator = StoryGenerator()
 
-def generate_story_from_image(image):
-    """ Generate a story from an image."""
-    temp_image_path = "assets/image.jpg"
-    image.save(temp_image_path)
 
-    scenario = img2txt(temp_image_path)
-    story = generate_story(scenario)
-
-    os.remove(temp_image_path)
-
-    return story
+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"),
+    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 and get a kids story based on it!",
-    examples=[["assets/image.jpg"]],
+    description="Upload an image, choose a model, and get a kids story based on it!",
+    examples=[["assets/image.jpg", "Mistral-7B"]],
 )
 
 if __name__ == "__main__":