Upload DMAE1d

config.json

@@ -1,7 +1,12 @@
 {
+  "architectures": [
+    "DMAE1d"
+  ],
   "auto_map": {
-    "AutoConfig": "dmae_config.DMAE1dConfig"
+    "AutoConfig": "dmae_config.DMAE1dConfig",
+    "AutoModel": "dmae.DMAE1d"
   },
   "model_type": "archinetai/dmae1d-ATC64-v2",
+  "torch_dtype": "float32",
   "transformers_version": "4.24.0"
 }

dmae.py

@@ -0,0 +1,52 @@
+import torch 
+from transformers import PreTrainedModel
+from .dmae_config import DMAE1dConfig
+from audio_encoders_pytorch import ME1d, TanhBottleneck # pip install audio_encoders_pytorch==0.0.20
+from audio_diffusion_pytorch.unets import UNetV0, LTPlugin # pip install -U git+https://github.com/archinetai/audio-diffusion-pytorch.git@nightly   # v0.0.2 
+from audio_diffusion_pytorch.models import DiffusionAE 
+
+class DMAE1d(PreTrainedModel):
+
+    config_class = DMAE1dConfig
+
+    def __init__(self, config: DMAE1dConfig):
+        super().__init__(config)
+
+        UNet = LTPlugin(
+            UNetV0,
+            num_filters=128,
+            window_length=64,
+            stride=64,
+        )
+
+        self.model = DiffusionAE(
+            net_t=UNet, 
+            dim=1, 
+            in_channels=2,
+            channels=[256, 512, 512, 512, 1024, 1024, 1024],
+            factors=[1, 2, 2, 2, 2, 2, 2],
+            items=[1, 2, 2, 2, 2, 2, 2],
+            encoder=ME1d(
+                in_channels=2,
+                channels=512,
+                multipliers=[1, 1, 1],
+                factors=[2, 2],
+                num_blocks=[4, 8],
+                stft_num_fft=1023,
+                stft_hop_length=256, 
+                out_channels=32,
+                bottleneck=TanhBottleneck()
+            ),
+            inject_depth=4
+        )
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def encode(self, *args, **kwargs):
+        return self.model.encode(*args, **kwargs)
+    
+    @torch.no_grad()
+    def decode(self, *args, **kwargs):
+        return self.model.decode(*args, **kwargs)
+

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9ad49fb4b5ba60c7db2774eebee21590731c9b2d423efff47d8e57119982f20
+size 740732261