import gradio as gr
import os
import shutil
import spaces
import sys

# we will clone the repo and install the dependencies
# NOTE: Still fixing bugs, not release, do not try :) !
os.system('pip install -r qa_mdt/requirements.txt')
os.system('pip install xformers==0.0.26.post1')
os.system('pip install torchlibrosa==0.0.9 librosa==0.9.2')
os.system('pip install -q pytorch_lightning==2.1.3 torchlibrosa==0.0.9 librosa==0.9.2 ftfy==6.1.1 braceexpand')
os.system('pip install torch==2.3.0+cu121 torchvision==0.18.0+cu121 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121')
os.system('pip install numpy==1.23.5')

# only then import the necessary modules from qa_mdt
from qa_mdt.pipeline import MOSDiffusionPipeline


pipe = MOSDiffusionPipeline()

# this runs the pipeline with user input and saves the output as 'awesome.wav'
@spaces.GPU(duration=120)
def generate_waveform(description):
    pipe(description)

    generated_file_path = "./awesome.wav"

    if os.path.exists(generated_file_path):
        return generated_file_path
    else:
        return "Error: Failed to generate the waveform."

# gradio interface
iface = gr.Interface(
    fn=generate_waveform,
    inputs=gr.Textbox(lines=2, placeholder="Enter a music description here..."),
    outputs=gr.Audio(label="Download the Music 🎼"),
    title="Flux Music Diffusion Pipeline",
    description="Enter a music description, and the model will generate a corresponding audio waveform. Download the output as 'awesome.wav'.",
    examples=[
        ["a guitar like pop"],
        ["a man singing jazz"]
        ],
    cache_examples=True
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()