MetaのMusicGenで音楽生成するUI

musicgen_ui.py

from transformers import AutoProcessor, MusicgenForConditionalGeneration
# import scipy

# model_name = "facebook/musicgen-small"
model_name = "facebook/musicgen-medium"
# model_name = "facebook/musicgen-large"

processor = AutoProcessor.from_pretrained(model_name)
model = MusicgenForConditionalGeneration.from_pretrained(model_name).to("cuda")

sampling_rate = model.config.audio_encoder.sampling_rate

def generate(text, token_count):
  inputs = processor(
      text=text,
      padding=True,
      return_tensors="pt",
  ).to(model.device)
  audio_values = model.generate(**inputs, max_new_tokens=token_count)
  #scipy.io.wavfile.write("musicgen.wav", rate=sampling_rate, data=audio_values.cpu().numpy())
  return sampling_rate, audio_values.cpu().numpy()

import gradio as gr

demo = gr.Interface(generate,
  inputs= ["text", gr.Slider(128, 1028, step=32)],
  outputs=gr.Audio(autoplay = True),
  examples=[["90s J-POP like Komuro", 512], 
    ["80s pop track with bassy drums and synth", 256],
    ["90s rock song with loud guitars and heavy drums",256]],
  title=model_name)

demo.launch()