Update app.py

app.py

@@ -4,11 +4,10 @@ import torch
 from PIL import Image
 import io
 import numpy as np
-from kokoro import KPipeline # for text-to-speech
-from IPython.display import display, Audio
-import soundfile as sf 
+from kokoro import KPipeline  # For text-to-speech
+import soundfile as sf
 
-# Load models
+# Load models globally to avoid reloading them repeatedly
 # Image-to-Text model
 processor = AutoProcessor.from_pretrained("Ertugrul/Qwen2-VL-7B-Captioner-Relaxed")
 caption_model = AutoModelForCausalLM.from_pretrained("Ertugrul/Qwen2-VL-7B-Captioner-Relaxed")
@@ -16,55 +15,46 @@ caption_model = AutoModelForCausalLM.from_pretrained("Ertugrul/Qwen2-VL-7B-Capti
 # Text-to-Story model
 story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-14B")
 
-# Load the text-to-speech model
+# Text-to-Speech model
+audio_pipeline = KPipeline(lang_code='a')
 
-for i, (gs, ps, audio) in enumerate(audio_generator):
-    print(i)  # i => index
-    print(gs) # gs => graphemes/text
-    print(ps) # ps => phonemes
-    display(Audio(data=audio, rate=24000, autoplay=i==0))
-    sf.write(f'{i}.wav', audio, 24000) # save each audio file
+# Function to generate a caption from an image
+def generate_caption(image_bytes):
+    image = Image.open(io.BytesIO(image_bytes))
+    inputs = processor(images=image, text="Generate a caption:", return_tensors="pt")
+    outputs = caption_model.generate(**inputs)
+    caption = processor.decode(outputs[0], skip_special_tokens=True)
+    return caption
 
+# Function to generate a story from a caption
+def generate_story(caption):
+    prompt = f"Based on the description '{caption}', tell a short story for children aged 3 to 10 in no more than 100 words."
+    story_output = story_generator(prompt, max_length=150, num_return_sequences=1)
+    story = story_output[0]["generated_text"]
+    # Truncate to 100 words if necessary
+    story_words = story.split()
+    if len(story_words) > 100:
+        story = " ".join(story_words[:100])
+    return story
 
-def generate_text(image_bytes): 
-  # load image-to-text model
-  processor = AutoProcessor.from_pretrained("Ertugrul/Qwen2-VL-7B-Captioner-Relaxed")
-  caption_model = AutoModelForCausalLM.from_pretrained("Ertugrul/Qwen2-VL-7B-Captioner-Relaxed")
-  
-  # Convert bytes to PIL Image
-  image = Image.open(io.BytesIO(image_bytes))
-  
-  # Step 1: Generate text from image
-  inputs = processor(images=image, text="Generate a caption:", return_tensors="pt")
-  outputs = caption_model.generate(**inputs)
-  text = processor.decode(outputs[0], skip_special_tokens=True)
-  return text
-
-def generate_story(text): 
-  # load text-to-story model
-  story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-14B")
-
-  # Step 2: Generate story from caption
-  prompt = f"Based on the description '{text}', tell a short story for children aged 3 to 10 in no more than 100 words."
-  story_output = story_generator(prompt, max_length=150, num_return_sequences=1)
-  story = story_output[0]["generated_text"]
-  return story
-
-
-def generate_audio(story): 
-  audio_pipeline = KPipeline(lang_code='a')
-  audio_generator = audio_pipeline(
-      story, voice='af_heart', # <= change voice here
-      speed=1, split_pattern=r'\n+'
-      )
-  for i, (gs, ps, audio) in enumerate(audio_generator):
-    print(i)  # i => index
-    print(gs) # gs => graphemes/text
-    print(ps) # ps => phonemes
-    display(Audio(data=audio, rate=24000, autoplay=i==0))
-    sf.write(f'{i}.wav', audio, 24000) # save each audio file
-  
-
+# Function to generate audio from a story
+def generate_audio(story):
+    audio_generator = audio_pipeline(
+        story, voice='af_heart', speed=1, split_pattern=r'\n+'
+    )
+    audio_segments = []
+    # Collect all audio segments
+    for i, (gs, ps, audio) in enumerate(audio_generator):
+        audio_segments.append(audio)
+    if not audio_segments:
+        return None
+    # Concatenate audio segments into a single array
+    concatenated_audio = np.concatenate(audio_segments)
+    # Write to a BytesIO buffer instead of saving to disk
+    audio_buffer = io.BytesIO()
+    sf.write(audio_buffer, concatenated_audio, 24000, format='WAV')
+    audio_buffer.seek(0)
+    return audio_buffer
 
 # Streamlit UI
 st.title("Image to Story Audio Generator")
@@ -76,8 +66,28 @@ if uploaded_file is not None:
     image_bytes = uploaded_file.read()
     st.image(image_bytes, caption="Uploaded Image", use_column_width=True)
 
-    with st.spinner("Generating story audio..."):
-        #audio, sample_rate = generate_story_audio(image_bytes)
-        text = generate_text(image_bytes)
-        story = generate_story(text)
-        generate_audio(story)
+    # Generate and display caption
+    with st.spinner("Generating caption..."):
+        caption = generate_caption(image_bytes)
+    st.write("**Generated Caption:**")
+    st.write(caption)
+
+    # Generate and display story
+    with st.spinner("Generating story..."):
+        story = generate_story(caption)
+    st.write("**Generated Story:**")
+    st.write(story)
+
+    # Generate and display audio
+    with st.spinner("Generating audio..."):
+        audio_buffer = generate_audio(story)
+        if audio_buffer:
+            st.audio(audio_buffer, format="audio/wav")
+            st.download_button(
+                label="Download Story Audio",
+                data=audio_buffer,
+                file_name="story_audio.wav",
+                mime="audio/wav"
+            )
+        else:
+            st.error("Failed to generate audio.")