AudioLDM2 と OpenVINO™ によるサウンド生成

この Jupyter ノートブックは、ローカルへのインストール後にのみ起動できます。

GitHub

AudioLDM 2 は、任意のテキスト入力からリアルなオーディオサンプルを生成できる潜在的なテキストからオーディオへの拡散モデルです。

AudioLDM 2 は、Haohe Liu 氏らによる論文 AudioLDM 2: Learning Holistic Audio Generation with Self-supervised Pretraining で提案されました。

モデルはテキストプロンプトを入力として受け取り、対応するオーディオを予測します。テキスト条件付きサウンド効果、人間の音声、音楽を生成できます。

image0

このチュートリアルでは、パイプラインを試し、それをサポートするモデルを 1 つずつ変換し、Gradio でインタラクティブなアプリを実行します。

目次

必要条件

%pip uninstall -q -y "openvino-dev" "openvino" "openvino-nightly"
%pip install -q accelerate "diffusers>=0.21.0" transformers torch gradio --extra-index-url https://download.pytorch.org/whl/cpu
%pip install -q "openvino-nightly"
Note: you may need to restart the kernel to use updated packages.
Note: you may need to restart the kernel to use updated packages.

生成パイプラインのインスタンス化

サリー大学視覚・音声・信号処理センターAudioLDM 2 を操作するには、Hugging Face Diffusers パッケージを使用します。Diffusers パッケージは AudioLDM2Pipeline クラスを公開し、モデルのインスタンス化と重みの読み込みを簡素化します。以下のコードは、AudioLDM2Pipeline を作成し、テキスト条件付きのサウンドサンプルを生成する方法を示しています。

from collections import namedtuple
from functools import partial
import gc
from pathlib import Path

from diffusers import AudioLDM2Pipeline
from IPython.display import Audio
import numpy as np
import openvino as ov
import torch

MODEL_ID = "cvssp/audioldm2"
pipe = AudioLDM2Pipeline.from_pretrained(MODEL_ID)

prompt = "birds singing in the forest"
negative_prompt = "Low quality"
audio = pipe(
    prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=150,
    audio_length_in_s=3.0
).audios[0]

sampling_rate = 16000
Audio(audio, rate=sampling_rate)
Loading pipeline components...:   0%|          | 0/11 [00:00<?, ?it/s]
0%|          | 0/100 [00:00<?, ?it/s]

モデルを OpenVINO 中間表現 (IR) 形式に変換

モデル変換 API を使用すると、パイプラインをバックアップする PyTorch モデルを直接変換できます。OpenVINO ov.Model オブジェクト・インスタンスを取得するには、モデル・オブジェクト、モデルトレース用の入力データを ov.convert_model 関数に提供する必要があります。ov.save_model 関数を使用して、次回のデプロイのためにモデルをディスクに保存できます。

パイプラインは 7 つの重要なパーツで構成されます。

  • テキストプロンプトからサウンドを生成する条件を作成するための T5 および CLAP テキスト・エンコーダー。

  • 2 つのテキスト・エンコーダーからの出力をマージする投影モデル。

  • 2 つのテキスト・エンコーダーからの投影された出力に基づいて、一連の隠し状態を生成する GPT-2 言語モデル。

  • メル・スペクトログラム潜在変数を最終的なオーディオ波形に変換する Vocoder。

  • 段階的にノイズを除去する潜像表現のための Unet。

  • 潜在空間を画像にデコードするオート・エンコーダー (VAE)。

models_base_folder = Path("models")

def cleanup_torchscript_cache():
    """
    Helper for removing cached model representation
    """
    torch._C._jit_clear_class_registry()
    torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()
    torch.jit._state._clear_class_state()

CLAP テキスト・エンコーダー変換

最初のフリーズされたテキスト・エンコーダー。AudioLDM2 は、音声とテキストの結合埋め込みモデル CLAP、具体的には laion/clap-htsat-unfused バリアントを使用します。テキストブランチは、テキストプロンプトをプロンプト埋め込みにエンコードするために使用されます。完全な音声テキストモデルは、類似性スコアを計算して、生成された波形をテキストプロンプトに対してランク付けするのに使用されます。

class ClapEncoderWrapper(torch.nn.Module):
    def __init__(self, encoder):
        super().__init__()
        encoder.eval()
        self.encoder = encoder

    def forward(self, input_ids, attention_mask):
        return self.encoder.get_text_features(input_ids, attention_mask)

clap_text_encoder_ir_path = models_base_folder / "clap_text_encoder.xml"

