app.py

import os
import re
import smtplib
import streamlit as st
from transformers import pipeline
from typing import Dict
from gtts import gTTS
from together import Together
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.audio import MIMEAudio
#from dotenv import load_dotenv

# Load environment variables
#load_dotenv()

# Image-to-text function
def img2txt(url: str) -> str:
    print("Initializing captioning model...")
    captioning_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
    
    print("Generating text from the image...")
    text = captioning_model(url, max_new_tokens=20)[0]["generated_text"]
    
    print(text)
    return text

# Text-to-story generation function
def txt2story(prompt: str, top_k: int, top_p: float, temperature: float) -> str:
    client = Together(api_key=os.environ.get("TOGETHER_API_KEY"))

    story_prompt = f"Write a short story of no more than 250 words based on the following prompt: {prompt}"

    stream = client.chat.completions.create(
        model="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
        messages=[
            {"role": "system", "content": '''As an experienced short story writer, write a meaningful story influenced by the provided prompt.
            Ensure the story does not exceed 250 words.'''},
            {"role": "user", "content": story_prompt}
        ],
        top_k=top_k,
        top_p=top_p,
        temperature=temperature,
        stream=True
    )

    story = ''
    for chunk in stream:
        story += chunk.choices[0].delta.content

    return story

# Text-to-speech function
def txt2speech(text: str) -> None:
    print("Converting text to speech using gTTS...")
    tts = gTTS(text=text, lang='en')
    tts.save("audio_story.mp3")

# Get user preferences function
def get_user_preferences() -> Dict[str, str]:
    preferences = {
        'continent': st.selectbox("Continent", ["North America", "Europe", "Asia", "Africa", "Australia"]),
        'genre': st.selectbox("Genre", ["Science Fiction", "Fantasy", "Mystery", "Romance"]),
        'setting': st.selectbox("Setting", ["Future", "Medieval times", "Modern day", "Alternate reality"]),
        'plot': st.selectbox("Plot", ["Hero's journey", "Solving a mystery", "Love story", "Survival"]),
        'tone': st.selectbox("Tone", ["Serious", "Light-hearted", "Humorous", "Dark"]),
        'theme': st.selectbox("Theme", ["Self-discovery", "Redemption", "Love", "Justice"]),
        'conflict': st.selectbox("Conflict Type", ["Person vs. Society", "Internal struggle", "Person vs. Nature", "Person vs. Person"]),
        'twist': st.selectbox("Mystery/Twist", ["Plot twist", "Hidden identity", "Unexpected ally/enemy", "Time paradox"]),
        'ending': st.selectbox("Ending", ["Happy", "Bittersweet", "Open-ended", "Tragic"])
    }
    return preferences

# Function to send the story via email using smtplib
def send_story_email(recipient_email: str, story_text: str, audio_file_path: str) -> bool:
    try:
        # Email configuration
        smtp_server = os.environ.get("SMTP_SERVER")  # e.g., "smtp.gmail.com"
        smtp_port = int(os.environ.get("SMTP_PORT", 587))  # e.g., 587 for TLS
        sender_email = os.environ.get("SENDER_EMAIL")  # Your email
        sender_password = os.environ.get("SENDER_PASSWORD")  # Your email password
        
        # Create the email
        msg = MIMEMultipart()
        msg['From'] = sender_email
        msg['To'] = recipient_email
        msg['Subject'] = "Your Generated Story"

        # Attach the story text
        msg.attach(MIMEText(f"Here's your generated story:\n\n{story_text}\n\nEnjoy!", 'plain'))

        # Attach the audio file
        with open(audio_file_path, 'rb') as audio_file:
            audio_part = MIMEAudio(audio_file.read(), _subtype='mp3')
            audio_part.add_header('Content-Disposition', 'attachment', filename=os.path.basename(audio_file_path))
            msg.attach(audio_part)

        # Send the email
        with smtplib.SMTP(smtp_server, smtp_port) as server:
            server.starttls()  # Upgrade to a secure connection
            server.login(sender_email, sender_password)
            server.send_message(msg)
        
        return True
        
    except Exception as e:
        print(f"Error sending email: {str(e)}")
        return False

