forked from JarodMica/rvc-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathinfer.py
More file actions
156 lines (134 loc) · 5.41 KB
/
infer.py
File metadata and controls
156 lines (134 loc) · 5.41 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
import os
from glob import glob
import soundfile as sf
from scipy.io import wavfile
from rvc_python.modules.vc.modules import VC
from rvc_python.configs.config import Config
from rvc_python.download_model import download_rvc_models
class RVCInference:
def __init__(self, models_dir="rvc_models",
device="cpu:0",
f0method = "rmvpe",
f0up_key = 0,
index_rate = 0.5,
filter_radius = 3,
resample_sr = 48000,
rms_mix_rate = 1,
protect = 0.33
):
self.models_dir = models_dir
self.device = device
self.lib_dir = os.path.dirname(os.path.abspath(__file__))
self.config = Config(self.lib_dir, self.device)
self.vc = VC(self.lib_dir, self.config)
self.current_model = None
self.models = self._load_available_models()
# Default parameters
self.f0method = f0method
self.f0up_key = f0up_key
self.index_rate = index_rate
self.filter_radius = filter_radius
self.resample_sr = resample_sr
self.rms_mix_rate = rms_mix_rate
self.protect = protect
# Download Models
download_rvc_models(self.lib_dir)
def _load_available_models(self):
"""Loads a list of available models from the directory."""
models = {}
for model_dir in glob(os.path.join(self.models_dir, "*")):
if os.path.isdir(model_dir):
model_name = os.path.basename(model_dir)
pth_file = glob(os.path.join(model_dir, "*.pth"))
index_file = glob(os.path.join(model_dir, "*.index"))
if pth_file:
models[model_name] = {
"pth": pth_file[0],
"index": index_file[0] if index_file else None
}
return models
def set_models_dir(self, new_models_dir):
if not os.path.isdir(new_models_dir):
raise ValueError(f"Directory {new_models_dir} does not exist")
self.models_dir = new_models_dir
self.models = self._load_available_models()
def list_models(self):
"""Returns a list of available models."""
return list(self.models.keys())
def load_model(self, model_name, version="v2"):
"""Loads a model into memory."""
if model_name not in self.models:
raise ValueError(f"Model {model_name} not found.")
model_path = self.models[model_name]["pth"]
self.vc.get_vc(model_path, version)
self.current_model = model_name
print(f"Model {model_name} loaded.")
def unload_model(self):
"""Unloads the current model from memory."""
if self.current_model:
self.vc = VC(self.lib_dir, self.config)
self.current_model = None
print("Model unloaded from memory.")
else:
print("No model loaded.")
def set_params(self, **kwargs):
"""Sets parameters for generation."""
valid_params = [
"index_rate", "filter_radius", "resample_sr",
"rms_mix_rate", "protect", "f0up_key", "f0method"
]
for key, value in kwargs.items():
if key in valid_params:
setattr(self, key, value)
else:
print(f"Warning: parameter {key} not recognized and will be ignored.")
def infer_file(self, input_path, output_path):
"""Processes a single file."""
if not self.current_model:
raise ValueError("Please load a model first.")
wav_opt = self.vc.vc_single(
sid=0,
input_audio_path=input_path,
f0_up_key=self.f0up_key,
f0_method=self.f0method,
file_index=self.models[self.current_model].get("index", ""),
index_rate=self.index_rate,
filter_radius=self.filter_radius,
resample_sr=self.resample_sr,
rms_mix_rate=self.rms_mix_rate,
protect=self.protect,
f0_file="",
file_index2=""
)
if isinstance(wav_opt, tuple):
err, _ = wav_opt
raise Exception(err)
wavfile.write(output_path, self.vc.tgt_sr, wav_opt)
return output_path
def infer_dir(self, input_dir, output_dir):
"""Processes all files in a directory."""
if not self.current_model:
raise ValueError("Please load a model first.")
os.makedirs(output_dir, exist_ok=True)
audio_files = glob(os.path.join(input_dir, '*.*'))
processed_files = []
for input_audio_path in audio_files:
output_filename = os.path.splitext(os.path.basename(input_audio_path))[0] + '.wav'
output_path = os.path.join(output_dir, output_filename)
self.infer_file(input_audio_path, output_path)
processed_files.append(output_path)
return processed_files
def set_device(self, device):
"""Sets the device for computations."""
self.device = device
self.config.device = device
self.vc.device = device
# Usage example:
if __name__ == "__main__":
rvc = RVCInference(device="cuda:0")
print("Available models:", rvc.list_models())
rvc.load_model("example_model")
rvc.set_params(f0up_key=2, protect=0.5)
rvc.infer_file("input.wav", "output.wav")
rvc.infer_dir("input_dir", "output_dir")
rvc.unload_model()