Update app.py

app.py

@@ -1,119 +1,166 @@
+# import part
 import streamlit as st
 from transformers import pipeline
+import torch
 from PIL import Image
 import io
-from gtts import gTTS
-import time
+import os
+from huggingface_hub import InferenceClient
+import numpy as np
+import base64
 
-# Set page title
-st.set_page_config(page_title="Kids Story Generator")
+# function part
+# img2text
+def img2text(image_path):
+    image_to_text = pipeline("image-to-text", model="noamrot/FuseCap_Image_Captioning")
+    text = image_to_text(image_path)[0]["generated_text"]
+    return text
 
-# Title and introduction
-st.title("Kids Story Generator")
-st.write("Upload a picture and let's create a magical story!")
+# text2story
+def text2story(text):
+    # Using Llama model through API to avoid GGUF format complexities in Streamlit
+    client = InferenceClient(model="MaziyarPanahi/Llama-3.2-1B-Instruct-GGUF")
+    
+    # Create a prompt for the story generation
+    prompt = f"""Write a fun, engaging children's story of about 100 words based on this caption: 
+    "{text}"
+    
+    The story should be suitable for kids aged 3-10 years old, with simple language, positive themes, and a clear beginning, middle, and end.
+    """
+    
+    # Generate the story
+    story_text = client.text_generation(
+        prompt,
+        max_new_tokens=250,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.2
+    )
+    
+    return story_text
 
-# Initialize models
-@st.cache_resource
-def load_models():
-    image_to_text = pipeline("image-to-text", model="microsoft/git-base-coco")
-    story_generator = pipeline("text-generation", model="gpt2")
-    return image_to_text, story_generator
+# text2audio
+def text2audio(story_text):
+    # Using Bark text-to-speech model
+    tts = pipeline("text-to-speech", model="suno/bark")
+    
+    # Generate audio with a voice suitable for children's stories
+    audio_output = tts(
+        text=story_text,
+        forward_params={"speaker": "v2/en_speaker_6", "text_temp": 0.7}
+    )
+    
+    return {
+        "audio": audio_output["audio"],
+        "sampling_rate": audio_output["sampling_rate"]
+    }
 
-image_to_text, story_generator = load_models()
+# Function to save temporary image file
+def save_uploaded_image(uploaded_file):
+    # Create a temp directory if it doesn't exist
+    if not os.path.exists("temp"):
+        os.makedirs("temp")
+    
+    # Define the path to save the image
+    image_path = os.path.join("temp", uploaded_file.name)
+    
+    # Save the image
+    with open(image_path, "wb") as f:
+        f.write(uploaded_file.getvalue())
+    
+    return image_path
 
-# Function to generate caption from image
-def generate_caption(image):
-    try:
-        caption = image_to_text(image)[0]['generated_text']
-        return caption
-    except Exception as e:
-        st.error(f"Error generating caption: {str(e)}")
-        return "children playing in a colorful park"
+# main part
+st.set_page_config(
+    page_title="Kids Storytelling Magic",
+    page_icon="📚",
+    layout="centered"
+)
 
-# Function to generate story from caption (less than 100 words)
-def generate_story(caption):
-    try:
-        prompt = f"Once upon a time, {caption} "
-        
-        # Set parameters for faster generation while keeping quality
-        story = story_generator(
-            prompt, 
-            max_length=100, 
-            do_sample=True,
-            temperature=0.7,  # Lower temperature for faster generation
-            top_p=0.9,
-            num_return_sequences=1
-        )[0]['generated_text']
-        
-        # Ensure story doesn't exceed 100 words
-        words = story.split()
-        if len(words) > 100:
-            words = words[:100]
-            story = " ".join(words)
-            # Add period to the end if needed
-            if not story.endswith(('.', '!', '?')):
-                story += '.'
-        
-        return story
-    except Exception as e:
-        st.error(f"Error generating story: {str(e)}")
-        return f"Once upon a time, {caption}. Something magical happened and everyone lived happily ever after."
+# Add some CSS for a child-friendly interface
+st.markdown("""
+    <style>
+    .main {
+        background-color: #f0f8ff;
+    }
+    h1, h2, h3 {
+        color: #1e90ff;
+    }
+    .stButton>button {
+        background-color: #ff6b6b;
+        color: white;
+        font-size: 1.2rem;
+        border-radius: 10px;
+        padding: 0.5rem 1rem;
+    }
+    </style>
+    """, unsafe_allow_html=True)
 
