MusicGen と OpenVINO による制御可能な音楽生成#

%pip install -q "openvino>=2023.3.0" 
%pip install -q "torch>=2.1" "gradio>=4.19" "transformers" packaging --extra-index-url https://download.pytorch.org/whl/cpu

from collections import namedtuple 
from functools import partial 
import gc 
from pathlib import Path 
from typing import Optional, Tuple 
import warnings 

from IPython.display import Audio 
import openvino as ov 
import numpy as np 
import torch 
from torch.jit import TracerWarning 
from transformers import AutoProcessor, MusicgenForConditionalGeneration 
from transformers.modeling_outputs import ( 
    BaseModelOutputWithPastAndCrossAttentions, 
    CausalLMOutputWithCrossAttentions, 
) 

# トレース警告を無視 
warnings.filterwarnings("ignore", category=TracerWarning)

processor = AutoProcessor.from_pretrained("facebook/musicgen-small") 

inputs = processor( 
    text=["80s pop track with bassy drums and synth"], 
    return_tensors="pt", 
) 

audio_values = model.generate(**inputs, do_sample=True, guidance_scale=3, max_new_tokens=n_tokens) 

Audio(audio_values[0].cpu().numpy(), rate=sampling_rate)

models_dir = Path("./models") 
t5_ir_path = models_dir / "t5.xml" 
musicgen_0_ir_path = models_dir / "mg_0.xml" 
musicgen_ir_path = models_dir / "mg.xml" 
audio_decoder_ir_path = models_dir / "encodec.xml"

if not t5_ir_path.exists(): 
    t5_ov = ov.convert_model(model.text_encoder, example_input={"input_ids": inputs["input_ids"]}) 

    ov.save_model(t5_ov, t5_ir_path) 
    del t5_ov 
    gc.collect()

# モデル設定 `torchscript` を True に設定すると、モデルは出力としてタプルを返します。 
model.decoder.config.torchscript = True 

if not musicgen_0_ir_path.exists(): 
    decoder_input = { 
        "input_ids": torch.ones(8, 1, dtype=torch.int64), 
        "encoder_hidden_states": torch.ones(2, 12, 1024, dtype=torch.float32), 
        "encoder_attention_mask": torch.ones(2, 12, dtype=torch.int64), 
    } 
    mg_ov_0_step = ov.convert_model(model.decoder, example_input=decoder_input) 

    ov.save_model(mg_ov_0_step, musicgen_0_ir_path) 
    del mg_ov_0_step 
    gc.collect()

# example_input 辞書に引数を追加 
if not musicgen_ir_path.exists():
    # 変換されたモデル署名に `past_key_values` を追加 
    decoder_input["past_key_values"] = tuple( 
        [ 
            ( 
                torch.ones(2, 16, 1, 64, dtype=torch.float32), 
                torch.ones(2, 16, 1, 64, dtype=torch.float32), 
                torch.ones(2, 16, 12, 64, dtype=torch.float32), 
                torch.ones(2, 16, 12, 64, dtype=torch.float32), 
            ) 
        ] 
        * 24
     ) 

    mg_ov = ov.convert_model(model.decoder, example_input=decoder_input) 
    for input in mg_ov.inputs[3:]: 
        input.get_node().set_partial_shape(ov.PartialShape([-1, 16, -1, 64])) 
        input.get_node().set_element_type(ov.Type.f32) 

    mg_ov.validate_nodes_and_infer_types() 

    ov.save_model(mg_ov, musicgen_ir_path) 
    del mg_ov 
    gc.collect()

if not audio_decoder_ir_path.exists(): 

    class AudioDecoder(torch.nn.Module): 
        def __init__(self, model): 
            super().__init__() 
            self.model = model 

        def forward(self, output_ids): 
            return self.model.decode(output_ids, [None]) 

    audio_decoder_input = {"output_ids": torch.ones((1, 1, 4, n_tokens - 3), dtype=torch.int64)} 

    with torch.no_grad(): 
        audio_decoder_ov = ov.convert_model(AudioDecoder(model.audio_encoder), example_input=audio_decoder_input) 
    ov.save_model(audio_decoder_ov, audio_decoder_ir_path) 
    del audio_decoder_ov 
    gc.collect()

class TextEncoderWrapper(torch.nn.Module): 
    def __init__(self, encoder_ir, config): 
        super().__init__() 
        self.encoder = core.compile_model(encoder_ir, device.value) 
        self.config = config 

    def forward(self, input_ids, **kwargs): 
        last_hidden_state = self.encoder(input_ids)[self.encoder.outputs[0]] 
        last_hidden_state = torch.tensor(last_hidden_state) 
        return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=last_hidden_state) 

