# 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.
import numpy as np
import threading
import time
from contextlib import nullcontext
import uuid
from inspiremusic.music_tokenizer.vqvae import VQVAE
from inspiremusic.wavtokenizer.decoder.pretrained import WavTokenizer
from torch.cuda.amp import autocast
import logging
import torch
import os
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
class InspireMusicModel:
def __init__(self,
llm: torch.nn.Module,
flow: torch.nn.Module,
music_tokenizer: torch.nn.Module,
wavtokenizer: torch.nn.Module,
fast: bool = False,
fp16: bool = True,
):
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.llm = llm
self.flow = flow
self.music_tokenizer = music_tokenizer
self.wavtokenizer = wavtokenizer
self.fp16 = fp16
self.token_min_hop_len = 100
self.token_max_hop_len = 200
self.token_overlap_len = 20
# mel fade in out
self.mel_overlap_len = 34
self.mel_window = np.hamming(2 * self.mel_overlap_len)
# hift cache
self.mel_cache_len = 20
self.source_cache_len = int(self.mel_cache_len * 256)
# rtf and decoding related
self.stream_scale_factor = 1
assert self.stream_scale_factor >= 1, 'stream_scale_factor should be greater than 1, change it according to your actual rtf'
self.llm_context = torch.cuda.stream(torch.cuda.Stream(self.device)) if torch.cuda.is_available() else nullcontext()
self.lock = threading.Lock()
# dict used to store session related variable
self.music_token_dict = {}
self.llm_end_dict = {}
self.mel_overlap_dict = {}
self.fast = fast
self.generator = "hifi"
def load(self, llm_model, flow_model, hift_model, wavtokenizer_model):
if llm_model is not None:
self.llm.load_state_dict(torch.load(llm_model, map_location=self.device))
self.llm.to(self.device).eval()
else:
self.llm = None
if flow_model is not None:
self.flow.load_state_dict(torch.load(flow_model, map_location=self.device))
self.flow.to(self.device).eval()
if hift_model is not None:
if ".pt" not in hift_model:
self.music_tokenizer = VQVAE( hift_model + '/config.json',
hift_model + '/model.pt', with_encoder=True)
else:
self.music_tokenizer = VQVAE(os.path.dirname(hift_model) + '/config.json',
hift_model, with_encoder=True)
self.music_tokenizer.to(self.device).eval()
if wavtokenizer_model is not None:
if ".pt" not in wavtokenizer_model:
self.wavtokenizer = WavTokenizer.from_pretrained_feat( wavtokenizer_model + '/config.yaml',
wavtokenizer_model + '/model.pt')
else:
self.wavtokenizer = WavTokenizer.from_pretrained_feat( os.path.dirname(wavtokenizer_model) + '/config.yaml',
wavtokenizer_model )
self.wavtokenizer.to(self.device)
def load_jit(self, llm_text_encoder_model, llm_llm_model, flow_encoder_model):
assert self.fp16 is True, "we only provide fp16 jit model, set fp16=True if you want to use jit model"
llm_text_encoder = torch.jit.load(llm_text_encoder_model, map_location=self.device)
self.llm.text_encoder = llm_text_encoder
llm_llm = torch.jit.load(llm_llm_model)
self.llm.llm = llm_llm
flow_encoder = torch.jit.load(flow_encoder_model)
self.flow.encoder = flow_encoder
def load_onnx(self, flow_decoder_estimator_model):
import onnxruntime
option = onnxruntime.SessionOptions()
option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
option.intra_op_num_threads = 1
providers = ['CUDAExecutionProvider' if torch.cuda.is_available() else 'CPUExecutionProvider']
del self.flow.decoder.estimator
self.flow.decoder.estimator = onnxruntime.InferenceSession(flow_decoder_estimator_model, sess_options=option, providers=providers)
def llm_job(self, text, audio_token, audio_token_len, prompt_text, llm_prompt_audio_token, embeddings, uuid, duration_to_gen, task):
with self.llm_context:
local_res = []
with autocast(enabled=self.fp16):
inference_kwargs = {
'text': text.to(self.device),
'text_len': torch.tensor([text.shape[1]], dtype=torch.int32).to(self.device),
'prompt_text': prompt_text.to(self.device),
'prompt_text_len': torch.tensor([prompt_text.shape[1]], dtype=torch.int32).to(self.device),
'prompt_audio_token': llm_prompt_audio_token.to(self.device),
'prompt_audio_token_len': torch.tensor([llm_prompt_audio_token.shape[1]], dtype=torch.int32).to(self.device),
'embeddings': embeddings,
'duration_to_gen': duration_to_gen,
'task': task
}
if audio_token is not None:
inference_kwargs['audio_token'] = audio_token.to(self.device)
else:
inference_kwargs['audio_token'] = torch.Tensor([0]).to(self.device)
if audio_token_len is not None:
inference_kwargs['audio_token_len'] = audio_token_len.to(self.device)
else:
inference_kwargs['audio_token_len'] = torch.Tensor([0]).to(self.device)
for i in self.llm.inference(**inference_kwargs):
local_res.append(i)
self.music_token_dict[uuid] = local_res
self.llm_end_dict[uuid] = True
# def token2wav(self, token, token_len, text, text_len, uuid, sample_rate, finalize=False):
def token2wav(self, token, token_len, uuid, sample_rate, finalize=False, flow_cfg=None):
# if self.flow is not None:
# if isinstance(self.flow,MaskedDiffWithText):
# codec_embed = self.flow.inference(token=token.to(self.device),
# token_len=token_len.to(self.device),
# text_token=text,
# text_token_len=text_len,
# )
# else:
if flow_cfg is not None:
codec_embed = self.flow.inference_cfg(token=token.to(self.device),
token_len=token_len.to(self.device),
sample_rate=sample_rate
)
else:
codec_embed = self.flow.inference(token=token.to(self.device),
token_len=token_len.to(self.device),
sample_rate=sample_rate
)
# use music_tokenizer decoder
wav = self.music_tokenizer.generator(codec_embed)
wav = wav.squeeze(0).cpu().detach()
return wav
def acoustictoken2wav(self, token):
# use music_tokenizer to generate waveform from token
token = token.view(token.size(0), -1, 4)
# codec = token.view(1, -1, 4)
codec_embed = self.music_tokenizer.quantizer.embed(torch.tensor(token).long().to(self.device)).cuda()
wav = self.music_tokenizer.generator(codec_embed)
wav = wav.squeeze(0).cpu().detach()
return wav
def semantictoken2wav(self, token):
# fast mode, use wavtokenizer decoder
new_tensor = torch.tensor(token.to(self.device)).unsqueeze(0)
features = self.wavtokenizer.codes_to_features(new_tensor)
bandwidth_id = torch.tensor([0]).to(self.device)
wav = self.wavtokenizer.to(self.device).decode(features, bandwidth_id=bandwidth_id)
wav = wav.cpu().detach()
return wav
@torch.inference_mode()
def inference(self, text, audio_token, audio_token_len, text_token, text_token_len, embeddings=None,
prompt_text=torch.zeros(1, 0, dtype=torch.int32),
llm_prompt_audio_token=torch.zeros(1, 0, dtype=torch.int32),
flow_prompt_audio_token=torch.zeros(1, 0, dtype=torch.int32),
prompt_audio_feat=torch.zeros(1, 0, 80), sample_rate=48000, duration_to_gen = 30, task="continuation", trim = True, stream=False, **kwargs):
# this_uuid is used to track variables related to this inference thread
# support tasks:
# text to music task
# music continuation task
# require either audio input only or text and audio inputs
this_uuid = str(uuid.uuid1())
if self.llm:
with self.lock:
self.music_token_dict[this_uuid], self.llm_end_dict[this_uuid] = [], False
p = threading.Thread(target=self.llm_job, args=(text_token, audio_token, audio_token_len, prompt_text, llm_prompt_audio_token, embeddings, this_uuid, duration_to_gen, task))
p.start()
if stream is True:
token_hop_len = self.token_min_hop_len
while True:
time.sleep(0.1)
if len(self.music_token_dict[this_uuid]) >= token_hop_len + self.token_overlap_len:
this_music_audio = self.token2wav(token=text_token,
token_len=text_token_len,
uuid=this_uuid,
sample_rate=sample_rate,
finalize=False)
yield {'music_audio': this_music_audio.cpu()}
with self.lock:
self.music_token_dict[this_uuid] = self.music_token_dict[this_uuid][token_hop_len:]
# increase token_hop_len for better audio quality
token_hop_len = min(self.token_max_hop_len, int(token_hop_len * self.stream_scale_factor))
if self.llm_end_dict[this_uuid] is True and len(self.music_token_dict[this_uuid]) < token_hop_len + self.token_overlap_len:
break
p.join()
# deal with remain tokens, make sure inference remain token len equals token_hop_len when cache_speech is not None
this_music_token = torch.concat(self.music_token_dict[this_uuid], dim=1)
with self.flow_hift_context:
this_music_audio = self.token2wav(token=this_music_token,
prompt_token=flow_prompt_audio_token,
prompt_feat=prompt_audio_feat,
embedding=flow_embedding,
uuid=this_uuid,
sample_rate=sample_rate,
finalize=True)
yield {'music_audio': this_music_audio.cpu()}
else:
# deal with all tokens
if self.fast:
if task == "reconstruct":
assert audio_token is None
this_music_token = audio_token
this_music_audio = self.acoustictoken2wav(token=this_music_token)
else:
if self.llm:
p.join()
print(len(self.music_token_dict[this_uuid]))
this_music_token = torch.concat(self.music_token_dict[this_uuid], dim=1)
print(this_music_token.shape)
else:
this_music_token = text_token
logging.info("using wavtokenizer generator without flow matching")
this_music_audio = self.semantictoken2wav(token=this_music_token)
print(this_music_audio.shape)
else:
if self.llm:
p.join()
if len(self.music_token_dict[this_uuid]) != 0:
this_music_token = torch.concat(self.music_token_dict[this_uuid], dim=1)
else:
print(f"The list of tensors is empty for UUID: {this_uuid}")
else:
this_music_token = text_token
logging.info(f"LLM generated audio token length: {this_music_token.shape[1]}")
logging.info(f"using flow matching and {self.generator} generator")
if self.generator == "hifi":
if (embeddings[1] - embeddings[0]) <= duration_to_gen:
if trim:
trim_length = (int((embeddings[1] - embeddings[0])*75))
this_music_token = this_music_token[:, :trim_length]
logging.info(f"After trimmed, generated audio token length: {this_music_token.shape[1]}")
elif (embeddings[1] - embeddings[0]) < 1:
logging.info(f"Given audio length={(embeddings[1] - embeddings[0])}, which is too short, please give a longer audio length.")
this_music_audio = self.token2wav(token=this_music_token,
token_len=torch.LongTensor([this_music_token.size(1)]),
uuid=this_uuid,
sample_rate=sample_rate,
finalize=True)
logging.info(f"Generated audio sequence length: {this_music_audio.shape[1]}")
elif self.generator == "wavtokenizer":
if (embeddings[1] - embeddings[0]) < duration_to_gen:
if trim:
trim_length = (int((embeddings[1] - embeddings[0])*75))
this_music_token = this_music_token[:,:trim_length]
logging.info(f"After trimmed, generated audio token length: {this_music_token.shape[1]}")
elif (embeddings[1] - embeddings[0]) < 1:
logging.info(f"Given audio length={(embeddings[1] - embeddings[0])}, which is too short, please give a longer audio length.")
this_music_audio = self.semantictoken2wav(token=this_music_token)
yield {'music_audio': this_music_audio.cpu()}
torch.cuda.synchronize()