init

SongGeneration/codeclm/tokenizer/Flow1dVAE/generate_1rvq.py

@@ -15,15 +15,14 @@ from safetensors.torch import load_file
 class Tango:
     def __init__(self, \
         model_path, \
-        vae_config="",
-        vae_model="",
+        vae_config,
+        vae_model,
         layer_num=6, \
-        device="cuda:0"):
+        device="cuda"):
         
         self.sample_rate = 48000
         scheduler_name = "custom_nodes/ComfyUI_SongGeneration/SongGeneration/codeclm/tokenizer/Flow1dVAE/our_MERT_BESTRQ/mert_fairseq/configs/scheduler/stable_diffusion_2.1_largenoise_sample.json"
         self.device = device
-
         self.vae = get_model(vae_config, vae_model)
         self.vae = self.vae.to(device)
         self.vae=self.vae.eval()
@@ -43,7 +42,7 @@ class Tango:
             main_weights = torch.load(model_path, map_location=device)
         self.model.load_state_dict(main_weights, strict=False)
         #print ("Successfully loaded checkpoint from:", model_path)
-        
+        del main_weights
         self.model.eval()
         self.model.init_device_dtype(torch.device(device), torch.float32)
         #print("scaling factor: ", self.model.normfeat.std)
@@ -92,6 +91,7 @@ class Tango:
     #         output = torch.cat([saved_samples.detach().cpu(),audio[0].detach().cpu()],0)
     #     return output
 
+
     @torch.no_grad()
     @torch.autocast(device_type="cuda", dtype=torch.float32)
     def sound2code(self, orig_samples, batch_size=3):

SongGeneration/codeclm/tokenizer/Flow1dVAE/generate_septoken.py

@@ -73,15 +73,14 @@ class Tango:
         vae_model,
         layer_vocal=7,\
         layer_bgm=3,\
-        device="cuda:0"):
+        device="cuda"
+        ):
         
         self.sample_rate = 48000
         scheduler_name = "configs/scheduler/stable_diffusion_2.1_largenoise_sample.json"
         self.device = device
-
         self.vae = get_model(vae_config, vae_model)
-        self.vae = self.vae.to(device)
-        self.vae=self.vae.eval()
+        self.vae=self.vae.to(device).eval()
         self.layer_vocal=layer_vocal
         self.layer_bgm=layer_bgm
 
@@ -99,7 +98,7 @@ class Tango:
             main_weights = torch.load(model_path, map_location=device)
         self.model.load_state_dict(main_weights, strict=False)
         #print ("Successfully loaded checkpoint from:", model_path)
-        
+        del main_weights
         self.model.eval()
         self.model.init_device_dtype(torch.device(device), torch.float32)
         #print("scaling factor: ", self.model.normfeat.std)
@@ -110,7 +109,7 @@ class Tango:
         #     scheduler_name, subfolder="scheduler")
         #print("Successfully loaded inference scheduler from {}".format(scheduler_name))
 
-
+         
     @torch.no_grad()
     @torch.autocast(device_type="cuda", dtype=torch.float32)
     def sound2code(self, orig_vocal, orig_bgm, batch_size=8):

SongGeneration/codeclm/tokenizer/Flow1dVAE/model_septoken.py

@@ -617,8 +617,6 @@ class PromptCondAudioDiffusion(nn.Module):
         if(scenario=='other_seg'):
             latent_masks[:,0:incontext_length] = 1
 
-        
-
         quantized_bestrq_emb = (latent_masks > 0.5).unsqueeze(-1) * quantized_bestrq_emb \
             + (latent_masks < 0.5).unsqueeze(-1) * self.zero_cond_embedding1.reshape(1,1,1024)
         quantized_bestrq_emb_bgm = (latent_masks > 0.5).unsqueeze(-1) * quantized_bestrq_emb_bgm \

SongGeneration/codeclm/tokenizer/Flow1dVAE/models_gpt/models/gpt2_rope2_time_new_correct_mask_noncasual_reflow.py

@@ -42,6 +42,7 @@ try:
     from transformers.models.gpt2.modeling_gpt2 import GPT2SequenceSummary as SequenceSummary
 except:
     from transformers.modeling_utils import SequenceSummary
+    
 from transformers.pytorch_utils import Conv1D, find_pruneable_heads_and_indices, prune_conv1d_layer
 from transformers.utils import (
     ModelOutput,

SongGeneration/codeclm/tokenizer/Flow1dVAE/tools/get_1dvae_large.py

@@ -12,4 +12,5 @@ def get_model(model_config, path):
     state_dict = torch.load(path, map_location='cpu')
     model = create_autoencoder_from_config(model_config)
     model.load_state_dict(state_dict['state_dict'], strict=False)
+    del state_dict
     return model

SongGeneration/codeclm/tokenizer/audio_tokenizer.py

@@ -75,9 +75,9 @@ class AudioTokenizer(ABC, nn.Module):
     @staticmethod
     def get_pretrained(
             name: str, 
-            vae_config: str,
-            vae_model: str,
-            device: tp.Union[torch.device, str] = 'cpu', 
+            vae_config,
+            vae_model,
+            device: tp.Union[torch.device, str] = 'cuda', 
             mode='extract'
             ) -> 'AudioTokenizer':
         """Instantiate a AudioTokenizer model from a given pretrained model.
@@ -91,11 +91,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)
         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)
         else:
             raise NotImplementedError("{} is not implemented in models/audio_tokenizer.py".format(
                 name))
@@ -106,19 +106,21 @@ class Flow1dVAE1rvq(AudioTokenizer):
     def __init__(
         self, 
         model_type: str = "model_2_fixed.safetensors",
-        vae_config: str = "",
-        vae_model: str = "",
+        vae_config:str = "",
+        vae_model:str = "",
         ):
         super().__init__()
 
         from .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='cuda')
         #print ("Successfully loaded checkpoint from:", model_path)
 
             
         self.n_quantizers = 1
 
+
+
     def forward(self, x: torch.Tensor) :
         # We don't support training with this.
         raise NotImplementedError("Forward and training with DAC not supported.")
@@ -181,19 +183,21 @@ class Flow1dVAESeparate(AudioTokenizer):
     def __init__(
         self, 
         model_type: str = "model_2.safetensors",
-        vae_config: str = "",
-        vae_model: str = "",
+        vae_config:str = "",
+        vae_model:str = "",
         ):
         super().__init__()
 
         from .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='cuda')
         #print ("Successfully loaded checkpoint from:", model_path)
 
             
         self.n_quantizers = 1
 
+
+        
     def forward(self, x: torch.Tensor) :
         # We don't support training with this.
         raise NotImplementedError("Forward and training with DAC not supported.")