RealTime (Stream) implementation

[rep-stosw]

Hi @lucadellalib,

Thanks for best VLBR codec! 👍

I try switch this codec into Stream mode for realtime using: feed fixed portions of audio-samples. But output stream made very joke and laugth, because there is no strict syncronization in tokens level.

How to proper use this codec in stream mode?

Below some code of my tests:

#...............
#some import from modules...
#...............

coder=WavLM()
decoder=Vocos()
compressor=FocalEncoder()
decompressor=FocalDecoder()
quantizer=BinarySphericalQuantizer()


#...............
#then load models data: state dicts...
#...............

coder.eval().requires_grad_(False)
decoder.eval().requires_grad_(False)
compressor.eval().requires_grad_(False)
decompressor.eval().requires_grad_(False)
quantizer.eval().requires_grad_(False)

sig,sample_rate=torchaudio.load("input.wav")

L=1280*8 #16

N=len(sig[0,:])//L

dec=np.array(range(L*N)).reshape((1,L*N)).astype('float32')
dec=torch.from_numpy(dec)

for n in range(N):
    
    print("\n")
    print(n)
    
    Frame=sig[:,n*L:(n+1)*L]
    print(Frame.shape)

    feats=coder(Frame)
    print(feats.shape)

    lats=compressor(feats)
    lats=nn.functional.normalize(lats,dim=-1)
    print(lats.shape)

    codes=quantizer.lats_to_codes(lats)
    print(codes.shape)

    toks=quantizer.codes_to_toks(codes)
    print(toks.shape)

    print(toks)

    #

    codes=quantizer.toks_to_codes(toks)
    print(codes.shape)

    qfeats=decompressor(codes)
    print(qfeats.shape)

    dsig=decoder(qfeats)
    print(dsig.shape)
    
    #dec[:,n*L:(n+1)*L]=dsig
    dec[:,n*L:(n+1)*L]=dsig.detach() #.clone()
    
    print(dec[:,n*L:(n+1)*L].shape)
    
torchaudio.save("reconstruction.wav",dec,sample_rate)

[lucadellalib]

Hi @rep-stosw, thank you for your interest in our work! You can find an example of online resynthesis in focalcodec/codec.py at line 1032 (see below). Note, however, that for acceptable performance you should split your audio into chunks of at least 500 ms, ideally 1-2 seconds. In general, the longer the chunks, the better the results.

