Nighthawk/llm-tts-factory
1
0
Fork 0
mirror of https://github.com/Nighthawk42/llm-tts-factory.git synced 2026-08-30 07:22:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: Centralize config, add OS-aware audio pipeline, and migrate to uv

Browse Source 
Major framework modernization and quality-of-life improvements:

- Centralized Configuration: Replaced scattered, hardcoded hyperparameters and paths across all training/inference scripts with a single, documented `config.yaml` and `config_loader.py`.
- OS-Aware Audio Pipeline: Introduced `utils/audio_utils.py` to handle cross-platform audio loading. Automatically routes Windows to a robust `ffmpeg` subprocess to bypass unstable Python audio bindings, while keeping `torchaudio` for Linux.
- Dependency Management: Migrated from `requirements.txt` to `uv` with a fully configured `pyproject.toml`. Explicitly targets Python 3.12 and pulls PyTorch `cu128` wheels by default.
- Dataset Fixes: Restored the missing `dataset_e2e.py` required for proper STFT/GAN decoder training and updated all dataloaders to utilize the new AudioPipeline.
- Documentation & Housekeeping: Overhauled `README.md` with updated workflows, Windows instructions, and `uv` setup. Added a comprehensive `.gitignore` for virtual environments, model weights, and cache files.
This commit is contained in:
Nighthawk
2026-02-26 13:58:35 -05:00
parent 5670af427b
commit e7ee43de41
25 changed files with 942 additions and 748 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+49 -3
View File
@@ -1,6 +1,52 @@
# =========================
# Virtual Environments
# =========================
.venv/
venv/
env/
node_modules/
# =========================
# Python Cache & Build
# =========================
__pycache__/
test.py
*.json
*.pth
*.py[cod]
*$py.class
build/
dist/
*.egg-info/
# =========================
# Model Weights & Checkpoints
# =========================
*.pth
*.pt
*.safetensors
*.bin
*.ckpt
*.onnx
# =========================
# Specific Files
# =========================
test.py
*.json
*.jsonl
code_digest.txt
# =========================
# Data, Logs, & Outputs
# =========================
wandb/
logs/
*.wav
*.flac
*.mp3
# =========================
# OS & Editor Files
# =========================
.DS_Store
.env
.vscode/
.idea/
README.md
+87 -21
View File
@@ -1,19 +1,18 @@
# llm-tts-factory: End-to-End LLM-Backbone TTS Training Framework
llm-tts-factory is a full suite of end-to-end training scripts designed for building llm backbone based TTS model from scratch.
llm-tts-factory is a full suite of end-to-end training scripts designed for building an LLM-backbone-based TTS model from scratch.
To begin with, taking inspiration from [Soprano](https://huggingface.co/ekwek/Soprano-1.1-80M), this repository allows you to train a Soprano style TTS model from the ground up. Because of its architecture, it features an **extra Decoder training step**. Instead of generating audio directly from discrete tokens, this model uses the **hidden states of the LLM** as inputs to generate high-fidelity audio.
Taking inspiration from [Soprano](https://huggingface.co/ekwek/Soprano-1.1-80M), this repository allows you to train a Soprano-style TTS model from the ground up. Because of its architecture, it features an **extra Decoder training step**. Instead of generating audio directly from discrete tokens, this model uses the **hidden states of the LLM** as inputs to generate high-fidelity audio.
I think this can be nicely setup for anybody to train llm backbone TTS models from scratch. Can spend some time to make it more modular and user friendly. Would love to hear from interested people on this!
---
## Architecture
## 🏗️ Architecture
The framework is divided into three core stages:
### 1. Codec
- **Description:** Encodes raw audio into discrete units and decodes them back.
- **Current State:** A very naive codec encoder and decoder. There is a ton of scope for improvements here (e.g., swapping in RVQ, DAC, or EnCodec).
- **Current State:** A naive codec encoder and decoder. There is a ton of scope for improvements here (e.g., swapping in RVQ, DAC, or EnCodec).
### 2. LLM Backbone
- **Model:** Qwen-based causal language model.
@@ -22,43 +21,110 @@ The framework is divided into three core stages:
### 3. Decoder
- **Model:** Vocos-based decoder.
- **Description:** A dedicated vocoder trained with Multi-Resolution STFT and GAN losses to synthesize the final audio waveform directly from the LLM's continuous hidden states.
- **Training strategy:** Trained in multiple stages, first to nail reconstruction and then the perception quality using gan losses.
- **Training strategy:** Trained in multiple stages, first to nail reconstruction and then the perception quality using GAN losses.
---
## Training & Inference Commands
## 📦 Installation & Setup
You can train the entire pipeline from scratch using the following step-by-step commands.
This project uses [uv](https://github.com/astral-sh/uv) for lightning-fast dependency management. We target Python 3.12 and PyTorch with CUDA 12.8 support by default.
1. **Install `uv`** (if you haven't already):
```bash
pip install uv
```
1. **Initialize the Virtual Environment & Install Dependencies**:
```bash
# Create a seeded Python 3.12 environment
uv venv --python 3.12 --seed
# Activate the environment (Windows)
.venv\Scripts\activate
# Or on Linux/macOS: source .venv/bin/activate
# Sync all dependencies (automatically pulls CUDA 12.8 PyTorch wheels)
uv sync
```
## ⚙️ Configuration
Say goodbye to messy command-line arguments! The entire framework is centrally managed by the `config.yaml` file located in the root directory.
Before running any scripts, open `config.yaml` to set your:
- Dataset and checkpoint paths (relative or absolute).
- Training hyperparameters (batch size, learning rate, max steps, etc.).
- Global settings (device selection, random seed, Weights & Biases logging).
---
## 🪟 Windows Compatibility & Audio Setup
Audio library bindings (`torchaudio`, `torchcodec`) can be problematic and crash-prone on Windows. This framework includes a custom OS-aware `AudioPipeline` (`utils/audio_utils.py`) that detects your operating system.
If you are on Windows, the pipeline will automatically use **ffmpeg** to safely decode, convert to mono, and resample audio into raw PCM float32 arrays without requiring complex C++ bindings.
**Windows Requirements:**
- Ensure `ffmpeg` is added to your system `PATH`, **OR**
- Place the `ffmpeg.exe` binary directly in your project folder at: `./tools/ffmpeg/ffmpeg.exe`.
- *Fallback:* If `ffmpeg` is not found, the pipeline will attempt to fallback to `soundfile` and `scipy`, but `ffmpeg` is highly recommended for speed and format compatibility.
*(Linux users: The pipeline will natively use torchaudio as intended.)*
---
## 🚀 Data Preparation, Training & Inference
Make sure your paths are set in `config.yaml`, then execute the steps in order:
### 0. Data Preparation
Convert your LJSpeech-formatted dataset into audio tokens:
```bash
python generate_dataset.py
# OR
python generate_dataset_from_lists.py
```
### 1. Codec Stage
Train the audio codec encoder and decoder:
```bash
python codec_train.py
```
### 2. LLM Stage
Train the Qwen-based causal LLM to learn the mapping from text to audio representations:
```bash
python train_llm.py \
--input-dir <input_dir> \
--save-dir <save_dir> \
--from-scratch
python train_llm.py
```
### 3. Decoder Stage
Freeze the LLM and train the Vocos decoder to reconstruct the audio from the LLM's hidden states:
```bash
python train_decoder.py \
--input-dir <input_dir> \
--save-dir <save_dir>
python train_decoder.py
```
### 4. Inference
Run end-to-end inference using your trained LLM and Decoder pair to generate TTS:
Run end-to-end inference. (Text is passed via CLI, but generation params like temperature and paths are pulled from `config.yaml`):
```bash
python simple_inference.py \
--text "hello, my name is soma siddhartha" \
--llm-path <llm_path> \
--decoder-path <decoder_path> \
--out simple_inf_out.wav
python simple_inference.py --text "hello, my name is soma siddhartha" --out simple_inf_out.wav
```
codec/codec_decoder/__pycache__/decoder.cpython-311.pyc
BIN
View File
Binary file not shown.
codec/encoder/__pycache__/codec.cpython-311.pyc
BIN
View File
Binary file not shown.
codec/encoder/__pycache__/quantizer.cpython-311.pyc
BIN
View File
Binary file not shown.
codec_dataset.py
+12 -26
View File
@@ -1,9 +1,9 @@
import torch
import torchaudio
from torch.utils.data import Dataset
import os
import json
from utils.audio_utils import AudioPipeline
class LJSpeechDataset(Dataset):
def __init__(self, root, sample_rate=32000, mode='train'):
@@ -12,42 +12,28 @@ class LJSpeechDataset(Dataset):
"""
self.root = root
self.sample_rate = sample_rate
# Write code here to handle modes. Take wave files only from mode json.
self.mode = mode
mode_json = os.path.join(root, f"{mode}.json")
if not os.path.exists(mode_json):
raise FileNotFoundError(f"Dataset JSON not found: {mode_json}")
with open(mode_json, 'r') as f:
self.dataset = json.load(f)
# self.wav_dir = os.path.join(root, "wavs")
# self.wav_files = sorted(
# [f for f in os.listdir(self.wav_dir) if f.endswith(".wav")]
# )
# assert len(self.wav_files) > 0, "No wav files found!"
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
# wav_path = os.path.join(self.wav_dir, self.wav_files[idx])
# Using train and val json files
item = self.dataset[idx]
text, audio_tokens, wav_path = item
wav, sr = torchaudio.load(wav_path)
# mono
if wav.shape[0] > 1:
wav = wav.mean(dim=0, keepdim=True)
# resample if needed
if sr != self.sample_rate:
wav = torchaudio.functional.resample(
wav, orig_freq=sr, new_freq=self.sample_rate
)
# print("wav shape is: ", wav.shape, idx)
# Use the robust OS-aware pipeline to load, convert to mono, and resample
try:
wav, _ = AudioPipeline.load_audio(wav_path, target_sr=self.sample_rate)
except Exception as e:
print(f"Error loading {wav_path}: {e}")
# Fallback to 1 second of silence to prevent the dataloader from crashing the entire run
wav = torch.zeros((1, self.sample_rate))
return wav
codec_train.py
+47 -55
View File
@@ -12,27 +12,8 @@ from codec_model import FSQAutoEncoder
from codec_dataset import LJSpeechDataset
from codec.codec_decoder.decoder import SimpleDecoder
# -----------------------------------------------------------------------------
# Global Configuration
# -----------------------------------------------------------------------------
DATASET_ROOT = "/home/ubuntu/soma/data/lj_speech/LJSpeech-1.1"
SAMPLE_RATE = 32000
BATCH_SIZE = 16
NUM_EPOCHS = 100
LEARNING_RATE = 1e-4
NUM_WORKERS = 8
CKPT_DIR = "/home/ubuntu/soma/ckpt/suprano/suprano_codec/codec_v3"
# Freeze params
FREEZE_ENCODER = False
PRETRAINED_MODEL_PATH = "/home/ubuntu/soma/ckpt/suprano/suprano_codec/codec_1/step_42000.pt"
WANDB_PROJECT = "soprano-codec"
USE_WANDB = True
# -----------------------------------------------------------------------------
# Import the config loader
from config_loader import load_config
def pad_collate(batch):
"""
@@ -67,7 +48,7 @@ def save_mel_plot(original_mel, reconstructed_mel, original_title, reconstructed
plt.close()
def evaluate(model, val_loader_it, val_loader, device, step, plot_dir, val_steps=10):
def evaluate(model, val_loader_it, val_loader, device, step, plot_dir, use_wandb, val_steps=10):
val_loss = 0.0
model.eval()
with torch.no_grad():
@@ -92,7 +73,7 @@ def evaluate(model, val_loader_it, val_loader, device, step, plot_dir, val_steps
val_loss /= val_steps
print(f"step {step} | val_loss {val_loss:.4f}")
if USE_WANDB:
if use_wandb:
wandb.log({"val/loss": val_loss}, step=step)
# Val Plotting
@@ -102,7 +83,7 @@ def evaluate(model, val_loader_it, val_loader, device, step, plot_dir, val_steps
plot_path = os.path.join(plot_dir, f"val_step_{step:05d}.png")
save_mel_plot(vmel_np, vmel_hat_np, "Val Original Mel", "Val Reconstructed Mel", plot_path)
if USE_WANDB:
if use_wandb:
wandb.log({"val/reconstruction": wandb.Image(plot_path)}, step=step)
model.train()
@@ -110,39 +91,47 @@ def evaluate(model, val_loader_it, val_loader, device, step, plot_dir, val_steps
def main():
if USE_WANDB:
wandb.init(project=WANDB_PROJECT)
# ------------------
# Load Configuration
# ------------------
config = load_config("config.yaml")
cfg_global = config["global"]
cfg_paths = config["paths"]
cfg_codec = config["codec"]
if cfg_global["use_wandb"]:
wandb.init(project=cfg_global["wandb_project"], config=config)
# ------------------
# Data Setup
# ------------------
dataset = LJSpeechDataset(
root=DATASET_ROOT,
sample_rate=SAMPLE_RATE,
root=cfg_paths["dataset_root"],
sample_rate=cfg_codec["sample_rate"],
)
loader = DataLoader(
dataset,
batch_size=BATCH_SIZE,
batch_size=cfg_codec["batch_size"],
shuffle=True,
drop_last=True,
num_workers=NUM_WORKERS,
num_workers=cfg_global["num_workers"],
pin_memory=True,
collate_fn=pad_collate
)
val_dataset = LJSpeechDataset(
root=DATASET_ROOT,
sample_rate=SAMPLE_RATE,
root=cfg_paths["dataset_root"],
sample_rate=cfg_codec["sample_rate"],
mode='val'
)
val_loader = DataLoader(
val_dataset,
batch_size=BATCH_SIZE,
batch_size=cfg_codec["batch_size"],
shuffle=False,
drop_last=True,
num_workers=NUM_WORKERS,
num_workers=cfg_global["num_workers"],
pin_memory=True,
collate_fn=pad_collate
)
@@ -154,24 +143,24 @@ def main():
encoder_cfg = dict(
num_input_mels=50,
mel_hop_length=512,
encoder_dim=768,
encoder_num_layers=8,
encoder_dim=cfg_codec["encoder_dim"],
encoder_num_layers=cfg_codec["encoder_num_layers"],
fsq_levels=[8, 8, 5, 5, 5],
)
decoder_cfg = dict(
n_mels=50,
encoder_dim=768,
bottleneck_channels=5,
num_layers=8,
encoder_dim=cfg_codec["encoder_dim"],
bottleneck_channels=cfg_codec["bottleneck_channels"],
num_layers=cfg_codec["decoder_num_layers"],
upsample_scale=2048 // 512,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
device = cfg_global["device"] if torch.cuda.is_available() else "cpu"
# Token rate sanity check
total_hop = 2048 # This is the token hop rate. The mel hop rate is 512
print(f"Token rate: {SAMPLE_RATE / total_hop:.2f} Hz")
print(f"Token rate: {cfg_codec['sample_rate'] / total_hop:.2f} Hz")
# ------------------
# Model Setup
@@ -180,12 +169,13 @@ def main():
# Here, we freeze the encoder and only train the decoder from scratch on learned encoder representation.
# This is to test whether the encoder has learned how to represent the audio well enough.
if FREEZE_ENCODER:
if os.path.exists(PRETRAINED_MODEL_PATH):
print(f"Loading model from {PRETRAINED_MODEL_PATH}")
model.load_state_dict(torch.load(PRETRAINED_MODEL_PATH, map_location=device))
if cfg_codec["freeze_encoder"]:
pretrained_path = cfg_paths["pretrained_codec_path"]
if pretrained_path and os.path.exists(pretrained_path):
print(f"Loading model from {pretrained_path}")
model.load_state_dict(torch.load(pretrained_path, map_location=device))
else:
print(f"Warning: Pretrained model path {PRETRAINED_MODEL_PATH} not found.")
print(f"Warning: Pretrained model path {pretrained_path} not found.")
# fix encoder; train only the decoder; reset the decoder weights.
for param in model.encoder.parameters():
@@ -197,10 +187,12 @@ def main():
model.decoder = SimpleDecoder(**decoder_cfg).to(device)
optimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LEARNING_RATE)
optimizer = Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=float(cfg_codec["learning_rate"]))
plot_dir = os.path.join(CKPT_DIR, "plots")
os.makedirs(CKPT_DIR, exist_ok=True)
# Setup Checkpoint Directories
ckpt_dir = os.path.join(cfg_paths["save_dir"], "codec")
plot_dir = os.path.join(ckpt_dir, "plots")
os.makedirs(ckpt_dir, exist_ok=True)
os.makedirs(plot_dir, exist_ok=True)
step = 0
@@ -208,7 +200,7 @@ def main():
# ------------------
# Training Loop
# ------------------
for epoch in range(NUM_EPOCHS):
for epoch in range(cfg_codec["num_epochs"]):
for epoch_step, data in tqdm(enumerate(loader), total=len(loader)):
@@ -232,13 +224,13 @@ def main():
if step % 100 == 0:
print(f"step {step} | loss {loss.item():.4f}")
if USE_WANDB:
if cfg_global["use_wandb"]:
wandb.log({"train/loss": loss.item()}, step=step)
if step % 400 == 0:
val_loader_it = evaluate(
model, val_loader_it, val_loader, device, step, plot_dir,
val_steps=10
cfg_global["use_wandb"], val_steps=10
)
with torch.no_grad():
@@ -249,7 +241,7 @@ def main():
unique_bins = len(torch.unique(indices))
print(f"[FSQ] unique bins: {unique_bins} / {total_bins}. Lens: {indices.shape} {z.shape}")
if USE_WANDB:
if cfg_global["use_wandb"]:
wandb.log({"train/unique_bins": unique_bins}, step=step)
if step % 200 == 0:
@@ -259,11 +251,11 @@ def main():
plot_path = os.path.join(plot_dir, f"step_{step:05d}.png")
save_mel_plot(mel_np, mel_hat_np, "Original Mel", "Reconstructed Mel", plot_path)
if USE_WANDB:
if cfg_global["use_wandb"]:
wandb.log({"train/reconstruction": wandb.Image(plot_path)}, step=step)
if step % 1000 == 0:
ckpt_path = os.path.join(CKPT_DIR, f"step_{step:05d}.pt")
ckpt_path = os.path.join(ckpt_dir, f"step_{step:05d}.pt")
torch.save(model.state_dict(), ckpt_path)
print(f"Saved checkpoint to {ckpt_path}")
config.yaml
+97
View File
@@ -0,0 +1,97 @@
# ==============================================================================
# LLM-TTS-Factory Configuration
# ==============================================================================
# ------------------------------------------------------------------------------
# Global Settings
# ------------------------------------------------------------------------------
global:
seed: 1337
device: "cuda:0"
num_workers: 4
use_wandb: true
wandb_project: "soprano-tts"
# ------------------------------------------------------------------------------
# File Paths
# Use forward slashes (/) for both Windows and Linux, or standard OS paths.
# Relative paths are evaluated from the directory where the script is run.
#
# Examples:
# Linux: "/home/ubuntu/data/lj_speech/LJSpeech-1.1"
# Windows: "C:/Users/Name/Documents/datasets/LJSpeech-1.1"
# Relative: "./datasets/LJSpeech-1.1"
# ------------------------------------------------------------------------------
paths:
dataset_root: "./data/LJSpeech-1.1"
# Base directory to save all checkpoints and logs
save_dir: "./checkpoints"
# Pretrained model paths (set to null if training from scratch)
# Linux ex: "/home/ubuntu/soma/ckpt/suprano/codec/step_42000.pt"
pretrained_codec_path: null
pretrained_llm_path: null
pretrained_decoder_path: null
pretrained_discriminator_path: null
# ------------------------------------------------------------------------------
# Codec Training Configuration
# ------------------------------------------------------------------------------
codec:
sample_rate: 32000
batch_size: 16
num_epochs: 100
learning_rate: 1.0e-4
freeze_encoder: false
# ------------------------------------------------------------------------------
# LLM Training Configuration
# ------------------------------------------------------------------------------
llm:
from_scratch: false
batch_size: 64
max_steps: 150000
max_lr: 2.0e-5
min_lr_ratio: 0.3
warmup_ratio: 0.3
cooldown_ratio: 0.1
grad_accum_steps: 1
seq_len: 1024
val_freq: 250
save_freq: 5000
text_factor: 0.5
betas: [0.9, 0.95]
weight_decay: 0.1
# ------------------------------------------------------------------------------
# Decoder (Vocos) Training Configuration
# ------------------------------------------------------------------------------
decoder:
use_discriminator: true
batch_size: 64
max_steps: 200000
max_lr: 2.0e-4
min_lr_ratio: 0.1
warmup_ratio: 0.2
cooldown_ratio: 0.1
grad_accum_steps: 1
seq_len: 1024
segment_size_samples: 32768 # ~1 sec (16 tokens)
val_freq: 250
text_factor: 0.0
betas: [0.8, 0.99]
weight_decay: 0.1
# Loss Weights
lambda_mel: 45.0
lambda_fm: 2.0
lambda_gen: 1.0
lambda_stft: 1.0
# ------------------------------------------------------------------------------
# Data Generation Configuration (generate_dataset*.py)
# ------------------------------------------------------------------------------
data_generation:
val_prop: 0.1
val_max: 512
dataset_e2e.py
+36
View File
@@ -0,0 +1,36 @@
import json
import torch
from torch.utils.data import Dataset
from utils.audio_utils import AudioPipeline
# The codec downsamples audio by a factor of 2048.
# At 32kHz, 1 token = 2048 raw audio samples.
SAMPLES_PER_TOKEN = 2048
class AudioDataset(Dataset):
def __init__(self, path, target_sr=32000):
with open(path, encoding='utf-8') as f:
self.dataset = json.load(f)
self.target_sr = target_sr
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
# The JSON format from generate_dataset.py is:
# [transcript, audio_tokens (list), audio_path]
text, audio_tokens, audio_path = self.dataset[idx]
# Use our robust OS-aware pipeline to load the audio
try:
wav, _ = AudioPipeline.load_audio(audio_path, target_sr=self.target_sr)
# Squeeze out the channel dimension so it's a 1D tensor (T,)
# as expected by train_decoder.py's alignment logic
wav = wav.squeeze(0)
except Exception as e:
print(f"Error loading {audio_path}: {e}")
# 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)
decoder/__pycache__/decoder.cpython-311.pyc
BIN
View File
Binary file not shown.
decoder/__pycache__/discriminator.cpython-311.pyc
BIN
View File
Binary file not shown.
decoder/__pycache__/heads.cpython-311.pyc
BIN
View File
Binary file not shown.
decoder/__pycache__/losses.cpython-311.pyc
BIN
View File
Binary file not shown.
decoder/__pycache__/models.cpython-311.pyc
BIN
View File
Binary file not shown.
decoder/__pycache__/modules.cpython-311.pyc
BIN
View File
Binary file not shown.
decoder/__pycache__/spectral_ops.cpython-311.pyc
BIN
View File
Binary file not shown.
generate_dataset.py
+94 -106
View File
@@ -1,133 +1,121 @@
"""
Converts a dataset in LJSpeech format into audio tokens that can be used to train/fine-tune Soprano.
This script creates two JSON files for train and test splits in the provided directory.
Converts a dataset in LJSpeech format into audio tokens for Soprano, using pre-defined train/val lists.
Usage:
python generate_dataset.py --input-dir path/to/files
Args:
--input-dir: Path to directory of LJSpeech-style dataset. If none is provided this defaults to the provided example dataset.
python generate_dataset_from_lists.py
"""
import argparse
import pathlib
import random
import json
import torchaudio
import os
import torch
from tqdm import tqdm
from huggingface_hub import hf_hub_download
from encoder.codec import Encoder
from config_loader import load_config
from utils.audio_utils import AudioPipeline
SAMPLE_RATE = 32000
SEED = 42
VAL_PROP = 0.1
VAL_MAX = 512
def load_metadata(input_dir):
print("Reading metadata...")
meta_map = {}
meta_path = input_dir / 'metadata_orig.csv'
if not meta_path.exists():
meta_path = input_dir / 'metadata.csv'
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--input-dir",
required=False,
default="./example_dataset",
type=pathlib.Path
)
return parser.parse_args()
with open(meta_path, encoding='utf-8') as f:
for line in f:
if not line.strip(): continue
parts = line.strip().split('|')
filename = parts[0]
transcript = parts[-1]
meta_map[filename] = transcript
return meta_map
def main():
args = get_args()
input_dir = args.input_dir
use_custom_model = True
print("Loading model.")
encoder = Encoder()
if not use_custom_model:
encoder_path = hf_hub_download(repo_id='ekwek/Soprano-Encoder', filename='encoder.pth')
encoder.load_state_dict(torch.load(encoder_path))
else:
speech_autoencoder_path = "/home/ubuntu/soma/ckpt/suprano/suprano_codec/codec_1/step_40000.pt"
print("Loading model using custom model path!", speech_autoencoder_path)
full_ckpt = torch.load(speech_autoencoder_path)
encoder_state_dict = {k.replace("encoder.", ""): v for k, v in full_ckpt.items() if k.startswith("encoder.")}
encoder_state_dict = {}
for k,v in full_ckpt.items():
if k.startswith("encoder."):
# replace the first occurance of 'encoder.' only
new_k = k.replace("encoder.", "", 1)
encoder_state_dict[new_k] = v
encoder.load_state_dict(encoder_state_dict)
print("Model loaded.")
import pdb;pdb.set_trace()
print("Reading metadata.")
files = []
with open(f'{input_dir}/metadata_orig.csv', encoding='utf-8') as f:
data = f.read().split('\n')
for line in data:
# import pdb;pdb.set_trace()
out = line.split("|", maxsplit=1)
filename = out[0]
transcript = out[-1].split('|')[-1]
files.append((filename, transcript))
print(f'{len(files)} samples located in directory.')
import pdb;pdb.set_trace()
print("Encoding audio.")
def process_list(list_file, meta_map, encoder, target_sr):
dataset = []
for sample in tqdm(files):
filename, transcript = sample
# sr, audio = wavfile.read(f'{input_dir}/wavs/{filename}.wav')
# audio = torch.from_numpy(audio)
print(f"Processing {list_file}...")
with open(list_file, 'r') as f:
lines = [l.strip() for l in f if l.strip()]
import pdb;pdb.set_trace()
for line in tqdm(lines):
path_obj = pathlib.Path(line)
filename = path_obj.stem # LJxxx
if filename not in meta_map:
print(f"Warning: {filename} not found in metadata. Skipping.")
continue
transcript = meta_map[filename]
wav_path = str(path_obj)
# Load and Encode with OS-aware pipeline
try:
audio, sr = torchaudio.load(f'{input_dir}/wavs/{filename}.wav')
except:
print("Error loading audio: ", filename)
audio, _ = AudioPipeline.load_audio(wav_path, target_sr)
except Exception as e:
print(f"Error loading {wav_path}: {e}")
continue
# import pdb;pdb.set_trace()
if sr != SAMPLE_RATE:
audio = torchaudio.functional.resample(audio, sr, SAMPLE_RATE)
# audio = audio.unsqueeze(0)
with torch.no_grad():
audio_tokens = encoder(audio)
# Save absolute path to audio for loading in training
audio_path = str(pathlib.Path(f'{input_dir}/wavs/{filename}.wav').resolve())
dataset.append([transcript, audio_tokens.squeeze(0).tolist(), audio_path])
print("Generating train/test splits.")
random.seed(SEED)
random.shuffle(dataset)
num_val = min(int(VAL_PROP * len(dataset)) + 1, VAL_MAX)
train_dataset = dataset[num_val:]
val_dataset = dataset[:num_val]
print(f'# train samples: {len(train_dataset)}')
print(f'# val samples: {len(val_dataset)}')
dataset.append([transcript, audio_tokens.squeeze(0).tolist(), wav_path])
print("Saving datasets.")
with open(f'{input_dir}/train.json', 'w', encoding='utf-8') as f:
json.dump(train_dataset, f, indent=2)
with open(f'{input_dir}/val.json', 'w', encoding='utf-8') as f:
json.dump(val_dataset, f, indent=2)
print("Datasets saved.")
return dataset
def main():
config = load_config("config.yaml")
cfg_paths = config["paths"]
cfg_codec = config["codec"]
input_dir = pathlib.Path(cfg_paths["dataset_root"])
output_dir = pathlib.Path(cfg_paths["save_dir"]) / "dataset_lists"
os.makedirs(output_dir, exist_ok=True)
target_sr = cfg_codec["sample_rate"]
device = config["global"]["device"] if torch.cuda.is_available() else 'cpu'
# Load Encoder
print("Loading Encoder...")
encoder = Encoder()
speech_autoencoder_path = cfg_paths["pretrained_codec_path"]
if not speech_autoencoder_path or not os.path.exists(speech_autoencoder_path):
raise FileNotFoundError(f"pretrained_codec_path not found: {speech_autoencoder_path}")
print(f"Loading weights from {speech_autoencoder_path}")
full_ckpt = torch.load(speech_autoencoder_path, map_location='cpu')
encoder_state_dict = {}
for k, v in full_ckpt.items():
if k.startswith("encoder."):
new_k = k.replace("encoder.", "", 1)
encoder_state_dict[new_k] = v
encoder.load_state_dict(encoder_state_dict)
encoder.to(device)
encoder.eval()
print("Encoder Loaded.")
meta_map = load_metadata(input_dir)
# Process Train List
train_list_path = input_dir / 'train_list.txt'
if train_list_path.exists():
train_data = process_list(train_list_path, meta_map, encoder, target_sr)
with open(output_dir / 'train.json', 'w') as f:
json.dump(train_data, f, indent=2)
print(f"Saved {len(train_data)} train samples to {output_dir}/train.json")
else:
print(f"Error: {train_list_path} not found.")
# Process Val List
val_list_path = input_dir / 'val_list.txt'
if val_list_path.exists():
val_data = process_list(val_list_path, meta_map, encoder, target_sr)
with open(output_dir / 'val.json', 'w') as f:
json.dump(val_data, f, indent=2)
print(f"Saved {len(val_data)} val samples to {output_dir}/val.json")
else:
print(f"Error: {val_list_path} not found.")
if __name__ == '__main__':
main()
generate_dataset_from_lists.py
+28 -43
View File
@@ -2,58 +2,41 @@
Converts a dataset in LJSpeech format into audio tokens for Soprano, using pre-defined train/val lists.
Usage:
python generate_dataset_from_lists.py --input-dir /home/ubuntu/soma/data/lj_speech/LJSpeech-1.1 --output-dir ./dataset_lists
python generate_dataset_from_lists.py
"""
import argparse
import pathlib
import json
import os
import torchaudio
import torch
from tqdm import tqdm
from encoder.codec import Encoder
SAMPLE_RATE = 32000
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--input-dir",
required=True,
type=pathlib.Path,
help="Path to LJSpeech-1.1 directory containing wavs/, metadata.csv, train_list.txt, val_list.txt"
)
parser.add_argument("--output-dir",
required=True,
type=pathlib.Path,
help="Directory to save new train.json and val.json"
)
return parser.parse_args()
from config_loader import load_config
from utils.audio_utils import AudioPipeline
def load_metadata(input_dir):
print("Reading metadata.")
print("Reading metadata...")
meta_map = {}
# Check for metadata.csv or metadata_orig.csv
meta_path = input_dir / 'metadata_orig.csv'
if not meta_path.exists():
meta_path = input_dir / 'metadata.csv'
with open(meta_path, encoding='utf-8') as f:
for line in f:
if not line.strip(): continue
parts = line.strip().split('|')
filename = parts[0]
transcript = parts[-1] # formatting might vary, usually last field
transcript = parts[-1]
meta_map[filename] = transcript
return meta_map
def process_list(list_file, meta_map, encoder, input_dir):
def process_list(list_file, meta_map, encoder, target_sr, device):
dataset = []
print(f"Processing {list_file}...")
with open(list_file, 'r') as f:
lines = [l.strip() for l in f if l.strip()]
for line in tqdm(lines):
# line is full path: /home/ubuntu/.../wavs/LJxxx.wav
# Extract filename without extension
path_obj = pathlib.Path(line)
filename = path_obj.stem # LJxxx
@@ -64,45 +47,47 @@ def process_list(list_file, meta_map, encoder, input_dir):
transcript = meta_map[filename]
wav_path = str(path_obj)
# Load and Encode
# Load and Encode with OS-aware pipeline
try:
audio, sr = torchaudio.load(wav_path)
audio, _ = AudioPipeline.load_audio(wav_path, target_sr)
except Exception as e:
print(f"Error loading {wav_path}: {e}")
continue
if sr != SAMPLE_RATE:
audio = torchaudio.functional.resample(audio, sr, SAMPLE_RATE)
# Mono check
if audio.shape[0] > 1:
audio = audio.mean(dim=0, keepdim=True)
audio = audio.to(device)
with torch.no_grad():
audio_tokens = encoder(audio) # Encoder expects (B, C, T) or (1, T)?
audio_tokens = encoder(audio)
# If it returns indices directly.
dataset.append([transcript, audio_tokens.squeeze(0).tolist(), wav_path])
return dataset
def main():
args = get_args()
input_dir = args.input_dir
output_dir = args.output_dir
config = load_config("config.yaml")
cfg_paths = config["paths"]
cfg_codec = config["codec"]
input_dir = pathlib.Path(cfg_paths["dataset_root"])
# Save lists into the configured save_dir
output_dir = pathlib.Path(cfg_paths["save_dir"]) / "dataset_lists"
os.makedirs(output_dir, exist_ok=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
target_sr = cfg_codec["sample_rate"]
device = config["global"]["device"] if torch.cuda.is_available() else 'cpu'
# Load Encoder
print("Loading Encoder...")
encoder = Encoder()
speech_autoencoder_path = "/home/ubuntu/soma/ckpt/suprano/suprano_codec/codec_v2/step_30000.pt"
speech_autoencoder_path = cfg_paths["pretrained_codec_path"]
if not speech_autoencoder_path or not os.path.exists(speech_autoencoder_path):
raise FileNotFoundError(f"pretrained_codec_path not found: {speech_autoencoder_path}")
print(f"Loading weights from {speech_autoencoder_path}")
full_ckpt = torch.load(speech_autoencoder_path, map_location='cpu')
# Extract encoder weights
encoder_state_dict = {}
for k, v in full_ckpt.items():
if k.startswith("encoder."):
@@ -110,16 +95,16 @@ def main():
encoder_state_dict[new_k] = v
encoder.load_state_dict(encoder_state_dict)
encoder.to(device)
encoder.eval()
print("Encoder Loaded.")
# Load Metadata
meta_map = load_metadata(input_dir)
# Process Train List
train_list_path = input_dir / 'train_list.txt'
if train_list_path.exists():
train_data = process_list(train_list_path, meta_map, encoder, input_dir)
train_data = process_list(train_list_path, meta_map, encoder, target_sr, device)
with open(output_dir / 'train.json', 'w') as f:
json.dump(train_data, f, indent=2)
print(f"Saved {len(train_data)} train samples to {output_dir}/train.json")
@@ -129,7 +114,7 @@ def main():
# Process Val List
val_list_path = input_dir / 'val_list.txt'
if val_list_path.exists():
val_data = process_list(val_list_path, meta_map, encoder, input_dir)
val_data = process_list(val_list_path, meta_map, encoder, target_sr, device)
with open(output_dir / 'val.json', 'w') as f:
json.dump(val_data, f, indent=2)
print(f"Saved {len(val_data)} val samples to {output_dir}/val.json")
pyproject.toml
+35
View File
@@ -0,0 +1,35 @@
[project]
name = "llm-tts-factory"
version = "0.1.0"
description = "End-to-End LLM-Backbone TTS Training Framework"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"einops",
"huggingface-hub",
"matplotlib",
"numpy",
"pyyaml",
"safetensors",
"scipy",
"soundfile",
"torch",
"torchaudio",
"tqdm",
"transformers",
"wandb",
]
[tool.uv]
# Tells uv this is an application/scripts directory, not a library to be built
package = false
[tool.uv.sources]
# Explicitly pull these from the custom PyTorch index
torch = { index = "pytorch" }
torchaudio = { index = "pytorch" }
[[tool.uv.index]]
name = "pytorch"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
requirements.txt
-8
View File
@@ -1,8 +0,0 @@
einops
huggingface_hub
numpy
scipy
torch
torchaudio
tqdm
transformers
simple_inference.py
+31 -35
View File
@@ -4,10 +4,17 @@ import argparse
import os
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
from safetensors.torch import load_file
# Ensure decoder module is importable
from decoder.decoder import SopranoDecoder
from config_loader import load_config
def load_models(llm_path, decoder_path, device='cuda'):
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):
raise FileNotFoundError(f"Decoder path invalid or not found: {decoder_path}. Please check your config.yaml.")
print(f"Loading LLM from {llm_path}...")
# Load LLM Config & Model
@@ -16,7 +23,6 @@ def load_models(llm_path, decoder_path, device='cuda'):
# Load LLM weights
if llm_path.endswith('.safetensors'):
print("loading llm from custom trained model: ", llm_path)
state_dict = load_file(llm_path)
llm.load_state_dict(state_dict)
else:
@@ -26,20 +32,17 @@ def load_models(llm_path, decoder_path, device='cuda'):
llm.to(device).eval()
print(f"Loading Decoder from {decoder_path}...")
# Instantiate Decoder with defaults (matching train_decoder.py)
# If training used non-defaults, user must manually edit this line.
# Instantiate Decoder with defaults
decoder = SopranoDecoder()
# Load Decoder weights
# Map to cpu first to avoid OOM or device mismatch during load
print("Loading decoder from: ", decoder_path)
decoder_state = torch.load(decoder_path, map_location='cpu')
decoder.load_state_dict(decoder_state)
decoder.to(device).eval()
return llm, decoder
def generate_audio(text, llm, decoder, tokenizer, device='cuda', save_path="output.wav"):
def generate_audio(text, llm, decoder, tokenizer, cfg_inf, device='cuda', save_path="output.wav"):
# 1. Format Prompt
prompt = f"[TEXT]{text}[START]"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
@@ -48,47 +51,35 @@ def generate_audio(text, llm, decoder, tokenizer, device='cuda', save_path="outp
print("Generating Tokens & Extracting Hidden States...")
# 2. Generate with Hidden States Extraction
if 'token_type_ids' in inputs: del inputs['token_type_ids']
if 'token_type_ids' in inputs:
del inputs['token_type_ids']
with torch.no_grad():
outputs = llm.generate(
input_ids=inputs['input_ids'],
attention_mask=inputs['attention_mask'] if 'attention_mask' in inputs else None,
max_new_tokens=512,
attention_mask=inputs.get('attention_mask'),
max_new_tokens=cfg_inf["max_new_tokens"],
do_sample=True,
temperature=0.8,
top_k=50,
top_p=0.95,
temperature=cfg_inf["temperature"],
top_k=cfg_inf["top_k"],
top_p=cfg_inf["top_p"],
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
return_dict_in_generate=True,
output_hidden_states=True,
repetition_penalty=1.2
repetition_penalty=cfg_inf["repetition_penalty"]
)
# import pdb;pdb.set_trace()
# 3. Process Hidden States
# outputs.hidden_states is a tuple of tuples.
# Outer tuple: one per generation step.
# Inner tuple: one per layer.
hidden_states_list = []
# We only care about the last layer
# outputs.hidden_states is 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.
# last_layer_state = step_states[-1][-1]
last_layer_state = step_states[-1][0, -1, :]
# Shape: (Batch, 1, Dim)
# print("shape of last layer state is: ", last_layer_state.shape)
hidden_states_list.append(last_layer_state)
import pdb;pdb.set_trace()
# Concatenate along time dimension
# Result: (Batch, T_gen, Dim)
audio_hidden = torch.stack(hidden_states_list).unsqueeze(0) # (B, T, D)
@@ -122,22 +113,27 @@ def generate_audio(text, llm, decoder, tokenizer, device='cuda', save_path="outp
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--text", type=str, required=True, help="Text to synthesize")
parser.add_argument("--llm-path", type=str, required=True, help="Path to LLM checkpoint")
parser.add_argument("--decoder-path", type=str, required=True, help="Path to Decoder checkpoint")
parser.add_argument("--out", type=str, default="output.wav", help="Output content")
parser.add_argument("--out", type=str, default="output.wav", help="Output audio file name")
args = parser.parse_args()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Load configuration
config = load_config("config.yaml")
cfg_global = config["global"]
cfg_paths = config["paths"]
cfg_inf = config["inference"]
device = cfg_global["device"] if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
tokenizer.eos_token_id = 3
llm, decoder = load_models(args.llm_path, args.decoder_path, device)
llm_path = cfg_paths["pretrained_llm_path"]
decoder_path = cfg_paths["pretrained_decoder_path"]
generate_audio(args.text, llm, decoder, tokenizer, device, args.out)
llm, decoder = load_models(llm_path, decoder_path, device)
generate_audio(args.text, llm, decoder, tokenizer, cfg_inf, device, args.out)
if __name__ == "__main__":
main()
train_decoder.py
+166 -255
View File
@@ -2,98 +2,41 @@
Training script for Soprano Decoder (Vocos).
Freezes LLM and trains Decoder with GAN loss.
"""
import argparse
import pathlib
import random
import time
import os
import wandb
import matplotlib.pyplot as plt
import io
import numpy as np
import torch
import torchaudio
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from dataset_e2e import AudioDataset, SAMPLES_PER_TOKEN
from decoder.decoder import SopranoDecoder
from decoder.discriminator import Discriminator
from decoder.losses import MelSpectrogramWrapper, feature_matching_loss, discriminator_loss, generator_loss, MultiResolutionSTFTLoss
# Initialize Mel Spectrogram Wrapper
mel_fn = MelSpectrogramWrapper().to('cuda')
from config_loader import load_config
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--input-dir",
required=False,
default="./example_dataset",
type=pathlib.Path
)
parser.add_argument("--save-dir",
required=True,
type=pathlib.Path
)
parser.add_argument("--use-disc",
action="store_true",
help="Whether to use discriminator and GAN losses. Defaults to False (only reconstruction loss)."
)
return parser.parse_args()
args = get_args()
# training hyperparameters
device = 'cuda:0'
seed = 1337
max_lr = 2e-4 # Lower LR for GAN usually
warmup_ratio = 0.2
cooldown_ratio = 0.1
min_lr = 0.1 * max_lr
batch_size = 64 # 64
grad_accum_steps = 1
seq_len = 1024
SEGMENT_SIZE_SAMPLES = 32768 # ~1 sec (16 tokens)
val_freq = 250
text_factor = 0.0
max_steps = 200000
betas = (0.8, 0.99) # GAN betas
weight_decay = 0.1
train_dataset_path = f'{args.input_dir}/train.json'
val_dataset_path = f'{args.input_dir}/val.json'
save_path = args.save_dir
os.makedirs(save_path, exist_ok=True)
# Loss Weights
lambda_mel = 45.0
lambda_fm = 2.0
lambda_gen = 1.0
lambda_stft = 1.0
# Global Tokenizer for collate function
tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
tokenizer.padding_side = 'right' # Essential for training!
def worker_seed_init(_):
worker_seed = torch.initial_seed() % (2**32-1)
np.random.seed(worker_seed)
random.seed(worker_seed)
def get_lr(it): # WSD schedule
if it<warmup_steps:
return max_lr * (it+1) / warmup_steps
if it<max_steps-cooldown_steps:
def get_lr(it, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps): # WSD schedule
if it < warmup_steps:
return max_lr * (it + 1) / warmup_steps
if it < max_steps - cooldown_steps:
return max_lr
return min_lr + (max_lr-min_lr) * ((max_steps-it) / cooldown_steps)
# ... (Previous code)
# Initialize Mel Spectrogram Wrapper
# Place this after device definition (Line 66)
# But here we are replacing lines 26-43 and 102-130.
# I will do this in chunks.
return min_lr + (max_lr - min_lr) * ((max_steps - it) / cooldown_steps)
def collate_pack(batch_in):
# batch_in is list of (text, wav)
@@ -101,21 +44,14 @@ def collate_pack(batch_in):
wavs = [x[1] for x in batch_in]
aud_token_lens = [x[2] for x in batch_in]
tokens_batch = tokenizer(texts, padding=True, return_tensors='pt')
input_ids = tokens_batch['input_ids'] # (B, T)
# We need to process each sample to align audio
# Since lengths vary, we process list then pad
batch_tokens_list = []
batch_audio_list = []
for i in range(len(texts)):
# Get raw tokens without padding for alignment logic
raw_tokens = tokenizer(texts[i], padding=False, truncation=False)['input_ids']
tokens = torch.tensor(raw_tokens, dtype=torch.long) # remove last token as per original logic? Or keep?
# Taking raw tokens
tokens = torch.tensor(raw_tokens, dtype=torch.long)
wav = wavs[i]
@@ -125,12 +61,6 @@ def collate_pack(batch_in):
is_audio = (tokens > 3) & (tokens <= 8003)
audio_indices = torch.where(is_audio)[0]
if len(audio_indices) != num_aud_tokens:
print(f"Audio token count mismatch: {len(audio_indices)} vs {num_aud_tokens}")
print(texts[i])
print(raw_tokens)
print(is_audio)
print(audio_indices)
assert len(audio_indices) == num_aud_tokens, f"Audio token count mismatch: {len(audio_indices)} vs {num_aud_tokens}"
for pos, idx in enumerate(audio_indices):
@@ -144,70 +74,95 @@ def collate_pack(batch_in):
batch_audio_list.append(aligned_audio)
# Pad Tokens
# We need to return x and y. So we need sequence length T.
# tokens are length T+1 (input + target).
# Pad to max length
batch_tokens = torch.nn.utils.rnn.pad_sequence(batch_tokens_list, batch_first=True, padding_value=0)
# Pad Audio
# Audio length = tokens_length * SAMPLES_PER_TOKEN
# Since 0 token -> 2048 zeros, padding matches.
batch_audio = torch.nn.utils.rnn.pad_sequence(batch_audio_list, batch_first=True, padding_value=0.0)
x = batch_tokens[:, :-1]
y = batch_tokens[:, 1:]
# Audio matches x (length T)
# batch_audio was created for length T+1 (full tokens).
# We need audio corresponding to x.
# Dimensions: batch_audio is (B, (T+1)*2048)
# We want up to T*2048
# Calculate max seq len of x
max_len_x = x.size(1)
gt_audio = batch_audio[:, :max_len_x * SAMPLES_PER_TOKEN]
# Create Attention Mask
# Start with x mask
# x_mask = (x != 0) # Assumes 0 is padding value used above
# Create Audio Mask (True where token is audio)
audio_mask = (y > 3) & (y <= 8003) # Compute output based on y. I should start collecting outputs from the start_token in x.
audio_mask = (y > 3) & (y <= 8003)
return x, y, gt_audio, audio_mask
tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
tokenizer.padding_side = 'right' # Essential for training!
if __name__ == '__main__':
# ------------------
# Load Configuration
# ------------------
config = load_config("config.yaml")
cfg_global = config["global"]
cfg_paths = config["paths"]
cfg_decoder = config["decoder"]
device = cfg_global["device"]
seed = cfg_global["seed"]
device_type = "cuda" if device.startswith("cuda") else "cpu"
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.set_float32_matmul_precision('high')
# Setup directories
train_dataset_path = os.path.join(cfg_paths["dataset_root"], "train.json")
val_dataset_path = os.path.join(cfg_paths["dataset_root"], "val.json")
save_path = os.path.join(cfg_paths["save_dir"], "decoder")
os.makedirs(save_path, exist_ok=True)
print(f"Save Path: {save_path}")
# lr schedule
warmup_steps = int(max_steps * warmup_ratio)
cooldown_steps = int(max_steps * cooldown_ratio)
if cfg_global["use_wandb"]:
wandb.init(project=cfg_global["wandb_project"], config=config)
# Initialize WandB
wandb.init(project="soprano-decoder-only", config=vars(args))
# Initialize Mel Spectrogram Wrapper dynamically
mel_fn = MelSpectrogramWrapper().to(device)
# ------------------
# Hyperparameters
# ------------------
max_steps = cfg_decoder["max_steps"]
max_lr = float(cfg_decoder["max_lr"])
min_lr = cfg_decoder["min_lr_ratio"] * max_lr
warmup_steps = int(max_steps * cfg_decoder["warmup_ratio"])
cooldown_steps = int(max_steps * cfg_decoder["cooldown_ratio"])
batch_size = cfg_decoder["batch_size"]
segment_size_samples = cfg_decoder["segment_size_samples"]
val_freq = cfg_decoder["val_freq"]
save_freq = cfg_decoder["save_freq"]
betas = tuple(cfg_decoder["betas"])
weight_decay = cfg_decoder["weight_decay"]
start_step = cfg_decoder.get("start_step", 0)
# Loss Weights
lambda_mel = cfg_decoder["lambda_mel"]
lambda_fm = cfg_decoder["lambda_fm"]
lambda_gen = cfg_decoder["lambda_gen"]
lambda_stft = cfg_decoder["lambda_stft"]
# ------------------
# 1. Load LLM and Freeze
# Load custom trained checkpoint
print("Loading LLM from custom checkpoint...")
config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
model = AutoModelForCausalLM.from_config(config)
# ------------------
print("Loading LLM...")
llm_config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
model = AutoModelForCausalLM.from_config(llm_config)
from safetensors.torch import load_file
# ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_llm/codec_1/model.safetensors"
# ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_llm/codec_1_2/checkpoint-10000/model.safetensors"
# ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_llm/codec_v2/v2/model.safetensors"
ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_llm/codec_v2/v2/checkpoint-40000/model.safetensors"
# ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_llm/codec_v2/v2/checkpoint-135000/model.safetensors"
state_dict = load_file(ckpt_path)
model.load_state_dict(state_dict)
pretrained_llm_path = cfg_paths["pretrained_llm_path"]
if pretrained_llm_path and os.path.exists(pretrained_llm_path):
print(f"Loading custom LLM checkpoint from {pretrained_llm_path}")
if pretrained_llm_path.endswith('.safetensors'):
state_dict = load_file(pretrained_llm_path)
model.load_state_dict(state_dict)
else:
model = AutoModelForCausalLM.from_pretrained(pretrained_llm_path)
else:
print("Warning: No pretrained LLM provided for Decoder training. Using random init (not recommended).")
model.to(torch.bfloat16).to(device)
model.eval()
@@ -215,13 +170,18 @@ if __name__ == '__main__':
param.requires_grad = False
print("LLM Frozen.")
# 2. Decoder
# ------------------
# 2. Load Decoder
# ------------------
print("Loading Decoder...")
decoder = SopranoDecoder()
# decoder_ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_deocoder/codec_1_fix_lens/decoder_step_10000.pth"
# decoder_ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_deocoder/codec_v2/v2_40k_w_stft_loss/decoder_trained.pth"
decoder_ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_deocoder/codec_v2/v2_40k_w_stft_loss_w_disc/decoder_step_141000.pth"
decoder.load_state_dict(torch.load(decoder_ckpt_path, map_location='cpu'))
pretrained_decoder_path = cfg_paths["pretrained_decoder_path"]
if pretrained_decoder_path and os.path.exists(pretrained_decoder_path):
print(f"Loading custom Decoder checkpoint from {pretrained_decoder_path}")
decoder.load_state_dict(torch.load(pretrained_decoder_path, map_location='cpu'))
else:
print("Training Decoder from scratch.")
decoder.to(device)
decoder.train()
@@ -230,27 +190,33 @@ if __name__ == '__main__':
# Initialize MR-STFT Loss
mr_stft = MultiResolutionSTFTLoss().to(device)
# 3. Discriminator
# ------------------
# 3. Load Discriminator
# ------------------
discriminator = None
if args.use_disc:
if cfg_decoder["use_discriminator"]:
print("Initializing Discriminator...")
discriminator = Discriminator()
pretrained_disc_path = cfg_paths["pretrained_discriminator_path"]
disc_model_path = "/home/ubuntu/soma/ckpt/suprano/suprano_deocoder/codec_v2/v2_40k_w_stft_loss_w_disc/discriminator_step_141000.pth"
discriminator.load_state_dict(torch.load(disc_model_path, map_location='cpu'))
if pretrained_disc_path and os.path.exists(pretrained_disc_path):
print(f"Loading custom Discriminator checkpoint from {pretrained_disc_path}")
discriminator.load_state_dict(torch.load(pretrained_disc_path, map_location='cpu'))
discriminator.to(device)
discriminator.train()
else:
print("Training WITHOUT Discriminator (Reconstruction only).")
# 4. Dataset
# ------------------
# 4. Dataset Setup
# ------------------
dataset = AudioDataset(train_dataset_path)
dataloader = DataLoader(dataset,
batch_size=8, # change this back
dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4,
num_workers=cfg_global["num_workers"],
pin_memory=True,
persistent_workers=True,
worker_init_fn=worker_seed_init,
@@ -259,10 +225,11 @@ if __name__ == '__main__':
dataloader_it = iter(dataloader)
val_dataset = AudioDataset(val_dataset_path)
val_dataloader = DataLoader(val_dataset,
batch_size=16,
val_dataloader = DataLoader(
val_dataset,
batch_size=max(1, batch_size // 4), # Reduce val batch size to prevent OOM
shuffle=False,
num_workers=4,
num_workers=max(1, cfg_global["num_workers"] // 2),
pin_memory=True,
persistent_workers=True,
worker_init_fn=worker_seed_init,
@@ -270,16 +237,19 @@ if __name__ == '__main__':
)
val_dataloader_it = iter(val_dataloader)
# import pdb;pdb.set_trace()
# ------------------
# 5. Optimizers
# ------------------
opt_g = torch.optim.AdamW(decoder.parameters(), max_lr, betas=betas, weight_decay=weight_decay)
opt_d = None
if args.use_disc:
if cfg_decoder["use_discriminator"]:
opt_d = torch.optim.AdamW(discriminator.parameters(), max_lr, betas=betas, weight_decay=weight_decay)
start_step = 141001
pbar = tqdm(range(start_step, max_steps), ncols=200, dynamic_ncols=True)
# ------------------
# Training Loop
# ------------------
pbar = tqdm(range(start_step + 1, max_steps + 1), ncols=200, dynamic_ncols=True)
for step in pbar:
start = time.time()
@@ -289,10 +259,7 @@ if __name__ == '__main__':
if batch_data[0] is None:
dataloader_it = iter(dataloader)
batch_data = next(dataloader_it)
x, y, gt_audio, audio_mask = batch_data # Audio mask is the audio token mask for token prediction loss here.
# import pdb;pdb.set_trace()
x, y, gt_audio, audio_mask = batch_data
except StopIteration:
dataloader_it = iter(dataloader)
batch_data = next(dataloader_it)
@@ -310,9 +277,6 @@ if __name__ == '__main__':
hidden_states = hidden_states.to(torch.float32)
# GATHER AUDIO LATENTS Logic
# We need to extract only the hidden states where audio_mask is True.
# Since number of audio tokens varies per sample, we gather and then Pad.
gathered_states_list = []
for b_idx in range(hidden_states.size(0)):
mask = audio_mask[b_idx]
@@ -321,7 +285,6 @@ if __name__ == '__main__':
decoder_in_padded = torch.nn.utils.rnn.pad_sequence(gathered_states_list, batch_first=True, padding_value=0.0)
bsz = decoder_in_padded.size(0)
max_aud_len = decoder_in_padded.size(1)
audio_loss_mask = torch.zeros((bsz, max_aud_len), dtype=torch.bool, device=device)
@@ -330,26 +293,21 @@ if __name__ == '__main__':
audio_loss_mask[b_idx, :length] = True
# ---------------------
# Train Discriminator
# Generator Forward
# ---------------------
d_loss_item = 0.0
if args.use_disc:
opt_d.zero_grad()
# Generator Forward (Detach for D training)
decoder_in = decoder_in_padded.transpose(1, 2) # (B, C, T)
fake_audio = decoder(decoder_in) # (B, 1, T_audio_gen)
if fake_audio.size(1) == 1: fake_audio = fake_audio.squeeze(1)
decoder_in = decoder_in_padded.transpose(1, 2) # (B, C, T)
fake_audio = decoder(decoder_in) # (B, 1, T_audio_gen)
if fake_audio.size(1) == 1: fake_audio = fake_audio.squeeze(1)
min_len = min(fake_audio.size(1), gt_audio.size(1))
fake_audio = fake_audio[:, :min_len]
real_audio = gt_audio[:, :min_len]
# ---------------------
# Train Discriminator (on Crops)
# Train Discriminator
# ---------------------
d_loss_item = 0.0
if args.use_disc:
if cfg_decoder["use_discriminator"]:
opt_d.zero_grad()
# --- Random Cropping Logic ---
@@ -357,33 +315,29 @@ if __name__ == '__main__':
fake_crop_list = []
for b_idx in range(bsz):
# Calculate valid audio length for this sample
valid_len = gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN
valid_len = min(valid_len, min_len) # Clamp to generated length
valid_len = min(valid_len, min_len)
if valid_len <= SEGMENT_SIZE_SAMPLES:
# Pad if shorter
pad_len = SEGMENT_SIZE_SAMPLES - valid_len
if valid_len <= segment_size_samples:
pad_len = segment_size_samples - valid_len
r_c = torch.nn.functional.pad(real_audio[b_idx, :valid_len], (0, pad_len))
f_c = torch.nn.functional.pad(fake_audio[b_idx, :valid_len], (0, pad_len))
else:
# Random Crop
start_idx = random.randint(0, valid_len - SEGMENT_SIZE_SAMPLES)
r_c = real_audio[b_idx, start_idx : start_idx + SEGMENT_SIZE_SAMPLES]
f_c = fake_audio[b_idx, start_idx : start_idx + SEGMENT_SIZE_SAMPLES]
start_idx = random.randint(0, valid_len - segment_size_samples)
r_c = real_audio[b_idx, start_idx : start_idx + segment_size_samples]
f_c = fake_audio[b_idx, start_idx : start_idx + segment_size_samples]
real_crop_list.append(r_c)
fake_crop_list.append(f_c)
real_crops = torch.stack(real_crop_list).unsqueeze(1) # (B, 1, T_seg)
fake_crops = torch.stack(fake_crop_list).unsqueeze(1).detach() # Detach for D update
real_crops = torch.stack(real_crop_list).unsqueeze(1)
fake_crops = torch.stack(fake_crop_list).unsqueeze(1).detach()
# Disc Forward
y_d_rs, y_d_gs, _, _ = discriminator(real_crops, fake_crops)
d_loss, _, _ = discriminator_loss(y_d_rs, y_d_gs)
d_loss.backward()
norm_d = torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1.0)
torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1.0)
opt_d.step()
d_loss_item = d_loss.item()
@@ -392,36 +346,22 @@ if __name__ == '__main__':
# ---------------------
opt_g.zero_grad()
# Re-run generator (or reuse graph if not detached incorrectly)
decoder_in = decoder_in_padded.transpose(1, 2)
fake_audio = decoder(decoder_in)
if fake_audio.size(1) == 1: fake_audio = fake_audio.squeeze(1)
min_len = min(fake_audio.size(1), gt_audio.size(1))
fake_audio = fake_audio[:, :min_len]
real_audio = gt_audio[:, :min_len]
# Re-Crop for Generator (Same logic, new random crop or reuse?)
# Ideally reuse same crops if we didn't update D? But we did.
# Or standard GAN: Update D, then update G.
# Usually we crop again (stochasticity helps).
# We need "fake_crops_g" (with grad).
# We need "fake_crops_g" (with grad) for generator loss
real_crop_list_g = []
fake_crop_list_g = []
if args.use_disc:
if cfg_decoder["use_discriminator"]:
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)
if valid_len <= SEGMENT_SIZE_SAMPLES:
pad_len = SEGMENT_SIZE_SAMPLES - valid_len
if valid_len <= segment_size_samples:
pad_len = segment_size_samples - valid_len
r_c = torch.nn.functional.pad(real_audio[b_idx, :valid_len], (0, pad_len))
f_c = torch.nn.functional.pad(fake_audio[b_idx, :valid_len], (0, pad_len))
else:
start_idx = random.randint(0, valid_len - SEGMENT_SIZE_SAMPLES)
r_c = real_audio[b_idx, start_idx : start_idx + SEGMENT_SIZE_SAMPLES]
f_c = fake_audio[b_idx, start_idx : start_idx + SEGMENT_SIZE_SAMPLES]
start_idx = random.randint(0, valid_len - segment_size_samples)
r_c = real_audio[b_idx, start_idx : start_idx + segment_size_samples]
f_c = fake_audio[b_idx, start_idx : start_idx + segment_size_samples]
real_crop_list_g.append(r_c)
fake_crop_list_g.append(f_c)
@@ -429,54 +369,31 @@ if __name__ == '__main__':
real_crops_g = torch.stack(real_crop_list_g).unsqueeze(1)
fake_crops_g = torch.stack(fake_crop_list_g).unsqueeze(1)
# Losses
# Create Mel Mask
# We need mask corresponding to decoder_in (Audio Only)
# 1 Token = SAMPLES_PER_TOKEN audio samples
# Mel Hop Length = 512
# Ratio = SAMPLES_PER_TOKEN / 512 = 4 frames per token
# Mel Loss
frames_per_token = SAMPLES_PER_TOKEN // 512
# Expand audio_loss_mask: (B, T_aud) -> (B, T_aud*4)
mel_mask = audio_loss_mask.repeat_interleave(frames_per_token, dim=1)
pred_mel = mel_fn(fake_audio)
gt_mel = mel_fn(real_audio)
# Crop to min length
min_mel_len = min(pred_mel.size(2), gt_mel.size(2), mel_mask.size(1))
pred_mel = pred_mel[:, :, :min_mel_len]
gt_mel = gt_mel[:, :, :min_mel_len]
mel_mask = mel_mask[:, :min_mel_len]
loss_mel_raw = torch.nn.functional.l1_loss(pred_mel, gt_mel, reduction='none')
# Apply mask
loss_mel = (loss_mel_raw * mel_mask.unsqueeze(1)).sum() / (mel_mask.sum() * pred_mel.size(1) + 1e-6)
# Multi-Resolution STFT Loss
# Masking: Apply sample_mask to inputs
# Create sample_mask
sample_mask = audio_loss_mask.repeat_interleave(SAMPLES_PER_TOKEN, dim=1)
# Crop to min_len
sample_mask = sample_mask[:, :min_len]
# Apply mask (Assuming padding is zeros, but ensuring it)
# This prevents the model from being penalized for non-zero output in padding region (if target is zero)
# However, standard MR-STFT computes loss on the full spectrogram.
# If we zeroes out padding, STFT of zero is zero. Loss over zero-zero region is zero.
# But we divide by total elements (implicitly in mean).
# MR-STFT implementation uses averaging over batch/time.
# Ideally we should use a masked reduction.
# But for now, let's just zero out inputs.
sc_loss, mag_loss = mr_stft(fake_audio * sample_mask, real_audio * sample_mask)
loss_fm = torch.tensor(0.0, device=device)
loss_gen = torch.tensor(0.0, device=device)
if args.use_disc:
# Disc Forward Again (No Detach, using Crops)
if cfg_decoder["use_discriminator"]:
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)
@@ -487,9 +404,9 @@ if __name__ == '__main__':
norm_g = torch.nn.utils.clip_grad_norm_(decoder.parameters(), 1.0)
# LR Update
lr = get_lr(step)
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 args.use_disc:
if cfg_decoder["use_discriminator"]:
for param_group in opt_d.param_groups: param_group['lr'] = lr / 2
opt_g.step()
@@ -512,7 +429,9 @@ if __name__ == '__main__':
"train/loss_mag": mag_loss.item()
}
# Validation Loop & Logging
# ---------------------
# Validation Loop
# ---------------------
if step % val_freq == 0:
decoder.eval()
if discriminator: discriminator.eval()
@@ -523,7 +442,7 @@ if __name__ == '__main__':
val_d_loss_accum = 0.0
val_sc_loss_accum = 0.0
val_mag_loss_accum = 0.0
val_steps = 10 # Check 10 batches for speed
val_steps = 10
with torch.no_grad():
for _ in range(val_steps):
@@ -542,7 +461,6 @@ if __name__ == '__main__':
voutputs = model(vx, output_hidden_states=True)
v_hidden = voutputs.hidden_states[-1].to(torch.float32)
# GATHER AUDIO LATENTS Logic
v_gathered_states_list = []
for b_idx in range(v_hidden.size(0)):
mask = vaudio_mask[b_idx]
@@ -551,7 +469,6 @@ if __name__ == '__main__':
v_in_padded = torch.nn.utils.rnn.pad_sequence(v_gathered_states_list, batch_first=True, padding_value=0.0)
# Calculate Audio-specific Mask for Loss
v_bsz = v_in_padded.size(0)
v_max_aud_len = v_in_padded.size(1)
v_audio_loss_mask = torch.zeros((v_bsz, v_max_aud_len), dtype=torch.bool, device=device)
@@ -567,7 +484,6 @@ if __name__ == '__main__':
v_fake_audio = v_fake_audio[:, :min_len_v]
v_real_audio = vgt_audio[:, :min_len_v]
# Val Mel Loss
frames_per_token_v = SAMPLES_PER_TOKEN // 512
v_mel_mask = v_audio_loss_mask.repeat_interleave(frames_per_token_v, dim=1)
@@ -581,36 +497,30 @@ if __name__ == '__main__':
v_mel_loss_raw = torch.nn.functional.l1_loss(v_pred_mel, v_gt_mel, reduction='none')
v_mel_loss = (v_mel_loss_raw * v_mel_mask.unsqueeze(1)).sum() / (v_mel_mask.sum() * v_pred_mel.size(1) + 1e-6)
val_mel_loss_accum += v_mel_loss.item()
# Val MR-STFT
v_sample_mask = v_audio_loss_mask.repeat_interleave(SAMPLES_PER_TOKEN, dim=1)[:, :min_len_v]
v_sc_loss, v_mag_loss = mr_stft(v_fake_audio * v_sample_mask, v_real_audio * v_sample_mask)
val_sc_loss_accum += v_sc_loss.item()
val_mag_loss_accum += v_mag_loss.item()
if args.use_disc:
# Disc Forward with Cropping (Similar to Training)
if cfg_decoder["use_discriminator"]:
v_real_crop_list = []
v_fake_crop_list = []
v_min_len = min(v_fake_audio.size(1), v_real_audio.size(1))
for b_idx in range(v_bsz):
v_valid_len = v_gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN
v_valid_len = min(v_valid_len, v_min_len)
if v_valid_len <= SEGMENT_SIZE_SAMPLES:
v_pad_len = SEGMENT_SIZE_SAMPLES - v_valid_len
if v_valid_len <= segment_size_samples:
v_pad_len = segment_size_samples - v_valid_len
vr_c = torch.nn.functional.pad(v_real_audio[b_idx, :v_valid_len], (0, v_pad_len))
vf_c = torch.nn.functional.pad(v_fake_audio[b_idx, :v_valid_len], (0, v_pad_len))
else:
# Random Crop (or fixed center crop for determinism in val?)
# Random is fine as it averages out over batches/epochs.
v_start_idx = random.randint(0, v_valid_len - SEGMENT_SIZE_SAMPLES)
vr_c = v_real_audio[b_idx, v_start_idx : v_start_idx + SEGMENT_SIZE_SAMPLES]
vf_c = v_fake_audio[b_idx, v_start_idx : v_start_idx + SEGMENT_SIZE_SAMPLES]
v_start_idx = random.randint(0, v_valid_len - segment_size_samples)
vr_c = v_real_audio[b_idx, v_start_idx : v_start_idx + segment_size_samples]
vf_c = v_fake_audio[b_idx, v_start_idx : v_start_idx + segment_size_samples]
v_real_crop_list.append(vr_c)
v_fake_crop_list.append(vf_c)
@@ -639,38 +549,39 @@ if __name__ == '__main__':
log_dict.update(val_log)
# 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()
if cfg_global["use_wandb"]:
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()
# Create Plot
fig, ax = plt.subplots(2, 1, figsize=(10, 6))
ax[0].imshow(gt_mel, aspect='auto', origin='lower')
ax[0].set_title("Ground Truth Mel")
ax[1].imshow(gen_mel, aspect='auto', origin='lower')
ax[1].set_title("Generated Mel (Val)")
plt.tight_layout()
fig, ax = plt.subplots(2, 1, figsize=(10, 6))
ax[0].imshow(gt_mel, aspect='auto', origin='lower')
ax[0].set_title("Ground Truth Mel")
ax[1].imshow(gen_mel, aspect='auto', origin='lower')
ax[1].set_title("Generated Mel (Val)")
plt.tight_layout()
# Log to WandB
log_dict["val/mel_spectrograms"] = wandb.Image(fig)
plt.close(fig)
log_dict["val/mel_spectrograms"] = wandb.Image(fig)
plt.close(fig)
# Return to train mode
decoder.train()
if discriminator: discriminator.train()
# Save Checkpoint
if step > 0 and step % 3000 == 0:
print(f"Saving checkpoint at step {step}...")
if step > 0 and step % save_freq == 0:
print(f"\nSaving checkpoint at step {step} to {save_path}...")
ckpt_name_dec = f"decoder_step_{step}.pth"
ckpt_name_disc = f"discriminator_step_{step}.pth"
torch.save(decoder.state_dict(), save_path / ckpt_name_dec)
torch.save(decoder.state_dict(), os.path.join(save_path, ckpt_name_dec))
if discriminator:
torch.save(discriminator.state_dict(), save_path / ckpt_name_disc)
torch.save(discriminator.state_dict(), os.path.join(save_path, ckpt_name_disc))
wandb.log(log_dict, step=step)
if cfg_global["use_wandb"]:
wandb.log(log_dict, step=step)
print(f"Training complete. Saving model at {save_path}")
torch.save(decoder.state_dict(), save_path / "decoder_trained.pth")
torch.save(decoder.state_dict(), os.path.join(save_path, "decoder_trained.pth"))
if discriminator:
torch.save(discriminator.state_dict(), save_path / "discriminator_trained.pth")
wandb.finish()
torch.save(discriminator.state_dict(), os.path.join(save_path, "discriminator_trained.pth"))
if cfg_global["use_wandb"]:
wandb.finish()
train_llm.py
+133 -175
View File
@@ -1,18 +1,12 @@
"""
Training script for Soprano.
Training script for Soprano LLM backbone.
Usage:
python train.py --input-dir path/to/files --save-dir path/to/weights
Args:
--input-dir: Path to directory of LJSpeech-style dataset. If none is provided this defaults to the provided example dataset.
--save-dir: Path to directory to save weights
python train_llm.py
Adapted from https://github.com/karpathy/nanoGPT
"""
import os
import argparse
import pathlib
import random
import time
import wandb
@@ -25,88 +19,25 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from safetensors.torch import load_file
from dataset import AudioDataset
from config_loader import load_config
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--input-dir",
required=False,
default="./example_dataset",
type=pathlib.Path
)
parser.add_argument("--save-dir",
required=True,
type=pathlib.Path
)
parser.add_argument("--from-scratch",
action="store_true",
help="If set, train from scratch using the config, instead of loading pretrained weights."
)
return parser.parse_args()
args = get_args()
# training hyperparameters
device = 'cuda:0'
seed = 1337
max_lr = 2e-5
warmup_ratio = 0.3
cooldown_ratio = 0.1
min_lr = 0.3 * max_lr
batch_size = 64
grad_accum_steps = 1
seq_len = 1024
val_freq = 250
save_freq = 5000
text_factor = 0.5 # currently does not train on text inputs, you can increase to change this
max_steps = 150000
betas = (0.9, 0.95)
weight_decay = 0.1
train_dataset_path = f'{args.input_dir}/train.json'
val_dataset_path = f'{args.input_dir}/val.json'
save_path = args.save_dir
os.makedirs(save_path, exist_ok=True)
# Initialize tokenizer globally so it can be used in the collate functions
tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
tokenizer.padding_side = 'right' # Essential for training!
def worker_seed_init(_):
worker_seed = torch.initial_seed() % (2**32-1)
np.random.seed(worker_seed)
random.seed(worker_seed)
def get_lr(it): # WSD schedule
if it<warmup_steps:
return max_lr * (it+1) / warmup_steps
if it<max_steps-cooldown_steps:
def get_lr(it, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps): # WSD schedule
if it < warmup_steps:
return max_lr * (it + 1) / warmup_steps
if it < max_steps - cooldown_steps:
return max_lr
return min_lr + (max_lr-min_lr) * ((max_steps-it) / cooldown_steps)
def collate_pack(texts):
tokens_batch = tokenizer(texts, padding=False, truncation=False)
batch = []
cur_sample, cur_size = [], 0
for i in range(len(texts)):
tokens = torch.tensor(tokens_batch['input_ids'][i][:-1], dtype=torch.long)
cur_size += tokens.size(0)
cur_sample.append(tokens)
if cur_size >= seq_len + 1:
batch.append(torch.cat(cur_sample)[: seq_len + 1])
cur_sample, cur_size = [], 0
if len(batch) == batch_size:
break
if cur_sample and not batch: # add partial sample if there isn't enough data
batch_item = torch.cat(cur_sample + [torch.zeros(seq_len, dtype=torch.long)])[: seq_len + 1]
batch.append(batch_item)
if len(batch) < batch_size:
# pad up to batch_size for consistency
pad = batch[-1]
while len(batch) < batch_size:
batch.append(pad)
batch = torch.stack(batch)
x = batch[:, :-1]
y = batch[:, 1:]
return x, y
return min_lr + (max_lr - min_lr) * ((max_steps - it) / cooldown_steps)
def collate_dynamic(texts):
# Dynamic Batching: Pad to the longest in this batch (max 2048 safety)
tokenized = tokenizer(texts, padding=True, truncation=True, max_length=2048, return_tensors='pt', add_special_tokens=False)
batch = tokenized['input_ids']
@@ -120,38 +51,7 @@ def collate_dynamic(texts):
return x, y, attn_mask
def collate_pack_val(texts):
# max_length=seq_len+1 because we need to shift for x, y
out = tokenizer(texts, padding=True, truncation=True, max_length=seq_len+1, return_tensors='pt', asdd_special_tokens=False)
batch = out['input_ids']
# Ensure fixed length padding to seq_len + 1
if batch.size(1) < seq_len + 1:
pad_len = seq_len + 1 - batch.size(1)
batch = torch.nn.functional.pad(batch, (0, pad_len), value=tokenizer.pad_token_id)
x = batch[:, :-1]
y = batch[:, 1:]
return x, y
# def get_audio_filter(x):
# audio_token_start = 4
# audio_token_end = 8003
# audio_start_token = 8195
# audio_end_token = 8196
# audio_tokens_list = list(range(audio_token_start, audio_token_end+1))
# audio_tokens_list.append(audio_start_token)
# audio_tokens_list.append(audio_end_token)
# return torch.isin(x, audio_tokens_list)
def compute_loss(x, logits, y, num_steps, mask=None):
pred = logits.view(-1, logits.size(-1))
labels = y.reshape(-1)
loss = torch.nn.functional.cross_entropy(pred, labels, reduction='none')
@@ -182,7 +82,6 @@ def compute_loss(x, logits, y, num_steps, mask=None):
# Acc: only on non-masked tokens.
# Current logic: (logits.argmax(dim=-1) == y).view(-1)[audio_mask]
# This correctly calculates accuracy only on valid audio tokens.
acc = (logits.argmax(dim=-1).view(-1) == labels).view(-1)[audio_mask].to(torch.float32).mean()
if torch.isnan(acc): acc = torch.tensor(0.0, device=loss.device)
@@ -191,7 +90,7 @@ def compute_loss(x, logits, y, num_steps, mask=None):
acc = acc / num_steps
return audio_loss, text_loss, acc
def evaluate(val_dataloader, step):
def evaluate(model, val_dataloader, step, device, use_wandb):
model.eval()
val_dataloader_it = iter(val_dataloader)
with torch.no_grad():
@@ -199,147 +98,206 @@ def evaluate(val_dataloader, step):
val_text_loss_accum = torch.tensor(0.0).to(device)
val_acc_accum = torch.tensor(0.0).to(device)
val_loss_steps = len(val_dataloader)
for _ in range(val_loss_steps):
x, y, attn_mask = next(val_dataloader_it)
x, y, attn_mask = x.to(device), y.to(device), attn_mask.to(device)
# with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
logits = model(x, attention_mask=attn_mask).logits
audio_loss, text_loss, acc = compute_loss(x, logits, y, val_loss_steps, mask=attn_mask)
val_audio_loss_accum += audio_loss.detach()
val_text_loss_accum += text_loss.detach()
val_acc_accum += acc.detach()
print(f"validation text loss: {val_text_loss_accum.item():.4f}\tvalidation audio loss: {val_audio_loss_accum.item():.4f}\tvalidation acc: {val_acc_accum.item():.4f}")
wandb.log({
"val/text_loss": val_text_loss_accum.item(),
"val/audio_loss": val_audio_loss_accum.item(),
"val/acc": val_acc_accum.item()
}, step=step)
if use_wandb:
wandb.log({
"val/text_loss": val_text_loss_accum.item(),
"val/audio_loss": val_audio_loss_accum.item(),
"val/acc": val_acc_accum.item()
}, step=step)
model.train()
tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
tokenizer.padding_side = 'right' # Essential for training!
if __name__ == '__main__':
# ------------------
# Load Configuration
# ------------------
config = load_config("config.yaml")
cfg_global = config["global"]
cfg_paths = config["paths"]
cfg_llm = config["llm"]
device = cfg_global["device"]
seed = cfg_global["seed"]
device_type = "cuda" if device.startswith("cuda") else "cpu"
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.set_float32_matmul_precision('high')
# Setup directories
train_dataset_path = os.path.join(cfg_paths["dataset_root"], "train.json")
val_dataset_path = os.path.join(cfg_paths["dataset_root"], "val.json")
save_path = os.path.join(cfg_paths["save_dir"], "llm")
os.makedirs(save_path, exist_ok=True)
print(f"Save Path: {save_path}")
wandb.init(project="soprano-llm", config=vars(args))
if cfg_global["use_wandb"]:
wandb.init(project=cfg_global["wandb_project"], config=config)
# lr schedule
warmup_steps = int(max_steps * warmup_ratio)
cooldown_steps = int(max_steps * cooldown_ratio)
# ------------------
# Hyperparameters
# ------------------
max_steps = cfg_llm["max_steps"]
max_lr = float(cfg_llm["max_lr"])
min_lr = cfg_llm["min_lr_ratio"] * max_lr
warmup_steps = int(max_steps * cfg_llm["warmup_ratio"])
cooldown_steps = int(max_steps * cfg_llm["cooldown_ratio"])
# model
# model
if args.from_scratch:
batch_size = cfg_llm["batch_size"]
grad_accum_steps = cfg_llm["grad_accum_steps"]
seq_len = cfg_llm["seq_len"]
val_freq = cfg_llm["val_freq"]
save_freq = cfg_llm["save_freq"]
text_factor = cfg_llm["text_factor"]
betas = tuple(cfg_llm["betas"])
weight_decay = cfg_llm["weight_decay"]
# ------------------
# Model Setup
# ------------------
if cfg_llm["from_scratch"]:
print("Initializing model from scratch (random weights)...")
config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
model = AutoModelForCausalLM.from_config(config)
m_config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
model = AutoModelForCausalLM.from_config(m_config)
else:
print("Loading pretrained model weights...")
model = AutoModelForCausalLM.from_pretrained('ekwek/Soprano-80M')
# ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_llm/codec_v2/v2/model.safetensors"
ckpt_path = "/home/ubuntu/soma/ckpt/suprano/suprano_llm/codec_v2/v2/checkpoint-40000/model.safetensors"
state_dict = load_file(ckpt_path)
model.load_state_dict(state_dict)
pretrained_path = cfg_paths["pretrained_llm_path"]
if pretrained_path and os.path.exists(pretrained_path):
print(f"Loading pretrained model weights from {pretrained_path}...")
m_config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
model = AutoModelForCausalLM.from_config(m_config)
state_dict = load_file(pretrained_path)
model.load_state_dict(state_dict)
else:
print("Loading default pretrained model weights from HF...")
model = AutoModelForCausalLM.from_pretrained('ekwek/Soprano-80M')
model.to(device)
model.train()
# dataset
# ------------------
# Dataset Setup
# ------------------
dataset = AudioDataset(train_dataset_path)
# Using dynamic batching now
dataloader = DataLoader(dataset,
dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4,
num_workers=cfg_global["num_workers"],
pin_memory=True,
persistent_workers=True,
worker_init_fn=worker_seed_init,
collate_fn=collate_dynamic,
)
dataloader_it = iter(dataloader)
val_dataset = AudioDataset(val_dataset_path)
val_dataloader = DataLoader(val_dataset,
val_dataloader = DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=2,
num_workers=max(1, cfg_global["num_workers"] // 2),
pin_memory=True,
persistent_workers=True,
worker_init_fn=worker_seed_init,
collate_fn=collate_dynamic,
)
import pdb;pdb.set_trace()
# optimizer
# ------------------
# Optimizer
# ------------------
opt = torch.optim.AdamW(model.parameters(), max_lr, betas=betas, weight_decay=weight_decay, fused=True)
pbar = tqdm(range(40001, max_steps), ncols=200, dynamic_ncols=True)
# ------------------
# Training Loop
# ------------------
# Determine start step based on loaded checkpoint if needed, defaulting to 1 for new runs
start_step = 1
pbar = tqdm(range(start_step, max_steps + 1), ncols=200, dynamic_ncols=True)
for step in pbar:
start = time.time()
if val_freq>0 and step != 0 and (step % val_freq == 0 or step==max_steps-1):
evaluate(val_dataloader, step)
if val_freq > 0 and step != start_step and (step % val_freq == 0 or step == max_steps):
evaluate(model, val_dataloader, step, device, cfg_global["use_wandb"])
if save_freq > 0 and step % save_freq == 0:
ckpt_path = os.path.join(save_path, f"checkpoint-{step}")
print(f"Saving checkpoint to {ckpt_path}")
model.save_pretrained(ckpt_path)
tokenizer.save_pretrained(ckpt_path)
ckpt_dir = os.path.join(save_path, f"checkpoint-{step}")
print(f"\nSaving checkpoint to {ckpt_dir}")
model.save_pretrained(ckpt_dir)
tokenizer.save_pretrained(ckpt_dir)
opt.zero_grad()
audio_loss_accum = 0.0
text_loss_accum = 0.0
acc_accum = 0.0
for micro_step in range(grad_accum_steps):
try:
x, y, attn_mask = next(dataloader_it)
except:
except StopIteration:
dataloader_it = iter(dataloader)
x, y, attn_mask = next(dataloader_it)
x, y, attn_mask = x.to(device), y.to(device), attn_mask.to(device)
# with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
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()
total_loss = audio_loss + text_factor*text_loss
total_loss = audio_loss + text_factor * text_loss
total_loss.backward()
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
lr = get_lr(step)
lr = get_lr(step, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps)
for param_group in opt.param_groups:
param_group['lr'] = lr
opt.step()
torch.cuda.synchronize()
total_tokens = step * batch_size*seq_len*grad_accum_steps
if device_type == "cuda":
torch.cuda.synchronize()
end = time.time()
dt = (end-start)*1000
tokens_per_second = (batch_size*seq_len*grad_accum_steps) / (end-start)
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'
pbar.set_description(tqdm_log)
wandb.log({
"train/text_loss": text_loss_accum.item(),
"train/audio_loss": audio_loss_accum.item(),
"train/acc": acc_accum.item(),
"train/lr": lr,
"train/grad_norm": norm,
"train/dt": dt,
"train/tokens_per_sec": tokens_per_second
}, step=step)
print(f"Training complete. Saving model at {save_path}")
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/lr": lr,
"train/grad_norm": norm,
"train/dt": dt,
"train/tokens_per_sec": tokens_per_second
}, step=step)
print(f"\nTraining complete. Saving final model at {save_path}")
model.save_pretrained(save_path)
tokenizer.save_pretrained(save_path)
print("Saving done.")
wandb.finish()
if cfg_global["use_wandb"]:
wandb.finish()
utils/audio_utils.py
+85
View File
@@ -0,0 +1,85 @@
import platform
import os
import subprocess
import numpy as np
import torch
import soundfile as sf
import torchaudio
class AudioPipeline:
@staticmethod
def load_audio(file_path, target_sr=32000):
"""OS-aware audio loading and resampling."""
system = platform.system()
if system == "Windows":
return AudioPipeline._load_windows(file_path, target_sr)
else:
return AudioPipeline._load_linux(file_path, target_sr)
@staticmethod
def _load_linux(file_path, target_sr):
# Linux handles torchaudio beautifully
audio, sr = torchaudio.load(file_path)
if sr != target_sr:
audio = torchaudio.functional.resample(audio, orig_freq=sr, new_freq=target_sr)
# Ensure Mono
if audio.shape[0] > 1:
audio = audio.mean(dim=0, keepdim=True)
return audio, target_sr
@staticmethod
def _load_windows(file_path, target_sr):
# Check for local ffmpeg first, then fallback to system PATH
ffmpeg_cmd = "ffmpeg"
if os.path.exists("./tools/ffmpeg/ffmpeg.exe"):
ffmpeg_cmd = "./tools/ffmpeg/ffmpeg.exe"
# Force sample rate, mono channel, and output to raw PCM float32
command = [
ffmpeg_cmd,
"-i", str(file_path),
"-ac", "1", # Force Mono
"-ar", str(target_sr), # Target sample rate
"-f", "f32le", # Format: float32 little endian
"-hide_banner",
"-loglevel", "error",
"-" # Output to stdout
]
try:
process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
out, err = process.communicate()
if process.returncode != 0:
print(f"ffmpeg error for {file_path}, falling back to soundfile. Error: {err.decode('utf-8')}")
return AudioPipeline._load_windows_fallback(file_path, target_sr)
# Read raw bytes directly into a numpy array, then to torch tensor
audio_np = np.frombuffer(out, dtype=np.float32).copy()
audio_tensor = torch.from_numpy(audio_np).unsqueeze(0) # Shape: (1, T)
return audio_tensor, target_sr
except FileNotFoundError:
print(f"ffmpeg not found in PATH or ./tools/ffmpeg/. Falling back to soundfile for {file_path}")
return AudioPipeline._load_windows_fallback(file_path, target_sr)
@staticmethod
def _load_windows_fallback(file_path, target_sr):
# Soundfile doesn't natively resample, so we rely on scipy
import scipy.signal
audio_np, sr = sf.read(file_path)
# Convert to mono if needed
if len(audio_np.shape) > 1:
audio_np = audio_np.mean(axis=1)
if sr != target_sr:
num_samples = int(round(len(audio_np) * float(target_sr) / sr))
audio_np = scipy.signal.resample(audio_np, num_samples)
# Shape to (1, T)
audio_tensor = torch.from_numpy(audio_np).float().unsqueeze(0)
return audio_tensor, target_sr
utils/config_loader.py
+21
View File
@@ -0,0 +1,21 @@
import yaml
import os
from pathlib import Path
def load_config(config_path="config.yaml"):
"""Loads the YAML config and resolves paths to absolute paths."""
if not os.path.exists(config_path):
raise FileNotFoundError(f"Configuration file not found at {config_path}")
with open(config_path, "r", encoding="utf-8") as f:
config = yaml.safe_load(f)
# Resolve paths relative to the current working directory
if "paths" in config:
for key, val in config["paths"].items():
if val is not None and isinstance(val, str):
# Expand ~ for Linux/Mac and resolve to absolute paths
resolved_path = Path(val).expanduser().resolve()
config["paths"][key] = str(resolved_path)
return config