import streamlit as st
from transformers import AutoProcessor, AutoModelForCausalLM
from audiocraft.models import MusicGen

@st.cache_resource
def load_musicgen_model():
    model = MusicGen.get_pretrained("facebook/musicgen-large")  
    model.set_generation_params(duration=10)  
    return model

st.title("Mood-Based Music Generator 🎶")
st.write("Generate a personalized song based on how you feel today.")

mood = st.selectbox(
    "Select your mood today:",
    ("Happy", "Sad", "Energetic", "Calm", "Romantic", "Mysterious")
)

prompt = st.text_input("What's on your mind?", "Describe a scene or feeling")

model = load_musicgen_model()

if st.button("Generate Song"):
    description = f"{mood.lower()} {prompt.lower()}"
    
    audio_output = model.generate([description])
    
    st.audio(audio_output[0].cpu().detach().numpy(), format="audio/wav")

st.write("Enjoy your personalized music experience! 🎼")