Update app.py

app.py

@@ -1,11 +1,12 @@
 import streamlit as st
-import transformers
 from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
 import torch
 from PIL import Image
 import io
 import numpy as np
-from kokoro import KokoroTTS  # for text-to-speech
+from kokoro import KPipeline # for text-to-speech
+from IPython.display import display, Audio
+import soundfile as sf 
 
 # Load models
 # Image-to-Text model
@@ -16,40 +17,54 @@ caption_model = AutoModelForCausalLM.from_pretrained("Ertugrul/Qwen2-VL-7B-Capti
 story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-14B")
 
 # Load the text-to-speech model
-tts_model = KokoroTTS("hexgrad/Kokoro-82M")
 
-def generate_story_audio(image_bytes):
-    """
-    Process an image to generate story audio.
+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
+
+
+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
+  
 
-    Args:
-        image_bytes: Bytes of the uploaded image
-
-    Returns:
-        audio (numpy array): Audio waveform
-        sample_rate (int): Sample rate of the audio
-    """
-    # Convert bytes to PIL Image
-    image = Image.open(io.BytesIO(image_bytes))
-
-    # Step 1: Generate caption from image
-    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)
-
-    # Step 2: Generate story from 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])
-
-    # Step 3: Generate audio from story using Kokoro TTS
-    audio, sample_rate = tts_model.generate(story)
-    return audio, sample_rate
 
 # Streamlit UI
 st.title("Image to Story Audio Generator")
@@ -62,35 +77,7 @@ if uploaded_file is not None:
     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)
-
-        # Save audio to a BytesIO object for Streamlit compatibility
-        audio_buffer = io.BytesIO()
-        # Simple WAV header for mono 32-bit float audio (minimal implementation)
-        def write_wav_header(buffer, data, sample_rate):
-            buffer.write(b'RIFF')
-            buffer.write((36 + len(data) * 4).to_bytes(4, 'little'))  # Chunk size
-            buffer.write(b'WAVE')
-            buffer.write(b'fmt ')
-            buffer.write((16).to_bytes(4, 'little'))  # Subchunk1 size
-            buffer.write((3).to_bytes(2, 'little'))  # Audio format (3 = IEEE float)
-            buffer.write((1).to_bytes(2, 'little'))  # Num channels (mono)
-            buffer.write(sample_rate.to_bytes(4, 'little'))  # Sample rate
-            buffer.write((sample_rate * 4).to_bytes(4, 'little'))  # Byte rate
-            buffer.write((4).to_bytes(2, 'little'))  # Block align
-            buffer.write((32).to_bytes(2, 'little'))  # Bits per sample
-            buffer.write(b'data')
-            buffer.write((len(data) * 4).to_bytes(4, 'little'))  # Data size
-            data.tofile(buffer)  # Write audio data
-
-        write_wav_header(audio_buffer, audio, sample_rate)
-        audio_buffer.seek(0)
-
-        # Provide audio playback and download
-        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"
-        )
+        #audio, sample_rate = generate_story_audio(image_bytes)
+        text = generate_text(image_bytes)
+        story = generate_story(text)
+        generate_audio(story)