@torch.no_grad()
def test_online(
    config: "str" = "lucadellalib/focalcodec_50hz",
) -> "None":

    def overlap_add(
        chunks: "Sequence[Tensor]", chunk_size: "int", overlap_size: "int"
    ) -> "Tensor":
        device = chunks[0].device
        step_size = chunk_size - overlap_size
        batch_size = chunks[0].shape[0]
        num_chunks = len(chunks)
        total_length = step_size * (num_chunks - 1) + chunk_size

        # Stack all processed chunks into a tensor of shape (batch_size, num_chunks, chunk_size)
        chunks_tensor = torch.stack(chunks, dim=1)

        # Create cross-fading window
        fade_in = torch.linspace(0, 1, overlap_size, device=device)
        fade_out = 1 - fade_in
        window = torch.cat(
            [
                fade_out,
                torch.ones(chunk_size - 2 * overlap_size, device=device),
                fade_in,
            ]
        )

        # Reshape the window for broadcasting
        window = window.reshape(1, 1, -1)  # Shape (1, 1, chunk_size)

        # Apply the window to all chunks
        weighted_chunks = (
            chunks_tensor * window
        )  # Shape (batch_size, num_chunks, chunk_size)

        # Initialize output and overlap count
        output = torch.zeros(batch_size, total_length, device=device)
        overlap_count = torch.zeros(batch_size, total_length, device=device)

        # Place weighted chunks in their respective positions
        for i in range(num_chunks):
            start = i * step_size
            end = start + chunk_size
            output[:, start:end] += weighted_chunks[:, i, :]
            overlap_count[:, start:end] += window[0]

        # Normalize by overlap count to handle overlapping regions
        output /= overlap_count
        return output

    try:
        import torchaudio
    except ImportError:
        raise ImportError("`pip install torchaudio` to run this script")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = FocalCodec.from_pretrained(config).to(device)
    model.eval()

    root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    paths = [
        os.path.join("librispeech-dev-clean", "84"),
    ]
    matching_set = {k: [] for k in paths}
    for path in paths:
        for filename in os.listdir(os.path.join(root_dir, "audios", path)):
            filepath = os.path.join(root_dir, "audios", path, filename)
            sig, sample_rate = torchaudio.load(filepath)
            sig = torchaudio.functional.resample(sig, sample_rate, 16000)
            sig = sig.to(device)
            feats = model.sig_to_feats(sig)
            matching_set[path].append(feats[0])
        matching_set[path] = torch.cat(matching_set[path])

    sig, sample_rate = torchaudio.load(
        os.path.join(
            root_dir, "audios", "librispeech-dev-clean", "251-118436-0003.wav"
        )
    )
    sig = torchaudio.functional.resample(sig, sample_rate, 16000)
    sig = sig.to(device)

    chunk_size = 8000
    initial_context = 8000
    context_size = 48000
    overlap_size = 1000

    _toks = model.sig_to_toks(torch.randn(1, chunk_size, device=device))
    _rec_sig = model.toks_to_sig(_toks)
    sig_from_toks_chunk_size = _rec_sig.shape[-1]
    _feats = model.sig_to_feats(torch.randn(1, chunk_size, device=device))
    _rec_sig = model.feats_to_sig(_feats)
    sig_from_feats_chunk_size = _rec_sig.shape[-1]

    sig_from_toks_overlap_size = int(
        overlap_size * sig_from_toks_chunk_size / chunk_size
    )
    sig_from_feats_overlap_size = int(
        overlap_size * sig_from_feats_chunk_size / chunk_size
    )

    all_sig_from_feats = []
    all_sig_from_lats = []
    all_sig_from_toks = []
    all_sig_from_toks_vc = []
    frame_idx = 0
    while frame_idx < sig.shape[1]:
        if frame_idx >= initial_context:
            chunk = sig[:, max(0, frame_idx - context_size) : frame_idx]
            feats = model.sig_to_feats(chunk)
            lats = model.feats_to_lats(feats)
            toks = model.lats_to_toks(lats)
            sig_from_feats = model.feats_to_sig(feats)[:, -sig_from_feats_chunk_size:]
            sig_from_lats = model.codes_to_sig(lats)[:, -sig_from_toks_chunk_size:]
            sig_from_toks = model.toks_to_sig(toks)[:, -sig_from_toks_chunk_size:]
            sig_from_toks_vc = model.toks_to_sig(
                toks,
                matching_set["librispeech-dev-clean/84"],
            )[:, -sig_from_toks_chunk_size:]

            all_sig_from_feats.append(sig_from_feats)
            all_sig_from_lats.append(sig_from_lats)
            all_sig_from_toks.append(sig_from_toks)
            all_sig_from_toks_vc.append(sig_from_toks_vc)
        frame_idx += chunk_size - overlap_size

    all_sig_from_feats = overlap_add(
        all_sig_from_feats, sig_from_feats_chunk_size, sig_from_feats_overlap_size
    )
    all_sig_from_lats = overlap_add(
        all_sig_from_lats, sig_from_toks_chunk_size, sig_from_toks_overlap_size
    )
    all_sig_from_toks = overlap_add(
        all_sig_from_toks, sig_from_toks_chunk_size, sig_from_toks_overlap_size
    )
    all_sig_from_toks_vc = overlap_add(
        all_sig_from_toks_vc, sig_from_toks_chunk_size, sig_from_toks_overlap_size
    )

    output_dir = os.path.join(root_dir, "reconstructions")
    os.makedirs(output_dir, exist_ok=True)
    torchaudio.save(
        os.path.join(output_dir, "sig.wav"),
        sig.cpu(),
        16000,
    )
    torchaudio.save(
        os.path.join(output_dir, "sig_from_feats_online.wav"),
        all_sig_from_feats.cpu(),
        16000,
    )
    torchaudio.save(
        os.path.join(output_dir, "sig_from_lats_online.wav"),
        all_sig_from_lats.cpu(),
        16000,
    )
    torchaudio.save(
        os.path.join(output_dir, "sig_from_toks_online.wav"),
        all_sig_from_toks.cpu(),
        16000,
    )
    torchaudio.save(
        os.path.join(output_dir, "sig_from_toks_online_vc.wav"),
        all_sig_from_toks_vc.cpu(),
        16000,
    )

[rep-stosw]

Hi, @lucadellalib, thanks for the answer!

I did to make a streaming playback of the codec. Since my main goal was to run this codec on ARM Cortex-A7 CPU, I had to reduce the neural network models and retrain them:

