# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. from functools import partial import logging import math import typing as tp import torch from torch import nn from ..utils import utils from ..modules.transformer import StreamingTransformer, create_norm_fn from ..utils.renoise import noise_regularization from ..modules.conditioners import ( ConditionFuser, ClassifierFreeGuidanceDropout, ConditioningProvider, ConditioningAttributes, ConditionType, ) from ..modules.activations import get_activation_fn logger = logging.getLogger(__name__) ConditionTensors = tp.Dict[str, ConditionType] CFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]] def get_init_fn( method: str, input_dim: int, init_depth: tp.Optional[int] = None ) -> partial[torch.Tensor]: """LM layer initialization. Inspired from xlformers: https://github.com/fairinternal/xlformers Args: method (str): Method name for init function. Valid options are: 'gaussian', 'uniform'. input_dim (int): Input dimension of the initialized module. init_depth (Optional[int]): Optional init depth value used to rescale the standard deviation if defined. """ # Compute std std = 1 / math.sqrt(input_dim) # Rescale with depth if init_depth is not None: std = std / math.sqrt(2 * init_depth) if method == "gaussian": return partial( torch.nn.init.trunc_normal_, mean=0.0, std=std, a=-3 * std, b=3 * std ) elif method == "uniform": bound = math.sqrt(3) * std # ensure the standard deviation is `std` return partial(torch.nn.init.uniform_, a=-bound, b=bound) else: raise ValueError("Unsupported layer initialization method") def init_layer( m: torch.nn.Module, method: str, init_depth: tp.Optional[int] = None, zero_bias_init: bool = False, ) -> None: """Wrapper around ``get_init_fn`` for proper initialization of LM modules. Args: m (nn.Module): Module to initialize. method (str): Method name for the init function. init_depth (Optional[int]): Optional init depth value used to rescale the standard deviation if defined. zero_bias_init (bool): Whether to initialize the bias to 0 or not. """ if isinstance(m, torch.nn.Linear): init_fn = get_init_fn(method, m.in_features, init_depth=init_depth) init_fn(m.weight) if zero_bias_init and m.bias is not None: torch.nn.init.constant_(m.bias, 0) elif isinstance(m, torch.nn.Embedding): init_fn = get_init_fn(method, m.embedding_dim, init_depth=None) init_fn(m.weight) class TimestepEmbedding(torch.nn.Module): """ Embeds scalar timesteps into vector representations. """ def __init__(self, hidden_size: int, frequency_embedding_size: int = 256, bias_proj: bool = False, **kwargs) -> None: super().__init__() self.mlp = torch.nn.Sequential( torch.nn.Linear(frequency_embedding_size, hidden_size, bias=bias_proj, **kwargs), torch.nn.SiLU(), torch.nn.Linear(hidden_size, hidden_size, bias=bias_proj, **kwargs), ) self.frequency_embedding_size: int = frequency_embedding_size @staticmethod def timestep_embedding( t: torch.Tensor, dim: int, max_period: int = 10000 ) -> torch.Tensor: """ Create sinusoidal timestep embeddings. :param t: a 1-D Tensor of N indices, one per batch element. These may be fractional. :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. :return: an (N, D) Tensor of positional embeddings. """ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half ) args = t[:, 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 forward(self, t: torch.Tensor) -> torch.Tensor: t_freq = self.timestep_embedding(t, self.frequency_embedding_size) t_emb = self.mlp(t_freq) return t_emb class FlowModel(torch.nn.Module): """Transformer-based flow model operating on continuous audio latents. Args: condition_provider (MusicConditioningProvider): Conditioning provider from metadata. fuser (ConditionFuser): Fuser handling the fusing of conditions with language model input. dim (int): Dimension of the transformer encoder. latent_dim (int): Dimension of the latent audio representation. num_heads (int): Number of heads for the transformer encoder. hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder. norm (str): Normalization method. norm_first (bool): Use pre-norm instead of post-norm. bias_proj (bool): Use bias for output projections. weight_init (str, optional): Method for weight initialization. depthwise_init (str, optional): Method for depthwise weight initialization. zero_bias_init (bool): If true and bias in Linears, initialize bias to zeros. cfg_coef (float): Classifier-free guidance coefficient. **kwargs: Additional parameters for the transformer encoder. """ def __init__( self, condition_provider: ConditioningProvider, fuser: ConditionFuser, dim: int = 128, latent_dim: int = 128, num_heads: int = 8, hidden_scale: int = 4, norm: str = 'layer_norm', norm_first: bool = False, bias_proj: bool = False, weight_init: tp.Optional[str] = None, depthwise_init: tp.Optional[str] = None, zero_bias_init: bool = False, cfg_coef: float = 1.0, **kwargs ): super().__init__() # self.cfg = cfg # self.device = device self.latent_dim = latent_dim self.timestep_embedder = TimestepEmbedding(dim, bias_proj=bias_proj, device=kwargs["device"]) if 'activation' in kwargs: kwargs['activation'] = get_activation_fn(kwargs['activation']) self.transformer = StreamingTransformer( d_model=dim, num_heads=num_heads, dim_feedforward=int(hidden_scale * dim), norm=norm, norm_first=norm_first, **kwargs) self.condition_provider: ConditioningProvider = condition_provider self.fuser: ConditionFuser = fuser self.out_norm: tp.Optional[nn.Module] = None if norm_first: self.out_norm = create_norm_fn(norm, dim, device=kwargs["device"]) self.in_proj = torch.nn.Linear(latent_dim, dim, bias=bias_proj, device=kwargs["device"]) self.out_proj = torch.nn.Linear(dim, latent_dim, bias=bias_proj, device=kwargs["device"]) self._init_weights(weight_init, depthwise_init, zero_bias_init) self.cfg_coef = cfg_coef # statistics of latent audio representations (tracked during codec training) self.register_buffer( "latent_mean", torch.zeros(1, latent_dim, 1, device=kwargs["device"]) ) self.register_buffer( "latent_std", torch.ones(1, latent_dim, 1, device=kwargs["device"]) ) def _init_weights( self, weight_init: tp.Optional[str], depthwise_init: tp.Optional[str], zero_bias_init: bool, ) -> None: """Initialization of the transformer module weights. This initialization schema is still experimental and may be subject to changes. Args: weight_init (Optional[str]): Weight initialization strategy. See ``get_init_fn`` for valid options. depthwise_init (Optional[str]): Depwthwise initialization strategy. The following options are valid: 'current' where the depth corresponds to the current layer index or 'global' where the total number of layer is used as depth. If not set, no depthwise initialization strategy is used. zero_bias_init (bool): Whether to initialize bias to zero or not. """ assert depthwise_init is None or depthwise_init in ["current", "global"] assert ( depthwise_init is None or weight_init is not None ), "If 'depthwise_init' is defined, a 'weight_init' method should be provided." assert ( not zero_bias_init or weight_init is not None ), "If 'zero_bias_init', a 'weight_init' method should be provided" if weight_init is None: return # for emb_layer in self.emb: init_layer( self.in_proj, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init, ) for layer_idx, tr_layer in enumerate(self.transformer.layers): depth = None if depthwise_init == "current": depth = layer_idx + 1 elif depthwise_init == "global": depth = len(self.transformer.layers) init_fn = partial( init_layer, method=weight_init, init_depth=depth, zero_bias_init=zero_bias_init, ) tr_layer.apply(init_fn) # for linear in self.linears: init_layer( self.out_proj, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init, ) for layer in self.timestep_embedder.mlp: if isinstance(layer, torch.nn.Linear): init_layer( layer, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init, ) for linear in self.transformer.skip_projections: init_layer( linear, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init, ) def forward( self, z: torch.Tensor, t: torch.Tensor, condition_src: torch.Tensor, condition_mask: torch.Tensor, ) -> torch.Tensor: # Rescale input so that its std remains 1 whatever the flow step t. x = z / torch.sqrt(torch.pow(t.unsqueeze(-1), 2) + torch.pow(1 - t.unsqueeze(-1), 2)) x = self.in_proj(x.permute(0, 2, 1)) x = self.transformer( x, cross_attention_src=condition_src, cross_attention_mask=condition_mask, timestep_embedding=self.timestep_embedder(t), ) if self.out_norm is not None: x = self.out_norm(x) # Rescale output so that the DiT output std remains 1. x = self.out_proj(x).permute(0, 2, 1) * math.sqrt(2) return x @torch.no_grad() def generate(self, prompt: tp.Optional[torch.Tensor] = None, conditions: tp.List[ConditioningAttributes] = [], solver: str = "euler", steps: int = 16, num_samples: tp.Optional[int] = None, max_gen_len: int = 256, source_flowstep: int = 0.0, target_flowstep: int = 1.0, regularize: bool = False, regularize_iters: int = 4, keep_last_k_iters: int = 2, lambda_kl: float = 0.2, callback: tp.Optional[tp.Callable[[int, int], None]] = None, sway_coefficient: float = -0.8, ) -> torch.Tensor: """Generate tokens by running the ODE solver from source flow step to target flow step, either given a prompt or unconditionally. Args: prompt (torch.Tensor, optional): Prompt tokens of shape [B, K, T] (for editing). conditions (list of ConditioningAttributes, optional): List of conditions. solver (str): ODE solver (either euler or midpoint) steps (int): number of solver steps. num_samples (int, optional): Number of samples to generate when no prompt and no conditions are given. max_gen_len (int): Maximum generation length. source_flowstep (float): Source flow step (0 for generation, 1 for inversion). target_flowstep (float): Target flow step (1 for generation, 0 for inversion). regularize (bool): Regularize each solver step. regularize_iters (int): Number of regularization iterations. keep_last_k_iters (int): Number of meaningful regularization iterations for moving average computation. lambda_kl (float): KL regularization loss weight. sway_coefficient (float): sway sampling coefficient (https://arxiv.org/pdf/2410.06885), should be set to 0 for uniform. callback (Callback, optional): Callback function to report generation progress. Returns: torch.Tensor: Generated tokens. """ assert not self.training, "generation shouldn't be used in training mode." first_param = next(iter(self.parameters())) device = first_param.device assert solver in ["euler", "midpoint"], "Supported ODE solvers are either euler or midpoint!" assert not (regularize and solver == "midpoint"), "Latent regularization is only supported with euler solver!" assert keep_last_k_iters <= regularize_iters # Checking all input shapes are consistent. possible_num_samples = [] if num_samples is not None: possible_num_samples.append(num_samples) elif prompt is not None: possible_num_samples.append(prompt.shape[0]) elif conditions: possible_num_samples.append(len(conditions)) else: possible_num_samples.append(1) assert [x == possible_num_samples[0] for x in possible_num_samples], "Inconsistent inputs shapes" num_samples = possible_num_samples[0] if prompt is not None and self.latent_mean.shape[1] != prompt.shape[1]: # tokens directly emanate from the VAE encoder mean, scale = prompt.chunk(2, dim=1) gen_tokens = utils.vae_sample(mean, scale) prompt = (gen_tokens - self.latent_mean) / (self.latent_std + 1e-5) # below we create set of conditions: one conditional and one unconditional # to do that we merge the regular condition together with the null condition # we then do 1 forward pass instead of 2. # the reason for that is two-fold: # 1. it is about x2 faster than doing 2 forward passes # 2. avoid the streaming API treating the 2 passes as part of different time steps # We also support doing two different passes, in particular to ensure that # the padding structure is exactly the same between train and test. # With a batch size of 1, this can be slower though. cfg_conditions: CFGConditions if conditions: null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions) conditions = conditions + null_conditions tokenized = self.condition_provider.tokenize(conditions) cfg_conditions = self.condition_provider(tokenized) else: cfg_conditions = {} gen_sequence = torch.randn(num_samples, self.latent_dim, max_gen_len, device=device) if prompt is None else prompt B = gen_sequence.shape[0] if solver == "midpoint": assert steps % 2 == 0, "Midpoint solver can only run with even number of steps" next_sequence: tp.Optional[torch.Tensor] = None if not regularize: regularize_iters = 1 keep_last_k_iters = 0 else: avg_regularized_velocity = torch.zeros_like(gen_sequence) regularization_iters_threshold = regularize_iters - keep_last_k_iters cfg_coef = 0.0 if target_flowstep < source_flowstep else self.cfg_coef # prevent divergence during latent inversion if target_flowstep > source_flowstep: schedule = torch.arange(0, 1+1e-5, 1/steps) else: schedule = torch.arange(1, 0-1e-5, -1/steps) schedule = schedule + sway_coefficient * (torch.cos(math.pi * 0.5 * schedule) - 1 + schedule) if target_flowstep > source_flowstep: schedule = schedule * (target_flowstep-source_flowstep) + source_flowstep else: schedule = schedule * (source_flowstep-target_flowstep) + target_flowstep for idx, current_flowstep in enumerate(schedule[:-1]): delta_t = schedule[idx+1] - current_flowstep input_sequence = gen_sequence if solver == "midpoint" and idx % 2 == 1: input_sequence = next_sequence for jdx in range(regularize_iters): should_compute_kl = jdx >= regularization_iters_threshold if should_compute_kl: regularizing_sequence = ( 1 - schedule[idx+1] ) * torch.randn_like(input_sequence) + schedule[idx+1] * prompt + 1e-5 * torch.randn_like(input_sequence) input_sequence = torch.cat( [ input_sequence, regularizing_sequence, ], dim=0, ) if cfg_coef == 0.0: predicted_velocity = self.forward(input_sequence, torch.tensor([schedule[idx+1] if regularize else current_flowstep], device=device).expand(B * (1+should_compute_kl), 1), cfg_conditions['description'][0][:B].repeat(1+should_compute_kl, 1, 1), torch.log(cfg_conditions['description'][1][:B].unsqueeze(1).unsqueeze(1).repeat(1+should_compute_kl, 1, 1, 1))) velocity = predicted_velocity[:B] if should_compute_kl: regularizing_velocity = predicted_velocity[B:] else: predicted_velocity = self.forward(input_sequence.repeat(2, 1, 1), torch.tensor([schedule[idx+1] if regularize else current_flowstep], device=device).expand(B*2*(1 + should_compute_kl), 1), cfg_conditions['description'][0].repeat_interleave(1 + should_compute_kl, dim=0), torch.log(cfg_conditions['description'][1].unsqueeze(1).unsqueeze(1)).repeat_interleave(1 + should_compute_kl, dim=0)) if should_compute_kl: velocity = (1 + cfg_coef) * predicted_velocity[:B] - cfg_coef * predicted_velocity[2*B:3*B] regularizing_velocity = (1 + cfg_coef) * predicted_velocity[B:2*B] - cfg_coef * predicted_velocity[3*B:4*B] else: velocity = (1 + cfg_coef) * predicted_velocity[:B] - cfg_coef * predicted_velocity[B:2*B] if should_compute_kl: regularized_velocity = noise_regularization( velocity, regularizing_velocity, lambda_kl=lambda_kl, lambda_ac=0.0, num_reg_steps=4, num_ac_rolls=5, ) avg_regularized_velocity += ( regularized_velocity * jdx / ( keep_last_k_iters * regularization_iters_threshold + sum(range(keep_last_k_iters)) ) ) if should_compute_kl: input_sequence = gen_sequence + regularized_velocity * delta_t else: input_sequence = gen_sequence + velocity * delta_t if callback is not None: callback(1 + idx * regularize_iters + jdx, steps * regularize_iters) if regularize: velocity = avg_regularized_velocity avg_regularized_velocity = torch.zeros_like(gen_sequence) if solver == "midpoint": if idx % 2 == 0: next_sequence = gen_sequence + velocity * delta_t else: gen_sequence = gen_sequence + velocity * (schedule[idx+1] - schedule[idx-1]) else: gen_sequence = gen_sequence + velocity * delta_t return gen_sequence