From a53a8b3982844bdfefe141a4d203045561dc5297 Mon Sep 17 00:00:00 2001 From: Nighthawk Date: Fri, 27 Feb 2026 10:23:04 -0500 Subject: [PATCH] Bug fixes, resuming support. --- .gitignore | 3 +- dataset_e2e.py | 8 +++- resume_training.py | 95 +++++++++++++++++++++++++++++++++++++++++++++ simple_inference.py | 24 +++++------- train_decoder.py | 40 ++++++++++--------- train_llm.py | 22 ++++++----- 6 files changed, 147 insertions(+), 45 deletions(-) create mode 100644 resume_training.py diff --git a/.gitignore b/.gitignore index b1fcde0..408f60b 100644 --- a/.gitignore +++ b/.gitignore @@ -33,7 +33,8 @@ test.py *.json *.jsonl code_digest.txt -uv.lock +*.lock +*.bak # ========================= # Data, Logs, & Outputs diff --git a/dataset_e2e.py b/dataset_e2e.py index 0ebd746..dfb81fc 100644 --- a/dataset_e2e.py +++ b/dataset_e2e.py @@ -21,6 +21,10 @@ class AudioDataset(Dataset): # [transcript, audio_tokens (list), audio_path] text, audio_tokens, audio_path = self.dataset[idx] + # CRITICAL FIX: We must format the string with the audio tokens physically embedded + # so train_decoder.py can tokenize it and find the audio indices to align the waveform! + formatted_text = f"[TEXT]{text}[START]{''.join(list(map(lambda x: f'[{x}]', audio_tokens)))}[STOP]" + # Use our robust OS-aware pipeline to load the audio try: wav, _ = AudioPipeline.load_audio(audio_path, target_sr=self.target_sr) @@ -32,5 +36,5 @@ class AudioDataset(Dataset): # Fallback to silence to prevent dataloader crashes wav = torch.zeros(len(audio_tokens) * SAMPLES_PER_TOKEN) - # Return the exact tuple expected by collate_pack in train_decoder.py - return text, wav, len(audio_tokens) \ No newline at end of file + # Return the formatted string, the waveform, and the token count + return formatted_text, wav, len(audio_tokens) \ No newline at end of file diff --git a/resume_training.py b/resume_training.py new file mode 100644 index 0000000..2deb33a --- /dev/null +++ b/resume_training.py @@ -0,0 +1,95 @@ +import os +import glob +import subprocess +import argparse +import yaml +from pathlib import Path + +def get_latest_checkpoint(base_path, pattern): + """Finds the latest checkpoint file/folder based on step number.""" + checkpoints = glob.glob(os.path.join(base_path, pattern)) + if not checkpoints: + return None + + # Extract step number and sort + # For LLM: 'checkpoint-3000' -> 3000 + # For Decoder: 'decoder_step_5000.pth' -> 5000 + try: + if "checkpoint-" in checkpoints[0]: + checkpoints.sort(key=lambda x: int(x.split('-')[-1])) + else: + import re + checkpoints.sort(key=lambda x: int(re.findall(r'\d+', os.path.basename(x))[0])) + except Exception: + checkpoints.sort() + + return checkpoints[-1] + +def update_config(config_path, updates): + with open(config_path, 'r') as f: + config = yaml.safe_load(f) + + # Deep update logic + for section, values in updates.items(): + if section in config: + config[section].update(values) + else: + config[section] = values + + with open(config_path, 'w') as f: + yaml.safe_dump(config, f, default_flow_style=False) + +def main(): + parser = argparse.ArgumentParser(description="Resume training Phase 2 or Phase 3.") + parser.add_argument("phase", choices=["llm", "decoder"], help="Which phase to resume") + args = parser.parse_args() + + with open("config.yaml", 'r') as f: + config = yaml.safe_load(f) + + save_dir = config["paths"]["save_dir"] + + if args.phase == "llm": + ckpt_dir = os.path.join(save_dir, "llm") + latest = get_latest_checkpoint(ckpt_dir, "checkpoint-*") + + if latest: + print(f"Found latest LLM checkpoint: {latest}") + # HuggingFace expects the path to the folder or the specific safetensors file + model_file = os.path.join(latest, "model.safetensors") + update_config("config.yaml", { + "paths": {"pretrained_llm_path": model_file}, + "llm": {"from_scratch": False} + }) + print("Config updated. Resuming LLM training...") + subprocess.run(["python", "train_llm.py"]) + else: + print("No LLM checkpoints found to resume from.") + + elif args.phase == "decoder": + ckpt_dir = os.path.join(save_dir, "decoder") + latest_dec = get_latest_checkpoint(ckpt_dir, "decoder_step_*.pth") + + if latest_dec: + print(f"Found latest Decoder checkpoint: {latest_dec}") + import re + step = int(re.findall(r'\d+', os.path.basename(latest_dec))[0]) + + # Find matching discriminator if it exists + latest_disc = latest_dec.replace("decoder_step_", "discriminator_step_") + disc_path = latest_disc if os.path.exists(latest_disc) else None + + update_config("config.yaml", { + "paths": { + "pretrained_decoder_path": latest_dec, + "pretrained_discriminator_path": disc_path + }, + "decoder": {"start_step": step} + }) + print(f"Config updated to step {step}. Resuming Decoder training...") + subprocess.run(["python", "train_decoder.py"]) + else: + print("No Decoder checkpoints found to resume from.") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/simple_inference.py b/simple_inference.py index 607555b..20492f4 100644 --- a/simple_inference.py +++ b/simple_inference.py @@ -2,14 +2,15 @@ import torch import torchaudio import argparse import os -from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer +from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer, PreTrainedModel +from torch import nn from safetensors.torch import load_file # Ensure decoder module is importable from decoder.decoder import SopranoDecoder from utils.config_loader import load_config -def load_models(llm_path, decoder_path, device='cuda'): +def load_models(llm_path: str, decoder_path: str, tokenizer_name: str, device: str = 'cuda') -> tuple[PreTrainedModel, nn.Module]: if not llm_path or not os.path.exists(llm_path): raise FileNotFoundError(f"LLM path invalid or not found: {llm_path}. Please check your config.yaml.") if not decoder_path or not os.path.exists(decoder_path): @@ -18,7 +19,7 @@ def load_models(llm_path, decoder_path, device='cuda'): print(f"Loading LLM from {llm_path}...") # Load LLM Config & Model - config = AutoConfig.from_pretrained('ekwek/Soprano-80M') + config = AutoConfig.from_pretrained(tokenizer_name) llm = AutoModelForCausalLM.from_config(config) # Load LLM weights @@ -32,7 +33,6 @@ def load_models(llm_path, decoder_path, device='cuda'): llm.to(device).eval() print(f"Loading Decoder from {decoder_path}...") - # Instantiate Decoder with defaults decoder = SopranoDecoder() # Load Decoder weights @@ -42,7 +42,7 @@ def load_models(llm_path, decoder_path, device='cuda'): return llm, decoder -def generate_audio(text, llm, decoder, tokenizer, cfg_inf, device='cuda', save_path="output.wav"): +def generate_audio(text: str, llm: PreTrainedModel, decoder: nn.Module, tokenizer, cfg_inf: dict, device: str = 'cuda', save_path: str = "output.wav"): # 1. Format Prompt prompt = f"[TEXT]{text}[START]" inputs = tokenizer(prompt, return_tensors="pt").to(device) @@ -73,7 +73,7 @@ def generate_audio(text, llm, decoder, tokenizer, cfg_inf, device='cuda', save_p # 3. Process Hidden States hidden_states_list = [] - # outputs.hidden_states is tuple of generated steps. + # outputs.hidden_states is a tuple of generated steps. # The first element is the Prompt (prefill) hidden states. We skip it. for i, step_states in enumerate(outputs.hidden_states): # step_states is tuple of layers. Get last layer. @@ -81,10 +81,7 @@ def generate_audio(text, llm, decoder, tokenizer, cfg_inf, device='cuda', save_p hidden_states_list.append(last_layer_state) # Concatenate along time dimension - # Result: (Batch, T_gen, Dim) audio_hidden = torch.stack(hidden_states_list).unsqueeze(0) # (B, T, D) - - # Ensure float32 audio_hidden = audio_hidden.to(torch.float32) num_audio_tokens = audio_hidden.size(1) @@ -95,15 +92,13 @@ def generate_audio(text, llm, decoder, tokenizer, cfg_inf, device='cuda', save_p return # 4. Decode - # Decoder expects (B, Channels, T) decoder_input = audio_hidden.transpose(1, 2) - print(f"Decoding shape: {decoder_input.shape}...") with torch.no_grad(): audio = decoder(decoder_input) # 5. Save - audio = audio.squeeze().cpu() # (Samples,) or (1, Samples) + audio = audio.squeeze().cpu() if audio.dim() == 1: audio = audio.unsqueeze(0) @@ -123,15 +118,16 @@ def main(): cfg_inf = config["inference"] device = cfg_global["device"] if torch.cuda.is_available() else "cpu" + tokenizer_name = cfg_global.get("tokenizer_name", "ekwek/Soprano-80M") print(f"Using device: {device}") - tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M') + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) tokenizer.eos_token_id = 3 llm_path = cfg_paths["pretrained_llm_path"] decoder_path = cfg_paths["pretrained_decoder_path"] - llm, decoder = load_models(llm_path, decoder_path, device) + llm, decoder = load_models(llm_path, decoder_path, tokenizer_name, device) generate_audio(args.text, llm, decoder, tokenizer, cfg_inf, device, args.out) diff --git a/train_decoder.py b/train_decoder.py index 1a0ef91..2eaee02 100644 --- a/train_decoder.py +++ b/train_decoder.py @@ -3,15 +3,14 @@ Training script for Soprano Decoder (Vocos). Freezes LLM and trains Decoder with GAN loss. """ import os +from functools import partial import random import time -import io import wandb import matplotlib.pyplot as plt import numpy as np import torch -import torchaudio from torch.utils.data import DataLoader from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig @@ -116,7 +115,8 @@ def evaluate(step, val_dataloader_it, val_dataloader, model, decoder, discrimina vaudio_mask = vaudio_mask.to(device) with torch.autocast(device_type=device_type, dtype=torch.bfloat16): - voutputs = model(vx, output_hidden_states=True) + # Explicitly name input_ids to satisfy strict typing + voutputs = model(input_ids=vx, output_hidden_states=True) v_hidden = voutputs.hidden_states[-1].to(torch.float32) v_gathered_states_list = [] @@ -210,7 +210,6 @@ def evaluate(step, val_dataloader_it, val_dataloader, model, decoder, discrimina }) if use_wandb: - # Generate Mel Images (from last val batch) gen_mel = mel_fn(v_fake_audio[0:1]).squeeze(0).cpu().numpy() gt_mel = mel_fn(v_real_audio[0:1]).squeeze(0).cpu().numpy() @@ -329,6 +328,8 @@ def main(): # 3. Load Discriminator # ------------------ discriminator = None + opt_d = None + if use_disc: print("Initializing Discriminator...") discriminator = Discriminator() @@ -342,6 +343,7 @@ def main(): discriminator.to(device) discriminator.train() + opt_d = torch.optim.AdamW(discriminator.parameters(), max_lr, betas=betas, weight_decay=weight_decay) else: print("Training WITHOUT Discriminator (Reconstruction only).") @@ -357,7 +359,7 @@ def main(): pin_memory=True, persistent_workers=True, worker_init_fn=worker_seed_init, - collate_fn=lambda batch_in: collate_pack(batch_in, tokenizer), + collate_fn=partial(collate_pack, tokenizer=tokenizer), ) dataloader_it = iter(dataloader) @@ -370,14 +372,11 @@ def main(): pin_memory=True, persistent_workers=True, worker_init_fn=worker_seed_init, - collate_fn=lambda batch_in: collate_pack(batch_in, tokenizer), + collate_fn=partial(collate_pack, tokenizer=tokenizer), ) val_dataloader_it = iter(val_dataloader) opt_g = torch.optim.AdamW(decoder.parameters(), max_lr, betas=betas, weight_decay=weight_decay) - opt_d = None - if use_disc: - opt_d = torch.optim.AdamW(discriminator.parameters(), max_lr, betas=betas, weight_decay=weight_decay) # ------------------ # Training Loop @@ -404,7 +403,7 @@ def main(): with torch.no_grad(): with torch.autocast(device_type=device_type, dtype=torch.bfloat16): - outputs = model(x, output_hidden_states=True) + outputs = model(input_ids=x, output_hidden_states=True) hidden_states = outputs.hidden_states[-1] hidden_states = hidden_states.to(torch.float32) @@ -425,7 +424,7 @@ def main(): d_loss_item = 0.0 - if use_disc: + if use_disc and opt_d is not None and discriminator is not None: opt_d.zero_grad() decoder_in = decoder_in_padded.transpose(1, 2) @@ -478,8 +477,9 @@ def main(): real_crop_list_g = [] fake_crop_list_g = [] - if use_disc: - for b_idx in range(bsz): + + if use_disc and discriminator is not None: + for b_idx in range(bsz): valid_len = gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN valid_len = min(valid_len, min_len) @@ -495,8 +495,8 @@ def main(): real_crop_list_g.append(r_c) fake_crop_list_g.append(f_c) - real_crops_g = torch.stack(real_crop_list_g).unsqueeze(1) - fake_crops_g = torch.stack(fake_crop_list_g).unsqueeze(1) + real_crops_g = torch.stack(real_crop_list_g).unsqueeze(1) + fake_crops_g = torch.stack(fake_crop_list_g).unsqueeze(1) frames_per_token = SAMPLES_PER_TOKEN // 512 mel_mask = audio_loss_mask.repeat_interleave(frames_per_token, dim=1) @@ -518,7 +518,7 @@ def main(): loss_fm = torch.tensor(0.0, device=device) loss_gen = torch.tensor(0.0, device=device) - if use_disc: + if use_disc and discriminator is not None: y_d_rs, y_d_gs, fmap_rs, fmap_gs = discriminator(real_crops_g, fake_crops_g) loss_fm = feature_matching_loss(fmap_rs, fmap_gs) loss_gen, _ = generator_loss(y_d_gs) @@ -530,7 +530,8 @@ def main(): lr = get_lr(step, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps) for param_group in opt_g.param_groups: param_group['lr'] = lr - if use_disc: + + if use_disc and opt_d is not None: for param_group in opt_d.param_groups: param_group['lr'] = lr / 2 opt_g.step() @@ -538,6 +539,7 @@ def main(): end = time.time() dt = (end-start)*1000 + # Pylance fix: loss_fm and loss_gen are tensors initialized to 0.0, so .item() works. d_loss_item is a standard float. tqdm_log = f'mel: {loss_mel.item():.3f} | gen: {loss_gen.item():.3f} | sc: {sc_loss.item():.3f} | mag: {mag_loss.item():.3f} | fm: {loss_fm.item():.3f} | d: {d_loss_item:.3f} | lr: {lr:.2e} | time: {dt:.2f} ms' pbar.set_description(tqdm_log) @@ -576,7 +578,7 @@ def main(): ckpt_name_dec = f"decoder_step_{step}.pth" ckpt_name_disc = f"discriminator_step_{step}.pth" torch.save(decoder.state_dict(), os.path.join(save_path, ckpt_name_dec)) - if discriminator: + if discriminator is not None: torch.save(discriminator.state_dict(), os.path.join(save_path, ckpt_name_disc)) if use_wandb: @@ -584,7 +586,7 @@ def main(): print(f"\nTraining complete. Saving model at {save_path}") torch.save(decoder.state_dict(), os.path.join(save_path, "decoder_trained.pth")) - if discriminator: + if discriminator is not None: torch.save(discriminator.state_dict(), os.path.join(save_path, "discriminator_trained.pth")) if use_wandb: diff --git a/train_llm.py b/train_llm.py index dd85e78..7ee3178 100644 --- a/train_llm.py +++ b/train_llm.py @@ -5,8 +5,11 @@ Usage: python train_llm.py """ import os +from functools import partial +from pyexpat import model import random import time +from matplotlib.pyplot import step import wandb import numpy as np @@ -221,7 +224,7 @@ def main(): pin_memory=True, persistent_workers=True, worker_init_fn=worker_seed_init, - collate_fn=lambda texts: collate_dynamic(texts, tokenizer), + collate_fn=partial(collate_dynamic, tokenizer=tokenizer), ) dataloader_it = iter(dataloader) @@ -234,7 +237,7 @@ def main(): pin_memory=True, persistent_workers=True, worker_init_fn=worker_seed_init, - collate_fn=lambda texts: collate_dynamic(texts, tokenizer), + collate_fn=partial(collate_dynamic, tokenizer=tokenizer), ) opt = torch.optim.AdamW(model.parameters(), max_lr, betas=betas, weight_decay=weight_decay, fused=True) @@ -271,9 +274,10 @@ def main(): logits = model(x, attention_mask=attn_mask).logits audio_loss, text_loss, acc = compute_loss(x, logits, y, grad_accum_steps, mask=attn_mask) - audio_loss_accum += audio_loss.detach() - text_loss_accum += text_loss.detach() - acc_accum += acc.detach() + # CRITICAL FIX: Extract the float value immediately + audio_loss_accum += audio_loss.item() + text_loss_accum += text_loss.item() + acc_accum += acc.item() total_loss = audio_loss + text_factor * text_loss total_loss.backward() @@ -292,14 +296,14 @@ def main(): dt = (end - start) * 1000 tokens_per_second = (batch_size * seq_len * grad_accum_steps) / (end - start) - tqdm_log = f'text loss: {text_loss_accum.item():.3f} | audio loss: {audio_loss_accum.item():.3f} | acc: {acc_accum.item():.4f} | lr: {lr:.2e} | norm: {norm:.3f} | time: {dt:.2f} ms | {tokens_per_second:.2f} t/s' + tqdm_log = f'text loss: {text_loss_accum:.3f} | audio loss: {audio_loss_accum:.3f} | acc: {acc_accum:.4f} | lr: {lr:.2e} | norm: {norm:.3f} | time: {dt:.2f} ms | {tokens_per_second:.2f} t/s' pbar.set_description(tqdm_log) if cfg_global["use_wandb"]: wandb.log({ - "train/text_loss": text_loss_accum.item(), - "train/audio_loss": audio_loss_accum.item(), - "train/acc": acc_accum.item(), + "train/text_loss": text_loss_accum, + "train/audio_loss": audio_loss_accum, + "train/acc": acc_accum, "train/lr": lr, "train/grad_norm": norm, "train/dt": dt,