import torch import torch.nn as nn import torch.nn.functional as F import math import numpy as np from typing import Optional, Dict, Any, List from soulxsinger.models.modules.vocoder import Vocoder from soulxsinger.models.modules.decoder import CFMDecoder from soulxsinger.models.modules.convnext import ConvNeXtV2Block from soulxsinger.models.modules.mel_transform import MelSpectrogramEncoder class SoulXSinger(nn.Module): """ SoulXSinger model. """ def __init__(self, config: Dict): super(SoulXSinger, self).__init__() audio_cfg = config.audio enc_cfg = config.model.encoder cfm_cfg = config.model.flow_matching self.note_text_encoder = nn.Embedding(enc_cfg["vocab_size"], enc_cfg["text_dim"]) self.note_pitch_encoder = nn.Embedding(256, enc_cfg["pitch_dim"]) self.note_type_encoder = nn.Embedding(256, enc_cfg["type_dim"]) self.f0_encoder = nn.Embedding(enc_cfg["f0_bin"], enc_cfg["f0_dim"]) self.preflow = nn.Sequential( *[ConvNeXtV2Block(enc_cfg["text_dim"], enc_cfg["text_dim"] * 2) for _ in range(enc_cfg["num_layers"])] ) self.cfm_decoder = CFMDecoder(cfm_cfg) if audio_cfg is None and isinstance(enc_cfg, dict): audio_cfg = enc_cfg.get("audio_config") self.mel = MelSpectrogramEncoder(audio_cfg) self.vocoder = Vocoder() @staticmethod def expand_states(h, mel2token): """ Expand the states to the mel-scale. args: h: states, shape: [B, T, H] mel2token: mel2token, shape: [B, F] returns: h: expanded states, shape: [B, F, H] """ try: assert mel2token.max() <= h.size(1) - 1 except: print(f"Warning: mel2token.max() ({mel2token.max()}) is greater than h.size(1) - 1 ({h.size(1) - 1})") mel2token = torch.clamp(mel2token, 0, h.size(1)-1) mel2token_ = mel2token[..., None].repeat([1, 1, h.shape[-1]]) h = torch.gather(h, 1, mel2token_) # [B, T, H] return h @staticmethod def f0_to_coarse(f0, f0_bin=361, f0_min=32.7031956625, f0_shift=0): """ Convert continuous F0 values to discrete F0 bins (SIL and C1 - B6, 361 bins). args: f0: continuous F0 values f0_bin: number of F0 bins f0_min: minimum F0 value f0_shift: shift value for F0 bins returns: f0_coarse: discrete F0 bins """ is_torch = isinstance(f0, torch.Tensor) uv_mask = f0 <= 0 if is_torch: f0_safe = torch.maximum(f0, torch.tensor(f0_min)) f0_cents = 1200 * torch.log2(f0_safe / f0_min) else: f0_safe = np.maximum(f0, f0_min) f0_cents = 1200 * np.log2(f0_safe / f0_min) f0_coarse = (f0_cents / 20) + 1 if is_torch: f0_coarse = torch.round(f0_coarse).long() f0_coarse = torch.clamp(f0_coarse, min=1, max=f0_bin - 1) else: f0_coarse = np.rint(f0_coarse).astype(int) f0_coarse = np.clip(f0_coarse, 1, f0_bin - 1) f0_coarse[uv_mask] = 0 if f0_shift != 0: if is_torch: voiced = f0_coarse > 0 if voiced.any(): shifted = f0_coarse[voiced] + f0_shift f0_coarse[voiced] = torch.clamp(shifted, 1, f0_bin - 1) else: voiced = f0_coarse > 0 if np.any(voiced): shifted = f0_coarse[voiced] + f0_shift f0_coarse[voiced] = np.clip(shifted, 1, f0_bin - 1) return f0_coarse def infer(self, meta: dict, auto_shift=False, pitch_shift=0, n_steps=32, cfg=3, control="melody"): gt_note_text = meta['target']['phoneme'] gt_mel2note = meta['target']['mel2note'] gt_note_type = meta['target']['note_type'] pt_wav = meta['prompt']['waveform'] pt_note_text = meta['prompt']['phoneme'] pt_mel2note = meta['prompt']['mel2note'] pt_note_type = meta['prompt']['note_type'] if control == "score": gt_note_pitch = meta['target']['note_pitch'] pt_note_pitch = meta['prompt']['note_pitch'] gt_f0 = None pt_f0 = None elif control == "melody": gt_f0 = meta['target']['f0'] pt_f0 = meta['prompt']['f0'] gt_note_pitch = None pt_note_pitch = None else: raise ValueError(f"Unknown control mode: {control}") # calculate auto pitch shift if auto_shift and pitch_shift == 0: if gt_note_pitch != None and pt_note_pitch != None: gt_median = torch.median(gt_note_pitch[gt_note_pitch >= 1]) pt_median = torch.median(pt_note_pitch[pt_note_pitch >= 1]) f0_shift = torch.round(pt_median - gt_median).int().item() elif gt_f0 != None and pt_f0 != None: gt_f0_median = torch.median(gt_f0[gt_f0 > 0]) pt_f0_median = torch.median(pt_f0[pt_f0 > 0]) f0_shift = torch.round(torch.log2(pt_f0_median / gt_f0_median) * 1200 / 100).int().item() else: print("Warning: pitch_shift is True but note_pitch or f0 is None. Set f0_shift to 0.") f0_shift = 0 else: f0_shift = 0 if gt_f0 is None or pt_f0 is None: gt_f0, pt_f0 = torch.zeros_like(gt_mel2note).float().to(gt_mel2note.device), torch.zeros_like(pt_mel2note).float().to(pt_mel2note.device) if gt_note_pitch is None or pt_note_pitch is None: gt_note_pitch, pt_note_pitch = torch.zeros_like(gt_note_type).int().to(gt_note_type.device), torch.zeros_like(pt_note_type).int().to(pt_note_type.device) # convert prompt waveform to mel spectrogram pt_mel = self.mel(pt_wav) len_prompt = pt_note_pitch.shape[1] len_prompt_mel = pt_f0.shape[1] note_pitch = torch.cat([pt_note_pitch, gt_note_pitch], 1) note_text = torch.cat([pt_note_text, gt_note_text], 1) note_type = torch.cat([pt_note_type, gt_note_type], 1) mel2note = torch.cat([pt_mel2note, gt_mel2note + len_prompt], 1) f0_course_pt = self.f0_to_coarse(pt_f0) f0_course_gt = self.f0_to_coarse(gt_f0, f0_shift=f0_shift * 5) f0_course = torch.cat([f0_course_pt, f0_course_gt], 1) note_pitch[note_pitch > 0] = note_pitch[note_pitch > 0] + f0_shift note_pitch = torch.clamp(note_pitch, 0, 255) features = self.note_pitch_encoder(note_pitch) + self.note_type_encoder(note_type) + self.note_text_encoder(note_text) features = self.preflow(features) features = self.expand_states(features, mel2note) features = features + self.f0_encoder(f0_course) gt_decoder_inp = features[:, len_prompt_mel:, :] pt_decoder_inp = features[:, :len_prompt_mel, :] generated_mel = self.cfm_decoder.reverse_diffusion( pt_mel, pt_decoder_inp, gt_decoder_inp, n_timesteps=n_steps, cfg=cfg ) generated_audio = self.vocoder(generated_mel.transpose(1, 2)[0:1, ...]) return generated_audio