Update app.py

app.py

@@ -1,64 +1,42 @@
 import streamlit as st
-from transformers import pipeline
 from PIL import Image
-import io
-
-# Stage 1: Image to Text (Captioning)
-@st.cache_resource
-def load_image_caption_model():
-    return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-
-def generate_caption(image):
-    caption_model = load_image_caption_model()
-    result = caption_model(image)
-    return result[0]['generated_text']
-
-# Stage 2: Text to Story (Children-friendly)
-@st.cache_resource
-def load_story_generator():
-    return pipeline("text2text-generation", model="google/flan-t5-base", max_length=100)
-
-def text2story(description):
-    story_gen = load_story_generator()
-    prompt = f"Generate a short and imaginative children's story about: {description}"
-    story = story_gen(prompt)[0]['generated_text']
-    return story
-
-# Stage 3: Story to Speech (Lightweight & Compatible)
-@st.cache_resource
-def load_tts():
-    return pipeline("text-to-speech", model="suno/bark-small")
-
-def story_to_audio(story_text):
-    tts = load_tts()
-    output = tts(story_text)
-    audio = output["audio"]
-    sample_rate = output["sampling_rate"]
-    return audio, sample_rate
-
-# Streamlit App UI
-def main():
-    st.set_page_config(page_title="Kids Story Creator", layout="centered")
-    st.title("🧒 Kids Story Creator 📖")
-
-    uploaded_image = st.file_uploader("Upload an image (jpg/jpeg/png):", type=["jpg", "jpeg", "png"])
-
-    if uploaded_image is not None:
-        image = Image.open(uploaded_image).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-
-        with st.spinner("Generating description..."):
-            caption = generate_caption(image)
-        st.success(f"Caption: {caption}")
-
-        with st.spinner("Generating story from caption..."):
-            story = text2story(caption)
-        st.success("Here's your story:")
-        st.write(story)
-
-        with st.spinner("Converting story to audio..."):
-            audio, sample_rate = story_to_audio(story)
-        st.audio(audio, format="audio/wav", sample_rate=sample_rate)
+import tempfile
+import pyttsx3
+from transformers import pipeline
 
-if __name__ == "__main__":
-    main()
+st.set_page_config(page_title="Image Storytelling App", layout="centered")
+
+# Title
+st.title("🧒 Children's Image Storytelling App")
+
+# Stage 1: Image Upload & Caption
+st.header("Step 1: Upload an Image")
+uploaded_file = st.file_uploader("Upload an illustration", type=["jpg", "jpeg", "png"])
+
+if uploaded_file:
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    # Generate image description (simulated for this example)
+    st.subheader("Image Description")
+    description = "Children playing in the park"  # You can use BLIP or similar models here
+    st.success(f"Caption: {description}")
+
+    # Stage 2: Description to Story
+    st.header("Step 2: Generate a Short Story")
+    story_generator = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", max_new_tokens=150)
+
+    prompt = (
+        f"Write a creative short story (50-100 words) for children (ages 3–10) "
+        f"based on this image description: '{description}'. Make it imaginative, safe, and fun."
+    )
+
+    if st.button("Generate Story"):
+        with st.spinner("Generating story..."):
+            try:
+                result = story_generator(prompt)
+                story = result[0]['generated_text']
+                st.subheader("Generated Story")
+                st.write(story)
+            except Exception as e:
+                st.error(f"Story generation