import torchaudio
import gradio as gr
import torch
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write

model = MusicGen.get_pretrained('melody', device='cuda' if torch.cuda.is_available() else 'cpu')

def generate(prompt, melody_path):
    wav, sr = torchaudio.load(melody_path)
    wav = torchaudio.functional.resample(wav, sr, 32000)
    wav = wav.mean(dim=0).unsqueeze(0)  # mono
    wav = wav[:, :32000 * 30]  # clip to 30s

    model.set_generation_params(duration=30)
    output = model.generate_with_chroma([prompt], wav)

    output_path = "output_gen"
    audio_write(output_path, output[0].cpu(), 32000, format="wav")
    return f"{output_path}.wav"

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Audio(source="upload", type="filepath", label="Melody Input")
    ],
    outputs=gr.Audio(label="Generated Track"),
    title="🎵 MusicGen-Melody Generator"
)

demo.launch()