SOME: Singing-Oriented MIDI Extractor.
import importlib
import pathlib

import click

import preprocessing
from utils.config_utils import read_full_config, print_config


@click.command(help='Process the raw dataset into binary dataset')
@click.option('--config', required=True, metavar='FILE', help='Path to the configuration file')
def binarize(config):
    config = pathlib.Path(config)
    config = read_full_config(config)
    print_config(config)
    binarizer_cls = config['binarizer_cls']
    pkg = ".".join(binarizer_cls.split(".")[:-1])
    cls_name = binarizer_cls.split(".")[-1]
    binarizer_cls = getattr(importlib.import_module(pkg), cls_name)
    assert issubclass(binarizer_cls, preprocessing.BaseBinarizer), \
        f'Binarizer class {binarizer_cls} is not a subclass of {preprocessing.BaseBinarizer}.'
    print("| Binarizer: ", binarizer_cls)
    binarizer_cls(config=config).process()


if __name__ == '__main__':
    binarize()

binarize.py
openvpi/SOME
from .base_onnx_module import BaseONNXModule
from .me_onnx_module import MIDIExtractionONNXModule
from .me_quant_onnx_module import QuantizedMIDIExtractionONNXModule

task_module_mapping = {
    'training.MIDIExtractionTask': 'deployment.MIDIExtractionONNXModule',
    'training.QuantizedMIDIExtractionTask': 'deployment.QuantizedMIDIExtractionONNXModule',
}

deployment/__init__.py
openvpi/SOME
import pathlib
from collections import OrderedDict

from librosa.filters import mel
import torch
from torch import nn

from utils import build_object_from_class_name


class BaseONNXModule(nn.Module):
    def __init__(self, config: dict, model_path: pathlib.Path, device=None):
        super().__init__()
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.config = config
        self.model_path = model_path
        self.device = device
        self.timestep = self.config['hop_size'] / self.config['audio_sample_rate']
        self.model: torch.nn.Module = self.build_model()

    def build_model(self) -> nn.Module:
        model: nn.Module = build_object_from_class_name(
            self.config['model_cls'], nn.Module, config=self.config
        ).eval().to(self.device)
        state_dict = torch.load(self.model_path, map_location=self.device)['state_dict']
        prefix_in_ckpt = 'model'
        state_dict = OrderedDict({
            k[len(prefix_in_ckpt) + 1:]: v
            for k, v in state_dict.items() if k.startswith(f'{prefix_in_ckpt}.')
        })
        model.load_state_dict(state_dict, strict=True)
        print(f'| load \'{prefix_in_ckpt}\' from \'{self.model_path}\'.')
        return model


class MelSpectrogram_ONNX(nn.Module):
    def __init__(
            self,
            n_mel_channels,
            sampling_rate,
            win_length,
            hop_length,
            n_fft=None,
            mel_fmin=0,
            mel_fmax=None,
            clamp=1e-5
    ):
        super().__init__()
        n_fft = win_length if n_fft is None else n_fft
        mel_basis = mel(
            sr=sampling_rate,
            n_fft=n_fft,
            n_mels=n_mel_channels,
            fmin=mel_fmin,
            fmax=mel_fmax,
            htk=True)
        mel_basis = torch.from_numpy(mel_basis).float()
        self.register_buffer("mel_basis", mel_basis)
        self.n_fft = win_length if n_fft is None else n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.sampling_rate = sampling_rate
        self.n_mel_channels = n_mel_channels
        self.clamp = clamp

    def forward(self, audio, center=True):
        fft = torch.stft(
            audio,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            win_length=self.win_length,
            window=torch.hann_window(self.win_length, device=audio.device),
            center=center,
            return_complex=False
        )
        magnitude = torch.sqrt(torch.sum(fft ** 2, dim=-1))
        mel_output = torch.matmul(self.mel_basis, magnitude)
        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
        return log_mel_spec

deployment/base_onnx_module.py
openvpi/SOME
import pathlib

import torch

from utils.infer_utils import decode_bounds_to_alignment, decode_gaussian_blurred_probs, decode_note_sequence
from .base_onnx_module import BaseONNXModule, MelSpectrogram_ONNX


class MIDIExtractionONNXModule(BaseONNXModule):
    def __init__(self, config: dict, model_path: pathlib.Path, device=None):
        super().__init__(config, model_path, device=device)
        self.mel_extractor = MelSpectrogram_ONNX(
            n_mel_channels=self.config['units_dim'], sampling_rate=self.config['audio_sample_rate'],
            win_length=self.config['win_size'], hop_length=self.config['hop_size'],
            mel_fmin=self.config['fmin'], mel_fmax=self.config['fmax']
        ).to(self.device)
        self.rmvpe = None
        self.midi_min = self.config['midi_min']
        self.midi_max = self.config['midi_max']
        self.midi_deviation = self.config['midi_prob_deviation']
        self.rest_threshold = self.config['rest_threshold']

    def forward(self, waveform: torch.Tensor):
        units = self.mel_extractor(waveform).transpose(1, 2)
        pitch = torch.zeros(units.shape[:2], dtype=torch.float32, device=self.device)
        masks = torch.ones_like(pitch, dtype=torch.bool)
        probs, bounds = self.model(x=units, f0=pitch, mask=masks, sig=True)
        probs *= masks[..., None]
        bounds *= masks
        unit2note_pred = decode_bounds_to_alignment(bounds, use_diff=False) * masks
        midi_pred, rest_pred = decode_gaussian_blurred_probs(
            probs, vmin=self.midi_min, vmax=self.midi_max,
            deviation=self.midi_deviation, threshold=self.rest_threshold
        )
        note_midi_pred, note_dur_pred, note_mask_pred = decode_note_sequence(
            unit2note_pred, midi_pred, ~rest_pred & masks
        )
        note_rest_pred = ~note_mask_pred
        return note_midi_pred, note_rest_pred, note_dur_pred * self.timestep

deployment/me_onnx_module.py
openvpi/SOME
import pathlib

import torch

from utils.infer_utils import decode_bounds_to_alignment, decode_note_sequence
from .base_onnx_module import BaseONNXModule, MelSpectrogram_ONNX


class QuantizedMIDIExtractionONNXModule(BaseONNXModule):
    def __init__(self, config: dict, model_path: pathlib.Path, device=None):
        super().__init__(config, model_path, device=device)
        self.mel_extractor = MelSpectrogram_ONNX(
            n_mel_channels=self.config['units_dim'], sampling_rate=self.config['audio_sample_rate'],
            win_length=self.config['win_size'], hop_length=self.config['hop_size'],
            mel_fmin=self.config['fmin'], mel_fmax=self.config['fmax']
        ).to(self.device)
        self.rmvpe = None

    def forward(self, waveform: torch.Tensor):
        units = self.mel_extractor(waveform).transpose(1, 2)
        pitch = torch.zeros(units.shape[:2], dtype=torch.float32, device=self.device)
        masks = torch.ones_like(pitch, dtype=torch.bool)
        probs, bounds = self.model(x=units, f0=pitch, mask=masks, sig=True)
        probs *= masks[..., None]
        bounds *= masks
        unit2note_pred = decode_bounds_to_alignment(bounds) * masks
        midi_pred = probs.argmax(dim=-1)
        rest_pred = midi_pred == 128
        note_midi_pred, note_dur_pred, note_mask_pred = decode_note_sequence(
            unit2note_pred, midi_pred.clip(min=0, max=127), ~rest_pred & masks
        )
        note_rest_pred = ~note_mask_pred
        return note_midi_pred, note_rest_pred, note_dur_pred * self.timestep

deployment/me_quant_onnx_module.py
openvpi/SOME
import importlib
import pathlib
from typing import Dict, Tuple, Union

import click
import onnx
import onnxsim
import torch
import yaml

import deployment
from utils.config_utils import print_config


def onnx_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

    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)


@click.command(help='Run inference with a trained model')
@click.option('--model', required=True, metavar='CKPT_PATH', help='Path to the model checkpoint (*.ckpt)')
@click.option('--out', required=False, metavar='ONNX_PATH', help='Path to the output model (*.onnx)')
def export(model, out):
    model_path = pathlib.Path(model)
    with open(model_path.with_name('config.yaml'), 'r', encoding='utf8') as f:
        config = yaml.safe_load(f)
    print_config(config)
    module_cls = deployment.task_module_mapping[config['task_cls']]

    pkg = ".".join(module_cls.split(".")[:-1])
    cls_name = module_cls.split(".")[-1]
    module_cls = getattr(importlib.import_module(pkg), cls_name)
    assert issubclass(module_cls, deployment.BaseONNXModule), \
        f'Module class {module_cls} is not a subclass of {deployment.BaseONNXModule}.'
    module_ins = module_cls(config=config, model_path=model_path)

    waveform = torch.randn((1, 114514), dtype=torch.float32, device=module_ins.device)
    out_path = pathlib.Path(out) if out is not None else model_path.with_suffix('.onnx')
    torch.onnx.export(
        module_ins,
        waveform,
        out_path,
        input_names=['waveform'],
        output_names=[
            'note_midi',
            'note_rest',
            'note_dur'
        ],
        dynamic_axes={
            'waveform': {
                1: 'n_samples'
            },
            'note_midi': {
                1: 'n_notes'
            },
            'note_rest': {
                1: 'n_notes'
            },
            'note_dur': {
                1: 'n_notes'
            },
        },
        opset_version=17
    )
    onnx_model = onnx.load(out_path.as_posix())
    onnx_override_io_shapes(onnx_model, output_shapes={
        'note_midi': (1, 'n_notes'),
        'note_rest': (1, 'n_notes'),
        'note_dur': (1, 'n_notes'),
    })
    print('Running ONNX Simplifier...')
    onnx_model, check = onnxsim.simplify(
        onnx_model,
        include_subgraph=True
    )
    assert check, 'Simplified ONNX model could not be validated'
    onnx.save(onnx_model, out_path)


if __name__ == '__main__':
    export()

export.py
openvpi/SOME
import importlib
import pathlib

import click
import librosa
import yaml

import inference
from utils.config_utils import print_config
from utils.infer_utils import build_midi_file
from utils.slicer2 import Slicer


@click.command(help='Run inference with a trained model')
@click.option('--model', required=True, metavar='CKPT_PATH', help='Path to the model checkpoint (*.ckpt)')
@click.option('--wav', required=True, metavar='WAV_PATH', help='Path to the input wav file (*.wav)')
@click.option('--midi', required=False, metavar='MIDI_PATH', help='Path to the output MIDI file (*.mid)')
@click.option('--tempo', required=False, type=float, default=120, metavar='TEMPO', help='Specify tempo in the output MIDI')
def infer(model, wav, midi, tempo):
    model_path = pathlib.Path(model)
    with open(model_path.with_name('config.yaml'), 'r', encoding='utf8') as f:
        config = yaml.safe_load(f)
    print_config(config)
    infer_cls = inference.task_inference_mapping[config['task_cls']]

    pkg = ".".join(infer_cls.split(".")[:-1])
    cls_name = infer_cls.split(".")[-1]
    infer_cls = getattr(importlib.import_module(pkg), cls_name)
    assert issubclass(infer_cls, inference.BaseInference), \
        f'Inference class {infer_cls} is not a subclass of {inference.BaseInference}.'
    infer_ins = infer_cls(config=config, model_path=model_path)

    wav_path = pathlib.Path(wav)
    waveform, _ = librosa.load(wav_path, sr=config['audio_sample_rate'], mono=True)
    slicer = Slicer(sr=config['audio_sample_rate'], max_sil_kept=1000)
    chunks = slicer.slice(waveform)
    midis = infer_ins.infer([c['waveform'] for c in chunks])

    midi_file = build_midi_file([c['offset'] for c in chunks], midis, tempo=tempo)

    midi_path = pathlib.Path(midi) if midi is not None else wav_path.with_suffix('.mid')
    midi_file.save(midi_path)
    print(f'MIDI file saved at: \'{midi_path}\'')


if __name__ == '__main__':
    infer()

infer.py
openvpi/SOME
from .base_infer import BaseInference
from .me_infer import MIDIExtractionInference
from .me_quant_infer import QuantizedMIDIExtractionInference

task_inference_mapping = {
    'training.MIDIExtractionTask': 'inference.MIDIExtractionInference',
    'training.QuantizedMIDIExtractionTask': 'inference.QuantizedMIDIExtractionInference',
}

inference/__init__.py
openvpi/SOME
import pathlib
from collections import OrderedDict
from typing import Dict, List

import numpy as np
import torch
import tqdm
from torch import nn

from utils import build_object_from_class_name


class BaseInference:
    def __init__(self, config: dict, model_path: pathlib.Path, device=None):
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.config = config
        self.model_path = model_path
        self.device = device
        self.timestep = self.config['hop_size'] / self.config['audio_sample_rate']
        self.model: torch.nn.Module = self.build_model()

    def build_model(self) -> nn.Module:
        model: nn.Module = build_object_from_class_name(
            self.config['model_cls'], nn.Module, config=self.config
        ).eval().to(self.device)
        state_dict = torch.load(self.model_path, map_location=self.device)['state_dict']
        prefix_in_ckpt = 'model'
        state_dict = OrderedDict({
            k[len(prefix_in_ckpt) + 1:]: v
            for k, v in state_dict.items() if k.startswith(f'{prefix_in_ckpt}.')
        })
        model.load_state_dict(state_dict, strict=True)
        print(f'| load \'{prefix_in_ckpt}\' from \'{self.model_path}\'.')
        return model

    def preprocess(self, waveform: np.ndarray) -> Dict[str, torch.Tensor]:
        raise NotImplementedError()

    def forward_model(self, sample: Dict[str, torch.Tensor]):
        raise NotImplementedError()

    def postprocess(self, results: Dict[str, torch.Tensor]) -> List[Dict[str, np.ndarray]]:
        raise NotImplementedError()

    def infer(self, waveforms: List[np.ndarray]) -> List[Dict[str, np.ndarray]]:
        results = []
        for w in tqdm.tqdm(waveforms):
            model_in = self.preprocess(w)
            model_out = self.forward_model(model_in)
            res = self.postprocess(model_out)
            results.append(res)
        return results

inference/base_infer.py
openvpi/SOME
import pathlib
from typing import Dict, List

import librosa
import numpy as np
import torch

import modules.rmvpe
from utils.binarizer_utils import get_pitch_parselmouth
from utils.infer_utils import decode_bounds_to_alignment, decode_gaussian_blurred_probs, decode_note_sequence
from utils.pitch_utils import resample_align_curve
from .base_infer import BaseInference


class MIDIExtractionInference(BaseInference):
    def __init__(self, config: dict, model_path: pathlib.Path, device=None):
        super().__init__(config, model_path, device=device)
        self.mel_spec = modules.rmvpe.MelSpectrogram(
            n_mel_channels=self.config['units_dim'], sampling_rate=self.config['audio_sample_rate'],
            win_length=self.config['win_size'], hop_length=self.config['hop_size'],
            mel_fmin=self.config['fmin'], mel_fmax=self.config['fmax']
        ).to(self.device)
        self.rmvpe = None
        self.midi_min = self.config['midi_min']
        self.midi_max = self.config['midi_max']
        self.midi_deviation = self.config['midi_prob_deviation']
        self.rest_threshold = self.config['rest_threshold']

    def preprocess(self, waveform: np.ndarray) -> Dict[str, torch.Tensor]:
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0).to(self.device)
        units = self.mel_spec(wav_tensor).transpose(1, 2)
        length = units.shape[1]
        # f0_algo = self.config['pe']
        # if f0_algo == 'parselmouth':
        #     f0, _ = get_pitch_parselmouth(
        #         waveform, sample_rate=self.config['audio_sample_rate'],
        #         hop_size=self.config['hop_size'], length=length, interp_uv=True
        #     )
        # elif f0_algo == 'rmvpe':
        #     if self.rmvpe is None:
        #         self.rmvpe = modules.rmvpe.RMVPE(self.config['pe_ckpt'], device=self.device)
        #     f0, _ = self.rmvpe.get_pitch(
        #         waveform, sample_rate=self.config['audio_sample_rate'],
        #         hop_size=self.rmvpe.mel_extractor.hop_length,
        #         length=(waveform.shape[0] + self.rmvpe.mel_extractor.hop_length - 1) // self.rmvpe.mel_extractor.hop_length,
        #         interp_uv=True
        #     )
        #     f0 = resample_align_curve(
        #         f0,
        #         original_timestep=self.rmvpe.mel_extractor.hop_length / self.config['audio_sample_rate'],
        #         target_timestep=self.config['hop_size'] / self.config['audio_sample_rate'],
        #         align_length=length
        #     )
        # else:
        #     raise NotImplementedError(f'Invalid pitch extractor: {f0_algo}')
        # pitch = librosa.hz_to_midi(f0)
        # pitch = torch.from_numpy(pitch).unsqueeze(0).to(self.device)
        pitch = torch.zeros(units.shape[:2], dtype=torch.float32, device=self.device)
        return {
            'units': units,
            'pitch': pitch,
            'masks': torch.ones_like(pitch, dtype=torch.bool)
        }

    @torch.no_grad()
    def forward_model(self, sample: Dict[str, torch.Tensor]):



        probs, bounds = self.model(x=sample['units'], f0=sample['pitch'], mask=sample['masks'],sig=True)

        return {
            'probs': probs,
            'bounds': bounds,
            'masks': sample['masks'],
        }

    def postprocess(self, results: Dict[str, torch.Tensor]) -> List[Dict[str, np.ndarray]]:
        probs = results['probs']
        bounds = results['bounds']
        masks = results['masks']
        probs *= masks[..., None]
        bounds *= masks
        unit2note_pred = decode_bounds_to_alignment(bounds) * masks
        midi_pred, rest_pred = decode_gaussian_blurred_probs(
            probs, vmin=self.midi_min, vmax=self.midi_max,
            deviation=self.midi_deviation, threshold=self.rest_threshold
        )
        note_midi_pred, note_dur_pred, note_mask_pred = decode_note_sequence(
            unit2note_pred, midi_pred, ~rest_pred & masks
        )
        note_rest_pred = ~note_mask_pred
        return {
            'note_midi': note_midi_pred.squeeze(0).cpu().numpy(),
            'note_dur': note_dur_pred.squeeze(0).cpu().numpy() * self.timestep,
            'note_rest': note_rest_pred.squeeze(0).cpu().numpy()
        }

inference/me_infer.py
openvpi/SOME
from typing import Dict, List

import numpy as np
import torch

from utils.infer_utils import decode_bounds_to_alignment, decode_note_sequence
from .me_infer import MIDIExtractionInference