if not clap_text_encoder_ir_path.exists():
    with torch.no_grad():
        ov_model = ov.convert_model(
            ClapEncoderWrapper(pipe.text_encoder),  # model instance
            example_input={
                "input_ids": torch.ones((1, 512), dtype=torch.long),
                "attention_mask": torch.ones((1, 512), dtype=torch.long),
            },  # inputs for model tracing
        )
    ov.save_model(ov_model, clap_text_encoder_ir_path)
    del ov_model
    cleanup_torchscript_cache()
    gc.collect()
    print("Text Encoder successfully converted to IR")
else:
    print(f"Text Encoder will be loaded from {clap_text_encoder_ir_path}")
Text Encoder will be loaded from clap_text_encoder.xml

T5 テキスト・エンコーダー変換

2 番目の固定テキスト・エンコーダーとして、AudioLDM2 は T5、具体的には google/flan-t5-large バリアントを使用します。

テキスト・エンコーダーは、入力プロンプト (例えば、“森で歌う鳥”) を U-Net が理解できる埋め込み空間に変換する役割を担います。これは通常、入力トークンのシーケンスを潜在テキスト埋め込みのシーケンスにマッピングする単純なトランスフォーマー・ベースのエンコーダーです。

テキスト・エンコーダーの入力はテンソル input_ids です。これには、トークナイザーによって処理され、モデルによって受け入れられる最大長までパディングされたテキストからのトークン・インデックスが含まれます。モデルの出力は 2 つのテンソルです。

  • last_hidden_state - モデル内の最後の MultiHeadtention レイヤーからの非表示状態。
  • pooler_out - モデル全体の非表示状態のプールされた出力。
t5_text_encoder_ir_path = models_base_folder / "t5_text_encoder.xml"

if not t5_text_encoder_ir_path.exists():
    pipe.text_encoder_2.eval()
    with torch.no_grad():
        ov_model = ov.convert_model(
            pipe.text_encoder_2,  # model instance
            example_input=torch.ones((1, 7), dtype=torch.long),  # inputs for model tracing
        )
    ov.save_model(ov_model, t5_text_encoder_ir_path)
    del ov_model
    cleanup_torchscript_cache()
    gc.collect()
    print("Text Encoder successfully converted to IR")
else:
    print(f"Text Encoder will be loaded from {t5_text_encoder_ir_path}")
Text Encoder will be loaded from t5_text_encoder.xml

投影モデルの変換

最初のテキスト・エンコーダー・モデルと 2 番目のテキスト・エンコーダー・モデルからの隠し状態を線形投影し、学習したシーケンスの開始トークンとシーケンスの終了トークンの埋め込みを挿入するために使用されるトレーニング済みモデル。2 つのテキスト・エンコーダーから投影された隠し状態が連結され、言語モデルに入力が与えられます。

projection_model_ir_path = models_base_folder / "projection_model.xml"

projection_model_inputs = {
    "hidden_states": torch.randn((1, 1, 512), dtype=torch.float32),
    "hidden_states_1": torch.randn((1, 7, 1024), dtype=torch.float32),
    "attention_mask": torch.ones((1, 1), dtype=torch.int64),
    "attention_mask_1": torch.ones((1, 7), dtype=torch.int64),
}

if not projection_model_ir_path.exists():
    pipe.projection_model.eval()
    with torch.no_grad():
        ov_model = ov.convert_model(
            pipe.projection_model,  # model instance
            example_input=projection_model_inputs,  # inputs for model tracing
        )
    ov.save_model(ov_model, projection_model_ir_path)
    del ov_model
    cleanup_torchscript_cache()
    gc.collect()
    print("The Projection Model successfully converted to IR")
else:
    print(f"The Projection Model will be loaded from {projection_model_ir_path}")
The Projection Model will be loaded from projection_model.xml

GPT-2 変換

GPT-2 は、2 つのテキスト・エンコーダーからの投影された出力に基づいて、一連の隠し状態を生成するのに使用される自己回帰言語モデルです。

language_model_ir_path = models_base_folder / "language_model.xml"

language_model_inputs = {
    "inputs_embeds": torch.randn((1, 12, 768), dtype=torch.float32),
    "attention_mask": torch.ones((1, 12), dtype=torch.int64),
}

