add lowmem mode

codeclm/models/builders.py

@@ -29,13 +29,29 @@ def get_audio_tokenizer_model(checkpoint_path: str, cfg: omegaconf.DictConfig):
         return None
     if checkpoint_path.startswith('//pretrained/'):
         name = checkpoint_path.split('/', 3)[-1]
-        return AudioTokenizer.get_pretrained(name, cfg.vae_config, cfg.vae_model, 'cpu', mode=cfg.mode)
+        return AudioTokenizer.get_pretrained(name, cfg.vae_config, cfg.vae_model, 'cuda', mode=cfg.mode)
     elif checkpoint_path == "":
         return None
     else:
         name = checkpoint_path
-        return AudioTokenizer.get_pretrained(name, cfg.vae_config, cfg.vae_model, 'cpu', mode=cfg.mode)
+        return AudioTokenizer.get_pretrained(name, cfg.vae_config, cfg.vae_model, 'cuda', mode=cfg.mode)
     
+
+def get_audio_tokenizer_model_cpu(checkpoint_path: str, cfg: omegaconf.DictConfig):
+    from codeclm.tokenizer.audio_tokenizer import AudioTokenizer
+    """Instantiate a compression model."""
+    if checkpoint_path is None:
+        return None
+    if checkpoint_path.startswith('//pretrained/'):
+        name = checkpoint_path.split('/', 3)[-1]
+        return AudioTokenizer.get_pretrained(name, cfg.vae_config, cfg.vae_model, 'cpu', mode=cfg.mode, tango_device='cpu')
+    elif checkpoint_path == "":
+        return None
+    else:
+        name = checkpoint_path
+        return AudioTokenizer.get_pretrained(name, cfg.vae_config, cfg.vae_model, 'cpu', mode=cfg.mode, tango_device='cpu')
+
+
 def get_lm_model(cfg: omegaconf.DictConfig): #-> LMModel:
     """Instantiate a LM."""    
     lm_kwargs = dict_from_config(getattr(cfg, 'lm'))

codeclm/models/codeclm.py

@@ -271,21 +271,24 @@ class CodecLM:
         return gen_tokens
 
     @torch.no_grad()
-    def generate_audio(self, gen_tokens: torch.Tensor, prompt=None, vocal_prompt=None, bgm_prompt=None, chunked=False, gen_type="all"):
+    def generate_audio(self, gen_tokens: torch.Tensor, prompt=None, vocal_prompt=None, bgm_prompt=None, chunked=False, chunk_size=128, gen_type='mixed'):
         """Generate Audio from tokens"""
         assert gen_tokens.dim() == 3
         if self.seperate_tokenizer is not None:
             gen_tokens_song = gen_tokens[:, [0], :]
             gen_tokens_vocal = gen_tokens[:, [1], :]
             gen_tokens_bgm = gen_tokens[:, [2], :]
-            if gen_type == "bgm":
+            if gen_type == 'bgm':
                 gen_tokens_vocal = torch.full_like(gen_tokens_vocal, 3142)
-                vocal_prompt = None
-            elif gen_type == "vocal":
+                if vocal_prompt is not None:
+                    vocal_prompt = torch.zeros_like(vocal_prompt)
+            elif gen_type == 'vocal':
                 gen_tokens_bgm = torch.full_like(gen_tokens_bgm, 9670)
-                bgm_prompt = None
-            # gen_audio_song = self.audiotokenizer.decode(gen_tokens_song, prompt)
-            gen_audio_seperate = self.seperate_tokenizer.decode([gen_tokens_vocal, gen_tokens_bgm], vocal_prompt, bgm_prompt, chunked=chunked)
+                if bgm_prompt is not None:
+                    bgm_prompt = torch.zeros_like(bgm_prompt)
+            else:
+                assert gen_type == 'mixed', f"gen_type {gen_type} not supported"
+            gen_audio_seperate = self.seperate_tokenizer.decode([gen_tokens_vocal, gen_tokens_bgm], vocal_prompt, bgm_prompt, chunked=chunked, chunk_size=chunk_size)
             return gen_audio_seperate
         else:
             gen_audio = self.audiotokenizer.decode(gen_tokens, prompt)

codeclm/tokenizer/Flow1dVAE/generate_1rvq.py

@@ -46,7 +46,6 @@ class Tango:
         
         self.model.eval()
         self.model.init_device_dtype(torch.device(device), torch.float32)
-        print("scaling factor: ", self.model.normfeat.std)
         
         # self.scheduler = DDIMScheduler.from_pretrained( \
         #     scheduler_name, subfolder="scheduler")
@@ -281,3 +280,11 @@ class Tango:
             else:
                 output = torch.cat([output, cur_output], -1)
         return output
+
+    def to(self, device=None, dtype=None, non_blocking=False):
+        if device is not None:
+            self.device = device
+            self.model.device = device
+        self.vae = self.vae.to(device, dtype, non_blocking)
+        self.model = self.model.to(device, dtype, non_blocking)
+        return self

codeclm/tokenizer/Flow1dVAE/generate_2rvq.py

@@ -51,7 +51,6 @@ class Tango:
         
         self.model.eval()
         self.model.init_device_dtype(torch.device(device), torch.float32)
-        print("scaling factor: ", self.model.normfeat.std)
         
         # self.scheduler = DDIMScheduler.from_pretrained( \
         #     scheduler_name, subfolder="scheduler")

codeclm/tokenizer/Flow1dVAE/generate_4rvq.py

@@ -50,7 +50,6 @@ class Tango:
         
         self.model.eval()
         self.model.init_device_dtype(torch.device(device), torch.float32)
-        print("scaling factor: ", self.model.normfeat.std)
         
         # self.scheduler = DDIMScheduler.from_pretrained( \
         #     scheduler_name, subfolder="scheduler")

codeclm/tokenizer/Flow1dVAE/generate_septoken.py

@@ -102,7 +102,6 @@ class Tango:
         
         self.model.eval()
         self.model.init_device_dtype(torch.device(device), torch.float32)
-        print("scaling factor: ", self.model.normfeat.std)
         
         # self.scheduler = DDIMScheduler.from_pretrained( \
         #     scheduler_name, subfolder="scheduler")
@@ -173,7 +172,7 @@ class Tango:
         return codes_vocal, codes_bgm
 
     @torch.no_grad()
-    def code2sound(self, codes, prompt_vocal=None, prompt_bgm=None, duration=40, guidance_scale=1.5, num_steps=20, disable_progress=False, chunked=False):
+    def code2sound(self, codes, prompt_vocal=None, prompt_bgm=None, duration=40, guidance_scale=1.5, num_steps=20, disable_progress=False, chunked=False, chunk_size=128):
         codes_vocal,codes_bgm = codes
         codes_vocal = codes_vocal.to(self.device)
         codes_bgm = codes_bgm.to(self.device)
@@ -188,7 +187,7 @@ class Tango:
         first_latent_codes_length = 0
 
 
-        if (isinstance(prompt_vocal, torch.Tensor)) and (isinstance(prompt_bgm, torch.Tensor)):
+        if(isinstance(prompt_vocal, torch.Tensor) and isinstance(prompt_bgm, torch.Tensor)):
             # prepare prompt
             prompt_vocal = prompt_vocal.to(self.device)
             prompt_bgm = prompt_bgm.to(self.device)
@@ -273,7 +272,7 @@ class Tango:
             output = None
             for i in range(len(latent_list)):
                 latent = latent_list[i]
-                cur_output = self.vae.decode_audio(latent, chunked=chunked)[0].detach().cpu()
+                cur_output = self.vae.decode_audio(latent, chunked=chunked, chunk_size=chunk_size)[0].detach().cpu()
 
                 if output is None:
                     output = cur_output
@@ -301,3 +300,11 @@ class Tango:
         codes=[codes_vocal, codes_bgm]
         wave = self.code2sound(codes, prompt_vocal,prompt_bgm, guidance_scale=1.5, num_steps=steps, disable_progress=disable_progress)
         return wave
+    
+    def to(self, device=None, dtype=None, non_blocking=False):
+        if device is not None:
+            self.device = device
+            self.model.device = device
+        self.vae = self.vae.to(device, dtype, non_blocking)
+        self.model = self.model.to(device, dtype, non_blocking)
+        return self

codeclm/tokenizer/Flow1dVAE/our_MERT_BESTRQ/mert_fairseq/models/musicfm/model/musicfm_25hz.py

@@ -78,7 +78,6 @@ class MusicFM25Hz(nn.Module):
             with open(stat_path, "r") as f:
                 self.stat = json.load(f)
         else:
