chore: import upstream snapshot with attribution

utils/__init__.py

@@ -0,0 +1,360 @@
+from __future__ import annotations
+
+import pathlib
+import re
+import time
+import types
+from collections import OrderedDict
+from fnmatch import fnmatch
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from basics.base_module import CategorizedModule
+from utils.hparams import hparams
+from utils.training_utils import get_latest_checkpoint_path
+
+
+def tensors_to_scalars(metrics):
+    new_metrics = {}
+    for k, v in metrics.items():
+        if isinstance(v, torch.Tensor):
+            v = v.item()
+        if type(v) is dict:
+            v = tensors_to_scalars(v)
+        new_metrics[k] = v
+    return new_metrics
+
+
+def collate_nd(values, pad_value=0, max_len=None):
+    """
+    Pad a list of Nd tensors on their first dimension and stack them into a (N+1)d tensor.
+    """
+    size = ((max(v.size(0) for v in values) if max_len is None else max_len), *values[0].shape[1:])
+    res = torch.full((len(values), *size), fill_value=pad_value, dtype=values[0].dtype, device=values[0].device)
+
+    for i, v in enumerate(values):
+        res[i, :len(v), ...] = v
+    return res
+
+
+def random_continuous_masks(*shape: int, dim: int, device: str | torch.device = 'cpu'):
+    start, end = torch.sort(
+        torch.randint(
+            low=0, high=shape[dim] + 1, size=(*shape[:dim], 2, *((1,) * (len(shape) - dim - 1))), device=device
+        ).expand(*((-1,) * (dim + 1)), *shape[dim + 1:]), dim=dim
+    )[0].split(1, dim=dim)
+    idx = torch.arange(
+        0, shape[dim], dtype=torch.long, device=device
+    ).reshape(*((1,) * dim), shape[dim], *((1,) * (len(shape) - dim - 1)))
+    masks = (idx >= start) & (idx < end)
+    return masks
+
+
+def _is_batch_full(batch, num_frames, max_batch_frames, max_batch_size):
+    if len(batch) == 0:
+        return 0
+    if len(batch) == max_batch_size:
+        return 1
+    if num_frames > max_batch_frames:
+        return 1
+    return 0
+
+
+def batch_by_size(
+        indices, num_frames_fn, max_batch_frames=80000, max_batch_size=48,
+        required_batch_size_multiple=1
+):
+    """
+    Yield mini-batches of indices bucketed by size. Batches may contain
+    sequences of different lengths.
+
+    Args:
+        indices (List[int]): ordered list of dataset indices
+        num_frames_fn (callable): function that returns the number of frames at
+            a given index
+        max_batch_frames (int, optional): max number of frames in each batch
+            (default: 80000).
+        max_batch_size (int, optional): max number of sentences in each
+            batch (default: 48).
+        required_batch_size_multiple: require the batch size to be multiple
+            of a given number
+    """
+    bsz_mult = required_batch_size_multiple
+
+    if isinstance(indices, types.GeneratorType):
+        indices = np.fromiter(indices, dtype=np.int64, count=-1)
+
+    sample_len = 0
+    sample_lens = []
+    batch = []
+    batches = []
+    for i in range(len(indices)):
+        idx = indices[i]
+        num_frames = num_frames_fn(idx)
+        sample_lens.append(num_frames)
+        sample_len = max(sample_len, num_frames)
+        assert sample_len <= max_batch_frames, (
+            "sentence at index {} of size {} exceeds max_batch_samples "
+            "limit of {}!".format(idx, sample_len, max_batch_frames)
+        )
+        num_frames = (len(batch) + 1) * sample_len
+
+        if _is_batch_full(batch, num_frames, max_batch_frames, max_batch_size):
+            mod_len = max(
+                bsz_mult * (len(batch) // bsz_mult),
+                len(batch) % bsz_mult,
+            )
+            batches.append(batch[:mod_len])
+            batch = batch[mod_len:]
+            sample_lens = sample_lens[mod_len:]
+            sample_len = max(sample_lens) if len(sample_lens) > 0 else 0
+        batch.append(idx)
+    if len(batch) > 0:
+        batches.append(batch)
+    return batches
+
+
+def make_positions(tensor, padding_idx):
+    """Replace non-padding symbols with their position numbers.
+
+    Position numbers begin at padding_idx+1. Padding symbols are ignored.
+    """
+    # The series of casts and type-conversions here are carefully
+    # balanced to both work with ONNX export and XLA. In particular XLA
+    # prefers ints, cumsum defaults to output longs, and ONNX doesn't know
+    # how to handle the dtype kwarg in cumsum.
+    mask = tensor.ne(padding_idx).int()
+    return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + padding_idx
+
+
+def softmax(x, dim):
+    return F.softmax(x, dim=dim, dtype=torch.float32)
+
+
+def unpack_dict_to_list(samples):
+    samples_ = []
+    bsz = samples.get('outputs').size(0)
+    for i in range(bsz):
+        res = {}
+        for k, v in samples.items():
+            try:
+                res[k] = v[i]
+            except (IndexError, TypeError):
+                pass
+        samples_.append(res)
+    return samples_
+
+
+def filter_kwargs(dict_to_filter, kwarg_obj):
+    import inspect
+
+    sig = inspect.signature(kwarg_obj)
+    if any(param.kind == param.VAR_KEYWORD for param in sig.parameters.values()):
+        # the signature contains definitions like **kwargs, so there is no need to filter
+        return dict_to_filter.copy()
+    filter_keys = [
+        param.name
+        for param in sig.parameters.values()
+        if param.kind == param.POSITIONAL_OR_KEYWORD or param.kind == param.KEYWORD_ONLY
+    ]
+    filtered_dict = {filter_key: dict_to_filter[filter_key] for filter_key in filter_keys if
+                     filter_key in dict_to_filter}
+    return filtered_dict
+
+
+def load_ckpt(
+        cur_model, ckpt_base_dir, ckpt_steps=None,
+        prefix_in_ckpt='model', exclude_key_patterns=None, key_in_ckpt='state_dict',
+        strict=True, device='cpu'
+):
+    if exclude_key_patterns is None:
+        # Pop all RoPE buffers from some old checkpoints,
+        # Because these buffers are all computed during initialization now.
+        # TODO: this is a legacy handling and should be removed in the future.
+        exclude_key_patterns = ['*.rotary_embed.*']
+    if not isinstance(ckpt_base_dir, pathlib.Path):
+        ckpt_base_dir = pathlib.Path(ckpt_base_dir)
+    if ckpt_base_dir.is_file():
+        checkpoint_path = [ckpt_base_dir]
+    elif ckpt_steps is not None:
+        checkpoint_path = [ckpt_base_dir / f'model_ckpt_steps_{int(ckpt_steps)}.ckpt']
+    else:
+        base_dir = ckpt_base_dir
+        checkpoint_path = sorted(
+            [
+                ckpt_file
+                for ckpt_file in base_dir.iterdir()
+                if ckpt_file.is_file() and re.fullmatch(r'model_ckpt_steps_\d+\.ckpt', ckpt_file.name)
+            ],
+            key=lambda x: int(re.search(r'\d+', x.name).group(0))
+        )
+    assert len(checkpoint_path) > 0, f'| ckpt not found in {ckpt_base_dir}.'
+    checkpoint_path = checkpoint_path[-1]
+    ckpt_loaded = torch.load(checkpoint_path, map_location=device)
+    if isinstance(cur_model, CategorizedModule):
+        cur_model.check_category(ckpt_loaded.get('category'))
+    if key_in_ckpt is None:
+        state_dict = ckpt_loaded
+    else:
+        state_dict = ckpt_loaded[key_in_ckpt]
+    if prefix_in_ckpt is not None:
+        old_state_dict = state_dict
+        state_dict = OrderedDict()
+        for k, v in old_state_dict.items():
+            if not k.startswith(f'{prefix_in_ckpt}.'):
+                continue
+            k = k[len(prefix_in_ckpt) + 1:]
+            excluded = False
+            for pat in exclude_key_patterns:
+                if fnmatch(k, pat):
+                    excluded = True
+                    break
+            if excluded:
+                continue
+            state_dict[k] = v
+
+    # Manual self-attention (MultiheadSelfAttentionWithRoPE) uses 'in_proj.weight',
+    # while older checkpoints saved from torch.nn.MultiheadAttention use 'in_proj_weight'.
+    # The two tensors have identical shape and semantics (Q/K/V stacked along dim 0),
+    # so a key rename is sufficient to load legacy ckpts.
+    renamed = OrderedDict()
+    for k, v in state_dict.items():
+        if k.endswith('.self_attn.in_proj_weight'):
+            k = k[:-len('in_proj_weight')] + 'in_proj.weight'
+        renamed[k] = v
+    state_dict = renamed
+
+    if not strict:
+        cur_model_state_dict = cur_model.state_dict()
+        unmatched_keys = []
+        for key, param in state_dict.items():
+            if key in cur_model_state_dict:
+                new_param = cur_model_state_dict[key]
+                if new_param.shape != param.shape:
+                    unmatched_keys.append(key)
+                    print('| Unmatched keys: ', key, new_param.shape, param.shape)
+        for key in unmatched_keys:
+            del state_dict[key]
+    cur_model.load_state_dict(state_dict, strict=strict)
+    shown_model_name = 'state dict'
+    if prefix_in_ckpt is not None:
+        shown_model_name = f'\'{prefix_in_ckpt}\''
+    elif key_in_ckpt is not None:
+        shown_model_name = f'\'{key_in_ckpt}\''
+    print(f'| load {shown_model_name} from \'{checkpoint_path}\'.')
+
+
+def remove_padding(x, padding_idx=0):
+    if x is None:
+        return None
+    assert len(x.shape) in [1, 2]
+    if len(x.shape) == 2:  # [T, H]
+        return x[np.abs(x).sum(-1) != padding_idx]
+    elif len(x.shape) == 1:  # [T]
+        return x[x != padding_idx]
+
+
+class Timer:
+    timer_map = {}
+
+    def __init__(self, name, print_time=False):
+        if name not in Timer.timer_map:
+            Timer.timer_map[name] = 0
+        self.name = name
+        self.print_time = print_time
+
+    def __enter__(self):
+        self.t = time.time()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        Timer.timer_map[self.name] += time.time() - self.t
+        if self.print_time:
+            print(self.name, Timer.timer_map[self.name])
+
+
+def print_arch(model, model_name='model'):
+    print(f"| {model_name} Arch: ", model)
+    # num_params(model, model_name=model_name)
+
+
+def num_params(model, print_out=True, model_name="model"):
+    parameters = filter(lambda p: p.requires_grad, model.parameters())
+    parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
+    if print_out:
+        print(f'| {model_name} Trainable Parameters: %.3fM' % parameters)
+    return parameters
+
+
+def build_object_from_class_name(cls_str, parent_cls, *args, **kwargs):
+    import importlib
+
+    pkg = ".".join(cls_str.split(".")[:-1])
+    cls_name = cls_str.split(".")[-1]
+    cls_type = getattr(importlib.import_module(pkg), cls_name)
+    if parent_cls is not None:
+        assert issubclass(cls_type, parent_cls), f'| {cls_type} is not subclass of {parent_cls}.'
+
+    return cls_type(*args, **filter_kwargs(kwargs, cls_type))
+
+
+def build_lr_scheduler_from_config(optimizer, scheduler_args):
+    try:
+        # PyTorch 2.0+
+        from torch.optim.lr_scheduler import LRScheduler as LRScheduler
+    except ImportError:
+        # PyTorch 1.X
+        from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+
+    def helper(params):
+        if isinstance(params, list):
+            return [helper(s) for s in params]
+        elif isinstance(params, dict):
+            resolved = {k: helper(v) for k, v in params.items()}
+            if 'cls' in resolved:
+                if (
+                    resolved["cls"] == "torch.optim.lr_scheduler.ChainedScheduler"
+                    and scheduler_args["scheduler_cls"] == "torch.optim.lr_scheduler.SequentialLR"
+                ):
+                    raise ValueError("ChainedScheduler cannot be part of a SequentialLR.")
+                resolved['optimizer'] = optimizer
+                obj = build_object_from_class_name(
+                    resolved['cls'],
+                    LRScheduler,
+                    **resolved
+                )
+                return obj
+            return resolved
+        else:
+            return params
+
+    resolved = helper(scheduler_args)
+    resolved['optimizer'] = optimizer
+    return build_object_from_class_name(
+        scheduler_args['scheduler_cls'],
+        LRScheduler,
+        **resolved
+    )
+
+
+def simulate_lr_scheduler(optimizer_args, scheduler_args, step_count, num_param_groups=1):
+    optimizer_cls = optimizer_args['optimizer_cls']
+    optimizer = build_object_from_class_name(
+        'torch.optim.AdamW' if optimizer_cls == 'modules.optimizer.muon.Muon_AdamW' else optimizer_cls,
+        torch.optim.Optimizer,
+        [{'params': torch.nn.Parameter(), 'initial_lr': optimizer_args['lr']} for _ in range(num_param_groups)],
+        **optimizer_args
+    )
+    scheduler = build_lr_scheduler_from_config(optimizer, scheduler_args)
+    scheduler.optimizer._step_count = 1
+    for _ in range(step_count):
+        scheduler.step()
+    return scheduler.state_dict()
+
+
+def remove_suffix(string: str, suffix: str):
+    #  Just for Python 3.8 compatibility, since `str.removesuffix()` API of is available since Python 3.9
+    if string.endswith(suffix):
+        string = string[:-len(suffix)]
+    return string

utils/binarizer_utils.py

@@ -0,0 +1,229 @@
+from typing import Union
+
+import librosa
+import numpy as np
+import parselmouth
+import torch
+
+from modules.nsf_hifigan.nvSTFT import STFT
+from utils.decomposed_waveform import DecomposedWaveform
+from utils.pitch_utils import interp_f0
+
+
+def get_mel_torch(
+        waveform, samplerate,
+        *,
+        num_mel_bins=128, hop_size=512, win_size=2048, fft_size=2048,
+        fmin=40, fmax=16000,
+        keyshift=0, speed=1, device=None
+):
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    stft = STFT(samplerate, num_mel_bins, fft_size, win_size, hop_size, fmin, fmax, device=device)
+    with torch.no_grad():
+        wav_torch = torch.from_numpy(waveform).to(device)
+        mel_torch = stft.get_mel(wav_torch.unsqueeze(0), keyshift=keyshift, speed=speed).squeeze(0).T
+        return mel_torch.cpu().numpy()
+
+
+@torch.no_grad()
+def get_mel2ph_torch(lr, durs, length, timestep, device='cpu'):
+    ph_acc = torch.round(torch.cumsum(durs.to(device), dim=0) / timestep + 0.5).long()
+    ph_dur = torch.diff(ph_acc, dim=0, prepend=torch.LongTensor([0]).to(device))
+    mel2ph = lr(ph_dur[None])[0]
+    num_frames = mel2ph.shape[0]
+    if num_frames < length:
+        mel2ph = torch.cat((mel2ph, torch.full((length - num_frames,), fill_value=mel2ph[-1], device=device)), dim=0)
+    elif num_frames > length:
+        mel2ph = mel2ph[:length]
+    return mel2ph
+
+
+def get_pitch_parselmouth(
+        waveform, samplerate, length,
+        *, hop_size, f0_min=65, f0_max=1100,
+        speed=1, interp_uv=False
+):
+    """
+
+    :param waveform: [T]
+    :param samplerate: sampling rate
+    :param length: Expected number of frames
+    :param hop_size: Frame width, in number of samples
+    :param f0_min: Minimum f0 in Hz
+    :param f0_max: Maximum f0 in Hz
+    :param speed: Change the speed
+    :param interp_uv: Interpolate unvoiced parts
+    :return: f0, uv
+    """
+    hop_size = int(np.round(hop_size * speed))
+    time_step = hop_size / samplerate
+
+    l_pad = int(np.ceil(1.5 / f0_min * samplerate))
+    r_pad = hop_size * ((len(waveform) - 1) // hop_size + 1) - len(waveform) + l_pad + 1
+    waveform = np.pad(waveform, (l_pad, r_pad))
+
+    # noinspection PyArgumentList
+    s = parselmouth.Sound(waveform, sampling_frequency=samplerate).to_pitch_ac(
+        time_step=time_step, voicing_threshold=0.6,
+        pitch_floor=f0_min, pitch_ceiling=f0_max
+    )
+    assert np.abs(s.t1 - 1.5 / f0_min) < 0.001
+    f0 = s.selected_array['frequency'].astype(np.float32)
+    if len(f0) < length:
+        f0 = np.pad(f0, (0, length - len(f0)))
+    f0 = f0[: length]
+    uv = f0 == 0
+    if interp_uv:
+        f0, uv = interp_f0(f0, uv)
+    return f0, uv
+
+
+def get_energy_librosa(waveform, length, *, hop_size, win_size, domain='db'):
+    """
+    Definition of energy: RMS of the waveform, in dB representation
+    :param waveform: [T]
+    :param length: Expected number of frames
+    :param hop_size: Frame width, in number of samples
+    :param win_size: Window size, in number of samples
+    :param domain: db or amplitude
+    :return: energy
+    """
+    energy = librosa.feature.rms(y=waveform, frame_length=win_size, hop_length=hop_size)[0]
+    if len(energy) < length:
+        energy = np.pad(energy, (0, length - len(energy)))
+    energy = energy[: length]
+    if domain == 'db':
+        energy = librosa.amplitude_to_db(energy)
+    elif domain == 'amplitude':
+        pass
+    else:
+        raise ValueError(f'Invalid domain: {domain}')
+    return energy
+
+
+def get_breathiness(
+        waveform: Union[np.ndarray, DecomposedWaveform],
+        samplerate, f0, length,
+        *, hop_size=None, fft_size=None, win_size=None
+):
+    """
+    Definition of breathiness: RMS of the aperiodic part, in dB representation
+    :param waveform: All other analysis parameters will not take effect if a DeconstructedWaveform is given
+    :param samplerate: sampling rate
+    :param f0: reference f0
+    :param length: Expected number of frames
+    :param hop_size: Frame width, in number of samples
+    :param fft_size: Number of fft bins
+    :param win_size: Window size, in number of samples
+    :return: breathiness
+    """
+    if not isinstance(waveform, DecomposedWaveform):
+        waveform = DecomposedWaveform(
+            waveform=waveform, samplerate=samplerate, f0=f0,
+            hop_size=hop_size, fft_size=fft_size, win_size=win_size
+        )
+    waveform_ap = waveform.aperiodic()
+    breathiness = get_energy_librosa(
+        waveform_ap, length=length,
+        hop_size=waveform.hop_size, win_size=waveform.win_size
+    )
+    return breathiness
+
+
+def get_voicing(
+        waveform: Union[np.ndarray, DecomposedWaveform],
+        samplerate, f0, length,
+        *, hop_size=None, fft_size=None, win_size=None
+):
+    """
+    Definition of voicing: RMS of the harmonic part, in dB representation
+    :param waveform: All other analysis parameters will not take effect if a DeconstructedWaveform is given
+    :param samplerate: sampling rate
+    :param f0: reference f0
+    :param length: Expected number of frames
+    :param hop_size: Frame width, in number of samples
+    :param fft_size: Number of fft bins
+    :param win_size: Window size, in number of samples
+    :return: voicing
+    """
+    if not isinstance(waveform, DecomposedWaveform):
+        waveform = DecomposedWaveform(
+            waveform=waveform, samplerate=samplerate, f0=f0,
+            hop_size=hop_size, fft_size=fft_size, win_size=win_size
+        )
+    waveform_sp = waveform.harmonic()
+    voicing = get_energy_librosa(
+        waveform_sp, length=length,
+        hop_size=waveform.hop_size, win_size=waveform.win_size
+    )
+    return voicing
+
+
+def get_tension_base_harmonic(
+        waveform: Union[np.ndarray, DecomposedWaveform],
+        samplerate, f0, length,
+        *, hop_size=None, fft_size=None, win_size=None,
+        domain='logit'
+):
+    """
+    Definition of tension: radio of the real harmonic part (harmonic part except the base harmonic)
+    to the full harmonic part.
+    :param waveform: All other analysis parameters will not take effect if a DeconstructedWaveform is given
+    :param samplerate: sampling rate
+    :param f0: reference f0
+    :param length: Expected number of frames
+    :param hop_size: Frame width, in number of samples
+    :param fft_size: Number of fft bins
+    :param win_size: Window size, in number of samples
+    :param domain: The domain of the final ratio representation.
+     Can be 'ratio' (the raw ratio), 'db' (log decibel) or 'logit' (the reverse function of sigmoid)
+    :return: tension
+    """
+    if not isinstance(waveform, DecomposedWaveform):
+        waveform = DecomposedWaveform(
+            waveform=waveform, samplerate=samplerate, f0=f0,
+            hop_size=hop_size, fft_size=fft_size, win_size=win_size
+        )
+    waveform_h = waveform.harmonic()
+    waveform_base_h = waveform.harmonic(0)
+    energy_base_h = get_energy_librosa(
+        waveform_base_h, length,
+        hop_size=waveform.hop_size, win_size=waveform.win_size,
+        domain='amplitude'
+    )
+    energy_h = get_energy_librosa(
+        waveform_h, length,
+        hop_size=waveform.hop_size, win_size=waveform.win_size,
+        domain='amplitude'
+    )
+    tension = np.sqrt(np.clip(energy_h ** 2 - energy_base_h ** 2, a_min=0, a_max=None)) / (energy_h + 1e-5)
+    if domain == 'ratio':
+        tension = np.clip(tension, a_min=0, a_max=1)
+    elif domain == 'db':
+        tension = np.clip(tension, a_min=1e-5, a_max=1)
+        tension = librosa.amplitude_to_db(tension)
+    elif domain == 'logit':
+        tension = np.clip(tension, a_min=1e-4, a_max=1 - 1e-4)
+        tension = np.log(tension / (1 - tension))
+    return tension
+
+
+class SinusoidalSmoothingConv1d(torch.nn.Conv1d):
+    def __init__(self, kernel_size):
+        super().__init__(
+            in_channels=1,
+            out_channels=1,
+            kernel_size=max(kernel_size, 1),
+            bias=False,
+            padding='same',
+            padding_mode='replicate'
+        )
+        if kernel_size > 1:
+            smooth_kernel = torch.sin(torch.from_numpy(
+                np.linspace(0, 1, kernel_size).astype(np.float32) * np.pi
+            ))
+            smooth_kernel /= smooth_kernel.sum()
+        else:
+            smooth_kernel = torch.tensor([1.0], dtype=torch.float32)
+        self.weight.data = smooth_kernel[None, None]

utils/decomposed_waveform.py

@@ -0,0 +1,282 @@
+from typing import Dict
+
+import numpy as np
+import pyworld as pw
+import torch
+from torch.nn import functional as F
+
+from modules.hnsep.vr import load_sep_model
+from utils import hparams
+from utils.pitch_utils import interp_f0
+
+
+class DecomposedWaveform:
+    def __new__(
+            cls, waveform, samplerate, f0,
+            *,
+            hop_size=None, fft_size=None, win_size=None,
+            algorithm='world', device=None
+    ):
+        if algorithm == 'world':
+            obj = object.__new__(DecomposedWaveformPyWorld)
+            # noinspection PyProtectedMember
+            obj._init(
+                waveform=waveform, samplerate=samplerate, f0=f0,
+                hop_size=hop_size, fft_size=fft_size, win_size=win_size,
+                device=device
+            )
+        elif algorithm == 'vr':
+            obj = object.__new__(DecomposedWaveformVocalRemover)
+            hnsep_ckpt = hparams['hnsep_ckpt']
+            # noinspection PyProtectedMember
+            obj._init(
+                waveform=waveform, samplerate=samplerate, f0=f0, hop_size=hop_size,
+                fft_size=fft_size, win_size=win_size, model_path=hnsep_ckpt,
+                device=device
+            )
+        else:
+            raise ValueError(f" [x] Unknown harmonic-noise separator: {algorithm}")
+        return obj
+
+    @property
+    def samplerate(self):
+        raise NotImplementedError()
+
+    @property
+    def hop_size(self):
+        raise NotImplementedError()
+
+    @property
+    def fft_size(self):
+        raise NotImplementedError()
+
+    @property
+    def win_size(self):
+        raise NotImplementedError()
+
+    def harmonic(self, k: int = None) -> np.ndarray:
+        raise NotImplementedError()
+
+    def aperiodic(self) -> np.ndarray:
+        raise NotImplementedError()
+
+
+class DecomposedWaveformPyWorld(DecomposedWaveform):
+    def _init(
+            self, waveform, samplerate, f0,  # basic parameters
+            *,
+            hop_size=None, fft_size=None, win_size=None, base_harmonic_radius=3.5,  # analysis parameters
+            device=None  # computation parameters
+    ):
+        # the source components
+        self._waveform = waveform
+        self._samplerate = samplerate
+        self._f0 = f0
+        # extraction parameters
+        self._hop_size = hop_size
+        self._fft_size = fft_size if fft_size is not None else win_size
+        self._win_size = win_size if win_size is not None else win_size
+        self._time_step = hop_size / samplerate
+        self._half_width = base_harmonic_radius
+        self._device = ('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
+        # intermediate variables
+        self._f0_world = None
+        self._sp = None
+        self._ap = None
+        # final components
+        self._harmonic_part: np.ndarray = None
+        self._aperiodic_part: np.ndarray = None
+        self._harmonics: Dict[int, np.ndarray] = {}
+
+    @property
+    def samplerate(self):
+        return self._samplerate
+
+    @property
+    def hop_size(self):
+        return self._hop_size
+
+    @property
+    def fft_size(self):
+        return self._fft_size
+
+    @property
+    def win_size(self):
+        return self._win_size
+
+    def _world_extraction(self):
+        # Add a tiny noise to the signal to avoid NaN results of D4C in rare edge cases
+        # References:
+        #   - https://github.com/JeremyCCHsu/Python-Wrapper-for-World-Vocoder/issues/50
+        #   - https://github.com/mmorise/World/issues/116
+        x = self._waveform.astype(np.double) + np.random.randn(*self._waveform.shape) * 1e-5
+        samplerate = self._samplerate
+        f0 = self._f0.astype(np.double)
+
+        hop_size = self._hop_size
+        fft_size = self._fft_size
+
+        wav_frames = (x.shape[0] + hop_size - 1) // hop_size
+        f0_frames = f0.shape[0]
+        if f0_frames < wav_frames:
+            f0 = np.pad(f0, (0, wav_frames - f0_frames), mode='constant', constant_values=(f0[0], f0[-1]))
+        elif f0_frames > wav_frames:
+            f0 = f0[:wav_frames]
+
+        time_step = hop_size / samplerate
+        t = np.arange(0, wav_frames) * time_step
+        self._f0_world = f0
+        self._sp = pw.cheaptrick(x, f0, t, samplerate, fft_size=fft_size)  # extract smoothed spectrogram
+        self._ap = pw.d4c(x, f0, t, samplerate, fft_size=fft_size)  # extract aperiodicity
+
+    def _kth_harmonic(self, k: int) -> np.ndarray:
+        """
+        Extract the Kth harmonic (starting from 0) from the waveform. Author: @yxlllc
+        :param k: a non-negative integer
+        :return: kth_harmonic float32[T]
+        """
+        if k in self._harmonics:
+            return self._harmonics[k]
+
+        hop_size = self._hop_size
+        win_size = self._win_size
+        samplerate = self._samplerate
+        half_width = self._half_width
+        device = self._device
+
+        waveform = torch.from_numpy(self.harmonic()).unsqueeze(0).to(device)  # [B, n_samples]
+        n_samples = waveform.shape[1]       
+        f0 = self._f0 * (k + 1)
+        pad_size = int(n_samples // hop_size) - len(f0) + 1
+        if pad_size > 0:
+            f0 = np.pad(f0, (0, pad_size), mode='constant', constant_values=(f0[0], f0[-1]))
+        
+        f0, _ = interp_f0(f0, uv=f0 == 0)
+        f0 = torch.from_numpy(f0).to(device)[None, :, None]  # [B, n_frames, 1]
+        n_f0_frames = f0.shape[1]
+
+        phase = torch.arange(win_size, dtype=waveform.dtype, device=device) / win_size * 2 * np.pi
+        nuttall_window = (
+                0.355768
+                - 0.487396 * torch.cos(phase)
+                + 0.144232 * torch.cos(2 * phase)
+                - 0.012604 * torch.cos(3 * phase)
+        )
+        spec = torch.stft(
+            waveform,
+            n_fft=win_size,
+            win_length=win_size,
+            hop_length=hop_size,
+            window=nuttall_window,
+            center=True,
+            return_complex=True
+        ).permute(0, 2, 1)  # [B, n_frames, n_spec]
+        n_spec_frames, n_specs = spec.shape[1:]
+        idx = torch.arange(n_specs).unsqueeze(0).unsqueeze(0).to(f0)  # [1, 1, n_spec]
+        center = f0 * win_size / samplerate
+        start = torch.clip(center - half_width, min=0)
+        end = torch.clip(center + half_width, max=n_specs)
+        idx_mask = (center >= 1) & (idx >= start) & (idx < end)  # [B, n_frames, n_spec]
+        if n_f0_frames < n_spec_frames:
+            idx_mask = F.pad(idx_mask, [0, 0, 0, n_spec_frames - n_f0_frames])
+        spec = spec * idx_mask[:, :n_spec_frames, :]
+        self._harmonics[k] = torch.istft(
+            spec.permute(0, 2, 1),
+            n_fft=win_size,
+            win_length=win_size,
+            hop_length=hop_size,
+            window=nuttall_window,
+            center=True,
+            length=n_samples
+        ).squeeze(0).cpu().numpy()
+
+        return self._harmonics[k]
+
+    def harmonic(self, k: int = None) -> np.ndarray:
+        """
+        Extract the full harmonic part, or the Kth harmonic if `k` is not None, from the waveform.
+        :param k: an integer representing the harmonic index, starting from 0
+        :return: full_harmonics float32[T] or kth_harmonic float32[T]
+        """
+        if k is not None:
+            return self._kth_harmonic(k)
+        if self._harmonic_part is not None:
+            return self._harmonic_part
+        if self._sp is None or self._ap is None:
+            self._world_extraction()
+        # noinspection PyAttributeOutsideInit
+        self._harmonic_part = pw.synthesize(
+            self._f0_world,
+            np.clip(self._sp * (1 - self._ap * self._ap), a_min=1e-16, a_max=None),  # clip to avoid zeros
+            np.zeros_like(self._ap),
+            self._samplerate, frame_period=self._time_step * 1000
+        ).astype(np.float32)  # synthesize the harmonic part using the parameters
+        return self._harmonic_part
+
+    def aperiodic(self) -> np.ndarray:
+        """
+        Extract the aperiodic part from the waveform.
+        :return: aperiodic_part float32[T]
+        """
+        if self._aperiodic_part is not None:
+            return self._aperiodic_part
+        if self._sp is None or self._ap is None:
+            self._world_extraction()
+        # noinspection PyAttributeOutsideInit
+        self._aperiodic_part = pw.synthesize(
+            self._f0_world, self._sp * self._ap * self._ap, np.ones_like(self._ap),
+            self._samplerate, frame_period=self._time_step * 1000
+        ).astype(np.float32)  # synthesize the aperiodic part using the parameters
+        return self._aperiodic_part
+
+
+SEP_MODEL = None
+
+
+class DecomposedWaveformVocalRemover(DecomposedWaveformPyWorld):
+    def _init(
+            self, waveform, samplerate, f0,
+            hop_size=None, fft_size=None, win_size=None, base_harmonic_radius=3.5,
+            model_path=None, device=None
+    ):
+        super()._init(
+            waveform, samplerate, f0, hop_size=hop_size, fft_size=fft_size,
+            win_size=win_size, base_harmonic_radius=base_harmonic_radius, device=device
+        )
+        global SEP_MODEL
+        if SEP_MODEL is None:
+            SEP_MODEL = load_sep_model(model_path, self._device)
+        self.sep_model = SEP_MODEL
+
+    def _infer(self):
+        with torch.no_grad():
+            x = torch.from_numpy(self._waveform).to(self._device).reshape(1, 1, -1)
+            if not self.sep_model.is_mono:
+                x = x.repeat(1, 2, 1)
+            x = self.sep_model.predict_from_audio(x)
+            x = torch.mean(x, dim=1)
+            self._harmonic_part = x.squeeze().cpu().numpy()
+            self._aperiodic_part = self._waveform - self._harmonic_part
+
+    def harmonic(self, k: int = None) -> np.ndarray:
+        """
+        Extract the full harmonic part, or the Kth harmonic if `k` is not None, from the waveform.
+        :param k: an integer representing the harmonic index, starting from 0
+        :return: full_harmonics float32[T] or kth_harmonic float32[T]
+        """
+        if k is not None:
+            return self._kth_harmonic(k)
+        if self._harmonic_part is not None:
+            return self._harmonic_part
+        self._infer()
+        return self._harmonic_part
+
+    def aperiodic(self) -> np.ndarray:
+        """
+        Extract the aperiodic part from the waveform.
+        :return: aperiodic_part float32[T]
+        """
+        if self._aperiodic_part is not None:
+            return self._aperiodic_part
+        self._infer()
+        return self._aperiodic_part

utils/hparams.py

@@ -0,0 +1,146 @@
+import argparse
+import os
+import yaml
+
+try:
+    from lightning.pytorch.utilities.rank_zero import rank_zero_only
+except ModuleNotFoundError:
+    def rank_zero_only(f):
+        return f
+
+from utils.multiprocess_utils import is_main_process as mp_is_main_process
+global_print_hparams = True
+hparams = {}
+
+
+class Args:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            self.__setattr__(k, v)
+
+
+def override_config(old_config: dict, new_config: dict):
+    for k, v in new_config.items():
+        if isinstance(v, dict) and k in old_config:
+            override_config(old_config[k], new_config[k])
+        else:
+            old_config[k] = v
+
+
+def set_hparams(config='', exp_name='', hparams_str='', print_hparams=True, global_hparams=True):
+    """
+        Load hparams from multiple sources:
+        1. config chain (i.e. first load base_config, then load config);
+        2. if reset == True, load from the (auto-saved) complete config file ('config.yaml')
+           which contains all settings and do not rely on base_config;
+        3. load from argument --hparams or hparams_str, as temporary modification.
+    """
+    if config == '':
+        parser = argparse.ArgumentParser(description='neural music')
+        parser.add_argument('--config', type=str, default='',
+                            help='location of the data corpus')
+        parser.add_argument('--exp_name', type=str, default='', help='exp_name')
+        parser.add_argument('--hparams', type=str, default='',
+                            help='location of the data corpus')
+        parser.add_argument('--infer', action='store_true', help='infer')
+        parser.add_argument('--reset', action='store_true', help='reset hparams')
+        args, unknown = parser.parse_known_args()
+        
+        tmp_args_hparams = args.hparams.split(',') if args.hparams.strip() != '' else []
+        tmp_args_hparams.extend(hparams_str.split(',') if hparams_str.strip() != '' else [])
+        args.hparams = ','.join(tmp_args_hparams)
+    else:
+        args = Args(config=config, exp_name=exp_name, hparams=hparams_str,
+                    infer=False, reset=False)
+
+    args_work_dir = ''
+    if args.exp_name != '':
+        args.work_dir = args.exp_name
+        args_work_dir = os.path.join('checkpoints', args.work_dir)
+
+    config_chains = []
+    loaded_config = set()
+
+    def load_config(config_fn):  # deep first
+        with open(config_fn, encoding='utf-8') as f:
+            hparams_ = yaml.safe_load(f)
+        loaded_config.add(config_fn)
+        if 'base_config' in hparams_:
+            ret_hparams = {}
+            if not isinstance(hparams_['base_config'], list):
+                hparams_['base_config'] = [hparams_['base_config']]
+            for c in hparams_['base_config']:
+                if c not in loaded_config:
+                    if c.startswith('.'):
+                        c = f'{os.path.dirname(config_fn)}/{c}'
+                        c = os.path.normpath(c)
+                    override_config(ret_hparams, load_config(c))
+            override_config(ret_hparams, hparams_)
+        else:
+            ret_hparams = hparams_
+        config_chains.append(config_fn)
+        return ret_hparams
+
+    global hparams
+    assert args.config != '' or args_work_dir != '', 'Either config or exp name should be specified.'
+    saved_hparams = {}
+    ckpt_config_path = os.path.join(args_work_dir, 'config.yaml')
+    if args_work_dir != '' and os.path.exists(ckpt_config_path):
+        with open(ckpt_config_path, encoding='utf-8') as f:
+            saved_hparams.update(yaml.safe_load(f))
+
+    hparams_ = {}
+    if args.config != '':
+        hparams_.update(load_config(args.config))
+
+    if not args.reset:
+        hparams_.update(saved_hparams)
+    hparams_['work_dir'] = args_work_dir
+
+    if args.hparams != "":
+        for new_hparam in args.hparams.split(","):
+            if new_hparam.strip() == "":
+                continue
+            k, v = new_hparam.split("=")
+            if k not in hparams_:
+                hparams_[k] = eval(v)
+            if v in ['True', 'False'] or type(hparams_[k]) == bool:
+                hparams_[k] = eval(v)
+            else:
+                hparams_[k] = type(hparams_[k])(v)
+
+    @rank_zero_only
+    def dump_hparams():
+        if args_work_dir != '' and (not os.path.exists(ckpt_config_path) or args.reset) and not args.infer:
+            os.makedirs(hparams_['work_dir'], exist_ok=True)
+            if mp_is_main_process:
+                # Only the main process will save the config file
+                with open(ckpt_config_path, 'w', encoding='utf-8') as f:
+                    hparams_non_recursive = hparams_.copy()
+                    hparams_non_recursive['base_config'] = []
+                    yaml.safe_dump(hparams_non_recursive, f, allow_unicode=True, encoding='utf-8')
+    dump_hparams()
+
+    hparams_['infer'] = args.infer
+    if global_hparams:
+        hparams.clear()
+        hparams.update(hparams_)
+    
+    if hparams.get('exp_name') is None:
+        hparams['exp_name'] = args.exp_name
+    if hparams_.get('exp_name') is None:
+        hparams_['exp_name'] = args.exp_name
+    
+    @rank_zero_only
+    def print_out_hparams():
+        global global_print_hparams
+        if mp_is_main_process and print_hparams and global_print_hparams and global_hparams:
+            print('| Hparams chains: ', config_chains)
+            print('| Hparams: ')
+            for i, (k, v) in enumerate(sorted(hparams_.items())):
+                print(f"\033[0;33m{k}\033[0m: {v}, ", end="\n" if i % 5 == 4 else "")
+            print("")
+            global_print_hparams = False
+    print_out_hparams()
+    
+    return hparams_

utils/indexed_datasets.py

@@ -0,0 +1,96 @@
+import pathlib
+import multiprocessing
+from collections import deque
+
+import h5py
+import torch
+import numpy as np
+
+
+class IndexedDataset:
+    def __init__(self, path, prefix, num_cache=0):
+        super().__init__()
+        self.path = pathlib.Path(path) / f'{prefix}.data'
+        if not self.path.exists():
+            raise FileNotFoundError(f'IndexedDataset not found: {self.path}')
+        self.dset = None
+        self.cache = deque(maxlen=num_cache)
+        self.num_cache = num_cache
+
+    def check_index(self, i):
+        if i < 0 or i >= len(self.dset):
+            raise IndexError('index out of range')
+
+    def __del__(self):
+        if self.dset:
+            self.dset.close()
+
+    def __getitem__(self, i):
+        if self.dset is None:
+            self.dset = h5py.File(self.path, 'r')
+        self.check_index(i)
+        if self.num_cache > 0:
+            for c in self.cache:
+                if c[0] == i:
+                    return c[1]
+        item = {k: v[()].item() if v.shape == () else torch.from_numpy(v[()]) for k, v in self.dset[str(i)].items()}
+        if self.num_cache > 0:
+            self.cache.appendleft((i, item))
+        return item
+
+    def __len__(self):
+        if self.dset is None:
+            self.dset = h5py.File(self.path, 'r')
+        return len(self.dset)
+
+
+class IndexedDatasetBuilder:
+    def __init__(self, path, prefix, allowed_attr=None, auto_increment=True):
+        self.path = pathlib.Path(path) / f'{prefix}.data'
+        self.prefix = prefix
+        self.dset = h5py.File(self.path, 'w')
+        self.counter = 0
+        self.auto_increment = auto_increment
+        if allowed_attr is not None:
+            self.allowed_attr = set(allowed_attr)
+        else:
+            self.allowed_attr = None
+
+    def add_item(self, item, item_no=None):
+        if self.auto_increment and item_no is not None or not self.auto_increment and item_no is None:
+            raise ValueError('auto_increment and provided item_no are mutually exclusive')
+        if self.allowed_attr is not None:
+            item = {
+                k: item[k]
+                for k in self.allowed_attr
+                if k in item
+            }
+        if self.auto_increment:
+            item_no = self.counter
+            self.counter += 1
+        for k, v in item.items():
+            if v is None:
+                continue
+            self.dset.create_dataset(f'{item_no}/{k}', data=v)
+        return item_no
+
+    def finalize(self):
+        self.dset.close()
+
+
+if __name__ == "__main__":
+    import random
+    from tqdm import tqdm
+
+    ds_path = './checkpoints/indexed_ds_example'
+    size = 100
+    items = [{"a": np.random.normal(size=[10000, 10]),
+              "b": np.random.normal(size=[10000, 10])} for i in range(size)]
+    builder = IndexedDatasetBuilder(ds_path, 'example')
+    for i in tqdm(range(size)):
+        builder.add_item(items[i])
+    builder.finalize()
+    ds = IndexedDataset(ds_path, 'example')
+    for i in tqdm(range(10000)):
+        idx = random.randint(0, size - 1)
+        assert (ds[idx]['a'] == items[idx]['a']).all()

utils/infer_utils.py

@@ -0,0 +1,104 @@
+import re
+
+import librosa
+import numpy as np
+from scipy.io import wavfile
+
+
+def trans_f0_seq(feature_pit, transform):
+    feature_pit = feature_pit * 2 ** (transform / 12)
+    return round(feature_pit, 1)
+
+
+def trans_key(raw_data, key):
+    warning_tag = False
+    for i in raw_data:
+        note_seq_list = i["note_seq"].split(" ")
+        new_note_seq_list = []
+        for note_seq in note_seq_list:
+            if note_seq != "rest":
+                new_note_seq = librosa.midi_to_note(librosa.note_to_midi(note_seq) + key, unicode=False)
+                # new_note_seq = move_key(note_seq, key)
+                new_note_seq_list.append(new_note_seq)
+            else:
+                new_note_seq_list.append(note_seq)
+        i["note_seq"] = " ".join(new_note_seq_list)
+        if i.get("f0_seq"):
+            f0_seq_list = i["f0_seq"].split(" ")
+            f0_seq_list = [float(x) for x in f0_seq_list]
+            new_f0_seq_list = []
+            for f0_seq in f0_seq_list:
+                new_f0_seq = trans_f0_seq(f0_seq, key)
+                new_f0_seq_list.append(str(new_f0_seq))
+            i["f0_seq"] = " ".join(new_f0_seq_list)
+        else:
+            warning_tag = True
+    if warning_tag:
+        print("Warning: parts of f0_seq do not exist, please freeze the pitch line in the editor.\r\n")
+    return raw_data
+
+
+def resample_align_curve(points: np.ndarray, original_timestep: float, target_timestep: float, align_length: int):
+    t_max = (len(points) - 1) * original_timestep
+    curve_interp = np.interp(
+        np.arange(0, t_max, target_timestep),
+        original_timestep * np.arange(len(points)),
+        points
+    ).astype(points.dtype)
+    delta_l = align_length - len(curve_interp)
+    if delta_l < 0:
+        curve_interp = curve_interp[:align_length]
+    elif delta_l > 0:
+        curve_interp = np.concatenate((curve_interp, np.full(delta_l, fill_value=curve_interp[-1])), axis=0)
+    return curve_interp
+
+
+def parse_commandline_spk_mix(mix: str) -> dict:
+    """
+    Parse speaker mix info from commandline
+    :param mix: Input like "opencpop" or "opencpop|qixuan" or "opencpop:0.5|qixuan:0.5"
+    :return: A dict whose keys are speaker names and values are proportions
+    """
+    name_pattern = r'[0-9A-Za-z_-]+'
+    proportion_pattern = r'\d+(\.\d+)?'
+    single_pattern = rf'{name_pattern}(:{proportion_pattern})?'
+    assert re.fullmatch(rf'{single_pattern}(\|{single_pattern})*', mix) is not None, f'Invalid mix pattern: {mix}'
+    without_proportion = set()
+    proportion_map = {}
+    for component in mix.split('|'):
+        # If already exists
+        name_and_proportion = component.split(':')
+        assert name_and_proportion[0] not in without_proportion and name_and_proportion[0] not in proportion_map, \
+            f'Duplicate speaker name: {name_and_proportion[0]}'
+        if ':' in component:
+            proportion_map[name_and_proportion[0]] = float(name_and_proportion[1])
+        else:
+            without_proportion.add(name_and_proportion[0])
+    sum_given_proportions = sum(proportion_map.values())
+    assert sum_given_proportions < 1 or len(without_proportion) == 0, \
+        'Proportion of all speakers should be specified if the sum of all given proportions are larger than 1.'
+    for name in without_proportion:
+        proportion_map[name] = (1 - sum_given_proportions) / len(without_proportion)
+    sum_all_proportions = sum(proportion_map.values())
+    assert sum_all_proportions > 0, 'Sum of all proportions should be positive.'
+    for name in proportion_map:
+        proportion_map[name] /= sum_all_proportions
+    return proportion_map
+
+
+def cross_fade(a: np.ndarray, b: np.ndarray, idx: int):
+    result = np.zeros(idx + b.shape[0])
+    fade_len = a.shape[0] - idx
+    np.copyto(dst=result[:idx], src=a[:idx])
+    k = np.linspace(0, 1.0, num=fade_len, endpoint=True)
+    result[idx: a.shape[0]] = (1 - k) * a[idx:] + k * b[: fade_len]
+    np.copyto(dst=result[a.shape[0]:], src=b[fade_len:])
+    return result
+
+
+def save_wav(wav, path, sr, norm=False):
+    if norm:
+        wav = wav / np.abs(wav).max()
+    wav *= 32767
+    # proposed by @dsmiller
+    wavfile.write(path, sr, wav.astype(np.int16))

utils/multiprocess_utils.py

@@ -0,0 +1,52 @@
+import platform
+import re
+import traceback
+
+from torch.multiprocessing import Manager, Process, current_process, get_context
+
+is_main_process = not bool(re.match(r'((.*Process)|(SyncManager)|(.*PoolWorker))-\d+', current_process().name))
+
+
+def main_process_print(self, *args, sep=' ', end='\n', file=None):
+    if is_main_process:
+        print(self, *args, sep=sep, end=end, file=file)
+
+
+def chunked_worker_run(map_func, args, results_queue=None):
+    for a in args:
+        # noinspection PyBroadException
+        try:
+            res = map_func(*a)
+            results_queue.put(res)
+        except KeyboardInterrupt:
+            break
+        except Exception:
+            traceback.print_exc()
+            results_queue.put(None)
+
+
+def chunked_multiprocess_run(map_func, args, num_workers, q_max_size=1000):
+    num_jobs = len(args)
+    if num_jobs < num_workers:
+        num_workers = num_jobs
+
+    queues = [Manager().Queue(maxsize=q_max_size // num_workers) for _ in range(num_workers)]
+    if platform.system().lower() != 'windows':
+        process_creation_func = get_context('spawn').Process
+    else:
+        process_creation_func = Process
+
+    workers = []
+    for i in range(num_workers):
+        worker = process_creation_func(
+            target=chunked_worker_run, args=(map_func, args[i::num_workers], queues[i]), daemon=True
+        )
+        workers.append(worker)
+        worker.start()
+
+    for i in range(num_jobs):
+        yield queues[i % num_workers].get()
+
+    for worker in workers:
+        worker.join()
+        worker.close()

utils/onnx_helper.py

@@ -0,0 +1,421 @@
+import inspect
+import re
+from typing import Dict, Tuple, Union, Literal
+
+import onnx
+import torch
+from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
+from onnx import GraphProto, ModelProto, NodeProto, ValueInfoProto
+
+__verbose__: bool = True
+"""
+Whether log information of successful operations
+"""
+
+
+# Whether the running torch.onnx.export() exposes the `dynamo` keyword.
+# Its presence is fragmented across PyTorch versions: introduced as a
+# separate `torch.onnx.dynamo_export` API in 2.1, added as a kwarg on
+# `torch.onnx.export` in 2.4, removed in some intermediate releases, and
+# reinstated (with a True default) in 2.9. We probe the signature once at
+# import time and only pass `dynamo=False` when the kwarg actually exists.
+# All ONNX graph surgery in this module is written against the TorchScript
+# exporter, so we want to stay on it whenever the choice is offered.
+TORCHSCRIPT_EXPORT_KWARGS: Dict[str, object] = (
+    {'dynamo': False}
+    if 'dynamo' in inspect.signature(torch.onnx.export).parameters
+    else {}
+)
+
+
+def _verbose(self, *args, sep=' ', end='\n', file=None):
+    if __verbose__:
+        print(self, *args, sep=sep, end=end, file=file)
+
+
+def model_override_io_shapes(
+        model: ModelProto,
+        input_shapes: Dict[str, Tuple[Union[str, int]]] = None,
+        output_shapes: Dict[str, Tuple[Union[str, int]]] = None,
+):
+    """
+    Override the shapes of inputs/outputs of the model graph (in-place operation).
+    :param model: model to perform the operation on
+    :param input_shapes: a dict with keys as input/output names and values as shape tuples
+    :param output_shapes: the same as input_shapes
+    """
+    def _override_shapes(
+            shape_list_old: RepeatedCompositeFieldContainer[ValueInfoProto],
+            shape_dict_new: Dict[str, Tuple[Union[str, int]]]):
+        for value_info in shape_list_old:
+            if value_info.name in shape_dict_new:
+                name = value_info.name
+                dims = value_info.type.tensor_type.shape.dim
+                assert len(shape_dict_new[name]) == len(dims), \
+                    f'Number of given and existing dimensions mismatch: {name}'
+                for i, dim in enumerate(shape_dict_new[name]):
+                    if isinstance(dim, int):
+                        dims[i].dim_param = ''
+                        dims[i].dim_value = dim
+                    else:
+                        dims[i].dim_value = 0
+                        dims[i].dim_param = dim
+                _verbose(f'| override shape of \'{name}\' with {shape_dict_new[name]}')
+
+    if input_shapes is not None:
+        _override_shapes(model.graph.input, input_shapes)
+    if output_shapes is not None:
+        _override_shapes(model.graph.output, output_shapes)
+
+
+def model_reorder_io_list(
+        model: ModelProto,
+        input_or_output: Literal['input', 'output'],
+        target_name: str,
+        insert_after_name: str,
+):
+    """
+    Reorder the input of the model graph by moving the target input after the specified input (in-place operation).
+    If the given names are not found, the operation will be ignored.
+    :param model: model to perform the operation on
+    :param input_or_output: 'input' or 'output' to specify the list to reorder
+    :param target_name: the name of the input to be reordered
+    :param insert_after_name: the name of the input to be inserted after (None for the first)
+    """
+    def _reorder_input(input_list: RepeatedCompositeFieldContainer[ValueInfoProto]):
+        nonlocal input_or_output
+        target_idx = -1
+        insert_after_idx = -1
+        for i, value_info in enumerate(input_list):
+            if value_info.name == target_name:
+                target_idx = i
+            if value_info.name == insert_after_name:
+                insert_after_idx = i
+        if target_idx != -1 and insert_after_idx != -1:
+            target = input_list.pop(target_idx)
+            input_list.insert(insert_after_idx + 1, target)
+            _verbose(f'| reorder {input_or_output}: \'{target_name}\' after \'{insert_after_name}\'')
+
+    if input_or_output == 'input':
+        _reorder_input(model.graph.input)
+    elif input_or_output == 'output':
+        _reorder_input(model.graph.output)
+    else:
+        raise ValueError('Argument \'input_or_output\' should be either \'input\' or \'output\'.')
+
+
+def model_add_prefixes(
+        model: ModelProto,
+        initializer_prefix=None,
+        value_info_prefix=None,
+        node_prefix=None,
+        dim_prefix=None,
+        ignored_pattern=None,
+):
+    """
+    Adds prefixes to names inside the given ONNX model graph, including sub-graphs (in-place operation).
+    This method is a complete version of the official onnx.compose.add_prefix API, which does not consider sub-graphs.
+    """
+    initializers = set()
+    value_infos = set()
+
+    def _record_initializers_and_value_infos_recursive(subgraph):
+        # Record names in current graph
+        for initializer in subgraph.initializer:
+            if ignored_pattern is not None and re.match(ignored_pattern, initializer.name):
+                continue
+            initializers.add(initializer.name)
+        for value_info in subgraph.value_info:
+            if ignored_pattern is not None and re.match(ignored_pattern, value_info.name):
+                continue
+            value_infos.add(value_info.name)
+        for node in subgraph.node:
+            # For 'If' and 'Loop' nodes, do recording recursively
+            if node.op_type == 'If':
+                for attr in node.attribute:
+                    branch = onnx.helper.get_attribute_value(attr)
+                    _record_initializers_and_value_infos_recursive(branch)
+            elif node.op_type == 'Loop':
+                for attr in node.attribute:
+                    if attr.name == 'body':
+                        body = onnx.helper.get_attribute_value(attr)
+                        _record_initializers_and_value_infos_recursive(body)
+
+    def _add_prefixes_recursive(subgraph):
+        # Add prefixes in current graph
+        if initializer_prefix is not None:
+            for initializer in subgraph.initializer:
+                if ignored_pattern is not None and re.match(ignored_pattern, initializer.name):
+                    continue
+                new_name = initializer_prefix + initializer.name
+                _verbose('| add prefix:', initializer.name, '->', new_name)
+                initializer.name = new_name
+
+        for value_info in subgraph.value_info:
+            if dim_prefix is not None:
+                for dim in value_info.type.tensor_type.shape.dim:
+                    if dim.dim_param is None or dim.dim_param == '' or \
+                            ignored_pattern is not None and re.match(ignored_pattern, dim.dim_param):
+                        continue
+                    new_dim_param = dim_prefix + dim.dim_param
+                    _verbose('| add prefix:', dim.dim_param, '->', new_dim_param)
+                    dim.dim_param = new_dim_param
+
+            if value_info_prefix is None or \
+                    ignored_pattern is not None and re.match(ignored_pattern, value_info.name):
+                continue
+            new_name = value_info_prefix + value_info.name
+            _verbose('| add prefix:', value_info.name, '->', new_name)
+            value_info.name = new_name
+
+        if node_prefix is not None:
+            for node in subgraph.node:
+                if ignored_pattern is not None and re.match(ignored_pattern, node.name):
+                    continue
+                new_name = node_prefix + node.name
+                _verbose('| add prefix:', node.name, '->', new_name)
+                node.name = new_name
+
+        for node in subgraph.node:
+            # For 'If' and 'Loop' nodes, add prefixes recursively
+            if node.op_type == 'If':
+                for attr in node.attribute:
+                    branch = onnx.helper.get_attribute_value(attr)
+                    _add_prefixes_recursive(branch)
+            elif node.op_type == 'Loop':
+                for attr in node.attribute:
+                    if attr.name == 'body':
+                        body = onnx.helper.get_attribute_value(attr)
+                        _add_prefixes_recursive(body)
+
+            # For each node, rename its inputs and outputs
+            for io_list in [node.input, node.output]:
+                for i, io_value in enumerate(io_list):
+                    if io_value in initializers and initializer_prefix is not None:
+                        new_value = initializer_prefix + io_value
+                        _verbose('| add prefix:', io_value, '->', new_value)
+                        io_list[i] = new_value
+                    if io_value in value_infos and value_info_prefix is not None:
+                        new_value = value_info_prefix + io_value
+                        _verbose('| add prefix:', io_value, '->', new_value)
+                        io_list[i] = new_value
+
+    _record_initializers_and_value_infos_recursive(model.graph)
+    _add_prefixes_recursive(model.graph)
+
+
+def graph_fold_back_to_squeeze(graph: GraphProto):
+    """
+    Fold the substructures of 'Shape', 'Gather', 'Equal', 'If' to one single 'Squeeze' node.
+    This can unify the different behaviors between aten::squeeze and onnx:Squeeze.
+    """
+    def _graph_fold_back_to_squeeze_recursive(subgraph: GraphProto):
+        # Do folding in sub-graphs recursively.
+        for node in subgraph.node:
+            if node.op_type == 'If':
+                for attr in node.attribute:
+                    branch = onnx.helper.get_attribute_value(attr)
+                    _graph_fold_back_to_squeeze_recursive(branch)
+            elif node.op_type == 'Loop':
+                for attr in node.attribute:
+                    if attr.name == 'body':
+                        body = onnx.helper.get_attribute_value(attr)
+                        _graph_fold_back_to_squeeze_recursive(body)
+
+        # Do folding in current graph.
+        i_shape = 0
+        while i_shape < len(subgraph.node):
+            if subgraph.node[i_shape].op_type == 'Shape':
+                shape_node = subgraph.node[i_shape]
+                shape_out = shape_node.output[0]
+                i_gather = i_shape + 1
+                while i_gather < len(subgraph.node):
+                    if subgraph.node[i_gather].op_type == 'Gather' and subgraph.node[i_gather].input[0] == shape_out:
+                        gather_node = subgraph.node[i_gather]
+                        gather_out = gather_node.output[0]
+                        i_equal = i_gather + 1
+                        while i_equal < len(subgraph.node):
+                            if subgraph.node[i_equal].op_type == 'Equal' and (
+                                    subgraph.node[i_equal].input[0] == gather_out
+                                    or subgraph.node[i_equal].input[1] == gather_out):
+                                equal_node = subgraph.node[i_equal]
+                                equal_out = equal_node.output[0]
+                                i_if = i_equal + 1
+                                while i_if < len(subgraph.node):
+                                    if subgraph.node[i_if].op_type == 'If' \
+                                            and subgraph.node[i_if].input[0] == equal_out:
+                                        # Found the substructure to be folded.
+                                        if_node = subgraph.node[i_if]
+                                        # Create 'Squeeze' node.
+                                        squeeze_node = onnx.helper.make_node(
+                                            op_type='Squeeze',
+                                            inputs=[
+                                                *list(shape_node.input),
+                                                # For ONNX opset >= 13, axes should be an input instead of an attribute.
+                                                gather_node.input[1]  # Use 'indices' input of 'Gather'
+                                            ],
+                                            outputs=if_node.output,
+                                            name=shape_node.name.replace('Shape', 'Squeeze')
+                                        )
+                                        # Replace 'Shape', 'Gather', 'Equal', 'If' with 'Squeeze'.
+                                        subgraph.node.insert(i_shape, squeeze_node)
+                                        subgraph.node.remove(shape_node)
+                                        subgraph.node.remove(gather_node)
+                                        subgraph.node.remove(equal_node)
+                                        subgraph.node.remove(if_node)
+                                        _verbose(
+                                            f'| fold nodes: [\'{shape_node.name}\', \'{gather_node.name}\', '
+                                            f'\'{equal_node.name}\', \'{if_node.name}\'] -> \'{squeeze_node.name}\'')
+                                        break
+                                    i_if += 1
+                                else:
+                                    break
+                            i_equal += 1
+                        else:
+                            break
+                    i_gather += 1
+                else:
+                    break
+            i_shape += 1
+
+    _graph_fold_back_to_squeeze_recursive(graph)
+
+
+def graph_extract_conditioner_projections(
+        graph: GraphProto,
+        op_type: str,
+        weight_pattern: str,
+        alias_prefix: str
+):
+    """
+    Extract conditioner projection nodes out of the backbone wrapped by diffusion.
+    These nodes only need to be calculated once before entering the main denoising loop,
+    and can be reused inside the loop. This optimizes the performance of ONNX inference.
+
+    :param graph: graph to perform the operation on
+    :param op_type: the ONNX operator type of the conditioner projections (usually 'Conv' or 'Gemm')
+    :param weight_pattern: a regular expression as pattern of the conditioner projection weight keys
+    :param alias_prefix: add prefixes to the outputs of extracted projection nodes
+    """
+    node_dict: Dict[str, Tuple[str, NodeProto]] = {}  # key: pattern match, value: (alias, node)
+
+    def _extract_conv_nodes_recursive(subgraph: GraphProto):
+        to_be_removed = []
+        for sub_node in subgraph.node:
+            if sub_node.op_type == 'If':
+                for attr in sub_node.attribute:
+                    branch = onnx.helper.get_attribute_value(attr)
+                    _extract_conv_nodes_recursive(branch)
+            elif sub_node.op_type == 'Loop':
+                for attr in sub_node.attribute:
+                    if attr.name == 'body':
+                        body = onnx.helper.get_attribute_value(attr)
+                        _extract_conv_nodes_recursive(body)
+            elif sub_node.op_type == op_type and re.match(weight_pattern, sub_node.input[1]):
+                # Found node to extract
+                cached = node_dict.get(sub_node.input[1])
+                if cached is None:
+                    out_alias = f'{alias_prefix}.{len(node_dict)}'
+                    node_dict[sub_node.input[1]] = (out_alias, sub_node)
+                else:
+                    out_alias = cached[0]
+                out = sub_node.output[0]
+                # Search for nodes downstream the extracted node and match them to the renamed output.
+                for dep_node in subgraph.node:
+                    for dep_idx, dep_input in enumerate(dep_node.input):
+                        if dep_input == out:
+                            dep_node.input.remove(out)
+                            dep_node.input.insert(dep_idx, out_alias)
+                # Add the node to the remove list.
+                to_be_removed.append(sub_node)
+        [subgraph.node.remove(_n) for _n in to_be_removed]
+
+    toplevel_entry_node_idx = toplevel_entry_node = None
+    # Find the **last** If node in toplevel graph
+    for i, n in enumerate(graph.node):
+        if n.op_type == 'If':
+            toplevel_entry_node_idx = i
+            toplevel_entry_node = n
+    # If not found, find the **last** Loop node in toplevel graph
+    if toplevel_entry_node is None:
+        for i, n in enumerate(graph.node):
+            if n.op_type == 'Loop':
+                toplevel_entry_node_idx = i
+                toplevel_entry_node = n
+    if toplevel_entry_node is not None:
+        for a in toplevel_entry_node.attribute:
+            # Apply to all sub-graphs
+            v = onnx.helper.get_attribute_value(a)
+            if isinstance(v, GraphProto):
+                _extract_conv_nodes_recursive(v)
+
+        # Insert the extracted nodes before the first 'If' node which carries the main denoising loop.
+        for key in reversed(node_dict):
+            alias, node = node_dict[key]
+            # Rename output of the node.
+            out_name = node.output[0]
+            node.output.remove(node.output[0])
+            node.output.insert(0, alias)
+            # Insert node into the main graph.
+            graph.node.insert(toplevel_entry_node_idx, node)
+            # Rename value info of the output.
+            for v in graph.value_info:
+                if v.name == out_name:
+                    v.name = alias
+                    break
+            _verbose(f'| extract conditioner projection: \'{node.name}\'')
+
+
+def graph_remove_unused_values(graph: GraphProto):
+    used_values = set()
+
+    def _record_usage_recursive(subgraph: GraphProto):
+        for node in subgraph.node:
+            # For 'If' and 'Loop' nodes, do recording recursively
+            if node.op_type == 'If':
+                for attr in node.attribute:
+                    branch = onnx.helper.get_attribute_value(attr)
+                    _record_usage_recursive(branch)
+            elif node.op_type == 'Loop':
+                for attr in node.attribute:
+                    if attr.name == 'body':
+                        body = onnx.helper.get_attribute_value(attr)
+                        _record_usage_recursive(body)
+            # For each node, record its inputs and outputs
+            for io_list in [node.input, node.output]:
+                for io_value in io_list:
+                    used_values.add(io_value)
+
+    def _clean_unused_recursively(subgraph):
+        # Do cleaning in sub-graphs recursively.
+        for node in subgraph.node:
+            if node.op_type == 'If':
+                for attr in node.attribute:
+                    branch = onnx.helper.get_attribute_value(attr)
+                    _clean_unused_recursively(branch)
+            elif node.op_type == 'Loop':
+                for attr in node.attribute:
+                    if attr.name == 'body':
+                        body = onnx.helper.get_attribute_value(attr)
+                        _clean_unused_recursively(body)
+
+        # Do cleaning in current graph.
+        i = 0
+        while i < len(subgraph.initializer):
+            name = subgraph.initializer[i].name
+            if name not in used_values:
+                subgraph.initializer.pop(i)
+                _verbose(f'| remove unused initializer: {name}')
+            else:
+                i += 1
+        i = 0
+        while i < len(subgraph.value_info):
+            name = subgraph.value_info[i].name
+            if name not in used_values:
+                subgraph.value_info.pop(i)
+                _verbose(f'| remove unused value info: {name}')
+            else:
+                i += 1
+
+    _record_usage_recursive(graph)
+    _clean_unused_recursively(graph)

utils/phoneme_utils.py

@@ -0,0 +1,210 @@
+import json
+import pathlib
+from typing import Dict, List, Union
+
+from utils.hparams import hparams
+
+PAD_INDEX = 0
+
+
+class PhonemeDictionary:
+    def __init__(
+            self,
+            dictionaries: Dict[str, pathlib.Path],
+            extra_phonemes: List[str] = None,
+            merged_groups: List[List[str]] = None
+    ):
+        # Step 1: Collect all phonemes
+        all_phonemes = {'AP', 'SP'}
+        if extra_phonemes:
+            for ph in extra_phonemes:
+                if '/' in ph:
+                    lang, name = ph.split('/', maxsplit=1)
+                    if lang not in dictionaries:
+                        raise ValueError(
+                            f"Invalid phoneme tag '{ph}' in extra phonemes: "
+                            f"unrecognized language name '{lang}'."
+                        )
+                    if name in all_phonemes:
+                        raise ValueError(
+                            f"Invalid phoneme tag '{ph}' in extra phonemes: "
+                            f"short name conflicts with existing tag."
+                        )
+                all_phonemes.add(ph)
+        self._multi_langs = len(dictionaries) > 1
+        for lang, dict_path in dictionaries.items():
+            with open(dict_path, 'r', encoding='utf8') as dict_file:
+                for line in dict_file:
+                    _, phonemes = line.strip().split('\t')
+                    phonemes = phonemes.split()
+                    for phoneme in phonemes:
+                        if '/' in phoneme:
+                            raise ValueError(
+                                f"Invalid phoneme tag '{phoneme}' in dictionary '{dict_path}': "
+                                f"should not contain the reserved character '/'."
+                            )
+                        if phoneme in all_phonemes:
+                            continue
+                        if self._multi_langs:
+                            all_phonemes.add(f'{lang}/{phoneme}')
+                        else:
+                            all_phonemes.add(phoneme)
+        # Step 2: Parse merged phoneme groups
+        if merged_groups is None:
+            merged_groups = []
+        else:
+            _merged_groups = []
+            for group in merged_groups:
+                _group = []
+                for phoneme in group:
+                    if '/' in phoneme:
+                        lang, name = phoneme.split('/', maxsplit=1)
+                        if lang not in dictionaries:
+                            raise ValueError(
+                                f"Invalid phoneme tag '{phoneme}' in merged group: "
+                                f"unrecognized language name '{lang}'."
+                            )
+                        if self._multi_langs:
+                            element = phoneme
+                        else:
+                            element = name
+                    else:
+                        element = phoneme
+                    if element not in all_phonemes:
+                        raise ValueError(
+                            f"Invalid phoneme tag '{phoneme}' in merged group: "
+                            f"not found in phoneme set."
+                        )
+                    _group.append(element)
+                _merged_groups.append(_group)
+            merged_groups = [set(phones) for phones in _merged_groups if len(phones) > 1]
+        # Step 3: Build phoneme index
+        merged_phonemes_inverted_index = {}
+        for idx, group in enumerate(merged_groups):
+            other_idx = None
+            for phoneme in group:
+                if phoneme in merged_phonemes_inverted_index:
+                    other_idx = merged_phonemes_inverted_index[phoneme]
+                    break
+            target_idx = idx if other_idx is None else other_idx
+            for phoneme in group:
+                merged_phonemes_inverted_index[phoneme] = target_idx
+            if other_idx is not None:
+                merged_groups[other_idx] |= group
+                group.clear()
+        phone_to_id = {}
+        id_to_phone = []
+        cross_lingual_phonemes = set()
+        idx = 1
+        for phoneme in sorted(all_phonemes):
+            if phoneme in merged_phonemes_inverted_index:
+                has_assigned = True
+                for alias in merged_groups[merged_phonemes_inverted_index[phoneme]]:
+                    if alias not in phone_to_id:
+                        has_assigned = False
+                        phone_to_id[alias] = idx
+                if not has_assigned:
+                    merged_group = sorted(merged_groups[merged_phonemes_inverted_index[phoneme]])
+                    merged_from_langs = {
+                        (alias.split('/', maxsplit=1)[0] if '/' in alias else None)
+                        for alias in merged_group
+                    }
+                    id_to_phone.append(tuple(merged_group))
+                    idx += 1
+                    if len(merged_from_langs) > 1:
+                        cross_lingual_phonemes.update(ph for ph in merged_group if '/' in ph)
+            else:
+                phone_to_id[phoneme] = idx
+                id_to_phone.append(phoneme)
+                idx += 1
+        self._phone_to_id: Dict[str, int] = phone_to_id
+        self._id_to_phone: List[Union[str, tuple]] = id_to_phone
+        self._cross_lingual_phonemes = frozenset(cross_lingual_phonemes)
+
+    @property
+    def vocab_size(self):
+        return len(self._id_to_phone) + 1
+
+    def __len__(self):
+        return self.vocab_size
+
+    @property
+    def cross_lingual_phonemes(self):
+        return self._cross_lingual_phonemes
+
+    def is_cross_lingual(self, phone):
+        return phone in self._cross_lingual_phonemes
+
+    def encode_one(self, phone, lang=None):
+        if '/' in phone:
+            lang, phone = phone.split('/', maxsplit=1)
+        if lang is None or not self._multi_langs or phone in self._phone_to_id:
+            return self._phone_to_id[phone]
+        if '/' not in phone:
+            phone = f'{lang}/{phone}'
+        return self._phone_to_id[phone]
+
+    def encode(self, sentence, lang=None):
+        phones = sentence.strip().split() if isinstance(sentence, str) else sentence
+        return [self.encode_one(phone, lang=lang) for phone in phones]
+
+    def decode_one(self, idx, lang=None, scalar=True):
+        if idx <= 0:
+            return None
+        phone = self._id_to_phone[idx - 1]
+        if not scalar or isinstance(phone, str):
+            return phone
+        if lang is None or not self._multi_langs:
+            return phone[0]
+        for alias in phone:
+            if alias.startswith(f'{lang}/'):
+                return alias
+        return phone[0]
+
+    def decode(self, ids, lang=None, scalar=True):
+        ids = list(ids)
+        return ' '.join([
+            self.decode_one(i, lang=lang, scalar=scalar)
+            for i in ids
+            if i >= 1
+        ])
+
+    def dump(self, filename):
+        with open(filename, 'w', encoding='utf8') as fp:
+            json.dump(self._phone_to_id, fp, ensure_ascii=False, indent=2)
+
+
+_dictionary = None
+
+
+def load_phoneme_dictionary() -> PhonemeDictionary:
+    if _dictionary is not None:
+        return _dictionary
+    config_dicts = hparams.get('dictionaries')
+    if config_dicts is not None:
+        dicts = {}
+        for lang, config_dict_path in config_dicts.items():
+            dict_path = pathlib.Path(hparams['work_dir']) / f'dictionary-{lang}.txt'
+            if not dict_path.exists():
+                dict_path = pathlib.Path(config_dict_path)
+            if not dict_path.exists():
+                raise FileNotFoundError(
+                    f"Could not locate dictionary for language '{lang}'."
+                )
+            dicts[lang] = dict_path
+    else:
+        dict_path = pathlib.Path(hparams['work_dir']) / 'dictionary.txt'
+        if not dict_path.exists():
+            dict_path = pathlib.Path(hparams['dictionary'])
+        if not dict_path.exists():
+            raise FileNotFoundError(
+                f"Could not locate dictionary file."
+            )
+        dicts = {
+            'default': dict_path
+        }
+    return PhonemeDictionary(
+        dictionaries=dicts,
+        extra_phonemes=hparams.get('extra_phonemes'),
+        merged_groups=hparams.get('merged_phoneme_groups')
+    )

utils/pitch_utils.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+
+def norm_f0(f0, uv=None):
+    if uv is None:
+        uv = f0 == 0
+    f0 = np.log2(f0 + uv)  # avoid arithmetic error
+    f0[uv] = -np.inf
+    return f0
+
+
+def interp_f0(f0, uv=None):
+    if uv is None:
+        uv = f0 == 0
+    f0 = norm_f0(f0, uv)
+    if uv.any() and not uv.all():
+        f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv])
+    return denorm_f0(f0, uv=None), uv
+
+
+def denorm_f0(f0, uv, pitch_padding=None):
+    f0 = 2 ** f0
+    if uv is not None:
+        f0[uv > 0] = 0
+    if pitch_padding is not None:
+        f0[pitch_padding] = 0
+    return f0

utils/plot.py

@@ -0,0 +1,122 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.ticker import MultipleLocator
+
+
+def spec_to_figure(spec, vmin=None, vmax=None, title=None):
+    if isinstance(spec, torch.Tensor):
+        spec = spec.cpu().numpy()
+    fig = plt.figure(figsize=(12, 9))
+    plt.pcolormesh(spec.T, vmin=vmin, vmax=vmax)
+    if title is not None:
+        plt.title(title, fontsize=15)
+    plt.tight_layout()
+    return fig
+
+
+def dur_to_figure(dur_gt, dur_pred, txt, title=None):
+    if isinstance(dur_gt, torch.Tensor):
+        dur_gt = dur_gt.cpu().numpy()
+    if isinstance(dur_pred, torch.Tensor):
+        dur_pred = dur_pred.cpu().numpy()
+    dur_gt = dur_gt.astype(np.int64)
+    dur_pred = dur_pred.astype(np.int64)
+    dur_gt = np.cumsum(dur_gt)
+    dur_pred = np.cumsum(dur_pred)
+    width = max(12, min(48, len(txt) // 2))
+    fig = plt.figure(figsize=(width, 8))
+    plt.vlines(dur_pred, 12, 22, colors='r', label='pred')
+    plt.vlines(dur_gt, 0, 10, colors='b', label='gt')
+    for i in range(len(txt)):
+        shift = (i % 8) + 1
+        plt.text((dur_pred[i-1] + dur_pred[i]) / 2 if i > 0 else dur_pred[i] / 2, 12 + shift, txt[i],
+                 size=16, horizontalalignment='center')
+        plt.text((dur_gt[i-1] + dur_gt[i]) / 2 if i > 0 else dur_gt[i] / 2, shift, txt[i],
+                 size=16, horizontalalignment='center')
+        plt.plot([dur_pred[i], dur_gt[i]], [12, 10], color='black', linewidth=2, linestyle=':')
+    plt.yticks([])
+    plt.xlim(0, max(dur_pred[-1], dur_gt[-1]))
+    plt.legend()
+    if title is not None:
+        plt.title(title, fontsize=15)
+    plt.tight_layout()
+    return fig
+
+
+def pitch_note_to_figure(pitch_gt, pitch_pred=None, note_midi=None, note_dur=None, note_rest=None, title=None):
+    if isinstance(pitch_gt, torch.Tensor):
+        pitch_gt = pitch_gt.cpu().numpy()
+    if isinstance(pitch_pred, torch.Tensor):
+        pitch_pred = pitch_pred.cpu().numpy()
+    if isinstance(note_midi, torch.Tensor):
+        note_midi = note_midi.cpu().numpy()
+    if isinstance(note_dur, torch.Tensor):
+        note_dur = note_dur.cpu().numpy()
+    if isinstance(note_rest, torch.Tensor):
+        note_rest = note_rest.cpu().numpy()
+    fig = plt.figure()
+    if note_midi is not None and note_dur is not None:
+        note_dur_acc = np.cumsum(note_dur)
+        if note_rest is None:
+            note_rest = np.zeros_like(note_midi, dtype=np.bool_)
+        for i in range(len(note_midi)):
+            # if note_rest[i]:
+            #     continue
+            plt.gca().add_patch(
+                plt.Rectangle(
+                    xy=(note_dur_acc[i-1] if i > 0 else 0, note_midi[i] - 0.5),
+                    width=note_dur[i], height=1,
+                    edgecolor='grey', fill=False,
+                    linewidth=1.5, linestyle='--' if note_rest[i] else '-'
+                )
+            )
+    plt.plot(pitch_gt, color='b', label='gt')
+    if pitch_pred is not None:
+        plt.plot(pitch_pred, color='r', label='pred')
+    plt.gca().yaxis.set_major_locator(MultipleLocator(1))
+    plt.grid(axis='y')
+    plt.legend()
+    if title is not None:
+        plt.title(title, fontsize=15)
+    plt.tight_layout()
+    return fig
+
+
+def curve_to_figure(curve_gt, curve_pred=None, curve_base=None, grid=None, title=None):
+    if isinstance(curve_gt, torch.Tensor):
+        curve_gt = curve_gt.cpu().numpy()
+    if isinstance(curve_pred, torch.Tensor):
+        curve_pred = curve_pred.cpu().numpy()
+    if isinstance(curve_base, torch.Tensor):
+        curve_base = curve_base.cpu().numpy()
+    fig = plt.figure()
+    if curve_base is not None:
+        plt.plot(curve_base, color='g', label='base')
+    plt.plot(curve_gt, color='b', label='gt')
+    if curve_pred is not None:
+        plt.plot(curve_pred, color='r', label='pred')
+    if grid is not None:
+        plt.gca().yaxis.set_major_locator(MultipleLocator(grid))
+    plt.grid(axis='y')
+    plt.legend()
+    if title is not None:
+        plt.title(title, fontsize=15)
+    plt.tight_layout()
+    return fig
+
+
+def distribution_to_figure(title, x_label, y_label, items: list, values: list, zoom=0.8, rotate=False):
+    fig = plt.figure(figsize=(int(len(items) * zoom), 10))
+    plt.bar(x=items, height=values)
+    plt.tick_params(labelsize=15)
+    plt.xlim(-1, len(items))
+    for a, b in zip(items, values):
+        plt.text(a, b, b, ha='center', va='bottom', fontsize=15)
+    plt.grid()
+    plt.title(title, fontsize=30)
+    plt.xlabel(x_label, fontsize=20)
+    plt.ylabel(y_label, fontsize=20)
+    if rotate:
+        fig.autofmt_xdate(rotation=45)
+    return fig

utils/training_utils.py

@@ -0,0 +1,447 @@
+import math
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Dict
+
+import lightning.pytorch as pl
+import numpy as np
+import torch
+from lightning.fabric.loggers.tensorboard import _TENSORBOARD_AVAILABLE
+from lightning.pytorch.callbacks import ModelCheckpoint, TQDMProgressBar
+from lightning.pytorch.loggers import TensorBoardLogger
+from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_only
+from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data.distributed import Sampler
+
+import utils
+from utils.hparams import hparams
+
+
+# ==========LR schedulers==========
+
+class RSQRTSchedule(object):
+    def __init__(self, optimizer):
+        super().__init__()
+        self.optimizer = optimizer
+        self.constant_lr = hparams['lr']
+        self.warmup_updates = hparams['warmup_updates']
+        self.hidden_size = hparams['hidden_size']
+        self.lr = hparams['lr']
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = self.lr
+        self.step(0)
+
+    def step(self, num_updates):
+        constant_lr = self.constant_lr
+        warmup = min(num_updates / self.warmup_updates, 1.0)
+        rsqrt_decay = max(self.warmup_updates, num_updates) ** -0.5
+        rsqrt_hidden = self.hidden_size ** -0.5
+        self.lr = max(constant_lr * warmup * rsqrt_decay * rsqrt_hidden, 1e-7)
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = self.lr
+        return self.lr
+
+    def get_lr(self):
+        return self.optimizer.param_groups[0]['lr']
+
+
+class WarmupCosineSchedule(LambdaLR):
+    """ Linear warmup and then cosine decay.
+        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
+        Decreases learning rate from 1. to 0. over remaining `t_total - warmup_steps` steps following a cosine curve.
+        If `cycles` (default=0.5) is different from default, learning rate follows cosine function after warmup.
+        `eta_min` (default=0.0) corresponds to the minimum learning rate reached by the scheduler.
+    """
+
+    def __init__(self, optimizer, warmup_steps, t_total, warmup_min=0.0, eta_min=0.0, cycles=.5, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.t_total = t_total
+        self.eta_min = eta_min
+        self.cycles = cycles
+        self.warmup_min = warmup_min
+        super(WarmupCosineSchedule, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)
+
+    def lr_lambda(self, step):
+        if step < self.warmup_steps:
+            progress = step / max(1.0, self.warmup_steps)
+            return self.warmup_min + progress * (1.0 - self.warmup_min)
+        # progress after warmup
+        progress = (step - self.warmup_steps) / max(1, self.t_total - self.warmup_steps)
+        return max(self.eta_min, 0.5 * (1. + math.cos(math.pi * self.cycles * 2.0 * progress)))
+
+
+# ==========Torch samplers==========
+
+class DsBatchSampler(Sampler):
+    def __init__(self, dataset, max_batch_frames, max_batch_size, sub_indices=None,
+                 num_replicas=None, rank=None,
+                 required_batch_count_multiple=1, batch_by_size=True, sort_by_similar_size=True,
+                 size_reversed=False, shuffle_sample=False, shuffle_batch=False,
+                 disallow_empty_batch=True, pad_batch_assignment=True, seed=0, drop_last=False) -> None:
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                f"Invalid rank {rank}, rank should be in the interval [0, {num_replicas - 1}]")
+        self.dataset = dataset
+        self.max_batch_frames = max_batch_frames
+        self.max_batch_size = max_batch_size
+        self.sub_indices = sub_indices
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.required_batch_count_multiple = required_batch_count_multiple
+        self.batch_by_size = batch_by_size
+        self.sort_by_similar_size = sort_by_similar_size
+        self.size_reversed = size_reversed
+        self.shuffle_sample = shuffle_sample
+        self.shuffle_batch = shuffle_batch
+        self.disallow_empty_batch = disallow_empty_batch
+        self.pad_batch_assignment = pad_batch_assignment
+        self.seed = seed
+        self.drop_last = drop_last
+        self.epoch = 0
+        self.batches = None
+        self.formed = None
+
+    def __form_batches(self):
+        if self.formed == self.epoch + self.seed:
+            return
+        rng = np.random.default_rng()
+        # Create indices
+        if self.shuffle_sample:
+            if self.sub_indices is not None:
+                rng.shuffle(self.sub_indices)
+                indices = np.array(self.sub_indices)
+            else:
+                indices = rng.permutation(len(self.dataset))
+
+            if self.sort_by_similar_size:
+                grid = int(hparams['sampler_frame_count_grid'])
+                assert grid > 0
+                sizes = (np.round(np.array(self.dataset.sizes)[indices] / grid) * grid).clip(grid, None)
+                sizes *= (-1 if self.size_reversed else 1)
+                indices = indices[np.argsort(sizes, kind='mergesort')]
+
+            indices = indices.tolist()
+        else:
+            indices = self.sub_indices if self.sub_indices is not None else list(range(len(self.dataset)))
+
+        # Batching
+        if self.batch_by_size:
+            batches = utils.batch_by_size(
+                indices, self.dataset.num_frames,
+                max_batch_frames=self.max_batch_frames,
+                max_batch_size=self.max_batch_size
+            )
+        else:
+            batches = [indices[i:i + self.max_batch_size] for i in range(0, len(indices), self.max_batch_size)]
+        if len(batches) < self.num_replicas and self.disallow_empty_batch:
+            raise RuntimeError("There is not enough batch to assign to each node.")
+
+        # Either drop_last or separate the leftovers.
+        floored_total_batch_count = (len(batches) // self.num_replicas) * self.num_replicas
+        if self.drop_last and len(batches) > floored_total_batch_count:
+            batches = batches[:floored_total_batch_count]
+            leftovers = []
+            if len(batches) == 0:
+                raise RuntimeError("There is no batch left after dropping the last batch.")
+        elif self.shuffle_batch:
+            leftovers = (rng.permutation(len(batches) - floored_total_batch_count) + floored_total_batch_count).tolist()
+        else:
+            leftovers = list(range(floored_total_batch_count, len(batches)))
+
+        # Initial batch assignment to current rank.
+        batch_assignment = np.arange(floored_total_batch_count).reshape(-1, self.num_replicas).transpose()
+        if self.shuffle_batch:
+            batch_assignment = rng.permuted(batch_assignment, axis=0)[self.rank].tolist()
+        else:
+            batch_assignment = batch_assignment[self.rank].tolist()
+
+        # Assign leftovers or pad the batch assignment.
+        floored_batch_count = len(batch_assignment)
+        if self.rank < len(leftovers):
+            batch_assignment.append(leftovers[self.rank])
+            floored_batch_count += 1
+        elif len(leftovers) > 0 and self.pad_batch_assignment:
+            if not batch_assignment:
+                raise RuntimeError("Cannot pad empty batch assignment.")
+            batch_assignment.append(batch_assignment[self.epoch % floored_batch_count])
+        # Ensure the batch count is multiple of required_batch_count_multiple.
+        if self.required_batch_count_multiple > 1 and len(batch_assignment) % self.required_batch_count_multiple != 0:
+            ceiled_batch_count = math.ceil(
+                len(batch_assignment) / self.required_batch_count_multiple
+            ) * self.required_batch_count_multiple
+            for i in range(ceiled_batch_count - len(batch_assignment)):
+                batch_assignment.append(
+                    batch_assignment[(i + self.epoch * self.required_batch_count_multiple) % floored_batch_count])
+
+        if batch_assignment:
+            self.batches = [deepcopy(batches[i]) for i in batch_assignment]
+        else:
+            self.batches = [[]]
+        self.formed = self.epoch + self.seed
+
+        del indices
+        del batches
+        del batch_assignment
+
+    def __iter__(self):
+        self.__form_batches()
+        return iter(self.batches)
+
+    def __len__(self):
+        self.__form_batches()
+        if self.batches is None:
+            raise RuntimeError("Batches are not initialized. Call __form_batches first.")
+        return len(self.batches)
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+        self.__form_batches()
+
+
+
+# ==========PL related==========
+
+class DsModelCheckpoint(ModelCheckpoint):
+    def __init__(
+            self,
+            *args,
+            permanent_ckpt_start,
+            permanent_ckpt_interval,
+            **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.permanent_ckpt_start = permanent_ckpt_start or 0
+        self.permanent_ckpt_interval = permanent_ckpt_interval or 0
+        self.enable_permanent_ckpt = self.permanent_ckpt_start > 0 and self.permanent_ckpt_interval > 9
+
+        self._verbose = self.verbose
+        self.verbose = False
+
+    def state_dict(self):
+        ret = super().state_dict()
+        ret.pop('dirpath')
+        return ret
+
+    def load_state_dict(self, state_dict) -> None:
+        super().load_state_dict(state_dict)
+
+    def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        if trainer.lightning_module.skip_immediate_ckpt_save:
+            trainer.lightning_module.skip_immediate_ckpt_save = False
+            return
+        self.last_val_step = trainer.global_step
+        super().on_validation_end(trainer, pl_module)
+
+    def _update_best_and_save(
+            self, current: torch.Tensor, trainer: "pl.Trainer", monitor_candidates: Dict[str, torch.Tensor]
+    ) -> None:
+        k = len(self.best_k_models) + 1 if self.save_top_k == -1 else self.save_top_k
+
+        del_filepath = None
+        _op = max if self.mode == "min" else min
+        while len(self.best_k_models) > k and k > 0:
+            self.kth_best_model_path = _op(self.best_k_models, key=self.best_k_models.get)  # type: ignore[arg-type]
+            self.kth_value = self.best_k_models[self.kth_best_model_path]
+
+            del_filepath = self.kth_best_model_path
+            self.best_k_models.pop(del_filepath)
+            filepath = self._get_metric_interpolated_filepath_name(monitor_candidates, trainer, del_filepath)
+            if del_filepath is not None and filepath != del_filepath:
+                self._remove_checkpoint(trainer, del_filepath)
+
+        if len(self.best_k_models) == k and k > 0:
+            self.kth_best_model_path = _op(self.best_k_models, key=self.best_k_models.get)  # type: ignore[arg-type]
+            self.kth_value = self.best_k_models[self.kth_best_model_path]
+
+        super()._update_best_and_save(current, trainer, monitor_candidates)
+
+    def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+        filepath = (Path(self.dirpath) / Path(filepath).name).resolve()
+        super()._save_checkpoint(trainer, str(filepath))
+        if self._verbose:
+            relative_path = filepath
+            # Avoid using `is_relative_to` because Python 3.8 does not support this
+            if Path('.').resolve() in filepath.parents:
+                relative_path = filepath.relative_to(Path('.').resolve())
+            rank_zero_info(f'Checkpoint {relative_path} saved.')
+
+    def _remove_checkpoint(self, trainer: "pl.Trainer", filepath: str):
+        filepath = (Path(self.dirpath) / Path(filepath).name).resolve()
+        relative_path = filepath
+        # Avoid using `is_relative_to` because Python 3.8 does not support this
+        if Path('.').resolve() in filepath.parents:
+            relative_path = filepath.relative_to(Path('.').resolve())
+        search = re.search(r'steps_\d+', relative_path.stem)
+        if search:
+            step = int(search.group(0)[6:])
+            if self.enable_permanent_ckpt and \
+                    step >= self.permanent_ckpt_start and \
+                    (step - self.permanent_ckpt_start) % self.permanent_ckpt_interval == 0:
+                rank_zero_info(f'Checkpoint {relative_path} is now permanent.')
+                return
+        super()._remove_checkpoint(trainer, filepath)
+        if self._verbose:
+            rank_zero_info(f'Removed checkpoint {relative_path}.')
+
+
+def get_latest_checkpoint_path(work_dir):
+    if not isinstance(work_dir, Path):
+        work_dir = Path(work_dir)
+    if not work_dir.exists():
+        return None
+
+    last_step = -1
+    last_ckpt_name = None
+
+    for ckpt in work_dir.glob('model_ckpt_steps_*.ckpt'):
+        search = re.search(r'steps_\d+', ckpt.name)
+        if search:
+            step = int(search.group(0)[6:])
+            if step > last_step:
+                last_step = step
+                last_ckpt_name = str(ckpt)
+
+    return last_ckpt_name if last_ckpt_name is not None else None
+
+
+class DsTQDMProgressBar(TQDMProgressBar):
+    def __init__(self, refresh_rate: int = 1, process_position: int = 0, show_steps: bool = True):
+        super().__init__(refresh_rate, process_position)
+        self.show_steps = show_steps
+
+    def get_metrics(self, trainer, model):
+        items = super().get_metrics(trainer, model)
+        if 'batch_size' in items:
+            items['batch_size'] = int(items['batch_size'])
+        if self.show_steps:
+            items['steps'] = str(trainer.global_step)
+        for k, v in items.items():
+            if isinstance(v, float):
+                if np.isnan(v):
+                    items[k] = 'nan'
+                elif 0.001 <= v < 10:
+                    items[k] = np.format_float_positional(v, unique=True, precision=5, trim='-')
+                elif 0.00001 <= v < 0.001:
+                    if len(np.format_float_positional(v, unique=True, precision=8, trim='-')) > 8:
+                        items[k] = np.format_float_scientific(v, precision=3, unique=True, min_digits=2, trim='-')
+                    else:
+                        items[k] = np.format_float_positional(v, unique=True, precision=5, trim='-')
+                elif v < 0.00001:
+                    items[k] = np.format_float_scientific(v, precision=3, unique=True, min_digits=2, trim='-')
+        items.pop("v_num", None)
+        return items
+
+
+class DsTensorBoardLogger(TensorBoardLogger):
+    @property
+    def all_rank_experiment(self):
+        if rank_zero_only.rank == 0:
+            return self.experiment
+        if hasattr(self, "_all_rank_experiment") and self._all_rank_experiment is not None:
+            return self._all_rank_experiment
+
+        assert rank_zero_only.rank != 0
+        if self.root_dir:
+            self._fs.makedirs(self.root_dir, exist_ok=True)
+
+        if _TENSORBOARD_AVAILABLE:
+            from torch.utils.tensorboard import SummaryWriter
+        else:
+            from tensorboardX import SummaryWriter  # type: ignore[no-redef]
+
+        self._all_rank_experiment = SummaryWriter(log_dir=self.log_dir, **self._kwargs)
+        return self._all_rank_experiment
+
+    def finalize(self, status: str) -> None:
+        if rank_zero_only.rank == 0:
+            super().finalize(status)
+        elif hasattr(self, "_all_rank_experiment") and self._all_rank_experiment is not None:
+            self.all_rank_experiment.flush()
+            self.all_rank_experiment.close()
+
+    def __getstate__(self):
+        state = super().__getstate__()
+        if "_all_rank_experiment" in state:
+            del state["_all_rank_experiment"]
+        return state
+
+def get_strategy(
+    devices="auto",
+    num_nodes=1,
+    accelerator="auto",
+    strategy={"name": "auto"},
+    precision=None,
+):
+    from lightning.fabric.utilities.device_parser import _determine_root_gpu_device
+    from lightning.pytorch.accelerators import AcceleratorRegistry
+    from lightning.pytorch.accelerators.cuda import CUDAAccelerator
+    from lightning.pytorch.accelerators.mps import MPSAccelerator
+    from lightning.pytorch.strategies import Strategy, SingleDeviceStrategy, StrategyRegistry
+    from lightning.pytorch.trainer.connectors import accelerator_connector
+    from lightning.pytorch.utilities.rank_zero import rank_zero_warn
+    class _DsAcceleratorConnector(accelerator_connector._AcceleratorConnector):
+        def __init__(self) -> None:
+            accelerator_connector._register_external_accelerators_and_strategies()
+            self._registered_strategies = StrategyRegistry.available_strategies()
+            self._accelerator_types = AcceleratorRegistry.available_accelerators()
+            self._parallel_devices = []
+            self._check_config_and_set_final_flags(
+                strategy=strategy["name"],
+                accelerator=accelerator,
+                precision=precision,
+                plugins=[],
+                sync_batchnorm=False,
+            )
+            if self._accelerator_flag == "auto":
+                self._accelerator_flag = self._choose_auto_accelerator()
+            elif self._accelerator_flag == "gpu":
+                self._accelerator_flag = self._choose_gpu_accelerator_backend()
+            self._check_device_config_and_set_final_flags(devices=devices, num_nodes=num_nodes)
+            self._set_parallel_devices_and_init_accelerator()
+            if self._strategy_flag == "auto":
+                self._strategy_flag = self._choose_strategy()
+            self._check_strategy_and_fallback()
+            self._init_strategy()
+            for k in ["colossalai", "bagua", "hpu", "hpu_parallel", "hpu_single", "ipu", "ipu_strategy"]:
+                if k in StrategyRegistry:
+                    StrategyRegistry.remove(k)
+
+        def _init_strategy(self) -> None:
+            assert isinstance(self._strategy_flag, (str, Strategy))
+            if isinstance(self._strategy_flag, str):
+                if self._strategy_flag not in StrategyRegistry:
+                    available_names = ", ".join(sorted(StrategyRegistry.available_strategies())) or "none"
+                    raise KeyError(f"Invalid strategy name {strategy['name']}. Available names: {available_names}")
+                data = StrategyRegistry[self._strategy_flag]
+                params = {}
+                # Replicate additional logic for _choose_strategy when dealing with single device strategies
+                if issubclass(data["strategy"], SingleDeviceStrategy):
+                    if self._accelerator_flag == "hpu":
+                        params = {"device": torch.device("hpu")}
+                    elif self._accelerator_flag == "tpu":
+                        params = {"device": self._parallel_devices[0]}
+                    elif data["strategy"] is SingleDeviceStrategy:
+                        if isinstance(self._accelerator_flag, (CUDAAccelerator, MPSAccelerator)) or (
+                            isinstance(self._accelerator_flag, str) and self._accelerator_flag in ("cuda", "gpu", "mps")
+                        ):
+                            params = {"device": _determine_root_gpu_device(self._parallel_devices)}
+                        else:
+                            params = {"device": "cpu"}
+                    else:
+                        raise NotImplementedError
+                params.update(data["init_params"])
+                params.update({k: v for k, v in strategy.items() if k != "name"})
+                self.strategy = data["strategy"](**utils.filter_kwargs(params, data["strategy"]))
+            elif isinstance(self._strategy_flag, SingleDeviceStrategy):
+                params = {"device": self._strategy_flag.root_device}
+                params.update({k: v for k, v in strategy.items() if k != "name"})
+                self.strategy = self._strategy_flag.__class__(**utils.filter_kwargs(params, self._strategy_flag.__class__))
+            else:
+                rank_zero_warn(
+                    f"Inferred strategy {self._strategy_flag.__class__.__name__} cannot take custom configurations."
+                    f"To use custom configurations, please specify the strategy name explicitly."
+                )
+                self.strategy = self._strategy_flag
+
+    return _DsAcceleratorConnector().strategy