if not language_model_ir_path.exists():
    pipe.language_model.config.torchscript = True
    pipe.language_model.eval()
    pipe.language_model.__call__ = partial(pipe.language_model.__call__, kwargs={
        "past_key_values": None,
        "use_cache": False,
        "return_dict": False})
    with torch.no_grad():
        ov_model = ov.convert_model(
            pipe.language_model,  # model instance
            example_input=language_model_inputs,  # inputs for model tracing
        )

    ov_model.inputs[0].get_node().set_partial_shape(ov.PartialShape([1, -1]))
    ov_model.inputs[0].get_node().set_element_type(ov.Type.i64)
    ov_model.inputs[1].get_node().set_partial_shape(ov.PartialShape([1, -1, 768]))
    ov_model.inputs[1].get_node().set_element_type(ov.Type.f32)

    ov_model.validate_nodes_and_infer_types()

    ov.save_model(ov_model, language_model_ir_path)
    del ov_model
    cleanup_torchscript_cache()
    gc.collect()
    print("The Projection Model successfully converted to IR")
else:
    print(f"The Projection Model will be loaded from {language_model_ir_path}")
The Projection Model will be loaded from language_model.xml

Vocoder 変換

SpeechT5 HiFi-GAN Vocoder は、メル・スペクトログラム潜在変数を最終的なオーディオ波形に変換するのに使用されます。

vocoder_ir_path = models_base_folder / "vocoder.xml"

if not vocoder_ir_path.exists():
    pipe.vocoder.eval()
    with torch.no_grad():
        ov_model = ov.convert_model(
            pipe.vocoder,  # model instance
            example_input=torch.ones((1, 700, 64), dtype=torch.float32),  # inputs for model tracing
        )
    ov.save_model(ov_model, vocoder_ir_path)
    del ov_model
    cleanup_torchscript_cache()
    gc.collect()
    print("The Vocoder successfully converted to IR")
else:
    print(f"The Vocoder will be loaded from {vocoder_ir_path}")
The Vocoder will be loaded from vocoder.xml

UNet 変換

UNet モデルは、エンコードされたオーディオ潜在音のノイズ除去に使用されます。UNet モデル変換のプロセスは、元の Stable Diffusion モデルと同様です。

unet_ir_path = models_base_folder / "unet.xml"

pipe.unet.eval()
unet_inputs = {
    "sample": torch.randn((2, 8, 75, 16), dtype=torch.float32),
    "timestep": torch.tensor(1, dtype=torch.int64),
    "encoder_hidden_states": torch.randn((2, 8, 768), dtype=torch.float32),
    "encoder_hidden_states_1": torch.randn((2, 7, 1024), dtype=torch.float32),
    "encoder_attention_mask_1": torch.ones((2, 7), dtype=torch.int64),
}

if not unet_ir_path.exists():
    with torch.no_grad():
        ov_model = ov.convert_model(pipe.unet, example_input=unet_inputs)

    ov_model.inputs[0].get_node().set_partial_shape(ov.PartialShape((2, 8, -1, 16)))
    ov_model.inputs[2].get_node().set_partial_shape(ov.PartialShape((2, 8, 768)))
    ov_model.inputs[3].get_node().set_partial_shape(ov.PartialShape((2, -1, 1024)))
    ov_model.inputs[4].get_node().set_partial_shape(ov.PartialShape((2, -1)))
    ov_model.validate_nodes_and_infer_types()

    ov.save_model(ov_model, unet_ir_path)

    del ov_model
    cleanup_torchscript_cache()
    gc.collect()
    print("Unet successfully converted to IR")
else:
    print(f"Unet will be loaded from {unet_ir_path}")
Unet will be loaded from unet.xml

VAE デコーダー変換

VAE モデルには、エンコーダーとデコーダーの 2 つのパーツがあります。エンコーダーは、画像を低次元の潜在表現に変換するのに使用され、これが U-Net モデルの入力となります。逆に、デコーダーは潜在表現を変換して画像に戻します。

潜在拡散トレーニング中、エンコーダーは、順拡散プロセス用の画像の潜在表現 (潜在) を取得するために使用され、各ステップでより多くのノイズが適用されます。論中、逆拡散プロセスによって生成されたノイズ除去された潜在は、VAE デコーダーによって画像に変換されます。推論中に、VAE デコーダーのみが必要であることが分かります。エンコーダー部分を変換する方法については、安定拡散のノートブックに記載されています。

vae_ir_path = models_base_folder / "vae.xml"

class VAEDecoderWrapper(torch.nn.Module):
    def __init__(self, vae):
        super().__init__()
        vae.eval()
        self.vae = vae

    def forward(self, latents):
        return self.vae.decode(latents)

