init

2025-02-24 11:56:50 +08:00
parent 32817bed27
commit ec5da63f6a
2 changed files with 296 additions and 0 deletions
@@ -0,0 +1,84 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch.nn as nn
+
+from .modules.decoder import Decoder
+from .modules.encoder import Encoder
+from .modules.projector import Projector
+from .modules.quantizer import Quantizer
+
+
+class RepCodec(nn.Module):
+    def __init__(
+            self,
+            input_channels=768,
+            output_channels=768,
+            encode_channels=768,
+            decode_channels=768,
+            code_dim=768,
+            codebook_num=1,
+            codebook_size=1024,
+            bias=True,
+            enc_ratios=(1, 1),
+            dec_ratios=(1, 1),
+            enc_strides=(1, 1),
+            dec_strides=(1, 1),
+            enc_kernel_size=3,
+            dec_kernel_size=3,
+            enc_block_dilations=(1, 1),
+            enc_block_kernel_size=3,
+            dec_block_dilations=(1, 1),
+            dec_block_kernel_size=3
+    ):
+        super().__init__()
+
+        self.input_channels = input_channels
+
+        self.encoder = Encoder(
+            input_channels=input_channels,
+            encode_channels=encode_channels,
+            channel_ratios=enc_ratios,
+            strides=enc_strides,
+            kernel_size=enc_kernel_size,
+            bias=bias,
+            block_dilations=enc_block_dilations,
+            unit_kernel_size=enc_block_kernel_size
+        )
+
+        self.decoder = Decoder(
+            code_dim=code_dim,
+            output_channels=output_channels,
+            decode_channels=decode_channels,
+            channel_ratios=dec_ratios,
+            strides=dec_strides,
+            kernel_size=dec_kernel_size,
+            bias=bias,
+            block_dilations=dec_block_dilations,
+            unit_kernel_size=dec_block_kernel_size
+        )
+
+        self.projector = Projector(
+            input_channels=self.encoder.out_channels,
+            code_dim=code_dim,
+            kernel_size=3,
+            stride=1,
+            bias=False
+        )
+
+        self.quantizer = Quantizer(
+            code_dim=code_dim,
+            codebook_num=codebook_num,
+            codebook_size=codebook_size
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        z = self.projector(x)
+        zq, vqloss, perplexity = self.quantizer(z)
+        y = self.decoder(zq)
+        return y, zq, z, vqloss, perplexity