Update app.py

app.py

@@ -48,24 +48,26 @@ def clean_generated_story(raw_story: str) -> str:
     clean_story = " ".join(filtered_words[:100])
     return clean_story
 
-def generate_content(image):
+def get_caption(image) -> str:
+    """
+    Takes an image and returns a generated caption.
+    """
     pil_image = Image.open(image)
-    
-    # Generate caption from the image
     caption = captioner(pil_image)[0]["generated_text"]
     st.write("**🌟 What's in the picture: 🌟**")
     st.write(caption)
+    return caption
 
-    # Create prompt for the story
-    # Notice there’s no need to include the extra cleaning instructions in this prompt,
-    # because our code handles them later.
+def get_story(caption: str) -> str:
+    """
+    Takes a caption and returns a funny, bright, and playful story targeted toward young children.
+    """
     prompt = (
-        f"Write a funny,bright,and playful story for young children precisely centered on this scene {caption}\nStory:"
-        f"mention the exact place, location or venue within {caption}"
-        f"Make the story magical and exciting, with lots of friendly descriptions that little ones can enjoy."
+        f"Write a funny, bright, and playful story for young children precisely centered on this scene {caption}\nStory: "
+        f"mention the exact place, location or venue within {caption}. "
+        "Make the story magical and exciting, with lots of friendly descriptions that little ones can enjoy."
     )
     
-    # Generate raw story from the model
     raw = storyer(
         prompt,
         max_new_tokens=150,
@@ -75,22 +77,33 @@ def generate_content(image):
         return_full_text=False
     )[0]["generated_text"].strip()
     
-    # Clean the raw story using our custom function
     story = clean_generated_story(raw)
-    
     st.write("**📖 Your funny story: 📖**")
     st.write(story)
+    return story
 
-    # Generate audio for the story
+def generate_audio(story: str) -> str:
+    """
+    Converts a text story into speech audio and returns the file path for the audio.
+    """
     chunks = textwrap.wrap(story, width=200)
     audio = np.concatenate([tts(chunk)["audio"].squeeze() for chunk in chunks])
-
+    
     # Save audio to a temporary file
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
         sf.write(temp_file.name, audio, tts.model.config.sampling_rate)
         temp_file_path = temp_file.name
+    return temp_file_path
 
-    return caption, story, temp_file_path
+def generate_content(image):
+    """
+    Pipeline that takes an image, generates a caption, a story based on that caption,
+    and produces an audio file from the story.
+    """
+    caption = get_caption(image)
+    story = get_story(caption)
+    audio_path = generate_audio(story)
+    return caption, story, audio_path
 
 # Streamlit UI section
 st.title("✨ Magic Story Maker ✨")
@@ -111,4 +124,4 @@ if st.button("✨ Make My Story! ✨"):
             st.audio(audio_path, format="audio/wav")
             os.remove(audio_path)
     else:
-        st.warning("Please upload a picture first! 📸")
+        st.warning("Please upload a picture first! 📸")