Initial upload: Nepali Chatterbox fine-tuning kit

94336a6 verified 26 days ago

8.22 kB

import os
import torch
import numpy as np
import soundfile as sf
import random
import re
from safetensors.torch import load_file

from src.model import resize_and_load_t3_weights
from src.utils import setup_logger, trim_silence_with_vad
from src.config import TrainConfig
from src.chatterbox_.tts import ChatterboxTTS
from src.chatterbox_.tts_turbo import ChatterboxTurboTTS
from src.chatterbox_.models.t3.t3 import T3


logger = setup_logger("Chatterbox-Inference")

cfg = TrainConfig()

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IS_TURBO = cfg.is_turbo
IS_LORA = cfg.is_lora
BASE_MODEL_DIR = cfg.model_dir
OUTPUT_DIR = cfg.output_dir

# --- LoRA path ---
LORA_CHECKPOINT_PATH = os.path.join(OUTPUT_DIR, "new_lang_adapter")

# --- Checkpoints and Generation Params ---
if IS_TURBO:
    if IS_LORA:
        FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_turbo_finetuned_merged.safetensors")
        PARAMS = {
            "temperature": 0.8,
            "repetition_penalty": 1.2,
        }   
    else:
        FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_turbo_finetuned.safetensors")
        PARAMS = {
            "temperature": 0.8,
            "repetition_penalty": 1.2,
        }
else:
    if IS_LORA:
        FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_finetuned_merged.safetensors")
        PARAMS = {
            "temperature": 0.8,
            "repetition_penalty": 1.2,
        }
    else:
        FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_finetuned.safetensors")
        PARAMS = {
            "temperature": 0.8,
            "exaggeration": 0.5,
            "cfg_weight": 0.5,
            "repetition_penalty": 1.2,
        }

TEXT_TO_SAY = "[hi]नेपाली भाषा नेपालको आधिकारिक भाषा हो र यो देवनागरी लिपिमा लेखिन्छ। यो हिमालयको काखमा बसोबास गर्ने लाखौं मानिसहरूको मातृभाषा हो।"
AUDIO_PROMPT = "./speaker_reference/reference.wav"
OUTPUT_FILE = "./output.wav"


# -----------------------------------------------------------------------------
# Engine loaders
# -----------------------------------------------------------------------------

def load_finetuned_engine_lora(device):
    """Loads the engine and restores a LoRA-finetuned T3 (with resized vocab)."""
    from peft import PeftModel

    logger.info(f"Loading in {'TURBO' if IS_TURBO else 'NORMAL'} mode  [LoRA]")
    logger.info(f"Loading base model from: {BASE_MODEL_DIR}")

    EngineClass = ChatterboxTurboTTS if IS_TURBO else ChatterboxTTS

    # 1. Load base model temporarily to grab pretrained weights + config
    temp_engine = EngineClass.from_local(BASE_MODEL_DIR, device="cpu")
    pretrained_state = temp_engine.t3.state_dict()
    original_config = temp_engine.t3.hp

    # 2. Build new T3 with expanded vocab
    logger.info(f"Initializing new T3 with vocab size: {cfg.new_vocab_size}")
    new_t3_config = original_config
    new_t3_config.text_tokens_dict_size = cfg.new_vocab_size
    if hasattr(new_t3_config, "use_cache"):
        new_t3_config.use_cache = False

    new_t3 = T3(hp=new_t3_config)

    # 3. Transfer & resize pretrained weights into new T3
    logger.info("Resizing base weights to fit new vocab...")
    new_t3 = resize_and_load_t3_weights(new_t3, pretrained_state)

    del temp_engine
    del pretrained_state

    # 4. Turbo: remove wte + monkey-patch (required by PEFT)
    if IS_TURBO:
        logger.info("Turbo Mode: Deleting 'wte' and applying monkey patch.")
        if hasattr(new_t3.tfmr, "wte"):
            del new_t3.tfmr.wte


    # 5. Prepare fresh engine and inject rebuilt T3
    tts_engine = EngineClass.from_local(BASE_MODEL_DIR, device="cpu")
    tts_engine.t3 = new_t3

    # 6. Load LoRA adapters + trained embeddings (modules_to_save)
    if not os.path.exists(LORA_CHECKPOINT_PATH):
        raise FileNotFoundError(f"Adapter path not found: {LORA_CHECKPOINT_PATH}")

    logger.info(f"Loading LoRA adapters & embeddings from: {LORA_CHECKPOINT_PATH}")
    tts_engine.t3 = PeftModel.from_pretrained(
        tts_engine.t3,
        LORA_CHECKPOINT_PATH,
        is_trainable=False,
    )

    tts_engine.t3.to(device).eval()
    tts_engine.s3gen.to(device).eval()
    tts_engine.ve.to(device).eval()
    tts_engine.device = device

    return tts_engine


