Update app.py

app.py

@@ -9,7 +9,7 @@ st.write("Upload an image and watch as it’s captioned, turned into a short sto
 
 @st.cache_resource
 def load_captioner():
-    return pipeline("image-to-text", model="unography/blip-large-long-cap")
+    return pipeline("image-captioning", model="nlpconnect/vit-gpt2-image-captioning")
 
 @st.cache_resource
 def load_story_gen():
@@ -18,19 +18,19 @@ def load_story_gen():
 captioner = load_captioner()
 story_gen = load_story_gen()
 
-# 1) Upload (key='image' gives us st.session_state.image)
+# 1) Upload
 uploaded = st.file_uploader("Upload an image", type=["png","jpg","jpeg"], key="image")
 if uploaded:
     img = Image.open(uploaded)
     st.image(img, use_column_width=True)
 
-    # 2) Caption (once per upload)
+    # 2) Caption
     if "caption" not in st.session_state:
         with st.spinner("Generating caption…"):
             st.session_state.caption = captioner(img)[0]["generated_text"]
     st.write("**Caption:**", st.session_state.caption)
 
-    # 3) Story (once per upload)
+    # 3) Story
     if "story" not in st.session_state:
         with st.spinner("Spinning up a story…"):
             out = story_gen(
@@ -43,15 +43,15 @@ if uploaded:
             st.session_state.story = out[0]["generated_text"]
     st.write("**Story:**", st.session_state.story)
 
-    # 4) Pre-generate audio buffer (once per upload)
-    if "audio_buffer" not in st.session_state:
+    # 4) Pre-generate raw MP3 bytes
+    if "audio_bytes" not in st.session_state:
         with st.spinner("Generating audio…"):
             tts = gTTS(text=st.session_state.story, lang="en")
             buf = io.BytesIO()
             tts.write_to_fp(buf)
-            buf.seek(0)
-            st.session_state.audio_buffer = buf.read()
+            st.session_state.audio_bytes = buf.getvalue()
 
     # 5) Play on demand
     if st.button("🔊 Play Story Audio"):
-        st.audio(st.session_state.audio_buffer, format="audio/mp3")
+        audio_buffer = io.BytesIO(st.session_state.audio_bytes)
+        st.audio(audio_buffer, format="audio/mp3")