Formatting.

Author: LaurentMazare

moshi_mlx/moshi_mlx/modules/conv.py

@@ -6,6 +6,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 
+
 class Conv1d(nn.Module):
     def __init__(
         self,
@@ -48,6 +49,7 @@ def __call__(self, xs: mx.array) -> mx.array:
             y = y + self.bias
         return y.swapaxes(-1, -2)
 
+
 class ConvTranspose1d(nn.Module):
     def __init__(
         self,
@@ -86,6 +88,7 @@ def __call__(self, xs: mx.array) -> mx.array:
             y = y + self.bias
         return y
 
+
 class NormConv1d(nn.Module):
     def __init__(
         self,
@@ -112,6 +115,7 @@ def __init__(
     def __call__(self, xs: mx.array) -> mx.array:
         return self.conv(xs)
 
+
 class NormConvTranspose1d(nn.Module):
     def __init__(
         self,
@@ -136,23 +140,27 @@ def __init__(
     def __call__(self, xs: mx.array) -> mx.array:
         return self.convtr(xs)
 
+
 def get_extra_padding_for_conv1d(
     xs: mx.array,
     ksize: int,
-    stride: int, 
+    stride: int,
     padding_total: int,
 ) -> int:
     l = xs.shape[-1]
     nframes = max(l + padding_total - ksize, 0) / stride + 1.0
     ideal_len = (int(math.ceil(nframes)) - 1) * stride + ksize - padding_total
     return max(0, ideal_len - l)
 
+
 def unpad1d(xs: mx.array, unpad_l: int, unpad_r: int) -> mx.array:
     left = unpad_l
     right = xs.shape[-1] - unpad_r
     return xs[..., left:right]
 
 # TODO(laurent): add a streaming module abstract class?
+
+
 class StreamableConv1d(nn.Module):
     def __init__(
         self,
@@ -218,6 +226,7 @@ def __call__(self, xs: mx.array) -> mx.array:
             self._prev_xs = xs
             return mx.zeros((b, self._out_channels, 0))
 
+
 class StreamableConvTranspose1d(nn.Module):
     def __init__(
         self,
@@ -259,9 +268,10 @@ def __call__(self, xs: mx.array) -> mx.array:
             ys1, ys2 = ys[..., :pt] + prev_ys, ys[..., pt:]
             ys = mx.concat([ys1, ys2], axis=-1)
         invalid_steps = self._ksize - stride
-        ys, self._prev_ys = ys[..., :ot-invalid_steps], ys[..., ot-invalid_steps]
+        ys, self._prev_ys = ys[..., :ot - invalid_steps], ys[..., ot - invalid_steps]
         return ys
 
+
 class ConvDownsample1d(nn.Module):
     def __init__(
         self,
@@ -272,7 +282,7 @@ def __init__(
         self.conv = StreamableConv1d(
             in_channels=dim,
             out_channels=dim,
-            ksize=2*stride,
+            ksize=2 * stride,
             stride=stride,
             dilation=1,
             groups=1,
@@ -287,6 +297,7 @@ def reset_state(self):
     def __call__(self, xs: mx.array) -> mx.array:
         return self.conv(xs)
 
+
 class ConvTrUpsample1d(nn.Module):
     def __init__(
         self,
@@ -297,9 +308,9 @@ def __init__(
         self.convtr = StreamableConvTranspose1d(
             in_channels=dim,
             out_channels=dim,
-            ksize=2*stride,
+            ksize=2 * stride,
             stride=stride,
-            groups=dim, # TODO: hopefully someday this will be fixed.
+            groups=dim,  # TODO: hopefully someday this will be fixed.
             bias=False,
             causal=causal,
         )

moshi_mlx/moshi_mlx/modules/quantization.py

@@ -7,6 +7,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 
+
 class EuclideanCodebook(nn.Module):
     def __init__(self, dim: int, codebook_size: int):
         super().__init__()
@@ -36,6 +37,7 @@ def decode(self, xs: mx.array) -> mx.array:
         target_shape = list(xs.shape) + [self._dim]
         return mx.take(self.embedding, xs.flatten()).reshape(target_shape)
 
+
 class VectorQuantization(nn.Module):
     def __init__(self, dim: int, codebook_size: int, codebook_dim: int | None):
         super().__init__()
@@ -60,6 +62,7 @@ def decode(self, xs: mx.array) -> mx.array:
             xs = self.project_out(xs)
         return xs.swapaxes(-1, -2)
 
+
 class ResidualVectorQuantization(nn.Module):
     def __init__(self, nq: int, dim: int, codebook_size: int, codebook_dim: int | None):
         super().__init__()
@@ -90,6 +93,7 @@ def decode(self, xs: mx.array) -> mx.array:
             quantized = quantized + self.layers[i].decode(xs[i])
         return quantized
 
+
 class ResidualVectorQuantizer(nn.Module):
     def __init__(
         self,
@@ -112,10 +116,10 @@ def __init__(
         else:
             self.output_proj = Conv1d(dim, output_dim, 1, bias=False)
         self.vq = ResidualVectorQuantization(
-                nq=nq,
-                dim=dim,
-                codebook_size=bins,
-                codebook_dim=None,
+            nq=nq,
+            dim=dim,
+            codebook_size=bins,
+            codebook_dim=None,
         )
 
     def encode(self, xs: mx.array) -> mx.array:
@@ -130,6 +134,7 @@ def decode(self, xs: mx.array) -> mx.array:
             quantized = self.output_proj(quantized)
         return quantized
 
+
 class SplitResidualVectorQuantizer(nn.Module):
     def __init__(
         self,
@@ -153,7 +158,7 @@ def __init__(
             dim=dim,
             input_dim=input_dim,
             output_dim=output_dim,
-            nq=nq-1,
+            nq=nq - 1,
             bins=bins,
             force_projection=True
         )

moshi_mlx/moshi_mlx/modules/seanet.py

@@ -25,6 +25,7 @@ class SeanetConfig:
     true_skip: bool
     compress: int
 
+
 class SeanetResnetBlock(nn.Module):
     def __init__(self, cfg: SeanetConfig, dim: int, ksizes_and_dilations: list):
         super().__init__()
@@ -79,6 +80,7 @@ def __call__(self, xs: mx.array) -> mx.array:
             xs = xs + self.shortcut(residual)
         return xs
 
+
 class EncoderLayer(nn.Module):
     def __init__(self, cfg: SeanetConfig, ratio: int, mult: int):
         super().__init__()
@@ -115,6 +117,7 @@ def __call__(self, xs: mx.array) -> mx.array:
             xs = r(xs)
         return self.downsample(nn.elu(xs, alpha=1.0))
 
+
 class SeanetEncoder(nn.Module):
     def __init__(self, cfg: SeanetConfig):
         super().__init__()
@@ -160,6 +163,7 @@ def __call__(self, xs: mx.array) -> mx.array:
         xs = nn.elu(xs, alpha=1.0)
         return self.final_conv1d(xs)
 
+
 class DecoderLayer(nn.Module):
     def __init__(self, cfg: SeanetConfig, ratio: int, mult: int):
         super().__init__()
@@ -185,6 +189,7 @@ def __call__(self, xs: mx.array) -> mx.array:
             xs = r(xs)
         return xs
 
+
 class SeanetDecoder(nn.Module):
     def __init__(self, cfg: SeanetConfig):
         super().__init__()
@@ -230,6 +235,7 @@ def __call__(self, xs: mx.array) -> mx.array:
         xs = nn.elu(xs, alpha=1.0)
         return self.final_conv1d(xs)
 
+
 class Seanet(nn.Module):
     def __init__(self, cfg: SeanetConfig):
         super().__init__()