Modules/diffusion/audio_diffusion_pytorch/models.py

from abc import ABC, abstractmethod
from math import floor
from typing import Any, Callable, Optional, Sequence, Tuple, Union

import torch
from einops import pack, rearrange, unpack
from torch import Generator, Tensor, nn

from .components import AppendChannelsPlugin, MelSpectrogram
from .diffusion import ARVDiffusion, ARVSampler, VDiffusion, VSampler
from .utils import (
    closest_power_2,
    default,
    downsample,
    exists,
    groupby,
    randn_like,
    upsample,
)


class DiffusionModel(nn.Module):
    def __init__(
        self,
        net_t: Callable,
        diffusion_t: Callable = VDiffusion,
        sampler_t: Callable = VSampler,
        loss_fn: Callable = torch.nn.functional.mse_loss,
        dim: int = 1,
        **kwargs,
    ):
        super().__init__()
        diffusion_kwargs, kwargs = groupby("diffusion_", kwargs)
        sampler_kwargs, kwargs = groupby("sampler_", kwargs)

        self.net = net_t(dim=dim, **kwargs)
        self.diffusion = diffusion_t(net=self.net, loss_fn=loss_fn, **diffusion_kwargs)
        self.sampler = sampler_t(net=self.net, **sampler_kwargs)

    def forward(self, *args, **kwargs) -> Tensor:
        return self.diffusion(*args, **kwargs)

    @torch.no_grad()
    def sample(self, *args, **kwargs) -> Tensor:
        return self.sampler(*args, **kwargs)


class EncoderBase(nn.Module, ABC):
    """Abstract class for DiffusionAE encoder"""

    @abstractmethod
    def __init__(self):
        super().__init__()
        self.out_channels = None
        self.downsample_factor = None


class AdapterBase(nn.Module, ABC):
    """Abstract class for DiffusionAE encoder"""

    @abstractmethod
    def encode(self, x: Tensor) -> Tensor:
        pass

    @abstractmethod
    def decode(self, x: Tensor) -> Tensor:
        pass


class DiffusionAE(DiffusionModel):
    """Diffusion Auto Encoder"""

    def __init__(
        self,
        in_channels: int,
        channels: Sequence[int],
        encoder: EncoderBase,
        inject_depth: int,
        latent_factor: Optional[int] = None,
        adapter: Optional[AdapterBase] = None,
        **kwargs,
    ):
        context_channels = [0] * len(channels)
        context_channels[inject_depth] = encoder.out_channels
        super().__init__(
            in_channels=in_channels,
            channels=channels,
            context_channels=context_channels,
            **kwargs,
        )
        self.in_channels = in_channels
        self.encoder = encoder
        self.inject_depth = inject_depth
        # Optional custom latent factor and adapter
        self.latent_factor = default(latent_factor, self.encoder.downsample_factor)
        self.adapter = adapter.requires_grad_(False) if exists(adapter) else None

    def forward(  # type: ignore
        self, x: Tensor, with_info: bool = False, **kwargs
    ) -> Union[Tensor, Tuple[Tensor, Any]]:
        # Encode input to latent channels
        latent, info = self.encode(x, with_info=True)
        channels = [None] * self.inject_depth + [latent]
        # Adapt input to diffusion if adapter provided
        x = self.adapter.encode(x) if exists(self.adapter) else x
        # Compute diffusion loss
        loss = super().forward(x, channels=channels, **kwargs)
        return (loss, info) if with_info else loss

    def encode(self, *args, **kwargs):
        return self.encoder(*args, **kwargs)

    @torch.no_grad()
    def decode(
        self, latent: Tensor, generator: Optional[Generator] = None, **kwargs
    ) -> Tensor:
        b = latent.shape[0]
        noise_length = closest_power_2(latent.shape[2] * self.latent_factor)
        # Compute noise by inferring shape from latent length
        noise = torch.randn(
            (b, self.in_channels, noise_length),
            device=latent.device,
            dtype=latent.dtype,
            generator=generator,
        )
        # Compute context from latent
        channels = [None] * self.inject_depth + [latent]  # type: ignore
        # Decode by sampling while conditioning on latent channels
        out = super().sample(noise, channels=channels, **kwargs)
        # Decode output with adapter if provided
        return self.adapter.decode(out) if exists(self.adapter) else out


