29 lines
1021 B
Python
29 lines
1021 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
from soulxsinger.models.modules.flow_matching import FlowMatchingTransformer
|
|
|
|
|
|
class CFMDecoder(nn.Module):
|
|
def __init__(self, config):
|
|
super(CFMDecoder, self).__init__()
|
|
self.model = FlowMatchingTransformer(cfg=config, **config)
|
|
|
|
def forward(self, mel, x_mask, decoder_inp, is_prompt):
|
|
outputs = self.model(mel, x_mask, decoder_inp, is_prompt)
|
|
|
|
noise, x, flow_pred, final_mask, prompt_len = outputs["output"]
|
|
return noise, x, flow_pred, final_mask, prompt_len
|
|
|
|
def reverse_diffusion(self, pt_mel, pt_decoder_inp, gt_decoder_inp, n_timesteps=32, cfg=1):
|
|
diffusion_cond = torch.cat([pt_decoder_inp, gt_decoder_inp], dim=1)
|
|
diffusion_cond_emb = self.model.cond_emb(diffusion_cond)
|
|
diffusion_prompt = pt_mel
|
|
|
|
generated = self.model.reverse_diffusion(
|
|
diffusion_cond_emb,
|
|
diffusion_prompt,
|
|
n_timesteps=n_timesteps,
|
|
cfg=cfg
|
|
)
|
|
return generated |