Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from gtts import gTTS
 import io
+from PIL import Image
 
 # Load the image captioning model
 caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
@@ -18,7 +19,7 @@ def generate_caption(image):
 def generate_story(caption):
     # Generate the story based on the caption
     input_ids = tokenizer.encode(caption, return_tensors="pt")
-    output = text_generation_model.generate(input_ids, max_length=200, num_return_sequences=1)
+    output = text_generation_model.generate(input_ids, max_length=100, num_return_sequences=1)
     story = tokenizer.decode(output[0], skip_special_tokens=True)
     return story
 
@@ -33,15 +34,18 @@ def convert_to_audio(story):
 def main():
     st.title("Storytelling Application")
     
-    # File uploader for the image
-    uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+    # File uploader for the image (restricted to JPG)
+    uploaded_image = st.file_uploader("Upload an image", type=["jpg"])
     
     if uploaded_image is not None:
+        # Convert the uploaded image to PIL image
+        image = Image.open(uploaded_image)
+        
         # Display the uploaded image
-        st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
         
         # Generate the caption for the image
-        caption = generate_caption(uploaded_image)
+        caption = generate_caption(image)
         st.subheader("Generated Caption:")
         st.write(caption)