# Import necessary libraries
from transformers import AutoTokenizer, AutoModelForTextToWaveform
import torch
import scipy.io.wavfile as wavfile
import numpy as np

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("facebook/musicgen-small")
model = AutoModelForTextToWaveform.from_pretrained("facebook/musicgen-small")

def generate_music(prompt, duration_s=10):
    # Tokenize the input prompt
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # Generate audio
    audio_values = model.generate(**inputs, max_new_tokens=int(duration_s * model.config.sample_rate / model.config.hop_length))
    
    # Convert to numpy array and scale
    audio_np = audio_values[0].cpu().numpy()
    audio_np = (audio_np * 32767).astype(np.int16)
    
    return audio_np

# Example usage
prompt = "A catchy electronic beat with a groovy bassline"
generated_audio = generate_music(prompt)

# Save the generated audio
wavfile.write("generated_music.wav", model.config.sample_rate, generated_audio)

print(f"Music generated and saved as 'generated_music.wav'")