Update app.py

app.py

@@ -2,57 +2,52 @@ import os
 import streamlit as st
 import requests
 from transformers import pipeline
-from together import Together
 from typing import Dict
 
 # Image-to-text
 def img2txt(url: str) -> str:
-    st.info("Initializing captioning model...")
+    print("Initializing captioning model...")
     captioning_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
     
-    st.info("Generating text from the image...")
+    print("Generating text from the image...")
     text = captioning_model(url, max_new_tokens=20)[0]["generated_text"]
     
+    print(text)
     return text
 
-# Text-to-story
+# Text-to-story generation with LLM model
 def txt2story(prompt: str, top_k: int, top_p: float, temperature: float) -> str:
+    # Load the Together API client
     client = Together(api_key=os.environ.get("TOGETHER_API_KEY"))
+
+    # Modify the prompt based on user inputs and ensure a 250-word limit
+    story_prompt = f"Write a short story of no more than 250 words based on the following prompt: {prompt}"
+
+    # Call the LLM model
     stream = client.chat.completions.create(
         model="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
         messages=[
-            {"role": "system", "content": '''You are a creative story writer. Create a meaningful, positive, and inspirational story based on the provided prompt. Ensure it stays under 250 words and ends on a happy note.'''},
-            {"role": "user", "content": prompt}
+            {"role": "system", "content": '''As an experienced short story writer, write a meaningful story influenced by the provided prompt.
+            Ensure the story is full of positive inspiration & enthusiasm and concludes with a happy ending. 
+            Ensure the story does not exceed 250 words.'''},
+            {"role": "user", "content": story_prompt}
         ],
         top_k=top_k,
         top_p=top_p,
         temperature=temperature,
         stream=True
     )
-    
+
+    # Concatenate story chunks
     story = ''
     for chunk in stream:
         story += chunk.choices[0].delta.content
 
-    # Enforce 250-word limit
-    story_words = story.split()
-    if len(story_words) > 250:
-        story = ' '.join(story_words[:250]) + '...'
-    
-    return story
-
-# Translate story
-def translate_story(story: str, target_language: str) -> str:
-    if target_language != "English":
-        st.info(f"Translating story to {target_language}...")
-        translator = pipeline("text2text-generation", model="SnypzZz/Llama2-13b-Language-translate")
-        translated_story = translator(story, forced_bos_token_id=target_language)
-        return translated_story[0]["generated_text"]
     return story
 
 # Text-to-speech
 def txt2speech(text: str) -> None:
-    st.info("Initializing text-to-speech conversion...")
+    print("Initializing text-to-speech conversion...")
     API_URL = "https://api-inference.huggingface.co/models/espnet/kan-bayashi_ljspeech_vits"
     headers = {"Authorization": f"Bearer {os.environ['HUGGINGFACEHUB_API_TOKEN']}"}
     payloads = {'inputs': text}
@@ -62,7 +57,17 @@ def txt2speech(text: str) -> None:
     with open('audio_story.mp3', 'wb') as file:
         file.write(response.content)
 
-# User preferences
+# Story translation function
+def translate_story(story: str, target_language: str) -> str:
+    # Translation pipeline
+    translation_model = pipeline("translation", model="SnypzZz/Llama2-13b-Language-translate")
+    
+    print(f"Translating the story to {target_language}...")
+    translated_story = translation_model(story, max_length=400, tgt_lang=target_language)[0]['translation_text']
+    
+    return translated_story
+
+# Get user preferences for the story
 def get_user_preferences() -> Dict[str, str]:
     preferences = {}
 
@@ -86,10 +91,10 @@ def main():
     st.title("Turn the Image into Audio Story")
 
     # Allows users to upload an image file
-    uploaded_file = st.file_uploader("📷 Upload an image...", type=["jpg", "jpeg", "png"])
+    uploaded_file = st.file_uploader("# 📷 Upload an image...", type=["jpg", "jpeg", "png"])
 
     # Parameters for LLM model (in the sidebar)
-    st.sidebar.markdown("## LLM Inference Configuration Parameters")
+    st.sidebar.markdown("# LLM Inference Configuration Parameters")
     top_k = st.sidebar.number_input("Top-K", min_value=1, max_value=100, value=5)
     top_p = st.sidebar.number_input("Top-P", min_value=0.0, max_value=1.0, value=0.8)
     temperature = st.sidebar.number_input("Temperature", min_value=0.1, max_value=2.0, value=1.5)
@@ -107,21 +112,23 @@ def main():
         st.image(uploaded_file, caption='🖼️ Uploaded Image', use_column_width=True)
 
         # Initiates AI processing and story generation
-        with st.spinner("🤖 AI is at Work! "):
+        with st.spinner("## 🤖 AI is at Work! "):
             scenario = img2txt("uploaded_image.jpg")  # Extracts text from the image
             
             # Modify the prompt to include user preferences
-            prompt = (f"Based on the image description: '{scenario}', create a {preferences['genre']} story set in {preferences['setting']}. "
-                      f"The story should have a {preferences['tone']} tone and explore the theme of {preferences['theme']}. "
-                      f"The main conflict should be {preferences['conflict']}. "
-                      f"Include {preferences['magic_tech']} as a key element. "
-                      f"The story should have a {preferences['twist']} and end with a {preferences['ending']} ending.")
+            prompt = f"Based on the image description: '{scenario}', create a {preferences['genre']} story set in {preferences['setting']}. " \
+                     f"The story should have a {preferences['tone']} tone and explore the theme of {preferences['theme']}. " \
+                     f"The main conflict should be {preferences['conflict']}. " \
+                     f"Include {preferences['magic_tech']} as a key element. " \
+                     f"The story should have a {preferences['twist']} and end with a {preferences['ending']} ending."
             
             story = txt2story(prompt, top_k, top_p, temperature)  # Generates a story based on the image text, LLM params, and user preferences
             
-            # Translate story based on user-selected language
-            translated_story = translate_story(story, preferences['language'])
-            txt2speech(translated_story)  # Converts the translated story to audio
+            # Translate the story if the user selected a non-English language
+            if preferences['language'] != "English":
+                story = translate_story(story, preferences['language'])
+            
+            txt2speech(story)  # Converts the story to audio
 
             st.markdown("---")
             st.markdown("## 📜 Image Caption")
@@ -129,7 +136,7 @@ def main():
 
             st.markdown("---")
             st.markdown("## 📖 Story")
-            st.write(translated_story)
+            st.write(story)
 
             st.markdown("---")
             st.markdown("## 🎧 Audio Story")