inspiremusic/bin/inference.py

# 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 __future__ import print_function

import argparse
import logging

logging.getLogger('matplotlib').setLevel(logging.WARNING)
import os
import torch
from torch.utils.data import DataLoader
import torchaudio
from hyperpyyaml import load_hyperpyyaml
from tqdm import tqdm
from inspiremusic.cli.model import InspireMusicModel
from inspiremusic.dataset.dataset import Dataset
import time
from inspiremusic.utils.audio_utils import trim_audio, fade_out, process_audio
from inspiremusic.utils.common import MUSIC_STRUCTURE_LABELS

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def get_args():
    parser = argparse.ArgumentParser(description='inference only with your model')
    parser.add_argument('--config', required=True, help='config file')
    parser.add_argument('--prompt_data', required=True, help='prompt data file')
    parser.add_argument('--flow_model', default=None, required=False, help='flow model file')
    parser.add_argument('--llm_model', default=None,required=False, help='flow model file')
    parser.add_argument('--music_tokenizer', required=True, help='music tokenizer model file')
    parser.add_argument('--wavtokenizer', required=True, help='wavtokenizer model file')
    parser.add_argument('--chorus', default="random",required=False, help='chorus tag generation mode, eg. random, verse, chorus, intro.')
    parser.add_argument('--fast', action='store_true', required=False, help='True: fast inference mode, without flow matching for fast inference. False: normal inference mode, with flow matching for high quality.')
    parser.add_argument('--fp16', default=True, type=bool, required=False, help='inference with fp16 model')
    parser.add_argument('--fade_out', default=True, type=bool, required=False, help='add fade out effect to generated audio')
    parser.add_argument('--fade_out_duration', default=1.0, type=float, required=False, help='fade out duration in seconds')
    parser.add_argument('--trim', default=False, type=bool, required=False, help='trim the silence ending of generated audio')
    parser.add_argument('--format', type=str, default="wav", required=False,
                        choices=["wav", "mp3", "m4a", "flac"],
                        help='sampling rate of input audio')
    parser.add_argument('--sample_rate', type=int, default=24000, required=False,
                        help='sampling rate of input audio')
    parser.add_argument('--output_sample_rate', type=int, default=48000, required=False, choices=[24000, 48000],
                        help='sampling rate of generated output audio')
    parser.add_argument('--min_generate_audio_seconds', type=float, default=10.0, required=False,
                        help='the minimum generated audio length in seconds')
    parser.add_argument('--max_generate_audio_seconds', type=float, default=30.0, required=False,
                        help='the maximum generated audio length in seconds')
    parser.add_argument('--gpu',
                        type=int,
                        default=0,
                        help='gpu id for this rank, -1 for cpu')
    parser.add_argument('--task',
                        default='text-to-music',
                        choices=['text-to-music', 'continuation', "reconstruct", "super_resolution"],
                        help='choose inference task type. text-to-music: text-to-music task. continuation: music continuation task. reconstruct: reconstruction of original music. super_resolution: convert original 24kHz music into 48kHz music.')
    parser.add_argument('--result_dir', required=True, help='asr result file')
    args = parser.parse_args()
    print(args)
    return args