class QuantizedMIDIExtractionInference(MIDIExtractionInference):
    @torch.no_grad()
    def forward_model(self, sample: Dict[str, torch.Tensor]):
        probs, bounds = self.model(x=sample['units'], f0=sample['pitch'], mask=sample['masks'], softmax=True)

        return {
            'probs': probs,
            'bounds': bounds,
            'masks': sample['masks'],
        }

    def postprocess(self, results: Dict[str, torch.Tensor]) -> List[Dict[str, np.ndarray]]:
        probs = results['probs']
        bounds = results['bounds']
        masks = results['masks']
        probs *= masks[..., None]
        bounds *= masks
        unit2note_pred = decode_bounds_to_alignment(bounds) * masks
        midi_pred = probs.argmax(dim=-1)
        rest_pred = midi_pred == 128
        note_midi_pred, note_dur_pred, note_mask_pred = decode_note_sequence(
            unit2note_pred, midi_pred.clip(min=0, max=127), ~rest_pred & masks
        )
        note_rest_pred = ~note_mask_pred
        return {
            'note_midi': note_midi_pred.squeeze(0).cpu().numpy(),
            'note_dur': note_dur_pred.squeeze(0).cpu().numpy() * self.timestep,
            'note_rest': note_rest_pred.squeeze(0).cpu().numpy()
        }

inference/me_quant_infer.py
openvpi/SOME
from typing import Union

import numpy as np
import torch
from torch.optim.lr_scheduler import _LRScheduler

# from typeguard import check_argument_types


class WarmupLR(_LRScheduler):
    """The WarmupLR scheduler

    This scheduler is almost same as NoamLR Scheduler except for following
    difference:

    NoamLR:
        lr = optimizer.lr * model_size ** -0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)
    WarmupLR:
        lr = optimizer.lr * warmup_step ** 0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)

    Note that the maximum lr equals to optimizer.lr in this scheduler.

    """

    def __init__(
            self,
            optimizer: torch.optim.Optimizer,
            warmup_steps: Union[int, float] = 5000,
            min_lr=2e-5,
            last_epoch: int = -1,
    ):
        # assert check_argument_types()
        self.warmup_steps = warmup_steps
        self.min_lr = min_lr
        super().__init__(optimizer, last_epoch)

    def __repr__(self):
        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps}, lr={self.base_lr}, min_lr={self.min_lr}, last_epoch={self.last_epoch})"

    def get_lr(self):
        step_num = self.last_epoch + 1
        if self.warmup_steps == 0:
            lrs = []
            for lr in self.base_lrs:
                lr = lr * step_num ** -0.5
                if lr < self.min_lr:
                    lr = self.min_lr
                lrs.append(lr)
            return lrs
        else:
            lrs = []
            for lr in self.base_lrs:
                lr = lr * self.warmup_steps ** 0.5 * min(step_num ** -0.5, step_num * self.warmup_steps ** -1.5)
                if lr < self.min_lr and step_num > self.warmup_steps:
                    lr = self.min_lr
                lrs.append(lr)
            return lrs

    def set_step(self, step: int):
        self.last_epoch = step

class SGDRLR(_LRScheduler):
    """The WarmupLR scheduler

    This scheduler is almost same as NoamLR Scheduler except for following
    difference:

    NoamLR:
        lr = optimizer.lr * model_size ** -0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)
    WarmupLR:
        lr = optimizer.lr * warmup_step ** 0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)

    Note that the maximum lr equals to optimizer.lr in this scheduler.

    """

    def __init__(
            self,
            optimizer: torch.optim.Optimizer,
            warmup_steps: Union[int, float] = 25000,
            min_lr=1e-5,
            last_epoch: int = -1, T_0=1500, eta_max=0.1, eta_min=0.,T_mul=2,T_mult=2
    ):
        # assert check_argument_types()
        self.warmup_steps = warmup_steps
        self.min_lr = min_lr
        self.eta_min = eta_min
        self.T_0 = T_0
        self.eta_max = eta_max
        self.T_mul = T_mul
        self.T_mult = T_mult

        super().__init__(optimizer, last_epoch)

    def __repr__(self):
        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps}, lr={self.base_lr}, min_lr={self.min_lr}, last_epoch={self.last_epoch})"

    def adjust_lr(self,):
        step_num = self.last_epoch + 1
        if self.T_mul == 2:
            i = np.log2(step_num / self.T_0 + 1).astype(np.int32)
            T_cur = step_num - self.T_0 * (self.T_mult ** (i) - 1)
            T_i = (self.T_0 * self.T_mult ** i)
        elif self.T_mul == 1:
            T_cur = step_num % self.T_0
            T_i = self.T_0
        cur_lr = self.eta_min + 0.5 * (self.eta_max - self.eta_min) * (1 + np.cos(np.pi * T_cur / T_i))
        return cur_lr


    def get_lr(self):
        # step_num = self.last_epoch + 1
        if self.warmup_steps == 0:
            lrs = []
            for lr in self.base_lrs:
                lr = self.adjust_lr()
                lrs.append(lr)
            return lrs
        else:
            lrs = []
            for lr in self.base_lrs:
                lr = self.adjust_lr()
                lrs.append(lr)
            return lrs

    def set_step(self, step: int):
        self.last_epoch = step
class LSGDRLR(_LRScheduler):
    """The WarmupLR scheduler

    This scheduler is almost same as NoamLR Scheduler except for following
    difference:

    NoamLR:
        lr = optimizer.lr * model_size ** -0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)
    WarmupLR:
        lr = optimizer.lr * warmup_step ** 0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)

    Note that the maximum lr equals to optimizer.lr in this scheduler.

    """

    def __init__(
            self,
            optimizer: torch.optim.Optimizer,
            warmup_steps: Union[int, float] = 25000,
            min_lr=1e-5,
            last_epoch: int = -1, T_0=1500, eta_max=0.1, eta_min=0.,T_mul=2,T_mult=0.9999
    ):
        # assert check_argument_types()
        self.warmup_steps = warmup_steps
        self.min_lr = min_lr
        self.eta_min = eta_min
        self.T_0 = T_0
        self.eta_max = eta_max
        self.T_mul = T_mul
        self.T_mult = T_mult

        super().__init__(optimizer, last_epoch)

    def __repr__(self):
        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps}, lr={self.base_lr}, min_lr={self.min_lr}, last_epoch={self.last_epoch})"

    def adjust_lr(self,):
        step_num = self.last_epoch + 1

        cur_lr = self.eta_min* self.T_mult ** step_num +  np.cos(np.pi * step_num / self.T_0 )
        return cur_lr


    def get_lr(self):
        # step_num = self.last_epoch + 1
        if self.warmup_steps == 0:
            lrs = []
            for lr in self.base_lrs:
                lr = self.adjust_lr()
                lrs.append(lr)
            return lrs
        else:
            lrs = []
            for lr in self.base_lrs:
                lr = self.adjust_lr()
                lrs.append(lr)
            return lrs

    def set_step(self, step: int):
        self.last_epoch = step

class V2LSGDRLR(_LRScheduler):
    """The WarmupLR scheduler

    This scheduler is almost same as NoamLR Scheduler except for following
    difference:

    NoamLR:
        lr = optimizer.lr * model_size ** -0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)
    WarmupLR:
        lr = optimizer.lr * warmup_step ** 0.5
             * min(step ** -0.5, step * warmup_step ** -1.5)

    Note that the maximum lr equals to optimizer.lr in this scheduler.

    """

    def __init__(
            self,
            optimizer: torch.optim.Optimizer,
            warmup_steps: Union[int, float] = 25000,
            min_lr=1e-5,
            last_epoch: int = -1, T_0=1500, eta_max=0.1, eta_min=0.,T_mul=2,T_mult=0.9999
    ):
        # assert check_argument_types()
        self.warmup_steps = warmup_steps
        self.min_lr = min_lr
        self.eta_min = eta_min
        self.T_0 = T_0
        self.eta_max = eta_max
        self.T_mul = T_mul
        self.T_mult = T_mult

        super().__init__(optimizer, last_epoch)

    def __repr__(self):
        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps}, lr={self.base_lr}, min_lr={self.min_lr}, last_epoch={self.last_epoch})"

    def ctxadjust_lr(self,T_mul = 1,T_0=15000,T_mult=1.5,eta_min=0.0000001,eta_max=0.00006,tmctx=0.99,ws=8000):
        step_num = self.last_epoch+1
        if T_mul == 2:
            i = np.log2(step_num / T_0 + 1).astype(np.int32)
            T_cur = step_num - T_0 * (T_mult ** (i) - 1)
            T_i = (T_0 * T_mult ** i)
        elif T_mul == 1:
            T_cur = (step_num + ws) % T_0
            T_i = T_0
            T_curX = (step_num + ws) // T_0


        cur_lr = eta_min + 0.5 * (eta_max *(tmctx**T_curX)- eta_min*(tmctx**T_curX)) * (1 + np.cos(np.pi * T_cur / T_i))
        if ws>step_num:
            cur_lr=step_num*(eta_max/ws)

        return cur_lr
class V3LSGDRLR(_LRScheduler):
    """The WarmupLR schedulerA
        This scheduler is almost same as NoamLR Scheduler except for following
        difference:
        NoamLR:
            lr = optimizer.lr * model_size ** -0.5
                 * min(step ** -0.5, step * warmup_step ** -1.5)
        WarmupLR:
            lr = optimizer.lr * warmup_step ** 0.5
                 * min(step ** -0.5, step * warmup_step ** -1.5)
        Note that the maximum lr equals to optimizer.lr in this scheduler.
        """
    def __init__(self,optimizer: torch.optim.Optimizer,warmup_steps: Union[int, float] = 25000,min_lr=1e-5,last_epoch: int = -1, T_0=1500, eta_max=0.1, eta_min=0., T_mul=2, T_mult=0.9999):
            # assert check_argument_types()
        self.warmup_steps = warmup_steps
        self.min_lr = min_lr
        self.eta_min = eta_min
        self.T_0 = T_0
        self.eta_max = eta_max
        self.T_mul = T_mul
        self.T_mult = T_mult
        super().__init__(optimizer, last_epoch)
    def __repr__(self):
        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps}, lr={self.base_lr}, min_lr={self.min_lr}, last_epoch={self.last_epoch})"
    def ctxadjust_lr(self, T_0=15000, eta_min=0.00006, eta_max=0.00009, tmctx=0.98, ws=5000):
        step_num = self.last_epoch + 1 #+360000
        T_cur = (step_num + ws) % T_0
        T_i = T_0
        T_curX = (step_num + ws) // T_0
        cur_lr = eta_min * (tmctx ** T_curX) + 0.5 * (eta_max * (tmctx ** T_curX) - eta_min * (tmctx ** T_curX)) * (
                    1 + np.cos(np.pi * T_cur / T_i))
        if ws > step_num:
            cur_lr = step_num * (eta_max / ws)
        return cur_lr



    def get_lr(self):
        # step_num = self.last_epoch + 1
        if self.warmup_steps == 0:
            lrs = []
            for lr in self.base_lrs:
                lr = self.ctxadjust_lr()
                lrs.append(lr)
            return lrs
        else:
            lrs = []
            for lr in self.base_lrs:
                lr = self.ctxadjust_lr()
                lrs.append(lr)
            return lrs

    def set_step(self, step: int):
        self.last_epoch = step



class NoamHoldAnnealing(_LRScheduler):
    def __init__(self, optimizer, max_steps=175680, warmup_steps=None, warmup_ratio=0.2, hold_steps=None,
                 hold_ratio=0.3, decay_rate=1.0, min_lr=1.e-5, last_epoch=-1):
        """
        From Nemo:
        Implementation of the Noam Hold Annealing policy from the SqueezeFormer paper.

        Unlike NoamAnnealing, the peak learning rate can be explicitly set for this scheduler.
        The schedule first performs linear warmup,
        then holds the peak LR, then decays with some schedule for
        the remainder of the steps.
        Therefore the min-lr is still dependent on the hyper parameters selected.

        It's schedule is determined by three factors-

        Warmup Steps: Initial stage, where linear warmup
            occurs uptil the peak LR is reached. Unlike NoamAnnealing,
            the peak LR is explicitly stated here instead of a scaling factor.

        Hold Steps: Intermediate stage, where the peak LR
            is maintained for some number of steps. In this region,
            the high peak LR allows the model to converge faster
            if training is stable. However the high LR
            may also cause instability during training.
            Should usually be a significant fraction of training
            steps (around 30-40% of the entire training steps).

        Decay Steps: Final stage, where the LR rapidly decays
            with some scaling rate (set by decay rate).
            To attain Noam decay, use 0.5,
            for Squeezeformer recommended decay, use 1.0.
            The fast decay after prolonged high LR during
            hold phase allows for rapid convergence.

        References:
            - [Squeezeformer:
            An Efficient Transformer for Automatic Speech Recognition]
            (https://arxiv.org/abs/2206.00888)

        Args:
            optimizer: Pytorch compatible Optimizer object.
            warmup_steps: Number of training steps in warmup stage
            warmup_ratio: Ratio of warmup steps to total steps
            hold_steps: Number of training steps to hold the learning rate after warm up
            hold_ratio: Ratio of hold steps to total steps
            max_steps: Total number of steps while training or `None` for infinite training
            decay_rate: Float value describing the polynomial decay
                        after the hold period. Default value of 0.5 corresponds to Noam decay.
            min_lr: Minimum learning rate.
        """
        self.decay_rate = decay_rate
        self.min_lr = min_lr
        self._last_warmup_lr = 0.0

        # Necessary to duplicate as class attributes are hidden in inner class
        self.max_steps = max_steps
        if warmup_steps is not None:
            self.warmup_steps = warmup_steps
        elif warmup_ratio is not None:
            self.warmup_steps = int(warmup_ratio * max_steps)
        else:
            self.warmup_steps = 0

        if hold_steps is not None:
            self.hold_steps = hold_steps + self.warmup_steps
        elif hold_ratio is not None:
            self.hold_steps = int(hold_ratio * max_steps) + self.warmup_steps
        else:
            self.hold_steps = 0

        super().__init__(optimizer, last_epoch)

    def _get_warmup_lr(self, step):
        lr_val = (step + 1) / (self.warmup_steps + 1)
        return [initial_lr * lr_val for initial_lr in self.base_lrs]

    def get_lr(self):
        step = self.last_epoch

        # Warmup phase
        if step <= self.warmup_steps and self.warmup_steps > 0:
            return self._get_warmup_lr(step)

        # Hold phase
        if (step >= self.warmup_steps) and (step < self.hold_steps):
            return self.base_lrs

        if step > self.max_steps:
            return [self.min_lr for _ in self.base_lrs]

        return self._get_lr(step)

    @staticmethod
    def _noam_hold_annealing(initial_lr, step, warmup_steps, hold_steps, decay_rate, min_lr):
        # hold_steps = total number of steps
        # to hold the LR, not the warmup + hold steps.
        T_warmup_decay = max(1, warmup_steps ** decay_rate)
        T_hold_decay = max(1, (step - hold_steps) ** decay_rate)
        lr = (initial_lr * T_warmup_decay) / T_hold_decay
        lr = max(lr, min_lr)
        return lr

    def _get_lr(self, step):
        if self.warmup_steps is None or self.warmup_steps == 0:
            raise ValueError("Noam scheduler cannot be used without warmup steps")

        if self.hold_steps > 0:
            hold_steps = self.hold_steps - self.warmup_steps
        else:
            hold_steps = 0

        new_lrs = [
            self._noam_hold_annealing(initial_lr=initial_lr,
                                      step=step,
                                      warmup_steps=self.warmup_steps,
                                      hold_steps=hold_steps,
                                      decay_rate=self.decay_rate,
                                      min_lr=self.min_lr)
            for initial_lr in self.base_lrs
        ]
        return new_lrs

    def set_step(self, step: int):
        self.last_epoch = step

lr_scheduler/scheduler.py
openvpi/SOME

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange


class Attention(nn.Module):
    def __init__(self, dim, heads=4, dim_head=32, conditiondim=None):
        super().__init__()
        if conditiondim is None:
            conditiondim = dim

        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_q = nn.Linear(dim, hidden_dim, bias=False)
        self.to_kv = nn.Linear(conditiondim, hidden_dim * 2, bias=False)

        self.to_out = nn.Sequential(nn.Linear(hidden_dim, dim, ),
                                    )

    def forward(self, q, kv=None, mask=None):
        # b, c, h, w = x.shape
        if kv is None:
            kv = q
        # q, kv = map(
        #     lambda t: rearrange(t, "b c t -> b t c", ), (q, kv)
        # )

        q = self.to_q(q)
        k, v = self.to_kv(kv).chunk(2, dim=2)

        q, k, v = map(
            lambda t: rearrange(t, "b t (h c) -> b h t c", h=self.heads), (q, k, v)
        )

        if mask is not None:
            mask = mask.unsqueeze(1).unsqueeze(1)

        with torch.backends.cuda.sdp_kernel(enable_math=False
                                            ):
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

        out = rearrange(out, "b h t c -> b t (h c) ", h=self.heads, )
        return self.to_out(out)

modules/attention/base_attention.py
openvpi/SOME
from .tts_modules import LengthRegulator

modules/commons/__init__.py
openvpi/SOME
import torch
import torch.nn.functional as F


class LengthRegulator(torch.nn.Module):
    # noinspection PyMethodMayBeStatic
    def forward(self, dur, dur_padding=None, alpha=None):
        """
        Example (no batch dim version):
            1. dur = [2,2,3]
            2. token_idx = [[1],[2],[3]], dur_cumsum = [2,4,7], dur_cumsum_prev = [0,2,4]
            3. token_mask = [[1,1,0,0,0,0,0],
                             [0,0,1,1,0,0,0],
                             [0,0,0,0,1,1,1]]
            4. token_idx * token_mask = [[1,1,0,0,0,0,0],
                                         [0,0,2,2,0,0,0],
                                         [0,0,0,0,3,3,3]]
            5. (token_idx * token_mask).sum(0) = [1,1,2,2,3,3,3]

        :param dur: Batch of durations of each frame (B, T_txt)
        :param dur_padding: Batch of padding of each frame (B, T_txt)
        :param alpha: duration rescale coefficient
        :return:
            mel2ph (B, T_speech)
        """
        assert alpha is None or alpha > 0
        if alpha is not None:
            dur = torch.round(dur.float() * alpha).long()
        if dur_padding is not None:
            dur = dur * (1 - dur_padding.long())
        token_idx = torch.arange(1, dur.shape[1] + 1)[None, :, None].to(dur.device)
        dur_cumsum = torch.cumsum(dur, 1)
        dur_cumsum_prev = F.pad(dur_cumsum, [1, -1], mode='constant', value=0)

        pos_idx = torch.arange(dur.sum(-1).max())[None, None].to(dur.device)
        token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
        mel2ph = (token_idx * token_mask.long()).sum(1)
        return mel2ph

