Pine/VoxCPM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: decode chunksize for audiovae_v2

Browse Source
This commit is contained in:
Labmem-Zhouyx
2026-04-08 16:31:36 +08:00
parent 8f95d13073
commit 82d77d445c
2 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/voxcpm/model/voxcpm2.py
+9 -8
View File
@@ -246,6 +246,7 @@ class VoxCPM2Model(nn.Module):
# Audio VAE # Audio VAE
self.audio_vae = audio_vae self.audio_vae = audio_vae
self.chunk_size = audio_vae.chunk_size self.chunk_size = audio_vae.chunk_size
self._decode_chunk_size = getattr(audio_vae, "decode_chunk_size", audio_vae.chunk_size)
self._encode_sample_rate = audio_vae.sample_rate self._encode_sample_rate = audio_vae.sample_rate
self.sample_rate = getattr(audio_vae, "out_sample_rate", audio_vae.sample_rate) self.sample_rate = getattr(audio_vae, "out_sample_rate", audio_vae.sample_rate)
@@ -656,10 +657,10 @@ class VoxCPM2Model(nn.Module):
streaming_prefix_len=streaming_prefix_len, streaming_prefix_len=streaming_prefix_len,
) )
if streaming: if streaming:
patch_len = self.patch_size * self.chunk_size decode_patch_len = self.patch_size * self._decode_chunk_size
for latent_pred, _ in inference_result: for latent_pred, _ in inference_result:
decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu() decode_audio = decode_audio[..., -decode_patch_len:].squeeze(1).cpu()
yield decode_audio yield decode_audio
break break
else: else:
@@ -679,10 +680,10 @@ class VoxCPM2Model(nn.Module):
if not streaming: if not streaming:
decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
patch_len = self.patch_size * self.chunk_size decode_patch_len = self.patch_size * self._decode_chunk_size
has_continuation = bool(prompt_wav_path) has_continuation = bool(prompt_wav_path)
if has_continuation: if has_continuation:
decode_audio = decode_audio[..., patch_len * (streaming_prefix_len - 1):].squeeze(1).cpu() decode_audio = decode_audio[..., decode_patch_len * (streaming_prefix_len - 1):].squeeze(1).cpu()
else: else:
decode_audio = decode_audio.squeeze(1).cpu() decode_audio = decode_audio.squeeze(1).cpu()
yield decode_audio yield decode_audio
@@ -944,10 +945,10 @@ class VoxCPM2Model(nn.Module):
streaming_prefix_len=streaming_prefix_len, streaming_prefix_len=streaming_prefix_len,
) )
if streaming: if streaming:
patch_len = self.patch_size * self.chunk_size decode_patch_len = self.patch_size * self._decode_chunk_size
for latent_pred, pred_audio_feat in inference_result: for latent_pred, pred_audio_feat in inference_result:
decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu() decode_audio = decode_audio[..., -decode_patch_len:].squeeze(1).cpu()
yield (decode_audio, target_text_token, pred_audio_feat) yield (decode_audio, target_text_token, pred_audio_feat)
break break
else: else:
@@ -966,9 +967,9 @@ class VoxCPM2Model(nn.Module):
break break
if not streaming: if not streaming:
decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
patch_len = self.patch_size * self.chunk_size decode_patch_len = self.patch_size * self._decode_chunk_size
if mode in ("continuation", "ref_continuation"): if mode in ("continuation", "ref_continuation"):
decode_audio = decode_audio[..., patch_len * (streaming_prefix_len - 1) :].squeeze(1).cpu() decode_audio = decode_audio[..., decode_patch_len * (streaming_prefix_len - 1) :].squeeze(1).cpu()
else: else:
decode_audio = decode_audio[..., :].squeeze(1).cpu() decode_audio = decode_audio[..., :].squeeze(1).cpu()
yield (decode_audio, target_text_token, pred_audio_feat) yield (decode_audio, target_text_token, pred_audio_feat)
src/voxcpm/modules/audiovae/audio_vae_v2.py
+1
View File
@@ -436,6 +436,7 @@ class AudioVAE(nn.Module):
self.out_sample_rate = out_sample_rate self.out_sample_rate = out_sample_rate
self.sr_bin_boundaries = sr_bin_boundaries self.sr_bin_boundaries = sr_bin_boundaries
self.chunk_size = math.prod(encoder_rates) self.chunk_size = math.prod(encoder_rates)
self.decode_chunk_size = math.prod(decoder_rates)
def preprocess(self, audio_data, sample_rate): def preprocess(self, audio_data, sample_rate):
if sample_rate is None: if sample_rate is None: