Surunrun/SoulX-Singer

1import os2import torch3import json4import argparse5from tqdm import tqdm6import numpy as np7import soundfile as sf8from collections import OrderedDict9from omegaconf import DictConfig10 11from soulxsinger.utils.file_utils import load_config12from soulxsinger.models.soulxsinger_svc import SoulXSingerSVC13from soulxsinger.utils.audio_utils import load_wav14 15 16def build_model(17    model_path: str,18    config: DictConfig,19    device: str = "cuda",20):21    """22    Build the model from the pre-trained model path and model configuration.23 24    Args:25        model_path (str): Path to the checkpoint file.26        config (DictConfig): Model configuration.27        device (str, optional): Device to use. Defaults to "cuda".28 29    Returns:30        Tuple[torch.nn.Module, torch.nn.Module]: The initialized model and vocoder.31    """32 33    if not os.path.isfile(model_path):34        raise FileNotFoundError(35            f"Model checkpoint not found: {model_path}. "36            "Please download the pretrained model and place it at the path, or set --model_path."37        )38    model = SoulXSingerSVC(config).to(device)39    print("Model initialized.")40    print("Model parameters:", sum(p.numel() for p in model.parameters()) / 1e6, "M")41    42    checkpoint = torch.load(model_path, weights_only=False, map_location=device)43    if "state_dict" not in checkpoint:44        raise KeyError(45            f"Checkpoint at {model_path} has no 'state_dict' key. "46            "Expected a checkpoint saved with model.state_dict()."47        )48    model.load_state_dict(checkpoint["state_dict"], strict=True)49    50    model.eval()51    model.to(device)52    print("Model checkpoint loaded.")53 54    return model55 56 57def process(args, config, model: torch.nn.Module):58    """Run the full inference pipeline given a data_processor and model.59    """60 61    os.makedirs(args.save_dir, exist_ok=True)62    pt_wav = load_wav(args.prompt_wav_path, config.audio.sample_rate).to(args.device)63    gt_wav = load_wav(args.target_wav_path, config.audio.sample_rate).to(args.device)64    pt_f0 = torch.from_numpy(np.load(args.prompt_f0_path)).unsqueeze(0).to(args.device)65    gt_f0 = torch.from_numpy(np.load(args.target_f0_path)).unsqueeze(0).to(args.device)66 67    n_step = args.n_steps if hasattr(args, "n_steps") else config.infer.n_steps68    cfg = args.cfg if hasattr(args, "cfg") else config.infer.cfg69 70    generated_audio, generated_shift = model.infer(71        pt_wav=pt_wav,72        gt_wav=gt_wav,73        pt_f0=pt_f0,74        gt_f0=gt_f0,75        auto_shift=args.auto_shift, 76        pitch_shift=args.pitch_shift, 77        n_steps=n_step, 78        cfg=cfg,79        use_fp16=args.use_fp16,80    )81    generated_audio = generated_audio.squeeze().float().cpu().numpy()82    if args.pitch_shift != generated_shift:83        args.pitch_shift = generated_shift84        # print(f"Applied pitch shift of {generated_shift} semitones to match GT F0 contour.")85 86    sf.write(os.path.join(args.save_dir, "generated.wav"), generated_audio, config.audio.sample_rate)87    print(f"Generated audio saved to {os.path.join(args.save_dir, 'generated.wav')}")88 89 90def main(args, config):91    model = build_model(92        model_path=args.model_path,93        config=config,94        device=args.device,95    )96    process(args, config, model)97 98if __name__ == "__main__":99    parser = argparse.ArgumentParser()100    parser.add_argument("--device", type=str, default="cuda")101    parser.add_argument("--model_path", type=str, default='pretrained_models/soulx-singer/model.pt')102    parser.add_argument("--config", type=str, default='soulxsinger/config/soulxsinger.yaml')103    parser.add_argument("--prompt_wav_path", type=str, default='example/audio/zh_prompt.wav')104    parser.add_argument("--target_wav_path", type=str, default='example/audio/zh_target.wav')105    parser.add_argument("--prompt_f0_path", type=str, default='example/audio/zh_prompt_f0.npy')106    parser.add_argument("--target_f0_path", type=str, default='example/audio/zh_target_f0.npy')107    parser.add_argument("--save_dir", type=str, default='outputs')108    parser.add_argument("--auto_shift", action="store_true")109    parser.add_argument("--pitch_shift", type=int, default=0)110    parser.add_argument("--n_steps", type=int, default=32)111    parser.add_argument("--cfg", type=float, default=3.0)112    parser.add_argument(113        "--fp16",114        action="store_true",115        default=False,116        help="Use FP16 inference (faster on GPU)",117    )118    args = parser.parse_args()119 120    config = load_config(args.config)121    args.use_fp16 = args.fp16122    main(args, config)123