import streamlit as st
import torch
import julius
import soundfile as sf
import io

from model import load_model, invert_audio

# Load the model and processor
preloaded = {}
preloaded["model"], preloaded["processor"] = load_model()
model = preloaded["model"]
processor = preloaded["processor"]

st.title("Audio Inversion with HuggingFace & Streamlit")

uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "flac"])

if uploaded_file:
    # Play the uploaded audio
    st.audio(uploaded_file, format="audio/wav")
    
    # Read the audio file
    audio, sr = sf.read(io.BytesIO(uploaded_file.getvalue()))

    # Convert audio to tensor
    audio_tensor = torch.tensor(audio).float()

    with st.spinner("Inverting audio..."):
        # Invert the audio using the modified function
        inverted_audio_tensor = invert_audio(model, processor, audio_tensor, sr)
        
        # Convert tensor back to numpy
        inverted_audio_np = inverted_audio_tensor.numpy()

        # Convert numpy audio to bytes and play
        with io.BytesIO() as out_io:
            sf.write(out_io, inverted_audio_np, sr, format="wav")
            st.audio(out_io.getvalue(), format="audio/wav")

    # Offer a download button for the inverted audio
    if st.button("Download Inverted Audio"):
        with io.BytesIO() as out_io:
            sf.write(out_io, inverted_audio_np, sr, format="wav")
            st.download_button("Download Inverted Audio", data=out_io.getvalue(), file_name="inverted_output.wav", mime="audio/wav")