feat: update diffusion autoencoder to multiencoder conditioning

Author: flavioschneider

.pre-commit-config.yaml

@@ -30,6 +30,7 @@ repos:
         args: [
             '--per-file-ignores=__init__.py:F401',
             '--max-line-length=88',
+            '--ignore=E203,W503'
         ]
 
 # Checks types

audio_diffusion_pytorch/__init__.py

@@ -21,4 +21,4 @@
     DiffusionUpsampler1d,
     Model1d,
 )
-from .modules import Encoder1d, UNet1d, UNetConditional1d
+from .modules import MultiEncoder1d, UNet1d, UNetConditional1d

audio_diffusion_pytorch/model.py

@@ -14,8 +14,8 @@
     Sampler,
     Schedule,
 )
-from .modules import Encoder1d, ResnetBlock1d, UNet1d, UNetConditional1d
-from .utils import default, exists, prod, to_list
+from .modules import MultiEncoder1d, UNet1d, UNetConditional1d
+from .utils import default, exists, to_list
 
 """
 Diffusion Classes (generic for 1d data)
@@ -117,51 +117,45 @@ def __init__(
         kernel_multiplier_downsample: int,
         encoder_depth: int,
         encoder_channels: int,
-        context_channels: int,
         bottleneck: Optional[Bottleneck] = None,
         encoder_num_blocks: Optional[Sequence[int]] = None,
         **kwargs
     ):
-        super().__init__(
+        self.in_channels = in_channels
+        encoder_num_blocks = default(encoder_num_blocks, num_blocks)
+        assert_message = "The number of encoder_num_blocks must match encoder_depth"
+        assert len(encoder_num_blocks) >= encoder_depth, assert_message
+
+        multiencoder = MultiEncoder1d(
             in_channels=in_channels,
             channels=channels,
             patch_size=patch_size,
-            kernel_sizes_init=kernel_sizes_init,
+            num_layers=encoder_depth,
+            latent_channels=encoder_channels,
             multipliers=multipliers,
             factors=factors,
-            num_blocks=num_blocks,
-            resnet_groups=resnet_groups,
+            num_blocks=encoder_num_blocks,
+            kernel_sizes_init=kernel_sizes_init,
             kernel_multiplier_downsample=kernel_multiplier_downsample,
-            context_channels=[0] * encoder_depth + [context_channels],
-            **kwargs,
+            resnet_groups=resnet_groups,
         )
 
-        self.in_channels = in_channels
-        self.encoder_factor = patch_size * prod(factors[0:encoder_depth])
-        self.bottleneck = bottleneck
-
-        encoder_num_blocks = default(encoder_num_blocks, num_blocks)
-        assert_message = "The number of encoder_num_blocks must match encoder_depth"
-        assert len(encoder_num_blocks) >= encoder_depth, assert_message
-
-        self.encoder = Encoder1d(
+        super().__init__(
             in_channels=in_channels,
             channels=channels,
             patch_size=patch_size,
             kernel_sizes_init=kernel_sizes_init,
             multipliers=multipliers,
             factors=factors,
-            num_blocks=encoder_num_blocks,
+            num_blocks=num_blocks,
             resnet_groups=resnet_groups,
             kernel_multiplier_downsample=kernel_multiplier_downsample,
-            extract_channels=[0] * (encoder_depth - 1) + [encoder_channels],
+            context_channels=multiencoder.channels_list,
+            **kwargs,
         )
 
-        self.to_context_channels = ResnetBlock1d(
-            in_channels=encoder_channels,
-            out_channels=context_channels,
-            num_groups=resnet_groups,
-        )
+        self.bottleneck = bottleneck
+        self.multiencoder = multiencoder
 
     def forward(  # type: ignore
         self, x: Tensor, with_info: bool = False, **kwargs
@@ -171,28 +165,28 @@ def forward(  # type: ignore
         else:
             latent = self.encode(x)
 
-        channels = self.to_context_channels(latent)
-        loss = self.diffusion(x, channels_list=[channels], **kwargs)
+        channels_list = self.multiencoder.decode(latent)
+        loss = self.diffusion(x, channels_list=channels_list, **kwargs)
         return (loss, info) if with_info else loss
 
     def encode(
         self, x: Tensor, with_info: bool = False
     ) -> Union[Tensor, Tuple[Tensor, Any]]:
-        x = self.encoder(x)[-1]
-        latent = torch.tanh(x)
+        latent = self.multiencoder.encode(x)
+        latent = torch.tanh(latent)
         # Apply bottleneck if provided (e.g. quantization module)
         if exists(self.bottleneck):
             latent, info = self.bottleneck(latent)
             return (latent, info) if with_info else latent
         return latent
 
     def decode(self, latent: Tensor, **kwargs) -> Tensor:
-        b, length = latent.shape[0], latent.shape[2] * self.encoder_factor
+        b, length = latent.shape[0], latent.shape[2] * self.multiencoder.factor
         # Compute noise by inferring shape from latent length
         noise = torch.randn(b, self.in_channels, length).to(latent)
         # Compute context form latent
-        channels = self.to_context_channels(latent)
-        default_kwargs = dict(channels_list=[channels])
+        channels_list = self.multiencoder.decode(latent)
+        default_kwargs = dict(channels_list=channels_list)
         # Decode by sampling while conditioning on latent channels
         return super().sample(noise, **{**default_kwargs, **kwargs})  # type: ignore
 
@@ -257,10 +251,7 @@ def sample(self, *args, **kwargs):
 class AudioDiffusionAutoencoder(DiffusionAutoencoder1d):
     def __init__(self, *args, **kwargs):
         default_kwargs = dict(
-            **get_default_model_kwargs(),
-            encoder_depth=4,
-            encoder_channels=32,
-            context_channels=512,
+            **get_default_model_kwargs(), encoder_depth=4, encoder_channels=64
         )
         super().__init__(*args, **{**default_kwargs, **kwargs})  # type: ignore
 

audio_diffusion_pytorch/modules.py

@@ -9,7 +9,7 @@
 from einops_exts.torch import EinopsToAndFrom
 from torch import Tensor, einsum
 
-from .utils import default, exists
+from .utils import default, exists, prod
 
 """
 Convolutional Blocks
