Update app.py

app.py

@@ -1,4 +1,3 @@
-# import part
 import streamlit as st
 from transformers import pipeline
 import textwrap
@@ -8,6 +7,7 @@ import tempfile
 import os
 from PIL import Image
 import string
+import re
 
 # Initialize pipelines with caching
 @st.cache_resource
@@ -19,7 +19,11 @@ def load_pipelines():
 
 captioner, storyer, tts = load_pipelines()
 
-# Function part
+# Function to clean raw story
+def clean_generated_story(raw):
+    allowed_pattern = re.compile(r'[a-zA-Z0-9.,!?"\'-]+\b(?<!\b\w\b)')
+    return ' '.join(word for word in re.findall(allowed_pattern, raw) if len(word) > 1)
+
 # Function to generate content from an image
 def generate_content(image):
     pil_image = Image.open(image)
@@ -29,36 +33,34 @@ def generate_content(image):
     st.write("**🌟 What's in the picture: 🌟**")
     st.write(caption)
 
-# Create prompt for story
-prompt = (
-    f"Write a funny, interesting children's story centered on this scene: {caption}\n"
-    f"Story in third-person narrative, describing this scene exactly: {caption} "
-    f"Mention the exact place, location, or venue within {caption}. "
-    f"Avoid numbers, random letter combinations, and single-letter words.")
+    # Create prompt for story
+    prompt = (
+        f"Write a funny, interesting children's story centered on this scene: {caption}\n"
+        f"Story in third-person narrative, describing this scene exactly: {caption} "
+        f"Mention the exact place, location, or venue within {caption}. "
+        f"Avoid numbers, random letter combinations, and single-letter words.")
 
-# Generate raw story with optimized parameters
-raw = storyer(
-    prompt,
-    max_new_tokens=100,  
-    temperature=0.6,    
-    top_p=0.85,          
-    no_repeat_ngram_size=0, 
-    return_full_text=False
-)[0]["generated_text"].strip()
+    # Generate raw story
+    raw = storyer(
+        prompt,
+        max_new_tokens=100,  
+        temperature=0.6,    
+        top_p=0.85,          
+        no_repeat_ngram_size=0, 
+        return_full_text=False
+    )[0]["generated_text"].strip()
 
-# Combine cleaning and word trimming in one step
-# Use regex to keep only allowed characters and remove single-letter words
-allowed_pattern = re.compile(r'[a-zA-Z0-9.,!?"\'-]+\b(?<!\b\w\b)')
-clean_raw = ' '.join(word for word in re.findall(allowed_pattern, raw) if len(word) > 1)
+    # Generate story and audio
+    caption, story, audio_path = generate_story(raw, caption, tts)
+    return caption, story, audio_path
 
+# Function to generate story and audio
 def generate_story(raw, caption, tts):  
-    # Split into words and trim to 100 words
-    words = raw.split()  
+    # Clean and trim story
+    story = clean_generated_story(raw)
+    words = story.split()
     story = " ".join(words[:100])
     
-    # Clean the story using clean_generated_story
-    story = clean_generated_story(raw)  
-    
     # Display story in Streamlit
     st.write("**📖 Your funny story: 📖**")
     st.write(story)
@@ -75,23 +77,29 @@ def generate_story(raw, caption, tts):
     return caption, story, temp_file_path
 
 # Streamlit UI
-st.title("😎Story Maker")
+st.title("😎 Story Maker")
 st.markdown("Upload a picture, I will generate a story for you")
 
 uploaded_image = st.file_uploader("Choose your picture", type=["jpg", "jpeg", "png"])
 
-# Streamlit UI (modified image display section)
+# Display image
 if uploaded_image is None:
-    st.image("https://example.com/placeholder_image.jpg", caption="Upload your picture here!", use_container_width=True)
+    st.image("https://via.placeholder.com/300", caption="Upload your picture here!", use_container_width=True)
 else:
-    st.image(uploaded_image, caption="Your Picture ", use_container_width=True)
+    st.image(uploaded_image, caption="Your Picture", use_container_width=True)
 
 if st.button("Generate a story"):
     if uploaded_image is not None:
         with st.spinner("Processing"):
-            caption, story, audio_path = generate_content(uploaded_image)
-            st.success(" Your story is ready!😊")
-            st.audio(audio_path, format="audio/wav")
-            os.remove(audio_path)
+            try:
+                caption, story, audio_path = generate_content(uploaded_image)
+                st.success("Your story is ready! 😊")
+                st.audio(audio_path, format="audio/wav")
+                try:
+                    os.remove(audio_path)
+                except OSError as e:
+                    st.warning(f"Failed to delete temporary audio file: {e}")
+            except Exception as e:
+                st.error(f"Error generating story: {e}")
     else:
-        st.warning("Please upload a picture first! ")
+        st.warning("Please upload a picture first!")