from transformers import MusicgenForConditionalGeneration
from transformers import AutoProcessor
import scipy
import streamlit as st

# set seed
import torch
torch.manual_seed(2912)

st.set_page_config(
    page_title="Plant Orchestra with GenAI",
    page_icon="🎵"
)

# initialise model
@st.cache_resource
def initialise_model():
    try:
        processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
        model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
        return processor, model
    except Exception as e:
        st.error(f"Error initializing the model: {str(e)}")
        return None, None
    
processor, model = initialise_model()
    

# Generate audio with given prompt
def generate_audio(processor, model, prompt):
    if processor is not None and model is not None:
        try:
            inputs = processor(
                text=[prompt],
                padding=True,
                return_tensors="pt",
            )
            audio_values = model.generate(**inputs.to("cpu"), do_sample=True, guidance_scale=3, max_new_tokens=256)
            return audio_values
        except Exception as e:
            st.error(f"Error generating audio: {str(e)}")
    return None

# save audio file with scipy
def save_file(model, audio_values, filename):
    sampling_rate = model.config.audio_encoder.sampling_rate
    scipy.io.wavfile.write(filename, rate=sampling_rate, data=audio_values[0, 0].cpu().numpy())



st.title("Plant Orchestra 🌻")
st.markdown("Generate music based on your own plant orchestra.")


prompt = st.text_input(label='Prompt:', value='Sunflower, roses')
if st.button("Generate Music"):
    if processor is not None and model is not None:
        with st.spinner("Generating audio..."):
            results = generate_audio(processor, model, prompt)
        if results is not None:
            sampling_rate = model.config.audio_encoder.sampling_rate
            st.write("Listen to the generated music:")
            st.audio(sample_rate=sampling_rate, data=results[0, 0].cpu().numpy(), format="audio/wav")
            
            
# Sidebar: How-to-use and Samples
st.sidebar.header("How to Use:")
st.sidebar.write("1. Enter a plant and condition (optional) in the text input. E.g. moss, 30C")
st.sidebar.write("2. Click the 'Generate Music' button to create music based on the provided input.")
st.sidebar.write("3. You can listen to the generated music and download it.")
st.sidebar.write()
st.sidebar.header('Samples 🎵')
st.sidebar.write('Holland moss')
st.sidebar.audio('sound/holland_moss.wav')
st.sidebar.write('Nerve plant')
st.sidebar.audio('sound/nerve_plant.wav')
st.sidebar.write('Artillery plant')
st.sidebar.audio('sound/artillery_plant.wav')
st.sidebar.write('Malayan moss')
st.sidebar.audio('sound/malayan_moss.wav')
st.sidebar.write('Pilea')
st.sidebar.audio('sound/pilea.wav')
st.sidebar.write('Hydrocotyle tripartita')
st.sidebar.audio('sound/hydrocotyle_tripartita.wav')
st.sidebar.write('Oak leaf fig')
st.sidebar.audio('sound/oak_leaf_fig.wav')


# Footer
st.write()
st.write()
st.write()
st.markdown("---")
st.markdown("Created with ❤️ by HS2912 W4 Group 2")