-# Function to convert text to speech
-def text_to_speech(text):
-    try:
-        tts = gTTS(text=text, lang='en', slow=False)
-        audio_file = "story_audio.mp3"
-        tts.save(audio_file)
-        return audio_file
-    except Exception as e:
-        st.error(f"Error generating audio: {str(e)}")
-        return None
+st.title("🧸 Kids Storytelling Magic 🦄")
+st.subheader("Upload a picture and hear a magical story!")
 
-# File uploader
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Choose a fun picture...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    try:
+    # Display a loading spinner
+    with st.spinner("Working on your magical story..."):
         # Display the uploaded image
-        image = Image.open(uploaded_file)
-        st.image(image, caption='Uploaded Image', use_container_width=True)
+        st.image(uploaded_file, caption="Your magical picture", use_column_width=True)
+        
+        # Save the image temporarily
+        image_path = save_uploaded_image(uploaded_file)
+        
+        # Stage 1: Image to Text
+        with st.spinner("Looking at your picture..."):
+            caption = img2text(image_path)
+            st.markdown("### 📝 I see...")
+            st.write(caption)
         
-        # Generate button
-        if st.button("Generate Story"):
-            # Use progress bar for better UX
-            progress_bar = st.progress(0)
-            
-            # Generate caption
-            progress_bar.progress(25)
-            st.text('Analyzing image...')
-            caption = generate_caption(image)
-            st.write("Image caption:", caption)
-            
-            # Generate story
-            progress_bar.progress(50)
-            st.text('Creating story...')
-            story = generate_story(caption)
-            word_count = len(story.split())
-            st.write(f"### Your Story ({word_count} words)")
+        # Stage 2: Text to Story
+        with st.spinner("Creating your story..."):
+            story = text2story(caption)
+            st.markdown("### 📖 Your Story")
             st.write(story)
-            
-            # Generate audio
-            progress_bar.progress(75)
-            st.text('Generating audio...')
-            audio_file = text_to_speech(story)
-            
-            # Display audio
-            if audio_file:
-                progress_bar.progress(100)
-                st.write("### Listen to your story")
-                st.audio(audio_file)
-            
-            # Clear progress when done
-            progress_bar.empty()
+        
+        # Stage 3: Story to Audio data
+        with st.spinner("Making your story speak..."):
+            try:
+                audio_data = text2audio(story)
+                
+                # Add a play button with cute icon
+                st.markdown("### 🔊 Listen to your story")
+                if st.button("🎵 Play Story"):
+                    st.audio(
+                        audio_data["audio"],
+                        format="audio/wav",
+                        start_time=0,
+                        sample_rate=audio_data["sampling_rate"]
+                    )
+            except Exception as e:
+                st.error(f"Oops! Something went wrong with the audio: {str(e)}")
+                st.write("But you can still read the story above!")
+        
+        # Clean up - delete the temporary image
+        try:
+            os.remove(image_path)
+        except:
+            pass
+
+else:
+    # Show instructions with a friendly message
+    st.markdown("""
+    ### How to use:
+    1. Click the button above to upload a picture
+    2. Wait for the magical story to appear
+    3. Press play to hear your story!
+    
+    Try pictures of animals, nature, toys, or anything fun!
+    """)
     
-    except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
+    # Show a placeholder image
+    st.image("https://placehold.co/600x400/9370db/ffffff?text=Upload+an+image+to+start+the+magic!", 
+             caption="Ready for your picture!", use_column_width=True)
 
+# Add a footer
 st.markdown("---")
-st.write("Created for ISOM5240 Assignment")
+st.markdown("Made for kids to enjoy the stories")