Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image
 # Load models
 def load_models():
     image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    storyteller = pipeline("text-generation", model="EleutherAI/gpt-neo-2.7B", max_new_tokens=200)
+    storyteller = pipeline("text-generation", model="databricks/dolly-v2-3b", max_new_tokens=300)
     return image_to_text, storyteller
 
 # Process image to text
@@ -17,8 +17,8 @@ def generate_caption(image, image_to_text):
 
 # Generate a narrative story
 def generate_story(text, storyteller):
-    prompt = f"Write a short, engaging story based on this image description: {text}"
-    story = storyteller(prompt, do_sample=True, temperature=0.7, max_new_tokens=200)
+    prompt = f"Write a short and engaging story inspired by this image description: {text}"
+    story = storyteller(prompt, do_sample=True, temperature=0.7, max_new_tokens=300)
     return story[0]["generated_text"] if story else "No story generated."
 
 # Convert text to speech
@@ -36,11 +36,8 @@ def main():
     uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
 
     if uploaded_file is not None:
-        # Convert uploaded file to a PIL image
         image = Image.open(uploaded_file)
-
-        # Display the uploaded image
-        st.image(image, caption="Uploaded Image", use_column_width=True)
+        st.image(image, caption="Uploaded Image", use_container_width=True)
 
         with st.spinner("Generating caption..."):
             caption = generate_caption(image, image_to_text)