# Basic email validation function
def validate_email(email: str) -> bool:
    pattern = r'^[\w\.-]+@[\w\.-]+\.\w+$'
    return re.match(pattern, email) is not None

# Main Streamlit application
def main():
    st.set_page_config(
        page_title="🎨 Image-to-Audio Story 🎧",
        page_icon="🖼️",
        layout="wide"
    )
    st.title("Turn the Image into Audio Story")

    # Initialize session state variables
    if "story" not in st.session_state:
        st.session_state.story = ""
    if "audio_file_path" not in st.session_state:
        st.session_state.audio_file_path = ""
    if "caption" not in st.session_state:
        st.session_state.caption = ""

    # Main content area
    col1, col2 = st.columns([2, 3])

    with col1:
        # Image upload section
        st.markdown("## 📷 Upload Image")
        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

        # Story preferences section
        st.markdown("## 🎭 Story Preferences")
        preferences = get_user_preferences()

    with col2:
        if uploaded_file is not None:
            # Reset session state variables when a new image is uploaded
            st.session_state.story = ""
            st.session_state.audio_file_path = ""

            # Display uploaded image
            st.markdown("## 🖼️ Your Image")
            bytes_data = uploaded_file.read()
            with open("uploaded_image.jpg", "wb") as file:
                file.write(bytes_data)
            st.image(uploaded_file, use_column_width=True)

            # Generate image caption as soon as image is uploaded
            with st.spinner("📝 Generating image caption..."):
                try:
                    scenario = img2txt("uploaded_image.jpg")
                    st.session_state.caption = scenario  # Store caption in session state
                    st.success("Image caption generated.")
                except Exception as e:
                    st.error(f"An error occurred while generating image caption: {str(e)}")
                    st.warning("Please try again or contact support if the problem persists.")

            # Display image caption
            if st.session_state.caption:
                st.markdown("## 📜 Image Caption")
                st.write(st.session_state.caption)

            # Generate Story button
            if st.session_state.caption:
                if st.button("🎨 Generate Story"):
                    with st.spinner("📝 Generating story..."):
                        try:
                            # Create story prompt using st.session_state.caption
                            prompt = f"""Based on the image description: '{st.session_state.caption}', 
                            create a {preferences['genre']} story set in {preferences['setting']} 
                            in {preferences['continent']}. The story should have a {preferences['tone']} 
                            tone and explore the theme of {preferences['theme']}. The main conflict 
                            should be {preferences['conflict']}. The story should have a {preferences['twist']} 
                            and end with a {preferences['ending']} ending."""
                            
                            # Generate story
                            story = txt2story(prompt, top_k=5, top_p=0.8, temperature=1.5)
                            st.session_state.story = story  # Store story in session state
                            st.success("Story generated.")

                            # Convert to audio
                            with st.spinner("🔊 Generating audio story..."):
                                txt2speech(story)
                                st.session_state.audio_file_path = "audio_story.mp3"  # Store audio path in session state
                                st.success("Audio story generated.")

                        except Exception as e:
                            st.error(f"An error occurred: {str(e)}")
                            st.warning("Please try again or contact support if the problem persists.")

        # Display results if story exists in session state
        if st.session_state.story:
            st.markdown("---")
            
            # Story text
            with st.expander("📖 Generated Story", expanded=True):
                st.write(st.session_state.story)
            
            # Audio player
            with st.expander("🎧 Audio Version", expanded=True):
                st.audio(st.session_state.audio_file_path)

            # Email section
            st.markdown("---")
            st.markdown("## 📧 Get Story via Email")
            email = st.text_input(
                "Enter your email address:",
                help="We'll send you the story text and audio file"
            )
            
            if st.button("📤 Send to Email"):
                if not email:
                    st.warning("Please enter an email address.")
                elif not validate_email(email):
                    st.error("Please enter a valid email address.")
                else:
                    with st.spinner("📨 Sending email..."):
                        if send_story_email(email, st.session_state.story, st.session_state.audio_file_path):
                            st.success(f"Email sent to: {email}")
                        else:
                            st.error("❌ Failed to send email. Please try again.")

if __name__ == '__main__':
    main()