modules/commons/tts_modules.py
openvpi/SOME
from math import sqrt

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

from modules.attention.base_attention import Attention
from modules.conv.base_conv import conform_conv
class GLU(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        out, gate = x.chunk(2, dim=self.dim)

        return out * gate.sigmoid()

class conform_ffn(nn.Module):
    def __init__(self, dim, DropoutL1: float = 0.1, DropoutL2: float = 0.1):
        super().__init__()
        self.ln1 = nn.Linear(dim, dim * 4)
        self.ln2 = nn.Linear(dim * 4, dim)
        self.drop1 = nn.Dropout(DropoutL1) if DropoutL1 > 0. else nn.Identity()
        self.drop2 = nn.Dropout(DropoutL2) if DropoutL2 > 0. else nn.Identity()
        self.act = nn.SiLU()

    def forward(self, x):
        x = self.ln1(x)
        x = self.act(x)
        x = self.drop1(x)
        x = self.ln2(x)
        return self.drop2(x)


class conform_blocke(nn.Module):
    def __init__(self, dim: int, kernel_size: int = 31, conv_drop: float = 0.1, ffn_latent_drop: float = 0.1,
                 ffn_out_drop: float = 0.1, attention_drop: float = 0.1, attention_heads: int = 4,
                 attention_heads_dim: int = 64):
        super().__init__()
        self.ffn1 = conform_ffn(dim, ffn_latent_drop, ffn_out_drop)
        self.ffn2 = conform_ffn(dim, ffn_latent_drop, ffn_out_drop)
        self.att = Attention(dim, heads=attention_heads, dim_head=attention_heads_dim)
        self.attdrop = nn.Dropout(attention_drop) if attention_drop > 0. else nn.Identity()
        self.conv = conform_conv(dim, kernel_size=kernel_size,

                                 DropoutL=conv_drop, )
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)
        self.norm4 = nn.LayerNorm(dim)
        self.norm5 = nn.LayerNorm(dim)


    def forward(self, x, mask=None,):
        x = self.ffn1(self.norm1(x)) * 0.5 + x


        x = self.attdrop(self.att(self.norm2(x), mask=mask)) + x
        x = self.conv(self.norm3(x)) + x
        x = self.ffn2(self.norm4(x)) * 0.5 + x
        return self.norm5(x)

        # return x


class Gcf(nn.Module):
    def __init__(self,dim: int, kernel_size: int = 31, conv_drop: float = 0.1, ffn_latent_drop: float = 0.1,
                 ffn_out_drop: float = 0.1, attention_drop: float = 0.1, attention_heads: int = 4,
                 attention_heads_dim: int = 64):
        super().__init__()
        self.att1=conform_blocke(dim=dim, kernel_size=kernel_size, conv_drop=conv_drop, ffn_latent_drop=ffn_latent_drop,
                            ffn_out_drop=ffn_out_drop, attention_drop=attention_drop, attention_heads=attention_heads,
                            attention_heads_dim=attention_heads_dim)
        self.att2 = conform_blocke(dim=dim, kernel_size=kernel_size, conv_drop=conv_drop, ffn_latent_drop=ffn_latent_drop,
                            ffn_out_drop=ffn_out_drop, attention_drop=attention_drop, attention_heads=attention_heads,
                            attention_heads_dim=attention_heads_dim)
        self.glu1=nn.Sequential(nn.Linear(dim, dim*2),GLU(2) )
        self.glu2 = nn.Sequential(nn.Linear(dim, dim * 2), GLU(2))

    def forward(self, midi,bound):
        midi=self.att1(midi)
        bound=self.att2(bound)
        midis=self.glu1(midi)
        bounds=self.glu2(bound)
        return midi+bounds,bound+midis




class Gmidi_conform(nn.Module):
    def __init__(self, lay: int, dim: int, indim: int, outdim: int, use_lay_skip: bool, kernel_size: int = 31,
                 conv_drop: float = 0.1,
                 ffn_latent_drop: float = 0.1,
                 ffn_out_drop: float = 0.1, attention_drop: float = 0.1, attention_heads: int = 4,
                 attention_heads_dim: int = 64):
        super().__init__()

        self.inln = nn.Linear(indim, dim)
        self.inln1 = nn.Linear(indim, dim)
        self.outln = nn.Linear(dim, outdim)
        self.cutheard = nn.Linear(dim, 1)
        # self.cutheard = nn.Linear(dim, outdim)
        self.lay = lay
        self.use_lay_skip = use_lay_skip
        self.cf_lay = nn.ModuleList(
            [Gcf(dim=dim, kernel_size=kernel_size, conv_drop=conv_drop, ffn_latent_drop=ffn_latent_drop,
                            ffn_out_drop=ffn_out_drop, attention_drop=attention_drop, attention_heads=attention_heads,
                            attention_heads_dim=attention_heads_dim) for _ in range(lay)])
        self.att1=conform_blocke(dim=dim, kernel_size=kernel_size, conv_drop=conv_drop, ffn_latent_drop=ffn_latent_drop,
                            ffn_out_drop=ffn_out_drop, attention_drop=attention_drop, attention_heads=attention_heads,
                            attention_heads_dim=attention_heads_dim)
        self.att2 = conform_blocke(dim=dim, kernel_size=kernel_size, conv_drop=conv_drop, ffn_latent_drop=ffn_latent_drop,
                            ffn_out_drop=ffn_out_drop, attention_drop=attention_drop, attention_heads=attention_heads,
                            attention_heads_dim=attention_heads_dim)


    def forward(self, x, pitch, mask=None):

        # torch.masked_fill()
        x1=x.clone()

        x = self.inln(x )
        x1=self.inln1(x1)
        if mask is not None:
            x = x.masked_fill(~mask.unsqueeze(-1), 0)
        for idx, i in enumerate(self.cf_lay):
            x,x1 = i(x,x1)

            if mask is not None:
                x = x.masked_fill(~mask.unsqueeze(-1), 0)
        x,x1=self.att1(x),self.att2(x1)

        cutprp = self.cutheard(x1)
        midiout = self.outln(x)
        cutprp = torch.sigmoid(cutprp)
        cutprp = torch.squeeze(cutprp, -1)

        return midiout, cutprp

modules/conform/Gconform.py
openvpi/SOME
import torch
from fairseq import checkpoint_utils


class ContentVec768L12(torch.nn.Module):
    def __init__(self, path, h_sample_rate=16000, h_hop_size=320, device='cpu'):
        super().__init__()
        self.device = device
        models, self.saved_cfg, self.task = checkpoint_utils.load_model_ensemble_and_task([path], suffix="")
        self.hubert = models[0].to(self.device).eval()

    def forward(self, waveform):  # B, T
        feats = waveform.view(1, -1)
        padding_mask = torch.BoolTensor(feats.shape).fill_(False)
        inputs = {
            "source": feats.to(waveform.device),
            "padding_mask": padding_mask.to(waveform.device),
            "output_layer": 9,  # layer 9
        }
        with torch.no_grad():
            logits = self.hubert.extract_features(**inputs)
            feats = logits[0]
        units = feats  # .transpose(2, 1)
        return units

modules/contentvec/__init__.py
openvpi/SOME
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange


class GLU(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        out, gate = x.chunk(2, dim=self.dim)

        return out * gate.sigmoid()


class conform_conv(nn.Module):
    def __init__(self, channels: int,
                 kernel_size: int = 31,

                 DropoutL=0.1,

                 bias: bool = True):
        super().__init__()
        self.act2 = nn.SiLU()
        self.act1 = GLU(1)

        self.pointwise_conv1 = nn.Conv1d(
            channels,
            2 * channels,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=bias)

        # self.lorder is used to distinguish if it's a causal convolution,
        # if self.lorder > 0:
        #    it's a causal convolution, the input will be padded with
        #    `self.lorder` frames on the left in forward (causal conv impl).
        # else: it's a symmetrical convolution

        assert (kernel_size - 1) % 2 == 0
        padding = (kernel_size - 1) // 2

        self.depthwise_conv = nn.Conv1d(channels, channels,  kernel_size,
                                        stride=1,
                                        padding=padding,
                                        groups=channels,
                                        bias=bias)


        self.norm = nn.BatchNorm1d(channels)


        self.pointwise_conv2 = nn.Conv1d(channels,
                                         channels,
                                         kernel_size=1,
                                         stride=1,
                                         padding=0,
                                         bias=bias)
        self.drop=nn.Dropout(DropoutL) if DropoutL>0. else nn.Identity()
    def forward(self,x):
        x=x.transpose(1,2)
        x=self.act1(self.pointwise_conv1(x))
        x=self.depthwise_conv (x)
        x=self.norm(x)
        x=self.act2(x)
        x=self.pointwise_conv2(x)
        return self.drop(x).transpose(1,2)


modules/conv/base_conv.py
openvpi/SOME
from .bound_loss import BinaryEMDLoss, BoundaryLoss

modules/losses/__init__.py
openvpi/SOME
import math

import torch.nn


class BinaryEMDLoss(torch.nn.Module):
    def __init__(self, bidirectional=False):
        super().__init__()
        self.loss = torch.nn.L1Loss()
        self.bidirectional = bidirectional

    def forward(self, pred, gt):
        # pred, gt: [B, T]
        scale = math.sqrt(gt.shape[1])
        loss = self.loss(pred.cumsum(dim=1) / scale, gt.cumsum(dim=1) / scale)
        if self.bidirectional:
            loss += self.loss(pred.flip(1).cumsum(dim=1) / scale, gt.flip(1).cumsum(dim=1) / scale)
            loss /= 2
        return loss


class BoundaryLoss(torch.nn.Module):
    def __init__(self, lambda_bce=0.1):
        super().__init__()
        self.emd = BinaryEMDLoss(bidirectional=False)
        self.bce = torch.nn.BCELoss()
        self.lambda_bce = lambda_bce

    def forward(self, pred, gt):
        # pred, gt: [B, T]
        emd_loss = self.emd(pred, gt)
        bce_loss = self.bce(pred, gt)
        return emd_loss + self.lambda_bce * bce_loss

modules/losses/bound_loss.py
openvpi/SOME
from .midi_acc import MIDIAccuracy

modules/metrics/__init__.py
openvpi/SOME
import torch
import torchmetrics
from torch import Tensor


class MIDIAccuracy(torchmetrics.Metric):
    def __init__(self, *, tolerance, **kwargs):
        super().__init__(**kwargs)
        self.tolerance = tolerance
        self.add_state('correct', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
        self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')

    def update(self, midi_pred: Tensor, rest_pred: Tensor, midi_gt: Tensor, rest_gt: Tensor, mask=None) -> None:
        """

        :param midi_pred: predicted MIDI
        :param rest_pred: predict rest flags
        :param midi_gt: reference MIDI
        :param rest_gt: reference rest flags
        :param mask: valid or non-padding mask
        """
        assert midi_gt.shape == rest_gt.shape == midi_pred.shape == rest_pred.shape, \
            (f'shapes of pred and gt mismatch: '
             f'{midi_pred.shape}, {rest_pred.shape}, {midi_gt.shape}, {rest_gt.shape}')
        if mask is not None:
            assert midi_gt.shape == mask.shape, \
                f'shapes of pred, target and mask mismatch: {midi_pred.shape}, {rest_pred.shape}, {mask.shape}'
        midi_close = ~rest_pred & ~rest_gt & (torch.abs(midi_pred - midi_gt) <= self.tolerance)
        rest_correct = rest_pred == rest_gt
        overall = midi_close & rest_correct
        if mask is not None:
            overall &= mask

        self.correct += overall.sum()
        self.total += midi_gt.numel() if mask is None else mask.sum()

    def compute(self) -> Tensor:
        return self.correct / self.total

modules/metrics/midi_acc.py
openvpi/SOME
import torch
import torch.nn as nn
import torch.nn.functional as F
from modules.conform.Gconform import Gmidi_conform



class midi_loss(nn.Module):
    def __init__(self):
        super().__init__()
        self.loss = nn.BCELoss()

    def forward(self, x, target):
        midiout, cutp = x
        midi_target, cutp_target = target

        cutploss = self.loss(cutp, cutp_target)
        midiloss = self.loss(midiout, midi_target)
        return midiloss, cutploss


class midi_conforms(nn.Module):
    def __init__(self, config):
        super().__init__()

        cfg = config['midi_extractor_args']
        cfg.update({'indim': config['units_dim'], 'outdim': config['midi_num_bins']})
        self.model = Gmidi_conform(**cfg)

    def forward(self, x, f0, mask=None,softmax=False,sig=False):

        midi,bound=self.model(x, f0, mask)
        if  sig:
            midi = torch.sigmoid(midi)

        if softmax:
            midi=F.softmax(midi,dim=2)


        return midi,bound

    def get_loss(self):
        return midi_loss()

modules/model/Gmidi_conform.py
openvpi/SOME
from .constants import *
from .model import E2E0
from .utils import to_local_average_f0, to_viterbi_f0
from .inference import RMVPE
from .spec import MelSpectrogram

modules/rmvpe/__init__.py
openvpi/SOME
SAMPLE_RATE = 16000

N_CLASS = 360

N_MELS = 128
MEL_FMIN = 30
MEL_FMAX = 8000
WINDOW_LENGTH = 1024
CONST = 1997.3794084376191

modules/rmvpe/constants.py
openvpi/SOME
import torch
import torch.nn as nn
from .constants import N_MELS


class ConvBlockRes(nn.Module):
    def __init__(self, in_channels, out_channels, momentum=0.01):
        super(ConvBlockRes, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=in_channels,
                      out_channels=out_channels,
                      kernel_size=(3, 3),
                      stride=(1, 1),
                      padding=(1, 1),
                      bias=False),
            nn.BatchNorm2d(out_channels, momentum=momentum),
            nn.ReLU(),

            nn.Conv2d(in_channels=out_channels,
                      out_channels=out_channels,
                      kernel_size=(3, 3),
                      stride=(1, 1),
                      padding=(1, 1),
                      bias=False),
            nn.BatchNorm2d(out_channels, momentum=momentum),
            nn.ReLU(),
        )
        if in_channels != out_channels:
            self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
            self.is_shortcut = True
        else:
            self.is_shortcut = False

    def forward(self, x):
        if self.is_shortcut:
            return self.conv(x) + self.shortcut(x)
        else:
            return self.conv(x) + x


class ResEncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01):
        super(ResEncoderBlock, self).__init__()
        self.n_blocks = n_blocks
        self.conv = nn.ModuleList()
        self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
        for i in range(n_blocks - 1):
            self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
        self.kernel_size = kernel_size
        if self.kernel_size is not None:
            self.pool = nn.AvgPool2d(kernel_size=kernel_size)

    def forward(self, x):
        for i in range(self.n_blocks):
            x = self.conv[i](x)
        if self.kernel_size is not None:
            return x, self.pool(x)
        else:
            return x


class ResDecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
        super(ResDecoderBlock, self).__init__()
        out_padding = (0, 1) if stride == (1, 2) else (1, 1)
        self.n_blocks = n_blocks
        self.conv1 = nn.Sequential(
            nn.ConvTranspose2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3),
                               stride=stride,
                               padding=(1, 1),
                               output_padding=out_padding,
                               bias=False),
            nn.BatchNorm2d(out_channels, momentum=momentum),
            nn.ReLU(),
        )
        self.conv2 = nn.ModuleList()
        self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
        for i in range(n_blocks-1):
            self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))

    def forward(self, x, concat_tensor):
        x = self.conv1(x)
        x = torch.cat((x, concat_tensor), dim=1)
        for i in range(self.n_blocks):
            x = self.conv2[i](x)
        return x


class Encoder(nn.Module):
    def __init__(self, in_channels, in_size, n_encoders, kernel_size, n_blocks, out_channels=16, momentum=0.01):
        super(Encoder, self).__init__()
        self.n_encoders = n_encoders
        self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
        self.layers = nn.ModuleList()
        self.latent_channels = []
        for i in range(self.n_encoders):
            self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum))
            self.latent_channels.append([out_channels, in_size])
            in_channels = out_channels
            out_channels *= 2
            in_size //= 2
        self.out_size = in_size
        self.out_channel = out_channels

    def forward(self, x):
        concat_tensors = []
        x = self.bn(x)
        for i in range(self.n_encoders):
            _, x = self.layers[i](x)
            concat_tensors.append(_)
        return x, concat_tensors


class Intermediate(nn.Module):
    def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
        super(Intermediate, self).__init__()
        self.n_inters = n_inters
        self.layers = nn.ModuleList()
        self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum))
        for i in range(self.n_inters-1):
            self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum))

    def forward(self, x):
        for i in range(self.n_inters):
            x = self.layers[i](x)
        return x


class Decoder(nn.Module):
    def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
        super(Decoder, self).__init__()
        self.layers = nn.ModuleList()
        self.n_decoders = n_decoders
        for i in range(self.n_decoders):
            out_channels = in_channels // 2
            self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum))
            in_channels = out_channels

    def forward(self, x, concat_tensors):
        for i in range(self.n_decoders):
            x = self.layers[i](x, concat_tensors[-1-i])
        return x


class TimbreFilter(nn.Module):
    def __init__(self, latent_rep_channels):
        super(TimbreFilter, self).__init__()
        self.layers = nn.ModuleList()
        for latent_rep in latent_rep_channels:
            self.layers.append(ConvBlockRes(latent_rep[0], latent_rep[0]))

    def forward(self, x_tensors):
        out_tensors = []
        for i, layer in enumerate(self.layers):
            out_tensors.append(layer(x_tensors[i]))
        return out_tensors


class DeepUnet0(nn.Module):
    def __init__(self, kernel_size, n_blocks, en_de_layers=5, inter_layers=4, in_channels=1, en_out_channels=16):
        super(DeepUnet0, self).__init__()
        self.encoder = Encoder(in_channels, N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels)
        self.intermediate = Intermediate(self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks)
        self.tf = TimbreFilter(self.encoder.latent_channels)
        self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks)

    def forward(self, x):
        x, concat_tensors = self.encoder(x)
        x = self.intermediate(x)
        x = self.decoder(x, concat_tensors)
        return x

modules/rmvpe/deepunet.py
openvpi/SOME
import numpy as np
import torch
import torch.nn.functional as F
from torchaudio.transforms import Resample

from utils.pitch_utils import interp_f0, resample_align_curve
from .constants import *
from .model import E2E0
from .spec import MelSpectrogram
from .utils import to_local_average_f0, to_viterbi_f0


