from dataclasses import dataclass import torch from tqdm.auto import trange import typing as tp from einops import rearrange from torch import nn from .autoencoders import AudioAutoencoder from .conditioners import MultiConditioner, create_multi_conditioner_from_conditioning_config from .factory import create_pretransform_from_config from .lm_backbone import AudioLMBackbone, XTransformersAudioLMBackbone, ContinuousTransformerAudioLMBackbone from .pretransforms import Pretransform, AutoencoderPretransform, PretrainedDACPretransform, AudiocraftCompressionPretransform from .utils import multinomial, sample_top_k, sample_top_p from .codebooks_patterns import ( CodebooksPatternProvider, DelayedPatternProvider, MusicLMPattern, ParallelPatternProvider, UnrolledPatternProvider, VALLEPattern, ) # Copied and modified from https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/models/lm.py under MIT license # License can be found in LICENSES/LICENSE_META.txt @dataclass class LMOutput: # The logits are already re-aligned with the input codes # hence no extra shift is required, e.g. when computing CE logits: torch.Tensor # [B, K, T, card] mask: torch.Tensor # [B, K, T] # Wrapper for a multi-codebook language model # Handles patterns and quantizer heads class AudioLanguageModel(nn.Module): def __init__( self, pattern_provider: CodebooksPatternProvider, backbone: AudioLMBackbone, num_quantizers: int, codebook_size: int ): super().__init__() self.pattern_provider = pattern_provider self.backbone = backbone self.num_quantizers = num_quantizers self.codebook_size = codebook_size self.masked_token_id = codebook_size # Per-quantizer embedders # Add one for the mask embed self.embeds = nn.ModuleList([nn.Embedding(codebook_size + 1, backbone.embed_dim) for _ in range(num_quantizers)]) # Per-quantizer output heads self.quantizer_heads = nn.ModuleList([ nn.Linear(backbone.embed_dim, codebook_size) for _ in range(num_quantizers) ]) def forward(self, sequence: torch.Tensor, #[batch, seq_len, prepend_cond=None, #[batch, seq, channels] prepend_cond_mask=None, cross_attn_cond=None, #[batch, seq, channels], **kwargs ): batch, num_quantizers, seq_len = sequence.shape assert num_quantizers == self.num_quantizers, "Number of quantizers in sequence must match number of quantizers in model" backbone_input = sum([self.embeds[i](sequence[:, i]) for i in range(num_quantizers)]) # [batch, seq_len, embed_dim] output = self.backbone( backbone_input, cross_attn_cond=cross_attn_cond, prepend_cond=prepend_cond, prepend_cond_mask=prepend_cond_mask, **kwargs ) # [batch, seq_len, embed_dim] # Run output through quantizer heads logits = torch.stack([self.quantizer_heads[i](output) for i in range(num_quantizers)], dim=1) # [batch, num_quantizers, seq_len, codebook_size] return logits def compute_logits( self, codes, #[batch, num_quantizers, seq_len] **kwargs): """ Compute logits for a batch of codes, optionally conditioning on cross-attention and prepend conditioning Handles translation between input sequence and pattern-shifted sequence Only used during training """ batch, _, seq_len = codes.shape pattern = self.pattern_provider.get_pattern(seq_len) # Apply the token pattern to the codes, shifting the codes as needed and masking out invalid steps shifted_codes, _, _ = pattern.build_pattern_sequence( codes, self.masked_token_id, keep_only_valid_steps=True ) # Run the model to get logits for each quantizer [batch, num_quantizers, seq_len, codebook_size] logits = self(shifted_codes, **kwargs) # Rearrange logits to prepare to revert pattern logits = rearrange(logits, "b n s c -> b c n s") # Revert sequence logits back to original sequence length, removing masked steps logits, _, logits_mask = pattern.revert_pattern_logits( logits, float('nan'), keep_only_valid_steps=True ) logits = rearrange(logits, "b c n t -> b n t c") logits_mask = logits_mask[None, :, :].expand(batch, -1, -1) # [batch, num_quantizers, seq_len] return LMOutput(logits=logits, mask=logits_mask) # Conditioning and generation wrapper for a multi-codebook language model # Handles conditioning, CFG, generation, and encoding/decoding class AudioLanguageModelWrapper(nn.Module): def __init__( self, pretransform: Pretransform, lm: AudioLanguageModel, sample_rate: int, min_input_length: int, conditioner: MultiConditioner = None, cross_attn_cond_ids: tp.List[str] = [], prepend_cond_ids: tp.List[str] = [], global_cond_ids: tp.List[str] = [] ): super().__init__() assert pretransform.is_discrete, "Pretransform must be discrete" self.pretransform = pretransform self.pretransform.requires_grad_(False) self.pretransform.eval() if isinstance(self.pretransform, AutoencoderPretransform): self.num_quantizers = self.pretransform.model.bottleneck.num_quantizers self.codebook_size = self.pretransform.model.bottleneck.codebook_size elif isinstance(self.pretransform, PretrainedDACPretransform): self.num_quantizers = self.pretransform.model.num_quantizers self.codebook_size = self.pretransform.model.codebook_size elif isinstance(self.pretransform, AudiocraftCompressionPretransform): self.num_quantizers = self.pretransform.num_quantizers self.codebook_size = self.pretransform.codebook_size else: raise NotImplementedError(f"Unrecognized pretransform type {type(self.pretransform)}") self.conditioner = conditioner self.lm = lm self.sample_rate = sample_rate self.min_input_length = min_input_length self.cross_attn_cond_ids = cross_attn_cond_ids self.prepend_cond_ids = prepend_cond_ids self.global_cond_ids = global_cond_ids def get_conditioning_inputs(self, cond: tp.Dict[str, tp.Any], negative=False): cross_attention_input = None prepend_cond = None prepend_cond_mask = None global_cond = 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 = torch.cat([cond[key][0] for key in self.cross_attn_cond_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_cond = torch.cat([cond[key][0] for key in self.prepend_cond_ids], dim=1) prepend_cond_mask = torch.cat([cond[key][1] for key in self.prepend_cond_ids], 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_cond = torch.cat([cond[key][0] for key in self.global_cond_ids], dim=-1) if len(global_cond.shape) == 3: global_cond = global_cond.squeeze(1) if negative: return { "negative_cross_attn_cond": cross_attention_input, "negative_prepend_cond": prepend_cond, "negative_prepend_cond_mask": prepend_cond_mask, "negative_global_cond": global_cond } else: return { "cross_attn_cond": cross_attention_input, "prepend_cond": prepend_cond, "prepend_cond_mask": prepend_cond_mask, "global_cond": global_cond } def compute_logits( self, codes, condition_tensors=None, cfg_dropout_prob=0.0, **kwargs ): """ Compute logits for a batch of codes, and translates from conditioning inputs to model inputs Handles CFG dropout """ if condition_tensors is None: condition_tensors = {} conditioning_inputs = self.get_conditioning_inputs(condition_tensors) cross_attn_cond = conditioning_inputs["cross_attn_cond"] prepend_cond = conditioning_inputs["prepend_cond"] prepend_cond_mask = conditioning_inputs["prepend_cond_mask"] global_cond = conditioning_inputs["global_cond"] 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 global_cond is not None: null_embed = torch.zeros_like(global_cond, device=global_cond.device) dropout_mask = torch.bernoulli(torch.full((global_cond.shape[0], 1), cfg_dropout_prob, device=global_cond.device)).to(torch.bool) global_cond = torch.where(dropout_mask, null_embed, global_cond) return self.lm.compute_logits(codes, cross_attn_cond=cross_attn_cond, prepend_cond=prepend_cond, prepend_cond_mask=prepend_cond_mask, global_cond=global_cond, **kwargs) def _sample_next_token( self, sequence, #[batch, num_quantizers, seq_len] conditioning_tensors=None, cross_attn_use_cfg=True, prepend_use_cfg=True, global_use_cfg=True, cfg_scale=1.0, top_k=250, top_p=0.0, temp=1.0, **kwargs ): """ Sample the next token for a batch of codes, and translates from conditioning inputs to model inputs Handles CFG inference """ if conditioning_tensors is None: conditioning_tensors = {} conditioning_inputs = self.get_conditioning_inputs(conditioning_tensors) cross_attn_cond = conditioning_inputs["cross_attn_cond"] prepend_cond = conditioning_inputs["prepend_cond"] prepend_cond_mask = conditioning_inputs["prepend_cond_mask"] global_cond = conditioning_inputs["global_cond"] if cfg_scale != 1.0: # Batch size is doubled to account for negative samples sequence = torch.cat([sequence, sequence], dim=0) if cross_attn_cond is not None and cross_attn_use_cfg: null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device) cross_attn_cond = torch.cat([cross_attn_cond, null_embed], dim=0) if prepend_cond is not None and prepend_use_cfg: null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device) prepend_cond = torch.cat([prepend_cond, null_embed], dim=0) if prepend_cond_mask is not None: prepend_cond_mask = torch.cat([prepend_cond_mask, prepend_cond_mask], dim=0) if global_cond is not None and global_use_cfg: null_embed = torch.zeros_like(global_cond, device=global_cond.device) global_cond = torch.cat([global_cond, null_embed], dim=0) logits = self.lm(sequence, cross_attn_cond=cross_attn_cond, prepend_cond=prepend_cond, prepend_cond_mask=prepend_cond_mask, global_cond=global_cond, **kwargs) if cfg_scale != 1.0: cond_logits, uncond_logits = logits.chunk(2, dim=0) logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale logits = rearrange(logits, "b n s c -> b n c s") # [batch, num_quantizers, codebook_size, seq_len] # Grab the logits for the last step logits = logits[:, :, :, -1] # [batch, num_quantizers, codebook_size] # Apply top-k or top-p sampling if temp > 0: probs = torch.softmax(logits / temp, dim=-1) if top_p > 0.0: next_token = sample_top_p(probs, p=top_p) elif top_k > 0: next_token = sample_top_k(probs, k=top_k) else: next_token = multinomial(probs, num_samples=1) else: next_token = torch.argmax(logits, dim=-1, keepdim=True) # [batch, num_quantizers, 1] return next_token @torch.no_grad() def generate( self, max_gen_len: int = 256, batch_size: tp.Optional[int] = None, init_data: tp.Optional[torch.Tensor] = None, conditioning: tp.Optional[tp.Dict[str, tp.Any]] = None, conditioning_tensors: tp.Optional[tp.Dict[str, tp.Any]] = None, callback: tp.Optional[tp.Callable[[int, int], None]] = None, use_cache: bool = True, cfg_scale: float = 1.0, **kwargs ): device = next(self.parameters()).device if conditioning_tensors is None and conditioning is not None: # Convert conditioning inputs to conditioning tensors conditioning_tensors = self.conditioner(conditioning, device) # Check that batch size is consistent across inputs possible_batch_sizes = [] if batch_size is not None: possible_batch_sizes.append(batch_size) elif init_data is not None: possible_batch_sizes.append(init_data.shape[0]) elif conditioning_tensors is not None: # Assume that the first conditioning tensor has the batch dimension possible_batch_sizes.append(conditioning_tensors[list(conditioning_tensors.keys())[0]][0].shape[0]) else: possible_batch_sizes.append(1) assert [x == possible_batch_sizes[0] for x in possible_batch_sizes], "Batch size must be consistent across inputs" batch_size = possible_batch_sizes[0] if init_data is None: # Initialize with zeros assert batch_size > 0 init_data = torch.zeros((batch_size, self.num_quantizers, 0), device=device, dtype=torch.long) batch_size, num_quantizers, seq_len = init_data.shape start_offset = seq_len assert start_offset < max_gen_len, "init data longer than max gen length" pattern = self.lm.pattern_provider.get_pattern(max_gen_len) unknown_token = -1 # Initialize the generated codes with the init data, padded with unknown tokens gen_codes = torch.full((batch_size, num_quantizers, max_gen_len), unknown_token, device=device, dtype=torch.long) gen_codes[:, :, :start_offset] = init_data # [batch, num_quantizers, max_gen_len] gen_sequence, _, mask = pattern.build_pattern_sequence(gen_codes, self.lm.masked_token_id) # [batch, num_quantizers, gen_sequence_len] start_offset_sequence = pattern.get_first_step_with_timesteps(start_offset) assert start_offset_sequence is not None # Generation prev_offset = 0 gen_sequence_len = gen_sequence.shape[-1] # Reset generation cache if use_cache and self.lm.backbone.use_generation_cache: self.lm.backbone.reset_generation_cache(max_gen_len, batch_size if cfg_scale == 1.0 else batch_size * 2) for offset in trange(start_offset_sequence, gen_sequence_len): # Get the full sequence up to the current offset curr_sequence = gen_sequence[..., prev_offset:offset] next_token = self._sample_next_token( curr_sequence, conditioning_tensors=conditioning_tensors, use_cache=use_cache, cfg_scale=cfg_scale, **kwargs ) valid_mask = mask[..., offset:offset+1].expand(batch_size, -1, -1) next_token[~valid_mask] = self.lm.masked_token_id # Update the generated sequence with the next token gen_sequence[..., offset:offset+1] = torch.where( gen_sequence[..., offset:offset+1] == unknown_token, next_token, gen_sequence[..., offset:offset+1] ) if use_cache and self.lm.backbone.use_generation_cache: # Only update the offset if caching is being used prev_offset = offset self.lm.backbone.update_generation_cache(offset) if callback is not None: # Callback to report progress # Pass in the offset relative to the start of the sequence, and the length of the current sequence callback(1 + offset - start_offset_sequence, gen_sequence_len - start_offset_sequence) assert not (gen_sequence == unknown_token).any(), "Unknown tokens in generated sequence" out_codes, _, out_mask = pattern.revert_pattern_sequence(gen_sequence, special_token=unknown_token) # sanity checks over the returned codes and corresponding masks assert (out_codes[..., :max_gen_len] != unknown_token).all() assert (out_mask[..., :max_gen_len] == 1).all() #out_codes = out_codes[..., 0:max_gen_len] return out_codes def generate_audio( self, **kwargs ): """ Generate audio from a batch of codes """ codes = self.generate(**kwargs) audio = self.pretransform.decode_tokens(codes) return audio def create_audio_lm_from_config(config): model_config = config.get('model', None) assert model_config is not None, 'model config must be specified in config' sample_rate = config.get('sample_rate', None) assert sample_rate is not None, "Must specify sample_rate in config" lm_config = model_config.get('lm', None) assert lm_config is not None, 'lm config must be specified in model config' codebook_pattern = lm_config.get("codebook_pattern", "delay") pattern_providers = { 'parallel': ParallelPatternProvider, 'delay': DelayedPatternProvider, 'unroll': UnrolledPatternProvider, 'valle': VALLEPattern, 'musiclm': MusicLMPattern, } pretransform_config = model_config.get("pretransform", None) pretransform = create_pretransform_from_config(pretransform_config, sample_rate) assert pretransform.is_discrete, "Pretransform must be discrete" min_input_length = pretransform.downsampling_ratio pattern_provider = pattern_providers[codebook_pattern](n_q=pretransform.num_quantizers) 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_attn_cond_ids = lm_config.get('cross_attention_cond_ids', []) prepend_cond_ids = lm_config.get('prepend_cond_ids', []) global_cond_ids = lm_config.get('global_cond_ids', []) lm_type = lm_config.get("type", None) lm_model_config = lm_config.get("config", None) assert lm_type is not None, "Must specify lm type in lm config" assert lm_model_config is not None, "Must specify lm model config in lm config" if lm_type == "x-transformers": backbone = XTransformersAudioLMBackbone(**lm_model_config) elif lm_type == "continuous_transformer": backbone = ContinuousTransformerAudioLMBackbone(**lm_model_config) else: raise NotImplementedError(f"Unrecognized lm type {lm_type}") lm = AudioLanguageModel( pattern_provider=pattern_provider, backbone=backbone, num_quantizers=pretransform.num_quantizers, codebook_size=pretransform.codebook_size ) model = AudioLanguageModelWrapper( pretransform=pretransform, lm=lm, conditioner=conditioner, sample_rate=sample_rate, min_input_length=min_input_length, cross_attn_cond_ids=cross_attn_cond_ids, prepend_cond_ids=prepend_cond_ids, global_cond_ids=global_cond_ids ) return model