Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| from einops import rearrange | |
| from torch import nn | |
| class Pretransform(nn.Module): | |
| def __init__(self, enable_grad, io_channels, is_discrete): | |
| super().__init__() | |
| self.is_discrete = is_discrete | |
| self.io_channels = io_channels | |
| self.encoded_channels = None | |
| self.downsampling_ratio = None | |
| self.enable_grad = enable_grad | |
| def encode(self, x): | |
| raise NotImplementedError | |
| def decode(self, z): | |
| raise NotImplementedError | |
| def tokenize(self, x): | |
| raise NotImplementedError | |
| def decode_tokens(self, tokens): | |
| raise NotImplementedError | |
| class AutoencoderPretransform(Pretransform): | |
| def __init__(self, model, scale=1.0, model_half=False, iterate_batch=False, chunked=False): | |
| super().__init__(enable_grad=False, io_channels=model.io_channels, is_discrete=model.bottleneck is not None and model.bottleneck.is_discrete) | |
| self.model = model | |
| self.model.requires_grad_(False).eval() | |
| self.scale=scale | |
| self.downsampling_ratio = model.downsampling_ratio | |
| self.io_channels = model.io_channels | |
| self.sample_rate = model.sample_rate | |
| self.model_half = model_half | |
| self.iterate_batch = iterate_batch | |
| self.encoded_channels = model.latent_dim | |
| self.chunked = chunked | |
| self.num_quantizers = model.bottleneck.num_quantizers if model.bottleneck is not None and model.bottleneck.is_discrete else None | |
| self.codebook_size = model.bottleneck.codebook_size if model.bottleneck is not None and model.bottleneck.is_discrete else None | |
| if self.model_half: | |
| self.model.half() | |
| def encode(self, x, **kwargs): | |
| if self.model_half: | |
| x = x.half() | |
| self.model.to(torch.float16) | |
| encoded = self.model.encode_audio(x, chunked=self.chunked, iterate_batch=self.iterate_batch, **kwargs) | |
| if self.model_half: | |
| encoded = encoded.float() | |
| return encoded / self.scale | |
| def decode(self, z, **kwargs): | |
| z = z * self.scale | |
| if self.model_half: | |
| z = z.half() | |
| self.model.to(torch.float16) | |
| decoded = self.model.decode_audio(z, chunked=self.chunked, iterate_batch=self.iterate_batch, **kwargs) | |
| if self.model_half: | |
| decoded = decoded.float() | |
| return decoded | |
| def tokenize(self, x, **kwargs): | |
| assert self.model.is_discrete, "Cannot tokenize with a continuous model" | |
| _, info = self.model.encode(x, return_info = True, **kwargs) | |
| return info[self.model.bottleneck.tokens_id] | |
| def decode_tokens(self, tokens, **kwargs): | |
| assert self.model.is_discrete, "Cannot decode tokens with a continuous model" | |
| return self.model.decode_tokens(tokens, **kwargs) | |
| def load_state_dict(self, state_dict, strict=True): | |
| self.model.load_state_dict(state_dict, strict=strict) | |
| class WaveletPretransform(Pretransform): | |
| def __init__(self, channels, levels, wavelet): | |
| super().__init__(enable_grad=False, io_channels=channels, is_discrete=False) | |
| from .wavelets import WaveletEncode1d, WaveletDecode1d | |
| self.encoder = WaveletEncode1d(channels, levels, wavelet) | |
| self.decoder = WaveletDecode1d(channels, levels, wavelet) | |
| self.downsampling_ratio = 2 ** levels | |
| self.io_channels = channels | |
| self.encoded_channels = channels * self.downsampling_ratio | |
| def encode(self, x): | |
| return self.encoder(x) | |
| def decode(self, z): | |
| return self.decoder(z) | |
| class PQMFPretransform(Pretransform): | |
| def __init__(self, attenuation=100, num_bands=16): | |
| # TODO: Fix PQMF to take in in-channels | |
| super().__init__(enable_grad=False, io_channels=1, is_discrete=False) | |
| from .pqmf import PQMF | |
| self.pqmf = PQMF(attenuation, num_bands) | |
| def encode(self, x): | |
| # x is (Batch x Channels x Time) | |
| x = self.pqmf.forward(x) | |
| # pqmf.forward returns (Batch x Channels x Bands x Time) | |
| # but Pretransform needs Batch x Channels x Time | |
| # so concatenate channels and bands into one axis | |
| return rearrange(x, "b c n t -> b (c n) t") | |
| def decode(self, x): | |
| # x is (Batch x (Channels Bands) x Time), convert back to (Batch x Channels x Bands x Time) | |
| x = rearrange(x, "b (c n) t -> b c n t", n=self.pqmf.num_bands) | |
| # returns (Batch x Channels x Time) | |
| return self.pqmf.inverse(x) | |
| class PretrainedDACPretransform(Pretransform): | |
| def __init__(self, model_type="44khz", model_bitrate="8kbps", scale=1.0, quantize_on_decode: bool = True, chunked=True): | |
| super().__init__(enable_grad=False, io_channels=1, is_discrete=True) | |
| import dac | |
| model_path = dac.utils.download(model_type=model_type, model_bitrate=model_bitrate) | |
| self.model = dac.DAC.load(model_path) | |
| self.quantize_on_decode = quantize_on_decode | |
| if model_type == "44khz": | |
| self.downsampling_ratio = 512 | |
| else: | |
| self.downsampling_ratio = 320 | |
| self.io_channels = 1 | |
| self.scale = scale | |
| self.chunked = chunked | |
| self.encoded_channels = self.model.latent_dim | |
| self.num_quantizers = self.model.n_codebooks | |
| self.codebook_size = self.model.codebook_size | |
| def encode(self, x): | |
| latents = self.model.encoder(x) | |
| if self.quantize_on_decode: | |
| output = latents | |
| else: | |
| z, _, _, _, _ = self.model.quantizer(latents, n_quantizers=self.model.n_codebooks) | |
| output = z | |
| if self.scale != 1.0: | |
| output = output / self.scale | |
| return output | |
| def decode(self, z): | |
| if self.scale != 1.0: | |
| z = z * self.scale | |
| if self.quantize_on_decode: | |
| z, _, _, _, _ = self.model.quantizer(z, n_quantizers=self.model.n_codebooks) | |
| return self.model.decode(z) | |
| def tokenize(self, x): | |
| return self.model.encode(x)[1] | |
| def decode_tokens(self, tokens): | |
| latents = self.model.quantizer.from_codes(tokens) | |
| return self.model.decode(latents) | |
| class AudiocraftCompressionPretransform(Pretransform): | |
| def __init__(self, model_type="facebook/encodec_32khz", scale=1.0, quantize_on_decode: bool = True): | |
| super().__init__(enable_grad=False, io_channels=1, is_discrete=True) | |
| try: | |
| from audiocraft.models import CompressionModel | |
| except ImportError: | |
| raise ImportError("Audiocraft is not installed. Please install audiocraft to use Audiocraft models.") | |
| self.model = CompressionModel.get_pretrained(model_type) | |
| self.quantize_on_decode = quantize_on_decode | |
| self.downsampling_ratio = round(self.model.sample_rate / self.model.frame_rate) | |
| self.sample_rate = self.model.sample_rate | |
| self.io_channels = self.model.channels | |
| self.scale = scale | |
| #self.encoded_channels = self.model.latent_dim | |
| self.num_quantizers = self.model.num_codebooks | |
| self.codebook_size = self.model.cardinality | |
| self.model.to(torch.float16).eval().requires_grad_(False) | |
| def encode(self, x): | |
| assert False, "Audiocraft compression models do not support continuous encoding" | |
| # latents = self.model.encoder(x) | |
| # if self.quantize_on_decode: | |
| # output = latents | |
| # else: | |
| # z, _, _, _, _ = self.model.quantizer(latents, n_quantizers=self.model.n_codebooks) | |
| # output = z | |
| # if self.scale != 1.0: | |
| # output = output / self.scale | |
| # return output | |
| def decode(self, z): | |
| assert False, "Audiocraft compression models do not support continuous decoding" | |
| # if self.scale != 1.0: | |
| # z = z * self.scale | |
| # if self.quantize_on_decode: | |
| # z, _, _, _, _ = self.model.quantizer(z, n_quantizers=self.model.n_codebooks) | |
| # return self.model.decode(z) | |
| def tokenize(self, x): | |
| with torch.cuda.amp.autocast(enabled=False): | |
| return self.model.encode(x.to(torch.float16))[0] | |
| def decode_tokens(self, tokens): | |
| with torch.cuda.amp.autocast(enabled=False): | |
| return self.model.decode(tokens) | |