mirror of
https://github.com/Nighthawk42/soprano-factory.git
synced 2026-08-30 04:30:21 +00:00
114 lines
4.1 KiB
Python
114 lines
4.1 KiB
Python
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()
|