Update app.py

app.py

@@ -7,18 +7,19 @@ from PIL import Image
 # Load models
 def load_models():
     image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-    return image_to_text, summarizer
+    storyteller = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.1", max_new_tokens=200)
+    return image_to_text, storyteller
 
 # Process image to text
 def generate_caption(image, image_to_text):
     result = image_to_text(image)
     return result[0]["generated_text"] if result else "No caption generated."
 
-# Summarize text
-def summarize_text(text, summarizer):
-    summary = summarizer(text, max_length=30, min_length=10, do_sample=False)
-    return summary[0]["summary_text"] if summary else "No summary generated."
+# Generate a narrative story
+def generate_story(text, storyteller):
+    prompt = f"Write a short, engaging story based on the following description: {text}"
+    story = storyteller(prompt)
+    return story[0]["generated_text"] if story else "No story generated."
 
 # Convert text to speech
 def text_to_speech(text, filename="output.mp3"):
@@ -28,9 +29,9 @@ def text_to_speech(text, filename="output.mp3"):
 
 # Main Streamlit app
 def main():
-    st.title("AI-Powered Image Captioning, Summarization, and Speech")
-    
-    image_to_text, summarizer = load_models()
+    st.title("AI-Powered Image Captioning and Storytelling")
+
+    image_to_text, storyteller = load_models()
 
     uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
 
@@ -46,13 +47,13 @@ def main():
             st.write("### Image Caption:")
             st.write(caption)
 
-        with st.spinner("Summarizing caption..."):
-            summary = summarize_text(caption, summarizer)
-            st.write("### Summary:")
-            st.write(summary)
+        with st.spinner("Generating story..."):
+            story = generate_story(caption, storyteller)
+            st.write("### Generated Story:")
+            st.write(story)
 
         with st.spinner("Generating speech..."):
-            audio_file = text_to_speech(summary)
+            audio_file = text_to_speech(story)
             st.audio(audio_file, format="audio/mp3")
 
 if __name__ == "__main__":