From dfad28c3cf94230a3836f26b1cef134570378ec8 Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Mon, 2 Feb 2026 12:27:22 +0100
Subject: [PATCH 1/8] do smth

---
 src/transformers/generation/utils.py          | 130 ++++++-------
 .../models/janus/modeling_janus.py            |   2 +-
 .../models/janus/modular_janus.py             | 137 +-------------
 .../modeling_kyutai_speech_to_text.py         | 173 +-----------------
 .../modular_kyutai_speech_to_text.py          | 170 +----------------
 5 files changed, 65 insertions(+), 547 deletions(-)

diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
index c929f88f184a..420d61eaf3f0 100644
--- a/src/transformers/generation/utils.py
+++ b/src/transformers/generation/utils.py
@@ -1858,7 +1858,9 @@ def _get_initial_cache_position(self, seq_length, device, model_kwargs):
         model_kwargs["cache_position"] = cache_position
         return model_kwargs
 
-    def _get_cache(self, cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs) -> Cache:
+    def _prepare_static_cache(
+        self, cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs
+    ) -> Cache:
         """
         Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a
         new `generate` call requires a larger cache or uses a different batch size.
@@ -1936,12 +1938,8 @@ def _prepare_cache_for_generation(
         instantiated, writes it to `model_kwargs`, under the name expected by the model.
         """
 
-        is_hybrid_cache = any(class_name in self.__class__.__name__.lower() for class_name in ["mamba", "falconh1"])
-        cache_name = "past_key_values" if not is_hybrid_cache else "cache_params"
-
-        requires_cross_attention_cache = (
-            self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
-        )
+        is_linear_cache = "mamba" in self.__class__.__name__.lower()
+        cache_name = "past_key_values" if not is_linear_cache else "cache_params"
 
         # Quick escape route 1: if the user specifies a cache, we only need to check for conflicting `generate` arguments
         user_defined_cache = model_kwargs.get(cache_name)