class DiffusionUpsampler(DiffusionModel):
    def __init__(
        self,
        in_channels: int,
        upsample_factor: int,
        net_t: Callable,
        **kwargs,
    ):
        self.upsample_factor = upsample_factor
        super().__init__(
            net_t=AppendChannelsPlugin(net_t, channels=in_channels),
            in_channels=in_channels,
            **kwargs,
        )

    def reupsample(self, x: Tensor) -> Tensor:
        x = x.clone()
        x = downsample(x, factor=self.upsample_factor)
        x = upsample(x, factor=self.upsample_factor)
        return x

    def forward(self, x: Tensor, *args, **kwargs) -> Tensor:  # type: ignore
        reupsampled = self.reupsample(x)
        return super().forward(x, *args, append_channels=reupsampled, **kwargs)

    @torch.no_grad()
    def sample(  # type: ignore
        self, downsampled: Tensor, generator: Optional[Generator] = None, **kwargs
    ) -> Tensor:
        reupsampled = upsample(downsampled, factor=self.upsample_factor)
        noise = randn_like(reupsampled, generator=generator)
        return super().sample(noise, append_channels=reupsampled, **kwargs)


class DiffusionVocoder(DiffusionModel):
    def __init__(
        self,
        net_t: Callable,
        mel_channels: int,
        mel_n_fft: int,
        mel_hop_length: Optional[int] = None,
        mel_win_length: Optional[int] = None,
        in_channels: int = 1,  # Ignored: channels are automatically batched.
        **kwargs,
    ):
        mel_hop_length = default(mel_hop_length, floor(mel_n_fft) // 4)
        mel_win_length = default(mel_win_length, mel_n_fft)
        mel_kwargs, kwargs = groupby("mel_", kwargs)
        super().__init__(
            net_t=AppendChannelsPlugin(net_t, channels=1),
            in_channels=1,
            **kwargs,
        )
        self.to_spectrogram = MelSpectrogram(
            n_fft=mel_n_fft,
            hop_length=mel_hop_length,
            win_length=mel_win_length,
            n_mel_channels=mel_channels,
            **mel_kwargs,
        )
        self.to_flat = nn.ConvTranspose1d(
            in_channels=mel_channels,
            out_channels=1,
            kernel_size=mel_win_length,
            stride=mel_hop_length,
            padding=(mel_win_length - mel_hop_length) // 2,
            bias=False,
        )

    def forward(self, x: Tensor, *args, **kwargs) -> Tensor:  # type: ignore
        # Get spectrogram, pack channels and flatten
        spectrogram = rearrange(self.to_spectrogram(x), "b c f l -> (b c) f l")
        spectrogram_flat = self.to_flat(spectrogram)
        # Pack wave channels
        x = rearrange(x, "b c t -> (b c) 1 t")
        return super().forward(x, *args, append_channels=spectrogram_flat, **kwargs)

    @torch.no_grad()
    def sample(  # type: ignore
        self, spectrogram: Tensor, generator: Optional[Generator] = None, **kwargs
    ) -> Tensor:  # type: ignore
        # Pack channels and flatten spectrogram
        spectrogram, ps = pack([spectrogram], "* f l")
        spectrogram_flat = self.to_flat(spectrogram)
        # Get start noise and sample
        noise = randn_like(spectrogram_flat, generator=generator)
        waveform = super().sample(noise, append_channels=spectrogram_flat, **kwargs)
        # Unpack wave channels
        waveform = rearrange(waveform, "... 1 t -> ... t")
        waveform = unpack(waveform, ps, "* t")[0]
        return waveform


class DiffusionAR(DiffusionModel):
    def __init__(
        self,
        in_channels: int,
        length: int,
        num_splits: int,
        diffusion_t: Callable = ARVDiffusion,
        sampler_t: Callable = ARVSampler,
        **kwargs,
    ):
        super().__init__(
            in_channels=in_channels + 1,
            out_channels=in_channels,
            diffusion_t=diffusion_t,
            diffusion_length=length,
            diffusion_num_splits=num_splits,
            sampler_t=sampler_t,
            sampler_in_channels=in_channels,
            sampler_length=length,
            sampler_num_splits=num_splits,
            use_time_conditioning=False,
            use_modulation=False,
            **kwargs,
        )