feat: add with_info during forward in autoencoder

Author: flavioschneider

audio_diffusion_pytorch/model.py

@@ -191,10 +191,17 @@ def __init__(
             num_groups=resnet_groups,
         )
 
-    def forward(self, x: Tensor, **kwargs) -> Tensor:
-        latent = self.encode(x)
+    def forward(  # type: ignore
+        self, x: Tensor, with_info: bool = False, **kwargs
+    ) -> Union[Tensor, Tuple[Tensor, Any]]:
+        if with_info:
+            latent, info = self.encode(x, with_info=True)
+        else:
+            latent = self.encode(x)
+
         context = self.to_context(latent)
-        return self.diffusion(x, context=[context], **kwargs)
+        loss = self.diffusion(x, context=[context], **kwargs)
+        return (loss, info) if with_info else loss
 
     def encode(
         self, x: Tensor, with_info: bool = False

setup.py

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