Clicko777/RVC_HFv2

1import torch2import numpy as np3from tqdm import tqdm4import json5 6 7def load_data(file_name: str = "./lib/uvr5_pack/name_params.json") -> dict:8    with open(file_name, "r") as f:9        data = json.load(f)10 11    return data12 13 14def make_padding(width, cropsize, offset):15    left = offset16    roi_size = cropsize - left * 217    if roi_size == 0:18        roi_size = cropsize19    right = roi_size - (width % roi_size) + left20 21    return left, right, roi_size22 23 24def inference(X_spec, device, model, aggressiveness, data):25    """26    data ： dic configs27    """28 29    def _execute(30        X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half=True31    ):32        model.eval()33        with torch.no_grad():34            preds = []35 36            iterations = [n_window]37 38            total_iterations = sum(iterations)39            for i in tqdm(range(n_window)):40                start = i * roi_size41                X_mag_window = X_mag_pad[42                    None, :, :, start : start + data["window_size"]43                ]44                X_mag_window = torch.from_numpy(X_mag_window)45                if is_half:46                    X_mag_window = X_mag_window.half()47                X_mag_window = X_mag_window.to(device)48 49                pred = model.predict(X_mag_window, aggressiveness)50 51                pred = pred.detach().cpu().numpy()52                preds.append(pred[0])53 54            pred = np.concatenate(preds, axis=2)55        return pred56 57    def preprocess(X_spec):58        X_mag = np.abs(X_spec)59        X_phase = np.angle(X_spec)60 61        return X_mag, X_phase62 63    X_mag, X_phase = preprocess(X_spec)64 65    coef = X_mag.max()66    X_mag_pre = X_mag / coef67 68    n_frame = X_mag_pre.shape[2]69    pad_l, pad_r, roi_size = make_padding(n_frame, data["window_size"], model.offset)70    n_window = int(np.ceil(n_frame / roi_size))71 72    X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")73 74    if list(model.state_dict().values())[0].dtype == torch.float16:75        is_half = True76    else:77        is_half = False78    pred = _execute(79        X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half80    )81    pred = pred[:, :, :n_frame]82 83    if data["tta"]:84        pad_l += roi_size // 285        pad_r += roi_size // 286        n_window += 187 88        X_mag_pad = np.pad(X_mag_pre, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")89 90        pred_tta = _execute(91            X_mag_pad, roi_size, n_window, device, model, aggressiveness, is_half92        )93        pred_tta = pred_tta[:, :, roi_size // 2 :]94        pred_tta = pred_tta[:, :, :n_frame]95 96        return (pred + pred_tta) * 0.5 * coef, X_mag, np.exp(1.0j * X_phase)97    else:98        return pred * coef, X_mag, np.exp(1.0j * X_phase)99 100 101def _get_name_params(model_path, model_hash):102    data = load_data()103    flag = False104    ModelName = model_path105    for type in list(data):106        for model in list(data[type][0]):107            for i in range(len(data[type][0][model])):108                if str(data[type][0][model][i]["hash_name"]) == model_hash:109                    flag = True110                elif str(data[type][0][model][i]["hash_name"]) in ModelName:111                    flag = True112 113                if flag:114                    model_params_auto = data[type][0][model][i]["model_params"]115                    param_name_auto = data[type][0][model][i]["param_name"]116                    if type == "equivalent":117                        return param_name_auto, model_params_auto118                    else:119                        flag = False120    return param_name_auto, model_params_auto121