Update app.py

app.py

@@ -1,6 +1,6 @@
 # import part
-from transformers import pipeline
 import streamlit as st
+from transformers import pipeline
 
 # function part
 # img2text
@@ -11,66 +11,51 @@ def img2text(url):
 
 # text2story
 def text2story(text):
-    story_prompt = (
-        f"You are an expert in storytelling for children ages 3-10."
-        f"Based on the image description below:{text}\n"
-        f"Please create a funny, heartwarming, imaginative fairy tale of no less than 50 words and less than 100 words."
-        f"Output only the body of the story, do not include other content."
-    )
-    
-    story_text = pipeline("text-generation", model="meta-llama/Meta-Llama-3-8B",
-                              max_length=100, 
-                              do_sample=True, 
-                              top_p=0.95, 
-                              temperature=0.7)
-    generation = story_text(story_prompt)
-    raw_output = generation[0]['generated_text']
-    
-    if raw_output.startswith(story_prompt):
-        story_text = raw_output[len(story_prompt):].strip()
-    else:
-        story_text = raw_output
+    pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
+    story_text = pipe(text)[0]['generated_text']
     return story_text
 
 # text2audio
 def text2audio(story_text):
-    audio_data = pipeline("text-to-audio", model="facebook/musicgen-medium")
+    pipe = pipeline("text-to-audio", model="Matthijs/mms-tts-eng")
+    audio_data = pipe(story_text)
     return audio_data
 
-# main part
-st.set_page_config(page_title="Your Image to Audio Story",
-                   page_icon="🦜")
-st.header("Turn Your Image to Audio Story")
 
-# Upload image here
-uploaded_file = st.file_uploader("Select an Image...")
+def main():
+    st.set_page_config(page_title="Your Image to Audio Story", page_icon="🦜")
+    st.header("Turn Your Image to Audio Story")
+    uploaded_file = st.file_uploader("Select an Image...")
+
+    if uploaded_file is not None:
+        print(uploaded_file)
+        bytes_data = uploaded_file.getvalue()
+        with open(uploaded_file.name, "wb") as file:
+            file.write(bytes_data)
+        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+
+
+        #Stage 1: Image to Text
+        st.text('Processing img2text...')
+        scenario = img2text(uploaded_file.name)
+        st.write(scenario)
 
-if uploaded_file is not None:
-    print(uploaded_file)
-    bytes_data = uploaded_file.getvalue()
-    with open(uploaded_file.name, "wb") as file:
-        file.write(bytes_data)
-    st.image(uploaded_file, caption="Uploaded Image",
-             use_container_width=True)
+        #Stage 2: Text to Story
+        st.text('Generating a story...')
+        story = text2story(scenario)
+        st.write(story)
 
-    #Stage 1: Image to Text
-    st.text('Processing img2text...')
-    scenario = img2text(uploaded_file.name)
-    st.write(scenario)
+        #Stage 3: Story to Audio data
+        st.text('Generating audio data...')
+        audio_data =text2audio(story)
 
-    #Stage 2: Text to Story
-    st.text('Generating a story...')
-    story = text2story(scenario)
-    st.write(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'])
 
-    #Stage 3: Story to Audio data
-    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("kids_playing_audio.wav")
+if __name__ == "__main__":
+    main()