import gradio as gr
from audiocraft.models import MusicGen

# Load the MusicGen model
model = MusicGen.from_pretrained("facebook/musicgen-medium")

def generate_music(prompt):
    # Generate music based on the prompt
    generated_audio = model.generate(prompt, max_length=30)  # Adjust the length as needed
    return generated_audio

def process_song(input_audio):
    # Process the input audio file and generate extended music
    prompt = "Generate a continuation for the song based on its style."  # Example prompt
    extended_music = generate_music(prompt)
    return extended_music

# Create the Gradio interface
iface = gr.Interface(
    fn=process_song,
    inputs=gr.inputs.Audio(label="Upload Your Song", type="filepath"),
    outputs=gr.outputs.Audio(label="Extended Song"),
    title="MusicGen Song Extender",
    description="Upload a song to extend it using MusicGen. The model generates additional music based on the input."
)

if __name__ == "__main__":
    iface.launch()