if not vae_ir_path.exists():
    vae_decoder = VAEDecoderWrapper(pipe.vae)
    latents = torch.zeros((1, 8, 175, 16))

    vae_decoder.eval()
    with torch.no_grad():
        ov_model = ov.convert_model(vae_decoder, example_input=latents)
        ov.save_model(ov_model, vae_ir_path)
    del ov_model
    cleanup_torchscript_cache()
    gc.collect()
    print("VAE decoder successfully converted to IR")
else:
    print(f"VAE decoder will be loaded from {vae_ir_path}")
VAE decoder will be loaded from vae.xml

Stable Diffusion パイプライン用の推論デバイスの選択

OpenVINO を使用して推論を実行するためにドロップダウン・リストからデバイスを選択します

import ipywidgets as widgets

core = ov.Core()

DEVICE = widgets.Dropdown(
    options=core.available_devices + ["AUTO"],
    value="CPU",
    description="Device:",
    disabled=False,
)

DEVICE
Dropdown(description='Device:', options=('CPU', 'AUTO'), value='CPU')

OpenVINO モデルを元のパイプラインに適合

ここでは、元の推論パイプラインに埋め込む 3 つの OpenVINO モデルすべてに対してラッパークラスを作成します。OV モデルを適応させる際に考慮すべき事項をいくつか示します。

  • 元のパイプラインから渡されたパラメーターがコンパイルされた OV モデルに適切に転送されることを確認します。OV モデルでは入力引数の一部のみが使用され、一部は無視される場合があり、引数を別のデータタイプに変換したり、タプルや辞書などの一部のデータ構造をアンラップしたりする必要がある場合があります。
  • ラッパー クラスが期待どおりの形式でパイプラインに結果を返すことを保証します。以下の例では、OV モデルの出力を特別な名前付きタプルにパックしてパイプラインに適合させる方法を確認できます。
  • モデルを呼び出すため元のパイプラインで使用されるモデルメソッドに注意してください。これは forward メソッドではない可能性があります。OV モデル推論を get_text_features メソッドにラップする方法は、OVClapEncoderWrapper を参照してください。
class OVClapEncoderWrapper:
    def __init__(self, encoder_ir, config):
        self.encoder = core.compile_model(encoder_ir, DEVICE.value)
        self.config = config

    def get_text_features(self, input_ids, attention_mask, **_):
        last_hidden_state = self.encoder([input_ids, attention_mask])[0]
        return torch.from_numpy(last_hidden_state)

class OVT5EncoderWrapper:
    def __init__(self, encoder_ir, config):
        self.encoder = core.compile_model(encoder_ir, DEVICE.value)
        self.config = config
        self.dtype = self.config.torch_dtype

    def __call__(self, input_ids, **_):
        last_hidden_state = self.encoder(input_ids)[0]
        return torch.from_numpy(last_hidden_state)[None, ...]

class OVVocoderWrapper:
    def __init__(self, vocoder_ir, config):
        self.vocoder = core.compile_model(vocoder_ir, DEVICE.value)
        self.config = config

    def __call__(self, mel_spectrogram, **_):
        waveform = self.vocoder(mel_spectrogram)[0]
        return torch.from_numpy(waveform)

class OVProjectionModelWrapper:
    def __init__(self, proj_model_ir, config):
        self.proj_model = core.compile_model(proj_model_ir, DEVICE.value)
        self.config = config
        self.output_type = namedtuple("ProjectionOutput", ["hidden_states", "attention_mask"])

    def __call__(
        self, hidden_states,
        hidden_states_1,
        attention_mask,
        attention_mask_1, **_
    ):
        output = self.proj_model({
            "hidden_states": hidden_states,
            "hidden_states_1": hidden_states_1,
            "attention_mask": attention_mask,
            "attention_mask_1": attention_mask_1,
        })
        return self.output_type(torch.from_numpy(output[0]), torch.from_numpy(output[1]))

class OVUnetWrapper:
    def __init__(self, unet_ir, config):
        self.unet = core.compile_model(unet_ir, DEVICE.value)
        self.config = config

    def __call__(
        self, sample,
        timestep,
        encoder_hidden_states,
        encoder_hidden_states_1,
        encoder_attention_mask_1, **_
    ):
        output = self.unet({
            "sample": sample,
            "timestep": timestep,
            "encoder_hidden_states": encoder_hidden_states,
            "encoder_hidden_states_1": encoder_hidden_states_1,
            "encoder_attention_mask_1": encoder_attention_mask_1,
        })
        return (torch.from_numpy(output[0]), )

