OpenAI Jukeboxの使い方

「OpenAI Jukebox」の使い方をまとめました。

1. OpenAI Jukebox

「OpenAI Jukebox」は、「アーティスト」「ジャンル」「歌詞」「音楽サンプルの長さ」を指定するだけで、自動的に作曲し、歌唱するwavファイルを生成するツールです。

Jukebox

以下のColabを参考にして、「OpenAI Jukebox」を試してみました。

Google Colaboratory

2. Colab Proの準備

「OpenAI Jukebox」を動かすには、無料の「Google Colab」ではメモリが足りないので、「Colab Pro」を使います。

メニュー「編集→ノートブックの設定」で「GPU」と「ハイメモリ」を選択し、GPUは「P100」または「V100」で、メモリ25GBになっていることを確認してください。

!nvidia-smi -L

GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-25768654-9b1b-da57-6ff2-88cde38a3db7)

3. データの永続化

Googleドライブに作業フォルダを作成することで、データを永続化し、インスタンスがリセットされても再開できるようにします。

# データの永続化
from google.colab import drive 
drive.mount('/content/drive')
!mkdir -p '/content/drive/My Drive/jukebox/'
%cd '/content/drive/My Drive/jukebox/'

4. OpenAI Jukeboxのインストール

以下のコマンドで、「OpenAI Jukebox」をインストールします。

# OpenAI Jukeboxのインストール
!pip install git+https://github.com/openai/jukebox.git

5. パッケージのインポート

各種パッケージをインポートします。

# パッケージのインポート
import jukebox
import torch as t
import librosa
import os
from IPython.display import Audio
from jukebox.make_models import make_vqvae, make_prior, MODELS, make_model
from jukebox.hparams import Hyperparams, setup_hparams
from jukebox.sample import sample_single_window, _sample, sample_partial_window, upsample, load_prompts
from jukebox.utils.dist_utils import setup_dist_from_mpi
from jukebox.utils.torch_utils import empty_cache
rank, local_rank, device = setup_dist_from_mpi()

6. ハイパーパラメータの準備

「model」には、モデルを指定します。「5b_lyrics」「5b」「1b_lyrics」のいずれかを指定でき、今回は一番小さな「1b_lyrics」を指定します。「hps.name」には、ニューラルネットワークの出力先となるフォルダを指定します。

実行すると、「音楽のベースジェネレータ」のダウンロードが開始されます。完了すると「0: Loading prior in eval mode」と表示されます。

# ハイパーパラメータの設定
model = '1b_lyrics' # '5b_lyrics' or '5b' or '1b_lyrics'
hps = Hyperparams()
hps.sr = 44100
hps.n_samples = 3 if model in ('5b', '5b_lyrics') else 8
hps.name = '/content/drive/My Drive/jukebox' # ニューラルネットワークの出力先
chunk_size = 16 if model in ('5b', '5b_lyrics') else 32
max_batch_size = 3 if model in ('5b', '5b_lyrics') else 16
hps.levels = 3
hps.hop_fraction = [.5,.5,.125]
vqvae, *priors = MODELS[model]
vqvae = make_vqvae(setup_hparams(vqvae, dict(sample_length = 1048576)), device)
top_prior = make_prior(setup_hparams(priors[-1], dict()), vqvae, device)

7. 音楽サンプルを使用しない音楽生成の準備

音楽生成するモードには、次の3つがあります。

(1) 音楽サンプルを使用しない音楽生成
(2) 音楽サンプルをベースに音楽生成
(3) 以前に生成した音楽をアップスケーリング

今回は「音楽サンプルを使用しない音楽生成」を行います。ランダム生成のため、うまく生成されない場合もあります。

(1) 「mode = 'ancestral'」の指定。

# 音楽サンプルを使用しない音楽生成
mode = 'ancestral'
codes_file=None
audio_file=None
prompt_length_in_seconds=None

(2) ハイパーパラメータのオブジェクトの準備。

# ハイパーパラメータのオブジェクトの準備
sample_hps = Hyperparams(dict(mode=mode, codes_file=codes_file, audio_file=audio_file, prompt_length_in_seconds=prompt_length_in_seconds))

(3) 音楽サンプルの長さの指定。
生成する音楽サンプルの全長-1〜4分の範囲の曲が適切に機能し、生成時間はサンプルの長さに比例します。この全長は、モデルが歌詞を進む速度に影響します。

