diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..0a4d38c2fbc28be317f91e87807f00086eeff07b --- /dev/null +++ b/.gitattributes @@ -0,0 +1,47 @@ +*.jpg filter=lfs diff=lfs merge=lfs -text +*.gif filter=lfs diff=lfs merge=lfs -text +*.mp4 filter=lfs diff=lfs merge=lfs -text +*.avi filter=lfs diff=lfs merge=lfs -text +*.dylib filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.svg filter=lfs diff=lfs merge=lfs -text +*.wav filter=lfs diff=lfs merge=lfs -text +*.m4a filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.mkv filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.docx filter=lfs diff=lfs merge=lfs -text +*.ppt filter=lfs diff=lfs merge=lfs -text +*.pptx filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.jpeg filter=lfs diff=lfs merge=lfs -text +*.ico filter=lfs diff=lfs merge=lfs -text +*.flac filter=lfs diff=lfs merge=lfs -text +*.ogg filter=lfs diff=lfs merge=lfs -text +*.xls filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.webm filter=lfs diff=lfs merge=lfs -text +*.doc filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.png filter=lfs diff=lfs merge=lfs -text +*.mp3 filter=lfs diff=lfs merge=lfs -text +*.7z filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.pdf filter=lfs diff=lfs merge=lfs -text +*.xlsx filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +*.dll filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.bmp filter=lfs diff=lfs merge=lfs -text +*.mov filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.so filter=lfs diff=lfs merge=lfs -text +*.hdf5 filter=lfs diff=lfs merge=lfs -text +*.ttf filter=lfs diff=lfs merge=lfs -text +src/YingMusicSinger/utils/f5_tts/g2p/sources/chinese_lexicon.txt filter=lfs diff=lfs merge=lfs -text diff --git a/src/YingMusicSinger/infer/YingMusicSinger.py b/src/YingMusicSinger/infer/YingMusicSinger.py new file mode 100644 index 0000000000000000000000000000000000000000..e63189229efcf9ca116485888619933a88218207 --- /dev/null +++ b/src/YingMusicSinger/infer/YingMusicSinger.py @@ -0,0 +1,263 @@ +import hydra +import torch +import torch.nn as nn +import torchaudio +from einops import rearrange +from ema_pytorch import EMA +from huggingface_hub import PyTorchModelHubMixin +from omegaconf import OmegaConf + +from src.YingMusicSinger.melody.midi_extractor import MIDIExtractor +from src.YingMusicSinger.models.model import Singer +from src.YingMusicSinger.utils.cnen_tokenizer import CNENTokenizer +from src.YingMusicSinger.utils.lrc_align import ( + align_lrc_put_to_front, + align_lrc_sentence_level, +) +from src.YingMusicSinger.utils.mel_spectrogram import MelodySpectrogram +from src.YingMusicSinger.utils.stable_audio_tools.vae_copysyn import StableAudioInfer + + +class YingMusicSinger(nn.Module, PyTorchModelHubMixin): + def __init__( + self, + model_cfg_path, + ckpt_path=None, + vae_config_path=None, + vae_ckpt_path=None, + midi_teacher_ckpt_path=None, + is_distilled=False, + use_ema=True, + ): + super().__init__() + self.cfg = OmegaConf.load(model_cfg_path) + model_cls = hydra.utils.get_class( + f"src.YingMusicSinger.models.{self.cfg.model.backbone}" + ) + self.melody_input_source = self.cfg.model.melody_input_source + self.is_tts_pretrain = self.cfg.model.is_tts_pretrain + + self.model = Singer( + transformer=model_cls( + **self.cfg.model.arch, + text_num_embeds=self.cfg.datasets_cfg.text_num_embeds, + mel_dim=self.cfg.model.mel_spec.n_mel_channels, + use_guidance_scale_embed=is_distilled, + ), + mel_spec_kwargs=self.cfg.model.mel_spec, + is_tts_pretrain=self.is_tts_pretrain, + melody_input_source=self.melody_input_source, + cka_disabled=self.cfg.model.cka_disabled, + num_channels=None, + extra_parameters=self.cfg.extra_parameters, + distill_stage=1, + use_guidance_scale_embed=is_distilled, + ) + + self.vae = StableAudioInfer( + model_config_path=vae_config_path, + model_ckpt_path=vae_ckpt_path, + ) + + self._need_midi = self.melody_input_source in { + "some_pretrain", + "some_pretrain_fuzzdisturb", + "some_pretrain_postprocess_embedding", + } + self.midi_teacher = None + if self._need_midi: + self.midi_teacher = MIDIExtractor() + if midi_teacher_ckpt_path is not None: + self.midi_teacher._load_form_ckpt(midi_teacher_ckpt_path) + for p in self.midi_teacher.parameters(): + p.requires_grad = False + + self.melody_spectrogram_extract = MelodySpectrogram() + + self.vae_frame_rate = 44100 / 2048 + + if ckpt_path is not None: + ckpt = torch.load(ckpt_path, map_location="cpu") + if use_ema: + ema_model = EMA(self.model, include_online_model=False) + ema_model.load_state_dict(ckpt["ema_model_state_dict"]) + + self.model = ema_model.ema_model + else: + self.model.load_state_dict(ckpt["model_state_dict"]) + + self.cnen_tokenizer = CNENTokenizer() + + @property + def device(self): + return next(self.parameters()).device + + def prepare_input( + self, + ref_audio_path, + melody_audio_path, + ref_text, + target_text, + sil_len_to_end, + lrc_align_mode, + ): + ref_audio, ref_audio_sr = torchaudio.load(ref_audio_path) + silence = torch.zeros(ref_audio.shape[0], int(ref_audio_sr * sil_len_to_end)) + ref_wav = torch.cat([ref_audio, silence], dim=1) + ref_latent = self.vae.encode_audio(ref_wav, in_sr=ref_audio_sr).transpose( + 1, 2 + ) # [B, T, D] + + melody_wav, melody_sr = torchaudio.load(melody_audio_path) + melody_latent = self.vae.encode_audio(melody_wav, in_sr=melody_sr).transpose( + 1, 2 + ) # [B, T, D] + + midi_in = torch.cat([ref_latent, melody_latent], dim=1) + if self.is_tts_pretrain: + midi_in = torch.zeros_like(midi_in) + + ref_latent_len = ref_latent.shape[1] + total_len = int(ref_latent.shape[1] + melody_latent.shape[1]) + + if self._need_midi: + ref_mel = self.melody_spectrogram_extract(audio=ref_wav, sr=ref_audio_sr) + melody_mel = self.melody_spectrogram_extract(audio=melody_wav, sr=melody_sr) + melody_mel_spec = torch.cat([ref_mel, melody_mel], dim=2) + else: + raise NotImplementedError() + + assert isinstance(ref_text, str) and isinstance(target_text, str) + text_list = [ref_text] + [target_text] + + if lrc_align_mode == "put_to_front": + lrc_token, _ = align_lrc_put_to_front( + tokenizer=self.cnen_tokenizer, + lrc_start_times=None, + lrc_lines=text_list, + total_lens=total_len, + ) + elif lrc_align_mode == "sentence_level": + lrc_token, _ = align_lrc_sentence_level( + tokenizer=self.cnen_tokenizer, + lrc_start_times=[0.0, ref_latent_len / self.vae_frame_rate], + lrc_lines=text_list, + total_lens=total_len, + vae_frame_rate=self.vae_frame_rate, + ) + else: + raise ValueError(f"Unsupported lrc_align_mode: {lrc_align_mode}") + + text_tokens = ( + torch.tensor(lrc_token, dtype=torch.int64).unsqueeze(0).to(self.device) + ) + + midi_p, bound_p = None, None + if self._need_midi: + with torch.no_grad(): + midi_p, bound_p = self.midi_teacher(melody_mel_spec.transpose(1, 2)) + + return ( + ref_latent, + ref_latent_len, + text_tokens, + total_len, + midi_in, + midi_p, + bound_p, + ) + + def forward( + self, + ref_audio_path, + melody_audio_path, + ref_text, + target_text, + lrc_align_mode: str = "sentence_level", + sil_len_to_end: float = 0.5, + t_shift: float = 0.5, + nfe_step: int = 32, + cfg_strength: float = 3.0, + seed: int = 666, + is_tts_pretrain: bool = False, + ): + """ + Args: + ref_audio_path: Path to the reference audio (for timbre) + melody_audio_path: Path to the melody reference audio (provides target duration and melody information) + ref_text: Text corresponding to the reference audio + target_text: Target text to be synthesized + lrc_align_mode: Lyric alignment mode "sentence_level" | "put_to_front" + sil_len_to_end: Duration of silence appended to the end of the reference audio (seconds) + t_shift: Sampling time offset + nfe_step: ODE sampling steps + cfg_strength: CFG strength + seed: Random seed + is_tts_pretrain: If True, melody is not provided (TTS mode) + """ + ref_latent, ref_latent_len, text_tokens, total_len, midi_in, midi_p, bound_p = ( + self.prepare_input( + ref_audio_path=ref_audio_path, + melody_audio_path=melody_audio_path, + ref_text=ref_text, + target_text=target_text, + sil_len_to_end=sil_len_to_end, + lrc_align_mode=lrc_align_mode, + ) + ) + + assert midi_p is not None and bound_p is not None + with torch.inference_mode(): + generated_latent, _ = self.model.sample( + cond=ref_latent, + midi_in=midi_in, + text=text_tokens, + duration=total_len, + steps=nfe_step, + cfg_strength=cfg_strength, + sway_sampling_coef=None, + use_epss=False, + seed=seed, + midi_p=midi_p, + t_shift=t_shift, + bound_p=bound_p, + guidance_scale=cfg_strength, + ) + generated_latent = generated_latent.to(torch.float32) + generated_latent = generated_latent[:, ref_latent_len:, :] + generated_latent = generated_latent.permute(0, 2, 1) # [B, D, T] + + generated_audio = self.vae.decode_audio(generated_latent) + audio = rearrange(generated_audio, "b d n -> d (b n)") + + audio = audio.to(torch.float32).cpu() + + return audio, 44100 + + +if __name__ == "__main__": + # === Export to HuggingFace safetensors (optional) === + # model = YingMusicSinger( + # model_cfg_path="src/YingMusicSinger/config/YingMusic_Singer.yaml", + # ckpt_path="ckpts/YingMusicSinger_model.pt", + # vae_config_path="src/YingMusicSinger/config/stable_audio_2_0_vae_20hz_official.json", + # vae_ckpt_path="ckpts/stable_audio_2_0_vae_20hz_official.ckpt", + # midi_teacher_ckpt_path="ckpts/model_ckpt_steps_100000_simplified.ckpt", + # ) + # model.save_pretrained("path/to/save") + + # === Inference Example === + model = YingMusicSinger.from_pretrained("ASLP-lab/YingMusic-Singer") + model.to("cuda:0") + model.eval() + + waveform, sample_rate = model( + ref_audio_path="path/to/ref_audio", # Timbre reference audio + melody_audio_path="path/to/melody_audio", # Melody-providing singing clip + ref_text="oh the reason i hold on", # Lyrics corresponding to ref_audio + target_text="oldest book broken watch|bare feet in grassy spot", # Modified target lyrics + seed=42, + ) + + torchaudio.save("output.wav", waveform, sample_rate=sample_rate) + print("Saved to output.wav") diff --git a/src/YingMusicSinger/melody/Gconform.py b/src/YingMusicSinger/melody/Gconform.py new file mode 100644 index 0000000000000000000000000000000000000000..9ca184ea96b6b871933507fffc66ad3295df766e --- /dev/null +++ b/src/YingMusicSinger/melody/Gconform.py @@ -0,0 +1,298 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + + +class GLU(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + out, gate = x.chunk(2, dim=self.dim) + + return out * gate.sigmoid() + + +class conform_conv(nn.Module): + def __init__( + self, channels: int, kernel_size: int = 31, DropoutL=0.1, bias: bool = True + ): + super().__init__() + self.act2 = nn.SiLU() + self.act1 = GLU(1) + + self.pointwise_conv1 = nn.Conv1d( + channels, 2 * channels, kernel_size=1, stride=1, padding=0, bias=bias + ) + + # self.lorder is used to distinguish if it's a causal convolution, + # if self.lorder > 0: + # it's a causal convolution, the input will be padded with + # `self.lorder` frames on the left in forward (causal conv impl). + # else: it's a symmetrical convolution + + assert (kernel_size - 1) % 2 == 0 + padding = (kernel_size - 1) // 2 + + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=padding, + groups=channels, + bias=bias, + ) + + self.norm = nn.BatchNorm1d(channels) + + self.pointwise_conv2 = nn.Conv1d( + channels, channels, kernel_size=1, stride=1, padding=0, bias=bias + ) + self.drop = nn.Dropout(DropoutL) if DropoutL > 0.0 else nn.Identity() + + def forward(self, x): + x = x.transpose(1, 2) + x = self.act1(self.pointwise_conv1(x)) + x = self.depthwise_conv(x) + x = self.norm(x) + x = self.act2(x) + x = self.pointwise_conv2(x) + return self.drop(x).transpose(1, 2) + + +class Attention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32, conditiondim=None): + super().__init__() + if conditiondim is None: + conditiondim = dim + + self.scale = dim_head**-0.5 + self.heads = heads + hidden_dim = dim_head * heads + self.to_q = nn.Linear(dim, hidden_dim, bias=False) + self.to_kv = nn.Linear(conditiondim, hidden_dim * 2, bias=False) + + self.to_out = nn.Sequential( + nn.Linear( + hidden_dim, + dim, + ), + ) + + def forward(self, q, kv=None, mask=None): + # b, c, h, w = x.shape + if kv is None: + kv = q + # q, kv = map( + # lambda t: rearrange(t, "b c t -> b t c", ), (q, kv) + # ) + + q = self.to_q(q) + k, v = self.to_kv(kv).chunk(2, dim=2) + + q, k, v = map( + lambda t: rearrange(t, "b t (h c) -> b h t c", h=self.heads), (q, k, v) + ) + + if mask is not None: + mask = mask.unsqueeze(1).unsqueeze(1) + + with torch.backends.cuda.sdp_kernel(enable_math=False): + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + + out = rearrange( + out, + "b h t c -> b t (h c) ", + h=self.heads, + ) + return self.to_out(out) + + +class conform_ffn(nn.Module): + def __init__(self, dim, DropoutL1: float = 0.1, DropoutL2: float = 0.1): + super().__init__() + self.ln1 = nn.Linear(dim, dim * 4) + self.ln2 = nn.Linear(dim * 4, dim) + self.drop1 = nn.Dropout(DropoutL1) if DropoutL1 > 0.0 else nn.Identity() + self.drop2 = nn.Dropout(DropoutL2) if DropoutL2 > 0.0 else nn.Identity() + self.act = nn.SiLU() + + def forward(self, x): + x = self.ln1(x) + x = self.act(x) + x = self.drop1(x) + x = self.ln2(x) + return self.drop2(x) + + +class conform_blocke(nn.Module): + def __init__( + self, + dim: int, + kernel_size: int = 31, + conv_drop: float = 0.1, + ffn_latent_drop: float = 0.1, + ffn_out_drop: float = 0.1, + attention_drop: float = 0.1, + attention_heads: int = 4, + attention_heads_dim: int = 64, + ): + super().__init__() + self.ffn1 = conform_ffn(dim, ffn_latent_drop, ffn_out_drop) + self.ffn2 = conform_ffn(dim, ffn_latent_drop, ffn_out_drop) + self.att = Attention(dim, heads=attention_heads, dim_head=attention_heads_dim) + self.attdrop = ( + nn.Dropout(attention_drop) if attention_drop > 0.0 else nn.Identity() + ) + self.conv = conform_conv( + dim, + kernel_size=kernel_size, + DropoutL=conv_drop, + ) + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + self.norm3 = nn.LayerNorm(dim) + self.norm4 = nn.LayerNorm(dim) + self.norm5 = nn.LayerNorm(dim) + + def forward( + self, + x, + mask=None, + ): + x = self.ffn1(self.norm1(x)) * 0.5 + x + + x = self.attdrop(self.att(self.norm2(x), mask=mask)) + x + x = self.conv(self.norm3(x)) + x + x = self.ffn2(self.norm4(x)) * 0.5 + x + return self.norm5(x) + + # return x + + +class Gcf(nn.Module): + def __init__( + self, + dim: int, + kernel_size: int = 31, + conv_drop: float = 0.1, + ffn_latent_drop: float = 0.1, + ffn_out_drop: float = 0.1, + attention_drop: float = 0.1, + attention_heads: int = 4, + attention_heads_dim: int = 64, + ): + super().__init__() + self.att1 = conform_blocke( + dim=dim, + kernel_size=kernel_size, + conv_drop=conv_drop, + ffn_latent_drop=ffn_latent_drop, + ffn_out_drop=ffn_out_drop, + attention_drop=attention_drop, + attention_heads=attention_heads, + attention_heads_dim=attention_heads_dim, + ) + self.att2 = conform_blocke( + dim=dim, + kernel_size=kernel_size, + conv_drop=conv_drop, + ffn_latent_drop=ffn_latent_drop, + ffn_out_drop=ffn_out_drop, + attention_drop=attention_drop, + attention_heads=attention_heads, + attention_heads_dim=attention_heads_dim, + ) + self.glu1 = nn.Sequential(nn.Linear(dim, dim * 2), GLU(2)) + self.glu2 = nn.Sequential(nn.Linear(dim, dim * 2), GLU(2)) + + def forward(self, midi, bound): + midi = self.att1(midi) + bound = self.att2(bound) + midis = self.glu1(midi) + bounds = self.glu2(bound) + return midi + bounds, bound + midis + + +class Gmidi_conform(nn.Module): + def __init__( + self, + lay: int, + dim: int, + indim: int, + outdim: int, + use_lay_skip: bool, + kernel_size: int = 31, + conv_drop: float = 0.1, + ffn_latent_drop: float = 0.1, + ffn_out_drop: float = 0.1, + attention_drop: float = 0.1, + attention_heads: int = 4, + attention_heads_dim: int = 64, + ): + super().__init__() + + self.inln = nn.Linear(indim, dim) + self.inln1 = nn.Linear(indim, dim) + self.outln = nn.Linear(dim, outdim) + self.cutheard = nn.Linear(dim, 1) + # self.cutheard = nn.Linear(dim, outdim) + self.lay = lay + self.use_lay_skip = use_lay_skip + self.cf_lay = nn.ModuleList( + [ + Gcf( + dim=dim, + kernel_size=kernel_size, + conv_drop=conv_drop, + ffn_latent_drop=ffn_latent_drop, + ffn_out_drop=ffn_out_drop, + attention_drop=attention_drop, + attention_heads=attention_heads, + attention_heads_dim=attention_heads_dim, + ) + for _ in range(lay) + ] + ) + self.att1 = conform_blocke( + dim=dim, + kernel_size=kernel_size, + conv_drop=conv_drop, + ffn_latent_drop=ffn_latent_drop, + ffn_out_drop=ffn_out_drop, + attention_drop=attention_drop, + attention_heads=attention_heads, + attention_heads_dim=attention_heads_dim, + ) + self.att2 = conform_blocke( + dim=dim, + kernel_size=kernel_size, + conv_drop=conv_drop, + ffn_latent_drop=ffn_latent_drop, + ffn_out_drop=ffn_out_drop, + attention_drop=attention_drop, + attention_heads=attention_heads, + attention_heads_dim=attention_heads_dim, + ) + + def forward(self, x, mask=None): + x1 = x.clone() + + x = self.inln(x) + x1 = self.inln1(x1) + if mask is not None: + x = x.masked_fill(~mask.unsqueeze(-1), 0) + for idx, i in enumerate(self.cf_lay): + x, x1 = i(x, x1) + + if mask is not None: + x = x.masked_fill(~mask.unsqueeze(-1), 0) + x, x1 = self.att1(x), self.att2(x1) + + cutprp = self.cutheard(x1) + midiout = self.outln(x) + + return midiout, cutprp diff --git a/src/YingMusicSinger/melody/Gconv.py b/src/YingMusicSinger/melody/Gconv.py new file mode 100644 index 0000000000000000000000000000000000000000..b14f48c351ac6d7007a161c7dab5bd8e58fbaa78 --- /dev/null +++ b/src/YingMusicSinger/melody/Gconv.py @@ -0,0 +1,60 @@ +import torch.nn as nn + + +class GLU(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + out, gate = x.chunk(2, dim=self.dim) + + return out * gate.sigmoid() + + +class conform_conv(nn.Module): + def __init__( + self, channels: int, kernel_size: int = 31, DropoutL=0.1, bias: bool = True + ): + super().__init__() + self.act2 = nn.SiLU() + self.act1 = GLU(1) + + self.pointwise_conv1 = nn.Conv1d( + channels, 2 * channels, kernel_size=1, stride=1, padding=0, bias=bias + ) + + # self.lorder is used to distinguish if it's a causal convolution, + # if self.lorder > 0: + # it's a causal convolution, the input will be padded with + # `self.lorder` frames on the left in forward (causal conv impl). + # else: it's a symmetrical convolution + + assert (kernel_size - 1) % 2 == 0 + padding = (kernel_size - 1) // 2 + + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=padding, + groups=channels, + bias=bias, + ) + + self.norm = nn.BatchNorm1d(channels) + + self.pointwise_conv2 = nn.Conv1d( + channels, channels, kernel_size=1, stride=1, padding=0, bias=bias + ) + self.drop = nn.Dropout(DropoutL) if DropoutL > 0.0 else nn.Identity() + + def forward(self, x): + x = x.transpose(1, 2) + x = self.act1(self.pointwise_conv1(x)) + x = self.depthwise_conv(x) + x = self.norm(x) + x = self.act2(x) + x = self.pointwise_conv2(x) + return self.drop(x).transpose(1, 2) diff --git a/src/YingMusicSinger/melody/SmoothMelody.py b/src/YingMusicSinger/melody/SmoothMelody.py new file mode 100644 index 0000000000000000000000000000000000000000..fd7ebb593c66086a8cac4d08acf09f0a00b33d5b --- /dev/null +++ b/src/YingMusicSinger/melody/SmoothMelody.py @@ -0,0 +1,144 @@ +import torch +import torch.nn as nn + + +class MIDIFuzzDisturb(nn.Module): + """Applies fuzzing perturbations to MIDI latent representations. + + The raw MIDI teacher model output preserves good prosody but causes + pronunciation interference. This module mitigates that by applying + blur, temporal dropout, and noise to the melody latent. + """ + + def __init__( + self, dim=128, drop_prob=0.3, noise_scale=0.1, blur_kernel=3, drop_type="random" + ): + super().__init__() + self.blur = None + self.drop_prob = None + self.noise_scale = None + self.dim = dim + self.drop_type = drop_type + + assert drop_prob is not None + assert drop_type is not None + if drop_type == "random": + # drop_prob is a float + if drop_prob != 0: + self.drop_prob = drop_prob + elif drop_type == "equal_space": + # drop_prob is a [drop, keep] list, e.g., [1, 1] means 1 frame drop, 1 frame keep + self.drop_prob = drop_prob + else: + raise ValueError(f"Unknown drop_type: {drop_type}") + + if noise_scale != 0: + self.noise_scale = noise_scale + if blur_kernel != 0: + assert blur_kernel % 2 == 1, f"blur_kernel {blur_kernel} must be odd" + self.blur = nn.AvgPool1d( + kernel_size=blur_kernel, stride=1, padding=blur_kernel // 2 + ) + + def _create_equal_space_mask(self, batch_size, seq_len, device): + """Create an equally-spaced mask cycling [drop, keep] frames.""" + drop_frames, keep_frames = self.drop_prob + cycle_len = drop_frames + keep_frames + + # Pattern: first drop_frames are 0 (drop), next keep_frames are 1 (keep) + pattern = torch.cat( + [ + torch.zeros(drop_frames, device=device), + torch.ones(keep_frames, device=device), + ] + ) + + # Repeat pattern to cover the full sequence length + num_repeats = (seq_len + cycle_len - 1) // cycle_len + mask = pattern.repeat(num_repeats)[:seq_len] # [T] + + # Expand to [B, T, 1] + mask = mask.view(1, seq_len, 1).expand(batch_size, -1, -1) + + return mask + + def forward(self, x): + # x: [B, T, D=128], pre-sigmoid logits + x = torch.sigmoid(x) + + assert x.shape[-1] == self.dim, ( + f"MIDIFuzzDisturb: expected dim={self.dim}, got {x.shape[-1]}" + ) + + if self.blur: + x = self.blur(x.transpose(1, 2)).transpose(1, 2) + + if self.drop_prob: + if self.drop_type == "random": + time_mask = ( + torch.rand(x.shape[0], x.shape[1], 1, device=x.device) + > self.drop_prob + ) + x = x * time_mask.float() + elif self.drop_type == "equal_space": + time_mask = self._create_equal_space_mask( + x.shape[0], x.shape[1], x.device + ) + x = x * time_mask.float() + else: + raise ValueError(f"Unknown drop_type: {self.drop_type}") + + if self.noise_scale: + noise = torch.randn_like(x) * self.noise_scale + x = x + noise + + return x + + +class MIDIDigitalEmbedding(nn.Module): + """Embeds continuous MIDI values into discrete token embeddings. + + Continuous MIDI values in [0, 127] are quantized at a configurable + resolution (mark_distinguish_scale) and mapped to learned embeddings. + """ + + def __init__(self, embed_dim=128, num_classes=128, mark_distinguish_scale=2): + super().__init__() + + # num_classes covers the input range [0, 127] plus 2 special tokens + self.num_classes = num_classes + 2 + self.mark_distinguish_scale = mark_distinguish_scale + self.embedding_input_num_class = self.num_classes * self.mark_distinguish_scale + self.embedding = nn.Embedding(self.embedding_input_num_class, embed_dim) + + def midi_to_class(self, midi_values): + """Map continuous MIDI values to discrete class indices. + + Args: + midi_values: [B, T] continuous MIDI values, roughly in [0, 127] + + Returns: + class_indices: [B, T] discrete class indices + """ + # Round to nearest quantization step + # e.g., with scale=2: 0->0, 0.3->1, 0.5->1, 0.8->2, 1.0->2, ... + class_indices = torch.round(midi_values * self.mark_distinguish_scale).long() + + # Clamp to valid range + class_indices = torch.clamp( + class_indices, 0, self.embedding_input_num_class - 1 + ) + + return class_indices + + def forward(self, midi_values): + """ + Args: + midi_values: [B, T] continuous MIDI values + + Returns: + embeddings: [B, T, embed_dim] embedding vectors + """ + class_indices = self.midi_to_class(midi_values) + embeddings = self.embedding(class_indices) + return embeddings diff --git a/src/YingMusicSinger/melody/midi_extractor.py b/src/YingMusicSinger/melody/midi_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..01be279a41d914071161dde9e02a2168d743b0b0 --- /dev/null +++ b/src/YingMusicSinger/melody/midi_extractor.py @@ -0,0 +1,208 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence + +from src.YingMusicSinger.melody.Gconform import Gmidi_conform + +# midi decoding utils + + +def decode_gaussian_blurred_probs(probs, vmin, vmax, deviation, threshold): + num_bins = int(probs.shape[-1]) + interval = (vmax - vmin) / (num_bins - 1) + width = int(3 * deviation / interval) # 3 * sigma + idx = torch.arange(num_bins, device=probs.device)[None, None, :] # [1, 1, N] + idx_values = idx * interval + vmin + center = torch.argmax(probs, dim=-1, keepdim=True) # [B, T, 1] + start = torch.clip(center - width, min=0) # [B, T, 1] + end = torch.clip(center + width + 1, max=num_bins) # [B, T, 1] + idx_masks = (idx >= start) & (idx < end) # [B, T, N] + weights = probs * idx_masks # [B, T, N] + product_sum = torch.sum(weights * idx_values, dim=2) # [B, T] + weight_sum = torch.sum(weights, dim=2) # [B, T] + values = product_sum / ( + weight_sum + (weight_sum == 0) + ) # avoid dividing by zero, [B, T] + rest = probs.max(dim=-1)[0] < threshold # [B, T] + return values, rest + + +def decode_bounds_to_alignment(bounds, use_diff=True): + bounds_step = bounds.cumsum(dim=1).round().long() + if use_diff: + bounds_inc = ( + torch.diff( + bounds_step, + dim=1, + prepend=torch.full( + (bounds.shape[0], 1), + fill_value=-1, + dtype=bounds_step.dtype, + device=bounds_step.device, + ), + ) + > 0 + ) + else: + bounds_inc = F.pad( + (bounds_step[:, 1:] > bounds_step[:, :-1]), [1, 0], value=True + ) + frame2item = bounds_inc.long().cumsum(dim=1) + return frame2item + + +def decode_note_sequence(frame2item, values, masks, threshold=0.5): + """ + + :param frame2item: [1, 1, 1, 1, 2, 2, 3, 3, 3] + :param values: + :param masks: + :param threshold: minimum ratio of unmasked frames required to be regarded as an unmasked item + :return: item_values, item_dur, item_masks + """ + b = frame2item.shape[0] + space = frame2item.max() + 1 + + item_dur = frame2item.new_zeros(b, space, dtype=frame2item.dtype).scatter_add( + 1, frame2item, torch.ones_like(frame2item) + )[:, 1:] + item_unmasked_dur = frame2item.new_zeros( + b, space, dtype=frame2item.dtype + ).scatter_add(1, frame2item, masks.long())[:, 1:] + item_masks = item_unmasked_dur / item_dur >= threshold + + values_quant = values.round().long() + histogram = ( + frame2item.new_zeros(b, space * 128, dtype=frame2item.dtype) + .scatter_add( + 1, frame2item * 128 + values_quant, torch.ones_like(frame2item) * masks + ) + .unflatten(1, [space, 128])[:, 1:, :] + ) + item_values_center = histogram.float().argmax(dim=2).to(dtype=values.dtype) + values_center = torch.gather(F.pad(item_values_center, [1, 0]), 1, frame2item) + values_near_center = ( + masks & (values >= values_center - 0.5) & (values <= values_center + 0.5) + ) + item_valid_dur = frame2item.new_zeros(b, space, dtype=frame2item.dtype).scatter_add( + 1, frame2item, values_near_center.long() + )[:, 1:] + item_values = values.new_zeros(b, space, dtype=values.dtype).scatter_add( + 1, frame2item, values * values_near_center + )[:, 1:] / (item_valid_dur + (item_valid_dur == 0)) + + return item_values, item_dur, item_masks + + +def expand_batch_padded(feature_tensor, counts_tensor, padding_value=0.0): + assert feature_tensor.dim() == 2 and counts_tensor.dim() == 2 + + lengths = torch.sum(counts_tensor, dim=1) + + feature_tensor = feature_tensor.reshape(-1) + counts_tensor = counts_tensor.reshape(-1) + expanded_flat = torch.repeat_interleave(feature_tensor, counts_tensor) + + ragged_list = torch.split(expanded_flat, lengths.tolist()) + + padded_tensor = pad_sequence( + ragged_list, batch_first=True, padding_value=padding_value + ) + + return padded_tensor, lengths + + +class midi_loss(nn.Module): + def __init__(self): + super().__init__() + self.loss = nn.BCELoss() + + def forward(self, x, target): + midiout, cutp = x + midi_target, cutp_target = target + + cutploss = self.loss(cutp, cutp_target) + midiloss = self.loss(midiout, midi_target) + return midiloss, cutploss + + +class MIDIExtractor(nn.Module): + def __init__(self, in_dim=None, out_dim=None): + super().__init__() + + cfg = { + "attention_drop": 0.1, + "attention_heads": 8, + "attention_heads_dim": 64, + "conv_drop": 0.1, + "dim": 512, + "ffn_latent_drop": 0.1, + "ffn_out_drop": 0.1, + "kernel_size": 31, + "lay": 8, + "use_lay_skip": True, + "indim": 80, + "outdim": 128, + } + if in_dim is not None: + cfg["indim"] = in_dim + if out_dim is not None: + cfg["outdim"] = out_dim + + self.midi_conform = Gmidi_conform(**cfg) + + self.midi_min = 0 + self.midi_max = 127 + self.midi_deviation = 1.0 + self.rest_threshold = 0.1 + + def _load_form_ckpt(self, ckpt_path, device="cpu"): + from collections import OrderedDict + + if ckpt_path is None: + raise ValueError("midi_extractor_path is required") + + state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"] + prefix_in_ckpt = "model.model" + state_dict = OrderedDict( + { + k.replace(f"{prefix_in_ckpt}.", "midi_conform."): v + for k, v in state_dict.items() + if k.startswith(f"{prefix_in_ckpt}.") + } + ) + self.load_state_dict(state_dict, strict=True) + # self.to(device) + + def forward(self, x, mask=None): + midi, bound = self.midi_conform(x, mask) + + return midi, bound + + def postprocess(self, midi, bounds, with_expand=False): + probs = torch.sigmoid(midi) + + bound_probs = torch.sigmoid(bounds) + bound_probs = torch.squeeze(bound_probs, -1) + + masks = torch.ones_like(bound_probs).bool() + # Avoid in-place ops on tensors needed for autograd (outputs of SigmoidBackward) + probs = probs * masks[..., None] + bound_probs = bound_probs * masks + unit2note_pred = decode_bounds_to_alignment(bound_probs) * masks + midi_pred, rest_pred = decode_gaussian_blurred_probs( + probs, + vmin=self.midi_min, + vmax=self.midi_max, + deviation=self.midi_deviation, + threshold=self.rest_threshold, + ) + note_midi_pred, note_dur_pred, note_mask_pred = decode_note_sequence( + unit2note_pred, midi_pred, ~rest_pred & masks + ) + if not with_expand: + return note_midi_pred, note_dur_pred + + note_midi_expand, _ = expand_batch_padded(note_midi_pred, note_dur_pred) + return note_midi_expand, None diff --git a/src/YingMusicSinger/models/__init__.py b/src/YingMusicSinger/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c4dcb5fbfd482efeb8e96c881c855d4b368eddff --- /dev/null +++ b/src/YingMusicSinger/models/__init__.py @@ -0,0 +1 @@ +from .dit import DiT diff --git a/src/YingMusicSinger/models/dit.py b/src/YingMusicSinger/models/dit.py new file mode 100644 index 0000000000000000000000000000000000000000..676ab120a25ea7de6a3f0be94ba093e84a89177a --- /dev/null +++ b/src/YingMusicSinger/models/dit.py @@ -0,0 +1,472 @@ +""" +ein notation: +b - batch +n - sequence +nt - text sequence +nw - raw wave length +d - dimension +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn +from x_transformers.x_transformers import RotaryEmbedding + +from src.YingMusicSinger.models.modules import ( + AdaLayerNorm_Final, + ConvNeXtV2Block, + ConvPositionEmbedding, + DiTBlock, + TimestepGuidanceEmbedding, + get_pos_embed_indices, + precompute_freqs_cis, +) + + +# Text embedding + + +class TextEmbedding(nn.Module): + def __init__( + self, + text_num_embeds, + text_dim, + mask_padding=False, + average_upsampling=False, + conv_layers=0, + conv_mult=2, + ): + super().__init__() + self.text_embed = nn.Embedding( + text_num_embeds + 1, text_dim + ) # index 0 reserved as filler token + + self.mask_padding = mask_padding + self.average_upsampling = average_upsampling # ZipVoice-style late average upsampling (after text encoder) + if average_upsampling: + assert mask_padding, ( + "text_embedding_average_upsampling requires text_mask_padding to be True" + ) + + if conv_layers > 0: + self.extra_modeling = True + self.precompute_max_pos = 4096 # ~44s of 24kHz audio + self.register_buffer( + "freqs_cis", + precompute_freqs_cis(text_dim, self.precompute_max_pos), + persistent=False, + ) + self.text_blocks = nn.Sequential( + *[ + ConvNeXtV2Block(text_dim, text_dim * conv_mult) + for _ in range(conv_layers) + ] + ) + else: + self.extra_modeling = False + + print( + f"[info] TextEmbedding: mask_padding={mask_padding}, average_upsampling={average_upsampling}, conv_layers={conv_layers}" + ) + + def average_upsample_text_by_mask(self, text, text_mask, audio_mask): + batch, text_len, text_dim = text.shape + + if audio_mask is None: + audio_mask = torch.ones_like(text_mask, dtype=torch.bool) + valid_mask = audio_mask & text_mask + audio_lens = audio_mask.sum(dim=1) # [batch] + valid_lens = valid_mask.sum(dim=1) # [batch] + + upsampled_text = torch.zeros_like(text) + + for i in range(batch): + audio_len = audio_lens[i].item() + valid_len = valid_lens[i].item() + + if valid_len == 0: + continue + + valid_ind = torch.where(valid_mask[i])[0] + valid_data = text[i, valid_ind, :] # [valid_len, text_dim] + + base_repeat = audio_len // valid_len + remainder = audio_len % valid_len + + indices = [] + for j in range(valid_len): + repeat_count = base_repeat + (1 if j >= valid_len - remainder else 0) + indices.extend([j] * repeat_count) + + indices = torch.tensor( + indices[:audio_len], device=text.device, dtype=torch.long + ) + upsampled = valid_data[indices] # [audio_len, text_dim] + + upsampled_text[i, :audio_len, :] = upsampled + + return upsampled_text + + def forward( + self, + text: int["b nt"], + seq_len, + drop_text=False, + audio_mask: bool["b n"] | None = None, + ): # noqa: F722 + # Text tokens start from 0; shift by 1 so that 0 is never a valid token + text = text + 1 + # Note: 1 is used as the PAD token + text = text[ + :, :seq_len + ] # Truncate if text tokens exceed mel spectrogram length + batch, text_len = text.shape[0], text.shape[1] + text = F.pad(text, (0, seq_len - text_len), value=1) + + if self.mask_padding: + text_mask = text == 1 + else: + text_mask = torch.zeros( + (batch, seq_len), device=text.device, dtype=torch.bool + ) + + if drop_text: # CFG for text + text = torch.zeros_like(text) + + text = self.text_embed(text) # b n -> b n d + + # Optional extra modeling + if self.extra_modeling: + # Sinusoidal positional embedding + batch_start = torch.zeros((batch,), device=text.device, dtype=torch.long) + pos_idx = get_pos_embed_indices( + batch_start, seq_len, max_pos=self.precompute_max_pos + ) + text_pos_embed = self.freqs_cis[pos_idx] + text = text + text_pos_embed + + # ConvNeXtV2 blocks + if self.mask_padding: + text = text.masked_fill( + text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0 + ) + for block in self.text_blocks: + text = block(text) + text = text.masked_fill( + text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0 + ) + else: + text = self.text_blocks(text) + + if self.average_upsampling: + text = self.average_upsample_text_by_mask(text, ~text_mask, audio_mask) + + return text, text_mask + + +# Noised input audio and context mixing embedding + + +class InputEmbedding(nn.Module): + def __init__(self, mel_dim, text_dim, out_dim, midi_dim=128): + super().__init__() + self.proj = nn.Linear(mel_dim * 2 + text_dim + midi_dim, out_dim) + self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim) + self.midi_proj = nn.Linear(128, 128) + + def forward( + self, + x: float["b n d"], # noqa: F722 + cond: float["b n d"], # noqa: F722 + text_embed: float["b n d"], # noqa: F722 + midi, + drop_audio_cond=False, + drop_midi=False, + ): + if drop_audio_cond: # CFG for conditioning audio + cond = torch.zeros_like(cond) + + midi = self.midi_proj(midi) + + if drop_midi: # CFG for melody + midi = torch.zeros_like(midi) + + x = self.proj(torch.cat((x, cond, text_embed, midi), dim=-1)) + x = self.conv_pos_embed(x) + x + return x + + +# Transformer backbone using DiT blocks + + +class DiT(nn.Module): + def __init__( + self, + *, + dim, + depth=8, + heads=8, + dim_head=64, + dropout=0.1, + ff_mult=4, + mel_dim=100, + text_num_embeds=256, + text_dim=None, + n_f0_bins=512, + text_mask_padding=True, + text_embedding_average_upsampling=False, + qk_norm=None, + conv_layers=0, + pe_attn_head=None, + attn_backend="torch", # "torch" | "flash_attn" + attn_mask_enabled=False, + long_skip_connection=False, + checkpoint_activations=False, + use_guidance_scale_embed: bool = False, + guidance_scale_embed_dim: int = 192, + ): + super().__init__() + + self.time_embed = TimestepGuidanceEmbedding( + dim, + use_guidance_scale_embed=use_guidance_scale_embed, + guidance_scale_embed_dim=guidance_scale_embed_dim, + ) + if text_dim is None: + text_dim = mel_dim + self.text_embed_p = TextEmbedding( + text_num_embeds, + text_dim, + mask_padding=text_mask_padding, + average_upsampling=text_embedding_average_upsampling, + conv_layers=conv_layers, + ) + self.text_cond, self.text_uncond = None, None # text cache + self.input_embed_with_midi = InputEmbedding(mel_dim, text_dim, dim) + + self.rotary_embed = RotaryEmbedding(dim_head) + self.use_guidance_scale_embed = use_guidance_scale_embed + + self.dim = dim + self.depth = depth + + self.transformer_blocks = nn.ModuleList( + [ + DiTBlock( + dim=dim, + heads=heads, + dim_head=dim_head, + ff_mult=ff_mult, + dropout=dropout, + qk_norm=qk_norm, + pe_attn_head=pe_attn_head, + attn_backend=attn_backend, + attn_mask_enabled=attn_mask_enabled, + ) + for _ in range(depth) + ] + ) + self.long_skip_connection = ( + nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None + ) + + self.norm_out = AdaLayerNorm_Final(dim) # Final modulation + self.proj_out = nn.Linear(dim, mel_dim) + + self.checkpoint_activations = checkpoint_activations + + self.initialize_weights() + + def initialize_weights(self): + # Zero-out AdaLN layers in DiT blocks + for block in self.transformer_blocks: + nn.init.constant_(block.attn_norm.linear.weight, 0) + nn.init.constant_(block.attn_norm.linear.bias, 0) + + # Zero-out output layers + nn.init.constant_(self.norm_out.linear.weight, 0) + nn.init.constant_(self.norm_out.linear.bias, 0) + nn.init.constant_(self.proj_out.weight, 0) + nn.init.constant_(self.proj_out.bias, 0) + + nn.init.zeros_(self.input_embed_with_midi.midi_proj.weight) + nn.init.zeros_(self.input_embed_with_midi.midi_proj.bias) + + def ckpt_wrapper(self, module): + # Ref: https://github.com/chuanyangjin/fast-DiT/blob/main/models.py + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward + + def get_input_embed( + self, + x, # b n d + cond, # b n d + text, # b nt + midi, # b n + drop_audio_cond: bool = False, + drop_text: bool = False, + drop_midi: bool = False, + cache: bool = True, + audio_mask: bool["b n"] | None = None, # noqa: F722 + ): + seq_len = x.shape[1] + + if cache: + if drop_text: + if self.text_uncond is None: + self.text_uncond, _ = self.text_embed_p( + text, seq_len, drop_text=True, audio_mask=audio_mask + ) + text_embed = self.text_uncond + else: + if self.text_cond is None: + self.text_cond, _ = self.text_embed_p( + text, seq_len, drop_text=False, audio_mask=audio_mask + ) + text_embed = self.text_cond + else: + text_embed, text_mask = self.text_embed_p( + text, seq_len, drop_text=drop_text, audio_mask=audio_mask + ) + + if midi is None: + midi = torch.zeros( + (x.size(0), x.size(1)), device=x.device, dtype=torch.long + ) + + x = self.input_embed_with_midi( + x, + cond, + text_embed, + midi, + drop_audio_cond=drop_audio_cond, + drop_midi=drop_midi, + ) + + return x, None + + def clear_cache(self): + self.text_cond, self.text_uncond = None, None + + def forward( + self, + x: float["b n d"], # Noised input audio # noqa: F722 + cond: float["b n d"], # Masked conditioning audio # noqa: F722 + text: int["b nt"], # Text tokens # noqa: F722 + time: float["b"] | float[""], # Timestep # noqa: F821 F722 + midi: float["b n"] | None = None, # Melody latent # noqa: F722 + mask: bool["b n"] | None = None, # noqa: F722 + drop_audio_cond: bool = False, # CFG for conditioning audio + drop_text: bool = False, # CFG for text + drop_midi: bool = False, # CFG for melody + cfg_infer: bool = False, # CFG inference: pack cond & uncond forward + cache: bool = False, + guidance_scale=None, + cfg_infer_ids=None, # tuple(bool): (x_cond, x_uncond, x_uncond_cc, x_drop_all_cond) + ): + batch, seq_len = x.shape[0], x.shape[1] + if time.ndim == 0: + time = time.repeat(batch) + + # Timestep embedding (with optional distillation guidance scale) + t = self.time_embed(time, guidance_scale=guidance_scale) + + if cfg_infer: # Pack cond & uncond forward: b n d -> Kb n d + x_cond, x_uncond, x_uncond_cc, x_drop_all_cond = None, None, None, None + if cfg_infer_ids is None or cfg_infer_ids[0]: + x_cond, _ = self.get_input_embed( + x, + cond, + text, + midi, + drop_audio_cond=False, + drop_text=False, + drop_midi=False, + cache=cache, + audio_mask=mask, + ) + if cfg_infer_ids is None or cfg_infer_ids[1]: + x_uncond, _ = self.get_input_embed( + x, + cond, + text, + midi, + drop_audio_cond=True, + drop_text=False, + drop_midi=False, + cache=cache, + audio_mask=mask, + ) + if cfg_infer_ids is None or cfg_infer_ids[2]: + x_uncond_cc, _ = self.get_input_embed( + x, + cond, + text, + midi, + drop_audio_cond=False, + drop_text=True, + drop_midi=True, + cache=cache, + audio_mask=mask, + ) + if cfg_infer_ids is None or cfg_infer_ids[3]: + x_drop_all_cond, _ = self.get_input_embed( + x, + cond, + text, + midi, + drop_audio_cond=True, + drop_text=True, + drop_midi=True, + cache=cache, + audio_mask=mask, + ) + + # Concatenate only non-None tensors + x_list = [ + xi + for xi in [x_cond, x_uncond, x_uncond_cc, x_drop_all_cond] + if xi is not None + ] + x = torch.cat(x_list, dim=0) + t = torch.cat([t] * len(x_list), dim=0) + mask = torch.cat([mask] * len(x_list), dim=0) if mask is not None else None + else: + x, text_inner_sim_matrix = self.get_input_embed( + x, + cond, + text, + midi, + drop_audio_cond=drop_audio_cond, + drop_text=drop_text, + drop_midi=drop_midi, + cache=cache, + audio_mask=mask, + ) + + rope = self.rotary_embed.forward_from_seq_len(seq_len) + + if self.long_skip_connection is not None: + residual = x + + # Mask is all zeros during inference + for block in self.transformer_blocks: + if self.checkpoint_activations: + x = torch.utils.checkpoint.checkpoint( + self.ckpt_wrapper(block), x, t, mask, rope, use_reentrant=False + ) + else: + x = block(x, t, mask=mask, rope=rope) + + if self.long_skip_connection is not None: + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + + x = self.norm_out(x, t) + output = self.proj_out(x) + + return output, text_inner_sim_matrix if not cfg_infer else None diff --git a/src/YingMusicSinger/models/model.py b/src/YingMusicSinger/models/model.py new file mode 100644 index 0000000000000000000000000000000000000000..6b5c0e8cc6dd4903b357804e2ed51ffb5b948cfb --- /dev/null +++ b/src/YingMusicSinger/models/model.py @@ -0,0 +1,423 @@ +from __future__ import annotations + +from typing import Callable + +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.utils.rnn import pad_sequence +from torchdiffeq import odeint + +from src.YingMusicSinger.melody.midi_extractor import MIDIExtractor +from src.YingMusicSinger.utils.common import ( + default, + exists, + get_epss_timesteps, + lens_to_mask, +) + + +def interpolation_midi_continuous(midi_p, bound_p, total_len): + """Temporally interpolate 3D melody latent to match target length.""" + if midi_p.shape[1] != total_len: + midi = ( + F.interpolate( + midi_p.clone().detach().transpose(1, 2), + size=total_len, + mode="linear", + align_corners=False, + ) + .transpose(1, 2) + .clone() + .detach() + ) + if bound_p is not None: + midi_bound = ( + F.interpolate( + bound_p.clone().detach().transpose(1, 2), + size=total_len, + mode="linear", + align_corners=False, + ) + .transpose(1, 2) + .clone() + .detach() + ) + else: + midi = midi_p.clone().detach() + if bound_p is not None: + midi_bound = bound_p.clone().detach() + if bound_p is not None: + return midi, midi_bound + else: + return midi + + +def interpolation_midi_continuous_2_dim(midi_p, bound_p, total_len): + """Temporally interpolate 2D melody latent to match target length.""" + assert len(midi_p.shape) == 2 + + if midi_p.shape[1] != total_len: + midi = ( + F.interpolate( + midi_p.unsqueeze(2).clone().detach().transpose(1, 2), + size=total_len, + mode="linear", + align_corners=False, + ) + .transpose(1, 2) + .clone() + .detach() + ) + if bound_p: + midi_bound = ( + F.interpolate( + bound_p.unsqueeze(2).clone().detach().transpose(1, 2), + size=total_len, + mode="linear", + align_corners=False, + ) + .transpose(1, 2) + .clone() + .detach() + ) + else: + midi = midi_p.clone().detach() + if bound_p: + midi_bound = bound_p.clone().detach() + if bound_p: + return midi.squeeze(2), midi_bound.squeeze(2) + else: + return midi.squeeze(2) + + +class Singer(nn.Module): + def __init__( + self, + transformer: nn.Module, + is_tts_pretrain, + melody_input_source, + cka_disabled, + distill_stage, + use_guidance_scale_embed, + sigma=0.0, + odeint_kwargs: dict = dict(method="euler"), + audio_drop_prob=0.3, + cond_drop_prob=0.2, + num_channels=None, + mel_spec_module: nn.Module | None = None, + mel_spec_kwargs: dict = dict(), + frac_lengths_mask: tuple[float, float] = (0.7, 1.0), + extra_parameters=None, + ): + super().__init__() + + self.is_tts_pretrain = is_tts_pretrain + + if distill_stage is None: + self.distill_stage = 0 + else: + self.distill_stage = int(distill_stage) + + self.use_guidance_scale_embed = use_guidance_scale_embed + + assert melody_input_source in { + "student_model", + "some_pretrain", + "some_pretrain_fuzzdisturb", + "some_pretrain_postprocess_embedding", + "none", + } + from src.YingMusicSinger.melody.SmoothMelody import MIDIFuzzDisturb + + if melody_input_source == "some_pretrain_fuzzdisturb": + self.smoothMelody_MIDIFuzzDisturb = MIDIFuzzDisturb( + dim=extra_parameters.some_pretrain_fuzzdisturb.dim, + drop_prob=extra_parameters.some_pretrain_fuzzdisturb.drop_prob, + noise_scale=extra_parameters.some_pretrain_fuzzdisturb.noise_scale, + blur_kernel=extra_parameters.some_pretrain_fuzzdisturb.blur_kernel, + drop_type=extra_parameters.some_pretrain_fuzzdisturb.drop_type, + ) + from src.YingMusicSinger.melody.SmoothMelody import MIDIDigitalEmbedding + + if melody_input_source == "some_pretrain_postprocess_embedding": + self.smoothMelody_MIDIDigitalEmbedding = MIDIDigitalEmbedding( + embed_dim=extra_parameters.some_pretrain_postprocess_embedding.embed_dim, + num_classes=extra_parameters.some_pretrain_postprocess_embedding.num_classes, + mark_distinguish_scale=extra_parameters.some_pretrain_postprocess_embedding.mark_distinguish_scale, + ) + + self.melody_input_source = melody_input_source + self.cka_disabled = cka_disabled + + self.frac_lengths_mask = frac_lengths_mask + + num_channels = default(num_channels, mel_spec_kwargs.n_mel_channels) + self.num_channels = num_channels + + # Classifier-free guidance drop probabilities + self.audio_drop_prob = audio_drop_prob + self.cond_drop_prob = cond_drop_prob + + # Transformer backbone + self.transformer = transformer + dim = transformer.dim + self.dim = dim + + # Conditional flow matching + self.sigma = sigma + self.odeint_kwargs = odeint_kwargs + + # Melody extractor + self.midi_extractor = MIDIExtractor(in_dim=num_channels) + + @property + def device(self): + return next(self.parameters()).device + + @torch.no_grad() + def sample( + self, + cond: float["b n d"] | float["b nw"], # noqa: F722 + text: int["b nt"] | list[str], # noqa: F722 + duration: int | int["b"] | None = None, # noqa: F821 + *, + midi_in: float["b n d"] | None = None, + lens: int["b"] | None = None, # noqa: F821 + steps=32, + cfg_strength=1.0, + sway_sampling_coef=None, + seed: int | None = None, + max_duration=4096, # Maximum total length (including ICL prompt), ~190s + vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None, # noqa: F722 + use_epss=True, + no_ref_audio=False, + duplicate_test=False, + t_inter=0.1, + t_shift=1.0, # Sampling timestep shift (ZipVoice-style) + guidance_scale=None, + edit_mask=None, + midi_p=None, + bound_p=None, + enable_melody_control=True, + ): + self.eval() + + assert isinstance(cond, torch.Tensor) + assert not edit_mask, "edit_mask is not supported in this mode" + assert not duplicate_test, "duplicate_test is not supported in this mode" + + if self.melody_input_source == "student_model": + assert midi_p is None and bound_p is None + elif self.melody_input_source in { + "some_pretrain", + "some_pretrain_fuzzdisturb", + "some_pretrain_postprocess_embedding", + }: + assert midi_p is not None and bound_p is not None + elif self.melody_input_source == "none": + assert midi_p is None and bound_p is None + else: + raise ValueError( + f"Unsupported melody_input_source: {self.melody_input_source}" + ) + + # duration is the total latent sequence length + assert duration + + cond = cond.to(next(self.parameters()).dtype) + + # Extract or interpolate melody representation + if self.melody_input_source == "student_model": + midi, midi_bound = self.midi_extractor(midi_in) + + elif self.melody_input_source == "some_pretrain": + midi, midi_bound = interpolation_midi_continuous( + midi_p=midi_p, bound_p=bound_p, total_len=text.shape[1] + ) + elif self.melody_input_source == "some_pretrain_fuzzdisturb": + midi, midi_bound = interpolation_midi_continuous( + midi_p=midi_p, bound_p=bound_p, total_len=text.shape[1] + ) + midi = self.smoothMelody_MIDIFuzzDisturb(midi) + + elif self.melody_input_source == "some_pretrain_postprocess_embedding": + midi_after_postprocess, _ = self.midi_extractor.postprocess( + midi=midi_p, bounds=bound_p, with_expand=True + ) + midi = interpolation_midi_continuous_2_dim( + midi_p=midi_after_postprocess, bound_p=None, total_len=text.shape[1] + ) + midi = self.smoothMelody_MIDIDigitalEmbedding(midi) + midi_bound = None + + elif self.melody_input_source == "none": + midi = torch.zeros( + text.shape[0], text.shape[1], 128, dtype=cond.dtype, device=text.device + ) + midi_bound = None + else: + raise NotImplementedError() + + batch, cond_seq_len, device = *cond.shape[:2], cond.device + if not exists(lens): + lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long) + + assert isinstance(text, torch.Tensor) + + cond_mask = lens_to_mask(lens) + + if edit_mask is not None: + cond_mask = cond_mask & edit_mask + + if isinstance(duration, int): + duration = torch.full((batch,), duration, device=device, dtype=torch.long) + + # Duration must be at least max(text_len, audio_prompt_len) + 1 + duration = torch.maximum( + torch.maximum((text != 0).sum(dim=-1), lens) + 1, duration + ) + duration = duration.clamp(max=max_duration) + + max_duration = duration.amax() + + # Duplicate test: interpolate between noise and conditioning + if duplicate_test: + test_cond = F.pad( + cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0 + ) + + # Zero-pad conditioning latent to max_duration + cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0) + + if no_ref_audio: + cond = torch.zeros_like(cond) + + cond_mask = F.pad( + cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False + ) + cond_mask = cond_mask.unsqueeze(-1) + step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond)) + + assert max_duration == midi.shape[1] + + # Zero out melody in prompt region; optionally disable melody control entirely + if enable_melody_control: + midi = torch.where(cond_mask, torch.zeros_like(midi), midi) + else: + midi = torch.zeros_like(midi) + + if self.is_tts_pretrain: + midi = torch.zeros_like(midi) + + # For batched inference, explicit mask prevents causal attention fallback + if batch > 1: + mask = lens_to_mask(duration) + else: + mask = None + + # ODE velocity function + def fn(t, x): + if cfg_strength < 1e-5: + # No classifier-free guidance + pred, _ = self.transformer( + x=x, + cond=step_cond, + text=text, + midi=midi, + time=t, + mask=mask, + drop_audio_cond=False, + drop_text=False, + drop_midi=not enable_melody_control, + cache=False, + ) + return pred + else: + if self.use_guidance_scale_embed: + # Distilled model with built-in CFG + assert enable_melody_control + pred_cfg, _ = self.transformer( + x=x, + cond=step_cond, + text=text, + midi=midi, + time=t, + mask=mask, + drop_audio_cond=False, + drop_text=False, + drop_midi=not enable_melody_control, + cache=False, + guidance_scale=torch.tensor([guidance_scale], device=device), + ) + print( + f"CFG 参数调节无作用！ 蒸馏之后的，输入CFG为 guidance_scale={guidance_scale}" + ) + return pred_cfg + else: + # Standard CFG: cond + uncond forward + # BUG If enable_melody_control is False, there might be a slight issue here + assert guidance_scale is not None + pred_cfg, _ = self.transformer( + x=x, + cond=step_cond, + text=text, + midi=midi, + time=t, + mask=mask, + cfg_infer=True, + cache=False, + cfg_infer_ids=(True, False, False, True), + ) + + pred, pred_drop_all_cond = torch.chunk(pred_cfg, 2, dim=0) + return pred + (pred - pred_drop_all_cond) * float(guidance_scale) + + # Generate initial noise (per-sample seeding for batch consistency) + y0 = [] + for dur in duration: + if exists(seed): + torch.manual_seed(seed) + y0.append( + torch.randn( + dur, self.num_channels, device=self.device, dtype=step_cond.dtype + ) + ) + y0 = pad_sequence(y0, padding_value=0, batch_first=True) + + t_start = 0 + + if duplicate_test: + t_start = t_inter + y0 = (1 - t_start) * y0 + t_start * test_cond + steps = int(steps * (1 - t_start)) + + # Build timestep schedule + assert not use_epss and sway_sampling_coef is None, ( + "Use timestep shift instead of the strategy in F5" + ) + if t_start == 0 and use_epss: + # Empirically Pruned Step Sampling for low NFE + t = get_epss_timesteps(steps, device=self.device, dtype=step_cond.dtype) + else: + t = torch.linspace( + t_start, 1, steps + 1, device=self.device, dtype=step_cond.dtype + ) + + if sway_sampling_coef is not None: + t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) + + # Apply timestep shift + t = t_shift * t / (1 + (t_shift - 1) * t) + + trajectory = odeint(fn, y0, t, **self.odeint_kwargs) + self.transformer.clear_cache() + + sampled = trajectory[-1] + out = sampled + + if exists(vocoder): + out = out.permute(0, 2, 1) + out = vocoder(out) + + return out, trajectory diff --git a/src/YingMusicSinger/models/modules.py b/src/YingMusicSinger/models/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..b694ed56e589de93f0672896339f41703a1bc9af --- /dev/null +++ b/src/YingMusicSinger/models/modules.py @@ -0,0 +1,961 @@ +""" +ein notation: +b - batch +n - sequence +nt - text sequence +nw - raw wave length +d - dimension +""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchaudio +from librosa.filters import mel as librosa_mel_fn +from x_transformers.x_transformers import apply_rotary_pos_emb + +from src.YingMusicSinger.utils.common import is_package_available + +# raw wav to mel spec + + +mel_basis_cache = {} +hann_window_cache = {} + + +def get_bigvgan_mel_spectrogram( + waveform, + n_fft=1024, + n_mel_channels=100, + target_sample_rate=24000, + hop_length=256, + win_length=1024, + fmin=0, + fmax=None, + center=False, +): # Copy from https://github.com/NVIDIA/BigVGAN/tree/main + device = waveform.device + key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}" + + if key not in mel_basis_cache: + mel = librosa_mel_fn( + sr=target_sample_rate, + n_fft=n_fft, + n_mels=n_mel_channels, + fmin=fmin, + fmax=fmax, + ) + mel_basis_cache[key] = ( + torch.from_numpy(mel).float().to(device) + ) # TODO: why they need .float()? + hann_window_cache[key] = torch.hann_window(win_length).to(device) + + mel_basis = mel_basis_cache[key] + hann_window = hann_window_cache[key] + + padding = (n_fft - hop_length) // 2 + waveform = torch.nn.functional.pad( + waveform.unsqueeze(1), (padding, padding), mode="reflect" + ).squeeze(1) + + spec = torch.stft( + waveform, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) + + mel_spec = torch.matmul(mel_basis, spec) + mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5)) + + return mel_spec + + +def get_vocos_mel_spectrogram( + waveform, + n_fft=1024, + n_mel_channels=100, + target_sample_rate=24000, + hop_length=256, + win_length=1024, +): + mel_stft = torchaudio.transforms.MelSpectrogram( + sample_rate=target_sample_rate, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + n_mels=n_mel_channels, + power=1, + center=True, + normalized=False, + norm=None, + ).to(waveform.device) + if len(waveform.shape) == 3: + waveform = waveform.squeeze(1) # 'b 1 nw -> b nw' + + assert len(waveform.shape) == 2 + + mel = mel_stft(waveform) + mel = mel.clamp(min=1e-5).log() + return mel + + +class MelSpec(nn.Module): + def __init__( + self, + n_fft=1024, + hop_length=256, + win_length=1024, + n_mel_channels=100, + target_sample_rate=24_000, + mel_spec_type="vocos", + ): + super().__init__() + assert mel_spec_type in ["vocos", "bigvgan"], print( + "We only support two extract mel backend: vocos or bigvgan" + ) + + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + self.n_mel_channels = n_mel_channels + self.target_sample_rate = target_sample_rate + + if mel_spec_type == "vocos": + self.extractor = get_vocos_mel_spectrogram + elif mel_spec_type == "bigvgan": + self.extractor = get_bigvgan_mel_spectrogram + + self.register_buffer("dummy", torch.tensor(0), persistent=False) + + def forward(self, wav): + if self.dummy.device != wav.device: + self.to(wav.device) + + mel = self.extractor( + waveform=wav, + n_fft=self.n_fft, + n_mel_channels=self.n_mel_channels, + target_sample_rate=self.target_sample_rate, + hop_length=self.hop_length, + win_length=self.win_length, + ) + + return mel + + +# sinusoidal position embedding + + +class SinusPositionEmbedding(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x, scale=1000): + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +# convolutional position embedding + + +class ConvPositionEmbedding(nn.Module): + def __init__(self, dim, kernel_size=31, groups=16): + super().__init__() + assert kernel_size % 2 != 0 + self.conv1d = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2), + nn.Mish(), + ) + + def forward(self, x: float["b n d"], mask: bool["b n"] | None = None): + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) + + x = x.permute(0, 2, 1) + x = self.conv1d(x) + out = x.permute(0, 2, 1) + + if mask is not None: + out = out.masked_fill(~mask, 0.0) + + return out + + +# rotary positional embedding related + + +def precompute_freqs_cis( + dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0 +): + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ + # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py + theta *= theta_rescale_factor ** (dim / (dim - 2)) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) # type: ignore + freqs = torch.outer(t, freqs).float() # type: ignore + freqs_cos = torch.cos(freqs) # real part + freqs_sin = torch.sin(freqs) # imaginary part + return torch.cat([freqs_cos, freqs_sin], dim=-1) + + +def get_pos_embed_indices(start, length, max_pos, scale=1.0): + # length = length if isinstance(length, int) else length.max() + scale = scale * torch.ones_like( + start, dtype=torch.float32 + ) # in case scale is a scalar + pos = ( + start.unsqueeze(1) + + ( + torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) + * scale.unsqueeze(1) + ).long() + ) + # avoid extra long error. + pos = torch.where(pos < max_pos, pos, max_pos - 1) + return pos + + +# Global Response Normalization layer (Instance Normalization ?) + + +class GRN(nn.Module): + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=1, keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py +# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108 + + +class ConvNeXtV2Block(nn.Module): + def __init__( + self, + dim: int, + intermediate_dim: int, + dilation: int = 1, + ): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d( + dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation + ) # depthwise conv + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear( + dim, intermediate_dim + ) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = x.transpose(1, 2) # b n d -> b d n + x = self.dwconv(x) + x = x.transpose(1, 2) # b d n -> b n d + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + return residual + x + + +# RMSNorm + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + self.native_rms_norm = float(torch.__version__[:3]) >= 2.4 + + def forward(self, x): + if self.native_rms_norm: + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = F.rms_norm( + x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps + ) + else: + variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = x * self.weight + + return x + + +# AdaLayerNorm +# return with modulated x for attn input, and params for later mlp modulation + + +class AdaLayerNorm(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb=None): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk( + emb, 6, dim=1 + ) + + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +# AdaLayerNorm for final layer +# return only with modulated x for attn input, cuz no more mlp modulation + + +class AdaLayerNorm_Final(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb): + emb = self.linear(self.silu(emb)) + scale, shift = torch.chunk(emb, 2, dim=1) + + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +# FeedForward + + +class FeedForward(nn.Module): + def __init__( + self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none" + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential( + project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out) + ) + + def forward(self, x): + return self.ff(x) + + +# Attention with possible joint part +# modified from diffusers/src/diffusers/models/attention_processor.py + + +class Attention(nn.Module): + def __init__( + self, + processor: JointAttnProcessor | AttnProcessor, + dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + context_dim: Optional[int] = None, # if not None -> joint attention + context_pre_only: bool = False, + qk_norm: Optional[str] = None, + ): + super().__init__() + + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." + ) + + self.processor = processor + + self.dim = dim + self.heads = heads + self.inner_dim = dim_head * heads + self.dropout = dropout + + self.context_dim = context_dim + self.context_pre_only = context_pre_only + + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + + if qk_norm is None: + self.q_norm = None + self.k_norm = None + elif qk_norm == "rms_norm": + self.q_norm = RMSNorm(dim_head, eps=1e-6) + self.k_norm = RMSNorm(dim_head, eps=1e-6) + else: + raise ValueError(f"Unimplemented qk_norm: {qk_norm}") + + if self.context_dim is not None: + self.to_q_c = nn.Linear(context_dim, self.inner_dim) + self.to_k_c = nn.Linear(context_dim, self.inner_dim) + self.to_v_c = nn.Linear(context_dim, self.inner_dim) + if qk_norm is None: + self.c_q_norm = None + self.c_k_norm = None + elif qk_norm == "rms_norm": + self.c_q_norm = RMSNorm(dim_head, eps=1e-6) + self.c_k_norm = RMSNorm(dim_head, eps=1e-6) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, dim)) + self.to_out.append(nn.Dropout(dropout)) + + if self.context_dim is not None and not self.context_pre_only: + self.to_out_c = nn.Linear(self.inner_dim, context_dim) + + def forward( + self, + x: float["b n d"], # noised input x + c: float["b n d"] = None, # context c + mask: bool["b n"] | None = None, + rope=None, # rotary position embedding for x + c_rope=None, # rotary position embedding for c + ) -> torch.Tensor: + if c is not None: + return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope) + else: + return self.processor(self, x, mask=mask, rope=rope) + + +# Attention processor + +if is_package_available("flash_attn"): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import pad_input, unpad_input + + +class AttnProcessor: + def __init__( + self, + pe_attn_head: int + | None = None, # number of attention head to apply rope, None for all + attn_backend: str = "torch", # "torch" or "flash_attn" + attn_mask_enabled: bool = True, + ): + if attn_backend == "flash_attn": + assert is_package_available("flash_attn"), ( + "Please install flash-attn first." + ) + + self.pe_attn_head = pe_attn_head + self.attn_backend = attn_backend + self.attn_mask_enabled = attn_mask_enabled + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x + mask: bool["b n"] | None = None, + rope=None, # rotary position embedding + ) -> torch.FloatTensor: + batch_size = x.shape[0] + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + + # apply rotary position embedding + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = ( + (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + ) + + if self.pe_attn_head is not None: + pn = self.pe_attn_head + query[:, :pn, :, :] = apply_rotary_pos_emb( + query[:, :pn, :, :], freqs, q_xpos_scale + ) + key[:, :pn, :, :] = apply_rotary_pos_emb( + key[:, :pn, :, :], freqs, k_xpos_scale + ) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + if self.attn_backend == "torch": + # mask. e.g. inference got a batch with different target durations, mask out the padding + if self.attn_mask_enabled and mask is not None: + attn_mask = mask + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' + attn_mask = attn_mask.expand( + batch_size, attn.heads, query.shape[-2], key.shape[-2] + ) + else: + attn_mask = None + x = F.scaled_dot_product_attention( + query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False + ) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + + elif self.attn_backend == "flash_attn": + query = query.transpose(1, 2) # [b, h, n, d] -> [b, n, h, d] + key = key.transpose(1, 2) + value = value.transpose(1, 2) + if self.attn_mask_enabled and mask is not None: + query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input( + query, mask + ) + key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(key, mask) + value, _, _, _, _ = unpad_input(value, mask) + x = flash_attn_varlen_func( + query, + key, + value, + q_cu_seqlens, + k_cu_seqlens, + q_max_seqlen_in_batch, + k_max_seqlen_in_batch, + ) + x = pad_input(x, indices, batch_size, q_max_seqlen_in_batch) + x = x.reshape(batch_size, -1, attn.heads * head_dim) + else: + x = flash_attn_func(query, key, value, dropout_p=0.0, causal=False) + x = x.reshape(batch_size, -1, attn.heads * head_dim) + + x = x.to(query.dtype) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + + +# Joint Attention processor for MM-DiT +# modified from diffusers/src/diffusers/models/attention_processor.py + + +class JointAttnProcessor: + def __init__(self): + pass + + def __call__( + self, + attn: Attention, + x: float["b n d"], # noised input x + c: float["b nt d"] = None, # context c, here text + mask: bool["b n"] | None = None, + rope=None, # rotary position embedding for x + c_rope=None, # rotary position embedding for c + ) -> torch.FloatTensor: + residual = x + + batch_size = c.shape[0] + + # `sample` projections + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + + # `context` projections + c_query = attn.to_q_c(c) + c_key = attn.to_k_c(c) + c_value = attn.to_v_c(c) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # qk norm + if attn.q_norm is not None: + query = attn.q_norm(query) + if attn.k_norm is not None: + key = attn.k_norm(key) + if attn.c_q_norm is not None: + c_query = attn.c_q_norm(c_query) + if attn.c_k_norm is not None: + c_key = attn.c_k_norm(c_key) + + # apply rope for context and noised input independently + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = ( + (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + ) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + if c_rope is not None: + freqs, xpos_scale = c_rope + q_xpos_scale, k_xpos_scale = ( + (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + ) + c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale) + c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale) + + # joint attention + query = torch.cat([query, c_query], dim=2) + key = torch.cat([key, c_key], dim=2) + value = torch.cat([value, c_value], dim=2) + + # mask. e.g. inference got a batch with different target durations, mask out the padding + if mask is not None: + attn_mask = F.pad(mask, (0, c.shape[1]), value=True) # no mask for c (text) + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n' + attn_mask = attn_mask.expand( + batch_size, attn.heads, query.shape[-2], key.shape[-2] + ) + else: + attn_mask = None + + x = F.scaled_dot_product_attention( + query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False + ) + x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + x = x.to(query.dtype) + + # Split the attention outputs. + x, c = ( + x[:, : residual.shape[1]], + x[:, residual.shape[1] :], + ) + + # linear proj + x = attn.to_out[0](x) + # dropout + x = attn.to_out[1](x) + if not attn.context_pre_only: + c = attn.to_out_c(c) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + # c = c.masked_fill(~mask, 0.) # no mask for c (text) + + return x, c + + +# DiT Block + + +class DiTBlock(nn.Module): + def __init__( + self, + dim, + heads, + dim_head, + ff_mult=4, + dropout=0.1, + qk_norm=None, + pe_attn_head=None, + attn_backend="torch", # "torch" or "flash_attn" + attn_mask_enabled=True, + ): + super().__init__() + + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention( + processor=AttnProcessor( + pe_attn_head=pe_attn_head, + attn_backend=attn_backend, + attn_mask_enabled=attn_mask_enabled, + ), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + qk_norm=qk_norm, + ) + + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward( + dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh" + ) + + def forward(self, x, t, mask=None, rope=None): # x: noised input, t: time embedding + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # attention + attn_output = self.attn(x=norm, mask=mask, rope=rope) + + # process attention output for input x + x = x + gate_msa.unsqueeze(1) * attn_output + + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x + + +# MMDiT Block https://arxiv.org/abs/2403.03206 + + +class MMDiTBlock(nn.Module): + r""" + modified from diffusers/src/diffusers/models/attention.py + + notes. + _c: context related. text, cond, etc. (left part in sd3 fig2.b) + _x: noised input related. (right part) + context_pre_only: last layer only do prenorm + modulation cuz no more ffn + """ + + def __init__( + self, + dim, + heads, + dim_head, + ff_mult=4, + dropout=0.1, + context_dim=None, + context_pre_only=False, + qk_norm=None, + ): + super().__init__() + if context_dim is None: + context_dim = dim + self.context_pre_only = context_pre_only + + self.attn_norm_c = ( + AdaLayerNorm_Final(context_dim) + if context_pre_only + else AdaLayerNorm(context_dim) + ) + self.attn_norm_x = AdaLayerNorm(dim) + self.attn = Attention( + processor=JointAttnProcessor(), + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + context_dim=context_dim, + context_pre_only=context_pre_only, + qk_norm=qk_norm, + ) + + if not context_pre_only: + self.ff_norm_c = nn.LayerNorm( + context_dim, elementwise_affine=False, eps=1e-6 + ) + self.ff_c = FeedForward( + dim=context_dim, mult=ff_mult, dropout=dropout, approximate="tanh" + ) + else: + self.ff_norm_c = None + self.ff_c = None + self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_x = FeedForward( + dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh" + ) + + def forward( + self, x, c, t, mask=None, rope=None, c_rope=None + ): # x: noised input, c: context, t: time embedding + # pre-norm & modulation for attention input + if self.context_pre_only: + norm_c = self.attn_norm_c(c, t) + else: + norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c( + c, emb=t + ) + norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x( + x, emb=t + ) + + # attention + x_attn_output, c_attn_output = self.attn( + x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope + ) + + # process attention output for context c + if self.context_pre_only: + c = None + else: # if not last layer + c = c + c_gate_msa.unsqueeze(1) * c_attn_output + + norm_c = ( + self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + ) + c_ff_output = self.ff_c(norm_c) + c = c + c_gate_mlp.unsqueeze(1) * c_ff_output + + # process attention output for input x + x = x + x_gate_msa.unsqueeze(1) * x_attn_output + + norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None] + x_ff_output = self.ff_x(norm_x) + x = x + x_gate_mlp.unsqueeze(1) * x_ff_output + + return c, x + + +# time step conditioning embedding + + +# class TimestepEmbedding(nn.Module): +# def __init__(self, dim, freq_embed_dim=256): +# super().__init__() +# self.time_embed = SinusPositionEmbedding(freq_embed_dim) +# self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + +# def forward(self, timestep: float["b"]): +# time_hidden = self.time_embed(timestep) +# time_hidden = time_hidden.to(timestep.dtype) +# time = self.time_mlp(time_hidden) # b d +# return time + + +def zipvoice_timestep_embedding(timesteps, dim, max_period=10000): + """Create sinusoidal timestep embeddings. + + :param timesteps: shape of (N) or (N, T) + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an Tensor of positional embeddings. shape of (N, dim) or (T, N, dim) + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) + / half + ) + + if timesteps.dim() == 2: + timesteps = timesteps.transpose(0, 1) # (N, T) -> (T, N) + + args = timesteps[..., None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[..., :1])], dim=-1) + return embedding + + +def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear: + """ + Behaves like a constructor of a modified version of nn.Linear + that gives an easy way to set the default initial parameter scale. + + Args: + Accepts the standard args and kwargs that nn.Linear accepts + e.g. in_features, out_features, bias=False. + + initial_scale: you can override this if you want to increase + or decrease the initial magnitude of the module's output + (affects the initialization of weight_scale and bias_scale). + Another option, if you want to do something like this, is + to re-initialize the parameters. + """ + ans = nn.Linear(*args, **kwargs) + with torch.no_grad(): + ans.weight[:] *= initial_scale + if ans.bias is not None: + torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale) + return ans + + +# 在蒸馏的时候使用！ +class TimestepGuidanceEmbedding(nn.Module): + def __init__( + self, + dim, + freq_embed_dim=256, + use_guidance_scale_embed=False, + guidance_scale_embed_dim=192, + ): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential( + nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim) + ) + if use_guidance_scale_embed: + self.guidance_scale_embed = ScaledLinear( + guidance_scale_embed_dim, + freq_embed_dim, + bias=False, + initial_scale=0.1, + ) + self.guidance_scale_embed_dim = guidance_scale_embed_dim + else: + self.guidance_scale_embed = None + + def forward(self, timestep: float["b"], guidance_scale=None): + # import pdb + + # pdb.set_trace() + time_hidden = self.time_embed(timestep) + + if self.guidance_scale_embed: + assert guidance_scale is not None + guidance_scale_emb = self.guidance_scale_embed( + zipvoice_timestep_embedding( + guidance_scale, self.guidance_scale_embed_dim + ) + ) + time_hidden = time_hidden + guidance_scale_emb + else: + assert guidance_scale is None + + time_hidden = time_hidden.to(timestep.dtype) + time = self.time_mlp(time_hidden) # b d + return time diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/__init__.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..993ec9a5f226e59c6ff3c91b0016b5b626e0f049 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/__init__.py @@ -0,0 +1,91 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import json +import re + +from tokenizers import Tokenizer + +from src.YingMusicSinger.utils.f5_tts.g2p.g2p import cleaners +from src.YingMusicSinger.utils.f5_tts.g2p.g2p.text_tokenizers import TextTokenizer + +# import LangSegment +from src.YingMusicSinger.utils.f5_tts.thirdparty.LangSegment import LangSegment + + +class PhonemeBpeTokenizer: + def __init__( + self, vacab_path="./src/YingMusicSinger/utils/f5_tts/g2p/g2p/vocab.json" + ): + self.lang2backend = { + "zh": "cmn", + "ja": "ja", + "en": "en-us", + "fr": "fr-fr", + "ko": "ko", + "de": "de", + } + self.text_tokenizers = {} + self.int_text_tokenizers() + + with open(vacab_path, "r") as f: + json_data = f.read() + data = json.loads(json_data) + self.vocab = data["vocab"] + LangSegment.setfilters(["en", "zh", "ja", "ko", "fr", "de"]) + + def int_text_tokenizers(self): + for key, value in self.lang2backend.items(): + self.text_tokenizers[key] = TextTokenizer(language=value) + + def tokenize(self, text, sentence, language): + # 1. convert text to phoneme + phonemes = [] + if language == "auto": + seglist = LangSegment.getTexts(text) + tmp_ph = [] + for seg in seglist: + tmp_ph.append( + self._clean_text( + seg["text"], sentence, seg["lang"], ["cjekfd_cleaners"] + ) + ) + phonemes = "|_|".join(tmp_ph) + else: + phonemes = self._clean_text(text, sentence, language, ["cjekfd_cleaners"]) + # print('clean text: ', phonemes) + + # 2. tokenize phonemes + phoneme_tokens = self.phoneme2token(phonemes) + # print('encode: ', phoneme_tokens) + + # # 3. decode tokens [optional] + # decoded_text = self.tokenizer.decode(phoneme_tokens) + # print('decoded: ', decoded_text) + + return phonemes, phoneme_tokens + + def _clean_text(self, text, sentence, language, cleaner_names): + for name in cleaner_names: + cleaner = getattr(cleaners, name) + if not cleaner: + raise Exception("Unknown cleaner: %s" % name) + text = cleaner(text, sentence, language, self.text_tokenizers) + return text + + def phoneme2token(self, phonemes): + tokens = [] + if isinstance(phonemes, list): + for phone in phonemes: + phone = phone.split("\t")[0] + phonemes_split = phone.split("|") + tokens.append( + [self.vocab[p] for p in phonemes_split if p in self.vocab] + ) + else: + phonemes = phonemes.split("\t")[0] + phonemes_split = phonemes.split("|") + tokens = [self.vocab[p] for p in phonemes_split if p in self.vocab] + return tokens diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/chinese_model_g2p.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/chinese_model_g2p.py new file mode 100644 index 0000000000000000000000000000000000000000..1866557b7d0553c7053be68590b0d171687ea654 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/chinese_model_g2p.py @@ -0,0 +1,209 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import json +import os + +import numpy as np +import torch +from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions +from torch.utils.data import DataLoader, Dataset +from transformers import BertTokenizer +from transformers.models.bert.modeling_bert import * + + +class PolyDataset(Dataset): + def __init__(self, words, labels, word_pad_idx=0, label_pad_idx=-1): + self.dataset = self.preprocess(words, labels) + self.word_pad_idx = word_pad_idx + self.label_pad_idx = label_pad_idx + + def preprocess(self, origin_sentences, origin_labels): + """ + Maps tokens and tags to their indices and stores them in the dict data. + examples: + word:['[CLS]', '浙', '商', '银', '行', '企', '业', '信', '贷', '部'] + sentence:([101, 3851, 1555, 7213, 6121, 821, 689, 928, 6587, 6956], + array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])) + label:[3, 13, 13, 13, 0, 0, 0, 0, 0] + """ + data = [] + labels = [] + sentences = [] + # tokenize + for line in origin_sentences: + # replace each token by its index + # we can not use encode_plus because our sentences are aligned to labels in list type + words = [] + word_lens = [] + for token in line: + words.append(token) + word_lens.append(1) + token_start_idxs = 1 + np.cumsum([0] + word_lens[:-1]) + sentences.append(((words, token_start_idxs), 0)) + ### + for tag in origin_labels: + labels.append(tag) + + for sentence, label in zip(sentences, labels): + data.append((sentence, label)) + return data + + def __getitem__(self, idx): + """sample data to get batch""" + word = self.dataset[idx][0] + label = self.dataset[idx][1] + return [word, label] + + def __len__(self): + """get dataset size""" + return len(self.dataset) + + def collate_fn(self, batch): + sentences = [x[0][0] for x in batch] + ori_sents = [x[0][1] for x in batch] + labels = [x[1] for x in batch] + batch_len = len(sentences) + + # compute length of longest sentence in batch + max_len = max([len(s[0]) for s in sentences]) + max_label_len = 0 + batch_data = np.ones((batch_len, max_len)) + batch_label_starts = [] + + # padding and aligning + for j in range(batch_len): + cur_len = len(sentences[j][0]) + batch_data[j][:cur_len] = sentences[j][0] + label_start_idx = sentences[j][-1] + label_starts = np.zeros(max_len) + label_starts[[idx for idx in label_start_idx if idx < max_len]] = 1 + batch_label_starts.append(label_starts) + max_label_len = max(int(sum(label_starts)), max_label_len) + + # padding label + batch_labels = self.label_pad_idx * np.ones((batch_len, max_label_len)) + batch_pmasks = self.label_pad_idx * np.ones((batch_len, max_label_len)) + for j in range(batch_len): + cur_tags_len = len(labels[j]) + batch_labels[j][:cur_tags_len] = labels[j] + batch_pmasks[j][:cur_tags_len] = [ + 1 if item > 0 else 0 for item in labels[j] + ] + + # convert data to torch LongTensors + batch_data = torch.tensor(batch_data, dtype=torch.long) + batch_label_starts = torch.tensor(batch_label_starts, dtype=torch.long) + batch_labels = torch.tensor(batch_labels, dtype=torch.long) + batch_pmasks = torch.tensor(batch_pmasks, dtype=torch.long) + return [batch_data, batch_label_starts, batch_labels, batch_pmasks, ori_sents] + + +class BertPolyPredict: + def __init__(self, bert_model, jsonr_file, json_file): + self.tokenizer = BertTokenizer.from_pretrained(bert_model, do_lower_case=True) + with open(jsonr_file, "r", encoding="utf8") as fp: + self.pron_dict = json.load(fp) + with open(json_file, "r", encoding="utf8") as fp: + self.pron_dict_id_2_pinyin = json.load(fp) + self.num_polyphone = len(self.pron_dict) + self.device = "cpu" + self.polydataset = PolyDataset + options = SessionOptions() # initialize session options + options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL + print(os.path.join(bert_model, "poly_bert_model.onnx")) + self.session = InferenceSession( + os.path.join(bert_model, "poly_bert_model.onnx"), + sess_options=options, + providers=[ + "CPUExecutionProvider", + "CUDAExecutionProvider", + ], # CPUExecutionProvider #CUDAExecutionProvider + ) + # self.session.set_providers(['CUDAExecutionProvider', "CPUExecutionProvider"], [ {'device_id': 0}]) + + # disable session.run() fallback mechanism, it prevents for a reset of the execution provider + self.session.disable_fallback() + + def predict_process(self, txt_list): + word_test, label_test, texts_test = self.get_examples_po(txt_list) + data = self.polydataset(word_test, label_test) + predict_loader = DataLoader( + data, batch_size=1, shuffle=False, collate_fn=data.collate_fn + ) + pred_tags = self.predict_onnx(predict_loader) + return pred_tags + + def predict_onnx(self, dev_loader): + pred_tags = [] + with torch.no_grad(): + for idx, batch_samples in enumerate(dev_loader): + # [batch_data, batch_label_starts, batch_labels, batch_pmasks, ori_sents] + batch_data, batch_label_starts, batch_labels, batch_pmasks, _ = ( + batch_samples + ) + # shift tensors to GPU if available + batch_data = batch_data.to(self.device) + batch_label_starts = batch_label_starts.to(self.device) + batch_labels = batch_labels.to(self.device) + batch_pmasks = batch_pmasks.to(self.device) + batch_data = np.asarray(batch_data, dtype=np.int32) + batch_pmasks = np.asarray(batch_pmasks, dtype=np.int32) + # batch_output = self.session.run(output_names=['outputs'], input_feed={"input_ids":batch_data, "input_pmasks": batch_pmasks})[0][0] + batch_output = self.session.run( + output_names=["outputs"], input_feed={"input_ids": batch_data} + )[0] + label_masks = batch_pmasks == 1 + batch_labels = batch_labels.to("cpu").numpy() + for i, indices in enumerate(np.argmax(batch_output, axis=2)): + for j, idx in enumerate(indices): + if label_masks[i][j]: + # pred_tag.append(idx) + pred_tags.append(self.pron_dict_id_2_pinyin[str(idx + 1)]) + return pred_tags + + def get_examples_po(self, text_list): + word_list = [] + label_list = [] + sentence_list = [] + id = 0 + for line in [text_list]: + sentence = line[0] + words = [] + tokens = line[0] + index = line[-1] + front = index + back = len(tokens) - index - 1 + labels = [0] * front + [1] + [0] * back + words = ["[CLS]"] + [item for item in sentence] + words = self.tokenizer.convert_tokens_to_ids(words) + word_list.append(words) + label_list.append(labels) + sentence_list.append(sentence) + + id += 1 + # mask_list.append(masks) + assert len(labels) + 1 == len(words), print( + ( + poly, + sentence, + words, + labels, + sentence, + len(sentence), + len(words), + len(labels), + ) + ) + assert len(labels) + 1 == len(words), ( + "Number of labels does not match number of words" + ) + assert len(labels) == len(sentence), ( + "Number of labels does not match number of sentences" + ) + assert len(word_list) == len(label_list), ( + "Number of label sentences does not match number of word sentences" + ) + return word_list, label_list, text_list diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/cleaners.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/cleaners.py new file mode 100644 index 0000000000000000000000000000000000000000..a7fdded53fecda9cac89688f04037b7b00bdaf97 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/cleaners.py @@ -0,0 +1,28 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from src.YingMusicSinger.utils.f5_tts.g2p.g2p.english import english_to_ipa +from src.YingMusicSinger.utils.f5_tts.g2p.g2p.french import french_to_ipa +from src.YingMusicSinger.utils.f5_tts.g2p.g2p.german import german_to_ipa +from src.YingMusicSinger.utils.f5_tts.g2p.g2p.korean import korean_to_ipa +from src.YingMusicSinger.utils.f5_tts.g2p.g2p.mandarin import chinese_to_ipa + + +def cjekfd_cleaners(text, sentence, language, text_tokenizers): + if language == "zh": + return chinese_to_ipa(text, sentence, text_tokenizers["zh"]) + elif language == "ja": + return japanese_to_ipa(text, text_tokenizers["ja"]) + elif language == "en": + return english_to_ipa(text, text_tokenizers["en"]) + elif language == "fr": + return french_to_ipa(text, text_tokenizers["fr"]) + elif language == "ko": + return korean_to_ipa(text, text_tokenizers["ko"]) + elif language == "de": + return german_to_ipa(text, text_tokenizers["de"]) + else: + raise Exception("Unknown language: %s" % language) + return None diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/english.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/english.py new file mode 100644 index 0000000000000000000000000000000000000000..1dbee6beaab0edac241fcdbfba2c6a53eef90b29 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/english.py @@ -0,0 +1,202 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re + +import inflect + +""" + Text clean time +""" +_inflect = inflect.engine() +_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])") +_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)") +_percent_number_re = re.compile(r"([0-9\.\,]*[0-9]+%)") +_pounds_re = re.compile(r"£([0-9\,]*[0-9]+)") +_dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)") +_fraction_re = re.compile(r"([0-9]+)/([0-9]+)") +_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)") +_number_re = re.compile(r"[0-9]+") + +# List of (regular expression, replacement) pairs for abbreviations: +_abbreviations = [ + (re.compile("\\b%s\\b" % x[0], re.IGNORECASE), x[1]) + for x in [ + ("mrs", "misess"), + ("mr", "mister"), + ("dr", "doctor"), + ("st", "saint"), + ("co", "company"), + ("jr", "junior"), + ("maj", "major"), + ("gen", "general"), + ("drs", "doctors"), + ("rev", "reverend"), + ("lt", "lieutenant"), + ("hon", "honorable"), + ("sgt", "sergeant"), + ("capt", "captain"), + ("esq", "esquire"), + ("ltd", "limited"), + ("col", "colonel"), + ("ft", "fort"), + ("etc", "et cetera"), + ("btw", "by the way"), + ] +] + +_special_map = [ + ("t|ɹ", "tɹ"), + ("d|ɹ", "dɹ"), + ("t|s", "ts"), + ("d|z", "dz"), + ("ɪ|ɹ", "ɪɹ"), + ("ɐ", "ɚ"), + ("ᵻ", "ɪ"), + ("əl", "l"), + ("x", "k"), + ("ɬ", "l"), + ("ʔ", "t"), + ("n̩", "n"), + ("oː|ɹ", "oːɹ"), +] + + +def expand_abbreviations(text): + for regex, replacement in _abbreviations: + text = re.sub(regex, replacement, text) + return text + + +def _remove_commas(m): + return m.group(1).replace(",", "") + + +def _expand_decimal_point(m): + return m.group(1).replace(".", " point ") + + +def _expand_percent(m): + return m.group(1).replace("%", " percent ") + + +def _expand_dollars(m): + match = m.group(1) + parts = match.split(".") + if len(parts) > 2: + return " " + match + " dollars " # Unexpected format + dollars = int(parts[0]) if parts[0] else 0 + cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0 + if dollars and cents: + dollar_unit = "dollar" if dollars == 1 else "dollars" + cent_unit = "cent" if cents == 1 else "cents" + return " %s %s, %s %s " % (dollars, dollar_unit, cents, cent_unit) + elif dollars: + dollar_unit = "dollar" if dollars == 1 else "dollars" + return " %s %s " % (dollars, dollar_unit) + elif cents: + cent_unit = "cent" if cents == 1 else "cents" + return " %s %s " % (cents, cent_unit) + else: + return " zero dollars " + + +def fraction_to_words(numerator, denominator): + if numerator == 1 and denominator == 2: + return " one half " + if numerator == 1 and denominator == 4: + return " one quarter " + if denominator == 2: + return " " + _inflect.number_to_words(numerator) + " halves " + if denominator == 4: + return " " + _inflect.number_to_words(numerator) + " quarters " + return ( + " " + + _inflect.number_to_words(numerator) + + " " + + _inflect.ordinal(_inflect.number_to_words(denominator)) + + " " + ) + + +def _expand_fraction(m): + numerator = int(m.group(1)) + denominator = int(m.group(2)) + return fraction_to_words(numerator, denominator) + + +def _expand_ordinal(m): + return " " + _inflect.number_to_words(m.group(0)) + " " + + +def _expand_number(m): + num = int(m.group(0)) + if num > 1000 and num < 3000: + if num == 2000: + return " two thousand " + elif num > 2000 and num < 2010: + return " two thousand " + _inflect.number_to_words(num % 100) + " " + elif num % 100 == 0: + return " " + _inflect.number_to_words(num // 100) + " hundred " + else: + return ( + " " + + _inflect.number_to_words(num, andword="", zero="oh", group=2).replace( + ", ", " " + ) + + " " + ) + else: + return " " + _inflect.number_to_words(num, andword="") + " " + + +# Normalize numbers pronunciation +def normalize_numbers(text): + text = re.sub(_comma_number_re, _remove_commas, text) + text = re.sub(_pounds_re, r"\1 pounds", text) + text = re.sub(_dollars_re, _expand_dollars, text) + text = re.sub(_fraction_re, _expand_fraction, text) + text = re.sub(_decimal_number_re, _expand_decimal_point, text) + text = re.sub(_percent_number_re, _expand_percent, text) + text = re.sub(_ordinal_re, _expand_ordinal, text) + text = re.sub(_number_re, _expand_number, text) + return text + + +def _english_to_ipa(text): + # text = unidecode(text).lower() + text = expand_abbreviations(text) + text = normalize_numbers(text) + return text + + +# special map +def special_map(text): + for regex, replacement in _special_map: + regex = regex.replace("|", "\|") + while re.search(r"(^|[_|]){}([_|]|$)".format(regex), text): + text = re.sub( + r"(^|[_|]){}([_|]|$)".format(regex), r"\1{}\2".format(replacement), text + ) + # text = re.sub(r'([,.!?])', r'|\1', text) + return text + + +# Add some special operation +def english_to_ipa(text, text_tokenizer): + if type(text) == str: + text = _english_to_ipa(text) + else: + text = [_english_to_ipa(t) for t in text] + phonemes = text_tokenizer(text) + if phonemes[-1] in "p⁼ʰmftnlkxʃs`ɹaoəɛɪeɑʊŋiuɥwæjː": + phonemes += "|_" + if type(text) == str: + return special_map(phonemes) + else: + result_ph = [] + for phone in phonemes: + result_ph.append(special_map(phone)) + return result_ph diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/french.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/french.py new file mode 100644 index 0000000000000000000000000000000000000000..bd9400cdfc6598e7d642480cbfc1f990fc78cddf --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/french.py @@ -0,0 +1,149 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re + +""" + Text clean time +""" +# List of (regular expression, replacement) pairs for abbreviations in french: +_abbreviations = [ + (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1]) + for x in [ + ("M", "monsieur"), + ("Mlle", "mademoiselle"), + ("Mlles", "mesdemoiselles"), + ("Mme", "Madame"), + ("Mmes", "Mesdames"), + ("N.B", "nota bene"), + ("M", "monsieur"), + ("p.c.q", "parce que"), + ("Pr", "professeur"), + ("qqch", "quelque chose"), + ("rdv", "rendez-vous"), + ("max", "maximum"), + ("min", "minimum"), + ("no", "numéro"), + ("adr", "adresse"), + ("dr", "docteur"), + ("st", "saint"), + ("co", "companie"), + ("jr", "junior"), + ("sgt", "sergent"), + ("capt", "capitain"), + ("col", "colonel"), + ("av", "avenue"), + ("av. J.-C", "avant Jésus-Christ"), + ("apr. J.-C", "après Jésus-Christ"), + ("art", "article"), + ("boul", "boulevard"), + ("c.-à-d", "c’est-à-dire"), + ("etc", "et cetera"), + ("ex", "exemple"), + ("excl", "exclusivement"), + ("boul", "boulevard"), + ] +] + [ + (re.compile("\\b%s" % x[0]), x[1]) + for x in [ + ("Mlle", "mademoiselle"), + ("Mlles", "mesdemoiselles"), + ("Mme", "Madame"), + ("Mmes", "Mesdames"), + ] +] + +rep_map = { + "：": ",", + "；": ",", + "，": ",", + "。": ".", + "！": "!", + "？": "?", + "\n": ".", + "·": ",", + "、": ",", + "...": ".", + "…": ".", + "$": ".", + "“": "", + "”": "", + "‘": "", + "’": "", + "（": "", + "）": "", + "(": "", + ")": "", + "《": "", + "》": "", + "【": "", + "】": "", + "[": "", + "]": "", + "—": "", + "～": "-", + "~": "-", + "「": "", + "」": "", + "¿": "", + "¡": "", +} + + +def collapse_whitespace(text): + # Regular expression matching whitespace: + _whitespace_re = re.compile(r"\s+") + return re.sub(_whitespace_re, " ", text).strip() + + +def remove_punctuation_at_begin(text): + return re.sub(r"^[,.!?]+", "", text) + + +def remove_aux_symbols(text): + text = re.sub(r"[\<\>\(\)\[\]\"\«\»]+", "", text) + return text + + +def replace_symbols(text): + text = text.replace(";", ",") + text = text.replace("-", " ") + text = text.replace(":", ",") + text = text.replace("&", " et ") + return text + + +def expand_abbreviations(text): + for regex, replacement in _abbreviations: + text = re.sub(regex, replacement, text) + return text + + +def replace_punctuation(text): + pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys())) + replaced_text = pattern.sub(lambda x: rep_map[x.group()], text) + return replaced_text + + +def text_normalize(text): + text = expand_abbreviations(text) + text = replace_punctuation(text) + text = replace_symbols(text) + text = remove_aux_symbols(text) + text = remove_punctuation_at_begin(text) + text = collapse_whitespace(text) + text = re.sub(r"([^\.,!\?\-…])$", r"\1", text) + return text + + +def french_to_ipa(text, text_tokenizer): + if type(text) == str: + text = text_normalize(text) + phonemes = text_tokenizer(text) + return phonemes + else: + for i, t in enumerate(text): + text[i] = text_normalize(t) + return text_tokenizer(text) diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/german.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/german.py new file mode 100644 index 0000000000000000000000000000000000000000..bd82eeabc44cc891acd98daa982cd2be1e991e3a --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/german.py @@ -0,0 +1,94 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re + +""" + Text clean time +""" +rep_map = { + "：": ",", + "；": ",", + "，": ",", + "。": ".", + "！": "!", + "？": "?", + "\n": ".", + "·": ",", + "、": ",", + "...": ".", + "…": ".", + "$": ".", + "“": "", + "”": "", + "‘": "", + "’": "", + "（": "", + "）": "", + "(": "", + ")": "", + "《": "", + "》": "", + "【": "", + "】": "", + "[": "", + "]": "", + "—": "", + "～": "-", + "~": "-", + "「": "", + "」": "", + "¿": "", + "¡": "", +} + + +def collapse_whitespace(text): + # Regular expression matching whitespace: + _whitespace_re = re.compile(r"\s+") + return re.sub(_whitespace_re, " ", text).strip() + + +def remove_punctuation_at_begin(text): + return re.sub(r"^[,.!?]+", "", text) + + +def remove_aux_symbols(text): + text = re.sub(r"[\<\>\(\)\[\]\"\«\»]+", "", text) + return text + + +def replace_symbols(text): + text = text.replace(";", ",") + text = text.replace("-", " ") + text = text.replace(":", ",") + return text + + +def replace_punctuation(text): + pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys())) + replaced_text = pattern.sub(lambda x: rep_map[x.group()], text) + return replaced_text + + +def text_normalize(text): + text = replace_punctuation(text) + text = replace_symbols(text) + text = remove_aux_symbols(text) + text = remove_punctuation_at_begin(text) + text = collapse_whitespace(text) + text = re.sub(r"([^\.,!\?\-…])$", r"\1", text) + return text + + +def german_to_ipa(text, text_tokenizer): + if type(text) == str: + text = text_normalize(text) + phonemes = text_tokenizer(text) + return phonemes + else: + for i, t in enumerate(text): + text[i] = text_normalize(t) + return text_tokenizer(text) diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/korean.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/korean.py new file mode 100644 index 0000000000000000000000000000000000000000..c7c540b47d98ccf6e0db5f938e52834abf679b59 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/korean.py @@ -0,0 +1,81 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re + +""" + Text clean time +""" +english_dictionary = { + "KOREA": "코리아", + "IDOL": "아이돌", + "IT": "아이티", + "IQ": "아이큐", + "UP": "업", + "DOWN": "다운", + "PC": "피씨", + "CCTV": "씨씨티비", + "SNS": "에스엔에스", + "AI": "에이아이", + "CEO": "씨이오", + "A": "에이", + "B": "비", + "C": "씨", + "D": "디", + "E": "이", + "F": "에프", + "G": "지", + "H": "에이치", + "I": "아이", + "J": "제이", + "K": "케이", + "L": "엘", + "M": "엠", + "N": "엔", + "O": "오", + "P": "피", + "Q": "큐", + "R": "알", + "S": "에스", + "T": "티", + "U": "유", + "V": "브이", + "W": "더블유", + "X": "엑스", + "Y": "와이", + "Z": "제트", +} + + +def normalize(text): + text = text.strip() + text = re.sub( + "[⺀-⺙⺛-⻳⼀-⿕々〇〡-〩〸-〺〻㐀-䶵一-鿃豈-鶴侮-頻並-龎]", "", text + ) + text = normalize_english(text) + text = text.lower() + return text + + +def normalize_english(text): + def fn(m): + word = m.group() + if word in english_dictionary: + return english_dictionary.get(word) + return word + + text = re.sub("([A-Za-z]+)", fn, text) + return text + + +def korean_to_ipa(text, text_tokenizer): + if type(text) == str: + text = normalize(text) + phonemes = text_tokenizer(text) + return phonemes + else: + for i, t in enumerate(text): + text[i] = normalize(t) + return text_tokenizer(text) diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/mandarin.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/mandarin.py new file mode 100644 index 0000000000000000000000000000000000000000..caf57d42216cb1fafaf8f0adaa429472f94d9a29 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/mandarin.py @@ -0,0 +1,603 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import re +from typing import List + +import cn2an +import jieba +from pypinyin import BOPOMOFO, lazy_pinyin + +from src.YingMusicSinger.utils.f5_tts.g2p.g2p.chinese_model_g2p import BertPolyPredict +from src.YingMusicSinger.utils.f5_tts.g2p.utils.front_utils import * + +# from g2pw import G2PWConverter + + +# set blank level, {0："none",1:"char", 2:"word"} +BLANK_LEVEL = 0 + +# conv = G2PWConverter(style='pinyin', enable_non_tradional_chinese=True) +resource_path = r"./src/YingMusicSinger/utils/f5_tts/g2p" +poly_all_class_path = os.path.join( + resource_path, "sources", "g2p_chinese_model", "polychar.txt" +) +if not os.path.exists(poly_all_class_path): + print( + "Incorrect path for polyphonic character class dictionary: {}, please check...".format( + poly_all_class_path + ) + ) + exit() +poly_dict = generate_poly_lexicon(poly_all_class_path) + +# Set up G2PW model parameters +g2pw_poly_model_path = os.path.join(resource_path, "sources", "g2p_chinese_model") +if not os.path.exists(g2pw_poly_model_path): + print( + "Incorrect path for g2pw polyphonic character model: {}, please check...".format( + g2pw_poly_model_path + ) + ) + exit() + +json_file_path = os.path.join( + resource_path, "sources", "g2p_chinese_model", "polydict.json" +) +if not os.path.exists(json_file_path): + print( + "Incorrect path for g2pw id to pinyin dictionary: {}, please check...".format( + json_file_path + ) + ) + exit() + +jsonr_file_path = os.path.join( + resource_path, "sources", "g2p_chinese_model", "polydict_r.json" +) +if not os.path.exists(jsonr_file_path): + print( + "Incorrect path for g2pw pinyin to id dictionary: {}, please check...".format( + jsonr_file_path + ) + ) + exit() + +g2pw_poly_predict = BertPolyPredict( + g2pw_poly_model_path, jsonr_file_path, json_file_path +) + + +""" + Text clean time +""" +# List of (Latin alphabet, bopomofo) pairs: +_latin_to_bopomofo = [ + (re.compile("%s" % x[0], re.IGNORECASE), x[1]) + for x in [ + ("a", "ㄟˉ"), + ("b", "ㄅㄧˋ"), + ("c", "ㄙㄧˉ"), + ("d", "ㄉㄧˋ"), + ("e", "ㄧˋ"), + ("f", "ㄝˊㄈㄨˋ"), + ("g", "ㄐㄧˋ"), + ("h", "ㄝˇㄑㄩˋ"), + ("i", "ㄞˋ"), + ("j", "ㄐㄟˋ"), + ("k", "ㄎㄟˋ"), + ("l", "ㄝˊㄛˋ"), + ("m", "ㄝˊㄇㄨˋ"), + ("n", "ㄣˉ"), + ("o", "ㄡˉ"), + ("p", "ㄆㄧˉ"), + ("q", "ㄎㄧㄡˉ"), + ("r", "ㄚˋ"), + ("s", "ㄝˊㄙˋ"), + ("t", "ㄊㄧˋ"), + ("u", "ㄧㄡˉ"), + ("v", "ㄨㄧˉ"), + ("w", "ㄉㄚˋㄅㄨˋㄌㄧㄡˋ"), + ("x", "ㄝˉㄎㄨˋㄙˋ"), + ("y", "ㄨㄞˋ"), + ("z", "ㄗㄟˋ"), + ] +] + +# List of (bopomofo, ipa) pairs: +_bopomofo_to_ipa = [ + (re.compile("%s" % x[0]), x[1]) + for x in [ + ("ㄅㄛ", "p⁼wo"), + ("ㄆㄛ", "pʰwo"), + ("ㄇㄛ", "mwo"), + ("ㄈㄛ", "fwo"), + ("ㄧㄢ", "|jɛn"), + ("ㄩㄢ", "|ɥæn"), + ("ㄧㄣ", "|in"), + ("ㄩㄣ", "|ɥn"), + ("ㄧㄥ", "|iŋ"), + ("ㄨㄥ", "|ʊŋ"), + ("ㄩㄥ", "|jʊŋ"), + # Add + ("ㄧㄚ", "|ia"), + ("ㄧㄝ", "|iɛ"), + ("ㄧㄠ", "|iɑʊ"), + ("ㄧㄡ", "|ioʊ"), + ("ㄧㄤ", "|iɑŋ"), + ("ㄨㄚ", "|ua"), + ("ㄨㄛ", "|uo"), + ("ㄨㄞ", "|uaɪ"), + ("ㄨㄟ", "|ueɪ"), + ("ㄨㄢ", "|uan"), + ("ㄨㄣ", "|uən"), + ("ㄨㄤ", "|uɑŋ"), + ("ㄩㄝ", "|ɥɛ"), + # End + ("ㄅ", "p⁼"), + ("ㄆ", "pʰ"), + ("ㄇ", "m"), + ("ㄈ", "f"), + ("ㄉ", "t⁼"), + ("ㄊ", "tʰ"), + ("ㄋ", "n"), + ("ㄌ", "l"), + ("ㄍ", "k⁼"), + ("ㄎ", "kʰ"), + ("ㄏ", "x"), + ("ㄐ", "tʃ⁼"), + ("ㄑ", "tʃʰ"), + ("ㄒ", "ʃ"), + ("ㄓ", "ts`⁼"), + ("ㄔ", "ts`ʰ"), + ("ㄕ", "s`"), + ("ㄖ", "ɹ`"), + ("ㄗ", "ts⁼"), + ("ㄘ", "tsʰ"), + ("ㄙ", "|s"), + ("ㄚ", "|a"), + ("ㄛ", "|o"), + ("ㄜ", "|ə"), + ("ㄝ", "|ɛ"), + ("ㄞ", "|aɪ"), + ("ㄟ", "|eɪ"), + ("ㄠ", "|ɑʊ"), + ("ㄡ", "|oʊ"), + ("ㄢ", "|an"), + ("ㄣ", "|ən"), + ("ㄤ", "|ɑŋ"), + ("ㄥ", "|əŋ"), + ("ㄦ", "əɹ"), + ("ㄧ", "|i"), + ("ㄨ", "|u"), + ("ㄩ", "|ɥ"), + ("ˉ", "→|"), + ("ˊ", "↑|"), + ("ˇ", "↓↑|"), + ("ˋ", "↓|"), + ("˙", "|"), + ] +] +must_not_er_words = {"女儿", "老儿", "男儿", "少儿", "小儿"} + +word_pinyin_dict = {} +with open( + r"src/YingMusicSinger/utils/f5_tts/g2p/sources/chinese_lexicon.txt", + "r", + encoding="utf-8", +) as fread: + txt_list = fread.readlines() + for txt in txt_list: + word, pinyin = txt.strip().split("\t") + word_pinyin_dict[word] = pinyin + fread.close() + +pinyin_2_bopomofo_dict = {} +with open( + r"./src/YingMusicSinger/utils/f5_tts/g2p/sources/pinyin_2_bpmf.txt", + "r", + encoding="utf-8", +) as fread: + txt_list = fread.readlines() + for txt in txt_list: + pinyin, bopomofo = txt.strip().split("\t") + pinyin_2_bopomofo_dict[pinyin] = bopomofo + fread.close() + +tone_dict = { + "0": "˙", + "5": "˙", + "1": "", + "2": "ˊ", + "3": "ˇ", + "4": "ˋ", +} + +bopomofos2pinyin_dict = {} +with open( + r"./src/YingMusicSinger/utils/f5_tts/g2p/sources/bpmf_2_pinyin.txt", + "r", + encoding="utf-8", +) as fread: + txt_list = fread.readlines() + for txt in txt_list: + v, k = txt.strip().split("\t") + bopomofos2pinyin_dict[k] = v + fread.close() + + +def bpmf_to_pinyin(text): + bopomofo_list = text.split("|") + pinyin_list = [] + for info in bopomofo_list: + pinyin = "" + for c in info: + if c in bopomofos2pinyin_dict: + pinyin += bopomofos2pinyin_dict[c] + if len(pinyin) == 0: + continue + if pinyin[-1] not in "01234": + pinyin += "1" + if pinyin[:-1] == "ve": + pinyin = "y" + pinyin + if pinyin[:-1] == "sh": + pinyin = pinyin[:-1] + "i" + pinyin[-1] + if pinyin == "sh": + pinyin = pinyin[:-1] + "i" + if pinyin[:-1] == "s": + pinyin = "si" + pinyin[-1] + if pinyin[:-1] == "c": + pinyin = "ci" + pinyin[-1] + if pinyin[:-1] == "i": + pinyin = "yi" + pinyin[-1] + if pinyin[:-1] == "iou": + pinyin = "you" + pinyin[-1] + if pinyin[:-1] == "ien": + pinyin = "yin" + pinyin[-1] + if "iou" in pinyin and pinyin[-4:-1] == "iou": + pinyin = pinyin[:-4] + "iu" + pinyin[-1] + if "uei" in pinyin: + if pinyin[:-1] == "uei": + pinyin = "wei" + pinyin[-1] + elif pinyin[-4:-1] == "uei": + pinyin = pinyin[:-4] + "ui" + pinyin[-1] + if "uen" in pinyin and pinyin[-4:-1] == "uen": + if pinyin[:-1] == "uen": + pinyin = "wen" + pinyin[-1] + elif pinyin[-4:-1] == "uei": + pinyin = pinyin[:-4] + "un" + pinyin[-1] + if "van" in pinyin and pinyin[-4:-1] == "van": + if pinyin[:-1] == "van": + pinyin = "yuan" + pinyin[-1] + elif pinyin[-4:-1] == "van": + pinyin = pinyin[:-4] + "uan" + pinyin[-1] + if "ueng" in pinyin and pinyin[-5:-1] == "ueng": + pinyin = pinyin[:-5] + "ong" + pinyin[-1] + if pinyin[:-1] == "veng": + pinyin = "yong" + pinyin[-1] + if "veng" in pinyin and pinyin[-5:-1] == "veng": + pinyin = pinyin[:-5] + "iong" + pinyin[-1] + if pinyin[:-1] == "ieng": + pinyin = "ying" + pinyin[-1] + if pinyin[:-1] == "u": + pinyin = "wu" + pinyin[-1] + if pinyin[:-1] == "v": + pinyin = "yv" + pinyin[-1] + if pinyin[:-1] == "ing": + pinyin = "ying" + pinyin[-1] + if pinyin[:-1] == "z": + pinyin = "zi" + pinyin[-1] + if pinyin[:-1] == "zh": + pinyin = "zhi" + pinyin[-1] + if pinyin[0] == "u": + pinyin = "w" + pinyin[1:] + if pinyin[0] == "i": + pinyin = "y" + pinyin[1:] + pinyin = pinyin.replace("ien", "in") + + pinyin_list.append(pinyin) + return " ".join(pinyin_list) + + +# Convert numbers to Chinese pronunciation +def number_to_chinese(text): + # numbers = re.findall(r'\d+(?:\.?\d+)?', text) + # for number in numbers: + # text = text.replace(number, cn2an.an2cn(number), 1) + text = cn2an.transform(text, "an2cn") + return text + + +def normalization(text): + text = text.replace("，", ",") + text = text.replace("。", ".") + text = text.replace("！", "!") + text = text.replace("？", "?") + text = text.replace("；", ";") + text = text.replace("：", ":") + text = text.replace("、", ",") + text = text.replace("‘", "'") + text = text.replace("’", "'") + text = text.replace("⋯", "…") + text = text.replace("···", "…") + text = text.replace("・・・", "…") + text = text.replace("...", "…") + text = re.sub(r"\s+", "", text) + text = re.sub(r"[^\u4e00-\u9fff\s_,\.\?!;:\'…]", "", text) + text = re.sub(r"\s*([,\.\?!;:\'…])\s*", r"\1", text) + return text + + +def change_tone(bopomofo: str, tone: str) -> str: + if bopomofo[-1] not in "˙ˊˇˋ": + bopomofo = bopomofo + tone + else: + bopomofo = bopomofo[:-1] + tone + return bopomofo + + +def er_sandhi(word: str, bopomofos: List[str]) -> List[str]: + if len(word) > 1 and word[-1] == "儿" and word not in must_not_er_words: + bopomofos[-1] = change_tone(bopomofos[-1], "˙") + return bopomofos + + +def bu_sandhi(word: str, bopomofos: List[str]) -> List[str]: + valid_char = set(word) + if len(valid_char) == 1 and "不" in valid_char: + pass + elif word in ["不字"]: + pass + elif len(word) == 3 and word[1] == "不" and bopomofos[1][:-1] == "ㄅㄨ": + bopomofos[1] = bopomofos[1][:-1] + "˙" + else: + for i, char in enumerate(word): + if ( + i + 1 < len(bopomofos) + and char == "不" + and i + 1 < len(word) + and 0 < len(bopomofos[i + 1]) + and bopomofos[i + 1][-1] == "ˋ" + ): + bopomofos[i] = bopomofos[i][:-1] + "ˊ" + return bopomofos + + +def yi_sandhi(word: str, bopomofos: List[str]) -> List[str]: + punc = "：，；。？！“”‘’':,;.?!()（）{}【】[]-~`、 " + if word.find("一") != -1 and any( + [item.isnumeric() for item in word if item != "一"] + ): + for i in range(len(word)): + if ( + i == 0 + and word[0] == "一" + and len(word) > 1 + and word[1] + not in [ + "零", + "一", + "二", + "三", + "四", + "五", + "六", + "七", + "八", + "九", + "十", + ] + ): + if len(bopomofos[0]) > 0 and bopomofos[1][-1] in ["ˋ", "˙"]: + bopomofos[0] = change_tone(bopomofos[0], "ˊ") + else: + bopomofos[0] = change_tone(bopomofos[0], "ˋ") + elif word[i] == "一": + bopomofos[i] = change_tone(bopomofos[i], "") + return bopomofos + elif len(word) == 3 and word[1] == "一" and word[0] == word[-1]: + bopomofos[1] = change_tone(bopomofos[1], "˙") + elif word.startswith("第一"): + bopomofos[1] = change_tone(bopomofos[1], "") + elif word.startswith("一月") or word.startswith("一日") or word.startswith("一号"): + bopomofos[0] = change_tone(bopomofos[0], "") + else: + for i, char in enumerate(word): + if char == "一" and i + 1 < len(word): + if ( + len(bopomofos) > i + 1 + and len(bopomofos[i + 1]) > 0 + and bopomofos[i + 1][-1] in {"ˋ"} + ): + bopomofos[i] = change_tone(bopomofos[i], "ˊ") + else: + if word[i + 1] not in punc: + bopomofos[i] = change_tone(bopomofos[i], "ˋ") + else: + pass + return bopomofos + + +def merge_bu(seg: List) -> List: + new_seg = [] + last_word = "" + for word in seg: + if word != "不": + if last_word == "不": + word = last_word + word + new_seg.append(word) + last_word = word + return new_seg + + +def merge_er(seg: List) -> List: + new_seg = [] + for i, word in enumerate(seg): + if i - 1 >= 0 and word == "儿": + new_seg[-1] = new_seg[-1] + seg[i] + else: + new_seg.append(word) + return new_seg + + +def merge_yi(seg: List) -> List: + new_seg = [] + # function 1 + for i, word in enumerate(seg): + if ( + i - 1 >= 0 + and word == "一" + and i + 1 < len(seg) + and seg[i - 1] == seg[i + 1] + ): + if i - 1 < len(new_seg): + new_seg[i - 1] = new_seg[i - 1] + "一" + new_seg[i - 1] + else: + new_seg.append(word) + new_seg.append(seg[i + 1]) + else: + if i - 2 >= 0 and seg[i - 1] == "一" and seg[i - 2] == word: + continue + else: + new_seg.append(word) + seg = new_seg + new_seg = [] + isnumeric_flag = False + for i, word in enumerate(seg): + if all([item.isnumeric() for item in word]) and not isnumeric_flag: + isnumeric_flag = True + new_seg.append(word) + else: + new_seg.append(word) + seg = new_seg + new_seg = [] + # function 2 + for i, word in enumerate(seg): + if new_seg and new_seg[-1] == "一": + new_seg[-1] = new_seg[-1] + word + else: + new_seg.append(word) + return new_seg + + +# Word Segmentation, and convert Chinese pronunciation to pinyin (bopomofo) +def chinese_to_bopomofo(text_short, sentence): + # bopomofos = conv(text_short) + words = jieba.lcut(text_short, cut_all=False) + words = merge_yi(words) + words = merge_bu(words) + words = merge_er(words) + text = "" + + char_index = 0 + for word in words: + bopomofos = [] + if word in word_pinyin_dict and word not in poly_dict: + pinyin = word_pinyin_dict[word] + for py in pinyin.split(" "): + if py[:-1] in pinyin_2_bopomofo_dict and py[-1] in tone_dict: + bopomofos.append( + pinyin_2_bopomofo_dict[py[:-1]] + tone_dict[py[-1]] + ) + if BLANK_LEVEL == 1: + bopomofos.append("_") + else: + bopomofos_lazy = lazy_pinyin(word, BOPOMOFO) + bopomofos += bopomofos_lazy + if BLANK_LEVEL == 1: + bopomofos.append("_") + else: + for i in range(len(word)): + c = word[i] + if c in poly_dict: + poly_pinyin = g2pw_poly_predict.predict_process( + [text_short, char_index + i] + )[0] + py = poly_pinyin[2:-1] + bopomofos.append( + pinyin_2_bopomofo_dict[py[:-1]] + tone_dict[py[-1]] + ) + if BLANK_LEVEL == 1: + bopomofos.append("_") + elif c in word_pinyin_dict: + py = word_pinyin_dict[c] + bopomofos.append( + pinyin_2_bopomofo_dict[py[:-1]] + tone_dict[py[-1]] + ) + if BLANK_LEVEL == 1: + bopomofos.append("_") + else: + bopomofos.append(c) + if BLANK_LEVEL == 1: + bopomofos.append("_") + if BLANK_LEVEL == 2: + bopomofos.append("_") + char_index += len(word) + + if ( + len(word) == 3 + and bopomofos[0][-1] == "ˇ" + and bopomofos[1][-1] == "ˇ" + and bopomofos[-1][-1] == "ˇ" + ): + bopomofos[0] = bopomofos[0] + "ˊ" + bopomofos[1] = bopomofos[1] + "ˊ" + if len(word) == 2 and bopomofos[0][-1] == "ˇ" and bopomofos[-1][-1] == "ˇ": + bopomofos[0] = bopomofos[0][:-1] + "ˊ" + bopomofos = bu_sandhi(word, bopomofos) + bopomofos = yi_sandhi(word, bopomofos) + bopomofos = er_sandhi(word, bopomofos) + if not re.search("[\u4e00-\u9fff]", word): + text += "|" + word + continue + for i in range(len(bopomofos)): + bopomofos[i] = re.sub(r"([\u3105-\u3129])$", r"\1ˉ", bopomofos[i]) + if text != "": + text += "|" + text += "|".join(bopomofos) + return text + + +# Convert latin pronunciation to pinyin (bopomofo) +def latin_to_bopomofo(text): + for regex, replacement in _latin_to_bopomofo: + text = re.sub(regex, replacement, text) + return text + + +# Convert pinyin (bopomofo) to IPA +def bopomofo_to_ipa(text): + for regex, replacement in _bopomofo_to_ipa: + text = re.sub(regex, replacement, text) + return text + + +def _chinese_to_ipa(text, sentence): + text = number_to_chinese(text.strip()) + text = normalization(text) + text = chinese_to_bopomofo(text, sentence) + # pinyin = bpmf_to_pinyin(text) + text = latin_to_bopomofo(text) + text = bopomofo_to_ipa(text) + text = re.sub("([sɹ]`[⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ\2", text) + text = re.sub("([s][⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ\2", text) + text = re.sub(r"^\||[^\w\s_,\.\?!;:\'…\|→↓↑⁼ʰ`]", "", text) + text = re.sub(r"([,\.\?!;:\'…])", r"|\1|", text) + text = re.sub(r"\|+", "|", text) + text = text.rstrip("|") + return text + + +# Convert Chinese to IPA +def chinese_to_ipa(text, sentence, text_tokenizer): + # phonemes = text_tokenizer(text.strip()) + if type(text) == str: + return _chinese_to_ipa(text, sentence) + else: + result_ph = [] + for t in text: + result_ph.append(_chinese_to_ipa(t, sentence)) + return result_ph diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/text_tokenizers.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/text_tokenizers.py new file mode 100644 index 0000000000000000000000000000000000000000..7e0ffa6b14b4716b9674f27a837e0d040ca5558a --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/text_tokenizers.py @@ -0,0 +1,82 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re +from typing import List, Union + +from phonemizer.backend import EspeakBackend +from phonemizer.backend.espeak.language_switch import LanguageSwitch +from phonemizer.backend.espeak.words_mismatch import WordMismatch +from phonemizer.separator import Separator +from phonemizer.utils import list2str, str2list + + +class TextTokenizer: + """Phonemize Text.""" + + def __init__( + self, + language="en-us", + backend="espeak", + separator=Separator(word="|_|", syllable="-", phone="|"), + preserve_punctuation=True, + with_stress: bool = False, + tie: Union[bool, str] = False, + language_switch: LanguageSwitch = "remove-flags", + words_mismatch: WordMismatch = "ignore", + ) -> None: + self.preserve_punctuation_marks = ",.?!;:'…" + self.backend = EspeakBackend( + language, + punctuation_marks=self.preserve_punctuation_marks, + preserve_punctuation=preserve_punctuation, + with_stress=with_stress, + tie=tie, + language_switch=language_switch, + words_mismatch=words_mismatch, + ) + + self.separator = separator + + # convert chinese punctuation to english punctuation + def convert_chinese_punctuation(self, text: str) -> str: + text = text.replace("，", ",") + text = text.replace("。", ".") + text = text.replace("！", "!") + text = text.replace("？", "?") + text = text.replace("；", ";") + text = text.replace("：", ":") + text = text.replace("、", ",") + text = text.replace("‘", "'") + text = text.replace("’", "'") + text = text.replace("⋯", "…") + text = text.replace("···", "…") + text = text.replace("・・・", "…") + text = text.replace("...", "…") + return text + + def __call__(self, text, strip=True) -> List[str]: + text_type = type(text) + normalized_text = [] + for line in str2list(text): + line = self.convert_chinese_punctuation(line.strip()) + line = re.sub(r"[^\w\s_,\.\?!;:\'…]", "", line) + line = re.sub(r"\s*([,\.\?!;:\'…])\s*", r"\1", line) + line = re.sub(r"\s+", " ", line) + normalized_text.append(line) + # print("Normalized test: ", normalized_text[0]) + phonemized = self.backend.phonemize( + normalized_text, separator=self.separator, strip=strip, njobs=1 + ) + if text_type == str: + phonemized = re.sub(r"([,\.\?!;:\'…])", r"|\1|", list2str(phonemized)) + phonemized = re.sub(r"\|+", "|", phonemized) + phonemized = phonemized.rstrip("|") + else: + for i in range(len(phonemized)): + phonemized[i] = re.sub(r"([,\.\?!;:\'…])", r"|\1|", phonemized[i]) + phonemized[i] = re.sub(r"\|+", "|", phonemized[i]) + phonemized[i] = phonemized[i].rstrip("|") + return phonemized diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p/vocab.json b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/vocab.json new file mode 100644 index 0000000000000000000000000000000000000000..28d32aaf01881c6ff5449aaaf942d94b753a4e91 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p/vocab.json @@ -0,0 +1,372 @@ +{ + "vocab": { + ",": 0, + ".": 1, + "?": 2, + "!": 3, + "_": 4, + "iː": 5, + "ɪ": 6, + "ɜː": 7, + "ɚ": 8, + "oːɹ": 9, + "ɔː": 10, + "ɔːɹ": 11, + "ɑː": 12, + "uː": 13, + "ʊ": 14, + "ɑːɹ": 15, + "ʌ": 16, + "ɛ": 17, + "æ": 18, + "eɪ": 19, + "aɪ": 20, + "ɔɪ": 21, + "aʊ": 22, + "oʊ": 23, + "ɪɹ": 24, + "ɛɹ": 25, + "ʊɹ": 26, + "p": 27, + "b": 28, + "t": 29, + "d": 30, + "k": 31, + "ɡ": 32, + "f": 33, + "v": 34, + "θ": 35, + "ð": 36, + "s": 37, + "z": 38, + "ʃ": 39, + "ʒ": 40, + "h": 41, + "tʃ": 42, + "dʒ": 43, + "m": 44, + "n": 45, + "ŋ": 46, + "j": 47, + "w": 48, + "ɹ": 49, + "l": 50, + "tɹ": 51, + "dɹ": 52, + "ts": 53, + "dz": 54, + "i": 55, + "ɔ": 56, + "ə": 57, + "ɾ": 58, + "iə": 59, + "r": 60, + "u": 61, + "oː": 62, + "ɛː": 63, + "ɪː": 64, + "aɪə": 65, + "aɪɚ": 66, + "ɑ̃": 67, + "ç": 68, + "ɔ̃": 69, + "ææ": 70, + "ɐɐ": 71, + "ɡʲ": 72, + "nʲ": 73, + "iːː": 74, + + "p⁼": 75, + "pʰ": 76, + "t⁼": 77, + "tʰ": 78, + "k⁼": 79, + "kʰ": 80, + "x": 81, + "tʃ⁼": 82, + "tʃʰ": 83, + "ts`⁼": 84, + "ts`ʰ": 85, + "s`": 86, + "ɹ`": 87, + "ts⁼": 88, + "tsʰ": 89, + "p⁼wo": 90, + "p⁼wo→": 91, + "p⁼wo↑": 92, + "p⁼wo↓↑": 93, + "p⁼wo↓": 94, + "pʰwo": 95, + "pʰwo→": 96, + "pʰwo↑": 97, + "pʰwo↓↑": 98, + "pʰwo↓": 99, + "mwo": 100, + "mwo→": 101, + "mwo↑": 102, + "mwo↓↑": 103, + "mwo↓": 104, + "fwo": 105, + "fwo→": 106, + "fwo↑": 107, + "fwo↓↑": 108, + "fwo↓": 109, + "jɛn": 110, + "jɛn→": 111, + "jɛn↑": 112, + "jɛn↓↑": 113, + "jɛn↓": 114, + "ɥæn": 115, + "ɥæn→": 116, + "ɥæn↑": 117, + "ɥæn↓↑": 118, + "ɥæn↓": 119, + "in": 120, + "in→": 121, + "in↑": 122, + "in↓↑": 123, + "in↓": 124, + "ɥn": 125, + "ɥn→": 126, + "ɥn↑": 127, + "ɥn↓↑": 128, + "ɥn↓": 129, + "iŋ": 130, + "iŋ→": 131, + "iŋ↑": 132, + "iŋ↓↑": 133, + "iŋ↓": 134, + "ʊŋ": 135, + "ʊŋ→": 136, + "ʊŋ↑": 137, + "ʊŋ↓↑": 138, + "ʊŋ↓": 139, + "jʊŋ": 140, + "jʊŋ→": 141, + "jʊŋ↑": 142, + "jʊŋ↓↑": 143, + "jʊŋ↓": 144, + "ia": 145, + "ia→": 146, + "ia↑": 147, + "ia↓↑": 148, + "ia↓": 149, + "iɛ": 150, + "iɛ→": 151, + "iɛ↑": 152, + "iɛ↓↑": 153, + "iɛ↓": 154, + "iɑʊ": 155, + "iɑʊ→": 156, + "iɑʊ↑": 157, + "iɑʊ↓↑": 158, + "iɑʊ↓": 159, + "ioʊ": 160, + "ioʊ→": 161, + "ioʊ↑": 162, + "ioʊ↓↑": 163, + "ioʊ↓": 164, + "iɑŋ": 165, + "iɑŋ→": 166, + "iɑŋ↑": 167, + "iɑŋ↓↑": 168, + "iɑŋ↓": 169, + "ua": 170, + "ua→": 171, + "ua↑": 172, + "ua↓↑": 173, + "ua↓": 174, + "uo": 175, + "uo→": 176, + "uo↑": 177, + "uo↓↑": 178, + "uo↓": 179, + "uaɪ": 180, + "uaɪ→": 181, + "uaɪ↑": 182, + "uaɪ↓↑": 183, + "uaɪ↓": 184, + "ueɪ": 185, + "ueɪ→": 186, + "ueɪ↑": 187, + "ueɪ↓↑": 188, + "ueɪ↓": 189, + "uan": 190, + "uan→": 191, + "uan↑": 192, + "uan↓↑": 193, + "uan↓": 194, + "uən": 195, + "uən→": 196, + "uən↑": 197, + "uən↓↑": 198, + "uən↓": 199, + "uɑŋ": 200, + "uɑŋ→": 201, + "uɑŋ↑": 202, + "uɑŋ↓↑": 203, + "uɑŋ↓": 204, + "ɥɛ": 205, + "ɥɛ→": 206, + "ɥɛ↑": 207, + "ɥɛ↓↑": 208, + "ɥɛ↓": 209, + "a": 210, + "a→": 211, + "a↑": 212, + "a↓↑": 213, + "a↓": 214, + "o": 215, + "o→": 216, + "o↑": 217, + "o↓↑": 218, + "o↓": 219, + "ə→": 220, + "ə↑": 221, + "ə↓↑": 222, + "ə↓": 223, + "ɛ→": 224, + "ɛ↑": 225, + "ɛ↓↑": 226, + "ɛ↓": 227, + "aɪ→": 228, + "aɪ↑": 229, + "aɪ↓↑": 230, + "aɪ↓": 231, + "eɪ→": 232, + "eɪ↑": 233, + "eɪ↓↑": 234, + "eɪ↓": 235, + "ɑʊ": 236, + "ɑʊ→": 237, + "ɑʊ↑": 238, + "ɑʊ↓↑": 239, + "ɑʊ↓": 240, + "oʊ→": 241, + "oʊ↑": 242, + "oʊ↓↑": 243, + "oʊ↓": 244, + "an": 245, + "an→": 246, + "an↑": 247, + "an↓↑": 248, + "an↓": 249, + "ən": 250, + "ən→": 251, + "ən↑": 252, + "ən↓↑": 253, + "ən↓": 254, + "ɑŋ": 255, + "ɑŋ→": 256, + "ɑŋ↑": 257, + "ɑŋ↓↑": 258, + "ɑŋ↓": 259, + "əŋ": 260, + "əŋ→": 261, + "əŋ↑": 262, + "əŋ↓↑": 263, + "əŋ↓": 264, + "əɹ": 265, + "əɹ→": 266, + "əɹ↑": 267, + "əɹ↓↑": 268, + "əɹ↓": 269, + "i→": 270, + "i↑": 271, + "i↓↑": 272, + "i↓": 273, + "u→": 274, + "u↑": 275, + "u↓↑": 276, + "u↓": 277, + "ɥ": 278, + "ɥ→": 279, + "ɥ↑": 280, + "ɥ↓↑": 281, + "ɥ↓": 282, + "ts`⁼ɹ": 283, + "ts`⁼ɹ→": 284, + "ts`⁼ɹ↑": 285, + "ts`⁼ɹ↓↑": 286, + "ts`⁼ɹ↓": 287, + "ts`ʰɹ": 288, + "ts`ʰɹ→": 289, + "ts`ʰɹ↑": 290, + "ts`ʰɹ↓↑": 291, + "ts`ʰɹ↓": 292, + "s`ɹ": 293, + "s`ɹ→": 294, + "s`ɹ↑": 295, + "s`ɹ↓↑": 296, + "s`ɹ↓": 297, + "ɹ`ɹ": 298, + "ɹ`ɹ→": 299, + "ɹ`ɹ↑": 300, + "ɹ`ɹ↓↑": 301, + "ɹ`ɹ↓": 302, + "ts⁼ɹ": 303, + "ts⁼ɹ→": 304, + "ts⁼ɹ↑": 305, + "ts⁼ɹ↓↑": 306, + "ts⁼ɹ↓": 307, + "tsʰɹ": 308, + "tsʰɹ→": 309, + "tsʰɹ↑": 310, + "tsʰɹ↓↑": 311, + "tsʰɹ↓": 312, + "sɹ": 313, + "sɹ→": 314, + "sɹ↑": 315, + "sɹ↓↑": 316, + "sɹ↓": 317, + + "ɯ": 318, + "e": 319, + "aː": 320, + "ɯː": 321, + "eː": 322, + "ç": 323, + "ɸ": 324, + "ɰᵝ": 325, + "ɴ": 326, + "g": 327, + "dʑ": 328, + "q": 329, + "ː": 330, + "bj": 331, + "tɕ": 332, + "dej": 333, + "tej": 334, + "gj": 335, + "gɯ": 336, + "çj": 337, + "kj": 338, + "kɯ": 339, + "mj": 340, + "nj": 341, + "pj": 342, + "ɾj": 343, + "ɕ": 344, + "tsɯ": 345, + + "ɐ": 346, + "ɑ": 347, + "ɒ": 348, + "ɜ": 349, + "ɫ": 350, + "ʑ": 351, + "ʲ": 352, + + "y": 353, + "ø": 354, + "œ": 355, + "ʁ": 356, + "̃": 357, + "ɲ": 358, + + ":": 359, + ";": 360, + "'": 361, + "…": 362 + } +} diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/g2p_generation.py b/src/YingMusicSinger/utils/f5_tts/g2p/g2p_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..1e54fbb719d18e0671fb656377f498378da7e930 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/g2p_generation.py @@ -0,0 +1,129 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +import json +from typing import List + +from src.YingMusicSinger.utils.f5_tts.g2p.g2p import PhonemeBpeTokenizer +from src.YingMusicSinger.utils.f5_tts.g2p.utils.g2p import phonemizer_g2p + + +def ph_g2p(text, language): + return phonemizer_g2p(text=text, language=language) + + +def g2p(text, sentence, language): + return text_tokenizer.tokenize(text=text, sentence=sentence, language=language) + + +def is_chinese(char): + if char >= "\u4e00" and char <= "\u9fa5": + return True + else: + return False + + +def is_alphabet(char): + if (char >= "\u0041" and char <= "\u005a") or ( + char >= "\u0061" and char <= "\u007a" + ): + return True + else: + return False + + +def is_other(char): + if not (is_chinese(char) or is_alphabet(char)): + return True + else: + return False + + +def get_segment(text: str) -> List[str]: + # sentence --> [ch_part, en_part, ch_part, ...] + segments = [] + types = [] + flag = 0 + temp_seg = "" + temp_lang = "" + + # Determine the type of each character. type: blank, chinese, alphabet, number, unk and point. + for i, ch in enumerate(text): + if is_chinese(ch): + types.append("zh") + elif is_alphabet(ch): + types.append("en") + else: + types.append("other") + + assert len(types) == len(text) + + for i in range(len(types)): + # find the first char of the seg + if flag == 0: + temp_seg += text[i] + temp_lang = types[i] + flag = 1 + else: + if temp_lang == "other": + if types[i] == temp_lang: + temp_seg += text[i] + else: + temp_seg += text[i] + temp_lang = types[i] + else: + if types[i] == temp_lang: + temp_seg += text[i] + elif types[i] == "other": + temp_seg += text[i] + else: + segments.append((temp_seg, temp_lang)) + temp_seg = text[i] + temp_lang = types[i] + flag = 1 + + segments.append((temp_seg, temp_lang)) + return segments + + +def chn_eng_g2p(text: str): + # now only en and ch + segments = get_segment(text) + all_phoneme = "" + all_tokens = [] + + for index in range(len(segments)): + seg = segments[index] + phoneme, token = g2p(seg[0], text, seg[1]) + all_phoneme += phoneme + "|" + all_tokens += token + + if seg[1] == "en" and index == len(segments) - 1 and all_phoneme[-2] == "_": + all_phoneme = all_phoneme[:-2] + all_tokens = all_tokens[:-1] + return all_phoneme, all_tokens + + +text_tokenizer = PhonemeBpeTokenizer() +with open("./src/YingMusicSinger/utils/f5_tts/g2p/g2p/vocab.json", "r") as f: + json_data = f.read() +data = json.loads(json_data) +vocab = data["vocab"] + +if __name__ == "__main__": + phone, token = chn_eng_g2p("你好，hello world") + phone, token = chn_eng_g2p( + "你好，hello world, Bonjour, 테스트 해 보겠습니다, 五月雨緑" + ) + print(phone) + print(token) + + # phone, token = text_tokenizer.tokenize("你好，hello world, Bonjour, 테스트 해 보겠습니다, 五月雨緑", "", "auto") + phone, token = text_tokenizer.tokenize("緑", "", "auto") + # phone, token = text_tokenizer.tokenize("आइए इसका परीक्षण करें", "", "auto") + # phone, token = text_tokenizer.tokenize("आइए इसका परीक्षण करें", "", "other") + print(phone) + print(token) diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/infer_dpo.py b/src/YingMusicSinger/utils/f5_tts/g2p/infer_dpo.py new file mode 100644 index 0000000000000000000000000000000000000000..3abd7b3b77df3dc13fdf7a7a1e8808a9352a29a8 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/infer_dpo.py @@ -0,0 +1,277 @@ +import argparse +import json +import os +import random + +import numpy as np +import torch +from f5_tts.infer.utils_infer import load_checkpoint +from f5_tts.model import CFM, DiT +from f5_tts.model.alsp_lance.data.npydata import FloatData +from f5_tts.model.alsp_lance.tools import LanceReader, LanceWriter +from tqdm import tqdm + +filter_keyword_list = [ + "纯音乐", + "编曲", + "作词", + "作曲", + "调音", + "制作人", + "录音师", +] + +filter_full_list = ["music", "end"] + + +def check_lyric(time: float, lyric: str): + if time < 0.1: + return False + for filter_keyword in filter_keyword_list: + if filter_keyword in lyric: + return False + for filter_full in filter_full_list: + if filter_full == lyric.strip().lower(): + return False + if len(lyric) == 0: + return False + return True + + +def parse_lyrics(lyrics: str): + lyrics_with_time = [] + lyrics = lyrics.strip() + for line in lyrics.split("\n"): + try: + time, lyric = line[1:9], line[10:] + lyric = lyric.strip() + mins, secs = time.split(":") + secs = int(mins) * 60 + float(secs) + # print(lyric, check_lyric(secs, lyric)) + if not check_lyric(secs, lyric): + continue + lyrics_with_time.append((secs, lyric)) + except: + # traceback.print_exc() + continue + # print("error", line) + return lyrics_with_time + + +class CNENTokenizer: + def __init__(self): + with open("./src/YingMusicSinger/utils/f5_tts/g2p/g2p/vocab.json", "r") as file: + self.phone2id: dict = json.load(file)["vocab"] + self.id2phone = {v: k for (k, v) in self.phone2id.items()} + from f5_tts.g2p.g2p_generation import chn_eng_g2p + + self.tokenizer = chn_eng_g2p + + def encode(self, text): + phone, token = self.tokenizer(text) + token = [x + 1 for x in token] + return token + + def decode(self, token): + return "|".join([self.id2phone[x - 1] for x in token]) + + +def inference( + model, + cond, + text, + duration, + style_prompt, + style, + output_dir, + song_name, + ckpt_step, + start_time, + latent_pred_start_frame, + latent_pred_end_frame, + epoch, + cfg_strength, +): + # import pdb; pdb.set_trace() + with torch.inference_mode(): + generated, _ = model.sample( + cond=cond, + text=text, + duration=duration, + style_prompt=style_prompt, + steps=32, + cfg_strength=cfg_strength, + sway_sampling_coef=None, + start_time=start_time, + latent_pred_start_frame=latent_pred_start_frame, + latent_pred_end_frame=latent_pred_end_frame, + ) + + generated = generated.to(torch.float32) # [b t d] + latent = generated.transpose(1, 2) # [b d t] + latent = latent.detach().cpu.numpy() + + return latent + + +def get_style_prompt(device, song_name, song_name2ref_npy): + mulan_style_path = song_name2ref_npy[song_name] + mulan_stlye = np.load(mulan_style_path) + + style_prompt = torch.from_numpy(mulan_stlye).to(device) # [1, 512] + style_prompt = style_prompt.half() + + return style_prompt + + +def get_lrc_prompt(text, tokenizer, dit_model, max_secs): + max_frames = 2048 + lyrics_shift = 2 + sampling_rate = 44100 + downsample_rate = 2048 + + pad_token_id = 0 + comma_token_id = 1 + period_token_id = 2 + + fsmin = -10 + fsmax = 10 + + lrc_with_time = parse_lyrics(text) + + modified_lrc_with_time = [] + for i in range(len(lrc_with_time)): + time, line = lrc_with_time[i] + # line_token = self.tokenizer.encode(line) + line_token = tokenizer.encode(line) + modified_lrc_with_time.append((time, line_token)) + + lrc_with_time = modified_lrc_with_time + + lrc_with_time = [ + (time_start, line) + for (time_start, line) in lrc_with_time + if time_start < max_secs + ] + # latent_end_time = lrc_with_time[-1][0] if len(lrc_with_time) >= 1 else -1 + lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time + + normalized_start_time = 0.0 + + lrc = torch.zeros((max_frames,), dtype=torch.long) + + tokens_count = 0 + last_end_pos = 0 + for time_start, line in lrc_with_time: + tokens = [ + token if token != period_token_id else comma_token_id for token in line + ] + [period_token_id] + tokens = torch.tensor(tokens, dtype=torch.long) + num_tokens = tokens.shape[0] + + gt_frame_start = int(time_start * sampling_rate / downsample_rate) + + frame_shift = random.randint(int(fsmin), int(fsmax)) + + frame_start = max(gt_frame_start - frame_shift, last_end_pos) + frame_len = min(num_tokens, max_frames - frame_start) + + # print(gt_frame_start, frame_shift, frame_start, frame_len, tokens_count, last_end_pos, full_pos_emb.shape) + + lrc[frame_start : frame_start + frame_len] = tokens[:frame_len] + + tokens_count += num_tokens + last_end_pos = frame_start + frame_len + + lrc_emb = lrc.unsqueeze(0).to(dit_model.device) + + normalized_start_time = ( + torch.tensor(normalized_start_time).unsqueeze(0).to(dit_model.device) + ) + + return lrc_emb, normalized_start_time + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--model-config", type=str, default=None) + parser.add_argument("--ckpt-path", type=str, default=None) + parser.add_argument("--output-dir", type=str, default=None) # lance + parser.add_argument("--lrc-path", type=str, default=None) + parser.add_argument("--mulan-style-path", type=str, default=None) # lance + parser.add_argument("--cfg-strength", type=float, default=None) + + args = parser.parse_args() + + lrc_path = args.lrc_path + cfg_strength = args.cfg_strength + style_path = args.mulan_style_path + + with open(args.model_config) as f: + model_config = json.load(f) + + model_cls = DiT + ckpt_path = args.ckpt_path + device = "cuda" + use_style_prompt = True + dit_model = CFM( + transformer=model_cls( + **model_config["model"], use_style_prompt=use_style_prompt + ), + num_channels=model_config["model"]["mel_dim"], + use_style_prompt=use_style_prompt, + ) + dit_model = dit_model.to(device) + dit_model = load_checkpoint(dit_model, ckpt_path, device=device, use_ema=True) + + lrc_tokenizer = CNENTokenizer() + + sampling_rate = 44100 + downsample_rate = 2048 + max_frames = 2048 + max_secs = max_frames / (sampling_rate / downsample_rate) + + output_dir = args.output_dir + writer = LanceWriter(output_dir, target_cls=FloatData) + + reader = LanceReader(style_path, target_cls=FloatData) + + WRITE_INTERVAL = 500 + + latent_data = [] + for id in tqdm(reader.get_ids()): + item = reader.get_datas_by_rowids(row_ids=[id._rowid])[0] + data_id = item.data_id + style_prompt = torch.from_numpy(item.data).to(device) + stlye_prompt = style_prompt.half() + + lrc_path = os.path.join(lrc_path, f"{data_id}.lrc") + with (open(lrc_path), "r") as f: + lrc = [line.strip() for line in f.readlines()] + lrc_prompt, start_time = get_lrc_prompt(lrc, lrc_tokenizer, dit_model, max_secs) + + latent_prompt = torch.zeros(1, max_frames, 64).to(device) + sf = 0 + ef = max_frames + + generated_latent = inference( + model=dit_model, + cond=latent_prompt, + text=lrc_prompt, + duration=max_frames, + style_prompt=style_prompt, + output_dir=output_dir, + start_time=start_time, + latent_pred_start_frame=sf, + latent_pred_end_frame=ef, + cfg_strength=cfg_strength, + ) # [b d t] numpy + + latent_data.append(generated_latent) + + if len(latent_data) > WRITE_INTERVAL: + writer.write_parallel(latent_data) + latent_data = [] + + writer.write_parallel(latent_data) diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/bpmf_2_pinyin.txt b/src/YingMusicSinger/utils/f5_tts/g2p/sources/bpmf_2_pinyin.txt new file mode 100644 index 0000000000000000000000000000000000000000..474529e5d347b94a80e5052de0065347ff14b95e --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/bpmf_2_pinyin.txt @@ -0,0 +1,41 @@ +b ㄅ +p ㄆ +m ㄇ +f ㄈ +d ㄉ +t ㄊ +n ㄋ +l ㄌ +g ㄍ +k ㄎ +h ㄏ +j ㄐ +q ㄑ +x ㄒ +zh ㄓ +ch ㄔ +sh ㄕ +r ㄖ +z ㄗ +c ㄘ +s ㄙ +i ㄧ +u ㄨ +v ㄩ +a ㄚ +o ㄛ +e ㄜ +e ㄝ +ai ㄞ +ei ㄟ +ao ㄠ +ou ㄡ +an ㄢ +en ㄣ +ang ㄤ +eng ㄥ +er ㄦ +2 ˊ +3 ˇ +4 ˋ +0 ˙ diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/chinese_lexicon.txt b/src/YingMusicSinger/utils/f5_tts/g2p/sources/chinese_lexicon.txt new file mode 100644 index 0000000000000000000000000000000000000000..4d7dbf347a29d3b87c199d0e56ef7f1dbf28a6ee --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/chinese_lexicon.txt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a3a7685d1c3e68eb2fa304bfc63e90c90c3c1a1948839a5b1b507b2131b3e2fb +size 14779443 diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/config.json b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/config.json new file mode 100644 index 0000000000000000000000000000000000000000..5fb70ca91db27a4ad73b58a0c500a903be9bc1a9 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/config.json @@ -0,0 +1,819 @@ +{ + "_name_or_path": "/BERT-POLY-v2/pretrained_models/mini_bert", + "architectures": [ + "BertPoly" + ], + "attention_probs_dropout_prob": 0.1, + "classifier_dropout": null, + "directionality": "bidi", + "gradient_checkpointing": false, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.1, + "hidden_size": 384, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1", + "2": "LABEL_2", + "3": "LABEL_3", + "4": "LABEL_4", + "5": "LABEL_5", + "6": "LABEL_6", + "7": "LABEL_7", + "8": "LABEL_8", + "9": "LABEL_9", + "10": "LABEL_10", + "11": "LABEL_11", + "12": "LABEL_12", + "13": "LABEL_13", + "14": "LABEL_14", + "15": "LABEL_15", + "16": "LABEL_16", + "17": "LABEL_17", + "18": "LABEL_18", + "19": "LABEL_19", + "20": "LABEL_20", + "21": "LABEL_21", + "22": "LABEL_22", + "23": "LABEL_23", + "24": "LABEL_24", + "25": "LABEL_25", + "26": "LABEL_26", + "27": "LABEL_27", + "28": "LABEL_28", + "29": "LABEL_29", + "30": "LABEL_30", + "31": "LABEL_31", + "32": "LABEL_32", + "33": "LABEL_33", + "34": "LABEL_34", + "35": "LABEL_35", + "36": "LABEL_36", + "37": "LABEL_37", + "38": "LABEL_38", + "39": "LABEL_39", + "40": "LABEL_40", + "41": "LABEL_41", + "42": "LABEL_42", + "43": "LABEL_43", + "44": "LABEL_44", + "45": "LABEL_45", + "46": "LABEL_46", + "47": "LABEL_47", + "48": "LABEL_48", + "49": "LABEL_49", + "50": "LABEL_50", + "51": "LABEL_51", + "52": "LABEL_52", + "53": "LABEL_53", + "54": "LABEL_54", + "55": "LABEL_55", + "56": "LABEL_56", + "57": "LABEL_57", + "58": "LABEL_58", + "59": "LABEL_59", + "60": "LABEL_60", + "61": "LABEL_61", + "62": "LABEL_62", + "63": "LABEL_63", + "64": "LABEL_64", + "65": "LABEL_65", + "66": "LABEL_66", + "67": "LABEL_67", + "68": "LABEL_68", + "69": "LABEL_69", + "70": "LABEL_70", + "71": "LABEL_71", + "72": "LABEL_72", + "73": "LABEL_73", + "74": "LABEL_74", + "75": "LABEL_75", + "76": "LABEL_76", + "77": "LABEL_77", + "78": "LABEL_78", + "79": "LABEL_79", + "80": "LABEL_80", + "81": "LABEL_81", + "82": "LABEL_82", + "83": "LABEL_83", + "84": "LABEL_84", + "85": "LABEL_85", + "86": "LABEL_86", + "87": "LABEL_87", + "88": "LABEL_88", + "89": "LABEL_89", + "90": "LABEL_90", + "91": "LABEL_91", + "92": "LABEL_92", + "93": "LABEL_93", + "94": "LABEL_94", + "95": "LABEL_95", + "96": "LABEL_96", + "97": "LABEL_97", + "98": "LABEL_98", + "99": "LABEL_99", + "100": "LABEL_100", + "101": "LABEL_101", + "102": "LABEL_102", + "103": "LABEL_103", + "104": "LABEL_104", + "105": "LABEL_105", + "106": "LABEL_106", + "107": "LABEL_107", + "108": "LABEL_108", + "109": "LABEL_109", + "110": "LABEL_110", + "111": "LABEL_111", + "112": "LABEL_112", + "113": "LABEL_113", + "114": "LABEL_114", + "115": "LABEL_115", + "116": "LABEL_116", + "117": "LABEL_117", + "118": "LABEL_118", + "119": "LABEL_119", + "120": "LABEL_120", + "121": "LABEL_121", + "122": "LABEL_122", + "123": "LABEL_123", + "124": "LABEL_124", + "125": "LABEL_125", + "126": "LABEL_126", + "127": "LABEL_127", + "128": "LABEL_128", + "129": "LABEL_129", + "130": "LABEL_130", + "131": "LABEL_131", + "132": "LABEL_132", + "133": "LABEL_133", + "134": "LABEL_134", + "135": "LABEL_135", + "136": "LABEL_136", + "137": "LABEL_137", + "138": "LABEL_138", + "139": "LABEL_139", + "140": "LABEL_140", + "141": "LABEL_141", + "142": "LABEL_142", + "143": "LABEL_143", + "144": "LABEL_144", + "145": "LABEL_145", + "146": "LABEL_146", + "147": "LABEL_147", + "148": "LABEL_148", + "149": "LABEL_149", + "150": "LABEL_150", + "151": "LABEL_151", + "152": "LABEL_152", + "153": "LABEL_153", + "154": "LABEL_154", + "155": "LABEL_155", + "156": "LABEL_156", + "157": "LABEL_157", + "158": "LABEL_158", + "159": "LABEL_159", + "160": "LABEL_160", + "161": "LABEL_161", + "162": "LABEL_162", + "163": "LABEL_163", + "164": "LABEL_164", + "165": "LABEL_165", + "166": "LABEL_166", + "167": "LABEL_167", + "168": "LABEL_168", + "169": "LABEL_169", + "170": "LABEL_170", + "171": "LABEL_171", + "172": "LABEL_172", + "173": "LABEL_173", + "174": "LABEL_174", + "175": "LABEL_175", + "176": "LABEL_176", + "177": "LABEL_177", + "178": "LABEL_178", + "179": "LABEL_179", + "180": "LABEL_180", + "181": "LABEL_181", + "182": "LABEL_182", + "183": "LABEL_183", + "184": "LABEL_184", + "185": "LABEL_185", + "186": "LABEL_186", + "187": "LABEL_187", + "188": "LABEL_188", + "189": "LABEL_189", + "190": "LABEL_190", + "191": "LABEL_191", + "192": "LABEL_192", + "193": "LABEL_193", + "194": "LABEL_194", + "195": "LABEL_195", + "196": "LABEL_196", + "197": "LABEL_197", + "198": "LABEL_198", + "199": "LABEL_199", + "200": "LABEL_200", + "201": "LABEL_201", + "202": "LABEL_202", + "203": "LABEL_203", + "204": "LABEL_204", + "205": "LABEL_205", + "206": "LABEL_206", + "207": "LABEL_207", + "208": "LABEL_208", + "209": "LABEL_209", + "210": "LABEL_210", + "211": "LABEL_211", + "212": "LABEL_212", + "213": "LABEL_213", + "214": "LABEL_214", + "215": "LABEL_215", + "216": "LABEL_216", + "217": "LABEL_217", + "218": "LABEL_218", + "219": "LABEL_219", + "220": "LABEL_220", + "221": "LABEL_221", + "222": "LABEL_222", + "223": "LABEL_223", + "224": "LABEL_224", + "225": "LABEL_225", + "226": "LABEL_226", + "227": "LABEL_227", + "228": "LABEL_228", + "229": "LABEL_229", + "230": "LABEL_230", + "231": "LABEL_231", + "232": "LABEL_232", + "233": "LABEL_233", + "234": "LABEL_234", + "235": "LABEL_235", + "236": "LABEL_236", + "237": "LABEL_237", + "238": "LABEL_238", + "239": "LABEL_239", + "240": "LABEL_240", + "241": "LABEL_241", + "242": "LABEL_242", + "243": "LABEL_243", + "244": "LABEL_244", + "245": "LABEL_245", + "246": "LABEL_246", + "247": "LABEL_247", + "248": "LABEL_248", + "249": "LABEL_249", + "250": "LABEL_250", + "251": "LABEL_251", + "252": "LABEL_252", + "253": "LABEL_253", + "254": "LABEL_254", + "255": "LABEL_255", + "256": "LABEL_256", + "257": "LABEL_257", + "258": "LABEL_258", + "259": "LABEL_259", + "260": "LABEL_260", + "261": "LABEL_261", + "262": "LABEL_262", + "263": "LABEL_263", + "264": "LABEL_264", + "265": "LABEL_265", + "266": "LABEL_266", + "267": "LABEL_267", + "268": "LABEL_268", + "269": "LABEL_269", + "270": "LABEL_270", + "271": "LABEL_271", + "272": "LABEL_272", + "273": "LABEL_273", + "274": "LABEL_274", + "275": "LABEL_275", + "276": "LABEL_276", + "277": "LABEL_277", + "278": "LABEL_278", + "279": "LABEL_279", + "280": "LABEL_280", + "281": "LABEL_281", + "282": "LABEL_282", + "283": "LABEL_283", + "284": "LABEL_284", + "285": "LABEL_285", + "286": "LABEL_286", + "287": "LABEL_287", + "288": "LABEL_288", + "289": "LABEL_289", + "290": "LABEL_290", + "291": "LABEL_291", + "292": "LABEL_292", + "293": "LABEL_293", + "294": "LABEL_294", + "295": "LABEL_295", + "296": "LABEL_296", + "297": "LABEL_297", + "298": "LABEL_298", + "299": "LABEL_299", + "300": "LABEL_300", + "301": "LABEL_301", + "302": "LABEL_302", + "303": "LABEL_303", + "304": "LABEL_304", + "305": "LABEL_305", + "306": "LABEL_306", + "307": "LABEL_307", + "308": "LABEL_308", + "309": "LABEL_309", + "310": "LABEL_310", + "311": "LABEL_311", + "312": "LABEL_312", + "313": "LABEL_313", + "314": "LABEL_314", + "315": "LABEL_315", + "316": "LABEL_316", + "317": "LABEL_317", + "318": "LABEL_318", + "319": "LABEL_319", + "320": "LABEL_320", + "321": "LABEL_321", + "322": "LABEL_322", + "323": "LABEL_323", + "324": "LABEL_324", + "325": "LABEL_325", + "326": "LABEL_326", + "327": "LABEL_327", + "328": "LABEL_328", + "329": "LABEL_329", + "330": "LABEL_330", + "331": "LABEL_331", + "332": "LABEL_332", + "333": "LABEL_333", + "334": "LABEL_334", + "335": "LABEL_335", + "336": "LABEL_336", + "337": "LABEL_337", + "338": "LABEL_338", + "339": "LABEL_339", + "340": "LABEL_340", + "341": "LABEL_341", + "342": "LABEL_342", + "343": "LABEL_343", + "344": "LABEL_344", + "345": "LABEL_345", + "346": "LABEL_346", + "347": "LABEL_347", + "348": "LABEL_348", + "349": "LABEL_349", + "350": "LABEL_350", + "351": "LABEL_351", + "352": "LABEL_352", + "353": "LABEL_353", + "354": "LABEL_354", + "355": "LABEL_355", + "356": "LABEL_356", + "357": "LABEL_357", + "358": "LABEL_358", + "359": "LABEL_359", + "360": "LABEL_360", + "361": "LABEL_361", + "362": "LABEL_362", + "363": "LABEL_363", + "364": "LABEL_364", + "365": "LABEL_365", + "366": "LABEL_366", + "367": "LABEL_367", + "368": "LABEL_368", + "369": "LABEL_369", + "370": "LABEL_370", + "371": "LABEL_371", + "372": "LABEL_372", + "373": "LABEL_373", + "374": "LABEL_374", + "375": "LABEL_375", + "376": "LABEL_376", + "377": "LABEL_377", + "378": "LABEL_378", + "379": "LABEL_379", + "380": "LABEL_380", + "381": "LABEL_381", + "382": "LABEL_382", + "383": "LABEL_383", + "384": "LABEL_384", + "385": "LABEL_385", + "386": "LABEL_386", + "387": "LABEL_387", + "388": "LABEL_388", + "389": "LABEL_389", + "390": "LABEL_390" + }, + "initializer_range": 0.02, + "intermediate_size": 1536, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1, + "LABEL_10": 10, + "LABEL_100": 100, + "LABEL_101": 101, + "LABEL_102": 102, + "LABEL_103": 103, + "LABEL_104": 104, + "LABEL_105": 105, + "LABEL_106": 106, + "LABEL_107": 107, + "LABEL_108": 108, + "LABEL_109": 109, + "LABEL_11": 11, + "LABEL_110": 110, + "LABEL_111": 111, + "LABEL_112": 112, + "LABEL_113": 113, + "LABEL_114": 114, + "LABEL_115": 115, + "LABEL_116": 116, + "LABEL_117": 117, + "LABEL_118": 118, + "LABEL_119": 119, + "LABEL_12": 12, + "LABEL_120": 120, + "LABEL_121": 121, + "LABEL_122": 122, + "LABEL_123": 123, + "LABEL_124": 124, + "LABEL_125": 125, + "LABEL_126": 126, + "LABEL_127": 127, + "LABEL_128": 128, + "LABEL_129": 129, + "LABEL_13": 13, + "LABEL_130": 130, + "LABEL_131": 131, + "LABEL_132": 132, + "LABEL_133": 133, + "LABEL_134": 134, + "LABEL_135": 135, + "LABEL_136": 136, + "LABEL_137": 137, + "LABEL_138": 138, + "LABEL_139": 139, + "LABEL_14": 14, + "LABEL_140": 140, + "LABEL_141": 141, + "LABEL_142": 142, + "LABEL_143": 143, + "LABEL_144": 144, + "LABEL_145": 145, + "LABEL_146": 146, + "LABEL_147": 147, + "LABEL_148": 148, + "LABEL_149": 149, + "LABEL_15": 15, + "LABEL_150": 150, + "LABEL_151": 151, + "LABEL_152": 152, + "LABEL_153": 153, + "LABEL_154": 154, + "LABEL_155": 155, + "LABEL_156": 156, + "LABEL_157": 157, + "LABEL_158": 158, + "LABEL_159": 159, + "LABEL_16": 16, + "LABEL_160": 160, + "LABEL_161": 161, + "LABEL_162": 162, + "LABEL_163": 163, + "LABEL_164": 164, + "LABEL_165": 165, + "LABEL_166": 166, + "LABEL_167": 167, + "LABEL_168": 168, + "LABEL_169": 169, + "LABEL_17": 17, + "LABEL_170": 170, + "LABEL_171": 171, + "LABEL_172": 172, + "LABEL_173": 173, + "LABEL_174": 174, + "LABEL_175": 175, + "LABEL_176": 176, + "LABEL_177": 177, + "LABEL_178": 178, + "LABEL_179": 179, + "LABEL_18": 18, + "LABEL_180": 180, + "LABEL_181": 181, + "LABEL_182": 182, + "LABEL_183": 183, + "LABEL_184": 184, + "LABEL_185": 185, + "LABEL_186": 186, + "LABEL_187": 187, + "LABEL_188": 188, + "LABEL_189": 189, + "LABEL_19": 19, + "LABEL_190": 190, + "LABEL_191": 191, + "LABEL_192": 192, + "LABEL_193": 193, + "LABEL_194": 194, + "LABEL_195": 195, + "LABEL_196": 196, + "LABEL_197": 197, + "LABEL_198": 198, + "LABEL_199": 199, + "LABEL_2": 2, + "LABEL_20": 20, + "LABEL_200": 200, + "LABEL_201": 201, + "LABEL_202": 202, + "LABEL_203": 203, + "LABEL_204": 204, + "LABEL_205": 205, + "LABEL_206": 206, + "LABEL_207": 207, + "LABEL_208": 208, + "LABEL_209": 209, + "LABEL_21": 21, + "LABEL_210": 210, + "LABEL_211": 211, + "LABEL_212": 212, + "LABEL_213": 213, + "LABEL_214": 214, + "LABEL_215": 215, + "LABEL_216": 216, + "LABEL_217": 217, + "LABEL_218": 218, + "LABEL_219": 219, + "LABEL_22": 22, + "LABEL_220": 220, + "LABEL_221": 221, + "LABEL_222": 222, + "LABEL_223": 223, + "LABEL_224": 224, + "LABEL_225": 225, + "LABEL_226": 226, + "LABEL_227": 227, + "LABEL_228": 228, + "LABEL_229": 229, + "LABEL_23": 23, + "LABEL_230": 230, + "LABEL_231": 231, + "LABEL_232": 232, + "LABEL_233": 233, + "LABEL_234": 234, + "LABEL_235": 235, + "LABEL_236": 236, + "LABEL_237": 237, + "LABEL_238": 238, + "LABEL_239": 239, + "LABEL_24": 24, + "LABEL_240": 240, + "LABEL_241": 241, + "LABEL_242": 242, + "LABEL_243": 243, + "LABEL_244": 244, + "LABEL_245": 245, + "LABEL_246": 246, + "LABEL_247": 247, + "LABEL_248": 248, + "LABEL_249": 249, + "LABEL_25": 25, + "LABEL_250": 250, + "LABEL_251": 251, + "LABEL_252": 252, + "LABEL_253": 253, + "LABEL_254": 254, + "LABEL_255": 255, + "LABEL_256": 256, + "LABEL_257": 257, + "LABEL_258": 258, + "LABEL_259": 259, + "LABEL_26": 26, + "LABEL_260": 260, + "LABEL_261": 261, + "LABEL_262": 262, + "LABEL_263": 263, + "LABEL_264": 264, + "LABEL_265": 265, + "LABEL_266": 266, + "LABEL_267": 267, + "LABEL_268": 268, + "LABEL_269": 269, + "LABEL_27": 27, + "LABEL_270": 270, + "LABEL_271": 271, + "LABEL_272": 272, + "LABEL_273": 273, + "LABEL_274": 274, + "LABEL_275": 275, + "LABEL_276": 276, + "LABEL_277": 277, + "LABEL_278": 278, + "LABEL_279": 279, + "LABEL_28": 28, + "LABEL_280": 280, + "LABEL_281": 281, + "LABEL_282": 282, + "LABEL_283": 283, + "LABEL_284": 284, + "LABEL_285": 285, + "LABEL_286": 286, + "LABEL_287": 287, + "LABEL_288": 288, + "LABEL_289": 289, + "LABEL_29": 29, + "LABEL_290": 290, + "LABEL_291": 291, + "LABEL_292": 292, + "LABEL_293": 293, + "LABEL_294": 294, + "LABEL_295": 295, + "LABEL_296": 296, + "LABEL_297": 297, + "LABEL_298": 298, + "LABEL_299": 299, + "LABEL_3": 3, + "LABEL_30": 30, + "LABEL_300": 300, + "LABEL_301": 301, + "LABEL_302": 302, + "LABEL_303": 303, + "LABEL_304": 304, + "LABEL_305": 305, + "LABEL_306": 306, + "LABEL_307": 307, + "LABEL_308": 308, + "LABEL_309": 309, + "LABEL_31": 31, + "LABEL_310": 310, + "LABEL_311": 311, + "LABEL_312": 312, + "LABEL_313": 313, + "LABEL_314": 314, + "LABEL_315": 315, + "LABEL_316": 316, + "LABEL_317": 317, + "LABEL_318": 318, + "LABEL_319": 319, + "LABEL_32": 32, + "LABEL_320": 320, + "LABEL_321": 321, + "LABEL_322": 322, + "LABEL_323": 323, + "LABEL_324": 324, + "LABEL_325": 325, + "LABEL_326": 326, + "LABEL_327": 327, + "LABEL_328": 328, + "LABEL_329": 329, + "LABEL_33": 33, + "LABEL_330": 330, + "LABEL_331": 331, + "LABEL_332": 332, + "LABEL_333": 333, + "LABEL_334": 334, + "LABEL_335": 335, + "LABEL_336": 336, + "LABEL_337": 337, + "LABEL_338": 338, + "LABEL_339": 339, + "LABEL_34": 34, + "LABEL_340": 340, + "LABEL_341": 341, + "LABEL_342": 342, + "LABEL_343": 343, + "LABEL_344": 344, + "LABEL_345": 345, + "LABEL_346": 346, + "LABEL_347": 347, + "LABEL_348": 348, + "LABEL_349": 349, + "LABEL_35": 35, + "LABEL_350": 350, + "LABEL_351": 351, + "LABEL_352": 352, + "LABEL_353": 353, + "LABEL_354": 354, + "LABEL_355": 355, + "LABEL_356": 356, + "LABEL_357": 357, + "LABEL_358": 358, + "LABEL_359": 359, + "LABEL_36": 36, + "LABEL_360": 360, + "LABEL_361": 361, + "LABEL_362": 362, + "LABEL_363": 363, + "LABEL_364": 364, + "LABEL_365": 365, + "LABEL_366": 366, + "LABEL_367": 367, + "LABEL_368": 368, + "LABEL_369": 369, + "LABEL_37": 37, + "LABEL_370": 370, + "LABEL_371": 371, + "LABEL_372": 372, + "LABEL_373": 373, + "LABEL_374": 374, + "LABEL_375": 375, + "LABEL_376": 376, + "LABEL_377": 377, + "LABEL_378": 378, + "LABEL_379": 379, + "LABEL_38": 38, + "LABEL_380": 380, + "LABEL_381": 381, + "LABEL_382": 382, + "LABEL_383": 383, + "LABEL_384": 384, + "LABEL_385": 385, + "LABEL_386": 386, + "LABEL_387": 387, + "LABEL_388": 388, + "LABEL_389": 389, + "LABEL_39": 39, + "LABEL_390": 390, + "LABEL_4": 4, + "LABEL_40": 40, + "LABEL_41": 41, + "LABEL_42": 42, + "LABEL_43": 43, + "LABEL_44": 44, + "LABEL_45": 45, + "LABEL_46": 46, + "LABEL_47": 47, + "LABEL_48": 48, + "LABEL_49": 49, + "LABEL_5": 5, + "LABEL_50": 50, + "LABEL_51": 51, + "LABEL_52": 52, + "LABEL_53": 53, + "LABEL_54": 54, + "LABEL_55": 55, + "LABEL_56": 56, + "LABEL_57": 57, + "LABEL_58": 58, + "LABEL_59": 59, + "LABEL_6": 6, + "LABEL_60": 60, + "LABEL_61": 61, + "LABEL_62": 62, + "LABEL_63": 63, + "LABEL_64": 64, + "LABEL_65": 65, + "LABEL_66": 66, + "LABEL_67": 67, + "LABEL_68": 68, + "LABEL_69": 69, + "LABEL_7": 7, + "LABEL_70": 70, + "LABEL_71": 71, + "LABEL_72": 72, + "LABEL_73": 73, + "LABEL_74": 74, + "LABEL_75": 75, + "LABEL_76": 76, + "LABEL_77": 77, + "LABEL_78": 78, + "LABEL_79": 79, + "LABEL_8": 8, + "LABEL_80": 80, + "LABEL_81": 81, + "LABEL_82": 82, + "LABEL_83": 83, + "LABEL_84": 84, + "LABEL_85": 85, + "LABEL_86": 86, + "LABEL_87": 87, + "LABEL_88": 88, + "LABEL_89": 89, + "LABEL_9": 9, + "LABEL_90": 90, + "LABEL_91": 91, + "LABEL_92": 92, + "LABEL_93": 93, + "LABEL_94": 94, + "LABEL_95": 95, + "LABEL_96": 96, + "LABEL_97": 97, + "LABEL_98": 98, + "LABEL_99": 99 + }, + "layer_norm_eps": 1e-12, + "max_position_embeddings": 512, + "model_type": "bert", + "num_attention_heads": 12, + "num_hidden_layers": 6, + "num_relation_heads": 32, + "pad_token_id": 0, + "pooler_fc_size": 768, + "pooler_num_attention_heads": 12, + "pooler_num_fc_layers": 3, + "pooler_size_per_head": 128, + "pooler_type": "first_token_transform", + "position_embedding_type": "absolute", + "torch_dtype": "float32", + "transformers_version": "4.44.1", + "type_vocab_size": 2, + "use_cache": true, + "vocab_size": 21128 +} diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/poly_bert_model.onnx b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/poly_bert_model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..6b952b9717eb71bb5a7aa2492478095f117858dd --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/poly_bert_model.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8765d835ffdf9811c832d4dc7b6a552757aa8615c01d1184db716a50c20aebbc +size 76583333 diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polychar.txt b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polychar.txt new file mode 100644 index 0000000000000000000000000000000000000000..819f6249a661134128c7a4bc72a1059ebe133d20 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polychar.txt @@ -0,0 +1,159 @@ +丧 +中 +为 +乌 +乐 +了 +什 +仔 +令 +任 +会 +传 +佛 +供 +便 +倒 +假 +兴 +冠 +冲 +几 +分 +切 +划 +创 +剥 +勒 +区 +华 +单 +卜 +占 +卡 +卷 +厦 +参 +发 +只 +号 +同 +吐 +和 +喝 +圈 +地 +塞 +壳 +处 +奇 +奔 +好 +宁 +宿 +将 +少 +尽 +岗 +差 +巷 +帖 +干 +应 +度 +弹 +强 +当 +待 +得 +恶 +扁 +扇 +扎 +扫 +担 +挑 +据 +撒 +教 +散 +数 +斗 +晃 +曝 +曲 +更 +曾 +朝 +朴 +杆 +查 +校 +模 +横 +没 +泡 +济 +混 +漂 +炸 +熟 +燕 +片 +率 +畜 +的 +盛 +相 +省 +看 +着 +矫 +禁 +种 +称 +空 +答 +粘 +糊 +系 +累 +纤 +结 +给 +缝 +肖 +背 +脏 +舍 +色 +落 +蒙 +薄 +藏 +血 +行 +要 +观 +觉 +角 +解 +说 +调 +踏 +车 +转 +载 +还 +遂 +都 +重 +量 +钻 +铺 +长 +间 +降 +难 +露 +鲜 diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polydict.json b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polydict.json new file mode 100644 index 0000000000000000000000000000000000000000..903fd018067b185c8cb8cd8a5b6cf07822512989 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polydict.json @@ -0,0 +1,393 @@ +{ + "1": "丧{sang1}", + "2": "丧{sang4}", + "3": "中{zhong1}", + "4": "中{zhong4}", + "5": "为{wei2}", + "6": "为{wei4}", + "7": "乌{wu1}", + "8": "乌{wu4}", + "9": "乐{lao4}", + "10": "乐{le4}", + "11": "乐{le5}", + "12": "乐{yao4}", + "13": "乐{yve4}", + "14": "了{le5}", + "15": "了{liao3}", + "16": "了{liao5}", + "17": "什{shen2}", + "18": "什{shi2}", + "19": "仔{zai3}", + "20": "仔{zai5}", + "21": "仔{zi3}", + "22": "仔{zi5}", + "23": "令{ling2}", + "24": "令{ling4}", + "25": "任{ren2}", + "26": "任{ren4}", + "27": "会{hui4}", + "28": "会{hui5}", + "29": "会{kuai4}", + "30": "传{chuan2}", + "31": "传{zhuan4}", + "32": "佛{fo2}", + "33": "佛{fu2}", + "34": "供{gong1}", + "35": "供{gong4}", + "36": "便{bian4}", + "37": "便{pian2}", + "38": "倒{dao3}", + "39": "倒{dao4}", + "40": "假{jia3}", + "41": "假{jia4}", + "42": "兴{xing1}", + "43": "兴{xing4}", + "44": "冠{guan1}", + "45": "冠{guan4}", + "46": "冲{chong1}", + "47": "冲{chong4}", + "48": "几{ji1}", + "49": "几{ji2}", + "50": "几{ji3}", + "51": "分{fen1}", + "52": "分{fen4}", + "53": "分{fen5}", + "54": "切{qie1}", + "55": "切{qie4}", + "56": "划{hua2}", + "57": "划{hua4}", + "58": "划{hua5}", + "59": "创{chuang1}", + "60": "创{chuang4}", + "61": "剥{bao1}", + "62": "剥{bo1}", + "63": "勒{le4}", + "64": "勒{le5}", + "65": "勒{lei1}", + "66": "区{ou1}", + "67": "区{qu1}", + "68": "华{hua2}", + "69": "华{hua4}", + "70": "单{chan2}", + "71": "单{dan1}", + "72": "单{shan4}", + "73": "卜{bo5}", + "74": "卜{bu3}", + "75": "占{zhan1}", + "76": "占{zhan4}", + "77": "卡{ka2}", + "78": "卡{ka3}", + "79": "卡{qia3}", + "80": "卷{jvan3}", + "81": "卷{jvan4}", + "82": "厦{sha4}", + "83": "厦{xia4}", + "84": "参{can1}", + "85": "参{cen1}", + "86": "参{shen1}", + "87": "发{fa1}", + "88": "发{fa4}", + "89": "发{fa5}", + "90": "只{zhi1}", + "91": "只{zhi3}", + "92": "号{hao2}", + "93": "号{hao4}", + "94": "号{hao5}", + "95": "同{tong2}", + "96": "同{tong4}", + "97": "同{tong5}", + "98": "吐{tu2}", + "99": "吐{tu3}", + "100": "吐{tu4}", + "101": "和{he2}", + "102": "和{he4}", + "103": "和{he5}", + "104": "和{huo2}", + "105": "和{huo4}", + "106": "和{huo5}", + "107": "喝{he1}", + "108": "喝{he4}", + "109": "圈{jvan4}", + "110": "圈{qvan1}", + "111": "圈{qvan5}", + "112": "地{de5}", + "113": "地{di4}", + "114": "地{di5}", + "115": "塞{sai1}", + "116": "塞{sai2}", + "117": "塞{sai4}", + "118": "塞{se4}", + "119": "壳{ke2}", + "120": "壳{qiao4}", + "121": "处{chu3}", + "122": "处{chu4}", + "123": "奇{ji1}", + "124": "奇{qi2}", + "125": "奔{ben1}", + "126": "奔{ben4}", + "127": "好{hao3}", + "128": "好{hao4}", + "129": "好{hao5}", + "130": "宁{ning2}", + "131": "宁{ning4}", + "132": "宁{ning5}", + "133": "宿{su4}", + "134": "宿{xiu3}", + "135": "宿{xiu4}", + "136": "将{jiang1}", + "137": "将{jiang4}", + "138": "少{shao3}", + "139": "少{shao4}", + "140": "尽{jin3}", + "141": "尽{jin4}", + "142": "岗{gang1}", + "143": "岗{gang3}", + "144": "差{cha1}", + "145": "差{cha4}", + "146": "差{chai1}", + "147": "差{ci1}", + "148": "巷{hang4}", + "149": "巷{xiang4}", + "150": "帖{tie1}", + "151": "帖{tie3}", + "152": "帖{tie4}", + "153": "干{gan1}", + "154": "干{gan4}", + "155": "应{ying1}", + "156": "应{ying4}", + "157": "应{ying5}", + "158": "度{du4}", + "159": "度{du5}", + "160": "度{duo2}", + "161": "弹{dan4}", + "162": "弹{tan2}", + "163": "弹{tan5}", + "164": "强{jiang4}", + "165": "强{qiang2}", + "166": "强{qiang3}", + "167": "当{dang1}", + "168": "当{dang4}", + "169": "当{dang5}", + "170": "待{dai1}", + "171": "待{dai4}", + "172": "得{de2}", + "173": "得{de5}", + "174": "得{dei3}", + "175": "得{dei5}", + "176": "恶{e3}", + "177": "恶{e4}", + "178": "恶{wu4}", + "179": "扁{bian3}", + "180": "扁{pian1}", + "181": "扇{shan1}", + "182": "扇{shan4}", + "183": "扎{za1}", + "184": "扎{zha1}", + "185": "扎{zha2}", + "186": "扫{sao3}", + "187": "扫{sao4}", + "188": "担{dan1}", + "189": "担{dan4}", + "190": "担{dan5}", + "191": "挑{tiao1}", + "192": "挑{tiao3}", + "193": "据{jv1}", + "194": "据{jv4}", + "195": "撒{sa1}", + "196": "撒{sa3}", + "197": "撒{sa5}", + "198": "教{jiao1}", + "199": "教{jiao4}", + "200": "散{san3}", + "201": "散{san4}", + "202": "散{san5}", + "203": "数{shu3}", + "204": "数{shu4}", + "205": "数{shu5}", + "206": "斗{dou3}", + "207": "斗{dou4}", + "208": "晃{huang3}", + "209": "曝{bao4}", + "210": "曲{qu1}", + "211": "曲{qu3}", + "212": "更{geng1}", + "213": "更{geng4}", + "214": "曾{ceng1}", + "215": "曾{ceng2}", + "216": "曾{zeng1}", + "217": "朝{chao2}", + "218": "朝{zhao1}", + "219": "朴{piao2}", + "220": "朴{pu2}", + "221": "朴{pu3}", + "222": "杆{gan1}", + "223": "杆{gan3}", + "224": "查{cha2}", + "225": "查{zha1}", + "226": "校{jiao4}", + "227": "校{xiao4}", + "228": "模{mo2}", + "229": "模{mu2}", + "230": "横{heng2}", + "231": "横{heng4}", + "232": "没{mei2}", + "233": "没{mo4}", + "234": "泡{pao1}", + "235": "泡{pao4}", + "236": "泡{pao5}", + "237": "济{ji3}", + "238": "济{ji4}", + "239": "混{hun2}", + "240": "混{hun3}", + "241": "混{hun4}", + "242": "混{hun5}", + "243": "漂{piao1}", + "244": "漂{piao3}", + "245": "漂{piao4}", + "246": "炸{zha2}", + "247": "炸{zha4}", + "248": "熟{shou2}", + "249": "熟{shu2}", + "250": "燕{yan1}", + "251": "燕{yan4}", + "252": "片{pian1}", + "253": "片{pian4}", + "254": "率{lv4}", + "255": "率{shuai4}", + "256": "畜{chu4}", + "257": "畜{xu4}", + "258": "的{de5}", + "259": "的{di1}", + "260": "的{di2}", + "261": "的{di4}", + "262": "的{di5}", + "263": "盛{cheng2}", + "264": "盛{sheng4}", + "265": "相{xiang1}", + "266": "相{xiang4}", + "267": "相{xiang5}", + "268": "省{sheng3}", + "269": "省{xing3}", + "270": "看{kan1}", + "271": "看{kan4}", + "272": "看{kan5}", + "273": "着{zhao1}", + "274": "着{zhao2}", + "275": "着{zhao5}", + "276": "着{zhe5}", + "277": "着{zhuo2}", + "278": "着{zhuo5}", + "279": "矫{jiao3}", + "280": "禁{jin1}", + "281": "禁{jin4}", + "282": "种{zhong3}", + "283": "种{zhong4}", + "284": "称{chen4}", + "285": "称{cheng1}", + "286": "空{kong1}", + "287": "空{kong4}", + "288": "答{da1}", + "289": "答{da2}", + "290": "粘{nian2}", + "291": "粘{zhan1}", + "292": "糊{hu2}", + "293": "糊{hu5}", + "294": "系{ji4}", + "295": "系{xi4}", + "296": "系{xi5}", + "297": "累{lei2}", + "298": "累{lei3}", + "299": "累{lei4}", + "300": "累{lei5}", + "301": "纤{qian4}", + "302": "纤{xian1}", + "303": "结{jie1}", + "304": "结{jie2}", + "305": "结{jie5}", + "306": "给{gei3}", + "307": "给{gei5}", + "308": "给{ji3}", + "309": "缝{feng2}", + "310": "缝{feng4}", + "311": "缝{feng5}", + "312": "肖{xiao1}", + "313": "肖{xiao4}", + "314": "背{bei1}", + "315": "背{bei4}", + "316": "脏{zang1}", + "317": "脏{zang4}", + "318": "舍{she3}", + "319": "舍{she4}", + "320": "色{se4}", + "321": "色{shai3}", + "322": "落{lao4}", + "323": "落{luo4}", + "324": "蒙{meng1}", + "325": "蒙{meng2}", + "326": "蒙{meng3}", + "327": "薄{bao2}", + "328": "薄{bo2}", + "329": "薄{bo4}", + "330": "藏{cang2}", + "331": "藏{zang4}", + "332": "血{xie3}", + "333": "血{xue4}", + "334": "行{hang2}", + "335": "行{hang5}", + "336": "行{heng5}", + "337": "行{xing2}", + "338": "行{xing4}", + "339": "要{yao1}", + "340": "要{yao4}", + "341": "观{guan1}", + "342": "观{guan4}", + "343": "觉{jiao4}", + "344": "觉{jiao5}", + "345": "觉{jve2}", + "346": "角{jiao3}", + "347": "角{jve2}", + "348": "解{jie3}", + "349": "解{jie4}", + "350": "解{xie4}", + "351": "说{shui4}", + "352": "说{shuo1}", + "353": "调{diao4}", + "354": "调{tiao2}", + "355": "踏{ta1}", + "356": "踏{ta4}", + "357": "车{che1}", + "358": "车{jv1}", + "359": "转{zhuan3}", + "360": "转{zhuan4}", + "361": "载{zai3}", + "362": "载{zai4}", + "363": "还{hai2}", + "364": "还{huan2}", + "365": "遂{sui2}", + "366": "遂{sui4}", + "367": "都{dou1}", + "368": "都{du1}", + "369": "重{chong2}", + "370": "重{zhong4}", + "371": "量{liang2}", + "372": "量{liang4}", + "373": "量{liang5}", + "374": "钻{zuan1}", + "375": "钻{zuan4}", + "376": "铺{pu1}", + "377": "铺{pu4}", + "378": "长{chang2}", + "379": "长{chang3}", + "380": "长{zhang3}", + "381": "间{jian1}", + "382": "间{jian4}", + "383": "降{jiang4}", + "384": "降{xiang2}", + "385": "难{nan2}", + "386": "难{nan4}", + "387": "难{nan5}", + "388": "露{lou4}", + "389": "露{lu4}", + "390": "鲜{xian1}", + "391": "鲜{xian3}" +} \ No newline at end of file diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polydict_r.json b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polydict_r.json new file mode 100644 index 0000000000000000000000000000000000000000..aabbe6257493eaee7d3f0b77f78f0cb006e89fb6 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/polydict_r.json @@ -0,0 +1,393 @@ +{ + "丧{sang1}": 1, + "丧{sang4}": 2, + "中{zhong1}": 3, + "中{zhong4}": 4, + "为{wei2}": 5, + "为{wei4}": 6, + "乌{wu1}": 7, + "乌{wu4}": 8, + "乐{lao4}": 9, + "乐{le4}": 10, + "乐{le5}": 11, + "乐{yao4}": 12, + "乐{yve4}": 13, + "了{le5}": 14, + "了{liao3}": 15, + "了{liao5}": 16, + "什{shen2}": 17, + "什{shi2}": 18, + "仔{zai3}": 19, + "仔{zai5}": 20, + "仔{zi3}": 21, + "仔{zi5}": 22, + "令{ling2}": 23, + "令{ling4}": 24, + "任{ren2}": 25, + "任{ren4}": 26, + "会{hui4}": 27, + "会{hui5}": 28, + "会{kuai4}": 29, + "传{chuan2}": 30, + "传{zhuan4}": 31, + "佛{fo2}": 32, + "佛{fu2}": 33, + "供{gong1}": 34, + "供{gong4}": 35, + "便{bian4}": 36, + "便{pian2}": 37, + "倒{dao3}": 38, + "倒{dao4}": 39, + "假{jia3}": 40, + "假{jia4}": 41, + "兴{xing1}": 42, + "兴{xing4}": 43, + "冠{guan1}": 44, + "冠{guan4}": 45, + "冲{chong1}": 46, + "冲{chong4}": 47, + "几{ji1}": 48, + "几{ji2}": 49, + "几{ji3}": 50, + "分{fen1}": 51, + "分{fen4}": 52, + "分{fen5}": 53, + "切{qie1}": 54, + "切{qie4}": 55, + "划{hua2}": 56, + "划{hua4}": 57, + "划{hua5}": 58, + "创{chuang1}": 59, + "创{chuang4}": 60, + "剥{bao1}": 61, + "剥{bo1}": 62, + "勒{le4}": 63, + "勒{le5}": 64, + "勒{lei1}": 65, + "区{ou1}": 66, + "区{qu1}": 67, + "华{hua2}": 68, + "华{hua4}": 69, + "单{chan2}": 70, + "单{dan1}": 71, + "单{shan4}": 72, + "卜{bo5}": 73, + "卜{bu3}": 74, + "占{zhan1}": 75, + "占{zhan4}": 76, + "卡{ka2}": 77, + "卡{ka3}": 78, + "卡{qia3}": 79, + "卷{jvan3}": 80, + "卷{jvan4}": 81, + "厦{sha4}": 82, + "厦{xia4}": 83, + "参{can1}": 84, + "参{cen1}": 85, + "参{shen1}": 86, + "发{fa1}": 87, + "发{fa4}": 88, + "发{fa5}": 89, + "只{zhi1}": 90, + "只{zhi3}": 91, + "号{hao2}": 92, + "号{hao4}": 93, + "号{hao5}": 94, + "同{tong2}": 95, + "同{tong4}": 96, + "同{tong5}": 97, + "吐{tu2}": 98, + "吐{tu3}": 99, + "吐{tu4}": 100, + "和{he2}": 101, + "和{he4}": 102, + "和{he5}": 103, + "和{huo2}": 104, + "和{huo4}": 105, + "和{huo5}": 106, + "喝{he1}": 107, + "喝{he4}": 108, + "圈{jvan4}": 109, + "圈{qvan1}": 110, + "圈{qvan5}": 111, + "地{de5}": 112, + "地{di4}": 113, + "地{di5}": 114, + "塞{sai1}": 115, + "塞{sai2}": 116, + "塞{sai4}": 117, + "塞{se4}": 118, + "壳{ke2}": 119, + "壳{qiao4}": 120, + "处{chu3}": 121, + "处{chu4}": 122, + "奇{ji1}": 123, + "奇{qi2}": 124, + "奔{ben1}": 125, + "奔{ben4}": 126, + "好{hao3}": 127, + "好{hao4}": 128, + "好{hao5}": 129, + "宁{ning2}": 130, + "宁{ning4}": 131, + "宁{ning5}": 132, + "宿{su4}": 133, + "宿{xiu3}": 134, + "宿{xiu4}": 135, + "将{jiang1}": 136, + "将{jiang4}": 137, + "少{shao3}": 138, + "少{shao4}": 139, + "尽{jin3}": 140, + "尽{jin4}": 141, + "岗{gang1}": 142, + "岗{gang3}": 143, + "差{cha1}": 144, + "差{cha4}": 145, + "差{chai1}": 146, + "差{ci1}": 147, + "巷{hang4}": 148, + "巷{xiang4}": 149, + "帖{tie1}": 150, + "帖{tie3}": 151, + "帖{tie4}": 152, + "干{gan1}": 153, + "干{gan4}": 154, + "应{ying1}": 155, + "应{ying4}": 156, + "应{ying5}": 157, + "度{du4}": 158, + "度{du5}": 159, + "度{duo2}": 160, + "弹{dan4}": 161, + "弹{tan2}": 162, + "弹{tan5}": 163, + "强{jiang4}": 164, + "强{qiang2}": 165, + "强{qiang3}": 166, + "当{dang1}": 167, + "当{dang4}": 168, + "当{dang5}": 169, + "待{dai1}": 170, + "待{dai4}": 171, + "得{de2}": 172, + "得{de5}": 173, + "得{dei3}": 174, + "得{dei5}": 175, + "恶{e3}": 176, + "恶{e4}": 177, + "恶{wu4}": 178, + "扁{bian3}": 179, + "扁{pian1}": 180, + "扇{shan1}": 181, + "扇{shan4}": 182, + "扎{za1}": 183, + "扎{zha1}": 184, + "扎{zha2}": 185, + "扫{sao3}": 186, + "扫{sao4}": 187, + "担{dan1}": 188, + "担{dan4}": 189, + "担{dan5}": 190, + "挑{tiao1}": 191, + "挑{tiao3}": 192, + "据{jv1}": 193, + "据{jv4}": 194, + "撒{sa1}": 195, + "撒{sa3}": 196, + "撒{sa5}": 197, + "教{jiao1}": 198, + "教{jiao4}": 199, + "散{san3}": 200, + "散{san4}": 201, + "散{san5}": 202, + "数{shu3}": 203, + "数{shu4}": 204, + "数{shu5}": 205, + "斗{dou3}": 206, + "斗{dou4}": 207, + "晃{huang3}": 208, + "曝{bao4}": 209, + "曲{qu1}": 210, + "曲{qu3}": 211, + "更{geng1}": 212, + "更{geng4}": 213, + "曾{ceng1}": 214, + "曾{ceng2}": 215, + "曾{zeng1}": 216, + "朝{chao2}": 217, + "朝{zhao1}": 218, + "朴{piao2}": 219, + "朴{pu2}": 220, + "朴{pu3}": 221, + "杆{gan1}": 222, + "杆{gan3}": 223, + "查{cha2}": 224, + "查{zha1}": 225, + "校{jiao4}": 226, + "校{xiao4}": 227, + "模{mo2}": 228, + "模{mu2}": 229, + "横{heng2}": 230, + "横{heng4}": 231, + "没{mei2}": 232, + "没{mo4}": 233, + "泡{pao1}": 234, + "泡{pao4}": 235, + "泡{pao5}": 236, + "济{ji3}": 237, + "济{ji4}": 238, + "混{hun2}": 239, + "混{hun3}": 240, + "混{hun4}": 241, + "混{hun5}": 242, + "漂{piao1}": 243, + "漂{piao3}": 244, + "漂{piao4}": 245, + "炸{zha2}": 246, + "炸{zha4}": 247, + "熟{shou2}": 248, + "熟{shu2}": 249, + "燕{yan1}": 250, + "燕{yan4}": 251, + "片{pian1}": 252, + "片{pian4}": 253, + "率{lv4}": 254, + "率{shuai4}": 255, + "畜{chu4}": 256, + "畜{xu4}": 257, + "的{de5}": 258, + "的{di1}": 259, + "的{di2}": 260, + "的{di4}": 261, + "的{di5}": 262, + "盛{cheng2}": 263, + "盛{sheng4}": 264, + "相{xiang1}": 265, + "相{xiang4}": 266, + "相{xiang5}": 267, + "省{sheng3}": 268, + "省{xing3}": 269, + "看{kan1}": 270, + "看{kan4}": 271, + "看{kan5}": 272, + "着{zhao1}": 273, + "着{zhao2}": 274, + "着{zhao5}": 275, + "着{zhe5}": 276, + "着{zhuo2}": 277, + "着{zhuo5}": 278, + "矫{jiao3}": 279, + "禁{jin1}": 280, + "禁{jin4}": 281, + "种{zhong3}": 282, + "种{zhong4}": 283, + "称{chen4}": 284, + "称{cheng1}": 285, + "空{kong1}": 286, + "空{kong4}": 287, + "答{da1}": 288, + "答{da2}": 289, + "粘{nian2}": 290, + "粘{zhan1}": 291, + "糊{hu2}": 292, + "糊{hu5}": 293, + "系{ji4}": 294, + "系{xi4}": 295, + "系{xi5}": 296, + "累{lei2}": 297, + "累{lei3}": 298, + "累{lei4}": 299, + "累{lei5}": 300, + "纤{qian4}": 301, + "纤{xian1}": 302, + "结{jie1}": 303, + "结{jie2}": 304, + "结{jie5}": 305, + "给{gei3}": 306, + "给{gei5}": 307, + "给{ji3}": 308, + "缝{feng2}": 309, + "缝{feng4}": 310, + "缝{feng5}": 311, + "肖{xiao1}": 312, + "肖{xiao4}": 313, + "背{bei1}": 314, + "背{bei4}": 315, + "脏{zang1}": 316, + "脏{zang4}": 317, + "舍{she3}": 318, + "舍{she4}": 319, + "色{se4}": 320, + "色{shai3}": 321, + "落{lao4}": 322, + "落{luo4}": 323, + "蒙{meng1}": 324, + "蒙{meng2}": 325, + "蒙{meng3}": 326, + "薄{bao2}": 327, + "薄{bo2}": 328, + "薄{bo4}": 329, + "藏{cang2}": 330, + "藏{zang4}": 331, + "血{xie3}": 332, + "血{xue4}": 333, + "行{hang2}": 334, + "行{hang5}": 335, + "行{heng5}": 336, + "行{xing2}": 337, + "行{xing4}": 338, + "要{yao1}": 339, + "要{yao4}": 340, + "观{guan1}": 341, + "观{guan4}": 342, + "觉{jiao4}": 343, + "觉{jiao5}": 344, + "觉{jve2}": 345, + "角{jiao3}": 346, + "角{jve2}": 347, + "解{jie3}": 348, + "解{jie4}": 349, + "解{xie4}": 350, + "说{shui4}": 351, + "说{shuo1}": 352, + "调{diao4}": 353, + "调{tiao2}": 354, + "踏{ta1}": 355, + "踏{ta4}": 356, + "车{che1}": 357, + "车{jv1}": 358, + "转{zhuan3}": 359, + "转{zhuan4}": 360, + "载{zai3}": 361, + "载{zai4}": 362, + "还{hai2}": 363, + "还{huan2}": 364, + "遂{sui2}": 365, + "遂{sui4}": 366, + "都{dou1}": 367, + "都{du1}": 368, + "重{chong2}": 369, + "重{zhong4}": 370, + "量{liang2}": 371, + "量{liang4}": 372, + "量{liang5}": 373, + "钻{zuan1}": 374, + "钻{zuan4}": 375, + "铺{pu1}": 376, + "铺{pu4}": 377, + "长{chang2}": 378, + "长{chang3}": 379, + "长{zhang3}": 380, + "间{jian1}": 381, + "间{jian4}": 382, + "降{jiang4}": 383, + "降{xiang2}": 384, + "难{nan2}": 385, + "难{nan4}": 386, + "难{nan5}": 387, + "露{lou4}": 388, + "露{lu4}": 389, + "鲜{xian1}": 390, + "鲜{xian3}": 391 +} \ No newline at end of file diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/vocab.txt b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/vocab.txt new file mode 100644 index 0000000000000000000000000000000000000000..ca4f9781030019ab9b253c6dcb8c7878b6dc87a5 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/g2p_chinese_model/vocab.txt @@ -0,0 +1,21128 @@ +[PAD] +[unused1] +[unused2] +[unused3] +[unused4] +[unused5] +[unused6] +[unused7] +[unused8] +[unused9] +[unused10] +[unused11] +[unused12] +[unused13] +[unused14] +[unused15] +[unused16] +[unused17] +[unused18] +[unused19] +[unused20] +[unused21] +[unused22] +[unused23] +[unused24] +[unused25] +[unused26] +[unused27] +[unused28] +[unused29] +[unused30] +[unused31] +[unused32] +[unused33] +[unused34] +[unused35] +[unused36] +[unused37] +[unused38] +[unused39] +[unused40] +[unused41] +[unused42] +[unused43] +[unused44] +[unused45] +[unused46] +[unused47] +[unused48] +[unused49] +[unused50] +[unused51] +[unused52] +[unused53] +[unused54] +[unused55] +[unused56] +[unused57] +[unused58] +[unused59] +[unused60] +[unused61] +[unused62] +[unused63] +[unused64] +[unused65] +[unused66] +[unused67] +[unused68] +[unused69] +[unused70] +[unused71] +[unused72] +[unused73] +[unused74] +[unused75] +[unused76] +[unused77] +[unused78] +[unused79] +[unused80] +[unused81] +[unused82] +[unused83] +[unused84] +[unused85] +[unused86] +[unused87] +[unused88] +[unused89] +[unused90] +[unused91] +[unused92] +[unused93] +[unused94] +[unused95] +[unused96] +[unused97] +[unused98] +[unused99] +[UNK] +[CLS] +[SEP] +[MASK] + + +! +" +# +$ +% +& +' +( +) +* ++ +, +- +. +/ +0 +1 +2 +3 +4 +5 +6 +7 +8 +9 +: +; +< += +> +? +@ +[ +\ +] +^ +_ +a +b +c +d +e +f +g +h +i +j +k +l +m +n +o +p +q +r +s +t +u +v +w +x +y +z +{ +| +} +~ +£ +¤ +¥ +§ +© +« +® +° +± +² +³ +µ +· +¹ +º +» +¼ +× +ß +æ +÷ +ø +đ +ŋ +ɔ +ə +ɡ +ʰ +ˇ +ˈ +ˊ +ˋ +ˍ +ː +˙ +˚ +ˢ +α +β +γ +δ +ε +η +θ +ι +κ +λ +μ +ν +ο +π +ρ +ς +σ +τ +υ +φ +χ +ψ +ω +а +б +в +г +д +е +ж +з +и +к +л +м +н +о +п +р +с +т +у +ф +х +ц +ч +ш +ы +ь +я +і +ا +ب +ة +ت +د +ر +س +ع +ل +م +ن +ه +و +ي +۩ +ก +ง +น +ม +ย +ร +อ +า +เ +๑ +་ +ღ +ᄀ +ᄁ +ᄂ +ᄃ +ᄅ +ᄆ +ᄇ +ᄈ +ᄉ +ᄋ +ᄌ +ᄎ +ᄏ +ᄐ +ᄑ +ᄒ +ᅡ +ᅢ +ᅣ +ᅥ +ᅦ +ᅧ +ᅨ +ᅩ +ᅪ +ᅬ +ᅭ +ᅮ +ᅯ +ᅲ +ᅳ +ᅴ +ᅵ +ᆨ +ᆫ +ᆯ +ᆷ +ᆸ +ᆺ +ᆻ +ᆼ +ᗜ +ᵃ +ᵉ +ᵍ +ᵏ +ᵐ +ᵒ +ᵘ +‖ +„ +† +• +‥ +‧ +
 +‰ +′ +″ +‹ +› +※ +‿ +⁄ +ⁱ +⁺ +ⁿ +₁ +₂ +₃ +₄ +€ +℃ +№ +™ +ⅰ +ⅱ +ⅲ +ⅳ +ⅴ +← +↑ +→ +↓ +↔ +↗ +↘ +⇒ +∀ +− +∕ +∙ +√ +∞ +∟ +∠ +∣ +∥ +∩ +∮ +∶ +∼ +∽ +≈ +≒ +≡ +≤ +≥ +≦ +≧ +≪ +≫ +⊙ +⋅ +⋈ +⋯ +⌒ +① +② +③ +④ +⑤ +⑥ +⑦ +⑧ +⑨ +⑩ +⑴ +⑵ +⑶ +⑷ +⑸ +⒈ +⒉ +⒊ +⒋ +ⓒ +ⓔ +ⓘ +─ +━ +│ +┃ +┅ +┆ +┊ +┌ +└ +├ +┣ +═ +║ +╚ +╞ +╠ +╭ +╮ +╯ +╰ +╱ +╳ +▂ +▃ +▅ +▇ +█ +▉ +▋ +▌ +▍ +▎ +■ +□ +▪ +▫ +▬ +▲ +△ +▶ +► +▼ +▽ +◆ +◇ +○ +◎ +● +◕ +◠ +◢ +◤ +☀ +★ +☆ +☕ +☞ +☺ +☼ +♀ +♂ +♠ +♡ +♣ +♥ +♦ +♪ +♫ +♬ +✈ +✔ +✕ +✖ +✦ +✨ +✪ +✰ +✿ +❀ +❤ +➜ +➤ +⦿ +、 +。 +〃 +々 +〇 +〈 +〉 +《 +》 +「 +」 +『 +』 +【 +】 +〓 +〔 +〕 +〖 +〗 +〜 +〝 +〞 +ぁ +あ +ぃ +い +う +ぇ +え +お +か +き +く +け +こ +さ +し +す +せ +そ +た +ち +っ +つ +て +と +な +に +ぬ +ね +の +は +ひ +ふ +へ +ほ +ま +み +む +め +も +ゃ +や +ゅ +ゆ +ょ +よ +ら +り +る +れ +ろ +わ +を +ん +゜ +ゝ +ァ +ア +ィ +イ +ゥ +ウ +ェ +エ +ォ +オ +カ +キ +ク +ケ +コ +サ +シ +ス +セ +ソ +タ +チ +ッ +ツ +テ +ト +ナ +ニ +ヌ +ネ +ノ +ハ +ヒ +フ +ヘ +ホ +マ +ミ +ム +メ +モ +ャ +ヤ +ュ +ユ +ョ +ヨ +ラ +リ +ル +レ +ロ +ワ +ヲ +ン +ヶ +・ +ー +ヽ +ㄅ +ㄆ +ㄇ +ㄉ +ㄋ +ㄌ +ㄍ +ㄎ +ㄏ +ㄒ +ㄚ +ㄛ +ㄞ +ㄟ +ㄢ +ㄤ +ㄥ +ㄧ +ㄨ +ㆍ +㈦ +㊣ +㎡ +㗎 +一 +丁 +七 +万 +丈 +三 +上 +下 +不 +与 +丐 +丑 +专 +且 +丕 +世 +丘 +丙 +业 +丛 +东 +丝 +丞 +丟 +両 +丢 +两 +严 +並 +丧 +丨 +个 +丫 +中 +丰 +串 +临 +丶 +丸 +丹 +为 +主 +丼 +丽 +举 +丿 +乂 +乃 +久 +么 +义 +之 +乌 +乍 +乎 +乏 +乐 +乒 +乓 +乔 +乖 +乗 +乘 +乙 +乜 +九 +乞 +也 +习 +乡 +书 +乩 +买 +乱 +乳 +乾 +亀 +亂 +了 +予 +争 +事 +二 +于 +亏 +云 +互 +五 +井 +亘 +亙 +亚 +些 +亜 +亞 +亟 +亡 +亢 +交 +亥 +亦 +产 +亨 +亩 +享 +京 +亭 +亮 +亲 +亳 +亵 +人 +亿 +什 +仁 +仃 +仄 +仅 +仆 +仇 +今 +介 +仍 +从 +仏 +仑 +仓 +仔 +仕 +他 +仗 +付 +仙 +仝 +仞 +仟 +代 +令 +以 +仨 +仪 +们 +仮 +仰 +仲 +件 +价 +任 +份 +仿 +企 +伉 +伊 +伍 +伎 +伏 +伐 +休 +伕 +众 +优 +伙 +会 +伝 +伞 +伟 +传 +伢 +伤 +伦 +伪 +伫 +伯 +估 +伴 +伶 +伸 +伺 +似 +伽 +佃 +但 +佇 +佈 +位 +低 +住 +佐 +佑 +体 +佔 +何 +佗 +佘 +余 +佚 +佛 +作 +佝 +佞 +佟 +你 +佢 +佣 +佤 +佥 +佩 +佬 +佯 +佰 +佳 +併 +佶 +佻 +佼 +使 +侃 +侄 +來 +侈 +例 +侍 +侏 +侑 +侖 +侗 +供 +依 +侠 +価 +侣 +侥 +侦 +侧 +侨 +侬 +侮 +侯 +侵 +侶 +侷 +便 +係 +促 +俄 +俊 +俎 +俏 +俐 +俑 +俗 +俘 +俚 +保 +俞 +俟 +俠 +信 +俨 +俩 +俪 +俬 +俭 +修 +俯 +俱 +俳 +俸 +俺 +俾 +倆 +倉 +個 +倌 +倍 +倏 +們 +倒 +倔 +倖 +倘 +候 +倚 +倜 +借 +倡 +値 +倦 +倩 +倪 +倫 +倬 +倭 +倶 +债 +值 +倾 +偃 +假 +偈 +偉 +偌 +偎 +偏 +偕 +做 +停 +健 +側 +偵 +偶 +偷 +偻 +偽 +偿 +傀 +傅 +傍 +傑 +傘 +備 +傚 +傢 +傣 +傥 +储 +傩 +催 +傭 +傲 +傳 +債 +傷 +傻 +傾 +僅 +働 +像 +僑 +僕 +僖 +僚 +僥 +僧 +僭 +僮 +僱 +僵 +價 +僻 +儀 +儂 +億 +儆 +儉 +儋 +儒 +儕 +儘 +償 +儡 +優 +儲 +儷 +儼 +儿 +兀 +允 +元 +兄 +充 +兆 +兇 +先 +光 +克 +兌 +免 +児 +兑 +兒 +兔 +兖 +党 +兜 +兢 +入 +內 +全 +兩 +八 +公 +六 +兮 +兰 +共 +兲 +关 +兴 +兵 +其 +具 +典 +兹 +养 +兼 +兽 +冀 +内 +円 +冇 +冈 +冉 +冊 +册 +再 +冏 +冒 +冕 +冗 +写 +军 +农 +冠 +冢 +冤 +冥 +冨 +冪 +冬 +冯 +冰 +冲 +决 +况 +冶 +冷 +冻 +冼 +冽 +冾 +净 +凄 +准 +凇 +凈 +凉 +凋 +凌 +凍 +减 +凑 +凛 +凜 +凝 +几 +凡 +凤 +処 +凪 +凭 +凯 +凰 +凱 +凳 +凶 +凸 +凹 +出 +击 +函 +凿 +刀 +刁 +刃 +分 +切 +刈 +刊 +刍 +刎 +刑 +划 +列 +刘 +则 +刚 +创 +初 +删 +判 +別 +刨 +利 +刪 +别 +刮 +到 +制 +刷 +券 +刹 +刺 +刻 +刽 +剁 +剂 +剃 +則 +剉 +削 +剋 +剌 +前 +剎 +剐 +剑 +剔 +剖 +剛 +剜 +剝 +剣 +剤 +剥 +剧 +剩 +剪 +副 +割 +創 +剷 +剽 +剿 +劃 +劇 +劈 +劉 +劊 +劍 +劏 +劑 +力 +劝 +办 +功 +加 +务 +劣 +动 +助 +努 +劫 +劭 +励 +劲 +劳 +労 +劵 +効 +劾 +势 +勁 +勃 +勇 +勉 +勋 +勐 +勒 +動 +勖 +勘 +務 +勛 +勝 +勞 +募 +勢 +勤 +勧 +勳 +勵 +勸 +勺 +勻 +勾 +勿 +匀 +包 +匆 +匈 +匍 +匐 +匕 +化 +北 +匙 +匝 +匠 +匡 +匣 +匪 +匮 +匯 +匱 +匹 +区 +医 +匾 +匿 +區 +十 +千 +卅 +升 +午 +卉 +半 +卍 +华 +协 +卑 +卒 +卓 +協 +单 +卖 +南 +単 +博 +卜 +卞 +卟 +占 +卡 +卢 +卤 +卦 +卧 +卫 +卮 +卯 +印 +危 +即 +却 +卵 +卷 +卸 +卻 +卿 +厂 +厄 +厅 +历 +厉 +压 +厌 +厕 +厘 +厚 +厝 +原 +厢 +厥 +厦 +厨 +厩 +厭 +厮 +厲 +厳 +去 +县 +叁 +参 +參 +又 +叉 +及 +友 +双 +反 +収 +发 +叔 +取 +受 +变 +叙 +叛 +叟 +叠 +叡 +叢 +口 +古 +句 +另 +叨 +叩 +只 +叫 +召 +叭 +叮 +可 +台 +叱 +史 +右 +叵 +叶 +号 +司 +叹 +叻 +叼 +叽 +吁 +吃 +各 +吆 +合 +吉 +吊 +吋 +同 +名 +后 +吏 +吐 +向 +吒 +吓 +吕 +吖 +吗 +君 +吝 +吞 +吟 +吠 +吡 +否 +吧 +吨 +吩 +含 +听 +吭 +吮 +启 +吱 +吳 +吴 +吵 +吶 +吸 +吹 +吻 +吼 +吽 +吾 +呀 +呂 +呃 +呆 +呈 +告 +呋 +呎 +呐 +呓 +呕 +呗 +员 +呛 +呜 +呢 +呤 +呦 +周 +呱 +呲 +味 +呵 +呷 +呸 +呻 +呼 +命 +咀 +咁 +咂 +咄 +咆 +咋 +和 +咎 +咏 +咐 +咒 +咔 +咕 +咖 +咗 +咘 +咙 +咚 +咛 +咣 +咤 +咦 +咧 +咨 +咩 +咪 +咫 +咬 +咭 +咯 +咱 +咲 +咳 +咸 +咻 +咽 +咿 +哀 +品 +哂 +哄 +哆 +哇 +哈 +哉 +哋 +哌 +响 +哎 +哏 +哐 +哑 +哒 +哔 +哗 +哟 +員 +哥 +哦 +哧 +哨 +哩 +哪 +哭 +哮 +哲 +哺 +哼 +哽 +唁 +唄 +唆 +唇 +唉 +唏 +唐 +唑 +唔 +唠 +唤 +唧 +唬 +售 +唯 +唰 +唱 +唳 +唷 +唸 +唾 +啃 +啄 +商 +啉 +啊 +問 +啓 +啕 +啖 +啜 +啞 +啟 +啡 +啤 +啥 +啦 +啧 +啪 +啫 +啬 +啮 +啰 +啱 +啲 +啵 +啶 +啷 +啸 +啻 +啼 +啾 +喀 +喂 +喃 +善 +喆 +喇 +喉 +喊 +喋 +喎 +喏 +喔 +喘 +喙 +喚 +喜 +喝 +喟 +喧 +喪 +喫 +喬 +單 +喰 +喱 +喲 +喳 +喵 +営 +喷 +喹 +喺 +喻 +喽 +嗅 +嗆 +嗇 +嗎 +嗑 +嗒 +嗓 +嗔 +嗖 +嗚 +嗜 +嗝 +嗟 +嗡 +嗣 +嗤 +嗦 +嗨 +嗪 +嗬 +嗯 +嗰 +嗲 +嗳 +嗶 +嗷 +嗽 +嘀 +嘅 +嘆 +嘈 +嘉 +嘌 +嘍 +嘎 +嘔 +嘖 +嘗 +嘘 +嘚 +嘛 +嘜 +嘞 +嘟 +嘢 +嘣 +嘤 +嘧 +嘩 +嘭 +嘮 +嘯 +嘰 +嘱 +嘲 +嘴 +嘶 +嘸 +嘹 +嘻 +嘿 +噁 +噌 +噎 +噓 +噔 +噗 +噙 +噜 +噠 +噢 +噤 +器 +噩 +噪 +噬 +噱 +噴 +噶 +噸 +噹 +噻 +噼 +嚀 +嚇 +嚎 +嚏 +嚐 +嚓 +嚕 +嚟 +嚣 +嚥 +嚨 +嚮 +嚴 +嚷 +嚼 +囂 +囉 +囊 +囍 +囑 +囔 +囗 +囚 +四 +囝 +回 +囟 +因 +囡 +团 +団 +囤 +囧 +囪 +囫 +园 +困 +囱 +囲 +図 +围 +囹 +固 +国 +图 +囿 +圃 +圄 +圆 +圈 +國 +圍 +圏 +園 +圓 +圖 +團 +圜 +土 +圣 +圧 +在 +圩 +圭 +地 +圳 +场 +圻 +圾 +址 +坂 +均 +坊 +坍 +坎 +坏 +坐 +坑 +块 +坚 +坛 +坝 +坞 +坟 +坠 +坡 +坤 +坦 +坨 +坪 +坯 +坳 +坵 +坷 +垂 +垃 +垄 +型 +垒 +垚 +垛 +垠 +垢 +垣 +垦 +垩 +垫 +垭 +垮 +垵 +埂 +埃 +埋 +城 +埔 +埕 +埗 +域 +埠 +埤 +埵 +執 +埸 +培 +基 +埼 +堀 +堂 +堃 +堅 +堆 +堇 +堑 +堕 +堙 +堡 +堤 +堪 +堯 +堰 +報 +場 +堵 +堺 +堿 +塊 +塌 +塑 +塔 +塗 +塘 +塚 +塞 +塢 +塩 +填 +塬 +塭 +塵 +塾 +墀 +境 +墅 +墉 +墊 +墒 +墓 +増 +墘 +墙 +墜 +增 +墟 +墨 +墩 +墮 +墳 +墻 +墾 +壁 +壅 +壆 +壇 +壊 +壑 +壓 +壕 +壘 +壞 +壟 +壢 +壤 +壩 +士 +壬 +壮 +壯 +声 +売 +壳 +壶 +壹 +壺 +壽 +处 +备 +変 +复 +夏 +夔 +夕 +外 +夙 +多 +夜 +够 +夠 +夢 +夥 +大 +天 +太 +夫 +夭 +央 +夯 +失 +头 +夷 +夸 +夹 +夺 +夾 +奂 +奄 +奇 +奈 +奉 +奋 +奎 +奏 +奐 +契 +奔 +奕 +奖 +套 +奘 +奚 +奠 +奢 +奥 +奧 +奪 +奬 +奮 +女 +奴 +奶 +奸 +她 +好 +如 +妃 +妄 +妆 +妇 +妈 +妊 +妍 +妒 +妓 +妖 +妘 +妙 +妝 +妞 +妣 +妤 +妥 +妨 +妩 +妪 +妮 +妲 +妳 +妹 +妻 +妾 +姆 +姉 +姊 +始 +姍 +姐 +姑 +姒 +姓 +委 +姗 +姚 +姜 +姝 +姣 +姥 +姦 +姨 +姪 +姫 +姬 +姹 +姻 +姿 +威 +娃 +娄 +娅 +娆 +娇 +娉 +娑 +娓 +娘 +娛 +娜 +娟 +娠 +娣 +娥 +娩 +娱 +娲 +娴 +娶 +娼 +婀 +婁 +婆 +婉 +婊 +婕 +婚 +婢 +婦 +婧 +婪 +婭 +婴 +婵 +婶 +婷 +婺 +婿 +媒 +媚 +媛 +媞 +媧 +媲 +媳 +媽 +媾 +嫁 +嫂 +嫉 +嫌 +嫑 +嫔 +嫖 +嫘 +嫚 +嫡 +嫣 +嫦 +嫩 +嫲 +嫵 +嫻 +嬅 +嬉 +嬌 +嬗 +嬛 +嬢 +嬤 +嬪 +嬰 +嬴 +嬷 +嬸 +嬿 +孀 +孃 +子 +孑 +孔 +孕 +孖 +字 +存 +孙 +孚 +孛 +孜 +孝 +孟 +孢 +季 +孤 +学 +孩 +孪 +孫 +孬 +孰 +孱 +孳 +孵 +學 +孺 +孽 +孿 +宁 +它 +宅 +宇 +守 +安 +宋 +完 +宏 +宓 +宕 +宗 +官 +宙 +定 +宛 +宜 +宝 +实 +実 +宠 +审 +客 +宣 +室 +宥 +宦 +宪 +宫 +宮 +宰 +害 +宴 +宵 +家 +宸 +容 +宽 +宾 +宿 +寂 +寄 +寅 +密 +寇 +富 +寐 +寒 +寓 +寛 +寝 +寞 +察 +寡 +寢 +寥 +實 +寧 +寨 +審 +寫 +寬 +寮 +寰 +寵 +寶 +寸 +对 +寺 +寻 +导 +対 +寿 +封 +専 +射 +将 +將 +專 +尉 +尊 +尋 +對 +導 +小 +少 +尔 +尕 +尖 +尘 +尚 +尝 +尤 +尧 +尬 +就 +尴 +尷 +尸 +尹 +尺 +尻 +尼 +尽 +尾 +尿 +局 +屁 +层 +屄 +居 +屆 +屈 +屉 +届 +屋 +屌 +屍 +屎 +屏 +屐 +屑 +展 +屜 +属 +屠 +屡 +屢 +層 +履 +屬 +屯 +山 +屹 +屿 +岀 +岁 +岂 +岌 +岐 +岑 +岔 +岖 +岗 +岘 +岙 +岚 +岛 +岡 +岩 +岫 +岬 +岭 +岱 +岳 +岷 +岸 +峇 +峋 +峒 +峙 +峡 +峤 +峥 +峦 +峨 +峪 +峭 +峯 +峰 +峴 +島 +峻 +峽 +崁 +崂 +崆 +崇 +崎 +崑 +崔 +崖 +崗 +崙 +崛 +崧 +崩 +崭 +崴 +崽 +嵇 +嵊 +嵋 +嵌 +嵐 +嵘 +嵩 +嵬 +嵯 +嶂 +嶄 +嶇 +嶋 +嶙 +嶺 +嶼 +嶽 +巅 +巍 +巒 +巔 +巖 +川 +州 +巡 +巢 +工 +左 +巧 +巨 +巩 +巫 +差 +己 +已 +巳 +巴 +巷 +巻 +巽 +巾 +巿 +币 +市 +布 +帅 +帆 +师 +希 +帐 +帑 +帕 +帖 +帘 +帚 +帛 +帜 +帝 +帥 +带 +帧 +師 +席 +帮 +帯 +帰 +帳 +帶 +帷 +常 +帼 +帽 +幀 +幂 +幄 +幅 +幌 +幔 +幕 +幟 +幡 +幢 +幣 +幫 +干 +平 +年 +并 +幸 +幹 +幺 +幻 +幼 +幽 +幾 +广 +庁 +広 +庄 +庆 +庇 +床 +序 +庐 +库 +应 +底 +庖 +店 +庙 +庚 +府 +庞 +废 +庠 +度 +座 +庫 +庭 +庵 +庶 +康 +庸 +庹 +庾 +廁 +廂 +廃 +廈 +廉 +廊 +廓 +廖 +廚 +廝 +廟 +廠 +廢 +廣 +廬 +廳 +延 +廷 +建 +廿 +开 +弁 +异 +弃 +弄 +弈 +弊 +弋 +式 +弑 +弒 +弓 +弔 +引 +弗 +弘 +弛 +弟 +张 +弥 +弦 +弧 +弩 +弭 +弯 +弱 +張 +強 +弹 +强 +弼 +弾 +彅 +彆 +彈 +彌 +彎 +归 +当 +录 +彗 +彙 +彝 +形 +彤 +彥 +彦 +彧 +彩 +彪 +彫 +彬 +彭 +彰 +影 +彷 +役 +彻 +彼 +彿 +往 +征 +径 +待 +徇 +很 +徉 +徊 +律 +後 +徐 +徑 +徒 +従 +徕 +得 +徘 +徙 +徜 +從 +徠 +御 +徨 +復 +循 +徬 +微 +徳 +徴 +徵 +德 +徹 +徼 +徽 +心 +必 +忆 +忌 +忍 +忏 +忐 +忑 +忒 +忖 +志 +忘 +忙 +応 +忠 +忡 +忤 +忧 +忪 +快 +忱 +念 +忻 +忽 +忿 +怀 +态 +怂 +怅 +怆 +怎 +怏 +怒 +怔 +怕 +怖 +怙 +怜 +思 +怠 +怡 +急 +怦 +性 +怨 +怪 +怯 +怵 +总 +怼 +恁 +恃 +恆 +恋 +恍 +恐 +恒 +恕 +恙 +恚 +恢 +恣 +恤 +恥 +恨 +恩 +恪 +恫 +恬 +恭 +息 +恰 +恳 +恵 +恶 +恸 +恺 +恻 +恼 +恿 +悄 +悅 +悉 +悌 +悍 +悔 +悖 +悚 +悟 +悠 +患 +悦 +您 +悩 +悪 +悬 +悯 +悱 +悲 +悴 +悵 +悶 +悸 +悻 +悼 +悽 +情 +惆 +惇 +惊 +惋 +惑 +惕 +惘 +惚 +惜 +惟 +惠 +惡 +惦 +惧 +惨 +惩 +惫 +惬 +惭 +惮 +惯 +惰 +惱 +想 +惴 +惶 +惹 +惺 +愁 +愆 +愈 +愉 +愍 +意 +愕 +愚 +愛 +愜 +感 +愣 +愤 +愧 +愫 +愷 +愿 +慄 +慈 +態 +慌 +慎 +慑 +慕 +慘 +慚 +慟 +慢 +慣 +慧 +慨 +慫 +慮 +慰 +慳 +慵 +慶 +慷 +慾 +憂 +憊 +憋 +憎 +憐 +憑 +憔 +憚 +憤 +憧 +憨 +憩 +憫 +憬 +憲 +憶 +憾 +懂 +懇 +懈 +應 +懊 +懋 +懑 +懒 +懦 +懲 +懵 +懶 +懷 +懸 +懺 +懼 +懾 +懿 +戀 +戈 +戊 +戌 +戍 +戎 +戏 +成 +我 +戒 +戕 +或 +战 +戚 +戛 +戟 +戡 +戦 +截 +戬 +戮 +戰 +戲 +戳 +戴 +戶 +户 +戸 +戻 +戾 +房 +所 +扁 +扇 +扈 +扉 +手 +才 +扎 +扑 +扒 +打 +扔 +払 +托 +扛 +扣 +扦 +执 +扩 +扪 +扫 +扬 +扭 +扮 +扯 +扰 +扱 +扳 +扶 +批 +扼 +找 +承 +技 +抄 +抉 +把 +抑 +抒 +抓 +投 +抖 +抗 +折 +抚 +抛 +抜 +択 +抟 +抠 +抡 +抢 +护 +报 +抨 +披 +抬 +抱 +抵 +抹 +押 +抽 +抿 +拂 +拄 +担 +拆 +拇 +拈 +拉 +拋 +拌 +拍 +拎 +拐 +拒 +拓 +拔 +拖 +拗 +拘 +拙 +拚 +招 +拜 +拟 +拡 +拢 +拣 +拥 +拦 +拧 +拨 +择 +括 +拭 +拮 +拯 +拱 +拳 +拴 +拷 +拼 +拽 +拾 +拿 +持 +挂 +指 +挈 +按 +挎 +挑 +挖 +挙 +挚 +挛 +挝 +挞 +挟 +挠 +挡 +挣 +挤 +挥 +挨 +挪 +挫 +振 +挲 +挹 +挺 +挽 +挾 +捂 +捅 +捆 +捉 +捋 +捌 +捍 +捎 +捏 +捐 +捕 +捞 +损 +捡 +换 +捣 +捧 +捨 +捩 +据 +捱 +捲 +捶 +捷 +捺 +捻 +掀 +掂 +掃 +掇 +授 +掉 +掌 +掏 +掐 +排 +掖 +掘 +掙 +掛 +掠 +採 +探 +掣 +接 +控 +推 +掩 +措 +掬 +掰 +掲 +掳 +掴 +掷 +掸 +掺 +揀 +揃 +揄 +揆 +揉 +揍 +描 +提 +插 +揖 +揚 +換 +握 +揣 +揩 +揪 +揭 +揮 +援 +揶 +揸 +揹 +揽 +搀 +搁 +搂 +搅 +損 +搏 +搐 +搓 +搔 +搖 +搗 +搜 +搞 +搡 +搪 +搬 +搭 +搵 +搶 +携 +搽 +摀 +摁 +摄 +摆 +摇 +摈 +摊 +摒 +摔 +摘 +摞 +摟 +摧 +摩 +摯 +摳 +摸 +摹 +摺 +摻 +撂 +撃 +撅 +撇 +撈 +撐 +撑 +撒 +撓 +撕 +撚 +撞 +撤 +撥 +撩 +撫 +撬 +播 +撮 +撰 +撲 +撵 +撷 +撸 +撻 +撼 +撿 +擀 +擁 +擂 +擄 +擅 +擇 +擊 +擋 +操 +擎 +擒 +擔 +擘 +據 +擞 +擠 +擡 +擢 +擦 +擬 +擰 +擱 +擲 +擴 +擷 +擺 +擼 +擾 +攀 +攏 +攒 +攔 +攘 +攙 +攜 +攝 +攞 +攢 +攣 +攤 +攥 +攪 +攫 +攬 +支 +收 +攸 +改 +攻 +放 +政 +故 +效 +敌 +敍 +敎 +敏 +救 +敕 +敖 +敗 +敘 +教 +敛 +敝 +敞 +敢 +散 +敦 +敬 +数 +敲 +整 +敵 +敷 +數 +斂 +斃 +文 +斋 +斌 +斎 +斐 +斑 +斓 +斗 +料 +斛 +斜 +斟 +斡 +斤 +斥 +斧 +斩 +斫 +斬 +断 +斯 +新 +斷 +方 +於 +施 +旁 +旃 +旅 +旋 +旌 +旎 +族 +旖 +旗 +无 +既 +日 +旦 +旧 +旨 +早 +旬 +旭 +旮 +旱 +时 +旷 +旺 +旻 +昀 +昂 +昆 +昇 +昉 +昊 +昌 +明 +昏 +易 +昔 +昕 +昙 +星 +映 +春 +昧 +昨 +昭 +是 +昱 +昴 +昵 +昶 +昼 +显 +晁 +時 +晃 +晉 +晋 +晌 +晏 +晒 +晓 +晔 +晕 +晖 +晗 +晚 +晝 +晞 +晟 +晤 +晦 +晨 +晩 +普 +景 +晰 +晴 +晶 +晷 +智 +晾 +暂 +暄 +暇 +暈 +暉 +暌 +暐 +暑 +暖 +暗 +暝 +暢 +暧 +暨 +暫 +暮 +暱 +暴 +暸 +暹 +曄 +曆 +曇 +曉 +曖 +曙 +曜 +曝 +曠 +曦 +曬 +曰 +曲 +曳 +更 +書 +曹 +曼 +曾 +替 +最 +會 +月 +有 +朋 +服 +朐 +朔 +朕 +朗 +望 +朝 +期 +朦 +朧 +木 +未 +末 +本 +札 +朮 +术 +朱 +朴 +朵 +机 +朽 +杀 +杂 +权 +杆 +杈 +杉 +李 +杏 +材 +村 +杓 +杖 +杜 +杞 +束 +杠 +条 +来 +杨 +杭 +杯 +杰 +東 +杳 +杵 +杷 +杼 +松 +板 +极 +构 +枇 +枉 +枋 +析 +枕 +林 +枚 +果 +枝 +枢 +枣 +枪 +枫 +枭 +枯 +枰 +枱 +枳 +架 +枷 +枸 +柄 +柏 +某 +柑 +柒 +染 +柔 +柘 +柚 +柜 +柞 +柠 +柢 +查 +柩 +柬 +柯 +柱 +柳 +柴 +柵 +査 +柿 +栀 +栃 +栄 +栅 +标 +栈 +栉 +栋 +栎 +栏 +树 +栓 +栖 +栗 +校 +栩 +株 +样 +核 +根 +格 +栽 +栾 +桀 +桁 +桂 +桃 +桅 +框 +案 +桉 +桌 +桎 +桐 +桑 +桓 +桔 +桜 +桠 +桡 +桢 +档 +桥 +桦 +桧 +桨 +桩 +桶 +桿 +梁 +梅 +梆 +梏 +梓 +梗 +條 +梟 +梢 +梦 +梧 +梨 +梭 +梯 +械 +梳 +梵 +梶 +检 +棂 +棄 +棉 +棋 +棍 +棒 +棕 +棗 +棘 +棚 +棟 +棠 +棣 +棧 +森 +棱 +棲 +棵 +棹 +棺 +椁 +椅 +椋 +植 +椎 +椒 +検 +椪 +椭 +椰 +椹 +椽 +椿 +楂 +楊 +楓 +楔 +楚 +楝 +楞 +楠 +楣 +楨 +楫 +業 +楮 +極 +楷 +楸 +楹 +楼 +楽 +概 +榄 +榆 +榈 +榉 +榔 +榕 +榖 +榛 +榜 +榨 +榫 +榭 +榮 +榱 +榴 +榷 +榻 +槁 +槃 +構 +槌 +槍 +槎 +槐 +槓 +様 +槛 +槟 +槤 +槭 +槲 +槳 +槻 +槽 +槿 +樁 +樂 +樊 +樑 +樓 +標 +樞 +樟 +模 +樣 +権 +横 +樫 +樯 +樱 +樵 +樸 +樹 +樺 +樽 +樾 +橄 +橇 +橋 +橐 +橘 +橙 +機 +橡 +橢 +橫 +橱 +橹 +橼 +檀 +檄 +檎 +檐 +檔 +檗 +檜 +檢 +檬 +檯 +檳 +檸 +檻 +櫃 +櫚 +櫛 +櫥 +櫸 +櫻 +欄 +權 +欒 +欖 +欠 +次 +欢 +欣 +欧 +欲 +欸 +欺 +欽 +款 +歆 +歇 +歉 +歌 +歎 +歐 +歓 +歙 +歛 +歡 +止 +正 +此 +步 +武 +歧 +歩 +歪 +歯 +歲 +歳 +歴 +歷 +歸 +歹 +死 +歼 +殁 +殃 +殆 +殇 +殉 +殊 +残 +殒 +殓 +殖 +殘 +殞 +殡 +殤 +殭 +殯 +殲 +殴 +段 +殷 +殺 +殼 +殿 +毀 +毁 +毂 +毅 +毆 +毋 +母 +毎 +每 +毒 +毓 +比 +毕 +毗 +毘 +毙 +毛 +毡 +毫 +毯 +毽 +氈 +氏 +氐 +民 +氓 +气 +氖 +気 +氙 +氛 +氟 +氡 +氢 +氣 +氤 +氦 +氧 +氨 +氪 +氫 +氮 +氯 +氰 +氲 +水 +氷 +永 +氹 +氾 +汀 +汁 +求 +汆 +汇 +汉 +汎 +汐 +汕 +汗 +汙 +汛 +汝 +汞 +江 +池 +污 +汤 +汨 +汩 +汪 +汰 +汲 +汴 +汶 +汹 +決 +汽 +汾 +沁 +沂 +沃 +沅 +沈 +沉 +沌 +沏 +沐 +沒 +沓 +沖 +沙 +沛 +沟 +没 +沢 +沣 +沥 +沦 +沧 +沪 +沫 +沭 +沮 +沱 +河 +沸 +油 +治 +沼 +沽 +沾 +沿 +況 +泄 +泉 +泊 +泌 +泓 +法 +泗 +泛 +泞 +泠 +泡 +波 +泣 +泥 +注 +泪 +泫 +泮 +泯 +泰 +泱 +泳 +泵 +泷 +泸 +泻 +泼 +泽 +泾 +洁 +洄 +洋 +洒 +洗 +洙 +洛 +洞 +津 +洩 +洪 +洮 +洱 +洲 +洵 +洶 +洸 +洹 +活 +洼 +洽 +派 +流 +浃 +浄 +浅 +浆 +浇 +浊 +测 +济 +浏 +浑 +浒 +浓 +浔 +浙 +浚 +浜 +浣 +浦 +浩 +浪 +浬 +浮 +浯 +浴 +海 +浸 +涂 +涅 +涇 +消 +涉 +涌 +涎 +涓 +涔 +涕 +涙 +涛 +涝 +涞 +涟 +涠 +涡 +涣 +涤 +润 +涧 +涨 +涩 +涪 +涮 +涯 +液 +涵 +涸 +涼 +涿 +淀 +淄 +淅 +淆 +淇 +淋 +淌 +淑 +淒 +淖 +淘 +淙 +淚 +淞 +淡 +淤 +淦 +淨 +淩 +淪 +淫 +淬 +淮 +深 +淳 +淵 +混 +淹 +淺 +添 +淼 +清 +済 +渉 +渊 +渋 +渍 +渎 +渐 +渔 +渗 +渙 +渚 +減 +渝 +渠 +渡 +渣 +渤 +渥 +渦 +温 +測 +渭 +港 +渲 +渴 +游 +渺 +渾 +湃 +湄 +湊 +湍 +湖 +湘 +湛 +湟 +湧 +湫 +湮 +湯 +湳 +湾 +湿 +満 +溃 +溅 +溉 +溏 +源 +準 +溜 +溝 +溟 +溢 +溥 +溧 +溪 +溫 +溯 +溱 +溴 +溶 +溺 +溼 +滁 +滂 +滄 +滅 +滇 +滋 +滌 +滑 +滓 +滔 +滕 +滙 +滚 +滝 +滞 +滟 +满 +滢 +滤 +滥 +滦 +滨 +滩 +滬 +滯 +滲 +滴 +滷 +滸 +滾 +滿 +漁 +漂 +漆 +漉 +漏 +漓 +演 +漕 +漠 +漢 +漣 +漩 +漪 +漫 +漬 +漯 +漱 +漲 +漳 +漸 +漾 +漿 +潆 +潇 +潋 +潍 +潑 +潔 +潘 +潛 +潜 +潞 +潟 +潢 +潤 +潦 +潧 +潭 +潮 +潰 +潴 +潸 +潺 +潼 +澀 +澄 +澆 +澈 +澍 +澎 +澗 +澜 +澡 +澤 +澧 +澱 +澳 +澹 +激 +濁 +濂 +濃 +濑 +濒 +濕 +濘 +濛 +濟 +濠 +濡 +濤 +濫 +濬 +濮 +濯 +濱 +濺 +濾 +瀅 +瀆 +瀉 +瀋 +瀏 +瀑 +瀕 +瀘 +瀚 +瀛 +瀝 +瀞 +瀟 +瀧 +瀨 +瀬 +瀰 +瀾 +灌 +灏 +灑 +灘 +灝 +灞 +灣 +火 +灬 +灭 +灯 +灰 +灵 +灶 +灸 +灼 +災 +灾 +灿 +炀 +炁 +炅 +炉 +炊 +炎 +炒 +炔 +炕 +炖 +炙 +炜 +炫 +炬 +炭 +炮 +炯 +炳 +炷 +炸 +点 +為 +炼 +炽 +烁 +烂 +烃 +烈 +烊 +烏 +烘 +烙 +烛 +烟 +烤 +烦 +烧 +烨 +烩 +烫 +烬 +热 +烯 +烷 +烹 +烽 +焉 +焊 +焕 +焖 +焗 +焘 +焙 +焚 +焜 +無 +焦 +焯 +焰 +焱 +然 +焼 +煅 +煉 +煊 +煌 +煎 +煒 +煖 +煙 +煜 +煞 +煤 +煥 +煦 +照 +煨 +煩 +煮 +煲 +煸 +煽 +熄 +熊 +熏 +熒 +熔 +熙 +熟 +熠 +熨 +熬 +熱 +熵 +熹 +熾 +燁 +燃 +燄 +燈 +燉 +燊 +燎 +燒 +燔 +燕 +燙 +燜 +營 +燥 +燦 +燧 +燭 +燮 +燴 +燻 +燼 +燿 +爆 +爍 +爐 +爛 +爪 +爬 +爭 +爰 +爱 +爲 +爵 +父 +爷 +爸 +爹 +爺 +爻 +爽 +爾 +牆 +片 +版 +牌 +牍 +牒 +牙 +牛 +牝 +牟 +牠 +牡 +牢 +牦 +牧 +物 +牯 +牲 +牴 +牵 +特 +牺 +牽 +犀 +犁 +犄 +犊 +犍 +犒 +犢 +犧 +犬 +犯 +状 +犷 +犸 +犹 +狀 +狂 +狄 +狈 +狎 +狐 +狒 +狗 +狙 +狞 +狠 +狡 +狩 +独 +狭 +狮 +狰 +狱 +狸 +狹 +狼 +狽 +猎 +猕 +猖 +猗 +猙 +猛 +猜 +猝 +猥 +猩 +猪 +猫 +猬 +献 +猴 +猶 +猷 +猾 +猿 +獄 +獅 +獎 +獐 +獒 +獗 +獠 +獣 +獨 +獭 +獰 +獲 +獵 +獷 +獸 +獺 +獻 +獼 +獾 +玄 +率 +玉 +王 +玑 +玖 +玛 +玟 +玠 +玥 +玩 +玫 +玮 +环 +现 +玲 +玳 +玷 +玺 +玻 +珀 +珂 +珅 +珈 +珉 +珊 +珍 +珏 +珐 +珑 +珙 +珞 +珠 +珣 +珥 +珩 +珪 +班 +珮 +珲 +珺 +現 +球 +琅 +理 +琇 +琉 +琊 +琍 +琏 +琐 +琛 +琢 +琥 +琦 +琨 +琪 +琬 +琮 +琰 +琲 +琳 +琴 +琵 +琶 +琺 +琼 +瑀 +瑁 +瑄 +瑋 +瑕 +瑗 +瑙 +瑚 +瑛 +瑜 +瑞 +瑟 +瑠 +瑣 +瑤 +瑩 +瑪 +瑯 +瑰 +瑶 +瑾 +璀 +璁 +璃 +璇 +璉 +璋 +璎 +璐 +璜 +璞 +璟 +璧 +璨 +環 +璽 +璿 +瓊 +瓏 +瓒 +瓜 +瓢 +瓣 +瓤 +瓦 +瓮 +瓯 +瓴 +瓶 +瓷 +甄 +甌 +甕 +甘 +甙 +甚 +甜 +生 +產 +産 +甥 +甦 +用 +甩 +甫 +甬 +甭 +甯 +田 +由 +甲 +申 +电 +男 +甸 +町 +画 +甾 +畀 +畅 +界 +畏 +畑 +畔 +留 +畜 +畝 +畢 +略 +畦 +番 +畫 +異 +畲 +畳 +畴 +當 +畸 +畹 +畿 +疆 +疇 +疊 +疏 +疑 +疔 +疖 +疗 +疙 +疚 +疝 +疟 +疡 +疣 +疤 +疥 +疫 +疮 +疯 +疱 +疲 +疳 +疵 +疸 +疹 +疼 +疽 +疾 +痂 +病 +症 +痈 +痉 +痊 +痍 +痒 +痔 +痕 +痘 +痙 +痛 +痞 +痠 +痢 +痣 +痤 +痧 +痨 +痪 +痫 +痰 +痱 +痴 +痹 +痺 +痼 +痿 +瘀 +瘁 +瘋 +瘍 +瘓 +瘘 +瘙 +瘟 +瘠 +瘡 +瘢 +瘤 +瘦 +瘧 +瘩 +瘪 +瘫 +瘴 +瘸 +瘾 +療 +癇 +癌 +癒 +癖 +癜 +癞 +癡 +癢 +癣 +癥 +癫 +癬 +癮 +癱 +癲 +癸 +発 +登 +發 +白 +百 +皂 +的 +皆 +皇 +皈 +皋 +皎 +皑 +皓 +皖 +皙 +皚 +皮 +皰 +皱 +皴 +皺 +皿 +盂 +盃 +盅 +盆 +盈 +益 +盎 +盏 +盐 +监 +盒 +盔 +盖 +盗 +盘 +盛 +盜 +盞 +盟 +盡 +監 +盤 +盥 +盧 +盪 +目 +盯 +盱 +盲 +直 +相 +盹 +盼 +盾 +省 +眈 +眉 +看 +県 +眙 +眞 +真 +眠 +眦 +眨 +眩 +眯 +眶 +眷 +眸 +眺 +眼 +眾 +着 +睁 +睇 +睏 +睐 +睑 +睛 +睜 +睞 +睡 +睢 +督 +睥 +睦 +睨 +睪 +睫 +睬 +睹 +睽 +睾 +睿 +瞄 +瞅 +瞇 +瞋 +瞌 +瞎 +瞑 +瞒 +瞓 +瞞 +瞟 +瞠 +瞥 +瞧 +瞩 +瞪 +瞬 +瞭 +瞰 +瞳 +瞻 +瞼 +瞿 +矇 +矍 +矗 +矚 +矛 +矜 +矢 +矣 +知 +矩 +矫 +短 +矮 +矯 +石 +矶 +矽 +矾 +矿 +码 +砂 +砌 +砍 +砒 +研 +砖 +砗 +砚 +砝 +砣 +砥 +砧 +砭 +砰 +砲 +破 +砷 +砸 +砺 +砼 +砾 +础 +硅 +硐 +硒 +硕 +硝 +硫 +硬 +确 +硯 +硼 +碁 +碇 +碉 +碌 +碍 +碎 +碑 +碓 +碗 +碘 +碚 +碛 +碟 +碣 +碧 +碩 +碰 +碱 +碳 +碴 +確 +碼 +碾 +磁 +磅 +磊 +磋 +磐 +磕 +磚 +磡 +磨 +磬 +磯 +磲 +磷 +磺 +礁 +礎 +礙 +礡 +礦 +礪 +礫 +礴 +示 +礼 +社 +祀 +祁 +祂 +祇 +祈 +祉 +祎 +祐 +祕 +祖 +祗 +祚 +祛 +祜 +祝 +神 +祟 +祠 +祢 +祥 +票 +祭 +祯 +祷 +祸 +祺 +祿 +禀 +禁 +禄 +禅 +禍 +禎 +福 +禛 +禦 +禧 +禪 +禮 +禱 +禹 +禺 +离 +禽 +禾 +禿 +秀 +私 +秃 +秆 +秉 +秋 +种 +科 +秒 +秘 +租 +秣 +秤 +秦 +秧 +秩 +秭 +积 +称 +秸 +移 +秽 +稀 +稅 +程 +稍 +税 +稔 +稗 +稚 +稜 +稞 +稟 +稠 +稣 +種 +稱 +稲 +稳 +稷 +稹 +稻 +稼 +稽 +稿 +穀 +穂 +穆 +穌 +積 +穎 +穗 +穢 +穩 +穫 +穴 +究 +穷 +穹 +空 +穿 +突 +窃 +窄 +窈 +窍 +窑 +窒 +窓 +窕 +窖 +窗 +窘 +窜 +窝 +窟 +窠 +窥 +窦 +窨 +窩 +窪 +窮 +窯 +窺 +窿 +竄 +竅 +竇 +竊 +立 +竖 +站 +竜 +竞 +竟 +章 +竣 +童 +竭 +端 +競 +竹 +竺 +竽 +竿 +笃 +笆 +笈 +笋 +笏 +笑 +笔 +笙 +笛 +笞 +笠 +符 +笨 +第 +笹 +笺 +笼 +筆 +等 +筊 +筋 +筍 +筏 +筐 +筑 +筒 +答 +策 +筛 +筝 +筠 +筱 +筲 +筵 +筷 +筹 +签 +简 +箇 +箋 +箍 +箏 +箐 +箔 +箕 +算 +箝 +管 +箩 +箫 +箭 +箱 +箴 +箸 +節 +篁 +範 +篆 +篇 +築 +篑 +篓 +篙 +篝 +篠 +篡 +篤 +篩 +篪 +篮 +篱 +篷 +簇 +簌 +簍 +簡 +簦 +簧 +簪 +簫 +簷 +簸 +簽 +簾 +簿 +籁 +籃 +籌 +籍 +籐 +籟 +籠 +籤 +籬 +籮 +籲 +米 +类 +籼 +籽 +粄 +粉 +粑 +粒 +粕 +粗 +粘 +粟 +粤 +粥 +粧 +粪 +粮 +粱 +粲 +粳 +粵 +粹 +粼 +粽 +精 +粿 +糅 +糊 +糍 +糕 +糖 +糗 +糙 +糜 +糞 +糟 +糠 +糧 +糬 +糯 +糰 +糸 +系 +糾 +紀 +紂 +約 +紅 +紉 +紊 +紋 +納 +紐 +紓 +純 +紗 +紘 +紙 +級 +紛 +紜 +素 +紡 +索 +紧 +紫 +紮 +累 +細 +紳 +紹 +紺 +終 +絃 +組 +絆 +経 +結 +絕 +絞 +絡 +絢 +給 +絨 +絮 +統 +絲 +絳 +絵 +絶 +絹 +綁 +綏 +綑 +經 +継 +続 +綜 +綠 +綢 +綦 +綫 +綬 +維 +綱 +網 +綴 +綵 +綸 +綺 +綻 +綽 +綾 +綿 +緊 +緋 +総 +緑 +緒 +緘 +線 +緝 +緞 +締 +緣 +編 +緩 +緬 +緯 +練 +緹 +緻 +縁 +縄 +縈 +縛 +縝 +縣 +縫 +縮 +縱 +縴 +縷 +總 +績 +繁 +繃 +繆 +繇 +繋 +織 +繕 +繚 +繞 +繡 +繩 +繪 +繫 +繭 +繳 +繹 +繼 +繽 +纂 +續 +纍 +纏 +纓 +纔 +纖 +纜 +纠 +红 +纣 +纤 +约 +级 +纨 +纪 +纫 +纬 +纭 +纯 +纰 +纱 +纲 +纳 +纵 +纶 +纷 +纸 +纹 +纺 +纽 +纾 +线 +绀 +练 +组 +绅 +细 +织 +终 +绊 +绍 +绎 +经 +绑 +绒 +结 +绔 +绕 +绘 +给 +绚 +绛 +络 +绝 +绞 +统 +绡 +绢 +绣 +绥 +绦 +继 +绩 +绪 +绫 +续 +绮 +绯 +绰 +绳 +维 +绵 +绶 +绷 +绸 +绻 +综 +绽 +绾 +绿 +缀 +缄 +缅 +缆 +缇 +缈 +缉 +缎 +缓 +缔 +缕 +编 +缘 +缙 +缚 +缜 +缝 +缠 +缢 +缤 +缥 +缨 +缩 +缪 +缭 +缮 +缰 +缱 +缴 +缸 +缺 +缽 +罂 +罄 +罌 +罐 +网 +罔 +罕 +罗 +罚 +罡 +罢 +罩 +罪 +置 +罰 +署 +罵 +罷 +罹 +羁 +羅 +羈 +羊 +羌 +美 +羔 +羚 +羞 +羟 +羡 +羣 +群 +羥 +羧 +羨 +義 +羯 +羲 +羸 +羹 +羽 +羿 +翁 +翅 +翊 +翌 +翎 +習 +翔 +翘 +翟 +翠 +翡 +翦 +翩 +翰 +翱 +翳 +翹 +翻 +翼 +耀 +老 +考 +耄 +者 +耆 +耋 +而 +耍 +耐 +耒 +耕 +耗 +耘 +耙 +耦 +耨 +耳 +耶 +耷 +耸 +耻 +耽 +耿 +聂 +聆 +聊 +聋 +职 +聒 +联 +聖 +聘 +聚 +聞 +聪 +聯 +聰 +聲 +聳 +聴 +聶 +職 +聽 +聾 +聿 +肃 +肄 +肅 +肆 +肇 +肉 +肋 +肌 +肏 +肓 +肖 +肘 +肚 +肛 +肝 +肠 +股 +肢 +肤 +肥 +肩 +肪 +肮 +肯 +肱 +育 +肴 +肺 +肽 +肾 +肿 +胀 +胁 +胃 +胄 +胆 +背 +胍 +胎 +胖 +胚 +胛 +胜 +胝 +胞 +胡 +胤 +胥 +胧 +胫 +胭 +胯 +胰 +胱 +胳 +胴 +胶 +胸 +胺 +能 +脂 +脅 +脆 +脇 +脈 +脉 +脊 +脍 +脏 +脐 +脑 +脓 +脖 +脘 +脚 +脛 +脣 +脩 +脫 +脯 +脱 +脲 +脳 +脸 +脹 +脾 +腆 +腈 +腊 +腋 +腌 +腎 +腐 +腑 +腓 +腔 +腕 +腥 +腦 +腩 +腫 +腭 +腮 +腰 +腱 +腳 +腴 +腸 +腹 +腺 +腻 +腼 +腾 +腿 +膀 +膈 +膊 +膏 +膑 +膘 +膚 +膛 +膜 +膝 +膠 +膦 +膨 +膩 +膳 +膺 +膻 +膽 +膾 +膿 +臀 +臂 +臃 +臆 +臉 +臊 +臍 +臓 +臘 +臟 +臣 +臥 +臧 +臨 +自 +臬 +臭 +至 +致 +臺 +臻 +臼 +臾 +舀 +舂 +舅 +舆 +與 +興 +舉 +舊 +舌 +舍 +舎 +舐 +舒 +舔 +舖 +舗 +舛 +舜 +舞 +舟 +航 +舫 +般 +舰 +舱 +舵 +舶 +舷 +舸 +船 +舺 +舾 +艇 +艋 +艘 +艙 +艦 +艮 +良 +艰 +艱 +色 +艳 +艷 +艹 +艺 +艾 +节 +芃 +芈 +芊 +芋 +芍 +芎 +芒 +芙 +芜 +芝 +芡 +芥 +芦 +芩 +芪 +芫 +芬 +芭 +芮 +芯 +花 +芳 +芷 +芸 +芹 +芻 +芽 +芾 +苁 +苄 +苇 +苋 +苍 +苏 +苑 +苒 +苓 +苔 +苕 +苗 +苛 +苜 +苞 +苟 +苡 +苣 +若 +苦 +苫 +苯 +英 +苷 +苹 +苻 +茁 +茂 +范 +茄 +茅 +茉 +茎 +茏 +茗 +茜 +茧 +茨 +茫 +茬 +茭 +茯 +茱 +茲 +茴 +茵 +茶 +茸 +茹 +茼 +荀 +荃 +荆 +草 +荊 +荏 +荐 +荒 +荔 +荖 +荘 +荚 +荞 +荟 +荠 +荡 +荣 +荤 +荥 +荧 +荨 +荪 +荫 +药 +荳 +荷 +荸 +荻 +荼 +荽 +莅 +莆 +莉 +莊 +莎 +莒 +莓 +莖 +莘 +莞 +莠 +莢 +莧 +莪 +莫 +莱 +莲 +莴 +获 +莹 +莺 +莽 +莿 +菀 +菁 +菅 +菇 +菈 +菊 +菌 +菏 +菓 +菖 +菘 +菜 +菟 +菠 +菡 +菩 +華 +菱 +菲 +菸 +菽 +萁 +萃 +萄 +萊 +萋 +萌 +萍 +萎 +萘 +萝 +萤 +营 +萦 +萧 +萨 +萩 +萬 +萱 +萵 +萸 +萼 +落 +葆 +葉 +著 +葚 +葛 +葡 +董 +葦 +葩 +葫 +葬 +葭 +葯 +葱 +葳 +葵 +葷 +葺 +蒂 +蒋 +蒐 +蒔 +蒙 +蒜 +蒞 +蒟 +蒡 +蒨 +蒲 +蒸 +蒹 +蒻 +蒼 +蒿 +蓁 +蓄 +蓆 +蓉 +蓋 +蓑 +蓓 +蓖 +蓝 +蓟 +蓦 +蓬 +蓮 +蓼 +蓿 +蔑 +蔓 +蔔 +蔗 +蔘 +蔚 +蔡 +蔣 +蔥 +蔫 +蔬 +蔭 +蔵 +蔷 +蔺 +蔻 +蔼 +蔽 +蕁 +蕃 +蕈 +蕉 +蕊 +蕎 +蕙 +蕤 +蕨 +蕩 +蕪 +蕭 +蕲 +蕴 +蕻 +蕾 +薄 +薅 +薇 +薈 +薊 +薏 +薑 +薔 +薙 +薛 +薦 +薨 +薩 +薪 +薬 +薯 +薰 +薹 +藉 +藍 +藏 +藐 +藓 +藕 +藜 +藝 +藤 +藥 +藩 +藹 +藻 +藿 +蘆 +蘇 +蘊 +蘋 +蘑 +蘚 +蘭 +蘸 +蘼 +蘿 +虎 +虏 +虐 +虑 +虔 +處 +虚 +虛 +虜 +虞 +號 +虢 +虧 +虫 +虬 +虱 +虹 +虻 +虽 +虾 +蚀 +蚁 +蚂 +蚊 +蚌 +蚓 +蚕 +蚜 +蚝 +蚣 +蚤 +蚩 +蚪 +蚯 +蚱 +蚵 +蛀 +蛆 +蛇 +蛊 +蛋 +蛎 +蛐 +蛔 +蛙 +蛛 +蛟 +蛤 +蛭 +蛮 +蛰 +蛳 +蛹 +蛻 +蛾 +蜀 +蜂 +蜃 +蜆 +蜇 +蜈 +蜊 +蜍 +蜒 +蜓 +蜕 +蜗 +蜘 +蜚 +蜜 +蜡 +蜢 +蜥 +蜱 +蜴 +蜷 +蜻 +蜿 +蝇 +蝈 +蝉 +蝌 +蝎 +蝕 +蝗 +蝙 +蝟 +蝠 +蝦 +蝨 +蝴 +蝶 +蝸 +蝼 +螂 +螃 +融 +螞 +螢 +螨 +螯 +螳 +螺 +蟀 +蟄 +蟆 +蟋 +蟎 +蟑 +蟒 +蟠 +蟬 +蟲 +蟹 +蟻 +蟾 +蠅 +蠍 +蠔 +蠕 +蠛 +蠟 +蠡 +蠢 +蠣 +蠱 +蠶 +蠹 +蠻 +血 +衄 +衅 +衆 +行 +衍 +術 +衔 +街 +衙 +衛 +衝 +衞 +衡 +衢 +衣 +补 +表 +衩 +衫 +衬 +衮 +衰 +衲 +衷 +衹 +衾 +衿 +袁 +袂 +袄 +袅 +袈 +袋 +袍 +袒 +袖 +袜 +袞 +袤 +袪 +被 +袭 +袱 +裁 +裂 +装 +裆 +裊 +裏 +裔 +裕 +裘 +裙 +補 +裝 +裟 +裡 +裤 +裨 +裱 +裳 +裴 +裸 +裹 +製 +裾 +褂 +複 +褐 +褒 +褓 +褔 +褚 +褥 +褪 +褫 +褲 +褶 +褻 +襁 +襄 +襟 +襠 +襪 +襬 +襯 +襲 +西 +要 +覃 +覆 +覇 +見 +規 +覓 +視 +覚 +覦 +覧 +親 +覬 +観 +覷 +覺 +覽 +觀 +见 +观 +规 +觅 +视 +览 +觉 +觊 +觎 +觐 +觑 +角 +觞 +解 +觥 +触 +觸 +言 +訂 +計 +訊 +討 +訓 +訕 +訖 +託 +記 +訛 +訝 +訟 +訣 +訥 +訪 +設 +許 +訳 +訴 +訶 +診 +註 +証 +詆 +詐 +詔 +評 +詛 +詞 +詠 +詡 +詢 +詣 +試 +詩 +詫 +詬 +詭 +詮 +詰 +話 +該 +詳 +詹 +詼 +誅 +誇 +誉 +誌 +認 +誓 +誕 +誘 +語 +誠 +誡 +誣 +誤 +誥 +誦 +誨 +說 +説 +読 +誰 +課 +誹 +誼 +調 +諄 +談 +請 +諏 +諒 +論 +諗 +諜 +諡 +諦 +諧 +諫 +諭 +諮 +諱 +諳 +諷 +諸 +諺 +諾 +謀 +謁 +謂 +謄 +謊 +謎 +謐 +謔 +謗 +謙 +講 +謝 +謠 +謨 +謬 +謹 +謾 +譁 +證 +譎 +譏 +識 +譙 +譚 +譜 +警 +譬 +譯 +議 +譲 +譴 +護 +譽 +讀 +變 +讓 +讚 +讞 +计 +订 +认 +讥 +讧 +讨 +让 +讪 +讫 +训 +议 +讯 +记 +讲 +讳 +讴 +讶 +讷 +许 +讹 +论 +讼 +讽 +设 +访 +诀 +证 +诃 +评 +诅 +识 +诈 +诉 +诊 +诋 +词 +诏 +译 +试 +诗 +诘 +诙 +诚 +诛 +话 +诞 +诟 +诠 +诡 +询 +诣 +诤 +该 +详 +诧 +诩 +诫 +诬 +语 +误 +诰 +诱 +诲 +说 +诵 +诶 +请 +诸 +诺 +读 +诽 +课 +诿 +谀 +谁 +调 +谄 +谅 +谆 +谈 +谊 +谋 +谌 +谍 +谎 +谏 +谐 +谑 +谒 +谓 +谔 +谕 +谗 +谘 +谙 +谚 +谛 +谜 +谟 +谢 +谣 +谤 +谥 +谦 +谧 +谨 +谩 +谪 +谬 +谭 +谯 +谱 +谲 +谴 +谶 +谷 +豁 +豆 +豇 +豈 +豉 +豊 +豌 +豎 +豐 +豔 +豚 +象 +豢 +豪 +豫 +豬 +豹 +豺 +貂 +貅 +貌 +貓 +貔 +貘 +貝 +貞 +負 +財 +貢 +貧 +貨 +販 +貪 +貫 +責 +貯 +貰 +貳 +貴 +貶 +買 +貸 +費 +貼 +貽 +貿 +賀 +賁 +賂 +賃 +賄 +資 +賈 +賊 +賑 +賓 +賜 +賞 +賠 +賡 +賢 +賣 +賤 +賦 +質 +賬 +賭 +賴 +賺 +購 +賽 +贅 +贈 +贊 +贍 +贏 +贓 +贖 +贛 +贝 +贞 +负 +贡 +财 +责 +贤 +败 +账 +货 +质 +贩 +贪 +贫 +贬 +购 +贮 +贯 +贰 +贱 +贲 +贴 +贵 +贷 +贸 +费 +贺 +贻 +贼 +贾 +贿 +赁 +赂 +赃 +资 +赅 +赈 +赊 +赋 +赌 +赎 +赏 +赐 +赓 +赔 +赖 +赘 +赚 +赛 +赝 +赞 +赠 +赡 +赢 +赣 +赤 +赦 +赧 +赫 +赭 +走 +赳 +赴 +赵 +赶 +起 +趁 +超 +越 +趋 +趕 +趙 +趟 +趣 +趨 +足 +趴 +趵 +趸 +趺 +趾 +跃 +跄 +跆 +跋 +跌 +跎 +跑 +跖 +跚 +跛 +距 +跟 +跡 +跤 +跨 +跩 +跪 +路 +跳 +践 +跷 +跹 +跺 +跻 +踉 +踊 +踌 +踏 +踐 +踝 +踞 +踟 +踢 +踩 +踪 +踮 +踱 +踴 +踵 +踹 +蹂 +蹄 +蹇 +蹈 +蹉 +蹊 +蹋 +蹑 +蹒 +蹙 +蹟 +蹣 +蹤 +蹦 +蹩 +蹬 +蹭 +蹲 +蹴 +蹶 +蹺 +蹼 +蹿 +躁 +躇 +躉 +躊 +躋 +躍 +躏 +躪 +身 +躬 +躯 +躲 +躺 +軀 +車 +軋 +軌 +軍 +軒 +軟 +転 +軸 +軼 +軽 +軾 +較 +載 +輒 +輓 +輔 +輕 +輛 +輝 +輟 +輩 +輪 +輯 +輸 +輻 +輾 +輿 +轄 +轅 +轆 +轉 +轍 +轎 +轟 +车 +轧 +轨 +轩 +转 +轭 +轮 +软 +轰 +轲 +轴 +轶 +轻 +轼 +载 +轿 +较 +辄 +辅 +辆 +辇 +辈 +辉 +辊 +辍 +辐 +辑 +输 +辕 +辖 +辗 +辘 +辙 +辛 +辜 +辞 +辟 +辣 +辦 +辨 +辩 +辫 +辭 +辮 +辯 +辰 +辱 +農 +边 +辺 +辻 +込 +辽 +达 +迁 +迂 +迄 +迅 +过 +迈 +迎 +运 +近 +返 +还 +这 +进 +远 +违 +连 +迟 +迢 +迤 +迥 +迦 +迩 +迪 +迫 +迭 +述 +迴 +迷 +迸 +迹 +迺 +追 +退 +送 +适 +逃 +逅 +逆 +选 +逊 +逍 +透 +逐 +递 +途 +逕 +逗 +這 +通 +逛 +逝 +逞 +速 +造 +逢 +連 +逮 +週 +進 +逵 +逶 +逸 +逻 +逼 +逾 +遁 +遂 +遅 +遇 +遊 +運 +遍 +過 +遏 +遐 +遑 +遒 +道 +達 +違 +遗 +遙 +遛 +遜 +遞 +遠 +遢 +遣 +遥 +遨 +適 +遭 +遮 +遲 +遴 +遵 +遶 +遷 +選 +遺 +遼 +遽 +避 +邀 +邁 +邂 +邃 +還 +邇 +邈 +邊 +邋 +邏 +邑 +邓 +邕 +邛 +邝 +邢 +那 +邦 +邨 +邪 +邬 +邮 +邯 +邰 +邱 +邳 +邵 +邸 +邹 +邺 +邻 +郁 +郅 +郊 +郎 +郑 +郜 +郝 +郡 +郢 +郤 +郦 +郧 +部 +郫 +郭 +郴 +郵 +郷 +郸 +都 +鄂 +鄉 +鄒 +鄔 +鄙 +鄞 +鄢 +鄧 +鄭 +鄰 +鄱 +鄲 +鄺 +酉 +酊 +酋 +酌 +配 +酐 +酒 +酗 +酚 +酝 +酢 +酣 +酥 +酩 +酪 +酬 +酮 +酯 +酰 +酱 +酵 +酶 +酷 +酸 +酿 +醃 +醇 +醉 +醋 +醍 +醐 +醒 +醚 +醛 +醜 +醞 +醣 +醪 +醫 +醬 +醮 +醯 +醴 +醺 +釀 +釁 +采 +釉 +释 +釋 +里 +重 +野 +量 +釐 +金 +釗 +釘 +釜 +針 +釣 +釦 +釧 +釵 +鈀 +鈉 +鈍 +鈎 +鈔 +鈕 +鈞 +鈣 +鈦 +鈪 +鈴 +鈺 +鈾 +鉀 +鉄 +鉅 +鉉 +鉑 +鉗 +鉚 +鉛 +鉤 +鉴 +鉻 +銀 +銃 +銅 +銑 +銓 +銖 +銘 +銜 +銬 +銭 +銮 +銳 +銷 +銹 +鋁 +鋅 +鋒 +鋤 +鋪 +鋰 +鋸 +鋼 +錄 +錐 +錘 +錚 +錠 +錢 +錦 +錨 +錫 +錮 +錯 +録 +錳 +錶 +鍊 +鍋 +鍍 +鍛 +鍥 +鍰 +鍵 +鍺 +鍾 +鎂 +鎊 +鎌 +鎏 +鎔 +鎖 +鎗 +鎚 +鎧 +鎬 +鎮 +鎳 +鏈 +鏖 +鏗 +鏘 +鏞 +鏟 +鏡 +鏢 +鏤 +鏽 +鐘 +鐮 +鐲 +鐳 +鐵 +鐸 +鐺 +鑄 +鑊 +鑑 +鑒 +鑣 +鑫 +鑰 +鑲 +鑼 +鑽 +鑾 +鑿 +针 +钉 +钊 +钎 +钏 +钒 +钓 +钗 +钙 +钛 +钜 +钝 +钞 +钟 +钠 +钡 +钢 +钣 +钤 +钥 +钦 +钧 +钨 +钩 +钮 +钯 +钰 +钱 +钳 +钴 +钵 +钺 +钻 +钼 +钾 +钿 +铀 +铁 +铂 +铃 +铄 +铅 +铆 +铉 +铎 +铐 +铛 +铜 +铝 +铠 +铡 +铢 +铣 +铤 +铨 +铩 +铬 +铭 +铮 +铰 +铲 +铵 +银 +铸 +铺 +链 +铿 +销 +锁 +锂 +锄 +锅 +锆 +锈 +锉 +锋 +锌 +锏 +锐 +锑 +错 +锚 +锟 +锡 +锢 +锣 +锤 +锥 +锦 +锭 +键 +锯 +锰 +锲 +锵 +锹 +锺 +锻 +镀 +镁 +镂 +镇 +镉 +镌 +镍 +镐 +镑 +镕 +镖 +镗 +镛 +镜 +镣 +镭 +镯 +镰 +镳 +镶 +長 +长 +門 +閃 +閉 +開 +閎 +閏 +閑 +閒 +間 +閔 +閘 +閡 +関 +閣 +閥 +閨 +閩 +閱 +閲 +閹 +閻 +閾 +闆 +闇 +闊 +闌 +闍 +闔 +闕 +闖 +闘 +關 +闡 +闢 +门 +闪 +闫 +闭 +问 +闯 +闰 +闲 +间 +闵 +闷 +闸 +闹 +闺 +闻 +闽 +闾 +阀 +阁 +阂 +阅 +阆 +阇 +阈 +阉 +阎 +阐 +阑 +阔 +阕 +阖 +阙 +阚 +阜 +队 +阡 +阪 +阮 +阱 +防 +阳 +阴 +阵 +阶 +阻 +阿 +陀 +陂 +附 +际 +陆 +陇 +陈 +陋 +陌 +降 +限 +陕 +陛 +陝 +陞 +陟 +陡 +院 +陣 +除 +陨 +险 +陪 +陰 +陲 +陳 +陵 +陶 +陷 +陸 +険 +陽 +隅 +隆 +隈 +隊 +隋 +隍 +階 +随 +隐 +隔 +隕 +隘 +隙 +際 +障 +隠 +隣 +隧 +隨 +險 +隱 +隴 +隶 +隸 +隻 +隼 +隽 +难 +雀 +雁 +雄 +雅 +集 +雇 +雉 +雋 +雌 +雍 +雎 +雏 +雑 +雒 +雕 +雖 +雙 +雛 +雜 +雞 +離 +難 +雨 +雪 +雯 +雰 +雲 +雳 +零 +雷 +雹 +電 +雾 +需 +霁 +霄 +霆 +震 +霈 +霉 +霊 +霍 +霎 +霏 +霑 +霓 +霖 +霜 +霞 +霧 +霭 +霰 +露 +霸 +霹 +霽 +霾 +靂 +靄 +靈 +青 +靓 +靖 +静 +靚 +靛 +靜 +非 +靠 +靡 +面 +靥 +靦 +革 +靳 +靴 +靶 +靼 +鞅 +鞋 +鞍 +鞏 +鞑 +鞘 +鞠 +鞣 +鞦 +鞭 +韆 +韋 +韌 +韓 +韜 +韦 +韧 +韩 +韬 +韭 +音 +韵 +韶 +韻 +響 +頁 +頂 +頃 +項 +順 +須 +頌 +預 +頑 +頒 +頓 +頗 +領 +頜 +頡 +頤 +頫 +頭 +頰 +頷 +頸 +頹 +頻 +頼 +顆 +題 +額 +顎 +顏 +顔 +願 +顛 +類 +顧 +顫 +顯 +顱 +顴 +页 +顶 +顷 +项 +顺 +须 +顼 +顽 +顾 +顿 +颁 +颂 +预 +颅 +领 +颇 +颈 +颉 +颊 +颌 +颍 +颐 +频 +颓 +颔 +颖 +颗 +题 +颚 +颛 +颜 +额 +颞 +颠 +颡 +颢 +颤 +颦 +颧 +風 +颯 +颱 +颳 +颶 +颼 +飄 +飆 +风 +飒 +飓 +飕 +飘 +飙 +飚 +飛 +飞 +食 +飢 +飨 +飩 +飪 +飯 +飲 +飼 +飽 +飾 +餃 +餅 +餉 +養 +餌 +餐 +餒 +餓 +餘 +餚 +餛 +餞 +餡 +館 +餮 +餵 +餾 +饅 +饈 +饋 +饌 +饍 +饑 +饒 +饕 +饗 +饞 +饥 +饨 +饪 +饬 +饭 +饮 +饯 +饰 +饱 +饲 +饴 +饵 +饶 +饷 +饺 +饼 +饽 +饿 +馀 +馁 +馄 +馅 +馆 +馈 +馋 +馍 +馏 +馒 +馔 +首 +馗 +香 +馥 +馨 +馬 +馭 +馮 +馳 +馴 +駁 +駄 +駅 +駆 +駐 +駒 +駕 +駛 +駝 +駭 +駱 +駿 +騁 +騎 +騏 +験 +騙 +騨 +騰 +騷 +驀 +驅 +驊 +驍 +驒 +驕 +驗 +驚 +驛 +驟 +驢 +驥 +马 +驭 +驮 +驯 +驰 +驱 +驳 +驴 +驶 +驷 +驸 +驹 +驻 +驼 +驾 +驿 +骁 +骂 +骄 +骅 +骆 +骇 +骈 +骊 +骋 +验 +骏 +骐 +骑 +骗 +骚 +骛 +骜 +骞 +骠 +骡 +骤 +骥 +骧 +骨 +骯 +骰 +骶 +骷 +骸 +骼 +髂 +髅 +髋 +髏 +髒 +髓 +體 +髖 +高 +髦 +髪 +髮 +髯 +髻 +鬃 +鬆 +鬍 +鬓 +鬚 +鬟 +鬢 +鬣 +鬥 +鬧 +鬱 +鬼 +魁 +魂 +魄 +魅 +魇 +魍 +魏 +魔 +魘 +魚 +魯 +魷 +鮑 +鮨 +鮪 +鮭 +鮮 +鯉 +鯊 +鯖 +鯛 +鯨 +鯰 +鯽 +鰍 +鰓 +鰭 +鰲 +鰻 +鰾 +鱈 +鱉 +鱔 +鱗 +鱷 +鱸 +鱼 +鱿 +鲁 +鲈 +鲍 +鲑 +鲛 +鲜 +鲟 +鲢 +鲤 +鲨 +鲫 +鲱 +鲲 +鲶 +鲷 +鲸 +鳃 +鳄 +鳅 +鳌 +鳍 +鳕 +鳖 +鳗 +鳝 +鳞 +鳥 +鳩 +鳳 +鳴 +鳶 +鴉 +鴕 +鴛 +鴦 +鴨 +鴻 +鴿 +鵑 +鵜 +鵝 +鵡 +鵬 +鵰 +鵲 +鶘 +鶩 +鶯 +鶴 +鷗 +鷲 +鷹 +鷺 +鸚 +鸞 +鸟 +鸠 +鸡 +鸢 +鸣 +鸥 +鸦 +鸨 +鸪 +鸭 +鸯 +鸳 +鸵 +鸽 +鸾 +鸿 +鹂 +鹃 +鹄 +鹅 +鹈 +鹉 +鹊 +鹌 +鹏 +鹑 +鹕 +鹘 +鹜 +鹞 +鹤 +鹦 +鹧 +鹫 +鹭 +鹰 +鹳 +鹵 +鹹 +鹼 +鹽 +鹿 +麂 +麋 +麒 +麓 +麗 +麝 +麟 +麥 +麦 +麩 +麴 +麵 +麸 +麺 +麻 +麼 +麽 +麾 +黃 +黄 +黍 +黎 +黏 +黑 +黒 +黔 +默 +黛 +黜 +黝 +點 +黠 +黨 +黯 +黴 +鼋 +鼎 +鼐 +鼓 +鼠 +鼬 +鼹 +鼻 +鼾 +齁 +齊 +齋 +齐 +齒 +齡 +齢 +齣 +齦 +齿 +龄 +龅 +龈 +龊 +龋 +龌 +龍 +龐 +龔 +龕 +龙 +龚 +龛 +龜 +龟 +︰ +︱ +︶ +︿ +﹁ +﹂ +﹍ +﹏ +﹐ +﹑ +﹒ +﹔ +﹕ +﹖ +﹗ +﹙ +﹚ +﹝ +﹞ +﹡ +﹣ +！ +＂ +＃ +＄ +％ +＆ +＇ +（ +） +＊ +＋ +， +－ +． +／ +０ +１ +２ +３ +４ +５ +６ +７ +８ +９ +： +； +＜ +＝ +＞ +？ +＠ +［ +＼ +］ +＾ +＿ +｀ +ａ +ｂ +ｃ +ｄ +ｅ +ｆ +ｇ +ｈ +ｉ +ｊ +ｋ +ｌ +ｍ +ｎ +ｏ +ｐ +ｑ +ｒ +ｓ +ｔ +ｕ +ｖ +ｗ +ｘ +ｙ +ｚ +｛ +｜ +｝ +～ +｡ +｢ +｣ +､ +･ +ｯ +ｰ +ｲ +ｸ +ｼ +ｽ +ﾄ +ﾉ +ﾌ +ﾗ +ﾙ +ﾝ +ﾞ +ﾟ +￣ +￥ +👍 +🔥 +😂 +😎 +... +yam +10 +2017 +12 +11 +2016 +20 +30 +15 +06 +lofter +##s +2015 +by +16 +14 +18 +13 +24 +17 +2014 +21 +##0 +22 +19 +25 +23 +com +100 +00 +05 +2013 +##a +03 +09 +08 +28 +##2 +50 +01 +04 +##1 +27 +02 +2012 +##3 +26 +##e +07 +##8 +##5 +##6 +##4 +##9 +##7 +29 +2011 +40 +##t +2010 +##o +##d +##i +2009 +##n +app +www +the +##m +31 +##c +##l +##y +##r +##g +2008 +60 +http +200 +qq +##p +80 +##f +google +pixnet +90 +cookies +tripadvisor +500 +##er +##k +35 +##h +facebook +2007 +2000 +70 +##b +of +##x +##u +45 +300 +iphone +32 +1000 +2006 +48 +ip +36 +in +38 +3d +##w +##ing +55 +ctrip +##on +##v +33 +##の +to +34 +400 +id +2005 +it +37 +windows +llc +top +99 +42 +39 +000 +led +at +##an +41 +51 +52 +46 +49 +43 +53 +44 +##z +android +58 +and +59 +2004 +56 +vr +##か +5000 +2003 +47 +blogthis +twitter +54 +##le +150 +ok +2018 +57 +75 +cn +no +ios +##in +##mm +##00 +800 +on +te +3000 +65 +2001 +360 +95 +ig +lv +120 +##ng +##を +##us +##に +pc +てす +── +600 +##te +85 +2002 +88 +##ed +html +ncc +wifi +email +64 +blog +is +##10 +##て +mail +online +##al +dvd +##ic +studio +##は +##℃ +##ia +##と +line +vip +72 +##q +98 +##ce +##en +for +##is +##ra +##es +##j +usb +net +cp +1999 +asia +4g +##cm +diy +new +3c +##お +ta +66 +language +vs +apple +tw +86 +web +##ne +ipad +62 +you +##re +101 +68 +##tion +ps +de +bt +pony +atm +##2017 +1998 +67 +##ch +ceo +##or +go +##na +av +pro +cafe +96 +pinterest +97 +63 +pixstyleme3c +##ta +more +said +##2016 +1997 +mp3 +700 +##ll +nba +jun +##20 +92 +tv +1995 +pm +61 +76 +nbsp +250 +##ie +linux +##ma +cd +110 +hd +##17 +78 +##ion +77 +6000 +am +##th +##st +94 +##se +##et +69 +180 +gdp +my +105 +81 +abc +89 +flash +79 +one +93 +1990 +1996 +##ck +gps +##も +##ly +web885 +106 +2020 +91 +##ge +4000 +1500 +xd +boss +isbn +1994 +org +##ry +me +love +##11 +0fork +73 +##12 +3g +##ter +##ar +71 +82 +##la +hotel +130 +1970 +pk +83 +87 +140 +ie +##os +##30 +##el +74 +##50 +seo +cpu +##ml +p2p +84 +may +##る +sun +tue +internet +cc +posted +youtube +##at +##ン +##man +ii +##ル +##15 +abs +nt +pdf +yahoo +ago +1980 +##it +news +mac +104 +##てす +##me +##り +java +1992 +spa +##de +##nt +hk +all +plus +la +1993 +##mb +##16 +##ve +west +##da +160 +air +##い +##ps +から +##to +1989 +logo +htc +php +https +fi +momo +##son +sat +##ke +##80 +ebd +suv +wi +day +apk +##88 +##um +mv +galaxy +wiki +or +brake +##ス +1200 +する +this +1991 +mon +##こ +❤2017 +po +##ない +javascript +life +home +june +##ss +system +900 +##ー +##０ +pp +1988 +world +fb +4k +br +##as +ic +ai +leonardo +safari +##60 +live +free +xx +wed +win7 +kiehl +##co +lg +o2o +##go +us +235 +1949 +mm +しい +vfm +kanye +##90 +##2015 +##id +jr +##ey +123 +rss +##sa +##ro +##am +##no +thu +fri +350 +##sh +##ki +103 +comments +name +##のて +##pe +##ine +max +1987 +8000 +uber +##mi +##ton +wordpress +office +1986 +1985 +##ment +107 +bd +win10 +##ld +##li +gmail +bb +dior +##rs +##ri +##rd +##ます +up +cad +##® +dr +して +read +##21 +をお +##io +##99 +url +1984 +pvc +paypal +show +policy +##40 +##ty +##18 +with +##★ +##01 +txt +102 +##ba +dna +from +post +mini +ar +taiwan +john +##ga +privacy +agoda +##13 +##ny +word +##24 +##22 +##by +##ur +##hz +1982 +##ang +265 +cookie +netscape +108 +##ka +##～ +##ad +house +share +note +ibm +code +hello +nike +sim +survey +##016 +1979 +1950 +wikia +##32 +##017 +5g +cbc +##tor +##kg +1983 +##rt +##14 +campaign +store +2500 +os +##ct +##ts +##° +170 +api +##ns +365 +excel +##な +##ao +##ら +##し +～～ +##nd +university +163 +には +518 +##70 +##ya +##il +##25 +pierre +ipo +0020 +897 +##23 +hotels +##ian +のお +125 +years +6606 +##ers +##26 +high +##day +time +##ay +bug +##line +##く +##す +##be +xp +talk2yam +yamservice +10000 +coco +##dy +sony +##ies +1978 +microsoft +david +people +##ha +1960 +instagram +intel +その +##ot +iso +1981 +##va +115 +##mo +##land +xxx +man +co +ltxsw +##ation +baby +220 +##pa +##ol +1945 +7000 +tag +450 +##ue +msn +##31 +oppo +##ト +##ca +control +##om +st +chrome +##ure +##ん +be +##き +lol +##19 +した +##bo +240 +lady +##100 +##way +##から +4600 +##ko +##do +##un +4s +corporation +168 +##ni +herme +##28 +ｃｐ +978 +##up +##06 +ui +##ds +ppt +admin +three +します +bbc +re +128 +##48 +ca +##015 +##35 +hp +##ee +tpp +##た +##ive +×× +root +##cc +##ました +##ble +##ity +adobe +park +114 +et +oled +city +##ex +##ler +##ap +china +##book +20000 +view +##ice +global +##km +your +hong +##mg +out +##ms +ng +ebay +##29 +menu +ubuntu +##cy +rom +##view +open +ktv +do +server +##lo +if +english +##ね +##５ +##oo +1600 +##02 +step1 +kong +club +135 +july +inc +1976 +mr +hi +##net +touch +##ls +##ii +michael +lcd +##05 +##33 +phone +james +step2 +1300 +ios9 +##box +dc +##２ +##ley +samsung +111 +280 +pokemon +css +##ent +##les +いいえ +##１ +s8 +atom +play +bmw +##said +sa +etf +ctrl +♥yoyo♥ +##55 +2025 +##2014 +##66 +adidas +amazon +1958 +##ber +##ner +visa +##77 +##der +1800 +connectivity +##hi +firefox +109 +118 +hr +so +style +mark +pop +ol +skip +1975 +as +##27 +##ir +##61 +190 +mba +##う +##ai +le +##ver +1900 +cafe2017 +lte +super +113 +129 +##ron +amd +like +##☆ +are +##ster +we +##sk +paul +data +international +##ft +longchamp +ssd +good +##ート +##ti +reply +##my +↓↓↓ +apr +star +##ker +source +136 +js +112 +get +force +photo +##one +126 +##2013 +##ow +link +bbs +1972 +goods +##lin +python +119 +##ip +game +##ics +##ません +blue +##● +520 +##45 +page +itunes +##03 +1955 +260 +1968 +gt +gif +618 +##ff +##47 +group +くたさい +about +bar +ganji +##nce +music +lee +not +1977 +1971 +1973 +##per +an +faq +comment +##って +days +##ock +116 +##bs +1974 +1969 +v1 +player +1956 +xbox +sql +fm +f1 +139 +##ah +210 +##lv +##mp +##000 +melody +1957 +##３ +550 +17life +199 +1966 +xml +market +##au +##71 +999 +##04 +what +gl +##95 +##age +tips +##68 +book +##ting +mysql +can +1959 +230 +##ung +wonderland +watch +10℃ +##ction +9000 +mar +mobile +1946 +1962 +article +##db +part +▲top +party +って +1967 +1964 +1948 +##07 +##ore +##op +この +dj +##78 +##38 +010 +main +225 +1965 +##ong +art +320 +ad +134 +020 +##73 +117 +pm2 +japan +228 +##08 +ts +1963 +##ica +der +sm +##36 +2019 +##wa +ct +##７ +##や +##64 +1937 +homemesh +search +##85 +##れは +##tv +##di +macbook +##９ +##くたさい +service +##♥ +type +った +750 +##ier +##si +##75 +##います +##ok +best +##ット +goris +lock +##った +cf +3m +big +##ut +ftp +carol +##vi +１０ +1961 +happy +sd +##ac +122 +anti +pe +cnn +iii +1920 +138 +##ラ +1940 +esp +jan +tags +##98 +##51 +august +vol +##86 +154 +##™ +##fs +##れ +##sion +design +ac +##ム +press +jordan +ppp +that +key +check +##６ +##tt +##㎡ +1080p +##lt +power +##42 +1952 +##bc +vivi +##ック +he +133 +121 +jpg +##rry +201 +175 +3500 +1947 +nb +##ted +##rn +しています +1954 +usd +##t00 +master +##ンク +001 +model +##58 +al +##09 +1953 +##34 +ram +goo +ても +##ui +127 +1930 +red +##ary +rpg +item +##pm +##41 +270 +##za +project +##2012 +hot +td +blogabstract +##ger +##62 +650 +##44 +gr2 +##します +##ｍ +black +electronic +nfc +year +asus +また +html5 +cindy +##hd +m3 +132 +esc +##od +booking +##53 +fed +tvb +##81 +##ina +mit +165 +##いる +chan +192 +distribution +next +になる +peter +bios +steam +cm +1941 +にも +pk10 +##ix +##65 +##91 +dec +nasa +##ana +icecat +00z +b1 +will +##46 +li +se +##ji +##み +##ard +oct +##ain +jp +##ze +##bi +cio +##56 +smart +h5 +##39 +##port +curve +vpn +##nm +##dia +utc +##あり +12345678910 +##52 +rmvb +chanel +a4 +miss +##and +##im +media +who +##63 +she +girl +5s +124 +vera +##して +class +vivo +king +##フ +##ei +national +ab +1951 +5cm +888 +145 +ipod +ap +1100 +5mm +211 +ms +2756 +##69 +mp4 +msci +##po +##89 +131 +mg +index +380 +##bit +##out +##zz +##97 +##67 +158 +apec +##８ +photoshop +opec +￥799 +ては +##96 +##tes +##ast +2g +○○ +##ール +￥2899 +##ling +##よ +##ory +1938 +##ical +kitty +content +##43 +step3 +##cn +win8 +155 +vc +1400 +iphone7 +robert +##した +tcl +137 +beauty +##87 +en +dollars +##ys +##oc +step +pay +yy +a1 +##2011 +##lly +##ks +##♪ +1939 +188 +download +1944 +sep +exe +ph +います +school +gb +center +pr +street +##board +uv +##37 +##lan +winrar +##que +##ua +##com +1942 +1936 +480 +gpu +##４ +ettoday +fu +tom +##54 +##ren +##via +149 +##72 +b2b +144 +##79 +##tch +rose +arm +mb +##49 +##ial +##nn +nvidia +step4 +mvp +00㎡ +york +156 +##イ +how +cpi +591 +2765 +gov +kg +joe +##xx +mandy +pa +##ser +copyright +fashion +1935 +don +##け +ecu +##ist +##art +erp +wap +have +##lm +talk +##ek +##ning +##if +ch +##ite +video +1943 +cs +san +iot +look +##84 +##2010 +##ku +october +##ux +trump +##hs +##ide +box +141 +first +##ins +april +##ight +##83 +185 +angel +protected +aa +151 +162 +x1 +m2 +##fe +##× +##ho +size +143 +min +ofo +fun +gomaji +ex +hdmi +food +dns +march +chris +kevin +##のか +##lla +##pp +##ec +ag +ems +6s +720p +##rm +##ham +off +##92 +asp +team +fandom +ed +299 +▌♥ +##ell +info +されています +##82 +sina +4066 +161 +##able +##ctor +330 +399 +315 +dll +rights +ltd +idc +jul +3kg +1927 +142 +ma +surface +##76 +##ク +～～～ +304 +mall +eps +146 +green +##59 +map +space +donald +v2 +sodu +##light +1931 +148 +1700 +まて +310 +reserved +htm +##han +##57 +2d +178 +mod +##ise +##tions +152 +ti +##shi +doc +1933 +icp +055 +wang +##ram +shopping +aug +##pi +##well +now +wam +b2 +からお +##hu +236 +1928 +##gb +266 +f2 +##93 +153 +mix +##ef +##uan +bwl +##plus +##res +core +##ess +tea +5℃ +hktvmall +nhk +##ate +list +##ese +301 +feb +4m +inn +ての +nov +159 +12345 +daniel +##ci +pass +##bet +##nk +coffee +202 +ssl +airbnb +##ute +fbi +woshipm +skype +ea +cg +sp +##fc +##www +yes +edge +alt +007 +##94 +fpga +##ght +##gs +iso9001 +さい +##ile +##wood +##uo +image +lin +icon +american +##em +1932 +set +says +##king +##tive +blogger +##74 +なと +256 +147 +##ox +##zy +##red +##ium +##lf +nokia +claire +##リ +##ding +november +lohas +##500 +##tic +##マ +##cs +##ある +##che +##ire +##gy +##ult +db +january +win +##カ +166 +road +ptt +##ま +##つ +198 +##fa +##mer +anna +pchome +はい +udn +ef +420 +##time +##tte +2030 +##ア +g20 +white +かかります +1929 +308 +garden +eleven +di +##おります +chen +309b +777 +172 +young +cosplay +ちてない +4500 +bat +##123 +##tra +##ては +kindle +npc +steve +etc +##ern +##｜ +call +xperia +ces +travel +sk +s7 +##ous +1934 +##int +みいたたけます +183 +edu +file +cho +qr +##car +##our +186 +##ant +##ｄ +eric +1914 +rends +##jo +##する +mastercard +##2000 +kb +##min +290 +##ino +vista +##ris +##ud +jack +2400 +##set +169 +pos +1912 +##her +##ou +taipei +しく +205 +beta +##ませんか +232 +##fi +express +255 +body +##ill +aphojoy +user +december +meiki +##ick +tweet +richard +##av +##ᆫ +iphone6 +##dd +ちてすか +views +##mark +321 +pd +##００ +times +##▲ +level +##ash +10g +point +5l +##ome +208 +koreanmall +##ak +george +q2 +206 +wma +tcp +##200 +スタッフ +full +mlb +##lle +##watch +tm +run +179 +911 +smith +business +##und +1919 +color +##tal +222 +171 +##less +moon +4399 +##rl +update +pcb +shop +499 +157 +little +なし +end +##mhz +van +dsp +easy +660 +##house +##key +history +##ｏ +oh +##001 +##hy +##web +oem +let +was +##2009 +##gg +review +##wan +182 +##°c +203 +uc +title +##val +united +233 +2021 +##ons +doi +trivago +overdope +sbs +##ance +##ち +grand +special +573032185 +imf +216 +wx17house +##so +##ーム +audi +##he +london +william +##rp +##ake +science +beach +cfa +amp +ps4 +880 +##800 +##link +##hp +crm +ferragamo +bell +make +##eng +195 +under +zh +photos +2300 +##style +##ント +via +176 +da +##gi +company +i7 +##ray +thomas +370 +ufo +i5 +##max +plc +ben +back +research +8g +173 +mike +##pc +##ッフ +september +189 +##ace +vps +february +167 +pantos +wp +lisa +1921 +★★ +jquery +night +long +offer +##berg +##news +1911 +##いて +ray +fks +wto +せます +over +164 +340 +##all +##rus +1924 +##888 +##works +blogtitle +loftpermalink +##→ +187 +martin +test +ling +km +##め +15000 +fda +v3 +##ja +##ロ +ｗedding +かある +outlet +family +##ea +をこ +##top +story +##ness +salvatore +##lu +204 +swift +215 +room +している +oracle +##ul +1925 +sam +b2c +week +pi +rock +##のは +##ａ +##けと +##ean +##300 +##gle +cctv +after +chinese +##back +powered +x2 +##tan +1918 +##nes +##イン +canon +only +181 +##zi +##las +say +##oe +184 +##sd +221 +##bot +##world +##zo +sky +made +top100 +just +1926 +pmi +802 +234 +gap +##vr +177 +les +174 +▲topoct +ball +vogue +vi +ing +ofweek +cos +##list +##ort +▲topmay +##なら +##lon +として +last +##tc +##of +##bus +##gen +real +eva +##コ +a3 +nas +##lie +##ria +##coin +##bt +▲topapr +his +212 +cat +nata +vive +health +⋯⋯ +drive +sir +▲topmar +du +cup +##カー +##ook +##よう +##sy +alex +msg +tour +しました +3ce +##word +193 +ebooks +r8 +block +318 +##より +2200 +nice +pvp +207 +months +1905 +rewards +##ther +1917 +0800 +##xi +##チ +##sc +micro +850 +gg +blogfp +op +1922 +daily +m1 +264 +true +##bb +ml +##tar +##のお +##ky +anthony +196 +253 +##yo +state +218 +##ara +##aa +##rc +##tz +##ston +より +gear +##eo +##ade +ge +see +1923 +##win +##ura +ss +heart +##den +##ita +down +##sm +el +png +2100 +610 +rakuten +whatsapp +bay +dream +add +##use +680 +311 +pad +gucci +mpv +##ode +##fo +island +▲topjun +##▼ +223 +jason +214 +chicago +##❤ +しの +##hone +io +##れる +##ことか +sogo +be2 +##ology +990 +cloud +vcd +##con +2～3 +##ford +##joy +##kb +##こさいます +##rade +but +##ach +docker +##ful +rfid +ul +##ase +hit +ford +##star +580 +##○ +１１ +a2 +sdk +reading +edited +##are +cmos +##mc +238 +siri +light +##ella +##ため +bloomberg +##read +pizza +##ison +jimmy +##vm +college +node +journal +ba +18k +##play +245 +##cer +２０ +magic +##yu +191 +jump +288 +tt +##ings +asr +##lia +3200 +step5 +network +##cd +mc +いします +1234 +pixstyleme +273 +##600 +2800 +money +★★★★★ +1280 +１２ +430 +bl +みの +act +##tus +tokyo +##rial +##life +emba +##ae +saas +tcs +##rk +##wang +summer +##sp +ko +##ving +390 +premium +##その +netflix +##ヒ +uk +mt +##lton +right +frank +two +209 +える +##ple +##cal +021 +##んな +##sen +##ville +hold +nexus +dd +##ius +てお +##mah +##なく +tila +zero +820 +ce +##tin +resort +##ws +charles +old +p10 +5d +report +##360 +##ru +##には +bus +vans +lt +##est +pv +##レ +links +rebecca +##ツ +##dm +azure +##365 +きな +limited +bit +4gb +##mon +1910 +moto +##eam +213 +1913 +var +eos +なとの +226 +blogspot +された +699 +e3 +dos +dm +fc +##ments +##ik +##kw +boy +##bin +##ata +960 +er +##せ +219 +##vin +##tu +##ula +194 +##∥ +station +##ろ +##ature +835 +files +zara +hdr +top10 +nature +950 +magazine +s6 +marriott +##シ +avira +case +##っと +tab +##ran +tony +##home +oculus +im +##ral +jean +saint +cry +307 +rosie +##force +##ini +ice +##bert +のある +##nder +##mber +pet +2600 +##◆ +plurk +▲topdec +##sis +00kg +▲topnov +720 +##ence +tim +##ω +##nc +##ても +##name +log +ips +great +ikea +malaysia +unix +##イト +3600 +##ncy +##nie +12000 +akb48 +##ye +##oid +404 +##chi +##いた +oa +xuehai +##1000 +##orm +##rf +275 +さん +##ware +##リー +980 +ho +##pro +text +##era +560 +bob +227 +##ub +##2008 +8891 +scp +avi +##zen +2022 +mi +wu +museum +qvod +apache +lake +jcb +▲topaug +★★★ +ni +##hr +hill +302 +ne +weibo +490 +ruby +##ーシ +##ヶ +##row +4d +▲topjul +iv +##ish +github +306 +mate +312 +##スト +##lot +##ane +andrew +のハイト +##tina +t1 +rf +ed2k +##vel +##900 +way +final +りの +ns +5a +705 +197 +##メ +sweet +bytes +##ene +▲topjan +231 +##cker +##2007 +##px +100g +topapp +229 +helpapp +rs +low +14k +g4g +care +630 +ldquo +あり +##fork +leave +rm +edition +##gan +##zon +##qq +▲topsep +##google +##ism +gold +224 +explorer +##zer +toyota +category +select +visual +##labels +restaurant +##md +posts +s1 +##ico +もっと +angelababy +123456 +217 +sports +s3 +mbc +1915 +してくたさい +shell +x86 +candy +##new +kbs +face +xl +470 +##here +4a +swissinfo +v8 +▲topfeb +dram +##ual +##vice +3a +##wer +sport +q1 +ios10 +public +int +card +##ｃ +ep +au +rt +##れた +1080 +bill +##mll +kim +３０ +460 +wan +##uk +##ミ +x3 +298 +0t +scott +##ming +239 +e5 +##3d +h7n9 +worldcat +brown +##あります +##vo +##led +##580 +##ax +249 +410 +##ert +paris +##～6 +polo +925 +##lr +599 +##ナ +capital +##hing +bank +cv +1g +##chat +##ｓ +##たい +adc +##ule +2m +##ｅ +digital +hotmail +268 +##pad +870 +bbq +quot +##ring +before +wali +##まて +mcu +2k +2b +という +costco +316 +north +333 +switch +##city +##ｐ +philips +##mann +management +panasonic +##cl +##vd +##ping +##rge +alice +##lk +##ましょう +css3 +##ney +vision +alpha +##ular +##400 +##tter +lz +にお +##ありません +mode +gre +1916 +pci +##tm +237 +1～2 +##yan +##そ +について +##let +##キ +work +war +coach +ah +mary +##ᅵ +huang +##pt +a8 +pt +follow +##berry +1895 +##ew +a5 +ghost +##ション +##wn +##og +south +##code +girls +##rid +action +villa +git +r11 +table +games +##cket +error +##anonymoussaid +##ag +here +##ame +##gc +qa +##■ +##lis +gmp +##gin +vmalife +##cher +yu +wedding +##tis +demo +dragon +530 +soho +social +bye +##rant +river +orz +acer +325 +##↑ +##ース +##ats +261 +del +##ven +440 +ups +##ように +##ター +305 +value +macd +yougou +##dn +661 +##ano +ll +##urt +##rent +continue +script +##wen +##ect +paper +263 +319 +shift +##chel +##フト +##cat +258 +x5 +fox +243 +##さん +car +aaa +##blog +loading +##yn +##tp +kuso +799 +si +sns +イカせるテンマ +ヒンクテンマ3 +rmb +vdc +forest +central +prime +help +ultra +##rmb +##ような +241 +square +688 +##しい +のないフロクに +##field +##reen +##ors +##ju +c1 +start +510 +##air +##map +cdn +##wo +cba +stephen +m8 +100km +##get +opera +##base +##ood +vsa +com™ +##aw +##ail +251 +なのて +count +t2 +##ᅡ +##een +2700 +hop +##gp +vsc +tree +##eg +##ose +816 +285 +##ories +##shop +alphago +v4 +1909 +simon +##ᆼ +fluke62max +zip +スホンサー +##sta +louis +cr +bas +##～10 +bc +##yer +hadoop +##ube +##wi +1906 +0755 +hola +##low +place +centre +5v +d3 +##fer +252 +##750 +##media +281 +540 +0l +exchange +262 +series +##ハー +##san +eb +##bank +##ｋ +q3 +##nge +##mail +take +##lp +259 +1888 +client +east +cache +event +vincent +##ールを +きを +##nse +sui +855 +adchoice +##и +##stry +##なたの +246 +##zone +ga +apps +sea +##ab +248 +cisco +##タ +##rner +kymco +##care +dha +##pu +##yi +minkoff +royal +p1 +への +annie +269 +collection +kpi +playstation +257 +になります +866 +bh +##bar +queen +505 +radio +1904 +andy +armani +##xy +manager +iherb +##ery +##share +spring +raid +johnson +1908 +##ob +volvo +hall +##ball +v6 +our +taylor +##hk +bi +242 +##cp +kate +bo +water +technology +##rie +サイトは +277 +##ona +##sl +hpv +303 +gtx +hip +rdquo +jayz +stone +##lex +##rum +namespace +##やり +620 +##ale +##atic +des +##erson +##ql +##ves +##type +enter +##この +##てきます +d2 +##168 +##mix +##bian +との +a9 +jj +ky +##lc +access +movie +##hc +リストに +tower +##ration +##mit +ます +##nch +ua +tel +prefix +##o2 +1907 +##point +1901 +ott +～10 +##http +##ury +baidu +##ink +member +##logy +bigbang +nownews +##js +##shot +##tb +##こと +247 +eba +##tics +##lus +ける +v5 +spark +##ama +there +##ions +god +##lls +##down +hiv +##ress +burberry +day2 +##kv +◆◆ +jeff +related +film +edit +joseph +283 +##ark +cx +32gb +order +g9 +30000 +##ans +##tty +s5 +##bee +かあります +thread +xr +buy +sh +005 +land +spotify +mx +##ari +276 +##verse +×email +sf +why +##ことて +244 +7headlines +nego +sunny +dom +exo +401 +666 +positioning +fit +rgb +##tton +278 +kiss +alexa +adam +lp +みリストを +##ｇ +mp +##ties +##llow +amy +##du +np +002 +institute +271 +##rth +##lar +2345 +590 +##des +sidebar +１５ +imax +site +##cky +##kit +##ime +##009 +season +323 +##fun +##ンター +##ひ +gogoro +a7 +pu +lily +fire +twd600 +##ッセーシを +いて +##vis +30ml +##cture +##をお +information +##オ +close +friday +##くれる +yi +nick +てすか +##tta +##tel +6500 +##lock +cbd +economy +254 +かお +267 +tinker +double +375 +8gb +voice +##app +oops +channel +today +985 +##right +raw +xyz +##＋ +jim +edm +##cent +7500 +supreme +814 +ds +##its +##asia +dropbox +##てすか +##tti +books +272 +100ml +##tle +##ller +##ken +##more +##boy +sex +309 +##dom +t3 +##ider +##なります +##unch +1903 +810 +feel +5500 +##かった +##put +により +s2 +mo +##gh +men +ka +amoled +div +##tr +##n1 +port +howard +##tags +ken +dnf +##nus +adsense +##а +ide +##へ +buff +thunder +##town +##ique +has +##body +auto +pin +##erry +tee +てした +295 +number +##the +##013 +object +psp +cool +udnbkk +16gb +##mic +miui +##tro +most +r2 +##alk +##nity +1880 +±0 +##いました +428 +s4 +law +version +##oa +n1 +sgs +docomo +##tf +##ack +henry +fc2 +##ded +##sco +##014 +##rite +286 +0mm +linkedin +##ada +##now +wii +##ndy +ucbug +##◎ +sputniknews +legalminer +##ika +##xp +2gb +##bu +q10 +oo +b6 +come +##rman +cheese +ming +maker +##gm +nikon +##fig +ppi +kelly +##ります +jchere +てきます +ted +md +003 +fgo +tech +##tto +dan +soc +##gl +##len +hair +earth +640 +521 +img +##pper +##a1 +##てきる +##ロク +acca +##ition +##ference +suite +##ig +outlook +##mond +##cation +398 +##pr +279 +101vip +358 +##999 +282 +64gb +3800 +345 +airport +##over +284 +##おり +jones +##ith +lab +##su +##いるのて +co2 +town +piece +##llo +no1 +vmware +24h +##qi +focus +reader +##admin +##ora +tb +false +##log +1898 +know +lan +838 +##ces +f4 +##ume +motel +stop +##oper +na +flickr +netcomponents +##af +##─ +pose +williams +local +##ound +##cg +##site +##iko +いお +274 +5m +gsm +con +##ath +1902 +friends +##hip +cell +317 +##rey +780 +cream +##cks +012 +##dp +facebooktwitterpinterestgoogle +sso +324 +shtml +song +swiss +##mw +##キンク +lumia +xdd +string +tiffany +522 +marc +られた +insee +russell +sc +dell +##ations +ｏｋ +camera +289 +##vs +##flow +##late +classic +287 +##nter +stay +g1 +mtv +512 +##ever +##lab +##nger +qe +sata +ryan +d1 +50ml +cms +##cing +su +292 +3300 +editor +296 +##nap +security +sunday +association +##ens +##700 +##bra +acg +##かり +sofascore +とは +mkv +##ign +jonathan +gary +build +labels +##oto +tesla +moba +qi +gohappy +general +ajax +1024 +##かる +サイト +society +##test +##urs +wps +fedora +##ich +mozilla +328 +##480 +##dr +usa +urn +##lina +##ｒ +grace +##die +##try +##ader +1250 +##なり +elle +570 +##chen +##ᆯ +price +##ten +uhz +##ough +eq +##hen +states +push +session +balance +wow +506 +##cus +##py +when +##ward +##ep +34e +wong +library +prada +##サイト +##cle +running +##ree +313 +ck +date +q4 +##ctive +##ool +##＞ +mk +##ira +##163 +388 +die +secret +rq +dota +buffet +は１ヶ +e6 +##ez +pan +368 +ha +##card +##cha +2a +##さ +alan +day3 +eye +f3 +##end +france +keep +adi +rna +tvbs +##ala +solo +nova +##え +##tail +##ょう +support +##ries +##なる +##ved +base +copy +iis +fps +##ways +hero +hgih +profile +fish +mu +ssh +entertainment +chang +##wd +click +cake +##ond +pre +##tom +kic +pixel +##ov +##fl +product +6a +##pd +dear +##gate +es +yumi +audio +##² +##sky +echo +bin +where +##ture +329 +##ape +find +sap +isis +##なと +nand +##101 +##load +##ream +band +a6 +525 +never +##post +festival +50cm +##we +555 +guide +314 +zenfone +##ike +335 +gd +forum +jessica +strong +alexander +##ould +software +allen +##ious +program +360° +else +lohasthree +##gar +することかてきます +please +##れます +rc +##ggle +##ric +bim +50000 +##own +eclipse +355 +brian +3ds +##side +061 +361 +##other +##ける +##tech +##ator +485 +engine +##ged +##ｔ +plaza +##fit +cia +ngo +westbrook +shi +tbs +50mm +##みませんか +sci +291 +reuters +##ily +contextlink +##hn +af +##cil +bridge +very +##cel +1890 +cambridge +##ize +15g +##aid +##data +790 +frm +##head +award +butler +##sun +meta +##mar +america +ps3 +puma +pmid +##すか +lc +670 +kitchen +##lic +オーフン5 +きなしソフトサーヒス +そして +day1 +future +★★★★ +##text +##page +##rris +pm1 +##ket +fans +##っています +1001 +christian +bot +kids +trackback +##hai +c3 +display +##hl +n2 +1896 +idea +さんも +##sent +airmail +##ug +##men +pwm +けます +028 +##lution +369 +852 +awards +schemas +354 +asics +wikipedia +font +##tional +##vy +c2 +293 +##れている +##dget +##ein +っている +contact +pepper +スキル +339 +##～5 +294 +##uel +##ument +730 +##hang +みてす +q5 +##sue +rain +##ndi +wei +swatch +##cept +わせ +331 +popular +##ste +##tag +p2 +501 +trc +1899 +##west +##live +justin +honda +ping +messenger +##rap +v9 +543 +##とは +unity +appqq +はすへて +025 +leo +##tone +##テ +##ass +uniqlo +##010 +502 +her +jane +memory +moneydj +##tical +human +12306 +していると +##m2 +coc +miacare +##mn +tmt +##core +vim +kk +##may +fan +target +use +too +338 +435 +2050 +867 +737 +fast +##2c +services +##ope +omega +energy +##わ +pinkoi +1a +##なから +##rain +jackson +##ement +##シャンルの +374 +366 +そんな +p9 +rd +##ᆨ +1111 +##tier +##vic +zone +##│ +385 +690 +dl +isofix +cpa +m4 +322 +kimi +めて +davis +##lay +lulu +##uck +050 +weeks +qs +##hop +920 +##ｎ +ae +##ear +～5 +eia +405 +##fly +korea +jpeg +boost +##ship +small +##リア +1860 +eur +297 +425 +valley +##iel +simple +##ude +rn +k2 +##ena +されます +non +patrick +しているから +##ナー +feed +5757 +30g +process +well +qqmei +##thing +they +aws +lu +pink +##ters +##kin +または +board +##vertisement +wine +##ien +unicode +##dge +r1 +359 +##tant +いを +##twitter +##3c +cool1 +される +##れて +##ｌ +isp +##012 +standard +45㎡2 +402 +##150 +matt +##fu +326 +##iner +googlemsn +pixnetfacebookyahoo +##ラン +x7 +886 +##uce +メーカー +sao +##ev +##きました +##file +9678 +403 +xddd +shirt +6l +##rio +##hat +3mm +givenchy +ya +bang +##lio +monday +crystal +ロクイン +##abc +336 +head +890 +ubuntuforumwikilinuxpastechat +##vc +##～20 +##rity +cnc +7866 +ipv6 +null +1897 +##ost +yang +imsean +tiger +##fet +##ンス +352 +##＝ +dji +327 +ji +maria +##come +##んて +foundation +3100 +##beth +##なった +1m +601 +active +##aft +##don +3p +sr +349 +emma +##khz +living +415 +353 +1889 +341 +709 +457 +sas +x6 +##face +pptv +x4 +##mate +han +sophie +##jing +337 +fifa +##mand +other +sale +inwedding +##gn +てきちゃいます +##mmy +##pmlast +bad +nana +nbc +してみてくたさいね +なとはお +##wu +##かあります +##あ +note7 +single +##340 +せからこ +してくたさい♪この +しにはとんとんワークケートを +するとあなたにもっとマッチした +ならワークケートへ +もみつかっちゃうかも +ワークケートの +##bel +window +##dio +##ht +union +age +382 +１４ +##ivity +##ｙ +コメント +domain +neo +##isa +##lter +5k +f5 +steven +##cts +powerpoint +tft +self +g2 +ft +##テル +zol +##act +mwc +381 +343 +もう +nbapop +408 +てある +eds +ace +##room +previous +author +tomtom +il +##ets +hu +financial +☆☆☆ +っています +bp +5t +chi +1gb +##hg +fairmont +cross +008 +gay +h2 +function +##けて +356 +also +1b +625 +##ータ +##raph +1894 +3～5 +##ils +i3 +334 +avenue +##host +による +##bon +##tsu +message +navigation +50g +fintech +h6 +##ことを +8cm +##ject +##vas +##firm +credit +##wf +xxxx +form +##nor +##space +huawei +plan +json +sbl +##dc +machine +921 +392 +wish +##120 +##sol +windows7 +edward +##ために +development +washington +##nsis +lo +818 +##sio +##ym +##bor +planet +##～8 +##wt +ieee +gpa +##めて +camp +ann +gm +##tw +##oka +connect +##rss +##work +##atus +wall +chicken +soul +2mm +##times +fa +##ather +##cord +009 +##eep +hitachi +gui +harry +##pan +e1 +disney +##press +##ーション +wind +386 +frigidaire +##tl +liu +hsu +332 +basic +von +ev +いた +てきる +スホンサーサイト +learning +##ull +expedia +archives +change +##wei +santa +cut +ins +6gb +turbo +brand +cf1 +508 +004 +return +747 +##rip +h1 +##nis +##をこ +128gb +##にお +3t +application +しており +emc +rx +##oon +384 +quick +412 +15058 +wilson +wing +chapter +##bug +beyond +##cms +##dar +##oh +zoom +e2 +trip +sb +##nba +rcep +342 +aspx +ci +080 +gc +gnu +める +##count +advanced +dance +dv +##url +##ging +367 +8591 +am09 +shadow +battle +346 +##ｉ +##cia +##という +emily +##のてす +##tation +host +ff +techorz +sars +##mini +##mporary +##ering +nc +4200 +798 +##next +cma +##mbps +##gas +##ift +##dot +##ィ +455 +##～17 +amana +##りの +426 +##ros +ir +00㎡1 +##eet +##ible +##↓ +710 +ˋ▽ˊ +##aka +dcs +iq +##ｖ +l1 +##lor +maggie +##011 +##iu +588 +##～1 +830 +##gt +1tb +articles +create +##burg +##iki +database +fantasy +##rex +##cam +dlc +dean +##you +hard +path +gaming +victoria +maps +cb +##lee +##itor +overchicstoretvhome +systems +##xt +416 +p3 +sarah +760 +##nan +407 +486 +x9 +install +second +626 +##ann +##ph +##rcle +##nic +860 +##nar +ec +##とう +768 +metro +chocolate +##rian +～4 +##table +##しています +skin +##sn +395 +mountain +##0mm +inparadise +6m +7x24 +ib +4800 +##jia +eeworld +creative +g5 +g3 +357 +parker +ecfa +village +からの +18000 +sylvia +サーヒス +hbl +##ques +##onsored +##x2 +##きます +##v4 +##tein +ie6 +383 +##stack +389 +ver +##ads +##baby +sound +bbe +##110 +##lone +##uid +ads +022 +gundam +351 +thinkpad +006 +scrum +match +##ave +mems +##470 +##oy +##なりました +##talk +glass +lamigo +span +##eme +job +##a5 +jay +wade +kde +498 +##lace +ocean +tvg +##covery +##r3 +##ners +##rea +junior +think +##aine +cover +##ision +##sia +↓↓ +##bow +msi +413 +458 +406 +##love +711 +801 +soft +z2 +##pl +456 +1840 +mobil +mind +##uy +427 +nginx +##oi +めた +##rr +6221 +##mple +##sson +##ーシてす +371 +##nts +91tv +comhd +crv3000 +##uard +1868 +397 +deep +lost +field +gallery +##bia +rate +spf +redis +traction +930 +icloud +011 +なら +fe +jose +372 +##tory +into +sohu +fx +899 +379 +kicstart2 +##hia +すく +##～3 +##sit +ra +２４ +##walk +##xure +500g +##pact +pacific +xa +natural +carlo +##250 +##walker +1850 +##can +cto +gigi +516 +##サー +pen +##hoo +ob +matlab +##ｂ +##yy +13913459 +##iti +mango +##bbs +sense +c5 +oxford +##ニア +walker +jennifer +##ola +course +##bre +701 +##pus +##rder +lucky +075 +##ぁ +ivy +なお +##nia +sotheby +side +##ugh +joy +##orage +##ush +##bat +##dt +364 +r9 +##2d +##gio +511 +country +wear +##lax +##～7 +##moon +393 +seven +study +411 +348 +lonzo +8k +##ェ +evolution +##イフ +##kk +gs +kd +##レス +arduino +344 +b12 +##lux +arpg +##rdon +cook +##x5 +dark +five +##als +##ida +とても +sign +362 +##ちの +something +20mm +##nda +387 +##posted +fresh +tf +1870 +422 +cam +##mine +##skip +##form +##ssion +education +394 +##tee +dyson +stage +##jie +want +##night +epson +pack +あります +##ppy +テリヘル +##█ +wd +##eh +##rence +left +##lvin +golden +mhz +discovery +##trix +##n2 +loft +##uch +##dra +##sse +speed +～1 +1mdb +sorry +welcome +##urn +wave +gaga +##lmer +teddy +##160 +トラックハック +せよ +611 +##f2016 +378 +rp +##sha +rar +##あなたに +##きた +840 +holiday +##ュー +373 +074 +##vg +##nos +##rail +gartner +gi +6p +##dium +kit +488 +b3 +eco +##ろう +20g +sean +##stone +autocad +nu +##np +f16 +write +029 +m5 +##ias +images +atp +##dk +fsm +504 +1350 +ve +52kb +##xxx +##のに +##cake +414 +unit +lim +ru +1v +##ification +published +angela +16g +analytics +ak +##ｑ +##nel +gmt +##icon +again +##₂ +##bby +ios11 +445 +かこさいます +waze +いてす +##ハ +9985 +##ust +##ティー +framework +##007 +iptv +delete +52sykb +cl +wwdc +027 +30cm +##fw +##ての +1389 +##xon +brandt +##ses +##dragon +tc +vetements +anne +monte +modern +official +##へて +##ere +##nne +##oud +もちろん +５０ +etnews +##a2 +##graphy +421 +863 +##ちゃん +444 +##rtex +##てお +l2 +##gma +mount +ccd +たと +archive +morning +tan +ddos +e7 +##ホ +day4 +##ウ +gis +453 +its +495 +factory +bruce +pg +##ito +ってくたさい +guest +cdma +##lling +536 +n3 +しかし +3～4 +mega +eyes +ro +１３ +women +dac +church +##jun +singapore +##facebook +6991 +starbucks +##tos +##stin +##shine +zen +##mu +tina +20℃ +1893 +##たけて +503 +465 +request +##gence +qt +##っ +1886 +347 +363 +q7 +##zzi +diary +##tore +409 +##ead +468 +cst +##osa +canada +agent +va +##jiang +##ちは +##ーク +##lam +sg +##nix +##sday +##よって +g6 +##master +bing +##zl +charlie +１６ +8mm +nb40 +##ーン +thai +##ルフ +ln284ct +##itz +##2f +bonnie +##food +##lent +originals +##stro +##lts +418 +∟∣ +##bscribe +children +ntd +yesstyle +##かも +hmv +##tment +d5 +2cm +arts +sms +##pn +##я +##いい +topios9 +539 +lifestyle +virtual +##ague +xz +##deo +muji +024 +unt +##nnis +##ᅩ +faq1 +1884 +396 +##ette +fly +64㎡ +はしめまして +441 +curry +##pop +のこ +release +##← +##◆◆ +##cast +073 +ありな +500ml +##ews +5c +##stle +ios7 +##ima +787 +dog +lenovo +##r4 +roger +013 +cbs +vornado +100m +417 +##desk +##クok +##ald +1867 +9595 +2900 +##van +oil +##ｘ +some +break +common +##jy +##lines +g7 +twice +419 +ella +nano +belle +にこ +##mes +##self +##note +jb +##ことかてきます +benz +##との +##ova +451 +save +##wing +##ますのて +kai +りは +##hua +##rect +rainer +##unge +448 +##0m +adsl +##かな +guestname +##uma +##kins +##zu +tokichoi +##price +county +##med +##mus +rmk +391 +address +vm +えて +openload +##group +##hin +##iginal +amg +urban +##oz +jobs +emi +##public +beautiful +##sch +album +##dden +##bell +jerry +works +hostel +miller +##drive +##rmin +##１０ +376 +boot +828 +##370 +##fx +##cm～ +1885 +##nome +##ctionary +##oman +##lish +##cr +##hm +433 +##how +432 +francis +xi +c919 +b5 +evernote +##uc +vga +##3000 +coupe +##urg +##cca +##uality +019 +6g +れる +multi +##また +##ett +em +hey +##ani +##tax +##rma +inside +than +740 +leonnhurt +##jin +ict +れた +bird +notes +200mm +くの +##dical +##lli +result +442 +iu +ee +438 +smap +gopro +##last +yin +pure +998 +32g +けた +5kg +##dan +##rame +mama +##oot +bean +marketing +##hur +2l +bella +sync +xuite +##ground +515 +discuz +##getrelax +##ince +##bay +##5s +cj +##イス +gmat +apt +##pass +jing +##rix +c4 +rich +##とても +niusnews +##ello +bag +770 +##eting +##mobile +１８ +culture +015 +##のてすか +377 +1020 +area +##ience +616 +details +gp +universal +silver +dit +はお +private +ddd +u11 +kanshu +##ified +fung +##nny +dx +##520 +tai +475 +023 +##fr +##lean +3s +##pin +429 +##rin +25000 +ly +rick +##bility +usb3 +banner +##baru +##gion +metal +dt +vdf +1871 +karl +qualcomm +bear +1010 +oldid +ian +jo +##tors +population +##ernel +1882 +mmorpg +##mv +##bike +603 +##© +ww +friend +##ager +exhibition +##del +##pods +fpx +structure +##free +##tings +kl +##rley +##copyright +##mma +california +3400 +orange +yoga +4l +canmake +honey +##anda +##コメント +595 +nikkie +##ルハイト +dhl +publishing +##mall +##gnet +20cm +513 +##クセス +##┅ +e88 +970 +##dog +fishbase +##! +##" +### +##$ +##% +##& +##' +##( +##) +##* +##+ +##, +##- +##. +##/ +##: +##; +##< +##= +##> +##? +##@ +##[ +##\ +##] +##^ +##_ +##{ +##| +##} +##~ +##£ +##¤ +##¥ +##§ +##« +##± +##³ +##µ +##· +##¹ +##º +##» +##¼ +##ß +##æ +##÷ +##ø +##đ +##ŋ +##ɔ +##ə +##ɡ +##ʰ +##ˇ +##ˈ +##ˊ +##ˋ +##ˍ +##ː +##˙ +##˚ +##ˢ +##α +##β +##γ +##δ +##ε +##η +##θ +##ι +##κ +##λ +##μ +##ν +##ο +##π +##ρ +##ς +##σ +##τ +##υ +##φ +##χ +##ψ +##б +##в +##г +##д +##е +##ж +##з +##к +##л +##м +##н +##о +##п +##р +##с +##т +##у +##ф +##х +##ц +##ч +##ш +##ы +##ь +##і +##ا +##ب +##ة +##ت +##د +##ر +##س +##ع +##ل +##م +##ن +##ه +##و +##ي +##۩ +##ก +##ง +##น +##ม +##ย +##ร +##อ +##า +##เ +##๑ +##་ +##ღ +##ᄀ +##ᄁ +##ᄂ +##ᄃ +##ᄅ +##ᄆ +##ᄇ +##ᄈ +##ᄉ +##ᄋ +##ᄌ +##ᄎ +##ᄏ +##ᄐ +##ᄑ +##ᄒ +##ᅢ +##ᅣ +##ᅥ +##ᅦ +##ᅧ +##ᅨ +##ᅪ +##ᅬ +##ᅭ +##ᅮ +##ᅯ +##ᅲ +##ᅳ +##ᅴ +##ᆷ +##ᆸ +##ᆺ +##ᆻ +##ᗜ +##ᵃ +##ᵉ +##ᵍ +##ᵏ +##ᵐ +##ᵒ +##ᵘ +##‖ +##„ +##† +##• +##‥ +##‧ +##
 +##‰ +##′ +##″ +##‹ +##› +##※ +##‿ +##⁄ +##ⁱ +##⁺ +##ⁿ +##₁ +##₃ +##₄ +##€ +##№ +##ⅰ +##ⅱ +##ⅲ +##ⅳ +##ⅴ +##↔ +##↗ +##↘ +##⇒ +##∀ +##− +##∕ +##∙ +##√ +##∞ +##∟ +##∠ +##∣ +##∩ +##∮ +##∶ +##∼ +##∽ +##≈ +##≒ +##≡ +##≤ +##≥ +##≦ +##≧ +##≪ +##≫ +##⊙ +##⋅ +##⋈ +##⋯ +##⌒ +##① +##② +##③ +##④ +##⑤ +##⑥ +##⑦ +##⑧ +##⑨ +##⑩ +##⑴ +##⑵ +##⑶ +##⑷ +##⑸ +##⒈ +##⒉ +##⒊ +##⒋ +##ⓒ +##ⓔ +##ⓘ +##━ +##┃ +##┆ +##┊ +##┌ +##└ +##├ +##┣ +##═ +##║ +##╚ +##╞ +##╠ +##╭ +##╮ +##╯ +##╰ +##╱ +##╳ +##▂ +##▃ +##▅ +##▇ +##▉ +##▋ +##▌ +##▍ +##▎ +##□ +##▪ +##▫ +##▬ +##△ +##▶ +##► +##▽ +##◇ +##◕ +##◠ +##◢ +##◤ +##☀ +##☕ +##☞ +##☺ +##☼ +##♀ +##♂ +##♠ +##♡ +##♣ +##♦ +##♫ +##♬ +##✈ +##✔ +##✕ +##✖ +##✦ +##✨ +##✪ +##✰ +##✿ +##❀ +##➜ +##➤ +##⦿ +##、 +##。 +##〃 +##々 +##〇 +##〈 +##〉 +##《 +##》 +##「 +##」 +##『 +##』 +##【 +##】 +##〓 +##〔 +##〕 +##〖 +##〗 +##〜 +##〝 +##〞 +##ぃ +##ぇ +##ぬ +##ふ +##ほ +##む +##ゃ +##ゅ +##ゆ +##ょ +##゜ +##ゝ +##ァ +##ゥ +##エ +##ォ +##ケ +##サ +##セ +##ソ +##ッ +##ニ +##ヌ +##ネ +##ノ +##ヘ +##モ +##ャ +##ヤ +##ュ +##ユ +##ョ +##ヨ +##ワ +##ヲ +##・ +##ヽ +##ㄅ +##ㄆ +##ㄇ +##ㄉ +##ㄋ +##ㄌ +##ㄍ +##ㄎ +##ㄏ +##ㄒ +##ㄚ +##ㄛ +##ㄞ +##ㄟ +##ㄢ +##ㄤ +##ㄥ +##ㄧ +##ㄨ +##ㆍ +##㈦ +##㊣ +##㗎 +##一 +##丁 +##七 +##万 +##丈 +##三 +##上 +##下 +##不 +##与 +##丐 +##丑 +##专 +##且 +##丕 +##世 +##丘 +##丙 +##业 +##丛 +##东 +##丝 +##丞 +##丟 +##両 +##丢 +##两 +##严 +##並 +##丧 +##丨 +##个 +##丫 +##中 +##丰 +##串 +##临 +##丶 +##丸 +##丹 +##为 +##主 +##丼 +##丽 +##举 +##丿 +##乂 +##乃 +##久 +##么 +##义 +##之 +##乌 +##乍 +##乎 +##乏 +##乐 +##乒 +##乓 +##乔 +##乖 +##乗 +##乘 +##乙 +##乜 +##九 +##乞 +##也 +##习 +##乡 +##书 +##乩 +##买 +##乱 +##乳 +##乾 +##亀 +##亂 +##了 +##予 +##争 +##事 +##二 +##于 +##亏 +##云 +##互 +##五 +##井 +##亘 +##亙 +##亚 +##些 +##亜 +##亞 +##亟 +##亡 +##亢 +##交 +##亥 +##亦 +##产 +##亨 +##亩 +##享 +##京 +##亭 +##亮 +##亲 +##亳 +##亵 +##人 +##亿 +##什 +##仁 +##仃 +##仄 +##仅 +##仆 +##仇 +##今 +##介 +##仍 +##从 +##仏 +##仑 +##仓 +##仔 +##仕 +##他 +##仗 +##付 +##仙 +##仝 +##仞 +##仟 +##代 +##令 +##以 +##仨 +##仪 +##们 +##仮 +##仰 +##仲 +##件 +##价 +##任 +##份 +##仿 +##企 +##伉 +##伊 +##伍 +##伎 +##伏 +##伐 +##休 +##伕 +##众 +##优 +##伙 +##会 +##伝 +##伞 +##伟 +##传 +##伢 +##伤 +##伦 +##伪 +##伫 +##伯 +##估 +##伴 +##伶 +##伸 +##伺 +##似 +##伽 +##佃 +##但 +##佇 +##佈 +##位 +##低 +##住 +##佐 +##佑 +##体 +##佔 +##何 +##佗 +##佘 +##余 +##佚 +##佛 +##作 +##佝 +##佞 +##佟 +##你 +##佢 +##佣 +##佤 +##佥 +##佩 +##佬 +##佯 +##佰 +##佳 +##併 +##佶 +##佻 +##佼 +##使 +##侃 +##侄 +##來 +##侈 +##例 +##侍 +##侏 +##侑 +##侖 +##侗 +##供 +##依 +##侠 +##価 +##侣 +##侥 +##侦 +##侧 +##侨 +##侬 +##侮 +##侯 +##侵 +##侶 +##侷 +##便 +##係 +##促 +##俄 +##俊 +##俎 +##俏 +##俐 +##俑 +##俗 +##俘 +##俚 +##保 +##俞 +##俟 +##俠 +##信 +##俨 +##俩 +##俪 +##俬 +##俭 +##修 +##俯 +##俱 +##俳 +##俸 +##俺 +##俾 +##倆 +##倉 +##個 +##倌 +##倍 +##倏 +##們 +##倒 +##倔 +##倖 +##倘 +##候 +##倚 +##倜 +##借 +##倡 +##値 +##倦 +##倩 +##倪 +##倫 +##倬 +##倭 +##倶 +##债 +##值 +##倾 +##偃 +##假 +##偈 +##偉 +##偌 +##偎 +##偏 +##偕 +##做 +##停 +##健 +##側 +##偵 +##偶 +##偷 +##偻 +##偽 +##偿 +##傀 +##傅 +##傍 +##傑 +##傘 +##備 +##傚 +##傢 +##傣 +##傥 +##储 +##傩 +##催 +##傭 +##傲 +##傳 +##債 +##傷 +##傻 +##傾 +##僅 +##働 +##像 +##僑 +##僕 +##僖 +##僚 +##僥 +##僧 +##僭 +##僮 +##僱 +##僵 +##價 +##僻 +##儀 +##儂 +##億 +##儆 +##儉 +##儋 +##儒 +##儕 +##儘 +##償 +##儡 +##優 +##儲 +##儷 +##儼 +##儿 +##兀 +##允 +##元 +##兄 +##充 +##兆 +##兇 +##先 +##光 +##克 +##兌 +##免 +##児 +##兑 +##兒 +##兔 +##兖 +##党 +##兜 +##兢 +##入 +##內 +##全 +##兩 +##八 +##公 +##六 +##兮 +##兰 +##共 +##兲 +##关 +##兴 +##兵 +##其 +##具 +##典 +##兹 +##养 +##兼 +##兽 +##冀 +##内 +##円 +##冇 +##冈 +##冉 +##冊 +##册 +##再 +##冏 +##冒 +##冕 +##冗 +##写 +##军 +##农 +##冠 +##冢 +##冤 +##冥 +##冨 +##冪 +##冬 +##冯 +##冰 +##冲 +##决 +##况 +##冶 +##冷 +##冻 +##冼 +##冽 +##冾 +##净 +##凄 +##准 +##凇 +##凈 +##凉 +##凋 +##凌 +##凍 +##减 +##凑 +##凛 +##凜 +##凝 +##几 +##凡 +##凤 +##処 +##凪 +##凭 +##凯 +##凰 +##凱 +##凳 +##凶 +##凸 +##凹 +##出 +##击 +##函 +##凿 +##刀 +##刁 +##刃 +##分 +##切 +##刈 +##刊 +##刍 +##刎 +##刑 +##划 +##列 +##刘 +##则 +##刚 +##创 +##初 +##删 +##判 +##別 +##刨 +##利 +##刪 +##别 +##刮 +##到 +##制 +##刷 +##券 +##刹 +##刺 +##刻 +##刽 +##剁 +##剂 +##剃 +##則 +##剉 +##削 +##剋 +##剌 +##前 +##剎 +##剐 +##剑 +##剔 +##剖 +##剛 +##剜 +##剝 +##剣 +##剤 +##剥 +##剧 +##剩 +##剪 +##副 +##割 +##創 +##剷 +##剽 +##剿 +##劃 +##劇 +##劈 +##劉 +##劊 +##劍 +##劏 +##劑 +##力 +##劝 +##办 +##功 +##加 +##务 +##劣 +##动 +##助 +##努 +##劫 +##劭 +##励 +##劲 +##劳 +##労 +##劵 +##効 +##劾 +##势 +##勁 +##勃 +##勇 +##勉 +##勋 +##勐 +##勒 +##動 +##勖 +##勘 +##務 +##勛 +##勝 +##勞 +##募 +##勢 +##勤 +##勧 +##勳 +##勵 +##勸 +##勺 +##勻 +##勾 +##勿 +##匀 +##包 +##匆 +##匈 +##匍 +##匐 +##匕 +##化 +##北 +##匙 +##匝 +##匠 +##匡 +##匣 +##匪 +##匮 +##匯 +##匱 +##匹 +##区 +##医 +##匾 +##匿 +##區 +##十 +##千 +##卅 +##升 +##午 +##卉 +##半 +##卍 +##华 +##协 +##卑 +##卒 +##卓 +##協 +##单 +##卖 +##南 +##単 +##博 +##卜 +##卞 +##卟 +##占 +##卡 +##卢 +##卤 +##卦 +##卧 +##卫 +##卮 +##卯 +##印 +##危 +##即 +##却 +##卵 +##卷 +##卸 +##卻 +##卿 +##厂 +##厄 +##厅 +##历 +##厉 +##压 +##厌 +##厕 +##厘 +##厚 +##厝 +##原 +##厢 +##厥 +##厦 +##厨 +##厩 +##厭 +##厮 +##厲 +##厳 +##去 +##县 +##叁 +##参 +##參 +##又 +##叉 +##及 +##友 +##双 +##反 +##収 +##发 +##叔 +##取 +##受 +##变 +##叙 +##叛 +##叟 +##叠 +##叡 +##叢 +##口 +##古 +##句 +##另 +##叨 +##叩 +##只 +##叫 +##召 +##叭 +##叮 +##可 +##台 +##叱 +##史 +##右 +##叵 +##叶 +##号 +##司 +##叹 +##叻 +##叼 +##叽 +##吁 +##吃 +##各 +##吆 +##合 +##吉 +##吊 +##吋 +##同 +##名 +##后 +##吏 +##吐 +##向 +##吒 +##吓 +##吕 +##吖 +##吗 +##君 +##吝 +##吞 +##吟 +##吠 +##吡 +##否 +##吧 +##吨 +##吩 +##含 +##听 +##吭 +##吮 +##启 +##吱 +##吳 +##吴 +##吵 +##吶 +##吸 +##吹 +##吻 +##吼 +##吽 +##吾 +##呀 +##呂 +##呃 +##呆 +##呈 +##告 +##呋 +##呎 +##呐 +##呓 +##呕 +##呗 +##员 +##呛 +##呜 +##呢 +##呤 +##呦 +##周 +##呱 +##呲 +##味 +##呵 +##呷 +##呸 +##呻 +##呼 +##命 +##咀 +##咁 +##咂 +##咄 +##咆 +##咋 +##和 +##咎 +##咏 +##咐 +##咒 +##咔 +##咕 +##咖 +##咗 +##咘 +##咙 +##咚 +##咛 +##咣 +##咤 +##咦 +##咧 +##咨 +##咩 +##咪 +##咫 +##咬 +##咭 +##咯 +##咱 +##咲 +##咳 +##咸 +##咻 +##咽 +##咿 +##哀 +##品 +##哂 +##哄 +##哆 +##哇 +##哈 +##哉 +##哋 +##哌 +##响 +##哎 +##哏 +##哐 +##哑 +##哒 +##哔 +##哗 +##哟 +##員 +##哥 +##哦 +##哧 +##哨 +##哩 +##哪 +##哭 +##哮 +##哲 +##哺 +##哼 +##哽 +##唁 +##唄 +##唆 +##唇 +##唉 +##唏 +##唐 +##唑 +##唔 +##唠 +##唤 +##唧 +##唬 +##售 +##唯 +##唰 +##唱 +##唳 +##唷 +##唸 +##唾 +##啃 +##啄 +##商 +##啉 +##啊 +##問 +##啓 +##啕 +##啖 +##啜 +##啞 +##啟 +##啡 +##啤 +##啥 +##啦 +##啧 +##啪 +##啫 +##啬 +##啮 +##啰 +##啱 +##啲 +##啵 +##啶 +##啷 +##啸 +##啻 +##啼 +##啾 +##喀 +##喂 +##喃 +##善 +##喆 +##喇 +##喉 +##喊 +##喋 +##喎 +##喏 +##喔 +##喘 +##喙 +##喚 +##喜 +##喝 +##喟 +##喧 +##喪 +##喫 +##喬 +##單 +##喰 +##喱 +##喲 +##喳 +##喵 +##営 +##喷 +##喹 +##喺 +##喻 +##喽 +##嗅 +##嗆 +##嗇 +##嗎 +##嗑 +##嗒 +##嗓 +##嗔 +##嗖 +##嗚 +##嗜 +##嗝 +##嗟 +##嗡 +##嗣 +##嗤 +##嗦 +##嗨 +##嗪 +##嗬 +##嗯 +##嗰 +##嗲 +##嗳 +##嗶 +##嗷 +##嗽 +##嘀 +##嘅 +##嘆 +##嘈 +##嘉 +##嘌 +##嘍 +##嘎 +##嘔 +##嘖 +##嘗 +##嘘 +##嘚 +##嘛 +##嘜 +##嘞 +##嘟 +##嘢 +##嘣 +##嘤 +##嘧 +##嘩 +##嘭 +##嘮 +##嘯 +##嘰 +##嘱 +##嘲 +##嘴 +##嘶 +##嘸 +##嘹 +##嘻 +##嘿 +##噁 +##噌 +##噎 +##噓 +##噔 +##噗 +##噙 +##噜 +##噠 +##噢 +##噤 +##器 +##噩 +##噪 +##噬 +##噱 +##噴 +##噶 +##噸 +##噹 +##噻 +##噼 +##嚀 +##嚇 +##嚎 +##嚏 +##嚐 +##嚓 +##嚕 +##嚟 +##嚣 +##嚥 +##嚨 +##嚮 +##嚴 +##嚷 +##嚼 +##囂 +##囉 +##囊 +##囍 +##囑 +##囔 +##囗 +##囚 +##四 +##囝 +##回 +##囟 +##因 +##囡 +##团 +##団 +##囤 +##囧 +##囪 +##囫 +##园 +##困 +##囱 +##囲 +##図 +##围 +##囹 +##固 +##国 +##图 +##囿 +##圃 +##圄 +##圆 +##圈 +##國 +##圍 +##圏 +##園 +##圓 +##圖 +##團 +##圜 +##土 +##圣 +##圧 +##在 +##圩 +##圭 +##地 +##圳 +##场 +##圻 +##圾 +##址 +##坂 +##均 +##坊 +##坍 +##坎 +##坏 +##坐 +##坑 +##块 +##坚 +##坛 +##坝 +##坞 +##坟 +##坠 +##坡 +##坤 +##坦 +##坨 +##坪 +##坯 +##坳 +##坵 +##坷 +##垂 +##垃 +##垄 +##型 +##垒 +##垚 +##垛 +##垠 +##垢 +##垣 +##垦 +##垩 +##垫 +##垭 +##垮 +##垵 +##埂 +##埃 +##埋 +##城 +##埔 +##埕 +##埗 +##域 +##埠 +##埤 +##埵 +##執 +##埸 +##培 +##基 +##埼 +##堀 +##堂 +##堃 +##堅 +##堆 +##堇 +##堑 +##堕 +##堙 +##堡 +##堤 +##堪 +##堯 +##堰 +##報 +##場 +##堵 +##堺 +##堿 +##塊 +##塌 +##塑 +##塔 +##塗 +##塘 +##塚 +##塞 +##塢 +##塩 +##填 +##塬 +##塭 +##塵 +##塾 +##墀 +##境 +##墅 +##墉 +##墊 +##墒 +##墓 +##増 +##墘 +##墙 +##墜 +##增 +##墟 +##墨 +##墩 +##墮 +##墳 +##墻 +##墾 +##壁 +##壅 +##壆 +##壇 +##壊 +##壑 +##壓 +##壕 +##壘 +##壞 +##壟 +##壢 +##壤 +##壩 +##士 +##壬 +##壮 +##壯 +##声 +##売 +##壳 +##壶 +##壹 +##壺 +##壽 +##处 +##备 +##変 +##复 +##夏 +##夔 +##夕 +##外 +##夙 +##多 +##夜 +##够 +##夠 +##夢 +##夥 +##大 +##天 +##太 +##夫 +##夭 +##央 +##夯 +##失 +##头 +##夷 +##夸 +##夹 +##夺 +##夾 +##奂 +##奄 +##奇 +##奈 +##奉 +##奋 +##奎 +##奏 +##奐 +##契 +##奔 +##奕 +##奖 +##套 +##奘 +##奚 +##奠 +##奢 +##奥 +##奧 +##奪 +##奬 +##奮 +##女 +##奴 +##奶 +##奸 +##她 +##好 +##如 +##妃 +##妄 +##妆 +##妇 +##妈 +##妊 +##妍 +##妒 +##妓 +##妖 +##妘 +##妙 +##妝 +##妞 +##妣 +##妤 +##妥 +##妨 +##妩 +##妪 +##妮 +##妲 +##妳 +##妹 +##妻 +##妾 +##姆 +##姉 +##姊 +##始 +##姍 +##姐 +##姑 +##姒 +##姓 +##委 +##姗 +##姚 +##姜 +##姝 +##姣 +##姥 +##姦 +##姨 +##姪 +##姫 +##姬 +##姹 +##姻 +##姿 +##威 +##娃 +##娄 +##娅 +##娆 +##娇 +##娉 +##娑 +##娓 +##娘 +##娛 +##娜 +##娟 +##娠 +##娣 +##娥 +##娩 +##娱 +##娲 +##娴 +##娶 +##娼 +##婀 +##婁 +##婆 +##婉 +##婊 +##婕 +##婚 +##婢 +##婦 +##婧 +##婪 +##婭 +##婴 +##婵 +##婶 +##婷 +##婺 +##婿 +##媒 +##媚 +##媛 +##媞 +##媧 +##媲 +##媳 +##媽 +##媾 +##嫁 +##嫂 +##嫉 +##嫌 +##嫑 +##嫔 +##嫖 +##嫘 +##嫚 +##嫡 +##嫣 +##嫦 +##嫩 +##嫲 +##嫵 +##嫻 +##嬅 +##嬉 +##嬌 +##嬗 +##嬛 +##嬢 +##嬤 +##嬪 +##嬰 +##嬴 +##嬷 +##嬸 +##嬿 +##孀 +##孃 +##子 +##孑 +##孔 +##孕 +##孖 +##字 +##存 +##孙 +##孚 +##孛 +##孜 +##孝 +##孟 +##孢 +##季 +##孤 +##学 +##孩 +##孪 +##孫 +##孬 +##孰 +##孱 +##孳 +##孵 +##學 +##孺 +##孽 +##孿 +##宁 +##它 +##宅 +##宇 +##守 +##安 +##宋 +##完 +##宏 +##宓 +##宕 +##宗 +##官 +##宙 +##定 +##宛 +##宜 +##宝 +##实 +##実 +##宠 +##审 +##客 +##宣 +##室 +##宥 +##宦 +##宪 +##宫 +##宮 +##宰 +##害 +##宴 +##宵 +##家 +##宸 +##容 +##宽 +##宾 +##宿 +##寂 +##寄 +##寅 +##密 +##寇 +##富 +##寐 +##寒 +##寓 +##寛 +##寝 +##寞 +##察 +##寡 +##寢 +##寥 +##實 +##寧 +##寨 +##審 +##寫 +##寬 +##寮 +##寰 +##寵 +##寶 +##寸 +##对 +##寺 +##寻 +##导 +##対 +##寿 +##封 +##専 +##射 +##将 +##將 +##專 +##尉 +##尊 +##尋 +##對 +##導 +##小 +##少 +##尔 +##尕 +##尖 +##尘 +##尚 +##尝 +##尤 +##尧 +##尬 +##就 +##尴 +##尷 +##尸 +##尹 +##尺 +##尻 +##尼 +##尽 +##尾 +##尿 +##局 +##屁 +##层 +##屄 +##居 +##屆 +##屈 +##屉 +##届 +##屋 +##屌 +##屍 +##屎 +##屏 +##屐 +##屑 +##展 +##屜 +##属 +##屠 +##屡 +##屢 +##層 +##履 +##屬 +##屯 +##山 +##屹 +##屿 +##岀 +##岁 +##岂 +##岌 +##岐 +##岑 +##岔 +##岖 +##岗 +##岘 +##岙 +##岚 +##岛 +##岡 +##岩 +##岫 +##岬 +##岭 +##岱 +##岳 +##岷 +##岸 +##峇 +##峋 +##峒 +##峙 +##峡 +##峤 +##峥 +##峦 +##峨 +##峪 +##峭 +##峯 +##峰 +##峴 +##島 +##峻 +##峽 +##崁 +##崂 +##崆 +##崇 +##崎 +##崑 +##崔 +##崖 +##崗 +##崙 +##崛 +##崧 +##崩 +##崭 +##崴 +##崽 +##嵇 +##嵊 +##嵋 +##嵌 +##嵐 +##嵘 +##嵩 +##嵬 +##嵯 +##嶂 +##嶄 +##嶇 +##嶋 +##嶙 +##嶺 +##嶼 +##嶽 +##巅 +##巍 +##巒 +##巔 +##巖 +##川 +##州 +##巡 +##巢 +##工 +##左 +##巧 +##巨 +##巩 +##巫 +##差 +##己 +##已 +##巳 +##巴 +##巷 +##巻 +##巽 +##巾 +##巿 +##币 +##市 +##布 +##帅 +##帆 +##师 +##希 +##帐 +##帑 +##帕 +##帖 +##帘 +##帚 +##帛 +##帜 +##帝 +##帥 +##带 +##帧 +##師 +##席 +##帮 +##帯 +##帰 +##帳 +##帶 +##帷 +##常 +##帼 +##帽 +##幀 +##幂 +##幄 +##幅 +##幌 +##幔 +##幕 +##幟 +##幡 +##幢 +##幣 +##幫 +##干 +##平 +##年 +##并 +##幸 +##幹 +##幺 +##幻 +##幼 +##幽 +##幾 +##广 +##庁 +##広 +##庄 +##庆 +##庇 +##床 +##序 +##庐 +##库 +##应 +##底 +##庖 +##店 +##庙 +##庚 +##府 +##庞 +##废 +##庠 +##度 +##座 +##庫 +##庭 +##庵 +##庶 +##康 +##庸 +##庹 +##庾 +##廁 +##廂 +##廃 +##廈 +##廉 +##廊 +##廓 +##廖 +##廚 +##廝 +##廟 +##廠 +##廢 +##廣 +##廬 +##廳 +##延 +##廷 +##建 +##廿 +##开 +##弁 +##异 +##弃 +##弄 +##弈 +##弊 +##弋 +##式 +##弑 +##弒 +##弓 +##弔 +##引 +##弗 +##弘 +##弛 +##弟 +##张 +##弥 +##弦 +##弧 +##弩 +##弭 +##弯 +##弱 +##張 +##強 +##弹 +##强 +##弼 +##弾 +##彅 +##彆 +##彈 +##彌 +##彎 +##归 +##当 +##录 +##彗 +##彙 +##彝 +##形 +##彤 +##彥 +##彦 +##彧 +##彩 +##彪 +##彫 +##彬 +##彭 +##彰 +##影 +##彷 +##役 +##彻 +##彼 +##彿 +##往 +##征 +##径 +##待 +##徇 +##很 +##徉 +##徊 +##律 +##後 +##徐 +##徑 +##徒 +##従 +##徕 +##得 +##徘 +##徙 +##徜 +##從 +##徠 +##御 +##徨 +##復 +##循 +##徬 +##微 +##徳 +##徴 +##徵 +##德 +##徹 +##徼 +##徽 +##心 +##必 +##忆 +##忌 +##忍 +##忏 +##忐 +##忑 +##忒 +##忖 +##志 +##忘 +##忙 +##応 +##忠 +##忡 +##忤 +##忧 +##忪 +##快 +##忱 +##念 +##忻 +##忽 +##忿 +##怀 +##态 +##怂 +##怅 +##怆 +##怎 +##怏 +##怒 +##怔 +##怕 +##怖 +##怙 +##怜 +##思 +##怠 +##怡 +##急 +##怦 +##性 +##怨 +##怪 +##怯 +##怵 +##总 +##怼 +##恁 +##恃 +##恆 +##恋 +##恍 +##恐 +##恒 +##恕 +##恙 +##恚 +##恢 +##恣 +##恤 +##恥 +##恨 +##恩 +##恪 +##恫 +##恬 +##恭 +##息 +##恰 +##恳 +##恵 +##恶 +##恸 +##恺 +##恻 +##恼 +##恿 +##悄 +##悅 +##悉 +##悌 +##悍 +##悔 +##悖 +##悚 +##悟 +##悠 +##患 +##悦 +##您 +##悩 +##悪 +##悬 +##悯 +##悱 +##悲 +##悴 +##悵 +##悶 +##悸 +##悻 +##悼 +##悽 +##情 +##惆 +##惇 +##惊 +##惋 +##惑 +##惕 +##惘 +##惚 +##惜 +##惟 +##惠 +##惡 +##惦 +##惧 +##惨 +##惩 +##惫 +##惬 +##惭 +##惮 +##惯 +##惰 +##惱 +##想 +##惴 +##惶 +##惹 +##惺 +##愁 +##愆 +##愈 +##愉 +##愍 +##意 +##愕 +##愚 +##愛 +##愜 +##感 +##愣 +##愤 +##愧 +##愫 +##愷 +##愿 +##慄 +##慈 +##態 +##慌 +##慎 +##慑 +##慕 +##慘 +##慚 +##慟 +##慢 +##慣 +##慧 +##慨 +##慫 +##慮 +##慰 +##慳 +##慵 +##慶 +##慷 +##慾 +##憂 +##憊 +##憋 +##憎 +##憐 +##憑 +##憔 +##憚 +##憤 +##憧 +##憨 +##憩 +##憫 +##憬 +##憲 +##憶 +##憾 +##懂 +##懇 +##懈 +##應 +##懊 +##懋 +##懑 +##懒 +##懦 +##懲 +##懵 +##懶 +##懷 +##懸 +##懺 +##懼 +##懾 +##懿 +##戀 +##戈 +##戊 +##戌 +##戍 +##戎 +##戏 +##成 +##我 +##戒 +##戕 +##或 +##战 +##戚 +##戛 +##戟 +##戡 +##戦 +##截 +##戬 +##戮 +##戰 +##戲 +##戳 +##戴 +##戶 +##户 +##戸 +##戻 +##戾 +##房 +##所 +##扁 +##扇 +##扈 +##扉 +##手 +##才 +##扎 +##扑 +##扒 +##打 +##扔 +##払 +##托 +##扛 +##扣 +##扦 +##执 +##扩 +##扪 +##扫 +##扬 +##扭 +##扮 +##扯 +##扰 +##扱 +##扳 +##扶 +##批 +##扼 +##找 +##承 +##技 +##抄 +##抉 +##把 +##抑 +##抒 +##抓 +##投 +##抖 +##抗 +##折 +##抚 +##抛 +##抜 +##択 +##抟 +##抠 +##抡 +##抢 +##护 +##报 +##抨 +##披 +##抬 +##抱 +##抵 +##抹 +##押 +##抽 +##抿 +##拂 +##拄 +##担 +##拆 +##拇 +##拈 +##拉 +##拋 +##拌 +##拍 +##拎 +##拐 +##拒 +##拓 +##拔 +##拖 +##拗 +##拘 +##拙 +##拚 +##招 +##拜 +##拟 +##拡 +##拢 +##拣 +##拥 +##拦 +##拧 +##拨 +##择 +##括 +##拭 +##拮 +##拯 +##拱 +##拳 +##拴 +##拷 +##拼 +##拽 +##拾 +##拿 +##持 +##挂 +##指 +##挈 +##按 +##挎 +##挑 +##挖 +##挙 +##挚 +##挛 +##挝 +##挞 +##挟 +##挠 +##挡 +##挣 +##挤 +##挥 +##挨 +##挪 +##挫 +##振 +##挲 +##挹 +##挺 +##挽 +##挾 +##捂 +##捅 +##捆 +##捉 +##捋 +##捌 +##捍 +##捎 +##捏 +##捐 +##捕 +##捞 +##损 +##捡 +##换 +##捣 +##捧 +##捨 +##捩 +##据 +##捱 +##捲 +##捶 +##捷 +##捺 +##捻 +##掀 +##掂 +##掃 +##掇 +##授 +##掉 +##掌 +##掏 +##掐 +##排 +##掖 +##掘 +##掙 +##掛 +##掠 +##採 +##探 +##掣 +##接 +##控 +##推 +##掩 +##措 +##掬 +##掰 +##掲 +##掳 +##掴 +##掷 +##掸 +##掺 +##揀 +##揃 +##揄 +##揆 +##揉 +##揍 +##描 +##提 +##插 +##揖 +##揚 +##換 +##握 +##揣 +##揩 +##揪 +##揭 +##揮 +##援 +##揶 +##揸 +##揹 +##揽 +##搀 +##搁 +##搂 +##搅 +##損 +##搏 +##搐 +##搓 +##搔 +##搖 +##搗 +##搜 +##搞 +##搡 +##搪 +##搬 +##搭 +##搵 +##搶 +##携 +##搽 +##摀 +##摁 +##摄 +##摆 +##摇 +##摈 +##摊 +##摒 +##摔 +##摘 +##摞 +##摟 +##摧 +##摩 +##摯 +##摳 +##摸 +##摹 +##摺 +##摻 +##撂 +##撃 +##撅 +##撇 +##撈 +##撐 +##撑 +##撒 +##撓 +##撕 +##撚 +##撞 +##撤 +##撥 +##撩 +##撫 +##撬 +##播 +##撮 +##撰 +##撲 +##撵 +##撷 +##撸 +##撻 +##撼 +##撿 +##擀 +##擁 +##擂 +##擄 +##擅 +##擇 +##擊 +##擋 +##操 +##擎 +##擒 +##擔 +##擘 +##據 +##擞 +##擠 +##擡 +##擢 +##擦 +##擬 +##擰 +##擱 +##擲 +##擴 +##擷 +##擺 +##擼 +##擾 +##攀 +##攏 +##攒 +##攔 +##攘 +##攙 +##攜 +##攝 +##攞 +##攢 +##攣 +##攤 +##攥 +##攪 +##攫 +##攬 +##支 +##收 +##攸 +##改 +##攻 +##放 +##政 +##故 +##效 +##敌 +##敍 +##敎 +##敏 +##救 +##敕 +##敖 +##敗 +##敘 +##教 +##敛 +##敝 +##敞 +##敢 +##散 +##敦 +##敬 +##数 +##敲 +##整 +##敵 +##敷 +##數 +##斂 +##斃 +##文 +##斋 +##斌 +##斎 +##斐 +##斑 +##斓 +##斗 +##料 +##斛 +##斜 +##斟 +##斡 +##斤 +##斥 +##斧 +##斩 +##斫 +##斬 +##断 +##斯 +##新 +##斷 +##方 +##於 +##施 +##旁 +##旃 +##旅 +##旋 +##旌 +##旎 +##族 +##旖 +##旗 +##无 +##既 +##日 +##旦 +##旧 +##旨 +##早 +##旬 +##旭 +##旮 +##旱 +##时 +##旷 +##旺 +##旻 +##昀 +##昂 +##昆 +##昇 +##昉 +##昊 +##昌 +##明 +##昏 +##易 +##昔 +##昕 +##昙 +##星 +##映 +##春 +##昧 +##昨 +##昭 +##是 +##昱 +##昴 +##昵 +##昶 +##昼 +##显 +##晁 +##時 +##晃 +##晉 +##晋 +##晌 +##晏 +##晒 +##晓 +##晔 +##晕 +##晖 +##晗 +##晚 +##晝 +##晞 +##晟 +##晤 +##晦 +##晨 +##晩 +##普 +##景 +##晰 +##晴 +##晶 +##晷 +##智 +##晾 +##暂 +##暄 +##暇 +##暈 +##暉 +##暌 +##暐 +##暑 +##暖 +##暗 +##暝 +##暢 +##暧 +##暨 +##暫 +##暮 +##暱 +##暴 +##暸 +##暹 +##曄 +##曆 +##曇 +##曉 +##曖 +##曙 +##曜 +##曝 +##曠 +##曦 +##曬 +##曰 +##曲 +##曳 +##更 +##書 +##曹 +##曼 +##曾 +##替 +##最 +##會 +##月 +##有 +##朋 +##服 +##朐 +##朔 +##朕 +##朗 +##望 +##朝 +##期 +##朦 +##朧 +##木 +##未 +##末 +##本 +##札 +##朮 +##术 +##朱 +##朴 +##朵 +##机 +##朽 +##杀 +##杂 +##权 +##杆 +##杈 +##杉 +##李 +##杏 +##材 +##村 +##杓 +##杖 +##杜 +##杞 +##束 +##杠 +##条 +##来 +##杨 +##杭 +##杯 +##杰 +##東 +##杳 +##杵 +##杷 +##杼 +##松 +##板 +##极 +##构 +##枇 +##枉 +##枋 +##析 +##枕 +##林 +##枚 +##果 +##枝 +##枢 +##枣 +##枪 +##枫 +##枭 +##枯 +##枰 +##枱 +##枳 +##架 +##枷 +##枸 +##柄 +##柏 +##某 +##柑 +##柒 +##染 +##柔 +##柘 +##柚 +##柜 +##柞 +##柠 +##柢 +##查 +##柩 +##柬 +##柯 +##柱 +##柳 +##柴 +##柵 +##査 +##柿 +##栀 +##栃 +##栄 +##栅 +##标 +##栈 +##栉 +##栋 +##栎 +##栏 +##树 +##栓 +##栖 +##栗 +##校 +##栩 +##株 +##样 +##核 +##根 +##格 +##栽 +##栾 +##桀 +##桁 +##桂 +##桃 +##桅 +##框 +##案 +##桉 +##桌 +##桎 +##桐 +##桑 +##桓 +##桔 +##桜 +##桠 +##桡 +##桢 +##档 +##桥 +##桦 +##桧 +##桨 +##桩 +##桶 +##桿 +##梁 +##梅 +##梆 +##梏 +##梓 +##梗 +##條 +##梟 +##梢 +##梦 +##梧 +##梨 +##梭 +##梯 +##械 +##梳 +##梵 +##梶 +##检 +##棂 +##棄 +##棉 +##棋 +##棍 +##棒 +##棕 +##棗 +##棘 +##棚 +##棟 +##棠 +##棣 +##棧 +##森 +##棱 +##棲 +##棵 +##棹 +##棺 +##椁 +##椅 +##椋 +##植 +##椎 +##椒 +##検 +##椪 +##椭 +##椰 +##椹 +##椽 +##椿 +##楂 +##楊 +##楓 +##楔 +##楚 +##楝 +##楞 +##楠 +##楣 +##楨 +##楫 +##業 +##楮 +##極 +##楷 +##楸 +##楹 +##楼 +##楽 +##概 +##榄 +##榆 +##榈 +##榉 +##榔 +##榕 +##榖 +##榛 +##榜 +##榨 +##榫 +##榭 +##榮 +##榱 +##榴 +##榷 +##榻 +##槁 +##槃 +##構 +##槌 +##槍 +##槎 +##槐 +##槓 +##様 +##槛 +##槟 +##槤 +##槭 +##槲 +##槳 +##槻 +##槽 +##槿 +##樁 +##樂 +##樊 +##樑 +##樓 +##標 +##樞 +##樟 +##模 +##樣 +##権 +##横 +##樫 +##樯 +##樱 +##樵 +##樸 +##樹 +##樺 +##樽 +##樾 +##橄 +##橇 +##橋 +##橐 +##橘 +##橙 +##機 +##橡 +##橢 +##橫 +##橱 +##橹 +##橼 +##檀 +##檄 +##檎 +##檐 +##檔 +##檗 +##檜 +##檢 +##檬 +##檯 +##檳 +##檸 +##檻 +##櫃 +##櫚 +##櫛 +##櫥 +##櫸 +##櫻 +##欄 +##權 +##欒 +##欖 +##欠 +##次 +##欢 +##欣 +##欧 +##欲 +##欸 +##欺 +##欽 +##款 +##歆 +##歇 +##歉 +##歌 +##歎 +##歐 +##歓 +##歙 +##歛 +##歡 +##止 +##正 +##此 +##步 +##武 +##歧 +##歩 +##歪 +##歯 +##歲 +##歳 +##歴 +##歷 +##歸 +##歹 +##死 +##歼 +##殁 +##殃 +##殆 +##殇 +##殉 +##殊 +##残 +##殒 +##殓 +##殖 +##殘 +##殞 +##殡 +##殤 +##殭 +##殯 +##殲 +##殴 +##段 +##殷 +##殺 +##殼 +##殿 +##毀 +##毁 +##毂 +##毅 +##毆 +##毋 +##母 +##毎 +##每 +##毒 +##毓 +##比 +##毕 +##毗 +##毘 +##毙 +##毛 +##毡 +##毫 +##毯 +##毽 +##氈 +##氏 +##氐 +##民 +##氓 +##气 +##氖 +##気 +##氙 +##氛 +##氟 +##氡 +##氢 +##氣 +##氤 +##氦 +##氧 +##氨 +##氪 +##氫 +##氮 +##氯 +##氰 +##氲 +##水 +##氷 +##永 +##氹 +##氾 +##汀 +##汁 +##求 +##汆 +##汇 +##汉 +##汎 +##汐 +##汕 +##汗 +##汙 +##汛 +##汝 +##汞 +##江 +##池 +##污 +##汤 +##汨 +##汩 +##汪 +##汰 +##汲 +##汴 +##汶 +##汹 +##決 +##汽 +##汾 +##沁 +##沂 +##沃 +##沅 +##沈 +##沉 +##沌 +##沏 +##沐 +##沒 +##沓 +##沖 +##沙 +##沛 +##沟 +##没 +##沢 +##沣 +##沥 +##沦 +##沧 +##沪 +##沫 +##沭 +##沮 +##沱 +##河 +##沸 +##油 +##治 +##沼 +##沽 +##沾 +##沿 +##況 +##泄 +##泉 +##泊 +##泌 +##泓 +##法 +##泗 +##泛 +##泞 +##泠 +##泡 +##波 +##泣 +##泥 +##注 +##泪 +##泫 +##泮 +##泯 +##泰 +##泱 +##泳 +##泵 +##泷 +##泸 +##泻 +##泼 +##泽 +##泾 +##洁 +##洄 +##洋 +##洒 +##洗 +##洙 +##洛 +##洞 +##津 +##洩 +##洪 +##洮 +##洱 +##洲 +##洵 +##洶 +##洸 +##洹 +##活 +##洼 +##洽 +##派 +##流 +##浃 +##浄 +##浅 +##浆 +##浇 +##浊 +##测 +##济 +##浏 +##浑 +##浒 +##浓 +##浔 +##浙 +##浚 +##浜 +##浣 +##浦 +##浩 +##浪 +##浬 +##浮 +##浯 +##浴 +##海 +##浸 +##涂 +##涅 +##涇 +##消 +##涉 +##涌 +##涎 +##涓 +##涔 +##涕 +##涙 +##涛 +##涝 +##涞 +##涟 +##涠 +##涡 +##涣 +##涤 +##润 +##涧 +##涨 +##涩 +##涪 +##涮 +##涯 +##液 +##涵 +##涸 +##涼 +##涿 +##淀 +##淄 +##淅 +##淆 +##淇 +##淋 +##淌 +##淑 +##淒 +##淖 +##淘 +##淙 +##淚 +##淞 +##淡 +##淤 +##淦 +##淨 +##淩 +##淪 +##淫 +##淬 +##淮 +##深 +##淳 +##淵 +##混 +##淹 +##淺 +##添 +##淼 +##清 +##済 +##渉 +##渊 +##渋 +##渍 +##渎 +##渐 +##渔 +##渗 +##渙 +##渚 +##減 +##渝 +##渠 +##渡 +##渣 +##渤 +##渥 +##渦 +##温 +##測 +##渭 +##港 +##渲 +##渴 +##游 +##渺 +##渾 +##湃 +##湄 +##湊 +##湍 +##湖 +##湘 +##湛 +##湟 +##湧 +##湫 +##湮 +##湯 +##湳 +##湾 +##湿 +##満 +##溃 +##溅 +##溉 +##溏 +##源 +##準 +##溜 +##溝 +##溟 +##溢 +##溥 +##溧 +##溪 +##溫 +##溯 +##溱 +##溴 +##溶 +##溺 +##溼 +##滁 +##滂 +##滄 +##滅 +##滇 +##滋 +##滌 +##滑 +##滓 +##滔 +##滕 +##滙 +##滚 +##滝 +##滞 +##滟 +##满 +##滢 +##滤 +##滥 +##滦 +##滨 +##滩 +##滬 +##滯 +##滲 +##滴 +##滷 +##滸 +##滾 +##滿 +##漁 +##漂 +##漆 +##漉 +##漏 +##漓 +##演 +##漕 +##漠 +##漢 +##漣 +##漩 +##漪 +##漫 +##漬 +##漯 +##漱 +##漲 +##漳 +##漸 +##漾 +##漿 +##潆 +##潇 +##潋 +##潍 +##潑 +##潔 +##潘 +##潛 +##潜 +##潞 +##潟 +##潢 +##潤 +##潦 +##潧 +##潭 +##潮 +##潰 +##潴 +##潸 +##潺 +##潼 +##澀 +##澄 +##澆 +##澈 +##澍 +##澎 +##澗 +##澜 +##澡 +##澤 +##澧 +##澱 +##澳 +##澹 +##激 +##濁 +##濂 +##濃 +##濑 +##濒 +##濕 +##濘 +##濛 +##濟 +##濠 +##濡 +##濤 +##濫 +##濬 +##濮 +##濯 +##濱 +##濺 +##濾 +##瀅 +##瀆 +##瀉 +##瀋 +##瀏 +##瀑 +##瀕 +##瀘 +##瀚 +##瀛 +##瀝 +##瀞 +##瀟 +##瀧 +##瀨 +##瀬 +##瀰 +##瀾 +##灌 +##灏 +##灑 +##灘 +##灝 +##灞 +##灣 +##火 +##灬 +##灭 +##灯 +##灰 +##灵 +##灶 +##灸 +##灼 +##災 +##灾 +##灿 +##炀 +##炁 +##炅 +##炉 +##炊 +##炎 +##炒 +##炔 +##炕 +##炖 +##炙 +##炜 +##炫 +##炬 +##炭 +##炮 +##炯 +##炳 +##炷 +##炸 +##点 +##為 +##炼 +##炽 +##烁 +##烂 +##烃 +##烈 +##烊 +##烏 +##烘 +##烙 +##烛 +##烟 +##烤 +##烦 +##烧 +##烨 +##烩 +##烫 +##烬 +##热 +##烯 +##烷 +##烹 +##烽 +##焉 +##焊 +##焕 +##焖 +##焗 +##焘 +##焙 +##焚 +##焜 +##無 +##焦 +##焯 +##焰 +##焱 +##然 +##焼 +##煅 +##煉 +##煊 +##煌 +##煎 +##煒 +##煖 +##煙 +##煜 +##煞 +##煤 +##煥 +##煦 +##照 +##煨 +##煩 +##煮 +##煲 +##煸 +##煽 +##熄 +##熊 +##熏 +##熒 +##熔 +##熙 +##熟 +##熠 +##熨 +##熬 +##熱 +##熵 +##熹 +##熾 +##燁 +##燃 +##燄 +##燈 +##燉 +##燊 +##燎 +##燒 +##燔 +##燕 +##燙 +##燜 +##營 +##燥 +##燦 +##燧 +##燭 +##燮 +##燴 +##燻 +##燼 +##燿 +##爆 +##爍 +##爐 +##爛 +##爪 +##爬 +##爭 +##爰 +##爱 +##爲 +##爵 +##父 +##爷 +##爸 +##爹 +##爺 +##爻 +##爽 +##爾 +##牆 +##片 +##版 +##牌 +##牍 +##牒 +##牙 +##牛 +##牝 +##牟 +##牠 +##牡 +##牢 +##牦 +##牧 +##物 +##牯 +##牲 +##牴 +##牵 +##特 +##牺 +##牽 +##犀 +##犁 +##犄 +##犊 +##犍 +##犒 +##犢 +##犧 +##犬 +##犯 +##状 +##犷 +##犸 +##犹 +##狀 +##狂 +##狄 +##狈 +##狎 +##狐 +##狒 +##狗 +##狙 +##狞 +##狠 +##狡 +##狩 +##独 +##狭 +##狮 +##狰 +##狱 +##狸 +##狹 +##狼 +##狽 +##猎 +##猕 +##猖 +##猗 +##猙 +##猛 +##猜 +##猝 +##猥 +##猩 +##猪 +##猫 +##猬 +##献 +##猴 +##猶 +##猷 +##猾 +##猿 +##獄 +##獅 +##獎 +##獐 +##獒 +##獗 +##獠 +##獣 +##獨 +##獭 +##獰 +##獲 +##獵 +##獷 +##獸 +##獺 +##獻 +##獼 +##獾 +##玄 +##率 +##玉 +##王 +##玑 +##玖 +##玛 +##玟 +##玠 +##玥 +##玩 +##玫 +##玮 +##环 +##现 +##玲 +##玳 +##玷 +##玺 +##玻 +##珀 +##珂 +##珅 +##珈 +##珉 +##珊 +##珍 +##珏 +##珐 +##珑 +##珙 +##珞 +##珠 +##珣 +##珥 +##珩 +##珪 +##班 +##珮 +##珲 +##珺 +##現 +##球 +##琅 +##理 +##琇 +##琉 +##琊 +##琍 +##琏 +##琐 +##琛 +##琢 +##琥 +##琦 +##琨 +##琪 +##琬 +##琮 +##琰 +##琲 +##琳 +##琴 +##琵 +##琶 +##琺 +##琼 +##瑀 +##瑁 +##瑄 +##瑋 +##瑕 +##瑗 +##瑙 +##瑚 +##瑛 +##瑜 +##瑞 +##瑟 +##瑠 +##瑣 +##瑤 +##瑩 +##瑪 +##瑯 +##瑰 +##瑶 +##瑾 +##璀 +##璁 +##璃 +##璇 +##璉 +##璋 +##璎 +##璐 +##璜 +##璞 +##璟 +##璧 +##璨 +##環 +##璽 +##璿 +##瓊 +##瓏 +##瓒 +##瓜 +##瓢 +##瓣 +##瓤 +##瓦 +##瓮 +##瓯 +##瓴 +##瓶 +##瓷 +##甄 +##甌 +##甕 +##甘 +##甙 +##甚 +##甜 +##生 +##產 +##産 +##甥 +##甦 +##用 +##甩 +##甫 +##甬 +##甭 +##甯 +##田 +##由 +##甲 +##申 +##电 +##男 +##甸 +##町 +##画 +##甾 +##畀 +##畅 +##界 +##畏 +##畑 +##畔 +##留 +##畜 +##畝 +##畢 +##略 +##畦 +##番 +##畫 +##異 +##畲 +##畳 +##畴 +##當 +##畸 +##畹 +##畿 +##疆 +##疇 +##疊 +##疏 +##疑 +##疔 +##疖 +##疗 +##疙 +##疚 +##疝 +##疟 +##疡 +##疣 +##疤 +##疥 +##疫 +##疮 +##疯 +##疱 +##疲 +##疳 +##疵 +##疸 +##疹 +##疼 +##疽 +##疾 +##痂 +##病 +##症 +##痈 +##痉 +##痊 +##痍 +##痒 +##痔 +##痕 +##痘 +##痙 +##痛 +##痞 +##痠 +##痢 +##痣 +##痤 +##痧 +##痨 +##痪 +##痫 +##痰 +##痱 +##痴 +##痹 +##痺 +##痼 +##痿 +##瘀 +##瘁 +##瘋 +##瘍 +##瘓 +##瘘 +##瘙 +##瘟 +##瘠 +##瘡 +##瘢 +##瘤 +##瘦 +##瘧 +##瘩 +##瘪 +##瘫 +##瘴 +##瘸 +##瘾 +##療 +##癇 +##癌 +##癒 +##癖 +##癜 +##癞 +##癡 +##癢 +##癣 +##癥 +##癫 +##癬 +##癮 +##癱 +##癲 +##癸 +##発 +##登 +##發 +##白 +##百 +##皂 +##的 +##皆 +##皇 +##皈 +##皋 +##皎 +##皑 +##皓 +##皖 +##皙 +##皚 +##皮 +##皰 +##皱 +##皴 +##皺 +##皿 +##盂 +##盃 +##盅 +##盆 +##盈 +##益 +##盎 +##盏 +##盐 +##监 +##盒 +##盔 +##盖 +##盗 +##盘 +##盛 +##盜 +##盞 +##盟 +##盡 +##監 +##盤 +##盥 +##盧 +##盪 +##目 +##盯 +##盱 +##盲 +##直 +##相 +##盹 +##盼 +##盾 +##省 +##眈 +##眉 +##看 +##県 +##眙 +##眞 +##真 +##眠 +##眦 +##眨 +##眩 +##眯 +##眶 +##眷 +##眸 +##眺 +##眼 +##眾 +##着 +##睁 +##睇 +##睏 +##睐 +##睑 +##睛 +##睜 +##睞 +##睡 +##睢 +##督 +##睥 +##睦 +##睨 +##睪 +##睫 +##睬 +##睹 +##睽 +##睾 +##睿 +##瞄 +##瞅 +##瞇 +##瞋 +##瞌 +##瞎 +##瞑 +##瞒 +##瞓 +##瞞 +##瞟 +##瞠 +##瞥 +##瞧 +##瞩 +##瞪 +##瞬 +##瞭 +##瞰 +##瞳 +##瞻 +##瞼 +##瞿 +##矇 +##矍 +##矗 +##矚 +##矛 +##矜 +##矢 +##矣 +##知 +##矩 +##矫 +##短 +##矮 +##矯 +##石 +##矶 +##矽 +##矾 +##矿 +##码 +##砂 +##砌 +##砍 +##砒 +##研 +##砖 +##砗 +##砚 +##砝 +##砣 +##砥 +##砧 +##砭 +##砰 +##砲 +##破 +##砷 +##砸 +##砺 +##砼 +##砾 +##础 +##硅 +##硐 +##硒 +##硕 +##硝 +##硫 +##硬 +##确 +##硯 +##硼 +##碁 +##碇 +##碉 +##碌 +##碍 +##碎 +##碑 +##碓 +##碗 +##碘 +##碚 +##碛 +##碟 +##碣 +##碧 +##碩 +##碰 +##碱 +##碳 +##碴 +##確 +##碼 +##碾 +##磁 +##磅 +##磊 +##磋 +##磐 +##磕 +##磚 +##磡 +##磨 +##磬 +##磯 +##磲 +##磷 +##磺 +##礁 +##礎 +##礙 +##礡 +##礦 +##礪 +##礫 +##礴 +##示 +##礼 +##社 +##祀 +##祁 +##祂 +##祇 +##祈 +##祉 +##祎 +##祐 +##祕 +##祖 +##祗 +##祚 +##祛 +##祜 +##祝 +##神 +##祟 +##祠 +##祢 +##祥 +##票 +##祭 +##祯 +##祷 +##祸 +##祺 +##祿 +##禀 +##禁 +##禄 +##禅 +##禍 +##禎 +##福 +##禛 +##禦 +##禧 +##禪 +##禮 +##禱 +##禹 +##禺 +##离 +##禽 +##禾 +##禿 +##秀 +##私 +##秃 +##秆 +##秉 +##秋 +##种 +##科 +##秒 +##秘 +##租 +##秣 +##秤 +##秦 +##秧 +##秩 +##秭 +##积 +##称 +##秸 +##移 +##秽 +##稀 +##稅 +##程 +##稍 +##税 +##稔 +##稗 +##稚 +##稜 +##稞 +##稟 +##稠 +##稣 +##種 +##稱 +##稲 +##稳 +##稷 +##稹 +##稻 +##稼 +##稽 +##稿 +##穀 +##穂 +##穆 +##穌 +##積 +##穎 +##穗 +##穢 +##穩 +##穫 +##穴 +##究 +##穷 +##穹 +##空 +##穿 +##突 +##窃 +##窄 +##窈 +##窍 +##窑 +##窒 +##窓 +##窕 +##窖 +##窗 +##窘 +##窜 +##窝 +##窟 +##窠 +##窥 +##窦 +##窨 +##窩 +##窪 +##窮 +##窯 +##窺 +##窿 +##竄 +##竅 +##竇 +##竊 +##立 +##竖 +##站 +##竜 +##竞 +##竟 +##章 +##竣 +##童 +##竭 +##端 +##競 +##竹 +##竺 +##竽 +##竿 +##笃 +##笆 +##笈 +##笋 +##笏 +##笑 +##笔 +##笙 +##笛 +##笞 +##笠 +##符 +##笨 +##第 +##笹 +##笺 +##笼 +##筆 +##等 +##筊 +##筋 +##筍 +##筏 +##筐 +##筑 +##筒 +##答 +##策 +##筛 +##筝 +##筠 +##筱 +##筲 +##筵 +##筷 +##筹 +##签 +##简 +##箇 +##箋 +##箍 +##箏 +##箐 +##箔 +##箕 +##算 +##箝 +##管 +##箩 +##箫 +##箭 +##箱 +##箴 +##箸 +##節 +##篁 +##範 +##篆 +##篇 +##築 +##篑 +##篓 +##篙 +##篝 +##篠 +##篡 +##篤 +##篩 +##篪 +##篮 +##篱 +##篷 +##簇 +##簌 +##簍 +##簡 +##簦 +##簧 +##簪 +##簫 +##簷 +##簸 +##簽 +##簾 +##簿 +##籁 +##籃 +##籌 +##籍 +##籐 +##籟 +##籠 +##籤 +##籬 +##籮 +##籲 +##米 +##类 +##籼 +##籽 +##粄 +##粉 +##粑 +##粒 +##粕 +##粗 +##粘 +##粟 +##粤 +##粥 +##粧 +##粪 +##粮 +##粱 +##粲 +##粳 +##粵 +##粹 +##粼 +##粽 +##精 +##粿 +##糅 +##糊 +##糍 +##糕 +##糖 +##糗 +##糙 +##糜 +##糞 +##糟 +##糠 +##糧 +##糬 +##糯 +##糰 +##糸 +##系 +##糾 +##紀 +##紂 +##約 +##紅 +##紉 +##紊 +##紋 +##納 +##紐 +##紓 +##純 +##紗 +##紘 +##紙 +##級 +##紛 +##紜 +##素 +##紡 +##索 +##紧 +##紫 +##紮 +##累 +##細 +##紳 +##紹 +##紺 +##終 +##絃 +##組 +##絆 +##経 +##結 +##絕 +##絞 +##絡 +##絢 +##給 +##絨 +##絮 +##統 +##絲 +##絳 +##絵 +##絶 +##絹 +##綁 +##綏 +##綑 +##經 +##継 +##続 +##綜 +##綠 +##綢 +##綦 +##綫 +##綬 +##維 +##綱 +##網 +##綴 +##綵 +##綸 +##綺 +##綻 +##綽 +##綾 +##綿 +##緊 +##緋 +##総 +##緑 +##緒 +##緘 +##線 +##緝 +##緞 +##締 +##緣 +##編 +##緩 +##緬 +##緯 +##練 +##緹 +##緻 +##縁 +##縄 +##縈 +##縛 +##縝 +##縣 +##縫 +##縮 +##縱 +##縴 +##縷 +##總 +##績 +##繁 +##繃 +##繆 +##繇 +##繋 +##織 +##繕 +##繚 +##繞 +##繡 +##繩 +##繪 +##繫 +##繭 +##繳 +##繹 +##繼 +##繽 +##纂 +##續 +##纍 +##纏 +##纓 +##纔 +##纖 +##纜 +##纠 +##红 +##纣 +##纤 +##约 +##级 +##纨 +##纪 +##纫 +##纬 +##纭 +##纯 +##纰 +##纱 +##纲 +##纳 +##纵 +##纶 +##纷 +##纸 +##纹 +##纺 +##纽 +##纾 +##线 +##绀 +##练 +##组 +##绅 +##细 +##织 +##终 +##绊 +##绍 +##绎 +##经 +##绑 +##绒 +##结 +##绔 +##绕 +##绘 +##给 +##绚 +##绛 +##络 +##绝 +##绞 +##统 +##绡 +##绢 +##绣 +##绥 +##绦 +##继 +##绩 +##绪 +##绫 +##续 +##绮 +##绯 +##绰 +##绳 +##维 +##绵 +##绶 +##绷 +##绸 +##绻 +##综 +##绽 +##绾 +##绿 +##缀 +##缄 +##缅 +##缆 +##缇 +##缈 +##缉 +##缎 +##缓 +##缔 +##缕 +##编 +##缘 +##缙 +##缚 +##缜 +##缝 +##缠 +##缢 +##缤 +##缥 +##缨 +##缩 +##缪 +##缭 +##缮 +##缰 +##缱 +##缴 +##缸 +##缺 +##缽 +##罂 +##罄 +##罌 +##罐 +##网 +##罔 +##罕 +##罗 +##罚 +##罡 +##罢 +##罩 +##罪 +##置 +##罰 +##署 +##罵 +##罷 +##罹 +##羁 +##羅 +##羈 +##羊 +##羌 +##美 +##羔 +##羚 +##羞 +##羟 +##羡 +##羣 +##群 +##羥 +##羧 +##羨 +##義 +##羯 +##羲 +##羸 +##羹 +##羽 +##羿 +##翁 +##翅 +##翊 +##翌 +##翎 +##習 +##翔 +##翘 +##翟 +##翠 +##翡 +##翦 +##翩 +##翰 +##翱 +##翳 +##翹 +##翻 +##翼 +##耀 +##老 +##考 +##耄 +##者 +##耆 +##耋 +##而 +##耍 +##耐 +##耒 +##耕 +##耗 +##耘 +##耙 +##耦 +##耨 +##耳 +##耶 +##耷 +##耸 +##耻 +##耽 +##耿 +##聂 +##聆 +##聊 +##聋 +##职 +##聒 +##联 +##聖 +##聘 +##聚 +##聞 +##聪 +##聯 +##聰 +##聲 +##聳 +##聴 +##聶 +##職 +##聽 +##聾 +##聿 +##肃 +##肄 +##肅 +##肆 +##肇 +##肉 +##肋 +##肌 +##肏 +##肓 +##肖 +##肘 +##肚 +##肛 +##肝 +##肠 +##股 +##肢 +##肤 +##肥 +##肩 +##肪 +##肮 +##肯 +##肱 +##育 +##肴 +##肺 +##肽 +##肾 +##肿 +##胀 +##胁 +##胃 +##胄 +##胆 +##背 +##胍 +##胎 +##胖 +##胚 +##胛 +##胜 +##胝 +##胞 +##胡 +##胤 +##胥 +##胧 +##胫 +##胭 +##胯 +##胰 +##胱 +##胳 +##胴 +##胶 +##胸 +##胺 +##能 +##脂 +##脅 +##脆 +##脇 +##脈 +##脉 +##脊 +##脍 +##脏 +##脐 +##脑 +##脓 +##脖 +##脘 +##脚 +##脛 +##脣 +##脩 +##脫 +##脯 +##脱 +##脲 +##脳 +##脸 +##脹 +##脾 +##腆 +##腈 +##腊 +##腋 +##腌 +##腎 +##腐 +##腑 +##腓 +##腔 +##腕 +##腥 +##腦 +##腩 +##腫 +##腭 +##腮 +##腰 +##腱 +##腳 +##腴 +##腸 +##腹 +##腺 +##腻 +##腼 +##腾 +##腿 +##膀 +##膈 +##膊 +##膏 +##膑 +##膘 +##膚 +##膛 +##膜 +##膝 +##膠 +##膦 +##膨 +##膩 +##膳 +##膺 +##膻 +##膽 +##膾 +##膿 +##臀 +##臂 +##臃 +##臆 +##臉 +##臊 +##臍 +##臓 +##臘 +##臟 +##臣 +##臥 +##臧 +##臨 +##自 +##臬 +##臭 +##至 +##致 +##臺 +##臻 +##臼 +##臾 +##舀 +##舂 +##舅 +##舆 +##與 +##興 +##舉 +##舊 +##舌 +##舍 +##舎 +##舐 +##舒 +##舔 +##舖 +##舗 +##舛 +##舜 +##舞 +##舟 +##航 +##舫 +##般 +##舰 +##舱 +##舵 +##舶 +##舷 +##舸 +##船 +##舺 +##舾 +##艇 +##艋 +##艘 +##艙 +##艦 +##艮 +##良 +##艰 +##艱 +##色 +##艳 +##艷 +##艹 +##艺 +##艾 +##节 +##芃 +##芈 +##芊 +##芋 +##芍 +##芎 +##芒 +##芙 +##芜 +##芝 +##芡 +##芥 +##芦 +##芩 +##芪 +##芫 +##芬 +##芭 +##芮 +##芯 +##花 +##芳 +##芷 +##芸 +##芹 +##芻 +##芽 +##芾 +##苁 +##苄 +##苇 +##苋 +##苍 +##苏 +##苑 +##苒 +##苓 +##苔 +##苕 +##苗 +##苛 +##苜 +##苞 +##苟 +##苡 +##苣 +##若 +##苦 +##苫 +##苯 +##英 +##苷 +##苹 +##苻 +##茁 +##茂 +##范 +##茄 +##茅 +##茉 +##茎 +##茏 +##茗 +##茜 +##茧 +##茨 +##茫 +##茬 +##茭 +##茯 +##茱 +##茲 +##茴 +##茵 +##茶 +##茸 +##茹 +##茼 +##荀 +##荃 +##荆 +##草 +##荊 +##荏 +##荐 +##荒 +##荔 +##荖 +##荘 +##荚 +##荞 +##荟 +##荠 +##荡 +##荣 +##荤 +##荥 +##荧 +##荨 +##荪 +##荫 +##药 +##荳 +##荷 +##荸 +##荻 +##荼 +##荽 +##莅 +##莆 +##莉 +##莊 +##莎 +##莒 +##莓 +##莖 +##莘 +##莞 +##莠 +##莢 +##莧 +##莪 +##莫 +##莱 +##莲 +##莴 +##获 +##莹 +##莺 +##莽 +##莿 +##菀 +##菁 +##菅 +##菇 +##菈 +##菊 +##菌 +##菏 +##菓 +##菖 +##菘 +##菜 +##菟 +##菠 +##菡 +##菩 +##華 +##菱 +##菲 +##菸 +##菽 +##萁 +##萃 +##萄 +##萊 +##萋 +##萌 +##萍 +##萎 +##萘 +##萝 +##萤 +##营 +##萦 +##萧 +##萨 +##萩 +##萬 +##萱 +##萵 +##萸 +##萼 +##落 +##葆 +##葉 +##著 +##葚 +##葛 +##葡 +##董 +##葦 +##葩 +##葫 +##葬 +##葭 +##葯 +##葱 +##葳 +##葵 +##葷 +##葺 +##蒂 +##蒋 +##蒐 +##蒔 +##蒙 +##蒜 +##蒞 +##蒟 +##蒡 +##蒨 +##蒲 +##蒸 +##蒹 +##蒻 +##蒼 +##蒿 +##蓁 +##蓄 +##蓆 +##蓉 +##蓋 +##蓑 +##蓓 +##蓖 +##蓝 +##蓟 +##蓦 +##蓬 +##蓮 +##蓼 +##蓿 +##蔑 +##蔓 +##蔔 +##蔗 +##蔘 +##蔚 +##蔡 +##蔣 +##蔥 +##蔫 +##蔬 +##蔭 +##蔵 +##蔷 +##蔺 +##蔻 +##蔼 +##蔽 +##蕁 +##蕃 +##蕈 +##蕉 +##蕊 +##蕎 +##蕙 +##蕤 +##蕨 +##蕩 +##蕪 +##蕭 +##蕲 +##蕴 +##蕻 +##蕾 +##薄 +##薅 +##薇 +##薈 +##薊 +##薏 +##薑 +##薔 +##薙 +##薛 +##薦 +##薨 +##薩 +##薪 +##薬 +##薯 +##薰 +##薹 +##藉 +##藍 +##藏 +##藐 +##藓 +##藕 +##藜 +##藝 +##藤 +##藥 +##藩 +##藹 +##藻 +##藿 +##蘆 +##蘇 +##蘊 +##蘋 +##蘑 +##蘚 +##蘭 +##蘸 +##蘼 +##蘿 +##虎 +##虏 +##虐 +##虑 +##虔 +##處 +##虚 +##虛 +##虜 +##虞 +##號 +##虢 +##虧 +##虫 +##虬 +##虱 +##虹 +##虻 +##虽 +##虾 +##蚀 +##蚁 +##蚂 +##蚊 +##蚌 +##蚓 +##蚕 +##蚜 +##蚝 +##蚣 +##蚤 +##蚩 +##蚪 +##蚯 +##蚱 +##蚵 +##蛀 +##蛆 +##蛇 +##蛊 +##蛋 +##蛎 +##蛐 +##蛔 +##蛙 +##蛛 +##蛟 +##蛤 +##蛭 +##蛮 +##蛰 +##蛳 +##蛹 +##蛻 +##蛾 +##蜀 +##蜂 +##蜃 +##蜆 +##蜇 +##蜈 +##蜊 +##蜍 +##蜒 +##蜓 +##蜕 +##蜗 +##蜘 +##蜚 +##蜜 +##蜡 +##蜢 +##蜥 +##蜱 +##蜴 +##蜷 +##蜻 +##蜿 +##蝇 +##蝈 +##蝉 +##蝌 +##蝎 +##蝕 +##蝗 +##蝙 +##蝟 +##蝠 +##蝦 +##蝨 +##蝴 +##蝶 +##蝸 +##蝼 +##螂 +##螃 +##融 +##螞 +##螢 +##螨 +##螯 +##螳 +##螺 +##蟀 +##蟄 +##蟆 +##蟋 +##蟎 +##蟑 +##蟒 +##蟠 +##蟬 +##蟲 +##蟹 +##蟻 +##蟾 +##蠅 +##蠍 +##蠔 +##蠕 +##蠛 +##蠟 +##蠡 +##蠢 +##蠣 +##蠱 +##蠶 +##蠹 +##蠻 +##血 +##衄 +##衅 +##衆 +##行 +##衍 +##術 +##衔 +##街 +##衙 +##衛 +##衝 +##衞 +##衡 +##衢 +##衣 +##补 +##表 +##衩 +##衫 +##衬 +##衮 +##衰 +##衲 +##衷 +##衹 +##衾 +##衿 +##袁 +##袂 +##袄 +##袅 +##袈 +##袋 +##袍 +##袒 +##袖 +##袜 +##袞 +##袤 +##袪 +##被 +##袭 +##袱 +##裁 +##裂 +##装 +##裆 +##裊 +##裏 +##裔 +##裕 +##裘 +##裙 +##補 +##裝 +##裟 +##裡 +##裤 +##裨 +##裱 +##裳 +##裴 +##裸 +##裹 +##製 +##裾 +##褂 +##複 +##褐 +##褒 +##褓 +##褔 +##褚 +##褥 +##褪 +##褫 +##褲 +##褶 +##褻 +##襁 +##襄 +##襟 +##襠 +##襪 +##襬 +##襯 +##襲 +##西 +##要 +##覃 +##覆 +##覇 +##見 +##規 +##覓 +##視 +##覚 +##覦 +##覧 +##親 +##覬 +##観 +##覷 +##覺 +##覽 +##觀 +##见 +##观 +##规 +##觅 +##视 +##览 +##觉 +##觊 +##觎 +##觐 +##觑 +##角 +##觞 +##解 +##觥 +##触 +##觸 +##言 +##訂 +##計 +##訊 +##討 +##訓 +##訕 +##訖 +##託 +##記 +##訛 +##訝 +##訟 +##訣 +##訥 +##訪 +##設 +##許 +##訳 +##訴 +##訶 +##診 +##註 +##証 +##詆 +##詐 +##詔 +##評 +##詛 +##詞 +##詠 +##詡 +##詢 +##詣 +##試 +##詩 +##詫 +##詬 +##詭 +##詮 +##詰 +##話 +##該 +##詳 +##詹 +##詼 +##誅 +##誇 +##誉 +##誌 +##認 +##誓 +##誕 +##誘 +##語 +##誠 +##誡 +##誣 +##誤 +##誥 +##誦 +##誨 +##說 +##説 +##読 +##誰 +##課 +##誹 +##誼 +##調 +##諄 +##談 +##請 +##諏 +##諒 +##論 +##諗 +##諜 +##諡 +##諦 +##諧 +##諫 +##諭 +##諮 +##諱 +##諳 +##諷 +##諸 +##諺 +##諾 +##謀 +##謁 +##謂 +##謄 +##謊 +##謎 +##謐 +##謔 +##謗 +##謙 +##講 +##謝 +##謠 +##謨 +##謬 +##謹 +##謾 +##譁 +##證 +##譎 +##譏 +##識 +##譙 +##譚 +##譜 +##警 +##譬 +##譯 +##議 +##譲 +##譴 +##護 +##譽 +##讀 +##變 +##讓 +##讚 +##讞 +##计 +##订 +##认 +##讥 +##讧 +##讨 +##让 +##讪 +##讫 +##训 +##议 +##讯 +##记 +##讲 +##讳 +##讴 +##讶 +##讷 +##许 +##讹 +##论 +##讼 +##讽 +##设 +##访 +##诀 +##证 +##诃 +##评 +##诅 +##识 +##诈 +##诉 +##诊 +##诋 +##词 +##诏 +##译 +##试 +##诗 +##诘 +##诙 +##诚 +##诛 +##话 +##诞 +##诟 +##诠 +##诡 +##询 +##诣 +##诤 +##该 +##详 +##诧 +##诩 +##诫 +##诬 +##语 +##误 +##诰 +##诱 +##诲 +##说 +##诵 +##诶 +##请 +##诸 +##诺 +##读 +##诽 +##课 +##诿 +##谀 +##谁 +##调 +##谄 +##谅 +##谆 +##谈 +##谊 +##谋 +##谌 +##谍 +##谎 +##谏 +##谐 +##谑 +##谒 +##谓 +##谔 +##谕 +##谗 +##谘 +##谙 +##谚 +##谛 +##谜 +##谟 +##谢 +##谣 +##谤 +##谥 +##谦 +##谧 +##谨 +##谩 +##谪 +##谬 +##谭 +##谯 +##谱 +##谲 +##谴 +##谶 +##谷 +##豁 +##豆 +##豇 +##豈 +##豉 +##豊 +##豌 +##豎 +##豐 +##豔 +##豚 +##象 +##豢 +##豪 +##豫 +##豬 +##豹 +##豺 +##貂 +##貅 +##貌 +##貓 +##貔 +##貘 +##貝 +##貞 +##負 +##財 +##貢 +##貧 +##貨 +##販 +##貪 +##貫 +##責 +##貯 +##貰 +##貳 +##貴 +##貶 +##買 +##貸 +##費 +##貼 +##貽 +##貿 +##賀 +##賁 +##賂 +##賃 +##賄 +##資 +##賈 +##賊 +##賑 +##賓 +##賜 +##賞 +##賠 +##賡 +##賢 +##賣 +##賤 +##賦 +##質 +##賬 +##賭 +##賴 +##賺 +##購 +##賽 +##贅 +##贈 +##贊 +##贍 +##贏 +##贓 +##贖 +##贛 +##贝 +##贞 +##负 +##贡 +##财 +##责 +##贤 +##败 +##账 +##货 +##质 +##贩 +##贪 +##贫 +##贬 +##购 +##贮 +##贯 +##贰 +##贱 +##贲 +##贴 +##贵 +##贷 +##贸 +##费 +##贺 +##贻 +##贼 +##贾 +##贿 +##赁 +##赂 +##赃 +##资 +##赅 +##赈 +##赊 +##赋 +##赌 +##赎 +##赏 +##赐 +##赓 +##赔 +##赖 +##赘 +##赚 +##赛 +##赝 +##赞 +##赠 +##赡 +##赢 +##赣 +##赤 +##赦 +##赧 +##赫 +##赭 +##走 +##赳 +##赴 +##赵 +##赶 +##起 +##趁 +##超 +##越 +##趋 +##趕 +##趙 +##趟 +##趣 +##趨 +##足 +##趴 +##趵 +##趸 +##趺 +##趾 +##跃 +##跄 +##跆 +##跋 +##跌 +##跎 +##跑 +##跖 +##跚 +##跛 +##距 +##跟 +##跡 +##跤 +##跨 +##跩 +##跪 +##路 +##跳 +##践 +##跷 +##跹 +##跺 +##跻 +##踉 +##踊 +##踌 +##踏 +##踐 +##踝 +##踞 +##踟 +##踢 +##踩 +##踪 +##踮 +##踱 +##踴 +##踵 +##踹 +##蹂 +##蹄 +##蹇 +##蹈 +##蹉 +##蹊 +##蹋 +##蹑 +##蹒 +##蹙 +##蹟 +##蹣 +##蹤 +##蹦 +##蹩 +##蹬 +##蹭 +##蹲 +##蹴 +##蹶 +##蹺 +##蹼 +##蹿 +##躁 +##躇 +##躉 +##躊 +##躋 +##躍 +##躏 +##躪 +##身 +##躬 +##躯 +##躲 +##躺 +##軀 +##車 +##軋 +##軌 +##軍 +##軒 +##軟 +##転 +##軸 +##軼 +##軽 +##軾 +##較 +##載 +##輒 +##輓 +##輔 +##輕 +##輛 +##輝 +##輟 +##輩 +##輪 +##輯 +##輸 +##輻 +##輾 +##輿 +##轄 +##轅 +##轆 +##轉 +##轍 +##轎 +##轟 +##车 +##轧 +##轨 +##轩 +##转 +##轭 +##轮 +##软 +##轰 +##轲 +##轴 +##轶 +##轻 +##轼 +##载 +##轿 +##较 +##辄 +##辅 +##辆 +##辇 +##辈 +##辉 +##辊 +##辍 +##辐 +##辑 +##输 +##辕 +##辖 +##辗 +##辘 +##辙 +##辛 +##辜 +##辞 +##辟 +##辣 +##辦 +##辨 +##辩 +##辫 +##辭 +##辮 +##辯 +##辰 +##辱 +##農 +##边 +##辺 +##辻 +##込 +##辽 +##达 +##迁 +##迂 +##迄 +##迅 +##过 +##迈 +##迎 +##运 +##近 +##返 +##还 +##这 +##进 +##远 +##违 +##连 +##迟 +##迢 +##迤 +##迥 +##迦 +##迩 +##迪 +##迫 +##迭 +##述 +##迴 +##迷 +##迸 +##迹 +##迺 +##追 +##退 +##送 +##适 +##逃 +##逅 +##逆 +##选 +##逊 +##逍 +##透 +##逐 +##递 +##途 +##逕 +##逗 +##這 +##通 +##逛 +##逝 +##逞 +##速 +##造 +##逢 +##連 +##逮 +##週 +##進 +##逵 +##逶 +##逸 +##逻 +##逼 +##逾 +##遁 +##遂 +##遅 +##遇 +##遊 +##運 +##遍 +##過 +##遏 +##遐 +##遑 +##遒 +##道 +##達 +##違 +##遗 +##遙 +##遛 +##遜 +##遞 +##遠 +##遢 +##遣 +##遥 +##遨 +##適 +##遭 +##遮 +##遲 +##遴 +##遵 +##遶 +##遷 +##選 +##遺 +##遼 +##遽 +##避 +##邀 +##邁 +##邂 +##邃 +##還 +##邇 +##邈 +##邊 +##邋 +##邏 +##邑 +##邓 +##邕 +##邛 +##邝 +##邢 +##那 +##邦 +##邨 +##邪 +##邬 +##邮 +##邯 +##邰 +##邱 +##邳 +##邵 +##邸 +##邹 +##邺 +##邻 +##郁 +##郅 +##郊 +##郎 +##郑 +##郜 +##郝 +##郡 +##郢 +##郤 +##郦 +##郧 +##部 +##郫 +##郭 +##郴 +##郵 +##郷 +##郸 +##都 +##鄂 +##鄉 +##鄒 +##鄔 +##鄙 +##鄞 +##鄢 +##鄧 +##鄭 +##鄰 +##鄱 +##鄲 +##鄺 +##酉 +##酊 +##酋 +##酌 +##配 +##酐 +##酒 +##酗 +##酚 +##酝 +##酢 +##酣 +##酥 +##酩 +##酪 +##酬 +##酮 +##酯 +##酰 +##酱 +##酵 +##酶 +##酷 +##酸 +##酿 +##醃 +##醇 +##醉 +##醋 +##醍 +##醐 +##醒 +##醚 +##醛 +##醜 +##醞 +##醣 +##醪 +##醫 +##醬 +##醮 +##醯 +##醴 +##醺 +##釀 +##釁 +##采 +##釉 +##释 +##釋 +##里 +##重 +##野 +##量 +##釐 +##金 +##釗 +##釘 +##釜 +##針 +##釣 +##釦 +##釧 +##釵 +##鈀 +##鈉 +##鈍 +##鈎 +##鈔 +##鈕 +##鈞 +##鈣 +##鈦 +##鈪 +##鈴 +##鈺 +##鈾 +##鉀 +##鉄 +##鉅 +##鉉 +##鉑 +##鉗 +##鉚 +##鉛 +##鉤 +##鉴 +##鉻 +##銀 +##銃 +##銅 +##銑 +##銓 +##銖 +##銘 +##銜 +##銬 +##銭 +##銮 +##銳 +##銷 +##銹 +##鋁 +##鋅 +##鋒 +##鋤 +##鋪 +##鋰 +##鋸 +##鋼 +##錄 +##錐 +##錘 +##錚 +##錠 +##錢 +##錦 +##錨 +##錫 +##錮 +##錯 +##録 +##錳 +##錶 +##鍊 +##鍋 +##鍍 +##鍛 +##鍥 +##鍰 +##鍵 +##鍺 +##鍾 +##鎂 +##鎊 +##鎌 +##鎏 +##鎔 +##鎖 +##鎗 +##鎚 +##鎧 +##鎬 +##鎮 +##鎳 +##鏈 +##鏖 +##鏗 +##鏘 +##鏞 +##鏟 +##鏡 +##鏢 +##鏤 +##鏽 +##鐘 +##鐮 +##鐲 +##鐳 +##鐵 +##鐸 +##鐺 +##鑄 +##鑊 +##鑑 +##鑒 +##鑣 +##鑫 +##鑰 +##鑲 +##鑼 +##鑽 +##鑾 +##鑿 +##针 +##钉 +##钊 +##钎 +##钏 +##钒 +##钓 +##钗 +##钙 +##钛 +##钜 +##钝 +##钞 +##钟 +##钠 +##钡 +##钢 +##钣 +##钤 +##钥 +##钦 +##钧 +##钨 +##钩 +##钮 +##钯 +##钰 +##钱 +##钳 +##钴 +##钵 +##钺 +##钻 +##钼 +##钾 +##钿 +##铀 +##铁 +##铂 +##铃 +##铄 +##铅 +##铆 +##铉 +##铎 +##铐 +##铛 +##铜 +##铝 +##铠 +##铡 +##铢 +##铣 +##铤 +##铨 +##铩 +##铬 +##铭 +##铮 +##铰 +##铲 +##铵 +##银 +##铸 +##铺 +##链 +##铿 +##销 +##锁 +##锂 +##锄 +##锅 +##锆 +##锈 +##锉 +##锋 +##锌 +##锏 +##锐 +##锑 +##错 +##锚 +##锟 +##锡 +##锢 +##锣 +##锤 +##锥 +##锦 +##锭 +##键 +##锯 +##锰 +##锲 +##锵 +##锹 +##锺 +##锻 +##镀 +##镁 +##镂 +##镇 +##镉 +##镌 +##镍 +##镐 +##镑 +##镕 +##镖 +##镗 +##镛 +##镜 +##镣 +##镭 +##镯 +##镰 +##镳 +##镶 +##長 +##长 +##門 +##閃 +##閉 +##開 +##閎 +##閏 +##閑 +##閒 +##間 +##閔 +##閘 +##閡 +##関 +##閣 +##閥 +##閨 +##閩 +##閱 +##閲 +##閹 +##閻 +##閾 +##闆 +##闇 +##闊 +##闌 +##闍 +##闔 +##闕 +##闖 +##闘 +##關 +##闡 +##闢 +##门 +##闪 +##闫 +##闭 +##问 +##闯 +##闰 +##闲 +##间 +##闵 +##闷 +##闸 +##闹 +##闺 +##闻 +##闽 +##闾 +##阀 +##阁 +##阂 +##阅 +##阆 +##阇 +##阈 +##阉 +##阎 +##阐 +##阑 +##阔 +##阕 +##阖 +##阙 +##阚 +##阜 +##队 +##阡 +##阪 +##阮 +##阱 +##防 +##阳 +##阴 +##阵 +##阶 +##阻 +##阿 +##陀 +##陂 +##附 +##际 +##陆 +##陇 +##陈 +##陋 +##陌 +##降 +##限 +##陕 +##陛 +##陝 +##陞 +##陟 +##陡 +##院 +##陣 +##除 +##陨 +##险 +##陪 +##陰 +##陲 +##陳 +##陵 +##陶 +##陷 +##陸 +##険 +##陽 +##隅 +##隆 +##隈 +##隊 +##隋 +##隍 +##階 +##随 +##隐 +##隔 +##隕 +##隘 +##隙 +##際 +##障 +##隠 +##隣 +##隧 +##隨 +##險 +##隱 +##隴 +##隶 +##隸 +##隻 +##隼 +##隽 +##难 +##雀 +##雁 +##雄 +##雅 +##集 +##雇 +##雉 +##雋 +##雌 +##雍 +##雎 +##雏 +##雑 +##雒 +##雕 +##雖 +##雙 +##雛 +##雜 +##雞 +##離 +##難 +##雨 +##雪 +##雯 +##雰 +##雲 +##雳 +##零 +##雷 +##雹 +##電 +##雾 +##需 +##霁 +##霄 +##霆 +##震 +##霈 +##霉 +##霊 +##霍 +##霎 +##霏 +##霑 +##霓 +##霖 +##霜 +##霞 +##霧 +##霭 +##霰 +##露 +##霸 +##霹 +##霽 +##霾 +##靂 +##靄 +##靈 +##青 +##靓 +##靖 +##静 +##靚 +##靛 +##靜 +##非 +##靠 +##靡 +##面 +##靥 +##靦 +##革 +##靳 +##靴 +##靶 +##靼 +##鞅 +##鞋 +##鞍 +##鞏 +##鞑 +##鞘 +##鞠 +##鞣 +##鞦 +##鞭 +##韆 +##韋 +##韌 +##韓 +##韜 +##韦 +##韧 +##韩 +##韬 +##韭 +##音 +##韵 +##韶 +##韻 +##響 +##頁 +##頂 +##頃 +##項 +##順 +##須 +##頌 +##預 +##頑 +##頒 +##頓 +##頗 +##領 +##頜 +##頡 +##頤 +##頫 +##頭 +##頰 +##頷 +##頸 +##頹 +##頻 +##頼 +##顆 +##題 +##額 +##顎 +##顏 +##顔 +##願 +##顛 +##類 +##顧 +##顫 +##顯 +##顱 +##顴 +##页 +##顶 +##顷 +##项 +##顺 +##须 +##顼 +##顽 +##顾 +##顿 +##颁 +##颂 +##预 +##颅 +##领 +##颇 +##颈 +##颉 +##颊 +##颌 +##颍 +##颐 +##频 +##颓 +##颔 +##颖 +##颗 +##题 +##颚 +##颛 +##颜 +##额 +##颞 +##颠 +##颡 +##颢 +##颤 +##颦 +##颧 +##風 +##颯 +##颱 +##颳 +##颶 +##颼 +##飄 +##飆 +##风 +##飒 +##飓 +##飕 +##飘 +##飙 +##飚 +##飛 +##飞 +##食 +##飢 +##飨 +##飩 +##飪 +##飯 +##飲 +##飼 +##飽 +##飾 +##餃 +##餅 +##餉 +##養 +##餌 +##餐 +##餒 +##餓 +##餘 +##餚 +##餛 +##餞 +##餡 +##館 +##餮 +##餵 +##餾 +##饅 +##饈 +##饋 +##饌 +##饍 +##饑 +##饒 +##饕 +##饗 +##饞 +##饥 +##饨 +##饪 +##饬 +##饭 +##饮 +##饯 +##饰 +##饱 +##饲 +##饴 +##饵 +##饶 +##饷 +##饺 +##饼 +##饽 +##饿 +##馀 +##馁 +##馄 +##馅 +##馆 +##馈 +##馋 +##馍 +##馏 +##馒 +##馔 +##首 +##馗 +##香 +##馥 +##馨 +##馬 +##馭 +##馮 +##馳 +##馴 +##駁 +##駄 +##駅 +##駆 +##駐 +##駒 +##駕 +##駛 +##駝 +##駭 +##駱 +##駿 +##騁 +##騎 +##騏 +##験 +##騙 +##騨 +##騰 +##騷 +##驀 +##驅 +##驊 +##驍 +##驒 +##驕 +##驗 +##驚 +##驛 +##驟 +##驢 +##驥 +##马 +##驭 +##驮 +##驯 +##驰 +##驱 +##驳 +##驴 +##驶 +##驷 +##驸 +##驹 +##驻 +##驼 +##驾 +##驿 +##骁 +##骂 +##骄 +##骅 +##骆 +##骇 +##骈 +##骊 +##骋 +##验 +##骏 +##骐 +##骑 +##骗 +##骚 +##骛 +##骜 +##骞 +##骠 +##骡 +##骤 +##骥 +##骧 +##骨 +##骯 +##骰 +##骶 +##骷 +##骸 +##骼 +##髂 +##髅 +##髋 +##髏 +##髒 +##髓 +##體 +##髖 +##高 +##髦 +##髪 +##髮 +##髯 +##髻 +##鬃 +##鬆 +##鬍 +##鬓 +##鬚 +##鬟 +##鬢 +##鬣 +##鬥 +##鬧 +##鬱 +##鬼 +##魁 +##魂 +##魄 +##魅 +##魇 +##魍 +##魏 +##魔 +##魘 +##魚 +##魯 +##魷 +##鮑 +##鮨 +##鮪 +##鮭 +##鮮 +##鯉 +##鯊 +##鯖 +##鯛 +##鯨 +##鯰 +##鯽 +##鰍 +##鰓 +##鰭 +##鰲 +##鰻 +##鰾 +##鱈 +##鱉 +##鱔 +##鱗 +##鱷 +##鱸 +##鱼 +##鱿 +##鲁 +##鲈 +##鲍 +##鲑 +##鲛 +##鲜 +##鲟 +##鲢 +##鲤 +##鲨 +##鲫 +##鲱 +##鲲 +##鲶 +##鲷 +##鲸 +##鳃 +##鳄 +##鳅 +##鳌 +##鳍 +##鳕 +##鳖 +##鳗 +##鳝 +##鳞 +##鳥 +##鳩 +##鳳 +##鳴 +##鳶 +##鴉 +##鴕 +##鴛 +##鴦 +##鴨 +##鴻 +##鴿 +##鵑 +##鵜 +##鵝 +##鵡 +##鵬 +##鵰 +##鵲 +##鶘 +##鶩 +##鶯 +##鶴 +##鷗 +##鷲 +##鷹 +##鷺 +##鸚 +##鸞 +##鸟 +##鸠 +##鸡 +##鸢 +##鸣 +##鸥 +##鸦 +##鸨 +##鸪 +##鸭 +##鸯 +##鸳 +##鸵 +##鸽 +##鸾 +##鸿 +##鹂 +##鹃 +##鹄 +##鹅 +##鹈 +##鹉 +##鹊 +##鹌 +##鹏 +##鹑 +##鹕 +##鹘 +##鹜 +##鹞 +##鹤 +##鹦 +##鹧 +##鹫 +##鹭 +##鹰 +##鹳 +##鹵 +##鹹 +##鹼 +##鹽 +##鹿 +##麂 +##麋 +##麒 +##麓 +##麗 +##麝 +##麟 +##麥 +##麦 +##麩 +##麴 +##麵 +##麸 +##麺 +##麻 +##麼 +##麽 +##麾 +##黃 +##黄 +##黍 +##黎 +##黏 +##黑 +##黒 +##黔 +##默 +##黛 +##黜 +##黝 +##點 +##黠 +##黨 +##黯 +##黴 +##鼋 +##鼎 +##鼐 +##鼓 +##鼠 +##鼬 +##鼹 +##鼻 +##鼾 +##齁 +##齊 +##齋 +##齐 +##齒 +##齡 +##齢 +##齣 +##齦 +##齿 +##龄 +##龅 +##龈 +##龊 +##龋 +##龌 +##龍 +##龐 +##龔 +##龕 +##龙 +##龚 +##龛 +##龜 +##龟 +##︰ +##︱ +##︶ +##︿ +##﹁ +##﹂ +##﹍ +##﹏ +##﹐ +##﹑ +##﹒ +##﹔ +##﹕ +##﹖ +##﹗ +##﹙ +##﹚ +##﹝ +##﹞ +##﹡ +##﹣ +##！ +##＂ +##＃ +##＄ +##％ +##＆ +##＇ +##（ +##） +##＊ +##， +##－ +##． +##／ +##： +##； +##＜ +##？ +##＠ +##［ +##＼ +##］ +##＾ +##＿ +##｀ +##ｆ +##ｈ +##ｊ +##ｕ +##ｗ +##ｚ +##｛ +##｝ +##｡ +##｢ +##｣ +##､ +##･ +##ｯ +##ｰ +##ｲ +##ｸ +##ｼ +##ｽ +##ﾄ +##ﾉ +##ﾌ +##ﾗ +##ﾙ +##ﾝ +##ﾞ +##ﾟ +##￣ +##￥ +##👍 +##🔥 +##😂 +##😎 diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/sources/pinyin_2_bpmf.txt b/src/YingMusicSinger/utils/f5_tts/g2p/sources/pinyin_2_bpmf.txt new file mode 100644 index 0000000000000000000000000000000000000000..af74dc687a547ed7822dacc77b7491924a8dcf1b --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/sources/pinyin_2_bpmf.txt @@ -0,0 +1,429 @@ +a ㄚ +ai ㄞ +an ㄢ +ang ㄤ +ao ㄠ +ba ㄅㄚ +bai ㄅㄞ +ban ㄅㄢ +bang ㄅㄤ +bao ㄅㄠ +bei ㄅㄟ +ben ㄅㄣ +beng ㄅㄥ +bi ㄅㄧ +bian ㄅㄧㄢ +biang ㄅㄧㄤ +biao ㄅㄧㄠ +bie ㄅㄧㄝ +bin ㄅㄧㄣ +bing ㄅㄧㄥ +bo ㄅㄛ +bu ㄅㄨ +ca ㄘㄚ +cai ㄘㄞ +can ㄘㄢ +cang ㄘㄤ +cao ㄘㄠ +ce ㄘㄜ +cen ㄘㄣ +ceng ㄘㄥ +cha ㄔㄚ +chai ㄔㄞ +chan ㄔㄢ +chang ㄔㄤ +chao ㄔㄠ +che ㄔㄜ +chen ㄔㄣ +cheng ㄔㄥ +chi ㄔ +chong ㄔㄨㄥ +chou ㄔㄡ +chu ㄔㄨ +chua ㄔㄨㄚ +chuai ㄔㄨㄞ +chuan ㄔㄨㄢ +chuang ㄔㄨㄤ +chui ㄔㄨㄟ +chun ㄔㄨㄣ +chuo ㄔㄨㄛ +ci ㄘ +cong ㄘㄨㄥ +cou ㄘㄡ +cu ㄘㄨ +cuan ㄘㄨㄢ +cui ㄘㄨㄟ +cun ㄘㄨㄣ +cuo ㄘㄨㄛ +da ㄉㄚ +dai ㄉㄞ +dan ㄉㄢ +dang ㄉㄤ +dao ㄉㄠ +de ㄉㄜ +dei ㄉㄟ +den ㄉㄣ +deng ㄉㄥ +di ㄉㄧ +dia ㄉㄧㄚ +dian ㄉㄧㄢ +diao ㄉㄧㄠ +die ㄉㄧㄝ +din ㄉㄧㄣ +ding ㄉㄧㄥ +diu ㄉㄧㄡ +dong ㄉㄨㄥ +dou ㄉㄡ +du ㄉㄨ +duan ㄉㄨㄢ +dui ㄉㄨㄟ +dun ㄉㄨㄣ +duo ㄉㄨㄛ +e ㄜ +ei ㄟ +en ㄣ +eng ㄥ +er ㄦ +fa ㄈㄚ +fan ㄈㄢ +fang ㄈㄤ +fei ㄈㄟ +fen ㄈㄣ +feng ㄈㄥ +fo ㄈㄛ +fou ㄈㄡ +fu ㄈㄨ +ga ㄍㄚ +gai ㄍㄞ +gan ㄍㄢ +gang ㄍㄤ +gao ㄍㄠ +ge ㄍㄜ +gei ㄍㄟ +gen ㄍㄣ +geng ㄍㄥ +gong ㄍㄨㄥ +gou ㄍㄡ +gu ㄍㄨ +gua ㄍㄨㄚ +guai ㄍㄨㄞ +guan ㄍㄨㄢ +guang ㄍㄨㄤ +gui ㄍㄨㄟ +gun ㄍㄨㄣ +guo ㄍㄨㄛ +ha ㄏㄚ +hai ㄏㄞ +han ㄏㄢ +hang ㄏㄤ +hao ㄏㄠ +he ㄏㄜ +hei ㄏㄟ +hen ㄏㄣ +heng ㄏㄥ +hm ㄏㄇ +hong ㄏㄨㄥ +hou ㄏㄡ +hu ㄏㄨ +hua ㄏㄨㄚ +huai ㄏㄨㄞ +huan ㄏㄨㄢ +huang ㄏㄨㄤ +hui ㄏㄨㄟ +hun ㄏㄨㄣ +huo ㄏㄨㄛ +ji ㄐㄧ +jia ㄐㄧㄚ +jian ㄐㄧㄢ +jiang ㄐㄧㄤ +jiao ㄐㄧㄠ +jie ㄐㄧㄝ +jin ㄐㄧㄣ +jing ㄐㄧㄥ +jiong ㄐㄩㄥ +jiu ㄐㄧㄡ +ju ㄐㄩ +jv ㄐㄩ +juan ㄐㄩㄢ +jvan ㄐㄩㄢ +jue ㄐㄩㄝ +jve ㄐㄩㄝ +jun ㄐㄩㄣ +ka ㄎㄚ +kai ㄎㄞ +kan ㄎㄢ +kang ㄎㄤ +kao ㄎㄠ +ke ㄎㄜ +kei ㄎㄟ +ken ㄎㄣ +keng ㄎㄥ +kong ㄎㄨㄥ +kou ㄎㄡ +ku ㄎㄨ +kua ㄎㄨㄚ +kuai ㄎㄨㄞ +kuan ㄎㄨㄢ +kuang ㄎㄨㄤ +kui ㄎㄨㄟ +kun ㄎㄨㄣ +kuo ㄎㄨㄛ +la ㄌㄚ +lai ㄌㄞ +lan ㄌㄢ +lang ㄌㄤ +lao ㄌㄠ +le ㄌㄜ +lei ㄌㄟ +leng ㄌㄥ +li ㄌㄧ +lia ㄌㄧㄚ +lian ㄌㄧㄢ +liang ㄌㄧㄤ +liao ㄌㄧㄠ +lie ㄌㄧㄝ +lin ㄌㄧㄣ +ling ㄌㄧㄥ +liu ㄌㄧㄡ +lo ㄌㄛ +long ㄌㄨㄥ +lou ㄌㄡ +lu ㄌㄨ +luan ㄌㄨㄢ +lue ㄌㄩㄝ +lun ㄌㄨㄣ +luo ㄌㄨㄛ +lv ㄌㄩ +lve ㄌㄩㄝ +m ㄇㄨ +ma ㄇㄚ +mai ㄇㄞ +man ㄇㄢ +mang ㄇㄤ +mao ㄇㄠ +me ㄇㄜ +mei ㄇㄟ +men ㄇㄣ +meng ㄇㄥ +mi ㄇㄧ +mian ㄇㄧㄢ +miao ㄇㄧㄠ +mie ㄇㄧㄝ +min ㄇㄧㄣ +ming ㄇㄧㄥ +miu ㄇㄧㄡ +mo ㄇㄛ +mou ㄇㄡ +mu ㄇㄨ +n ㄣ +na ㄋㄚ +nai ㄋㄞ +nan ㄋㄢ +nang ㄋㄤ +nao ㄋㄠ +ne ㄋㄜ +nei ㄋㄟ +nen ㄋㄣ +neng ㄋㄥ +ng ㄣ +ni ㄋㄧ +nian ㄋㄧㄢ +niang ㄋㄧㄤ +niao ㄋㄧㄠ +nie ㄋㄧㄝ +nin ㄋㄧㄣ +ning ㄋㄧㄥ +niu ㄋㄧㄡ +nong ㄋㄨㄥ +nou ㄋㄡ +nu ㄋㄨ +nuan ㄋㄨㄢ +nue ㄋㄩㄝ +nun ㄋㄨㄣ +nuo ㄋㄨㄛ +nv ㄋㄩ +nve ㄋㄩㄝ +o ㄛ +ou ㄡ +pa ㄆㄚ +pai ㄆㄞ +pan ㄆㄢ +pang ㄆㄤ +pao ㄆㄠ +pei ㄆㄟ +pen ㄆㄣ +peng ㄆㄥ +pi ㄆㄧ +pian ㄆㄧㄢ +piao ㄆㄧㄠ +pie ㄆㄧㄝ +pin ㄆㄧㄣ +ping ㄆㄧㄥ +po ㄆㄛ +pou ㄆㄡ +pu ㄆㄨ +qi ㄑㄧ +qia ㄑㄧㄚ +qian ㄑㄧㄢ +qiang ㄑㄧㄤ +qiao ㄑㄧㄠ +qie ㄑㄧㄝ +qin ㄑㄧㄣ +qing ㄑㄧㄥ +qiong ㄑㄩㄥ +qiu ㄑㄧㄡ +qu ㄑㄩ +quan ㄑㄩㄢ +qvan ㄑㄩㄢ +que ㄑㄩㄝ +qun ㄑㄩㄣ +ran ㄖㄢ +rang ㄖㄤ +rao ㄖㄠ +re ㄖㄜ +ren ㄖㄣ +reng ㄖㄥ +ri ㄖ +rong ㄖㄨㄥ +rou ㄖㄡ +ru ㄖㄨ +rua ㄖㄨㄚ +ruan ㄖㄨㄢ +rui ㄖㄨㄟ +run ㄖㄨㄣ +ruo ㄖㄨㄛ +sa ㄙㄚ +sai ㄙㄞ +san ㄙㄢ +sang ㄙㄤ +sao ㄙㄠ +se ㄙㄜ +sen ㄙㄣ +seng ㄙㄥ +sha ㄕㄚ +shai ㄕㄞ +shan ㄕㄢ +shang ㄕㄤ +shao ㄕㄠ +she ㄕㄜ +shei ㄕㄟ +shen ㄕㄣ +sheng ㄕㄥ +shi ㄕ +shou ㄕㄡ +shu ㄕㄨ +shua ㄕㄨㄚ +shuai ㄕㄨㄞ +shuan ㄕㄨㄢ +shuang ㄕㄨㄤ +shui ㄕㄨㄟ +shun ㄕㄨㄣ +shuo ㄕㄨㄛ +si ㄙ +song ㄙㄨㄥ +sou ㄙㄡ +su ㄙㄨ +suan ㄙㄨㄢ +sui ㄙㄨㄟ +sun ㄙㄨㄣ +suo ㄙㄨㄛ +ta ㄊㄚ +tai ㄊㄞ +tan ㄊㄢ +tang ㄊㄤ +tao ㄊㄠ +te ㄊㄜ +tei ㄊㄟ +teng ㄊㄥ +ti ㄊㄧ +tian ㄊㄧㄢ +tiao ㄊㄧㄠ +tie ㄊㄧㄝ +ting ㄊㄧㄥ +tong ㄊㄨㄥ +tou ㄊㄡ +tsuo ㄘㄨㄛ +tu ㄊㄨ +tuan ㄊㄨㄢ +tui ㄊㄨㄟ +tun ㄊㄨㄣ +tuo ㄊㄨㄛ +tzan ㄗㄢ +wa ㄨㄚ +wai ㄨㄞ +wan ㄨㄢ +wang ㄨㄤ +wei ㄨㄟ +wen ㄨㄣ +weng ㄨㄥ +wo ㄨㄛ +wong ㄨㄥ +wu ㄨ +xi ㄒㄧ +xia ㄒㄧㄚ +xian ㄒㄧㄢ +xiang ㄒㄧㄤ +xiao ㄒㄧㄠ +xie ㄒㄧㄝ +xin ㄒㄧㄣ +xing ㄒㄧㄥ +xiong ㄒㄩㄥ +xiu ㄒㄧㄡ +xu ㄒㄩ +xuan ㄒㄩㄢ +xue ㄒㄩㄝ +xun ㄒㄩㄣ +ya ㄧㄚ +yai ㄧㄞ +yan ㄧㄢ +yang ㄧㄤ +yao ㄧㄠ +ye ㄧㄝ +yi ㄧ +yin ㄧㄣ +ying ㄧㄥ +yo ㄧㄛ +yong ㄩㄥ +you ㄧㄡ +yu ㄩ +yuan ㄩㄢ +yue ㄩㄝ +yve ㄩㄝ +yun ㄩㄣ +za ㄗㄚ +zai ㄗㄞ +zan ㄗㄢ +zang ㄗㄤ +zao ㄗㄠ +ze ㄗㄜ +zei ㄗㄟ +zen ㄗㄣ +zeng ㄗㄥ +zha ㄓㄚ +zhai ㄓㄞ +zhan ㄓㄢ +zhang ㄓㄤ +zhao ㄓㄠ +zhe ㄓㄜ +zhei ㄓㄟ +zhen ㄓㄣ +zheng ㄓㄥ +zhi ㄓ +zhong ㄓㄨㄥ +zhou ㄓㄡ +zhu ㄓㄨ +zhua ㄓㄨㄚ +zhuai ㄓㄨㄞ +zhuan ㄓㄨㄢ +zhuang ㄓㄨㄤ +zhui ㄓㄨㄟ +zhun ㄓㄨㄣ +zhuo ㄓㄨㄛ +zi ㄗ +zong ㄗㄨㄥ +zou ㄗㄡ +zu ㄗㄨ +zuan ㄗㄨㄢ +zui ㄗㄨㄟ +zun ㄗㄨㄣ +zuo ㄗㄨㄛ diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/utils/front_utils.py b/src/YingMusicSinger/utils/f5_tts/g2p/utils/front_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7f0de1aca8d42995e20bc6d33dc1475b5e1ae150 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/utils/front_utils.py @@ -0,0 +1,18 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +def generate_poly_lexicon(file_path: str): + """Generate poly char lexicon for Mandarin Chinese.""" + poly_dict = {} + + with open(file_path, "r", encoding="utf-8") as readf: + txt_list = readf.readlines() + for txt in txt_list: + word = txt.strip("\n") + if word not in poly_dict: + poly_dict[word] = 1 + readf.close() + return poly_dict diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/utils/g2p.py b/src/YingMusicSinger/utils/f5_tts/g2p/utils/g2p.py new file mode 100644 index 0000000000000000000000000000000000000000..dbf19e6673ecb385f6b343aa2d2fc2acfedd55f8 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/utils/g2p.py @@ -0,0 +1,139 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import json +import os +from typing import List, Union + +from phonemizer.backend import EspeakBackend +from phonemizer.separator import Separator +from phonemizer.utils import list2str, str2list + +# separator=Separator(phone=' ', word=' _ ', syllable='|'), +separator = Separator(word=" _ ", syllable="|", phone=" ") + +phonemizer_zh = EspeakBackend( + "cmn", preserve_punctuation=False, with_stress=False, language_switch="remove-flags" +) +# phonemizer_zh.separator = separator + +phonemizer_en = EspeakBackend( + "en-us", + preserve_punctuation=False, + with_stress=False, + language_switch="remove-flags", +) +# phonemizer_en.separator = separator + +phonemizer_ja = EspeakBackend( + "ja", preserve_punctuation=False, with_stress=False, language_switch="remove-flags" +) +# phonemizer_ja.separator = separator + +phonemizer_ko = EspeakBackend( + "ko", preserve_punctuation=False, with_stress=False, language_switch="remove-flags" +) +# phonemizer_ko.separator = separator + +phonemizer_fr = EspeakBackend( + "fr-fr", + preserve_punctuation=False, + with_stress=False, + language_switch="remove-flags", +) +# phonemizer_fr.separator = separator + +phonemizer_de = EspeakBackend( + "de", preserve_punctuation=False, with_stress=False, language_switch="remove-flags" +) +# phonemizer_de.separator = separator + + +lang2backend = { + "zh": phonemizer_zh, + "ja": phonemizer_ja, + "en": phonemizer_en, + "fr": phonemizer_fr, + "ko": phonemizer_ko, + "de": phonemizer_de, +} + +with open("./src/YingMusicSinger/utils/f5_tts/g2p/utils/mls_en.json", "r") as f: + json_data = f.read() +token = json.loads(json_data) + + +def phonemizer_g2p(text, language): + langbackend = lang2backend[language] + phonemes = _phonemize( + langbackend, + text, + separator, + strip=True, + njobs=1, + prepend_text=False, + preserve_empty_lines=False, + ) + token_id = [] + if isinstance(phonemes, list): + for phone in phonemes: + phonemes_split = phone.split(" ") + token_id.append([token[p] for p in phonemes_split if p in token]) + else: + phonemes_split = phonemes.split(" ") + token_id = [token[p] for p in phonemes_split if p in token] + return phonemes, token_id + + +def _phonemize( # pylint: disable=too-many-arguments + backend, + text: Union[str, List[str]], + separator: Separator, + strip: bool, + njobs: int, + prepend_text: bool, + preserve_empty_lines: bool, +): + """Auxiliary function to phonemize() + + Does the phonemization and returns the phonemized text. Raises a + RuntimeError on error. + + """ + # remember the text type for output (either list or string) + text_type = type(text) + + # force the text as a list + text = [line.strip(os.linesep) for line in str2list(text)] + + # if preserving empty lines, note the index of each empty line + if preserve_empty_lines: + empty_lines = [n for n, line in enumerate(text) if not line.strip()] + + # ignore empty lines + text = [line for line in text if line.strip()] + + if text: + # phonemize the text + phonemized = backend.phonemize( + text, separator=separator, strip=strip, njobs=njobs + ) + else: + phonemized = [] + + # if preserving empty lines, reinsert them into text and phonemized lists + if preserve_empty_lines: + for i in empty_lines: # noqa + if prepend_text: + text.insert(i, "") + phonemized.insert(i, "") + + # at that point, the phonemized text is a list of str. Format it as + # expected by the parameters + if prepend_text: + return list(zip(text, phonemized)) + if text_type == str: + return list2str(phonemized) + return phonemized diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/utils/log.py b/src/YingMusicSinger/utils/f5_tts/g2p/utils/log.py new file mode 100644 index 0000000000000000000000000000000000000000..d10b887ef2e9292bd79c628e9ed7881c7a91bf52 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/utils/log.py @@ -0,0 +1,52 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +import functools +import logging + +__all__ = [ + "logger", +] + + +class Logger(object): + def __init__(self, name: str = None): + name = "PaddleSpeech" if not name else name + self.logger = logging.getLogger(name) + + log_config = { + "DEBUG": 10, + "INFO": 20, + "TRAIN": 21, + "EVAL": 22, + "WARNING": 30, + "ERROR": 40, + "CRITICAL": 50, + "EXCEPTION": 100, + } + for key, level in log_config.items(): + logging.addLevelName(level, key) + if key == "EXCEPTION": + self.__dict__[key.lower()] = self.logger.exception + else: + self.__dict__[key.lower()] = functools.partial(self.__call__, level) + + self.format = logging.Formatter( + fmt="[%(asctime)-15s] [%(levelname)8s] - %(message)s" + ) + + self.handler = logging.StreamHandler() + self.handler.setFormatter(self.format) + + self.logger.addHandler(self.handler) + self.logger.setLevel(logging.INFO) + self.logger.propagate = False + + def __call__(self, log_level: str, msg: str): + self.logger.log(log_level, msg) + + +logger = Logger() diff --git a/src/YingMusicSinger/utils/f5_tts/g2p/utils/mls_en.json b/src/YingMusicSinger/utils/f5_tts/g2p/utils/mls_en.json new file mode 100644 index 0000000000000000000000000000000000000000..f3aadbf144427af10ec06ca3cab8c4a2c461925d --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/g2p/utils/mls_en.json @@ -0,0 +1,335 @@ +{ + "[UNK]": 0, + "_": 1, + "b": 2, + "d": 3, + "f": 4, + "h": 5, + "i": 6, + "j": 7, + "k": 8, + "l": 9, + "m": 10, + "n": 11, + "p": 12, + "r": 13, + "s": 14, + "t": 15, + "v": 16, + "w": 17, + "x": 18, + "z": 19, + "æ": 20, + "ç": 21, + "ð": 22, + "ŋ": 23, + "ɐ": 24, + "ɔ": 25, + "ə": 26, + "ɚ": 27, + "ɛ": 28, + "ɡ": 29, + "ɪ": 30, + "ɬ": 31, + "ɹ": 32, + "ɾ": 33, + "ʃ": 34, + "ʊ": 35, + "ʌ": 36, + "ʒ": 37, + "ʔ": 38, + "θ": 39, + "ᵻ": 40, + "aɪ": 41, + "aʊ": 42, + "dʒ": 43, + "eɪ": 44, + "iə": 45, + "iː": 46, + "n̩": 47, + "oʊ": 48, + "oː": 49, + "tʃ": 50, + "uː": 51, + "ææ": 52, + "ɐɐ": 53, + "ɑː": 54, + "ɑ̃": 55, + "ɔɪ": 56, + "ɔː": 57, + "ɔ̃": 58, + "əl": 59, + "ɛɹ": 60, + "ɜː": 61, + "ɡʲ": 62, + "ɪɹ": 63, + "ʊɹ": 64, + "aɪə": 65, + "aɪɚ": 66, + "iːː": 67, + "oːɹ": 68, + "ɑːɹ": 69, + "ɔːɹ": 70, + + "1": 71, + "a": 72, + "e": 73, + "o": 74, + "q": 75, + "u": 76, + "y": 77, + "ɑ": 78, + "ɒ": 79, + "ɕ": 80, + "ɣ": 81, + "ɫ": 82, + "ɯ": 83, + "ʐ": 84, + "ʲ": 85, + "a1": 86, + "a2": 87, + "a5": 88, + "ai": 89, + "aɜ": 90, + "aː": 91, + "ei": 92, + "eə": 93, + "i.": 94, + "i1": 95, + "i2": 96, + "i5": 97, + "io": 98, + "iɑ": 99, + "iɛ": 100, + "iɜ": 101, + "i̪": 102, + "kh": 103, + "nʲ": 104, + "o1": 105, + "o2": 106, + "o5": 107, + "ou": 108, + "oɜ": 109, + "ph": 110, + "s.": 111, + "th": 112, + "ts": 113, + "tɕ": 114, + "u1": 115, + "u2": 116, + "u5": 117, + "ua": 118, + "uo": 119, + "uə": 120, + "uɜ": 121, + "y1": 122, + "y2": 123, + "y5": 124, + "yu": 125, + "yæ": 126, + "yə": 127, + "yɛ": 128, + "yɜ": 129, + "ŋɜ": 130, + "ŋʲ": 131, + "ɑ1": 132, + "ɑ2": 133, + "ɑ5": 134, + "ɑu": 135, + "ɑɜ": 136, + "ɑʲ": 137, + "ə1": 138, + "ə2": 139, + "ə5": 140, + "ər": 141, + "əɜ": 142, + "əʊ": 143, + "ʊə": 144, + "ai1": 145, + "ai2": 146, + "ai5": 147, + "aiɜ": 148, + "ei1": 149, + "ei2": 150, + "ei5": 151, + "eiɜ": 152, + "i.1": 153, + "i.2": 154, + "i.5": 155, + "i.ɜ": 156, + "io5": 157, + "iou": 158, + "iɑ1": 159, + "iɑ2": 160, + "iɑ5": 161, + "iɑɜ": 162, + "iɛ1": 163, + "iɛ2": 164, + "iɛ5": 165, + "iɛɜ": 166, + "i̪1": 167, + "i̪2": 168, + "i̪5": 169, + "i̪ɜ": 170, + "onɡ": 171, + "ou1": 172, + "ou2": 173, + "ou5": 174, + "ouɜ": 175, + "ts.": 176, + "tsh": 177, + "tɕh": 178, + "u5ʲ": 179, + "ua1": 180, + "ua2": 181, + "ua5": 182, + "uai": 183, + "uaɜ": 184, + "uei": 185, + "uo1": 186, + "uo2": 187, + "uo5": 188, + "uoɜ": 189, + "uə1": 190, + "uə2": 191, + "uə5": 192, + "uəɜ": 193, + "yiɜ": 194, + "yu2": 195, + "yu5": 196, + "yæ2": 197, + "yæ5": 198, + "yæɜ": 199, + "yə2": 200, + "yə5": 201, + "yəɜ": 202, + "yɛ1": 203, + "yɛ2": 204, + "yɛ5": 205, + "yɛɜ": 206, + "ɑu1": 207, + "ɑu2": 208, + "ɑu5": 209, + "ɑuɜ": 210, + "ər1": 211, + "ər2": 212, + "ər5": 213, + "ərɜ": 214, + "əː1": 215, + "iou1": 216, + "iou2": 217, + "iou5": 218, + "iouɜ": 219, + "onɡ1": 220, + "onɡ2": 221, + "onɡ5": 222, + "onɡɜ": 223, + "ts.h": 224, + "uai2": 225, + "uai5": 226, + "uaiɜ": 227, + "uei1": 228, + "uei2": 229, + "uei5": 230, + "ueiɜ": 231, + "uoɜʲ": 232, + "yɛ5ʲ": 233, + "ɑu2ʲ": 234, + + "2": 235, + "5": 236, + "ɜ": 237, + "ʂ": 238, + "dʑ": 239, + "iɪ": 240, + "uɪ": 241, + "xʲ": 242, + "ɑt": 243, + "ɛɜ": 244, + "ɛː": 245, + "ɪː": 246, + "phʲ": 247, + "ɑ5ʲ": 248, + "ɑuʲ": 249, + "ərə": 250, + "uozʰ": 251, + "ər1ʲ": 252, + "tɕhtɕh": 253, + + "c": 254, + "ʋ": 255, + "ʍ": 256, + "ʑ": 257, + "ː": 258, + "aə": 259, + "eː": 260, + "hʲ": 261, + "iʊ": 262, + "kʲ": 263, + "lʲ": 264, + "oə": 265, + "oɪ": 266, + "oʲ": 267, + "pʲ": 268, + "sʲ": 269, + "u4": 270, + "uʲ": 271, + "yi": 272, + "yʲ": 273, + "ŋ2": 274, + "ŋ5": 275, + "ŋ̩": 276, + "ɑɪ": 277, + "ɑʊ": 278, + "ɕʲ": 279, + "ət": 280, + "əə": 281, + "əɪ": 282, + "əʲ": 283, + "ɛ1": 284, + "ɛ5": 285, + "aiə": 286, + "aiɪ": 287, + "azʰ": 288, + "eiə": 289, + "eiɪ": 290, + "eiʊ": 291, + "i.ə": 292, + "i.ɪ": 293, + "i.ʊ": 294, + "ioɜ": 295, + "izʰ": 296, + "iɑə": 297, + "iɑʊ": 298, + "iɑʲ": 299, + "iɛə": 300, + "iɛɪ": 301, + "iɛʊ": 302, + "i̪ə": 303, + "i̪ʊ": 304, + "khʲ": 305, + "ouʲ": 306, + "tsʲ": 307, + "u2ʲ": 308, + "uoɪ": 309, + "uzʰ": 310, + "uɜʲ": 311, + "yæɪ": 312, + "yəʊ": 313, + "ərt": 314, + "ərɪ": 315, + "ərʲ": 316, + "əːt": 317, + "iouə": 318, + "iouʊ": 319, + "iouʲ": 320, + "iɛzʰ": 321, + "onɡə": 322, + "onɡɪ": 323, + "onɡʊ": 324, + "ouzʰ": 325, + "uai1": 326, + "ueiɪ": 327, + "ɑuzʰ": 328, + "iouzʰ": 329 +} \ No newline at end of file diff --git a/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/LangSegment.py b/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/LangSegment.py new file mode 100755 index 0000000000000000000000000000000000000000..b5d5730fba92cbe3b2be88b0a85f48e31f1f9e26 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/LangSegment.py @@ -0,0 +1,1251 @@ +""" +This file bundles language identification functions. + +Modifications (fork): Copyright (c) 2021, Adrien Barbaresi. + +Original code: Copyright (c) 2011 Marco Lui . +Based on research by Marco Lui and Tim Baldwin. + +See LICENSE file for more info. +https://github.com/adbar/py3langid + +Projects: +https://github.com/juntaosun/LangSegment +""" + +import re +from collections import Counter, defaultdict + +import numpy as np + +# import langid +# import py3langid as langid +# pip install py3langid==0.2.2 +# 启用语言预测概率归一化，概率预测的分数。因此，实现重新规范化 产生 0-1 范围内的输出。 +# langid disables probability normalization by default. For command-line usages of , it can be enabled by passing the flag. +# For probability normalization in library use, the user must instantiate their own . An example of such usage is as follows: +from py3langid.langid import MODEL_FILE, LanguageIdentifier + +langid = LanguageIdentifier.from_pickled_model(MODEL_FILE, norm_probs=True) + +# Digital processing +try: + from src.YingMusicSinger.utils.f5_tts.thirdparty.LangSegment.utils.num import ( + num2str, + ) +except ImportError: + try: + from thirdparty.LangSegment.utils.num import num2str + except ImportError as e: + raise e + +# ----------------------------------- +# 更新日志：新版本分词更加精准。 +# Changelog: The new version of the word segmentation is more accurate. +# チェンジログ:新しいバージョンの単語セグメンテーションはより正確です。 +# Changelog: 분할이라는 단어의 새로운 버전이 더 정확합니다. +# ----------------------------------- + + +# Word segmentation function: +# automatically identify and split the words (Chinese/English/Japanese/Korean) in the article or sentence according to different languages, +# making it more suitable for TTS processing. +# This code is designed for front-end text multi-lingual mixed annotation distinction, multi-language mixed training and inference of various TTS projects. +# This processing result is mainly for (Chinese = zh, Japanese = ja, English = en, Korean = ko), and can actually support up to 97 different language mixing processing. + +# =========================================================================================================== +# 分かち書き機能:文章や文章の中の例えば（中国語/英語/日本語/韓国語）を、異なる言語で自動的に認識して分割し、TTS処理により適したものにします。 +# このコードは、さまざまなTTSプロジェクトのフロントエンドテキストの多言語混合注釈区別、多言語混合トレーニング、および推論のために特別に作成されています。 +# =========================================================================================================== +# (1)自動分詞:「韓国語では何を読むのですかあなたの体育の先生は誰ですか?今回の発表会では、iPhone 15シリーズの4機種が登場しました」 +# （2）手动分词:“あなたの名前は佐々木ですか?ですか?” +# この処理結果は主に（中国語=ja、日本語=ja、英語=en、韓国語=ko）を対象としており、実際には最大97の異なる言語の混合処理をサポートできます。 +# =========================================================================================================== + +# =========================================================================================================== +# 단어 분할 기능: 기사 또는 문장에서 단어(중국어/영어/일본어/한국어)를 다른 언어에 따라 자동으로 식별하고 분할하여 TTS 처리에 더 적합합니다. +# 이 코드는 프런트 엔드 텍스트 다국어 혼합 주석 분화, 다국어 혼합 교육 및 다양한 TTS 프로젝트의 추론을 위해 설계되었습니다. +# =========================================================================================================== +# (1) 자동 단어 분할: "한국어로 무엇을 읽습니까? 스포츠 씨? 이 컨퍼런스는 4개의 iPhone 15 시리즈 모델을 제공합니다." +# (2) 수동 참여: "이름이 Saki입니까? ?" +# 이 처리 결과는 주로 (중국어 = zh, 일본어 = ja, 영어 = en, 한국어 = ko)를 위한 것이며 실제로 혼합 처리를 위해 최대 97개의 언어를 지원합니다. +# =========================================================================================================== + +# =========================================================================================================== +# 分词功能：将文章或句子里的例如（中/英/日/韩），按不同语言自动识别并拆分，让它更适合TTS处理。 +# 本代码专为各种 TTS 项目的前端文本多语种混合标注区分，多语言混合训练和推理而编写。 +# =========================================================================================================== +# （1）自动分词：“韩语中的오빠读什么呢？あなたの体育の先生は誰ですか? 此次发布会带来了四款iPhone 15系列机型” +# （2）手动分词：“你的名字叫佐々木？吗？” +# 本处理结果主要针对（中文=zh , 日文=ja , 英文=en , 韩语=ko）, 实际上可支持多达 97 种不同的语言混合处理。 +# =========================================================================================================== + + +# 手动分词标签规范：<语言标签>文本内容 +# 수동 단어 분할 태그 사양: <언어 태그> 텍스트 내용 +# Manual word segmentation tag specification: text content +# 手動分詞タグ仕様:<言語タグ>テキスト内容 +# =========================================================================================================== +# For manual word segmentation, labels need to appear in pairs, such as: +# 如需手动分词，标签需要成对出现，例如：“佐々木” 或者 “佐々木” +# 错误示范：“你的名字叫佐々木。” 此句子中出现的单个标签将被忽略，不会处理。 +# Error demonstration: "Your name is 佐々木。" Single tags that appear in this sentence will be ignored and will not be processed. +# =========================================================================================================== + + +# =========================================================================================================== +# 语音合成标记语言 SSML , 这里只支持它的标签（非 XML）Speech Synthesis Markup Language SSML, only its tags are supported here (not XML) +# 想支持更多的 SSML 标签？欢迎 PR！ Want to support more SSML tags? PRs are welcome! +# 说明：除了中文以外，它也可改造成支持多语种 SSML ，不仅仅是中文。 +# Note: In addition to Chinese, it can also be modified to support multi-language SSML, not just Chinese. +# =========================================================================================================== +# 中文实现：Chinese implementation: +# 【SSML】=中文大写数字读法（单字） +# 【SSML】=数字转成中文电话号码大写汉字（单字） +# 【SSML】=按金额发音。 +# 【SSML】=按日期发音。支持 2024年08月24, 2024/8/24, 2024-08, 08-24, 24 等输入。 +# =========================================================================================================== +class LangSSML: + # 纯数字 + _zh_numerals_number = { + "0": "零", + "1": "一", + "2": "二", + "3": "三", + "4": "四", + "5": "五", + "6": "六", + "7": "七", + "8": "八", + "9": "九", + } + + # 将2024/8/24, 2024-08, 08-24, 24 标准化“年月日” + # Standardize 2024/8/24, 2024-08, 08-24, 24 to "year-month-day" + def _format_chinese_data(date_str: str): + # 处理日期格式 + input_date = date_str + if date_str is None or date_str.strip() == "": + return "" + date_str = re.sub(r"[\/\._|年|月]", "-", date_str) + date_str = re.sub(r"日", r"", date_str) + date_arrs = date_str.split(" ") + if len(date_arrs) == 1 and ":" in date_arrs[0]: + time_str = date_arrs[0] + date_arrs = [] + else: + time_str = date_arrs[1] if len(date_arrs) >= 2 else "" + + def nonZero(num, cn, func=None): + if func is not None: + num = func(num) + return f"{num}{cn}" if num is not None and num != "" and num != "0" else "" + + f_number = LangSSML.to_chinese_number + f_currency = LangSSML.to_chinese_currency + # year, month, day + year_month_day = "" + if len(date_arrs) > 0: + year, month, day = "", "", "" + parts = date_arrs[0].split("-") + if len(parts) == 3: # 格式为 YYYY-MM-DD + year, month, day = parts + elif len(parts) == 2: # 格式为 MM-DD 或 YYYY-MM + if len(parts[0]) == 4: # 年-月 + year, month = parts + else: + month, day = parts # 月-日 + elif len(parts[0]) > 0: # 仅有月-日或年 + if len(parts[0]) == 4: + year = parts[0] + else: + day = parts[0] + year, month, day = ( + nonZero(year, "年", f_number), + nonZero(month, "月", f_currency), + nonZero(day, "日", f_currency), + ) + year_month_day = re.sub(r"([年|月|日])+", r"\1", f"{year}{month}{day}") + # hours, minutes, seconds + time_str = re.sub(r"[\/\.\-：_]", ":", time_str) + time_arrs = time_str.split(":") + hours, minutes, seconds = "", "", "" + if len(time_arrs) == 3: # H/M/S + hours, minutes, seconds = time_arrs + elif len(time_arrs) == 2: # H/M + hours, minutes = time_arrs + elif len(time_arrs[0]) > 0: + hours = f"{time_arrs[0]}点" # H + if len(time_arrs) > 1: + hours, minutes, seconds = ( + nonZero(hours, "点", f_currency), + nonZero(minutes, "分", f_currency), + nonZero(seconds, "秒", f_currency), + ) + hours_minutes_seconds = re.sub( + r"([点|分|秒])+", r"\1", f"{hours}{minutes}{seconds}" + ) + output_date = f"{year_month_day}{hours_minutes_seconds}" + return output_date + + # 【SSML】number=中文大写数字读法（单字） + # Chinese Numbers(single word) + def to_chinese_number(num: str): + pattern = r"(\d+)" + zh_numerals = LangSSML._zh_numerals_number + arrs = re.split(pattern, num) + output = "" + for item in arrs: + if re.match(pattern, item): + output += "".join( + zh_numerals[digit] if digit in zh_numerals else "" + for digit in str(item) + ) + else: + output += item + output = output.replace(".", "点") + return output + + # 【SSML】telephone=数字转成中文电话号码大写汉字（单字） + # Convert numbers to Chinese phone numbers in uppercase Chinese characters(single word) + def to_chinese_telephone(num: str): + output = LangSSML.to_chinese_number(num.replace("+86", "")) # zh +86 + output = output.replace("一", "幺") + return output + + # 【SSML】currency=按金额发音。 + # Digital processing from GPT_SoVITS num.py （thanks） + def to_chinese_currency(num: str): + pattern = r"(\d+)" + arrs = re.split(pattern, num) + output = "" + for item in arrs: + if re.match(pattern, item): + output += num2str(item) + else: + output += item + output = output.replace(".", "点") + return output + + # 【SSML】date=按日期发音。支持 2024年08月24, 2024/8/24, 2024-08, 08-24, 24 等输入。 + def to_chinese_date(num: str): + chinese_date = LangSSML._format_chinese_data(num) + return chinese_date + + +class LangSegment: + _text_cache = None + _text_lasts = None + _text_langs = None + _lang_count = None + _lang_eos = None + + # 可自定义语言匹配标签：カスタマイズ可能な言語対応タグ:사용자 지정 가능한 언어 일치 태그: + # Customizable language matching tags: These are supported，이 표현들은 모두 지지합니다 + # 你好 , 佐々木 , OK , 오빠 这些写法均支持 + SYMBOLS_PATTERN = r"(<([a-zA-Z|-]*)>(.*?)<\/*[a-zA-Z|-]*>)" + + # 语言过滤组功能, 可以指定保留语言。不在过滤组中的语言将被清除。您可随心搭配TTS语音合成所支持的语言。 + # 언어 필터 그룹 기능을 사용하면 예약된 언어를 지정할 수 있습니다. 필터 그룹에 없는 언어는 지워집니다. TTS 텍스트에서 지원하는 언어를 원하는 대로 일치시킬 수 있습니다. + # 言語フィルターグループ機能では、予約言語を指定できます。フィルターグループに含まれていない言語はクリアされます。TTS音声合成がサポートする言語を自由に組み合わせることができます。 + # The language filter group function allows you to specify reserved languages. + # Languages not in the filter group will be cleared. You can match the languages supported by TTS Text To Speech as you like. + # 排名越前，优先级越高，The higher the ranking, the higher the priority，ランキングが上位になるほど、優先度が高くなります。 + + # 系统默认过滤器。System default filter。(ISO 639-1 codes given) + # ---------------------------------------------------------------------------------------------------------------------------------- + # "zh"中文=Chinese ,"en"英语=English ,"ja"日语=Japanese ,"ko"韩语=Korean ,"fr"法语=French ,"vi"越南语=Vietnamese , "ru"俄语=Russian + # "th"泰语=Thai + # ---------------------------------------------------------------------------------------------------------------------------------- + DEFAULT_FILTERS = ["zh", "ja", "ko", "en"] + + # 用户可自定义过滤器。User-defined filters + Langfilters = DEFAULT_FILTERS[:] # 创建副本 + + # 合并文本 + isLangMerge = True + + # 试验性支持：您可自定义添加："fr"法语 , "vi"越南语。Experimental: You can customize to add: "fr" French, "vi" Vietnamese. + # 请使用API启用：LangSegment.setfilters(["zh", "en", "ja", "ko", "fr", "vi" , "ru" , "th"]) # 您可自定义添加，如："fr"法语 , "vi"越南语。 + + # 预览版功能，自动启用或禁用，无需设置 + # Preview feature, automatically enabled or disabled, no settings required + EnablePreview = False + + # 除此以外，它支持简写过滤器，只需按不同语种任意组合即可。 + # In addition to that, it supports abbreviation filters, allowing for any combination of different languages. + # 示例：您可以任意指定多种组合，进行过滤 + # Example: You can specify any combination to filter + + # 中/日语言优先级阀值（评分范围为 0 ~ 1）:评分低于设定阀值 <0.89 时，启用 filters 中的优先级。\n + # 중/일본어 우선 순위 임계값(점수 범위 0-1): 점수가 설정된 임계값 <0.89보다 낮을 때 필터에서 우선 순위를 활성화합니다. + # 中国語/日本語の優先度しきい値（スコア範囲0〜1）:スコアが設定されたしきい値<0.89未満の場合、フィルターの優先度が有効になります。\n + # Chinese and Japanese language priority threshold (score range is 0 ~ 1): The default threshold is 0.89. \n + # Only the common characters between Chinese and Japanese are processed with confidence and priority. \n + LangPriorityThreshold = 0.89 + + # Langfilters = ["zh"] # 按中文识别 + # Langfilters = ["en"] # 按英文识别 + # Langfilters = ["ja"] # 按日文识别 + # Langfilters = ["ko"] # 按韩文识别 + # Langfilters = ["zh_ja"] # 中日混合识别 + # Langfilters = ["zh_en"] # 中英混合识别 + # Langfilters = ["ja_en"] # 日英混合识别 + # Langfilters = ["zh_ko"] # 中韩混合识别 + # Langfilters = ["ja_ko"] # 日韩混合识别 + # Langfilters = ["en_ko"] # 英韩混合识别 + # Langfilters = ["zh_ja_en"] # 中日英混合识别 + # Langfilters = ["zh_ja_en_ko"] # 中日英韩混合识别 + + # 更多过滤组合，请您随意。。。For more filter combinations, please feel free to...... + # より多くのフィルターの組み合わせ、お気軽に。。。더 많은 필터 조합을 원하시면 자유롭게 해주세요. ..... + + # 可选保留：支持中文数字拼音格式，更方便前端实现拼音音素修改和推理，默认关闭 False 。 + # 开启后 True ，括号内的数字拼音格式均保留，并识别输出为："zh"中文。 + keepPinyin = False + + # DEFINITION + PARSE_TAG = re.compile(r"(⑥\$*\d+[\d]{6,}⑥)") + + @staticmethod + def _clears(): + LangSegment._text_cache = None + LangSegment._text_lasts = None + LangSegment._text_langs = None + LangSegment._text_waits = None + LangSegment._lang_count = None + LangSegment._lang_eos = None + pass + + @staticmethod + def _is_english_word(word): + return bool(re.match(r"^[a-zA-Z]+$", word)) + + @staticmethod + def _is_chinese(word): + for char in word: + if "\u4e00" <= char <= "\u9fff": + return True + return False + + @staticmethod + def _is_japanese_kana(word): + pattern = re.compile(r"[\u3040-\u309F\u30A0-\u30FF]+") + matches = pattern.findall(word) + return len(matches) > 0 + + @staticmethod + def _insert_english_uppercase(word): + modified_text = re.sub(r"(? 0 else None + if symbol is not None: + pass + elif preData is not None and preData["symbol"] is None: + if len(clear_text) == 0: + language = preData["lang"] + elif is_number == True: + language = preData["lang"] + _, pre_is_number = LangSegment._clear_text_number(preData["text"]) + if preData["lang"] == language: + LangSegment._statistics(preData["lang"], text) + text = preData["text"] + text + preData["text"] = text + return preData + elif pre_is_number == True: + text = f"{preData['text']}{text}" + words.pop() + elif is_number == True: + priority_language = LangSegment._get_filters_string()[:2] + if priority_language in "ja-zh-en-ko-fr-vi": + language = priority_language + data = {"lang": language, "text": text, "score": score, "symbol": symbol} + filters = LangSegment.Langfilters + if ( + filters is None + or len(filters) == 0 + or "?" in language + or language in filters + or language in filters[0] + or filters[0] == "*" + or filters[0] in "alls-mixs-autos" + ): + words.append(data) + LangSegment._statistics(data["lang"], data["text"]) + return data + + @staticmethod + def _addwords(words, language, text, score, symbol=None): + if text == "\n": + pass # Keep Line Breaks + elif text is None or len(text.strip()) == 0: + return True + if language is None: + language = "" + language = language.lower() + if language == "en": + text = LangSegment._insert_english_uppercase(text) + # text = re.sub(r'[(（）)]', ',' , text) # Keep it. + text_waits = LangSegment._text_waits + ispre_waits = len(text_waits) > 0 + preResult = text_waits.pop() if ispre_waits else None + if preResult is None: + preResult = words[-1] if len(words) > 0 else None + if preResult and ("|" in preResult["lang"]): + pre_lang = preResult["lang"] + if language in pre_lang: + preResult["lang"] = language = language.split("|")[0] + else: + preResult["lang"] = pre_lang.split("|")[0] + if ispre_waits: + preResult = LangSegment._saveData( + words, + preResult["lang"], + preResult["text"], + preResult["score"], + preResult["symbol"], + ) + pre_lang = preResult["lang"] if preResult else None + if ("|" in language) and ( + pre_lang and pre_lang not in language and "…" not in language + ): + language = language.split("|")[0] + if "|" in language: + LangSegment._text_waits.append( + {"lang": language, "text": text, "score": score, "symbol": symbol} + ) + else: + LangSegment._saveData(words, language, text, score, symbol) + return False + + @staticmethod + def _get_prev_data(words): + data = words[-1] if words and len(words) > 0 else None + if data: + return (data["lang"], data["text"]) + return (None, "") + + @staticmethod + def _match_ending(input, index): + if input is None or len(input) == 0: + return False, None + input = re.sub(r"\s+", "", input) + if len(input) == 0 or abs(index) > len(input): + return False, None + ending_pattern = re.compile(r'([「」“”‘’"\':：。.！!?．？])') + return ending_pattern.match(input[index]), input[index] + + @staticmethod + def _cleans_text(cleans_text): + cleans_text = re.sub(r"(.*?)([^\w]+)", r"\1 ", cleans_text) + cleans_text = re.sub(r"(.)\1+", r"\1", cleans_text) + return cleans_text.strip() + + @staticmethod + def _mean_processing(text: str): + if text is None or (text.strip()) == "": + return None, 0.0 + arrs = LangSegment._split_camel_case(text).split(" ") + langs = [] + for t in arrs: + if len(t.strip()) <= 3: + continue + language, score = langid.classify(t) + langs.append({"lang": language}) + if len(langs) == 0: + return None, 0.0 + return Counter([item["lang"] for item in langs]).most_common(1)[0][0], 1.0 + + @staticmethod + def _lang_classify(cleans_text): + language, score = langid.classify(cleans_text) + # fix: Huggingface is np.float32 + if ( + score is not None + and isinstance(score, np.generic) + and hasattr(score, "item") + ): + score = score.item() + score = round(score, 3) + return language, score + + @staticmethod + def _get_filters_string(): + filters = LangSegment.Langfilters + return "-".join(filters).lower().strip() if filters is not None else "" + + @staticmethod + def _parse_language(words, segment): + LANG_JA = "ja" + LANG_ZH = "zh" + LANG_ZH_JA = f"{LANG_ZH}|{LANG_JA}" + LANG_JA_ZH = f"{LANG_JA}|{LANG_ZH}" + language = LANG_ZH + regex_pattern = re.compile(r"([^\w\s]+)") + lines = regex_pattern.split(segment) + lines_max = len(lines) + LANG_EOS = LangSegment._lang_eos + for index, text in enumerate(lines): + if len(text) == 0: + continue + EOS = index >= (lines_max - 1) + nextId = index + 1 + nextText = lines[nextId] if not EOS else "" + nextPunc = ( + len(re.sub(regex_pattern, "", re.sub(r"\n+", "", nextText)).strip()) + == 0 + ) + textPunc = ( + len(re.sub(regex_pattern, "", re.sub(r"\n+", "", text)).strip()) == 0 + ) + if not EOS and ( + textPunc == True or (len(nextText.strip()) >= 0 and nextPunc == True) + ): + lines[nextId] = f"{text}{nextText}" + continue + number_tags = re.compile(r"(⑥\d{6,}⑥)") + cleans_text = re.sub(number_tags, "", text) + cleans_text = re.sub(r"\d+", "", cleans_text) + cleans_text = LangSegment._cleans_text(cleans_text) + # fix:Langid's recognition of short sentences is inaccurate, and it is spliced longer. + if not EOS and len(cleans_text) <= 2: + lines[nextId] = f"{text}{nextText}" + continue + language, score = LangSegment._lang_classify(cleans_text) + prev_language, prev_text = LangSegment._get_prev_data(words) + if language != LANG_ZH and all( + "\u4e00" <= c <= "\u9fff" for c in re.sub(r"\s", "", cleans_text) + ): + language, score = LANG_ZH, 1 + if len(cleans_text) <= 5 and LangSegment._is_chinese(cleans_text): + filters_string = LangSegment._get_filters_string() + if ( + score < LangSegment.LangPriorityThreshold + and len(filters_string) > 0 + ): + index_ja, index_zh = ( + filters_string.find(LANG_JA), + filters_string.find(LANG_ZH), + ) + if index_ja != -1 and index_ja < index_zh: + language = LANG_JA + elif index_zh != -1 and index_zh < index_ja: + language = LANG_ZH + if LangSegment._is_japanese_kana(cleans_text): + language = LANG_JA + elif len(cleans_text) > 2 and score > 0.90: + pass + elif EOS and LANG_EOS: + language = LANG_ZH if len(cleans_text) <= 1 else language + else: + LANG_UNKNOWN = ( + LANG_ZH_JA + if language == LANG_ZH + or (len(cleans_text) <= 2 and prev_language == LANG_ZH) + else LANG_JA_ZH + ) + match_end, match_char = LangSegment._match_ending(text, -1) + referen = ( + prev_language in LANG_UNKNOWN or LANG_UNKNOWN in prev_language + if prev_language + else False + ) + if match_char in "。.": + language = ( + prev_language if referen and len(words) > 0 else language + ) + else: + language = f"{LANG_UNKNOWN}|…" + text, *_ = re.subn(number_tags, LangSegment._restore_number, text) + LangSegment._addwords(words, language, text, score) + pass + pass + + # ---------------------------------------------------------- + # 【SSML】中文数字处理：Chinese Number Processing (SSML support) + # 这里默认都是中文，用于处理 SSML 中文标签。当然可以支持任意语言，例如： + # The default here is Chinese, which is used to process SSML Chinese tags. Of course, any language can be supported, for example: + # 中文电话号码：1234567 + # 中文数字号码：1234567 + @staticmethod + def _process_symbol_SSML(words, data): + tag, match = data + language = SSML = match[1] + text = match[2] + score = 1.0 + if SSML == "telephone": + # 中文-电话号码 + language = "zh" + text = LangSSML.to_chinese_telephone(text) + pass + elif SSML == "number": + # 中文-数字读法 + language = "zh" + text = LangSSML.to_chinese_number(text) + pass + elif SSML == "currency": + # 中文-按金额发音 + language = "zh" + text = LangSSML.to_chinese_currency(text) + pass + elif SSML == "date": + # 中文-按金额发音 + language = "zh" + text = LangSSML.to_chinese_date(text) + pass + LangSegment._addwords(words, language, text, score, SSML) + pass + + # ---------------------------------------------------------- + + @staticmethod + def _restore_number(matche): + value = matche.group(0) + text_cache = LangSegment._text_cache + if value in text_cache: + process, data = text_cache[value] + tag, match = data + value = match + return value + + @staticmethod + def _pattern_symbols(item, text): + if text is None: + return text + tag, pattern, process = item + matches = pattern.findall(text) + if len(matches) == 1 and "".join(matches[0]) == text: + return text + for i, match in enumerate(matches): + key = f"⑥{tag}{i:06d}⑥" + text = re.sub(pattern, key, text, count=1) + LangSegment._text_cache[key] = (process, (tag, match)) + return text + + @staticmethod + def _process_symbol(words, data): + tag, match = data + language = match[1] + text = match[2] + score = 1.0 + filters = LangSegment._get_filters_string() + if language not in filters: + LangSegment._process_symbol_SSML(words, data) + else: + LangSegment._addwords(words, language, text, score, True) + pass + + @staticmethod + def _process_english(words, data): + tag, match = data + text = match[0] + filters = LangSegment._get_filters_string() + priority_language = filters[:2] + # Preview feature, other language segmentation processing + enablePreview = LangSegment.EnablePreview + if enablePreview == True: + # Experimental: Other language support + regex_pattern = re.compile(r"(.*?[。.?？!！]+[\n]{,1})") + lines = regex_pattern.split(text) + for index, text in enumerate(lines): + if len(text.strip()) == 0: + continue + cleans_text = LangSegment._cleans_text(text) + language, score = LangSegment._lang_classify(cleans_text) + if language not in filters: + language, score = LangSegment._mean_processing(cleans_text) + if language is None or score <= 0.0: + continue + elif language in filters: + pass # pass + elif score >= 0.95: + continue # High score, but not in the filter, excluded. + elif score <= 0.15 and filters[:2] == "fr": + language = priority_language + else: + language = "en" + LangSegment._addwords(words, language, text, score) + else: + # Default is English + language, score = "en", 1.0 + LangSegment._addwords(words, language, text, score) + pass + + @staticmethod + def _process_Russian(words, data): + tag, match = data + text = match[0] + language = "ru" + score = 1.0 + LangSegment._addwords(words, language, text, score) + pass + + @staticmethod + def _process_Thai(words, data): + tag, match = data + text = match[0] + language = "th" + score = 1.0 + LangSegment._addwords(words, language, text, score) + pass + + @staticmethod + def _process_korean(words, data): + tag, match = data + text = match[0] + language = "ko" + score = 1.0 + LangSegment._addwords(words, language, text, score) + pass + + @staticmethod + def _process_quotes(words, data): + tag, match = data + text = "".join(match) + childs = LangSegment.PARSE_TAG.findall(text) + if len(childs) > 0: + LangSegment._process_tags(words, text, False) + else: + cleans_text = LangSegment._cleans_text(match[1]) + if len(cleans_text) <= 5: + LangSegment._parse_language(words, text) + else: + language, score = LangSegment._lang_classify(cleans_text) + LangSegment._addwords(words, language, text, score) + pass + + @staticmethod + def _process_pinyin(words, data): + tag, match = data + text = match + language = "zh" + score = 1.0 + LangSegment._addwords(words, language, text, score) + pass + + @staticmethod + def _process_number(words, data): # "$0" process only + """ + Numbers alone cannot accurately identify language. + Because numbers are universal in all languages. + So it won't be executed here, just for testing. + """ + tag, match = data + language = words[0]["lang"] if len(words) > 0 else "zh" + text = match + score = 0.0 + LangSegment._addwords(words, language, text, score) + pass + + @staticmethod + def _process_tags(words, text, root_tag): + text_cache = LangSegment._text_cache + segments = re.split(LangSegment.PARSE_TAG, text) + segments_len = len(segments) - 1 + for index, text in enumerate(segments): + if root_tag: + LangSegment._lang_eos = index >= segments_len + if LangSegment.PARSE_TAG.match(text): + process, data = text_cache[text] + if process: + process(words, data) + else: + LangSegment._parse_language(words, text) + pass + return words + + @staticmethod + def _merge_results(words): + new_word = [] + for index, cur_data in enumerate(words): + if "symbol" in cur_data: + del cur_data["symbol"] + if index == 0: + new_word.append(cur_data) + else: + pre_data = new_word[-1] + if cur_data["lang"] == pre_data["lang"]: + pre_data["text"] = f"{pre_data['text']}{cur_data['text']}" + else: + new_word.append(cur_data) + return new_word + + @staticmethod + def _parse_symbols(text): + TAG_NUM = "00" # "00" => default channels , "$0" => testing channel + TAG_S1, TAG_S2, TAG_P1, TAG_P2, TAG_EN, TAG_KO, TAG_RU, TAG_TH = ( + "$1", + "$2", + "$3", + "$4", + "$5", + "$6", + "$7", + "$8", + ) + TAG_BASE = re.compile(r'(([【《（(“‘"\']*[LANGUAGE]+[\W\s]*)+)') + # Get custom language filter + filters = LangSegment.Langfilters + filters = filters if filters is not None else "" + # ======================================================================================================= + # Experimental: Other language support.Thử nghiệm: Hỗ trợ ngôn ngữ khác.Expérimental : prise en charge d’autres langues. + # 相关语言字符如有缺失，熟悉相关语言的朋友，可以提交把缺失的发音符号补全。 + # If relevant language characters are missing, friends who are familiar with the relevant languages can submit a submission to complete the missing pronunciation symbols. + # S'il manque des caractères linguistiques pertinents, les amis qui connaissent les langues concernées peuvent soumettre une soumission pour compléter les symboles de prononciation manquants. + # Nếu thiếu ký tự ngôn ngữ liên quan, những người bạn quen thuộc với ngôn ngữ liên quan có thể gửi bài để hoàn thành các ký hiệu phát âm còn thiếu. + # ------------------------------------------------------------------------------------------------------- + # Preview feature, other language support + enablePreview = LangSegment.EnablePreview + if "fr" in filters or "vi" in filters: + enablePreview = True + LangSegment.EnablePreview = enablePreview + # 实验性：法语字符支持。Prise en charge des caractères français + RE_FR = "" if not enablePreview else "àáâãäåæçèéêëìíîïðñòóôõöùúûüýþÿ" + # 实验性：越南语字符支持。Hỗ trợ ký tự tiếng Việt + RE_VI = ( + "" + if not enablePreview + else "đơưăáàảãạắằẳẵặấầẩẫậéèẻẽẹếềểễệíìỉĩịóòỏõọốồổỗộớờởỡợúùủũụứừửữựôâêơưỷỹ" + ) + # ------------------------------------------------------------------------------------------------------- + # Basic options: + process_list = [ + ( + TAG_S1, + re.compile(LangSegment.SYMBOLS_PATTERN), + LangSegment._process_symbol, + ), # Symbol Tag + ( + TAG_KO, + re.compile(re.sub(r"LANGUAGE", "\uac00-\ud7a3", TAG_BASE.pattern)), + LangSegment._process_korean, + ), # Korean words + ( + TAG_TH, + re.compile(re.sub(r"LANGUAGE", "\u0e00-\u0e7f", TAG_BASE.pattern)), + LangSegment._process_Thai, + ), # Thai words support. + ( + TAG_RU, + re.compile(re.sub(r"LANGUAGE", "А-Яа-яЁё", TAG_BASE.pattern)), + LangSegment._process_Russian, + ), # Russian words support. + ( + TAG_NUM, + re.compile(r"(\W*\d+\W+\d*\W*\d*)"), + LangSegment._process_number, + ), # Number words, Universal in all languages, Ignore it. + ( + TAG_EN, + re.compile( + re.sub(r"LANGUAGE", f"a-zA-Z{RE_FR}{RE_VI}", TAG_BASE.pattern) + ), + LangSegment._process_english, + ), # English words + Other language support. + ( + TAG_P1, + re.compile(r'(["\'])(.*?)(\1)'), + LangSegment._process_quotes, + ), # Regular quotes + ( + TAG_P2, + re.compile( + r"([\n]*[【《（(“‘])([^【《（(“‘’”)）》】]{3,})([’”)）》】][\W\s]*[\n]{,1})" + ), + LangSegment._process_quotes, + ), # Special quotes, There are left and right. + ] + # Extended options: Default False + if LangSegment.keepPinyin == True: + process_list.insert( + 1, + ( + TAG_S2, + re.compile(r"([\(（{](?:\s*\w*\d\w*\s*)+[}）\)])"), + LangSegment._process_pinyin, + ), # Chinese Pinyin Tag. + ) + # ------------------------------------------------------------------------------------------------------- + words = [] + lines = re.findall(r".*\n*", re.sub(LangSegment.PARSE_TAG, "", text)) + for index, text in enumerate(lines): + if len(text.strip()) == 0: + continue + LangSegment._lang_eos = False + LangSegment._text_cache = {} + for item in process_list: + text = LangSegment._pattern_symbols(item, text) + cur_word = LangSegment._process_tags([], text, True) + if len(cur_word) == 0: + continue + cur_data = cur_word[0] if len(cur_word) > 0 else None + pre_data = words[-1] if len(words) > 0 else None + if ( + cur_data + and pre_data + and cur_data["lang"] == pre_data["lang"] + and cur_data["symbol"] == False + and pre_data["symbol"] + ): + cur_data["text"] = f"{pre_data['text']}{cur_data['text']}" + words.pop() + words += cur_word + if LangSegment.isLangMerge == True: + words = LangSegment._merge_results(words) + lang_count = LangSegment._lang_count + if lang_count and len(lang_count) > 0: + lang_count = dict( + sorted(lang_count.items(), key=lambda x: x[1], reverse=True) + ) + lang_count = list(lang_count.items()) + LangSegment._lang_count = lang_count + return words + + @staticmethod + def setfilters(filters): + # 当过滤器更改时，清除缓存 + # 필터가 변경되면 캐시를 지웁니다. + # フィルタが変更されると、キャッシュがクリアされます + # When the filter changes, clear the cache + if LangSegment.Langfilters != filters: + LangSegment._clears() + LangSegment.Langfilters = filters + pass + + @staticmethod + def getfilters(): + return LangSegment.Langfilters + + @staticmethod + def setPriorityThreshold(threshold: float): + LangSegment.LangPriorityThreshold = threshold + pass + + @staticmethod + def getPriorityThreshold(): + return LangSegment.LangPriorityThreshold + + @staticmethod + def getCounts(): + lang_count = LangSegment._lang_count + if lang_count is not None: + return lang_count + text_langs = LangSegment._text_langs + if text_langs is None or len(text_langs) == 0: + return [("zh", 0)] + lang_counts = defaultdict(int) + for d in text_langs: + lang_counts[d["lang"]] += ( + int(len(d["text"]) * 2) if d["lang"] == "zh" else len(d["text"]) + ) + lang_counts = dict( + sorted(lang_counts.items(), key=lambda x: x[1], reverse=True) + ) + lang_counts = list(lang_counts.items()) + LangSegment._lang_count = lang_counts + return lang_counts + + @staticmethod + def getTexts(text: str): + if text is None or len(text.strip()) == 0: + LangSegment._clears() + return [] + # lasts + text_langs = LangSegment._text_langs + if LangSegment._text_lasts == text and text_langs is not None: + return text_langs + # parse + LangSegment._text_waits = [] + LangSegment._lang_count = None + LangSegment._text_lasts = text + text = LangSegment._parse_symbols(text) + LangSegment._text_langs = text + return text + + @staticmethod + def classify(text: str): + return LangSegment.getTexts(text) + + +def setLangMerge(value: bool): + """是否优化合并结果""" + LangSegment.isLangMerge = value + pass + + +def getLangMerge(): + """是否优化合并结果""" + return LangSegment.isLangMerge + + +def setfilters(filters): + """ + 功能：语言过滤组功能, 可以指定保留语言。不在过滤组中的语言将被清除。您可随心搭配TTS语音合成所支持的语言。 + 기능: 언어 필터 그룹 기능, 예약된 언어를 지정할 수 있습니다. 필터 그룹에 없는 언어는 지워집니다. TTS 텍스트에서 지원하는 언어를 원하는 대로 일치시킬 수 있습니다. + 機能:言語フィルターグループ機能で、予約言語を指定できます。フィルターグループに含まれていない言語はクリアされます。TTS音声合成がサポートする言語を自由に組み合わせることができます。 + Function: Language filter group function, you can specify reserved languages. \n + Languages not in the filter group will be cleared. You can match the languages supported by TTS Text To Speech as you like.\n + Args: + filters (list): ["zh", "en", "ja", "ko"] 排名越前，优先级越高 + """ + LangSegment.setfilters(filters) + pass + + +def getfilters(): + """ + 功能：语言过滤组功能, 可以指定保留语言。不在过滤组中的语言将被清除。您可随心搭配TTS语音合成所支持的语言。 + 기능: 언어 필터 그룹 기능, 예약된 언어를 지정할 수 있습니다. 필터 그룹에 없는 언어는 지워집니다. TTS 텍스트에서 지원하는 언어를 원하는 대로 일치시킬 수 있습니다. + 機能:言語フィルターグループ機能で、予約言語を指定できます。フィルターグループに含まれていない言語はクリアされます。TTS音声合成がサポートする言語を自由に組み合わせることができます。 + Function: Language filter group function, you can specify reserved languages. \n + Languages not in the filter group will be cleared. You can match the languages supported by TTS Text To Speech as you like.\n + Args: + filters (list): ["zh", "en", "ja", "ko"] 排名越前，优先级越高 + """ + return LangSegment.getfilters() + + +# # @Deprecated：Use shorter setfilters +# def setLangfilters(filters): +# """ +# >0.1.9废除：使用更简短的setfilters +# """ +# setfilters(filters) +# # @Deprecated：Use shorter getfilters +# def getLangfilters(): +# """ +# >0.1.9废除：使用更简短的getfilters +# """ +# return getfilters() + + +def setKeepPinyin(value: bool): + """ + 可选保留：支持中文数字拼音格式，更方便前端实现拼音音素修改和推理，默认关闭 False 。\n + 开启后 True ，括号内的数字拼音格式均保留，并识别输出为："zh"中文。 + """ + LangSegment.keepPinyin = value + pass + + +def getKeepPinyin(): + """ + 可选保留：支持中文数字拼音格式，更方便前端实现拼音音素修改和推理，默认关闭 False 。\n + 开启后 True ，括号内的数字拼音格式均保留，并识别输出为："zh"中文。 + """ + return LangSegment.keepPinyin + + +def setEnablePreview(value: bool): + """ + 启用预览版功能（默认关闭） + Enable preview functionality (off by default) + Args: + value (bool): True=开启， False=关闭 + """ + LangSegment.EnablePreview = value == True + pass + + +def getEnablePreview(): + """ + 启用预览版功能（默认关闭） + Enable preview functionality (off by default) + Args: + value (bool): True=开启， False=关闭 + """ + return LangSegment.EnablePreview == True + + +def setPriorityThreshold(threshold: float): + """ + 中/日语言优先级阀值（评分范围为 0 ~ 1）:评分低于设定阀值 <0.89 时，启用 filters 中的优先级。\n + 中国語/日本語の優先度しきい値（スコア範囲0〜1）:スコアが設定されたしきい値<0.89未満の場合、フィルターの優先度が有効になります。\n + 중/일본어 우선 순위 임계값(점수 범위 0-1): 점수가 설정된 임계값 <0.89보다 낮을 때 필터에서 우선 순위를 활성화합니다. + Chinese and Japanese language priority threshold (score range is 0 ~ 1): The default threshold is 0.89. \n + Only the common characters between Chinese and Japanese are processed with confidence and priority. \n + Args: + threshold:float (score range is 0 ~ 1) + """ + LangSegment.setPriorityThreshold(threshold) + pass + + +def getPriorityThreshold(): + """ + 中/日语言优先级阀值（评分范围为 0 ~ 1）:评分低于设定阀值 <0.89 时，启用 filters 中的优先级。\n + 中国語/日本語の優先度しきい値（スコア範囲0〜1）:スコアが設定されたしきい値<0.89未満の場合、フィルターの優先度が有効になります。\n + 중/일본어 우선 순위 임계값(점수 범위 0-1): 점수가 설정된 임계값 <0.89보다 낮을 때 필터에서 우선 순위를 활성화합니다. + Chinese and Japanese language priority threshold (score range is 0 ~ 1): The default threshold is 0.89. \n + Only the common characters between Chinese and Japanese are processed with confidence and priority. \n + Args: + threshold:float (score range is 0 ~ 1) + """ + return LangSegment.getPriorityThreshold() + + +def getTexts(text: str): + """ + 功能：对输入的文本进行多语种分词\n + 기능: 입력 텍스트의 다국어 분할 \n + 機能:入力されたテキストの多言語セグメンテーション\n + Feature: Tokenizing multilingual text input.\n + 参数-Args: + text (str): Text content,文本内容\n + 返回-Returns: + list: 示例结果：[{'lang':'zh','text':'?'},...]\n + lang=语种 , text=内容\n + """ + return LangSegment.getTexts(text) + + +def getCounts(): + """ + 功能：分词结果统计，按语种字数降序，用于确定其主要语言\n + 기능: 주요 언어를 결정하는 데 사용되는 언어별 단어 수 내림차순으로 단어 분할 결과의 통계 \n + 機能:主な言語を決定するために使用される、言語の単語数の降順による単語分割結果の統計\n + Function: Tokenizing multilingual text input.\n + 返回-Returns: + list: 示例结果：[('zh', 5), ('ja', 2), ('en', 1)] = [(语种,字数含标点)]\n + """ + return LangSegment.getCounts() + + +def classify(text: str): + """ + 功能：兼容接口实现 + Function: Compatible interface implementation + """ + return LangSegment.classify(text) + + +def printList(langlist): + """ + 功能：打印数组结果 + 기능: 어레이 결과 인쇄 + 機能:配列結果を印刷 + Function: Print array results + """ + print("\n===================【打印结果】===================") + if langlist is None or len(langlist) == 0: + print("无内容结果,No content result") + return + for line in langlist: + print(line) + pass + + +def main(): + # ----------------------------------- + # 更新日志：新版本分词更加精准。 + # Changelog: The new version of the word segmentation is more accurate. + # チェンジログ:新しいバージョンの単語セグメンテーションはより正確です。 + # Changelog: 분할이라는 단어의 새로운 버전이 더 정확합니다. + # ----------------------------------- + + # 输入示例1：（包含日文，中文）Input Example 1: (including Japanese, Chinese) + # text = "“昨日は雨が降った，音楽、映画。。。”你今天学习日语了吗？春は桜の季節です。语种分词是语音合成必不可少的环节。言語分詞は音声合成に欠かせない環節である！" + + # 输入示例2：（包含日文，中文）Input Example 1: (including Japanese, Chinese) + # text = "欢迎来玩。東京，は日本の首都です。欢迎来玩. 太好了!" + + # 输入示例3：（包含日文，中文）Input Example 1: (including Japanese, Chinese) + # text = "明日、私たちは海辺にバカンスに行きます。你会说日语吗：“中国語、話せますか” 你的日语真好啊！" + + # 输入示例4：（包含日文，中文，韩语，英文）Input Example 4: (including Japanese, Chinese, Korean, English) + # text = "你的名字叫佐々木？吗？韩语中的안녕 오빠读什么呢？あなたの体育の先生は誰ですか? 此次发布会带来了四款iPhone 15系列机型和三款Apple Watch等一系列新品，这次的iPad Air采用了LCD屏幕" + + # 试验性支持："fr"法语 , "vi"越南语 , "ru"俄语 , "th"泰语。Experimental: Other language support. + LangSegment.setfilters(["fr", "vi", "ja", "zh", "ko", "en", "ru", "th"]) + text = """ +我喜欢在雨天里听音乐。 +I enjoy listening to music on rainy days. +雨の日に音楽を聴くのが好きです。 +비 오는 날에 음악을 듣는 것을 즐깁니다。 +J'aime écouter de la musique les jours de pluie. +Tôi thích nghe nhạc vào những ngày mưa. +Мне нравится слушать музыку в дождливую погоду. +ฉันชอบฟังเพลงในวันที่ฝนตก +""" + + # 进行分词：（接入TTS项目仅需一行代码调用）Segmentation: (Only one line of code is required to access the TTS project) + langlist = LangSegment.getTexts(text) + printList(langlist) + + # 语种统计:Language statistics: + print("\n===================【语种统计】===================") + # 获取所有语种数组结果，根据内容字数降序排列 + # Get the array results in all languages, sorted in descending order according to the number of content words + langCounts = LangSegment.getCounts() + print(langCounts, "\n") + + # 根据结果获取内容的主要语种 (语言，字数含标点) + # Get the main language of content based on the results (language, word count including punctuation) + lang, count = langCounts[0] + print(f"输入内容的主要语言为 = {lang} ，字数 = {count}") + print("==================================================\n") + + # 分词输出：lang=语言，text=内容。Word output: lang = language, text = content + # ===================【打印结果】=================== + # {'lang': 'zh', 'text': '你的名字叫'} + # {'lang': 'ja', 'text': '佐々木？'} + # {'lang': 'zh', 'text': '吗？韩语中的'} + # {'lang': 'ko', 'text': '안녕 오빠'} + # {'lang': 'zh', 'text': '读什么呢？'} + # {'lang': 'ja', 'text': 'あなたの体育の先生は誰ですか?'} + # {'lang': 'zh', 'text': ' 此次发布会带来了四款'} + # {'lang': 'en', 'text': 'i Phone '} + # {'lang': 'zh', 'text': '15系列机型和三款'} + # {'lang': 'en', 'text': 'Apple Watch '} + # {'lang': 'zh', 'text': '等一系列新品，这次的'} + # {'lang': 'en', 'text': 'i Pad Air '} + # {'lang': 'zh', 'text': '采用了'} + # {'lang': 'en', 'text': 'L C D '} + # {'lang': 'zh', 'text': '屏幕'} + # ===================【语种统计】=================== + + # ===================【语种统计】=================== + # [('zh', 51), ('ja', 19), ('en', 18), ('ko', 5)] + + # 输入内容的主要语言为 = zh ，字数 = 51 + # ================================================== + # The main language of the input content is = zh, word count = 51 + + +if __name__ == "__main__": + main() diff --git a/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/__init__.py b/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..be4ce2d65d0c3a228f1a36934000d3d682a79602 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/__init__.py @@ -0,0 +1,24 @@ +from .LangSegment import ( + LangSegment, + classify, + getCounts, + getEnablePreview, + getfilters, + getKeepPinyin, + getLangMerge, + getPriorityThreshold, + getTexts, + printList, + setEnablePreview, + setfilters, + setKeepPinyin, + setLangMerge, + setPriorityThreshold, +) + +# release +__version__ = "0.3.5" + + +# develop +__develop__ = "dev-0.0.1" diff --git a/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/utils/__init__.py b/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/utils/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/utils/num.py b/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/utils/num.py new file mode 100755 index 0000000000000000000000000000000000000000..80b2c3b779242847350477ea7fe9beed11afde24 --- /dev/null +++ b/src/YingMusicSinger/utils/f5_tts/thirdparty/LangSegment/utils/num.py @@ -0,0 +1,332 @@ +# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Digital processing from GPT_SoVITS num.py （thanks） +""" +Rules to verbalize numbers into Chinese characters. +https://zh.wikipedia.org/wiki/中文数字#現代中文 +""" + +import re +from collections import OrderedDict +from typing import List + +DIGITS = {str(i): tran for i, tran in enumerate("零一二三四五六七八九")} +UNITS = OrderedDict( + { + 1: "十", + 2: "百", + 3: "千", + 4: "万", + 8: "亿", + } +) + +COM_QUANTIFIERS = "(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)" + +# 分数表达式 +RE_FRAC = re.compile(r"(-?)(\d+)/(\d+)") + + +def replace_frac(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + sign = match.group(1) + nominator = match.group(2) + denominator = match.group(3) + sign: str = "负" if sign else "" + nominator: str = num2str(nominator) + denominator: str = num2str(denominator) + result = f"{sign}{denominator}分之{nominator}" + return result + + +# 百分数表达式 +RE_PERCENTAGE = re.compile(r"(-?)(\d+(\.\d+)?)%") + + +def replace_percentage(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + sign = match.group(1) + percent = match.group(2) + sign: str = "负" if sign else "" + percent: str = num2str(percent) + result = f"{sign}百分之{percent}" + return result + + +# 整数表达式 +# 带负号的整数 -10 +RE_INTEGER = re.compile(r"(-)" r"(\d+)") + + +def replace_negative_num(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + sign = match.group(1) + number = match.group(2) + sign: str = "负" if sign else "" + number: str = num2str(number) + result = f"{sign}{number}" + return result + + +# 编号-无符号整形 +# 00078 +RE_DEFAULT_NUM = re.compile(r"\d{3}\d*") + + +def replace_default_num(match): + """ + Args: + match (re.Match) + Returns: + str + """ + number = match.group(0) + return verbalize_digit(number, alt_one=True) + + +# 加减乘除 +# RE_ASMD = re.compile( +# r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))') +RE_ASMD = re.compile( + r"((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))" +) + +asmd_map = {"+": "加", "-": "减", "×": "乘", "÷": "除", "=": "等于"} + + +def replace_asmd(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + result = match.group(1) + asmd_map[match.group(8)] + match.group(9) + return result + + +# 次方专项 +RE_POWER = re.compile(r"[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+") + +power_map = { + "⁰": "0", + "¹": "1", + "²": "2", + "³": "3", + "⁴": "4", + "⁵": "5", + "⁶": "6", + "⁷": "7", + "⁸": "8", + "⁹": "9", + "ˣ": "x", + "ʸ": "y", + "ⁿ": "n", +} + + +def replace_power(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + power_num = "" + for m in match.group(0): + power_num += power_map[m] + result = "的" + power_num + "次方" + return result + + +# 数字表达式 +# 纯小数 +RE_DECIMAL_NUM = re.compile(r"(-?)((\d+)(\.\d+))" r"|(\.(\d+))") +# 正整数 + 量词 +RE_POSITIVE_QUANTIFIERS = re.compile(r"(\d+)([多余几\+])?" + COM_QUANTIFIERS) +RE_NUMBER = re.compile(r"(-?)((\d+)(\.\d+)?)" r"|(\.(\d+))") + + +def replace_positive_quantifier(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + number = match.group(1) + match_2 = match.group(2) + if match_2 == "+": + match_2 = "多" + match_2: str = match_2 if match_2 else "" + quantifiers: str = match.group(3) + number: str = num2str(number) + result = f"{number}{match_2}{quantifiers}" + return result + + +def replace_number(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + sign = match.group(1) + number = match.group(2) + pure_decimal = match.group(5) + if pure_decimal: + result = num2str(pure_decimal) + else: + sign: str = "负" if sign else "" + number: str = num2str(number) + result = f"{sign}{number}" + return result + + +# 范围表达式 +# match.group(1) and match.group(8) are copy from RE_NUMBER + +RE_RANGE = re.compile( + r""" + (? str: + """ + Args: + match (re.Match) + Returns: + str + """ + first, second = match.group(1), match.group(6) + first = RE_NUMBER.sub(replace_number, first) + second = RE_NUMBER.sub(replace_number, second) + result = f"{first}到{second}" + return result + + +# ~至表达式 +RE_TO_RANGE = re.compile( + r"((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)" +) + + +def replace_to_range(match) -> str: + """ + Args: + match (re.Match) + Returns: + str + """ + result = match.group(0).replace("~", "至") + return result + + +def _get_value(value_string: str, use_zero: bool = True) -> List[str]: + stripped = value_string.lstrip("0") + if len(stripped) == 0: + return [] + elif len(stripped) == 1: + if use_zero and len(stripped) < len(value_string): + return [DIGITS["0"], DIGITS[stripped]] + else: + return [DIGITS[stripped]] + else: + largest_unit = next( + power for power in reversed(UNITS.keys()) if power < len(stripped) + ) + first_part = value_string[:-largest_unit] + second_part = value_string[-largest_unit:] + return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(second_part) + + +def verbalize_cardinal(value_string: str) -> str: + if not value_string: + return "" + + # 000 -> '零' , 0 -> '零' + value_string = value_string.lstrip("0") + if len(value_string) == 0: + return DIGITS["0"] + + result_symbols = _get_value(value_string) + # verbalized number starting with '一十*' is abbreviated as `十*` + if ( + len(result_symbols) >= 2 + and result_symbols[0] == DIGITS["1"] + and result_symbols[1] == UNITS[1] + ): + result_symbols = result_symbols[1:] + return "".join(result_symbols) + + +def verbalize_digit(value_string: str, alt_one=False) -> str: + result_symbols = [DIGITS[digit] for digit in value_string] + result = "".join(result_symbols) + if alt_one: + result = result.replace("一", "幺") + return result + + +def num2str(value_string: str) -> str: + integer_decimal = value_string.split(".") + if len(integer_decimal) == 1: + integer = integer_decimal[0] + decimal = "" + elif len(integer_decimal) == 2: + integer, decimal = integer_decimal + else: + raise ValueError( + f"The value string: '${value_string}' has more than one point in it." + ) + + result = verbalize_cardinal(integer) + + decimal = decimal.rstrip("0") + if decimal: + # '.22' is verbalized as '零点二二' + # '3.20' is verbalized as '三点二 + result = result if result else "零" + result += "点" + verbalize_digit(decimal) + return result + + +if __name__ == "__main__": + text = "" + text = num2str(text) + print(text) + pass diff --git a/src/YingMusicSinger/utils/stable_audio_tools/__init__.py b/src/YingMusicSinger/utils/stable_audio_tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/YingMusicSinger/utils/stable_audio_tools/adp.py b/src/YingMusicSinger/utils/stable_audio_tools/adp.py new file mode 100755 index 0000000000000000000000000000000000000000..56b7bbbb277665f2ec96c4319bc0db8e7fa1022c --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/adp.py @@ -0,0 +1,1686 @@ +# Copied and modified from https://github.com/archinetai/audio-diffusion-pytorch/blob/v0.0.94/audio_diffusion_pytorch/modules.py under MIT License +# License can be found in LICENSES/LICENSE_ADP.txt + +import math +from inspect import isfunction +from math import ceil, floor, log, log2, pi +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union + +import torch +import torch.nn as nn +from dac.nn.layers import Snake1d +from einops import rearrange, reduce, repeat +from einops.layers.torch import Rearrange +from einops_exts import rearrange_many +from packaging import version +from torch import Tensor, einsum +from torch.backends.cuda import sdp_kernel +from torch.nn import functional as F + +""" +Utils +""" + + +class ConditionedSequential(nn.Module): + def __init__(self, *modules): + super().__init__() + self.module_list = nn.ModuleList(*modules) + + def forward(self, x: Tensor, mapping: Optional[Tensor] = None): + for module in self.module_list: + x = module(x, mapping) + return x + + +T = TypeVar("T") + + +def default(val: Optional[T], d: Union[Callable[..., T], T]) -> T: + if exists(val): + return val + return d() if isfunction(d) else d + + +def exists(val: Optional[T]) -> T: + return val is not None + + +def closest_power_2(x: float) -> int: + exponent = log2(x) + distance_fn = lambda z: abs(x - 2**z) # noqa + exponent_closest = min((floor(exponent), ceil(exponent)), key=distance_fn) + return 2 ** int(exponent_closest) + + +def group_dict_by_prefix(prefix: str, d: Dict) -> Tuple[Dict, Dict]: + return_dicts: Tuple[Dict, Dict] = ({}, {}) + for key in d.keys(): + no_prefix = int(not key.startswith(prefix)) + return_dicts[no_prefix][key] = d[key] + return return_dicts + + +def groupby(prefix: str, d: Dict, keep_prefix: bool = False) -> Tuple[Dict, Dict]: + kwargs_with_prefix, kwargs = group_dict_by_prefix(prefix, d) + if keep_prefix: + return kwargs_with_prefix, kwargs + kwargs_no_prefix = {k[len(prefix) :]: v for k, v in kwargs_with_prefix.items()} + return kwargs_no_prefix, kwargs + + +""" +Convolutional Blocks +""" +import typing as tp + +# Copied from https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/modules/conv.py under MIT License +# License available in LICENSES/LICENSE_META.txt + + +def get_extra_padding_for_conv1d( + x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 +) -> int: + """See `pad_for_conv1d`.""" + length = x.shape[-1] + n_frames = (length - kernel_size + padding_total) / stride + 1 + ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) + return ideal_length - length + + +def pad_for_conv1d( + x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 +): + """Pad for a convolution to make sure that the last window is full. + Extra padding is added at the end. This is required to ensure that we can rebuild + an output of the same length, as otherwise, even with padding, some time steps + might get removed. + For instance, with total padding = 4, kernel size = 4, stride = 2: + 0 0 1 2 3 4 5 0 0 # (0s are padding) + 1 2 3 # (output frames of a convolution, last 0 is never used) + 0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding) + 1 2 3 4 # once you removed padding, we are missing one time step ! + """ + extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) + return F.pad(x, (0, extra_padding)) + + +def pad1d( + x: torch.Tensor, + paddings: tp.Tuple[int, int], + mode: str = "constant", + value: float = 0.0, +): + """Tiny wrapper around F.pad, just to allow for reflect padding on small input. + If this is the case, we insert extra 0 padding to the right before the reflection happen. + """ + length = x.shape[-1] + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + if mode == "reflect": + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + x = F.pad(x, (0, extra_pad)) + padded = F.pad(x, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + else: + return F.pad(x, paddings, mode, value) + + +def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]): + """Remove padding from x, handling properly zero padding. Only for 1d!""" + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + assert (padding_left + padding_right) <= x.shape[-1] + end = x.shape[-1] - padding_right + return x[..., padding_left:end] + + +class Conv1d(nn.Conv1d): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, x: Tensor, causal=False) -> Tensor: + kernel_size = self.kernel_size[0] + stride = self.stride[0] + dilation = self.dilation[0] + kernel_size = ( + kernel_size - 1 + ) * dilation + 1 # effective kernel size with dilations + padding_total = kernel_size - stride + extra_padding = get_extra_padding_for_conv1d( + x, kernel_size, stride, padding_total + ) + if causal: + # Left padding for causal + x = pad1d(x, (padding_total, extra_padding)) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + x = pad1d(x, (padding_left, padding_right + extra_padding)) + return super().forward(x) + + +class ConvTranspose1d(nn.ConvTranspose1d): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, x: Tensor, causal=False) -> Tensor: + kernel_size = self.kernel_size[0] + stride = self.stride[0] + padding_total = kernel_size - stride + + y = super().forward(x) + + # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be + # removed at the very end, when keeping only the right length for the output, + # as removing it here would require also passing the length at the matching layer + # in the encoder. + if causal: + padding_right = ceil(padding_total) + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + return y + + +def Downsample1d( + in_channels: int, out_channels: int, factor: int, kernel_multiplier: int = 2 +) -> nn.Module: + assert kernel_multiplier % 2 == 0, "Kernel multiplier must be even" + + return Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=factor * kernel_multiplier + 1, + stride=factor, + ) + + +def Upsample1d( + in_channels: int, out_channels: int, factor: int, use_nearest: bool = False +) -> nn.Module: + if factor == 1: + return Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=3) + + if use_nearest: + return nn.Sequential( + nn.Upsample(scale_factor=factor, mode="nearest"), + Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=3), + ) + else: + return ConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=factor * 2, + stride=factor, + ) + + +class ConvBlock1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + *, + kernel_size: int = 3, + stride: int = 1, + dilation: int = 1, + num_groups: int = 8, + use_norm: bool = True, + use_snake: bool = False, + ) -> None: + super().__init__() + + self.groupnorm = ( + nn.GroupNorm(num_groups=num_groups, num_channels=in_channels) + if use_norm + else nn.Identity() + ) + + if use_snake: + self.activation = Snake1d(in_channels) + else: + self.activation = nn.SiLU() + + self.project = Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + ) + + def forward( + self, + x: Tensor, + scale_shift: Optional[Tuple[Tensor, Tensor]] = None, + causal=False, + ) -> Tensor: + x = self.groupnorm(x) + if exists(scale_shift): + scale, shift = scale_shift + x = x * (scale + 1) + shift + x = self.activation(x) + return self.project(x, causal=causal) + + +class MappingToScaleShift(nn.Module): + def __init__( + self, + features: int, + channels: int, + ): + super().__init__() + + self.to_scale_shift = nn.Sequential( + nn.SiLU(), + nn.Linear(in_features=features, out_features=channels * 2), + ) + + def forward(self, mapping: Tensor) -> Tuple[Tensor, Tensor]: + scale_shift = self.to_scale_shift(mapping) + scale_shift = rearrange(scale_shift, "b c -> b c 1") + scale, shift = scale_shift.chunk(2, dim=1) + return scale, shift + + +class ResnetBlock1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + *, + kernel_size: int = 3, + stride: int = 1, + dilation: int = 1, + use_norm: bool = True, + use_snake: bool = False, + num_groups: int = 8, + context_mapping_features: Optional[int] = None, + ) -> None: + super().__init__() + + self.use_mapping = exists(context_mapping_features) + + self.block1 = ConvBlock1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + use_norm=use_norm, + num_groups=num_groups, + use_snake=use_snake, + ) + + if self.use_mapping: + assert exists(context_mapping_features) + self.to_scale_shift = MappingToScaleShift( + features=context_mapping_features, channels=out_channels + ) + + self.block2 = ConvBlock1d( + in_channels=out_channels, + out_channels=out_channels, + use_norm=use_norm, + num_groups=num_groups, + use_snake=use_snake, + ) + + self.to_out = ( + Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1) + if in_channels != out_channels + else nn.Identity() + ) + + def forward( + self, x: Tensor, mapping: Optional[Tensor] = None, causal=False + ) -> Tensor: + assert_message = "context mapping required if context_mapping_features > 0" + assert not (self.use_mapping ^ exists(mapping)), assert_message + + h = self.block1(x, causal=causal) + + scale_shift = None + if self.use_mapping: + scale_shift = self.to_scale_shift(mapping) + + h = self.block2(h, scale_shift=scale_shift, causal=causal) + + return h + self.to_out(x) + + +class Patcher(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + patch_size: int, + context_mapping_features: Optional[int] = None, + use_snake: bool = False, + ): + super().__init__() + assert_message = f"out_channels must be divisible by patch_size ({patch_size})" + assert out_channels % patch_size == 0, assert_message + self.patch_size = patch_size + + self.block = ResnetBlock1d( + in_channels=in_channels, + out_channels=out_channels // patch_size, + num_groups=1, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + + def forward( + self, x: Tensor, mapping: Optional[Tensor] = None, causal=False + ) -> Tensor: + x = self.block(x, mapping, causal=causal) + x = rearrange(x, "b c (l p) -> b (c p) l", p=self.patch_size) + return x + + +class Unpatcher(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + patch_size: int, + context_mapping_features: Optional[int] = None, + use_snake: bool = False, + ): + super().__init__() + assert_message = f"in_channels must be divisible by patch_size ({patch_size})" + assert in_channels % patch_size == 0, assert_message + self.patch_size = patch_size + + self.block = ResnetBlock1d( + in_channels=in_channels // patch_size, + out_channels=out_channels, + num_groups=1, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + + def forward( + self, x: Tensor, mapping: Optional[Tensor] = None, causal=False + ) -> Tensor: + x = rearrange(x, " b (c p) l -> b c (l p) ", p=self.patch_size) + x = self.block(x, mapping, causal=causal) + return x + + +""" +Attention Components +""" + + +def FeedForward(features: int, multiplier: int) -> nn.Module: + mid_features = features * multiplier + return nn.Sequential( + nn.Linear(in_features=features, out_features=mid_features), + nn.GELU(), + nn.Linear(in_features=mid_features, out_features=features), + ) + + +def add_mask(sim: Tensor, mask: Tensor) -> Tensor: + b, ndim = sim.shape[0], mask.ndim + if ndim == 3: + mask = rearrange(mask, "b n m -> b 1 n m") + if ndim == 2: + mask = repeat(mask, "n m -> b 1 n m", b=b) + max_neg_value = -torch.finfo(sim.dtype).max + sim = sim.masked_fill(~mask, max_neg_value) + return sim + + +def causal_mask(q: Tensor, k: Tensor) -> Tensor: + b, i, j, device = q.shape[0], q.shape[-2], k.shape[-2], q.device + mask = ~torch.ones((i, j), dtype=torch.bool, device=device).triu(j - i + 1) + mask = repeat(mask, "n m -> b n m", b=b) + return mask + + +class AttentionBase(nn.Module): + def __init__( + self, + features: int, + *, + head_features: int, + num_heads: int, + out_features: Optional[int] = None, + ): + super().__init__() + self.scale = head_features**-0.5 + self.num_heads = num_heads + mid_features = head_features * num_heads + out_features = default(out_features, features) + + self.to_out = nn.Linear(in_features=mid_features, out_features=out_features) + + self.use_flash = torch.cuda.is_available() and version.parse( + torch.__version__ + ) >= version.parse("2.0.0") + + if not self.use_flash: + return + + device_properties = torch.cuda.get_device_properties(torch.device("cuda")) + + if device_properties.major == 8 and device_properties.minor == 0: + # Use flash attention for A100 GPUs + self.sdp_kernel_config = (True, False, False) + else: + # Don't use flash attention for other GPUs + self.sdp_kernel_config = (False, True, True) + + def forward( + self, + q: Tensor, + k: Tensor, + v: Tensor, + mask: Optional[Tensor] = None, + is_causal: bool = False, + ) -> Tensor: + # Split heads + q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads) + + if not self.use_flash: + if is_causal and not mask: + # Mask out future tokens for causal attention + mask = causal_mask(q, k) + + # Compute similarity matrix and add eventual mask + sim = einsum("... n d, ... m d -> ... n m", q, k) * self.scale + sim = add_mask(sim, mask) if exists(mask) else sim + + # Get attention matrix with softmax + attn = sim.softmax(dim=-1, dtype=torch.float32) + + # Compute values + out = einsum("... n m, ... m d -> ... n d", attn, v) + else: + with sdp_kernel(*self.sdp_kernel_config): + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, is_causal=is_causal + ) + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +class Attention(nn.Module): + def __init__( + self, + features: int, + *, + head_features: int, + num_heads: int, + out_features: Optional[int] = None, + context_features: Optional[int] = None, + causal: bool = False, + ): + super().__init__() + self.context_features = context_features + self.causal = causal + mid_features = head_features * num_heads + context_features = default(context_features, features) + + self.norm = nn.LayerNorm(features) + self.norm_context = nn.LayerNorm(context_features) + self.to_q = nn.Linear( + in_features=features, out_features=mid_features, bias=False + ) + self.to_kv = nn.Linear( + in_features=context_features, out_features=mid_features * 2, bias=False + ) + self.attention = AttentionBase( + features, + num_heads=num_heads, + head_features=head_features, + out_features=out_features, + ) + + def forward( + self, + x: Tensor, # [b, n, c] + context: Optional[Tensor] = None, # [b, m, d] + context_mask: Optional[Tensor] = None, # [b, m], false is masked, + causal: Optional[bool] = False, + ) -> Tensor: + assert_message = "You must provide a context when using context_features" + assert not self.context_features or exists(context), assert_message + # Use context if provided + context = default(context, x) + # Normalize then compute q from input and k,v from context + x, context = self.norm(x), self.norm_context(context) + + q, k, v = (self.to_q(x), *torch.chunk(self.to_kv(context), chunks=2, dim=-1)) + + if exists(context_mask): + # Mask out cross-attention for padding tokens + mask = repeat(context_mask, "b m -> b m d", d=v.shape[-1]) + k, v = k * mask, v * mask + + # Compute and return attention + return self.attention(q, k, v, is_causal=self.causal or causal) + + +def FeedForward(features: int, multiplier: int) -> nn.Module: + mid_features = features * multiplier + return nn.Sequential( + nn.Linear(in_features=features, out_features=mid_features), + nn.GELU(), + nn.Linear(in_features=mid_features, out_features=features), + ) + + +""" +Transformer Blocks +""" + + +class TransformerBlock(nn.Module): + def __init__( + self, + features: int, + num_heads: int, + head_features: int, + multiplier: int, + context_features: Optional[int] = None, + ): + super().__init__() + + self.use_cross_attention = exists(context_features) and context_features > 0 + + self.attention = Attention( + features=features, num_heads=num_heads, head_features=head_features + ) + + if self.use_cross_attention: + self.cross_attention = Attention( + features=features, + num_heads=num_heads, + head_features=head_features, + context_features=context_features, + ) + + self.feed_forward = FeedForward(features=features, multiplier=multiplier) + + def forward( + self, + x: Tensor, + *, + context: Optional[Tensor] = None, + context_mask: Optional[Tensor] = None, + causal: Optional[bool] = False, + ) -> Tensor: + x = self.attention(x, causal=causal) + x + if self.use_cross_attention: + x = self.cross_attention(x, context=context, context_mask=context_mask) + x + x = self.feed_forward(x) + x + return x + + +""" +Transformers +""" + + +class Transformer1d(nn.Module): + def __init__( + self, + num_layers: int, + channels: int, + num_heads: int, + head_features: int, + multiplier: int, + context_features: Optional[int] = None, + ): + super().__init__() + + self.to_in = nn.Sequential( + nn.GroupNorm(num_groups=32, num_channels=channels, eps=1e-6, affine=True), + Conv1d( + in_channels=channels, + out_channels=channels, + kernel_size=1, + ), + Rearrange("b c t -> b t c"), + ) + + self.blocks = nn.ModuleList( + [ + TransformerBlock( + features=channels, + head_features=head_features, + num_heads=num_heads, + multiplier=multiplier, + context_features=context_features, + ) + for i in range(num_layers) + ] + ) + + self.to_out = nn.Sequential( + Rearrange("b t c -> b c t"), + Conv1d( + in_channels=channels, + out_channels=channels, + kernel_size=1, + ), + ) + + def forward( + self, + x: Tensor, + *, + context: Optional[Tensor] = None, + context_mask: Optional[Tensor] = None, + causal=False, + ) -> Tensor: + x = self.to_in(x) + for block in self.blocks: + x = block(x, context=context, context_mask=context_mask, causal=causal) + x = self.to_out(x) + return x + + +""" +Time Embeddings +""" + + +class SinusoidalEmbedding(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.dim = dim + + def forward(self, x: Tensor) -> Tensor: + device, half_dim = x.device, self.dim // 2 + emb = torch.tensor(log(10000) / (half_dim - 1), device=device) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = rearrange(x, "i -> i 1") * rearrange(emb, "j -> 1 j") + return torch.cat((emb.sin(), emb.cos()), dim=-1) + + +class LearnedPositionalEmbedding(nn.Module): + """Used for continuous time""" + + def __init__(self, dim: int): + super().__init__() + assert (dim % 2) == 0 + half_dim = dim // 2 + self.weights = nn.Parameter(torch.randn(half_dim)) + + def forward(self, x: Tensor) -> Tensor: + x = rearrange(x, "b -> b 1") + freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * pi + fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1) + fouriered = torch.cat((x, fouriered), dim=-1) + return fouriered + + +def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module: + return nn.Sequential( + LearnedPositionalEmbedding(dim), + nn.Linear(in_features=dim + 1, out_features=out_features), + ) + + +""" +Encoder/Decoder Components +""" + + +class DownsampleBlock1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + *, + factor: int, + num_groups: int, + num_layers: int, + kernel_multiplier: int = 2, + use_pre_downsample: bool = True, + use_skip: bool = False, + use_snake: bool = False, + extract_channels: int = 0, + context_channels: int = 0, + num_transformer_blocks: int = 0, + attention_heads: Optional[int] = None, + attention_features: Optional[int] = None, + attention_multiplier: Optional[int] = None, + context_mapping_features: Optional[int] = None, + context_embedding_features: Optional[int] = None, + ): + super().__init__() + self.use_pre_downsample = use_pre_downsample + self.use_skip = use_skip + self.use_transformer = num_transformer_blocks > 0 + self.use_extract = extract_channels > 0 + self.use_context = context_channels > 0 + + channels = out_channels if use_pre_downsample else in_channels + + self.downsample = Downsample1d( + in_channels=in_channels, + out_channels=out_channels, + factor=factor, + kernel_multiplier=kernel_multiplier, + ) + + self.blocks = nn.ModuleList( + [ + ResnetBlock1d( + in_channels=channels + context_channels if i == 0 else channels, + out_channels=channels, + num_groups=num_groups, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + for i in range(num_layers) + ] + ) + + if self.use_transformer: + assert (exists(attention_heads) or exists(attention_features)) and exists( + attention_multiplier + ) + + if attention_features is None and attention_heads is not None: + attention_features = channels // attention_heads + + if attention_heads is None and attention_features is not None: + attention_heads = channels // attention_features + + self.transformer = Transformer1d( + num_layers=num_transformer_blocks, + channels=channels, + num_heads=attention_heads, + head_features=attention_features, + multiplier=attention_multiplier, + context_features=context_embedding_features, + ) + + if self.use_extract: + num_extract_groups = min(num_groups, extract_channels) + self.to_extracted = ResnetBlock1d( + in_channels=out_channels, + out_channels=extract_channels, + num_groups=num_extract_groups, + use_snake=use_snake, + ) + + def forward( + self, + x: Tensor, + *, + mapping: Optional[Tensor] = None, + channels: Optional[Tensor] = None, + embedding: Optional[Tensor] = None, + embedding_mask: Optional[Tensor] = None, + causal: Optional[bool] = False, + ) -> Union[Tuple[Tensor, List[Tensor]], Tensor]: + if self.use_pre_downsample: + x = self.downsample(x) + + if self.use_context and exists(channels): + x = torch.cat([x, channels], dim=1) + + skips = [] + for block in self.blocks: + x = block(x, mapping=mapping, causal=causal) + skips += [x] if self.use_skip else [] + + if self.use_transformer: + x = self.transformer( + x, context=embedding, context_mask=embedding_mask, causal=causal + ) + skips += [x] if self.use_skip else [] + + if not self.use_pre_downsample: + x = self.downsample(x) + + if self.use_extract: + extracted = self.to_extracted(x) + return x, extracted + + return (x, skips) if self.use_skip else x + + +class UpsampleBlock1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + *, + factor: int, + num_layers: int, + num_groups: int, + use_nearest: bool = False, + use_pre_upsample: bool = False, + use_skip: bool = False, + use_snake: bool = False, + skip_channels: int = 0, + use_skip_scale: bool = False, + extract_channels: int = 0, + num_transformer_blocks: int = 0, + attention_heads: Optional[int] = None, + attention_features: Optional[int] = None, + attention_multiplier: Optional[int] = None, + context_mapping_features: Optional[int] = None, + context_embedding_features: Optional[int] = None, + ): + super().__init__() + + self.use_extract = extract_channels > 0 + self.use_pre_upsample = use_pre_upsample + self.use_transformer = num_transformer_blocks > 0 + self.use_skip = use_skip + self.skip_scale = 2**-0.5 if use_skip_scale else 1.0 + + channels = out_channels if use_pre_upsample else in_channels + + self.blocks = nn.ModuleList( + [ + ResnetBlock1d( + in_channels=channels + skip_channels, + out_channels=channels, + num_groups=num_groups, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + for _ in range(num_layers) + ] + ) + + if self.use_transformer: + assert (exists(attention_heads) or exists(attention_features)) and exists( + attention_multiplier + ) + + if attention_features is None and attention_heads is not None: + attention_features = channels // attention_heads + + if attention_heads is None and attention_features is not None: + attention_heads = channels // attention_features + + self.transformer = Transformer1d( + num_layers=num_transformer_blocks, + channels=channels, + num_heads=attention_heads, + head_features=attention_features, + multiplier=attention_multiplier, + context_features=context_embedding_features, + ) + + self.upsample = Upsample1d( + in_channels=in_channels, + out_channels=out_channels, + factor=factor, + use_nearest=use_nearest, + ) + + if self.use_extract: + num_extract_groups = min(num_groups, extract_channels) + self.to_extracted = ResnetBlock1d( + in_channels=out_channels, + out_channels=extract_channels, + num_groups=num_extract_groups, + use_snake=use_snake, + ) + + def add_skip(self, x: Tensor, skip: Tensor) -> Tensor: + return torch.cat([x, skip * self.skip_scale], dim=1) + + def forward( + self, + x: Tensor, + *, + skips: Optional[List[Tensor]] = None, + mapping: Optional[Tensor] = None, + embedding: Optional[Tensor] = None, + embedding_mask: Optional[Tensor] = None, + causal: Optional[bool] = False, + ) -> Union[Tuple[Tensor, Tensor], Tensor]: + if self.use_pre_upsample: + x = self.upsample(x) + + for block in self.blocks: + x = self.add_skip(x, skip=skips.pop()) if exists(skips) else x + x = block(x, mapping=mapping, causal=causal) + + if self.use_transformer: + x = self.transformer( + x, context=embedding, context_mask=embedding_mask, causal=causal + ) + + if not self.use_pre_upsample: + x = self.upsample(x) + + if self.use_extract: + extracted = self.to_extracted(x) + return x, extracted + + return x + + +class BottleneckBlock1d(nn.Module): + def __init__( + self, + channels: int, + *, + num_groups: int, + num_transformer_blocks: int = 0, + attention_heads: Optional[int] = None, + attention_features: Optional[int] = None, + attention_multiplier: Optional[int] = None, + context_mapping_features: Optional[int] = None, + context_embedding_features: Optional[int] = None, + use_snake: bool = False, + ): + super().__init__() + self.use_transformer = num_transformer_blocks > 0 + + self.pre_block = ResnetBlock1d( + in_channels=channels, + out_channels=channels, + num_groups=num_groups, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + + if self.use_transformer: + assert (exists(attention_heads) or exists(attention_features)) and exists( + attention_multiplier + ) + + if attention_features is None and attention_heads is not None: + attention_features = channels // attention_heads + + if attention_heads is None and attention_features is not None: + attention_heads = channels // attention_features + + self.transformer = Transformer1d( + num_layers=num_transformer_blocks, + channels=channels, + num_heads=attention_heads, + head_features=attention_features, + multiplier=attention_multiplier, + context_features=context_embedding_features, + ) + + self.post_block = ResnetBlock1d( + in_channels=channels, + out_channels=channels, + num_groups=num_groups, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + + def forward( + self, + x: Tensor, + *, + mapping: Optional[Tensor] = None, + embedding: Optional[Tensor] = None, + embedding_mask: Optional[Tensor] = None, + causal: Optional[bool] = False, + ) -> Tensor: + x = self.pre_block(x, mapping=mapping, causal=causal) + if self.use_transformer: + x = self.transformer( + x, context=embedding, context_mask=embedding_mask, causal=causal + ) + x = self.post_block(x, mapping=mapping, causal=causal) + return x + + +""" +UNet +""" + + +class UNet1d(nn.Module): + def __init__( + self, + in_channels: int, + channels: int, + multipliers: Sequence[int], + factors: Sequence[int], + num_blocks: Sequence[int], + attentions: Sequence[int], + patch_size: int = 1, + resnet_groups: int = 8, + use_context_time: bool = True, + kernel_multiplier_downsample: int = 2, + use_nearest_upsample: bool = False, + use_skip_scale: bool = True, + use_snake: bool = False, + use_stft: bool = False, + use_stft_context: bool = False, + out_channels: Optional[int] = None, + context_features: Optional[int] = None, + context_features_multiplier: int = 4, + context_channels: Optional[Sequence[int]] = None, + context_embedding_features: Optional[int] = None, + **kwargs, + ): + super().__init__() + out_channels = default(out_channels, in_channels) + context_channels = list(default(context_channels, [])) + num_layers = len(multipliers) - 1 + use_context_features = exists(context_features) + use_context_channels = len(context_channels) > 0 + context_mapping_features = None + + attention_kwargs, kwargs = groupby("attention_", kwargs, keep_prefix=True) + + self.num_layers = num_layers + self.use_context_time = use_context_time + self.use_context_features = use_context_features + self.use_context_channels = use_context_channels + self.use_stft = use_stft + self.use_stft_context = use_stft_context + + self.context_features = context_features + context_channels_pad_length = num_layers + 1 - len(context_channels) + context_channels = context_channels + [0] * context_channels_pad_length + self.context_channels = context_channels + self.context_embedding_features = context_embedding_features + + if use_context_channels: + has_context = [c > 0 for c in context_channels] + self.has_context = has_context + self.channels_ids = [sum(has_context[:i]) for i in range(len(has_context))] + + assert ( + len(factors) == num_layers + and len(attentions) >= num_layers + and len(num_blocks) == num_layers + ) + + if use_context_time or use_context_features: + context_mapping_features = channels * context_features_multiplier + + self.to_mapping = nn.Sequential( + nn.Linear(context_mapping_features, context_mapping_features), + nn.GELU(), + nn.Linear(context_mapping_features, context_mapping_features), + nn.GELU(), + ) + + if use_context_time: + assert exists(context_mapping_features) + self.to_time = nn.Sequential( + TimePositionalEmbedding( + dim=channels, out_features=context_mapping_features + ), + nn.GELU(), + ) + + if use_context_features: + assert exists(context_features) and exists(context_mapping_features) + self.to_features = nn.Sequential( + nn.Linear( + in_features=context_features, out_features=context_mapping_features + ), + nn.GELU(), + ) + + if use_stft: + stft_kwargs, kwargs = groupby("stft_", kwargs) + assert "num_fft" in stft_kwargs, "stft_num_fft required if use_stft=True" + stft_channels = (stft_kwargs["num_fft"] // 2 + 1) * 2 + in_channels *= stft_channels + out_channels *= stft_channels + context_channels[0] *= stft_channels if use_stft_context else 1 + assert exists(in_channels) and exists(out_channels) + self.stft = STFT(**stft_kwargs) + + assert not kwargs, f"Unknown arguments: {', '.join(list(kwargs.keys()))}" + + self.to_in = Patcher( + in_channels=in_channels + context_channels[0], + out_channels=channels * multipliers[0], + patch_size=patch_size, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + + self.downsamples = nn.ModuleList( + [ + DownsampleBlock1d( + in_channels=channels * multipliers[i], + out_channels=channels * multipliers[i + 1], + context_mapping_features=context_mapping_features, + context_channels=context_channels[i + 1], + context_embedding_features=context_embedding_features, + num_layers=num_blocks[i], + factor=factors[i], + kernel_multiplier=kernel_multiplier_downsample, + num_groups=resnet_groups, + use_pre_downsample=True, + use_skip=True, + use_snake=use_snake, + num_transformer_blocks=attentions[i], + **attention_kwargs, + ) + for i in range(num_layers) + ] + ) + + self.bottleneck = BottleneckBlock1d( + channels=channels * multipliers[-1], + context_mapping_features=context_mapping_features, + context_embedding_features=context_embedding_features, + num_groups=resnet_groups, + num_transformer_blocks=attentions[-1], + use_snake=use_snake, + **attention_kwargs, + ) + + self.upsamples = nn.ModuleList( + [ + UpsampleBlock1d( + in_channels=channels * multipliers[i + 1], + out_channels=channels * multipliers[i], + context_mapping_features=context_mapping_features, + context_embedding_features=context_embedding_features, + num_layers=num_blocks[i] + (1 if attentions[i] else 0), + factor=factors[i], + use_nearest=use_nearest_upsample, + num_groups=resnet_groups, + use_skip_scale=use_skip_scale, + use_pre_upsample=False, + use_skip=True, + use_snake=use_snake, + skip_channels=channels * multipliers[i + 1], + num_transformer_blocks=attentions[i], + **attention_kwargs, + ) + for i in reversed(range(num_layers)) + ] + ) + + self.to_out = Unpatcher( + in_channels=channels * multipliers[0], + out_channels=out_channels, + patch_size=patch_size, + context_mapping_features=context_mapping_features, + use_snake=use_snake, + ) + + def get_channels( + self, channels_list: Optional[Sequence[Tensor]] = None, layer: int = 0 + ) -> Optional[Tensor]: + """Gets context channels at `layer` and checks that shape is correct""" + use_context_channels = self.use_context_channels and self.has_context[layer] + if not use_context_channels: + return None + assert exists(channels_list), "Missing context" + # Get channels index (skipping zero channel contexts) + channels_id = self.channels_ids[layer] + # Get channels + channels = channels_list[channels_id] + message = f"Missing context for layer {layer} at index {channels_id}" + assert exists(channels), message + # Check channels + num_channels = self.context_channels[layer] + message = f"Expected context with {num_channels} channels at idx {channels_id}" + assert channels.shape[1] == num_channels, message + # STFT channels if requested + channels = self.stft.encode1d(channels) if self.use_stft_context else channels # type: ignore # noqa + return channels + + def get_mapping( + self, time: Optional[Tensor] = None, features: Optional[Tensor] = None + ) -> Optional[Tensor]: + """Combines context time features and features into mapping""" + items, mapping = [], None + # Compute time features + if self.use_context_time: + assert_message = "use_context_time=True but no time features provided" + assert exists(time), assert_message + items += [self.to_time(time)] + # Compute features + if self.use_context_features: + assert_message = "context_features exists but no features provided" + assert exists(features), assert_message + items += [self.to_features(features)] + # Compute joint mapping + if self.use_context_time or self.use_context_features: + mapping = reduce(torch.stack(items), "n b m -> b m", "sum") + mapping = self.to_mapping(mapping) + return mapping + + def forward( + self, + x: Tensor, + time: Optional[Tensor] = None, + *, + features: Optional[Tensor] = None, + channels_list: Optional[Sequence[Tensor]] = None, + embedding: Optional[Tensor] = None, + embedding_mask: Optional[Tensor] = None, + causal: Optional[bool] = False, + ) -> Tensor: + channels = self.get_channels(channels_list, layer=0) + # Apply stft if required + x = self.stft.encode1d(x) if self.use_stft else x # type: ignore + # Concat context channels at layer 0 if provided + x = torch.cat([x, channels], dim=1) if exists(channels) else x + # Compute mapping from time and features + mapping = self.get_mapping(time, features) + x = self.to_in(x, mapping, causal=causal) + skips_list = [x] + + for i, downsample in enumerate(self.downsamples): + channels = self.get_channels(channels_list, layer=i + 1) + x, skips = downsample( + x, + mapping=mapping, + channels=channels, + embedding=embedding, + embedding_mask=embedding_mask, + causal=causal, + ) + skips_list += [skips] + + x = self.bottleneck( + x, + mapping=mapping, + embedding=embedding, + embedding_mask=embedding_mask, + causal=causal, + ) + + for i, upsample in enumerate(self.upsamples): + skips = skips_list.pop() + x = upsample( + x, + skips=skips, + mapping=mapping, + embedding=embedding, + embedding_mask=embedding_mask, + causal=causal, + ) + + x += skips_list.pop() + x = self.to_out(x, mapping, causal=causal) + x = self.stft.decode1d(x) if self.use_stft else x + + return x + + +""" Conditioning Modules """ + + +class FixedEmbedding(nn.Module): + def __init__(self, max_length: int, features: int): + super().__init__() + self.max_length = max_length + self.embedding = nn.Embedding(max_length, features) + + def forward(self, x: Tensor) -> Tensor: + batch_size, length, device = *x.shape[0:2], x.device + assert_message = "Input sequence length must be <= max_length" + assert length <= self.max_length, assert_message + position = torch.arange(length, device=device) + fixed_embedding = self.embedding(position) + fixed_embedding = repeat(fixed_embedding, "n d -> b n d", b=batch_size) + return fixed_embedding + + +def rand_bool(shape: Any, proba: float, device: Any = None) -> Tensor: + if proba == 1: + return torch.ones(shape, device=device, dtype=torch.bool) + elif proba == 0: + return torch.zeros(shape, device=device, dtype=torch.bool) + else: + return torch.bernoulli(torch.full(shape, proba, device=device)).to(torch.bool) + + +class UNetCFG1d(UNet1d): + """UNet1d with Classifier-Free Guidance""" + + def __init__( + self, + context_embedding_max_length: int, + context_embedding_features: int, + use_xattn_time: bool = False, + **kwargs, + ): + super().__init__( + context_embedding_features=context_embedding_features, **kwargs + ) + + self.use_xattn_time = use_xattn_time + + if use_xattn_time: + assert exists(context_embedding_features) + self.to_time_embedding = nn.Sequential( + TimePositionalEmbedding( + dim=kwargs["channels"], out_features=context_embedding_features + ), + nn.GELU(), + ) + + context_embedding_max_length += 1 # Add one for time embedding + + self.fixed_embedding = FixedEmbedding( + max_length=context_embedding_max_length, features=context_embedding_features + ) + + def forward( # type: ignore + self, + x: Tensor, + time: Tensor, + *, + embedding: Tensor, + embedding_mask: Optional[Tensor] = None, + embedding_scale: float = 1.0, + embedding_mask_proba: float = 0.0, + batch_cfg: bool = False, + rescale_cfg: bool = False, + scale_phi: float = 0.4, + negative_embedding: Optional[Tensor] = None, + negative_embedding_mask: Optional[Tensor] = None, + **kwargs, + ) -> Tensor: + b, device = embedding.shape[0], embedding.device + + if self.use_xattn_time: + embedding = torch.cat( + [embedding, self.to_time_embedding(time).unsqueeze(1)], dim=1 + ) + + if embedding_mask is not None: + embedding_mask = torch.cat( + [embedding_mask, torch.ones((b, 1), device=device)], dim=1 + ) + + fixed_embedding = self.fixed_embedding(embedding) + + if embedding_mask_proba > 0.0: + # Randomly mask embedding + batch_mask = rand_bool( + shape=(b, 1, 1), proba=embedding_mask_proba, device=device + ) + embedding = torch.where(batch_mask, fixed_embedding, embedding) + + if embedding_scale != 1.0: + if batch_cfg: + batch_x = torch.cat([x, x], dim=0) + batch_time = torch.cat([time, time], dim=0) + + if negative_embedding is not None: + if negative_embedding_mask is not None: + negative_embedding_mask = negative_embedding_mask.to( + torch.bool + ).unsqueeze(2) + + negative_embedding = torch.where( + negative_embedding_mask, negative_embedding, fixed_embedding + ) + + batch_embed = torch.cat([embedding, negative_embedding], dim=0) + + else: + batch_embed = torch.cat([embedding, fixed_embedding], dim=0) + + batch_mask = None + if embedding_mask is not None: + batch_mask = torch.cat([embedding_mask, embedding_mask], dim=0) + + batch_features = None + features = kwargs.pop("features", None) + if self.use_context_features: + batch_features = torch.cat([features, features], dim=0) + + batch_channels = None + channels_list = kwargs.pop("channels_list", None) + if self.use_context_channels: + batch_channels = [] + for channels in channels_list: + batch_channels += [torch.cat([channels, channels], dim=0)] + + # Compute both normal and fixed embedding outputs + batch_out = super().forward( + batch_x, + batch_time, + embedding=batch_embed, + embedding_mask=batch_mask, + features=batch_features, + channels_list=batch_channels, + **kwargs, + ) + out, out_masked = batch_out.chunk(2, dim=0) + + else: + # Compute both normal and fixed embedding outputs + out = super().forward( + x, + time, + embedding=embedding, + embedding_mask=embedding_mask, + **kwargs, + ) + out_masked = super().forward( + x, + time, + embedding=fixed_embedding, + embedding_mask=embedding_mask, + **kwargs, + ) + + out_cfg = out_masked + (out - out_masked) * embedding_scale + + if rescale_cfg: + out_std = out.std(dim=1, keepdim=True) + out_cfg_std = out_cfg.std(dim=1, keepdim=True) + + return ( + scale_phi * (out_cfg * (out_std / out_cfg_std)) + + (1 - scale_phi) * out_cfg + ) + + else: + return out_cfg + + else: + return super().forward( + x, time, embedding=embedding, embedding_mask=embedding_mask, **kwargs + ) + + +class UNetNCCA1d(UNet1d): + """UNet1d with Noise Channel Conditioning Augmentation""" + + def __init__(self, context_features: int, **kwargs): + super().__init__(context_features=context_features, **kwargs) + self.embedder = NumberEmbedder(features=context_features) + + def expand(self, x: Any, shape: Tuple[int, ...]) -> Tensor: + x = x if torch.is_tensor(x) else torch.tensor(x) + return x.expand(shape) + + def forward( # type: ignore + self, + x: Tensor, + time: Tensor, + *, + channels_list: Sequence[Tensor], + channels_augmentation: Union[ + bool, Sequence[bool], Sequence[Sequence[bool]], Tensor + ] = False, + channels_scale: Union[ + float, Sequence[float], Sequence[Sequence[float]], Tensor + ] = 0, + **kwargs, + ) -> Tensor: + b, n = x.shape[0], len(channels_list) + channels_augmentation = self.expand(channels_augmentation, shape=(b, n)).to(x) + channels_scale = self.expand(channels_scale, shape=(b, n)).to(x) + + # Augmentation (for each channel list item) + for i in range(n): + scale = channels_scale[:, i] * channels_augmentation[:, i] + scale = rearrange(scale, "b -> b 1 1") + item = channels_list[i] + channels_list[i] = torch.randn_like(item) * scale + item * (1 - scale) # type: ignore # noqa + + # Scale embedding (sum reduction if more than one channel list item) + channels_scale_emb = self.embedder(channels_scale) + channels_scale_emb = reduce(channels_scale_emb, "b n d -> b d", "sum") + + return super().forward( + x=x, + time=time, + channels_list=channels_list, + features=channels_scale_emb, + **kwargs, + ) + + +class UNetAll1d(UNetCFG1d, UNetNCCA1d): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, *args, **kwargs): # type: ignore + return UNetCFG1d.forward(self, *args, **kwargs) + + +def XUNet1d(type: str = "base", **kwargs) -> UNet1d: + if type == "base": + return UNet1d(**kwargs) + elif type == "all": + return UNetAll1d(**kwargs) + elif type == "cfg": + return UNetCFG1d(**kwargs) + elif type == "ncca": + return UNetNCCA1d(**kwargs) + else: + raise ValueError(f"Unknown XUNet1d type: {type}") + + +class NumberEmbedder(nn.Module): + def __init__( + self, + features: int, + dim: int = 256, + ): + super().__init__() + self.features = features + self.embedding = TimePositionalEmbedding(dim=dim, out_features=features) + + def forward(self, x: Union[List[float], Tensor]) -> Tensor: + if not torch.is_tensor(x): + device = next(self.embedding.parameters()).device + x = torch.tensor(x, device=device) + assert isinstance(x, Tensor) + shape = x.shape + x = rearrange(x, "... -> (...)") + embedding = self.embedding(x) + x = embedding.view(*shape, self.features) + return x # type: ignore + + +""" +Audio Transforms +""" + + +class STFT(nn.Module): + """Helper for torch stft and istft""" + + def __init__( + self, + num_fft: int = 1023, + hop_length: int = 256, + window_length: Optional[int] = None, + length: Optional[int] = None, + use_complex: bool = False, + ): + super().__init__() + self.num_fft = num_fft + self.hop_length = default(hop_length, floor(num_fft // 4)) + self.window_length = default(window_length, num_fft) + self.length = length + self.register_buffer("window", torch.hann_window(self.window_length)) + self.use_complex = use_complex + + def encode(self, wave: Tensor) -> Tuple[Tensor, Tensor]: + b = wave.shape[0] + wave = rearrange(wave, "b c t -> (b c) t") + + stft = torch.stft( + wave, + n_fft=self.num_fft, + hop_length=self.hop_length, + win_length=self.window_length, + window=self.window, # type: ignore + return_complex=True, + normalized=True, + ) + + if self.use_complex: + # Returns real and imaginary + stft_a, stft_b = stft.real, stft.imag + else: + # Returns magnitude and phase matrices + magnitude, phase = torch.abs(stft), torch.angle(stft) + stft_a, stft_b = magnitude, phase + + return rearrange_many((stft_a, stft_b), "(b c) f l -> b c f l", b=b) + + def decode(self, stft_a: Tensor, stft_b: Tensor) -> Tensor: + b, l = stft_a.shape[0], stft_a.shape[-1] # noqa + length = closest_power_2(l * self.hop_length) + + stft_a, stft_b = rearrange_many((stft_a, stft_b), "b c f l -> (b c) f l") + + if self.use_complex: + real, imag = stft_a, stft_b + else: + magnitude, phase = stft_a, stft_b + real, imag = magnitude * torch.cos(phase), magnitude * torch.sin(phase) + + stft = torch.stack([real, imag], dim=-1) + + wave = torch.istft( + stft, + n_fft=self.num_fft, + hop_length=self.hop_length, + win_length=self.window_length, + window=self.window, # type: ignore + length=default(self.length, length), + normalized=True, + ) + + return rearrange(wave, "(b c) t -> b c t", b=b) + + def encode1d( + self, wave: Tensor, stacked: bool = True + ) -> Union[Tensor, Tuple[Tensor, Tensor]]: + stft_a, stft_b = self.encode(wave) + stft_a, stft_b = rearrange_many((stft_a, stft_b), "b c f l -> b (c f) l") + return torch.cat((stft_a, stft_b), dim=1) if stacked else (stft_a, stft_b) + + def decode1d(self, stft_pair: Tensor) -> Tensor: + f = self.num_fft // 2 + 1 + stft_a, stft_b = stft_pair.chunk(chunks=2, dim=1) + stft_a, stft_b = rearrange_many((stft_a, stft_b), "b (c f) l -> b c f l", f=f) + return self.decode(stft_a, stft_b) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/autoencoders.py b/src/YingMusicSinger/utils/stable_audio_tools/autoencoders.py new file mode 100755 index 0000000000000000000000000000000000000000..d4c36f093a866a66708a9e4fe778935665e60857 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/autoencoders.py @@ -0,0 +1,975 @@ +import math +from typing import Any, Dict, Literal + +import numpy as np +import torch +from alias_free_torch import Activation1d +from dac.nn.layers import WNConv1d, WNConvTranspose1d +from torch import nn +from torch.nn import functional as F +from torchaudio import transforms as T + +# from ..inference.sampling import sample +# from ..inference.utils import prepare_audio +from .blocks import SnakeBeta +from .bottleneck import Bottleneck, DiscreteBottleneck +from .diffusion import ( + ConditionedDiffusionModel, + DAU1DCondWrapper, + DiTWrapper, + UNet1DCondWrapper, +) +from .factory import create_bottleneck_from_config, create_pretransform_from_config +from .pretransforms import Pretransform + + +def checkpoint(function, *args, **kwargs): + kwargs.setdefault("use_reentrant", False) + return torch.utils.checkpoint.checkpoint(function, *args, **kwargs) + + +def get_activation( + activation: Literal["elu", "snake", "none"], antialias=False, channels=None +) -> nn.Module: + if activation == "elu": + act = nn.ELU() + elif activation == "snake": + act = SnakeBeta(channels) + elif activation == "none": + act = nn.Identity() + else: + raise ValueError(f"Unknown activation {activation}") + + if antialias: + act = Activation1d(act) + + return act + + +class ResidualUnit(nn.Module): + def __init__( + self, + in_channels, + out_channels, + dilation, + use_snake=False, + antialias_activation=False, + ): + super().__init__() + + self.dilation = dilation + + padding = (dilation * (7 - 1)) // 2 + + self.layers = nn.Sequential( + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=out_channels, + ), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=7, + dilation=dilation, + padding=padding, + ), + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=out_channels, + ), + WNConv1d( + in_channels=out_channels, out_channels=out_channels, kernel_size=1 + ), + ) + + def forward(self, x): + res = x + + # x = checkpoint(self.layers, x) + x = self.layers(x) + + return x + res + + +class EncoderBlock(nn.Module): + def __init__( + self, + in_channels, + out_channels, + stride, + use_snake=False, + antialias_activation=False, + ): + super().__init__() + + self.layers = nn.Sequential( + ResidualUnit( + in_channels=in_channels, + out_channels=in_channels, + dilation=1, + use_snake=use_snake, + ), + ResidualUnit( + in_channels=in_channels, + out_channels=in_channels, + dilation=3, + use_snake=use_snake, + ), + ResidualUnit( + in_channels=in_channels, + out_channels=in_channels, + dilation=9, + use_snake=use_snake, + ), + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=in_channels, + ), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ) + + def forward(self, x): + return self.layers(x) + + +class DecoderBlock(nn.Module): + def __init__( + self, + in_channels, + out_channels, + stride, + use_snake=False, + antialias_activation=False, + use_nearest_upsample=False, + ): + super().__init__() + + if use_nearest_upsample: + upsample_layer = nn.Sequential( + nn.Upsample(scale_factor=stride, mode="nearest"), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=1, + bias=False, + padding="same", + ), + ) + else: + upsample_layer = WNConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + + self.layers = nn.Sequential( + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=in_channels, + ), + upsample_layer, + ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=1, + use_snake=use_snake, + ), + ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=3, + use_snake=use_snake, + ), + ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=9, + use_snake=use_snake, + ), + ) + + def forward(self, x): + return self.layers(x) + + +class OobleckEncoder(nn.Module): + def __init__( + self, + in_channels=2, + channels=128, + latent_dim=32, + c_mults=[1, 2, 4, 8], + strides=[2, 4, 8, 8], + use_snake=False, + antialias_activation=False, + ): + super().__init__() + + c_mults = [1] + c_mults + + self.depth = len(c_mults) + + layers = [ + WNConv1d( + in_channels=in_channels, + out_channels=c_mults[0] * channels, + kernel_size=7, + padding=3, + ) + ] + + for i in range(self.depth - 1): + layers += [ + EncoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i + 1] * channels, + stride=strides[i], + use_snake=use_snake, + ) + ] + + layers += [ + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=c_mults[-1] * channels, + ), + WNConv1d( + in_channels=c_mults[-1] * channels, + out_channels=latent_dim, + kernel_size=3, + padding=1, + ), + ] + + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class OobleckDecoder(nn.Module): + def __init__( + self, + out_channels=2, + channels=128, + latent_dim=32, + c_mults=[1, 2, 4, 8], + strides=[2, 4, 8, 8], + use_snake=False, + antialias_activation=False, + use_nearest_upsample=False, + final_tanh=True, + ): + super().__init__() + + c_mults = [1] + c_mults + + self.depth = len(c_mults) + + layers = [ + WNConv1d( + in_channels=latent_dim, + out_channels=c_mults[-1] * channels, + kernel_size=7, + padding=3, + ), + ] + + for i in range(self.depth - 1, 0, -1): + layers += [ + DecoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i - 1] * channels, + stride=strides[i - 1], + use_snake=use_snake, + antialias_activation=antialias_activation, + use_nearest_upsample=use_nearest_upsample, + ) + ] + + layers += [ + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=c_mults[0] * channels, + ), + WNConv1d( + in_channels=c_mults[0] * channels, + out_channels=out_channels, + kernel_size=7, + padding=3, + bias=False, + ), + nn.Tanh() if final_tanh else nn.Identity(), + ] + + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class DACEncoderWrapper(nn.Module): + def __init__(self, in_channels=1, **kwargs): + super().__init__() + + from dac.model.dac import Encoder as DACEncoder + + latent_dim = kwargs.pop("latent_dim", None) + + encoder_out_dim = kwargs["d_model"] * (2 ** len(kwargs["strides"])) + self.encoder = DACEncoder(d_latent=encoder_out_dim, **kwargs) + self.latent_dim = latent_dim + + # Latent-dim support was added to DAC after this was first written, and implemented differently, so this is for backwards compatibility + self.proj_out = ( + nn.Conv1d(self.encoder.enc_dim, latent_dim, kernel_size=1) + if latent_dim is not None + else nn.Identity() + ) + + if in_channels != 1: + self.encoder.block[0] = WNConv1d( + in_channels, kwargs.get("d_model", 64), kernel_size=7, padding=3 + ) + + def forward(self, x): + x = self.encoder(x) + x = self.proj_out(x) + return x + + +class DACDecoderWrapper(nn.Module): + def __init__(self, latent_dim, out_channels=1, **kwargs): + super().__init__() + + from dac.model.dac import Decoder as DACDecoder + + self.decoder = DACDecoder( + **kwargs, input_channel=latent_dim, d_out=out_channels + ) + + self.latent_dim = latent_dim + + def forward(self, x): + return self.decoder(x) + + +class AudioAutoencoder(nn.Module): + def __init__( + self, + encoder, + decoder, + latent_dim, + downsampling_ratio, + sample_rate, + io_channels=2, + bottleneck: Bottleneck = None, + pretransform: Pretransform = None, + in_channels=None, + out_channels=None, + soft_clip=False, + ): + super().__init__() + + self.downsampling_ratio = downsampling_ratio + self.sample_rate = sample_rate + + self.latent_dim = latent_dim + self.io_channels = io_channels + self.in_channels = io_channels + self.out_channels = io_channels + + self.min_length = self.downsampling_ratio + + if in_channels is not None: + self.in_channels = in_channels + + if out_channels is not None: + self.out_channels = out_channels + + self.bottleneck = bottleneck + + self.encoder = encoder + + self.decoder = decoder + + self.pretransform = pretransform + + self.soft_clip = soft_clip + + self.is_discrete = self.bottleneck is not None and self.bottleneck.is_discrete + + def encode( + self, + audio, + return_info=False, + skip_pretransform=False, + iterate_batch=False, + **kwargs, + ): + info = {} + + if self.pretransform is not None and not skip_pretransform: + if self.pretransform.enable_grad: + if iterate_batch: + audios = [] + for i in range(audio.shape[0]): + audios.append(self.pretransform.encode(audio[i : i + 1])) + audio = torch.cat(audios, dim=0) + else: + audio = self.pretransform.encode(audio) + else: + with torch.no_grad(): + if iterate_batch: + audios = [] + for i in range(audio.shape[0]): + audios.append(self.pretransform.encode(audio[i : i + 1])) + audio = torch.cat(audios, dim=0) + else: + audio = self.pretransform.encode(audio) + + if self.encoder is not None: + if iterate_batch: + latents = [] + for i in range(audio.shape[0]): + latents.append(self.encoder(audio[i : i + 1])) + latents = torch.cat(latents, dim=0) + else: + latents = self.encoder(audio) + else: + latents = audio + + if self.bottleneck is not None: + # TODO: Add iterate batch logic, needs to merge the info dicts + latents, bottleneck_info = self.bottleneck.encode( + latents, return_info=True, **kwargs + ) + + info.update(bottleneck_info) + + if return_info: + return latents, info + + return latents + + def decode(self, latents, iterate_batch=False, **kwargs): + if self.bottleneck is not None: + if iterate_batch: + decoded = [] + for i in range(latents.shape[0]): + decoded.append(self.bottleneck.decode(latents[i : i + 1])) + latents = torch.cat(decoded, dim=0) + else: + latents = self.bottleneck.decode(latents) + + if iterate_batch: + decoded = [] + for i in range(latents.shape[0]): + decoded.append(self.decoder(latents[i : i + 1])) + decoded = torch.cat(decoded, dim=0) + else: + decoded = self.decoder(latents, **kwargs) + + if self.pretransform is not None: + if self.pretransform.enable_grad: + if iterate_batch: + decodeds = [] + for i in range(decoded.shape[0]): + decodeds.append(self.pretransform.decode(decoded[i : i + 1])) + decoded = torch.cat(decodeds, dim=0) + else: + decoded = self.pretransform.decode(decoded) + else: + with torch.no_grad(): + if iterate_batch: + decodeds = [] + for i in range(latents.shape[0]): + decodeds.append( + self.pretransform.decode(decoded[i : i + 1]) + ) + decoded = torch.cat(decodeds, dim=0) + else: + decoded = self.pretransform.decode(decoded) + + if self.soft_clip: + decoded = torch.tanh(decoded) + + return decoded + + def decode_tokens(self, tokens, **kwargs): + """ + Decode discrete tokens to audio + Only works with discrete autoencoders + """ + + assert isinstance(self.bottleneck, DiscreteBottleneck), ( + "decode_tokens only works with discrete autoencoders" + ) + + latents = self.bottleneck.decode_tokens(tokens, **kwargs) + + return self.decode(latents, **kwargs) + + def preprocess_audio_for_encoder(self, audio, in_sr): + """ + Preprocess single audio tensor (Channels x Length) to be compatible with the encoder. + If the model is mono, stereo audio will be converted to mono. + Audio will be silence-padded to be a multiple of the model's downsampling ratio. + Audio will be resampled to the model's sample rate. + The output will have batch size 1 and be shape (1 x Channels x Length) + """ + return self.preprocess_audio_list_for_encoder([audio], [in_sr]) + + def preprocess_audio_list_for_encoder(self, audio_list, in_sr_list): + """ + Preprocess a [list] of audio (Channels x Length) into a batch tensor to be compatable with the encoder. + The audio in that list can be of different lengths and channels. + in_sr can be an integer or list. If it's an integer it will be assumed it is the input sample_rate for every audio. + All audio will be resampled to the model's sample rate. + Audio will be silence-padded to the longest length, and further padded to be a multiple of the model's downsampling ratio. + If the model is mono, all audio will be converted to mono. + The output will be a tensor of shape (Batch x Channels x Length) + """ + batch_size = len(audio_list) + if isinstance(in_sr_list, int): + in_sr_list = [in_sr_list] * batch_size + assert len(in_sr_list) == batch_size, ( + "list of sample rates must be the same length of audio_list" + ) + new_audio = [] + max_length = 0 + # resample & find the max length + for i in range(batch_size): + audio = audio_list[i] + in_sr = in_sr_list[i] + if len(audio.shape) == 3 and audio.shape[0] == 1: + # batchsize 1 was given by accident. Just squeeze it. + audio = audio.squeeze(0) + elif len(audio.shape) == 1: + # Mono signal, channel dimension is missing, unsqueeze it in + audio = audio.unsqueeze(0) + assert len(audio.shape) == 2, ( + "Audio should be shape (Channels x Length) with no batch dimension" + ) + # Resample audio + if in_sr != self.sample_rate: + resample_tf = T.Resample(in_sr, self.sample_rate).to(audio.device) + audio = resample_tf(audio) + new_audio.append(audio) + if audio.shape[-1] > max_length: + max_length = audio.shape[-1] + # Pad every audio to the same length, multiple of model's downsampling ratio + padded_audio_length = ( + max_length + + (self.min_length - (max_length % self.min_length)) % self.min_length + ) + for i in range(batch_size): + # Pad it & if necessary, mixdown/duplicate stereo/mono channels to support model + new_audio[i] = prepare_audio( + new_audio[i], + in_sr=in_sr, + target_sr=in_sr, + target_length=padded_audio_length, + target_channels=self.in_channels, + device=new_audio[i].device, + ).squeeze(0) + # convert to tensor + return torch.stack(new_audio) + + def encode_audio(self, audio, chunked=False, overlap=32, chunk_size=128, **kwargs): + """ + Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder. + If chunked is True, split the audio into chunks of a given maximum size chunk_size, with given overlap. + Overlap and chunk_size params are both measured in number of latents (not audio samples) + # and therefore you likely could use the same values with decode_audio. + A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. + Every autoencoder will have a different receptive field size, and thus ideal overlap. + You can determine it empirically by diffing unchunked vs chunked output and looking at maximum diff. + The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks. + Smaller chunk_size uses less memory, but more compute. + The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version + For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks + """ + if not chunked: + # default behavior. Encode the entire audio in parallel + return self.encode(audio, **kwargs) + else: + # CHUNKED ENCODING + # samples_per_latent is just the downsampling ratio (which is also the upsampling ratio) + samples_per_latent = self.downsampling_ratio + total_size = audio.shape[2] # in samples + batch_size = audio.shape[0] + chunk_size *= samples_per_latent # converting metric in latents to samples + overlap *= samples_per_latent # converting metric in latents to samples + hop_size = chunk_size - overlap + chunks = [] + for i in range(0, total_size - chunk_size + 1, hop_size): + chunk = audio[:, :, i : i + chunk_size] + chunks.append(chunk) + if i + chunk_size != total_size: + # Final chunk + chunk = audio[:, :, -chunk_size:] + chunks.append(chunk) + chunks = torch.stack(chunks) + num_chunks = chunks.shape[0] + # Note: y_size might be a different value from the latent length used in diffusion training + # because we can encode audio of varying lengths + # However, the audio should've been padded to a multiple of samples_per_latent by now. + y_size = total_size // samples_per_latent + # Create an empty latent, we will populate it with chunks as we encode them + y_final = torch.zeros((batch_size, self.latent_dim, y_size)).to( + audio.device + ) + for i in range(num_chunks): + x_chunk = chunks[i, :] + # encode the chunk + y_chunk = self.encode(x_chunk) + # figure out where to put the audio along the time domain + if i == num_chunks - 1: + # final chunk always goes at the end + t_end = y_size + t_start = t_end - y_chunk.shape[2] + else: + t_start = i * hop_size // samples_per_latent + t_end = t_start + chunk_size // samples_per_latent + # remove the edges of the overlaps + ol = overlap // samples_per_latent // 2 + chunk_start = 0 + chunk_end = y_chunk.shape[2] + if i > 0: + # no overlap for the start of the first chunk + t_start += ol + chunk_start += ol + if i < num_chunks - 1: + # no overlap for the end of the last chunk + t_end -= ol + chunk_end -= ol + # paste the chunked audio into our y_final output audio + y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end] + return y_final + + def decode_audio( + self, latents, chunked=False, overlap=32, chunk_size=128, **kwargs + ): + """ + Decode latents to audio. + If chunked is True, split the latents into chunks of a given maximum size chunk_size, with given overlap, both of which are measured in number of latents. + A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. + Every autoencoder will have a different receptive field size, and thus ideal overlap. + You can determine it empirically by diffing unchunked vs chunked audio and looking at maximum diff. + The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks. + Smaller chunk_size uses less memory, but more compute. + The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version + For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks + """ + if not chunked: + # default behavior. Decode the entire latent in parallel + return self.decode(latents, **kwargs) + else: + # chunked decoding + hop_size = chunk_size - overlap + total_size = latents.shape[2] + batch_size = latents.shape[0] + chunks = [] + if total_size < chunk_size: + # pad the latents to be at least chunk_size + # 如果在这里pad之后，那么之后的生成歌曲就变噪音了 + pad_size = chunk_size - total_size + 1 + latents = F.pad(latents, (0, pad_size), mode="replicate") + total_size = latents.shape[2] + # import pdb; pdb.set_trace() + for i in range(0, total_size - chunk_size + 1, hop_size): + chunk = latents[:, :, i : i + chunk_size] + chunks.append(chunk) + if i + chunk_size != total_size: + # Final chunk + chunk = latents[:, :, -chunk_size:] + chunks.append(chunk) + chunks = torch.stack(chunks) + num_chunks = chunks.shape[0] + # samples_per_latent is just the downsampling ratio + samples_per_latent = self.downsampling_ratio + # Create an empty waveform, we will populate it with chunks as decode them + y_size = total_size * samples_per_latent + y_final = torch.zeros((batch_size, self.out_channels, y_size)).to( + latents.device + ) + for i in range(num_chunks): + x_chunk = chunks[i, :] + # decode the chunk + y_chunk = self.decode(x_chunk) + # figure out where to put the audio along the time domain + if i == num_chunks - 1: + # final chunk always goes at the end + t_end = y_size + t_start = t_end - y_chunk.shape[2] + else: + t_start = i * hop_size * samples_per_latent + t_end = t_start + chunk_size * samples_per_latent + # remove the edges of the overlaps + ol = (overlap // 2) * samples_per_latent + chunk_start = 0 + chunk_end = y_chunk.shape[2] + if i > 0: + # no overlap for the start of the first chunk + t_start += ol + chunk_start += ol + if i < num_chunks - 1: + # no overlap for the end of the last chunk + t_end -= ol + chunk_end -= ol + # paste the chunked audio into our y_final output audio + y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end] + return y_final + + +class DiffusionAutoencoder(AudioAutoencoder): + def __init__( + self, + diffusion: ConditionedDiffusionModel, + diffusion_downsampling_ratio, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + self.diffusion = diffusion + + self.min_length = self.downsampling_ratio * diffusion_downsampling_ratio + + if self.encoder is not None: + # Shrink the initial encoder parameters to avoid saturated latents + with torch.no_grad(): + for param in self.encoder.parameters(): + param *= 0.5 + + def decode(self, latents, steps=100): + upsampled_length = latents.shape[2] * self.downsampling_ratio + + if self.bottleneck is not None: + latents = self.bottleneck.decode(latents) + + if self.decoder is not None: + latents = self.decode(latents) + + # Upsample latents to match diffusion length + if latents.shape[2] != upsampled_length: + latents = F.interpolate(latents, size=upsampled_length, mode="nearest") + + noise = torch.randn( + latents.shape[0], self.io_channels, upsampled_length, device=latents.device + ) + decoded = sample(self.diffusion, noise, steps, 0, input_concat_cond=latents) + + if self.pretransform is not None: + if self.pretransform.enable_grad: + decoded = self.pretransform.decode(decoded) + else: + with torch.no_grad(): + decoded = self.pretransform.decode(decoded) + + return decoded + + +# AE factories + + +def create_encoder_from_config(encoder_config: Dict[str, Any]): + encoder_type = encoder_config.get("type", None) + assert encoder_type is not None, "Encoder type must be specified" + + if encoder_type == "oobleck": + encoder = OobleckEncoder(**encoder_config["config"]) + + elif encoder_type == "seanet": + from encodec.modules import SEANetEncoder + + seanet_encoder_config = encoder_config["config"] + + # SEANet encoder expects strides in reverse order + seanet_encoder_config["ratios"] = list( + reversed(seanet_encoder_config.get("ratios", [2, 2, 2, 2, 2])) + ) + encoder = SEANetEncoder(**seanet_encoder_config) + elif encoder_type == "dac": + dac_config = encoder_config["config"] + + encoder = DACEncoderWrapper(**dac_config) + elif encoder_type == "local_attn": + from .local_attention import TransformerEncoder1D + + local_attn_config = encoder_config["config"] + + encoder = TransformerEncoder1D(**local_attn_config) + else: + raise ValueError(f"Unknown encoder type {encoder_type}") + + requires_grad = encoder_config.get("requires_grad", True) + if not requires_grad: + for param in encoder.parameters(): + param.requires_grad = False + + return encoder + + +def create_decoder_from_config(decoder_config: Dict[str, Any]): + decoder_type = decoder_config.get("type", None) + assert decoder_type is not None, "Decoder type must be specified" + + if decoder_type == "oobleck": + decoder = OobleckDecoder(**decoder_config["config"]) + elif decoder_type == "seanet": + from encodec.modules import SEANetDecoder + + decoder = SEANetDecoder(**decoder_config["config"]) + elif decoder_type == "dac": + dac_config = decoder_config["config"] + + decoder = DACDecoderWrapper(**dac_config) + elif decoder_type == "local_attn": + from .local_attention import TransformerDecoder1D + + local_attn_config = decoder_config["config"] + + decoder = TransformerDecoder1D(**local_attn_config) + else: + raise ValueError(f"Unknown decoder type {decoder_type}") + + requires_grad = decoder_config.get("requires_grad", True) + if not requires_grad: + for param in decoder.parameters(): + param.requires_grad = False + + return decoder + + +def create_autoencoder_from_config(config: Dict[str, Any]): + ae_config = config["model"] + + encoder = create_encoder_from_config(ae_config["encoder"]) + decoder = create_decoder_from_config(ae_config["decoder"]) + + bottleneck = ae_config.get("bottleneck", None) + + latent_dim = ae_config.get("latent_dim", None) + assert latent_dim is not None, "latent_dim must be specified in model config" + downsampling_ratio = ae_config.get("downsampling_ratio", None) + assert downsampling_ratio is not None, ( + "downsampling_ratio must be specified in model config" + ) + io_channels = ae_config.get("io_channels", None) + assert io_channels is not None, "io_channels must be specified in model config" + sample_rate = config.get("sample_rate", None) + assert sample_rate is not None, "sample_rate must be specified in model config" + + in_channels = ae_config.get("in_channels", None) + out_channels = ae_config.get("out_channels", None) + + pretransform = ae_config.get("pretransform", None) + + if pretransform is not None: + pretransform = create_pretransform_from_config(pretransform, sample_rate) + + if bottleneck is not None: + bottleneck = create_bottleneck_from_config(bottleneck) + + soft_clip = ae_config["decoder"].get("soft_clip", False) + + return AudioAutoencoder( + encoder, + decoder, + io_channels=io_channels, + latent_dim=latent_dim, + downsampling_ratio=downsampling_ratio, + sample_rate=sample_rate, + bottleneck=bottleneck, + pretransform=pretransform, + in_channels=in_channels, + out_channels=out_channels, + soft_clip=soft_clip, + ) + + +def create_diffAE_from_config(config: Dict[str, Any]): + diffae_config = config["model"] + + if "encoder" in diffae_config: + encoder = create_encoder_from_config(diffae_config["encoder"]) + else: + encoder = None + + if "decoder" in diffae_config: + decoder = create_decoder_from_config(diffae_config["decoder"]) + else: + decoder = None + + diffusion_model_type = diffae_config["diffusion"]["type"] + + if diffusion_model_type == "DAU1d": + diffusion = DAU1DCondWrapper(**diffae_config["diffusion"]["config"]) + elif diffusion_model_type == "adp_1d": + diffusion = UNet1DCondWrapper(**diffae_config["diffusion"]["config"]) + elif diffusion_model_type == "dit": + diffusion = DiTWrapper(**diffae_config["diffusion"]["config"]) + + latent_dim = diffae_config.get("latent_dim", None) + assert latent_dim is not None, "latent_dim must be specified in model config" + downsampling_ratio = diffae_config.get("downsampling_ratio", None) + assert downsampling_ratio is not None, ( + "downsampling_ratio must be specified in model config" + ) + io_channels = diffae_config.get("io_channels", None) + assert io_channels is not None, "io_channels must be specified in model config" + sample_rate = config.get("sample_rate", None) + assert sample_rate is not None, "sample_rate must be specified in model config" + + bottleneck = diffae_config.get("bottleneck", None) + + pretransform = diffae_config.get("pretransform", None) + + if pretransform is not None: + pretransform = create_pretransform_from_config(pretransform, sample_rate) + + if bottleneck is not None: + bottleneck = create_bottleneck_from_config(bottleneck) + + diffusion_downsampling_ratio = (None,) + + if diffusion_model_type == "DAU1d": + diffusion_downsampling_ratio = np.prod( + diffae_config["diffusion"]["config"]["strides"] + ) + elif diffusion_model_type == "adp_1d": + diffusion_downsampling_ratio = np.prod( + diffae_config["diffusion"]["config"]["factors"] + ) + elif diffusion_model_type == "dit": + diffusion_downsampling_ratio = 1 + + return DiffusionAutoencoder( + encoder=encoder, + decoder=decoder, + diffusion=diffusion, + io_channels=io_channels, + sample_rate=sample_rate, + latent_dim=latent_dim, + downsampling_ratio=downsampling_ratio, + diffusion_downsampling_ratio=diffusion_downsampling_ratio, + bottleneck=bottleneck, + pretransform=pretransform, + ) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/blocks.py b/src/YingMusicSinger/utils/stable_audio_tools/blocks.py new file mode 100755 index 0000000000000000000000000000000000000000..b2b2c1e41d7d73dfc6fe5024aab2a52492a46070 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/blocks.py @@ -0,0 +1,398 @@ +import math +from functools import reduce + +import numpy as np +import torch +from dac.nn.layers import Snake1d +from packaging import version +from torch import nn +from torch.backends.cuda import sdp_kernel +from torch.nn import functional as F + + +class ResidualBlock(nn.Module): + def __init__(self, main, skip=None): + super().__init__() + self.main = nn.Sequential(*main) + self.skip = skip if skip else nn.Identity() + + def forward(self, input): + return self.main(input) + self.skip(input) + + +class ResConvBlock(ResidualBlock): + def __init__( + self, + c_in, + c_mid, + c_out, + is_last=False, + kernel_size=5, + conv_bias=True, + use_snake=False, + ): + skip = None if c_in == c_out else nn.Conv1d(c_in, c_out, 1, bias=False) + super().__init__( + [ + nn.Conv1d( + c_in, c_mid, kernel_size, padding=kernel_size // 2, bias=conv_bias + ), + nn.GroupNorm(1, c_mid), + Snake1d(c_mid) if use_snake else nn.GELU(), + nn.Conv1d( + c_mid, c_out, kernel_size, padding=kernel_size // 2, bias=conv_bias + ), + nn.GroupNorm(1, c_out) if not is_last else nn.Identity(), + (Snake1d(c_out) if use_snake else nn.GELU()) + if not is_last + else nn.Identity(), + ], + skip, + ) + + +class SelfAttention1d(nn.Module): + def __init__(self, c_in, n_head=1, dropout_rate=0.0): + super().__init__() + assert c_in % n_head == 0 + self.norm = nn.GroupNorm(1, c_in) + self.n_head = n_head + self.qkv_proj = nn.Conv1d(c_in, c_in * 3, 1) + self.out_proj = nn.Conv1d(c_in, c_in, 1) + self.dropout = nn.Dropout(dropout_rate, inplace=True) + + self.use_flash = torch.cuda.is_available() and version.parse( + torch.__version__ + ) >= version.parse("2.0.0") + + if not self.use_flash: + return + + device_properties = torch.cuda.get_device_properties(torch.device("cuda")) + + if device_properties.major == 8 and device_properties.minor == 0: + # Use flash attention for A100 GPUs + self.sdp_kernel_config = (True, False, False) + else: + # Don't use flash attention for other GPUs + self.sdp_kernel_config = (False, True, True) + + def forward(self, input): + n, c, s = input.shape + qkv = self.qkv_proj(self.norm(input)) + qkv = qkv.view([n, self.n_head * 3, c // self.n_head, s]).transpose(2, 3) + q, k, v = qkv.chunk(3, dim=1) + scale = k.shape[3] ** -0.25 + + if self.use_flash: + with sdp_kernel(*self.sdp_kernel_config): + y = ( + F.scaled_dot_product_attention(q, k, v, is_causal=False) + .contiguous() + .view([n, c, s]) + ) + else: + att = ((q * scale) @ (k.transpose(2, 3) * scale)).softmax(3) + y = (att @ v).transpose(2, 3).contiguous().view([n, c, s]) + + return input + self.dropout(self.out_proj(y)) + + +class SkipBlock(nn.Module): + def __init__(self, *main): + super().__init__() + self.main = nn.Sequential(*main) + + def forward(self, input): + return torch.cat([self.main(input), input], dim=1) + + +class FourierFeatures(nn.Module): + def __init__(self, in_features, out_features, std=1.0): + super().__init__() + assert out_features % 2 == 0 + self.weight = nn.Parameter(torch.randn([out_features // 2, in_features]) * std) + + def forward(self, input): + f = 2 * math.pi * input @ self.weight.T + return torch.cat([f.cos(), f.sin()], dim=-1) + + +def expand_to_planes(input, shape): + return input[..., None].repeat([1, 1, shape[2]]) + + +_kernels = { + "linear": [1 / 8, 3 / 8, 3 / 8, 1 / 8], + "cubic": [ + -0.01171875, + -0.03515625, + 0.11328125, + 0.43359375, + 0.43359375, + 0.11328125, + -0.03515625, + -0.01171875, + ], + "lanczos3": [ + 0.003689131001010537, + 0.015056144446134567, + -0.03399861603975296, + -0.066637322306633, + 0.13550527393817902, + 0.44638532400131226, + 0.44638532400131226, + 0.13550527393817902, + -0.066637322306633, + -0.03399861603975296, + 0.015056144446134567, + 0.003689131001010537, + ], +} + + +class Downsample1d(nn.Module): + def __init__(self, kernel="linear", pad_mode="reflect", channels_last=False): + super().__init__() + self.pad_mode = pad_mode + kernel_1d = torch.tensor(_kernels[kernel]) + self.pad = kernel_1d.shape[0] // 2 - 1 + self.register_buffer("kernel", kernel_1d) + self.channels_last = channels_last + + def forward(self, x): + if self.channels_last: + x = x.permute(0, 2, 1) + x = F.pad(x, (self.pad,) * 2, self.pad_mode) + weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0]]) + indices = torch.arange(x.shape[1], device=x.device) + weight[indices, indices] = self.kernel.to(weight) + x = F.conv1d(x, weight, stride=2) + if self.channels_last: + x = x.permute(0, 2, 1) + return x + + +class Upsample1d(nn.Module): + def __init__(self, kernel="linear", pad_mode="reflect", channels_last=False): + super().__init__() + self.pad_mode = pad_mode + kernel_1d = torch.tensor(_kernels[kernel]) * 2 + self.pad = kernel_1d.shape[0] // 2 - 1 + self.register_buffer("kernel", kernel_1d) + self.channels_last = channels_last + + def forward(self, x): + if self.channels_last: + x = x.permute(0, 2, 1) + x = F.pad(x, ((self.pad + 1) // 2,) * 2, self.pad_mode) + weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0]]) + indices = torch.arange(x.shape[1], device=x.device) + weight[indices, indices] = self.kernel.to(weight) + x = F.conv_transpose1d(x, weight, stride=2, padding=self.pad * 2 + 1) + if self.channels_last: + x = x.permute(0, 2, 1) + return x + + +def Downsample1d_2( + in_channels: int, out_channels: int, factor: int, kernel_multiplier: int = 2 +) -> nn.Module: + assert kernel_multiplier % 2 == 0, "Kernel multiplier must be even" + + return nn.Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=factor * kernel_multiplier + 1, + stride=factor, + padding=factor * (kernel_multiplier // 2), + ) + + +def Upsample1d_2( + in_channels: int, out_channels: int, factor: int, use_nearest: bool = False +) -> nn.Module: + if factor == 1: + return nn.Conv1d( + in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1 + ) + + if use_nearest: + return nn.Sequential( + nn.Upsample(scale_factor=factor, mode="nearest"), + nn.Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + padding=1, + ), + ) + else: + return nn.ConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=factor * 2, + stride=factor, + padding=factor // 2 + factor % 2, + output_padding=factor % 2, + ) + + +def zero_init(layer): + nn.init.zeros_(layer.weight) + if layer.bias is not None: + nn.init.zeros_(layer.bias) + return layer + + +def rms_norm(x, scale, eps): + dtype = reduce(torch.promote_types, (x.dtype, scale.dtype, torch.float32)) + mean_sq = torch.mean(x.to(dtype) ** 2, dim=-1, keepdim=True) + scale = scale.to(dtype) * torch.rsqrt(mean_sq + eps) + return x * scale.to(x.dtype) + + +# rms_norm = torch.compile(rms_norm) + + +class AdaRMSNorm(nn.Module): + def __init__(self, features, cond_features, eps=1e-6): + super().__init__() + self.eps = eps + self.linear = zero_init(nn.Linear(cond_features, features, bias=False)) + + def extra_repr(self): + return f"eps={self.eps}," + + def forward(self, x, cond): + return rms_norm(x, self.linear(cond)[:, None, :] + 1, self.eps) + + +def normalize(x, eps=1e-4): + dim = list(range(1, x.ndim)) + n = torch.linalg.vector_norm(x, dim=dim, keepdim=True) + alpha = np.sqrt(n.numel() / x.numel()) + return x / torch.add(eps, n, alpha=alpha) + + +class ForcedWNConv1d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size=1): + super().__init__() + self.weight = nn.Parameter( + torch.randn([out_channels, in_channels, kernel_size]) + ) + + def forward(self, x): + if self.training: + with torch.no_grad(): + self.weight.copy_(normalize(self.weight)) + + fan_in = self.weight[0].numel() + + w = normalize(self.weight) / math.sqrt(fan_in) + + return F.conv1d(x, w, padding="same") + + +# Kernels + +use_compile = True + + +def compile(function, *args, **kwargs): + if not use_compile: + return function + try: + return torch.compile(function, *args, **kwargs) + except RuntimeError: + return function + + +@compile +def linear_geglu(x, weight, bias=None): + x = x @ weight.mT + if bias is not None: + x = x + bias + x, gate = x.chunk(2, dim=-1) + return x * F.gelu(gate) + + +@compile +def rms_norm(x, scale, eps): + dtype = reduce(torch.promote_types, (x.dtype, scale.dtype, torch.float32)) + mean_sq = torch.mean(x.to(dtype) ** 2, dim=-1, keepdim=True) + scale = scale.to(dtype) * torch.rsqrt(mean_sq + eps) + return x * scale.to(x.dtype) + + +# Layers + + +class LinearGEGLU(nn.Linear): + def __init__(self, in_features, out_features, bias=True): + super().__init__(in_features, out_features * 2, bias=bias) + self.out_features = out_features + + def forward(self, x): + return linear_geglu(x, self.weight, self.bias) + + +class RMSNorm(nn.Module): + def __init__(self, shape, fix_scale=False, eps=1e-6): + super().__init__() + self.eps = eps + + if fix_scale: + self.register_buffer("scale", torch.ones(shape)) + else: + self.scale = nn.Parameter(torch.ones(shape)) + + def extra_repr(self): + return f"shape={tuple(self.scale.shape)}, eps={self.eps}" + + def forward(self, x): + return rms_norm(x, self.scale, self.eps) + + +def snake_beta(x, alpha, beta): + return x + (1.0 / (beta + 0.000000001)) * pow(torch.sin(x * alpha), 2) + + +# try: +# snake_beta = torch.compile(snake_beta) +# except RuntimeError: +# pass + + +# Adapted from https://github.com/NVIDIA/BigVGAN/blob/main/activations.py under MIT license +# License available in LICENSES/LICENSE_NVIDIA.txt +class SnakeBeta(nn.Module): + def __init__( + self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True + ): + super(SnakeBeta, self).__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) + self.beta = nn.Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = nn.Parameter(torch.ones(in_features) * alpha) + self.beta = nn.Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + x = snake_beta(x, alpha, beta) + + return x diff --git a/src/YingMusicSinger/utils/stable_audio_tools/bottleneck copy.py b/src/YingMusicSinger/utils/stable_audio_tools/bottleneck copy.py new file mode 100755 index 0000000000000000000000000000000000000000..c3fb58d7a401b356e066e3653ce0ad66f7c573d5 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/bottleneck copy.py @@ -0,0 +1,393 @@ +import numpy as np +import torch +from dac.nn.quantize import ResidualVectorQuantize as DACResidualVQ +from einops import rearrange +from torch import nn +from torch.nn import functional as F +from vector_quantize_pytorch import FSQ, ResidualVQ + + +class Bottleneck(nn.Module): + def __init__(self, is_discrete: bool = False): + super().__init__() + + self.is_discrete = is_discrete + + def encode(self, x, return_info=False, **kwargs): + raise NotImplementedError + + def decode(self, x): + raise NotImplementedError + + +class DiscreteBottleneck(Bottleneck): + def __init__(self, num_quantizers, codebook_size, tokens_id): + super().__init__(is_discrete=True) + + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.tokens_id = tokens_id + + def decode_tokens(self, codes, **kwargs): + raise NotImplementedError + + +class TanhBottleneck(Bottleneck): + def __init__(self): + super().__init__(is_discrete=False) + self.tanh = nn.Tanh() + + def encode(self, x, return_info=False): + info = {} + + x = torch.tanh(x) + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + +def vae_sample(mean, scale): + stdev = nn.functional.softplus(scale) + 1e-4 + var = stdev * stdev + logvar = torch.log(var) + latents = torch.randn_like(mean) * stdev + mean + + kl = (mean * mean + var - logvar - 1).sum(1).mean() + + return latents, kl + + +class VAEBottleneck(Bottleneck): + def __init__(self): + super().__init__(is_discrete=False) + + def encode(self, x, return_info=False, **kwargs): + info = {} + + mean, scale = x.chunk(2, dim=1) + + x, kl = vae_sample(mean, scale) + + info["kl"] = kl + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + +def compute_mean_kernel(x, y): + kernel_input = (x[:, None] - y[None]).pow(2).mean(2) / x.shape[-1] + return torch.exp(-kernel_input).mean() + + +def compute_mmd(latents): + latents_reshaped = latents.permute(0, 2, 1).reshape(-1, latents.shape[1]) + noise = torch.randn_like(latents_reshaped) + + latents_kernel = compute_mean_kernel(latents_reshaped, latents_reshaped) + noise_kernel = compute_mean_kernel(noise, noise) + latents_noise_kernel = compute_mean_kernel(latents_reshaped, noise) + + mmd = latents_kernel + noise_kernel - 2 * latents_noise_kernel + return mmd.mean() + + +class WassersteinBottleneck(Bottleneck): + def __init__(self, noise_augment_dim: int = 0, bypass_mmd: bool = False): + super().__init__(is_discrete=False) + + self.noise_augment_dim = noise_augment_dim + self.bypass_mmd = bypass_mmd + + def encode(self, x, return_info=False): + info = {} + + if self.training and return_info: + if self.bypass_mmd: + mmd = torch.tensor(0.0) + else: + mmd = compute_mmd(x) + + info["mmd"] = mmd + + if return_info: + return x, info + + return x + + def decode(self, x): + if self.noise_augment_dim > 0: + noise = torch.randn( + x.shape[0], self.noise_augment_dim, x.shape[-1] + ).type_as(x) + x = torch.cat([x, noise], dim=1) + + return x + + +class L2Bottleneck(Bottleneck): + def __init__(self): + super().__init__(is_discrete=False) + + def encode(self, x, return_info=False): + info = {} + + x = F.normalize(x, dim=1) + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return F.normalize(x, dim=1) + + +class RVQBottleneck(DiscreteBottleneck): + def __init__(self, **quantizer_kwargs): + super().__init__( + num_quantizers=quantizer_kwargs["num_quantizers"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="quantizer_indices", + ) + self.quantizer = ResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["num_quantizers"] + + def encode(self, x, return_info=False, **kwargs): + info = {} + + x = rearrange(x, "b c n -> b n c") + x, indices, loss = self.quantizer(x) + x = rearrange(x, "b n c -> b c n") + + info["quantizer_indices"] = indices + info["quantizer_loss"] = loss.mean() + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + def decode_tokens(self, codes, **kwargs): + latents = self.quantizer.get_outputs_from_indices(codes) + + return self.decode(latents, **kwargs) + + +class RVQVAEBottleneck(DiscreteBottleneck): + def __init__(self, **quantizer_kwargs): + super().__init__( + num_quantizers=quantizer_kwargs["num_quantizers"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="quantizer_indices", + ) + self.quantizer = ResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["num_quantizers"] + + def encode(self, x, return_info=False): + info = {} + + x, kl = vae_sample(*x.chunk(2, dim=1)) + + info["kl"] = kl + + x = rearrange(x, "b c n -> b n c") + x, indices, loss = self.quantizer(x) + x = rearrange(x, "b n c -> b c n") + + info["quantizer_indices"] = indices + info["quantizer_loss"] = loss.mean() + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + def decode_tokens(self, codes, **kwargs): + latents = self.quantizer.get_outputs_from_indices(codes) + + return self.decode(latents, **kwargs) + + +class DACRVQBottleneck(DiscreteBottleneck): + def __init__( + self, quantize_on_decode=False, noise_augment_dim=0, **quantizer_kwargs + ): + super().__init__( + num_quantizers=quantizer_kwargs["n_codebooks"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="codes", + ) + self.quantizer = DACResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["n_codebooks"] + self.quantize_on_decode = quantize_on_decode + self.noise_augment_dim = noise_augment_dim + + def encode(self, x, return_info=False, **kwargs): + info = {} + + info["pre_quantizer"] = x + + if self.quantize_on_decode: + return x, info if return_info else x + + z, codes, latents, commitment_loss, codebook_loss = self.quantizer(x, **kwargs) + + output = { + "z": z, + "codes": codes, + "latents": latents, + "vq/commitment_loss": commitment_loss, + "vq/codebook_loss": codebook_loss, + } + + output["vq/commitment_loss"] /= self.num_quantizers + output["vq/codebook_loss"] /= self.num_quantizers + + info.update(output) + + if return_info: + return output["z"], info + + return output["z"] + + def decode(self, x): + if self.quantize_on_decode: + x = self.quantizer(x)[0] + + if self.noise_augment_dim > 0: + noise = torch.randn( + x.shape[0], self.noise_augment_dim, x.shape[-1] + ).type_as(x) + x = torch.cat([x, noise], dim=1) + + return x + + def decode_tokens(self, codes, **kwargs): + latents, _, _ = self.quantizer.from_codes(codes) + + return self.decode(latents, **kwargs) + + +class DACRVQVAEBottleneck(DiscreteBottleneck): + def __init__(self, quantize_on_decode=False, **quantizer_kwargs): + super().__init__( + num_quantizers=quantizer_kwargs["n_codebooks"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="codes", + ) + self.quantizer = DACResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["n_codebooks"] + self.quantize_on_decode = quantize_on_decode + + def encode(self, x, return_info=False, n_quantizers: int = None): + info = {} + + mean, scale = x.chunk(2, dim=1) + + x, kl = vae_sample(mean, scale) + + info["pre_quantizer"] = x + info["kl"] = kl + + if self.quantize_on_decode: + return x, info if return_info else x + + z, codes, latents, commitment_loss, codebook_loss = self.quantizer( + x, n_quantizers=n_quantizers + ) + + output = { + "z": z, + "codes": codes, + "latents": latents, + "vq/commitment_loss": commitment_loss, + "vq/codebook_loss": codebook_loss, + } + + output["vq/commitment_loss"] /= self.num_quantizers + output["vq/codebook_loss"] /= self.num_quantizers + + info.update(output) + + if return_info: + return output["z"], info + + return output["z"] + + def decode(self, x): + if self.quantize_on_decode: + x = self.quantizer(x)[0] + + return x + + def decode_tokens(self, codes, **kwargs): + latents, _, _ = self.quantizer.from_codes(codes) + + return self.decode(latents, **kwargs) + + +class FSQBottleneck(DiscreteBottleneck): + def __init__(self, noise_augment_dim=0, **kwargs): + super().__init__( + num_quantizers=kwargs.get("num_codebooks", 1), + codebook_size=np.prod(kwargs["levels"]), + tokens_id="quantizer_indices", + ) + + self.noise_augment_dim = noise_augment_dim + + self.quantizer = FSQ( + **kwargs, allowed_dtypes=[torch.float16, torch.float32, torch.float64] + ) + + def encode(self, x, return_info=False): + info = {} + + orig_dtype = x.dtype + x = x.float() + + x = rearrange(x, "b c n -> b n c") + x, indices = self.quantizer(x) + x = rearrange(x, "b n c -> b c n") + + x = x.to(orig_dtype) + + # Reorder indices to match the expected format + indices = rearrange(indices, "b n q -> b q n") + + info["quantizer_indices"] = indices + + if return_info: + return x, info + else: + return x + + def decode(self, x): + if self.noise_augment_dim > 0: + noise = torch.randn( + x.shape[0], self.noise_augment_dim, x.shape[-1] + ).type_as(x) + x = torch.cat([x, noise], dim=1) + + return x + + def decode_tokens(self, tokens, **kwargs): + latents = self.quantizer.indices_to_codes(tokens) + + return self.decode(latents, **kwargs) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/bottleneck.py b/src/YingMusicSinger/utils/stable_audio_tools/bottleneck.py new file mode 100755 index 0000000000000000000000000000000000000000..c3fb58d7a401b356e066e3653ce0ad66f7c573d5 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/bottleneck.py @@ -0,0 +1,393 @@ +import numpy as np +import torch +from dac.nn.quantize import ResidualVectorQuantize as DACResidualVQ +from einops import rearrange +from torch import nn +from torch.nn import functional as F +from vector_quantize_pytorch import FSQ, ResidualVQ + + +class Bottleneck(nn.Module): + def __init__(self, is_discrete: bool = False): + super().__init__() + + self.is_discrete = is_discrete + + def encode(self, x, return_info=False, **kwargs): + raise NotImplementedError + + def decode(self, x): + raise NotImplementedError + + +class DiscreteBottleneck(Bottleneck): + def __init__(self, num_quantizers, codebook_size, tokens_id): + super().__init__(is_discrete=True) + + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.tokens_id = tokens_id + + def decode_tokens(self, codes, **kwargs): + raise NotImplementedError + + +class TanhBottleneck(Bottleneck): + def __init__(self): + super().__init__(is_discrete=False) + self.tanh = nn.Tanh() + + def encode(self, x, return_info=False): + info = {} + + x = torch.tanh(x) + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + +def vae_sample(mean, scale): + stdev = nn.functional.softplus(scale) + 1e-4 + var = stdev * stdev + logvar = torch.log(var) + latents = torch.randn_like(mean) * stdev + mean + + kl = (mean * mean + var - logvar - 1).sum(1).mean() + + return latents, kl + + +class VAEBottleneck(Bottleneck): + def __init__(self): + super().__init__(is_discrete=False) + + def encode(self, x, return_info=False, **kwargs): + info = {} + + mean, scale = x.chunk(2, dim=1) + + x, kl = vae_sample(mean, scale) + + info["kl"] = kl + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + +def compute_mean_kernel(x, y): + kernel_input = (x[:, None] - y[None]).pow(2).mean(2) / x.shape[-1] + return torch.exp(-kernel_input).mean() + + +def compute_mmd(latents): + latents_reshaped = latents.permute(0, 2, 1).reshape(-1, latents.shape[1]) + noise = torch.randn_like(latents_reshaped) + + latents_kernel = compute_mean_kernel(latents_reshaped, latents_reshaped) + noise_kernel = compute_mean_kernel(noise, noise) + latents_noise_kernel = compute_mean_kernel(latents_reshaped, noise) + + mmd = latents_kernel + noise_kernel - 2 * latents_noise_kernel + return mmd.mean() + + +class WassersteinBottleneck(Bottleneck): + def __init__(self, noise_augment_dim: int = 0, bypass_mmd: bool = False): + super().__init__(is_discrete=False) + + self.noise_augment_dim = noise_augment_dim + self.bypass_mmd = bypass_mmd + + def encode(self, x, return_info=False): + info = {} + + if self.training and return_info: + if self.bypass_mmd: + mmd = torch.tensor(0.0) + else: + mmd = compute_mmd(x) + + info["mmd"] = mmd + + if return_info: + return x, info + + return x + + def decode(self, x): + if self.noise_augment_dim > 0: + noise = torch.randn( + x.shape[0], self.noise_augment_dim, x.shape[-1] + ).type_as(x) + x = torch.cat([x, noise], dim=1) + + return x + + +class L2Bottleneck(Bottleneck): + def __init__(self): + super().__init__(is_discrete=False) + + def encode(self, x, return_info=False): + info = {} + + x = F.normalize(x, dim=1) + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return F.normalize(x, dim=1) + + +class RVQBottleneck(DiscreteBottleneck): + def __init__(self, **quantizer_kwargs): + super().__init__( + num_quantizers=quantizer_kwargs["num_quantizers"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="quantizer_indices", + ) + self.quantizer = ResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["num_quantizers"] + + def encode(self, x, return_info=False, **kwargs): + info = {} + + x = rearrange(x, "b c n -> b n c") + x, indices, loss = self.quantizer(x) + x = rearrange(x, "b n c -> b c n") + + info["quantizer_indices"] = indices + info["quantizer_loss"] = loss.mean() + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + def decode_tokens(self, codes, **kwargs): + latents = self.quantizer.get_outputs_from_indices(codes) + + return self.decode(latents, **kwargs) + + +class RVQVAEBottleneck(DiscreteBottleneck): + def __init__(self, **quantizer_kwargs): + super().__init__( + num_quantizers=quantizer_kwargs["num_quantizers"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="quantizer_indices", + ) + self.quantizer = ResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["num_quantizers"] + + def encode(self, x, return_info=False): + info = {} + + x, kl = vae_sample(*x.chunk(2, dim=1)) + + info["kl"] = kl + + x = rearrange(x, "b c n -> b n c") + x, indices, loss = self.quantizer(x) + x = rearrange(x, "b n c -> b c n") + + info["quantizer_indices"] = indices + info["quantizer_loss"] = loss.mean() + + if return_info: + return x, info + else: + return x + + def decode(self, x): + return x + + def decode_tokens(self, codes, **kwargs): + latents = self.quantizer.get_outputs_from_indices(codes) + + return self.decode(latents, **kwargs) + + +class DACRVQBottleneck(DiscreteBottleneck): + def __init__( + self, quantize_on_decode=False, noise_augment_dim=0, **quantizer_kwargs + ): + super().__init__( + num_quantizers=quantizer_kwargs["n_codebooks"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="codes", + ) + self.quantizer = DACResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["n_codebooks"] + self.quantize_on_decode = quantize_on_decode + self.noise_augment_dim = noise_augment_dim + + def encode(self, x, return_info=False, **kwargs): + info = {} + + info["pre_quantizer"] = x + + if self.quantize_on_decode: + return x, info if return_info else x + + z, codes, latents, commitment_loss, codebook_loss = self.quantizer(x, **kwargs) + + output = { + "z": z, + "codes": codes, + "latents": latents, + "vq/commitment_loss": commitment_loss, + "vq/codebook_loss": codebook_loss, + } + + output["vq/commitment_loss"] /= self.num_quantizers + output["vq/codebook_loss"] /= self.num_quantizers + + info.update(output) + + if return_info: + return output["z"], info + + return output["z"] + + def decode(self, x): + if self.quantize_on_decode: + x = self.quantizer(x)[0] + + if self.noise_augment_dim > 0: + noise = torch.randn( + x.shape[0], self.noise_augment_dim, x.shape[-1] + ).type_as(x) + x = torch.cat([x, noise], dim=1) + + return x + + def decode_tokens(self, codes, **kwargs): + latents, _, _ = self.quantizer.from_codes(codes) + + return self.decode(latents, **kwargs) + + +class DACRVQVAEBottleneck(DiscreteBottleneck): + def __init__(self, quantize_on_decode=False, **quantizer_kwargs): + super().__init__( + num_quantizers=quantizer_kwargs["n_codebooks"], + codebook_size=quantizer_kwargs["codebook_size"], + tokens_id="codes", + ) + self.quantizer = DACResidualVQ(**quantizer_kwargs) + self.num_quantizers = quantizer_kwargs["n_codebooks"] + self.quantize_on_decode = quantize_on_decode + + def encode(self, x, return_info=False, n_quantizers: int = None): + info = {} + + mean, scale = x.chunk(2, dim=1) + + x, kl = vae_sample(mean, scale) + + info["pre_quantizer"] = x + info["kl"] = kl + + if self.quantize_on_decode: + return x, info if return_info else x + + z, codes, latents, commitment_loss, codebook_loss = self.quantizer( + x, n_quantizers=n_quantizers + ) + + output = { + "z": z, + "codes": codes, + "latents": latents, + "vq/commitment_loss": commitment_loss, + "vq/codebook_loss": codebook_loss, + } + + output["vq/commitment_loss"] /= self.num_quantizers + output["vq/codebook_loss"] /= self.num_quantizers + + info.update(output) + + if return_info: + return output["z"], info + + return output["z"] + + def decode(self, x): + if self.quantize_on_decode: + x = self.quantizer(x)[0] + + return x + + def decode_tokens(self, codes, **kwargs): + latents, _, _ = self.quantizer.from_codes(codes) + + return self.decode(latents, **kwargs) + + +class FSQBottleneck(DiscreteBottleneck): + def __init__(self, noise_augment_dim=0, **kwargs): + super().__init__( + num_quantizers=kwargs.get("num_codebooks", 1), + codebook_size=np.prod(kwargs["levels"]), + tokens_id="quantizer_indices", + ) + + self.noise_augment_dim = noise_augment_dim + + self.quantizer = FSQ( + **kwargs, allowed_dtypes=[torch.float16, torch.float32, torch.float64] + ) + + def encode(self, x, return_info=False): + info = {} + + orig_dtype = x.dtype + x = x.float() + + x = rearrange(x, "b c n -> b n c") + x, indices = self.quantizer(x) + x = rearrange(x, "b n c -> b c n") + + x = x.to(orig_dtype) + + # Reorder indices to match the expected format + indices = rearrange(indices, "b n q -> b q n") + + info["quantizer_indices"] = indices + + if return_info: + return x, info + else: + return x + + def decode(self, x): + if self.noise_augment_dim > 0: + noise = torch.randn( + x.shape[0], self.noise_augment_dim, x.shape[-1] + ).type_as(x) + x = torch.cat([x, noise], dim=1) + + return x + + def decode_tokens(self, tokens, **kwargs): + latents = self.quantizer.indices_to_codes(tokens) + + return self.decode(latents, **kwargs) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/conditioners.py b/src/YingMusicSinger/utils/stable_audio_tools/conditioners.py new file mode 100755 index 0000000000000000000000000000000000000000..bb9638f8b6d5f950a5307ffd451758e72370fa4b --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/conditioners.py @@ -0,0 +1,664 @@ +# Heavily influenced by https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/modules/conditioners.py + +import gc +import logging +import string +import typing as tp +import warnings + +import torch +from torch import nn + +from .adp import NumberEmbedder + +# from ..inference.utils import set_audio_channels +from .factory import create_pretransform_from_config +from .pretransforms import Pretransform + +# from ..training.utils import copy_state_dict +from .utils import load_ckpt_state_dict + + +class Conditioner(nn.Module): + def __init__(self, dim: int, output_dim: int, project_out: bool = False): + super().__init__() + + self.dim = dim + self.output_dim = output_dim + self.proj_out = ( + nn.Linear(dim, output_dim) + if (dim != output_dim or project_out) + else nn.Identity() + ) + + def forward(self, x: tp.Any) -> tp.Any: + raise NotImplementedError() + + +class IntConditioner(Conditioner): + def __init__(self, output_dim: int, min_val: int = 0, max_val: int = 512): + super().__init__(output_dim, output_dim) + + self.min_val = min_val + self.max_val = max_val + self.int_embedder = nn.Embedding( + max_val - min_val + 1, output_dim + ).requires_grad_(True) + + def forward(self, ints: tp.List[int], device=None) -> tp.Any: + # self.int_embedder.to(device) + + ints = torch.tensor(ints).to(device) + ints = ints.clamp(self.min_val, self.max_val) + + int_embeds = self.int_embedder(ints).unsqueeze(1) + + return [int_embeds, torch.ones(int_embeds.shape[0], 1).to(device)] + + +class NumberConditioner(Conditioner): + """ + Conditioner that takes a list of floats, normalizes them for a given range, and returns a list of embeddings + """ + + def __init__(self, output_dim: int, min_val: float = 0, max_val: float = 1): + super().__init__(output_dim, output_dim) + + self.min_val = min_val + self.max_val = max_val + + self.embedder = NumberEmbedder(features=output_dim) + + def forward(self, floats: tp.List[float], device=None) -> tp.Any: + # Cast the inputs to floats + floats = [float(x) for x in floats] + + floats = torch.tensor(floats).to(device) + + floats = floats.clamp(self.min_val, self.max_val) + + normalized_floats = (floats - self.min_val) / (self.max_val - self.min_val) + + # Cast floats to same type as embedder + embedder_dtype = next(self.embedder.parameters()).dtype + normalized_floats = normalized_floats.to(embedder_dtype) + + float_embeds = self.embedder(normalized_floats).unsqueeze(1) + + return [float_embeds, torch.ones(float_embeds.shape[0], 1).to(device)] + + +class CLAPTextConditioner(Conditioner): + def __init__( + self, + output_dim: int, + clap_ckpt_path, + use_text_features=False, + feature_layer_ix: int = -1, + audio_model_type="HTSAT-base", + enable_fusion=True, + project_out: bool = False, + finetune: bool = False, + ): + super().__init__( + 768 if use_text_features else 512, output_dim, project_out=project_out + ) + + self.use_text_features = use_text_features + self.feature_layer_ix = feature_layer_ix + self.finetune = finetune + + # Suppress logging from transformers + previous_level = logging.root.manager.disable + logging.disable(logging.ERROR) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + import laion_clap + from laion_clap.clap_module.factory import ( + load_state_dict as clap_load_state_dict, + ) + + model = laion_clap.CLAP_Module( + enable_fusion=enable_fusion, amodel=audio_model_type, device="cpu" + ) + + if self.finetune: + self.model = model + else: + self.__dict__["model"] = model + + state_dict = clap_load_state_dict(clap_ckpt_path) + self.model.model.load_state_dict(state_dict, strict=False) + + if self.finetune: + self.model.model.text_branch.requires_grad_(True) + self.model.model.text_branch.train() + else: + self.model.model.text_branch.requires_grad_(False) + self.model.model.text_branch.eval() + + finally: + logging.disable(previous_level) + + del self.model.model.audio_branch + + gc.collect() + torch.cuda.empty_cache() + + def get_clap_features(self, prompts, layer_ix=-2, device: tp.Any = "cuda"): + prompt_tokens = self.model.tokenizer(prompts) + attention_mask = prompt_tokens["attention_mask"].to( + device=device, non_blocking=True + ) + prompt_features = self.model.model.text_branch( + input_ids=prompt_tokens["input_ids"].to(device=device, non_blocking=True), + attention_mask=attention_mask, + output_hidden_states=True, + )["hidden_states"][layer_ix] + + return prompt_features, attention_mask + + def forward(self, texts: tp.List[str], device: tp.Any = "cuda") -> tp.Any: + self.model.to(device) + + if self.use_text_features: + if len(texts) == 1: + text_features, text_attention_mask = self.get_clap_features( + [texts[0], ""], layer_ix=self.feature_layer_ix, device=device + ) + text_features = text_features[:1, ...] + text_attention_mask = text_attention_mask[:1, ...] + else: + text_features, text_attention_mask = self.get_clap_features( + texts, layer_ix=self.feature_layer_ix, device=device + ) + return [self.proj_out(text_features), text_attention_mask] + + # Fix for CLAP bug when only one text is passed + if len(texts) == 1: + text_embedding = self.model.get_text_embedding( + [texts[0], ""], use_tensor=True + )[:1, ...] + else: + text_embedding = self.model.get_text_embedding(texts, use_tensor=True) + + text_embedding = text_embedding.unsqueeze(1).to(device) + + return [ + self.proj_out(text_embedding), + torch.ones(text_embedding.shape[0], 1).to(device), + ] + + +class CLAPAudioConditioner(Conditioner): + def __init__( + self, + output_dim: int, + clap_ckpt_path, + audio_model_type="HTSAT-base", + enable_fusion=True, + project_out: bool = False, + ): + super().__init__(512, output_dim, project_out=project_out) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Suppress logging from transformers + previous_level = logging.root.manager.disable + logging.disable(logging.ERROR) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + import laion_clap + from laion_clap.clap_module.factory import ( + load_state_dict as clap_load_state_dict, + ) + + model = laion_clap.CLAP_Module( + enable_fusion=enable_fusion, amodel=audio_model_type, device="cpu" + ) + + if self.finetune: + self.model = model + else: + self.__dict__["model"] = model + + state_dict = clap_load_state_dict(clap_ckpt_path) + self.model.model.load_state_dict(state_dict, strict=False) + + if self.finetune: + self.model.model.audio_branch.requires_grad_(True) + self.model.model.audio_branch.train() + else: + self.model.model.audio_branch.requires_grad_(False) + self.model.model.audio_branch.eval() + + finally: + logging.disable(previous_level) + + del self.model.model.text_branch + + gc.collect() + torch.cuda.empty_cache() + + def forward( + self, + audios: tp.Union[torch.Tensor, tp.List[torch.Tensor], tp.Tuple[torch.Tensor]], + device: tp.Any = "cuda", + ) -> tp.Any: + self.model.to(device) + + if isinstance(audios, list) or isinstance(audios, tuple): + audios = torch.cat(audios, dim=0) + + # Convert to mono + mono_audios = audios.mean(dim=1) + + with torch.cuda.amp.autocast(enabled=False): + audio_embedding = self.model.get_audio_embedding_from_data( + mono_audios.float(), use_tensor=True + ) + + audio_embedding = audio_embedding.unsqueeze(1).to(device) + + return [ + self.proj_out(audio_embedding), + torch.ones(audio_embedding.shape[0], 1).to(device), + ] + + +class T5Conditioner(Conditioner): + T5_MODELS = [ + "t5-small", + "t5-base", + "t5-large", + "t5-3b", + "t5-11b", + "google/flan-t5-small", + "google/flan-t5-base", + "google/flan-t5-large", + "google/flan-t5-xl", + "google/flan-t5-xxl", + ] + + T5_MODEL_DIMS = { + "t5-small": 512, + "t5-base": 768, + "t5-large": 1024, + "t5-3b": 1024, + "t5-11b": 1024, + "t5-xl": 2048, + "t5-xxl": 4096, + "google/flan-t5-small": 512, + "google/flan-t5-base": 768, + "google/flan-t5-large": 1024, + "google/flan-t5-3b": 1024, + "google/flan-t5-11b": 1024, + "google/flan-t5-xl": 2048, + "google/flan-t5-xxl": 4096, + } + + def __init__( + self, + output_dim: int, + t5_model_name: str = "t5-base", + max_length: str = 128, + enable_grad: bool = False, + project_out: bool = False, + ): + assert t5_model_name in self.T5_MODELS, ( + f"Unknown T5 model name: {t5_model_name}" + ) + super().__init__( + self.T5_MODEL_DIMS[t5_model_name], output_dim, project_out=project_out + ) + + from transformers import AutoTokenizer, T5EncoderModel + + self.max_length = max_length + self.enable_grad = enable_grad + + # Suppress logging from transformers + previous_level = logging.root.manager.disable + logging.disable(logging.ERROR) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + # self.tokenizer = T5Tokenizer.from_pretrained(t5_model_name, model_max_length = max_length) + # model = T5EncoderModel.from_pretrained(t5_model_name, max_length=max_length).train(enable_grad).requires_grad_(enable_grad) + self.tokenizer = AutoTokenizer.from_pretrained(t5_model_name) + model = ( + T5EncoderModel.from_pretrained(t5_model_name) + .train(enable_grad) + .requires_grad_(enable_grad) + .to(torch.float16) + ) + finally: + logging.disable(previous_level) + + if self.enable_grad: + self.model = model + else: + self.__dict__["model"] = model + + def forward( + self, texts: tp.List[str], device: tp.Union[torch.device, str] + ) -> tp.Tuple[torch.Tensor, torch.Tensor]: + self.model.to(device) + self.proj_out.to(device) + + encoded = self.tokenizer( + texts, + truncation=True, + max_length=self.max_length, + padding="max_length", + return_tensors="pt", + ) + + input_ids = encoded["input_ids"].to(device) + attention_mask = encoded["attention_mask"].to(device).to(torch.bool) + + self.model.eval() + + with torch.cuda.amp.autocast(dtype=torch.float16) and torch.set_grad_enabled( + self.enable_grad + ): + embeddings = self.model(input_ids=input_ids, attention_mask=attention_mask)[ + "last_hidden_state" + ] + + embeddings = self.proj_out(embeddings.float()) + + embeddings = embeddings * attention_mask.unsqueeze(-1).float() + + return embeddings, attention_mask + + +class PhonemeConditioner(Conditioner): + """ + A conditioner that turns text into phonemes and embeds them using a lookup table + Only works for English text + + Args: + output_dim: the dimension of the output embeddings + max_length: the maximum number of phonemes to embed + project_out: whether to add another linear projection to the output embeddings + """ + + def __init__( + self, + output_dim: int, + max_length: int = 1024, + project_out: bool = False, + ): + super().__init__(output_dim, output_dim, project_out=project_out) + + from g2p_en import G2p + + self.max_length = max_length + + self.g2p = G2p() + + # Reserving 0 for padding, 1 for ignored + self.phoneme_embedder = nn.Embedding(len(self.g2p.phonemes) + 2, output_dim) + + def forward( + self, texts: tp.List[str], device: tp.Union[torch.device, str] + ) -> tp.Tuple[torch.Tensor, torch.Tensor]: + self.phoneme_embedder.to(device) + self.proj_out.to(device) + + batch_phonemes = [ + self.g2p(text) for text in texts + ] # shape [batch_size, length] + + phoneme_ignore = [" ", *string.punctuation] + + # Remove ignored phonemes and cut to max length + batch_phonemes = [ + [p if p not in phoneme_ignore else "_" for p in phonemes] + for phonemes in batch_phonemes + ] + + # Convert to ids + phoneme_ids = [ + [self.g2p.p2idx[p] + 2 if p in self.g2p.p2idx else 1 for p in phonemes] + for phonemes in batch_phonemes + ] + + # Pad to match longest and make a mask tensor for the padding + longest = max([len(ids) for ids in phoneme_ids]) + phoneme_ids = [ids + [0] * (longest - len(ids)) for ids in phoneme_ids] + + phoneme_ids = torch.tensor(phoneme_ids).to(device) + + # Convert to embeddings + phoneme_embeds = self.phoneme_embedder(phoneme_ids) + + phoneme_embeds = self.proj_out(phoneme_embeds) + + return phoneme_embeds, torch.ones( + phoneme_embeds.shape[0], phoneme_embeds.shape[1] + ).to(device) + + +class TokenizerLUTConditioner(Conditioner): + """ + A conditioner that embeds text using a lookup table on a pretrained tokenizer's vocabulary + + Args: + tokenizer_name: the name of the tokenizer from the Hugging Face transformers library + output_dim: the dimension of the output embeddings + max_length: the maximum length of the text to embed + project_out: whether to add another linear projection to the output embeddings + """ + + def __init__( + self, + tokenizer_name: str, # Name of a tokenizer from the Hugging Face transformers library + output_dim: int, + max_length: int = 1024, + project_out: bool = False, + ): + super().__init__(output_dim, output_dim, project_out=project_out) + + from transformers import AutoTokenizer + + # Suppress logging from transformers + previous_level = logging.root.manager.disable + logging.disable(logging.ERROR) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) + finally: + logging.disable(previous_level) + + self.max_length = max_length + + self.token_embedder = nn.Embedding(len(self.tokenizer), output_dim) + + def forward( + self, texts: tp.List[str], device: tp.Union[torch.device, str] + ) -> tp.Tuple[torch.Tensor, torch.Tensor]: + self.proj_out.to(device) + + encoded = self.tokenizer( + texts, + truncation=True, + max_length=self.max_length, + padding="max_length", + return_tensors="pt", + ) + + input_ids = encoded["input_ids"].to(device) + attention_mask = encoded["attention_mask"].to(device).to(torch.bool) + + embeddings = self.token_embedder(input_ids) + + embeddings = self.proj_out(embeddings) + + embeddings = embeddings * attention_mask.unsqueeze(-1).float() + + return embeddings, attention_mask + + +class PretransformConditioner(Conditioner): + """ + A conditioner that uses a pretransform's encoder for conditioning + + Args: + pretransform: an instantiated pretransform to use for conditioning + output_dim: the dimension of the output embeddings + """ + + def __init__(self, pretransform: Pretransform, output_dim: int): + super().__init__(pretransform.encoded_channels, output_dim) + + self.pretransform = pretransform + + def forward( + self, + audio: tp.Union[torch.Tensor, tp.List[torch.Tensor], tp.Tuple[torch.Tensor]], + device: tp.Union[torch.device, str], + ) -> tp.Tuple[torch.Tensor, torch.Tensor]: + self.pretransform.to(device) + self.proj_out.to(device) + + if isinstance(audio, list) or isinstance(audio, tuple): + audio = torch.cat(audio, dim=0) + + # Convert audio to pretransform input channels + audio = set_audio_channels(audio, self.pretransform.io_channels) + + latents = self.pretransform.encode(audio) + + latents = self.proj_out(latents) + + return [ + latents, + torch.ones(latents.shape[0], latents.shape[2]).to(latents.device), + ] + + +class MultiConditioner(nn.Module): + """ + A module that applies multiple conditioners to an input dictionary based on the keys + + Args: + conditioners: a dictionary of conditioners with keys corresponding to the keys of the conditioning input dictionary (e.g. "prompt") + default_keys: a dictionary of default keys to use if the key is not in the input dictionary (e.g. {"prompt_t5": "prompt"}) + """ + + def __init__( + self, + conditioners: tp.Dict[str, Conditioner], + default_keys: tp.Dict[str, str] = {}, + ): + super().__init__() + + self.conditioners = nn.ModuleDict(conditioners) + self.default_keys = default_keys + + def forward( + self, + batch_metadata: tp.List[tp.Dict[str, tp.Any]], + device: tp.Union[torch.device, str], + ) -> tp.Dict[str, tp.Any]: + output = {} + + for key, conditioner in self.conditioners.items(): + condition_key = key + + conditioner_inputs = [] + + for x in batch_metadata: + if condition_key not in x: + if condition_key in self.default_keys: + condition_key = self.default_keys[condition_key] + else: + raise ValueError( + f"Conditioner key {condition_key} not found in batch metadata" + ) + + # Unwrap the condition info if it's a single-element list or tuple, this is to support collation functions that wrap everything in a list + if ( + isinstance(x[condition_key], list) + or isinstance(x[condition_key], tuple) + and len(x[condition_key]) == 1 + ): + conditioner_input = x[condition_key][0] + + else: + conditioner_input = x[condition_key] + + conditioner_inputs.append(conditioner_input) + + output[key] = conditioner(conditioner_inputs, device) + + return output + + +def create_multi_conditioner_from_conditioning_config( + config: tp.Dict[str, tp.Any], +) -> MultiConditioner: + """ + Create a MultiConditioner from a conditioning config dictionary + + Args: + config: the conditioning config dictionary + device: the device to put the conditioners on + """ + conditioners = {} + cond_dim = config["cond_dim"] + + default_keys = config.get("default_keys", {}) + + for conditioner_info in config["configs"]: + id = conditioner_info["id"] + + conditioner_type = conditioner_info["type"] + + conditioner_config = {"output_dim": cond_dim} + + conditioner_config.update(conditioner_info["config"]) + + if conditioner_type == "t5": + conditioners[id] = T5Conditioner(**conditioner_config) + elif conditioner_type == "clap_text": + conditioners[id] = CLAPTextConditioner(**conditioner_config) + elif conditioner_type == "clap_audio": + conditioners[id] = CLAPAudioConditioner(**conditioner_config) + elif conditioner_type == "int": + conditioners[id] = IntConditioner(**conditioner_config) + elif conditioner_type == "number": + conditioners[id] = NumberConditioner(**conditioner_config) + elif conditioner_type == "phoneme": + conditioners[id] = PhonemeConditioner(**conditioner_config) + elif conditioner_type == "lut": + conditioners[id] = TokenizerLUTConditioner(**conditioner_config) + elif conditioner_type == "pretransform": + sample_rate = conditioner_config.pop("sample_rate", None) + assert sample_rate is not None, ( + "Sample rate must be specified for pretransform conditioners" + ) + + pretransform = create_pretransform_from_config( + conditioner_config.pop("pretransform_config"), sample_rate=sample_rate + ) + + if conditioner_config.get("pretransform_ckpt_path", None) is not None: + pretransform.load_state_dict( + load_ckpt_state_dict( + conditioner_config.pop("pretransform_ckpt_path") + ) + ) + + conditioners[id] = PretransformConditioner( + pretransform, **conditioner_config + ) + else: + raise ValueError(f"Unknown conditioner type: {conditioner_type}") + + return MultiConditioner(conditioners, default_keys=default_keys) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/diffusion.py b/src/YingMusicSinger/utils/stable_audio_tools/diffusion.py new file mode 100755 index 0000000000000000000000000000000000000000..2aef74d564f53450db86be8befcbb9a005b882d0 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/diffusion.py @@ -0,0 +1,740 @@ +import typing as tp +from functools import partial +from time import time + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F + +# from ..inference.generation import generate_diffusion_cond +from .adp import UNet1d, UNetCFG1d +from .blocks import ( + Downsample1d, + Downsample1d_2, + FourierFeatures, + ResConvBlock, + SelfAttention1d, + SkipBlock, + Upsample1d, + Upsample1d_2, + expand_to_planes, +) +from .conditioners import ( + MultiConditioner, + create_multi_conditioner_from_conditioning_config, +) +from .dit import DiffusionTransformer +from .factory import create_pretransform_from_config +from .pretransforms import Pretransform + + +class Profiler: + def __init__(self): + self.ticks = [[time(), None]] + + def tick(self, msg): + self.ticks.append([time(), msg]) + + def __repr__(self): + rep = 80 * "=" + "\n" + for i in range(1, len(self.ticks)): + msg = self.ticks[i][1] + ellapsed = self.ticks[i][0] - self.ticks[i - 1][0] + rep += msg + f": {ellapsed * 1000:.2f}ms\n" + rep += 80 * "=" + "\n\n\n" + return rep + + +class DiffusionModel(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, x, t, **kwargs): + raise NotImplementedError() + + +class DiffusionModelWrapper(nn.Module): + def __init__( + self, + model: DiffusionModel, + io_channels, + sample_size, + sample_rate, + min_input_length, + pretransform: tp.Optional[Pretransform] = None, + ): + super().__init__() + self.io_channels = io_channels + self.sample_size = sample_size + self.sample_rate = sample_rate + self.min_input_length = min_input_length + + self.model = model + + if pretransform is not None: + self.pretransform = pretransform + else: + self.pretransform = None + + def forward(self, x, t, **kwargs): + return self.model(x, t, **kwargs) + + +class ConditionedDiffusionModel(nn.Module): + def __init__( + self, + *args, + supports_cross_attention: bool = False, + supports_input_concat: bool = False, + supports_global_cond: bool = False, + supports_prepend_cond: bool = False, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.supports_cross_attention = supports_cross_attention + self.supports_input_concat = supports_input_concat + self.supports_global_cond = supports_global_cond + self.supports_prepend_cond = supports_prepend_cond + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + cross_attn_cond: torch.Tensor = None, + cross_attn_mask: torch.Tensor = None, + input_concat_cond: torch.Tensor = None, + global_embed: torch.Tensor = None, + prepend_cond: torch.Tensor = None, + prepend_cond_mask: torch.Tensor = None, + cfg_scale: float = 1.0, + cfg_dropout_prob: float = 0.0, + batch_cfg: bool = False, + rescale_cfg: bool = False, + **kwargs, + ): + raise NotImplementedError() + + +class ConditionedDiffusionModelWrapper(nn.Module): + """ + A diffusion model that takes in conditioning + """ + + def __init__( + self, + model: ConditionedDiffusionModel, + conditioner: MultiConditioner, + io_channels, + sample_rate, + min_input_length: int, + diffusion_objective: tp.Literal["v", "rectified_flow"] = "v", + pretransform: tp.Optional[Pretransform] = None, + cross_attn_cond_ids: tp.List[str] = [], + global_cond_ids: tp.List[str] = [], + input_concat_ids: tp.List[str] = [], + prepend_cond_ids: tp.List[str] = [], + ): + super().__init__() + + self.model = model + self.conditioner = conditioner + self.io_channels = io_channels + self.sample_rate = sample_rate + self.diffusion_objective = diffusion_objective + self.pretransform = pretransform + self.cross_attn_cond_ids = cross_attn_cond_ids + self.global_cond_ids = global_cond_ids + self.input_concat_ids = input_concat_ids + self.prepend_cond_ids = prepend_cond_ids + self.min_input_length = min_input_length + + def get_conditioning_inputs( + self, conditioning_tensors: tp.Dict[str, tp.Any], negative=False + ): + cross_attention_input = None + cross_attention_masks = None + global_cond = None + input_concat_cond = None + prepend_cond = None + prepend_cond_mask = None + + if len(self.cross_attn_cond_ids) > 0: + # Concatenate all cross-attention inputs over the sequence dimension + # Assumes that the cross-attention inputs are of shape (batch, seq, channels) + cross_attention_input = [] + cross_attention_masks = [] + + for key in self.cross_attn_cond_ids: + cross_attn_in, cross_attn_mask = conditioning_tensors[key] + + # Add sequence dimension if it's not there + if len(cross_attn_in.shape) == 2: + cross_attn_in = cross_attn_in.unsqueeze(1) + cross_attn_mask = cross_attn_mask.unsqueeze(1) + + cross_attention_input.append(cross_attn_in) + cross_attention_masks.append(cross_attn_mask) + + cross_attention_input = torch.cat(cross_attention_input, dim=1) + cross_attention_masks = torch.cat(cross_attention_masks, dim=1) + + if len(self.global_cond_ids) > 0: + # Concatenate all global conditioning inputs over the channel dimension + # Assumes that the global conditioning inputs are of shape (batch, channels) + global_conds = [] + for key in self.global_cond_ids: + global_cond_input = conditioning_tensors[key][0] + + global_conds.append(global_cond_input) + + # Concatenate over the channel dimension + global_cond = torch.cat(global_conds, dim=-1) + + if len(global_cond.shape) == 3: + global_cond = global_cond.squeeze(1) + + if len(self.input_concat_ids) > 0: + # Concatenate all input concat conditioning inputs over the channel dimension + # Assumes that the input concat conditioning inputs are of shape (batch, channels, seq) + input_concat_cond = torch.cat( + [conditioning_tensors[key][0] for key in self.input_concat_ids], dim=1 + ) + + if len(self.prepend_cond_ids) > 0: + # Concatenate all prepend conditioning inputs over the sequence dimension + # Assumes that the prepend conditioning inputs are of shape (batch, seq, channels) + prepend_conds = [] + prepend_cond_masks = [] + + for key in self.prepend_cond_ids: + prepend_cond_input, prepend_cond_mask = conditioning_tensors[key] + prepend_conds.append(prepend_cond_input) + prepend_cond_masks.append(prepend_cond_mask) + + prepend_cond = torch.cat(prepend_conds, dim=1) + prepend_cond_mask = torch.cat(prepend_cond_masks, dim=1) + + if negative: + return { + "negative_cross_attn_cond": cross_attention_input, + "negative_cross_attn_mask": cross_attention_masks, + "negative_global_cond": global_cond, + "negative_input_concat_cond": input_concat_cond, + } + else: + return { + "cross_attn_cond": cross_attention_input, + "cross_attn_mask": cross_attention_masks, + "global_cond": global_cond, + "input_concat_cond": input_concat_cond, + "prepend_cond": prepend_cond, + "prepend_cond_mask": prepend_cond_mask, + } + + def forward( + self, x: torch.Tensor, t: torch.Tensor, cond: tp.Dict[str, tp.Any], **kwargs + ): + return self.model(x, t, **self.get_conditioning_inputs(cond), **kwargs) + + def generate(self, *args, **kwargs): + return generate_diffusion_cond(self, *args, **kwargs) + + +class UNetCFG1DWrapper(ConditionedDiffusionModel): + def __init__(self, *args, **kwargs): + super().__init__( + supports_cross_attention=True, + supports_global_cond=True, + supports_input_concat=True, + ) + + self.model = UNetCFG1d(*args, **kwargs) + + with torch.no_grad(): + for param in self.model.parameters(): + param *= 0.5 + + def forward( + self, + x, + t, + cross_attn_cond=None, + cross_attn_mask=None, + input_concat_cond=None, + global_cond=None, + cfg_scale=1.0, + cfg_dropout_prob: float = 0.0, + batch_cfg: bool = False, + rescale_cfg: bool = False, + negative_cross_attn_cond=None, + negative_cross_attn_mask=None, + negative_global_cond=None, + negative_input_concat_cond=None, + prepend_cond=None, + prepend_cond_mask=None, + **kwargs, + ): + p = Profiler() + + p.tick("start") + + channels_list = None + if input_concat_cond is not None: + channels_list = [input_concat_cond] + + outputs = self.model( + x, + t, + embedding=cross_attn_cond, + embedding_mask=cross_attn_mask, + features=global_cond, + channels_list=channels_list, + embedding_scale=cfg_scale, + embedding_mask_proba=cfg_dropout_prob, + batch_cfg=batch_cfg, + rescale_cfg=rescale_cfg, + negative_embedding=negative_cross_attn_cond, + negative_embedding_mask=negative_cross_attn_mask, + **kwargs, + ) + + p.tick("UNetCFG1D forward") + + # print(f"Profiler: {p}") + return outputs + + +class UNet1DCondWrapper(ConditionedDiffusionModel): + def __init__(self, *args, **kwargs): + super().__init__( + supports_cross_attention=False, + supports_global_cond=True, + supports_input_concat=True, + ) + + self.model = UNet1d(*args, **kwargs) + + with torch.no_grad(): + for param in self.model.parameters(): + param *= 0.5 + + def forward( + self, + x, + t, + input_concat_cond=None, + global_cond=None, + cross_attn_cond=None, + cross_attn_mask=None, + prepend_cond=None, + prepend_cond_mask=None, + cfg_scale=1.0, + cfg_dropout_prob: float = 0.0, + batch_cfg: bool = False, + rescale_cfg: bool = False, + negative_cross_attn_cond=None, + negative_cross_attn_mask=None, + negative_global_cond=None, + negative_input_concat_cond=None, + **kwargs, + ): + channels_list = None + if input_concat_cond is not None: + # Interpolate input_concat_cond to the same length as x + if input_concat_cond.shape[2] != x.shape[2]: + input_concat_cond = F.interpolate( + input_concat_cond, (x.shape[2],), mode="nearest" + ) + + channels_list = [input_concat_cond] + + outputs = self.model( + x, t, features=global_cond, channels_list=channels_list, **kwargs + ) + + return outputs + + +class UNet1DUncondWrapper(DiffusionModel): + def __init__(self, in_channels, *args, **kwargs): + super().__init__() + + self.model = UNet1d(in_channels=in_channels, *args, **kwargs) + + self.io_channels = in_channels + + with torch.no_grad(): + for param in self.model.parameters(): + param *= 0.5 + + def forward(self, x, t, **kwargs): + return self.model(x, t, **kwargs) + + +class DAU1DCondWrapper(ConditionedDiffusionModel): + def __init__(self, *args, **kwargs): + super().__init__( + supports_cross_attention=False, + supports_global_cond=False, + supports_input_concat=True, + ) + + self.model = DiffusionAttnUnet1D(*args, **kwargs) + + with torch.no_grad(): + for param in self.model.parameters(): + param *= 0.5 + + def forward( + self, + x, + t, + input_concat_cond=None, + cross_attn_cond=None, + cross_attn_mask=None, + global_cond=None, + cfg_scale=1.0, + cfg_dropout_prob: float = 0.0, + batch_cfg: bool = False, + rescale_cfg: bool = False, + negative_cross_attn_cond=None, + negative_cross_attn_mask=None, + negative_global_cond=None, + negative_input_concat_cond=None, + prepend_cond=None, + **kwargs, + ): + return self.model(x, t, cond=input_concat_cond) + + +class DiffusionAttnUnet1D(nn.Module): + def __init__( + self, + io_channels=2, + depth=14, + n_attn_layers=6, + channels=[128, 128, 256, 256] + [512] * 10, + cond_dim=0, + cond_noise_aug=False, + kernel_size=5, + learned_resample=False, + strides=[2] * 13, + conv_bias=True, + use_snake=False, + ): + super().__init__() + + self.cond_noise_aug = cond_noise_aug + + self.io_channels = io_channels + + if self.cond_noise_aug: + self.rng = torch.quasirandom.SobolEngine(1, scramble=True) + + self.timestep_embed = FourierFeatures(1, 16) + + attn_layer = depth - n_attn_layers + + strides = [1] + strides + + block = nn.Identity() + + conv_block = partial( + ResConvBlock, + kernel_size=kernel_size, + conv_bias=conv_bias, + use_snake=use_snake, + ) + + for i in range(depth, 0, -1): + c = channels[i - 1] + stride = strides[i - 1] + if stride > 2 and not learned_resample: + raise ValueError("Must have stride 2 without learned resampling") + + if i > 1: + c_prev = channels[i - 2] + add_attn = i >= attn_layer and n_attn_layers > 0 + block = SkipBlock( + Downsample1d_2(c_prev, c_prev, stride) + if (learned_resample or stride == 1) + else Downsample1d("cubic"), + conv_block(c_prev, c, c), + SelfAttention1d(c, c // 32) if add_attn else nn.Identity(), + conv_block(c, c, c), + SelfAttention1d(c, c // 32) if add_attn else nn.Identity(), + conv_block(c, c, c), + SelfAttention1d(c, c // 32) if add_attn else nn.Identity(), + block, + conv_block(c * 2 if i != depth else c, c, c), + SelfAttention1d(c, c // 32) if add_attn else nn.Identity(), + conv_block(c, c, c), + SelfAttention1d(c, c // 32) if add_attn else nn.Identity(), + conv_block(c, c, c_prev), + SelfAttention1d(c_prev, c_prev // 32) + if add_attn + else nn.Identity(), + Upsample1d_2(c_prev, c_prev, stride) + if learned_resample + else Upsample1d(kernel="cubic"), + ) + else: + cond_embed_dim = 16 if not self.cond_noise_aug else 32 + block = nn.Sequential( + conv_block((io_channels + cond_dim) + cond_embed_dim, c, c), + conv_block(c, c, c), + conv_block(c, c, c), + block, + conv_block(c * 2, c, c), + conv_block(c, c, c), + conv_block(c, c, io_channels, is_last=True), + ) + self.net = block + + with torch.no_grad(): + for param in self.net.parameters(): + param *= 0.5 + + def forward(self, x, t, cond=None, cond_aug_scale=None): + timestep_embed = expand_to_planes(self.timestep_embed(t[:, None]), x.shape) + + inputs = [x, timestep_embed] + + if cond is not None: + if cond.shape[2] != x.shape[2]: + cond = F.interpolate( + cond, (x.shape[2],), mode="linear", align_corners=False + ) + + if self.cond_noise_aug: + # Get a random number between 0 and 1, uniformly sampled + if cond_aug_scale is None: + aug_level = self.rng.draw(cond.shape[0])[:, 0].to(cond) + else: + aug_level = ( + torch.tensor([cond_aug_scale]).repeat([cond.shape[0]]).to(cond) + ) + + # Add noise to the conditioning signal + cond = cond + torch.randn_like(cond) * aug_level[:, None, None] + + # Get embedding for noise cond level, reusing timestamp_embed + aug_level_embed = expand_to_planes( + self.timestep_embed(aug_level[:, None]), x.shape + ) + + inputs.append(aug_level_embed) + + inputs.append(cond) + + outputs = self.net(torch.cat(inputs, dim=1)) + + return outputs + + +class DiTWrapper(ConditionedDiffusionModel): + def __init__(self, *args, **kwargs): + super().__init__( + supports_cross_attention=True, + supports_global_cond=False, + supports_input_concat=False, + ) + + self.model = DiffusionTransformer(*args, **kwargs) + + with torch.no_grad(): + for param in self.model.parameters(): + param *= 0.5 + + def forward( + self, + x, + t, + cross_attn_cond=None, + cross_attn_mask=None, + negative_cross_attn_cond=None, + negative_cross_attn_mask=None, + input_concat_cond=None, + negative_input_concat_cond=None, + global_cond=None, + negative_global_cond=None, + prepend_cond=None, + prepend_cond_mask=None, + cfg_scale=1.0, + cfg_dropout_prob: float = 0.0, + batch_cfg: bool = True, + rescale_cfg: bool = False, + scale_phi: float = 0.0, + **kwargs, + ): + assert batch_cfg, "batch_cfg must be True for DiTWrapper" + # assert negative_input_concat_cond is None, "negative_input_concat_cond is not supported for DiTWrapper" + + return self.model( + x, + t, + cross_attn_cond=cross_attn_cond, + cross_attn_cond_mask=cross_attn_mask, + negative_cross_attn_cond=negative_cross_attn_cond, + negative_cross_attn_mask=negative_cross_attn_mask, + input_concat_cond=input_concat_cond, + prepend_cond=prepend_cond, + prepend_cond_mask=prepend_cond_mask, + cfg_scale=cfg_scale, + cfg_dropout_prob=cfg_dropout_prob, + scale_phi=scale_phi, + global_embed=global_cond, + **kwargs, + ) + + +class DiTUncondWrapper(DiffusionModel): + def __init__(self, in_channels, *args, **kwargs): + super().__init__() + + self.model = DiffusionTransformer(io_channels=in_channels, *args, **kwargs) + + self.io_channels = in_channels + + with torch.no_grad(): + for param in self.model.parameters(): + param *= 0.5 + + def forward(self, x, t, **kwargs): + return self.model(x, t, **kwargs) + + +def create_diffusion_uncond_from_config(config: tp.Dict[str, tp.Any]): + diffusion_uncond_config = config["model"] + + model_type = diffusion_uncond_config.get("type", None) + + diffusion_config = diffusion_uncond_config.get("config", {}) + + assert model_type is not None, "Must specify model type in config" + + pretransform = diffusion_uncond_config.get("pretransform", None) + + sample_size = config.get("sample_size", None) + assert sample_size is not None, "Must specify sample size in config" + + sample_rate = config.get("sample_rate", None) + assert sample_rate is not None, "Must specify sample rate in config" + + if pretransform is not None: + pretransform = create_pretransform_from_config(pretransform, sample_rate) + min_input_length = pretransform.downsampling_ratio + else: + min_input_length = 1 + + if model_type == "DAU1d": + model = DiffusionAttnUnet1D(**diffusion_config) + + elif model_type == "adp_uncond_1d": + model = UNet1DUncondWrapper(**diffusion_config) + + elif model_type == "dit": + model = DiTUncondWrapper(**diffusion_config) + + else: + raise NotImplementedError(f"Unknown model type: {model_type}") + + return DiffusionModelWrapper( + model, + io_channels=model.io_channels, + sample_size=sample_size, + sample_rate=sample_rate, + pretransform=pretransform, + min_input_length=min_input_length, + ) + + +def create_diffusion_cond_from_config(config: tp.Dict[str, tp.Any]): + model_config = config["model"] + + model_type = config["model_type"] + + diffusion_config = model_config.get("diffusion", None) + assert diffusion_config is not None, "Must specify diffusion config" + + diffusion_model_type = diffusion_config.get("type", None) + assert diffusion_model_type is not None, "Must specify diffusion model type" + + diffusion_model_config = diffusion_config.get("config", None) + assert diffusion_model_config is not None, "Must specify diffusion model config" + + if diffusion_model_type == "adp_cfg_1d": + diffusion_model = UNetCFG1DWrapper(**diffusion_model_config) + elif diffusion_model_type == "adp_1d": + diffusion_model = UNet1DCondWrapper(**diffusion_model_config) + elif diffusion_model_type == "dit": + diffusion_model = DiTWrapper(**diffusion_model_config) + + io_channels = model_config.get("io_channels", None) + assert io_channels is not None, "Must specify io_channels in model config" + + sample_rate = config.get("sample_rate", None) + assert sample_rate is not None, "Must specify sample_rate in config" + + diffusion_objective = diffusion_config.get("diffusion_objective", "v") + + conditioning_config = model_config.get("conditioning", None) + + conditioner = None + if conditioning_config is not None: + conditioner = create_multi_conditioner_from_conditioning_config( + conditioning_config + ) + + cross_attention_ids = diffusion_config.get("cross_attention_cond_ids", []) + global_cond_ids = diffusion_config.get("global_cond_ids", []) + input_concat_ids = diffusion_config.get("input_concat_ids", []) + prepend_cond_ids = diffusion_config.get("prepend_cond_ids", []) + + pretransform = model_config.get("pretransform", None) + + if pretransform is not None: + pretransform = create_pretransform_from_config(pretransform, sample_rate) + min_input_length = pretransform.downsampling_ratio + else: + min_input_length = 1 + + if diffusion_model_type == "adp_cfg_1d" or diffusion_model_type == "adp_1d": + min_input_length *= np.prod(diffusion_model_config["factors"]) + elif diffusion_model_type == "dit": + min_input_length *= diffusion_model.model.patch_size + + # Get the proper wrapper class + + extra_kwargs = {} + + if model_type == "diffusion_cond" or model_type == "diffusion_cond_inpaint": + wrapper_fn = ConditionedDiffusionModelWrapper + + extra_kwargs["diffusion_objective"] = diffusion_objective + + elif model_type == "diffusion_prior": + prior_type = model_config.get("prior_type", None) + assert prior_type is not None, ( + "Must specify prior_type in diffusion prior model config" + ) + + if prior_type == "mono_stereo": + from .diffusion_prior import MonoToStereoDiffusionPrior + + wrapper_fn = MonoToStereoDiffusionPrior + + return wrapper_fn( + diffusion_model, + conditioner, + min_input_length=min_input_length, + sample_rate=sample_rate, + cross_attn_cond_ids=cross_attention_ids, + global_cond_ids=global_cond_ids, + input_concat_ids=input_concat_ids, + prepend_cond_ids=prepend_cond_ids, + pretransform=pretransform, + io_channels=io_channels, + **extra_kwargs, + ) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/dit.py b/src/YingMusicSinger/utils/stable_audio_tools/dit.py new file mode 100755 index 0000000000000000000000000000000000000000..17361d073ded5d6347a2261bfade62249c69b457 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/dit.py @@ -0,0 +1,451 @@ +import typing as tp + +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F +from x_transformers import ContinuousTransformerWrapper, Encoder + +from .blocks import FourierFeatures +from .transformer import ContinuousTransformer + + +class DiffusionTransformer(nn.Module): + def __init__( + self, + io_channels=32, + patch_size=1, + embed_dim=768, + cond_token_dim=0, + project_cond_tokens=True, + global_cond_dim=0, + project_global_cond=True, + input_concat_dim=0, + prepend_cond_dim=0, + depth=12, + num_heads=8, + transformer_type: tp.Literal[ + "x-transformers", "continuous_transformer" + ] = "x-transformers", + global_cond_type: tp.Literal["prepend", "adaLN"] = "prepend", + **kwargs, + ): + super().__init__() + + self.cond_token_dim = cond_token_dim + + # Timestep embeddings + timestep_features_dim = 256 + + self.timestep_features = FourierFeatures(1, timestep_features_dim) + + self.to_timestep_embed = nn.Sequential( + nn.Linear(timestep_features_dim, embed_dim, bias=True), + nn.SiLU(), + nn.Linear(embed_dim, embed_dim, bias=True), + ) + + if cond_token_dim > 0: + # Conditioning tokens + + cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim + self.to_cond_embed = nn.Sequential( + nn.Linear(cond_token_dim, cond_embed_dim, bias=False), + nn.SiLU(), + nn.Linear(cond_embed_dim, cond_embed_dim, bias=False), + ) + else: + cond_embed_dim = 0 + + if global_cond_dim > 0: + # Global conditioning + global_embed_dim = global_cond_dim if not project_global_cond else embed_dim + self.to_global_embed = nn.Sequential( + nn.Linear(global_cond_dim, global_embed_dim, bias=False), + nn.SiLU(), + nn.Linear(global_embed_dim, global_embed_dim, bias=False), + ) + + if prepend_cond_dim > 0: + # Prepend conditioning + self.to_prepend_embed = nn.Sequential( + nn.Linear(prepend_cond_dim, embed_dim, bias=False), + nn.SiLU(), + nn.Linear(embed_dim, embed_dim, bias=False), + ) + + self.input_concat_dim = input_concat_dim + + dim_in = io_channels + self.input_concat_dim + + self.patch_size = patch_size + + # Transformer + + self.transformer_type = transformer_type + + self.global_cond_type = global_cond_type + + if self.transformer_type == "x-transformers": + self.transformer = ContinuousTransformerWrapper( + dim_in=dim_in * patch_size, + dim_out=io_channels * patch_size, + max_seq_len=0, # Not relevant without absolute positional embeds + attn_layers=Encoder( + dim=embed_dim, + depth=depth, + heads=num_heads, + attn_flash=True, + cross_attend=cond_token_dim > 0, + dim_context=None if cond_embed_dim == 0 else cond_embed_dim, + zero_init_branch_output=True, + use_abs_pos_emb=False, + rotary_pos_emb=True, + ff_swish=True, + ff_glu=True, + **kwargs, + ), + ) + + elif self.transformer_type == "continuous_transformer": + global_dim = None + + if self.global_cond_type == "adaLN": + # The global conditioning is projected to the embed_dim already at this point + global_dim = embed_dim + + self.transformer = ContinuousTransformer( + dim=embed_dim, + depth=depth, + dim_heads=embed_dim // num_heads, + dim_in=dim_in * patch_size, + dim_out=io_channels * patch_size, + cross_attend=cond_token_dim > 0, + cond_token_dim=cond_embed_dim, + global_cond_dim=global_dim, + **kwargs, + ) + + else: + raise ValueError(f"Unknown transformer type: {self.transformer_type}") + + self.preprocess_conv = nn.Conv1d(dim_in, dim_in, 1, bias=False) + nn.init.zeros_(self.preprocess_conv.weight) + self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) + nn.init.zeros_(self.postprocess_conv.weight) + + def _forward( + self, + x, + t, + mask=None, + cross_attn_cond=None, + cross_attn_cond_mask=None, + input_concat_cond=None, + global_embed=None, + prepend_cond=None, + prepend_cond_mask=None, + return_info=False, + **kwargs, + ): + if cross_attn_cond is not None: + cross_attn_cond = self.to_cond_embed(cross_attn_cond) + + if global_embed is not None: + # Project the global conditioning to the embedding dimension + global_embed = self.to_global_embed(global_embed) + + prepend_inputs = None + prepend_mask = None + prepend_length = 0 + if prepend_cond is not None: + # Project the prepend conditioning to the embedding dimension + prepend_cond = self.to_prepend_embed(prepend_cond) + + prepend_inputs = prepend_cond + if prepend_cond_mask is not None: + prepend_mask = prepend_cond_mask + + if input_concat_cond is not None: + # Interpolate input_concat_cond to the same length as x + if input_concat_cond.shape[2] != x.shape[2]: + input_concat_cond = F.interpolate( + input_concat_cond, (x.shape[2],), mode="nearest" + ) + + x = torch.cat([x, input_concat_cond], dim=1) + + # Get the batch of timestep embeddings + timestep_embed = self.to_timestep_embed( + self.timestep_features(t[:, None]) + ) # (b, embed_dim) + + # Timestep embedding is considered a global embedding. Add to the global conditioning if it exists + if global_embed is not None: + global_embed = global_embed + timestep_embed + else: + global_embed = timestep_embed + + # Add the global_embed to the prepend inputs if there is no global conditioning support in the transformer + if self.global_cond_type == "prepend": + if prepend_inputs is None: + # Prepend inputs are just the global embed, and the mask is all ones + prepend_inputs = global_embed.unsqueeze(1) + prepend_mask = torch.ones( + (x.shape[0], 1), device=x.device, dtype=torch.bool + ) + else: + # Prepend inputs are the prepend conditioning + the global embed + prepend_inputs = torch.cat( + [prepend_inputs, global_embed.unsqueeze(1)], dim=1 + ) + prepend_mask = torch.cat( + [ + prepend_mask, + torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool), + ], + dim=1, + ) + + prepend_length = prepend_inputs.shape[1] + + x = self.preprocess_conv(x) + x + + x = rearrange(x, "b c t -> b t c") + + extra_args = {} + + if self.global_cond_type == "adaLN": + extra_args["global_cond"] = global_embed + + if self.patch_size > 1: + x = rearrange(x, "b (t p) c -> b t (c p)", p=self.patch_size) + + if self.transformer_type == "x-transformers": + output = self.transformer( + x, + prepend_embeds=prepend_inputs, + context=cross_attn_cond, + context_mask=cross_attn_cond_mask, + mask=mask, + prepend_mask=prepend_mask, + **extra_args, + **kwargs, + ) + elif self.transformer_type == "continuous_transformer": + output = self.transformer( + x, + prepend_embeds=prepend_inputs, + context=cross_attn_cond, + context_mask=cross_attn_cond_mask, + mask=mask, + prepend_mask=prepend_mask, + return_info=return_info, + **extra_args, + **kwargs, + ) + + if return_info: + output, info = output + elif self.transformer_type == "mm_transformer": + output = self.transformer( + x, + context=cross_attn_cond, + mask=mask, + context_mask=cross_attn_cond_mask, + **extra_args, + **kwargs, + ) + + output = rearrange(output, "b t c -> b c t")[:, :, prepend_length:] + + if self.patch_size > 1: + output = rearrange(output, "b (c p) t -> b c (t p)", p=self.patch_size) + + output = self.postprocess_conv(output) + output + + if return_info: + return output, info + + return output + + def forward( + self, + x, + t, + cross_attn_cond=None, + cross_attn_cond_mask=None, + negative_cross_attn_cond=None, + negative_cross_attn_mask=None, + input_concat_cond=None, + global_embed=None, + negative_global_embed=None, + prepend_cond=None, + prepend_cond_mask=None, + cfg_scale=1.0, + cfg_dropout_prob=0.0, + causal=False, + scale_phi=0.0, + mask=None, + return_info=False, + **kwargs, + ): + assert causal == False, "Causal mode is not supported for DiffusionTransformer" + + if cross_attn_cond_mask is not None: + cross_attn_cond_mask = cross_attn_cond_mask.bool() + + cross_attn_cond_mask = None # Temporarily disabling conditioning masks due to kernel issue for flash attention + + if prepend_cond_mask is not None: + prepend_cond_mask = prepend_cond_mask.bool() + + # CFG dropout + if cfg_dropout_prob > 0.0: + if cross_attn_cond is not None: + null_embed = torch.zeros_like( + cross_attn_cond, device=cross_attn_cond.device + ) + dropout_mask = torch.bernoulli( + torch.full( + (cross_attn_cond.shape[0], 1, 1), + cfg_dropout_prob, + device=cross_attn_cond.device, + ) + ).to(torch.bool) + cross_attn_cond = torch.where(dropout_mask, null_embed, cross_attn_cond) + + if prepend_cond is not None: + null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device) + dropout_mask = torch.bernoulli( + torch.full( + (prepend_cond.shape[0], 1, 1), + cfg_dropout_prob, + device=prepend_cond.device, + ) + ).to(torch.bool) + prepend_cond = torch.where(dropout_mask, null_embed, prepend_cond) + + if cfg_scale != 1.0 and ( + cross_attn_cond is not None or prepend_cond is not None + ): + # Classifier-free guidance + # Concatenate conditioned and unconditioned inputs on the batch dimension + batch_inputs = torch.cat([x, x], dim=0) + batch_timestep = torch.cat([t, t], dim=0) + + if global_embed is not None: + batch_global_cond = torch.cat([global_embed, global_embed], dim=0) + else: + batch_global_cond = None + + if input_concat_cond is not None: + batch_input_concat_cond = torch.cat( + [input_concat_cond, input_concat_cond], dim=0 + ) + else: + batch_input_concat_cond = None + + batch_cond = None + batch_cond_masks = None + + # Handle CFG for cross-attention conditioning + if cross_attn_cond is not None: + null_embed = torch.zeros_like( + cross_attn_cond, device=cross_attn_cond.device + ) + + # For negative cross-attention conditioning, replace the null embed with the negative cross-attention conditioning + if negative_cross_attn_cond is not None: + # If there's a negative cross-attention mask, set the masked tokens to the null embed + if negative_cross_attn_mask is not None: + negative_cross_attn_mask = negative_cross_attn_mask.to( + torch.bool + ).unsqueeze(2) + + negative_cross_attn_cond = torch.where( + negative_cross_attn_mask, + negative_cross_attn_cond, + null_embed, + ) + + batch_cond = torch.cat( + [cross_attn_cond, negative_cross_attn_cond], dim=0 + ) + + else: + batch_cond = torch.cat([cross_attn_cond, null_embed], dim=0) + + if cross_attn_cond_mask is not None: + batch_cond_masks = torch.cat( + [cross_attn_cond_mask, cross_attn_cond_mask], dim=0 + ) + + batch_prepend_cond = None + batch_prepend_cond_mask = None + + if prepend_cond is not None: + null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device) + + batch_prepend_cond = torch.cat([prepend_cond, null_embed], dim=0) + + if prepend_cond_mask is not None: + batch_prepend_cond_mask = torch.cat( + [prepend_cond_mask, prepend_cond_mask], dim=0 + ) + + if mask is not None: + batch_masks = torch.cat([mask, mask], dim=0) + else: + batch_masks = None + + batch_output = self._forward( + batch_inputs, + batch_timestep, + cross_attn_cond=batch_cond, + cross_attn_cond_mask=batch_cond_masks, + mask=batch_masks, + input_concat_cond=batch_input_concat_cond, + global_embed=batch_global_cond, + prepend_cond=batch_prepend_cond, + prepend_cond_mask=batch_prepend_cond_mask, + return_info=return_info, + **kwargs, + ) + + if return_info: + batch_output, info = batch_output + + cond_output, uncond_output = torch.chunk(batch_output, 2, dim=0) + cfg_output = uncond_output + (cond_output - uncond_output) * cfg_scale + + # CFG Rescale + if scale_phi != 0.0: + cond_out_std = cond_output.std(dim=1, keepdim=True) + out_cfg_std = cfg_output.std(dim=1, keepdim=True) + output = ( + scale_phi * (cfg_output * (cond_out_std / out_cfg_std)) + + (1 - scale_phi) * cfg_output + ) + else: + output = cfg_output + + if return_info: + return output, info + + return output + + else: + return self._forward( + x, + t, + cross_attn_cond=cross_attn_cond, + cross_attn_cond_mask=cross_attn_cond_mask, + input_concat_cond=input_concat_cond, + global_embed=global_embed, + prepend_cond=prepend_cond, + prepend_cond_mask=prepend_cond_mask, + mask=mask, + return_info=return_info, + **kwargs, + ) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/factory.py b/src/YingMusicSinger/utils/stable_audio_tools/factory.py new file mode 100755 index 0000000000000000000000000000000000000000..49dc7ae8669e73b5e2221755643c94f3578c5298 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/factory.py @@ -0,0 +1,185 @@ +import json + + +def create_model_from_config(model_config): + model_type = model_config.get("model_type", None) + + assert model_type is not None, "model_type must be specified in model config" + + if model_type == "autoencoder": + from .autoencoders import create_autoencoder_from_config + + return create_autoencoder_from_config(model_config) + elif model_type == "diffusion_uncond": + from .diffusion import create_diffusion_uncond_from_config + + return create_diffusion_uncond_from_config(model_config) + elif ( + model_type == "diffusion_cond" + or model_type == "diffusion_cond_inpaint" + or model_type == "diffusion_prior" + ): + from .diffusion import create_diffusion_cond_from_config + + return create_diffusion_cond_from_config(model_config) + elif model_type == "diffusion_autoencoder": + from .autoencoders import create_diffAE_from_config + + return create_diffAE_from_config(model_config) + elif model_type == "lm": + from .lm import create_audio_lm_from_config + + return create_audio_lm_from_config(model_config) + else: + raise NotImplementedError(f"Unknown model type: {model_type}") + + +def create_model_from_config_path(model_config_path): + with open(model_config_path) as f: + model_config = json.load(f) + + return create_model_from_config(model_config) + + +def create_pretransform_from_config(pretransform_config, sample_rate): + pretransform_type = pretransform_config.get("type", None) + + assert pretransform_type is not None, ( + "type must be specified in pretransform config" + ) + + if pretransform_type == "autoencoder": + from .autoencoders import create_autoencoder_from_config + from .pretransforms import AutoencoderPretransform + + # Create fake top-level config to pass sample rate to autoencoder constructor + # This is a bit of a hack but it keeps us from re-defining the sample rate in the config + autoencoder_config = { + "sample_rate": sample_rate, + "model": pretransform_config["config"], + } + autoencoder = create_autoencoder_from_config(autoencoder_config) + + scale = pretransform_config.get("scale", 1.0) + model_half = pretransform_config.get("model_half", False) + iterate_batch = pretransform_config.get("iterate_batch", False) + chunked = pretransform_config.get("chunked", False) + + pretransform = AutoencoderPretransform( + autoencoder, + scale=scale, + model_half=model_half, + iterate_batch=iterate_batch, + chunked=chunked, + ) + elif pretransform_type == "wavelet": + from .pretransforms import WaveletPretransform + + wavelet_config = pretransform_config["config"] + channels = wavelet_config["channels"] + levels = wavelet_config["levels"] + wavelet = wavelet_config["wavelet"] + + pretransform = WaveletPretransform(channels, levels, wavelet) + elif pretransform_type == "pqmf": + from .pretransforms import PQMFPretransform + + pqmf_config = pretransform_config["config"] + pretransform = PQMFPretransform(**pqmf_config) + elif pretransform_type == "dac_pretrained": + from .pretransforms import PretrainedDACPretransform + + pretrained_dac_config = pretransform_config["config"] + pretransform = PretrainedDACPretransform(**pretrained_dac_config) + elif pretransform_type == "audiocraft_pretrained": + from .pretransforms import AudiocraftCompressionPretransform + + audiocraft_config = pretransform_config["config"] + pretransform = AudiocraftCompressionPretransform(**audiocraft_config) + else: + raise NotImplementedError(f"Unknown pretransform type: {pretransform_type}") + + enable_grad = pretransform_config.get("enable_grad", False) + pretransform.enable_grad = enable_grad + + pretransform.eval().requires_grad_(pretransform.enable_grad) + + return pretransform + + +def create_bottleneck_from_config(bottleneck_config): + bottleneck_type = bottleneck_config.get("type", None) + + assert bottleneck_type is not None, "type must be specified in bottleneck config" + + if bottleneck_type == "tanh": + from .bottleneck import TanhBottleneck + + bottleneck = TanhBottleneck() + elif bottleneck_type == "vae": + from .bottleneck import VAEBottleneck + + bottleneck = VAEBottleneck() + elif bottleneck_type == "rvq": + from .bottleneck import RVQBottleneck + + quantizer_params = { + "dim": 128, + "codebook_size": 1024, + "num_quantizers": 8, + "decay": 0.99, + "kmeans_init": True, + "kmeans_iters": 50, + "threshold_ema_dead_code": 2, + } + + quantizer_params.update(bottleneck_config["config"]) + + bottleneck = RVQBottleneck(**quantizer_params) + elif bottleneck_type == "dac_rvq": + from .bottleneck import DACRVQBottleneck + + bottleneck = DACRVQBottleneck(**bottleneck_config["config"]) + + elif bottleneck_type == "rvq_vae": + from .bottleneck import RVQVAEBottleneck + + quantizer_params = { + "dim": 128, + "codebook_size": 1024, + "num_quantizers": 8, + "decay": 0.99, + "kmeans_init": True, + "kmeans_iters": 50, + "threshold_ema_dead_code": 2, + } + + quantizer_params.update(bottleneck_config["config"]) + + bottleneck = RVQVAEBottleneck(**quantizer_params) + + elif bottleneck_type == "dac_rvq_vae": + from .bottleneck import DACRVQVAEBottleneck + + bottleneck = DACRVQVAEBottleneck(**bottleneck_config["config"]) + elif bottleneck_type == "l2_norm": + from .bottleneck import L2Bottleneck + + bottleneck = L2Bottleneck() + elif bottleneck_type == "wasserstein": + from .bottleneck import WassersteinBottleneck + + bottleneck = WassersteinBottleneck(**bottleneck_config.get("config", {})) + elif bottleneck_type == "fsq": + from .bottleneck import FSQBottleneck + + bottleneck = FSQBottleneck(**bottleneck_config["config"]) + else: + raise NotImplementedError(f"Unknown bottleneck type: {bottleneck_type}") + + requires_grad = bottleneck_config.get("requires_grad", True) + if not requires_grad: + for param in bottleneck.parameters(): + param.requires_grad = False + + return bottleneck diff --git a/src/YingMusicSinger/utils/stable_audio_tools/pretransforms.py b/src/YingMusicSinger/utils/stable_audio_tools/pretransforms.py new file mode 100755 index 0000000000000000000000000000000000000000..ecdce4a71e62ab9f3e87d8c3025b7661583cc33a --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/pretransforms.py @@ -0,0 +1,425 @@ +import torch +from einops import rearrange +from torch import nn + + +class Pretransform(nn.Module): + def __init__(self, enable_grad, io_channels, is_discrete): + super().__init__() + + self.is_discrete = is_discrete + self.io_channels = io_channels + self.encoded_channels = None + self.downsampling_ratio = None + + self.enable_grad = enable_grad + + def encode(self, x): + raise NotImplementedError + + def decode(self, z): + raise NotImplementedError + + def tokenize(self, x): + raise NotImplementedError + + def decode_tokens(self, tokens): + raise NotImplementedError + + +class AutoencoderPretransform(Pretransform): + def __init__( + self, model, scale=1.0, model_half=False, iterate_batch=False, chunked=False + ): + super().__init__( + enable_grad=False, + io_channels=model.io_channels, + is_discrete=model.bottleneck is not None and model.bottleneck.is_discrete, + ) + self.model = model + self.model.requires_grad_(False).eval() + self.scale = scale + self.downsampling_ratio = model.downsampling_ratio + self.io_channels = model.io_channels + self.sample_rate = model.sample_rate + + self.model_half = model_half + self.iterate_batch = iterate_batch + + self.encoded_channels = model.latent_dim + self.latent_dim = model.latent_dim + + self.chunked = chunked + self.num_quantizers = ( + model.bottleneck.num_quantizers + if model.bottleneck is not None and model.bottleneck.is_discrete + else None + ) + self.codebook_size = ( + model.bottleneck.codebook_size + if model.bottleneck is not None and model.bottleneck.is_discrete + else None + ) + + if self.model_half: + self.model.half() + + def encode(self, x, **kwargs): + if self.model_half: + x = x.half() + self.model.to(torch.float16) + + encoded = self.model.encode_audio( + x, chunked=self.chunked, iterate_batch=self.iterate_batch, **kwargs + ) + + if self.model_half: + encoded = encoded.float() + + return encoded / self.scale + + def encode_audio(self, audio, chunked=False, overlap=32, chunk_size=128, **kwargs): + """ + Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder. + If chunked is True, split the audio into chunks of a given maximum size chunk_size, with given overlap. + Overlap and chunk_size params are both measured in number of latents (not audio samples) + # and therefore you likely could use the same values with decode_audio. + A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. + Every autoencoder will have a different receptive field size, and thus ideal overlap. + You can determine it empirically by diffing unchunked vs chunked output and looking at maximum diff. + The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks. + Smaller chunk_size uses less memory, but more compute. + The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version + For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks + """ + if not chunked: + # default behavior. Encode the entire audio in parallel + return self.encode(audio, **kwargs) + else: + # CHUNKED ENCODING + # samples_per_latent is just the downsampling ratio (which is also the upsampling ratio) + samples_per_latent = self.downsampling_ratio + total_size = audio.shape[2] # in samples + batch_size = audio.shape[0] + chunk_size *= samples_per_latent # converting metric in latents to samples + overlap *= samples_per_latent # converting metric in latents to samples + hop_size = chunk_size - overlap + chunks = [] + for i in range(0, total_size - chunk_size + 1, hop_size): + chunk = audio[:, :, i : i + chunk_size] + chunks.append(chunk) + if i + chunk_size != total_size: + # Final chunk + chunk = audio[:, :, -chunk_size:] + chunks.append(chunk) + chunks = torch.stack(chunks) + num_chunks = chunks.shape[0] + # Note: y_size might be a different value from the latent length used in diffusion training + # because we can encode audio of varying lengths + # However, the audio should've been padded to a multiple of samples_per_latent by now. + y_size = total_size // samples_per_latent + # Create an empty latent, we will populate it with chunks as we encode them + y_final = torch.zeros((batch_size, self.latent_dim, y_size)).to( + audio.device + ) + for i in range(num_chunks): + x_chunk = chunks[i, :] + # encode the chunk + y_chunk = self.encode(x_chunk) + # figure out where to put the audio along the time domain + if i == num_chunks - 1: + # final chunk always goes at the end + t_end = y_size + t_start = t_end - y_chunk.shape[2] + else: + t_start = i * hop_size // samples_per_latent + t_end = t_start + chunk_size // samples_per_latent + # remove the edges of the overlaps + ol = overlap // samples_per_latent // 2 + chunk_start = 0 + chunk_end = y_chunk.shape[2] + if i > 0: + # no overlap for the start of the first chunk + t_start += ol + chunk_start += ol + if i < num_chunks - 1: + # no overlap for the end of the last chunk + t_end -= ol + chunk_end -= ol + # paste the chunked audio into our y_final output audio + y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end] + return y_final + + def decode(self, z, **kwargs): + z = z * self.scale + + if self.model_half: + z = z.half() + self.model.to(torch.float16) + + decoded = self.model.decode_audio( + z, chunked=self.chunked, iterate_batch=self.iterate_batch, **kwargs + ) + + if self.model_half: + decoded = decoded.float() + + return decoded + + def decode_audio( + self, latents, chunked=False, overlap=32, chunk_size=128, **kwargs + ): + if not chunked: + # default behavior. Decode the entire latent in parallel + return self.decode(latents, **kwargs) + else: + # chunked decoding + hop_size = chunk_size - overlap + total_size = latents.shape[2] + batch_size = latents.shape[0] + chunks = [] + i = 0 + for i in range(0, total_size - chunk_size + 1, hop_size): + chunk = latents[:, :, i : i + chunk_size] + chunks.append(chunk) + if i + chunk_size != total_size: + # Final chunk + chunk = latents[:, :, -chunk_size:] + chunks.append(chunk) + chunks = torch.stack(chunks) + num_chunks = chunks.shape[0] + # samples_per_latent is just the downsampling ratio + samples_per_latent = self.downsampling_ratio + # Create an empty waveform, we will populate it with chunks as decode them + y_size = total_size * samples_per_latent + y_final = torch.zeros((batch_size, self.io_channels, y_size)).to( + latents.device + ) + for i in range(num_chunks): + x_chunk = chunks[i, :] + # decode the chunk + y_chunk = self.decode(x_chunk) + # figure out where to put the audio along the time domain + if i == num_chunks - 1: + # final chunk always goes at the end + t_end = y_size + t_start = t_end - y_chunk.shape[2] + else: + t_start = i * hop_size * samples_per_latent + t_end = t_start + chunk_size * samples_per_latent + # remove the edges of the overlaps + ol = (overlap // 2) * samples_per_latent + chunk_start = 0 + chunk_end = y_chunk.shape[2] + if i > 0: + # no overlap for the start of the first chunk + t_start += ol + chunk_start += ol + if i < num_chunks - 1: + # no overlap for the end of the last chunk + t_end -= ol + chunk_end -= ol + # paste the chunked audio into our y_final output audio + y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end] + return y_final + + def tokenize(self, x, **kwargs): + assert self.model.is_discrete, "Cannot tokenize with a continuous model" + + _, info = self.model.encode(x, return_info=True, **kwargs) + + return info[self.model.bottleneck.tokens_id] + + def decode_tokens(self, tokens, **kwargs): + assert self.model.is_discrete, "Cannot decode tokens with a continuous model" + + return self.model.decode_tokens(tokens, **kwargs) + + def load_state_dict(self, state_dict, strict=True): + self.model.load_state_dict(state_dict, strict=strict) + + +class WaveletPretransform(Pretransform): + def __init__(self, channels, levels, wavelet): + super().__init__(enable_grad=False, io_channels=channels, is_discrete=False) + + from .wavelets import WaveletDecode1d, WaveletEncode1d + + self.encoder = WaveletEncode1d(channels, levels, wavelet) + self.decoder = WaveletDecode1d(channels, levels, wavelet) + + self.downsampling_ratio = 2**levels + self.io_channels = channels + self.encoded_channels = channels * self.downsampling_ratio + + def encode(self, x): + return self.encoder(x) + + def decode(self, z): + return self.decoder(z) + + +class PQMFPretransform(Pretransform): + def __init__(self, attenuation=100, num_bands=16): + # TODO: Fix PQMF to take in in-channels + super().__init__(enable_grad=False, io_channels=1, is_discrete=False) + from .pqmf import PQMF + + self.pqmf = PQMF(attenuation, num_bands) + + def encode(self, x): + # x is (Batch x Channels x Time) + x = self.pqmf.forward(x) + # pqmf.forward returns (Batch x Channels x Bands x Time) + # but Pretransform needs Batch x Channels x Time + # so concatenate channels and bands into one axis + return rearrange(x, "b c n t -> b (c n) t") + + def decode(self, x): + # x is (Batch x (Channels Bands) x Time), convert back to (Batch x Channels x Bands x Time) + x = rearrange(x, "b (c n) t -> b c n t", n=self.pqmf.num_bands) + # returns (Batch x Channels x Time) + return self.pqmf.inverse(x) + + +class PretrainedDACPretransform(Pretransform): + def __init__( + self, + model_type="44khz", + model_bitrate="8kbps", + scale=1.0, + quantize_on_decode: bool = True, + chunked=True, + ): + super().__init__(enable_grad=False, io_channels=1, is_discrete=True) + + import dac + + model_path = dac.utils.download( + model_type=model_type, model_bitrate=model_bitrate + ) + + self.model = dac.DAC.load(model_path) + + self.quantize_on_decode = quantize_on_decode + + if model_type == "44khz": + self.downsampling_ratio = 512 + else: + self.downsampling_ratio = 320 + + self.io_channels = 1 + + self.scale = scale + + self.chunked = chunked + + self.encoded_channels = self.model.latent_dim + + self.num_quantizers = self.model.n_codebooks + + self.codebook_size = self.model.codebook_size + + def encode(self, x): + latents = self.model.encoder(x) + + if self.quantize_on_decode: + output = latents + else: + z, _, _, _, _ = self.model.quantizer( + latents, n_quantizers=self.model.n_codebooks + ) + output = z + + if self.scale != 1.0: + output = output / self.scale + + return output + + def decode(self, z): + if self.scale != 1.0: + z = z * self.scale + + if self.quantize_on_decode: + z, _, _, _, _ = self.model.quantizer(z, n_quantizers=self.model.n_codebooks) + + return self.model.decode(z) + + def tokenize(self, x): + return self.model.encode(x)[1] + + def decode_tokens(self, tokens): + latents = self.model.quantizer.from_codes(tokens) + return self.model.decode(latents) + + +class AudiocraftCompressionPretransform(Pretransform): + def __init__( + self, + model_type="facebook/encodec_32khz", + scale=1.0, + quantize_on_decode: bool = True, + ): + super().__init__(enable_grad=False, io_channels=1, is_discrete=True) + + try: + from audiocraft.models import CompressionModel + except ImportError: + raise ImportError( + "Audiocraft is not installed. Please install audiocraft to use Audiocraft models." + ) + + self.model = CompressionModel.get_pretrained(model_type) + + self.quantize_on_decode = quantize_on_decode + + self.downsampling_ratio = round(self.model.sample_rate / self.model.frame_rate) + + self.sample_rate = self.model.sample_rate + + self.io_channels = self.model.channels + + self.scale = scale + + # self.encoded_channels = self.model.latent_dim + + self.num_quantizers = self.model.num_codebooks + + self.codebook_size = self.model.cardinality + + self.model.to(torch.float16).eval().requires_grad_(False) + + def encode(self, x): + assert False, "Audiocraft compression models do not support continuous encoding" + + # latents = self.model.encoder(x) + + # if self.quantize_on_decode: + # output = latents + # else: + # z, _, _, _, _ = self.model.quantizer(latents, n_quantizers=self.model.n_codebooks) + # output = z + + # if self.scale != 1.0: + # output = output / self.scale + + # return output + + def decode(self, z): + assert False, "Audiocraft compression models do not support continuous decoding" + + # if self.scale != 1.0: + # z = z * self.scale + + # if self.quantize_on_decode: + # z, _, _, _, _ = self.model.quantizer(z, n_quantizers=self.model.n_codebooks) + + # return self.model.decode(z) + + def tokenize(self, x): + with torch.cuda.amp.autocast(enabled=False): + return self.model.encode(x.to(torch.float16))[0] + + def decode_tokens(self, tokens): + with torch.cuda.amp.autocast(enabled=False): + return self.model.decode(tokens) diff --git a/src/YingMusicSinger/utils/stable_audio_tools/transformer.py b/src/YingMusicSinger/utils/stable_audio_tools/transformer.py new file mode 100755 index 0000000000000000000000000000000000000000..d0570a1454eb79d6caf091cdab4e1c3c9ccc2de7 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/transformer.py @@ -0,0 +1,910 @@ +from functools import reduce +from typing import Callable, Literal + +import torch +import torch.nn.functional as F +from einops import rearrange +from einops.layers.torch import Rearrange +from packaging import version +from torch import einsum, nn +from torch.cuda.amp import autocast + +try: + from flash_attn import flash_attn_func, flash_attn_kvpacked_func +except ImportError as e: + print(e) + print("flash_attn not installed, disabling Flash Attention") + flash_attn_kvpacked_func = None + flash_attn_func = None + +try: + import natten +except ImportError: + natten = None + + +def checkpoint(function, *args, **kwargs): + kwargs.setdefault("use_reentrant", False) + return torch.utils.checkpoint.checkpoint(function, *args, **kwargs) + + +# Copied and modified from https://github.com/lucidrains/x-transformers/blob/main/x_transformers/attend.py under MIT License +# License can be found in LICENSES/LICENSE_XTRANSFORMERS.txt + + +def create_causal_mask(i, j, device): + return torch.ones((i, j), device=device, dtype=torch.bool).triu(j - i + 1) + + +def or_reduce(masks): + head, *body = masks + for rest in body: + head = head | rest + return head + + +# positional embeddings + + +class AbsolutePositionalEmbedding(nn.Module): + def __init__(self, dim, max_seq_len): + super().__init__() + self.scale = dim**-0.5 + self.max_seq_len = max_seq_len + self.emb = nn.Embedding(max_seq_len, dim) + + def forward(self, x, pos=None, seq_start_pos=None): + seq_len, device = x.shape[1], x.device + assert seq_len <= self.max_seq_len, ( + f"you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}" + ) + + if pos is None: + pos = torch.arange(seq_len, device=device) + + if seq_start_pos is not None: + pos = (pos - seq_start_pos[..., None]).clamp(min=0) + + pos_emb = self.emb(pos) + pos_emb = pos_emb * self.scale + return pos_emb + + +class ScaledSinusoidalEmbedding(nn.Module): + def __init__(self, dim, theta=10000): + super().__init__() + assert (dim % 2) == 0, "dimension must be divisible by 2" + self.scale = nn.Parameter(torch.ones(1) * dim**-0.5) + + half_dim = dim // 2 + freq_seq = torch.arange(half_dim).float() / half_dim + inv_freq = theta**-freq_seq + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, x, pos=None, seq_start_pos=None): + seq_len, device = x.shape[1], x.device + + if pos is None: + pos = torch.arange(seq_len, device=device) + + if seq_start_pos is not None: + pos = pos - seq_start_pos[..., None] + + emb = einsum("i, j -> i j", pos, self.inv_freq) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb * self.scale + + +class RotaryEmbedding(nn.Module): + def __init__( + self, + dim, + use_xpos=False, + scale_base=512, + interpolation_factor=1.0, + base=10000, + base_rescale_factor=1.0, + ): + super().__init__() + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ + base *= base_rescale_factor ** (dim / (dim - 2)) + + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer("inv_freq", inv_freq) + + assert interpolation_factor >= 1.0 + self.interpolation_factor = interpolation_factor + + if not use_xpos: + self.register_buffer("scale", None) + return + + scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) + + self.scale_base = scale_base + self.register_buffer("scale", scale) + + def forward_from_seq_len(self, seq_len): + device = self.inv_freq.device + + t = torch.arange(seq_len, device=device) + return self.forward(t) + + @autocast(enabled=False) + def forward(self, t): + device = self.inv_freq.device + + t = t.to(torch.float32) + + t = t / self.interpolation_factor + + freqs = torch.einsum("i , j -> i j", t, self.inv_freq) + freqs = torch.cat((freqs, freqs), dim=-1) + + if self.scale is None: + return freqs, 1.0 + + power = ( + torch.arange(seq_len, device=device) - (seq_len // 2) + ) / self.scale_base + scale = self.scale ** rearrange(power, "n -> n 1") + scale = torch.cat((scale, scale), dim=-1) + + return freqs, scale + + +def rotate_half(x): + x = rearrange(x, "... (j d) -> ... j d", j=2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((-x2, x1), dim=-1) + + +@autocast(enabled=False) +def apply_rotary_pos_emb(t, freqs, scale=1): + out_dtype = t.dtype + + # cast to float32 if necessary for numerical stability + dtype = reduce(torch.promote_types, (t.dtype, freqs.dtype, torch.float32)) + rot_dim, seq_len = freqs.shape[-1], t.shape[-2] + freqs, t = freqs.to(dtype), t.to(dtype) + freqs = freqs[-seq_len:, :] + + if t.ndim == 4 and freqs.ndim == 3: + freqs = rearrange(freqs, "b n d -> b 1 n d") + + # partial rotary embeddings, Wang et al. GPT-J + t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:] + t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale) + + t, t_unrotated = t.to(out_dtype), t_unrotated.to(out_dtype) + + return torch.cat((t, t_unrotated), dim=-1) + + +# norms +class LayerNorm(nn.Module): + def __init__(self, dim, bias=False, fix_scale=False): + """ + bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less + """ + super().__init__() + + if fix_scale: + self.register_buffer("gamma", torch.ones(dim)) + else: + self.gamma = nn.Parameter(torch.ones(dim)) + + if bias: + self.beta = nn.Parameter(torch.zeros(dim)) + else: + self.register_buffer("beta", torch.zeros(dim)) + + def forward(self, x): + return F.layer_norm(x, x.shape[-1:], weight=self.gamma, bias=self.beta) + + +# feedforward + + +class GLU(nn.Module): + def __init__( + self, + dim_in, + dim_out, + activation: Callable, + use_conv=False, + conv_kernel_size=3, + ): + super().__init__() + self.act = activation + self.proj = ( + nn.Linear(dim_in, dim_out * 2) + if not use_conv + else nn.Conv1d( + dim_in, dim_out * 2, conv_kernel_size, padding=(conv_kernel_size // 2) + ) + ) + self.use_conv = use_conv + + def forward(self, x): + if self.use_conv: + x = rearrange(x, "b n d -> b d n") + x = self.proj(x) + x = rearrange(x, "b d n -> b n d") + else: + x = self.proj(x) + + x, gate = x.chunk(2, dim=-1) + return x * self.act(gate) + + +class FeedForward(nn.Module): + def __init__( + self, + dim, + dim_out=None, + mult=4, + no_bias=False, + glu=True, + use_conv=False, + conv_kernel_size=3, + zero_init_output=True, + ): + super().__init__() + inner_dim = int(dim * mult) + + # Default to SwiGLU + + activation = nn.SiLU() + + dim_out = dim if dim_out is None else dim_out + + if glu: + linear_in = GLU(dim, inner_dim, activation) + else: + linear_in = nn.Sequential( + Rearrange("b n d -> b d n") if use_conv else nn.Identity(), + nn.Linear(dim, inner_dim, bias=not no_bias) + if not use_conv + else nn.Conv1d( + dim, + inner_dim, + conv_kernel_size, + padding=(conv_kernel_size // 2), + bias=not no_bias, + ), + Rearrange("b n d -> b d n") if use_conv else nn.Identity(), + activation, + ) + + linear_out = ( + nn.Linear(inner_dim, dim_out, bias=not no_bias) + if not use_conv + else nn.Conv1d( + inner_dim, + dim_out, + conv_kernel_size, + padding=(conv_kernel_size // 2), + bias=not no_bias, + ) + ) + + # init last linear layer to 0 + if zero_init_output: + nn.init.zeros_(linear_out.weight) + if not no_bias: + nn.init.zeros_(linear_out.bias) + + self.ff = nn.Sequential( + linear_in, + Rearrange("b d n -> b n d") if use_conv else nn.Identity(), + linear_out, + Rearrange("b n d -> b d n") if use_conv else nn.Identity(), + ) + + def forward(self, x): + return self.ff(x) + + +class Attention(nn.Module): + def __init__( + self, + dim, + dim_heads=64, + dim_context=None, + causal=False, + zero_init_output=True, + qk_norm: Literal["l2", "ln", "none"] = "none", + natten_kernel_size=None, + ): + super().__init__() + self.dim = dim + self.dim_heads = dim_heads + self.causal = causal + + dim_kv = dim_context if dim_context is not None else dim + + self.num_heads = dim // dim_heads + self.kv_heads = dim_kv // dim_heads + + if dim_context is not None: + self.to_q = nn.Linear(dim, dim, bias=False) + self.to_kv = nn.Linear(dim_kv, dim_kv * 2, bias=False) + else: + self.to_qkv = nn.Linear(dim, dim * 3, bias=False) + + self.to_out = nn.Linear(dim, dim, bias=False) + + if zero_init_output: + nn.init.zeros_(self.to_out.weight) + + self.qk_norm = qk_norm + + if self.qk_norm == "ln": + self.q_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6) + self.k_norm = nn.LayerNorm(dim_heads, elementwise_affine=True, eps=1.0e-6) + + # Using 1d neighborhood attention + self.natten_kernel_size = natten_kernel_size + if natten_kernel_size is not None: + return + + self.use_pt_flash = torch.cuda.is_available() and version.parse( + torch.__version__ + ) >= version.parse("2.0.0") + + self.use_fa_flash = torch.cuda.is_available() and flash_attn_func is not None + + self.sdp_kwargs = dict( + enable_flash=True, enable_math=True, enable_mem_efficient=True + ) + + def flash_attn(self, q, k, v, mask=None, causal=None): + batch, heads, q_len, _, k_len, device = *q.shape, k.shape[-2], q.device + kv_heads = k.shape[1] + # Recommended for multi-query single-key-value attention by Tri Dao + # kv shape torch.Size([1, 512, 64]) -> torch.Size([1, 8, 512, 64]) + + if heads != kv_heads: + # Repeat interleave kv_heads to match q_heads + heads_per_kv_head = heads // kv_heads + k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim=1), (k, v)) + + if k.ndim == 3: + k = rearrange(k, "b ... -> b 1 ...").expand_as(q) + + if v.ndim == 3: + v = rearrange(v, "b ... -> b 1 ...").expand_as(q) + + causal = self.causal if causal is None else causal + + if q_len == 1 and causal: + causal = False + + if mask is not None: + assert mask.ndim == 4 + mask = mask.expand(batch, heads, q_len, k_len) + + # handle kv cache - this should be bypassable in updated flash attention 2 + + if k_len > q_len and causal: + causal_mask = self.create_causal_mask(q_len, k_len, device=device) + if mask is None: + mask = ~causal_mask + else: + mask = mask & ~causal_mask + causal = False + + # manually handle causal mask, if another mask was given + + row_is_entirely_masked = None + + if mask is not None and causal: + causal_mask = self.create_causal_mask(q_len, k_len, device=device) + mask = mask & ~causal_mask + + # protect against an entire row being masked out + + row_is_entirely_masked = ~mask.any(dim=-1) + mask[..., 0] = mask[..., 0] | row_is_entirely_masked + + causal = False + + with torch.backends.cuda.sdp_kernel(**self.sdp_kwargs): + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, is_causal=causal + ) + + # for a row that is entirely masked out, should zero out the output of that row token + + if row_is_entirely_masked is not None: + out = out.masked_fill(row_is_entirely_masked[..., None], 0.0) + + return out + + def forward( + self, + x, + context=None, + mask=None, + context_mask=None, + rotary_pos_emb=None, + causal=None, + ): + h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None + + kv_input = context if has_context else x + + if hasattr(self, "to_q"): + # Use separate linear projections for q and k/v + q = self.to_q(x) + q = rearrange(q, "b n (h d) -> b h n d", h=h) + + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=kv_h), (k, v)) + else: + # Use fused linear projection + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + q, k, v = map( + lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v) + ) + + # Normalize q and k for cosine sim attention + if self.qk_norm == "l2": + q = F.normalize(q, dim=-1) + k = F.normalize(k, dim=-1) + elif self.qk_norm == "ln": + q = self.q_norm(q) + k = self.k_norm(k) + + if rotary_pos_emb is not None and not has_context: + freqs, _ = rotary_pos_emb + + q_dtype = q.dtype + k_dtype = k.dtype + + q = q.to(torch.float32) + k = k.to(torch.float32) + freqs = freqs.to(torch.float32) + + q = apply_rotary_pos_emb(q, freqs) + k = apply_rotary_pos_emb(k, freqs) + + q = q.to(q_dtype) + k = k.to(k_dtype) + + input_mask = context_mask + + if input_mask is None and not has_context: + input_mask = mask + + # determine masking + masks = [] + final_attn_mask = None # The mask that will be applied to the attention matrix, taking all masks into account + + if input_mask is not None: + input_mask = rearrange(input_mask, "b j -> b 1 1 j") + masks.append(~input_mask) + + # Other masks will be added here later + + if len(masks) > 0: + final_attn_mask = ~or_reduce(masks) + + n, device = q.shape[-2], q.device + + causal = self.causal if causal is None else causal + + if n == 1 and causal: + causal = False + + if self.natten_kernel_size is not None: + if natten is None: + raise ImportError( + "natten not installed, please install natten to use neighborhood attention" + ) + + dtype_in = q.dtype + q, k, v = map(lambda t: t.to(torch.float32), (q, k, v)) + + attn = natten.functional.natten1dqk( + q, k, kernel_size=self.natten_kernel_size, dilation=1 + ) + + if final_attn_mask is not None: + attn = attn.masked_fill(final_attn_mask, -torch.finfo(attn.dtype).max) + + attn = F.softmax(attn, dim=-1, dtype=torch.float32) + + out = natten.functional.natten1dav( + attn, v, kernel_size=self.natten_kernel_size, dilation=1 + ).to(dtype_in) + + # Prioritize Flash Attention 2 + elif self.use_fa_flash: + assert final_attn_mask is None, ( + "masking not yet supported for Flash Attention 2" + ) + # Flash Attention 2 requires FP16 inputs + fa_dtype_in = q.dtype + q, k, v = map( + lambda t: rearrange(t, "b h n d -> b n h d").to(torch.float16), + (q, k, v), + ) + + out = flash_attn_func(q, k, v, causal=causal) + + out = rearrange(out.to(fa_dtype_in), "b n h d -> b h n d") + + # Fall back to PyTorch implementation + elif self.use_pt_flash: + out = self.flash_attn(q, k, v, causal=causal, mask=final_attn_mask) + + else: + # Fall back to custom implementation + + if h != kv_h: + # Repeat interleave kv_heads to match q_heads + heads_per_kv_head = h // kv_h + k, v = map( + lambda t: t.repeat_interleave(heads_per_kv_head, dim=1), (k, v) + ) + + scale = 1.0 / (q.shape[-1] ** 0.5) + + kv_einsum_eq = "b j d" if k.ndim == 3 else "b h j d" + + dots = einsum(f"b h i d, {kv_einsum_eq} -> b h i j", q, k) * scale + + i, j, dtype = *dots.shape[-2:], dots.dtype + + mask_value = -torch.finfo(dots.dtype).max + + if final_attn_mask is not None: + dots = dots.masked_fill(~final_attn_mask, mask_value) + + if causal: + causal_mask = self.create_causal_mask(i, j, device=device) + dots = dots.masked_fill(causal_mask, mask_value) + + attn = F.softmax(dots, dim=-1, dtype=torch.float32) + attn = attn.type(dtype) + + out = einsum(f"b h i j, {kv_einsum_eq} -> b h i d", attn, v) + + # merge heads + out = rearrange(out, " b h n d -> b n (h d)") + + # Communicate between heads + + # with autocast(enabled = False): + # out_dtype = out.dtype + # out = out.to(torch.float32) + # out = self.to_out(out).to(out_dtype) + out = self.to_out(out) + + if mask is not None: + mask = rearrange(mask, "b n -> b n 1") + out = out.masked_fill(~mask, 0.0) + + return out + + +class ConformerModule(nn.Module): + def __init__( + self, + dim, + norm_kwargs={}, + ): + super().__init__() + + self.dim = dim + + self.in_norm = LayerNorm(dim, **norm_kwargs) + self.pointwise_conv = nn.Conv1d(dim, dim, kernel_size=1, bias=False) + self.glu = GLU(dim, dim, nn.SiLU()) + self.depthwise_conv = nn.Conv1d( + dim, dim, kernel_size=17, groups=dim, padding=8, bias=False + ) + self.mid_norm = LayerNorm( + dim, **norm_kwargs + ) # This is a batch norm in the original but I don't like batch norm + self.swish = nn.SiLU() + self.pointwise_conv_2 = nn.Conv1d(dim, dim, kernel_size=1, bias=False) + + def forward(self, x): + x = self.in_norm(x) + x = rearrange(x, "b n d -> b d n") + x = self.pointwise_conv(x) + x = rearrange(x, "b d n -> b n d") + x = self.glu(x) + x = rearrange(x, "b n d -> b d n") + x = self.depthwise_conv(x) + x = rearrange(x, "b d n -> b n d") + x = self.mid_norm(x) + x = self.swish(x) + x = rearrange(x, "b n d -> b d n") + x = self.pointwise_conv_2(x) + x = rearrange(x, "b d n -> b n d") + + return x + + +class TransformerBlock(nn.Module): + def __init__( + self, + dim, + dim_heads=64, + cross_attend=False, + dim_context=None, + global_cond_dim=None, + causal=False, + zero_init_branch_outputs=True, + conformer=False, + layer_ix=-1, + remove_norms=False, + attn_kwargs={}, + ff_kwargs={}, + norm_kwargs={}, + ): + super().__init__() + self.dim = dim + self.dim_heads = dim_heads + self.cross_attend = cross_attend + self.dim_context = dim_context + self.causal = causal + + self.pre_norm = ( + LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity() + ) + + self.self_attn = Attention( + dim, + dim_heads=dim_heads, + causal=causal, + zero_init_output=zero_init_branch_outputs, + **attn_kwargs, + ) + + if cross_attend: + self.cross_attend_norm = ( + LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity() + ) + self.cross_attn = Attention( + dim, + dim_heads=dim_heads, + dim_context=dim_context, + causal=causal, + zero_init_output=zero_init_branch_outputs, + **attn_kwargs, + ) + + self.ff_norm = ( + LayerNorm(dim, **norm_kwargs) if not remove_norms else nn.Identity() + ) + self.ff = FeedForward( + dim, zero_init_output=zero_init_branch_outputs, **ff_kwargs + ) + + self.layer_ix = layer_ix + + self.conformer = ( + ConformerModule(dim, norm_kwargs=norm_kwargs) if conformer else None + ) + + self.global_cond_dim = global_cond_dim + + if global_cond_dim is not None: + self.to_scale_shift_gate = nn.Sequential( + nn.SiLU(), nn.Linear(global_cond_dim, dim * 6, bias=False) + ) + + nn.init.zeros_(self.to_scale_shift_gate[1].weight) + # nn.init.zeros_(self.to_scale_shift_gate_self[1].bias) + + def forward( + self, + x, + context=None, + global_cond=None, + mask=None, + context_mask=None, + rotary_pos_emb=None, + ): + if ( + self.global_cond_dim is not None + and self.global_cond_dim > 0 + and global_cond is not None + ): + scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = ( + self.to_scale_shift_gate(global_cond).unsqueeze(1).chunk(6, dim=-1) + ) + + # self-attention with adaLN + residual = x + x = self.pre_norm(x) + x = x * (1 + scale_self) + shift_self + x = self.self_attn(x, mask=mask, rotary_pos_emb=rotary_pos_emb) + x = x * torch.sigmoid(1 - gate_self) + x = x + residual + + if context is not None: + x = x + self.cross_attn( + self.cross_attend_norm(x), + context=context, + context_mask=context_mask, + ) + + if self.conformer is not None: + x = x + self.conformer(x) + + # feedforward with adaLN + residual = x + x = self.ff_norm(x) + x = x * (1 + scale_ff) + shift_ff + x = self.ff(x) + x = x * torch.sigmoid(1 - gate_ff) + x = x + residual + + else: + x = x + self.self_attn( + self.pre_norm(x), mask=mask, rotary_pos_emb=rotary_pos_emb + ) + + if context is not None: + x = x + self.cross_attn( + self.cross_attend_norm(x), + context=context, + context_mask=context_mask, + ) + + if self.conformer is not None: + x = x + self.conformer(x) + + x = x + self.ff(self.ff_norm(x)) + + return x + + +class ContinuousTransformer(nn.Module): + def __init__( + self, + dim, + depth, + *, + dim_in=None, + dim_out=None, + dim_heads=64, + cross_attend=False, + cond_token_dim=None, + global_cond_dim=None, + causal=False, + rotary_pos_emb=True, + zero_init_branch_outputs=True, + conformer=False, + use_sinusoidal_emb=False, + use_abs_pos_emb=False, + abs_pos_emb_max_length=10000, + **kwargs, + ): + super().__init__() + + self.dim = dim + self.depth = depth + self.causal = causal + self.layers = nn.ModuleList([]) + + self.project_in = ( + nn.Linear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity() + ) + self.project_out = ( + nn.Linear(dim, dim_out, bias=False) + if dim_out is not None + else nn.Identity() + ) + + if rotary_pos_emb: + self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32)) + else: + self.rotary_pos_emb = None + + self.use_sinusoidal_emb = use_sinusoidal_emb + if use_sinusoidal_emb: + self.pos_emb = ScaledSinusoidalEmbedding(dim) + + self.use_abs_pos_emb = use_abs_pos_emb + if use_abs_pos_emb: + self.pos_emb = AbsolutePositionalEmbedding(dim, abs_pos_emb_max_length) + + for i in range(depth): + self.layers.append( + TransformerBlock( + dim, + dim_heads=dim_heads, + cross_attend=cross_attend, + dim_context=cond_token_dim, + global_cond_dim=global_cond_dim, + causal=causal, + zero_init_branch_outputs=zero_init_branch_outputs, + conformer=conformer, + layer_ix=i, + **kwargs, + ) + ) + + def forward( + self, + x, + mask=None, + prepend_embeds=None, + prepend_mask=None, + global_cond=None, + return_info=False, + **kwargs, + ): + batch, seq, device = *x.shape[:2], x.device + + info = { + "hidden_states": [], + } + + x = self.project_in(x) + + if prepend_embeds is not None: + prepend_length, prepend_dim = prepend_embeds.shape[1:] + + assert prepend_dim == x.shape[-1], ( + "prepend dimension must match sequence dimension" + ) + + x = torch.cat((prepend_embeds, x), dim=-2) + + if prepend_mask is not None or mask is not None: + mask = ( + mask + if mask is not None + else torch.ones((batch, seq), device=device, dtype=torch.bool) + ) + prepend_mask = ( + prepend_mask + if prepend_mask is not None + else torch.ones( + (batch, prepend_length), device=device, dtype=torch.bool + ) + ) + + mask = torch.cat((prepend_mask, mask), dim=-1) + + # Attention layers + + if self.rotary_pos_emb is not None: + rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1]) + else: + rotary_pos_emb = None + + if self.use_sinusoidal_emb or self.use_abs_pos_emb: + x = x + self.pos_emb(x) + + # Iterate over the transformer layers + for layer in self.layers: + # x = layer(x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, **kwargs) + x = checkpoint( + layer, + x, + rotary_pos_emb=rotary_pos_emb, + global_cond=global_cond, + **kwargs, + ) + + if return_info: + info["hidden_states"].append(x) + + x = self.project_out(x) + + if return_info: + return x, info + + return x diff --git a/src/YingMusicSinger/utils/stable_audio_tools/utils.py b/src/YingMusicSinger/utils/stable_audio_tools/utils.py new file mode 100755 index 0000000000000000000000000000000000000000..cd556344a261975f5540db8d18a6465504f1f9ba --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/utils.py @@ -0,0 +1,98 @@ +import torch +from safetensors.torch import load_file +from torch.nn.utils import remove_weight_norm + + +def load_ckpt_state_dict(ckpt_path): + if ckpt_path.endswith(".safetensors"): + state_dict = load_file(ckpt_path) + else: + state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"] + + return state_dict + + +def remove_weight_norm_from_model(model): + for module in model.modules(): + if hasattr(module, "weight"): + print(f"Removing weight norm from {module}") + remove_weight_norm(module) + + return model + + +# Sampling functions copied from https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/utils/utils.py under MIT license +# License can be found in LICENSES/LICENSE_META.txt + + +def multinomial( + input: torch.Tensor, num_samples: int, replacement=False, *, generator=None +): + """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension. + + Args: + input (torch.Tensor): The input tensor containing probabilities. + num_samples (int): Number of samples to draw. + replacement (bool): Whether to draw with replacement or not. + Keywords args: + generator (torch.Generator): A pseudorandom number generator for sampling. + Returns: + torch.Tensor: Last dimension contains num_samples indices + sampled from the multinomial probability distribution + located in the last dimension of tensor input. + """ + + if num_samples == 1: + q = torch.empty_like(input).exponential_(1, generator=generator) + return torch.argmax(input / q, dim=-1, keepdim=True).to(torch.int64) + + input_ = input.reshape(-1, input.shape[-1]) + output_ = torch.multinomial( + input_, num_samples=num_samples, replacement=replacement, generator=generator + ) + output = output_.reshape(*list(input.shape[:-1]), -1) + return output + + +def sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor: + """Sample next token from top K values along the last dimension of the input probs tensor. + + Args: + probs (torch.Tensor): Input probabilities with token candidates on the last dimension. + k (int): The k in “top-k”. + Returns: + torch.Tensor: Sampled tokens. + """ + top_k_value, _ = torch.topk(probs, k, dim=-1) + min_value_top_k = top_k_value[..., [-1]] + probs *= (probs >= min_value_top_k).float() + probs.div_(probs.sum(dim=-1, keepdim=True)) + next_token = multinomial(probs, num_samples=1) + return next_token + + +def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor: + """Sample next token from top P probabilities along the last dimension of the input probs tensor. + + Args: + probs (torch.Tensor): Input probabilities with token candidates on the last dimension. + p (int): The p in “top-p”. + Returns: + torch.Tensor: Sampled tokens. + """ + probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True) + probs_sum = torch.cumsum(probs_sort, dim=-1) + mask = probs_sum - probs_sort > p + probs_sort *= (~mask).float() + probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True)) + next_token = multinomial(probs_sort, num_samples=1) + next_token = torch.gather(probs_idx, -1, next_token) + return next_token + + +def next_power_of_two(n): + return 2 ** (n - 1).bit_length() + + +def next_multiple_of_64(n): + return ((n + 63) // 64) * 64 diff --git a/src/YingMusicSinger/utils/stable_audio_tools/vae_copysyn.py b/src/YingMusicSinger/utils/stable_audio_tools/vae_copysyn.py new file mode 100755 index 0000000000000000000000000000000000000000..d4c2de5736146ec714366a4f26d045a740594cf4 --- /dev/null +++ b/src/YingMusicSinger/utils/stable_audio_tools/vae_copysyn.py @@ -0,0 +1,164 @@ +import json + +import torch +import torchaudio.transforms as T +from torch import nn + +from .autoencoders import create_autoencoder_from_config +from .utils import load_ckpt_state_dict + + +class PadCrop(nn.Module): + def __init__(self, n_samples, randomize=True): + super().__init__() + self.n_samples = n_samples + self.randomize = randomize + + def __call__(self, signal): + n, s = signal.shape + start = ( + 0 + if (not self.randomize) + else torch.randint(0, max(0, s - self.n_samples) + 1, []).item() + ) + end = start + self.n_samples + output = signal.new_zeros([n, self.n_samples]) + output[:, : min(s, self.n_samples)] = signal[:, start:end] + return output + + +def set_audio_channels(audio, target_channels): + if target_channels == 1: + audio = audio.mean(1, keepdim=True) + elif target_channels == 2: + if audio.shape[1] == 1: + audio = audio.repeat(1, 2, 1) + elif audio.shape[1] > 2: + audio = audio[:, :2, :] + return audio + + +def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device): + audio = audio.to(device) + + if in_sr != target_sr: + resample_tf = T.Resample(in_sr, target_sr).to(device) + audio = resample_tf(audio) + + assert target_length is None + if target_length is None: + target_length = audio.shape[-1] + + audio = PadCrop(target_length, randomize=False)(audio) + + # Add batch dimension + if audio.dim() == 1: + audio = audio.unsqueeze(0).unsqueeze(0) + elif audio.dim() == 2: + audio = audio.unsqueeze(0) + + audio = set_audio_channels(audio, target_channels) + + return audio + + +class StableAudioInfer(nn.Module): + def __init__(self, model_config_path, model_ckpt_path=None): + super().__init__() + + with open(model_config_path) as f: + self.model_config = json.load(f) + + self.model = create_autoencoder_from_config(self.model_config) + if model_ckpt_path is not None: + self.model.load_state_dict(load_ckpt_state_dict(model_ckpt_path)) + + self.sample_rate = self.model_config["sample_rate"] + self.sample_size = self.model_config["sample_size"] + self.io_channels = self.model.io_channels + self.sample_size = 24576 + + @property + def device(self): + return next(self.parameters()).device + + def normalize_audio(self, y, target_dbfs=0): + """Normalize audio to a specific dBFS level.""" + max_amplitude = torch.max(torch.abs(y)) + target_amplitude = 10.0 ** (target_dbfs / 20.0) + scale_factor = target_amplitude / max_amplitude + return y * scale_factor + + def encode_audio(self, input_audio, in_sr): + """Encode audio waveform into VAE latent representation. + + Args: + input_audio: Input audio tensor. + in_sr: Input sample rate. + + Returns: + Latent tensor from the VAE encoder. + """ + input_audio = prepare_audio( + input_audio, + in_sr=in_sr, + target_sr=self.model.sample_rate, + target_length=None, # Determined after resampling + target_channels=self.io_channels, + device=self.device, + ) + input_audio = self.normalize_audio(input_audio, -6) + + with torch.no_grad(): + # Use chunked encoding for long audio to save memory + if input_audio.shape[-1] > (128 + 10) * self.model.sample_rate: + latent = self.model.encode_audio(input_audio, chunked=True) + else: + latent = self.model.encode_audio(input_audio, chunked=False) + + return latent + + def decode_audio(self, latent): + """Decode VAE latent back to audio waveform. + + Args: + latent: Latent tensor. + + Returns: + Decoded audio tensor. + """ + with torch.no_grad(): + # Use chunked decoding for long latents to save memory + if latent.shape[-1] > 128 + 10: + output = self.model.decode_audio(latent, chunked=True) + else: + output = self.model.decode_audio(latent, chunked=False) + return output + + def forward(self, func_type, x, sr=None): + x = x.to(next(self.parameters()).device) + if func_type == "encode": + assert sr is not None, "sr is required for encoding" + return self.encode_audio(input_audio=x, in_sr=sr) + elif func_type == "decode": + return self.decode_audio(x) + else: + raise ValueError(f"Unknown func_type: {func_type}") + + +if __name__ == "__main__": + import torchaudio + + device = "cuda" + vae_model = StableAudioInfer( + model_config_path="config/stable_audio_2_0_vae_20hz_official.json", + model_ckpt_path="ckpts/stable_audio_2_0_vae_20hz_official.ckpt", + ) + vae_model = vae_model.eval().to(device) + + input_audio, in_sr = torchaudio.load("path/to/input.wav") + latent = vae_model(func_type="encode", x=input_audio, sr=in_sr) + + output_audio = vae_model(func_type="decode", x=latent, sr=None) + output_audio = output_audio.squeeze(0).cpu() + torchaudio.save("output.wav", output_audio, sample_rate=44100) diff --git a/src/third_party/MusicSourceSeparationTraining/.gitignore b/src/third_party/MusicSourceSeparationTraining/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..fc3be620893c26bfb77941c6b8cb8e66228abca8 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/.gitignore @@ -0,0 +1,76 @@ +__pycache__ +.DS_Store +*.py[cod] +*$py.class + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +Lib/site-packages/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +share/man/man1/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Jupyter Notebook +.ipynb_checkpoints +share/jupyter +etc/jupyter + +# IPython +profile_default/ +ipython_config.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +pyvenv.cfg +Scripts/ + +*.code-workspace + +results/ +wandb/ \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/LICENSE b/src/third_party/MusicSourceSeparationTraining/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..9d7186e88bca9975edd65956cd499fa60bd04251 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Roman Solovyev (ZFTurbo) + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/third_party/MusicSourceSeparationTraining/README.md b/src/third_party/MusicSourceSeparationTraining/README.md new file mode 100644 index 0000000000000000000000000000000000000000..1652f0b85533aea5c6d7e5f5ac1238aef772beb1 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/README.md @@ -0,0 +1,126 @@ +# Music Source Separation Universal Training Code + +Repository for training models for music source separation. Repository is based on [kuielab code](https://github.com/kuielab/sdx23/tree/mdx_AB/my_submission/src) for [SDX23 challenge](https://github.com/kuielab/sdx23/tree/mdx_AB/my_submission/src). The main idea of this repository is to create training code, which is easy to modify for experiments. Brought to you by [MVSep.com](https://mvsep.com). + +## Models + +Model can be chosen with `--model_type` arg. + +Available models for training: + +* MDX23C based on [KUIELab TFC TDF v3 architecture](https://github.com/kuielab/sdx23/). Key: `mdx23c`. +* Demucs4HT [[Paper](https://arxiv.org/abs/2211.08553)]. Key: `htdemucs`. +* VitLarge23 based on [Segmentation Models Pytorch](https://github.com/qubvel/segmentation_models.pytorch). Key: `segm_models`. +* TorchSeg based on [TorchSeg module](https://github.com/qubvel/segmentation_models.pytorch). Key: `torchseg`. +* Band Split RoFormer [[Paper](https://arxiv.org/abs/2309.02612), [Repository](https://github.com/lucidrains/BS-RoFormer)] . Key: `bs_roformer`. +* Mel-Band RoFormer [[Paper](https://arxiv.org/abs/2310.01809), [Repository](https://github.com/lucidrains/BS-RoFormer)]. Key: `mel_band_roformer`. +* Swin Upernet [[Paper](https://arxiv.org/abs/2103.14030)] Key: `swin_upernet`. +* BandIt Plus [[Paper](https://arxiv.org/abs/2309.02539), [Repository](https://github.com/karnwatcharasupat/bandit)] Key: `bandit`. +* SCNet [[Paper](https://arxiv.org/abs/2401.13276), [Official Repository](https://github.com/starrytong/SCNet), [Unofficial Repository](https://github.com/amanteur/SCNet-PyTorch)] Key: `scnet`. +* BandIt v2 [[Paper](https://arxiv.org/abs/2407.07275), [Repository](https://github.com/kwatcharasupat/bandit-v2)] Key: `bandit_v2`. +* Apollo [[Paper](https://arxiv.org/html/2409.08514v1), [Repository](https://github.com/JusperLee/Apollo)] Key: `apollo`. +* BSMamba2 [[Paper](https://arxiv.org/abs/2508.14556), [Repository](https://github.com/EuiYeonKim/BSMamba2)] Key: `bs_mamba2`. +* Conformer [[Paper](https://arxiv.org/abs/2005.08100), [Repository](https://github.com/lucidrains/conformer)] Key: `conformer`. +* BS Conformer Key: `bs_conformer` +* SCNet Tran Key: `scnet_tran`. +* SCNet Masked Key: `scnet_masked`. + +1. **Note 1**: For `segm_models` there are many different encoders is possible. [Look here](https://github.com/qubvel/segmentation_models.pytorch#encoders-). +2. **Note 2**: Thanks to [@lucidrains](https://github.com/lucidrains) for recreating the RoFormer models based on papers. +3. **Note 3**: For `torchseg` gives access to more than 800 encoders from `timm` module. It's similar to `segm_models`. + +## How to: Train + +To train model you need to: + +1) Choose model type with option `--model_type`, including: `mdx23c`, `htdemucs`, `segm_models`, `mel_band_roformer`, `bs_roformer`. +2) Choose location of config for model `--config_path` ``. You can find examples of configs in [configs folder](configs/). Prefixes `config_musdb18_` are examples for [MUSDB18 dataset](https://sigsep.github.io/datasets/musdb.html). +3) If you have a check-point from the same model or from another similar model you can use it with option: `--start_check_point` `` +4) Choose path where to store results of training `--results_path` `` + +### Training example + +```bash +python train.py \ + --model_type mel_band_roformer \ + --config_path configs/config_mel_band_roformer_vocals.yaml \ + --start_check_point results/model.ckpt \ + --results_path results/ \ + --data_path 'datasets/dataset1' 'datasets/dataset2' \ + --valid_path datasets/musdb18hq/test \ + --num_workers 4 \ + --device_ids 0 +``` + +All training parameters are [here](https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/utils/settings.py#L20). + +### Training with LoRA + +Look here: [LoRA training](docs/LoRA.md) + +## How to: Inference + +### Inference example + +```bash +python inference.py \ + --model_type mdx23c \ + --config_path configs/config_mdx23c_musdb18.yaml \ + --start_check_point results/last_mdx23c.ckpt \ + --input_folder input/wavs/ \ + --store_dir separation_results/ +``` + +All inference parameters are [here](https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/utils/settings.py#L130). +Convert models to ONNX and TensorRT formats [here](https://github.com/ZFTurbo/MSS_ONNX_TensorRT). + +## Useful notes + +* All batch sizes in config are adjusted to use with single NVIDIA A6000 48GB. If you have less memory please adjust correspodningly in model config `training.batch_size` and `training.gradient_accumulation_steps`. +* It's usually always better to start with old weights even if shapes not fully match. Code supports loading weights for not fully same models (but it must have the same architecture). Training will be much faster. + +## Code description + +* `configs/config_*.yaml` - configuration files for models +* `models/*` - set of available models for training and inference +* `dataset.py` - dataset which creates new samples for training +* `gui-wx.py` - GUI interface for code +* `inference.py` - process folder with music files and separate them +* `train.py` - main training code for single GPU +* `train_ddp.py` - training code for Multi GPU config. Faster than `train.py`. Use it for 2 or more GPUs. +* `utils.py` - common functions used by train/valid +* `valid.py` - validation of model with metrics +* `ensemble.py` - useful script to ensemble results of different models to make results better (see [docs](docs/ensemble.md)). + +## Pre-trained models + +Look here: [List of Pre-trained models](docs/pretrained_models.md) + +If you trained some good models, please, share them. You can post config and model weights [in this issue](https://github.com/ZFTurbo/Music-Source-Separation-Training/issues/1). + +## Dataset types + +Look here: [Dataset types](docs/dataset_types.md) + +## Augmentations + +Look here: [Augmentations](docs/augmentations.md) + +## Graphical user interface + +Look here: [GUI documentation](docs/gui.md) or see tutorial on [Youtube](https://youtu.be/M8JKFeN7HfU) + +## Citation + +* [arxiv paper](https://arxiv.org/abs/2305.07489) + +```text +@misc{solovyev2023benchmarks, + title={Benchmarks and leaderboards for sound demixing tasks}, + author={Roman Solovyev and Alexander Stempkovskiy and Tatiana Habruseva}, + year={2023}, + eprint={2305.07489}, + archivePrefix={arXiv}, + primaryClass={cs.SD} +} +``` diff --git a/src/third_party/MusicSourceSeparationTraining/configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml b/src/third_party/MusicSourceSeparationTraining/configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f7f511025573acaa7a53f75c957b878c0aaa8205 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml @@ -0,0 +1,72 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 6 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 4 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 1.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + augmentation: false # enable augmentations by audiomentations and pedalboard + augmentation_type: null + use_mp3_compress: false # Deprecated + augmentation_mix: false # Mix several stems of the same type with some probability + augmentation_loudness: false # randomly change loudness of each stem + augmentation_loudness_type: 1 # Type 1 or 2 + augmentation_loudness_min: 0 + augmentation_loudness_max: 0 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +inference: + batch_size: 4 + dim_t: 256 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_apollo.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_apollo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..290547b44d0c5d62d3d3315e7c2e444728139f22 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_apollo.yaml @@ -0,0 +1,33 @@ +audio: + chunk_size: 132300 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.0 + +model: + sr: 44100 + win: 20 + feature_dim: 256 + layer: 6 + +training: + instruments: ['restored', 'addition'] + target_instrument: 'restored' + batch_size: 2 + num_steps: 1000 + num_epochs: 1000 + optimizer: 'prodigy' + lr: 1.0 + patience: 2 + reduce_factor: 0.95 + coarse_loss_clip: true + grad_clip: 0 + q: 0.95 + use_amp: true + +augmentations: + enable: false # enable or disable all augmentations (to fast disable if needed) + +inference: + batch_size: 4 + num_overlap: 4 diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_bs_mamba2_vocals.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_bs_mamba2_vocals.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f054e8ca861eae6a370eca28eaf2c409ec35f347 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_bs_mamba2_vocals.yaml @@ -0,0 +1,104 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 801 # don't work (use in model) + hop_length: 441 # don't work (use in model) + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 256 + depth: 6 + stereo: true + time_module_depth: 1 + freq_module_depth: 1 + mask_estimator_depth: 2 + module_type: mamba2 + dim_head: 48 + heads: 8 + ff_dropout: 0.1 + attn_dropout: 0.1 + flash_attn: true + mamba_gmlp: false + num_stems: 1 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: false + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: false + +training: + batch_size: 6 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: ['vocals', 'other'] + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + # optimizer: prodigy + # lr: 1.0 + optimizer: adam + lr: 1.0e-5 + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + - 0.002 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + vocals: + pitch_shift: 0.01 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.01 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.01 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + other: + pitch_shift: 0.01 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.01 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + +inference: + chunk_size: 352800 + batch_size: 4 + dim_t: 1101 + num_overlap: 2 + normalize: false \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_dnr_bandit_bsrnn_multi_mus64.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_dnr_bandit_bsrnn_multi_mus64.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2392ca496e498e57a99d70f1f28f73fe3dd7c432 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_dnr_bandit_bsrnn_multi_mus64.yaml @@ -0,0 +1,78 @@ +name: "MultiMaskMultiSourceBandSplitRNN" +audio: + chunk_size: 264600 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + in_channel: 1 + stems: ['speech', 'music', 'effects'] + band_specs: "musical" + n_bands: 64 + fs: 44100 + require_no_overlap: false + require_no_gap: true + normalize_channel_independently: false + treat_channel_as_feature: true + n_sqm_modules: 8 + emb_dim: 128 + rnn_dim: 256 + bidirectional: true + rnn_type: "GRU" + mlp_dim: 512 + hidden_activation: "Tanh" + hidden_activation_kwargs: null + complex_mask: true + n_fft: 2048 + win_length: 2048 + hop_length: 512 + window_fn: "hann_window" + wkwargs: null + power: null + center: true + normalized: true + pad_mode: "constant" + onesided: true + +training: + batch_size: 4 + gradient_accumulation_steps: 4 + grad_clip: 0 + instruments: + - speech + - music + - effects + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_dnr_bandit_v2_mus64.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_dnr_bandit_v2_mus64.yaml new file mode 100644 index 0000000000000000000000000000000000000000..db74fee27426b6e2204d459070603abcf846e3a6 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_dnr_bandit_v2_mus64.yaml @@ -0,0 +1,77 @@ +cls: Bandit + +audio: + chunk_size: 384000 + num_channels: 2 + sample_rate: 48000 + min_mean_abs: 0.000 + +kwargs: + in_channels: 1 + stems: ['speech', 'music', 'sfx'] + band_type: musical + n_bands: 64 + normalize_channel_independently: false + treat_channel_as_feature: true + n_sqm_modules: 8 + emb_dim: 128 + rnn_dim: 256 + bidirectional: true + rnn_type: "GRU" + mlp_dim: 512 + hidden_activation: "Tanh" + hidden_activation_kwargs: null + complex_mask: true + use_freq_weights: true + n_fft: 2048 + win_length: 2048 + hop_length: 512 + window_fn: "hann_window" + wkwargs: null + power: null + center: true + normalized: true + pad_mode: "reflect" + onesided: true + +training: + batch_size: 4 + gradient_accumulation_steps: 4 + grad_clip: 0 + instruments: + - speech + - music + - sfx + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + +inference: + batch_size: 8 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_drumsep.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_drumsep.yaml new file mode 100644 index 0000000000000000000000000000000000000000..687b5ff0639a476a57a6e24552964759c8ce1ff5 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_drumsep.yaml @@ -0,0 +1,72 @@ +audio: + chunk_size: 1764000 # samplerate * segment + min_mean_abs: 0.000 + hop_length: 1024 + +training: + batch_size: 8 + gradient_accumulation_steps: 1 + grad_clip: 0 + segment: 40 + shift: 1 + samplerate: 44100 + channels: 2 + normalize: true + instruments: ['kick', 'snare', 'cymbals', 'toms'] + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + optimizer: adam + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: false # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + +inference: + num_overlap: 4 + batch_size: 8 + +model: hdemucs + +hdemucs: # see demucs/hdemucs.py for a detailed description + channels: 48 + channels_time: null + growth: 2 + nfft: 4096 + wiener_iters: 0 + end_iters: 0 + wiener_residual: False + cac: True + depth: 6 + rewrite: True + hybrid: True + hybrid_old: False + multi_freqs: [] + multi_freqs_depth: 3 + freq_emb: 0.2 + emb_scale: 10 + emb_smooth: True + kernel_size: 8 + stride: 4 + time_stride: 2 + context: 1 + context_enc: 0 + norm_starts: 4 + norm_groups: 4 + dconv_mode: 1 + dconv_depth: 2 + dconv_comp: 4 + dconv_attn: 4 + dconv_lstm: 4 + dconv_init: 0.001 + rescale: 0.1 diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_htdemucs_6stems.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_htdemucs_6stems.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d06a489ec66794414dedd4c143f6e937b26ce666 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_htdemucs_6stems.yaml @@ -0,0 +1,127 @@ +audio: + chunk_size: 485100 # samplerate * segment + min_mean_abs: 0.001 + hop_length: 1024 + +training: + batch_size: 8 + gradient_accumulation_steps: 1 + grad_clip: 0 + segment: 11 + shift: 1 + samplerate: 44100 + channels: 2 + normalize: true + instruments: ['drums', 'bass', 'other', 'vocals', 'guitar', 'piano'] + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + optimizer: adam + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: [0.2, 0.02] + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + +inference: + num_overlap: 4 + batch_size: 8 + +model: htdemucs + +htdemucs: # see demucs/htdemucs.py for a detailed description + # Channels + channels: 48 + channels_time: + growth: 2 + # STFT + num_subbands: 1 + nfft: 4096 + wiener_iters: 0 + end_iters: 0 + wiener_residual: false + cac: true + # Main structure + depth: 4 + rewrite: true + # Frequency Branch + multi_freqs: [] + multi_freqs_depth: 3 + freq_emb: 0.2 + emb_scale: 10 + emb_smooth: true + # Convolutions + kernel_size: 8 + stride: 4 + time_stride: 2 + context: 1 + context_enc: 0 + # normalization + norm_starts: 4 + norm_groups: 4 + # DConv residual branch + dconv_mode: 3 + dconv_depth: 2 + dconv_comp: 8 + dconv_init: 1e-3 + # Before the Transformer + bottom_channels: 0 + # CrossTransformer + # ------ Common to all + # Regular parameters + t_layers: 5 + t_hidden_scale: 4.0 + t_heads: 8 + t_dropout: 0.0 + t_layer_scale: True + t_gelu: True + # ------------- Positional Embedding + t_emb: sin + t_max_positions: 10000 # for the scaled embedding + t_max_period: 10000.0 + t_weight_pos_embed: 1.0 + t_cape_mean_normalize: True + t_cape_augment: True + t_cape_glob_loc_scale: [5000.0, 1.0, 1.4] + t_sin_random_shift: 0 + # ------------- norm before a transformer encoder + t_norm_in: True + t_norm_in_group: False + # ------------- norm inside the encoder + t_group_norm: False + t_norm_first: True + t_norm_out: True + # ------------- optim + t_weight_decay: 0.0 + t_lr: + # ------------- sparsity + t_sparse_self_attn: False + t_sparse_cross_attn: False + t_mask_type: diag + t_mask_random_seed: 42 + t_sparse_attn_window: 400 + t_global_window: 100 + t_sparsity: 0.95 + t_auto_sparsity: False + # Cross Encoder First (False) + t_cross_first: False + # Weight init + rescale: 0.1 + diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_mamba2.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_mamba2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..89154451ca213849592ff4aa4a2d21e132b13cf2 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_mamba2.yaml @@ -0,0 +1,58 @@ +audio: + chunk_size: 132300 # samplerate * segment + hop_length: 1024 + min_mean_abs: 0.0 + +training: + batch_size: 8 + gradient_accumulation_steps: 1 + grad_clip: 0 + segment: 11 + shift: 1 + samplerate: 44100 + channels: 2 + normalize: true + instruments: ['drums', 'bass', 'other', 'vocals'] + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + optimizer: prodigy + lr: 1.0 + patience: 2 + reduce_factor: 0.95 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + read_metadata_procs: 8 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +model: + sr: 44100 + win: 2048 + stride: 512 + feature_dim: 128 + num_repeat_mask: 8 + num_repeat_map: 4 + num_output: 4 + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: + !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + +inference: + num_overlap: 2 + batch_size: 8 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ff17998d201e7f3d894cdc80671b9ac330023541 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer.yaml @@ -0,0 +1,137 @@ +audio: + chunk_size: 131584 + dim_f: 1024 + dim_t: 256 + hop_length: 512 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + dim: 192 + depth: 6 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 512 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 4 # Probably too big (requires a lot of memory for weights) + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + +training: + batch_size: 10 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer_384_8_2_485100_sage.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer_384_8_2_485100_sage.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9824f0e8ff86660edbf7eb1b1089f1cc17af7235 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer_384_8_2_485100_sage.yaml @@ -0,0 +1,197 @@ +audio: + chunk_size: 485100 + dim_f: 1024 + dim_t: 801 # don't work (use in model) + hop_length: 441 # don't work (use in model) + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 8 + stereo: true + num_stems: 4 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 2 + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + sage_attention: True + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: ['drums', 'bass', 'other', 'vocals'] + patience: 3 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + augmentation: false # enable augmentations by audiomentations and pedalboard + augmentation_type: simple1 + use_mp3_compress: false # Deprecated + augmentation_mix: true # Mix several stems of the same type with some probability + augmentation_loudness: true # randomly change loudness of each stem + augmentation_loudness_type: 1 # Type 1 or 2 + augmentation_loudness_min: 0.5 + augmentation_loudness_max: 1.5 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + # optimizer: prodigy + optimizer: adam + # lr: 1.0 + lr: 1.0e-5 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.1 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 2 + dim_t: 1101 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer_with_lora.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer_with_lora.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6c6dcdbcf7ca2a1c820d6649e3be056c270b8788 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_bs_roformer_with_lora.yaml @@ -0,0 +1,205 @@ +audio: + chunk_size: 485100 + dim_f: 1024 + dim_t: 801 # don't work (use in model) + hop_length: 441 # don't work (use in model) + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +lora: + r: 8 + lora_alpha: 16 # alpha / rank > 1 + lora_dropout: 0.05 + merge_weights: False + fan_in_fan_out: False + enable_lora: [True, False, True] # This for QKV + # enable_lora: [True] # For non-Roformers architectures + +model: + dim: 384 + depth: 8 + stereo: true + num_stems: 4 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 2 + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: ['drums', 'bass', 'other', 'vocals'] + patience: 3 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + augmentation: false # enable augmentations by audiomentations and pedalboard + augmentation_type: simple1 + use_mp3_compress: false # Deprecated + augmentation_mix: true # Mix several stems of the same type with some probability + augmentation_loudness: true # randomly change loudness of each stem + augmentation_loudness_type: 1 # Type 1 or 2 + augmentation_loudness_min: 0.5 + augmentation_loudness_max: 1.5 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + # optimizer: prodigy + optimizer: adam + # lr: 1.0 + lr: 1.0e-5 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.1 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 2 + dim_t: 1101 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_demucs3_mmi.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_demucs3_mmi.yaml new file mode 100644 index 0000000000000000000000000000000000000000..08c25c50f8f747d0e4af7acae68b1e47a01f3d0c --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_demucs3_mmi.yaml @@ -0,0 +1,72 @@ +audio: + chunk_size: 485100 # samplerate * segment + min_mean_abs: 0.000 + hop_length: 1024 + +training: + batch_size: 8 + gradient_accumulation_steps: 1 + grad_clip: 0 + segment: 11 + shift: 1 + samplerate: 44100 + channels: 2 + normalize: true + instruments: ['drums', 'bass', 'other', 'vocals'] + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + optimizer: adam + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: false # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + +inference: + num_overlap: 4 + batch_size: 8 + +model: hdemucs + +hdemucs: # see demucs/hdemucs.py for a detailed description + channels: 48 + channels_time: null + growth: 2 + nfft: 4096 + wiener_iters: 0 + end_iters: 0 + wiener_residual: False + cac: True + depth: 6 + rewrite: True + hybrid: True + hybrid_old: False + multi_freqs: [] + multi_freqs_depth: 3 + freq_emb: 0.2 + emb_scale: 10 + emb_smooth: True + kernel_size: 8 + stride: 4 + time_stride: 2 + context: 1 + context_enc: 0 + norm_starts: 4 + norm_groups: 4 + dconv_mode: 1 + dconv_depth: 2 + dconv_comp: 4 + dconv_attn: 4 + dconv_lstm: 4 + dconv_init: 0.001 + rescale: 0.1 diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_htdemucs.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_htdemucs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ba635367baca0b58a977fa4bb38a1cec99579ca9 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_htdemucs.yaml @@ -0,0 +1,119 @@ +audio: + chunk_size: 485100 # samplerate * segment + min_mean_abs: 0.001 + hop_length: 1024 + +training: + batch_size: 8 + gradient_accumulation_steps: 1 + grad_clip: 0 + segment: 11 + shift: 1 + samplerate: 44100 + channels: 2 + normalize: true + instruments: ['drums', 'bass', 'other', 'vocals'] + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + optimizer: adam + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + +inference: + num_overlap: 4 + batch_size: 8 + +model: htdemucs + +htdemucs: # see demucs/htdemucs.py for a detailed description + # Channels + channels: 48 + channels_time: + growth: 2 + # STFT + num_subbands: 1 + nfft: 4096 + wiener_iters: 0 + end_iters: 0 + wiener_residual: false + cac: true + # Main structure + depth: 4 + rewrite: true + # Frequency Branch + multi_freqs: [] + multi_freqs_depth: 3 + freq_emb: 0.2 + emb_scale: 10 + emb_smooth: true + # Convolutions + kernel_size: 8 + stride: 4 + time_stride: 2 + context: 1 + context_enc: 0 + # normalization + norm_starts: 4 + norm_groups: 4 + # DConv residual branch + dconv_mode: 3 + dconv_depth: 2 + dconv_comp: 8 + dconv_init: 1e-3 + # Before the Transformer + bottom_channels: 512 + # CrossTransformer + # ------ Common to all + # Regular parameters + t_layers: 5 + t_hidden_scale: 4.0 + t_heads: 8 + t_dropout: 0.0 + t_layer_scale: True + t_gelu: True + # ------------- Positional Embedding + t_emb: sin + t_max_positions: 10000 # for the scaled embedding + t_max_period: 10000.0 + t_weight_pos_embed: 1.0 + t_cape_mean_normalize: True + t_cape_augment: True + t_cape_glob_loc_scale: [5000.0, 1.0, 1.4] + t_sin_random_shift: 0 + # ------------- norm before a transformer encoder + t_norm_in: True + t_norm_in_group: False + # ------------- norm inside the encoder + t_group_norm: False + t_norm_first: True + t_norm_out: True + # ------------- optim + t_weight_decay: 0.0 + t_lr: + # ------------- sparsity + t_sparse_self_attn: False + t_sparse_cross_attn: False + t_mask_type: diag + t_mask_random_seed: 42 + t_sparse_attn_window: 400 + t_global_window: 100 + t_sparsity: 0.95 + t_auto_sparsity: False + # Cross Encoder First (False) + t_cross_first: False + # Weight init + rescale: 0.1 + diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mdx23c.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mdx23c.yaml new file mode 100644 index 0000000000000000000000000000000000000000..73631f7293c8db94c55c7e1db9fdfc79c712d6e0 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mdx23c.yaml @@ -0,0 +1,182 @@ +audio: + chunk_size: 261120 + dim_f: 4096 + dim_t: 256 + hop_length: 1024 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + act: gelu + bottleneck_factor: 4 + growth: 128 + norm: InstanceNorm + num_blocks_per_scale: 2 + num_channels: 128 + num_scales: 5 + num_subbands: 4 + scale: + - 2 + - 2 + +training: + batch_size: 6 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + # pedalboard reverb block + pedalboard_reverb: 0.01 + pedalboard_reverb_room_size_min: 0.1 + pedalboard_reverb_room_size_max: 0.9 + pedalboard_reverb_damping_min: 0.1 + pedalboard_reverb_damping_max: 0.9 + pedalboard_reverb_wet_level_min: 0.1 + pedalboard_reverb_wet_level_max: 0.9 + pedalboard_reverb_dry_level_min: 0.1 + pedalboard_reverb_dry_level_max: 0.9 + pedalboard_reverb_width_min: 0.9 + pedalboard_reverb_width_max: 1.0 + + # pedalboard chorus block + pedalboard_chorus: 0.01 + pedalboard_chorus_rate_hz_min: 1.0 + pedalboard_chorus_rate_hz_max: 7.0 + pedalboard_chorus_depth_min: 0.25 + pedalboard_chorus_depth_max: 0.95 + pedalboard_chorus_centre_delay_ms_min: 3 + pedalboard_chorus_centre_delay_ms_max: 10 + pedalboard_chorus_feedback_min: 0.0 + pedalboard_chorus_feedback_max: 0.5 + pedalboard_chorus_mix_min: 0.1 + pedalboard_chorus_mix_max: 0.9 + + # pedalboard phazer block + pedalboard_phazer: 0.01 + pedalboard_phazer_rate_hz_min: 1.0 + pedalboard_phazer_rate_hz_max: 10.0 + pedalboard_phazer_depth_min: 0.25 + pedalboard_phazer_depth_max: 0.95 + pedalboard_phazer_centre_frequency_hz_min: 200 + pedalboard_phazer_centre_frequency_hz_max: 12000 + pedalboard_phazer_feedback_min: 0.0 + pedalboard_phazer_feedback_max: 0.5 + pedalboard_phazer_mix_min: 0.1 + pedalboard_phazer_mix_max: 0.9 + + # pedalboard distortion block + pedalboard_distortion: 0.01 + pedalboard_distortion_drive_db_min: 1.0 + pedalboard_distortion_drive_db_max: 25.0 + + # pedalboard pitch shift block + pedalboard_pitch_shift: 0.01 + pedalboard_pitch_shift_semitones_min: -7 + pedalboard_pitch_shift_semitones_max: 7 + + # pedalboard resample block + pedalboard_resample: 0.01 + pedalboard_resample_target_sample_rate_min: 4000 + pedalboard_resample_target_sample_rate_max: 44100 + + # pedalboard bitcrash block + pedalboard_bitcrash: 0.01 + pedalboard_bitcrash_bit_depth_min: 4 + pedalboard_bitcrash_bit_depth_max: 16 + + # pedalboard mp3 compressor block + pedalboard_mp3_compressor: 0.01 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_min: 0 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_max: 9.999 + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.2 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.33 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.33 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mdx23c_stht.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mdx23c_stht.yaml new file mode 100644 index 0000000000000000000000000000000000000000..73631f7293c8db94c55c7e1db9fdfc79c712d6e0 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mdx23c_stht.yaml @@ -0,0 +1,182 @@ +audio: + chunk_size: 261120 + dim_f: 4096 + dim_t: 256 + hop_length: 1024 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + act: gelu + bottleneck_factor: 4 + growth: 128 + norm: InstanceNorm + num_blocks_per_scale: 2 + num_channels: 128 + num_scales: 5 + num_subbands: 4 + scale: + - 2 + - 2 + +training: + batch_size: 6 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + # pedalboard reverb block + pedalboard_reverb: 0.01 + pedalboard_reverb_room_size_min: 0.1 + pedalboard_reverb_room_size_max: 0.9 + pedalboard_reverb_damping_min: 0.1 + pedalboard_reverb_damping_max: 0.9 + pedalboard_reverb_wet_level_min: 0.1 + pedalboard_reverb_wet_level_max: 0.9 + pedalboard_reverb_dry_level_min: 0.1 + pedalboard_reverb_dry_level_max: 0.9 + pedalboard_reverb_width_min: 0.9 + pedalboard_reverb_width_max: 1.0 + + # pedalboard chorus block + pedalboard_chorus: 0.01 + pedalboard_chorus_rate_hz_min: 1.0 + pedalboard_chorus_rate_hz_max: 7.0 + pedalboard_chorus_depth_min: 0.25 + pedalboard_chorus_depth_max: 0.95 + pedalboard_chorus_centre_delay_ms_min: 3 + pedalboard_chorus_centre_delay_ms_max: 10 + pedalboard_chorus_feedback_min: 0.0 + pedalboard_chorus_feedback_max: 0.5 + pedalboard_chorus_mix_min: 0.1 + pedalboard_chorus_mix_max: 0.9 + + # pedalboard phazer block + pedalboard_phazer: 0.01 + pedalboard_phazer_rate_hz_min: 1.0 + pedalboard_phazer_rate_hz_max: 10.0 + pedalboard_phazer_depth_min: 0.25 + pedalboard_phazer_depth_max: 0.95 + pedalboard_phazer_centre_frequency_hz_min: 200 + pedalboard_phazer_centre_frequency_hz_max: 12000 + pedalboard_phazer_feedback_min: 0.0 + pedalboard_phazer_feedback_max: 0.5 + pedalboard_phazer_mix_min: 0.1 + pedalboard_phazer_mix_max: 0.9 + + # pedalboard distortion block + pedalboard_distortion: 0.01 + pedalboard_distortion_drive_db_min: 1.0 + pedalboard_distortion_drive_db_max: 25.0 + + # pedalboard pitch shift block + pedalboard_pitch_shift: 0.01 + pedalboard_pitch_shift_semitones_min: -7 + pedalboard_pitch_shift_semitones_max: 7 + + # pedalboard resample block + pedalboard_resample: 0.01 + pedalboard_resample_target_sample_rate_min: 4000 + pedalboard_resample_target_sample_rate_max: 44100 + + # pedalboard bitcrash block + pedalboard_bitcrash: 0.01 + pedalboard_bitcrash_bit_depth_min: 4 + pedalboard_bitcrash_bit_depth_max: 16 + + # pedalboard mp3 compressor block + pedalboard_mp3_compressor: 0.01 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_min: 0 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_max: 9.999 + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.2 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.33 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.33 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_conformer.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_conformer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..290b5b228aface2e29a1575ab5b9549087588eff --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_conformer.yaml @@ -0,0 +1,92 @@ +audio: + chunk_size: 131584 + dim_f: 1024 + dim_t: 256 + hop_length: 512 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + dim: 192 + depth: 8 + stereo: true + num_stems: 1 + + time_conformer_depth: 1 + freq_conformer_depth: 1 + + # band-splitting + num_bands: 60 + + # attention/width + dim_head: 64 + heads: 8 + ff_mult: 4 + + # conformer conv sub-block + conv_expansion_factor: 2 + conv_kernel_size: 31 + + # dropouts + attn_dropout: 0.0 + ff_dropout: 0.0 + conv_dropout: 0.0 + + # STFT + dim_freqs_in: 1025 + sample_rate: 44100 + stft_n_fft: 2048 + stft_hop_length: 512 + stft_win_length: 2048 + stft_normalized: False + + # mask estimator + multi-res STFT loss + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + + use_torch_checkpoint: False + skip_connection: False + match_input_audio_length: False + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false + use_amp: true + +augmentations: + enable: false + loudness: true + loudness_min: 0.5 + loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_conformer_all_stems.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_conformer_all_stems.yaml new file mode 100644 index 0000000000000000000000000000000000000000..52684c5c58552f3ebb3458b1c5247f106fa7fe81 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_conformer_all_stems.yaml @@ -0,0 +1,109 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 6 + stereo: true + num_stems: 4 + + time_conformer_depth: 1 + freq_conformer_depth: 1 + + # band-splitting + num_bands: 60 + + dim_head: 64 + heads: 8 + ff_mult: 4 + + # conformer conv sub-block + conv_expansion_factor: 2 + conv_kernel_size: 31 + + # dropouts + attn_dropout: 0.0 + ff_dropout: 0.0 + conv_dropout: 0.0 + + # STFT + dim_freqs_in: 1025 # = n_fft//2 + 1 + sample_rate: 44100 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + + # mask estimator + multi-res STFT loss + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + + use_torch_checkpoint: False + skip_connection: False + match_input_audio_length: False + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - drums + - bass + - other + - vocals + lr: 1.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + augmentation: false + augmentation_type: null + use_mp3_compress: false + augmentation_mix: false + augmentation_loudness: false + augmentation_loudness_type: 1 + augmentation_loudness_min: 0 + augmentation_loudness_max: 0 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adam + other_fix: false + use_amp: true + +augmentations: + enable: true + loudness: true + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true + mixup_probs: !!python/tuple + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + all: + channel_shuffle: 0.5 + random_inverse: 0.1 + random_polarity: 0.5 + +inference: + batch_size: 4 + dim_t: 256 + num_overlap: 2 diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4f5b9e05544f74b53a61dd8256b29d91704ca4fc --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer.yaml @@ -0,0 +1,76 @@ +audio: + chunk_size: 131584 + dim_f: 1024 + dim_t: 256 + hop_length: 512 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + dim: 192 + depth: 8 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 512 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 4 # Probably too big (requires a lot of memory for weights) + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + +training: + batch_size: 7 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer_all_stems.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer_all_stems.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9f7f323d40ca96c8359089d3258745078a6da2a9 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer_all_stems.yaml @@ -0,0 +1,97 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 6 + stereo: true + num_stems: 4 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 4 # Probably too big (requires a lot of memory for weights) + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - drums + - bass + - other + - vocals + lr: 1.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + augmentation: false # enable augmentations by audiomentations and pedalboard + augmentation_type: null + use_mp3_compress: false # Deprecated + augmentation_mix: false # Mix several stems of the same type with some probability + augmentation_loudness: false # randomly change loudness of each stem + augmentation_loudness_type: 1 # Type 1 or 2 + augmentation_loudness_min: 0 + augmentation_loudness_max: 0 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: + !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + +inference: + batch_size: 4 + dim_t: 256 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer_sage.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer_sage.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0a36bb5f8c95d594822dcb70dbf03e5b0b9e8e97 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_mel_band_roformer_sage.yaml @@ -0,0 +1,194 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 6 + stereo: true + num_stems: 4 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple [4096, 2048, 1024, 512, 256] + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 2 + sage_attention: True + +training: + batch_size: 2 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: ['drums', 'bass', 'other', 'vocals'] + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adam + lr: 1.0e-04 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + # pedalboard reverb block + pedalboard_reverb: 0.01 + pedalboard_reverb_room_size_min: 0.1 + pedalboard_reverb_room_size_max: 0.9 + pedalboard_reverb_damping_min: 0.1 + pedalboard_reverb_damping_max: 0.9 + pedalboard_reverb_wet_level_min: 0.1 + pedalboard_reverb_wet_level_max: 0.9 + pedalboard_reverb_dry_level_min: 0.1 + pedalboard_reverb_dry_level_max: 0.9 + pedalboard_reverb_width_min: 0.9 + pedalboard_reverb_width_max: 1.0 + + # pedalboard chorus block + pedalboard_chorus: 0.01 + pedalboard_chorus_rate_hz_min: 1.0 + pedalboard_chorus_rate_hz_max: 7.0 + pedalboard_chorus_depth_min: 0.25 + pedalboard_chorus_depth_max: 0.95 + pedalboard_chorus_centre_delay_ms_min: 3 + pedalboard_chorus_centre_delay_ms_max: 10 + pedalboard_chorus_feedback_min: 0.0 + pedalboard_chorus_feedback_max: 0.5 + pedalboard_chorus_mix_min: 0.1 + pedalboard_chorus_mix_max: 0.9 + + # pedalboard phazer block + pedalboard_phazer: 0.01 + pedalboard_phazer_rate_hz_min: 1.0 + pedalboard_phazer_rate_hz_max: 10.0 + pedalboard_phazer_depth_min: 0.25 + pedalboard_phazer_depth_max: 0.95 + pedalboard_phazer_centre_frequency_hz_min: 200 + pedalboard_phazer_centre_frequency_hz_max: 12000 + pedalboard_phazer_feedback_min: 0.0 + pedalboard_phazer_feedback_max: 0.5 + pedalboard_phazer_mix_min: 0.1 + pedalboard_phazer_mix_max: 0.9 + + # pedalboard distortion block + pedalboard_distortion: 0.01 + pedalboard_distortion_drive_db_min: 1.0 + pedalboard_distortion_drive_db_max: 25.0 + + # pedalboard pitch shift block + pedalboard_pitch_shift: 0.01 + pedalboard_pitch_shift_semitones_min: -7 + pedalboard_pitch_shift_semitones_max: 7 + + # pedalboard resample block + pedalboard_resample: 0.01 + pedalboard_resample_target_sample_rate_min: 4000 + pedalboard_resample_target_sample_rate_max: 44100 + + # pedalboard bitcrash block + pedalboard_bitcrash: 0.01 + pedalboard_bitcrash_bit_depth_min: 4 + pedalboard_bitcrash_bit_depth_max: 16 + + # pedalboard mp3 compressor block + pedalboard_mp3_compressor: 0.01 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_min: 0 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_max: 9.999 + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.2 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.33 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.33 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 4 + dim_t: 256 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e7dcafdd8d023938f3f680c8e107a18dba6c892b --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet.yaml @@ -0,0 +1,83 @@ +audio: + chunk_size: 485100 # 44100 * 11 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + sources: + - drums + - bass + - other + - vocals + audio_channels: 2 + dims: + - 4 + - 32 + - 64 + - 128 + nfft: 4096 + hop_size: 1024 + win_size: 4096 + normalized: True + band_SR: + - 0.175 + - 0.392 + - 0.433 + band_stride: + - 1 + - 4 + - 16 + band_kernel: + - 3 + - 4 + - 16 + conv_depths: + - 3 + - 2 + - 1 + compress: 4 + conv_kernel: 3 + num_dplayer: 6 + expand: 1 + +training: + batch_size: 10 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - drums + - bass + - other + - vocals + lr: 5.0e-04 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: + !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 8 + dim_t: 256 + num_overlap: 4 + normalize: true diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet_large.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet_large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..939ba190c8bb18ea782326c5c90b1c26f460cd36 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet_large.yaml @@ -0,0 +1,83 @@ +audio: + chunk_size: 485100 # 44100 * 11 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + sources: + - drums + - bass + - other + - vocals + audio_channels: 2 + dims: + - 4 + - 64 + - 128 + - 256 + nfft: 4096 + hop_size: 1024 + win_size: 4096 + normalized: True + band_SR: + - 0.225 + - 0.372 + - 0.403 + band_stride: + - 1 + - 4 + - 16 + band_kernel: + - 3 + - 4 + - 16 + conv_depths: + - 3 + - 2 + - 1 + compress: 4 + conv_kernel: 3 + num_dplayer: 6 + expand: 1 + +training: + batch_size: 6 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - drums + - bass + - other + - vocals + lr: 5.0e-04 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: + !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 8 + dim_t: 256 + num_overlap: 4 + normalize: false diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet_tran.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet_tran.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4d879d166629406935548b0ebd9f8788cc74bbd0 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_scnet_tran.yaml @@ -0,0 +1,154 @@ +audio: + chunk_size: 485100 # 44100 * 11 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + sources: ['drums', 'bass', 'other', 'vocals'] + audio_channels: 2 + dims: [4, 32, 64, 128] + nfft: 4096 + hop_size: 1024 + win_size: 4096 + normalized: True + band_SR: [0.175, 0.392, 0.433] + band_stride: [1, 4, 16] + band_kernel: [3, 4, 16] + conv_depths: [3, 2, 1] + compress: 4 + conv_kernel: 3 + num_dplayer: 6 + expand: 1 + tran_rotary_embedding_dim: 64 + tran_depth: 1 + tran_heads: 8 + tran_dim_head: 64 + tran_attn_dropout: 0.0 + tran_ff_dropout: 0.0 + tran_flash_attn: False + +training: + batch_size: 5 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: ['drums', 'bass', 'other', 'vocals'] + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + lr: 5.0e-05 + # optimizer: prodigy + # lr: 1.0 + normalize: false # perform normalization on input of model (use the same for inference!) + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + - 0.002 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.01 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + bass: + pitch_shift: 0.01 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.01 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.01 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + time_stretch: 0.1 + time_stretch_min_rate: 0.9 + time_stretch_max_rate: 1.1 + drums: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + +inference: + batch_size: 2 + dim_t: 256 + num_overlap: 2 + normalize: false + +loss_multistft: + fft_sizes: + - 1024 + - 2048 + - 4096 + hop_sizes: + - 147 + - 256 + - 512 + win_lengths: + - 1024 + - 2048 + - 4096 + window: "hann_window" + scale: "mel" + n_bins: 128 + sample_rate: 44100 + perceptual_weighting: true + w_sc: 1.0 + w_log_mag: 1.0 + w_lin_mag: 0.0 + w_phs: 0.0 + mag_distance: "L1" \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_segm_models.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_segm_models.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cbec03910a628bd83c6f42f3984f5d9ba732a9fd --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_segm_models.yaml @@ -0,0 +1,92 @@ +audio: + chunk_size: 261632 + dim_f: 4096 + dim_t: 512 + hop_length: 512 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + encoder_name: tu-maxvit_large_tf_512 # look here for possibilities: https://github.com/qubvel/segmentation_models.pytorch#encoders- + decoder_type: unet # unet, fpn + act: gelu + num_channels: 128 + num_subbands: 8 + +training: + batch_size: 7 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 2000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adamw + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 1 + dim_t: 512 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_torchseg.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_torchseg.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8de81fccd55e4946d6180a1382603cad27e2a7c0 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb18_torchseg.yaml @@ -0,0 +1,92 @@ +audio: + chunk_size: 261632 + dim_f: 4096 + dim_t: 512 + hop_length: 512 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + encoder_name: maxvit_tiny_tf_512 # look with torchseg.list_encoders(). Currently 858 available + decoder_type: unet # unet, fpn + act: gelu + num_channels: 128 + num_subbands: 8 + +training: + batch_size: 18 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 2000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adamw + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 1 + dim_t: 512 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_musdb_conformers.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb_conformers.yaml new file mode 100644 index 0000000000000000000000000000000000000000..41f8886eaba5259147886a394f8ac4f22e3701ec --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_musdb_conformers.yaml @@ -0,0 +1,57 @@ +audio: + chunk_size: 485100 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.00 + +model: + in_channels: 2 + sources: 4 + freq_bins: 2049 + embed_dim: 512 + depth: 8 + dim_head: 64 + heads: 8 + ff_mult: 4 + conv_expansion_factor: 2 + conv_kernel_size: 31 + attn_dropout: 0.1 + ff_dropout: 0.1, + conv_dropout: 0.1 + +training: + batch_size: 30 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - bass + - drums + - other + lr: 1.0e-5 + optimizer: adam + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +stft: + n_fft: 4096 + hop_length: 1024 + win_length: 4096 + center: true + +augmentations: + enable: false # enable or disable all augmentations (to fast disable if needed) + + +inference: + batch_size: 2 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_saxophone_conformers.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_saxophone_conformers.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e8af5a26e583586dd9111b88f307fca5ffbf61ce --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_saxophone_conformers.yaml @@ -0,0 +1,54 @@ +audio: + chunk_size: 485100 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + in_channels: 2 + sources: 2 + freq_bins: 2049 + embed_dim: 512 + depth: 8 + dim_head: 64 + heads: 8 + ff_mult: 4 + conv_expansion_factor: 2 + conv_kernel_size: 31 + attn_dropout: 0.1 + ff_dropout: 0.1, + conv_dropout: 0.1 + +training: + batch_size: 10 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - saxophone + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +stft: + n_fft: 4096 + hop_length: 1024 + win_length: 4096 + center: true + +augmentations: + enable: false # enable or disable all augmentations (to fast disable if needed) + + +inference: + batch_size: 1 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bandit_bsrnn_multi_mus64.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bandit_bsrnn_multi_mus64.yaml new file mode 100644 index 0000000000000000000000000000000000000000..432ae32c19e6136806a718ca882afc516f2aa1f4 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bandit_bsrnn_multi_mus64.yaml @@ -0,0 +1,73 @@ +name: "MultiMaskMultiSourceBandSplitRNN" +audio: + chunk_size: 264600 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + in_channel: 1 + stems: ['vocals', 'other'] + band_specs: "musical" + n_bands: 64 + fs: 44100 + require_no_overlap: false + require_no_gap: true + normalize_channel_independently: false + treat_channel_as_feature: true + n_sqm_modules: 8 + emb_dim: 128 + rnn_dim: 256 + bidirectional: true + rnn_type: "GRU" + mlp_dim: 512 + hidden_activation: "Tanh" + hidden_activation_kwargs: null + complex_mask: true + n_fft: 2048 + win_length: 2048 + hop_length: 512 + window_fn: "hann_window" + wkwargs: null + power: null + center: true + normalized: true + pad_mode: "constant" + onesided: true + +training: + batch_size: 4 + gradient_accumulation_steps: 4 + grad_clip: 0 + instruments: + - vocals + - other + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_conformer.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_conformer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e0341df1eeee5f2fac7d2e79a8d92e2c8d751919 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_conformer.yaml @@ -0,0 +1,133 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 256 + depth: 12 + stereo: true + num_stems: 1 + time_conformer_depth: 1 + freq_conformer_depth: 1 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: false + mlp_expansion_factor: 4 + ff_mult: 4 + conv_expansion_factor: 2 + conv_kernel_size: 31 + use_torch_checkpoint: false + skip_connection: false + sage_attention: false + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0.0 + instruments: + - vocals + - other + lr: 0.0001 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adamw + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true + +inference: + batch_size: 1 + dim_t: 801 + num_overlap: 2 + normalize: false \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_mamba2.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_mamba2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..967e01172c1f99306f61053238100f8e18d34963 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_mamba2.yaml @@ -0,0 +1,51 @@ +audio: + chunk_size: 132300 # samplerate * segment + hop_length: 1024 + min_mean_abs: 0.0 + +training: + batch_size: 8 + gradient_accumulation_steps: 1 + grad_clip: 0 + segment: 11 + shift: 1 + samplerate: 44100 + channels: 2 + normalize: true + instruments: ['vocals', 'other'] + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + optimizer: prodigy + lr: 1.0 + patience: 2 + reduce_factor: 0.95 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + read_metadata_procs: 8 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +model: + sr: 44100 + win: 2048 + stride: 512 + feature_dim: 128 + num_repeat_mask: 8 + num_repeat_map: 4 + num_output: 2 + +augmentations: + enable: false # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: [0.2, 0.02] + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + num_overlap: 2 + batch_size: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_roformer.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_roformer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..98fcb290f86ff37f81f7e45b98c2ed1c14c02c2d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_bs_roformer.yaml @@ -0,0 +1,141 @@ +audio: + chunk_size: 131584 + dim_f: 1024 + dim_t: 256 + hop_length: 512 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + dim: 192 + depth: 6 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 512 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 4 # Probably too big (requires a lot of memory for weights) + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + +training: + batch_size: 10 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_htdemucs.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_htdemucs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..796004a5cbd8a841963b5b41616ffd5cf8b247ea --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_htdemucs.yaml @@ -0,0 +1,123 @@ +audio: + chunk_size: 485100 # samplerate * segment + min_mean_abs: 0.001 + hop_length: 1024 + +training: + batch_size: 10 + gradient_accumulation_steps: 1 + grad_clip: 0 + segment: 11 + shift: 1 + samplerate: 44100 + channels: 2 + normalize: true + instruments: ['vocals', 'other'] + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + optimizer: adam + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: [0.2, 0.02] + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + num_overlap: 2 + batch_size: 8 + +model: htdemucs + +htdemucs: # see demucs/htdemucs.py for a detailed description + # Channels + channels: 48 + channels_time: + growth: 2 + # STFT + num_subbands: 1 + nfft: 4096 + wiener_iters: 0 + end_iters: 0 + wiener_residual: false + cac: true + # Main structure + depth: 4 + rewrite: true + # Frequency Branch + multi_freqs: [] + multi_freqs_depth: 3 + freq_emb: 0.2 + emb_scale: 10 + emb_smooth: true + # Convolutions + kernel_size: 8 + stride: 4 + time_stride: 2 + context: 1 + context_enc: 0 + # normalization + norm_starts: 4 + norm_groups: 4 + # DConv residual branch + dconv_mode: 3 + dconv_depth: 2 + dconv_comp: 8 + dconv_init: 1e-3 + # Before the Transformer + bottom_channels: 512 + # CrossTransformer + # ------ Common to all + # Regular parameters + t_layers: 5 + t_hidden_scale: 4.0 + t_heads: 8 + t_dropout: 0.0 + t_layer_scale: True + t_gelu: True + # ------------- Positional Embedding + t_emb: sin + t_max_positions: 10000 # for the scaled embedding + t_max_period: 10000.0 + t_weight_pos_embed: 1.0 + t_cape_mean_normalize: True + t_cape_augment: True + t_cape_glob_loc_scale: [5000.0, 1.0, 1.4] + t_sin_random_shift: 0 + # ------------- norm before a transformer encoder + t_norm_in: True + t_norm_in_group: False + # ------------- norm inside the encoder + t_group_norm: False + t_norm_first: True + t_norm_out: True + # ------------- optim + t_weight_decay: 0.0 + t_lr: + # ------------- sparsity + t_sparse_self_attn: False + t_sparse_cross_attn: False + t_mask_type: diag + t_mask_random_seed: 42 + t_sparse_attn_window: 400 + t_global_window: 100 + t_sparsity: 0.95 + t_auto_sparsity: False + # Cross Encoder First (False) + t_cross_first: False + # Weight init + rescale: 0.1 + diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mdx23c.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mdx23c.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b241ad5da8e4e4cdeca43a3f09ec64961321ce13 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mdx23c.yaml @@ -0,0 +1,96 @@ +audio: + chunk_size: 261120 + dim_f: 4096 + dim_t: 256 + hop_length: 1024 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + act: gelu + bottleneck_factor: 4 + growth: 128 + norm: InstanceNorm + num_blocks_per_scale: 2 + num_channels: 128 + num_scales: 5 + num_subbands: 4 + scale: + - 2 + - 2 + +training: + batch_size: 6 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 9.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + read_metadata_procs: 8 # Number of processes to use during metadata reading for dataset. Can speed up metadata generation + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mel_band_conformer.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mel_band_conformer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1fb524519a2eddc51421de30824f92a14f285d0f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mel_band_conformer.yaml @@ -0,0 +1,71 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 256 + depth: 12 + stereo: true + num_stems: 1 + time_conformer_depth: 1 + freq_conformer_depth: 1 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: true + dim_freqs_in: 1025 + sample_rate: 44100 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 4 + ff_mult: 4 + conv_expansion_factor: 2 + conv_kernel_size: 31 + use_torch_checkpoint: false + skip_connection: false + sage_attention: false + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0.0 + instruments: + - vocals + - other + lr: 0.0001 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adamw + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true + +inference: + batch_size: 1 + dim_t: 801 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mel_band_roformer.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mel_band_roformer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4fa4385aab64f0cb2c6a424bb6c63bd954f780e0 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_mel_band_roformer.yaml @@ -0,0 +1,80 @@ +audio: + chunk_size: 131584 + dim_f: 1024 + dim_t: 256 + hop_length: 512 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + dim: 192 + depth: 8 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 512 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 4 # Probably too big (requires a lot of memory for weights) + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + +training: + batch_size: 7 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 256 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1210f9654a6e09b0965937b31e9dcaeaaf2257a0 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet.yaml @@ -0,0 +1,79 @@ +audio: + chunk_size: 485100 # 44100 * 11 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + sources: + - vocals + - other + audio_channels: 2 + dims: + - 4 + - 32 + - 64 + - 128 + nfft: 4096 + hop_size: 1024 + win_size: 4096 + normalized: True + band_SR: + - 0.175 + - 0.392 + - 0.433 + band_stride: + - 1 + - 4 + - 16 + band_kernel: + - 3 + - 4 + - 16 + conv_depths: + - 3 + - 2 + - 1 + compress: 4 + conv_kernel: 3 + num_dplayer: 6 + expand: 1 + +training: + batch_size: 10 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 5.0e-04 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 10 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: + !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 8 + dim_t: 256 + num_overlap: 4 + normalize: false diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet_large.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet_large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f34450eb84d2d2b072577504edbf2948efb94158 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet_large.yaml @@ -0,0 +1,79 @@ +audio: + chunk_size: 485100 # 44100 * 11 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + sources: + - vocals + - other + audio_channels: 2 + dims: + - 4 + - 64 + - 128 + - 256 + nfft: 4096 + hop_size: 1024 + win_size: 4096 + normalized: True + band_SR: + - 0.225 + - 0.372 + - 0.403 + band_stride: + - 1 + - 4 + - 16 + band_kernel: + - 3 + - 4 + - 16 + conv_depths: + - 3 + - 2 + - 1 + compress: 4 + conv_kernel: 3 + num_dplayer: 6 + expand: 1 + +training: + batch_size: 6 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 1.0e-04 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: false # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: + !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 8 + dim_t: 256 + num_overlap: 4 + normalize: false diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet_unofficial.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet_unofficial.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2d3e604e4992e1f4090da91227c9ecc5e66e9117 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_scnet_unofficial.yaml @@ -0,0 +1,62 @@ +audio: + chunk_size: 264600 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dims: [4, 32, 64, 128] + bandsplit_ratios: [.175, .392, .433] + downsample_strides: [1, 4, 16] + n_conv_modules: [3, 2, 1] + n_rnn_layers: 6 + rnn_hidden_dim: 128 + n_sources: 2 + + n_fft: 4096 + hop_length: 1024 + win_length: 4096 + stft_normalized: false + + use_mamba: false + d_state: 16 + d_conv: 4 + d_expand: 2 + +training: + batch_size: 10 + gradient_accumulation_steps: 2 + grad_clip: 0 + instruments: + - vocals + - other + lr: 5.0e-04 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: + !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 8 + dim_t: 256 + num_overlap: 4 diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_segm_models.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_segm_models.yaml new file mode 100644 index 0000000000000000000000000000000000000000..44711a0658a95289c8d3745a6d78114b937df1fa --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_segm_models.yaml @@ -0,0 +1,78 @@ +audio: + chunk_size: 261632 + dim_f: 4096 + dim_t: 512 + hop_length: 512 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + encoder_name: tu-maxvit_large_tf_512 # look here for possibilities: https://github.com/qubvel/segmentation_models.pytorch#encoders- + decoder_type: unet # unet, fpn + act: gelu + num_channels: 128 + num_subbands: 8 + +loss_multistft: + fft_sizes: + - 1024 + - 2048 + - 4096 + hop_sizes: + - 512 + - 1024 + - 2048 + win_lengths: + - 1024 + - 2048 + - 4096 + window: "hann_window" + scale: "mel" + n_bins: 128 + sample_rate: 44100 + perceptual_weighting: true + w_sc: 1.0 + w_log_mag: 1.0 + w_lin_mag: 0.0 + w_phs: 0.0 + mag_distance: "L1" + + +training: + batch_size: 8 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 2000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adamw + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 512 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_swin_upernet.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_swin_upernet.yaml new file mode 100644 index 0000000000000000000000000000000000000000..28a837346ea70fa02a1df5245c5048a4676b63c7 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_swin_upernet.yaml @@ -0,0 +1,51 @@ +audio: + chunk_size: 261632 + dim_f: 4096 + dim_t: 512 + hop_length: 512 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + act: gelu + num_channels: 16 + num_subbands: 8 + +training: + batch_size: 14 + gradient_accumulation_steps: 4 + grad_clip: 0 + instruments: + - vocals + - other + lr: 3.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adamw + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 1 + dim_t: 512 + num_overlap: 4 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_torchseg.yaml b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_torchseg.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1ebbeae2770d68c4b2dbdcfd6125a5ca387e6d9b --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/config_vocals_torchseg.yaml @@ -0,0 +1,58 @@ +audio: + chunk_size: 261632 + dim_f: 4096 + dim_t: 512 + hop_length: 512 + n_fft: 8192 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + encoder_name: maxvit_tiny_tf_512 # look with torchseg.list_encoders(). Currently 858 available + decoder_type: unet # unet, fpn + act: gelu + num_channels: 128 + num_subbands: 8 + +training: + batch_size: 18 + gradient_accumulation_steps: 1 + grad_clip: 1.0 + instruments: + - vocals + - other + lr: 1.0e-04 + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: radam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: false # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + +inference: + batch_size: 8 + dim_t: 512 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/sage/config_musdb18_bs_roformer_384_8_2_485100_sage.yaml b/src/third_party/MusicSourceSeparationTraining/configs/sage/config_musdb18_bs_roformer_384_8_2_485100_sage.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9824f0e8ff86660edbf7eb1b1089f1cc17af7235 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/sage/config_musdb18_bs_roformer_384_8_2_485100_sage.yaml @@ -0,0 +1,197 @@ +audio: + chunk_size: 485100 + dim_f: 1024 + dim_t: 801 # don't work (use in model) + hop_length: 441 # don't work (use in model) + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 8 + stereo: true + num_stems: 4 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 2 + use_torch_checkpoint: False # it allows to greatly reduce GPU memory consumption during training (not fully tested) + skip_connection: False # Enable skip connection between transformer blocks - can solve problem with gradients and probably faster training + sage_attention: True + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: ['drums', 'bass', 'other', 'vocals'] + patience: 3 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + augmentation: false # enable augmentations by audiomentations and pedalboard + augmentation_type: simple1 + use_mp3_compress: false # Deprecated + augmentation_mix: true # Mix several stems of the same type with some probability + augmentation_loudness: true # randomly change loudness of each stem + augmentation_loudness_type: 1 # Type 1 or 2 + augmentation_loudness_min: 0.5 + augmentation_loudness_max: 1.5 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + # optimizer: prodigy + optimizer: adam + # lr: 1.0 + lr: 1.0e-5 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.1 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.1 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 2 + dim_t: 1101 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/sage/config_musdb18_mel_band_roformer_sage.yaml b/src/third_party/MusicSourceSeparationTraining/configs/sage/config_musdb18_mel_band_roformer_sage.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0a36bb5f8c95d594822dcb70dbf03e5b0b9e8e97 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/sage/config_musdb18_mel_band_roformer_sage.yaml @@ -0,0 +1,194 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 256 + hop_length: 441 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 6 + stereo: true + num_stems: 4 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0 + ff_dropout: 0 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple [4096, 2048, 1024, 512, 256] + multi_stft_hop_size: 147 + multi_stft_normalized: False + mlp_expansion_factor: 2 + sage_attention: True + +training: + batch_size: 2 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: ['drums', 'bass', 'other', 'vocals'] + patience: 2 + reduce_factor: 0.95 + target_instrument: null + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: false + ema_momentum: 0.999 + optimizer: adam + lr: 1.0e-04 + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + # pedalboard reverb block + pedalboard_reverb: 0.01 + pedalboard_reverb_room_size_min: 0.1 + pedalboard_reverb_room_size_max: 0.9 + pedalboard_reverb_damping_min: 0.1 + pedalboard_reverb_damping_max: 0.9 + pedalboard_reverb_wet_level_min: 0.1 + pedalboard_reverb_wet_level_max: 0.9 + pedalboard_reverb_dry_level_min: 0.1 + pedalboard_reverb_dry_level_max: 0.9 + pedalboard_reverb_width_min: 0.9 + pedalboard_reverb_width_max: 1.0 + + # pedalboard chorus block + pedalboard_chorus: 0.01 + pedalboard_chorus_rate_hz_min: 1.0 + pedalboard_chorus_rate_hz_max: 7.0 + pedalboard_chorus_depth_min: 0.25 + pedalboard_chorus_depth_max: 0.95 + pedalboard_chorus_centre_delay_ms_min: 3 + pedalboard_chorus_centre_delay_ms_max: 10 + pedalboard_chorus_feedback_min: 0.0 + pedalboard_chorus_feedback_max: 0.5 + pedalboard_chorus_mix_min: 0.1 + pedalboard_chorus_mix_max: 0.9 + + # pedalboard phazer block + pedalboard_phazer: 0.01 + pedalboard_phazer_rate_hz_min: 1.0 + pedalboard_phazer_rate_hz_max: 10.0 + pedalboard_phazer_depth_min: 0.25 + pedalboard_phazer_depth_max: 0.95 + pedalboard_phazer_centre_frequency_hz_min: 200 + pedalboard_phazer_centre_frequency_hz_max: 12000 + pedalboard_phazer_feedback_min: 0.0 + pedalboard_phazer_feedback_max: 0.5 + pedalboard_phazer_mix_min: 0.1 + pedalboard_phazer_mix_max: 0.9 + + # pedalboard distortion block + pedalboard_distortion: 0.01 + pedalboard_distortion_drive_db_min: 1.0 + pedalboard_distortion_drive_db_max: 25.0 + + # pedalboard pitch shift block + pedalboard_pitch_shift: 0.01 + pedalboard_pitch_shift_semitones_min: -7 + pedalboard_pitch_shift_semitones_max: 7 + + # pedalboard resample block + pedalboard_resample: 0.01 + pedalboard_resample_target_sample_rate_min: 4000 + pedalboard_resample_target_sample_rate_max: 44100 + + # pedalboard bitcrash block + pedalboard_bitcrash: 0.01 + pedalboard_bitcrash_bit_depth_min: 4 + pedalboard_bitcrash_bit_depth_max: 16 + + # pedalboard mp3 compressor block + pedalboard_mp3_compressor: 0.01 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_min: 0 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_max: 9.999 + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.2 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.33 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.33 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 + + +inference: + batch_size: 4 + dim_t: 256 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_bs_roformer_ep_317_sdr_12.9755.yaml b/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_bs_roformer_ep_317_sdr_12.9755.yaml new file mode 100644 index 0000000000000000000000000000000000000000..135a051897dee27285ac46ee350afe1e1ec02011 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_bs_roformer_ep_317_sdr_12.9755.yaml @@ -0,0 +1,126 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 801 # don't work (use in model) + hop_length: 441 # don't work (use in model) + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 512 + depth: 12 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 2 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 1.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: true # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +inference: + batch_size: 4 + dim_t: 801 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_bs_roformer_ep_937_sdr_10.5309.yaml b/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_bs_roformer_ep_937_sdr_10.5309.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d5e9a0b670759dd378af60e09e0a5e3c650cbf7c --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_bs_roformer_ep_937_sdr_10.5309.yaml @@ -0,0 +1,138 @@ +audio: + chunk_size: 131584 + dim_f: 1024 + dim_t: 256 + hop_length: 512 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.001 + +model: + dim: 384 + depth: 12 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: true + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 512 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 4 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 5.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: other + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + +inference: + batch_size: 8 + dim_t: 512 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_mel_band_roformer_ep_3005_sdr_11.4360.yaml b/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_mel_band_roformer_ep_3005_sdr_11.4360.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7cb922c9c06076e382826decc017ca9d760b9623 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/configs/viperx/model_mel_band_roformer_ep_3005_sdr_11.4360.yaml @@ -0,0 +1,65 @@ +audio: + chunk_size: 352800 + dim_f: 1024 + dim_t: 801 # don't work (use in model) + hop_length: 441 # don't work (use in model) + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 12 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + num_bands: 60 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: True + dim_freqs_in: 1025 + sample_rate: 44100 # needed for mel filter bank from librosa + stft_n_fft: 2048 + stft_hop_length: 441 + stft_win_length: 2048 + stft_normalized: False + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 1 + gradient_accumulation_steps: 8 + grad_clip: 0 + instruments: + - vocals + - other + lr: 4.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true + +inference: + batch_size: 4 + dim_t: 801 + num_overlap: 2 \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/docs/LoRA.md b/src/third_party/MusicSourceSeparationTraining/docs/LoRA.md new file mode 100644 index 0000000000000000000000000000000000000000..52c2bc86606323d98c7530acb1a2b2b58ff6cfe4 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/LoRA.md @@ -0,0 +1,114 @@ +## Training with LoRA + +### What is LoRA? + +LoRA (Low-Rank Adaptation) is a technique designed to reduce the computational and memory cost of fine-tuning large-scale neural networks. Instead of fine-tuning all the model parameters, LoRA introduces small trainable low-rank matrices that are injected into the network. This allows significant reductions in the number of trainable parameters, making it more efficient to adapt pre-trained models to new tasks. For more details, you can refer to the original paper. + +### Enabling LoRA in Training + +To include LoRA in your training pipeline, you need to: + +Add the `--train_lora` flag to the training command. + +Add the following configuration for LoRA in your config file: + +Example: +``` +lora: + r: 8 + lora_alpha: 16 # alpha / rank > 1 + lora_dropout: 0.05 + merge_weights: False + fan_in_fan_out: False + enable_lora: [True] +``` + +Configuration Parameters Explained: + +* `r` (Rank): This determines the rank of the low-rank adaptation matrices. A smaller rank reduces memory usage and file size but may limit the model's adaptability to new tasks. Common values are 4, 8, or 16. + +* `lora_alpha`: Scaling factor for the LoRA weights. The ratio lora_alpha / r should generally be greater than 1 to ensure sufficient expressive power. For example, with r=8 and lora_alpha=16, the scaling factor is 2. + +* `lora_dropout`: Dropout rate applied to LoRA layers. It helps regularize the model and prevent overfitting, especially for smaller datasets. Typical values are in the range [0.0, 0.1]. + +* `merge_weights`: Whether to merge the LoRA weights into the original model weights during inference. Set this to True only if you want to save the final model with merged weights for deployment. + +* `fan_in_fan_out`: Defines the weight initialization convention. Leave this as False for most scenarios unless your model uses a specific convention requiring it. + +* `enable_lora`: A list of booleans specifying whether LoRA should be applied to certain layers. + * For example, `[True, False, True]` enables LoRA for the 1st and 3rd layers but not the 2nd. + * The number of output neurons in the layer must be divisible by the length of enable_lora to ensure proper distribution of LoRA parameters across layers. + * For transformer architectures such as GPT models, `enable_lora = [True, False, True]` is typically used to apply LoRA to the Query (Q) and Value (V) projection matrices while skipping the Key (K) projection matrix. This setup allows efficient fine-tuning of the attention mechanism while maintaining computational efficiency. + +### Benefits of Using LoRA + +* File Size Reduction: With LoRA, only the LoRA layer weights are saved, which significantly reduces the size of the saved model. + +* Flexible Fine-Tuning: You can fine-tune the LoRA layers while keeping the base model frozen, preserving the original model's generalization capabilities. + +* Using Pretrained Weights with LoRA + +### To train a model using both pretrained weights and LoRA weights, you need to: + +1. Include the `--lora_checkpoint` parameter in the training command. + +2. Specify the path to the LoRA checkpoint file. + +### Validating and Inferencing with LoRA + +When using a model fine-tuned with LoRA for validation or inference, you must provide the LoRA checkpoint using the `--lora_checkpoint` parameter. + +### Example Commands + +* Training with LoRA + +``` +python train.py --model_type scnet \ + --config_path configs/config_musdb18_scnet_large_starrytong.yaml \ + --start_check_point weights/last_scnet.ckpt \ + --results_path results/ \ + --data_path datasets/moisesdb/train_tracks \ + --valid_path datasets/moisesdb/valid \ + --device_ids 0 \ + --metrics neg_log_wmse l1_freq sdr \ + --metric_for_scheduler neg_log_wmse \ + --train_lora +``` + +* Validating with LoRA +``` +python valid.py --model_type scnet \ + --config_path configs/config_musdb18_scnet_large_starrytong.yaml \ + --start_check_point weights/last_scnet.ckpt \ + --store_dir results_store/ \ + --valid_path datasets/moisesdb/valid \ + --device_ids 0 \ + --metrics neg_log_wmse l1_freq si_sdr sdr aura_stft aura_mrstft bleedless fullness +``` + +* Inference with LoRA +``` +python inference.py --lora_checkpoint weights/lora_last_scnet.ckpt \ + --model_type scnet \ + --config_path configs/config_musdb18_scnet_large_starrytong.yaml \ + --start_check_point weights/last_scnet.ckpt \ + --store_dir inference_results/ \ + --input_folder datasets/moisesdb/mixtures_for_inference \ + --device_ids 0 +``` + +### Train example with BSRoformer and LoRA + +You can use this [config](../configs/config_musdb18_bs_roformer_with_lora.yaml) and this [weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.12/model_bs_roformer_ep_17_sdr_9.6568.ckpt) to finetune BSRoformer on your dataset. + +``` +python train.py --model_type bs_roformer \ + --config_path configs/config_musdb18_bs_roformer_with_lora.yaml \ + --start_check_point weights/model_bs_roformer_ep_17_sdr_9.6568.ckpt \ + --results_path results/ \ + --data_path musdb18hq/train \ + --valid_path musdb18hq/test \ + --device_ids 0 \ + --metrics sdr \ + --train_lora +``` diff --git a/src/third_party/MusicSourceSeparationTraining/docs/augmentations.md b/src/third_party/MusicSourceSeparationTraining/docs/augmentations.md new file mode 100644 index 0000000000000000000000000000000000000000..41d03585111340cc2c05dfc603753e5af1348a7e --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/augmentations.md @@ -0,0 +1,146 @@ +### Augmentations + +Augmentations allows to change stems on the fly increasing the size of dataset by creating new samples from old samples. +Now control for augmentations is done from config file. Below you can find the example of full config, +which includes all available augmentations: + +```config +augmentations: + enable: true # enable or disable all augmentations (to fast disable if needed) + loudness: true # randomly change loudness of each stem on the range (loudness_min; loudness_max) + loudness_min: 0.5 + loudness_max: 1.5 + mixup: true # mix several stems of same type with some probability (only works for dataset types: 1, 2, 3) + mixup_probs: !!python/tuple # 2 additional stems of the same type (1st with prob 0.2, 2nd with prob 0.02) + - 0.2 + - 0.02 + mixup_loudness_min: 0.5 + mixup_loudness_max: 1.5 + + # apply mp3 compression to mixture only (emulate downloading mp3 from internet) + mp3_compression_on_mixture: 0.01 + mp3_compression_on_mixture_bitrate_min: 32 + mp3_compression_on_mixture_bitrate_max: 320 + mp3_compression_on_mixture_backend: "lameenc" + + all: + channel_shuffle: 0.5 # Set 0 or lower to disable + random_inverse: 0.1 # inverse track (better lower probability) + random_polarity: 0.5 # polarity change (multiply waveform to -1) + mp3_compression: 0.01 + mp3_compression_min_bitrate: 32 + mp3_compression_max_bitrate: 320 + mp3_compression_backend: "lameenc" + + # pedalboard reverb block + pedalboard_reverb: 0.01 + pedalboard_reverb_room_size_min: 0.1 + pedalboard_reverb_room_size_max: 0.9 + pedalboard_reverb_damping_min: 0.1 + pedalboard_reverb_damping_max: 0.9 + pedalboard_reverb_wet_level_min: 0.1 + pedalboard_reverb_wet_level_max: 0.9 + pedalboard_reverb_dry_level_min: 0.1 + pedalboard_reverb_dry_level_max: 0.9 + pedalboard_reverb_width_min: 0.9 + pedalboard_reverb_width_max: 1.0 + + # pedalboard chorus block + pedalboard_chorus: 0.01 + pedalboard_chorus_rate_hz_min: 1.0 + pedalboard_chorus_rate_hz_max: 7.0 + pedalboard_chorus_depth_min: 0.25 + pedalboard_chorus_depth_max: 0.95 + pedalboard_chorus_centre_delay_ms_min: 3 + pedalboard_chorus_centre_delay_ms_max: 10 + pedalboard_chorus_feedback_min: 0.0 + pedalboard_chorus_feedback_max: 0.5 + pedalboard_chorus_mix_min: 0.1 + pedalboard_chorus_mix_max: 0.9 + + # pedalboard phazer block + pedalboard_phazer: 0.01 + pedalboard_phazer_rate_hz_min: 1.0 + pedalboard_phazer_rate_hz_max: 10.0 + pedalboard_phazer_depth_min: 0.25 + pedalboard_phazer_depth_max: 0.95 + pedalboard_phazer_centre_frequency_hz_min: 200 + pedalboard_phazer_centre_frequency_hz_max: 12000 + pedalboard_phazer_feedback_min: 0.0 + pedalboard_phazer_feedback_max: 0.5 + pedalboard_phazer_mix_min: 0.1 + pedalboard_phazer_mix_max: 0.9 + + # pedalboard distortion block + pedalboard_distortion: 0.01 + pedalboard_distortion_drive_db_min: 1.0 + pedalboard_distortion_drive_db_max: 25.0 + + # pedalboard pitch shift block + pedalboard_pitch_shift: 0.01 + pedalboard_pitch_shift_semitones_min: -7 + pedalboard_pitch_shift_semitones_max: 7 + + # pedalboard resample block + pedalboard_resample: 0.01 + pedalboard_resample_target_sample_rate_min: 4000 + pedalboard_resample_target_sample_rate_max: 44100 + + # pedalboard bitcrash block + pedalboard_bitcrash: 0.01 + pedalboard_bitcrash_bit_depth_min: 4 + pedalboard_bitcrash_bit_depth_max: 16 + + # pedalboard mp3 compressor block + pedalboard_mp3_compressor: 0.01 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_min: 0 + pedalboard_mp3_compressor_pedalboard_mp3_compressor_max: 9.999 + + vocals: + pitch_shift: 0.1 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.1 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.7 + bass: + pitch_shift: 0.1 + pitch_shift_min_semitones: -2 + pitch_shift_max_semitones: 2 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -3 + seven_band_parametric_eq_max_gain_db: 6 + tanh_distortion: 0.2 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.5 + drums: + pitch_shift: 0.33 + pitch_shift_min_semitones: -5 + pitch_shift_max_semitones: 5 + seven_band_parametric_eq: 0.25 + seven_band_parametric_eq_min_gain_db: -9 + seven_band_parametric_eq_max_gain_db: 9 + tanh_distortion: 0.33 + tanh_distortion_min: 0.1 + tanh_distortion_max: 0.6 + other: + pitch_shift: 0.1 + pitch_shift_min_semitones: -4 + pitch_shift_max_semitones: 4 + gaussian_noise: 0.1 + gaussian_noise_min_amplitude: 0.001 + gaussian_noise_max_amplitude: 0.015 + time_stretch: 0.01 + time_stretch_min_rate: 0.8 + time_stretch_max_rate: 1.25 +``` + +You can copypaste it into your config to use augmentations. +Notes: +* To completely disable all augmentations you can either remove `augmentations` section from config or set `enable` to `false`. +* If you want to disable some augmentation, just set it to zero. +* Augmentations in `all` subsections applied to all stems +* Augmentations in `vocals`, `bass` etc subsections applied only to corresponding stems. You can create such subsections for all stems which are given in `training.instruments`. \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/docs/bs_roformer_info.md b/src/third_party/MusicSourceSeparationTraining/docs/bs_roformer_info.md new file mode 100644 index 0000000000000000000000000000000000000000..ad7bfc9f8f57e54de1be42cdcdb14775811ebe36 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/bs_roformer_info.md @@ -0,0 +1,145 @@ +### Batch sizes for BSRoformer + +You can use table below to choose BS Roformer `batch_size` parameter for training based on your GPUs. Batch size values provided for single GPU. If you have several GPUs you need to multiply value on number of GPUs. + +| chunk_size | dim | depth | batch_size (A6000 48GB) | batch_size (3090/4090 24GB) | batch_size (16GB) | +|:----------:|:---:|:-----:|:-----------------------:|:---------------------------:|:-----------------:| +| 131584 | 128 | 6 | 10 | 5 | 3 | +| 131584 | 256 | 6 | 8 | 4 | 2 | +| 131584 | 384 | 6 | 7 | 3 | 2 | +| 131584 | 512 | 6 | 6 | 3 | 2 | +| 131584 | 256 | 8 | 6 | 3 | 2 | +| 131584 | 256 | 12 | 4 | 2 | 1 | +| 263168 | 128 | 6 | 4 | 2 | 1 | +| 263168 | 256 | 6 | 3 | 1 | 1 | +| 352800 | 128 | 6 | 2 | 1 | - | +| 352800 | 256 | 6 | 2 | 1 | - | +| 352800 | 384 | 12 | 1 | - | - | +| 352800 | 512 | 12 | - | - | - | + + +Parameters obtained with initial config: + +``` +audio: + chunk_size: 131584 + dim_f: 1024 + dim_t: 515 + hop_length: 512 + n_fft: 2048 + num_channels: 2 + sample_rate: 44100 + min_mean_abs: 0.000 + +model: + dim: 384 + depth: 12 + stereo: true + num_stems: 1 + time_transformer_depth: 1 + freq_transformer_depth: 1 + linear_transformer_depth: 0 + freqs_per_bands: !!python/tuple + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 2 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 4 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 12 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 24 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 48 + - 128 + - 129 + dim_head: 64 + heads: 8 + attn_dropout: 0.1 + ff_dropout: 0.1 + flash_attn: false + dim_freqs_in: 1025 + stft_n_fft: 2048 + stft_hop_length: 512 + stft_win_length: 2048 + stft_normalized: false + mask_estimator_depth: 2 + multi_stft_resolution_loss_weight: 1.0 + multi_stft_resolutions_window_sizes: !!python/tuple + - 4096 + - 2048 + - 1024 + - 512 + - 256 + multi_stft_hop_size: 147 + multi_stft_normalized: False + +training: + batch_size: 1 + gradient_accumulation_steps: 1 + grad_clip: 0 + instruments: + - vocals + - other + lr: 3.0e-05 + patience: 2 + reduce_factor: 0.95 + target_instrument: vocals + num_epochs: 1000 + num_steps: 1000 + q: 0.95 + coarse_loss_clip: true + ema_momentum: 0.999 + optimizer: adam + other_fix: false # it's needed for checking on multisong dataset if other is actually instrumental + use_amp: true # enable or disable usage of mixed precision (float16) - usually it must be true +``` diff --git a/src/third_party/MusicSourceSeparationTraining/docs/changes.md b/src/third_party/MusicSourceSeparationTraining/docs/changes.md new file mode 100644 index 0000000000000000000000000000000000000000..9aeba78f94f0b569849ac6560c40f68c82f6206d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/changes.md @@ -0,0 +1,20 @@ +### Changes + +#### v1.0.2 + +* Added multi GPU validation (earlier validation was performed on single GPU) +* `training.batch_size` in config now must be set for single GPU (if you use multiple GPUs it will be automatically multiplied by number of GPUs) + +#### v1.0.3 + +* Added "spawn" fix for multiprocessing +* Function `get_model_from_config` now takes path of config as input. +* On latest version of pytorch some problems with torch.backends.cudnn.benchmark = True - big slow down. Fixed version 2.0.1 in requirements.txt +* `--valid_path` parameter for train.py now can accept several validation folders instead of one. Added warning if validation folder is empty. +* Small fix for AMP usage in Demucs models taken from config +* Support for Demucs3 mmi model was added +* GPU memory consumption was reduced during inference and validation. +* Some changes to repair click problems on the edges of segment. +* Added support to train on FLAC files. Some more error checks added. +* viperx's Roformer weights and configs added +* `--extract_instrumental` argument added to inference.py \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/docs/dataset_types.md b/src/third_party/MusicSourceSeparationTraining/docs/dataset_types.md new file mode 100644 index 0000000000000000000000000000000000000000..5dc954c4dbd2c10ddbf96f3eb9171862c60549d3 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/dataset_types.md @@ -0,0 +1,152 @@ + Dataset types for training + +### Type 1 (MUSDB) + +Different folders. Each folder contains all needed stems in format +`.wav` (or `flac` in latest releases). + +The structure is the same as in MUSDBHQ18. + +Example: +``` +--- Song 1: +------ vocals.wav +------ bass.wav +------ drums.wav +------ other.wav +--- Song 2: +------ vocals.wav +------ bass.wav +------ drums.wav +------ other.wav +--- Song 3: +........... +``` + +### Type 2 (Stems) + +Each folder represents a single stem name. +The folder contains audio files consisting only of that stem. + +Example: +``` +--- vocals: +------ vocals_1.wav +------ vocals_2.wav +------ vocals_3.wav +------ vocals_4.wav +------ ... +--- bass: +------ bass_1.wav +------ bass_2.wav +------ bass_3.wav +------ bass_4.wav +------ ... +........... +``` + +### Type 3 (CSV file) + +You can provide a CSV file (or a list of CSV files) with the following structure: + +``` +instrum,path +vocals,/path/to/dataset/vocals_1.wav +vocals,/path/to/dataset2/vocals_v2.wav +vocals,/path/to/dataset3/vocals_some.wav +... +drums,/path/to/dataset/drums_good.wav +... +``` + +### Type 4 (MUSDB Aligned) + +The same structure as Type 1, but during training all instruments are loaded from the same position of the song. + +### Type 5 (Precomputed Chunks) + +The same structure as Type 1, but all tracks are pre-split into chunks with 50% overlap. + +### Type 6 (MUSDB Aligned + Explicit Mixture) + +An extension of Type 4, designed for scenarios where the **mixture is treated as a separate signal**, rather than always being reconstructed as `sum(stems)`. + +Structure is the same as Type 1 / Type 4, but it is recommended that each song folder contains `mixture.wav`. + +Example: +``` +--- Song 1: +------ vocals.wav +------ bass.wav +------ drums.wav +------ other.wav +------ mixture.wav +``` + +Key properties: +- All stems are loaded aligned from the same song position +- `mixture.wav` is loaded explicitly if present +- If `mixture.wav` is missing, mixture is computed as sum of stems +- Supports precomputed random chunks (same logic as Type 4) + +Typical use cases: +- Teacher–student or distillation training +- Consistency losses +- Training with real mixes not equal to sum of stems + + +### Type 7 (Class-Balanced Aligned Dataset) + +A class-balanced aligned dataset designed to reduce class frequency bias and improve learning of rare instruments. + +Structure is the same as Type 6: +``` +--- Song 1: +------ flute.wav +------ violin.wav +------ mixture.wav +``` + +How it works: +1. A random instrument (class) is selected +2. A random track containing this instrument is chosen +3. An aligned chunk is loaded from that track +4. The dataset returns which stems are actually present + +Class frequency filtering: +- For each instrument, the ratio of tracks where it appears is computed +- Instruments appearing in more than `max_class_presence_ratio` + (default: 0.4) of tracks are excluded +- Prevents dominant classes (e.g. vocals) from overwhelming training + +Returned values: +- Stems tensor +- Mixture tensor (from `mixture.wav` if available, otherwise sum of stems) +- `active_stem_ids`: indices of instruments present in the current sample + +Typical use cases: +- Training on sparse multi-instrument datasets +- Improving performance on rare instruments +- Conditional or multi-head source separation models + +### Dataset for validation + +* The validation dataset must be the same structure as type 1 datasets (regardless of what type of dataset you're using for training), but also each folder must include `mixture.wav` for each song. `mixture.wav` - is the sum of all stems for song. + +Example: +``` +--- Song 1: +------ vocals.wav +------ bass.wav +------ drums.wav +------ other.wav +------ mixture.wav +--- Song 2: +------ vocals.wav +------ bass.wav +------ drums.wav +------ other.wav +------ mixture.wav +--- Song 3: +........... +``` diff --git a/src/third_party/MusicSourceSeparationTraining/docs/ensemble.md b/src/third_party/MusicSourceSeparationTraining/docs/ensemble.md new file mode 100644 index 0000000000000000000000000000000000000000..bac03d5edd222ce71138c766f066e232dfdda8e9 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/ensemble.md @@ -0,0 +1,30 @@ +### Ensemble usage + +Repository contains `ensemble.py` script which can be used to ensemble results of different algorithms. + +Arguments: +* `--files` - Path to all audio-files to ensemble +* `--type` - Method to do ensemble. One of avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft. Default: avg_wave. +* `--weights` - Weights to create ensemble. Number of weights must be equal to number of files +* `--output` - Path to wav file where ensemble result will be stored (Default: res.wav) + +Example: +``` +ensemble.py --files ./results_tracks/vocals1.wav ./results_tracks/vocals2.wav --weights 2 1 --type max_fft --output out.wav +``` + +### Ensemble types: + +* `avg_wave` - ensemble on 1D variant, find average for every sample of waveform independently +* `median_wave` - ensemble on 1D variant, find median value for every sample of waveform independently +* `min_wave` - ensemble on 1D variant, find minimum absolute value for every sample of waveform independently +* `max_wave` - ensemble on 1D variant, find maximum absolute value for every sample of waveform independently +* `avg_fft` - ensemble on spectrogram (Short-time Fourier transform (STFT), 2D variant), find average for every pixel of spectrogram independently. After averaging use inverse STFT to obtain original 1D-waveform back. +* `median_fft` - the same as avg_fft but use median instead of mean (only useful for ensembling of 3 or more sources). +* `min_fft` - the same as avg_fft but use minimum function instead of mean (reduce aggressiveness). +* `max_fft` - the same as avg_fft but use maximum function instead of mean (the most aggressive). + +### Notes +* `min_fft` can be used to do more conservative ensemble - it will reduce influence of more aggressive models. +* It's better to ensemble models which are of equal quality - in this case it will give gain. If one of model is bad - it will reduce overall quality. +* In my experiments `avg_wave` was always better or equal in SDR score comparing with other methods. diff --git a/src/third_party/MusicSourceSeparationTraining/docs/gui.md b/src/third_party/MusicSourceSeparationTraining/docs/gui.md new file mode 100644 index 0000000000000000000000000000000000000000..04252c3b19ae92974c47e04f54c756698d3bbddd --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/gui.md @@ -0,0 +1,31 @@ +## GUI for MSST code + +GUI was prepared by **Bas Curtiz** and is based on [wxpython](https://en.wikipedia.org/wiki/WxPython) module. + +![Window example](https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/gui/wx_msst_screen.png) + +### How to + +How to use GUI with ZFTurbo's Music Source Separation Universal Training Code: + +1. Install Python: https://www.python.org/ftp/python/3.11.6/python-3.11.6-amd64.exe +2. Install Microsoft Visual C++ 2015-2022 (x64): https://aka.ms/vs/17/release/vc_redist.x64.exe +3. Install Microsoft C++ Build Tools: https://visualstudio.microsoft.com/visual-cpp-build-tools/ +Select Desktop development with C++ +4. Install PyTorch: https://pytorch.org/get-started/locally/ +5. Download and unzip Music Source Separation Universal Training Code: +https://github.com/ZFTurbo/Music-Source-Separation-Training +6. Open up CMD inside the folder and enter: `pip install -r requirements.txt` +7. Enter: `python gui-wx.py` +8. Download models - assign the config (.yaml) and checkpoint (.bin, .ckpt, or .th) + +Video guide on [Youtube](https://youtu.be/M8JKFeN7HfU) (~6.5 minutes). + +[![Tutorial screenshot](https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/gui/tutorial_screenshot.jpg)](https://youtu.be/M8JKFeN7HfU) + +Also you can use GUI as EXE-file on Windows: [Link](https://mega.nz/file/xAAzTCzR#2IapG3RJ9Vew3oC8l9H2zrw1vwUtZSqsUdJAjmARmPs). Put it inside the root folder. This way you can make a shortcut to the exe on your desktop to run it with a double-click. + +### Other links + +* You can try [non-official GUI to MSST](https://github.com/SUC-DriverOld/MSST-WebUI). +* The one more version [by SiftedSand](https://github.com/SiftedSand/MusicSepGUI) \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/docs/mel_roformer_experiments.md b/src/third_party/MusicSourceSeparationTraining/docs/mel_roformer_experiments.md new file mode 100644 index 0000000000000000000000000000000000000000..3d68b92d4c8725bcc9ce724922e5802d0a60e36e --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/mel_roformer_experiments.md @@ -0,0 +1,21 @@ +## Mel Roformer models + +All experiments were made using MUSDB18HQ dataset. All metrics were measured using 'test' set. Training was made using 'train' set. + +### Experiments table + +| Average SDR Score | Chunk size | Depth | Dim | mlp expansion factor | Skip connection | Hop size | FFT Size | Dropout | Batch Size | DL Checkpoint | Comment | +|:-----------------:|:-------------:|:-----------------:|:---:|:--------------------:|:-----:|:-----:|:-----:|:-----:|:----------:|:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|:-----------------------------:| +| 5.1235 | 88200 | 2 | 64 | 1 | No | 441 | 2048 | 0/0 | 32 (48 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_53_sdr_5.1235_config_mel_64_2_1_88200_experimental.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_53_sdr_5.1235.ckpt) | | +| 6.4698 | 88200 | 4 | 128 | 1 | No | 441 | 2048 | 0.1/0.1 | 28 (80 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_38_sdr_6.4698.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_38_sdr_6.4698.ckpt) | | +| 6.7022 | 88200 | 4 | 128 | 1 | No | 882 | 4096 | 0/0 | 20 (80 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_166_sdr_6.7022_config_mel_128_4_1_88200_big_fft_4096.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_166_sdr_6.7022.ckpt) | | +| 7.8127 | 88200 | 6 | 256 | 1 | Yes | 441 | 2048 | 0.1/0.1 | 16 (80 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_168_sdr_7.8127_config_mel_256_6_1_88200.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_168_sdr_7.8127.ckpt) | | +| 6.4908 | 176400 | 4 | 128 | 1 | Yes | 441 | 2048 | 0.1/0.1 | 8 (48 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_15_sdr_6.4908.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_15_sdr_6.4908.ckpt) | | +| 6.5224 | 176400 | 4 | 128 | 2 | Yes | 441 | 2048 | 0.1/0.1 | 8 (48 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_9_sdr_6.5254.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_9_sdr_6.5254.ckpt) | | +| 7.0412 | 352800 | 4 | 128 | 1 | No | 882 | 4096 | 0/0 | 5 (80 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_48_sdr_7.0412_config_mel_128_4_1_352800_big_fft_4096.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_experimental_ep_48_sdr_7.0412.ckpt) | | +| 8.2175 | 352800 | 4 | 256 | 1 | No | 441 | 2048 | 0/0 | 5 (80 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_1_sdr_8.2175.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_1_sdr_8.2175.ckpt) | Trained longer on different strategies. Looks like it a bit overfit in the end | +| 1.0557 | 352800 | 4 | 128 | 1 | No | 882 | 2048 | 0/0 | 6 (48 GB) | --- | Looks like big hop size is not great | +| 6.8652 | 485100 | 4 | 128 | 1 | No | 441 | 2048 | 0.1/0.1 | 5 (48 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_7_sdr_6.8652.yaml) / [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_7_sdr_6.8652.ckpt) | | +| 8.9400* | 485100 | 8 | 384 | 4 | Yes | 882 | 4096 | 0/0 | 2 (80 GB) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_5_sdr_8.9443_config_mel_384_8_4_485100_big_fft_4096_skip_connect.yaml) / Weights ([part 1](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_5_sdr_8.9443.zip.001), [part2](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.11/model_mel_band_roformer_ep_5_sdr_8.9443.zip.002)) | Very big file with weights > 3GB. Continue to increase metrics | + +* Note 1: Some models probably undertrained \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/docs/pretrained_models.md b/src/third_party/MusicSourceSeparationTraining/docs/pretrained_models.md new file mode 100644 index 0000000000000000000000000000000000000000..a156ea57c99e21090a999538e846c39bf3f08d6a --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/pretrained_models.md @@ -0,0 +1,73 @@ +## Pre-trained models + +If you trained some good models, please, share them. You can post config and model weights [in this issue](https://github.com/ZFTurbo/Music-Source-Separation-Training/issues/1). + +### Vocal models + +| Model Type | Instruments | Metrics (SDR) | Config | Checkpoint | +|:---------------------------------------------------------------------------------:|:-------------:|:-----------------:|:-----:|:-----:| +| MDX23C | vocals / other | SDR vocals: 10.17 | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.0/config_vocals_mdx23c.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.0/model_vocals_mdx23c_sdr_10.17.ckpt) | +| HTDemucs4 (MVSep finetuned) | vocals / other | SDR vocals: 8.78 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_vocals_htdemucs.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.0/model_vocals_htdemucs_sdr_8.78.ckpt) | +| Segm Models (VitLarge23) | vocals / other | SDR vocals: 9.77 | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.0/config_vocals_segm_models.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.0/model_vocals_segm_models_sdr_9.77.ckpt) | +| Swin Upernet | vocals / other | SDR vocals: 7.57 | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.2/config_vocals_swin_upernet.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.2/model_swin_upernet_ep_56_sdr_10.6703.ckpt) | +| BS Mamba2 (Trained on MUSDB18 only) | vocals / other | SDR vocals: 8.82 | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.19/config_bs_mamba2_vocals.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.19/bs_mamba2_vocals.ckpt) | +| BS Roformer ([viperx](https://github.com/playdasegunda) edition) | vocals / other | SDR vocals: 10.87 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/viperx/model_bs_roformer_ep_317_sdr_12.9755.yaml) | [Weights](https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/model_bs_roformer_ep_317_sdr_12.9755.ckpt) | +| MelBand Roformer ([viperx](https://github.com/playdasegunda) edition) | vocals / other | SDR vocals: 9.67 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/viperx/model_mel_band_roformer_ep_3005_sdr_11.4360.yaml) | [Weights](https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt) | +| MelBand Roformer ([KimberleyJensen](https://github.com/KimberleyJensen/) edition) | vocals / other | SDR vocals: 10.98 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml) | [Weights](https://huggingface.co/KimberleyJSN/melbandroformer/resolve/main/MelBandRoformer.ckpt) | + +**Note**: Metrics measured on [Multisong Dataset](https://mvsep.com/en/quality_checker). + +### Single stem models + +| Model Type | Instruments | Metrics (SDR) | Config | Checkpoint | +|:-------------------------------------------------------------------------------------------------------------:|:-----------:|:----------------:|:-----:|:-------------------------------------------------------------------------------------------------------------------------------------------------------:| +| HTDemucs4 FT Drums | drums | SDR drums: 11.13 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_musdb18_htdemucs.yaml) | [Weights](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/f7e0c4bc-ba3fe64a.th) | +| HTDemucs4 FT Bass | bass | SDR bass: 11.96 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_musdb18_htdemucs.yaml) | [Weights](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/d12395a8-e57c48e6.th) | +| HTDemucs4 FT Other | other | SDR other: 5.85 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_musdb18_htdemucs.yaml) | [Weights](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/92cfc3b6-ef3bcb9c.th) | +| HTDemucs4 FT Vocals (Official repository) | vocals | SDR vocals: 8.38 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_musdb18_htdemucs.yaml) | [Weights](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/04573f0d-f3cf25b2.th) | +| BS Roformer ([viperx](https://github.com/playdasegunda) edition) | other | SDR other: 6.85 | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/viperx/model_bs_roformer_ep_937_sdr_10.5309.yaml) | [Weights](https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/model_bs_roformer_ep_937_sdr_10.5309.ckpt) | +| MelBand Roformer ([aufr33](https://github.com/aufr33) and [viperx](https://github.com/playdasegunda) edition) | crowd | SDR crowd: 5.99 | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.4/model_mel_band_roformer_crowd.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.4/mel_band_roformer_crowd_aufr33_viperx_sdr_8.7144.ckpt) | +| MelBand Roformer ([anvuew](https://github.com/anvuew) edition) | dereverb | --- | [Config](https://huggingface.co/anvuew/dereverb_mel_band_roformer/resolve/main/dereverb_mel_band_roformer_anvuew.yaml) | [Weights](https://huggingface.co/anvuew/dereverb_mel_band_roformer/resolve/main/dereverb_mel_band_roformer_anvuew_sdr_19.1729.ckpt) | +| MelBand Roformer Denoise (by [aufr33](https://github.com/aufr33)) | denoise | --- | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.7/model_mel_band_roformer_denoise.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.7/denoise_mel_band_roformer_aufr33_sdr_27.9959.ckpt) | +| MelBand Roformer Denoise Aggressive (by [aufr33](https://github.com/aufr33)) | denoise | --- | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.7/model_mel_band_roformer_denoise.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.7/denoise_mel_band_roformer_aufr33_aggr_sdr_27.9768.ckpt) | +| Apollo LQ MP3 restoration (by [JusperLee](https://github.com/JusperLee)) | restored | --- | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/configs/config_apollo.yaml) | [Weights](https://huggingface.co/JusperLee/Apollo/resolve/main/pytorch_model.bin) | +| MelBand Roformer Aspiration (by [SUC-DriverOld](https://github.com/SUC-DriverOld)) | aspiration | SDR: 9.85 | [Config](https://huggingface.co/Sucial/Aspiration_Mel_Band_Roformer/blob/main/config_aspiration_mel_band_roformer.yaml) | [Weights](https://huggingface.co/Sucial/Aspiration_Mel_Band_Roformer/blob/main/aspiration_mel_band_roformer_sdr_18.9845.ckpt) | +| MDX23C Phantom Centre extraction (by [wesleyr36](https://github.com/wesleyr36)) | similarity | L1Freq: 72.23 | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.10/config_mdx23c_similarity.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.10/model_mdx23c_ep_271_l1_freq_72.2383.ckpt) | +| MelBand Roformer Vocals DeReverb/DeEcho (by [SUC-DriverOld](https://github.com/SUC-DriverOld)) | dry | SDR: 10.01 | [Config](https://huggingface.co/Sucial/Dereverb-Echo_Mel_Band_Roformer/resolve/main/config_dereverb-echo_mel_band_roformer.yaml) | [Weights](https://huggingface.co/Sucial/Dereverb-Echo_Mel_Band_Roformer/resolve/main/dereverb-echo_mel_band_roformer_sdr_10.0169.ckpt) | + +**Note**: All HTDemucs4 FT models output 4 stems, but quality is best only on target stem (all other stems are dummy). + +### Multi-stem models + +| Model Type | Instruments | Metrics (SDR) | Config | Checkpoint | +|:---------------------------------------------------------------------------------------------------:|:----------------------------------------------:|:---------------------------------------------------------------------------------------------:|:-----:|:-----:| +| BandIt Plus | speech / music / effects | DnR test avg: 11.50 (speech: 15.64, music: 9.18 effects: 9.69) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.3/config_dnr_bandit_bsrnn_multi_mus64.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.3/model_bandit_plus_dnr_sdr_11.47.chpt) | +| HTDemucs4 | bass / drums / vocals / other | Multisong avg: 9.16 (bass: 11.76, drums: 10.88 vocals: 8.24 other: 5.74) | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_musdb18_htdemucs.yaml) | [Weights](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/955717e8-8726e21a.th) | +| HTDemucs4 (6 stems) | bass / drums / vocals / other / piano / guitar | Multisong (bass: 11.22, drums: 10.22 vocals: 8.05 other: --- piano: --- guitar: ---) | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_htdemucs_6stems.yaml) | [Weights](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/5c90dfd2-34c22ccb.th) | +| Demucs3 mmi | bass / drums / vocals / other | Multisong avg: 8.88 (bass: 11.17, drums: 10.70 vocals: 8.22 other: 5.42) | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_musdb18_demucs3_mmi.yaml) | [Weights](https://dl.fbaipublicfiles.com/demucs/hybrid_transformer/75fc33f5-1941ce65.th) | +| DrumSep htdemucs (by [inagoy](https://github.com/inagoy)) | kick / snare / cymbals / toms | DrumSep test (kick: 10.52, snare: 6.05, toms: 4.68, hh-cymbals: 5.03) | [Config](https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/main/configs/config_drumsep.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.5/model_drumsep.th) | +| DrumSep mdx23c (by [aufr33](https://github.com/aufr33) and [jarredou](https://github.com/jarredou)) | kick / snare / toms / hh / ride / crash | DrumSep test (kick: 14.54, snare: 9.79, toms: 10.63, hh: 3.19, cymbals (ride + crash): 6.08) | [Config](https://github.com/jarredou/models/releases/download/aufr33-jarredou_MDX23C_DrumSep_model_v0.1/aufr33-jarredou_DrumSep_model_mdx23c_ep_141_sdr_10.8059.yaml) | [Weights](https://github.com/jarredou/models/releases/download/aufr33-jarredou_MDX23C_DrumSep_model_v0.1/aufr33-jarredou_DrumSep_model_mdx23c_ep_141_sdr_10.8059.ckpt) | +| DrumSep mdx23c (by [jarredou](https://github.com/jarredou)) | kick / snare / toms / hh / cymbals | DrumSep test (kick: 16.66, snare: 11.53, toms: 12.33, hh: 4.04, cymbals (ride + crash): 6.36) | [Config](https://github.com/jarredou/models/releases/download/DrumSep/config_mdx23c.yaml) | [Weights](https://github.com/jarredou/models/releases/download/DrumSep/drumsep_5stems_mdx23c_jarredou.ckpt) | + +### Multi-stem models (MUSDB18HQ) + +* Models in this list were trained only on MUSDB18HQ dataset (100 songs train data). These weights are useful for fine-tuning. +* Instruments: bass / drums / vocals / other + +| Model Type | Metrics (SDR) | Config | Checkpoint | +|:------------------------------------------------------------:|:-------------------------------------------------------------------------------------------------------------------------------------------------------:|:-------------------------------------------------------------------------------------------------------------------------------------------:|:-------------------------------------------------------------------------------------------------------------------------------------------------:| +| MDX23C | MUSDB test avg: 7.15 (bass: 5.77, drums: 7.93 vocals: 9.23 other: 5.68)
Multisong avg: 7.02 (bass: 8.40, drums: 7.73 vocals: 7.36 other: 4.57) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.1/config_musdb18_mdx23c.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.1/model_mdx23c_ep_168_sdr_7.0207.ckpt) | +| SCNet Small (by [starrytong](https://github.com/starrytong)) | MUSDB test avg: 9.03 (bass: 8.89, drums: 10.44 vocals: 9.90 other: 6.89)
Multisong avg: 8.87 (bass: 11.07, drums: 10.79 vocals: 8.27 other: 5.34) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.6/config_musdb18_scnet.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v.1.0.6/scnet_checkpoint_musdb18.ckpt) | +| SCNet Tran Small | MUSDB test avg: 8.92 (bass: 8.07, drums: 10.81 vocals: 9.97 other: 6.84)
Multisong avg: 8.97 (bass: 10.99, drums: 10.87 vocals: 8.42 other: 5.63) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.14/config_musdb18_scnet_tran.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.14/model_scnet_tran_sdr_8.9272.ckpt) | +| SCNet Masked Small | MUSDB test avg: 8.81 (bass: 8.32, drums: 10.48 vocals: 10.05 other: 6.40)
Multisong avg: 8.86 (bass: 11.00, drums: 10.57 vocals: 8.35 other: 5.51) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.16/config_musdb18_scnet_small.yaml ) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.16/model_scnet_masked_ep_156_sdr_8.8149.ckpt) | +| SCNet Large | MUSDB test avg: 9.32 (bass: 8.63, drums: 10.89 vocals: 10.69 other: 7.06)
Multisong avg: 9.19 (bass: 11.15, drums: 11.04 vocals: 8.94 other: 5.62) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.8/config_musdb18_scnet_large.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.8/model_scnet_sdr_9.3244.ckpt) | +| SCNet Large (by [starrytong](https://github.com/starrytong)) | MUSDB test avg: 9.70 (bass: 9.38, drums: 11.15 vocals: 10.94 other: 7.31)
Multisong avg: 9.28 (bass: 11.27, drums: 11.23 vocals: 9.05 other: 5.57) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.9/config_musdb18_scnet_large_starrytong.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.9/SCNet-large_starrytong_fixed.ckpt) | +| SCNet XL | MUSDB test avg: 9.80 (bass: 9.23, drums: 11.51 vocals: 11.05 other: 7.41)
Multisong avg: 9.72 (bass: 11.87, drums: 11.49 vocals: 9.32 other: 6.19) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.13/config_musdb18_scnet_xl.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.13/model_scnet_ep_54_sdr_9.8051.ckpt) | +| SCNet XL IHF | MUSDB test avg: 10.08 (bass: 9.23, drums: 11.81 vocals: 11.42 other: 7.88)
Multisong avg: 9.92 (bass: 11.94, drums: 11.58 vocals: 9.68 other: 6.48) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.15/config_musdb18_scnet_xl_more_wide_v5.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.15/model_scnet_ep_36_sdr_10.0891.ckpt) | +| SCNet Masked XL IHF | MUSDB test avg: 9.82 (bass: 8.91, drums: 11.62 vocals: 11.06 other: 7.70)
Multisong avg: 9.67 (bass: 11.57, drums: 11.38 vocals: 9.43 other: 6.30) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.17/config_musdb18_scnet_xl_ihf.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.17/model_scnet_masked_ep_111_sdr_9.8286.ckpt) | +| BS Roformer | MUSDB test avg: 9.65 (bass: 8.48, drums: 11.61 vocals: 11.08 other: 7.44)
Multisong avg: 9.38 (bass: 11.08, drums: 11.29 vocals: 9.19 other: 5.96) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.12/config_bs_roformer_384_8_2_485100.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.12/model_bs_roformer_ep_17_sdr_9.6568.ckpt) | +| BS Conformer (Medium) | MUSDB test avg: 9.18 (bass: 8.11, drums: 10.96 vocals: 10.63 other: 7.03)
Multisong avg: 8.84 (bass: 10.36, drums: 10.73 vocals: 8.75 other: 5.53) | [Config](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.18/config_musdb18_bs_conformer_infer.yaml) | [Weights](https://github.com/ZFTurbo/Music-Source-Separation-Training/releases/download/v1.0.18/fused_model_bs_conformer_sdr_9.18.ckpt) | + +### MelRoformer models + +[Table of Mel Band Roformers with different paramers](mel_roformer_experiments.md) \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/docs/test.md b/src/third_party/MusicSourceSeparationTraining/docs/test.md new file mode 100644 index 0000000000000000000000000000000000000000..0fb5ea259e079087e923c14d9bd8cd31aa960baa --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/docs/test.md @@ -0,0 +1,150 @@ +`tests` Documentation +======================== + +Overview +-------- + +The `tests.py` script is designed to verify the functionality of a specific configuration, model weights, and dataset before proceeding with training, validation, or inference. Additionally, it allows the specification of other parameters, which can be passed either through the command line or via the `base_args` variable in the script. + +Usage +----- + +To use `tests.py`, provide the desired arguments via the command line using the `--` prefix. It is mandatory to specify the following arguments: + +* `--model_type` + +* `--config_path` + +* `--start_check_point` + +* `--data_path` + +* `--valid_path` + + +For example: + +``` +python tests.py --check_train \ +--config_path config.yaml \ +--model_type scnet \ +--data_path /path/to/data \ +--valid_path /path/to/valid +``` + +Alternatively, you can define default arguments in the `base_args` variable directly in the script. + +Arguments +--------- + +The script accepts the following arguments: + +* `--check_train`: Check training functionality. + +* `--check_valid`: Check validation functionality. + +* `--check_inference`: Check inference functionality. + +* `--device_ids`: Specify device IDs for training or inference. + +* `--model_type`: Specify the type of model to use. + +* `--start_check_point`: Path to the checkpoint to start from. + +* `--config_path`: Path to the configuration file. + +* `--data_path`: Path to the training data. + +* `--valid_path`: Path to the validation data. + +* `--results_path`: Path to save training results. + +* `--store_dir`: Path to store validation or inference results. + +* `--input_folder`: Path to the input folder for inference. + +* `--metrics`: List of metrics to evaluate, provided as space-separated values. + +* `--max_folders`: Maximum number of folders to process. + +* `--dataset_type`: Dataset type. Must be one of: 1, 2, 3, or 4. Default is 1. + +* `--num_workers`: Number of workers for the dataloader. Default is 0. + +* `--pin_memory`: Use pinned memory in the dataloader. + +* `--seed`: Random seed for reproducibility. Default is 0. + +* `--use_multistft_loss`: Use MultiSTFT Loss from the auraloss package. + +* `--use_mse_loss`: Use Mean Squared Error (MSE) loss. + +* `--use_l1_loss`: Use L1 loss. + +* `--wandb_key`: API Key for Weights and Biases (wandb). Default is an empty string. + +* `--pre_valid`: Run validation before training. + +* `--metric_for_scheduler`: Metric to be used for the learning rate scheduler. Choices are `sdr`, `l1_freq`, `si_sdr`, `neg_log_wmse`, `aura_stft`, `aura_mrstft`, `bleedless`, or `fullness`. Default is `sdr`. + +* `--train_lora`: Enable training with LoRA. + +* `--lora_checkpoint`: Path to the initial LoRA weights checkpoint. Default is an empty string. + +* `--extension`: File extension for validation. Default is `wav`. + +* `--use_tta`: Enable test-time augmentation during inference. This triples runtime but improves prediction quality. + +* `--extract_instrumental`: Invert vocals to obtain instrumental output if available. + +* `--disable_detailed_pbar`: Disable the detailed progress bar. + +* `--force_cpu`: Force the use of the CPU, even if CUDA is available. + +* `--flac_file`: Output FLAC files instead of WAV. + +* `--pcm_type`: PCM type for FLAC files. Choices are `PCM_16` or `PCM_24`. Default is `PCM_24`. + +* `--draw_spectro`: Generate spectrograms for the resulting stems. Specify the value in seconds of the track. Requires `--store_dir` to be set. Default is 0. + + +Example +------- + +To check train, validate and inference with a configuration file with a specific dataset and checkpoint we can use: + +``` +python tests/test.py \ +--check_train \ +--check_valid \ +--check_inference \ +--model_type scnet \ +--config_path configs/config_musdb18_scnet_large_starrytong.yaml \ +--start_check_point weights/model_scnet_ep_30_neg_log_wmse_-11.8688.ckpt \ +--data_path datasets/moisesdb/train_tracks \ +--valid_path datasets/moisesdb/valid \ +--use_tta \ +--use_mse_loss +``` + +This command validates the setup by: + +* Specifying `scnet` as the model type. + +* Loading the configuration from `configs/config_musdb18_scnet_large_starrytong.yaml`. + +* Using the dataset located at `datasets/moisesdb/train_tracks` for training. + +* Using `datasets/moisesdb/valid` for validation. + +* Starting from the checkpoint at `weights/model_scnet_ep_30_neg_log_wmse_-11.8688.ckpt`. + +* Enabling test-time augmentation and using MSE loss. + + +Additional Script: `admin_test.py` +---------------------------------- + +The `admin_test.py` script provides a way to verify the functionality of all configurations and models without specifying model weights or datasets. By default, it performs validation and inference. The configurations and corresponding parameters can be modified using the `MODEL_CONFIGS` variable in the script. + +This script is useful for bulk testing and ensuring that multiple configurations are correctly set up. It can help identify potential issues with configurations or models before proceeding to detailed testing with `tests.py` or full-scale training. diff --git a/src/third_party/MusicSourceSeparationTraining/ensemble.py b/src/third_party/MusicSourceSeparationTraining/ensemble.py new file mode 100644 index 0000000000000000000000000000000000000000..a65cef2b22f769f4454ab468b0cb6cd6ccc5040f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/ensemble.py @@ -0,0 +1,184 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" + +import argparse +import os + +import librosa +import numpy as np +import soundfile as sf + + +def stft(wave, nfft, hl): + wave_left = np.asfortranarray(wave[0]) + wave_right = np.asfortranarray(wave[1]) + spec_left = librosa.stft(wave_left, n_fft=nfft, hop_length=hl) + spec_right = librosa.stft(wave_right, n_fft=nfft, hop_length=hl) + spec = np.asfortranarray([spec_left, spec_right]) + return spec + + +def istft(spec, hl, length): + spec_left = np.asfortranarray(spec[0]) + spec_right = np.asfortranarray(spec[1]) + wave_left = librosa.istft(spec_left, hop_length=hl, length=length) + wave_right = librosa.istft(spec_right, hop_length=hl, length=length) + wave = np.asfortranarray([wave_left, wave_right]) + return wave + + +def absmax(a, *, axis): + dims = list(a.shape) + dims.pop(axis) + indices = np.ogrid[tuple(slice(0, d) for d in dims)] + argmax = np.abs(a).argmax(axis=axis) + indices.insert((len(a.shape) + axis) % len(a.shape), argmax) + return a[tuple(indices)] + + +def absmin(a, *, axis): + dims = list(a.shape) + dims.pop(axis) + indices = np.ogrid[tuple(slice(0, d) for d in dims)] + argmax = np.abs(a).argmin(axis=axis) + indices.insert((len(a.shape) + axis) % len(a.shape), argmax) + return a[tuple(indices)] + + +def lambda_max(arr, axis=None, key=None, keepdims=False): + idxs = np.argmax(key(arr), axis) + if axis is not None: + idxs = np.expand_dims(idxs, axis) + result = np.take_along_axis(arr, idxs, axis) + if not keepdims: + result = np.squeeze(result, axis=axis) + return result + else: + return arr.flatten()[idxs] + + +def lambda_min(arr, axis=None, key=None, keepdims=False): + idxs = np.argmin(key(arr), axis) + if axis is not None: + idxs = np.expand_dims(idxs, axis) + result = np.take_along_axis(arr, idxs, axis) + if not keepdims: + result = np.squeeze(result, axis=axis) + return result + else: + return arr.flatten()[idxs] + + +def average_waveforms(pred_track, weights, algorithm): + """ + :param pred_track: shape = (num, channels, length) + :param weights: shape = (num, ) + :param algorithm: One of avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft + :return: averaged waveform in shape (channels, length) + """ + + pred_track = np.array(pred_track) + final_length = pred_track.shape[-1] + + mod_track = [] + for i in range(pred_track.shape[0]): + if algorithm == "avg_wave": + mod_track.append(pred_track[i] * weights[i]) + elif algorithm in ["median_wave", "min_wave", "max_wave"]: + mod_track.append(pred_track[i]) + elif algorithm in ["avg_fft", "min_fft", "max_fft", "median_fft"]: + spec = stft(pred_track[i], nfft=2048, hl=1024) + if algorithm in ["avg_fft"]: + mod_track.append(spec * weights[i]) + else: + mod_track.append(spec) + pred_track = np.array(mod_track) + + if algorithm in ["avg_wave"]: + pred_track = pred_track.sum(axis=0) + pred_track /= np.array(weights).sum().T + elif algorithm in ["median_wave"]: + pred_track = np.median(pred_track, axis=0) + elif algorithm in ["min_wave"]: + pred_track = np.array(pred_track) + pred_track = lambda_min(pred_track, axis=0, key=np.abs) + elif algorithm in ["max_wave"]: + pred_track = np.array(pred_track) + pred_track = lambda_max(pred_track, axis=0, key=np.abs) + elif algorithm in ["avg_fft"]: + pred_track = pred_track.sum(axis=0) + pred_track /= np.array(weights).sum() + pred_track = istft(pred_track, 1024, final_length) + elif algorithm in ["min_fft"]: + pred_track = np.array(pred_track) + pred_track = lambda_min(pred_track, axis=0, key=np.abs) + pred_track = istft(pred_track, 1024, final_length) + elif algorithm in ["max_fft"]: + pred_track = np.array(pred_track) + pred_track = absmax(pred_track, axis=0) + pred_track = istft(pred_track, 1024, final_length) + elif algorithm in ["median_fft"]: + pred_track = np.array(pred_track) + pred_track = np.median(pred_track, axis=0) + pred_track = istft(pred_track, 1024, final_length) + return pred_track + + +def ensemble_files(args): + parser = argparse.ArgumentParser() + parser.add_argument( + "--files", + type=str, + required=True, + nargs="+", + help="Path to all audio-files to ensemble", + ) + parser.add_argument( + "--type", + type=str, + default="avg_wave", + help="One of avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft", + ) + parser.add_argument( + "--weights", + type=float, + nargs="+", + help="Weights to create ensemble. Number of weights must be equal to number of files", + ) + parser.add_argument( + "--output", + default="res.wav", + type=str, + help="Path to wav file where ensemble result will be stored", + ) + if args is None: + args = parser.parse_args() + else: + args = parser.parse_args(args) + + print("Ensemble type: {}".format(args.type)) + print("Number of input files: {}".format(len(args.files))) + if args.weights is not None: + weights = args.weights + else: + weights = np.ones(len(args.files)) + print("Weights: {}".format(weights)) + print("Output file: {}".format(args.output)) + data = [] + for f in args.files: + if not os.path.isfile(f): + print("Error. Can't find file: {}. Check paths.".format(f)) + exit() + print("Reading file: {}".format(f)) + wav, sr = librosa.load(f, sr=None, mono=False) + # wav, sr = sf.read(f) + print("Waveform shape: {} sample rate: {}".format(wav.shape, sr)) + data.append(wav) + data = np.array(data) + res = average_waveforms(data, weights, args.type) + print("Result shape: {}".format(res.shape)) + sf.write(args.output, res.T, sr, "FLOAT") + + +if __name__ == "__main__": + ensemble_files(None) diff --git a/src/third_party/MusicSourceSeparationTraining/gui/Poppins Bold 700.ttf b/src/third_party/MusicSourceSeparationTraining/gui/Poppins Bold 700.ttf new file mode 100644 index 0000000000000000000000000000000000000000..c2b12d278444c1f28a1947cc77756096aea7cf48 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/gui/Poppins Bold 700.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6b566b380759dd19554228564c6fb2dc01fcf6bfe5cdc5ba28f1ce3b360a9fb5 +size 142316 diff --git a/src/third_party/MusicSourceSeparationTraining/gui/Poppins Regular 400.ttf b/src/third_party/MusicSourceSeparationTraining/gui/Poppins Regular 400.ttf new file mode 100644 index 0000000000000000000000000000000000000000..dd471c31051310d1e8c92431581c93941adf89a1 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/gui/Poppins Regular 400.ttf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4bb722fdfadc6ca417a99d2b89ce6743795d1a59a86762962b03dcaf0ea70a2f +size 146204 diff --git a/src/third_party/MusicSourceSeparationTraining/gui/favicon.ico b/src/third_party/MusicSourceSeparationTraining/gui/favicon.ico new file mode 100644 index 0000000000000000000000000000000000000000..a6e14209a01f53ea9b46766360158381ca98c916 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/gui/favicon.ico @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a72b0d3cdeeb1935d0127fd8d4df842a0b65d7c3fd3d832baebcd43e074e9a48 +size 10462 diff --git a/src/third_party/MusicSourceSeparationTraining/gui/gui-wx.py b/src/third_party/MusicSourceSeparationTraining/gui/gui-wx.py new file mode 100644 index 0000000000000000000000000000000000000000..90b6f21fc54a53bb88845fdd67a6095d04e9939f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/gui/gui-wx.py @@ -0,0 +1,862 @@ +import json +import os +import queue +import subprocess +import sys +import threading +import webbrowser + +import wx +import wx.adv +import wx.html +import wx.html2 + + +def run_subprocess(cmd, output_queue): + try: + process = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True + ) + for line in process.stdout: + output_queue.put(line) + process.wait() + if process.returncode == 0: + output_queue.put("Process completed successfully!") + else: + output_queue.put(f"Process failed with return code {process.returncode}") + except Exception as e: + output_queue.put(f"An error occurred: {str(e)}") + + +def update_output(output_text, output_queue): + try: + while True: + line = output_queue.get_nowait() + if wx.Window.FindWindowById(output_text.GetId()): + wx.CallAfter(output_text.AppendText, line) + else: + return # Exit if the text control no longer exists + except queue.Empty: + pass + except RuntimeError: + return # Exit if a RuntimeError occurs (e.g., window closed) + wx.CallLater(100, update_output, output_text, output_queue) + + +def open_store_folder(folder_path): + if os.path.exists(folder_path): + os.startfile(folder_path) + else: + wx.MessageBox( + f"The folder {folder_path} does not exist.", "Error", wx.OK | wx.ICON_ERROR + ) + + +class DarkThemedTextCtrl(wx.TextCtrl): + def __init__(self, parent, id=wx.ID_ANY, value="", style=0): + super().__init__(parent, id, value, style=style | wx.NO_BORDER) + self.SetBackgroundColour(wx.Colour(0, 0, 0)) + self.SetForegroundColour(wx.WHITE) + + +class CollapsiblePanel(wx.Panel): + def __init__(self, parent, title, *args, **kwargs): + wx.Panel.__init__(self, parent, *args, **kwargs) + self.SetBackgroundColour(parent.GetBackgroundColour()) + + self.toggle_button = wx.Button(self, label=title, style=wx.NO_BORDER) + self.toggle_button.SetBackgroundColour(self.GetBackgroundColour()) + self.toggle_button.Bind(wx.EVT_BUTTON, self.on_toggle) + + self.content_panel = wx.Panel(self) + self.content_panel.SetBackgroundColour(self.GetBackgroundColour()) + + self.main_sizer = wx.BoxSizer(wx.VERTICAL) + self.main_sizer.Add(self.toggle_button, 0, wx.EXPAND | wx.ALL, 5) + self.main_sizer.Add(self.content_panel, 0, wx.EXPAND | wx.ALL, 5) + + self.SetSizer(self.main_sizer) + self.collapsed = True + self.toggle_button.SetLabel(f"▶ {title}") + self.content_panel.Hide() + + def on_toggle(self, event): + self.collapsed = not self.collapsed + self.toggle_button.SetLabel( + f"{'▶' if self.collapsed else '▼'} {self.toggle_button.GetLabel()[2:]}" + ) + self.content_panel.Show(not self.collapsed) + self.Layout() + self.GetParent().Layout() + + def get_content_panel(self): + return self.content_panel + + +class CustomToolTip(wx.PopupWindow): + def __init__(self, parent, text): + wx.PopupWindow.__init__(self, parent) + + # Main panel for tooltip + panel = wx.Panel(self) + self.st = wx.StaticText(panel, 1, text, pos=(10, 10)) + + font = wx.Font( + 8, + wx.FONTFAMILY_DEFAULT, + wx.FONTSTYLE_NORMAL, + wx.FONTWEIGHT_NORMAL, + False, + "Poppins", + ) + self.st.SetFont(font) + + size = self.st.GetBestSize() + self.SetSize((size.width + 20, size.height + 20)) + + # Adjust the panel size + panel.SetSize(self.GetSize()) + panel.SetBackgroundColour(wx.Colour(255, 255, 255)) + + # Bind paint event to draw border + panel.Bind(wx.EVT_PAINT, self.on_paint) + + def on_paint(self, event): + # Get the device context for the panel (not self) + panel = event.GetEventObject() # Get the panel triggering the paint event + dc = wx.PaintDC(panel) # Use panel as the target of the PaintDC + dc.SetPen(wx.Pen(wx.Colour(210, 210, 210), 1)) # Border color + dc.SetBrush(wx.Brush(wx.Colour(255, 255, 255))) # Fill with white + + size = panel.GetSize() + dc.DrawRectangle(0, 0, size.width, size.height) # Draw border around panel + + +class MainFrame(wx.Frame): + def __init__(self): + super().__init__( + parent=None, title="Music Source Separation Training & Inference GUI" + ) + self.SetSize(994, 670) + self.SetBackgroundColour(wx.Colour(247, 248, 250)) # #F7F8FA + + icon = wx.Icon("gui/favicon.ico", wx.BITMAP_TYPE_ICO) + self.SetIcon(icon) + + self.saved_combinations = {} + + # Center the window on the screen + self.Center() + + # Set Poppins font for the entire application + font_path = "gui/Poppins Regular 400.ttf" + bold_font_path = "gui/Poppins Bold 700.ttf" + wx.Font.AddPrivateFont(font_path) + wx.Font.AddPrivateFont(bold_font_path) + self.font = wx.Font( + 9, + wx.FONTFAMILY_DEFAULT, + wx.FONTSTYLE_NORMAL, + wx.FONTWEIGHT_NORMAL, + False, + "Poppins", + ) + self.bold_font = wx.Font( + 10, + wx.FONTFAMILY_DEFAULT, + wx.FONTSTYLE_NORMAL, + wx.FONTWEIGHT_BOLD, + False, + "Poppins", + ) + self.SetFont(self.font) + + panel = wx.Panel(self) + main_sizer = wx.BoxSizer(wx.VERTICAL) + + # Add image (with error handling) + try: + img = wx.Image("gui/mvsep.png", wx.BITMAP_TYPE_PNG) + img_bitmap = wx.Bitmap(img) + img_ctrl = wx.StaticBitmap(panel, -1, img_bitmap) + main_sizer.Add(img_ctrl, 0, wx.ALIGN_CENTER | wx.TOP, 20) + except: + print("Failed to load image: gui/mvsep.png") + + # Add title text + title_text = wx.StaticText( + panel, label="Music Source Separation Training && Inference GUI" + ) + title_text.SetFont(self.bold_font) + title_text.SetForegroundColour(wx.BLACK) + main_sizer.Add(title_text, 0, wx.ALIGN_CENTER | wx.TOP, 10) + + # Add subtitle text + subtitle_text = wx.StaticText( + panel, label="Code by ZFTurbo / GUI by Bas Curtiz" + ) + subtitle_text.SetForegroundColour(wx.BLACK) + main_sizer.Add(subtitle_text, 0, wx.ALIGN_CENTER | wx.TOP, 5) + + # Add GitHub link + github_link = wx.adv.HyperlinkCtrl( + panel, + -1, + "GitHub Repository", + "https://github.com/ZFTurbo/Music-Source-Separation-Training", + ) + github_link.SetNormalColour(wx.Colour(1, 118, 179)) # #0176B3 + github_link.SetHoverColour(wx.Colour(86, 91, 123)) # #565B7B + main_sizer.Add(github_link, 0, wx.ALIGN_CENTER | wx.TOP, 10) + + # Add Download models button on a new line with 10px bottom margin + download_models_btn = self.create_styled_button( + panel, "Download Models", self.on_download_models + ) + main_sizer.Add(download_models_btn, 0, wx.ALIGN_CENTER | wx.TOP | wx.BOTTOM, 10) + + # Training Configuration + self.training_panel = CollapsiblePanel(panel, "Training Configuration") + self.training_panel.toggle_button.SetFont(self.bold_font) + self.create_training_controls(self.training_panel.get_content_panel()) + main_sizer.Add(self.training_panel, 0, wx.EXPAND | wx.ALL, 10) + + # Inference Configuration + self.inference_panel = CollapsiblePanel(panel, "Inference Configuration") + self.inference_panel.toggle_button.SetFont(self.bold_font) + self.create_inference_controls(self.inference_panel.get_content_panel()) + main_sizer.Add(self.inference_panel, 0, wx.EXPAND | wx.ALL, 10) + + panel.SetSizer(main_sizer) + self.load_settings() + + def create_styled_button(self, parent, label, handler): + btn = wx.Button(parent, label=label, style=wx.BORDER_NONE) + btn.SetBackgroundColour(wx.Colour(1, 118, 179)) # #0176B3 + btn.SetForegroundColour(wx.WHITE) + btn.SetFont(self.bold_font) + + def on_enter(event): + btn.SetBackgroundColour(wx.Colour(86, 91, 123)) + event.Skip() + + def on_leave(event): + btn.SetBackgroundColour(wx.Colour(1, 118, 179)) + event.Skip() + + def on_click(event): + btn.SetBackgroundColour(wx.Colour(86, 91, 123)) + handler(event) + wx.CallLater(100, lambda: btn.SetBackgroundColour(wx.Colour(1, 118, 179))) + + btn.Bind(wx.EVT_ENTER_WINDOW, on_enter) + btn.Bind(wx.EVT_LEAVE_WINDOW, on_leave) + btn.Bind(wx.EVT_BUTTON, on_click) + + return btn + + def create_training_controls(self, panel): + sizer = wx.BoxSizer(wx.VERTICAL) + + # Model Type + model_type_sizer = wx.BoxSizer(wx.HORIZONTAL) + model_type_sizer.Add( + wx.StaticText(panel, label="Model Type:"), 0, wx.ALIGN_CENTER_VERTICAL + ) + self.model_type = wx.Choice( + panel, + choices=[ + "apollo", + "bandit", + "bandit_v2", + "bs_roformer", + "htdemucs", + "mdx23c", + "mel_band_roformer", + "scnet", + "scnet_unofficial", + "segm_models", + "swin_upernet", + "torchseg", + ], + ) + self.model_type.SetFont(self.font) + model_type_sizer.Add(self.model_type, 0, wx.LEFT, 5) + sizer.Add(model_type_sizer, 0, wx.EXPAND | wx.ALL, 5) + + # Config File + self.config_entry = self.add_browse_control( + panel, sizer, "Config File:", is_folder=False, is_config=True + ) + + # Start Checkpoint + self.checkpoint_entry = self.add_browse_control( + panel, sizer, "Checkpoint:", is_folder=False, is_checkpoint=True + ) + + # Results Path + self.result_path_entry = self.add_browse_control( + panel, sizer, "Results Path:", is_folder=True + ) + + # Data Paths + self.data_entry = self.add_browse_control( + panel, sizer, "Data Paths (separated by ';'):", is_folder=True + ) + + # Validation Paths + self.valid_entry = self.add_browse_control( + panel, sizer, "Validation Paths (separated by ';'):", is_folder=True + ) + + # Number of Workers and Device IDs + workers_device_sizer = wx.BoxSizer(wx.HORIZONTAL) + + workers_sizer = wx.BoxSizer(wx.HORIZONTAL) + workers_sizer.Add( + wx.StaticText(panel, label="Number of Workers:"), + 0, + wx.ALIGN_CENTER_VERTICAL, + ) + self.workers_entry = wx.TextCtrl(panel, value="4") + self.workers_entry.SetFont(self.font) + workers_sizer.Add(self.workers_entry, 0, wx.LEFT, 5) + workers_device_sizer.Add(workers_sizer, 0, wx.EXPAND) + + device_sizer = wx.BoxSizer(wx.HORIZONTAL) + device_sizer.Add( + wx.StaticText(panel, label="Device IDs (comma-separated):"), + 0, + wx.ALIGN_CENTER_VERTICAL | wx.LEFT, + 20, + ) + self.device_entry = wx.TextCtrl(panel, value="0") + self.device_entry.SetFont(self.font) + device_sizer.Add(self.device_entry, 0, wx.LEFT, 5) + workers_device_sizer.Add(device_sizer, 0, wx.EXPAND) + + sizer.Add(workers_device_sizer, 0, wx.EXPAND | wx.ALL, 5) + + # Run Training Button + self.run_button = self.create_styled_button( + panel, "Run Training", self.run_training + ) + sizer.Add(self.run_button, 0, wx.ALIGN_CENTER | wx.ALL, 10) + + panel.SetSizer(sizer) + + def create_inference_controls(self, panel): + sizer = wx.BoxSizer(wx.VERTICAL) + + # Model Type and Saved Combinations + infer_model_type_sizer = wx.BoxSizer(wx.HORIZONTAL) + infer_model_type_sizer.Add( + wx.StaticText(panel, label="Model Type:"), 0, wx.ALIGN_CENTER_VERTICAL + ) + self.infer_model_type = wx.Choice( + panel, + choices=[ + "apollo", + "bandit", + "bandit_v2", + "bs_roformer", + "htdemucs", + "mdx23c", + "mel_band_roformer", + "scnet", + "scnet_unofficial", + "segm_models", + "swin_upernet", + "torchseg", + ], + ) + self.infer_model_type.SetFont(self.font) + infer_model_type_sizer.Add(self.infer_model_type, 0, wx.LEFT, 5) + + # Add "Preset:" label + infer_model_type_sizer.Add( + wx.StaticText(panel, label="Preset:"), + 0, + wx.ALIGN_CENTER_VERTICAL | wx.LEFT, + 20, + ) + + # Add dropdown for saved combinations + self.saved_combinations_dropdown = wx.Choice(panel, choices=[]) + self.saved_combinations_dropdown.SetFont(self.font) + self.saved_combinations_dropdown.Bind( + wx.EVT_CHOICE, self.on_combination_selected + ) + + # Set the width to 200px and an appropriate height + self.saved_combinations_dropdown.SetMinSize( + (358, -1) + ) # -1 keeps the height unchanged + + # Add to sizer + infer_model_type_sizer.Add(self.saved_combinations_dropdown, 0, wx.LEFT, 5) + + # Add plus button + plus_button = self.create_styled_button(panel, "+", self.on_save_combination) + plus_button.SetMinSize((30, 30)) + infer_model_type_sizer.Add(plus_button, 0, wx.LEFT, 5) + + # Add help button with custom tooltip + help_button = wx.StaticText(panel, label="?") + help_button.SetFont(self.bold_font) + help_button.SetForegroundColour(wx.Colour(1, 118, 179)) # 0176B3 + tooltip_text = ( + "How to add a preset?\n\n" + "1. Click Download Models\n" + "2. Right-click a model's Config && Checkpoint\n" + "3. Save link as && select a proper destination\n" + "4. Copy the Model name\n" + "5. Close Download Models\n\n" + "6. Browse for the Config file\n" + "7. Browse for the Checkpoint\n" + "8. Select the Model Type\n" + "9. Click the + button\n" + "10. Paste the Model name && click OK\n\n" + "On next use, just select it from the Preset dropdown." + ) + + self.tooltip = CustomToolTip(self, tooltip_text) + self.tooltip.Hide() + + def on_help_enter(event): + self.tooltip.Position( + help_button.ClientToScreen((0, help_button.GetSize().height)), (0, 0) + ) + self.tooltip.Show() + + def on_help_leave(event): + self.tooltip.Hide() + + help_button.Bind(wx.EVT_ENTER_WINDOW, on_help_enter) + help_button.Bind(wx.EVT_LEAVE_WINDOW, on_help_leave) + + infer_model_type_sizer.Add( + help_button, 0, wx.LEFT | wx.ALIGN_CENTER_VERTICAL, 5 + ) + + sizer.Add(infer_model_type_sizer, 0, wx.EXPAND | wx.ALL, 5) + + # Config File + self.infer_config_entry = self.add_browse_control( + panel, sizer, "Config File:", is_folder=False, is_config=True + ) + + # Start Checkpoint + self.infer_checkpoint_entry = self.add_browse_control( + panel, sizer, "Checkpoint:", is_folder=False, is_checkpoint=True + ) + + # Input Folder + self.infer_input_entry = self.add_browse_control( + panel, sizer, "Input Folder:", is_folder=True + ) + + # Store Directory + self.infer_store_entry = self.add_browse_control( + panel, sizer, "Output Folder:", is_folder=True + ) + + # Extract Instrumental Checkbox + self.extract_instrumental_checkbox = wx.CheckBox( + panel, label="Extract Instrumental" + ) + self.extract_instrumental_checkbox.SetFont(self.font) + sizer.Add(self.extract_instrumental_checkbox, 0, wx.EXPAND | wx.ALL, 5) + + # Run Inference Button + self.run_infer_button = self.create_styled_button( + panel, "Run Inference", self.run_inference + ) + sizer.Add(self.run_infer_button, 0, wx.ALIGN_CENTER | wx.ALL, 10) + + panel.SetSizer(sizer) + + def add_browse_control( + self, panel, sizer, label, is_folder=False, is_config=False, is_checkpoint=False + ): + browse_sizer = wx.BoxSizer(wx.HORIZONTAL) + browse_sizer.Add(wx.StaticText(panel, label=label), 0, wx.ALIGN_CENTER_VERTICAL) + entry = wx.TextCtrl(panel) + entry.SetFont(self.font) + browse_sizer.Add(entry, 1, wx.EXPAND | wx.LEFT, 5) + browse_button = self.create_styled_button( + panel, + "Browse", + lambda event, + entry=entry, + is_folder=is_folder, + is_config=is_config, + is_checkpoint=is_checkpoint: self.browse( + event, entry, is_folder, is_config, is_checkpoint + ), + ) + browse_sizer.Add(browse_button, 0, wx.LEFT, 5) + sizer.Add(browse_sizer, 0, wx.EXPAND | wx.ALL, 5) + return entry + + def browse( + self, event, entry, is_folder=False, is_config=False, is_checkpoint=False + ): + if is_folder: + dialog = wx.DirDialog( + self, + "Choose a directory", + style=wx.DD_DEFAULT_STYLE | wx.DD_DIR_MUST_EXIST, + ) + else: + wildcard = "All files (*.*)|*.*" + if is_config: + wildcard = "YAML files (*.yaml)|*.yaml" + elif is_checkpoint: + wildcard = "Checkpoint files (*.bin;*.chpt;*.ckpt;*.th)|*.bin;*.chpt;*.ckpt;*.th" + + dialog = wx.FileDialog( + self, + "Choose a file", + style=wx.FD_OPEN | wx.FD_FILE_MUST_EXIST, + wildcard=wildcard, + ) + + dialog.SetFont(self.font) + if dialog.ShowModal() == wx.ID_OK: + entry.SetValue(dialog.GetPath()) + dialog.Destroy() + + def create_output_window(self, title, folder_path): + output_frame = wx.Frame( + self, title=title, style=wx.DEFAULT_FRAME_STYLE | wx.STAY_ON_TOP + ) + output_frame.SetIcon(self.GetIcon()) + output_frame.SetSize(994, 670) + output_frame.SetBackgroundColour(wx.Colour(0, 0, 0)) + output_frame.SetFont(self.font) + + # Set the position of the output frame to match the main frame + output_frame.SetPosition(self.GetPosition()) + + output_title = wx.StaticText(output_frame, label=title) + output_title.SetFont(self.bold_font) + output_title.SetForegroundColour(wx.WHITE) + + output_text = DarkThemedTextCtrl( + output_frame, style=wx.TE_MULTILINE | wx.TE_READONLY + ) + output_text.SetFont(self.font) + + open_folder_button = self.create_styled_button( + output_frame, + "Open Output Folder", + lambda event: open_store_folder(folder_path), + ) + + sizer = wx.BoxSizer(wx.VERTICAL) + sizer.Add(output_title, 0, wx.ALIGN_CENTER | wx.TOP, 10) + sizer.Add(output_text, 1, wx.EXPAND | wx.ALL, 10) + sizer.Add(open_folder_button, 0, wx.ALIGN_CENTER | wx.BOTTOM, 10) + output_frame.SetSizer(sizer) + + return output_frame, output_text + + def run_training(self, event): + model_type = self.model_type.GetStringSelection() + config_path = self.config_entry.GetValue() + start_checkpoint = self.checkpoint_entry.GetValue() + results_path = self.result_path_entry.GetValue() + data_paths = self.data_entry.GetValue() + valid_paths = self.valid_entry.GetValue() + num_workers = self.workers_entry.GetValue() + device_ids = self.device_entry.GetValue() + + if not model_type: + wx.MessageBox( + "Please select a model type.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + if not config_path: + wx.MessageBox( + "Please select a config file.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + if not results_path: + wx.MessageBox( + "Please specify a results path.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + if not data_paths: + wx.MessageBox( + "Please specify data paths.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + if not valid_paths: + wx.MessageBox( + "Please specify validation paths.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + + cmd = [ + sys.executable, + "train.py", + "--model_type", + model_type, + "--config_path", + config_path, + "--results_path", + results_path, + "--data_path", + *data_paths.split(";"), + "--valid_path", + *valid_paths.split(";"), + "--num_workers", + num_workers, + "--device_ids", + device_ids, + ] + + if start_checkpoint: + cmd += ["--start_check_point", start_checkpoint] + + output_queue = queue.Queue() + threading.Thread( + target=run_subprocess, args=(cmd, output_queue), daemon=True + ).start() + + output_frame, output_text = self.create_output_window( + "Training Output", results_path + ) + output_frame.Show() + update_output(output_text, output_queue) + + self.save_settings() + + def run_inference(self, event): + model_type = self.infer_model_type.GetStringSelection() + config_path = self.infer_config_entry.GetValue() + start_checkpoint = self.infer_checkpoint_entry.GetValue() + input_folder = self.infer_input_entry.GetValue() + store_dir = self.infer_store_entry.GetValue() + extract_instrumental = self.extract_instrumental_checkbox.GetValue() + + if not model_type: + wx.MessageBox( + "Please select a model type.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + if not config_path: + wx.MessageBox( + "Please select a config file.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + if not input_folder: + wx.MessageBox( + "Please specify an input folder.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + if not store_dir: + wx.MessageBox( + "Please specify an output folder.", "Input Error", wx.OK | wx.ICON_ERROR + ) + return + + cmd = [ + sys.executable, + "inference.py", + "--model_type", + model_type, + "--config_path", + config_path, + "--input_folder", + input_folder, + "--store_dir", + store_dir, + ] + + if start_checkpoint: + cmd += ["--start_check_point", start_checkpoint] + + if extract_instrumental: + cmd += ["--extract_instrumental"] + + output_queue = queue.Queue() + threading.Thread( + target=run_subprocess, args=(cmd, output_queue), daemon=True + ).start() + + output_frame, output_text = self.create_output_window( + "Inference Output", store_dir + ) + output_frame.Show() + update_output(output_text, output_queue) + + self.save_settings() + + def save_settings(self): + settings = { + "model_type": self.model_type.GetStringSelection(), + "config_path": self.config_entry.GetValue(), + "start_checkpoint": self.checkpoint_entry.GetValue(), + "results_path": self.result_path_entry.GetValue(), + "data_paths": self.data_entry.GetValue(), + "valid_paths": self.valid_entry.GetValue(), + "num_workers": self.workers_entry.GetValue(), + "device_ids": self.device_entry.GetValue(), + "infer_model_type": self.infer_model_type.GetStringSelection(), + "infer_config_path": self.infer_config_entry.GetValue(), + "infer_start_checkpoint": self.infer_checkpoint_entry.GetValue(), + "infer_input_folder": self.infer_input_entry.GetValue(), + "infer_store_dir": self.infer_store_entry.GetValue(), + "extract_instrumental": self.extract_instrumental_checkbox.GetValue(), + "saved_combinations": self.saved_combinations, + } + with open("settings.json", "w") as f: + json.dump(settings, f, indent=2, ensure_ascii=False) + + def load_settings(self): + try: + with open("settings.json", "r") as f: + settings = json.load(f) + + self.model_type.SetStringSelection(settings.get("model_type", "")) + self.config_entry.SetValue(settings.get("config_path", "")) + self.checkpoint_entry.SetValue(settings.get("start_checkpoint", "")) + self.result_path_entry.SetValue(settings.get("results_path", "")) + self.data_entry.SetValue(settings.get("data_paths", "")) + self.valid_entry.SetValue(settings.get("valid_paths", "")) + self.workers_entry.SetValue(settings.get("num_workers", "4")) + self.device_entry.SetValue(settings.get("device_ids", "0")) + + self.infer_model_type.SetStringSelection( + settings.get("infer_model_type", "") + ) + self.infer_config_entry.SetValue(settings.get("infer_config_path", "")) + self.infer_checkpoint_entry.SetValue( + settings.get("infer_start_checkpoint", "") + ) + self.infer_input_entry.SetValue(settings.get("infer_input_folder", "")) + self.infer_store_entry.SetValue(settings.get("infer_store_dir", "")) + self.extract_instrumental_checkbox.SetValue( + settings.get("extract_instrumental", False) + ) + self.saved_combinations = settings.get("saved_combinations", {}) + + self.update_saved_combinations() + except FileNotFoundError: + pass # If the settings file doesn't exist, use default values + + def on_download_models(self, event): + DownloadModelsFrame(self).Show() + + def on_save_combination(self, event): + dialog = wx.TextEntryDialog( + self, "Enter a name for this preset:", "Save Preset" + ) + if dialog.ShowModal() == wx.ID_OK: + name = dialog.GetValue() + if name: + combination = { + "model_type": self.infer_model_type.GetStringSelection(), + "config_path": self.infer_config_entry.GetValue(), + "checkpoint": self.infer_checkpoint_entry.GetValue(), + } + self.saved_combinations[name] = combination + self.update_saved_combinations() + self.save_settings() + dialog.Destroy() + + def on_combination_selected(self, event): + name = self.saved_combinations_dropdown.GetStringSelection() + if name: + combination = self.saved_combinations.get(name) + if combination: + self.infer_model_type.SetStringSelection(combination["model_type"]) + self.infer_config_entry.SetValue(combination["config_path"]) + self.infer_checkpoint_entry.SetValue(combination["checkpoint"]) + + def update_saved_combinations(self): + self.saved_combinations_dropdown.Clear() + for name in self.saved_combinations.keys(): + self.saved_combinations_dropdown.Append(name) + + +class DownloadModelsFrame(wx.Frame): + def __init__(self, parent): + super().__init__( + parent, + title="Download Models", + size=(994, 670), + style=wx.DEFAULT_FRAME_STYLE & ~(wx.RESIZE_BORDER | wx.MAXIMIZE_BOX), + ) + self.SetBackgroundColour(wx.Colour(247, 248, 250)) # #F7F8FA + self.SetFont( + wx.Font( + 9, + wx.FONTFAMILY_DEFAULT, + wx.FONTSTYLE_NORMAL, + wx.FONTWEIGHT_NORMAL, + False, + "Poppins", + ) + ) + + # Set the position of the Download Models frame to match the main frame + self.SetPosition(parent.GetPosition()) + + # Set the icon for the Download Models frame + icon = wx.Icon("gui/favicon.ico", wx.BITMAP_TYPE_ICO) + self.SetIcon(icon) + + panel = wx.Panel(self) + sizer = wx.BoxSizer(wx.VERTICAL) + + # Add WebView + self.web_view = wx.html2.WebView.New(panel) + self.web_view.LoadURL( + "https://bascurtiz.x10.mx/models-checkpoint-config-urls.html" + ) + self.web_view.Bind(wx.html2.EVT_WEBVIEW_NAVIGATING, self.on_link_click) + self.web_view.Bind(wx.html2.EVT_WEBVIEW_NAVIGATED, self.on_page_load) + sizer.Add(self.web_view, 1, wx.EXPAND) + + panel.SetSizer(sizer) + + def on_link_click(self, event): + url = event.GetURL() + if not url.startswith("https://bascurtiz.x10.mx"): + event.Veto() # Prevent the WebView from navigating + webbrowser.open(url) # Open the link in the default browser + + def on_page_load(self, event): + self.inject_custom_css() + + def inject_custom_css(self): + css = """ + body { + margin: 0; + padding: 0; + } + ::-webkit-scrollbar { + width: 12px; + } + ::-webkit-scrollbar-track { + background: #f1f1f1; + } + ::-webkit-scrollbar-thumb { + background: #888; + } + ::-webkit-scrollbar-thumb:hover { + background: #555; + } + """ + js = f"var style = document.createElement('style'); style.textContent = `{css}`; document.head.appendChild(style);" + self.web_view.RunScript(js) + + +if __name__ == "__main__": + app = wx.App() + frame = MainFrame() + frame.Show() + app.MainLoop() diff --git a/src/third_party/MusicSourceSeparationTraining/gui/mvsep.png b/src/third_party/MusicSourceSeparationTraining/gui/mvsep.png new file mode 100644 index 0000000000000000000000000000000000000000..2f98c5f145a908cd8b848e1d8f80adb0fd855bce --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/gui/mvsep.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d5e00fdd7ccf7259bf07434d44919e063cc494486690f1251826cb0b6fce133 +size 4336 diff --git a/src/third_party/MusicSourceSeparationTraining/gui/tutorial_screenshot.jpg b/src/third_party/MusicSourceSeparationTraining/gui/tutorial_screenshot.jpg new file mode 100644 index 0000000000000000000000000000000000000000..21f4ac393ee4089c8988ddcab2401140046f0397 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/gui/tutorial_screenshot.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f23831d64d11fa833e573200d4341ed024772d2d6f468602315d5426de6e960c +size 64729 diff --git a/src/third_party/MusicSourceSeparationTraining/gui/wx_msst_screen.png b/src/third_party/MusicSourceSeparationTraining/gui/wx_msst_screen.png new file mode 100644 index 0000000000000000000000000000000000000000..6175c82e7a0474bb9ed9541a4236a2a44776898a --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/gui/wx_msst_screen.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:102e87a5bbf54a96dde2acf1f0bcb5b8f3c324970fa832b1e9ab6218149a7eb4 +size 29416 diff --git a/src/third_party/MusicSourceSeparationTraining/inference.py b/src/third_party/MusicSourceSeparationTraining/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..b40900e99c0bab38d6dc4fa5e466505fce455816 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/inference.py @@ -0,0 +1,234 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" + +import glob +import os +import sys +import time + +import librosa +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +from tqdm.auto import tqdm + +# Using the embedded version of Python can also correctly import the utils module. +current_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(current_dir) + +import warnings + +from utils.audio_utils import denormalize_audio, draw_spectrogram, normalize_audio +from utils.model_utils import ( + apply_tta, + demix, + load_start_checkpoint, + prefer_target_instrument, +) +from utils.settings import get_model_from_config, parse_args_inference + +warnings.filterwarnings("ignore") + + +def run_folder( + model: "torch.nn.Module", + args: "argparse.Namespace", + config: dict, + device: "torch.device", + verbose: bool = False, +) -> None: + """ + Process a folder of audio files for source separation. + + Parameters: + ---------- + model : torch.nn.Module + Pre-trained model for source separation. + args : argparse.Namespace + Arguments containing input folder, output folder, and processing options. + config : dict + Configuration object with audio and inference settings. + device : torch.device + Device for model inference (CPU or CUDA). + verbose : bool, optional + If True, prints detailed information during processing. Default is False. + """ + + start_time = time.time() + model.eval() + + # Recursively collect all files from input directory + mixture_paths = sorted( + glob.glob(os.path.join(args.input_folder, "**/*.*"), recursive=True) + ) + mixture_paths = [p for p in mixture_paths if os.path.isfile(p)] + + sample_rate: int = getattr(config.audio, "sample_rate", 44100) + + print(f"Total files found: {len(mixture_paths)}. Using sample rate: {sample_rate}") + + instruments: list[str] = prefer_target_instrument(config)[:] + os.makedirs(args.store_dir, exist_ok=True) + + # Wrap paths with progress bar if not in verbose mode + if not verbose: + mixture_paths = tqdm(mixture_paths, desc="Total progress") + + # Determine whether to use detailed progress bar + if args.disable_detailed_pbar: + detailed_pbar = False + else: + detailed_pbar = True + + for path in mixture_paths: + # Get relative path from input folder + relative_path: str = os.path.relpath(path, args.input_folder) + # Extract directory and file name + dir_name: str = os.path.dirname(relative_path) + file_name: str = os.path.splitext(os.path.basename(path))[0] + + try: + mix, sr = librosa.load(path, sr=sample_rate, mono=False) + except Exception as e: + print(f"Cannot read track: {format(path)}") + print(f"Error message: {str(e)}") + continue + + # Convert mono audio to expected channel format if needed + if len(mix.shape) == 1: + mix = np.expand_dims(mix, axis=0) + if "num_channels" in config.audio: + if config.audio["num_channels"] == 2: + print("Convert mono track to stereo...") + mix = np.concatenate([mix, mix], axis=0) + + mix_orig = mix.copy() + + # Normalize input audio if enabled + if "normalize" in config.inference: + if config.inference["normalize"] is True: + mix, norm_params = normalize_audio(mix) + + # Perform source separation + waveforms_orig = demix( + config, model, mix, device, model_type=args.model_type, pbar=detailed_pbar + ) + + # Apply test-time augmentation if enabled + if args.use_tta: + waveforms_orig = apply_tta( + config, model, mix, waveforms_orig, device, args.model_type + ) + + # Extract instrumental track if requested + if args.extract_instrumental: + instr = "vocals" if "vocals" in instruments else instruments[0] + waveforms_orig["instrumental"] = mix_orig - waveforms_orig[instr] + if "instrumental" not in instruments: + instruments.append("instrumental") + + for instr in instruments: + estimates = waveforms_orig[instr] + + # Denormalize output audio if normalization was applied + if "normalize" in config.inference: + if config.inference["normalize"] is True: + estimates = denormalize_audio(estimates, norm_params) + + peak: float = float(np.abs(estimates).max()) + if peak <= 1.0 and args.pcm_type != "FLOAT": + codec = "flac" + else: + codec = "wav" + + subtype = args.pcm_type + + # Generate output directory structure using relative paths + dirnames, fname = format_filename( + args.filename_template, + instr=instr, + start_time=int(start_time), + file_name=file_name, + dir_name=dir_name, + model_type=args.model_type, + model=os.path.splitext(os.path.basename(args.start_check_point))[0], + ) + + # Create output directory + output_dir: str = os.path.join(args.store_dir, *dirnames) + os.makedirs(output_dir, exist_ok=True) + + output_path: str = os.path.join(output_dir, f"{fname}.{codec}") + sf.write(output_path, estimates.T, sr, subtype=subtype) + + # Draw and save spectrogram if enabled + if args.draw_spectro > 0: + output_img_path = os.path.join(output_dir, f"{fname}.jpg") + draw_spectrogram(estimates.T, sr, args.draw_spectro, output_img_path) + print("Wrote file:", output_img_path) + + print(f"Elapsed time: {time.time() - start_time:.2f} seconds.") + + +def format_filename(template, **kwargs): + """ + Formats a filename from a template. e.g "{file_name}/{instr}" + Using slashes ('/') in template will result in directories being created + Returns [dirnames, fname], i.e. an array of dir names and a single file name + """ + result = template + for k, v in kwargs.items(): + result = result.replace(f"{{{k}}}", str(v)) + *dirnames, fname = result.split("/") + return dirnames, fname + + +def proc_folder(dict_args): + args = parse_args_inference(dict_args) + device = "cpu" + if args.force_cpu: + device = "cpu" + elif torch.cuda.is_available(): + print("CUDA is available, use --force_cpu to disable it.") + device = ( + f"cuda:{args.device_ids[0]}" + if isinstance(args.device_ids, list) + else f"cuda:{args.device_ids}" + ) + elif torch.backends.mps.is_available(): + device = "mps" + + print("Using device: ", device) + + model_load_start_time = time.time() + torch.backends.cudnn.benchmark = True + + model, config = get_model_from_config(args.model_type, args.config_path) + if "model_type" in config.training: + args.model_type = config.training.model_type + if args.start_check_point: + checkpoint = torch.load( + args.start_check_point, weights_only=False, map_location="cpu" + ) + load_start_checkpoint(args, model, checkpoint, type_="inference") + + print("Instruments: {}".format(config.training.instruments)) + + # in case multiple CUDA GPUs are used and --device_ids arg is passed + if ( + isinstance(args.device_ids, list) + and len(args.device_ids) > 1 + and not args.force_cpu + ): + model = nn.DataParallel(model, device_ids=args.device_ids) + + model = model.to(device) + + print("Model load time: {:.2f} sec".format(time.time() - model_load_start_time)) + + run_folder(model, args, config, device, verbose=True) + + +if __name__ == "__main__": + proc_folder(None) diff --git a/src/third_party/MusicSourceSeparationTraining/inference_api.py b/src/third_party/MusicSourceSeparationTraining/inference_api.py new file mode 100644 index 0000000000000000000000000000000000000000..f507cdd46da33dce24625d5c895100fc3d2800ef --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/inference_api.py @@ -0,0 +1,119 @@ +import os +import sys +import warnings + +import numpy as np +import torch +import torchaudio + +sys.path.append(os.path.dirname(os.path.abspath(__file__))) + +from utils.audio_utils import denormalize_audio, normalize_audio +from utils.model_utils import demix, load_start_checkpoint +from utils.settings import get_model_from_config + +warnings.filterwarnings("ignore") + + +class Separator: + def __init__( + self, + config_path: str, + checkpoint_path: str, + model_type: str = "mel_band_roformer", + device: str = "auto", + ): + if device == "auto": + device = "cuda:0" if torch.cuda.is_available() else "cpu" + self.device = torch.device(device) + self.model_type = model_type + + torch.backends.cudnn.benchmark = True + self.model, self.config = get_model_from_config(model_type, config_path) + + if "model_type" in self.config.training: + self.model_type = self.config.training.model_type + + from argparse import Namespace + + fake_args = Namespace( + model_type=self.model_type, + config_path=config_path, + start_check_point=checkpoint_path, + device="auto", + output_dir="./output", + use_tta=False, + extract_instrumental=True, + pcm_type="FLOAT", + lora_checkpoint_loralib="", + draw_spectro=False, + ) + ckpt = torch.load(checkpoint_path, weights_only=False, map_location="cpu") + load_start_checkpoint(fake_args, self.model, ckpt, type_="inference") + + self.model = self.model.to(self.device) + self.model.eval() + self.sample_rate = getattr(self.config.audio, "sample_rate", 44100) + + def separate(self, wav: torch.Tensor, sr: int): + """ + Args: + wav: Waveform returned by torchaudio.load, shape (channels, samples) + sr: Sample rate + Returns: + vocal_wav: np.ndarray, shape (channels, samples) + inst_wav: np.ndarray, shape (channels, samples) + sr: int, output sample rate + """ + # Resample if needed + if sr != self.sample_rate: + wav = torchaudio.transforms.Resample(sr, self.sample_rate)(wav) + sr = self.sample_rate + + mix = wav.numpy() + + # Convert mono to stereo + if mix.shape[0] == 1 and getattr(self.config.audio, "num_channels", 1) == 2: + mix = np.concatenate([mix, mix], axis=0) + + mix_orig = mix.copy() + + # Normalize + norm_params = None + if getattr(self.config.inference, "normalize", False): + mix, norm_params = normalize_audio(mix) + + # Separate + waveforms = demix( + self.config, + self.model, + mix, + self.device, + model_type=self.model_type, + pbar=True, + ) + + # Extract vocals + vocal_wav = waveforms.get("vocals", list(waveforms.values())[0]) + if norm_params is not None: + vocal_wav = denormalize_audio(vocal_wav, norm_params) + + # Instrumental = original mix - vocals + inst_wav = mix_orig - vocal_wav + + return vocal_wav, inst_wav, sr + + +# ---- Example Usage ---- +if __name__ == "__main__": + sep = Separator( + config_path="ckpts/config_vocals_mel_band_roformer_kj.yaml", + checkpoint_path="ckpts/MelBandRoformer.ckpt", + device="cuda:0", + ) + + wav, sr = torchaudio.load("path/to/input.mp3") + vocal_wav, inst_wav, sr = sep.separate(wav, sr) + + torchaudio.save("output_vocals.wav", torch.from_numpy(vocal_wav), sr) + torchaudio.save("output_instrumental.wav", torch.from_numpy(inst_wav), sr) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dd7b0fb4edea9e6be3eb367709d4a8461e6eed12 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/__init__.py @@ -0,0 +1,678 @@ +import os.path +from collections import defaultdict +from itertools import chain, combinations +from typing import Any, Dict, Iterator, Mapping, Optional, Tuple, Type, TypedDict + +import pytorch_lightning as pl +import torch +import torchaudio as ta +import torchmetrics as tm +from asteroid import losses as asteroid_losses +from models.bandit.core import loss, model +from models.bandit.core import metrics as metrics_ +from models.bandit.core.data._types import BatchedDataDict +from models.bandit.core.data.augmentation import BaseAugmentor, StemAugmentor +from models.bandit.core.utils import audio as audio_ +from models.bandit.core.utils.audio import BaseFader + +# from deepspeed.ops.adam import DeepSpeedCPUAdam +# from geoopt import optim as gooptim +from pytorch_lightning.utilities.types import STEP_OUTPUT +from torch import nn, optim +from torch.optim import lr_scheduler +from torch.optim.lr_scheduler import LRScheduler + +# from pandas.io.json._normalize import nested_to_record + +ConfigDict = TypedDict("ConfigDict", {"name": str, "kwargs": Dict[str, Any]}) + + +class SchedulerConfigDict(ConfigDict): + monitor: str + + +OptimizerSchedulerConfigDict = TypedDict( + "OptimizerSchedulerConfigDict", + {"optimizer": ConfigDict, "scheduler": SchedulerConfigDict}, + total=False, +) + + +class LRSchedulerReturnDict(TypedDict, total=False): + scheduler: LRScheduler + monitor: str + + +class ConfigureOptimizerReturnDict(TypedDict, total=False): + optimizer: torch.optim.Optimizer + lr_scheduler: LRSchedulerReturnDict + + +OutputType = Dict[str, Any] +MetricsType = Dict[str, torch.Tensor] + + +def get_optimizer_class(name: str) -> Type[optim.Optimizer]: + if name == "DeepSpeedCPUAdam": + return DeepSpeedCPUAdam + + for module in [optim, gooptim]: + if name in module.__dict__: + return module.__dict__[name] + + raise NameError + + +def parse_optimizer_config( + config: OptimizerSchedulerConfigDict, parameters: Iterator[nn.Parameter] +) -> ConfigureOptimizerReturnDict: + optim_class = get_optimizer_class(config["optimizer"]["name"]) + optimizer = optim_class(parameters, **config["optimizer"]["kwargs"]) + + optim_dict: ConfigureOptimizerReturnDict = { + "optimizer": optimizer, + } + + if "scheduler" in config: + lr_scheduler_class_ = config["scheduler"]["name"] + lr_scheduler_class = lr_scheduler.__dict__[lr_scheduler_class_] + lr_scheduler_dict: LRSchedulerReturnDict = { + "scheduler": lr_scheduler_class(optimizer, **config["scheduler"]["kwargs"]) + } + + if lr_scheduler_class_ == "ReduceLROnPlateau": + lr_scheduler_dict["monitor"] = config["scheduler"]["monitor"] + + optim_dict["lr_scheduler"] = lr_scheduler_dict + + return optim_dict + + +def parse_model_config(config: ConfigDict) -> Any: + name = config["name"] + + for module in [model]: + if name in module.__dict__: + return module.__dict__[name](**config["kwargs"]) + + raise NameError + + +_LEGACY_LOSS_NAMES = ["HybridL1Loss"] + + +def _parse_legacy_loss_config(config: ConfigDict) -> nn.Module: + name = config["name"] + + if name == "HybridL1Loss": + return loss.TimeFreqL1Loss(**config["kwargs"]) + + raise NameError + + +def parse_loss_config(config: ConfigDict) -> nn.Module: + name = config["name"] + + if name in _LEGACY_LOSS_NAMES: + return _parse_legacy_loss_config(config) + + for module in [loss, nn.modules.loss, asteroid_losses]: + if name in module.__dict__: + # print(config["kwargs"]) + return module.__dict__[name](**config["kwargs"]) + + raise NameError + + +def get_metric(config: ConfigDict) -> tm.Metric: + name = config["name"] + + for module in [tm, metrics_]: + if name in module.__dict__: + return module.__dict__[name](**config["kwargs"]) + raise NameError + + +def parse_metric_config(config: Dict[str, ConfigDict]) -> tm.MetricCollection: + metrics = {} + + for metric in config: + metrics[metric] = get_metric(config[metric]) + + return tm.MetricCollection(metrics) + + +def parse_fader_config(config: ConfigDict) -> BaseFader: + name = config["name"] + + for module in [audio_]: + if name in module.__dict__: + return module.__dict__[name](**config["kwargs"]) + + raise NameError + + +class LightningSystem(pl.LightningModule): + _VOX_STEMS = ["speech", "vocals"] + _BG_STEMS = ["background", "effects", "mne"] + + def __init__( + self, config: Dict, loss_adjustment: float = 1.0, attach_fader: bool = False + ) -> None: + super().__init__() + self.optimizer_config = config["optimizer"] + self.model = parse_model_config(config["model"]) + self.loss = parse_loss_config(config["loss"]) + self.metrics = nn.ModuleDict( + { + stem: parse_metric_config(config["metrics"]["dev"]) + for stem in self.model.stems + } + ) + + self.metrics.disallow_fsdp = True + + self.test_metrics = nn.ModuleDict( + { + stem: parse_metric_config(config["metrics"]["test"]) + for stem in self.model.stems + } + ) + + self.test_metrics.disallow_fsdp = True + + self.fs = config["model"]["kwargs"]["fs"] + + self.fader_config = config["inference"]["fader"] + if attach_fader: + self.fader = parse_fader_config(config["inference"]["fader"]) + else: + self.fader = None + + self.augmentation: Optional[BaseAugmentor] + if config.get("augmentation", None) is not None: + self.augmentation = StemAugmentor(**config["augmentation"]) + else: + self.augmentation = None + + self.predict_output_path: Optional[str] = None + self.loss_adjustment = loss_adjustment + + self.val_prefix = None + self.test_prefix = None + + def configure_optimizers(self) -> Any: + return parse_optimizer_config( + self.optimizer_config, self.trainer.model.parameters() + ) + + def compute_loss( + self, batch: BatchedDataDict, output: OutputType + ) -> Dict[str, torch.Tensor]: + return {"loss": self.loss(output, batch)} + + def update_metrics( + self, batch: BatchedDataDict, output: OutputType, mode: str + ) -> None: + if mode == "test": + metrics = self.test_metrics + else: + metrics = self.metrics + + for stem, metric in metrics.items(): + if stem == "mne:+": + stem = "mne" + + # print(f"matching for {stem}") + if mode == "train": + metric.update( + output["audio"][stem], # .cpu(), + batch["audio"][stem], # .cpu() + ) + else: + if stem not in batch["audio"]: + matched = False + if stem in self._VOX_STEMS: + for bstem in self._VOX_STEMS: + if bstem in batch["audio"]: + batch["audio"][stem] = batch["audio"][bstem] + matched = True + break + elif stem in self._BG_STEMS: + for bstem in self._BG_STEMS: + if bstem in batch["audio"]: + batch["audio"][stem] = batch["audio"][bstem] + matched = True + break + else: + matched = True + + # print(batch["audio"].keys()) + + if matched: + # print(f"matched {stem}!") + if stem == "mne" and "mne" not in output["audio"]: + output["audio"]["mne"] = ( + output["audio"]["music"] + output["audio"]["effects"] + ) + + metric.update( + output["audio"][stem], # .cpu(), + batch["audio"][stem], # .cpu(), + ) + + # print(metric.compute()) + + def compute_metrics(self, mode: str = "dev") -> Dict[str, torch.Tensor]: + if mode == "test": + metrics = self.test_metrics + else: + metrics = self.metrics + + metric_dict = {} + + for stem, metric in metrics.items(): + md = metric.compute() + metric_dict.update({f"{stem}/{k}": v for k, v in md.items()}) + + self.log_dict(metric_dict, prog_bar=True, logger=False) + + return metric_dict + + def reset_metrics(self, test_mode: bool = False) -> None: + if test_mode: + metrics = self.test_metrics + else: + metrics = self.metrics + + for _, metric in metrics.items(): + metric.reset() + + def forward(self, batch: BatchedDataDict) -> Any: + batch, output = self.model(batch) + + return batch, output + + def common_step(self, batch: BatchedDataDict, mode: str) -> Any: + batch, output = self.forward(batch) + # print(batch) + # print(output) + loss_dict = self.compute_loss(batch, output) + + with torch.no_grad(): + self.update_metrics(batch, output, mode=mode) + + if mode == "train": + self.log("loss", loss_dict["loss"], prog_bar=True) + + return output, loss_dict + + def training_step(self, batch: BatchedDataDict) -> Dict[str, Any]: + if self.augmentation is not None: + with torch.no_grad(): + batch = self.augmentation(batch) + + _, loss_dict = self.common_step(batch, mode="train") + + with torch.inference_mode(): + self.log_dict_with_prefix( + loss_dict, "train", batch_size=batch["audio"]["mixture"].shape[0] + ) + + loss_dict["loss"] *= self.loss_adjustment + + return loss_dict + + def on_train_batch_end( + self, outputs: STEP_OUTPUT, batch: BatchedDataDict, batch_idx: int + ) -> None: + metric_dict = self.compute_metrics() + self.log_dict_with_prefix(metric_dict, "train") + self.reset_metrics() + + def validation_step( + self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0 + ) -> Dict[str, Any]: + with torch.inference_mode(): + curr_val_prefix = f"val{dataloader_idx}" if dataloader_idx > 0 else "val" + + if curr_val_prefix != self.val_prefix: + # print(f"Switching to validation dataloader {dataloader_idx}") + if self.val_prefix is not None: + self._on_validation_epoch_end() + self.val_prefix = curr_val_prefix + _, loss_dict = self.common_step(batch, mode="val") + + self.log_dict_with_prefix( + loss_dict, + self.val_prefix, + batch_size=batch["audio"]["mixture"].shape[0], + prog_bar=True, + add_dataloader_idx=False, + ) + + return loss_dict + + def on_validation_epoch_end(self) -> None: + self._on_validation_epoch_end() + + def _on_validation_epoch_end(self) -> None: + metric_dict = self.compute_metrics() + self.log_dict_with_prefix( + metric_dict, self.val_prefix, prog_bar=True, add_dataloader_idx=False + ) + # self.logger.save() + # print(self.val_prefix, "Validation metrics:", metric_dict) + self.reset_metrics() + + def old_predtest_step( + self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0 + ) -> Tuple[BatchedDataDict, OutputType]: + audio_batch = batch["audio"]["mixture"] + track_batch = batch.get("track", ["" for _ in range(len(audio_batch))]) + + output_list_of_dicts = [ + self.fader(audio[None, ...], lambda a: self.test_forward(a, track)) + for audio, track in zip(audio_batch, track_batch) + ] + + output_dict_of_lists = defaultdict(list) + + for output_dict in output_list_of_dicts: + for stem, audio in output_dict.items(): + output_dict_of_lists[stem].append(audio) + + output = { + "audio": { + stem: torch.concat(output_list, dim=0) + for stem, output_list in output_dict_of_lists.items() + } + } + + return batch, output + + def predtest_step( + self, batch: BatchedDataDict, batch_idx: int = -1, dataloader_idx: int = 0 + ) -> Tuple[BatchedDataDict, OutputType]: + if getattr(self.model, "bypass_fader", False): + batch, output = self.model(batch) + else: + audio_batch = batch["audio"]["mixture"] + output = self.fader( + audio_batch, lambda a: self.test_forward(a, "", batch=batch) + ) + + return batch, output + + def test_forward( + self, audio: torch.Tensor, track: str = "", batch: BatchedDataDict = None + ) -> torch.Tensor: + if self.fader is None: + self.attach_fader() + + cond = batch.get("condition", None) + + if cond is not None and cond.shape[0] == 1: + cond = cond.repeat(audio.shape[0], 1) + + _, output = self.forward( + { + "audio": {"mixture": audio}, + "track": track, + "condition": cond, + } + ) # TODO: support track properly + + return output["audio"] + + def on_test_epoch_start(self) -> None: + self.attach_fader(force_reattach=True) + + def test_step( + self, batch: BatchedDataDict, batch_idx: int, dataloader_idx: int = 0 + ) -> Any: + curr_test_prefix = f"test{dataloader_idx}" + + # print(batch["audio"].keys()) + + if curr_test_prefix != self.test_prefix: + # print(f"Switching to test dataloader {dataloader_idx}") + if self.test_prefix is not None: + self._on_test_epoch_end() + self.test_prefix = curr_test_prefix + + with torch.inference_mode(): + _, output = self.predtest_step(batch, batch_idx, dataloader_idx) + # print(output) + self.update_metrics(batch, output, mode="test") + + return output + + def on_test_epoch_end(self) -> None: + self._on_test_epoch_end() + + def _on_test_epoch_end(self) -> None: + metric_dict = self.compute_metrics(mode="test") + self.log_dict_with_prefix( + metric_dict, self.test_prefix, prog_bar=True, add_dataloader_idx=False + ) + # self.logger.save() + # print(self.test_prefix, "Test metrics:", metric_dict) + self.reset_metrics() + + def predict_step( + self, + batch: BatchedDataDict, + batch_idx: int = 0, + dataloader_idx: int = 0, + include_track_name: Optional[bool] = None, + get_no_vox_combinations: bool = True, + get_residual: bool = False, + treat_batch_as_channels: bool = False, + fs: Optional[int] = None, + ) -> Any: + assert self.predict_output_path is not None + + batch_size = batch["audio"]["mixture"].shape[0] + + if include_track_name is None: + include_track_name = batch_size > 1 + + with torch.inference_mode(): + batch, output = self.predtest_step(batch, batch_idx, dataloader_idx) + print("Pred test finished...") + torch.cuda.empty_cache() + metric_dict = {} + + if get_residual: + mixture = batch["audio"]["mixture"] + extracted = sum([output["audio"][stem] for stem in output["audio"]]) + residual = mixture - extracted + print(extracted.shape, mixture.shape, residual.shape) + + output["audio"]["residual"] = residual + + if get_no_vox_combinations: + no_vox_stems = [ + stem for stem in output["audio"] if stem not in self._VOX_STEMS + ] + no_vox_combinations = chain.from_iterable( + combinations(no_vox_stems, r) for r in range(2, len(no_vox_stems) + 1) + ) + + for combination in no_vox_combinations: + combination_ = list(combination) + output["audio"]["+".join(combination_)] = sum( + [output["audio"][stem] for stem in combination_] + ) + + if treat_batch_as_channels: + for stem in output["audio"]: + output["audio"][stem] = output["audio"][stem].reshape( + 1, -1, output["audio"][stem].shape[-1] + ) + batch_size = 1 + + for b in range(batch_size): + print("!!", b) + for stem in output["audio"]: + print(f"Saving audio for {stem} to {self.predict_output_path}") + track_name = batch["track"][b].split("/")[-1] + + if batch.get("audio", {}).get(stem, None) is not None: + self.test_metrics[stem].reset() + metrics = self.test_metrics[stem]( + batch["audio"][stem][[b], ...], output["audio"][stem][[b], ...] + ) + snr = metrics["snr"] + sisnr = metrics["sisnr"] + sdr = metrics["sdr"] + metric_dict[stem] = metrics + print( + track_name, + f"snr={snr:2.2f} dB", + f"sisnr={sisnr:2.2f}", + f"sdr={sdr:2.2f} dB", + ) + filename = f"{stem} - snr={snr:2.2f}dB - sdr={sdr:2.2f}dB.wav" + else: + filename = f"{stem}.wav" + + if include_track_name: + output_dir = os.path.join(self.predict_output_path, track_name) + else: + output_dir = self.predict_output_path + + os.makedirs(output_dir, exist_ok=True) + + if fs is None: + fs = self.fs + + ta.save( + os.path.join(output_dir, filename), + output["audio"][stem][b, ...].cpu(), + fs, + ) + + return metric_dict + + def get_stems( + self, + batch: BatchedDataDict, + batch_idx: int = 0, + dataloader_idx: int = 0, + include_track_name: Optional[bool] = None, + get_no_vox_combinations: bool = True, + get_residual: bool = False, + treat_batch_as_channels: bool = False, + fs: Optional[int] = None, + ) -> Any: + assert self.predict_output_path is not None + + batch_size = batch["audio"]["mixture"].shape[0] + + if include_track_name is None: + include_track_name = batch_size > 1 + + with torch.inference_mode(): + batch, output = self.predtest_step(batch, batch_idx, dataloader_idx) + torch.cuda.empty_cache() + metric_dict = {} + + if get_residual: + mixture = batch["audio"]["mixture"] + extracted = sum([output["audio"][stem] for stem in output["audio"]]) + residual = mixture - extracted + # print(extracted.shape, mixture.shape, residual.shape) + + output["audio"]["residual"] = residual + + if get_no_vox_combinations: + no_vox_stems = [ + stem for stem in output["audio"] if stem not in self._VOX_STEMS + ] + no_vox_combinations = chain.from_iterable( + combinations(no_vox_stems, r) for r in range(2, len(no_vox_stems) + 1) + ) + + for combination in no_vox_combinations: + combination_ = list(combination) + output["audio"]["+".join(combination_)] = sum( + [output["audio"][stem] for stem in combination_] + ) + + if treat_batch_as_channels: + for stem in output["audio"]: + output["audio"][stem] = output["audio"][stem].reshape( + 1, -1, output["audio"][stem].shape[-1] + ) + batch_size = 1 + + result = {} + for b in range(batch_size): + for stem in output["audio"]: + track_name = batch["track"][b].split("/")[-1] + + if batch.get("audio", {}).get(stem, None) is not None: + self.test_metrics[stem].reset() + metrics = self.test_metrics[stem]( + batch["audio"][stem][[b], ...], output["audio"][stem][[b], ...] + ) + snr = metrics["snr"] + sisnr = metrics["sisnr"] + sdr = metrics["sdr"] + metric_dict[stem] = metrics + print( + track_name, + f"snr={snr:2.2f} dB", + f"sisnr={sisnr:2.2f}", + f"sdr={sdr:2.2f} dB", + ) + filename = f"{stem} - snr={snr:2.2f}dB - sdr={sdr:2.2f}dB.wav" + else: + filename = f"{stem}.wav" + + if include_track_name: + output_dir = os.path.join(self.predict_output_path, track_name) + else: + output_dir = self.predict_output_path + + os.makedirs(output_dir, exist_ok=True) + + if fs is None: + fs = self.fs + + result[stem] = output["audio"][stem][b, ...].cpu().numpy() + + return result + + def load_state_dict( + self, state_dict: Mapping[str, Any], strict: bool = False + ) -> Any: + return super().load_state_dict(state_dict, strict=False) + + def set_predict_output_path(self, path: str) -> None: + self.predict_output_path = path + os.makedirs(self.predict_output_path, exist_ok=True) + + self.attach_fader() + + def attach_fader(self, force_reattach=False) -> None: + if self.fader is None or force_reattach: + self.fader = parse_fader_config(self.fader_config) + self.fader.to(self.device) + + def log_dict_with_prefix( + self, + dict_: Dict[str, torch.Tensor], + prefix: str, + batch_size: Optional[int] = None, + **kwargs: Any, + ) -> None: + self.log_dict( + {f"{prefix}/{k}": v for k, v in dict_.items()}, + batch_size=batch_size, + logger=True, + sync_dist=True, + **kwargs, + ) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9d4d672bd3b6ad90a26e19ee6c26e02ee3be84c --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/__init__.py @@ -0,0 +1,2 @@ +from .dnr.datamodule import DivideAndRemasterDataModule +from .musdb.datamodule import MUSDB18DataModule diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/_types.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/_types.py new file mode 100644 index 0000000000000000000000000000000000000000..65e4607a558e6b6a65ee68de883b69e282f8fcf4 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/_types.py @@ -0,0 +1,17 @@ +from typing import Dict, Sequence, TypedDict + +import torch + +AudioDict = Dict[str, torch.Tensor] + +DataDict = TypedDict("DataDict", {"audio": AudioDict, "track": str}) + +BatchedDataDict = TypedDict( + "BatchedDataDict", {"audio": AudioDict, "track": Sequence[str]} +) + + +class DataDictWithLanguage(TypedDict): + audio: AudioDict + track: str + language: str diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/augmentation.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/augmentation.py new file mode 100644 index 0000000000000000000000000000000000000000..f6c48b9a4e2f7161aefb2a68b21244ed087d1516 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/augmentation.py @@ -0,0 +1,100 @@ +from abc import ABC +from typing import Any, Dict, Union + +import torch +import torch_audiomentations as tam +from models.bandit.core.data._types import BatchedDataDict, DataDict +from torch import nn + + +class BaseAugmentor(nn.Module, ABC): + def forward( + self, item: Union[DataDict, BatchedDataDict] + ) -> Union[DataDict, BatchedDataDict]: + raise NotImplementedError + + +class StemAugmentor(BaseAugmentor): + def __init__( + self, + audiomentations: Dict[str, Dict[str, Any]], + fix_clipping: bool = True, + scaler_margin: float = 0.5, + apply_both_default_and_common: bool = False, + ) -> None: + super().__init__() + + augmentations = {} + + self.has_default = "[default]" in audiomentations + self.has_common = "[common]" in audiomentations + self.apply_both_default_and_common = apply_both_default_and_common + + for stem in audiomentations: + if audiomentations[stem]["name"] == "Compose": + augmentations[stem] = getattr(tam, audiomentations[stem]["name"])( + [ + getattr(tam, aug["name"])(**aug["kwargs"]) + for aug in audiomentations[stem]["kwargs"]["transforms"] + ], + **audiomentations[stem]["kwargs"]["kwargs"], + ) + else: + augmentations[stem] = getattr(tam, audiomentations[stem]["name"])( + **audiomentations[stem]["kwargs"] + ) + + self.augmentations = nn.ModuleDict(augmentations) + self.fix_clipping = fix_clipping + self.scaler_margin = scaler_margin + + def check_and_fix_clipping( + self, item: Union[DataDict, BatchedDataDict] + ) -> Union[DataDict, BatchedDataDict]: + max_abs = [] + + for stem in item["audio"]: + max_abs.append(item["audio"][stem].abs().max().item()) + + if max(max_abs) > 1.0: + scaler = 1.0 / ( + max(max_abs) + + torch.rand((1,), device=item["audio"]["mixture"].device) + * self.scaler_margin + ) + + for stem in item["audio"]: + item["audio"][stem] *= scaler + + return item + + def forward( + self, item: Union[DataDict, BatchedDataDict] + ) -> Union[DataDict, BatchedDataDict]: + for stem in item["audio"]: + if stem == "mixture": + continue + + if self.has_common: + item["audio"][stem] = self.augmentations["[common]"]( + item["audio"][stem] + ).samples + + if stem in self.augmentations: + item["audio"][stem] = self.augmentations[stem]( + item["audio"][stem] + ).samples + elif self.has_default: + if not self.has_common or self.apply_both_default_and_common: + item["audio"][stem] = self.augmentations["[default]"]( + item["audio"][stem] + ).samples + + item["audio"]["mixture"] = sum( + [item["audio"][stem] for stem in item["audio"] if stem != "mixture"] + ) # type: ignore[call-overload, assignment] + + if self.fix_clipping: + item = self.check_and_fix_clipping(item) + + return item diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/augmented.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/augmented.py new file mode 100644 index 0000000000000000000000000000000000000000..3c0524409bf99b009605989eba5d5f46f0560f2e --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/augmented.py @@ -0,0 +1,34 @@ +import warnings +from typing import Dict, Optional, Union + +import torch +from torch import nn +from torch.utils import data + + +class AugmentedDataset(data.Dataset): + def __init__( + self, + dataset: data.Dataset, + augmentation: nn.Module = nn.Identity(), + target_length: Optional[int] = None, + ) -> None: + warnings.warn( + "This class is no longer used. Attach augmentation to " + "the LightningSystem instead.", + DeprecationWarning, + ) + + self.dataset = dataset + self.augmentation = augmentation + + self.ds_length: int = len(dataset) # type: ignore[arg-type] + self.length = target_length if target_length is not None else self.ds_length + + def __getitem__(self, index: int) -> Dict[str, Union[str, Dict[str, torch.Tensor]]]: + item = self.dataset[index % self.ds_length] + item = self.augmentation(item) + return item + + def __len__(self) -> int: + return self.length diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/base.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1039023e76c4e390140a4b60f619372e4dd94ecd --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/base.py @@ -0,0 +1,53 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +import torch +from models.bandit.core.data._types import AudioDict +from torch.utils import data + + +class BaseSourceSeparationDataset(data.Dataset, ABC): + def __init__( + self, + split: str, + stems: List[str], + files: List[str], + data_path: str, + fs: int, + npy_memmap: bool, + recompute_mixture: bool, + ): + self.split = split + self.stems = stems + self.stems_no_mixture = [s for s in stems if s != "mixture"] + self.files = files + self.data_path = data_path + self.fs = fs + self.npy_memmap = npy_memmap + self.recompute_mixture = recompute_mixture + + @abstractmethod + def get_stem(self, *, stem: str, identifier: Dict[str, Any]) -> torch.Tensor: + raise NotImplementedError + + def _get_audio(self, stems, identifier: Dict[str, Any]): + audio = {} + for stem in stems: + audio[stem] = self.get_stem(stem=stem, identifier=identifier) + + return audio + + def get_audio(self, identifier: Dict[str, Any]) -> AudioDict: + if self.recompute_mixture: + audio = self._get_audio(self.stems_no_mixture, identifier=identifier) + audio["mixture"] = self.compute_mixture(audio) + return audio + else: + return self._get_audio(self.stems, identifier=identifier) + + @abstractmethod + def get_identifier(self, index: int) -> Dict[str, Any]: + pass + + def compute_mixture(self, audio: AudioDict) -> torch.Tensor: + return sum(audio[stem] for stem in audio if stem != "mixture") diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/datamodule.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/datamodule.py new file mode 100644 index 0000000000000000000000000000000000000000..ba26707e891d35940672c8e1ca2b523b14f73f4c --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/datamodule.py @@ -0,0 +1,68 @@ +import os +from typing import Mapping, Optional + +import pytorch_lightning as pl + +from .dataset import ( + DivideAndRemasterDataset, + DivideAndRemasterDeterministicChunkDataset, + DivideAndRemasterRandomChunkDataset, + DivideAndRemasterRandomChunkDatasetWithSpeechReverb, +) + + +def DivideAndRemasterDataModule( + data_root: str = "$DATA_ROOT/DnR/v2", + batch_size: int = 2, + num_workers: int = 8, + train_kwargs: Optional[Mapping] = None, + val_kwargs: Optional[Mapping] = None, + test_kwargs: Optional[Mapping] = None, + datamodule_kwargs: Optional[Mapping] = None, + use_speech_reverb: bool = False, + # augmentor=None +) -> pl.LightningDataModule: + if train_kwargs is None: + train_kwargs = {} + + if val_kwargs is None: + val_kwargs = {} + + if test_kwargs is None: + test_kwargs = {} + + if datamodule_kwargs is None: + datamodule_kwargs = {} + + if num_workers is None: + num_workers = os.cpu_count() + + if num_workers is None: + num_workers = 32 + + num_workers = min(num_workers, 64) + + if use_speech_reverb: + train_cls = DivideAndRemasterRandomChunkDatasetWithSpeechReverb + else: + train_cls = DivideAndRemasterRandomChunkDataset + + train_dataset = train_cls(data_root, "train", **train_kwargs) + + # if augmentor is not None: + # train_dataset = AugmentedDataset(train_dataset, augmentor) + + datamodule = pl.LightningDataModule.from_datasets( + train_dataset=train_dataset, + val_dataset=DivideAndRemasterDeterministicChunkDataset( + data_root, "val", **val_kwargs + ), + test_dataset=DivideAndRemasterDataset(data_root, "test", **test_kwargs), + batch_size=batch_size, + num_workers=num_workers, + **datamodule_kwargs, + ) + + datamodule.predict_dataloader = datamodule.test_dataloader # type: ignore[method-assign] + + return datamodule diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/dataset.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..553c0dd3036491766989d912119e93bb787e63f6 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/dataset.py @@ -0,0 +1,356 @@ +import os +from abc import ABC +from typing import Any, Dict, List, Optional + +import numpy as np +import pedalboard as pb +import torch +import torchaudio as ta +from models.bandit.core.data._types import DataDict +from models.bandit.core.data.base import BaseSourceSeparationDataset + + +class DivideAndRemasterBaseDataset(BaseSourceSeparationDataset, ABC): + ALLOWED_STEMS = ["mixture", "speech", "music", "effects", "mne"] + STEM_NAME_MAP = { + "mixture": "mix", + "speech": "speech", + "music": "music", + "effects": "sfx", + } + SPLIT_NAME_MAP = {"train": "tr", "val": "cv", "test": "tt"} + + FULL_TRACK_LENGTH_SECOND = 60 + FULL_TRACK_LENGTH_SAMPLES = FULL_TRACK_LENGTH_SECOND * 44100 + + def __init__( + self, + split: str, + stems: List[str], + files: List[str], + data_path: str, + fs: int = 44100, + npy_memmap: bool = True, + recompute_mixture: bool = False, + ) -> None: + super().__init__( + split=split, + stems=stems, + files=files, + data_path=data_path, + fs=fs, + npy_memmap=npy_memmap, + recompute_mixture=recompute_mixture, + ) + + def get_stem(self, *, stem: str, identifier: Dict[str, Any]) -> torch.Tensor: + if stem == "mne": + return self.get_stem(stem="music", identifier=identifier) + self.get_stem( + stem="effects", identifier=identifier + ) + + track = identifier["track"] + path = os.path.join(self.data_path, track) + + if self.npy_memmap: + audio = np.load( + os.path.join(path, f"{self.STEM_NAME_MAP[stem]}.npy"), mmap_mode="r" + ) + else: + # noinspection PyUnresolvedReferences + audio, _ = ta.load(os.path.join(path, f"{self.STEM_NAME_MAP[stem]}.wav")) + + return audio + + def get_identifier(self, index): + return dict(track=self.files[index]) + + def __getitem__(self, index: int) -> DataDict: + identifier = self.get_identifier(index) + audio = self.get_audio(identifier) + + return {"audio": audio, "track": f"{self.split}/{identifier['track']}"} + + +class DivideAndRemasterDataset(DivideAndRemasterBaseDataset): + def __init__( + self, + data_root: str, + split: str, + stems: Optional[List[str]] = None, + fs: int = 44100, + npy_memmap: bool = True, + ) -> None: + if stems is None: + stems = self.ALLOWED_STEMS + self.stems = stems + + data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split]) + + files = sorted(os.listdir(data_path)) + files = [ + f + for f in files + if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f)) + ] + # pprint(list(enumerate(files))) + if split == "train": + assert len(files) == 3406, len(files) + elif split == "val": + assert len(files) == 487, len(files) + elif split == "test": + assert len(files) == 973, len(files) + + self.n_tracks = len(files) + + super().__init__( + data_path=data_path, + split=split, + stems=stems, + files=files, + fs=fs, + npy_memmap=npy_memmap, + ) + + def __len__(self) -> int: + return self.n_tracks + + +class DivideAndRemasterRandomChunkDataset(DivideAndRemasterBaseDataset): + def __init__( + self, + data_root: str, + split: str, + target_length: int, + chunk_size_second: float, + stems: Optional[List[str]] = None, + fs: int = 44100, + npy_memmap: bool = True, + ) -> None: + if stems is None: + stems = self.ALLOWED_STEMS + self.stems = stems + + data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split]) + + files = sorted(os.listdir(data_path)) + files = [ + f + for f in files + if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f)) + ] + + if split == "train": + assert len(files) == 3406, len(files) + elif split == "val": + assert len(files) == 487, len(files) + elif split == "test": + assert len(files) == 973, len(files) + + self.n_tracks = len(files) + + self.target_length = target_length + self.chunk_size = int(chunk_size_second * fs) + + super().__init__( + data_path=data_path, + split=split, + stems=stems, + files=files, + fs=fs, + npy_memmap=npy_memmap, + ) + + def __len__(self) -> int: + return self.target_length + + def get_identifier(self, index): + return super().get_identifier(index % self.n_tracks) + + def get_stem( + self, + *, + stem: str, + identifier: Dict[str, Any], + chunk_here: bool = False, + ) -> torch.Tensor: + stem = super().get_stem(stem=stem, identifier=identifier) + + if chunk_here: + start = np.random.randint( + 0, self.FULL_TRACK_LENGTH_SAMPLES - self.chunk_size + ) + end = start + self.chunk_size + + stem = stem[:, start:end] + + return stem + + def __getitem__(self, index: int) -> DataDict: + identifier = self.get_identifier(index) + # self.index_lock = index + audio = self.get_audio(identifier) + # self.index_lock = None + + start = np.random.randint(0, self.FULL_TRACK_LENGTH_SAMPLES - self.chunk_size) + end = start + self.chunk_size + + audio = {k: v[:, start:end] for k, v in audio.items()} + + return {"audio": audio, "track": f"{self.split}/{identifier['track']}"} + + +class DivideAndRemasterDeterministicChunkDataset(DivideAndRemasterBaseDataset): + def __init__( + self, + data_root: str, + split: str, + chunk_size_second: float, + hop_size_second: float, + stems: Optional[List[str]] = None, + fs: int = 44100, + npy_memmap: bool = True, + ) -> None: + if stems is None: + stems = self.ALLOWED_STEMS + self.stems = stems + + data_path = os.path.join(data_root, self.SPLIT_NAME_MAP[split]) + + files = sorted(os.listdir(data_path)) + files = [ + f + for f in files + if (not f.startswith(".")) and os.path.isdir(os.path.join(data_path, f)) + ] + # pprint(list(enumerate(files))) + if split == "train": + assert len(files) == 3406, len(files) + elif split == "val": + assert len(files) == 487, len(files) + elif split == "test": + assert len(files) == 973, len(files) + + self.n_tracks = len(files) + + self.chunk_size = int(chunk_size_second * fs) + self.hop_size = int(hop_size_second * fs) + self.n_chunks_per_track = int( + (self.FULL_TRACK_LENGTH_SECOND - chunk_size_second) / hop_size_second + ) + + self.length = self.n_tracks * self.n_chunks_per_track + + super().__init__( + data_path=data_path, + split=split, + stems=stems, + files=files, + fs=fs, + npy_memmap=npy_memmap, + ) + + def get_identifier(self, index): + return super().get_identifier(index % self.n_tracks) + + def __len__(self) -> int: + return self.length + + def __getitem__(self, item: int) -> DataDict: + index = item % self.n_tracks + chunk = item // self.n_tracks + + data_ = super().__getitem__(index) + + audio = data_["audio"] + + start = chunk * self.hop_size + end = start + self.chunk_size + + for stem in self.stems: + data_["audio"][stem] = audio[stem][:, start:end] + + return data_ + + +class DivideAndRemasterRandomChunkDatasetWithSpeechReverb( + DivideAndRemasterRandomChunkDataset +): + def __init__( + self, + data_root: str, + split: str, + target_length: int, + chunk_size_second: float, + stems: Optional[List[str]] = None, + fs: int = 44100, + npy_memmap: bool = True, + ) -> None: + if stems is None: + stems = self.ALLOWED_STEMS + + stems_no_mixture = [s for s in stems if s != "mixture"] + + super().__init__( + data_root=data_root, + split=split, + target_length=target_length, + chunk_size_second=chunk_size_second, + stems=stems_no_mixture, + fs=fs, + npy_memmap=npy_memmap, + ) + + self.stems = stems + self.stems_no_mixture = stems_no_mixture + + def __getitem__(self, index: int) -> DataDict: + data_ = super().__getitem__(index) + + dry = data_["audio"]["speech"][:] + n_samples = dry.shape[-1] + + wet_level = np.random.rand() + + speech = pb.Reverb( + room_size=np.random.rand(), + damping=np.random.rand(), + wet_level=wet_level, + dry_level=(1 - wet_level), + width=np.random.rand(), + ).process(dry, self.fs, buffer_size=8192 * 4)[..., :n_samples] + + data_["audio"]["speech"] = speech + + data_["audio"]["mixture"] = sum( + [data_["audio"][s] for s in self.stems_no_mixture] + ) + + return data_ + + def __len__(self) -> int: + return super().__len__() + + +if __name__ == "__main__": + from pprint import pprint + + from tqdm.auto import tqdm + + for split_ in ["train", "val", "test"]: + ds = DivideAndRemasterRandomChunkDatasetWithSpeechReverb( + data_root="$DATA_ROOT/DnR/v2np", + split=split_, + target_length=100, + chunk_size_second=6.0, + ) + + print(split_, len(ds)) + + for track_ in tqdm(ds): # type: ignore + pprint(track_) + track_["audio"] = {k: v.shape for k, v in track_["audio"].items()} + pprint(track_) + # break + + break diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/preprocess.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/preprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..18d68b18fbe963647df1253190625ea639035572 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/dnr/preprocess.py @@ -0,0 +1,51 @@ +import glob +import os +from typing import Tuple + +import numpy as np +import torchaudio as ta +from tqdm.contrib.concurrent import process_map + + +def process_one(inputs: Tuple[str, str, int]) -> None: + infile, outfile, target_fs = inputs + + dir = os.path.dirname(outfile) + os.makedirs(dir, exist_ok=True) + + data, fs = ta.load(infile) + + if fs != target_fs: + data = ta.functional.resample( + data, fs, target_fs, resampling_method="sinc_interp_kaiser" + ) + fs = target_fs + + data = data.numpy() + data = data.astype(np.float32) + + if os.path.exists(outfile): + data_ = np.load(outfile) + if np.allclose(data, data_): + return + + np.save(outfile, data) + + +def preprocess(data_path: str, output_path: str, fs: int) -> None: + files = glob.glob(os.path.join(data_path, "**", "*.wav"), recursive=True) + print(files) + outfiles = [ + f.replace(data_path, output_path).replace(".wav", ".npy") for f in files + ] + + os.makedirs(output_path, exist_ok=True) + inputs = list(zip(files, outfiles, [fs] * len(files))) + + process_map(process_one, inputs, chunksize=32) + + +if __name__ == "__main__": + import fire + + fire.Fire() diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/datamodule.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/datamodule.py new file mode 100644 index 0000000000000000000000000000000000000000..bb97c5a8c5c8828ea49ca0b3f5a0e4494d606f49 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/datamodule.py @@ -0,0 +1,74 @@ +import os.path +from typing import Mapping, Optional + +import pytorch_lightning as pl +from models.bandit.core.data.musdb.dataset import ( + MUSDB18BaseDataset, + MUSDB18FullTrackDataset, + MUSDB18SadDataset, + MUSDB18SadOnTheFlyAugmentedDataset, +) + + +def MUSDB18DataModule( + data_root: str = "$DATA_ROOT/MUSDB18/HQ", + target_stem: str = "vocals", + batch_size: int = 2, + num_workers: int = 8, + train_kwargs: Optional[Mapping] = None, + val_kwargs: Optional[Mapping] = None, + test_kwargs: Optional[Mapping] = None, + datamodule_kwargs: Optional[Mapping] = None, + use_on_the_fly: bool = True, + npy_memmap: bool = True, +) -> pl.LightningDataModule: + if train_kwargs is None: + train_kwargs = {} + + if val_kwargs is None: + val_kwargs = {} + + if test_kwargs is None: + test_kwargs = {} + + if datamodule_kwargs is None: + datamodule_kwargs = {} + + train_dataset: MUSDB18BaseDataset + + if use_on_the_fly: + train_dataset = MUSDB18SadOnTheFlyAugmentedDataset( + data_root=os.path.join(data_root, "saded-np"), + split="train", + target_stem=target_stem, + **train_kwargs, + ) + else: + train_dataset = MUSDB18SadDataset( + data_root=os.path.join(data_root, "saded-np"), + split="train", + target_stem=target_stem, + **train_kwargs, + ) + + datamodule = pl.LightningDataModule.from_datasets( + train_dataset=train_dataset, + val_dataset=MUSDB18SadDataset( + data_root=os.path.join(data_root, "saded-np"), + split="val", + target_stem=target_stem, + **val_kwargs, + ), + test_dataset=MUSDB18FullTrackDataset( + data_root=os.path.join(data_root, "canonical"), split="test", **test_kwargs + ), + batch_size=batch_size, + num_workers=num_workers, + **datamodule_kwargs, + ) + + datamodule.predict_dataloader = ( # type: ignore[method-assign] + datamodule.test_dataloader + ) + + return datamodule diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/dataset.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..e0db657dfd858546dda85090b6f2d65516cb3e93 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/dataset.py @@ -0,0 +1,265 @@ +import os +from abc import ABC +from typing import List, Optional, Tuple + +import numpy as np +import torch +import torchaudio as ta +from models.bandit.core.data._types import DataDict +from models.bandit.core.data.base import BaseSourceSeparationDataset + + +class MUSDB18BaseDataset(BaseSourceSeparationDataset, ABC): + ALLOWED_STEMS = ["mixture", "vocals", "bass", "drums", "other"] + + def __init__( + self, + split: str, + stems: List[str], + files: List[str], + data_path: str, + fs: int = 44100, + npy_memmap=False, + ) -> None: + super().__init__( + split=split, + stems=stems, + files=files, + data_path=data_path, + fs=fs, + npy_memmap=npy_memmap, + recompute_mixture=False, + ) + + def get_stem(self, *, stem: str, identifier) -> torch.Tensor: + track = identifier["track"] + path = os.path.join(self.data_path, track) + # noinspection PyUnresolvedReferences + + if self.npy_memmap: + audio = np.load(os.path.join(path, f"{stem}.wav.npy"), mmap_mode="r") + else: + audio, _ = ta.load(os.path.join(path, f"{stem}.wav")) + + return audio + + def get_identifier(self, index): + return dict(track=self.files[index]) + + def __getitem__(self, index: int) -> DataDict: + identifier = self.get_identifier(index) + audio = self.get_audio(identifier) + + return {"audio": audio, "track": f"{self.split}/{identifier['track']}"} + + +class MUSDB18FullTrackDataset(MUSDB18BaseDataset): + N_TRAIN_TRACKS = 100 + N_TEST_TRACKS = 50 + VALIDATION_FILES = [ + "Actions - One Minute Smile", + "Clara Berry And Wooldog - Waltz For My Victims", + "Johnny Lokke - Promises & Lies", + "Patrick Talbot - A Reason To Leave", + "Triviul - Angelsaint", + "Alexander Ross - Goodbye Bolero", + "Fergessen - Nos Palpitants", + "Leaf - Summerghost", + "Skelpolu - Human Mistakes", + "Young Griffo - Pennies", + "ANiMAL - Rockshow", + "James May - On The Line", + "Meaxic - Take A Step", + "Traffic Experiment - Sirens", + ] + + def __init__( + self, data_root: str, split: str, stems: Optional[List[str]] = None + ) -> None: + if stems is None: + stems = self.ALLOWED_STEMS + self.stems = stems + + if split == "test": + subset = "test" + elif split in ["train", "val"]: + subset = "train" + else: + raise NameError + + data_path = os.path.join(data_root, subset) + + files = sorted(os.listdir(data_path)) + files = [f for f in files if not f.startswith(".")] + # pprint(list(enumerate(files))) + if subset == "train": + assert len(files) == 100, len(files) + if split == "train": + files = [f for f in files if f not in self.VALIDATION_FILES] + assert len(files) == 100 - len(self.VALIDATION_FILES) + else: + files = [f for f in files if f in self.VALIDATION_FILES] + assert len(files) == len(self.VALIDATION_FILES) + else: + split = "test" + assert len(files) == 50 + + self.n_tracks = len(files) + + super().__init__(data_path=data_path, split=split, stems=stems, files=files) + + def __len__(self) -> int: + return self.n_tracks + + +class MUSDB18SadDataset(MUSDB18BaseDataset): + def __init__( + self, + data_root: str, + split: str, + target_stem: str, + stems: Optional[List[str]] = None, + target_length: Optional[int] = None, + npy_memmap=False, + ) -> None: + if stems is None: + stems = self.ALLOWED_STEMS + + data_path = os.path.join(data_root, target_stem, split) + + files = sorted(os.listdir(data_path)) + files = [f for f in files if not f.startswith(".")] + + super().__init__( + data_path=data_path, + split=split, + stems=stems, + files=files, + npy_memmap=npy_memmap, + ) + self.n_segments = len(files) + self.target_stem = target_stem + self.target_length = ( + target_length if target_length is not None else self.n_segments + ) + + def __len__(self) -> int: + return self.target_length + + def __getitem__(self, index: int) -> DataDict: + index = index % self.n_segments + + return super().__getitem__(index) + + def get_identifier(self, index): + return super().get_identifier(index % self.n_segments) + + +class MUSDB18SadOnTheFlyAugmentedDataset(MUSDB18SadDataset): + def __init__( + self, + data_root: str, + split: str, + target_stem: str, + stems: Optional[List[str]] = None, + target_length: int = 20000, + apply_probability: Optional[float] = None, + chunk_size_second: float = 3.0, + random_scale_range_db: Tuple[float, float] = (-10, 10), + drop_probability: float = 0.1, + rescale: bool = True, + ) -> None: + super().__init__(data_root, split, target_stem, stems) + + if apply_probability is None: + apply_probability = (target_length - self.n_segments) / target_length + + self.apply_probability = apply_probability + self.drop_probability = drop_probability + self.chunk_size_second = chunk_size_second + self.random_scale_range_db = random_scale_range_db + self.rescale = rescale + + self.chunk_size_sample = int(self.chunk_size_second * self.fs) + self.target_length = target_length + + def __len__(self) -> int: + return self.target_length + + def __getitem__(self, index: int) -> DataDict: + index = index % self.n_segments + + # if np.random.rand() > self.apply_probability: + # return super().__getitem__(index) + + audio = {} + identifier = self.get_identifier(index) + + # assert self.target_stem in self.stems_no_mixture + for stem in self.stems_no_mixture: + if stem == self.target_stem: + identifier_ = identifier + else: + if np.random.rand() < self.apply_probability: + index_ = np.random.randint(self.n_segments) + identifier_ = self.get_identifier(index_) + else: + identifier_ = identifier + + audio[stem] = self.get_stem(stem=stem, identifier=identifier_) + + # if stem == self.target_stem: + + if self.chunk_size_sample < audio[stem].shape[-1]: + chunk_start = np.random.randint( + audio[stem].shape[-1] - self.chunk_size_sample + ) + else: + chunk_start = 0 + + if np.random.rand() < self.drop_probability: + # db_scale = "-inf" + linear_scale = 0.0 + else: + db_scale = np.random.uniform(*self.random_scale_range_db) + linear_scale = np.power(10, db_scale / 20) + # db_scale = f"{db_scale:+2.1f}" + # print(linear_scale) + audio[stem][..., chunk_start : chunk_start + self.chunk_size_sample] = ( + linear_scale + * audio[stem][..., chunk_start : chunk_start + self.chunk_size_sample] + ) + + audio["mixture"] = self.compute_mixture(audio) + + if self.rescale: + max_abs_val = max( + [torch.max(torch.abs(audio[stem])) for stem in self.stems] + ) # type: ignore[type-var] + if max_abs_val > 1: + audio = {k: v / max_abs_val for k, v in audio.items()} + + track = identifier["track"] + + return {"audio": audio, "track": f"{self.split}/{track}"} + + +# if __name__ == "__main__": +# +# from pprint import pprint +# from tqdm.auto import tqdm +# +# for split_ in ["train", "val", "test"]: +# ds = MUSDB18SadOnTheFlyAugmentedDataset( +# data_root="$DATA_ROOT/MUSDB18/HQ/saded", +# split=split_, +# target_stem="vocals" +# ) +# +# print(split_, len(ds)) +# +# for track_ in tqdm(ds): +# track_["audio"] = { +# k: v.shape for k, v in track_["audio"].items() +# } +# pprint(track_) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/preprocess.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/preprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..5a52371223c566d57ef8b41fe0b628d4e0468a32 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/preprocess.py @@ -0,0 +1,223 @@ +import glob +import os + +import numpy as np +import pyloudnorm as pyln +import torch +import torchaudio as ta +from torch import nn +from torch.nn import functional as F +from tqdm.contrib.concurrent import process_map + +from core.data._types import DataDict +from core.data.musdb.dataset import MUSDB18FullTrackDataset + + +class SourceActivityDetector(nn.Module): + def __init__( + self, + analysis_stem: str, + output_path: str, + fs: int = 44100, + segment_length_second: float = 6.0, + hop_length_second: float = 3.0, + n_chunks: int = 10, + chunk_epsilon: float = 1e-5, + energy_threshold_quantile: float = 0.15, + segment_epsilon: float = 1e-3, + salient_proportion_threshold: float = 0.5, + target_lufs: float = -24, + ) -> None: + super().__init__() + + self.fs = fs + self.segment_length = int(segment_length_second * self.fs) + self.hop_length = int(hop_length_second * self.fs) + self.n_chunks = n_chunks + assert self.segment_length % self.n_chunks == 0 + self.chunk_size = self.segment_length // self.n_chunks + self.chunk_epsilon = chunk_epsilon + self.energy_threshold_quantile = energy_threshold_quantile + self.segment_epsilon = segment_epsilon + self.salient_proportion_threshold = salient_proportion_threshold + self.analysis_stem = analysis_stem + + self.meter = pyln.Meter(self.fs) + self.target_lufs = target_lufs + + self.output_path = output_path + + def forward(self, data: DataDict) -> None: + stem_ = self.analysis_stem if (self.analysis_stem != "none") else "mixture" + + x = data["audio"][stem_] + + xnp = x.numpy() + loudness = self.meter.integrated_loudness(xnp.T) + + for stem in data["audio"]: + s = data["audio"][stem] + s = pyln.normalize.loudness(s.numpy().T, loudness, self.target_lufs).T + s = torch.as_tensor(s) + data["audio"][stem] = s + + if x.ndim == 3: + assert x.shape[0] == 1 + x = x[0] + + n_chan, n_samples = x.shape + + n_segments = ( + int(np.ceil((n_samples - self.segment_length) / self.hop_length)) + 1 + ) + + segments = torch.zeros((n_segments, n_chan, self.segment_length)) + for i in range(n_segments): + start = i * self.hop_length + end = start + self.segment_length + end = min(end, n_samples) + + xseg = x[:, start:end] + + if end - start < self.segment_length: + xseg = F.pad( + xseg, pad=(0, self.segment_length - (end - start)), value=torch.nan + ) + + segments[i, :, :] = xseg + + chunks = segments.reshape((n_segments, n_chan, self.n_chunks, self.chunk_size)) + + if self.analysis_stem != "none": + chunk_energies = torch.mean(torch.square(chunks), dim=(1, 3)) + chunk_energies = torch.nan_to_num(chunk_energies, nan=0) + chunk_energies[chunk_energies == 0] = self.chunk_epsilon + + energy_threshold = torch.nanquantile( + chunk_energies, q=self.energy_threshold_quantile + ) + + if energy_threshold < self.segment_epsilon: + energy_threshold = self.segment_epsilon # type: ignore[assignment] + + chunks_above_threshold = chunk_energies > energy_threshold + n_chunks_above_threshold = torch.mean( + chunks_above_threshold.to(torch.float), dim=-1 + ) + + segment_above_threshold = ( + n_chunks_above_threshold > self.salient_proportion_threshold + ) + + if torch.sum(segment_above_threshold) == 0: + return + + else: + segment_above_threshold = torch.ones((n_segments,)) + + for i in range(n_segments): + if not segment_above_threshold[i]: + continue + + outpath = os.path.join( + self.output_path, + self.analysis_stem, + f"{data['track']} - {self.analysis_stem}{i:03d}", + ) + os.makedirs(outpath, exist_ok=True) + + for stem in data["audio"]: + if stem == self.analysis_stem: + segment = torch.nan_to_num(segments[i, :, :], nan=0) + else: + start = i * self.hop_length + end = start + self.segment_length + end = min(n_samples, end) + + segment = data["audio"][stem][:, start:end] + + if end - start < self.segment_length: + segment = F.pad( + segment, (0, self.segment_length - (end - start)) + ) + + assert segment.shape[-1] == self.segment_length, segment.shape + + # ta.save(os.path.join(outpath, f"{stem}.wav"), segment, self.fs) + + np.save(os.path.join(outpath, f"{stem}.wav"), segment) + + +def preprocess( + analysis_stem: str, + output_path: str = "/data/MUSDB18/HQ/saded-np", + fs: int = 44100, + segment_length_second: float = 6.0, + hop_length_second: float = 3.0, + n_chunks: int = 10, + chunk_epsilon: float = 1e-5, + energy_threshold_quantile: float = 0.15, + segment_epsilon: float = 1e-3, + salient_proportion_threshold: float = 0.5, +) -> None: + sad = SourceActivityDetector( + analysis_stem=analysis_stem, + output_path=output_path, + fs=fs, + segment_length_second=segment_length_second, + hop_length_second=hop_length_second, + n_chunks=n_chunks, + chunk_epsilon=chunk_epsilon, + energy_threshold_quantile=energy_threshold_quantile, + segment_epsilon=segment_epsilon, + salient_proportion_threshold=salient_proportion_threshold, + ) + + for split in ["train", "val", "test"]: + ds = MUSDB18FullTrackDataset( + data_root="/data/MUSDB18/HQ/canonical", + split=split, + ) + + tracks = [] + for i, track in enumerate(tqdm(ds, total=len(ds))): + if i % 32 == 0 and tracks: + process_map(sad, tracks, max_workers=8) + tracks = [] + tracks.append(track) + process_map(sad, tracks, max_workers=8) + + +def loudness_norm_one(inputs): + infile, outfile, target_lufs = inputs + + audio, fs = ta.load(infile) + audio = audio.mean(dim=0, keepdim=True).numpy().T + + meter = pyln.Meter(fs) + loudness = meter.integrated_loudness(audio) + audio = pyln.normalize.loudness(audio, loudness, target_lufs) + + os.makedirs(os.path.dirname(outfile), exist_ok=True) + np.save(outfile, audio.T) + + +def loudness_norm( + data_path: str, + # output_path: str, + target_lufs=-17.0, +): + files = glob.glob(os.path.join(data_path, "**", "*.wav"), recursive=True) + + outfiles = [f.replace(".wav", ".npy").replace("saded", "saded-np") for f in files] + + files = [(f, o, target_lufs) for f, o in zip(files, outfiles)] + + process_map(loudness_norm_one, files, chunksize=2) + + +if __name__ == "__main__": + import fire + from tqdm.auto import tqdm + + fire.Fire() diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/validation.yaml b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/validation.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2f8752478d285d1d13d5e842225af1de95cae57a --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/data/musdb/validation.yaml @@ -0,0 +1,15 @@ +validation: + - 'Actions - One Minute Smile' + - 'Clara Berry And Wooldog - Waltz For My Victims' + - 'Johnny Lokke - Promises & Lies' + - 'Patrick Talbot - A Reason To Leave' + - 'Triviul - Angelsaint' + - 'Alexander Ross - Goodbye Bolero' + - 'Fergessen - Nos Palpitants' + - 'Leaf - Summerghost' + - 'Skelpolu - Human Mistakes' + - 'Young Griffo - Pennies' + - 'ANiMAL - Rockshow' + - 'James May - On The Line' + - 'Meaxic - Take A Step' + - 'Traffic Experiment - Sirens' \ No newline at end of file diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..993be521fa7ab8f06a2a012beabdb9fdd6cd0a80 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/__init__.py @@ -0,0 +1,8 @@ +from ._multistem import MultiStemWrapperFromConfig +from ._timefreq import ( + ReImL1Loss, + ReImL2Loss, + TimeFreqL1Loss, + TimeFreqL2Loss, + TimeFreqSignalNoisePNormRatioLoss, +) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_complex.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_complex.py new file mode 100644 index 0000000000000000000000000000000000000000..d3a20e478230fa614f9eb2247900425b5472888f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_complex.py @@ -0,0 +1,26 @@ +from typing import Any + +import torch +from torch.nn.modules import loss as _loss +from torch.nn.modules.loss import _Loss + + +class ReImLossWrapper(_Loss): + def __init__(self, module: _Loss) -> None: + super().__init__() + self.module = module + + def forward(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + return self.module(torch.view_as_real(preds), torch.view_as_real(target)) + + +class ReImL1Loss(ReImLossWrapper): + def __init__(self, **kwargs: Any) -> None: + l1_loss = _loss.L1Loss(**kwargs) + super().__init__(module=(l1_loss)) + + +class ReImL2Loss(ReImLossWrapper): + def __init__(self, **kwargs: Any) -> None: + l2_loss = _loss.MSELoss(**kwargs) + super().__init__(module=(l2_loss)) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_multistem.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_multistem.py new file mode 100644 index 0000000000000000000000000000000000000000..5b283919bfac33c1ac8f071d3e85a3ec8610cf72 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_multistem.py @@ -0,0 +1,42 @@ +from typing import Any, Dict + +import torch +from asteroid import losses as asteroid_losses +from torch import nn +from torch.nn.modules.loss import _Loss + +from . import snr + + +def parse_loss(name: str, kwargs: Dict[str, Any]) -> _Loss: + for module in [nn.modules.loss, snr, asteroid_losses, asteroid_losses.sdr]: + if name in module.__dict__: + return module.__dict__[name](**kwargs) + + raise NameError + + +class MultiStemWrapper(_Loss): + def __init__(self, module: _Loss, modality: str = "audio") -> None: + super().__init__() + self.loss = module + self.modality = modality + + def forward( + self, + preds: Dict[str, Dict[str, torch.Tensor]], + target: Dict[str, Dict[str, torch.Tensor]], + ) -> torch.Tensor: + loss = { + stem: self.loss(preds[self.modality][stem], target[self.modality][stem]) + for stem in preds[self.modality] + if stem in target[self.modality] + } + + return sum(list(loss.values())) + + +class MultiStemWrapperFromConfig(MultiStemWrapper): + def __init__(self, name: str, kwargs: Any, modality: str = "audio") -> None: + loss = parse_loss(name, kwargs) + super().__init__(module=loss, modality=modality) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_timefreq.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_timefreq.py new file mode 100644 index 0000000000000000000000000000000000000000..c9f7874027c12e000f456b1b668a8fb86004c33d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/_timefreq.py @@ -0,0 +1,93 @@ +from typing import Any, Dict, Optional + +import torch +from models.bandit.core.loss._complex import ReImL1Loss, ReImL2Loss +from models.bandit.core.loss._multistem import MultiStemWrapper +from models.bandit.core.loss.snr import SignalNoisePNormRatio +from torch import nn +from torch.nn.modules.loss import _Loss + + +class TimeFreqWrapper(_Loss): + def __init__( + self, + time_module: _Loss, + freq_module: Optional[_Loss] = None, + time_weight: float = 1.0, + freq_weight: float = 1.0, + multistem: bool = True, + ) -> None: + super().__init__() + + if freq_module is None: + freq_module = time_module + + if multistem: + time_module = MultiStemWrapper(time_module, modality="audio") + freq_module = MultiStemWrapper(freq_module, modality="spectrogram") + + self.time_module = time_module + self.freq_module = freq_module + + self.time_weight = time_weight + self.freq_weight = freq_weight + + # TODO: add better type hints + def forward(self, preds: Any, target: Any) -> torch.Tensor: + return self.time_weight * self.time_module( + preds, target + ) + self.freq_weight * self.freq_module(preds, target) + + +class TimeFreqL1Loss(TimeFreqWrapper): + def __init__( + self, + time_weight: float = 1.0, + freq_weight: float = 1.0, + tkwargs: Optional[Dict[str, Any]] = None, + fkwargs: Optional[Dict[str, Any]] = None, + multistem: bool = True, + ) -> None: + if tkwargs is None: + tkwargs = {} + if fkwargs is None: + fkwargs = {} + time_module = nn.L1Loss(**tkwargs) + freq_module = ReImL1Loss(**fkwargs) + super().__init__(time_module, freq_module, time_weight, freq_weight, multistem) + + +class TimeFreqL2Loss(TimeFreqWrapper): + def __init__( + self, + time_weight: float = 1.0, + freq_weight: float = 1.0, + tkwargs: Optional[Dict[str, Any]] = None, + fkwargs: Optional[Dict[str, Any]] = None, + multistem: bool = True, + ) -> None: + if tkwargs is None: + tkwargs = {} + if fkwargs is None: + fkwargs = {} + time_module = nn.MSELoss(**tkwargs) + freq_module = ReImL2Loss(**fkwargs) + super().__init__(time_module, freq_module, time_weight, freq_weight, multistem) + + +class TimeFreqSignalNoisePNormRatioLoss(TimeFreqWrapper): + def __init__( + self, + time_weight: float = 1.0, + freq_weight: float = 1.0, + tkwargs: Optional[Dict[str, Any]] = None, + fkwargs: Optional[Dict[str, Any]] = None, + multistem: bool = True, + ) -> None: + if tkwargs is None: + tkwargs = {} + if fkwargs is None: + fkwargs = {} + time_module = SignalNoisePNormRatio(**tkwargs) + freq_module = SignalNoisePNormRatio(**fkwargs) + super().__init__(time_module, freq_module, time_weight, freq_weight, multistem) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/snr.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/snr.py new file mode 100644 index 0000000000000000000000000000000000000000..17d76792151237cb97a95625e903fa1370595db0 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/loss/snr.py @@ -0,0 +1,137 @@ +import torch +from torch.nn.modules.loss import _Loss + + +class SignalNoisePNormRatio(_Loss): + def __init__( + self, + p: float = 1.0, + scale_invariant: bool = False, + zero_mean: bool = False, + take_log: bool = True, + reduction: str = "mean", + EPS: float = 1e-3, + ) -> None: + assert reduction != "sum", NotImplementedError + super().__init__(reduction=reduction) + assert not zero_mean + + self.p = p + + self.EPS = EPS + self.take_log = take_log + + self.scale_invariant = scale_invariant + + def forward(self, est_target: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + target_ = target + if self.scale_invariant: + ndim = target.ndim + dot = torch.sum(est_target * torch.conj(target), dim=-1, keepdim=True) + s_target_energy = torch.sum( + target * torch.conj(target), dim=-1, keepdim=True + ) + + if ndim > 2: + dot = torch.sum(dot, dim=list(range(1, ndim)), keepdim=True) + s_target_energy = torch.sum( + s_target_energy, dim=list(range(1, ndim)), keepdim=True + ) + + target_scaler = (dot + 1e-8) / (s_target_energy + 1e-8) + target = target_ * target_scaler + + if torch.is_complex(est_target): + est_target = torch.view_as_real(est_target) + target = torch.view_as_real(target) + + batch_size = est_target.shape[0] + est_target = est_target.reshape(batch_size, -1) + target = target.reshape(batch_size, -1) + # target_ = target_.reshape(batch_size, -1) + + if self.p == 1: + e_error = torch.abs(est_target - target).mean(dim=-1) + e_target = torch.abs(target).mean(dim=-1) + elif self.p == 2: + e_error = torch.square(est_target - target).mean(dim=-1) + e_target = torch.square(target).mean(dim=-1) + else: + raise NotImplementedError + + if self.take_log: + loss = 10 * ( + torch.log10(e_error + self.EPS) - torch.log10(e_target + self.EPS) + ) + else: + loss = (e_error + self.EPS) / (e_target + self.EPS) + + if self.reduction == "mean": + loss = loss.mean() + elif self.reduction == "sum": + loss = loss.sum() + + return loss + + +class MultichannelSingleSrcNegSDR(_Loss): + def __init__( + self, + sdr_type: str, + p: float = 2.0, + zero_mean: bool = True, + take_log: bool = True, + reduction: str = "mean", + EPS: float = 1e-8, + ) -> None: + assert reduction != "sum", NotImplementedError + super().__init__(reduction=reduction) + + assert sdr_type in ["snr", "sisdr", "sdsdr"] + self.sdr_type = sdr_type + self.zero_mean = zero_mean + self.take_log = take_log + self.EPS = 1e-8 + + self.p = p + + def forward(self, est_target: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + if target.size() != est_target.size() or target.ndim != 3: + raise TypeError( + f"Inputs must be of shape [batch, time], got {target.size()} and {est_target.size()} instead" + ) + # Step 1. Zero-mean norm + if self.zero_mean: + mean_source = torch.mean(target, dim=[1, 2], keepdim=True) + mean_estimate = torch.mean(est_target, dim=[1, 2], keepdim=True) + target = target - mean_source + est_target = est_target - mean_estimate + # Step 2. Pair-wise SI-SDR. + if self.sdr_type in ["sisdr", "sdsdr"]: + # [batch, 1] + dot = torch.sum(est_target * target, dim=[1, 2], keepdim=True) + # [batch, 1] + s_target_energy = torch.sum(target**2, dim=[1, 2], keepdim=True) + self.EPS + # [batch, time] + scaled_target = dot * target / s_target_energy + else: + # [batch, time] + scaled_target = target + if self.sdr_type in ["sdsdr", "snr"]: + e_noise = est_target - target + else: + e_noise = est_target - scaled_target + # [batch] + + if self.p == 2.0: + losses = torch.sum(scaled_target**2, dim=[1, 2]) / ( + torch.sum(e_noise**2, dim=[1, 2]) + self.EPS + ) + else: + losses = torch.norm(scaled_target, p=self.p, dim=[1, 2]) / ( + torch.linalg.vector_norm(e_noise, p=self.p, dim=[1, 2]) + self.EPS + ) + if self.take_log: + losses = 10 * torch.log10(losses + self.EPS) + losses = losses.mean() if self.reduction == "mean" else losses + return -losses diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c638b4df585ad6c3c6490d9e67b7fc197f0d06f4 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/__init__.py @@ -0,0 +1,9 @@ +from .snr import ( + ChunkMedianScaleInvariantSignalDistortionRatio, + ChunkMedianScaleInvariantSignalNoiseRatio, + ChunkMedianSignalDistortionRatio, + ChunkMedianSignalNoiseRatio, + SafeSignalDistortionRatio, +) + +# from .mushra import EstimatedMushraScore diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/_squim.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/_squim.py new file mode 100644 index 0000000000000000000000000000000000000000..056c78099798acd84c68fd4e986d82f277908ff5 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/_squim.py @@ -0,0 +1,440 @@ +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchaudio._internal import load_state_dict_from_url + + +def transform_wb_pesq_range(x: float) -> float: + """The metric defined by ITU-T P.862 is often called 'PESQ score', which is defined + for narrow-band signals and has a value range of [-0.5, 4.5] exactly. Here, we use the metric + defined by ITU-T P.862.2, commonly known as 'wide-band PESQ' and will be referred to as "PESQ score". + + Args: + x (float): Narrow-band PESQ score. + + Returns: + (float): Wide-band PESQ score. + """ + return 0.999 + (4.999 - 0.999) / (1 + math.exp(-1.3669 * x + 3.8224)) + + +PESQRange: Tuple[float, float] = ( + 1.0, # P.862.2 uses a different input filter than P.862, and the lower bound of + # the raw score is not -0.5 anymore. It's hard to figure out the true lower bound. + # We are using 1.0 as a reasonable approximation. + transform_wb_pesq_range(4.5), +) + + +class RangeSigmoid(nn.Module): + def __init__(self, val_range: Tuple[float, float] = (0.0, 1.0)) -> None: + super(RangeSigmoid, self).__init__() + assert isinstance(val_range, tuple) and len(val_range) == 2 + self.val_range: Tuple[float, float] = val_range + self.sigmoid: nn.modules.Module = nn.Sigmoid() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = ( + self.sigmoid(x) * (self.val_range[1] - self.val_range[0]) + + self.val_range[0] + ) + return out + + +class Encoder(nn.Module): + """Encoder module that transform 1D waveform to 2D representations. + + Args: + feat_dim (int, optional): The feature dimension after Encoder module. (Default: 512) + win_len (int, optional): kernel size in the Conv1D layer. (Default: 32) + """ + + def __init__(self, feat_dim: int = 512, win_len: int = 32) -> None: + super(Encoder, self).__init__() + + self.conv1d = nn.Conv1d(1, feat_dim, win_len, stride=win_len // 2, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply waveforms to convolutional layer and ReLU layer. + + Args: + x (torch.Tensor): Input waveforms. Tensor with dimensions `(batch, time)`. + + Returns: + (torch,Tensor): Feature Tensor with dimensions `(batch, channel, frame)`. + """ + out = x.unsqueeze(dim=1) + out = F.relu(self.conv1d(out)) + return out + + +class SingleRNN(nn.Module): + def __init__( + self, rnn_type: str, input_size: int, hidden_size: int, dropout: float = 0.0 + ) -> None: + super(SingleRNN, self).__init__() + + self.rnn_type = rnn_type + self.input_size = input_size + self.hidden_size = hidden_size + + self.rnn: nn.modules.Module = getattr(nn, rnn_type)( + input_size, + hidden_size, + 1, + dropout=dropout, + batch_first=True, + bidirectional=True, + ) + + self.proj = nn.Linear(hidden_size * 2, input_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # input shape: batch, seq, dim + out, _ = self.rnn(x) + out = self.proj(out) + return out + + +class DPRNN(nn.Module): + """*Dual-path recurrent neural networks (DPRNN)* :cite:`luo2020dual`. + + Args: + feat_dim (int, optional): The feature dimension after Encoder module. (Default: 64) + hidden_dim (int, optional): Hidden dimension in the RNN layer of DPRNN. (Default: 128) + num_blocks (int, optional): Number of DPRNN layers. (Default: 6) + rnn_type (str, optional): Type of RNN in DPRNN. Valid options are ["RNN", "LSTM", "GRU"]. (Default: "LSTM") + d_model (int, optional): The number of expected features in the input. (Default: 256) + chunk_size (int, optional): Chunk size of input for DPRNN. (Default: 100) + chunk_stride (int, optional): Stride of chunk input for DPRNN. (Default: 50) + """ + + def __init__( + self, + feat_dim: int = 64, + hidden_dim: int = 128, + num_blocks: int = 6, + rnn_type: str = "LSTM", + d_model: int = 256, + chunk_size: int = 100, + chunk_stride: int = 50, + ) -> None: + super(DPRNN, self).__init__() + + self.num_blocks = num_blocks + + self.row_rnn = nn.ModuleList([]) + self.col_rnn = nn.ModuleList([]) + self.row_norm = nn.ModuleList([]) + self.col_norm = nn.ModuleList([]) + for _ in range(num_blocks): + self.row_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim)) + self.col_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim)) + self.row_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8)) + self.col_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8)) + self.conv = nn.Sequential( + nn.Conv2d(feat_dim, d_model, 1), + nn.PReLU(), + ) + self.chunk_size = chunk_size + self.chunk_stride = chunk_stride + + def pad_chunk(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]: + # input shape: (B, N, T) + seq_len = x.shape[-1] + + rest = ( + self.chunk_size + - (self.chunk_stride + seq_len % self.chunk_size) % self.chunk_size + ) + out = F.pad(x, [self.chunk_stride, rest + self.chunk_stride]) + + return out, rest + + def chunking(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]: + out, rest = self.pad_chunk(x) + batch_size, feat_dim, seq_len = out.shape + + segments1 = ( + out[:, :, : -self.chunk_stride] + .contiguous() + .view(batch_size, feat_dim, -1, self.chunk_size) + ) + segments2 = ( + out[:, :, self.chunk_stride :] + .contiguous() + .view(batch_size, feat_dim, -1, self.chunk_size) + ) + out = torch.cat([segments1, segments2], dim=3) + out = ( + out.view(batch_size, feat_dim, -1, self.chunk_size) + .transpose(2, 3) + .contiguous() + ) + + return out, rest + + def merging(self, x: torch.Tensor, rest: int) -> torch.Tensor: + batch_size, dim, _, _ = x.shape + out = ( + x.transpose(2, 3) + .contiguous() + .view(batch_size, dim, -1, self.chunk_size * 2) + ) + out1 = ( + out[:, :, :, : self.chunk_size] + .contiguous() + .view(batch_size, dim, -1)[:, :, self.chunk_stride :] + ) + out2 = ( + out[:, :, :, self.chunk_size :] + .contiguous() + .view(batch_size, dim, -1)[:, :, : -self.chunk_stride] + ) + out = out1 + out2 + if rest > 0: + out = out[:, :, :-rest] + out = out.contiguous() + return out + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, rest = self.chunking(x) + batch_size, _, dim1, dim2 = x.shape + out = x + for row_rnn, row_norm, col_rnn, col_norm in zip( + self.row_rnn, self.row_norm, self.col_rnn, self.col_norm + ): + row_in = ( + out.permute(0, 3, 2, 1) + .contiguous() + .view(batch_size * dim2, dim1, -1) + .contiguous() + ) + row_out = row_rnn(row_in) + row_out = ( + row_out.view(batch_size, dim2, dim1, -1) + .permute(0, 3, 2, 1) + .contiguous() + ) + row_out = row_norm(row_out) + out = out + row_out + + col_in = ( + out.permute(0, 2, 3, 1) + .contiguous() + .view(batch_size * dim1, dim2, -1) + .contiguous() + ) + col_out = col_rnn(col_in) + col_out = ( + col_out.view(batch_size, dim1, dim2, -1) + .permute(0, 3, 1, 2) + .contiguous() + ) + col_out = col_norm(col_out) + out = out + col_out + out = self.conv(out) + out = self.merging(out, rest) + out = out.transpose(1, 2).contiguous() + return out + + +class AutoPool(nn.Module): + def __init__(self, pool_dim: int = 1) -> None: + super(AutoPool, self).__init__() + self.pool_dim: int = pool_dim + self.softmax: nn.modules.Module = nn.Softmax(dim=pool_dim) + self.register_parameter("alpha", nn.Parameter(torch.ones(1))) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + weight = self.softmax(torch.mul(x, self.alpha)) + out = torch.sum(torch.mul(x, weight), dim=self.pool_dim) + return out + + +class SquimObjective(nn.Module): + """Speech Quality and Intelligibility Measures (SQUIM) model that predicts **objective** metric scores + for speech enhancement (e.g., STOI, PESQ, and SI-SDR). + + Args: + encoder (torch.nn.Module): Encoder module to transform 1D waveform to 2D feature representation. + dprnn (torch.nn.Module): DPRNN module to model sequential feature. + branches (torch.nn.ModuleList): Transformer branches in which each branch estimate one objective metirc score. + """ + + def __init__( + self, + encoder: nn.Module, + dprnn: nn.Module, + branches: nn.ModuleList, + ): + super(SquimObjective, self).__init__() + self.encoder = encoder + self.dprnn = dprnn + self.branches = branches + + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + """ + Args: + x (torch.Tensor): Input waveforms. Tensor with dimensions `(batch, time)`. + + Returns: + List(torch.Tensor): List of score Tenosrs. Each Tensor is with dimension `(batch,)`. + """ + if x.ndim != 2: + raise ValueError( + f"The input must be a 2D Tensor. Found dimension {x.ndim}." + ) + x = x / (torch.mean(x**2, dim=1, keepdim=True) ** 0.5 * 20) + out = self.encoder(x) + out = self.dprnn(out) + scores = [] + for branch in self.branches: + scores.append(branch(out).squeeze(dim=1)) + return scores + + +def _create_branch(d_model: int, nhead: int, metric: str) -> nn.modules.Module: + """Create branch module after DPRNN model for predicting metric score. + + Args: + d_model (int): The number of expected features in the input. + nhead (int): Number of heads in the multi-head attention model. + metric (str): The metric name to predict. + + Returns: + (nn.Module): Returned module to predict corresponding metric score. + """ + layer1 = nn.TransformerEncoderLayer( + d_model, nhead, d_model * 4, dropout=0.0, batch_first=True + ) + layer2 = AutoPool() + if metric == "stoi": + layer3 = nn.Sequential( + nn.Linear(d_model, d_model), + nn.PReLU(), + nn.Linear(d_model, 1), + RangeSigmoid(), + ) + elif metric == "pesq": + layer3 = nn.Sequential( + nn.Linear(d_model, d_model), + nn.PReLU(), + nn.Linear(d_model, 1), + RangeSigmoid(val_range=PESQRange), + ) + else: + layer3: nn.modules.Module = nn.Sequential( + nn.Linear(d_model, d_model), nn.PReLU(), nn.Linear(d_model, 1) + ) + return nn.Sequential(layer1, layer2, layer3) + + +def squim_objective_model( + feat_dim: int, + win_len: int, + d_model: int, + nhead: int, + hidden_dim: int, + num_blocks: int, + rnn_type: str, + chunk_size: int, + chunk_stride: Optional[int] = None, +) -> SquimObjective: + """Build a custome :class:`torchaudio.prototype.models.SquimObjective` model. + + Args: + feat_dim (int, optional): The feature dimension after Encoder module. + win_len (int): Kernel size in the Encoder module. + d_model (int): The number of expected features in the input. + nhead (int): Number of heads in the multi-head attention model. + hidden_dim (int): Hidden dimension in the RNN layer of DPRNN. + num_blocks (int): Number of DPRNN layers. + rnn_type (str): Type of RNN in DPRNN. Valid options are ["RNN", "LSTM", "GRU"]. + chunk_size (int): Chunk size of input for DPRNN. + chunk_stride (int or None, optional): Stride of chunk input for DPRNN. + """ + if chunk_stride is None: + chunk_stride = chunk_size // 2 + encoder = Encoder(feat_dim, win_len) + dprnn = DPRNN( + feat_dim, hidden_dim, num_blocks, rnn_type, d_model, chunk_size, chunk_stride + ) + branches = nn.ModuleList( + [ + _create_branch(d_model, nhead, "stoi"), + _create_branch(d_model, nhead, "pesq"), + _create_branch(d_model, nhead, "sisdr"), + ] + ) + return SquimObjective(encoder, dprnn, branches) + + +def squim_objective_base() -> SquimObjective: + """Build :class:`torchaudio.prototype.models.SquimObjective` model with default arguments.""" + return squim_objective_model( + feat_dim=256, + win_len=64, + d_model=256, + nhead=4, + hidden_dim=256, + num_blocks=2, + rnn_type="LSTM", + chunk_size=71, + ) + + +@dataclass +class SquimObjectiveBundle: + _path: str + _sample_rate: float + + def _get_state_dict(self, dl_kwargs): + url = f"https://download.pytorch.org/torchaudio/models/{self._path}" + dl_kwargs = {} if dl_kwargs is None else dl_kwargs + state_dict = load_state_dict_from_url(url, **dl_kwargs) + return state_dict + + def get_model(self, *, dl_kwargs=None) -> SquimObjective: + """Construct the SquimObjective model, and load the pretrained weight. + + The weight file is downloaded from the internet and cached with + :func:`torch.hub.load_state_dict_from_url` + + Args: + dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`. + + Returns: + Variation of :py:class:`~torchaudio.models.SquimObjective`. + """ + model = squim_objective_base() + model.load_state_dict(self._get_state_dict(dl_kwargs)) + model.eval() + return model + + @property + def sample_rate(self): + """Sample rate of the audio that the model is trained on. + + :type: float + """ + return self._sample_rate + + +SQUIM_OBJECTIVE = SquimObjectiveBundle( + "squim_objective_dns2020.pth", + _sample_rate=16000, +) +SQUIM_OBJECTIVE.__doc__ = """SquimObjective pipeline trained using approach described in + :cite:`kumar2023torchaudio` on the *DNS 2020 Dataset* :cite:`reddy2020interspeech`. + + The underlying model is constructed by :py:func:`torchaudio.models.squim_objective_base`. + The weights are under `Creative Commons Attribution 4.0 International License + `__. + + Please refer to :py:class:`SquimObjectiveBundle` for usage instructions. + """ diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/snr.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/snr.py new file mode 100644 index 0000000000000000000000000000000000000000..75e92866befd15f5fc41b470555e7daad279415d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/metrics/snr.py @@ -0,0 +1,126 @@ +from typing import Any, Callable + +import numpy as np +import torch +import torchmetrics as tm +from torch._C import _LinAlgError +from torchmetrics import functional as tmF + + +class SafeSignalDistortionRatio(tm.SignalDistortionRatio): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + + def update(self, *args, **kwargs) -> Any: + try: + super().update(*args, **kwargs) + except: + pass + + def compute(self) -> Any: + if self.total == 0: + return torch.tensor(torch.nan) + return super().compute() + + +class BaseChunkMedianSignalRatio(tm.Metric): + def __init__( + self, + func: Callable, + window_size: int, + hop_size: int = None, + zero_mean: bool = False, + ) -> None: + super().__init__() + + # self.zero_mean = zero_mean + self.func = func + self.window_size = window_size + if hop_size is None: + hop_size = window_size + self.hop_size = hop_size + + self.add_state("sum_snr", default=torch.tensor(0.0), dist_reduce_fx="sum") + self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum") + + def update(self, preds: torch.Tensor, target: torch.Tensor) -> None: + n_samples = target.shape[-1] + + n_chunks = int(np.ceil((n_samples - self.window_size) / self.hop_size) + 1) + + snr_chunk = [] + + for i in range(n_chunks): + start = i * self.hop_size + + if n_samples - start < self.window_size: + continue + + end = start + self.window_size + + try: + chunk_snr = self.func(preds[..., start:end], target[..., start:end]) + + # print(preds.shape, chunk_snr.shape) + + if torch.all(torch.isfinite(chunk_snr)): + snr_chunk.append(chunk_snr) + except _LinAlgError: + pass + + snr_chunk = torch.stack(snr_chunk, dim=-1) + snr_batch, _ = torch.nanmedian(snr_chunk, dim=-1) + + self.sum_snr += snr_batch.sum() + self.total += snr_batch.numel() + + def compute(self) -> Any: + return self.sum_snr / self.total + + +class ChunkMedianSignalNoiseRatio(BaseChunkMedianSignalRatio): + def __init__( + self, window_size: int, hop_size: int = None, zero_mean: bool = False + ) -> None: + super().__init__( + func=tmF.signal_noise_ratio, + window_size=window_size, + hop_size=hop_size, + zero_mean=zero_mean, + ) + + +class ChunkMedianScaleInvariantSignalNoiseRatio(BaseChunkMedianSignalRatio): + def __init__( + self, window_size: int, hop_size: int = None, zero_mean: bool = False + ) -> None: + super().__init__( + func=tmF.scale_invariant_signal_noise_ratio, + window_size=window_size, + hop_size=hop_size, + zero_mean=zero_mean, + ) + + +class ChunkMedianSignalDistortionRatio(BaseChunkMedianSignalRatio): + def __init__( + self, window_size: int, hop_size: int = None, zero_mean: bool = False + ) -> None: + super().__init__( + func=tmF.signal_distortion_ratio, + window_size=window_size, + hop_size=hop_size, + zero_mean=zero_mean, + ) + + +class ChunkMedianScaleInvariantSignalDistortionRatio(BaseChunkMedianSignalRatio): + def __init__( + self, window_size: int, hop_size: int = None, zero_mean: bool = False + ) -> None: + super().__init__( + func=tmF.scale_invariant_signal_distortion_ratio, + window_size=window_size, + hop_size=hop_size, + zero_mean=zero_mean, + ) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..54ac48eb69d6f844ba5b73b213eae4cfab157cac --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/__init__.py @@ -0,0 +1,3 @@ +from .bsrnn.wrapper import ( + MultiMaskMultiSourceBandSplitRNNSimple, +) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/_spectral.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/_spectral.py new file mode 100644 index 0000000000000000000000000000000000000000..6af5cbd0dcb6ed0a4babd6b8554184d91c406655 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/_spectral.py @@ -0,0 +1,54 @@ +from typing import Dict, Optional + +import torch +import torchaudio as ta +from torch import nn + + +class _SpectralComponent(nn.Module): + def __init__( + self, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + **kwargs, + ) -> None: + super().__init__() + + assert power is None + + window_fn = torch.__dict__[window_fn] + + self.stft = ta.transforms.Spectrogram( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + pad_mode=pad_mode, + pad=0, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + normalized=normalized, + center=center, + onesided=onesided, + ) + + self.istft = ta.transforms.InverseSpectrogram( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + pad_mode=pad_mode, + pad=0, + window_fn=window_fn, + wkwargs=wkwargs, + normalized=normalized, + center=center, + onesided=onesided, + ) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0f3f78526a0e1bad99aa4635471d472833cc5337 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/__init__.py @@ -0,0 +1,22 @@ +from abc import ABC +from typing import Iterable, Mapping, Union + +from models.bandit.core.model.bsrnn.bandsplit import BandSplitModule +from models.bandit.core.model.bsrnn.tfmodel import ( + SeqBandModellingModule, + TransformerTimeFreqModule, +) +from torch import nn + + +class BandsplitCoreBase(nn.Module, ABC): + band_split: nn.Module + tf_model: nn.Module + mask_estim: Union[nn.Module, Mapping[str, nn.Module], Iterable[nn.Module]] + + def __init__(self) -> None: + super().__init__() + + @staticmethod + def mask(x, m): + return x * m diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/bandsplit.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/bandsplit.py new file mode 100644 index 0000000000000000000000000000000000000000..cf4e50b9c44f40131f59fc32a7be344c47e2b712 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/bandsplit.py @@ -0,0 +1,134 @@ +from typing import List, Tuple + +import torch +from models.bandit.core.model.bsrnn.utils import ( + band_widths_from_specs, + check_no_gap, + check_no_overlap, + check_nonzero_bandwidth, +) +from torch import nn + + +class NormFC(nn.Module): + def __init__( + self, + emb_dim: int, + bandwidth: int, + in_channel: int, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + ) -> None: + super().__init__() + + self.treat_channel_as_feature = treat_channel_as_feature + + if normalize_channel_independently: + raise NotImplementedError + + reim = 2 + + self.norm = nn.LayerNorm(in_channel * bandwidth * reim) + + fc_in = bandwidth * reim + + if treat_channel_as_feature: + fc_in *= in_channel + else: + assert emb_dim % in_channel == 0 + emb_dim = emb_dim // in_channel + + self.fc = nn.Linear(fc_in, emb_dim) + + def forward(self, xb): + # xb = (batch, n_time, in_chan, reim * band_width) + + batch, n_time, in_chan, ribw = xb.shape + xb = self.norm(xb.reshape(batch, n_time, in_chan * ribw)) + # (batch, n_time, in_chan * reim * band_width) + + if not self.treat_channel_as_feature: + xb = xb.reshape(batch, n_time, in_chan, ribw) + # (batch, n_time, in_chan, reim * band_width) + + zb = self.fc(xb) + # (batch, n_time, emb_dim) + # OR + # (batch, n_time, in_chan, emb_dim_per_chan) + + if not self.treat_channel_as_feature: + batch, n_time, in_chan, emb_dim_per_chan = zb.shape + # (batch, n_time, in_chan, emb_dim_per_chan) + zb = zb.reshape((batch, n_time, in_chan * emb_dim_per_chan)) + + return zb # (batch, n_time, emb_dim) + + +class BandSplitModule(nn.Module): + def __init__( + self, + band_specs: List[Tuple[float, float]], + emb_dim: int, + in_channel: int, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + ) -> None: + super().__init__() + + check_nonzero_bandwidth(band_specs) + + if require_no_gap: + check_no_gap(band_specs) + + if require_no_overlap: + check_no_overlap(band_specs) + + self.band_specs = band_specs + # list of [fstart, fend) in index. + # Note that fend is exclusive. + self.band_widths = band_widths_from_specs(band_specs) + self.n_bands = len(band_specs) + self.emb_dim = emb_dim + + self.norm_fc_modules = nn.ModuleList( + [ # type: ignore + ( + NormFC( + emb_dim=emb_dim, + bandwidth=bw, + in_channel=in_channel, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + ) + ) + for bw in self.band_widths + ] + ) + + def forward(self, x: torch.Tensor): + # x = complex spectrogram (batch, in_chan, n_freq, n_time) + + batch, in_chan, _, n_time = x.shape + + z = torch.zeros( + size=(batch, self.n_bands, n_time, self.emb_dim), device=x.device + ) + + xr = torch.view_as_real(x) # batch, in_chan, n_freq, n_time, 2 + xr = torch.permute(xr, (0, 3, 1, 4, 2)) # batch, n_time, in_chan, 2, n_freq + batch, n_time, in_chan, reim, band_width = xr.shape + for i, nfm in enumerate(self.norm_fc_modules): + # print(f"bandsplit/band{i:02d}") + fstart, fend = self.band_specs[i] + xb = xr[..., fstart:fend] + # (batch, n_time, in_chan, reim, band_width) + xb = torch.reshape(xb, (batch, n_time, in_chan, -1)) + # (batch, n_time, in_chan, reim * band_width) + # z.append(nfm(xb)) # (batch, n_time, emb_dim) + z[:, i, :, :] = nfm(xb.contiguous()) + + # z = torch.stack(z, dim=1) + + return z diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/core.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/core.py new file mode 100644 index 0000000000000000000000000000000000000000..c958114511f21638136092387cabb04a094eb624 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/core.py @@ -0,0 +1,645 @@ +from typing import Dict, List, Optional, Tuple + +import torch +from models.bandit.core.model.bsrnn import BandsplitCoreBase +from models.bandit.core.model.bsrnn.bandsplit import BandSplitModule +from models.bandit.core.model.bsrnn.maskestim import ( + MaskEstimationModule, + OverlappingMaskEstimationModule, +) +from models.bandit.core.model.bsrnn.tfmodel import ( + ConvolutionalTimeFreqModule, + SeqBandModellingModule, + TransformerTimeFreqModule, +) +from torch import nn +from torch.nn import functional as F + + +class MultiMaskBandSplitCoreBase(BandsplitCoreBase): + def __init__(self) -> None: + super().__init__() + + def forward(self, x, cond=None, compute_residual: bool = True): + # x = complex spectrogram (batch, in_chan, n_freq, n_time) + # print(x.shape) + batch, in_chan, n_freq, n_time = x.shape + x = torch.reshape(x, (-1, 1, n_freq, n_time)) + + z = self.band_split(x) # (batch, emb_dim, n_band, n_time) + + # if torch.any(torch.isnan(z)): + # raise ValueError("z nan") + + # print(z) + q = self.tf_model(z) # (batch, emb_dim, n_band, n_time) + # print(q) + + # if torch.any(torch.isnan(q)): + # raise ValueError("q nan") + + out = {} + + for stem, mem in self.mask_estim.items(): + m = mem(q, cond=cond) + + # if torch.any(torch.isnan(m)): + # raise ValueError("m nan", stem) + + s = self.mask(x, m) + s = torch.reshape(s, (batch, in_chan, n_freq, n_time)) + out[stem] = s + + return {"spectrogram": out} + + def instantiate_mask_estim( + self, + in_channel: int, + stems: List[str], + band_specs: List[Tuple[float, float]], + emb_dim: int, + mlp_dim: int, + cond_dim: int, + hidden_activation: str, + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + overlapping_band: bool = False, + freq_weights: Optional[List[torch.Tensor]] = None, + n_freq: Optional[int] = None, + use_freq_weights: bool = True, + mult_add_mask: bool = False, + ): + if hidden_activation_kwargs is None: + hidden_activation_kwargs = {} + + if "mne:+" in stems: + stems = [s for s in stems if s != "mne:+"] + + if overlapping_band: + assert freq_weights is not None + assert n_freq is not None + + if mult_add_mask: + self.mask_estim = nn.ModuleDict( + { + stem: MultAddMaskEstimationModule( + band_specs=band_specs, + freq_weights=freq_weights, + n_freq=n_freq, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + in_channel=in_channel, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + use_freq_weights=use_freq_weights, + ) + for stem in stems + } + ) + else: + self.mask_estim = nn.ModuleDict( + { + stem: OverlappingMaskEstimationModule( + band_specs=band_specs, + freq_weights=freq_weights, + n_freq=n_freq, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + in_channel=in_channel, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + use_freq_weights=use_freq_weights, + ) + for stem in stems + } + ) + else: + self.mask_estim = nn.ModuleDict( + { + stem: MaskEstimationModule( + band_specs=band_specs, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + in_channel=in_channel, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + for stem in stems + } + ) + + def instantiate_bandsplit( + self, + in_channel: int, + band_specs: List[Tuple[float, float]], + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + emb_dim: int = 128, + ): + self.band_split = BandSplitModule( + in_channel=in_channel, + band_specs=band_specs, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + + +class SingleMaskBandsplitCoreBase(BandsplitCoreBase): + def __init__(self, **kwargs) -> None: + super().__init__() + + def forward(self, x): + # x = complex spectrogram (batch, in_chan, n_freq, n_time) + z = self.band_split(x) # (batch, emb_dim, n_band, n_time) + q = self.tf_model(z) # (batch, emb_dim, n_band, n_time) + m = self.mask_estim(q) # (batch, in_chan, n_freq, n_time) + + s = self.mask(x, m) + + return s + + +class SingleMaskBandsplitCoreRNN( + SingleMaskBandsplitCoreBase, +): + def __init__( + self, + in_channel: int, + band_specs: List[Tuple[float, float]], + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + ) -> None: + super().__init__() + self.band_split = BandSplitModule( + in_channel=in_channel, + band_specs=band_specs, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + self.tf_model = SeqBandModellingModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ) + self.mask_estim = MaskEstimationModule( + in_channel=in_channel, + band_specs=band_specs, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + + +class SingleMaskBandsplitCoreTransformer( + SingleMaskBandsplitCoreBase, +): + def __init__( + self, + in_channel: int, + band_specs: List[Tuple[float, float]], + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + tf_dropout: float = 0.0, + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + ) -> None: + super().__init__() + self.band_split = BandSplitModule( + in_channel=in_channel, + band_specs=band_specs, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + self.tf_model = TransformerTimeFreqModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + dropout=tf_dropout, + ) + self.mask_estim = MaskEstimationModule( + in_channel=in_channel, + band_specs=band_specs, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + + +class MultiSourceMultiMaskBandSplitCoreRNN(MultiMaskBandSplitCoreBase): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: List[Tuple[float, float]], + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + cond_dim: int = 0, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + overlapping_band: bool = False, + freq_weights: Optional[List[torch.Tensor]] = None, + n_freq: Optional[int] = None, + use_freq_weights: bool = True, + mult_add_mask: bool = False, + ) -> None: + super().__init__() + self.instantiate_bandsplit( + in_channel=in_channel, + band_specs=band_specs, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + + self.tf_model = SeqBandModellingModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ) + + self.mult_add_mask = mult_add_mask + + self.instantiate_mask_estim( + in_channel=in_channel, + stems=stems, + band_specs=band_specs, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=overlapping_band, + freq_weights=freq_weights, + n_freq=n_freq, + use_freq_weights=use_freq_weights, + mult_add_mask=mult_add_mask, + ) + + @staticmethod + def _mult_add_mask(x, m): + assert m.ndim == 5 + + mm = m[..., 0] + am = m[..., 1] + + # print(mm.shape, am.shape, x.shape, m.shape) + + return x * mm + am + + def mask(self, x, m): + if self.mult_add_mask: + return self._mult_add_mask(x, m) + else: + return super().mask(x, m) + + +class MultiSourceMultiMaskBandSplitCoreTransformer( + MultiMaskBandSplitCoreBase, +): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: List[Tuple[float, float]], + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + tf_dropout: float = 0.0, + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + overlapping_band: bool = False, + freq_weights: Optional[List[torch.Tensor]] = None, + n_freq: Optional[int] = None, + use_freq_weights: bool = True, + rnn_type: str = "LSTM", + cond_dim: int = 0, + mult_add_mask: bool = False, + ) -> None: + super().__init__() + self.instantiate_bandsplit( + in_channel=in_channel, + band_specs=band_specs, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + self.tf_model = TransformerTimeFreqModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + dropout=tf_dropout, + ) + + self.instantiate_mask_estim( + in_channel=in_channel, + stems=stems, + band_specs=band_specs, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=overlapping_band, + freq_weights=freq_weights, + n_freq=n_freq, + use_freq_weights=use_freq_weights, + mult_add_mask=mult_add_mask, + ) + + +class MultiSourceMultiMaskBandSplitCoreConv( + MultiMaskBandSplitCoreBase, +): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: List[Tuple[float, float]], + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + tf_dropout: float = 0.0, + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + overlapping_band: bool = False, + freq_weights: Optional[List[torch.Tensor]] = None, + n_freq: Optional[int] = None, + use_freq_weights: bool = True, + rnn_type: str = "LSTM", + cond_dim: int = 0, + mult_add_mask: bool = False, + ) -> None: + super().__init__() + self.instantiate_bandsplit( + in_channel=in_channel, + band_specs=band_specs, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + self.tf_model = ConvolutionalTimeFreqModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + dropout=tf_dropout, + ) + + self.instantiate_mask_estim( + in_channel=in_channel, + stems=stems, + band_specs=band_specs, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=overlapping_band, + freq_weights=freq_weights, + n_freq=n_freq, + use_freq_weights=use_freq_weights, + mult_add_mask=mult_add_mask, + ) + + +class PatchingMaskBandsplitCoreBase(MultiMaskBandSplitCoreBase): + def __init__(self) -> None: + super().__init__() + + def mask(self, x, m): + # x.shape = (batch, n_channel, n_freq, n_time) + # m.shape = (kernel_freq, kernel_time, batch, n_channel, n_freq, n_time) + + _, n_channel, kernel_freq, kernel_time, n_freq, n_time = m.shape + padding = ((kernel_freq - 1) // 2, (kernel_time - 1) // 2) + + xf = F.unfold( + x, + kernel_size=(kernel_freq, kernel_time), + padding=padding, + stride=(1, 1), + ) + + xf = xf.view( + -1, + n_channel, + kernel_freq, + kernel_time, + n_freq, + n_time, + ) + + sf = xf * m + + sf = sf.view( + -1, + n_channel * kernel_freq * kernel_time, + n_freq * n_time, + ) + + s = F.fold( + sf, + output_size=(n_freq, n_time), + kernel_size=(kernel_freq, kernel_time), + padding=padding, + stride=(1, 1), + ).view( + -1, + n_channel, + n_freq, + n_time, + ) + + return s + + def old_mask(self, x, m): + # x.shape = (batch, n_channel, n_freq, n_time) + # m.shape = (kernel_freq, kernel_time, batch, n_channel, n_freq, n_time) + + s = torch.zeros_like(x) + + _, n_channel, n_freq, n_time = x.shape + kernel_freq, kernel_time, _, _, _, _ = m.shape + + # print(x.shape, m.shape) + + kernel_freq_half = (kernel_freq - 1) // 2 + kernel_time_half = (kernel_time - 1) // 2 + + for ifreq in range(kernel_freq): + for itime in range(kernel_time): + df, dt = kernel_freq_half - ifreq, kernel_time_half - itime + x = x.roll(shifts=(df, dt), dims=(2, 3)) + + # if `df` > 0: + # x[:, :, :df, :] = 0 + # elif `df` < 0: + # x[:, :, df:, :] = 0 + + # if `dt` > 0: + # x[:, :, :, :dt] = 0 + # elif `dt` < 0: + # x[:, :, :, dt:] = 0 + + fslice = slice(max(0, df), min(n_freq, n_freq + df)) + tslice = slice(max(0, dt), min(n_time, n_time + dt)) + + s[:, :, fslice, tslice] += ( + x[:, :, fslice, tslice] * m[ifreq, itime, :, :, fslice, tslice] + ) + + return s + + +class MultiSourceMultiPatchingMaskBandSplitCoreRNN(PatchingMaskBandsplitCoreBase): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: List[Tuple[float, float]], + mask_kernel_freq: int, + mask_kernel_time: int, + conv_kernel_freq: int, + conv_kernel_time: int, + kernel_norm_mlp_version: int, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + overlapping_band: bool = False, + freq_weights: Optional[List[torch.Tensor]] = None, + n_freq: Optional[int] = None, + ) -> None: + super().__init__() + self.band_split = BandSplitModule( + in_channel=in_channel, + band_specs=band_specs, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + + self.tf_model = SeqBandModellingModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ) + + if hidden_activation_kwargs is None: + hidden_activation_kwargs = {} + + if overlapping_band: + assert freq_weights is not None + assert n_freq is not None + self.mask_estim = nn.ModuleDict( + { + stem: PatchingMaskEstimationModule( + band_specs=band_specs, + freq_weights=freq_weights, + n_freq=n_freq, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + in_channel=in_channel, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + mask_kernel_freq=mask_kernel_freq, + mask_kernel_time=mask_kernel_time, + conv_kernel_freq=conv_kernel_freq, + conv_kernel_time=conv_kernel_time, + kernel_norm_mlp_version=kernel_norm_mlp_version, + ) + for stem in stems + } + ) + else: + raise NotImplementedError diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/maskestim.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/maskestim.py new file mode 100644 index 0000000000000000000000000000000000000000..69ef4f846415f0ff7f451c92cba96132f30d8823 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/maskestim.py @@ -0,0 +1,347 @@ +from typing import Dict, List, Optional, Tuple, Type + +import torch +from models.bandit.core.model.bsrnn.utils import ( + band_widths_from_specs, + check_no_gap, + check_no_overlap, + check_nonzero_bandwidth, +) +from torch import nn +from torch.nn.modules import activation + + +class BaseNormMLP(nn.Module): + def __init__( + self, + emb_dim: int, + mlp_dim: int, + bandwidth: int, + in_channel: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs=None, + complex_mask: bool = True, + ): + super().__init__() + if hidden_activation_kwargs is None: + hidden_activation_kwargs = {} + self.hidden_activation_kwargs = hidden_activation_kwargs + self.norm = nn.LayerNorm(emb_dim) + self.hidden = torch.jit.script( + nn.Sequential( + nn.Linear(in_features=emb_dim, out_features=mlp_dim), + activation.__dict__[hidden_activation](**self.hidden_activation_kwargs), + ) + ) + + self.bandwidth = bandwidth + self.in_channel = in_channel + + self.complex_mask = complex_mask + self.reim = 2 if complex_mask else 1 + self.glu_mult = 2 + + +class NormMLP(BaseNormMLP): + def __init__( + self, + emb_dim: int, + mlp_dim: int, + bandwidth: int, + in_channel: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs=None, + complex_mask: bool = True, + ) -> None: + super().__init__( + emb_dim=emb_dim, + mlp_dim=mlp_dim, + bandwidth=bandwidth, + in_channel=in_channel, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + + self.output = torch.jit.script( + nn.Sequential( + nn.Linear( + in_features=mlp_dim, + out_features=bandwidth * in_channel * self.reim * 2, + ), + nn.GLU(dim=-1), + ) + ) + + def reshape_output(self, mb): + # print(mb.shape) + batch, n_time, _ = mb.shape + if self.complex_mask: + mb = mb.reshape( + batch, n_time, self.in_channel, self.bandwidth, self.reim + ).contiguous() + # print(mb.shape) + mb = torch.view_as_complex(mb) # (batch, n_time, in_channel, bandwidth) + else: + mb = mb.reshape(batch, n_time, self.in_channel, self.bandwidth) + + mb = torch.permute(mb, (0, 2, 3, 1)) # (batch, in_channel, bandwidth, n_time) + + return mb + + def forward(self, qb): + # qb = (batch, n_time, emb_dim) + + # if torch.any(torch.isnan(qb)): + # raise ValueError("qb0") + + qb = self.norm(qb) # (batch, n_time, emb_dim) + + # if torch.any(torch.isnan(qb)): + # raise ValueError("qb1") + + qb = self.hidden(qb) # (batch, n_time, mlp_dim) + # if torch.any(torch.isnan(qb)): + # raise ValueError("qb2") + mb = self.output(qb) # (batch, n_time, bandwidth * in_channel * reim) + # if torch.any(torch.isnan(qb)): + # raise ValueError("mb") + mb = self.reshape_output(mb) # (batch, in_channel, bandwidth, n_time) + + return mb + + +class MultAddNormMLP(NormMLP): + def __init__( + self, + emb_dim: int, + mlp_dim: int, + bandwidth: int, + in_channel: "int | None", + hidden_activation: str = "Tanh", + hidden_activation_kwargs=None, + complex_mask: bool = True, + ) -> None: + super().__init__( + emb_dim, + mlp_dim, + bandwidth, + in_channel, + hidden_activation, + hidden_activation_kwargs, + complex_mask, + ) + + self.output2 = torch.jit.script( + nn.Sequential( + nn.Linear( + in_features=mlp_dim, + out_features=bandwidth * in_channel * self.reim * 2, + ), + nn.GLU(dim=-1), + ) + ) + + def forward(self, qb): + qb = self.norm(qb) # (batch, n_time, emb_dim) + qb = self.hidden(qb) # (batch, n_time, mlp_dim) + mmb = self.output(qb) # (batch, n_time, bandwidth * in_channel * reim) + mmb = self.reshape_output(mmb) # (batch, in_channel, bandwidth, n_time) + amb = self.output2(qb) # (batch, n_time, bandwidth * in_channel * reim) + amb = self.reshape_output(amb) # (batch, in_channel, bandwidth, n_time) + + return mmb, amb + + +class MaskEstimationModuleSuperBase(nn.Module): + pass + + +class MaskEstimationModuleBase(MaskEstimationModuleSuperBase): + def __init__( + self, + band_specs: List[Tuple[float, float]], + emb_dim: int, + mlp_dim: int, + in_channel: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Dict = None, + complex_mask: bool = True, + norm_mlp_cls: Type[nn.Module] = NormMLP, + norm_mlp_kwargs: Dict = None, + ) -> None: + super().__init__() + + self.band_widths = band_widths_from_specs(band_specs) + self.n_bands = len(band_specs) + + if hidden_activation_kwargs is None: + hidden_activation_kwargs = {} + + if norm_mlp_kwargs is None: + norm_mlp_kwargs = {} + + self.norm_mlp = nn.ModuleList( + [ + ( + norm_mlp_cls( + bandwidth=self.band_widths[b], + emb_dim=emb_dim, + mlp_dim=mlp_dim, + in_channel=in_channel, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + **norm_mlp_kwargs, + ) + ) + for b in range(self.n_bands) + ] + ) + + def compute_masks(self, q): + batch, n_bands, n_time, emb_dim = q.shape + + masks = [] + + for b, nmlp in enumerate(self.norm_mlp): + # print(f"maskestim/{b:02d}") + qb = q[:, b, :, :] + mb = nmlp(qb) + masks.append(mb) + + return masks + + +class OverlappingMaskEstimationModule(MaskEstimationModuleBase): + def __init__( + self, + in_channel: int, + band_specs: List[Tuple[float, float]], + freq_weights: List[torch.Tensor], + n_freq: int, + emb_dim: int, + mlp_dim: int, + cond_dim: int = 0, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Dict = None, + complex_mask: bool = True, + norm_mlp_cls: Type[nn.Module] = NormMLP, + norm_mlp_kwargs: Dict = None, + use_freq_weights: bool = True, + ) -> None: + check_nonzero_bandwidth(band_specs) + check_no_gap(band_specs) + + # if cond_dim > 0: + # raise NotImplementedError + + super().__init__( + band_specs=band_specs, + emb_dim=emb_dim + cond_dim, + mlp_dim=mlp_dim, + in_channel=in_channel, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + norm_mlp_cls=norm_mlp_cls, + norm_mlp_kwargs=norm_mlp_kwargs, + ) + + self.n_freq = n_freq + self.band_specs = band_specs + self.in_channel = in_channel + + if freq_weights is not None: + for i, fw in enumerate(freq_weights): + self.register_buffer(f"freq_weights/{i}", fw) + + self.use_freq_weights = use_freq_weights + else: + self.use_freq_weights = False + + self.cond_dim = cond_dim + + def forward(self, q, cond=None): + # q = (batch, n_bands, n_time, emb_dim) + + batch, n_bands, n_time, emb_dim = q.shape + + if cond is not None: + print(cond) + if cond.ndim == 2: + cond = cond[:, None, None, :].expand(-1, n_bands, n_time, -1) + elif cond.ndim == 3: + assert cond.shape[1] == n_time + else: + raise ValueError(f"Invalid cond shape: {cond.shape}") + + q = torch.cat([q, cond], dim=-1) + elif self.cond_dim > 0: + cond = torch.ones( + (batch, n_bands, n_time, self.cond_dim), + device=q.device, + dtype=q.dtype, + ) + q = torch.cat([q, cond], dim=-1) + else: + pass + + mask_list = self.compute_masks( + q + ) # [n_bands * (batch, in_channel, bandwidth, n_time)] + + masks = torch.zeros( + (batch, self.in_channel, self.n_freq, n_time), + device=q.device, + dtype=mask_list[0].dtype, + ) + + for im, mask in enumerate(mask_list): + fstart, fend = self.band_specs[im] + if self.use_freq_weights: + fw = self.get_buffer(f"freq_weights/{im}")[:, None] + mask = mask * fw + masks[:, :, fstart:fend, :] += mask + + return masks + + +class MaskEstimationModule(OverlappingMaskEstimationModule): + def __init__( + self, + band_specs: List[Tuple[float, float]], + emb_dim: int, + mlp_dim: int, + in_channel: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Dict = None, + complex_mask: bool = True, + **kwargs, + ) -> None: + check_nonzero_bandwidth(band_specs) + check_no_gap(band_specs) + check_no_overlap(band_specs) + super().__init__( + in_channel=in_channel, + band_specs=band_specs, + freq_weights=None, + n_freq=None, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + + def forward(self, q, cond=None): + # q = (batch, n_bands, n_time, emb_dim) + + masks = self.compute_masks( + q + ) # [n_bands * (batch, in_channel, bandwidth, n_time)] + + # TODO: currently this requires band specs to have no gap and no overlap + masks = torch.concat(masks, dim=2) # (batch, in_channel, n_freq, n_time) + + return masks diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/tfmodel.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/tfmodel.py new file mode 100644 index 0000000000000000000000000000000000000000..a79390c0ae0afd103072310a999bb5baf99658f9 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/tfmodel.py @@ -0,0 +1,317 @@ +import warnings + +import torch +import torch.backends.cuda +from torch import nn +from torch.nn.modules import rnn + + +class TimeFrequencyModellingModule(nn.Module): + def __init__(self) -> None: + super().__init__() + + +class ResidualRNN(nn.Module): + def __init__( + self, + emb_dim: int, + rnn_dim: int, + bidirectional: bool = True, + rnn_type: str = "LSTM", + use_batch_trick: bool = True, + use_layer_norm: bool = True, + ) -> None: + # n_group is the size of the 2nd dim + super().__init__() + + self.use_layer_norm = use_layer_norm + if use_layer_norm: + self.norm = nn.LayerNorm(emb_dim) + else: + self.norm = nn.GroupNorm(num_groups=emb_dim, num_channels=emb_dim) + + self.rnn = rnn.__dict__[rnn_type]( + input_size=emb_dim, + hidden_size=rnn_dim, + num_layers=1, + batch_first=True, + bidirectional=bidirectional, + ) + + self.fc = nn.Linear( + in_features=rnn_dim * (2 if bidirectional else 1), out_features=emb_dim + ) + + self.use_batch_trick = use_batch_trick + if not self.use_batch_trick: + warnings.warn("NOT USING BATCH TRICK IS EXTREMELY SLOW!!") + + def forward(self, z): + # z = (batch, n_uncrossed, n_across, emb_dim) + + z0 = torch.clone(z) + + # print(z.device) + + if self.use_layer_norm: + z = self.norm(z) # (batch, n_uncrossed, n_across, emb_dim) + else: + z = torch.permute( + z, (0, 3, 1, 2) + ) # (batch, emb_dim, n_uncrossed, n_across) + + z = self.norm(z) # (batch, emb_dim, n_uncrossed, n_across) + + z = torch.permute( + z, (0, 2, 3, 1) + ) # (batch, n_uncrossed, n_across, emb_dim) + + batch, n_uncrossed, n_across, emb_dim = z.shape + + if self.use_batch_trick: + z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim)) + + z = self.rnn(z.contiguous())[ + 0 + ] # (batch * n_uncrossed, n_across, dir_rnn_dim) + + z = torch.reshape(z, (batch, n_uncrossed, n_across, -1)) + # (batch, n_uncrossed, n_across, dir_rnn_dim) + else: + # Note: this is EXTREMELY SLOW + zlist = [] + for i in range(n_uncrossed): + zi = self.rnn(z[:, i, :, :])[0] # (batch, n_across, emb_dim) + zlist.append(zi) + + z = torch.stack(zlist, dim=1) # (batch, n_uncrossed, n_across, dir_rnn_dim) + + z = self.fc(z) # (batch, n_uncrossed, n_across, emb_dim) + + z = z + z0 + + return z + + +class SeqBandModellingModule(TimeFrequencyModellingModule): + def __init__( + self, + n_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + parallel_mode=False, + ) -> None: + super().__init__() + self.seqband = nn.ModuleList([]) + + if parallel_mode: + for _ in range(n_modules): + self.seqband.append( + nn.ModuleList( + [ + ResidualRNN( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ), + ResidualRNN( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ), + ] + ) + ) + else: + for _ in range(2 * n_modules): + self.seqband.append( + ResidualRNN( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ) + ) + + self.parallel_mode = parallel_mode + + def forward(self, z): + # z = (batch, n_bands, n_time, emb_dim) + + if self.parallel_mode: + for sbm_pair in self.seqband: + # z: (batch, n_bands, n_time, emb_dim) + sbm_t, sbm_f = sbm_pair[0], sbm_pair[1] + zt = sbm_t(z) # (batch, n_bands, n_time, emb_dim) + zf = sbm_f(z.transpose(1, 2)) # (batch, n_time, n_bands, emb_dim) + z = zt + zf.transpose(1, 2) + else: + for sbm in self.seqband: + z = sbm(z) + z = z.transpose(1, 2) + + # (batch, n_bands, n_time, emb_dim) + # --> (batch, n_time, n_bands, emb_dim) + # OR + # (batch, n_time, n_bands, emb_dim) + # --> (batch, n_bands, n_time, emb_dim) + + q = z + return q # (batch, n_bands, n_time, emb_dim) + + +class ResidualTransformer(nn.Module): + def __init__( + self, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + dropout: float = 0.0, + ) -> None: + # n_group is the size of the 2nd dim + super().__init__() + + self.tf = nn.TransformerEncoderLayer( + d_model=emb_dim, nhead=4, dim_feedforward=rnn_dim, batch_first=True + ) + + self.is_causal = not bidirectional + self.dropout = dropout + + def forward(self, z): + batch, n_uncrossed, n_across, emb_dim = z.shape + z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim)) + z = self.tf( + z, is_causal=self.is_causal + ) # (batch, n_uncrossed, n_across, emb_dim) + z = torch.reshape(z, (batch, n_uncrossed, n_across, emb_dim)) + + return z + + +class TransformerTimeFreqModule(TimeFrequencyModellingModule): + def __init__( + self, + n_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.norm = nn.LayerNorm(emb_dim) + self.seqband = nn.ModuleList([]) + + for _ in range(2 * n_modules): + self.seqband.append( + ResidualTransformer( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + dropout=dropout, + ) + ) + + def forward(self, z): + # z = (batch, n_bands, n_time, emb_dim) + z = self.norm(z) # (batch, n_bands, n_time, emb_dim) + + for sbm in self.seqband: + z = sbm(z) + z = z.transpose(1, 2) + + # (batch, n_bands, n_time, emb_dim) + # --> (batch, n_time, n_bands, emb_dim) + # OR + # (batch, n_time, n_bands, emb_dim) + # --> (batch, n_bands, n_time, emb_dim) + + q = z + return q # (batch, n_bands, n_time, emb_dim) + + +class ResidualConvolution(nn.Module): + def __init__( + self, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + dropout: float = 0.0, + ) -> None: + # n_group is the size of the 2nd dim + super().__init__() + self.norm = nn.InstanceNorm2d(emb_dim, affine=True) + + self.conv = nn.Sequential( + nn.Conv2d( + in_channels=emb_dim, + out_channels=rnn_dim, + kernel_size=(3, 3), + padding="same", + stride=(1, 1), + ), + nn.Tanhshrink(), + ) + + self.is_causal = not bidirectional + self.dropout = dropout + + self.fc = nn.Conv2d( + in_channels=rnn_dim, + out_channels=emb_dim, + kernel_size=(1, 1), + padding="same", + stride=(1, 1), + ) + + def forward(self, z): + # z = (batch, n_uncrossed, n_across, emb_dim) + + z0 = torch.clone(z) + + z = self.norm(z) # (batch, n_uncrossed, n_across, emb_dim) + z = self.conv(z) # (batch, n_uncrossed, n_across, emb_dim) + z = self.fc(z) # (batch, n_uncrossed, n_across, emb_dim) + z = z + z0 + + return z + + +class ConvolutionalTimeFreqModule(TimeFrequencyModellingModule): + def __init__( + self, + n_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.seqband = torch.jit.script( + nn.Sequential( + *[ + ResidualConvolution( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + dropout=dropout, + ) + for _ in range(2 * n_modules) + ] + ) + ) + + def forward(self, z): + # z = (batch, n_bands, n_time, emb_dim) + + z = torch.permute(z, (0, 3, 1, 2)) # (batch, emb_dim, n_bands, n_time) + + z = self.seqband(z) # (batch, emb_dim, n_bands, n_time) + + z = torch.permute(z, (0, 2, 3, 1)) # (batch, n_bands, n_time, emb_dim) + + return z diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/utils.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..645074f07afb6b9618ee21f721a5500d70740b1f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/utils.py @@ -0,0 +1,522 @@ +import os +from abc import abstractmethod +from typing import Callable + +import numpy as np +import torch +from librosa import hz_to_midi, midi_to_hz +from spafe.fbanks import bark_fbanks +from spafe.utils.converters import hz2bark, hz2erb +from torch import Tensor +from torchaudio import functional as taF +from torchaudio.functional.functional import _create_triangular_filterbank + + +def band_widths_from_specs(band_specs): + return [e - i for i, e in band_specs] + + +def check_nonzero_bandwidth(band_specs): + # pprint(band_specs) + for fstart, fend in band_specs: + if fend - fstart <= 0: + raise ValueError("Bands cannot be zero-width") + + +def check_no_overlap(band_specs): + fend_prev = -1 + for fstart_curr, fend_curr in band_specs: + if fstart_curr <= fend_prev: + raise ValueError("Bands cannot overlap") + + +def check_no_gap(band_specs): + fstart, _ = band_specs[0] + assert fstart == 0 + + fend_prev = -1 + for fstart_curr, fend_curr in band_specs: + if fstart_curr - fend_prev > 1: + raise ValueError("Bands cannot leave gap") + fend_prev = fend_curr + + +class BandsplitSpecification: + def __init__(self, nfft: int, fs: int) -> None: + self.fs = fs + self.nfft = nfft + self.nyquist = fs / 2 + self.max_index = nfft // 2 + 1 + + self.split500 = self.hertz_to_index(500) + self.split1k = self.hertz_to_index(1000) + self.split2k = self.hertz_to_index(2000) + self.split4k = self.hertz_to_index(4000) + self.split8k = self.hertz_to_index(8000) + self.split16k = self.hertz_to_index(16000) + self.split20k = self.hertz_to_index(20000) + + self.above20k = [(self.split20k, self.max_index)] + self.above16k = [(self.split16k, self.split20k)] + self.above20k + + def index_to_hertz(self, index: int): + return index * self.fs / self.nfft + + def hertz_to_index(self, hz: float, round: bool = True): + index = hz * self.nfft / self.fs + + if round: + index = int(np.round(index)) + + return index + + def get_band_specs_with_bandwidth(self, start_index, end_index, bandwidth_hz): + band_specs = [] + lower = start_index + + while lower < end_index: + upper = int(np.floor(lower + self.hertz_to_index(bandwidth_hz))) + upper = min(upper, end_index) + + band_specs.append((lower, upper)) + lower = upper + + return band_specs + + @abstractmethod + def get_band_specs(self): + raise NotImplementedError + + +class VocalBandsplitSpecification(BandsplitSpecification): + def __init__(self, nfft: int, fs: int, version: str = "7") -> None: + super().__init__(nfft=nfft, fs=fs) + + self.version = version + + def get_band_specs(self): + return getattr(self, f"version{self.version}")() + + @property + def version1(self): + return self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.max_index, bandwidth_hz=1000 + ) + + def version2(self): + below16k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split16k, bandwidth_hz=1000 + ) + below20k = self.get_band_specs_with_bandwidth( + start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000 + ) + + return below16k + below20k + self.above20k + + def version3(self): + below8k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + + return below8k + below16k + self.above16k + + def version4(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + + return below1k + below8k + below16k + self.above16k + + def version5(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split16k, bandwidth_hz=1000 + ) + below20k = self.get_band_specs_with_bandwidth( + start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000 + ) + return below1k + below16k + below20k + self.above20k + + def version6(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + return below1k + below4k + below8k + below16k + self.above16k + + def version7(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split4k, bandwidth_hz=250 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000 + ) + below20k = self.get_band_specs_with_bandwidth( + start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000 + ) + return below1k + below4k + below8k + below16k + below20k + self.above20k + + +class OtherBandsplitSpecification(VocalBandsplitSpecification): + def __init__(self, nfft: int, fs: int) -> None: + super().__init__(nfft=nfft, fs=fs, version="7") + + +class BassBandsplitSpecification(BandsplitSpecification): + def __init__(self, nfft: int, fs: int, version: str = "7") -> None: + super().__init__(nfft=nfft, fs=fs) + + def get_band_specs(self): + below500 = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split500, bandwidth_hz=50 + ) + below1k = self.get_band_specs_with_bandwidth( + start_index=self.split500, end_index=self.split1k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + above16k = [(self.split16k, self.max_index)] + + return below500 + below1k + below4k + below8k + below16k + above16k + + +class DrumBandsplitSpecification(BandsplitSpecification): + def __init__(self, nfft: int, fs: int) -> None: + super().__init__(nfft=nfft, fs=fs) + + def get_band_specs(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=50 + ) + below2k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split2k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split2k, end_index=self.split4k, bandwidth_hz=250 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000 + ) + above16k = [(self.split16k, self.max_index)] + + return below1k + below2k + below4k + below8k + below16k + above16k + + +class PerceptualBandsplitSpecification(BandsplitSpecification): + def __init__( + self, + nfft: int, + fs: int, + fbank_fn: Callable[[int, int, float, float, int], torch.Tensor], + n_bands: int, + f_min: float = 0.0, + f_max: float = None, + ) -> None: + super().__init__(nfft=nfft, fs=fs) + self.n_bands = n_bands + if f_max is None: + f_max = fs / 2 + + self.filterbank = fbank_fn(n_bands, fs, f_min, f_max, self.max_index) + + weight_per_bin = torch.sum(self.filterbank, dim=0, keepdim=True) # (1, n_freqs) + normalized_mel_fb = self.filterbank / weight_per_bin # (n_mels, n_freqs) + + freq_weights = [] + band_specs = [] + for i in range(self.n_bands): + active_bins = torch.nonzero(self.filterbank[i, :]).squeeze().tolist() + if isinstance(active_bins, int): + active_bins = (active_bins, active_bins) + if len(active_bins) == 0: + continue + start_index = active_bins[0] + end_index = active_bins[-1] + 1 + band_specs.append((start_index, end_index)) + freq_weights.append(normalized_mel_fb[i, start_index:end_index]) + + self.freq_weights = freq_weights + self.band_specs = band_specs + + def get_band_specs(self): + return self.band_specs + + def get_freq_weights(self): + return self.freq_weights + + def save_to_file(self, dir_path: str) -> None: + os.makedirs(dir_path, exist_ok=True) + + import pickle + + with open(os.path.join(dir_path, "mel_bandsplit_spec.pkl"), "wb") as f: + pickle.dump( + { + "band_specs": self.band_specs, + "freq_weights": self.freq_weights, + "filterbank": self.filterbank, + }, + f, + ) + + +def mel_filterbank(n_bands, fs, f_min, f_max, n_freqs): + fb = taF.melscale_fbanks( + n_mels=n_bands, + sample_rate=fs, + f_min=f_min, + f_max=f_max, + n_freqs=n_freqs, + ).T + + fb[0, 0] = 1.0 + + return fb + + +class MelBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=mel_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +def musical_filterbank(n_bands, fs, f_min, f_max, n_freqs, scale="constant"): + nfft = 2 * (n_freqs - 1) + df = fs / nfft + # init freqs + f_max = f_max or fs / 2 + f_min = f_min or 0 + f_min = fs / nfft + + n_octaves = np.log2(f_max / f_min) + n_octaves_per_band = n_octaves / n_bands + bandwidth_mult = np.power(2.0, n_octaves_per_band) + + low_midi = max(0, hz_to_midi(f_min)) + high_midi = hz_to_midi(f_max) + midi_points = np.linspace(low_midi, high_midi, n_bands) + hz_pts = midi_to_hz(midi_points) + + low_pts = hz_pts / bandwidth_mult + high_pts = hz_pts * bandwidth_mult + + low_bins = np.floor(low_pts / df).astype(int) + high_bins = np.ceil(high_pts / df).astype(int) + + fb = np.zeros((n_bands, n_freqs)) + + for i in range(n_bands): + fb[i, low_bins[i] : high_bins[i] + 1] = 1.0 + + fb[0, : low_bins[0]] = 1.0 + fb[-1, high_bins[-1] + 1 :] = 1.0 + + return torch.as_tensor(fb) + + +class MusicalBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=musical_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +def bark_filterbank(n_bands, fs, f_min, f_max, n_freqs): + nfft = 2 * (n_freqs - 1) + fb, _ = bark_fbanks.bark_filter_banks( + nfilts=n_bands, + nfft=nfft, + fs=fs, + low_freq=f_min, + high_freq=f_max, + scale="constant", + ) + + return torch.as_tensor(fb) + + +class BarkBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=bark_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +def triangular_bark_filterbank(n_bands, fs, f_min, f_max, n_freqs): + all_freqs = torch.linspace(0, fs // 2, n_freqs) + + # calculate mel freq bins + m_min = hz2bark(f_min) + m_max = hz2bark(f_max) + + m_pts = torch.linspace(m_min, m_max, n_bands + 2) + f_pts = 600 * torch.sinh(m_pts / 6) + + # create filterbank + fb = _create_triangular_filterbank(all_freqs, f_pts) + + fb = fb.T + + first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0] + first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0] + + fb[first_active_band, :first_active_bin] = 1.0 + + return fb + + +class TriangularBarkBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=triangular_bark_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +def minibark_filterbank(n_bands, fs, f_min, f_max, n_freqs): + fb = bark_filterbank(n_bands, fs, f_min, f_max, n_freqs) + + fb[fb < np.sqrt(0.5)] = 0.0 + + return fb + + +class MiniBarkBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=minibark_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +def erb_filterbank( + n_bands: int, + fs: int, + f_min: float, + f_max: float, + n_freqs: int, +) -> Tensor: + # freq bins + A = (1000 * np.log(10)) / (24.7 * 4.37) + all_freqs = torch.linspace(0, fs // 2, n_freqs) + + # calculate mel freq bins + m_min = hz2erb(f_min) + m_max = hz2erb(f_max) + + m_pts = torch.linspace(m_min, m_max, n_bands + 2) + f_pts = (torch.pow(10, (m_pts / A)) - 1) / 0.00437 + + # create filterbank + fb = _create_triangular_filterbank(all_freqs, f_pts) + + fb = fb.T + + first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0] + first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0] + + fb[first_active_band, :first_active_bin] = 1.0 + + return fb + + +class EquivalentRectangularBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=erb_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +if __name__ == "__main__": + import pandas as pd + + band_defs = [] + + for bands in [VocalBandsplitSpecification]: + band_name = bands.__name__.replace("BandsplitSpecification", "") + + mbs = bands(nfft=2048, fs=44100).get_band_specs() + + for i, (f_min, f_max) in enumerate(mbs): + band_defs.append( + {"band": band_name, "band_index": i, "f_min": f_min, "f_max": f_max} + ) + + df = pd.DataFrame(band_defs) + df.to_csv("vox7bands.csv", index=False) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/wrapper.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..2c825151f7c92a83d5eba4aa64c86dfcf160d51e --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/model/bsrnn/wrapper.py @@ -0,0 +1,827 @@ +from typing import Dict, List, Optional, Tuple, Union + +import pytorch_lightning as pl +import torch +from models.bandit.core.model._spectral import _SpectralComponent +from models.bandit.core.model.bsrnn.utils import ( + BarkBandsplitSpecification, + BassBandsplitSpecification, + DrumBandsplitSpecification, + EquivalentRectangularBandsplitSpecification, + MelBandsplitSpecification, + MusicalBandsplitSpecification, + OtherBandsplitSpecification, + TriangularBarkBandsplitSpecification, + VocalBandsplitSpecification, +) +from torch import nn + +from .core import ( + MultiSourceMultiMaskBandSplitCoreConv, + MultiSourceMultiMaskBandSplitCoreRNN, + MultiSourceMultiMaskBandSplitCoreTransformer, + MultiSourceMultiPatchingMaskBandSplitCoreRNN, + SingleMaskBandsplitCoreRNN, + SingleMaskBandsplitCoreTransformer, +) + + +def get_band_specs(band_specs, n_fft, fs, n_bands=None): + if band_specs in ["dnr:speech", "dnr:vox7", "musdb:vocals", "musdb:vox7"]: + bsm = VocalBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs() + freq_weights = None + overlapping_band = False + elif "tribark" in band_specs: + assert n_bands is not None + specs = TriangularBarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands) + bsm = specs.get_band_specs() + freq_weights = specs.get_freq_weights() + overlapping_band = True + elif "bark" in band_specs: + assert n_bands is not None + specs = BarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands) + bsm = specs.get_band_specs() + freq_weights = specs.get_freq_weights() + overlapping_band = True + elif "erb" in band_specs: + assert n_bands is not None + specs = EquivalentRectangularBandsplitSpecification( + nfft=n_fft, fs=fs, n_bands=n_bands + ) + bsm = specs.get_band_specs() + freq_weights = specs.get_freq_weights() + overlapping_band = True + elif "musical" in band_specs: + assert n_bands is not None + specs = MusicalBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands) + bsm = specs.get_band_specs() + freq_weights = specs.get_freq_weights() + overlapping_band = True + elif band_specs == "dnr:mel" or "mel" in band_specs: + assert n_bands is not None + specs = MelBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands) + bsm = specs.get_band_specs() + freq_weights = specs.get_freq_weights() + overlapping_band = True + else: + raise NameError + + return bsm, freq_weights, overlapping_band + + +def get_band_specs_map(band_specs_map, n_fft, fs, n_bands=None): + if band_specs_map == "musdb:all": + bsm = { + "vocals": VocalBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(), + "drums": DrumBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(), + "bass": BassBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(), + "other": OtherBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs(), + } + freq_weights = None + overlapping_band = False + elif band_specs_map == "dnr:vox7": + bsm_, freq_weights, overlapping_band = get_band_specs( + "dnr:speech", n_fft, fs, n_bands + ) + bsm = {"speech": bsm_, "music": bsm_, "effects": bsm_} + elif "dnr:vox7:" in band_specs_map: + stem = band_specs_map.split(":")[-1] + bsm_, freq_weights, overlapping_band = get_band_specs( + "dnr:speech", n_fft, fs, n_bands + ) + bsm = {stem: bsm_} + else: + raise NameError + + return bsm, freq_weights, overlapping_band + + +class BandSplitWrapperBase(pl.LightningModule): + bsrnn: nn.Module + + def __init__(self, **kwargs): + super().__init__() + + +class SingleMaskMultiSourceBandSplitBase(BandSplitWrapperBase, _SpectralComponent): + def __init__( + self, + band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]], + fs: int = 44100, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + ) -> None: + super().__init__( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + ) + + if isinstance(band_specs_map, str): + self.band_specs_map, self.freq_weights, self.overlapping_band = ( + get_band_specs_map(band_specs_map, n_fft, fs, n_bands=n_bands) + ) + + self.stems = list(self.band_specs_map.keys()) + + def forward(self, batch): + audio = batch["audio"] + + with torch.no_grad(): + batch["spectrogram"] = {stem: self.stft(audio[stem]) for stem in audio} + + X = batch["spectrogram"]["mixture"] + length = batch["audio"]["mixture"].shape[-1] + + output = {"spectrogram": {}, "audio": {}} + + for stem, bsrnn in self.bsrnn.items(): + S = bsrnn(X) + s = self.istft(S, length) + output["spectrogram"][stem] = S + output["audio"][stem] = s + + return batch, output + + +class MultiMaskMultiSourceBandSplitBase(BandSplitWrapperBase, _SpectralComponent): + def __init__( + self, + stems: List[str], + band_specs: Union[str, List[Tuple[float, float]]], + fs: int = 44100, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + ) -> None: + super().__init__( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + ) + + if isinstance(band_specs, str): + self.band_specs, self.freq_weights, self.overlapping_band = get_band_specs( + band_specs, n_fft, fs, n_bands + ) + + self.stems = stems + + def forward(self, batch): + # with torch.no_grad(): + audio = batch["audio"] + cond = batch.get("condition", None) + with torch.no_grad(): + batch["spectrogram"] = {stem: self.stft(audio[stem]) for stem in audio} + + X = batch["spectrogram"]["mixture"] + length = batch["audio"]["mixture"].shape[-1] + + output = self.bsrnn(X, cond=cond) + output["audio"] = {} + + for stem, S in output["spectrogram"].items(): + s = self.istft(S, length) + output["audio"][stem] = s + + return batch, output + + +class MultiMaskMultiSourceBandSplitBaseSimple(BandSplitWrapperBase, _SpectralComponent): + def __init__( + self, + stems: List[str], + band_specs: Union[str, List[Tuple[float, float]]], + fs: int = 44100, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + ) -> None: + super().__init__( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + ) + + if isinstance(band_specs, str): + self.band_specs, self.freq_weights, self.overlapping_band = get_band_specs( + band_specs, n_fft, fs, n_bands + ) + + self.stems = stems + + def forward(self, batch): + with torch.no_grad(): + X = self.stft(batch) + length = batch.shape[-1] + output = self.bsrnn(X, cond=None) + res = [] + for stem, S in output["spectrogram"].items(): + s = self.istft(S, length) + res.append(s) + res = torch.stack(res, dim=1) + return res + + +class SingleMaskMultiSourceBandSplitRNN(SingleMaskMultiSourceBandSplitBase): + def __init__( + self, + in_channel: int, + band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]], + fs: int = 44100, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + ) -> None: + super().__init__( + band_specs_map=band_specs_map, + fs=fs, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + ) + + self.bsrnn = nn.ModuleDict( + { + src: SingleMaskBandsplitCoreRNN( + band_specs=specs, + in_channel=in_channel, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + for src, specs in self.band_specs_map.items() + } + ) + + +class SingleMaskMultiSourceBandSplitTransformer(SingleMaskMultiSourceBandSplitBase): + def __init__( + self, + in_channel: int, + band_specs_map: Union[str, Dict[str, List[Tuple[float, float]]]], + fs: int = 44100, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + tf_dropout: float = 0.0, + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + ) -> None: + super().__init__( + band_specs_map=band_specs_map, + fs=fs, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + ) + + self.bsrnn = nn.ModuleDict( + { + src: SingleMaskBandsplitCoreTransformer( + band_specs=specs, + in_channel=in_channel, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + tf_dropout=tf_dropout, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + for src, specs in self.band_specs_map.items() + } + ) + + +class MultiMaskMultiSourceBandSplitRNN(MultiMaskMultiSourceBandSplitBase): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: Union[str, List[Tuple[float, float]]], + fs: int = 44100, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + cond_dim: int = 0, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + use_freq_weights: bool = True, + normalize_input: bool = False, + mult_add_mask: bool = False, + freeze_encoder: bool = False, + ) -> None: + super().__init__( + stems=stems, + band_specs=band_specs, + fs=fs, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + n_bands=n_bands, + ) + + self.bsrnn = MultiSourceMultiMaskBandSplitCoreRNN( + stems=stems, + band_specs=self.band_specs, + in_channel=in_channel, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=self.overlapping_band, + freq_weights=self.freq_weights, + n_freq=n_fft // 2 + 1, + use_freq_weights=use_freq_weights, + mult_add_mask=mult_add_mask, + ) + + self.normalize_input = normalize_input + self.cond_dim = cond_dim + + if freeze_encoder: + for param in self.bsrnn.band_split.parameters(): + param.requires_grad = False + + for param in self.bsrnn.tf_model.parameters(): + param.requires_grad = False + + +class MultiMaskMultiSourceBandSplitRNNSimple(MultiMaskMultiSourceBandSplitBaseSimple): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: Union[str, List[Tuple[float, float]]], + fs: int = 44100, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + cond_dim: int = 0, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + use_freq_weights: bool = True, + normalize_input: bool = False, + mult_add_mask: bool = False, + freeze_encoder: bool = False, + ) -> None: + super().__init__( + stems=stems, + band_specs=band_specs, + fs=fs, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + n_bands=n_bands, + ) + + self.bsrnn = MultiSourceMultiMaskBandSplitCoreRNN( + stems=stems, + band_specs=self.band_specs, + in_channel=in_channel, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=self.overlapping_band, + freq_weights=self.freq_weights, + n_freq=n_fft // 2 + 1, + use_freq_weights=use_freq_weights, + mult_add_mask=mult_add_mask, + ) + + self.normalize_input = normalize_input + self.cond_dim = cond_dim + + if freeze_encoder: + for param in self.bsrnn.band_split.parameters(): + param.requires_grad = False + + for param in self.bsrnn.tf_model.parameters(): + param.requires_grad = False + + +class MultiMaskMultiSourceBandSplitTransformer(MultiMaskMultiSourceBandSplitBase): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: Union[str, List[Tuple[float, float]]], + fs: int = 44100, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + cond_dim: int = 0, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + use_freq_weights: bool = True, + normalize_input: bool = False, + mult_add_mask: bool = False, + ) -> None: + super().__init__( + stems=stems, + band_specs=band_specs, + fs=fs, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + n_bands=n_bands, + ) + + self.bsrnn = MultiSourceMultiMaskBandSplitCoreTransformer( + stems=stems, + band_specs=self.band_specs, + in_channel=in_channel, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=self.overlapping_band, + freq_weights=self.freq_weights, + n_freq=n_fft // 2 + 1, + use_freq_weights=use_freq_weights, + mult_add_mask=mult_add_mask, + ) + + +class MultiMaskMultiSourceBandSplitConv(MultiMaskMultiSourceBandSplitBase): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: Union[str, List[Tuple[float, float]]], + fs: int = 44100, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + cond_dim: int = 0, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + use_freq_weights: bool = True, + normalize_input: bool = False, + mult_add_mask: bool = False, + ) -> None: + super().__init__( + stems=stems, + band_specs=band_specs, + fs=fs, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + n_bands=n_bands, + ) + + self.bsrnn = MultiSourceMultiMaskBandSplitCoreConv( + stems=stems, + band_specs=self.band_specs, + in_channel=in_channel, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + mlp_dim=mlp_dim, + cond_dim=cond_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=self.overlapping_band, + freq_weights=self.freq_weights, + n_freq=n_fft // 2 + 1, + use_freq_weights=use_freq_weights, + mult_add_mask=mult_add_mask, + ) + + +class PatchingMaskMultiSourceBandSplitRNN(MultiMaskMultiSourceBandSplitBase): + def __init__( + self, + in_channel: int, + stems: List[str], + band_specs: Union[str, List[Tuple[float, float]]], + kernel_norm_mlp_version: int = 1, + mask_kernel_freq: int = 3, + mask_kernel_time: int = 3, + conv_kernel_freq: int = 1, + conv_kernel_time: int = 1, + fs: int = 44100, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + n_bands: int = None, + ) -> None: + super().__init__( + stems=stems, + band_specs=band_specs, + fs=fs, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + n_bands=n_bands, + ) + + self.bsrnn = MultiSourceMultiPatchingMaskBandSplitCoreRNN( + stems=stems, + band_specs=self.band_specs, + in_channel=in_channel, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + overlapping_band=self.overlapping_band, + freq_weights=self.freq_weights, + n_freq=n_fft // 2 + 1, + mask_kernel_freq=mask_kernel_freq, + mask_kernel_time=mask_kernel_time, + conv_kernel_freq=conv_kernel_freq, + conv_kernel_time=conv_kernel_time, + kernel_norm_mlp_version=kernel_norm_mlp_version, + ) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/utils/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/core/utils/audio.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/utils/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..2eec9532be2b59815b53a5c307782b4b162dcf24 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/core/utils/audio.py @@ -0,0 +1,406 @@ +from collections import defaultdict +from typing import Callable, Dict, Tuple + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F + + +@torch.jit.script +def merge( + combined: torch.Tensor, + original_batch_size: int, + n_channel: int, + n_chunks: int, + chunk_size: int, +): + combined = torch.reshape( + combined, (original_batch_size, n_chunks, n_channel, chunk_size) + ) + combined = torch.permute(combined, (0, 2, 3, 1)).reshape( + original_batch_size * n_channel, chunk_size, n_chunks + ) + + return combined + + +@torch.jit.script +def unfold( + padded_audio: torch.Tensor, + original_batch_size: int, + n_channel: int, + chunk_size: int, + hop_size: int, +) -> torch.Tensor: + unfolded_input = F.unfold( + padded_audio[:, :, None, :], kernel_size=(1, chunk_size), stride=(1, hop_size) + ) + + _, _, n_chunks = unfolded_input.shape + unfolded_input = unfolded_input.view( + original_batch_size, n_channel, chunk_size, n_chunks + ) + unfolded_input = torch.permute(unfolded_input, (0, 3, 1, 2)).reshape( + original_batch_size * n_chunks, n_channel, chunk_size + ) + + return unfolded_input + + +@torch.jit.script +# @torch.compile +def merge_chunks_all( + combined: torch.Tensor, + original_batch_size: int, + n_channel: int, + n_samples: int, + n_padded_samples: int, + n_chunks: int, + chunk_size: int, + hop_size: int, + edge_frame_pad_sizes: Tuple[int, int], + standard_window: torch.Tensor, + first_window: torch.Tensor, + last_window: torch.Tensor, +): + combined = merge(combined, original_batch_size, n_channel, n_chunks, chunk_size) + + combined = combined * standard_window[:, None].to(combined.device) + + combined = F.fold( + combined.to(torch.float32), + output_size=(1, n_padded_samples), + kernel_size=(1, chunk_size), + stride=(1, hop_size), + ) + + combined = combined.view(original_batch_size, n_channel, n_padded_samples) + + pad_front, pad_back = edge_frame_pad_sizes + combined = combined[..., pad_front:-pad_back] + + combined = combined[..., :n_samples] + + return combined + + # @torch.jit.script + + +def merge_chunks_edge( + combined: torch.Tensor, + original_batch_size: int, + n_channel: int, + n_samples: int, + n_padded_samples: int, + n_chunks: int, + chunk_size: int, + hop_size: int, + edge_frame_pad_sizes: Tuple[int, int], + standard_window: torch.Tensor, + first_window: torch.Tensor, + last_window: torch.Tensor, +): + combined = merge(combined, original_batch_size, n_channel, n_chunks, chunk_size) + + combined[..., 0] = combined[..., 0] * first_window + combined[..., -1] = combined[..., -1] * last_window + combined[..., 1:-1] = combined[..., 1:-1] * standard_window[:, None] + + combined = F.fold( + combined, + output_size=(1, n_padded_samples), + kernel_size=(1, chunk_size), + stride=(1, hop_size), + ) + + combined = combined.view(original_batch_size, n_channel, n_padded_samples) + + combined = combined[..., :n_samples] + + return combined + + +class BaseFader(nn.Module): + def __init__( + self, + chunk_size_second: float, + hop_size_second: float, + fs: int, + fade_edge_frames: bool, + batch_size: int, + ) -> None: + super().__init__() + + self.chunk_size = int(chunk_size_second * fs) + self.hop_size = int(hop_size_second * fs) + self.overlap_size = self.chunk_size - self.hop_size + self.fade_edge_frames = fade_edge_frames + self.batch_size = batch_size + + # @torch.jit.script + def prepare(self, audio): + if self.fade_edge_frames: + audio = F.pad(audio, self.edge_frame_pad_sizes, mode="reflect") + + n_samples = audio.shape[-1] + n_chunks = int(np.ceil((n_samples - self.chunk_size) / self.hop_size) + 1) + + padded_size = (n_chunks - 1) * self.hop_size + self.chunk_size + pad_size = padded_size - n_samples + + padded_audio = F.pad(audio, (0, pad_size)) + + return padded_audio, n_chunks + + def forward( + self, + audio: torch.Tensor, + model_fn: Callable[[torch.Tensor], Dict[str, torch.Tensor]], + ): + original_dtype = audio.dtype + original_device = audio.device + + audio = audio.to("cpu") + + original_batch_size, n_channel, n_samples = audio.shape + padded_audio, n_chunks = self.prepare(audio) + del audio + n_padded_samples = padded_audio.shape[-1] + + if n_channel > 1: + padded_audio = padded_audio.view( + original_batch_size * n_channel, 1, n_padded_samples + ) + + unfolded_input = unfold( + padded_audio, original_batch_size, n_channel, self.chunk_size, self.hop_size + ) + + n_total_chunks, n_channel, chunk_size = unfolded_input.shape + + n_batch = np.ceil(n_total_chunks / self.batch_size).astype(int) + + chunks_in = [ + unfolded_input[b * self.batch_size : (b + 1) * self.batch_size, ...].clone() + for b in range(n_batch) + ] + + all_chunks_out = defaultdict( + lambda: torch.zeros_like(unfolded_input, device="cpu") + ) + + # for b, cin in enumerate(tqdm(chunks_in)): + for b, cin in enumerate(chunks_in): + if torch.allclose(cin, torch.tensor(0.0)): + del cin + continue + + chunks_out = model_fn(cin.to(original_device)) + del cin + for s, c in chunks_out.items(): + all_chunks_out[s][ + b * self.batch_size : (b + 1) * self.batch_size, ... + ] = c.cpu() + del chunks_out + + del unfolded_input + del padded_audio + + if self.fade_edge_frames: + fn = merge_chunks_all + else: + fn = merge_chunks_edge + outputs = {} + + torch.cuda.empty_cache() + + for s, c in all_chunks_out.items(): + combined: torch.Tensor = fn( + c, + original_batch_size, + n_channel, + n_samples, + n_padded_samples, + n_chunks, + self.chunk_size, + self.hop_size, + self.edge_frame_pad_sizes, + self.standard_window, + self.__dict__.get("first_window", self.standard_window), + self.__dict__.get("last_window", self.standard_window), + ) + + outputs[s] = combined.to(dtype=original_dtype, device=original_device) + + return {"audio": outputs} + + # + # def old_forward( + # self, + # audio: torch.Tensor, + # model_fn: Callable[[torch.Tensor], Dict[str, torch.Tensor]], + # ): + # + # n_samples = audio.shape[-1] + # original_batch_size = audio.shape[0] + # + # padded_audio, n_chunks = self.prepare(audio) + # + # ndim = padded_audio.ndim + # broadcaster = [1 for _ in range(ndim - 1)] + [self.chunk_size] + # + # outputs = defaultdict( + # lambda: torch.zeros_like( + # padded_audio, device=audio.device, dtype=torch.float64 + # ) + # ) + # + # all_chunks_out = [] + # len_chunks_in = [] + # + # batch_size_ = int(self.batch_size // original_batch_size) + # for b in range(int(np.ceil(n_chunks / batch_size_))): + # chunks_in = [] + # for j in range(batch_size_): + # i = b * batch_size_ + j + # if i == n_chunks: + # break + # + # start = i * hop_size + # end = start + self.chunk_size + # chunk_in = padded_audio[..., start:end] + # chunks_in.append(chunk_in) + # + # chunks_in = torch.concat(chunks_in, dim=0) + # chunks_out = model_fn(chunks_in) + # all_chunks_out.append(chunks_out) + # len_chunks_in.append(len(chunks_in)) + # + # for b, (chunks_out, lci) in enumerate( + # zip(all_chunks_out, len_chunks_in) + # ): + # for stem in chunks_out: + # for j in range(lci // original_batch_size): + # i = b * batch_size_ + j + # + # if self.fade_edge_frames: + # window = self.standard_window + # else: + # if i == 0: + # window = self.first_window + # elif i == n_chunks - 1: + # window = self.last_window + # else: + # window = self.standard_window + # + # start = i * hop_size + # end = start + self.chunk_size + # + # chunk_out = chunks_out[stem][j * original_batch_size: (j + 1) * original_batch_size, + # ...] + # contrib = window.view(*broadcaster) * chunk_out + # outputs[stem][..., start:end] = ( + # outputs[stem][..., start:end] + contrib + # ) + # + # if self.fade_edge_frames: + # pad_front, pad_back = self.edge_frame_pad_sizes + # outputs = {k: v[..., pad_front:-pad_back] for k, v in + # outputs.items()} + # + # outputs = {k: v[..., :n_samples].to(audio.dtype) for k, v in + # outputs.items()} + # + # return { + # "audio": outputs + # } + + +class LinearFader(BaseFader): + def __init__( + self, + chunk_size_second: float, + hop_size_second: float, + fs: int, + fade_edge_frames: bool = False, + batch_size: int = 1, + ) -> None: + assert hop_size_second >= chunk_size_second / 2 + + super().__init__( + chunk_size_second=chunk_size_second, + hop_size_second=hop_size_second, + fs=fs, + fade_edge_frames=fade_edge_frames, + batch_size=batch_size, + ) + + in_fade = torch.linspace(0.0, 1.0, self.overlap_size + 1)[:-1] + out_fade = torch.linspace(1.0, 0.0, self.overlap_size + 1)[1:] + center_ones = torch.ones(self.chunk_size - 2 * self.overlap_size) + inout_ones = torch.ones(self.overlap_size) + + # using nn.Parameters allows lightning to take care of devices for us + self.register_buffer( + "standard_window", torch.concat([in_fade, center_ones, out_fade]) + ) + + self.fade_edge_frames = fade_edge_frames + self.edge_frame_pad_size = (self.overlap_size, self.overlap_size) + + if not self.fade_edge_frames: + self.first_window = nn.Parameter( + torch.concat([inout_ones, center_ones, out_fade]), requires_grad=False + ) + self.last_window = nn.Parameter( + torch.concat([in_fade, center_ones, inout_ones]), requires_grad=False + ) + + +class OverlapAddFader(BaseFader): + def __init__( + self, + window_type: str, + chunk_size_second: float, + hop_size_second: float, + fs: int, + batch_size: int = 1, + ) -> None: + assert (chunk_size_second / hop_size_second) % 2 == 0 + assert int(chunk_size_second * fs) % 2 == 0 + + super().__init__( + chunk_size_second=chunk_size_second, + hop_size_second=hop_size_second, + fs=fs, + fade_edge_frames=True, + batch_size=batch_size, + ) + + self.hop_multiplier = self.chunk_size / (2 * self.hop_size) + # print(f"hop multiplier: {self.hop_multiplier}") + + self.edge_frame_pad_sizes = (2 * self.overlap_size, 2 * self.overlap_size) + + self.register_buffer( + "standard_window", + torch.windows.__dict__[window_type]( + self.chunk_size, + sym=False, # dtype=torch.float64 + ) + / self.hop_multiplier, + ) + + +if __name__ == "__main__": + import torchaudio as ta + + fs = 44100 + ola = OverlapAddFader("hann", 6.0, 1.0, fs, batch_size=16) + audio_, _ = ta.load( + "$DATA_ROOT/MUSDB18/HQ/canonical/test/BKS - Too Much/vocals.wav" + ) + audio_ = audio_[None, ...] + out = ola(audio_, lambda x: {"stem": x})["audio"]["stem"] + print(torch.allclose(out, audio_)) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit/model_from_config.py b/src/third_party/MusicSourceSeparationTraining/models/bandit/model_from_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6c374a78c7955707d9e420378ba8c0b1624ac4d6 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit/model_from_config.py @@ -0,0 +1,30 @@ +import os.path +import sys + +import torch + +code_path = os.path.dirname(os.path.abspath(__file__)) + "/" +sys.path.append(code_path) + +import yaml +from ml_collections import ConfigDict + +torch.set_float32_matmul_precision("medium") + + +def get_model( + config_path, + weights_path, + device, +): + from models.bandit.core.model import MultiMaskMultiSourceBandSplitRNNSimple + + f = open(config_path) + config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader)) + f.close() + + model = MultiMaskMultiSourceBandSplitRNNSimple(**config.model) + d = torch.load(code_path + "model_bandit_plus_dnr_sdr_11.47.chpt") + model.load_state_dict(d) + model.to(device) + return model, config diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/bandit.py b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/bandit.py new file mode 100644 index 0000000000000000000000000000000000000000..3b32c11a79ae2671231084945748dac78f8e92a5 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/bandit.py @@ -0,0 +1,363 @@ +from typing import Dict, List, Optional + +import pytorch_lightning as pl +import torch +import torchaudio as ta +from torch import nn + +from .bandsplit import BandSplitModule +from .maskestim import OverlappingMaskEstimationModule +from .tfmodel import SeqBandModellingModule +from .utils import MusicalBandsplitSpecification + + +class BaseEndToEndModule(pl.LightningModule): + def __init__( + self, + ) -> None: + super().__init__() + + +class BaseBandit(BaseEndToEndModule): + def __init__( + self, + in_channels: int, + fs: int, + band_type: str = "musical", + n_bands: int = 64, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + ): + super().__init__() + + self.in_channels = in_channels + + self.instantitate_spectral( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + normalized=normalized, + center=center, + pad_mode=pad_mode, + onesided=onesided, + ) + + self.instantiate_bandsplit( + in_channels=in_channels, + band_type=band_type, + n_bands=n_bands, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + n_fft=n_fft, + fs=fs, + ) + + self.instantiate_tf_modelling( + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ) + + def instantitate_spectral( + self, + n_fft: int = 2048, + win_length: Optional[int] = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Optional[Dict] = None, + power: Optional[int] = None, + normalized: bool = True, + center: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + ): + assert power is None + + window_fn = torch.__dict__[window_fn] + + self.stft = ta.transforms.Spectrogram( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + pad_mode=pad_mode, + pad=0, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + normalized=normalized, + center=center, + onesided=onesided, + ) + + self.istft = ta.transforms.InverseSpectrogram( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + pad_mode=pad_mode, + pad=0, + window_fn=window_fn, + wkwargs=wkwargs, + normalized=normalized, + center=center, + onesided=onesided, + ) + + def instantiate_bandsplit( + self, + in_channels: int, + band_type: str = "musical", + n_bands: int = 64, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + emb_dim: int = 128, + n_fft: int = 2048, + fs: int = 44100, + ): + assert band_type == "musical" + + self.band_specs = MusicalBandsplitSpecification( + nfft=n_fft, fs=fs, n_bands=n_bands + ) + + self.band_split = BandSplitModule( + in_channels=in_channels, + band_specs=self.band_specs.get_band_specs(), + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + emb_dim=emb_dim, + ) + + def instantiate_tf_modelling( + self, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + ): + try: + self.tf_model = torch.compile( + SeqBandModellingModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ), + disable=True, + ) + except Exception: + self.tf_model = SeqBandModellingModule( + n_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ) + + def mask(self, x, m): + return x * m + + def forward(self, batch, mode="train"): + # Model takes mono as input we give stereo, so we do process of each channel independently + init_shape = batch.shape + if not isinstance(batch, dict): + mono = batch.view(-1, 1, batch.shape[-1]) + batch = {"mixture": {"audio": mono}} + + with torch.no_grad(): + mixture = batch["mixture"]["audio"] + + x = self.stft(mixture) + batch["mixture"]["spectrogram"] = x + + if "sources" in batch.keys(): + for stem in batch["sources"].keys(): + s = batch["sources"][stem]["audio"] + s = self.stft(s) + batch["sources"][stem]["spectrogram"] = s + + batch = self.separate(batch) + + if 1: + b = [] + for s in self.stems: + # We need to obtain stereo again + r = batch["estimates"][s]["audio"].view( + -1, init_shape[1], init_shape[2] + ) + b.append(r) + # And we need to return back tensor and not independent stems + batch = torch.stack(b, dim=1) + return batch + + def encode(self, batch): + x = batch["mixture"]["spectrogram"] + length = batch["mixture"]["audio"].shape[-1] + + z = self.band_split(x) # (batch, emb_dim, n_band, n_time) + q = self.tf_model(z) # (batch, emb_dim, n_band, n_time) + + return x, q, length + + def separate(self, batch): + raise NotImplementedError + + +class Bandit(BaseBandit): + def __init__( + self, + in_channels: int, + stems: List[str], + band_type: str = "musical", + n_bands: int = 64, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + n_sqm_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + mlp_dim: int = 512, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Dict | None = None, + complex_mask: bool = True, + use_freq_weights: bool = True, + n_fft: int = 2048, + win_length: int | None = 2048, + hop_length: int = 512, + window_fn: str = "hann_window", + wkwargs: Dict | None = None, + power: int | None = None, + center: bool = True, + normalized: bool = True, + pad_mode: str = "constant", + onesided: bool = True, + fs: int = 44100, + stft_precisions="32", + bandsplit_precisions="bf16", + tf_model_precisions="bf16", + mask_estim_precisions="bf16", + ): + super().__init__( + in_channels=in_channels, + band_type=band_type, + n_bands=n_bands, + require_no_overlap=require_no_overlap, + require_no_gap=require_no_gap, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + n_sqm_modules=n_sqm_modules, + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + window_fn=window_fn, + wkwargs=wkwargs, + power=power, + center=center, + normalized=normalized, + pad_mode=pad_mode, + onesided=onesided, + fs=fs, + ) + + self.stems = stems + + self.instantiate_mask_estim( + in_channels=in_channels, + stems=stems, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + n_freq=n_fft // 2 + 1, + use_freq_weights=use_freq_weights, + ) + + def instantiate_mask_estim( + self, + in_channels: int, + stems: List[str], + emb_dim: int, + mlp_dim: int, + hidden_activation: str, + hidden_activation_kwargs: Optional[Dict] = None, + complex_mask: bool = True, + n_freq: Optional[int] = None, + use_freq_weights: bool = False, + ): + if hidden_activation_kwargs is None: + hidden_activation_kwargs = {} + + assert n_freq is not None + + self.mask_estim = nn.ModuleDict( + { + stem: OverlappingMaskEstimationModule( + band_specs=self.band_specs.get_band_specs(), + freq_weights=self.band_specs.get_freq_weights(), + n_freq=n_freq, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + in_channels=in_channels, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + use_freq_weights=use_freq_weights, + ) + for stem in stems + } + ) + + def separate(self, batch): + batch["estimates"] = {} + + x, q, length = self.encode(batch) + + for stem, mem in self.mask_estim.items(): + m = mem(q) + + s = self.mask(x, m.to(x.dtype)) + s = torch.reshape(s, x.shape) + batch["estimates"][stem] = { + "audio": self.istft(s, length), + "spectrogram": s, + } + + return batch diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/bandsplit.py b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/bandsplit.py new file mode 100644 index 0000000000000000000000000000000000000000..b9bfe0f6b1b1ca61301beabbe56ff6f77a3fc848 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/bandsplit.py @@ -0,0 +1,130 @@ +from typing import List, Tuple + +import torch +from torch import nn +from torch.utils.checkpoint import checkpoint_sequential + +from .utils import ( + band_widths_from_specs, + check_no_gap, + check_no_overlap, + check_nonzero_bandwidth, +) + + +class NormFC(nn.Module): + def __init__( + self, + emb_dim: int, + bandwidth: int, + in_channels: int, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + ) -> None: + super().__init__() + + if not treat_channel_as_feature: + raise NotImplementedError + + self.treat_channel_as_feature = treat_channel_as_feature + + if normalize_channel_independently: + raise NotImplementedError + + reim = 2 + + norm = nn.LayerNorm(in_channels * bandwidth * reim) + + fc_in = bandwidth * reim + + if treat_channel_as_feature: + fc_in *= in_channels + else: + assert emb_dim % in_channels == 0 + emb_dim = emb_dim // in_channels + + fc = nn.Linear(fc_in, emb_dim) + + self.combined = nn.Sequential(norm, fc) + + def forward(self, xb): + return checkpoint_sequential(self.combined, 1, xb, use_reentrant=False) + + +class BandSplitModule(nn.Module): + def __init__( + self, + band_specs: List[Tuple[float, float]], + emb_dim: int, + in_channels: int, + require_no_overlap: bool = False, + require_no_gap: bool = True, + normalize_channel_independently: bool = False, + treat_channel_as_feature: bool = True, + ) -> None: + super().__init__() + + check_nonzero_bandwidth(band_specs) + + if require_no_gap: + check_no_gap(band_specs) + + if require_no_overlap: + check_no_overlap(band_specs) + + self.band_specs = band_specs + # list of [fstart, fend) in index. + # Note that fend is exclusive. + self.band_widths = band_widths_from_specs(band_specs) + self.n_bands = len(band_specs) + self.emb_dim = emb_dim + + try: + self.norm_fc_modules = nn.ModuleList( + [ # type: ignore + torch.compile( + NormFC( + emb_dim=emb_dim, + bandwidth=bw, + in_channels=in_channels, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + ), + disable=True, + ) + for bw in self.band_widths + ] + ) + except Exception: + self.norm_fc_modules = nn.ModuleList( + [ # type: ignore + NormFC( + emb_dim=emb_dim, + bandwidth=bw, + in_channels=in_channels, + normalize_channel_independently=normalize_channel_independently, + treat_channel_as_feature=treat_channel_as_feature, + ) + for bw in self.band_widths + ] + ) + + def forward(self, x: torch.Tensor): + # x = complex spectrogram (batch, in_chan, n_freq, n_time) + + batch, in_chan, band_width, n_time = x.shape + + z = torch.zeros( + size=(batch, self.n_bands, n_time, self.emb_dim), device=x.device + ) + + x = torch.permute(x, (0, 3, 1, 2)).contiguous() + + for i, nfm in enumerate(self.norm_fc_modules): + fstart, fend = self.band_specs[i] + xb = x[:, :, :, fstart:fend] + xb = torch.view_as_real(xb) + xb = torch.reshape(xb, (batch, n_time, -1)) + z[:, i, :, :] = nfm(xb) + + return z diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/film.py b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/film.py new file mode 100644 index 0000000000000000000000000000000000000000..9b4aac248b2f04b634bd5c39f203a4bf4f77c180 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/film.py @@ -0,0 +1,21 @@ +from torch import nn + + +class FiLM(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, gamma, beta): + return gamma * x + beta + + +class BTFBroadcastedFiLM(nn.Module): + def __init__(self): + super().__init__() + self.film = FiLM() + + def forward(self, x, gamma, beta): + gamma = gamma[None, None, None, :] + beta = beta[None, None, None, :] + + return self.film(x, gamma, beta) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/maskestim.py b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/maskestim.py new file mode 100644 index 0000000000000000000000000000000000000000..d05923944f63fe6c3ff4719c16d031ae3f33f22a --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/maskestim.py @@ -0,0 +1,281 @@ +from typing import Dict, List, Optional, Tuple, Type + +import torch +from torch import nn +from torch.nn.modules import activation +from torch.utils.checkpoint import checkpoint_sequential + +from .utils import ( + band_widths_from_specs, + check_no_gap, + check_no_overlap, + check_nonzero_bandwidth, +) + + +class BaseNormMLP(nn.Module): + def __init__( + self, + emb_dim: int, + mlp_dim: int, + bandwidth: int, + in_channels: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs=None, + complex_mask: bool = True, + ): + super().__init__() + if hidden_activation_kwargs is None: + hidden_activation_kwargs = {} + self.hidden_activation_kwargs = hidden_activation_kwargs + self.norm = nn.LayerNorm(emb_dim) + self.hidden = nn.Sequential( + nn.Linear(in_features=emb_dim, out_features=mlp_dim), + activation.__dict__[hidden_activation](**self.hidden_activation_kwargs), + ) + + self.bandwidth = bandwidth + self.in_channels = in_channels + + self.complex_mask = complex_mask + self.reim = 2 if complex_mask else 1 + self.glu_mult = 2 + + +class NormMLP(BaseNormMLP): + def __init__( + self, + emb_dim: int, + mlp_dim: int, + bandwidth: int, + in_channels: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs=None, + complex_mask: bool = True, + ) -> None: + super().__init__( + emb_dim=emb_dim, + mlp_dim=mlp_dim, + bandwidth=bandwidth, + in_channels=in_channels, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + + self.output = nn.Sequential( + nn.Linear( + in_features=mlp_dim, + out_features=bandwidth * in_channels * self.reim * 2, + ), + nn.GLU(dim=-1), + ) + + try: + self.combined = torch.compile( + nn.Sequential(self.norm, self.hidden, self.output), disable=True + ) + except Exception: + self.combined = nn.Sequential(self.norm, self.hidden, self.output) + + def reshape_output(self, mb): + # print(mb.shape) + batch, n_time, _ = mb.shape + if self.complex_mask: + mb = mb.reshape( + batch, n_time, self.in_channels, self.bandwidth, self.reim + ).contiguous() + # print(mb.shape) + mb = torch.view_as_complex(mb) # (batch, n_time, in_channels, bandwidth) + else: + mb = mb.reshape(batch, n_time, self.in_channels, self.bandwidth) + + mb = torch.permute(mb, (0, 2, 3, 1)) # (batch, in_channels, bandwidth, n_time) + + return mb + + def forward(self, qb): + # qb = (batch, n_time, emb_dim) + # qb = self.norm(qb) # (batch, n_time, emb_dim) + # qb = self.hidden(qb) # (batch, n_time, mlp_dim) + # mb = self.output(qb) # (batch, n_time, bandwidth * in_channels * reim) + + mb = checkpoint_sequential(self.combined, 2, qb, use_reentrant=False) + mb = self.reshape_output(mb) # (batch, in_channels, bandwidth, n_time) + + return mb + + +class MaskEstimationModuleSuperBase(nn.Module): + pass + + +class MaskEstimationModuleBase(MaskEstimationModuleSuperBase): + def __init__( + self, + band_specs: List[Tuple[float, float]], + emb_dim: int, + mlp_dim: int, + in_channels: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Dict = None, + complex_mask: bool = True, + norm_mlp_cls: Type[nn.Module] = NormMLP, + norm_mlp_kwargs: Dict = None, + ) -> None: + super().__init__() + + self.band_widths = band_widths_from_specs(band_specs) + self.n_bands = len(band_specs) + + if hidden_activation_kwargs is None: + hidden_activation_kwargs = {} + + if norm_mlp_kwargs is None: + norm_mlp_kwargs = {} + + self.norm_mlp = nn.ModuleList( + [ + norm_mlp_cls( + bandwidth=self.band_widths[b], + emb_dim=emb_dim, + mlp_dim=mlp_dim, + in_channels=in_channels, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + **norm_mlp_kwargs, + ) + for b in range(self.n_bands) + ] + ) + + def compute_masks(self, q): + batch, n_bands, n_time, emb_dim = q.shape + + masks = [] + + for b, nmlp in enumerate(self.norm_mlp): + # print(f"maskestim/{b:02d}") + qb = q[:, b, :, :] + mb = nmlp(qb) + masks.append(mb) + + return masks + + def compute_mask(self, q, b): + batch, n_bands, n_time, emb_dim = q.shape + qb = q[:, b, :, :] + mb = self.norm_mlp[b](qb) + return mb + + +class OverlappingMaskEstimationModule(MaskEstimationModuleBase): + def __init__( + self, + in_channels: int, + band_specs: List[Tuple[float, float]], + freq_weights: List[torch.Tensor], + n_freq: int, + emb_dim: int, + mlp_dim: int, + cond_dim: int = 0, + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Dict = None, + complex_mask: bool = True, + norm_mlp_cls: Type[nn.Module] = NormMLP, + norm_mlp_kwargs: Dict = None, + use_freq_weights: bool = False, + ) -> None: + check_nonzero_bandwidth(band_specs) + check_no_gap(band_specs) + + if cond_dim > 0: + raise NotImplementedError + + super().__init__( + band_specs=band_specs, + emb_dim=emb_dim + cond_dim, + mlp_dim=mlp_dim, + in_channels=in_channels, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + norm_mlp_cls=norm_mlp_cls, + norm_mlp_kwargs=norm_mlp_kwargs, + ) + + self.n_freq = n_freq + self.band_specs = band_specs + self.in_channels = in_channels + + if freq_weights is not None and use_freq_weights: + for i, fw in enumerate(freq_weights): + self.register_buffer(f"freq_weights/{i}", fw) + + self.use_freq_weights = use_freq_weights + else: + self.use_freq_weights = False + + def forward(self, q): + # q = (batch, n_bands, n_time, emb_dim) + + batch, n_bands, n_time, emb_dim = q.shape + + masks = torch.zeros( + (batch, self.in_channels, self.n_freq, n_time), + device=q.device, + dtype=torch.complex64, + ) + + for im in range(n_bands): + fstart, fend = self.band_specs[im] + + mask = self.compute_mask(q, im) + + if self.use_freq_weights: + fw = self.get_buffer(f"freq_weights/{im}")[:, None] + mask = mask * fw + masks[:, :, fstart:fend, :] += mask + + return masks + + +class MaskEstimationModule(OverlappingMaskEstimationModule): + def __init__( + self, + band_specs: List[Tuple[float, float]], + emb_dim: int, + mlp_dim: int, + in_channels: Optional[int], + hidden_activation: str = "Tanh", + hidden_activation_kwargs: Dict = None, + complex_mask: bool = True, + **kwargs, + ) -> None: + check_nonzero_bandwidth(band_specs) + check_no_gap(band_specs) + check_no_overlap(band_specs) + super().__init__( + in_channels=in_channels, + band_specs=band_specs, + freq_weights=None, + n_freq=None, + emb_dim=emb_dim, + mlp_dim=mlp_dim, + hidden_activation=hidden_activation, + hidden_activation_kwargs=hidden_activation_kwargs, + complex_mask=complex_mask, + ) + + def forward(self, q, cond=None): + # q = (batch, n_bands, n_time, emb_dim) + + masks = self.compute_masks( + q + ) # [n_bands * (batch, in_channels, bandwidth, n_time)] + + # TODO: currently this requires band specs to have no gap and no overlap + masks = torch.concat(masks, dim=2) # (batch, in_channels, n_freq, n_time) + + return masks diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/tfmodel.py b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/tfmodel.py new file mode 100644 index 0000000000000000000000000000000000000000..21aef03d1f0e814c20db05fe7d14f8019f07713b --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/tfmodel.py @@ -0,0 +1,145 @@ +import warnings + +import torch +import torch.backends.cuda +from torch import nn +from torch.nn.modules import rnn +from torch.utils.checkpoint import checkpoint_sequential + + +class TimeFrequencyModellingModule(nn.Module): + def __init__(self) -> None: + super().__init__() + + +class ResidualRNN(nn.Module): + def __init__( + self, + emb_dim: int, + rnn_dim: int, + bidirectional: bool = True, + rnn_type: str = "LSTM", + use_batch_trick: bool = True, + use_layer_norm: bool = True, + ) -> None: + # n_group is the size of the 2nd dim + super().__init__() + + assert use_layer_norm + assert use_batch_trick + + self.use_layer_norm = use_layer_norm + self.norm = nn.LayerNorm(emb_dim) + self.rnn = rnn.__dict__[rnn_type]( + input_size=emb_dim, + hidden_size=rnn_dim, + num_layers=1, + batch_first=True, + bidirectional=bidirectional, + ) + + self.fc = nn.Linear( + in_features=rnn_dim * (2 if bidirectional else 1), out_features=emb_dim + ) + + self.use_batch_trick = use_batch_trick + if not self.use_batch_trick: + warnings.warn("NOT USING BATCH TRICK IS EXTREMELY SLOW!!") + + def forward(self, z): + # z = (batch, n_uncrossed, n_across, emb_dim) + + z0 = torch.clone(z) + z = self.norm(z) + + batch, n_uncrossed, n_across, emb_dim = z.shape + z = torch.reshape(z, (batch * n_uncrossed, n_across, emb_dim)) + z = self.rnn(z)[0] + z = torch.reshape(z, (batch, n_uncrossed, n_across, -1)) + + z = self.fc(z) # (batch, n_uncrossed, n_across, emb_dim) + + z = z + z0 + + return z + + +class Transpose(nn.Module): + def __init__(self, dim0: int, dim1: int) -> None: + super().__init__() + self.dim0 = dim0 + self.dim1 = dim1 + + def forward(self, z): + return z.transpose(self.dim0, self.dim1) + + +class SeqBandModellingModule(TimeFrequencyModellingModule): + def __init__( + self, + n_modules: int = 12, + emb_dim: int = 128, + rnn_dim: int = 256, + bidirectional: bool = True, + rnn_type: str = "LSTM", + parallel_mode=False, + ) -> None: + super().__init__() + + self.n_modules = n_modules + + if parallel_mode: + self.seqband = nn.ModuleList([]) + for _ in range(n_modules): + self.seqband.append( + nn.ModuleList( + [ + ResidualRNN( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ), + ResidualRNN( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ), + ] + ) + ) + else: + seqband = [] + for _ in range(2 * n_modules): + seqband += [ + ResidualRNN( + emb_dim=emb_dim, + rnn_dim=rnn_dim, + bidirectional=bidirectional, + rnn_type=rnn_type, + ), + Transpose(1, 2), + ] + + self.seqband = nn.Sequential(*seqband) + + self.parallel_mode = parallel_mode + + def forward(self, z): + # z = (batch, n_bands, n_time, emb_dim) + + if self.parallel_mode: + for sbm_pair in self.seqband: + # z: (batch, n_bands, n_time, emb_dim) + sbm_t, sbm_f = sbm_pair[0], sbm_pair[1] + zt = sbm_t(z) # (batch, n_bands, n_time, emb_dim) + zf = sbm_f(z.transpose(1, 2)) # (batch, n_time, n_bands, emb_dim) + z = zt + zf.transpose(1, 2) + else: + z = checkpoint_sequential( + self.seqband, self.n_modules, z, use_reentrant=False + ) + + q = z + return q # (batch, n_bands, n_time, emb_dim) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/utils.py b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ad4eab5d8c5b5396ed717f5b9c365a6900eddd2f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bandit_v2/utils.py @@ -0,0 +1,523 @@ +import os +from abc import abstractmethod +from typing import Callable + +import numpy as np +import torch +from librosa import hz_to_midi, midi_to_hz +from torchaudio import functional as taF + +# from spafe.fbanks import bark_fbanks +# from spafe.utils.converters import erb2hz, hz2bark, hz2erb + + +def band_widths_from_specs(band_specs): + return [e - i for i, e in band_specs] + + +def check_nonzero_bandwidth(band_specs): + # pprint(band_specs) + for fstart, fend in band_specs: + if fend - fstart <= 0: + raise ValueError("Bands cannot be zero-width") + + +def check_no_overlap(band_specs): + fend_prev = -1 + for fstart_curr, fend_curr in band_specs: + if fstart_curr <= fend_prev: + raise ValueError("Bands cannot overlap") + + +def check_no_gap(band_specs): + fstart, _ = band_specs[0] + assert fstart == 0 + + fend_prev = -1 + for fstart_curr, fend_curr in band_specs: + if fstart_curr - fend_prev > 1: + raise ValueError("Bands cannot leave gap") + fend_prev = fend_curr + + +class BandsplitSpecification: + def __init__(self, nfft: int, fs: int) -> None: + self.fs = fs + self.nfft = nfft + self.nyquist = fs / 2 + self.max_index = nfft // 2 + 1 + + self.split500 = self.hertz_to_index(500) + self.split1k = self.hertz_to_index(1000) + self.split2k = self.hertz_to_index(2000) + self.split4k = self.hertz_to_index(4000) + self.split8k = self.hertz_to_index(8000) + self.split16k = self.hertz_to_index(16000) + self.split20k = self.hertz_to_index(20000) + + self.above20k = [(self.split20k, self.max_index)] + self.above16k = [(self.split16k, self.split20k)] + self.above20k + + def index_to_hertz(self, index: int): + return index * self.fs / self.nfft + + def hertz_to_index(self, hz: float, round: bool = True): + index = hz * self.nfft / self.fs + + if round: + index = int(np.round(index)) + + return index + + def get_band_specs_with_bandwidth(self, start_index, end_index, bandwidth_hz): + band_specs = [] + lower = start_index + + while lower < end_index: + upper = int(np.floor(lower + self.hertz_to_index(bandwidth_hz))) + upper = min(upper, end_index) + + band_specs.append((lower, upper)) + lower = upper + + return band_specs + + @abstractmethod + def get_band_specs(self): + raise NotImplementedError + + +class VocalBandsplitSpecification(BandsplitSpecification): + def __init__(self, nfft: int, fs: int, version: str = "7") -> None: + super().__init__(nfft=nfft, fs=fs) + + self.version = version + + def get_band_specs(self): + return getattr(self, f"version{self.version}")() + + @property + def version1(self): + return self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.max_index, bandwidth_hz=1000 + ) + + def version2(self): + below16k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split16k, bandwidth_hz=1000 + ) + below20k = self.get_band_specs_with_bandwidth( + start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000 + ) + + return below16k + below20k + self.above20k + + def version3(self): + below8k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + + return below8k + below16k + self.above16k + + def version4(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + + return below1k + below8k + below16k + self.above16k + + def version5(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split16k, bandwidth_hz=1000 + ) + below20k = self.get_band_specs_with_bandwidth( + start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000 + ) + return below1k + below16k + below20k + self.above20k + + def version6(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + return below1k + below4k + below8k + below16k + self.above16k + + def version7(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split4k, bandwidth_hz=250 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000 + ) + below20k = self.get_band_specs_with_bandwidth( + start_index=self.split16k, end_index=self.split20k, bandwidth_hz=2000 + ) + return below1k + below4k + below8k + below16k + below20k + self.above20k + + +class OtherBandsplitSpecification(VocalBandsplitSpecification): + def __init__(self, nfft: int, fs: int) -> None: + super().__init__(nfft=nfft, fs=fs, version="7") + + +class BassBandsplitSpecification(BandsplitSpecification): + def __init__(self, nfft: int, fs: int, version: str = "7") -> None: + super().__init__(nfft=nfft, fs=fs) + + def get_band_specs(self): + below500 = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split500, bandwidth_hz=50 + ) + below1k = self.get_band_specs_with_bandwidth( + start_index=self.split500, end_index=self.split1k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split4k, bandwidth_hz=500 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=1000 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=2000 + ) + above16k = [(self.split16k, self.max_index)] + + return below500 + below1k + below4k + below8k + below16k + above16k + + +class DrumBandsplitSpecification(BandsplitSpecification): + def __init__(self, nfft: int, fs: int) -> None: + super().__init__(nfft=nfft, fs=fs) + + def get_band_specs(self): + below1k = self.get_band_specs_with_bandwidth( + start_index=0, end_index=self.split1k, bandwidth_hz=50 + ) + below2k = self.get_band_specs_with_bandwidth( + start_index=self.split1k, end_index=self.split2k, bandwidth_hz=100 + ) + below4k = self.get_band_specs_with_bandwidth( + start_index=self.split2k, end_index=self.split4k, bandwidth_hz=250 + ) + below8k = self.get_band_specs_with_bandwidth( + start_index=self.split4k, end_index=self.split8k, bandwidth_hz=500 + ) + below16k = self.get_band_specs_with_bandwidth( + start_index=self.split8k, end_index=self.split16k, bandwidth_hz=1000 + ) + above16k = [(self.split16k, self.max_index)] + + return below1k + below2k + below4k + below8k + below16k + above16k + + +class PerceptualBandsplitSpecification(BandsplitSpecification): + def __init__( + self, + nfft: int, + fs: int, + fbank_fn: Callable[[int, int, float, float, int], torch.Tensor], + n_bands: int, + f_min: float = 0.0, + f_max: float = None, + ) -> None: + super().__init__(nfft=nfft, fs=fs) + self.n_bands = n_bands + if f_max is None: + f_max = fs / 2 + + self.filterbank = fbank_fn(n_bands, fs, f_min, f_max, self.max_index) + + weight_per_bin = torch.sum(self.filterbank, dim=0, keepdim=True) # (1, n_freqs) + normalized_mel_fb = self.filterbank / weight_per_bin # (n_mels, n_freqs) + + freq_weights = [] + band_specs = [] + for i in range(self.n_bands): + active_bins = torch.nonzero(self.filterbank[i, :]).squeeze().tolist() + if isinstance(active_bins, int): + active_bins = (active_bins, active_bins) + if len(active_bins) == 0: + continue + start_index = active_bins[0] + end_index = active_bins[-1] + 1 + band_specs.append((start_index, end_index)) + freq_weights.append(normalized_mel_fb[i, start_index:end_index]) + + self.freq_weights = freq_weights + self.band_specs = band_specs + + def get_band_specs(self): + return self.band_specs + + def get_freq_weights(self): + return self.freq_weights + + def save_to_file(self, dir_path: str) -> None: + os.makedirs(dir_path, exist_ok=True) + + import pickle + + with open(os.path.join(dir_path, "mel_bandsplit_spec.pkl"), "wb") as f: + pickle.dump( + { + "band_specs": self.band_specs, + "freq_weights": self.freq_weights, + "filterbank": self.filterbank, + }, + f, + ) + + +def mel_filterbank(n_bands, fs, f_min, f_max, n_freqs): + fb = taF.melscale_fbanks( + n_mels=n_bands, + sample_rate=fs, + f_min=f_min, + f_max=f_max, + n_freqs=n_freqs, + ).T + + fb[0, 0] = 1.0 + + return fb + + +class MelBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=mel_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +def musical_filterbank(n_bands, fs, f_min, f_max, n_freqs, scale="constant"): + nfft = 2 * (n_freqs - 1) + df = fs / nfft + # init freqs + f_max = f_max or fs / 2 + f_min = f_min or 0 + f_min = fs / nfft + + n_octaves = np.log2(f_max / f_min) + n_octaves_per_band = n_octaves / n_bands + bandwidth_mult = np.power(2.0, n_octaves_per_band) + + low_midi = max(0, hz_to_midi(f_min)) + high_midi = hz_to_midi(f_max) + midi_points = np.linspace(low_midi, high_midi, n_bands) + hz_pts = midi_to_hz(midi_points) + + low_pts = hz_pts / bandwidth_mult + high_pts = hz_pts * bandwidth_mult + + low_bins = np.floor(low_pts / df).astype(int) + high_bins = np.ceil(high_pts / df).astype(int) + + fb = np.zeros((n_bands, n_freqs)) + + for i in range(n_bands): + fb[i, low_bins[i] : high_bins[i] + 1] = 1.0 + + fb[0, : low_bins[0]] = 1.0 + fb[-1, high_bins[-1] + 1 :] = 1.0 + + return torch.as_tensor(fb) + + +class MusicalBandsplitSpecification(PerceptualBandsplitSpecification): + def __init__( + self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None + ) -> None: + super().__init__( + fbank_fn=musical_filterbank, + nfft=nfft, + fs=fs, + n_bands=n_bands, + f_min=f_min, + f_max=f_max, + ) + + +# def bark_filterbank( +# n_bands, fs, f_min, f_max, n_freqs +# ): +# nfft = 2 * (n_freqs -1) +# fb, _ = bark_fbanks.bark_filter_banks( +# nfilts=n_bands, +# nfft=nfft, +# fs=fs, +# low_freq=f_min, +# high_freq=f_max, +# scale="constant" +# ) + +# return torch.as_tensor(fb) + +# class BarkBandsplitSpecification(PerceptualBandsplitSpecification): +# def __init__( +# self, +# nfft: int, +# fs: int, +# n_bands: int, +# f_min: float = 0.0, +# f_max: float = None +# ) -> None: +# super().__init__(fbank_fn=bark_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max) + + +# def triangular_bark_filterbank( +# n_bands, fs, f_min, f_max, n_freqs +# ): + +# all_freqs = torch.linspace(0, fs // 2, n_freqs) + +# # calculate mel freq bins +# m_min = hz2bark(f_min) +# m_max = hz2bark(f_max) + +# m_pts = torch.linspace(m_min, m_max, n_bands + 2) +# f_pts = 600 * torch.sinh(m_pts / 6) + +# # create filterbank +# fb = _create_triangular_filterbank(all_freqs, f_pts) + +# fb = fb.T + +# first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0] +# first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0] + +# fb[first_active_band, :first_active_bin] = 1.0 + +# return fb + +# class TriangularBarkBandsplitSpecification(PerceptualBandsplitSpecification): +# def __init__( +# self, +# nfft: int, +# fs: int, +# n_bands: int, +# f_min: float = 0.0, +# f_max: float = None +# ) -> None: +# super().__init__(fbank_fn=triangular_bark_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max) + + +# def minibark_filterbank( +# n_bands, fs, f_min, f_max, n_freqs +# ): +# fb = bark_filterbank( +# n_bands, +# fs, +# f_min, +# f_max, +# n_freqs +# ) + +# fb[fb < np.sqrt(0.5)] = 0.0 + +# return fb + +# class MiniBarkBandsplitSpecification(PerceptualBandsplitSpecification): +# def __init__( +# self, +# nfft: int, +# fs: int, +# n_bands: int, +# f_min: float = 0.0, +# f_max: float = None +# ) -> None: +# super().__init__(fbank_fn=minibark_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max) + + +# def erb_filterbank( +# n_bands: int, +# fs: int, +# f_min: float, +# f_max: float, +# n_freqs: int, +# ) -> Tensor: +# # freq bins +# A = (1000 * np.log(10)) / (24.7 * 4.37) +# all_freqs = torch.linspace(0, fs // 2, n_freqs) + +# # calculate mel freq bins +# m_min = hz2erb(f_min) +# m_max = hz2erb(f_max) + +# m_pts = torch.linspace(m_min, m_max, n_bands + 2) +# f_pts = (torch.pow(10, (m_pts / A)) - 1)/ 0.00437 + +# # create filterbank +# fb = _create_triangular_filterbank(all_freqs, f_pts) + +# fb = fb.T + + +# first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0] +# first_active_bin = torch.nonzero(fb[first_active_band, :])[0, 0] + +# fb[first_active_band, :first_active_bin] = 1.0 + +# return fb + + +# class EquivalentRectangularBandsplitSpecification(PerceptualBandsplitSpecification): +# def __init__( +# self, +# nfft: int, +# fs: int, +# n_bands: int, +# f_min: float = 0.0, +# f_max: float = None +# ) -> None: +# super().__init__(fbank_fn=erb_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max) + +if __name__ == "__main__": + import pandas as pd + + band_defs = [] + + for bands in [VocalBandsplitSpecification]: + band_name = bands.__name__.replace("BandsplitSpecification", "") + + mbs = bands(nfft=2048, fs=44100).get_band_specs() + + for i, (f_min, f_max) in enumerate(mbs): + band_defs.append( + {"band": band_name, "band_index": i, "f_min": f_min, "f_max": f_max} + ) + + df = pd.DataFrame(band_defs) + df.to_csv("vox7bands.csv", index=False) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_mamba2_code/attend_mamba.py b/src/third_party/MusicSourceSeparationTraining/models/bs_mamba2_code/attend_mamba.py new file mode 100644 index 0000000000000000000000000000000000000000..ea90538ea90f1438ad0bc5e0d24f6addde2bcb85 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_mamba2_code/attend_mamba.py @@ -0,0 +1,128 @@ +from collections import namedtuple +from functools import wraps + +import torch +import torch.nn.functional as F +from packaging import version +from torch import einsum, nn + +# constants + +FlashAttentionConfig = namedtuple( + "FlashAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"] +) + +# helpers + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def once(fn): + called = False + + @wraps(fn) + def inner(x): + nonlocal called + if called: + return + called = True + return fn(x) + + return inner + + +print_once = once(print) + +# main class + + +class Attend(nn.Module): + def __init__(self, dropout=0.0, flash=False, scale=None, idx=None): + super().__init__() + self.scale = scale + self.dropout = dropout + self.attn_dropout = nn.Dropout(dropout) + + self.flash = flash + assert not ( + flash and version.parse(torch.__version__) < version.parse("2.0.0") + ), "in order to use flash attention, you must be using pytorch 2.0 or above" + + # determine efficient attention configs for cuda and cpu + self.cpu_config = FlashAttentionConfig(True, True, True) + self.cuda_config = None + self.idx = idx + + if not torch.cuda.is_available() or not flash: + return + + device_properties = torch.cuda.get_device_properties(torch.device("cuda")) + + if device_properties.major == 8 and device_properties.minor == 0: + print_once( + "A100 GPU detected, using flash attention if input tensor is on cuda" + ) + self.cuda_config = FlashAttentionConfig(True, False, False) + else: + print_once( + "Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda" + ) + self.cuda_config = FlashAttentionConfig(False, True, True) + + def flash_attn(self, q, k, v): + _, heads, q_len, _, k_len, is_cuda, device = ( + *q.shape, + k.shape[-2], + q.is_cuda, + q.device, + ) + + if exists(self.scale): + default_scale = q.shape[-1] ** -0.5 + q = q * (self.scale / default_scale) + + # Check if there is a compatible device for flash attention + config = self.cuda_config if is_cuda else self.cpu_config + + # pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale + with torch.backends.cuda.sdp_kernel(**config._asdict()): + out = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.dropout if self.training else 0.0 + ) + + return out + + def forward(self, q, k, v): + """ + einstein notation + b - batch + h - heads + n, i, j - sequence length (base sequence length, source, target) + d - feature dimension + """ + + q_len, k_len, device = q.shape[-2], k.shape[-2], q.device + + scale = default(self.scale, q.shape[-1] ** -0.5) + + if self.flash: + return self.flash_attn(q, k, v) + + # similarity + + sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale + + # attention + attn = sim.softmax(dim=-1) + attn = self.attn_dropout(attn) + + # aggregate values + out = einsum("b h i j, b h j d -> b h i d", attn, v) + + return out diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_mamba2_code/bs_mamba2.py b/src/third_party/MusicSourceSeparationTraining/models/bs_mamba2_code/bs_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..d0d69aa668fb168c07bb2098054a8367e50fa1f5 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_mamba2_code/bs_mamba2.py @@ -0,0 +1,758 @@ +if __name__ == "__main__": + import os + + gpu_use = "0" + print("GPU use: {}".format(gpu_use)) + os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + os.environ["CUDA_VISIBLE_DEVICES"] = "{}".format(gpu_use) + + +import sys + +sys.path.append("..") +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Tuple +from einops import pack, rearrange, unpack +from mamba_ssm.models.mixer_seq_simple import _init_weights +from mamba_ssm.modules.block import Block +from mamba_ssm.modules.mamba2 import Mamba2 +from mamba_ssm.modules.mamba_simple import Mamba +from mamba_ssm.modules.mlp import GatedMLP +from mamba_ssm.ops.triton.layer_norm import RMSNorm as MRMSNorm +from models.bs_mamba2_code.attend_mamba import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch import nn +from torch.nn import Module, ModuleList + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +class Rearrange(Module): + def __init__(self, pattern): + super().__init__() + self.pattern = pattern + + def forward(self, x): + return rearrange(x, self.pattern) + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class MambaBlock(nn.Module): + def __init__( + self, + in_channels, + n_layer=1, + bidirectional=False, + mamba_type=Mamba, + mlp_cls=nn.Identity, + ): + super(MambaBlock, self).__init__() + self.forward_blocks = nn.ModuleList([]) + for i in range(n_layer): + self.forward_blocks.append( + Block( + in_channels, + mixer_cls=partial(mamba_type, layer_idx=i, expand=4), + norm_cls=partial(MRMSNorm, eps=1e-5), + fused_add_norm=False, + mlp_cls=mlp_cls, + ) + ) + if bidirectional: + self.backward_blocks = nn.ModuleList([]) + for i in range(n_layer): + self.backward_blocks.append( + Block( + in_channels, + mixer_cls=partial(mamba_type, layer_idx=i, expand=4), + norm_cls=partial(MRMSNorm, eps=1e-5), + fused_add_norm=False, + mlp_cls=nn.Identity, + ) + ) + self.linear = nn.ConvTranspose1d( + in_channels=in_channels * 2, + out_channels=in_channels, + kernel_size=1, + stride=1, + ) + self.apply(partial(_init_weights, n_layer=n_layer)) + + def forward(self, input): + for_residual = None + forward_f = input.clone() + for block in self.forward_blocks: + forward_f, for_residual = block( + forward_f, for_residual, inference_params=None + ) + residual = (forward_f + for_residual) if for_residual is not None else forward_f + + if self.backward_blocks is not None: + back_residual = None + backward_f = torch.flip(input, [1]) + for block in self.backward_blocks: + backward_f, back_residual = block( + backward_f, back_residual, inference_params=None + ) + back_residual = ( + (backward_f + back_residual) + if back_residual is not None + else backward_f + ) + + back_residual = torch.flip(back_residual, [1]) + residual = torch.cat([residual, back_residual], -1) + residual = self.linear(residual.transpose(1, 2)).transpose(1, 2).contiguous() + + return residual + input + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + flash=True, + idx=None, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = Attend(flash=flash, dropout=dropout, idx=idx) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + + return self.to_out(out) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + flash_attn=True, + idx=None, + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + self.layers.append( + ModuleList( + [ + Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + idx=idx, + ), + FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout), + ] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + + return self.norm(x) + + +# bandsplit module +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + return torch.stack(outs, dim=-2) + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + + net = [] + dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out) + + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + + if is_last: + continue + + net.append(activation()) + + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + + for dim_in in dim_inputs: + net = [] + + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + + outs = [] + + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + + return torch.cat(outs, dim=-1) + + +DEFAULT_FREQS_PER_BANDS = ( + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 128, + 129, +) + + +class BSMamba2Model(Module): + @beartype + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_module_depth=1, + freq_module_depth=1, + freqs_per_bands=DEFAULT_FREQS_PER_BANDS, + # in the paper, they divide into ~60 bands, test with 1 for starters + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + flash_attn=True, + stft_n_fft=2048, + stft_hop_length=512, + stft_win_length=2048, + stft_normalized=False, + mask_estimator_depth=2, + multi_stft_resolution_loss_weight=1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size=147, + multi_stft_normalized=False, + module_type="transformer", + mamba_gmlp=False, + audio_cfg=None, + ): + super().__init__() + + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.dim = dim + self.stft_n_fft = stft_n_fft + self.stft_hop_length = stft_hop_length + self.stft_win_length = stft_win_length + self.stft_normalized = stft_normalized + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = torch.hann_window + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, normalized=multi_stft_normalized + ) + + self.layers = ModuleList([]) + + self.stft_window_fn = partial(torch.hann_window, 2048) + + if module_type == "transformer" or module_type == "mamba2-roformer": + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn, + norm_output=False, + ) + + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + + for i in range(depth): + if module_type == "transformer": + self.layers.append( + nn.ModuleList( + [ + Transformer( + depth=time_module_depth, + rotary_embed=time_rotary_embed, + **transformer_kwargs, + idx=i, + ), + Transformer( + depth=freq_module_depth, + rotary_embed=freq_rotary_embed, + **transformer_kwargs, + idx=i, + ), + ] + ) + ) + elif module_type == "mamba": + if mamba_gmlp: + mlp_cls = partial(GatedMLP, out_features=dim) + else: + mlp_cls = nn.Identity + + self.layers.append( + nn.ModuleList( + [ + MambaBlock( + in_channels=dim, + n_layer=1, + bidirectional=True, + mamba_type=Mamba, + mlp_cls=mlp_cls, + ), + MambaBlock( + in_channels=dim, + n_layer=1, + bidirectional=True, + mamba_type=Mamba, + mlp_cls=mlp_cls, + ), + ] + ) + ) + elif module_type == "mamba2": + if mamba_gmlp: + mlp_cls = partial(GatedMLP, out_features=dim) + else: + mlp_cls = nn.Identity + + self.layers.append( + nn.ModuleList( + [ + # just make sure d_model * expand / headdim = multiple of 8 + MambaBlock( + in_channels=dim, + n_layer=1, + bidirectional=True, + mamba_type=Mamba2, + mlp_cls=mlp_cls, + ), + MambaBlock( + in_channels=dim, + n_layer=1, + bidirectional=True, + mamba_type=Mamba2, + mlp_cls=mlp_cls, + ), + ] + ) + ) + elif module_type == "mamba2-roformer": + if mamba_gmlp: + mlp_cls = partial(GatedMLP, out_features=dim) + else: + mlp_cls = nn.Identity + + self.layers.append( + nn.ModuleList( + [ + # just make sure d_model * expand / headdim = multiple of 8 + MambaBlock( + in_channels=dim, + n_layer=1, + bidirectional=True, + mamba_type=Mamba2, + mlp_cls=mlp_cls, + ), + Transformer( + depth=freq_module_depth, + rotary_embed=freq_rotary_embed, + **transformer_kwargs, + idx=i, + ), + ] + ) + ) + else: + print("module_type must be mamba or transformer") + exit() + + self.final_norm = RMSNorm(dim) + + assert len(freqs_per_bands) > 1 + + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in freqs_per_bands + ) + + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList([]) + + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + ) + + self.mask_estimators.append(mask_estimator) + + def forward( + self, + x: torch.Tensor, + target=None, + active_stem_ids=None, + return_loss_breakdown=False, + ): + """ + Input and output are T-F complex-valued features. + Input shape: batch_size, n_channels, freq, time] + Output shape: batch_size, n_channels, freq, time] + """ + + raw_audio_length = x.shape[-1] + device = x.device + + if x.ndim == 2: + x = rearrange(x, "b t -> b 1 t") + + x, batch_audio_channel_packed_shape = pack_one(x, "* t") + + self.stft_kwargs = dict( + n_fft=self.stft_n_fft, + hop_length=self.stft_hop_length, + win_length=self.stft_win_length, + normalized=self.stft_normalized, + center=True, + ) + + stft_window = self.stft_window_fn(device=device) + + x = torch.stft(x, **self.stft_kwargs, window=stft_window, return_complex=True) + + x = unpack_one(x, batch_audio_channel_packed_shape, "* f t") + + # x: b c f t + stft_repr = torch.view_as_real(x) + + # x: b c f t 2 + stft_repr = rearrange(stft_repr, "b c f t p -> b (f c) t p") + + x = rearrange(stft_repr, "b f t p -> b t (f p)") + x = self.band_split(x) + + # axial / hierarchical attention + for time_module, freq_module in self.layers: + x = rearrange(x, "b t f d -> b f t d") + x, ps = pack([x], "* t d") + x = time_module(x) + + (x,) = unpack(x, ps, "* t d") + x = rearrange(x, "b f t d -> b t f d") + x, ps = pack([x], "* f d") + + x = freq_module(x) + + (x,) = unpack(x, ps, "* f d") + + x = self.final_norm(x) + + if active_stem_ids is None: + heads = self.mask_estimators + stem_ids = list(range(len(self.mask_estimators))) + else: + heads = [self.mask_estimators[i] for i in active_stem_ids] + stem_ids = active_stem_ids + + num_stems = len(heads) + + mask = torch.stack([fn(x) for fn in heads], dim=1) + + mask = rearrange(mask, "b n t (f p) -> b n f t p", p=2) + + stft_repr = rearrange(stft_repr, "b f t p -> b 1 f t p") + + stft_repr = torch.view_as_complex(stft_repr) + mask = torch.view_as_complex(mask.to(torch.float32)) + + stft_repr = stft_repr * mask + + stft_repr = rearrange( + stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + recon_audio = torch.istft( + stft_repr, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=raw_audio_length, + ) + + recon_audio = rearrange( + recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems + ) + + if not exists(target): + return recon_audio + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[ + ..., : recon_audio.shape[-1] + ] # protect against lost length on istft + + target_sel = target[:, stem_ids] + loss = F.l1_loss(recon_audio, target_sel) + + multi_stft_resolution_loss = 0.0 + + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max( + window_size, self.multi_stft_n_fft + ), # not sure what n_fft is across multi resolution stft + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + + recon_Y = torch.stft( + rearrange(recon_audio, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target_sel, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss( + recon_Y, target_Y + ) + + weighted_multi_resolution_loss = ( + multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + ) + + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) + + +if __name__ == "__main__": + batch_size, n_channels, freq, time = 1, 2, 1025, 800 + in_features = torch.rand( + batch_size, n_channels, freq, time, dtype=torch.cfloat + ).cuda() + cfg = { + "depth": 2, + "dim": 256, + "stereo": True, + "ff_dropout": 0.1, + "attn_dropout": 0.1, + "flash_attn": True, + "module_type": "transformer", + "dim_head": 48, + "num_stems": 3, + } + model = BSModel(**cfg).cuda() + print( + f"Total number of parameters: {sum([p.numel() for p in model.mask_estimators.parameters()])}" + ) + + with torch.cuda.amp.autocast(): + out_features = model(in_features) + + print(f"In shape: {in_features.shape}\nOut shape: {out_features.shape}") + print(f"In dtype: {in_features.dtype}\nOut dtype: {out_features.dtype}") + exit() diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..78c5647922550682b06125eac43d1adf9bdcfcc3 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/__init__.py @@ -0,0 +1,4 @@ +from models.bs_roformer.bs_conformer import BSConformer +from models.bs_roformer.bs_roformer import BSRoformer +from models.bs_roformer.mel_band_conformer import MelBandConformer +from models.bs_roformer.mel_band_roformer import MelBandRoformer diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/attend.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/attend.py new file mode 100644 index 0000000000000000000000000000000000000000..92f861a4a6d86a512ff65a6530673e6cc5a01073 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/attend.py @@ -0,0 +1,142 @@ +import os +from collections import namedtuple +from functools import wraps + +import torch +import torch.nn.functional as F +from packaging import version +from torch import einsum, nn + +# constants + +FlashAttentionConfig = namedtuple( + "FlashAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"] +) + +# helpers + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def once(fn): + called = False + + @wraps(fn) + def inner(x): + nonlocal called + if called: + return + called = True + return fn(x) + + return inner + + +print_once = once(print) + +# main class + + +class Attend(nn.Module): + def __init__(self, dropout=0.0, flash=False, scale=None): + super().__init__() + self.scale = scale + self.dropout = dropout + self.attn_dropout = nn.Dropout(dropout) + + self.flash = flash + assert not ( + flash and version.parse(torch.__version__) < version.parse("2.0.0") + ), "in order to use flash attention, you must be using pytorch 2.0 or above" + + # determine efficient attention configs for cuda and cpu + + self.cpu_config = FlashAttentionConfig(True, True, True) + self.cuda_config = None + + if not torch.cuda.is_available() or not flash: + return + + device_properties = torch.cuda.get_device_properties(torch.device("cuda")) + device_version = version.parse( + f"{device_properties.major}.{device_properties.minor}" + ) + + if device_version >= version.parse("8.0"): + if os.name == "nt": + print_once( + "Windows OS detected, using math or mem efficient attention if input tensor is on cuda" + ) + self.cuda_config = FlashAttentionConfig(False, True, True) + else: + print_once( + "GPU Compute Capability equal or above 8.0, using flash attention if input tensor is on cuda" + ) + self.cuda_config = FlashAttentionConfig(True, False, False) + else: + print_once( + "GPU Compute Capability below 8.0, using math or mem efficient attention if input tensor is on cuda" + ) + self.cuda_config = FlashAttentionConfig(False, True, True) + + def flash_attn(self, q, k, v): + _, heads, q_len, _, k_len, is_cuda, device = ( + *q.shape, + k.shape[-2], + q.is_cuda, + q.device, + ) + + if exists(self.scale): + default_scale = q.shape[-1] ** -0.5 + q = q * (self.scale / default_scale) + + # Check if there is a compatible device for flash attention + + config = self.cuda_config if is_cuda else self.cpu_config + + # pytorch 2.0 flash attn: q, k, v, mask, dropout, softmax_scale + + with torch.backends.cuda.sdp_kernel(**config._asdict()): + out = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.dropout if self.training else 0.0 + ) + + return out + + def forward(self, q, k, v): + """ + einstein notation + b - batch + h - heads + n, i, j - sequence length (base sequence length, source, target) + d - feature dimension + """ + + q_len, k_len, device = q.shape[-2], k.shape[-2], q.device + + scale = default(self.scale, q.shape[-1] ** -0.5) + + if self.flash: + return self.flash_attn(q, k, v) + + # similarity + + sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale + + # attention + + attn = sim.softmax(dim=-1) + attn = self.attn_dropout(attn) + + # aggregate values + + out = einsum("b h i j, b h j d -> b h i d", attn, v) + + return out diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_conformer.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_conformer.py new file mode 100644 index 0000000000000000000000000000000000000000..d104287aae7ca08721e476a0fb6f29b6322418e3 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_conformer.py @@ -0,0 +1,820 @@ +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, Optional, Tuple +from einops import pack, rearrange, unpack +from einops.layers.torch import Rearrange +from models.bs_roformer.attend import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch import nn, tensor +from torch.nn import Module, ModuleList +from torch.utils.checkpoint import checkpoint + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +def pad_at_dim(t, pad, dim=-1, value=0.0): + dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) + zeros = (0, 0) * dims_from_right + return F.pad(t, (*zeros, *pad), value=value) + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +# norm + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# feedforward + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +# attention + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + flash=True, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = Attend(flash=flash, dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +# optional linear attention block + + +class LinearAttention(Module): + """ + https://arxiv.org/abs/2106.09681 (El-Nouby et al.) + """ + + @beartype + def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): + super().__init__() + dim_inner = dim_head * heads + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Sequential( + nn.Linear(dim, dim_inner * 3, bias=False), + Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads), + ) + + self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) + + self.attend = Attend(scale=scale, dropout=dropout, flash=flash) + + self.to_out = nn.Sequential( + Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False) + ) + + def forward(self, x): + x = self.norm(x) + q, k, v = self.to_qkv(x) + q, k = map(l2norm, (q, k)) + q = q * self.temperature.exp() + out = self.attend(q, k, v) + return self.to_out(out) + + +# transformer (kept for optional initial linear blocks) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + flash_attn=True, + linear_attn=False, + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + if linear_attn: + attn = LinearAttention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + flash=flash_attn, + ) + else: + attn = Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + ) + + self.layers.append( + ModuleList( + [attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + return self.norm(x) + + +# conformer + + +class MacaronFF(nn.Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + self.ff = FeedForward(dim=dim, mult=mult, dropout=dropout) + self.scale = 0.5 + + def forward(self, x): + return self.ff(x) * self.scale + + +class ConformerConvModule(nn.Module): + def __init__(self, dim, expansion_factor=2, kernel_size=31, dropout=0.0): + super().__init__() + inner = dim * expansion_factor + assert (kernel_size - 1) % 2 == 0, "kernel_size must be odd" + self.net = nn.Sequential( + RMSNorm(dim), + Rearrange("b n d -> b d n"), + nn.Conv1d(dim, inner * 2, 1), + nn.GLU(dim=1), + nn.Conv1d( + inner, inner, kernel_size, padding=(kernel_size - 1) // 2, groups=inner + ), + nn.BatchNorm1d(inner), + nn.SiLU(inplace=True), + nn.Conv1d(inner, dim, 1), + Rearrange("b d n -> b n d"), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class ConformerBlock(nn.Module): + def __init__( + self, + *, + dim, + heads=8, + dim_head=64, + ff_mult=4, + attn_dropout=0.0, + ff_dropout=0.0, + conv_expansion_factor=2, + conv_kernel_size=31, + rotary_embed=None, + flash_attn=True, + ): + super().__init__() + self.ff1 = MacaronFF(dim=dim, mult=ff_mult, dropout=ff_dropout) + self.attn = Attention( + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + ) + self.conv = ConformerConvModule( + dim=dim, + expansion_factor=conv_expansion_factor, + kernel_size=conv_kernel_size, + dropout=ff_dropout, + ) + self.ff2 = MacaronFF(dim=dim, mult=ff_mult, dropout=ff_dropout) + self.out_norm = RMSNorm(dim) + + def forward(self, x): + x = x + self.ff1(x) + x = x + self.attn(x) + x = x + self.conv(x) + x = x + self.ff2(x) + return self.out_norm(x) + + +class Conformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + rotary_embed=None, + flash_attn=True, + conv_expansion_factor=2, + conv_kernel_size=31, + norm_output=True, + ): + super().__init__() + self.layers = ModuleList( + [ + ConformerBlock( + dim=dim, + heads=heads, + dim_head=dim_head, + ff_mult=ff_mult, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + rotary_embed=rotary_embed, + flash_attn=flash_attn, + ) + for _ in range(depth) + ] + ) + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for block in self.layers: + x = block(x) + return self.norm(x) + + +# bandsplit module + + +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + return torch.stack(outs, dim=-2) + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + net = [] + dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out) + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + if is_last: + continue + net.append(activation()) + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + for dim_in in dim_inputs: + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + outs = [] + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + return torch.cat(outs, dim=-1) + + +# main class + +DEFAULT_FREQS_PER_BANDS = ( + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 128, + 129, +) + + +class BSConformer(Module): + @beartype + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_conformer_depth=2, + freq_conformer_depth=2, + linear_conformer_depth=0, + freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + flash_attn=True, + dim_freqs_in=1025, + stft_n_fft=2048, + stft_hop_length=512, + stft_win_length=2048, + stft_normalized=False, + stft_window_fn: Optional[Callable] = None, + zero_dc=True, + mask_estimator_depth=2, + multi_stft_resolution_loss_weight=1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size=147, + multi_stft_normalized=False, + multi_stft_window_fn: Callable = torch.hann_window, + mlp_expansion_factor=4, + use_torch_checkpoint=False, + skip_connection=False, + # conformer-specific + ff_mult=4, + conv_expansion_factor=2, + conv_kernel_size=31, + ): + super().__init__() + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.use_torch_checkpoint = use_torch_checkpoint + self.skip_connection = skip_connection + + self.layers = ModuleList([]) + + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn, + norm_output=False, + ) + + # rotary embeddings per axis + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + + # build per-depth blocks: optional linear -> time conformer -> freq conformer + for _ in range(depth): + modules = [] + + if linear_conformer_depth > 0: + modules.append( + Transformer( + depth=linear_conformer_depth, + linear_attn=True, + **transformer_kwargs, + ) + ) + + modules.append( + Conformer( + depth=time_conformer_depth, + rotary_embed=time_rotary_embed, + ff_mult=ff_mult, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + **transformer_kwargs, + ) + ) + + modules.append( + Conformer( + depth=freq_conformer_depth, + rotary_embed=freq_rotary_embed, + ff_mult=ff_mult, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + **transformer_kwargs, + ) + ) + + self.layers.append(nn.ModuleList(modules)) + + self.final_norm = RMSNorm(dim) + + # STFT parameters + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized, + ) + self.stft_window_fn = partial( + default(stft_window_fn, torch.hann_window), stft_win_length + ) + + # derive number of freq bins from STFT to validate band split + freqs = torch.stft( + torch.randn(1, 4096), + **self.stft_kwargs, + window=torch.ones(stft_win_length), + return_complex=True, + ).shape[1] + + assert len(freqs_per_bands) > 1 + assert sum(freqs_per_bands) == freqs, ( + f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}" + ) + + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in freqs_per_bands + ) + + # band split and mask estimator + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList([]) + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + mlp_expansion_factor=mlp_expansion_factor, + ) + self.mask_estimators.append(mask_estimator) + + # options and multi-res STFT loss settings + self.zero_dc = zero_dc + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, normalized=multi_stft_normalized + ) + + def forward( + self, raw_audio, target=None, active_stem_ids=None, return_loss_breakdown=False + ): + """ + einops + b - batch + f - freq + t - time + s - audio channel (1 mono, 2 stereo) + n - number of stems + c - complex (2) + d - feature dimension + """ + device = raw_audio.device + x_is_mps = True if device.type == "mps" else False + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, "b t -> b 1 t") + + channels = raw_audio.shape[1] + assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), ( + "stereo True requires 2 channels; mono requires 1 channel" + ) + + # STFT + raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t") + stft_window = self.stft_window_fn(device=device) + + try: + stft_repr = torch.stft( + raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True + ) + except: + stft_repr = torch.stft( + raw_audio.cpu() if x_is_mps else raw_audio, + **self.stft_kwargs, + window=stft_window.cpu() if x_is_mps else stft_window, + return_complex=True, + ).to(device) + + stft_repr = torch.view_as_real(stft_repr) + stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c") + + # merge stereo / mono into frequency leading dimension for band splitting + stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c") + + # flatten complex into "frequency" + x = rearrange(stft_repr, "b f t c -> b t (f c)") + + # band split to per-band features + if self.use_torch_checkpoint: + x = checkpoint(self.band_split, x, use_reentrant=False) + else: + x = self.band_split(x) + + # axial hierarchical encoding: [optional linear] -> time conformer -> freq conformer + store = [None] * len(self.layers) + + for i, block in enumerate(self.layers): + if len(block) == 3: + linear_transformer, time_encoder, freq_encoder = block + + x, ft_ps = pack([x], "b * d") + if self.use_torch_checkpoint: + x = checkpoint(linear_transformer, x, use_reentrant=False) + else: + x = linear_transformer(x) + (x,) = unpack(x, ft_ps, "b * d") + else: + time_encoder, freq_encoder = block + + if self.skip_connection: + for j in range(i): + x = x + store[j] + + # time-axis + x = rearrange(x, "b t f d -> b f t d") + x, ps = pack([x], "* t d") + if self.use_torch_checkpoint: + x = checkpoint(time_encoder, x, use_reentrant=False) + else: + x = time_encoder(x) + (x,) = unpack(x, ps, "* t d") + + # freq-axis + x = rearrange(x, "b f t d -> b t f d") + x, ps = pack([x], "* f d") + if self.use_torch_checkpoint: + x = checkpoint(freq_encoder, x, use_reentrant=False) + else: + x = freq_encoder(x) + (x,) = unpack(x, ps, "* f d") + + if self.skip_connection: + store[i] = x + + x = self.final_norm(x) + + # masks + if active_stem_ids is None: + heads = self.mask_estimators + stem_ids = list(range(len(self.mask_estimators))) + else: + heads = [self.mask_estimators[i] for i in active_stem_ids] + stem_ids = active_stem_ids + + num_stems = len(heads) + + if self.use_torch_checkpoint: + mask = torch.stack( + [checkpoint(fn, x, use_reentrant=False) for fn in heads], dim=1 + ) + else: + mask = torch.stack([fn(x) for fn in heads], dim=1) + + mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2) + + # modulate complex STFT with masks + stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c") + stft_repr = torch.view_as_complex(stft_repr) + mask = torch.view_as_complex(mask) + stft_repr = stft_repr * mask + + # iSTFT + stft_repr = rearrange( + stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + if self.zero_dc: + stft_repr = stft_repr.index_fill(1, tensor(0, device=device), 0.0) + + try: + recon_audio = torch.istft( + stft_repr, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=raw_audio.shape[-1], + ) + except: + recon_audio = torch.istft( + stft_repr.cpu() if x_is_mps else stft_repr, + **self.stft_kwargs, + window=stft_window.cpu() if x_is_mps else stft_window, + return_complex=False, + length=raw_audio.shape[-1], + ).to(device) + + recon_audio = rearrange( + recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems + ) + if num_stems == 1: + recon_audio = rearrange(recon_audio, "b 1 s t -> b s t") + + # optional loss + if not exists(target): + return recon_audio + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[..., : recon_audio.shape[-1]] + target_sel = target[:, stem_ids] + + loss = F.l1_loss(recon_audio, target_sel) + + multi_stft_resolution_loss = 0.0 + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max(window_size, self.multi_stft_n_fft), + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + recon_Y = torch.stft( + rearrange(recon_audio, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target_sel, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss( + recon_Y, target_Y + ) + + weighted_multi_resolution_loss = ( + multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + ) + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_roformer.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_roformer.py new file mode 100644 index 0000000000000000000000000000000000000000..24e28dd5487a86a3862f1bd53b881957ac5efc55 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_roformer.py @@ -0,0 +1,744 @@ +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, Optional, Tuple +from einops import pack, rearrange, unpack +from einops.layers.torch import Rearrange +from models.bs_roformer.attend import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch import nn, tensor +from torch.nn import Module, ModuleList +from torch.utils.checkpoint import checkpoint + +try: + from PoPE_pytorch import PoPE, flash_attn_with_pope + + _HAS_POPE = True +except Exception: + PoPE = None + flash_attn_with_pope = None + _HAS_POPE = False + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +# norm + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# attention + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + flash=True, + pope_embed=None, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + self.pope_embed = pope_embed + assert not (self.rotary_embed is not None and self.pope_embed is not None), ( + "cannot have both rotary and pope embeddings" + ) + + self.attend = Attend(flash=flash, dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + if exists(self.pope_embed): + assert _HAS_POPE, "PoPE requested but PoPE_pytorch is not installed" + out = flash_attn_with_pope( + q, k, v, pos_emb=self.pope_embed(q.shape[-2]), softmax_scale=self.scale + ) + elif exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + out = self.attend(q, k, v) + else: + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +class LinearAttention(Module): + """ + this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al. + """ + + @beartype + def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): + super().__init__() + dim_inner = dim_head * heads + self.norm = RMSNorm(dim) + + self.to_qkv = nn.Sequential( + nn.Linear(dim, dim_inner * 3, bias=False), + Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads), + ) + + self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) + + self.attend = Attend(scale=scale, dropout=dropout, flash=flash) + + self.to_out = nn.Sequential( + Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = self.to_qkv(x) + + q, k = map(l2norm, (q, k)) + q = q * self.temperature.exp() + + out = self.attend(q, k, v) + + return self.to_out(out) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + pope_embed=None, + flash_attn=True, + linear_attn=False, + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + if linear_attn: + attn = LinearAttention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + flash=flash_attn, + ) + else: + attn = Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + pope_embed=pope_embed, + flash=flash_attn, + ) + + self.layers.append( + ModuleList( + [attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + + return self.norm(x) + + +# bandsplit module + + +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + return torch.stack(outs, dim=-2) + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + + net = [] + dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out) + + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + + if is_last: + continue + + net.append(activation()) + + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + + for dim_in in dim_inputs: + net = [] + + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + + outs = [] + + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + + return torch.cat(outs, dim=-1) + + +# main class + +DEFAULT_FREQS_PER_BANDS = ( + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 128, + 129, +) + + +class BSRoformer(Module): + @beartype + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_transformer_depth=2, + freq_transformer_depth=2, + linear_transformer_depth=0, + freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS, + # in the paper, they divide into ~60 bands, test with 1 for starters + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + flash_attn=True, + dim_freqs_in=1025, + stft_n_fft=2048, + stft_hop_length=512, + # 10ms at 44100Hz, from sections 4.1, 4.4 in the paper - @faroit recommends // 2 or // 4 for better reconstruction + stft_win_length=2048, + stft_normalized=False, + stft_window_fn: Optional[Callable] = None, + zero_dc=True, + mask_estimator_depth=2, + multi_stft_resolution_loss_weight=1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size=147, + multi_stft_normalized=False, + multi_stft_window_fn: Callable = torch.hann_window, + mlp_expansion_factor=4, + use_torch_checkpoint=False, + skip_connection=False, + use_pope: bool = False, + ): + super().__init__() + + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.use_torch_checkpoint = use_torch_checkpoint + self.skip_connection = skip_connection + + self.layers = ModuleList([]) + + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn, + norm_output=False, + ) + + if use_pope: + time_pope_embed = PoPE(dim=dim_head, heads=heads) + freq_pope_embed = PoPE(dim=dim_head, heads=heads) + time_rotary_embed = None + freq_rotary_embed = None + else: + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + time_pope_embed = freq_pope_embed = None + + for _ in range(depth): + tran_modules = [] + if linear_transformer_depth > 0: + tran_modules.append( + Transformer( + depth=linear_transformer_depth, + linear_attn=True, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=time_transformer_depth, + rotary_embed=time_rotary_embed, + pope_embed=time_pope_embed, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=freq_transformer_depth, + rotary_embed=freq_rotary_embed, + pope_embed=freq_pope_embed, + **transformer_kwargs, + ) + ) + self.layers.append(nn.ModuleList(tran_modules)) + + self.final_norm = RMSNorm(dim) + + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized, + ) + + self.stft_window_fn = partial( + default(stft_window_fn, torch.hann_window), stft_win_length + ) + + freqs = torch.stft( + torch.randn(1, 4096), + **self.stft_kwargs, + window=torch.ones(stft_win_length), + return_complex=True, + ).shape[1] + + assert len(freqs_per_bands) > 1 + assert sum(freqs_per_bands) == freqs, ( + f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}" + ) + + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in freqs_per_bands + ) + + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList([]) + + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + mlp_expansion_factor=mlp_expansion_factor, + ) + + self.mask_estimators.append(mask_estimator) + + # whether to zero out dc + + self.zero_dc = zero_dc + + # for the multi-resolution stft loss + + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, normalized=multi_stft_normalized + ) + + def forward( + self, raw_audio, target=None, active_stem_ids=None, return_loss_breakdown=False + ): + """ + einops + + b - batch + f - freq + t - time + s - audio channel (1 for mono, 2 for stereo) + n - number of 'stems' + c - complex (2) + d - feature dimension + """ + + device = raw_audio.device + x_is_mps = True if device.type == "mps" else False + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, "b t -> b 1 t") + + channels = raw_audio.shape[1] + assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), ( + "stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2)." + " also need to be False if mono (channel dimension of 1)" + ) + + # to stft + + raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t") + + stft_window = self.stft_window_fn(device=device) + + try: + stft_repr = torch.stft( + raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True + ) + except: + stft_repr = torch.stft( + raw_audio.cpu() if x_is_mps else raw_audio, + **self.stft_kwargs, + window=stft_window.cpu() if x_is_mps else stft_window, + return_complex=True, + ).to(device) + stft_repr = torch.view_as_real(stft_repr) + + stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c") + + # merge stereo / mono into the frequency, with frequency leading dimension, for band splitting + stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c") + + x = rearrange(stft_repr, "b f t c -> b t (f c)") + + if self.use_torch_checkpoint: + x = checkpoint(self.band_split, x, use_reentrant=False) + else: + x = self.band_split(x) + + # axial / hierarchical attention + + store = [None] * len(self.layers) + for i, transformer_block in enumerate(self.layers): + if len(transformer_block) == 3: + linear_transformer, time_transformer, freq_transformer = ( + transformer_block + ) + + x, ft_ps = pack([x], "b * d") + if self.use_torch_checkpoint: + x = checkpoint(linear_transformer, x, use_reentrant=False) + else: + x = linear_transformer(x) + (x,) = unpack(x, ft_ps, "b * d") + else: + time_transformer, freq_transformer = transformer_block + + if self.skip_connection: + # Sum all previous + for j in range(i): + x = x + store[j] + + x = rearrange(x, "b t f d -> b f t d") + x, ps = pack([x], "* t d") + + if self.use_torch_checkpoint: + x = checkpoint(time_transformer, x, use_reentrant=False) + else: + x = time_transformer(x) + + (x,) = unpack(x, ps, "* t d") + x = rearrange(x, "b f t d -> b t f d") + x, ps = pack([x], "* f d") + + if self.use_torch_checkpoint: + x = checkpoint(freq_transformer, x, use_reentrant=False) + else: + x = freq_transformer(x) + + (x,) = unpack(x, ps, "* f d") + + if self.skip_connection: + store[i] = x + + x = self.final_norm(x) + + if active_stem_ids is None: + heads = self.mask_estimators + stem_ids = list(range(len(self.mask_estimators))) + else: + heads = [self.mask_estimators[i] for i in active_stem_ids] + stem_ids = active_stem_ids + + num_stems = len(heads) + + if self.use_torch_checkpoint: + mask = torch.stack( + [checkpoint(fn, x, use_reentrant=False) for fn in heads], dim=1 + ) + else: + mask = torch.stack([fn(x) for fn in heads], dim=1) + mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2) + + # modulate frequency representation + + stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c") + + # complex number multiplication + + stft_repr = torch.view_as_complex(stft_repr) + mask = torch.view_as_complex(mask) + + stft_repr = stft_repr * mask + + # istft + + stft_repr = rearrange( + stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + if self.zero_dc: + # whether to dc filter + stft_repr = stft_repr.index_fill(1, tensor(0, device=device), 0.0) + + try: + recon_audio = torch.istft( + stft_repr, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=raw_audio.shape[-1], + ) + except: + recon_audio = torch.istft( + stft_repr.cpu() if x_is_mps else stft_repr, + **self.stft_kwargs, + window=stft_window.cpu() if x_is_mps else stft_window, + return_complex=False, + length=raw_audio.shape[-1], + ).to(device) + + recon_audio = rearrange( + recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems + ) + + if not exists(target): + return recon_audio + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[ + ..., : recon_audio.shape[-1] + ] # protect against lost length on istft + + target_sel = target[:, stem_ids] + loss = F.l1_loss(recon_audio, target_sel) + + multi_stft_resolution_loss = 0.0 + + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max( + window_size, self.multi_stft_n_fft + ), # not sure what n_fft is across multi resolution stft + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + + recon_Y = torch.stft( + rearrange(recon_audio, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target_sel, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss( + recon_Y, target_Y + ) + + weighted_multi_resolution_loss = ( + multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + ) + + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_roformer_experimental.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_roformer_experimental.py new file mode 100644 index 0000000000000000000000000000000000000000..71d7ac9a33c3ec3478af2747d700ed78e6f6a1d9 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/bs_roformer_experimental.py @@ -0,0 +1,795 @@ +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, Optional, Tuple +from einops import pack, rearrange, unpack +from einops.layers.torch import Rearrange +from hyper_connections import get_init_and_expand_reduce_stream_functions +from models.bs_roformer.attend import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch import nn +from torch.nn import Module, ModuleList +from torch.utils.checkpoint import checkpoint + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +# norm + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# attention + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + flash=True, + learned_value_residual_mix=False, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = Attend(flash=flash, dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + self.to_value_residual_mix = ( + nn.Linear(dim, heads) if learned_value_residual_mix else None + ) + + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x, value_residual=None): + x = self.norm(x) + + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + orig_v = v + + if exists(self.to_value_residual_mix): + mix = self.to_value_residual_mix(x) + mix = rearrange(mix, "b n h -> b h n 1").sigmoid() + + assert exists(value_residual) + v = v.lerp(value_residual, mix) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out), orig_v + + +class LinearAttention(Module): + """ + this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al. + """ + + @beartype + def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): + super().__init__() + dim_inner = dim_head * heads + self.norm = RMSNorm(dim) + + self.to_qkv = nn.Sequential( + nn.Linear(dim, dim_inner * 3, bias=False), + Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads), + ) + + self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) + + self.attend = Attend(scale=scale, dropout=dropout, flash=flash) + + self.to_out = nn.Sequential( + Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = self.to_qkv(x) + + q, k = map(l2norm, (q, k)) + q = q * self.temperature.exp() + + out = self.attend(q, k, v) + + return self.to_out(out) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + flash_attn=True, + linear_attn=False, + add_value_residual=False, + num_residual_streams=1, + ): + super().__init__() + self.layers = ModuleList([]) + + init_hyper_conn, *_ = get_init_and_expand_reduce_stream_functions( + num_residual_streams, disable=num_residual_streams == 1 + ) + + for _ in range(depth): + if linear_attn: + attn = LinearAttention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + flash=flash_attn, + ) + else: + if num_residual_streams != 1: + attn = init_hyper_conn( + dim=dim, + branch=Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + learned_value_residual_mix=add_value_residual, + ), + ) + else: + attn = Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + learned_value_residual_mix=add_value_residual, + ) + if num_residual_streams != 1: + ff = init_hyper_conn( + dim=dim, + branch=FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout), + ) + else: + ff = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout) + + self.layers.append(ModuleList([attn, ff])) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x, value_residual=None): + first_values = None + if value_residual is not None: + for attn, ff in self.layers: + x, next_values = attn(x, value_residual=value_residual) + first_values = default(first_values, next_values) + x = ff(x) + else: + # Compatibility with old weights + for attn, ff in self.layers: + attn_out, next_values = attn(x, value_residual=None) + first_values = default(first_values, next_values) + x = attn_out + x + x = ff(x) + x + + return self.norm(x), first_values + + +# bandsplit module + + +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + return torch.stack(outs, dim=-2) + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + + net = [] + dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out) + + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + + if is_last: + continue + + net.append(activation()) + + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + + for dim_in in dim_inputs: + net = [] + + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + + outs = [] + + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + + return torch.cat(outs, dim=-1) + + +# main class + +DEFAULT_FREQS_PER_BANDS = ( + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 2, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 4, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 12, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 24, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 48, + 128, + 129, +) + + +class BSRoformer(Module): + @beartype + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_transformer_depth=2, + freq_transformer_depth=2, + linear_transformer_depth=0, + freqs_per_bands: Tuple[int, ...] = DEFAULT_FREQS_PER_BANDS, + # in the paper, they divide into ~60 bands, test with 1 for starters + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + flash_attn=True, + dim_freqs_in=1025, + stft_n_fft=2048, + stft_hop_length=512, + # 10ms at 44100Hz, from sections 4.1, 4.4 in the paper - @faroit recommends // 2 or // 4 for better reconstruction + stft_win_length=2048, + stft_normalized=False, + stft_window_fn: Optional[Callable] = None, + mask_estimator_depth=2, + multi_stft_resolution_loss_weight=1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size=147, + multi_stft_normalized=False, + multi_stft_window_fn: Callable = torch.hann_window, + mlp_expansion_factor=4, + use_torch_checkpoint=False, + skip_connection=False, + use_value_residual_learning=False, + num_residual_streams=1, # set to 1. to disable hyper connections (Default in original is 4) + ): + super().__init__() + + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.use_torch_checkpoint = use_torch_checkpoint + self.skip_connection = skip_connection + self.num_residual_streams = num_residual_streams + + _, self.expand_stream, self.reduce_stream = ( + get_init_and_expand_reduce_stream_functions( + num_residual_streams, disable=num_residual_streams == 1 + ) + ) + + self.layers = ModuleList([]) + + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn, + norm_output=False, + num_residual_streams=num_residual_streams, + ) + + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + + for layer_index in range(depth): + if use_value_residual_learning: + is_first = layer_index == 0 + else: + is_first = True + + tran_modules = [] + if linear_transformer_depth > 0: + tran_modules.append( + Transformer( + depth=linear_transformer_depth, + linear_attn=True, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=time_transformer_depth, + rotary_embed=time_rotary_embed, + add_value_residual=not is_first, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=freq_transformer_depth, + rotary_embed=freq_rotary_embed, + add_value_residual=not is_first, + **transformer_kwargs, + ) + ) + self.layers.append(nn.ModuleList(tran_modules)) + + self.final_norm = RMSNorm(dim) + + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized, + ) + + self.stft_window_fn = partial( + default(stft_window_fn, torch.hann_window), stft_win_length + ) + + freqs = torch.stft( + torch.randn(1, 4096), + **self.stft_kwargs, + window=torch.ones(stft_win_length), + return_complex=True, + ).shape[1] + + assert len(freqs_per_bands) > 1 + assert sum(freqs_per_bands) == freqs, ( + f"the number of freqs in the bands must equal {freqs} based on the STFT settings, but got {sum(freqs_per_bands)}" + ) + + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in freqs_per_bands + ) + + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList([]) + + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + mlp_expansion_factor=mlp_expansion_factor, + ) + + self.mask_estimators.append(mask_estimator) + + # for the multi-resolution stft loss + + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, normalized=multi_stft_normalized + ) + + def forward(self, raw_audio, target=None, return_loss_breakdown=False): + """ + einops + + b - batch + f - freq + t - time + s - audio channel (1 for mono, 2 for stereo) + n - number of 'stems' + c - complex (2) + d - feature dimension + """ + + device = raw_audio.device + + # defining whether model is loaded on MPS (MacOS GPU accelerator) + x_is_mps = True if device.type == "mps" else False + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, "b t -> b 1 t") + + channels = raw_audio.shape[1] + assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), ( + "stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)" + ) + + # to stft + + raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t") + + stft_window = self.stft_window_fn(device=device) + + # RuntimeError: FFT operations are only supported on MacOS 14+ + # Since it's tedious to define whether we're on correct MacOS version - simple try-catch is used + try: + stft_repr = torch.stft( + raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True + ) + except: + stft_repr = torch.stft( + raw_audio.cpu() if x_is_mps else raw_audio, + **self.stft_kwargs, + window=stft_window.cpu() if x_is_mps else stft_window, + return_complex=True, + ).to(device) + stft_repr = torch.view_as_real(stft_repr) + + stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c") + + # merge stereo / mono into the frequency, with frequency leading dimension, for band splitting + stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c") + + x = rearrange(stft_repr, "b f t c -> b t (f c)") + + if self.use_torch_checkpoint: + x = checkpoint(self.band_split, x, use_reentrant=False) + else: + x = self.band_split(x) + + # value residuals + + time_v_residual = None + freq_v_residual = None + + # maybe expand residual streams + if self.num_residual_streams != 1: + x = self.expand_stream(x) + + # axial / hierarchical attention + + store = [None] * len(self.layers) + for i, transformer_block in enumerate(self.layers): + if len(transformer_block) == 3: + linear_transformer, time_transformer, freq_transformer = ( + transformer_block + ) + + x, ft_ps = pack([x], "b * d") + if self.use_torch_checkpoint: + x = checkpoint(linear_transformer, x, use_reentrant=False) + else: + x = linear_transformer(x) + (x,) = unpack(x, ft_ps, "b * d") + else: + time_transformer, freq_transformer = transformer_block + + if self.skip_connection: + # Sum all previous + for j in range(i): + x = x + store[j] + + x = rearrange(x, "b t f d -> b f t d") + x, ps = pack([x], "* t d") + + if self.use_torch_checkpoint: + x, next_time_v_residual = checkpoint( + time_transformer, x, time_v_residual, use_reentrant=False + ) + else: + x, next_time_v_residual = time_transformer( + x, value_residual=time_v_residual + ) + time_v_residual = default(time_v_residual, next_time_v_residual) + + (x,) = unpack(x, ps, "* t d") + x = rearrange(x, "b f t d -> b t f d") + x, ps = pack([x], "* f d") + + if self.use_torch_checkpoint: + x, next_freq_v_residual = checkpoint( + freq_transformer, x, freq_v_residual, use_reentrant=False + ) + else: + x, next_freq_v_residual = freq_transformer( + x, value_residual=freq_v_residual + ) + freq_v_residual = default(freq_v_residual, next_freq_v_residual) + + (x,) = unpack(x, ps, "* f d") + + if self.skip_connection: + store[i] = x + + # maybe reduce residual streams + if self.num_residual_streams != 1: + x = self.reduce_stream(x) + + x = self.final_norm(x) + + num_stems = len(self.mask_estimators) + + if self.use_torch_checkpoint: + mask = torch.stack( + [checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators], + dim=1, + ) + else: + mask = torch.stack([fn(x) for fn in self.mask_estimators], dim=1) + mask = rearrange(mask, "b n t (f c) -> b n f t c", c=2) + + # modulate frequency representation + + stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c") + + # complex number multiplication + + stft_repr = torch.view_as_complex(stft_repr) + mask = torch.view_as_complex(mask) + + stft_repr = stft_repr * mask + + # istft + + stft_repr = rearrange( + stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + # same as torch.stft() fix for MacOS MPS above + try: + recon_audio = torch.istft( + stft_repr, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=raw_audio.shape[-1], + ) + except: + recon_audio = torch.istft( + stft_repr.cpu() if x_is_mps else stft_repr, + **self.stft_kwargs, + window=stft_window.cpu() if x_is_mps else stft_window, + return_complex=False, + length=raw_audio.shape[-1], + ).to(device) + + recon_audio = rearrange( + recon_audio, "(b n s) t -> b n s t", s=self.audio_channels, n=num_stems + ) + + if num_stems == 1: + recon_audio = rearrange(recon_audio, "b 1 s t -> b s t") + + # if a target is passed in, calculate loss for learning + + if not exists(target): + return recon_audio + + if self.num_stems > 1: + assert target.ndim == 4 and target.shape[1] == self.num_stems + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[ + ..., : recon_audio.shape[-1] + ] # protect against lost length on istft + + loss = F.l1_loss(recon_audio, target) + + multi_stft_resolution_loss = 0.0 + + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max( + window_size, self.multi_stft_n_fft + ), # not sure what n_fft is across multi resolution stft + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + + recon_Y = torch.stft( + rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs + ) + + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss( + recon_Y, target_Y + ) + + weighted_multi_resolution_loss = ( + multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + ) + + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_conformer.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_conformer.py new file mode 100644 index 0000000000000000000000000000000000000000..c3b63bf555cbda87941c165f0ace04abf97ee622 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_conformer.py @@ -0,0 +1,768 @@ +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, Optional, Tuple +from einops import pack, rearrange, reduce, repeat, unpack +from einops.layers.torch import Rearrange +from librosa import filters +from models.bs_roformer.attend import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch import nn +from torch.nn import Module, ModuleList +from torch.utils.checkpoint import checkpoint + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +def pad_at_dim(t, pad, dim=-1, value=0.0): + dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) + zeros = (0, 0) * dims_from_right + return F.pad(t, (*zeros, *pad), value=value) + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +# norm + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# feedforward + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +# attention + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + flash=True, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = Attend(flash=flash, dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +# optional linear attention block + + +class LinearAttention(Module): + """ + https://arxiv.org/abs/2106.09681 (El-Nouby et al.) + """ + + @beartype + def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): + super().__init__() + dim_inner = dim_head * heads + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Sequential( + nn.Linear(dim, dim_inner * 3, bias=False), + Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads), + ) + + self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) + + self.attend = Attend(scale=scale, dropout=dropout, flash=flash) + + self.to_out = nn.Sequential( + Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False) + ) + + def forward(self, x): + x = self.norm(x) + q, k, v = self.to_qkv(x) + q, k = map(l2norm, (q, k)) + q = q * self.temperature.exp() + out = self.attend(q, k, v) + return self.to_out(out) + + +# transformer (kept for optional initial linear blocks) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + flash_attn=True, + linear_attn=False, + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + if linear_attn: + attn = LinearAttention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + flash=flash_attn, + ) + else: + attn = Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + ) + + self.layers.append( + ModuleList( + [attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + return self.norm(x) + + +# conformer + + +class MacaronFF(nn.Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + self.ff = FeedForward(dim=dim, mult=mult, dropout=dropout) + self.scale = 0.5 + + def forward(self, x): + return self.ff(x) * self.scale + + +class ConformerConvModule(nn.Module): + def __init__(self, dim, expansion_factor=2, kernel_size=31, dropout=0.0): + super().__init__() + inner = dim * expansion_factor + assert (kernel_size - 1) % 2 == 0, "kernel_size must be odd" + self.net = nn.Sequential( + RMSNorm(dim), + Rearrange("b n d -> b d n"), + nn.Conv1d(dim, inner * 2, 1), + nn.GLU(dim=1), + nn.Conv1d( + inner, inner, kernel_size, padding=(kernel_size - 1) // 2, groups=inner + ), + nn.BatchNorm1d(inner), + nn.SiLU(inplace=True), + nn.Conv1d(inner, dim, 1), + Rearrange("b d n -> b n d"), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class ConformerBlock(nn.Module): + def __init__( + self, + *, + dim, + heads=8, + dim_head=64, + ff_mult=4, + attn_dropout=0.0, + ff_dropout=0.0, + conv_expansion_factor=2, + conv_kernel_size=31, + rotary_embed=None, + flash_attn=True, + ): + super().__init__() + self.ff1 = MacaronFF(dim=dim, mult=ff_mult, dropout=ff_dropout) + self.attn = Attention( + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + ) + self.conv = ConformerConvModule( + dim=dim, + expansion_factor=conv_expansion_factor, + kernel_size=conv_kernel_size, + dropout=ff_dropout, + ) + self.ff2 = MacaronFF(dim=dim, mult=ff_mult, dropout=ff_dropout) + self.out_norm = RMSNorm(dim) + + def forward(self, x): + x = x + self.ff1(x) + x = x + self.attn(x) + x = x + self.conv(x) + x = x + self.ff2(x) + return self.out_norm(x) + + +class Conformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + rotary_embed=None, + flash_attn=True, + conv_expansion_factor=2, + conv_kernel_size=31, + norm_output=True, + ): + super().__init__() + self.layers = ModuleList( + [ + ConformerBlock( + dim=dim, + heads=heads, + dim_head=dim_head, + ff_mult=ff_mult, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + rotary_embed=rotary_embed, + flash_attn=flash_attn, + ) + for _ in range(depth) + ] + ) + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for block in self.layers: + x = block(x) + return self.norm(x) + + +# bandsplit module + + +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + return torch.stack(outs, dim=-2) + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + net = [] + dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out) + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + if is_last: + continue + net.append(activation()) + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + for dim_in in dim_inputs: + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + outs = [] + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + return torch.cat(outs, dim=-1) + + +# main class + + +class MelBandConformer(Module): + @beartype + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_conformer_depth=2, + freq_conformer_depth=2, + linear_conformer_depth=0, + num_bands=60, + dim_head=64, + heads=8, + attn_dropout=0.1, + ff_dropout=0.1, + flash_attn=True, + dim_freqs_in=1025, + sample_rate=44100, + stft_n_fft=2048, + stft_hop_length=512, + stft_win_length=2048, + stft_normalized=False, + stft_window_fn: Optional[Callable] = None, + zero_dc=True, + mask_estimator_depth=1, + multi_stft_resolution_loss_weight=1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size=147, + multi_stft_normalized=False, + multi_stft_window_fn: Callable = torch.hann_window, + match_input_audio_length=False, + mlp_expansion_factor=4, + use_torch_checkpoint=False, + skip_connection=False, + # conformer-specific + ff_mult=4, + conv_expansion_factor=2, + conv_kernel_size=31, + ): + super().__init__() + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.use_torch_checkpoint = use_torch_checkpoint + self.skip_connection = skip_connection + + self.layers = ModuleList([]) + + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn, + norm_output=False, + ) + + # rotary embeddings per axis + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + + # build per-depth blocks: optional linear -> time conformer -> freq conformer + for _ in range(depth): + modules = [] + + if linear_conformer_depth > 0: + modules.append( + Transformer( + depth=linear_conformer_depth, + linear_attn=True, + **transformer_kwargs, + ) + ) + + modules.append( + Conformer( + depth=time_conformer_depth, + rotary_embed=time_rotary_embed, + ff_mult=ff_mult, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + **transformer_kwargs, + ) + ) + + modules.append( + Conformer( + depth=freq_conformer_depth, + rotary_embed=freq_rotary_embed, + ff_mult=ff_mult, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + **transformer_kwargs, + ) + ) + + self.layers.append(nn.ModuleList(modules)) + + self.stft_window_fn = partial( + default(stft_window_fn, torch.hann_window), stft_win_length + ) + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized, + ) + + freqs = stft_n_fft // 2 + 1 + + mel_filter_bank_numpy = filters.mel( + sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands + ) + mel_filter_bank = torch.from_numpy(mel_filter_bank_numpy) + + mel_filter_bank[0, 0] = mel_filter_bank[0, 1] * 0.25 + mel_filter_bank[-1, -1] = mel_filter_bank[-1, -2] * 0.25 + + freqs_per_band = mel_filter_bank > 0 + assert freqs_per_band.any(dim=0).all(), ( + "all frequencies need to be covered by all bands for now" + ) + + repeated_freq_indices = repeat(torch.arange(freqs), "f -> b f", b=num_bands) + freq_indices = repeated_freq_indices[freqs_per_band] + + if stereo: + freq_indices = repeat(freq_indices, "f -> f s", s=2) + freq_indices = freq_indices * 2 + torch.arange(2) + freq_indices = rearrange(freq_indices, "f s -> (f s)") + + self.register_buffer("freq_indices", freq_indices, persistent=False) + self.register_buffer("freqs_per_band", freqs_per_band, persistent=False) + + num_freqs_per_band = reduce(freqs_per_band, "b f -> b", "sum") + num_bands_per_freq = reduce(freqs_per_band, "b f -> f", "sum") + + self.register_buffer("num_freqs_per_band", num_freqs_per_band, persistent=False) + self.register_buffer("num_bands_per_freq", num_bands_per_freq, persistent=False) + + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in num_freqs_per_band.tolist() + ) + + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList([]) + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + mlp_expansion_factor=mlp_expansion_factor, + ) + self.mask_estimators.append(mask_estimator) + + self.zero_dc = zero_dc + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + self.multi_stft_hop_length = multi_stft_hop_size + self.multi_stft_normalized = multi_stft_normalized + self.match_input_audio_length = match_input_audio_length + + def forward( + self, raw_audio, target=None, active_stem_ids=None, return_loss_breakdown=False + ): + """ + einops dims: + b - batch + f - freq + t - time + s - audio channel (1 mono, 2 stereo) + n - stems + c - complex (2) + d - feature dim + """ + device = raw_audio.device + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, "b t -> b 1 t") + + batch, channels, raw_audio_length = raw_audio.shape + istft_length = raw_audio_length if self.match_input_audio_length else None + + assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), ( + "stereo True requires 2 channels; mono requires 1 channel" + ) + + raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t") + stft_window = self.stft_window_fn(device=device) + stft_repr = torch.stft( + raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True + ) + stft_repr = torch.view_as_real(stft_repr) + stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c") + + stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c") + + batch_arange = torch.arange(batch, device=device)[..., None] + x = stft_repr[batch_arange, self.freq_indices] + + x = rearrange(x, "b f t c -> b t (f c)") + + if self.use_torch_checkpoint: + x = checkpoint(self.band_split, x, use_reentrant=False) + else: + x = self.band_split(x) + + store = [None] * len(self.layers) + + for i, block in enumerate(self.layers): + if len(block) == 3: + linear_transformer, time_encoder, freq_encoder = block + + x, ft_ps = pack([x], "b * d") + if self.use_torch_checkpoint: + x = checkpoint(linear_transformer, x, use_reentrant=False) + else: + x = linear_transformer(x) + (x,) = unpack(x, ft_ps, "b * d") + else: + time_encoder, freq_encoder = block + + if self.skip_connection: + for j in range(i): + x = x + store[j] + + x = rearrange(x, "b t f d -> b f t d") + x, ps = pack([x], "* t d") + if self.use_torch_checkpoint: + x = checkpoint(time_encoder, x, use_reentrant=False) + else: + x = time_encoder(x) + (x,) = unpack(x, ps, "* t d") + + x = rearrange(x, "b f t d -> b t f d") + x, ps = pack([x], "* f d") + if self.use_torch_checkpoint: + x = checkpoint(freq_encoder, x, use_reentrant=False) + else: + x = freq_encoder(x) + (x,) = unpack(x, ps, "* f d") + + if self.skip_connection: + store[i] = x + + if active_stem_ids is None: + heads = self.mask_estimators + stem_ids = list(range(len(self.mask_estimators))) + else: + heads = [self.mask_estimators[i] for i in active_stem_ids] + stem_ids = active_stem_ids + + num_stems = len(heads) + + if self.use_torch_checkpoint: + masks = torch.stack( + [checkpoint(fn, x, use_reentrant=False) for fn in heads], dim=1 + ) + else: + masks = torch.stack([fn(x) for fn in heads], dim=1) + + masks = rearrange(masks, "b n t (f c) -> b n f t c", c=2) + + stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c") + stft_repr = torch.view_as_complex(stft_repr) + masks = torch.view_as_complex(masks) + masks = masks.type(stft_repr.dtype) + + scatter_indices = repeat( + self.freq_indices, + "f -> b n f t", + b=batch, + n=num_stems, + t=stft_repr.shape[-1], + ) + stft_repr_expanded_stems = repeat(stft_repr, "b 1 ... -> b n ...", n=num_stems) + masks_summed = torch.zeros_like(stft_repr_expanded_stems).scatter_add_( + 2, scatter_indices, masks + ) + + denom = repeat(self.num_bands_per_freq, "f -> (f r) 1", r=channels) + masks_averaged = masks_summed / denom.clamp(min=1e-8) + + stft_repr = stft_repr * masks_averaged + + stft_repr = rearrange( + stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + if self.zero_dc: + stft_repr = stft_repr.index_fill(1, torch.tensor(0, device=device), 0.0) + + recon_audio = torch.istft( + stft_repr, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=istft_length, + ) + + recon_audio = rearrange( + recon_audio, + "(b n s) t -> b n s t", + b=batch, + s=self.audio_channels, + n=num_stems, + ) + if num_stems == 1: + recon_audio = rearrange(recon_audio, "b 1 s t -> b s t") + + if not exists(target): + return recon_audio + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[..., : recon_audio.shape[-1]] + target_sel = target[:, stem_ids] + + loss = F.l1_loss(recon_audio, target_sel) + + multi_stft_resolution_loss = 0.0 + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max(window_size, self.multi_stft_n_fft), + win_length=window_size, + hop_length=max(self.multi_stft_hop_length, window_size // 4), + normalized=self.multi_stft_normalized, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + ) + + recon_Y = torch.stft( + rearrange(recon_audio, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target_sel, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss( + recon_Y, target_Y + ) + + weighted_multi_resolution_loss = ( + multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + ) + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_roformer.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_roformer.py new file mode 100644 index 0000000000000000000000000000000000000000..95fa0df3d1a394c5349eebcd7fc6446fa3d05787 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_roformer.py @@ -0,0 +1,753 @@ +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, Optional, Tuple +from einops import pack, rearrange, reduce, repeat, unpack +from einops.layers.torch import Rearrange +from librosa import filters +from models.bs_roformer.attend import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch import nn, tensor +from torch.nn import Module, ModuleList +from torch.utils.checkpoint import checkpoint + +try: + from PoPE_pytorch import PoPE, flash_attn_with_pope + + _HAS_POPE = True +except Exception: + PoPE = None + flash_attn_with_pope = None + _HAS_POPE = False + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +def pad_at_dim(t, pad, dim=-1, value=0.0): + dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) + zeros = (0, 0) * dims_from_right + return F.pad(t, (*zeros, *pad), value=value) + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +# norm + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# attention + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + flash=True, + pope_embed=None, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + self.pope_embed = pope_embed + assert not (self.rotary_embed is not None and self.pope_embed is not None), ( + "cannot have both rotary and pope embeddings" + ) + + self.attend = Attend(flash=flash, dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + if exists(self.pope_embed): + assert _HAS_POPE, "PoPE requested but PoPE_pytorch is not installed" + out = flash_attn_with_pope( + q, k, v, pos_emb=self.pope_embed(q.shape[-2]), softmax_scale=self.scale + ) + elif exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + out = self.attend(q, k, v) + else: + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +class LinearAttention(Module): + """ + this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al. + """ + + @beartype + def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): + super().__init__() + dim_inner = dim_head * heads + self.norm = RMSNorm(dim) + + self.to_qkv = nn.Sequential( + nn.Linear(dim, dim_inner * 3, bias=False), + Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads), + ) + + self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) + + self.attend = Attend(scale=scale, dropout=dropout, flash=flash) + + self.to_out = nn.Sequential( + Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = self.to_qkv(x) + + q, k = map(l2norm, (q, k)) + q = q * self.temperature.exp() + + out = self.attend(q, k, v) + + return self.to_out(out) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + pope_embed=None, + flash_attn=True, + linear_attn=False, + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + if linear_attn: + attn = LinearAttention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + flash=flash_attn, + ) + else: + attn = Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + pope_embed=pope_embed, + flash=flash_attn, + ) + + self.layers.append( + ModuleList( + [attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + + return self.norm(x) + + +# bandsplit module + + +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + return torch.stack(outs, dim=-2) + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + + net = [] + dims = (dim_in, *((dim_hidden,) * depth), dim_out) + + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + + if is_last: + continue + + net.append(activation()) + + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + + for dim_in in dim_inputs: + net = [] + + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + + outs = [] + + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + + return torch.cat(outs, dim=-1) + + +# main class + + +class MelBandRoformer(Module): + @beartype + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_transformer_depth=2, + freq_transformer_depth=2, + linear_transformer_depth=0, + num_bands=60, + dim_head=64, + heads=8, + attn_dropout=0.1, + ff_dropout=0.1, + flash_attn=True, + dim_freqs_in=1025, + sample_rate=44100, # needed for mel filter bank from librosa + stft_n_fft=2048, + stft_hop_length=512, + # 10ms at 44100Hz, from sections 4.1, 4.4 in the paper - @faroit recommends // 2 or // 4 for better reconstruction + stft_win_length=2048, + stft_normalized=False, + stft_window_fn: Optional[Callable] = None, + zero_dc=True, + mask_estimator_depth=1, + multi_stft_resolution_loss_weight=1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size=147, + multi_stft_normalized=False, + multi_stft_window_fn: Callable = torch.hann_window, + match_input_audio_length=False, # if True, pad output tensor to match length of input tensor + mlp_expansion_factor=4, + use_torch_checkpoint=False, + skip_connection=False, + use_pope: bool = False, + ): + super().__init__() + + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.use_torch_checkpoint = use_torch_checkpoint + self.skip_connection = skip_connection + + self.layers = ModuleList([]) + + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn, + ) + + if use_pope: + time_pope_embed = PoPE(dim=dim_head, heads=heads) + freq_pope_embed = PoPE(dim=dim_head, heads=heads) + time_rotary_embed = None + freq_rotary_embed = None + else: + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + time_pope_embed = freq_pope_embed = None + + for _ in range(depth): + tran_modules = [] + if linear_transformer_depth > 0: + tran_modules.append( + Transformer( + depth=linear_transformer_depth, + linear_attn=True, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=time_transformer_depth, + rotary_embed=time_rotary_embed, + pope_embed=time_pope_embed, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=freq_transformer_depth, + rotary_embed=freq_rotary_embed, + pope_embed=freq_pope_embed, + **transformer_kwargs, + ) + ) + self.layers.append(nn.ModuleList(tran_modules)) + + self.stft_window_fn = partial( + default(stft_window_fn, torch.hann_window), stft_win_length + ) + + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized, + ) + + freqs = torch.stft( + torch.randn(1, 4096), + **self.stft_kwargs, + window=torch.ones(stft_n_fft), + return_complex=True, + ).shape[1] + + # create mel filter bank + # with librosa.filters.mel as in section 2 of paper + + mel_filter_bank_numpy = filters.mel( + sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands + ) + + mel_filter_bank = torch.from_numpy(mel_filter_bank_numpy) + + # for some reason, it doesn't include the first freq? just force a value for now + + mel_filter_bank[0][0] = 1.0 + + # In some systems/envs we get 0.0 instead of ~1.9e-18 in the last position, + # so let's force a positive value + + mel_filter_bank[-1, -1] = 1.0 + + # binary as in paper (then estimated masks are averaged for overlapping regions) + + freqs_per_band = mel_filter_bank > 0 + assert freqs_per_band.any(dim=0).all(), ( + "all frequencies need to be covered by all bands for now" + ) + + repeated_freq_indices = repeat(torch.arange(freqs), "f -> b f", b=num_bands) + freq_indices = repeated_freq_indices[freqs_per_band] + + if stereo: + freq_indices = repeat(freq_indices, "f -> f s", s=2) + freq_indices = freq_indices * 2 + torch.arange(2) + freq_indices = rearrange(freq_indices, "f s -> (f s)") + + self.register_buffer("freq_indices", freq_indices, persistent=False) + self.register_buffer("freqs_per_band", freqs_per_band, persistent=False) + + num_freqs_per_band = reduce(freqs_per_band, "b f -> b", "sum") + num_bands_per_freq = reduce(freqs_per_band, "b f -> f", "sum") + + self.register_buffer("num_freqs_per_band", num_freqs_per_band, persistent=False) + self.register_buffer("num_bands_per_freq", num_bands_per_freq, persistent=False) + + # band split and mask estimator + + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in num_freqs_per_band.tolist() + ) + + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList([]) + + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + mlp_expansion_factor=mlp_expansion_factor, + ) + + self.mask_estimators.append(mask_estimator) + + # whether to zero out dc + + self.zero_dc = zero_dc + + # for the multi-resolution stft loss + + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, normalized=multi_stft_normalized + ) + + self.match_input_audio_length = match_input_audio_length + + def forward( + self, raw_audio, target=None, active_stem_ids=None, return_loss_breakdown=False + ): + """ + einops + + b - batch + f - freq + t - time + s - audio channel (1 for mono, 2 for stereo) + n - number of 'stems' + c - complex (2) + d - feature dimension + """ + + device = raw_audio.device + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, "b t -> b 1 t") + + batch, channels, raw_audio_length = raw_audio.shape + + istft_length = raw_audio_length if self.match_input_audio_length else None + + assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), ( + "stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)" + ) + + # to stft + + raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t") + + stft_window = self.stft_window_fn(device=device) + + stft_repr = torch.stft( + raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True + ) + stft_repr = torch.view_as_real(stft_repr) + + stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c") + + # merge stereo / mono into the frequency, with frequency leading dimension, for band splitting + stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c") + + # index out all frequencies for all frequency ranges across bands ascending in one go + + batch_arange = torch.arange(batch, device=device)[..., None] + + # account for stereo + + x = stft_repr[batch_arange, self.freq_indices] + + # fold the complex (real and imag) into the frequencies dimension + + x = rearrange(x, "b f t c -> b t (f c)") + + if self.use_torch_checkpoint: + x = checkpoint(self.band_split, x, use_reentrant=False) + else: + x = self.band_split(x) + + # axial / hierarchical attention + + store = [None] * len(self.layers) + for i, transformer_block in enumerate(self.layers): + if len(transformer_block) == 3: + linear_transformer, time_transformer, freq_transformer = ( + transformer_block + ) + + x, ft_ps = pack([x], "b * d") + if self.use_torch_checkpoint: + x = checkpoint(linear_transformer, x, use_reentrant=False) + else: + x = linear_transformer(x) + (x,) = unpack(x, ft_ps, "b * d") + else: + time_transformer, freq_transformer = transformer_block + + if self.skip_connection: + # Sum all previous + for j in range(i): + x = x + store[j] + + x = rearrange(x, "b t f d -> b f t d") + x, ps = pack([x], "* t d") + + if self.use_torch_checkpoint: + x = checkpoint(time_transformer, x, use_reentrant=False) + else: + x = time_transformer(x) + + (x,) = unpack(x, ps, "* t d") + x = rearrange(x, "b f t d -> b t f d") + x, ps = pack([x], "* f d") + + if self.use_torch_checkpoint: + x = checkpoint(freq_transformer, x, use_reentrant=False) + else: + x = freq_transformer(x) + + (x,) = unpack(x, ps, "* f d") + + if self.skip_connection: + store[i] = x + + if active_stem_ids is None: + heads = self.mask_estimators + stem_ids = list(range(len(self.mask_estimators))) + else: + heads = [self.mask_estimators[i] for i in active_stem_ids] + stem_ids = active_stem_ids + + num_stems = len(heads) + + if self.use_torch_checkpoint: + masks = torch.stack( + [checkpoint(fn, x, use_reentrant=False) for fn in heads], dim=1 + ) + else: + masks = torch.stack([fn(x) for fn in heads], dim=1) + masks = rearrange(masks, "b n t (f c) -> b n f t c", c=2) + + # modulate frequency representation + + stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c") + + # complex number multiplication + + stft_repr = torch.view_as_complex(stft_repr) + masks = torch.view_as_complex(masks) + + masks = masks.type(stft_repr.dtype) + + # need to average the estimated mask for the overlapped frequencies + + scatter_indices = repeat( + self.freq_indices, + "f -> b n f t", + b=batch, + n=num_stems, + t=stft_repr.shape[-1], + ) + + stft_repr_expanded_stems = repeat(stft_repr, "b 1 ... -> b n ...", n=num_stems) + masks_summed = torch.zeros_like(stft_repr_expanded_stems).scatter_add_( + 2, scatter_indices, masks + ) + + denom = repeat(self.num_bands_per_freq, "f -> (f r) 1", r=channels) + + masks_averaged = masks_summed / denom.clamp(min=1e-8) + + # modulate stft repr with estimated mask + + stft_repr = stft_repr * masks_averaged + + # istft + + stft_repr = rearrange( + stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + if self.zero_dc: + # whether to dc filter + stft_repr = stft_repr.index_fill(1, tensor(0, device=device), 0.0) + + recon_audio = torch.istft( + stft_repr, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=istft_length, + ) + + recon_audio = rearrange( + recon_audio, + "(b n s) t -> b n s t", + b=batch, + s=self.audio_channels, + n=num_stems, + ) + + if num_stems == 1: + recon_audio = rearrange(recon_audio, "b 1 s t -> b s t") + + # if a target is passed in, calculate loss for learning + + if not exists(target): + return recon_audio + + if self.num_stems > 1: + assert target.ndim == 4 and target.shape[1] == self.num_stems + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[ + ..., : recon_audio.shape[-1] + ] # protect against lost length on istft + + target_sel = target[:, stem_ids] + + loss = F.l1_loss(recon_audio, target_sel) + + multi_stft_resolution_loss = 0.0 + + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max( + window_size, self.multi_stft_n_fft + ), # not sure what n_fft is across multi resolution stft + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + + recon_Y = torch.stft( + rearrange(recon_audio, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target_sel, "b n s t -> (b n s) t"), **res_stft_kwargs + ) + + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss( + recon_Y, target_Y + ) + + weighted_multi_resolution_loss = ( + multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + ) + + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_roformer_experimental.py b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_roformer_experimental.py new file mode 100644 index 0000000000000000000000000000000000000000..060d9a867d4308679bae38e7105bd05a9f6f7771 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/bs_roformer/mel_band_roformer_experimental.py @@ -0,0 +1,794 @@ +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, Optional, Tuple +from einops import pack, rearrange, reduce, repeat, unpack +from einops.layers.torch import Rearrange +from hyper_connections import get_init_and_expand_reduce_stream_functions +from librosa import filters +from models.bs_roformer.attend import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch import nn +from torch.nn import Module, ModuleList +from torch.utils.checkpoint import checkpoint + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +def pad_at_dim(t, pad, dim=-1, value=0.0): + dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) + zeros = (0, 0) * dims_from_right + return F.pad(t, (*zeros, *pad), value=value) + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +# norm + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# attention + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + flash=True, + learned_value_residual_mix=False, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = Attend(flash=flash, dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + self.to_value_residual_mix = ( + nn.Linear(dim, heads) if learned_value_residual_mix else None + ) + + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x, value_residual=None): + x = self.norm(x) + + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + orig_v = v + + if exists(self.to_value_residual_mix): + mix = self.to_value_residual_mix(x) + mix = rearrange(mix, "b n h -> b h n 1").sigmoid() + + assert exists(value_residual) + v = v.lerp(value_residual, mix) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out), orig_v + + +class LinearAttention(Module): + """ + this flavor of linear attention proposed in https://arxiv.org/abs/2106.09681 by El-Nouby et al. + """ + + @beartype + def __init__(self, *, dim, dim_head=32, heads=8, scale=8, flash=False, dropout=0.0): + super().__init__() + dim_inner = dim_head * heads + self.norm = RMSNorm(dim) + + self.to_qkv = nn.Sequential( + nn.Linear(dim, dim_inner * 3, bias=False), + Rearrange("b n (qkv h d) -> qkv b h d n", qkv=3, h=heads), + ) + + self.temperature = nn.Parameter(torch.zeros(heads, 1, 1)) + + self.attend = Attend(scale=scale, dropout=dropout, flash=flash) + + self.to_out = nn.Sequential( + Rearrange("b h d n -> b n (h d)"), nn.Linear(dim_inner, dim, bias=False) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = self.to_qkv(x) + + q, k = map(l2norm, (q, k)) + q = q * self.temperature.exp() + + out = self.attend(q, k, v) + + return self.to_out(out) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + flash_attn=True, + linear_attn=False, + add_value_residual=False, + num_residual_streams=1, + ): + super().__init__() + self.layers = ModuleList([]) + + init_hyper_conn, *_ = get_init_and_expand_reduce_stream_functions( + num_residual_streams, disable=num_residual_streams == 1 + ) + + for _ in range(depth): + if linear_attn: + attn = LinearAttention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + flash=flash_attn, + ) + else: + if num_residual_streams != 1: + attn = init_hyper_conn( + dim=dim, + branch=Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + learned_value_residual_mix=add_value_residual, + ), + ) + else: + attn = Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + learned_value_residual_mix=add_value_residual, + ) + + if num_residual_streams != 1: + ff = init_hyper_conn( + dim=dim, + branch=FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout), + ) + else: + ff = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout) + self.layers.append( + ModuleList( + [ + attn, + ff, + ] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x, value_residual=None): + first_values = None + if value_residual is not None: + for attn, ff in self.layers: + x, next_values = attn(x, value_residual=value_residual) + first_values = default(first_values, next_values) + x = ff(x) + else: + # Compatibility with old weights + for attn, ff in self.layers: + attn_out, next_values = attn(x, value_residual=None) + first_values = default(first_values, next_values) + x = attn_out + x + x = ff(x) + x + + return self.norm(x), first_values + + +# bandsplit module + + +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + return torch.stack(outs, dim=-2) + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + + net = [] + dims = (dim_in, *((dim_hidden,) * depth), dim_out) + + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + + if is_last: + continue + + net.append(activation()) + + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + + for dim_in in dim_inputs: + net = [] + + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + + self.to_freqs.append(mlp) + + def forward(self, x): + x = x.unbind(dim=-2) + + outs = [] + + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + + return torch.cat(outs, dim=-1) + + +# main class + + +class MelBandRoformer(Module): + @beartype + def __init__( + self, + dim, + *, + depth, + stereo=False, + num_stems=1, + time_transformer_depth=2, + freq_transformer_depth=2, + linear_transformer_depth=0, + num_bands=60, + dim_head=64, + heads=8, + attn_dropout=0.1, + ff_dropout=0.1, + flash_attn=True, + dim_freqs_in=1025, + sample_rate=44100, # needed for mel filter bank from librosa + stft_n_fft=2048, + stft_hop_length=512, + # 10ms at 44100Hz, from sections 4.1, 4.4 in the paper - @faroit recommends // 2 or // 4 for better reconstruction + stft_win_length=2048, + stft_normalized=False, + stft_window_fn: Optional[Callable] = None, + mask_estimator_depth=1, + multi_stft_resolution_loss_weight=1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size=147, + multi_stft_normalized=False, + multi_stft_window_fn: Callable = torch.hann_window, + match_input_audio_length=False, # if True, pad output tensor to match length of input tensor + mlp_expansion_factor=4, + use_torch_checkpoint=False, + skip_connection=False, + use_value_residual_learning=False, + num_residual_streams=1, # set to 1. to disable hyper connections (Default in original is 4) + ): + super().__init__() + + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.use_torch_checkpoint = use_torch_checkpoint + self.skip_connection = skip_connection + self.num_residual_streams = num_residual_streams + + _, self.expand_stream, self.reduce_stream = ( + get_init_and_expand_reduce_stream_functions( + num_residual_streams, disable=num_residual_streams == 1 + ) + ) + + self.layers = ModuleList([]) + + transformer_kwargs = dict( + dim=dim, + heads=heads, + dim_head=dim_head, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + flash_attn=flash_attn, + num_residual_streams=num_residual_streams, + ) + + time_rotary_embed = RotaryEmbedding(dim=dim_head) + freq_rotary_embed = RotaryEmbedding(dim=dim_head) + + for layer_index in range(depth): + if use_value_residual_learning: + is_first = layer_index == 0 + else: + is_first = True + + tran_modules = [] + if linear_transformer_depth > 0: + tran_modules.append( + Transformer( + depth=linear_transformer_depth, + linear_attn=True, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=time_transformer_depth, + rotary_embed=time_rotary_embed, + add_value_residual=not is_first, + **transformer_kwargs, + ) + ) + tran_modules.append( + Transformer( + depth=freq_transformer_depth, + rotary_embed=freq_rotary_embed, + add_value_residual=not is_first, + **transformer_kwargs, + ) + ) + self.layers.append(nn.ModuleList(tran_modules)) + + self.stft_window_fn = partial( + default(stft_window_fn, torch.hann_window), stft_win_length + ) + + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized, + ) + + freqs = torch.stft( + torch.randn(1, 4096), + **self.stft_kwargs, + window=torch.ones(stft_n_fft), + return_complex=True, + ).shape[1] + + # create mel filter bank + # with librosa.filters.mel as in section 2 of paper + + mel_filter_bank_numpy = filters.mel( + sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands + ) + + mel_filter_bank = torch.from_numpy(mel_filter_bank_numpy) + + # for some reason, it doesn't include the first freq? just force a value for now + + mel_filter_bank[0][0] = 1.0 + + # In some systems/envs we get 0.0 instead of ~1.9e-18 in the last position, + # so let's force a positive value + + mel_filter_bank[-1, -1] = 1.0 + + # binary as in paper (then estimated masks are averaged for overlapping regions) + + freqs_per_band = mel_filter_bank > 0 + assert freqs_per_band.any(dim=0).all(), ( + "all frequencies need to be covered by all bands for now" + ) + + repeated_freq_indices = repeat(torch.arange(freqs), "f -> b f", b=num_bands) + freq_indices = repeated_freq_indices[freqs_per_band] + + if stereo: + freq_indices = repeat(freq_indices, "f -> f s", s=2) + freq_indices = freq_indices * 2 + torch.arange(2) + freq_indices = rearrange(freq_indices, "f s -> (f s)") + + self.register_buffer("freq_indices", freq_indices, persistent=False) + self.register_buffer("freqs_per_band", freqs_per_band, persistent=False) + + num_freqs_per_band = reduce(freqs_per_band, "b f -> b", "sum") + num_bands_per_freq = reduce(freqs_per_band, "b f -> f", "sum") + + self.register_buffer("num_freqs_per_band", num_freqs_per_band, persistent=False) + self.register_buffer("num_bands_per_freq", num_bands_per_freq, persistent=False) + + # band split and mask estimator + + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in num_freqs_per_band.tolist() + ) + + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList([]) + + for _ in range(num_stems): + mask_estimator = MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + mlp_expansion_factor=mlp_expansion_factor, + ) + + self.mask_estimators.append(mask_estimator) + + # for the multi-resolution stft loss + + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, normalized=multi_stft_normalized + ) + + self.match_input_audio_length = match_input_audio_length + + def forward(self, raw_audio, target=None, return_loss_breakdown=False): + """ + einops + + b - batch + f - freq + t - time + s - audio channel (1 for mono, 2 for stereo) + n - number of 'stems' + c - complex (2) + d - feature dimension + """ + + device = raw_audio.device + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, "b t -> b 1 t") + + batch, channels, raw_audio_length = raw_audio.shape + + istft_length = raw_audio_length if self.match_input_audio_length else None + + assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), ( + "stereo needs to be set to True if passing in audio signal that is stereo (channel dimension of 2). also need to be False if mono (channel dimension of 1)" + ) + + # to stft + + raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, "* t") + + stft_window = self.stft_window_fn(device=device) + + stft_repr = torch.stft( + raw_audio, **self.stft_kwargs, window=stft_window, return_complex=True + ) + stft_repr = torch.view_as_real(stft_repr) + + stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, "* f t c") + + # merge stereo / mono into the frequency, with frequency leading dimension, for band splitting + stft_repr = rearrange(stft_repr, "b s f t c -> b (f s) t c") + + # index out all frequencies for all frequency ranges across bands ascending in one go + + batch_arange = torch.arange(batch, device=device)[..., None] + + # account for stereo + + x = stft_repr[batch_arange, self.freq_indices] + + # fold the complex (real and imag) into the frequencies dimension + + x = rearrange(x, "b f t c -> b t (f c)") + + if self.use_torch_checkpoint: + x = checkpoint(self.band_split, x, use_reentrant=False) + else: + x = self.band_split(x) + + # value residuals + time_v_residual = None + freq_v_residual = None + + # maybe expand residual streams + if self.num_residual_streams != 1: + x = self.expand_stream(x) + + # axial / hierarchical attention + + store = [None] * len(self.layers) + for i, transformer_block in enumerate(self.layers): + if len(transformer_block) == 3: + linear_transformer, time_transformer, freq_transformer = ( + transformer_block + ) + + x, ft_ps = pack([x], "b * d") + if self.use_torch_checkpoint: + x = checkpoint(linear_transformer, x, use_reentrant=False) + else: + x = linear_transformer(x) + (x,) = unpack(x, ft_ps, "b * d") + else: + time_transformer, freq_transformer = transformer_block + + if self.skip_connection: + # Sum all previous + for j in range(i): + x = x + store[j] + + x = rearrange(x, "b t f d -> b f t d") + x, ps = pack([x], "* t d") + + if self.use_torch_checkpoint: + x, next_time_v_residual = checkpoint( + time_transformer, x, time_v_residual, use_reentrant=False + ) + else: + x, next_time_v_residual = time_transformer(x, time_v_residual) + time_v_residual = default(time_v_residual, next_time_v_residual) + + (x,) = unpack(x, ps, "* t d") + x = rearrange(x, "b f t d -> b t f d") + x, ps = pack([x], "* f d") + + if self.use_torch_checkpoint: + x, next_freq_v_residual = checkpoint( + freq_transformer, x, freq_v_residual, use_reentrant=False + ) + else: + x, next_freq_v_residual = freq_transformer( + x, value_residual=freq_v_residual + ) + freq_v_residual = default(freq_v_residual, next_freq_v_residual) + + (x,) = unpack(x, ps, "* f d") + + if self.skip_connection: + store[i] = x + + # maybe reduce residual streams + if self.num_residual_streams != 1: + x = self.reduce_stream(x) + + num_stems = len(self.mask_estimators) + if self.use_torch_checkpoint: + masks = torch.stack( + [checkpoint(fn, x, use_reentrant=False) for fn in self.mask_estimators], + dim=1, + ) + else: + masks = torch.stack([fn(x) for fn in self.mask_estimators], dim=1) + masks = rearrange(masks, "b n t (f c) -> b n f t c", c=2) + + # modulate frequency representation + + stft_repr = rearrange(stft_repr, "b f t c -> b 1 f t c") + + # complex number multiplication + + stft_repr = torch.view_as_complex(stft_repr) + masks = torch.view_as_complex(masks) + + masks = masks.type(stft_repr.dtype) + + # need to average the estimated mask for the overlapped frequencies + + scatter_indices = repeat( + self.freq_indices, + "f -> b n f t", + b=batch, + n=num_stems, + t=stft_repr.shape[-1], + ) + + stft_repr_expanded_stems = repeat(stft_repr, "b 1 ... -> b n ...", n=num_stems) + masks_summed = torch.zeros_like(stft_repr_expanded_stems).scatter_add_( + 2, scatter_indices, masks + ) + + denom = repeat(self.num_bands_per_freq, "f -> (f r) 1", r=channels) + + masks_averaged = masks_summed / denom.clamp(min=1e-8) + + # modulate stft repr with estimated mask + + stft_repr = stft_repr * masks_averaged + + # istft + + stft_repr = rearrange( + stft_repr, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + recon_audio = torch.istft( + stft_repr, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=istft_length, + ) + + recon_audio = rearrange( + recon_audio, + "(b n s) t -> b n s t", + b=batch, + s=self.audio_channels, + n=num_stems, + ) + + if num_stems == 1: + recon_audio = rearrange(recon_audio, "b 1 s t -> b s t") + + # if a target is passed in, calculate loss for learning + + if not exists(target): + return recon_audio + + if self.num_stems > 1: + assert target.ndim == 4 and target.shape[1] == self.num_stems + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[ + ..., : recon_audio.shape[-1] + ] # protect against lost length on istft + + loss = F.l1_loss(recon_audio, target) + + multi_stft_resolution_loss = 0.0 + + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max( + window_size, self.multi_stft_n_fft + ), # not sure what n_fft is across multi resolution stft + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + + recon_Y = torch.stft( + rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs + ) + + multi_stft_resolution_loss = multi_stft_resolution_loss + F.l1_loss( + recon_Y, target_Y + ) + + weighted_multi_resolution_loss = ( + multi_stft_resolution_loss * self.multi_stft_resolution_loss_weight + ) + + total_loss = loss + weighted_multi_resolution_loss + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/src/third_party/MusicSourceSeparationTraining/models/conformer_model.py b/src/third_party/MusicSourceSeparationTraining/models/conformer_model.py new file mode 100644 index 0000000000000000000000000000000000000000..b92ea4c92fd8735d1a2b35ab4bb0d85211adae78 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/conformer_model.py @@ -0,0 +1,193 @@ +from typing import Optional + +import torch +import torch.nn as nn +from conformer import Conformer + + +class NeuralModel(nn.Module): + """ + Принимает |X| STFT: (B, C, F, T_spec) и предсказывает комплексные маски + в свернутом виде: (B, 2 * (sources*channels), F, T_spec) + где 2 — это [real, imag]. + """ + + def __init__( + self, + in_channels: int = 2, + sources: int = 2, + freq_bins: int = 2049, + embed_dim: int = 512, + depth: int = 8, + dim_head: int = 64, + heads: int = 8, + ff_mult: int = 4, + conv_expansion_factor: int = 2, + conv_kernel_size: int = 31, + attn_dropout: float = 0.1, + ff_dropout: float = 0.1, + conv_dropout: float = 0.1, + ): + super().__init__() + self.freq_bins = freq_bins + self.in_channels = in_channels + self.sources = sources + self.out_masks = sources * in_channels + self.embed_dim = embed_dim + + self.input_proj_stft = nn.Linear(freq_bins * in_channels, embed_dim) + self.model = Conformer( + dim=embed_dim, + depth=depth, + dim_head=dim_head, + heads=heads, + ff_mult=ff_mult, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + conv_dropout=conv_dropout, + ) + # 2 = [real, imag] + self.output_proj = nn.Linear(embed_dim, freq_bins * self.out_masks * 2) + + def forward(self, x_stft_mag: torch.Tensor) -> torch.Tensor: + """ + x_stft_mag: (B, C, F, T_spec) + returns: (B, 2 * (sources*channels), F, T_spec) — real/imag масок + """ + assert x_stft_mag.dim() == 4, ( + f"Expected (B,C,F,T), got {tuple(x_stft_mag.shape)}" + ) + B, C, F, T_spec = x_stft_mag.shape + # (B, T_spec, C*F) + x_stft_mag = x_stft_mag.permute(0, 3, 1, 2).contiguous().view(B, T_spec, C * F) + + x = self.input_proj_stft(x_stft_mag) # (B, T_spec, E) + x = self.model(x) # (B, T_spec, E) + x = torch.tanh(x) # стабилизируем + x = self.output_proj(x) # (B, T_spec, F * out_masks * 2) + + # back to (B, 2*out_masks, F, T_spec) + x = x.reshape(B, T_spec, self.out_masks * 2, F).permute(0, 2, 3, 1).contiguous() + return x + + +class ConformerMSS(nn.Module): + """ + Совместимо с твоим train: + forward(x: (B, C, T)) -> y_hat: (B, S, C, T) + где S = число источников (sources). + Внутри: STFT -> NeuralModel -> комплексные маски -> iSTFT. + """ + + def __init__( + self, + core: NeuralModel, + n_fft: int = 4096, + hop_length: int = 1024, + win_length: Optional[int] = None, + center: bool = True, + ): + super().__init__() + self.core = core + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length if win_length is not None else n_fft + self.center = center + + window = torch.hann_window(self.win_length) + # окно — буфер, чтобы таскалось на .to(device) + self.register_buffer("window", window, persistent=False) + + # sanity-check: freq_bins у core должен совпадать с n_fft//2 + 1 + expected_bins = n_fft // 2 + 1 + assert core.freq_bins == expected_bins, ( + f"NeuralModel.freq_bins={core.freq_bins} != n_fft//2+1={expected_bins}. " + f"Поставь freq_bins={expected_bins} при создании core." + ) + + def _stft(self, x: torch.Tensor) -> torch.Tensor: + """ + x: (B, C, T) -> spec: complex (B, C, F, TT) + """ + assert x.dim() == 3, f"Expected (B,C,T), got {tuple(x.shape)}" + B, C, T = x.shape + x_bc_t = x.reshape(B * C, T) + spec = torch.stft( + x_bc_t, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + window=self.window.to(x.device), + center=self.center, + return_complex=True, + ) # (B*C, F, TT) + F, TT = spec.shape[-2], spec.shape[-1] + spec = spec.reshape(B, C, F, TT) + return spec + + def _istft(self, spec: torch.Tensor, length: int) -> torch.Tensor: + """ + spec: complex (B, C, F, TT) -> audio: (B, C, T) + """ + B, C, F, TT = spec.shape + spec_bc = spec.reshape(B * C, F, TT) + y_bc_t = torch.istft( + spec_bc, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + window=self.window.to(spec.device), + center=self.center, + length=length, + ) + return y_bc_t.reshape(B, C, -1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + x: (B, C, T) (микс в волне) + returns y_hat: (B, S, C, T) — предсказанные источники в волне + """ + B, C, T = x.shape + # 1) STFT + mix_spec = self._stft(x) # (B, C, F, TT) + mix_mag = mix_spec.abs() # (B, C, F, TT) + + # 2) Прогон через core -> real/imag масок + mask_ri = self.core(mix_mag) # (B, 2*(S*C), F, TT2) + _, two_sc, F, TT2 = mask_ri.shape + + S = self.core.sources + assert two_sc == 2 * (S * C), ( + f"core вернул {two_sc} каналов масок, ожидалось {2 * (S * C)} " + f"(2*[real/imag]*[sources*channels]). Проверь in_channels/sources." + ) + + # 3) Синхронизация по времени (если вдруг TT != TT2) + TT = mix_spec.shape[-1] + TT_min = min(TT, TT2) + if TT != TT_min: + mix_spec = mix_spec[..., :TT_min] + if TT2 != TT_min: + mask_ri = mask_ri[..., :TT_min] + TT = TT_min + # теперь у обоих время = TT + + # 4) Преобразуем к (B, 2, S, C, F, TT) + mask_ri = mask_ri.view(B, 2, S, C, F, TT).contiguous() + mask_real = mask_ri[:, 0] # (B, S, C, F, TT) + mask_imag = mask_ri[:, 1] # (B, S, C, F, TT) + masks_c = torch.complex(mask_real, mask_imag) + + # 5) Применяем маски к комплексному спектру микса + mix_spec_bc = mix_spec.unsqueeze(1) # (B, 1, C, F, TT) + est_specs = masks_c * mix_spec_bc # (B, S, C, F, TT) + + # 6) iSTFT по каждому источнику + outs = [] + for s in range(S): + y_s = self._istft(est_specs[:, s], length=T) # (B, C, T) + outs.append(y_s) + y_hat = torch.stack(outs, dim=1) # (B, S, C, T) + return y_hat diff --git a/src/third_party/MusicSourceSeparationTraining/models/demucs4ht.py b/src/third_party/MusicSourceSeparationTraining/models/demucs4ht.py new file mode 100644 index 0000000000000000000000000000000000000000..b26cf085df6f8e236fa37664147c18a765a03fa8 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/demucs4ht.py @@ -0,0 +1,708 @@ +import math +from fractions import Fraction + +import torch +import torch.nn as nn +import torch.nn.functional as F +from demucs.demucs import Demucs, rescale_module +from demucs.hdemucs import ( + HDecLayer, + HDemucs, + HEncLayer, + MultiWrap, + ScaledEmbedding, + pad1d, +) +from demucs.spec import ispectro, spectro +from demucs.states import capture_init +from demucs.transformer import CrossTransformerEncoder +from einops import rearrange +from omegaconf import OmegaConf +from openunmix.filtering import wiener + + +class HTDemucs(nn.Module): + """ + Spectrogram and hybrid Demucs model. + The spectrogram model has the same structure as Demucs, except the first few layers are over the + frequency axis, until there is only 1 frequency, and then it moves to time convolutions. + Frequency layers can still access information across time steps thanks to the DConv residual. + + Hybrid model have a parallel time branch. At some layer, the time branch has the same stride + as the frequency branch and then the two are combined. The opposite happens in the decoder. + + Models can either use naive iSTFT from masking, Wiener filtering ([Ulhih et al. 2017]), + or complex as channels (CaC) [Choi et al. 2020]. Wiener filtering is based on + Open Unmix implementation [Stoter et al. 2019]. + + The loss is always on the temporal domain, by backpropagating through the above + output methods and iSTFT. This allows to define hybrid models nicely. However, this breaks + a bit Wiener filtering, as doing more iteration at test time will change the spectrogram + contribution, without changing the one from the waveform, which will lead to worse performance. + I tried using the residual option in OpenUnmix Wiener implementation, but it didn't improve. + CaC on the other hand provides similar performance for hybrid, and works naturally with + hybrid models. + + This model also uses frequency embeddings are used to improve efficiency on convolutions + over the freq. axis, following [Isik et al. 2020] (https://arxiv.org/pdf/2008.04470.pdf). + + Unlike classic Demucs, there is no resampling here, and normalization is always applied. + """ + + @capture_init + def __init__( + self, + sources, + # Channels + audio_channels=2, + channels=48, + channels_time=None, + growth=2, + # STFT + nfft=4096, + num_subbands=1, + wiener_iters=0, + end_iters=0, + wiener_residual=False, + cac=True, + # Main structure + depth=4, + rewrite=True, + # Frequency branch + multi_freqs=None, + multi_freqs_depth=3, + freq_emb=0.2, + emb_scale=10, + emb_smooth=True, + # Convolutions + kernel_size=8, + time_stride=2, + stride=4, + context=1, + context_enc=0, + # Normalization + norm_starts=4, + norm_groups=4, + # DConv residual branch + dconv_mode=1, + dconv_depth=2, + dconv_comp=8, + dconv_init=1e-3, + # Before the Transformer + bottom_channels=0, + # Transformer + t_layers=5, + t_emb="sin", + t_hidden_scale=4.0, + t_heads=8, + t_dropout=0.0, + t_max_positions=10000, + t_norm_in=True, + t_norm_in_group=False, + t_group_norm=False, + t_norm_first=True, + t_norm_out=True, + t_max_period=10000.0, + t_weight_decay=0.0, + t_lr=None, + t_layer_scale=True, + t_gelu=True, + t_weight_pos_embed=1.0, + t_sin_random_shift=0, + t_cape_mean_normalize=True, + t_cape_augment=True, + t_cape_glob_loc_scale=[5000.0, 1.0, 1.4], + t_sparse_self_attn=False, + t_sparse_cross_attn=False, + t_mask_type="diag", + t_mask_random_seed=42, + t_sparse_attn_window=500, + t_global_window=100, + t_sparsity=0.95, + t_auto_sparsity=False, + # ------ Particuliar parameters + t_cross_first=False, + # Weight init + rescale=0.1, + # Metadata + samplerate=44100, + segment=10, + use_train_segment=False, + ): + """ + Args: + sources (list[str]): list of source names. + audio_channels (int): input/output audio channels. + channels (int): initial number of hidden channels. + channels_time: if not None, use a different `channels` value for the time branch. + growth: increase the number of hidden channels by this factor at each layer. + nfft: number of fft bins. Note that changing this require careful computation of + various shape parameters and will not work out of the box for hybrid models. + wiener_iters: when using Wiener filtering, number of iterations at test time. + end_iters: same but at train time. For a hybrid model, must be equal to `wiener_iters`. + wiener_residual: add residual source before wiener filtering. + cac: uses complex as channels, i.e. complex numbers are 2 channels each + in input and output. no further processing is done before ISTFT. + depth (int): number of layers in the encoder and in the decoder. + rewrite (bool): add 1x1 convolution to each layer. + multi_freqs: list of frequency ratios for splitting frequency bands with `MultiWrap`. + multi_freqs_depth: how many layers to wrap with `MultiWrap`. Only the outermost + layers will be wrapped. + freq_emb: add frequency embedding after the first frequency layer if > 0, + the actual value controls the weight of the embedding. + emb_scale: equivalent to scaling the embedding learning rate + emb_smooth: initialize the embedding with a smooth one (with respect to frequencies). + kernel_size: kernel_size for encoder and decoder layers. + stride: stride for encoder and decoder layers. + time_stride: stride for the final time layer, after the merge. + context: context for 1x1 conv in the decoder. + context_enc: context for 1x1 conv in the encoder. + norm_starts: layer at which group norm starts being used. + decoder layers are numbered in reverse order. + norm_groups: number of groups for group norm. + dconv_mode: if 1: dconv in encoder only, 2: decoder only, 3: both. + dconv_depth: depth of residual DConv branch. + dconv_comp: compression of DConv branch. + dconv_attn: adds attention layers in DConv branch starting at this layer. + dconv_lstm: adds a LSTM layer in DConv branch starting at this layer. + dconv_init: initial scale for the DConv branch LayerScale. + bottom_channels: if >0 it adds a linear layer (1x1 Conv) before and after the + transformer in order to change the number of channels + t_layers: number of layers in each branch (waveform and spec) of the transformer + t_emb: "sin", "cape" or "scaled" + t_hidden_scale: the hidden scale of the Feedforward parts of the transformer + for instance if C = 384 (the number of channels in the transformer) and + t_hidden_scale = 4.0 then the intermediate layer of the FFN has dimension + 384 * 4 = 1536 + t_heads: number of heads for the transformer + t_dropout: dropout in the transformer + t_max_positions: max_positions for the "scaled" positional embedding, only + useful if t_emb="scaled" + t_norm_in: (bool) norm before addinf positional embedding and getting into the + transformer layers + t_norm_in_group: (bool) if True while t_norm_in=True, the norm is on all the + timesteps (GroupNorm with group=1) + t_group_norm: (bool) if True, the norms of the Encoder Layers are on all the + timesteps (GroupNorm with group=1) + t_norm_first: (bool) if True the norm is before the attention and before the FFN + t_norm_out: (bool) if True, there is a GroupNorm (group=1) at the end of each layer + t_max_period: (float) denominator in the sinusoidal embedding expression + t_weight_decay: (float) weight decay for the transformer + t_lr: (float) specific learning rate for the transformer + t_layer_scale: (bool) Layer Scale for the transformer + t_gelu: (bool) activations of the transformer are GeLU if True, ReLU else + t_weight_pos_embed: (float) weighting of the positional embedding + t_cape_mean_normalize: (bool) if t_emb="cape", normalisation of positional embeddings + see: https://arxiv.org/abs/2106.03143 + t_cape_augment: (bool) if t_emb="cape", must be True during training and False + during the inference, see: https://arxiv.org/abs/2106.03143 + t_cape_glob_loc_scale: (list of 3 floats) if t_emb="cape", CAPE parameters + see: https://arxiv.org/abs/2106.03143 + t_sparse_self_attn: (bool) if True, the self attentions are sparse + t_sparse_cross_attn: (bool) if True, the cross-attentions are sparse (don't use it + unless you designed really specific masks) + t_mask_type: (str) can be "diag", "jmask", "random", "global" or any combination + with '_' between: i.e. "diag_jmask_random" (note that this is permutation + invariant i.e. "diag_jmask_random" is equivalent to "jmask_random_diag") + t_mask_random_seed: (int) if "random" is in t_mask_type, controls the seed + that generated the random part of the mask + t_sparse_attn_window: (int) if "diag" is in t_mask_type, for a query (i), and + a key (j), the mask is True id |i-j|<=t_sparse_attn_window + t_global_window: (int) if "global" is in t_mask_type, mask[:t_global_window, :] + and mask[:, :t_global_window] will be True + t_sparsity: (float) if "random" is in t_mask_type, t_sparsity is the sparsity + level of the random part of the mask. + t_cross_first: (bool) if True cross attention is the first layer of the + transformer (False seems to be better) + rescale: weight rescaling trick + use_train_segment: (bool) if True, the actual size that is used during the + training is used during inference. + """ + super().__init__() + self.num_subbands = num_subbands + self.cac = cac + self.wiener_residual = wiener_residual + self.audio_channels = audio_channels + self.sources = sources + self.kernel_size = kernel_size + self.context = context + self.stride = stride + self.depth = depth + self.bottom_channels = bottom_channels + self.channels = channels + self.samplerate = samplerate + self.segment = segment + self.use_train_segment = use_train_segment + self.nfft = nfft + self.hop_length = nfft // 4 + self.wiener_iters = wiener_iters + self.end_iters = end_iters + self.freq_emb = None + assert wiener_iters == end_iters + + self.encoder = nn.ModuleList() + self.decoder = nn.ModuleList() + + self.tencoder = nn.ModuleList() + self.tdecoder = nn.ModuleList() + + chin = audio_channels + chin_z = chin # number of channels for the freq branch + if self.cac: + chin_z *= 2 + if self.num_subbands > 1: + chin_z *= self.num_subbands + chout = channels_time or channels + chout_z = channels + freqs = nfft // 2 + + for index in range(depth): + norm = index >= norm_starts + freq = freqs > 1 + stri = stride + ker = kernel_size + if not freq: + assert freqs == 1 + ker = time_stride * 2 + stri = time_stride + + pad = True + last_freq = False + if freq and freqs <= kernel_size: + ker = freqs + pad = False + last_freq = True + + kw = { + "kernel_size": ker, + "stride": stri, + "freq": freq, + "pad": pad, + "norm": norm, + "rewrite": rewrite, + "norm_groups": norm_groups, + "dconv_kw": { + "depth": dconv_depth, + "compress": dconv_comp, + "init": dconv_init, + "gelu": True, + }, + } + kwt = dict(kw) + kwt["freq"] = 0 + kwt["kernel_size"] = kernel_size + kwt["stride"] = stride + kwt["pad"] = True + kw_dec = dict(kw) + multi = False + if multi_freqs and index < multi_freqs_depth: + multi = True + kw_dec["context_freq"] = False + + if last_freq: + chout_z = max(chout, chout_z) + chout = chout_z + + enc = HEncLayer( + chin_z, chout_z, dconv=dconv_mode & 1, context=context_enc, **kw + ) + if freq: + tenc = HEncLayer( + chin, + chout, + dconv=dconv_mode & 1, + context=context_enc, + empty=last_freq, + **kwt, + ) + self.tencoder.append(tenc) + + if multi: + enc = MultiWrap(enc, multi_freqs) + self.encoder.append(enc) + if index == 0: + chin = self.audio_channels * len(self.sources) + chin_z = chin + if self.cac: + chin_z *= 2 + if self.num_subbands > 1: + chin_z *= self.num_subbands + dec = HDecLayer( + chout_z, + chin_z, + dconv=dconv_mode & 2, + last=index == 0, + context=context, + **kw_dec, + ) + if multi: + dec = MultiWrap(dec, multi_freqs) + if freq: + tdec = HDecLayer( + chout, + chin, + dconv=dconv_mode & 2, + empty=last_freq, + last=index == 0, + context=context, + **kwt, + ) + self.tdecoder.insert(0, tdec) + self.decoder.insert(0, dec) + + chin = chout + chin_z = chout_z + chout = int(growth * chout) + chout_z = int(growth * chout_z) + if freq: + if freqs <= kernel_size: + freqs = 1 + else: + freqs //= stride + if index == 0 and freq_emb: + self.freq_emb = ScaledEmbedding( + freqs, chin_z, smooth=emb_smooth, scale=emb_scale + ) + self.freq_emb_scale = freq_emb + + if rescale: + rescale_module(self, reference=rescale) + + transformer_channels = channels * growth ** (depth - 1) + if bottom_channels: + self.channel_upsampler = nn.Conv1d(transformer_channels, bottom_channels, 1) + self.channel_downsampler = nn.Conv1d( + bottom_channels, transformer_channels, 1 + ) + self.channel_upsampler_t = nn.Conv1d( + transformer_channels, bottom_channels, 1 + ) + self.channel_downsampler_t = nn.Conv1d( + bottom_channels, transformer_channels, 1 + ) + + transformer_channels = bottom_channels + + if t_layers > 0: + self.crosstransformer = CrossTransformerEncoder( + dim=transformer_channels, + emb=t_emb, + hidden_scale=t_hidden_scale, + num_heads=t_heads, + num_layers=t_layers, + cross_first=t_cross_first, + dropout=t_dropout, + max_positions=t_max_positions, + norm_in=t_norm_in, + norm_in_group=t_norm_in_group, + group_norm=t_group_norm, + norm_first=t_norm_first, + norm_out=t_norm_out, + max_period=t_max_period, + weight_decay=t_weight_decay, + lr=t_lr, + layer_scale=t_layer_scale, + gelu=t_gelu, + sin_random_shift=t_sin_random_shift, + weight_pos_embed=t_weight_pos_embed, + cape_mean_normalize=t_cape_mean_normalize, + cape_augment=t_cape_augment, + cape_glob_loc_scale=t_cape_glob_loc_scale, + sparse_self_attn=t_sparse_self_attn, + sparse_cross_attn=t_sparse_cross_attn, + mask_type=t_mask_type, + mask_random_seed=t_mask_random_seed, + sparse_attn_window=t_sparse_attn_window, + global_window=t_global_window, + sparsity=t_sparsity, + auto_sparsity=t_auto_sparsity, + ) + else: + self.crosstransformer = None + + def _spec(self, x): + hl = self.hop_length + nfft = self.nfft + x0 = x # noqa + + # We re-pad the signal in order to keep the property + # that the size of the output is exactly the size of the input + # divided by the stride (here hop_length), when divisible. + # This is achieved by padding by 1/4th of the kernel size (here nfft). + # which is not supported by torch.stft. + # Having all convolution operations follow this convention allow to easily + # align the time and frequency branches later on. + assert hl == nfft // 4 + le = int(math.ceil(x.shape[-1] / hl)) + pad = hl // 2 * 3 + x = pad1d(x, (pad, pad + le * hl - x.shape[-1]), mode="reflect") + + z = spectro(x, nfft, hl)[..., :-1, :] + assert z.shape[-1] == le + 4, (z.shape, x.shape, le) + z = z[..., 2 : 2 + le] + return z + + def _ispec(self, z, length=None, scale=0): + hl = self.hop_length // (4**scale) + z = F.pad(z, (0, 0, 0, 1)) + z = F.pad(z, (2, 2)) + pad = hl // 2 * 3 + le = hl * int(math.ceil(length / hl)) + 2 * pad + x = ispectro(z, hl, length=le) + x = x[..., pad : pad + length] + return x + + def _magnitude(self, z): + # return the magnitude of the spectrogram, except when cac is True, + # in which case we just move the complex dimension to the channel one. + if self.cac: + B, C, Fr, T = z.shape + m = torch.view_as_real(z).permute(0, 1, 4, 2, 3) + m = m.reshape(B, C * 2, Fr, T) + else: + m = z.abs() + return m + + def _mask(self, z, m): + # Apply masking given the mixture spectrogram `z` and the estimated mask `m`. + # If `cac` is True, `m` is actually a full spectrogram and `z` is ignored. + niters = self.wiener_iters + if self.cac: + B, S, C, Fr, T = m.shape + out = m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3) + out = torch.view_as_complex(out.contiguous()) + return out + if self.training: + niters = self.end_iters + if niters < 0: + z = z[:, None] + return z / (1e-8 + z.abs()) * m + else: + return self._wiener(m, z, niters) + + def _wiener(self, mag_out, mix_stft, niters): + # apply wiener filtering from OpenUnmix. + init = mix_stft.dtype + wiener_win_len = 300 + residual = self.wiener_residual + + B, S, C, Fq, T = mag_out.shape + mag_out = mag_out.permute(0, 4, 3, 2, 1) + mix_stft = torch.view_as_real(mix_stft.permute(0, 3, 2, 1)) + + outs = [] + for sample in range(B): + pos = 0 + out = [] + for pos in range(0, T, wiener_win_len): + frame = slice(pos, pos + wiener_win_len) + z_out = wiener( + mag_out[sample, frame], + mix_stft[sample, frame], + niters, + residual=residual, + ) + out.append(z_out.transpose(-1, -2)) + outs.append(torch.cat(out, dim=0)) + out = torch.view_as_complex(torch.stack(outs, 0)) + out = out.permute(0, 4, 3, 2, 1).contiguous() + if residual: + out = out[:, :-1] + assert list(out.shape) == [B, S, C, Fq, T] + return out.to(init) + + def valid_length(self, length: int): + """ + Return a length that is appropriate for evaluation. + In our case, always return the training length, unless + it is smaller than the given length, in which case this + raises an error. + """ + if not self.use_train_segment: + return length + training_length = int(self.segment * self.samplerate) + if training_length < length: + raise ValueError( + f"Given length {length} is longer than " + f"training length {training_length}" + ) + return training_length + + def cac2cws(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c, k, f // k, t) + x = x.reshape(b, c * k, f // k, t) + return x + + def cws2cac(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c // k, k, f, t) + x = x.reshape(b, c // k, f * k, t) + return x + + def forward(self, mix): + length = mix.shape[-1] + length_pre_pad = None + if self.use_train_segment: + if self.training: + self.segment = Fraction(mix.shape[-1], self.samplerate) + else: + training_length = int(self.segment * self.samplerate) + # print('Training length: {} Segment: {} Sample rate: {}'.format(training_length, self.segment, self.samplerate)) + if mix.shape[-1] < training_length: + length_pre_pad = mix.shape[-1] + mix = F.pad(mix, (0, training_length - length_pre_pad)) + # print("Mix: {}".format(mix.shape)) + # print("Length: {}".format(length)) + z = self._spec(mix) + # print("Z: {} Type: {}".format(z.shape, z.dtype)) + mag = self._magnitude(z) + x = mag + # print("MAG: {} Type: {}".format(x.shape, x.dtype)) + + if self.num_subbands > 1: + x = self.cac2cws(x) + # print("After SUBBANDS: {} Type: {}".format(x.shape, x.dtype)) + + B, C, Fq, T = x.shape + + # unlike previous Demucs, we always normalize because it is easier. + mean = x.mean(dim=(1, 2, 3), keepdim=True) + std = x.std(dim=(1, 2, 3), keepdim=True) + x = (x - mean) / (1e-5 + std) + # x will be the freq. branch input. + + # Prepare the time branch input. + xt = mix + meant = xt.mean(dim=(1, 2), keepdim=True) + stdt = xt.std(dim=(1, 2), keepdim=True) + xt = (xt - meant) / (1e-5 + stdt) + + # print("XT: {}".format(xt.shape)) + + # okay, this is a giant mess I know... + saved = [] # skip connections, freq. + saved_t = [] # skip connections, time. + lengths = [] # saved lengths to properly remove padding, freq branch. + lengths_t = [] # saved lengths for time branch. + for idx, encode in enumerate(self.encoder): + lengths.append(x.shape[-1]) + inject = None + if idx < len(self.tencoder): + # we have not yet merged branches. + lengths_t.append(xt.shape[-1]) + tenc = self.tencoder[idx] + xt = tenc(xt) + # print("Encode XT {}: {}".format(idx, xt.shape)) + if not tenc.empty: + # save for skip connection + saved_t.append(xt) + else: + # tenc contains just the first conv., so that now time and freq. + # branches have the same shape and can be merged. + inject = xt + x = encode(x, inject) + # print("Encode X {}: {}".format(idx, x.shape)) + if idx == 0 and self.freq_emb is not None: + # add frequency embedding to allow for non equivariant convolutions + # over the frequency axis. + frs = torch.arange(x.shape[-2], device=x.device) + emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x) + x = x + self.freq_emb_scale * emb + + saved.append(x) + if self.crosstransformer: + if self.bottom_channels: + b, c, f, t = x.shape + x = rearrange(x, "b c f t-> b c (f t)") + x = self.channel_upsampler(x) + x = rearrange(x, "b c (f t)-> b c f t", f=f) + xt = self.channel_upsampler_t(xt) + + x, xt = self.crosstransformer(x, xt) + # print("Cross Tran X {}, XT: {}".format(x.shape, xt.shape)) + + if self.bottom_channels: + x = rearrange(x, "b c f t-> b c (f t)") + x = self.channel_downsampler(x) + x = rearrange(x, "b c (f t)-> b c f t", f=f) + xt = self.channel_downsampler_t(xt) + + for idx, decode in enumerate(self.decoder): + skip = saved.pop(-1) + x, pre = decode(x, skip, lengths.pop(-1)) + # print('Decode {} X: {}'.format(idx, x.shape)) + # `pre` contains the output just before final transposed convolution, + # which is used when the freq. and time branch separate. + + offset = self.depth - len(self.tdecoder) + if idx >= offset: + tdec = self.tdecoder[idx - offset] + length_t = lengths_t.pop(-1) + if tdec.empty: + assert pre.shape[2] == 1, pre.shape + pre = pre[:, :, 0] + xt, _ = tdec(pre, None, length_t) + else: + skip = saved_t.pop(-1) + xt, _ = tdec(xt, skip, length_t) + # print('Decode {} XT: {}'.format(idx, xt.shape)) + + # Let's make sure we used all stored skip connections. + assert len(saved) == 0 + assert len(lengths_t) == 0 + assert len(saved_t) == 0 + + S = len(self.sources) + + if self.num_subbands > 1: + x = x.view(B, -1, Fq, T) + # print("X view 1: {}".format(x.shape)) + x = self.cws2cac(x) + # print("X view 2: {}".format(x.shape)) + + x = x.view(B, S, -1, Fq * self.num_subbands, T) + x = x * std[:, None] + mean[:, None] + # print("X returned: {}".format(x.shape)) + + zout = self._mask(z, x) + if self.use_train_segment: + if self.training: + x = self._ispec(zout, length) + else: + x = self._ispec(zout, training_length) + else: + x = self._ispec(zout, length) + + if self.use_train_segment: + if self.training: + xt = xt.view(B, S, -1, length) + else: + xt = xt.view(B, S, -1, training_length) + else: + xt = xt.view(B, S, -1, length) + xt = xt * stdt[:, None] + meant[:, None] + x = xt + x + if length_pre_pad: + x = x[..., :length_pre_pad] + return x + + +def get_model(args): + extra = { + "sources": list(args.training.instruments), + "audio_channels": args.training.channels, + "samplerate": args.training.samplerate, + # 'segment': args.model_segment or 4 * args.dset.segment, + "segment": args.training.segment, + } + klass = { + "demucs": Demucs, + "hdemucs": HDemucs, + "htdemucs": HTDemucs, + }[args.model] + kw = OmegaConf.to_container(getattr(args, args.model), resolve=True) + model = klass(**extra, **kw) + return model diff --git a/src/third_party/MusicSourceSeparationTraining/models/ex_bi_mamba2.py b/src/third_party/MusicSourceSeparationTraining/models/ex_bi_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..af30b6bb8dbe4489f308568c4e24024655f4eb4d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/ex_bi_mamba2.py @@ -0,0 +1,405 @@ +# https://github.com/Human9000/nd-Mamba2-torch + +from abc import abstractmethod + +import torch +from torch import Tensor, nn +from torch.nn import functional as F + + +def silu(x): + return x * F.sigmoid(x) + + +class RMSNorm(nn.Module): + def __init__(self, d: int, eps: float = 1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(d)) + + def forward(self, x, z): + x = x * silu(z) + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight + + +class Mamba2(nn.Module): + def __init__( + self, + d_model: int, # model dimension (D) + n_layer: int = 24, # number of Mamba-2 layers in the language model + d_state: int = 128, # state dimension (N) + d_conv: int = 4, # convolution kernel size + expand: int = 2, # expansion factor (E) + headdim: int = 64, # head dimension (P) + chunk_size: int = 64, # matrix partition size (Q) + ): + super().__init__() + self.n_layer = n_layer + self.d_state = d_state + self.headdim = headdim + # self.chunk_size = torch.tensor(chunk_size, dtype=torch.int32) + self.chunk_size = chunk_size + + self.d_inner = expand * d_model + assert self.d_inner % self.headdim == 0, ( + "self.d_inner must be divisible by self.headdim" + ) + self.nheads = self.d_inner // self.headdim + + d_in_proj = 2 * self.d_inner + 2 * self.d_state + self.nheads + self.in_proj = nn.Linear(d_model, d_in_proj, bias=False) + + conv_dim = self.d_inner + 2 * d_state + self.conv1d = nn.Conv1d( + conv_dim, + conv_dim, + d_conv, + groups=conv_dim, + padding=d_conv - 1, + ) + self.dt_bias = nn.Parameter( + torch.empty( + self.nheads, + ) + ) + self.A_log = nn.Parameter( + torch.empty( + self.nheads, + ) + ) + self.D = nn.Parameter( + torch.empty( + self.nheads, + ) + ) + self.norm = RMSNorm( + self.d_inner, + ) + self.out_proj = nn.Linear( + self.d_inner, + d_model, + bias=False, + ) + + def forward(self, u: Tensor): + A = -torch.exp(self.A_log) # (nheads,) + zxbcdt = self.in_proj(u) # (batch, seqlen, d_in_proj) + z, xBC, dt = torch.split( + zxbcdt, + [ + self.d_inner, + self.d_inner + 2 * self.d_state, + self.nheads, + ], + dim=-1, + ) + dt = F.softplus(dt + self.dt_bias) # (batch, seqlen, nheads) + + # Pad or truncate xBC seqlen to d_conv + xBC = silu( + self.conv1d(xBC.transpose(1, 2)).transpose(1, 2)[:, : u.shape[1], :] + ) # (batch, seqlen, d_inner + 2 * d_state)) + x, B, C = torch.split(xBC, [self.d_inner, self.d_state, self.d_state], dim=-1) + + _b, _l, _hp = x.shape + _h = _hp // self.headdim + _p = self.headdim + x = x.reshape(_b, _l, _h, _p) + + y = self.ssd( + x * dt.unsqueeze(-1), + A * dt, + B.unsqueeze(2), + C.unsqueeze(2), + ) + + y = y + x * self.D.unsqueeze(-1) + + _b, _l, _h, _p = y.shape + y = y.reshape(_b, _l, _h * _p) + + y = self.norm(y, z) + y = self.out_proj(y) + + return y + + def segsum(self, x: Tensor) -> Tensor: + T = x.size(-1) + device = x.device + x = x[..., None].repeat(1, 1, 1, 1, T) + mask = torch.tril( + torch.ones(T, T, dtype=torch.bool, device=device), diagonal=-1 + ) + x = x.masked_fill(~mask, 0) + x_segsum = torch.cumsum(x, dim=-2) + mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=0) + x_segsum = x_segsum.masked_fill(~mask, -torch.inf) + return x_segsum + + def ssd(self, x, A, B, C): + chunk_size = self.chunk_size + # if x.shape[1] % chunk_size == 0: + # + x = x.reshape( + x.shape[0], + x.shape[1] // chunk_size, + chunk_size, + x.shape[2], + x.shape[3], + ) + B = B.reshape( + B.shape[0], + B.shape[1] // chunk_size, + chunk_size, + B.shape[2], + B.shape[3], + ) + C = C.reshape( + C.shape[0], + C.shape[1] // chunk_size, + chunk_size, + C.shape[2], + C.shape[3], + ) + A = A.reshape(A.shape[0], A.shape[1] // chunk_size, chunk_size, A.shape[2]) + A = A.permute(0, 3, 1, 2) + A_cumsum = torch.cumsum(A, dim=-1) + + # 1. Compute the output for each intra-chunk (diagonal blocks) + L = torch.exp(self.segsum(A)) + Y_diag = torch.einsum("bclhn, bcshn, bhcls, bcshp -> bclhp", C, B, L, x) + + # 2. Compute the state for each intra-chunk + # (right term of low-rank factorization of off-diagonal blocks; B terms) + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + states = torch.einsum("bclhn, bhcl, bclhp -> bchpn", B, decay_states, x) + + # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries + # (middle term of factorization of off-diag blocks; A terms) + + initial_states = torch.zeros_like(states[:, :1]) + states = torch.cat([initial_states, states], dim=1) + + decay_chunk = torch.exp(self.segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))[0] + new_states = torch.einsum("bhzc, bchpn -> bzhpn", decay_chunk, states) + states = new_states[:, :-1] + + # 4. Compute state -> output conversion per chunk + # (left term of low-rank factorization of off-diagonal blocks; C terms) + state_decay_out = torch.exp(A_cumsum) + Y_off = torch.einsum("bclhn, bchpn, bhcl -> bclhp", C, states, state_decay_out) + + # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks) + # Y = rearrange(Y_diag + Y_off, "b c l h p -> b (c l) h p") + Y = Y_diag + Y_off + Y = Y.reshape( + Y.shape[0], + Y.shape[1] * Y.shape[2], + Y.shape[3], + Y.shape[4], + ) + + return Y + + +class _BiMamba2(nn.Module): + def __init__( + self, + cin: int, + cout: int, + d_model: int, # model dimension (D) + n_layer: int = 24, # number of Mamba-2 layers in the language model + d_state: int = 128, # state dimension (N) + d_conv: int = 4, # convolution kernel size + expand: int = 2, # expansion factor (E) + headdim: int = 64, # head dimension (P) + chunk_size: int = 64, # matrix partition size (Q) + ): + super().__init__() + self.fc_in = nn.Linear(cin, d_model, bias=False) # 调整通道数到cmid + self.mamba2_for = Mamba2( + d_model, + n_layer, + d_state, + d_conv, + expand, + headdim, + chunk_size, + ) # 正向 + self.mamba2_back = Mamba2( + d_model, + n_layer, + d_state, + d_conv, + expand, + headdim, + chunk_size, + ) # 负向 + self.fc_out = nn.Linear(d_model, cout, bias=False) # 调整通道数到cout + self.chunk_size = chunk_size + + @abstractmethod + def forward(self, x): + pass + + +class BiMamba2_1D(_BiMamba2): + def __init__(self, cin, cout, d_model, **mamba2_args): + super().__init__(cin, cout, d_model, **mamba2_args) + + def forward(self, x): + l = x.shape[2] + x = F.pad( + x, (0, (64 - x.shape[2] % 64) % 64) + ) # 将 l , pad到4的倍数, [b, c64,l4] + x = x.transpose(1, 2) # 转成 1d 信号 [b, c64, d4*w4*h4] + x = self.fc_in(x) # 调整通道数为目标通道数 + x1 = self.mamba2_for(x) + x2 = self.mamba2_back(x.flip(1)).flip(1) + x = x1 + x2 + x = self.fc_out(x) # 调整通道数为目标通道数 + x = x.transpose(1, 2) # 转成 1d 信号 [b, c64, d4*w4*h4] ] + x = x[:, :, :l] # 截取原图大小 + return x + + +class BiMamba2_2D(_BiMamba2): + def __init__(self, cin, cout, d_model, **mamba2_args): + super().__init__(cin, cout, d_model, **mamba2_args) + + def forward(self, x): + h, w = x.shape[2:] + x = F.pad( + x, (0, (8 - x.shape[3] % 8) % 8, 0, (8 - x.shape[2] % 8) % 8) + ) # 将 h , w pad到8的倍数, [b, c64, h8, w8] + _b, _c, _h, _w = x.shape + x = x.permute(0, 2, 3, 1).reshape(_b, _h * _w, _c) + x = self.fc_in(x) # 调整通道数为目标通道数 + x1 = self.mamba2_for(x) + x2 = self.mamba2_back(x.flip(1)).flip(1) + x = x1 + x2 + x = self.fc_out(x) # 调整通道数为目标通道数 + x = x.reshape( + _b, + _h, + _w, + -1, + ) + x = x.permute(0, 3, 1, 2) + x = x.reshape( + _b, + -1, + _h, + _w, + ) + x = x[:, :, :h, :w] # 截取原图大小 + return x + + +class BiMamba2_3D(_BiMamba2): + def __init__(self, cin, cout, d_model, **mamba2_args): + super().__init__(cin, cout, d_model, **mamba2_args) + + def forward(self, x): + d, h, w = x.shape[2:] + x = F.pad( + x, + ( + 0, + (4 - x.shape[4] % 4) % 4, + 0, + (4 - x.shape[3] % 4) % 4, + 0, + (4 - x.shape[2] % 4) % 4, + ), + ) # 将 d, h, w , pad到4的倍数, [b, c64,d4, h4, w4] + _b, _c, _d, _h, _w = x.shape + x = x.permute(0, 2, 3, 4, 1).reshape(_b, _d * _h * _w, _c) + x = self.fc_in(x) # 调整通道数为目标通道数 + x1 = self.mamba2_for(x) + x2 = self.mamba2_back(x.flip(1)).flip(1) + x = x1 + x2 + x = self.fc_out(x) # 调整通道数为目标通道数 + x = x.reshape(_b, _d, _h, _w, -1) + x = x.permute(0, 4, 1, 2, 3) + x = x.reshape( + _b, + -1, + _d, + _h, + _w, + ) + x = x[:, :, :d, :h, :w] # 截取原图大小 + return x + + +class BiMamba2(_BiMamba2): + def __init__(self, cin, cout, d_model, **mamba2_args): + super().__init__(cin, cout, d_model, **mamba2_args) + + def forward(self, x): + size = x.shape[2:] + out_size = list(x.shape) + out_size[1] = -1 + + x = torch.flatten(x, 2) # b c size + l = x.shape[2] + _s = self.chunk_size + x = F.pad( + x, [0, (_s - x.shape[2] % _s) % _s] + ) # 将 l, pad到chunk_size的倍数, [b, c64,l4] + x = x.transpose(1, 2) # 转成 1d 信号 + x = self.fc_in(x) # 调整通道数为目标通道数 + x1 = self.mamba2_for(x) + x2 = self.mamba2_back(x.flip(1)).flip(1) + x = x1 + x2 + x = self.fc_out(x) # 调整通道数为目标通道数 + x = x.transpose(1, 2) # 转成 1d 信号 + x = x[:, :, :l] # 截取原图大小 + x = x.reshape(out_size) + + return x + + +def test_export_jit_script(net, x): + y = net(x) + net_script = torch.jit.script(net) + torch.jit.save(net_script, "net.jit.script") + net2 = torch.jit.load("net.jit.script") + y = net2(x) + print(y.shape) + + +def test_export_onnx(net, x): + torch.onnx.export( + net, + x, + "net.onnx", # 输出的 ONNX 文件名 + export_params=True, # 存储训练参数 + opset_version=14, # 指定 ONNX 操作集版本 + do_constant_folding=False, # 是否执行常量折叠优化 + input_names=["input"], # 输入张量的名称 + output_names=["output"], # 输出张量的名称 + dynamic_axes={ + "input": {0: "batch_size"}, # 可变维度的字典 + "output": {0: "batch_size"}, + }, + ) + + +if __name__ == "__main__": + # 通用的多维度双向mamba2 + from torchnssd import ( + export_jit_script, + export_onnx, + statistics, + test_run, + ) + + net_n = BiMamba2_1D(61, 128, 32).cuda() + net_n.eval() + x = torch.randn(1, 61, 63).cuda() + export_jit_script(net_n) + export_onnx(net_n, x) + test_run(net_n, x) + statistics(net_n, (61, 63)) diff --git a/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f4378f0e488f4eeb2b433767b86dfda9621e9465 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/__init__.py @@ -0,0 +1,45 @@ +### +# Author: Kai Li +# Date: 2022-02-12 15:16:35 +# Email: lk21@mails.tsinghua.edu.cn +# LastEditTime: 2022-10-04 16:24:53 +### +from .apollo import Apollo +from .base_model import BaseModel + +__all__ = ["BaseModel", "GullFullband", "Apollo"] + + +def register_model(custom_model): + """Register a custom model, gettable with `models.get`. + + Args: + custom_model: Custom model to register. + + """ + if ( + custom_model.__name__ in globals().keys() + or custom_model.__name__.lower() in globals().keys() + ): + raise ValueError( + f"Model {custom_model.__name__} already exists. Choose another name." + ) + globals().update({custom_model.__name__: custom_model}) + + +def get(identifier): + """Returns an model class from a string (case-insensitive). + + Args: + identifier (str): the model name. + + Returns: + :class:`torch.nn.Module` + """ + if isinstance(identifier, str): + to_get = {k.lower(): v for k, v in globals().items()} + cls = to_get.get(identifier.lower()) + if cls is None: + raise ValueError(f"Could not interpret model name : {str(identifier)}") + return cls + raise ValueError(f"Could not interpret model name : {str(identifier)}") diff --git a/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/apollo.py b/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/apollo.py new file mode 100644 index 0000000000000000000000000000000000000000..f60533cb2116f94746f3996d0c1b78fb11cf5990 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/apollo.py @@ -0,0 +1,409 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .base_model import BaseModel + + +class RMSNorm(nn.Module): + def __init__(self, dimension, groups=1): + super().__init__() + + self.weight = nn.Parameter(torch.ones(dimension)) + self.groups = groups + self.eps = 1e-5 + + def forward(self, input): + # input size: (B, N, T) + B, N, T = input.shape + assert N % self.groups == 0 + + input_float = input.reshape(B, self.groups, -1, T).float() + input_norm = input_float * torch.rsqrt( + input_float.pow(2).mean(-2, keepdim=True) + self.eps + ) + + return input_norm.type_as(input).reshape(B, N, T) * self.weight.reshape( + 1, -1, 1 + ) + + +class RMVN(nn.Module): + """ + Rescaled MVN. + """ + + def __init__(self, dimension, groups=1): + super(RMVN, self).__init__() + + self.mean = nn.Parameter(torch.zeros(dimension)) + self.std = nn.Parameter(torch.ones(dimension)) + self.groups = groups + self.eps = 1e-5 + + def forward(self, input): + # input size: (B, N, *) + B, N = input.shape[:2] + assert N % self.groups == 0 + input_reshape = input.reshape(B, self.groups, N // self.groups, -1) + T = input_reshape.shape[-1] + + input_norm = (input_reshape - input_reshape.mean(2).unsqueeze(2)) / ( + input_reshape.var(2).unsqueeze(2) + self.eps + ).sqrt() + input_norm = input_norm.reshape(B, N, T) * self.std.reshape( + 1, -1, 1 + ) + self.mean.reshape(1, -1, 1) + + return input_norm.reshape(input.shape) + + +class Roformer(nn.Module): + """ + Transformer with rotary positional embedding. + """ + + def __init__( + self, + input_size, + hidden_size, + num_head=8, + theta=10000, + window=10000, + input_drop=0.0, + attention_drop=0.0, + causal=True, + ): + super().__init__() + + self.input_size = input_size + self.hidden_size = hidden_size // num_head + self.num_head = num_head + self.theta = theta # base frequency for RoPE + self.window = window + # pre-calculate rotary embeddings + cos_freq, sin_freq = self._calc_rotary_emb() + self.register_buffer("cos_freq", cos_freq) # win, N + self.register_buffer("sin_freq", sin_freq) # win, N + + self.attention_drop = attention_drop + self.causal = causal + self.eps = 1e-5 + + self.input_norm = RMSNorm(self.input_size) + self.input_drop = nn.Dropout(p=input_drop) + self.weight = nn.Conv1d( + self.input_size, self.hidden_size * self.num_head * 3, 1, bias=False + ) + self.output = nn.Conv1d( + self.hidden_size * self.num_head, self.input_size, 1, bias=False + ) + + self.MLP = nn.Sequential( + RMSNorm(self.input_size), + nn.Conv1d(self.input_size, self.input_size * 8, 1, bias=False), + nn.SiLU(), + ) + self.MLP_output = nn.Conv1d(self.input_size * 4, self.input_size, 1, bias=False) + + def _calc_rotary_emb(self): + freq = 1.0 / ( + self.theta + ** ( + torch.arange(0, self.hidden_size, 2)[: (self.hidden_size // 2)] + / self.hidden_size + ) + ) # theta_i + freq = freq.reshape(1, -1) # 1, N//2 + pos = torch.arange(0, self.window).reshape(-1, 1) # win, 1 + cos_freq = torch.cos(pos * freq) # win, N//2 + sin_freq = torch.sin(pos * freq) # win, N//2 + cos_freq = torch.stack([cos_freq] * 2, -1).reshape( + self.window, self.hidden_size + ) # win, N + sin_freq = torch.stack([sin_freq] * 2, -1).reshape( + self.window, self.hidden_size + ) # win, N + + return cos_freq, sin_freq + + def _add_rotary_emb(self, feature, pos): + # feature shape: ..., N + N = feature.shape[-1] + + feature_reshape = feature.reshape(-1, N) + pos = min(pos, self.window - 1) + cos_freq = self.cos_freq[pos] + sin_freq = self.sin_freq[pos] + reverse_sign = ( + torch.from_numpy(np.asarray([-1, 1])).to(feature.device).type(feature.dtype) + ) + feature_reshape_neg = ( + torch.flip(feature_reshape.reshape(-1, N // 2, 2), [-1]) + * reverse_sign.reshape(1, 1, 2) + ).reshape(-1, N) + feature_rope = feature_reshape * cos_freq.unsqueeze( + 0 + ) + feature_reshape_neg * sin_freq.unsqueeze(0) + + return feature_rope.reshape(feature.shape) + + def _add_rotary_sequence(self, feature): + # feature shape: ..., T, N + T, N = feature.shape[-2:] + feature_reshape = feature.reshape(-1, T, N) + + cos_freq = self.cos_freq[:T] + sin_freq = self.sin_freq[:T] + reverse_sign = ( + torch.from_numpy(np.asarray([-1, 1])).to(feature.device).type(feature.dtype) + ) + feature_reshape_neg = ( + torch.flip(feature_reshape.reshape(-1, N // 2, 2), [-1]) + * reverse_sign.reshape(1, 1, 2) + ).reshape(-1, T, N) + feature_rope = feature_reshape * cos_freq.unsqueeze( + 0 + ) + feature_reshape_neg * sin_freq.unsqueeze(0) + + return feature_rope.reshape(feature.shape) + + def forward(self, input): + # input shape: B, N, T + + B, _, T = input.shape + + weight = ( + self.weight(self.input_drop(self.input_norm(input))) + .reshape(B, self.num_head, self.hidden_size * 3, T) + .mT + ) + Q, K, V = torch.split(weight, self.hidden_size, dim=-1) # B, num_head, T, N + + # rotary positional embedding + Q_rot = self._add_rotary_sequence(Q) + K_rot = self._add_rotary_sequence(K) + + attention_output = F.scaled_dot_product_attention( + Q_rot.contiguous(), + K_rot.contiguous(), + V.contiguous(), + dropout_p=self.attention_drop, + is_causal=self.causal, + ) # B, num_head, T, N + attention_output = attention_output.mT.reshape(B, -1, T) + output = self.output(attention_output) + input + + gate, z = self.MLP(output).chunk(2, dim=1) + output = output + self.MLP_output(F.silu(gate) * z) + + return output, (K_rot, V) + + +class ConvActNorm1d(nn.Module): + def __init__(self, in_channel, hidden_channel, kernel=7, causal=False): + super(ConvActNorm1d, self).__init__() + + self.in_channel = in_channel + self.kernel = kernel + self.causal = causal + if not causal: + self.conv = nn.Sequential( + nn.Conv1d( + in_channel, + in_channel, + kernel, + padding=(kernel - 1) // 2, + groups=in_channel, + ), + RMSNorm(in_channel), + nn.Conv1d(in_channel, hidden_channel, 1), + nn.SiLU(), + nn.Conv1d(hidden_channel, in_channel, 1), + ) + else: + self.conv = nn.Sequential( + nn.Conv1d( + in_channel, + in_channel, + kernel, + padding=kernel - 1, + groups=in_channel, + ), + RMSNorm(in_channel), + nn.Conv1d(in_channel, hidden_channel, 1), + nn.SiLU(), + nn.Conv1d(hidden_channel, in_channel, 1), + ) + + def forward(self, input): + output = self.conv(input) + if self.causal: + output = output[..., : -self.kernel + 1] + return input + output + + +class ICB(nn.Module): + def __init__(self, in_channel, kernel=7, causal=False): + super(ICB, self).__init__() + + self.blocks = nn.Sequential( + ConvActNorm1d(in_channel, in_channel * 4, kernel, causal=causal), + ConvActNorm1d(in_channel, in_channel * 4, kernel, causal=causal), + ConvActNorm1d(in_channel, in_channel * 4, kernel, causal=causal), + ) + + def forward(self, input): + return self.blocks(input) + + +class BSNet(nn.Module): + def __init__(self, feature_dim, kernel=7): + super(BSNet, self).__init__() + + self.feature_dim = feature_dim + + self.band_net = Roformer( + self.feature_dim, self.feature_dim, num_head=8, window=100, causal=False + ) + self.seq_net = ICB(self.feature_dim, kernel=kernel) + + def forward(self, input): + # input shape: B, nband, N, T + + B, nband, N, T = input.shape + + # band comm + band_input = input.permute(0, 3, 2, 1).reshape(B * T, -1, nband) + band_output, _ = self.band_net(band_input) + band_output = band_output.reshape(B, T, -1, nband).permute(0, 3, 2, 1) + + # sequence modeling + output = self.seq_net(band_output.reshape(B * nband, -1, T)).reshape( + B, nband, -1, T + ) # B, nband, N, T + + return output + + +class Apollo(BaseModel): + def __init__(self, sr: int, win: int, feature_dim: int, layer: int): + super().__init__(sample_rate=sr) + + self.sr = sr + self.win = int(sr * win // 1000) + self.stride = self.win // 2 + self.enc_dim = self.win // 2 + 1 + self.feature_dim = feature_dim + self.eps = torch.finfo(torch.float32).eps + + # 80 bands + bandwidth = int(self.win / 160) + self.band_width = [bandwidth] * 79 + self.band_width.append(self.enc_dim - np.sum(self.band_width)) + self.nband = len(self.band_width) + print(self.band_width, self.nband) + + self.BN = nn.ModuleList([]) + for i in range(self.nband): + self.BN.append( + nn.Sequential( + RMSNorm(self.band_width[i] * 2 + 1), + nn.Conv1d(self.band_width[i] * 2 + 1, self.feature_dim, 1), + ) + ) + + self.net = [] + for _ in range(layer): + self.net.append(BSNet(self.feature_dim)) + self.net = nn.Sequential(*self.net) + + self.output = nn.ModuleList([]) + for i in range(self.nband): + self.output.append( + nn.Sequential( + RMSNorm(self.feature_dim), + nn.Conv1d(self.feature_dim, self.band_width[i] * 4, 1), + nn.GLU(dim=1), + ) + ) + + def spec_band_split(self, input): + B, nch, nsample = input.shape + + spec = torch.stft( + input.view(B * nch, nsample), + n_fft=self.win, + hop_length=self.stride, + window=torch.hann_window(self.win).to(input.device), + return_complex=True, + ) + + subband_spec = [] + subband_spec_norm = [] + subband_power = [] + band_idx = 0 + for i in range(self.nband): + this_spec = spec[:, band_idx : band_idx + self.band_width[i]] + subband_spec.append(this_spec) # B, BW, T + subband_power.append( + (this_spec.abs().pow(2).sum(1) + self.eps).sqrt().unsqueeze(1) + ) # B, 1, T + subband_spec_norm.append( + torch.complex( + this_spec.real / subband_power[-1], + this_spec.imag / subband_power[-1], + ) + ) # B, BW, T + band_idx += self.band_width[i] + subband_power = torch.cat(subband_power, 1) # B, nband, T + + return subband_spec_norm, subband_power + + def feature_extractor(self, input): + subband_spec_norm, subband_power = self.spec_band_split(input) + + # normalization and bottleneck + subband_feature = [] + for i in range(self.nband): + concat_spec = torch.cat( + [ + subband_spec_norm[i].real, + subband_spec_norm[i].imag, + torch.log(subband_power[:, i].unsqueeze(1)), + ], + 1, + ) + subband_feature.append(self.BN[i](concat_spec)) + subband_feature = torch.stack(subband_feature, 1) # B, nband, N, T + + return subband_feature + + def forward(self, input): + B, nch, nsample = input.shape + + subband_feature = self.feature_extractor(input) + feature = self.net(subband_feature) + + est_spec = [] + for i in range(self.nband): + this_RI = self.output[i](feature[:, i]).view( + B * nch, 2, self.band_width[i], -1 + ) + est_spec.append(torch.complex(this_RI[:, 0], this_RI[:, 1])) + est_spec = torch.cat(est_spec, 1) + est_spec = est_spec.to(dtype=torch.complex64) + output = torch.istft( + est_spec, + n_fft=self.win, + hop_length=self.stride, + window=torch.hann_window(self.win).to(input.device), + length=nsample, + ).view(B, nch, -1) + + return output + + def get_model_args(self): + model_args = {"n_sample_rate": 2} + return model_args diff --git a/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/base_model.py b/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/base_model.py new file mode 100644 index 0000000000000000000000000000000000000000..80701cabe029e6b60ccf4f3d19354d52b3d2e486 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/look2hear/models/base_model.py @@ -0,0 +1,104 @@ +### +# Author: Kai Li +# Date: 2021-06-17 23:08:32 +# LastEditors: Please set LastEditors +# LastEditTime: 2022-05-26 18:06:22 +### +import torch +import torch.nn as nn + + +def _unsqueeze_to_3d(x): + """Normalize shape of `x` to [batch, n_chan, time].""" + if x.ndim == 1: + return x.reshape(1, 1, -1) + elif x.ndim == 2: + return x.unsqueeze(1) + else: + return x + + +def pad_to_appropriate_length(x, lcm): + values_to_pad = int(x.shape[-1]) % lcm + if values_to_pad: + appropriate_shape = x.shape + padded_x = torch.zeros( + list(appropriate_shape[:-1]) + + [appropriate_shape[-1] + lcm - values_to_pad], + dtype=torch.float32, + ).to(x.device) + padded_x[..., : x.shape[-1]] = x + return padded_x + return x + + +class BaseModel(nn.Module): + def __init__(self, sample_rate, in_chan=1): + super().__init__() + self._sample_rate = sample_rate + self._in_chan = in_chan + + def forward(self, *args, **kwargs): + raise NotImplementedError + + def sample_rate( + self, + ): + return self._sample_rate + + @staticmethod + def load_state_dict_in_audio(model, pretrained_dict): + model_dict = model.state_dict() + update_dict = {} + for k, v in pretrained_dict.items(): + if "audio_model" in k: + update_dict[k[12:]] = v + model_dict.update(update_dict) + model.load_state_dict(model_dict) + return model + + @staticmethod + def from_pretrain(pretrained_model_conf_or_path, *args, **kwargs): + from . import get + + conf = torch.load( + pretrained_model_conf_or_path, map_location="cpu" + ) # Attempt to find the model and instantiate it. + + model_class = get(conf["model_name"]) + # model_class = get("Conv_TasNet") + model = model_class(*args, **kwargs) + model.load_state_dict(conf["state_dict"]) + return model + + def apollo(*args, **kwargs): + from . import get + + model_class = get("Apollo") + model = model_class(*args, **kwargs) + return model + + def serialize(self): + import pytorch_lightning as pl # Not used in torch.hub + + model_conf = dict( + model_name=self.__class__.__name__, + state_dict=self.get_state_dict(), + model_args=self.get_model_args(), + ) + # Additional infos + infos = dict() + infos["software_versions"] = dict( + torch_version=torch.__version__, + pytorch_lightning_version=pl.__version__, + ) + model_conf["infos"] = infos + return model_conf + + def get_state_dict(self): + """In case the state dict needs to be modified before sharing the model.""" + return self.state_dict() + + def get_model_args(self): + """Should return args to re-instantiate the class.""" + raise NotImplementedError diff --git a/src/third_party/MusicSourceSeparationTraining/models/mdx23c_tfc_tdf_v3.py b/src/third_party/MusicSourceSeparationTraining/models/mdx23c_tfc_tdf_v3.py new file mode 100644 index 0000000000000000000000000000000000000000..ea06256cce7f3c49f7e9463c3685b813c5449b6e --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/mdx23c_tfc_tdf_v3.py @@ -0,0 +1,257 @@ +from functools import partial + +import torch +import torch.nn as nn +from utils.model_utils import prefer_target_instrument + + +class STFT: + def __init__(self, config): + self.n_fft = config.n_fft + self.hop_length = config.hop_length + self.window = torch.hann_window(window_length=self.n_fft, periodic=True) + self.dim_f = config.dim_f + + def __call__(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-2] + c, t = x.shape[-2:] + x = x.reshape([-1, t]) + x = torch.stft( + x, + n_fft=self.n_fft, + hop_length=self.hop_length, + window=window, + center=True, + return_complex=True, + ) + x = torch.view_as_real(x) + x = x.permute([0, 3, 1, 2]) + x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape( + [*batch_dims, c * 2, -1, x.shape[-1]] + ) + return x[..., : self.dim_f, :] + + def inverse(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-3] + c, f, t = x.shape[-3:] + n = self.n_fft // 2 + 1 + f_pad = torch.zeros([*batch_dims, c, n - f, t]).to(x.device) + x = torch.cat([x, f_pad], -2) + x = x.reshape([*batch_dims, c // 2, 2, n, t]).reshape([-1, 2, n, t]) + x = x.permute([0, 2, 3, 1]) + x = x[..., 0] + x[..., 1] * 1.0j + x = torch.istft( + x, n_fft=self.n_fft, hop_length=self.hop_length, window=window, center=True + ) + x = x.reshape([*batch_dims, 2, -1]) + return x + + +def get_norm(norm_type): + def norm(c, norm_type): + if norm_type == "BatchNorm": + return nn.BatchNorm2d(c) + elif norm_type == "InstanceNorm": + return nn.InstanceNorm2d(c, affine=True) + elif "GroupNorm" in norm_type: + g = int(norm_type.replace("GroupNorm", "")) + return nn.GroupNorm(num_groups=g, num_channels=c) + else: + return nn.Identity() + + return partial(norm, norm_type=norm_type) + + +def get_act(act_type): + if act_type == "gelu": + return nn.GELU() + elif act_type == "relu": + return nn.ReLU() + elif act_type[:3] == "elu": + alpha = float(act_type.replace("elu", "")) + return nn.ELU(alpha) + else: + raise Exception + + +class Upscale(nn.Module): + def __init__(self, in_c, out_c, scale, norm, act): + super().__init__() + self.conv = nn.Sequential( + norm(in_c), + act, + nn.ConvTranspose2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=scale, + stride=scale, + bias=False, + ), + ) + + def forward(self, x): + return self.conv(x) + + +class Downscale(nn.Module): + def __init__(self, in_c, out_c, scale, norm, act): + super().__init__() + self.conv = nn.Sequential( + norm(in_c), + act, + nn.Conv2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=scale, + stride=scale, + bias=False, + ), + ) + + def forward(self, x): + return self.conv(x) + + +class TFC_TDF(nn.Module): + def __init__(self, in_c, c, l, f, bn, norm, act): + super().__init__() + + self.blocks = nn.ModuleList() + for i in range(l): + block = nn.Module() + + block.tfc1 = nn.Sequential( + norm(in_c), + act, + nn.Conv2d(in_c, c, 3, 1, 1, bias=False), + ) + block.tdf = nn.Sequential( + norm(c), + act, + nn.Linear(f, f // bn, bias=False), + norm(c), + act, + nn.Linear(f // bn, f, bias=False), + ) + block.tfc2 = nn.Sequential( + norm(c), + act, + nn.Conv2d(c, c, 3, 1, 1, bias=False), + ) + block.shortcut = nn.Conv2d(in_c, c, 1, 1, 0, bias=False) + + self.blocks.append(block) + in_c = c + + def forward(self, x): + for block in self.blocks: + s = block.shortcut(x) + x = block.tfc1(x) + x = x + block.tdf(x) + x = block.tfc2(x) + x = x + s + return x + + +class TFC_TDF_net(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + norm = get_norm(norm_type=config.model.norm) + act = get_act(act_type=config.model.act) + + self.num_target_instruments = len(prefer_target_instrument(config)) + self.num_subbands = config.model.num_subbands + + dim_c = self.num_subbands * config.audio.num_channels * 2 + n = config.model.num_scales + scale = config.model.scale + l = config.model.num_blocks_per_scale + c = config.model.num_channels + g = config.model.growth + bn = config.model.bottleneck_factor + f = config.audio.dim_f // self.num_subbands + + self.first_conv = nn.Conv2d(dim_c, c, 1, 1, 0, bias=False) + + self.encoder_blocks = nn.ModuleList() + for i in range(n): + block = nn.Module() + block.tfc_tdf = TFC_TDF(c, c, l, f, bn, norm, act) + block.downscale = Downscale(c, c + g, scale, norm, act) + f = f // scale[1] + c += g + self.encoder_blocks.append(block) + + self.bottleneck_block = TFC_TDF(c, c, l, f, bn, norm, act) + + self.decoder_blocks = nn.ModuleList() + for i in range(n): + block = nn.Module() + block.upscale = Upscale(c, c - g, scale, norm, act) + f = f * scale[1] + c -= g + block.tfc_tdf = TFC_TDF(2 * c, c, l, f, bn, norm, act) + self.decoder_blocks.append(block) + + self.final_conv = nn.Sequential( + nn.Conv2d(c + dim_c, c, 1, 1, 0, bias=False), + act, + nn.Conv2d(c, self.num_target_instruments * dim_c, 1, 1, 0, bias=False), + ) + + self.stft = STFT(config.audio) + + def cac2cws(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c, k, f // k, t) + x = x.reshape(b, c * k, f // k, t) + return x + + def cws2cac(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c // k, k, f, t) + x = x.reshape(b, c // k, f * k, t) + return x + + def forward(self, x): + x = self.stft(x) + + mix = x = self.cac2cws(x) + + first_conv_out = x = self.first_conv(x) + + x = x.transpose(-1, -2) + + encoder_outputs = [] + for block in self.encoder_blocks: + x = block.tfc_tdf(x) + encoder_outputs.append(x) + x = block.downscale(x) + + x = self.bottleneck_block(x) + + for block in self.decoder_blocks: + x = block.upscale(x) + x = torch.cat([x, encoder_outputs.pop()], 1) + x = block.tfc_tdf(x) + + x = x.transpose(-1, -2) + + x = x * first_conv_out # reduce artifacts + + x = self.final_conv(torch.cat([mix, x], 1)) + + x = self.cws2cac(x) + + b, c, f, t = x.shape + x = x.reshape(b, self.num_target_instruments, -1, f, t) + + x = self.stft.inverse(x) + + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/mdx23c_tfc_tdf_v3_with_STHT.py b/src/third_party/MusicSourceSeparationTraining/models/mdx23c_tfc_tdf_v3_with_STHT.py new file mode 100644 index 0000000000000000000000000000000000000000..f2a636572662f4c6600f8f4056f9ee4240ee9c80 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/mdx23c_tfc_tdf_v3_with_STHT.py @@ -0,0 +1,350 @@ +from functools import partial + +import torch +import torch.nn as nn +import torch.nn.functional as F +from utils.model_utils import prefer_target_instrument + + +class ShortTimeHartleyTransform: + def __init__( + self, + *, + n_fft: int, + hop_length: int, + center: bool = True, + pad_mode: str = "reflect", + ) -> None: + self.n_fft = n_fft + self.hop_length = hop_length + self.center = center + self.pad_mode = pad_mode + self.window = torch.hamming_window(self.n_fft) + + @staticmethod + def _hartley_transform(x: torch.Tensor) -> torch.Tensor: + fft = torch.fft.fft(x) + return fft.real - fft.imag + + @staticmethod + def _inverse_hartley_transform(X: torch.Tensor) -> torch.Tensor: + N = X.size(-1) + return ShortTimeHartleyTransform._hartley_transform(X) / N + + def transform(self, *, signal: torch.Tensor) -> torch.Tensor: + assert signal.dim() == 3, ( + "Signal must be a 3D tensor (batch_size, channel, samples)" + ) + self.window = self.window.to(signal.device) + batch_size, channels, samples = signal.shape + + # Apply padding if center=True + if self.center: + pad_length = self.n_fft // 2 + signal = F.pad(signal, (pad_length, pad_length), mode=self.pad_mode) + else: + pad_length = 0 + + # print( + # f"samples={samples}\n" + # f"self.hop_length={self.hop_length}\n" + # f"pad_length={pad_length}\n" + # f"signal_padded={signal.size(2)}" + # ) + + # Compute number of frames + num_frames = (signal.size(2) - self.n_fft) // self.hop_length + 1 + + # Apply window and compute Hartley transform + window = self.window.to(signal.device, signal.dtype).unsqueeze(0).unsqueeze(0) + stht_coeffs = [] + + for i in range(num_frames): + start = i * self.hop_length + end = start + self.n_fft + frame = signal[:, :, start:end] * window + stht_coeffs.append(self._hartley_transform(frame)) + + return torch.stack(stht_coeffs, dim=-1) + + def inverse_transform( + self, *, stht_coeffs: torch.Tensor, length: int + ) -> torch.Tensor: + self.window = self.window.to(stht_coeffs.device) + # print(stht_coeffs.shape) + batch_size, channels, n_fft, num_frames = stht_coeffs.shape + signal_length = length + + # Initialize reconstruction + reconstructed_signal = torch.zeros( + (batch_size, channels, signal_length + (self.n_fft if self.center else 0)), + device=stht_coeffs.device, + dtype=stht_coeffs.dtype, + ) + normalization = torch.zeros( + signal_length + (self.n_fft if self.center else 0), + device=stht_coeffs.device, + dtype=stht_coeffs.dtype, + ) + + window = ( + self.window.to(stht_coeffs.device, stht_coeffs.dtype) + .unsqueeze(0) + .unsqueeze(0) + ) + + for i in range(num_frames): + start = i * self.hop_length + end = start + self.n_fft + + # Reconstruct frame and add to signal + frame = self._inverse_hartley_transform(stht_coeffs[:, :, :, i]) * window + reconstructed_signal[:, :, start:end] += frame + normalization[start:end] += (window**2).squeeze() + + # Normalize the overlapping regions + eps = torch.finfo(normalization.dtype).eps + normalization = torch.clamp(normalization, min=eps) + reconstructed_signal /= normalization.unsqueeze(0).unsqueeze(0) + + # Remove padding if center=True + if self.center: + pad_length = self.n_fft // 2 + reconstructed_signal = reconstructed_signal[:, :, pad_length:-pad_length] + + # Trim to the specified length + return reconstructed_signal[:, :, :signal_length] + + +def get_norm(norm_type): + def norm(c, norm_type): + if norm_type == "BatchNorm": + return nn.BatchNorm2d(c) + elif norm_type == "InstanceNorm": + return nn.InstanceNorm2d(c, affine=True) + elif "GroupNorm" in norm_type: + g = int(norm_type.replace("GroupNorm", "")) + return nn.GroupNorm(num_groups=g, num_channels=c) + else: + return nn.Identity() + + return partial(norm, norm_type=norm_type) + + +def get_act(act_type): + if act_type == "gelu": + return nn.GELU() + elif act_type == "relu": + return nn.ReLU() + elif act_type[:3] == "elu": + alpha = float(act_type.replace("elu", "")) + return nn.ELU(alpha) + else: + raise Exception + + +class Upscale(nn.Module): + def __init__(self, in_c, out_c, scale, norm, act): + super().__init__() + self.conv = nn.Sequential( + norm(in_c), + act, + nn.ConvTranspose2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=scale, + stride=scale, + bias=False, + ), + ) + + def forward(self, x): + return self.conv(x) + + +class Downscale(nn.Module): + def __init__(self, in_c, out_c, scale, norm, act): + super().__init__() + self.conv = nn.Sequential( + norm(in_c), + act, + nn.Conv2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=scale, + stride=scale, + bias=False, + ), + ) + + def forward(self, x): + return self.conv(x) + + +class TFC_TDF(nn.Module): + def __init__(self, in_c, c, l, f, bn, norm, act): + super().__init__() + + self.blocks = nn.ModuleList() + for i in range(l): + block = nn.Module() + + block.tfc1 = nn.Sequential( + norm(in_c), + act, + nn.Conv2d(in_c, c, 3, 1, 1, bias=False), + ) + block.tdf = nn.Sequential( + norm(c), + act, + nn.Linear(f, f // bn, bias=False), + norm(c), + act, + nn.Linear(f // bn, f, bias=False), + ) + block.tfc2 = nn.Sequential( + norm(c), + act, + nn.Conv2d(c, c, 3, 1, 1, bias=False), + ) + block.shortcut = nn.Conv2d(in_c, c, 1, 1, 0, bias=False) + + self.blocks.append(block) + in_c = c + + def forward(self, x): + for block in self.blocks: + s = block.shortcut(x) + x = block.tfc1(x) + x = x + block.tdf(x) + x = block.tfc2(x) + x = x + s + return x + + +class TFC_TDF_net(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + norm = get_norm(norm_type=config.model.norm) + act = get_act(act_type=config.model.act) + + self.num_target_instruments = len(prefer_target_instrument(config)) + self.num_subbands = config.model.num_subbands + + # dim_c = self.num_subbands * config.audio.num_channels * 2 + dim_c = self.num_subbands * config.audio.num_channels + n = config.model.num_scales + scale = config.model.scale + l = config.model.num_blocks_per_scale + c = config.model.num_channels + g = config.model.growth + bn = config.model.bottleneck_factor + f = config.audio.dim_f // (self.num_subbands // 2) + + self.first_conv = nn.Conv2d(dim_c, c, 1, 1, 0, bias=False) + + self.encoder_blocks = nn.ModuleList() + for i in range(n): + block = nn.Module() + block.tfc_tdf = TFC_TDF(c, c, l, f, bn, norm, act) + block.downscale = Downscale(c, c + g, scale, norm, act) + f = f // scale[1] + c += g + self.encoder_blocks.append(block) + + self.bottleneck_block = TFC_TDF(c, c, l, f, bn, norm, act) + + self.decoder_blocks = nn.ModuleList() + for i in range(n): + block = nn.Module() + block.upscale = Upscale(c, c - g, scale, norm, act) + f = f * scale[1] + c -= g + block.tfc_tdf = TFC_TDF(2 * c, c, l, f, bn, norm, act) + self.decoder_blocks.append(block) + + self.final_conv = nn.Sequential( + nn.Conv2d(c + dim_c, c, 1, 1, 0, bias=False), + act, + nn.Conv2d(c, self.num_target_instruments * dim_c, 1, 1, 0, bias=False), + ) + + self.stft = ShortTimeHartleyTransform( + n_fft=config.audio.n_fft, hop_length=config.audio.hop_length + ) + + def cac2cws(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c, k, f // k, t) + x = x.reshape(b, c * k, f // k, t) + return x + + def cws2cac(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c // k, k, f, t) + x = x.reshape(b, c // k, f * k, t) + return x + + def forward(self, x): + length = x.shape[-1] + # print(x.shape) + x = self.stft.transform(signal=x) + # print(x.shape) + + mix = x = self.cac2cws(x) + + # print(x.shape) + + first_conv_out = x = self.first_conv(x) + + # print(x.shape) + + x = x.transpose(-1, -2) + + # print(x.shape) + + encoder_outputs = [] + for block in self.encoder_blocks: + # print(x.shape) + x = block.tfc_tdf(x) + # print(x.shape) + encoder_outputs.append(x) + x = block.downscale(x) + # print(x.shape) + + x = self.bottleneck_block(x) + # print(x.shape) + + for block in self.decoder_blocks: + # print(x.shape) + x = block.upscale(x) + # print(x.shape) + x = torch.cat([x, encoder_outputs.pop()], 1) + # print(x.shape) + x = block.tfc_tdf(x) + # print(x.shape) + + x = x.transpose(-1, -2) + # print(x.shape) + + x = x * first_conv_out # reduce artifacts + + # print(x.shape) + + x = self.final_conv(torch.cat([mix, x], 1)) + + x = self.cws2cac(x) + + if self.num_target_instruments > 1: + b, c, f, t = x.shape + x = x.reshape(b * self.num_target_instruments, -1, f, t) + x = self.stft.inverse_transform(stht_coeffs=x, length=length) + x = x.reshape(b, self.num_target_instruments, x.shape[-2], x.shape[-1]) + else: + x = self.stft.inverse_transform(stht_coeffs=x, length=length) + # print("!!!", x.shape) + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/mel_band_conformer.py b/src/third_party/MusicSourceSeparationTraining/models/mel_band_conformer.py new file mode 100644 index 0000000000000000000000000000000000000000..37051f5b0ddbebfbb750a9c8456a73d60f678e98 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/mel_band_conformer.py @@ -0,0 +1,458 @@ +from functools import partial + +import torch +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, Optional, Tuple +from conformer import Conformer +from einops import rearrange, reduce, repeat +from librosa import filters +from torch import nn +from torch.nn import Module, ModuleList + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# attention + + +def MLP(dim_in, dim_out, dim_hidden=None, depth=1, activation=nn.Tanh): + dim_hidden = default(dim_hidden, dim_in) + + net = [] + dims = (dim_in, *((dim_hidden,) * depth), dim_out) + + for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): + is_last = ind == (len(dims) - 2) + + net.append(nn.Linear(layer_dim_in, layer_dim_out)) + + if is_last: + continue + + net.append(activation()) + + return nn.Sequential(*net) + + +class MaskEstimator(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...], depth, mlp_expansion_factor=4): + super().__init__() + self.dim_inputs = dim_inputs + self.to_freqs = ModuleList([]) + dim_hidden = dim * mlp_expansion_factor + + for dim_in in dim_inputs: + net = [] + + mlp = nn.Sequential( + MLP(dim, dim_in * 2, dim_hidden=dim_hidden, depth=depth), nn.GLU(dim=-1) + ) + + self.to_freqs.append(mlp) + + def forward(self, x): + # split along band dimension and run per-band MLP + x = x.unbind(dim=-2) + + outs = [] + + for band_features, mlp in zip(x, self.to_freqs): + freq_out = mlp(band_features) + outs.append(freq_out) + + return torch.cat(outs, dim=-1) + + +class BandSplit(Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + + self.to_features.append(net) + + def forward(self, x): + # split input into predefined frequency-band chunks + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + # stack back as (bands) axis + return torch.stack(outs, dim=-2) + + +class MelBandConformer(nn.Module): + def __init__( + self, + dim: int, + *, + depth: int, + stereo: bool = False, + num_stems: int = 1, + time_conformer_depth: int = 2, + freq_conformer_depth: int = 2, + num_bands: int = 60, + dim_head: int = 64, + heads: int = 8, + # Conformer params + ff_mult: int = 4, + conv_expansion_factor: int = 2, + conv_kernel_size: int = 31, + attn_dropout: float = 0.0, + ff_dropout: float = 0.0, + conv_dropout: float = 0.0, + # STFT + dim_freqs_in: int = 1025, + sample_rate: int = 44100, + stft_n_fft: int = 2048, + stft_hop_length: int = 512, + stft_win_length: int = 2048, + stft_normalized: bool = False, + stft_window_fn: Optional[Callable] = None, + # Loss + mask_estimator_depth: int = 1, + multi_stft_resolution_loss_weight: float = 1.0, + multi_stft_resolutions_window_sizes: Tuple[int, ...] = ( + 4096, + 2048, + 1024, + 512, + 256, + ), + multi_stft_hop_size: int = 147, + multi_stft_normalized: bool = False, + multi_stft_window_fn: Callable = torch.hann_window, + match_input_audio_length: bool = False, + use_torch_checkpoint: bool = False, + skip_connection: bool = False, + ): + super().__init__() + + self.stereo = stereo + self.audio_channels = 2 if stereo else 1 + self.num_stems = num_stems + self.use_torch_checkpoint = use_torch_checkpoint + self.skip_connection = skip_connection + + self.layers = nn.ModuleList([]) + + # Layers per block: [ time-Conformer, freq-Conformer ] + conformer_kwargs = dict( + dim=dim, + dim_head=dim_head, + heads=heads, + ff_mult=ff_mult, + conv_expansion_factor=conv_expansion_factor, + conv_kernel_size=conv_kernel_size, + attn_dropout=attn_dropout, + ff_dropout=ff_dropout, + conv_dropout=conv_dropout, + ) + + for _ in range(depth): + time_block = Conformer(depth=time_conformer_depth, **conformer_kwargs) + freq_block = Conformer(depth=freq_conformer_depth, **conformer_kwargs) + self.layers.append(nn.ModuleList([time_block, freq_block])) + + self.stft_window_fn = partial( + stft_window_fn or torch.hann_window, stft_win_length + ) + + self.stft_kwargs = dict( + n_fft=stft_n_fft, + hop_length=stft_hop_length, + win_length=stft_win_length, + normalized=stft_normalized, + ) + + # number of frequency bins produced by STFT (ignoring complex axis) + freqs = torch.stft( + torch.randn(1, 4096), + **self.stft_kwargs, + window=torch.ones(stft_n_fft), + return_complex=True, + ).shape[1] + + # build mel filter bank to define band grouping + mel_filter_bank_numpy = filters.mel( + sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands + ) + mel_filter_bank = torch.from_numpy(mel_filter_bank_numpy) + # ensure coverage at the boundaries + mel_filter_bank[0][0] = 1.0 + mel_filter_bank[-1, -1] = 1.0 + + freqs_per_band = mel_filter_bank > 0 + assert freqs_per_band.any(dim=0).all(), ( + "all frequency bins must be covered by bands" + ) + + repeated_freq_indices = repeat(torch.arange(freqs), "f -> b f", b=num_bands) + freq_indices = repeated_freq_indices[freqs_per_band] + + if stereo: + # duplicate indices for stereo by interleaving channels along the freq axis + freq_indices = repeat(freq_indices, "f -> f s", s=2) + freq_indices = freq_indices * 2 + torch.arange(2) + freq_indices = rearrange(freq_indices, "f s -> (f s)") + + self.register_buffer("freq_indices", freq_indices, persistent=False) + self.register_buffer("freqs_per_band", freqs_per_band, persistent=False) + + num_freqs_per_band = reduce(freqs_per_band, "b f -> b", "sum") + num_bands_per_freq = reduce(freqs_per_band, "b f -> f", "sum") + + self.register_buffer("num_freqs_per_band", num_freqs_per_band, persistent=False) + self.register_buffer("num_bands_per_freq", num_bands_per_freq, persistent=False) + + # BandSplit and MaskEstimator — same structure as your original + freqs_per_bands_with_complex = tuple( + 2 * f * self.audio_channels for f in num_freqs_per_band.tolist() + ) + + self.band_split = BandSplit(dim=dim, dim_inputs=freqs_per_bands_with_complex) + + self.mask_estimators = nn.ModuleList( + [ + MaskEstimator( + dim=dim, + dim_inputs=freqs_per_bands_with_complex, + depth=mask_estimator_depth, + mlp_expansion_factor=4, # could be exposed as a parameter + ) + for _ in range(num_stems) + ] + ) + + # multi-resolution STFT loss setup + self.multi_stft_resolution_loss_weight = multi_stft_resolution_loss_weight + self.multi_stft_resolutions_window_sizes = multi_stft_resolutions_window_sizes + self.multi_stft_n_fft = stft_n_fft + self.multi_stft_window_fn = multi_stft_window_fn + + self.multi_stft_kwargs = dict( + hop_length=multi_stft_hop_size, normalized=multi_stft_normalized + ) + + self.match_input_audio_length = match_input_audio_length + + def forward( + self, + raw_audio: torch.Tensor, + target: Optional[torch.Tensor] = None, + return_loss_breakdown: bool = False, + ): + """ + b - batch + f - freq + t - time + s - audio channel (1 mono / 2 stereo) + n - stems + c - complex (2) + d - feature dim + """ + device = raw_audio.device + + if raw_audio.ndim == 2: + raw_audio = rearrange(raw_audio, "b t -> b 1 t") + + batch, channels, raw_audio_length = raw_audio.shape + istft_length = raw_audio_length if self.match_input_audio_length else None + + assert (not self.stereo and channels == 1) or (self.stereo and channels == 2), ( + "set stereo=True for stereo input (C=2), stereo=False for mono (C=1)" + ) + + # --- STFT --- + raw_audio_flat, packed_shape = ( + raw_audio.reshape(-1, raw_audio.shape[-1]), + raw_audio.shape[:2], + ) + stft_window = self.stft_window_fn(device=device) + + stft_repr = torch.stft( + raw_audio_flat, **self.stft_kwargs, window=stft_window, return_complex=True + ) + stft_repr = torch.view_as_real(stft_repr) # (B*C, F, T, 2) + stft_repr = stft_repr.view( + *packed_shape, *stft_repr.shape[1:] + ) # (b, s, f, t, c) + + # fold channel into frequency axis (as in your setup) + stft_repr_fs = rearrange(stft_repr, "b s f t c -> b (f s) t c") + + # index frequencies by mel bands + b_idx = torch.arange(batch, device=device)[..., None] + x = stft_repr_fs[b_idx, self.freq_indices] # (b, sum(freqs_in_bands), t, c) + x = rearrange(x, "b f t c -> b t (f c)") # flatten complex axis into features + + # --- BandSplit -> (b, t, bands, dim) --- + if self.use_torch_checkpoint: + x = torch.utils.checkpoint.checkpoint( + self.band_split, x, use_reentrant=False + ) + else: + x = self.band_split(x) + + # --- Axial Conformer (time, then freq) --- + store = [None] * len(self.layers) + + for i, (time_conf, freq_conf) in enumerate(self.layers): + # Time axis: (b, t, bands, d) -> ((b*bands), t, d) + bsz, tlen, bands, d = x.shape + x_time = rearrange(x, "b t f d -> (b f) t d") + + if self.use_torch_checkpoint: + x_time = torch.utils.checkpoint.checkpoint( + time_conf, x_time, use_reentrant=False + ) + else: + x_time = time_conf(x_time) + + x = rearrange(x_time, "(b f) t d -> b t f d", b=bsz, f=bands) + + # Freq axis: (b, t, f, d) -> ((b*t), f, d) + bsz, tlen, bands, d = x.shape + x_freq = rearrange(x, "b t f d -> (b t) f d") + + if self.use_torch_checkpoint: + x_freq = torch.utils.checkpoint.checkpoint( + freq_conf, x_freq, use_reentrant=False + ) + else: + x_freq = freq_conf(x_freq) + + x = rearrange(x_freq, "(b t) f d -> b t f d", b=bsz, t=tlen) + + if self.skip_connection: + store[i] = x if store[i] is None else store[i] + x + + # --- Mask estimation --- + # (b, t, f_bands, d) -> per-stem MLP over bands + if self.use_torch_checkpoint: + masks = torch.stack( + [ + torch.utils.checkpoint.checkpoint(fn, x, use_reentrant=False) + for fn in self.mask_estimators + ], + dim=1, + ) + else: + masks = torch.stack([fn(x) for fn in self.mask_estimators], dim=1) + masks = rearrange(masks, "b n t (f c) -> b n f t c", c=2) + + # --- Complex modulation --- + stft_repr_c = rearrange(stft_repr, "b s f t c -> b 1 (f s) t c") + stft_repr_c = torch.view_as_complex(stft_repr_c) # (b, 1, F*S, T) + masks_c = torch.view_as_complex(masks) # (b, n, F*S, T) + + masks_c = masks_c.type(stft_repr_c.dtype) + + scatter_idx = repeat( + self.freq_indices, + "f -> b n f t", + b=batch, + n=self.num_stems, + t=stft_repr_c.shape[-1], + ) + stft_repr_expanded = repeat(stft_repr_c, "b 1 ... -> b n ...", n=self.num_stems) + + masks_summed = torch.zeros_like(stft_repr_expanded).scatter_add_( + 2, scatter_idx, masks_c + ) + denom = repeat(self.num_bands_per_freq, "f -> (f r) 1", r=self.audio_channels) + + masks_averaged = masks_summed / denom.clamp(min=1e-8) + stft_mod = stft_repr_c * masks_averaged + + # --- iSTFT --- + stft_mod = rearrange( + stft_mod, "b n (f s) t -> (b n s) f t", s=self.audio_channels + ) + + recon_audio = torch.istft( + stft_mod, + **self.stft_kwargs, + window=stft_window, + return_complex=False, + length=istft_length, + ) + recon_audio = rearrange( + recon_audio, + "(b n s) t -> b n s t", + b=batch, + s=self.audio_channels, + n=self.num_stems, + ) + + if self.num_stems == 1: + recon_audio = rearrange(recon_audio, "b 1 s t -> b s t") + + # Loss + if target is None: + return recon_audio + + if self.num_stems > 1: + assert target.ndim == 4 and target.shape[1] == self.num_stems + + if target.ndim == 2: + target = rearrange(target, "... t -> ... 1 t") + + target = target[..., : recon_audio.shape[-1]] + + loss = F.l1_loss(recon_audio, target) + + multi_stft_resolution_loss = 0.0 + for window_size in self.multi_stft_resolutions_window_sizes: + res_stft_kwargs = dict( + n_fft=max(window_size, self.multi_stft_n_fft), + win_length=window_size, + return_complex=True, + window=self.multi_stft_window_fn(window_size, device=device), + **self.multi_stft_kwargs, + ) + + recon_Y = torch.stft( + rearrange(recon_audio, "... s t -> (... s) t"), **res_stft_kwargs + ) + target_Y = torch.stft( + rearrange(target, "... s t -> (... s) t"), **res_stft_kwargs + ) + + multi_stft_resolution_loss += F.l1_loss(recon_Y, target_Y) + + total_loss = ( + loss + self.multi_stft_resolution_loss_weight * multi_stft_resolution_loss + ) + + if not return_loss_breakdown: + return total_loss + + return total_loss, (loss, multi_stft_resolution_loss) diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/scnet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5f6ecefede9345237623066dd21ebd8253af1c60 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet/__init__.py @@ -0,0 +1 @@ +from .scnet import SCNet diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet.py b/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet.py new file mode 100644 index 0000000000000000000000000000000000000000..d6f26943f0936309070298035a66ce53ca352099 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet.py @@ -0,0 +1,420 @@ +import math +from collections import deque + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .separation import SeparationNet + + +class Swish(nn.Module): + def forward(self, x): + return x * x.sigmoid() + + +class ConvolutionModule(nn.Module): + """ + Convolution Module in SD block. + + Args: + channels (int): input/output channels. + depth (int): number of layers in the residual branch. Each layer has its own + compress (float): amount of channel compression. + kernel (int): kernel size for the convolutions. + """ + + def __init__(self, channels, depth=2, compress=4, kernel=3): + super().__init__() + assert kernel % 2 == 1 + self.depth = abs(depth) + hidden_size = int(channels / compress) + norm = lambda d: nn.GroupNorm(1, d) + self.layers = nn.ModuleList([]) + for _ in range(self.depth): + padding = kernel // 2 + mods = [ + norm(channels), + nn.Conv1d(channels, hidden_size * 2, kernel, padding=padding), + nn.GLU(1), + nn.Conv1d( + hidden_size, + hidden_size, + kernel, + padding=padding, + groups=hidden_size, + ), + norm(hidden_size), + Swish(), + nn.Conv1d(hidden_size, channels, 1), + ] + layer = nn.Sequential(*mods) + self.layers.append(layer) + + def forward(self, x): + for layer in self.layers: + x = x + layer(x) + return x + + +class FusionLayer(nn.Module): + """ + A FusionLayer within the decoder. + + Args: + - channels (int): Number of input channels. + - kernel_size (int, optional): Kernel size for the convolutional layer, defaults to 3. + - stride (int, optional): Stride for the convolutional layer, defaults to 1. + - padding (int, optional): Padding for the convolutional layer, defaults to 1. + """ + + def __init__(self, channels, kernel_size=3, stride=1, padding=1): + super(FusionLayer, self).__init__() + self.conv = nn.Conv2d( + channels * 2, channels * 2, kernel_size, stride=stride, padding=padding + ) + + def forward(self, x, skip=None): + if skip is not None: + x += skip + x = x.repeat(1, 2, 1, 1) + x = self.conv(x) + x = F.glu(x, dim=1) + return x + + +class SDlayer(nn.Module): + """ + Implements a Sparse Down-sample Layer for processing different frequency bands separately. + + Args: + - channels_in (int): Input channel count. + - channels_out (int): Output channel count. + - band_configs (dict): A dictionary containing configuration for each frequency band. + Keys are 'low', 'mid', 'high' for each band, and values are + dictionaries with keys 'SR', 'stride', and 'kernel' for proportion, + stride, and kernel size, respectively. + """ + + def __init__(self, channels_in, channels_out, band_configs): + super(SDlayer, self).__init__() + + # Initializing convolutional layers for each band + self.convs = nn.ModuleList() + self.strides = [] + self.kernels = [] + for config in band_configs.values(): + self.convs.append( + nn.Conv2d( + channels_in, + channels_out, + (config["kernel"], 1), + (config["stride"], 1), + (0, 0), + ) + ) + self.strides.append(config["stride"]) + self.kernels.append(config["kernel"]) + + # Saving rate proportions for determining splits + self.SR_low = band_configs["low"]["SR"] + self.SR_mid = band_configs["mid"]["SR"] + + def forward(self, x): + B, C, Fr, T = x.shape + # Define splitting points based on sampling rates + splits = [ + (0, math.ceil(Fr * self.SR_low)), + (math.ceil(Fr * self.SR_low), math.ceil(Fr * (self.SR_low + self.SR_mid))), + (math.ceil(Fr * (self.SR_low + self.SR_mid)), Fr), + ] + + # Processing each band with the corresponding convolution + outputs = [] + original_lengths = [] + for conv, stride, kernel, (start, end) in zip( + self.convs, self.strides, self.kernels, splits + ): + extracted = x[:, :, start:end, :] + original_lengths.append(end - start) + current_length = extracted.shape[2] + + # padding + if stride == 1: + total_padding = kernel - stride + else: + total_padding = (stride - current_length % stride) % stride + pad_left = total_padding // 2 + pad_right = total_padding - pad_left + + padded = F.pad(extracted, (0, 0, pad_left, pad_right)) + + output = conv(padded) + outputs.append(output) + + return outputs, original_lengths + + +class SUlayer(nn.Module): + """ + Implements a Sparse Up-sample Layer in decoder. + + Args: + - channels_in: The number of input channels. + - channels_out: The number of output channels. + - convtr_configs: Dictionary containing the configurations for transposed convolutions. + """ + + def __init__(self, channels_in, channels_out, band_configs): + super(SUlayer, self).__init__() + + # Initializing convolutional layers for each band + self.convtrs = nn.ModuleList( + [ + nn.ConvTranspose2d( + channels_in, + channels_out, + [config["kernel"], 1], + [config["stride"], 1], + ) + for _, config in band_configs.items() + ] + ) + + def forward(self, x, lengths, origin_lengths): + B, C, Fr, T = x.shape + # Define splitting points based on input lengths + splits = [ + (0, lengths[0]), + (lengths[0], lengths[0] + lengths[1]), + (lengths[0] + lengths[1], None), + ] + # Processing each band with the corresponding convolution + outputs = [] + for idx, (convtr, (start, end)) in enumerate(zip(self.convtrs, splits)): + out = convtr(x[:, :, start:end, :]) + # Calculate the distance to trim the output symmetrically to original length + current_Fr_length = out.shape[2] + dist = abs(origin_lengths[idx] - current_Fr_length) // 2 + + # Trim the output to the original length symmetrically + trimmed_out = out[:, :, dist : dist + origin_lengths[idx], :] + + outputs.append(trimmed_out) + + # Concatenate trimmed outputs along the frequency dimension to return the final tensor + x = torch.cat(outputs, dim=2) + + return x + + +class SDblock(nn.Module): + """ + Implements a simplified Sparse Down-sample block in encoder. + + Args: + - channels_in (int): Number of input channels. + - channels_out (int): Number of output channels. + - band_config (dict): Configuration for the SDlayer specifying band splits and convolutions. + - conv_config (dict): Configuration for convolution modules applied to each band. + - depths (list of int): List specifying the convolution depths for low, mid, and high frequency bands. + """ + + def __init__( + self, + channels_in, + channels_out, + band_configs={}, + conv_config={}, + depths=[3, 2, 1], + kernel_size=3, + ): + super(SDblock, self).__init__() + self.SDlayer = SDlayer(channels_in, channels_out, band_configs) + + # Dynamically create convolution modules for each band based on depths + self.conv_modules = nn.ModuleList( + [ConvolutionModule(channels_out, depth, **conv_config) for depth in depths] + ) + # Set the kernel_size to an odd number. + self.globalconv = nn.Conv2d( + channels_out, channels_out, kernel_size, 1, (kernel_size - 1) // 2 + ) + + def forward(self, x): + bands, original_lengths = self.SDlayer(x) + # B, C, f, T = band.shape + bands = [ + F.gelu( + conv(band.permute(0, 2, 1, 3).reshape(-1, band.shape[1], band.shape[3])) + .view(band.shape[0], band.shape[2], band.shape[1], band.shape[3]) + .permute(0, 2, 1, 3) + ) + for conv, band in zip(self.conv_modules, bands) + ] + lengths = [band.size(-2) for band in bands] + full_band = torch.cat(bands, dim=2) + skip = full_band + + output = self.globalconv(full_band) + + return output, skip, lengths, original_lengths + + +class SCNet(nn.Module): + """ + The implementation of SCNet: Sparse Compression Network for Music Source Separation. Paper: https://arxiv.org/abs/2401.13276.pdf + + Args: + - sources (List[str]): List of sources to be separated. + - audio_channels (int): Number of audio channels. + - nfft (int): Number of FFTs to determine the frequency dimension of the input. + - hop_size (int): Hop size for the STFT. + - win_size (int): Window size for STFT. + - normalized (bool): Whether to normalize the STFT. + - dims (List[int]): List of channel dimensions for each block. + - band_SR (List[float]): The proportion of each frequency band. + - band_stride (List[int]): The down-sampling ratio of each frequency band. + - band_kernel (List[int]): The kernel sizes for down-sampling convolution in each frequency band + - conv_depths (List[int]): List specifying the number of convolution modules in each SD block. + - compress (int): Compression factor for convolution module. + - conv_kernel (int): Kernel size for convolution layer in convolution module. + - num_dplayer (int): Number of dual-path layers. + - expand (int): Expansion factor in the dual-path RNN, default is 1. + + """ + + def __init__( + self, + sources=["drums", "bass", "other", "vocals"], + audio_channels=2, + # Main structure + dims=[4, 32, 64, 128], # dims = [4, 64, 128, 256] in SCNet-large + # STFT + nfft=4096, + hop_size=1024, + win_size=4096, + normalized=True, + # SD/SU layer + band_SR=[0.175, 0.392, 0.433], + band_stride=[1, 4, 16], + band_kernel=[3, 4, 16], + # Convolution Module + conv_depths=[3, 2, 1], + compress=4, + conv_kernel=3, + # Dual-path RNN + num_dplayer=6, + expand=1, + ): + super().__init__() + self.sources = sources + self.audio_channels = audio_channels + self.dims = dims + band_keys = ["low", "mid", "high"] + self.band_configs = { + band_keys[i]: { + "SR": band_SR[i], + "stride": band_stride[i], + "kernel": band_kernel[i], + } + for i in range(len(band_keys)) + } + self.hop_length = hop_size + self.conv_config = { + "compress": compress, + "kernel": conv_kernel, + } + + self.stft_config = { + "n_fft": nfft, + "hop_length": hop_size, + "win_length": win_size, + "center": True, + "normalized": normalized, + } + + self.encoder = nn.ModuleList() + self.decoder = nn.ModuleList() + + for index in range(len(dims) - 1): + enc = SDblock( + channels_in=dims[index], + channels_out=dims[index + 1], + band_configs=self.band_configs, + conv_config=self.conv_config, + depths=conv_depths, + ) + self.encoder.append(enc) + + dec = nn.Sequential( + FusionLayer(channels=dims[index + 1]), + SUlayer( + channels_in=dims[index + 1], + channels_out=dims[index] + if index != 0 + else dims[index] * len(sources), + band_configs=self.band_configs, + ), + ) + self.decoder.insert(0, dec) + + self.separation_net = SeparationNet( + channels=dims[-1], + expand=expand, + num_layers=num_dplayer, + ) + + def forward(self, x): + # B, C, L = x.shape + B = x.shape[0] + # In the initial padding, ensure that the number of frames after the STFT (the length of the T dimension) is even, + # so that the RFFT operation can be used in the separation network. + padding = self.hop_length - x.shape[-1] % self.hop_length + if (x.shape[-1] + padding) // self.hop_length % 2 == 0: + padding += self.hop_length + x = F.pad(x, (0, padding)) + + # STFT + L = x.shape[-1] + x = x.reshape(-1, L) + x = torch.stft(x, **self.stft_config, return_complex=True) + x = torch.view_as_real(x) + x = x.permute(0, 3, 1, 2).reshape( + x.shape[0] // self.audio_channels, + x.shape[3] * self.audio_channels, + x.shape[1], + x.shape[2], + ) + + B, C, Fr, T = x.shape + + save_skip = deque() + save_lengths = deque() + save_original_lengths = deque() + # encoder + for sd_layer in self.encoder: + x, skip, lengths, original_lengths = sd_layer(x) + save_skip.append(skip) + save_lengths.append(lengths) + save_original_lengths.append(original_lengths) + + # separation + x = self.separation_net(x) + + # decoder + for fusion_layer, su_layer in self.decoder: + x = fusion_layer(x, save_skip.pop()) + x = su_layer(x, save_lengths.pop(), save_original_lengths.pop()) + + # output + n = self.dims[0] + x = x.view(B, n, -1, Fr, T) + x = x.reshape(-1, 2, Fr, T).permute(0, 2, 3, 1) + x = torch.view_as_complex(x.contiguous()) + x = torch.istft(x, **self.stft_config) + x = x.reshape(B, len(self.sources), self.audio_channels, -1) + + x = x[:, :, :, :-padding] + + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet_masked.py b/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet_masked.py new file mode 100644 index 0000000000000000000000000000000000000000..26e98b43d5462dd61c0ff77bed7eeba236f49f40 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet_masked.py @@ -0,0 +1,466 @@ +import math +from collections import deque + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .separation import SeparationNet + + +class Swish(nn.Module): + def forward(self, x): + return x * x.sigmoid() + + +class ConvolutionModule(nn.Module): + """ + Convolution Module in SD block. + + Args: + channels (int): input/output channels. + depth (int): number of layers in the residual branch. Each layer has its own + compress (float): amount of channel compression. + kernel (int): kernel size for the convolutions. + """ + + def __init__(self, channels, depth=2, compress=4, kernel=3): + super().__init__() + assert kernel % 2 == 1 + self.depth = abs(depth) + hidden_size = int(channels / compress) + norm = lambda d: nn.GroupNorm(1, d) + self.layers = nn.ModuleList([]) + for _ in range(self.depth): + padding = kernel // 2 + mods = [ + norm(channels), + nn.Conv1d(channels, hidden_size * 2, kernel, padding=padding), + nn.GLU(1), + nn.Conv1d( + hidden_size, + hidden_size, + kernel, + padding=padding, + groups=hidden_size, + ), + norm(hidden_size), + Swish(), + nn.Conv1d(hidden_size, channels, 1), + ] + layer = nn.Sequential(*mods) + self.layers.append(layer) + + def forward(self, x): + for layer in self.layers: + x = x + layer(x) + return x + + +class FusionLayer(nn.Module): + """ + A FusionLayer within the decoder. + + Args: + - channels (int): Number of input channels. + - kernel_size (int, optional): Kernel size for the convolutional layer, defaults to 3. + - stride (int, optional): Stride for the convolutional layer, defaults to 1. + - padding (int, optional): Padding for the convolutional layer, defaults to 1. + """ + + def __init__(self, channels, kernel_size=3, stride=1, padding=1): + super(FusionLayer, self).__init__() + self.conv = nn.Conv2d( + channels * 2, channels * 2, kernel_size, stride=stride, padding=padding + ) + + def forward(self, x, skip=None): + if skip is not None: + x += skip + x = x.repeat(1, 2, 1, 1) + x = self.conv(x) + x = F.glu(x, dim=1) + return x + + +class SDlayer(nn.Module): + """ + Implements a Sparse Down-sample Layer for processing different frequency bands separately. + + Args: + - channels_in (int): Input channel count. + - channels_out (int): Output channel count. + - band_configs (dict): A dictionary containing configuration for each frequency band. + Keys are 'low', 'mid', 'high' for each band, and values are + dictionaries with keys 'SR', 'stride', and 'kernel' for proportion, + stride, and kernel size, respectively. + """ + + def __init__(self, channels_in, channels_out, band_configs): + super(SDlayer, self).__init__() + + # Initializing convolutional layers for each band + self.convs = nn.ModuleList() + self.strides = [] + self.kernels = [] + for config in band_configs.values(): + self.convs.append( + nn.Conv2d( + channels_in, + channels_out, + (config["kernel"], 1), + (config["stride"], 1), + (0, 0), + ) + ) + self.strides.append(config["stride"]) + self.kernels.append(config["kernel"]) + + # Saving rate proportions for determining splits + self.SR_low = band_configs["low"]["SR"] + self.SR_mid = band_configs["mid"]["SR"] + + def forward(self, x): + B, C, Fr, T = x.shape + # Define splitting points based on sampling rates + splits = [ + (0, math.ceil(Fr * self.SR_low)), + (math.ceil(Fr * self.SR_low), math.ceil(Fr * (self.SR_low + self.SR_mid))), + (math.ceil(Fr * (self.SR_low + self.SR_mid)), Fr), + ] + + # Processing each band with the corresponding convolution + outputs = [] + original_lengths = [] + for conv, stride, kernel, (start, end) in zip( + self.convs, self.strides, self.kernels, splits + ): + extracted = x[:, :, start:end, :] + original_lengths.append(end - start) + current_length = extracted.shape[2] + + # padding + if stride == 1: + total_padding = kernel - stride + else: + total_padding = (stride - current_length % stride) % stride + pad_left = total_padding // 2 + pad_right = total_padding - pad_left + + padded = F.pad(extracted, (0, 0, pad_left, pad_right)) + + output = conv(padded) + outputs.append(output) + + return outputs, original_lengths + + +class SUlayer(nn.Module): + """ + Implements a Sparse Up-sample Layer in decoder. + + Args: + - channels_in: The number of input channels. + - channels_out: The number of output channels. + - convtr_configs: Dictionary containing the configurations for transposed convolutions. + """ + + def __init__(self, channels_in, channels_out, band_configs): + super(SUlayer, self).__init__() + + # Initializing convolutional layers for each band + self.convtrs = nn.ModuleList( + [ + nn.ConvTranspose2d( + channels_in, + channels_out, + [config["kernel"], 1], + [config["stride"], 1], + ) + for _, config in band_configs.items() + ] + ) + + def forward(self, x, lengths, origin_lengths): + B, C, Fr, T = x.shape + # Define splitting points based on input lengths + splits = [ + (0, lengths[0]), + (lengths[0], lengths[0] + lengths[1]), + (lengths[0] + lengths[1], None), + ] + # Processing each band with the corresponding convolution + outputs = [] + for idx, (convtr, (start, end)) in enumerate(zip(self.convtrs, splits)): + out = convtr(x[:, :, start:end, :]) + # Calculate the distance to trim the output symmetrically to original length + current_Fr_length = out.shape[2] + dist = abs(origin_lengths[idx] - current_Fr_length) // 2 + + # Trim the output to the original length symmetrically + trimmed_out = out[:, :, dist : dist + origin_lengths[idx], :] + + outputs.append(trimmed_out) + + # Concatenate trimmed outputs along the frequency dimension to return the final tensor + x = torch.cat(outputs, dim=2) + + return x + + +class SDblock(nn.Module): + """ + Implements a simplified Sparse Down-sample block in encoder. + + Args: + - channels_in (int): Number of input channels. + - channels_out (int): Number of output channels. + - band_config (dict): Configuration for the SDlayer specifying band splits and convolutions. + - conv_config (dict): Configuration for convolution modules applied to each band. + - depths (list of int): List specifying the convolution depths for low, mid, and high frequency bands. + """ + + def __init__( + self, + channels_in, + channels_out, + band_configs={}, + conv_config={}, + depths=[3, 2, 1], + kernel_size=3, + ): + super(SDblock, self).__init__() + self.SDlayer = SDlayer(channels_in, channels_out, band_configs) + + # Dynamically create convolution modules for each band based on depths + self.conv_modules = nn.ModuleList( + [ConvolutionModule(channels_out, depth, **conv_config) for depth in depths] + ) + # Set the kernel_size to an odd number. + self.globalconv = nn.Conv2d( + channels_out, channels_out, kernel_size, 1, (kernel_size - 1) // 2 + ) + + def forward(self, x): + bands, original_lengths = self.SDlayer(x) + # B, C, f, T = band.shape + bands = [ + F.gelu( + conv(band.permute(0, 2, 1, 3).reshape(-1, band.shape[1], band.shape[3])) + .view(band.shape[0], band.shape[2], band.shape[1], band.shape[3]) + .permute(0, 2, 1, 3) + ) + for conv, band in zip(self.conv_modules, bands) + ] + lengths = [band.size(-2) for band in bands] + full_band = torch.cat(bands, dim=2) + skip = full_band + + output = self.globalconv(full_band) + + return output, skip, lengths, original_lengths + + +class SCNet(nn.Module): + """ + The implementation of SCNet: Sparse Compression Network for Music Source Separation. Paper: https://arxiv.org/abs/2401.13276.pdf + + Args: + - sources (List[str]): List of sources to be separated. + - audio_channels (int): Number of audio channels. + - nfft (int): Number of FFTs to determine the frequency dimension of the input. + - hop_size (int): Hop size for the STFT. + - win_size (int): Window size for STFT. + - normalized (bool): Whether to normalize the STFT. + - dims (List[int]): List of channel dimensions for each block. + - band_SR (List[float]): The proportion of each frequency band. + - band_stride (List[int]): The down-sampling ratio of each frequency band. + - band_kernel (List[int]): The kernel sizes for down-sampling convolution in each frequency band + - conv_depths (List[int]): List specifying the number of convolution modules in each SD block. + - compress (int): Compression factor for convolution module. + - conv_kernel (int): Kernel size for convolution layer in convolution module. + - num_dplayer (int): Number of dual-path layers. + - expand (int): Expansion factor in the dual-path RNN, default is 1. + + """ + + def __init__( + self, + sources=["drums", "bass", "other", "vocals"], + audio_channels=2, + # Main structure + dims=[4, 32, 64, 128], # dims = [4, 64, 128, 256] in SCNet-large + # STFT + nfft=4096, + hop_size=1024, + win_size=4096, + normalized=True, + # SD/SU layer + band_SR=[0.175, 0.392, 0.433], + band_stride=[1, 4, 16], + band_kernel=[3, 4, 16], + # Convolution Module + conv_depths=[3, 2, 1], + compress=4, + conv_kernel=3, + # Dual-path RNN + num_dplayer=6, + expand=1, + ): + super().__init__() + self.sources = sources + self.audio_channels = audio_channels + self.dims = dims + band_keys = ["low", "mid", "high"] + self.band_configs = { + band_keys[i]: { + "SR": band_SR[i], + "stride": band_stride[i], + "kernel": band_kernel[i], + } + for i in range(len(band_keys)) + } + self.hop_length = hop_size + self.conv_config = { + "compress": compress, + "kernel": conv_kernel, + } + + self.embed_dim = dims[0] + self.max_f = nfft // 2 + 1 + self.pos_embed_f = nn.Parameter(torch.zeros(1, self.embed_dim, self.max_f, 1)) + nn.init.trunc_normal_(self.pos_embed_f, std=0.02) + + window = torch.hann_window(window_length=nfft, periodic=True) + self.register_buffer("window", window, persistent=False) + + self.stft_config = { + "n_fft": nfft, + "hop_length": hop_size, + "win_length": win_size, + "center": True, + "normalized": normalized, + } + + self.encoder = nn.ModuleList() + self.decoder = nn.ModuleList() + + for index in range(len(dims) - 1): + enc = SDblock( + channels_in=dims[index], + channels_out=dims[index + 1], + band_configs=self.band_configs, + conv_config=self.conv_config, + depths=conv_depths, + ) + self.encoder.append(enc) + + dec = nn.Sequential( + FusionLayer(channels=dims[index + 1]), + SUlayer( + channels_in=dims[index + 1], + channels_out=dims[index] + if index != 0 + else dims[index] * len(sources), + band_configs=self.band_configs, + ), + ) + self.decoder.insert(0, dec) + + self.separation_net = SeparationNet( + channels=dims[-1], + expand=expand, + num_layers=num_dplayer, + ) + + self.mask_layer = nn.Sequential( + nn.Conv2d(4 * len(self.sources), 64, kernel_size=3, padding="same"), + nn.GELU(), + nn.Conv2d( + 64, + 4 * len(self.sources), + kernel_size=1, + padding="same", + ), + nn.Tanh(), + ) + + def forward(self, x): + # B, C, L = x.shape + B = x.shape[0] + # In the initial padding, ensure that the number of frames after the STFT (the length of the T dimension) is even, + # so that the RFFT operation can be used in the separation network. + padding = self.hop_length - x.shape[-1] % self.hop_length + if (x.shape[-1] + padding) // self.hop_length % 2 == 0: + padding += self.hop_length + x = F.pad(x, (0, padding)) + + # STFT + L = x.shape[-1] + x = x.reshape(-1, L) + stft_opts = {**self.stft_config, "window": self.window.to(x.device)} + x = torch.stft(x, **stft_opts, return_complex=True) + x = torch.view_as_real(x) + x = x.permute(0, 3, 1, 2).reshape( + x.shape[0] // self.audio_channels, + x.shape[3] * self.audio_channels, + x.shape[1], + x.shape[2], + ) + + B, C, Fr, T = x.shape + + assert C == self.embed_dim, ( + f"Input channel dimension {C} after STFT/reshape doesn't match self.embed_dim {self.embed_dim}" + ) + mixture = x.repeat(1, len(self.sources), 1, 1) + + if Fr > self.max_f: + print( + f"Warning: Input frequency dim {Fr} > max_f {self.max_f}. Positional embedding will be truncated/repeated." + ) + repeats = math.ceil(Fr / self.max_f) + pos_f = self.pos_embed_f.repeat(1, 1, repeats, 1)[:, :, :Fr, :] + else: + pos_f = self.pos_embed_f[:, :, :Fr, :] + x = x + pos_f + + save_skip = deque() + save_lengths = deque() + save_original_lengths = deque() + # encoder + for sd_layer in self.encoder: + x, skip, lengths, original_lengths = sd_layer(x) + save_skip.append(skip) + save_lengths.append(lengths) + save_original_lengths.append(original_lengths) + + # separation + x = self.separation_net(x) + + # decoder + for fusion_layer, su_layer in self.decoder: + x = fusion_layer(x, save_skip.pop()) + x = su_layer(x, save_lengths.pop(), save_original_lengths.pop()) + + mask = self.mask_layer(x) + + # output + n = self.dims[0] + + mixture = mixture.view(B, n, -1, Fr, T) + mixture = mixture.reshape(-1, 2, Fr, T).permute(0, 2, 3, 1) + mixture = torch.view_as_complex(mixture.contiguous()) + + mask = mask.view(B, n, -1, Fr, T) + mask = mask.reshape(-1, 2, Fr, T).permute(0, 2, 3, 1) + mask = torch.view_as_complex(mask.contiguous()) + + x = mixture * mask + + x = torch.istft(x, **stft_opts) + x = x.reshape(B, len(self.sources), self.audio_channels, -1) + + x = x[:, :, :, :-padding] + + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet_tran.py b/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet_tran.py new file mode 100644 index 0000000000000000000000000000000000000000..695734b5c3d5f0cd9c05e6934414327d2bb60981 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet/scnet_tran.py @@ -0,0 +1,725 @@ +import math +from collections import deque + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import pack, rearrange, unpack +from models.bs_roformer.attend import Attend +from rotary_embedding_torch import RotaryEmbedding +from torch.nn import Module, ModuleList + +# helper functions + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +def pad_at_dim(t, pad, dim=-1, value=0.0): + dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) + zeros = (0, 0) * dims_from_right + return F.pad(t, (*zeros, *pad), value=value) + + +def l2norm(t): + return F.normalize(t, dim=-1, p=2) + + +# norm + + +class RMSNorm(Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +# attention + + +class FeedForward(Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + dim_inner = int(dim * mult) + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(dim_inner, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(Module): + def __init__( + self, dim, heads=8, dim_head=64, dropout=0.0, rotary_embed=None, flash=True + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = Attend(flash=flash, dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + self.to_gates = nn.Linear(dim, heads) + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=64, + heads=8, + attn_dropout=0.0, + ff_dropout=0.0, + ff_mult=4, + norm_output=True, + rotary_embed=None, + flash_attn=True, + linear_attn=False, + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + attn = Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + flash=flash_attn, + ) + + self.layers.append( + ModuleList( + [attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + + return self.norm(x) + + +class FeatureConversion(nn.Module): + """ + Integrates into the adjacent Dual-Path layer. + + Args: + channels (int): Number of input channels. + inverse (bool): If True, uses ifft; otherwise, uses rfft. + """ + + def __init__(self, channels, inverse): + super().__init__() + self.inverse = inverse + self.channels = channels + + def forward(self, x): + # B, C, F, T = x.shape + if self.inverse: + x = x.float() + x_r = x[:, : self.channels // 2, :, :] + x_i = x[:, self.channels // 2 :, :, :] + x = torch.complex(x_r, x_i) + x = torch.fft.irfft(x, dim=3, norm="ortho") + else: + x = x.float() + x = torch.fft.rfft(x, dim=3, norm="ortho") + x_real = x.real + x_imag = x.imag + x = torch.cat([x_real, x_imag], dim=1) + return x + + +class DualPathTran(nn.Module): + """ + Dual-Path Transformer in Separation Network. + + Args: + d_model (int): The number of expected features in the input (input_size). + expand (int): Expansion factor used to calculate the hidden_size of LSTM. + bidirectional (bool): If True, becomes a bidirectional LSTM. + """ + + def __init__(self, d_model, time_rotary_embed, freq_rotary_embed, tran_params): + super(DualPathTran, self).__init__() + + self.d_model = d_model + + transformer_kwargs = dict( + dim=d_model, + heads=tran_params["heads"], + dim_head=tran_params["dim_head"], + attn_dropout=tran_params["attn_dropout"], + ff_dropout=tran_params["ff_dropout"], + flash_attn=tran_params["flash_attn"], + ) + self.norm_layers = nn.ModuleList([nn.GroupNorm(1, d_model) for _ in range(2)]) + self.time_layer = Transformer( + depth=tran_params["depth"], + rotary_embed=time_rotary_embed, + **transformer_kwargs, + ) + self.freq_layer = Transformer( + depth=tran_params["depth"], + rotary_embed=freq_rotary_embed, + **transformer_kwargs, + ) + + def forward(self, x): + B, C, F, T = x.shape + + # Process dual-path rnn + original_x = x + # Frequency-path + x = self.norm_layers[0](x) + x = x.transpose(1, 3).contiguous().view(B * T, F, C) + # print('XXX', x.shape) + x = self.freq_layer(x) + x = x.view(B, T, F, C).transpose(1, 3) + x = x + original_x + + original_x = x + # Time-path + x = self.norm_layers[1](x) + x = x.transpose(1, 2).contiguous().view(B * F, C, T).transpose(1, 2) + # print('RRR', x.shape) + x = self.time_layer(x) + x = x.transpose(1, 2).contiguous().view(B, F, C, T).transpose(1, 2) + x = x + original_x + + return x + + +class SeparationNetTran(nn.Module): + """ + Implements a simplified Sparse Down-sample block in an encoder architecture. + + Args: + - channels (int): Number input channels. + - expand (int): Expansion factor used to calculate the hidden_size of LSTM. + - num_layers (int): Number of dual-path layers. + """ + + def __init__(self, channels, expand=1, num_layers=6, tran_params=None): + super(SeparationNetTran, self).__init__() + + self.num_layers = num_layers + + time_rotary_embed = RotaryEmbedding(dim=tran_params["rotary_embedding_dim"]) + freq_rotary_embed = RotaryEmbedding(dim=tran_params["rotary_embedding_dim"]) + + modules = [] + for i in range(num_layers): + m = DualPathTran( + channels * (2 if i % 2 == 1 else 1), + time_rotary_embed, + freq_rotary_embed, + tran_params, + ) + modules.append(m) + self.dp_modules = nn.ModuleList(modules) + + self.feature_conversion = nn.ModuleList( + [ + FeatureConversion(channels * 2, inverse=False if i % 2 == 0 else True) + for i in range(num_layers) + ] + ) + + def forward(self, x): + for i in range(self.num_layers): + x = self.dp_modules[i](x) + x = self.feature_conversion[i](x) + return x + + +class Swish(nn.Module): + def forward(self, x): + return x * x.sigmoid() + + +class ConvolutionModule(nn.Module): + """ + Convolution Module in SD block. + + Args: + channels (int): input/output channels. + depth (int): number of layers in the residual branch. Each layer has its own + compress (float): amount of channel compression. + kernel (int): kernel size for the convolutions. + """ + + def __init__(self, channels, depth=2, compress=4, kernel=3): + super().__init__() + assert kernel % 2 == 1 + self.depth = abs(depth) + hidden_size = int(channels / compress) + norm = lambda d: nn.GroupNorm(1, d) + self.layers = nn.ModuleList([]) + for _ in range(self.depth): + padding = kernel // 2 + mods = [ + norm(channels), + nn.Conv1d(channels, hidden_size * 2, kernel, padding=padding), + nn.GLU(1), + nn.Conv1d( + hidden_size, + hidden_size, + kernel, + padding=padding, + groups=hidden_size, + ), + norm(hidden_size), + Swish(), + nn.Conv1d(hidden_size, channels, 1), + ] + layer = nn.Sequential(*mods) + self.layers.append(layer) + + def forward(self, x): + for layer in self.layers: + x = x + layer(x) + return x + + +class FusionLayer(nn.Module): + """ + A FusionLayer within the decoder. + + Args: + - channels (int): Number of input channels. + - kernel_size (int, optional): Kernel size for the convolutional layer, defaults to 3. + - stride (int, optional): Stride for the convolutional layer, defaults to 1. + - padding (int, optional): Padding for the convolutional layer, defaults to 1. + """ + + def __init__(self, channels, kernel_size=3, stride=1, padding=1): + super(FusionLayer, self).__init__() + self.conv = nn.Conv2d( + channels * 2, channels * 2, kernel_size, stride=stride, padding=padding + ) + + def forward(self, x, skip=None): + if skip is not None: + x += skip + x = x.repeat(1, 2, 1, 1) + x = self.conv(x) + x = F.glu(x, dim=1) + return x + + +class SDlayer(nn.Module): + """ + Implements a Sparse Down-sample Layer for processing different frequency bands separately. + + Args: + - channels_in (int): Input channel count. + - channels_out (int): Output channel count. + - band_configs (dict): A dictionary containing configuration for each frequency band. + Keys are 'low', 'mid', 'high' for each band, and values are + dictionaries with keys 'SR', 'stride', and 'kernel' for proportion, + stride, and kernel size, respectively. + """ + + def __init__(self, channels_in, channels_out, band_configs): + super(SDlayer, self).__init__() + + # Initializing convolutional layers for each band + self.convs = nn.ModuleList() + self.strides = [] + self.kernels = [] + for config in band_configs.values(): + self.convs.append( + nn.Conv2d( + channels_in, + channels_out, + (config["kernel"], 1), + (config["stride"], 1), + (0, 0), + ) + ) + self.strides.append(config["stride"]) + self.kernels.append(config["kernel"]) + + # Saving rate proportions for determining splits + self.SR_low = band_configs["low"]["SR"] + self.SR_mid = band_configs["mid"]["SR"] + + def forward(self, x): + B, C, Fr, T = x.shape + # Define splitting points based on sampling rates + splits = [ + (0, math.ceil(Fr * self.SR_low)), + (math.ceil(Fr * self.SR_low), math.ceil(Fr * (self.SR_low + self.SR_mid))), + (math.ceil(Fr * (self.SR_low + self.SR_mid)), Fr), + ] + + # Processing each band with the corresponding convolution + outputs = [] + original_lengths = [] + for conv, stride, kernel, (start, end) in zip( + self.convs, self.strides, self.kernels, splits + ): + extracted = x[:, :, start:end, :] + original_lengths.append(end - start) + current_length = extracted.shape[2] + + # padding + if stride == 1: + total_padding = kernel - stride + else: + total_padding = (stride - current_length % stride) % stride + pad_left = total_padding // 2 + pad_right = total_padding - pad_left + + padded = F.pad(extracted, (0, 0, pad_left, pad_right)) + + output = conv(padded) + outputs.append(output) + + return outputs, original_lengths + + +class SUlayer(nn.Module): + """ + Implements a Sparse Up-sample Layer in decoder. + + Args: + - channels_in: The number of input channels. + - channels_out: The number of output channels. + - convtr_configs: Dictionary containing the configurations for transposed convolutions. + """ + + def __init__(self, channels_in, channels_out, band_configs): + super(SUlayer, self).__init__() + + # Initializing convolutional layers for each band + self.convtrs = nn.ModuleList( + [ + nn.ConvTranspose2d( + channels_in, + channels_out, + [config["kernel"], 1], + [config["stride"], 1], + ) + for _, config in band_configs.items() + ] + ) + + def forward(self, x, lengths, origin_lengths): + B, C, Fr, T = x.shape + # Define splitting points based on input lengths + splits = [ + (0, lengths[0]), + (lengths[0], lengths[0] + lengths[1]), + (lengths[0] + lengths[1], None), + ] + # Processing each band with the corresponding convolution + outputs = [] + for idx, (convtr, (start, end)) in enumerate(zip(self.convtrs, splits)): + out = convtr(x[:, :, start:end, :]) + # Calculate the distance to trim the output symmetrically to original length + current_Fr_length = out.shape[2] + dist = abs(origin_lengths[idx] - current_Fr_length) // 2 + + # Trim the output to the original length symmetrically + trimmed_out = out[:, :, dist : dist + origin_lengths[idx], :] + + outputs.append(trimmed_out) + + # Concatenate trimmed outputs along the frequency dimension to return the final tensor + x = torch.cat(outputs, dim=2) + + return x + + +class SDblock(nn.Module): + """ + Implements a simplified Sparse Down-sample block in encoder. + + Args: + - channels_in (int): Number of input channels. + - channels_out (int): Number of output channels. + - band_config (dict): Configuration for the SDlayer specifying band splits and convolutions. + - conv_config (dict): Configuration for convolution modules applied to each band. + - depths (list of int): List specifying the convolution depths for low, mid, and high frequency bands. + """ + + def __init__( + self, + channels_in, + channels_out, + band_configs={}, + conv_config={}, + depths=[3, 2, 1], + kernel_size=3, + ): + super(SDblock, self).__init__() + self.SDlayer = SDlayer(channels_in, channels_out, band_configs) + + # Dynamically create convolution modules for each band based on depths + self.conv_modules = nn.ModuleList( + [ConvolutionModule(channels_out, depth, **conv_config) for depth in depths] + ) + # Set the kernel_size to an odd number. + self.globalconv = nn.Conv2d( + channels_out, channels_out, kernel_size, 1, (kernel_size - 1) // 2 + ) + + def forward(self, x): + bands, original_lengths = self.SDlayer(x) + # B, C, f, T = band.shape + bands = [ + F.gelu( + conv(band.permute(0, 2, 1, 3).reshape(-1, band.shape[1], band.shape[3])) + .view(band.shape[0], band.shape[2], band.shape[1], band.shape[3]) + .permute(0, 2, 1, 3) + ) + for conv, band in zip(self.conv_modules, bands) + ] + lengths = [band.size(-2) for band in bands] + full_band = torch.cat(bands, dim=2) + skip = full_band + + output = self.globalconv(full_band) + + return output, skip, lengths, original_lengths + + +class SCNet_Tran(nn.Module): + """ + The implementation of SCNet: Sparse Compression Network for Music Source Separation. Paper: https://arxiv.org/abs/2401.13276.pdf + LSTM layers replaced with transformer layers + + Args: + - sources (List[str]): List of sources to be separated. + - audio_channels (int): Number of audio channels. + - nfft (int): Number of FFTs to determine the frequency dimension of the input. + - hop_size (int): Hop size for the STFT. + - win_size (int): Window size for STFT. + - normalized (bool): Whether to normalize the STFT. + - dims (List[int]): List of channel dimensions for each block. + - band_SR (List[float]): The proportion of each frequency band. + - band_stride (List[int]): The down-sampling ratio of each frequency band. + - band_kernel (List[int]): The kernel sizes for down-sampling convolution in each frequency band + - conv_depths (List[int]): List specifying the number of convolution modules in each SD block. + - compress (int): Compression factor for convolution module. + - conv_kernel (int): Kernel size for convolution layer in convolution module. + - num_dplayer (int): Number of dual-path layers. + - expand (int): Expansion factor in the dual-path RNN, default is 1. + + """ + + def __init__( + self, + sources=("drums", "bass", "other", "vocals"), + audio_channels=2, + # Main structure + dims=(4, 32, 64, 128), # dims = [4, 64, 128, 256] in SCNet-large + # STFT + nfft=4096, + hop_size=1024, + win_size=4096, + normalized=True, + # SD/SU layer + band_SR=(0.175, 0.392, 0.433), + band_stride=(1, 4, 16), + band_kernel=(3, 4, 16), + # Convolution Module + conv_depths=(3, 2, 1), + compress=4, + conv_kernel=3, + # Dual-path RNN + num_dplayer=6, + expand=1, + tran_rotary_embedding_dim=64, + tran_depth=1, + tran_heads=8, + tran_dim_head=64, + tran_attn_dropout=0.0, + tran_ff_dropout=0.0, + tran_flash_attn=False, + ): + super().__init__() + self.sources = sources + self.audio_channels = audio_channels + self.dims = dims + band_keys = ["low", "mid", "high"] + self.band_configs = { + band_keys[i]: { + "SR": band_SR[i], + "stride": band_stride[i], + "kernel": band_kernel[i], + } + for i in range(len(band_keys)) + } + self.hop_length = hop_size + self.conv_config = { + "compress": compress, + "kernel": conv_kernel, + } + self.tran_params = { + "rotary_embedding_dim": tran_rotary_embedding_dim, + "depth": tran_depth, + "heads": tran_heads, + "dim_head": tran_dim_head, + "attn_dropout": tran_attn_dropout, + "ff_dropout": tran_ff_dropout, + "flash_attn": tran_flash_attn, + } + + self.stft_config = { + "n_fft": nfft, + "hop_length": hop_size, + "win_length": win_size, + "center": True, + "normalized": normalized, + } + + self.first_conv = nn.Conv2d(dims[0], dims[0], 1, 1, 0, bias=False) + + self.encoder = nn.ModuleList() + self.decoder = nn.ModuleList() + + for index in range(len(dims) - 1): + enc = SDblock( + channels_in=dims[index], + channels_out=dims[index + 1], + band_configs=self.band_configs, + conv_config=self.conv_config, + depths=conv_depths, + ) + self.encoder.append(enc) + + dec = nn.Sequential( + FusionLayer(channels=dims[index + 1]), + SUlayer( + channels_in=dims[index + 1], + channels_out=dims[index] + if index != 0 + else dims[index] * len(sources), + band_configs=self.band_configs, + ), + ) + self.decoder.insert(0, dec) + + self.separation_net = SeparationNetTran( + channels=dims[-1], + expand=expand, + num_layers=num_dplayer, + tran_params=self.tran_params, + ) + + def forward(self, x): + # B, C, L = x.shape + B = x.shape[0] + # In the initial padding, ensure that the number of frames after the STFT (the length of the T dimension) is even, + # so that the RFFT operation can be used in the separation network. + padding = self.hop_length - x.shape[-1] % self.hop_length + if (x.shape[-1] + padding) // self.hop_length % 2 == 0: + padding += self.hop_length + x = F.pad(x, (0, padding)) + + # STFT + L = x.shape[-1] + x = x.reshape(-1, L) + x = torch.stft(x, **self.stft_config, return_complex=True) + x = torch.view_as_real(x) + x = x.permute(0, 3, 1, 2).reshape( + x.shape[0] // self.audio_channels, + x.shape[3] * self.audio_channels, + x.shape[1], + x.shape[2], + ) + + B, C, Fr, T = x.shape + + save_skip = deque() + save_lengths = deque() + save_original_lengths = deque() + # encoder + for sd_layer in self.encoder: + x, skip, lengths, original_lengths = sd_layer(x) + save_skip.append(skip) + save_lengths.append(lengths) + save_original_lengths.append(original_lengths) + + # separation + x = self.separation_net(x) + + # decoder + for fusion_layer, su_layer in self.decoder: + x = fusion_layer(x, save_skip.pop()) + x = su_layer(x, save_lengths.pop(), save_original_lengths.pop()) + + # output + n = self.dims[0] + x = x.view(B, n, -1, Fr, T) + + x = x.reshape(-1, 2, Fr, T).permute(0, 2, 3, 1) + x = torch.view_as_complex(x.contiguous()) + x = torch.istft(x, **self.stft_config) + x = x.reshape(B, len(self.sources), self.audio_channels, -1) + + x = x[:, :, :, :-padding] + + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet/separation.py b/src/third_party/MusicSourceSeparationTraining/models/scnet/separation.py new file mode 100644 index 0000000000000000000000000000000000000000..8965e2c8b14fa2c1fb6a2766e840c45128e1303d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet/separation.py @@ -0,0 +1,129 @@ +import torch +import torch.nn as nn +from torch.nn.modules.rnn import LSTM + + +class FeatureConversion(nn.Module): + """ + Integrates into the adjacent Dual-Path layer. + + Args: + channels (int): Number of input channels. + inverse (bool): If True, uses ifft; otherwise, uses rfft. + """ + + def __init__(self, channels, inverse): + super().__init__() + self.inverse = inverse + self.channels = channels + + def forward(self, x): + # B, C, F, T = x.shape + if self.inverse: + x = x.float() + x_r = x[:, : self.channels // 2, :, :] + x_i = x[:, self.channels // 2 :, :, :] + x = torch.complex(x_r, x_i) + x = torch.fft.irfft(x, dim=3, norm="ortho") + else: + x = x.float() + x = torch.fft.rfft(x, dim=3, norm="ortho") + x_real = x.real + x_imag = x.imag + x = torch.cat([x_real, x_imag], dim=1) + return x + + +class DualPathRNN(nn.Module): + """ + Dual-Path RNN in Separation Network. + + Args: + d_model (int): The number of expected features in the input (input_size). + expand (int): Expansion factor used to calculate the hidden_size of LSTM. + bidirectional (bool): If True, becomes a bidirectional LSTM. + """ + + def __init__(self, d_model, expand, bidirectional=True): + super(DualPathRNN, self).__init__() + + self.d_model = d_model + self.hidden_size = d_model * expand + self.bidirectional = bidirectional + # Initialize LSTM layers and normalization layers + self.lstm_layers = nn.ModuleList( + [self._init_lstm_layer(self.d_model, self.hidden_size) for _ in range(2)] + ) + self.linear_layers = nn.ModuleList( + [nn.Linear(self.hidden_size * 2, self.d_model) for _ in range(2)] + ) + self.norm_layers = nn.ModuleList([nn.GroupNorm(1, d_model) for _ in range(2)]) + + def _init_lstm_layer(self, d_model, hidden_size): + return LSTM( + d_model, + hidden_size, + num_layers=1, + bidirectional=self.bidirectional, + batch_first=True, + ) + + def forward(self, x): + B, C, F, T = x.shape + + # Process dual-path rnn + original_x = x + # Frequency-path + x = self.norm_layers[0](x) + x = x.transpose(1, 3).contiguous().view(B * T, F, C) + x, _ = self.lstm_layers[0](x) + x = self.linear_layers[0](x) + x = x.view(B, T, F, C).transpose(1, 3) + x = x + original_x + + original_x = x + # Time-path + x = self.norm_layers[1](x) + x = x.transpose(1, 2).contiguous().view(B * F, C, T).transpose(1, 2) + x, _ = self.lstm_layers[1](x) + x = self.linear_layers[1](x) + x = x.transpose(1, 2).contiguous().view(B, F, C, T).transpose(1, 2) + x = x + original_x + + return x + + +class SeparationNet(nn.Module): + """ + Implements a simplified Sparse Down-sample block in an encoder architecture. + + Args: + - channels (int): Number input channels. + - expand (int): Expansion factor used to calculate the hidden_size of LSTM. + - num_layers (int): Number of dual-path layers. + """ + + def __init__(self, channels, expand=1, num_layers=6): + super(SeparationNet, self).__init__() + + self.num_layers = num_layers + + self.dp_modules = nn.ModuleList( + [ + DualPathRNN(channels * (2 if i % 2 == 1 else 1), expand) + for i in range(num_layers) + ] + ) + + self.feature_conversion = nn.ModuleList( + [ + FeatureConversion(channels * 2, inverse=False if i % 2 == 0 else True) + for i in range(num_layers) + ] + ) + + def forward(self, x): + for i in range(self.num_layers): + x = self.dp_modules[i](x) + x = self.feature_conversion[i](x) + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..298d9939f5177c6b24cca743c83a351a84a6ffce --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/__init__.py @@ -0,0 +1 @@ +from models.scnet_unofficial.scnet import SCNet diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/__init__.py b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..69617bb15044d9bbfd0211fcdfa0fa605b01c048 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/__init__.py @@ -0,0 +1,3 @@ +from models.scnet_unofficial.modules.dualpath_rnn import DualPathRNN +from models.scnet_unofficial.modules.sd_encoder import SDBlock +from models.scnet_unofficial.modules.su_decoder import SUBlock diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/dualpath_rnn.py b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/dualpath_rnn.py new file mode 100644 index 0000000000000000000000000000000000000000..07ab55f7a61c53c0cd6e692b573444ca416db12f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/dualpath_rnn.py @@ -0,0 +1,236 @@ +import torch +import torch.nn as nn +import torch.nn.functional as Func + + +class RMSNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return Func.normalize(x, dim=-1) * self.scale * self.gamma + + +class MambaModule(nn.Module): + def __init__(self, d_model, d_state, d_conv, d_expand): + super().__init__() + self.norm = RMSNorm(dim=d_model) + self.mamba = Mamba( + d_model=d_model, d_state=d_state, d_conv=d_conv, d_expand=d_expand + ) + + def forward(self, x): + x = x + self.mamba(self.norm(x)) + return x + + +class RNNModule(nn.Module): + """ + RNNModule class implements a recurrent neural network module with LSTM cells. + + Args: + - input_dim (int): Dimensionality of the input features. + - hidden_dim (int): Dimensionality of the hidden state of the LSTM. + - bidirectional (bool, optional): If True, uses bidirectional LSTM. Defaults to True. + + Shapes: + - Input: (B, T, D) where + B is batch size, + T is sequence length, + D is input dimensionality. + - Output: (B, T, D) where + B is batch size, + T is sequence length, + D is input dimensionality. + """ + + def __init__(self, input_dim: int, hidden_dim: int, bidirectional: bool = True): + """ + Initializes RNNModule with input dimension, hidden dimension, and bidirectional flag. + """ + super().__init__() + self.groupnorm = nn.GroupNorm(num_groups=1, num_channels=input_dim) + self.rnn = nn.LSTM( + input_dim, hidden_dim, batch_first=True, bidirectional=bidirectional + ) + self.fc = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, input_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the RNNModule. + + Args: + - x (torch.Tensor): Input tensor of shape (B, T, D). + + Returns: + - torch.Tensor: Output tensor of shape (B, T, D). + """ + x = x.transpose(1, 2) + x = self.groupnorm(x) + x = x.transpose(1, 2) + + x, (hidden, _) = self.rnn(x) + x = self.fc(x) + return x + + +class RFFTModule(nn.Module): + """ + RFFTModule class implements a module for performing real-valued Fast Fourier Transform (FFT) + or its inverse on input tensors. + + Args: + - inverse (bool, optional): If False, performs forward FFT. If True, performs inverse FFT. Defaults to False. + + Shapes: + - Input: (B, F, T, D) where + B is batch size, + F is the number of features, + T is sequence length, + D is input dimensionality. + - Output: (B, F, T // 2 + 1, D * 2) if performing forward FFT. + (B, F, T, D // 2, 2) if performing inverse FFT. + """ + + def __init__(self, inverse: bool = False): + """ + Initializes RFFTModule with inverse flag. + """ + super().__init__() + self.inverse = inverse + + def forward(self, x: torch.Tensor, time_dim: int) -> torch.Tensor: + """ + Performs forward or inverse FFT on the input tensor x. + + Args: + - x (torch.Tensor): Input tensor of shape (B, F, T, D). + - time_dim (int): Input size of time dimension. + + Returns: + - torch.Tensor: Output tensor after FFT or its inverse operation. + """ + dtype = x.dtype + B, F, T, D = x.shape + + # RuntimeError: cuFFT only supports dimensions whose sizes are powers of two when computing in half precision + x = x.float() + + if not self.inverse: + x = torch.fft.rfft(x, dim=2) + x = torch.view_as_real(x) + x = x.reshape(B, F, T // 2 + 1, D * 2) + else: + x = x.reshape(B, F, T, D // 2, 2) + x = torch.view_as_complex(x) + x = torch.fft.irfft(x, n=time_dim, dim=2) + + x = x.to(dtype) + return x + + def extra_repr(self) -> str: + """ + Returns extra representation string with module's configuration. + """ + return f"inverse={self.inverse}" + + +class DualPathRNN(nn.Module): + """ + DualPathRNN class implements a neural network with alternating layers of RNNModule and RFFTModule. + + Args: + - n_layers (int): Number of layers in the network. + - input_dim (int): Dimensionality of the input features. + - hidden_dim (int): Dimensionality of the hidden state of the RNNModule. + + Shapes: + - Input: (B, F, T, D) where + B is batch size, + F is the number of features (frequency dimension), + T is sequence length (time dimension), + D is input dimensionality (channel dimension). + - Output: (B, F, T, D) where + B is batch size, + F is the number of features (frequency dimension), + T is sequence length (time dimension), + D is input dimensionality (channel dimension). + """ + + def __init__( + self, + n_layers: int, + input_dim: int, + hidden_dim: int, + use_mamba: bool = False, + d_state: int = 16, + d_conv: int = 4, + d_expand: int = 2, + ): + """ + Initializes DualPathRNN with the specified number of layers, input dimension, and hidden dimension. + """ + super().__init__() + + if use_mamba: + net = MambaModule + dkwargs = { + "d_model": input_dim, + "d_state": d_state, + "d_conv": d_conv, + "d_expand": d_expand, + } + ukwargs = { + "d_model": input_dim * 2, + "d_state": d_state, + "d_conv": d_conv, + "d_expand": d_expand * 2, + } + else: + net = RNNModule + dkwargs = {"input_dim": input_dim, "hidden_dim": hidden_dim} + ukwargs = {"input_dim": input_dim * 2, "hidden_dim": hidden_dim * 2} + + self.layers = nn.ModuleList() + for i in range(1, n_layers + 1): + kwargs = dkwargs if i % 2 == 1 else ukwargs + layer = nn.ModuleList( + [ + net(**kwargs), + net(**kwargs), + RFFTModule(inverse=(i % 2 == 0)), + ] + ) + self.layers.append(layer) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the DualPathRNN. + + Args: + - x (torch.Tensor): Input tensor of shape (B, F, T, D). + + Returns: + - torch.Tensor: Output tensor of shape (B, F, T, D). + """ + + time_dim = x.shape[2] + + for time_layer, freq_layer, rfft_layer in self.layers: + B, F, T, D = x.shape + + x = x.reshape((B * F), T, D) + x = time_layer(x) + x = x.reshape(B, F, T, D) + x = x.permute(0, 2, 1, 3) + + x = x.reshape((B * T), F, D) + x = freq_layer(x) + x = x.reshape(B, T, F, D) + x = x.permute(0, 2, 1, 3) + + x = rfft_layer(x, time_dim) + + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/sd_encoder.py b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/sd_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..eab3e931b1d0889e5c021d41cf64269061acdd59 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/sd_encoder.py @@ -0,0 +1,284 @@ +from typing import List, Tuple + +import torch +import torch.nn as nn +from models.scnet_unofficial.utils import create_intervals + + +class Downsample(nn.Module): + """ + Downsample class implements a module for downsampling input tensors using 2D convolution. + + Args: + - input_dim (int): Dimensionality of the input channels. + - output_dim (int): Dimensionality of the output channels. + - stride (int): Stride value for the convolution operation. + + Shapes: + - Input: (B, C_in, F, T) where + B is batch size, + C_in is the number of input channels, + F is the frequency dimension, + T is the time dimension. + - Output: (B, C_out, F // stride, T) where + B is batch size, + C_out is the number of output channels, + F // stride is the downsampled frequency dimension. + + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + stride: int, + ): + """ + Initializes Downsample with input dimension, output dimension, and stride. + """ + super().__init__() + self.conv = nn.Conv2d(input_dim, output_dim, 1, (stride, 1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the Downsample module. + + Args: + - x (torch.Tensor): Input tensor of shape (B, C_in, F, T). + + Returns: + - torch.Tensor: Downsampled tensor of shape (B, C_out, F // stride, T). + """ + return self.conv(x) + + +class ConvolutionModule(nn.Module): + """ + ConvolutionModule class implements a module with a sequence of convolutional layers similar to Conformer. + + Args: + - input_dim (int): Dimensionality of the input features. + - hidden_dim (int): Dimensionality of the hidden features. + - kernel_sizes (List[int]): List of kernel sizes for the convolutional layers. + - bias (bool, optional): If True, adds a learnable bias to the output. Default is False. + + Shapes: + - Input: (B, T, D) where + B is batch size, + T is sequence length, + D is input dimensionality. + - Output: (B, T, D) where + B is batch size, + T is sequence length, + D is input dimensionality. + """ + + def __init__( + self, + input_dim: int, + hidden_dim: int, + kernel_sizes: List[int], + bias: bool = False, + ) -> None: + """ + Initializes ConvolutionModule with input dimension, hidden dimension, kernel sizes, and bias. + """ + super().__init__() + self.sequential = nn.Sequential( + nn.GroupNorm(num_groups=1, num_channels=input_dim), + nn.Conv1d( + input_dim, + 2 * hidden_dim, + kernel_sizes[0], + stride=1, + padding=(kernel_sizes[0] - 1) // 2, + bias=bias, + ), + nn.GLU(dim=1), + nn.Conv1d( + hidden_dim, + hidden_dim, + kernel_sizes[1], + stride=1, + padding=(kernel_sizes[1] - 1) // 2, + groups=hidden_dim, + bias=bias, + ), + nn.GroupNorm(num_groups=1, num_channels=hidden_dim), + nn.SiLU(), + nn.Conv1d( + hidden_dim, + input_dim, + kernel_sizes[2], + stride=1, + padding=(kernel_sizes[2] - 1) // 2, + bias=bias, + ), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the ConvolutionModule. + + Args: + - x (torch.Tensor): Input tensor of shape (B, T, D). + + Returns: + - torch.Tensor: Output tensor of shape (B, T, D). + """ + x = x.transpose(1, 2) + x = x + self.sequential(x) + x = x.transpose(1, 2) + return x + + +class SDLayer(nn.Module): + """ + SDLayer class implements a subband decomposition layer with downsampling and convolutional modules. + + Args: + - subband_interval (Tuple[float, float]): Tuple representing the frequency interval for subband decomposition. + - input_dim (int): Dimensionality of the input channels. + - output_dim (int): Dimensionality of the output channels after downsampling. + - downsample_stride (int): Stride value for the downsampling operation. + - n_conv_modules (int): Number of convolutional modules. + - kernel_sizes (List[int]): List of kernel sizes for the convolutional layers. + - bias (bool, optional): If True, adds a learnable bias to the convolutional layers. Default is True. + + Shapes: + - Input: (B, Fi, T, Ci) where + B is batch size, + Fi is the number of input subbands, + T is sequence length, and + Ci is the number of input channels. + - Output: (B, Fi+1, T, Ci+1) where + B is batch size, + Fi+1 is the number of output subbands, + T is sequence length, + Ci+1 is the number of output channels. + """ + + def __init__( + self, + subband_interval: Tuple[float, float], + input_dim: int, + output_dim: int, + downsample_stride: int, + n_conv_modules: int, + kernel_sizes: List[int], + bias: bool = True, + ): + """ + Initializes SDLayer with subband interval, input dimension, + output dimension, downsample stride, number of convolutional modules, kernel sizes, and bias. + """ + super().__init__() + self.subband_interval = subband_interval + self.downsample = Downsample(input_dim, output_dim, downsample_stride) + self.activation = nn.GELU() + conv_modules = [ + ConvolutionModule( + input_dim=output_dim, + hidden_dim=output_dim // 4, + kernel_sizes=kernel_sizes, + bias=bias, + ) + for _ in range(n_conv_modules) + ] + self.conv_modules = nn.Sequential(*conv_modules) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the SDLayer. + + Args: + - x (torch.Tensor): Input tensor of shape (B, Fi, T, Ci). + + Returns: + - torch.Tensor: Output tensor of shape (B, Fi+1, T, Ci+1). + """ + B, F, T, C = x.shape + x = x[:, int(self.subband_interval[0] * F) : int(self.subband_interval[1] * F)] + x = x.permute(0, 3, 1, 2) + x = self.downsample(x) + x = self.activation(x) + x = x.permute(0, 2, 3, 1) + + B, F, T, C = x.shape + x = x.reshape((B * F), T, C) + x = self.conv_modules(x) + x = x.reshape(B, F, T, C) + + return x + + +class SDBlock(nn.Module): + """ + SDBlock class implements a block with subband decomposition layers and global convolution. + + Args: + - input_dim (int): Dimensionality of the input channels. + - output_dim (int): Dimensionality of the output channels. + - bandsplit_ratios (List[float]): List of ratios for splitting the frequency bands. + - downsample_strides (List[int]): List of stride values for downsampling in each subband layer. + - n_conv_modules (List[int]): List specifying the number of convolutional modules in each subband layer. + - kernel_sizes (List[int], optional): List of kernel sizes for the convolutional layers. Default is None. + + Shapes: + - Input: (B, Fi, T, Ci) where + B is batch size, + Fi is the number of input subbands, + T is sequence length, + Ci is the number of input channels. + - Output: (B, Fi+1, T, Ci+1) where + B is batch size, + Fi+1 is the number of output subbands, + T is sequence length, + Ci+1 is the number of output channels. + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + bandsplit_ratios: List[float], + downsample_strides: List[int], + n_conv_modules: List[int], + kernel_sizes: List[int] = None, + ): + """ + Initializes SDBlock with input dimension, output dimension, band split ratios, downsample strides, number of convolutional modules, and kernel sizes. + """ + super().__init__() + if kernel_sizes is None: + kernel_sizes = [3, 3, 1] + assert sum(bandsplit_ratios) == 1, "The split ratios must sum up to 1." + subband_intervals = create_intervals(bandsplit_ratios) + self.sd_layers = nn.ModuleList( + SDLayer( + input_dim=input_dim, + output_dim=output_dim, + subband_interval=sbi, + downsample_stride=dss, + n_conv_modules=ncm, + kernel_sizes=kernel_sizes, + ) + for sbi, dss, ncm in zip( + subband_intervals, downsample_strides, n_conv_modules + ) + ) + self.global_conv2d = nn.Conv2d(output_dim, output_dim, 1, 1) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Performs forward pass through the SDBlock. + + Args: + - x (torch.Tensor): Input tensor of shape (B, Fi, T, Ci). + + Returns: + - Tuple[torch.Tensor, torch.Tensor]: Output tensor and skip connection tensor. + """ + x_skip = torch.concat([layer(x) for layer in self.sd_layers], dim=1) + x = self.global_conv2d(x_skip.permute(0, 3, 1, 2)).permute(0, 2, 3, 1) + return x, x_skip diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/su_decoder.py b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/su_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..4f92221b4051c08f2038ecaf314e4e5bc54b8adc --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/modules/su_decoder.py @@ -0,0 +1,240 @@ +from typing import List, Tuple + +import torch +import torch.nn as nn +from models.scnet_unofficial.utils import get_convtranspose_output_padding + + +class FusionLayer(nn.Module): + """ + FusionLayer class implements a module for fusing two input tensors using convolutional operations. + + Args: + - input_dim (int): Dimensionality of the input channels. + - kernel_size (int, optional): Kernel size for the convolutional layer. Default is 3. + - stride (int, optional): Stride value for the convolutional layer. Default is 1. + - padding (int, optional): Padding value for the convolutional layer. Default is 1. + + Shapes: + - Input: (B, F, T, C) and (B, F, T, C) where + B is batch size, + F is the number of features, + T is sequence length, + C is input dimensionality. + - Output: (B, F, T, C) where + B is batch size, + F is the number of features, + T is sequence length, + C is input dimensionality. + """ + + def __init__( + self, input_dim: int, kernel_size: int = 3, stride: int = 1, padding: int = 1 + ): + """ + Initializes FusionLayer with input dimension, kernel size, stride, and padding. + """ + super().__init__() + self.conv = nn.Conv2d( + input_dim * 2, + input_dim * 2, + kernel_size=(kernel_size, 1), + stride=(stride, 1), + padding=(padding, 0), + ) + self.activation = nn.GLU() + + def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the FusionLayer. + + Args: + - x1 (torch.Tensor): First input tensor of shape (B, F, T, C). + - x2 (torch.Tensor): Second input tensor of shape (B, F, T, C). + + Returns: + - torch.Tensor: Output tensor of shape (B, F, T, C). + """ + x = x1 + x2 + x = x.repeat(1, 1, 1, 2) + x = self.conv(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1) + x = self.activation(x) + return x + + +class Upsample(nn.Module): + """ + Upsample class implements a module for upsampling input tensors using transposed 2D convolution. + + Args: + - input_dim (int): Dimensionality of the input channels. + - output_dim (int): Dimensionality of the output channels. + - stride (int): Stride value for the transposed convolution operation. + - output_padding (int): Output padding value for the transposed convolution operation. + + Shapes: + - Input: (B, C_in, F, T) where + B is batch size, + C_in is the number of input channels, + F is the frequency dimension, + T is the time dimension. + - Output: (B, C_out, F * stride + output_padding, T) where + B is batch size, + C_out is the number of output channels, + F * stride + output_padding is the upsampled frequency dimension. + """ + + def __init__( + self, input_dim: int, output_dim: int, stride: int, output_padding: int + ): + """ + Initializes Upsample with input dimension, output dimension, stride, and output padding. + """ + super().__init__() + self.conv = nn.ConvTranspose2d( + input_dim, output_dim, 1, (stride, 1), output_padding=(output_padding, 0) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the Upsample module. + + Args: + - x (torch.Tensor): Input tensor of shape (B, C_in, F, T). + + Returns: + - torch.Tensor: Output tensor of shape (B, C_out, F * stride + output_padding, T). + """ + return self.conv(x) + + +class SULayer(nn.Module): + """ + SULayer class implements a subband upsampling layer using transposed convolution. + + Args: + - input_dim (int): Dimensionality of the input channels. + - output_dim (int): Dimensionality of the output channels. + - upsample_stride (int): Stride value for the upsampling operation. + - subband_shape (int): Shape of the subband. + - sd_interval (Tuple[int, int]): Start and end indices of the subband interval. + + Shapes: + - Input: (B, F, T, C) where + B is batch size, + F is the number of features, + T is sequence length, + C is input dimensionality. + - Output: (B, F, T, C) where + B is batch size, + F is the number of features, + T is sequence length, + C is input dimensionality. + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + upsample_stride: int, + subband_shape: int, + sd_interval: Tuple[int, int], + ): + """ + Initializes SULayer with input dimension, output dimension, upsample stride, subband shape, and subband interval. + """ + super().__init__() + sd_shape = sd_interval[1] - sd_interval[0] + upsample_output_padding = get_convtranspose_output_padding( + input_shape=sd_shape, output_shape=subband_shape, stride=upsample_stride + ) + self.upsample = Upsample( + input_dim=input_dim, + output_dim=output_dim, + stride=upsample_stride, + output_padding=upsample_output_padding, + ) + self.sd_interval = sd_interval + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the SULayer. + + Args: + - x (torch.Tensor): Input tensor of shape (B, F, T, C). + + Returns: + - torch.Tensor: Output tensor of shape (B, F, T, C). + """ + x = x[:, self.sd_interval[0] : self.sd_interval[1]] + x = x.permute(0, 3, 1, 2) + x = self.upsample(x) + x = x.permute(0, 2, 3, 1) + return x + + +class SUBlock(nn.Module): + """ + SUBlock class implements a block with fusion layer and subband upsampling layers. + + Args: + - input_dim (int): Dimensionality of the input channels. + - output_dim (int): Dimensionality of the output channels. + - upsample_strides (List[int]): List of stride values for the upsampling operations. + - subband_shapes (List[int]): List of shapes for the subbands. + - sd_intervals (List[Tuple[int, int]]): List of intervals for subband decomposition. + + Shapes: + - Input: (B, Fi-1, T, Ci-1) and (B, Fi-1, T, Ci-1) where + B is batch size, + Fi-1 is the number of input subbands, + T is sequence length, + Ci-1 is the number of input channels. + - Output: (B, Fi, T, Ci) where + B is batch size, + Fi is the number of output subbands, + T is sequence length, + Ci is the number of output channels. + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + upsample_strides: List[int], + subband_shapes: List[int], + sd_intervals: List[Tuple[int, int]], + ): + """ + Initializes SUBlock with input dimension, output dimension, + upsample strides, subband shapes, and subband intervals. + """ + super().__init__() + self.fusion_layer = FusionLayer(input_dim=input_dim) + self.su_layers = nn.ModuleList( + SULayer( + input_dim=input_dim, + output_dim=output_dim, + upsample_stride=uss, + subband_shape=sbs, + sd_interval=sdi, + ) + for i, (uss, sbs, sdi) in enumerate( + zip(upsample_strides, subband_shapes, sd_intervals) + ) + ) + + def forward(self, x: torch.Tensor, x_skip: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the SUBlock. + + Args: + - x (torch.Tensor): Input tensor of shape (B, Fi-1, T, Ci-1). + - x_skip (torch.Tensor): Input skip connection tensor of shape (B, Fi-1, T, Ci-1). + + Returns: + - torch.Tensor: Output tensor of shape (B, Fi, T, Ci). + """ + x = self.fusion_layer(x, x_skip) + x = torch.concat([layer(x) for layer in self.su_layers], dim=1) + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/scnet.py b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/scnet.py new file mode 100644 index 0000000000000000000000000000000000000000..ca140912799e7a6277cfbf0cd34a2f8f712fce19 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/scnet.py @@ -0,0 +1,242 @@ +""" +SCNet - great paper, great implementation +https://arxiv.org/pdf/2401.13276.pdf +https://github.com/amanteur/SCNet-PyTorch +""" + +from functools import partial +from typing import List + +import torch +import torch.nn as nn +import torch.nn.functional as F +from beartype import beartype +from beartype.typing import Callable, List, Optional, Tuple +from einops import pack, rearrange, unpack +from models.scnet_unofficial.modules import DualPathRNN, SDBlock, SUBlock +from models.scnet_unofficial.utils import compute_gcr, compute_sd_layer_shapes + + +def exists(val): + return val is not None + + +def default(v, d): + return v if exists(v) else d + + +def pack_one(t, pattern): + return pack([t], pattern) + + +def unpack_one(t, ps, pattern): + return unpack(t, ps, pattern)[0] + + +class RMSNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return F.normalize(x, dim=-1) * self.scale * self.gamma + + +class BandSplit(nn.Module): + @beartype + def __init__(self, dim, dim_inputs: Tuple[int, ...]): + super().__init__() + self.dim_inputs = dim_inputs + self.to_features = ModuleList([]) + + for dim_in in dim_inputs: + net = nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) + + self.to_features.append(net) + + def forward(self, x): + x = x.split(self.dim_inputs, dim=-1) + + outs = [] + for split_input, to_feature in zip(x, self.to_features): + split_output = to_feature(split_input) + outs.append(split_output) + + return torch.stack(outs, dim=-2) + + +class SCNet(nn.Module): + """ + SCNet class implements a source separation network, + which explicitly split the spectrogram of the mixture into several subbands + and introduce a sparsity-based encoder to model different frequency bands. + + Paper: "SCNET: SPARSE COMPRESSION NETWORK FOR MUSIC SOURCE SEPARATION" + Authors: Weinan Tong, Jiaxu Zhu et al. + Link: https://arxiv.org/abs/2401.13276.pdf + + Args: + - n_fft (int): Number of FFTs to determine the frequency dimension of the input. + - dims (List[int]): List of channel dimensions for each block. + - bandsplit_ratios (List[float]): List of ratios for splitting the frequency bands. + - downsample_strides (List[int]): List of stride values for downsampling in each block. + - n_conv_modules (List[int]): List specifying the number of convolutional modules in each block. + - n_rnn_layers (int): Number of recurrent layers in the dual path RNN. + - rnn_hidden_dim (int): Dimensionality of the hidden state in the dual path RNN. + - n_sources (int, optional): Number of sources to be separated. Default is 4. + + Shapes: + - Input: (B, C, T) where + B is batch size, + C is channel dim (mono / stereo), + T is time dim + - Output: (B, N, C, T) where + B is batch size, + N is the number of sources. + C is channel dim (mono / stereo), + T is sequence length, + """ + + @beartype + def __init__( + self, + n_fft: int, + dims: List[int], + bandsplit_ratios: List[float], + downsample_strides: List[int], + n_conv_modules: List[int], + n_rnn_layers: int, + rnn_hidden_dim: int, + n_sources: int = 4, + hop_length: int = 1024, + win_length: int = 4096, + stft_window_fn: Optional[Callable] = None, + stft_normalized: bool = False, + **kwargs, + ): + """ + Initializes SCNet with input parameters. + """ + super().__init__() + self.assert_input_data( + bandsplit_ratios, + downsample_strides, + n_conv_modules, + ) + + n_blocks = len(dims) - 1 + n_freq_bins = n_fft // 2 + 1 + subband_shapes, sd_intervals = compute_sd_layer_shapes( + input_shape=n_freq_bins, + bandsplit_ratios=bandsplit_ratios, + downsample_strides=downsample_strides, + n_layers=n_blocks, + ) + self.sd_blocks = nn.ModuleList( + SDBlock( + input_dim=dims[i], + output_dim=dims[i + 1], + bandsplit_ratios=bandsplit_ratios, + downsample_strides=downsample_strides, + n_conv_modules=n_conv_modules, + ) + for i in range(n_blocks) + ) + self.dualpath_blocks = DualPathRNN( + n_layers=n_rnn_layers, + input_dim=dims[-1], + hidden_dim=rnn_hidden_dim, + **kwargs, + ) + self.su_blocks = nn.ModuleList( + SUBlock( + input_dim=dims[i + 1], + output_dim=dims[i] if i != 0 else dims[i] * n_sources, + subband_shapes=subband_shapes[i], + sd_intervals=sd_intervals[i], + upsample_strides=downsample_strides, + ) + for i in reversed(range(n_blocks)) + ) + self.gcr = compute_gcr(subband_shapes) + + self.stft_kwargs = dict( + n_fft=n_fft, + hop_length=hop_length, + win_length=win_length, + normalized=stft_normalized, + ) + + self.stft_window_fn = partial( + default(stft_window_fn, torch.hann_window), win_length + ) + self.n_sources = n_sources + self.hop_length = hop_length + + @staticmethod + def assert_input_data(*args): + """ + Asserts that the shapes of input features are equal. + """ + for arg1 in args: + for arg2 in args: + if len(arg1) != len(arg2): + raise ValueError( + f"Shapes of input features {arg1} and {arg2} are not equal." + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Performs forward pass through the SCNet. + + Args: + - x (torch.Tensor): Input tensor of shape (B, C, T). + + Returns: + - torch.Tensor: Output tensor of shape (B, N, C, T). + """ + + device = x.device + stft_window = self.stft_window_fn(device=device) + + if x.ndim == 2: + x = rearrange(x, "b t -> b 1 t") + + c = x.shape[1] + + stft_pad = self.hop_length - x.shape[-1] % self.hop_length + x = F.pad(x, (0, stft_pad)) + + # stft + x, ps = pack_one(x, "* t") + x = torch.stft(x, **self.stft_kwargs, window=stft_window, return_complex=True) + x = torch.view_as_real(x) + x = unpack_one(x, ps, "* c f t") + x = rearrange(x, "b c f t r -> b f t (c r)") + + # encoder part + x_skips = [] + for sd_block in self.sd_blocks: + x, x_skip = sd_block(x) + x_skips.append(x_skip) + + # separation part + x = self.dualpath_blocks(x) + + # decoder part + for su_block, x_skip in zip(self.su_blocks, reversed(x_skips)): + x = su_block(x, x_skip) + + # istft + x = rearrange(x, "b f t (c r n) -> b n c f t r", c=c, n=self.n_sources, r=2) + x = x.contiguous() + + x = torch.view_as_complex(x) + x = rearrange(x, "b n c f t -> (b n c) f t") + x = torch.istft(x, **self.stft_kwargs, window=stft_window, return_complex=False) + x = rearrange(x, "(b n c) t -> b n c t", c=c, n=self.n_sources) + + x = x[..., :-stft_pad] + + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/utils.py b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d236d499322a5db4ae4a813e75901dcfe28f7993 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/scnet_unofficial/utils.py @@ -0,0 +1,135 @@ +""" +SCNet - great paper, great implementation +https://arxiv.org/pdf/2401.13276.pdf +https://github.com/amanteur/SCNet-PyTorch +""" + +from typing import List, Tuple, Union + +import torch + + +def create_intervals( + splits: List[Union[float, int]], +) -> List[Union[Tuple[float, float], Tuple[int, int]]]: + """ + Create intervals based on splits provided. + + Args: + - splits (List[Union[float, int]]): List of floats or integers representing splits. + + Returns: + - List[Union[Tuple[float, float], Tuple[int, int]]]: List of tuples representing intervals. + """ + start = 0 + return [(start, start := start + split) for split in splits] + + +def get_conv_output_shape( + input_shape: int, + kernel_size: int = 1, + padding: int = 0, + dilation: int = 1, + stride: int = 1, +) -> int: + """ + Compute the output shape of a convolutional layer. + + Args: + - input_shape (int): Input shape. + - kernel_size (int, optional): Kernel size of the convolution. Default is 1. + - padding (int, optional): Padding size. Default is 0. + - dilation (int, optional): Dilation factor. Default is 1. + - stride (int, optional): Stride value. Default is 1. + + Returns: + - int: Output shape. + """ + return int( + (input_shape + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1 + ) + + +def get_convtranspose_output_padding( + input_shape: int, + output_shape: int, + kernel_size: int = 1, + padding: int = 0, + dilation: int = 1, + stride: int = 1, +) -> int: + """ + Compute the output padding for a convolution transpose operation. + + Args: + - input_shape (int): Input shape. + - output_shape (int): Desired output shape. + - kernel_size (int, optional): Kernel size of the convolution. Default is 1. + - padding (int, optional): Padding size. Default is 0. + - dilation (int, optional): Dilation factor. Default is 1. + - stride (int, optional): Stride value. Default is 1. + + Returns: + - int: Output padding. + """ + return ( + output_shape + - (input_shape - 1) * stride + + 2 * padding + - dilation * (kernel_size - 1) + - 1 + ) + + +def compute_sd_layer_shapes( + input_shape: int, + bandsplit_ratios: List[float], + downsample_strides: List[int], + n_layers: int, +) -> Tuple[List[List[int]], List[List[Tuple[int, int]]]]: + """ + Compute the shapes for the subband layers. + + Args: + - input_shape (int): Input shape. + - bandsplit_ratios (List[float]): Ratios for splitting the frequency bands. + - downsample_strides (List[int]): Strides for downsampling in each layer. + - n_layers (int): Number of layers. + + Returns: + - Tuple[List[List[int]], List[List[Tuple[int, int]]]]: Tuple containing subband shapes and convolution shapes. + """ + bandsplit_shapes_list = [] + conv2d_shapes_list = [] + for _ in range(n_layers): + bandsplit_intervals = create_intervals(bandsplit_ratios) + bandsplit_shapes = [ + int(right * input_shape) - int(left * input_shape) + for left, right in bandsplit_intervals + ] + conv2d_shapes = [ + get_conv_output_shape(bs, stride=ds) + for bs, ds in zip(bandsplit_shapes, downsample_strides) + ] + input_shape = sum(conv2d_shapes) + bandsplit_shapes_list.append(bandsplit_shapes) + conv2d_shapes_list.append(create_intervals(conv2d_shapes)) + + return bandsplit_shapes_list, conv2d_shapes_list + + +def compute_gcr(subband_shapes: List[List[int]]) -> float: + """ + Compute the global compression ratio. + + Args: + - subband_shapes (List[List[int]]): List of subband shapes. + + Returns: + - float: Global compression ratio. + """ + t = torch.Tensor(subband_shapes) + gcr = torch.stack( + [(1 - t[i + 1] / t[i]).mean() for i in range(0, len(t) - 1)] + ).mean() + return float(gcr) diff --git a/src/third_party/MusicSourceSeparationTraining/models/segm_models.py b/src/third_party/MusicSourceSeparationTraining/models/segm_models.py new file mode 100644 index 0000000000000000000000000000000000000000..a531a26ec5abcf928f50a4cf3b9e889c022af7fc --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/segm_models.py @@ -0,0 +1,253 @@ +import segmentation_models_pytorch as smp +import torch +import torch.nn as nn +from utils.model_utils import prefer_target_instrument + + +class STFT: + def __init__(self, config): + self.n_fft = config.n_fft + self.hop_length = config.hop_length + self.window = torch.hann_window(window_length=self.n_fft, periodic=True) + self.dim_f = config.dim_f + + def __call__(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-2] + c, t = x.shape[-2:] + x = x.reshape([-1, t]) + x = torch.stft( + x, + n_fft=self.n_fft, + hop_length=self.hop_length, + window=window, + center=True, + return_complex=True, + ) + x = torch.view_as_real(x) + x = x.permute([0, 3, 1, 2]) + x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape( + [*batch_dims, c * 2, -1, x.shape[-1]] + ) + return x[..., : self.dim_f, :] + + def inverse(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-3] + c, f, t = x.shape[-3:] + n = self.n_fft // 2 + 1 + f_pad = torch.zeros([*batch_dims, c, n - f, t]).to(x.device) + x = torch.cat([x, f_pad], -2) + x = x.reshape([*batch_dims, c // 2, 2, n, t]).reshape([-1, 2, n, t]) + x = x.permute([0, 2, 3, 1]) + x = x[..., 0] + x[..., 1] * 1.0j + x = torch.istft( + x, n_fft=self.n_fft, hop_length=self.hop_length, window=window, center=True + ) + x = x.reshape([*batch_dims, 2, -1]) + return x + + +def get_act(act_type): + if act_type == "gelu": + return nn.GELU() + elif act_type == "relu": + return nn.ReLU() + elif act_type[:3] == "elu": + alpha = float(act_type.replace("elu", "")) + return nn.ELU(alpha) + else: + raise Exception + + +def get_decoder(config, c): + decoder = None + decoder_options = dict() + if config.model.decoder_type == "unet": + try: + decoder_options = dict(config.decoder_unet) + except: + pass + decoder = smp.Unet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "fpn": + try: + decoder_options = dict(config.decoder_fpn) + except: + pass + decoder = smp.FPN( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "unet++": + try: + decoder_options = dict(config.decoder_unet_plus_plus) + except: + pass + decoder = smp.UnetPlusPlus( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "manet": + try: + decoder_options = dict(config.decoder_manet) + except: + pass + decoder = smp.MAnet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "linknet": + try: + decoder_options = dict(config.decoder_linknet) + except: + pass + decoder = smp.Linknet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "pspnet": + try: + decoder_options = dict(config.decoder_pspnet) + except: + pass + decoder = smp.PSPNet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "pspnet": + try: + decoder_options = dict(config.decoder_pspnet) + except: + pass + decoder = smp.PSPNet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "pan": + try: + decoder_options = dict(config.decoder_pan) + except: + pass + decoder = smp.PAN( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "deeplabv3": + try: + decoder_options = dict(config.decoder_deeplabv3) + except: + pass + decoder = smp.DeepLabV3( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "deeplabv3plus": + try: + decoder_options = dict(config.decoder_deeplabv3plus) + except: + pass + decoder = smp.DeepLabV3Plus( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + return decoder + + +class Segm_Models_Net(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + act = get_act(act_type=config.model.act) + + self.num_target_instruments = len(prefer_target_instrument(config)) + self.num_subbands = config.model.num_subbands + + dim_c = self.num_subbands * config.audio.num_channels * 2 + c = config.model.num_channels + f = config.audio.dim_f // self.num_subbands + + self.first_conv = nn.Conv2d(dim_c, c, 1, 1, 0, bias=False) + + self.unet_model = get_decoder(config, c) + + self.final_conv = nn.Sequential( + nn.Conv2d(c + dim_c, c, 1, 1, 0, bias=False), + act, + nn.Conv2d(c, self.num_target_instruments * dim_c, 1, 1, 0, bias=False), + ) + + self.stft = STFT(config.audio) + + def cac2cws(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c, k, f // k, t) + x = x.reshape(b, c * k, f // k, t) + return x + + def cws2cac(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c // k, k, f, t) + x = x.reshape(b, c // k, f * k, t) + return x + + def forward(self, x): + x = self.stft(x) + + mix = x = self.cac2cws(x) + + first_conv_out = x = self.first_conv(x) + + x = x.transpose(-1, -2) + + x = self.unet_model(x) + + x = x.transpose(-1, -2) + + x = x * first_conv_out # reduce artifacts + + x = self.final_conv(torch.cat([mix, x], 1)) + + x = self.cws2cac(x) + + if self.num_target_instruments > 1: + b, c, f, t = x.shape + x = x.reshape(b, self.num_target_instruments, -1, f, t) + + x = self.stft.inverse(x) + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/torchseg_models.py b/src/third_party/MusicSourceSeparationTraining/models/torchseg_models.py new file mode 100644 index 0000000000000000000000000000000000000000..42056d01e4f219522d99d8c6e6e4fdb7858856ca --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/torchseg_models.py @@ -0,0 +1,253 @@ +import torch +import torch.nn as nn +import torchseg as smp +from utils.model_utils import prefer_target_instrument + + +class STFT: + def __init__(self, config): + self.n_fft = config.n_fft + self.hop_length = config.hop_length + self.window = torch.hann_window(window_length=self.n_fft, periodic=True) + self.dim_f = config.dim_f + + def __call__(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-2] + c, t = x.shape[-2:] + x = x.reshape([-1, t]) + x = torch.stft( + x, + n_fft=self.n_fft, + hop_length=self.hop_length, + window=window, + center=True, + return_complex=True, + ) + x = torch.view_as_real(x) + x = x.permute([0, 3, 1, 2]) + x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape( + [*batch_dims, c * 2, -1, x.shape[-1]] + ) + return x[..., : self.dim_f, :] + + def inverse(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-3] + c, f, t = x.shape[-3:] + n = self.n_fft // 2 + 1 + f_pad = torch.zeros([*batch_dims, c, n - f, t]).to(x.device) + x = torch.cat([x, f_pad], -2) + x = x.reshape([*batch_dims, c // 2, 2, n, t]).reshape([-1, 2, n, t]) + x = x.permute([0, 2, 3, 1]) + x = x[..., 0] + x[..., 1] * 1.0j + x = torch.istft( + x, n_fft=self.n_fft, hop_length=self.hop_length, window=window, center=True + ) + x = x.reshape([*batch_dims, 2, -1]) + return x + + +def get_act(act_type): + if act_type == "gelu": + return nn.GELU() + elif act_type == "relu": + return nn.ReLU() + elif act_type[:3] == "elu": + alpha = float(act_type.replace("elu", "")) + return nn.ELU(alpha) + else: + raise Exception + + +def get_decoder(config, c): + decoder = None + decoder_options = dict() + if config.model.decoder_type == "unet": + try: + decoder_options = dict(config.decoder_unet) + except: + pass + decoder = smp.Unet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "fpn": + try: + decoder_options = dict(config.decoder_fpn) + except: + pass + decoder = smp.FPN( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "unet++": + try: + decoder_options = dict(config.decoder_unet_plus_plus) + except: + pass + decoder = smp.UnetPlusPlus( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "manet": + try: + decoder_options = dict(config.decoder_manet) + except: + pass + decoder = smp.MAnet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "linknet": + try: + decoder_options = dict(config.decoder_linknet) + except: + pass + decoder = smp.Linknet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "pspnet": + try: + decoder_options = dict(config.decoder_pspnet) + except: + pass + decoder = smp.PSPNet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "pspnet": + try: + decoder_options = dict(config.decoder_pspnet) + except: + pass + decoder = smp.PSPNet( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "pan": + try: + decoder_options = dict(config.decoder_pan) + except: + pass + decoder = smp.PAN( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "deeplabv3": + try: + decoder_options = dict(config.decoder_deeplabv3) + except: + pass + decoder = smp.DeepLabV3( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + elif config.model.decoder_type == "deeplabv3plus": + try: + decoder_options = dict(config.decoder_deeplabv3plus) + except: + pass + decoder = smp.DeepLabV3Plus( + encoder_name=config.model.encoder_name, + encoder_weights="imagenet", + in_channels=c, + classes=c, + **decoder_options, + ) + return decoder + + +class Torchseg_Net(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + act = get_act(act_type=config.model.act) + + self.num_target_instruments = len(prefer_target_instrument(config)) + self.num_subbands = config.model.num_subbands + + dim_c = self.num_subbands * config.audio.num_channels * 2 + c = config.model.num_channels + f = config.audio.dim_f // self.num_subbands + + self.first_conv = nn.Conv2d(dim_c, c, 1, 1, 0, bias=False) + + self.unet_model = get_decoder(config, c) + + self.final_conv = nn.Sequential( + nn.Conv2d(c + dim_c, c, 1, 1, 0, bias=False), + act, + nn.Conv2d(c, self.num_target_instruments * dim_c, 1, 1, 0, bias=False), + ) + + self.stft = STFT(config.audio) + + def cac2cws(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c, k, f // k, t) + x = x.reshape(b, c * k, f // k, t) + return x + + def cws2cac(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c // k, k, f, t) + x = x.reshape(b, c // k, f * k, t) + return x + + def forward(self, x): + x = self.stft(x) + + mix = x = self.cac2cws(x) + + first_conv_out = x = self.first_conv(x) + + x = x.transpose(-1, -2) + + x = self.unet_model(x) + + x = x.transpose(-1, -2) + + x = x * first_conv_out # reduce artifacts + + x = self.final_conv(torch.cat([mix, x], 1)) + + x = self.cws2cac(x) + + if self.num_target_instruments > 1: + b, c, f, t = x.shape + x = x.reshape(b, self.num_target_instruments, -1, f, t) + + x = self.stft.inverse(x) + return x diff --git a/src/third_party/MusicSourceSeparationTraining/models/ts_bs_mamba2.py b/src/third_party/MusicSourceSeparationTraining/models/ts_bs_mamba2.py new file mode 100644 index 0000000000000000000000000000000000000000..9324a19cb38c28a62d1cef551b2fd4fdb53b01a1 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/ts_bs_mamba2.py @@ -0,0 +1,467 @@ +# https://github.com/Human9000/nd-Mamba2-torch + +from __future__ import print_function + +import numpy as np +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint_sequential + +try: + from mamba_ssm.modules.mamba2 import Mamba2 +except Exception as e: + print("Exception during load Mamba2 modules: {}".format(str(e))) + print("Load local torch implementation!") + from .ex_bi_mamba2 import Mamba2 + + +class MambaBlock(nn.Module): + def __init__(self, in_channels): + super(MambaBlock, self).__init__() + self.forward_mamba2 = Mamba2( + d_model=in_channels, + d_state=128, + d_conv=4, + expand=4, + headdim=64, + ) + + self.backward_mamba2 = Mamba2( + d_model=in_channels, + d_state=128, + d_conv=4, + expand=4, + headdim=64, + ) + + def forward(self, input): + forward_f = input + forward_f_output = self.forward_mamba2(forward_f) + backward_f = torch.flip(input, [1]) + backward_f_output = self.backward_mamba2(backward_f) + backward_f_output2 = torch.flip(backward_f_output, [1]) + output = torch.cat([forward_f_output + input, backward_f_output2 + input], -1) + return output + + +class TAC(nn.Module): + """ + A transform-average-concatenate (TAC) module. + """ + + def __init__(self, input_size, hidden_size): + super(TAC, self).__init__() + + self.input_size = input_size + self.eps = torch.finfo(torch.float32).eps + + self.input_norm = nn.GroupNorm(1, input_size, self.eps) + self.TAC_input = nn.Sequential(nn.Linear(input_size, hidden_size), nn.Tanh()) + self.TAC_mean = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.Tanh()) + self.TAC_output = nn.Sequential( + nn.Linear(hidden_size * 2, input_size), nn.Tanh() + ) + + def forward(self, input): + # input shape: batch, group, N, * + + batch_size, G, N = input.shape[:3] + output = self.input_norm(input.view(batch_size * G, N, -1)).view( + batch_size, G, N, -1 + ) + T = output.shape[-1] + + # transform + group_input = output # B, G, N, T + group_input = ( + group_input.permute(0, 3, 1, 2).contiguous().view(-1, N) + ) # B*T*G, N + group_output = self.TAC_input(group_input).view( + batch_size, T, G, -1 + ) # B, T, G, H + + # mean pooling + group_mean = group_output.mean(2).view(batch_size * T, -1) # B*T, H + group_mean = ( + self.TAC_mean(group_mean) + .unsqueeze(1) + .expand(batch_size * T, G, group_mean.shape[-1]) + .contiguous() + ) # B*T, G, H + + # concate + group_output = group_output.view(batch_size * T, G, -1) # B*T, G, H + group_output = torch.cat([group_output, group_mean], 2) # B*T, G, 2H + group_output = self.TAC_output( + group_output.view(-1, group_output.shape[-1]) + ) # B*T*G, N + group_output = ( + group_output.view(batch_size, T, G, -1).permute(0, 2, 3, 1).contiguous() + ) # B, G, N, T + output = input + group_output.view(input.shape) + + return output + + +class ResMamba(nn.Module): + def __init__(self, input_size, hidden_size, dropout=0.0, bidirectional=True): + super(ResMamba, self).__init__() + + self.input_size = input_size + self.hidden_size = hidden_size + self.eps = torch.finfo(torch.float32).eps + + self.norm = nn.GroupNorm(1, input_size, self.eps) + self.dropout = nn.Dropout(p=dropout) + self.rnn = MambaBlock(input_size) + self.proj = nn.Linear(input_size * 2, input_size) + # linear projection layer + + def forward(self, input): + # input shape: batch, dim, seq + rnn_output = self.rnn( + self.dropout(self.norm(input)).transpose(1, 2).contiguous() + ) + rnn_output = self.proj( + rnn_output.contiguous().view(-1, rnn_output.shape[2]) + ).view(input.shape[0], input.shape[2], input.shape[1]) + + return input + rnn_output.transpose(1, 2).contiguous() + + +class BSNet(nn.Module): + def __init__(self, in_channel, nband=7): + super(BSNet, self).__init__() + + self.nband = nband + self.feature_dim = in_channel // nband + + self.band_rnn = ResMamba(self.feature_dim, self.feature_dim * 2) + self.band_comm = ResMamba(self.feature_dim, self.feature_dim * 2) + self.channel_comm = TAC(self.feature_dim, self.feature_dim * 3) + + def forward(self, input): + # input shape: B, nch, nband*N, T + B, nch, N, T = input.shape + + band_output = self.band_rnn( + input.view(B * nch * self.nband, self.feature_dim, -1) + ).view(B * nch, self.nband, -1, T) + + # band comm + band_output = ( + band_output.permute(0, 3, 2, 1) + .contiguous() + .view(B * nch * T, -1, self.nband) + ) + output = ( + self.band_comm(band_output) + .view(B * nch, T, -1, self.nband) + .permute(0, 3, 2, 1) + .contiguous() + ) + + # channel comm + output = ( + output.view(B, nch, self.nband, -1, T) + .transpose(1, 2) + .contiguous() + .view(B * self.nband, nch, -1, T) + ) + output = ( + self.channel_comm(output) + .view(B, self.nband, nch, -1, T) + .transpose(1, 2) + .contiguous() + ) + + return output.view(B, nch, N, T) + + +class Separator(nn.Module): + def __init__( + self, + sr=44100, + win=2048, + stride=512, + feature_dim=128, + num_repeat_mask=8, + num_repeat_map=4, + num_output=4, + ): + super(Separator, self).__init__() + + self.sr = sr + self.win = win + self.stride = stride + self.group = self.win // 2 + self.enc_dim = self.win // 2 + 1 + self.feature_dim = feature_dim + self.num_output = num_output + self.eps = torch.finfo(torch.float32).eps + + # 0-1k (50 hop), 1k-2k (100 hop), 2k-4k (250 hop), 4k-8k (500 hop), 8k-16k (1k hop), 16k-20k (2k hop), 20k-inf + bandwidth_50 = int(np.floor(50 / (sr / 2.0) * self.enc_dim)) + bandwidth_100 = int(np.floor(100 / (sr / 2.0) * self.enc_dim)) + bandwidth_250 = int(np.floor(250 / (sr / 2.0) * self.enc_dim)) + bandwidth_500 = int(np.floor(500 / (sr / 2.0) * self.enc_dim)) + bandwidth_1k = int(np.floor(1000 / (sr / 2.0) * self.enc_dim)) + bandwidth_2k = int(np.floor(2000 / (sr / 2.0) * self.enc_dim)) + self.band_width = [bandwidth_50] * 20 + self.band_width += [bandwidth_100] * 10 + self.band_width += [bandwidth_250] * 8 + self.band_width += [bandwidth_500] * 8 + self.band_width += [bandwidth_1k] * 8 + self.band_width += [bandwidth_2k] * 2 + self.band_width.append(self.enc_dim - np.sum(self.band_width)) + self.nband = len(self.band_width) + print(self.band_width) + + self.BN_mask = nn.ModuleList([]) + for i in range(self.nband): + self.BN_mask.append( + nn.Sequential( + nn.GroupNorm(1, self.band_width[i] * 2, self.eps), + nn.Conv1d(self.band_width[i] * 2, self.feature_dim, 1), + ) + ) + + self.BN_map = nn.ModuleList([]) + for i in range(self.nband): + self.BN_map.append( + nn.Sequential( + nn.GroupNorm(1, self.band_width[i] * 2, self.eps), + nn.Conv1d(self.band_width[i] * 2, self.feature_dim, 1), + ) + ) + + self.separator_mask = [] + for i in range(num_repeat_mask): + self.separator_mask.append(BSNet(self.nband * self.feature_dim, self.nband)) + self.separator_mask = nn.Sequential(*self.separator_mask) + + self.separator_map = [] + for i in range(num_repeat_map): + self.separator_map.append(BSNet(self.nband * self.feature_dim, self.nband)) + self.separator_map = nn.Sequential(*self.separator_map) + + self.in_conv = nn.Conv1d(self.feature_dim * 2, self.feature_dim, 1) + self.Tanh = nn.Tanh() + self.mask = nn.ModuleList([]) + self.map = nn.ModuleList([]) + for i in range(self.nband): + self.mask.append( + nn.Sequential( + nn.GroupNorm(1, self.feature_dim, torch.finfo(torch.float32).eps), + nn.Conv1d( + self.feature_dim, self.feature_dim * 1 * self.num_output, 1 + ), + nn.Tanh(), + nn.Conv1d( + self.feature_dim * 1 * self.num_output, + self.feature_dim * 1 * self.num_output, + 1, + groups=self.num_output, + ), + nn.Tanh(), + nn.Conv1d( + self.feature_dim * 1 * self.num_output, + self.band_width[i] * 4 * self.num_output, + 1, + groups=self.num_output, + ), + ) + ) + self.map.append( + nn.Sequential( + nn.GroupNorm(1, self.feature_dim, torch.finfo(torch.float32).eps), + nn.Conv1d( + self.feature_dim, self.feature_dim * 1 * self.num_output, 1 + ), + nn.Tanh(), + nn.Conv1d( + self.feature_dim * 1 * self.num_output, + self.feature_dim * 1 * self.num_output, + 1, + groups=self.num_output, + ), + nn.Tanh(), + nn.Conv1d( + self.feature_dim * 1 * self.num_output, + self.band_width[i] * 4 * self.num_output, + 1, + groups=self.num_output, + ), + ) + ) + + def pad_input(self, input, window, stride): + """ + Zero-padding input according to window/stride size. + """ + batch_size, nsample = input.shape + + # pad the signals at the end for matching the window/stride size + rest = window - (stride + nsample % window) % window + if rest > 0: + pad = torch.zeros(batch_size, rest).type(input.type()) + input = torch.cat([input, pad], 1) + pad_aux = torch.zeros(batch_size, stride).type(input.type()) + input = torch.cat([pad_aux, input, pad_aux], 1) + + return input, rest + + def forward(self, input): + # input shape: (B, C, T) + + batch_size, nch, nsample = input.shape + input = input.view(batch_size * nch, -1) + + # frequency-domain separation + spec = torch.stft( + input, + n_fft=self.win, + hop_length=self.stride, + window=torch.hann_window(self.win).to(input.device).type(input.type()), + return_complex=True, + ) + + # concat real and imag, split to subbands + spec_RI = torch.stack([spec.real, spec.imag], 1) # B*nch, 2, F, T + subband_spec_RI = [] + subband_spec = [] + band_idx = 0 + for i in range(len(self.band_width)): + subband_spec_RI.append( + spec_RI[:, :, band_idx : band_idx + self.band_width[i]].contiguous() + ) + subband_spec.append( + spec[:, band_idx : band_idx + self.band_width[i]] + ) # B*nch, BW, T + band_idx += self.band_width[i] + + # normalization and bottleneck + subband_feature_mask = [] + for i in range(len(self.band_width)): + subband_feature_mask.append( + self.BN_mask[i]( + subband_spec_RI[i].view( + batch_size * nch, self.band_width[i] * 2, -1 + ) + ) + ) + subband_feature_mask = torch.stack(subband_feature_mask, 1) # B, nband, N, T + + subband_feature_map = [] + for i in range(len(self.band_width)): + subband_feature_map.append( + self.BN_map[i]( + subband_spec_RI[i].view( + batch_size * nch, self.band_width[i] * 2, -1 + ) + ) + ) + subband_feature_map = torch.stack(subband_feature_map, 1) # B, nband, N, T + # separator + sep_output = checkpoint_sequential( + self.separator_mask, + 2, + subband_feature_mask.view( + batch_size, nch, self.nband * self.feature_dim, -1 + ), + ) # B, nband*N, T + sep_output = sep_output.view(batch_size * nch, self.nband, self.feature_dim, -1) + combined = torch.cat((subband_feature_map, sep_output), dim=2) + combined1 = combined.reshape( + batch_size * nch * self.nband, self.feature_dim * 2, -1 + ) + combined2 = self.Tanh(self.in_conv(combined1)) + combined3 = combined2.reshape( + batch_size * nch, self.nband, self.feature_dim, -1 + ) + sep_output2 = checkpoint_sequential( + self.separator_map, + 2, + combined3.view(batch_size, nch, self.nband * self.feature_dim, -1), + ) # 1B, nband*N, T + sep_output2 = sep_output2.view( + batch_size * nch, self.nband, self.feature_dim, -1 + ) + + sep_subband_spec = [] + sep_subband_spec_mask = [] + for i in range(self.nband): + this_output = self.mask[i](sep_output[:, i]).view( + batch_size * nch, 2, 2, self.num_output, self.band_width[i], -1 + ) + this_mask = this_output[:, 0] * torch.sigmoid( + this_output[:, 1] + ) # B*nch, 2, K, BW, T + this_mask_real = this_mask[:, 0] # B*nch, K, BW, T + this_mask_imag = this_mask[:, 1] # B*nch, K, BW, T + # force mask sum to 1 + this_mask_real_sum = this_mask_real.sum(1).unsqueeze(1) # B*nch, 1, BW, T + this_mask_imag_sum = this_mask_imag.sum(1).unsqueeze(1) # B*nch, 1, BW, T + this_mask_real = this_mask_real - (this_mask_real_sum - 1) / self.num_output + this_mask_imag = this_mask_imag - this_mask_imag_sum / self.num_output + est_spec_real = ( + subband_spec[i].real.unsqueeze(1) * this_mask_real + - subband_spec[i].imag.unsqueeze(1) * this_mask_imag + ) # B*nch, K, BW, T + est_spec_imag = ( + subband_spec[i].real.unsqueeze(1) * this_mask_imag + + subband_spec[i].imag.unsqueeze(1) * this_mask_real + ) # B*nch, K, BW, T + + ################################## + this_output2 = self.map[i](sep_output2[:, i]).view( + batch_size * nch, 2, 2, self.num_output, self.band_width[i], -1 + ) + this_map = this_output2[:, 0] * torch.sigmoid( + this_output2[:, 1] + ) # B*nch, 2, K, BW, T + this_map_real = this_map[:, 0] # B*nch, K, BW, T + this_map_imag = this_map[:, 1] # B*nch, K, BW, T + est_spec_real2 = est_spec_real + this_map_real + est_spec_imag2 = est_spec_imag + this_map_imag + + sep_subband_spec.append(torch.complex(est_spec_real2, est_spec_imag2)) + sep_subband_spec_mask.append(torch.complex(est_spec_real, est_spec_imag)) + + sep_subband_spec = torch.cat(sep_subband_spec, 2) + est_spec_mask = torch.cat(sep_subband_spec_mask, 2) + + output = torch.istft( + sep_subband_spec.view(batch_size * nch * self.num_output, self.enc_dim, -1), + n_fft=self.win, + hop_length=self.stride, + window=torch.hann_window(self.win).to(input.device).type(input.type()), + length=nsample, + ) + output_mask = torch.istft( + est_spec_mask.view(batch_size * nch * self.num_output, self.enc_dim, -1), + n_fft=self.win, + hop_length=self.stride, + window=torch.hann_window(self.win).to(input.device).type(input.type()), + length=nsample, + ) + + output = ( + output.view(batch_size, nch, self.num_output, -1) + .transpose(1, 2) + .contiguous() + ) + output_mask = ( + output_mask.view(batch_size, nch, self.num_output, -1) + .transpose(1, 2) + .contiguous() + ) + # return output, output_mask + return output + + +if __name__ == "__main__": + model = Separator().cuda() + arr = np.zeros((1, 2, 3 * 44100), dtype=np.float32) + x = torch.from_numpy(arr).cuda() + res = model(x) diff --git a/src/third_party/MusicSourceSeparationTraining/models/upernet_swin_transformers.py b/src/third_party/MusicSourceSeparationTraining/models/upernet_swin_transformers.py new file mode 100644 index 0000000000000000000000000000000000000000..b530b2346ef5114c5eee12901b515f0756cd108e --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/models/upernet_swin_transformers.py @@ -0,0 +1,249 @@ +from functools import partial + +import torch +import torch.nn as nn +from transformers import UperNetForSemanticSegmentation +from utils.model_utils import prefer_target_instrument + + +class STFT: + def __init__(self, config): + self.n_fft = config.n_fft + self.hop_length = config.hop_length + self.window = torch.hann_window(window_length=self.n_fft, periodic=True) + self.dim_f = config.dim_f + + def __call__(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-2] + c, t = x.shape[-2:] + x = x.reshape([-1, t]) + x = torch.stft( + x, + n_fft=self.n_fft, + hop_length=self.hop_length, + window=window, + center=True, + return_complex=True, + ) + x = torch.view_as_real(x) + x = x.permute([0, 3, 1, 2]) + x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape( + [*batch_dims, c * 2, -1, x.shape[-1]] + ) + return x[..., : self.dim_f, :] + + def inverse(self, x): + window = self.window.to(x.device) + batch_dims = x.shape[:-3] + c, f, t = x.shape[-3:] + n = self.n_fft // 2 + 1 + f_pad = torch.zeros([*batch_dims, c, n - f, t]).to(x.device) + x = torch.cat([x, f_pad], -2) + x = x.reshape([*batch_dims, c // 2, 2, n, t]).reshape([-1, 2, n, t]) + x = x.permute([0, 2, 3, 1]) + x = x[..., 0] + x[..., 1] * 1.0j + x = torch.istft( + x, n_fft=self.n_fft, hop_length=self.hop_length, window=window, center=True + ) + x = x.reshape([*batch_dims, 2, -1]) + return x + + +def get_norm(norm_type): + def norm(c, norm_type): + if norm_type == "BatchNorm": + return nn.BatchNorm2d(c) + elif norm_type == "InstanceNorm": + return nn.InstanceNorm2d(c, affine=True) + elif "GroupNorm" in norm_type: + g = int(norm_type.replace("GroupNorm", "")) + return nn.GroupNorm(num_groups=g, num_channels=c) + else: + return nn.Identity() + + return partial(norm, norm_type=norm_type) + + +def get_act(act_type): + if act_type == "gelu": + return nn.GELU() + elif act_type == "relu": + return nn.ReLU() + elif act_type[:3] == "elu": + alpha = float(act_type.replace("elu", "")) + return nn.ELU(alpha) + else: + raise Exception + + +class Upscale(nn.Module): + def __init__(self, in_c, out_c, scale, norm, act): + super().__init__() + self.conv = nn.Sequential( + norm(in_c), + act, + nn.ConvTranspose2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=scale, + stride=scale, + bias=False, + ), + ) + + def forward(self, x): + return self.conv(x) + + +class Downscale(nn.Module): + def __init__(self, in_c, out_c, scale, norm, act): + super().__init__() + self.conv = nn.Sequential( + norm(in_c), + act, + nn.Conv2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=scale, + stride=scale, + bias=False, + ), + ) + + def forward(self, x): + return self.conv(x) + + +class TFC_TDF(nn.Module): + def __init__(self, in_c, c, l, f, bn, norm, act): + super().__init__() + + self.blocks = nn.ModuleList() + for i in range(l): + block = nn.Module() + + block.tfc1 = nn.Sequential( + norm(in_c), + act, + nn.Conv2d(in_c, c, 3, 1, 1, bias=False), + ) + block.tdf = nn.Sequential( + norm(c), + act, + nn.Linear(f, f // bn, bias=False), + norm(c), + act, + nn.Linear(f // bn, f, bias=False), + ) + block.tfc2 = nn.Sequential( + norm(c), + act, + nn.Conv2d(c, c, 3, 1, 1, bias=False), + ) + block.shortcut = nn.Conv2d(in_c, c, 1, 1, 0, bias=False) + + self.blocks.append(block) + in_c = c + + def forward(self, x): + for block in self.blocks: + s = block.shortcut(x) + x = block.tfc1(x) + x = x + block.tdf(x) + x = block.tfc2(x) + x = x + s + return x + + +class Swin_UperNet_Model(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + act = get_act(act_type=config.model.act) + + self.num_target_instruments = len(prefer_target_instrument(config)) + self.num_subbands = config.model.num_subbands + + dim_c = self.num_subbands * config.audio.num_channels * 2 + c = config.model.num_channels + f = config.audio.dim_f // self.num_subbands + + self.first_conv = nn.Conv2d(dim_c, c, 1, 1, 0, bias=False) + + self.swin_upernet_model = UperNetForSemanticSegmentation.from_pretrained( + "openmmlab/upernet-swin-large" + ) + + self.swin_upernet_model.auxiliary_head.classifier = nn.Conv2d( + 256, c, kernel_size=(1, 1), stride=(1, 1) + ) + self.swin_upernet_model.decode_head.classifier = nn.Conv2d( + 512, c, kernel_size=(1, 1), stride=(1, 1) + ) + self.swin_upernet_model.backbone.embeddings.patch_embeddings.projection = ( + nn.Conv2d(c, 192, kernel_size=(4, 4), stride=(4, 4)) + ) + + self.final_conv = nn.Sequential( + nn.Conv2d(c + dim_c, c, 1, 1, 0, bias=False), + act, + nn.Conv2d(c, self.num_target_instruments * dim_c, 1, 1, 0, bias=False), + ) + + self.stft = STFT(config.audio) + + def cac2cws(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c, k, f // k, t) + x = x.reshape(b, c * k, f // k, t) + return x + + def cws2cac(self, x): + k = self.num_subbands + b, c, f, t = x.shape + x = x.reshape(b, c // k, k, f, t) + x = x.reshape(b, c // k, f * k, t) + return x + + def forward(self, x): + x = self.stft(x) + + mix = x = self.cac2cws(x) + + first_conv_out = x = self.first_conv(x) + + x = x.transpose(-1, -2) + + x = self.swin_upernet_model(x).logits + + x = x.transpose(-1, -2) + + x = x * first_conv_out # reduce artifacts + + x = self.final_conv(torch.cat([mix, x], 1)) + + x = self.cws2cac(x) + + if self.num_target_instruments > 1: + b, c, f, t = x.shape + x = x.reshape(b, self.num_target_instruments, -1, f, t) + + x = self.stft.inverse(x) + return x + + +if __name__ == "__main__": + model = UperNetForSemanticSegmentation.from_pretrained( + "./results/", ignore_mismatched_sizes=True + ) + print(model) + print(model.auxiliary_head.classifier) + print(model.decode_head.classifier) + + x = torch.zeros((2, 16, 512, 512), dtype=torch.float32) + res = model(x) + print(res.logits.shape) + model.save_pretrained("./results/") diff --git a/src/third_party/MusicSourceSeparationTraining/requirements.txt b/src/third_party/MusicSourceSeparationTraining/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..b9ce4b3cee795c9495c42377d044a5bbd3ae54a2 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/requirements.txt @@ -0,0 +1,40 @@ +torch>=2.0.1 +torchaudio +numpy +pandas +scipy +soundfile +ml_collections +tqdm +segmentation_models_pytorch==0.3.3 +timm==0.9.2 +audiomentations==0.24.0 +pedalboard~=0.8.1 +omegaconf==2.2.3 +beartype==0.14.1 +rotary_embedding_torch==0.3.5 +einops==0.8.1 +librosa +demucs==4.0.0 +transformers~=4.35.0 +torchmetrics==0.11.4 +spafe==0.3.2 +protobuf==3.20.3 +torch_audiomentations +asteroid==0.7.0 +auraloss +torchseg +bitsandbytes +wandb +accelerate +huggingface-hub>=0.23.0 +prodigyopt +torch_log_wmse>=0.3.1 +torch_l1_snr>=0.1.2 +loralib +pyaudio +wxpython==4.2.2 +keyboard +matplotlib +hyper_connections==0.1.11 +sageattention==1.0.6 diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/do_metadata.py b/src/third_party/MusicSourceSeparationTraining/scripts/do_metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..f53a50e83e0099ff22101c910008c23e69ebd59b --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/do_metadata.py @@ -0,0 +1,9 @@ +from utils.dataset import prepare_data +from utils.settings import get_model_from_config, parse_args_train + +args = parse_args_train(None) +_, config = get_model_from_config(args.model_type, args.config_path) + +batch_size = config.training.batch_size * args.device_ids + +prepare_data(config, args, batch_size) diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/moises_to_musdb.py b/src/third_party/MusicSourceSeparationTraining/scripts/moises_to_musdb.py new file mode 100644 index 0000000000000000000000000000000000000000..03e8afc1180728141db43eb54dccf31a0947c1bd --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/moises_to_musdb.py @@ -0,0 +1,459 @@ +import argparse +import os +import shutil +import time +from multiprocessing import Pool +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import soundfile as sf +from tqdm import tqdm + + +def combine_audio_files(files: List[str]) -> Tuple[np.ndarray, int]: + """ + Combines multiple audio files into one by overlaying them. + Parameters: + - files (List[str]): List of file paths to be combined. + Returns: + - Tuple[np.ndarray, int]: A Tuple containing the combined audio data array and sample rate. + """ + combined_data, sample_rate = sf.read(files[0]) + + for file in files[1:]: + data, sr = sf.read(file) + if len(data) > len(combined_data): + combined_data = np.pad( + combined_data, ((0, len(data) - len(combined_data)), (0, 0)), "constant" + ) + elif len(combined_data) > len(data): + data = np.pad( + data, ((0, len(combined_data) - len(data)), (0, 0)), "constant" + ) + combined_data += data + + return combined_data, sample_rate + + +# Finds all .wav files located in folders that do not contain specified categories in their folder names +def files_to_categories(src_folder: str, categories: List[str]) -> Dict[str, List[str]]: + """ + Finds all .wav files located in folders that do not contain specified categories + in their folder names, within the given src_folder directory. + Parameters: + - src_folder (str): Path to the main directory containing subdirectories with files. + - categories (List[str]): Keywords that should not be part of the folder's name. + Returns: + - Dict[str, List[str]]: A Dict with keys as categories, values as lists of paths to .wav files found. + """ + files = {category: [] for category in categories + ["other"]} + + for folder in os.listdir(src_folder): + folder_path = os.path.join(src_folder, folder) + if os.path.isdir(folder_path): + if folder.lower() in categories: + stem = folder.lower() + else: + stem = "other" + for f in os.listdir(folder_path): + if f.endswith(".wav"): + files[stem].append(os.path.join(folder_path, f)) + return files + + +# Processes a folder containing audio tracks, copying and combining necessary files into the target structure +def process_folder( + src_folder: str, dest_folder: str, stems: List[str], trim: bool = False +) -> None: + """ + Processes a folder containing audio tracks, copying and combining necessary files into the target structure. + Parameters: + - src_folder (str): Path to the source folder of MoisesDB. + - dest_folder (str): Path to the target folder for MUSDB18. + - stems (List[str]): List of stem categories to process. + - trim (bool): If True, trim all stems to the length of the shortest one. + """ + + if not os.path.exists(dest_folder): + os.makedirs(dest_folder) + + categories = stems + + if trim: + # First pass: load all stems and find the minimum length + stem_data = {} + sample_rate = None + min_length = float("inf") + all_files = files_to_categories(src_folder, categories) + + # Using tqdm to display progress for categories + for category in tqdm( + categories, desc=f"Processing categories in {os.path.basename(src_folder)}" + ): + files = all_files[category] + if files: + if len(files) > 1: + combined_data, sr = combine_audio_files(files) + else: + combined_data, sr = sf.read(files[0]) + + stem_data[category] = combined_data + sample_rate = sr + min_length = min(min_length, len(combined_data)) + + # Process 'other' files + other_files = all_files["other"] + if other_files: + other_combined_data, sr = combine_audio_files(other_files) + stem_data["other"] = other_combined_data + sample_rate = sr + min_length = min(min_length, len(other_combined_data)) + + # If no stems were found, set a default sample rate + if sample_rate is None: + sample_rate = 44100 + min_length = 0 + + # Second pass: trim all stems to min_length and save + for category in categories: + if category in stem_data: + trimmed_data = stem_data[category][:min_length] + sf.write( + os.path.join(dest_folder, f"{category}.wav"), + trimmed_data, + sample_rate, + ) + else: + # Create silence with min_length + silence = np.zeros((min_length, 2), dtype=np.float32) + sf.write( + os.path.join(dest_folder, f"{category}.wav"), silence, sample_rate + ) + + # Save 'other' stem + if "other" in stem_data: + trimmed_other = stem_data["other"][:min_length] + sf.write(os.path.join(dest_folder, "other.wav"), trimmed_other, sample_rate) + else: + silence = np.zeros((min_length, 2), dtype=np.float32) + sf.write(os.path.join(dest_folder, "other.wav"), silence, sample_rate) + + # Create mixture.wav from all files, then trim to min_length + all_files_list = [file for sublist in all_files.values() for file in sublist] + if all_files_list: + mixture_data, sample_rate = combine_audio_files(all_files_list) + mixture_data = mixture_data[:min_length] + sf.write( + os.path.join(dest_folder, "mixture.wav"), mixture_data, sample_rate + ) + else: + # If no files at all, create silent mixture + silence = np.zeros((min_length, 2), dtype=np.float32) + sf.write(os.path.join(dest_folder, "mixture.wav"), silence, sample_rate) + else: + # Original behavior: If the required stem does not exist in the source folder (src_folder), + # we add silence instead of the file with the same duration as the standard file. + problem_categories = [] + duration = 0 + all_files = files_to_categories(src_folder, categories) + + # Using tqdm to display progress for categories + for category in tqdm( + categories, desc=f"Processing categories in {os.path.basename(src_folder)}" + ): + files = all_files[category] + if files: + if len(files) > 1: + combined_data, sample_rate = combine_audio_files(files) + else: + combined_data, sample_rate = sf.read(files[0]) + + sf.write( + os.path.join(dest_folder, f"{category}.wav"), + combined_data, + sample_rate, + ) + duration = max(duration, len(combined_data) / sample_rate) + else: + problem_categories.append(category) + + other_files = all_files["other"] + if other_files: + other_combined_data, sample_rate = combine_audio_files(other_files) + sf.write( + os.path.join(dest_folder, "other.wav"), other_combined_data, sample_rate + ) + else: + problem_categories.append("other") + + for category in problem_categories: + silence = np.zeros((int(duration * sample_rate), 2), dtype=np.float32) + sf.write(os.path.join(dest_folder, f"{category}.wav"), silence, sample_rate) + # mixture.wav + all_files_list = [file for sublist in all_files.values() for file in sublist] + mixture_data, sample_rate = combine_audio_files(all_files_list) + sf.write(os.path.join(dest_folder, "mixture.wav"), mixture_data, sample_rate) + + +# Wrapper function for 'process_folder' that unpacks the arguments +def process_folder_wrapper(args: Tuple[str, str, List[str], bool]) -> None: + """ + A wrapper function for 'process_folder' that unpacks the arguments. + Parameters: + - args (Tuple[str, str, List[str], bool]): A Tuple containing the source folder, destination folder paths, stems, and trim flag. + """ + src_folder, dest_folder, stems, trim = args + return process_folder(src_folder, dest_folder, stems, trim) + + +# Converts MoisesDB dataset to MUSDB18 format for a specified number of folders +def convert_dataset( + src_root: str, + dest_root: str, + stems: List[str], + max_folders: int = 240, + num_workers: int = 4, + trim: bool = False, +) -> None: + """ + Converts MoisesDB dataset to MUSDB18 format for a specified number of folders. + Parameters: + - src_root (str): Root directory of the MoisesDB dataset. + - dest_root (str): Root directory where the new dataset will be saved. + - max_folders (int): Maximum number of folders to process. + - num_workers (int): Number of parallel workers for processing. + - trim (bool): If True, trim all stems to the length of the shortest one. + """ + folders_to_process = [] + for folder in os.listdir(src_root): + if len(folders_to_process) >= max_folders: + break + + src_folder = os.path.join(src_root, folder) + dest_folder = os.path.join(dest_root, folder) + + if os.path.isdir(src_folder): + folders_to_process.append((src_folder, dest_folder, stems, trim)) + else: + print(f"Skip {src_folder} — not dir") + + with Pool(num_workers) as pool: + pool.map(process_folder_wrapper, folders_to_process) + + +# Count number of subfolders in a folder +def count_folders_in_folder(args_to_func) -> Dict[str, int]: + """ + Counts the number of subfolders inside a given folder. + + Parameters: + - folder_path (str): Path to the folder where the count is needed. + + Returns: + - Dict[str, int]: A dictionary with folder paths as keys and subfolder counts as values. + """ + folder_count = 0 + folder_path, stems = args_to_func + if os.path.isdir(folder_path): + # Count subfolders in stems + folder_count = len( + [ + f + for f in os.listdir(folder_path) + if os.path.isdir(os.path.join(folder_path, f)) and f in stems + ] + ) + # For other.wav + if any( + os.path.isdir(os.path.join(folder_path, f)) and f not in stems + for f in os.listdir(folder_path) + ): + folder_count += 1 + + return {folder_path: folder_count} + + +# Parallel count of subfolders in each folder inside src_folder +def count_folders_parallel( + src_folder: str, stems, num_workers: int = 4 +) -> Dict[str, int]: + """ + Parallelly counts the number of subfolders in each folder inside src_folder. + + Parameters: + - src_folder (str): Root folder containing subfolders to count. + + Returns: + - Dict[str, int]: Dictionary with folder paths as keys and subfolder counts as values. + """ + # Get list of all folders inside src_folder + folders_to_process = [ + os.path.join(src_folder, folder) + for folder in os.listdir(src_folder) + if os.path.isdir(os.path.join(src_folder, folder)) + ] + + args_to_func = [(folder, stems) for folder in folders_to_process] + + # Parallelly process each folder using pool.map + with Pool(num_workers) as pool: + results = pool.map(count_folders_in_folder, args_to_func) + + # Merge results from different processes + merged_counts = {} + for result in results: + merged_counts.update(result) + + return merged_counts + + +def parse_args(dict_args: Union[Dict, None]) -> argparse.Namespace: + """ + Parse command-line arguments for configuring the model, dataset, and training parameters. + + Args: + dict_args: Dict of command-line arguments. If None, arguments will be parsed from sys.argv. + + Returns: + Namespace object containing parsed arguments and their values. + """ + parser = argparse.ArgumentParser( + description="Copy mixture files from VALID_DIR to INFERENCE_DIR" + ) + parser.add_argument( + "--src_dir", + type=str, + required=True, + help="Source directory with MoisesDB tracks", + ) + parser.add_argument( + "--dest_dir", + type=str, + required=True, + help="Directory to save tracks in MUSDB18", + ) + parser.add_argument( + "--num_workers", type=int, default=os.cpu_count(), help="Num of processors" + ) + parser.add_argument( + "--max_folders", type=int, default=240, help="Num of folders to use" + ) + parser.add_argument( + "--create_valid", action="store_true", help="Create valid folders or not" + ) + parser.add_argument( + "--valid_dir", type=str, default=r"\valid", help="Directory for valid" + ) + parser.add_argument( + "--valid_size", type=int, default=10, help="Num of folders to use in valitd" + ) + parser.add_argument( + "--stems", + nargs="+", + type=str, + default=["bass", "drums", "vocals"], + choices=[ + "drums", + "guitar", + "vocals", + "bass", + "other_keys", + "piano", + "percussion", + "bowed_strings", + "wind", + "other_plucked", + ], + help="List of stems to use.", + ) + parser.add_argument( + "--mixture_name", type=str, default="mixture.wav", help="Name of mixture tracks" + ) + parser.add_argument( + "--trim", + action="store_true", + help="Trim all stems to the length of the shortest one", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + return args + + +def main(args: Optional[argparse.Namespace] = None) -> None: + start = time.time() + + args = parse_args(args) + + source_directory = args.src_dir + destination_directory = args.dest_dir + num_workers = args.num_workers + stems = args.stems + max_folders = args.max_folders + trim = args.trim + + print(f"num_workers: {num_workers}, categories: {stems + ['other']}, trim: {trim}") + + convert_dataset( + source_directory, + destination_directory, + stems, + max_folders=max_folders, + num_workers=num_workers, + trim=trim, + ) + + print( + f"All {max_folders} files have been processed, time: {time.time() - start:.2f} sec" + ) + + if args.create_valid: + # Count folders in the MoisesDB dataset + result = count_folders_parallel(source_directory, stems) + result = dict(sorted(result.items(), key=lambda item: item[1], reverse=True)) + # valid_size = min(args.valid_size, max_folders) + valid_size = 10 + list_folders = list(result.keys()) + print(f"Top {valid_size} folders:") + for track in list(result.items())[:valid_size]: + print(track) + + valid_folder = args.valid_dir + train_tracks_folder = destination_directory + + # Create valid folder if not exists + if not os.path.exists(valid_folder): + os.makedirs(valid_folder) + + num_val = 0 + + # Copy folders from train_tracks to valid folder + for folder in list_folders: + folder_name = os.path.basename(folder) # Get folder name + + # Form the path to the folder in train_tracks + source_folder = os.path.join(train_tracks_folder, folder_name) + + # If the folder exists in train_tracks, copy it to valid + if os.path.exists(source_folder): + destination = os.path.join(valid_folder, folder_name) + shutil.copytree(source_folder, destination) + shutil.rmtree(source_folder) + num_val += 1 + print(f"Folder: {folder}, num_stems: {result[folder]}") + if num_val >= valid_size: + break + else: + print(f"Folder {source_folder} not found.") + + print("The end!") + + +if __name__ == "__main__": + main(None) diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/prepare_weights_for_inference.py b/src/third_party/MusicSourceSeparationTraining/scripts/prepare_weights_for_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..5b21fb8ac70df7685bf68e1ee8511891e4ba2d4f --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/prepare_weights_for_inference.py @@ -0,0 +1,44 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" + +import argparse + +import torch + + +def clean_weights(args): + weights = torch.load(args.checkpoint, map_location="cpu") + print("Keys: {}".format(list(weights.keys()))) + if "model_state_dict" in list(weights.keys()): + weights = weights["model_state_dict"] + if args.float16: + for el in weights: + weights[el] = weights[el].to(torch.float16) + torch.save(weights, args.output_file) + + +if __name__ == "__main__": + dict_args = None + parser = argparse.ArgumentParser( + description="Clean all except weights from checkpoint file. Optionally converts to float16." + ) + parser.add_argument("--checkpoint", type=str, help="Input checkpoint to clean") + parser.add_argument( + "--output_file", type=str, help="File to save cleaned checkpoint" + ) + parser.add_argument( + "--float16", + action="store_true", + help="Convert weights to float16 instead of float32. Reduce weights size two times.", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + print(args) + else: + args = parser.parse_args() + + clean_weights(args) diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/print_list_of_model_layers.py b/src/third_party/MusicSourceSeparationTraining/scripts/print_list_of_model_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..ed07ef3b14d2031221cd553871823d8a3f71683e --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/print_list_of_model_layers.py @@ -0,0 +1,51 @@ +""" +Script to list layers of model for possible freeze +""" + +import argparse + +from utils.settings import get_model_from_config + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_type", + required=True, + type=str, + default="mdx23c", + help="One of mdx23c, htdemucs, segm_models, mel_band_roformer, bs_roformer, swin_upernet, bandit, etc", + ) + parser.add_argument( + "--config_path", required=True, type=str, help="path to config file" + ) + parser.add_argument( + "--output_file", + type=str, + default="layers.txt", + help="path to results file with list of model layers", + ) + parser.add_argument( + "--layer_mask", + nargs="+", + type=str, + help="mask to print layer names containing mask. Can be several masks", + ) + args = parser.parse_args() + return args + + +def main(): + args = parse_args() + model, config = get_model_from_config(args.model_type, args.config_path) + for name, module in model.named_modules(): + if args.layer_mask is not None: + for mask in args.layer_mask: + if mask in name: + print(name) + else: + print(name) + + +if __name__ == "__main__": + main() diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/redact_config.py b/src/third_party/MusicSourceSeparationTraining/scripts/redact_config.py new file mode 100644 index 0000000000000000000000000000000000000000..f4039662287c944bac3015d3436809f1d702d777 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/redact_config.py @@ -0,0 +1,138 @@ +import argparse +import os +import sys +from typing import Dict, Union + +import yaml +from ml_collections import ConfigDict +from omegaconf import OmegaConf + + +def save_config(config: Union[ConfigDict, OmegaConf], save_path: str): + """ + Save a configuration object (ConfigDict or OmegaConf) to a file. + + Parameters: + ---------- + config : Union[ConfigDict, OmegaConf] + The configuration object to save. + save_path : str + The path where the configuration file will be saved. + + Raises: + ------ + ValueError: + If the configuration type is not supported. + """ + + os.makedirs(os.path.dirname(save_path), exist_ok=True) + + try: + with open(save_path, "w") as f: + if isinstance(config, ConfigDict): + yaml.dump( + config.to_dict(), + f, + default_flow_style=False, + sort_keys=False, + allow_unicode=True, + ) + elif isinstance(config, OmegaConf): + OmegaConf.save(config, save_path) + else: + OmegaConf.save(config, save_path) + except Exception as e: + raise ValueError( + f"Error saving configuration: {e}." + f"Unsupported configuration type. Supported types: ConfigDict, OmegaConf." + f"Config type is {type(config)}" + ) + + +def create_test_config( + original_config_path: str, new_config_path: str, model_type: str +): + """ + Create a test configuration file based on an existing configuration. + + Parameters: + ---------- + original_config_path : str + Path to the original configuration file. + new_config_path : str + Path where the new configuration file will be saved. + model_type : str + The type of model (e.g., 'scnet', 'htdemucs'). + + Returns: + ------- + None + """ + from utils.settings import load_config + + config = load_config(model_type=model_type, config_path=original_config_path) + + config["inference"]["batch_size"] = 1 + config["training"]["batch_size"] = 1 + config["training"]["gradient_accumulation_steps"] = 1 + config["training"]["num_epochs"] = 2 + config["training"]["num_steps"] = 3 + + save_config(config, new_config_path) + print(f"Test config created at: {new_config_path}") + + +def parse_args(dict_args: Union[Dict, None]) -> argparse.Namespace: + """ + Parse command-line arguments for configuring the model, dataset, and training parameters. + + Args: + dict_args: Dict of command-line arguments. If None, arguments will be parsed from sys.argv. + + Returns: + Namespace object containing parsed arguments and their values. + """ + + parser = argparse.ArgumentParser() + parser.add_argument( + "--orig_config", type=str, default="", help="Path to the original config file." + ) + parser.add_argument("--model_type", type=str, default="", help="Model type") + parser.add_argument( + "--new_config", + type=str, + default="", + help="Path to save the new test configuration file.", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + # Determine the default path for the new configuration if not provided + if not args.new_config: + original_dir = os.path.dirname(args.orig_config) + tests_dir = os.path.join("tests_cache", original_dir) + os.makedirs(tests_dir, exist_ok=True) + args.new_config = os.path.join(tests_dir, os.path.basename(args.orig_config)) + + return args + + +def redact_config(args): + # Ensure proper imports for utilities + sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + + args = parse_args(args) + + # Create the test configuration + create_test_config(args.orig_config, args.new_config, args.model_type) + return args.new_config + + +if __name__ == "__main__": + redact_config(None) diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/stream.py b/src/third_party/MusicSourceSeparationTraining/scripts/stream.py new file mode 100644 index 0000000000000000000000000000000000000000..53daa0de72160597a74a3cebc58fc6469f490b10 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/stream.py @@ -0,0 +1,443 @@ +import argparse +import asyncio +import os +import queue +import sys +import threading +import time +import warnings +from typing import Any, Dict, Tuple, Union + +import keyboard +import numpy as np +import pyaudio +import soundfile as sf +import torch +from torch import nn + +warnings.filterwarnings( + "ignore", category=UserWarning, message="TypedStorage is deprecated" +) + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from utils.model_utils import apply_tta, demix, load_start_checkpoint +from utils.settings import get_model_from_config + +RATE: int = 44100 # Sampling rate (44.1 kHz) +FORMAT: int = pyaudio.paFloat32 # Audio format + +base_dir = str(os.path.dirname(os.path.abspath(__file__))) + + +def parse_args(dict_args: Union[Dict[str, Any], None]) -> argparse.Namespace: + """ + Parse command-line arguments. + + Args: + dict_args (Union[Dict[str, Any], None]): Optional dictionary of arguments + to override the command-line input. + + Returns: + argparse.Namespace: Parsed arguments as a namespace object. + """ + parser = argparse.ArgumentParser() + parser.add_argument( + "--type", + type=int, + required=True, + choices=[1, 2], + help="Choose script type: 1 or 2", + ) + parser.add_argument( + "--model_type", + type=str, + default="mdx23c", + help="One of bandit, bandit_v2, bs_roformer, htdemucs, mdx23c, mel_band_roformer," + " scnet, scnet_unofficial, segm_models, swin_upernet, torchseg", + ) + parser.add_argument("--config_path", type=str, help="Path to config file") + parser.add_argument( + "--start_check_point", + type=str, + default="", + help="Initial checkpoint to valid weights", + ) + parser.add_argument( + "--out_dir", + type=str, + default="stream_dir", + help="Path to directory with results as wav file", + ) + parser.add_argument( + "--out_name", + type=str, + default="final", + help="Path to directory with results as wav file", + ) + parser.add_argument( + "--device_ids", nargs="+", type=int, default=0, help="List of GPU IDs" + ) + parser.add_argument( + "--force_cpu", + action="store_true", + help="Force the use of CPU even if CUDA is available", + ) + parser.add_argument( + "--use_tta", + action="store_true", + help="Flag adds test time augmentation during inference (polarity and channel inverse)." + "While this triples the runtime, it reduces noise and slightly improves prediction quality.", + ) + parser.add_argument( + "--lora_checkpoint", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + return args + + +def initialize_device(args: argparse.Namespace) -> str: + """ + Initialize device based on the provided arguments. + + Args: + args (argparse.Namespace): Parsed arguments. + + Returns: + str: Device type ('cpu', 'cuda', or 'mps'). + """ + if args.force_cpu: + return "cpu" + elif torch.cuda.is_available(): + return ( + f"cuda:{args.device_ids[0]}" + if isinstance(args.device_ids, list) + else f"cuda:{args.device_ids}" + ) + elif torch.backends.mps.is_available(): + return "mps" + return "cpu" + + +def load_model(args: argparse.Namespace, device: str) -> Tuple[nn.Module, Any]: + """ + Load the model and configuration from the given arguments. + + Args: + args (argparse.Namespace): Parsed arguments. + device (str): Device to load the model on ('cpu', 'cuda', 'mps'). + + Returns: + Tuple[nn.Module, Any]: The loaded model and its configuration. + """ + torch.backends.cudnn.benchmark = True + model, config = get_model_from_config(args.model_type, args.config_path) + + if args.start_check_point: + checkpoint = torch.load( + args.start_check_point, weights_only=False, map_location="cpu" + ) + load_start_checkpoint(args, model, checkpoint, type_="inference") + + if ( + isinstance(args.device_ids, list) + and len(args.device_ids) > 1 + and not args.force_cpu + ): + model = nn.DataParallel(model, device_ids=args.device_ids) + + model = model.to(device) + return model, config + + +def initialize_audio_streams( + chunk_size: int, +) -> Tuple[pyaudio.PyAudio, pyaudio.Stream, pyaudio.Stream]: + """ + Initialize input and output audio streams. + + Args: + chunk_size (int): Size of each audio chunk. + + Returns: + Tuple[pyaudio.PyAudio, pyaudio.Stream, pyaudio.Stream]: PyAudio instance and the input/output streams. + """ + p = pyaudio.PyAudio() + stream_input = p.open( + format=FORMAT, channels=1, rate=RATE, input=True, frames_per_buffer=chunk_size + ) + stream_output = p.open( + format=FORMAT, channels=2, rate=RATE, output=True, frames_per_buffer=chunk_size + ) + return p, stream_input, stream_output + + +def close_audio_streams( + stream_input: pyaudio.Stream, stream_output: pyaudio.Stream, p: pyaudio.PyAudio +) -> None: + """ + Close input and output audio streams. + + Args: + stream_input (pyaudio.Stream): Input audio stream. + stream_output (pyaudio.Stream): Output audio stream. + p (pyaudio.PyAudio): PyAudio instance. + """ + stream_input.stop_stream() + stream_input.close() + stream_output.stop_stream() + stream_output.close() + p.terminate() + + +# Implementation for type 1 +def type_1_main( + args: argparse.Namespace, + device: str, + model: nn.Module, + config: Any, + chunk_size: int, +) -> None: + """ + Main function for type 1 script, which records audio, processes it, and saves the output. + + Args: + args (argparse.Namespace): Parsed arguments. + device (str): Device to run the model on. + model (nn.Module): The model to use for audio separation. + config (Any): Configuration used for demixing. + chunk_size (int): Size of each audio chunk. + """ + + def record_audio(stream_input: pyaudio.Stream) -> bytes: + """ + Record audio from the input stream until the user presses 'S' to stop. + + Args: + stream_input (pyaudio.Stream): The input audio stream. + + Returns: + bytes: Recorded audio data. + """ + stop_flag = threading.Event() + + def check_stop_key() -> None: + while not stop_flag.is_set(): + if keyboard.is_pressed("s"): + print("Recording stopped...") + stop_flag.set() + + keyboard_thread = threading.Thread(target=check_stop_key, daemon=True) + keyboard_thread.start() + + audio_queue = queue.Queue() + print("Recording started... Press 'S' to stop.") + + while not stop_flag.is_set(): + try: + audio_data = stream_input.read(44100, exception_on_overflow=False) + audio_queue.put(audio_data) + except IOError as e: + print("Error reading audio data:", e) + break + + audio_data = b"".join(list(audio_queue.queue)) + with audio_queue.mutex: + audio_queue.queue.clear() + return audio_data + + def process_audio( + audio_data: bytes, + model: nn.Module, + args: argparse.Namespace, + config: Any, + device: str, + stream_output: pyaudio.Stream, + ) -> None: + """ + Process the recorded audio using the model and save the output. + + Args: + audio_data (bytes): The recorded audio data. + model (nn.Module): The model to use for audio separation. + args (argparse.Namespace): Parsed arguments. + config (Any): Configuration for the model. + device (str): Device to run the model on. + stream_output (pyaudio.Stream): The output audio stream. + """ + audio_array = np.frombuffer(audio_data, dtype=np.float32) + audio_array = np.expand_dims(audio_array, axis=0) + audio_array = np.concatenate([audio_array, audio_array], axis=0) + + output_path = os.path.abspath(os.path.join(args.out_dir, "raw_audio.wav")) + sf.write(output_path, audio_array.T, RATE, "FLOAT") + + waveforms_orig = demix( + config, model, audio_array, device, model_type=args.model_type + ) + if args.use_tta: + waveforms_orig = apply_tta( + config, model, audio_array, waveforms_orig, device, args.model_type + ) + waveforms_orig = waveforms_orig["vocals"] + + output_path = os.path.abspath( + os.path.join(args.out_dir, f"{args.out_name}.wav") + ) + sf.write(output_path, waveforms_orig.T, RATE, "FLOAT") + print(f"Processing completed. Output saved to {output_path}.") + stream_output.write(waveforms_orig.T.tobytes()) + + os.makedirs(args.out_dir, exist_ok=True) + p, stream_input, stream_output = initialize_audio_streams(chunk_size) + audio_data = record_audio(stream_input) + process_audio(audio_data, model, args, config, device, stream_output) + close_audio_streams(stream_input, stream_output, p) + print("End.") + + +# Implementation for type 2 +async def type_2_main( + args: argparse.Namespace, + device: str, + model: nn.Module, + config: Any, + chunk_size: int, +) -> None: + """ + Main function for type 2 script, which continuously records audio, processes it, + and streams the output in real-time. + + Args: + args (argparse.Namespace): Parsed arguments. + device (str): Device to run the model on. + model (nn.Module): The model to use for audio separation. + config (Any): Configuration used for demixing. + chunk_size (int): Size of each audio chunk. + """ + + def fake( + config: Any, model: nn.Module, device: str, args: argparse.Namespace + ) -> None: + """ + Fake initialization function to simulate processing and model loading. + + Args: + config (Any): Configuration used for demixing. + model (nn.Module): The model to use for audio separation. + device (str): Device to run the model on. + args (argparse.Namespace): Parsed arguments. + """ + print("Please wait...") + audio_array = np.random.randn(chunk_size) + audio_array = np.expand_dims(audio_array, axis=0) + audio_array = np.concatenate([audio_array, audio_array], axis=0) + demix(config, model, audio_array, device, model_type=args.model_type) + print("Model initialized. Speak...") + + async def output_to_speakers( + waveforms_orig: np.ndarray, stream_output: pyaudio.Stream + ) -> None: + """ + Stream the output audio to the speakers. + + Args: + waveforms_orig (Dict[str, np.ndarray]): Dictionary containing the separated audio sources. + stream_output (pyaudio.Stream): The output audio stream. + """ + stream_output.write(waveforms_orig.T.tobytes()) + + async def record_and_process_audio( + stream_input: pyaudio.Stream, + stream_output: pyaudio.Stream, + model: nn.Module, + args: argparse.Namespace, + config: Any, + device: str, + ) -> None: + """ + Continuously record audio from the input stream, process it using the model, + and stream the output to the speakers in real-time. + + Args: + stream_input (pyaudio.Stream): The input audio stream. + stream_output (pyaudio.Stream): The output audio stream. + model (nn.Module): The model to use for audio separation. + args (argparse.Namespace): Parsed arguments. + config (Any): Configuration used for demixing. + device (str): Device to run the model on. + """ + while True: + try: + audio_data = stream_input.read(chunk_size, exception_on_overflow=False) + start_time = time.time() + audio_array = np.frombuffer(audio_data, dtype=np.float32) + audio_array = np.expand_dims(audio_array, axis=0) + audio_array = np.concatenate([audio_array, audio_array], axis=0) + + waveforms_orig = demix( + config, model, audio_array, device, model_type=args.model_type + ) + if args.use_tta: + waveforms_orig = apply_tta( + config, + model, + audio_array, + waveforms_orig, + device, + args.model_type, + ) + + waveforms_orig = waveforms_orig["vocals"] + if time.time() - start_time > 0.7: + print(f"Elapsed time: {time.time() - start_time:.2f}") + await output_to_speakers(waveforms_orig, stream_output) + except IOError as e: + print("Error reading audio data:", e) + break + + fake(config, model, device, args) + p, stream_input, stream_output = initialize_audio_streams(chunk_size) + await record_and_process_audio( + stream_input, stream_output, model, args, config, device + ) + close_audio_streams(stream_input, stream_output, p) + + +if __name__ == "__main__": + args = parse_args(None) + args.config_path = str( + os.path.abspath(os.path.join(base_dir, "..", args.config_path)) + ) + args.start_check_point = str( + os.path.abspath(os.path.join(base_dir, "..", args.start_check_point)) + ) + args.out_dir = str(os.path.abspath(os.path.join(base_dir, "..", args.out_dir))) + + device = initialize_device(args) + model, config = load_model(args, device) + + if args.model_type == "htdemucs": + chunk_size = config.training.samplerate * config.training.segment + else: + chunk_size = config.audio.chunk_size + + if args.type == 1: + type_1_main(args, device, model, config, chunk_size) + elif args.type == 2: + try: + asyncio.run(type_2_main(args, device, model, config, chunk_size)) + finally: + print("End!") + exit(0) diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/trim.py b/src/third_party/MusicSourceSeparationTraining/scripts/trim.py new file mode 100644 index 0000000000000000000000000000000000000000..77b8754938e638da938ca270f8b88602baca014d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/trim.py @@ -0,0 +1,108 @@ +import argparse +import os +import time +from typing import Dict, Union + +import soundfile as sf + + +def parse_args(dict_args: Union[Dict, None]) -> argparse.Namespace: + """ + Parse command-line arguments for configuring the model, dataset, and training parameters. + + Args: + dict_args: Dict of command-line arguments. If None, arguments will be parsed from sys.argv. + + Returns: + Namespace object containing parsed arguments and their values. + """ + + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_directory", type=str, help="Path to the input directory." + ) + parser.add_argument( + "--output_directory", type=str, help="Path to the output directory." + ) + parser.add_argument("--start_sec", type=float, default=20.0) + parser.add_argument("--end_sec", type=float, default=30.0) + parser.add_argument("--codec", type=str, default="wav") + parser.add_argument( + "--max_folders", + type=int, + default=float("inf"), + help="Maximum number of folders to process.", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + if not args.output_directory: + original_dir = os.path.dirname(args.input_directory) + tests_dir = os.path.join("tests", original_dir) + os.makedirs(tests_dir, exist_ok=True) + args.output_directory = os.path.join( + tests_dir, os.path.basename(args.input_directory) + ) + + return args + + +def trim_wav( + input_file: str, output_file: str, start_sec: float, end_sec: float, codec: str +): + data, samplerate = sf.read(input_file) + start_sample = int(start_sec * samplerate) + end_sample = int(end_sec * samplerate) + trimmed_data = data[start_sample:end_sample] + sf.write(output_file, trimmed_data, samplerate, format=codec) + + +def trim_directory(args): + args = parse_args(args) + input_directory = args.input_directory + output_directory = args.output_directory + start_sec = args.start_sec + end_sec = args.end_sec + codec = args.codec + max_folder = args.max_folders + + folder_count = 0 + start_time = time.time() + + for root, dirs, files in os.walk(input_directory): + if folder_count >= max_folder: + break + if os.path.relpath(root, input_directory) != ".": + print(f"Processing folder: {os.path.relpath(root, input_directory)}") + + relative_path = os.path.relpath(root, input_directory) + target_folder = os.path.join(output_directory, relative_path) + os.makedirs(target_folder, exist_ok=True) + + for filename in files: + if filename.endswith(f".{codec}"): + input_file = os.path.join(root, filename) + output_file = os.path.join( + target_folder, filename.replace(f".{codec}", f".{codec}") + ) + try: + trim_wav(input_file, output_file, start_sec, end_sec, codec) + except Exception as e: + print(f"Error processing {filename} in folder {root}: {e}") + + folder_count += 1 + + end_time = time.time() + total_time = end_time - start_time + print(f"Processing complete. Total time: {total_time:.2f} seconds.") + return output_directory + + +if __name__ == "__main__": + trim_directory(None) diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/valid_to_inference.py b/src/third_party/MusicSourceSeparationTraining/scripts/valid_to_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..5fd798d2be9eb2260adf65ef204aeae5dc4d4db7 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/valid_to_inference.py @@ -0,0 +1,93 @@ +import argparse +import os +import shutil +from typing import Dict, Optional, Union + + +def parse_args(dict_args: Union[Dict, None]) -> argparse.Namespace: + """ + Parse command-line arguments for configuring the model, dataset, and training parameters. + + Args: + dict_args: Dict of command-line arguments. If None, arguments will be parsed from sys.argv. + + Returns: + Namespace object containing parsed arguments and their values. + """ + + parser = argparse.ArgumentParser( + description="Copy mixture files from VALID_PATH to INFERENCE_DIR" + ) + parser.add_argument("--valid_path", type=str, help="Directory with valid tracks") + parser.add_argument( + "--inference_dir", type=str, help="Directory to save inference tracks" + ) + parser.add_argument( + "--mixture_name", + type=str, + default="mixture.wav", + help="Name of mixture tracks (default: 'mixture.wav')", + ) + parser.add_argument( + "--max_mixtures", + type=int, + default=float("inf"), + help="Maximum number of mixtures to process.", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + return args + + +def copying_files(args: Optional[argparse.Namespace] = None) -> None: + """ + Main function to copy mixture files from valid directory to inference directory. + + Parameters: + ---------- + args : Optional[argparse.Namespace] + The parsed arguments containing valid_path, inference_dir, and mixture_name. + """ + args = parse_args(args) + + valid_path = args.valid_path + inference_dir = args.inference_dir + mixture_name = args.mixture_name + max_mixtures = args.max_mixtures + # Create the inference directory if it doesn't exist + os.makedirs(inference_dir, exist_ok=True) + mixture_count = 0 + # Walk through the valid directory to find and copy mixture files + for root, _, files in os.walk(valid_path): + if mixture_count >= max_mixtures: + break + for file in files: + if file == mixture_name: + mixture_count += 1 + # Full path to the valid file + source_path = os.path.join(root, file) + + # Track ID from the parent directory name + track_id = os.path.basename(os.path.dirname(source_path)) + + # Define target file path + target_filename = f"{track_id}.wav" + target_path = os.path.join(inference_dir, target_filename) + + # Copy the file to the inference directory + shutil.copy2(source_path, target_path) + + print(f"Has copied: {source_path} -> {target_path}") + + print("Copying ends.") + + +if __name__ == "__main__": + copying_files(None) diff --git a/src/third_party/MusicSourceSeparationTraining/scripts/watch_all_metrics.py b/src/third_party/MusicSourceSeparationTraining/scripts/watch_all_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..0d208091877bb3e90e4121d32347f70e337ce740 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/scripts/watch_all_metrics.py @@ -0,0 +1,69 @@ +""" +Script to print metrics for checkpoint file of new format +""" + +import argparse + +import torch + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--start_check_point", + type=str, + default="", + help="Initial checkpoint to start training", + ) + args = parser.parse_args() + return args + + +def main(): + args = parse_args() + checkpoint = torch.load( + args.start_check_point, weights_only=False, map_location="cpu" + ) + + all_metrics = checkpoint["all_metrics"] + + from typing import Any, Dict, List + + import numpy as np + + def _fmt_list(xs: List[float], n: int, p: int) -> str: + if not xs: + return "[]" + xs = [float(x) for x in xs] + head = ", ".join(f"{v:.{p}f}" for v in xs[:n]) + if len(xs) > n: + return f"[{head}, … {len(xs) - n} more]" + return f"[{head}]" + + def pretty_metrics( + all_time_all_metrics: Dict[str, Any], *, precision: int = 4, show_items: int = 6 + ) -> str: + lines = [] + for epoch_key in sorted( + all_time_all_metrics.keys(), key=lambda k: int(k.split("_")[-1]) + ): + m = all_time_all_metrics[epoch_key] + lines.append(f"\n=== {epoch_key} ===") + for metric_name, per_instr in m.items(): + lines.append(f"{metric_name}:") + for instr, values in per_instr.items(): + arr = np.array(values, dtype=float) + mean = np.mean(arr) if arr.size else float("nan") + std = np.std(arr) if arr.size else float("nan") + preview = _fmt_list(values, show_items, precision) + lines.append( + f" {instr:>10s}: mean={mean:.{precision}f} std={std:.{precision}f} " + f"n={arr.size} values={preview}" + ) + return "\n".join(lines) + + print(pretty_metrics(all_metrics, precision=4, show_items=6)) + + +if __name__ == "__main__": + main() diff --git a/src/third_party/MusicSourceSeparationTraining/seprate_test.py b/src/third_party/MusicSourceSeparationTraining/seprate_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0d1c37c2e7cc2ed7e7e3526a840e235c8a02ae05 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/seprate_test.py @@ -0,0 +1,380 @@ +# coding: utf-8 +""" +MelBand RoFormer 音频分离器 —— 类封装版本 +基于 ZFTurbo 的 Music-Source-Separation-Training 推理脚本改写。 + +用法示例: + from separator import MelBandSeparator + + sep = MelBandSeparator( + model_type="mel_band_roformer", + config_path="ckpts/config_vocals_mel_band_roformer_kj.yaml", + checkpoint_path="ckpts/MelBandRoformer.ckpt", + device="cuda:0", # 或 "cpu" / "mps" + ) + + results = sep.separate("song.wav", output_dir="output/") + # results: dict[str, str] —— {instrument_name: output_file_path, ...} +""" + +from __future__ import annotations + +import os +import sys +import time +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +import torchaudio + +# --------------------------------------------------------------------------- +# 让 utils 模块可以被正确导入（兼容嵌入式 Python 等场景） +# --------------------------------------------------------------------------- +_CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) +if _CURRENT_DIR not in sys.path: + sys.path.append(_CURRENT_DIR) + +import warnings + +from utils.audio_utils import denormalize_audio, draw_spectrogram, normalize_audio +from utils.model_utils import ( + apply_tta, + demix, + load_start_checkpoint, + prefer_target_instrument, +) +from utils.settings import get_model_from_config + +warnings.filterwarnings("ignore") + + +# --------------------------------------------------------------------------- +# 数据类：分离结果 +# --------------------------------------------------------------------------- +@dataclass +class SeparationResult: + """单个音频文件的分离结果。""" + + instrument: str + audio: np.ndarray # shape: (channels, samples) + sample_rate: int + output_path: Optional[str] = None + + +# --------------------------------------------------------------------------- +# 核心类 +# --------------------------------------------------------------------------- +class MelBandSeparator: + """ + 基于 MelBand RoFormer 的音频源分离器。 + + Parameters + ---------- + model_type : str + 模型类型，例如 ``"mel_band_roformer"``、``"bs_roformer"``、``"mdx23c"`` 等。 + config_path : str + 模型配置文件 (.yaml) 路径。 + checkpoint_path : str + 模型权重文件 (.ckpt) 路径。 + device : str | torch.device + 推理设备，例如 ``"cuda:0"``、``"cpu"``、``"mps"``。 + 默认为 ``"auto"``，会自动选择可用的 GPU > MPS > CPU。 + device_ids : list[int] | None + 多 GPU DataParallel 时使用的 GPU id 列表。 + 仅当 device 为 cuda 且提供多个 id 时生效。 + use_tta : bool + 是否启用 Test-Time Augmentation（推理时数据增强），可略微提升质量但更慢。 + extract_instrumental : bool + 是否额外提取伴奏轨（原始混音 - 人声）。 + pcm_type : str + 输出音频的 PCM 子类型，例如 ``"FLOAT"``、``"PCM_16"``、``"PCM_24"``。 + """ + + def __init__( + self, + args, + model_type: str, + config_path: str, + checkpoint_path: str, + device: str = "auto", + device_ids: list[int] | None = None, + use_tta: bool = False, + extract_instrumental: bool = False, + pcm_type: str = "FLOAT", + ) -> None: + self.model_type = model_type + self.config_path = config_path + self.checkpoint_path = checkpoint_path + self.use_tta = use_tta + self.extract_instrumental = extract_instrumental + self.pcm_type = pcm_type + + # ---- 选择设备 ---- + self.device = self._resolve_device(device) + self.device_ids = device_ids + print(f"[MelBandSeparator] Using device: {self.device}") + + # ---- 加载模型 ---- + t0 = time.time() + torch.backends.cudnn.benchmark = True + + self.model, self.config = get_model_from_config(model_type, config_path) + + # 覆盖 model_type（部分 config 里会自带） + if "model_type" in self.config.training: + self.model_type = self.config.training.model_type + + # 加载权重 + checkpoint = torch.load(checkpoint_path, weights_only=False, map_location="cpu") + load_start_checkpoint(args, self.model, checkpoint, type_="inference") + + # 多 GPU + if ( + device_ids is not None + and len(device_ids) > 1 + and "cuda" in str(self.device) + ): + self.model = nn.DataParallel(self.model, device_ids=device_ids) + + self.model = self.model.to(self.device) + self.model.eval() + + self.sample_rate: int = getattr(self.config.audio, "sample_rate", 44100) + self.instruments: list[str] = prefer_target_instrument(self.config)[:] + + print(f"[MelBandSeparator] Instruments: {self.instruments}") + print(f"[MelBandSeparator] Model loaded in {time.time() - t0:.2f}s") + + # ------------------------------------------------------------------ + # 公共方法 + # ------------------------------------------------------------------ + def separate( + self, + mix_audio, + mix_audio_sr, + output_dir: str | None = None, + draw_spectro: bool = False, + ) -> list[SeparationResult]: + """ + 对单个音频文件进行源分离。 + + Parameters + ---------- + audio_path : str + 输入音频文件路径（支持 wav / flac / mp3 等 torchaudio 能读取的格式）。 + output_dir : str | None + 输出目录。为 ``None`` 时不写文件，仅返回内存中的结果。 + draw_spectro : bool + 是否保存频谱图（需要 output_dir 不为 None）。 + + Returns + ------- + list[SeparationResult] + 每个乐器/人声轨对应一个 ``SeparationResult``。 + """ + # ---- 1. 使用 torchaudio 读取音频 ---- + mix = mix_audio + sr = mix_audio_sr + # 重采样到模型要求的采样率 + if sr != self.sample_rate: + resampler = torchaudio.transforms.Resample( + orig_freq=sr, new_freq=self.sample_rate + ) + mix = resampler(mix) + sr = self.sample_rate + + # 转 numpy: (channels, samples) + mix: np.ndarray = mix.numpy() + + # 单声道 → 立体声（如果模型需要） + if mix.shape[0] == 1: + num_channels = getattr(self.config.audio, "num_channels", 1) + if num_channels == 2: + print("[MelBandSeparator] Converting mono to stereo...") + mix = np.concatenate([mix, mix], axis=0) + + mix_orig = mix.copy() + + # ---- 2. 归一化 ---- + norm_params = None + if getattr(self.config.inference, "normalize", False): + mix, norm_params = normalize_audio(mix) + + # ---- 3. 分离 ---- + waveforms = demix( + self.config, + self.model, + mix, + self.device, + model_type=self.model_type, + pbar=True, + ) + + # ---- 4. TTA ---- + if self.use_tta: + waveforms = apply_tta( + self.config, + self.model, + mix, + waveforms, + self.device, + self.model_type, + ) + + # ---- 5. 伴奏提取 ---- + instruments = self.instruments[:] + if self.extract_instrumental: + instr_key = "vocals" if "vocals" in instruments else instruments[0] + waveforms["instrumental"] = mix_orig - waveforms[instr_key] + if "instrumental" not in instruments: + instruments.append("instrumental") + + # ---- 6. 反归一化 & 收集结果 ---- + results: list[SeparationResult] = [] + file_stem = 111 + + for instr in instruments: + estimates = waveforms[instr] + if norm_params is not None: + estimates = denormalize_audio(estimates, norm_params) + + result = SeparationResult( + instrument=instr, + audio=estimates, + sample_rate=sr, + ) + + # ---- 写文件 ---- + if output_dir is not None: + os.makedirs(output_dir, exist_ok=True) + + peak = float(np.abs(estimates).max()) + codec = "flac" if (peak <= 1.0 and self.pcm_type != "FLOAT") else "wav" + + out_path = os.path.join(output_dir, f"{file_stem}_{instr}.{codec}") + sf.write(out_path, estimates.T, sr, subtype=self.pcm_type) + result.output_path = out_path + print(f"[MelBandSeparator] Saved: {out_path}") + + # out_path = os.path.join(output_dir, f"_inst.wav") + # sf.write(out_path, estimates.T, sr, subtype=self.pcm_type) + # result.output_path = out_path + # print(f"[MelBandSeparator] Saved: {out_path}") + + if draw_spectro: + img_path = os.path.join(output_dir, f"{file_stem}_{instr}.jpg") + draw_spectrogram(estimates.T, sr, 1, img_path) + print(f"[MelBandSeparator] Spectrogram: {img_path}") + + results.append(result) + + return results + + # ------------------------------------------------------------------ + # 内部方法 + # ------------------------------------------------------------------ + @staticmethod + def _resolve_device(device: str) -> torch.device: + """自动选择设备。""" + if device != "auto": + return torch.device(device) + + if torch.cuda.is_available(): + print("[MelBandSeparator] CUDA detected.") + return torch.device("cuda:0") + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + print("[MelBandSeparator] MPS detected.") + return torch.device("mps") + else: + return torch.device("cpu") + + +# --------------------------------------------------------------------------- +# CLI 入口（可选） +# --------------------------------------------------------------------------- +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="MelBand RoFormer Separator") + parser.add_argument("--model_type", default="mel_band_roformer") + parser.add_argument("--config_path", required=True, help="Path to config YAML") + parser.add_argument( + "--start_check_point", required=True, help="Path to model checkpoint" + ) + parser.add_argument( + "--device", default="auto", help="Device: auto / cpu / cuda:0 / mps" + ) + parser.add_argument("--output_dir", default="./output", help="Output directory") + parser.add_argument("--use_tta", action="store_true", help="Enable TTA") + parser.add_argument("--extract_instrumental", action="store_true") + parser.add_argument( + "--pcm_type", default="FLOAT", choices=["FLOAT", "PCM_16", "PCM_24"] + ) + # parser.add_argument("--model_type", type=str, default='mdx23c', + # help="One of bandit, bandit_v2, bs_roformer, htdemucs, mdx23c, mel_band_roformer," + # " scnet, scnet_unofficial, segm_models, swin_upernet, torchseg") + # parser.add_argument("--config_path", type=str, help="path to config file") + # parser.add_argument("--start_check_point", type=str, default='', help="Initial checkpoint to valid weights") + # parser.add_argument("--input_folder", type=str, help="folder with mixtures to process") + # parser.add_argument("--store_dir", type=str, default="", help="path to store results as wav file") + # parser.add_argument("--draw_spectro", type=float, default=0, + # help="Code will generate spectrograms for resulted stems." + # " Value defines for how many seconds os track spectrogram will be generated.") + # parser.add_argument("--device_ids", nargs='+', type=int, default=0, help='list of gpu ids') + # parser.add_argument( + # "--extract_instrumental", + # action="store_true", + # help="invert vocals to get instrumental if provided", + # ) + # parser.add_argument("--disable_detailed_pbar", action='store_true', help="disable detailed progress bar") + # parser.add_argument("--force_cpu", action='store_true', help="Force the use of CPU even if CUDA is available") + # parser.add_argument("--flac_file", action='store_true', help="Output flac file instead of wav") + # parser.add_argument("--pcm_type", type=str, choices=['PCM_16', 'PCM_24', 'FLOAT'], default='FLOAT', + # help="PCM type for FLAC files (PCM_16 or PCM_24)") + # parser.add_argument("--use_tta", action='store_true', + # help="Flag adds test time augmentation during inference (polarity and channel inverse)." + # "While this triples the runtime, it reduces noise and slightly improves prediction quality.") + # parser.add_argument("--lora_checkpoint_peft", type=str, default='', help="Initial checkpoint to LoRA weights") + # parser.add_argument("--filename_template", type=str, default='{file_name}/{instr}', + # help="Output filename template, without extension, using '/' for subdirectories. Default: '{file_name}/{instr}'") + parser.add_argument( + "--lora_checkpoint_loralib", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + parser.add_argument("--draw_spectro", action="store_true") + + cli_args = parser.parse_args() + + sep = MelBandSeparator( + cli_args, + model_type=cli_args.model_type, + config_path=cli_args.config_path, + checkpoint_path=cli_args.start_check_point, + device=cli_args.device, + use_tta=cli_args.use_tta, + extract_instrumental=cli_args.extract_instrumental, + pcm_type=cli_args.pcm_type, + ) + + mix, sr = torchaudio.load( + "/user-fs/chenzihao/aslp_music/haochunbo/final/张韶轩-隐形的翅膀.mp3" + ) # (channels, samples) + + results = sep.separate( + mix_audio=mix, + mix_audio_sr=sr, + output_dir="/user-fs/chenzihao/aslp_music/haochunbo/final", + draw_spectro=False, + ) + + for r in results: + print(f" {r.instrument}: {r.output_path or '(in memory only)'}") +""" +PYTHONPATH=. python /user-fs/chenzihao/aslp_music/haochunbo/final/YingMusic-Singer/src/third_party/Music-Source-Separation-Training/seprate_test.py --config_path ckpts/config_vocals_mel_band_roformer_kj.yaml --start_check_point ckpts/MelBandRoformer.ckpt --extract_instrumental +""" diff --git a/src/third_party/MusicSourceSeparationTraining/tests/admin_test.py b/src/third_party/MusicSourceSeparationTraining/tests/admin_test.py new file mode 100644 index 0000000000000000000000000000000000000000..fb61176dad73a5ad1b9f4fd3560f7b1304e2ff90 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/tests/admin_test.py @@ -0,0 +1,192 @@ +import os +from pathlib import Path +from typing import Dict, List + +import numpy as np +import soundfile as sf +from scripts.redact_config import redact_config +from test import test_settings +from utils.settings import load_config + +MODEL_CONFIGS = { + "config_apollo.yaml": {"model_type": "apollo"}, + "config_dnr_bandit_bsrnn_multi_mus64.yaml": {"model_type": "bandit"}, + "config_dnr_bandit_v2_mus64.yaml": {"model_type": "bandit_v2"}, + "config_drumsep.yaml": {"model_type": "htdemucs"}, + "config_htdemucs_6stems.yaml": {"model_type": "htdemucs"}, + "config_musdb18_bs_roformer.yaml": {"model_type": "bs_roformer"}, + "config_musdb18_demucs3_mmi.yaml": {"model_type": "htdemucs"}, + "config_musdb18_htdemucs.yaml": {"model_type": "htdemucs"}, + "config_musdb18_mdx23c.yaml": {"model_type": "mdx23c"}, + "config_musdb18_mel_band_roformer.yaml": {"model_type": "mel_band_roformer"}, + "config_musdb18_mel_band_roformer_all_stems.yaml": { + "model_type": "mel_band_roformer" + }, + "config_musdb18_scnet.yaml": {"model_type": "scnet"}, + "config_musdb18_scnet_large.yaml": {"model_type": "scnet"}, + # 'config_musdb18_scnet_large_starrytong.yaml': {'model_type': 'scnet'}, + "config_vocals_bandit_bsrnn_multi_mus64.yaml": {"model_type": "bandit"}, + "config_vocals_bs_roformer.yaml": {"model_type": "bs_roformer"}, + "config_vocals_htdemucs.yaml": {"model_type": "htdemucs"}, + "config_vocals_mdx23c.yaml": {"model_type": "mdx23c"}, + "config_vocals_mel_band_roformer.yaml": {"model_type": "mel_band_roformer"}, + "config_vocals_scnet.yaml": {"model_type": "scnet"}, + "config_vocals_scnet_large.yaml": {"model_type": "scnet"}, + "config_vocals_scnet_unofficial.yaml": {"model_type": "scnet_unofficial"}, + "config_vocals_segm_models.yaml": {"model_type": "segm_models"}, + # 'config_vocals_swin_upernet.yaml': {'model_type': 'swin_upernet'}, + # 'config_musdb18_torchseg.yaml': {'model_type': 'torchseg'}, + # 'config_musdb18_segm_models.yaml': {'model_type': 'segm_models'}, + # 'config_musdb18_bs_mamba2.yaml': {'model_type': 'bs_mamba2'}, + # 'config_vocals_bs_mamba2.yaml': {'model_type': 'bs_mamba2'}, + # 'config_vocals_torchseg.yaml': {'model_type': 'torchseg'} +} + + +# Folders for tests +ROOT_DIR = Path(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +CONFIGS_DIR = ROOT_DIR / "configs/" +TEST_DIR = ROOT_DIR / "tests_cache/" +TRAIN_DIR = TEST_DIR / "train_tracks/" +VALID_DIR = TEST_DIR / "valid_tracks/" + + +def create_dummy_tracks( + directory: Path, + num_tracks: int, + instruments: List[str], + duration: float = 5.0, + sample_rate: int = 44100, +) -> None: + """ + Generates random audio tracks for stems in two subdirectories within the specified directory. + + Parameters: + ---------- + directory : Path + Path to the directory where the tracks will be saved. + num_tracks : int + Number of tracks to generate in each folder. + instruments : List[str] + List of instrument names (stems) to create. + duration : float, optional + Duration of each track in seconds. Default is 5.0. + sample_rate : int, optional + Sampling rate of the generated audio. Default is 44100 Hz. + + Returns: + ------- + None + """ + + os.makedirs(directory, exist_ok=True) + + for folder_name in [str(i) for i in range(1, num_tracks + 1)]: + folder_path = directory / folder_name + os.makedirs(folder_path, exist_ok=True) + for instrument in instruments: + # Generate random noice for each track + samples = int(duration * sample_rate) + track = np.random.uniform(-1.0, 1.0, (2, samples)).astype(np.float32) + file_path = folder_path / f"{instrument}.wav" + sf.write(file_path, track.T, sample_rate) + + +def cleanup_test_tracks() -> None: + """ + Removes all cached test tracks. + + This function deletes the entire directory specified by the global `TEST_DIR` variable + if it exists. + + Returns: + ------- + None + This function does not return a value. It performs cleanup of test data. + """ + + +def modify_configs() -> Dict[str, Path]: + """ + Updates configuration files in the `configs` directory for use with test data. + + This function processes configuration files defined in the global `MODEL_CONFIGS` dictionary, + modifies them to be compatible with test scenarios, and saves the updated configurations + in a test-specific directory. + + Returns: + ------- + Dict[str, Path] + A dictionary where the keys are the original configuration file names, and the values + are the paths to the updated configuration files. + """ + config_dir = CONFIGS_DIR + updated_configs = {} + for config, args in MODEL_CONFIGS.items(): + model_type = args["model_type"] + config_path = config_dir / config + updated_config_path = redact_config( + { + "orig_config": str(config_path), + "model_type": model_type, + "new_config": str(TEST_DIR / "configs" / config), + } + ) + updated_configs[config] = updated_config_path + return updated_configs + + +def run_tests() -> None: + """ + Executes validation tests for all configurations. + + This function updates configurations, generates random dummy data for testing, + and runs a series of tests (training, validation, and inference checks) for each + model configuration specified in the global `MODEL_CONFIGS` dictionary. + + Returns: + ------- + None + """ + + updated_configs = modify_configs() + + # For every config + for config, args in MODEL_CONFIGS.items(): + model_type = args["model_type"] + cfg = load_config( + model_type=model_type, config_path=TEST_DIR / "configs" / config + ) + # Random tracks + create_dummy_tracks( + TRAIN_DIR, instruments=cfg.training.instruments + ["mixture"], num_tracks=2 + ) + create_dummy_tracks( + VALID_DIR, instruments=cfg.training.instruments + ["mixture"], num_tracks=2 + ) + + print(f"\nRunning tests for model: {model_type} (config: {config})") + + test_args = { + "check_train": False, + "check_valid": True, + "check_inference": True, + "config_path": updated_configs[config], + "data_path": str(TRAIN_DIR), + "valid_path": str(VALID_DIR), + "results_path": str(TEST_DIR / "results" / model_type), + "store_dir": str(TEST_DIR / "inference_results" / model_type), + "metrics": ["sdr", "si_sdr", "l1_freq"], + } + + test_args.update(args) + + test_settings(test_args, "admin") + print(f"Tests for model {model_type} completed successfully.") + + # Remove test_cache + cleanup_test_tracks() + + +if __name__ == "__main__": + run_tests() diff --git a/src/third_party/MusicSourceSeparationTraining/tests/test.py b/src/third_party/MusicSourceSeparationTraining/tests/test.py new file mode 100644 index 0000000000000000000000000000000000000000..c426fe481991254a5645fb9a89524a9e1cadc0ae --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/tests/test.py @@ -0,0 +1,317 @@ +import argparse +import os +import sys + +# Добавляем корень репозитория в системный путь +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from inference import proc_folder +from scripts.redact_config import redact_config +from scripts.trim import trim_directory +from scripts.valid_to_inference import copying_files +from train import train_model +from valid import check_validation + +base_args = { + "device_ids": "0", + "model_type": "", + "start_check_point": "", + "config_path": "", + "data_path": "", + "valid_path": "", + "results_path": "tests/train_results", + "store_dir": "tests/valid_inference_result", + "input_folder": "", + "metrics": [ + "neg_log_wmse", + "l1_freq", + "si_sdr", + "sdr", + "aura_stft", + "aura_mrstft", + "bleedless", + "fullness", + ], + "max_folders": 2, +} + + +def parse_args(dict_args): + parser = argparse.ArgumentParser() + parser.add_argument("--check_train", action="store_true", help="Check train or not") + parser.add_argument("--check_valid", action="store_true", help="Check train or not") + parser.add_argument( + "--check_inference", action="store_true", help="Check train or not" + ) + parser.add_argument( + "--device_ids", type=str, help="Device IDs for training/inference" + ) + parser.add_argument("--model_type", type=str, help="Model type") + parser.add_argument( + "--start_check_point", type=str, help="Path to the checkpoint to start from" + ) + parser.add_argument( + "--config_path", type=str, help="Path to the configuration file" + ) + parser.add_argument("--data_path", type=str, help="Path to the training data") + parser.add_argument("--valid_path", type=str, help="Path to the validation data") + parser.add_argument( + "--results_path", type=str, help="Path to save training results" + ) + parser.add_argument( + "--store_dir", type=str, help="Path to store validation/inference results" + ) + parser.add_argument( + "--input_folder", type=str, help="Path to the input folder for inference" + ) + parser.add_argument("--metrics", nargs="+", help="List of metrics to evaluate") + parser.add_argument( + "--max_folders", type=str, help="Maximum number of folders to process" + ) + parser.add_argument( + "--dataset_type", + type=int, + default=1, + help="Dataset type. Must be one of: 1, 2, 3 or 4.", + ) + parser.add_argument( + "--num_workers", type=int, default=0, help="dataloader num_workers" + ) + parser.add_argument( + "--pin_memory", action="store_true", help="dataloader pin_memory" + ) + parser.add_argument("--seed", type=int, default=0, help="random seed") + parser.add_argument( + "--use_multistft_loss", + action="store_true", + help="Use MultiSTFT Loss (from auraloss package)", + ) + parser.add_argument( + "--use_mse_loss", action="store_true", help="Use default MSE loss" + ) + parser.add_argument("--use_l1_loss", action="store_true", help="Use L1 loss") + parser.add_argument("--wandb_key", type=str, default="", help="wandb API Key") + parser.add_argument( + "--pre_valid", action="store_true", help="Run validation before training" + ) + parser.add_argument( + "--metric_for_scheduler", + default="sdr", + choices=[ + "sdr", + "l1_freq", + "si_sdr", + "neg_log_wmse", + "aura_stft", + "aura_mrstft", + "bleedless", + "fullness", + ], + help="Metric which will be used for scheduler.", + ) + parser.add_argument("--train_lora", action="store_true", help="Train with LoRA") + parser.add_argument( + "--lora_checkpoint", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + parser.add_argument( + "--extension", type=str, default="wav", help="Choose extension for validation" + ) + parser.add_argument( + "--use_tta", + action="store_true", + help="Flag adds test time augmentation during inference (polarity and channel inverse)." + " While this triples the runtime, it reduces noise and slightly improves prediction quality.", + ) + parser.add_argument( + "--extract_instrumental", + action="store_true", + help="invert vocals to get instrumental if provided", + ) + parser.add_argument( + "--disable_detailed_pbar", + action="store_true", + help="disable detailed progress bar", + ) + parser.add_argument( + "--force_cpu", + action="store_true", + help="Force the use of CPU even if CUDA is available", + ) + parser.add_argument( + "--flac_file", action="store_true", help="Output flac file instead of wav" + ) + parser.add_argument( + "--pcm_type", + type=str, + choices=["PCM_16", "PCM_24"], + default="PCM_24", + help="PCM type for FLAC files (PCM_16 or PCM_24)", + ) + parser.add_argument( + "--draw_spectro", + type=float, + default=0, + help="If --store_dir is set then code will generate spectrograms for resulted stems as well." + " Value defines for how many seconds os track spectrogram will be generated.", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + return args + + +def test_settings(dict_args, test_type): + # Parse from cmd + cli_args = parse_args(dict_args) + + # If args from cmd, add or replace in base_args + for key, value in vars(cli_args).items(): + if value is not None: + base_args[key] = value + + if test_type == "user": + # Check required arguments + missing_args = [ + arg + for arg in [ + "model_type", + "config_path", + "start_check_point", + "data_path", + "valid_path", + ] + if not base_args[arg] + ] + if missing_args: + missing_args_str = ", ".join(f"--{arg}" for arg in missing_args) + raise ValueError( + f"The following arguments are required but missing: {missing_args_str}." + f" Please specify them either via command-line arguments or directly in `base_args`." + ) + + # Replace config + base_args["config_path"] = redact_config( + { + "orig_config": base_args["config_path"], + "model_type": base_args["model_type"], + "new_config": "", + } + ) + + # Trim train + trim_args_train = { + "input_directory": base_args["data_path"], + "max_folders": base_args["max_folders"], + } + base_args["data_path"] = trim_directory(trim_args_train) + # Trim valid + trim_args_valid = { + "input_directory": base_args["valid_path"], + "max_folders": base_args["max_folders"], + } + base_args["valid_path"] = trim_directory(trim_args_valid) + # Valid to inference + if not base_args["input_folder"]: + tests_dir = os.path.join( + os.path.dirname(base_args["valid_path"]), "for_inference" + ) + base_args["input_folder"] = tests_dir + val_to_inf_args = { + "valid_path": base_args["valid_path"], + "inference_dir": base_args["input_folder"], + "max_mixtures": 1, + } + copying_files(val_to_inf_args) + + if base_args["check_valid"]: + valid_args = { + key: base_args[key] + for key in [ + "model_type", + "config_path", + "start_check_point", + "store_dir", + "device_ids", + "num_workers", + "pin_memory", + "extension", + "use_tta", + "metrics", + "lora_checkpoint", + "draw_spectro", + ] + } + valid_args["valid_path"] = [base_args["valid_path"]] + print("Start validation.") + check_validation(valid_args) + print(f"Validation ended. See results in {base_args['store_dir']}") + + if base_args["check_inference"]: + inference_args = { + key: base_args[key] + for key in [ + "model_type", + "config_path", + "start_check_point", + "input_folder", + "store_dir", + "device_ids", + "extract_instrumental", + "disable_detailed_pbar", + "force_cpu", + "flac_file", + "pcm_type", + "use_tta", + "lora_checkpoint", + "draw_spectro", + ] + } + + print("Start inference.") + proc_folder(inference_args) + print(f"Inference ended. See results in {base_args['store_dir']}") + + if base_args["check_train"]: + train_args = { + key: base_args[key] + for key in [ + "model_type", + "config_path", + "start_check_point", + "results_path", + "data_path", + "dataset_type", + "valid_path", + "num_workers", + "pin_memory", + "seed", + "device_ids", + "use_multistft_loss", + "use_mse_loss", + "use_l1_loss", + "wandb_key", + "pre_valid", + "metrics", + "metric_for_scheduler", + "train_lora", + "lora_checkpoint", + ] + } + + print("Start train.") + train_model(train_args) + + print("End!") + + +if __name__ == "__main__": + test_settings(None, "user") diff --git a/src/third_party/MusicSourceSeparationTraining/train.py b/src/third_party/MusicSourceSeparationTraining/train.py new file mode 100644 index 0000000000000000000000000000000000000000..4a47789f0c10c302a00ebe90730f7f9122d37751 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/train.py @@ -0,0 +1,599 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" +__version__ = "1.0.5" + +import argparse +import sys +import warnings +from typing import Callable, List, Union + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +import wandb +from ml_collections import ConfigDict +from tqdm.auto import tqdm +from utils.model_utils import ( + initialize_model_and_device, + normalize_batch, + save_last_weights, + save_weights, +) +from utils.settings import ( + get_model_from_config, + get_scheduler, + initialize_environment, + initialize_environment_ddp, + parse_args_train, + wandb_init, +) +from valid import valid, valid_multi_gpu + +warnings.filterwarnings("ignore") + + +def forward_step( + x, y, active_stem_ids, get_internal_loss, model, multi_loss, device_ids +): + if get_internal_loss: + loss = model(x, y, active_stem_ids=active_stem_ids) + if isinstance(device_ids, (list, tuple)): + loss = loss.mean() + return loss + else: + y_ = model(x) + return multi_loss(y_, y, x) + + +def train_one_epoch( + model: torch.nn.Module, + config: ConfigDict, + args: argparse.Namespace, + optimizer: torch.optim.Optimizer, + device: torch.device, + device_ids: List[int], + epoch: int, + use_amp: bool, + scaler: torch.cuda.amp.GradScaler, + scheduler, + gradient_accumulation_steps: int, + train_loader: torch.utils.data.DataLoader, + multi_loss: Callable[ + [ + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + torch.Tensor, + ], + all_losses=None, + world_size=None, + ema_model=None, + safe_mode=None, +) -> None: + """ + Train the model for one epoch. + + Args: + world_size: + scheduler: + model: The model to train. + config: Configuration object containing training parameters. + args: Command-line arguments with specific settings (e.g., model type). + optimizer: Optimizer used for training. + device: Device to run the model on (CPU or GPU). + device_ids: List of GPU device IDs if using multiple GPUs. + epoch: The current epoch number. + use_amp: Whether to use automatic mixed precision (AMP) for training. + scaler: Scaler for AMP to manage gradient scaling. + gradient_accumulation_steps: Number of gradient accumulation steps before updating the optimizer. + train_loader: DataLoader for the training dataset. + multi_loss: The loss function to use during training. + + Returns: + None + """ + ddp = True if world_size else False + should_print = not dist.is_initialized() or dist.get_rank() == 0 + model.train() + if not ddp: + model.to(device) + if should_print: + print(f"Train epoch: {epoch} Learning rate: {optimizer.param_groups[0]['lr']}") + sys.stdout.flush() + loss_val = 0.0 + total = 0 + all_losses[f"epoch_{epoch}"] = [] + + normalize = getattr(config.training, "normalize", False) + + get_internal_loss = ( + args.model_type + in ( + "mel_band_roformer", + "bs_roformer", + "bs_mamba2", + "mel_band_conformer", + "bs_conformer", + ) + and not args.use_standard_loss + ) + + if ddp: + pbar = ( + tqdm(train_loader, dynamic_ncols=True) + if dist.get_rank() == 0 + else train_loader + ) + else: + pbar = tqdm(train_loader) + + for i, data in enumerate(pbar): + if len(data) == 3: + batch, mixes, active_stem_ids = data + elif len(data) == 2: + batch, mixes = data + active_stem_ids = None + else: + raise ValueError(f"len data is {len(data)}") + x = mixes.to(device) + y = batch.to(device) + + if normalize: + x, y = normalize_batch(x, y) + if safe_mode: + try: + with torch.cuda.amp.autocast(enabled=use_amp): + loss = forward_step( + x, + y, + active_stem_ids, + get_internal_loss, + model, + multi_loss, + device_ids, + ) + except Exception as e: + print(f"Error: {e}") + continue + else: + with torch.cuda.amp.autocast(enabled=use_amp): + loss = forward_step( + x, + y, + active_stem_ids, + get_internal_loss, + model, + multi_loss, + device_ids, + ) + loss /= gradient_accumulation_steps + scaler.scale(loss).backward() + + if ((i + 1) % gradient_accumulation_steps == 0) or (i == len(train_loader) - 1): + scaler.unscale_(optimizer) + + if config.training.grad_clip: + nn.utils.clip_grad_norm_(model.parameters(), config.training.grad_clip) + + scaler.step(optimizer) + scaler.update() + + if ema_model is not None: + if ddp: + ema_model.update_parameters(model.module) + else: + ema_model.update_parameters(model) + + if scheduler.name in ["linear_scheduler"]: + scheduler.step() + optimizer.zero_grad(set_to_none=True) + if ddp: + with torch.no_grad(): + loss_copy = loss.detach().clone() + dist.all_reduce(loss_copy, op=dist.ReduceOp.SUM) + loss_copy /= dist.get_world_size() + if dist.get_rank() == 0: + li = loss_copy.item() * gradient_accumulation_steps + all_losses[f"epoch_{epoch}"].append(li) + loss_val += li + total += 1 + pbar.set_postfix( + {"loss": 100 * li, "avg_loss": 100 * loss_val / (i + 1)} + ) + sys.stdout.flush() + wandb.log( + {"loss": 100 * li, "avg_loss": 100 * loss_val / (i + 1), "i": i} + ) + else: + li = loss.item() * gradient_accumulation_steps + all_losses[f"epoch_{epoch}"].append(li) + loss_val += li + total += 1 + pbar.set_postfix({"loss": 100 * li, "avg_loss": 100 * loss_val / (i + 1)}) + wandb.log({"loss": 100 * li, "avg_loss": 100 * loss_val / (i + 1), "i": i}) + loss.detach() + + if should_print: + print(f"Training loss: {loss_val / total}") + wandb.log( + { + "train_loss": loss_val / total, + "epoch": epoch, + "learning_rate": optimizer.param_groups[0]["lr"], + } + ) + + +def compute_epoch_metrics( + model: torch.nn.Module, + args: argparse.Namespace, + config: ConfigDict, + device: torch.device, + device_ids: List[int], + best_metric: float, + epoch: int, + scheduler: torch.optim.lr_scheduler, + optimizer, + all_time_all_metrics, + all_losses, + world_size=None, + metrics_avg=None, + all_metrics=None, +) -> float: + """ + Compute and log the metrics for the current epoch, and save model weights if the metric improves. + + Args: + all_losses: + all_metrics: + metrics_avg: + world_size: + model: The model to evaluate. + args: Command-line arguments containing configuration paths and other settings. + config: Configuration dictionary containing training settings. + device: The device (CPU or GPU) used for evaluation. + device_ids: List of GPU device IDs when using multiple GPUs. + best_metric: The best metric value seen so far. + epoch: The current epoch number. + scheduler: The learning rate scheduler to adjust the learning rate. + optimizer: + all_time_all_metrics: + Returns: + The updated best_metric. + """ + + ddp = True if world_size else False + should_print = not dist.is_initialized() or dist.get_rank() == 0 + if not ddp: + if torch.cuda.is_available() and len(device_ids) > 1: + metrics_avg, all_metrics = valid_multi_gpu( + model, args, config, args.device_ids, verbose=False + ) + else: + metrics_avg, all_metrics = valid(model, args, config, device, verbose=False) + all_time_all_metrics[f"epoch_{epoch}"] = all_metrics + + metric_avg = metrics_avg[args.metric_for_scheduler] + if metric_avg > best_metric: + if args.each_metrics_in_name: + stem_parts = [] + for stem_name, values in all_metrics[args.metric_for_scheduler].items(): + stem_values = np.array(values) + mean_val = stem_values.mean() + std_val = stem_values.std() + stem_parts.append( + f"{stem_name}_{args.metric_for_scheduler}_{mean_val:.4f}_std_{std_val:.4f}" + ) + stem_info = "__".join(stem_parts) + store_path = f"{args.results_path}/model_{args.model_type}_ep_{epoch}_{stem_info}.ckpt" + else: + store_path = f"{args.results_path}/model_{args.model_type}_ep_{epoch}_{args.metric_for_scheduler}_{metric_avg:.4f}.ckpt" + if should_print: + print(f"Store weights: {store_path}") + save_weights( + store_path=store_path, + model=model, + device_ids=device_ids, + optimizer=optimizer, + epoch=epoch, + all_time_all_metrics=all_time_all_metrics, + all_losses=all_losses, + best_metric=best_metric, + args=args, + scheduler=scheduler, + ) + best_metric = metric_avg + + if args.save_weights_every_epoch: + metric_string = "" + for m in metrics_avg: + metric_string += "_{}_{:.4f}".format(m, metrics_avg[m]) + store_path = f"{args.results_path}/model_{args.model_type}_ep_{epoch}{metric_string}.ckpt" + save_weights( + store_path=store_path, + model=model, + device_ids=device_ids, + optimizer=optimizer, + epoch=epoch, + all_time_all_metrics=all_time_all_metrics, + all_losses=all_losses, + best_metric=best_metric, + args=args, + scheduler=scheduler, + ) + + if scheduler.name in ["ReduceLROnPlateau"]: + scheduler.step(metric_avg) + + if should_print: + wandb.log({"metric_main": metric_avg, "best_metric": best_metric}) + for metric_name in metrics_avg: + wandb.log({f"metric_{metric_name}": metrics_avg[metric_name]}) + + return best_metric + + +def train_model( + args: Union[argparse.Namespace, None], rank=None, world_size=None +) -> None: + """ + Trains the model based on the provided arguments, including data preparation, optimizer setup, + and loss calculation. The model is trained for multiple epochs with logging via wandb. + + Args: + world_size: + rank: + args: Command-line arguments containing configuration paths, hyperparameters, and other settings. + + Returns: + None + """ + + from torch.cuda.amp.grad_scaler import GradScaler + from utils.dataset import prepare_data + from utils.losses import choice_loss + from utils.model_utils import ( + get_lora, + get_optimizer, + load_start_checkpoint, + log_model_info, + ) + + args = parse_args_train(args) + ddp = True if world_size else False + if ddp: + initialize_environment_ddp(rank, world_size, args.seed, args.results_path) + else: + initialize_environment(args.seed, args.results_path) + model, config = get_model_from_config(args.model_type, args.config_path) + if "model_type" in config.training: + args.model_type = config.training.model_type + use_amp = getattr(config.training, "use_amp", True) + device_ids = args.device_ids + if ddp: + batch_size = config.training.batch_size + else: + batch_size = config.training.batch_size * len(device_ids) + + if not dist.is_initialized() or dist.get_rank() == 0: + wandb_init(args, config, batch_size) + + train_loader = prepare_data(config, args, batch_size) + + if args.start_check_point: + checkpoint = torch.load( + args.start_check_point, weights_only=False, map_location="cpu" + ) + load_start_checkpoint(args, model, checkpoint, type_="train") + model = get_lora(args, config, model) + + if args.freeze_layers is not None: + freeze_layers = [] + train_layers = [] + for name, param in model.named_parameters(): + if any(name.startswith(prefix) for prefix in args.freeze_layers): + freeze_layers.append(name) + print("Freezing layer:", name) + param.requires_grad = False + else: + train_layers.append(name) + print("Trainable layers: {}".format(len(train_layers))) + print("Frozen layers: {}".format(len(freeze_layers))) + + if ddp: + device = torch.device(f"cuda:{rank}") + model.to(device) + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[rank], find_unused_parameters=True + ) + model_module = model.module + else: + device, model = initialize_model_and_device(model, args.device_ids) + # If model is DataParallel, get underlying module + model_module = model.module if hasattr(model, "module") else model + + ema_model = None + if hasattr(config.training, "ema_momentum") and config.training.ema_momentum > 0: + from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"Initializing EMA with decay: {config.training.ema_momentum}") + ema_model = AveragedModel( + model_module, + multi_avg_fn=get_ema_multi_avg_fn(config.training.ema_momentum), + ) + + if args.pre_valid: + model_to_valid = ema_model if ema_model is not None else model + if ddp: + valid_multi_gpu( + model_to_valid, args, config, args.device_ids, verbose=False + ) + else: + if torch.cuda.is_available() and len(args.device_ids) > 1: + valid_multi_gpu( + model_to_valid, args, config, args.device_ids, verbose=True + ) + else: + valid(model_to_valid, args, config, device, verbose=True) + + gradient_accumulation_steps = int( + getattr(config.training, "gradient_accumulation_steps", 1) + ) + + # load optimizer + optimizer = get_optimizer(config, model) + scheduler = get_scheduler(config, optimizer) + + if ( + args.start_check_point + and "optimizer_state_dict" in checkpoint + and args.load_optimizer + ): + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + + if ( + args.start_check_point + and "scheduler_state_dict" in checkpoint + and args.load_scheduler + ): + scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) + + # load num epoch + if args.start_check_point and "epoch" in checkpoint and args.load_epoch: + start_epoch = checkpoint["epoch"] + 1 + else: + start_epoch = 0 + + if args.start_check_point and "best_metric" in checkpoint and args.load_best_metric: + best_metric = checkpoint["best_metric"] + else: + best_metric = float("-inf") + + if args.start_check_point and "all_metrics" in checkpoint and args.load_all_metrics: + all_time_all_metrics = checkpoint["all_metrics"] + else: + all_time_all_metrics = {} + + if args.start_check_point and "all_losses" in checkpoint and args.load_all_losses: + all_losses = checkpoint["all_losses"] + else: + all_losses = {} + + multi_loss = choice_loss(args, config) + scaler = GradScaler() + + if args.set_per_process_memory_fraction: + torch.cuda.set_per_process_memory_fraction(1.0) + torch.cuda.empty_cache() + + safe_mode = args.safe_mode + + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + if should_print: + if world_size: + batch_size = config.training.batch_size + ef_batch_size = batch_size * gradient_accumulation_steps * world_size + num_gpu = world_size + else: + device_ids = args.device_ids + batch_size = config.training.batch_size * len(device_ids) + ef_batch_size = batch_size * gradient_accumulation_steps + num_gpu = len(device_ids) + + print( + f"Instruments: {config.training.instruments}\n" + f"Metrics for training: {args.metrics}. Metric for scheduler: {args.metric_for_scheduler}\n" + f"Patience: {config.training.patience} " + f"Reduce factor: {config.training.reduce_factor}\n" + f"Batch size: {batch_size} " + f"Grad accum steps: {gradient_accumulation_steps} " + f"Num gpus: {num_gpu} " + f"Effective batch size: {ef_batch_size}\n" + f"Dataset type: {args.dataset_type}\n" + f"Optimizer: {config.training.optimizer}" + ) + + print(f"Train for: {config.training.num_epochs} epochs") + log_model_info(model, args.results_path) + + for epoch in range(start_epoch, config.training.num_epochs): + if ddp: + train_loader.sampler.set_epoch(epoch) + + train_one_epoch( + model, + config, + args, + optimizer, + device, + device_ids, + epoch, + use_amp, + scaler, + scheduler, + gradient_accumulation_steps, + train_loader, + multi_loss, + all_losses, + world_size, + ema_model=ema_model, + safe_mode=safe_mode, + ) + + model_to_valid = ema_model if ema_model is not None else model + + if should_print: + save_last_weights( + args, + model, + device_ids, + optimizer, + epoch, + all_time_all_metrics, + best_metric, + scheduler, + ) + if ddp: + metrics_avg, all_metrics = valid_multi_gpu( + model, args, config, args.device_ids, verbose=False + ) + if rank == 0: + all_time_all_metrics[f"epoch_{epoch}"] = all_metrics + best_metric = compute_epoch_metrics( + model=model, + args=args, + config=config, + device=device, + device_ids=device_ids, + best_metric=best_metric, + epoch=epoch, + scheduler=scheduler, + optimizer=optimizer, + all_time_all_metrics=all_time_all_metrics, + all_losses=all_losses, + world_size=world_size, + metrics_avg=metrics_avg, + all_metrics=all_metrics, + ) + else: + best_metric = compute_epoch_metrics( + model=model, + args=args, + config=config, + device=device, + device_ids=device_ids, + best_metric=best_metric, + epoch=epoch, + scheduler=scheduler, + optimizer=optimizer, + all_time_all_metrics=all_time_all_metrics, + all_losses=all_losses, + ) + + +if __name__ == "__main__": + train_model(None) diff --git a/src/third_party/MusicSourceSeparationTraining/train_accelerate.py b/src/third_party/MusicSourceSeparationTraining/train_accelerate.py new file mode 100644 index 0000000000000000000000000000000000000000..1b470873f86b959b514942be2b561cb8675f0de4 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/train_accelerate.py @@ -0,0 +1,505 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" +__version__ = "1.0.3" + +# Read more here: +# https://huggingface.co/docs/accelerate/index + +import argparse +import glob +import os +import time +import warnings + +import auraloss +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +import torch.nn.functional as F +import wandb +from accelerate import Accelerator +from torch.optim import SGD, Adam, AdamW, RAdam, RMSprop +from torch.optim.lr_scheduler import ReduceLROnPlateau +from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn +from torch.utils.data import DataLoader +from tqdm.auto import tqdm +from utils.dataset import MSSDataset +from utils.losses import masked_loss +from utils.metrics import sdr +from utils.model_utils import ( + demix, + load_not_compatible_weights, + prefer_target_instrument, +) +from utils.settings import get_model_from_config, manual_seed + +warnings.filterwarnings("ignore") + + +def valid(model, valid_loader, args, config, device, verbose=False): + instruments = prefer_target_instrument(config) + + all_sdr = dict() + for instr in instruments: + all_sdr[instr] = [] + + all_mixtures_path = valid_loader + if verbose: + all_mixtures_path = tqdm(valid_loader) + + pbar_dict = {} + for path_list in all_mixtures_path: + path = path_list[0] + mix, sr = sf.read(path) + folder = os.path.dirname(path) + res = demix(config, model, mix.T, device, model_type=args.model_type) # mix.T + for instr in instruments: + if instr != "other" or config.training.other_fix is False: + track, sr1 = sf.read(folder + "/{}.wav".format(instr)) + else: + # other is actually instrumental + track, sr1 = sf.read(folder + "/{}.wav".format("vocals")) + track = mix - track + # sf.write("{}.wav".format(instr), res[instr].T, sr, subtype='FLOAT') + references = np.expand_dims(track, axis=0) + estimates = np.expand_dims(res[instr].T, axis=0) + sdr_val = sdr(references, estimates)[0] + single_val = torch.from_numpy(np.array([sdr_val])).to(device) + all_sdr[instr].append(single_val) + pbar_dict["sdr_{}".format(instr)] = sdr_val + if verbose: + all_mixtures_path.set_postfix(pbar_dict) + + return all_sdr + + +class MSSValidationDataset(torch.utils.data.Dataset): + def __init__(self, args): + all_mixtures_path = [] + for valid_path in args.valid_path: + part = sorted(glob.glob(valid_path + "/*/mixture.wav")) + if len(part) == 0: + print("No validation data found in: {}".format(valid_path)) + all_mixtures_path += part + + self.list_of_files = all_mixtures_path + + def __len__(self): + return len(self.list_of_files) + + def __getitem__(self, index): + return self.list_of_files[index] + + +def train_model(args): + accelerator = Accelerator() + device = accelerator.device + + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_type", + type=str, + default="mdx23c", + help="One of mdx23c, htdemucs, segm_models, mel_band_roformer, bs_roformer, swin_upernet, bandit", + ) + parser.add_argument("--config_path", type=str, help="path to config file") + parser.add_argument( + "--start_check_point", + type=str, + default="", + help="Initial checkpoint to start training", + ) + parser.add_argument( + "--results_path", + type=str, + help="path to folder where results will be stored (weights, metadata)", + ) + parser.add_argument( + "--data_path", + nargs="+", + type=str, + help="Dataset data paths. You can provide several folders.", + ) + parser.add_argument( + "--dataset_type", + type=int, + default=1, + help="Dataset type. Must be one of: 1, 2, 3 or 4. Details here: https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/docs/dataset_types.md", + ) + parser.add_argument( + "--valid_path", + nargs="+", + type=str, + help="validation data paths. You can provide several folders.", + ) + parser.add_argument( + "--num_workers", type=int, default=0, help="dataloader num_workers" + ) + parser.add_argument( + "--pin_memory", type=bool, default=False, help="dataloader pin_memory" + ) + parser.add_argument("--seed", type=int, default=0, help="random seed") + parser.add_argument( + "--device_ids", nargs="+", type=int, default=[0], help="list of gpu ids" + ) + parser.add_argument( + "--use_multistft_loss", + action="store_true", + help="Use MultiSTFT Loss (from auraloss package)", + ) + parser.add_argument( + "--use_mse_loss", action="store_true", help="Use default MSE loss" + ) + parser.add_argument("--use_l1_loss", action="store_true", help="Use L1 loss") + parser.add_argument("--wandb_key", type=str, default="", help="wandb API Key") + parser.add_argument( + "--pre_valid", action="store_true", help="Run validation before training" + ) + if args is None: + args = parser.parse_args() + else: + args = parser.parse_args(args) + + manual_seed(args.seed + int(time.time())) + # torch.backends.cudnn.benchmark = True + torch.backends.cudnn.deterministic = ( + False # Fix possible slow down with dilation convolutions + ) + torch.multiprocessing.set_start_method("spawn") + + model, config = get_model_from_config(args.model_type, args.config_path) + if "model_type" in config.training: + args.model_type = config.training.model_type + accelerator.print("Instruments: {}".format(config.training.instruments)) + + os.makedirs(args.results_path, exist_ok=True) + + device_ids = args.device_ids + batch_size = config.training.batch_size + + # wandb + if ( + accelerator.is_main_process + and args.wandb_key is not None + and args.wandb_key.strip() != "" + ): + wandb.login(key=args.wandb_key) + wandb.init( + project="msst-accelerate", + config={ + "config": config, + "args": args, + "device_ids": device_ids, + "batch_size": batch_size, + }, + ) + else: + wandb.init(mode="disabled") + + # Fix for num of steps + config.training.num_steps *= accelerator.num_processes + + trainset = MSSDataset( + config, + args.data_path, + batch_size=batch_size, + metadata_path=os.path.join( + args.results_path, "metadata_{}.pkl".format(args.dataset_type) + ), + dataset_type=args.dataset_type, + verbose=accelerator.is_main_process, + ) + + train_loader = DataLoader( + trainset, + batch_size=batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=args.pin_memory, + ) + + validset = MSSValidationDataset(args) + valid_dataset_length = len(validset) + + valid_loader = DataLoader( + validset, + batch_size=1, + shuffle=False, + ) + + valid_loader = accelerator.prepare(valid_loader) + + if args.start_check_point != "": + accelerator.print("Start from checkpoint: {}".format(args.start_check_point)) + if 1: + load_not_compatible_weights(model, args.start_check_point, verbose=False) + else: + model.load_state_dict(torch.load(args.start_check_point)) + + optim_params = dict() + if "optimizer" in config: + optim_params = dict(config["optimizer"]) + accelerator.print("Optimizer params from config:\n{}".format(optim_params)) + + if config.training.optimizer == "adam": + optimizer = Adam(model.parameters(), lr=config.training.lr, **optim_params) + elif config.training.optimizer == "adamw": + optimizer = AdamW(model.parameters(), lr=config.training.lr, **optim_params) + elif config.training.optimizer == "radam": + optimizer = RAdam(model.parameters(), lr=config.training.lr, **optim_params) + elif config.training.optimizer == "rmsprop": + optimizer = RMSprop(model.parameters(), lr=config.training.lr, **optim_params) + elif config.training.optimizer == "prodigy": + from prodigyopt import Prodigy + + # you can choose weight decay value based on your problem, 0 by default + # We recommend using lr=1.0 (default) for all networks. + optimizer = Prodigy(model.parameters(), lr=config.training.lr, **optim_params) + elif config.training.optimizer == "adamw8bit": + import bitsandbytes as bnb + + optimizer = bnb.optim.AdamW8bit( + model.parameters(), lr=config.training.lr, **optim_params + ) + elif config.training.optimizer == "sgd": + accelerator.print("Use SGD optimizer") + optimizer = SGD(model.parameters(), lr=config.training.lr, **optim_params) + else: + accelerator.print("Unknown optimizer: {}".format(config.training.optimizer)) + exit() + + if accelerator.is_main_process: + print("Processes GPU: {}".format(accelerator.num_processes)) + print( + "Patience: {} Reduce factor: {} Batch size: {} Optimizer: {}".format( + config.training.patience, + config.training.reduce_factor, + batch_size, + config.training.optimizer, + ) + ) + # Reduce LR if no SDR improvements for several epochs + scheduler = ReduceLROnPlateau( + optimizer, + "max", + # patience=accelerator.num_processes * config.training.patience, # This is strange place... + patience=config.training.patience, + factor=config.training.reduce_factor, + ) + + if args.use_multistft_loss: + try: + loss_options = dict(config.loss_multistft) + except: + loss_options = dict() + accelerator.print("Loss options: {}".format(loss_options)) + loss_multistft = auraloss.freq.MultiResolutionSTFTLoss(**loss_options) + + model, optimizer, train_loader, scheduler = accelerator.prepare( + model, optimizer, train_loader, scheduler + ) + + ema_model = None + if hasattr(config.training, "ema_momentum") and config.training.ema_momentum > 0: + accelerator.print( + f"Initializing EMA with decay: {config.training.ema_momentum}" + ) + ema_model = AveragedModel( + accelerator.unwrap_model(model), + multi_avg_fn=get_ema_multi_avg_fn(config.training.ema_momentum), + ) + ema_model.to(device) + + if args.pre_valid: + model_to_valid = ema_model if ema_model is not None else model + sdr_list = valid( + model_to_valid, + valid_loader, + args, + config, + device, + verbose=accelerator.is_main_process, + ) + sdr_list = accelerator.gather(sdr_list) + accelerator.wait_for_everyone() + + # print(sdr_list) + + sdr_avg = 0.0 + instruments = prefer_target_instrument(config) + + for instr in instruments: + # print(sdr_list[instr]) + sdr_data = torch.cat(sdr_list[instr], dim=0).cpu().numpy() + sdr_val = sdr_data.mean() + accelerator.print("Valid length: {}".format(valid_dataset_length)) + accelerator.print( + "Instr SDR {}: {:.4f} Debug: {}".format(instr, sdr_val, len(sdr_data)) + ) + sdr_val = sdr_data[:valid_dataset_length].mean() + accelerator.print( + "Instr SDR {}: {:.4f} Debug: {}".format(instr, sdr_val, len(sdr_data)) + ) + sdr_avg += sdr_val + sdr_avg /= len(instruments) + if len(instruments) > 1: + accelerator.print("SDR Avg: {:.4f}".format(sdr_avg)) + sdr_list = None + + accelerator.print("Train for: {}".format(config.training.num_epochs)) + best_sdr = -100 + for epoch in range(config.training.num_epochs): + model.train().to(device) + accelerator.print( + "Train epoch: {} Learning rate: {}".format( + epoch, optimizer.param_groups[0]["lr"] + ) + ) + loss_val = 0.0 + total = 0 + + pbar = tqdm(train_loader, disable=not accelerator.is_main_process) + for i, (batch, mixes) in enumerate(pbar): + y = batch + x = mixes + + if args.model_type in [ + "mel_band_roformer", + "bs_roformer", + "bs_mamba2", + "mel_band_conformer", + "bs_conformer", + ]: + # loss is computed in forward pass + loss = model(x, y) + else: + y_ = model(x) + if args.use_multistft_loss: + y1_ = torch.reshape( + y_, (y_.shape[0], y_.shape[1] * y_.shape[2], y_.shape[3]) + ) + y1 = torch.reshape( + y, (y.shape[0], y.shape[1] * y.shape[2], y.shape[3]) + ) + loss = loss_multistft(y1_, y1) + # We can use many losses at the same time + if args.use_mse_loss: + loss += 1000 * nn.MSELoss()(y1_, y1) + if args.use_l1_loss: + loss += 1000 * F.l1_loss(y1_, y1) + elif args.use_mse_loss: + loss = nn.MSELoss()(y_, y) + elif args.use_l1_loss: + loss = F.l1_loss(y_, y) + else: + loss = masked_loss( + y_, + y, + q=config.training.q, + coarse=config.training.coarse_loss_clip, + ) + + accelerator.backward(loss) + if config.training.grad_clip: + accelerator.clip_grad_norm_( + model.parameters(), config.training.grad_clip + ) + + optimizer.step() + optimizer.zero_grad() + + if ema_model is not None: + ema_model.update_parameters(accelerator.unwrap_model(model)) + + li = loss.item() + loss_val += li + total += 1 + if accelerator.is_main_process: + wandb.log( + { + "loss": 100 * li, + "avg_loss": 100 * loss_val / (i + 1), + "total": total, + "loss_val": loss_val, + "i": i, + } + ) + pbar.set_postfix( + {"loss": 100 * li, "avg_loss": 100 * loss_val / (i + 1)} + ) + + if accelerator.is_main_process: + print("Training loss: {:.6f}".format(loss_val / total)) + wandb.log({"train_loss": loss_val / total, "epoch": epoch}) + + # Save last + store_path = args.results_path + "/last_{}.ckpt".format(args.model_type) + accelerator.wait_for_everyone() + if accelerator.is_main_process: + if ema_model is not None: + accelerator.save(ema_model.module.state_dict(), store_path) + else: + unwrapped_model = accelerator.unwrap_model(model) + accelerator.save(unwrapped_model.state_dict(), store_path) + + # Validation + model_to_valid = ema_model if ema_model is not None else model + sdr_list = valid( + model_to_valid, + valid_loader, + args, + config, + device, + verbose=accelerator.is_main_process, + ) + sdr_list = accelerator.gather(sdr_list) + accelerator.wait_for_everyone() + + sdr_avg = 0.0 + instruments = prefer_target_instrument(config) + + for instr in instruments: + if accelerator.is_main_process and 0: + print(sdr_list[instr]) + sdr_data = torch.cat(sdr_list[instr], dim=0).cpu().numpy() + # sdr_val = sdr_data.mean() + sdr_val = sdr_data[:valid_dataset_length].mean() + if accelerator.is_main_process: + print( + "Instr SDR {}: {:.4f} Debug: {}".format( + instr, sdr_val, len(sdr_data) + ) + ) + wandb.log({f"{instr}_sdr": sdr_val}) + sdr_avg += sdr_val + sdr_avg /= len(instruments) + if len(instruments) > 1: + if accelerator.is_main_process: + print("SDR Avg: {:.4f}".format(sdr_avg)) + wandb.log({"sdr_avg": sdr_avg, "best_sdr": best_sdr}) + + if accelerator.is_main_process: + if sdr_avg > best_sdr: + store_path = ( + args.results_path + + "/model_{}_ep_{}_sdr_{:.4f}.ckpt".format( + args.model_type, epoch, sdr_avg + ) + ) + print("Store weights: {}".format(store_path)) + if ema_model is not None: + accelerator.save(ema_model.module.state_dict(), store_path) + else: + unwrapped_model = accelerator.unwrap_model(model) + accelerator.save(unwrapped_model.state_dict(), store_path) + best_sdr = sdr_avg + + scheduler.step(sdr_avg) + + sdr_list = None + accelerator.wait_for_everyone() + + +if __name__ == "__main__": + train_model(None) diff --git a/src/third_party/MusicSourceSeparationTraining/train_ddp.py b/src/third_party/MusicSourceSeparationTraining/train_ddp.py new file mode 100644 index 0000000000000000000000000000000000000000..4182f14614c416ab3b3336fb21798dc441e4e572 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/train_ddp.py @@ -0,0 +1,43 @@ +# coding: utf-8 +__author__ = "Ilya Kiselev (kiselecheck): https://github.com/kiselecheck" +__version__ = "1.0.1" + +import warnings + +import torch +import torch.multiprocessing as mp +from train import train_model +from utils.settings import cleanup_ddp + +warnings.filterwarnings("ignore") + + +def train_model_single(rank: int, world_size: int, args=None): + """ + Trains the model based on the provided arguments, including data preparation, optimizer setup, + and loss calculation. The model is trained for multiple epochs with logging via wandb. + + Args: + world_size: + rank: + args: Command-line arguments containing configuration paths, hyperparameters, and other settings. + + Returns: + None + """ + train_model(args, rank, world_size) # Close DDP + + +def train_model_ddp(args=None): + world_size = torch.cuda.device_count() + try: + mp.spawn( + train_model_single, args=(world_size, args), nprocs=world_size, join=True + ) + except Exception as e: + cleanup_ddp() + raise e + + +if __name__ == "__main__": + train_model_ddp() diff --git a/src/third_party/MusicSourceSeparationTraining/utils/audio_utils.py b/src/third_party/MusicSourceSeparationTraining/utils/audio_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a7855287dd2d2c89063232fc0ec9d51b89430e9b --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/utils/audio_utils.py @@ -0,0 +1,372 @@ +import os +from typing import Dict, Optional, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import soundfile as sf +import torch.distributed as dist + + +def read_audio_transposed( + path: str, instr: Optional[str] = None, skip_err: bool = False +) -> Tuple[Optional[np.ndarray], Optional[int]]: + """ + Read an audio file and return transposed waveform data with channels first. + + Loads the audio file from `path`, converts mono signals to 2D format, and + transposes the array so that its shape is (channels, length). In case of + errors, either raises an exception or skips gracefully depending on + `skip_err`. + + Args: + path (str): Path to the audio file to load. + instr (Optional[str], optional): Instrument name, used for informative + messages when `skip_err` is True. Defaults to None. + skip_err (bool, optional): If True, skip files with read errors and + return `(None, None)` instead of raising. Defaults to False. + + Returns: + Tuple[Optional[np.ndarray], Optional[int]]: A tuple containing: + - NumPy array of shape (channels, length), or None if skipped. + - Sampling rate as an integer, or None if skipped. + """ + + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + try: + mix, sr = sf.read(path) + except Exception as e: + if skip_err: + if should_print: + print(f"No stem {instr}: skip!") + return None, None + else: + raise RuntimeError(f"Error reading the file at {path}: {e}") + else: + if len(mix.shape) == 1: # For mono audio + mix = np.expand_dims(mix, axis=-1) + return mix.T, sr + + +def normalize_audio(audio: np.ndarray) -> Tuple[np.ndarray, Dict[str, float]]: + """ + Normalize an audio signal using mean and standard deviation. + + Computes the mean and standard deviation from the mono mix of the input + signal, then applies normalization to each channel. + + Args: + audio (np.ndarray): Input audio array of shape (channels, time) or (time,). + + Returns: + Tuple[np.ndarray, Dict[str, float]]: A tuple containing: + - Normalized audio with the same shape as the input. + - A dictionary with keys "mean" and "std" from the original audio. + """ + + mono = audio.mean(0) + mean, std = mono.mean(), mono.std() + return (audio - mean) / std, {"mean": mean, "std": std} + + +def denormalize_audio(audio: np.ndarray, norm_params: Dict[str, float]) -> np.ndarray: + """ + Reverse normalization on an audio signal. + + Applies the stored mean and standard deviation to restore the original + scale of a previously normalized signal. + + Args: + audio (np.ndarray): Normalized audio array to be denormalized. + norm_params (Dict[str, float]): Dictionary containing the keys + "mean" and "std" used during normalization. + + Returns: + np.ndarray: Denormalized audio with the same shape as the input. + """ + + return audio * norm_params["std"] + norm_params["mean"] + + +def draw_spectrogram( + waveform: np.ndarray, sample_rate: int, length: float, output_file: str +) -> None: + """ + Generate and save a spectrogram image from an audio waveform. + + Converts the provided waveform into a mono signal, computes its Short-Time + Fourier Transform (STFT), converts the amplitude spectrogram to dB scale, + and plots it using a plasma colormap. + + Args: + waveform (np.ndarray): Input audio waveform array of shape (time, channels) + or (time,). + sample_rate (int): Sampling rate of the waveform in Hz. + length (float): Duration (in seconds) of the waveform to include in the + spectrogram. + output_file (str): Path to save the resulting spectrogram image. + + Returns: + None + """ + + import librosa.display + + # Cut only required part of spectorgram + x = waveform[: int(length * sample_rate), :] + X = librosa.stft( + x.mean(axis=-1) + ) # perform short-term fourier transform on mono signal + Xdb = librosa.amplitude_to_db( + np.abs(X), ref=np.max + ) # convert an amplitude spectrogram to dB-scaled spectrogram. + fig, ax = plt.subplots() + # plt.figure(figsize=(30, 10)) # initialize the fig size + img = librosa.display.specshow( + Xdb, cmap="plasma", sr=sample_rate, x_axis="time", y_axis="linear", ax=ax + ) + ax.set(title="File: " + os.path.basename(output_file)) + fig.colorbar(img, ax=ax, format="%+2.f dB") + if output_file is not None: + plt.savefig(output_file) + + +def draw_2_mel_spectrogram( + estimates_waveform: np.ndarray, + track_waveform: np.ndarray, + sample_rate: int, + length: float, + output_base: str, +) -> None: + """ + Generate and save separate images for spectrograms and waveforms + for both estimated and original audio. + + Creates two separate images: + - One with mel-spectrograms (estimated vs original) + - One with waveforms (estimated vs original) + + Args: + estimates_waveform (np.ndarray): Estimated audio waveform + track_waveform (np.ndarray): Original audio waveform + sample_rate (int): Sampling rate in Hz + length (float): Duration in seconds to include + output_base (str): Base path for output files (without extension) + + Returns: + None + """ + import librosa.display + + # Prepare both waveforms + waveforms = [estimates_waveform, track_waveform] + titles = ["Estimated", "Original"] + + # Store processed (mono, possibly decimated) waveforms + processed_waveforms: list[tuple[np.ndarray, int]] = [] + + for waveform in waveforms: + # Convert to mono if multi-channel + mono_signal = waveform.mean(axis=-1) if len(waveform.shape) > 1 else waveform + + # Apply decimation for long audio signals + if len(mono_signal) > 60 * sample_rate: + # Decimation: take every second sample + mono_signal = mono_signal[::2] + effective_sr = sample_rate // 2 + else: + effective_sr = sample_rate + + processed_waveforms.append((mono_signal, effective_sr)) + + # Create mel-spectrograms figure + fig_spec, axes_spec = plt.subplots(2, 1, figsize=(16, 10)) + + for i, ((mono_signal, effective_sr), title) in enumerate( + zip(processed_waveforms, titles) + ): + # Compute mel-spectrogram with reduced number of mel bins + S = librosa.feature.melspectrogram(y=mono_signal, sr=effective_sr, n_mels=128) + S_db = librosa.power_to_db(S, ref=np.max) + + # Plot mel-spectrogram + img = librosa.display.specshow( + S_db, + cmap="plasma", + sr=effective_sr, + x_axis="time", + y_axis="mel", + ax=axes_spec[i], + ) + axes_spec[i].set_title( + f"Mel-spectrogram: {title}", fontsize=14, fontweight="bold" + ) + axes_spec[i].set_xlabel("Time (seconds)", fontsize=12) + axes_spec[i].set_ylabel("Frequency (Mel)", fontsize=12) + + # Colorbar intentionally disabled + # fig_spec.colorbar(img, ax=axes_spec, format="%+2.f dB", + # shrink=0.8, pad=0.02, location="right") + + # Set global title for spectrograms + fig_spec.suptitle( + f"Mel-spectrograms: {os.path.basename(output_base)}", + fontsize=16, + fontweight="bold", + y=0.98, + ) + + plt.tight_layout() + plt.subplots_adjust(top=0.94, hspace=0.4, right=0.88) + + # Save spectrograms image with reduced DPI + spec_output = f"{output_base}_spectrograms.jpg" + plt.savefig(spec_output, dpi=150, bbox_inches="tight") + plt.close(fig_spec) + + # Create waveforms figure + fig_wave, axes_wave = plt.subplots(2, 1, figsize=(16, 8)) + + for i, ((mono_signal, effective_sr), title) in enumerate( + zip(processed_waveforms, titles) + ): + # Generate time axis + time = np.linspace(0, len(mono_signal) / effective_sr, len(mono_signal)) + + # Plot simplified waveform for very long signals + if len(mono_signal) > 100000: + # Take every 10th sample for plotting + plot_indices = np.arange(0, len(mono_signal), 10) + axes_wave[i].plot( + time[plot_indices], + mono_signal[plot_indices], + color="#00ff88", + alpha=0.9, + linewidth=0.5, + ) + else: + axes_wave[i].plot( + time, mono_signal, color="#00ff88", alpha=0.9, linewidth=0.8 + ) + + axes_wave[i].fill_between(time, mono_signal, alpha=0.3, color="#00ff8833") + axes_wave[i].set_xlabel("Time (seconds)", fontsize=12) + axes_wave[i].set_ylabel("Amplitude", fontsize=12) + axes_wave[i].set_title(f"Waveform: {title}", fontsize=14, fontweight="bold") + axes_wave[i].grid(True, alpha=0.3, color="gray") + axes_wave[i].set_xlim(0, time[-1]) + + # Set global title for waveforms + fig_wave.suptitle( + f"Waveforms: {os.path.basename(output_base)}", + fontsize=16, + fontweight="bold", + y=0.98, + ) + + plt.tight_layout() + plt.subplots_adjust(top=0.94, hspace=0.4) + + # Save waveforms image + wave_output = f"{output_base}_waveforms.jpg" + plt.savefig(wave_output, dpi=150, bbox_inches="tight") + plt.close(fig_wave) + + +def draw_mel_spectrogram( + waveform: np.ndarray, sample_rate: int, length: float, output_file: str +) -> None: + """ + Generate and save a spectrogram image from an audio waveform. + + Converts the provided waveform into a mono signal, computes its Short-Time + Fourier Transform (STFT), converts the amplitude spectrogram to dB scale, + and plots it using a plasma colormap. + + Args: + waveform (np.ndarray): Input audio waveform array of shape (time, channels) + or (time,). + sample_rate (int): Sampling rate of the waveform in Hz. + length (float): Duration (in seconds) of the waveform to include in the + spectrogram. + output_file (str): Path to save the resulting spectrogram image. + + Returns: + None + """ + + import librosa.display + + # Cut only required part of spectrogram + x = waveform + + # Compute mel-spectrogram instead of STFT + S = librosa.feature.melspectrogram( + y=x.mean(axis=-1), # mono signal + sr=sample_rate, + ) + + # Convert to dB scale + S_db = librosa.power_to_db(S, ref=np.max) + + fig, ax = plt.subplots() + try: + img = librosa.display.specshow( + S_db, cmap="plasma", sr=sample_rate, x_axis="time", y_axis="mel", ax=ax + ) + ax.set(title="Mel-spectrogram: " + os.path.basename(output_file)) + fig.colorbar(img, ax=ax, format="%+2.f dB") + if output_file is not None: + plt.savefig(output_file) + finally: + plt.close(fig) + + plot_waveform_basic( + waveform, sample_rate, output_file.replace(".jpg", "_waveform.jpg") + ) + + +def plot_waveform_basic(waveform, samplerate, output_path=None, theme="dark"): + data = waveform + if len(data.shape) > 1: + data = np.mean(data, axis=1) + try: + themes = { + "dark": {"bg": "#0f0f0f", "wave": "#00ff88", "fill": "#00ff8833"}, + "light": {"bg": "white", "wave": "#2563eb", "fill": "#3b82f633"}, + "purple": {"bg": "#1a1a2e", "wave": "#e94560", "fill": "#e9456033"}, + } + + colors = themes.get(theme, themes["dark"]) + + fig, ax = plt.subplots(figsize=(12, 3), facecolor=colors["bg"]) + + time = np.linspace(0, len(data) / samplerate, len(data)) + + ax.plot(time, data, color=colors["wave"], alpha=0.9, linewidth=0.8) + ax.fill_between(time, data, alpha=0.3, color=colors["fill"]) + + ax.set_facecolor(colors["bg"]) + if theme == "dark" or theme == "purple": + ax.tick_params(colors="white", labelsize=8) + ax.set_xlabel("Time (seconds)", color="white", fontsize=10) + ax.set_ylabel("Amplitude", color="white", fontsize=10) + else: + ax.tick_params(colors="black", labelsize=8) + + ax.grid(True, alpha=0.2, color="gray") + ax.set_xlim(0, time[-1]) + + plt.tight_layout() + + if output_path: + plt.savefig( + output_path, + dpi=200, + bbox_inches="tight", + facecolor=colors["bg"], + edgecolor="none", + ) + + finally: + plt.close() diff --git a/src/third_party/MusicSourceSeparationTraining/utils/dataset.py b/src/third_party/MusicSourceSeparationTraining/utils/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..46b333f7c150260bbbf09317158ee6f8e17f1081 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/utils/dataset.py @@ -0,0 +1,1415 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" + + +import itertools +import multiprocessing +import os +import pickle +import random +import warnings +from concurrent.futures import ThreadPoolExecutor, as_completed +from glob import glob +from typing import Union + +import audiomentations as AU +import numpy as np +import pedalboard as PB +import soundfile as sf +import torch +import torch.distributed as dist +from ml_collections import ConfigDict +from omegaconf import OmegaConf +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from tqdm import tqdm +from tqdm.auto import tqdm + +warnings.filterwarnings("ignore") +import argparse + + +def prepare_data( + config: Union[ConfigDict, OmegaConf], args: argparse.Namespace, batch_size: int +) -> DataLoader: + """ + Build the training DataLoader. If torch.distributed.is_initialized() is True, + construct a DDP DataLoader with DistributedSampler; otherwise, construct a regular DataLoader. + + Args: + config: Dataset configuration passed to MSSDataset. + args: Must provide data_path, results_path, dataset_type, and DataLoader settings. + batch_size: Per-process mini-batch size. + + Returns: + Configured DataLoader for the training split. + """ + # DDP + if dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + + if args.dataset_type != 5: + ddp_batch = ( + batch_size * world_size + ) # maintain "num_steps" semantics across the whole world + else: + ddp_batch = batch_size + + trainset = MSSDataset( + config, + args.data_path, + batch_size=ddp_batch, + metadata_path=os.path.join( + args.results_path, f"metadata_{args.dataset_type}.pkl" + ), + dataset_type=args.dataset_type, + ) + + sampler = DistributedSampler( + trainset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True + ) + + train_loader = DataLoader( + trainset, + batch_size=batch_size, # per-process batch size + sampler=sampler, # sampler handles shuffling in DDP + num_workers=args.num_workers, + pin_memory=args.pin_memory, + persistent_workers=args.persistent_workers, + prefetch_factor=args.prefetch_factor, + ) + else: + trainset = MSSDataset( + config, + args.data_path, + batch_size=batch_size, + metadata_path=os.path.join( + args.results_path, f"metadata_{args.dataset_type}.pkl" + ), + dataset_type=args.dataset_type, + ) + + train_loader = DataLoader( + trainset, + batch_size=batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=args.pin_memory, + persistent_workers=args.persistent_workers, + prefetch_factor=args.prefetch_factor, + ) + + return train_loader + + +def load_chunk(path, length, chunk_size, offset=None, target_channels=2): + """ + Returns array with shape (target_channels, chunk_size) + """ + + if chunk_size <= length: + if offset is None: + start = np.random.randint(length - chunk_size + 1) + else: + start = offset + x = sf.read(path, dtype="float32", start=start, frames=chunk_size)[0] + else: + if offset is None: + start = 0 + else: + start = offset + frames_to_read = length + x = sf.read(path, dtype="float32", start=start, frames=frames_to_read)[0] + + if x.ndim == 1: + x = x[:, None] + + if x.shape[0] < chunk_size: + pad = np.zeros((chunk_size - x.shape[0], x.shape[1]), dtype=np.float32) + x = np.concatenate([x, pad], axis=0) + elif x.shape[0] > chunk_size: + x = x[:chunk_size] + + ch = x.shape[1] + if ch == target_channels: + pass + elif ch > target_channels: + x = x[:, :target_channels] + elif ch == 1: + x = np.repeat(x, 2, axis=1) + else: + raise ValueError(f"Path: {path}, num_channels: {ch}") + + return x.T + + +def get_track_set_length(params): + path, instruments, file_types, dataset_type = params + should_print = ( + not dist.is_initialized() or dist.get_rank() == 0 + ) and dataset_type != 7 + # Check lengths of all instruments (it can be different in some cases) + lengths_arr = [] + for instr in instruments: + length = -1 + for extension in file_types: + path_to_audio_file = path + "/{}.{}".format(instr, extension) + if os.path.isfile(path_to_audio_file): + length = sf.info(path_to_audio_file).frames + break + if length == -1: + if should_print: + print('Cant find file "{}" in folder {}'.format(instr, path)) + continue + lengths_arr.append(length) + lengths_arr = np.array(lengths_arr) + if lengths_arr.min() != lengths_arr.max() and should_print: + print( + f"Warning: lengths of stems are different for path: {path}. ({lengths_arr.min()} != {lengths_arr.max()})" + ) + # We use minimum to allow overflow for soundfile read in non-equal length cases + return path, lengths_arr.min() + + +# For multiprocessing +def get_track_length(params): + path = params + length = sf.info(path).frames + return (path, length) + + +def process_chunk_worker(args): + task, instruments, file_types, min_mean_abs, default_chunk_size = args + track_path, track_length, offset, chunk_size = task + + try: + for instrument in instruments: + instrument_loud_enough = False + for extension in file_types: + path_to_audio_file = track_path + "/{}.{}".format(instrument, extension) + if os.path.isfile(path_to_audio_file): + try: + source = load_chunk( + path_to_audio_file, + length=track_length, + offset=offset, + chunk_size=chunk_size, + ) + if np.abs(source).mean() >= min_mean_abs: + instrument_loud_enough = True + break + except Exception: + return (track_path, offset, False) + + if not instrument_loud_enough: + return (track_path, offset, False) + + return (track_path, offset, True) + + except Exception: + return (track_path, offset, False) + + +class MSSDataset(torch.utils.data.Dataset): + def __init__( + self, + config, + data_path, + metadata_path="metadata.pkl", + dataset_type=1, + batch_size=None, + verbose=True, + ): + self.verbose = verbose + self.config = config + self.dataset_type = dataset_type # 1, 2, 3, 4 or 5 + self.data_path = data_path + self.instruments = instruments = config.training.instruments + if batch_size is None: + batch_size = config.training.batch_size + self.batch_size = batch_size + self.file_types = ["wav", "flac"] + self.metadata_path = metadata_path + + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + # Augmentation block + self.aug = False + if "augmentations" in config: + if config["augmentations"].enable is True: + if self.verbose and should_print: + print("Use augmentation for training") + self.aug = True + else: + if self.verbose and should_print: + print( + "There is no augmentations block in config. Augmentations disabled for training..." + ) + + metadata = self.get_metadata() + + if self.dataset_type in [1, 4, 5, 6, 7]: + if len(metadata) > 0: + if self.verbose and should_print: + print("Found tracks in dataset: {}".format(len(metadata))) + else: + if should_print: + print("No tracks found for training. Check paths you provided!") + exit() + else: + for instr in self.instruments: + if self.verbose and should_print: + print( + "Found tracks for {} in dataset: {}".format( + instr, len(metadata[instr]) + ) + ) + self.metadata = metadata + self.chunk_size = config.audio.chunk_size + self.min_mean_abs = config.audio.min_mean_abs + self.do_chunks = ( + config.training.get("precompute_chunks", False) + and float(self.min_mean_abs) > 0 + ) + # For dataset_type 5 - precompute all chunks + if ( + self.dataset_type == 5 + or (self.dataset_type == 4 or self.dataset_type == 6) + and self.do_chunks + ): + self._initialize_chunks_metadata() + if self.dataset_type == 7: + self._build_class_to_tracks() + + def __len__(self): + if self.dataset_type == 5: + return len(self.chunks_metadata) + return self.config.training.num_steps * self.batch_size + + def __getitem__(self, index): + if self.dataset_type == 7: + res, mix, active_stem_ids = self.load_class_balanced_aligned() + elif self.dataset_type == 5: + track_path, offset = self.chunks_metadata[index] + res = self._load_chunk_by_offset(track_path, offset) + elif self.dataset_type in [1, 2, 3]: + res = self.load_random_mix() + else: # type 4 or 6 + if self.do_chunks: + track_path, offset = self.chunks_metadata[ + np.random.randint(len(self.chunks_metadata)) + ] + res = self._load_chunk_by_offset(track_path, offset) + else: + if self.dataset_type == 6: + res, mix = self.load_aligned_data() + else: + res, _ = self.load_aligned_data() + + # Randomly change loudness of each stem + if self.aug: + if "loudness" in self.config["augmentations"]: + if self.config["augmentations"]["loudness"]: + loud_values = np.random.uniform( + low=self.config["augmentations"]["loudness_min"], + high=self.config["augmentations"]["loudness_max"], + size=(len(res),), + ) + loud_values = torch.tensor(loud_values, dtype=torch.float32) + res *= loud_values[:, None, None] + if self.dataset_type != 6 and self.dataset_type != 7: + mix = res.sum(0) + + if self.aug: + if "mp3_compression_on_mixture" in self.config["augmentations"]: + apply_aug = AU.Mp3Compression( + min_bitrate=self.config["augmentations"][ + "mp3_compression_on_mixture_bitrate_min" + ], + max_bitrate=self.config["augmentations"][ + "mp3_compression_on_mixture_bitrate_max" + ], + backend=self.config["augmentations"][ + "mp3_compression_on_mixture_backend" + ], + p=self.config["augmentations"]["mp3_compression_on_mixture"], + ) + mix_conv = mix.cpu().numpy().astype(np.float32) + required_shape = mix_conv.shape + mix = apply_aug(samples=mix_conv, sample_rate=44100) + # Sometimes it gives longer audio (so we cut) + if mix.shape != required_shape: + mix = mix[..., : required_shape[-1]] + mix = torch.tensor(mix, dtype=torch.float32) + + # If we need to optimize only given stem + if self.config.training.target_instrument is not None: + index = self.config.training.instruments.index( + self.config.training.target_instrument + ) + return res[index : index + 1], mix + + if self.dataset_type == 7: + return res, mix, active_stem_ids + + return res, mix + + def _build_class_to_tracks(self): + import json + + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + cache_path = "class_to_tracks_cache.json" + + total_tracks = len(self.metadata) + max_ratio = self.config.training.get("max_class_presence_ratio", 0.4) + + if os.path.isfile(cache_path): + if should_print: + print("[dataset_type=7] Loading class_to_tracks from cache") + + with open(cache_path, "r", encoding="utf8") as f: + cache = json.load(f) + + if ( + cache.get("total_tracks") == total_tracks + and cache.get("max_ratio") == max_ratio + ): + self.class_to_tracks = cache["class_to_tracks"] + self.available_classes = list(self.class_to_tracks.keys()) + + if should_print: + print( + f"[dataset_type=7] Loaded {len(self.available_classes)} classes from cache" + ) + return + else: + if should_print: + print("[dataset_type=7] Cache invalid, rebuilding") + + class_to_tracks = {instr: [] for instr in self.instruments} + + track_iter = self.metadata + if should_print: + track_iter = tqdm( + self.metadata, + desc="[dataset_type=7] Building class_to_tracks", + total=total_tracks, + ) + + for track_path, _ in track_iter: + for instr in self.instruments: + for ext in self.file_types: + path = f"{track_path}/{instr}.{ext}" + if os.path.isfile(path): + class_to_tracks[instr].append(track_path) + break + + filtered_class_to_tracks = {} + + for instr, tracks in class_to_tracks.items(): + count = len(tracks) + ratio = count / total_tracks + + if count == 0: + continue + + if ratio > max_ratio: + if should_print: + print( + f"[dataset_type=7] Skip frequent stem '{instr}': " + f"{count}/{total_tracks} ({ratio:.1%})" + ) + continue + + filtered_class_to_tracks[instr] = tracks + + if len(filtered_class_to_tracks) == 0: + raise RuntimeError( + "dataset_type 7: all classes were filtered out by frequency threshold" + ) + + self.class_to_tracks = filtered_class_to_tracks + self.available_classes = list(filtered_class_to_tracks.keys()) + + if should_print: + print("[dataset_type=7] Saving class_to_tracks cache") + + with open(cache_path, "w", encoding="utf8") as f: + json.dump( + { + "total_tracks": total_tracks, + "max_ratio": max_ratio, + "class_to_tracks": filtered_class_to_tracks, + }, + f, + indent=2, + ) + + if should_print: + print( + f"[dataset_type=7] Using {len(self.available_classes)} balanced classes " + f"out of {len(self.instruments)} instruments" + ) + + def load_class_balanced_aligned(self): + """ + 1) Randomly choose instrument (class) + 2) Randomly choose track containing this instrument + 3) Load aligned chunk from this track + """ + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + instr = random.choice(self.available_classes) + track_path = random.choice(self.class_to_tracks[instr]) + + # Find track length + track_length = None + for path, length in self.metadata: + if path == track_path: + track_length = length + break + + if track_length is None: + raise RuntimeError(f"Track length not found: {track_path}") + + if track_length >= self.chunk_size: + offset = np.random.randint(track_length - self.chunk_size + 1) + else: + offset = None + + mix = None + for extension in self.file_types: + path_to_mix_file = f"{track_path}/mixture.{extension}" + if os.path.isfile(path_to_mix_file): + try: + mix = load_chunk( + path_to_mix_file, track_length, self.chunk_size, offset=offset + ) + break + except Exception as e: + print(e) + res = [] + active_stem_ids = [] + + for idx, instr in enumerate(self.instruments): + found = False + for extension in self.file_types: + path_to_audio_file = f"{track_path}/{instr}.{extension}" + if os.path.isfile(path_to_audio_file): + try: + source = load_chunk( + path_to_audio_file, + track_length, + self.chunk_size, + offset=offset, + ) + active_stem_ids.append(idx) + found = True + break + except Exception as e: + print(e) + + if not found: + source = np.zeros((2, self.chunk_size), dtype=np.float32) + + res.append(source) + + res = np.stack(res, axis=0) + + if mix is None: + mix = np.sum(res, axis=0) + + if self.aug: + for i, instr in enumerate(self.instruments): + res[i] = self.augm_data(res[i], instr) + + return ( + torch.tensor(res, dtype=torch.float32), + torch.tensor(mix, dtype=torch.float32), + active_stem_ids, + ) + + def _initialize_chunks_metadata(self): + should_print = not dist.is_initialized() or dist.get_rank() == 0 + chunks_cache_path = self.metadata_path.replace(".pkl", "_chunks.pkl") + current_config = { + "chunk_size": self.chunk_size, + "min_mean_abs": self.min_mean_abs, + "instruments": sorted(self.instruments), + } + if os.path.exists(chunks_cache_path): + try: + cached_chunks = pickle.load(open(chunks_cache_path, "rb")) + cached_config = cached_chunks.get("config", {}) + config_matches = ( + cached_config.get("chunk_size") == current_config["chunk_size"] + and cached_config.get("min_mean_abs") + == current_config["min_mean_abs"] + and cached_config.get("instruments") + == current_config["instruments"] + ) + if config_matches: + self.chunks_metadata = cached_chunks["chunks_metadata"] + if self.verbose and should_print: + print( + f"Loaded {len(self.chunks_metadata)} cached chunks from {chunks_cache_path}" + ) + else: + if self.verbose and should_print: + print("Config changed, recomputing chunks...") + print(f"Cached config: {cached_config}") + print(f"Current config: {current_config}") + self.chunks_metadata = self._precompute_and_cache_chunks( + chunks_cache_path, current_config + ) + except Exception as e: + if self.verbose and should_print: + print(f"Chunks cache corrupted ({e}), recomputing...") + self.chunks_metadata = self._precompute_and_cache_chunks( + chunks_cache_path, current_config + ) + else: + self.chunks_metadata = self._precompute_and_cache_chunks( + chunks_cache_path, current_config + ) + + if self.verbose and should_print: + print(f"Precomputed {len(self.chunks_metadata)} chunks") + + def _precompute_and_cache_chunks(self, cache_path, config): + """Precompute all chunks and save to cache with config""" + if self.dataset_type == 4 or self.dataset_type == 6: + chunks_metadata = self._precompute_random_chunks() + elif self.dataset_type == 5: + chunks_metadata = self._precompute_chunks() + else: + raise "Only dataset type 4, 5 can be precomputed" + cache_data = {"chunks_metadata": chunks_metadata, "config": config} + pickle.dump(cache_data, open(cache_path, "wb")) + + return chunks_metadata + + def _precompute_chunks(self): + """Precompute all chunks for dataset_type 5 with overlap 2 using multiprocessing""" + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + tasks = [] + for track_path, track_length in self.metadata: + if track_length < self.chunk_size: + tasks.append((track_path, track_length, 0, track_length)) + else: + step = self.chunk_size // 2 + num_chunks = (track_length - self.chunk_size) // step + 1 + for i in range(num_chunks): + offset = i * step + tasks.append((track_path, track_length, offset, self.chunk_size)) + + if should_print: + print(f"Total tasks to process: {len(tasks)}") + + if multiprocessing.cpu_count() > 1: + chunks_metadata = self._process_tasks_parallel(tasks, should_print) + else: + chunks_metadata = self._process_tasks_sequential(tasks, should_print) + + if self.verbose and should_print: + print( + f"Created {len(chunks_metadata)} good chunks from {len(self.metadata)} tracks" + ) + + return chunks_metadata + + def _precompute_random_chunks(self): + """Precompute exact number of good chunks""" + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + target_count = self.config.training.get( + "num_precompute_chunks", + self.config.training.num_steps + * self.batch_size + * self.config.training.num_epochs, + ) + chunks_metadata = [] + + if should_print: + print(f"Generating exactly {target_count} good chunks...") + + with tqdm(total=target_count, desc="Progress good chunks") as pbar: + while len(chunks_metadata) < target_count: + batch_size = self.config.training.get( + "precompute_batch_for_chunks", 500 + ) + tasks = [] + need = target_count - len(chunks_metadata) + for i in range(batch_size): + track_path, track_length = random.choice(self.metadata) + if track_length < self.chunk_size: + tasks.append((track_path, track_length, 0, track_length)) + else: + offset = np.random.randint(track_length - self.chunk_size + 1) + tasks.append( + (track_path, track_length, offset, self.chunk_size) + ) + + if multiprocessing.cpu_count() > 1: + good_chunks = self._process_tasks_parallel(tasks, False) + else: + good_chunks = self._process_tasks_sequential(tasks, False) + + chunks_metadata.extend(good_chunks) + pbar.update(min(len(good_chunks), need)) + + chunks_metadata = chunks_metadata[:target_count] + + return chunks_metadata + + def _process_tasks_sequential(self, tasks, should_print): + chunks_metadata = [] + + pbar = tqdm(tasks, desc="Processing chunks") if should_print else tasks + for task in pbar: + track_path, track_length, offset, chunk_size = task + if self._is_chunk_loud_enough(track_path, offset, chunk_size, track_length): + chunks_metadata.append((track_path, offset)) + + return chunks_metadata + + def _process_tasks_parallel(self, tasks, should_print): + chunks_metadata = [] + + with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool: + worker_args = [ + ( + task, + self.instruments, + self.file_types, + self.min_mean_abs, + self.chunk_size, + ) + for task in tasks + ] + + results = [] + if should_print: + with tqdm(total=len(tasks), desc="Processing chunks") as pbar: + for i, result in enumerate( + pool.imap_unordered(process_chunk_worker, worker_args) + ): + results.append(result) + pbar.update(1) + else: + for result in pool.imap_unordered(process_chunk_worker, worker_args): + results.append(result) + + for result in results: + track_path, offset, is_loud_enough = result + if is_loud_enough: + chunks_metadata.append((track_path, offset)) + + return chunks_metadata + + def _is_chunk_loud_enough(self, track_path, offset, chunk_size, track_length): + try: + for instrument in self.instruments: + instrument_loud_enough = False + for extension in self.file_types: + path_to_audio_file = track_path + "/{}.{}".format( + instrument, extension + ) + if os.path.isfile(path_to_audio_file): + try: + source = load_chunk( + path_to_audio_file, + length=track_length, + offset=offset, + chunk_size=chunk_size, + ) + if np.abs(source).mean() >= self.min_mean_abs: + instrument_loud_enough = True + break + except Exception as e: + if not dist.is_initialized() or dist.get_rank() == 0: + print( + "Error loading: {} Path: {}".format( + e, path_to_audio_file + ) + ) + return False + + if not instrument_loud_enough: + return False + + return True + + except Exception as e: + if not dist.is_initialized() or dist.get_rank() == 0: + print( + "Error checking chunk loudness: {} Path: {}".format(e, track_path) + ) + return False + + def read_from_metadata_cache(self, track_paths, instr=None): + should_print = not dist.is_initialized() or dist.get_rank() == 0 + metadata = [] + if os.path.isfile(self.metadata_path): + if self.verbose and should_print: + print("Found metadata cache file: {}".format(self.metadata_path)) + old_metadata = pickle.load(open(self.metadata_path, "rb")) + else: + return track_paths, metadata + + if instr: + old_metadata = old_metadata[instr] + + # We will not re-read tracks existed in old metadata file + track_paths_set = set(track_paths) + for old_path, file_size in old_metadata: + if old_path in track_paths_set: + metadata.append([old_path, file_size]) + track_paths_set.remove(old_path) + track_paths = list(track_paths_set) + if len(metadata) > 0 and should_print: + print("Old metadata was used for {} tracks.".format(len(metadata))) + return track_paths, metadata + + def get_metadata(self): + read_metadata_procs = multiprocessing.cpu_count() - 2 + should_print = not dist.is_initialized() or dist.get_rank() == 0 + if "read_metadata_procs" in self.config["training"]: + read_metadata_procs = int(self.config["training"]["read_metadata_procs"]) + + if self.verbose and should_print: + print( + "Dataset type:", + self.dataset_type, + "Processes to use:", + read_metadata_procs, + "\nCollecting metadata for", + str(self.data_path), + ) + + if self.dataset_type in [1, 4, 5, 6, 7]: # Added type 7 + track_paths = [] + if type(self.data_path) == list: + for tp in self.data_path: + tracks_for_folder = sorted(glob(tp + "/*")) + if len(tracks_for_folder) == 0 and should_print: + print( + "Warning: no tracks found in folder '{}'. Please check it!".format( + tp + ) + ) + track_paths += tracks_for_folder + else: + track_paths += sorted(glob(self.data_path + "/*")) + + track_paths = [ + path + for path in track_paths + if os.path.basename(path)[0] != "." and os.path.isdir(path) + ] + track_paths, metadata = self.read_from_metadata_cache(track_paths, None) + + if read_metadata_procs <= 1: + pbar = tqdm(track_paths) if should_print else track_paths + for path in pbar: + track_path, track_length = get_track_set_length( + (path, self.instruments, self.file_types, self.dataset_type) + ) + metadata.append((track_path, track_length)) + else: + with ThreadPoolExecutor(max_workers=read_metadata_procs) as executor: + futures = [ + executor.submit(get_track_set_length, args) + for args in zip( + track_paths, + itertools.repeat(self.instruments), + itertools.repeat(self.file_types), + itertools.repeat(self.dataset_type), + ) + ] + + if should_print: + for f in tqdm(as_completed(futures), total=len(futures)): + track_path, track_length = f.result() + metadata.append((track_path, track_length)) + else: + for f in as_completed(futures): + metadata.append(f.result()) + + elif self.dataset_type == 2: + metadata = dict() + for instr in self.instruments: + metadata[instr] = [] + track_paths = [] + if type(self.data_path) == list: + for tp in self.data_path: + track_paths += sorted(glob(tp + "/{}/*.wav".format(instr))) + track_paths += sorted(glob(tp + "/{}/*.flac".format(instr))) + else: + track_paths += sorted( + glob(self.data_path + "/{}/*.wav".format(instr)) + ) + track_paths += sorted( + glob(self.data_path + "/{}/*.flac".format(instr)) + ) + + track_paths, metadata[instr] = self.read_from_metadata_cache( + track_paths, instr + ) + + if read_metadata_procs <= 1: + pbar = tqdm(track_paths) if should_print else track_paths + for path in pbar: + length = sf.info(path).frames + metadata[instr].append((path, length)) + else: + p = multiprocessing.Pool(processes=read_metadata_procs) + track_iter = p.imap(get_track_length, track_paths) + if should_print: + track_iter = tqdm(track_iter, total=len(track_paths)) + + for out in track_iter: + metadata[instr].append(out) + p.close() + + elif self.dataset_type == 3: + import pandas as pd + + if type(self.data_path) != list: + data_path = [self.data_path] + + metadata = dict() + for i in range(len(self.data_path)): + if self.verbose and should_print: + print("Reading tracks from: {}".format(self.data_path[i])) + df = pd.read_csv(self.data_path[i]) + + skipped = 0 + for instr in self.instruments: + part = df[df["instrum"] == instr].copy() + if should_print: + print("Tracks found for {}: {}".format(instr, len(part))) + for instr in self.instruments: + part = df[df["instrum"] == instr].copy() + metadata[instr] = [] + track_paths = list(part["path"].values) + track_paths, metadata[instr] = self.read_from_metadata_cache( + track_paths, instr + ) + + pbar = tqdm(track_paths) if should_print else track_paths + for path in pbar: + if not os.path.isfile(path): + if should_print: + print("Cant find track: {}".format(path)) + skipped += 1 + continue + # print(path) + try: + length = sf.info(path).frames + except: + if should_print: + print("Problem with path: {}".format(path)) + skipped += 1 + continue + metadata[instr].append((path, length)) + if skipped > 0 and should_print: + print("Missing tracks: {} from {}".format(skipped, len(df))) + else: + if should_print: + print( + "Unknown dataset type: {}. Must be 1, 2, 3, 4, 5 or 6".format( + self.dataset_type + ) + ) + exit() + + # Save metadata + pickle.dump(metadata, open(self.metadata_path, "wb")) + return metadata + + def load_source(self, metadata, instr): + should_print = not dist.is_initialized() or dist.get_rank() == 0 + while True: + if self.dataset_type in [1, 4, 5, 6, 7]: + track_path, track_length = random.choice(metadata) + for extension in self.file_types: + path_to_audio_file = track_path + "/{}.{}".format(instr, extension) + if os.path.isfile(path_to_audio_file): + try: + source = load_chunk( + path_to_audio_file, track_length, self.chunk_size + ) + except Exception as e: + # Sometimes error during FLAC reading, catch it and use zero stem + if should_print: + print( + "Error: {} Path: {}".format(e, path_to_audio_file) + ) + source = np.zeros((2, self.chunk_size), dtype=np.float32) + break + else: + track_path, track_length = random.choice(metadata[instr]) + try: + source = load_chunk(track_path, track_length, self.chunk_size) + except Exception as e: + # Sometimes error during FLAC reading, catch it and use zero stem + if should_print: + print("Error: {} Path: {}".format(e, track_path)) + source = np.zeros((2, self.chunk_size), dtype=np.float32) + + if np.abs(source).mean() >= self.min_mean_abs: # remove quiet chunks + break + if self.aug: + source = self.augm_data(source, instr) + return torch.tensor(source, dtype=torch.float32) + + def load_random_mix(self): + res = [] + for instr in self.instruments: + s1 = self.load_source(self.metadata, instr) + # Mixup augmentation. Multiple mix of same type of stems + if self.aug: + if "mixup" in self.config["augmentations"]: + if self.config["augmentations"].mixup: + mixup = [s1] + for prob in self.config.augmentations.mixup_probs: + if random.uniform(0, 1) < prob: + s2 = self.load_source(self.metadata, instr) + mixup.append(s2) + mixup = torch.stack(mixup, dim=0) + loud_values = np.random.uniform( + low=self.config.augmentations.loudness_min, + high=self.config.augmentations.loudness_max, + size=(len(mixup),), + ) + loud_values = torch.tensor(loud_values, dtype=torch.float32) + mixup *= loud_values[:, None, None] + s1 = mixup.mean(dim=0, dtype=torch.float32) + res.append(s1) + res = torch.stack(res) + return res + + def _load_chunk_by_offset(self, track_path, offset): + """Load specific chunk by track path and offset""" + should_print = not dist.is_initialized() or dist.get_rank() == 0 + res = [] + + for instr in self.instruments: + for extension in self.file_types: + path_to_audio_file = track_path + "/{}.{}".format(instr, extension) + if os.path.isfile(path_to_audio_file): + try: + # Get track length from metadata + track_length = None + for path, length in self.metadata: + if path == track_path: + track_length = length + break + + if track_length is None: + source = np.zeros((2, self.chunk_size), dtype=np.float32) + else: + source = load_chunk( + path_to_audio_file, + track_length, + self.chunk_size, + offset=offset, + ) + except Exception as e: + if should_print: + print("Error: {} Path: {}".format(e, path_to_audio_file)) + source = np.zeros((2, self.chunk_size), dtype=np.float32) + break + else: + source = np.zeros((2, self.chunk_size), dtype=np.float32) + + res.append(source) + + res = np.stack(res, axis=0) + + if self.aug: + for i, instr in enumerate(self.instruments): + res[i] = self.augm_data(res[i], instr) + + return torch.tensor(res, dtype=torch.float32) + + def load_aligned_data(self): + track_path, track_length = random.choice(self.metadata) + should_print = not dist.is_initialized() or dist.get_rank() == 0 + attempts = 10 + while attempts: + if track_length >= self.chunk_size: + common_offset = np.random.randint(track_length - self.chunk_size + 1) + else: + common_offset = None + res = [] + silent_chunks = 0 + for i in self.instruments: + found = False + for extension in self.file_types: + path_to_audio_file = f"{track_path}/{i}.{extension}" + if os.path.isfile(path_to_audio_file): + found = True + try: + source = load_chunk( + path_to_audio_file, + track_length, + self.chunk_size, + offset=common_offset, + ) + except Exception as e: + if should_print: + print(f"Error: {e} Path: {path_to_audio_file}") + source = np.zeros((2, self.chunk_size), dtype=np.float32) + break + + if not found: + source = np.zeros((2, self.chunk_size), dtype=np.float32) + + res.append(source) + if np.abs(source).mean() < self.min_mean_abs: # remove quiet chunks + silent_chunks += 1 + + mix = None + for extension in self.file_types: + path_to_mix_file = track_path + "/mixture.{}".format(extension) + if os.path.isfile(path_to_mix_file): + try: + mix = load_chunk( + path_to_mix_file, + track_length, + self.chunk_size, + offset=common_offset, + ) + except Exception as e: + if should_print: + print( + "Error loading mix: {} Path: {}".format( + e, path_to_mix_file + ) + ) + break + + if silent_chunks == 0: + break + + attempts -= 1 + if attempts <= 0 and should_print: + print("Attempts max!", track_path) + if common_offset is None: + break + + try: + res = np.stack(res, axis=0) + except Exception as e: + print( + "Error during stacking stems: {} Track Length: {} Track path: {}".format( + str(e), track_length, track_path + ) + ) + res = np.zeros( + (len(self.instruments), 2, self.chunk_size), dtype=np.float32 + ) + if mix is None: + mix = res.sum(0) + if self.aug: + for i, instr in enumerate(self.instruments): + res[i] = self.augm_data(res[i], instr) + return torch.tensor(res, dtype=torch.float32), torch.tensor( + mix, dtype=torch.float32 + ) + + def augm_data(self, source, instr): + # source.shape = (2, 261120) - first channels, second length + source_shape = source.shape + applied_augs = [] + if "all" in self.config["augmentations"]: + augs = self.config["augmentations"]["all"] + else: + augs = dict() + + # We need to add to all augmentations specific augs for stem. And rewrite values if needed + if instr in self.config["augmentations"]: + for el in self.config["augmentations"][instr]: + augs[el] = self.config["augmentations"][instr][el] + + # Channel shuffle + if "channel_shuffle" in augs: + if augs["channel_shuffle"] > 0: + if random.uniform(0, 1) < augs["channel_shuffle"]: + source = source[::-1].copy() + applied_augs.append("channel_shuffle") + # Random inverse + if "random_inverse" in augs: + if augs["random_inverse"] > 0: + if random.uniform(0, 1) < augs["random_inverse"]: + source = source[:, ::-1].copy() + applied_augs.append("random_inverse") + # Random polarity (multiply -1) + if "random_polarity" in augs: + if augs["random_polarity"] > 0: + if random.uniform(0, 1) < augs["random_polarity"]: + source = -source.copy() + applied_augs.append("random_polarity") + # Random pitch shift + if "pitch_shift" in augs: + if augs["pitch_shift"] > 0: + if random.uniform(0, 1) < augs["pitch_shift"]: + apply_aug = AU.PitchShift( + min_semitones=augs["pitch_shift_min_semitones"], + max_semitones=augs["pitch_shift_max_semitones"], + p=1.0, + ) + source = apply_aug(samples=source, sample_rate=44100) + applied_augs.append("pitch_shift") + # Random seven band parametric eq + if "seven_band_parametric_eq" in augs: + if augs["seven_band_parametric_eq"] > 0: + if random.uniform(0, 1) < augs["seven_band_parametric_eq"]: + apply_aug = AU.SevenBandParametricEQ( + min_gain_db=augs["seven_band_parametric_eq_min_gain_db"], + max_gain_db=augs["seven_band_parametric_eq_max_gain_db"], + p=1.0, + ) + source = apply_aug(samples=source, sample_rate=44100) + applied_augs.append("seven_band_parametric_eq") + # Random tanh distortion + if "tanh_distortion" in augs: + if augs["tanh_distortion"] > 0: + if random.uniform(0, 1) < augs["tanh_distortion"]: + apply_aug = AU.TanhDistortion( + min_distortion=augs["tanh_distortion_min"], + max_distortion=augs["tanh_distortion_max"], + p=1.0, + ) + source = apply_aug(samples=source, sample_rate=44100) + applied_augs.append("tanh_distortion") + # Random MP3 Compression + if "mp3_compression" in augs: + if augs["mp3_compression"] > 0: + if random.uniform(0, 1) < augs["mp3_compression"]: + apply_aug = AU.Mp3Compression( + min_bitrate=augs["mp3_compression_min_bitrate"], + max_bitrate=augs["mp3_compression_max_bitrate"], + backend=augs["mp3_compression_backend"], + p=1.0, + ) + source = apply_aug(samples=source, sample_rate=44100) + applied_augs.append("mp3_compression") + # Random AddGaussianNoise + if "gaussian_noise" in augs: + if augs["gaussian_noise"] > 0: + if random.uniform(0, 1) < augs["gaussian_noise"]: + apply_aug = AU.AddGaussianNoise( + min_amplitude=augs["gaussian_noise_min_amplitude"], + max_amplitude=augs["gaussian_noise_max_amplitude"], + p=1.0, + ) + source = apply_aug(samples=source, sample_rate=44100) + applied_augs.append("gaussian_noise") + # Random TimeStretch + if "time_stretch" in augs: + if augs["time_stretch"] > 0: + if random.uniform(0, 1) < augs["time_stretch"]: + apply_aug = AU.TimeStretch( + min_rate=augs["time_stretch_min_rate"], + max_rate=augs["time_stretch_max_rate"], + leave_length_unchanged=True, + p=1.0, + ) + source = apply_aug(samples=source, sample_rate=44100) + applied_augs.append("time_stretch") + + # Possible fix of shape + if source_shape != source.shape: + source = source[..., : source_shape[-1]] + + # Random Reverb + if "pedalboard_reverb" in augs: + if augs["pedalboard_reverb"] > 0: + if random.uniform(0, 1) < augs["pedalboard_reverb"]: + room_size = random.uniform( + augs["pedalboard_reverb_room_size_min"], + augs["pedalboard_reverb_room_size_max"], + ) + damping = random.uniform( + augs["pedalboard_reverb_damping_min"], + augs["pedalboard_reverb_damping_max"], + ) + wet_level = random.uniform( + augs["pedalboard_reverb_wet_level_min"], + augs["pedalboard_reverb_wet_level_max"], + ) + dry_level = random.uniform( + augs["pedalboard_reverb_dry_level_min"], + augs["pedalboard_reverb_dry_level_max"], + ) + width = random.uniform( + augs["pedalboard_reverb_width_min"], + augs["pedalboard_reverb_width_max"], + ) + board = PB.Pedalboard( + [ + PB.Reverb( + room_size=room_size, # 0.1 - 0.9 + damping=damping, # 0.1 - 0.9 + wet_level=wet_level, # 0.1 - 0.9 + dry_level=dry_level, # 0.1 - 0.9 + width=width, # 0.9 - 1.0 + freeze_mode=0.0, + ) + ] + ) + source = board(source, 44100) + applied_augs.append("pedalboard_reverb") + + # Random Chorus + if "pedalboard_chorus" in augs: + if augs["pedalboard_chorus"] > 0: + if random.uniform(0, 1) < augs["pedalboard_chorus"]: + rate_hz = random.uniform( + augs["pedalboard_chorus_rate_hz_min"], + augs["pedalboard_chorus_rate_hz_max"], + ) + depth = random.uniform( + augs["pedalboard_chorus_depth_min"], + augs["pedalboard_chorus_depth_max"], + ) + centre_delay_ms = random.uniform( + augs["pedalboard_chorus_centre_delay_ms_min"], + augs["pedalboard_chorus_centre_delay_ms_max"], + ) + feedback = random.uniform( + augs["pedalboard_chorus_feedback_min"], + augs["pedalboard_chorus_feedback_max"], + ) + mix = random.uniform( + augs["pedalboard_chorus_mix_min"], + augs["pedalboard_chorus_mix_max"], + ) + board = PB.Pedalboard( + [ + PB.Chorus( + rate_hz=rate_hz, + depth=depth, + centre_delay_ms=centre_delay_ms, + feedback=feedback, + mix=mix, + ) + ] + ) + source = board(source, 44100) + applied_augs.append("pedalboard_chorus") + + # Random Phazer + if "pedalboard_phazer" in augs: + if augs["pedalboard_phazer"] > 0: + if random.uniform(0, 1) < augs["pedalboard_phazer"]: + rate_hz = random.uniform( + augs["pedalboard_phazer_rate_hz_min"], + augs["pedalboard_phazer_rate_hz_max"], + ) + depth = random.uniform( + augs["pedalboard_phazer_depth_min"], + augs["pedalboard_phazer_depth_max"], + ) + centre_frequency_hz = random.uniform( + augs["pedalboard_phazer_centre_frequency_hz_min"], + augs["pedalboard_phazer_centre_frequency_hz_max"], + ) + feedback = random.uniform( + augs["pedalboard_phazer_feedback_min"], + augs["pedalboard_phazer_feedback_max"], + ) + mix = random.uniform( + augs["pedalboard_phazer_mix_min"], + augs["pedalboard_phazer_mix_max"], + ) + board = PB.Pedalboard( + [ + PB.Phaser( + rate_hz=rate_hz, + depth=depth, + centre_frequency_hz=centre_frequency_hz, + feedback=feedback, + mix=mix, + ) + ] + ) + source = board(source, 44100) + applied_augs.append("pedalboard_phazer") + + # Random Distortion + if "pedalboard_distortion" in augs: + if augs["pedalboard_distortion"] > 0: + if random.uniform(0, 1) < augs["pedalboard_distortion"]: + drive_db = random.uniform( + augs["pedalboard_distortion_drive_db_min"], + augs["pedalboard_distortion_drive_db_max"], + ) + board = PB.Pedalboard( + [ + PB.Distortion( + drive_db=drive_db, + ) + ] + ) + source = board(source, 44100) + applied_augs.append("pedalboard_distortion") + + # Random PitchShift + if "pedalboard_pitch_shift" in augs: + if augs["pedalboard_pitch_shift"] > 0: + if random.uniform(0, 1) < augs["pedalboard_pitch_shift"]: + semitones = random.uniform( + augs["pedalboard_pitch_shift_semitones_min"], + augs["pedalboard_pitch_shift_semitones_max"], + ) + board = PB.Pedalboard([PB.PitchShift(semitones=semitones)]) + source = board(source, 44100) + applied_augs.append("pedalboard_pitch_shift") + + # Random Resample + if "pedalboard_resample" in augs: + if augs["pedalboard_resample"] > 0: + if random.uniform(0, 1) < augs["pedalboard_resample"]: + target_sample_rate = random.uniform( + augs["pedalboard_resample_target_sample_rate_min"], + augs["pedalboard_resample_target_sample_rate_max"], + ) + board = PB.Pedalboard( + [PB.Resample(target_sample_rate=target_sample_rate)] + ) + source = board(source, 44100) + applied_augs.append("pedalboard_resample") + + # Random Bitcrash + if "pedalboard_bitcrash" in augs: + if augs["pedalboard_bitcrash"] > 0: + if random.uniform(0, 1) < augs["pedalboard_bitcrash"]: + bit_depth = random.uniform( + augs["pedalboard_bitcrash_bit_depth_min"], + augs["pedalboard_bitcrash_bit_depth_max"], + ) + board = PB.Pedalboard([PB.Bitcrush(bit_depth=bit_depth)]) + source = board(source, 44100) + applied_augs.append("pedalboard_bitcrash") + + # Random MP3Compressor + if "pedalboard_mp3_compressor" in augs: + if augs["pedalboard_mp3_compressor"] > 0: + if random.uniform(0, 1) < augs["pedalboard_mp3_compressor"]: + vbr_quality = random.uniform( + augs["pedalboard_mp3_compressor_pedalboard_mp3_compressor_min"], + augs["pedalboard_mp3_compressor_pedalboard_mp3_compressor_max"], + ) + board = PB.Pedalboard([PB.MP3Compressor(vbr_quality=vbr_quality)]) + source = board(source, 44100) + applied_augs.append("pedalboard_mp3_compressor") + + # print(applied_augs) + return source diff --git a/src/third_party/MusicSourceSeparationTraining/utils/losses.py b/src/third_party/MusicSourceSeparationTraining/utils/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..443fbd1a9487b05db97ac40e34ccd92d6db3f4f1 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/utils/losses.py @@ -0,0 +1,403 @@ +import argparse +from typing import Any, Callable, Optional, Union + +import auraloss +import torch +import torch.nn.functional as F +from ml_collections import ConfigDict +from torch import nn +from torch_log_wmse import LogWMSE + + +def multistft_loss( + y_: torch.Tensor, + y: torch.Tensor, + loss_multistft: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], +) -> torch.Tensor: + """ + Compute a (multi-resolution) STFT-based loss on waveforms. + + Reshapes inputs to (B, C*T, L) when needed and delegates to a provided + multi-resolution STFT criterion (e.g., `auraloss.freq.MultiResolutionSTFTLoss`), + a widely used spectral loss for audio synthesis/enhancement that compares + magnitudes across multiple STFT settings. + See: Steinmetz & Reiss, 2020, “auraloss: Audio-focused loss functions in PyTorch”. + + Args: + y_ (torch.Tensor): Predicted waveform tensor of shape (B, C, T) or (B, S, C, T). + y (torch.Tensor): Target waveform tensor with a compatible shape. + loss_multistft (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): + A callable implementing the MR-STFT loss. + + Returns: + torch.Tensor: Scalar loss tensor. + """ + + if len(y_.shape) == 4: + y1_ = y_.reshape(y_.shape[0], y_.shape[1] * y_.shape[2], y_.shape[3]) + elif len(y_.shape) == 3: + y1_ = y_ + if len(y.shape) == 4: + y1 = y.reshape(y.shape[0], y.shape[1] * y.shape[2], y.shape[3]) + elif len(y_.shape) == 3: + y1 = y + if len(y_.shape) not in [3, 4]: + raise ValueError( + f"Invalid shape for predicted array: {y_.shape}. Expected 3 or 4 dimensions." + ) + return loss_multistft(y1_, y1) + + +def masked_loss( + y_: torch.Tensor, y: torch.Tensor, q: float, coarse: bool = True +) -> torch.Tensor: + """ + Robust, quantile-masked MSE (“trimmed” MSE). + + Computes an elementwise MSE, optionally averages spatial dims (“coarse”), + then masks out the largest residuals by keeping values below the `q`-quantile. + This yields robustness to outliers akin to trimmed/robust regression losses. + See classical robust estimation: Huber, 1964; Rousseeuw & Leroy, 1987. + + Args: + y_ (torch.Tensor): Predicted tensor matching `y`'s shape. + y (torch.Tensor): Ground-truth tensor. + q (float): Quantile in (0, 1] used to keep low-error elements. + coarse (bool, optional): If True, average over last two dims before masking. + Defaults to True. + + Returns: + torch.Tensor: Scalar loss tensor. + """ + + loss = torch.nn.MSELoss(reduction="none")(y_, y).transpose(0, 1) + if coarse: + loss = loss.mean(dim=(-1, -2)) + loss = loss.reshape(loss.shape[0], -1) + quantile = torch.quantile( + loss.detach(), q, interpolation="linear", dim=1, keepdim=True + ) + mask = loss < quantile + return (loss * mask).mean() + + +def spec_rmse_loss( + estimate: torch.Tensor, sources: torch.Tensor, stft_config: dict, eps: float = 1e-8 +) -> torch.Tensor: + """ + RMSE in the complex STFT domain. + + Computes STFT for prediction and target, represents complex values as + real+imag pairs, and applies RMSE (L2) over the spectral representation. + Spectral-domain L2/RMSE losses are common in speech/music enhancement. + See, e.g., Steinmetz & Reiss, 2020; Yamamoto et al., 2020 (Parallel WaveGAN). + + Args: + estimate (torch.Tensor): Predicted time-domain signal(s), e.g., (B, S, C, T). + sources (torch.Tensor): Target time-domain signal(s), matching shape. + stft_config (dict): Parameters for `torch.stft` (e.g., n_fft, hop_length, win_length). + + Returns: + torch.Tensor: Scalar loss tensor. + """ + + lenc = estimate.shape[-1] + spec_estimate = estimate.view(-1, lenc) + spec_sources = sources.view(-1, lenc) + + spec_estimate = torch.stft(spec_estimate, **stft_config, return_complex=True) + spec_sources = torch.stft(spec_sources, **stft_config, return_complex=True) + + spec_estimate = torch.view_as_real(spec_estimate) + spec_sources = torch.view_as_real(spec_sources) + + new_shape = estimate.shape[:-1] + spec_estimate.shape[-3:] + spec_estimate = spec_estimate.view(*new_shape) + spec_sources = spec_sources.view(*new_shape) + + loss = F.mse_loss(spec_estimate, spec_sources, reduction="none") + + dims = tuple(range(2, loss.dim())) + loss = (loss.mean(dims) + eps).sqrt().mean(dim=(0, 1)) + + return loss + + +def spec_masked_loss( + estimate: torch.Tensor, + sources: torch.Tensor, + stft_config: dict, + q: float = 0.9, + coarse: bool = True, +) -> torch.Tensor: + """ + Quantile-masked MSE in the complex STFT domain. + + Computes a complex STFT for prediction and target, forms an elementwise MSE + in the spectral domain, optionally averages spatial/frequency dims (“coarse”), + and masks out the highest-error elements using the `q`-quantile threshold for + robustness to outliers. Related to trimmed/robust spectral losses. + See: Huber, 1964; Rousseeuw & Leroy, 1987; spectral losses as in Steinmetz & Reiss, 2020. + + Args: + estimate (torch.Tensor): Predicted time-domain signal(s), e.g., (B, S, C, T). + sources (torch.Tensor): Target time-domain signal(s), matching shape. + stft_config (dict): Parameters for `torch.stft`. + q (float, optional): Quantile in (0, 1] to keep low-error elements. Defaults to 0.9. + coarse (bool, optional): If True, average over spectral dims before masking. Defaults to True. + + Returns: + torch.Tensor: Scalar loss tensor. + """ + + lenc = estimate.shape[-1] + spec_estimate = estimate.view(-1, lenc) + spec_sources = sources.view(-1, lenc) + + spec_estimate = torch.stft(spec_estimate, **stft_config, return_complex=True) + spec_sources = torch.stft(spec_sources, **stft_config, return_complex=True) + + spec_estimate = torch.view_as_real(spec_estimate) + spec_sources = torch.view_as_real(spec_sources) + + new_shape = estimate.shape[:-1] + spec_estimate.shape[-3:] + spec_estimate = spec_estimate.view(*new_shape) + spec_sources = spec_sources.view(*new_shape) + + loss = F.mse_loss(spec_estimate, spec_sources, reduction="none") + + if coarse: + loss = loss.mean(dim=(-3, -2)) + + loss = loss.reshape(loss.shape[0], -1) + + quantile = torch.quantile( + loss.detach(), q, interpolation="linear", dim=1, keepdim=True + ) + + mask = loss < quantile + + masked_loss = (loss * mask).mean() + + return masked_loss + + +def l1_snr_loss(y_: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """ + L1-SNR loss in time domain. + + L1-based signal-to-noise ratio loss (without additional regularization). + From torch-l1-snr package. + + Args: + y_ (torch.Tensor): Predicted waveform tensor. + y (torch.Tensor): Target waveform tensor. + + Returns: + torch.Tensor: Scalar loss tensor. + """ + from torch_l1_snr import L1SNRLoss + + loss_fn = L1SNRLoss(name="l1_snr") + return loss_fn(y_, y) + + +def l1_snr_db_loss(y_: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """ + L1-SNR loss with dB-scale level regularization. + + Extends L1-SNR with adaptive level-matching regularization in dB scale. + From torch-l1-snr package. + + Args: + y_ (torch.Tensor): Predicted waveform tensor. + y (torch.Tensor): Target waveform tensor. + + Returns: + torch.Tensor: Scalar loss tensor. + """ + from torch_l1_snr import L1SNRDBLoss + + loss_fn = L1SNRDBLoss(name="l1_snr_db") + return loss_fn(y_, y) + + +def stft_l1_snr_db_loss(y_: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """ + L1-SNR loss in multi-resolution STFT domain. + + Applies L1-SNR to complex STFT (real/imaginary) across multiple resolutions. + From torch-l1-snr package. + + Args: + y_ (torch.Tensor): Predicted waveform tensor. + y (torch.Tensor): Target waveform tensor. + + Returns: + torch.Tensor: Scalar loss tensor. + """ + from torch_l1_snr import STFTL1SNRDBLoss + + loss_fn = STFTL1SNRDBLoss(name="stft_l1_snr_db") + return loss_fn(y_, y) + + +def multi_l1_snr_db_loss(y_: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """ + Combined time + STFT domain L1-SNR loss. + + Balances time-domain and spectral-domain L1-SNR with optional regularization. + This is the recommended loss from torch-l1-snr for most use cases. + From torch-l1-snr package. + + Args: + y_ (torch.Tensor): Predicted waveform tensor. + y (torch.Tensor): Target waveform tensor. + + Returns: + torch.Tensor: Scalar loss tensor. + """ + from torch_l1_snr import MultiL1SNRDBLoss + + loss_fn = MultiL1SNRDBLoss(name="multi_l1_snr_db") + return loss_fn(y_, y) + + +def choice_loss( + args: argparse.Namespace, config: ConfigDict +) -> Callable[[Any, Any, Union[Any, None]], torch.Tensor]: + """ + Build a composite loss from CLI/config options. + + Returns a callable that sums enabled terms (with per-term coefficients): + - `masked_loss`: robust, quantile-masked MSE (trimmed MSE; Huber, 1964; Rousseeuw & Leroy, 1987). + - `mse_loss`: standard mean squared error. + - `l1_loss`: mean absolute error. + - `multistft_loss`: multi-resolution STFT magnitude loss (Steinmetz & Reiss, 2020). + - `log_wmse_loss`: weighted MSE operating in a log/spectral perceptual space (log-weighted MSE). + - `l1_snr_loss`: L1-SNR loss in time domain (Watcharasupat et al., 2023). + - `l1_snr_db_loss`: L1-SNR with dB-scale level regularization. + - `stft_l1_snr_db_loss`: L1-SNR in multi-resolution STFT domain. + - `multi_l1_snr_db_loss`: combined time + STFT domain L1-SNR (recommended). + - `spec_rmse_loss`: RMSE in complex STFT domain. + - `spec_masked_loss`: quantile-masked spectral MSE (robust spectral loss). + + Args: + args (argparse.Namespace): Parsed arguments specifying which losses are active + and their coefficients. + config (ConfigDict): Configuration with loss hyperparameters (e.g., STFT settings, + quantile `q`, coarse masking flag). + + Returns: + Callable[[Any, Any, Optional[Any]], torch.Tensor]: A function `loss(y_pred, y_true, x=None)` + that computes the weighted sum of the selected loss terms. + """ + + loss_fns = [] + + if "masked_loss" in args.loss: + loss_fns.append( + lambda y_pred, y_true, x=None: masked_loss( + y_pred, + y_true, + q=config["training"]["q"], + coarse=config["training"]["coarse_loss_clip"], + ) + * args.masked_loss_coef + ) + + if "mse_loss" in args.loss: + mse = nn.MSELoss() + loss_fns.append( + lambda y_pred, y_true, x=None: mse(y_pred, y_true) * args.mse_loss_coef + ) + + if "l1_loss" in args.loss: + loss_fns.append( + lambda y_pred, y_true, x=None: F.l1_loss(y_pred, y_true) * args.l1_loss_coef + ) + + if "multistft_loss" in args.loss: + loss_options = dict(config.get("loss_multistft", {})) + stft_loss = auraloss.freq.MultiResolutionSTFTLoss(**loss_options) + loss_fns.append( + lambda y_pred, y_true, x=None: multistft_loss(y_pred, y_true, stft_loss) + * args.multistft_loss_coef + ) + + if "log_wmse_loss" in args.loss: + log_wmse = LogWMSE( + audio_length=int(getattr(config.audio, "chunk_size", 485100)) + // int(getattr(config.audio, "sample_rate", 44100)), + sample_rate=int(getattr(config.audio, "sample_rate", 44100)), + return_as_loss=True, + bypass_filter=getattr(config.training, "bypass_filter", False), + ) + loss_fns.append( + lambda y_pred, y_true, x: log_wmse(x, y_pred, y_true) + * args.log_wmse_loss_coef + ) + + if "l1_snr_loss" in args.loss: + loss_fns.append( + lambda y_pred, y_true, x=None: l1_snr_loss(y_pred, y_true) + * args.l1_snr_loss_coef + ) + + if "l1_snr_db_loss" in args.loss: + loss_fns.append( + lambda y_pred, y_true, x=None: l1_snr_db_loss(y_pred, y_true) + * args.l1_snr_db_loss_coef + ) + + if "stft_l1_snr_db_loss" in args.loss: + loss_fns.append( + lambda y_pred, y_true, x=None: stft_l1_snr_db_loss(y_pred, y_true) + * args.stft_l1_snr_db_loss_coef + ) + + if "multi_l1_snr_db_loss" in args.loss: + loss_fns.append( + lambda y_pred, y_true, x=None: multi_l1_snr_db_loss(y_pred, y_true) + * args.multi_l1_snr_db_loss_coef + ) + + if "spec_rmse_loss" in args.loss: + stft_config = { + "n_fft": getattr(config.model, "nfft", 4096), + "hop_length": getattr(config.model, "hop_size", 1024), + "win_length": getattr(config.model, "win_size", 4096), + "center": True, + "normalized": getattr(config.model, "normalized", True), + } + loss_fns.append( + lambda y_pred, y_true, x=None: spec_rmse_loss(y_pred, y_true, stft_config) + * args.spec_rmse_loss_coef + ) + + if "spec_masked_loss" in args.loss: + stft_config = { + "n_fft": getattr(config.model, "nfft", 4096), + "hop_length": getattr(config.model, "hop_size", 1024), + "win_length": getattr(config.model, "win_size", 4096), + "center": True, + "normalized": getattr(config.model, "normalized", True), + } + loss_fns.append( + lambda y_pred, y_true, x=None: spec_masked_loss( + y_pred, + y_true, + stft_config, + q=config["training"]["q"], + coarse=config["training"]["coarse_loss_clip"], + ) + * args.spec_masked_loss_coef + ) + + def multi_loss(y_pred: Any, y_true: Any, x: Optional[Any] = None) -> torch.Tensor: + total = 0 + for fn in loss_fns: + total = total + fn(y_pred, y_true, x) + return total + + return multi_loss diff --git a/src/third_party/MusicSourceSeparationTraining/utils/metrics.py b/src/third_party/MusicSourceSeparationTraining/utils/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..d59556dcc999bcc6af8fde79777bec5cf518cbfb --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/utils/metrics.py @@ -0,0 +1,496 @@ +import math +from typing import Dict, List, Tuple + +import librosa +import numpy as np +import torch +import torch.nn.functional as F + + +def sdr(references: np.ndarray, estimates: np.ndarray) -> float: + """ + Compute Signal-to-Distortion Ratio (SDR) for one or more audio tracks. + + SDR is a measure of how well the predicted source (estimate) matches the reference source. + It is calculated as the ratio of the energy of the reference signal to the energy of the error (difference between reference and estimate). + Return SDR in decibels (dB) + Parameters: + ---------- + references : np.ndarray + A 3D numpy array of shape (num_sources, num_channels, num_samples), where num_sources is the number of sources, + num_channels is the number of channels (e.g., 1 for mono, 2 for stereo), and num_samples is the length of the audio signal. + + estimates : np.ndarray + A 3D numpy array of shape (num_sources, num_channels, num_samples) representing the estimated sources. + + Returns: + ------- + np.ndarray + A 1D numpy array containing the SDR values for each source. + """ + eps = 1e-8 # to avoid numerical errors + num = np.sum(np.square(references), axis=(1, 2)) + den = np.sum(np.square(references - estimates), axis=(1, 2)) + num += eps + den += eps + return 10 * np.log10(num / den) + + +def k_sdr(sdr: float, K: float = 10.0) -> float: + """Normalize SDR value using bounded logarithmic scaling.""" + sdr = max(min(sdr, K), -K + 1e-6) + return 100.0 * math.log1p(sdr + K) / math.log1p(2 * K) + + +def si_sdr(reference: np.ndarray, estimate: np.ndarray) -> float: + """ + Compute Scale-Invariant Signal-to-Distortion Ratio (SI-SDR) for one or more audio tracks. + + SI-SDR is a variant of the SDR metric that is invariant to the scaling of the estimate relative to the reference. + It is calculated by scaling the estimate to match the reference signal and then computing the SDR. + + Parameters: + ---------- + reference : np.ndarray + A 3D numpy array of shape (num_sources, num_channels, num_samples), where num_sources is the number of sources, + num_channels is the number of channels (e.g., 1 for mono, 2 for stereo), and num_samples is the length of the audio signal. + + estimate : np.ndarray + A 3D numpy array of shape (num_sources, num_channels, num_samples) representing the estimated sources. + + Returns: + ------- + float + The SI-SDR value for the source. It is a scalar representing the Signal-to-Distortion Ratio in decibels (dB). + """ + eps = 1e-8 # To avoid numerical errors + scale = np.sum(estimate * reference + eps, axis=(0, 1)) / np.sum( + reference**2 + eps, axis=(0, 1) + ) + scale = np.expand_dims(scale, axis=(0, 1)) # Reshape to [num_sources, 1] + + reference = reference * scale + si_sdr = np.mean( + 10 + * np.log10( + np.sum(reference**2, axis=(0, 1)) + / (np.sum((reference - estimate) ** 2, axis=(0, 1)) + eps) + + eps + ) + ) + + return si_sdr + + +def L1Freq_metric( + reference: np.ndarray, + estimate: np.ndarray, + fft_size: int = 2048, + hop_size: int = 1024, + device: str = "cpu", +) -> float: + """ + Compute the L1 Frequency Metric between the reference and estimated audio signals. + + This metric compares the magnitude spectrograms of the reference and estimated audio signals + using the Short-Time Fourier Transform (STFT) and calculates the L1 loss between them. The result + is scaled to the range [0, 100] where a higher value indicates better performance. + + Parameters: + ---------- + reference : np.ndarray + A 2D numpy array of shape (num_channels, num_samples) representing the reference (ground truth) audio signal. + + estimate : np.ndarray + A 2D numpy array of shape (num_channels, num_samples) representing the estimated (predicted) audio signal. + + fft_size : int, optional + The size of the FFT (Short-Time Fourier Transform). Default is 2048. + + hop_size : int, optional + The hop size between STFT frames. Default is 1024. + + device : str, optional + The device to run the computation on ('cpu' or 'cuda'). Default is 'cpu'. + + Returns: + ------- + float + The L1 Frequency Metric in the range [0, 100], where higher values indicate better performance. + """ + + reference = torch.from_numpy(reference).to(device) + estimate = torch.from_numpy(estimate).to(device) + + reference_stft = torch.stft(reference, fft_size, hop_size, return_complex=True) + estimated_stft = torch.stft(estimate, fft_size, hop_size, return_complex=True) + + reference_mag = torch.abs(reference_stft) + estimate_mag = torch.abs(estimated_stft) + + loss = 10 * F.l1_loss(estimate_mag, reference_mag) + + ret = 100 / (1.0 + float(loss.cpu().numpy())) + + return ret + + +def LogWMSE_metric( + reference: np.ndarray, + estimate: np.ndarray, + mixture: np.ndarray, + device: str = "cpu", +) -> float: + """ + Calculate the Log-WMSE (Logarithmic Weighted Mean Squared Error) between the reference, estimate, and mixture signals. + + This metric evaluates the quality of the estimated signal compared to the reference signal in the + context of audio source separation. The result is given in logarithmic scale, which helps in evaluating + signals with large amplitude differences. + + Parameters: + ---------- + reference : np.ndarray + The ground truth audio signal of shape (channels, time), where channels is the number of audio channels + (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. + + estimate : np.ndarray + The estimated audio signal of shape (channels, time). + + mixture : np.ndarray + The mixed audio signal of shape (channels, time). + + device : str, optional + The device to run the computation on, either 'cpu' or 'cuda'. Default is 'cpu'. + + Returns: + ------- + float + The Log-WMSE value, which quantifies the difference between the reference and estimated signal on a logarithmic scale. + """ + from torch_log_wmse import LogWMSE + + log_wmse = LogWMSE( + audio_length=reference.shape[-1] / 44100, # audio length in seconds + sample_rate=44100, # sample rate of 44100 Hz + return_as_loss=False, # return as loss (False means return as metric) + bypass_filter=False, # bypass frequency filtering (False means apply filter) + ) + + reference = torch.from_numpy(reference).unsqueeze(0).unsqueeze(0).to(device) + estimate = torch.from_numpy(estimate).unsqueeze(0).unsqueeze(0).to(device) + mixture = torch.from_numpy(mixture).unsqueeze(0).to(device) + + res = log_wmse(mixture, reference, estimate) + return float(res.cpu().numpy()) + + +def MultiL1SNRDB_metric( + reference: np.ndarray, + estimate: np.ndarray, + device: str = "cpu", +) -> float: + """ + Calculate L1-SNR metric (higher is better). + Returns negative loss value for scheduler compatibility (mode='max'). + """ + from torch_l1_snr import MultiL1SNRDBLoss + + l1_snr = MultiL1SNRDBLoss( + name="l1_snr_metric", + weight=1.0, + spec_weight=0.5, + l1_weight=0.0, + use_time_regularization=True, + use_spec_regularization=False, + ) + + reference_t = torch.from_numpy(reference).unsqueeze(0).to(device) + estimate_t = torch.from_numpy(estimate).unsqueeze(0).to(device) + + with torch.no_grad(): + res = l1_snr(estimate_t, reference_t) + + return -float(res.cpu().numpy()) + + +def AuraSTFT_metric( + reference: np.ndarray, + estimate: np.ndarray, + device: str = "cpu", +) -> float: + """ + Calculate the AuraSTFT metric, which evaluates the spectral difference between the reference and estimated + audio signals using Short-Time Fourier Transform (STFT) loss. + + The AuraSTFT metric computes the STFT loss in both logarithmic and linear magnitudes, and it is commonly used + to assess the quality of audio separation tasks. The result is returned as a value scaled to the range [0, 100]. + + Parameters: + ---------- + reference : np.ndarray + The ground truth audio signal of shape (channels, time), where channels is the number of audio channels + (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. + + estimate : np.ndarray + The estimated audio signal of shape (channels, time). + + device : str, optional + The device to run the computation on, either 'cpu' or 'cuda'. Default is 'cpu'. + + Returns: + ------- + float + The AuraSTFT metric value, scaled to the range [0, 100], which quantifies the difference between + the reference and estimated signal in the spectral domain. + """ + + from auraloss.freq import STFTLoss + + stft_loss = STFTLoss( + w_log_mag=1.0, # weight for log magnitude + w_lin_mag=0.0, # weight for linear magnitude + w_sc=1.0, # weight for spectral centroid + device=device, + ) + + reference = torch.from_numpy(reference).unsqueeze(0).to(device) + estimate = torch.from_numpy(estimate).unsqueeze(0).to(device) + + res = 100 / (1.0 + 10 * stft_loss(reference, estimate)) + return float(res.cpu().numpy()) + + +def AuraMRSTFT_metric( + reference: np.ndarray, + estimate: np.ndarray, + device: str = "cpu", +) -> float: + """ + Calculate the AuraMRSTFT metric, which evaluates the spectral difference between the reference and estimated + audio signals using Multi-Resolution Short-Time Fourier Transform (STFT) loss. + + The AuraMRSTFT metric uses multi-resolution STFT analysis, which allows better representation of both + low- and high-frequency components in the audio signals. The result is returned as a value scaled to the range [0, 100]. + + Parameters: + ---------- + reference : np.ndarray + The ground truth audio signal of shape (channels, time), where channels is the number of audio channels + (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. + + estimate : np.ndarray + The estimated audio signal of shape (channels, time). + + device : str, optional + The device to run the computation on, either 'cpu' or 'cuda'. Default is 'cpu'. + + Returns: + ------- + float + The AuraMRSTFT metric value, scaled to the range [0, 100], which quantifies the difference between + the reference and estimated signal in the multi-resolution spectral domain. + """ + + from auraloss.freq import MultiResolutionSTFTLoss + + mrstft_loss = MultiResolutionSTFTLoss( + fft_sizes=[1024, 2048, 4096], + hop_sizes=[256, 512, 1024], + win_lengths=[1024, 2048, 4096], + scale="mel", # mel scale for frequency resolution + n_bins=128, # number of bins for mel scale + sample_rate=44100, + perceptual_weighting=True, # apply perceptual weighting + device=device, + ) + + reference = torch.from_numpy(reference).unsqueeze(0).float().to(device) + estimate = torch.from_numpy(estimate).unsqueeze(0).float().to(device) + + res = 100 / (1.0 + 10 * mrstft_loss(reference, estimate)) + return float(res.cpu().numpy()) + + +def bleed_full( + reference: np.ndarray, + estimate: np.ndarray, + sr: int = 44100, + n_fft: int = 4096, + hop_length: int = 1024, + n_mels: int = 512, + device: str = "cpu", +) -> Tuple[float, float]: + """ + Calculate the 'bleed' and 'fullness' metrics between a reference and an estimated audio signal. + + The 'bleed' metric measures how much the estimated signal bleeds into the reference signal, + while the 'fullness' metric measures how much the estimated signal retains its distinctiveness + in relation to the reference signal, both using mel spectrograms and decibel scaling. + + Parameters: + ---------- + reference : np.ndarray + The reference audio signal, shape (channels, time), where channels is the number of audio channels + (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. + + estimate : np.ndarray + The estimated audio signal, shape (channels, time). + + sr : int, optional + The sample rate of the audio signals. Default is 44100 Hz. + + n_fft : int, optional + The FFT size used to compute the STFT. Default is 4096. + + hop_length : int, optional + The hop length for STFT computation. Default is 1024. + + n_mels : int, optional + The number of mel frequency bins. Default is 512. + + device : str, optional + The device for computation, either 'cpu' or 'cuda'. Default is 'cpu'. + + Returns: + ------- + tuple + A tuple containing two values: + - `bleedless` (float): A score indicating how much 'bleeding' the estimated signal has (higher is better). + - `fullness` (float): A score indicating how 'full' the estimated signal is (higher is better). + """ + + from torchaudio.transforms import AmplitudeToDB + + reference = torch.from_numpy(reference).float().to(device) + estimate = torch.from_numpy(estimate).float().to(device) + + window = torch.hann_window(n_fft).to(device) + + # Compute STFTs with the Hann window + D1 = torch.abs( + torch.stft( + reference, + n_fft=n_fft, + hop_length=hop_length, + window=window, + return_complex=True, + pad_mode="constant", + ) + ) + D2 = torch.abs( + torch.stft( + estimate, + n_fft=n_fft, + hop_length=hop_length, + window=window, + return_complex=True, + pad_mode="constant", + ) + ) + + mel_basis = librosa.filters.mel(sr=sr, n_fft=n_fft, n_mels=n_mels) + mel_filter_bank = torch.from_numpy(mel_basis).to(device) + + S1_mel = torch.matmul(mel_filter_bank, D1) + S2_mel = torch.matmul(mel_filter_bank, D2) + + S1_db = AmplitudeToDB(stype="magnitude", top_db=80)(S1_mel) + S2_db = AmplitudeToDB(stype="magnitude", top_db=80)(S2_mel) + + diff = S2_db - S1_db + + positive_diff = diff[diff > 0] + negative_diff = diff[diff < 0] + + average_positive = ( + torch.mean(positive_diff) + if positive_diff.numel() > 0 + else torch.tensor(0.0).to(device) + ) + average_negative = ( + torch.mean(negative_diff) + if negative_diff.numel() > 0 + else torch.tensor(0.0).to(device) + ) + + bleedless = 100 * 1 / (average_positive + 1) + fullness = 100 * 1 / (-average_negative + 1) + + return bleedless.cpu().numpy(), fullness.cpu().numpy() + + +def get_metrics( + metrics: List[str], + reference: np.ndarray, + estimate: np.ndarray, + mix: np.ndarray, + device: str = "cpu", + k: float = 10, +) -> Dict[str, float]: + """ + Calculate a list of metrics to evaluate the performance of audio source separation models. + + The function computes the specified metrics based on the reference, estimate, and mixture. + + Parameters: + ---------- + metrics : List[str] + A list of metric names to compute (e.g., ['sdr', 'si_sdr', 'l1_freq']). + + reference : np.ndarray + The reference audio (true signal) with shape (channels, length). + + estimate : np.ndarray + The estimated audio (predicted signal) with shape (channels, length). + + mix : np.ndarray + The mixed audio signal with shape (channels, length). + + device : str, optional, default='cpu' + The device ('cpu' or 'cuda') to perform the calculations on. + + Returns: + ------- + Dict[str, float] + A dictionary containing the computed metric values. + """ + result = dict() + + # Adjust the length to be the same across all inputs + min_length = min(reference.shape[1], estimate.shape[1]) + reference = reference[..., :min_length] + estimate = estimate[..., :min_length] + mix = mix[..., :min_length] + + if "sdr" in metrics or "k_sdr" in metrics: + references = np.expand_dims(reference, axis=0) + estimates = np.expand_dims(estimate, axis=0) + result["sdr"] = float(sdr(references, estimates)) + result["k_sdr"] = k_sdr(float(sdr(references, estimates)), k) + if "si_sdr" in metrics: + result["si_sdr"] = float(si_sdr(reference, estimate)) + + if "l1_freq" in metrics: + result["l1_freq"] = L1Freq_metric(reference, estimate, device=device) + + if "log_wmse" in metrics: + result["log_wmse"] = LogWMSE_metric(reference, estimate, mix, device) + + if "aura_stft" in metrics: + result["aura_stft"] = AuraSTFT_metric(reference, estimate, device) + + if "aura_mrstft" in metrics: + result["aura_mrstft"] = AuraMRSTFT_metric(reference, estimate, device) + + if "l1_snr" in metrics: + result["l1_snr"] = MultiL1SNRDB_metric(reference, estimate, device) + + if "bleedless" in metrics or "fullness" in metrics: + bleedless, fullness = bleed_full(reference, estimate, device=device) + if "bleedless" in metrics: + result["bleedless"] = float(bleedless) + if "fullness" in metrics: + result["fullness"] = float(fullness) + + return result diff --git a/src/third_party/MusicSourceSeparationTraining/utils/model_utils.py b/src/third_party/MusicSourceSeparationTraining/utils/model_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4a8d5a962d8f7865733ea7378be4952e4a1880e3 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/utils/model_utils.py @@ -0,0 +1,904 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" + +import argparse +import json +import os +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple, Union + +import loralib as lora +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +from ml_collections import ConfigDict +from torch.optim import Adam, AdamW, RAdam, RMSprop +from tqdm.auto import tqdm + + +def demix( + config: ConfigDict, + model: torch.nn.Module, + mix: torch.Tensor, + device: torch.device, + model_type: str, + pbar: bool = False, +) -> Union[Dict[str, np.ndarray], np.ndarray]: + """ + Perform audio source separation with a given model. + + Supports both Demucs-specific and generic processing modes, including + overlapping chunk-based inference with optional progress bar display. + Handles padding, fading, and batching to reduce artifacts during separation. + + Args: + config (ConfigDict): Configuration object with audio and inference + parameters (chunk size, overlap, batch size, etc.). + model (torch.nn.Module): Source separation model for inference. + mix (torch.Tensor): Input audio tensor of shape (channels, time). + device (torch.device): Device on which to run inference (CPU or CUDA). + model_type (str): Type of model (e.g., 'htdemucs', 'mdx23c') that + determines processing mode. + pbar (bool, optional): If True, show a progress bar during chunk + processing. Defaults to False. + + Returns: + Union[Dict[str, np.ndarray], np.ndarray]: + - Dictionary mapping instrument names to separated waveforms if + multiple instruments are predicted. + - NumPy array of separated audio if only a single instrument is + present (Demucs mode). + """ + + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + mix = torch.tensor(mix, dtype=torch.float32) + + if model_type == "htdemucs": + mode = "demucs" + else: + mode = "generic" + # Define processing parameters based on the mode + if mode == "demucs": + chunk_size = config.training.samplerate * config.training.segment + num_instruments = len(config.training.instruments) + num_overlap = config.inference.num_overlap + step = chunk_size // num_overlap + else: + if "chunk_size" in config.inference: + chunk_size = config.inference.chunk_size + else: + chunk_size = config.audio.chunk_size + num_instruments = len(prefer_target_instrument(config)) + num_overlap = config.inference.num_overlap + + fade_size = chunk_size // 10 + step = chunk_size // num_overlap + border = chunk_size - step + length_init = mix.shape[-1] + windowing_array = _getWindowingArray(chunk_size, fade_size) + # Add padding for generic mode to handle edge artifacts + if length_init > 2 * border and border > 0: + mix = nn.functional.pad(mix, (border, border), mode="reflect") + + batch_size = config.inference.batch_size + + use_amp = getattr(config.training, "use_amp", True) + + with torch.cuda.amp.autocast(enabled=use_amp): + with torch.inference_mode(): + # Initialize result and counter tensors + req_shape = (num_instruments,) + mix.shape + result = torch.zeros(req_shape, dtype=torch.float32) + counter = torch.zeros(req_shape, dtype=torch.float32) + + i = 0 + batch_data = [] + batch_locations = [] + if pbar and should_print: + progress_bar = tqdm( + total=mix.shape[1], desc="Processing audio chunks", leave=False + ) + else: + progress_bar = None + + while i < mix.shape[1]: + # Extract chunk and apply padding if necessary + part = mix[:, i : i + chunk_size].to(device) + chunk_len = part.shape[-1] + if mode == "generic" and chunk_len > chunk_size // 2: + pad_mode = "reflect" + else: + pad_mode = "constant" + part = nn.functional.pad( + part, (0, chunk_size - chunk_len), mode=pad_mode, value=0 + ) + + batch_data.append(part) + batch_locations.append((i, chunk_len)) + i += step + + # Process batch if it's full or the end is reached + if len(batch_data) >= batch_size or i >= mix.shape[1]: + arr = torch.stack(batch_data, dim=0) + x = model(arr) + + if mode == "generic": + window = windowing_array.clone() # using clone() fixes the clicks at chunk edges when using batch_size=1 + if i - step == 0: # First audio chunk, no fadein + window[:fade_size] = 1 + elif i >= mix.shape[1]: # Last audio chunk, no fadeout + window[-fade_size:] = 1 + + for j, (start, seg_len) in enumerate(batch_locations): + if mode == "generic": + result[..., start : start + seg_len] += ( + x[j, ..., :seg_len].cpu() * window[..., :seg_len] + ) + counter[..., start : start + seg_len] += window[ + ..., :seg_len + ] + else: + result[..., start : start + seg_len] += x[ + j, ..., :seg_len + ].cpu() + counter[..., start : start + seg_len] += 1.0 + + batch_data.clear() + batch_locations.clear() + + if progress_bar: + progress_bar.update(step) + + if progress_bar: + progress_bar.close() + + # Compute final estimated sources + estimated_sources = result / counter + estimated_sources = estimated_sources.cpu().numpy() + np.nan_to_num(estimated_sources, copy=False, nan=0.0) + + # Remove padding for generic mode + if mode == "generic": + if length_init > 2 * border and border > 0: + estimated_sources = estimated_sources[..., border:-border] + + # Return the result as a dictionary or a single array + if mode == "demucs": + instruments = config.training.instruments + else: + instruments = prefer_target_instrument(config) + + ret_data = {k: v for k, v in zip(instruments, estimated_sources)} + + if mode == "demucs" and num_instruments <= 1: + return estimated_sources + else: + return ret_data + + +def initialize_model_and_device( + model: torch.nn.Module, device_ids: List[int] +) -> Tuple[Union[torch.device, str], torch.nn.Module]: + """ + Move a model to the correct computation device and wrap with DataParallel if needed. + + Selects GPU(s) if CUDA is available; otherwise defaults to CPU. If multiple + GPU IDs are provided, wraps the model with `nn.DataParallel` for multi-GPU + execution. + + Args: + model (torch.nn.Module): PyTorch model to be initialized. + device_ids (List[int]): List of GPU device IDs to use. If length > 1, + the model will be wrapped with DataParallel. + + Returns: + Tuple[Union[torch.device, str], torch.nn.Module]: A tuple containing: + - The computation device (`torch.device` or "cpu"). + - The model moved to that device (wrapped in DataParallel if applicable). + """ + + if torch.cuda.is_available(): + if len(device_ids) <= 1: + device = torch.device(f"cuda:{device_ids[0]}") + model = model.to(device) + else: + device = torch.device(f"cuda:{device_ids[0]}") + model = nn.DataParallel(model, device_ids=device_ids).to(device) + else: + device = "cpu" + model = model.to(device) + print("CUDA is not available. Running on CPU.") + + return device, model + + +def get_optimizer(config: ConfigDict, model: torch.nn.Module) -> torch.optim.Optimizer: + """ + Create and configure an optimizer for training. + + Selects the optimizer type based on `config.training.optimizer` and applies + the corresponding parameters, including support for advanced optimizers + such as Muon, Prodigy, and 8-bit AdamW. Handles parameter group separation + for specialized optimizers (e.g., Muon vs. Adam parameters). + + Args: + config (ConfigDict): Training configuration containing optimizer type, + learning rate, and optional optimizer-specific parameters. + model (torch.nn.Module): Model whose parameters will be optimized. + + Returns: + torch.optim.Optimizer: Initialized optimizer ready for training. + + Raises: + ValueError: If required optimizer configuration is missing (e.g., for Muon). + SystemExit: If an unknown optimizer name is encountered. + """ + + should_print = not dist.is_initialized() or dist.get_rank() == 0 + optim_params = dict() + if "optimizer" in config: + optim_params = dict(config["optimizer"]) + if config.training.optimizer != "muon" and should_print: + print(f"Optimizer params from config:\n{optim_params}") + + name_optimizer = getattr(config.training, "optimizer", "No optimizer in config") + + if name_optimizer == "adam": + optimizer = Adam(model.parameters(), lr=config.training.lr, **optim_params) + elif name_optimizer == "adamw": + optimizer = AdamW(model.parameters(), lr=config.training.lr, **optim_params) + elif name_optimizer == "radam": + optimizer = RAdam(model.parameters(), lr=config.training.lr, **optim_params) + elif name_optimizer == "rmsprop": + optimizer = RMSprop(model.parameters(), lr=config.training.lr, **optim_params) + elif name_optimizer == "prodigy": + from prodigyopt import Prodigy + + # you can choose weight decay value based on your problem, 0 by default + # We recommend using lr=1.0 (default) for all networks. + optimizer = Prodigy(model.parameters(), lr=config.training.lr, **optim_params) + elif name_optimizer == "adamw8bit": + import bitsandbytes as bnb + + optimizer = bnb.optim.AdamW8bit( + model.parameters(), lr=config.training.lr, **optim_params + ) + elif name_optimizer == "muon": + from .muon import Muon as Muon + + if should_print: + print("Using Muon optimizer with AdamW-like branch for non-muon params.") + muon_params = [p for p in model.parameters() if p.ndim >= 2] + adam_params = [p for p in model.parameters() if p.ndim < 2] + + if ( + not hasattr(config, "optimizer") + or "muon_group" not in config.optimizer + or "adam_group" not in config.optimizer + ): + raise ValueError( + "For the 'muon' optimizer, the config must have an 'optimizer' section " + "with 'muon_group' and 'adam_group' dictionaries." + ) + + muon_group_config = dict(config.optimizer.muon_group) + adam_group_config = dict(config.optimizer.adam_group) + + muon_group_config.setdefault("weight_decouple", True) + adam_group_config.setdefault("weight_decouple", True) + + if should_print: + print(f"Muon group params: {muon_group_config}") + print(f"Adam group params: {adam_group_config}") + + param_groups = [ + dict(params=muon_params, use_muon=True, **muon_group_config), + dict(params=adam_params, use_muon=False, **adam_group_config), + ] + optimizer = Muon(param_groups) + elif name_optimizer == "adago": + from .muon import AdaGO as AdaGO + + if should_print: + print("Using AdaGO optimizer with AdamW-like branch for non-muon params.") + muon_params = [p for p in model.parameters() if p.ndim >= 2] + adam_params = [p for p in model.parameters() if p.ndim < 2] + + if ( + not hasattr(config, "optimizer") + or "muon_group" not in config.optimizer + or "adam_group" not in config.optimizer + ): + raise ValueError( + "For 'adago', the config must have an 'optimizer' section with 'muon_group' and 'adam_group' dictionaries." + ) + + muon_group_config = dict(config.optimizer.muon_group) + adam_group_config = dict(config.optimizer.adam_group) + + muon_group_config.setdefault("weight_decouple", True) + adam_group_config.setdefault("weight_decouple", True) + + if should_print: + print(f"AdaGO muon group params: {muon_group_config}") + print(f"AdaGO adam group params: {adam_group_config}") + + param_groups = [ + dict(params=muon_params, use_muon=True, **muon_group_config), + dict(params=adam_params, use_muon=False, **adam_group_config), + ] + optimizer = AdaGO(param_groups) + return optimizer + + +def normalize_batch( + x: torch.Tensor, y: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply mean-variance normalization to a pair of tensors. + + Computes the mean and standard deviation from `x` and normalizes both `x` + and `y` using those statistics. This ensures the two tensors are scaled + consistently. + + Args: + x (torch.Tensor): Input tensor used to compute normalization statistics. + y (torch.Tensor): Input tensor normalized using the same statistics as `x`. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Normalized tensors `(x, y)`. + """ + + mean = x.mean() + std = x.std() + if std != 0: + x = (x - mean) / std + y = (y - mean) / std + return x, y + + +def apply_tta( + config, + model: torch.nn.Module, + mix: torch.Tensor, + waveforms_orig: Union[dict[str, np.ndarray], np.ndarray], + device: torch.device, + model_type: str, +) -> Union[dict[str, np.ndarray], np.ndarray]: + """ + Enhance source separation results using Test-Time Augmentation (TTA). + + Applies augmentations such as channel reversal and polarity inversion to + the input mixture, reprocesses with the model, and combines the results + with the original predictions by averaging. + + Args: + config: Configuration object with model and inference parameters. + model (torch.nn.Module): Trained source separation model. + mix (torch.Tensor): Input mixture tensor of shape (channels, time). + waveforms_orig (Dict[str, torch.Tensor]): Dictionary of separated + sources before augmentation. + device (torch.device): Computation device (CPU or CUDA). + model_type (str): Model type identifier used for demixing. + + Returns: + Dict[str, torch.Tensor]: Dictionary of separated sources after applying TTA. + """ + + # Create augmentations: channel inversion and polarity inversion + track_proc_list = [mix[::-1].copy(), -1.0 * mix.copy()] + + # Process each augmented mixture + for i, augmented_mix in enumerate(track_proc_list): + waveforms = demix(config, model, augmented_mix, device, model_type=model_type) + for el in waveforms: + if i == 0: + waveforms_orig[el] += waveforms[el][::-1].copy() + else: + waveforms_orig[el] -= waveforms[el] + + # Average the results across augmentations + for el in waveforms_orig: + waveforms_orig[el] /= len(track_proc_list) + 1 + + return waveforms_orig + + +def _getWindowingArray(window_size: int, fade_size: int) -> torch.Tensor: + """ + Generate a windowing array with a linear fade-in at the beginning and a fade-out at the end. + + This function creates a window of size `window_size` where the first `fade_size` elements + linearly increase from 0 to 1 (fade-in) and the last `fade_size` elements linearly decrease + from 1 to 0 (fade-out). The middle part of the window is filled with ones. + + Parameters: + ---------- + window_size : int + The total size of the window. + fade_size : int + The size of the fade-in and fade-out regions. + + Returns: + ------- + torch.Tensor + A tensor of shape (window_size,) containing the generated windowing array. + + Example: + ------- + If `window_size=10` and `fade_size=3`, the output will be: + tensor([0.0000, 0.5000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.5000, 0.0000]) + """ + + fadein = torch.linspace(0, 1, fade_size) + fadeout = torch.linspace(1, 0, fade_size) + + window = torch.ones(window_size) + window[-fade_size:] = fadeout + window[:fade_size] = fadein + return window + + +def prefer_target_instrument(config: ConfigDict) -> List[str]: + """ + Return the list of target instruments based on the configuration. + If a specific target instrument is specified in the configuration, + it returns a list with that instrument. Otherwise, it returns the list of instruments. + + Parameters: + ---------- + config : ConfigDict + Configuration object containing the list of instruments or the target instrument. + + Returns: + ------- + List[str] + A list of target instruments. + """ + if getattr(config.training, "target_instrument", None): + return [config.training.target_instrument] + else: + return config.training.instruments + + +def load_not_compatible_weights( + model: torch.nn.Module, old_model: dict, verbose: bool = False +) -> None: + """ + Load a possibly incompatible state dict into `model` with best-effort matching. + + Accepts either a raw state_dict or a checkpoint dict with weights under "state" or "state_dict". + For each param/buffer in `model`: if the name exists and shapes match → copy; + if ndim matches but shapes differ → zero-pad/crop the source to fit the target; + if the name is missing or ndim differs → skip. Optional logging on rank 0 when `verbose=True`. + + Args: + model: Target PyTorch module. + old_model: Source weights (state_dict or checkpoint dict). + verbose: Print brief load decisions. + + Returns: + None + """ + + should_print = verbose and (not dist.is_initialized() or dist.get_rank() == 0) + + new_model = model.state_dict() + + if "state" in old_model: + # Fix for htdemucs weights loading + old_model = old_model["state"] + if "state_dict" in old_model: + # Fix for apollo weights loading + old_model = old_model["state_dict"] + if "model_state_dict" in old_model: + # Fix for full_check_point + old_model = old_model["model_state_dict"] + + for el in new_model: + if el in old_model: + if should_print: + print(f"Match found for {el}!") + if new_model[el].shape == old_model[el].shape: + if should_print: + print("Action: Just copy weights!") + new_model[el] = old_model[el] + else: + if ( + len(new_model[el].shape) != len(old_model[el].shape) + and should_print + ): + print( + "Action: Different dimension! Too lazy to write the code... Skip it" + ) + else: + if should_print: + print( + f"Shape is different: {tuple(new_model[el].shape)} != {tuple(old_model[el].shape)}" + ) + ln = len(new_model[el].shape) + max_shape = [] + slices_old = [] + slices_new = [] + for i in range(ln): + max_shape.append( + max(new_model[el].shape[i], old_model[el].shape[i]) + ) + slices_old.append(slice(0, old_model[el].shape[i])) + slices_new.append(slice(0, new_model[el].shape[i])) + # print(max_shape) + # print(slices_old, slices_new) + slices_old = tuple(slices_old) + slices_new = tuple(slices_new) + max_matrix = np.zeros(max_shape, dtype=np.float32) + for i in range(ln): + max_matrix[slices_old] = old_model[el].cpu().numpy() + max_matrix = torch.from_numpy(max_matrix) + new_model[el] = max_matrix[slices_new] + else: + if should_print: + print(f"Match not found for {el}!") + model.load_state_dict(new_model) + + +def load_lora_weights( + model: torch.nn.Module, lora_path: str, device: str = "cpu" +) -> None: + """ + Load LoRA weights into a model. + This function updates the given model with LoRA-specific weights from the specified checkpoint file. + It does not require the checkpoint to match the model's full state dictionary, as only LoRA layers are updated. + + Parameters: + ---------- + model : Module + The PyTorch model into which the LoRA weights will be loaded. + lora_path : str + Path to the LoRA checkpoint file. + device : str, optional + The device to load the weights onto, by default 'cpu'. Common values are 'cpu' or 'cuda'. + + Returns: + ------- + None + The model is updated in place. + """ + lora_state_dict = torch.load(lora_path, map_location=device) + model.load_state_dict(lora_state_dict, strict=False) + + +def get_lora(args, config, model): + if args.train_lora_loralib: + model = bind_lora_to_model(config, model) + lora.mark_only_lora_as_trainable(model) + if args.train_lora_peft: + if args.lora_checkpoint_peft: + from peft import PeftModel + + model = PeftModel.from_pretrained(model, args.lora_checkpoint_peft) + for name, param in model.named_parameters(): + if "lora" in name.lower(): + param.requires_grad = True + else: + from peft import LoraConfig, get_peft_model + + lora_config = LoraConfig(**config["lora"]) + model = get_peft_model(model, lora_config) + return model + + +def load_start_checkpoint( + args: argparse.Namespace, model: torch.nn.Module, old_model, type_: str = "train" +) -> None: + """ + Load an initial checkpoint into `model`. + + For `type_ == "train"`, performs a tolerant load using `old_model` (a state dict or a + checkpoint dict) via `load_not_compatible_weights`, allowing partial shape mismatches. + For other modes, loads a strict state dict from `args.start_check_point`, with special + handling for HTDemucs/Apollo checkpoints (keys under "state"/"state_dict"). If + `args.lora_checkpoint` is set, LoRA weights are applied after the base load. + + Args: + args: Namespace with at least `start_check_point`, `model_type`, and optionally `lora_checkpoint`. + model: Target PyTorch module to receive weights. + old_model: Source weights for tolerant loading in train mode (state dict or checkpoint dict). + type_: Loading strategy; "train" uses tolerant loading, otherwise strict loading from path. + + Returns: + None + """ + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + if should_print: + print(f"Start from checkpoint: {args.start_check_point}") + if type_ in ["train"]: + if not args.load_only_compatible_weights: + load_not_compatible_weights(model, old_model, verbose=False) + else: + model.load_state_dict(torch.load(args.start_check_point)) + else: + device = "cpu" + if args.model_type in ["htdemucs", "apollo"]: + old_model = torch.load( + args.start_check_point, map_location=device, weights_only=False + ) + # Fix for htdemucs pretrained models + if "state" in old_model: + old_model = old_model["state"] + # Fix for apollo pretrained models + if "state_dict" in old_model: + old_model = old_model["state_dict"] + else: + if "state" in old_model: + # Fix for htdemucs weights loading + old_model = old_model["state"] + if "state_dict" in old_model: + # Fix for apollo weights loading + old_model = old_model["state_dict"] + if "model_state_dict" in old_model: + # Fix for full_check_point + old_model = old_model["model_state_dict"] + model.load_state_dict(old_model) + + if args.lora_checkpoint_loralib: + if should_print: + print(f"Loading LoRA weights from: {args.lora_checkpoint_loralib}") + load_lora_weights(model, args.lora_checkpoint_loralib) + + +def bind_lora_to_model(config: Dict[str, Any], model: nn.Module) -> nn.Module: + """ + Replaces specific layers in the model with LoRA-extended versions. + + Parameters: + ---------- + config : Dict[str, Any] + Configuration containing parameters for LoRA. It should include a 'lora' key with parameters for `MergedLinear`. + model : nn.Module + The original model in which the layers will be replaced. + + Returns: + ------- + nn.Module + The modified model with the replaced layers. + """ + + if "lora" not in config: + raise ValueError( + "Configuration must contain the 'lora' key with parameters for LoRA." + ) + + replaced_layers = 0 # Counter for replaced layers + should_print = not dist.is_initialized() or dist.get_rank() == 0 + + for name, module in model.named_modules(): + hierarchy = name.split(".") + layer_name = hierarchy[-1] + + # Check if this is the target layer to replace (and layer_name == 'to_qkv') + if isinstance(module, nn.Linear): + try: + # Get the parent module + parent_module = model + for submodule_name in hierarchy[:-1]: + parent_module = getattr(parent_module, submodule_name) + + # Replace the module with LoRA-enabled layer + setattr( + parent_module, + layer_name, + lora.Linear( + in_features=module.in_features, + out_features=module.out_features, + bias=module.bias is not None, + **config["lora"], + ), + ) + replaced_layers += 1 # Increment the counter + + except Exception as e: + if should_print: + print(f"Error replacing layer {name}: {e}") + + if replaced_layers == 0 and should_print: + print( + "Warning: No layers were replaced. Check the model structure and configuration." + ) + elif should_print: + print(f"Number of layers replaced with LoRA: {replaced_layers}") + + return model + + +def log_model_info(model: torch.nn.Module, results_path=None): + """Log comprehensive model information""" + model_info = { + "timestamp": datetime.now().isoformat(), + "model_class": model.__class__.__name__, + "model_module": model.__class__.__module__, + } + + # Count parameters + total_params = sum(p.numel() for p in model.parameters()) + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + + model_info["parameters"] = { + "total": total_params, + "trainable": trainable_params, + "non_trainable": total_params - trainable_params, + "total_millions": round(total_params / 1e6, 2), + "trainable_millions": round(trainable_params / 1e6, 2), + } + + # Get model size in memory + param_size = 0 + buffer_size = 0 + + for param in model.parameters(): + param_size += param.nelement() * param.element_size() + + for buffer in model.buffers(): + buffer_size += buffer.nelement() * buffer.element_size() + + model_size_mb = (param_size + buffer_size) / 1024 / 1024 + + model_info["memory"] = { + "parameters_mb": round(param_size / 1024 / 1024, 2), + "buffers_mb": round(buffer_size / 1024 / 1024, 2), + "total_mb": round(model_size_mb, 2), + } + + # Log layer information + layer_info = [] + for name, module in model.named_modules(): + if len(list(module.children())) == 0: # Only leaf modules + layer_params = sum(p.numel() for p in module.parameters()) + if layer_params > 0: + layer_info.append( + { + "name": name, + "type": module.__class__.__name__, + "parameters": layer_params, + } + ) + + model_info["layers"] = layer_info + + if results_path: + path = os.path.join(results_path, "model_info.json") + # Save model info + with open(path, "w") as f: + json.dump(model_info, f, indent=2) + + # Log summary + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"Model: {model_info['model_class']}") + print( + f"Total parameters: {model_info['parameters']['total']:,} ({model_info['parameters']['total_millions']}M)" + ) + print( + f"Trainable parameters: {model_info['parameters']['trainable']:,} ({model_info['parameters']['trainable_millions']}M)" + ) + print(f"Model size: {model_info['memory']['total_mb']:.2f} MB") + print(f"Number of layers: {len(layer_info)}") + + +def save_weights( + store_path: str, + model: nn.Module, + device_ids: List[int], + optimizer: torch.optim.Optimizer, + epoch: int, + all_time_all_metrics, + all_losses, + best_metric: float, + args, + scheduler: Optional[torch.optim.lr_scheduler.ReduceLROnPlateau] = None, +) -> None: + """ + Save a training checkpoint containing model weights, optimizer/scheduler states, and metadata. + + Behavior: + - In Distributed Data Parallel (DDP), only rank 0 writes the file to avoid conflicts. + - If `train_lora` is True, saves only LoRA adapter weights (`lora_state_dict`); otherwise saves the full model. + - Uses `model.module.state_dict()` when the model is wrapped by DDP/DataParallel. + - Stores `epoch` and `best_metric` alongside optimizer/scheduler states. + + Args: + all_losses: + args: + store_path: Destination file path for the checkpoint (will be overwritten). + model: The model whose weights are being saved (may be wrapped by DDP/DataParallel). + device_ids: List of GPU device IDs used during training (used to detect DP wrapping in non-DDP runs). + optimizer: Optimizer whose state will be saved. + epoch: Current training epoch to record in the checkpoint. + all_time_all_metrics: + best_metric: Best validation metric achieved so far. + scheduler: Optional learning rate scheduler; its state is saved if provided. + + Returns: + None + """ + + checkpoint: Dict[str, Any] = { + "epoch": epoch, + "optimizer_name": optimizer.__class__.__name__, + "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": scheduler.state_dict() if scheduler else None, + "best_metric": best_metric, + "all_metrics": all_time_all_metrics, + "all_losses": all_losses, + } + + # Save model weights + if args.train_lora_peft: + model.save_pretrained(store_path + "_lora_") + elif args.train_lora_loralib: + checkpoint["model_state_dict"] = lora.lora_state_dict(model) + else: + if dist.is_initialized(): + # In DDP, use .module + checkpoint["model_state_dict"] = model.module.state_dict() + else: + checkpoint["model_state_dict"] = ( + model.state_dict() + if len(device_ids) <= 1 + else model.module.state_dict() + ) + + # Save only on rank 0 (or if not using DDP) + if not dist.is_initialized() or dist.get_rank() == 0: + torch.save(checkpoint, store_path) + + +def save_last_weights( + args: argparse.Namespace, + model: nn.Module, + device_ids: List[int], + optimizer: torch.optim.Optimizer, + epoch: int, + all_time_all_metrics, + all_losses, + best_metric: float, + scheduler: Optional[torch.optim.lr_scheduler.ReduceLROnPlateau] = None, +) -> None: + """ + Save the latest training checkpoint for continuation or recovery. + + The checkpoint is always written to: + {args.results_path}/last_{args.model_type}.ckpt + + This wraps `save_weights` and ensures the latest model/optimizer/scheduler + states are recorded, along with the current epoch and best metric. In DDP, + only rank 0 performs the save. Supports both standard and LoRA training. + + Args: + all_time_all_metrics: + args: Training arguments. Must define `results_path`, `model_type`, + and `train_lora`. + model: Model instance (may be wrapped by DDP/DataParallel). + device_ids: List of GPU IDs used for training. + optimizer: Optimizer whose state will be saved. + epoch: Current training epoch. + best_metric: Current best validation metric. + scheduler: Optional learning rate scheduler to save state for. + + Returns: + None + """ + store_path = f"{args.results_path}/last_{args.model_type}.ckpt" + save_weights( + store_path=store_path, + model=model, + device_ids=device_ids, + optimizer=optimizer, + epoch=epoch, + all_time_all_metrics=all_time_all_metrics, + all_losses=all_losses, + best_metric=best_metric, + args=args, + scheduler=scheduler, + ) diff --git a/src/third_party/MusicSourceSeparationTraining/utils/muon.py b/src/third_party/MusicSourceSeparationTraining/utils/muon.py new file mode 100644 index 0000000000000000000000000000000000000000..6be9d415aaae47d27dfb568fa8bc71084ebe541d --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/utils/muon.py @@ -0,0 +1,923 @@ +import math +from typing import List, Tuple + +import torch +from pytorch_optimizer.base.exception import ( + NoComplexParameterError, + NoSparseGradientError, +) +from pytorch_optimizer.base.optimizer import BaseOptimizer +from pytorch_optimizer.base.type import Betas, Closure, Loss, Parameters, ParamGroup +from pytorch_optimizer.optimizer.shampoo_utils import zero_power_via_newton_schulz_5 +from torch import nn +from torch.distributed import all_gather, get_rank, get_world_size +from torch.optim import Optimizer + + +def get_adjusted_lr( + lr: float, param_shape: Tuple[float, ...], use_adjusted_lr: bool = False +) -> float: + r"""Get the adjust learning rate.""" + output_shape, *input_shape = param_shape + input_shape = math.prod(input_shape) + + ratio: float = ( + math.pow(max(1.0, output_shape / input_shape), 0.5) + if use_adjusted_lr + else 0.2 * math.sqrt(max(output_shape, input_shape)) + ) + + return lr * ratio + + +class Muon(BaseOptimizer): + """Momentum Orthogonalized by Newton-schulz. + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-processing step, in which + each 2D parameter's update is replaced with the nearest orthogonal matrix. To efficiently orthogonalize each + update, we use a Newton-Schulz iteration, which has the advantage that it can be stably run in bfloat16 on the GPU. + + Muon is intended to optimize only the internal ≥2D parameters of a network. Embeddings, classifier heads, and + scalar or vector parameters should be optimized using AdamW. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for fine-tuning pretrained models, but we haven't tested this. + + Args: + params (Parameters): The parameters to be optimized by Muon. + lr (float): Learning rate. + momentum (float): The momentum used by the internal SGD. + weight_decay (float): Weight decay (L2 penalty). + weight_decouple (bool): The optimizer uses decoupled weight decay as in AdamW. + nesterov (bool): Whether to use nesterov momentum. + ns_steps (int): The number of Newton-Schulz iterations to run. (5 is probably always enough) + use_adjusted_lr (bool): Whether to use adjusted learning rate, which is from the Moonlight. + Reference: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py + adamw_lr (float): The learning rate for the internal AdamW. + adamw_betas (tuple): The betas for the internal AdamW. + adamw_wd (float): The weight decay for the internal AdamW. + adamw_eps (float): The epsilon for the internal AdamW. + maximize (bool): Maximize the objective with respect to the params, instead of minimizing. + + Example: + from pytorch_optimizer import Muon + + hidden_weights = [p for p in model.body.parameters() if p.ndim >= 2] + hidden_gains_biases = [p for p in model.body.parameters() if p.ndim < 2] + non_hidden_params = [*model.head.parameters(), *model.embed.parameters()] + + param_groups = [ + dict(params=hidden_weights, lr=0.02, weight_decay=0.01, use_muon=True), + dict( + params=hidden_gains_biases + non_hidden_params, + lr=3e-4, + betas=(0.9, 0.95), + weight_decay=0.01, + use_muon=False, + ), + ] + + optimizer = Muon(param_groups) + """ + + def __init__( + self, + params: Parameters, + lr: float = 2e-2, + momentum: float = 0.95, + weight_decay: float = 0.0, + weight_decouple: bool = True, + nesterov: bool = True, + ns_steps: int = 5, + use_adjusted_lr: bool = False, + adamw_lr: float = 3e-4, + adamw_betas: Betas = (0.9, 0.95), + adamw_wd: float = 0.0, + adamw_eps: float = 1e-10, + maximize: bool = False, + **kwargs, + ): + self.validate_learning_rate(lr) + self.validate_learning_rate(adamw_lr) + self.validate_non_negative(weight_decay, "weight_decay") + self.validate_range(momentum, "momentum", 0.0, 1.0, range_type="[)") + self.validate_positive(ns_steps, "ns_steps") + self.validate_betas(adamw_betas) + self.validate_non_negative(adamw_wd, "adamw_wd") + self.validate_non_negative(adamw_eps, "adamw_eps") + + self.maximize = maximize + + for group in params: + if "use_muon" not in group: + raise ValueError("`use_muon` must be set.") + + if group["use_muon"]: + group["lr"] = group.get("lr", lr) + group["momentum"] = group.get("momentum", momentum) + group["nesterov"] = group.get("nesterov", nesterov) + group["weight_decay"] = group.get("weight_decay", weight_decay) + group["ns_steps"] = group.get("ns_steps", ns_steps) + group["use_adjusted_lr"] = group.get("use_adjusted_lr", use_adjusted_lr) + else: + group["lr"] = group.get("lr", adamw_lr) + group["betas"] = group.get("betas", adamw_betas) + group["eps"] = group.get("eps", adamw_eps) + group["weight_decay"] = group.get("weight_decay", adamw_wd) + + group["weight_decouple"] = group.get("weight_decouple", weight_decouple) + + super().__init__(params, kwargs) + + def __str__(self) -> str: + return "Muon" + + def init_group(self, group: ParamGroup, **kwargs) -> None: + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad + if grad.is_sparse: + raise NoSparseGradientError(str(self)) + + if torch.is_complex(p): + raise NoComplexParameterError(str(self)) + + state = self.state[p] + + if len(state) == 0: + if group["use_muon"]: + state["momentum_buffer"] = torch.zeros_like(p) + else: + state["exp_avg"] = torch.zeros_like(p) + state["exp_avg_sq"] = torch.zeros_like(p) + + @torch.no_grad() + def step(self, closure: Closure = None) -> Loss: + loss: Loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + if "step" not in group: + self.init_group(group) + group["step"] = 1 + else: + group["step"] += 1 + + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad + + self.maximize_gradient(grad, maximize=self.maximize) + + state = self.state[p] + + self.apply_weight_decay( + p, + grad=grad, + lr=group["lr"], + weight_decay=group["weight_decay"], + weight_decouple=group["weight_decouple"], + fixed_decay=False, + ) + + if group["use_muon"]: + buf = state["momentum_buffer"] + buf.lerp_(grad, weight=1.0 - group["momentum"]) + + update = ( + grad.lerp_(buf, weight=group["momentum"]) + if group["nesterov"] + else buf + ) + if update.ndim > 2: + update = update.view(len(update), -1) + + update = zero_power_via_newton_schulz_5( + update, num_steps=group["ns_steps"] + ) + + if group.get("cautious"): + self.apply_cautious(update, grad) + + lr: float = get_adjusted_lr( + group["lr"], p.size(), use_adjusted_lr=group["use_adjusted_lr"] + ) + + p.add_(update.reshape(p.shape), alpha=-lr) + else: + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + + beta1, beta2 = group["betas"] + + bias_correction1: float = self.debias(beta1, group["step"]) + bias_correction2_sq: float = math.sqrt( + self.debias(beta2, group["step"]) + ) + + exp_avg.lerp_(grad, weight=1.0 - beta1) + exp_avg_sq.lerp_(grad.square(), weight=1.0 - beta2) + + de_nom = ( + exp_avg_sq.sqrt().add_(group["eps"]).div_(bias_correction2_sq) + ) + + p.addcdiv_(exp_avg / bias_correction1, de_nom, value=-group["lr"]) + + return loss + + +class DistributedMuon(BaseOptimizer): # pragma: no cover + """Momentum Orthogonalized by Newton-schulz. + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-processing step, in which + each 2D parameter's update is replaced with the nearest orthogonal matrix. To efficiently orthogonalize each + update, we use a Newton-Schulz iteration, which has the advantage that it can be stably run in bfloat16 on the GPU. + + Muon is intended to optimize only the internal ≥2D parameters of a network. Embeddings, classifier heads, and + scalar or vector parameters should be optimized using AdamW. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for fine-tuning pretrained models, but we haven't tested this. + + Args: + params (Parameters): The parameters to be optimized by Muon. + lr (float): Learning rate. + momentum (float): The momentum used by the internal SGD. + weight_decay (float): Weight decay (L2 penalty). + weight_decouple (bool): The optimizer uses decoupled weight decay as in AdamW. + nesterov (bool): Whether to use nesterov momentum. + ns_steps (int): The number of Newton-Schulz iterations to run. (5 is probably always enough) + use_adjusted_lr (bool): Whether to use adjusted learning rate, which is from the Moonlight. + Reference: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py + adamw_lr (float): The learning rate for the internal AdamW. + adamw_betas (tuple): The betas for the internal AdamW. + adamw_wd (float): The weight decay for the internal AdamW. + adamw_eps (float): The epsilon for the internal AdamW. + maximize (bool): Maximize the objective with respect to the params, instead of minimizing. + + Example: + from pytorch_optimizer import DistributedMuon + + hidden_weights = [p for p in model.body.parameters() if p.ndim >= 2] + hidden_gains_biases = [p for p in model.body.parameters() if p.ndim < 2] + non_hidden_params = [*model.head.parameters(), *model.embed.parameters()] + + param_groups = [ + dict(params=hidden_weights, lr=0.02, weight_decay=0.01, use_muon=True), + dict( + params=hidden_gains_biases + non_hidden_params, + lr=3e-4, + betas=(0.9, 0.95), + weight_decay=0.01, + use_muon=False, + ), + ] + + optimizer = DistributedMuon(param_groups) + """ + + def __init__( + self, + params: Parameters, + lr: float = 2e-2, + momentum: float = 0.95, + weight_decay: float = 0.0, + weight_decouple: bool = True, + nesterov: bool = True, + ns_steps: int = 5, + use_adjusted_lr: bool = False, + adamw_lr: float = 3e-4, + adamw_betas: Betas = (0.9, 0.95), + adamw_wd: float = 0.0, + adamw_eps: float = 1e-10, + maximize: bool = False, + **kwargs, + ): + self.validate_learning_rate(lr) + self.validate_learning_rate(adamw_lr) + self.validate_non_negative(weight_decay, "weight_decay") + self.validate_range(momentum, "momentum", 0.0, 1.0, range_type="[)") + self.validate_positive(ns_steps, "ns_steps") + self.validate_betas(adamw_betas) + self.validate_non_negative(adamw_wd, "adamw_wd") + self.validate_non_negative(adamw_eps, "adamw_eps") + + self.maximize = maximize + + self.world_size: int = get_world_size() + self.rank: int = get_rank() + + for group in params: + if "use_muon" not in group: + raise ValueError("`use_muon` must be set.") + + if group["use_muon"]: + group["lr"] = group.get("lr", lr) + group["momentum"] = group.get("momentum", momentum) + group["nesterov"] = group.get("nesterov", nesterov) + group["weight_decay"] = group.get("weight_decay", weight_decay) + group["ns_steps"] = group.get("ns_steps", ns_steps) + group["use_adjusted_lr"] = group.get("use_adjusted_lr", use_adjusted_lr) + else: + group["lr"] = group.get("lr", adamw_lr) + group["betas"] = group.get("betas", adamw_betas) + group["eps"] = group.get("eps", adamw_eps) + group["weight_decay"] = group.get("weight_decay", adamw_wd) + + group["weight_decouple"] = group.get("weight_decouple", weight_decouple) + + super().__init__(params, kwargs) + + def __str__(self) -> str: + return "DistributedMuon" + + def init_group(self, group: ParamGroup, **kwargs) -> None: + for p in group["params"]: + if p.grad is None: + p.grad = torch.zeros_like(p) + + grad = p.grad + if grad.is_sparse: + raise NoSparseGradientError(str(self)) + + if torch.is_complex(p): + raise NoComplexParameterError(str(self)) + + state = self.state[p] + + if len(state) == 0 and not group["use_muon"]: + state["exp_avg"] = torch.zeros_like(p) + state["exp_avg_sq"] = torch.zeros_like(p) + + @torch.no_grad() + def step(self, closure: Closure = None) -> Loss: + loss: Loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + if "step" not in group: + self.init_group(group) + group["step"] = 1 + else: + group["step"] += 1 + + if group["use_muon"]: + params = group["params"] + padded_params = params + [torch.empty_like(params[-1])] * ( + self.world_size - len(params) % self.world_size + ) + + for i in range(len(params))[:: self.world_size]: + if i + self.rank < len(params): + p = params[i + self.rank] + + grad = p.grad + + self.maximize_gradient(grad, maximize=self.maximize) + + state = self.state[p] + if len(state) == 0: + state["momentum_buffer"] = torch.zeros_like(p) + + self.apply_weight_decay( + p, + grad=grad, + lr=group["lr"], + weight_decay=group["weight_decay"], + weight_decouple=group["weight_decouple"], + fixed_decay=False, + ) + + buf = state["momentum_buffer"] + buf.lerp_(grad, weight=1.0 - group["momentum"]) + + update = ( + grad.lerp_(buf, weight=group["momentum"]) + if group["nesterov"] + else buf + ) + if update.ndim > 2: + update = update.view(len(update), -1) + + update = zero_power_via_newton_schulz_5( + update, num_steps=group["ns_steps"] + ) + + if group.get("cautious"): + self.apply_cautious(update, grad) + + lr: float = get_adjusted_lr( + group["lr"], + p.size(), + use_adjusted_lr=group["use_adjusted_lr"], + ) + + p.add_(update.reshape(p.shape), alpha=-lr) + + all_gather(padded_params[i:i + self.world_size], padded_params[i:i + self.rank]) # fmt: skip + else: + for p in group["params"]: + grad = p.grad + + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + + beta1, beta2 = group["betas"] + + bias_correction1: float = self.debias(beta1, group["step"]) + bias_correction2_sq: float = math.sqrt( + self.debias(beta2, group["step"]) + ) + + exp_avg.lerp_(grad, weight=1.0 - beta1) + exp_avg_sq.lerp_(grad.square(), weight=1.0 - beta2) + + de_nom = ( + exp_avg_sq.sqrt().add_(group["eps"]).div_(bias_correction2_sq) + ) + + p.addcdiv_(exp_avg / bias_correction1, de_nom, value=-group["lr"]) + + return loss + + +class AdaMuon(BaseOptimizer): + """Adaptive Muon optimizer. + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-processing step, in which + each 2D parameter's update is replaced with the nearest orthogonal matrix. To efficiently orthogonalize each + update, we use a Newton-Schulz iteration, which has the advantage that it can be stably run in bfloat16 on the GPU. + + Muon is intended to optimize only the internal ≥2D parameters of a network. Embeddings, classifier heads, and + scalar or vector parameters should be optimized using AdamW. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for fine-tuning pretrained models, but we haven't tested this. + + Args: + params (Parameters): The parameters to be optimized by Muon. + lr (float): Learning rate. + betas (tuple): Coefficients used for computing running averages of gradient and the squared Hessian trace. + weight_decay (float): Weight decay (L2 penalty). + weight_decouple (bool): The optimizer uses decoupled weight decay as in AdamW. + ns_steps (int): The number of Newton-Schulz iterations to run. (5 is probably always enough) + use_adjusted_lr (bool): Whether to use adjusted learning rate, which is from the Moonlight. + Reference: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py + adamw_lr (float): The learning rate for the internal AdamW. + adamw_betas (tuple): The betas for the internal AdamW. + adamw_wd (float): The weight decay for the internal AdamW. + eps (float): Term added to the denominator to improve numerical stability. + maximize (bool): Maximize the objective with respect to the params, instead of minimizing. + + Example: + from pytorch_optimizer import AdaMuon + + hidden_weights = [p for p in model.body.parameters() if p.ndim >= 2] + hidden_gains_biases = [p for p in model.body.parameters() if p.ndim < 2] + non_hidden_params = [*model.head.parameters(), *model.embed.parameters()] + + param_groups = [ + dict(params=hidden_weights, lr=0.02, weight_decay=0.01, use_muon=True), + dict( + params=hidden_gains_biases + non_hidden_params, + lr=3e-4, + betas=(0.9, 0.95), + weight_decay=0.01, + use_muon=False, + ), + ] + + optimizer = AdaMuon(param_groups) + """ + + def __init__( + self, + params: Parameters, + lr: float = 2e-2, + betas: Betas = (0.9, 0.95), + weight_decay: float = 0.0, + weight_decouple: bool = True, + ns_steps: int = 5, + use_adjusted_lr: bool = False, + adamw_lr: float = 3e-4, + adamw_betas: Betas = (0.9, 0.999), + adamw_wd: float = 0.0, + eps: float = 1e-10, + maximize: bool = False, + **kwargs, + ): + self.validate_learning_rate(lr) + self.validate_learning_rate(adamw_lr) + self.validate_non_negative(weight_decay, "weight_decay") + self.validate_positive(ns_steps, "ns_steps") + self.validate_betas(betas) + self.validate_betas(adamw_betas) + self.validate_non_negative(adamw_wd, "adamw_wd") + self.validate_non_negative(eps, "eps") + + self.maximize = maximize + + for group in params: + if "use_muon" not in group: + raise ValueError("`use_muon` must be set.") + + if group["use_muon"]: + group["lr"] = group.get("lr", lr) + group["betas"] = group.get("betas", betas) + group["weight_decay"] = group.get("weight_decay", weight_decay) + group["ns_steps"] = group.get("ns_steps", ns_steps) + group["use_adjusted_lr"] = group.get("use_adjusted_lr", use_adjusted_lr) + else: + group["lr"] = group.get("lr", adamw_lr) + group["betas"] = group.get("betas", adamw_betas) + group["weight_decay"] = group.get("weight_decay", adamw_wd) + + group["weight_decouple"] = group.get("weight_decouple", weight_decouple) + group["eps"] = group.get("eps", eps) + + super().__init__(params, kwargs) + + def __str__(self) -> str: + return "AdaMuon" + + def init_group(self, group: ParamGroup, **kwargs) -> None: + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad + if grad.is_sparse: + raise NoSparseGradientError(str(self)) + + if torch.is_complex(p): + raise NoComplexParameterError(str(self)) + + state = self.state[p] + + if len(state) == 0: + if group["use_muon"]: + state["m"] = torch.zeros_like(p) + state["v"] = torch.zeros_like(p.flatten()) + else: + state["exp_avg"] = torch.zeros_like(p) + state["exp_avg_sq"] = torch.zeros_like(p) + + @torch.no_grad() + def step(self, closure: Closure = None) -> Loss: + loss: Loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + if "step" not in group: + self.init_group(group) + group["step"] = 1 + else: + group["step"] += 1 + + beta1, beta2 = group["betas"] + + bias_correction1: float = self.debias(beta1, group["step"]) + bias_correction2: float = self.debias(beta2, group["step"]) + + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad + + self.maximize_gradient(grad, maximize=self.maximize) + + state = self.state[p] + + self.apply_weight_decay( + p, + grad=grad, + lr=group["lr"], + weight_decay=group["weight_decay"], + weight_decouple=group["weight_decouple"], + fixed_decay=False, + ) + + if group["use_muon"]: + m = state["m"] + m.lerp_(grad, weight=1.0 - beta1) + + update = m.clone() + + if update.ndim > 2: + update = update.view(len(update), -1) + + update = zero_power_via_newton_schulz_5( + update, num_steps=group["ns_steps"] + ).flatten() + + v = state["v"] + v.mul_(beta2).addcmul_(update, update, value=1.0 - beta2) + + update.div_((v / bias_correction2).sqrt_().add_(group["eps"])) + update = update.reshape(p.size()) + + update.mul_(0.2 * math.sqrt(p.numel())).div_( + update.norm().add_(group["eps"]) + ) + + lr: float = get_adjusted_lr( + group["lr"], p.size(), use_adjusted_lr=group["use_adjusted_lr"] + ) + + p.add_(update, alpha=-lr) + else: + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + + exp_avg.lerp_(grad, weight=1.0 - beta1) + exp_avg_sq.lerp_(grad.square(), weight=1.0 - beta2) + + de_nom = ( + exp_avg_sq.sqrt() + .add_(group["eps"]) + .div_(math.sqrt(bias_correction2)) + ) + + p.addcdiv_(exp_avg / bias_correction1, de_nom, value=-group["lr"]) + + return loss + + +class AdaGO(BaseOptimizer): + """AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates. + + Args: + params (Parameters): The parameters to be optimized by Muon. + lr (float): Learning rate. + momentum (float): The momentum used by the internal SGD. + weight_decay (float): Weight decay (L2 penalty). + weight_decouple (bool): The optimizer uses decoupled weight decay as in AdamW. + nesterov (bool): Whether to use nesterov momentum. + gamma (float): Gamma factor. Empirically, AdaGO performs robustly across a wide range of gamma values. + eps (float): Epsilon value. Lower bound eps > 0 on the stepsizes. + ns_steps (int): The number of Newton-Schulz iterations to run. (5 is probably always enough) + use_adjusted_lr (bool): Whether to use adjusted learning rate, which is from the Moonlight. + Reference: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py + adamw_lr (float): The learning rate for the internal AdamW. + adamw_betas (tuple): The betas for the internal AdamW. + adamw_wd (float): The weight decay for the internal AdamW. + adamw_eps (float): The epsilon for the internal AdamW. + maximize (bool): Maximize the objective with respect to the params, instead of minimizing. + + Example: + from pytorch_optimizer import AdaGO + + hidden_weights = [p for p in model.body.parameters() if p.ndim >= 2] + hidden_gains_biases = [p for p in model.body.parameters() if p.ndim < 2] + non_hidden_params = [*model.head.parameters(), *model.embed.parameters()] + + param_groups = [ + dict(params=hidden_weights, lr=0.02, weight_decay=0.01, use_muon=True), + dict( + params=hidden_gains_biases + non_hidden_params, + lr=3e-4, + betas=(0.9, 0.95), + weight_decay=0.01, + use_muon=False, + ), + ] + + optimizer = AdaGO(param_groups) + """ + + def __init__( + self, + params: Parameters, + lr: float = 5e-2, + momentum: float = 0.95, + weight_decay: float = 0.0, + weight_decouple: bool = True, + gamma: float = 10.0, + eps: float = 5e-4, + v: float = 1e-6, + nesterov: bool = False, + ns_steps: int = 5, + use_adjusted_lr: bool = False, + adamw_lr: float = 3e-4, + adamw_betas: Betas = (0.9, 0.95), + adamw_wd: float = 0.0, + adamw_eps: float = 1e-10, + maximize: bool = False, + **kwargs, + ): + self.validate_learning_rate(lr) + self.validate_learning_rate(adamw_lr) + self.validate_non_negative(weight_decay, "weight_decay") + self.validate_range(momentum, "momentum", 0.0, 1.0, range_type="[)") + self.validate_positive(ns_steps, "ns_steps") + self.validate_positive(gamma, "gamma") + self.validate_positive(eps, "eps") + self.validate_positive(v, "v") + self.validate_betas(adamw_betas) + self.validate_non_negative(adamw_wd, "adamw_wd") + self.validate_non_negative(adamw_eps, "adamw_eps") + + self.maximize = maximize + + for group in params: + if "use_muon" not in group: + raise ValueError("`use_muon` must be set.") + + if group["use_muon"]: + group["lr"] = group.get("lr", lr) + group["momentum"] = group.get("momentum", momentum) + group["nesterov"] = group.get("nesterov", nesterov) + group["weight_decay"] = group.get("weight_decay", weight_decay) + group["ns_steps"] = group.get("ns_steps", ns_steps) + group["gamma"] = group.get("gamma", gamma) + group["eps"] = group.get("eps", eps) + group["v"] = group.get("v", v) + group["use_adjusted_lr"] = group.get("use_adjusted_lr", use_adjusted_lr) + else: + group["lr"] = group.get("lr", adamw_lr) + group["betas"] = group.get("betas", adamw_betas) + group["eps"] = group.get("eps", adamw_eps) + group["weight_decay"] = group.get("weight_decay", adamw_wd) + + group["weight_decouple"] = group.get("weight_decouple", weight_decouple) + + super().__init__(params, kwargs) + + def __str__(self) -> str: + return "AdaGO" + + def init_group(self, group: ParamGroup, **kwargs) -> None: + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad + if grad.is_sparse: + raise NoSparseGradientError(str(self)) + + if torch.is_complex(p): + raise NoComplexParameterError(str(self)) + + state = self.state[p] + + if len(state) == 0: + if group["use_muon"]: + state["momentum_buffer"] = torch.zeros_like(p) + state["v"] = torch.tensor( + group["v"], dtype=p.dtype, device=p.device + ) + else: + state["exp_avg"] = torch.zeros_like(p) + state["exp_avg_sq"] = torch.zeros_like(p) + + @torch.no_grad() + def step(self, closure: Closure = None) -> Loss: + loss: Loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + if "step" not in group: + self.init_group(group) + group["step"] = 1 + else: + group["step"] += 1 + + for p in group["params"]: + if p.grad is None: + continue + + grad = p.grad + + self.maximize_gradient(grad, maximize=self.maximize) + + state = self.state[p] + + self.apply_weight_decay( + p, + grad=grad, + lr=group["lr"], + weight_decay=group["weight_decay"], + weight_decouple=group["weight_decouple"], + fixed_decay=False, + ) + + if group["use_muon"]: + buf, v = state["momentum_buffer"], state["v"] + buf.lerp_(grad, weight=1.0 - group["momentum"]) + + v.add_(min(grad.norm(p=2.0).pow(2), group["gamma"] ** 2)) + + update = ( + grad.lerp_(buf, weight=group["momentum"]) + if group["nesterov"] + else buf + ) + if update.ndim > 2: + update = update.view(len(update), -1) + + update = zero_power_via_newton_schulz_5( + update, num_steps=group["ns_steps"] + ) + + if group.get("cautious"): + self.apply_cautious(update, grad) + + lr: float = get_adjusted_lr( + group["lr"], p.size(), use_adjusted_lr=group["use_adjusted_lr"] + ) + + p.add_( + update.reshape(p.shape), + alpha=-max( + group["eps"], + (lr * min(grad.norm(2), group["gamma"]) / v).item(), + ), + ) + else: + exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] + + beta1, beta2 = group["betas"] + + bias_correction1: float = self.debias(beta1, group["step"]) + bias_correction2_sq: float = math.sqrt( + self.debias(beta2, group["step"]) + ) + + exp_avg.lerp_(grad, weight=1.0 - beta1) + exp_avg_sq.lerp_(grad.square(), weight=1.0 - beta2) + + de_nom = ( + exp_avg_sq.sqrt().add_(group["eps"]).div_(bias_correction2_sq) + ) + + p.addcdiv_(exp_avg / bias_correction1, de_nom, value=-group["lr"]) + + return loss + + +def prepare_muon_parameters( + model: nn.Module, + optimizer_name: str, + lr: float, + weight_decay: float, + adamw_lr: float = 3e-4, + adamw_wd: float = 0.0, + **kwargs, +) -> Optimizer: + """Prepare the parameters for Muon optimizer. + + Be careful at using this function to prepare the parameters for Muon optimizer. It's not likely acting perfectly + for all cases. So, highly recommend you to create the Muon optimizer manually following by the given example in the + docstring. + """ + muon_parameters: List[str] = [] + non_muon_params: List[str] = [] + + for _, module in model.named_modules(): + for name, param in module.named_parameters(recurse=False): + if ( + isinstance(module, (nn.Linear, nn.Conv1d, nn.LSTM, nn.Conv2d)) + and param.ndim >= 2 + and "head" not in name + ): + muon_parameters.append(param) + else: + non_muon_params.append(param) + + param_groups: Parameters = [ + { + "params": muon_parameters, + "lr": lr, + "weight_decay": weight_decay, + "use_muon": True, + }, + { + "params": non_muon_params, + "lr": adamw_lr, + "weight_decay": adamw_wd, + "use_muon": False, + }, + ] + + optimizer_name = optimizer_name.lower() + + if optimizer_name == "adamuon": + return AdaMuon(param_groups, **kwargs) + if optimizer_name == "adago": + return AdaGO(param_groups, **kwargs) + + return Muon(param_groups, **kwargs) diff --git a/src/third_party/MusicSourceSeparationTraining/utils/settings.py b/src/third_party/MusicSourceSeparationTraining/utils/settings.py new file mode 100644 index 0000000000000000000000000000000000000000..7085a9e4e52612e54cb3bbc1f01ebb5fd505d9d4 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/utils/settings.py @@ -0,0 +1,963 @@ +import argparse +import os +import random +import socket +import time +from typing import Dict, List, Tuple, Union + +import numpy as np +import soundfile as sf +import torch +import torch.distributed as dist +import wandb +import yaml +from ml_collections import ConfigDict +from omegaconf import OmegaConf +from torch import nn + + +def parse_args_train( + dict_args: Union[argparse.Namespace, Dict, None], +) -> argparse.Namespace: + """ + Parse command-line arguments for training configuration. + + This function constructs an argument parser for model, dataset, training, and logging + options, merges overrides from a provided dictionary (if any), and returns the parsed + arguments. If `dict_args` is None, the arguments are parsed from `sys.argv`. + + Args: + dict_args (Dict | None): Optional dictionary of argument overrides. Keys should + match the defined CLI options. + + Returns: + argparse.Namespace: Parsed arguments namespace containing all configuration + values required for training. + """ + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_type", + type=str, + default="mdx23c", + help="One of mdx23c, htdemucs, segm_models, mel_band_roformer, bs_roformer, swin_upernet, bandit", + ) + parser.add_argument("--config_path", type=str, help="path to config file") + parser.add_argument( + "--start_check_point", + type=str, + default="", + help="Initial checkpoint to start training", + ) + parser.add_argument( + "--load_optimizer", + action="store_true", + help="Load optimizer state from checkpoint (if available)", + ) + parser.add_argument( + "--load_scheduler", + action="store_true", + help="Load scheduler state from checkpoint (if available)", + ) + parser.add_argument( + "--load_epoch", + action="store_true", + help="Load epoch number from checkpoint (if available)", + ) + parser.add_argument( + "--load_best_metric", + action="store_true", + help="Load best metric from checkpoint (if available)", + ) + parser.add_argument( + "--load_all_metrics", + action="store_true", + help="Load all metrics from checkpoint (if available)", + ) + parser.add_argument( + "--load_all_losses", + action="store_true", + help="Load all losses from checkpoint (if available)", + ) + parser.add_argument( + "--safe_mode", action="store_true", help="Ignore forward errors" + ) + parser.add_argument( + "--results_path", + type=str, + help="path to folder where results will be stored (weights, metadata)", + ) + parser.add_argument( + "--data_path", + nargs="+", + type=str, + help="Dataset data paths. You can provide several folders.", + ) + parser.add_argument( + "--dataset_type", + type=int, + default=1, + help="Dataset type. Must be one of: 1, 2, 3, 4, 5, 6, 7. Details here: https://github.com/ZFTurbo/Music-Source-Separation-Training/blob/main/docs/dataset_types.md", + ) + parser.add_argument( + "--valid_path", + nargs="+", + type=str, + help="validation data paths. You can provide several folders.", + ) + parser.add_argument( + "--num_workers", type=int, default=0, help="dataloader num_workers" + ) + parser.add_argument( + "--pin_memory", action="store_true", help="dataloader pin_memory" + ) + parser.add_argument("--seed", type=int, default=0, help="random seed") + parser.add_argument( + "--device_ids", nargs="+", type=int, default=[0], help="list of gpu ids" + ) + parser.add_argument( + "--loss", + type=str, + nargs="+", + choices=[ + "masked_loss", + "mse_loss", + "l1_loss", + "multistft_loss", + "spec_masked_loss", + "spec_rmse_loss", + "log_wmse_loss", + "l1_snr_loss", + "l1_snr_db_loss", + "stft_l1_snr_db_loss", + "multi_l1_snr_db_loss", + ], + default=["masked_loss"], + help="List of loss functions to use", + ) + parser.add_argument( + "--masked_loss_coef", type=float, default=1.0, help="Coef for loss" + ) + parser.add_argument( + "--mse_loss_coef", type=float, default=1.0, help="Coef for loss" + ) + parser.add_argument("--l1_loss_coef", type=float, default=1.0, help="Coef for loss") + parser.add_argument( + "--log_wmse_loss_coef", type=float, default=1.0, help="Coef for loss" + ) + parser.add_argument( + "--multistft_loss_coef", type=float, default=0.001, help="Coef for loss" + ) + parser.add_argument( + "--spec_masked_loss_coef", type=float, default=1, help="Coef for loss" + ) + parser.add_argument( + "--spec_rmse_loss_coef", type=float, default=1, help="Coef for loss" + ) + parser.add_argument( + "--l1_snr_loss_coef", type=float, default=1.0, help="Coef for L1-SNR loss" + ) + parser.add_argument( + "--l1_snr_db_loss_coef", type=float, default=1.0, help="Coef for L1-SNR-DB loss" + ) + parser.add_argument( + "--stft_l1_snr_db_loss_coef", + type=float, + default=1.0, + help="Coef for STFT-L1-SNR-DB loss", + ) + parser.add_argument( + "--multi_l1_snr_db_loss_coef", + type=float, + default=1.0, + help="Coef for Multi-L1-SNR-DB loss", + ) + parser.add_argument("--wandb_key", type=str, default="", help="wandb API Key") + parser.add_argument("--wandb_offline", action="store_true", help="local wandb") + parser.add_argument( + "--pre_valid", action="store_true", help="Run validation before training" + ) + parser.add_argument( + "--metrics", + nargs="+", + type=str, + default=["sdr"], + choices=[ + "k_sdr", + "sdr", + "l1_freq", + "si_sdr", + "log_wmse", + "aura_stft", + "aura_mrstft", + "bleedless", + "fullness", + "l1_snr", + ], + help="List of metrics to use.", + ) + parser.add_argument( + "--metric_for_scheduler", + default="sdr", + choices=[ + "k_sdr", + "sdr", + "l1_freq", + "si_sdr", + "log_wmse", + "aura_stft", + "aura_mrstft", + "bleedless", + "fullness", + "l1_snr", + ], + help="Metric which will be used for scheduler.", + ) + parser.add_argument( + "--train_lora_peft", action="store_true", help="Training with LoRA from peft" + ) + parser.add_argument( + "--train_lora_loralib", + action="store_true", + help="Training with LoRA from loralib", + ) + parser.add_argument( + "--lora_checkpoint_peft", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + parser.add_argument( + "--lora_checkpoint_loralib", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + parser.add_argument( + "--each_metrics_in_name", + action="store_true", + help="All stems in naming checkpoints", + ) + parser.add_argument( + "--use_standard_loss", + action="store_true", + help="Roformers will use provided loss instead of internal", + ) + parser.add_argument( + "--save_weights_every_epoch", + action="store_true", + help="Weights will be saved every epoch with all metric values", + ) + parser.add_argument( + "--persistent_workers", + action="store_true", + help="dataloader persistent_workers", + ) + parser.add_argument( + "--prefetch_factor", type=int, default=None, help="dataloader prefetch_factor" + ) + parser.add_argument( + "--set_per_process_memory_fraction", + action="store_true", + help="using only VRAM, no RAM", + ) + parser.add_argument( + "--load_only_compatible_weights", + action="store_true", + help="using only VRAM, no RAM", + ) + parser.add_argument( + "--freeze_layers", + nargs="+", + type=str, + help="List of layers to freeze. Use prefixes e.g. layer1 - will freeze all layers whose names " + "starts with layer1. You can set mulitple parameters.", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + if args.metric_for_scheduler not in args.metrics: + args.metrics += [args.metric_for_scheduler] + + get_internal_loss = ( + args.model_type in ("mel_band_conformer",) or "roformer" in args.model_type + ) and not args.use_standard_loss + if get_internal_loss: + args.loss = [f"{args.model_type}_loss"] + return args + + +def parse_args_valid(dict_args: Union[Dict, None]) -> argparse.Namespace: + """ + Parse command-line arguments for validation configuration. + + Builds the CLI for model selection, configuration paths, validation data + locations, output/spectrogram saving options, device/runtime settings, and + evaluation metrics. If `dict_args` is provided, its key–value pairs override + or set the parsed arguments; otherwise arguments are read from `sys.argv`. + + Args: + dict_args (Union[Dict, None]): Optional mapping of argument names to values + used to override or supply CLI options programmatically. + + Returns: + argparse.Namespace: Parsed arguments namespace containing all validation + configuration values. + """ + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_type", + type=str, + default="mdx23c", + help="One of mdx23c, htdemucs, segm_models, mel_band_roformer," + " bs_roformer, swin_upernet, bandit", + ) + parser.add_argument("--config_path", type=str, help="Path to config file") + parser.add_argument( + "--start_check_point", + type=str, + default="", + help="Initial checkpoint to valid weights", + ) + parser.add_argument("--valid_path", nargs="+", type=str, help="Validate path") + parser.add_argument( + "--store_dir", type=str, default="", help="Path to store results as wav file" + ) + parser.add_argument( + "--draw_spectro", + type=float, + default=0, + help="If --store_dir is set then code will generate spectrograms for resulted stems as well." + " Value defines for how many seconds os track spectrogram will be generated.", + ) + parser.add_argument( + "--device_ids", nargs="+", type=int, default=[0], help="List of gpu ids" + ) + parser.add_argument( + "--num_workers", type=int, default=0, help="Dataloader num_workers" + ) + parser.add_argument( + "--pin_memory", action="store_true", help="Dataloader pin_memory" + ) + parser.add_argument( + "--extension", type=str, default="wav", help="Choose extension for validation" + ) + parser.add_argument( + "--use_tta", + action="store_true", + help="Flag adds test time augmentation during inference (polarity and channel inverse)." + "While this triples the runtime, it reduces noise and slightly improves prediction quality.", + ) + parser.add_argument( + "--metrics", + nargs="+", + type=str, + default=["sdr"], + choices=[ + "k_sdr", + "sdr", + "l1_freq", + "si_sdr", + "neg_log_wmse", + "aura_stft", + "aura_mrstft", + "bleedless", + "fullness", + "l1_snr", + ], + help="List of metrics to use.", + ) + parser.add_argument( + "--lora_checkpoint_peft", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + parser.add_argument( + "--lora_checkpoint_loralib", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + + return args + + +def parse_args_inference(dict_args: Union[Dict, None]) -> argparse.Namespace: + """ + Parse command-line arguments for inference configuration. + + Builds the CLI for model selection, configuration path, input/output handling, + device/runtime options, test-time augmentation, and optional LoRA checkpoints. + If `dict_args` is provided, its key–value pairs override or supply CLI options + programmatically; otherwise, arguments are read from `sys.argv`. + + Args: + dict_args (Union[Dict, None]): Optional mapping of argument names to values + used to override or supply CLI options programmatically. + + Returns: + argparse.Namespace: Parsed arguments namespace containing all inference + configuration values. + """ + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_type", + type=str, + default="mdx23c", + help="One of bandit, bandit_v2, bs_roformer, htdemucs, mdx23c, mel_band_roformer," + " scnet, scnet_unofficial, segm_models, swin_upernet, torchseg", + ) + parser.add_argument("--config_path", type=str, help="path to config file") + parser.add_argument( + "--start_check_point", + type=str, + default="", + help="Initial checkpoint to valid weights", + ) + parser.add_argument( + "--input_folder", type=str, help="folder with mixtures to process" + ) + parser.add_argument( + "--store_dir", type=str, default="", help="path to store results as wav file" + ) + parser.add_argument( + "--draw_spectro", + type=float, + default=0, + help="Code will generate spectrograms for resulted stems." + " Value defines for how many seconds os track spectrogram will be generated.", + ) + parser.add_argument( + "--device_ids", nargs="+", type=int, default=0, help="list of gpu ids" + ) + parser.add_argument( + "--extract_instrumental", + action="store_true", + help="invert vocals to get instrumental if provided", + ) + parser.add_argument( + "--disable_detailed_pbar", + action="store_true", + help="disable detailed progress bar", + ) + parser.add_argument( + "--force_cpu", + action="store_true", + help="Force the use of CPU even if CUDA is available", + ) + parser.add_argument( + "--flac_file", action="store_true", help="Output flac file instead of wav" + ) + parser.add_argument( + "--pcm_type", + type=str, + choices=["PCM_16", "PCM_24", "FLOAT"], + default="FLOAT", + help="PCM type for FLAC files (PCM_16 or PCM_24)", + ) + parser.add_argument( + "--use_tta", + action="store_true", + help="Flag adds test time augmentation during inference (polarity and channel inverse)." + "While this triples the runtime, it reduces noise and slightly improves prediction quality.", + ) + parser.add_argument( + "--lora_checkpoint_peft", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + parser.add_argument( + "--filename_template", + type=str, + default="{file_name}/{instr}", + help="Output filename template, without extension, using '/' for subdirectories. Default: '{file_name}/{instr}'", + ) + parser.add_argument( + "--lora_checkpoint_loralib", + type=str, + default="", + help="Initial checkpoint to LoRA weights", + ) + if dict_args is not None: + args = parser.parse_args([]) + args_dict = vars(args) + args_dict.update(dict_args) + args = argparse.Namespace(**args_dict) + else: + args = parser.parse_args() + args.pcm_type = validate_sndfile_subtype(args) + + return args + + +def validate_sndfile_subtype(args): + codec = "flac" if getattr(args, "flac_file", False) else "wav" + subtype = args.pcm_type + if subtype in sf.available_subtypes(codec): + return subtype + default = sf.default_subtype(codec) + print( + f"WARNING: codec {codec} doesn't support subtype {subtype}, defaulting to {default}" + ) + return default + + +def load_config(model_type: str, config_path: str) -> Union[ConfigDict, OmegaConf]: + """ + Load a model configuration from a file. + + Based on `model_type`, returns either an OmegaConf (e.g., for 'htdemucs') + or a YAML-parsed ConfigDict for other models. + + Args: + model_type (str): Model identifier that determines the loader behavior + (e.g., 'htdemucs', 'mdx23c', etc.). + config_path (str): Path to the configuration file (YAML/OmegaConf). + + Returns: + Union[ConfigDict, OmegaConf]: Loaded configuration object. + + Raises: + FileNotFoundError: If `config_path` does not point to an existing file. + ValueError: If the configuration cannot be parsed or is otherwise invalid. + """ + try: + with open(config_path, "r") as f: + if model_type == "htdemucs": + config = OmegaConf.load(config_path) + else: + config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader)) + return config + except FileNotFoundError: + raise FileNotFoundError(f"Configuration file not found at {config_path}") + except Exception as e: + raise ValueError(f"Error loading configuration: {e}") + + +def get_model_from_config( + model_type: str, config_path: str +) -> Tuple[nn.Module, Union[ConfigDict, OmegaConf]]: + """ + Load and instantiate a model using a configuration file. + + Given a `model_type` and a path to a configuration, this function loads the + configuration (YAML or OmegaConf) and constructs the corresponding model. + + Args: + model_type (str): Identifier of the model family (e.g., 'mdx23c', 'htdemucs', + 'scnet', 'mel_band_conformer', etc.). + config_path (str): Filesystem path to the configuration file used to + initialize the model. + + Returns: + Tuple[nn.Module, Union[ConfigDict, OmegaConf]]: A tuple containing the + initialized PyTorch model and the loaded configuration object. + + Raises: + ValueError: If `model_type` is unknown or model initialization fails. + FileNotFoundError: If `config_path` does not exist (may be raised by the + underlying config loader). + """ + + config = load_config(model_type, config_path) + if "model_type" in config.training: + model_type = config.training.model_type + if model_type == "mdx23c": + from models.mdx23c_tfc_tdf_v3 import TFC_TDF_net + + model = TFC_TDF_net(config) + elif model_type == "htdemucs": + from models.demucs4ht import get_model + + model = get_model(config) + elif model_type == "segm_models": + from models.segm_models import Segm_Models_Net + + model = Segm_Models_Net(config) + elif model_type == "torchseg": + from models.torchseg_models import Torchseg_Net + + model = Torchseg_Net(config) + elif model_type == "mel_band_roformer": + from models.bs_roformer import MelBandRoformer + + model = MelBandRoformer(**dict(config.model)) + elif model_type == "mel_band_conformer": + from models.bs_roformer import MelBandConformer + + model = MelBandConformer(**dict(config.model)) + elif model_type == "mel_band_roformer_experimental": + from models.bs_roformer.mel_band_roformer_experimental import MelBandRoformer + + model = MelBandRoformer(**dict(config.model)) + elif model_type == "bs_roformer": + from models.bs_roformer import BSRoformer + + model = BSRoformer(**dict(config.model)) + elif model_type == "bs_conformer": + from models.bs_roformer import BSConformer + + model = BSConformer(**dict(config.model)) + elif model_type == "bs_roformer_experimental": + from models.bs_roformer.bs_roformer_experimental import BSRoformer + + model = BSRoformer(**dict(config.model)) + elif model_type == "bs_mamba2": + from models.bs_mamba2_code.bs_mamba2 import BSMamba2Model + + model = BSMamba2Model(**dict(config.model)) + elif model_type == "swin_upernet": + from models.upernet_swin_transformers import Swin_UperNet_Model + + model = Swin_UperNet_Model(config) + elif model_type == "bandit": + from models.bandit.core.model import MultiMaskMultiSourceBandSplitRNNSimple + + model = MultiMaskMultiSourceBandSplitRNNSimple(**config.model) + elif model_type == "bandit_v2": + from models.bandit_v2.bandit import Bandit + + model = Bandit(**config.kwargs) + elif model_type == "scnet_unofficial": + from models.scnet_unofficial import SCNet + + model = SCNet(**config.model) + elif model_type == "scnet": + from models.scnet import SCNet + + model = SCNet(**config.model) + elif model_type == "scnet_tran": + from models.scnet.scnet_tran import SCNet_Tran + + model = SCNet_Tran(**config.model) + elif model_type == "apollo": + from models.look2hear.models import BaseModel + + model = BaseModel.apollo(**config.model) + elif model_type == "experimental_mdx23c_stht": + from models.mdx23c_tfc_tdf_v3_with_STHT import TFC_TDF_net + + model = TFC_TDF_net(config) + elif model_type == "scnet_masked": + from models.scnet.scnet_masked import SCNet + + model = SCNet(**config.model) + elif model_type == "conformer": + from models.conformer_model import ConformerMSS, NeuralModel + + model = ConformerMSS( + core=NeuralModel(**config.model), + n_fft=config.stft.n_fft, + hop_length=config.stft.hop_length, + win_length=getattr(config.stft, "win_length", config.stft.n_fft), + center=config.stft.center, + ) + elif model_type == "mel_band_conformer": + from models.mel_band_conformer import MelBandConformer + + model = MelBandConformer(**config.model) + else: + raise ValueError(f"Unknown model type: {model_type}") + + return model, config + + +def get_scheduler(config, optimizer): + scheduler_name = config.training.get("scheduler", "ReduceLROnPlateau") + if scheduler_name == "linear_scheduler": + from transformers import get_linear_schedule_with_warmup + + num_training_steps = config.training.num_epochs * config.training.num_steps + num_warmup_steps = config.training.num_warmup_steps + scheduler = get_linear_schedule_with_warmup( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + ) + elif scheduler_name == "ReduceLROnPlateau": + from torch.optim.lr_scheduler import ReduceLROnPlateau + + scheduler = ReduceLROnPlateau( + optimizer, + "max", + patience=config.training.patience, + factor=config.training.reduce_factor, + ) + else: + available_schedulers = ["linear_scheduler", "ReduceLROnPlateau"] + raise ValueError( + f"Unknown scheduler '{scheduler_name}'. " + f"Available options: {available_schedulers}. " + f"Check your config.training.scheduler setting." + ) + scheduler.name = scheduler_name + return scheduler + + +def logging( + logs: List[str], text: str, verbose_logging: bool = False +) -> Union[List[str], None]: + """ + Print a log message and optionally append it to an in-memory list. + + In Distributed Data Parallel (DDP) contexts, the message is printed only on + rank 0; when DDP is uninitialized, it prints unconditionally. If + `verbose_logging` is True, the message is also appended to `logs`. + + Args: + logs (List[str]): Mutable list to which the message is appended when + `verbose_logging` is True. + text (str): The log message to print (rank 0 only under DDP) and + optionally store. + verbose_logging (bool, optional): If True, append `text` to `logs`. + Defaults to False. + + Returns: + List[str]: The function prints and may mutate `logs` in place. + """ + if not dist.is_initialized() or dist.get_rank() == 0: + print(text) + if verbose_logging: + logs.append(text) + return logs + + +def write_results_in_file(store_dir: str, logs: List[str]) -> None: + """ + Write accumulated log messages to a results file. + + Creates (or overwrites) a `results.txt` file inside `store_dir` and writes + each entry from `logs` as a separate line. In Distributed Data Parallel (DDP) + scenarios, writing is intended to occur only on rank 0. + + Args: + store_dir (str): Directory path where `results.txt` will be saved. + logs (List[str]): Ordered collection of log lines to write. + + Returns: + None + """ + if not dist.is_initialized() or dist.get_rank() == 0: + with open(f"{store_dir}/results.txt", "w") as out: + for item in logs: + out.write(item + "\n") + + +def manual_seed(seed: int) -> None: + """ + Initialize random seeds for reproducibility. + + Sets the seed across Python's `random`, NumPy, and PyTorch (CPU and CUDA) + libraries, and updates the `PYTHONHASHSEED` environment variable. This helps + ensure deterministic behavior where possible, though some GPU operations + may still introduce nondeterminism. + + Args: + seed (int): The seed value to use for all random number generators. + + Returns: + None + """ + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) # if multi-GPU + torch.backends.cudnn.deterministic = False + os.environ["PYTHONHASHSEED"] = str(seed) + + +def initialize_environment(seed: int, results_path: str) -> None: + """ + Initialize runtime environment settings. + + Sets random seeds for reproducibility, adjusts PyTorch cuDNN behavior, + configures multiprocessing with the 'spawn' start method, and ensures + the results directory exists. + + Args: + seed (int): Random seed value for deterministic initialization. + results_path (str): Filesystem path to create for saving results. + + Returns: + None + """ + + manual_seed(seed) + torch.backends.cudnn.deterministic = False + try: + torch.multiprocessing.set_start_method("spawn") + except Exception: + pass + os.makedirs(results_path, exist_ok=True) + + +def initialize_environment_ddp( + rank: int, world_size: int, seed: int = 0, resuls_path: str = None +) -> None: + """ + Initialize environment for Distributed Data Parallel (DDP) training/validation. + + Sets up the DDP process group, seeds random number generators, configures + multiprocessing to use the 'spawn' method, and creates a results directory + if provided. + + Args: + rank (int): Rank of the current process within the DDP group. + world_size (int): Total number of processes participating in DDP. + seed (int, optional): Random seed for reproducibility. Defaults to 0. + resuls_path (str, optional): Directory path to create for storing results. + If None, no directory is created. Defaults to None. + + Returns: + None + """ + seed = (seed + int(time.time())) % 55535 + 10000 + setup_ddp(rank, world_size, seed) + manual_seed(seed) + + try: + torch.multiprocessing.set_start_method( + "spawn", force=True + ) # force=True prevent errors + except RuntimeError as e: + if "context has already been set" not in str(e): + raise e + if resuls_path is not None: + os.makedirs(resuls_path, exist_ok=True) + + +def gen_wandb_name(args, config) -> str: + """ + Generate a descriptive name for a Weights & Biases (wandb) run. + + Combines the model type, a dash-joined list of training instruments, + and the current date into a single string identifier. + + Args: + args: Parsed arguments namespace containing at least `model_type`. + config: Configuration object/dict with a `training.instruments` field. + + Returns: + str: Formatted run name in the form + "_[--...]_". + """ + + instrum = "-".join(config["training"]["instruments"]) + time_str = time.strftime("%Y-%m-%d") + name = "{}_[{}]_{}".format(args.model_type, instrum, time_str) + return name + + +def wandb_init( + args: argparse.Namespace, config: Union[ConfigDict, OmegaConf], batch_size: int +) -> None: + """ + Initialize Weights & Biases (wandb) for experiment tracking. + + Depending on the provided arguments, sets up wandb in one of three modes: + - Offline mode when `args.wandb_offline` is True. + - Disabled mode when no valid `wandb_key` is provided. + - Online mode with authentication using `args.wandb_key`. + + Args: + args (argparse.Namespace): Parsed arguments containing wandb options + (`wandb_offline`, `wandb_key`, `device_ids`). + config (Dict): Experiment configuration dictionary to log. + batch_size (int): Training batch size to include in the run configuration. + + Returns: + None + """ + + if args.wandb_offline: + wandb.init( + mode="offline", + project="msst", + name=gen_wandb_name(args, config), + config={ + "config": config, + "args": args, + "device_ids": args.device_ids, + "batch_size": batch_size, + }, + ) + elif args.wandb_key is None or args.wandb_key.strip() == "": + wandb.init(mode="disabled") + else: + wandb.login(key=args.wandb_key) + wandb.init( + project="msst", + name=gen_wandb_name(args, config), + config={ + "config": config, + "args": args, + "device_ids": args.device_ids, + "batch_size": batch_size, + }, + ) + + +def find_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) # 0 → OS chooses free port + return s.getsockname()[1] + + +def setup_ddp(rank: int, world_size: int, seed: int) -> None: + """ + Initialize a Distributed Data Parallel (DDP) process group. + + Configures environment variables for the DDP master node, attempts to + initialize the process group with the NCCL backend (preferred for GPUs), + and falls back to the Gloo backend if NCCL is unavailable. Also sets the + current CUDA device to match the process rank. + + Args: + rank (int): Rank of the current process in the DDP group. + world_size (int): Total number of processes participating in DDP. + seed: + Returns: + None + """ + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(seed) + os.environ["USE_LIBUV"] = "0" + try: + dist.init_process_group("nccl", rank=rank, world_size=world_size) + except: + dist.init_process_group("gloo", rank=rank, world_size=world_size) + if dist.get_rank() == 0: + print('NCCL are not available. Using "gloo" backend.') + + torch.cuda.set_device(rank) + + +def cleanup_ddp() -> None: + """ + Finalize and clean up a Distributed Data Parallel (DDP) process group. + + Calls `torch.distributed.destroy_process_group()` to release resources + associated with the current DDP environment. + + Returns: + None + """ + dist.destroy_process_group() diff --git a/src/third_party/MusicSourceSeparationTraining/valid.py b/src/third_party/MusicSourceSeparationTraining/valid.py new file mode 100644 index 0000000000000000000000000000000000000000..430258d3218aa02a36d2d09662c898861193b6d4 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/valid.py @@ -0,0 +1,942 @@ +# coding: utf-8 +__author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" + +import argparse +import os +import time +import warnings +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import librosa +import numpy as np +import soundfile as sf +import torch +import torch.distributed as dist +from ml_collections import ConfigDict +from tqdm.auto import tqdm +from utils.audio_utils import ( + denormalize_audio, + draw_2_mel_spectrogram, + normalize_audio, + read_audio_transposed, +) +from utils.metrics import get_metrics +from utils.model_utils import ( + apply_tta, + demix, + load_start_checkpoint, + prefer_target_instrument, +) +from utils.settings import ( + get_model_from_config, + logging, + parse_args_valid, + write_results_in_file, +) + +warnings.filterwarnings("ignore") + + +def get_mixture_paths( + args: argparse.Namespace, verbose: bool, config: ConfigDict, extension: str +) -> List[str]: + """ + Collect validation mixture file paths from one or more root directories. + + Scans each directory in `args.valid_path` for files matching the pattern + `/*/mixture.` and returns a sorted list of absolute paths. + In Distributed Data Parallel (DDP) runs, status messages are printed only + on rank 0; otherwise they are printed unconditionally when `verbose=True`. + + Args: + args (argparse.Namespace): Arguments with `valid_path` (str or List[str]) + specifying root directories to search. + verbose (bool): If True, print collection details and summary. + config (ConfigDict): Configuration used for informational logging + (e.g., `inference.num_overlap`, `inference.batch_size`). + extension (str): Audio file extension to match (with or without a leading dot). + + Returns: + List[str]: Sorted list of discovered mixture file paths. + """ + + ddp_mode = dist.is_initialized() + should_print = (not ddp_mode) or (dist.get_rank() == 0) + + # --- read & normalize args.valid_path --- + try: + valid_path = args.valid_path + except Exception as e: + if should_print: + print("No valid path in args") + raise e + + if isinstance(valid_path, str): + valid_paths: List[str] = [valid_path] + else: + valid_paths = list(valid_path) + + # --- collect mixture files --- + all_mixtures_path: List[str] = [] + + def find_mixture_files(root_dir): + from pathlib import Path + + root_path = Path(root_dir) + wav_files = list(root_path.rglob("mixture.wav")) + flac_files = list(root_path.rglob("mixture.flac")) + if extension not in ["wav", "flac"]: + ext_file = list(root_path.rglob(f"mixture.{extension}")) + else: + ext_file = [] + return wav_files + flac_files + ext_file + + for root in valid_paths: + part = find_mixture_files(root) + if not part and verbose and should_print: + print(f"No validation data found in: {root}") + all_mixtures_path.extend(part) + + # --- verbose summary --- + if verbose and should_print: + # be robust to dict-like or attribute-like config + inference = getattr(config, "inference", None) + if inference is None and isinstance(config, dict): + inference = config.get("inference", None) + + def _get(obj, name, default=None): + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(name, default) + return getattr(obj, name, default) + + num_overlap = _get(inference, "num_overlap", "?") + batch_size = _get(inference, "batch_size", "?") + + print(f"Total mixtures: {len(all_mixtures_path)}") + print(f"Overlap: {num_overlap} Batch size: {batch_size}") + + return all_mixtures_path + + +def update_metrics_and_pbar( + track_metrics: Dict[str, float], + all_metrics: Dict[str, Dict[str, Union[Dict[str, float], List[float]]]], + instr: str, + pbar_dict: Dict[str, float], + mixture_paths: Optional[Union[List[str], tqdm]], + verbose: bool = False, + path: Optional[str] = None, +) -> None: + """ + Update accumulated metrics and (optionally) a tqdm progress bar. + + In non-DDP runs, appends each metric value to `all_metrics[metric_name][instr]` + (a list). In DDP runs (when `torch.distributed` is initialized), stores values + as `all_metrics[metric_name][instr][path]` (a dict keyed by file `path`); + therefore `path` must be provided under DDP. When `verbose=True`, metric + values are printed only on rank 0. Also updates `pbar_dict` and, if a tqdm + instance is provided, calls `set_postfix` for live display. + + Args: + track_metrics (Dict[str, float]): Mapping from metric name to its value for + the current track/instrument. + all_metrics (Dict[str, Dict[str, Union[Dict[str, float], List[float]]]]): + Aggregator for all collected metrics, organized as + `{metric_name: {instrument: list_or_dict}}`, where the inner container + is a list (non-DDP) or dict keyed by `path` (DDP). + instr (str): Instrument name associated with the current metrics. + pbar_dict (Dict[str, float]): Dictionary holding the latest values to show + in the tqdm postfix (updated in place). + mixture_paths (Optional[Union[List[str], tqdm]]): If a tqdm progress bar is + supplied, its `set_postfix` is called with `pbar_dict`. + verbose (bool, optional): If True, print metric updates (rank 0 only in DDP). + Defaults to False. + path (Optional[str], optional): File path key required in DDP mode to index + per-track metrics for `instr`. Ignored in non-DDP. Defaults to None. + + Returns: + None + """ + + ddp_mode = dist.is_initialized() + should_print = (not ddp_mode) or (dist.get_rank() == 0) + + if ddp_mode and path is None: + raise ValueError( + "`path` must be provided when torch.distributed is initialized." + ) + + for metric_name, metric_value in track_metrics.items(): + if verbose and should_print: + print(f"Metric {metric_name:11s} value: {metric_value:.4f}") + + if metric_name not in all_metrics: + all_metrics[metric_name] = {} + if instr not in all_metrics[metric_name]: + all_metrics[metric_name][instr] = {} if ddp_mode else [] + + if ddp_mode: + all_metrics[metric_name][instr][path] = metric_value # type: ignore[index] + else: + all_metrics[metric_name][instr].append(metric_value) # type: ignore[union-attr] + + pbar_dict[f"{metric_name}_{instr}"] = metric_value + + if mixture_paths is not None and hasattr(mixture_paths, "set_postfix"): + try: + mixture_paths.set_postfix(pbar_dict) + except Exception: + pass + + +def process_audio_files( + mixture_paths: List[str], + model: torch.nn.Module, + args: Any, + config: ConfigDict, + device: torch.device, + verbose: bool = False, + is_tqdm: bool = True, +) -> Dict[str, Dict[str, Union[Dict[str, float], List[float]]]]: + """ + Run source separation on a list of mixtures and collect evaluation metrics. + + Performs optional resampling and normalization, demixes each track (with + optional Test-Time Augmentation), saves separated stems (FLAC PCM_16 when + peak ≤ 1.0 else WAV FLOAT), optionally renders spectrograms, computes the + requested metrics, and aggregates them in a nested dictionary. + + In non-DDP runs, metrics are stored as lists: + {metric_name: {instrument: [values...]}} + In DDP runs (when `torch.distributed` is initialized), metrics are stored as + dicts keyed by the track path: + {metric_name: {instrument: {path: value, ...}}} + + Args: + mixture_paths (List[str]): Absolute or relative paths to `mixture.` files. + model (torch.nn.Module): Trained separator model in eval mode. + args (Any): Runtime arguments (e.g., `metrics`, `model_type`, `use_tta`, + `store_dir`, `draw_spectro`, `extension`). + config (ConfigDict): Configuration with audio/inference/training settings + (e.g., `audio.sample_rate`, `inference.batch_size`, `inference.num_overlap`, + `inference.normalize`, `training.instruments`). + device (torch.device): Device for inference (CPU/CUDA). + verbose (bool, optional): Print per-track details and timings. Defaults to False. + is_tqdm (bool, optional): Show a tqdm progress bar (rank 0 only under DDP). Defaults to True. + + Returns: + Dict[str, Dict[str, Union[Dict[str, float], List[float]]]]: Aggregated metrics + per metric and instrument; inner container is a list (non-DDP) or a dict keyed + by track path (DDP). + """ + + ddp_mode = dist.is_initialized() + should_print = (not ddp_mode) or (dist.get_rank() == 0) + + instruments = prefer_target_instrument(config) + use_tta = getattr(args, "use_tta", False) + store_dir = getattr(args, "store_dir", "") + + # extension is used only for reading GT stems; outputs use FLAC/WAV rule unconditionally + if "inference" in config and "extension" in config["inference"]: + extension = config["inference"]["extension"] + else: + extension = getattr(args, "extension", "wav") + + # --- init metrics container --- + if ddp_mode: + # behave like first: dict of dicts + all_metrics: Dict[str, Dict[str, Dict]] = { + metric: {instr: {} for instr in config.training.instruments} + for metric in args.metrics + } + else: + # behave like second: dict of lists + all_metrics: Dict[str, Dict[str, List[float]]] = { + metric: {instr: [] for instr in config.training.instruments} + for metric in args.metrics + } + + # --- tqdm wrapping as requested --- + if is_tqdm and should_print: + mixture_paths = tqdm(mixture_paths) + + def get_instruments(path: str) -> dict[str, str]: + """Detect available instrument files and their extensions.""" + real_instruments: dict[str, str] = {} + + for instr in instruments: + # Check supported extensions for each instrument + for ext in [extension, "flac", "wav"]: + file_path = Path(path) / f"{instr}.{ext}" + if file_path.exists(): + real_instruments[instr] = ext + break + + return real_instruments + + for path in mixture_paths: + start_time = time.time() + mix, sr = read_audio_transposed(path) + mix_orig = mix.copy() + folder = os.path.dirname(path) + real_instruments = get_instruments(folder) + # resample input to config SR if needed + if "audio" in config and "sample_rate" in config.audio: + target_sr = config.audio["sample_rate"] + if sr != target_sr: + orig_length = mix.shape[-1] + if verbose and should_print: + print( + f"Warning: sample rate is different. In config: {target_sr} in file {path}: {sr}" + ) + mix = librosa.resample( + mix, orig_sr=sr, target_sr=target_sr, res_type="kaiser_best" + ) + + if verbose and should_print: + print(f"Song: {os.path.abspath(folder)} Shape: {mix.shape}") + + # optional normalize + if "inference" in config and config.inference.get("normalize", False): + mix, norm_params = normalize_audio(mix) + else: + norm_params = None + + waveforms_orig = demix( + config, model, mix.copy(), device, model_type=args.model_type + ) + + if use_tta: + waveforms_orig = apply_tta( + config, model, mix, waveforms_orig, device, args.model_type + ) + + pbar_dict = {} + + for instr, extension in real_instruments.items(): + if verbose and should_print: + print(f"Instr: {instr}") + + # read GT track + if instr != "other" or not getattr(config.training, "other_fix", False): + track, sr1 = read_audio_transposed( + f"{folder}/{instr}.{extension}", instr, skip_err=True + ) + if track is None: + continue + else: + # other = mix - vocals + track, sr1 = read_audio_transposed(f"{folder}/vocals.{extension}") + track = mix_orig - track + + estimates = waveforms_orig[instr] + + # back-resample estimates to original SR if input was resampled + if "audio" in config and "sample_rate" in config.audio: + target_sr = config.audio["sample_rate"] + if sr != target_sr: + estimates = librosa.resample( + estimates, + orig_sr=target_sr, + target_sr=sr, + res_type="kaiser_best", + ) + estimates = librosa.util.fix_length(estimates, size=orig_length) + + # denormalize if needed + if ( + norm_params is not None + and "inference" in config + and config.inference.get("normalize", False) + ): + estimates = denormalize_audio(estimates, norm_params) + + # --- saving (uniform rule) --- + if store_dir: + os.makedirs(store_dir, exist_ok=True) + base = f"{store_dir}/{os.path.basename(folder)}_{instr}" + peak = float(np.abs(estimates).max()) + if peak <= 1.0: + out_path = f"{base}.flac" + sf.write(out_path, estimates.T, sr, subtype="PCM_16") + else: + out_path = f"{base}.wav" + sf.write(out_path, estimates.T, sr, subtype="FLOAT") + + draw_spec = getattr(args, "draw_spectro", 0) + if draw_spec and draw_spec > 0: + draw_2_mel_spectrogram(estimates.T, track.T, sr, draw_spec, base) + + # --- metrics --- + k = config.training.get("k_sdr", 10) + track_metrics = get_metrics( + args.metrics, track, estimates, mix_orig, device=device, k=k + ) + + # --- update metrics + progress --- + if ddp_mode: + # behave like first: include path in call + update_metrics_and_pbar( + track_metrics, + all_metrics, + instr, + pbar_dict, + mixture_paths=mixture_paths, + verbose=verbose and should_print, + path=path, + ) + else: + # behave like second: no path argument + update_metrics_and_pbar( + track_metrics, + all_metrics, + instr, + pbar_dict, + mixture_paths=mixture_paths, + verbose=verbose and should_print, + ) + + if verbose and should_print: + print(f"Time for song: {time.time() - start_time:.2f} sec") + + return all_metrics + + +def compute_metric_avg( + store_dir: str, + args, + instruments: List[str], + config: ConfigDict, + all_metrics: Dict[str, Dict[str, Union[List[float], Dict[str, float]]]], + start_time: float, +) -> Dict[str, float]: + """ + Compute average metrics across instruments (DDP-aware) and optionally log to file. + + For each metric, computes the mean value per instrument from its collected values + (list in non-DDP, or dict-of-{path: value} in DDP), sums these instrument means, + and divides by `len(instruments)` to obtain the final average (legacy behavior). + Prints/logs only on rank 0 when `torch.distributed` is initialized; if `store_dir` + is non-empty, writes a `results.txt` with logs. + + Args: + store_dir (str): Directory to write `results.txt` when logging is enabled. + args: Run arguments included in the log header when `store_dir` is provided. + instruments (List[str]): Instruments to include in the averaging. + config (ConfigDict): Config used for informational logging (e.g., overlap). + all_metrics (Dict[str, Dict[str, Union[List[float], Dict[str, float]]]]): + Nested metrics container: + - non-DDP: {metric: {instrument: [values...]}} + - DDP: {metric: {instrument: {path: value, ...}}} + start_time (float): Timestamp for reporting elapsed time. + + Returns: + Dict[str, float]: Mapping from metric name to its average over instruments. + """ + + ddp_mode = dist.is_initialized() + should_print = (not ddp_mode) or (dist.get_rank() == 0) + + logs: List[str] = [] + verbose_logging = bool(store_dir) and should_print + if verbose_logging: + logs.append(str(args)) + + logs = logging( + logs, + text=f"Num overlap: {config.inference.num_overlap}", + verbose_logging=verbose_logging, + ) + + metric_sum: Dict[str, float] = {} + + for instr in instruments: + for metric_name in all_metrics: + per_instr_container = all_metrics[ + metric_name + ] # dict: instr -> (list | dict[path->val]) + + values_obj = ( + per_instr_container.get(instr, []) + if isinstance(per_instr_container, dict) + else [] + ) + if isinstance(values_obj, dict): + vals = list(values_obj.values()) + else: + vals = list(values_obj) + + arr = np.asarray(vals, dtype=float) + if arr.size == 0: + mean_val = float("nan") + std_val = float("nan") + else: + mean_val = float(arr.mean()) + std_val = float(arr.std()) + + logs = logging( + logs, + text=f"Instr {instr} {metric_name}: {mean_val:.4f} (Std: {std_val:.4f})", + verbose_logging=verbose_logging, + ) + + metric_sum[metric_name] = metric_sum.get(metric_name, 0.0) + mean_val + + metric_avg: Dict[str, float] = {} + denom = max(len(instruments), 1) + for metric_name in all_metrics: + metric_avg[metric_name] = metric_sum.get(metric_name, float("nan")) / denom + + if len(instruments) > 1: + for metric_name, avg in metric_avg.items(): + logs = logging( + logs, + text=f"Metric avg {metric_name:11s}: {avg:.4f}", + verbose_logging=verbose_logging, + ) + + logs = logging( + logs, + text=f"Elapsed time: {time.time() - start_time:.2f} sec", + verbose_logging=verbose_logging, + ) + + if store_dir: + write_results_in_file(store_dir, logs) + + return metric_avg + + +def valid( + model: torch.nn.Module, + args, + config: ConfigDict, + device: torch.device, + verbose: bool = False, +) -> Tuple[dict, dict]: + """ + Validate a trained model on a set of audio mixtures and compute metrics. + + This function performs validation by separating audio sources from mixtures, + computing evaluation metrics, and optionally saving results to a file. + + Parameters: + ---------- + model : torch.nn.Module + The trained model for source separation. + args : Namespace + Command-line arguments or equivalent object containing configurations. + config : dict + Configuration dictionary with model and processing parameters. + device : torch.device + The device (CPU or CUDA) to run the model on. + verbose : bool, optional + If True, enables verbose output during processing. Default is False. + + Returns: + ------- + dict + A dictionary of average metrics across all instruments. + """ + + start_time = time.time() + model.eval().to(device) + + # dir to save files, if empty no saving + store_dir = getattr(args, "store_dir", "") + # codec to save files + if "extension" in config["inference"]: + extension = config["inference"]["extension"] + else: + extension = getattr(args, "extension", "wav") + + all_mixtures_path = get_mixture_paths(args, verbose, config, extension) + all_metrics = process_audio_files( + all_mixtures_path, model, args, config, device, verbose, not verbose + ) + instruments = prefer_target_instrument(config) + + return compute_metric_avg( + store_dir, args, instruments, config, all_metrics, start_time + ), all_metrics + + +def validate_in_subprocess( + proc_id: int, + queue: torch.multiprocessing.Queue, + all_mixtures_path: List[str], + model: torch.nn.Module, + args, + config: ConfigDict, + device: str, + return_dict, +) -> None: + """ + Perform validation on a subprocess with multi-processing support. Each process handles inference on a subset of the mixture files + and updates the shared metrics dictionary. + + Parameters: + ---------- + proc_id : int + The process ID (used to assign metrics to the correct key in `return_dict`). + queue : torch.multiprocessing.Queue + Queue to receive paths to the mixture files for processing. + all_mixtures_path : List[str] + List of paths to the mixture files to be processed. + model : torch.nn.Module + The model to be used for inference. + args : dict + Dictionary containing various argument configurations (e.g., metrics to calculate). + config : ConfigDict + Configuration object containing model settings and training parameters. + device : str + The device to use for inference (e.g., 'cpu', 'cuda:0'). + return_dict : torch.multiprocessing.Manager().dict + Shared dictionary to store the results from each process. + + Returns: + ------- + None + The function modifies the `return_dict` in place, but does not return any value. + """ + + m1 = model.eval().to(device) + if proc_id == 0: + progress_bar = tqdm(total=len(all_mixtures_path)) + + # Initialize metrics dictionary + all_metrics = { + metric: {instr: [] for instr in config.training.instruments} + for metric in args.metrics + } + + while True: + current_step, path = queue.get() + if path is None: # check for sentinel value + break + single_metrics = process_audio_files( + [path], m1, args, config, device, False, False + ) + pbar_dict = {} + for instr in config.training.instruments: + for metric_name in all_metrics: + all_metrics[metric_name][instr] += single_metrics[metric_name][instr] + if len(single_metrics[metric_name][instr]) > 0: + pbar_dict[f"{metric_name}_{instr}"] = ( + f"{single_metrics[metric_name][instr][0]:.4f}" + ) + if proc_id == 0: + progress_bar.update(current_step - progress_bar.n) + progress_bar.set_postfix(pbar_dict) + # print(f"Inference on process {proc_id}", all_sdr) + return_dict[proc_id] = all_metrics + return + + +def run_parallel_validation( + verbose: bool, + all_mixtures_path: List[str], + config: ConfigDict, + model: torch.nn.Module, + device_ids: List[int], + args, + return_dict, +) -> None: + """ + Run parallel validation using multiple processes. Each process handles a subset of the mixture files and computes the metrics. + The results are stored in a shared dictionary. + + Parameters: + ---------- + verbose : bool + Flag to print detailed information about the validation process. + all_mixtures_path : List[str] + List of paths to the mixture files to be processed. + config : ConfigDict + Configuration object containing model settings and validation parameters. + model : torch.nn.Module + The model to be used for inference. + device_ids : List[int] + List of device IDs (for multi-GPU setups) to use for validation. + args : dict + Dictionary containing various argument configurations (e.g., metrics to calculate). + + Returns: + ------- + A shared dictionary containing the validation metrics from all processes. + """ + + model = model.to("cpu") + try: + # For multiGPU training extract single model + model = model.module + except: + pass + + queue = torch.multiprocessing.Queue() + processes = [] + + for i, device in enumerate(device_ids): + if torch.cuda.is_available(): + device = f"cuda:{device}" + else: + device = "cpu" + p = torch.multiprocessing.Process( + target=validate_in_subprocess, + args=( + i, + queue, + all_mixtures_path, + model, + args, + config, + device, + return_dict, + ), + ) + p.start() + processes.append(p) + for i, path in enumerate(all_mixtures_path): + queue.put((i, path)) + for _ in range(len(device_ids)): + queue.put((None, None)) # sentinel value to signal subprocesses to exit + for p in processes: + p.join() # wait for all subprocesses to finish + + return + + +def block_bounds(num_tracks: int, world_size: int, rank: int) -> Tuple[int, int]: + """ + Split a dataset of `num_tracks` items into `world_size` equal contiguous blocks + and return the half-open interval [start, end) assigned to the given `rank`. + + This function enforces exact divisibility: `num_tracks` must be divisible + by `world_size`, otherwise a ValueError is raised. + + Args: + num_tracks (int): Total number of items to split (must be ≥ 0). + world_size (int): Number of workers to divide the items into (must be > 0). + rank (int): Zero-based worker index (0 ≤ rank < world_size). + + Returns: + Tuple[int, int]: A pair `(start, end)` defining the block of indices for this rank. + + Raises: + ValueError: If `num_tracks` is not divisible by `world_size`. + + Example: + [block_bounds(12, 4, r) for r in range(4)] + [(0, 3), (3, 6), (6, 9), (9, 12)] + + block_bounds(8, 2, 1) + (4, 8) + + block_bounds(10, 3, 0) + Traceback (most recent call last): + ... + ValueError: n (10) must be divisible by world_size (3) + """ + if num_tracks % world_size != 0: + raise ValueError( + f"n ({num_tracks}) must be divisible by world_size ({world_size})" + ) + + chunk = num_tracks // world_size + start = rank * chunk + end = start + chunk + return start, end + + +def valid_multi_gpu( + model: torch.nn.Module, + args, + config: ConfigDict, + device_ids: Optional[List[int]] = None, + verbose: bool = False, +) -> Tuple[Dict[str, float], Dict]: + """ + Validate a separator model across multiple GPUs with a unified API. + + Runs validation either in Distributed Data Parallel (DDP) mode—detected via + `torch.distributed.is_initialized()`—or, if DDP is not active, via + multi-processing / single-GPU execution using the provided `device_ids`. + Collects per-track metrics, aggregates them into per-instrument/per-metric + arrays, and computes per-metric averages. + + Behavior: + * DDP mode: splits the dataset across ranks and gathers metrics; only rank 0 + returns results, while other ranks return `(None, None)`. + * Non-DDP: launches parallel workers when `len(device_ids) > 1`, otherwise + runs on a single device/CPU. + + Args: + model (torch.nn.Module): Trained model to evaluate. + args: Runtime arguments (e.g., metrics list, store dir). + config (ConfigDict): Configuration with inference/training settings. + device_ids (Optional[List[int]]): GPU device IDs for non-DDP parallelism. + If None or length is 1, runs on a single device. + verbose (bool, optional): If True, print progress/logs. Defaults to False. + + Returns: + Tuple[Dict[str, float], Dict]: A pair `(metric_avg, all_metrics)` where + - `metric_avg` maps metric name to its average score, + - `all_metrics` is a nested dict `{metric: {instrument: List[float]}}`. + In DDP mode, non-zero ranks return `(None, None)`. + """ + + start_time = time.time() + + inference = getattr(config, "inference", None) + if inference is None and isinstance(config, dict): + inference = config.get("inference", {}) + extension = getattr(inference, "extension", None) + if extension is None: + if isinstance(inference, dict): + extension = inference.get("extension", getattr(args, "extension", "wav")) + else: + extension = getattr(args, "extension", "wav") + + all_mixtures_path = get_mixture_paths(args, verbose, config, extension) + + ddp_mode = dist.is_initialized() + + if ddp_mode: + rank = dist.get_rank() + world_size = dist.get_world_size() + + device = torch.device(f"cuda:{rank}") + model.to(device) + model.eval() + + num_tracks = len(all_mixtures_path) + pad_needed = (-num_tracks) % world_size + if pad_needed and num_tracks > 0: + all_mixtures_path += all_mixtures_path[:pad_needed] + padded_num_tracks = len(all_mixtures_path) + target_len = padded_num_tracks // world_size + start, end = block_bounds(padded_num_tracks, world_size, rank) + per_rank_data = all_mixtures_path[start:end] + + local_metrics = { + metric: {instr: [] for instr in config.training.instruments} + for metric in args.metrics + } + + with torch.no_grad(): + single_metrics = process_audio_files( + per_rank_data, model, args, config, device, verbose=verbose + ) + for instr in config.training.instruments: + for metric_name in args.metrics: + local_metrics[metric_name][instr] = single_metrics[metric_name][ + instr + ] + + all_metrics: Dict[str, Dict[str, List[float]]] = {m: {} for m in args.metrics} + for metric in args.metrics: + for instr in config.training.instruments: + all_metrics[metric][instr] = [] + per_instr = local_metrics[metric][instr] + if isinstance(per_instr, dict): + local_data = list(per_instr.values()) + else: + local_data = list(per_instr) + + if len(local_data) == 0: + local_tensor = torch.zeros( + target_len, dtype=torch.float32, device=device + ) + else: + if len(local_data) < target_len: + local_data = local_data + [0.0] * (target_len - len(local_data)) + local_tensor = torch.tensor( + local_data, dtype=torch.float32, device=device + ) + + gathered_list = [ + torch.zeros_like(local_tensor) for _ in range(world_size) + ] + dist.all_gather(gathered_list, local_tensor) + + cat_vals = torch.cat(gathered_list).tolist()[:num_tracks] + all_metrics[metric][instr] = cat_vals + + if dist.get_rank() == 0: + instruments = prefer_target_instrument(config) + metric_avg = compute_metric_avg( + getattr(args, "store_dir", ""), + args, + instruments, + config, + all_metrics, + start_time, + ) + return metric_avg, all_metrics + + return None, None + + # Not DDP + store_dir = getattr(args, "store_dir", "") + + return_dict = torch.multiprocessing.Manager().dict() + run_parallel_validation( + verbose, all_mixtures_path, config, model, device_ids, args, return_dict + ) + + all_metrics: Dict[str, Dict[str, List[float]]] = {m: {} for m in args.metrics} + for metric in args.metrics: + for instr in config.training.instruments: + merged: List[float] = [] + for i in range(len(device_ids)): + merged += return_dict[i][metric][instr] + all_metrics[metric][instr] = merged + + instruments = prefer_target_instrument(config) + metric_avg = compute_metric_avg( + store_dir, args, instruments, config, all_metrics, start_time + ) + return metric_avg, all_metrics + + +def check_validation(dict_args): + args = parse_args_valid(dict_args) + torch.backends.cudnn.benchmark = True + try: + torch.multiprocessing.set_start_method("spawn") + except Exception: + pass + model, config = get_model_from_config(args.model_type, args.config_path) + if "model_type" in config.training: + args.model_type = config.training.model_type + if args.start_check_point: + checkpoint = torch.load( + args.start_check_point, weights_only=False, map_location="cpu" + ) + load_start_checkpoint(args, model, checkpoint, type_="valid") + if args.lora_checkpoint_peft: + from peft import PeftModel + + model = PeftModel.from_pretrained(model, args.lora_checkpoint_peft) + model = model.merge_and_unload() + print(f"Instruments: {config.training.instruments}") + + device_ids = args.device_ids + if torch.cuda.is_available(): + device = torch.device(f"cuda:{device_ids[0]}") + else: + device = "cpu" + print("CUDA is not available. Run validation on CPU. It will be very slow...") + + if torch.cuda.is_available() and len(device_ids) > 1: + valid_multi_gpu(model, args, config, device_ids, verbose=False) + else: + valid(model, args, config, device, verbose=True) + + +if __name__ == "__main__": + check_validation(None) diff --git a/src/third_party/MusicSourceSeparationTraining/valid_ddp.py b/src/third_party/MusicSourceSeparationTraining/valid_ddp.py new file mode 100644 index 0000000000000000000000000000000000000000..7c7f1dc1a9b1ae02844f0ed0111d0e413193db95 --- /dev/null +++ b/src/third_party/MusicSourceSeparationTraining/valid_ddp.py @@ -0,0 +1,52 @@ +# coding: utf-8 +__author__ = "Ilya Kiselev (kiselecheck): https://github.com/kiselecheck" +__version__ = "1.0.1" + + +import warnings + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from utils.model_utils import load_start_checkpoint +from utils.settings import ( + get_model_from_config, + initialize_environment_ddp, + parse_args_valid, +) +from valid import valid_multi_gpu + +warnings.filterwarnings("ignore") + + +def check_validation_single(rank: int, world_size: int, args=None): + args = parse_args_valid(args) + + initialize_environment_ddp(rank, world_size) + model, config = get_model_from_config(args.model_type, args.config_path) + + if args.start_check_point: + checkpoint = torch.load( + args.start_check_point, weights_only=False, map_location="cpu" + ) + load_start_checkpoint(args, model, checkpoint, type_="valid") + + if dist.get_rank() == 0: + print(f"Instruments: {config.training.instruments}") + + device = torch.device(f"cuda:{rank}") + model.to(device) + model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) + + valid_multi_gpu(model, args, config, args.device_ids, verbose=False) + + +def check_validation(args=None): + world_size = torch.cuda.device_count() + mp.spawn( + check_validation_single, args=(world_size, args), nprocs=world_size, join=True + ) + + +if __name__ == "__main__": + check_validation()