Moshi Generation Does Not Work as Expected
Author: SeungyounShin
Pull request: no
Labels: name: bug

### System Info

**🐛 Bug Report**

### Description
The provided Moshi example code does not function correctly with the Transformers library. The `generate` function fails when attempting to generate new tokens, and an issue arises with the expected input formats.


And here is `moshi_output.wav`

https://github.com/user-attachments/assets/191d8176-a846-4b72-8b7c-de5c15b8140b

I tried different temperature settings, generation configurations, and other samples, but it only produces a static 'chijijik...' sound.

cc. @ylacombe 

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

```python
from datasets import load_dataset, Audio
import torch, math
from transformers import MoshiForConditionalGeneration, AutoFeatureExtractor, AutoTokenizer
import soundfile as sf
import torch
import transformers
import os
import torch

# Disable all automatic compilation features
os.environ['TORCH_COMPILE'] = '0'
os.environ['TORCHDYNAMO_DISABLE'] = '1'  # Fully disables TorchDynamo
os.environ['TORCHDYNAMO_VERBOSE'] = '0'  # Suppresses unnecessary logs
os.environ['TORCHDYNAMO_RECOMPILE_LIMIT'] = '0'  # Avoid recompile limits

# Apply global config settings for eager mode
torch._dynamo.config.suppress_errors = True  # Avoids crashes and falls back to eager mode
torch._dynamo.config.cache_size_limit = 0  # Prevents recompilation limits
torch._dynamo.reset()  # Clears any cached compile traces


librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")

device = "cuda"
# prepare user input audio 
librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=24000))
audio_sample = librispeech_dummy[-1]["audio"]["array"] # (107520,)
# WAV_PATH = f"./audio/moshi_opening.wav"
# audio_sample, sample_rate = sf.read(WAV_PATH)
waveform_to_token_ratio = 1 / 1920

model = MoshiForConditionalGeneration.from_pretrained("kmhf/hf-moshiko", attn_implementation="eager", torch_dtype=torch.float16)
feature_extractor = AutoFeatureExtractor.from_pretrained("kmhf/hf-moshiko")
tokenizer = AutoTokenizer.from_pretrained("kmhf/hf-moshiko")
model = model.to(device)

user_input_values = feature_extractor(raw_audio=audio_sample, sampling_rate=24000, return_tensors="pt").to(device=device, dtype=torch.float16)

# prepare moshi input values - we suppose moshi didn't say anything while the user spoke
moshi_input_values = torch.zeros_like(user_input_values.input_values) # (1, 1, 107520)

# prepare moshi input ids - we suppose moshi didn't say anything while the user spoke
num_tokens = math.ceil(moshi_input_values.shape[-1] * waveform_to_token_ratio)
input_ids = torch.ones((1, num_tokens), device=device, dtype=torch.int64) * tokenizer.encode("<pad>")[0]

# Force disable torch.compile inside Transformers
transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.forward = torch._dynamo.disable(
    transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.forward
)
transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.generate = torch._dynamo.disable(
    transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.generate
)
transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.prepare_inputs_for_generation = torch._dynamo.disable(
    transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.prepare_inputs_for_generation
)

# generate 25 new tokens (around 2s of audio)
output = model.generate(
    input_ids=input_ids,
    user_input_values=user_input_values.input_values,
    moshi_input_values=moshi_input_values,
    max_new_tokens=50,
    temperature=0.8,
    do_sample=True,
)

text_tokens = output.sequences
# decode text tokens
text = tokenizer.decode(text_tokens[0], skip_special_tokens=True)
print(text)

# decode audio tokens
audio_waveforms = output.audio_sequences.squeeze(0).squeeze(0) # (L,)
audio_waveforms = audio_waveforms.double()

# cut audio for input length
audio_waveforms = audio_waveforms[:user_input_values.input_values.shape[-1]]

# save audio
sf.write("moshi_output.wav", audio_waveforms.cpu().numpy(), 24000)
```