sample_length_in_seconds = 50 # サンプルの長さ（秒）
hps.sample_length = (int(sample_length_in_seconds*hps.sr)//top_prior.raw_to_tokens)*top_prior.raw_to_tokens
assert hps.sample_length >= top_prior.n_ctx*top_prior.raw_to_tokens, f'Please choose a larger sampling rate'

(4) アーティスト、ジャンル、歌詞の指定。

# アーティスト、ジャンル、歌詞の指定
metas = [dict(artist = "Rick Astley",
    genre = "Pop",
    total_length = hps.sample_length,
    offset = 0,
    lyrics = """We're no strangers to love
You know the rules and so do I
A full commitment's what I'm thinking of
You wouldn't get this from any other guy

I just wanna tell you how I'm feeling
Gotta make you understand

Never gonna give you up
Never gonna let you down
Never gonna run around and desert you
Never gonna make you cry
Never gonna say goodbye
Never gonna tell a lie and hurt you

We've known each other for so long
Your heart's been aching, but
You're too shy to say it
Inside, we both know what's been going on
We know the game and we're gonna play it

And if you ask me how I'm feeling
Don't tell me you're too blind to see

Never gonna give you up
Never gonna let you down
Never gonna run around and desert you
Never gonna make you cry
Never gonna say goodbye
Never gonna tell a lie and hurt you

Never gonna give you up
Never gonna let you down
Never gonna run around and desert you
Never gonna make you cry
Never gonna say goodbye
Never gonna tell a lie and hurt you

(Ooh, give you up)
(Ooh, give you up)
Never gonna give, never gonna give
(Give you up)
Never gonna give, never gonna give
(Give you up)

We've known each other for so long
Your heart's been aching, but
You're too shy to say it
Inside, we both know what's been going on
We know the game and we're gonna play it

I just wanna tell you how I'm feeling
Gotta make you understand

Never gonna give you up
Never gonna let you down
Never gonna run around and desert you
Never gonna make you cry
Never gonna say goodbye
Never gonna tell a lie and hurt you

Never gonna give you up
Never gonna let you down
Never gonna run around and desert you
Never gonna make you cry
Never gonna say goodbye
Never gonna tell a lie and hurt you

Never gonna give you up
Never gonna let you down
Never gonna run around and desert you
Never gonna make you cry
Never gonna say goodbye
Never gonna tell a lie and hurt you
""",
           ),
         ] * hps.n_samples
labels = [None, None, top_prior.labeller.get_batch_labels(metas, 'cuda')]

サポートされているアーティストは、以下を参照。

openai/jukebox

サポートされているジャンルは以下を参照。

openai/jukebox

(5) サンプリング温度を調整。
.98または.99が良さそうです。

# サンプリング温度の調整
sampling_temperature = .98

lower_batch_size = 16
max_batch_size = 3 if model in ('5b', '5b_lyrics') else 16
lower_level_chunk_size = 32
chunk_size = 16 if model in ('5b', '5b_lyrics') else 32
sampling_kwargs = [dict(temp=.99, fp16=True, max_batch_size=lower_batch_size,
        chunk_size=lower_level_chunk_size),
    dict(temp=0.99, fp16=True, max_batch_size=lower_batch_size,
        chunk_size=lower_level_chunk_size),
    dict(temp=sampling_temperature, fp16=True, 
        max_batch_size=max_batch_size, chunk_size=chunk_size)]

8. サンプリング

サンプリングを実行します。最初にトップレベル（level_2）を生成し、入った曲ができそうであれば、、次に最初のアップサンプリング（level_1）、2番目のアップサンプリング（level_0）を生成します。

(1) トップレベル（level_2）の生成。
この処理には時間がかかります（音楽サンプル20秒あたり約10分）。

# トップレベル（level_2）の生成。
if sample_hps.mode == 'ancestral':
  zs = [t.zeros(hps.n_samples,0,dtype=t.long, device='cuda') for _ in range(len(priors))]
  zs = _sample(zs, labels, sampling_kwargs, [None, None, top_prior], [2], hps)
elif sample_hps.mode == 'upsample':
  assert sample_hps.codes_file is not None
  data = t.load(sample_hps.codes_file, map_location='cpu')
  zs = [z.cuda() for z in data['zs']]
  assert zs[-1].shape[0] == hps.n_samples, f"Expected bs = {hps.n_samples}, got {zs[-1].shape[0]}"
  del data
  print('Falling through to the upsample step later in the notebook.')
elif sample_hps.mode == 'primed':
  assert sample_hps.audio_file is not None
  audio_files = sample_hps.audio_file.split(',')
  duration = (int(sample_hps.prompt_length_in_seconds*hps.sr)//top_prior.raw_to_tokens)*top_prior.raw_to_tokens
  x = load_prompts(audio_files, duration, hps)
  zs = top_prior.encode(x, start_level=0, end_level=len(priors), bs_chunks=x.shape[0])
  zs = _sample(zs, labels, sampling_kwargs, [None, None, top_prior], [2], hps)
else:
  raise ValueError(f'Unknown sample mode {sample_hps.mode}.')

(2) トップレベル（level_2）の視聴。
アップサンプリング段階を実行するまで、これは非常にノイズが多いように聞こえます。 要求したバッチサイズによっては、より多くのサンプルが生成される場合があります。

# トップレベル（level_2）の視聴
Audio(f'{hps.name}/level_2/item_0.wav')

(3) アップサンプラーの読み込み。

# 十分なメモリがあるローカルマシンを使用している場合は、これを＃Falseに設定# これにより、アップサンプリング段階での歌詞の配置の視覚化
if True:
  del top_prior
  empty_cache()
  top_prior=None
upsamplers = [make_prior(setup_hparams(prior, dict()), vqvae, 'cpu') for prior in priors[:-1]]
labels[:2] = [prior.labeller.get_batch_labels(metas, 'cuda') for prior in upsamplers]

(4) アップサンプリングの実行。
この処理には数時間かかります。

# アップサンプリングの実行
zs = upsample(zs, labels, sampling_kwargs, [*upsamplers, top_prior], hps)

(5) 2番目のアップサンプリング（level_0）の視聴。

# 2番目のアップサンプリング（level_0）の視聴
Audio(f'{hps.name}/level_0/item_0.wav')

(6) アップサンプリングの解放。

# アップサンプリングの解放
del upsamplers
empty_cache()

50秒をサンプリングするのに、V100で10時間かかりました。

9. 参考

OpenAI Jukebox and Google Colab “Tutorial”

BERT/GPT-3/DALL-E 自然言語処理・画像処理・音声処理 人工知能プログラミング実践入門