move beartype onto methods

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',