Update extract_feature_print.py

extract_feature_print.py

@@ -3,6 +3,135 @@ from transformers import HubertModel
 import os
 from torch import nn
 
+import json
+version_config_paths = [
+    os.path.join("/kaggle/working/Mangio-RVC-Fork/configs", "32k.json"),
+    os.path.join("/kaggle/working/Mangio-RVC-Fork/configs", "40k.json"),
+    os.path.join("/kaggle/working/Mangio-RVC-Fork/configs", "48k.json"),
+    os.path.join("/kaggle/working/Mangio-RVC-Fork/configs", "48k_v2.json"),
+    os.path.join("/kaggle/working/Mangio-RVC-Fork/configs", "40k.json"),
+    os.path.join("/kaggle/working/Mangio-RVC-Fork/configs", "32k_v2.json"),
+]
+
+class Config:
+    def __init__(self):
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.is_half = self.device != "cpu"
+        self.gpu_name = (
+            torch.cuda.get_device_name(int(self.device.split(":")[-1]))
+            if self.device.startswith("cuda")
+            else None
+        )
+        self.json_config = self.load_config_json()
+        self.gpu_mem = None
+        self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()
+
+    def load_config_json(self) -> dict:
+        configs = {}
+        for config_file in version_config_paths:
+            config_path = os.path.join("rvc", "configs", config_file)
+            with open(config_path, "r") as f:
+                configs[config_file] = json.load(f)
+        return configs
+
+    def has_mps(self) -> bool:
+        # Check if Metal Performance Shaders are available - for macOS 12.3+.
+        return torch.backends.mps.is_available()
+
+    def has_xpu(self) -> bool:
+        # Check if XPU is available.
+        return hasattr(torch, "xpu") and torch.xpu.is_available()
+
+    def set_precision(self, precision):
+        if precision not in ["fp32", "fp16"]:
+            raise ValueError("Invalid precision type. Must be 'fp32' or 'fp16'.")
+
+        fp16_run_value = precision == "fp16"
+        preprocess_target_version = "3.7" if precision == "fp16" else "3.0"
+        preprocess_path = os.path.join(
+            os.path.dirname(__file__),
+            os.pardir,
+            "rvc",
+            "train",
+            "preprocess",
+            "preprocess.py",
+        )
+
+        for config_path in version_config_paths:
+            full_config_path = os.path.join("rvc", "configs", config_path)
+            try:
+                with open(full_config_path, "r") as f:
+                    config = json.load(f)
+                config["train"]["fp16_run"] = fp16_run_value
+                with open(full_config_path, "w") as f:
+                    json.dump(config, f, indent=4)
+            except FileNotFoundError:
+                print(f"File not found: {full_config_path}")
+
+        if os.path.exists(preprocess_path):
+            with open(preprocess_path, "r") as f:
+                preprocess_content = f.read()
+            preprocess_content = preprocess_content.replace(
+                "3.0" if precision == "fp16" else "3.7", preprocess_target_version
+            )
+            with open(preprocess_path, "w") as f:
+                f.write(preprocess_content)
+
+        return f"Overwritten preprocess and config.json to use {precision}."
+
+    def get_precision(self):
+        if not version_config_paths:
+            raise FileNotFoundError("No configuration paths provided.")
+
+        full_config_path = os.path.join("rvc", "configs", version_config_paths[0])
+        try:
+            with open(full_config_path, "r") as f:
+                config = json.load(f)
+            fp16_run_value = config["train"].get("fp16_run", False)
+            precision = "fp16" if fp16_run_value else "fp32"
+            return precision
+        except FileNotFoundError:
+            print(f"File not found: {full_config_path}")
+            return None
+
+    def device_config(self) -> tuple:
+        if self.device.startswith("cuda"):
+            self.set_cuda_config()
+        elif self.has_mps():
+            self.device = "mps"
+            self.is_half = False
+            self.set_precision("fp32")
+        else:
+            self.device = "cpu"
+            self.is_half = False
+            self.set_precision("fp32")
+
+        # Configuration for 6GB GPU memory
+        x_pad, x_query, x_center, x_max = (
+            (3, 10, 60, 65) if self.is_half else (1, 6, 38, 41)
+        )
+        if self.gpu_mem is not None and self.gpu_mem <= 4:
+            # Configuration for 5GB GPU memory
+            x_pad, x_query, x_center, x_max = (1, 5, 30, 32)
+
+        return x_pad, x_query, x_center, x_max
+
+    def set_cuda_config(self):
+        i_device = int(self.device.split(":")[-1])
+        self.gpu_name = torch.cuda.get_device_name(i_device)
+        low_end_gpus = ["16", "P40", "P10", "1060", "1070", "1080"]
+        if (
+            any(gpu in self.gpu_name for gpu in low_end_gpus)
+            and "V100" not in self.gpu_name.upper()
+        ):
+            self.is_half = False
+            self.set_precision("fp32")
+
+        self.gpu_mem = torch.cuda.get_device_properties(i_device).total_memory // (
+            1024**3
+        )
+Config()
+
 #HuggingFacePlaceHolder = None
 class HubertModelWithFinalProj(HubertModel):
     def __init__(self, config):
@@ -89,7 +218,8 @@ models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
 if Custom_Embed == False:
     model = models[0]
 else:
-    model = HubertModelWithFinalProj.from_pretrained("/kaggle/working/Mangio-RVC-Fork/Custom/")
+    dtype = torch.float16 if config.is_half and "cuda" in device else torch.float32
+    model = HubertModelWithFinalProj.from_pretrained("/kaggle/working/Mangio-RVC-Fork/Custom/").to(dtype).to(device)
 model = model.to(device)
 printt("move model to %s" % device)
 if device not in ["mps", "cpu"]:
@@ -97,7 +227,7 @@ if device not in ["mps", "cpu"]:
 model.eval()
 
 todo = sorted(list(os.listdir(wavPath)))[i_part::n_part]
-n = max(1, len(todo) // 10)  # 最多打印十条
+n = max(1, len(todo) // 10) 
 if len(todo) == 0:
     printt("no-feature-todo")
 else:
@@ -121,10 +251,16 @@ else:
                     "output_layer": 9 if version == "v1" else 12,  # layer 9
                 }
                 with torch.no_grad():
-                    logits = model.extract_features(**inputs)
-                    feats = (
-                        model.final_proj(logits[0]) if version == "v1" else logits[0]
-                    )
+                    if Custom_Embed == False:
+                        logits = model.extract_features(**inputs)
+                        feats = (
+                            model.final_proj(logits[0]) if version == "v1" else logits[0]
+                        )
+                    elif Custom_Embed == True:
+                        feats = model(feats)["last_hidden_state"]
+                        feats = (
+                            model.final_proj(feats[0]).unsqueeze(0) if version == "v1" else feats
+                        )
 
                 feats = feats.squeeze(0).float().cpu().numpy()
                 if np.isnan(feats).sum() == 0: