Spaces:
Sleeping
Sleeping
# 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 | |