MLXでSunoのBarkモデルの移植を試してみた

2024年7月24日

AppleのMLフレームワーク、MLXでのSunoのBarkモデルの移植を試してみました。

Barkは、音声や背景音楽などのオーディオを生成できるトランスフォーマーベースのテキストからオーディオへのモデルです。

MLXでのSunoのBarkモデルの移植方法（基本)

以下の手順で簡単に試すことができます。

準備

以下の手順で必要なライブラリやリポジトリをセットアップします。

# 必要なパッケージをインストール
conda create -n mlx_bark python=3.12
conda activate mlx_bark
pip install transformers huggingface_hub hf_transfer

# リポジトリをクローン
git clone https://github.com/j-csc/mlx_bark
cd mlx_bark
pip install -r requirements.txt

# モデルをダウンロード(数GBあります)
export HF_HUB_ENABLE_HF_TRANSFER=1
huggingface-cli download --local-dir-use-symlinks False --local-dir weights/ mlx-community/mlx_bark

実行例

以下のコマンドでモデルを実行し、テキストからオーディオを生成します。

python model.py --text="Hello world!" --path weights/ --model large

MLXでのSunoのBarkモデルの移植方法(発展)

上記まではHuggingfaceのあたりに書いていることです。

ここからは、コマンドで実行する以外の方法を紹介します。

準備まで終わっていると、以下のようなフォルダ構成になっていると思います。

mlx_bark/
├── requirements.txt     # 必要なパッケージリスト
├── generation.wav       # 音声生成サンプル
├── convert.py           # 変換スクリプト
├── model.py             # モデル実行スクリプト
├── README.md            # 説明書
├── .gitignore           # Git無視設定
├── torch_codec.py       # PyTorch用コーデック
├── generation.mp4       # 生成されたビデオサンプル
└── weights/             # モデルの重みファイル格納ディレクトリ

新たにtest.ipynbのようなファイルをmlx_bark配下に作ってください。

そこに以下のコードを貼り付けてください。

from scipy.io.wavfile import write as write_wav
from torch_codec import codec_decode
from model import load_model, generate_text_semantic, generate_coarse, generate_fine, SAMPLE_RATE

# モデルのパスとタイプを指定
model_path = "weights/"
model_type = "large"
text_input = "こんにちは！あなたの名前を教えてください！"

# モデルをロード
tokenizer, bark_coarse, bark_fine, bark_text = load_model(model_path, model_type)

# テキストからセマンティックトークンを生成
semantic_tokens = generate_text_semantic(bark_text, tokenizer, text_input, use_kv_caching=True)

# セマンティックトークンからコーストークンを生成
coarse_tokens = generate_coarse(bark_coarse, x_semantic=semantic_tokens, use_kv_caching=True)

# コーストークンからファイントークンを生成
fine_tokens = generate_fine(bark_fine, coarse_tokens, temp=0.5)

# ファイントークンをデコードして音声を生成
audio_arr = codec_decode(fine_tokens)

# 音声をファイルに保存
output_path = "generation.wav"
write_wav(output_path, SAMPLE_RATE, audio_arr)

print("Generation complete! Audio saved to:", output_path)

これで音声ファイルが保存できると思います。

エラー回避策

上記のコードを実行すると、もしかするとバージョンの問題で以下のようなエラーが出るかもしれません。

ValueError                                Traceback (most recent call last)
Cell In[3], line 11
      8 fine_tokens = generate_fine(bark_fine, coarse_tokens, temp=0.5)
     10 # ファイントークンをデコードして音声を生成
---> 11 audio_arr = codec_decode(fine_tokens)
     13 # 音声をファイルに保存
     14 output_path = "generation.wav"

File ~/Desktop/mlx_bark/torch_codec.py:23, in codec_decode(fine_tokens)
     21 def codec_decode(fine_tokens):
     22     codec = _load_codec_model("cpu")
---> 23     arr = torch.from_numpy(np.array(fine_tokens, copy=False, dtype=np.int32))[None]
     24     arr = arr.to("cpu")
     25     arr = arr.transpose(0, 1)

ValueError: Unable to avoid copy while creating an array as requested.
If using `np.array(obj, copy=False)` replace it with `np.asarray(obj)` to allow a copy when needed (no behavior change in NumPy 1.x).
For more details, see https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword.

このようなエラーが出た場合は、torch_codec.pyの23行目を次のように変更してください。

# mlx_bark/torch_codec.py

# Loads to torch Encodec model
def codec_decode(fine_tokens):
    codec = _load_codec_model("cpu")
    arr = torch.from_numpy(np.asarray(fine_tokens, dtype=np.int32))[None] # 修正
    arr = arr.to("cpu")
    arr = arr.transpose(0, 1)
    emb = codec.quantizer.decode(arr)
    out = codec.decoder(emb)
    audio_arr = out.detach().cpu().numpy().squeeze()
    del arr, emb, out
    return audio_arr

これでもう一度読み直せば、エラーはなくなると思います。