318 lines
14 KiB
Python
318 lines
14 KiB
Python
import sys
|
|
import os
|
|
import folder_paths
|
|
import time
|
|
import copy
|
|
import torch
|
|
import torchaudio
|
|
import numpy as np
|
|
import gc
|
|
from omegaconf import OmegaConf
|
|
from .SongGeneration.codeclm.models import builders
|
|
from .SongGeneration.codeclm.trainer.codec_song_pl import CodecLM_PL
|
|
from .SongGeneration.codeclm.models import CodecLM
|
|
from .SongGeneration.third_party.demucs.models.pretrained import get_model_from_yaml
|
|
current_node_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
auto_prompt_type = ['Pop', 'R&B', 'Dance', 'Jazz', 'Folk', 'Rock', 'Chinese Style', 'Chinese Tradition', 'Metal', 'Reggae', 'Chinese Opera', 'Auto']
|
|
|
|
class Separator():
|
|
def __init__(self, dm_model_path='third_party/demucs/ckpt/htdemucs.pth', dm_config_path='third_party/demucs/ckpt/htdemucs.yaml', gpu_id=0) -> None:
|
|
if torch.cuda.is_available() and gpu_id < torch.cuda.device_count():
|
|
self.device = torch.device(f"cuda:{gpu_id}")
|
|
else:
|
|
self.device = torch.device("cpu")
|
|
self.demucs_model = self.init_demucs_model(dm_model_path, dm_config_path)
|
|
|
|
def init_demucs_model(self, model_path, config_path):
|
|
model = get_model_from_yaml(config_path, model_path)
|
|
model.to(self.device)
|
|
model.eval()
|
|
return model
|
|
|
|
def load_audio(self, f):
|
|
a, fs = torchaudio.load(f)
|
|
if (fs != 48000):
|
|
a = torchaudio.functional.resample(a, fs, 48000)
|
|
if a.shape[-1] >= 48000*10:
|
|
a = a[..., :48000*10]
|
|
else:
|
|
a = torch.cat([a, a], -1)
|
|
return a[:, 0:48000*10]
|
|
|
|
def run(self, audio_path, output_dir='tmp', ext=".flac"):
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
name, _ = os.path.splitext(os.path.split(audio_path)[-1])
|
|
output_paths = []
|
|
|
|
for stem in self.demucs_model.sources:
|
|
output_path = os.path.join(output_dir, f"{name}_{stem}{ext}")
|
|
if os.path.exists(output_path):
|
|
output_paths.append(output_path)
|
|
if len(output_paths) == 1: # 4
|
|
vocal_path = output_paths[0]
|
|
else:
|
|
drums_path, bass_path, other_path, vocal_path = self.demucs_model.separate(audio_path, output_dir, device=self.device)
|
|
for path in [drums_path, bass_path, other_path]:
|
|
os.remove(path)
|
|
full_audio = self.load_audio(audio_path)
|
|
vocal_audio = self.load_audio(vocal_path)
|
|
bgm_audio = full_audio - vocal_audio
|
|
return full_audio, vocal_audio, bgm_audio
|
|
|
|
|
|
|
|
def build_model(Weigths_Path,infer_model_path):
|
|
torch.backends.cudnn.enabled = False
|
|
curent_dir = os.path.join(current_node_path,"SongGeneration")
|
|
RESOLVERS = {
|
|
"eval": lambda x: eval(x),
|
|
"concat": lambda *x: [xxx for xx in x for xxx in xx],
|
|
"get_fname": lambda: os.path.splitext(os.path.basename(sys.argv[1]))[0],
|
|
"load_yaml": lambda x: list(OmegaConf.load(os.path.join(curent_dir, x)))
|
|
}
|
|
|
|
for name, func in RESOLVERS.items():
|
|
if not OmegaConf.has_resolver(name):
|
|
OmegaConf.register_new_resolver(name, func)
|
|
np.random.seed(int(time.time()))
|
|
infer_model_type="new" if "new" in infer_model_path.lower() else "large" if "large" in infer_model_path.lower() else "full" if "full" in infer_model_path.lower() else "base"
|
|
|
|
cfg_path = os.path.join(current_node_path, f'SongGeneration/conf/{infer_model_type}_config.yaml')
|
|
|
|
cfg = OmegaConf.load(cfg_path)
|
|
cfg.mode = 'inference'
|
|
cfg.vae_config=f"{Weigths_Path}/vae/stable_audio_1920_vae.json"
|
|
cfg.vae_model=f"{Weigths_Path}/vae/autoencoder_music_1320k.ckpt"
|
|
cfg.audio_tokenizer_checkpoint=f"Flow1dVAE1rvq_{Weigths_Path}/model_1rvq/model_2_fixed.safetensors"
|
|
cfg.audio_tokenizer_checkpoint_sep=f"Flow1dVAESeparate_{Weigths_Path}/model_septoken/model_2.safetensors"
|
|
cfg.conditioners.type_info.QwTextTokenizer.token_path=os.path.join(current_node_path,"SongGeneration/third_party/Qwen2-7B")
|
|
|
|
audiolm = builders.get_lm_model(cfg)
|
|
checkpoint = torch.load(infer_model_path, map_location='cpu')
|
|
audiolm_state_dict = {k.replace('audiolm.', ''): v for k, v in checkpoint.items() if k.startswith('audiolm')}
|
|
audiolm.load_state_dict(audiolm_state_dict, strict=False)
|
|
audiolm = audiolm.eval()
|
|
#audiolm = audiolm.cuda().to(torch.float16)
|
|
del audiolm_state_dict,checkpoint
|
|
return audiolm,cfg
|
|
|
|
|
|
|
|
# def pre_data(Weigths_Path,dm_model_path,dm_config_path,save_dir,prompt_audio_path,auto_prompt_audio_type,infer_model_type,prompt_pt_path):
|
|
# torch.backends.cudnn.enabled = False
|
|
# curent_dir = os.path.join(current_node_path,"SongGeneration")
|
|
# RESOLVERS = {
|
|
# "eval": lambda x: eval(x),
|
|
# "concat": lambda *x: [xxx for xx in x for xxx in xx],
|
|
# "get_fname": lambda: os.path.splitext(os.path.basename(sys.argv[1]))[0],
|
|
# "load_yaml": lambda x: list(OmegaConf.load(os.path.join(curent_dir, x)))
|
|
# }
|
|
|
|
# for name, func in RESOLVERS.items():
|
|
# if not OmegaConf.has_resolver(name):
|
|
# OmegaConf.register_new_resolver(name, func)
|
|
# np.random.seed(int(time.time()))
|
|
|
|
# cfg_path = os.path.join(current_node_path, f'SongGeneration/conf/{infer_model_type}_config.yaml')
|
|
|
|
# cfg = OmegaConf.load(cfg_path)
|
|
# cfg.mode = 'inference'
|
|
|
|
# cfg.vae_config=f"{Weigths_Path}/vae/stable_audio_1920_vae.json"
|
|
# cfg.vae_model=f"{Weigths_Path}/vae/autoencoder_music_1320k.ckpt"
|
|
|
|
# cfg.audio_tokenizer_checkpoint=f"Flow1dVAE1rvq_{Weigths_Path}/model_1rvq/model_2_fixed.safetensors"
|
|
# cfg.audio_tokenizer_checkpoint_sep=f"Flow1dVAESeparate_{Weigths_Path}/model_septoken/model_2.safetensors"
|
|
# cfg.conditioners.type_info.QwTextTokenizer.token_path=os.path.join(current_node_path,"SongGeneration/third_party/Qwen2-7B")
|
|
# max_duration = cfg.max_dur
|
|
# vae_model=f"{Weigths_Path}/vae/autoencoder_music_1320k.ckpt"
|
|
# vae_config=os.path.join(current_node_path, f'SongGeneration/conf/stable_audio_1920_vae.json')
|
|
# auto_prompt = torch.load(prompt_pt_path,weights_only=False)
|
|
# merge_prompt = [x for sublist in auto_prompt.values() for x in sublist]
|
|
|
|
# if prompt_audio_path is not None:
|
|
# separator = Separator(dm_model_path, dm_config_path)
|
|
# audio_tokenizer = builders.get_audio_tokenizer_model(f"Flow1dVAE1rvq_{Weigths_Path}/model_1rvq/model_2_fixed.safetensors", vae_config,vae_model)
|
|
# audio_tokenizer = audio_tokenizer.eval().cuda()
|
|
|
|
# else:
|
|
# audio_tokenizer = None
|
|
# separator = None
|
|
|
|
# original_item=song_infer_lowram(cfg,separator,audio_tokenizer,merge_prompt,auto_prompt, save_dir,prompt_audio_path,auto_prompt_audio_type,)
|
|
# print("step1 is done.")
|
|
# return copy.deepcopy(original_item),max_duration,cfg
|
|
|
|
|
|
def infer_stage2(item,audiolm,max_duration,lyric,descriptions,cfg_coef = 1.5, temp = 0.9,top_k = 50,top_p = 0.0,record_tokens = True,record_window = 50):
|
|
#ckpt_path = os.path.join(Weigths_Path, 'songgeneration_base/model.pt')
|
|
|
|
item_copy = {
|
|
'pmt_wav': item['pmt_wav'], # 这些是引用，但安全因为后续设为None不影响原始
|
|
'vocal_wav': item['vocal_wav'],
|
|
'bgm_wav': item['bgm_wav'],
|
|
'melody_is_wav': item['melody_is_wav'],
|
|
'idx': item['idx'],
|
|
'wav_path': item['wav_path']
|
|
# 不包含 'tokens' 因为它将在 step2 中生成
|
|
}
|
|
# Define model or load pretrained model
|
|
# model_light = CodecLM_PL(cfg, ckpt_path)
|
|
# model_light = model_light.eval()
|
|
# model_light.audiolm.cfg = cfg
|
|
# model = CodecLM(name = "tmp",
|
|
# lm = model_light.audiolm,
|
|
# audiotokenizer = None,
|
|
# max_duration = max_duration,
|
|
# seperate_tokenizer = None,
|
|
# )
|
|
# del model_light
|
|
# audiolm = builders.get_lm_model(cfg)
|
|
# checkpoint = torch.load(ckpt_path, map_location='cpu')
|
|
# audiolm_state_dict = {k.replace('audiolm.', ''): v for k, v in checkpoint.items() if k.startswith('audiolm')}
|
|
# audiolm.load_state_dict(audiolm_state_dict, strict=False)
|
|
audiolm=audiolm.cuda().to(torch.float16)
|
|
torch.cuda.empty_cache()
|
|
model = CodecLM(name = "tmp",
|
|
lm = audiolm,
|
|
audiotokenizer = None,
|
|
max_duration = max_duration,
|
|
seperate_tokenizer = None,
|
|
)
|
|
|
|
model.set_generation_params(duration=max_duration, extend_stride=5, temperature=temp,
|
|
top_k=top_k, top_p=top_p,cfg_coef=cfg_coef, record_tokens=record_tokens, record_window=record_window)
|
|
|
|
print("model loaded,start inference step2")
|
|
items=inference_lowram_step2(model,lyric,descriptions,item_copy,)
|
|
audiolm = audiolm.cpu()
|
|
del audiolm
|
|
model=None
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return items
|
|
|
|
|
|
|
|
def inference_lowram_step2(model,lyric,descriptions,item,):
|
|
|
|
generate_inp = {
|
|
'lyrics': [lyric.replace(" ", " ")],
|
|
'descriptions': [descriptions],
|
|
'melody_wavs': item['pmt_wav'],
|
|
'vocal_wavs': item['vocal_wav'],
|
|
'bgm_wavs': item['bgm_wav'],
|
|
'melody_is_wav': item['melody_is_wav'],
|
|
}
|
|
with torch.autocast(device_type="cuda", dtype=torch.float16):
|
|
tokens = model.generate(**generate_inp, return_tokens=True)
|
|
item['tokens'] = tokens
|
|
|
|
return item
|
|
|
|
|
|
|
|
def inference_lowram_final(cfg,seperate_tokenizer,max_duration,item,save_dir,save_separate):
|
|
target_wav_name = f"{save_dir}/song_audios{time.strftime('%m%d%H%S')}.flac"
|
|
|
|
model = CodecLM(name = "tmp",
|
|
lm = None,
|
|
audiotokenizer = None,
|
|
max_duration = max_duration,
|
|
seperate_tokenizer = seperate_tokenizer,
|
|
)
|
|
print("model loaded,start inference final...")
|
|
|
|
with torch.no_grad():
|
|
if item["melody_is_wav"]:
|
|
if save_separate :
|
|
wav_seperate = model.generate_audio(item['tokens'], item['pmt_wav'], item['vocal_wav'], item['bgm_wav'], chunked=True, gen_type='mixed')
|
|
wav_vocal = model.generate_audio(item['tokens'], item['pmt_wav'], item['vocal_wav'], item['bgm_wav'], chunked=True, gen_type='vocal')
|
|
wav_bgm = model.generate_audio(item['tokens'], item['pmt_wav'], item['vocal_wav'], item['bgm_wav'], chunked=True, gen_type='bgm')
|
|
else:
|
|
if cfg.gen_type == 'mixed':
|
|
wav_seperate=model.generate_audio(item['tokens'], item['pmt_wav'], item['pmt_wav'], item['bgm_wav'], chunked=True, gen_type='mixed')
|
|
else:
|
|
wav_seperate = model.generate_audio(item['tokens'],chunked=True, gen_type=cfg.sample_rate)
|
|
else:
|
|
if save_separate :
|
|
wav_vocal = model.generate_audio(item['tokens'], chunked=True, gen_type='vocal')
|
|
wav_bgm = model.generate_audio(item['tokens'], chunked=True, gen_type='bgm')
|
|
wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type='mixed')
|
|
else:
|
|
wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type=cfg.gen_type)
|
|
|
|
if save_separate :
|
|
torchaudio.save(f"{save_dir}/vocal_audios{time.strftime('%m%d%H%S')}.flac", wav_vocal[0].cpu().float(), cfg.sample_rate)
|
|
torchaudio.save(f"{save_dir}/bgm_audios{time.strftime('%m%d%H%S')}.flac", wav_bgm[0].cpu().float(), cfg.sample_rate)
|
|
|
|
torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
|
|
# item['tokens']=None
|
|
# item['pmt_wav']=None
|
|
# item['vocal_wav']=None
|
|
# item['bgm_wav']=None
|
|
# item['melody_is_wav']=None
|
|
return {"waveform": wav_seperate[0].cpu().float().unsqueeze(0), "sample_rate": cfg.sample_rate}
|
|
|
|
|
|
|
|
def song_infer_lowram(seperate_tokenizer,separator,audio_tokenizer,prompt_pt_path, save_dir,prompt_audio_path,auto_prompt_audio_type): #item dict
|
|
item = {}
|
|
target_wav_name = f"{save_dir}/song_audios{time.strftime('%m%d%H%S')}.flac"
|
|
melody_is_wav = False
|
|
if prompt_audio_path is not None:
|
|
|
|
pmt_wav, vocal_wav, bgm_wav = separator.run(prompt_audio_path)
|
|
pmt_wav = pmt_wav.cuda()
|
|
vocal_wav = vocal_wav.cuda()
|
|
bgm_wav = bgm_wav.cuda()
|
|
|
|
audio_tokenizer = audio_tokenizer.eval().cuda()
|
|
with torch.no_grad():
|
|
pmt_wav, _ = audio_tokenizer.encode(pmt_wav)
|
|
audio_tokenizer=None
|
|
separator=None
|
|
gc.collect()
|
|
|
|
seperate_tokenizer = seperate_tokenizer.eval().cuda()
|
|
with torch.no_grad():
|
|
vocal_wav, bgm_wav = seperate_tokenizer.encode(vocal_wav, bgm_wav)
|
|
del seperate_tokenizer
|
|
gc.collect()
|
|
|
|
elif auto_prompt_audio_type:
|
|
assert prompt_pt_path is not None ,"prompt模型不能为空,need prmmpt model"
|
|
auto_prompt = torch.load(prompt_pt_path,weights_only=False)
|
|
#assert item["auto_prompt_audio_type"] in auto_prompt_type, f"auto_prompt_audio_type {item['auto_prompt_audio_type']} not found"
|
|
# if auto_prompt_audio_type == 'Auto':
|
|
# prompt_token = merge_prompt[np.random.randint(0, len(merge_prompt))]
|
|
# else:
|
|
# prompt_token = auto_prompt[auto_prompt_audio_type][np.random.randint(0, len(auto_prompt[auto_prompt_audio_type]))]
|
|
prompt_token = auto_prompt[auto_prompt_audio_type][np.random.randint(0, len(auto_prompt[auto_prompt_audio_type]))]
|
|
if torch.cuda.is_available():
|
|
prompt_token = prompt_token.cuda()
|
|
pmt_wav = prompt_token[:,[0],:]
|
|
vocal_wav = prompt_token[:,[1],:]
|
|
bgm_wav = prompt_token[:,[2],:]
|
|
else:
|
|
pmt_wav = None
|
|
vocal_wav = None
|
|
bgm_wav = None
|
|
melody_is_wav = True
|
|
|
|
item['pmt_wav'] = pmt_wav
|
|
item['vocal_wav'] = vocal_wav
|
|
item['bgm_wav'] = bgm_wav
|
|
item['melody_is_wav'] = melody_is_wav
|
|
item["idx"] = 0
|
|
item["wav_path"] = target_wav_name
|
|
|
|
|
|
return item
|
|
|
|
|
|
|
|
|
|
|