move beartype onto methods

Author: lucidrains

musiclm_pytorch/musiclm_pytorch.py

@@ -503,8 +503,8 @@ def forward(
 
 # text transformer
 
-@beartype
 class TextTransformer(nn.Module):
+    @beartype
     def __init__(
         self,
         dim,
@@ -546,6 +546,7 @@ def __init__(
     def device(self):
         return next(self.parameters()).device
 
+    @beartype
     def forward(
         self,
         x = None,
@@ -648,8 +649,8 @@ def forward(self, *, audio_layers, text_layers):
 
 # main classes
 
-@beartype
 class MuLaN(nn.Module):
+    @beartype
     def __init__(
         self,
         audio_transformer: AudioSpectrogramTransformer,
@@ -705,6 +706,7 @@ def get_audio_latents(
 
         return out, audio_layers
 
+    @beartype
     def get_text_latents(
         self,
         texts = None,
@@ -720,6 +722,7 @@ def get_text_latents(
 
         return out, text_layers
 
+    @beartype
     def forward(
         self,
         wavs,
@@ -766,8 +769,8 @@ def forward(
 
 # music lm
 
-@beartype
 class MuLaNEmbedQuantizer(AudioConditionerBase):
+    @beartype
     def __init__(
         self,
         mulan: MuLaN,
@@ -851,8 +854,8 @@ def forward(
         cond_embeddings = cond_embeddings.gather(2, indices)
         return rearrange(cond_embeddings, 'b q 1 d -> b q d')
 
-@beartype
 class MusicLM(nn.Module):
+    @beartype
     def __init__(
         self,
         audio_lm: AudioLM,

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'musiclm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.2.1',
+  version = '0.2.2',
   license='MIT',
   description = 'MusicLM - AudioLM + Audio CLIP to text to music synthesis',
   author = 'Phil Wang',