Upload DMAE1d

config.json

@@ -1,4 +1,7 @@
 {
+  "architectures": [
+    "DMAE1d"
+  ],
   "attentions": [
     0,
     0,
@@ -9,7 +12,8 @@
     0
   ],
   "auto_map": {
-    "AutoConfig": "dmae_config.DMAE1dConfig"
+    "AutoConfig": "dmae_config.DMAE1dConfig",
+    "AutoModel": "dmae.DMAE1d"
   },
   "bottleneck": "tanh",
   "channels": 512,
@@ -73,5 +77,6 @@
   "stft_hop_length": 256,
   "stft_num_fft": 1023,
   "stft_use_complex": true,
+  "torch_dtype": "float32",
   "transformers_version": "4.24.0"
 }

dmae.py

@@ -0,0 +1,55 @@
+from transformers import PreTrainedModel
+from audio_encoders_pytorch import TanhBottleneck
+from audio_diffusion_pytorch import UniformDistribution, LinearSchedule, VSampler, DiffusionMAE1d
+from .dmae_config import DMAE1dConfig
+
+bottleneck = { 'tanh': TanhBottleneck }
+
+class DMAE1d(PreTrainedModel):
+
+    config_class = DMAE1dConfig
+
+    def __init__(self, config: DMAE1dConfig):
+        super().__init__(config)
+
+        self.model = DiffusionMAE1d(
+            in_channels = config.in_channels,
+            channels = config.channels, 
+            multipliers = config.multipliers,
+            factors = config.factors, 
+            num_blocks = config.num_blocks, 
+            attentions = config.attentions, 
+            encoder_inject_depth = config.encoder_inject_depth,
+            encoder_channels = config.encoder_channels, 
+            encoder_factors = config.encoder_factors, 
+            encoder_multipliers = config.encoder_multipliers,
+            encoder_num_blocks = config.encoder_num_blocks, 
+            bottleneck = bottleneck[config.bottleneck](),
+            stft_use_complex = config.stft_use_complex, 
+            stft_num_fft = config.stft_num_fft,
+            stft_hop_length = config.stft_hop_length,
+            diffusion_type = 'v',
+            diffusion_sigma_distribution = UniformDistribution(),
+            resnet_groups=8,
+            kernel_multiplier_downsample=2,
+            use_nearest_upsample=False,
+            use_skip_scale=True,
+            use_context_time=True,
+            patch_factor=1,
+            patch_blocks=1,
+        )
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def encode(self, *args, **kwargs):
+        return self.model.encode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        default_kwargs = dict(
+            sigma_schedule=LinearSchedule(), 
+            sampler=VSampler(), 
+            clamp=True
+        )
+        return self.model.decode(*args, **{**default_kwargs, **kwargs})
+

pytorch_model.bin

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