Bump transformers version (#150)

* add mechanism to reconcile DACs

* add mimi compatibility

* correct batched generation

* update dac on the hub script

* update quantizes in training

* introduce fused lm heads

* update training code and modeling code with per codebook losses and smarter eval

* bump transformers version

* update transformers version

* update bump

* update transformers version

* update dependencies

---------

Co-authored-by: [email protected] <Yoach Lacombe>

Author: ylacombe

helpers/push_to_hub_scripts/push_dac_to_hub.py

@@ -5,7 +5,14 @@
 from transformers import AutoConfig, AutoModel
 from transformers import EncodecFeatureExtractor
 
-AutoConfig.register("dac", DACConfig)
+from importlib.metadata import version
+from packaging.version import Version
+
+if Version(version("transformers"))<= Version("4.44.2dev"):
+    AutoConfig.register("dac", DACConfig)
+else:
+    AutoConfig.register("dac_on_the_hub", DACConfig)
+
 AutoModel.register(DACConfig, DACModel)
 
 # Download a model

parler_tts/__init__.py

@@ -14,5 +14,12 @@
 
 from .streamer import ParlerTTSStreamer
 
-AutoConfig.register("dac", DACConfig)
+from importlib.metadata import version
+from packaging.version import Version
+
+if Version(version("transformers"))<= Version("4.44.2dev"):
+    AutoConfig.register("dac", DACConfig)
+else:
+    AutoConfig.register("dac_on_the_hub", DACConfig)
+
 AutoModel.register(DACConfig, DACModel)

parler_tts/configuration_parler_tts.py

@@ -17,6 +17,10 @@
 from transformers import AutoConfig, logging
 from transformers.configuration_utils import PretrainedConfig
 
+from importlib.metadata import version
+from packaging.version import Version
+
+use_dac_on_the_hub = Version(version("transformers")) > Version("4.44.2dev")
 
 logger = logging.get_logger(__name__)
 
@@ -91,6 +95,10 @@ class ParlerTTSDecoderConfig(PretrainedConfig):
             The base period of the RoPE embeddings.
         cross_attention_implementation_strategy (`str`, *optional*):
             If not specified, the cross-attention implementation will be the same as `_attn_implementation`. If `always_eager`, it will always be the eager implementation. If `always_sdpa`, it will always be the sdpa implementation.
+        use_fused_lm_heads(`bool`, *optional*, defaults to `False`):
+            Whether to fuse audio LM heads instead of applying them sequentially.
+        codebook_weights(`List[int]`, *optional*):
+            Weights applied to each codebook when computing the loss.
     """
 
     model_type = "parler_tts_decoder"
@@ -122,6 +130,8 @@ def __init__(
         rope_embeddings=False,
         rope_theta=10_000.0,
         cross_attention_implementation_strategy=None,
+        use_fused_lm_heads=False,
+        codebook_weights=None,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -148,7 +158,11 @@ def __init__(
         self.rope_embeddings = rope_embeddings
         self.rope_theta = rope_theta
         self.cross_attention_implementation_strategy = cross_attention_implementation_strategy
+        self.use_fused_lm_heads = use_fused_lm_heads
+        self.codebook_weights = codebook_weights
 
+        if codebook_weights is not None and len(codebook_weights) != num_codebooks:
+            raise ValueError(f"`codebook_weights` has length {len(codebook_weights)} when it should be of length {num_codebooks}.")
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,
@@ -234,6 +248,11 @@ def __init__(self, vocab_size=1024, prompt_cross_attention=False, **kwargs):
         audio_encoder_config = kwargs.pop("audio_encoder")
         audio_encoder_model_type = audio_encoder_config.pop("model_type")
 
+        model_version = kwargs.get("transformers_version", None)
+        if model_version is not None and Version(model_version) <= Version("4.44.2dev") and use_dac_on_the_hub and audio_encoder_model_type=="dac":
+            # here we have to manually change model type if DAC based on transformers version
+            audio_encoder_model_type = "dac_on_the_hub"
+
         decoder_config = kwargs.pop("decoder")
 
         self.vocab_size = vocab_size
@@ -269,22 +288,4 @@ def from_sub_models_config(
     @property
     # This is a property because you might want to change the codec model on the fly
     def sampling_rate(self):
-        return self.audio_encoder.sampling_rate
-
-    # Copy from musicgen
-    @property
-    def _attn_implementation(self):
-        # This property is made private for now (as it cannot be changed and a PreTrainedModel.use_attn_implementation method needs to be implemented.)
-        if hasattr(self, "_attn_implementation_internal"):
-            if self._attn_implementation_internal is None:
-                # `config.attn_implementation` should never be None, for backward compatibility.
-                return "eager"
-            else:
-                return self._attn_implementation_internal
-        else:
-            return "eager"
-
-    @_attn_implementation.setter
-    def _attn_implementation(self, value):
-        self._attn_implementation_internal = value
-        self.decoder._attn_implementation = value
+        return self.audio_encoder.sampling_rate

parler_tts/dac_wrapper/configuration_dac.py

@@ -1,9 +1,11 @@
 
 from transformers import PretrainedConfig
+from importlib.metadata import version
+from packaging.version import Version
 
 
 class DACConfig(PretrainedConfig):
-    model_type = "dac"
+    model_type = "dac" if Version(version("transformers"))<= Version("4.44.2dev") else "dac_on_the_hub"
 
     def __init__(
         self,

parler_tts/dac_wrapper/modeling_dac.py

@@ -1,5 +1,7 @@
 import torch
 from dac.model import DAC
+from torch import nn
+
 from transformers import PreTrainedModel
 from transformers.models.encodec.modeling_encodec import EncodecDecoderOutput, EncodecEncoderOutput
 
@@ -11,6 +13,7 @@
 
 class DACModel(PreTrainedModel):
     config_class = DACConfig
+    main_input_name = "input_values"
 
     # Set main input to 'input_values' for voice steering
     main_input_name = "input_values"
@@ -23,6 +26,9 @@ def __init__(self, config):
             latent_dim=config.latent_dim,
             codebook_size=config.codebook_size,
         )
+        
+        self.remove_weight_norm()
+        self.apply_weight_norm()
 
     def encode(
         self, input_values, padding_mask=None, bandwidth=None, return_dict=None, n_quantizers=None, sample_rate=None
@@ -137,3 +143,22 @@ def decode(
 
     def forward(self, tensor):
         raise ValueError("`DACModel.forward` not implemented yet")
+    
+
+    def apply_weight_norm(self):
+        weight_norm = nn.utils.weight_norm
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+
+        def _apply_weight_norm(module):
+            if isinstance(module, nn.Conv1d) or isinstance(module, nn.ConvTranspose1d):
+                weight_norm(module)
+
+        self.apply(_apply_weight_norm)
+
+
+    def remove_weight_norm(self):
+        def _remove_weight_norm(module):
+            if isinstance(module, nn.Conv1d) or isinstance(module, nn.ConvTranspose1d):
+                nn.utils.remove_weight_norm(module)
+        self.apply(_remove_weight_norm)