Upload DMAE1d

config.json

@@ -1,7 +1,12 @@
 {
+  "architectures": [
+    "DMAE1d"
+  ],
   "auto_map": {
-    "AutoConfig": "config.DMAE1dConfig"
+    "AutoConfig": "config.DMAE1dConfig",
+    "AutoModel": "model.DMAE1d"
   },
   "model_type": "archinetai/dmae1d-ATC32-v3",
+  "torch_dtype": "float32",
   "transformers_version": "4.24.0"
 }

model.py

@@ -0,0 +1,56 @@
+import torch 
+from torch import Tensor
+from transformers import PreTrainedModel
+from audio_encoders_pytorch import MelE1d, TanhBottleneck
+from audio_diffusion_pytorch import DiffusionAE, UNetV0, LTPlugin, VDiffusion, VSampler
+from .config import DMAE1dConfig 
+
+
+class DMAE1d(PreTrainedModel):
+
+    config_class = DMAE1dConfig
+
+    def __init__(self, config: DMAE1dConfig):
+        super().__init__(config)
+
+        UNet = LTPlugin(
+            UNetV0,
+            num_filters=128,
+            window_length=64,
+            stride=32,
+        )
+
+        self.model = DiffusionAE(
+            net_t=UNet, 
+            dim=1, 
+            in_channels=2,
+            channels=[256, 512, 512, 512, 512],
+            factors=[1, 2, 2, 2, 2],
+            items=[1, 2, 2, 2, 4],
+            inject_depth=4,
+            encoder=MelE1d(
+                in_channels=2,
+                channels=512,
+                multipliers=[1, 1],
+                factors=[2],
+                num_blocks=[12],
+                out_channels=32,
+                mel_channels=80,
+                mel_sample_rate=48000,
+                mel_normalize_log=True,
+                bottleneck=TanhBottleneck()
+            ),
+            diffusion_t=VDiffusion, 
+            sampler_t=VSampler
+        )
+    
+    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:302c66edf614e12da3004192f8517f0a9883bcd36e937ab471b3123ae66fdca6
+size 345503551