class OVVaeDecoderWrapper:
    def __init__(self, vae_ir, config):
        self.vae = core.compile_model(vae_ir, DEVICE.value)
        self.config = config
        self.output_type = namedtuple("VaeOutput", ["sample"])

    def decode(self, latents, **_):
        last_hidden_state = self.vae(latents)[0]
        return self.output_type(torch.from_numpy(last_hidden_state))

def generate_language_model(
    gpt_2: ov.CompiledModel,
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor,
    max_new_tokens: int = 8,
    **_
) -> torch.Tensor:
    """
    Generates a sequence of hidden-states from the language model, conditioned on the embedding inputs.
    """
    if not max_new_tokens:
        max_new_tokens = 8
    inputs_embeds = inputs_embeds.cpu().numpy()
    attention_mask = attention_mask.cpu().numpy()
    for _ in range(max_new_tokens):
        # forward pass to get next hidden states
        output = gpt_2({"inputs_embeds":inputs_embeds, "attention_mask":attention_mask})

        next_hidden_states = output[0]

        # Update the model input
        inputs_embeds = np.concatenate([inputs_embeds, next_hidden_states[:, -1:, :]], axis=1)
        attention_mask = np.concatenate([attention_mask, np.ones((attention_mask.shape[0], 1))], axis=1)
    return torch.from_numpy(inputs_embeds[:, -max_new_tokens:, :])

ラッパー・オブジェクトを初期化し、HF パイプラインにロードします。

pipe = AudioLDM2Pipeline.from_pretrained(MODEL_ID)
pipe.config.torchscript = True
pipe.config.return_dict = False

np.random.seed(0)
torch.manual_seed(0)

pipe.text_encoder = OVClapEncoderWrapper(clap_text_encoder_ir_path, pipe.text_encoder.config)
pipe.text_encoder_2 = OVT5EncoderWrapper(t5_text_encoder_ir_path, pipe.text_encoder_2.config)
pipe.projection_model = OVProjectionModelWrapper(projection_model_ir_path, pipe.projection_model.config)
pipe.vocoder = OVVocoderWrapper(vocoder_ir_path, pipe.vocoder.config)
pipe.unet = OVUnetWrapper(unet_ir_path, pipe.unet.config)
pipe.vae = OVVaeDecoderWrapper(vae_ir_path, pipe.vae.config)

pipe.generate_language_model = partial(generate_language_model, core.compile_model(language_model_ir_path, DEVICE.value))

gc.collect()

prompt = "birds singing in the forest"
negative_prompt = "Low quality"
audio = pipe(
    prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=150,
    audio_length_in_s=3.0
).audios[0]

sampling_rate = 16000
Audio(audio, rate=sampling_rate)
Loading pipeline components...:   0%|          | 0/11 [00:00<?, ?it/s]
0%|          | 0/200 [00:00<?, ?it/s]

変換されたパイプラインを試す

これで、生成の準備が整いました。生成プロセスを改善するために、否定プロンプトを提供する可能性も導入します。技術的には、ポジティブなプロンプトはそれに関連付けられた出力に向かって拡散を誘導し、ネガティブなプロンプトは拡散をそこから遠ざけるように誘導します。以下のデモアプリは Gradio パッケージを使用して作成されています

import gradio as gr

def _generate(prompt, negative_prompt, audio_length_in_s,
              num_inference_steps, _=gr.Progress(track_tqdm=True)):
    """Gradio backing function."""
    audio_values = pipe(
        prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        audio_length_in_s=audio_length_in_s
    )
    waveform = audio_values[0].squeeze() * 2**15
    return (sampling_rate, waveform.astype(np.int16))

demo = gr.Interface(
    _generate,
    inputs=[
        gr.Textbox(label="Text Prompt"),
        gr.Textbox(label="Negative Prompt", placeholder="Example: Low quality"),
        gr.Slider(
            minimum=1.0,
            maximum=15.0,
            step=0.25,
            value=7,
            label="Audio Length (s)",
        ),
        gr.Slider(label="Inference Steps", step=5, value=150, minimum=50, maximum=250)
    ],
    outputs=[
        "audio"
    ],
    examples=[
        ["birds singing in the forest", "Low quality", 7, 150],
        ["The sound of a hammer hitting a wooden surface", "", 4, 200],
    ],
)
try:
    demo.queue().launch(debug=False)
except Exception:
    demo.queue().launch(share=True, debug=False)

# If you are launching remotely, specify server_name and server_port
# EXAMPLE: `demo.launch(server_name="your server name", server_port="server port in int")`
# To learn more please refer to the Gradio docs: https://gradio.app/docs/