import streamlit as st
from transformers import pipeline
import scipy.io.wavfile
import tempfile
import numpy as np

# -----------------------------------------------------------
def load_css(file_path):
    with open(file_path) as f:
        css_to_load = f.read()
        st.markdown(f"<style>{css_to_load}</style>", unsafe_allow_html=True)

# -----------------------------------------------------------
def load_model():
    # Initializes and returns the pipeline for text-to-audio generation
    return pipeline("text-to-audio", model="facebook/musicgen-small")

# -----------------------------------------------------------
def main():
    try:
        st.set_page_config(page_title="Text to Music", page_icon="🎵")
        load_css("styles/styles.css")
        st.title("Turn Your Text Into Music")

        # Load the model
        synthesizer = load_model()
        user_text = st.text_area("What type of beat would you like to hear: ")

        # Generate and play music based on user input
        if user_text:
            # Generate music with the synthesizer
            music = synthesizer(user_text, forward_params={"do_sample": True})

            # Create a temporary file to store the generated audio
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                print(f"Temporary file created at: {tmp.name}")
                
                # Check if the audio data needs scaling and convert to int16
                if music["audio"].dtype != np.int16:
                    audio_data = np.int16(music["audio"] / np.max(np.abs(music["audio"])) * 32767)
                else:
                    audio_data = music["audio"]
                
                # Write the scaled audio data to the temporary WAV file
                scipy.io.wavfile.write(tmp.name, rate=music["sampling_rate"], data=audio_data)
                
                # Use Streamlit's audio widget to display the audio player
                st.audio(tmp.name, format="audio/wav")

    except Exception as e:
        st.error(f"An error occurred: {e}")

# -----------------------------------------------------------

if __name__ == "__main__":
    main()