import streamlit as st
from transformers import pipeline
import scipy.io.wavfile
import numpy as np
import io  # Import the io module for BytesIO

# -----------------------------------------------------------
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, "max_new_tokens": 500})

            # 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"]

            # Instead of writing to a temp file, write the audio data to a BytesIO object
            audio_bytes_io = io.BytesIO()
            scipy.io.wavfile.write(audio_bytes_io, rate=music["sampling_rate"], data=audio_data)
            audio_bytes_io.seek(0)  # Go to the beginning of the BytesIO object

            # Use Streamlit's audio widget to display the audio player with the BytesIO object
            st.audio(audio_bytes_io, format="audio/wav")

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

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

if __name__ == "__main__":
    main()