class MusicGenWrapper(torch.nn.Module): 
    def __init__( 
        self, 
        music_gen_lm_0_ir, 
        music_gen_lm_ir, 
        config, 
        num_codebooks, 
        build_delay_pattern_mask, 
        apply_delay_pattern_mask, 
    ): 
        super().__init__() 
        self.music_gen_lm_0 = core.compile_model(music_gen_lm_0_ir, device.value) 
        self.music_gen_lm = core.compile_model(music_gen_lm_ir, device.value) 
        self.config = config self.num_codebooks = num_codebooks 
        self.build_delay_pattern_mask = build_delay_pattern_mask 
        self.apply_delay_pattern_mask = apply_delay_pattern_mask 

    def forward( 
        self, 
        input_ids: torch.LongTensor = None, 
        encoder_hidden_states: torch.FloatTensor = None, 
        encoder_attention_mask: torch.LongTensor = None, 
        past_key_values: Optional[Tuple[torch.FloatTensor]] = None, 
        **kwargs 
    ): 
        if past_key_values is None: 
            model = self.music_gen_lm_0 
            arguments = (input_ids, encoder_hidden_states, encoder_attention_mask) 
        else: 
            model = self.music_gen_lm 
            arguments = ( 
                input_ids, 
                encoder_hidden_states, 
                encoder_attention_mask, 
                *past_key_values, 
            ) 

        output = model(arguments) 
        return CausalLMOutputWithCrossAttentions( 
            logits=torch.tensor(output[model.outputs[0]]), 
            past_key_values=tuple([output[model.outputs[i]] for i in range(1, 97)]), 
        ) 

class AudioDecoderWrapper(torch.nn.Module): 
    def __init__(self, decoder_ir, config): 
        super().__init__() 
        self.decoder = core.compile_model(decoder_ir, device.value) 
        self.config = config 
        self.output_type = namedtuple("AudioDecoderOutput", ["audio_values"]) 

    def decode(self, output_ids, audio_scales): 
        output = self.decoder(output_ids)[self.decoder.outputs[0]] 
        return self.output_type(audio_values=torch.tensor(output))

text_encode_ov = TextEncoderWrapper(t5_ir_path, model.text_encoder.config) 
musicgen_decoder_ov = MusicGenWrapper( 
    musicgen_0_ir_path, 
    musicgen_ir_path, 
    model.decoder.config, 
    model.decoder.num_codebooks, 
    model.decoder.build_delay_pattern_mask, 
    model.decoder.apply_delay_pattern_mask, 
) 
audio_encoder_ov = AudioDecoderWrapper(audio_decoder_ir_path, model.audio_encoder.config) 

del model.text_encoder 
del model.decoder 
del model.audio_encoder 
gc.collect() 

model.text_encoder = text_encode_ov 
model.decoder = musicgen_decoder_ov 
model.audio_encoder = audio_encoder_ov 

def prepare_inputs_for_generation( 
    self, 
    decoder_input_ids, 
    past_key_values=None, 
    attention_mask=None, 
    head_mask=None, 
    decoder_attention_mask=None, 
    decoder_head_mask=None, 
    cross_attn_head_mask=None, 
    use_cache=None, 
    encoder_outputs=None, decoder_delay_pattern_mask=None, 
    guidance_scale=None, 
    **kwargs, 
): 
    if decoder_delay_pattern_mask is None:         ( 
            decoder_input_ids, 
            decoder_delay_pattern_mask, 
        ) = self.decoder.build_delay_pattern_mask( 
            decoder_input_ids, 
            self.generation_config.pad_token_id, 
            max_length=self.generation_config.max_length, 
        ) 

    # 遅延パターンマスクを適用 
    decoder_input_ids = self.decoder.apply_delay_pattern_mask(decoder_input_ids, decoder_delay_pattern_mask) 

    if guidance_scale is not None and guidance_scale > 1:
        # 分類器フリーのガイダンスのため、デコーダー引数をバッチ dim 全体に複製する必要があります 
        # (サンプリング前にこれらを分割します) 
        decoder_input_ids = decoder_input_ids.repeat((2, 1)) 
    if decoder_attention_mask is not None: 
        decoder_attention_mask = decoder_attention_mask.repeat((2, 1)) 

    if past_key_values is not None:
        # past が使用されている場合は decoder_input_ids をカット 
        decoder_input_ids = decoder_input_ids[:, -1:] 

        return { "input_ids": None, # エンコーダー出力が定義されています。input_ids は必要ありません 
            "encoder_outputs": encoder_outputs, 
            "past_key_values": past_key_values, 
            "decoder_input_ids": decoder_input_ids, 
            "attention_mask": attention_mask, 
            "decoder_attention_mask": decoder_attention_mask, 
            "head_mask": head_mask, 
            "decoder_head_mask": decoder_head_mask, 
            "cross_attn_head_mask": cross_attn_head_mask, 
            "use_cache": use_cache, 
    } 

model.prepare_inputs_for_generation = partial(prepare_inputs_for_generation, model)

processor = AutoProcessor.from_pretrained("facebook/musicgen-small") 

inputs = processor( 
    text=["80s pop track with bassy drums and synth"], 
    return_tensors="pt", 
) 

audio_values = model.generate(**inputs, do_sample=True, guidance_scale=3, max_new_tokens=n_tokens) 

Audio(audio_values[0].cpu().numpy(), rate=sampling_rate)