Bug fixes, resuming support.

.gitignore

@@ -33,7 +33,8 @@ test.py
 *.json
 *.jsonl
 code_digest.txt
-uv.lock
+*.lock
+*.bak
 
 # =========================
 # Data, Logs, & Outputs

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)
+        # Return the formatted string, the waveform, and the token count
+        return formatted_text, wav, len(audio_tokens)

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()

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)
 

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:

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,