voice-changer/server/voice_changer/DiffusionSVC/DiffusionSVCModelSlotGenerator.py
2023-08-09 16:55:59 +09:00

57 lines
2.2 KiB
Python

import os
from data.ModelSlot import DiffusionSVCModelSlot, ModelSlot
from voice_changer.DiffusionSVC.inferencer.diffusion_svc_model.diffusion.unit2mel import load_model_vocoder_from_combo
from voice_changer.VoiceChangerParamsManager import VoiceChangerParamsManager
from voice_changer.utils.LoadModelParams import LoadModelParams
from voice_changer.utils.ModelSlotGenerator import ModelSlotGenerator
def get_divisors(n):
divisors = []
for i in range(1, int(n**0.5)+1):
if n % i == 0:
divisors.append(i)
if i != n // i:
divisors.append(n // i)
return sorted(divisors)
class DiffusionSVCModelSlotGenerator(ModelSlotGenerator):
@classmethod
def loadModel(cls, props: LoadModelParams):
slotInfo: DiffusionSVCModelSlot = DiffusionSVCModelSlot()
for file in props.files:
if file.kind == "diffusionSVCModel":
slotInfo.modelFile = file.name
slotInfo.defaultTune = 0
slotInfo.isONNX = slotInfo.modelFile.endswith(".onnx")
slotInfo.name = os.path.splitext(os.path.basename(slotInfo.modelFile))[0]
# slotInfo.iconFile = "/assets/icons/noimage.png"
slotInfo.embChannels = 768
slotInfo.slotIndex = props.slot
if slotInfo.isONNX:
slotInfo = cls._setInfoByONNX(slotInfo)
else:
slotInfo = cls._setInfoByPytorch(slotInfo)
return slotInfo
@classmethod
def _setInfoByPytorch(cls, slot: DiffusionSVCModelSlot):
vcparams = VoiceChangerParamsManager.get_instance().params
modelPath = os.path.join(vcparams.model_dir, str(slot.slotIndex), os.path.basename(slot.modelFile))
diff_model, diff_args, naive_model, naive_args = load_model_vocoder_from_combo(modelPath, device="cpu")
slot.kStepMax = diff_args.model.k_step_max
slot.nLayers = diff_args.model.n_layers
slot.nnLayers = naive_args.model.n_layers
slot.defaultKstep = slot.kStepMax
divs = get_divisors(slot.defaultKstep)
slot.defaultSpeedup = divs[-2]
slot.speakers = {(x+1): f"user{x+1}" for x in range(diff_args.model.n_spk)}
return slot
@classmethod
def _setInfoByONNX(cls, slot: ModelSlot):
return slot