# Commented out IPython magic to ensure Python compatibility.
# %%capture
# ! pip install git+https://github.com/facebookresearch/audiocraft
# ! pip install torchvision==0.16.0
# ! pip install hf-transfer
# ! pip install gradio

import os
import gradio as gr
import torchaudio
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
from huggingface_hub import hf_hub_download

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')

def generate_music(description, melody_path):
    model.set_generation_params(duration=5)
    wav = model.generate([description])
    for idx, one_wav in enumerate(wav):
        audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
    return f'{idx}.wav'

def remix_music(description, melody_path):
    melody, sr = torchaudio.load(melody_path)
    wav = model.generate_with_chroma([description], melody[None].expand(1, -1, -1), sr)
    for idx, one_wav in enumerate(wav):
        audio_write(f'{idx}_bach', one_wav.cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
    return f'{idx}_bach.wav'

examples = [
    ["acoustic, guitar, melody, trap, d minor, 90 bpm", None],
    ["piano, jazz, upbeat, c major, 120 bpm"],
    ["cinematic, orchestra, epic", None]
]

demo = gr.Interface(
    fn=lambda description, melody_path: remix_music(description, melody_path) if melody_path else generate_music(description, None),
    inputs=[
        gr.Textbox(label="Enter Description"),
        gr.File(label="Upload Melody")
    ],
    outputs=gr.Audio(label="Generated Music"),
    examples=examples,
    title="Music Generation SongStarter App"
)

demo.launch(debug=True)