add auto prompt and interface

app.py

@@ -1,196 +1,231 @@
-import os
 import gradio as gr
 import json
-import numpy as np
 from datetime import datetime
-import os
 import yaml
-import sys
-import librosa
 import time
 import os.path as op
-APP_DIR = op.dirname(op.abspath(__file__))
-
 from download import download_model
+from levo_inference import LeVoInference
+
+
 # 下载模型
+APP_DIR = op.dirname(op.abspath(__file__))
 download_model(APP_DIR)
 print("Successful downloaded model.")
 
-from levo_inference import LeVoInference
-
-MODEL = LeVoInference(op.join(APP_DIR, "conf/infer.yaml"))
+# 模型初始化
+MODEL = LeVoInference(op.join(APP_DIR, "ckpt/songgeneration_base_zn/"))
 
-EXAMPLE_DESC = """female, dark, pop, sad, piano and drums, the bpm is 125."""
 EXAMPLE_LYRICS = """
 [intro-short]
 
 [verse]
-夜晚的街灯闪烁.
-我漫步在熟悉的角落.
-回忆像潮水般涌来.
-你的笑容如此清晰.
-在心头无法抹去.
-那些曾经的甜蜜.
-如今只剩我独自回忆.
-
-[bridge]
-手机屏幕亮起.
-是你发来的消息.
-简单的几个字.
-却让我泪流满面.
-曾经的拥抱温暖.
-如今却变得遥远.
-我多想回到从前.
-重新拥有你的陪伴.
+雪花舞动在无尽的天际
+情缘如同雪花般轻轻逝去
+希望与真挚
+永不磨灭
+你的忧虑
+随风而逝
 
 [chorus]
-回忆的温度还在.
-你却已不在.
-我的心被爱填满.
-却又被思念刺痛.
-R&B的节奏奏响.
-我的心却在流浪.
-没有你的日子.
-我该如何继续向前.
+我怀抱着守护这片梦境
+在这世界中寻找爱与虚幻
+苦辣酸甜
+我们一起品尝
+在雪的光芒中
+紧紧相拥
+
+[inst-short]
+
+[verse]
+雪花再次在风中飘扬
+情愿如同雪花般消失无踪
+希望与真挚
+永不消失
+在痛苦与喧嚣中
+你找到解脱
+
+[chorus]
+我环绕着守护这片梦境
+在这世界中感受爱与虚假
+苦辣酸甜
+我们一起分享
+在白银的光芒中
+我们同在
 
 [outro-short]
 """.strip()
 
-with open('conf/vocab.yaml', 'r', encoding='utf-8') as file:
+with open(op.join(APP_DIR, 'conf/vocab.yaml'), 'r', encoding='utf-8') as file:
     STRUCTS = yaml.safe_load(file)
 
 
 # 模拟歌曲生成函数
-def generate_song(description, lyric, prompt_audio=None, cfg_coef=None, temperature=None, top_k=None, 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, progress=gr.Progress(track_tqdm=True)):
     global MODEL
     global STRUCTS
     params = {'cfg_coef':cfg_coef, 'temperature':temperature, 'top_k':top_k}
     params = {k:v for k,v in params.items() if v is not None}
     sample_rate = MODEL.cfg.sample_rate
-
-    # 生成过程
-    print(f"Generating song with description: {description}")
-    print(f"Lyrics provided: {lyric}")
     
     # 适配lyric格式
+    lyric = lyric.replace("[intro]", "[intro-short]").replace("[inst]", "[inst-short]").replace("[outro]", "[outro-short]")
     lyric = lyric.replace("\n\n", " ; ")
     for s in STRUCTS:
         lyric = lyric.replace(f"{s}\n", f"{s} ")
-    lyric = lyric.replace("\n", "")
+    lyric = lyric.replace("\n", ".")
     lyric = lyric.replace(". ; ", " ; ")
     
     # 适配prompt 
     if prompt_audio is not None:
-        print("Using prompt audio for generation")
-    else:
-        prompt_audio = op.join(APP_DIR, 'sample/prompt.wav')
+        genre = None
+        description = None
+    elif description is not None and description != "":
+        genre = None
 
     progress(0.0, "Start Generation")
     start = time.time()
     
-    audio_data = MODEL(lyric, description, prompt_audio, params).cpu().permute(1, 0).float().numpy()
+    audio_data = MODEL(lyric, description, prompt_audio, genre, op.join(APP_DIR, "ckpt/prompt.pt"), params).cpu().permute(1, 0).float().numpy()
 
     end = time.time()
     
     # 创建输入配置的JSON
     input_config = {
-        "description": description,
         "lyric": lyric,
+        "genre": genre,
         "prompt_audio": prompt_audio,
+        "description": description,
         "params": params,
         "inference_duration": end - start,
         "timestamp": datetime.now().isoformat(),
     }
+    print(input_config)
     
     return (sample_rate, audio_data), json.dumps(input_config, indent=2)
 
+
 # 创建Gradio界面
-with gr.Blocks(title="LeVo Demo Space") as demo:
-    gr.Markdown("# 🎵 LeVo Demo Space")
-    gr.Markdown("Demo interface for the LeVo song generation model. Provide a description, lyrics, and optionally an audio prompt to generate a custom song.")
+with gr.Blocks(title="SongGeration Demo Space") as demo:
+    gr.Markdown("# 🎵 SongGeration Demo Space")
+    gr.Markdown("Demo interface for the song generation model. Provide a lyrics, and optionally an audio or text prompt, to generate a custom song.")
     
     with gr.Row():
         with gr.Column():
-            description = gr.Textbox(
-                label="Song Description",
-                placeholder="Describe the style, mood, and characteristics of the song...",
-                lines=1,
-                max_lines=2,
-                value=EXAMPLE_DESC,
-            )
             lyric = gr.Textbox(
                 label="Lyrics",
-                placeholder="Enter the lyrics for the song...",
                 lines=5,
-                max_lines=8,
+                max_lines=15,
                 value=EXAMPLE_LYRICS,
+                info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics. Use [intro] [outro] [inst] to generate instrumental music.",
+                placeholder="""Lyric Format
+'''
+[structure tag]
+lyrics
+
+[structure tag]
+lyrics
+'''
+1. One paragraph represents one section, starting with a structure tag and ending with a blank line
+2. One line represents one lyric line, punctuation is not recommended inside the line
+3. Structure tags can be chosen from the following list
+    - '[verse]'
+    - '[chorus]'
+    - '[bridge]'
+    - '[intro-short]'
+    - '[intro-medium]'
+    - '[intro-long]'
+    - '[outro-short]'
+    - '[outro-medium]'
+    - '[outro-long]'
+    - '[inst-short]'
+    - '[inst-medium]'
+    - '[inst-long]'
+    - '[silence]'
+"""
             )
 
             with gr.Tabs(elem_id="extra-tabs"):
+                with gr.Tab("Genre Select"):
+                    genre = gr.Radio(
+                        choices=["Auto", "Pop", "R&B", "Dance", "Jazz", "Folk", "Rock", "Chinese Style", "Chinese Tradition", "Metal", "Reggae", "Chinese Opera"],
+                        label="Genre Select(Optional)",
+                        value="Auto",  # 默认选中第一个
+                        interactive=True,
+                        elem_id="single-select-radio"  # 便于自定义样式
+                    )
                 with gr.Tab("Audio Prompt"):
                     prompt_audio = gr.Audio(
                         label="Prompt Audio (Optional)",
                         type="filepath",
                         elem_id="audio-prompt"
                     )