### Expected behavior

should produce sounds

Comments:
Another approach  (using codes directly)

```python
from datasets import load_dataset, Audio
import torch, math
from transformers import MoshiForConditionalGeneration, AutoFeatureExtractor, AutoTokenizer, MimiModel
import soundfile as sf
import torch
import transformers
import os
import torch

# Disable all automatic compilation features
os.environ['TORCH_COMPILE'] = '0'
os.environ['TORCHDYNAMO_DISABLE'] = '1'  # Fully disables TorchDynamo
os.environ['TORCHDYNAMO_VERBOSE'] = '0'  # Suppresses unnecessary logs
os.environ['TORCHDYNAMO_RECOMPILE_LIMIT'] = '0'  # Avoid recompile limits

# Apply global config settings for eager mode
torch._dynamo.config.suppress_errors = True  # Avoids crashes and falls back to eager mode
torch._dynamo.config.cache_size_limit = 0  # Prevents recompilation limits
torch._dynamo.reset()  # Clears any cached compile traces


librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")

device = "cuda"
# prepare user input audio 
librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=24000))
audio_sample = librispeech_dummy[-1]["audio"]["array"] # (107520,)
# WAV_PATH = f"./audio/moshi_opening.wav"
# audio_sample, sample_rate = sf.read(WAV_PATH)
waveform_to_token_ratio = 1 / 1920

model = MoshiForConditionalGeneration.from_pretrained("kmhf/hf-moshiko", attn_implementation="eager", torch_dtype=torch.float16)
feature_extractor = AutoFeatureExtractor.from_pretrained("kmhf/hf-moshiko")
tokenizer = AutoTokenizer.from_pretrained("kmhf/hf-moshiko")
model = model.to(device)

user_input_values = feature_extractor(raw_audio=audio_sample, sampling_rate=24000, return_tensors="pt").to(device=device, dtype=torch.float16)

# prepare moshi input values - we suppose moshi didn't say anything while the user spoke
moshi_input_values = torch.zeros_like(user_input_values.input_values) # (1, 1, 107520)

# prepare moshi input ids - we suppose moshi didn't say anything while the user spoke
num_tokens = math.ceil(moshi_input_values.shape[-1] * waveform_to_token_ratio)
input_ids = torch.ones((1, num_tokens), device=device, dtype=torch.int64) * tokenizer.encode("<pad>")[0]

# Force disable torch.compile inside Transformers
transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.forward = torch._dynamo.disable(
    transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.forward
)
transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.generate = torch._dynamo.disable(
    transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.generate
)
transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.prepare_inputs_for_generation = torch._dynamo.disable(
    transformers.models.moshi.modeling_moshi.MoshiForConditionalGeneration.prepare_inputs_for_generation
)


input_ids = torch.tensor([[3, 3, 3, 3, 3, 0, 2295, 3, 667, 261, 3, 3, 3, 3, 3, 3, 3, 3, 3]]).to(device)
input_ids = torch.tensor([[3, 3, 3, 3, 3, 0, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]]).to(device)

# (batch_size, num_codebooks, sequence_length)
user_audio_codes = torch.tensor([
    [[1049],[1520],[1502],[1562],[1168],[1016],[ 976],[2008]],
    [[127],[441],[1873],[317],[1622],[757],[ 1686],[1621]],
    [[1272],[724],[665],[1417],[896],[1331],[ 1926],[1451]],
    [[962],[773],[720],[1869],[1944],[720],[ 1223],[467]],
    [[1707],[773],[1606],[966],[700],[1467],[ 886],[1349]],
    [[1707],[773],[1515],[1869],[38],[1878],[ 815],[1146]],
    [[1216],[677],[230],[1015],[1924],[1582],[ 595],[605]],
    [[1942],[1732],[440],[714],[1324],[1568],[ 41],[1938]],
    [[2005],[238],[739],[1381],[27],[87],[ 13],[959]],
    [[1293],[1628],[1744],[1263],[1051],[1686],[ 1007],[94]],
    [[70],[424],[683],[1552],[869],[19],[ 634],[1624]],
    [[1164],[1393],[71],[159],[1166],[1679],[ 1235],[919]],
    [[1957],[287],[370],[829],[928],[210],[1943],[ 144]],
    [[502],[403],[1626],[164],[1736],[1572],[ 976],[1551]],
    [[769],[243],[1149],[290],[481],[1030],[ 1238],[568]],
    [[1743],[243],[1559],[1348],[340],[347],[ 1238],[1744]],
    [[716],[1056],[1502],[1712],[306],[1030],[ 1238],[1744]],
    [[1544],[1056],[1149],[164],[1736],[1572],[ 1238],[1744]],
    [[1544],[1056],[1178],[1348],[340],[1443],[ 1238],[1744]],
]).to(device).permute(2,1,0)
moshi_audio_codes = torch.tensor([
    [[1049],[1700],[1626],[546],[306],[ 1443],[1871],[2008]],
    [[127],[243],[457],[290],[1019],[1030],[ 428],[152]],
    [[1880],[91],[1029],[390],[1335],[1569],[ 192],[424]],
    [[972],[91],[457],[164],[1809],[478],[ 1044],[1744]],
    [[972],[91],[457],[290],[1335],[1447],[ 428],[1053]],
    [[1766],[1783],[1626],[400],[1317],[332],[ 1433],[1439]],
    [[1853],[398],[1089],[325],[1246],[1256],[ 922],[514]], # Good
    [[1853],[161],[1762],[1589],[484],[1162],[ 273],[166]],
    [[662],[1919],[380],[610],[204],[39],[ 390],[391]],
    [[966],[833],[1832],[425],[1466],[ 1575],[2003],[245]], # day
    [[1575],[1068],[1588],[264],[92],[854],[ 1765],[2007]], # ,
    [[198],[168],[928],[229],[266],[612],[ 1164],[238]],
    [[198],[1635],[1653],[1709],[757],[889],[ 286],[410]],
    [[582],[669],[480],[826],[1335],[1707],[ 849],[1084]],
    [[978],[811],[1178],[164],[1217],[347],[ 1543],[1511]],
    [[729],[1056],[457],[290],[1335],[1572],[ 1433],[1744]],
    [[940],[811],[1178],[164],[1314],[347],[ 43],[1648]],
    [[940],[91],[1029],[546],[1335],[347],[ 1433],[424]],
    [[1098],[1519],[1697],[390],[1836],[1443],[ 1234],[2008]],
]).to(device).permute(2,1,0)

# generate 25 new tokens (around 2s of audio)
output = model.generate(
    input_ids=input_ids,
    user_audio_codes=user_audio_codes,
    moshi_audio_codes=moshi_audio_codes,
    max_new_tokens=50,
    temperature=1.0,
    do_sample=True,
)

text_tokens = output.sequences
# decode text tokens
text = tokenizer.decode(text_tokens[0], skip_special_tokens=False)
print(text)

# decode audio tokens
audio_waveforms = output.audio_sequences.squeeze(0).squeeze(0) # (L,)
audio_waveforms = audio_waveforms.double()

# save audio
sf.write("moshi_output_2.wav", audio_waveforms.cpu().numpy(), 24000)

# save user audio
mimi_model = MimiModel.from_pretrained("kyutai/mimi")
mimi_model = mimi_model.to(device)
user_audio_values = mimi_model.decode(user_audio_codes)
sf.write("moshi_user_audio_2.wav",  user_audio_values.audio_values.detach().cpu().squeeze().numpy(), 24000)

import pdb; pdb.set_trace()

```

This acoustic codes work on original [repo](https://github.com/kyutai-labs/moshi) but not work for `transformers`