Update app.py

app.py

@@ -5,22 +5,21 @@ 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")
+st.set_page_config(page_title="Image to Audio Story Generator")
 
 # Title and introduction
-st.title("Story Generator for Kids")
+st.title("Image to Audio Story Generator")
 st.write("Upload a picture and let's create a magical story!")
 
-# Initialize models with better performance
+# Initialize models with better error handling
 @st.cache_resource
 def load_models():
     try:
-        # Use smaller, faster models
         image_to_text = pipeline("image-to-text", model="microsoft/git-base-coco")
-        # Use distilgpt2 which is smaller and faster than gpt2
-        story_generator = pipeline("text-generation", model="distilgpt2")
+        story_generator = pipeline("text-generation", model="gpt2")
         return image_to_text, story_generator, None
     except Exception as e:
         return None, None, str(e)
@@ -30,6 +29,8 @@ with st.spinner("Loading models..."):
     image_to_text, story_generator, error = load_models()
     if error:
         st.error(f"Failed to load models: {error}")
+    else:
+        st.success("Models loaded successfully!")
 
 # Function to generate caption from image
 def generate_caption(image):
@@ -38,7 +39,7 @@ def generate_caption(image):
         if result and len(result) > 0:
             caption = result[0]['generated_text']
             return caption, None
-        return "An interesting image", None
+        return "An interesting image", "No caption generated"
     except Exception as e:
         return "An interesting image", str(e)
 
@@ -47,16 +48,21 @@ def generate_story(caption):
     try:
         prompt = f"Once upon a time, {caption} "
         
-        # Use more efficient parameters for faster generation
+        # Debug output
+        st.write(f"Prompt: {prompt}")
+        
+        # Generate with increased timeout and temperature
         result = story_generator(
             prompt, 
-            max_length=50,  # Reduce max length for faster generation
+            max_length=100, 
             do_sample=True, 
-            temperature=0.7,  # Lower temperature for faster results
-            top_p=0.9,
-            num_return_sequences=1
+            temperature=0.9,
+            top_p=0.95
         )
         
+        # Debug output
+        st.write(f"Generation result: {result}")
+        
         if result and len(result) > 0:
             story = result[0]['generated_text']
             
@@ -64,15 +70,17 @@ 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.", None
+        return "Story generation failed.", "No story generated"
     except Exception as e:
-        return f"Once upon a time, {caption}. The end.", str(e)
+        st.error(f"Error in story generation: {str(e)}")
+        st.error(traceback.format_exc())
+        return "Once upon a time... (Story generation failed)", str(e)
 
 # Function to convert text to speech
 def text_to_speech(text):
@@ -95,36 +103,32 @@ if uploaded_file is not None and image_to_text is not None and story_generator i
         
         # Generate button
         if st.button("Generate Story"):
-            # 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()
-            
+            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}")
+                st.write("Image caption:", caption)
+                
+                # Generate story
+                story, story_error = generate_story(caption)
+                if story_error:
+                    st.warning(f"Story generation issue: {story_error}")
+                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:
+                    st.warning(f"Audio generation issue: {audio_error}")
+                else:
+                    # Display audio
+                    st.write("### Listen to your story")
+                    st.audio(audio_file)
     except Exception as e:
-        st.error("An error occurred. Please try again with a different image.")
+        st.error(f"Error processing image: {str(e)}")
+        st.error(traceback.format_exc())
 
 st.markdown("---")
 st.write("Created for ISOM5240 Assignment 1")