@@ -666,6 +666,7 @@ def __init__(
         use_skip: bool = False,
         skip_channels: int = 0,
         use_skip_scale: bool = False,
+        extract_channels: int = 0,
         use_attention: bool = False,
         attention_heads: Optional[int] = None,
         attention_features: Optional[int] = None,
@@ -675,12 +676,7 @@ def __init__(
     ):
         super().__init__()
 
-        assert (not use_attention) or (
-            exists(attention_heads)
-            and exists(attention_features)
-            and exists(attention_multiplier)
-        )
-
+        self.use_extract = extract_channels > 0
         self.use_pre_upsample = use_pre_upsample
         self.use_attention = use_attention
         self.use_skip = use_skip
@@ -723,6 +719,14 @@ def __init__(
             use_nearest=use_nearest,
         )
 
+        if self.use_extract:
+            num_extract_groups = min(num_groups, extract_channels)
+            self.to_extracted = ResnetBlock1d(
+                in_channels=out_channels,
+                out_channels=extract_channels,
+                num_groups=num_extract_groups,
+            )
+
     def add_skip(self, x: Tensor, skip: Tensor) -> Tensor:
         return torch.cat([x, skip * self.skip_scale], dim=1)
 
@@ -732,7 +736,7 @@ def forward(
         skips: Optional[List[Tensor]] = None,
         mapping: Optional[Tensor] = None,
         embedding: Optional[Tensor] = None,
-    ) -> Tensor:
+    ) -> Union[Tuple[Tensor, Tensor], Tensor]:
 
         if self.use_pre_upsample:
             x = self.upsample(x)
