Update app.py

app.py

@@ -5,7 +5,6 @@ import io
 from gtts import gTTS
 import time
 import os
-import traceback
 
 # Set page title
 st.set_page_config(page_title="Story Generator for Kids")
@@ -14,12 +13,14 @@ st.set_page_config(page_title="Story Generator for Kids")
 st.title("Story Generator for Kids")
 st.write("Upload a picture and let's create a magical story!")
 
-# Initialize models with better error handling
+# Initialize models with better performance
 @st.cache_resource
 def load_models():
     try:
+        # Use smaller, faster models
         image_to_text = pipeline("image-to-text", model="microsoft/git-base-coco")
-        story_generator = pipeline("text-generation", model="gpt2")
+        # Use distilgpt2 which is smaller and faster than gpt2
+        story_generator = pipeline("text-generation", model="distilgpt2")
         return image_to_text, story_generator, None
     except Exception as e:
         return None, None, str(e)
@@ -37,7 +38,7 @@ def generate_caption(image):
         if result and len(result) > 0:
             caption = result[0]['generated_text']
             return caption, None
-        return "An interesting image", "No caption generated"
+        return "An interesting image", None
     except Exception as e:
         return "An interesting image", str(e)
 
@@ -46,13 +47,14 @@ def generate_story(caption):
     try:
         prompt = f"Once upon a time, {caption} "
         
-        # Generate with increased timeout and temperature
+        # Use more efficient parameters for faster generation
         result = story_generator(
             prompt, 
-            max_length=100, 
+            max_length=50,  # Reduce max length for faster generation
             do_sample=True, 
-            temperature=0.9,
-            top_p=0.95
+            temperature=0.7,  # Lower temperature for faster results
+            top_p=0.9,
+            num_return_sequences=1
         )
         
         if result and len(result) > 0:
@@ -62,15 +64,15 @@ def generate_story(caption):
             words = story.split()
             if len(words) > 100:
                 words = words[:100]
-                story = " ".join(words)
-                # Add period to the end if needed
-                if not story.endswith(('.', '!', '?')):
-                    story += '.'
+            story = " ".join(words)
+            # Add period to the end if needed
+            if not story.endswith(('.', '!', '?')):
+                story += '.'
             
             return story, None
-        return "Story generation failed.", "No story generated"
+        return "Story generation failed.", None
     except Exception as e:
-        return "Once upon a time... (Story generation failed)", str(e)
+        return f"Once upon a time, {caption}. The end.", str(e)
 
 # Function to convert text to speech
 def text_to_speech(text):
@@ -93,36 +95,36 @@ if uploaded_file is not None and image_to_text is not None and story_generator i
         
         # Generate button
         if st.button("Generate Story"):
-            with st.spinner("Generating your story..."):
-                # Generate caption
-                caption, caption_error = generate_caption(image)
-                if caption_error:
-                    st.warning(f"Caption generation issue: {caption_error}", icon="⚠️")
-                
-                # Display the caption (without debug information)
-                st.write("Image caption:", caption)
-                
-                # Generate story
-                story, story_error = generate_story(caption)
-                if story_error and not st.session_state.get("deployed", True):
-                    st.warning(f"Story generation issue: {story_error}", icon="⚠️")
-                
-                # Display the story (without debug information)
-                word_count = len(story.split())
-                st.write(f"### Your Story ({word_count} words)")
-                st.write(story)
-                
-                # Generate audio
-                audio_file, audio_error = text_to_speech(story)
-                if audio_error and not st.session_state.get("deployed", True):
-                    st.warning(f"Audio generation issue: {audio_error}", icon="⚠️")
-                elif audio_file:
-                    # Display audio
-                    st.write("### Listen to your story")
-                    st.audio(audio_file)
+            # Use progress indicator instead of spinner for better UX
+            progress_bar = st.progress(0)
+            
+            # Generate caption
+            progress_bar.progress(25)
+            caption, caption_error = generate_caption(image)
+            st.write("Image caption:", caption)
+            
+            # Generate story
+            progress_bar.progress(50)
+            story, story_error = generate_story(caption)
+            word_count = len(story.split())
+            st.write(f"### Your Story ({word_count} words)")
+            st.write(story)
+            
+            # Generate audio
+            progress_bar.progress(75)
+            audio_file, audio_error = text_to_speech(story)
+            
+            if audio_file:
+                # Display audio
+                progress_bar.progress(100)
+                st.write("### Listen to your story")
+                st.audio(audio_file)
+            
+            # Clear progress bar when done
+            progress_bar.empty()
+            
     except Exception as e:
-        if not st.session_state.get("deployed", True):
-            st.error(f"Error processing image: {str(e)}")
+        st.error("An error occurred. Please try again with a different image.")
 
 st.markdown("---")
 st.write("Created for ISOM5240 Assignment 1")