import sys import os import folder_paths import time import copy import torch import torchaudio import numpy as np 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 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 pre_data(Weigths_Path,dm_model_path,dm_config_path,save_dir,prompt_audio_path,auto_prompt_audio_type): torch.backends.cudnn.enabled = False try: OmegaConf.register_new_resolver("eval", lambda x: eval(x)) except: pass try: OmegaConf.register_new_resolver("concat", lambda *x: [xxx for xx in x for xxx in xx]) except: pass try: OmegaConf.register_new_resolver("get_fname", lambda: os.path.splitext(os.path.basename(sys.argv[1]))[0]) except: pass curent_dir = os.path.join(folder_paths.base_path,"custom_nodes/ComfyUI_SongGeneration/SongGeneration") try: OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(os.path.join(curent_dir,x)))) except: pass np.random.seed(int(time.time())) cfg_path = os.path.join(Weigths_Path, 'songgeneration_base/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(folder_paths.base_path,"custom_nodes/ComfyUI_SongGeneration/SongGeneration/third_party/Qwen2-7B") max_duration = cfg.max_dur separator = Separator(dm_model_path, dm_config_path) auto_prompt = torch.load(os.path.join(Weigths_Path,'prompt.pt'),weights_only=False) audio_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint, cfg) if "audio_tokenizer_checkpoint_sep" in cfg.keys(): seperate_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint_sep, cfg) else: seperate_tokenizer = None audio_tokenizer = audio_tokenizer.eval().cuda() if seperate_tokenizer is not None: seperate_tokenizer = seperate_tokenizer.eval().cuda() merge_prompt = [x for sublist in auto_prompt.values() for x in sublist] original_item=song_infer_lowram(seperate_tokenizer,separator,audio_tokenizer,merge_prompt,auto_prompt, save_dir,prompt_audio_path,auto_prompt_audio_type) return copy.deepcopy(original_item),max_duration,cfg def infer_stage2(item,cfg,Weigths_Path,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 torch.cuda.empty_cache() model.lm = model.lm.cuda().to(torch.float16) model.set_generation_params(duration=max_duration, extend_stride=5, temperature=temp, cfg_coef=cfg_coef, top_k=top_k, top_p=top_p, record_tokens=record_tokens, record_window=record_window) items=inference_lowram_step2(model,lyric,descriptions,item_copy) model=None torch.cuda.empty_cache() return items def inference_lowram_step2(model,lyric,descriptions,item): #print(item) pmt_wav = item['pmt_wav'] vocal_wav = item['vocal_wav'] bgm_wav = item['bgm_wav'] melody_is_wav = item['melody_is_wav'] generate_inp = { 'lyrics': [lyric.replace(" ", " ")], 'descriptions': [descriptions], 'melody_wavs': pmt_wav, 'vocal_wavs': vocal_wav, 'bgm_wavs': bgm_wav, 'melody_is_wav': 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,max_duration,item,save_dir): target_wav_name = f"{save_dir}/song_audios{time.strftime('%m%d%H%S')}.flac" seperate_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint_sep, cfg) seperate_tokenizer = seperate_tokenizer.eval().cuda() model = CodecLM(name = "tmp", lm = None, audiotokenizer = None, max_duration = max_duration, seperate_tokenizer = seperate_tokenizer, ) with torch.no_grad(): if 'raw_pmt_wav' in item: wav_seperate = model.generate_audio(item['tokens'], item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True) item['raw_pmt_wav']=None item['raw_vocal_wav']=None item['raw_bgm_wav']=None else: wav_seperate = model.generate_audio(item['tokens'], chunked=True) #torchaudio.save(item['wav_path'], wav_seperate[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,merge_prompt,auto_prompt, 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" if prompt_audio_path: # assert os.path.exists(item['prompt_audio_path']), f"prompt_audio_path {item['prompt_audio_path']} not found" # assert 'auto_prompt_audio_type' not in item, f"auto_prompt_audio_type and prompt_audio_path cannot be used together" pmt_wav, vocal_wav, bgm_wav = separator.run(prompt_audio_path) item['raw_pmt_wav'] = pmt_wav item['raw_vocal_wav'] = vocal_wav item['raw_bgm_wav'] = bgm_wav if pmt_wav.dim() == 2: pmt_wav = pmt_wav[None] if pmt_wav.dim() != 3: raise ValueError("Melody wavs should have a shape [B, C, T].") pmt_wav = list(pmt_wav) if vocal_wav.dim() == 2: vocal_wav = vocal_wav[None] if vocal_wav.dim() != 3: raise ValueError("Vocal wavs should have a shape [B, C, T].") vocal_wav = list(vocal_wav) if bgm_wav.dim() == 2: bgm_wav = bgm_wav[None] if bgm_wav.dim() != 3: raise ValueError("BGM wavs should have a shape [B, C, T].") bgm_wav = list(bgm_wav) if type(pmt_wav) == list: pmt_wav = torch.stack(pmt_wav, dim=0) if type(vocal_wav) == list: vocal_wav = torch.stack(vocal_wav, dim=0) if type(bgm_wav) == list: bgm_wav = torch.stack(bgm_wav, dim=0) pmt_wav = pmt_wav.cuda() vocal_wav = vocal_wav.cuda() bgm_wav = bgm_wav.cuda() pmt_wav, _ = audio_tokenizer.encode(pmt_wav) vocal_wav, bgm_wav = seperate_tokenizer.encode(vocal_wav, bgm_wav) melody_is_wav = False elif auto_prompt_audio_type: #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]))] pmt_wav = prompt_token[:,[0],:] vocal_wav = prompt_token[:,[1],:] bgm_wav = prompt_token[:,[2],:] melody_is_wav = False 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 del audio_tokenizer del seperate_tokenizer del separator return item