feat: add mapping context features, improve output block

Author: flavioschneider

README.md

@@ -139,6 +139,7 @@ unet = UNet1d(
     kernel_multiplier_downsample=2,
     use_nearest_upsample=False,
     use_skip_scale=True,
+    use_context_time=True,
 )
 
 x = torch.randn(3, 1, 2 ** 16)
@@ -236,6 +237,7 @@ y_long = composer(y, keep_start=True) # [1, 1, 98304]
 - [x] Add autoencoder bottleneck option for quantization.
 - [x] Add option to provide context tokens (resnet cross attention).
 - [x] Add conditional model with classifier-free guidance.
+- [x] Add option to provide context features mapping.
 
 ## Appreciation
 

audio_diffusion_pytorch/model.py

@@ -219,6 +219,7 @@ def get_default_model_kwargs():
         kernel_multiplier_downsample=2,
         use_nearest_upsample=False,
         use_skip_scale=True,
+        use_context_time=True,
         diffusion_sigma_distribution=LogNormalDistribution(mean=-3.0, std=1.0),
         diffusion_sigma_data=0.1,
         diffusion_dynamic_threshold=0.0,