mirror of
https://github.com/Nighthawk42/soprano-factory.git
synced 2026-08-30 04:30:21 +00:00
v2 Refactor
This commit is contained in:
Nighthawk
@@ -1,129 +1,98 @@
|
||||
# Soprano-Reforged
|
||||
# Soprano-Factory (Clean-Room Edition)
|
||||
|
||||
Soprano-Reforged is a clean, modular, and robust implementation of the Soprano text-to-speech architecture. It is designed for training custom high-fidelity speech models on Windows and Linux with minimal friction.
|
||||
A fully rebuilt, production-ready training pipeline for fine-tuning the [Soprano](https://github.com/ekwek1/soprano) Text-to-Speech model.
|
||||
|
||||
## Overview
|
||||
This repository solves the notorious "noise soup" issue present in earlier fine-tuning scripts by implementing **Hidden State Distillation**, strict `int16` audio normalization, and proper `[STOP]` token preservation.
|
||||
|
||||
This repository implements a 3-stage training pipeline:
|
||||
|
||||
1. **Stage 0 (Codec)**: Trains a custom neural audio codec (Encoder and Decoder) specific to your speaker's voice.
|
||||
2. **Stage 1 (Tokenization)**: Converts your audio dataset into discrete acoustic tokens.
|
||||
3. **Stage 2 (Joint Training)**: Fine-tunes the Soprano 1.1 LLM to generate these tokens from text.
|
||||
|
||||
## Installation
|
||||
## Why the Rebuild?
|
||||
Standard causal language model fine-tuning uses discrete Cross-Entropy loss. However, the Soprano Vocos decoder generates audio from the LLM's *continuous hidden states*. Standard fine-tuning causes these hidden states to drift, resulting in the frozen decoder outputting pure static.
|
||||
|
||||
This project is built with `uv` for fast, reliable dependency management.
|
||||
This clean-room pipeline fixes this by:
|
||||
1. **Hidden State Distillation**: A frozen "Teacher" model anchors the latent space. A trainable "Student" model learns the new tokens while an MSE loss penalty forces its hidden states to remain compatible with the frozen audio decoder.
|
||||
2. **Audio Normalization**: Safely scales 16-bit PCM WAV files to `[-1.0, 1.0]` float32 tensors, preventing the Mel Spectrogram energy from exploding and corrupting the discrete audio tokens.
|
||||
3. **Collation Fixes**: Safely packs sequences without accidentally slicing off the `[STOP]` token, allowing the model to learn when to terminate speech.
|
||||
|
||||
### 1. Clone the repository
|
||||
---
|
||||
|
||||
```bash
|
||||
git clone https://github.com/yourusername/soprano-reforged
|
||||
cd soprano-reforged
|
||||
```
|
||||
## 📦 Installation
|
||||
|
||||
### 2. Initialize environment and install dependencies
|
||||
This project strictly uses `uv` for dependency management to ensure reliable, reproducible builds with the correct CUDA 12.8 bindings.
|
||||
|
||||
```bash
|
||||
uv sync
|
||||
```
|
||||
1. **Install `uv`** (if you haven't already):
|
||||
```bash
|
||||
curl -LsSf [https://astral.sh/uv/install.sh](https://astral.sh/uv/install.sh) | sh
|
||||
|
||||
### Dependencies
|
||||
Clone and Sync:
|
||||
Bash
|
||||
|
||||
- **Core**: torch, torchaudio, transformers
|
||||
- **Audio**: soundfile, librosa
|
||||
- **GUI**: sv-ttk, darkdetect
|
||||
- **Utils**: pyyaml, tqdm
|
||||
git clone [https://github.com/yourusername/soprano-factory.git](https://github.com/yourusername/soprano-factory.git)
|
||||
cd soprano-factory
|
||||
uv sync
|
||||
|
||||
> **Note for Windows Users**: The pipeline includes automatic detection for local FFmpeg binaries in `tools/ffmpeg` to prevent path issues.
|
||||
Note: The included pyproject.toml automatically configures the PyTorch explicit index for CUDA 12.8.
|
||||
|
||||
## Quick Start (GUI)
|
||||
📂 Project Structure
|
||||
Plaintext
|
||||
|
||||
The easiest way to use Soprano-Reforged is the built-in Training Factory GUI.
|
||||
soprano-factory/
|
||||
├── pyproject.toml # uv dependencies and CUDA index
|
||||
├── config.yaml # Centralized hyperparameters and paths
|
||||
├── generate_dataset.py # Preprocesses WAVs into quantized tokens
|
||||
├── train.py # Main training loop with Distillation & TensorBoard
|
||||
├── test_inference.py # Generates test audio from the fine-tuned model
|
||||
├── dataset.py # PyTorch Dataset and PackedCollation logic
|
||||
├── text_normalizer.py # Text cleaning utility
|
||||
└── example_dataset/ # Your raw data goes here
|
||||
├── metadata.txt
|
||||
└── wavs/
|
||||
|
||||
```bash
|
||||
python gui.py
|
||||
```
|
||||
🚀 Usage Pipeline
|
||||
1. Prepare Your Data
|
||||
|
||||
This launches a modern, dark-mode interface where you can:
|
||||
Place your audio files in example_dataset/wavs/ and create an LJSpeech-formatted metadata.txt inside example_dataset/.
|
||||
|
||||
- Configure paths for your dataset
|
||||
- Monitor training logs in real-time
|
||||
- Execute all 3 stages of the pipeline without touching the command line
|
||||
Format: filename|Transcript text here.
|
||||
Plaintext
|
||||
|
||||
## Configuration
|
||||
example1|Soprano is an extremely lightweight text to speech model.
|
||||
example2|Gabagool? Ova here!
|
||||
|
||||
All hyperparameters are centralized in `config/settings.yaml`. You should edit this file to match your hardware capabilities (VRAM) and dataset size.
|
||||
2. Configure Training
|
||||
|
||||
### config/settings.yaml excerpt
|
||||
Edit config.yaml to set your desired batch size, accumulation steps, learning rate, and distillation weight.
|
||||
3. Generate the Dataset
|
||||
|
||||
```yaml
|
||||
common:
|
||||
device: "auto"
|
||||
Run the data generator to normalize the audio, clean the text, and encode the audio into discrete FSQ tokens.
|
||||
Bash
|
||||
|
||||
codec:
|
||||
batch_size: 8
|
||||
epochs: 100
|
||||
uv run generate_dataset.py --input-dir ./example_dataset
|
||||
|
||||
training:
|
||||
base_model: "ekwek/Soprano-1.1-80M"
|
||||
batch_size: 4
|
||||
epochs: 10
|
||||
```
|
||||
This will create train.json and val.json in your dataset directory.
|
||||
4. Train the Model
|
||||
|
||||
## The 3-Stage Pipeline (CLI Usage)
|
||||
Start the training loop. This script utilizes Hidden State Distillation and logs metrics, audio samples, and hidden-state visualizations directly to TensorBoard.
|
||||
Bash
|
||||
|
||||
If you prefer the command line, follow these steps.
|
||||
uv run train.py --config config.yaml
|
||||
|
||||
### Stage 0: Train the Codec
|
||||
Monitor Training Live:
|
||||
Open a new terminal and run:
|
||||
Bash
|
||||
|
||||
Learns to compress and reconstruct your specific audio data.
|
||||
uv run tensorboard --logdir=./runs/soprano_finetune
|
||||
|
||||
```bash
|
||||
python train_codec.py --wav-dir "./my_dataset/wavs/*.wav"
|
||||
```
|
||||
5. Test Inference
|
||||
|
||||
### Stage 1: Generate Dataset
|
||||
Once training completes, use the test script to verify that your model generates clear audio and stops correctly. This script automatically bundles the required decoder.pth into your fine-tuned directory.
|
||||
Bash
|
||||
|
||||
Uses the trained encoder to convert audio files into token sequences (`train.json`).
|
||||
uv run test_inference.py --model-dir ./finetuned_model --text "This is a test of my custom Soprano voice." --output my_test.wav
|
||||
|
||||
```bash
|
||||
python generate_dataset.py --input-dir ./my_dataset --encoder-ckpt ./weights/codec/encoder.pth
|
||||
```
|
||||
📄 License
|
||||
|
||||
### Stage 2: Joint Training
|
||||
This project is licensed under the Apache-2.0 license.
|
||||
|
||||
Trains the LLM and Decoder together. The LLM learns to predict tokens, and the Decoder learns to reconstruct audio from the LLM's hidden states, fixing the drift problem.
|
||||
|
||||
```bash
|
||||
python train.py --input-dir ./my_dataset
|
||||
```
|
||||
|
||||
## Inference
|
||||
|
||||
Once Stage 2 is complete, you can generate speech using your new model.
|
||||
|
||||
```bash
|
||||
python inference.py --text "Soprano Reforged is fully operational." --model-dir ./weights/model/final --output result.wav
|
||||
```
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
SopranoReforged/
|
||||
├── config/ # Configuration files (settings.yaml)
|
||||
├── model/ # Neural network architecture (Encoder, Decoder, Quantizer)
|
||||
├── training/ # Training utilities (Loss functions, Collators)
|
||||
├── utils/ # Helper scripts (Config loader, Text normalizer)
|
||||
├── gui.py # Main GUI application
|
||||
├── train_codec.py # Stage 0 script
|
||||
├── generate_dataset.py # Stage 1 script
|
||||
├── train.py # Stage 2 script
|
||||
└── inference.py # Generation script
|
||||
```
|
||||
|
||||
## Credits
|
||||
|
||||
- **Original Architecture**: Based on the Soprano model by [ekwek](https://huggingface.co/ekwek)
|
||||
- **Vocos Backbone**: Utilizes ConvNeXt blocks for high-fidelity audio reconstruction
|
||||
- **Reforged By**: Nighthawk42
|
||||
|
||||
The pipeline is officially clean, robust, and ready to go. Let me know if you'd like to dive into pushing this to a
|
||||
@@ -1,54 +0,0 @@
|
||||
# analyze_tokens.py - Diagnostic for LLM Audio Tokens
|
||||
import torch
|
||||
import argparse
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from utils.text_normalizer import normalize_text
|
||||
from utils.config import cfg
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--text", type=str, required=True)
|
||||
parser.add_argument("--model-dir", type=str, required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
||||
llm = AutoModelForCausalLM.from_pretrained(args.model_dir, torch_dtype=torch.float32).to(cfg.device)
|
||||
|
||||
norm_text = normalize_text(args.text)
|
||||
prompt = f"[STOP][TEXT]{norm_text}[START]"
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(cfg.device)
|
||||
|
||||
print(f"Analyzing LLM response for: {norm_text}")
|
||||
|
||||
with torch.no_grad():
|
||||
output_ids = llm.generate(
|
||||
inputs["input_ids"],
|
||||
max_new_tokens=128,
|
||||
do_sample=True,
|
||||
top_k=50
|
||||
)
|
||||
|
||||
# Convert to list and find the start of audio
|
||||
tokens = output_ids[0].tolist()
|
||||
start_token_id = tokenizer.convert_tokens_to_ids("[START]")
|
||||
|
||||
try:
|
||||
idx = tokens.index(start_token_id)
|
||||
audio_tokens = tokens[idx+1:]
|
||||
except:
|
||||
audio_tokens = tokens
|
||||
|
||||
print("-" * 30)
|
||||
print(f"Total Audio Tokens Generated: {len(audio_tokens)}")
|
||||
print(f"Unique Tokens: {len(set(audio_tokens))}")
|
||||
print(f"First 20 Tokens: {audio_tokens[:20]}")
|
||||
|
||||
if len(set(audio_tokens)) <= 3 and len(audio_tokens) > 10:
|
||||
print("CRITICAL: Token Collapse detected. The LLM is repeating itself.")
|
||||
print("Remedy: Increase training epochs or adjust temperature/top_k.")
|
||||
else:
|
||||
print("Status: LLM is producing a varied token stream. The issue is likely the Decoder's projection.")
|
||||
print("-" * 30)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1 @@
|
||||
# __init__.py - Initializes the Soprano Factory codec package, which includes the Encoder and Decoder classes for the TTS pipeline.
|
||||
@@ -0,0 +1,50 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from encoder.codec import VocosBackbone
|
||||
|
||||
|
||||
class SimpleDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_mels=50,
|
||||
encoder_dim=768,
|
||||
bottleneck_channels=5,
|
||||
num_layers=8,
|
||||
intermediate_dim=None,
|
||||
upsample_scale=4,
|
||||
dw_kernel=5,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
intermediate_dim = intermediate_dim or encoder_dim * 3
|
||||
self.upsample_scale = upsample_scale
|
||||
|
||||
# project FSQ channels back to model dim
|
||||
self.in_proj = nn.Linear(bottleneck_channels, encoder_dim)
|
||||
|
||||
# ConvNeXt backbone
|
||||
self.backbone = VocosBackbone(
|
||||
input_channels=encoder_dim,
|
||||
dim=encoder_dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
num_layers=num_layers,
|
||||
input_kernel_size=1,
|
||||
dw_kernel_size=dw_kernel,
|
||||
)
|
||||
|
||||
# output mel projection
|
||||
self.out_proj = nn.Conv1d(encoder_dim, n_mels, kernel_size=1)
|
||||
|
||||
def forward(self, z):
|
||||
"""
|
||||
z: (B, T_latent, bottleneck_channels)
|
||||
"""
|
||||
z = self.in_proj(z) # (B, T_latent, D)
|
||||
z = z.transpose(1, 2) # (B, D, T_latent)
|
||||
|
||||
# naive upsampling (good enough for now)
|
||||
z = z.repeat_interleave(self.upsample_scale, dim=2)
|
||||
|
||||
z = self.backbone(z) # (B, D, T_mel)
|
||||
mel_hat = self.out_proj(z) # (B, n_mels, T_mel)
|
||||
return mel_hat
|
||||
@@ -0,0 +1,213 @@
|
||||
"""
|
||||
Adapted from https://github.com/gemelo-ai/vocos
|
||||
"""
|
||||
from typing import Optional
|
||||
|
||||
import torchaudio
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .quantizer import FSQSTE
|
||||
|
||||
def safe_log(x: torch.Tensor, clip_val: float = 5e-3) -> torch.Tensor:
|
||||
return torch.log(torch.clip(x, min=clip_val))
|
||||
|
||||
|
||||
class SimpleMLP(nn.Module):
|
||||
def __init__(self,
|
||||
dim,
|
||||
intermediate_dim,
|
||||
):
|
||||
super().__init__()
|
||||
self.pwconv1 = nn.Linear(dim, intermediate_dim)
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(intermediate_dim, dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
"""ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
intermediate_dim (int): Dimensionality of the intermediate layer.
|
||||
layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
|
||||
Defaults to None.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
layer_scale_init_value: float,
|
||||
dw_kernel_size: int = 7,
|
||||
):
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv1d(dim, dim, kernel_size=dw_kernel_size, padding=dw_kernel_size//2, groups=dim) # depthwise conv
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.mlp = SimpleMLP(dim, intermediate_dim)
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.dwconv(x)
|
||||
x = x.transpose(1, 2) # (B, C, T) -> (B, T, C)
|
||||
x = self.norm(x)
|
||||
x = self.mlp(x)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
x = x.transpose(1, 2) # (B, T, C) -> (B, C, T)
|
||||
|
||||
x = residual + x
|
||||
return x
|
||||
|
||||
|
||||
class VocosBackbone(nn.Module):
|
||||
"""
|
||||
Vocos backbone module built with ConvNeXt blocks.
|
||||
|
||||
Args:
|
||||
input_channels (int): Number of input features channels.
|
||||
dim (int): Hidden dimension of the model.
|
||||
intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
|
||||
num_layers (int): Number of ConvNeXtBlock layers.
|
||||
layer_scale_init_value (float, optional): Initial value for layer scaling.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
num_layers: int,
|
||||
input_kernel_size: int = 7,
|
||||
dw_kernel_size: int = 7,
|
||||
layer_scale_init_value: Optional[float] = None,
|
||||
pad: str = 'zeros',
|
||||
):
|
||||
super().__init__()
|
||||
self.input_channels = input_channels
|
||||
self.dim = dim
|
||||
self.embed = nn.Conv1d(
|
||||
input_channels,
|
||||
dim,
|
||||
kernel_size=input_kernel_size,
|
||||
padding=input_kernel_size//2,
|
||||
padding_mode=pad
|
||||
)
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.convnext = nn.ModuleList([
|
||||
ConvNeXtBlock(
|
||||
dim=dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
dw_kernel_size=dw_kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value or 1 / num_layers**0.5,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.apply(self._init_weights)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
|
||||
C denotes output features, and L is the sequence length.
|
||||
|
||||
Returns:
|
||||
Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
|
||||
and H denotes the model dimension.
|
||||
"""
|
||||
x = self.embed(x) # (B, C, L)
|
||||
x = self.norm(x.transpose(1, 2))
|
||||
x = x.transpose(1, 2)
|
||||
for conv_block in self.convnext:
|
||||
x = conv_block(x)
|
||||
x = self.final_layer_norm(x.transpose(1, 2))
|
||||
x = x.transpose(1, 2)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self,
|
||||
num_input_mels=50,
|
||||
mel_hop_length=512,
|
||||
mel_hop_scale=0.25,
|
||||
encoder_num_layers=8,
|
||||
encoder_dim=768,
|
||||
encoder_intermediate_dim=None,
|
||||
fsq_levels=[8, 8, 5, 5, 5],
|
||||
dw_kernel=5,
|
||||
):
|
||||
super().__init__()
|
||||
self.downsample_scale = 2048 // mel_hop_length
|
||||
self.mel_hop_length = mel_hop_length
|
||||
self.mel_n_fft = int(mel_hop_length/mel_hop_scale)
|
||||
self.encoder_dim = encoder_dim
|
||||
self.encoder_intermediate_dim = encoder_intermediate_dim if encoder_intermediate_dim else encoder_dim*3
|
||||
self.encoder_num_layers = encoder_num_layers
|
||||
self.encoder_initial_channels = num_input_mels
|
||||
self.bottleneck_channels = 5
|
||||
|
||||
self.mel_spec = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=32000,
|
||||
n_fft=self.mel_n_fft,
|
||||
hop_length=self.mel_hop_length,
|
||||
n_mels=num_input_mels,
|
||||
center=True,
|
||||
power=1,
|
||||
)
|
||||
self.encoder = VocosBackbone(input_channels=self.encoder_initial_channels,
|
||||
dim=self.encoder_dim,
|
||||
intermediate_dim=self.encoder_intermediate_dim,
|
||||
num_layers=self.encoder_num_layers,
|
||||
input_kernel_size=1,
|
||||
dw_kernel_size=dw_kernel,
|
||||
pad='zeros'
|
||||
)
|
||||
self.downsampler = nn.Linear(self.encoder_dim, self.bottleneck_channels)
|
||||
self.quant = FSQSTE(levels=fsq_levels)
|
||||
|
||||
def encode(self, x):
|
||||
x = self.encoder(x)
|
||||
|
||||
# import pdb;pdb.set_trace()
|
||||
|
||||
x = x[:, :, ::self.downsample_scale] # What the heck is this? Brute force downsampling from mel -> tokens.
|
||||
x = x.transpose(1,2)
|
||||
x = self.downsampler(x)
|
||||
x = self.quant(x)
|
||||
return x
|
||||
|
||||
def preprocess(self, audio):
|
||||
if audio.dim() == 2: # raw audio
|
||||
x = self.mel_spec(audio)
|
||||
x = safe_log(x)
|
||||
elif audio.dim() == 3: # mel spectrogram
|
||||
x = audio
|
||||
return x
|
||||
|
||||
def forward(self, audio):
|
||||
x = self.preprocess(audio)
|
||||
|
||||
# print("done preprocessing: ",x.shape)
|
||||
# import pdb;pdb.set_trace()
|
||||
|
||||
x = self.encode(x)
|
||||
codes = self.quant.to_codebook_index(x)
|
||||
return codes
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Adapted from https://github.com/duchenzhuang/FSQ-pytorch/blob/main/quantizers/fsq.py#L41
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class FSQSTE(nn.Module):
|
||||
def __init__(self, levels):
|
||||
super().__init__()
|
||||
if levels:
|
||||
self.dim = len(levels)
|
||||
self._levels = torch.tensor(levels, dtype=torch.int32).view(1, 1, self.dim)
|
||||
else:
|
||||
self._levels = levels
|
||||
|
||||
_levels = self._levels
|
||||
self.register_buffer("levels", _levels, persistent=False)
|
||||
|
||||
_basis = torch.cumprod(torch.tensor([1] + levels[:-1]),
|
||||
dim=0,
|
||||
dtype=torch.int32)
|
||||
self.register_buffer("_basis", _basis, persistent=False)
|
||||
|
||||
def _scale_and_shift(self, zhat_normalized):
|
||||
half_width = self.levels // 2
|
||||
return (zhat_normalized * half_width) + half_width
|
||||
|
||||
def _scale_and_shift_inverse(self, zhat):
|
||||
half_width = self.levels // 2
|
||||
return (zhat - half_width) / half_width
|
||||
|
||||
def indices_to_level_indices(self, indices):
|
||||
""" Converts indices to indices at each level, perhaps needed for a transformer with factorized embeddings """
|
||||
indices = rearrange(indices, '... -> ... 1')
|
||||
codes_non_centered = (indices // self._basis) % self.levels
|
||||
return codes_non_centered
|
||||
|
||||
def to_codebook_index(self, zhat):
|
||||
""" Converts a `code` to an index in the codebook. """
|
||||
assert zhat.shape[-1] == self.dim
|
||||
zhat = self._scale_and_shift(zhat)
|
||||
indices = (zhat * self._basis).sum(dim = -1).round().to(torch.int32)
|
||||
return indices
|
||||
|
||||
def from_codebook_index(self, indices):
|
||||
""" Inverse of `codes_to_indices`. """
|
||||
level_indices = self.indices_to_level_indices(indices)
|
||||
codes = self._scale_and_shift_inverse(level_indices)
|
||||
return codes
|
||||
|
||||
def forward(self, x):
|
||||
if self.levels is not None:
|
||||
|
||||
half_levels = (self.levels - 1) * (1 - 1e-3) / 2
|
||||
offset = 0.5 - 0.5 * (self.levels % 2)
|
||||
shift = torch.tan(offset / half_levels)
|
||||
|
||||
x = torch.tanh(x + shift) * half_levels - offset
|
||||
x = x + (x.round() - x).detach()
|
||||
x = x / (self.levels // 2)
|
||||
return x
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
# ==============================================================================
|
||||
# S O P R A N O F A C T O R Y
|
||||
# ==============================================================================
|
||||
|
||||
dataset:
|
||||
input_dir: "./datasets/mio" # Must contain metadata.csv and wavs/
|
||||
sample_rate: 32000
|
||||
val_split_prop: 0.1 # 10% of data for validation
|
||||
|
||||
model:
|
||||
base_llm: "ekwek/Soprano-1.1-80M"
|
||||
llm_save_dir: "./finetuned_llm"
|
||||
decoder_save_dir: "./finetuned_decoder"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# STAGE 1: LLM AUTOREGRESSIVE TRAINING
|
||||
# ------------------------------------------------------------------------------
|
||||
llm_training:
|
||||
batch_size: 8
|
||||
grad_accum_steps: 4 # Effective Batch = 32
|
||||
seq_len: 1024
|
||||
max_steps: 25000 # Causal LLMs need time to learn grammar
|
||||
learning_rate: 0.0001
|
||||
weight_decay: 0.1
|
||||
val_freq: 250
|
||||
save_freq: 1000
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# STAGE 2: DECODER GAN TRAINING
|
||||
# ------------------------------------------------------------------------------
|
||||
decoder_training:
|
||||
batch_size: 4 # GAN training is memory intensive; keep small
|
||||
max_steps: 20000 # Pre-trained decoders adapt quickly
|
||||
learning_rate_g: 0.0002
|
||||
learning_rate_d: 0.0001
|
||||
segment_size_samples: 32768
|
||||
val_freq: 500
|
||||
save_freq: 2000
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# HARDWARE OPTIMIZATIONS
|
||||
# ------------------------------------------------------------------------------
|
||||
optimizations:
|
||||
attn_implementation: "sdpa"
|
||||
allow_tf32: true
|
||||
mixed_precision: "bfloat16"
|
||||
seed: 1337
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# TELEMETRY & LOGGING
|
||||
# ------------------------------------------------------------------------------
|
||||
logging:
|
||||
tensorboard_dir: "./runs"
|
||||
log_file: "training.log"
|
||||
log_audio_freq: 500
|
||||
@@ -1,55 +0,0 @@
|
||||
# config/config_5090.yaml
|
||||
|
||||
################################################################################
|
||||
# Soprano-Reforged Configuration
|
||||
# Optimized for NVIDIA RTX 5090 (32GB VRAM)
|
||||
# Extreme Performance Profile
|
||||
################################################################################
|
||||
|
||||
# --- Common Settings ---
|
||||
common:
|
||||
sample_rate: 32000 # Audio sample rate (Hz)
|
||||
seed: 1337 # Random seed for reproducibility
|
||||
device: "cuda" # Force CUDA
|
||||
|
||||
# --- Stage 0: Codec (Encoder/Decoder) Training ---
|
||||
codec:
|
||||
# Architecture
|
||||
input_mels: 80
|
||||
dim: 512
|
||||
layers: 8
|
||||
bottleneck: 5
|
||||
|
||||
# Training
|
||||
segment_size: 32768 # ~1 sec
|
||||
batch_size: 128 # 4x standard
|
||||
epochs: 100
|
||||
lr: 1.0e-4
|
||||
save_dir: "weights/codec"
|
||||
|
||||
# --- Stage 1: Dataset Generation ---
|
||||
dataset:
|
||||
val_split: 0.1
|
||||
|
||||
# --- Stage 2: Joint Training (LLM + Decoder) ---
|
||||
training:
|
||||
# Model
|
||||
base_model: "ekwek/Soprano-1.1-80M"
|
||||
seq_len: 1024 # Context window
|
||||
|
||||
# Hyperparameters - RTX 5090 EXTREME
|
||||
batch_size: 64 # 8x larger than default (32GB VRAM can handle this for 80M model)
|
||||
epochs: 1500 # Train longer by default
|
||||
base_lr: 1.0e-4 # Learning rate
|
||||
decoder_lr: 1.0e-4 # Learning rate for the Decoder
|
||||
weight_decay: 0.01
|
||||
grad_accum_steps: 1 # True gradients every step (No simulation needed)
|
||||
|
||||
# Joint Training Logic
|
||||
decoder_step_freq: 1 # Train the Decoder EVERY step (Critical for convergence)
|
||||
decoder_loss_weight: 1.0
|
||||
|
||||
# Checkpointing
|
||||
checkpoint_freq: 500 # Save more often since it's fast
|
||||
val_freq: 250 # Validate often
|
||||
save_dir: "weights/model"
|
||||
@@ -1,54 +0,0 @@
|
||||
# config/settings.yaml
|
||||
|
||||
################################################################################
|
||||
# Soprano-Reforged Configuration
|
||||
# Optimized for RTX 4070 (12GB) - Exclusive Mode (No Background Apps)
|
||||
################################################################################
|
||||
|
||||
# --- Common Settings ---
|
||||
common:
|
||||
sample_rate: 32000 # Audio sample rate (Hz)
|
||||
seed: 1337 # Random seed for reproducibility
|
||||
device: "auto" # "cuda", "cpu", or "auto" to detect
|
||||
|
||||
# --- Stage 0: Codec (Encoder/Decoder) Training ---
|
||||
codec:
|
||||
# Architecture
|
||||
input_mels: 80 # Number of Mel bands for spectrogram
|
||||
dim: 512 # Hidden dimension size
|
||||
layers: 8 # Number of ConvNeXt layers
|
||||
bottleneck: 5 # Bottleneck dimension (matches FSQ levels)
|
||||
|
||||
# Training
|
||||
segment_size: 32768 # Audio segment length (samples) per step (~1 sec)
|
||||
batch_size: 32 # UPGRADED: 4 -> 32 (Fast training for future runs)
|
||||
epochs: 100 # Total training epochs
|
||||
lr: 1.0e-4 # Learning rate
|
||||
save_dir: "weights/codec" # Where to save encoder.pth / decoder.pth
|
||||
|
||||
# --- Stage 1: Dataset Generation ---
|
||||
dataset:
|
||||
val_split: 0.1 # Percentage of data to use for validation (0.1 = 10%)
|
||||
|
||||
# --- Stage 2: Joint Training (LLM + Decoder) ---
|
||||
training:
|
||||
# Model
|
||||
base_model: "ekwek/Soprano-1.1-80M" # HuggingFace ID of the base LLM
|
||||
seq_len: 1024 # Context window size (text + audio tokens)
|
||||
|
||||
# Hyperparameters - RTX 4070 OPTIMIZED
|
||||
batch_size: 8 # The Safe Zone
|
||||
epochs: 10 # Total training epochs
|
||||
base_lr: 1.0e-4 # Learning rate for the LLM
|
||||
decoder_lr: 1.0e-4 # Learning rate for the Decoder
|
||||
weight_decay: 0.01 # Weight decay for optimizer
|
||||
grad_accum_steps: 4 # The Speed Hack (Simulates batch size 32)
|
||||
|
||||
# Joint Training Logic
|
||||
decoder_step_freq: 1 # Train the Decoder EVERY step (was 5)
|
||||
decoder_loss_weight: 1.0 # Weight of the audio reconstruction loss
|
||||
|
||||
# Checkpointing
|
||||
checkpoint_freq: 1000 # Save checkpoint every N steps
|
||||
val_freq: 500 # Validate every N steps
|
||||
save_dir: "weights/model" # Output directory for the final model
|
||||
@@ -1,96 +0,0 @@
|
||||
import sys
|
||||
import subprocess
|
||||
import os
|
||||
import json
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
def run_command(cmd, desc):
|
||||
print(f"\n[Step 1] Starting: {desc}")
|
||||
print(f"Command: {' '.join(cmd)}")
|
||||
try:
|
||||
subprocess.check_call(cmd)
|
||||
print(f"[Step 1] Completed: {desc}\n")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"[Step 1] Error during {desc}: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def verify_data_prep(output_dir):
|
||||
print("[Step 1] Verifying outputs...")
|
||||
train_json = output_dir / "train.json"
|
||||
val_json = output_dir / "val.json"
|
||||
|
||||
if not train_json.exists():
|
||||
print(f"[ERROR] train.json not found at {train_json}")
|
||||
sys.exit(1)
|
||||
|
||||
if not val_json.exists():
|
||||
print(f"[ERROR] val.json not found at {val_json}")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
with open(train_json, "r", encoding="utf-8") as f:
|
||||
train_data = json.load(f)
|
||||
with open(val_json, "r", encoding="utf-8") as f:
|
||||
val_data = json.load(f)
|
||||
|
||||
print(f"[SUCCESS] Data Prep Complete!")
|
||||
print(f" - Train Samples: {len(train_data)}")
|
||||
print(f" - Val Samples: {len(val_data)}")
|
||||
|
||||
if len(train_data) == 0:
|
||||
print("[WARNING] Train dataset is empty! Check your input metadata.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Failed to read JSON outputs: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Soprano Custom Pipeline - Step 1: Data Prep")
|
||||
parser.add_argument("--test", action="store_true", help="Run in test mode (limit to 50 samples)")
|
||||
args = parser.parse_args()
|
||||
|
||||
root_dir = Path(os.getcwd())
|
||||
dataset_dir = root_dir / "mio_dataset"
|
||||
custom_dir = root_dir / "custom"
|
||||
encoder_path = root_dir / "weights" / "codec" / "encoder.pth"
|
||||
|
||||
# Ensure custom directory exists
|
||||
custom_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Check inputs
|
||||
if not dataset_dir.exists():
|
||||
print(f"[ERROR] Dataset directory not found: {dataset_dir}")
|
||||
print("Please ensure 'mio_dataset' is in the project root.")
|
||||
sys.exit(1)
|
||||
|
||||
if not encoder_path.exists():
|
||||
print(f"[ERROR] Encoder checkpoint not found: {encoder_path}")
|
||||
sys.exit(1)
|
||||
|
||||
# Run Generation
|
||||
script_path = root_dir / "generate_dataset.py"
|
||||
if not script_path.exists():
|
||||
print(f"[ERROR] generate_dataset.py not found at {script_path}")
|
||||
sys.exit(1)
|
||||
|
||||
cmd = [
|
||||
sys.executable, str(script_path),
|
||||
"--input-dir", str(dataset_dir),
|
||||
"--output-dir", str(custom_dir),
|
||||
"--encoder-ckpt", str(encoder_path)
|
||||
]
|
||||
|
||||
if args.test:
|
||||
cmd.extend(["--limit", "50"])
|
||||
|
||||
run_command(cmd, "Data Preparation (Audio -> Tokens)")
|
||||
|
||||
# Verify
|
||||
verify_data_prep(custom_dir)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,77 +0,0 @@
|
||||
import sys
|
||||
import subprocess
|
||||
import os
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
def run_command(cmd, desc):
|
||||
print(f"\n[Step 2] Starting: {desc}")
|
||||
print(f"Command: {' '.join(cmd)}")
|
||||
try:
|
||||
subprocess.check_call(cmd)
|
||||
print(f"[Step 2] Completed: {desc}\n")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"[Step 2] Error during {desc}: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def verify_training(save_dir, epochs):
|
||||
print("[Step 2] Verifying training artifacts...")
|
||||
|
||||
# Check for epoch folder
|
||||
# Assuming epoch numbering starts at 0
|
||||
last_epoch_idx = epochs - 1
|
||||
epoch_dir = save_dir / f"epoch_{last_epoch_idx}"
|
||||
|
||||
decoder_pth = epoch_dir / "decoder.pth"
|
||||
config_json = epoch_dir / "config.json"
|
||||
|
||||
if epoch_dir.exists() and decoder_pth.exists() and config_json.exists():
|
||||
print(f"[SUCCESS] Training completed! Found checkpoint at: {epoch_dir}")
|
||||
print(f"Artifacts verified: decoder.pth, config.json")
|
||||
else:
|
||||
# Fallback check (sometimes epochs might be saved differently or crashed)
|
||||
print(f"[WARNING] Could not find specific epoch folder: {epoch_dir}")
|
||||
print(f"Listing {save_dir} contents:")
|
||||
for item in save_dir.glob("*"):
|
||||
print(f" - {item.name}")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Soprano Custom Pipeline - Step 2: Training")
|
||||
parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
|
||||
parser.add_argument("--test", action="store_true", help="Run in test mode (1 epoch)")
|
||||
args = parser.parse_args()
|
||||
|
||||
epochs = 1 if args.test else args.epochs
|
||||
|
||||
root_dir = Path(os.getcwd())
|
||||
custom_dir = root_dir / "custom"
|
||||
weights_dir = custom_dir / "weights" / "model"
|
||||
|
||||
# Check if Step 1 was run
|
||||
train_json = custom_dir / "train.json"
|
||||
if not train_json.exists():
|
||||
print(f"[ERROR] train.json not found in {custom_dir}. Please run Step 1 first.")
|
||||
sys.exit(1)
|
||||
|
||||
# Run Training
|
||||
train_script = root_dir / "train.py"
|
||||
if not train_script.exists():
|
||||
print(f"[ERROR] train.py not found at {train_script}")
|
||||
sys.exit(1)
|
||||
|
||||
cmd = [
|
||||
sys.executable, str(train_script),
|
||||
"--input-dir", str(custom_dir), # Read JSONs from custom dir
|
||||
"--save-dir", str(weights_dir),
|
||||
"--epochs", str(epochs)
|
||||
]
|
||||
run_command(cmd, "LLM Training (Stage 2)")
|
||||
|
||||
# Verify
|
||||
verify_training(weights_dir, epochs)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,82 +0,0 @@
|
||||
import sys
|
||||
import subprocess
|
||||
import os
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
def run_command(cmd, desc):
|
||||
print(f"\n[Step 3] Starting: {desc}")
|
||||
print(f"Command: {' '.join(cmd)}")
|
||||
try:
|
||||
subprocess.check_call(cmd)
|
||||
print(f"[Step 3] Completed: {desc}\n")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"[Step 3] Error during {desc}: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def verify_inference(output_path):
|
||||
print("[Step 3] Verifying output audio...")
|
||||
if not output_path.exists():
|
||||
print(f"[ERROR] Output file not found: {output_path}")
|
||||
sys.exit(1)
|
||||
|
||||
size = output_path.stat().st_size
|
||||
if size < 1000: # 1KB is suspiciously small for audio
|
||||
print(f"[WARNING] Output file is very small ({size} bytes). Generation might have failed or produced silence.")
|
||||
else:
|
||||
print(f"[SUCCESS] Audio generated at: {output_path}")
|
||||
print(f"File Size: {size / 1024:.2f} KB")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Soprano Custom Pipeline - Step 3: Inference")
|
||||
parser.add_argument("--text", type=str, default="Hello world, this is a test of the custom pipeline.", help="Text to synthesize")
|
||||
parser.add_argument("--epoch", type=int, default=None, help="Specific epoch number to use (default: latest)")
|
||||
args = parser.parse_args()
|
||||
|
||||
root_dir = Path(os.getcwd())
|
||||
custom_dir = root_dir / "custom"
|
||||
weights_dir = custom_dir / "weights" / "model"
|
||||
samples_dir = custom_dir / "samples"
|
||||
|
||||
samples_dir.mkdir(exist_ok=True)
|
||||
|
||||
# improved epoch finding logic
|
||||
if args.epoch is not None:
|
||||
model_path = weights_dir / f"epoch_{args.epoch}"
|
||||
else:
|
||||
# Find latest
|
||||
epochs = sorted([d for d in weights_dir.glob("epoch_*") if d.is_dir()],
|
||||
key=lambda x: int(x.name.split('_')[1]) if '_' in x.name else -1)
|
||||
if epochs:
|
||||
model_path = epochs[-1]
|
||||
print(f"[Info] Using latest checkpoint: {model_path.name}")
|
||||
else:
|
||||
print(f"[ERROR] No checkpoints found at {weights_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
if not model_path.exists():
|
||||
print(f"[ERROR] Model path does not exist: {model_path}")
|
||||
sys.exit(1)
|
||||
|
||||
output_wav = samples_dir / "output.wav"
|
||||
|
||||
script_path = root_dir / "inference.py"
|
||||
if not script_path.exists():
|
||||
print(f"[ERROR] inference.py not found at {script_path}")
|
||||
sys.exit(1)
|
||||
|
||||
cmd = [
|
||||
sys.executable, str(script_path),
|
||||
"--text", args.text,
|
||||
"--model-dir", str(model_path),
|
||||
"--output", str(output_wav)
|
||||
]
|
||||
run_command(cmd, "Inference Generation")
|
||||
|
||||
verify_inference(output_wav)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,113 +0,0 @@
|
||||
import argparse
|
||||
import subprocess
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
def run_command(cmd, desc):
|
||||
print(f"\n[Pipeline] Starting: {desc}")
|
||||
print(f"Command: {' '.join(cmd)}")
|
||||
try:
|
||||
subprocess.check_call(cmd)
|
||||
print(f"[Pipeline] Completed: {desc}\n")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"[Pipeline] Error during {desc}: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Soprano Custom Pipeline Orchestrator")
|
||||
parser.add_argument("--dataset-dir", type=str, default="mio_dataset", help="Input dataset directory")
|
||||
parser.add_argument("--custom-dir", type=str, default="custom", help="Directory for pipeline artifacts")
|
||||
|
||||
# Flags for stages
|
||||
parser.add_argument("--run-data-prep", action="store_true", help="Run generate_dataset.py")
|
||||
parser.add_argument("--run-codec", action="store_true", help="Run train_codec.py (Stage 0)")
|
||||
parser.add_argument("--run-training", action="store_true", help="Run train.py (Stage 2)")
|
||||
parser.add_argument("--run-inference", action="store_true", help="Run inference.py")
|
||||
|
||||
# Training Loop Args
|
||||
parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
|
||||
parser.add_argument("--text", type=str, default="Hello world, this is a test of the custom pipeline.", help="Inference text")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Setup Paths
|
||||
root_dir = Path(os.getcwd())
|
||||
dataset_dir = root_dir / args.dataset_dir
|
||||
custom_dir = root_dir / args.custom_dir
|
||||
|
||||
weights_dir = custom_dir / "weights"
|
||||
samples_dir = custom_dir / "samples"
|
||||
|
||||
# Ensure directories exist
|
||||
custom_dir.mkdir(exist_ok=True)
|
||||
weights_dir.mkdir(exist_ok=True)
|
||||
samples_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Check for Encoder (required for data prep and training)
|
||||
encoder_path = root_dir / "weights" / "codec" / "encoder.pth"
|
||||
if not encoder_path.exists():
|
||||
# Fallback to checking if it's in the standard location
|
||||
encoder_path = root_dir / "weights" / "codec" / "encoder.pth"
|
||||
if not encoder_path.exists():
|
||||
print(f"Warning: Encoder not found at {encoder_path}. Data prep might fail if not downloaded.")
|
||||
|
||||
# --- Stage 1: Data Preparation ---
|
||||
if args.run_data_prep:
|
||||
cmd = [
|
||||
sys.executable, "generate_dataset.py",
|
||||
"--input-dir", str(dataset_dir),
|
||||
"--output-dir", str(custom_dir), # Save JSONs to custom dir
|
||||
"--encoder-ckpt", str(encoder_path)
|
||||
]
|
||||
run_command(cmd, "Data Preparation")
|
||||
|
||||
# --- Stage 0: Codec Training (Optional) ---
|
||||
if args.run_codec:
|
||||
codec_save_dir = weights_dir / "codec"
|
||||
cmd = [
|
||||
sys.executable, "train_codec.py",
|
||||
"--wav-dir", str(dataset_dir / "wavs"),
|
||||
"--save-dir", str(codec_save_dir),
|
||||
"--epochs", str(args.epochs)
|
||||
]
|
||||
run_command(cmd, "Codec Training")
|
||||
|
||||
# --- Stage 2: LLM Training ---
|
||||
if args.run_training:
|
||||
model_save_dir = weights_dir / "model"
|
||||
cmd = [
|
||||
sys.executable, "train.py",
|
||||
"--input-dir", str(custom_dir), # Read JSONs from custom dir
|
||||
"--save-dir", str(model_save_dir),
|
||||
"--epochs", str(args.epochs)
|
||||
]
|
||||
run_command(cmd, "LLM Training")
|
||||
|
||||
# --- Inference ---
|
||||
if args.run_inference:
|
||||
# Find latest epoch or use default
|
||||
model_dir = weights_dir / "model"
|
||||
# Try to find the last epoch
|
||||
epochs = sorted([d for d in model_dir.glob("epoch_*") if d.is_dir()], key=lambda x: int(x.name.split('_')[1]))
|
||||
|
||||
if epochs:
|
||||
latest_model = epochs[-1]
|
||||
else:
|
||||
latest_model = model_dir # Fallback
|
||||
|
||||
output_wav = samples_dir / "inference_output.wav"
|
||||
|
||||
cmd = [
|
||||
sys.executable, "inference.py",
|
||||
"--text", args.text,
|
||||
"--model-dir", str(latest_model),
|
||||
"--output", str(output_wav)
|
||||
]
|
||||
run_command(cmd, "Inference")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+166
-30
@@ -1,42 +1,178 @@
|
||||
# dataset.py - Audio & Token Dataset Loader
|
||||
# dataset.py - Unified Dataset and Collators for the 2-Stage Soprano Pipeline.
|
||||
|
||||
import json
|
||||
from typing import List, Tuple, Dict
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
import torchaudio
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import PreTrainedTokenizer
|
||||
|
||||
class AudioDataset(Dataset):
|
||||
# Soprano Vocos ratio: 32000Hz / 2048 Total Hop Length
|
||||
SAMPLES_PER_TOKEN = 2048
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# STAGE 1: LLM AUTOREGRESSIVE DATASET & COLLATOR
|
||||
# ==============================================================================
|
||||
|
||||
class SopranoLLMDataset(Dataset):
|
||||
"""Dataset for Step 1: Training the Causal Language Model."""
|
||||
|
||||
def __init__(self, json_path: str):
|
||||
with open(json_path, 'r', encoding='utf-8') as f:
|
||||
self.data = json.load(f)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, idx: int) -> str:
|
||||
# JSON structure from generate_dataset.py: [text, audio_tokens, path]
|
||||
text, audio_tokens, _ = self.data[idx]
|
||||
|
||||
audio_str = "".join([f"[{token}]" for token in audio_tokens])
|
||||
formatted_sequence = f"[STOP][TEXT]{text}[START]{audio_str}[STOP]"
|
||||
return formatted_sequence
|
||||
|
||||
|
||||
class LLMCollator:
|
||||
"""Dynamically pads and formats sequences for the causal LLM."""
|
||||
|
||||
def __init__(self, tokenizer: PreTrainedTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def __call__(self, texts: List[str]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
# Dynamic Batching: Pad to the longest in this specific batch
|
||||
tokenized = self.tokenizer(
|
||||
texts,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=2048,
|
||||
return_tensors='pt',
|
||||
add_special_tokens=False
|
||||
)
|
||||
|
||||
batch = tokenized['input_ids']
|
||||
attn_mask = tokenized['attention_mask']
|
||||
|
||||
# Shift inputs and targets for causal language modeling
|
||||
x = batch[:, :-1]
|
||||
y = batch[:, 1:]
|
||||
attn_mask = attn_mask[:, :-1]
|
||||
|
||||
return x, y, attn_mask
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# STAGE 2: DECODER GAN DATASET & COLLATOR
|
||||
# ==============================================================================
|
||||
|
||||
class SopranoDecoderDataset(Dataset):
|
||||
"""Dataset for Step 2: Training the Vocos GAN to synthesize HD audio."""
|
||||
|
||||
def __init__(self, json_path: str, target_sr: int = 32000):
|
||||
with open(json_path, 'r', encoding='utf-8') as f:
|
||||
self.data = json.load(f)
|
||||
self.target_sr = target_sr
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, idx: int) -> Tuple[str, torch.Tensor, int]:
|
||||
text, audio_tokens, wav_path = self.data[idx]
|
||||
|
||||
# Format Text
|
||||
audio_str = "".join([f"[{token}]" for token in audio_tokens])
|
||||
formatted_sequence = f"[STOP][TEXT]{text}[START]{audio_str}[STOP]"
|
||||
|
||||
# Load Raw Audio using Soundfile (Bypasses Windows torchcodec errors)
|
||||
audio_np, sr = sf.read(wav_path, dtype='float32')
|
||||
|
||||
# Convert to proper Tensor shape: (Channels, Samples)
|
||||
if audio_np.ndim == 1:
|
||||
audio_np = np.expand_dims(audio_np, axis=0) # (1, T)
|
||||
else:
|
||||
audio_np = audio_np.T # (C, T)
|
||||
|
||||
wav = torch.from_numpy(audio_np)
|
||||
|
||||
# Mono conversion
|
||||
if wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
|
||||
# Resample if needed
|
||||
if sr != self.target_sr:
|
||||
wav = torchaudio.functional.resample(wav, orig_freq=sr, new_freq=self.target_sr)
|
||||
|
||||
wav = wav.squeeze(0) # Shape: (Samples,)
|
||||
num_audio_tokens = len(audio_tokens)
|
||||
|
||||
return formatted_sequence, wav, num_audio_tokens
|
||||
|
||||
|
||||
class DecoderCollator:
|
||||
"""
|
||||
Loads the JSON dataset generated by Stage 1.
|
||||
Formats inputs into the specific text-audio sequence expected by Soprano.
|
||||
Complex Collator for the GAN: Aligns raw audio samples with the LLM's
|
||||
discrete tokens so the STFT Discriminators can grade specific slices.
|
||||
"""
|
||||
def __init__(self, path):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
self.dataset = json.load(f)
|
||||
print(f"Loaded {len(self.dataset)} samples from {path}")
|
||||
|
||||
def __init__(self, tokenizer: PreTrainedTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
# Dynamically fetch boundaries so it doesn't break on tokenizer updates
|
||||
self.audio_min_id = int(self.tokenizer.convert_tokens_to_ids("[0]")) # type: ignore
|
||||
self.audio_max_id = int(self.tokenizer.convert_tokens_to_ids("[7999]")) # type: ignore
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
def __call__(self, batch_in: List[Tuple[str, torch.Tensor, int]]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
texts = [x[0] for x in batch_in]
|
||||
wavs = [x[1] for x in batch_in]
|
||||
aud_token_lens = [x[2] for x in batch_in]
|
||||
|
||||
batch_tokens_list = []
|
||||
batch_audio_list = []
|
||||
|
||||
for i in range(len(texts)):
|
||||
# Tokenize raw
|
||||
raw_tokens = self.tokenizer(texts[i], padding=False, truncation=False, add_special_tokens=False)['input_ids']
|
||||
tokens = torch.tensor(raw_tokens, dtype=torch.long)
|
||||
|
||||
wav = wavs[i]
|
||||
num_aud_tokens = aud_token_lens[i]
|
||||
|
||||
# Create a blank canvas for the aligned audio
|
||||
target_length = num_aud_tokens * SAMPLES_PER_TOKEN
|
||||
aligned_audio = torch.zeros(target_length, dtype=torch.float32)
|
||||
|
||||
# Vectorized Audio Slicing (Replaces the slow for-loop from the fork)
|
||||
actual_length = min(wav.size(0), target_length)
|
||||
aligned_audio[:actual_length] = wav[:actual_length]
|
||||
|
||||
# Validate integrity
|
||||
is_audio = (tokens >= self.audio_min_id) & (tokens <= self.audio_max_id)
|
||||
audio_indices = torch.where(is_audio)[0]
|
||||
|
||||
assert len(audio_indices) == num_aud_tokens, \
|
||||
f"Token mismatch! Found {len(audio_indices)} audio tokens but expected {num_aud_tokens}."
|
||||
|
||||
batch_tokens_list.append(tokens)
|
||||
batch_audio_list.append(aligned_audio)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
item = self.dataset[idx]
|
||||
text = item["text"]
|
||||
tokens = item["tokens"] # List of integers from Stage 1
|
||||
wav_path = item["wav_path"]
|
||||
# Pad sequences
|
||||
pad_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0
|
||||
batch_tokens = torch.nn.utils.rnn.pad_sequence(batch_tokens_list, batch_first=True, padding_value=pad_id)
|
||||
batch_audio = torch.nn.utils.rnn.pad_sequence(batch_audio_list, batch_first=True, padding_value=0.0)
|
||||
|
||||
# 1. Formatting the Token Sequence
|
||||
# Soprano uses a specific vocabulary where each audio code is represented
|
||||
# as a bracketed string: e.g., token 8500 becomes "[8500]".
|
||||
token_str = "".join([f"[{t}]" for t in tokens])
|
||||
# Shift inputs and targets
|
||||
x = batch_tokens[:, :-1]
|
||||
y = batch_tokens[:, 1:]
|
||||
|
||||
# 2. Constructing the Full Training Prompt
|
||||
# [STOP] is the separator, [TEXT] marks the prompt, [START] marks the audio segment.
|
||||
# This matches the inference prompt exactly: f"[STOP][TEXT]{norm_text}[START]"
|
||||
formatted_text = f"[STOP][TEXT]{text}[START]{token_str}[STOP]"
|
||||
# Crop audio to match the shifted 'x' tensor length
|
||||
max_len_x = x.size(1)
|
||||
gt_audio = batch_audio[:, :max_len_x * SAMPLES_PER_TOKEN]
|
||||
|
||||
# 3. Return Dictionary
|
||||
# We return the raw text because the Collator handles tokenization on the fly.
|
||||
# We return wav_path so the Decoder-step in Stage 2 can load ground-truth audio.
|
||||
return {
|
||||
"text": formatted_text,
|
||||
"wav_path": wav_path,
|
||||
}
|
||||
# Create Boolean Mask where output equals audio token
|
||||
audio_mask = (y >= self.audio_min_id) & (y <= self.audio_max_id)
|
||||
|
||||
return x, y, gt_audio, audio_mask
|
||||
@@ -0,0 +1,45 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .models import VocosBackbone
|
||||
from .heads import ISTFTHead
|
||||
|
||||
|
||||
class SopranoDecoder(nn.Module):
|
||||
def __init__(self,
|
||||
num_input_channels=512,
|
||||
decoder_num_layers=8,
|
||||
decoder_dim=768,
|
||||
decoder_intermediate_dim=None,
|
||||
hop_length=512,
|
||||
n_fft=2048,
|
||||
upscale=4,
|
||||
dw_kernel=3,
|
||||
):
|
||||
super().__init__()
|
||||
self.decoder_initial_channels = num_input_channels
|
||||
self.num_layers = decoder_num_layers
|
||||
self.dim = decoder_dim
|
||||
self.intermediate_dim = decoder_intermediate_dim if decoder_intermediate_dim else decoder_dim*3
|
||||
self.hop_length = hop_length
|
||||
self.n_fft = n_fft
|
||||
self.upscale = upscale
|
||||
self.dw_kernel = dw_kernel
|
||||
|
||||
self.decoder = VocosBackbone(input_channels=self.decoder_initial_channels,
|
||||
dim=self.dim,
|
||||
intermediate_dim=self.intermediate_dim,
|
||||
num_layers=self.num_layers,
|
||||
input_kernel_size=1,#dw_kernel,
|
||||
dw_kernel_size=dw_kernel,
|
||||
)
|
||||
self.head = ISTFTHead(dim=self.dim,
|
||||
n_fft=self.n_fft,
|
||||
hop_length=self.hop_length)
|
||||
|
||||
def forward(self, x):
|
||||
T = x.size(2)
|
||||
x = torch.nn.functional.interpolate(x, size=self.upscale*(T-1)+1, mode='linear', align_corners=True)
|
||||
x = self.decoder(x)
|
||||
reconstructed = self.head(x)
|
||||
return reconstructed
|
||||
@@ -0,0 +1,143 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from torch.nn.utils.parametrizations import weight_norm, spectral_norm
|
||||
|
||||
class DiscriminatorP(nn.Module):
|
||||
def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
|
||||
super(DiscriminatorP, self).__init__()
|
||||
self.period = period
|
||||
self.use_spectral_norm = use_spectral_norm
|
||||
norm_f = weight_norm if use_spectral_norm == False else spectral_norm
|
||||
self.convs = nn.ModuleList([
|
||||
norm_f(nn.Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(2, 0))),
|
||||
norm_f(nn.Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(2, 0))),
|
||||
norm_f(nn.Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(2, 0))),
|
||||
norm_f(nn.Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(2, 0))),
|
||||
norm_f(nn.Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
|
||||
])
|
||||
self.conv_post = norm_f(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
|
||||
|
||||
def forward(self, x):
|
||||
fmap = []
|
||||
|
||||
# 1d to 2d
|
||||
b, c, t = x.shape
|
||||
if t % self.period != 0: # pad valid
|
||||
n_pad = self.period - (t % self.period)
|
||||
x = F.pad(x, (0, n_pad), "reflect")
|
||||
t = t + n_pad
|
||||
x = x.view(b, c, t // self.period, self.period)
|
||||
|
||||
for l in self.convs:
|
||||
x = l(x)
|
||||
x = F.leaky_relu(x, 0.1)
|
||||
fmap.append(x)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
x = torch.flatten(x, 1, -1)
|
||||
|
||||
return x, fmap
|
||||
|
||||
|
||||
class DiscriminatorS(nn.Module):
|
||||
def __init__(self, use_spectral_norm=False):
|
||||
super(DiscriminatorS, self).__init__()
|
||||
norm_f = weight_norm if use_spectral_norm == False else spectral_norm
|
||||
self.convs = nn.ModuleList([
|
||||
norm_f(nn.Conv1d(1, 16, 15, 1, padding=7)),
|
||||
norm_f(nn.Conv1d(16, 64, 41, 4, groups=4, padding=20)),
|
||||
norm_f(nn.Conv1d(64, 256, 41, 4, groups=16, padding=20)),
|
||||
norm_f(nn.Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
|
||||
norm_f(nn.Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
|
||||
norm_f(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
|
||||
])
|
||||
self.conv_post = norm_f(nn.Conv1d(1024, 1, 3, 1, padding=1))
|
||||
|
||||
def forward(self, x):
|
||||
fmap = []
|
||||
for l in self.convs:
|
||||
x = l(x)
|
||||
x = F.leaky_relu(x, 0.1)
|
||||
fmap.append(x)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
x = torch.flatten(x, 1, -1)
|
||||
|
||||
return x, fmap
|
||||
|
||||
|
||||
class MultiPeriodDiscriminator(nn.Module):
|
||||
def __init__(self, use_spectral_norm=False):
|
||||
super(MultiPeriodDiscriminator, self).__init__()
|
||||
self.discriminators = nn.ModuleList([
|
||||
DiscriminatorP(2, use_spectral_norm=use_spectral_norm),
|
||||
DiscriminatorP(3, use_spectral_norm=use_spectral_norm),
|
||||
DiscriminatorP(5, use_spectral_norm=use_spectral_norm),
|
||||
DiscriminatorP(7, use_spectral_norm=use_spectral_norm),
|
||||
DiscriminatorP(11, use_spectral_norm=use_spectral_norm),
|
||||
])
|
||||
|
||||
def forward(self, y, y_hat):
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
for i, d in enumerate(self.discriminators):
|
||||
y_d_r, fmap_r = d(y)
|
||||
y_d_g, fmap_g = d(y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_rs.append(fmap_r)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
class MultiScaleDiscriminator(nn.Module):
|
||||
def __init__(self, use_spectral_norm=False):
|
||||
super(MultiScaleDiscriminator, self).__init__()
|
||||
self.discriminators = nn.ModuleList([
|
||||
DiscriminatorS(use_spectral_norm=use_spectral_norm),
|
||||
DiscriminatorS(use_spectral_norm=use_spectral_norm),
|
||||
DiscriminatorS(use_spectral_norm=use_spectral_norm),
|
||||
])
|
||||
self.meanpools = nn.ModuleList([
|
||||
nn.AvgPool1d(4, 2, padding=2),
|
||||
nn.AvgPool1d(4, 2, padding=2)
|
||||
])
|
||||
|
||||
def forward(self, y, y_hat):
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
for i, d in enumerate(self.discriminators):
|
||||
if i != 0:
|
||||
y = self.meanpools[i-1](y)
|
||||
y_hat = self.meanpools[i-1](y_hat)
|
||||
y_d_r, fmap_r = d(y)
|
||||
y_d_g, fmap_g = d(y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_rs.append(fmap_r)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.mpd = MultiPeriodDiscriminator()
|
||||
self.msd = MultiScaleDiscriminator()
|
||||
|
||||
def forward(self, y, y_hat):
|
||||
# y: real audio, y_hat: gen audio
|
||||
# Unsqueeze if needed (B, T) -> (B, 1, T)
|
||||
if y.ndim == 2: y = y.unsqueeze(1)
|
||||
if y_hat.ndim == 2: y_hat = y_hat.unsqueeze(1)
|
||||
|
||||
y_d_rs_p, y_d_gs_p, fmap_rs_p, fmap_gs_p = self.mpd(y, y_hat)
|
||||
y_d_rs_s, y_d_gs_s, fmap_rs_s, fmap_gs_s = self.msd(y, y_hat)
|
||||
|
||||
return (y_d_rs_p + y_d_rs_s, y_d_gs_p + y_d_gs_s, fmap_rs_p + fmap_rs_s, fmap_gs_p + fmap_gs_s)
|
||||
@@ -0,0 +1,50 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from .spectral_ops import ISTFT
|
||||
|
||||
|
||||
class ISTFTHead(nn.Module):
|
||||
"""
|
||||
ISTFT Head module for predicting STFT complex coefficients.
|
||||
|
||||
Args:
|
||||
dim (int): Hidden dimension of the model.
|
||||
n_fft (int): Size of Fourier transform.
|
||||
hop_length (int): The distance between neighboring sliding window frames, which should align with
|
||||
the resolution of the input features.
|
||||
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "center"):
|
||||
super().__init__()
|
||||
out_dim = n_fft + 2
|
||||
self.out = torch.nn.Linear(dim, out_dim)
|
||||
self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
|
||||
|
||||
@torch.compiler.disable
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass of the ISTFTHead module.
|
||||
|
||||
Args:
|
||||
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
|
||||
L is the sequence length, and H denotes the model dimension.
|
||||
|
||||
Returns:
|
||||
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
|
||||
"""
|
||||
x = self.out(x.transpose(1,2)).transpose(1, 2)
|
||||
mag, p = x.chunk(2, dim=1)
|
||||
mag = torch.exp(mag)
|
||||
mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes
|
||||
# wrapping happens here. These two lines produce real and imaginary value
|
||||
x = torch.cos(p)
|
||||
y = torch.sin(p)
|
||||
# recalculating phase here does not produce anything new
|
||||
# only costs time
|
||||
# phase = torch.atan2(y, x)
|
||||
# S = mag * torch.exp(phase * 1j)
|
||||
# better directly produce the complex value
|
||||
S = mag * (x + 1j * y)
|
||||
audio = self.istft(S)
|
||||
return audio
|
||||
@@ -0,0 +1,144 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchaudio
|
||||
|
||||
def dynamic_range_compression_torch(x, C=1, clip_val: float = 5e-3):
|
||||
return torch.log(torch.clamp(x, min=clip_val) * C)
|
||||
|
||||
def spectral_normalize_torch(magnitudes):
|
||||
output = dynamic_range_compression_torch(magnitudes)
|
||||
return output
|
||||
|
||||
# Mel Spectrogram initialization logic
|
||||
# Note: Should we initialize the transform here or pass it in?
|
||||
# For simplicity, we can provide a factory function or a class.
|
||||
# But original code used global `mel_transform`.
|
||||
# To keep it clean, let's wrap it in a class or function that returns the transform.
|
||||
|
||||
class MelSpectrogramWrapper(torch.nn.Module):
|
||||
def __init__(self, sample_rate=32000, n_fft=2048, hop_length=512, n_mels=50):
|
||||
super().__init__()
|
||||
self.mel_transform = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=sample_rate, n_fft=n_fft, hop_length=hop_length,
|
||||
n_mels=n_mels, center=True, power=1
|
||||
)
|
||||
|
||||
def forward(self, audio):
|
||||
mel_spec = self.mel_transform(audio)
|
||||
mel_spec = spectral_normalize_torch(mel_spec)
|
||||
return mel_spec
|
||||
|
||||
def feature_matching_loss(fmap_r, fmap_g):
|
||||
loss = 0
|
||||
for dr, dg in zip(fmap_r, fmap_g):
|
||||
for rl, gl in zip(dr, dg):
|
||||
loss += torch.mean(torch.abs(rl - gl))
|
||||
return loss * 2
|
||||
|
||||
def discriminator_loss(disc_real_outputs, disc_generated_outputs):
|
||||
loss = 0
|
||||
r_losses = []
|
||||
g_losses = []
|
||||
for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
|
||||
r_loss = torch.mean((1-dr)**2)
|
||||
g_loss = torch.mean(dg**2)
|
||||
loss += (r_loss + g_loss)
|
||||
r_losses.append(r_loss.item())
|
||||
g_losses.append(g_loss.item())
|
||||
|
||||
return loss, r_losses, g_losses
|
||||
|
||||
def generator_loss(disc_outputs):
|
||||
loss = 0
|
||||
gen_losses = []
|
||||
for dg in disc_outputs:
|
||||
l = torch.mean((1-dg)**2)
|
||||
gen_losses.append(l)
|
||||
loss += l
|
||||
|
||||
return loss, gen_losses
|
||||
|
||||
# -----------------
|
||||
# Multi-Resolution STFT Loss
|
||||
# -----------------
|
||||
|
||||
def stft(x, fft_size, hop_size, win_length, window):
|
||||
"""Perform STFT and return linear mag spectogram."""
|
||||
# x: (B, T)
|
||||
x_stft = torch.stft(x, fft_size, hop_size, win_length, window, center=True, return_complex=True)
|
||||
x_mag = torch.abs(x_stft)
|
||||
return x_mag
|
||||
|
||||
def spectral_convergence_loss(x_mag, y_mag):
|
||||
"""
|
||||
Spectral convergence loss.
|
||||
"""
|
||||
return torch.norm(y_mag - x_mag, p="fro") / (torch.norm(y_mag, p="fro") + 1e-7)
|
||||
|
||||
def log_magnitude_loss(x_mag, y_mag):
|
||||
"""
|
||||
Log-magnitude L1 loss.
|
||||
"""
|
||||
return F.l1_loss(torch.log(x_mag), torch.log(y_mag))
|
||||
|
||||
class STFTLoss(torch.nn.Module):
|
||||
"""
|
||||
STFT Loss module.
|
||||
"""
|
||||
def __init__(self, fft_size, hop_size, win_length, window="hann_window"):
|
||||
super(STFTLoss, self).__init__()
|
||||
self.fft_size = fft_size
|
||||
self.hop_size = hop_size
|
||||
self.win_length = win_length
|
||||
self.register_buffer("window", getattr(torch, window)(win_length))
|
||||
|
||||
def forward(self, x, y):
|
||||
"""
|
||||
Args:
|
||||
x (Tensor): Predicted audio (B, T).
|
||||
y (Tensor): Target audio (B, T).
|
||||
"""
|
||||
x_mag = stft(x, self.fft_size, self.hop_size, self.win_length, self.window)
|
||||
y_mag = stft(y, self.fft_size, self.hop_size, self.win_length, self.window)
|
||||
|
||||
# Add epsilon to prevent log(0)
|
||||
x_mag = torch.clamp(x_mag, min=1e-7)
|
||||
y_mag = torch.clamp(y_mag, min=1e-7)
|
||||
|
||||
sc_loss = spectral_convergence_loss(x_mag, y_mag)
|
||||
mag_loss = log_magnitude_loss(x_mag, y_mag)
|
||||
|
||||
return sc_loss, mag_loss
|
||||
|
||||
class MultiResolutionSTFTLoss(torch.nn.Module):
|
||||
"""
|
||||
Multi-Resolution STFT Loss module.
|
||||
"""
|
||||
def __init__(self,
|
||||
fft_sizes=[1024, 2048, 512],
|
||||
hop_sizes=[120, 240, 50],
|
||||
win_lengths=[600, 1200, 240], # Standard HiFiGAN config
|
||||
window="hann_window", factor_sc=1.0, factor_mag=1.0):
|
||||
super(MultiResolutionSTFTLoss, self).__init__()
|
||||
assert len(fft_sizes) == len(hop_sizes) == len(win_lengths)
|
||||
|
||||
self.stft_losses = torch.nn.ModuleList()
|
||||
for fs, hs, wl in zip(fft_sizes, hop_sizes, win_lengths):
|
||||
self.stft_losses.append(STFTLoss(fs, hs, wl, window))
|
||||
|
||||
self.factor_sc = factor_sc
|
||||
self.factor_mag = factor_mag
|
||||
|
||||
def forward(self, x, y):
|
||||
sc_loss = 0.0
|
||||
mag_loss = 0.0
|
||||
|
||||
for f in self.stft_losses:
|
||||
sc_l, mag_l = f(x, y)
|
||||
sc_loss += sc_l
|
||||
mag_loss += mag_l
|
||||
|
||||
sc_loss /= len(self.stft_losses)
|
||||
mag_loss /= len(self.stft_losses)
|
||||
|
||||
return sc_loss, mag_loss
|
||||
@@ -0,0 +1,61 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .modules import ConvNeXtBlock
|
||||
|
||||
class VocosBackbone(nn.Module):
|
||||
"""
|
||||
Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
|
||||
|
||||
Args:
|
||||
input_channels (int): Number of input features channels.
|
||||
dim (int): Hidden dimension of the model.
|
||||
intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
|
||||
num_layers (int): Number of ConvNeXtBlock layers.
|
||||
layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
num_layers: int,
|
||||
input_kernel_size: int = 9,
|
||||
dw_kernel_size: int = 9,
|
||||
layer_scale_init_value: Optional[float] = None,
|
||||
pad: str = 'zeros',
|
||||
):
|
||||
super().__init__()
|
||||
self.embed = nn.Conv1d(input_channels, dim, kernel_size=input_kernel_size, padding=input_kernel_size//2, padding_mode=pad)
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.convnext = nn.ModuleList(
|
||||
[
|
||||
ConvNeXtBlock(
|
||||
dim=dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
dw_kernel_size=dw_kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value or 1 / num_layers**0.5,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.apply(self._init_weights)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None: nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.embed(x) # (B, C, L)
|
||||
x = self.norm(x.transpose(1, 2))
|
||||
x = x.transpose(1, 2)
|
||||
for conv_block in self.convnext:
|
||||
x = conv_block(x)
|
||||
x = self.final_layer_norm(x.transpose(1, 2))
|
||||
x = x.transpose(1, 2)
|
||||
return x
|
||||
@@ -0,0 +1,47 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
"""ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
intermediate_dim (int): Dimensionality of the intermediate layer.
|
||||
layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
|
||||
Defaults to None.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
layer_scale_init_value: float,
|
||||
dw_kernel_size: int = 9,
|
||||
):
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv1d(dim, dim, kernel_size=dw_kernel_size, padding=dw_kernel_size//2, groups=dim) # depthwise conv
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(intermediate_dim, dim)
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.dwconv(x)
|
||||
x = x.transpose(1, 2) # (B, C, T) -> (B, T, C)
|
||||
x = self.norm(x)
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
x = x.transpose(1, 2) # (B, T, C) -> (B, C, T)
|
||||
|
||||
x = residual + x
|
||||
return x
|
||||
@@ -0,0 +1,74 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
class ISTFT(nn.Module):
|
||||
"""
|
||||
Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
|
||||
windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
|
||||
See issue: https://github.com/pytorch/pytorch/issues/62323
|
||||
Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
|
||||
The NOLA constraint is met as we trim padded samples anyway.
|
||||
|
||||
Args:
|
||||
n_fft (int): Size of Fourier transform.
|
||||
hop_length (int): The distance between neighboring sliding window frames.
|
||||
win_length (int): The size of window frame and STFT filter.
|
||||
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
|
||||
"""
|
||||
|
||||
def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
|
||||
super().__init__()
|
||||
if padding not in ["center", "same"]:
|
||||
raise ValueError("Padding must be 'center' or 'same'.")
|
||||
self.padding = padding
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
self.win_length = win_length
|
||||
window = torch.hann_window(win_length)
|
||||
self.register_buffer("window", window)
|
||||
|
||||
def forward(self, spec: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
|
||||
|
||||
Args:
|
||||
spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
|
||||
N is the number of frequency bins, and T is the number of time frames.
|
||||
|
||||
Returns:
|
||||
Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
|
||||
"""
|
||||
if self.padding == "center":
|
||||
spec[:,0] = 0 # fixes some strange bug where first/last freqs don't matter when bs<16 which causes exploding gradients
|
||||
spec[:,-1] = 0
|
||||
# Fallback to pytorch native implementation
|
||||
return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
|
||||
elif self.padding == "same":
|
||||
pad = (self.win_length - self.hop_length) // 2
|
||||
else:
|
||||
raise ValueError("Padding must be 'center' or 'same'.")
|
||||
|
||||
assert spec.dim() == 3, "Expected a 3D tensor as input"
|
||||
B, N, T = spec.shape
|
||||
|
||||
# Inverse FFT
|
||||
ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
|
||||
ifft = ifft * self.window[None, :, None]
|
||||
|
||||
# Overlap and Add
|
||||
output_size = (T - 1) * self.hop_length + self.win_length
|
||||
y = torch.nn.functional.fold(
|
||||
ifft, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
|
||||
)[:, 0, 0, pad:-pad]
|
||||
|
||||
# Window envelope
|
||||
window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
|
||||
window_envelope = torch.nn.functional.fold(
|
||||
window_sq, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
|
||||
).squeeze()[pad:-pad]
|
||||
|
||||
# Normalize
|
||||
assert (window_envelope > 1e-11).all()
|
||||
y = y / window_envelope
|
||||
|
||||
return y
|
||||
@@ -1,2 +0,0 @@
|
||||
example1|Soprano is an extremely lightweight text to speech model designed to produce highly realistic speech at unprecedented speed.
|
||||
example2|Gabagool? Ova here!
|
||||
Binary file not shown.
Binary file not shown.
+137
-139
@@ -1,152 +1,150 @@
|
||||
# generate_dataset.py
|
||||
# generate_dataset.py - Preprocesses raw audio and text data into quantized tokens for SopranoTTS training.
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import random
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
import torch
|
||||
import torchaudio
|
||||
import soundfile as sf
|
||||
from tqdm import tqdm
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
|
||||
# Import our components
|
||||
from model.encoder import Encoder
|
||||
from utils.text_normalizer import normalize_text
|
||||
# Import Config
|
||||
from utils.config import cfg
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchaudio
|
||||
from huggingface_hub import hf_hub_download
|
||||
from scipy.io import wavfile
|
||||
from tqdm import tqdm
|
||||
|
||||
from codec.encoder.codec import Encoder
|
||||
|
||||
|
||||
class DatasetGenerator:
|
||||
"""
|
||||
Handles the preprocessing of raw audio and text data into quantized
|
||||
tokens suitable for training the Soprano causal language model.
|
||||
"""
|
||||
|
||||
def __init__(self, config_path: str):
|
||||
with open(config_path, 'r') as f:
|
||||
self.config = yaml.safe_load(f)
|
||||
|
||||
self.input_dir = Path(self.config['dataset']['input_dir'])
|
||||
self.sample_rate = self.config['dataset']['sample_rate']
|
||||
self.val_prop = self.config['dataset']['val_split_prop']
|
||||
self.seed = self.config['optimizations']['seed']
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
self.encoder = self._load_encoder()
|
||||
|
||||
def _load_encoder(self) -> Encoder:
|
||||
"""Downloads and loads the frozen Soprano audio encoder."""
|
||||
print("Loading Pre-trained Base Encoder...")
|
||||
encoder = Encoder()
|
||||
encoder_path = hf_hub_download(repo_id='ekwek/Soprano-Encoder', filename='encoder.pth')
|
||||
encoder.load_state_dict(torch.load(encoder_path, map_location='cpu'))
|
||||
encoder.to(self.device)
|
||||
encoder.eval()
|
||||
return encoder
|
||||
|
||||
def _normalize_audio(self, audio_np: np.ndarray) -> torch.Tensor:
|
||||
"""Safely converts numpy audio to a normalized float32 tensor in [-1.0, 1.0]."""
|
||||
if audio_np.ndim == 2:
|
||||
audio_np = audio_np.T
|
||||
|
||||
if audio_np.dtype == np.int16:
|
||||
audio = torch.from_numpy(audio_np).to(torch.float32) / 32768.0
|
||||
elif audio_np.dtype == np.int32:
|
||||
audio = torch.from_numpy(audio_np).to(torch.float32) / 2147483648.0
|
||||
else:
|
||||
audio = torch.from_numpy(audio_np).to(torch.float32)
|
||||
|
||||
if audio.ndim == 1:
|
||||
audio = audio.unsqueeze(0)
|
||||
return audio
|
||||
|
||||
def generate(self) -> None:
|
||||
metadata_path = self.input_dir / 'metadata.csv'
|
||||
if not metadata_path.exists():
|
||||
raise FileNotFoundError(f"Metadata file not found at {metadata_path}")
|
||||
|
||||
print("Reading metadata...")
|
||||
samples = []
|
||||
with open(metadata_path, 'r', encoding='utf-8') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split('|')
|
||||
if len(parts) >= 2:
|
||||
filename = parts[0].strip()
|
||||
transcript = parts[1].strip()
|
||||
samples.append({"filename": filename, "text": transcript})
|
||||
|
||||
print(f"Processing {len(samples)} audio samples...")
|
||||
dataset = []
|
||||
|
||||
for sample in tqdm(samples):
|
||||
wav_path = self.input_dir / 'wavs' / f"{sample['filename']}.wav"
|
||||
|
||||
if not wav_path.exists():
|
||||
print(f"\n [!] Missing audio file: {wav_path}. Skipping.")
|
||||
continue
|
||||
|
||||
# Load, normalize, and resample
|
||||
try:
|
||||
sr, audio_np = wavfile.read(wav_path)
|
||||
except Exception as e:
|
||||
print(f"\n [!] Failed to read {wav_path}: {e}")
|
||||
continue
|
||||
|
||||
audio_tensor = self._normalize_audio(audio_np)
|
||||
if sr != self.sample_rate:
|
||||
audio_tensor = torchaudio.functional.resample(
|
||||
audio_tensor, orig_freq=sr, new_freq=self.sample_rate
|
||||
)
|
||||
|
||||
# Encode to discrete tokens
|
||||
audio_tensor = audio_tensor.to(self.device)
|
||||
with torch.no_grad():
|
||||
audio_tokens = self.encoder(audio_tensor)
|
||||
|
||||
# CRITICAL FOR 2-STAGE: Save absolute path for the Decoder GAN
|
||||
abs_audio_path = str(wav_path.resolve())
|
||||
|
||||
# Output format matches the fork's expected dataset format exactly
|
||||
dataset.append([
|
||||
sample["text"],
|
||||
audio_tokens.squeeze(0).cpu().tolist(),
|
||||
abs_audio_path
|
||||
])
|
||||
|
||||
print("Generating train/test splits...")
|
||||
random.seed(self.seed)
|
||||
random.shuffle(dataset)
|
||||
|
||||
num_val = int(self.val_prop * len(dataset))
|
||||
num_val = max(1, min(num_val, 512)) # Ensure at least 1 val sample, cap at 512
|
||||
|
||||
val_dataset = dataset[:num_val]
|
||||
train_dataset = dataset[num_val:]
|
||||
|
||||
print(f"Train samples: {len(train_dataset)}")
|
||||
print(f"Val samples: {len(val_dataset)}")
|
||||
|
||||
with open(self.input_dir / 'train.json', 'w', encoding='utf-8') as f:
|
||||
json.dump(train_dataset, f, indent=2)
|
||||
with open(self.input_dir / 'val.json', 'w', encoding='utf-8') as f:
|
||||
json.dump(val_dataset, f, indent=2)
|
||||
|
||||
print("Datasets saved successfully.")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input-dir", type=str, required=True, help="Path to LJSpeech-formatted dataset")
|
||||
parser.add_argument("--encoder-ckpt", type=str, required=True, help="Path to your trained encoder.pth")
|
||||
parser.add_argument("--output-dir", type=str, default=None, help="Where to save train.json/val.json")
|
||||
parser.add_argument("--limit", type=int, default=None, help="Limit number of samples (for testing)")
|
||||
parser = argparse.ArgumentParser(description="Generate Soprano TTS Dataset")
|
||||
parser.add_argument("--config", type=str, default="config.yaml", help="Path to config file")
|
||||
args = parser.parse_args()
|
||||
|
||||
# 1. Setup Paths
|
||||
input_dir = Path(args.input_dir)
|
||||
output_dir = Path(args.output_dir) if args.output_dir else input_dir
|
||||
wav_dir = input_dir / "wavs"
|
||||
meta_path = input_dir / "metadata.txt"
|
||||
generator = DatasetGenerator(config_path=args.config)
|
||||
generator.generate()
|
||||
|
||||
if not meta_path.exists():
|
||||
if (input_dir / "metadata.csv").exists():
|
||||
meta_path = input_dir / "metadata.csv"
|
||||
else:
|
||||
raise FileNotFoundError(f"Could not find metadata.txt or metadata.csv at {input_dir}")
|
||||
|
||||
# 2. Load Encoder (Strict Float32)
|
||||
device = cfg.device
|
||||
print(f"Stage 1: Tokenizing dataset using Float32 precision...")
|
||||
print(f"Loading Encoder: {args.encoder_ckpt}")
|
||||
|
||||
encoder = Encoder(
|
||||
num_input_mels=cfg.codec['input_mels'],
|
||||
encoder_dim=cfg.codec['dim'],
|
||||
encoder_layers=cfg.codec['layers'],
|
||||
bottleneck_channels=cfg.codec['bottleneck']
|
||||
).to(device).float() # Explicitly ensure Float32
|
||||
|
||||
encoder.load_state_dict(torch.load(args.encoder_ckpt, map_location=device, weights_only=True))
|
||||
encoder.eval()
|
||||
|
||||
# 3. Read Metadata
|
||||
print(f"Reading metadata from {meta_path.name}...")
|
||||
samples = []
|
||||
with open(meta_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line: continue
|
||||
parts = line.split("|")
|
||||
|
||||
if len(parts) >= 2:
|
||||
file_id = parts[0]
|
||||
raw_text = parts[1]
|
||||
clean_text = normalize_text(raw_text)
|
||||
|
||||
# Support common Windows path naming conventions
|
||||
wav_path = wav_dir / f"{file_id}.wav"
|
||||
if not wav_path.exists():
|
||||
wav_path = wav_dir / f"{file_id}.WAV"
|
||||
if not wav_path.exists():
|
||||
continue
|
||||
|
||||
samples.append({
|
||||
"id": file_id,
|
||||
"text": clean_text,
|
||||
"wav_path": str(wav_path.absolute())
|
||||
})
|
||||
|
||||
print(f"Found {len(samples)} valid samples.")
|
||||
|
||||
if args.limit:
|
||||
print(f"[Test Mode] Limiting to first {args.limit} samples.")
|
||||
samples = samples[:args.limit]
|
||||
|
||||
# 4. Process Audio -> Tokens
|
||||
dataset = []
|
||||
target_sr = cfg.common['sample_rate']
|
||||
|
||||
with torch.no_grad():
|
||||
for sample in tqdm(samples, desc="Encoding Audio"):
|
||||
try:
|
||||
# Load as Float32
|
||||
wav_np, sr = sf.read(sample["wav_path"])
|
||||
wav = torch.from_numpy(wav_np).float()
|
||||
|
||||
if wav.ndim == 1: wav = wav.unsqueeze(0)
|
||||
else: wav = wav.t()
|
||||
|
||||
# Resample if needed
|
||||
if sr != target_sr:
|
||||
wav = torchaudio.functional.resample(wav, sr, target_sr)
|
||||
|
||||
# Force Mono
|
||||
if wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
|
||||
wav = wav.to(device)
|
||||
|
||||
# Get integer indices for the LLM to learn
|
||||
tokens = encoder(wav, return_indices=True)
|
||||
token_list = tokens.squeeze().cpu().tolist()
|
||||
|
||||
# Ensure it's a list even for single-token results
|
||||
if isinstance(token_list, int):
|
||||
token_list = [token_list]
|
||||
|
||||
dataset.append({
|
||||
"text": sample["text"],
|
||||
"tokens": token_list,
|
||||
"wav_path": sample["wav_path"]
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {sample['id']}: {e}")
|
||||
|
||||
# 5. Split and Save
|
||||
random.seed(cfg.common['seed'])
|
||||
random.shuffle(dataset)
|
||||
|
||||
val_split = cfg.dataset['val_split']
|
||||
split_idx = int(len(dataset) * (1 - val_split))
|
||||
train_data = dataset[:split_idx]
|
||||
val_data = dataset[split_idx:]
|
||||
|
||||
print(f"Saving {len(train_data)} train and {len(val_data)} validation samples...")
|
||||
|
||||
with open(output_dir / "train.json", "w", encoding="utf-8") as f:
|
||||
json.dump(train_data, f, indent=2)
|
||||
|
||||
with open(output_dir / "val.json", "w", encoding="utf-8") as f:
|
||||
json.dump(val_data, f, indent=2)
|
||||
|
||||
print("Stage 1 Complete. Tokens are ready for Stage 2.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,190 +0,0 @@
|
||||
# gui.pyw - Soprano Reforged V3 Training Factory
|
||||
import tkinter as tk
|
||||
from tkinter import ttk, filedialog, messagebox
|
||||
import sv_ttk
|
||||
import subprocess
|
||||
import threading
|
||||
import sys
|
||||
import queue
|
||||
import os
|
||||
|
||||
class SopranoGUI(tk.Tk):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self.title("Soprano Reforged V3 - Training Factory")
|
||||
self.geometry("900x700")
|
||||
|
||||
# Queue for thread-safe GUI updates
|
||||
self.log_queue = queue.Queue()
|
||||
self.is_running = False
|
||||
|
||||
# --- Layout ---
|
||||
self.create_widgets()
|
||||
|
||||
# --- Theme ---
|
||||
sv_ttk.set_theme("dark")
|
||||
self.refresh_display()
|
||||
|
||||
# --- Log Loop ---
|
||||
self.after(100, self.process_logs)
|
||||
|
||||
def refresh_display(self):
|
||||
self.update_idletasks()
|
||||
self.wm_attributes("-alpha", 0.99)
|
||||
self.wm_attributes("-alpha", 1.0)
|
||||
|
||||
def create_widgets(self):
|
||||
main_frame = ttk.Frame(self)
|
||||
main_frame.pack(fill="both", expand=True, padx=10, pady=10)
|
||||
|
||||
# 1. Tabs for Stages
|
||||
self.notebook = ttk.Notebook(main_frame)
|
||||
self.notebook.pack(fill="both", expand=True, pady=(0, 10))
|
||||
|
||||
self.tab_codec = ttk.Frame(self.notebook, padding=20)
|
||||
self.tab_dataset = ttk.Frame(self.notebook, padding=20)
|
||||
self.tab_train = ttk.Frame(self.notebook, padding=20)
|
||||
self.tab_infer = ttk.Frame(self.notebook, padding=20)
|
||||
|
||||
self.notebook.add(self.tab_codec, text="Stage 0: Codec")
|
||||
self.notebook.add(self.tab_dataset, text="Stage 1: Dataset")
|
||||
self.notebook.add(self.tab_train, text="Stage 2: Training")
|
||||
self.notebook.add(self.tab_infer, text="Inference")
|
||||
|
||||
self.build_codec_tab()
|
||||
self.build_dataset_tab()
|
||||
self.build_train_tab()
|
||||
self.build_infer_tab()
|
||||
|
||||
# 2. Console Output
|
||||
console_frame = ttk.LabelFrame(main_frame, text="V3 System Output", padding=10)
|
||||
console_frame.pack(fill="both", expand=True)
|
||||
|
||||
self.console = tk.Text(console_frame, height=12, state="disabled", bg="#1c1c1c", fg="#f0f0f0", font=("Consolas", 10))
|
||||
self.console.pack(fill="both", expand=True, side="left")
|
||||
|
||||
scrollbar = ttk.Scrollbar(console_frame, command=self.console.yview)
|
||||
scrollbar.pack(side="right", fill="y")
|
||||
self.console.config(yscrollcommand=scrollbar.set)
|
||||
|
||||
# --- Tab Builders ---
|
||||
|
||||
def build_codec_tab(self):
|
||||
f = self.tab_codec
|
||||
ttk.Label(f, text="Stage 0: Audio Codec (Float32)", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
|
||||
ttk.Label(f, text="WAV Directory (e.g. ./mio_dataset/wavs/*.wav)").pack(anchor="w")
|
||||
self.codec_wav_entry = ttk.Entry(f)
|
||||
self.codec_wav_entry.pack(fill="x", pady=(5, 10))
|
||||
self.codec_wav_entry.insert(0, "./mio_dataset/wavs/*.wav")
|
||||
|
||||
btn_frame = ttk.Frame(f)
|
||||
btn_frame.pack(fill="x", pady=10)
|
||||
ttk.Button(btn_frame, text="Browse Folder", command=lambda: self.browse_folder(self.codec_wav_entry, suffix="/*.wav")).pack(side="left", padx=(0, 10))
|
||||
ttk.Button(btn_frame, text="Start Codec Training", style="Accent.TButton",
|
||||
command=lambda: self.run_script("train_codec.py", ["--wav-dir", self.codec_wav_entry.get()])).pack(side="left")
|
||||
|
||||
def build_dataset_tab(self):
|
||||
f = self.tab_dataset
|
||||
ttk.Label(f, text="Stage 1: Token Generation", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
|
||||
ttk.Label(f, text="Dataset Input Directory").pack(anchor="w")
|
||||
self.ds_input_entry = ttk.Entry(f)
|
||||
self.ds_input_entry.pack(fill="x", pady=(5, 10))
|
||||
self.ds_input_entry.insert(0, "./mio_dataset")
|
||||
|
||||
ttk.Label(f, text="Encoder Checkpoint").pack(anchor="w")
|
||||
self.ds_encoder_entry = ttk.Entry(f)
|
||||
self.ds_encoder_entry.pack(fill="x", pady=(5, 10))
|
||||
self.ds_encoder_entry.insert(0, "./weights/codec/encoder.pth")
|
||||
|
||||
ttk.Button(f, text="Generate Dataset", style="Accent.TButton",
|
||||
command=lambda: self.run_script("generate_dataset.py", [
|
||||
"--input-dir", self.ds_input_entry.get(),
|
||||
"--encoder-ckpt", self.ds_encoder_entry.get()
|
||||
])).pack(anchor="w", pady=10)
|
||||
|
||||
def build_train_tab(self):
|
||||
f = self.tab_train
|
||||
ttk.Label(f, text="Stage 2: Joint V3 Training", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
|
||||
ttk.Label(f, text="Dataset Directory (containing train.json)").pack(anchor="w")
|
||||
self.train_input_entry = ttk.Entry(f)
|
||||
self.train_input_entry.pack(fill="x", pady=(5, 10))
|
||||
self.train_input_entry.insert(0, "./mio_dataset")
|
||||
|
||||
ttk.Button(f, text="Start Speed Run", style="Accent.TButton",
|
||||
command=lambda: self.run_script("train.py", ["--input-dir", self.train_input_entry.get()])).pack(anchor="w")
|
||||
|
||||
def build_infer_tab(self):
|
||||
f = self.tab_infer
|
||||
ttk.Label(f, text="Inference: Generate Speech", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
|
||||
ttk.Label(f, text="Enter Text:").pack(anchor="w")
|
||||
self.infer_text = ttk.Entry(f)
|
||||
self.infer_text.pack(fill="x", pady=(5, 10))
|
||||
self.infer_text.insert(0, "Soprano V3 is now fully operational.")
|
||||
|
||||
ttk.Label(f, text="Model Directory (e.g. weights/model/epoch_9)").pack(anchor="w")
|
||||
self.infer_model_entry = ttk.Entry(f)
|
||||
self.infer_model_entry.pack(fill="x", pady=(5, 10))
|
||||
self.infer_model_entry.insert(0, "./weights/model/epoch_9")
|
||||
|
||||
ttk.Button(f, text="Generate .WAV", style="Accent.TButton",
|
||||
command=lambda: self.run_script("inference.py", [
|
||||
"--text", self.infer_text.get(),
|
||||
"--model-dir", self.infer_model_entry.get(),
|
||||
"--output", "v3_output.wav"
|
||||
])).pack(anchor="w", pady=10)
|
||||
|
||||
# --- Subprocess Logic ---
|
||||
|
||||
def browse_folder(self, entry_widget, suffix=""):
|
||||
path = filedialog.askdirectory()
|
||||
if path:
|
||||
entry_widget.delete(0, tk.END)
|
||||
entry_widget.insert(0, path + suffix)
|
||||
|
||||
def log(self, message):
|
||||
self.log_queue.put(message)
|
||||
|
||||
def process_logs(self):
|
||||
while not self.log_queue.empty():
|
||||
msg = self.log_queue.get()
|
||||
self.console.config(state="normal")
|
||||
self.console.insert(tk.END, msg)
|
||||
self.console.see(tk.END)
|
||||
self.console.config(state="disabled")
|
||||
self.after(100, self.process_logs)
|
||||
|
||||
def run_script(self, script_name, args):
|
||||
if self.is_running:
|
||||
messagebox.showwarning("Busy", "Process already running.")
|
||||
return
|
||||
|
||||
self.is_running = True
|
||||
self.console.config(state="normal")
|
||||
self.console.delete(1.0, tk.END)
|
||||
self.console.config(state="disabled")
|
||||
|
||||
cmd = [sys.executable, script_name] + args
|
||||
self.log(f"RUNNING: {' '.join(cmd)}\n" + "-"*50 + "\n")
|
||||
|
||||
thread = threading.Thread(target=self._execute_subprocess, args=(cmd,))
|
||||
thread.start()
|
||||
|
||||
def _execute_subprocess(self, cmd):
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||||
text=True, bufsize=1, universal_newlines=True
|
||||
)
|
||||
for line in process.stdout:
|
||||
self.log(line)
|
||||
process.wait()
|
||||
self.log(f"\nProcess Exit Code: {process.returncode}\n")
|
||||
except Exception as e:
|
||||
self.log(f"\nError: {e}\n")
|
||||
finally:
|
||||
self.is_running = False
|
||||
|
||||
if __name__ == "__main__":
|
||||
app = SopranoGUI()
|
||||
app.mainloop()
|
||||
-132
@@ -1,132 +0,0 @@
|
||||
# inference.py - Soprano Reforged V3 (Full Precision Inference)
|
||||
# EMOJI-FREE VERSION for Windows Compatibility
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import torch
|
||||
import soundfile as sf
|
||||
import numpy as np
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from model.decoder import Decoder
|
||||
from utils.text_normalizer import normalize_text
|
||||
from utils.config import cfg
|
||||
|
||||
def generate(text, llm, tokenizer, decoder, output_path="output.wav", top_k=50, temperature=0.7, device="cuda"):
|
||||
"""
|
||||
Core generation function with integrated Pulse-Check and Audio Normalization.
|
||||
"""
|
||||
print(f"--- Inference Start ---")
|
||||
print(f"Input Text: '{text}'")
|
||||
|
||||
# 1. Text Normalization
|
||||
norm_text = normalize_text(text)
|
||||
prompt = f"[STOP][TEXT]{norm_text}[START]"
|
||||
|
||||
# 2. Tokenization
|
||||
inputs = tokenizer(prompt, return_tensors="pt").to(device)
|
||||
input_ids = inputs["input_ids"]
|
||||
|
||||
# 3. LLM Generation
|
||||
print("Running LLM (Brain)...")
|
||||
with torch.no_grad():
|
||||
output_ids = llm.generate(
|
||||
input_ids,
|
||||
max_new_tokens=cfg.training.get('seq_len', 512),
|
||||
do_sample=True,
|
||||
top_k=top_k,
|
||||
temperature=temperature,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
|
||||
# 4. Token Slicing (Find the [START] marker)
|
||||
start_token_id = tokenizer.convert_tokens_to_ids("[START]")
|
||||
out_list = output_ids[0].tolist()
|
||||
|
||||
try:
|
||||
start_idx = out_list.index(start_token_id)
|
||||
# Slicing after [START] to get the audio tokens/hidden states
|
||||
audio_token_ids = output_ids[:, start_idx+1:]
|
||||
except ValueError:
|
||||
print("Warning: [START] token not found in output. Using full sequence.")
|
||||
start_idx = 0
|
||||
audio_token_ids = output_ids
|
||||
|
||||
if audio_token_ids.size(1) == 0:
|
||||
print("Error: Model generated no audio tokens.")
|
||||
return
|
||||
|
||||
# 5. Decoding (Mouth)
|
||||
print(f"Decoding {audio_token_ids.size(1)} tokens...")
|
||||
with torch.no_grad():
|
||||
# Get hidden states for the specific tokens generated
|
||||
llm_out = llm(output_ids, output_hidden_states=True)
|
||||
# We target the last layer's hidden states for the audio portion
|
||||
hidden_states = llm_out.hidden_states[-1][:, start_idx+1:, :]
|
||||
|
||||
# Pulse Check 1: LLM Strength
|
||||
llm_strength = hidden_states.abs().mean().item()
|
||||
print(f"LLM Output Strength (Mean Abs): {llm_strength:.6f}")
|
||||
|
||||
waveform = decoder(hidden_states)
|
||||
|
||||
# Pulse Check 2: Decoder Strength
|
||||
dec_strength = waveform.abs().mean().item()
|
||||
print(f"Decoder Output Strength (Mean Abs): {dec_strength:.6f}")
|
||||
|
||||
# 6. Post-Processing & Audio Boosting
|
||||
audio_np = waveform.squeeze().cpu().numpy()
|
||||
|
||||
max_amp = np.max(np.abs(audio_np))
|
||||
if max_amp < 0.01:
|
||||
print(f"Warning: Output is extremely weak (Max Amp: {max_amp:.6f}). This likely means the model hasn't learned yet.")
|
||||
print("Skipping normalization to avoid boosting static/noise.")
|
||||
else:
|
||||
# Boost to peak 0.9 for clarity without clipping
|
||||
audio_np = audio_np / max_amp * 0.9
|
||||
print(f"Normalization: Audio boosted by {0.9/max_amp:.2f}x (Max Amp: {max_amp:.4f})")
|
||||
|
||||
# 7. Save to File
|
||||
sf.write(output_path, audio_np, cfg.common['sample_rate'])
|
||||
print(f"--- Success: File saved to {output_path} ---")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Soprano Reforged V3 Inference")
|
||||
parser.add_argument("--text", type=str, required=True, help="Text to synthesize")
|
||||
parser.add_argument("--model-dir", type=str, required=True, help="Path to epoch folder")
|
||||
parser.add_argument("--output", type=str, default="inference_output.wav", help="Output filename")
|
||||
args = parser.parse_args()
|
||||
|
||||
device = cfg.device
|
||||
print(f"Loading Model from {args.model_dir} on {device}...")
|
||||
|
||||
# Load Tokenizer & LLM
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
|
||||
llm = AutoModelForCausalLM.from_pretrained(
|
||||
args.model_dir,
|
||||
torch_dtype=torch.float32,
|
||||
low_cpu_mem_usage=True
|
||||
).to(device)
|
||||
llm.eval()
|
||||
|
||||
# Load Decoder (Internalized Adapter)
|
||||
decoder_path = os.path.join(args.model_dir, "decoder.pth")
|
||||
if not os.path.exists(decoder_path):
|
||||
# Fallback to general weights directory if not in epoch folder
|
||||
decoder_path = os.path.join(cfg.training['save_dir'], "decoder.pth")
|
||||
|
||||
decoder = Decoder(
|
||||
input_channels=llm.config.hidden_size,
|
||||
decoder_dim=cfg.codec['dim'],
|
||||
decoder_layers=cfg.codec['layers']
|
||||
).to(device).float()
|
||||
|
||||
print(f"Loading Decoder weights from {decoder_path}...")
|
||||
decoder.load_state_dict(torch.load(decoder_path, map_location=device, weights_only=True))
|
||||
decoder.eval()
|
||||
|
||||
# Run Generation
|
||||
generate(args.text, llm, tokenizer, decoder, output_path=args.output, device=device)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,5 +0,0 @@
|
||||
@echo off
|
||||
cd /d "%~dp0"
|
||||
call .venv\Scripts\activate
|
||||
start "" pythonw gui.pyw
|
||||
exit
|
||||
-113
@@ -1,113 +0,0 @@
|
||||
# model/common.py
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
class SimpleMLP(nn.Module):
|
||||
"""
|
||||
Simple Multi-Layer Perceptron used within ConvNeXt blocks.
|
||||
Standardized for Float32 precision.
|
||||
"""
|
||||
def __init__(self, dim, intermediate_dim):
|
||||
super().__init__()
|
||||
self.pwconv1 = nn.Linear(dim, intermediate_dim)
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(intermediate_dim, dim)
|
||||
|
||||
def forward(self, x):
|
||||
return self.pwconv2(self.act(self.pwconv1(x)))
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
"""
|
||||
ConvNeXt Block adapted for 1D audio.
|
||||
Used by both Encoder and Decoder backbones.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
layer_scale_init_value: float = 1e-6,
|
||||
kernel_size: int = 7,
|
||||
):
|
||||
super().__init__()
|
||||
# Depthwise conv: Process spatial (time) info
|
||||
self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size//2, groups=dim)
|
||||
|
||||
# LayerNorm: Essential for stability in Float32
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
|
||||
# Pointwise MLP: Process channel info
|
||||
self.mlp = SimpleMLP(dim, intermediate_dim)
|
||||
|
||||
# Layer Scale: Helps with deep network convergence
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.dwconv(x)
|
||||
x = x.transpose(1, 2) # [B, C, T] -> [B, T, C] for Norm/MLP
|
||||
x = self.norm(x)
|
||||
x = self.mlp(x)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
x = x.transpose(1, 2) # [B, T, C] -> [B, C, T]
|
||||
x = residual + x
|
||||
return x
|
||||
|
||||
class VocosBackbone(nn.Module):
|
||||
"""
|
||||
Shared backbone architecture for Encoder/Decoder.
|
||||
Uses ConvNeXt blocks to extract/reconstruct high-fidelity features.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
num_layers: int,
|
||||
kernel_size: int = 7,
|
||||
layer_scale_init_value: float = None,
|
||||
):
|
||||
super().__init__()
|
||||
# Initial projection into hidden dimension
|
||||
self.embed = nn.Conv1d(
|
||||
input_channels, dim, kernel_size=kernel_size, padding=kernel_size//2
|
||||
)
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
|
||||
self.layers = nn.ModuleList([
|
||||
ConvNeXtBlock(
|
||||
dim=dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
kernel_size=kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value or 1 / num_layers**0.5,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
|
||||
self.final_norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.apply(self._init_weights)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# Initial Embedding
|
||||
x = self.embed(x)
|
||||
# Apply LayerNorm (requires transpose)
|
||||
x = self.norm(x.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
# Processing Blocks
|
||||
for layer in self.layers:
|
||||
x = layer(x)
|
||||
|
||||
# Final Norm
|
||||
x = self.final_norm(x.transpose(1, 2)).transpose(1, 2)
|
||||
return x
|
||||
@@ -1,117 +0,0 @@
|
||||
# model/decoder.py - Standardized for V3
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from .common import VocosBackbone
|
||||
|
||||
class ISTFTHead(nn.Module):
|
||||
"""
|
||||
Inverse STFT head for waveform generation.
|
||||
Converts hidden states -> complex spectrogram -> waveform.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
n_fft: int = 2048,
|
||||
hop_length: int = 512,
|
||||
padding: str = "same",
|
||||
):
|
||||
super().__init__()
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
out_dim = n_fft + 2
|
||||
self.out = nn.Linear(dim, out_dim)
|
||||
self.padding = padding
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [B, T, C] - hidden states
|
||||
Returns:
|
||||
audio: [B, 1, L] - waveform
|
||||
"""
|
||||
x = x.float() # Ensure float32
|
||||
x = self.out(x) # [B, T, n_fft + 2]
|
||||
x = x.transpose(1, 2) # [B, n_fft + 2, T]
|
||||
|
||||
mag, phase = x.chunk(2, dim=1) # each [B, (n_fft+2)//2, T]
|
||||
mag = torch.exp(mag)
|
||||
phase = torch.sin(phase)
|
||||
|
||||
# Construct complex spectrogram
|
||||
S = mag * torch.exp(1j * phase)
|
||||
|
||||
# Inverse STFT
|
||||
audio = torch.istft(
|
||||
S,
|
||||
n_fft=self.n_fft,
|
||||
hop_length=self.hop_length,
|
||||
window=torch.hann_window(self.n_fft).to(x.device),
|
||||
center=True,
|
||||
)
|
||||
|
||||
return audio.unsqueeze(1) # [B, 1, L]
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channels=512,
|
||||
decoder_dim=512,
|
||||
decoder_layers=8,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# 1. The Adapter (Projection)
|
||||
self.input_proj = nn.Linear(input_channels, decoder_dim)
|
||||
|
||||
# Standardization Layer
|
||||
self.post_proj_norm = nn.LayerNorm(decoder_dim)
|
||||
|
||||
# 2. Upsampler (Match Encoder's slice::4)
|
||||
# We need to upsample by 4x to match the Hop=512 resolution
|
||||
# Encoder: Hop=512 -> Slice(4) -> Effective Hop=2048
|
||||
# Decoder: Latent(2048) -> Upsample(4) -> Latent(512) -> ISTFT(512) -> Audio
|
||||
self.upsampler = nn.Sequential(
|
||||
nn.ConvTranspose1d(decoder_dim, decoder_dim, kernel_size=4, stride=4),
|
||||
nn.GELU()
|
||||
)
|
||||
|
||||
# 3. Backbone
|
||||
self.backbone = VocosBackbone(
|
||||
input_channels=decoder_dim,
|
||||
dim=decoder_dim,
|
||||
intermediate_dim=decoder_dim * 3,
|
||||
num_layers=decoder_layers,
|
||||
)
|
||||
|
||||
# 4. Output Head (ISTFT)
|
||||
self.head = ISTFTHead(
|
||||
dim=decoder_dim,
|
||||
n_fft=2048,
|
||||
hop_length=512
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: [Batch, Time, Channels]
|
||||
x = x.float() # Ensure float32 everywhere
|
||||
|
||||
# Project and Normalize
|
||||
x = self.input_proj(x)
|
||||
x = self.post_proj_norm(x)
|
||||
|
||||
# Upsample [B, T, C] -> [B, C, T] -> [B, C, T*4]
|
||||
x = x.transpose(1, 2)
|
||||
x = self.upsampler(x)
|
||||
|
||||
# Backbone ResBlocks [B, C, T*4] -> [B, C, T*4]
|
||||
# (VocosBackbone expects [B, C, T] and returns [B, C, T])
|
||||
x = self.backbone(x)
|
||||
|
||||
# Prepare for ISTFT Head [B, C, T] -> [B, T, C]
|
||||
x = x.transpose(1, 2)
|
||||
|
||||
# Generate Waveform
|
||||
x = self.head(x)
|
||||
|
||||
return x
|
||||
@@ -1,79 +0,0 @@
|
||||
# model/encoder.py - Standardized for V3
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torchaudio
|
||||
from .common import VocosBackbone
|
||||
from .quantizer import FSQSTE
|
||||
|
||||
def safe_log(x: torch.Tensor, clip_val: float = 1e-5) -> torch.Tensor:
|
||||
"""Prevents log(0) errors. Explicitly Float32."""
|
||||
return torch.log(torch.clip(x, min=clip_val))
|
||||
|
||||
class Encoder(nn.Module):
|
||||
"""
|
||||
Encodes audio into discrete quantized tokens.
|
||||
Standardized for Float32 precision and Windows compatibility.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
num_input_mels=80,
|
||||
encoder_dim=512,
|
||||
encoder_layers=8,
|
||||
bottleneck_channels=5,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# 1. Mel Spectrogram Transform
|
||||
self.mel_spec = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=32000,
|
||||
n_fft=2048,
|
||||
hop_length=512,
|
||||
n_mels=num_input_mels,
|
||||
normalized=True
|
||||
)
|
||||
|
||||
# 2. Backbone Network
|
||||
self.backbone = VocosBackbone(
|
||||
input_channels=num_input_mels,
|
||||
dim=encoder_dim,
|
||||
intermediate_dim=encoder_dim * 3,
|
||||
num_layers=encoder_layers,
|
||||
)
|
||||
|
||||
# 3. Downsampling Projector
|
||||
self.downsample_factor = 4
|
||||
self.down_proj = nn.Linear(encoder_dim, bottleneck_channels)
|
||||
|
||||
# 4. Quantizer
|
||||
self.quantizer = FSQSTE(levels=[8, 8, 5, 5, 5])
|
||||
|
||||
def preprocess(self, audio):
|
||||
"""Converts raw audio to log-mel spectrogram (Float32)."""
|
||||
if audio.dim() == 2:
|
||||
audio = audio.unsqueeze(1) # Ensure [B, 1, T]
|
||||
mel = self.mel_spec(audio.float()) # Force Float32 math
|
||||
return safe_log(mel).squeeze(1)
|
||||
|
||||
def forward(self, audio, return_indices=False):
|
||||
"""
|
||||
Audio -> Spectrogram -> Features -> Quantized Codes
|
||||
"""
|
||||
# 1. Mel Spec
|
||||
x = self.preprocess(audio)
|
||||
|
||||
# 2. Backbone
|
||||
x = self.backbone(x) # [B, Dim, Frames]
|
||||
|
||||
# 3. Downsample (Slicing)
|
||||
x = x[:, :, ::self.downsample_factor]
|
||||
|
||||
# 4. Project to Bottleneck
|
||||
x = x.transpose(1, 2) # [B, Frames//4, Dim]
|
||||
x = self.down_proj(x) # [B, Frames//4, 5]
|
||||
|
||||
# 5. Quantize
|
||||
if return_indices:
|
||||
return self.quantizer.to_codebook_index(x)
|
||||
|
||||
return self.quantizer(x)
|
||||
@@ -1,59 +0,0 @@
|
||||
# model/quantizer.py
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
class FSQSTE(nn.Module):
|
||||
"""
|
||||
Finite Scalar Quantization with Straight-Through Estimator.
|
||||
Unicode-safe and explicitly optimized for Float32 precision.
|
||||
"""
|
||||
def __init__(self, levels=[8, 8, 5, 5, 5]):
|
||||
super().__init__()
|
||||
# Convert levels to float32 for calculation stability
|
||||
levels_tensor = torch.tensor(levels, dtype=torch.float32)
|
||||
self.dim = len(levels)
|
||||
|
||||
self.register_buffer("_levels", levels_tensor)
|
||||
|
||||
# Calculate quantization parameters in full precision
|
||||
self.register_buffer("half_levels", (self._levels - 1) * (1 - 1e-3) / 2)
|
||||
self.register_buffer("offset", 0.5 - 0.5 * (self._levels % 2))
|
||||
self.register_buffer("shift", torch.tan(self.offset / self.half_levels))
|
||||
|
||||
# Basis for converting indices <-> vectors
|
||||
basis = torch.cumprod(torch.tensor([1] + levels[:-1], dtype=torch.float32), dim=0).to(torch.int32)
|
||||
self.register_buffer("_basis", basis)
|
||||
|
||||
def _scale_and_shift(self, zhat_normalized):
|
||||
half_width = self._levels // 2
|
||||
return (zhat_normalized * half_width) + half_width
|
||||
|
||||
def to_codebook_index(self, zhat):
|
||||
"""Converts continuous vectors to integer indices."""
|
||||
# Ensure input is float32
|
||||
zhat = zhat.float()
|
||||
zhat = self._scale_and_shift(zhat)
|
||||
indices = (zhat * self._basis).sum(dim=-1).round().to(torch.int32)
|
||||
return indices
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Quantizes input `x` using the implicit codebook via STE.
|
||||
Forces all math to Float32.
|
||||
"""
|
||||
# 1. Ensure float32 for mathematical stability
|
||||
x = x.float()
|
||||
|
||||
# 2. Quantize with tanh trick
|
||||
# Maps input to valid discrete levels
|
||||
x = torch.tanh(x + self.shift) * self.half_levels - self.offset
|
||||
|
||||
# 3. Straight-Through Estimator:
|
||||
# Forward pass uses rounded values; backward pass uses un-rounded gradients.
|
||||
x = x + (x.round() - x).detach()
|
||||
|
||||
# 4. Normalize to [-1, 1] for backbone stability
|
||||
x = x / (self._levels // 2)
|
||||
|
||||
return x
|
||||
-114
@@ -1,114 +0,0 @@
|
||||
# monitor.py - Soprano V3 System Observer
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import subprocess
|
||||
import requests
|
||||
import glob
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
# =================CONFIGURATION=================
|
||||
NTFY_TOPIC = "changeme"
|
||||
TEST_PROMPT = "This is Epoch {}. Checking systems. The vocal reconstruction is proceeding as planned."
|
||||
MODEL_DIR = "weights/model"
|
||||
INFERENCE_SCRIPT = "inference.py"
|
||||
LOG_FILE = "training.log"
|
||||
# ===============================================
|
||||
|
||||
print(f"Monitoring {MODEL_DIR} for new epochs...")
|
||||
print(f"Notifications: https://ntfy.sh/{NTFY_TOPIC}")
|
||||
|
||||
processed_epochs = set()
|
||||
|
||||
def get_latest_loss():
|
||||
try:
|
||||
if not os.path.exists(LOG_FILE): return "N/A"
|
||||
with open(LOG_FILE, 'r', encoding='utf-8') as f:
|
||||
lines = f.readlines()
|
||||
for line in reversed(lines):
|
||||
if "Ep" in line and "|" in line:
|
||||
return line.strip().split("|")[-1].strip()
|
||||
except: return "Unknown"
|
||||
return "N/A"
|
||||
|
||||
def send_notification(epoch_num, wav_path):
|
||||
loss_stat = get_latest_loss()
|
||||
print(f"Sending notification for Epoch {epoch_num}...")
|
||||
try:
|
||||
# 1. Send Text Update
|
||||
requests.post(f"https://ntfy.sh/{NTFY_TOPIC}",
|
||||
data=f"Epoch {epoch_num} Complete!\nLoss: {loss_stat}",
|
||||
headers={"Title": f"Epoch {epoch_num} Done", "Tags": "tada"})
|
||||
|
||||
# 2. Upload Audio File (Proper Header for local files)
|
||||
if os.path.exists(wav_path):
|
||||
with open(wav_path, 'rb') as f:
|
||||
requests.put(f"https://ntfy.sh/{NTFY_TOPIC}", data=f,
|
||||
headers={
|
||||
"Filename": f"ep{epoch_num}.wav",
|
||||
"Title": f"Audio Sample: Epoch {epoch_num}"
|
||||
})
|
||||
print("Notification Sent Successfully!")
|
||||
except Exception as e:
|
||||
print(f"Network error: {e}")
|
||||
|
||||
def run_inference(epoch_folder, epoch_num):
|
||||
output_wav = f"monitor_epoch_{epoch_num}.wav"
|
||||
prompt = TEST_PROMPT.format(epoch_num)
|
||||
|
||||
# Use the current python executable
|
||||
python_exe = sys.executable
|
||||
|
||||
# Command matches the new inference.py arguments
|
||||
cmd = [
|
||||
python_exe, INFERENCE_SCRIPT,
|
||||
"--model-dir", epoch_folder,
|
||||
"--text", prompt,
|
||||
"--output", output_wav
|
||||
]
|
||||
|
||||
# OPTIONAL: Force CPU for monitor inference to save 5090 VRAM for training
|
||||
# env = os.environ.copy()
|
||||
# env["CUDA_VISIBLE_DEVICES"] = ""
|
||||
|
||||
print(f"Generating sample for Epoch {epoch_num}...")
|
||||
try:
|
||||
# Run and capture output to prevent terminal clutter
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True
|
||||
)
|
||||
return output_wav
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Inference Failed for Epoch {epoch_num}")
|
||||
print(f"Error details: {e.stderr}")
|
||||
return None
|
||||
|
||||
# Main Loop
|
||||
while True:
|
||||
# Look for epoch folders (e.g., weights/model/epoch_0)
|
||||
epoch_dirs = glob.glob(os.path.join(MODEL_DIR, "epoch_*"))
|
||||
for d in epoch_dirs:
|
||||
try:
|
||||
folder_name = os.path.basename(d)
|
||||
epoch_num = int(folder_name.split("_")[1])
|
||||
|
||||
if epoch_num not in processed_epochs:
|
||||
# Wait 10 seconds to ensure the OS has finished writing the .pth files
|
||||
time.sleep(10)
|
||||
|
||||
print(f"\nProcessing Epoch {epoch_num}")
|
||||
wav_file = run_inference(d, epoch_num)
|
||||
|
||||
if wav_file:
|
||||
send_notification(epoch_num, wav_file)
|
||||
|
||||
processed_epochs.add(epoch_num)
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
time.sleep(60)
|
||||
+43
-38
@@ -1,48 +1,53 @@
|
||||
[project]
|
||||
name = "soprano-factory"
|
||||
version = "0.1.0"
|
||||
description = "Soprano-Reforged: Clean training environment"
|
||||
readme = "README.md"
|
||||
# We explicitly pin to 3.12 to match stable Torch/Audio wheels
|
||||
requires-python = ">=3.12, <3.13"
|
||||
dependencies = [
|
||||
# Core Data & ML
|
||||
"numpy>=1.26.0",
|
||||
"scipy",
|
||||
"tqdm",
|
||||
"einops",
|
||||
"huggingface_hub",
|
||||
"transformers",
|
||||
"accelerate",
|
||||
# Audio Processing (Stable Windows Stack)
|
||||
"librosa", # General audio manipulation
|
||||
"soundfile", # The stable backend for Torchaudio on Windows
|
||||
"ffmpeg-python", # Fallback/Helper (requires ffmpeg.exe in PATH)
|
||||
# Torch Ecosystem
|
||||
"torch>=2.6.0",
|
||||
"torchaudio>=2.6.0",
|
||||
"sv-ttk>=2.6.1",
|
||||
"darkdetect>=0.8.0",
|
||||
"pywinstyles>=1.8",
|
||||
"requests>=2.32.5",
|
||||
description = "A clean-room training pipeline for fine-tuning Soprano TTS models."
|
||||
authors = [
|
||||
{ name = "Your Name", email = "your.email@example.com" }
|
||||
]
|
||||
requires-python = ">=3.12"
|
||||
readme = "README.md"
|
||||
license = { text = "Apache-2.0" }
|
||||
|
||||
[tool.uv]
|
||||
# Force Windows to grab CUDA-enabled PyTorch
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cuda"
|
||||
url = "https://download.pytorch.org/whl/cu126"
|
||||
explicit = true
|
||||
|
||||
[tool.uv.sources]
|
||||
torch = [{ index = "pytorch-cuda" }]
|
||||
torchaudio = [{ index = "pytorch-cuda" }]
|
||||
torchvision = [{ index = "pytorch-cuda" }]
|
||||
# Core dependencies required for data generation, training, and inference testing
|
||||
dependencies = [
|
||||
"einops>=0.8.2",
|
||||
"hf-xet>=1.2.0",
|
||||
"huggingface-hub>=0.20.0",
|
||||
"matplotlib>=3.10.8",
|
||||
"numpy>=1.26.0",
|
||||
"pyyaml",
|
||||
"safetensors>=0.7.0",
|
||||
"scipy>=1.11.0",
|
||||
"soprano-tts>=0.2.0",
|
||||
"soundfile>=0.13.1",
|
||||
"tensorboard",
|
||||
"torch>=2.8.0",
|
||||
"torchaudio>=2.8.0",
|
||||
"tqdm>=4.66.0",
|
||||
"transformers>=4.51.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
# We tell hatchling that the code is in the 'model' folder
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["model"]
|
||||
# ------------------------------------------------------------------------------
|
||||
# uv-Specific Configuration
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
# Define the PyTorch CUDA 12.8 index
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cu128"
|
||||
url = "https://download.pytorch.org/whl/cu128"
|
||||
explicit = true
|
||||
|
||||
# Route torch-related packages exclusively to the CUDA 12.8 index
|
||||
[tool.uv.sources]
|
||||
torch = { index = "pytorch-cu128" }
|
||||
torchaudio = { index = "pytorch-cu128" }
|
||||
|
||||
[tool.uv]
|
||||
managed = true
|
||||
# This tells uv we just want an environment for scripts, not a buildable package
|
||||
package = false
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# start_pipeline.py - Interactive CLI to run Soprano Factory's 2-stage TTS training pipeline with color-coded output and error handling.
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class Colors:
|
||||
HEADER = '\033[95m'
|
||||
OKBLUE = '\033[94m'
|
||||
OKCYAN = '\033[96m'
|
||||
OKGREEN = '\033[92m'
|
||||
WARNING = '\033[93m'
|
||||
FAIL = '\033[91m'
|
||||
ENDC = '\033[0m'
|
||||
BOLD = '\033[1m'
|
||||
|
||||
|
||||
def print_banner():
|
||||
banner = f"""{Colors.OKCYAN}{Colors.BOLD}
|
||||
===================================================================
|
||||
S O P R A N O F A C T O R Y
|
||||
===================================================================
|
||||
{Colors.ENDC}"""
|
||||
print(banner)
|
||||
|
||||
|
||||
def check_prerequisites():
|
||||
"""Checks if the required config and dataset folders exist."""
|
||||
print(f"{Colors.HEADER}Checking Pipeline Prerequisites...{Colors.ENDC}")
|
||||
|
||||
if not Path("config.yaml").exists():
|
||||
print(f"{Colors.FAIL}[!] config.yaml not found! Please create it before starting.{Colors.ENDC}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"{Colors.OKGREEN}[✓] config.yaml found.{Colors.ENDC}")
|
||||
# Note: We aren't strictly checking the dataset path here because it relies on reading the config,
|
||||
# but the individual scripts will catch missing data safely.
|
||||
|
||||
|
||||
def run_script(script_name: str, description: str):
|
||||
"""Runs a python script via uv and handles interruptions gracefully."""
|
||||
print(f"\n{Colors.OKBLUE}{Colors.BOLD}>>> {description}{Colors.ENDC}")
|
||||
print(f"{Colors.WARNING}Press Ctrl+C to stop this step and return to the main menu.{Colors.ENDC}\n")
|
||||
|
||||
try:
|
||||
# Using 'uv run' to ensure it executes in the correct environment
|
||||
subprocess.run(["uv", "run", script_name], check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"\n{Colors.FAIL}[!] Process '{script_name}' exited with error code {e.returncode}.{Colors.ENDC}")
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n{Colors.WARNING}[!] Process '{script_name}' interrupted by user. Returning to menu...{Colors.ENDC}")
|
||||
|
||||
|
||||
def start_tensorboard():
|
||||
"""Spins up TensorBoard in the background or a new shell."""
|
||||
print(f"\n{Colors.OKBLUE}>>> Starting TensorBoard...{Colors.ENDC}")
|
||||
print("TensorBoard will be available at: http://localhost:6006/")
|
||||
print(f"{Colors.WARNING}(Note: You may need to open this URL in your browser manually. Press Ctrl+C in this terminal to stop TensorBoard.){Colors.ENDC}\n")
|
||||
|
||||
try:
|
||||
subprocess.run(["uv", "run", "tensorboard", "--logdir=./runs"])
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n{Colors.OKGREEN}Closed TensorBoard.{Colors.ENDC}")
|
||||
|
||||
|
||||
def interactive_menu():
|
||||
while True:
|
||||
print_banner()
|
||||
print("Please select an action:")
|
||||
print(f" {Colors.BOLD}1.{Colors.ENDC} Generate Dataset (Tokenization)")
|
||||
print(f" {Colors.BOLD}2.{Colors.ENDC} Train Stage 1: LLM (Autoregressive Text-to-Tokens)")
|
||||
print(f" {Colors.BOLD}3.{Colors.ENDC} Train Stage 2: Decoder (GAN HD Audio Synthesis)")
|
||||
print(f" {Colors.BOLD}4.{Colors.ENDC} Run Full Automated Pipeline (Steps 1 -> 2 -> 3)")
|
||||
print(f" {Colors.BOLD}5.{Colors.ENDC} Start TensorBoard Monitoring")
|
||||
print(f" {Colors.BOLD}6.{Colors.ENDC} Test Fine-Tuned Model Inference")
|
||||
print(f" {Colors.BOLD}0.{Colors.ENDC} Exit")
|
||||
|
||||
choice = input(f"\n{Colors.BOLD}Enter your choice (0-6): {Colors.ENDC}").strip()
|
||||
|
||||
if choice == '1':
|
||||
run_script("generate_dataset.py", "Generating Dataset")
|
||||
elif choice == '2':
|
||||
run_script("train_llm.py", "Running Stage 1: LLM Training")
|
||||
elif choice == '3':
|
||||
run_script("train_decoder.py", "Running Stage 2: Decoder GAN Training")
|
||||
elif choice == '4':
|
||||
print(f"\n{Colors.HEADER}Starting Full Automated Pipeline...{Colors.ENDC}")
|
||||
run_script("generate_dataset.py", "Step 1/3: Generating Dataset")
|
||||
run_script("train_llm.py", "Step 2/3: Stage 1 LLM Training")
|
||||
run_script("train_decoder.py", "Step 3/3: Stage 2 Decoder GAN Training")
|
||||
print(f"\n{Colors.OKGREEN}{Colors.BOLD}Full Pipeline Execution Complete!{Colors.ENDC}")
|
||||
elif choice == '5':
|
||||
start_tensorboard()
|
||||
elif choice == '6':
|
||||
run_script("test_inference.py", "Testing Model Inference")
|
||||
elif choice == '0':
|
||||
print(f"{Colors.OKGREEN}Exiting Soprano Factory. Goodbye!{Colors.ENDC}")
|
||||
break
|
||||
else:
|
||||
print(f"{Colors.FAIL}Invalid choice. Please enter a number between 0 and 6.{Colors.ENDC}")
|
||||
|
||||
input(f"\n{Colors.BOLD}Press Enter to return to the main menu...{Colors.ENDC}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Clear terminal screen for a clean UI
|
||||
os.system('cls' if os.name == 'nt' else 'clear')
|
||||
check_prerequisites()
|
||||
interactive_menu()
|
||||
-130
@@ -1,130 +0,0 @@
|
||||
# train_codec.py - Stage 0: Codec Training (Standardized for V3 - Unicode Safe)
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchaudio
|
||||
import soundfile as sf
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torch.optim import AdamW
|
||||
from tqdm import tqdm
|
||||
|
||||
# Import our architecture
|
||||
from model.encoder import Encoder
|
||||
from model.decoder import Decoder
|
||||
from utils.config import cfg
|
||||
|
||||
# --- Windows / Audio Backend Setup ---
|
||||
def setup_audio_backend():
|
||||
cwd = os.getcwd()
|
||||
local_ffmpeg = os.path.join(cwd, "tools", "ffmpeg")
|
||||
if os.path.exists(local_ffmpeg) and os.path.isdir(local_ffmpeg):
|
||||
if local_ffmpeg not in os.environ["PATH"]:
|
||||
print(f"[Setup] Found local FFmpeg at: {local_ffmpeg}")
|
||||
os.environ["PATH"] = local_ffmpeg + os.pathsep + os.environ["PATH"]
|
||||
|
||||
setup_audio_backend()
|
||||
|
||||
class SpectralLoss(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.mel = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=cfg.common['sample_rate'],
|
||||
n_mels=cfg.codec['input_mels'],
|
||||
n_fft=2048, hop_length=512
|
||||
)
|
||||
|
||||
def forward(self, pred, target):
|
||||
pred, target = pred.float(), target.float()
|
||||
min_len = min(pred.shape[-1], target.shape[-1])
|
||||
pred = pred[..., :min_len]
|
||||
target = target[..., :min_len]
|
||||
|
||||
loss_time = F.l1_loss(pred, target)
|
||||
|
||||
if self.mel.mel_scale.fb.device != pred.device:
|
||||
self.mel = self.mel.to(pred.device)
|
||||
|
||||
loss_mel = F.l1_loss(self.mel(pred), self.mel(target))
|
||||
return loss_time + loss_mel
|
||||
|
||||
class WavDataset(Dataset):
|
||||
def __init__(self, glob_pattern, segment_size):
|
||||
self.files = glob.glob(glob_pattern, recursive=True)
|
||||
self.segment_size = segment_size
|
||||
print(f"Stage 0: Found {len(self.files)} wav files for training.")
|
||||
|
||||
def __len__(self):
|
||||
return len(self.files)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
try:
|
||||
wav_np, sr = sf.read(self.files[idx])
|
||||
wav = torch.from_numpy(wav_np).float()
|
||||
if wav.ndim == 1: wav = wav.unsqueeze(0)
|
||||
else: wav = wav.t()
|
||||
if sr != cfg.common['sample_rate']:
|
||||
wav = torchaudio.functional.resample(wav, sr, cfg.common['sample_rate'])
|
||||
if wav.shape[0] > 1: wav = wav.mean(dim=0, keepdim=True)
|
||||
if wav.size(-1) < self.segment_size:
|
||||
wav = F.pad(wav, (0, self.segment_size - wav.size(-1)))
|
||||
if wav.size(-1) > self.segment_size:
|
||||
start = torch.randint(0, wav.size(-1) - self.segment_size, (1,))
|
||||
wav = wav[..., start : start + self.segment_size]
|
||||
return wav
|
||||
except Exception:
|
||||
return torch.zeros(1, self.segment_size)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--wav-dir", type=str, required=True, help="Path to wavs")
|
||||
parser.add_argument("--epochs", type=int, default=cfg.codec['epochs'])
|
||||
args = parser.parse_args()
|
||||
|
||||
save_dir = cfg.codec['save_dir']
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
device = cfg.device
|
||||
print(f"Launching Codec Training (Float32) on {device}...")
|
||||
|
||||
encoder = Encoder(
|
||||
num_input_mels=cfg.codec['input_mels'],
|
||||
encoder_dim=cfg.codec['dim'],
|
||||
encoder_layers=cfg.codec['layers'],
|
||||
bottleneck_channels=cfg.codec['bottleneck']
|
||||
).to(device).float()
|
||||
|
||||
decoder = Decoder(
|
||||
input_channels=cfg.codec['bottleneck'],
|
||||
decoder_dim=cfg.codec['dim'],
|
||||
decoder_layers=cfg.codec['layers']
|
||||
).to(device).float()
|
||||
|
||||
opt = AdamW(list(encoder.parameters()) + list(decoder.parameters()), lr=float(cfg.codec['lr']))
|
||||
criterion = SpectralLoss()
|
||||
|
||||
ds = WavDataset(args.wav_dir, segment_size=cfg.codec['segment_size'])
|
||||
dl = DataLoader(ds, batch_size=cfg.codec['batch_size'], shuffle=True, num_workers=0, pin_memory=True)
|
||||
|
||||
for epoch in range(args.epochs):
|
||||
encoder.train()
|
||||
decoder.train()
|
||||
pbar = tqdm(dl)
|
||||
for wav in pbar:
|
||||
wav = wav.to(device)
|
||||
z = encoder(wav)
|
||||
rec = decoder(z)
|
||||
loss = criterion(rec, wav)
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
pbar.set_description(f"Ep {epoch+1}/{args.epochs} | Loss: {loss.item():.4f}")
|
||||
|
||||
if (epoch + 1) % 5 == 0 or epoch == args.epochs - 1:
|
||||
torch.save(encoder.state_dict(), f"{save_dir}/encoder.pth")
|
||||
torch.save(decoder.state_dict(), f"{save_dir}/decoder.pth")
|
||||
print(f"Weights saved to {save_dir}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,141 @@
|
||||
import argparse
|
||||
import time
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
# Assumes the decoder folder is in your root directory
|
||||
from decoder.decoder import SopranoDecoder
|
||||
|
||||
|
||||
class SopranoInferencer:
|
||||
"""
|
||||
Stitches together the fine-tuned LLM (Stage 1) and the fine-tuned
|
||||
Vocos GAN Decoder (Stage 2) to generate custom audio from text.
|
||||
"""
|
||||
|
||||
def __init__(self, config_path: str, custom_text: str = None):
|
||||
with open(config_path, 'r') as f:
|
||||
self.config = yaml.safe_load(f)
|
||||
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.dtype = getattr(torch, self.config['optimizations']['mixed_precision'])
|
||||
|
||||
# Determine Paths
|
||||
self.llm_dir = Path(self.config['model']['llm_save_dir']) / "final_llm"
|
||||
self.decoder_path = Path(self.config['model']['decoder_save_dir']) / "final_decoder" / "decoder.pth"
|
||||
|
||||
# Load Components
|
||||
self.tokenizer = self._load_tokenizer()
|
||||
self.llm = self._load_llm()
|
||||
self.decoder = self._load_decoder()
|
||||
|
||||
# Token Boundaries
|
||||
self.audio_min_id = int(self.tokenizer.convert_tokens_to_ids("[0]")) # type: ignore
|
||||
self.audio_max_id = int(self.tokenizer.convert_tokens_to_ids("[7999]")) # type: ignore
|
||||
|
||||
def _load_tokenizer(self):
|
||||
print(f"Loading Tokenizer from {self.llm_dir}...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.llm_dir)
|
||||
tokenizer.padding_side = 'right'
|
||||
return tokenizer
|
||||
|
||||
def _load_llm(self):
|
||||
print(f"Loading Fine-Tuned LLM from {self.llm_dir}...")
|
||||
llm = AutoModelForCausalLM.from_pretrained(
|
||||
self.llm_dir,
|
||||
attn_implementation=self.config['optimizations']['attn_implementation']
|
||||
)
|
||||
llm.to(self.dtype).to(self.device)
|
||||
llm.eval()
|
||||
return llm
|
||||
|
||||
def _load_decoder(self):
|
||||
print(f"Loading Fine-Tuned GAN Decoder from {self.decoder_path}...")
|
||||
if not self.decoder_path.exists():
|
||||
raise FileNotFoundError(f"Decoder weights not found at {self.decoder_path}. Did Stage 2 finish?")
|
||||
|
||||
decoder = SopranoDecoder()
|
||||
decoder.load_state_dict(torch.load(self.decoder_path, map_location='cpu'))
|
||||
decoder.to(self.device)
|
||||
decoder.eval()
|
||||
return decoder
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, text: str, output_path: str = "custom_output.wav"):
|
||||
print(f"\n--- Generating Audio ---")
|
||||
print(f"Text: \"{text}\"")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# 1. Format the Prompt exactly as seen in training
|
||||
prompt = f"[TEXT]{text}[START]"
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
|
||||
|
||||
# 2. Autoregressive Generation (Generate discrete tokens)
|
||||
print("1. Generating discrete audio tokens...")
|
||||
generation_output = self.llm.generate(
|
||||
input_ids=inputs['input_ids'],
|
||||
attention_mask=inputs.get('attention_mask'),
|
||||
max_new_tokens=1024,
|
||||
do_sample=True,
|
||||
temperature=0.8, # Slightly creative but focused
|
||||
top_k=50,
|
||||
top_p=0.95,
|
||||
pad_token_id=self.tokenizer.pad_token_id,
|
||||
eos_token_id=self.tokenizer.eos_token_id,
|
||||
repetition_penalty=1.2
|
||||
)
|
||||
|
||||
# 3. Extract Continuous Hidden States (The "Single Forward Pass" Fix)
|
||||
print("2. Extracting continuous hidden states...")
|
||||
with torch.autocast(device_type=self.device.type, dtype=self.dtype):
|
||||
hidden_out = self.llm(generation_output, output_hidden_states=True).hidden_states[-1].to(torch.float32)
|
||||
|
||||
# 4. Filter for only the audio tokens
|
||||
sequence = generation_output[0]
|
||||
audio_mask = (sequence >= self.audio_min_id) & (sequence <= self.audio_max_id)
|
||||
|
||||
num_audio_tokens = audio_mask.sum().item()
|
||||
print(f" -> Found {num_audio_tokens} audio tokens.")
|
||||
|
||||
if num_audio_tokens == 0:
|
||||
print("[!] The LLM generated no audio tokens. Try tweaking generation temperature or training longer.")
|
||||
return
|
||||
|
||||
# Gather latents: Shape changes from (1, T_total, Dim) -> (1, Dim, T_audio) for Decoder
|
||||
audio_latents = hidden_out[0][audio_mask].unsqueeze(0).transpose(1, 2)
|
||||
|
||||
# 5. Synthesize Waveform via GAN Decoder
|
||||
print("3. Synthesizing HD waveform via Vocos GAN...")
|
||||
waveform = self.decoder(audio_latents).squeeze(1) # (1, Samples)
|
||||
waveform = waveform.cpu()
|
||||
|
||||
# 6. Save Audio
|
||||
torchaudio.save(output_path, waveform, self.config['dataset']['sample_rate'])
|
||||
|
||||
dt = time.time() - start_time
|
||||
print(f"\n[✓] Success! Audio saved to {output_path} in {dt:.2f} seconds.")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Test Soprano 2-Stage Generation")
|
||||
parser.add_argument("--config", type=str, default="config.yaml", help="Path to config file")
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="This is a test of the fully rebuilt, two stage Soprano factory.",
|
||||
help="The text you want the model to speak."
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="output.wav", help="Output filename")
|
||||
args = parser.parse_args()
|
||||
|
||||
inferencer = SopranoInferencer(config_path=args.config)
|
||||
inferencer.generate(text=args.text, output_path=args.output)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,81 +0,0 @@
|
||||
# verify_voice.py - Pure ASCII, Standardized for V3
|
||||
import argparse
|
||||
import torch
|
||||
import soundfile as sf
|
||||
import torchaudio
|
||||
import os
|
||||
from model.encoder import Encoder
|
||||
from model.decoder import Decoder
|
||||
from utils.config import cfg
|
||||
|
||||
def setup_audio_backend():
|
||||
cwd = os.getcwd()
|
||||
local_ffmpeg = os.path.join(cwd, "tools", "ffmpeg")
|
||||
if os.path.exists(local_ffmpeg) and os.path.isdir(local_ffmpeg):
|
||||
if local_ffmpeg not in os.environ["PATH"]:
|
||||
os.environ["PATH"] = local_ffmpeg + os.pathsep + os.environ["PATH"]
|
||||
|
||||
setup_audio_backend()
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Verify Codec Reconstruction")
|
||||
parser.add_argument("--input", type=str, required=True, help="Single wav file")
|
||||
parser.add_argument("--output", type=str, default="verify_out.wav")
|
||||
args = parser.parse_args()
|
||||
|
||||
device = cfg.device
|
||||
print(f"Running verification on {device}")
|
||||
|
||||
# Initialize Encoder
|
||||
encoder = Encoder(
|
||||
num_input_mels=cfg.codec['input_mels'],
|
||||
encoder_dim=cfg.codec['dim'],
|
||||
encoder_layers=cfg.codec['layers'],
|
||||
bottleneck_channels=cfg.codec['bottleneck']
|
||||
).to(device).float()
|
||||
|
||||
# Initialize Decoder (5 channels for Stage 0 weights)
|
||||
decoder = Decoder(
|
||||
input_channels=cfg.codec['bottleneck'],
|
||||
decoder_dim=cfg.codec['dim'],
|
||||
decoder_layers=cfg.codec['layers']
|
||||
).to(device).float()
|
||||
|
||||
# Load weights
|
||||
e_path = os.path.join(cfg.codec['save_dir'], "encoder.pth")
|
||||
d_path = os.path.join(cfg.codec['save_dir'], "decoder.pth")
|
||||
|
||||
if not os.path.exists(e_path) or not os.path.exists(d_path):
|
||||
print(f"Error: Weights not found in {cfg.codec['save_dir']}")
|
||||
return
|
||||
|
||||
encoder.load_state_dict(torch.load(e_path, map_location=device, weights_only=True))
|
||||
decoder.load_state_dict(torch.load(d_path, map_location=device, weights_only=True))
|
||||
|
||||
encoder.eval()
|
||||
decoder.eval()
|
||||
|
||||
# Load Audio
|
||||
wav, sr = sf.read(args.input)
|
||||
wav = torch.from_numpy(wav).float()
|
||||
if wav.ndim == 1: wav = wav.unsqueeze(0)
|
||||
else: wav = wav.t()
|
||||
|
||||
if sr != cfg.common['sample_rate']:
|
||||
wav = torchaudio.functional.resample(wav, sr, cfg.common['sample_rate'])
|
||||
|
||||
if wav.shape[0] > 1: wav = wav.mean(dim=0, keepdim=True)
|
||||
wav = wav.unsqueeze(0).to(device)
|
||||
|
||||
# Process
|
||||
with torch.no_grad():
|
||||
z = encoder(wav)
|
||||
out = decoder(z)
|
||||
|
||||
# Save
|
||||
out_np = out.squeeze().cpu().numpy()
|
||||
sf.write(args.output, out_np, cfg.common['sample_rate'])
|
||||
print(f"Verification complete: {args.output}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,217 +0,0 @@
|
||||
# train.py - Soprano Reforged V3 (Standardized Precision & Joint Training)
|
||||
# EMOJI-FREE VERSION for Windows Compatibility
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import random
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
import numpy as np
|
||||
import torch
|
||||
import soundfile as sf
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.optim import AdamW
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from tqdm import tqdm
|
||||
|
||||
# Import modular components
|
||||
from dataset import AudioDataset
|
||||
from model.decoder import Decoder
|
||||
from training.collator import SopranoCollator
|
||||
from training.loss import AudioReconstructionLoss
|
||||
from utils.config import cfg
|
||||
|
||||
def set_seed(seed):
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input-dir", type=str, default="./mio_dataset", help="Path containing train.json")
|
||||
parser.add_argument("--epochs", type=int, default=cfg.training['epochs'])
|
||||
parser.add_argument("--save-dir", type=str, default=None, help="Override default save directory")
|
||||
args = parser.parse_args()
|
||||
|
||||
set_seed(cfg.common['seed'])
|
||||
save_dir = args.save_dir if args.save_dir else cfg.training['save_dir']
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
os.makedirs(os.path.join(save_dir, "samples"), exist_ok=True)
|
||||
device = cfg.device
|
||||
|
||||
# 1. Load Tokenizer & LLM (FORCED FLOAT32)
|
||||
base_model = cfg.training['base_model']
|
||||
print(f"[System] Launching Stage 2 Training in Float32...")
|
||||
print(f"[Model] Loading Base LLM: {base_model}...")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(base_model)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
llm = AutoModelForCausalLM.from_pretrained(
|
||||
base_model,
|
||||
dtype=torch.float32,
|
||||
attn_implementation="sdpa"
|
||||
).to(device)
|
||||
llm.train()
|
||||
|
||||
# 2. Load Decoder
|
||||
print(f"[Model] Initializing Decoder for Joint Training (Dim: {llm.config.hidden_size})...")
|
||||
decoder = Decoder(
|
||||
input_channels=llm.config.hidden_size,
|
||||
decoder_dim=cfg.codec['dim'],
|
||||
decoder_layers=cfg.codec['layers']
|
||||
).to(device)
|
||||
|
||||
decoder_ckpt = f"{cfg.codec['save_dir']}/decoder.pth"
|
||||
if os.path.exists(decoder_ckpt):
|
||||
print(f"[Checkpoint] Loading pretrained backbone from Stage 0...")
|
||||
state_dict = torch.load(decoder_ckpt, map_location=device, weights_only=True)
|
||||
# Skip input_proj as it resizes from 5 to 512
|
||||
filtered_dict = {k: v for k, v in state_dict.items() if "input_proj" not in k}
|
||||
decoder.load_state_dict(filtered_dict, strict=False)
|
||||
|
||||
decoder.train()
|
||||
|
||||
# 3. Data Setup
|
||||
collator = SopranoCollator(tokenizer, seq_len=cfg.training['seq_len'])
|
||||
train_ds = AudioDataset(f"{args.input_dir}/train.json")
|
||||
|
||||
llm_loader = DataLoader(
|
||||
train_ds, batch_size=cfg.training['batch_size'], shuffle=True,
|
||||
collate_fn=collator.pack_for_llm, num_workers=0
|
||||
)
|
||||
|
||||
dec_loader = DataLoader(
|
||||
train_ds, batch_size=2, shuffle=True,
|
||||
collate_fn=collator.collate_for_decoder, num_workers=0
|
||||
)
|
||||
dec_iter = iter(dec_loader)
|
||||
|
||||
# 4. Optimizer & Loss
|
||||
optimizer = AdamW([
|
||||
{"params": llm.parameters(), "lr": float(cfg.training['base_lr'])},
|
||||
{"params": decoder.parameters(), "lr": float(cfg.training['decoder_lr'])}
|
||||
], weight_decay=cfg.training['weight_decay'])
|
||||
|
||||
dec_criterion = AudioReconstructionLoss(sample_rate=cfg.common['sample_rate'])
|
||||
|
||||
# 5. Training Loop
|
||||
global_step = 0
|
||||
grad_accum = cfg.training['grad_accum_steps']
|
||||
|
||||
for epoch in range(args.epochs):
|
||||
pbar = tqdm(llm_loader, ascii=True) # Ensure tqdm doesn't use unicode blocks
|
||||
for batch_idx, (input_ids, target_ids) in enumerate(pbar):
|
||||
if input_ids is None: continue
|
||||
|
||||
input_ids, target_ids = input_ids.to(device), target_ids.to(device)
|
||||
outputs = llm(input_ids, labels=target_ids)
|
||||
llm_loss = outputs.loss
|
||||
(llm_loss / grad_accum).backward()
|
||||
|
||||
# --- Decoder Step (Joint Training) ---
|
||||
dec_loss_val = 0.0
|
||||
if (global_step + 1) % cfg.training['decoder_step_freq'] == 0:
|
||||
try:
|
||||
try:
|
||||
d_input_ids, wav_paths = next(dec_iter)
|
||||
except StopIteration:
|
||||
dec_iter = iter(dec_loader)
|
||||
d_input_ids, wav_paths = next(dec_iter)
|
||||
|
||||
with torch.no_grad():
|
||||
llm_out = llm(d_input_ids.to(device), output_hidden_states=True)
|
||||
hidden_states_full = llm_out.hidden_states[-1]
|
||||
|
||||
# --- Slicing Logic: Only decode tokens AFTER [START] ---
|
||||
start_token_id = tokenizer.convert_tokens_to_ids("[START]")
|
||||
valid_hidden_states = []
|
||||
target_audio_list = []
|
||||
|
||||
# We need to process each sample in the batch individually because start_idx varies
|
||||
for i in range(len(d_input_ids)):
|
||||
row_ids = d_input_ids[i].tolist()
|
||||
try:
|
||||
# varying start position
|
||||
start_idx = row_ids.index(start_token_id)
|
||||
# Slice after [START]
|
||||
audio_latent = hidden_states_full[i, start_idx+1:, :]
|
||||
valid_hidden_states.append(audio_latent)
|
||||
except ValueError:
|
||||
# Start token not found? Should not happen with correct data prep
|
||||
# Fallback: use whole sequence (garbage in, garbage out, but avoids crash)
|
||||
valid_hidden_states.append(hidden_states_full[i])
|
||||
|
||||
# Load corresponding audio target
|
||||
path = wav_paths[i]
|
||||
w, _ = sf.read(path)
|
||||
w = torch.from_numpy(w).float().to(device)
|
||||
if w.ndim == 1: w = w.unsqueeze(0)
|
||||
else: w = w.t()
|
||||
target_audio_list.append(w)
|
||||
|
||||
# Pad the hidden states to create a batch for the decoder
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
# pad_sequence expects [L, D], and our audio_latent is [L, D]
|
||||
# batch_first=True -> [B, Max_L, D]
|
||||
dec_input = pad_sequence(valid_hidden_states, batch_first=True, padding_value=0.0)
|
||||
|
||||
pred_audio = decoder(dec_input)
|
||||
|
||||
# Align pred_audio and target_audio lengths
|
||||
# pred_audio: [B, 1, L_pred]
|
||||
# target_audio needs to be the same list-based loss or padded
|
||||
|
||||
# Because targets are variable length, we should compute loss per sample or pad targets
|
||||
# Let's pad targets to match pred_audio
|
||||
max_len = pred_audio.shape[-1]
|
||||
padded_targets = []
|
||||
for t in target_audio_list:
|
||||
if t.shape[-1] < max_len:
|
||||
t = torch.nn.functional.pad(t, (0, max_len - t.shape[-1]))
|
||||
elif t.shape[-1] > max_len:
|
||||
t = t[..., :max_len]
|
||||
padded_targets.append(t)
|
||||
|
||||
target_audio = torch.stack(padded_targets)
|
||||
|
||||
dec_loss = dec_criterion(pred_audio, target_audio)
|
||||
(dec_loss * cfg.training['decoder_loss_weight'] / grad_accum).backward()
|
||||
dec_loss_val = dec_loss.item()
|
||||
|
||||
if global_step % 100 == 0:
|
||||
print(f"[Debug] Decoder Step {global_step} Run. Loss: {dec_loss_val}")
|
||||
|
||||
except Exception as e:
|
||||
with open("training_errors.log", "a", encoding="utf-8") as f:
|
||||
f.write(f"Step {global_step} Error: {str(e)}\n")
|
||||
|
||||
# --- Optimization ---
|
||||
if (batch_idx + 1) % grad_accum == 0:
|
||||
torch.nn.utils.clip_grad_norm_(llm.parameters(), 1.0)
|
||||
torch.nn.utils.clip_grad_norm_(decoder.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
optimizer.zero_grad()
|
||||
global_step += 1
|
||||
|
||||
pbar.set_description(f"Epoch {epoch} | LLM Loss: {llm_loss.item():.4f} | Dec Loss: {dec_loss_val:.4f}")
|
||||
|
||||
if global_step % 500 == 0 and dec_loss_val > 0:
|
||||
s_path = f"{save_dir}/samples/step_{global_step}.wav"
|
||||
sf.write(s_path, pred_audio[0].detach().cpu().squeeze().numpy(), 32000)
|
||||
|
||||
epoch_path = f"{save_dir}/epoch_{epoch}"
|
||||
os.makedirs(epoch_path, exist_ok=True)
|
||||
llm.save_pretrained(epoch_path)
|
||||
tokenizer.save_pretrained(epoch_path)
|
||||
torch.save(decoder.state_dict(), f"{epoch_path}/decoder.pth")
|
||||
|
||||
print("[Done] Training Complete.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
-154
@@ -1,154 +0,0 @@
|
||||
# train_codec.py - Stage 0: Codec Training (Standardized for V3)
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Windows UTF-8 Support
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchaudio
|
||||
import soundfile as sf
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torch.optim import AdamW
|
||||
from tqdm import tqdm
|
||||
|
||||
# Import our architecture
|
||||
from model.encoder import Encoder
|
||||
from model.decoder import Decoder
|
||||
from utils.config import cfg
|
||||
|
||||
# --- Windows / Audio Backend Setup ---
|
||||
def setup_audio_backend():
|
||||
cwd = os.getcwd()
|
||||
local_ffmpeg = os.path.join(cwd, "tools", "ffmpeg")
|
||||
if os.path.exists(local_ffmpeg) and os.path.isdir(local_ffmpeg):
|
||||
if local_ffmpeg not in os.environ["PATH"]:
|
||||
os.environ["PATH"] = local_ffmpeg + os.pathsep + os.environ["PATH"]
|
||||
|
||||
setup_audio_backend()
|
||||
|
||||
class SpectralLoss(torch.nn.Module):
|
||||
"""
|
||||
Computes time-domain L1 loss and frequency-domain Mel loss.
|
||||
Forced to Float32 for numerical stability.
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.mel = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=cfg.common['sample_rate'],
|
||||
n_mels=cfg.codec['input_mels'],
|
||||
n_fft=2048, hop_length=512
|
||||
)
|
||||
|
||||
def forward(self, pred, target):
|
||||
pred, target = pred.float(), target.float()
|
||||
min_len = min(pred.shape[-1], target.shape[-1])
|
||||
pred = pred[..., :min_len]
|
||||
target = target[..., :min_len]
|
||||
|
||||
loss_time = F.l1_loss(pred, target)
|
||||
|
||||
if self.mel.mel_scale.fb.device != pred.device:
|
||||
self.mel = self.mel.to(pred.device)
|
||||
|
||||
loss_mel = F.l1_loss(self.mel(pred), self.mel(target))
|
||||
return loss_time + loss_mel
|
||||
|
||||
class WavDataset(Dataset):
|
||||
def __init__(self, glob_pattern, segment_size):
|
||||
self.files = glob.glob(glob_pattern, recursive=True)
|
||||
self.segment_size = segment_size
|
||||
print(f"Stage 0: Found {len(self.files)} wav files for training.")
|
||||
|
||||
def __len__(self):
|
||||
return len(self.files)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
try:
|
||||
wav_np, sr = sf.read(self.files[idx])
|
||||
wav = torch.from_numpy(wav_np).float()
|
||||
|
||||
if wav.ndim == 1: wav = wav.unsqueeze(0)
|
||||
else: wav = wav.t()
|
||||
|
||||
if sr != cfg.common['sample_rate']:
|
||||
wav = torchaudio.functional.resample(wav, sr, cfg.common['sample_rate'])
|
||||
|
||||
if wav.shape[0] > 1: wav = wav.mean(dim=0, keepdim=True)
|
||||
|
||||
if wav.size(-1) < self.segment_size:
|
||||
wav = F.pad(wav, (0, self.segment_size - wav.size(-1)))
|
||||
|
||||
if wav.size(-1) > self.segment_size:
|
||||
start = torch.randint(0, wav.size(-1) - self.segment_size, (1,))
|
||||
wav = wav[..., start : start + self.segment_size]
|
||||
|
||||
return wav
|
||||
except Exception as e:
|
||||
return torch.zeros(1, self.segment_size)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--wav-dir", type=str, required=True, help="Path to wavs")
|
||||
parser.add_argument("--epochs", type=int, default=cfg.codec['epochs'])
|
||||
parser.add_argument("--save-dir", type=str, default=None, help="Override default save directory")
|
||||
args = parser.parse_args()
|
||||
|
||||
save_dir = args.save_dir if args.save_dir else cfg.codec['save_dir']
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
device = cfg.device
|
||||
print(f"Launching Codec Training (Float32) on {device}...")
|
||||
|
||||
# 1. Initialize Encoder
|
||||
encoder = Encoder(
|
||||
num_input_mels=cfg.codec['input_mels'],
|
||||
encoder_dim=cfg.codec['dim'],
|
||||
encoder_layers=cfg.codec['layers'],
|
||||
bottleneck_channels=cfg.codec['bottleneck']
|
||||
).to(device).float()
|
||||
|
||||
# 2. Initialize Decoder (SET INPUT TO 5 FOR STAGE 0)
|
||||
decoder = Decoder(
|
||||
input_channels=cfg.codec['bottleneck'], # <--- Specifically uses the 5 bottleneck channels
|
||||
decoder_dim=cfg.codec['dim'],
|
||||
decoder_layers=cfg.codec['layers']
|
||||
).to(device).float()
|
||||
|
||||
opt = AdamW(list(encoder.parameters()) + list(decoder.parameters()), lr=float(cfg.codec['lr']))
|
||||
criterion = SpectralLoss()
|
||||
|
||||
ds = WavDataset(args.wav_dir, segment_size=cfg.codec['segment_size'])
|
||||
dl = DataLoader(ds, batch_size=cfg.codec['batch_size'], shuffle=True, num_workers=0, pin_memory=True)
|
||||
|
||||
# 3. Training Loop
|
||||
for epoch in range(args.epochs):
|
||||
encoder.train()
|
||||
decoder.train()
|
||||
pbar = tqdm(dl)
|
||||
|
||||
for wav in pbar:
|
||||
wav = wav.to(device)
|
||||
|
||||
# Forward
|
||||
z = encoder(wav)
|
||||
rec = decoder(z)
|
||||
loss = criterion(rec, wav)
|
||||
|
||||
# Backward
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
|
||||
pbar.set_description(f"Ep {epoch+1}/{args.epochs} | Loss: {loss.item():.4f}")
|
||||
|
||||
# Save checkpoints
|
||||
if (epoch + 1) % 5 == 0 or epoch == args.epochs - 1:
|
||||
torch.save(encoder.state_dict(), f"{save_dir}/encoder.pth")
|
||||
torch.save(decoder.state_dict(), f"{save_dir}/decoder.pth")
|
||||
print(f"Weights saved to {save_dir}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,370 @@
|
||||
# train_decoder.py - Trains the SopranoDecoder using a frozen LLM backbone with GAN and MR-STFT losses for high-fidelity audio synthesis.
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Tuple, Dict, Any, List
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import yaml
|
||||
from huggingface_hub import hf_hub_download
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer
|
||||
|
||||
from dataset import SopranoDecoderDataset, DecoderCollator, 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
|
||||
)
|
||||
|
||||
|
||||
def setup_logger(log_file: str):
|
||||
logger = logging.getLogger("SopranoDecoder")
|
||||
logger.setLevel(logging.INFO)
|
||||
if logger.hasHandlers():
|
||||
logger.handlers.clear()
|
||||
|
||||
formatter = logging.Formatter('%(asctime)s - %(message)s')
|
||||
fh = logging.FileHandler(log_file)
|
||||
fh.setFormatter(formatter)
|
||||
logger.addHandler(fh)
|
||||
|
||||
ch = logging.StreamHandler()
|
||||
ch.setFormatter(logging.Formatter('%(message)s'))
|
||||
logger.addHandler(ch)
|
||||
return logger
|
||||
|
||||
|
||||
class SopranoDecoderTrainer:
|
||||
"""
|
||||
Stage 2: Freezes the LLM, extracts hidden states, and trains a Vocos
|
||||
decoder with MR-STFT and GAN losses to synthesize high-fidelity audio.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any], use_disc: bool = True):
|
||||
self.config = config
|
||||
self.use_disc = use_disc
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.logger = setup_logger(self.config['logging'].get('log_file', 'decoder_training.log'))
|
||||
self.logger.info("Initializing Soprano Stage-2 Decoder GAN Trainer...")
|
||||
|
||||
# Setup Hardware & Seeds
|
||||
self.seed = self.config['optimizations']['seed']
|
||||
torch.manual_seed(self.seed)
|
||||
random.seed(self.seed)
|
||||
np.random.seed(self.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(self.seed)
|
||||
|
||||
self.allow_tf32 = self.config['optimizations']['allow_tf32']
|
||||
torch.set_float32_matmul_precision('high' if self.allow_tf32 else 'highest')
|
||||
|
||||
self.dtype = getattr(torch, self.config['optimizations']['mixed_precision'])
|
||||
|
||||
# Setup TensorBoard
|
||||
tb_dir = Path(self.config['logging']['tensorboard_dir']) / "decoder_stage"
|
||||
self.writer = SummaryWriter(log_dir=str(tb_dir))
|
||||
|
||||
# Load Components
|
||||
self.tokenizer = self._setup_tokenizer()
|
||||
self.llm, self.decoder, self.discriminator = self._setup_models()
|
||||
self.train_loader, self.val_loader = self._setup_dataloaders()
|
||||
|
||||
# Setup Losses
|
||||
self.mel_fn = MelSpectrogramWrapper().to(self.device)
|
||||
self.mr_stft = MultiResolutionSTFTLoss().to(self.device)
|
||||
|
||||
# Loss Weights
|
||||
self.lambda_mel = 45.0
|
||||
self.lambda_fm = 2.0
|
||||
self.lambda_gen = 1.0
|
||||
self.lambda_stft = 1.0
|
||||
|
||||
# Optimizers
|
||||
self.max_lr_g = self.config['decoder_training']['learning_rate_g']
|
||||
self.max_lr_d = self.config['decoder_training']['learning_rate_d']
|
||||
self.opt_g = torch.optim.AdamW(self.decoder.parameters(), lr=self.max_lr_g, betas=(0.8, 0.99), weight_decay=0.1)
|
||||
|
||||
if self.use_disc:
|
||||
self.opt_d = torch.optim.AdamW(self.discriminator.parameters(), lr=self.max_lr_d, betas=(0.8, 0.99), weight_decay=0.1)
|
||||
|
||||
def _setup_tokenizer(self) -> PreTrainedTokenizer:
|
||||
self.logger.info("Loading Tokenizer...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.config['model']['base_llm'], use_fast=True)
|
||||
tokenizer.padding_side = 'right'
|
||||
return tokenizer
|
||||
|
||||
def _setup_models(self):
|
||||
# 1. Load the FINE-TUNED LLM from Stage 1
|
||||
llm_path = Path(self.config['model']['llm_save_dir']) / "final_llm"
|
||||
if not llm_path.exists():
|
||||
# Fallback to the base model if the user skipped Stage 1 for some reason
|
||||
self.logger.warning(f"Fine-tuned LLM not found at {llm_path}. Falling back to base model.")
|
||||
llm_path = self.config['model']['base_llm']
|
||||
|
||||
self.logger.info(f"Loading LLM from {llm_path} (Frozen)...")
|
||||
llm = AutoModelForCausalLM.from_pretrained(llm_path, attn_implementation=self.config['optimizations']['attn_implementation'])
|
||||
llm.to(self.dtype).to(self.device)
|
||||
llm.eval()
|
||||
for param in llm.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
# 2. Load the base Vocos Decoder
|
||||
self.logger.info("Loading Base Vocos Decoder (Trainable)...")
|
||||
decoder = SopranoDecoder()
|
||||
decoder_path = hf_hub_download(repo_id=self.config['model']['base_llm'], filename='decoder.pth')
|
||||
decoder.load_state_dict(torch.load(decoder_path, map_location='cpu'))
|
||||
decoder.to(self.device)
|
||||
decoder.train()
|
||||
|
||||
# 3. Initialize GAN Discriminator
|
||||
discriminator = None
|
||||
if self.use_disc:
|
||||
self.logger.info("Initializing GAN Discriminator...")
|
||||
discriminator = Discriminator()
|
||||
discriminator.to(self.device)
|
||||
discriminator.train()
|
||||
|
||||
return llm, decoder, discriminator
|
||||
|
||||
def _setup_dataloaders(self) -> Tuple[DataLoader, DataLoader]:
|
||||
input_dir = Path(self.config['dataset']['input_dir'])
|
||||
batch_size = self.config['decoder_training']['batch_size']
|
||||
collator = DecoderCollator(tokenizer=self.tokenizer)
|
||||
|
||||
train_ds = SopranoDecoderDataset(str(input_dir / 'train.json'), target_sr=self.config['dataset']['sample_rate'])
|
||||
val_ds = SopranoDecoderDataset(str(input_dir / 'val.json'), target_sr=self.config['dataset']['sample_rate'])
|
||||
|
||||
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4, collate_fn=collator, drop_last=True)
|
||||
val_loader = DataLoader(val_ds, batch_size=max(1, batch_size // 2), shuffle=False, num_workers=2, collate_fn=collator, drop_last=True)
|
||||
return train_loader, val_loader
|
||||
|
||||
def _crop_for_gan(self, real_audio: torch.Tensor, fake_audio: torch.Tensor, valid_lens: List[int]) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Randomly crops audio segments for the Discriminator to grade."""
|
||||
seg_size = self.config['decoder_training']['segment_size_samples']
|
||||
bsz = real_audio.size(0)
|
||||
|
||||
real_crops, fake_crops = [], []
|
||||
|
||||
for b in range(bsz):
|
||||
v_len = min(valid_lens[b], real_audio.size(1), fake_audio.size(1))
|
||||
|
||||
if v_len <= seg_size:
|
||||
pad_len = seg_size - v_len
|
||||
r_c = F.pad(real_audio[b, :v_len], (0, pad_len))
|
||||
f_c = F.pad(fake_audio[b, :v_len], (0, pad_len))
|
||||
else:
|
||||
start_idx = random.randint(0, v_len - seg_size)
|
||||
r_c = real_audio[b, start_idx : start_idx + seg_size]
|
||||
f_c = fake_audio[b, start_idx : start_idx + seg_size]
|
||||
|
||||
real_crops.append(r_c)
|
||||
fake_crops.append(f_c)
|
||||
|
||||
real_tensor = torch.stack(real_crops).unsqueeze(1) # (B, 1, T)
|
||||
fake_tensor = torch.stack(fake_crops).unsqueeze(1)
|
||||
return real_tensor, fake_tensor
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(self, step: int):
|
||||
self.decoder.eval()
|
||||
if self.use_disc: self.discriminator.eval()
|
||||
|
||||
val_mel, val_stft = 0.0, 0.0
|
||||
val_steps = min(5, len(self.val_loader))
|
||||
val_iter = iter(self.val_loader)
|
||||
|
||||
for i in range(val_steps):
|
||||
x, y, gt_audio, audio_mask = next(val_iter)
|
||||
x, y, gt_audio, audio_mask = x.to(self.device), y.to(self.device), gt_audio.to(self.device), audio_mask.to(self.device)
|
||||
|
||||
with torch.autocast(device_type=self.device.type, dtype=self.dtype):
|
||||
outputs = self.llm(x, output_hidden_states=True)
|
||||
hidden_states = outputs.hidden_states[-1].to(torch.float32)
|
||||
|
||||
# Gather and pad latents
|
||||
gathered_list = [hidden_states[b][audio_mask[b]] for b in range(hidden_states.size(0))]
|
||||
decoder_in = torch.nn.utils.rnn.pad_sequence(gathered_list, batch_first=True).transpose(1, 2)
|
||||
|
||||
fake_audio = self.decoder(decoder_in).squeeze(1)
|
||||
min_len = min(fake_audio.size(1), gt_audio.size(1))
|
||||
fake_audio, real_audio = fake_audio[:, :min_len], gt_audio[:, :min_len]
|
||||
|
||||
# Generate Mel Images on first batch
|
||||
if i == 0:
|
||||
gen_mel = self.mel_fn(fake_audio[0:1]).squeeze(0).cpu().numpy()
|
||||
real_mel = self.mel_fn(real_audio[0:1]).squeeze(0).cpu().numpy()
|
||||
|
||||
fig, ax = plt.subplots(2, 1, figsize=(10, 6))
|
||||
ax[0].imshow(real_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(f"Generated Mel (Step {step})")
|
||||
plt.tight_layout()
|
||||
|
||||
self.writer.add_figure("Val/Mel_Spectrogram", fig, step)
|
||||
self.writer.add_audio("Val/Generated_Audio", fake_audio[0], step, sample_rate=self.config['dataset']['sample_rate'])
|
||||
self.writer.flush()
|
||||
|
||||
sc_loss, mag_loss = self.mr_stft(fake_audio, real_audio)
|
||||
val_stft += (sc_loss + mag_loss).item()
|
||||
|
||||
val_stft /= val_steps
|
||||
self.writer.add_scalar("Val/STFT_Loss", val_stft, step)
|
||||
self.logger.info(f"[Val Step {step}] MR-STFT Loss: {val_stft:.3f}")
|
||||
|
||||
self.decoder.train()
|
||||
if self.use_disc: self.discriminator.train()
|
||||
|
||||
def save_checkpoint(self, step: int, name: str = None):
|
||||
save_name = name if name else f"checkpoint_{step}"
|
||||
path = Path(self.config['model']['decoder_save_dir']) / save_name
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
torch.save(self.decoder.state_dict(), path / "decoder.pth")
|
||||
if self.use_disc:
|
||||
torch.save(self.discriminator.state_dict(), path / "discriminator.pth")
|
||||
self.logger.info(f"Saved Decoder checkpoint to {path}")
|
||||
|
||||
def train(self):
|
||||
max_steps = self.config['decoder_training']['max_steps']
|
||||
val_freq = self.config['decoder_training']['val_freq']
|
||||
save_freq = self.config['decoder_training']['save_freq']
|
||||
|
||||
self.logger.info(f"Starting Decoder GAN Training for {max_steps} steps...")
|
||||
train_iter = iter(self.train_loader)
|
||||
pbar = tqdm(range(max_steps), ncols=150, dynamic_ncols=True)
|
||||
|
||||
for step in pbar:
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
x, y, gt_audio, audio_mask = next(train_iter)
|
||||
except StopIteration:
|
||||
train_iter = iter(self.train_loader)
|
||||
x, y, gt_audio, audio_mask = next(train_iter)
|
||||
|
||||
x, y, gt_audio, audio_mask = x.to(self.device), y.to(self.device), gt_audio.to(self.device), audio_mask.to(self.device)
|
||||
|
||||
# 1. Forward LLM (Frozen)
|
||||
with torch.no_grad():
|
||||
with torch.autocast(device_type=self.device.type, dtype=self.dtype):
|
||||
hidden_states = self.llm(x, output_hidden_states=True).hidden_states[-1].to(torch.float32)
|
||||
|
||||
# 2. Gather active audio latent states and pad for Decoder
|
||||
gathered_list = []
|
||||
valid_lens = []
|
||||
for b_idx in range(hidden_states.size(0)):
|
||||
states = hidden_states[b_idx][audio_mask[b_idx]]
|
||||
gathered_list.append(states)
|
||||
valid_lens.append(states.size(0) * SAMPLES_PER_TOKEN)
|
||||
|
||||
decoder_in = torch.nn.utils.rnn.pad_sequence(gathered_list, batch_first=True).transpose(1, 2)
|
||||
|
||||
# =======================================================
|
||||
# TRAIN DISCRIMINATOR
|
||||
# =======================================================
|
||||
d_loss_item = 0.0
|
||||
if self.use_disc:
|
||||
self.opt_d.zero_grad()
|
||||
|
||||
with torch.no_grad(): # Detach generator graph
|
||||
fake_audio = self.decoder(decoder_in).squeeze(1)
|
||||
|
||||
min_len = min(fake_audio.size(1), gt_audio.size(1))
|
||||
fake_audio, real_audio = fake_audio[:, :min_len], gt_audio[:, :min_len]
|
||||
|
||||
real_crops, fake_crops = self._crop_for_gan(real_audio, fake_audio, valid_lens)
|
||||
|
||||
y_d_rs, y_d_gs, _, _ = self.discriminator(real_crops, fake_crops)
|
||||
d_loss, _, _ = discriminator_loss(y_d_rs, y_d_gs)
|
||||
|
||||
d_loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), 1.0)
|
||||
self.opt_d.step()
|
||||
d_loss_item = d_loss.item()
|
||||
|
||||
# =======================================================
|
||||
# TRAIN GENERATOR (DECODER)
|
||||
# =======================================================
|
||||
self.opt_g.zero_grad()
|
||||
fake_audio = self.decoder(decoder_in).squeeze(1)
|
||||
min_len = min(fake_audio.size(1), gt_audio.size(1))
|
||||
fake_audio, real_audio = fake_audio[:, :min_len], gt_audio[:, :min_len]
|
||||
|
||||
# Mel & STFT Losses (Audio reconstruction)
|
||||
pred_mel, gt_mel = self.mel_fn(fake_audio), self.mel_fn(real_audio)
|
||||
|
||||
# Masking for Mel Loss based on valid sequence length
|
||||
mel_loss = 0.0
|
||||
frames_per_token = SAMPLES_PER_TOKEN // 512
|
||||
for b in range(fake_audio.size(0)):
|
||||
v_mel_len = min((valid_lens[b] // 512), pred_mel.size(2))
|
||||
if v_mel_len > 0:
|
||||
mel_loss += F.l1_loss(pred_mel[b, :, :v_mel_len], gt_mel[b, :, :v_mel_len])
|
||||
mel_loss /= fake_audio.size(0)
|
||||
|
||||
sc_loss, mag_loss = self.mr_stft(fake_audio, real_audio)
|
||||
|
||||
# GAN Generator Losses
|
||||
loss_fm, loss_gen = torch.tensor(0.0, device=self.device), torch.tensor(0.0, device=self.device)
|
||||
if self.use_disc:
|
||||
real_crops_g, fake_crops_g = self._crop_for_gan(real_audio, fake_audio, valid_lens)
|
||||
y_d_rs, y_d_gs, fmap_rs, fmap_gs = self.discriminator(real_crops_g, fake_crops_g)
|
||||
|
||||
loss_fm = feature_matching_loss(fmap_rs, fmap_gs)
|
||||
loss_gen, _ = generator_loss(y_d_gs)
|
||||
|
||||
total_loss_g = (self.lambda_mel * mel_loss) + (self.lambda_stft * (sc_loss + mag_loss)) + (self.lambda_gen * loss_gen) + (self.lambda_fm * loss_fm)
|
||||
|
||||
total_loss_g.backward()
|
||||
torch.nn.utils.clip_grad_norm_(self.decoder.parameters(), 1.0)
|
||||
self.opt_g.step()
|
||||
|
||||
# Logging
|
||||
self.writer.add_scalar("Train/Mel_Loss", mel_loss.item(), step)
|
||||
self.writer.add_scalar("Train/MR_STFT_Loss", (sc_loss + mag_loss).item(), step)
|
||||
if self.use_disc:
|
||||
self.writer.add_scalar("Train/D_Loss", d_loss_item, step)
|
||||
self.writer.add_scalar("Train/G_Loss", loss_gen.item(), step)
|
||||
|
||||
dt_ms = (time.time() - start_time) * 1000
|
||||
pbar.set_description(
|
||||
f"STFT: {(sc_loss+mag_loss).item():.3f} | Mel: {mel_loss.item():.3f} | "
|
||||
f"G: {loss_gen.item():.3f} | D: {d_loss_item:.3f} | {dt_ms:.1f}ms"
|
||||
)
|
||||
|
||||
if step > 0 and step % val_freq == 0:
|
||||
self.evaluate(step)
|
||||
|
||||
if step > 0 and step % save_freq == 0:
|
||||
self.save_checkpoint(step)
|
||||
|
||||
self.save_checkpoint(max_steps, name="final_decoder")
|
||||
self.writer.close()
|
||||
self.logger.info("Stage 2 Decoder Training Complete.")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Train Soprano Vocos Decoder (GAN)")
|
||||
parser.add_argument("--config", type=Path, default=Path("config.yaml"))
|
||||
parser.add_argument("--no-disc", action="store_true", help="Disable GAN Discriminator")
|
||||
args = parser.parse_args()
|
||||
|
||||
with open(args.config, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
trainer = SopranoDecoderTrainer(config, use_disc=not args.no_disc)
|
||||
trainer.train()
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
+290
@@ -0,0 +1,290 @@
|
||||
# train_llm.py - Trains the SopranoLLM backbone to predict discrete audio tokens from text prompts.
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Tuple, Dict, Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import yaml
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer
|
||||
|
||||
from dataset import SopranoLLMDataset, LLMCollator
|
||||
|
||||
|
||||
def setup_logger(log_file: str):
|
||||
"""Configures a persistent file and console logger."""
|
||||
logger = logging.getLogger("SopranoLLM")
|
||||
logger.setLevel(logging.INFO)
|
||||
if logger.hasHandlers():
|
||||
logger.handlers.clear()
|
||||
|
||||
formatter = logging.Formatter('%(asctime)s - %(message)s')
|
||||
fh = logging.FileHandler(log_file)
|
||||
fh.setFormatter(formatter)
|
||||
logger.addHandler(fh)
|
||||
|
||||
ch = logging.StreamHandler()
|
||||
ch.setFormatter(logging.Formatter('%(message)s'))
|
||||
logger.addHandler(ch)
|
||||
return logger
|
||||
|
||||
|
||||
class SopranoLLMTrainer:
|
||||
"""
|
||||
Stage 1: Trains the Causal LLM to predict discrete audio tokens from text.
|
||||
Uses pure Autoregressive Cross-Entropy (No Teacher Distillation).
|
||||
"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
self.config = config
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.logger = setup_logger(self.config['logging'].get('log_file', 'llm_training.log'))
|
||||
self.logger.info("Initializing Soprano Stage-1 LLM Trainer...")
|
||||
|
||||
# Hardware Optimizations
|
||||
torch.manual_seed(self.config['optimizations']['seed'])
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(self.config['optimizations']['seed'])
|
||||
|
||||
self.allow_tf32 = self.config['optimizations']['allow_tf32']
|
||||
torch.set_float32_matmul_precision('high' if self.allow_tf32 else 'highest')
|
||||
if torch.cuda.is_available():
|
||||
torch.backends.cuda.matmul.allow_tf32 = self.allow_tf32
|
||||
torch.backends.cudnn.allow_tf32 = self.allow_tf32
|
||||
|
||||
self.dtype = getattr(torch, self.config['optimizations']['mixed_precision'])
|
||||
self.attn_impl = self.config['optimizations']['attn_implementation']
|
||||
|
||||
# Setup TensorBoard
|
||||
tb_dir = Path(self.config['logging']['tensorboard_dir']) / "llm_stage"
|
||||
self.writer = SummaryWriter(log_dir=str(tb_dir))
|
||||
|
||||
# Load Components
|
||||
self.tokenizer = self._setup_tokenizer()
|
||||
self.model = self._setup_model()
|
||||
|
||||
# Token Boundaries
|
||||
self.audio_min_id = int(self.tokenizer.convert_tokens_to_ids("[0]")) # type: ignore
|
||||
self.audio_max_id = int(self.tokenizer.convert_tokens_to_ids("[7999]")) # type: ignore
|
||||
self.stop_id = int(self.tokenizer.convert_tokens_to_ids("[STOP]")) # type: ignore
|
||||
|
||||
# Dataloaders
|
||||
self.train_loader, self.val_loader = self._setup_dataloaders()
|
||||
|
||||
# Optimizer
|
||||
self.max_lr = self.config['llm_training']['learning_rate']
|
||||
self.opt = torch.optim.AdamW(
|
||||
self.model.parameters(),
|
||||
lr=self.max_lr,
|
||||
weight_decay=self.config['llm_training']['weight_decay'],
|
||||
fused=(self.device.type == "cuda")
|
||||
)
|
||||
|
||||
# WSD Scheduler Params
|
||||
self.max_steps = self.config['llm_training']['max_steps']
|
||||
self.warmup_steps = int(self.max_steps * 0.1) # 10% warmup
|
||||
self.cooldown_steps = int(self.max_steps * 0.1) # 10% cooldown
|
||||
self.min_lr = 0.1 * self.max_lr
|
||||
|
||||
def _setup_tokenizer(self) -> PreTrainedTokenizer:
|
||||
self.logger.info("Loading Tokenizer...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.config['model']['base_llm'], use_fast=True)
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.padding_side = 'right'
|
||||
return tokenizer
|
||||
|
||||
def _setup_model(self) -> torch.nn.Module:
|
||||
self.logger.info(f"Loading Student LLM (Attention: {self.attn_impl}, Dtype: {self.dtype})...")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.config['model']['base_llm'],
|
||||
attn_implementation=self.attn_impl
|
||||
)
|
||||
model.to(self.dtype).to(self.device)
|
||||
model.train()
|
||||
return model
|
||||
|
||||
def _setup_dataloaders(self) -> Tuple[DataLoader, DataLoader]:
|
||||
self.logger.info("Initializing DataLoaders...")
|
||||
input_dir = Path(self.config['dataset']['input_dir'])
|
||||
batch_size = self.config['llm_training']['batch_size']
|
||||
|
||||
collator = LLMCollator(tokenizer=self.tokenizer)
|
||||
|
||||
train_dataset = SopranoLLMDataset(str(input_dir / 'train.json'))
|
||||
val_dataset = SopranoLLMDataset(str(input_dir / 'val.json'))
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset, batch_size=batch_size, shuffle=True,
|
||||
num_workers=4, collate_fn=collator, drop_last=True
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_dataset, batch_size=batch_size, shuffle=False,
|
||||
num_workers=2, collate_fn=collator, drop_last=True
|
||||
)
|
||||
return train_loader, val_loader
|
||||
|
||||
def get_lr(self, step: int) -> float:
|
||||
"""Warmup-Stable-Decay (WSD) Learning Rate Schedule."""
|
||||
if step < self.warmup_steps:
|
||||
return self.max_lr * (step + 1) / self.warmup_steps
|
||||
if step < self.max_steps - self.cooldown_steps:
|
||||
return self.max_lr
|
||||
|
||||
decay_ratio = (self.max_steps - step) / self.cooldown_steps
|
||||
return self.min_lr + (self.max_lr - self.min_lr) * decay_ratio
|
||||
|
||||
def compute_loss(self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor, text_factor: float = 0.5):
|
||||
"""Calculates split loss for Text vs Audio tokens."""
|
||||
pred = logits.view(-1, logits.size(-1))
|
||||
labels = targets.reshape(-1)
|
||||
flat_mask = mask.reshape(-1)
|
||||
|
||||
loss = F.cross_entropy(pred, labels, reduction='none') * flat_mask
|
||||
|
||||
# Audio mask includes all acoustic tokens and the STOP token
|
||||
is_audio_cond = (labels >= self.audio_min_id) & (labels <= self.audio_max_id) | (labels == self.stop_id)
|
||||
audio_mask = is_audio_cond & (flat_mask > 0)
|
||||
text_mask = (~is_audio_cond) & (flat_mask > 0)
|
||||
|
||||
# Averages
|
||||
audio_loss = loss[audio_mask].mean() if audio_mask.any() else torch.tensor(0.0, device=self.device)
|
||||
text_loss = loss[text_mask].mean() if text_mask.any() else torch.tensor(0.0, device=self.device)
|
||||
|
||||
# Accuracy (Audio only)
|
||||
predictions = logits.argmax(dim=-1).view(-1)
|
||||
acc = (predictions == labels)[audio_mask].to(torch.float32).mean() if audio_mask.any() else torch.tensor(0.0, device=self.device)
|
||||
|
||||
return audio_loss, text_loss, acc
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(self, step: int):
|
||||
self.model.eval()
|
||||
val_audio_loss, val_text_loss, val_acc = 0.0, 0.0, 0.0
|
||||
val_steps = 10
|
||||
|
||||
val_iter = iter(self.val_loader)
|
||||
for i in range(val_steps):
|
||||
try:
|
||||
x, y, attn_mask = next(val_iter)
|
||||
except StopIteration:
|
||||
break
|
||||
|
||||
x, y, attn_mask = x.to(self.device), y.to(self.device), attn_mask.to(self.device)
|
||||
|
||||
with torch.autocast(device_type=self.device.type, dtype=self.dtype):
|
||||
logits = self.model(x, attention_mask=attn_mask).logits
|
||||
a_loss, t_loss, acc = self.compute_loss(logits, y, attn_mask)
|
||||
|
||||
val_audio_loss += a_loss.item()
|
||||
val_text_loss += t_loss.item()
|
||||
val_acc += acc.item()
|
||||
|
||||
val_audio_loss /= val_steps
|
||||
val_text_loss /= val_steps
|
||||
val_acc /= val_steps
|
||||
|
||||
self.writer.add_scalar("Val/Audio_Loss", val_audio_loss, step)
|
||||
self.writer.add_scalar("Val/Text_Loss", val_text_loss, step)
|
||||
self.writer.add_scalar("Val/Accuracy", val_acc, step)
|
||||
|
||||
self.logger.info(f"[Val Step {step}] Audio CE: {val_audio_loss:.3f} | Text CE: {val_text_loss:.3f} | Acc: {val_acc:.4f}")
|
||||
self.model.train()
|
||||
|
||||
def save_checkpoint(self, step: int, name: str = None):
|
||||
save_name = name if name else f"checkpoint_{step}"
|
||||
path = Path(self.config['model']['llm_save_dir']) / save_name
|
||||
os.makedirs(path, exist_ok=True)
|
||||
self.logger.info(f"Saving LLM Checkpoint to: {path}")
|
||||
self.model.save_pretrained(path)
|
||||
self.tokenizer.save_pretrained(path)
|
||||
|
||||
def train(self):
|
||||
val_freq = self.config['llm_training']['val_freq']
|
||||
save_freq = self.config['llm_training']['save_freq']
|
||||
grad_accum = self.config['llm_training']['grad_accum_steps']
|
||||
text_factor = 0.5 # Weight text loss less than audio loss
|
||||
|
||||
self.logger.info(f"Starting LLM Stage 1 for {self.max_steps} steps...")
|
||||
train_iter = iter(self.train_loader)
|
||||
pbar = tqdm(range(self.max_steps), ncols=150, dynamic_ncols=True)
|
||||
|
||||
for step in pbar:
|
||||
start_time = time.time()
|
||||
self.opt.zero_grad(set_to_none=True)
|
||||
|
||||
accum_a_loss = torch.tensor(0.0, device=self.device)
|
||||
accum_t_loss = torch.tensor(0.0, device=self.device)
|
||||
accum_acc = torch.tensor(0.0, device=self.device)
|
||||
|
||||
for _ in range(grad_accum):
|
||||
try:
|
||||
x, y, attn_mask = next(train_iter)
|
||||
except StopIteration:
|
||||
train_iter = iter(self.train_loader)
|
||||
x, y, attn_mask = next(train_iter)
|
||||
|
||||
x, y, attn_mask = x.to(self.device), y.to(self.device), attn_mask.to(self.device)
|
||||
|
||||
with torch.autocast(device_type=self.device.type, dtype=self.dtype):
|
||||
logits = self.model(x, attention_mask=attn_mask).logits
|
||||
a_loss, t_loss, acc = self.compute_loss(logits, y, attn_mask, text_factor)
|
||||
total_loss = (a_loss + (text_factor * t_loss)) / grad_accum
|
||||
|
||||
total_loss.backward()
|
||||
|
||||
accum_a_loss += a_loss.detach() / grad_accum
|
||||
accum_t_loss += t_loss.detach() / grad_accum
|
||||
accum_acc += acc.detach() / grad_accum
|
||||
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
|
||||
|
||||
lr = self.get_lr(step)
|
||||
for param_group in self.opt.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
self.opt.step()
|
||||
|
||||
# Logging
|
||||
self.writer.add_scalar("Train/Audio_Loss", accum_a_loss.item(), step)
|
||||
self.writer.add_scalar("Train/Text_Loss", accum_t_loss.item(), step)
|
||||
self.writer.add_scalar("Train/Accuracy", accum_acc.item(), step)
|
||||
self.writer.add_scalar("Train/Learning_Rate", lr, step)
|
||||
|
||||
dt_ms = (time.time() - start_time) * 1000
|
||||
pbar.set_description(
|
||||
f"A-Loss: {accum_a_loss.item():.3f} | T-Loss: {accum_t_loss.item():.3f} | "
|
||||
f"Acc: {accum_acc.item():.4f} | {dt_ms:.1f}ms"
|
||||
)
|
||||
|
||||
if step > 0 and step % val_freq == 0:
|
||||
self.evaluate(step)
|
||||
|
||||
if step > 0 and step % save_freq == 0:
|
||||
self.save_checkpoint(step)
|
||||
|
||||
self.save_checkpoint(self.max_steps, name="final_llm")
|
||||
self.writer.close()
|
||||
self.logger.info("Stage 1 LLM Training Complete.")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Train Soprano LLM Backbone")
|
||||
parser.add_argument("--config", type=Path, default=Path("config.yaml"), help="Path to config")
|
||||
args = parser.parse_args()
|
||||
|
||||
with open(args.config, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
trainer = SopranoLLMTrainer(config)
|
||||
trainer.train()
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,41 +0,0 @@
|
||||
# training/collator.py
|
||||
import torch
|
||||
import random
|
||||
|
||||
class SopranoCollator:
|
||||
def __init__(self, tokenizer, seq_len=1024):
|
||||
self.tokenizer = tokenizer
|
||||
self.seq_len = seq_len
|
||||
self.pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id
|
||||
|
||||
def pack_for_llm(self, batch):
|
||||
texts = [item["text"] for item in batch]
|
||||
encodings = self.tokenizer(texts, add_special_tokens=False, padding=False, truncation=False)
|
||||
input_ids_list = encodings["input_ids"]
|
||||
|
||||
packed_batch = []
|
||||
buffer = []
|
||||
buffer_len = 0
|
||||
random.shuffle(input_ids_list)
|
||||
|
||||
for ids in input_ids_list:
|
||||
ids = torch.tensor(ids, dtype=torch.long)
|
||||
if buffer_len + len(ids) > self.seq_len:
|
||||
full_seq = torch.cat(buffer)
|
||||
if len(full_seq) < self.seq_len + 1:
|
||||
padding = torch.full((self.seq_len + 1 - len(full_seq),), self.pad_token_id, dtype=torch.long)
|
||||
full_seq = torch.cat([full_seq, padding])
|
||||
packed_batch.append(full_seq[:self.seq_len + 1])
|
||||
buffer, buffer_len = [], 0
|
||||
buffer.append(ids)
|
||||
buffer_len += len(ids)
|
||||
|
||||
if not packed_batch: return None, None
|
||||
batch_tensor = torch.stack(packed_batch)
|
||||
return batch_tensor[:, :-1], batch_tensor[:, 1:]
|
||||
|
||||
def collate_for_decoder(self, batch):
|
||||
texts = [item["text"] for item in batch]
|
||||
wav_paths = [item["wav_path"] for item in batch]
|
||||
encodings = self.tokenizer(texts, padding=True, truncation=True, max_length=self.seq_len, return_tensors="pt", add_special_tokens=False)
|
||||
return encodings["input_ids"], wav_paths
|
||||
@@ -1,41 +0,0 @@
|
||||
# training/loss.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchaudio
|
||||
|
||||
class MultiResolutionSTFTLoss(nn.Module):
|
||||
def __init__(self, fft_sizes=[1024, 2048, 512], hop_sizes=[120, 240, 50], win_lengths=[600, 1200, 240]):
|
||||
super().__init__()
|
||||
self.fft_sizes, self.hop_sizes, self.win_lengths = fft_sizes, hop_sizes, win_lengths
|
||||
|
||||
def stft(self, x, fft_size, hop_size, win_length):
|
||||
window = torch.hann_window(win_length).to(x.device).float()
|
||||
return torch.stft(
|
||||
x.float(), n_fft=fft_size, hop_length=hop_size, win_length=win_length,
|
||||
window=window, return_complex=True, center=True
|
||||
)
|
||||
|
||||
def forward(self, x, y):
|
||||
loss = 0.0
|
||||
for n_fft, hop, win in zip(self.fft_sizes, self.hop_sizes, self.win_lengths):
|
||||
x_stft, y_stft = self.stft(x, n_fft, hop, win), self.stft(y, n_fft, hop, win)
|
||||
sc_loss = (x_stft.abs() - y_stft.abs()).norm(p="fro") / (y_stft.abs().norm(p="fro") + 1e-7)
|
||||
mag_loss = (torch.log(x_stft.abs() + 1e-7) - torch.log(y_stft.abs() + 1e-7)).abs().mean()
|
||||
loss += sc_loss + mag_loss
|
||||
return loss / len(self.fft_sizes)
|
||||
|
||||
class AudioReconstructionLoss(nn.Module):
|
||||
def __init__(self, sample_rate=32000):
|
||||
super().__init__()
|
||||
self.mel = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate, n_mels=80, n_fft=2048, hop_length=512, normalized=True)
|
||||
self.stft = MultiResolutionSTFTLoss()
|
||||
|
||||
def forward(self, pred, target):
|
||||
pred, target = pred.squeeze(1).float(), target.squeeze(1).float()
|
||||
min_len = min(pred.shape[-1], target.shape[-1])
|
||||
pred, target = pred[..., :min_len], target[..., :min_len]
|
||||
if self.mel.mel_scale.fb.device != pred.device: self.mel = self.mel.to(pred.device)
|
||||
mel_pred = torch.log(self.mel(pred) + 1e-5)
|
||||
mel_target = torch.log(self.mel(target) + 1e-5)
|
||||
loss_mel = (mel_pred - mel_target).abs().mean()
|
||||
return loss_mel + self.stft(pred, target)
|
||||
@@ -1,40 +0,0 @@
|
||||
# utils/config.py
|
||||
import yaml
|
||||
import torch
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
class Config:
|
||||
def __init__(self, config_path=None):
|
||||
if config_path is None:
|
||||
config_path = os.environ.get("SOPRANO_CONFIG", "config/settings.yaml")
|
||||
|
||||
# Resolve the absolute path to avoid "File Not Found" errors when running from subdirectories
|
||||
self.path = Path(config_path).absolute()
|
||||
|
||||
if not self.path.exists():
|
||||
# Fallback for localized execution
|
||||
self.path = Path(os.getcwd()) / Path(config_path).name
|
||||
|
||||
if not self.path.exists():
|
||||
raise FileNotFoundError(f"Config not found at: {self.path}")
|
||||
|
||||
with open(self.path, "r", encoding="utf-8") as f:
|
||||
self.cfg = yaml.safe_load(f)
|
||||
|
||||
# Shortcuts for cleaner code access
|
||||
self.common = self.cfg.get("common", {})
|
||||
self.codec = self.cfg.get("codec", {})
|
||||
self.dataset = self.cfg.get("dataset", {})
|
||||
self.training = self.cfg.get("training", {})
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
"""Resolves 'auto' to 'cuda' (GPU) if available, otherwise 'cpu'."""
|
||||
d = self.common.get("device", "auto")
|
||||
if d == "auto":
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
return d
|
||||
|
||||
# Global instance for easy importing: from utils.config import cfg
|
||||
cfg = Config()
|
||||
@@ -1,29 +0,0 @@
|
||||
# utils/text_normalizer.py
|
||||
import re
|
||||
import unicodedata
|
||||
|
||||
def normalize_text(text: str) -> str:
|
||||
"""
|
||||
Standardizes text for training and inference.
|
||||
- Lowercases everything.
|
||||
- Removes accents (e.g., é -> e).
|
||||
- Strips unknown symbols.
|
||||
- Collapses extra whitespace.
|
||||
"""
|
||||
if not isinstance(text, str):
|
||||
return ""
|
||||
|
||||
# 1. Standardize Case
|
||||
text = text.lower()
|
||||
|
||||
# 2. Decompose characters and remove non-spacing marks (accents)
|
||||
text = unicodedata.normalize('NFD', text)
|
||||
text = "".join([c for c in text if unicodedata.category(c) != 'Mn'])
|
||||
|
||||
# 3. Filter characters: keep only basic English letters, numbers, spaces, and core punctuation
|
||||
text = re.sub(r"[^a-z0-9\s.,!?'-]", "", text)
|
||||
|
||||
# 4. Collapse multiple spaces into one
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
|
||||
return text.strip()
|
||||
@@ -1,64 +0,0 @@
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
from model.encoder import Encoder
|
||||
from model.decoder import Decoder
|
||||
from train_codec import SpectralLoss
|
||||
|
||||
def verify():
|
||||
print("Verifying Decoder Pipeline...")
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Using device: {device}")
|
||||
|
||||
# Initialize models
|
||||
try:
|
||||
encoder = Encoder().to(device).float()
|
||||
decoder = Decoder(
|
||||
input_channels=5, # Bottleneck channels
|
||||
decoder_dim=512,
|
||||
decoder_layers=8
|
||||
).to(device).float()
|
||||
print("Models initialized successfully.")
|
||||
except Exception as e:
|
||||
print(f"Failed to initialize models: {e}")
|
||||
return
|
||||
|
||||
# Create dummy audio [1, 1, 32000] (1 second at 32kHz)
|
||||
audio = torch.randn(1, 1, 32000).to(device).float()
|
||||
print(f"Input audio shape: {audio.shape}")
|
||||
|
||||
# Forward pass
|
||||
try:
|
||||
# Encoder
|
||||
z = encoder(audio)
|
||||
print(f"Encoder output shape: {z.shape}")
|
||||
|
||||
# Decoder
|
||||
rec = decoder(z)
|
||||
print(f"Decoder output shape: {rec.shape}")
|
||||
|
||||
# Check shapes
|
||||
if rec.shape != audio.shape:
|
||||
print(f"WARNING: Shape mismatch! Input: {audio.shape}, Output: {rec.shape}")
|
||||
# It might be slightly off due to padding/striding, but should be close.
|
||||
# ISTFT usually returns exact length if configured right, or slightly more.
|
||||
else:
|
||||
print("Shapes match exactly!")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Forward pass failed: {e}")
|
||||
return
|
||||
|
||||
# Loss computation
|
||||
try:
|
||||
criterion = SpectralLoss().to(device)
|
||||
loss = criterion(rec, audio)
|
||||
print(f"Loss computed successfully: {loss.item()}")
|
||||
except Exception as e:
|
||||
print(f"Loss computation failed: {e}")
|
||||
return
|
||||
|
||||
print("Verification passed!")
|
||||
|
||||
if __name__ == "__main__":
|
||||
verify()
|
||||
Reference in New Issue
Block a user