Error with Kyutai Speech-To-Text inference
Author: alien-0119
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.57.3
- Platform: Linux-4.15.0-213-generic-x86_64-with-glibc2.27
- Python version: 3.10.17
- Huggingface_hub version: 0.36.0
- Safetensors version: 0.6.2
- Accelerate version: 1.10.1
- Accelerate config: 	- compute_environment: LOCAL_MACHINE
	- distributed_type: NO
	- mixed_precision: fp16
	- use_cpu: False
	- debug: False
	- num_processes: 1
	- machine_rank: 0
	- num_machines: 1
	- gpu_ids: 0
	- rdzv_backend: static
	- same_network: True
	- main_training_function: main
	- enable_cpu_affinity: False
	- downcast_bf16: no
	- tpu_use_cluster: False
	- tpu_use_sudo: False
	- tpu_env: []
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.6.0+cu124 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-80GB


### Who can help?

@eustlb  @ebezzam  @vasqu 

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
import torch
from datasets import load_dataset, Audio
from transformers import infer_device, KyutaiSpeechToTextProcessor, KyutaiSpeechToTextForConditionalGeneration

# 1. load the model and the processor
torch_device = infer_device()
model_id = "kyutai/stt-2.6b-en-trfs"

processor = KyutaiSpeechToTextProcessor.from_pretrained(model_id)
model = KyutaiSpeechToTextForConditionalGeneration.from_pretrained(model_id, device_map=torch_device, dtype="auto")

# 2. load audio samples
ds = load_dataset(
    "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation"
)
ds = ds.cast_column("audio", Audio(sampling_rate=24000))

# 3. prepare the model inputs
inputs = processor(
    ds[0]["audio"]["array"],
)
inputs.to(model.device)

# 4. infer the model
output_tokens = model.generate(**inputs)

# 5. decode the generated tokens
print(processor.batch_decode(output_tokens, skip_special_tokens=True))
```


**ERROR**
```
Traceback (most recent call last):
  File "/srv/mindone/wty/test.py", line 25, in <module>
    output_tokens = model.generate(**inputs)
  File "/home/miniconda3/envs/vllm_wty/lib/python3.10/site-packages/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py", line 1352, in generate
    max_audio_frames = input_values.shape[-1] // self.config.codec_config.frame_size
AttributeError: 'NoneType' object has no attribute 'shape'
```

### Expected behavior

It's OK with transformers v4.54 but not OK with v4.57

Comments:
Thanks for the bug report! Can you do a `git bisect` to find where the issue started?

Hi i think I've found the fix .
In the `forward` method for `KyutaiSpeechToTextEmbeddings` we can just add a small conditional to distinguish between text tokens and audio codes .
```python
if input_ids.dim() == 2:
            return self.embed_tokens(input_ids)

```

And another in `generate` method for `KyutaiSpeechToTextForConditionalGeneration`
```python 
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)
        if "input_values" in kwargs:
            input_values = kwargs.get("input_values")
        elif len(args) > 0:
            input_values = args[0]
        elif "inputs" in kwargs:
            input_values = kwargs.get("inputs")
        else:
            input_values = None
        if input_values is not None:
            # 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
```
It does feel like a little redundant so I would appreciate advice on how to make cleaner .
Is this the expected fix @ArthurZucker ? 


output for the testing script 

```
The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization.
The tokenizer class you load from this checkpoint is 'PreTrainedTokenizerFast'.
The class this function is called from is 'TokenizersBackend'.
Fetching 2 files: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 21129.99it/s]
Download complete: : 0.00B [00:00, ?B/s]                                                                                               | 0/2 [00:00<?, ?it/s]
Loading weights: 100%|█████████████████████████████████████████████████████████████| 737/737 [00:03<00:00, 189.30it/s, Materializing param=model.norm.weight]
Using `cache_implementation='sliding_window' is deprecated. Please only use one of ('static', 'offloaded_static'), and the layer structure will be inferred automatically.
['Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.']
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.