Delete infer.py

infer.py

@@ -1,456 +0,0 @@
-import os
-import sys
-sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
-sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
-import argparse
-import torch
-import numpy as np
-import json
-from omegaconf import OmegaConf
-import torchaudio
-from torchaudio.transforms import Resample
-import soundfile as sf
-
-import uuid
-from tqdm import tqdm
-from einops import rearrange
-from codecmanipulator import CodecManipulator
-from mmtokenizer import _MMSentencePieceTokenizer
-from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
-import glob
-import time
-import copy
-from collections import Counter
-from models.soundstream_hubert_new import SoundStream
-from vocoder import build_codec_model, process_audio
-from post_process_audio import replace_low_freq_with_energy_matched
-import re
-
-
-parser = argparse.ArgumentParser()
-# Model Configuration:
-parser.add_argument("--stage1_model", type=str, default="m-a-p/YuE-s1-7B-anneal-en-cot", help="The model checkpoint path or identifier for the Stage 1 model.")
-parser.add_argument("--stage2_model", type=str, default="m-a-p/YuE-s2-1B-general", help="The model checkpoint path or identifier for the Stage 2 model.")
-parser.add_argument("--max_new_tokens", type=int, default=3000, help="The maximum number of new tokens to generate in one pass during text generation.")
-parser.add_argument("--run_n_segments", type=int, default=2, help="The number of segments to process during the generation.")
-parser.add_argument("--stage2_batch_size", type=int, default=4, help="The batch size used in Stage 2 inference.")
-# Prompt
-parser.add_argument("--genre_txt", type=str, required=True, help="The file path to a text file containing genre tags that describe the musical style or characteristics (e.g., instrumental, genre, mood, vocal timbre, vocal gender). This is used as part of the generation prompt.")
-parser.add_argument("--lyrics_txt", type=str, required=True, help="The file path to a text file containing the lyrics for the music generation. These lyrics will be processed and split into structured segments to guide the generation process.")
-parser.add_argument("--use_audio_prompt", action="store_true", help="If set, the model will use an audio file as a prompt during generation. The audio file should be specified using --audio_prompt_path.")
-parser.add_argument("--audio_prompt_path", type=str, default="", help="The file path to an audio file to use as a reference prompt when --use_audio_prompt is enabled.")
-parser.add_argument("--prompt_start_time", type=float, default=0.0, help="The start time in seconds to extract the audio prompt from the given audio file.")
-parser.add_argument("--prompt_end_time", type=float, default=30.0, help="The end time in seconds to extract the audio prompt from the given audio file.")
-# Output 
-parser.add_argument("--output_dir", type=str, default="./output", help="The directory where generated outputs will be saved.")
-parser.add_argument("--keep_intermediate", action="store_true", help="If set, intermediate outputs will be saved during processing.")
-parser.add_argument("--disable_offload_model", action="store_true", help="If set, the model will not be offloaded from the GPU to CPU after Stage 1 inference.")
-parser.add_argument("--cuda_idx", type=int, default=0)
-# Config for xcodec and upsampler
-parser.add_argument('--basic_model_config', default='./xcodec_mini_infer/final_ckpt/config.yaml', help='YAML files for xcodec configurations.')
-parser.add_argument('--resume_path', default='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', help='Path to the xcodec checkpoint.')
-parser.add_argument('--config_path', type=str, default='./xcodec_mini_infer/decoders/config.yaml', help='Path to Vocos config file.')
-parser.add_argument('--vocal_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_131000.pth', help='Path to Vocos decoder weights.')
-parser.add_argument('--inst_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_151000.pth', help='Path to Vocos decoder weights.')
-parser.add_argument('-r', '--rescale', action='store_true', help='Rescale output to avoid clipping.')
-
-
-args = parser.parse_args()
-if args.use_audio_prompt and not args.audio_prompt_path:
-    raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
-stage1_model = args.stage1_model
-stage2_model = args.stage2_model
-cuda_idx = args.cuda_idx
-max_new_tokens = args.max_new_tokens
-stage1_output_dir = os.path.join(args.output_dir, f"stage1")
-stage2_output_dir = stage1_output_dir.replace('stage1', 'stage2')
-os.makedirs(stage1_output_dir, exist_ok=True)
-os.makedirs(stage2_output_dir, exist_ok=True)
-
-# load tokenizer and model
-device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
-mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-model = AutoModelForCausalLM.from_pretrained(
-    stage1_model, 
-    torch_dtype=torch.bfloat16,
-    attn_implementation="flash_attention_2", # To enable flashattn, you have to install flash-attn
-    )
-# to device, if gpu is available
-model.to(device)
-model.eval()
-
-codectool = CodecManipulator("xcodec", 0, 1)
-codectool_stage2 = CodecManipulator("xcodec", 0, 8)
-model_config = OmegaConf.load(args.basic_model_config)
-codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
-parameter_dict = torch.load(args.resume_path, map_location='cpu')
-codec_model.load_state_dict(parameter_dict['codec_model'])
-codec_model.to(device)
-codec_model.eval()
-
-class BlockTokenRangeProcessor(LogitsProcessor):
-    def __init__(self, start_id, end_id):
-        self.blocked_token_ids = list(range(start_id, end_id))
-
-    def __call__(self, input_ids, scores):
-        scores[:, self.blocked_token_ids] = -float("inf")
-        return scores
-
-def load_audio_mono(filepath, sampling_rate=16000):
-    audio, sr = torchaudio.load(filepath)
-    # Convert to mono
-    audio = torch.mean(audio, dim=0, keepdim=True)
-    # Resample if needed
-    if sr != sampling_rate:
-        resampler = Resample(orig_freq=sr, new_freq=sampling_rate)
-        audio = resampler(audio)
-    return audio
-
-def split_lyrics(lyrics):
-    pattern = r"\[(\w+)\](.*?)\n(?=\[|\Z)"
-    segments = re.findall(pattern, lyrics, re.DOTALL)
-    structured_lyrics = [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
-    return structured_lyrics
-
-# Call the function and print the result
-stage1_output_set = []
-# Tips:
-# genre tags support instrumental，genre，mood，vocal timbr and vocal gender
-# all kinds of tags are needed
-with open(args.genre_txt) as f:
-    genres = f.read().strip()
-with open(args.lyrics_txt) as f:
-    lyrics = split_lyrics(f.read())
-# intruction
-full_lyrics = "\n".join(lyrics)
-prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"]
-prompt_texts += lyrics
-
-
-random_id = uuid.uuid4()
-output_seq = None
-# Here is suggested decoding config
-top_p = 0.93
-temperature = 1.0
-repetition_penalty = 1.2
-# special tokens
-start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
-end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
-# Format text prompt
-run_n_segments = min(args.run_n_segments+1, len(lyrics))
-for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
-    section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
-    guidance_scale = 1.5 if i <=1 else 1.2
-    if i==0:
-        continue
-    if i==1:
-        if args.use_audio_prompt:
-            audio_prompt = load_audio_mono(args.audio_prompt_path)
-            audio_prompt.unsqueeze_(0)
-            with torch.no_grad():
-                raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5)
-            raw_codes = raw_codes.transpose(0, 1)
-            raw_codes = raw_codes.cpu().numpy().astype(np.int16)
-            # Format audio prompt
-            code_ids = codectool.npy2ids(raw_codes[0])
-            audio_prompt_codec = code_ids[int(args.prompt_start_time *50): int(args.prompt_end_time *50)] # 50 is tps of xcodec
-            audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
-            sentence_ids = mmtokenizer.tokenize("[start_of_reference]") +  audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
-            head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
-        else:
-            head_id = mmtokenizer.tokenize(prompt_texts[0])
-        prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
-    else:
-        prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
-
-    prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device) 
-    input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids
-    # Use window slicing in case output sequence exceeds the context of model
-    max_context = 16384-max_new_tokens-1
-    if input_ids.shape[-1] > max_context:
-        print(f'Section {i}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
-        input_ids = input_ids[:, -(max_context):]
-    with torch.no_grad():
-        output_seq = model.generate(
-            input_ids=input_ids, 
-            max_new_tokens=max_new_tokens, 
-            min_new_tokens=100, 
-            do_sample=True, 
-            top_p=top_p,
-            temperature=temperature, 
-            repetition_penalty=repetition_penalty, 
-            eos_token_id=mmtokenizer.eoa,
-            pad_token_id=mmtokenizer.eoa,
-            logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
-            guidance_scale=guidance_scale,
-            )
-        if output_seq[0][-1].item() != mmtokenizer.eoa:
-            tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
-            output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
-    if i > 1:
-        raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
-    else:
-        raw_output = output_seq
-
-# save raw output and check sanity
-ids = raw_output[0].cpu().numpy()
-soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
-eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
-if len(soa_idx)!=len(eoa_idx):
-    raise ValueError(f'invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')
-
-vocals = []
-instrumentals = []
-range_begin = 1 if args.use_audio_prompt else 0
-for i in range(range_begin, len(soa_idx)):
-    codec_ids = ids[soa_idx[i]+1:eoa_idx[i]]
-    if codec_ids[0] == 32016:
-        codec_ids = codec_ids[1:]
-    codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
-    vocals_ids = codectool.ids2npy(rearrange(codec_ids,"(n b) -> b n", b=2)[0])
-    vocals.append(vocals_ids)
-    instrumentals_ids = codectool.ids2npy(rearrange(codec_ids,"(n b) -> b n", b=2)[1])
-    instrumentals.append(instrumentals_ids)
-vocals = np.concatenate(vocals, axis=1)
-instrumentals = np.concatenate(instrumentals, axis=1)
-vocal_save_path = os.path.join(stage1_output_dir, f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_vocal_{random_id}".replace('.', '@')+'.npy')
-inst_save_path = os.path.join(stage1_output_dir, f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_instrumental_{random_id}".replace('.', '@')+'.npy')
-np.save(vocal_save_path, vocals)
-np.save(inst_save_path, instrumentals)
-stage1_output_set.append(vocal_save_path)
-stage1_output_set.append(inst_save_path)
-
-
-# offload model
-if not args.disable_offload_model:
-    model.cpu()
-    del model
-    torch.cuda.empty_cache()
-
-print("Stage 2 inference...")
-model_stage2 = AutoModelForCausalLM.from_pretrained(
-    stage2_model, 
-    torch_dtype=torch.float16,
-    attn_implementation="flash_attention_2"
-    )
-model_stage2.to(device)
-model_stage2.eval()
-
-def stage2_generate(model, prompt, batch_size=16):
-    codec_ids = codectool.unflatten(prompt, n_quantizer=1)
-    codec_ids = codectool.offset_tok_ids(
-                    codec_ids, 
-                    global_offset=codectool.global_offset, 
-                    codebook_size=codectool.codebook_size, 
-                    num_codebooks=codectool.num_codebooks, 
-                ).astype(np.int32)
-    
-    # Prepare prompt_ids based on batch size or single input
-    if batch_size > 1:
-        codec_list = []
-        for i in range(batch_size):
-            idx_begin = i * 300
-            idx_end = (i + 1) * 300
-            codec_list.append(codec_ids[:, idx_begin:idx_end])
-
-        codec_ids = np.concatenate(codec_list, axis=0)
-        prompt_ids = np.concatenate(
-            [
-                np.tile([mmtokenizer.soa, mmtokenizer.stage_1], (batch_size, 1)),
-                codec_ids,
-                np.tile([mmtokenizer.stage_2], (batch_size, 1)),
-            ],
-            axis=1
-        )
-    else:
-        prompt_ids = np.concatenate([
-            np.array([mmtokenizer.soa, mmtokenizer.stage_1]),
-            codec_ids.flatten(),  # Flatten the 2D array to 1D
-            np.array([mmtokenizer.stage_2])
-        ]).astype(np.int32)
-        prompt_ids = prompt_ids[np.newaxis, ...]
-
-    codec_ids = torch.as_tensor(codec_ids).to(device)
-    prompt_ids = torch.as_tensor(prompt_ids).to(device)
-    len_prompt = prompt_ids.shape[-1]
-    
-    block_list = LogitsProcessorList([BlockTokenRangeProcessor(0, 46358), BlockTokenRangeProcessor(53526, mmtokenizer.vocab_size)])
-
-    # Teacher forcing generate loop
-    for frames_idx in range(codec_ids.shape[1]):
-        cb0 = codec_ids[:, frames_idx:frames_idx+1]
-        prompt_ids = torch.cat([prompt_ids, cb0], dim=1)
-        input_ids = prompt_ids
-
-        with torch.no_grad():
-            stage2_output = model.generate(input_ids=input_ids, 
-                min_new_tokens=7,
-                max_new_tokens=7,
-                eos_token_id=mmtokenizer.eoa,
-                pad_token_id=mmtokenizer.eoa,
-                logits_processor=block_list,
-            )
-        
-        assert stage2_output.shape[1] - prompt_ids.shape[1] == 7, f"output new tokens={stage2_output.shape[1]-prompt_ids.shape[1]}"
-        prompt_ids = stage2_output
-
-    # Return output based on batch size
-    if batch_size > 1:
-        output = prompt_ids.cpu().numpy()[:, len_prompt:]
-        output_list = [output[i] for i in range(batch_size)]
-        output = np.concatenate(output_list, axis=0)
-    else:
-        output = prompt_ids[0].cpu().numpy()[len_prompt:]
-
-    return output
-
-def stage2_inference(model, stage1_output_set, stage2_output_dir, batch_size=4):
-    stage2_result = []
-    for i in tqdm(range(len(stage1_output_set))):
-        output_filename = os.path.join(stage2_output_dir, os.path.basename(stage1_output_set[i]))
-        
-        if os.path.exists(output_filename):
-            print(f'{output_filename} stage2 has done.')
-            continue
-        
-        # Load the prompt
-        prompt = np.load(stage1_output_set[i]).astype(np.int32)
-        
-        # Only accept 6s segments
-        output_duration = prompt.shape[-1] // 50 // 6 * 6
-        num_batch = output_duration // 6
-        
-        if num_batch <= batch_size:
-            # If num_batch is less than or equal to batch_size, we can infer the entire prompt at once
-            output = stage2_generate(model, prompt[:, :output_duration*50], batch_size=num_batch)
-        else:
-            # If num_batch is greater than batch_size, process in chunks of batch_size
-            segments = []
-            num_segments = (num_batch // batch_size) + (1 if num_batch % batch_size != 0 else 0)
-
-            for seg in range(num_segments):
-                start_idx = seg * batch_size * 300
-                # Ensure the end_idx does not exceed the available length
-                end_idx = min((seg + 1) * batch_size * 300, output_duration*50)  # Adjust the last segment
-                current_batch_size = batch_size if seg != num_segments-1 or num_batch % batch_size == 0 else num_batch % batch_size
-                segment = stage2_generate(
-                    model,
-                    prompt[:, start_idx:end_idx],
-                    batch_size=current_batch_size
-                )
-                segments.append(segment)
-
-            # Concatenate all the segments
-            output = np.concatenate(segments, axis=0)
-        
-        # Process the ending part of the prompt
-        if output_duration*50 != prompt.shape[-1]:
-            ending = stage2_generate(model, prompt[:, output_duration*50:], batch_size=1)
-            output = np.concatenate([output, ending], axis=0)
-        output = codectool_stage2.ids2npy(output)
-
-        # Fix invalid codes (a dirty solution, which may harm the quality of audio)
-        # We are trying to find better one
-        fixed_output = copy.deepcopy(output)
-        for i, line in enumerate(output):
-            for j, element in enumerate(line):
-                if element < 0 or element > 1023:
-                    counter = Counter(line)
-                    most_frequant = sorted(counter.items(), key=lambda x: x[1], reverse=True)[0][0]
-                    fixed_output[i, j] = most_frequant
-        # save output
-        np.save(output_filename, fixed_output)
-        stage2_result.append(output_filename)
-    return stage2_result
-
-stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=args.stage2_batch_size)
-print(stage2_result)
-print('Stage 2 DONE.\n')
-# convert audio tokens to audio
-def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
-    folder_path = os.path.dirname(path)
-    if not os.path.exists(folder_path):
-        os.makedirs(folder_path)
-    limit = 0.99
-    max_val = wav.abs().max()
-    wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
-    torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)
-# reconstruct tracks
-recons_output_dir = os.path.join(args.output_dir, "recons")
-recons_mix_dir = os.path.join(recons_output_dir, 'mix')
-os.makedirs(recons_mix_dir, exist_ok=True)
-tracks = []
-for npy in stage2_result:
-    codec_result = np.load(npy)
-    decodec_rlt=[]
-    with torch.no_grad():
-        decoded_waveform = codec_model.decode(torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device))
-    decoded_waveform = decoded_waveform.cpu().squeeze(0)
-    decodec_rlt.append(torch.as_tensor(decoded_waveform))
-    decodec_rlt = torch.cat(decodec_rlt, dim=-1)
-    save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy))[0] + ".mp3")
-    tracks.append(save_path)
-    save_audio(decodec_rlt, save_path, 16000)
-# mix tracks
-for inst_path in tracks:
-    try:
-        if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
-            and 'instrumental' in inst_path:
-            # find pair
-            vocal_path = inst_path.replace('instrumental', 'vocal')
-            if not os.path.exists(vocal_path):
-                continue
-            # mix
-            recons_mix = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace('instrumental', 'mixed'))
-            vocal_stem, sr = sf.read(inst_path)
-            instrumental_stem, _ = sf.read(vocal_path)
-            mix_stem = (vocal_stem + instrumental_stem) / 1
-            sf.write(recons_mix, mix_stem, sr)
-    except Exception as e:
-        print(e)
-
-# vocoder to upsample audios
-vocal_decoder, inst_decoder = build_codec_model(args.config_path, args.vocal_decoder_path, args.inst_decoder_path)
-vocoder_output_dir = os.path.join(args.output_dir, 'vocoder')
-vocoder_stems_dir = os.path.join(vocoder_output_dir, 'stems')
-vocoder_mix_dir = os.path.join(vocoder_output_dir, 'mix')
-os.makedirs(vocoder_mix_dir, exist_ok=True)
-os.makedirs(vocoder_stems_dir, exist_ok=True)
-for npy in stage2_result:
-    if 'instrumental' in npy:
-        # Process instrumental
-        instrumental_output = process_audio(
-            npy,
-            os.path.join(vocoder_stems_dir, 'instrumental.mp3'),
-            args.rescale,
-            args,
-            inst_decoder,
-            codec_model
-        )
-    else:
-        # Process vocal
-        vocal_output = process_audio(
-            npy,
-            os.path.join(vocoder_stems_dir, 'vocal.mp3'),
-            args.rescale,
-            args,
-            vocal_decoder,
-            codec_model
-        )
-# mix tracks
-try:
-    mix_output = instrumental_output + vocal_output
-    vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
-    save_audio(mix_output, vocoder_mix, 44100, args.rescale)
-    print(f"Created mix: {vocoder_mix}")
-except RuntimeError as e:
-    print(e)
-    print(f"mix {vocoder_mix} failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
-
-# Post process
-replace_low_freq_with_energy_matched(
-    a_file=recons_mix,     # 16kHz
-    b_file=vocoder_mix,     # 48kHz
-    c_file=os.path.join(args.output_dir, os.path.basename(recons_mix)),
-    cutoff_freq=5500.0
-)