Parler-TTS と OpenVINO によるテキスト読み上げ (TTS)#

import os 

os.environ["GIT_CLONE_PROTECTION_ACTIVE"] = "false" 

%pip install -q "openvino>=2024.2.0" 
%pip install -q git+https://github.com/huggingface/parler-tts.git "gradio>=4.19" transformers "torch>=2.2" --extra-index-url https://download.pytorch.org/whl/cpu

import torch 
from parler_tts import ParlerTTSForConditionalGeneration 
from transformers import AutoTokenizer 
import soundfile as sf 

device = "cpu" 

repo_id = "parler-tts/parler_tts_mini_v0.1" 
model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id).to(device) 
tokenizer = AutoTokenizer.from_pretrained(repo_id) 

prompt = "Hey, how are you doing today?" 
description = "A female speaker with a slightly low-pitched voice delivers her words quite expressively, in a very confined sounding environment with clear audio quality.She speaks very fast." 

input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device) 
prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) 

generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids) 
audio_arr = generation.cpu().numpy().squeeze() 
sf.write("parler_tts_out.wav", audio_arr, model.config.sampling_rate)

import IPython.display as ipd 

ipd.Audio("parler_tts_out.wav")

import openvino as ov 

def convert(model: torch.nn.Module, xml_path: str, example_input): 
    xml_path = Path(xml_path) 
    if not xml_path.exists(): 
        xml_path.parent.mkdir(parents=True, exist_ok=True) 
        with torch.no_grad(): 
            converted_model = ov.convert_model(model, example_input=example_input) 

        ov.save_model(converted_model, xml_path) 

        # メモリーをクリーンアップ 
        torch._C._jit_clear_class_registry() 
        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore() 
        torch.jit._state._clear_class_state()

from pathlib import Path 

TEXT_ENCODER_OV_PATH = Path("models/text_encoder_ir.xml") 

example_input = { 
    "input_ids": torch.ones((1, 39), dtype=torch.int64), 
} 

text_encoder_ov_model = convert(model.text_encoder, TEXT_ENCODER_OV_PATH, example_input)

DECODER_STAGE_1_OV_PATH = Path("models/decoder_stage_1_ir.xml") 

class DecoderStage1Wrapper(torch.nn.Module): 
    def __init__(self, decoder): 
        super().__init__() 
        self.decoder = decoder 

    def forward(self, input_ids=None, encoder_hidden_states=None, encoder_attention_mask=None, prompt_hidden_states=None): 
        return self.decoder( 
            input_ids=input_ids, 
            return_dict=False, 
            encoder_hidden_states=encoder_hidden_states, 
            encoder_attention_mask=encoder_attention_mask, 
            prompt_hidden_states=prompt_hidden_states, 
        ) 

example_input = { 
    "input_ids": torch.ones((9, 1), dtype=torch.int64), 
    "encoder_hidden_states": torch.ones((1, 39, 1024), dtype=torch.float32), 
    "encoder_attention_mask": torch.ones((1, 39), dtype=torch.int64), 
    "prompt_hidden_states": torch.ones((1, 9, 1024), dtype=torch.float32), 
} 

decoder_1_ov_model = convert(DecoderStage1Wrapper(model.decoder.model.decoder), DECODER_STAGE_1_OV_PATH, example_input)

DECODER_STAGE_2_OV_PATH = Path("models/decoder_stage_2_ir.xml") 

class DecoderStage2Wrapper(torch.nn.Module): 
    def __init__(self, decoder): 
        super().__init__() 
        self.decoder = decoder 

    def forward(self, input_ids=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_values=None): 
        past_key_values = tuple(tuple(past_key_values[i : i + 4]) for i in range(0, len(past_key_values), 4)) 
        return self.decoder( 
            input_ids=input_ids, 
            return_dict=False, 
            encoder_hidden_states=encoder_hidden_states, 
            encoder_attention_mask=encoder_attention_mask, 
            past_key_values=past_key_values, 
        ) 

example_input = { 
    "input_ids": torch.ones((9, 1), dtype=torch.int64), 
    "encoder_hidden_states": torch.ones((1, 39, 1024), dtype=torch.float32), 
    "encoder_attention_mask": torch.ones((1, 39), dtype=torch.int64), 
    "past_key_values": ( 
        ( 
            torch.ones(1, 16, 10, 64, dtype=torch.float32), 
            torch.ones(1, 16, 10, 64, dtype=torch.float32), 
            torch.ones(1, 16, 39, 64, dtype=torch.float32), 
            torch.ones(1, 16, 39, 64, dtype=torch.float32), 
        ) 
        * 24 
    ), 
} 

decoder_2_ov_model = convert(DecoderStage2Wrapper(model.decoder.model.decoder), DECODER_STAGE_2_OV_PATH, example_input)

import ipywidgets as widgets 

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

device

from collections import namedtuple 

