Add audio format and extraction precision options to preprocessing and extraction scripts

core.py

@@ -417,6 +417,7 @@ def run_preprocess_script(
     chunk_len: float,
     overlap_len: float,
     normalization_mode: str = "none",
+    audio_format: str = "wav",
 ):
     preprocess_script_path = os.path.join("rvc", "train", "preprocess", "preprocess.py")
     command = [
@@ -436,6 +437,7 @@ def run_preprocess_script(
                 chunk_len,
                 overlap_len,
                 normalization_mode,
+                audio_format,
             ],
         ),
     ]
@@ -456,6 +458,8 @@ def run_extract_script(
     embedder_model: str,
     embedder_model_custom: str = None,
     include_mutes: int = 2,
+    extract_precision: str = "fp32",
+    remove_sliced_16k: bool = False,
 ):
     model_path = os.path.join(logs_path, model_name)
     extract = os.path.join("rvc", "train", "extract", "extract.py")
@@ -474,6 +478,8 @@ def run_extract_script(
                 embedder_model,
                 embedder_model_custom,
                 include_mutes,
+                extract_precision,
+                remove_sliced_16k,
             ],
         ),
     ]
@@ -1055,6 +1061,15 @@ def tts(**kwargs):
     default="none",
     help="Normalization mode.",
 )
+@click.option(
+    "--audio-format",
+    type=click.Choice(["wav", "flac"]),
+    default="wav",
+    help=(
+        "Format of the sliced audio files. 'flac' saves ~30% disk, but quantizes to 24-bit "
+        "and clips peaks above 0 dBFS - avoid it with --normalization-mode none."
+    ),
+)
 def preprocess(**kwargs):
     """Preprocess a dataset for training."""
     kwargs["sample_rate"] = int(kwargs["sample_rate"])
@@ -1074,6 +1089,7 @@ def preprocess(**kwargs):
         chunk_len=kwargs["chunk_len"],
         overlap_len=kwargs["overlap_len"],
         normalization_mode=kwargs["normalization_mode"],
+        audio_format=kwargs["audio_format"],
     )
     click.echo(result)
 
@@ -1124,6 +1140,24 @@ def preprocess(**kwargs):
     default=2,
     help="Number of silent files to include.",
 )
+@click.option(
+    "--extract-precision",
+    type=click.Choice(["fp32", "fp16"]),
+    default="fp32",
+    help=(
+        "Format of the extracted files. 'fp16' halves disk usage: coarse f0 as uint8 is "
+        "lossless, features as float16 lose precision - negligible, but irreversible."
+    ),
+)
+@click.option(
+    "--remove-sliced-16k",
+    is_flag=True,
+    default=False,
+    help=(
+        "Delete the sliced_audios_16k folder after extraction. Training does not read those "
+        "files, but re-extracting later will require preprocessing again."
+    ),
+)
 def extract(**kwargs):
     """Extract features from a preprocessed dataset."""
     kwargs["sample_rate"] = int(kwargs["sample_rate"])
@@ -1137,6 +1171,8 @@ def extract(**kwargs):
         embedder_model=kwargs["embedder_model"],
         embedder_model_custom=kwargs["embedder_model_custom"],
         include_mutes=kwargs["include_mutes"],
+        extract_precision=kwargs["extract_precision"],
+        remove_sliced_16k=kwargs["remove_sliced_16k"],
     )
     click.echo(result)
 

rvc/train/data_utils.py

