chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,109 @@
|
||||
"""Unit checks for pick_runtime_dtype / get_dtype consistency.
|
||||
|
||||
Loads src/voxcpm/model/utils.py directly to avoid the heavy voxcpm package
|
||||
init. Run with: `python scripts/test_pick_runtime_dtype.py`.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
|
||||
UTILS = str(REPO_ROOT / "src" / "voxcpm" / "model" / "utils.py")
|
||||
spec = importlib.util.spec_from_file_location("voxcpm_utils", UTILS)
|
||||
utils = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(utils)
|
||||
|
||||
_LOW_PRECISION_DTYPES = utils._LOW_PRECISION_DTYPES
|
||||
_VALID_DTYPE_OVERRIDES = utils._VALID_DTYPE_OVERRIDES
|
||||
get_dtype = utils.get_dtype
|
||||
pick_runtime_dtype = utils.pick_runtime_dtype
|
||||
|
||||
|
||||
def expect(actual, expected, label):
|
||||
ok = actual == expected
|
||||
mark = "OK " if ok else "FAIL"
|
||||
print(f"[{mark}] {label}: got={actual!r} expected={expected!r}")
|
||||
return ok
|
||||
|
||||
|
||||
def expect_raises(fn, exc_type, label):
|
||||
try:
|
||||
fn()
|
||||
except exc_type as e:
|
||||
print(f"[OK ] {label}: raised {exc_type.__name__}: {e}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"[FAIL] {label}: raised {type(e).__name__} not {exc_type.__name__}: {e}")
|
||||
return False
|
||||
print(f"[FAIL] {label}: no exception raised")
|
||||
return False
|
||||
|
||||
|
||||
results = []
|
||||
|
||||
print("=== override set sanity ===")
|
||||
results.append(expect("half" not in _VALID_DTYPE_OVERRIDES, True, "half removed from _VALID_DTYPE_OVERRIDES"))
|
||||
results.append(expect("half" not in _LOW_PRECISION_DTYPES, True, "half removed from _LOW_PRECISION_DTYPES"))
|
||||
|
||||
print("\n=== every accepted override parses through get_dtype ===")
|
||||
for dt in sorted(_VALID_DTYPE_OVERRIDES):
|
||||
try:
|
||||
torch_dtype = get_dtype(dt)
|
||||
print(f"[OK ] get_dtype({dt!r}) -> {torch_dtype}")
|
||||
results.append(True)
|
||||
except Exception as e:
|
||||
print(f"[FAIL] get_dtype({dt!r}) raised: {e}")
|
||||
results.append(False)
|
||||
|
||||
print("\n=== pick_runtime_dtype: non-mps is a no-op ===")
|
||||
results.append(expect(pick_runtime_dtype("cuda", "bfloat16"), "bfloat16", "cuda/bf16 untouched"))
|
||||
results.append(expect(pick_runtime_dtype("cpu", "float16"), "float16", "cpu/fp16 untouched"))
|
||||
results.append(expect(pick_runtime_dtype("cuda", "float32"), "float32", "cuda/fp32 untouched"))
|
||||
|
||||
print("\n=== pick_runtime_dtype: mps forces fp32 for low-precision ===")
|
||||
os.environ.pop("VOXCPM_MPS_DTYPE", None)
|
||||
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "float32", "mps/bf16 -> fp32"))
|
||||
results.append(expect(pick_runtime_dtype("mps", "bf16"), "float32", "mps/bf16-alias -> fp32"))
|
||||
results.append(expect(pick_runtime_dtype("mps", "float16"), "float32", "mps/fp16 -> fp32"))
|
||||
results.append(expect(pick_runtime_dtype("mps", "fp16"), "float32", "mps/fp16-alias -> fp32"))
|
||||
results.append(expect(pick_runtime_dtype("mps", "float32"), "float32", "mps/fp32 stays"))
|
||||
results.append(expect(pick_runtime_dtype("mps", "fp32"), "fp32", "mps/fp32-alias stays"))
|
||||
|
||||
print("\n=== pick_runtime_dtype: VOXCPM_MPS_DTYPE override ===")
|
||||
os.environ["VOXCPM_MPS_DTYPE"] = "bfloat16"
|
||||
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "bfloat16", "override bf16 honored"))
|
||||
|
||||
os.environ["VOXCPM_MPS_DTYPE"] = "FP16"
|
||||
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "fp16", "override is case-insensitive"))
|
||||
|
||||
os.environ["VOXCPM_MPS_DTYPE"] = " float32 "
|
||||
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "float32", "override is whitespace-trimmed"))
|
||||
|
||||
print("\n=== pick_runtime_dtype: 'half' is no longer a valid override ===")
|
||||
os.environ["VOXCPM_MPS_DTYPE"] = "half"
|
||||
results.append(
|
||||
expect_raises(
|
||||
lambda: pick_runtime_dtype("mps", "bfloat16"),
|
||||
ValueError,
|
||||
"override=half now rejected (was the bug)",
|
||||
)
|
||||
)
|
||||
|
||||
os.environ["VOXCPM_MPS_DTYPE"] = "garbage"
|
||||
results.append(
|
||||
expect_raises(
|
||||
lambda: pick_runtime_dtype("mps", "bfloat16"),
|
||||
ValueError,
|
||||
"override=garbage still rejected",
|
||||
)
|
||||
)
|
||||
|
||||
os.environ.pop("VOXCPM_MPS_DTYPE", None)
|
||||
|
||||
print("\n=== summary ===")
|
||||
passed = sum(results)
|
||||
total = len(results)
|
||||
print(f"{passed}/{total} passed")
|
||||
sys.exit(0 if passed == total else 1)
|
||||
@@ -0,0 +1,146 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Full finetune inference script (no LoRA).
|
||||
|
||||
Checkpoint directory contains complete model files (pytorch_model.bin, config.json, audiovae.pth, etc.),
|
||||
can be loaded directly via VoxCPM.
|
||||
|
||||
Usage:
|
||||
|
||||
python scripts/test_voxcpm_ft_infer.py \
|
||||
--ckpt_dir /path/to/checkpoints/step_0001000 \
|
||||
--text "Hello, I am the finetuned VoxCPM." \
|
||||
--output ft_test.wav
|
||||
|
||||
With voice cloning:
|
||||
|
||||
python scripts/test_voxcpm_ft_infer.py \
|
||||
--ckpt_dir /path/to/checkpoints/step_0001000 \
|
||||
--text "Hello, this is voice cloning result." \
|
||||
--prompt_audio path/to/ref.wav \
|
||||
--prompt_text "Reference audio transcript" \
|
||||
--seed 42 \
|
||||
--output ft_clone.wav
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
|
||||
from voxcpm.core import VoxCPM
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser("VoxCPM full-finetune inference test (no LoRA)")
|
||||
parser.add_argument(
|
||||
"--ckpt_dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Checkpoint directory (contains pytorch_model.bin, config.json, audiovae.pth, etc.)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Target text to synthesize",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt_audio",
|
||||
type=str,
|
||||
default="",
|
||||
help="Optional: reference audio path for voice cloning",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt_text",
|
||||
type=str,
|
||||
default="",
|
||||
help="Optional: transcript of reference audio",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default="ft_test.wav",
|
||||
help="Output wav file path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg_value",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="CFG scale (default: 2.0)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inference_timesteps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Diffusion inference steps (default: 10)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_len",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Max generation steps",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--normalize",
|
||||
action="store_true",
|
||||
help="Enable text normalization",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for generation (default: None)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
# Load model from checkpoint directory (no denoiser)
|
||||
print(f"[FT Inference] Loading model: {args.ckpt_dir}", file=sys.stderr)
|
||||
model = VoxCPM.from_pretrained(
|
||||
hf_model_id=args.ckpt_dir,
|
||||
load_denoiser=False,
|
||||
optimize=True,
|
||||
)
|
||||
|
||||
# Run inference
|
||||
prompt_wav_path = args.prompt_audio if args.prompt_audio else None
|
||||
prompt_text = args.prompt_text if args.prompt_text else None
|
||||
|
||||
print(f"[FT Inference] Synthesizing: text='{args.text}'", file=sys.stderr)
|
||||
if prompt_wav_path:
|
||||
print(f"[FT Inference] Using reference audio: {prompt_wav_path}", file=sys.stderr)
|
||||
print(f"[FT Inference] Reference text: {prompt_text}", file=sys.stderr)
|
||||
|
||||
if args.seed is not None:
|
||||
print(f"[FT Inference] Using seed: {args.seed}", file=sys.stderr)
|
||||
|
||||
audio_np = model.generate(
|
||||
text=args.text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=args.cfg_value,
|
||||
inference_timesteps=args.inference_timesteps,
|
||||
max_len=args.max_len,
|
||||
normalize=args.normalize,
|
||||
denoise=False,
|
||||
seed=args.seed,
|
||||
)
|
||||
|
||||
# Save audio
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
sf.write(str(out_path), audio_np, model.tts_model.sample_rate)
|
||||
|
||||
print(
|
||||
f"[FT Inference] Saved to: {out_path}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,276 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
LoRA inference test script.
|
||||
|
||||
Usage:
|
||||
|
||||
python scripts/test_voxcpm_lora_infer.py \
|
||||
--lora_ckpt checkpoints/step_0002000 \
|
||||
--text "Hello, this is LoRA finetuned result." \
|
||||
--output lora_test.wav
|
||||
|
||||
With voice cloning:
|
||||
|
||||
python scripts/test_voxcpm_lora_infer.py \
|
||||
--lora_ckpt checkpoints/step_0002000 \
|
||||
--text "This is voice cloning result." \
|
||||
--prompt_audio path/to/ref.wav \
|
||||
--prompt_text "Reference audio transcript" \
|
||||
--seed 42 \
|
||||
--output lora_clone.wav
|
||||
|
||||
Note: The script reads base_model path and lora_config from lora_config.json
|
||||
in the checkpoint directory (saved automatically during training).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
|
||||
from voxcpm.core import VoxCPM
|
||||
from voxcpm.model.voxcpm import LoRAConfig
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser("VoxCPM LoRA inference test")
|
||||
parser.add_argument(
|
||||
"--lora_ckpt",
|
||||
type=str,
|
||||
required=True,
|
||||
help="LoRA checkpoint directory (contains lora_weights.safetensors and lora_config.json)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--base_model",
|
||||
type=str,
|
||||
default="",
|
||||
help="Optional: override base model path (default: read from lora_config.json)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Target text to synthesize",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt_audio",
|
||||
type=str,
|
||||
default="",
|
||||
help="Optional: reference audio path for voice cloning",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt_text",
|
||||
type=str,
|
||||
default="",
|
||||
help="Optional: transcript of reference audio",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default="lora_test.wav",
|
||||
help="Output wav file path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg_value",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="CFG scale (default: 2.0)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inference_timesteps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Diffusion inference steps (default: 10)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_len",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Max generation steps",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--normalize",
|
||||
action="store_true",
|
||||
help="Enable text normalization",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for generation (default: None)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
# 1. Check LoRA checkpoint directory
|
||||
ckpt_dir = Path(args.lora_ckpt)
|
||||
if not ckpt_dir.exists():
|
||||
raise FileNotFoundError(f"LoRA checkpoint not found: {ckpt_dir}")
|
||||
|
||||
# 2. Load lora_config.json from checkpoint
|
||||
lora_config_path = ckpt_dir / "lora_config.json"
|
||||
if not lora_config_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"lora_config.json not found in {ckpt_dir}. "
|
||||
"Make sure the checkpoint was saved with the updated training script."
|
||||
)
|
||||
|
||||
with open(lora_config_path, "r", encoding="utf-8") as f:
|
||||
lora_info = json.load(f)
|
||||
|
||||
# Get base model path (command line arg overrides config)
|
||||
pretrained_path = args.base_model if args.base_model else lora_info.get("base_model")
|
||||
if not pretrained_path:
|
||||
raise ValueError("base_model not found in lora_config.json and --base_model not provided")
|
||||
|
||||
# Get LoRA config
|
||||
lora_cfg_dict = lora_info.get("lora_config", {})
|
||||
lora_cfg = LoRAConfig(**lora_cfg_dict) if lora_cfg_dict else None
|
||||
|
||||
print(f"Loaded config from: {lora_config_path}", file=sys.stderr)
|
||||
print(f" Base model: {pretrained_path}", file=sys.stderr)
|
||||
print(
|
||||
f" LoRA config: r={lora_cfg.r}, alpha={lora_cfg.alpha}" if lora_cfg else " LoRA config: None", file=sys.stderr
|
||||
)
|
||||
if args.seed is not None:
|
||||
print(f" Seed: {args.seed}", file=sys.stderr)
|
||||
|
||||
# 3. Load model with LoRA (no denoiser)
|
||||
print(f"\n[1/2] Loading model with LoRA: {pretrained_path}", file=sys.stderr)
|
||||
print(f" LoRA weights: {ckpt_dir}", file=sys.stderr)
|
||||
model = VoxCPM.from_pretrained(
|
||||
hf_model_id=pretrained_path,
|
||||
load_denoiser=False,
|
||||
optimize=True,
|
||||
lora_config=lora_cfg,
|
||||
lora_weights_path=str(ckpt_dir),
|
||||
)
|
||||
|
||||
# 4. Synthesize audio
|
||||
prompt_wav_path = args.prompt_audio if args.prompt_audio else None
|
||||
prompt_text = args.prompt_text if args.prompt_text else None
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print("\n[2/2] Starting synthesis tests...", file=sys.stderr)
|
||||
|
||||
# === Test 1: With LoRA ===
|
||||
print("\n [Test 1] Synthesize with LoRA...", file=sys.stderr)
|
||||
audio_np = model.generate(
|
||||
text=args.text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=args.cfg_value,
|
||||
inference_timesteps=args.inference_timesteps,
|
||||
max_len=args.max_len,
|
||||
normalize=args.normalize,
|
||||
denoise=False,
|
||||
seed=args.seed,
|
||||
)
|
||||
lora_output = out_path.with_stem(out_path.stem + "_with_lora")
|
||||
sf.write(str(lora_output), audio_np, model.tts_model.sample_rate)
|
||||
print(
|
||||
f" Saved: {lora_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# === Test 2: Disable LoRA (via set_lora_enabled) ===
|
||||
print("\n [Test 2] Disable LoRA (set_lora_enabled=False)...", file=sys.stderr)
|
||||
model.set_lora_enabled(False)
|
||||
audio_np = model.generate(
|
||||
text=args.text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=args.cfg_value,
|
||||
inference_timesteps=args.inference_timesteps,
|
||||
max_len=args.max_len,
|
||||
normalize=args.normalize,
|
||||
denoise=False,
|
||||
seed=args.seed,
|
||||
)
|
||||
disabled_output = out_path.with_stem(out_path.stem + "_lora_disabled")
|
||||
sf.write(str(disabled_output), audio_np, model.tts_model.sample_rate)
|
||||
print(
|
||||
f" Saved: {disabled_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# === Test 3: Re-enable LoRA ===
|
||||
print("\n [Test 3] Re-enable LoRA (set_lora_enabled=True)...", file=sys.stderr)
|
||||
model.set_lora_enabled(True)
|
||||
audio_np = model.generate(
|
||||
text=args.text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=args.cfg_value,
|
||||
inference_timesteps=args.inference_timesteps,
|
||||
max_len=args.max_len,
|
||||
normalize=args.normalize,
|
||||
denoise=False,
|
||||
seed=args.seed,
|
||||
)
|
||||
reenabled_output = out_path.with_stem(out_path.stem + "_lora_reenabled")
|
||||
sf.write(str(reenabled_output), audio_np, model.tts_model.sample_rate)
|
||||
print(
|
||||
f" Saved: {reenabled_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# === Test 4: Unload LoRA (reset_lora_weights) ===
|
||||
print("\n [Test 4] Unload LoRA (unload_lora)...", file=sys.stderr)
|
||||
model.unload_lora()
|
||||
audio_np = model.generate(
|
||||
text=args.text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=args.cfg_value,
|
||||
inference_timesteps=args.inference_timesteps,
|
||||
max_len=args.max_len,
|
||||
normalize=args.normalize,
|
||||
denoise=False,
|
||||
seed=args.seed,
|
||||
)
|
||||
reset_output = out_path.with_stem(out_path.stem + "_lora_reset")
|
||||
sf.write(str(reset_output), audio_np, model.tts_model.sample_rate)
|
||||
print(
|
||||
f" Saved: {reset_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# === Test 5: Hot-reload LoRA (load_lora) ===
|
||||
print("\n [Test 5] Hot-reload LoRA (load_lora)...", file=sys.stderr)
|
||||
loaded, skipped = model.load_lora(ckpt_dir)
|
||||
print(f" Reloaded {len(loaded)} parameters", file=sys.stderr)
|
||||
audio_np = model.generate(
|
||||
text=args.text,
|
||||
prompt_wav_path=prompt_wav_path,
|
||||
prompt_text=prompt_text,
|
||||
cfg_value=args.cfg_value,
|
||||
inference_timesteps=args.inference_timesteps,
|
||||
max_len=args.max_len,
|
||||
normalize=args.normalize,
|
||||
denoise=False,
|
||||
seed=args.seed,
|
||||
)
|
||||
reload_output = out_path.with_stem(out_path.stem + "_lora_reloaded")
|
||||
sf.write(str(reload_output), audio_np, model.tts_model.sample_rate)
|
||||
print(
|
||||
f" Saved: {reload_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
print("\n[Done] All tests completed!", file=sys.stderr)
|
||||
print(f" - with_lora: {lora_output}", file=sys.stderr)
|
||||
print(f" - lora_disabled: {disabled_output}", file=sys.stderr)
|
||||
print(f" - lora_reenabled: {reenabled_output}", file=sys.stderr)
|
||||
print(f" - lora_reset: {reset_output}", file=sys.stderr)
|
||||
print(f" - lora_reloaded: {reload_output}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,841 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
project_root = Path(__file__).parent.parent
|
||||
sys.path.insert(0, str(project_root / "src"))
|
||||
|
||||
import contextlib
|
||||
from typing import Dict
|
||||
|
||||
import argbind
|
||||
import torch
|
||||
from tensorboardX import SummaryWriter
|
||||
from torch.optim import AdamW
|
||||
from transformers import get_cosine_schedule_with_warmup
|
||||
import signal
|
||||
import os
|
||||
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
try:
|
||||
from safetensors.torch import save_file
|
||||
|
||||
SAFETENSORS_AVAILABLE = True
|
||||
except ImportError:
|
||||
SAFETENSORS_AVAILABLE = False
|
||||
print("Warning: safetensors not available, will use pytorch format", file=sys.stderr)
|
||||
|
||||
import json
|
||||
|
||||
from voxcpm.model import VoxCPMModel, VoxCPM2Model
|
||||
from voxcpm.model.voxcpm import LoRAConfig as LoRAConfigV1
|
||||
from voxcpm.model.voxcpm2 import LoRAConfig as LoRAConfigV2
|
||||
from voxcpm.training import (
|
||||
Accelerator,
|
||||
BatchProcessor,
|
||||
TrainingTracker,
|
||||
build_dataloader,
|
||||
load_audio_text_datasets,
|
||||
)
|
||||
|
||||
|
||||
@argbind.bind(without_prefix=True)
|
||||
def train(
|
||||
pretrained_path: str,
|
||||
train_manifest: str,
|
||||
val_manifest: str = "",
|
||||
sample_rate: int = 16_000,
|
||||
out_sample_rate: int = 0, # AudioVAE decoder output rate; used for TensorBoard audio logging
|
||||
batch_size: int = 1,
|
||||
grad_accum_steps: int = 1,
|
||||
num_workers: int = 2,
|
||||
num_iters: int = 100_000,
|
||||
log_interval: int = 100,
|
||||
valid_interval: int = 1_000,
|
||||
save_interval: int = 10_000,
|
||||
learning_rate: float = 1e-4,
|
||||
weight_decay: float = 1e-2,
|
||||
warmup_steps: int = 1_000,
|
||||
max_steps: int = 100_000,
|
||||
max_batch_tokens: int = 0,
|
||||
save_path: str = "checkpoints",
|
||||
tensorboard: str = "",
|
||||
lambdas: Dict[str, float] = {"loss/diff": 1.0, "loss/stop": 1.0},
|
||||
lora: dict = None,
|
||||
config_path: str = "",
|
||||
max_grad_norm: float = 0.0, # gradient clipping; 0 = disabled (backward compat)
|
||||
# Distribution options (for LoRA checkpoints)
|
||||
hf_model_id: str = "", # HuggingFace model ID (e.g., "openbmb/VoxCPM1.5")
|
||||
distribute: bool = False, # If True, save hf_model_id as base_model; otherwise save pretrained_path
|
||||
):
|
||||
_ = config_path
|
||||
|
||||
# Validate distribution options
|
||||
if lora is not None and distribute and not hf_model_id:
|
||||
raise ValueError("hf_model_id is required when distribute=True")
|
||||
|
||||
accelerator = Accelerator(amp=True)
|
||||
|
||||
save_dir = Path(save_path)
|
||||
tb_dir = Path(tensorboard) if tensorboard else save_dir / "logs"
|
||||
|
||||
# Only create directories on rank 0 to avoid race conditions
|
||||
if accelerator.rank == 0:
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
tb_dir.mkdir(parents=True, exist_ok=True)
|
||||
accelerator.barrier() # Wait for directory creation
|
||||
|
||||
writer = SummaryWriter(log_dir=str(tb_dir)) if accelerator.rank == 0 else None
|
||||
tracker = TrainingTracker(writer=writer, log_file=str(save_dir / "train.log"), rank=accelerator.rank)
|
||||
|
||||
# Auto-detect model architecture from config.json
|
||||
with open(os.path.join(pretrained_path, "config.json"), "r", encoding="utf-8") as _f:
|
||||
_arch = json.load(_f).get("architecture", "voxcpm").lower()
|
||||
_model_cls = VoxCPM2Model if _arch == "voxcpm2" else VoxCPMModel
|
||||
LoRAConfig = LoRAConfigV2 if _arch == "voxcpm2" else LoRAConfigV1
|
||||
if accelerator.rank == 0:
|
||||
print(f"Detected architecture: {_arch} -> {_model_cls.__name__}", file=sys.stderr)
|
||||
base_model = _model_cls.from_local(
|
||||
pretrained_path, optimize=False, training=True, lora_config=LoRAConfig(**lora) if lora else None
|
||||
)
|
||||
tokenizer = base_model.text_tokenizer
|
||||
|
||||
expected_sr = base_model.audio_vae.sample_rate
|
||||
assert sample_rate == expected_sr, (
|
||||
f"sample_rate mismatch: config says {sample_rate}, but the AudioVAE encoder expects {expected_sr}. "
|
||||
f"Please set sample_rate: {expected_sr} in your training config. "
|
||||
)
|
||||
|
||||
train_ds, val_ds = load_audio_text_datasets(
|
||||
train_manifest=train_manifest,
|
||||
val_manifest=val_manifest,
|
||||
sample_rate=sample_rate,
|
||||
)
|
||||
|
||||
def tokenize(batch):
|
||||
text_list = batch["text"]
|
||||
text_ids = [tokenizer(text) for text in text_list]
|
||||
return {"text_ids": text_ids}
|
||||
|
||||
train_ds = train_ds.map(tokenize, batched=True, remove_columns=["text"])
|
||||
# Save original validation texts for audio generation display
|
||||
val_texts = None
|
||||
if val_ds is not None:
|
||||
val_texts = list(val_ds["text"]) # Save original texts
|
||||
val_ds = val_ds.map(tokenize, batched=True, remove_columns=["text"])
|
||||
|
||||
dataset_cnt = int(max(train_ds["dataset_id"])) + 1 if "dataset_id" in train_ds.column_names else 1
|
||||
num_train_samples = len(train_ds)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Optional: filter samples by estimated token count to avoid OOM
|
||||
# Enabled when max_batch_tokens > 0:
|
||||
# max_sample_len = max_batch_tokens // batch_size
|
||||
# Samples exceeding this length will be dropped
|
||||
# ------------------------------------------------------------------ #
|
||||
if max_batch_tokens and max_batch_tokens > 0:
|
||||
from voxcpm.training.data import compute_sample_lengths
|
||||
|
||||
audio_vae_fps = base_model.audio_vae.sample_rate / base_model.audio_vae.hop_length
|
||||
est_lengths = compute_sample_lengths(
|
||||
train_ds,
|
||||
audio_vae_fps=audio_vae_fps,
|
||||
patch_size=base_model.config.patch_size,
|
||||
)
|
||||
max_sample_len = max_batch_tokens // batch_size if batch_size > 0 else max(est_lengths)
|
||||
keep_indices = [i for i, L in enumerate(est_lengths) if L <= max_sample_len]
|
||||
|
||||
if len(keep_indices) < len(train_ds) and accelerator.rank == 0:
|
||||
tracker.print(
|
||||
f"Filtering {len(train_ds) - len(keep_indices)} / {len(train_ds)} "
|
||||
f"training samples longer than {max_sample_len} tokens "
|
||||
f"(max_batch_tokens={max_batch_tokens})."
|
||||
)
|
||||
train_ds = train_ds.select(keep_indices)
|
||||
|
||||
train_loader = build_dataloader(
|
||||
train_ds,
|
||||
accelerator=accelerator,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers,
|
||||
drop_last=True,
|
||||
)
|
||||
val_loader = (
|
||||
build_dataloader(
|
||||
val_ds,
|
||||
accelerator=accelerator,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers,
|
||||
drop_last=False,
|
||||
)
|
||||
if val_ds is not None
|
||||
else None
|
||||
)
|
||||
|
||||
batch_processor = BatchProcessor(
|
||||
config=base_model.config,
|
||||
audio_vae=base_model.audio_vae,
|
||||
dataset_cnt=dataset_cnt,
|
||||
device=accelerator.device,
|
||||
)
|
||||
# Save audio_vae and output sample rate for audio generation.
|
||||
# Prefer model's actual output rate; fall back to YAML out_sample_rate or encode rate.
|
||||
audio_vae_for_gen = base_model.audio_vae
|
||||
out_sr = base_model.sample_rate # decoder output rate (e.g. 48000 for V2)
|
||||
if out_sr == 0 and out_sample_rate > 0:
|
||||
out_sr = out_sample_rate
|
||||
del base_model.audio_vae
|
||||
model = accelerator.prepare_model(base_model)
|
||||
unwrapped_model = accelerator.unwrap(model)
|
||||
unwrapped_model.train()
|
||||
|
||||
# Only print param info on rank 0 to avoid cluttered output
|
||||
if accelerator.rank == 0:
|
||||
for name, param in model.named_parameters():
|
||||
print(name, param.requires_grad, file=sys.stderr)
|
||||
|
||||
optimizer = AdamW(
|
||||
(p for p in model.parameters() if p.requires_grad),
|
||||
lr=learning_rate,
|
||||
weight_decay=weight_decay,
|
||||
)
|
||||
|
||||
# Cosine + warmup scheduler from transformers:
|
||||
# - num_warmup_steps: warmup steps
|
||||
# - num_training_steps: total training steps (outer step count)
|
||||
total_training_steps = max_steps if max_steps > 0 else num_iters
|
||||
scheduler = get_cosine_schedule_with_warmup(
|
||||
optimizer,
|
||||
num_warmup_steps=warmup_steps,
|
||||
num_training_steps=total_training_steps,
|
||||
)
|
||||
|
||||
# All ranks load the same checkpoint to keep model and optimizer state in sync.
|
||||
start_step = load_checkpoint(model, optimizer, scheduler, save_dir, rank=accelerator.rank)
|
||||
accelerator.barrier()
|
||||
|
||||
if start_step > 0 and accelerator.rank == 0:
|
||||
tracker.print(f"Resuming training from step {start_step}")
|
||||
|
||||
# Resume tracker for signal handler to read current step
|
||||
resume = {"step": start_step}
|
||||
|
||||
# Register signal handler to save checkpoint on termination (SIGTERM/SIGINT)
|
||||
def _signal_handler(
|
||||
signum,
|
||||
frame,
|
||||
_model=model,
|
||||
_optim=optimizer,
|
||||
_sched=scheduler,
|
||||
_save_dir=save_dir,
|
||||
_pretrained=pretrained_path,
|
||||
_hf_id=hf_model_id,
|
||||
_dist=distribute,
|
||||
_resume=resume,
|
||||
_rank=accelerator.rank,
|
||||
):
|
||||
try:
|
||||
cur_step = int(_resume.get("step", start_step))
|
||||
except Exception:
|
||||
cur_step = start_step
|
||||
if _rank == 0:
|
||||
print(f"Signal {signum} received. Saving checkpoint at step {cur_step} ...", file=sys.stderr)
|
||||
try:
|
||||
save_checkpoint(_model, _optim, _sched, _save_dir, cur_step, _pretrained, _hf_id, _dist)
|
||||
print("Checkpoint saved. Exiting.", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"Error saving checkpoint on signal: {e}", file=sys.stderr)
|
||||
os._exit(0)
|
||||
|
||||
signal.signal(signal.SIGTERM, _signal_handler)
|
||||
signal.signal(signal.SIGINT, _signal_handler)
|
||||
|
||||
# Manual epoch management instead of itertools.cycle to support DistributedSampler.set_epoch()
|
||||
grad_accum_steps = max(int(grad_accum_steps), 1)
|
||||
data_epoch = 0
|
||||
train_iter = iter(train_loader)
|
||||
|
||||
def get_next_batch():
|
||||
"""Get next batch, handles epoch boundary and DistributedSampler."""
|
||||
nonlocal train_iter, data_epoch
|
||||
try:
|
||||
return next(train_iter)
|
||||
except StopIteration:
|
||||
data_epoch += 1
|
||||
# Key: set DistributedSampler epoch to ensure different data order each epoch
|
||||
sampler = getattr(train_loader, "sampler", None)
|
||||
if hasattr(sampler, "set_epoch"):
|
||||
sampler.set_epoch(data_epoch)
|
||||
train_iter = iter(train_loader)
|
||||
return next(train_iter)
|
||||
|
||||
with tracker.live():
|
||||
for step in range(start_step, num_iters):
|
||||
# update resume step so signal handler can save current progress
|
||||
resume["step"] = step
|
||||
tracker.step = step
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
# Gradient accumulation: accumulate gradients over micro-batches before optimizer step
|
||||
loss_dict = {}
|
||||
for micro_step in range(grad_accum_steps):
|
||||
batch = get_next_batch()
|
||||
processed = batch_processor(batch)
|
||||
|
||||
# Only sync gradients on the last micro-batch
|
||||
# Use no_sync() for intermediate steps to reduce communication overhead
|
||||
is_last_micro_step = micro_step == grad_accum_steps - 1
|
||||
sync_context = contextlib.nullcontext() if is_last_micro_step else accelerator.no_sync()
|
||||
|
||||
with sync_context:
|
||||
with accelerator.autocast(dtype=torch.bfloat16):
|
||||
outputs = model(
|
||||
processed["text_tokens"],
|
||||
processed["text_mask"],
|
||||
processed["audio_feats"],
|
||||
processed["audio_mask"],
|
||||
processed["loss_mask"],
|
||||
processed["position_ids"],
|
||||
processed["labels"],
|
||||
progress=step / max(1, num_iters),
|
||||
)
|
||||
|
||||
total_loss = 0.0
|
||||
for key, value in outputs.items():
|
||||
if key.startswith("loss/"):
|
||||
weight = lambdas.get(key, 1.0)
|
||||
loss_value = value * weight / grad_accum_steps
|
||||
total_loss = total_loss + loss_value
|
||||
# Record raw loss from last micro-batch for logging
|
||||
loss_dict[key] = value.detach()
|
||||
|
||||
# Accumulate gradients (normalized by grad_accum_steps)
|
||||
accelerator.backward(total_loss)
|
||||
|
||||
# After all micro-batches, do unscale / grad_norm / step
|
||||
scaler = getattr(accelerator, "scaler", None)
|
||||
if scaler is not None:
|
||||
scaler.unscale_(optimizer)
|
||||
effective_max_norm = max_grad_norm if max_grad_norm > 0 else 1e9
|
||||
grad_norm = torch.nn.utils.clip_grad_norm_(unwrapped_model.parameters(), max_norm=effective_max_norm)
|
||||
|
||||
accelerator.step(optimizer)
|
||||
accelerator.update()
|
||||
scheduler.step()
|
||||
|
||||
if step % log_interval == 0 or step == num_iters - 1:
|
||||
loss_values = {k: v.item() if isinstance(v, torch.Tensor) else float(v) for k, v in loss_dict.items()}
|
||||
loss_values["lr"] = float(optimizer.param_groups[0]["lr"])
|
||||
# Account for all GPUs when converting steps to epochs.
|
||||
epoch = (step * grad_accum_steps * batch_size * accelerator.world_size) / max(1, num_train_samples)
|
||||
loss_values["epoch"] = float(epoch)
|
||||
loss_values["grad_norm"] = float(grad_norm)
|
||||
tracker.log_metrics(loss_values, split="train")
|
||||
|
||||
if val_loader is not None and (step % valid_interval == 0 or step == num_iters - 1):
|
||||
validate(
|
||||
model,
|
||||
val_loader,
|
||||
batch_processor,
|
||||
accelerator,
|
||||
tracker,
|
||||
lambdas,
|
||||
writer=writer,
|
||||
step=step,
|
||||
val_ds=val_ds,
|
||||
audio_vae=audio_vae_for_gen,
|
||||
sample_rate=sample_rate,
|
||||
out_sample_rate=out_sr,
|
||||
val_texts=val_texts,
|
||||
tokenizer=tokenizer,
|
||||
valid_interval=valid_interval,
|
||||
)
|
||||
|
||||
if (step % save_interval == 0 or step == num_iters - 1) and accelerator.rank == 0:
|
||||
save_checkpoint(model, optimizer, scheduler, save_dir, step, pretrained_path, hf_model_id, distribute)
|
||||
|
||||
if accelerator.rank == 0:
|
||||
save_checkpoint(model, optimizer, scheduler, save_dir, num_iters, pretrained_path, hf_model_id, distribute)
|
||||
if writer:
|
||||
writer.close()
|
||||
|
||||
|
||||
def validate(
|
||||
model,
|
||||
val_loader,
|
||||
batch_processor,
|
||||
accelerator,
|
||||
tracker,
|
||||
lambdas,
|
||||
writer=None,
|
||||
step=0,
|
||||
val_ds=None,
|
||||
audio_vae=None,
|
||||
sample_rate=22050,
|
||||
out_sample_rate=0,
|
||||
val_texts=None,
|
||||
tokenizer=None,
|
||||
valid_interval=1000,
|
||||
):
|
||||
"""Validate and generate sample audio"""
|
||||
import numpy as np # noqa: F401
|
||||
from collections import defaultdict
|
||||
|
||||
model.eval()
|
||||
total_losses = []
|
||||
sub_losses = defaultdict(list) # Track individual sub-losses
|
||||
num_batches = 0
|
||||
max_val_batches = 10
|
||||
|
||||
with torch.no_grad():
|
||||
for batch in val_loader:
|
||||
if num_batches >= max_val_batches:
|
||||
break
|
||||
processed = batch_processor(batch)
|
||||
with accelerator.autocast(dtype=torch.bfloat16):
|
||||
outputs = model(
|
||||
processed["text_tokens"],
|
||||
processed["text_mask"],
|
||||
processed["audio_feats"],
|
||||
processed["audio_mask"],
|
||||
processed["loss_mask"],
|
||||
processed["position_ids"],
|
||||
processed["labels"],
|
||||
progress=0.0,
|
||||
sample_generate=False,
|
||||
)
|
||||
total = 0.0
|
||||
for key, value in outputs.items():
|
||||
if key.startswith("loss/"):
|
||||
weighted_loss = lambdas.get(key, 1.0) * value
|
||||
total += weighted_loss
|
||||
sub_losses[key].append(value.detach())
|
||||
total_losses.append(total.detach())
|
||||
num_batches += 1
|
||||
|
||||
if total_losses:
|
||||
# Compute mean total loss
|
||||
mean_total_loss = torch.stack(total_losses).mean()
|
||||
accelerator.all_reduce(mean_total_loss)
|
||||
|
||||
# Compute mean of each sub-loss
|
||||
val_metrics = {"loss/total": mean_total_loss.item()}
|
||||
for key, values in sub_losses.items():
|
||||
mean_sub_loss = torch.stack(values).mean()
|
||||
accelerator.all_reduce(mean_sub_loss)
|
||||
val_metrics[key] = mean_sub_loss.item()
|
||||
|
||||
tracker.log_metrics(val_metrics, split="val")
|
||||
|
||||
# Generate sample audio for TensorBoard display
|
||||
if writer is not None and val_ds is not None and audio_vae is not None and accelerator.rank == 0:
|
||||
try:
|
||||
generate_sample_audio(
|
||||
model,
|
||||
val_ds,
|
||||
audio_vae,
|
||||
writer,
|
||||
step,
|
||||
accelerator,
|
||||
sample_rate,
|
||||
out_sample_rate=out_sample_rate,
|
||||
val_texts=val_texts,
|
||||
tokenizer=tokenizer,
|
||||
valid_interval=valid_interval,
|
||||
tracker=tracker,
|
||||
)
|
||||
except Exception as e:
|
||||
tracker.print(f"[Warning] Failed to generate sample audio: {e}")
|
||||
import traceback
|
||||
import io
|
||||
|
||||
buf = io.StringIO()
|
||||
traceback.print_exc(file=buf)
|
||||
tracker.print(buf.getvalue())
|
||||
else:
|
||||
# Log why audio generation was skipped
|
||||
missing = []
|
||||
if writer is None:
|
||||
missing.append("writer")
|
||||
if val_ds is None:
|
||||
missing.append("val_ds")
|
||||
if audio_vae is None:
|
||||
missing.append("audio_vae")
|
||||
if missing and accelerator.rank == 0:
|
||||
tracker.print(f"[Warning] Skip audio generation: missing {', '.join(missing)}")
|
||||
|
||||
model.train()
|
||||
|
||||
|
||||
def compute_mel_spectrogram(audio_np, sample_rate, n_mels=128):
|
||||
"""Compute Mel Spectrogram (dB) using librosa"""
|
||||
import numpy as np
|
||||
import librosa
|
||||
|
||||
audio_np = audio_np.flatten().astype(np.float32)
|
||||
mel = librosa.feature.melspectrogram(y=audio_np, sr=sample_rate, n_mels=n_mels, fmax=sample_rate // 2)
|
||||
return librosa.power_to_db(mel, ref=np.max)
|
||||
|
||||
|
||||
def create_mel_figure(gen_audio_np, gen_mel, sample_rate, step=None, ref_audio_np=None, ref_mel=None):
|
||||
"""
|
||||
Create mel spectrogram figure: show comparison if reference audio exists, otherwise show generated only
|
||||
"""
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import librosa.display
|
||||
|
||||
fmax = sample_rate // 2
|
||||
step_str = f" @ Step {step}" if step is not None else ""
|
||||
|
||||
if ref_audio_np is not None and ref_mel is not None:
|
||||
# Comparison mode: reference vs generated
|
||||
fig, (ax_ref, ax_gen) = plt.subplots(2, 1, figsize=(12, 8))
|
||||
|
||||
img_ref = librosa.display.specshow(
|
||||
ref_mel, sr=sample_rate, x_axis="time", y_axis="mel", fmax=fmax, cmap="viridis", ax=ax_ref
|
||||
)
|
||||
ax_ref.set_title(
|
||||
f"Reference (GT) - {len(ref_audio_np)/sample_rate:.2f}s{step_str}",
|
||||
fontsize=10,
|
||||
fontweight="bold",
|
||||
color="#28A745",
|
||||
)
|
||||
plt.colorbar(img_ref, ax=ax_ref, format="%+2.0f dB", pad=0.02)
|
||||
|
||||
img_gen = librosa.display.specshow(
|
||||
gen_mel, sr=sample_rate, x_axis="time", y_axis="mel", fmax=fmax, cmap="viridis", ax=ax_gen
|
||||
)
|
||||
ax_gen.set_title(
|
||||
f"Generated - {len(gen_audio_np)/sample_rate:.2f}s", fontsize=10, fontweight="bold", color="#DC3545"
|
||||
)
|
||||
plt.colorbar(img_gen, ax=ax_gen, format="%+2.0f dB", pad=0.02)
|
||||
else:
|
||||
# Single figure mode: show generated only
|
||||
fig, ax = plt.subplots(figsize=(12, 4))
|
||||
img = librosa.display.specshow(
|
||||
gen_mel, sr=sample_rate, x_axis="time", y_axis="mel", fmax=fmax, cmap="viridis", ax=ax
|
||||
)
|
||||
ax.set_title(f"Generated - {len(gen_audio_np)/sample_rate:.2f}s{step_str}", fontsize=11, fontweight="bold")
|
||||
plt.colorbar(img, ax=ax, format="%+2.0f dB", pad=0.02)
|
||||
|
||||
plt.tight_layout()
|
||||
return fig
|
||||
|
||||
|
||||
def normalize_audio(audio_np):
|
||||
"""Normalize audio to [-0.9, 0.9]"""
|
||||
import numpy as np
|
||||
|
||||
max_val = np.abs(audio_np).max()
|
||||
return audio_np / max_val * 0.9 if max_val > 0 else audio_np
|
||||
|
||||
|
||||
def generate_sample_audio(
|
||||
model,
|
||||
val_ds,
|
||||
audio_vae,
|
||||
writer,
|
||||
step,
|
||||
accelerator,
|
||||
sample_rate=22050,
|
||||
out_sample_rate=0,
|
||||
val_texts=None,
|
||||
tokenizer=None,
|
||||
pretrained_path=None,
|
||||
valid_interval=1000,
|
||||
tracker=None,
|
||||
):
|
||||
"""Select 2 fixed validation samples, generate audio and log to TensorBoard"""
|
||||
import numpy as np
|
||||
|
||||
log = tracker.print if tracker else print
|
||||
num_samples = min(2, len(val_ds))
|
||||
log(f"[Audio] Starting audio generation for {num_samples} samples at step {step}")
|
||||
|
||||
unwrapped_model = accelerator.unwrap(model)
|
||||
# Determine the correct output sample rate for generated audio.
|
||||
# out_sample_rate is the decoder output rate (e.g. 48kHz for V2);
|
||||
# sample_rate is the encoder input rate (e.g. 16kHz for V2).
|
||||
gen_sr = out_sample_rate if out_sample_rate > 0 else sample_rate
|
||||
|
||||
for i in range(num_samples):
|
||||
sample = val_ds[i]
|
||||
text = val_texts[i] if val_texts and i < len(val_texts) else "Hello, this is a test."
|
||||
|
||||
# Load reference audio
|
||||
ref_audio_np = None
|
||||
try:
|
||||
if "audio" in sample and isinstance(sample["audio"], dict) and "array" in sample["audio"]:
|
||||
ref_audio_np = np.array(sample["audio"]["array"], dtype=np.float32)
|
||||
ref_sr = sample["audio"].get("sampling_rate", sample_rate)
|
||||
if ref_sr != sample_rate:
|
||||
import torchaudio.functional as F
|
||||
|
||||
ref_audio_np = (
|
||||
F.resample(torch.from_numpy(ref_audio_np).unsqueeze(0), ref_sr, sample_rate).squeeze(0).numpy()
|
||||
)
|
||||
log(f"[Audio] Loaded reference audio for sample {i}: duration={len(ref_audio_np)/sample_rate:.2f}s")
|
||||
except Exception as e:
|
||||
log(f"[Warning] Failed to load reference audio: {e}")
|
||||
|
||||
# Preserve the original mode so validation failures do not leak into training.
|
||||
prev_training = unwrapped_model.training
|
||||
try:
|
||||
# Inference setup
|
||||
unwrapped_model.eval()
|
||||
# unwrapped_model.to(torch.bfloat16)
|
||||
unwrapped_model.audio_vae = audio_vae.to(torch.float32)
|
||||
|
||||
log(f"[Audio] Generating sample {i} with text: '{text[:50]}...'")
|
||||
autocast_ctx = (
|
||||
torch.autocast(device_type="cuda", dtype=torch.bfloat16)
|
||||
if torch.cuda.is_available()
|
||||
else contextlib.nullcontext()
|
||||
)
|
||||
with torch.no_grad():
|
||||
with autocast_ctx:
|
||||
generated = unwrapped_model.generate(
|
||||
target_text=text, inference_timesteps=10, cfg_value=2.0, seed=42
|
||||
)
|
||||
|
||||
# Restore training setup
|
||||
# unwrapped_model.to(torch.float32)
|
||||
# unwrapped_model.audio_vae = None
|
||||
|
||||
if generated is None or len(generated) == 0:
|
||||
log(f"[Warning] Generated audio is empty for sample {i}")
|
||||
continue
|
||||
|
||||
# Process generated audio
|
||||
gen_audio_np = (
|
||||
generated.cpu().float().numpy().flatten()
|
||||
if isinstance(generated, torch.Tensor)
|
||||
else np.array(generated, dtype=np.float32).flatten()
|
||||
)
|
||||
gen_audio_np = normalize_audio(gen_audio_np)
|
||||
|
||||
tag = f"val_sample_{i}"
|
||||
writer.add_audio(f"{tag}/generated_audio", gen_audio_np, global_step=step, sample_rate=gen_sr)
|
||||
log(f"[Audio] Generated audio for sample {i}: duration={len(gen_audio_np)/gen_sr:.2f}s")
|
||||
|
||||
# Log reference audio (at encoder input rate, which is what val_ds provides)
|
||||
if ref_audio_np is not None:
|
||||
writer.add_audio(
|
||||
f"{tag}/reference_audio", normalize_audio(ref_audio_np), global_step=step, sample_rate=sample_rate
|
||||
)
|
||||
|
||||
# Generate mel spectrogram figure
|
||||
try:
|
||||
mel_gen = compute_mel_spectrogram(gen_audio_np, gen_sr)
|
||||
mel_ref = compute_mel_spectrogram(ref_audio_np, sample_rate) if ref_audio_np is not None else None
|
||||
fig = create_mel_figure(gen_audio_np, mel_gen, gen_sr, step, ref_audio_np, mel_ref)
|
||||
writer.add_figure(f"{tag}/mel_spectrogram", fig, global_step=step)
|
||||
log(f"[Audio] Created mel spectrogram figure for sample {i}")
|
||||
except Exception as e:
|
||||
log(f"[Warning] Failed to create mel spectrogram: {e}")
|
||||
|
||||
except Exception as e:
|
||||
log(f"[Warning] Failed to generate audio for sample {i}: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
finally:
|
||||
# Always restore the training state, even if generation fails.
|
||||
try:
|
||||
# unwrapped_model.to(torch.float32)
|
||||
unwrapped_model.audio_vae = None
|
||||
if prev_training:
|
||||
unwrapped_model.train()
|
||||
else:
|
||||
unwrapped_model.eval()
|
||||
except Exception as e:
|
||||
log(f"[Warning] Failed to restore model state: {e}")
|
||||
|
||||
|
||||
def load_checkpoint(model, optimizer, scheduler, save_dir: Path, rank: int = 0):
|
||||
"""
|
||||
Load the latest checkpoint if it exists.
|
||||
Called by all ranks so that distributed state stays aligned.
|
||||
Returns the step number to resume from, or 0 if no checkpoint found.
|
||||
"""
|
||||
latest_folder = save_dir / "latest"
|
||||
if not latest_folder.exists():
|
||||
return 0
|
||||
|
||||
unwrapped = model.module if hasattr(model, "module") else model
|
||||
lora_cfg = unwrapped.lora_config
|
||||
|
||||
# Load model weights
|
||||
if lora_cfg is not None:
|
||||
# LoRA: load lora_weights
|
||||
lora_weights_path = latest_folder / "lora_weights.safetensors"
|
||||
if not lora_weights_path.exists():
|
||||
lora_weights_path = latest_folder / "lora_weights.ckpt"
|
||||
|
||||
if lora_weights_path.exists():
|
||||
if lora_weights_path.suffix == ".safetensors":
|
||||
from safetensors.torch import load_file
|
||||
|
||||
state_dict = load_file(str(lora_weights_path))
|
||||
else:
|
||||
ckpt = torch.load(lora_weights_path, map_location="cpu", weights_only=True)
|
||||
state_dict = ckpt.get("state_dict", ckpt)
|
||||
|
||||
unwrapped.load_state_dict(state_dict, strict=False)
|
||||
if rank == 0:
|
||||
print(f"Loaded LoRA weights from {lora_weights_path}", file=sys.stderr)
|
||||
else:
|
||||
# Full finetune: load model.safetensors or pytorch_model.bin
|
||||
model_path = latest_folder / "model.safetensors"
|
||||
if not model_path.exists():
|
||||
model_path = latest_folder / "pytorch_model.bin"
|
||||
|
||||
if model_path.exists():
|
||||
if model_path.suffix == ".safetensors":
|
||||
from safetensors.torch import load_file
|
||||
|
||||
state_dict = load_file(str(model_path))
|
||||
else:
|
||||
ckpt = torch.load(model_path, map_location="cpu", weights_only=True)
|
||||
state_dict = ckpt.get("state_dict", ckpt)
|
||||
|
||||
unwrapped.load_state_dict(state_dict, strict=False)
|
||||
if rank == 0:
|
||||
print(f"Loaded model weights from {model_path}", file=sys.stderr)
|
||||
|
||||
# Load optimizer state
|
||||
optimizer_path = latest_folder / "optimizer.pth"
|
||||
if optimizer_path.exists():
|
||||
optimizer.load_state_dict(torch.load(optimizer_path, map_location="cpu", weights_only=True))
|
||||
if rank == 0:
|
||||
print(f"Loaded optimizer state from {optimizer_path}", file=sys.stderr)
|
||||
|
||||
# Load scheduler state
|
||||
scheduler_path = latest_folder / "scheduler.pth"
|
||||
if scheduler_path.exists():
|
||||
scheduler.load_state_dict(torch.load(scheduler_path, map_location="cpu", weights_only=True))
|
||||
if rank == 0:
|
||||
print(f"Loaded scheduler state from {scheduler_path}", file=sys.stderr)
|
||||
|
||||
state_path = latest_folder / "training_state.json"
|
||||
if state_path.exists():
|
||||
with open(state_path, "r", encoding="utf-8") as f:
|
||||
state = json.load(f)
|
||||
resume_step = int(state.get("step", 0))
|
||||
if rank == 0:
|
||||
print(f"Resuming from step {resume_step}", file=sys.stderr)
|
||||
return resume_step
|
||||
|
||||
# Fallback for older checkpoints without metadata.
|
||||
step_folders = [d for d in save_dir.iterdir() if d.is_dir() and d.name.startswith("step_")]
|
||||
if step_folders:
|
||||
steps = [int(d.name.split("_")[1]) for d in step_folders]
|
||||
resume_step = max(steps)
|
||||
if rank == 0:
|
||||
print(f"Resuming from step {resume_step}", file=sys.stderr)
|
||||
return resume_step
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def save_checkpoint(
|
||||
model,
|
||||
optimizer,
|
||||
scheduler,
|
||||
save_dir: Path,
|
||||
step: int,
|
||||
pretrained_path: str = None,
|
||||
hf_model_id: str = "",
|
||||
distribute: bool = False,
|
||||
):
|
||||
"""
|
||||
Save checkpoint with different strategies for full finetune vs LoRA:
|
||||
- Full finetune: save non-vae weights to model.safetensors (or pytorch_model.bin if safetensors unavailable)
|
||||
- LoRA: save only lora weights to lora_weights.safetensors (or lora_weights.ckpt if safetensors unavailable)
|
||||
"""
|
||||
import shutil
|
||||
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
tag = f"step_{step:07d}"
|
||||
folder = save_dir / tag
|
||||
folder.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
unwrapped = model.module if hasattr(model, "module") else model
|
||||
full_state = unwrapped.state_dict()
|
||||
lora_cfg = unwrapped.lora_config
|
||||
|
||||
if lora_cfg is not None:
|
||||
# LoRA finetune: save only lora_A/lora_B weights
|
||||
state_dict = {k: v for k, v in full_state.items() if "lora_" in k}
|
||||
if SAFETENSORS_AVAILABLE:
|
||||
save_file(state_dict, folder / "lora_weights.safetensors")
|
||||
else:
|
||||
torch.save({"state_dict": state_dict}, folder / "lora_weights.ckpt")
|
||||
|
||||
# Save LoRA config and base model path to a separate JSON file
|
||||
# If distribute=True, save hf_model_id; otherwise save local pretrained_path
|
||||
base_model_to_save = hf_model_id if distribute else (str(pretrained_path) if pretrained_path else None)
|
||||
lora_info = {
|
||||
"base_model": base_model_to_save,
|
||||
"lora_config": lora_cfg.model_dump() if hasattr(lora_cfg, "model_dump") else vars(lora_cfg),
|
||||
}
|
||||
with open(folder / "lora_config.json", "w", encoding="utf-8") as f:
|
||||
json.dump(lora_info, f, indent=2, ensure_ascii=False)
|
||||
else:
|
||||
# Full finetune: save non-vae weights to model.safetensors
|
||||
state_dict = {k: v for k, v in full_state.items() if not k.startswith("audio_vae.")}
|
||||
if SAFETENSORS_AVAILABLE:
|
||||
save_file(state_dict, folder / "model.safetensors")
|
||||
else:
|
||||
torch.save({"state_dict": state_dict}, folder / "pytorch_model.bin")
|
||||
|
||||
# Copy config files from pretrained path
|
||||
if pretrained_path:
|
||||
pretrained_dir = Path(pretrained_path)
|
||||
files_to_copy = [
|
||||
"config.json",
|
||||
"audiovae.pth",
|
||||
"audiovae.safetensors",
|
||||
"tokenizer.json",
|
||||
"special_tokens_map.json",
|
||||
"tokenizer_config.json",
|
||||
]
|
||||
for fname in files_to_copy:
|
||||
src = pretrained_dir / fname
|
||||
if src.exists():
|
||||
shutil.copy2(src, folder / fname)
|
||||
|
||||
torch.save(optimizer.state_dict(), folder / "optimizer.pth")
|
||||
torch.save(scheduler.state_dict(), folder / "scheduler.pth")
|
||||
with open(folder / "training_state.json", "w", encoding="utf-8") as f:
|
||||
json.dump({"step": int(step)}, f)
|
||||
|
||||
# Update (or create) a `latest` folder by copying the most recent checkpoint
|
||||
latest_link = save_dir / "latest"
|
||||
try:
|
||||
if latest_link.exists():
|
||||
shutil.rmtree(latest_link)
|
||||
shutil.copytree(folder, latest_link)
|
||||
except Exception:
|
||||
print(f"Warning: failed to update latest checkpoint at {latest_link}", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from voxcpm.training.config import load_yaml_config
|
||||
|
||||
args = argbind.parse_args()
|
||||
config_file = args.get("config_path")
|
||||
# If YAML config provided, use YAML args to call train
|
||||
if config_file:
|
||||
yaml_args = load_yaml_config(config_file)
|
||||
train(**yaml_args)
|
||||
else:
|
||||
# Otherwise use command line args (parsed by argbind)
|
||||
with argbind.scope(args):
|
||||
train()
|
||||
Reference in New Issue
Block a user