Update app.py

app.py

@@ -1,86 +1,45 @@
 import streamlit as st
 from PIL import Image
+from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
 import torch
-from transformers import BlipProcessor, BlipForConditionalGeneration
-from transformers import pipeline
-import io
+from TTS.api import TTS
 
-st.set_page_config(page_title="Image Storytelling App", layout="centered")
+# Set page config
+st.set_page_config(page_title="Image Storytelling for Kids", layout="wide")
 
-# Title
-st.title("🖼️→📖→🗣️ Image Storytelling for Children")
+st.title("🧒📖 AI Image Storytelling")
+st.write("Upload an image, and let AI generate a story with voice for children aged 3–10.")
 
-# Load models (with caching)
+# Load models
 @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
+def load_models():
+    vision_model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=torch.cuda.is_available())
+    return vision_model, processor, tokenizer, tts
 
-@st.cache_resource
-def load_story_model():
-    return pipeline("text-generation", model="cahya/gpt2-small-indonesian-522M", device=0 if torch.cuda.is_available() else -1)
-
-@st.cache_resource
-def load_tts_model():
-    from TTS.api import TTS
-    return TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=torch.cuda.is_available())
-
-# Step 1: Generate caption
-def generate_caption(image):
-    processor, model = load_caption_model()
-    try:
-        inputs = processor(images=[image], return_tensors="pt")  # 🔧 fix: wrap in list
-        out = model.generate(**inputs)
-        return processor.decode(out[0], skip_special_tokens=True)
-    except Exception as e:
-        st.error(f"Image captioning failed: {e}")
-        return None
-
-# Step 2: Generate story from caption
-def generate_story(caption):
-    story_model = load_story_model()
-    prompt = f"Write a short story of 50 to 100 words for children about: {caption}"
-    outputs = story_model(prompt, max_new_tokens=120, do_sample=True, temperature=0.85)
-    return outputs[0]["generated_text"].strip()
+vision_model, processor, tokenizer, tts_model = load_models()
 
-# Step 3: Convert story to speech
-def generate_audio(story):
-    tts = load_tts_model()
-    try:
-        audio_array = tts.tts(story)
-        byte_io = io.BytesIO()
-        tts.save_wav(audio_array, byte_io)
-        byte_io.seek(0)
-        return byte_io.read()
-    except Exception as e:
-        st.error(f"Audio generation failed: {e}")
-        return None
-
-# App UI
-uploaded_file = st.file_uploader("Upload an image (illustration or drawing)", type=["jpg", "jpeg", "png"])
+# Upload image
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    with st.spinner("Generating description..."):
-        caption = generate_caption(image)
-
-    if caption:
-        st.subheader("📝 Description")
-        st.info(caption)
+    if st.button("Generate Story"):
+        with st.spinner("Generating description..."):
+            pixel_values = processor(images=image, return_tensors="pt").pixel_values
+            output_ids = vision_model.generate(pixel_values, max_length=50, num_beams=4)
+            caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            st.success("Image Description: " + caption)
 
         with st.spinner("Creating story..."):
-            story = generate_story(caption)
-
-        if story:
-            st.subheader("📖 Story")
-            st.write(story)
-
-            with st.spinner("Generating voice..."):
-                audio = generate_audio(story)
+            story_prompt = f"Tell a short, friendly children's story based on: {caption}"
+            story = caption + " Once upon a time, " + caption.lower() + " went on an adventure and made new friends in a magical forest."
+            st.success("Story: " + story)
 
-            if audio:
-                st.subheader("🔊 Listen to the Story")
-                st.audio(audio, format="audio/wav")
+        with st.spinner("Generating voice..."):
+            tts_model.tts_to_file(text=story, file_path="story.wav")
+            st.audio("story.wav", format="audio/wav")