-                with gr.Tab("Advanced Config"):
-                    cfg_coef = gr.Slider(
-                        label="CFG Coefficient",
-                        minimum=0.1,
-                        maximum=3.0,
-                        step=0.1,
-                        value=1.5,
-                        interactive=True,
-                        elem_id="cfg-coef",
-                    )
-                    temperature = gr.Slider(
-                        label="Temperature",
-                        minimum=0.1,
-                        maximum=2.0,
-                        step=0.1,
-                        value=1.0,
-                        interactive=True,
-                        elem_id="temperature",
-                    )
-                    top_k = gr.Slider(
-                        label="Top-K",
-                        minimum=1,
-                        maximum=100,
-                        step=1,
-                        value=50,
-                        interactive=True,
-                        elem_id="top_k",
+                with gr.Tab("Text Prompt"):
+                    description = gr.Textbox(
+                        label="Song Description (Optional)",
+                        info="Describe the gender, timbre, genre, emotion, instrument and bpm of the song",
+                        placeholder="female, dark, pop, sad, piano and drums, the bpm is 125.",
+                        lines=1,
+                        max_lines=2
                     )
+
+            with gr.Accordion("Advanced Config", open=False):
+                cfg_coef = gr.Slider(
+                    label="CFG Coefficient",
+                    minimum=0.1,
+                    maximum=3.0,
+                    step=0.1,
+                    value=1.5,
+                    interactive=True,
+                    elem_id="cfg-coef",
+                )
+                temperature = gr.Slider(
+                    label="Temperature",
+                    minimum=0.1,
+                    maximum=2.0,
+                    step=0.1,
+                    value=0.9,
+                    interactive=True,
+                    elem_id="temperature",
+                )
+                top_k = gr.Slider(
+                    label="Top-K",
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=50,
+                    interactive=True,
+                    elem_id="top_k",
+                )
             generate_btn = gr.Button("Generate Song", variant="primary")
         
         with gr.Column():
             output_audio = gr.Audio(label="Generated Song", type="numpy")
             output_json = gr.JSON(label="Input Configuration")
     
-    # 示例按钮
-    examples = gr.Examples(
-        examples=[
-            ["An uplifting pop song with catchy melodies"],
-            ["Melancholic piano ballad"],
-        ],
-        inputs=[description],
-        label="Description examples"
-    )
-
-    examples = gr.Examples(
-        examples=[
-            ["Shine bright like the stars above\nYou're the one that I'm dreaming of"],
-            ["The rain keeps falling on my window pane\nReminding me of love that's gone away"],
-        ],
-        inputs=[lyric],
-        label="Lyrics examples"
-    )
+        # # 示例按钮
+        # examples = gr.Examples(
+        #     examples=[
+        #         ["male, bright, rock, happy, electric guitar and drums, the bpm is 150."],
+        #         ["female, warm, jazz, romantic, synthesizer and piano, the bpm is 100."]
+        #     ],
+        #     inputs=[description],
+        #     label="Text Prompt examples"
+        # )
+
+        # examples = gr.Examples(
+        #     examples=[
+        #     "[intro-medium]\n\n[verse]\n在这个疯狂的世界里\n谁不渴望一点改变\n在爱情面前\n我们都显得那么不安全\n你紧紧抱着我\n告诉我再靠近一点\n别让这璀璨的夜晚白白浪费\n我那迷茫的眼睛\n看不见未来的路\n在情感消散之前\n我们对爱的渴望永不熄灭\n你给我留下一句誓言\n想知道我们的爱是否能持续到永远\n[chorus]\n\n约定在那最后的夜晚\n不管命运如何摆布\n我们的心是否依然如初\n我会穿上红衬衫\n带着摇滚的激情\n回到我们初遇的地方\n约定在那最后的夜晚\n就算全世界都变了样\n我依然坚守诺言\n铭记这一天\n你永远是我心中的爱恋\n\n[outro-medium]\n",
+        #     "[intro-short]\n\n[verse]\nThrough emerald canyons where fireflies dwell\nCerulean berries kiss morning's first swell\nCrystalline dew crowns each Vitamin Dawn's confection dissolves slowly on me\nAmbrosia breezes through honeycomb vines\nNature's own candy in Fibonacci lines\n[chorus] Blueberry fruit so sweet\n takes you higher\n can't be beat\n In your lungs\n it starts to swell\n You're under its spell\n [verse] Resin of sunlight in candied retreat\nMarmalade moonbeams melt under bare feet\nNectar spirals bloom chloroplast champagne\nPhotosynthesis sings through my veins\nChlorophyll rhythms pulse warm in my blood\nThe forest's green pharmacy floods every bud[chorus] Blueberry fruit so sweet\n takes you higher\n can't be beat\n In your lungs\n it starts to swell\n You're under its spell\n feel the buzz\n ride the wave\n Limey me\n blueberry\n your mind's enslaved\n In the haze\n lose all time\n floating free\n feeling fine\n Blueberry\n fruit so sweet\n takes you higher\n can't be beat\n In your lungs\n it starts to swell\n cry\n You're under its spell\n\n[outro-short]\n",
+        #     ],
+        #     inputs=[lyric],
+        #     label="Lyrics examples",
+        # )
     
     # 生成按钮点击事件
-
     generate_btn.click(
         fn=generate_song,
-        inputs=[description, lyric, prompt_audio, cfg_coef, temperature, top_k],
+        inputs=[lyric, description, prompt_audio, genre, cfg_coef, temperature, top_k],
         outputs=[output_audio, output_json]
     )
     

codeclm/models/builders.py

@@ -16,7 +16,6 @@ from codeclm.modules.conditioners import (
     BaseConditioner,
     QwTokenizerConditioner,
     QwTextConditioner,
-    PhonemeTokenizerConditioner,
     QuantizedEmbeddingConditioner,
     ConditionerProvider,
     ConditionFuser,
@@ -102,11 +101,6 @@ def get_conditioner_provider(output_dim: int, cfg: omegaconf.DictConfig) -> Cond
                 output_dim=output_dim,
                 **model_args
             )
-        elif model_type == 'PhonemeTokenizer':
-            conditioners[str(cond)] = PhonemeTokenizerConditioner(                                                 
-                output_dim=output_dim,
-                **model_args
-            )
         elif model_type == "qt_embedding":
             conditioners[str(cond)] = QuantizedEmbeddingConditioner(
                 dim=output_dim,

codeclm/models/codeclm.py

@@ -208,29 +208,29 @@ class CodecLM:
             elif melody_tokens.shape[-1] < target_melody_token_len:
                 melody_tokens = torch.cat([melody_tokens, torch.full((1,1,target_melody_token_len - melody_tokens.shape[-1]), 16385, device=self.device).long()], dim=-1)
         if self.seperate_tokenizer is not None:
-            if vocal_wavs is not None:
+            if bgm_wavs is None:
+                assert vocal_wavs is None, "vocal_wavs is not None when bgm_wavs is None"
+                bgm_tokens = torch.full((1,1,target_melody_token_len), 16385, device=self.device).long()
+                vocal_tokens = torch.full((1,1,target_melody_token_len), 16385, device=self.device).long()
+            else:
+                assert vocal_wavs is not None, "vocal_wavs is None when bgm_wavs is not None"
                 if type(vocal_wavs) == list:
                     vocal_wavs = torch.stack(vocal_wavs, dim=0)
-                if bgm_wavs is None:
-                    use_bgm = False
-                    bgm_wavs = torch.zeros_like(vocal_wavs)
-                    bgm_wavs[:, 0] = 1.0
-                    bgm_wavs[:, 1:] = torch.randn_like(bgm_wavs[:, 1:])* 0.0003
-                else:
-                    use_bgm = True
-                    if type(bgm_wavs) == list:
-                        bgm_wavs = torch.stack(bgm_wavs, dim=0)
+                if type(bgm_wavs) == list:
+                    bgm_wavs = torch.stack(bgm_wavs, dim=0)
                 vocal_wavs = vocal_wavs.to(self.device)
                 bgm_wavs = bgm_wavs.to(self.device)
-                vocal_tokens, bgm_tokens = self.seperate_tokenizer.encode(vocal_wavs, bgm_wavs)
+                if melody_is_wav:
+                    vocal_tokens, bgm_tokens = self.seperate_tokenizer.encode(vocal_wavs, bgm_wavs)
+                else:
+                    vocal_tokens = vocal_wavs
+                    bgm_tokens = bgm_wavs
                 assert len(vocal_tokens.shape) == len(bgm_tokens.shape) == 3, \
                     f"vocal and bgm tokens should have a shape [B, C, T]! " \
                     f"got vocal len={vocal_tokens.shape}, and bgm len={bgm_tokens.shape}"
                 assert vocal_tokens.shape[-1] == bgm_tokens.shape[-1], \
                     f"vocal and bgm tokens should have the same length! " \
                     f"got vocal len={vocal_tokens.shape[-1]}, and bgm len={bgm_tokens.shape[-1]}"
-                if not use_bgm:
-                    bgm_tokens = torch.full_like(bgm_tokens, 16385)
                 if bgm_tokens.shape[-1] > target_melody_token_len:
                     bgm_tokens = bgm_tokens[...,:target_melody_token_len]
                 elif bgm_tokens.shape[-1] < target_melody_token_len:
@@ -239,10 +239,6 @@ class CodecLM:
                     vocal_tokens = vocal_tokens[...,:target_melody_token_len]
                 elif vocal_tokens.shape[-1] < target_melody_token_len:
                     vocal_tokens = torch.cat([vocal_tokens, torch.full((1,1,target_melody_token_len - vocal_tokens.shape[-1]), 16385, device=self.device).long()], dim=-1)
-            else:
-                bgm_tokens = torch.full((1,1,target_melody_token_len), 16385, device=self.device).long()
-                vocal_tokens = torch.full((1,1,target_melody_token_len), 16385, device=self.device).long()
-        
             melody_tokens = torch.cat([melody_tokens, vocal_tokens, bgm_tokens], dim=1)
         assert melody_tokens.shape[-1] == target_melody_token_len
         audio_qt_embs = melody_tokens.long()

codeclm/models/lm_levo.py

@@ -66,13 +66,17 @@ class LmModel(StreamingModule):
                  intermediate_size: int = 4096,
                  num_heads: int = 8,
                  norm: str = 'layer_norm', norm_first: bool = False,
-                 bias_proj: bool = True,
                  weight_init: tp.Optional[str] = None, depthwise_init: tp.Optional[str] = None,
                  zero_bias_init: bool = False, cfg_dropout: float = 0, cfg_coef: float = 1.0,
                  attribute_dropout: tp.Dict[str, tp.Dict[str, float]] = {}, 
-                 lm_type = 'Llama', 
                  num_layers=16,
+                 max_position_embeddings: int = 8196,
+                 max_position_embeddings_sub: int = 10000,
+                 rope_theta: float = 100000.0,
+                 rope_theta_sub: float = 500000.0,
+                 num_layers_sub: int = 12,
                  cfg = None,
+                 use_flash_attn_2: bool = True,
                  **kwargs):
         super().__init__()
 
@@ -89,8 +93,6 @@ class LmModel(StreamingModule):
         self.cfg = cfg
         self.pattern_provider = pattern_provider
         self.emb = nn.ModuleList([nn.Embedding(input_emb_dim, dim)])
-        # if 'activation' in kwargs:
-        #     kwargs['activation'] = get_activation_fn(kwargs['activation'])
                 
         model_cfg = LlamaConfig(
             hidden_size=dim,
@@ -100,12 +102,10 @@ class LmModel(StreamingModule):
             num_key_value_heads = num_heads,
             vocab_size = self.code_size,
             use_cache=False,
-            max_position_embeddings=8196,
-            _flash_attn_2_enabled=True,
+            max_position_embeddings=max_position_embeddings,
             rms_norm_eps= 1e-5,
-            rope_theta= 100000.0,
-            use_flash_attn_2=True,
-            attn_implementation="flash_attention_2"
+            rope_theta= rope_theta,
+            _flash_attn_2_enabled=use_flash_attn_2,
         )
 
         self.transformer = CausalLM(model_cfg)
@@ -114,23 +114,22 @@ class LmModel(StreamingModule):
             nn.GELU(),
             nn.Linear(dim, dim)
         )
