Update app.py

app.py

@@ -66,7 +66,7 @@ def get_caption(image_bytes):
     and generates a caption.
     """
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    # Resize to speed up processing
+    # Resize image to 256x256 maximum for faster processing
     image.thumbnail((256, 256))
     caption = st.session_state.captioner(image)[0]["generated_text"]
     return caption
@@ -96,12 +96,32 @@ def get_story(caption):
 def get_audio(story):
     """
     Converts the generated story text into audio.
-    Splits the text into 300-character chunks to reduce repeated TTS calls,
-    concatenates the resulting audio chunks, and returns an in-memory WAV buffer.
+    Splits the text into 300-character chunks to reduce repeated TTS calls.
+    Checks each chunk, and if no valid audio is produced, creates a brief default silent audio.
     """
     chunks = textwrap.wrap(story, width=300)
-    audio_chunks = [st.session_state.tts(chunk)["audio"].squeeze() for chunk in chunks]
-    audio = np.concatenate(audio_chunks)
+    audio_chunks = []
+    for chunk in chunks:
+        try:
+            output = st.session_state.tts(chunk)
+            # Some pipelines return a list; if so, use the first element.
+            if isinstance(output, list):
+                output = output[0]
+            if "audio" in output:
+                # Ensure the audio is a numpy array and squeeze any extra dimensions.
+                audio_array = np.array(output["audio"]).squeeze()
+                audio_chunks.append(audio_array)
+        except Exception as e:
+            # Skip any chunk that raises an error.
+            continue
+
+    # If no audio was generated, produce 1 second of silence as a fallback.
+    if not audio_chunks:
+        sr = st.session_state.tts.model.config.sampling_rate
+        audio = np.zeros(sr, dtype=np.float32)
+    else:
+        audio = np.concatenate(audio_chunks)
+    
     buffer = io.BytesIO()
     sf.write(buffer, audio, st.session_state.tts.model.config.sampling_rate, format="WAV")
     buffer.seek(0)
@@ -111,7 +131,7 @@ def get_audio(story):
 uploaded_file = st.file_uploader("Choose a Picture...", type=["jpg", "jpeg", "png"])
 if uploaded_file is not None:
     try:
-        load_models()  # Make sure models are loaded
+        load_models()  # Ensure models are loaded
         image_bytes = uploaded_file.getvalue()
         # Display the uploaded image
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")