ai-content-maker/.venv/Lib/site-packages/encodec/model.py

326 lines
12 KiB
Python

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""EnCodec model implementation."""
import math
from pathlib import Path
import typing as tp
import numpy as np
import torch
from torch import nn
from . import quantization as qt
from . import modules as m
from .utils import _check_checksum, _linear_overlap_add, _get_checkpoint_url
ROOT_URL = 'https://dl.fbaipublicfiles.com/encodec/v0/'
EncodedFrame = tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]
class LMModel(nn.Module):
"""Language Model to estimate probabilities of each codebook entry.
We predict all codebooks in parallel for a given time step.
Args:
n_q (int): number of codebooks.
card (int): codebook cardinality.
dim (int): transformer dimension.
**kwargs: passed to `encodec.modules.transformer.StreamingTransformerEncoder`.
"""
def __init__(self, n_q: int = 32, card: int = 1024, dim: int = 200, **kwargs):
super().__init__()
self.card = card
self.n_q = n_q
self.dim = dim
self.transformer = m.StreamingTransformerEncoder(dim=dim, **kwargs)
self.emb = nn.ModuleList([nn.Embedding(card + 1, dim) for _ in range(n_q)])
self.linears = nn.ModuleList([nn.Linear(dim, card) for _ in range(n_q)])
def forward(self, indices: torch.Tensor,
states: tp.Optional[tp.List[torch.Tensor]] = None, offset: int = 0):
"""
Args:
indices (torch.Tensor): indices from the previous time step. Indices
should be 1 + actual index in the codebook. The value 0 is reserved for
when the index is missing (i.e. first time step). Shape should be
`[B, n_q, T]`.
states: state for the streaming decoding.
offset: offset of the current time step.
Returns a 3-tuple `(probabilities, new_states, new_offset)` with probabilities
with a shape `[B, card, n_q, T]`.
"""
B, K, T = indices.shape
input_ = sum([self.emb[k](indices[:, k]) for k in range(K)])
out, states, offset = self.transformer(input_, states, offset)
logits = torch.stack([self.linears[k](out) for k in range(K)], dim=1).permute(0, 3, 1, 2)
return torch.softmax(logits, dim=1), states, offset
class EncodecModel(nn.Module):
"""EnCodec model operating on the raw waveform.
Args:
target_bandwidths (list of float): Target bandwidths.
encoder (nn.Module): Encoder network.
decoder (nn.Module): Decoder network.
sample_rate (int): Audio sample rate.
channels (int): Number of audio channels.
normalize (bool): Whether to apply audio normalization.
segment (float or None): segment duration in sec. when doing overlap-add.
overlap (float): overlap between segment, given as a fraction of the segment duration.
name (str): name of the model, used as metadata when compressing audio.
"""
def __init__(self,
encoder: m.SEANetEncoder,
decoder: m.SEANetDecoder,
quantizer: qt.ResidualVectorQuantizer,
target_bandwidths: tp.List[float],
sample_rate: int,
channels: int,
normalize: bool = False,
segment: tp.Optional[float] = None,
overlap: float = 0.01,
name: str = 'unset'):
super().__init__()
self.bandwidth: tp.Optional[float] = None
self.target_bandwidths = target_bandwidths
self.encoder = encoder
self.quantizer = quantizer
self.decoder = decoder
self.sample_rate = sample_rate
self.channels = channels
self.normalize = normalize
self.segment = segment
self.overlap = overlap
self.frame_rate = math.ceil(self.sample_rate / np.prod(self.encoder.ratios))
self.name = name
self.bits_per_codebook = int(math.log2(self.quantizer.bins))
assert 2 ** self.bits_per_codebook == self.quantizer.bins, \
"quantizer bins must be a power of 2."
@property
def segment_length(self) -> tp.Optional[int]:
if self.segment is None:
return None
return int(self.segment * self.sample_rate)
@property
def segment_stride(self) -> tp.Optional[int]:
segment_length = self.segment_length
if segment_length is None:
return None
return max(1, int((1 - self.overlap) * segment_length))
def encode(self, x: torch.Tensor) -> tp.List[EncodedFrame]:
"""Given a tensor `x`, returns a list of frames containing
the discrete encoded codes for `x`, along with rescaling factors
for each segment, when `self.normalize` is True.
Each frames is a tuple `(codebook, scale)`, with `codebook` of
shape `[B, K, T]`, with `K` the number of codebooks.
"""
assert x.dim() == 3
_, channels, length = x.shape
assert channels > 0 and channels <= 2
segment_length = self.segment_length
if segment_length is None:
segment_length = length
stride = length
else:
stride = self.segment_stride # type: ignore
assert stride is not None
encoded_frames: tp.List[EncodedFrame] = []
for offset in range(0, length, stride):
frame = x[:, :, offset: offset + segment_length]
encoded_frames.append(self._encode_frame(frame))
return encoded_frames
def _encode_frame(self, x: torch.Tensor) -> EncodedFrame:
length = x.shape[-1]
duration = length / self.sample_rate
assert self.segment is None or duration <= 1e-5 + self.segment
if self.normalize:
mono = x.mean(dim=1, keepdim=True)
volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
scale = 1e-8 + volume
x = x / scale
scale = scale.view(-1, 1)
else:
scale = None
emb = self.encoder(x)
codes = self.quantizer.encode(emb, self.frame_rate, self.bandwidth)
codes = codes.transpose(0, 1)
# codes is [B, K, T], with T frames, K nb of codebooks.
return codes, scale
def decode(self, encoded_frames: tp.List[EncodedFrame]) -> torch.Tensor:
"""Decode the given frames into a waveform.
Note that the output might be a bit bigger than the input. In that case,
any extra steps at the end can be trimmed.
"""
segment_length = self.segment_length
if segment_length is None:
assert len(encoded_frames) == 1
return self._decode_frame(encoded_frames[0])
frames = [self._decode_frame(frame) for frame in encoded_frames]
return _linear_overlap_add(frames, self.segment_stride or 1)
def _decode_frame(self, encoded_frame: EncodedFrame) -> torch.Tensor:
codes, scale = encoded_frame
codes = codes.transpose(0, 1)
emb = self.quantizer.decode(codes)
out = self.decoder(emb)
if scale is not None:
out = out * scale.view(-1, 1, 1)
return out
def forward(self, x: torch.Tensor) -> torch.Tensor:
frames = self.encode(x)
return self.decode(frames)[:, :, :x.shape[-1]]
def set_target_bandwidth(self, bandwidth: float):
if bandwidth not in self.target_bandwidths:
raise ValueError(f"This model doesn't support the bandwidth {bandwidth}. "
f"Select one of {self.target_bandwidths}.")
self.bandwidth = bandwidth
def get_lm_model(self) -> LMModel:
"""Return the associated LM model to improve the compression rate.
"""
torch.manual_seed(1234) # todo remove: this
device = next(self.parameters()).device
lm = LMModel(self.quantizer.n_q, self.quantizer.bins, num_layers=5, dim=200,
past_context=int(3.5 * self.frame_rate)).to(device)
checkpoints = {
'encodec_24khz': 'encodec_lm_24khz-1608e3c0.th',
'encodec_48khz': 'encodec_lm_48khz-7add9fc3.th',
}
try:
checkpoint_name = checkpoints[self.name]
except KeyError:
raise RuntimeError("No LM pre-trained for the current Encodec model.")
url = _get_checkpoint_url(ROOT_URL, checkpoint_name)
state = torch.hub.load_state_dict_from_url(
url, map_location='cpu', check_hash=True) # type: ignore
lm.load_state_dict(state)
lm.eval()
return lm
@staticmethod
def _get_model(target_bandwidths: tp.List[float],
sample_rate: int = 24_000,
channels: int = 1,
causal: bool = True,
model_norm: str = 'weight_norm',
audio_normalize: bool = False,
segment: tp.Optional[float] = None,
name: str = 'unset'):
encoder = m.SEANetEncoder(channels=channels, norm=model_norm, causal=causal)
decoder = m.SEANetDecoder(channels=channels, norm=model_norm, causal=causal)
n_q = int(1000 * target_bandwidths[-1] // (math.ceil(sample_rate / encoder.hop_length) * 10))
quantizer = qt.ResidualVectorQuantizer(
dimension=encoder.dimension,
n_q=n_q,
bins=1024,
)
model = EncodecModel(
encoder,
decoder,
quantizer,
target_bandwidths,
sample_rate,
channels,
normalize=audio_normalize,
segment=segment,
name=name,
)
return model
@staticmethod
def _get_pretrained(checkpoint_name: str, repository: tp.Optional[Path] = None):
if repository is not None:
if not repository.is_dir():
raise ValueError(f"{repository} must exist and be a directory.")
file = repository / checkpoint_name
checksum = file.stem.split('-')[1]
_check_checksum(file, checksum)
return torch.load(file)
else:
url = _get_checkpoint_url(ROOT_URL, checkpoint_name)
return torch.hub.load_state_dict_from_url(url, map_location='cpu', check_hash=True) # type:ignore
@staticmethod
def encodec_model_24khz(pretrained: bool = True, repository: tp.Optional[Path] = None):
"""Return the pretrained causal 24khz model.
"""
if repository:
assert pretrained
target_bandwidths = [1.5, 3., 6, 12., 24.]
checkpoint_name = 'encodec_24khz-d7cc33bc.th'
sample_rate = 24_000
channels = 1
model = EncodecModel._get_model(
target_bandwidths, sample_rate, channels,
causal=True, model_norm='weight_norm', audio_normalize=False,
name='encodec_24khz' if pretrained else 'unset')
if pretrained:
state_dict = EncodecModel._get_pretrained(checkpoint_name, repository)
model.load_state_dict(state_dict)
model.eval()
return model
@staticmethod
def encodec_model_48khz(pretrained: bool = True, repository: tp.Optional[Path] = None):
"""Return the pretrained 48khz model.
"""
if repository:
assert pretrained
target_bandwidths = [3., 6., 12., 24.]
checkpoint_name = 'encodec_48khz-7e698e3e.th'
sample_rate = 48_000
channels = 2
model = EncodecModel._get_model(
target_bandwidths, sample_rate, channels,
causal=False, model_norm='time_group_norm', audio_normalize=True,
segment=1., name='encodec_48khz' if pretrained else 'unset')
if pretrained:
state_dict = EncodecModel._get_pretrained(checkpoint_name, repository)
model.load_state_dict(state_dict)
model.eval()
return model
def test():
from itertools import product
import torchaudio
bandwidths = [3, 6, 12, 24]
models = {
'encodec_24khz': EncodecModel.encodec_model_24khz,
'encodec_48khz': EncodecModel.encodec_model_48khz
}
for model_name, bw in product(models.keys(), bandwidths):
model = models[model_name]()
model.set_target_bandwidth(bw)
audio_suffix = model_name.split('_')[1][:3]
wav, sr = torchaudio.load(f"test_{audio_suffix}.wav")
wav = wav[:, :model.sample_rate * 2]
wav_in = wav.unsqueeze(0)
wav_dec = model(wav_in)[0]
assert wav.shape == wav_dec.shape, (wav.shape, wav_dec.shape)
if __name__ == '__main__':
test()