class RMVPE:
    def __init__(self, model_path, hop_length=160, device=None):
        self.resample_kernel = {}
        if device is None:
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = device
        self.model = E2E0(4, 1, (2, 2)).eval().to(self.device)
        ckpt = torch.load(model_path, map_location=self.device)
        self.model.load_state_dict(ckpt['model'], strict=False)
        self.mel_extractor = MelSpectrogram(
            N_MELS, SAMPLE_RATE, WINDOW_LENGTH, hop_length, None, MEL_FMIN, MEL_FMAX
        ).to(self.device)

    @torch.no_grad()
    def mel2hidden(self, mel):
        n_frames = mel.shape[-1]
        mel = F.pad(mel, (0, 32 * ((n_frames - 1) // 32 + 1) - n_frames), mode='constant')
        hidden = self.model(mel)
        return hidden[:, :n_frames]

    def decode(self, hidden, thred=0.03, use_viterbi=False):
        if use_viterbi:
            f0 = to_viterbi_f0(hidden, thred=thred)
        else:
            f0 = to_local_average_f0(hidden, thred=thred)
        return f0

    def infer_from_audio(self, audio, sample_rate=16000, thred=0.03, use_viterbi=False):
        audio = torch.from_numpy(audio).float().unsqueeze(0).to(self.device)
        if sample_rate == 16000:
            audio_res = audio
        else:
            key_str = str(sample_rate)
            if key_str not in self.resample_kernel:
                self.resample_kernel[key_str] = Resample(sample_rate, 16000, lowpass_filter_width=128)
            self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device)
            audio_res = self.resample_kernel[key_str](audio)
        mel = self.mel_extractor(audio_res, center=True)
        hidden = self.mel2hidden(mel)
        f0 = self.decode(hidden, thred=thred, use_viterbi=use_viterbi)
        return f0

    def get_pitch(self, waveform, sample_rate, hop_size, length, interp_uv=False):
        f0 = self.infer_from_audio(waveform, sample_rate=sample_rate)
        uv = f0 == 0
        f0, uv = interp_f0(f0, uv)

        time_step = hop_size / sample_rate
        f0_res = resample_align_curve(f0, 0.01, time_step, length)
        uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5
        if not interp_uv:
            f0_res[uv_res] = 0
        return f0_res, uv_res

modules/rmvpe/inference.py
openvpi/SOME
from torch import nn

from .constants import *
from .deepunet import DeepUnet0
from .seq import BiGRU


class E2E0(nn.Module):
    def __init__(self, n_blocks, n_gru, kernel_size, en_de_layers=5, inter_layers=4, in_channels=1,
                 en_out_channels=16):
        super(E2E0, self).__init__()
        self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels)
        self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
        if n_gru:
            self.fc = nn.Sequential(
                BiGRU(3 * N_MELS, 256, n_gru),
                nn.Linear(512, N_CLASS),
                nn.Dropout(0.25),
                nn.Sigmoid()
            )
        else:
            self.fc = nn.Sequential(
                nn.Linear(3 * N_MELS, N_CLASS),
                nn.Dropout(0.25),
                nn.Sigmoid()
            )

    def forward(self, mel):
        mel = mel.transpose(-1, -2).unsqueeze(1)
        x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
        x = self.fc(x)
        return x

modules/rmvpe/model.py
openvpi/SOME
import torch.nn as nn


class BiGRU(nn.Module):
    def __init__(self, input_features, hidden_features, num_layers):
        super(BiGRU, self).__init__()
        self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True)

    def forward(self, x):
        return self.gru(x)[0]

modules/rmvpe/seq.py
openvpi/SOME
import torch
import numpy as np
import torch.nn.functional as F
from librosa.filters import mel


class MelSpectrogram(torch.nn.Module):
    def __init__(
            self,
            n_mel_channels,
            sampling_rate,
            win_length,
            hop_length,
            n_fft=None,
            mel_fmin=0,
            mel_fmax=None,
            clamp=1e-5
    ):
        super().__init__()
        n_fft = win_length if n_fft is None else n_fft
        self.hann_window = {}
        mel_basis = mel(
            sr=sampling_rate,
            n_fft=n_fft,
            n_mels=n_mel_channels,
            fmin=mel_fmin,
            fmax=mel_fmax,
            htk=True)
        mel_basis = torch.from_numpy(mel_basis).float()
        self.register_buffer("mel_basis", mel_basis)
        self.n_fft = win_length if n_fft is None else n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.sampling_rate = sampling_rate
        self.n_mel_channels = n_mel_channels
        self.clamp = clamp

    def forward(self, audio, keyshift=0, speed=1, center=True):
        factor = 2 ** (keyshift / 12)
        n_fft_new = int(np.round(self.n_fft * factor))
        win_length_new = int(np.round(self.win_length * factor))
        hop_length_new = int(np.round(self.hop_length * speed))

        keyshift_key = str(keyshift) + '_' + str(audio.device)
        if keyshift_key not in self.hann_window:
            self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(audio.device)
        if center:
            pad_left = win_length_new // 2
            pad_right = (win_length_new + 1) // 2
            audio = F.pad(audio, (pad_left, pad_right))

        fft = torch.stft(
            audio,
            n_fft=n_fft_new,
            hop_length=hop_length_new,
            win_length=win_length_new,
            window=self.hann_window[keyshift_key],
            center=False,
            return_complex=True
        )
        magnitude = fft.abs()

        if keyshift != 0:
            size = self.n_fft // 2 + 1
            resize = magnitude.size(1)
            if resize < size:
                magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
            magnitude = magnitude[:, :size, :] * self.win_length / win_length_new

        mel_output = torch.matmul(self.mel_basis, magnitude)
        log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
        return log_mel_spec

modules/rmvpe/spec.py
openvpi/SOME
import librosa
import numpy as np
import torch

from .constants import *


def to_local_average_f0(hidden, center=None, thred=0.03):
    idx = torch.arange(N_CLASS, device=hidden.device)[None, None, :]  # [B=1, T=1, N]
    idx_cents = idx * 20 + CONST  # [B=1, N]
    if center is None:
        center = torch.argmax(hidden, dim=2, keepdim=True)  # [B, T, 1]
    start = torch.clip(center - 4, min=0)  # [B, T, 1]
    end = torch.clip(center + 5, max=N_CLASS)  # [B, T, 1]
    idx_mask = (idx >= start) & (idx < end)  # [B, T, N]
    weights = hidden * idx_mask  # [B, T, N]
    product_sum = torch.sum(weights * idx_cents, dim=2)  # [B, T]
    weight_sum = torch.sum(weights, dim=2)  # [B, T]
    cents = product_sum / (weight_sum + (weight_sum == 0))  # avoid dividing by zero, [B, T]
    f0 = 10 * 2 ** (cents / 1200)
    uv = hidden.max(dim=2)[0] < thred  # [B, T]
    f0 = f0 * ~uv
    return f0.squeeze(0).cpu().numpy()


def to_viterbi_f0(hidden, thred=0.03):
    # Create viterbi transition matrix
    if not hasattr(to_viterbi_f0, 'transition'):
        xx, yy = np.meshgrid(range(N_CLASS), range(N_CLASS))
        transition = np.maximum(30 - abs(xx - yy), 0)
        transition = transition / transition.sum(axis=1, keepdims=True)
        to_viterbi_f0.transition = transition

    # Convert to probability
    prob = hidden.squeeze(0).cpu().numpy()
    prob = prob.T
    prob = prob / prob.sum(axis=0)

    # Perform viterbi decoding
    path = librosa.sequence.viterbi(prob, to_viterbi_f0.transition).astype(np.int64)
    center = torch.from_numpy(path).unsqueeze(0).unsqueeze(-1).to(hidden.device)

    return to_local_average_f0(hidden, center=center, thred=thred)

modules/rmvpe/utils.py
openvpi/SOME
from .base_binarizer import BaseBinarizer
from .me_binarizer import MIDIExtractionBinarizer
from .me_quant_binarizer import QuantizedMIDIExtractionBinarizer

preprocessing/__init__.py
openvpi/SOME
import pathlib
import random
import warnings
from copy import deepcopy

import numpy as np
import torch
from tqdm import tqdm

from utils.indexed_datasets import IndexedDatasetBuilder
from utils.multiprocess_utils import chunked_multiprocess_run


class BinarizationError(Exception):
    pass


class BaseBinarizer:
    """
        Base class for data processing.
        1. *process* and *process_data_split*:
            process entire data, generate the train-test split (support parallel processing);
        2. *process_item*:
            process singe piece of data;
        3. *get_pitch*:
            infer the pitch using some algorithm;
        4. *get_align*:
            get the alignment using 'mel2ph' format (see https://arxiv.org/abs/1905.09263).
        5. phoneme encoder, voice encoder, etc.

        Subclasses should define:
        1. *load_metadata*:
            how to read multiple datasets from files;
        2. *train_item_names*, *valid_item_names*, *test_item_names*:
            how to split the dataset;
        3. load_ph_set:
            the phoneme set.
    """

    def __init__(self, config: dict, data_attrs=None):
        self.config = config
        self.raw_data_dirs = [pathlib.Path(d) for d in config['raw_data_dir']]
        self.binary_data_dir = pathlib.Path(self.config['binary_data_dir'])
        self.data_attrs = [] if data_attrs is None else data_attrs

        self.binarization_args = self.config['binarization_args']
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.items = {}
        self.item_names: list = None
        self._train_item_names: list = None
        self._valid_item_names: list = None

        self.timestep = self.config['hop_size'] / self.config['audio_sample_rate']

    def load_meta_data(self, raw_data_dir: pathlib.Path, ds_id):
        raise NotImplementedError()

    def split_train_valid_set(self):
        """
        Split the dataset into training set and validation set.
        :return: train_item_names, valid_item_names
        """
        prefixes = set([str(pr) for pr in self.config['test_prefixes']])
        valid_item_names = set()
        # Add prefixes that specified speaker index and matches exactly item name to test set
        for prefix in deepcopy(prefixes):
            if prefix in self.item_names:
                valid_item_names.add(prefix)
                prefixes.remove(prefix)
        # Add prefixes that exactly matches item name without speaker id to test set
        for prefix in deepcopy(prefixes):
            matched = False
            for name in self.item_names:
                if name.split(':')[-1] == prefix:
                    valid_item_names.add(name)
                    matched = True
            if matched:
                prefixes.remove(prefix)
        # Add names with one of the remaining prefixes to test set
        for prefix in deepcopy(prefixes):
            matched = False
            for name in self.item_names:
                if name.startswith(prefix):
                    valid_item_names.add(name)
                    matched = True
            if matched:
                prefixes.remove(prefix)
        for prefix in deepcopy(prefixes):
            matched = False
            for name in self.item_names:
                if name.split(':')[-1].startswith(prefix):
                    valid_item_names.add(name)
                    matched = True
            if matched:
                prefixes.remove(prefix)

        if len(prefixes) != 0:
            warnings.warn(
                f'The following rules in test_prefixes have no matching names in the dataset: {sorted(prefixes)}',
                category=UserWarning
            )
            warnings.filterwarnings('default')

        valid_item_names = sorted(list(valid_item_names))
        assert len(valid_item_names) > 0, 'Validation set is empty!'
        train_item_names = [x for x in self.item_names if x not in set(valid_item_names)]
        assert len(train_item_names) > 0, 'Training set is empty!'

        return train_item_names, valid_item_names

    @property
    def train_item_names(self):
        return self._train_item_names

    @property
    def valid_item_names(self):
        return self._valid_item_names

    def meta_data_iterator(self, prefix):
        if prefix == 'train':
            item_names = self.train_item_names
        else:
            item_names = self.valid_item_names
        for item_name in item_names:
            meta_data = self.items[item_name]
            yield item_name, meta_data

    def process(self):
        # load each dataset
        for ds_id, data_dir in zip(range(len(self.raw_data_dirs)), self.raw_data_dirs):
            self.load_meta_data(pathlib.Path(data_dir), ds_id=ds_id)
        self.item_names = sorted(list(self.items.keys()))
        self._train_item_names, self._valid_item_names = self.split_train_valid_set()

        if self.binarization_args['shuffle']:
            random.seed(self.config['seed'])
            random.shuffle(self.item_names)

        self.binary_data_dir.mkdir(parents=True, exist_ok=True)
        self.check_coverage()

        # Process valid set and train set
        try:
            self.process_dataset('valid')
            self.process_dataset(
                'train',
                num_workers=int(self.binarization_args['num_workers']),
                apply_augmentation=True
            )
        except KeyboardInterrupt:
            exit(-1)

    def check_coverage(self):
        pass

    def process_dataset(self, prefix, num_workers=0, apply_augmentation=False):
        args = []
        builder = IndexedDatasetBuilder(self.binary_data_dir, prefix=prefix, allowed_attr=self.data_attrs)
        lengths = []
        total_sec = 0
        total_raw_sec = 0

        for item_name, meta_data in self.meta_data_iterator(prefix):
            args.append([item_name, meta_data, apply_augmentation])

        def postprocess(_item, _is_raw=True):
            nonlocal total_sec, total_raw_sec
            if _item is None:
                return
            builder.add_item(_item)
            lengths.append(_item['length'])
            total_sec += _item['seconds']
            if _is_raw:
                total_raw_sec += _item['seconds']

        try:
            if num_workers > 0:
                # code for parallel processing
                for items in tqdm(
                        chunked_multiprocess_run(self.process_item, args, num_workers=num_workers),
                        total=len(list(self.meta_data_iterator(prefix)))
                ):
                    for i, item in enumerate(items):
                        postprocess(item, i == 0)
            else:
                # code for single cpu processing
                for a in tqdm(args):
                    items = self.process_item(*a)
                    for i, item in enumerate(items):
                        postprocess(item, i == 0)
        except KeyboardInterrupt:
            builder.finalize()
            raise

        builder.finalize()
        with open(self.binary_data_dir / f'{prefix}.lengths', 'wb') as f:
            # noinspection PyTypeChecker
            np.save(f, lengths)

        if apply_augmentation:
            print(f'| {prefix} total duration (before augmentation): {total_raw_sec:.2f}s')
            print(
                f'| {prefix} total duration (after augmentation): {total_sec:.2f}s ({total_sec / total_raw_sec:.2f}x)')
        else:
            print(f'| {prefix} total duration: {total_raw_sec:.2f}s')

    def process_item(self, item_name, meta_data, allow_aug=False):
        raise NotImplementedError()

preprocessing/base_binarizer.py
openvpi/SOME
import copy
import csv
import json
import os
import pathlib
import random

import librosa
import numpy as np
import torch
from scipy import interpolate

import modules.contentvec
import modules.rmvpe
from modules.commons import LengthRegulator
from utils.binarizer_utils import merge_slurs, merge_rests, get_mel2ph_torch, get_pitch_parselmouth
from utils.pitch_utils import resample_align_curve
from utils.plot import distribution_to_figure
from .base_binarizer import BaseBinarizer

os.environ["OMP_NUM_THREADS"] = "1"
MIDI_EXTRACTION_ITEM_ATTRIBUTES = [
    'units',  # contentvec units, float32[T_s, 256]
    'pitch',  # actual pitch in semitones, float32[T_s,]
    'note_midi',  # note-level MIDI pitch, float32[T_n,]
    'note_rest',  # flags for rest notes, bool[T_n,]
    'note_dur',  # durations of notes, in number of frames, int64[T_n,]
    'unit2note',  # mel2ph format for alignment between units and notes
]

# These modules are used as global variables due to a PyTorch shared memory bug on Windows platforms.
# See https://github.com/pytorch/pytorch/issues/100358
contentvec = None
mel_spec = None
rmvpe = None