def main():
	args = get_args()
	logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(message)s')
	os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)

	if args.fast:
		args.output_sample_rate = 24000

	min_generate_audio_length = int(args.output_sample_rate * args.min_generate_audio_seconds)
	max_generate_audio_length = int(args.output_sample_rate * args.max_generate_audio_seconds)
	assert args.min_generate_audio_seconds <= args.max_generate_audio_seconds

	# Init inspiremusic models from configs
	use_cuda = args.gpu >= 0 and torch.cuda.is_available()
	device = torch.device('cuda' if use_cuda else 'cpu')
	with open(args.config, 'r') as f:
		configs = load_hyperpyyaml(f)

	model = InspireMusicModel(configs['llm'], configs['flow'], configs['hift'], configs['wavtokenizer'], args.fast, args.fp16)

	model.load(args.llm_model, args.flow_model, args.music_tokenizer, args.wavtokenizer)

	if args.llm_model is None:
		model.llm = None
	else:
		model.llm = model.llm.to(torch.float32)

	if args.flow_model is None:
		model.flow = None

	test_dataset = Dataset(args.prompt_data, data_pipeline=configs['data_pipeline'], mode='inference', shuffle=True, partition=False)
	test_data_loader = DataLoader(test_dataset, batch_size=None, num_workers=0)

	del configs
	os.makedirs(args.result_dir, exist_ok=True)
	fn = os.path.join(args.result_dir, 'wav.scp')
	f = open(fn, 'w')
	caption_fn = os.path.join(args.result_dir, 'captions.txt')
	caption_f = open(caption_fn, 'w')

	with torch.no_grad():
		for _, batch in tqdm(enumerate(test_data_loader)):
			utts = batch["utts"]

			assert len(utts) == 1, "inference mode only support batchsize 1"
			text_token = batch["text_token"].to(device)
			text_token_len = batch["text_token_len"].to(device)

			if "time_start" not in batch.keys():
				batch["time_start"] = torch.randint(0, args.min_generate_audio_seconds, (1,)).to(torch.float64)

			if batch["time_start"].numpy()[0] > 300:
				batch["time_start"] = torch.Tensor([0]).to(torch.float64)

			if "time_end" not in batch.keys():
				batch["time_end"] = torch.randint(int(batch["time_start"].numpy()[0] + args.min_generate_audio_seconds), int(batch["time_start"].numpy()[0] + args.max_generate_audio_seconds), (1,)).to(torch.float64)
			else:
				if (batch["time_end"].numpy()[0] - batch["time_start"].numpy()[0]) < args.min_generate_audio_seconds:
					batch["time_end"] = torch.randint(int(batch["time_start"].numpy()[0] + args.min_generate_audio_seconds), int(batch["time_start"].numpy()[0] + args.max_generate_audio_seconds), (1,)).to(torch.float64)
				elif (batch["time_end"].numpy()[0] - batch["time_start"].numpy()[0]) > args.max_generate_audio_seconds:
					batch["time_end"] = torch.Tensor([(batch["time_start"].numpy()[0] + args.max_generate_audio_seconds)]).to(torch.float64)

			if "chorus" not in batch.keys():
				batch["chorus"] = torch.randint(1, 5, (1,))

			if args.chorus == "random":
				batch["chorus"] = torch.randint(1, 5, (1,))
			elif args.chorus == "intro":
				batch["chorus"] = torch.Tensor([0])
			elif "verse" in args.chorus:
				batch["chorus"] = torch.Tensor([1])
			elif args.chorus == "chorus":
				batch["chorus"] = torch.Tensor([2])
			elif args.chorus == "outro":
				batch["chorus"] = torch.Tensor([4])
			else:
				batch["chorus"] = batch["chorus"]

			time_start = batch["time_start"].to(device)
			time_end = batch["time_end"].to(device)
			chorus = batch["chorus"].to(torch.int)

			text_prompt = f"<|{batch['time_start'].numpy()[0]}|><|{MUSIC_STRUCTURE_LABELS[chorus.numpy()[0]]}|><|{batch['text'][0]}|><|{batch['time_end'].numpy()[0]}|>"
			chorus = chorus.to(device)

			if batch["acoustic_token"] is None:
				audio_token = None
				audio_token_len = None
			else:
				audio_token = batch["acoustic_token"].to(device)
				audio_token_len = batch["acoustic_token_len"].to(device)

			text = batch["text"]

			if "semantic_token" in batch:
				token = batch["semantic_token"].to(device)
				token_len = batch["semantic_token_len"].to(device)
			else:
				if audio_token is None:
					token = None
					token_len = None
				else:
					token = audio_token.view(audio_token.size(0), -1, 4)[:, :, 0]
					token_len = audio_token_len / 4

			if args.task in ['text-to-music', 'continuation']:
				# text to music, music continuation
				model_input = {"text": text, "audio_token": token,
							   "audio_token_len": token_len,
							   "text_token": text_token,
							   "text_token_len": text_token_len,
							   "embeddings": [time_start, time_end, chorus],
							   "raw_text": text,
							   "sample_rate": args.output_sample_rate,
							   "duration_to_gen": args.max_generate_audio_seconds,
							   "task": args.task}
			elif args.task in ['reconstruct', 'super_resolution']:
				# audio reconstruction, audio super resolution
				model_input = {"text": text, "audio_token": audio_token,
							   "audio_token_len": audio_token_len,
							   "text_token": text_token,
							   "text_token_len": text_token_len,
							   "embeddings": [time_start, time_end, chorus],
							   "raw_text": text,
							   "sample_rate": args.output_sample_rate,
							   "duration_to_gen": args.max_generate_audio_seconds,
							   "task": args.task}
			else:
				# zero-shot
				model_input = {'text'                       : text,
							   'text_len'                   : text_token_len,
							   'prompt_text'                : text_token,
							   'prompt_text_len'            : text_token_len,
							   'llm_prompt_audio_token'     : token,
							   'llm_prompt_audio_token_len' : token_len,
							   'flow_prompt_audio_token'    : audio_token,
							   'flow_prompt_audio_token_len': audio_token_len,
							   'prompt_audio_feat'          : audio_feat,
							   'prompt_audio_feat_len'      : audio_feat_len,
							   "embeddings"                 : [time_start,
															   time_end,
															   chorus]}

			music_key = utts[0]
			music_audios = []
			music_fn = os.path.join(args.result_dir, f'{music_key}.{args.format}')
			bench_start = time.time()

			for model_output in model.inference(**model_input):
				music_audios.append(model_output['music_audio'])
			bench_end = time.time()
			if args.trim:
				music_audio = trim_audio(music_audios[0],
										 sample_rate=args.output_sample_rate,
										 threshold=0.05,
										 min_silence_duration=0.8)
			else:
				music_audio = music_audios[0]
			if music_audio.shape[0] != 0:
				if music_audio.shape[1] > max_generate_audio_length:
					music_audio = music_audio[:, :max_generate_audio_length]
				if music_audio.shape[1] >= min_generate_audio_length:
					try:
						if args.fade_out:
							music_audio = fade_out(music_audio, args.output_sample_rate, args.fade_out_duration)
						music_audio = music_audio.repeat(2, 1)
						if args.format in ["wav", "flac"]:
							torchaudio.save(music_fn, music_audio, sample_rate=args.output_sample_rate, encoding="PCM_S", bits_per_sample=24)
						elif args.format in ["mp3", "m4a"]:
							torchaudio.backend.sox_io_backend.save(filepath=music_fn, src=music_audio, sample_rate=args.output_sample_rate, format=args.format)
						else:
							logging.info(f"Format is not supported. Please choose from wav, mp3, m4a, flac.")
					except Exception as e:
						logging.info(f"Error saving file: {e}")
						raise

					audio_duration = music_audio.shape[1] / args.output_sample_rate
					rtf = (bench_end - bench_start) / audio_duration
					logging.info(f"processing time: {int(bench_end - bench_start)}s, audio length: {int(audio_duration)}s, rtf: {rtf}, text prompt: {text_prompt}")
					f.write('{} {}\n'.format(music_key, music_fn))
					f.flush()
					caption_f.write('{}\t{}\n'.format(music_key, text_prompt))
					caption_f.flush()
				else:
					logging.info(f"Generate audio length {music_audio.shape[1]} is shorter than min_generate_audio_length.")
			else:
				logging.info(f"Generate audio is empty, dim = {music_audio.shape[0]}.")
	f.close()
	logging.info('Result wav.scp saved in {}'.format(fn))


if __name__ == '__main__':
	main()