def load_finetuned_engine_full(device):
    """Loads the engine and restores a fully fine-tuned T3 from safetensors."""

    logger.info(f"Loading in {'TURBO' if IS_TURBO else 'NORMAL'} mode  [Full Fine-Tune]")
    logger.info(f"Loading base model from: {BASE_MODEL_DIR}")

    EngineClass = ChatterboxTurboTTS if IS_TURBO else ChatterboxTTS

    tts_engine = EngineClass.from_local(BASE_MODEL_DIR, device="cpu")

    logger.info(f"Initializing new T3 with vocab size: {cfg.new_vocab_size}")
    t3_config = tts_engine.t3.hp
    t3_config.text_tokens_dict_size = cfg.new_vocab_size

    new_t3 = T3(hp=t3_config)

    if IS_TURBO:
        logger.info("Turbo Mode: Removing 'wte' layer from new T3 model.")
        if hasattr(new_t3.tfmr, "wte"):
            del new_t3.tfmr.wte

    if not os.path.exists(FINETUNED_WEIGHTS):
        logger.error(f"FATAL: Fine-tuned file not found at {FINETUNED_WEIGHTS}.")
        raise FileNotFoundError(FINETUNED_WEIGHTS)

    logger.info(f"Loading fine-tuned weights: {FINETUNED_WEIGHTS}")
    state_dict = load_file(FINETUNED_WEIGHTS, device="cpu")
    new_t3.load_state_dict(state_dict, strict=True)
    logger.info("Fine-tuned weights loaded 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

    return tts_engine


# -----------------------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------------------

def generate_sentence_audio(engine, text, prompt_path, **kwargs):
    """Generates audio for a single sentence and trims silence."""
    try:
        wav_tensor = engine.generate(text=text, audio_prompt_path=prompt_path, **kwargs)
        if isinstance(wav_tensor, tuple):
            wav_tensor = wav_tensor[0]
        wav_np = wav_tensor.squeeze().cpu().numpy()
        trimmed_wav = trim_silence_with_vad(wav_np, engine.sr)
        return engine.sr, trimmed_wav
    except Exception as e:
        logger.error(f"Error generating sentence '{text[:30]}...': {e}")
        return 24000, np.zeros(0)


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------

def main():

    device = "cuda" if torch.cuda.is_available() else "cpu"
    logger.info(f"Inference running on: {device}")
    logger.info(f"Training strategy: {'LoRA' if IS_LORA else 'Full Fine-Tune'}")

    if IS_LORA:
        engine = load_finetuned_engine_lora(device)
    else:
        engine = load_finetuned_engine_full(device)

    sentences = re.split(r'(?<=[.?!])\s+', TEXT_TO_SAY.strip())
    sentences = [s for s in sentences if s.strip()]

    logger.info(f"Found {len(sentences)} sentences to synthesize.")

    all_chunks = []
    sample_rate = 24000

    set_seed(42)

    for i, sent in enumerate(sentences):
        logger.info(f"Synthesizing ({i+1}/{len(sentences)}): {sent}")
        sr, audio_chunk = generate_sentence_audio(engine, sent, AUDIO_PROMPT, **PARAMS)

        if len(audio_chunk) > 0:
            all_chunks.append(audio_chunk)
            sample_rate = sr
            pause_samples = int(sr * 0.2)
            all_chunks.append(np.zeros(pause_samples, dtype=np.float32))

    if all_chunks:
        final_audio = np.concatenate(all_chunks)
        sf.write(OUTPUT_FILE, final_audio, sample_rate)
        logger.info(f"Result saved to: {OUTPUT_FILE}")
    else:
        logger.error("No audio was generated.")


if __name__ == "__main__":
    main()

	import os
	import torch
	import numpy as np
	import soundfile as sf
	import random
	import re
	from safetensors.torch import load_file

	from src.model import resize_and_load_t3_weights
	from src.utils import setup_logger, trim_silence_with_vad
	from src.config import TrainConfig
	from src.chatterbox_.tts import ChatterboxTTS
	from src.chatterbox_.tts_turbo import ChatterboxTurboTTS
	from src.chatterbox_.models.t3.t3 import T3


	logger = setup_logger("Chatterbox-Inference")

	cfg = TrainConfig()

	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
	IS_TURBO = cfg.is_turbo
	IS_LORA = cfg.is_lora
	BASE_MODEL_DIR = cfg.model_dir
	OUTPUT_DIR = cfg.output_dir

	# --- LoRA path ---
	LORA_CHECKPOINT_PATH = os.path.join(OUTPUT_DIR, "new_lang_adapter")

	# --- Checkpoints and Generation Params ---
	if IS_TURBO:
	if IS_LORA:
	FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_turbo_finetuned_merged.safetensors")
	PARAMS = {
	"temperature": 0.8,
	"repetition_penalty": 1.2,
	}
	else:
	FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_turbo_finetuned.safetensors")
	PARAMS = {
	"temperature": 0.8,
	"repetition_penalty": 1.2,
	}
	else:
	if IS_LORA:
	FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_finetuned_merged.safetensors")
	PARAMS = {
	"temperature": 0.8,
	"repetition_penalty": 1.2,
	}
	else:
	FINETUNED_WEIGHTS = os.path.join(OUTPUT_DIR, "t3_finetuned.safetensors")
	PARAMS = {
	"temperature": 0.8,
	"exaggeration": 0.5,
	"cfg_weight": 0.5,
	"repetition_penalty": 1.2,
	}

	TEXT_TO_SAY = "[hi]नेपाली भाषा नेपालको आधिकारिक भाषा हो र यो देवनागरी लिपिमा लेखिन्छ। यो हिमालयको काखमा बसोबास गर्ने लाखौं मानिसहरूको मातृभाषा हो।"
	AUDIO_PROMPT = "./speaker_reference/reference.wav"
	OUTPUT_FILE = "./output.wav"


	# -----------------------------------------------------------------------------
	# Engine loaders
	# -----------------------------------------------------------------------------

	def load_finetuned_engine_lora(device):
	"""Loads the engine and restores a LoRA-finetuned T3 (with resized vocab)."""
	from peft import PeftModel

	logger.info(f"Loading in {'TURBO' if IS_TURBO else 'NORMAL'} mode [LoRA]")
	logger.info(f"Loading base model from: {BASE_MODEL_DIR}")

	EngineClass = ChatterboxTurboTTS if IS_TURBO else ChatterboxTTS

	# 1. Load base model temporarily to grab pretrained weights + config
	temp_engine = EngineClass.from_local(BASE_MODEL_DIR, device="cpu")
	pretrained_state = temp_engine.t3.state_dict()
	original_config = temp_engine.t3.hp

	# 2. Build new T3 with expanded vocab
	logger.info(f"Initializing new T3 with vocab size: {cfg.new_vocab_size}")
	new_t3_config = original_config
	new_t3_config.text_tokens_dict_size = cfg.new_vocab_size
	if hasattr(new_t3_config, "use_cache"):
	new_t3_config.use_cache = False

	new_t3 = T3(hp=new_t3_config)

	# 3. Transfer & resize pretrained weights into new T3
	logger.info("Resizing base weights to fit new vocab...")
	new_t3 = resize_and_load_t3_weights(new_t3, pretrained_state)

	del temp_engine
	del pretrained_state

	# 4. Turbo: remove wte + monkey-patch (required by PEFT)
	if IS_TURBO:
	logger.info("Turbo Mode: Deleting 'wte' and applying monkey patch.")
	if hasattr(new_t3.tfmr, "wte"):
	del new_t3.tfmr.wte


	# 5. Prepare fresh engine and inject rebuilt T3
	tts_engine = EngineClass.from_local(BASE_MODEL_DIR, device="cpu")
	tts_engine.t3 = new_t3

	# 6. Load LoRA adapters + trained embeddings (modules_to_save)
	if not os.path.exists(LORA_CHECKPOINT_PATH):
	raise FileNotFoundError(f"Adapter path not found: {LORA_CHECKPOINT_PATH}")

	logger.info(f"Loading LoRA adapters & embeddings from: {LORA_CHECKPOINT_PATH}")
	tts_engine.t3 = PeftModel.from_pretrained(
	tts_engine.t3,
	LORA_CHECKPOINT_PATH,
	is_trainable=False,
	)

	tts_engine.t3.to(device).eval()
	tts_engine.s3gen.to(device).eval()
	tts_engine.ve.to(device).eval()
	tts_engine.device = device

	return tts_engine


	def load_finetuned_engine_full(device):
	"""Loads the engine and restores a fully fine-tuned T3 from safetensors."""

	logger.info(f"Loading in {'TURBO' if IS_TURBO else 'NORMAL'} mode [Full Fine-Tune]")
	logger.info(f"Loading base model from: {BASE_MODEL_DIR}")

	EngineClass = ChatterboxTurboTTS if IS_TURBO else ChatterboxTTS

	tts_engine = EngineClass.from_local(BASE_MODEL_DIR, device="cpu")

	logger.info(f"Initializing new T3 with vocab size: {cfg.new_vocab_size}")
	t3_config = tts_engine.t3.hp
	t3_config.text_tokens_dict_size = cfg.new_vocab_size

	new_t3 = T3(hp=t3_config)

	if IS_TURBO:
	logger.info("Turbo Mode: Removing 'wte' layer from new T3 model.")
	if hasattr(new_t3.tfmr, "wte"):
	del new_t3.tfmr.wte

	if not os.path.exists(FINETUNED_WEIGHTS):
	logger.error(f"FATAL: Fine-tuned file not found at {FINETUNED_WEIGHTS}.")
	raise FileNotFoundError(FINETUNED_WEIGHTS)

	logger.info(f"Loading fine-tuned weights: {FINETUNED_WEIGHTS}")
	state_dict = load_file(FINETUNED_WEIGHTS, device="cpu")
	new_t3.load_state_dict(state_dict, strict=True)
	logger.info("Fine-tuned weights loaded 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

	return tts_engine


	# -----------------------------------------------------------------------------
	# Helpers
	# -----------------------------------------------------------------------------

	def generate_sentence_audio(engine, text, prompt_path, **kwargs):
	"""Generates audio for a single sentence and trims silence."""
	try:
	wav_tensor = engine.generate(text=text, audio_prompt_path=prompt_path, **kwargs)
	if isinstance(wav_tensor, tuple):
	wav_tensor = wav_tensor[0]
	wav_np = wav_tensor.squeeze().cpu().numpy()
	trimmed_wav = trim_silence_with_vad(wav_np, engine.sr)
	return engine.sr, trimmed_wav
	except Exception as e:
	logger.error(f"Error generating sentence '{text[:30]}...': {e}")
	return 24000, np.zeros(0)


	def set_seed(seed):
	random.seed(seed)
	np.random.seed(seed)
	torch.manual_seed(seed)
	if torch.cuda.is_available():
	torch.cuda.manual_seed_all(seed)


	# -----------------------------------------------------------------------------
	# Main
	# -----------------------------------------------------------------------------

	def main():

	device = "cuda" if torch.cuda.is_available() else "cpu"
	logger.info(f"Inference running on: {device}")
	logger.info(f"Training strategy: {'LoRA' if IS_LORA else 'Full Fine-Tune'}")

	if IS_LORA:
	engine = load_finetuned_engine_lora(device)
	else:
	engine = load_finetuned_engine_full(device)

	sentences = re.split(r'(?<=[.?!])\s+', TEXT_TO_SAY.strip())
	sentences = [s for s in sentences if s.strip()]

	logger.info(f"Found {len(sentences)} sentences to synthesize.")

	all_chunks = []
	sample_rate = 24000

	set_seed(42)

	for i, sent in enumerate(sentences):
	logger.info(f"Synthesizing ({i+1}/{len(sentences)}): {sent}")
	sr, audio_chunk = generate_sentence_audio(engine, sent, AUDIO_PROMPT, **PARAMS)

	if len(audio_chunk) > 0:
	all_chunks.append(audio_chunk)
	sample_rate = sr
	pause_samples = int(sr * 0.2)
	all_chunks.append(np.zeros(pause_samples, dtype=np.float32))

	if all_chunks:
	final_audio = np.concatenate(all_chunks)
	sf.write(OUTPUT_FILE, final_audio, sample_rate)
	logger.info(f"Result saved to: {OUTPUT_FILE}")
	else:
	logger.error("No audio was generated.")


	if __name__ == "__main__":
	main()