# Copyright (c) 2024 Alibaba Inc # # 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. from typing import Dict, Optional, Callable, List, Generator import torch from torch import nn from torch.nn.utils.rnn import pad_sequence, unpad_sequence from inspiremusic.utils.common import IGNORE_ID from inspiremusic.transformer.label_smoothing_loss import LabelSmoothingLoss from inspiremusic.utils.common import th_accuracy from torch import Tensor from math import log from einops import rearrange, reduce, repeat import logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') 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).to(torch.float16) class LLM(torch.nn.Module): def __init__( self, text_encoder_input_size: int, llm_input_size: int, llm_output_size: int, audio_token_size: int, llm: torch.nn.Module, sampling: Callable, text_encoder_conf: Dict = None, length_normalized_loss: bool = True, lsm_weight: float = 0.0, frozen_input_embed: bool = False, **kwargs, ): super().__init__() self.llm_input_size = llm_input_size self.audio_token_size = audio_token_size # 1. build text token inputs related modules if llm is None: self.text_embedding = torch.nn.Embedding(text_token_size, text_encoder_input_size) else: self.text_embedding = llm.model.model.embed_tokens if frozen_input_embed: print("Freezing input embedding layer") for p in self.text_embedding.parameters(): p.requires_grad = False self.chorus_embedding = torch.nn.Embedding(5, llm_input_size) # intro, chorus, verse1, verse2 , outro self.text_encoder_conf = text_encoder_conf self.text_encoder = self.build_encoder(text_encoder_conf) self.infer_cfg_ratio = kwargs.get("infer_cfg_ratio", None) logging.info(f"infer_cfg_ratio: {self.infer_cfg_ratio}") self.train_cfg_ratio = kwargs.get("train_cfg_ratio", None) logging.info(f"train_cfg_ratio: {self.train_cfg_ratio}") # 2. build audio token language model related modules self.sos_eos = 0 self.task_id = 1 self.llm_embedding = torch.nn.Embedding(2, llm_input_size) self.llm = llm self.llm_decoder = nn.Linear(llm_output_size, audio_token_size + 1) self.criterion_ce = LabelSmoothingLoss( size=audio_token_size + 1, padding_idx=IGNORE_ID, smoothing=lsm_weight, normalize_length=length_normalized_loss, ) # 3. [Optional] build audio token related modules self.speech_embedding = torch.nn.Embedding(audio_token_size, llm_input_size) self.spk_embed_affine_layer = torch.nn.Linear(192, llm_input_size) self.num_codebooks = 4 # 4. sampling method self.sampling = sampling self.time_embedding = SinusoidalEmbedding(llm_input_size) def cfg_dropout(self, text_token, text_token_len, p): # Classifier-Free Guidance Dropout B = text_token.size(0) num_samples_to_mask = int(p * B) if num_samples_to_mask == 0: num_samples_to_mask = 1 indices_to_mask = torch.randperm(B, device=text_token.device)[:num_samples_to_mask] text_token[indices_to_mask] = 0 text_token_len[indices_to_mask] = 0 return text_token, text_token_len def build_encoder(self, encoder_conf=None): if encoder_conf is None: assert hasattr(self, "encoder_conf"), \ "function param encoder_conf is None and model doesn't has encoder_conf attribute either." encoder_conf = self.encoder_conf encoder_name = encoder_conf.pop("name", "transformer") model = None if encoder_name == "transformer": from inspiremusic.transformer.encoder.conformer_encoder import ConformerEncoder model = ConformerEncoder( **encoder_conf, input_size=self.input_size, use_cnn_module=False, macaron_style=False, ) elif encoder_name == "conformer": from inspiremusic.transformer.encoder.conformer_encoder import ConformerEncoder model = ConformerEncoder( **encoder_conf, input_size=self.input_size, ) elif encoder_name == "llama_encoder": from inspiremusic.transformer.encoder.llama_encoder import LlamaEncoder model = LlamaEncoder( **encoder_conf, input_size=self.input_size, ) elif encoder_name == "qwen2": from inspiremusic.transformer.encoder.qwen_encoder import QwenEncoder model = QwenEncoder( **encoder_conf, input_size=self.input_size, ) elif encoder_name == "qwen2.5": from inspiremusic.transformer.encoder.qwen_encoder import QwenEncoder model = QwenEncoder( **encoder_conf, input_size=self.input_size, ) encoder_conf["name"] = encoder_name return model def encode(self, text: torch.Tensor, text_lengths: torch.Tensor): if self.text_encoder is not None: encoder_out, encoder_mask = self.text_encoder(text, text_lengths, decoding_chunk_size=1, num_decoding_left_chunks=-1) encoder_out_lens = encoder_mask.squeeze(1).sum(1) encoder_out = self.text_encoder_affine_layer(encoder_out) else: encoder_out, encoder_out_lens = text, text_lengths return encoder_out, encoder_out_lens def pad_unpad_sequence(self, sos_eos_emb, embeddings, text_token, text_token_len, task_id_emb, audio_token, audio_token_len, seg_len): text_token = unpad_sequence(text_token, text_token_len.cpu(), batch_first=True) audio_token = unpad_sequence(audio_token, audio_token_len.cpu(), batch_first=True) for i in range(len(embeddings)): embeddings[i] = unpad_sequence(embeddings[i], seg_len.cpu(), batch_first=True) lm_input = [torch.concat([sos_eos_emb.squeeze(dim=0)] + [embedding[i] for embedding in embeddings] + [text_token[i], task_id_emb.squeeze(dim=0), audio_token[i]], dim=0) for i in range(len(text_token))] lm_input_len = torch.tensor([i.size(0) for i in lm_input], dtype=torch.int32) lm_input = pad_sequence(lm_input, batch_first=True, padding_value=IGNORE_ID) return lm_input, lm_input_len def forward( self, batch: dict, device: torch.device, ) -> Dict[str, Optional[torch.Tensor]]: """ Args: text: (B, L, D) text_lengths: (B,) audio: (B, T, N) or (B, T) audio_lengths: (B,) """ mask = True text_token = batch['text_token'].to(device) text_token_len = batch['text_token_len'].to(device) if "semantic_token" not in batch: audio_token = batch['acoustic_token'].to(device) audio_token_len = batch['acoustic_token_len'].to(device) audio_token = audio_token.view(audio_token.size(0), -1, self.num_codebooks) audio_token = audio_token[:, :, 0] audio_token_len = (audio_token_len / self.num_codebooks).long() else: audio_token = batch['semantic_token'].to(device) audio_token_len = batch['semantic_token_len'].to(device) time_start = batch['time_start'].to(device) time_end = batch['time_end'].to(device) chorus = batch['chorus'].to(device) # 1. encode text_token if self.train_cfg_ratio > 0: # Classifier-Free Guidance text_token, _ = self.cfg_dropout(text_token, text_token_len, self.train_cfg_ratio) # 2. Time Embedding & chorus embedding text_token = self.text_embedding(text_token) text_token, text_token_len = self.encode(text_token, text_token_len) if mask: time_mask = time_start != -1.0 seg_len = time_mask.sum(-1) time_start = time_start.masked_fill(~time_mask, 0.0) time_end = time_end.masked_fill(~time_mask, 0.0) chorus = chorus.masked_fill(~time_mask, 0) time_start_embed = self.time_embedding(time_start.view(-1)).to(text_token.dtype) time_end_embed = self.time_embedding(time_end.view(-1)).to(text_token.dtype) time_start_embed = time_start_embed.view(chorus.size(0), chorus.size(1), -1) time_end_embed = time_end_embed.view(chorus.size(0), chorus.size(1), -1) chorus_embed = self.chorus_embedding(chorus) lm_target = [torch.tensor([IGNORE_ID] * (1 + 3 * seg_len[i] + text_token_len[i]) + audio_token[i,:audio_token_len[i]].tolist() + [self.audio_token_size]) for i in range(text_token.size(0))] else: time_start_embed = self.time_embedding(time_start).to(text_token.dtype) time_end_embed = self.time_embedding(time_end).to(text_token.dtype) chorus_embed = self.chorus_embedding(chorus) lm_target = [torch.tensor( [IGNORE_ID] * (4 + text_token_len[i]) + audio_token[i,:audio_token_len[i]].tolist() + [self.audio_token_size]) for i in range(text_token.size(0))] lm_target = pad_sequence(lm_target, batch_first=True, padding_value=IGNORE_ID).to(device) # 3. eos and task_id sos_eos_emb = self.llm_embedding.weight[self.sos_eos].reshape(1, 1, -1) task_id_emb = self.llm_embedding.weight[self.task_id].reshape(1, 1, -1) # 4. encode audio_token audio_token = self.speech_embedding(audio_token) # 5. unpad and pad lm_input, lm_input_len = self.pad_unpad_sequence(sos_eos_emb, [time_start_embed, time_end_embed, chorus_embed], text_token, text_token_len, task_id_emb, audio_token, audio_token_len, seg_len) # 6. run lm forward lm_output, lm_output_mask = self.llm(lm_input, lm_input_len.to(device)) logits = self.llm_decoder(lm_output) loss = self.criterion_ce(logits, lm_target) acc = th_accuracy(logits.view(-1, self.audio_token_size + 1), lm_target, ignore_label=IGNORE_ID) return {'loss': loss, 'acc': acc} def sampling_ids( self, weighted_scores: torch.Tensor, decoded_tokens: List, ignore_eos: bool = True, ): top_ids = self.sampling(weighted_scores, decoded_tokens) return top_ids @torch.inference_mode() def inference( self, text: torch.Tensor, text_len: torch.Tensor, audio_token: torch.Tensor, audio_token_len: torch.Tensor, prompt_text: torch.Tensor, prompt_text_len: torch.Tensor, prompt_audio_token: torch.Tensor, prompt_audio_token_len: torch.Tensor, embeddings: List, duration_to_gen: float = 300, task: str = "continuation", token_rate: int = 75, limit_audio_prompt_len: int = 5, ) -> Generator[torch.Tensor, None, None]: device = text.device if text is not None: text = torch.concat([prompt_text, text], dim=1) text_len += prompt_text_len infer_cfg = self.infer_cfg_ratio >= 0.0 if infer_cfg: text_cfg = self.text_embedding(text.new_zeros(text.shape)) text = self.text_embedding(text) # 1. encode text text, text_len = self.encode(text, text_len) # 2. encode embedding if embeddings is not None: time_start, time_end, chorus = embeddings if len(chorus.shape) == 1: time_start_embed = self.time_embedding(time_start).reshape(1, 1, -1) # .half() time_end_embed = self.time_embedding(time_end).reshape(1, 1, -1) # .half() chorus_embed = self.chorus_embedding(chorus).reshape(1, 1, -1) # .half() else: time_start_embed = self.time_embedding( time_start.view(-1)).reshape(1, chorus.size(1), -1) # .half() time_end_embed = self.time_embedding(time_end.view(-1)).reshape(1, chorus.size(1), -1) # .half() chorus_embed = self.chorus_embedding(chorus) # .half() # 3. concat llm_input sos_eos_emb = self.llm_embedding.weight[self.sos_eos].reshape(1, 1, -1) task_id_emb = self.llm_embedding.weight[self.task_id].reshape(1, 1, -1) if audio_token_len: audio_token = audio_token[:, :(limit_audio_prompt_len * token_rate)] audio_token_emb = self.speech_embedding(audio_token) else: audio_token_emb = torch.zeros(1, 0, self.llm_input_size, dtype=text.dtype).to(device) if prompt_audio_token_len: prompt_audio_token_emb = self.speech_embedding(prompt_audio_token) else: prompt_audio_token_emb = torch.zeros(1, 0, self.llm_input_size, dtype=text.dtype).to(device) # Check if removing prompt audio token will fail decoding. if task == "continuation": lm_input = torch.concat( [sos_eos_emb, time_start_embed, time_end_embed, chorus_embed, text, task_id_emb, audio_token_emb], dim=1) if infer_cfg: audio_cfg = self.speech_embedding( audio_token.new_zeros(audio_token.shape)) lm_cf_input = torch.concat( [sos_eos_emb, torch.rand_like(time_start_embed), torch.rand_like(time_end_embed), torch.rand_like(chorus_embed), text_cfg, task_id_emb, audio_cfg], dim=1) lm_input = torch.cat([lm_input, lm_cf_input], 0) else: lm_input = torch.concat( [sos_eos_emb, time_start_embed, time_end_embed, chorus_embed, text, task_id_emb], dim=1) if infer_cfg: lm_cf_input = torch.concat( [sos_eos_emb, torch.rand_like(time_start_embed), torch.rand_like(time_end_embed), torch.rand_like(chorus_embed), text_cfg, task_id_emb], dim=1) lm_input = torch.cat([lm_input, lm_cf_input], 0) # 4. cal min/max_length min_len = 0.9 * duration_to_gen * token_rate max_len = duration_to_gen * token_rate logging.info( f"LLM generation sequence length: {max_len}, generate audio length {duration_to_gen}s.") # 5. step by step decode out_tokens = [] offset = 0 state = None for i in range(int(max_len)): y_pred, _, state = self.llm.forward_one_step(lm_input, torch.ones(lm_input.shape[0], lm_input.shape[1], device=lm_input.device).to(torch.bool), cache=state) logits = self.llm_decoder(y_pred[:, -1]) if infer_cfg: # perform context free guidance logits_cf = logits[1] logits = logits[0] infer_cfg_ratio = self.infer_cfg_ratio logits = infer_cfg_ratio * logits + (1 - infer_cfg_ratio) * logits_cf logp = logits.log_softmax(dim=-1) logp = logp.squeeze(dim=0) if i < int(min_len): logp[self.audio_token_size] = torch.tensor(float('-inf'), dtype=torch.float16) if i < int(min_len): logp[self.audio_token_size] = torch.tensor(float('-inf'), dtype=torch.float16) top_ids = self.sampling_ids(logp, out_tokens, ignore_eos=i < min_len).item() if top_ids == self.audio_token_size: break # # in stream mode, yield token one by one yield torch.tensor([[top_ids]], dtype=torch.int64, device=device) out_tokens.append(top_ids) offset += lm_input.size(1) lm_input = self.speech_embedding.weight[top_ids].reshape(1, 1, -1) if infer_cfg: lm_input = lm_input.repeat(2, 1, 1)