Fix Turbo TTS and Multilingual TTS compatibility issues
Fixes #45 - FL Chatterbox Turbo returns error (meanflow parameter) Fixes #44 - Multilingual TTS min_p parameter compatibility Changes: - Add meanflow support to S3Gen for Turbo model (s3gen.py, flow.py, flow_matching.py, decoder.py) - Add n_cfm_timesteps, noised_mels, speech_token_lens parameters throughout inference chain - Add basic_euler solver for meanflow mode in flow_matching.py - Add min_p parameter support to T3.inference() for Multilingual TTS (t3.py) - Add MinPLogitsWarper for min_p sampling - Fix Float vs Double type errors (librosa returns float64, PyTorch expects float32) - Fix int vs float type errors in mask.py and flow.py - Add save_audio_wav() helper with fallbacks for robust audio saving - Add intmeanflow time mixer and configs for meanflow support 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Fill
+39
-7
@@ -4,6 +4,38 @@ import torchaudio
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, Any
|
||||
import scipy.io.wavfile as wavfile
|
||||
|
||||
|
||||
def save_audio_wav(filepath: str, waveform: torch.Tensor, sample_rate: int):
|
||||
"""
|
||||
Save audio to WAV file with fallback for TorchCodec issues.
|
||||
|
||||
Args:
|
||||
filepath: Path to save the audio file
|
||||
waveform: Audio tensor of shape (channels, samples) or (samples,)
|
||||
sample_rate: Sample rate in Hz
|
||||
"""
|
||||
# Ensure waveform is float32 for consistency
|
||||
waveform = waveform.float()
|
||||
|
||||
try:
|
||||
# Try torchaudio first with explicit backend
|
||||
torchaudio.save(filepath, waveform, sample_rate, backend="soundfile")
|
||||
except Exception:
|
||||
try:
|
||||
# Fallback: try without specifying backend
|
||||
torchaudio.save(filepath, waveform, sample_rate)
|
||||
except Exception:
|
||||
# Final fallback: use scipy
|
||||
audio_np = waveform.cpu().numpy().astype(np.float32)
|
||||
# scipy expects (samples, channels) for stereo, or (samples,) for mono
|
||||
if len(audio_np.shape) == 2:
|
||||
audio_np = audio_np.T # Transpose from (channels, samples) to (samples, channels)
|
||||
# Normalize to int16 range for WAV file
|
||||
audio_np = (audio_np * 32767).astype(np.int16)
|
||||
wavfile.write(filepath, sample_rate, audio_np)
|
||||
|
||||
|
||||
# Import from the local chatterbox implementation
|
||||
from .local_chatterbox.chatterbox import ChatterboxTTS
|
||||
@@ -299,7 +331,7 @@ class FL_ChatterboxTTSNode(AudioNodeBase):
|
||||
|
||||
# Save the audio prompt to the temporary file
|
||||
prompt_waveform = audio_prompt['waveform'].squeeze(0)
|
||||
torchaudio.save(audio_prompt_path, prompt_waveform, audio_prompt['sample_rate'])
|
||||
save_audio_wav(audio_prompt_path, prompt_waveform, audio_prompt['sample_rate'])
|
||||
message += f"\nUsing provided audio prompt for voice cloning: {audio_prompt_path}"
|
||||
|
||||
# Debug: Check if the file exists and has content
|
||||
@@ -499,7 +531,7 @@ class FL_ChatterboxTurboTTSNode(AudioNodeBase):
|
||||
|
||||
# Save the audio prompt to the temporary file
|
||||
prompt_waveform = audio_prompt['waveform'].squeeze(0)
|
||||
torchaudio.save(audio_prompt_path, prompt_waveform, audio_prompt['sample_rate'])
|
||||
save_audio_wav(audio_prompt_path, prompt_waveform, audio_prompt['sample_rate'])
|
||||
message += f"\nUsing provided audio prompt for voice cloning"
|
||||
|
||||
# Check audio duration (Turbo requires min 5 seconds)
|
||||
@@ -712,7 +744,7 @@ class FL_ChatterboxMultilingualTTSNode(AudioNodeBase):
|
||||
|
||||
# Save the audio prompt to the temporary file
|
||||
prompt_waveform = audio_prompt['waveform'].squeeze(0)
|
||||
torchaudio.save(audio_prompt_path, prompt_waveform, audio_prompt['sample_rate'])
|
||||
save_audio_wav(audio_prompt_path, prompt_waveform, audio_prompt['sample_rate'])
|
||||
message += f"\nUsing provided audio prompt for voice cloning"
|
||||
|
||||
# Check audio duration (Multilingual requires min 6 seconds)
|
||||
@@ -901,16 +933,16 @@ class FL_ChatterboxVCNode(AudioNodeBase):
|
||||
|
||||
# Save the input audio to the temporary file
|
||||
input_waveform = input_audio['waveform'].squeeze(0)
|
||||
torchaudio.save(input_audio_path, input_waveform, input_audio['sample_rate'])
|
||||
|
||||
save_audio_wav(input_audio_path, input_waveform, input_audio['sample_rate'])
|
||||
|
||||
# Create a temporary file for the target voice
|
||||
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_target:
|
||||
target_voice_path = temp_target.name
|
||||
temp_files.append(target_voice_path)
|
||||
|
||||
|
||||
# Save the target voice to the temporary file
|
||||
target_waveform = target_voice['waveform'].squeeze(0)
|
||||
torchaudio.save(target_voice_path, target_waveform, target_voice['sample_rate'])
|
||||
save_audio_wav(target_voice_path, target_waveform, target_voice['sample_rate'])
|
||||
|
||||
vc_model = None
|
||||
pbar = ProgressBar(100) # Simple progress bar for overall process
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
from ..utils import AttrDict
|
||||
|
||||
CFM_PARAMS = AttrDict({
|
||||
"sigma_min": 1e-06,
|
||||
"solver": "euler",
|
||||
"t_scheduler": "cosine",
|
||||
"training_cfg_rate": 0.2,
|
||||
"inference_cfg_rate": 0.7,
|
||||
"reg_loss_type": "l1"
|
||||
})
|
||||
@@ -20,6 +20,7 @@ from .utils.mask import add_optional_chunk_mask
|
||||
from .matcha.decoder import SinusoidalPosEmb, Block1D, ResnetBlock1D, Downsample1D, \
|
||||
TimestepEmbedding, Upsample1D
|
||||
from .matcha.transformer import BasicTransformerBlock
|
||||
from .utils.intmeanflow import get_intmeanflow_time_mixer
|
||||
|
||||
|
||||
def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
|
||||
@@ -110,6 +111,7 @@ class ConditionalDecoder(nn.Module):
|
||||
num_mid_blocks=12,
|
||||
num_heads=8,
|
||||
act_fn="gelu",
|
||||
meanflow=False,
|
||||
):
|
||||
"""
|
||||
This decoder requires an input with the same shape of the target. So, if your text content
|
||||
@@ -117,6 +119,7 @@ class ConditionalDecoder(nn.Module):
|
||||
"""
|
||||
super().__init__()
|
||||
channels = tuple(channels)
|
||||
self.meanflow = meanflow
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.causal = causal
|
||||
@@ -215,6 +218,14 @@ class ConditionalDecoder(nn.Module):
|
||||
self.final_block = CausalBlock1D(channels[-1], channels[-1]) if self.causal else Block1D(channels[-1], channels[-1])
|
||||
self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1)
|
||||
self.initialize_weights()
|
||||
self.time_embed_mixer = None
|
||||
if self.meanflow:
|
||||
self.time_embed_mixer = get_intmeanflow_time_mixer(time_embed_dim)
|
||||
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self.final_proj.weight.dtype
|
||||
|
||||
def initialize_weights(self):
|
||||
for m in self.modules():
|
||||
@@ -230,7 +241,7 @@ class ConditionalDecoder(nn.Module):
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x, mask, mu, t, spks=None, cond=None):
|
||||
def forward(self, x, mask, mu, t, spks=None, cond=None, r=None):
|
||||
"""Forward pass of the UNet1DConditional model.
|
||||
|
||||
Args:
|
||||
@@ -239,6 +250,7 @@ class ConditionalDecoder(nn.Module):
|
||||
t (_type_): shape (batch_size)
|
||||
spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None.
|
||||
cond (_type_, optional): placeholder for future use. Defaults to None.
|
||||
r: end time for meanflow mode (shape (1,) tensor)
|
||||
|
||||
Raises:
|
||||
ValueError: _description_
|
||||
@@ -247,10 +259,15 @@ class ConditionalDecoder(nn.Module):
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
|
||||
t = self.time_embeddings(t).to(t.dtype)
|
||||
t = self.time_mlp(t)
|
||||
|
||||
if self.meanflow:
|
||||
r = self.time_embeddings(r).to(t.dtype)
|
||||
r = self.time_mlp(r)
|
||||
concat_embed = torch.cat([t, r], dim=1)
|
||||
t = self.time_embed_mixer(concat_embed)
|
||||
|
||||
x = pack([x, mu], "b * t")[0]
|
||||
|
||||
if spks is not None:
|
||||
|
||||
@@ -20,6 +20,24 @@ from torch.nn import functional as F
|
||||
from omegaconf import DictConfig
|
||||
from .utils.mask import make_pad_mask
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _repeat_batch_dim(x, B, ndim):
|
||||
"""Repeat tensor along batch dimension if needed."""
|
||||
if x is None:
|
||||
return x
|
||||
if x.size(0) == B:
|
||||
return x
|
||||
if x.size(0) == 1:
|
||||
if ndim == 1:
|
||||
return x.repeat(B)
|
||||
elif ndim == 2:
|
||||
return x.repeat(B, 1)
|
||||
elif ndim == 3:
|
||||
return x.repeat(B, 1, 1)
|
||||
raise ValueError(f"Cannot repeat tensor of size {x.size(0)} to batch size {B}")
|
||||
|
||||
|
||||
class MaskedDiffWithXvec(torch.nn.Module):
|
||||
def __init__(self,
|
||||
@@ -202,41 +220,62 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
|
||||
prompt_feat,
|
||||
prompt_feat_len,
|
||||
embedding,
|
||||
finalize):
|
||||
if self.fp16 is True:
|
||||
prompt_feat = prompt_feat.half()
|
||||
embedding = embedding.half()
|
||||
finalize,
|
||||
n_timesteps=10,
|
||||
noised_mels=None,
|
||||
meanflow=False):
|
||||
# token: (B, n_toks)
|
||||
# token_len: (B,)
|
||||
B = token.size(0)
|
||||
|
||||
assert token.shape[0] == 1
|
||||
# xvec projection
|
||||
embedding = torch.atleast_2d(embedding)
|
||||
embedding = F.normalize(embedding, dim=1)
|
||||
embedding = self.spk_embed_affine_layer(embedding)
|
||||
embedding = self.spk_embed_affine_layer(embedding) # (1 or B, emb_dim)
|
||||
|
||||
# adjust shapes (batching logic)
|
||||
prompt_token = _repeat_batch_dim(prompt_token, B, ndim=2) # (B, n_prompt)
|
||||
prompt_token_len = _repeat_batch_dim(prompt_token_len, B, ndim=1) # (B,)
|
||||
prompt_feat = _repeat_batch_dim(prompt_feat, B, ndim=3) # (B, n_feat, feat_dim=80)
|
||||
prompt_feat_len = _repeat_batch_dim(prompt_feat_len, B, ndim=1) # (B,) or None
|
||||
embedding = _repeat_batch_dim(embedding, B, ndim=2) # (B, emb_dim)
|
||||
|
||||
# concat text and prompt_text
|
||||
token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len
|
||||
token, token_len = torch.concat([prompt_token, token], dim=1), (prompt_token_len + token_len).long()
|
||||
mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding)
|
||||
token = self.input_embedding(torch.clamp(token, min=0)) * mask
|
||||
|
||||
if (token >= self.vocab_size).any():
|
||||
logger.error(f"{token.max()}>{self.vocab_size}\n out-of-range special tokens found in flow, fix inputs!")
|
||||
token = self.input_embedding(token.long()) * mask
|
||||
|
||||
# text encode
|
||||
h, h_lengths = self.encoder(token, token_len)
|
||||
h, h_masks = self.encoder(token, token_len)
|
||||
if finalize is False:
|
||||
h = h[:, :-self.pre_lookahead_len * self.token_mel_ratio]
|
||||
|
||||
h_lengths = h_masks.sum(dim=-1).squeeze(dim=-1).long()
|
||||
mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1]
|
||||
h = self.encoder_proj(h)
|
||||
|
||||
# get conditions
|
||||
conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype)
|
||||
# # get conditions
|
||||
conds = torch.zeros([B, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype)
|
||||
conds[:, :mel_len1] = prompt_feat
|
||||
conds = conds.transpose(1, 2)
|
||||
|
||||
mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h)
|
||||
mask = (~make_pad_mask(h_lengths)).unsqueeze(1).to(h)
|
||||
|
||||
if mask.shape[0] != B:
|
||||
mask = mask.repeat(B, 1, 1)
|
||||
|
||||
feat, _ = self.decoder(
|
||||
mu=h.transpose(1, 2).contiguous(),
|
||||
mask=mask.unsqueeze(1),
|
||||
mask=mask,
|
||||
spks=embedding,
|
||||
cond=conds,
|
||||
n_timesteps=10
|
||||
n_timesteps=n_timesteps,
|
||||
noised_mels=noised_mels,
|
||||
meanflow=meanflow,
|
||||
)
|
||||
feat = feat[:, :, mel_len1:]
|
||||
assert feat.shape[2] == mel_len2
|
||||
return feat.float(), None # NOTE jrm: why are they returning None here?
|
||||
return feat, None # NOTE jrm: why are they returning None here?
|
||||
|
||||
@@ -14,10 +14,16 @@
|
||||
import threading
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from tqdm import tqdm
|
||||
from .matcha.flow_matching import BASECFM
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
|
||||
def cast_all(*tensors, dtype):
|
||||
"""Cast all tensors to specified dtype."""
|
||||
return [t.to(dtype) if t is not None else None for t in tensors]
|
||||
|
||||
|
||||
CFM_PARAMS = OmegaConf.create({
|
||||
"sigma_min": 1e-06,
|
||||
"solver": "euler",
|
||||
@@ -198,10 +204,11 @@ class ConditionalCFM(BASECFM):
|
||||
class CausalConditionalCFM(ConditionalCFM):
|
||||
def __init__(self, in_channels=240, cfm_params=CFM_PARAMS, n_spks=1, spk_emb_dim=80, estimator=None):
|
||||
super().__init__(in_channels, cfm_params, n_spks, spk_emb_dim, estimator)
|
||||
self.rand_noise = torch.randn([1, 80, 50 * 300])
|
||||
# TODO: BAD BAD IDEA - IT'LL MESS UP DISTILLATION - SETTING TO NONE
|
||||
self.rand_noise = None
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None):
|
||||
def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, noised_mels=None, meanflow=False):
|
||||
"""Forward diffusion
|
||||
|
||||
Args:
|
||||
@@ -214,15 +221,43 @@ class CausalConditionalCFM(ConditionalCFM):
|
||||
spks (torch.Tensor, optional): speaker ids. Defaults to None.
|
||||
shape: (batch_size, spk_emb_dim)
|
||||
cond: Not used but kept for future purposes
|
||||
noised_mels: gt mels noised at time t
|
||||
meanflow: whether to use meanflow mode
|
||||
|
||||
Returns:
|
||||
sample: generated mel-spectrogram
|
||||
shape: (batch_size, n_feats, mel_timesteps)
|
||||
"""
|
||||
|
||||
z = self.rand_noise[:, :, :mu.size(2)].to(mu.device).to(mu.dtype) * temperature
|
||||
# fix prompt and overlap part mu and z
|
||||
B = mu.size(0)
|
||||
z = torch.randn_like(mu)
|
||||
|
||||
if noised_mels is not None:
|
||||
prompt_len = mu.size(2) - noised_mels.size(2)
|
||||
z[..., prompt_len:] = noised_mels
|
||||
|
||||
# time steps for reverse diffusion
|
||||
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
|
||||
if self.t_scheduler == 'cosine':
|
||||
if (not meanflow) and (self.t_scheduler == 'cosine'):
|
||||
t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
|
||||
|
||||
# NOTE: right now, the only meanflow models are also distilled models, which don't need CFG
|
||||
# because they were distilled with CFG outputs. We would need to add another hparam and
|
||||
# change the conditional logic here if we want to use CFG inference with a meanflow model.
|
||||
if meanflow:
|
||||
return self.basic_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None
|
||||
|
||||
return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None
|
||||
|
||||
def basic_euler(self, x, t_span, mu, mask, spks, cond):
|
||||
in_dtype = x.dtype
|
||||
x, t_span, mu, mask, spks, cond = cast_all(x, t_span, mu, mask, spks, cond, dtype=self.estimator.dtype)
|
||||
|
||||
print("S3 Token -> Mel Inference...")
|
||||
for t, r in tqdm(zip(t_span[..., :-1], t_span[..., 1:]), total=t_span.shape[-1] - 1):
|
||||
t, r = t[None], r[None]
|
||||
dxdt = self.estimator.forward(x, mask=mask, mu=mu, t=t, spks=spks, cond=cond, r=r)
|
||||
dt = r - t
|
||||
x = x + dt * dxdt
|
||||
|
||||
return x.to(in_dtype)
|
||||
|
||||
@@ -19,7 +19,6 @@ import torch
|
||||
import torchaudio as ta
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from ..s3tokenizer import S3_SR, SPEECH_VOCAB_SIZE, S3Tokenizer
|
||||
from .const import S3GEN_SR
|
||||
@@ -31,6 +30,7 @@ from .hifigan import HiFTGenerator
|
||||
from .transformer.upsample_encoder import UpsampleConformerEncoder
|
||||
from .flow_matching import CausalConditionalCFM
|
||||
from .decoder import ConditionalDecoder
|
||||
from .configs import CFM_PARAMS
|
||||
|
||||
|
||||
def drop_invalid_tokens(x):
|
||||
@@ -46,15 +46,20 @@ def get_resampler(src_sr, dst_sr, device):
|
||||
|
||||
class S3Token2Mel(torch.nn.Module):
|
||||
"""
|
||||
CosyVoice2's CFM decoder maps S3 speech tokens to mel-spectrograms.
|
||||
S3Gen's CFM decoder maps S3 speech tokens to mel-spectrograms.
|
||||
|
||||
TODO: make these modules configurable?
|
||||
"""
|
||||
def __init__(self):
|
||||
def __init__(self, meanflow=False):
|
||||
super().__init__()
|
||||
self.tokenizer = S3Tokenizer("speech_tokenizer_v2_25hz")
|
||||
self.mel_extractor = mel_spectrogram # TODO: make it a torch module?
|
||||
self.speaker_encoder = CAMPPlus() # use default args
|
||||
self.speaker_encoder = CAMPPlus(
|
||||
# NOTE: This doesn't affect inference. It turns off activation checkpointing
|
||||
# (a training optimization), which causes a crazy DDP error with accelerate
|
||||
memory_efficient=False,
|
||||
)
|
||||
self.meanflow = meanflow
|
||||
|
||||
encoder = UpsampleConformerEncoder(
|
||||
output_size=512,
|
||||
@@ -84,15 +89,9 @@ class S3Token2Mel(torch.nn.Module):
|
||||
num_mid_blocks=12,
|
||||
num_heads=8,
|
||||
act_fn='gelu',
|
||||
meanflow=self.meanflow,
|
||||
)
|
||||
cfm_params = DictConfig({
|
||||
"sigma_min": 1e-06,
|
||||
"solver": 'euler',
|
||||
"t_scheduler": 'cosine',
|
||||
"training_cfg_rate": 0.2,
|
||||
"inference_cfg_rate": 0.7,
|
||||
"reg_loss_type": 'l1',
|
||||
})
|
||||
cfm_params = CFM_PARAMS
|
||||
decoder = CausalConditionalCFM(
|
||||
spk_emb_dim=80,
|
||||
cfm_params=cfm_params,
|
||||
@@ -111,6 +110,11 @@ class S3Token2Mel(torch.nn.Module):
|
||||
params = self.tokenizer.parameters()
|
||||
return next(params).device
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
params = self.flow.parameters()
|
||||
return next(params).dtype
|
||||
|
||||
def embed_ref(
|
||||
self,
|
||||
ref_wav: torch.Tensor,
|
||||
@@ -171,7 +175,10 @@ class S3Token2Mel(torch.nn.Module):
|
||||
ref_sr: Optional[int],
|
||||
# pre-computed ref embedding (prod API)
|
||||
ref_dict: Optional[dict] = None,
|
||||
n_cfm_timesteps = None,
|
||||
finalize: bool = False,
|
||||
speech_token_lens=None,
|
||||
noised_mels=None,
|
||||
):
|
||||
"""
|
||||
Generate waveforms from S3 speech tokens and a reference waveform, which the speaker timbre is inferred from.
|
||||
@@ -199,18 +206,21 @@ class S3Token2Mel(torch.nn.Module):
|
||||
if isinstance(ref_dict[rk], np.ndarray):
|
||||
ref_dict[rk] = torch.from_numpy(ref_dict[rk])
|
||||
if torch.is_tensor(ref_dict[rk]):
|
||||
ref_dict[rk] = ref_dict[rk].to(self.device)
|
||||
ref_dict[rk] = ref_dict[rk].to(device=self.device, dtype=self.dtype)
|
||||
|
||||
if len(speech_tokens.shape) == 1:
|
||||
speech_tokens = speech_tokens.unsqueeze(0)
|
||||
speech_tokens = torch.atleast_2d(speech_tokens)
|
||||
|
||||
# assert speech_tokens.shape[0] == 1, "only batch size of one allowed for now"
|
||||
speech_token_lens = torch.LongTensor([speech_tokens.size(1)]).to(self.device)
|
||||
# backcompat
|
||||
if speech_token_lens is None:
|
||||
speech_token_lens = torch.LongTensor([st.size(-1) for st in speech_tokens]).to(self.device)
|
||||
|
||||
output_mels, _ = self.flow.inference(
|
||||
token=speech_tokens,
|
||||
token_len=speech_token_lens,
|
||||
finalize=finalize,
|
||||
noised_mels=noised_mels,
|
||||
n_timesteps=n_cfm_timesteps,
|
||||
meanflow=self.meanflow,
|
||||
**ref_dict,
|
||||
)
|
||||
return output_mels
|
||||
@@ -218,13 +228,15 @@ class S3Token2Mel(torch.nn.Module):
|
||||
|
||||
class S3Token2Wav(S3Token2Mel):
|
||||
"""
|
||||
The decoder of CosyVoice2 is a concat of token-to-mel (CFM) and a mel-to-waveform (HiFiGAN) modules.
|
||||
The decoder of S3Gen is a concat of token-to-mel (CFM) and a mel-to-waveform (HiFiGAN) modules.
|
||||
|
||||
TODO: make these modules configurable?
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
ignore_state_dict_missing = ("tokenizer._mel_filters", "tokenizer.window")
|
||||
|
||||
def __init__(self, meanflow=False):
|
||||
super().__init__(meanflow)
|
||||
|
||||
f0_predictor = ConvRNNF0Predictor()
|
||||
self.mel2wav = HiFTGenerator(
|
||||
@@ -241,6 +253,7 @@ class S3Token2Wav(S3Token2Mel):
|
||||
trim_fade = torch.zeros(2 * n_trim)
|
||||
trim_fade[n_trim:] = (torch.cos(torch.linspace(torch.pi, 0, n_trim)) + 1) / 2
|
||||
self.register_buffer("trim_fade", trim_fade, persistent=False) # (buffers get automatic device casting)
|
||||
self.estimator_dtype = "fp32"
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -250,9 +263,24 @@ class S3Token2Wav(S3Token2Mel):
|
||||
ref_sr: Optional[int],
|
||||
# pre-computed ref embedding (prod API)
|
||||
ref_dict: Optional[dict] = None,
|
||||
finalize: bool = False
|
||||
finalize: bool = False,
|
||||
speech_token_lens=None,
|
||||
skip_vocoder=False,
|
||||
n_cfm_timesteps=None,
|
||||
noised_mels=None,
|
||||
):
|
||||
output_mels = super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
|
||||
"""
|
||||
Generate waveforms from S3 speech tokens and a reference waveform, which the speaker timbre is inferred from.
|
||||
NOTE: used for sync synthesis only. Please use `S3GenStreamer` for streaming synthesis.
|
||||
"""
|
||||
output_mels = super().forward(
|
||||
speech_tokens, speech_token_lens=speech_token_lens, ref_wav=ref_wav,
|
||||
ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize,
|
||||
n_cfm_timesteps=n_cfm_timesteps, noised_mels=noised_mels,
|
||||
)
|
||||
|
||||
if skip_vocoder:
|
||||
return output_mels
|
||||
|
||||
# TODO jrm: ignoring the speed control (mel interpolation) and the HiFTGAN caching mechanisms for now.
|
||||
hift_cache_source = torch.zeros(1, 1, 0).to(self.device)
|
||||
@@ -274,14 +302,24 @@ class S3Token2Wav(S3Token2Mel):
|
||||
ref_sr: Optional[int] = None,
|
||||
# pre-computed ref embedding (prod API)
|
||||
ref_dict: Optional[dict] = None,
|
||||
n_cfm_timesteps = None,
|
||||
finalize: bool = False,
|
||||
speech_token_lens=None,
|
||||
):
|
||||
return super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
|
||||
n_cfm_timesteps = n_cfm_timesteps or (2 if self.meanflow else 10)
|
||||
noise = None
|
||||
if self.meanflow:
|
||||
noise = torch.randn(1, 80, speech_tokens.size(-1) * 2, dtype=self.dtype, device=self.device)
|
||||
output_mels = super().forward(
|
||||
speech_tokens, speech_token_lens=speech_token_lens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict,
|
||||
n_cfm_timesteps=n_cfm_timesteps, finalize=finalize, noised_mels=noise,
|
||||
)
|
||||
return output_mels
|
||||
|
||||
@torch.inference_mode()
|
||||
def hift_inference(self, speech_feat, cache_source: torch.Tensor = None):
|
||||
if cache_source is None:
|
||||
cache_source = torch.zeros(1, 1, 0).to(self.device)
|
||||
cache_source = torch.zeros(1, 1, 0).to(device=self.device, dtype=self.dtype)
|
||||
return self.mel2wav.inference(speech_feat=speech_feat, cache_source=cache_source)
|
||||
|
||||
@torch.inference_mode()
|
||||
@@ -293,11 +331,22 @@ class S3Token2Wav(S3Token2Mel):
|
||||
ref_sr: Optional[int] = None,
|
||||
# pre-computed ref embedding (prod API)
|
||||
ref_dict: Optional[dict] = None,
|
||||
cache_source: torch.Tensor = None, # NOTE: this arg is for streaming, it can probably be removed here
|
||||
finalize: bool = True,
|
||||
# left as a kwarg because this can change input/output size ratio
|
||||
drop_invalid_tokens=True,
|
||||
n_cfm_timesteps=None,
|
||||
speech_token_lens=None,
|
||||
):
|
||||
output_mels = self.flow_inference(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
|
||||
output_wavs, output_sources = self.hift_inference(output_mels, cache_source)
|
||||
output_mels = self.flow_inference(
|
||||
speech_tokens,
|
||||
speech_token_lens=speech_token_lens,
|
||||
ref_wav=ref_wav,
|
||||
ref_sr=ref_sr,
|
||||
ref_dict=ref_dict,
|
||||
n_cfm_timesteps=n_cfm_timesteps,
|
||||
finalize=True,
|
||||
)
|
||||
output_mels = output_mels.to(dtype=self.dtype) # FIXME (fp16 mode) is this still needed?
|
||||
output_wavs, output_sources = self.hift_inference(output_mels, None)
|
||||
|
||||
# NOTE: ad-hoc method to reduce "spillover" from the reference clip.
|
||||
output_wavs[:, :len(self.trim_fade)] *= self.trim_fade
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def get_intmeanflow_time_mixer(dims):
|
||||
""""
|
||||
Diagonal init as described in 3.3 https://arxiv.org/pdf/2510.07979
|
||||
"""
|
||||
layer = nn.Linear(dims * 2, dims, bias=False)
|
||||
|
||||
with torch.no_grad():
|
||||
target_weight = torch.zeros(dims, 2 * dims)
|
||||
target_weight[:, 0:dims] = torch.eye(dims)
|
||||
layer.weight.data = target_weight
|
||||
|
||||
return layer
|
||||
@@ -182,7 +182,7 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
|
||||
[0, 0, 1, 1, 1]]
|
||||
"""
|
||||
batch_size = lengths.size(0)
|
||||
max_len = max_len if max_len > 0 else lengths.max().item()
|
||||
max_len = max_len if max_len > 0 else int(lengths.max().item())
|
||||
seq_range = torch.arange(0,
|
||||
max_len,
|
||||
dtype=torch.int64,
|
||||
|
||||
@@ -14,6 +14,7 @@ from transformers.generation.logits_process import (
|
||||
RepetitionPenaltyLogitsProcessor,
|
||||
TemperatureLogitsWarper,
|
||||
TopKLogitsWarper,
|
||||
MinPLogitsWarper,
|
||||
)
|
||||
|
||||
from .modules.learned_pos_emb import LearnedPositionEmbeddings
|
||||
@@ -247,10 +248,11 @@ class T3(nn.Module):
|
||||
stop_on_eos=True,
|
||||
do_sample=True,
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
top_p=0.95,
|
||||
min_p=0.05,
|
||||
length_penalty=1.0,
|
||||
repetition_penalty=2.0,
|
||||
cfg_weight=0,
|
||||
repetition_penalty=1.2,
|
||||
cfg_weight=0.5,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -334,7 +336,8 @@ class T3(nn.Module):
|
||||
|
||||
# Instantiate the logits processors.
|
||||
top_p_warper = TopPLogitsWarper(top_p=top_p)
|
||||
repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)
|
||||
min_p_warper = MinPLogitsWarper(min_p=min_p)
|
||||
repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=float(repetition_penalty))
|
||||
|
||||
# ---- Initial Forward Pass (no kv_cache yet) ----
|
||||
output = self.patched_model(
|
||||
@@ -364,9 +367,11 @@ class T3(nn.Module):
|
||||
if temperature != 1.0:
|
||||
logits = logits / temperature
|
||||
|
||||
# Apply repetition penalty and top‑p filtering.
|
||||
logits = repetition_penalty_processor(generated_ids, logits)
|
||||
logits = top_p_warper(None, logits)
|
||||
# Apply repetition penalty, min_p, and top_p filtering.
|
||||
ids_for_proc = generated_ids[:1, ...] # batch = 1
|
||||
logits = repetition_penalty_processor(ids_for_proc, logits)
|
||||
logits = min_p_warper(ids_for_proc, logits)
|
||||
logits = top_p_warper(ids_for_proc, logits)
|
||||
|
||||
# Convert logits to probabilities and sample the next token.
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
class AttrDict(dict):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(AttrDict, self).__init__(*args, **kwargs)
|
||||
self.__dict__ = self
|
||||
@@ -257,14 +257,17 @@ class VoiceEncoder(nn.Module):
|
||||
|
||||
:param trim_top_db: this argument was only added for the sake of compatibility with metavoice's implementation
|
||||
"""
|
||||
# Ensure all wavs are float32 (librosa operations can return float64)
|
||||
wavs = [wav.astype(np.float32) if wav.dtype != np.float32 else wav for wav in wavs]
|
||||
|
||||
if sample_rate != self.hp.sample_rate:
|
||||
wavs = [
|
||||
librosa.resample(wav, orig_sr=sample_rate, target_sr=self.hp.sample_rate, res_type="kaiser_fast")
|
||||
librosa.resample(wav, orig_sr=sample_rate, target_sr=self.hp.sample_rate, res_type="kaiser_fast").astype(np.float32)
|
||||
for wav in wavs
|
||||
]
|
||||
|
||||
if trim_top_db:
|
||||
wavs = [librosa.effects.trim(wav, top_db=trim_top_db)[0] for wav in wavs]
|
||||
wavs = [librosa.effects.trim(wav, top_db=trim_top_db)[0].astype(np.float32) for wav in wavs]
|
||||
|
||||
if "rate" not in kwargs:
|
||||
kwargs["rate"] = 1.3 # Resemble's default value.
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import librosa
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
@@ -192,6 +193,8 @@ class ChatterboxTTS:
|
||||
def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
|
||||
## Load reference wav
|
||||
s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
|
||||
# Ensure float32 for PyTorch compatibility
|
||||
s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
|
||||
|
||||
ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import logging
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import librosa
|
||||
import torch
|
||||
from safetensors.torch import load_file
|
||||
@@ -223,13 +224,17 @@ class ChatterboxTurboTTS:
|
||||
def prepare_conditionals(self, wav_fpath, exaggeration=0.5, norm_loudness=True):
|
||||
"""Prepare voice conditionals from reference audio."""
|
||||
s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
|
||||
# Ensure float32 for PyTorch compatibility
|
||||
s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
|
||||
|
||||
assert len(s3gen_ref_wav) / _sr > 5.0, "Audio prompt must be longer than 5 seconds!"
|
||||
|
||||
if norm_loudness:
|
||||
s3gen_ref_wav = self.norm_loudness(s3gen_ref_wav, _sr)
|
||||
# Ensure float32 after loudness normalization (multiplication can promote to float64)
|
||||
s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
|
||||
|
||||
ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
|
||||
ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR).astype(np.float32)
|
||||
|
||||
s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
|
||||
s3gen_ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import librosa
|
||||
import torch
|
||||
|
||||
@@ -85,6 +86,8 @@ class ChatterboxVC:
|
||||
def set_target_voice(self, wav_fpath):
|
||||
## Load reference wav
|
||||
s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
|
||||
# Ensure float32 for PyTorch compatibility
|
||||
s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
|
||||
|
||||
s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
|
||||
self.ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)
|
||||
|
||||
Reference in New Issue
Block a user