File size: 10,404 Bytes
94336a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 | import os
import torch
import soundfile as sf
from transformers import TrainerCallback
from safetensors.torch import load_file
from src.chatterbox_.tts import ChatterboxTTS
from src.chatterbox_.tts_turbo import ChatterboxTurboTTS
from src.chatterbox_.models.t3.t3 import T3
from src.model import resize_and_load_t3_weights
from src.utils import setup_logger, trim_silence_with_vad
logger = setup_logger("InferenceCallback")
class InferenceCallback(TrainerCallback):
def __init__(self, config):
self.config = config
self.inference_dir = os.path.join(config.output_dir, "inference_samples")
os.makedirs(self.inference_dir, exist_ok=True)
if not hasattr(config, 'inference_prompt_path') or not config.inference_prompt_path:
logger.warning("The inference prompt path is not specified; sampling will be skipped.")
self.skip_inference = True
elif not hasattr(config, 'inference_test_text') or not config.inference_test_text:
logger.warning("The inference test text is not specified; the sample will be skipped.")
self.skip_inference = True
else:
self.skip_inference = False
logger.info(f"Inference Callback is ready. Examples will be saved here: {self.inference_dir}")
def on_save(self, args, state, control, **kwargs):
if self.skip_inference:
return
step = state.global_step
checkpoint_dir = os.path.join(args.output_dir, f"checkpoint-{step}")
is_lora = getattr(self.config, "is_lora", False)
if is_lora:
if not os.path.exists(checkpoint_dir):
logger.warning(f"Checkpoint directory could not be found: {checkpoint_dir}")
return
logger.info(f"Initializing inference for checkpoint-{step} (LoRA)...")
try:
logger.info(f"Saving PEFT adapters explicitly to {checkpoint_dir}...")
model_wrapper = kwargs.get('model')
peft_model_to_save = None
if hasattr(model_wrapper, 'model') and isinstance(model_wrapper.model, torch.nn.Module):
peft_model_to_save = model_wrapper.model
elif hasattr(model_wrapper, 't3'):
peft_model_to_save = model_wrapper.t3
else:
peft_model_to_save = model_wrapper
if hasattr(peft_model_to_save, 'save_pretrained'):
peft_model_to_save.save_pretrained(checkpoint_dir)
logger.info("Adapter config and weights saved successfully.")
else:
logger.warning("Could not find a save_pretrained method on the model.")
except Exception as e:
logger.error(f"Failed to force save PEFT adapters: {e}")
try:
output_path = os.path.join(self.inference_dir, f"checkpoint-{step}.wav")
self._generate_sample_lora(checkpoint_dir, output_path)
except Exception as e:
logger.error(f"An error occurred during LoRA inference (Step: {step}): {e}", exc_info=True)
else:
weights_path = os.path.join(checkpoint_dir, "model.safetensors")
if not os.path.exists(weights_path):
weights_path = os.path.join(checkpoint_dir, "pytorch_model.bin")
if not os.path.exists(weights_path):
logger.warning(f"Checkpoint weights could not be found: {checkpoint_dir}")
return
logger.info(f"Initializing inference for checkpoint-{step} (Full Fine-Tune)...")
try:
output_path = os.path.join(self.inference_dir, f"checkpoint-{step}.wav")
self._generate_sample_full(weights_path, output_path)
except Exception as e:
logger.error(f"An error occurred during inference (Step: {step}): {e}", exc_info=True)
# -------------------------------------------------------------------------
# LoRA inference
# -------------------------------------------------------------------------
def _generate_sample_lora(self, checkpoint_dir: str, output_path: str):
from peft import PeftModel
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
is_turbo = getattr(self.config, "is_turbo", False)
EngineClass = ChatterboxTurboTTS if is_turbo else ChatterboxTTS
inference_engine = None
new_t3 = None
try:
# Rebuild the base T3 with resized vocab
temp_original = EngineClass.from_local(self.config.model_dir, device="cpu")
pretrained_state = temp_original.t3.state_dict()
original_config = temp_original.t3.hp
new_config = original_config
new_config.text_tokens_dict_size = self.config.new_vocab_size
if hasattr(new_config, "use_cache"):
new_config.use_cache = False
new_t3 = T3(hp=new_config)
new_t3 = resize_and_load_t3_weights(new_t3, pretrained_state)
if is_turbo and hasattr(new_t3.tfmr, "wte"):
del new_t3.tfmr.wte
del temp_original
del pretrained_state
inference_engine = EngineClass.from_local(self.config.model_dir, device="cpu")
inference_engine.t3 = new_t3
logger.info(f"Loading LoRA adapters from: {checkpoint_dir}")
inference_engine.t3 = PeftModel.from_pretrained(
inference_engine.t3,
checkpoint_dir,
is_trainable=False,
)
inference_engine.t3.to(device).eval()
inference_engine.s3gen.to(device).eval()
inference_engine.ve.to(device).eval()
inference_engine.device = device
params = {"temperature": 0.8, "repetition_penalty": 1.2}
if not is_turbo:
params["cfg_weight"] = 0.5
params["exaggeration"] = 0.5
with torch.no_grad():
wav = inference_engine.generate(
text=self.config.inference_test_text,
audio_prompt_path=self.config.inference_prompt_path,
**params,
)
if isinstance(wav, tuple):
wav = wav[0]
wav_np = wav.squeeze().cpu().numpy()
trimmed_wav = trim_silence_with_vad(wav_np, inference_engine.sr)
sf.write(output_path, trimmed_wav, inference_engine.sr)
logger.info(f"Example saved: {output_path}")
except Exception as e:
logger.error(f"LoRA inference callback failed: {e}", exc_info=True)
finally:
if inference_engine:
del inference_engine
if new_t3:
del new_t3
torch.cuda.empty_cache()
logger.info("LoRA inference cleanup done.")
# -------------------------------------------------------------------------
# Full fine-tune inference
# -------------------------------------------------------------------------
def _generate_sample_full(self, checkpoint_path: str, output_path: str):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
is_turbo = getattr(self.config, "is_turbo", False)
EngineClass = ChatterboxTurboTTS if is_turbo else ChatterboxTTS
tts_engine = EngineClass.from_local(self.config.model_dir, device="cpu")
t3_config = tts_engine.t3.hp
if hasattr(self.config, 'new_vocab_size'):
t3_config.text_tokens_dict_size = self.config.new_vocab_size
new_t3 = T3(hp=t3_config)
if is_turbo and hasattr(new_t3.tfmr, "wte"):
del new_t3.tfmr.wte
if checkpoint_path.endswith(".safetensors"):
state_dict = load_file(checkpoint_path)
else:
state_dict = torch.load(checkpoint_path, map_location="cpu")
clean_state_dict = {}
for k, v in state_dict.items():
k_clean = k.replace("module.", "").replace("model.", "").replace("t3.", "")
if k_clean.startswith("t3."):
clean_state_dict[k_clean.replace("t3.", "")] = v
elif not any(x in k_clean for x in ["s3gen", "ve.", "tokenizer"]):
clean_state_dict[k_clean] = v
missing_keys, unexpected_keys = new_t3.load_state_dict(clean_state_dict, strict=False)
critical_missing = [k for k in missing_keys if "tfmr.layers" in k]
if len(critical_missing) > 0:
logger.error("[CRITICAL ERROR] Model weights COULD NOT BE LOADED!")
logger.error(f"Number of missing keys: {len(missing_keys)}")
logger.error(f"Examples of missing keys: {critical_missing[:3]}")
logger.error("The sound produced will be 100% NOISE. Check your checkpoint saving method.")
elif len(missing_keys) > 0:
non_wte_missing = [k for k in missing_keys if "wte" not in k]
if non_wte_missing:
logger.warning(f"Some weights are missing ({len(non_wte_missing)} keys): {non_wte_missing[:3]}...")
else:
logger.info("Weights loaded successfully (WTE missing is normal for Turbo).")
else:
logger.info("All weights loaded completely and successfully.")
tts_engine.t3 = new_t3
tts_engine.t3.to(device).eval()
tts_engine.s3gen.to(device).eval()
tts_engine.ve.to(device).eval()
tts_engine.device = device
params = {"temperature": 0.8, "repetition_penalty": 1.2}
if not is_turbo:
params["cfg_weight"] = 0.2
params["exaggeration"] = 1.2
with torch.no_grad():
wav = tts_engine.generate(
text=self.config.inference_test_text,
audio_prompt_path=self.config.inference_prompt_path,
**params,
)
if isinstance(wav, tuple):
wav = wav[0]
wav_np = wav.squeeze().cpu().numpy()
trimmed_wav = trim_silence_with_vad(wav_np, tts_engine.sr)
sf.write(output_path, trimmed_wav, tts_engine.sr)
logger.info(f"Example saved: {output_path}")
del tts_engine
del new_t3
del state_dict
del clean_state_dict
torch.cuda.empty_cache() |