-        self.layer2_emb = nn.ModuleList([nn.Embedding(input_emb_dim, dim) #, lr=emb_lr) 
+        self.layer2_emb = nn.ModuleList([nn.Embedding(input_emb_dim, dim)
                                   for _ in range(self.code_depth)])
         sub_model_cfg = LlamaConfig(
             hidden_size=dim,
             intermediate_size = intermediate_size,
             num_attention_heads = num_heads,
-            num_hidden_layers = 12,
+            num_hidden_layers = num_layers_sub,
             num_key_value_heads = num_heads,
             vocab_size = self.code_size,
             use_cache=False,
-            max_position_embeddings=10000,
+            max_position_embeddings=max_position_embeddings_sub,
             rms_norm_eps= 1e-5,
-            rope_theta= 500000.0,
-            _flash_attn_2_enabled=True,
-            use_flash_attn_2=True,
-            attn_implementation="flash_attention_2"
+            rope_theta= rope_theta_sub,
+            _flash_attn_2_enabled=use_flash_attn_2,
         )
+
         self.transformer2 = CausalLM(sub_model_cfg)
         self.out_norm: tp.Optional[nn.Module] = None
         if norm_first:
@@ -208,15 +207,9 @@ class LmModel(StreamingModule):
                     if descriptions is not None:
                         attr["text"]["type_info"] = descriptions[i]
                 attributes.append(attr)
-            # print("before cfg dropout", attributes)
             attributes = self.cfg_dropout(attributes)   # drop ALL conditions
-            # print("after cfg dropout", attributes)
             attributes = self.att_dropout(attributes)   # selectively drop some attributes (text, wav, or more fine-grained)
-            # print("after attribute dropout", attributes)
-            # attribute to discrete tokenized ids
             tokenized = self.condition_provider.tokenize(attributes)
-            # print("after tokenize", attributes)
-            # discrete tokenized ids to continuous embeddings
             condition_tensors = self.condition_provider(tokenized)
         else:
             conditions = []
@@ -418,6 +411,7 @@ class LmModel(StreamingModule):
         assert start_offset_sequence is not None
         is_end = torch.zeros((B, self.code_depth, 1)).bool().to(device)
         ignore_tokens = audio_qt_embs[0][0]
+        ignore_tokens = ignore_tokens[ignore_tokens < 16384]
         # 5) auto-regressive sampling
         with self.streaming():
             gen_sequence_len = gen_sequence.shape[-1]  # gen_sequence shape is [B, K, S]
@@ -457,7 +451,6 @@ class LmModel(StreamingModule):
                 if torch.all(is_end):
                     gen_sequence = gen_sequence[..., :offset+1]
                     break
-                
                 prev_offset = offset
                 
         # ensure sequence has been entirely filled
@@ -529,7 +522,7 @@ class LmModel(StreamingModule):
                 logits[:, q, :tmp] /= (1.1 ** q_count[:tmp])
 
         # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error.
-        if(ignore_tokens is not None):
+        if(ignore_tokens is not None and len(ignore_tokens) > 0):
             logits[0][0][ignore_tokens.to(torch.int)] = float('-inf')
         if use_sampling and temp > 0.0:
             probs = torch.softmax(logits / temp, dim=-1)

codeclm/modules/conditioners.py

@@ -107,173 +107,6 @@ class TextConditioner(BaseConditioner):
     ...
 
 
-class PhonemeTokenizerConditioner(TextConditioner):
-    def __init__(self, 
-                 output_dim: int, 
-                 vocab_list,
-                 max_len = 600, 
-                 max_sentence_per_structure = 50,
-                 structure_tokens=None,
-                 structure_split_tokens=[','],
-                 sentence_split_tokens=['.'],
-                 mode='sum',
-                 structure_output_dim = 64,
-                 sentence_output_dim = 64,
-                 max_duration = 120,
-                 ): 
-        
-        self.vocab_list = vocab_list
-        self.max_len = max_len
-        self.mode = mode
-        self.max_sentence_per_structure = max_sentence_per_structure
-        voc_size = len(self.vocab_list)
-
-        if structure_tokens is None:
-            structure_tokens = [i for i in vocab_list if len(i) > 1 and i[0] == '[' and i[-1] == ']']
-        self.structure_token_ids = [vocab_list.index(i) for i in structure_tokens if i in vocab_list]
-        self.structure_split_token_ids = [vocab_list.index(i) for i in structure_split_tokens]
-        self.sentence_split_token_ids = [vocab_list.index(i) for i in sentence_split_tokens]
-
-        # here initialize a output_proj (nn.Embedding) layer
-        # By default the first vocab is "" (null)
-        if mode == 'sum':
-            content_output_dim = output_dim
-            sentence_output_dim = output_dim
-            structure_output_dim = output_dim
-        else:   # concat'
-            raise NotImplementedError("concat 模式还未实现")    
-            # content_output_dim = output_dim - sentence_output_dim - structure_output_dim   # by default
-            
-        super().__init__(voc_size, content_output_dim, input_token=True, padding_idx=0)
-        self.special_emb = nn.Embedding(voc_size, structure_output_dim, padding_idx=0)
-        
-        self.blank_emb = nn.Parameter(torch.zeros(1, output_dim), requires_grad=False)
-
-        # the first index is "empty structure" token
-        self.sentence_idx_in_structure_emb = nn.Embedding(max_sentence_per_structure, sentence_output_dim) 
-        self.sentence_reidx_in_structure_emb = nn.Embedding(max_sentence_per_structure, sentence_output_dim)
-
-        print("max_len", self.max_len)
-        print(self.structure_token_ids)
-        
-        self.resolution = max_duration / max_len    # e.g., 120 / 600 = 0.2s 
-        print(self.__class__, f"resolution = {self.resolution}")
-    
-    def tokenize(self, x: tp.List[tp.Optional[str]]) -> tp.Dict[str, torch.Tensor]:
-        inputs = []
-        for xx in x:
-            xx = '' if xx is None else xx
-            vocab_id = [self.vocab_list.index(item) for item in xx.split(" ") if item in self.vocab_list]
-            inputs.append(torch.tensor(vocab_id).long()) # [T]
-        return inputs
-            
-            
-    def forward(self, batch_tokens: tp.List, structure_dur = None) -> ConditionType:
-        """
-        Encode token_id into three types of embeddings:
-        1) content embedding: phoneme only (or meaningful contents to be sung out) 
-        2) structure embedding: structure / separation embeddings, including structures (verse/chorus/...), separators (. / ,)
-        The two above share the same embedding layer, can be changed to separate embedding layers.
-        3) sentence_idx embedding (per structure): 
-        """
-        embeds_batch = []
-        for b in range(len(batch_tokens)):
-            tokens = batch_tokens[b]
-            content_tokens = torch.zeros_like(tokens)
-            special_tokens = torch.zeros_like(tokens)
-            sentence_idx_in_structure_tokens = torch.zeros_like(tokens) 
-            sentence_reidx_in_structure_tokens = torch.zeros_like(tokens)
-
-            current_sentence_in_structure_idx = 1
-            current_structure = 0
-            for i in range(tokens.shape[-1]):
-                token = tokens[i]
-                if token in self.structure_token_ids:       # structure token
-                    # only update structure token, leave content and sentence index token null (default 0)
-                    special_tokens[i] = token
-                    content_tokens[i] = token
-                    current_structure = token
-                    current_sentence_in_structure_idx = 1
-                    sentence_idx_in_structure_tokens[i] = 0
-
-                elif token in self.sentence_split_token_ids:    # utterance split token
-                    # only update structure token, leave content and sentence index token null (default 0)
-                    # add up sentence index
-                    special_tokens[i] = current_structure
-                    content_tokens[i] = token
-                    sentence_idx_in_structure_tokens[i] = min(current_sentence_in_structure_idx, self.max_sentence_per_structure - 1)
-                    current_sentence_in_structure_idx += 1
-
-                elif token in self.structure_split_token_ids:    # structure split token
-                    # update structure token (current structure), content token (current token), 
-                    # blank index token 
-                    content_tokens[i] = token
-                    special_tokens[i] = current_structure
-                    sentence_idx_in_structure_tokens[i] = sentence_idx_in_structure_tokens[i-1]
-                else:       # content tokens
-                    content_tokens[i] = token
-                    special_tokens[i] = current_structure
-                    sentence_idx_in_structure_tokens[i] = min(current_sentence_in_structure_idx, self.max_sentence_per_structure - 1)
-            # 反推
-            current_sentence_num = sentence_idx_in_structure_tokens[-1]
-            for i in range(tokens.shape[-1]-1,-1,-1):
-                if current_sentence_num != 0:
-                    sentence_reidx_in_structure_tokens[i] = min(current_sentence_num + 1 - sentence_idx_in_structure_tokens[i], self.max_sentence_per_structure - 1)
-                if sentence_idx_in_structure_tokens[i] == 0 and i > 0:
-                    current_sentence_num = sentence_idx_in_structure_tokens[i-1]
-
-            # print("tokens", tokens.max(), tokens.min())
-            # print("special tokens", special_tokens.max(), special_tokens.min())
-            # print("sentence idx in structure", sentence_idx_in_structure_tokens.max(), sentence_idx_in_structure_tokens.min())
-            device = self.output_proj.weight.device
-
-            # import pdb; pdb.set_trace()
-            content_embeds = self.output_proj(content_tokens.to(device))    # [T, N]
-            structure_embeds = self.output_proj(special_tokens.to(device))
-            # sentence_idx_embeds = self.sentence_idx_in_structure_emb(sentence_idx_in_structure_tokens.to(device))
-            sentence_idx_embeds = self.sentence_idx_in_structure_emb(sentence_idx_in_structure_tokens.to(device)) + self.sentence_reidx_in_structure_emb(sentence_reidx_in_structure_tokens.to(device))
-
-            if self.mode == 'sum':
-                embeds = content_embeds + structure_embeds + sentence_idx_embeds
-            else:
-                embeds = torch.cat((content_embeds, structure_embeds, sentence_idx_embeds), -1) # [T, N]
-            embeds_batch.append(embeds)
-
-        # set batch_size = 1, [B, T, N]
-        if self.max_len is not None:
-            max_len = self.max_len
-        else:
-            max_len = max([e.shape[0] for e in embeds_batch])
-        embeds, mask = self.pad_2d_tensor(embeds_batch, max_len)
-        
-        return embeds, embeds, mask
-    
-    
-    def pad_2d_tensor(self, xs, max_len):
-        new_tensor = []
-        new_mask = []
-        for x in xs:
-            seq_len, dim = x.size()
-            pad_len = max_len - seq_len
-
-            if pad_len > 0:
-                pad_tensor = self.blank_emb.repeat(pad_len, 1).to(x.device)  # T, D
-                padded_tensor = torch.cat([x, pad_tensor], dim=0)
-                mask = torch.cat((torch.ones_like(x[:, 0]), 
-                                  torch.zeros_like(pad_tensor[:, 0])), 0)   # T
-            elif pad_len < 0:
-                padded_tensor = x[:max_len]
-                mask = torch.ones_like(padded_tensor[:, 0])
-            else:
-                padded_tensor = x
-                mask = torch.ones_like(x[:, 0])
-
-            new_tensor.append(padded_tensor)
-            new_mask.append(mask)
-        # [B, T, D] & [B, T]
-        return torch.stack(new_tensor, 0), torch.stack(new_mask, 0)   
-
-
 class QwTokenizerConditioner(TextConditioner):
     def __init__(self, output_dim: int, 
                  token_path = "",

codeclm/tokenizer/audio_tokenizer.py

@@ -92,515 +92,16 @@ class AudioTokenizer(ABC, nn.Module):
             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)
-        elif name.split('_')[0] == 'FlowVocalAndMusicDecoderStereo':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = FlowVocalAndMusicDecoderStereo(model_type, mode=mode)
-        elif name.split('_')[0] == 'FlowVocalAndMusicDecoderStereoLayer7':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = FlowVocalAndMusicDecoderStereoLayer7(model_type, mode=mode)  
-        elif name.split('_')[0] == 'FlowVocalAndMusicDecoderStereoLayer11':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = FlowVocalAndMusicDecoderStereoLayer11(model_type, mode=mode)
-        elif name.split('_')[0] == 'FlowVocalAndMusicDecoderStereoASRTuneLayer7':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = FlowVocalAndMusicDecoderStereoASRTuneLayer7(model_type, mode=mode)
-        elif name.split('_')[0] == 'FlowVocalAndMusicDecoderStereoASRTuneLayer7Code2':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = FlowVocalAndMusicDecoderStereoASRTuneLayer7Code2(model_type, mode=mode)
-        elif name.split('_')[0] == 'FlowVocalAndMusicDecoderStereoASRTuneLayer7Code1':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = FlowVocalAndMusicDecoderStereoASRTuneLayer7Code1(model_type, mode=mode)
-        elif name.split('_')[0] == 'Flow1dVAE2rvq':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = Flow1dVAE2rvq(model_type)
-        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)
-        elif name.split('_')[0] == 'Flow1dVAE4rvq':
-            model_type = name.split('_', 1)[1]
-            logger.info("Getting pretrained compression model from semantic model %s", model_type)
-            model = Flow1dVAE4rvq(model_type)
-        else:
-            raise NotImplementedError("{} is not implemented in models/audio_tokenizer.py".format(
-                name))
-        return model.to(device).eval()
-    
-
-class FlowVocalAndMusicDecoderStereo(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str, 
-        sample_rate=48000, 
-        mode = 'extract',
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.FlowVocalAndMusicDecoderStereoV014.generate_stereo import Tango
-        model_path = model_type
-        self.mode = mode
-        if mode == 'extract':
-            self.model = Tango(model_path=model_path, layer_num=3, load_main_model=False, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-        elif mode == 'inference':
-            self.samplerate = sample_rate
-            self.model = Tango(model_path=model_path, layer_num=3, load_main_model=True, 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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, duration=40.96, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
-    
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
-
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
-
-class FlowVocalAndMusicDecoderStereoLayer7(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str = "pytorch_model_2.bin", 
-        sample_rate=48000, 
-        mode = 'extract',
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.FlowVocalAndMusicDecoderStereoV014.generate_stereo_layer7 import Tango
-        model_path = model_type
-        self.mode = mode
-        if mode == 'extract':
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=False, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-        elif mode == 'inference':
-            self.samplerate = sample_rate
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=True, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-            # print("Successfully loaded inference scheduler from {}".format(scheduler_name))
-            
-        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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, duration=40.96, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
-    
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
-
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
-
-class FlowVocalAndMusicDecoderStereoASRTuneLayer7(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str = "model_layer7_1x4.safetensors", 
-        sample_rate=48000, 
-        mode = 'extract',
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.FlowVocalAndMusicDecoderStereoV014.generate_stereo_7_1x4 import Tango
-        model_path = model_type
-        self.mode = mode
-        if mode == 'extract':
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=False, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-        elif mode == 'inference':
-            self.samplerate = sample_rate
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=True, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-            # print("Successfully loaded inference scheduler from {}".format(scheduler_name))
-            
-        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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, duration=40.96, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
-    
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
-
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
-class FlowVocalAndMusicDecoderStereoASRTuneLayer7Code2(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str = "model_layer7_1x2.safetensors", 
-        sample_rate=48000, 
-        mode = 'extract',
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.FlowVocalAndMusicDecoderStereoV014.generate_stereo_7_1x2 import Tango
-        model_path = model_type
-        self.mode = mode
-        if mode == 'extract':
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=False, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-        elif mode == 'inference':
-            self.samplerate = sample_rate
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=True, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-            # print("Successfully loaded inference scheduler from {}".format(scheduler_name))
-            
-        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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, duration=40.96, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
-    
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
-
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
-class FlowVocalAndMusicDecoderStereoASRTuneLayer7Code1(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str = "model_layer7_1x1.safetensors",
-        sample_rate=48000, 
-        mode = 'extract',
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.FlowVocalAndMusicDecoderStereoV014.generate_stereo_7_1x1 import Tango
-        model_path = model_type
-        self.mode = mode
-        if mode == 'extract':
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=False, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-        elif mode == 'inference':
-            self.samplerate = sample_rate
-            self.model = Tango(model_path=model_path, layer_num=7, load_main_model=True, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-            # print("Successfully loaded inference scheduler from {}".format(scheduler_name))
-            
-        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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, duration=40.96, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
-    
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
-
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
-class Flow1dVAE2rvq(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str = "model_2.safetensors",
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.Flow1dVAE.generate_2rvq import Tango
-        model_path = model_type
-        self.model = Tango(model_path=model_path, rvq_num=2, 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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
+        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)
+        else:
+            raise NotImplementedError("{} is not implemented in models/audio_tokenizer.py".format(
+                name))
+        return model.to(device).eval()
     
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
 
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
 class Flow1dVAE1rvq(AudioTokenizer):
     def __init__(
         self, 
@@ -674,78 +175,6 @@ class Flow1dVAE1rvq(AudioTokenizer):
         assert n >= 1
         assert n <= self.total_codebooks
         self.n_quantizers = n
-class Flow1dVAE4rvq(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str = "model_2.safetensors",
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.Flow1dVAE.generate_4rvq import Tango
-        model_path = model_type
-        self.model = Tango(model_path=model_path, rvq_num=4, 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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
-    
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
-
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
-
 
 
 class Flow1dVAESeparate(AudioTokenizer):
@@ -822,86 +251,3 @@ class Flow1dVAESeparate(AudioTokenizer):
         assert n >= 1
         assert n <= self.total_codebooks
         self.n_quantizers = n
-
-class FlowVocalAndMusicDecoderStereoLayer11(AudioTokenizer):
-    def __init__(
-        self, 
-        model_type: str = "layer11_ckpt.pth", 
-        sample_rate=48000, 
-        mode = 'extract',
-        ):
-        super().__init__()
-
-        from codeclm.tokenizer.FlowVocalAndMusicDecoderStereoV014.generate_stereo_11 import Tango
-        model_path = model_type
-        self.mode = mode
-        if mode == 'extract':
-            self.model = Tango(model_path=model_path, layer_num=11, load_main_model=False, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-        elif mode == 'inference':
-            self.samplerate = sample_rate
-            self.model = Tango(model_path=model_path, layer_num=11, load_main_model=True, device='cuda')
-            print ("Successfully loaded checkpoint from:", model_path)
-            # print("Successfully loaded inference scheduler from {}".format(scheduler_name))
-            
-        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.")
-    
-    @torch.no_grad()
-    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        if x.ndim == 2:
-            x = x.unsqueeze(1)
-        codes = self.model.sound2code(x) # [B T] -> [B N T]
-        return codes, None
-
-    
-    @torch.no_grad()    
-    def decode(self, codes: torch.Tensor, prompt = None, scale: tp.Optional[torch.Tensor] = None, ncodes=9):
-        wav = self.model.code2sound(codes, prompt=prompt, duration=40.96, guidance_scale=1.5, 
-                                    num_steps=50, disable_progress=False) # [B,N,T] -> [B,T]
-        return wav[None]
-
-    
-    @torch.no_grad()
-    def decode_latent(self, codes: torch.Tensor):
-        """Decode from the discrete codes to continuous latent space."""
-        # import pdb; pdb.set_trace()
-        return self.model.quantizer.from_codes(codes.transpose(1,2))[0]
-
-    @property
-    def channels(self) -> int:
-        return 2
-
-    @property
-    def frame_rate(self) -> float:
-        return 25
-
-    @property
-    def sample_rate(self) -> int:
-        return self.samplerate
-
-    @property
-    def cardinality(self) -> int:
-        return 10000
-
-    @property
-    def num_codebooks(self) -> int:
-        return self.n_quantizers
-
-    @property
-    def total_codebooks(self) -> int:
-        # return self.model.RVQ
-        return 1
-
-    def set_num_codebooks(self, n: int):
-        """Set the active number of codebooks used by the quantizer.
-        """
-        assert n >= 1
-        assert n <= self.total_codebooks
-        self.n_quantizers = n
-    
-    
-    

codeclm/trainer/codec_song_pl.py

@@ -26,7 +26,7 @@ os.environ['TOKENIZERS_PARALLELISM'] = "false"
 
 
 class CodecLM_PL(pl.LightningModule):
-    def __init__(self, cfg):
+    def __init__(self, cfg, ckpt_path):
         super().__init__()
 
         self.cfg = cfg
@@ -46,30 +46,12 @@ class CodecLM_PL(pl.LightningModule):
         # 2) Build LM
         self.audiolm = builders.get_lm_model(self.cfg)
         print(self.audiolm)
-        # 输出参数量
-        print('Number of parameters: ', sum(p.numel() for p in self.audiolm.parameters()))
         # 3) Load pretrained checkpoint (if any)
-        if self.cfg.use_pretrained == 'deepspeed':
-            checkpoint = torch.load(self.cfg.pretrained.deepspeed_checkpoint, 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(self.cfg.pretrained.deepspeed_checkpoint))
-            self.missing = missing
-        else:
-            self.missing = []
-        # 如果cfg参数中有lora
-        if hasattr(self.cfg, 'lora'):
-            perf_config = LoraConfig(
-                r = self.cfg.lora.r,
-                lora_alpha = self.cfg.lora.lora_alpha,
-                target_modules = self.cfg.lora.target_modules,
-                lora_dropout = self.cfg.lora.lora_dropout,
-                bias = self.cfg.lora.bias,
-                task_type = self.cfg.lora.task_type,
-            )
-            self.audiolm = get_peft_model(self.audiolm, perf_config)
-        
+        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))
         # 4) Build metrics
         self.val_steps = []
         self.train_slide_acc = []
@@ -113,32 +95,6 @@ class CodecLM_PL(pl.LightningModule):
         x = torch.where(mask_3d, x, end_id+1)
         return x, mask_3d
 
-    @torch.no_grad()
-    def preprocess_batch(self, batch):  # this function is usually called during training
-        # 处理 dataloader 返回的数据
-        audio, text_lyric, time_stamp, structure_dur, prompt_audio, structure_labels = batch
-
-        dur, valid_st, valid_et = zip(*time_stamp)
-        
-        if self.audio_tokenizer is not None:
-            # only used in inference
-            self.audio_tokenizer.eval()
-            with torch.no_grad():
-                with torch.cuda.amp.autocast(enabled=False):
-                    audio_tokens, scale = self.audio_tokenizer.encode(audio)
-                audio_tokens = audio_tokens[:,:self.cfg.lm.code_depth,:]
-                audio_tokens = audio_tokens.long()
-        else:
-            audio_tokens = audio.long()
-        
-        token_dur = (torch.Tensor(dur) * self.cfg.audio_tokenizer_frame_rate).int()
-        audio_tokens, audio_padding_mask = self.generate_mask_and_end_token(audio_tokens, token_dur, 
-                                                                            end_id=self.audiolm.eos_token_id)
-        condition_tensors = self.audiolm.prepare_condition_tensors(batch_size=len(text_lyric),
-                                                                   text=text_lyric, audio_qt_emb=prompt_audio)
-
-        return condition_tensors, audio_tokens, audio_padding_mask
-
     def get_time(self):
         # 获取当前的日期和时间
         now = datetime.now()
@@ -147,506 +103,6 @@ class CodecLM_PL(pl.LightningModule):
         formatted_now = now.strftime("%Y-%m-%d %H:%M:%S.%f")
         return formatted_now
 
-    def training_step(self, batch, batch_idx):
-        # 1) data processing
-        condition_tensors, audio_tokens, padding_mask = self.preprocess_batch(batch)
-        
-        # 2) compute model predictions (model forward)
-        model_output = self.audiolm.compute_predictions(audio_tokens, condition_tensors, 
-                                                        training_steps=self.global_step)  # this input can be ignored        
-        logits = model_output.logits.float()
-        mask = padding_mask & model_output.mask
-        
-        # 3) compute loss (float)
-        with torch.cuda.amp.autocast(enabled=False):
-            ce, ce_per_codebook = self._compute_cross_entropy(logits, audio_tokens, mask)
-        
-        total_loss = ce
-        if torch.isnan(total_loss):
-            print(self.trainer.global_rank, ce, padding_mask, batch[1])
-            print('--------------------------------------------------------------')
-            return None
-            # torchaudio.save("error_rank{}.wav".format(self.trainer.global_rank), batch[0][:,0].cpu(), 24000)
-            # import pdb; pdb.set_trace()
-        # 4) compute metrics and log
-        metrics = {}
-        self.log('ce', ce, prog_bar=True) 
-        metrics['ppl'] = torch.exp(ce)
-        for k, ce_q in enumerate(ce_per_codebook):
-            metrics[f'ce_q{k + 1}'] = ce_q
-            metrics[f'ppl_q{k + 1}'] = torch.exp(ce_q)
-
-        masked_labels = audio_tokens.masked_fill(~mask, value=self.cfg.lm.code_size)
-        metrics['acc'] = []
-        for k in range(self.audiolm.code_depth):
-            metrics['acc'].append(self.top1_acc_metric[k](logits[:, k].transpose(1,2).detach(), 
-                                                          masked_labels[:, k]).item())
-        metrics['acc'] = torch.mean(torch.Tensor(metrics['acc'])).item()
-
-        self.train_steps.append({'ce': ce.detach().cpu().item(), 'acc': metrics['acc']})        
-        self.log('train_acc', metrics['acc']+1e-8, prog_bar=True)
-        self.log('lr', self.trainer.optimizers[0].param_groups[0]['lr'], prog_bar=True)  
-        self.log_dict(metrics)
-
-        return total_loss
-    
-    @torch.no_grad()
-    def validation_step(self, batch, batch_idx):
-        # 1) data processing
-        condition_tensors, audio_tokens, padding_mask = self.preprocess_batch(batch)
-
-        # 2) compute model predictions
-        model_output = self.audiolm.compute_predictions(audio_tokens, condition_tensors)  
-        logits = model_output.logits
-        mask = padding_mask & model_output.mask
-        
-        # 3) compute loss and metrics
-        ce, ce_per_codebook = self._compute_cross_entropy(logits, audio_tokens, mask)
-        metrics = {}   
-        metrics['val_ce'] = ce
-        metrics['val_ppl'] = torch.exp(ce)
-        for k, ce_q in enumerate(ce_per_codebook):
-            metrics[f'val_ce_q{k + 1}'] = ce_q
-            metrics[f'val_ppl_q{k + 1}'] = torch.exp(ce_q)
-        masked_labels = audio_tokens.masked_fill(~mask, value=self.cfg.lm.code_size)
-
-        for k in range(self.audiolm.code_depth):
-            self.top1_acc_metric[k].update(logits[:, k].transpose(1,2).detach(), masked_labels[:,k]) #* total_length
-            self.top10_acc_metric[k].update(logits[:, k].transpose(1,2).detach(), masked_labels[:,k])
-        self.val_steps.append(metrics)
-            
-        metrics['acc'] = []
-        metrics['acc_top10'] = []
-        for k in range(self.audiolm.code_depth):
-            metrics['acc'].append(self.top1_acc_metric[k](logits[:, k].transpose(1,2).detach(), masked_labels[:,k]).item())
-            metrics['acc_top10'].append(self.top10_acc_metric[k](logits[:, k].transpose(1,2).detach(), masked_labels[:,k]).item())
-        metrics['acc'] = torch.mean(torch.Tensor(metrics['acc']))
-        metrics['acc_top10'] = torch.mean(torch.Tensor(metrics['acc_top10'])) 
-        
-        return metrics['acc']
-
-
-    def on_validation_epoch_end(self) -> None:        
-        final_metrics = {}
-        for i in self.val_steps:
-            for k in i:
-                final_metrics[k] = final_metrics.get(k, []) + [i[k]]
-        final_metrics = {k: sum(v) / len(v) for k,v in list(final_metrics.items())}
-        self.log_dict(final_metrics)
-
-        q_acc = []
-        q_acc10 = []
-        for i in range(self.audiolm.code_depth):
-            q_acc.append(self.top1_acc_metric[i].compute())
-            q_acc10.append(self.top10_acc_metric[i].compute())
-            self.log(f"val_Top1Acc_{i}", q_acc[-1])
-            self.log(f"val_Top10Acc_{i}", q_acc10[-1])
-            self.top1_acc_metric[i].reset()
-            self.top10_acc_metric[i].reset()
-        
-        self.log('val_Top1Acc', sum(q_acc) / self.audiolm.code_depth)
-        self.log('val_Top10Acc', sum(q_acc10) / self.audiolm.code_depth)
-
-        return super().on_validation_epoch_end()
-
-
-    def on_validation_epoch_start(self) -> None:
-        self.val_steps = []
-        for i in range(self.audiolm.code_depth):
-            self.top1_acc_metric[i].reset()
-            self.top10_acc_metric[i].reset()
-
-        if len(self.train_steps) > 0:
-            train_metrics = {}
-            for i in self.train_steps:
-                for k in i:
-                    train_metrics[k] = train_metrics.get(k, []) + [i[k]]
-            train_metrics = {k: sum(v) / len(v) for k,v in list(train_metrics.items())}
-            self.log('train_summary_Top1Acc', train_metrics['acc'])
-            self.log('train_summary_ce', train_metrics['ce'])
-            self.train_steps = []
-
-        return super().on_validation_epoch_start()
-
-
-    # 定义优化器
-    def configure_optimizers(self):
-        total_updates = self.cfg.optim.epochs * self.cfg.optim.updates_per_epoch
-        optim_dict = {}
-
-        param_groups = []
-        missing_params = []
-        other_params = []
-        cnt = 0
-        # 去掉开头的‘audiolm.'
-        print('before missing len', len(self.missing))
-        self.missing = [name.replace('audiolm.', '') for name in self.missing]
-        print('after missing len', len(self.missing))
-        for name, param in self.audiolm.named_parameters():
-            if name in self.missing:
-                cnt += 1
-                print(name)
-                missing_params.append(param)
-            else:
-                other_params.append(param)
-        print(cnt)
-        assert cnt == len(self.missing)
-        param_groups.append({'params': other_params, 'lr': self.cfg.optim.old_lr})
-        param_groups.append({
-            'params': missing_params,
-            'lr': self.cfg.optim.new_lr  # 为missing参数设置10倍的学习率,你可以调整这个倍数
-        })
-
-        if self.cfg.optim.optimizer == "adamw":
-            optim_dict['optimizer'] = torch.optim.AdamW(
-                param_groups,  # 使用参数分组替代原来的 self.audiolm.parameters()
-                betas=tuple(self.cfg.optim.adam.betas),
-                weight_decay=self.cfg.optim.adam.weight_decay,
-                eps=self.cfg.optim.adam.eps,
-            )
-        else:
-            raise NotImplementedError
-
-        if self.cfg.schedule is None:
-            pass
-        elif self.cfg.schedule.lr_scheduler == "cosine":
-            scheduler = CosineLRScheduler(optim_dict['optimizer'], 
-                                          total_steps=total_updates, 
-                                          warmup_steps=self.cfg.schedule.cosine.warmup,
-                                          lr_min_ratio=self.cfg.schedule.cosine.lr_min_ratio,
-                                          cycle_length=self.cfg.schedule.cosine.cycle_length,
-                                          )
-            optim_dict['lr_scheduler'] = {"scheduler": scheduler, "interval": "step"}
-        else:
-            raise NotImplementedError
-        
-        return optim_dict
-
-    
-    def _compute_cross_entropy(
-        self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor
-    ) -> tp.Tuple[torch.Tensor, tp.List[torch.Tensor]]:
-        """Compute cross entropy between multi-codebook targets and model's logits.
-        The cross entropy is computed per codebook to provide codebook-level cross entropy.
-        Valid timesteps for each of the codebook are pulled from the mask, where invalid
-        timesteps are set to 0.
-
-        Args:
-            logits (torch.Tensor): Model's logits of shape [B, K, T, card].
-            targets (torch.Tensor): Target codes, of shape [B, K, T].
-            mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T].
-        Returns:
-            ce (torch.Tensor): Cross entropy averaged over the codebooks
-            ce_per_codebook (list of torch.Tensor): Cross entropy per codebook (detached).
-        """
-        # import pdb; pdb.set_trace()
-        B, K, T = targets.shape
-        assert logits.shape[:-1] == targets.shape
-        assert mask.shape == targets.shape
-        ce = torch.zeros([], device=targets.device)
-        ce_per_codebook: tp.List[torch.Tensor] = []
-        for k in range(K):
-            logits_k = logits[:, k, ...].contiguous().view(-1, logits.size(-1))  # [B x T, card]
-            targets_k = targets[:, k, ...].contiguous().view(-1)  # [B x T]
-            mask_k = mask[:, k, ...].contiguous().view(-1)  # [B x T]
-            ce_targets = targets_k[mask_k]
-            ce_logits = logits_k[mask_k]
-            q_ce = F.cross_entropy(ce_logits, ce_targets)
-            ce += q_ce
-            ce_per_codebook.append(q_ce.detach())
-        # average cross entropy across codebooks
-        ce = ce / K
-        return ce, ce_per_codebook
-
-
-class CodecLM_PL_FT(pl.LightningModule):
-    def __init__(self, cfg):
-        super().__init__()
-
-        self.cfg = cfg
-        
-        # 1) Build audio tokenizer (usually None during training)
-        self.audio_tokenizer = builders.get_audio_tokenizer_model(self.cfg)
-        if self.audio_tokenizer is not None:
-            for param in self.audio_tokenizer.parameters():
-                param.requires_grad = False
-        
-        # 2) Build LM
-        self.audiolm = builders.get_lm_model(self.cfg)
-        
-        # 3) Load pretrained checkpoint (if any)
-        if self.cfg.use_pretrained == 'deepspeed':
-            checkpoint = torch.load(self.cfg.pretrained.deepspeed_checkpoint,  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(self.cfg.pretrained.deepspeed_checkpoint))
-
-        # 4) Build metrics
-        self.val_steps = []
-        self.train_slide_acc = []
-        self.train_steps = []
-        self.top1_acc_metric = nn.ModuleList([MulticlassAccuracy(
-            self.audiolm.code_size, 
-            top_k=1,
-            average="micro", multidim_average="global",
-            ignore_index=self.cfg.lm.code_size, # ignore EOS token prediction
-        ) for _ in range(self.audiolm.code_depth)])
-        self.top10_acc_metric = nn.ModuleList([MulticlassAccuracy(
-            self.audiolm.code_size,
-            top_k=10,
-            average="micro", multidim_average="global",
-            ignore_index=self.cfg.lm.code_size,
-        ) 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):
-        batch_size = sequence_lengths.size(0)
-        max_length = x.size(2)
-
-        # pad one frame, if the maximum sequence length is equal to the input length
-        if max_length == sequence_lengths.max():
-            x = F.pad(x, (0, 1), value=end_id)
-        max_length = x.size(2)
-
-        if max_length <= sequence_lengths.max() + 1:
-            sequence_lengths = sequence_lengths - (sequence_lengths.max()+1 - max_length)
-
-        # Add end token to x according to the sequence length
-        x[torch.arange(batch_size), :, sequence_lengths] = end_id
-        sequence_lengths += 1
-
-        mask = torch.arange(max_length).expand(batch_size, max_length) < sequence_lengths.unsqueeze(1)
-        mask = mask.to(x.device)
-        mask_3d = mask.unsqueeze(1).expand(batch_size, x.size(1), max_length)
-        x = torch.where(mask_3d, x, end_id+1)
-        return x, mask_3d
-
-    @torch.no_grad()
-    def preprocess_batch(self, batch):  # this function is usually called during training
-        # 处理 dataloader 返回的数据
-        audio, text_lyric, time_stamp, lang_type, prompt_audio = batch        
-        dur, valid_st, valid_et = zip(*time_stamp)
-        
-        if self.audio_tokenizer is not None:
-            # only used in inference
-            self.audio_tokenizer.eval()
-            with torch.no_grad():
-                with torch.cuda.amp.autocast(enabled=False):
-                    audio_tokens, scale = self.audio_tokenizer.encode(audio)
-                audio_tokens = audio_tokens[:,:self.cfg.lm.code_depth,:]
-                audio_tokens = audio_tokens.long()
-        else:
-            audio_tokens = audio.long()
-        
-        token_dur = (torch.Tensor(dur) * self.cfg.audio_tokenizer_frame_rate).int()
-
-        audio_tokens, audio_padding_mask = self.generate_mask_and_end_token(audio_tokens, token_dur, 
-                                                                            end_id=self.audiolm.eos_token_id)
-        condition_tensors = self.audiolm.prepare_condition_tensors(batch_size=len(text_lyric),
-                                                                   text=text_lyric, audio_qt_emb=prompt_audio)
-
-        return condition_tensors, audio_tokens, audio_padding_mask
-
-    def get_time(self):
-        # 获取当前的日期和时间
-        now = datetime.now()
-
-        # 使用strftime函数格式化日期和时间
-        formatted_now = now.strftime("%Y-%m-%d %H:%M:%S.%f")
-        return formatted_now
-
-    def training_step(self, batch, batch_idx):
-        # 1) data processing
-        condition_tensors, audio_tokens, padding_mask = self.preprocess_batch(batch)
-        
-        # 2) compute model predictions (model forward)
-        model_output = self.audiolm.compute_predictions(audio_tokens, condition_tensors, 
-                                                        training_steps=self.global_step)  # this input can be ignored        
-        logits = model_output.logits.float()
-        mask = padding_mask & model_output.mask
-        
-        # 3) compute loss (float)
-        with torch.cuda.amp.autocast(enabled=False):
-            ce, ce_per_codebook = self._compute_cross_entropy(logits, audio_tokens, mask)
-        
-        total_loss = ce
-        if torch.isnan(total_loss):
-            print(self.trainer.global_rank, ce, padding_mask, batch[1])
-            # print('------------------------------------------------------------------------')
-            torchaudio.save("error_rank{}.wav".format(self.trainer.global_rank), batch[0][:,0].cpu(), 24000)
-            import pdb; pdb.set_trace()
-            return None
-
-        # 4) compute metrics and log
-        metrics = {}
-        self.log('ce', ce, prog_bar=True) 
-        metrics['ppl'] = torch.exp(ce)
-        for k, ce_q in enumerate(ce_per_codebook):
-            metrics[f'ce_q{k + 1}'] = ce_q
-            metrics[f'ppl_q{k + 1}'] = torch.exp(ce_q)
-
-        masked_labels = audio_tokens.masked_fill(~mask, value=self.cfg.lm.code_size)
-        metrics['acc'] = []
-        for k in range(self.audiolm.code_depth):
-            metrics['acc'].append(self.top1_acc_metric[k](logits[:, k].transpose(1,2).detach(), 
-                                                          masked_labels[:, k]).item())
-        metrics['acc'] = torch.mean(torch.Tensor(metrics['acc'])).item()
-
-        self.train_steps.append({'ce': ce.detach().cpu().item(), 'acc': metrics['acc']})        
-        self.log('train_acc', metrics['acc']+1e-8, prog_bar=True)
-        self.log('lr', self.trainer.optimizers[0].param_groups[0]['lr'], prog_bar=True)  
-        self.log_dict(metrics)
-
-        return total_loss
-    
-    @torch.no_grad()
-    def validation_step(self, batch, batch_idx):
-        # 1) data processing
-        condition_tensors, audio_tokens, padding_mask = self.preprocess_batch(batch)
-
-        # 2) compute model predictions
-        model_output = self.audiolm.compute_predictions(audio_tokens, condition_tensors)  
-        logits = model_output.logits
-        mask = padding_mask & model_output.mask
-        
-        # 3) compute loss and metrics
-        ce, ce_per_codebook = self._compute_cross_entropy(logits, audio_tokens, mask)
-        metrics = {}   
-        metrics['val_ce'] = ce
-        metrics['val_ppl'] = torch.exp(ce)
-        for k, ce_q in enumerate(ce_per_codebook):
-            metrics[f'val_ce_q{k + 1}'] = ce_q
-            metrics[f'val_ppl_q{k + 1}'] = torch.exp(ce_q)
-        masked_labels = audio_tokens.masked_fill(~mask, value=self.cfg.lm.code_size)
-
-        for k in range(self.audiolm.code_depth):
-            self.top1_acc_metric[k].update(logits[:, k].transpose(1,2).detach(), masked_labels[:,k]) #* total_length
-            self.top10_acc_metric[k].update(logits[:, k].transpose(1,2).detach(), masked_labels[:,k])
-        self.val_steps.append(metrics)
-        metrics['acc'] = []
-        metrics['acc_top10'] = []
-        for k in range(self.audiolm.code_depth):
-            metrics['acc'].append(self.top1_acc_metric[k](logits[:, k].transpose(1,2).detach(), masked_labels[:,k]).item())
-            metrics['acc_top10'].append(self.top10_acc_metric[k](logits[:, k].transpose(1,2).detach(), masked_labels[:,k]).item())
-        metrics['acc'] = torch.mean(torch.Tensor(metrics['acc']))
-        metrics['acc_top10'] = torch.mean(torch.Tensor(metrics['acc_top10']))
-        
-        return metrics['acc']
-
-    def on_validation_epoch_end(self) -> None:        
-        final_metrics = {}
-        for i in self.val_steps:
-            for k in i:
-                final_metrics[k] = final_metrics.get(k, []) + [i[k]]
-        final_metrics = {k: sum(v) / len(v) for k,v in list(final_metrics.items())}
-        self.log_dict(final_metrics)
-
-        q_acc = []
-        q_acc10 = []
-        for i in range(self.audiolm.code_depth):
-            q_acc.append(self.top1_acc_metric[i].compute())
-            q_acc10.append(self.top10_acc_metric[i].compute())
-            self.log(f"val_Top1Acc_{i}", q_acc[-1])
-            self.log(f"val_Top10Acc_{i}", q_acc10[-1])
-            self.top1_acc_metric[i].reset()
-            self.top10_acc_metric[i].reset()
-        
-        self.log('val_Top1Acc', sum(q_acc) / self.audiolm.code_depth)
-        self.log('val_Top10Acc', sum(q_acc10) / self.audiolm.code_depth)
-
-        return super().on_validation_epoch_end()
-
-
-    def on_validation_epoch_start(self) -> None:
-        self.val_steps = []
-        for i in range(self.audiolm.code_depth):
-            self.top1_acc_metric[i].reset()
-            self.top10_acc_metric[i].reset()
-
-        if len(self.train_steps) > 0:
-            train_metrics = {}
-            for i in self.train_steps:
-                for k in i:
-                    train_metrics[k] = train_metrics.get(k, []) + [i[k]]
-            train_metrics = {k: sum(v) / len(v) for k,v in list(train_metrics.items())}
-            self.log('train_summary_Top1Acc', train_metrics['acc'])
-            self.log('train_summary_ce', train_metrics['ce'])
-            self.train_steps = []
-
-        return super().on_validation_epoch_start()
-
-
-    # 定义优化器
-    def configure_optimizers(self):
-        total_updates = self.cfg.optim.epochs * self.cfg.optim.updates_per_epoch
-        optim_dict = {}
-
-        if self.cfg.optim.optimizer == "adamw":
-            optim_dict['optimizer'] = torch.optim.AdamW(
-                self.audiolm.parameters(),
-                lr=self.cfg.optim.lr,
-                betas=tuple(self.cfg.optim.adam.betas),
-                weight_decay=self.cfg.optim.adam.weight_decay,
-                eps=self.cfg.optim.adam.eps,
-            )
-        else:
-            raise NotImplementedError
-
-        if self.cfg.schedule is None:
-            pass
-        elif self.cfg.schedule.lr_scheduler == "cosine":
-            scheduler = CosineLRScheduler(optim_dict['optimizer'], 
-                                          total_steps=total_updates, 
-                                          warmup_steps=self.cfg.schedule.cosine.warmup,
-                                          lr_min_ratio=self.cfg.schedule.cosine.lr_min_ratio,
-                                          cycle_length=self.cfg.schedule.cosine.cycle_length,
-                                          )
-            optim_dict['lr_scheduler'] = {"scheduler": scheduler, "interval": "step"}
-        else:
-            raise NotImplementedError
-        
-        return optim_dict
-
-    
-    def _compute_cross_entropy(
-        self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor
-    ) -> tp.Tuple[torch.Tensor, tp.List[torch.Tensor]]:
-        """Compute cross entropy between multi-codebook targets and model's logits.
-        The cross entropy is computed per codebook to provide codebook-level cross entropy.
-        Valid timesteps for each of the codebook are pulled from the mask, where invalid
-        timesteps are set to 0.
-
-        Args:
-            logits (torch.Tensor): Model's logits of shape [B, K, T, card].
-            targets (torch.Tensor): Target codes, of shape [B, K, T].
-            mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T].
-        Returns:
-            ce (torch.Tensor): Cross entropy averaged over the codebooks
-            ce_per_codebook (list of torch.Tensor): Cross entropy per codebook (detached).
-        """
-        # import pdb; pdb.set_trace()
-        B, K, T = targets.shape
-        assert logits.shape[:-1] == targets.shape
-        assert mask.shape == targets.shape
-        ce = torch.zeros([], device=targets.device)
-        ce_per_codebook: tp.List[torch.Tensor] = []
-        for k in range(K):
-            logits_k = logits[:, k, ...].contiguous().view(-1, logits.size(-1))  # [B x T, card]
-            targets_k = targets[:, k, ...].contiguous().view(-1)  # [B x T]
-            mask_k = mask[:, k, ...].contiguous().view(-1)  # [B x T]
-            ce_targets = targets_k[mask_k]
-            ce_logits = logits_k[mask_k]
-            q_ce = F.cross_entropy(ce_logits, ce_targets)
-            ce += q_ce
-            ce_per_codebook.append(q_ce.detach())
-        # average cross entropy across codebooks
-        ce = ce / K
-        return ce, ce_per_codebook
-
 class CosineLRScheduler(_LRScheduler):# 
     """Cosine LR scheduler.
 

conf/infer.yaml

@@ -1,152 +0,0 @@
-# ================ Logging ====================== #
-root_dir: exp/song/${get_fname:}
-
-# ================ Checkpoints ================== #
-use_pretrained: deepspeed # ['ddp', 'continue', 'deepspeed']
-pretrained:
-  ddp_checkpoint:
-  deepspeed_checkpoint: ./ckpt/60000_alnew.pt
-  continue_checkpoint: 
-  
-# ================ Data & loader ================== #
-prompt_select: random
-train_jsonl_list:
-- .jsonl
-val_jsonl_list:
-- .jsonl
-train_scp_list:
-- .scp
-val_scp_list:
-- .scp
-
-lyric_processor:
-max_dur: 150
-min_dur: 30
-batch_size: 2
-prompt_len: 10
-pad_to_max: true
-
-# ================ Training ======================= #
-accelerator: gpu
-devices: 8
-num_nodes: 4
-val_check_interval: 2500
-accumulate_grad_batches: 1
-strategy: 'deepspeed_stage_2' # ['ddp', 'fsdp', 'deepspeed_stage_2', 'ddp_find_unused_parameters_true']
-precision: 'bf16-mixed' # ['16-mixed', 'bf16-mixed']
-
-optim:
-  optimizer: adamw
-  updates_per_epoch: 1000
-  epochs: 100
-  old_lr: 0 # 1e-4
-  new_lr: 1e-4
-  max_norm: 0.5
-  adam:
-    betas:
-    - 0.9
-    - 0.95
-    weight_decay: 0.00001 # 0.1
-    eps: 1e-8
-
-schedule:
-  lr_scheduler: cosine
-  cosine:
-    warmup: 4000
-    lr_min_ratio: 0.0
-    cycle_length: 1.0
-
-# ================ Audio tokenzier ================ #
-audio_tokenizer_checkpoint: Flow1dVAE1rvq_./ckpt/model_1rvq/model_2_fixed.safetensors
-audio_tokenizer_frame_rate: 25
-audio_tokenizer_code_depth: 1
-sample_rate: 48000
-
-audio_tokenizer_checkpoint_sep: Flow1dVAESeparate_./ckpt/model_septoken/model_2.safetensors
-audio_tokenizer_frame_rate_sep: 25
-audio_tokenizer_code_depth_sep: 2
-sample_rate_sep: 48000
-
-# ================ VAE ================ #
-vae_config: ./ckpt/vae/stable_audio_1920_vae.json
-vae_model: ./ckpt/vae/autoencoder_music_1320k.ckpt
-
-# ================== LM =========================== #
-lm:
-  lm_type: Llama # [Llama]
-  dim: 1536
-  intermediate_size: 8960
-  num_heads: 12
-  num_layers: 28
-  code_depth: 3
-  code_size: 16384
-  dropout: 0.0
-  activation: gelu
-  norm_first: true
-  bias_ff: false
-  bias_attn: false
-  bias_proj: false
-  causal: true
-  custom: false
-  memory_efficient: true
-  attention_as_float32: false
-  layer_scale: null
-  positional_embedding: sin
-  xpos: false
-  checkpointing: torch
-  weight_init: gaussian
-  depthwise_init: current
-  zero_bias_init: true
-  norm: layer_norm
-  cross_attention: false
-  qk_layer_norm: false
-  qk_layer_norm_cross: false
-  attention_dropout: null
-  kv_repeat: 1
-
-codebooks_pattern:
-  modeling: delay
-  delay:
-    delays: [ 0, 250, 250 ]
-    flatten_first: 0
-    empty_initial: 0
-
-# ================ Conditioners ===================== #
-classifier_free_guidance:
-  # drop all conditions simultaneously
-  training_dropout: 0.15
-  inference_coef: 1.5
-
-attribute_dropout:
-  # drop each condition separately
-  args:
-    active_on_eval: false
-  text:
-    description: 0.0
-    type_info: 0.5
-  audio:
-    prompt_audio: 0.0
-
-use_text_training: True
-fuser:
-  sum: []
-  prepend: [ description, prompt_audio, type_info ] # this order is the SAME with the input concatenation order
-
-conditioners:
-  prompt_audio:
-    model: qt_embedding
-    qt_embedding:
-      code_size: 16384
-      code_depth: 3
-      max_len: ${eval:${prompt_len}*${audio_tokenizer_frame_rate}+2} # 25*10+2+1
-  description:
-    model: QwTokenizer
-    QwTokenizer:
-      token_path: third_party/Qwen2-7B
-      max_len: 300
-      add_token_list: ${load_yaml:conf/vocab.yaml}
-  type_info:
-    model: QwTextTokenizer
-    QwTextTokenizer:
-      token_path: third_party/Qwen2-7B
-      max_len: 50

generate.py

@@ -12,6 +12,7 @@ from codeclm.trainer.codec_song_pl import CodecLM_PL
 from codeclm.models import CodecLM
 from third_party.demucs.models.pretrained import get_model_from_yaml
 
+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:
@@ -58,21 +59,25 @@ class Separator:
         return full_audio, vocal_audio, bgm_audio
 
 
-def main_sep():
-    torch.backends.cudnn.enabled = False #taiji的某些傻呗node会报奇奇怪怪的错
+
+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)))
-    cfg = OmegaConf.load(sys.argv[1])
-    save_dir = sys.argv[2]
-    input_jsonl = sys.argv[3]
-    sidx = sys.argv[4]
+    np.random.seed(int(time.time()))    
+    ckpt_path = sys.argv[1]
+    input_jsonl = sys.argv[2]
+    save_dir = sys.argv[3]
+    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
     
     # Define model or load pretrained model
-    model_light = CodecLM_PL(cfg)
+    model_light = CodecLM_PL(cfg, ckpt_path)
 
     model_light = model_light.eval().cuda()
     model_light.audiolm.cfg = cfg
@@ -83,9 +88,10 @@ def main_sep():
         seperate_tokenizer = model_light.seperate_tokenizer,
     )
     separator = Separator()
-    
+    auto_prompt = torch.load('ckpt/prompt.pt')
+    merge_prompt = [item for sublist in auto_prompt.values() for item in sublist]
     cfg_coef = 1.5 #25
-    temp = 1.0
+    temp = 0.9
     top_k = 50
     top_p = 0.0
     record_tokens = True
@@ -93,7 +99,7 @@ def main_sep():
 
     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 + "/token", exist_ok=True)
+    os.makedirs(save_dir, exist_ok=True)
     os.makedirs(save_dir + "/audios", exist_ok=True)
     os.makedirs(save_dir + "/jsonl", exist_ok=True)
 
@@ -103,43 +109,58 @@ def main_sep():
     new_items = []
     for line in lines:
         item = json.loads(line)
-        target_name = f"{save_dir}/token/{item['idx']}_s{sidx}.npy"
-        target_wav_name = f"{save_dir}/audios/{item['idx']}_s{sidx}.flac"
-        descriptions = item["descriptions"]
+        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
         lyric = item["gt_lyric"]
-        
-        start_time = time.time()
-        pmt_wav, vocal_wav, bgm_wav = separator.run(item['prompt_audio_path'])
+        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
+        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
+            
         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,
         }
-
-        mid_time = time.time()
+        start_time = time.time()
         with torch.autocast(device_type="cuda", dtype=torch.float16):
             tokens = model.generate(**generate_inp, return_tokens=True)
-        end_time = time.time()
-        if tokens.shape[-1] > 3000:
-            tokens = tokens[..., :3000]
+        mid_time = time.time()
             
         with torch.no_grad():
-            wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav)
+            if melody_is_wav:   
+                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav)
+            else:
+                wav_seperate = model.generate_audio(tokens)
+        end_time = time.time()
         torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
-        np.save(target_name, tokens.cpu().squeeze(0).numpy())
-        print(f"process{item['idx']}, demucs cost {mid_time - start_time}s, lm cos {end_time - mid_time}")
+        print(f"process{item['idx']}, lm cost {mid_time - start_time}s, diffusion cost {end_time - mid_time}")
 
-        item["idx"] = f"{item['idx']}_s{sidx}"
-        item["tk_path"] = target_name
+        item["idx"] = f"{item['idx']}"
+        item["wav_path"] = target_wav_name
         new_items.append(item)
     
     src_jsonl_name = os.path.split(input_jsonl)[-1]
-    with open(f"{save_dir}/jsonl/{src_jsonl_name}-s{sidx}.jsonl", "w", encoding='utf-8') as fw:
+    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__":
-    main_sep()

generate.sh

@@ -4,9 +4,7 @@ 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
 
-CFG_FILE=conf/infer.yaml
-JSONL=$1
-SAVE_DIR=$2
-SIDX=0
-DEVICE=0
-OMP_NUM_THREADS=1 CUDA_VISIBLE_DEVICES=$DEVICE python3 generate.py $CFG_FILE $SAVE_DIR $JSONL $SIDX
+CKPT_PATH=$1
+JSONL=$2
+SAVE_DIR=$3
+python3 generate.py $CKPT_PATH $JSONL $SAVE_DIR

levo_inference.py

@@ -18,7 +18,7 @@ from separator import Separator
 
 
 class LeVoInference(torch.nn.Module):
-    def __init__(self, cfg_path):
+    def __init__(self, ckpt_path):
         super().__init__()
 
         torch.backends.cudnn.enabled = False 
@@ -27,12 +27,15 @@ class LeVoInference(torch.nn.Module):
         OmegaConf.register_new_resolver("get_fname", lambda: 'default')
         OmegaConf.register_new_resolver("load_yaml", lambda x: list(OmegaConf.load(x)))
 
+        cfg_path = os.path.join(ckpt_path, 'config.yaml')
+        pt_path = os.path.join(ckpt_path, 'model.pt')
+
         self.cfg = OmegaConf.load(cfg_path)
         self.cfg.mode = 'inference'
         self.max_duration = self.cfg.max_dur
 
         # Define model or load pretrained model
-        model_light = CodecLM_PL(self.cfg)
+        model_light = CodecLM_PL(self.cfg, pt_path)
 
         model_light = model_light.eval().cuda()
         model_light.audiolm.cfg = self.cfg
@@ -63,15 +66,28 @@ class LeVoInference(torch.nn.Module):
 
         self.model.set_generation_params(**self.default_params)
 
-
-    def forward(self, lyric: str, description: str, prompt_audio_path: os.PathLike = None, 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, params = dict()):
         params = {**self.default_params, **params}
         self.model.set_generation_params(**params)
 
-        if prompt_audio_path is None:
-            pmt_wav, vocal_wav, bgm_wav = None, None, None
-        else:
+        if prompt_audio_path is not None:
             pmt_wav, vocal_wav, bgm_wav = self.separator.run(prompt_audio_path)
+            melody_is_wav = True
+        elif genre is not None and auto_prompt_path is not None:
+            auto_prompt = torch.load(auto_prompt_path)
+            if genre == "Auto": 
+                prompt_token = merge_prompt[np.random.randint(0, len(merge_prompt))]
+            else:
+                prompt_token = auto_prompt[genre][np.random.randint(0, len(auto_prompt[genre]))]
+            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
 
         generate_inp = {
             'lyrics': [lyric.replace("  ", " ")],
@@ -79,6 +95,7 @@ class LeVoInference(torch.nn.Module):
             '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):
@@ -91,38 +108,3 @@ class LeVoInference(torch.nn.Module):
             wav_seperate = self.model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav)
 
         return wav_seperate[0]
-
-def build_levo_inference():
-    cfg_path = './conf/infer.yaml'
-    return LeVoInference(cfg_path)
-
-if __name__ == '__main__':
-    import sys
-    import os
-    import time
-    import json
-    import torchaudio
-
-    cfg_path = sys.argv[1]
-    save_dir = sys.argv[2]
-    input_jsonl = sys.argv[3]
-
-    model = LeVoInference(cfg_path)
-    
-    os.makedirs(save_dir + "/audios", exist_ok=True)
-
-    with open(input_jsonl, "r") as fp:
-        lines = fp.readlines()
-
-    for line in lines:
-        item = json.loads(line)
-        target_wav_name = f"{save_dir}/audios/{item['idx']}.flac"
-        descriptions = item["descriptions"]
-        lyric = item["gt_lyric"]
-        prompt_audio_path = item['prompt_audio_path']
-        
-        wav = model(lyric, descriptions, prompt_audio_path)
-
-        torchaudio.save(target_wav_name, wav.cpu().float(), model.cfg.sample_rate)
-
-

sample/description/emotion.txt

@@ -0,0 +1,8 @@
+sad
+emotional
+angry
+happy
+uplifting
+intense
+romantic
+melancholic

sample/description/gender.txt

@@ -0,0 +1,2 @@
+female
+male

sample/description/genre.txt

@@ -0,0 +1,27 @@
+pop
+electronic
+hip hop
+rock
+jazz
+blues
+classical
+rap
+country
+classic rock
+hard rock
+folk
+soul
+dance, electronic
+rockabilly
+dance, dancepop, house, pop
+reggae
+experimental
+dance, pop
+dance, deephouse, electronic
+k-pop
+experimental pop
+pop punk
+rock and roll
+R&B
+varies
+pop rock

sample/description/instrument.txt

@@ -0,0 +1,40 @@
+synthesizer and piano
+piano and drums
+piano and synthesizer
+synthesizer and drums
+piano and strings
+guitar and drums
+guitar and piano
+piano and double bass
+piano and guitar
+acoustic guitar and piano
+acoustic guitar and synthesizer
+synthesizer and guitar
+piano and saxophone
+saxophone and piano
+piano and violin
+electric guitar and drums
+acoustic guitar and drums
+synthesizer
+guitar and fiddle
+guitar and harmonica
+synthesizer and acoustic guitar
+beats
+piano
+acoustic guitar and fiddle
+brass and piano
+bass and drums
+violin
+acoustic guitar and harmonica
+piano and cello
+saxophone and trumpet
+guitar and banjo
+guitar and synthesizer
+saxophone
+violin and piano
+synthesizer and bass
+synthesizer and electric guitar
+electric guitar and piano
+beats and piano
+synthesizer and
+guitar

sample/description/timbre.txt

@@ -0,0 +1,7 @@
+dark
+bright
+warm
+rock
+varies
+soft
+vocal

sample/lyric.jsonl

@@ -1 +0,0 @@
-{"idx": "01_节奏蓝调", "descriptions": "female, dark, pop, sad, piano and drums, the bpm is 125.", "gt_lyric": "[intro-short] ; [verse] 夜晚的街灯闪烁.我漫步在熟悉的角落.回忆像潮水般涌来.你的笑容如此清晰.在心头无法抹去.那些曾经的甜蜜.如今只剩我独自回忆 ; [bridge] 手机屏幕亮起.是你发来的消息.简单的几个字.却让我泪流满面.曾经的拥抱温暖.如今却变得遥远.我多想回到从前.重新拥有你的陪伴 ; [chorus] 回忆的温度还在.你却已不在.我的心被爱填满.却又被思念刺痛.R&B的节奏奏响.我的心却在流浪.没有你的日子.我该如何继续向前 ; [outro-short]", "prompt_audio_path": "sample/prompt.wav"}

sample/lyrics.jsonl

@@ -0,0 +1,4 @@
+{"idx": "sample_01_autoprompt", "gt_lyric": "[intro-short] ; [verse] 雪花舞动在无尽的天际.情缘如同雪花般轻轻逝去.希望与真挚.永不磨灭.你的忧虑.随风而逝 ; [chorus] 我怀抱着守护这片梦境.在这世界中寻找爱与虚幻.苦辣酸甜.我们一起品尝.在雪的光芒中.紧紧相拥 ; [inst-short] ; [verse] 雪花再次在风中飘扬.情愿如同雪花般消失无踪.希望与真挚.永不消失.在痛苦与喧嚣中.你找到解脱 ; [chorus] 我环绕着守护这片梦境.在这世界中感受爱与虚假.苦辣酸甜.我们一起分享.在白银的光芒中.我们同在 ; [outro-short]", "auto_prompt_audio_type": "Auto"}
+{"idx": "sample_01_noprompt", "gt_lyric": "[intro-short] ; [verse] 雪花舞动在无尽的天际.情缘如同雪花般轻轻逝去.希望与真挚.永不磨灭.你的忧虑.随风而逝 ; [chorus] 我怀抱着守护这片梦境.在这世界中寻找爱与虚幻.苦辣酸甜.我们一起品尝.在雪的光芒中.紧紧相拥 ; [inst-short] ; [verse] 雪花再次在风中飘扬.情愿如同雪花般消失无踪.希望与真挚.永不消失.在痛苦与喧嚣中.你找到解脱 ; [chorus] 我环绕着守护这片梦境.在这世界中感受爱与虚假.苦辣酸甜.我们一起分享.在白银的光芒中.我们同在 ; [outro-short]"}
+{"idx": "sample_01_textprompt", "descriptions": "female, dark, pop, sad, piano and drums, the bpm is 125.", "gt_lyric": "[intro-short] ;  [verse] 雪花舞动在无尽的天际.情缘如同雪花般轻轻逝去.希望与真挚.永不磨灭.你的忧虑.随风而逝 ; [chorus] 我怀抱着守护这片梦境.在这世界中寻找爱与虚幻.苦辣酸甜.我们一起品尝.在雪的光芒中.紧紧相拥 ; [inst-short] ; [verse] 雪花再次在风中飘扬.情愿如同雪花般消失无踪.希望与真挚.永不消失.在痛苦与喧嚣中.你找到解脱 ; [chorus] 我环绕着守护这片梦境.在这世界中感受爱与虚假.苦辣酸甜.我们一起分享.在白银的光芒中.我们同在 ; [outro-short]"}
+{"idx": "sample_01_audioprompt", "gt_lyric": "[intro-short] ; [verse] 雪花舞动在无尽的天际.情缘如同雪花般轻轻逝去.希望与真挚.永不磨灭.你的忧虑.随风而逝 ; [chorus] 我怀抱着守护这片梦境.在这世界中寻找爱与虚幻.苦辣酸甜.我们一起品尝.在雪的光芒中.紧紧相拥 ; [inst-short] ; [verse] 雪花再次在风中飘扬.情愿如同雪花般消失无踪.希望与真挚.永不消失.在痛苦与喧嚣中.你找到解脱 ; [chorus] 我环绕着守护这片梦境.在这世界中感受爱与虚假.苦辣酸甜.我们一起分享.在白银的光芒中.我们同在 ; [outro-short]", "prompt_audio_path": "sample/sample_prompt_audio.wav"}

sample/sample_prompt_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2068592b00263f7c0b0f1d82a882d7738730ace3e04f2d889d06ff983ad6d618
+size 3845542