Upload folder using huggingface_hub

README.md

@@ -1,4 +1 @@
----
-license: mit
----
-Soundstream implementation. Sample rate is 16000.
+Soundstream implementation. Sample rate 16000.

config.json

@@ -1,6 +1,11 @@
 {
+  "model_type": "soundstream",
   "channels": 32,
-  "codebook_size": 1024,
   "latent_dim": 512,
-  "num_quantizers": 8
-}
+  "codebook_size": 1024,
+  "num_quantizers": 8,
+  "auto_map": {
+    "AutoConfig": "config_model.SoundStreamConfig",
+    "AutoModel": "model.SoundStreamCodec"
+  }
+}

config_model.py

@@ -0,0 +1,19 @@
+from transformers import PretrainedConfig
+
+
+class SoundStreamConfig(PretrainedConfig):
+    model_type = "soundstream"
+
+    def __init__(
+        self,
+        channels: int = 32,
+        latent_dim: int = 512,
+        codebook_size: int = 1024,
+        num_quantizers: int = 8,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.latent_dim = latent_dim
+        self.codebook_size = codebook_size
+        self.num_quantizers = num_quantizers

model.py

@@ -1,7 +1,9 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from huggingface_hub import PyTorchModelHubMixin
+from transformers import PreTrainedModel
+
+from .config_model import SoundStreamConfig
 
 
 class CausalConv1d(nn.Module):
@@ -102,7 +104,6 @@ class Encoder(nn.Module):
 
     def __init__(self, channels: int = 16, dim: int = 512):
         super().__init__()
-        # NB: attribute name "encoder" matches training checkpoint keys
         self.encoder = nn.Sequential(
             CausalConv1d(kernel_size=7, in_channels=1, out_channels=channels),
             EncoderBlock(channels=2 * channels, s=2),
@@ -120,7 +121,6 @@ class Decoder(nn.Module):
 
     def __init__(self, channels: int = 16, dim: int = 512):
         super().__init__()
-        # NB: attribute name "decoder" matches training checkpoint keys
         self.decoder = nn.Sequential(
             CausalConv1d(kernel_size=7, in_channels=dim, out_channels=16 * channels),
             DecoderBlock(channels=16 * channels, s=5),
@@ -325,32 +325,23 @@ class ResidualVectorQuantizer(nn.Module):
         return quantized
 
 
-class SoundStreamCodec(
-    nn.Module,
-    PyTorchModelHubMixin,
-    library_name="soundstream-impl",
-    license="mit",
-):
-    def __init__(
-        self,
-        channels: int = 32,
-        latent_dim: int = 512,
-        codebook_size: int = 1024,
-        num_quantizers: int = 8,
-    ):
-        super().__init__()
+class SoundStreamCodec(PreTrainedModel):
+    config_class = SoundStreamConfig
+
+    def __init__(self, config):
+        super().__init__(config)
         self.strides = (2, 4, 5, 5)
         self.downsampling_factor = 1
         for s in self.strides:
             self.downsampling_factor *= s
 
-        self.encoder = Encoder(channels=channels, dim=latent_dim)
+        self.encoder = Encoder(channels=config.channels, dim=config.latent_dim)
         self.quantizer = ResidualVectorQuantizer(
-            latent_dim=latent_dim,
-            codebook_size=codebook_size,
-            num_quantizers=num_quantizers,
+            latent_dim=config.latent_dim,
+            codebook_size=config.codebook_size,
+            num_quantizers=config.num_quantizers,
         )
-        self.decoder = Decoder(channels=channels, dim=latent_dim)
+        self.decoder = Decoder(channels=config.channels, dim=config.latent_dim)
 
     def forward(self, audio, **kwargs):
         original_length = audio.size(-1)