Update app.py

app.py

@@ -1,86 +1,86 @@
-# app.py
-
 import streamlit as st
 from PIL import Image
-from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM
 import torch
-import pyttsx3
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import pipeline
 import io
 
-# ----------- Stage 1: Image to Description -----------
+st.set_page_config(page_title="Image Storytelling App", layout="centered")
+
+# Title
+st.title("🖼️→📖→🗣️ Image Storytelling for Children")
 
+# Load models (with caching)
 @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 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)
-
-# ----------- Stage 2: Description to Story -----------
-
 @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 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)
+    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()
+
+# 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"])
+
+if uploaded_file:
+    image = Image.open(uploaded_file).convert("RGB")
     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)
-
-    # Stage 3
-    with st.spinner("Generating voice..."):
-        audio_data = generate_speech(story)
-    st.audio(audio_data, format='audio/mp3')
+        caption = generate_caption(image)
+
+    if caption:
+        st.subheader("📝 Description")
+        st.info(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)
+
+            if audio:
+                st.subheader("🔊 Listen to the Story")
+                st.audio(audio, format="audio/wav")