import torch
from transformers import AutoProcessor, MusicgenForConditionalGeneration
import torchaudio
import gradio as gr

# Load model dan processor
device = "cuda" if torch.cuda.is_available() else "cpu"
model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small").to(device)
processor = AutoProcessor.from_pretrained("facebook/musicgen-small")

# Fungsi untuk generate musik
def generate_music(prompt):
    inputs = processor(text=[prompt], padding=True, return_tensors="pt").to(device)
    audio_values = model.generate(**inputs, max_new_tokens=1024)
    
    sample_rate = model.config.audio_encoder.sampling_rate
    audio = audio_values[0].cpu()
    
    file_name = "output.wav"
    torchaudio.save(file_name, audio.unsqueeze(0), sample_rate)
    return file_name

# Gradio UI
iface = gr.Interface(
    fn=generate_music,
    inputs=gr.Textbox(lines=2, placeholder="Masukkan prompt musik..."),
    outputs=gr.Audio(type="filepath"),
    title="MusicGen Gradio App"
)

# Jalankan Gradio
if __name__ == "__main__":
    iface.launch(share=True)