src/singer/vae/__init__.py

import typing as tp
from abc import ABC, abstractmethod

import torch
import torchaudio
from torch import nn


class AbstractVAE(ABC, nn.Module):
    @property
    @abstractmethod
    def frame_rate(self) -> float: ...

    @property
    @abstractmethod
    def orig_sample_rate(self) -> int: ...

    @property
    @abstractmethod
    def channel_dim(self) -> int: ...

    @property
    @abstractmethod
    def split_bands(self) -> int: ...

    @property
    @abstractmethod
    def input_channel(self) -> int: ...

    def encode(self, wav) -> torch.Tensor: ...

    def decode(self, latents) -> torch.Tensor: ...


from .autoencoders import AudioAutoencoder, create_autoencoder_from_config


class StableVAE(AbstractVAE):
    def __init__(self, vae_ckpt, vae_cfg, sr=48000) -> None:
        super().__init__()
        import json

        with open(vae_cfg) as f:
            config = json.load(f)
        self.vae: AudioAutoencoder = create_autoencoder_from_config(config)
        self.vae.load_state_dict(torch.load(vae_ckpt)["state_dict"])
        self.sample_rate = sr
        self.rsp48k = (
            torchaudio.transforms.Resample(sr, self.orig_sample_rate) if sr != self.orig_sample_rate else nn.Identity()
        )

    @torch.no_grad()
    def encode(self, wav: torch.Tensor, sample=True) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
        wav = self.rsp48k(wav)
        if wav.shape[-1] < 2048:
            return torch.zeros((wav.shape[0], self.channel_dim, 0), device=wav.device, dtype=wav.dtype)
        if wav.ndim == 2:
            wav = wav.unsqueeze(1)
        if wav.shape[1] == 1:
            wav = wav.repeat(1, self.vae.in_channels, 1)
        latent = self.vae.encode_audio(wav)  # B, 64, T
        return latent

    def decode(self, latents: torch.Tensor, **kwargs):
        # B, 64, T
        with torch.no_grad():
            audio_recon = self.vae.decode_audio(latents, **kwargs)

        return audio_recon

    @property
    def frame_rate(self) -> float:
        return float(self.vae.sample_rate) / self.vae.downsampling_ratio

    @property
    def orig_sample_rate(self) -> int:
        return self.vae.sample_rate

    @property
    def channel_dim(self) -> int:
        return self.vae.latent_dim

    @property
    def split_bands(self) -> int:
        return 1

    @property
    def input_channel(self) -> int:
        return self.vae.in_channels