Merge branch 'nar' of https://github.com/Stability-AI/stable-audio-tools-dev into nar

Author: julian-parker

stable_audio_tools/models/transformer.py

@@ -8,7 +8,6 @@
 from torch import nn, einsum
 from torch.cuda.amp import autocast
 from typing import Callable, Literal
-from soft_moe_pytorch import SoftMoE
 
 try:
     from flash_attn import flash_attn_func, flash_attn_kvpacked_func
@@ -578,7 +577,6 @@ def __init__(
             conformer = False,
             layer_ix = -1,
             remove_norms = False,
-            number_of_experts = 1,
             attn_kwargs = {},
             ff_kwargs = {},
             norm_kwargs = {}
@@ -613,10 +611,7 @@ def __init__(
             )
 
         self.ff_norm = LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity()
-        if number_of_experts > 1:
-            self.ff = SoftMoE(dim = dim, seq_len = 1500, num_experts= number_of_experts, geglu = True)
-        else:
-            self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, **ff_kwargs)
+        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, **ff_kwargs)
 
         self.layer_ix = layer_ix
 
@@ -700,7 +695,6 @@ def __init__(
         use_sinusoidal_emb=False,
         use_abs_pos_emb=False,
         abs_pos_emb_max_length=10000,
-        number_of_experts = 1,
         **kwargs
         ):
 
@@ -739,7 +733,6 @@ def __init__(
                     zero_init_branch_outputs = zero_init_branch_outputs,
                     conformer=conformer,
                     layer_ix=i,
-                    number_of_experts=number_of_experts,
                     **kwargs
                 )
             )
@@ -787,4 +780,4 @@ def forward(
 
         x = self.project_out(x)
 
-        return x
+        return x