Update app.py

app.py

@@ -23,6 +23,15 @@ from modules import (
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
 
+def reset_state():
+    """
+    Clear generated caption, story, and audio when image source changes.
+    """
+    for key in ["caption", "story", "audio_data", "audio_sr"]:
+        if key in st.session_state:
+            del st.session_state[key]
+
+
 def main():
     st.title("🖼️ → 📖 Image-to-Story App for Kids")
     st.write(
@@ -30,6 +39,13 @@ def main():
         "with audio playback powered by Hugging Face pipelines!"
     )
 
+    # Choose image source with callback to reset state
+    source = st.radio(
+        "Choose image source:",
+        ("Upload my own image", "Use demo image"),
+        on_change=reset_state
+    )
+
     # Load pipelines (cached) with a friendly spinner
     if "models_loaded" not in st.session_state:
         with st.spinner("Loading AI models, please wait…"):
@@ -42,14 +58,8 @@ def main():
     story_gen = st.session_state.story_gen
     tts_pipe  = st.session_state.tts_pipe
 
-    # Choose image source
-    source = st.radio(
-        "Choose image source:",
-        ("Upload my own image", "Use demo image")
-    )
-
+    # Get image object
     if source == "Use demo image":
-        # Load the bundled demo image
         img = Image.open("test_kids_playing.jpg")
     else:
         uploaded = st.file_uploader(
@@ -64,15 +74,14 @@ def main():
     # Display image
     st.image(img, use_container_width=True)
 
-    # Generate caption once
-    if "caption" not in st.session_state or source != st.session_state.get("last_source"):
+    # Generate caption
+    if "caption" not in st.session_state:
         with st.spinner("Generating image caption…"):
             st.session_state.caption = generate_caption(captioner, img)
-            st.session_state.last_source = source
     st.markdown(f"**Caption:** {st.session_state.caption}")
 
-    # Generate story once
-    if "story" not in st.session_state or source != st.session_state.get("last_source"):
+    # Generate story
+    if "story" not in st.session_state:
         with st.spinner("Creating story for kids…"):
             st.session_state.story = generate_story(
                 story_gen,
@@ -80,16 +89,14 @@ def main():
                 min_words=50,
                 max_words=100
             )
-            st.session_state.last_source = source
     st.markdown(f"**Story:** {st.session_state.story}")
 
-    # Generate TTS audio once
-    if "audio_data" not in st.session_state or source != st.session_state.get("last_source"):
+    # Generate audio
+    if "audio_data" not in st.session_state:
         with st.spinner("Synthesizing speech…"):
             audio_array, sr = generate_audio(tts_pipe, st.session_state.story)
             st.session_state.audio_data = audio_array
             st.session_state.audio_sr   = sr
-            st.session_state.last_source = source
 
     # Audio playback button
     if st.button("🔊 Play Story Audio"):