@@ -747,6 +751,10 @@ def forward(
         if not self.use_pre_upsample:
             x = self.upsample(x)
 
+        if self.use_extract:
+            extracted = self.to_extracted(x)
+            return x, extracted
+
         return x
 
 
@@ -1144,11 +1152,11 @@ def forward(  # type: ignore
 
 
 """
-Encoder
+Encoders / Decoders
 """
 
 
-class Encoder1d(nn.Module):
+class MultiEncoder1d(nn.Module):
     def __init__(
         self,
         in_channels: int,
@@ -1157,18 +1165,17 @@ def __init__(
         resnet_groups: int,
         kernel_multiplier_downsample: int,
         kernel_sizes_init: Sequence[int],
+        num_layers: int,
+        latent_channels: int,
         multipliers: Sequence[int],
         factors: Sequence[int],
         num_blocks: Sequence[int],
-        extract_channels: Sequence[int],
     ):
         super().__init__()
-
-        num_layers = len(extract_channels)
-        self.num_layers = num_layers
-
-        use_extract = [channels > 0 for channels in extract_channels]
-        self.use_extract = use_extract
+        self.factor = patch_size * prod(factors[0:num_layers])
+        self.channels_list = self.get_channels_list(
+            in_channels, channels, multipliers, num_layers
+        )
 
         assert (
             len(multipliers) >= num_layers + 1
@@ -1195,21 +1202,78 @@ def __init__(
                     kernel_multiplier=kernel_multiplier_downsample,
                     num_groups=resnet_groups,
                     num_layers=num_blocks[i],
-                    extract_channels=extract_channels[i],
                 )
                 for i in range(num_layers)
             ]
         )
 
-    def forward(self, x: Tensor) -> List[Tensor]:
-        x = self.to_in(x)
-        channels_list = []
+        pre_latent_channels = channels * multipliers[num_layers]
+
+        self.to_latent = ResnetBlock1d(
+            in_channels=pre_latent_channels,
+            out_channels=latent_channels,
+            num_groups=resnet_groups,
+        )
+
+        self.from_latent = ResnetBlock1d(
+            in_channels=latent_channels,
+            out_channels=pre_latent_channels,
+            num_groups=resnet_groups,
+        )
 
-        for downsample, use_extract in zip(self.downsamples, self.use_extract):
-            if use_extract:
-                x, channels = downsample(x)
-                channels_list += [channels]
-            else:
-                x = downsample(x)
+        self.upsamples = nn.ModuleList(
+            [
+                UpsampleBlock1d(
+                    in_channels=channels * multipliers[i + 1],
+                    out_channels=channels * multipliers[i],
+                    factor=factors[i],
+                    num_groups=resnet_groups,
+                    num_layers=num_blocks[i],
+                    use_nearest=False,
+                    use_skip=False,
+                    extract_channels=channels * multipliers[i],
+                )
+                for i in reversed(range(num_layers))
+            ]
+        )
 
+        self.to_out = nn.Sequential(
+            ResnetBlock1d(
+                in_channels=channels, out_channels=channels, num_groups=resnet_groups
+            ),
+            Conv1d(
+                in_channels=channels,
+                out_channels=in_channels * patch_size,
+                kernel_size=1,
+            ),
+            Rearrange("b (c p) l -> b c (l p)", p=patch_size),
+        )
+
+    def get_channels_list(
+        self,
+        in_channels: int,
+        channels: int,
+        multipliers: Sequence[int],
+        num_layers: int,
+    ) -> List[int]:
+        channels_list = [in_channels]
+        channels_list += [channels * m for m in multipliers[1 : num_layers + 1]]
         return channels_list
+
+    def encode(self, x: Tensor) -> Tensor:
+        x = self.to_in(x)
+        for downsample in self.downsamples:
+            x = downsample(x)
+        latent = self.to_latent(x)
+        return latent
+
+    def decode(self, latent: Tensor) -> List[Tensor]:
+        x = self.from_latent(latent)
+        channels_list = []
+        channels = x
+        for upsample in self.upsamples:
+            channels_list += [channels]
+            x, channels = upsample(x)
+        x = self.to_out(x)
+        channels_list += [x]
+        return channels_list[::-1]

setup.py

@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.35",
+    version="0.0.36",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",