Create import

import

@@ -0,0 +1,31 @@
+# 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'")