1. Instead of WavLM, I use my model (MelLM), which uses a spectrum of 80 MEL bands instead of an audio wave. The output size of the encoder's tensor is not 1024, but 256 (features). I trained my MelLM model by distilling knowledge from pretrained WavLM, the output tensor was transformed from 256 to 1024 through a linear projection layer(low-to-high linear projection uses only in training phase). L2 loss.

2. FocalNET. Encoder and Decoder internal layers: 256/128/64. Trained on fixed speech segments of 64 frames (1280 ms). L2 loss. I noticed that this part of the codec depends on the language. And shorter segment sizes will fine for streaming versions vs. offline versions with full size of speech utterance.

3. Vocos vocoder is also reduced: the number of residual blocks is reduced to 5 (in the original has 8), the number of features in the internal blocks is reduced, the input tensor of dimension is 256 instead of 1024. Training is the same as in the official version of Vocos.

As a result, I did to run a reduced version of the codec on ARM Cortex 1200 MHz. I used onnxruntime as a neural network engine (ported this code to hardware), and I converted the models themselves to ONNX, then to ORT. The result is a real-time codec that encodes human speech from a microphone, decodes and outputs the decoded voice to headphones (with a delay of up to 1.3 s).

I take audio segments with overlap from the input. And during playback, I glued them together with overlap, matching using window functions (Hann or Bartlett works great here). The overlap coefficient was set - initially I took an overlap of 50%, but then I reduced the overlap percentage to 1/8. I reduced the quantized tensor dimension from 13 to 12 (after Fcocal Compressor or BSQ), since it is easier to pack the stream into bits: 2 tensors of 12 bits give as many as 3 bytes. Total 50 Hz * 12 = 600 bps, taking into account the overlap it turned out a little more: 600 * 16/15 = 640 bps. Sounds impressive.