class MIDIExtractionBinarizer(BaseBinarizer):
    def __init__(self, config: dict):
        super().__init__(config, data_attrs=MIDI_EXTRACTION_ITEM_ATTRIBUTES)
        self.lr = LengthRegulator().to(self.device)
        self.merge_rest = self.binarization_args['merge_rest']
        self.merge_slur = self.binarization_args['merge_slur']
        self.slur_tolerance = self.binarization_args.get('slur_tolerance')
        self.round_midi = self.binarization_args.get('round_midi', False)
        self.key_shift_min, self.key_shift_max = self.config['key_shift_range']

    def load_meta_data(self, raw_data_dir: pathlib.Path, ds_id):
        meta_data_dict = {}
        if (raw_data_dir / 'transcriptions.csv').exists():
            for utterance_label in csv.DictReader(
                    open(raw_data_dir / 'transcriptions.csv', 'r', encoding='utf-8')
            ):
                item_name = utterance_label['name']
                temp_dict = {
                    'wav_fn': str(raw_data_dir / 'wavs' / f'{item_name}.wav')
                }
                ds_path = raw_data_dir / 'wavs' / f'{item_name}.ds'
                with open(ds_path, 'r', encoding='utf8') as f:
                    ds = json.load(f)
                    if isinstance(ds, list):
                        ds = ds[0]
                # normalize
                note_seq = [
                    librosa.midi_to_note(
                        np.clip(
                            librosa.note_to_midi(n, round_midi=self.round_midi),
                            a_min=0, a_max=127
                        ),
                        cents=not self.round_midi, unicode=False
                    ) if n != 'rest' else 'rest'
                    for n in ds['note_seq'].split()
                ]
                note_slur = [bool(int(s)) for s in ds['note_slur'].split()]
                note_dur = [float(x) for x in ds['note_dur'].split()]

                # if not len(note_seq) == len(note_slur) == len(note_dur):
                #     continue
                assert len(note_seq) == len(note_slur) == len(note_dur), \
                    f'Lengths of note_seq, note_slur and note_dur mismatch in \'{item_name}\'.'
                assert any([note != 'rest' for note in note_seq]), \
                    f'All notes are rest in \'{item_name}\'.'

                if self.merge_slur:
                    # merge slurs with the same pitch
                    note_seq, note_dur = merge_slurs(note_seq, note_dur, note_slur, tolerance=self.slur_tolerance)

                if self.merge_rest:
                    # merge continuous rest notes
                    note_seq, note_dur = merge_rests(note_seq, note_dur)

                temp_dict['note_seq'] = note_seq
                temp_dict['note_dur'] = note_dur

                meta_data_dict[f'{ds_id}:{item_name}'] = temp_dict
        else:
            raise FileNotFoundError(
                f'transcriptions.csv not found in {raw_data_dir}.'
            )
        self.items.update(meta_data_dict)

    def check_coverage(self):
        super().check_coverage()
        # MIDI pitch distribution summary
        midi_map = {}
        for item_name in self.items:
            for midi in self.items[item_name]['note_seq']:
                if midi == 'rest':
                    continue
                midi = librosa.note_to_midi(midi, round_midi=True)
                if midi in midi_map:
                    midi_map[midi] += 1
                else:
                    midi_map[midi] = 1

        print('===== MIDI Pitch Distribution Summary =====')
        for i, key in enumerate(sorted(midi_map.keys())):
            if i == len(midi_map) - 1:
                end = '\n'
            elif i % 10 == 9:
                end = ',\n'
            else:
                end = ', '
            print(f'\'{librosa.midi_to_note(key, unicode=False)}\': {midi_map[key]}', end=end)

        # Draw graph.
        midis = sorted(midi_map.keys())
        notes = [librosa.midi_to_note(m, unicode=False) for m in range(midis[0], midis[-1] + 1)]
        plt = distribution_to_figure(
            title='MIDI Pitch Distribution Summary',
            x_label='MIDI Key', y_label='Number of occurrences',
            items=notes, values=[midi_map.get(m, 0) for m in range(midis[0], midis[-1] + 1)]
        )
        filename = self.binary_data_dir / 'midi_distribution.jpg'
        plt.savefig(fname=filename,
                    bbox_inches='tight',
                    pad_inches=0.25)
        print(f'| save summary to \'{filename}\'')

    def _process_item(self, waveform, meta_data, int_midi=False):
        wav_tensor = torch.from_numpy(waveform).to(self.device)
        units_encoder = self.config['units_encoder']
        if units_encoder == 'contentvec768l12':
            global contentvec
            if contentvec is None:
                contentvec = modules.contentvec.ContentVec768L12(self.config['units_encoder_ckpt'], device=self.device)
            units = contentvec(wav_tensor).squeeze(0).cpu().numpy()
        elif units_encoder == 'mel':
            global mel_spec
            if mel_spec is None:
                mel_spec = modules.rmvpe.MelSpectrogram(
                    n_mel_channels=self.config['units_dim'], sampling_rate=self.config['audio_sample_rate'],
                    win_length=self.config['win_size'], hop_length=self.config['hop_size'],
                    mel_fmin=self.config['fmin'], mel_fmax=self.config['fmax']
                ).to(self.device)
            units = mel_spec(wav_tensor.unsqueeze(0)).transpose(1, 2).squeeze(0).cpu().numpy()
        else:
            raise NotImplementedError(f'Invalid units encoder: {units_encoder}')
        assert len(units.shape) == 2 and units.shape[1] == self.config['units_dim'], \
            f'Shape of units must be [T, units_dim], but is {units.shape}.'
        length = units.shape[0]
        seconds = length * self.config['hop_size'] / self.config['audio_sample_rate']
        processed_input = {
            'seconds': seconds,
            'length': length,
            'units': units
        }

        f0_algo = self.config['pe']
        if f0_algo == 'parselmouth':
            f0, _ = get_pitch_parselmouth(
                waveform, sample_rate=self.config['audio_sample_rate'],
                hop_size=self.config['hop_size'], length=length, interp_uv=True
            )
        elif f0_algo == 'rmvpe':
            global rmvpe
            if rmvpe is None:
                rmvpe = modules.rmvpe.RMVPE(self.config['pe_ckpt'], device=self.device)
            f0, _ = rmvpe.get_pitch(
                waveform, sample_rate=self.config['audio_sample_rate'],
                hop_size=rmvpe.mel_extractor.hop_length,
                length=(waveform.shape[0] + rmvpe.mel_extractor.hop_length - 1) // rmvpe.mel_extractor.hop_length,
                interp_uv=True
            )
            f0 = resample_align_curve(
                f0,
                original_timestep=rmvpe.mel_extractor.hop_length / self.config['audio_sample_rate'],
                target_timestep=self.config['hop_size'] / self.config['audio_sample_rate'],
                align_length=length
            )
        else:
            raise NotImplementedError(f'Invalid pitch extractor: {f0_algo}')
        pitch = librosa.hz_to_midi(f0)
        processed_input['pitch'] = pitch

        note_midi = np.array(
            [(librosa.note_to_midi(n, round_midi=int_midi) if n != 'rest' else -1) for n in meta_data['note_seq']],
            dtype=np.int64 if int_midi else np.float32
        )
        note_rest = note_midi < 0
        interp_func = interpolate.interp1d(
            np.where(~note_rest)[0], note_midi[~note_rest],
            kind='nearest', fill_value='extrapolate'
        )
        note_midi[note_rest] = interp_func(np.where(note_rest)[0])
        processed_input['note_midi'] = note_midi
        processed_input['note_rest'] = note_rest

        note_dur_sec = torch.FloatTensor(meta_data['note_dur']).to(self.device)
        note_acc = torch.round(torch.cumsum(note_dur_sec, dim=0) / self.timestep + 0.5).long()
        note_dur = torch.diff(note_acc, dim=0, prepend=torch.LongTensor([0]).to(self.device))
        processed_input['note_dur'] = note_dur.cpu().numpy()
        unit2note = get_mel2ph_torch(
            self.lr, note_dur_sec, processed_input['length'], self.timestep, device=self.device
        )
        processed_input['unit2note'] = unit2note.cpu().numpy()
        return processed_input

    @torch.no_grad()
    def process_item(self, item_name, meta_data, allow_aug=False):
        waveform, _ = librosa.load(meta_data['wav_fn'], sr=self.config['audio_sample_rate'], mono=True)

        processed_input = self._process_item(waveform, meta_data, int_midi=False)
        items = [processed_input]
        if not allow_aug:
            return items

        wav_tensor = torch.from_numpy(waveform).to(self.device)
        for _ in range(self.config['key_shift_factor']):
            assert mel_spec is not None, 'Units encoder must be mel if augmentation is applied!'
            key_shift = random.random() * (self.key_shift_max - self.key_shift_min) + self.key_shift_min
            if self.round_midi:
                key_shift = round(key_shift)
            processed_input_aug = copy.deepcopy(processed_input)
            assert isinstance(mel_spec, modules.rmvpe.MelSpectrogram)
            processed_input_aug['units'] = mel_spec(
                wav_tensor.unsqueeze(0), keyshift=key_shift
            ).transpose(1, 2).squeeze(0).cpu().numpy()
            processed_input_aug['pitch'] += key_shift
            processed_input_aug['note_midi'] += key_shift
            items.append(processed_input_aug)

        return items

preprocessing/me_binarizer.py
openvpi/SOME
import copy
import os
import random

import librosa
import torch

import modules.contentvec
import modules.rmvpe
from .me_binarizer import MIDIExtractionBinarizer

os.environ["OMP_NUM_THREADS"] = "1"
QUANTIZED_MIDI_EXTRACTION_ITEM_ATTRIBUTES = [
    'units',  # contentvec units, float32[T_s, 256]
    'pitch',  # actual pitch in semitones, float32[T_s,]
    'note_midi',  # note-level MIDI pitch (0-127: MIDI, 128: rest) int64[T_n,]
    'note_dur',  # durations of notes, in number of frames, int64[T_n,]
    'unit2note',  # mel2ph format for alignment between units and notes
]


class QuantizedMIDIExtractionBinarizer(MIDIExtractionBinarizer):
    def __init__(self, config: dict):
        super().__init__(config)
        self.round_midi = True
        self.data_attrs = QUANTIZED_MIDI_EXTRACTION_ITEM_ATTRIBUTES

    def process_item(self, item_name, meta_data, allow_aug=False):
        waveform, _ = librosa.load(meta_data['wav_fn'], sr=self.config['audio_sample_rate'], mono=True)

        processed_input = self._process_item(waveform, meta_data, int_midi=True)
        processed_input['note_midi'][processed_input['note_rest']] = 128
        items = [processed_input]
        if not allow_aug:
            return items

        from .me_binarizer import mel_spec
        wav_tensor = torch.from_numpy(waveform).to(self.device)
        for _ in range(self.config['key_shift_factor']):
            assert mel_spec is not None, 'Units encoder must be mel if augmentation is applied!'
            key_shift = random.randint(int(self.key_shift_min), int(self.key_shift_max))
            processed_input_aug = copy.deepcopy(processed_input)
            assert isinstance(mel_spec, modules.rmvpe.MelSpectrogram)
            processed_input_aug['units'] = mel_spec(
                wav_tensor.unsqueeze(0), keyshift=key_shift
            ).transpose(1, 2).squeeze(0).cpu().numpy()
            processed_input_aug['pitch'] += key_shift
            processed_input_aug['note_midi'][~processed_input_aug['note_rest']] += key_shift
            items.append(processed_input_aug)

        return items

preprocessing/me_quant_binarizer.py
openvpi/SOME
import pathlib

import click
import torch


@click.command(help='Simplify a checkpoint file, dropping all useless keys for inference.')
@click.argument('input_ckpt', metavar='INPUT_CKPT')
@click.argument('output_ckpt', metavar='OUTPUT_CKPT')
def simplify(input_ckpt, output_ckpt):
    input_ckpt_path = pathlib.Path(input_ckpt)
    output_ckpt_path = pathlib.Path(output_ckpt)
    ckpt = torch.load(input_ckpt_path)
    ckpt = {
        'state_dict': ckpt['state_dict']
    }
    torch.save(ckpt, output_ckpt_path)


if __name__ == '__main__':
    simplify()

simplify.py
openvpi/SOME
import importlib
import logging
import os
import pathlib
import sys

import click
import lightning.pytorch as pl
import torch.utils.data
import yaml
from lightning.pytorch.loggers import TensorBoardLogger

import training.base_task
from utils.config_utils import read_full_config, print_config
from utils.training_utils import (
    DsModelCheckpoint, DsTQDMProgressBar,
    get_latest_checkpoint_path, get_strategy
)

torch.multiprocessing.set_sharing_strategy(os.getenv('TORCH_SHARE_STRATEGY', 'file_system'))

log_format = '%(asctime)s %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
                    format=log_format, datefmt='%m/%d %I:%M:%S %p')


@click.command(help='Train a SOME model')
@click.option('--config', required=True, metavar='FILE', help='Path to the configuration file')
@click.option('--exp_name', required=True, metavar='EXP', help='Name of the experiment')
@click.option('--work_dir', required=False, metavar='DIR', help='Directory to save the experiment')
def train(config, exp_name, work_dir):
    config = pathlib.Path(config)
    config = read_full_config(config)
    print_config(config)
    if work_dir is None:
        work_dir = pathlib.Path(__file__).parent / 'experiments'
    else:
        work_dir = pathlib.Path(work_dir)
    work_dir = work_dir / exp_name
    assert not work_dir.exists() or work_dir.is_dir(), f'Path \'{work_dir}\' is not a directory.'
    work_dir.mkdir(parents=True, exist_ok=True)
    with open(work_dir / 'config.yaml', 'w', encoding='utf8') as f:
        yaml.safe_dump(config, f)

    if config['ddp_backend'] == 'nccl_no_p2p':
        print("Disabling NCCL P2P")
        os.environ['NCCL_P2P_DISABLE'] = '1'

    pl.seed_everything(config['seed'], workers=True)
    assert config['task_cls'] != ''
    pkg = ".".join(config["task_cls"].split(".")[:-1])
    cls_name = config["task_cls"].split(".")[-1]
    task_cls = getattr(importlib.import_module(pkg), cls_name)
    assert issubclass(task_cls, training.BaseTask), f'Task class {task_cls} is not a subclass of {training.BaseTask}.'

    task = task_cls(config=config)

    # work_dir = pathlib.Path(config['work_dir'])
    trainer = pl.Trainer(
        accelerator=config['pl_trainer_accelerator'],
        devices=config['pl_trainer_devices'],
        num_nodes=config['pl_trainer_num_nodes'],
        strategy=get_strategy(
            accelerator=config['pl_trainer_accelerator'],
            devices=config['pl_trainer_devices'],
            num_nodes=config['pl_trainer_num_nodes'],
            strategy=config['pl_trainer_strategy'],
            backend=config['ddp_backend']
        ),
        precision=config['pl_trainer_precision'],
        callbacks=[
            DsModelCheckpoint(
                dirpath=work_dir,
                filename='model_ckpt_steps_{step}',
                auto_insert_metric_name=False,
                monitor='step',
                mode='max',
                save_last=False,
                # every_n_train_steps=config['val_check_interval'],
                save_top_k=config['num_ckpt_keep'],
                permanent_ckpt_start=config['permanent_ckpt_start'],
                permanent_ckpt_interval=config['permanent_ckpt_interval'],
                verbose=True
            ),
            # LearningRateMonitor(logging_interval='step'),
            DsTQDMProgressBar(),
        ],
        logger=TensorBoardLogger(
            save_dir=str(work_dir),
            name='lightning_logs',
            version='lastest'
        ),
        gradient_clip_val=config['clip_grad_norm'],
        val_check_interval=config['val_check_interval'] * config['accumulate_grad_batches'],
        # so this is global_steps
        check_val_every_n_epoch=None,
        log_every_n_steps=1,
        max_steps=config['max_updates'],
        use_distributed_sampler=False,
        num_sanity_val_steps=config['num_sanity_val_steps'],
        accumulate_grad_batches=config['accumulate_grad_batches']
    )
    trainer.fit(task, ckpt_path=get_latest_checkpoint_path(work_dir))


os.environ['TORCH_CUDNN_V8_API_ENABLED'] = '1'  # Prevent unacceptable slowdowns when using 16 precision


if __name__ == '__main__':
    train()

train.py
openvpi/SOME
from .base_task import BaseTask
from .me_task import MIDIExtractionTask
from .me_quant_task import QuantizedMIDIExtractionTask

training/__init__.py
openvpi/SOME
import logging
import os
import pathlib
import sys
from typing import Dict

import lightning.pytorch as pl
import matplotlib
import numpy as np
import torch.utils.data
from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_info, rank_zero_only
from torch.utils.data import Dataset
from torchmetrics import Metric, MeanMetric

import utils
from utils.indexed_datasets import IndexedDataset
from utils.training_utils import (
    DsBatchSampler, DsEvalBatchSampler,
    get_latest_checkpoint_path
)

matplotlib.use('Agg')

torch.multiprocessing.set_sharing_strategy(os.getenv('TORCH_SHARE_STRATEGY', 'file_system'))

log_format = '%(asctime)s %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
                    format=log_format, datefmt='%m/%d %I:%M:%S %p')


class BaseDataset(Dataset):
    """
        Base class for datasets.
        1. *sizes*:
            clipped length if "max_frames" is set;
        2. *num_frames*:
            unclipped length.

        Subclasses should define:
        1. *collate*:
            take the longest data, pad other data to the same length;
        2. *__getitem__*:
            the index function.
    """

    def __init__(self, config: dict, data_dir, prefix, allow_aug=False):
        super().__init__()
        self.config = config
        self.prefix = prefix
        self.data_dir = data_dir if isinstance(data_dir, pathlib.Path) else pathlib.Path(data_dir)
        self.sizes = np.load(self.data_dir / f'{self.prefix}.lengths')
        self.indexed_ds = IndexedDataset(self.data_dir, self.prefix)
        self.allow_aug = allow_aug

    @property
    def _sizes(self):
        return self.sizes

    def __getitem__(self, index):
        return self.indexed_ds[index]

    def __len__(self):
        return len(self._sizes)

    def num_frames(self, index):
        return self.size(index)

    def size(self, index):
        """Return an example's size as a float or tuple. This value is used when
        filtering a dataset with ``--max-positions``."""
        return self._sizes[index]

    def collater(self, samples):
        return {
            'size': len(samples)
        }