EncoderOutput = namedtuple("EncoderOutput", "last_hidden_state") 
DecoderOutput = namedtuple("DecoderOutput", ("last_hidden_state", "past_key_values", "hidden_states", "attentions", "cross_attentions")) 

class TextEncoderModelWrapper(torch.nn.Module): 
    def __init__(self, encoder_ir_path, config): 
        self.encoder = core.compile_model(encoder_ir_path, 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 EncoderOutput(torch.from_numpy(last_hidden_state)) 

class DecoderWrapper(torch.nn.Module): 
    def __init__(self, decoder_stage_1_ir_path, decoder_stage_2_ir_path): 
        super().__init__() 
        self.decoder_stage_1 = core.compile_model(decoder_stage_1_ir_path, device.value) 
        self.decoder_stage_2 = core.compile_model(decoder_stage_2_ir_path, device.value) 

    def __call__(self, input_ids=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_values=None, prompt_hidden_states=None, **kwargs): 
    inputs = {} 
    if input_ids is not None: 
        inputs["input_ids"] = input_ids 
    if encoder_hidden_states is not None: 
        inputs["encoder_hidden_states"] = encoder_hidden_states 
    if encoder_attention_mask is not None: 
        inputs["encoder_attention_mask"] = encoder_attention_mask 
    if prompt_hidden_states is not None: 
        inputs["prompt_hidden_states"] = prompt_hidden_states 
    if past_key_values is not None: 
        past_key_values = tuple(past_key_value for pkv_per_layer in past_key_values for past_key_value in pkv_per_layer) 
        inputs["past_key_values"] = past_key_values 
        arguments = ( 
            input_ids, 
            encoder_hidden_states, 
            encoder_attention_mask, 
            *past_key_values, 
        ) 
        outs = self.decoder_stage_2(arguments) 
    else: 
        outs = self.decoder_stage_1(inputs) 

    outs = [torch.from_numpy(out) for out in outs.values()] 
    past_key_values = list(list(outs[i : i + 4]) for i in range(1, len(outs), 4)) 

    return DecoderOutput(outs[0], past_key_values, None, None, None)

model.text_encoder = TextEncoderModelWrapper(TEXT_ENCODER_OV_PATH, 
model.text_encoder.config) 
model.decoder.model.decoder = DecoderWrapper(DECODER_STAGE_1_OV_PATH, DECODER_STAGE_2_OV_PATH)

generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids) 
audio_arr = generation.cpu().numpy().squeeze() 
sf.write("parler_tts_out.wav", audio_arr, model.config.sampling_rate)

import IPython.display as ipd 

ipd.Audio("parler_tts_out.wav")

import gradio as gr 
import numpy as np 
from transformers import AutoFeatureExtractor, set_seed 

title = "Text-to-speech (TTS) with Parler-TTS and OpenVINO" 
feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id) 
SAMPLE_RATE = feature_extractor.sampling_rate 

def infer(prompt, description, seed): 
    set_seed(seed) 

    input_ids = tokenizer(description, return_tensors="pt").input_ids.to("cpu") 
    prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cpu") 

    generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids) 
    audio_arr = generation.cpu().numpy().squeeze() 
    sr = SAMPLE_RATE 

    return sr, audio_arr 

demo = gr.Interface( 
    infer, 
    [ 
        gr.Text(label="Prompt"), 
        gr.Text(label="Description"), 
        gr.Slider( 
            label="Seed", 
            value=42, 
            step=1, 
            minimum=0, 
            maximum=np.iinfo(np.int32).max, 
        ), 
    ], 
    gr.Audio(label="Output Audio", type="numpy"), 
    title=title, 
    description=description, 
    examples=[ 
        [ 
            "Hey, how are you doing today?", 
            "A female speaker with a slightly low-pitched voice delivers her words quite expressively, in a very confined sounding environment with clear audio quality.She speaks very fast.", 
        ], 
        [ 
            "'This is the best time of my life, Bartley,' she said happily.", 
            "A female speaker with a slightly low-pitched, quite monotone voice delivers her words at a slightly faster-than-average pace in a confined space with very clear audio.", 
        ], 
        [ 
            "Montrose also, after having experienced still more variety of good and bad fortune, threw down his arms, and retired out of the kingdom.", 
            "A male speaker with a slightly high-pitched voice delivering his words at a slightly slow pace in a small, confined space with a touch of background noise and a quite monotone tone.", 
        ], 
        [ 
            "montrose also after having experienced still more variety of good and bad fortune threw down his arms and retired out of the kingdom", 
            "A male speaker with a low-pitched voice delivering his words at a fast pace in a small, confined space with a lot of background noise and an animated tone.", 
        ], 
    ], 
) 

try: 
    demo.queue().launch(debug=False) 
except Exception: 
    demo.queue().launch(share=True, debug=False) 
# リモートで起動する場合は、server_name と server_port を指定 
# demo.launch(server_name='your server name', server_port='server port in int') 
# 詳細については、ドキュメントをご覧ください: https://gradio.app/docs/