Update app.py

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 from PIL import Image
 from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
 from gtts import gTTS
-import os
 import tempfile
 
 # Load models
@@ -15,7 +14,8 @@ def load_models():
 
 processor, blip_model, gpt2 = load_models()
 
-# UI
+# Streamlit UI
+st.set_page_config(page_title="Kids Storyteller", page_icon="📖", layout="centered")
 st.title("🖼️📖 Storyteller for Kids")
 st.write("Upload an image and let the app create and read a magical story just for kids!")
 
@@ -25,6 +25,7 @@ if uploaded_file:
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
+    # Step 1: Generate Caption
     with st.spinner("Generating image caption..."):
         inputs = processor(images=image, return_tensors="pt")
         out = blip_model.generate(**inputs)
@@ -32,6 +33,7 @@ if uploaded_file:
         st.success("Caption generated!")
         st.write(f"**Caption:** {caption}")
 
+    # Step 2: Generate Story (remove prompt from output)
     with st.spinner("Writing a children's story..."):
         prompt = f"Write a short, imaginative story for children aged 3-10 about this: {caption}"
         story_output = gpt2(
@@ -46,12 +48,15 @@ if uploaded_file:
             pad_token_id=50256,
             eos_token_id=50256,
         )[0]["generated_text"]
-        story = story_output.strip().replace('\n', ' ')
-        # Truncate to ~100 words for safety
-        story = " ".join(story.split()[:100])
+
+        # Strip prompt from generated output
+        story = story_output[len(prompt):].strip() if story_output.startswith(prompt) else story_output.strip()
+        story = " ".join(story.split()[:100])  # Trim to ~100 words
+
         st.success("Story created!")
         st.write(f"**Story:**\n\n{story}")
 
+    # Step 3: Convert to Audio
     with st.spinner("Converting story to audio..."):
         try:
             tts = gTTS(text=story, lang='en')