-            print("No stats file found at `{}`, use default from msd.".format(stat_path))
             self.stat = {"spec_256_cnt": 14394344256, "spec_256_mean": -23.34296658431829, "spec_256_std": 26.189295587132637, "spec_512_cnt": 28677104448, "spec_512_mean": -21.31267396860235, "spec_512_std": 26.52644536245769, "spec_1024_cnt": 57242624832, "spec_1024_mean": -18.852271129208273, "spec_1024_std": 26.443154583585663, "spec_2048_cnt": 114373665600, "spec_2048_mean": -15.638743433896792, "spec_2048_std": 26.115825961611545, "spec_4096_cnt": 228635747136, "spec_4096_mean": -11.715532502794836, "spec_4096_std": 25.763972210234062, "melspec_256_cnt": 14282760192, "melspec_256_mean": -26.962600400166156, "melspec_256_std": 36.13614100912126, "melspec_512_cnt": 14282760192, "melspec_512_mean": -9.108344167718862, "melspec_512_std": 24.71910937988429, "melspec_1024_cnt": 14282760192, "melspec_1024_mean": 0.37302579246531126, "melspec_1024_std": 18.684082325919388, "melspec_2048_cnt": 14282760192, "melspec_2048_mean": 6.768444971712967, "melspec_2048_std": 18.417922652295623, "melspec_4096_cnt": 14282760192, "melspec_4096_mean": 13.617164614990036, "melspec_4096_std": 18.08552130124525, "cqt_cnt": 9373061376, "cqt_mean": 0.46341379757927165, "cqt_std": 0.9543998080910191, "mfcc_256_cnt": 1339008768, "mfcc_256_mean": -11.681755459447485, "mfcc_256_std": 29.183186444668316, "mfcc_512_cnt": 1339008768, "mfcc_512_mean": -2.540581461792183, "mfcc_512_std": 31.93752185832081, "mfcc_1024_cnt": 1339008768, "mfcc_1024_mean": 6.606636263169779, "mfcc_1024_std": 34.151644801729624, "mfcc_2048_cnt": 1339008768, "mfcc_2048_mean": 5.281600844245184, "mfcc_2048_std": 33.12784541220003, "mfcc_4096_cnt": 1339008768, "mfcc_4096_mean": 4.7616569480166095, "mfcc_4096_std": 32.61458906894133, "chromagram_256_cnt": 1339008768, "chromagram_256_mean": 55.15596556703181, "chromagram_256_std": 73.91858278719991, "chromagram_512_cnt": 1339008768, "chromagram_512_mean": 175.73092252759895, "chromagram_512_std": 248.48485148525953, "chromagram_1024_cnt": 1339008768, "chromagram_1024_mean": 589.2947481634608, "chromagram_1024_std": 913.857929063196, "chromagram_2048_cnt": 1339008768, "chromagram_2048_mean": 2062.286388327397, "chromagram_2048_std": 3458.92657915397, "chromagram_4096_cnt": 1339008768, "chromagram_4096_mean": 7673.039107997085, "chromagram_4096_std": 13009.883158267234}
 
         # feature extractor
@@ -90,40 +89,6 @@ class MusicFM25Hz(nn.Module):
         self.use_rvq_target = use_rvq_target
         
         seed = 142
