update new version

Author: stephenlei

README.md

@@ -14,7 +14,7 @@
 from codeclm.trainer.codec_song_pl import CodecLM_PL
 from codeclm.models import CodecLM
 from third_party.demucs.models.pretrained import get_model_from_yaml
-
+import re
 
 auto_prompt_type = ['Pop', 'R&B', 'Dance', 'Jazz', 'Folk', 'Rock', 'Chinese Style', 'Chinese Tradition', 'Metal', 'Reggae', 'Chinese Opera', 'Auto']
 
@@ -81,6 +81,7 @@ def parse_args():
     return parser.parse_args()
 
 def generate(args):
+    torch.set_num_threads(1)
     ckpt_path = args.ckpt_path
     input_jsonl = args.input_jsonl
     save_dir = args.save_dir
@@ -95,10 +96,9 @@ def generate(args):
 
 
     separator = Separator()
-    auto_prompt = torch.load('ckpt/prompt.pt')
+    auto_prompt = torch.load('tools/new_prompt.pt')
     audio_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint, cfg)
     audio_tokenizer = audio_tokenizer.eval().cuda()
-    merge_prompt = [item for sublist in auto_prompt.values() for item in sublist]
     with open(input_jsonl, "r") as fp:
         lines = fp.readlines()
 
@@ -145,10 +145,7 @@ def generate(args):
             melody_is_wav = False
         elif "auto_prompt_audio_type" in item:
             assert item["auto_prompt_audio_type"] in auto_prompt_type, f"auto_prompt_audio_type {item['auto_prompt_audio_type']} not found"
-            if item["auto_prompt_audio_type"] == "Auto": 
-                prompt_token = merge_prompt[np.random.randint(0, len(merge_prompt))]
-            else:
-                prompt_token = auto_prompt[item["auto_prompt_audio_type"]][np.random.randint(0, len(auto_prompt[item["auto_prompt_audio_type"]]))]
+            prompt_token = auto_prompt[item["auto_prompt_audio_type"]][np.random.randint(0, len(auto_prompt[item["auto_prompt_audio_type"]]))]
             pmt_wav = prompt_token[:,[0],:]
             vocal_wav = prompt_token[:,[1],:]
             bgm_wav = prompt_token[:,[2],:]
@@ -280,6 +277,7 @@ def generate(args):
             fw.writelines(json.dumps(item, ensure_ascii=False)+"\n")
 
 def generate_lowmem(args):
+    torch.set_num_threads(1)
     ckpt_path = args.ckpt_path
     input_jsonl = args.input_jsonl
     save_dir = args.save_dir
@@ -304,8 +302,7 @@ def generate_lowmem(args):
         separator = Separator()
         audio_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint, cfg)
         audio_tokenizer = audio_tokenizer.eval().cuda()
-    auto_prompt = torch.load('ckpt/prompt.pt')
-    merge_prompt = [item for sublist in auto_prompt.values() for item in sublist]
+    auto_prompt = torch.load('tools/new_prompt.pt')
     new_items = []
     for line in lines:
         item = json.loads(line)
@@ -345,10 +342,7 @@ def generate_lowmem(args):
             melody_is_wav = False
         elif "auto_prompt_audio_type" in item:
             assert item["auto_prompt_audio_type"] in auto_prompt_type, f"auto_prompt_audio_type {item['auto_prompt_audio_type']} not found"
-            if item["auto_prompt_audio_type"] == "Auto": 
-                prompt_token = merge_prompt[np.random.randint(0, len(merge_prompt))]
-            else:
-                prompt_token = auto_prompt[item["auto_prompt_audio_type"]][np.random.randint(0, len(auto_prompt[item["auto_prompt_audio_type"]]))]
+            prompt_token = auto_prompt[item["auto_prompt_audio_type"]][np.random.randint(0, len(auto_prompt[item["auto_prompt_audio_type"]]))]
             pmt_wav = prompt_token[:,[0],:]
             vocal_wav = prompt_token[:,[1],:]
             bgm_wav = prompt_token[:,[2],:]
@@ -471,7 +465,8 @@ def generate_lowmem(args):
     seperate_tokenizer.model.model.device = torch.device(device)
     seperate_tokenizer = seperate_tokenizer.eval()
 
-    offload_wav_tokenizer_diffusion =  True if 'offload' in cfg.keys() and 'wav_tokenizer_diffusion' in cfg.offload else False
+    # offload_wav_tokenizer_diffusion =  True if 'offload' in cfg.keys() and 'wav_tokenizer_diffusion' in cfg.offload else False
+    offload_wav_tokenizer_diffusion =  False
     if offload_wav_tokenizer_diffusion:
         sep_offload_param = OffloadParamParse.parse_config(seperate_tokenizer, cfg.offload.wav_tokenizer_diffusion)
         sep_offload_param.show()
@@ -533,6 +528,8 @@ def generate_lowmem(args):
 
 
 if __name__ == "__main__":
+    # 限制模型使用的显存为0.6
+    torch.cuda.set_per_process_memory_fraction(0.55)
     torch.backends.cudnn.enabled = False
     OmegaConf.register_new_resolver("eval", lambda x: eval(x))
     OmegaConf.register_new_resolver("concat", lambda *x: [xxx for xx in x for xxx in xx])
@@ -548,18 +545,29 @@ def generate_lowmem(args):
         res_mem = (total - reserved) / 1024 / 1024 / 1024
         print(f"reserved memory: {res_mem}GB")
 
-        model_name = args.ckpt_path.split("/")[-1]
-        assert model_name in ['songgeneration_base'], f'{model_name} is not supported, currently only songgeneration_base is supported'
-        if model_name == 'songgeneration_base':
+        model_name = args.ckpt_path.split("/")[-1].lower().replace('-', '_')
+        assert model_name in ['songgeneration_base', 'songgeneration_base_new', 'songgeneration_base_full', 'songgeneration_large'], f'{model_name} is not supported, currently only songgeneration_base, songgeneration_base_new, songgeneration_base_full, songgeneration_large are supported. Please download correct files and rename the folder to the corresponding version name.'
+        if model_name == 'songgeneration_base' or model_name == 'songgeneration_base_new' or model_name == 'songgeneration_base_full':
             if res_mem > 24 and not args.low_mem:
                 print("use generate")
                 generate(args)
             else:
                 from codeclm.utils.offload_profiler import OffloadProfiler, OffloadParamParse
                 print("use generate_lowmem")
                 generate_lowmem(args)
+        elif model_name == 'songgeneration_large':
+            if res_mem > 36 and not args.low_mem:
+                print("use generate")
+                generate(args)
+            else:                
+                print("use generate_lowmem")   
+                from codeclm.utils.offload_profiler import OffloadProfiler, OffloadParamParse
+                generate_lowmem(args)
+            
+
+        # elif model_name == 'songgeneration_base_full':
 
     else:
         print("CUDA is not available")
         exit()
-    
+

generate.py

@@ -3,6 +3,9 @@ export PYTHONDONTWRITEBYTECODE=1
 export TRANSFORMERS_CACHE="$(pwd)/third_party/hub"
 export NCCL_HOME=/usr/local/tccl
 export PYTHONPATH="$(pwd)/codeclm/tokenizer/":"$(pwd)":"$(pwd)/codeclm/tokenizer/Flow1dVAE/":"$(pwd)/codeclm/tokenizer/":$PYTHONPATH
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export CUDA_LAUNCH_BLOCKING=0
 
 CKPT_PATH=$1
 JSONL=$2