@@ -1973,74 +1971,68 @@ def _prepare_cache_for_generation(
 
         # Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation`
 
+        # Assisted decoding and contrastive search require cache rollback, which is incompatible with sliding layers.
+        # To handle this, we skip passing the model config to DynamicCache (forcing a full-layer cache).
+        # The "dynamic_full" option is a shortcut for generate() users to avoid sliding layers on their own.
         # TODO(joao): support static caches in assisted generation. assisted generation needs to roll back caches,
         # which is only supported in dynamic caches atm
-        if (
-            generation_mode == GenerationMode.ASSISTED_GENERATION
-            and generation_config.cache_implementation is not None
-        ):
-            logger.warning_once(
-                "An assistant model is provided, using a dynamic cache instead of a cache of type="
-                f"'{generation_config.cache_implementation}'."
+        if generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH):
+            if generation_config.cache_implementation is not None:
+                logger.warning_once(
+                    "An assistant model is provided, using a dynamic cache instead of a cache of type="
+                    f"'{generation_config.cache_implementation}'."
+                )
+            generation_config.cache_implementation = "dynamic_full"
+
+        dynamic_cache_kwargs = {}
+        if generation_config.cache_implementation != "dynamic_full":
+            dynamic_cache_kwargs["config"] = self.config.get_text_config(decoder=True)
+
+        if generation_config.cache_implementation == "offloaded":
+            dynamic_cache_kwargs["offloading"] = True
+
+        if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS:
+            if generation_config.cache_implementation in DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS:
+                logger.warning_once(
+                    f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated. "
+                    f"Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, and the layer structure will be "
+                    "inferred automatically."
+                )
+            model_kwargs["past_key_values"] = self._prepare_static_cache(
+                cache_implementation=generation_config.cache_implementation,
+                batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
+                max_cache_len=max_cache_length,
+                model_kwargs=model_kwargs,
             )
-            generation_config.cache_implementation = None
+        elif generation_config.cache_implementation == "quantized":
+            if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache():
+                raise ValueError(
+                    "This model does not support the quantized cache. If you want your model to support quantized "
+                    "cache, please open an issue and tag @zucchini-nlp."
+                )
 
-        # Assisted decoding and contrastive search require cache rollback, which is incompatible with sliding layers.
-        # To handle this, we skip passing the model config to DynamicCache (forcing a full-layer cache).
-        # The "dynamic_full" option is a shortcut for generate() users to avoid sliding layers on their own.
-        if (
-            generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH)
-            or generation_config.cache_implementation == "dynamic_full"
-        ):
-            dynamic_cache_kwargs = {}
-        else:
-            dynamic_cache_kwargs = {"config": self.config.get_text_config(decoder=True)}
-        if generation_config.cache_implementation is not None:
-            if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS:
-                if generation_config.cache_implementation in DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS:
-                    logger.warning_once(
-                        f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated. "
-                        f"Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, and the layer structure will be "
-                        "inferred automatically."
-                    )
-                model_kwargs[cache_name] = self._get_cache(
-                    cache_implementation=generation_config.cache_implementation,
-                    batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
-                    max_cache_len=max_cache_length,
-                    model_kwargs=model_kwargs,
+            cache_config = generation_config.cache_config if generation_config.cache_config is not None else {}
+            cache_config.setdefault("config", self.config.get_text_config(decoder=True))
+            backend = cache_config.pop("backend", "quanto")
+            if backend == "quanto" and not is_optimum_quanto_available():
+                raise ImportError(
+                    "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto "
+                    "backend. Please install it via  with `pip install optimum-quanto`"
                 )
-            elif generation_config.cache_implementation == "quantized":
-                if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache():
-                    raise ValueError(
-                        "This model does not support the quantized cache. If you want your model to support quantized "
-                        "cache, please open an issue and tag @zucchini-nlp."
-                    )
+            elif backend == "HQQ" and not is_hqq_available():
+                raise ImportError(
+                    "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. "
+                    "Please install it via  with `pip install hqq`"
+                )
+            model_kwargs["past_key_values"] = QuantizedCache(backend=backend, **cache_config)
+        # i.e. `cache_implementation` in [None, "dynamic", "offloaded", "dynamic_full"]
+        # TODO: prepare linear cache from a single API, instead of creating in modeling code
+        elif self._supports_default_dynamic_cache():
+            model_kwargs["past_key_values"] = DynamicCache(**dynamic_cache_kwargs)
 
-                cache_config = generation_config.cache_config if generation_config.cache_config is not None else {}
-                # Add the config if it was not provided, as it's a required argument
-                if "config" not in cache_config:
-                    cache_config["config"] = self.config.get_text_config()
-                # Pop the backend from the config (defaults to quanto if not defined)
-                backend = cache_config.pop("backend", "quanto")
-
-                if backend == "quanto" and not is_optimum_quanto_available():
-                    raise ImportError(
-                        "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto "
-                        "backend. Please install it via  with `pip install optimum-quanto`"
-                    )
-                elif backend == "HQQ" and not is_hqq_available():
-                    raise ImportError(
-                        "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. "
-                        "Please install it via  with `pip install hqq`"
-                    )
-                model_kwargs[cache_name] = QuantizedCache(backend=backend, **cache_config)
-            elif generation_config.cache_implementation == "offloaded":
-                model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs, offloading=True)
-            elif "dynamic" in generation_config.cache_implementation:
-                model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs)
-
-        # TODO (joao): this logic is incomplete, e.g. `offloaded` should apply to both caches. Refactor this function
-        # to correctly pass parameterization to both caches.
+        requires_cross_attention_cache = (
+            self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
+        )
         if (
             requires_cross_attention_cache
             and "past_key_values" in model_kwargs
diff --git a/src/transformers/models/janus/modeling_janus.py b/src/transformers/models/janus/modeling_janus.py
index 0c61a42ee790..08098e8e88ab 100644
--- a/src/transformers/models/janus/modeling_janus.py
+++ b/src/transformers/models/janus/modeling_janus.py
@@ -1314,7 +1314,7 @@ def generate(
 
         if model_kwargs.get("past_key_values", None) is None:
             # Prepare cache if not provided.
-            model_kwargs["past_key_values"] = self._get_cache(
+            model_kwargs["past_key_values"] = self._prepare_static_cache(
                 cache_implementation=generation_config.cache_implementation or "static",
                 # batch_size should account for both conditional/unconditional input; hence multiplied by 2.
                 batch_size=batch_size * 2,
diff --git a/src/transformers/models/janus/modular_janus.py b/src/transformers/models/janus/modular_janus.py
index 67b2f6c7b8a8..f676ad10f451 100644
--- a/src/transformers/models/janus/modular_janus.py
+++ b/src/transformers/models/janus/modular_janus.py
@@ -1255,7 +1255,7 @@ def generate(
 
         if model_kwargs.get("past_key_values", None) is None:
             # Prepare cache if not provided.
-            model_kwargs["past_key_values"] = self._get_cache(
+            model_kwargs["past_key_values"] = self._prepare_static_cache(
                 cache_implementation=generation_config.cache_implementation or "static",
                 # batch_size should account for both conditional/unconditional input; hence multiplied by 2.
                 batch_size=batch_size * 2,
@@ -1678,138 +1678,3 @@ def preprocess(
             # Expand and pad the images to obtain a square image of dimensions `size x size`
             images = [
                 self.pad_to_square(
-                    image=image,
-                    background_color=background_color,
-                    input_data_format=input_data_format,
-                )
-                for image in images
-            ]
-
-        if do_rescale:
-            images = [
-                self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
-                for image in images
-            ]
-
-        if do_normalize:
-            images = [
-                self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
-                for image in images
-            ]
-
-        images = [
-            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
-        ]
-
-        encoded_outputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)
-
-        return encoded_outputs
-
-    def postprocess(
-        self,
-        images: ImageInput,
-        do_rescale: bool | None = None,
-        rescale_factor: float | None = None,
-        do_normalize: bool | None = None,
-        image_mean: list[float] | None = None,
-        image_std: list[float] | None = None,
-        input_data_format: str | None = None,
-        return_tensors: str | None = None,
-    ):
-        """Applies post-processing to the decoded image tokens by reversing transformations applied during preprocessing."""
-        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
-        rescale_factor = 1.0 / self.rescale_factor if rescale_factor is None else rescale_factor
-        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
-        image_mean = image_mean if image_mean is not None else self.image_mean
-        image_std = image_std if image_std is not None else self.image_std
-
-        images = make_flat_list_of_images(images)  # Ensures input is a list
-
-        if isinstance(images[0], PIL.Image.Image):
-            return images if len(images) > 1 else images[0]
-
-        if input_data_format is None:
-            input_data_format = infer_channel_dimension_format(images[0])  # Determine format dynamically
-
-        pixel_values = []
-
-        for image in images:
-            image = to_numpy_array(image)  # Ensure NumPy format
-
-            if do_normalize:
-                image = self.unnormalize(
-                    image=image, image_mean=image_mean, image_std=image_std, input_data_format=input_data_format
-                )
-
-            if do_rescale:
-                image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
-                image = image.clip(0, 255).astype(np.uint8)
-
-            if do_normalize and do_rescale and return_tensors == "PIL.Image.Image":
-                image = to_channel_dimension_format(image, ChannelDimension.LAST, input_channel_dim=input_data_format)
-                image = PIL.Image.fromarray(image)
-
-            pixel_values.append(image)
-
-        data = {"pixel_values": pixel_values}
-        return_tensors = return_tensors if return_tensors != "PIL.Image.Image" else None
-
-        return BatchFeature(data=data, tensor_type=return_tensors)
-
-    def unnormalize(
-        self,
-        image: np.ndarray,
-        image_mean: float | Iterable[float],
-        image_std: float | Iterable[float],
-        input_data_format: str | ChannelDimension | None = None,
-    ) -> np.ndarray:
-        """
-        Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.
-        image = (image * image_std) + image_mean
-        Args:
-            image (`torch.Tensor` of shape `(batch_size, num_channels, image_size, image_size)` or `(num_channels, image_size, image_size)`):
-                Batch of pixel values to postprocess.
-            image_mean (`float` or `Iterable[float]`):
-                The mean to use for unnormalization.
-            image_std (`float` or `Iterable[float]`):
-                The standard deviation to use for unnormalization.
-            input_data_format (`ChannelDimension` or `str`, *optional*):
-                The channel dimension format for the input image. If unset, the channel dimension format is inferred
-                from the input image. Can be one of:
-                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
-                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
-                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
-        """
-        num_channels = 3
-
-        if isinstance(image_mean, Iterable):
-            if len(image_mean) != num_channels:
-                raise ValueError(f"mean must have {num_channels} elements if it is an iterable, got {len(image_mean)}")
-        else:
-            image_mean = [image_mean] * num_channels
-
-        if isinstance(image_std, Iterable):
-            if len(image_std) != num_channels:
-                raise ValueError(f"std must have {num_channels} elements if it is an iterable, got {len(image_std)}")
-        else:
-            image_std = [image_std] * num_channels
-
-        rev_image_mean = tuple(-mean / std for mean, std in zip(image_mean, image_std))
-        rev_image_std = tuple(1 / std for std in image_std)
-        image = self.normalize(
-            image=image, mean=rev_image_mean, std=rev_image_std, input_data_format=input_data_format
-        )
-        return image
-
-
-__all__ = [
-    "JanusImageProcessor",
-    "JanusPreTrainedModel",
-    "JanusForConditionalGeneration",
-    "JanusModel",
-    "JanusVQVAE",
-    "JanusVisionModel",
-    "JanusVQVAEConfig",
-    "JanusVisionConfig",
-    "JanusConfig",
-]
diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
index 785d60879cf1..167edfda2330 100644
--- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
+++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
@@ -1214,175 +1214,4 @@ def _prepare_model_inputs(
         )
 
         # let's use generate's cache preparation to prepare the cache for the codec model
-        temporary_model_kwargs = {}
-
-        # monkey patching the codec model with cache preparation methods since we don't want it to inherit fully from GenerationMixin
-        # Add cache-related methods from GenerationMixin to codec model
-        cache_methods = [
-            "_prepare_cache_for_generation",
-            "_get_cache",
-        ]
-        for method in cache_methods:
-            setattr(self.codec_model, method, types.MethodType(getattr(self, method).__func__, self.codec_model))
-
-        setattr(
-            self.codec_model, "_supports_default_dynamic_cache", types.MethodType(lambda x: True, self.codec_model)
-        )
-
-        self.codec_model.generation_config.cache_implementation = "dynamic"
-        self.codec_model._prepare_cache_for_generation(
-            generation_config=self.codec_model.generation_config,
-            model_kwargs=temporary_model_kwargs,
-            generation_mode=None,
-            batch_size=batch_size,
-            max_cache_length=self.config.codec_config.sliding_window,
-        )
-
-        if "past_key_values" in temporary_model_kwargs:
-            model_kwargs["encoder_past_key_values"] = temporary_model_kwargs["past_key_values"]
-
-        # initialize the padding cache for the codec model
-        per_layer_padding, per_layer_padding_mode, per_layer_in_channels = [], [], []
-        for layer_name in self.codec_model.encoder._mimiconv1d_layer_names:
-            per_layer_padding.append(self.codec_model.encoder.get_submodule(layer_name).padding_total)
-            per_layer_padding_mode.append(self.codec_model.encoder.get_submodule(layer_name).pad_mode)
-            per_layer_in_channels.append(self.codec_model.encoder.get_submodule(layer_name).in_channels)
-
-        # downsample layer
-        per_layer_padding.append(self.codec_model.downsample.padding_total)
-        per_layer_padding_mode.append(self.codec_model.downsample.pad_mode)
-        per_layer_in_channels.append(self.codec_model.downsample.in_channels)
-
-        model_kwargs["padding_cache"] = KyutaiSpeechToTextConv1dPaddingCache(
-            num_layers=len(self.codec_model.encoder._mimiconv1d_layer_names) + 1,
-            per_layer_padding=per_layer_padding,
-            per_layer_padding_mode=per_layer_padding_mode,
-            per_layer_in_channels=per_layer_in_channels,
-        )
-
-        return inputs, input_name, model_kwargs
-
-    def prepare_inputs_for_generation(
-        self,
-        *args,
-        audio_tokens: torch.LongTensor | None = None,
-        input_values: torch.FloatTensor | None = None,
-        padding_mask: torch.Tensor | None = None,
-        audio_window_size: int | None = None,
-        current_window: tuple[int, int] | None = None,
-        encoder_past_key_values: Cache | None = None,
-        padding_cache: KyutaiSpeechToTextConv1dPaddingCache | None = None,
-        **kwargs,
-    ):
-        model_inputs = super().prepare_inputs_for_generation(*args, **kwargs)
-
-        if input_values is not None:
-            cache_position = model_inputs["cache_position"]
-            start, end = current_window[0]
-
-            # first cache position is for bos token, so we need to offset by -1
-            if cache_position[-1] - 1 >= end:
-                # we need to encode the new audio tokens
-                with torch.no_grad():
-                    input_values_start_idx = start * self.config.frame_size
-                    input_values_end_idx = (start + audio_window_size) * self.config.frame_size
-                    current_input_values = input_values[..., input_values_start_idx:input_values_end_idx]
-                    codec_model_output = self.codec_model.encode(
-                        current_input_values,
-                        encoder_past_key_values=encoder_past_key_values,
-                        padding_cache=padding_cache,
-                    )
-                    new_audio_tokens = codec_model_output.audio_codes.transpose(1, 2)
-
-                audio_tokens.copy_(new_audio_tokens)
-
-                start = end.clone()
-                end = end + audio_window_size
-                current_window.copy_(
-                    torch.tensor([start, end], device=current_window.device).expand(current_window.shape[0], -1)
-                )
-
-            # first cache position is for bos token, so we need to offset by -1
-            current_audio_tokens_idxs = (cache_position - start - 1).clamp(min=0)
-            current_audio_tokens = audio_tokens[:, current_audio_tokens_idxs, :]
-
-            current_audio_tokens[:, cache_position == 0, :] = self.config.audio_bos_token_id
-
-            input_ids = model_inputs.pop("input_ids")
-            input_ids = torch.cat(
-                [input_ids.unsqueeze(2), current_audio_tokens],
-                dim=2,
-            )
-            model_inputs["input_ids"] = input_ids
-
-        return model_inputs
-
-    # TODO: @eustlb, this should be standardized
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        if kwargs.get("output_loading_info", False):
-            model, loading_info = super().from_pretrained(*args, **kwargs)
-        else:
-            model = super().from_pretrained(*args, **kwargs)
-
-        # copy depth decoder generation conf attr to the depth decoder generation config
-        prefix = "codec_"
-        prefix_len = len(prefix)
-        codec_model_attrs = {
-            attr[prefix_len:]: value
-            for attr, value in vars(model.generation_config).items()
-            if attr.startswith(prefix)
-        }
-
-        vars(model.codec_model.generation_config).update({"_from_model_config": False, **codec_model_attrs})
-
-        # remove the depth decoder generation conf attr from the model generation config
-        for attr in codec_model_attrs:
-            delattr(model.generation_config, prefix + attr)
-
-        if "output_loading_info" in kwargs:
-            return model, loading_info
-        else:
-            return model
-
-    # TODO: @eustlb, this should be standardized
-    def save_pretrained(self, *args, **kwargs):
-        prefix = "codec_"
-        codec_model_attrs = self.codec_model.generation_config.to_diff_dict()
-        codec_model_attrs.pop("transformers_version", None)
-        for attr, value in codec_model_attrs.items():
-            setattr(self.generation_config, prefix + attr, value)
-
-        super().save_pretrained(*args, **kwargs)
-
-    def generate(self, *args, **kwargs):
-        r"""
-        This method forwards all its arguments to GenerationMixin's [`~GenerationMixin.generate`]. Please refer to the docstring of this method for more information.
-        """
-        max_new_tokens = kwargs.pop("max_new_tokens", None)
-        input_values = kwargs.get("input_values")
-
-        # TODO: @eustlb, we should have per-batch-idx values
-        # here we do not use padding_mask to be aligned to what's done in the original codebase
-        max_audio_frames = input_values.shape[-1] // self.config.codec_config.frame_size
-
-        if max_new_tokens is None or max_new_tokens > max_audio_frames:
-            if max_new_tokens is not None:
-                logger.warning(
-                    f"`max_new_tokens` ({max_new_tokens}) is greater than the maximum number of audio frames ({max_audio_frames})."
-                    f"Setting `max_new_tokens` to {max_audio_frames}."
-                )
-            max_new_tokens = max_audio_frames
-
-        return super().generate(
-            *args,
-            max_new_tokens=max_new_tokens,
-            **kwargs,
-        )
-
-
-__all__ = [
-    "KyutaiSpeechToTextPreTrainedModel",
-    "KyutaiSpeechToTextModel",
-    "KyutaiSpeechToTextForConditionalGeneration",
-]
+        temporary_model_
\ No newline at end of file
diff --git a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
index 39bd26f3987d..e35c7df7d1d5 100644
--- a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
+++ b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
@@ -349,172 +349,4 @@ def _prepare_model_inputs(
 
         # monkey patching the codec model with cache preparation methods since we don't want it to inherit fully from GenerationMixin
         # Add cache-related methods from GenerationMixin to codec model
-        cache_methods = [
-            "_prepare_cache_for_generation",
-            "_get_cache",
-        ]
-        for method in cache_methods:
-            setattr(self.codec_model, method, types.MethodType(getattr(self, method).__func__, self.codec_model))
-
-        setattr(
-            self.codec_model, "_supports_default_dynamic_cache", types.MethodType(lambda x: True, self.codec_model)
-        )
-
-        self.codec_model.generation_config.cache_implementation = "dynamic"
-        self.codec_model._prepare_cache_for_generation(
-            generation_config=self.codec_model.generation_config,
-            model_kwargs=temporary_model_kwargs,
-            generation_mode=None,
-            batch_size=batch_size,
-            max_cache_length=self.config.codec_config.sliding_window,
-        )
-
-        if "past_key_values" in temporary_model_kwargs:
-            model_kwargs["encoder_past_key_values"] = temporary_model_kwargs["past_key_values"]
-
-        # initialize the padding cache for the codec model
-        per_layer_padding, per_layer_padding_mode, per_layer_in_channels = [], [], []
-        for layer_name in self.codec_model.encoder._mimiconv1d_layer_names:
-            per_layer_padding.append(self.codec_model.encoder.get_submodule(layer_name).padding_total)
-            per_layer_padding_mode.append(self.codec_model.encoder.get_submodule(layer_name).pad_mode)
-            per_layer_in_channels.append(self.codec_model.encoder.get_submodule(layer_name).in_channels)
-
-        # downsample layer
-        per_layer_padding.append(self.codec_model.downsample.padding_total)
-        per_layer_padding_mode.append(self.codec_model.downsample.pad_mode)
-        per_layer_in_channels.append(self.codec_model.downsample.in_channels)
-
-        model_kwargs["padding_cache"] = KyutaiSpeechToTextConv1dPaddingCache(
-            num_layers=len(self.codec_model.encoder._mimiconv1d_layer_names) + 1,
-            per_layer_padding=per_layer_padding,
-            per_layer_padding_mode=per_layer_padding_mode,
-            per_layer_in_channels=per_layer_in_channels,
-        )
-
-        return inputs, input_name, model_kwargs
-
-    def prepare_inputs_for_generation(
-        self,
-        *args,
-        audio_tokens: torch.LongTensor | None = None,
-        input_values: torch.FloatTensor | None = None,
-        padding_mask: torch.Tensor | None = None,
-        audio_window_size: int | None = None,
-        current_window: tuple[int, int] | None = None,
-        encoder_past_key_values: Cache | None = None,
-        padding_cache: KyutaiSpeechToTextConv1dPaddingCache | None = None,
-        **kwargs,
-    ):
-        model_inputs = GenerationMixin.prepare_inputs_for_generation(self, *args, **kwargs)
-
-        if input_values is not None:
-            cache_position = model_inputs["cache_position"]
-            start, end = current_window[0]
-
-            # first cache position is for bos token, so we need to offset by -1
-            if cache_position[-1] - 1 >= end:
-                # we need to encode the new audio tokens
-                with torch.no_grad():
-                    input_values_start_idx = start * self.config.frame_size
-                    input_values_end_idx = (start + audio_window_size) * self.config.frame_size
-                    current_input_values = input_values[..., input_values_start_idx:input_values_end_idx]
-                    codec_model_output = self.codec_model.encode(
-                        current_input_values,
-                        encoder_past_key_values=encoder_past_key_values,
-                        padding_cache=padding_cache,
-                    )
-                    new_audio_tokens = codec_model_output.audio_codes.transpose(1, 2)
-
-                audio_tokens.copy_(new_audio_tokens)
-
-                start = end.clone()
-                end = end + audio_window_size
-                current_window.copy_(
-                    torch.tensor([start, end], device=current_window.device).expand(current_window.shape[0], -1)
-                )
-
-            # first cache position is for bos token, so we need to offset by -1
-            current_audio_tokens_idxs = (cache_position - start - 1).clamp(min=0)
-            current_audio_tokens = audio_tokens[:, current_audio_tokens_idxs, :]
-
-            current_audio_tokens[:, cache_position == 0, :] = self.config.audio_bos_token_id
-
-            input_ids = model_inputs.pop("input_ids")
-            input_ids = torch.cat(
-                [input_ids.unsqueeze(2), current_audio_tokens],
-                dim=2,
-            )
-            model_inputs["input_ids"] = input_ids
-
-        return model_inputs
-
-    # TODO: @eustlb, this should be standardized
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        if kwargs.get("output_loading_info", False):
-            model, loading_info = PreTrainedModel.from_pretrained(*args, **kwargs)
-        else:
-            model = PreTrainedModel.from_pretrained(*args, **kwargs)
-
-        # copy depth decoder generation conf attr to the depth decoder generation config
-        prefix = "codec_"
-        prefix_len = len(prefix)
-        codec_model_attrs = {
-            attr[prefix_len:]: value
-            for attr, value in vars(model.generation_config).items()
-            if attr.startswith(prefix)
-        }
-
-        vars(model.codec_model.generation_config).update({"_from_model_config": False, **codec_model_attrs})
-
-        # remove the depth decoder generation conf attr from the model generation config
-        for attr in codec_model_attrs:
-            delattr(model.generation_config, prefix + attr)
-
-        if "output_loading_info" in kwargs:
-            return model, loading_info
-        else:
-            return model
-
-    # TODO: @eustlb, this should be standardized
-    def save_pretrained(self, *args, **kwargs):
-        prefix = "codec_"
-        codec_model_attrs = self.codec_model.generation_config.to_diff_dict()
-        codec_model_attrs.pop("transformers_version", None)
-        for attr, value in codec_model_attrs.items():
-            setattr(self.generation_config, prefix + attr, value)
-
-        PreTrainedModel.save_pretrained(self, *args, **kwargs)
-
-    def generate(self, *args, **kwargs):
-        r"""
-        This method forwards all its arguments to GenerationMixin's [`~GenerationMixin.generate`]. Please refer to the docstring of this method for more information.
-        """
-        max_new_tokens = kwargs.pop("max_new_tokens", None)
-        input_values = kwargs.get("input_values")
-
-        # TODO: @eustlb, we should have per-batch-idx values
-        # here we do not use padding_mask to be aligned to what's done in the original codebase
-        max_audio_frames = input_values.shape[-1] // self.config.codec_config.frame_size
-
-        if max_new_tokens is None or max_new_tokens > max_audio_frames:
-            if max_new_tokens is not None:
-                logger.warning(
-                    f"`max_new_tokens` ({max_new_tokens}) is greater than the maximum number of audio frames ({max_audio_frames})."
-                    f"Setting `max_new_tokens` to {max_audio_frames}."
-                )
-            max_new_tokens = max_audio_frames
-
-        return GenerationMixin.generate(
-            *args,
-            max_new_tokens=max_new_tokens,
-            **kwargs,
-        )
-
-
-__all__ = [
-    "KyutaiSpeechToTextPreTrainedModel",
-    "KyutaiSpeechToTextModel",
-    "KyutaiSpeechToTextForConditionalGeneration",
-    "KyutaiSpeechToTextFeatureExtractor",
-]
+        ca
\ No newline at end of file

From b2415576cd620d9ab5f6bed7fedbc53c73e0f984 Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Mon, 2 Feb 2026 14:24:04 +0100
Subject: [PATCH 2/8] how did this even happen O_o

---
 .../models/janus/modular_janus.py             | 135 ++++++++++++++
 .../modeling_kyutai_speech_to_text.py         | 173 +++++++++++++++++-
 .../modular_kyutai_speech_to_text.py          | 170 ++++++++++++++++-
 3 files changed, 476 insertions(+), 2 deletions(-)

diff --git a/src/transformers/models/janus/modular_janus.py b/src/transformers/models/janus/modular_janus.py
index f676ad10f451..902d903cdd68 100644
--- a/src/transformers/models/janus/modular_janus.py
+++ b/src/transformers/models/janus/modular_janus.py
@@ -1678,3 +1678,138 @@ def preprocess(
             # Expand and pad the images to obtain a square image of dimensions `size x size`
             images = [
                 self.pad_to_square(
+                    image=image,
+                    background_color=background_color,
+                    input_data_format=input_data_format,
+                )
+                for image in images
+            ]
+
+        if do_rescale:
+            images = [
+                self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
+                for image in images
+            ]
+
+        if do_normalize:
+            images = [
+                self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+                for image in images
+            ]
+
+        images = [
+            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
+        ]
+
+        encoded_outputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)
+
+        return encoded_outputs
+
+    def postprocess(
+        self,
+        images: ImageInput,
+        do_rescale: bool | None = None,
+        rescale_factor: float | None = None,
+        do_normalize: bool | None = None,
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        input_data_format: str | None = None,
+        return_tensors: str | None = None,
+    ):
+        """Applies post-processing to the decoded image tokens by reversing transformations applied during preprocessing."""
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = 1.0 / self.rescale_factor if rescale_factor is None else rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+
+        images = make_flat_list_of_images(images)  # Ensures input is a list
+
+        if isinstance(images[0], PIL.Image.Image):
+            return images if len(images) > 1 else images[0]
+
+        if input_data_format is None:
+            input_data_format = infer_channel_dimension_format(images[0])  # Determine format dynamically
+
+        pixel_values = []
+
+        for image in images:
+            image = to_numpy_array(image)  # Ensure NumPy format
+
+            if do_normalize:
+                image = self.unnormalize(
+                    image=image, image_mean=image_mean, image_std=image_std, input_data_format=input_data_format
+                )
+
+            if do_rescale:
+                image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
+                image = image.clip(0, 255).astype(np.uint8)
+
+            if do_normalize and do_rescale and return_tensors == "PIL.Image.Image":
+                image = to_channel_dimension_format(image, ChannelDimension.LAST, input_channel_dim=input_data_format)
+                image = PIL.Image.fromarray(image)
+
+            pixel_values.append(image)
+
+        data = {"pixel_values": pixel_values}
+        return_tensors = return_tensors if return_tensors != "PIL.Image.Image" else None
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def unnormalize(
+        self,
+        image: np.ndarray,
+        image_mean: float | Iterable[float],
+        image_std: float | Iterable[float],
+        input_data_format: str | ChannelDimension | None = None,
+    ) -> np.ndarray:
+        """
+        Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.
+        image = (image * image_std) + image_mean
+        Args:
+            image (`torch.Tensor` of shape `(batch_size, num_channels, image_size, image_size)` or `(num_channels, image_size, image_size)`):
+                Batch of pixel values to postprocess.
+            image_mean (`float` or `Iterable[float]`):
+                The mean to use for unnormalization.
+            image_std (`float` or `Iterable[float]`):
+                The standard deviation to use for unnormalization.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. If unset, the channel dimension format is inferred
+                from the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+        """
+        num_channels = 3
+
+        if isinstance(image_mean, Iterable):
+            if len(image_mean) != num_channels:
+                raise ValueError(f"mean must have {num_channels} elements if it is an iterable, got {len(image_mean)}")
+        else:
+            image_mean = [image_mean] * num_channels
+
+        if isinstance(image_std, Iterable):
+            if len(image_std) != num_channels:
+                raise ValueError(f"std must have {num_channels} elements if it is an iterable, got {len(image_std)}")
+        else:
+            image_std = [image_std] * num_channels
+
+        rev_image_mean = tuple(-mean / std for mean, std in zip(image_mean, image_std))
+        rev_image_std = tuple(1 / std for std in image_std)
+        image = self.normalize(
+            image=image, mean=rev_image_mean, std=rev_image_std, input_data_format=input_data_format
+        )
+        return image
+
+
+__all__ = [
+    "JanusImageProcessor",
+    "JanusPreTrainedModel",
+    "JanusForConditionalGeneration",
+    "JanusModel",
+    "JanusVQVAE",
+    "JanusVisionModel",
+    "JanusVQVAEConfig",
+    "JanusVisionConfig",
+    "JanusConfig",
+]
diff --git a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
index 167edfda2330..094ec2c0948f 100644
--- a/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
+++ b/src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py
@@ -1214,4 +1214,175 @@ def _prepare_model_inputs(
         )
 
         # let's use generate's cache preparation to prepare the cache for the codec model
-        temporary_model_
\ No newline at end of file
+        temporary_model_kwargs = {}
+
+        # monkey patching the codec model with cache preparation methods since we don't want it to inherit fully from GenerationMixin
+        # Add cache-related methods from GenerationMixin to codec model
+        cache_methods = [
+            "_prepare_cache_for_generation",
+            "_prepare_static_cache",
+        ]
+        for method in cache_methods:
+            setattr(self.codec_model, method, types.MethodType(getattr(self, method).__func__, self.codec_model))
+
+        setattr(
+            self.codec_model, "_supports_default_dynamic_cache", types.MethodType(lambda x: True, self.codec_model)
+        )
+
+        self.codec_model.generation_config.cache_implementation = "dynamic"
+        self.codec_model._prepare_cache_for_generation(
+            generation_config=self.codec_model.generation_config,
+            model_kwargs=temporary_model_kwargs,
+            generation_mode=None,
+            batch_size=batch_size,
+            max_cache_length=self.config.codec_config.sliding_window,
+        )
+
+        if "past_key_values" in temporary_model_kwargs:
+            model_kwargs["encoder_past_key_values"] = temporary_model_kwargs["past_key_values"]
+
+        # initialize the padding cache for the codec model
+        per_layer_padding, per_layer_padding_mode, per_layer_in_channels = [], [], []
+        for layer_name in self.codec_model.encoder._mimiconv1d_layer_names:
+            per_layer_padding.append(self.codec_model.encoder.get_submodule(layer_name).padding_total)
+            per_layer_padding_mode.append(self.codec_model.encoder.get_submodule(layer_name).pad_mode)
+            per_layer_in_channels.append(self.codec_model.encoder.get_submodule(layer_name).in_channels)
+
+        # downsample layer
+        per_layer_padding.append(self.codec_model.downsample.padding_total)
+        per_layer_padding_mode.append(self.codec_model.downsample.pad_mode)
+        per_layer_in_channels.append(self.codec_model.downsample.in_channels)
+
+        model_kwargs["padding_cache"] = KyutaiSpeechToTextConv1dPaddingCache(
+            num_layers=len(self.codec_model.encoder._mimiconv1d_layer_names) + 1,
+            per_layer_padding=per_layer_padding,
+            per_layer_padding_mode=per_layer_padding_mode,
+            per_layer_in_channels=per_layer_in_channels,
+        )
+
+        return inputs, input_name, model_kwargs
+
+    def prepare_inputs_for_generation(
+        self,
+        *args,
+        audio_tokens: torch.LongTensor | None = None,
+        input_values: torch.FloatTensor | None = None,
+        padding_mask: torch.Tensor | None = None,
+        audio_window_size: int | None = None,
+        current_window: tuple[int, int] | None = None,
+        encoder_past_key_values: Cache | None = None,
+        padding_cache: KyutaiSpeechToTextConv1dPaddingCache | None = None,
+        **kwargs,
+    ):
+        model_inputs = super().prepare_inputs_for_generation(*args, **kwargs)
+
+        if input_values is not None:
+            cache_position = model_inputs["cache_position"]
+            start, end = current_window[0]
+
+            # first cache position is for bos token, so we need to offset by -1
+            if cache_position[-1] - 1 >= end:
+                # we need to encode the new audio tokens
+                with torch.no_grad():
+                    input_values_start_idx = start * self.config.frame_size
+                    input_values_end_idx = (start + audio_window_size) * self.config.frame_size
+                    current_input_values = input_values[..., input_values_start_idx:input_values_end_idx]
+                    codec_model_output = self.codec_model.encode(
+                        current_input_values,
+                        encoder_past_key_values=encoder_past_key_values,
+                        padding_cache=padding_cache,
+                    )
+                    new_audio_tokens = codec_model_output.audio_codes.transpose(1, 2)
+
+                audio_tokens.copy_(new_audio_tokens)
+
+                start = end.clone()
+                end = end + audio_window_size
+                current_window.copy_(
+                    torch.tensor([start, end], device=current_window.device).expand(current_window.shape[0], -1)
+                )
+
+            # first cache position is for bos token, so we need to offset by -1
+            current_audio_tokens_idxs = (cache_position - start - 1).clamp(min=0)
+            current_audio_tokens = audio_tokens[:, current_audio_tokens_idxs, :]
+
+            current_audio_tokens[:, cache_position == 0, :] = self.config.audio_bos_token_id
+
+            input_ids = model_inputs.pop("input_ids")
+            input_ids = torch.cat(
+                [input_ids.unsqueeze(2), current_audio_tokens],
+                dim=2,
+            )
+            model_inputs["input_ids"] = input_ids
+
+        return model_inputs
+
+    # TODO: @eustlb, this should be standardized
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        if kwargs.get("output_loading_info", False):
+            model, loading_info = super().from_pretrained(*args, **kwargs)
+        else:
+            model = super().from_pretrained(*args, **kwargs)
+
+        # copy depth decoder generation conf attr to the depth decoder generation config
+        prefix = "codec_"
+        prefix_len = len(prefix)
+        codec_model_attrs = {
+            attr[prefix_len:]: value
+            for attr, value in vars(model.generation_config).items()
+            if attr.startswith(prefix)
+        }
+
+        vars(model.codec_model.generation_config).update({"_from_model_config": False, **codec_model_attrs})
+
+        # remove the depth decoder generation conf attr from the model generation config
+        for attr in codec_model_attrs:
+            delattr(model.generation_config, prefix + attr)
+
+        if "output_loading_info" in kwargs:
+            return model, loading_info
+        else:
+            return model
+
+    # TODO: @eustlb, this should be standardized
+    def save_pretrained(self, *args, **kwargs):
+        prefix = "codec_"
+        codec_model_attrs = self.codec_model.generation_config.to_diff_dict()
+        codec_model_attrs.pop("transformers_version", None)
+        for attr, value in codec_model_attrs.items():
+            setattr(self.generation_config, prefix + attr, value)
+
+        super().save_pretrained(*args, **kwargs)
+
+    def generate(self, *args, **kwargs):
+        r"""
+        This method forwards all its arguments to GenerationMixin's [`~GenerationMixin.generate`]. Please refer to the docstring of this method for more information.
+        """
+        max_new_tokens = kwargs.pop("max_new_tokens", None)
+        input_values = kwargs.get("input_values")
+
+        # TODO: @eustlb, we should have per-batch-idx values
+        # here we do not use padding_mask to be aligned to what's done in the original codebase
+        max_audio_frames = input_values.shape[-1] // self.config.codec_config.frame_size
+
+        if max_new_tokens is None or max_new_tokens > max_audio_frames:
+            if max_new_tokens is not None:
+                logger.warning(
+                    f"`max_new_tokens` ({max_new_tokens}) is greater than the maximum number of audio frames ({max_audio_frames})."
+                    f"Setting `max_new_tokens` to {max_audio_frames}."
+                )
+            max_new_tokens = max_audio_frames
+
+        return super().generate(
+            *args,
+            max_new_tokens=max_new_tokens,
+            **kwargs,
+        )
+
+
+__all__ = [
+    "KyutaiSpeechToTextPreTrainedModel",
+    "KyutaiSpeechToTextModel",
+    "KyutaiSpeechToTextForConditionalGeneration",
+]
diff --git a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
index e35c7df7d1d5..043663bb026c 100644
--- a/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
+++ b/src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
@@ -349,4 +349,172 @@ def _prepare_model_inputs(
 
         # monkey patching the codec model with cache preparation methods since we don't want it to inherit fully from GenerationMixin
         # Add cache-related methods from GenerationMixin to codec model
-        ca
\ No newline at end of file
+        cache_methods = [
+            "_prepare_cache_for_generation",
+            "_prepare_static_cache",
+        ]
+        for method in cache_methods:
+            setattr(self.codec_model, method, types.MethodType(getattr(self, method).__func__, self.codec_model))
+
+        setattr(
+            self.codec_model, "_supports_default_dynamic_cache", types.MethodType(lambda x: True, self.codec_model)
+        )
+
+        self.codec_model.generation_config.cache_implementation = "dynamic"
+        self.codec_model._prepare_cache_for_generation(
+            generation_config=self.codec_model.generation_config,
+            model_kwargs=temporary_model_kwargs,
+            generation_mode=None,
+            batch_size=batch_size,
+            max_cache_length=self.config.codec_config.sliding_window,
+        )
+
+        if "past_key_values" in temporary_model_kwargs:
+            model_kwargs["encoder_past_key_values"] = temporary_model_kwargs["past_key_values"]
+
+        # initialize the padding cache for the codec model
+        per_layer_padding, per_layer_padding_mode, per_layer_in_channels = [], [], []
+        for layer_name in self.codec_model.encoder._mimiconv1d_layer_names:
+            per_layer_padding.append(self.codec_model.encoder.get_submodule(layer_name).padding_total)
+            per_layer_padding_mode.append(self.codec_model.encoder.get_submodule(layer_name).pad_mode)
+            per_layer_in_channels.append(self.codec_model.encoder.get_submodule(layer_name).in_channels)
+
+        # downsample layer
+        per_layer_padding.append(self.codec_model.downsample.padding_total)
+        per_layer_padding_mode.append(self.codec_model.downsample.pad_mode)
+        per_layer_in_channels.append(self.codec_model.downsample.in_channels)
+
+        model_kwargs["padding_cache"] = KyutaiSpeechToTextConv1dPaddingCache(
+            num_layers=len(self.codec_model.encoder._mimiconv1d_layer_names) + 1,
+            per_layer_padding=per_layer_padding,
+            per_layer_padding_mode=per_layer_padding_mode,
+            per_layer_in_channels=per_layer_in_channels,
+        )
+
+        return inputs, input_name, model_kwargs
+
+    def prepare_inputs_for_generation(
+        self,
+        *args,
+        audio_tokens: torch.LongTensor | None = None,
+        input_values: torch.FloatTensor | None = None,
+        padding_mask: torch.Tensor | None = None,
+        audio_window_size: int | None = None,
+        current_window: tuple[int, int] | None = None,
+        encoder_past_key_values: Cache | None = None,
+        padding_cache: KyutaiSpeechToTextConv1dPaddingCache | None = None,
+        **kwargs,
+    ):
+        model_inputs = GenerationMixin.prepare_inputs_for_generation(self, *args, **kwargs)
+
+        if input_values is not None:
+            cache_position = model_inputs["cache_position"]
+            start, end = current_window[0]
+
+            # first cache position is for bos token, so we need to offset by -1
+            if cache_position[-1] - 1 >= end:
+                # we need to encode the new audio tokens
+                with torch.no_grad():
+                    input_values_start_idx = start * self.config.frame_size
+                    input_values_end_idx = (start + audio_window_size) * self.config.frame_size
+                    current_input_values = input_values[..., input_values_start_idx:input_values_end_idx]
+                    codec_model_output = self.codec_model.encode(
+                        current_input_values,
+                        encoder_past_key_values=encoder_past_key_values,
+                        padding_cache=padding_cache,
+                    )
+                    new_audio_tokens = codec_model_output.audio_codes.transpose(1, 2)
+
+                audio_tokens.copy_(new_audio_tokens)
+
+                start = end.clone()
+                end = end + audio_window_size
+                current_window.copy_(
+                    torch.tensor([start, end], device=current_window.device).expand(current_window.shape[0], -1)
+                )
+
+            # first cache position is for bos token, so we need to offset by -1
+            current_audio_tokens_idxs = (cache_position - start - 1).clamp(min=0)
+            current_audio_tokens = audio_tokens[:, current_audio_tokens_idxs, :]
+
+            current_audio_tokens[:, cache_position == 0, :] = self.config.audio_bos_token_id
+
+            input_ids = model_inputs.pop("input_ids")
+            input_ids = torch.cat(
+                [input_ids.unsqueeze(2), current_audio_tokens],
+                dim=2,
+            )
+            model_inputs["input_ids"] = input_ids
+
+        return model_inputs
+
+    # TODO: @eustlb, this should be standardized
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        if kwargs.get("output_loading_info", False):
+            model, loading_info = PreTrainedModel.from_pretrained(*args, **kwargs)
+        else:
+            model = PreTrainedModel.from_pretrained(*args, **kwargs)
+
+        # copy depth decoder generation conf attr to the depth decoder generation config
+        prefix = "codec_"
+        prefix_len = len(prefix)
+        codec_model_attrs = {
+            attr[prefix_len:]: value
+            for attr, value in vars(model.generation_config).items()
+            if attr.startswith(prefix)
+        }
+
+        vars(model.codec_model.generation_config).update({"_from_model_config": False, **codec_model_attrs})
+
+        # remove the depth decoder generation conf attr from the model generation config
+        for attr in codec_model_attrs:
+            delattr(model.generation_config, prefix + attr)
+
+        if "output_loading_info" in kwargs:
+            return model, loading_info
+        else:
+            return model
+
+    # TODO: @eustlb, this should be standardized
+    def save_pretrained(self, *args, **kwargs):
+        prefix = "codec_"
+        codec_model_attrs = self.codec_model.generation_config.to_diff_dict()
+        codec_model_attrs.pop("transformers_version", None)
+        for attr, value in codec_model_attrs.items():
+            setattr(self.generation_config, prefix + attr, value)
+
+        PreTrainedModel.save_pretrained(self, *args, **kwargs)
+
+    def generate(self, *args, **kwargs):
+        r"""
+        This method forwards all its arguments to GenerationMixin's [`~GenerationMixin.generate`]. Please refer to the docstring of this method for more information.
+        """
+        max_new_tokens = kwargs.pop("max_new_tokens", None)
+        input_values = kwargs.get("input_values")
+
+        # TODO: @eustlb, we should have per-batch-idx values
+        # here we do not use padding_mask to be aligned to what's done in the original codebase
+        max_audio_frames = input_values.shape[-1] // self.config.codec_config.frame_size
+
+        if max_new_tokens is None or max_new_tokens > max_audio_frames:
+            if max_new_tokens is not None:
+                logger.warning(
+                    f"`max_new_tokens` ({max_new_tokens}) is greater than the maximum number of audio frames ({max_audio_frames})."
+                    f"Setting `max_new_tokens` to {max_audio_frames}."
+                )
+            max_new_tokens = max_audio_frames
+
+        return GenerationMixin.generate(
+            *args,
+            max_new_tokens=max_new_tokens,
+            **kwargs,
+        )
+
+
+__all__ = [
+    "KyutaiSpeechToTextPreTrainedModel",
+    "KyutaiSpeechToTextModel",
+    "KyutaiSpeechToTextForConditionalGeneration",
+    "KyutaiSpeechToTextFeatureExtractor",
+]

From bd31bb2106f8077b65f4a79e50830c8c4fd31b0e Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Mon, 2 Feb 2026 15:13:27 +0100
Subject: [PATCH 3/8] fix models

---
 src/transformers/generation/utils.py        | 16 +++++-----------
 src/transformers/models/blt/modeling_blt.py | 13 ++++++++++---
 src/transformers/models/blt/modular_blt.py  | 17 +++++++++++++----
 tests/models/afmoe/test_modeling_afmoe.py   |  2 +-
 4 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
index 420d61eaf3f0..4b17b9504ba2 100644
--- a/src/transformers/generation/utils.py
+++ b/src/transformers/generation/utils.py
@@ -1867,13 +1867,10 @@ def _prepare_static_cache(
 
         Returns the resulting cache object.
         """
-        requires_cross_attention_cache = (
-            self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
-        )
         offload_cache = "offloaded" in cache_implementation
 
         if hasattr(self, "_cache"):
-            cache_to_check = self._cache.self_attention_cache if requires_cross_attention_cache else self._cache
+            cache_to_check = self._cache.self_attention_cache if self.config.is_encoder_decoder else self._cache
 
         need_new_cache = (
             not hasattr(self, "_cache")
@@ -1882,7 +1879,7 @@ def _prepare_static_cache(
             or cache_to_check.max_cache_len < max_cache_len
         )
 
-        if requires_cross_attention_cache and hasattr(self, "_cache"):
+        if self.config.is_encoder_decoder and hasattr(self, "_cache"):
             need_new_cache = (
                 need_new_cache
                 or self._cache.cross_attention_cache.max_cache_len != model_kwargs["encoder_outputs"][0].shape[1]
@@ -1895,7 +1892,7 @@ def _prepare_static_cache(
                 "offloading": offload_cache,
             }
             self._cache = StaticCache(**self_attention_cache_kwargs)
-            if requires_cross_attention_cache:
+            if self.config.is_encoder_decoder:
                 cross_attention_cache_kwargs = {
                     "config": self.config.get_text_config(decoder=True),
                     "max_cache_len": model_kwargs["encoder_outputs"][0].shape[1],
@@ -2027,14 +2024,11 @@ def _prepare_cache_for_generation(
             model_kwargs["past_key_values"] = QuantizedCache(backend=backend, **cache_config)
         # i.e. `cache_implementation` in [None, "dynamic", "offloaded", "dynamic_full"]
         # TODO: prepare linear cache from a single API, instead of creating in modeling code
-        elif self._supports_default_dynamic_cache():
+        else:
             model_kwargs["past_key_values"] = DynamicCache(**dynamic_cache_kwargs)
 
-        requires_cross_attention_cache = (
-            self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
-        )
         if (
-            requires_cross_attention_cache
+            self.config.is_encoder_decoder
             and "past_key_values" in model_kwargs
             and not isinstance(model_kwargs["past_key_values"], EncoderDecoderCache)
         ):
diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py
index 8cc3e28bbdf0..5b249f864574 100644
--- a/src/transformers/models/blt/modeling_blt.py
+++ b/src/transformers/models/blt/modeling_blt.py
@@ -1229,8 +1229,15 @@ def forward(
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
-        if use_cache and past_key_values is None:
-            past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = EncoderDecoderCache(
+                    DynamicCache(config=self.config), DynamicCache(config=self.config)
+                )
+            elif not isinstance(past_key_values, EncoderDecoderCache):
+                # BLT uses an encoder-decoder cache even though it is not en encoder-decoder model. Create a cross-cache
+                # if not yet created by the user
+                past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config))
 
         # Extract input embeddings as early as possible
         if inputs_embeds is not None:
@@ -1377,7 +1384,7 @@ class BltForCausalLM(BltPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"model.local_encoder.embed_tokens.weight": "lm_head.weight"}
 
     def __init__(self, config: BltConfig):
-        super().__init__(config.get_text_config())
+        super().__init__(config)
         self.text_config = config.get_text_config()
         self.vocab_size = config.vocab_size
         self.model = BltModel(config)
diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py
index 22fef55c96d2..eed2251329d7 100644
--- a/src/transformers/models/blt/modular_blt.py
+++ b/src/transformers/models/blt/modular_blt.py
@@ -22,6 +22,7 @@
 
 from ... import initialization as init
 from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
 from ...masking_utils import create_causal_mask
 from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
@@ -32,7 +33,6 @@
 from ..cohere2.modeling_cohere2 import rotate_half  # noqa: F401
 from ..llama.modeling_llama import LlamaRotaryEmbedding
 from ..mllama.modeling_mllama import (
-    MllamaForCausalLM,
     MllamaPreTrainedModel,
     MllamaSelfAttentionDecoderLayer,
     MllamaTextCrossAttention,
@@ -954,8 +954,15 @@ def forward(
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
-        if use_cache and past_key_values is None:
-            past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = EncoderDecoderCache(
+                    DynamicCache(config=self.config), DynamicCache(config=self.config)
+                )
+            elif not isinstance(past_key_values, EncoderDecoderCache):
+                # BLT uses an encoder-decoder cache even though it is not en encoder-decoder model. Create a cross-cache
+                # if not yet created by the user
+                past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config))
 
         # Extract input embeddings as early as possible
         if inputs_embeds is not None:
@@ -1090,7 +1097,7 @@ def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) ->
         return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1
 
 
-class BltForCausalLM(MllamaForCausalLM):
+class BltForCausalLM(BltPreTrainedModel, GenerationMixin):
     config: BltConfig
     _can_compile_fullgraph = False
     base_model_prefix = "model"
@@ -1098,9 +1105,11 @@ class BltForCausalLM(MllamaForCausalLM):
 
     def __init__(self, config: BltConfig):
         super().__init__(config)
+        self.text_config = config.get_text_config()
         self.vocab_size = config.vocab_size
         self.model = BltModel(config)
         self.lm_head = nn.Linear(config.decoder_config.hidden_size, config.vocab_size, bias=False)
+
         self.post_init()
 
     def forward(
diff --git a/tests/models/afmoe/test_modeling_afmoe.py b/tests/models/afmoe/test_modeling_afmoe.py
index ed954a4ba687..52af6880b552 100644
--- a/tests/models/afmoe/test_modeling_afmoe.py
+++ b/tests/models/afmoe/test_modeling_afmoe.py
@@ -32,7 +32,7 @@ def __init__(
         self,
         parent,
         batch_size=4,
-        seq_length=128,
+        seq_length=12,
         is_training=True,
         use_input_mask=True,
         use_token_type_ids=False,

From c042373d5de1ce52d81726412bca247088c9276e Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Mon, 2 Feb 2026 16:03:35 +0100
Subject: [PATCH 4/8] dia needs only one prefill, fix

---
 src/transformers/models/blt/modeling_blt.py   | 13 ++----
 src/transformers/models/blt/modular_blt.py    | 44 +++++++++++++++++++
 src/transformers/models/dia/generation_dia.py | 14 +-----
 3 files changed, 49 insertions(+), 22 deletions(-)

diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py
index 5b249f864574..c5e040b0f267 100644
--- a/src/transformers/models/blt/modeling_blt.py
+++ b/src/transformers/models/blt/modeling_blt.py
@@ -37,7 +37,7 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils import TransformersKwargs, auto_docstring
 from ...utils.generic import OutputRecorder, check_model_inputs, maybe_autocast
 from .configuration_blt import (
     BltConfig,
@@ -1372,11 +1372,6 @@ def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) ->
         return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1
 
 
-@auto_docstring(
-    custom_intro="""
-    The Blt Text Model with a language modeling head on top.
-    """
-)
 class BltForCausalLM(BltPreTrainedModel, GenerationMixin):
     config: BltConfig
     _can_compile_fullgraph = False
@@ -1392,8 +1387,6 @@ def __init__(self, config: BltConfig):
 
         self.post_init()
 
-    @can_return_tuple
-    @auto_docstring
     def forward(
         self,
         input_ids: torch.LongTensor | None = None,
@@ -1440,8 +1433,8 @@ def forward(
         ```python
         >>> from transformers import AutoTokenizer, BltForCausalLM
 
-        >>> model = BltForCausalLM.from_pretrained("Llama-3.2-11B-Vision")
-        >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision")
+        >>> model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf")
+        >>> tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")
 
         >>> prompt = "If I had to write a haiku, it would be:"
         >>> inputs = tokenizer(prompt, return_tensors="pt")
diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py
index eed2251329d7..bf79f0147046 100644
--- a/src/transformers/models/blt/modular_blt.py
+++ b/src/transformers/models/blt/modular_blt.py
@@ -1128,6 +1128,50 @@ def forward(
         logits_to_keep: int | torch.Tensor = 0,
         **kwargs: Unpack[TransformersKwargs],
     ) -> tuple | CausalLMOutputWithPast:
+        r"""
+        cross_attention_states (`torch.FloatTensor`, *optional*):
+            Output of the vision model, used for cross-attention. This tensor contains the processed image features that
+            the language model will attend to.
+        cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*):
+            Cross-attention mask to control the interaction between text tokens and image tiles.
+            This 4D tensor defines which image tiles each text token should attend to.
+
+            For each text token (in seq_length):
+            - 1 indicates the token **should attend** to the corresponding image tile
+            - 0 indicates the token **should not attend** to the corresponding image tile
+        full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*):
+            A tuple containing two tensors that mask out rows in the cross-attention mechanism:
+            - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1.
+              A value of 0 indicates that the corresponding text token's entire row in the cross-attention
+              matrix should be masked out (all image tokens ignored).
+            - The second tensor has the same shape and is used internally to apply the masking during
+              the forward pass of cross-attention layers.
+            This mask is derived from the cross_attention_mask and is used to handle cases where a text token
+            should not attend to any image token.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, BltForCausalLM
+
+        >>> model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf")
+        >>> tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")
+
+        >>> prompt = "If I had to write a haiku, it would be:"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6)
+        >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        >>> print(result)
+        If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful.
+        I love the idea of snowflakes gently falling, each one
+        ```
+        """
         # Call parent forward but exclude cross_attention_states from model call
         outputs = self.model(
             input_ids=input_ids,
diff --git a/src/transformers/models/dia/generation_dia.py b/src/transformers/models/dia/generation_dia.py
index d87eb7879b12..5a94f081af05 100644
--- a/src/transformers/models/dia/generation_dia.py
+++ b/src/transformers/models/dia/generation_dia.py
@@ -388,26 +388,16 @@ def _main_generate_loop(
         # Prepare inner 2D logic in generation loop
         input_ids = input_ids.reshape(-1, input_ids.shape[-1])
 
-        model_kwargs = self._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
-        # prepare model inputs
-        model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
-
-        # 10. Prefill
-        model_inputs.update({"output_attentions": generation_config.output_attentions})
-        model_inputs.update({"output_hidden_states": generation_config.output_hidden_states})
-        outputs = self(**model_inputs, return_dict=True)
-
-        # 11. expand input_ids with `num_return_sequences` additional sequences per batch
+        # 10. expand input_ids with `num_return_sequences` additional sequences per batch
         if generation_config.num_return_sequences > 1:
             raise ValueError("`num_return_sequences>1` is incompatible with Dia.")
 
-        # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
+        # 11. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
         return self._sample(
             input_ids,
             logits_processor=prepared_logits_processor,
             stopping_criteria=prepared_stopping_criteria,
             generation_config=generation_config,
-            prefill_outputs=outputs,
             **generation_mode_kwargs,
             **model_kwargs,
         )

From d58e5b74ac6745806f79b8d791ea55b1c6ea2af0 Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Mon, 2 Feb 2026 16:05:10 +0100
Subject: [PATCH 5/8] fxi repo again

---
 src/transformers/models/blt/modeling_blt.py | 5 +++++
 src/transformers/models/blt/modular_blt.py  | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py
index c5e040b0f267..47113d315cbf 100644
--- a/src/transformers/models/blt/modeling_blt.py
+++ b/src/transformers/models/blt/modeling_blt.py
@@ -1372,6 +1372,11 @@ def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) ->
         return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1
 
 
+@auto_docstring(
+    custom_intro="""
+    The Blt Text Model with a language modeling head on top.
+    """
+)
 class BltForCausalLM(BltPreTrainedModel, GenerationMixin):
     config: BltConfig
     _can_compile_fullgraph = False
diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py
index bf79f0147046..4c42816e9e49 100644
--- a/src/transformers/models/blt/modular_blt.py
+++ b/src/transformers/models/blt/modular_blt.py
@@ -1097,6 +1097,11 @@ def _patch_ids_from_lengths(self, patch_lengths: torch.Tensor, seq_len: int) ->
         return (patch_starts.unsqueeze(1) <= token_positions.unsqueeze(0).unsqueeze(-1)).sum(dim=-1) - 1
 
 
+@auto_docstring(
+    custom_intro="""
+    The Blt Text Model with a language modeling head on top.
+    """
+)
 class BltForCausalLM(BltPreTrainedModel, GenerationMixin):
     config: BltConfig
     _can_compile_fullgraph = False

From 524734f82f5e07f8db3c6ef75d82c118b6bd2608 Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Tue, 3 Feb 2026 11:23:50 +0100
Subject: [PATCH 6/8] comments, comments

---
 src/transformers/cache_utils.py             | 13 +++++++++--
 src/transformers/generation/utils.py        | 25 +++++----------------
 src/transformers/models/blt/modeling_blt.py |  4 +++-
 src/transformers/models/blt/modular_blt.py  |  4 +++-
 4 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py
index a79672d76e05..cb8df79f8079 100644
--- a/src/transformers/cache_utils.py
+++ b/src/transformers/cache_utils.py
@@ -7,6 +7,7 @@
 from .configuration_utils import PreTrainedConfig
 from .utils import (
     is_hqq_available,
+    is_optimum_quanto_available,
     is_quanto_greater,
     is_torch_greater_or_equal,
     is_torchdynamo_compiling,
@@ -584,7 +585,12 @@ def __init__(
         )
 
         # We need to import quanto here to avoid circular imports due to optimum/quanto/models/transformers_models.py
-        if is_quanto_greater("0.2.5", accept_dev=True):
+        if not is_optimum_quanto_available():
+            raise ImportError(
+                "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto "
+                "backend. Please install it via  with `pip install optimum-quanto`"
+            )
+        elif is_quanto_greater("0.2.5", accept_dev=True):
             from optimum.quanto import MaxOptimizer, qint2, qint4
         else:
             raise ImportError(
@@ -634,7 +640,10 @@ def __init__(
         )
 
         if not is_hqq_available():
-            raise ImportError("You need to install `hqq` to use `HQQQuantizedLayer`")
+            raise ImportError(
+                "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. "
+                "Please install it via  with `pip install hqq`"
+            )
 
         if self.nbits not in [1, 2, 3, 4, 8]:
             raise ValueError(
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
index 4b17b9504ba2..8d6fdc01fd53 100644
--- a/src/transformers/generation/utils.py
+++ b/src/transformers/generation/utils.py
@@ -48,8 +48,6 @@
     ModelOutput,
     TransformersKwargs,
     is_accelerate_available,
-    is_hqq_available,
-    is_optimum_quanto_available,
     is_torchdynamo_exporting,
     logging,
 )
@@ -1935,8 +1933,9 @@ def _prepare_cache_for_generation(
         instantiated, writes it to `model_kwargs`, under the name expected by the model.
         """
 
-        is_linear_cache = "mamba" in self.__class__.__name__.lower()
-        cache_name = "past_key_values" if not is_linear_cache else "cache_params"
+        # TODO @raushan, unify cache arg naming for all models
+        is_linear_attn_cache = "mamba" in self.__class__.__name__.lower()
+        cache_name = "past_key_values" if not is_linear_attn_cache else "cache_params"
 
         # Quick escape route 1: if the user specifies a cache, we only need to check for conflicting `generate` arguments
         user_defined_cache = model_kwargs.get(cache_name)
@@ -1971,8 +1970,6 @@ def _prepare_cache_for_generation(
         # Assisted decoding and contrastive search require cache rollback, which is incompatible with sliding layers.
         # To handle this, we skip passing the model config to DynamicCache (forcing a full-layer cache).
         # The "dynamic_full" option is a shortcut for generate() users to avoid sliding layers on their own.
-        # TODO(joao): support static caches in assisted generation. assisted generation needs to roll back caches,
-        # which is only supported in dynamic caches atm
         if generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH):
             if generation_config.cache_implementation is not None:
                 logger.warning_once(
@@ -1991,9 +1988,9 @@ def _prepare_cache_for_generation(
         if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS:
             if generation_config.cache_implementation in DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS:
                 logger.warning_once(
-                    f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated. "
-                    f"Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, and the layer structure will be "
-                    "inferred automatically."
+                    f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated "
+                    f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, "
+                    "and the layer structure will be inferred automatically."
                 )
             model_kwargs["past_key_values"] = self._prepare_static_cache(
                 cache_implementation=generation_config.cache_implementation,
@@ -2011,16 +2008,6 @@ def _prepare_cache_for_generation(
             cache_config = generation_config.cache_config if generation_config.cache_config is not None else {}
             cache_config.setdefault("config", self.config.get_text_config(decoder=True))
             backend = cache_config.pop("backend", "quanto")
-            if backend == "quanto" and not is_optimum_quanto_available():
-                raise ImportError(
-                    "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto "
-                    "backend. Please install it via  with `pip install optimum-quanto`"
-                )
-            elif backend == "HQQ" and not is_hqq_available():
-                raise ImportError(
-                    "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. "
-                    "Please install it via  with `pip install hqq`"
-                )
             model_kwargs["past_key_values"] = QuantizedCache(backend=backend, **cache_config)
         # i.e. `cache_implementation` in [None, "dynamic", "offloaded", "dynamic_full"]
         # TODO: prepare linear cache from a single API, instead of creating in modeling code
diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py
index 47113d315cbf..b52ff48c796c 100644
--- a/src/transformers/models/blt/modeling_blt.py
+++ b/src/transformers/models/blt/modeling_blt.py
@@ -37,7 +37,7 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
 from ...utils.generic import OutputRecorder, check_model_inputs, maybe_autocast
 from .configuration_blt import (
     BltConfig,
@@ -1392,6 +1392,8 @@ def __init__(self, config: BltConfig):
 
         self.post_init()
 
+    @can_return_tuple
+    @auto_docstring
     def forward(
         self,
         input_ids: torch.LongTensor | None = None,
diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py
index 4c42816e9e49..dfa86cf1408c 100644
--- a/src/transformers/models/blt/modular_blt.py
+++ b/src/transformers/models/blt/modular_blt.py
@@ -28,7 +28,7 @@
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 from ...processing_utils import Unpack
-from ...utils import TransformersKwargs, auto_docstring, logging
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 from ...utils.generic import OutputRecorder, check_model_inputs, maybe_autocast
 from ..cohere2.modeling_cohere2 import rotate_half  # noqa: F401
 from ..llama.modeling_llama import LlamaRotaryEmbedding
@@ -1117,6 +1117,8 @@ def __init__(self, config: BltConfig):
 
         self.post_init()
 
+    @can_return_tuple
+    @auto_docstring
     def forward(
         self,
         input_ids: torch.LongTensor | None = None,

From 4493bb2ee5d56bd4e229ebd27cd1147a3e337056 Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Tue, 3 Feb 2026 11:33:29 +0100
Subject: [PATCH 7/8] almost forgot

---
 src/transformers/models/afmoe/modeling_afmoe.py | 2 +-
 src/transformers/models/afmoe/modular_afmoe.py  | 2 +-
 src/transformers/models/blt/modeling_blt.py     | 2 +-
 src/transformers/models/blt/modular_blt.py      | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py
index bb66018d583b..6b27bd9c1d3d 100644
--- a/src/transformers/models/afmoe/modeling_afmoe.py
+++ b/src/transformers/models/afmoe/modeling_afmoe.py
@@ -576,7 +576,7 @@ def forward(
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
+            past_key_values = DynamicCache(self.config)
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py
index 3bd29f904f2e..4df8b412fc58 100644
--- a/src/transformers/models/afmoe/modular_afmoe.py
+++ b/src/transformers/models/afmoe/modular_afmoe.py
@@ -397,7 +397,7 @@ def forward(
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
+            past_key_values = DynamicCache(self.config)
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py
index b52ff48c796c..c86ffa1fcadd 100644
--- a/src/transformers/models/blt/modeling_blt.py
+++ b/src/transformers/models/blt/modeling_blt.py
@@ -930,7 +930,7 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
+            past_key_values = DynamicCache(self.config)
 
         if cache_position is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py
index dfa86cf1408c..9a00fcd9d3c8 100644
--- a/src/transformers/models/blt/modular_blt.py
+++ b/src/transformers/models/blt/modular_blt.py
@@ -818,7 +818,7 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
+            past_key_values = DynamicCache(self.config)
 
         if cache_position is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0

From e10455e839dc423ea4d5af69ce8576d1ffc95517 Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Tue, 3 Feb 2026 13:01:25 +0100
Subject: [PATCH 8/8] explicit keyword arg

---
 src/transformers/models/afmoe/modeling_afmoe.py | 2 +-
 src/transformers/models/afmoe/modular_afmoe.py  | 2 +-
 src/transformers/models/blt/modeling_blt.py     | 2 +-
 src/transformers/models/blt/modular_blt.py      | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py
index 6b27bd9c1d3d..b8b03e42fce4 100644
--- a/src/transformers/models/afmoe/modeling_afmoe.py
+++ b/src/transformers/models/afmoe/modeling_afmoe.py
@@ -576,7 +576,7 @@ def forward(
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(self.config)
+            past_key_values = DynamicCache(config=self.config)
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py
index 4df8b412fc58..56ab3efa6277 100644
--- a/src/transformers/models/afmoe/modular_afmoe.py
+++ b/src/transformers/models/afmoe/modular_afmoe.py
@@ -397,7 +397,7 @@ def forward(
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(self.config)
+            past_key_values = DynamicCache(config=self.config)
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py
index c86ffa1fcadd..2f9da1f0091e 100644
--- a/src/transformers/models/blt/modeling_blt.py
+++ b/src/transformers/models/blt/modeling_blt.py
@@ -930,7 +930,7 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(self.config)
+            past_key_values = DynamicCache(config=self.config)
 
         if cache_position is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py
index 9a00fcd9d3c8..e5b721927fb1 100644
--- a/src/transformers/models/blt/modular_blt.py
+++ b/src/transformers/models/blt/modular_blt.py
@@ -818,7 +818,7 @@ def forward(
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(self.config)
+            past_key_values = DynamicCache(config=self.config)
 
         if cache_position is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0