@@ -95,17 +95,18 @@ def get_labels(self, phone, pitch, pitchf):
             pitch (str): Path to pitch label file.
             pitchf (str): Path to pitchf label file.
         """
-        phone = np.load(phone)
+        # features may be float16 and coarse pitch uint8 on disk
+        phone = np.load(phone).astype(np.float32)
         phone = np.repeat(phone, 2, axis=0)
-        pitch = np.load(pitch)
-        pitchf = np.load(pitchf)
+        pitch = np.load(pitch).astype(np.int64)
+        pitchf = np.load(pitchf).astype(np.float32)
         n_num = min(phone.shape[0], 900)
         phone = phone[:n_num, :]
         pitch = pitch[:n_num]
         pitchf = pitchf[:n_num]
-        phone = torch.FloatTensor(phone)
-        pitch = torch.LongTensor(pitch)
-        pitchf = torch.FloatTensor(pitchf)
+        phone = torch.from_numpy(phone)
+        pitch = torch.from_numpy(pitch)
+        pitchf = torch.from_numpy(pitchf)
         return phone, pitch, pitchf
 
     def get_audio(self, filename):
@@ -122,7 +123,7 @@ def get_audio(self, filename):
             )
         audio_norm = audio
         audio_norm = audio_norm.unsqueeze(0)
-        spec_filename = filename.replace(".wav", ".spec.pt")
+        spec_filename = os.path.splitext(filename)[0] + ".spec.pt"
         if os.path.exists(spec_filename):
             try:
                 spec = torch.load(spec_filename, weights_only=True)

rvc/train/extract/extract.py

@@ -3,6 +3,7 @@
 import json
 import multiprocessing as mp
 import os
+import shutil
 import sys
 import time
 
@@ -26,7 +27,7 @@
 
 
 class FeatureInput:
-    def __init__(self, f0_method="rmvpe", device="cpu"):
+    def __init__(self, f0_method="rmvpe", device="cpu", quantize=False):
         self.hop_size = 160  # default
         self.sample_rate = 16000  # default
         self.f0_bin = 256
@@ -35,6 +36,7 @@ def __init__(self, f0_method="rmvpe", device="cpu"):
         self.f0_mel_min = 1127 * np.log(1 + self.f0_min / 700)
         self.f0_mel_max = 1127 * np.log(1 + self.f0_max / 700)
         self.device = device
+        self.quantize = quantize
         if f0_method in ("crepe", "crepe-tiny"):
             self.model = CREPE(
                 device=self.device, sample_rate=self.sample_rate, hop_size=self.hop_size
@@ -91,22 +93,25 @@ def process_file(self, file_info):
             feature_pit = self.compute_f0(np_arr)
             np.save(opt_path_full, feature_pit, allow_pickle=False)
             coarse_pit = self.coarse_f0(feature_pit)
+            if self.quantize:
+                # coarse pitch fits in [1, 255], so uint8 is lossless
+                coarse_pit = coarse_pit.astype(np.uint8)
             np.save(opt_path_coarse, coarse_pit, allow_pickle=False)
         except Exception as error:
             print(
                 f"An error occurred extracting file {inp_path} on {self.device}: {error}"
             )
 
 
-def process_files(files, f0_method, device, threads):
-    fe = FeatureInput(f0_method=f0_method, device=device)
+def process_files(files, f0_method, device, threads, quantize=False):
+    fe = FeatureInput(f0_method=f0_method, device=device, quantize=quantize)
     with tqdm.tqdm(total=len(files), leave=True) as pbar:
         for file_info in files:
             fe.process_file(file_info)
             pbar.update(1)
 
 
-def run_pitch_extraction(files, devices, f0_method, threads):
+def run_pitch_extraction(files, devices, f0_method, threads, quantize=False):
     devices_str = ", ".join(devices)
     print(f"Starting pitch extraction on {devices_str} using {f0_method}...")
     start_time = time.time()
@@ -119,6 +124,7 @@ def run_pitch_extraction(files, devices, f0_method, threads):
                 f0_method,
                 devices[i],
                 threads // len(devices),
+                quantize,
             )
             for i in range(len(devices))
         ]
@@ -128,7 +134,13 @@ def run_pitch_extraction(files, devices, f0_method, threads):
 
 
 def process_file_embedding(
-    files, embedder_model, embedder_model_custom, device_num, device, n_threads
+    files,
+    embedder_model,
+    embedder_model_custom,
+    device_num,
+    device,
+    n_threads,
+    half_precision=False,
 ):
     model = load_embedding(embedder_model, embedder_model_custom).to(device).float()
     model.eval()
@@ -143,6 +155,8 @@ def worker(file_info):
         with torch.no_grad():
             result = model(feats)["last_hidden_state"]
         feats_out = result.squeeze(0).float().cpu().numpy()
+        if half_precision:
+            feats_out = feats_out.astype(np.float16)
         if not np.isnan(feats_out).any():
             np.save(out_file_path, feats_out, allow_pickle=False)
         else:
@@ -156,7 +170,7 @@ def worker(file_info):
 
 
 def run_embedding_extraction(
-    files, devices, embedder_model, embedder_model_custom, threads
+    files, devices, embedder_model, embedder_model_custom, threads, half_precision=False
 ):
     devices_str = ", ".join(devices)
     print(
@@ -173,6 +187,7 @@ def run_embedding_extraction(
                 i,
                 devices[i],
                 threads // len(devices),
+                half_precision,
             )
             for i in range(len(devices))
         ]
@@ -190,6 +205,14 @@ def run_embedding_extraction(
     embedder_model = sys.argv[6]
     embedder_model_custom = sys.argv[7] if len(sys.argv) > 7 else None
     include_mutes = int(sys.argv[8]) if len(sys.argv) > 8 else 2
+    # fp16: features as float16, coarse f0 as uint8
+    precision = sys.argv[9].lower() if len(sys.argv) > 9 else "fp32"
+    remove_sliced_16k = (
+        sys.argv[10].lower() in ("yes", "true", "t", "y", "1")
+        if len(sys.argv) > 10
+        else False
+    )
+    half_precision = precision == "fp16"
 
     wav_path = os.path.join(exp_dir, "sliced_audios_16k")
 
@@ -217,13 +240,17 @@ def run_embedding_extraction(
         json.dump(data, f, indent=4)
 
     files = []
-    for file in glob.glob(os.path.join(wav_path, "*.wav")):
+    audio_files = sorted(
+        glob.glob(os.path.join(wav_path, "*.wav"))
+        + glob.glob(os.path.join(wav_path, "*.flac"))
+    )
+    for file in audio_files:
         file_name = os.path.basename(file)
         file_info = [
             file,
             os.path.join(exp_dir, "f0", file_name + ".npy"),
             os.path.join(exp_dir, "f0_voiced", file_name + ".npy"),
-            os.path.join(exp_dir, "extracted", file_name.replace("wav", "npy")),
+            os.path.join(exp_dir, "extracted", os.path.splitext(file_name)[0] + ".npy"),
         ]
         files.append(file_info)
 
@@ -235,11 +262,23 @@ def run_embedding_extraction(
 
     devices = ["cpu"] if gpus == "-" else [f"cuda:{idx}" for idx in gpus.split("-")]
 
-    run_pitch_extraction(files, devices, f0_method, num_processes)
+    run_pitch_extraction(files, devices, f0_method, num_processes, half_precision)
 
     run_embedding_extraction(
-        files, devices, embedder_model, embedder_model_custom, num_processes
+        files,
+        devices,
+        embedder_model,
+        embedder_model_custom,
+        num_processes,
+        half_precision,
     )
 
     generate_config(sample_rate, exp_dir)
     generate_filelist(exp_dir, sample_rate, include_mutes)
+
+    if remove_sliced_16k:
+        try:
+            shutil.rmtree(wav_path)
+            print(f"Removed {wav_path} to save disk space.")
+        except OSError as error:
+            print(f"Could not remove {wav_path}: {error}")

rvc/train/extract/preparing_files.py

@@ -23,12 +23,33 @@ def generate_filelist(model_path: str, sample_rate: int, include_mutes: int = 2)
     f0_dir = os.path.join(model_path, "f0")
     f0nsf_dir = os.path.join(model_path, "f0_voiced")
 
-    gt_wavs_files = set(name.split(".")[0] for name in os.listdir(gt_wavs_dir))
-    feature_files = set(name.split(".")[0] for name in os.listdir(feature_dir))
-
-    f0_files = set(name.split(".")[0] for name in os.listdir(f0_dir))
-    f0nsf_files = set(name.split(".")[0] for name in os.listdir(f0nsf_dir))
-    names = gt_wavs_files & feature_files & f0_files & f0nsf_files
+    # slices may be .wav or .flac, and the f0 files inherit that extension
+    gt_wavs_files = {
+        name.split(".")[0]: name
+        for name in os.listdir(gt_wavs_dir)
+        if name.lower().endswith((".wav", ".flac"))
+    }
+    feature_files = {
+        name.split(".")[0]: name
+        for name in os.listdir(feature_dir)
+        if name.lower().endswith(".npy")
+    }
+    f0_files = {
+        name.split(".")[0]: name
+        for name in os.listdir(f0_dir)
+        if name.lower().endswith(".npy")
+    }
+    f0nsf_files = {
+        name.split(".")[0]: name
+        for name in os.listdir(f0nsf_dir)
+        if name.lower().endswith(".npy")
+    }
+    names = (
+        gt_wavs_files.keys()
+        & feature_files.keys()
+        & f0_files.keys()
+        & f0nsf_files.keys()
+    )
 
     try:
         model_info_path = os.path.join(model_path, "model_info.json")
@@ -53,10 +74,10 @@ def generate_filelist(model_path: str, sample_rate: int, include_mutes: int = 2)
             sids.append(sid)
 
         # Calculate relative pathing
-        rel_wav = os.path.relpath(f"{os.path.join(gt_wavs_dir, name)}.wav")
-        rel_feat = os.path.relpath(f"{os.path.join(feature_dir, name)}.npy")
-        rel_f0 = os.path.relpath(f"{os.path.join(f0_dir, name)}.wav.npy")
-        rel_f0nsf = os.path.relpath(f"{os.path.join(f0nsf_dir, name)}.wav.npy")
+        rel_wav = os.path.relpath(os.path.join(gt_wavs_dir, gt_wavs_files[name]))
+        rel_feat = os.path.relpath(os.path.join(feature_dir, feature_files[name]))
+        rel_f0 = os.path.relpath(os.path.join(f0_dir, f0_files[name]))
+        rel_f0nsf = os.path.relpath(os.path.join(f0nsf_dir, f0nsf_files[name]))
 
         options.append(
             f"{rel_wav}|{rel_feat}|{rel_f0}|{rel_f0nsf}|{sid}".replace("\\", "/")