| import torch |
| from transformers import AutoProcessor, MusicgenForConditionalGeneration |
| import soundfile as sf |
| import uuid |
|
|
| MODEL_NAME = "facebook/musicgen-small" |
|
|
| processor = None |
| model = None |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| model_status = "not_loaded" |
| model_error = None |
|
|
|
|
| def _get_model(): |
| global processor, model, model_status, model_error |
|
|
| if processor is None or model is None: |
| try: |
| model_status = "loading" |
| model_error = None |
| print("Loading MusicGen model...") |
| processor = AutoProcessor.from_pretrained(MODEL_NAME) |
| model = MusicgenForConditionalGeneration.from_pretrained(MODEL_NAME) |
| model.to(device) |
| model_status = "ready" |
| except Exception as exc: |
| model_status = "failed" |
| model_error = str(exc) |
| raise |
|
|
| return processor, model |
|
|
| def generate_music(prompt, duration): |
| processor, model = _get_model() |
|
|
| inputs = processor( |
| text=[prompt], |
| padding=True, |
| return_tensors="pt" |
| ).to(device) |
|
|
| audio_values = model.generate( |
| **inputs, |
| max_new_tokens=int(duration * 50) |
| ) |
|
|
| filename = f"/tmp/{uuid.uuid4()}.wav" |
|
|
| sampling_rate = model.config.audio_encoder.sampling_rate |
|
|
| sf.write( |
| filename, |
| audio_values[0, 0].cpu().numpy(), |
| sampling_rate |
| ) |
|
|
| return filename |
|
|
