Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from PIL import Image
 from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
 from gtts import gTTS
+import os
 import tempfile
 
 # Load models
@@ -14,8 +15,7 @@ def load_models():
 
 processor, blip_model, gpt2 = load_models()
 
-# Streamlit UI
-st.set_page_config(page_title="Kids Storyteller", page_icon="📖", layout="centered")
+# UI
 st.title("🖼️📖 Storyteller for Kids")
 st.write("Upload an image and let the app create and read a magical story just for kids!")
 
@@ -25,7 +25,6 @@ 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)
@@ -33,10 +32,10 @@ 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(
+            prompt,
             max_length=100,
             num_return_sequences=1,
             do_sample=True,
@@ -47,15 +46,12 @@ if uploaded_file:
             pad_token_id=50256,
             eos_token_id=50256,
         )[0]["generated_text"]
-
-        # 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
-
+        story = story_output.strip().replace('\n', ' ')
+        # Truncate to ~100 words for safety
+        story = " ".join(story.split()[:100])
         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')