class BaseTask(pl.LightningModule):
    """
        Base class for training tasks.
        1. *load_ckpt*:
            load checkpoint;
        2. *training_step*:
            record and log the loss;
        3. *optimizer_step*:
            run backwards step;
        4. *start*:
            load training configs, backup code, log to tensorboard, start training;
        5. *configure_ddp* and *init_ddp_connection*:
            start parallel training.

        Subclasses should define:
        1. *build_model*, *build_optimizer*, *build_scheduler*:
            how to build the model, the optimizer and the training scheduler;
        2. *_training_step*:
            one training step of the model;
        3. *on_validation_end* and *_on_validation_end*:
            postprocess the validation output.
    """

    def __init__(self, config: dict, *args, **kwargs):
        # dataset configs
        super().__init__(*args, **kwargs)
        self.dataset_cls = None
        self.config = config
        self.max_batch_frames = self.config['max_batch_frames']
        self.max_batch_size = self.config['max_batch_size']
        self.max_val_batch_frames = self.config['max_val_batch_frames']
        self.max_val_batch_size = self.config['max_val_batch_size']

        self.training_sampler = None
        self.model = None
        self.skip_immediate_validation = False
        self.skip_immediate_ckpt_save = False

        self.valid_losses: Dict[str, Metric] = {
            'total_loss': MeanMetric()
        }
        self.valid_metric_names = set()

    ###########
    # Training, validation and testing
    ###########
    def setup(self, stage):
        self.model = self.build_model()
        self.unfreeze_all_params()
        if self.config['freezing_enabled']:
            self.freeze_params()
        if self.config['finetune_enabled'] and get_latest_checkpoint_path(
                pathlib.Path(self.config['work_dir'])) is None:
            self.load_finetune_ckpt(self.load_pre_train_model())
        self.print_arch()
        self.build_losses_and_metrics()
        self.train_dataset = self.dataset_cls(
            config=self.config, data_dir=self.config['binary_data_dir'],
            prefix=self.config['train_set_name'], allow_aug=True
        )
        self.valid_dataset = self.dataset_cls(
            config=self.config, data_dir=self.config['binary_data_dir'],
            prefix=self.config['valid_set_name'], allow_aug=False
        )

    def get_need_freeze_state_dict_key(self, model_state_dict) -> list:
        key_list = []
        for i in self.config['frozen_params']:
            for j in model_state_dict:
                if j.startswith(i):
                    key_list.append(j)
        return list(set(key_list))

    def freeze_params(self) -> None:
        model_state_dict = self.state_dict().keys()
        freeze_key = self.get_need_freeze_state_dict_key(model_state_dict=model_state_dict)

        for i in freeze_key:
            params = self.get_parameter(i)

            params.requires_grad = False

    def unfreeze_all_params(self) -> None:
        for i in self.model.parameters():
            i.requires_grad = True

    def load_finetune_ckpt(
            self, state_dict
    ) -> None:

        adapt_shapes = self.config['finetune_strict_shapes']
        if not adapt_shapes:
            cur_model_state_dict = self.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]
        self.load_state_dict(state_dict, strict=False)

    def load_pre_train_model(self):

        pre_train_ckpt_path = self.config.get('finetune_ckpt_path')
        blacklist = self.config.get('finetune_ignored_params')
        if blacklist is None:
            blacklist = []
        # if whitelist is  None:
        #     raise RuntimeError("")

        if pre_train_ckpt_path is not None:
            ckpt = torch.load(pre_train_ckpt_path)

            state_dict = {}
            for i in ckpt['state_dict']:
                # if 'diffusion' in i:
                # if i in rrrr:
                #     continue
                skip = False
                for b in blacklist:
                    if i.startswith(b):
                        skip = True
                        break

                if skip:
                    continue

                state_dict[i] = ckpt['state_dict'][i]
                print(i)
            return state_dict
        else:
            raise RuntimeError("")

    def build_model(self):
        raise NotImplementedError()

    @rank_zero_only
    def print_arch(self):
        utils.print_arch(self.model)

    def build_losses_and_metrics(self):
        raise NotImplementedError()

    def register_metric(self, name: str, metric: Metric):
        assert isinstance(metric, Metric)
        setattr(self, name, metric)
        self.valid_metric_names.add(name)

    def run_model(self, sample, infer=False):
        """
        steps:
            1. run the full model
            2. calculate losses if not infer
        """
        raise NotImplementedError()

    def on_train_epoch_start(self):
        if self.training_sampler is not None:
            self.training_sampler.set_epoch(self.current_epoch)

    def _training_step(self, sample):
        """
        :return: total loss: torch.Tensor, loss_log: dict, other_log: dict
        """
        losses = self.run_model(sample)
        total_loss = sum(losses.values())
        return total_loss, {**losses, 'batch_size': float(sample['size'])}

    def training_step(self, sample, batch_idx, optimizer_idx=-1):
        total_loss, log_outputs = self._training_step(sample)

        # logs to progress bar
        self.log_dict(log_outputs, prog_bar=True, logger=False, on_step=True, on_epoch=False)
        self.log('lr', self.lr_schedulers().get_last_lr()[0], prog_bar=True, logger=False, on_step=True, on_epoch=False)
        # logs to tensorboard
        if self.global_step % self.config['log_interval'] == 0:
            tb_log = {f'training/{k}': v for k, v in log_outputs.items()}
            tb_log['training/lr'] = self.lr_schedulers().get_last_lr()[0]
            self.logger.log_metrics(tb_log, step=self.global_step)

        return total_loss

    # def on_before_optimizer_step(self, *args, **kwargs):
    #     self.log_dict(grad_norm(self, norm_type=2))

    def _on_validation_start(self):
        pass

    def on_validation_start(self):
        self._on_validation_start()
        for metric in self.valid_losses.values():
            metric.to(self.device)
            metric.reset()

    def _validation_step(self, sample, batch_idx):
        """

        :param sample:
        :param batch_idx:
        :return: loss_log: dict, weight: int
        """
        raise NotImplementedError()

    def validation_step(self, sample, batch_idx):
        """

        :param sample:
        :param batch_idx:
        """
        if self.skip_immediate_validation:
            rank_zero_debug(f"Skip validation {batch_idx}")
            return {}
        with torch.autocast(self.device.type, enabled=False):
            losses, weight = self._validation_step(sample, batch_idx)
        losses = {
            'total_loss': sum(losses.values()),
            **losses
        }
        for k, v in losses.items():
            if k not in self.valid_losses:
                self.valid_losses[k] = MeanMetric().to(self.device)
            self.valid_losses[k].update(v, weight=weight)
        return losses

    def on_validation_epoch_end(self):
        if self.skip_immediate_validation:
            self.skip_immediate_validation = False
            self.skip_immediate_ckpt_save = True
            return
        loss_vals = {k: v.compute() for k, v in self.valid_losses.items()}
        self.log('val_loss', loss_vals['total_loss'], on_epoch=True, prog_bar=True, logger=False, sync_dist=True)
        self.logger.log_metrics({f'validation/{k}': v for k, v in loss_vals.items()}, step=self.global_step)
        for metric in self.valid_losses.values():
            metric.reset()
        metric_vals = {k: getattr(self, k).compute() for k in self.valid_metric_names}
        self.logger.log_metrics({f'metrics/{k}': v for k, v in metric_vals.items()}, step=self.global_step)
        for metric_name in self.valid_metric_names:
            getattr(self, metric_name).reset()

    # noinspection PyMethodMayBeStatic
    def build_scheduler(self, optimizer):
        from utils import build_lr_scheduler_from_config

        scheduler_args = self.config['lr_scheduler_args']
        assert scheduler_args['scheduler_cls'] != ''
        scheduler = build_lr_scheduler_from_config(optimizer, scheduler_args)
        return scheduler

    # noinspection PyMethodMayBeStatic
    def build_optimizer(self, model):
        from utils import build_object_from_class_name

        optimizer_args = self.config['optimizer_args']
        assert optimizer_args['optimizer_cls'] != ''
        if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
            optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])
        optimizer = build_object_from_class_name(
            optimizer_args['optimizer_cls'],
            torch.optim.Optimizer,
            model.parameters(),
            **optimizer_args
        )
        return optimizer

    def configure_optimizers(self):
        optm = self.build_optimizer(self.model)
        scheduler = self.build_scheduler(optm)
        if scheduler is None:
            return optm
        return {
            "optimizer": optm,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1
            }
        }

    def train_dataloader(self):
        self.training_sampler = DsBatchSampler(
            self.train_dataset,
            max_batch_frames=self.max_batch_frames,
            max_batch_size=self.max_batch_size,
            num_replicas=(self.trainer.distributed_sampler_kwargs or {}).get('num_replicas', 1),
            rank=(self.trainer.distributed_sampler_kwargs or {}).get('rank', 0),
            sort_by_similar_size=self.config['sort_by_len'],
            required_batch_count_multiple=self.config['accumulate_grad_batches'],
            frame_count_grid=self.config['sampler_frame_count_grid'],
            shuffle_sample=True,
            shuffle_batch=False,
            seed=self.config['seed']
        )
        return torch.utils.data.DataLoader(self.train_dataset,
                                           collate_fn=self.train_dataset.collater,
                                           batch_sampler=self.training_sampler,
                                           num_workers=self.config['ds_workers'],
                                           prefetch_factor=self.config['dataloader_prefetch_factor'],
                                           pin_memory=True,
                                           persistent_workers=True)

    def val_dataloader(self):
        sampler = DsEvalBatchSampler(
            self.valid_dataset,
            max_batch_frames=self.max_val_batch_frames,
            max_batch_size=self.max_val_batch_size,
            rank=(self.trainer.distributed_sampler_kwargs or {}).get('rank', 0),
            batch_by_size=False
        )
        return torch.utils.data.DataLoader(self.valid_dataset,
                                           collate_fn=self.valid_dataset.collater,
                                           batch_sampler=sampler,
                                           num_workers=self.config['ds_workers'],
                                           prefetch_factor=self.config['dataloader_prefetch_factor'],
                                           shuffle=False)

    def test_dataloader(self):
        return self.val_dataloader()

    def on_test_start(self):
        self.on_validation_start()

    def test_step(self, sample, batch_idx):
        return self.validation_step(sample, batch_idx)

    def on_test_end(self):
        return self.on_validation_end()

    def on_save_checkpoint(self, checkpoint):
        checkpoint['trainer_stage'] = self.trainer.state.stage.value

    def on_load_checkpoint(self, checkpoint):
        from lightning.pytorch.trainer.states import RunningStage
        from utils import simulate_lr_scheduler
        if checkpoint.get('trainer_stage', '') == RunningStage.VALIDATING.value:
            self.skip_immediate_validation = True

        optimizer_args = self.config['optimizer_args']
        scheduler_args = self.config['lr_scheduler_args']

        if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
            optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])

        if checkpoint.get('optimizer_states', None):
            opt_states = checkpoint['optimizer_states']
            assert len(opt_states) == 1  # only support one optimizer
            opt_state = opt_states[0]
            for param_group in opt_state['param_groups']:
                for k, v in optimizer_args.items():
                    if k in param_group and param_group[k] != v:
                        if 'lr_schedulers' in checkpoint and checkpoint['lr_schedulers'] and k == 'lr':
                            continue
                        rank_zero_info(f'| Overriding optimizer parameter {k} from checkpoint: {param_group[k]} -> {v}')
                        param_group[k] = v
                if 'initial_lr' in param_group and param_group['initial_lr'] != optimizer_args['lr']:
                    rank_zero_info(
                        f'| Overriding optimizer parameter initial_lr from checkpoint: {param_group["initial_lr"]} -> {optimizer_args["lr"]}'
                    )
                    param_group['initial_lr'] = optimizer_args['lr']

        if checkpoint.get('lr_schedulers', None):
            assert checkpoint.get('optimizer_states', False)
            assert len(checkpoint['lr_schedulers']) == 1  # only support one scheduler
            checkpoint['lr_schedulers'][0] = simulate_lr_scheduler(
                optimizer_args, scheduler_args,
                step_count=checkpoint['global_step'],
                num_param_groups=len(checkpoint['optimizer_states'][0]['param_groups'])
            )
            for param_group, new_lr in zip(
                    checkpoint['optimizer_states'][0]['param_groups'],
                    checkpoint['lr_schedulers'][0]['_last_lr'],
            ):
                if param_group['lr'] != new_lr:
                    rank_zero_info(
                        f'| Overriding optimizer parameter lr from checkpoint: {param_group["lr"]} -> {new_lr}')
                    param_group['lr'] = new_lr

training/base_task.py
openvpi/SOME
import torch
import torch.nn.functional as F
from torch import nn

import modules.losses
import modules.metrics
from utils import build_object_from_class_name, collate_nd
from utils.infer_utils import decode_bounds_to_alignment, decode_note_sequence
from .base_task import BaseDataset
from .me_task import MIDIExtractionTask


class QuantizedMIDIExtractionDataset(BaseDataset):
    def collater(self, samples):
        batch = super().collater(samples)
        batch['units'] = collate_nd([s['units'] for s in samples])  # [B, T_s, C]
        batch['pitch'] = collate_nd([s['pitch'] for s in samples])  # [B, T_s]
        batch['note_midi'] = collate_nd([s['note_midi'] for s in samples], pad_value=-1)  # [B, T_n]
        batch['note_dur'] = collate_nd([s['note_dur'] for s in samples])  # [B, T_n]
        unit2note = collate_nd([s['unit2note'] for s in samples])
        batch['unit2note'] = unit2note
        batch['midi_idx'] = torch.gather(F.pad(batch['note_midi'], [1, 0], value=-1), 1, unit2note)
        bounds = torch.diff(
            unit2note, dim=1, prepend=unit2note.new_zeros((batch['size'], 1))
        ) > 0
        batch['bounds'] = bounds.float()
        return batch


class QuantizedMIDIExtractionTask(MIDIExtractionTask):
    def __init__(self, config: dict):
        super().__init__(config)
        self.dataset_cls = QuantizedMIDIExtractionDataset
        self.config = config

    # noinspection PyAttributeOutsideInit
    def build_model(self):

        model = build_object_from_class_name(self.config['model_cls'], nn.Module, config=self.config)

        return model

    def build_losses_and_metrics(self):
        self.midi_loss = nn.CrossEntropyLoss(ignore_index=-1)
        self.bound_loss = modules.losses.BinaryEMDLoss(bidirectional=False)
        self.register_metric('midi_acc', modules.metrics.MIDIAccuracy(tolerance=0.5))

    def run_model(self, sample, infer=False):
        """
        steps:
            1. run the full model
            2. calculate losses if not infer
        """
        spec = sample['units']  # [B, T_ph]
        # target = (sample['probs'],sample['bounds'])  # [B, T_s, M]
        mask = sample['unit2note'] > 0
        # mask=None

        f0 = sample['pitch']
        probs, bounds = self.model(x=spec, f0=f0, mask=mask, softmax=infer)

        if infer:
            return probs, bounds
        else:
            losses = {}

            if  self.cfg['use_bound_loss']:
                bound_loss = self.bound_loss(bounds, sample['bounds'])

                losses['bound_loss'] = bound_loss
            if self.cfg['use_midi_loss']:
                midi_loss = self.midi_loss(probs.transpose(1, 2), sample['midi_idx'])

                losses['midi_loss'] = midi_loss

            return losses

    def _validation_step(self, sample, batch_idx):
        losses = self.run_model(sample, infer=False)
        if batch_idx < self.config['num_valid_plots']:
            probs, bounds = self.run_model(sample, infer=True)
            unit2note_gt = sample['unit2note']
            masks = unit2note_gt > 0
            probs *= masks[..., None]
            bounds *= masks
            # probs: [B, T, 129] => [B, T, 128]
            probs_pred = probs[:, :, :-1]
            probs_gt = F.one_hot(sample['midi_idx'], num_classes=129)[:, :, :-1]
            self.plot_prob(batch_idx, probs_gt, probs_pred)

            unit2note_pred = decode_bounds_to_alignment(bounds) * masks
            midi_pred = probs.argmax(dim=-1)
            rest_pred = midi_pred == 128
            note_midi_pred, note_dur_pred, note_mask_pred = decode_note_sequence(
                unit2note_pred, midi_pred.clip(min=0, max=127), ~rest_pred & masks
            )
            note_rest_pred = ~note_mask_pred
            self.plot_boundary(
                batch_idx, bounds_gt=sample['bounds'], bounds_pred=bounds,
                dur_gt=sample['note_dur'], dur_pred=note_dur_pred
            )
            self.plot_final(
                batch_idx, sample['note_midi'], sample['note_dur'], sample['note_midi'] == 128,
                note_midi_pred, note_dur_pred, note_rest_pred, sample['pitch']
            )

            midi_pred = midi_pred.float()
            midi_pred[rest_pred] = -torch.inf  # rest part is set to -inf
            note_midi_gt = sample['note_midi'].float()
            note_rest_gt = sample['note_midi'] == 128
            note_midi_gt[note_rest_gt] = -torch.inf
            midi_gt = torch.gather(F.pad(note_midi_gt, [1, 0], value=-torch.inf), 1, unit2note_gt)
            self.plot_midi_curve(
                batch_idx, midi_gt=midi_gt, midi_pred=midi_pred, pitch=sample['pitch']
            )
            self.midi_acc.update(
                midi_pred=midi_pred, rest_pred=rest_pred, midi_gt=midi_gt, rest_gt=midi_gt < 0, mask=masks
            )

        return losses, sample['size']

training/me_quant_task.py
openvpi/SOME
import torch
import torch.nn.functional as F
from torch import nn

import modules.losses
import modules.metrics
from utils import build_object_from_class_name, collate_nd
from utils.infer_utils import decode_gaussian_blurred_probs, decode_bounds_to_alignment, decode_note_sequence
from utils.plot import boundary_to_figure, curve_to_figure, spec_to_figure, pitch_notes_to_figure
from .base_task import BaseDataset, BaseTask


class MIDIExtractionDataset(BaseDataset):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.midi_min = self.config['midi_min']
        self.midi_max = self.config['midi_max']
        self.num_bins = self.config['midi_num_bins']
        self.midi_deviation = self.config['midi_prob_deviation']
        self.interval = (self.midi_max - self.midi_min) / (self.num_bins - 1)  # align with centers of bins
        self.sigma = self.midi_deviation / self.interval

    def midi_to_bin(self, midi):
        return (midi - self.midi_min) / self.interval

    def collater(self, samples):
        batch = super().collater(samples)
        batch['units'] = collate_nd([s['units'] for s in samples])  # [B, T_s, C]
        batch['pitch'] = collate_nd([s['pitch'] for s in samples])  # [B, T_s]
        batch['note_midi'] = collate_nd([s['note_midi'] for s in samples])  # [B, T_n]
        batch['note_rest'] = collate_nd([s['note_rest'] for s in samples])  # [B, T_n]
        batch['note_dur'] = collate_nd([s['note_dur'] for s in samples])  # [B, T_n]

        miu = self.midi_to_bin(batch['note_midi'])[:, :, None]  # [B, T_n, 1]
        x = torch.arange(self.num_bins).float().reshape(1, 1, -1).to(miu.device)  # [1, 1, N]
        probs = ((x - miu) / self.sigma).pow(2).div(-2).exp()  # gaussian blur, [B, T_n, N]
        note_mask = collate_nd([torch.ones_like(s['note_rest']) for s in samples], pad_value=False)
        probs *= (note_mask[..., None] & ~batch['note_rest'][..., None])

        probs = F.pad(probs, [0, 0, 1, 0])
        unit2note = collate_nd([s['unit2note'] for s in samples])
        unit2note_ = unit2note[..., None].repeat([1, 1, self.num_bins])
        probs = torch.gather(probs, 1, unit2note_)
        batch['probs'] = probs  # [B, T_s, N]
        batch['unit2note'] = unit2note
        bounds = torch.diff(
            unit2note, dim=1, prepend=unit2note.new_zeros((batch['size'], 1))
        ) > 0
        batch['bounds'] = bounds.float()  # [B, T_s]

        return batch


