Soul-AILab/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 import SoulXSinger13from soulxsinger.utils.data_processor import DataProcessor14 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 = SoulXSinger(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    if args.control not in ("melody", "score"):61        raise ValueError(f"control must be 'melody' or 'score', got: {args.control}")62 63    print(f"prompt_metadata_path: {args.prompt_metadata_path}")64    print(f"target_metadata_path: {args.target_metadata_path}")65 66    os.makedirs(args.save_dir, exist_ok=True)67    data_processor = DataProcessor(68        hop_size=config.audio.hop_size,69        sample_rate=config.audio.sample_rate,70        phoneset_path=args.phoneset_path,71        device=args.device,72    )73 74    with open(args.prompt_metadata_path, "r", encoding="utf-8") as f:75        prompt_meta_list = json.load(f)76    if not prompt_meta_list:77        raise ValueError("Prompt metadata is empty. Please run preprocess on prompt audio first.")78    prompt_meta = prompt_meta_list[0]  # load the first segment as the prompt79    with open(args.target_metadata_path, "r", encoding="utf-8") as f:80        target_meta_list = json.load(f)81    infer_prompt_data = data_processor.process(prompt_meta, args.prompt_wav_path)82 83    assert len(target_meta_list) > 0, "No target segments found in the target metadata."84    generated_len = int(target_meta_list[-1]["time"][1] / 1000 * config.audio.sample_rate)85    generated_merged = np.zeros(generated_len, dtype=np.float32)86 87    for idx, target_meta in enumerate(88        tqdm(target_meta_list, total=len(target_meta_list), desc="Inferring segments"),89    ):90        start_sample_idx = int(target_meta["time"][0] / 1000 * config.audio.sample_rate)91        end_sample_idx = int(target_meta["time"][1] / 1000 * config.audio.sample_rate)92        infer_target_data = data_processor.process(target_meta, None)93 94        infer_data = {95            "prompt": infer_prompt_data,96            "target": infer_target_data,97        }98 99        with torch.no_grad():100            generated_audio = model.infer(101                infer_data,102                auto_shift=args.auto_shift,103                pitch_shift=args.pitch_shift,104                n_steps=config.infer.n_steps,105                cfg=config.infer.cfg,106                control=args.control,107            )108 109        generated_audio = generated_audio.squeeze().cpu().numpy()110        generated_merged[start_sample_idx : start_sample_idx + generated_audio.shape[0]] = generated_audio111 112    merged_path = os.path.join(args.save_dir, "generated.wav")113    sf.write(merged_path, generated_merged, 24000)114    print(f"Generated audio saved to {merged_path}")115 116 117def main(args, config):118    model = build_model(119        model_path=args.model_path,120        config=config,121        device=args.device,122    )123    process(args, config, model)124 125if __name__ == "__main__":126    parser = argparse.ArgumentParser()127    parser.add_argument("--device", type=str, default="cuda")128    parser.add_argument("--model_path", type=str, default='pretrained_models/soulx-singer/model.pt')129    parser.add_argument("--config", type=str, default='soulxsinger/config/soulxsinger.yaml')130    parser.add_argument("--prompt_wav_path", type=str, default='example/audio/zh_prompt.wav')131    parser.add_argument("--prompt_metadata_path", type=str, default='example/metadata/zh_prompt.json')132    parser.add_argument("--target_metadata_path", type=str, default='example/metadata/zh_target.json')133    parser.add_argument("--phoneset_path", type=str, default='soulxsinger/utils/phoneme/phone_set.json')134    parser.add_argument("--save_dir", type=str, default='outputs')135    parser.add_argument("--auto_shift", action="store_true")136    parser.add_argument("--pitch_shift", type=int, default=0)137    parser.add_argument(138        "--control",139        type=str,140        default="melody",141        choices=["melody", "score"],142        help="Control mode: melody or score only",143    )144    args = parser.parse_args()145    146    config = load_config(args.config)147    main(args, config)148