64 lines
2.5 KiB
Python
64 lines
2.5 KiB
Python
# Copyright (c) 2023 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import torch
|
|
from diffusers import DDIMScheduler, DDPMScheduler, PNDMScheduler
|
|
|
|
from models.svc.base import SVCInference
|
|
from models.svc.diffusion.diffusion_inference_pipeline import DiffusionInferencePipeline
|
|
from models.svc.diffusion.diffusion_wrapper import DiffusionWrapper
|
|
from modules.encoder.condition_encoder import ConditionEncoder
|
|
|
|
|
|
class DiffusionInference(SVCInference):
|
|
def __init__(self, args=None, cfg=None, infer_type="from_dataset"):
|
|
SVCInference.__init__(self, args, cfg, infer_type)
|
|
|
|
settings = {
|
|
**cfg.model.diffusion.scheduler_settings,
|
|
**cfg.inference.diffusion.scheduler_settings,
|
|
}
|
|
settings.pop("num_inference_timesteps")
|
|
|
|
if cfg.inference.diffusion.scheduler.lower() == "ddpm":
|
|
self.scheduler = DDPMScheduler(**settings)
|
|
self.logger.info("Using DDPM scheduler.")
|
|
elif cfg.inference.diffusion.scheduler.lower() == "ddim":
|
|
self.scheduler = DDIMScheduler(**settings)
|
|
self.logger.info("Using DDIM scheduler.")
|
|
elif cfg.inference.diffusion.scheduler.lower() == "pndm":
|
|
self.scheduler = PNDMScheduler(**settings)
|
|
self.logger.info("Using PNDM scheduler.")
|
|
else:
|
|
raise NotImplementedError(
|
|
"Unsupported scheduler type: {}".format(
|
|
cfg.inference.diffusion.scheduler.lower()
|
|
)
|
|
)
|
|
|
|
self.pipeline = DiffusionInferencePipeline(
|
|
self.model[1],
|
|
self.scheduler,
|
|
args.diffusion_inference_steps,
|
|
)
|
|
|
|
def _build_model(self):
|
|
self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min
|
|
self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max
|
|
self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder)
|
|
self.acoustic_mapper = DiffusionWrapper(self.cfg)
|
|
model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper])
|
|
return model
|
|
|
|
def _inference_each_batch(self, batch_data):
|
|
device = self.accelerator.device
|
|
for k, v in batch_data.items():
|
|
batch_data[k] = v.to(device)
|
|
|
|
conditioner = self.model[0](batch_data)
|
|
noise = torch.randn_like(batch_data["mel"], device=device)
|
|
y_pred = self.pipeline(noise, conditioner)
|
|
return y_pred
|