Create app.py

app.py

@@ -0,0 +1,155 @@
+import streamlit as st
+from transformers import pipeline
+import torch
+from gtts import gTTS
+import io
+import time
+
+# Initialize session state
+if 'processed_data' not in st.session_state:
+    st.session_state.processed_data = {
+        'scenario': None,
+        'story': None,
+        'audio': None
+    }
+
+if 'progress' not in st.session_state:
+    st.session_state.progress = {
+        'current_stage': '',
+        'percent': 0,
+        'start_time': None
+    }
+
+# Page setup
+st.set_page_config(page_title="Your Image to Audio Story", page_icon="🦜")
+st.header("Turn Your Image to a Short Audio Story for Children")
+
+# Helper functions
+def update_progress(stage, percentage):
+    st.session_state.progress.update({
+        'current_stage': stage,
+        'percent': percentage,
+        'start_time': time.time() if percentage == 0 else st.session_state.progress['start_time']
+    })
+
+def get_elapsed_time():
+    elapsed = time.time() - st.session_state.progress['start_time']
+    return f"{int(elapsed//60):02d}:{int(elapsed%60):02d}"
+
+# Model loading
+@st.cache_resource
+def load_models():
+    return {
+        "img_model": pipeline("image-to-text", "cnmoro/tiny-image-captioning"),
+        "story_model": pipeline("text-generation", "Qwen/Qwen2.5-0.5B-Instruct")
+    }
+
+models = load_models()
+
+# Processing functions
+def img2text(url):
+    return models["img_model"](url)[0]["generated_text"]
+
+def text2story(text):
+    prompt = f"Generate a brief 100-word story about: {text}"
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": prompt}
+    ]
+    response = models["story_model"](
+        messages,
+        max_new_tokens=100,
+        do_sample=True,
+        temperature=0.7
+    )[0]["generated_text"]
+    return response[2]["content"]
+
+def text2audio(story_text):
+    audio_io = io.BytesIO()
+    tts = gTTS(text=story_text, lang='en', slow=False)
+    tts.write_to_fp(audio_io)
+    audio_io.seek(0)
+    return {'audio': audio_io, 'sampling_rate': 16000}
+
+# UI components
+uploaded_file = st.file_uploader("Select an Image After the Models are Loaded...")
+
+if uploaded_file is not None:
+    # Initialize progress containers
+    status_text = st.empty()
+    progress_bar = st.progress(0)
+    time_text = st.empty()
+
+    try:
+        # Save uploaded file
+        bytes_data = uploaded_file.getvalue()
+        with open(uploaded_file.name, "wb") as file:
+            file.write(bytes_data)
+
+        if st.session_state.get('current_file') != uploaded_file.name:
+            st.session_state.current_file = uploaded_file.name
+            update_progress("Starting processing...", 0)
+
+            # Display image
+            st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
+
+            # Stage 1: Image to Text
+            update_progress("🖼️ Analyzing image...", 0)
+            status_text.markdown(f"""
+                **{st.session_state.progress['current_stage']}**  
+                ⏱️ Elapsed: {get_elapsed_time()}
+            """)
+            st.session_state.processed_data['scenario'] = img2text(uploaded_file.name)
+            progress_bar.progress(33)
+
+            # Stage 2: Text to Story
+            update_progress("📖 Generating story...", 33)
+            status_text.markdown(f"""
+                **{st.session_state.progress['current_stage']}**  
+                ⏱️ Elapsed: {get_elapsed_time()}
+            """)
+            st.session_state.processed_data['story'] = text2story(
+                st.session_state.processed_data['scenario']
+            )
+            progress_bar.progress(66)
+
+            # Stage 3: Story to Audio
+            update_progress("🔊 Synthesizing audio...", 66)
+            status_text.markdown(f"""
+                **{st.session_state.progress['current_stage']}**  
+                ⏱️ Elapsed: {get_elapsed_time()}
+            """)
+            st.session_state.processed_data['audio'] = text2audio(
+                st.session_state.processed_data['story']
+            )
+            progress_bar.progress(100)
+
+            # Final status
+            status_text.success(f"""
+                ✅ Generation complete!  
+                Total time: {get_elapsed_time()}
+            """)
+
+        # Show results
+        st.subheader("Results")
+        st.write("**Caption:**", st.session_state.processed_data['scenario'])
+        st.write("**Story:**", st.session_state.processed_data['story'])
+
+    except Exception as e:
+        status_text.error(f"❌ Error: {str(e)}")
+        progress_bar.empty()
+        raise e
+
+    finally:
+        time_text.empty()
+
+# Audio playback
+if st.button("Play Audio of the Story Generated"):
+    if st.session_state.processed_data.get('audio'):
+        audio_data = st.session_state.processed_data['audio']
+        st.audio(
+            audio_data['audio'].getvalue(),
+            format="audio/mp3"
+        )
+    else:
+        st.warning("Please generate a story first!")