EnCodec と OpenVINO によるオーディオ圧縮#

%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu "openvino>=2023.3.0" "torch>=2.1" "torchaudio>=2.1" "encodec>=0.1.1" "gradio>=4.19" "librosa>=0.8.1" "matplotlib<=3.7" tqdm

from encodec import compress, decompress 
from encodec.utils import convert_audio, save_audio 
from encodec.compress import MODELS 
import torchaudio 
import torch 
import typing as tp 

model_id = "encodec_24khz" 
# 事前学習済み EnCodec モデルをインスタンス化 
model = MODELS[model_id]() 
model.set_target_bandwidth(6.0)

import librosa 
import matplotlib.pyplot as plt 
import librosa.display 
import IPython.display as ipd 

# `notebook_utils` モジュールを取得 
import requests 

r = requests.get( 
    url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py", 
) 

open("notebook_utils.py", "w").write(r.text) 
from notebook_utils import download_file 

test_data_url = "https://github.com/facebookresearch/encodec/raw/main/test_24k.wav" 

sample_file = "test_24k.wav" 
download_file(test_data_url, sample_file) 
audio, sr = librosa.load(sample_file) 
plt.figure(figsize=(14, 5)) 
librosa.display.waveshow(audio, sr=sr) 

ipd.Audio(sample_file)

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

    def forward(self, x: torch.Tensor): 
        codes, scale = self.model._encode_frame(x) 
        if not self.model.normalize: 
            return codes 
        return codes, scale

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

    def forward(self, codes, scale=None): 
        return model._decode_frame((codes, scale))

encoder = FrameEncoder(model) 
decoder = FrameDecoder(model)

import openvino as ov 

core = ov.Core() 

OV_ENCODER_PATH = Path("encodec_encoder.xml") 
if not OV_ENCODER_PATH.exists(): 
    encoder_ov = ov.convert_model(encoder, example_input=torch.zeros(1, 1, 480000), input=[[1, 1, 480000]]) 
    ov.save_model(encoder_ov, OV_ENCODER_PATH) 
else: 
    encoder_ov = core.read_model(OV_ENCODER_PATH)

OV_DECODER_PATH = Path("encodec_decoder.xml") 
if not OV_DECODER_PATH.exists(): 
    decoder_ov = ov.convert_model( 
        decoder, 
        example_input=torch.zeros([1, 8, 1500], dtype=torch.long), 
        input=[[1, 8, 1500]], 
    ) 
    ov.save_model(decoder_ov, OV_DECODER_PATH) 
else: 
    decoder_ov = core.read_model(OV_DECODER_PATH)

b = compress(model, wav, use_lm=False) 
out_file = Path("compressed_ov.ecdc") 
out_file.write_bytes(b)

out, out_sr = decompress(out_file.read_bytes())

ov_output_file = "decopressed_ov.wav" 
save_audio(out, ov_output_file, out_sr)

audio, sr = librosa.load(ov_output_file) 
plt.figure(figsize=(14, 5)) 
librosa.display.waveshow(audio, sr=sr) 

ipd.Audio(ov_output_file)

import gradio as gr 
from typing import Tuple 
import numpy as np 

def preprocess(input, sample_rate, model_sr, model_channels): 
    input = torch.tensor(input, dtype=torch.float32) 
    input = input / 2**15 # int16 スケールに調整 
    input = input.unsqueeze(0) 
    input = convert_audio(input, sample_rate, model_sr, model_channels) 

    return input 

def postprocess(output): 
    output = output.squeeze() 
    output = output * 2**15 #  [-1, 1] スケールに調整 
    output = output.numpy(force=True) 
    output = output.astype(np.int16) 

    return output 

def _compress(input: Tuple[int, np.ndarray]): 
    sample_rate, waveform = input 
    waveform = preprocess(waveform, sample_rate, model_sr, model_channels) 

    b = compress(model, waveform, use_lm=False) 
    out, out_sr = decompress(b) 
    out = postprocess(out) 
    return out_sr, out 

demo = gr.Interface(_compress, "audio", "audio", examples=["test_24k.wav"]) 

try: 
    demo.launch(debug=False) 
except Exception: 
    demo.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/