I will have to try checking at lower bitrates: 320, 160 and (I'm afraid to imagine) at 80 bps.

Below is the python code: the sound from the microphone is encoded and decoded with output to the headphones in real time (here used self-own MelLM encoder instead WavLM, and modified Vocos: IFFT routine was removed from model to general python code, because on target platform FFT and IFFT was optimized via ARM NEON pipeline):

import torch
import torchaudio
from torch import Tensor,nn

from wavlm import WavLM

from vocos import Vocos

from mel2wav.modules import Audio2Mel,MelLMT
from focalnet import FocalDecoder,FocalEncoder

import math

import numpy as np
from scipy.io import wavfile

import librosa

import pyaudio
import wave

# ---------------------------------------------------------------------------------------------------------------

device="cpu" #torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---------------------------------------------------------------------------------------------------------------

SR=16000
MEL=80
FEATURE=256
LATENT=12
HOP=320
SEGMENT=64
FFTW=1024

# ---------------------------------------------------------------------------------------------------------------

#only for Inference! For train need STE via SG.
def BSQ(z,qdim=LATENT):
    z=nn.functional.normalize(z,dim=-1)
    zq=torch.where(z>0.0,torch.tensor(1.0,dtype=z.dtype,device=z.device),torch.tensor(-1.0,dtype=z.dtype,device=z.device))
    q_scale=1.0/(qdim**0.5)
    return zq*q_scale

# ---------------------------------------------------------------------------------------------------------------

fft=Audio2Mel(n_mel_channels=MEL).to(device)

coder=MelLMT(input_size=MEL,output_size=FEATURE).to(device)
compressor=FocalEncoder(input_dim=FEATURE,output_dim=LATENT,hidden_dims=(FEATURE,FEATURE//2,FEATURE//4),downscale_factors=(1,1,1)).to(device)
decompressor=FocalDecoder(input_dim=LATENT,output_dim=FEATURE,hidden_dims=(FEATURE//4,FEATURE//2,FEATURE),upscale_factors=(1,1,1)).to(device)

#------------------------------------

def load_vocos(checkpoint_path,config_path,device):
    model=Vocos.from_hparams(config_path).to(device)
    raw_model=torch.load(checkpoint_path,map_location=device)
    raw_model=raw_model if 'state_dict' not in raw_model else raw_model['state_dict']
    model.load_state_dict(raw_model,strict=False)
    model.eval()
    return model

decoder=load_vocos("vocos_checkpoint_epoch=69_step=476700_val_loss=5.7498.ckpt","configs/vocos_mellm2.yaml",device)

#------------------------------------

coder.load_state_dict(torch.load("best_mellm.pt",weights_only=True))
compressor.load_state_dict(torch.load("best_compressor.pt",weights_only=True))
decompressor.load_state_dict(torch.load("best_decompressor.pt",weights_only=True))

coder.eval().requires_grad_(False)
compressor.eval().requires_grad_(False)
decompressor.eval().requires_grad_(False)
decoder.eval().requires_grad_(False)

#---------------------------------------------------------------------

K=16
STEP=(SEGMENT//K)*(K-1)
WIDTH=SEGMENT-STEP

#filt=torch.bartlett_window(WIDTH*2*HOP,device=device)
filt=torch.hann_window(WIDTH*2*HOP,device=device)

FILT=torch.ones(size=(SEGMENT*HOP,),dtype=torch.float32,device=device)

FILT[        :WIDTH*HOP]=filt[         :WIDTH*HOP]
FILT[STEP*HOP:         ]=filt[WIDTH*HOP:         ]

DATA=[] #Input Wav Buffer
OUT=[]  #Output Wav Buffer

PredIN =torch.zeros(size=(1,WIDTH*HOP),dtype=torch.float32,device=device)
PredOUT=torch.zeros(size=(1,WIDTH*HOP),dtype=torch.float32,device=device)

#--------------------------------------------------------------------------

audio=pyaudio.PyAudio()

stream_input=audio.open(
                        format=pyaudio.paInt16,
                        channels=1,
                        rate=SR,
                        frames_per_buffer=SEGMENT*HOP,
                        input=True
                       )

stream_loopback=audio.open(
                           format=pyaudio.paInt16,
                           channels=1,
                           rate=SR,
                           frames_per_buffer=SEGMENT*HOP,
                           output=True
                          )

#--------------------------------------------------------------------------

i=0

print('\nAudio recording...\n')

try:
    while True:

        data=stream_input.read(STEP*HOP)
        #stream_loopback.write(data)

        data=np.frombuffer(data,dtype=np.int16)

        DATA.append(data)

        data=data.astype(np.float32)/32768.0   #normalize [-1..1]
        data=torch.from_numpy(data).to(device)
        data=torch.unsqueeze(data,0)

        sigIN=torch.cat((PredIN,data),1)
        PredIN=data[:,(STEP-WIDTH)*HOP:]

        IN=fft(sigIN).transpose(1,2) #get MEL spectrogram (1,SEGMENT,MEL)

        feats=coder(IN)              #MelLMT
        lats=compressor(feats)       #FocalEncoder

        qlats=BSQ(lats)              #BSQ

        qfeats=decompressor(qlats);  #FocalDecoder
        #qfeats=feats                 #bypass FocalNET - direct MelLMT to Vocos

        pre=decoder.backbone(qfeats) #Vocos

        mag,pha=pre.chunk(2,dim=2)

        mag=torch.exp(mag)
        mag=torch.clip(mag,max=1e2)

        S=torch.polar(mag,pha)
        S=S.transpose(1,2).detach().cpu().numpy()

        ifft=librosa.istft(S,hop_length=HOP,win_length=FFTW,n_fft=FFTW,window='hann',center=True,length=SEGMENT*HOP)
        ifft=torch.from_numpy(ifft).to(device)
        ifft*=FILT

        sigOUT=ifft[:,:STEP*HOP]
        sigOUT[:,:WIDTH*HOP]+=PredOUT

        PredOUT=ifft[:,STEP*HOP:]

        sigOUT*=32768.0
        sigOUT=sigOUT.cpu().numpy()
        sigOUT=sigOUT.astype(np.int16)

        OUT.append(sigOUT)

        stream_loopback.write(sigOUT.tobytes(order='C'))

        print(i)
        i+=1

except KeyboardInterrupt:
    print("\nLoop Broken")

#--------------------------------------------------------------------

print('\nAudio saving...')

#Stop and close the stream 
stream_input.stop_stream()
stream_input.close()

#Terminate the PortAudio interface
audio.terminate()

wf=wave.open('in.wav','wb')
wf.setnchannels(1)
wf.setsampwidth(audio.get_sample_size(pyaudio.paInt16))
wf.setframerate(SR)
wf.writeframes(b''.join(DATA))
wf.close()

wf=wave.open('out.wav','wb')
wf.setnchannels(1)
wf.setsampwidth(audio.get_sample_size(pyaudio.paInt16))
wf.setframerate(SR)
wf.writeframes(b''.join(OUT))
wf.close()

print("\nOK!")

[lucadellalib]

@rep-stosw great work! Glad to hear the codec can be adapted to run on low-resource devices. Does the resynthesis quality degrade much in that setting?

[rep-stosw]

@lucadellalib , at a token frequency of 50 Hz (FocalNET dividers: [1,1,1] ) speech in streaming mode is quite intelligible and does not depend on the speaker's language.

I have attached several sample statements:

1. Original (251-118436-0003.wav):

251-118436-0003.wav

2. Direct connection of MelLM with Vocos (without FocalNET and without BSQ):

test(wo focal).wav

3. Full pipeline (MelLM => Focal Compressor => BSQ => Focal Decompressor => Vocos):

test(bsq).wav

This is 640 bps (600 bps *16/15 overlap)

Before that, I tried at 163.5 bps (token frequency of 12.5 Hz). In this mode, speech was recognizable, but sometimes it degenerated into childish babble. Syllables also sometimes get mixed up. But if speak slower and more clearly, then the recognition increases.

On the target platform (where I ported the codec), the dominant factor is the limited computing resource. And then only the limitation of the model sizes.

It turned out that to ensure operation in streaming mode on the target platform, the NN model sizes must be mercilessly reduced in order to have time to do one cycle of encoding and decoding.

As an encoder, I use this model of a lightweight network (MelLM):

Slope=0.2

class Residual(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_hiddens):
        super(Residual, self).__init__()
        self._block = nn.Sequential(

            nn.BatchNorm1d(in_channels),
            nn.LeakyReLU(negative_slope=Slope,inplace=True),
            nn.Conv1d(in_channels=in_channels,
                      out_channels=num_residual_hiddens,
                      kernel_size=3, stride=1, padding=1, bias=False),

            nn.BatchNorm1d(num_residual_hiddens),
            nn.LeakyReLU(negative_slope=Slope,inplace=True),
            nn.Conv1d(in_channels=num_residual_hiddens,
                      out_channels=num_hiddens,
                      kernel_size=1, stride=1, bias=False)
        )
    
    def forward(self, x):
        return x + self._block(x)

class ResidualStack(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
        super(ResidualStack, self).__init__()
        self._num_residual_layers = num_residual_layers
        self._layers = nn.ModuleList([Residual(in_channels, num_hiddens, num_residual_hiddens)
                             for _ in range(self._num_residual_layers)])

    def forward(self, x):
        for i in range(self._num_residual_layers):
            x = self._layers[i](x)
        return x

class MelLMT(nn.Module):
    def __init__(self,input_size,output_size):

        super().__init__()

        self.model=nn.Sequential(
            nn.Conv1d(input_size, output_size, kernel_size=3, stride=1, padding=1, bias=False),
            ResidualStack(in_channels=output_size,num_hiddens=output_size,num_residual_layers=1,num_residual_hiddens=output_size),
            nn.BatchNorm1d(output_size),
            nn.LeakyReLU(negative_slope=Slope,inplace=True),
        )        

        self.TransformerEncoder=nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=output_size,nhead=16,batch_first=True),num_layers=2,enable_nested_tensor=False)

    def forward(self,x):
        x=x.transpose(1,2)
        x=self.model(x)
        x=x.transpose(1,2)
        x=self.TransformerEncoder(x)
        return x

As an increase in the features tensor (from 256 to 1024), I use a linear layer without bias:

class Projection(nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()

        self.model = nn.Sequential(
            nn.Linear(input_size, output_size, bias=False),
        )

    def forward(self, x):
        return self.model(x)

Maybe, ideally, it would be possible to come up with a network more successful than this one. But for me, the main criterion was that all codec models could process one cycle in a strictly specified time.

In your opinion, what can be added or removed from the encoder model to make it work more efficiently?

I feed the MEL spectrogram (1,64,80) to the model input, and the model outputs features (1,64,256). The linear layer of the up-projection gives the tensor (1,64,1024) - in order to minimize Loss L2 when distilling knowledge from the previously trained WavLM.

I compare WavLM output tensor with output tensor after linear projection(80MEL spectra => MelLM => Linear Proj.) by loss L2. Is it correct?

P.S. I use features dim 256 instead 1024 because with 1024 models are huge and runs slow on my target platform.

[lucadellalib]

@rep-stosw it sounds correct!

I've just updated the repo with the latest code and checkpoints I've been working on over the past few months to make the model streamable. The new model is larger in size, but it can be streamed with chunks of 80 ms. It probably won't fit within the constraints of your system, but it could be a better starting point for distilling into a smaller model. Feel free to give it a try!

[rep-stosw]

@lucadellalib Excellent! I tested stream mode from file and from microphone: all works fine. And not depend from language. Thanks!

Now I need dig causal convolution, make reduced models and distill knowledges from big to small models.

What about causal low-bitrate modes with token rates: 25 / 12.5 / 6.25 Hz ?