Update app.py

app.py

@@ -7,13 +7,12 @@ from typing import Dict
 
 # Image-to-text
 def img2txt(url: str) -> str:
-    print("Initializing captioning model...")
+    st.info("Initializing captioning model...")
     captioning_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
     
-    print("Generating text from the image...")
+    st.info("Generating text from the image...")
     text = captioning_model(url, max_new_tokens=20)[0]["generated_text"]
     
-    print(text)
     return text
 
 # Text-to-story
@@ -22,8 +21,7 @@ def txt2story(prompt: str, top_k: int, top_p: float, temperature: float) -> str:
     stream = client.chat.completions.create(
         model="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
         messages=[
-            {"role": "system", "content": '''As an experienced short story writer, write a story title and then create a meaningful story influenced by the provided prompt.
-            Ensure the story is full of positive inspiration & enthusiasm and concludes with a happy ending within 250 words. Remember the story must end within 100 words'''},
+            {"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}
         ],
         top_k=top_k,
@@ -35,12 +33,26 @@ def txt2story(prompt: str, top_k: int, top_p: float, temperature: float) -> str:
     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:
-    print("Initializing text-to-speech conversion...")
+    st.info("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}
@@ -50,6 +62,7 @@ def txt2speech(text: str) -> None:
     with open('audio_story.mp3', 'wb') as file:
         file.write(response.content)
 
+# User preferences
 def get_user_preferences() -> Dict[str, str]:
     preferences = {}
 
@@ -67,15 +80,16 @@ def get_user_preferences() -> Dict[str, str]:
     
     return preferences
 
+# Main function
 def main():
     st.set_page_config(page_title="🎨 Image-to-Audio Story 🎧", page_icon="🖼️")
     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)
@@ -93,18 +107,21 @@ 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
-            txt2speech(story)  # Converts the story to audio
+            
+            # Translate story based on user-selected language
+            translated_story = translate_story(story, preferences['language'])
+            txt2speech(translated_story)  # Converts the translated story to audio
 
             st.markdown("---")
             st.markdown("## 📜 Image Caption")
@@ -112,7 +129,7 @@ def main():
 
             st.markdown("---")
             st.markdown("## 📖 Story")
-            st.write(story)
+            st.write(translated_story)
 
             st.markdown("---")
             st.markdown("## 🎧 Audio Story")