384 lines
16 KiB
Python
384 lines
16 KiB
Python
import sys
|
|
import os
|
|
import folder_paths
|
|
import time
|
|
import librosa
|
|
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
|
|
import re
|
|
current_node_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
auto_prompt_type = ['Pop', 'Latin', 'Rock', 'Electronic', 'Metal', 'Country', 'R&B/Soul', 'Ballad', 'Jazz', 'World', 'Hip-Hop', 'Funk', 'Soundtrack','Auto']
|
|
def check_language_by_text(text):
|
|
chinese_pattern = re.compile(r'[\u4e00-\u9fff]')
|
|
english_pattern = re.compile(r'[a-zA-Z]')
|
|
chinese_count = len(re.findall(chinese_pattern, text))
|
|
english_count = len(re.findall(english_pattern, text))
|
|
chinese_ratio = chinese_count / len(text)
|
|
english_ratio = english_count / len(text)
|
|
if chinese_ratio >= 0.2:
|
|
return "zh"
|
|
elif english_ratio >= 0.5:
|
|
return "en"
|
|
else:
|
|
return "en"
|
|
|
|
def load_audio_by_librosa(f):
|
|
a, fs= librosa.load(f, sr=48000)
|
|
a = torch.tensor(a).unsqueeze(0)
|
|
if (fs != 48000):
|
|
a = torchaudio.functional.resample(a, fs, 48000)
|
|
if a.shape[-1] >= 48000*10:
|
|
a = a[..., :48000*10]
|
|
return a[:, 0:48000*10], 48000
|
|
|
|
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):
|
|
try:
|
|
a, fs = torchaudio.load(f)
|
|
except:
|
|
a, fs = load_audio_by_librosa(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 load_gguf_checkpoint_vl(gguf_checkpoint_path):
|
|
from diffusers.utils import is_gguf_available, is_torch_available
|
|
if is_gguf_available() and is_torch_available():
|
|
import gguf
|
|
from gguf import GGUFReader
|
|
from diffusers.quantizers.gguf.utils import SUPPORTED_GGUF_QUANT_TYPES, GGUFParameter
|
|
else:
|
|
raise ImportError("Please install torch and gguf>=0.10.0 to load a GGUF checkpoint in PyTorch.")
|
|
|
|
reader = GGUFReader(gguf_checkpoint_path)
|
|
parsed_parameters = {}
|
|
|
|
for tensor in reader.tensors:
|
|
name = tensor.name
|
|
quant_type = tensor.tensor_type
|
|
|
|
# if the tensor is a torch supported dtype do not use GGUFParameter
|
|
is_gguf_quant = quant_type not in [gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]
|
|
if is_gguf_quant and quant_type not in SUPPORTED_GGUF_QUANT_TYPES:
|
|
_supported_quants_str = "\n".join([str(type) for type in SUPPORTED_GGUF_QUANT_TYPES])
|
|
raise ValueError(
|
|
(
|
|
f"{name} has a quantization type: {str(quant_type)} which is unsupported."
|
|
"\n\nCurrently the following quantization types are supported: \n\n"
|
|
f"{_supported_quants_str}"
|
|
"\n\nTo request support for this quantization type please open an issue here: https://github.com/huggingface/diffusers"
|
|
)
|
|
)
|
|
weights = torch.from_numpy(tensor.data.copy())
|
|
parsed_parameters[name] = GGUFParameter(weights, quant_type=quant_type) if is_gguf_quant else weights
|
|
|
|
del reader
|
|
gc.collect()
|
|
return parsed_parameters
|
|
|
|
|
|
def build_model(Weigths_Path,infer_model_path,version,use_flash_attn,offload_audiolm):
|
|
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') if version=="v1" else os.path.join(current_node_path, f'SongGeneration/conf/{infer_model_type}_config_v2.yaml')
|
|
print(cfg_path)
|
|
cfg = OmegaConf.load(cfg_path)
|
|
cfg.mode = 'inference'
|
|
cfg.lm.use_flash_attn_2 = use_flash_attn
|
|
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")
|
|
cfg.version = version
|
|
cfg.offload_audiolm = offload_audiolm
|
|
|
|
audiolm = builders.get_lm_model(cfg,version,offload_audiolm)
|
|
if not infer_model_path.endswith(".gguf"):
|
|
checkpoint = torch.load(infer_model_path, map_location='cpu',weights_only=False)
|
|
audiolm_state_dict = {k.replace('audiolm.', ''): v for k, v in checkpoint.items() if k.startswith('audiolm')}
|
|
del checkpoint
|
|
audiolm.load_state_dict(audiolm_state_dict, strict=False)
|
|
del audiolm_state_dict
|
|
else:
|
|
from diffusers import GGUFQuantizationConfig
|
|
from diffusers.quantizers.gguf import GGUFQuantizer
|
|
from diffusers.models.model_loading_utils import load_model_dict_into_meta
|
|
g_config = GGUFQuantizationConfig(compute_dtype=torch.float16)
|
|
hf_quantizer = GGUFQuantizer(quantization_config=g_config)
|
|
hf_quantizer.pre_quantized = True
|
|
model_state_dict=load_gguf_checkpoint_vl(infer_model_path)
|
|
gc.collect()
|
|
hf_quantizer._process_model_before_weight_loading(
|
|
audiolm,
|
|
device_map=None,
|
|
state_dict=model_state_dict
|
|
)
|
|
load_model_dict_into_meta(
|
|
audiolm,
|
|
model_state_dict,
|
|
hf_quantizer=hf_quantizer,
|
|
device_map=None,
|
|
dtype=torch.float16,
|
|
|
|
)
|
|
|
|
hf_quantizer._process_model_after_weight_loading(audiolm)
|
|
del model_state_dict
|
|
gc.collect()
|
|
audiolm.eval().to(torch.float16)
|
|
|
|
return audiolm,cfg
|
|
|
|
|
|
def infer_stage2(item,audiolm,max_duration,lyric,descriptions,gen_type,cfg,cfg_coef = 1.5, temp = 0.9,top_k = 50,top_p = 0.0,record_tokens = True,record_window = 50,offload_audiolm = False):
|
|
#ckpt_path = os.path.join(Weigths_Path, 'songgeneration_base/model.pt')
|
|
|
|
item_copy = {
|
|
'pmt_wav': item['pmt_wav'],
|
|
'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']
|
|
}
|
|
if offload_audiolm:
|
|
# from .SongGeneration.codeclm.utils.offload_profiler import OffloadProfiler, OffloadParamParse
|
|
# audiolm_offload_param = OffloadParamParse.parse_config(audiolm, cfg.offload.audiolm)
|
|
# audiolm_offload_param.show()
|
|
# offload_profiler = OffloadProfiler(device_index=0, **(audiolm_offload_param.init_param_dict()))
|
|
# offload_profiler.offload_layer(**(audiolm_offload_param.offload_layer_param_dict()))
|
|
# offload_profiler.clean_cache_wrapper(**(audiolm_offload_param.clean_cache_param_dict()))
|
|
audiolm.to_cuda("cuda")
|
|
else:
|
|
audiolm.cuda()
|
|
|
|
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,gen_type)
|
|
audiolm = audiolm.cpu()
|
|
del audiolm
|
|
model=None
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return items
|
|
|
|
|
|
def inference_lowram_step2(model,lyric,descriptions,item,gen_type):
|
|
|
|
generate_inp = {
|
|
'lyrics': [lyric.replace(" ", " ")] if gen_type != 'bgm' else '.',
|
|
'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 save_with_fallback(path, tensor, sample_rate):
|
|
|
|
if tensor.ndim == 1:
|
|
tensor = tensor.unsqueeze(0)
|
|
|
|
tensor = tensor.detach().cpu().float().contiguous()
|
|
|
|
try:
|
|
import torchaudio
|
|
torchaudio.save(path, tensor, sample_rate)
|
|
return
|
|
except Exception as e:
|
|
print(f"[WARN] torchaudio.save failed, using soundfile: {e}")
|
|
|
|
try:
|
|
audio = tensor.numpy().astype("float32")
|
|
import soundfile as sf
|
|
sf.write(path, audio.T if audio.shape[0] < audio.shape[1] else audio, sample_rate)
|
|
print(f"[INFO] salvo com soundfile: {path}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Falhou soundfile.write também. Path={path}. Erro: {e}")
|
|
|
|
|
|
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_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')
|
|
wav_mix = model.generate_audio(item['tokens'], item['pmt_wav'], item['vocal_wav'], item['bgm_wav'], chunked=True, gen_type='mixed')
|
|
else:
|
|
wav_mix = model.generate_audio(item['tokens'], item['pmt_wav'], item['pmt_wav'], item['bgm_wav'], chunked=True, gen_type=cfg.gen_type)
|
|
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_mix = model.generate_audio(item['tokens'], chunked=True, gen_type='mixed')
|
|
else:
|
|
wav_mix = model.generate_audio(item['tokens'], chunked=True, gen_type=cfg.gen_type)
|
|
|
|
def ensure_valid(name, t):
|
|
if t is None:
|
|
raise ValueError(f"{name} retornou None do modelo")
|
|
if not torch.is_tensor(t[0]):
|
|
raise ValueError(f"{name} retorno inválido: {type(t[0])}")
|
|
if t[0].numel() == 0:
|
|
raise ValueError(f"{name} retornou tensor vazio")
|
|
return t[0]
|
|
|
|
if save_separate:
|
|
save_with_fallback(f"{save_dir}/vocal_audios{time.strftime('%m%d%H%S')}.flac", ensure_valid("vocal", wav_vocal), cfg.sample_rate)
|
|
save_with_fallback(f"{save_dir}/bgm_audios{time.strftime('%m%d%H%S')}.flac", ensure_valid("bgm", wav_bgm), cfg.sample_rate)
|
|
|
|
save_with_fallback(target_wav_name, ensure_valid("mixed", wav_mix), cfg.sample_rate)
|
|
|
|
return {
|
|
"waveform": ensure_valid("mixed", wav_mix).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,lyric): #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:
|
|
print("auto_prompt_audio_type:",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)
|
|
lang = check_language_by_text(lyric)
|
|
#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][lang][np.random.randint(0, len(auto_prompt[auto_prompt_audio_type][lang]))]
|
|
del auto_prompt
|
|
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],:]
|
|
del prompt_token
|
|
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
|
|
item["gt_lyric"] = lyric
|
|
|
|
|
|
return item
|
|
|
|
|
|
|
|
|
|
|