-        if use_rvq_target:
-            try:
-                from .rvq_musicfm import ResidualVectorQuantize
-                
-            except:
-                import sys, os
-                sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-                from rvq_musicfm import ResidualVectorQuantize
-    
-            self.rvq = ResidualVectorQuantize(
-                input_dim = 128*4, 
-                n_codebooks = 8, 
-                codebook_size = 1024, 
-                codebook_dim = 16, 
-                quantizer_dropout = 0.0,
-                )
-            import os
-            if rvq_ckpt_path is not None and os.path.exists(rvq_ckpt_path):
-                state_dict = torch.load(rvq_ckpt_path, map_location="cpu")
-                self.rvq.load_state_dict(state_dict)
-            else:
-                print(f'Checkpoint for rvq `{rvq_ckpt_path}` not found. Using random initialization.')
-
-        else:
-            for feature in self.features:
-                for i in range(num_codebooks):
-                    setattr(
-                        self,
-                        f"quantizer_{feature}", # _{i}
-                        RandomProjectionQuantizer(
-                            n_mels * 4, codebook_dim, codebook_size, seed=seed + i
-                        ),
-                    )
-
         # two residual convolution layers + one projection layer
         self.conv = Conv2dSubsampling(
             1, conv_dim, encoder_dim, strides=[2, 2], n_bands=n_mels
@@ -247,16 +212,8 @@ class MusicFM25Hz(nn.Module):
     @torch.no_grad()
     def tokenize(self, x):
         out = {}
-        for key in x.keys():
-            if self.use_rvq_target:
-                self.rvq.eval()
-                quantized_prompt_embeds, codes, _, commitment_loss, codebook_loss, rvq_usage = self.rvq(x[key].permute((0, 2, 1)))
-                out[key] = torch.cat([codes[:, idx, :] for idx in range(int(self.codebook_size//1024))], dim=-1)
-            else:
-                layer = getattr(self, "quantizer_%s" % key)
-                out[key] = layer(x[key])
-        return out
-
+        raise NotImplementedError("tokenize is not implemented")
+        
     def get_targets(self, x):
         x = self.preprocessing(x, features=self.features) # -> {'melspec_2048': Tensor{Size([3, 128, 3000]) cuda:0 f32}}
         x = self.normalize(x)

codeclm/tokenizer/audio_tokenizer.py

@@ -78,7 +78,8 @@ class AudioTokenizer(ABC, nn.Module):
             vae_config: str,
             vae_model: str,
             device: tp.Union[torch.device, str] = 'cpu', 
-            mode='extract'
+            mode='extract',
+            tango_device:str='cuda'
             ) -> 'AudioTokenizer':
         """Instantiate a AudioTokenizer model from a given pretrained model.
 
@@ -91,11 +92,11 @@ class AudioTokenizer(ABC, nn.Module):
         if name.split('_')[0] == 'Flow1dVAESeparate':
             model_type = name.split('_', 1)[1]
             logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = Flow1dVAESeparate(model_type, vae_config, vae_model)
+            model = Flow1dVAESeparate(model_type, vae_config, vae_model, tango_device=tango_device)
         elif name.split('_')[0] == 'Flow1dVAE1rvq':
             model_type = name.split('_', 1)[1]
             logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = Flow1dVAE1rvq(model_type, vae_config, vae_model)
+            model = Flow1dVAE1rvq(model_type, vae_config, vae_model, tango_device=tango_device)
         else:
             raise NotImplementedError("{} is not implemented in models/audio_tokenizer.py".format(
                 name))
@@ -108,12 +109,13 @@ class Flow1dVAE1rvq(AudioTokenizer):
         model_type: str = "model_2_fixed.safetensors",
         vae_config: str = "",
         vae_model: str = "",
+        tango_device: str = "cuda"
         ):
         super().__init__()
 
         from codeclm.tokenizer.Flow1dVAE.generate_1rvq import Tango
         model_path = model_type
-        self.model = Tango(model_path=model_path, vae_config=vae_config, vae_model=vae_model, device='cuda')
+        self.model = Tango(model_path=model_path, vae_config=vae_config, vae_model=vae_model, device=tango_device)
         print ("Successfully loaded checkpoint from:", model_path)
 
             
@@ -176,6 +178,15 @@ class Flow1dVAE1rvq(AudioTokenizer):
         assert n <= self.total_codebooks
         self.n_quantizers = n
 
+    def to(self, device=None, dtype=None, non_blocking=False):
+        self = super(Flow1dVAE1rvq, self).to(device, dtype, non_blocking)
+        self.model = self.model.to(device, dtype, non_blocking)
+        return self
+    
+    def cuda(self, device=None):
+        if device is None:
+            device = 'cuda:0'
+        return super(Flow1dVAE1rvq, self).cuda(device)
 
 class Flow1dVAESeparate(AudioTokenizer):
     def __init__(
@@ -183,12 +194,13 @@ class Flow1dVAESeparate(AudioTokenizer):
         model_type: str = "model_2.safetensors",
         vae_config: str = "",
         vae_model: str = "",
+        tango_device: str = "cuda"
         ):
         super().__init__()
 
         from codeclm.tokenizer.Flow1dVAE.generate_septoken import Tango
         model_path = model_type
-        self.model = Tango(model_path=model_path, vae_config=vae_config, vae_model=vae_model, device='cuda')
+        self.model = Tango(model_path=model_path, vae_config=vae_config, vae_model=vae_model, device=tango_device)
         print ("Successfully loaded checkpoint from:", model_path)
 
             
@@ -208,9 +220,9 @@ class Flow1dVAESeparate(AudioTokenizer):
         return codes_vocal, codes_bgm
     
     @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt_vocal = None, prompt_bgm = None, chunked=False):
+    def decode(self, codes: torch.Tensor, prompt_vocal = None, prompt_bgm = None, chunked=False, chunk_size=128):
         wav = self.model.code2sound(codes, prompt_vocal=prompt_vocal, prompt_bgm=prompt_bgm, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False, chunked=chunked) # [B,N,T] -> [B,T]
+                                    num_steps=50, disable_progress=False, chunked=chunked, chunk_size=chunk_size) # [B,N,T] -> [B,T]
         return wav[None]
 
     
@@ -251,3 +263,14 @@ class Flow1dVAESeparate(AudioTokenizer):
         assert n >= 1
         assert n <= self.total_codebooks
         self.n_quantizers = n
+
+    def to(self, device=None, dtype=None, non_blocking=False):
+        self = super(Flow1dVAESeparate, self).to(device, dtype, non_blocking)
+        self.model = self.model.to(device, dtype, non_blocking)
+        return self
+
+    def cuda(self, device=None):
+        if device is None:
+            device = 'cuda:0'
+        self = super(Flow1dVAESeparate, self).cuda(device)
+        return self

codeclm/trainer/codec_song_pl.py

@@ -49,9 +49,7 @@ class CodecLM_PL(pl.LightningModule):
         # 3) Load pretrained checkpoint (if any)
         checkpoint = torch.load(ckpt_path, map_location='cpu')
         missing, unexpected = self.load_state_dict(checkpoint, strict=False)
-        print(f'-------------Missing--------------\n{missing}')
-        print(f'-------------Unexpected--------------\n{unexpected}')
-        print("successfully load deepspeed pretrained model {}".format(ckpt_path))
+        print("successfully load pretrained model {}".format(ckpt_path))
         # 4) Build metrics
         self.val_steps = []
         self.train_slide_acc = []
@@ -70,7 +68,6 @@ class CodecLM_PL(pl.LightningModule):
         ) for _ in range(self.audiolm.code_depth)])
 
         self.epoch = 0
-        print("++++++++++++++++ training <song> +++++++++++++++++")
 
     # TODO: move this part to loader
     def generate_mask_and_end_token(self, x, sequence_lengths, end_id=16384):

codeclm/utils/offload_profiler.py

@@ -0,0 +1,505 @@
+import torch
+from torch.func import functional_call
+import queue
+import threading
+from typing import Dict, List, Any
+import omegaconf
+from pydantic import BaseModel, validator
+from typing import Optional
+from functools import wraps
+
+def _callable_once(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        method_called_flag = f"_called_once_{func.__name__}"
+        if getattr(self, method_called_flag, False):
+            raise RuntimeError(f"{func.__name__} can only be called once.")
+        setattr(self, method_called_flag, True)
+        return func(self, *args, **kwargs)
+    return wrapper
+
+class OffloadCleanCacheWrapperParam(BaseModel):
+    module: Any 
+    method_name: str
+    diff_mem_gb_thre: float
+
+class OffloadParam(BaseModel):
+    offload_module: Any 
+    cpu_mem_gb: float
+    pre_copy_step: Optional[int] = None
+    clean_cache_after_forward: Optional[bool] = None
+    dtype: Optional[str] = None 
+    offload_layer_dict: Dict[str, int] = {}
+    ignore_layer_list: List[str] = []
+    clean_cache_wrapper: Optional[OffloadCleanCacheWrapperParam] = None
+    debug: Optional[bool] = None
+
+    @validator('dtype')
+    def parse_dtype(cls, value):
+        if value is None:
+            return None
+        dtype_map = {
+            'torch.float16': torch.float16,
+            'torch.float32': torch.float32,
+            'torch.float64': torch.float64,
+            'torch.int64': torch.int64,
+        }
+        if value not in dtype_map:
+            raise ValueError(f"Unsupported dtype: {value}")
+        return dtype_map[value]
+    
+    def init_param_dict(self):
+        param_dict = {}
+        param_dict['cpu_mem_gb'] = self.cpu_mem_gb
+        if self.pre_copy_step is not None:
+            param_dict['pre_copy_step'] = self.pre_copy_step
+        if self.clean_cache_after_forward is not None:
+            param_dict['clean_cache_after_forward'] = self.clean_cache_after_forward
+        if self.debug is not None:
+            param_dict['debug'] = self.debug
+        
+        return param_dict
+        
+    def offload_layer_param_dict(self):
+        param_dict = {}
+        param_dict['module'] = self.offload_module
+        param_dict['offload_layer_dict'] = self.offload_layer_dict
+        param_dict['ignore_layer_list'] = self.ignore_layer_list
+        param_dict['dtype'] = self.dtype
+
+        return param_dict
+    
+    def clean_cache_param_dict(self):
+        param_dict = {}
+        if self.clean_cache_wrapper is not None:
+            param_dict['module'] = self.clean_cache_wrapper.module
+            param_dict['method_name'] = self.clean_cache_wrapper.method_name
+            param_dict['diff_mem_gb_thre'] = self.clean_cache_wrapper.diff_mem_gb_thre
+
+        return param_dict
+    
+    @staticmethod
+    def recursive_print(model, indent=0):
+        for field_name, field_info in model.__fields__.items():
+            field_value = getattr(model, field_name)
+            print(" " * indent + f"{field_name}:")
+
+            if issubclass(type(field_value), BaseModel):
+                print(" " * (indent + 2) + f"--- Nested model: {field_value.__class__.__name__}")
+                OffloadParam.recursive_print(field_value, indent + 4) 
+            else:
+                print(" " * (indent + 2) + f"class: {field_value.__class__.__name__}")
+                if isinstance(field_value, torch.nn.Module):
+                    pass
+                else:
+                    print(" " * (indent + 2) + f"value: {field_value}")
+
+    def show(self):
+        print("-"*20 + "[OffloadParam]" + "-"*20)
+        OffloadParam.recursive_print(self)
+        print("-"*40)
+
+
+class OffloadParamParse:
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def _get_model(root_model: torch.nn.Module, model_dir: str):
+        assert(model_dir.startswith("self")), f"model_dir {model_dir} must startswith `self`"
+        model = root_model
+        for layer in model_dir.split('.'):
+            if layer == "self":
+                continue
+            assert(hasattr(model, layer)), f"model not has layer [{layer}]!"
+            model = getattr(model, layer)
+        return model
+
+    @staticmethod
+    def parse_config(root_model: torch.nn.Module, cfg: omegaconf.DictConfig)->OffloadParam:
+        assert(hasattr(cfg, "offload_module") and hasattr(cfg, "cpu_mem_gb") and hasattr(cfg, "dtype"))
+        
+        offload_module = OffloadParamParse._get_model(root_model, cfg.offload_module)
+        cpu_mem_gb = cfg.cpu_mem_gb
+        dtype = cfg.dtype
+
+        pre_copy_step = cfg.pre_copy_step \
+            if hasattr(cfg, "pre_copy_step") else None
+
+        clean_cache_after_forward = cfg.clean_cache_after_forward \
+            if hasattr(cfg, "clean_cache_after_forward") else None
+            
+        offload_layer_dict = {k: v for k, v in cfg.offload_layer_dict.items()} \
+            if hasattr(cfg, "offload_layer_dict") else {}
+
+        ignore_layer_list = cfg.ignore_layer_list \
+            if hasattr(cfg, "ignore_layer_list") else []
+        
+        debug = cfg.debug if hasattr(cfg, "debug") else None
+        
+        clean_cache_wrapper = None
+        if hasattr(cfg, "clean_cache_wrapper"):
+            clean_cache_cfg = cfg.clean_cache_wrapper
+            cc_module = OffloadParamParse._get_model(root_model, clean_cache_cfg.module)
+            cc_method_name = clean_cache_cfg.method_name
+            diff_mem_gb_thre = clean_cache_cfg.diff_mem_gb_thre
+            clean_cache_wrapper = OffloadCleanCacheWrapperParam(
+                                        module=cc_module, 
+                                        method_name=cc_method_name, 
+                                        diff_mem_gb_thre=diff_mem_gb_thre)
+        
+        return OffloadParam(
+            offload_module=offload_module,
+            cpu_mem_gb=cpu_mem_gb,
+            pre_copy_step=pre_copy_step,
+            clean_cache_after_forward=clean_cache_after_forward,
+            dtype=dtype,
+            offload_layer_dict=offload_layer_dict,
+            ignore_layer_list=ignore_layer_list,
+            clean_cache_wrapper=clean_cache_wrapper,
+            debug=debug
+            )
+
+
+class LayerParamStruct:
+    def __init__(self):
+        self.count = 0
+        self.device_state = None
+
+
+class OffloadProfiler:
+    def __init__(self, device_index=0, cpu_mem_gb=-1, pre_copy_step=1, clean_cache_after_forward=False, debug=False):
+        self.clean_cache_after_forward = clean_cache_after_forward
+        self.cpu_mem_gb = cpu_mem_gb
+        self.cpu_mem_b_count = 0
+        self.device_index = device_index
+        self.execution_order = []
+        self.execution_order_idx = {} 
+        self.pin_memory = False
+        test_data = torch.rand(1,1, device='cpu')
+        pin_data = test_data.pin_memory()
+        self.pin_memory = pin_data.is_pinned()
+        print(f"pin:{self.pin_memory}")
+        self.copy_stream = torch.cuda.Stream() 
+        self.copy_queue = queue.Queue() 
+        self.layer_param:Dict[str, LayerParamStruct] = {} 
+        self.model_map = {}
+        self.stop_flag = False
+        self.copy_condition = threading.Condition()
+        self.queue_condition = threading.Condition()
+        self.mem_line_b = 0
+
+        self.copy_thread = threading.Thread(target=self._copy_thread_fun)
+        self.copy_thread.daemon = True
+        self.copy_thread.start()
+
+        self.cur_copy_idx = 0 
+        self.execute_over = False
+        self.pre_copy_step = pre_copy_step
+
+        self.tmp_state_list = []
+        self.tmp_state_idx = 0
+        for i in range(pre_copy_step + 2):
+            self.tmp_state_list.append(None)
+
+        self.debug = debug
+
+    def stop(self):
+        self.stop_flag = True
+        with self.queue_condition:
+            self.queue_condition.notify()
+        self.copy_thread.join()
+
+        del self.layer_param
+        del self.model_map
+        del self.copy_stream
+
+    def _copy_thread_fun(self):
+        while self.stop_flag == False:
+            layer_name = "--"
+            with self.queue_condition:
+                while self.copy_queue.qsize() == 0 and self.stop_flag == False:
+                    self.queue_condition.wait()
+                if self.stop_flag == True:
+                    break
+                layer_name = self.copy_queue.get()
+            with torch.cuda.stream(self.copy_stream):
+                if layer_name in self.model_map:
+                    model = self.model_map[layer_name]
+                    self.tmp_state_list[self.tmp_state_idx] = {
+                        k: v.to(torch.device(f"cuda:{self.device_index}"), non_blocking=False)
+                        for k, v in model.state_dict().items()
+                    }
+                    self.copy_stream.synchronize()
+
+                    device_state = self.tmp_state_list[self.tmp_state_idx]
+                    self.tmp_state_idx = (self.tmp_state_idx + 1) % len(self.tmp_state_list)
+
+                    with self.copy_condition:
+                        if layer_name in self.layer_param:
+                            self.layer_param[layer_name].count += 1
+                        else:
+                            self.layer_param[layer_name] = LayerParamStruct()
+                            self.layer_param[layer_name].count = 1
+                        self.layer_param[layer_name].device_state = device_state
+                        self.copy_condition.notify()
+                else:
+                    print(f"get model error! {layer_name}")
+        print("copy thread stop..")
+
+    def _get_new_step_copy_begin_end(self, tag_name):
+        
+        pre_copy_step = self.pre_copy_step
+        pre_copy_step = min(pre_copy_step, len(self.execution_order) // 2)
+        
+        cur_exe_idx = self.execution_order_idx[tag_name]
+        copy_begin = self.cur_copy_idx
+        copy_end = cur_exe_idx + pre_copy_step + 1
+        if copy_end - copy_begin > len(self.execution_order):
+            copy_end %= len(self.execution_order)
+        if copy_end - copy_begin > pre_copy_step + 1 or copy_end - copy_begin < 0:
+            # jump
+            self.cur_copy_idx = cur_exe_idx
+            copy_begin, copy_end = self._get_new_step_copy_begin_end(tag_name=tag_name)
+        return copy_begin, copy_end
+    
+    def make_forward_wrapper(self, module, tag_name, ignore_layer_list=[]):
+        original_forward = module.forward
+        layer_param_size = 0
+        for name, param in module.named_parameters():
+            layer_param_size += param.data.numel() * param.data.element_size() / 1024 / 1024 #MB
+        
+        taget_cpu_mem_b = self.cpu_mem_gb * 1024 * 1024 * 1024
+        offload = False
+        for name, param in module.named_parameters():
+            p_name = f"{tag_name}.{name}" if tag_name else name
+            for i_layer in ignore_layer_list:
+                if p_name.startswith(i_layer):
+                    if self.debug:
+                        print(f"ignore layer param: {p_name}")
+                    continue
+
+            if taget_cpu_mem_b >= 0 and self.cpu_mem_b_count >= taget_cpu_mem_b:
+                break
+            cpu_data = torch.empty_strided(size=param.data.size(),
+                                        stride=param.data.stride(),
+                                        dtype=param.data.dtype,
+                                        layout=param.data.layout,
+                                        device='cpu',
+                                        pin_memory=self.pin_memory)
+            cpu_data.copy_(param.data)
+            param.data = cpu_data
+
+            param_size = param.data.numel() * param.data.element_size()
+            self.cpu_mem_b_count += param_size
+            offload = True
+        if self.debug:
+            print(f"layer: {tag_name}, type: {module.__class__.__name__}, size(MB): {layer_param_size}, offload: {offload}, sum_offload_size(MB): {self.cpu_mem_b_count/1024/1024}")
+        
+        if offload:
+            copy_condition = self.copy_condition
+            queue_condition = self.queue_condition
+            copy_queue = self.copy_queue
+            layer_param = self.layer_param
+            def forward_wrapper(*args, **kwargs):
+                module.forward = original_forward
+
+                execute_over = False if tag_name not in self.execution_order_idx else True
+                if execute_over == False:
+                    self.model_map[tag_name] = module
+                    self.execution_order.append(tag_name)
+                    self.execution_order_idx[tag_name] = len(self.execution_order) - 1
+                    copy_queue.put(tag_name)
+                    with queue_condition:
+                        queue_condition.notify()
+                else: 
+                
+                    copy_begin, copy_end = self._get_new_step_copy_begin_end(tag_name=tag_name)
+                    if copy_end > copy_begin:
+                        for idx in range(copy_begin, copy_end):
+                            idx = idx % len(self.execution_order)
+                            copy_tag_name = self.execution_order[idx]
+                            copy_queue.put(copy_tag_name)
+                            with queue_condition:
+                                queue_condition.notify()
+
+                        self.cur_copy_idx = copy_end % len(self.execution_order)
+                
+                run_state = None
+                with self.copy_condition:
+                    while tag_name not in self.layer_param:
+                        copy_condition.wait()
+                    run_state = self.layer_param[tag_name].device_state
+                    self.layer_param[tag_name].count -= 1
+                    
+                module.eval()
+                with torch.no_grad():
+                    output = functional_call(module, run_state, args=args, kwargs=kwargs)
+                with self.copy_condition:
+                    if self.layer_param[tag_name].count == 0:
+                        del self.layer_param[tag_name]
+                diff_mem_b_thre = 1 * (1024 ** 3)
+                if self.clean_cache_after_forward:
+                    reserved = torch.cuda.memory_reserved()
+                    if reserved > self.mem_line_b:
+                        torch.cuda.empty_cache()
+                        cur_reserved = torch.cuda.memory_reserved()
+                        diff_mem = reserved - cur_reserved
+                        if diff_mem > diff_mem_b_thre:
+                            self.mem_line_b = cur_reserved + (reserved - cur_reserved) / 2 + 10
+                        else:
+                            self.mem_line_b = reserved + 10
+                        if self.debug:
+                            print(f"child mem line update, clean cache:{reserved/1024/1024}, cur mem: {cur_reserved/1024/1024}  new limit: {self.mem_line_b / 1024 / 1024}, child name: {tag_name}")
+                    
+                module.forward = forward_wrapper
+                return output
+            module.forward = forward_wrapper
+        
+        torch.cuda.empty_cache()
+        return module
+    
+    def reset_empty_cache_mem_line(self):
+        self.mem_line_b = 0
+        torch.cuda.empty_cache()
+    
+    def clean_cache_wrapper(self, module, method_name='', diff_mem_gb_thre=1):
+        if not hasattr(module, method_name) or not callable(getattr(module, method_name)):
+            print(f"no this method {method_name}")
+            return module
+        
+        original_fun = getattr(module, method_name)
+        diff_mem_b_thre = diff_mem_gb_thre * (1024 ** 3)
+        self.reset_empty_cache_mem_line()
+
+        def clean_wrapper(*args, **kwargs):
+            setattr(module, method_name, original_fun)
+            output = original_fun(*args, **kwargs)
+            reserved = torch.cuda.memory_reserved()
+            if reserved > self.mem_line_b:
+                torch.cuda.empty_cache()
+                cur_reserved = torch.cuda.memory_reserved()
+                diff_mem = reserved - cur_reserved
+                if diff_mem > diff_mem_b_thre:
+                    self.mem_line_b = cur_reserved + (reserved - cur_reserved) / 2 + 10
+                else:
+                    self.mem_line_b = reserved + 10
+
+                if self.debug:
+                    print(f"mem line update, clean cache:{reserved/1024/1024}, cur mem: {cur_reserved/1024/1024}  new limit: {self.mem_line_b / 1024 / 1024}")
+            setattr(module, method_name, clean_wrapper)
+            return output
+        
+        setattr(module, method_name, clean_wrapper)
+        return module
+    
+    @_callable_once
+    def offload_layer(self, module, offload_layer_dict={},  ignore_layer_list=[], dtype:torch.dtype = None):
+        return self._offload_layer(
+                                    module=module,
+                                    tag="",
+                                    offload_layer_dict=offload_layer_dict,
+                                    ignore_layer_list=ignore_layer_list,
+                                    dtype=dtype
+                                    )
+    
+    def _offload_layer(self, module, tag="", offload_layer_dict={},  ignore_layer_list=[], dtype:torch.dtype = None):
+        """
+            Offload specific layers of a PyTorch model to a specified depth.
+            A model can only be offloaded once.
+
+            Args:
+                module (torch.nn.Module): 
+                    The PyTorch model containing the layers to offload. This is the model that will be modified in place.
+                
+                tag (str, optional): 
+                    A string identifier for the model. 
+                    Default is an empty string.
+                
+                offload_layer_dict (dict, optional): 
+                    A dictionary where keys are layer names and values represent the depth at which the offloading should occur. 
+                    For example, 
+                    ```offload_layer_dict = {'cfm_wrapper': 5, 'hubert': 4}``` means that the `cfm_wrapper` layer should 
+                    be offloaded at depth 5, and the `hubert` layer should be offloaded at depth 4.
+                    Default is an empty dictionary.
+                
+                ignore_layer_list (list, optional): 
+                    A list of layer names or parameter identifiers to be ignored during the offloading process. 
+                    Layers in this list will not be offloaded, even if they are present in the `offload_layer_dict`. 
+                     For example, 
+                    ```ignore_layer_list = ['cfm_wrapper.estimator.h', 'cfm_wrapper.estimator.adaln_single']```
+                    means that layers starting with `cfm_wrapper.estimator.h` or  'cfm_wrapper.estimator.adaln_single' will not be offload.
+                    Default is an empty list.
+                
+                dtype (torch.dtype, optional): 
+                    The data type (e.g., `torch.float16`, `torch.float32`) to which the offloaded layers should be converted. 
+                    If `None`, the data type of the layers will remain unchanged. Default is `None`.
+
+            Returns:
+                None
+        """
+        for p in module._parameters.values():
+            if p is not None:
+                p.data = p.data.to(torch.device(f"cuda:{self.device_index}"))
+                if dtype is not None:
+                    p.data = p.data.to(dtype)
+        for b in module._buffers.values():
+            if b is not None:
+                b.data = b.data.to(torch.device(f"cuda:{self.device_index}"))
+                if dtype is not None:
+                    b.data = b.data.to(dtype)
+        for attr_name, attr in module.__dict__.items():
+            if isinstance(attr, torch.Tensor) and not attr_name.startswith('_'):
+                attr.data = attr.data.to(torch.device(f"cuda:{self.device_index}"))
+                if dtype is not None:
+                    attr.data = attr.data.to(dtype)
+
+        for name, child in module.named_children():
+            current_tag = f"{tag}.{name}" if tag else name
+            child = child.to(torch.device(f"cuda:{self.device_index}"))
+            if dtype is not None:
+                child = child.to(dtype)
+
+            torch.cuda.empty_cache()
+            setattr(module, name, child)
+            pre_name = current_tag.split('.')[0]
+            if pre_name not in offload_layer_dict:
+                param_size = 0
+                for p in child.parameters():
+                    param_size += p.data.numel() * p.data.element_size()
+                param_size = param_size / 1024 / 1024
+                if self.debug:
+                    print(f"not offload layer {current_tag}, size: {param_size}MB")
+                continue
+            
+            has_children = any(child.named_children())
+            layer_count = current_tag.count('.') + 1
+            
+            layer_deep = offload_layer_dict[pre_name]
+            if layer_count >= layer_deep:
+                has_children = False 
+            
+            if has_children:
+                self._offload_layer(module=child, 
+                                   tag=current_tag, 
+                                   offload_layer_dict=offload_layer_dict, 
+                                   ignore_layer_list=ignore_layer_list,
+                                   dtype=dtype)
+                continue
+
+            ignore = False
+            for i_layer in ignore_layer_list:
+                if current_tag.startswith(i_layer):
+                    ignore = True
+                    if self.debug:
+                        print(f"ignore layer offload: {current_tag}")
+                    break
+    
+            if hasattr(child, "forward") and not ignore:
+                child = self.make_forward_wrapper(
+                    child, current_tag, ignore_layer_list=ignore_layer_list
+                )
+        return module
+    
+    def get_execution_order(self):
+        return self.execution_order

download.py

@@ -7,7 +7,7 @@ def download_model(local_dir):
     downloaded_path = snapshot_download(
         repo_id=repo_id,
         local_dir=local_dir,
-        revision="0c80d30",
+        revision="647f0a5",
         token=os.environ.get("HF_TOKEN"), 
         ignore_patterns=['.git*']
     )

generate.py

@@ -1,5 +1,7 @@
+from hmac import new
 import sys
 import os
+import argparse
 
 import time
 import json
@@ -7,11 +9,13 @@ import torch
 import torchaudio
 import numpy as np
 from omegaconf import OmegaConf
-
+from codeclm.models import builders
+import gc
 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
 
+
 auto_prompt_type = ['Pop', 'R&B', 'Dance', 'Jazz', 'Folk', 'Rock', 'Chinese Style', 'Chinese Tradition', 'Metal', 'Reggae', 'Chinese Opera', 'Auto']
 
 class Separator:
@@ -34,8 +38,6 @@ class Separator:
             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"):
@@ -59,38 +61,146 @@ class Separator:
         return full_audio, vocal_audio, bgm_audio
 
 
+def parse_args():
+    parser = argparse.ArgumentParser(description='Song Generation Script')
+    
+    # 必需参数
+    parser.add_argument('--ckpt_path', type=str, required=True,
+                      help='Path to the checkpoint directory containing config.yaml and model.pt')
+    parser.add_argument('--input_jsonl', type=str, required=True,
+                      help='Path to input JSONL file containing generation tasks')
+    parser.add_argument('--save_dir', type=str, required=True,
+                      help='Directory to save generated audio files and results')
+    # 可选参数
+    parser.add_argument('--generate_type', type=str, default='mixed',
+                      help='Type of generation: "vocal" or "bgm" or "separate" or "mixed" (default: "mixed")')
+    parser.add_argument('--use_flash_attn', action='store_true',
+                      help='Whether to use flash attention (default: False)')
+    parser.add_argument('--low_mem', action='store_true',
+                      help='Whether to use low memory mode (default: False)')
+    return parser.parse_args()
 
-if __name__ == "__main__":
-    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])
-    OmegaConf.register_new_resolver("get_fname", lambda: os.path.splitext(os.path.basename(sys.argv[1]))[0])
-    OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(x)))
-    np.random.seed(int(time.time()))    
-    ckpt_path = sys.argv[1]
-    input_jsonl = sys.argv[2]
-    save_dir = sys.argv[3]
-    gen_type = sys.argv[4] if len(sys.argv) > 4 else "all"
+def generate(args):
+    ckpt_path = args.ckpt_path
+    input_jsonl = args.input_jsonl
+    save_dir = args.save_dir
     cfg_path = os.path.join(ckpt_path, 'config.yaml')
     ckpt_path = os.path.join(ckpt_path, 'model.pt')
     cfg = OmegaConf.load(cfg_path)
+    cfg.lm.use_flash_attn_2 = args.use_flash_attn
+    print(f"use_flash_attn: {args.use_flash_attn}")
     cfg.mode = 'inference'
     max_duration = cfg.max_dur
+    gen_type = args.generate_type
     
-    # Define model or load pretrained model
-    model_light = CodecLM_PL(cfg, ckpt_path)
 
-    model_light = model_light.eval().cuda()
-    model_light.audiolm.cfg = cfg
-    model = CodecLM(name = "tmp",
-        lm = model_light.audiolm,
-        audiotokenizer = model_light.audio_tokenizer,
-        max_duration = max_duration,
-        seperate_tokenizer = model_light.seperate_tokenizer,
-    )
     separator = Separator()
     auto_prompt = torch.load('ckpt/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()
+
+        
+    new_items = []
+    for line in lines:
+        item = json.loads(line)
+        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
+        # get prompt audio
+        if "prompt_audio_path" in item:
+            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"
+            with torch.no_grad():
+                pmt_wav, vocal_wav, bgm_wav = separator.run(item['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
+            vocal_wav = vocal_wav
+            bgm_wav = bgm_wav
+            with torch.no_grad():
+                pmt_wav, _ = audio_tokenizer.encode(pmt_wav.cuda())
+            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"]]))]
+            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"] = f"{item['idx']}"
+        item["wav_path"] = target_wav_name
+        new_items.append(item)
+
+    del audio_tokenizer
+    del separator
+
+    torch.cuda.empty_cache()
+
+    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
+    
+    if seperate_tokenizer is not None:
+        seperate_tokenizer = seperate_tokenizer.eval().cuda()
+
+    for item in new_items:
+        if "prompt_audio_path" in item:
+            with torch.no_grad():
+                vocal_wav, bgm_wav = seperate_tokenizer.encode(item['vocal_wav'].cuda(), item['bgm_wav'].cuda())
+            item['vocal_wav'] = vocal_wav
+            item['bgm_wav'] = bgm_wav
+
+    torch.cuda.empty_cache()
+    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.eval()
+    audiolm = audiolm.cuda().to(torch.float16)
+
+    model = CodecLM(name = "tmp",
+        lm = audiolm,
+        audiotokenizer = None,
+        max_duration = max_duration,
+        seperate_tokenizer = seperate_tokenizer,
+    )
+
     cfg_coef = 1.5 #25
     temp = 0.9
     top_k = 50
@@ -104,21 +214,135 @@ if __name__ == "__main__":
     os.makedirs(save_dir + "/audios", exist_ok=True)
     os.makedirs(save_dir + "/jsonl", exist_ok=True)
 
+    for item in new_items:
+        lyric = item["gt_lyric"]
+        descriptions = item["descriptions"] if "descriptions" in item else None
+        pmt_wav = item['pmt_wav']
+        vocal_wav = item['vocal_wav']
+        bgm_wav = item['bgm_wav']
+        melody_is_wav = item['melody_is_wav']
+        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
+
+
+        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,
+        }
+        start_time = time.time()
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            with torch.no_grad():
+                tokens = model.generate(**generate_inp, return_tokens=True)
+        mid_time = time.time()
+
+        with torch.no_grad():
+            if 'raw_pmt_wav' in item:
+                if gen_type == 'separate':
+                    wav_seperate = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True, gen_type='mixed')
+                    wav_vocal = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True, gen_type='vocal')
+                    wav_bgm = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True, gen_type='bgm')
+                elif gen_type == 'mixed':
+                    wav_seperate = model.generate_audio(tokens, item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'],chunked=True, gen_type=gen_type)
+                else:
+                    wav_seperate = model.generate_audio(tokens,chunked=True, gen_type=gen_type)
+                del item['raw_pmt_wav']
+                del item['raw_vocal_wav']
+                del item['raw_bgm_wav']
+            else:
+                if gen_type == 'separate':
+                    wav_vocal = model.generate_audio(tokens, chunked=True, gen_type='vocal')
+                    wav_bgm = model.generate_audio(tokens, chunked=True, gen_type='bgm')
+                    wav_seperate = model.generate_audio(tokens, chunked=True, gen_type='mixed')
+                else:
+                    wav_seperate = model.generate_audio(tokens, chunked=True, gen_type=gen_type)
+        del item['pmt_wav']
+        del item['vocal_wav']
+        del item['bgm_wav']
+        del item['melody_is_wav']
+        end_time = time.time()
+        if gen_type == 'separate':
+            torchaudio.save(target_wav_name.replace('.flac', '_vocal.flac'), wav_vocal[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(target_wav_name.replace('.flac', '_bgm.flac'), wav_bgm[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
+        else:
+            torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
+
+        print(f"process{item['idx']}, lm cost {mid_time - start_time}s, diffusion cost {end_time - mid_time}")
+        item["idx"] = f"{item['idx']}"
+        item["wav_path"] = target_wav_name
+    
+    src_jsonl_name = os.path.split(input_jsonl)[-1]
+    with open(f"{save_dir}/jsonl/{src_jsonl_name}.jsonl", "w", encoding='utf-8') as fw:
+        for item in new_items:
+            fw.writelines(json.dumps(item, ensure_ascii=False)+"\n")
+
+def generate_lowmem(args):
+    ckpt_path = args.ckpt_path
+    input_jsonl = args.input_jsonl
+    save_dir = args.save_dir
+    cfg_path = os.path.join(ckpt_path, 'config.yaml')
+    ckpt_path = os.path.join(ckpt_path, 'model.pt')
+    cfg = OmegaConf.load(cfg_path)
+    cfg.lm.use_flash_attn_2 = args.use_flash_attn
+    print(f"use_flash_attn: {args.use_flash_attn}")
+    cfg.mode = 'inference'
+    max_duration = cfg.max_dur
+    gen_type = args.generate_type
+    chunk_size = 128
+    use_audio_tokenizer = False
     with open(input_jsonl, "r") as fp:
         lines = fp.readlines()
-
+    for line in lines:
+        item = json.loads(line)
+        if "prompt_audio_path" in item:
+            use_audio_tokenizer = True
+            break
+    if use_audio_tokenizer:
+        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]
     new_items = []
     for line in lines:
         item = json.loads(line)
         target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
-        lyric = item["gt_lyric"]
-        descriptions = item["descriptions"] if "descriptions" in item else None
         # get prompt audio
         if "prompt_audio_path" in item:
             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(item['prompt_audio_path'])
-            melody_is_wav = True
+            with torch.no_grad():
+                pmt_wav, vocal_wav, bgm_wav = separator.run(item['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)
+            with torch.no_grad():
+                pmt_wav, _ = audio_tokenizer.encode(pmt_wav.cuda())
+            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": 
@@ -134,6 +358,86 @@ if __name__ == "__main__":
             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"] = f"{item['idx']}"
+        item["wav_path"] = target_wav_name
+        new_items.append(item)
+
+    if use_audio_tokenizer:
+        del audio_tokenizer
+        del separator
+
+    torch.cuda.empty_cache()
+    
+    if "audio_tokenizer_checkpoint_sep" in cfg.keys() and use_audio_tokenizer:
+        seperate_tokenizer = builders.get_audio_tokenizer_model(cfg.audio_tokenizer_checkpoint_sep, cfg) 
+    else:
+        seperate_tokenizer = None
+    
+    if seperate_tokenizer is not None:
+        seperate_tokenizer = seperate_tokenizer.eval().cuda()
+
+    for item in new_items:
+        if "prompt_audio_path" in item:
+            with torch.no_grad():
+                vocal_wav, bgm_wav = seperate_tokenizer.encode(item['vocal_wav'].cuda(), item['bgm_wav'].cuda())
+            item['vocal_wav'] = vocal_wav
+            item['bgm_wav'] = bgm_wav
+
+    if use_audio_tokenizer:
+        del seperate_tokenizer
+
+    torch.cuda.empty_cache()
+
+    # Define model or load pretrained model
+    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.eval()
+
+    offload_audiolm = True if 'offload' in cfg.keys() and 'audiolm' in cfg.offload else False
+    if offload_audiolm:
+        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()))
+    else:
+        audiolm = audiolm.cuda().to(torch.float16)
+
+    model = CodecLM(name = "tmp",
+        lm = audiolm,
+        audiotokenizer = None,
+        max_duration = max_duration,
+        seperate_tokenizer = None,
+    )
+    
+    cfg_coef = 1.5 #25
+    temp = 0.9
+    top_k = 50
+    top_p = 0.0
+    record_tokens = True
+    record_window = 50
+    
+
+    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)
+    os.makedirs(save_dir, exist_ok=True)
+    os.makedirs(save_dir + "/audios", exist_ok=True)
+    os.makedirs(save_dir + "/jsonl", exist_ok=True)
+
+    
+    for item in new_items:
+        lyric = item["gt_lyric"]
+        descriptions = item["descriptions"] if "descriptions" in item else None
+        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("  ", " ")],
@@ -143,25 +447,119 @@ if __name__ == "__main__":
             'bgm_wavs': bgm_wav,
             'melody_is_wav': melody_is_wav,
         }
-        start_time = time.time()
         with torch.autocast(device_type="cuda", dtype=torch.float16):
-            tokens = model.generate(**generate_inp, return_tokens=True)
-        mid_time = time.time()
+            with torch.no_grad():
+                tokens = model.generate(**generate_inp, return_tokens=True)
+                if offload_audiolm:
+                    offload_profiler.reset_empty_cache_mem_line()
+        item['tokens'] = tokens
+    if offload_audiolm:
+        offload_profiler.stop()
+        del offload_profiler
+        del audiolm_offload_param
+    del model
+    audiolm = audiolm.cpu()
+    del audiolm
+    del checkpoint
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    seperate_tokenizer = builders.get_audio_tokenizer_model_cpu(cfg.audio_tokenizer_checkpoint_sep, cfg)
+    device = "cuda:0"
+    seperate_tokenizer.model.device = device
+    seperate_tokenizer.model.vae = seperate_tokenizer.model.vae.to(device)
+    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
+    if offload_wav_tokenizer_diffusion:
+        sep_offload_param = OffloadParamParse.parse_config(seperate_tokenizer, cfg.offload.wav_tokenizer_diffusion)
+        sep_offload_param.show()
+        sep_offload_profiler = OffloadProfiler(device_index=0, **(sep_offload_param.init_param_dict()))
+        sep_offload_profiler.offload_layer(**(sep_offload_param.offload_layer_param_dict()))
+        sep_offload_profiler.clean_cache_wrapper(**(sep_offload_param.clean_cache_param_dict()))
+    else:
+        seperate_tokenizer.model.model = seperate_tokenizer.model.model.to(device)
+
+    model = CodecLM(name = "tmp",
+        lm = None,
+        audiotokenizer = None,
+        max_duration = max_duration,
+        seperate_tokenizer = seperate_tokenizer,
+    )
 
+    for item in new_items:
         with torch.no_grad():
-            if melody_is_wav:   
-                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type)
+            if 'raw_pmt_wav' in item:
+                if gen_type == 'separate':
+                    wav_seperate = model.generate_audio(item['tokens'], item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'],chunked=True, gen_type='mixed')
+                    wav_vocal = model.generate_audio(item['tokens'],chunked=True, gen_type='vocal')
+                    wav_bgm = model.generate_audio(item['tokens'], chunked=True, gen_type='bgm')
+                elif gen_type == 'mixed':
+                    wav_seperate = model.generate_audio(item['tokens'], item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'],chunked=True, gen_type=gen_type)
+                else:
+                    wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type=gen_type)
+                del item['raw_pmt_wav']
+                del item['raw_vocal_wav']
+                del item['raw_bgm_wav']
             else:
-                wav_seperate = model.generate_audio(tokens, gen_type=gen_type)
-        end_time = time.time()
-        torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
-        print(f"process{item['idx']} {gen_type}, lm cost {mid_time - start_time}s, diffusion cost {end_time - mid_time}")
-
-        item["idx"] = f"{item['idx']}"
-        item["wav_path"] = target_wav_name
-        new_items.append(item)
+                if gen_type == '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=gen_type)
+        if gen_type == 'separate':
+            torchaudio.save(item['wav_path'].replace('.flac', '_vocal.flac'), wav_vocal[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(item['wav_path'].replace('.flac', '_bgm.flac'), wav_bgm[0].cpu().float(), cfg.sample_rate)
+            torchaudio.save(item['wav_path'], wav_seperate[0].cpu().float(), cfg.sample_rate)
+        else:
+            torchaudio.save(item['wav_path'], wav_seperate[0].cpu().float(), cfg.sample_rate)
+        del item['tokens']
+        del item['pmt_wav']
+        del item['vocal_wav']
+        del item['bgm_wav']
+        del item['melody_is_wav']
+        if offload_wav_tokenizer_diffusion:
+            sep_offload_profiler.reset_empty_cache_mem_line()
     
+    if offload_wav_tokenizer_diffusion:
+        sep_offload_profiler.stop()
+    torch.cuda.empty_cache()
     src_jsonl_name = os.path.split(input_jsonl)[-1]
     with open(f"{save_dir}/jsonl/{src_jsonl_name}.jsonl", "w", encoding='utf-8') as fw:
         for item in new_items:
             fw.writelines(json.dumps(item, ensure_ascii=False)+"\n")
+
+
+if __name__ == "__main__":
+    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])
+    OmegaConf.register_new_resolver("get_fname", lambda: os.path.splitext(os.path.basename(sys.argv[1]))[0])
+    OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(x)))
+    np.random.seed(int(time.time()))
+    # 解析命令行参数
+    args = parse_args()
+    if torch.cuda.is_available():
+        device = torch.cuda.current_device()
+        reserved = torch.cuda.memory_reserved(device)
+        total = torch.cuda.get_device_properties(device).total_memory
+        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':
+            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)
+
+    else:
+        print("CUDA is not available")
+        exit()
+    

generate.sh

@@ -7,5 +7,66 @@ export PYTHONPATH="$(pwd)/codeclm/tokenizer/":"$(pwd)":"$(pwd)/codeclm/tokenizer
 CKPT_PATH=$1
 JSONL=$2
 SAVE_DIR=$3
-GEN_TYEP=$4
-python3 generate.py $CKPT_PATH $JSONL $SAVE_DIR $GEN_TYEP
+USE_FLASH_ATTN="True"
+LOW_MEM="False"
+GENERATE_TYPE="mixed"
+for arg in "$@"; do
+    if [[ $arg == "--not_use_flash_attn" ]]; then
+        USE_FLASH_ATTN="False"
+    fi
+done
+for arg in "$@"; do
+    if [[ $arg == "--low_mem" ]]; then
+        LOW_MEM="True"
+    fi
+done
+for arg in "$@"; do
+    if [[ $arg == "--separate" ]]; then
+        GENERATE_TYPE="separate"
+    fi
+done
+for arg in "$@"; do
+    if [[ $arg == "--bgm" ]]; then
+        GENERATE_TYPE="bgm"
+    fi
+done
+for arg in "$@"; do
+    if [[ $arg == "--vocal" ]]; then
+        GENERATE_TYPE="vocal"
+    fi
+done
+
+
+if [ "$USE_FLASH_ATTN" == "True" ] && [ "$LOW_MEM" == "True" ]; then
+    echo "Use Flash Attention + Low Memory Mode"
+    python3 generate.py \
+        --ckpt_path $CKPT_PATH \
+        --input_jsonl $JSONL \
+        --save_dir $SAVE_DIR \
+        --generate_type $GENERATE_TYPE \
+        --use_flash_attn \
+        --low_mem 
+elif [ "$USE_FLASH_ATTN" == "True" ] && [ "$LOW_MEM" == "False" ]; then
+    echo "Use Flash Attention + Auto Memory Mode"
+    python3 generate.py \
+        --ckpt_path $CKPT_PATH \
+        --input_jsonl $JSONL \
+        --save_dir $SAVE_DIR \
+        --generate_type $GENERATE_TYPE \
+        --use_flash_attn 
+elif [ "$USE_FLASH_ATTN" == "False" ] && [ "$LOW_MEM" == "False" ]; then
+    echo "Not Use Flash Attention + Auto Memory Mode"
+    python3 generate.py \
+        --ckpt_path $CKPT_PATH \
+        --input_jsonl $JSONL \
+        --generate_type $GENERATE_TYPE \
+        --save_dir $SAVE_DIR 
+elif [ "$USE_FLASH_ATTN" == "False" ] && [ "$LOW_MEM" == "True" ]; then
+    echo "Not Use Flash Attention + Low Memory Mode"
+    python3 generate.py \
+        --ckpt_path $CKPT_PATH \
+        --input_jsonl $JSONL \
+        --save_dir $SAVE_DIR \
+        --generate_type $GENERATE_TYPE \
+        --low_mem 
+fi

generate_lowmem.py

@@ -1,241 +0,0 @@
-import sys
-import os
-
-import time
-import json
-import torch
-import torchaudio
-import numpy as np
-from omegaconf import OmegaConf
-from codeclm.models import builders
-
-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
-
-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
-
-
-
-if __name__ == "__main__":
-    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])
-    OmegaConf.register_new_resolver("get_fname", lambda: os.path.splitext(os.path.basename(sys.argv[1]))[0])
-    OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(x)))
-    np.random.seed(int(time.time()))    
-    ckpt_path = sys.argv[1]
-    input_jsonl = sys.argv[2]
-    save_dir = sys.argv[3]
-    gen_type = sys.argv[4] if len(sys.argv) > 4 else "all"
-    cfg_path = os.path.join(ckpt_path, 'config.yaml')
-    ckpt_path = os.path.join(ckpt_path, 'model.pt')
-    cfg = OmegaConf.load(cfg_path)
-    cfg.mode = 'inference'
-    max_duration = cfg.max_dur
-    
-    separator = Separator()
-    auto_prompt = torch.load('ckpt/prompt.pt')
-    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 = [item for sublist in auto_prompt.values() for item in sublist]
-    with open(input_jsonl, "r") as fp:
-        lines = fp.readlines()
-    new_items = []
-    for line in lines:
-        item = json.loads(line)
-        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
-        # get prompt audio
-        if "prompt_audio_path" in item:
-            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(item['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" 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"]]))]
-            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"] = f"{item['idx']}"
-        item["wav_path"] = target_wav_name
-        new_items.append(item)
-
-    del audio_tokenizer
-    del seperate_tokenizer
-    del separator
-
-    # 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
-    model.lm = model.lm.cuda().to(torch.float16)
-    
-    cfg_coef = 1.5 #25
-    temp = 0.9
-    top_k = 50
-    top_p = 0.0
-    record_tokens = True
-    record_window = 50
-
-    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)
-    os.makedirs(save_dir, exist_ok=True)
-    os.makedirs(save_dir + "/audios", exist_ok=True)
-    os.makedirs(save_dir + "/jsonl", exist_ok=True)
-
-    
-    for item in new_items:
-        lyric = item["gt_lyric"]
-        descriptions = item["descriptions"] if "descriptions" in item else None
-        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
-    
-    del model
-    torch.cuda.empty_cache()
-
-
-    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,
-    )
-    for item in new_items:
-        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, gen_type=gen_type)
-                del item['raw_pmt_wav']
-                del item['raw_vocal_wav']
-                del item['raw_bgm_wav']
-            else:
-                wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type=gen_type)
-        torchaudio.save(item['wav_path'], wav_seperate[0].cpu().float(), cfg.sample_rate)
-        del item['tokens']
-        del item['pmt_wav']
-        del item['vocal_wav']
-        del item['bgm_wav']
-        del item['melody_is_wav']
-        
-    torch.cuda.empty_cache()
-    src_jsonl_name = os.path.split(input_jsonl)[-1]
-    with open(f"{save_dir}/jsonl/{src_jsonl_name}.jsonl", "w", encoding='utf-8') as fw:
-        for item in new_items:
-            fw.writelines(json.dumps(item, ensure_ascii=False)+"\n")

generate_lowmem.sh

@@ -1,11 +0,0 @@
-export USER=root
-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
-
-CKPT_PATH=$1
-JSONL=$2
-SAVE_DIR=$3
-GEN_TYEP=$4
-python3 generate_lowmem.py $CKPT_PATH $JSONL $SAVE_DIR $GEN_TYEP

tools/gradio/app.py

@@ -49,7 +49,7 @@ with open(op.join(APP_DIR, 'conf/vocab.yaml'), 'r', encoding='utf-8') as file:
     STRUCTS = yaml.safe_load(file)
 
 
-def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_coef=None, temperature=None, top_k=None, gen_type="all", progress=gr.Progress(track_tqdm=True)):
+def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_coef=None, temperature=None, top_k=None, gen_type="mixed", progress=gr.Progress(track_tqdm=True)):
     global MODEL
     global STRUCTS
     params = {'cfg_coef':cfg_coef, 'temperature':temperature, 'top_k':top_k}
@@ -240,4 +240,3 @@ lyrics
 # 启动应用
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=8081)
-

tools/gradio/levo_inference.py

@@ -62,7 +62,7 @@ class LeVoInference(torch.nn.Module):
 
         self.model.set_generation_params(**self.default_params)
 
-    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "all", params = dict()):
+    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "mixed", params = dict()):
         params = {**self.default_params, **params}
         self.model.set_generation_params(**params)
 

tools/gradio/levo_inference_lowmem.py

@@ -1,4 +1,5 @@
 import os
+import gc
 import sys
 
 import torch
@@ -12,6 +13,7 @@ from codeclm.models import CodecLM
 from codeclm.models import builders
 
 from separator import Separator
+from codeclm.utils.offload_profiler import OffloadProfiler, OffloadParamParse
 
 
 class LeVoInference(torch.nn.Module):
@@ -40,24 +42,28 @@ class LeVoInference(torch.nn.Module):
         )
 
 
-    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "all", params = dict()):
+    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "mixed", params = dict()):
         if prompt_audio_path is not None and os.path.exists(prompt_audio_path):
             separator = Separator()
             audio_tokenizer = builders.get_audio_tokenizer_model(self.cfg.audio_tokenizer_checkpoint, self.cfg)
             audio_tokenizer = audio_tokenizer.eval().cuda()
-            seperate_tokenizer = builders.get_audio_tokenizer_model(self.cfg.audio_tokenizer_checkpoint_sep, self.cfg)
-            seperate_tokenizer = seperate_tokenizer.eval().cuda()
             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()
-            pmt_wav, _ = audio_tokenizer.encode(pmt_wav)
-            vocal_wav, bgm_wav = seperate_tokenizer.encode(vocal_wav, bgm_wav)
-            melody_is_wav = False
-            melody_is_wav = False
+            with torch.no_grad():
+                pmt_wav, _ = audio_tokenizer.encode(pmt_wav)
             del audio_tokenizer
-            del seperate_tokenizer
             del separator
+            torch.cuda.empty_cache()
+
+            seperate_tokenizer = builders.get_audio_tokenizer_model(self.cfg.audio_tokenizer_checkpoint_sep, self.cfg)
+            seperate_tokenizer = seperate_tokenizer.eval().cuda()
+            with torch.no_grad():
+                vocal_wav, bgm_wav = seperate_tokenizer.encode(vocal_wav, bgm_wav)
+            del seperate_tokenizer
+            melody_is_wav = False
+            torch.cuda.empty_cache()
         elif genre is not None and auto_prompt_path is not None:
             auto_prompt = torch.load(auto_prompt_path)
             merge_prompt = [item for sublist in auto_prompt.values() for item in sublist]
@@ -75,17 +81,28 @@ class LeVoInference(torch.nn.Module):
             bgm_wav = None
             melody_is_wav = True
 
-        model_light = CodecLM_PL(self.cfg, self.pt_path)
-        model_light = model_light.eval()
-        model_light.audiolm.cfg = self.cfg
+        audiolm = builders.get_lm_model(self.cfg)
+        checkpoint = torch.load(self.pt_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()
+
+        offload_audiolm = True if 'offload' in self.cfg.keys() and 'audiolm' in self.cfg.offload else False
+        if offload_audiolm:
+            audiolm_offload_param = OffloadParamParse.parse_config(audiolm, self.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()))
+        else:
+            audiolm = audiolm.cuda().to(torch.float16)
+
         model = CodecLM(name = "tmp",
-            lm = model_light.audiolm,
+            lm = audiolm,
             audiotokenizer = None,
             max_duration = self.max_duration,
             seperate_tokenizer = None,
         )
-        del model_light
-        model.lm = model.lm.cuda().to(torch.float16)
         params = {**self.default_params, **params}
         model.set_generation_params(**params)
 
@@ -99,28 +116,53 @@ class LeVoInference(torch.nn.Module):
         }
 
         with torch.autocast(device_type="cuda", dtype=torch.float16):
-            tokens = model.generate(**generate_inp, return_tokens=True)
-
+            with torch.no_grad():
+                tokens = model.generate(**generate_inp, return_tokens=True)
+                if offload_audiolm:
+                    offload_profiler.reset_empty_cache_mem_line()
+        offload_profiler.stop()
+        del offload_profiler
+        del audiolm_offload_param
         del model
+        audiolm = audiolm.cpu()
+        del audiolm
+        del checkpoint
+        gc.collect()
         torch.cuda.empty_cache()
 
-        seperate_tokenizer = builders.get_audio_tokenizer_model(self.cfg.audio_tokenizer_checkpoint_sep, self.cfg)
-        seperate_tokenizer = seperate_tokenizer.eval().cuda()
+        seperate_tokenizer = builders.get_audio_tokenizer_model_cpu(self.cfg.audio_tokenizer_checkpoint_sep, self.cfg)
+        device = "cuda:0"
+        seperate_tokenizer.model.device = device
+        seperate_tokenizer.model.vae = seperate_tokenizer.model.vae.to(device)
+        seperate_tokenizer.model.model.device = torch.device(device)
+        seperate_tokenizer = seperate_tokenizer.eval()
+
+        offload_wav_tokenizer_diffusion =  True if 'offload' in self.cfg.keys() and 'wav_tokenizer_diffusion' in self.cfg.offload else False
+        if offload_wav_tokenizer_diffusion:
+            sep_offload_param = OffloadParamParse.parse_config(seperate_tokenizer, self.cfg.offload.wav_tokenizer_diffusion)
+            sep_offload_param.show()
+            sep_offload_profiler = OffloadProfiler(device_index=0, **(sep_offload_param.init_param_dict()))
+            sep_offload_profiler.offload_layer(**(sep_offload_param.offload_layer_param_dict()))
+            sep_offload_profiler.clean_cache_wrapper(**(sep_offload_param.clean_cache_param_dict()))
+        else:
+            seperate_tokenizer.model.model = seperate_tokenizer.model.model.to(device)
+
         model = CodecLM(name = "tmp",
             lm = None,
             audiotokenizer = None,
             max_duration = self.max_duration,
             seperate_tokenizer = seperate_tokenizer,
         )
-            
+
         with torch.no_grad():
             if melody_is_wav:
-                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type)
+                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type, chunked=True)
             else:
-                wav_seperate = model.generate_audio(tokens, gen_type=gen_type)
+                wav_seperate = model.generate_audio(tokens, gen_type=gen_type, chunked=True)
 
-        del seperate_tokenizer
-        del model
+        if offload_wav_tokenizer_diffusion:
+            sep_offload_profiler.reset_empty_cache_mem_line()
+            sep_offload_profiler.stop()
         torch.cuda.empty_cache()
 
         return wav_seperate[0]