Fix compatibility with transformers 4.41+ and remove diffusers dependency

audio_tokenizer/modeling_audio_vae.py

@@ -1,5 +1,5 @@
-from diffusers.models.autoencoders.autoencoder_oobleck import OobleckDiagonalGaussianDistribution
 from transformers import PreTrainedModel
+from .oobleck_distribution import OobleckDiagonalGaussianDistribution
 import torch
 import torch.nn as nn
 

audio_tokenizer/oobleck_distribution.py

@@ -0,0 +1,65 @@
+# Extracted from diffusers.models.autoencoders.autoencoder_oobleck
+# To avoid dependency issues with diffusers and transformers version compatibility
+
+from typing import Optional
+import torch
+import torch.nn as nn
+
+
+def randn_tensor(
+    shape,
+    generator: Optional[torch.Generator] = None,
+    device: Optional[torch.device] = None,
+    dtype: Optional[torch.dtype] = None,
+):
+    """A helper function to create random tensors with the desired `generator`."""
+    if generator is None:
+        return torch.randn(shape, device=device, dtype=dtype)
+
+    # Handle generator properly
+    if device is not None and device.type == "cpu":
+        return torch.randn(shape, generator=generator, device=device, dtype=dtype)
+    else:
+        # For CUDA, we need to be careful with generators
+        latents = torch.randn(shape, generator=generator, device="cpu", dtype=dtype)
+        return latents.to(device)
+
+
+class OobleckDiagonalGaussianDistribution(object):
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        self.parameters = parameters
+        self.mean, self.scale = parameters.chunk(2, dim=1)
+        self.std = nn.functional.softplus(self.scale) + 1e-4
+        self.var = self.std * self.std
+        self.logvar = torch.log(self.var)
+        self.deterministic = deterministic
+
+    def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = randn_tensor(
+            self.mean.shape,
+            generator=generator,
+            device=self.parameters.device,
+            dtype=self.parameters.dtype,
+        )
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(self, other: "OobleckDiagonalGaussianDistribution" = None) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            if other is None:
+                return (self.mean * self.mean + self.var - self.logvar - 1.0).sum(1).mean()
+            else:
+                normalized_diff = torch.pow(self.mean - other.mean, 2) / other.var
+                var_ratio = self.var / other.var
+                logvar_diff = self.logvar - other.logvar
+
+                kl = normalized_diff + var_ratio + logvar_diff - 1
+
+                kl = kl.sum(1).mean()
+                return kl
+
+    def mode(self) -> torch.Tensor:
+        return self.mean

modeling_bailing_moe.py

@@ -662,7 +662,7 @@ def forward(
                     "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
                     "with a layer index."
                 )
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            kv_seq_len += past_key_value.get_seq_length()
 
         if self.config.rope_scaling is not None and self.config.rope_scaling["type"] == "3D":
             cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
@@ -771,7 +771,7 @@ def forward(
 
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            kv_seq_len += past_key_value.get_seq_length()
 
         if self.config.rope_scaling is not None and self.config.rope_scaling["type"] == "3D":
             cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
@@ -980,7 +980,7 @@ def forward(
 
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
-            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            kv_seq_len += past_key_value.get_seq_length()
 
         if self.config.rope_scaling is not None and self.config.rope_scaling["type"] == "3D":
             cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
@@ -1327,7 +1327,7 @@ def forward(
             use_legacy_cache = not isinstance(past_key_values, Cache)
             if use_legacy_cache:
                 past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-            past_key_values_length = past_key_values.get_usable_length(seq_length)
+            past_key_values_length = past_key_values.get_seq_length()
 
         if position_ids is None:
             device = input_ids.device if input_ids is not None else inputs_embeds.device
@@ -1600,7 +1600,7 @@ def prepare_inputs_for_generation(
         if past_key_values is not None:
             if isinstance(past_key_values, Cache):
                 cache_length = past_key_values.get_seq_length()
-                past_length = past_key_values.seen_tokens
+                past_length = cache_length  # Use cache_length instead of deprecated seen_tokens
                 max_cache_length = (
                     past_key_values.get_max_length()
                     if hasattr(past_key_values, "get_max_length")