Update app.py

app.py

@@ -1,6 +1,8 @@
 #Import part
 from transformers import pipeline
 import streamlit as st
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
+import torch
 
 # Use function for the implementation
 
@@ -14,13 +16,33 @@ def img2text(img):
 
 # text2story
 def text2story(text):
-    story_text = ""   # to be completed，见2025-02-22_class.ipynb
+    generator = pipeline("text-to-story",
+                     model="distilbert/distilgpt2")
+    story_text = generator(text,
+                            min_length=100,
+                            max_length=150,
+                            num_return_sequences=1)  
     return story_text
 
+
 # text2audio
 def text2audio(story_text):
-    audio_data = ""     # to be completed, 直接在task中指定
-    return audio_data
+
+    processor = SpeechT5Processor.from_pretrained("facebook/fastspeech2-en-ljspeech")
+    model = SpeechT5ForTextToSpeech.from_pretrained("facebook/fastspeech2-en-ljspeech")
+
+    inputs = processor(story_text, return_tensors="pt")
+    with torch.no_grad():
+        speech = model.generate_speech(inputs["input_ids"], model.config.vocoder)
+
+    audio_buffer = io.BytesIO()
+    sf.write(audio_buffer, speech.numpy(), samplerate=22050, format='WAV')
+    audio_buffer.seek(0)  
+
+    return {
+        'audio': audio_buffer.getvalue(), 
+        'sampling_rate': 22050           
+    }
 
 # program main part
 
@@ -45,19 +67,19 @@ if uploaded_file is not None:
 
     #Stage 2: Text to Story
     st.text('Generating a story...')
-    #story = text2story(scenario)
-    #st.write(story)
+    story = text2story(scenario)
+    st.write(story)
 
     #Stage 3: Story to Audio data
-    #st.text('Generating audio data...')
-    #audio_data =text2audio(story)
+    st.text('Generating audio data...')
+    audio_data =text2audio(story)
 
 
     # Play button
     if st.button("Play Audio"):
-        #st.audio(audio_data['audio'],
-        #            format="audio/wav",
-        #            start_time=0,
-        #            sample_rate = audio_data['sampling_rate'])
+        st.audio(audio_data['audio'],
+                   format="audio/wav",
+                   start_time=0,
+                   sample_rate = audio_data['sampling_rate'])
         st.audio("kids_playing_audio.wav")