# todo
class MIDIExtractionTask(BaseTask):
    def __init__(self, config: dict):
        super().__init__(config)
        self.midiloss = None
        self.dataset_cls = MIDIExtractionDataset
        self.midi_min = self.config['midi_min']
        self.midi_max = self.config['midi_max']
        self.midi_deviation = self.config['midi_prob_deviation']
        self.rest_threshold = self.config['rest_threshold']
        self.cfg=config

    def build_model(self):

        model = build_object_from_class_name(self.config['model_cls'], nn.Module, config=self.config)

        return model

    def build_losses_and_metrics(self):

        self.midi_loss = nn.BCEWithLogitsLoss()
        self.bound_loss = modules.losses.BinaryEMDLoss()
        # self.bound_loss = modules.losses.BinaryEMDLoss(bidirectional=True)
        self.register_metric('midi_acc', modules.metrics.MIDIAccuracy(tolerance=0.5))

    def run_model(self, sample, infer=False):
        """
        steps:
            1. run the full model
            2. calculate losses if not infer
        """
        spec = sample['units']  # [B, T_ph]
        # target = (sample['probs'],sample['bounds'])  # [B, T_s, M]
        mask = sample['unit2note'] > 0
        # mask=None

        f0 = sample['pitch']




        if infer:
            probs, bounds = self.model(x=spec, f0=f0, mask=mask, sig=True)
            return probs, bounds
        else:
            losses = {}
            probs, bounds = self.model(x=spec, f0=f0, mask=mask, sig=False)

            if  self.cfg['use_bound_loss']:
                bound_loss = self.bound_loss(bounds, sample['bounds'])

                losses['bound_loss'] = bound_loss
            if self.cfg['use_midi_loss']:
                midi_loss = self.midi_loss(probs, sample['probs'])

                losses['midi_loss'] = midi_loss

            return losses

        # raise NotImplementedError()

    def _validation_step(self, sample, batch_idx):
        losses = self.run_model(sample, infer=False)
        if batch_idx < self.config['num_valid_plots']:
            probs, bounds = self.run_model(sample, infer=True)
            unit2note_gt = sample['unit2note']
            masks = unit2note_gt > 0
            probs *= masks[..., None]
            bounds *= masks
            self.plot_prob(batch_idx, sample['probs'], probs)

            unit2note_pred = decode_bounds_to_alignment(bounds) * masks
            midi_pred, rest_pred = decode_gaussian_blurred_probs(
                probs, vmin=self.midi_min, vmax=self.midi_max,
                deviation=self.midi_deviation, threshold=self.rest_threshold
            )
            note_midi_pred, note_dur_pred, note_mask_pred = decode_note_sequence(
                unit2note_pred, midi_pred, ~rest_pred & masks
            )
            note_rest_pred = ~note_mask_pred
            self.plot_boundary(
                batch_idx, bounds_gt=sample['bounds'], bounds_pred=bounds,
                dur_gt=sample['note_dur'], dur_pred=note_dur_pred
            )
            self.plot_final(
                batch_idx, sample['note_midi'], sample['note_dur'], sample['note_rest'],
                note_midi_pred, note_dur_pred, note_rest_pred, sample['pitch']
            )

            midi_pred[rest_pred] = -torch.inf  # rest part is set to -inf
            note_midi_gt = sample['note_midi'].clone()
            note_midi_gt[sample['note_rest']] = -torch.inf
            midi_gt = torch.gather(F.pad(note_midi_gt, [1, 0], value=-torch.inf), 1, unit2note_gt)
            self.plot_midi_curve(
                batch_idx, midi_gt=midi_gt, midi_pred=midi_pred, pitch=sample['pitch']
            )
            self.midi_acc.update(
                midi_pred=midi_pred, rest_pred=rest_pred, midi_gt=midi_gt, rest_gt=midi_gt < 0, mask=masks
            )

        return losses, sample['size']

    ############
    # validation plots
    ############
    def plot_prob(self, batch_idx, probs_gt, probs_pred):
        name = f'prob/{batch_idx}'
        vmin, vmax = 0, 1
        spec_cat = torch.cat([(probs_pred - probs_gt).abs() + vmin, probs_gt, probs_pred], -1)
        self.logger.experiment.add_figure(name, spec_to_figure(spec_cat[0], vmin, vmax), self.global_step)

    def plot_boundary(self, batch_idx, bounds_gt, bounds_pred, dur_gt, dur_pred):
        name = f'boundary/{batch_idx}'
        bounds_gt = bounds_gt[0].cpu().numpy()
        bounds_pred = bounds_pred[0].cpu().numpy()
        dur_gt = dur_gt[0].cpu().numpy()
        dur_pred = dur_pred[0].cpu().numpy()
        self.logger.experiment.add_figure(name, boundary_to_figure(
            bounds_gt, bounds_pred, dur_gt, dur_pred
        ), self.global_step)

    def plot_midi_curve(self, batch_idx, midi_gt, midi_pred, pitch):
        name = f'midi/{batch_idx}'
        midi_gt = midi_gt[0].cpu().numpy()
        midi_pred = midi_pred[0].cpu().numpy()
        pitch = pitch[0].cpu().numpy()
        self.logger.experiment.add_figure(name, curve_to_figure(
            midi_gt, midi_pred, curve_base=pitch, grid=1, base_label='pitch'
        ), self.global_step)

    def plot_final(self, batch_idx, midi_gt, dur_gt, rest_gt, midi_pred, dur_pred, rest_pred, pitch):
        name = f'final/{batch_idx}'
        midi_gt = midi_gt[0].cpu().numpy()
        midi_pred = midi_pred[0].cpu().numpy()
        dur_gt = dur_gt[0].cpu().numpy()
        dur_pred = dur_pred[0].cpu().numpy()
        rest_gt = rest_gt[0].cpu().numpy()
        rest_pred = rest_pred[0].cpu().numpy()
        pitch = pitch[0].cpu().numpy()
        self.logger.experiment.add_figure(name, pitch_notes_to_figure(
            pitch=pitch, note_midi_gt=midi_gt, note_dur_gt=dur_gt, note_rest_gt=rest_gt,
            note_midi_pred=midi_pred, note_dur_pred=dur_pred, note_rest_pred=rest_pred
        ), self.global_step)

training/me_task.py
openvpi/SOME
from __future__ import annotations

import pathlib
import re
import types
from collections import OrderedDict

import numpy as np
import torch

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 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:
                pass
        samples_.append(res)
    return samples_


def filter_kwargs(dict_to_filter, kwarg_obj):
    import inspect

    sig = inspect.signature(kwarg_obj)
    filter_keys = [param.name for param in sig.parameters.values() if param.kind == param.POSITIONAL_OR_KEYWORD]
    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', key_in_ckpt='state_dict',
        strict=True, device='cpu'
):
    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 key_in_ckpt is None:
        state_dict = ckpt_loaded
    else:
        state_dict = ckpt_loaded[key_in_ckpt]
    if prefix_in_ckpt is not None:
        state_dict = OrderedDict({
            k[len(prefix_in_ckpt) + 1:]: v
            for k, v in state_dict.items() if k.startswith(f'{prefix_in_ckpt}.')
        })
    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]


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(f"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 = build_object_from_class_name(
        optimizer_args['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/__init__.py
openvpi/SOME
from typing import Tuple

import librosa
import numpy as np
import parselmouth
import torch

from utils.pitch_utils import interp_f0


def merge_slurs(note_seq: list, note_dur: list, note_slur: list, tolerance=None) -> Tuple[list, list]:
    """
    merge slurs with the similar pitch
    """
    note_midi = [librosa.note_to_midi(n, round_midi=False) if n != 'rest' else 'rest' for n in note_seq]
    prev_min = prev_max = None
    note_midi_merge_slur = [note_midi[0]]
    note_dur_merge_slur = [note_dur[0]]

    def can_be_merged(midi):
        if tolerance is None or midi == 'rest' or note_midi_merge_slur[-1] == 'rest':
            return note_midi_merge_slur[-1] == midi
        return (
            abs(midi - note_midi_merge_slur[-1]) <= tolerance
            and (prev_min is None or abs(midi - prev_min) <= tolerance)
            and (prev_max is None or abs(midi - prev_max) <= tolerance)
        )

    def get_merged_midi(midi1, dur1, midi2, dur2):
        if midi1 == midi2:
            return midi1
        return (midi1 * dur1 + midi2 * dur2) / (dur1 + dur2)

    for i in range(1, len(note_seq)):
        if note_slur[i] and can_be_merged(note_midi[i]):
            # update min and max
            prev_min = min(note_midi[i], note_midi_merge_slur[-1]) if prev_min is None else min(prev_min, note_midi[i])
            prev_max = max(note_midi[i], note_midi_merge_slur[-1]) if prev_max is None else max(prev_max, note_midi[i])
            note_midi_merge_slur[-1] = get_merged_midi(
                note_midi_merge_slur[-1], note_dur_merge_slur[-1], note_midi[i], note_dur[i]
            )
            note_dur_merge_slur[-1] += note_dur[i]
        else:
            note_midi_merge_slur.append(note_midi[i])
            note_dur_merge_slur.append(note_dur[i])
            prev_min = prev_max = None
    note_seq_merge_slur = [
        librosa.midi_to_note(n, cents=True, unicode=False) if n != 'rest' else 'rest' for n in note_midi_merge_slur
    ]
    return note_seq_merge_slur, note_dur_merge_slur


def merge_rests(note_seq: list, note_dur: list) -> Tuple[list, list]:
    i = 0
    note_seq_merge_rest = []
    note_dur_merge_rest = []
    while i < len(note_seq):
        if note_seq[i] != 'rest':
            note_seq_merge_rest.append(note_seq[i])
            note_dur_merge_rest.append(note_dur[i])
            i += 1
        else:
            j = i
            rest_dur = 0
            while j < len(note_seq) and note_seq[j] == 'rest':
                rest_dur += note_dur[j]
                j += 1
            note_seq_merge_rest.append('rest')
            note_dur_merge_rest.append(rest_dur)
            i = j
    return note_seq_merge_rest, note_dur_merge_rest


@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 pad_frames(frames, hop_size, n_samples, n_expect):
    n_frames = frames.shape[0]
    lpad = (int(n_samples // hop_size) - n_frames + 1) // 2
    rpad = n_expect - n_frames - lpad
    if rpad < 0:
        frames = frames[:rpad]
        rpad = 0
    if lpad > 0 or rpad > 0:
        frames = np.pad(frames, (lpad, rpad), mode='constant', constant_values=(frames[0], frames[-1]))
    return frames


def get_pitch_parselmouth(waveform, sample_rate, hop_size, length, interp_uv=False):
    """

    :param waveform: [T]
    :param hop_size: size of each frame
    :param sample_rate: sampling rate of waveform
    :param length: Expected number of frames
    :param interp_uv: Interpolate unvoiced parts
    :return: f0, uv
    """

    time_step = hop_size / sample_rate
    f0_min = 65
    f0_max = 800

    # noinspection PyArgumentList
    f0 = parselmouth.Sound(waveform, sampling_frequency=sample_rate).to_pitch_ac(
        time_step=time_step, voicing_threshold=0.6,
        pitch_floor=f0_min, pitch_ceiling=f0_max
    ).selected_array['frequency'].astype(np.float32)
    f0 = pad_frames(f0, hop_size, waveform.shape[0], length)
    uv = f0 == 0
    if interp_uv:
        f0, uv = interp_f0(f0, uv)
    return f0, uv


class SinusoidalSmoothingConv1d(torch.nn.Conv1d):
    def __init__(self, kernel_size):
        super().__init__(
            in_channels=1,
            out_channels=1,
            kernel_size=kernel_size,
            bias=False,
            padding='same',
            padding_mode='replicate'
        )
        smooth_kernel = torch.sin(torch.from_numpy(
            np.linspace(0, 1, kernel_size).astype(np.float32) * np.pi
        ))
        smooth_kernel /= smooth_kernel.sum()
        self.weight.data = smooth_kernel[None, None]

utils/binarizer_utils.py
openvpi/SOME
from __future__ import annotations

import pathlib

import lightning.pytorch.utilities
import yaml

loaded_config_files = {}


def override_dict(old_config: dict, new_config: dict):
    for k, v in new_config.items():
        if isinstance(v, dict) and k in old_config:
            override_dict(old_config[k], new_config[k])
        else:
            old_config[k] = v


def read_full_config(config_path: pathlib.Path) -> dict:
    config_path = config_path.resolve()
    config_path_str = config_path.as_posix()
    if config_path in loaded_config_files:
        return loaded_config_files[config_path_str]

    with open(config_path, 'r', encoding='utf8') as f:
        config = yaml.safe_load(f)
    if 'base_config' not in config:
        loaded_config_files[config_path_str] = config
        return config

    if not isinstance(config['base_config'], list):
        config['base_config'] = [config['base_config']]
    squashed_config = {}
    for base_config in config['base_config']:
        c_path = pathlib.Path(base_config)
        full_base_config = read_full_config(c_path)
        override_dict(squashed_config, full_base_config)
    override_dict(squashed_config, config)
    squashed_config.pop('base_config')
    loaded_config_files[config_path_str] = squashed_config
    return squashed_config


@lightning.pytorch.utilities.rank_zero.rank_zero_only
def print_config(config: dict):
    for i, (k, v) in enumerate(sorted(config.items())):
        print(f"\033[0;33m{k}\033[0m: {v}", end='')
        if i < len(config) - 1:
            print(", ", end="")
        if i % 5 == 4:
            print()
    print()

utils/config_utils.py
openvpi/SOME
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):
        self.path = pathlib.Path(path) / f'{prefix}.data'
        self.prefix = prefix
        self.dset = None
        self.counter = 0
        self.lock = multiprocessing.Lock()
        if allowed_attr is not None:
            self.allowed_attr = set(allowed_attr)
        else:
            self.allowed_attr = None

    def add_item(self, item):
        if self.dset is None:
            self.dset = h5py.File(self.path, 'w')
        if self.allowed_attr is not None:
            item = {
                k: item[k]
                for k in self.allowed_attr
                if k in item
            }
        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)

    def finalize(self):
        if self.dset is not None:
            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/indexed_datasets.py
openvpi/SOME
from typing import List, Dict

import mido
import numpy as np
import torch
import torch.nn.functional as F


def decode_gaussian_blurred_probs(probs, vmin, vmax, deviation, threshold):
    num_bins = int(probs.shape[-1])
    interval = (vmax - vmin) / (num_bins - 1)
    width = int(3 * deviation / interval)  # 3 * sigma
    idx = torch.arange(num_bins, device=probs.device)[None, None, :]  # [1, 1, N]
    idx_values = idx * interval + vmin
    center = torch.argmax(probs, dim=-1, keepdim=True)  # [B, T, 1]
    start = torch.clip(center - width, min=0)  # [B, T, 1]
    end = torch.clip(center + width + 1, max=num_bins)  # [B, T, 1]
    idx_masks = (idx >= start) & (idx < end)  # [B, T, N]
    weights = probs * idx_masks  # [B, T, N]
    product_sum = torch.sum(weights * idx_values, dim=2)  # [B, T]
    weight_sum = torch.sum(weights, dim=2)  # [B, T]
    values = product_sum / (weight_sum + (weight_sum == 0))  # avoid dividing by zero, [B, T]
    rest = probs.max(dim=-1)[0] < threshold  # [B, T]
    return values, rest


def decode_bounds_to_alignment(bounds, use_diff=True):
    bounds_step = bounds.cumsum(dim=1).round().long()
    if use_diff:
        bounds_inc = torch.diff(
            bounds_step, dim=1, prepend=torch.full(
                (bounds.shape[0], 1), fill_value=-1,
                dtype=bounds_step.dtype, device=bounds_step.device
            )
        ) > 0
    else:
        bounds_inc = F.pad((bounds_step[:, 1:] > bounds_step[:, :-1]), [1, 0], value=True)
    frame2item = bounds_inc.long().cumsum(dim=1)
    return frame2item


def decode_note_sequence(frame2item, values, masks, threshold=0.5):
    """

    :param frame2item: [1, 1, 1, 1, 2, 2, 3, 3, 3]
    :param values:
    :param masks:
    :param threshold: minimum ratio of unmasked frames required to be regarded as an unmasked item
    :return: item_values, item_dur, item_masks
    """
    b = frame2item.shape[0]
    space = frame2item.max() + 1

    item_dur = frame2item.new_zeros(b, space, dtype=frame2item.dtype).scatter_add(
        1, frame2item, torch.ones_like(frame2item)
    )[:, 1:]
    item_unmasked_dur = frame2item.new_zeros(b, space, dtype=frame2item.dtype).scatter_add(
        1, frame2item, masks.long()
    )[:, 1:]
    item_masks = item_unmasked_dur / item_dur >= threshold

    values_quant = values.round().long()
    histogram = frame2item.new_zeros(b, space * 128, dtype=frame2item.dtype).scatter_add(
        1, frame2item * 128 + values_quant, torch.ones_like(frame2item) * masks
    ).unflatten(1, [space, 128])[:, 1:, :]
    item_values_center = histogram.float().argmax(dim=2).to(dtype=values.dtype)
    values_center = torch.gather(F.pad(item_values_center, [1, 0]), 1, frame2item)
    values_near_center = masks & (values >= values_center - 0.5) & (values <= values_center + 0.5)
    item_valid_dur = frame2item.new_zeros(b, space, dtype=frame2item.dtype).scatter_add(
        1, frame2item, values_near_center.long()
    )[:, 1:]
    item_values = values.new_zeros(b, space, dtype=values.dtype).scatter_add(
        1, frame2item, values * values_near_center
    )[:, 1:] / (item_valid_dur + (item_valid_dur == 0))

    return item_values, item_dur, item_masks


def build_midi_file(offsets: List[float], segments: List[Dict[str, np.ndarray]], tempo=120) -> mido.MidiFile:
    midi_file = mido.MidiFile(charset='utf8')
    midi_track = mido.MidiTrack()
    midi_track.append(mido.MetaMessage('set_tempo', tempo=mido.bpm2tempo(tempo), time=0))
    last_time = 0
    offsets = [round(o * tempo * 8) for o in offsets]
    for i, (offset, segment) in enumerate(zip(offsets, segments)):
        note_midi = np.round(segment['note_midi']).astype(np.int64).tolist()
        note_tick = np.diff(np.round(np.cumsum(segment['note_dur']) * tempo * 8).astype(np.int64), prepend=0).tolist()
        note_rest = segment['note_rest'].tolist()
        start = offset
        for j in range(len(note_midi)):
            end = start + note_tick[j]
            if i < len(offsets) - 1 and end > offsets[i + 1]:
                end = offsets[i + 1]
            if start < end and not note_rest[j]:
                midi_track.append(mido.Message('note_on', note=note_midi[j], time=start - last_time))
                midi_track.append(mido.Message('note_off', note=note_midi[j], time=end - start))
                last_time = end
            start = end
    midi_file.tracks.append(midi_track)
    return midi_file


# if __name__ == '__main__':
#     frame2item = torch.LongTensor([
#         [1, 1, 1, 1, 2, 2, 3, 3, 3, 0, 0, 0, 0, 0],
#         [1, 1, 1, 2, 3, 3, 3, 3, 3, 4, 4, 0, 0, 0]
#     ])
#     values = torch.FloatTensor([
#         [60, 61, 60.5, 63, 57, 57, 50, 55, 54, 0, 0, 0, 0, 0],
#         [50, 51, 50.5, 53, 47, 47, 40, 45, 44, 38, 38, 0, 0, 0]
#     ])
#     masks = frame2item > 0
#     decode_note_sequence(frame2item, values, masks)

utils/infer_utils.py
openvpi/SOME
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/multiprocess_utils.py
openvpi/SOME
import numpy as np
import torch

f0_bin = 256
f0_max = 1100.0
f0_min = 50.0
f0_mel_min = 1127 * np.log(1 + f0_min / 700)
f0_mel_max = 1127 * np.log(1 + f0_max / 700)


def f0_to_coarse(f0):
    is_torch = isinstance(f0, torch.Tensor)
    f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)
    f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1

    f0_mel[f0_mel <= 1] = 1
    f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1
    f0_coarse = (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(np.int)
    assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (f0_coarse.max(), f0_coarse.min())
    return f0_coarse


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


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

utils/pitch_utils.py
openvpi/SOME
import math

import matplotlib.pyplot as plt
import numpy as np
import torch
from matplotlib.ticker import MultipleLocator

LINE_COLORS = ['w', 'r', 'y', 'cyan', 'm', 'b', 'lime']


def spec_to_figure(spec, vmin=None, vmax=None):
    if isinstance(spec, torch.Tensor):
        spec = spec.cpu().numpy()
    fig = plt.figure(figsize=(12, 15))
    plt.pcolor(spec.T, vmin=vmin, vmax=vmax)
    plt.tight_layout()
    return fig


def dur_to_figure(dur_gt, dur_pred, txt):
    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]))
    fig.legend()
    fig.tight_layout()
    return fig


