Update app.py

app.py

@@ -14,7 +14,7 @@ def load_models():
     return caption_model, story_model, tts_model  # Return all three models
 
 # Function to generate story from caption
-def generate_story(caption):
+def generate_story(caption, story_model):
     # Generate a story based on the caption
     story = story_model(caption, max_length=100, num_return_sequences=1)[0]['generated_text']  # Generate the story
     return story  # Return the generated story
@@ -25,16 +25,17 @@ def text_to_audio(text, tts_model):
     return audio  # Return the audio object
 
 # Function to process the uploaded image and generate a story
-def process_image(image, caption_model):
+def process_image(image, caption_model, story_model):
     # Generate a caption from the uploaded image
-    caption = caption_model(image)[0]['generated_text']  # Get the caption from the model
+    result = caption_model(image)  # Get the result from the model
+    caption = result[0]['generated_text']  # Access the generated caption
     # Generate a story from the caption
-    story = generate_story(caption)  # Call the story generation function
+    story = generate_story(caption, story_model)  # Call the story generation function
     return caption, story  # Return both caption and story
 
 # Main part
 def main():
-    st.set_page_config(page_title="Storytelling Friend", page_icon="🦦") # Title of the application
+    st.set_page_config(page_title="Storytelling Friend", page_icon="🦦")  # Title of the application
     st.write("Upload an image to generate a story!")  # Instructions for the user
 
     # Upload image section
@@ -46,10 +47,10 @@ def main():
     if uploaded_file is not None:
         # Open and read the uploaded image
         image = Image.open(uploaded_file)  # Open the uploaded image file
-        st.image(image, caption="Uploaded Image", use_column_width=True)  # Display the uploaded image
+        st.image(image, caption="Uploaded Image", use_container_width=True)  # Display the uploaded image
 
         # Process the image and generate story
-        caption, story = process_image(image, caption_model)  # Get caption and story
+        caption, story = process_image(image, caption_model, story_model)  # Get caption and story
         st.subheader("Generated Caption:")  # Subheader for caption
         st.write(caption)  # Display the caption
         st.subheader("Generated Story:")  # Subheader for story