model persistent

.gitignore

@@ -44,7 +44,7 @@ server/pretrain/
 server/weights/
 server/weights_/
 server/weights__/
-
+server/models/
 start_trainer.sh
 
 # venv

server/MMVCServerSIO.py

@@ -41,6 +41,7 @@ def setupArgParser():
         default=True,
         help="generate self-signed certificate",
     )
+    parser.add_argument("--model_dir", type=str, help="path to model files")
 
     parser.add_argument(
         "--content_vec_500", type=str, help="path to content_vec_500 model(pytorch)"
@@ -212,6 +213,7 @@ def localServer():
 if __name__ == "MMVCServerSIO":
     mp.freeze_support()
     voiceChangerParams = VoiceChangerParams(
+        model_dir=args.model_dir,
         content_vec_500=args.content_vec_500,
         content_vec_500_onnx=args.content_vec_500_onnx,
         content_vec_500_onnx_on=args.content_vec_500_onnx_on,

server/voice_changer/RVC/ModelSlotGenerator.py

@@ -4,9 +4,10 @@ from voice_changer.RVC.ModelSlot import ModelSlot
 import torch
 import onnxruntime
 import json
+import os
 
 
-def generateModelSlot(params):
+def generateModelSlot_(params):
     modelSlot = ModelSlot()
 
     modelSlot.modelFile = params["files"]["rvcModel"]
@@ -28,6 +29,41 @@ def generateModelSlot(params):
     return modelSlot
 
 
+def generateModelSlot(slotDir: str):
+    modelSlot = ModelSlot()
+    if os.path.exists(slotDir) == False:
+        return modelSlot
+    paramFile = os.path.join(slotDir, "params.json")
+    with open(paramFile, "r") as f:
+        params = json.load(f)
+
+    modelSlot.modelFile = os.path.join(
+        slotDir, os.path.basename(params["files"]["rvcModel"])
+    )
+    if "rvcFeature" in params["files"]:
+        modelSlot.featureFile = os.path.join(
+            slotDir, os.path.basename(params["files"]["rvcFeature"])
+        )
+    else:
+        modelSlot.featureFile = None
+    if "rvcIndex" in params["files"]:
+        modelSlot.indexFile = os.path.join(
+            slotDir, os.path.basename(params["files"]["rvcIndex"])
+        )
+    else:
+        modelSlot.indexFile = None
+
+    modelSlot.defaultTrans = params["trans"] if "trans" in params else 0
+
+    modelSlot.isONNX = modelSlot.modelFile.endswith(".onnx")
+
+    if modelSlot.isONNX:
+        _setInfoByONNX(modelSlot)
+    else:
+        _setInfoByPytorch(modelSlot)
+    return modelSlot
+
+
 def _setInfoByPytorch(slot: ModelSlot):
     cpt = torch.load(slot.modelFile, map_location="cpu")
     config_len = len(cpt["config"])

server/voice_changer/RVC/RVC.py

@@ -36,7 +36,8 @@ from voice_changer.RVC.pipeline.Pipeline import Pipeline
 
 from Exceptions import NoModeLoadedException
 from const import UPLOAD_DIR
-
+import shutil
+import json
 
 providers = [
     "OpenVINOExecutionProvider",
@@ -45,6 +46,9 @@ providers = [
     "CPUExecutionProvider",
 ]
 
+RVC_MODEL_DIRNAME = "rvc"
+RVC_MAX_SLOT_NUM = 6
+
 
 class RVC:
     initialLoad: bool = True
@@ -66,30 +70,64 @@ class RVC:
         )
         self.params = params
         EmbedderManager.initialize(params)
+        self.loadSlots()
         print("RVC initialization: ", params)
 
     def loadModel(self, props: LoadModelParams):
         target_slot_idx = props.slot
         params = props.params
 
-        modelSlot = generateModelSlot(params)
+        # modelName = os.path.splitext(os.path.basename(params["files"]["rvcModel"]))[0]
-        self.settings.modelSlots[target_slot_idx] = modelSlot
+        slotDir = os.path.join(
-        print(
+            self.params.model_dir, RVC_MODEL_DIRNAME, str(target_slot_idx)
-            f"[Voice Changer] RVC new model is uploaded,{target_slot_idx}",
-            asdict(modelSlot),
         )
+        files = [params["files"]["rvcModel"]]
+        if "rvcFeature" in params["files"]:
+            files.append(params["files"]["rvcFeature"])
+        if "rvcIndex" in params["files"]:
+            files.append(params["files"]["rvcIndex"])
 
-        # 初回のみロード
+        os.makedirs(slotDir, exist_ok=True)
-        if self.initialLoad:
+        for f in files:
-            self.prepareModel(target_slot_idx)
+            dst = os.path.join(slotDir, os.path.basename(f))
-            self.settings.modelSlotIndex = target_slot_idx
+            if os.path.exists(dst):
-            self.switchModel()
+                os.remove(dst)
-            self.initialLoad = False
+            shutil.move(f, dst)
-        elif target_slot_idx == self.currentSlot:
+        json.dump(params, open(os.path.join(slotDir, "params.json"), "w"))
-            self.prepareModel(target_slot_idx)
+        self.loadSlots()
-
         return self.get_info()
 
+    def loadSlots(self):
+        dirname = os.path.join(self.params.model_dir, RVC_MODEL_DIRNAME)
+        self.settings.modelSlots = []
+        if not os.path.exists(dirname):
+            return
+
+        for slot_idx in range(RVC_MAX_SLOT_NUM):
+            slotDir = os.path.join(
+                self.params.model_dir, RVC_MODEL_DIRNAME, str(slot_idx)
+            )
+            modelSlot = generateModelSlot(slotDir)
+            self.settings.modelSlots.append(modelSlot)
+
+        # modelSlot = generateModelSlot(params)
+        # self.settings.modelSlots[target_slot_idx] = modelSlot
+        # print(
+        #     f"[Voice Changer] RVC new model is uploaded,{target_slot_idx}",
+        #     asdict(modelSlot),
+        # )
+
+        # # 初回のみロード
+        # if self.initialLoad:
+        #     self.prepareModel(target_slot_idx)
+        #     self.settings.modelSlotIndex = target_slot_idx
+        #     self.switchModel()
+        #     self.initialLoad = False
+        # elif target_slot_idx == self.currentSlot:
+        #     self.prepareModel(target_slot_idx)
+
+        # return self.get_info()
+
     def update_settings(self, key: str, val: int | float | str):
         if key in self.settings.intData:
             # 設定前処理

server/voice_changer/RVC/RVCSettings.py

@@ -16,7 +16,14 @@ class RVCSettings:
 
     framework: str = "PyTorch"  # PyTorch or ONNX
     modelSlots: list[ModelSlot] = field(
-        default_factory=lambda: [ModelSlot(), ModelSlot(), ModelSlot(), ModelSlot()]
+        default_factory=lambda: [
+            ModelSlot(),  # 1
+            ModelSlot(),  # 2
+            ModelSlot(),  # 3
+            ModelSlot(),  # 4
+            ModelSlot(),  # 5
+            ModelSlot(),  # 6(merged)
+        ]
     )
     indexRatio: float = 0
     rvcQuality: int = 0

server/voice_changer/RVC/pipeline/PipelineGenerator.py

@@ -68,6 +68,7 @@ def _loadIndex(modelSlot: ModelSlot):
     print("[Voice Changer] Loading index...")
     # ファイル指定がない場合はNone
     if modelSlot.featureFile is None or modelSlot.indexFile is None:
+        print("[Voice Changer] Index is None, not used")
         return None, None
 
     # ファイル指定があってもファイルがない場合はNone

server/voice_changer/utils/VoiceChangerParams.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass
 
 @dataclass
 class VoiceChangerParams:
+    model_dir: str
     content_vec_500: str
     content_vec_500_onnx: str
     content_vec_500_onnx_on: bool