Update app.py

app.py

@@ -1,42 +1,86 @@
+# app.py
+
 import streamlit as st
 from PIL import Image
-import tempfile
+from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM
+import torch
 import pyttsx3
-from transformers import pipeline
-
-st.set_page_config(page_title="Image Storytelling App", layout="centered")
+import io
 
-# Title
-st.title("🧒 Children's Image Storytelling App")
+# ----------- Stage 1: Image to Description -----------
 
-# Stage 1: Image Upload & Caption
-st.header("Step 1: Upload an Image")
-uploaded_file = st.file_uploader("Upload an illustration", type=["jpg", "jpeg", "png"])
+@st.cache_resource
+def load_caption_model():
+    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+    return processor, model
 
-if uploaded_file:
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+def generate_caption(image):
+    processor, model = load_caption_model()
+    inputs = processor(images=image, return_tensors="pt")
+    out = model.generate(**inputs)
+    return processor.decode(out[0], skip_special_tokens=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 -----------
 
-    # 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)
+@st.cache_resource
+def load_story_model():
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5")
+    model = AutoModelForCausalLM.from_pretrained("microsoft/phi-1_5")
+    return tokenizer, model
 
+def generate_story(description):
+    tokenizer, model = load_story_model()
     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."
+        f"Write a short and fun story (50-100 words) for children based on the following: {description}\n\n"
+        "Story:"
     )
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    output = model.generate(**inputs, max_new_tokens=120, do_sample=True, top_k=50, top_p=0.95)
+    story = tokenizer.decode(output[0], skip_special_tokens=True)
+    return story.split("Story:")[-1].strip()
+
+# ----------- Stage 3: Story to Speech -----------
+
+def generate_speech(story):
+    engine = pyttsx3.init()
+    engine.setProperty('rate', 150)
+    engine.setProperty('volume', 0.9)
+
+    with io.BytesIO() as audio:
+        engine.save_to_file(story, 'temp.mp3')
+        engine.runAndWait()
+        with open('temp.mp3', 'rb') as f:
+            audio_bytes = f.read()
+        return audio_bytes
+
+# ----------- Streamlit Interface -----------
+
+st.set_page_config(page_title="Children's Story Generator", layout="centered")
+
+st.title("📖 Children's Storytelling from Images")
+st.markdown("Upload an illustration and we'll turn it into a fun story with voice narration!")
+
+uploaded_image = st.file_uploader("Upload a drawing or illustration", type=["jpg", "jpeg", "png"])
+
+if uploaded_image:
+    image = Image.open(uploaded_image)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    # Stage 1
+    with st.spinner("Generating description..."):
+        description = generate_caption(image)
+    st.success("✅ Description Generated!")
+    st.markdown(f"**Image Caption:** _{description}_")
+
+    # Stage 2
+    with st.spinner("Generating children's story..."):
+        story = generate_story(description)
+    st.success("✅ Story Generated!")
+    st.markdown("**Generated Story:**")
+    st.write(story)
 
-    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
+    # Stage 3
+    with st.spinner("Generating voice..."):
+        audio_data = generate_speech(story)
+    st.audio(audio_data, format='audio/mp3')