def boundary_to_figure(
        bounds_gt: np.ndarray, bounds_pred: np.ndarray,
        dur_gt: np.ndarray = None, dur_pred: np.ndarray = None
):
    fig = plt.figure(figsize=(12, 6))
    bounds_acc_gt = np.cumsum(bounds_gt)
    bounds_acc_pred = np.cumsum(bounds_pred)
    plt.plot(bounds_acc_gt, color='b', label='gt')
    plt.plot(bounds_acc_pred, color='r', label='pred')
    if dur_gt is not None and dur_pred is not None:
        height = math.ceil(max(bounds_acc_gt[-1], bounds_acc_pred[-1]))
        dur_acc_gt = np.cumsum(dur_gt)
        dur_acc_pred = np.cumsum(dur_pred)
        plt.vlines(dur_acc_gt[:-1], 0, height / 2, colors='b', linestyles='--')
        plt.vlines(dur_acc_pred[:-1], height / 2, height, colors='r', linestyles='--')
    plt.gca().yaxis.set_major_locator(MultipleLocator(1))
    plt.grid(axis='y')
    plt.legend()
    plt.tight_layout()
    return fig


def pitch_notes_to_figure(
        pitch, note_midi_gt, note_dur_gt, note_rest_gt,
        note_midi_pred=None, note_dur_pred=None, note_rest_pred=None
):
    fig = plt.figure()

    def draw_notes(note_midi, note_dur, note_rest, color, label):
        note_dur_acc = np.cumsum(note_dur)
        if note_rest is None:
            note_rest = np.zeros_like(note_midi, dtype=np.bool_)
        labeled = False
        for i in range(len(note_midi)):
            if note_rest[i]:
                continue
            x0 = note_dur_acc[i - 1] if i > 0 else 0
            y0 = note_midi[i] - 0.5
            rec = plt.Rectangle(
                xy=(x0, y0),
                width=note_dur[i], height=1,
                edgecolor=color, fill=False,
                linewidth=1.5, label=label if not labeled else None,
                # linestyle='--' if note_rest[i] else '-'
            )
            plt.gca().add_patch(rec)
            plt.fill_between([x0, x0 + note_dur[i]], y0, y0 + 1, color='none', facecolor=color, alpha=0.2)
            labeled = True

    draw_notes(note_midi_gt, note_dur_gt, note_rest_gt, color='b', label='gt')
    draw_notes(note_midi_pred, note_dur_pred, note_rest_pred, color='r', label='pred')
    plt.plot(pitch, color='grey', label='pitch')
    plt.gca().yaxis.set_major_locator(MultipleLocator(1))
    plt.grid(axis='y')
    plt.legend()
    plt.tight_layout()
    return fig


def curve_to_figure(curve_gt, curve_pred=None, curve_base=None, grid=None, base_label='base'):
    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='grey', label=base_label)
    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()
    plt.tight_layout()
    return fig


def distribution_to_figure(title, x_label, y_label, items: list, values: list, zoom=0.8):
    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)
    return fig

utils/plot.py
openvpi/SOME
import numpy as np


# This function is obtained from librosa.
def get_rms(
        y,
        *,
        frame_length=2048,
        hop_length=512,
        pad_mode="constant",
):
    padding = (int(frame_length // 2), int(frame_length // 2))
    y = np.pad(y, padding, mode=pad_mode)

    axis = -1
    # put our new within-frame axis at the end for now
    out_strides = y.strides + tuple([y.strides[axis]])
    # Reduce the shape on the framing axis
    x_shape_trimmed = list(y.shape)
    x_shape_trimmed[axis] -= frame_length - 1
    out_shape = tuple(x_shape_trimmed) + tuple([frame_length])
    xw = np.lib.stride_tricks.as_strided(
        y, shape=out_shape, strides=out_strides
    )
    if axis < 0:
        target_axis = axis - 1
    else:
        target_axis = axis + 1
    xw = np.moveaxis(xw, -1, target_axis)
    # Downsample along the target axis
    slices = [slice(None)] * xw.ndim
    slices[axis] = slice(0, None, hop_length)
    x = xw[tuple(slices)]

    # Calculate power
    power = np.mean(np.abs(x) ** 2, axis=-2, keepdims=True)

    return np.sqrt(power)


class Slicer:
    def __init__(self,
                 sr: int,
                 threshold: float = -40.,
                 min_length: int = 5000,
                 min_interval: int = 300,
                 hop_size: int = 20,
                 max_sil_kept: int = 5000):
        if not min_length >= min_interval >= hop_size:
            raise ValueError('The following condition must be satisfied: min_length >= min_interval >= hop_size')
        if not max_sil_kept >= hop_size:
            raise ValueError('The following condition must be satisfied: max_sil_kept >= hop_size')
        min_interval = sr * min_interval / 1000
        self.sr = sr
        self.threshold = 10 ** (threshold / 20.)
        self.hop_size = round(sr * hop_size / 1000)
        self.win_size = min(round(min_interval), 4 * self.hop_size)
        self.min_length = round(sr * min_length / 1000 / self.hop_size)
        self.min_interval = round(min_interval / self.hop_size)
        self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)

    def _apply_slice(self, waveform, begin, end):
        chunk = {
            'offset': begin * self.hop_size / self.sr
        }
        if len(waveform.shape) > 1:
            chunk['waveform'] = waveform[:, begin * self.hop_size: min(waveform.shape[1], end * self.hop_size)]
        else:
            chunk['waveform'] = waveform[begin * self.hop_size: min(waveform.shape[0], end * self.hop_size)]
        return chunk

    # @timeit
    def slice(self, waveform):
        if len(waveform.shape) > 1:
            samples = waveform.mean(axis=0)
        else:
            samples = waveform
        if (samples.shape[0] + self.hop_size - 1) // self.hop_size <= self.min_length:
            return [{'offset': 0, 'waveform': waveform}]
        rms_list = get_rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
        sil_tags = []
        silence_start = None
        clip_start = 0
        for i, rms in enumerate(rms_list):
            # Keep looping while frame is silent.
            if rms < self.threshold:
                # Record start of silent frames.
                if silence_start is None:
                    silence_start = i
                continue
            # Keep looping while frame is not silent and silence start has not been recorded.
            if silence_start is None:
                continue
            # Clear recorded silence start if interval is not enough or clip is too short
            is_leading_silence = silence_start == 0 and i > self.max_sil_kept
            need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length
            if not is_leading_silence and not need_slice_middle:
                silence_start = None
                continue
            # Need slicing. Record the range of silent frames to be removed.
            if i - silence_start <= self.max_sil_kept:
                pos = rms_list[silence_start: i + 1].argmin() + silence_start
                if silence_start == 0:
                    sil_tags.append((0, pos))
                else:
                    sil_tags.append((pos, pos))
                clip_start = pos
            elif i - silence_start <= self.max_sil_kept * 2:
                pos = rms_list[i - self.max_sil_kept: silence_start + self.max_sil_kept + 1].argmin()
                pos += i - self.max_sil_kept
                pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
                pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
                if silence_start == 0:
                    sil_tags.append((0, pos_r))
                    clip_start = pos_r
                else:
                    sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
                    clip_start = max(pos_r, pos)
            else:
                pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
                pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
                if silence_start == 0:
                    sil_tags.append((0, pos_r))
                else:
                    sil_tags.append((pos_l, pos_r))
                clip_start = pos_r
            silence_start = None
        # Deal with trailing silence.
        total_frames = rms_list.shape[0]
        if silence_start is not None and total_frames - silence_start >= self.min_interval:
            silence_end = min(total_frames, silence_start + self.max_sil_kept)
            pos = rms_list[silence_start: silence_end + 1].argmin() + silence_start
            sil_tags.append((pos, total_frames + 1))
        # Apply and return slices.
        if len(sil_tags) == 0:
            return [{'offset': 0, 'waveform': waveform}]
        else:
            chunks = []
            if sil_tags[0][0] > 0:
                chunks.append(self._apply_slice(waveform, 0, sil_tags[0][0]))
            for i in range(len(sil_tags) - 1):
                chunks.append(self._apply_slice(waveform, sil_tags[i][1], sil_tags[i + 1][0]))
            if sil_tags[-1][1] < total_frames:
                chunks.append(self._apply_slice(waveform, sil_tags[-1][1], total_frames))
            return chunks

utils/slicer2.py
openvpi/SOME
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.pytorch.callbacks import ModelCheckpoint, TQDMProgressBar
from lightning.pytorch.strategies import DDPStrategy
from lightning.pytorch.utilities.rank_zero import rank_zero_info
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data.distributed import Sampler

import utils


# ==========LR schedulers==========

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, 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
        super(WarmupCosineSchedule, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)

    def lr_lambda(self, step):
        if step < self.warmup_steps:
            return step / max(1.0, self.warmup_steps)
        # 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, frame_count_grid=200,
                 required_batch_count_multiple=1, batch_by_size=True, sort_by_similar_size=True,
                 shuffle_sample=False, shuffle_batch=False, seed=0, drop_last=False) -> None:
        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.frame_count_grid = frame_count_grid
        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.shuffle_sample = shuffle_sample
        self.shuffle_batch = shuffle_batch
        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(self.seed + self.epoch)
        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 = self.frame_count_grid
                assert grid > 0
                sizes = (np.round(np.array(self.dataset._sizes)[indices] / grid) * grid).clip(grid, None).astype(
                    np.int64)
                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)))

        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)]

        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 = []
        else:
            leftovers = (rng.permutation(len(batches) - floored_total_batch_count) + floored_total_batch_count).tolist()

        batch_assignment = rng.permuted(
            np.arange(floored_total_batch_count).reshape(-1, self.num_replicas).transpose(), axis=0
        )[self.rank].tolist()
        floored_batch_count = len(batch_assignment)
        ceiled_batch_count = floored_batch_count + (1 if len(leftovers) > 0 else 0)
        if self.rank < len(leftovers):
            batch_assignment.append(leftovers[self.rank])
        elif len(leftovers) > 0:
            batch_assignment.append(batch_assignment[self.epoch % floored_batch_count])
        if self.required_batch_count_multiple > 1 and ceiled_batch_count % self.required_batch_count_multiple != 0:
            # batch_assignment = batch_assignment[:((floored_batch_count \
            # // self.required_batch_count_multiple) * self.required_batch_count_multiple)]
            ceiled_batch_count = math.ceil(
                ceiled_batch_count / 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])

        self.batches = [deepcopy(batches[i]) for i in batch_assignment]

        if self.shuffle_batch:
            rng.shuffle(self.batches)

        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


class DsEvalBatchSampler(Sampler):
    def __init__(self, dataset, max_batch_frames, max_batch_size, rank=None, batch_by_size=True) -> None:
        self.dataset = dataset
        self.max_batch_frames = max_batch_frames
        self.max_batch_size = max_batch_size
        self.rank = rank
        self.batch_by_size = batch_by_size
        self.batches = None
        self.batch_size = max_batch_size
        self.drop_last = False

        if self.rank == 0:
            indices = list(range(len(self.dataset)))
            if self.batch_by_size:
                self.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:
                self.batches = [
                    indices[i:i + self.max_batch_size]
                    for i in range(0, len(indices), self.max_batch_size)
                ]
        else:
            self.batches = [[0]]

    def __iter__(self):
        return iter(self.batches)

    def __len__(self):
        return len(self.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.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.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


def get_strategy(accelerator, devices, num_nodes, strategy, backend):
    if accelerator != 'auto' and accelerator != 'gpu':
        return strategy

    from lightning.fabric.utilities.imports import _IS_INTERACTIVE
    from lightning.pytorch.accelerators import AcceleratorRegistry
    from lightning.pytorch.accelerators.cuda import CUDAAccelerator
    from lightning.pytorch.accelerators.hpu import HPUAccelerator
    from lightning.pytorch.accelerators.ipu import IPUAccelerator
    from lightning.pytorch.accelerators.mps import MPSAccelerator
    from lightning.pytorch.accelerators.tpu import TPUAccelerator
    from lightning.pytorch.utilities.exceptions import MisconfigurationException

    def _choose_auto_accelerator():
        if TPUAccelerator.is_available():
            return "tpu"
        if IPUAccelerator.is_available():
            return "ipu"
        if HPUAccelerator.is_available():
            return "hpu"
        if MPSAccelerator.is_available():
            return "mps"
        if CUDAAccelerator.is_available():
            return "cuda"
        return "cpu"

    def _choose_gpu_accelerator_backend():
        if MPSAccelerator.is_available():
            return "mps"
        if CUDAAccelerator.is_available():
            return "cuda"
        raise MisconfigurationException("No supported gpu backend found!")

    if accelerator == "auto":
        _accelerator_flag = _choose_auto_accelerator()
    elif accelerator == "gpu":
        _accelerator_flag = _choose_gpu_accelerator_backend()
    else:
        return strategy

    if _accelerator_flag != "mps" and _accelerator_flag != "cuda":
        return strategy

    _num_nodes_flag = int(num_nodes) if num_nodes is not None else 1
    _devices_flag = devices

    accelerator = AcceleratorRegistry.get(_accelerator_flag)
    accelerator_cls = accelerator.__class__

    if _devices_flag == "auto":
        _devices_flag = accelerator.auto_device_count()

    _devices_flag = accelerator_cls.parse_devices(_devices_flag)
    _parallel_devices = accelerator_cls.get_parallel_devices(_devices_flag)

    def get_ddp_strategy(_backend):
        if _backend == 'gloo':
            return DDPStrategy(process_group_backend='gloo', find_unused_parameters=False)
        elif _backend == 'nccl' or _backend == 'nccl_no_p2p':
            return DDPStrategy(process_group_backend='nccl', find_unused_parameters=False)
        else:
            raise ValueError(f'backend {_backend} is not valid.')

    if _num_nodes_flag > 1:
        return get_ddp_strategy(backend)
    if len(_parallel_devices) <= 1:
        return strategy
    if len(_parallel_devices) > 1 and _IS_INTERACTIVE:
        return strategy
    return get_ddp_strategy(backend)

utils/training_utils.py
openvpi/SOME
import importlib
import os
import pathlib
import time
from typing import Dict, Tuple

import click
import gradio as gr
import librosa
import yaml

import inference
from inference import BaseInference
from utils.infer_utils import build_midi_file
from utils.slicer2 import Slicer

_work_dir: pathlib.Path = None
_infer_instances: Dict[str, Tuple[BaseInference, dict]] = {}  # dict mapping model_rel_path to (infer_ins, config)


def infer(model_rel_path, input_audio_path, tempo_value):
    if not model_rel_path or not input_audio_path or tempo_value is None:
        return None, "Error: required inputs not specified."
    if model_rel_path not in _infer_instances:
        model_path = _work_dir / model_rel_path
        with open(model_path.with_name('config.yaml'), 'r', encoding='utf8') as f:
            config = yaml.safe_load(f)
        infer_cls = inference.task_inference_mapping[config['task_cls']]

        pkg = ".".join(infer_cls.split(".")[:-1])
        cls_name = infer_cls.split(".")[-1]
        infer_cls = getattr(importlib.import_module(pkg), cls_name)
        assert issubclass(infer_cls, inference.BaseInference), \
            f'Binarizer class {infer_cls} is not a subclass of {inference.BaseInference}.'
        infer_ins = infer_cls(config=config, model_path=model_path)
        print(f"Initialized: {infer_ins}")
        _infer_instances[model_rel_path] = (infer_ins, config)
    else:
        infer_ins, config = _infer_instances[model_rel_path]

    input_audio_path = pathlib.Path(input_audio_path)
    total_duration = librosa.get_duration(filename=input_audio_path)
    if total_duration > 20 * 60:  # 20 minutes
        return None, f"Error: the input audio is too long (>= 20 minutes)."

    try:
        waveform, _ = librosa.load(input_audio_path, sr=config['audio_sample_rate'], mono=True)
    except:
        return None, f"Error: unsupported or corrupt file format: {input_audio_path.name}"

    start_time = time.time()
    slicer = Slicer(sr=config['audio_sample_rate'], max_sil_kept=1000)
    chunks = slicer.slice(waveform)
    midis = infer_ins.infer([c['waveform'] for c in chunks])
    infer_time = time.time() - start_time
    rtf = infer_time / total_duration
    print(f'RTF: {rtf}')

    midi_file = build_midi_file([c['offset'] for c in chunks], midis, tempo=tempo_value)

    output_midi_path = input_audio_path.with_suffix('.mid')
    midi_file.save(output_midi_path)
    os.remove(input_audio_path)

    return output_midi_path, f"Cost {round(infer_time, 2)} s, RTF: {round(rtf, 3)}"


@click.command(help='Launch the web UI for inference')
@click.option('--port', type=int, default=7860, help='Server port')
@click.option('--work_dir', type=str, required=False, help='Directory to read the experiments')
def webui(port, work_dir):
    if work_dir is None:
        work_dir = pathlib.Path(__file__).with_name('experiments')
    else:
        work_dir = pathlib.Path(work_dir)
    assert work_dir.is_dir(), f'{work_dir} is not a directory.'
    global _work_dir
    _work_dir = work_dir
    choices = [
        p.relative_to(work_dir).as_posix()
        for p in work_dir.rglob('*.ckpt')
    ]
    if len(choices) == 0:
        raise FileNotFoundError(f'No checkpoints found in {work_dir}.')
    iface = gr.Interface(
        title="SOME: Singing-Oriented MIDI Extractor",
        description="Submit an audio file and download the extracted MIDI file.",
        theme="default",
        fn=infer,
        inputs=[
            gr.components.Dropdown(
                label="Model Checkpoint", choices=choices, value=choices[0],
                multiselect=False, allow_custom_value=False
            ),
            gr.components.Audio(label="Input Audio File", type="filepath"),
            gr.components.Number(label='Tempo Value', minimum=20, maximum=200, value=120),
        ],
        outputs=[
            gr.components.File(label="Output MIDI File", file_types=['.mid']),
            gr.components.Label(label="Inference Statistics"),
        ]
    )
    iface.queue(concurrency_count=10)
    iface.launch(server_port=port)


if __name__ == "__main__":
    webui()

webui.py
openvpi/SOME