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