update single track diffusion

app.py

@@ -56,7 +56,7 @@ with open(op.join(APP_DIR, 'conf/vocab.yaml'), 'r', encoding='utf-8') as file:
 
 
 # 模拟歌曲生成函数
-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)):
+def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_coef=None, temperature=None, top_k=None, gen_type="all", progress=gr.Progress(track_tqdm=True)):
     global MODEL
     global STRUCTS
     params = {'cfg_coef':cfg_coef, 'temperature':temperature, 'top_k':top_k}
@@ -105,7 +105,7 @@ def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_co
     progress(0.0, "Start Generation")
     start = time.time()
     
-    audio_data = MODEL(lyric_norm, description, prompt_audio, genre, op.join(APP_DIR, "ckpt/prompt.pt"), params).cpu().permute(1, 0).float().numpy()
+    audio_data = MODEL(lyric_norm, description, prompt_audio, genre, op.join(APP_DIR, "ckpt/prompt.pt"), gen_type, params).cpu().permute(1, 0).float().numpy()
 
     end = time.time()
     
@@ -204,7 +204,9 @@ lyrics
                     interactive=True,
                     elem_id="top_k",
                 )
-            generate_btn = gr.Button("Generate Song", variant="primary")
+            with gr.Row():
+                generate_btn = gr.Button("Generate Song", variant="primary")
+                generate_bgm_btn = gr.Button("Generate Pure Music", variant="primary")
         
         with gr.Column():
             output_audio = gr.Audio(label="Generated Song", type="numpy")
@@ -235,6 +237,11 @@ lyrics
         inputs=[lyric, description, prompt_audio, genre, cfg_coef, temperature, top_k],
         outputs=[output_audio, output_json]
     )
+    generate_bgm_btn.click(
+        fn=generate_song,
+        inputs=[lyric, description, prompt_audio, genre, cfg_coef, temperature, top_k, gr.State("bgm")],
+        outputs=[output_audio, output_json]
+    )
     
 
 # 启动应用

codeclm/models/codeclm.py

@@ -271,13 +271,19 @@ class CodecLM:
         return gen_tokens
 
     @torch.no_grad()
-    def generate_audio(self, gen_tokens: torch.Tensor, prompt=None, vocal_prompt=None, bgm_prompt=None, chunked=False):
+    def generate_audio(self, gen_tokens: torch.Tensor, prompt=None, vocal_prompt=None, bgm_prompt=None, chunked=False, gen_type="all"):
         """Generate Audio from tokens"""
         assert gen_tokens.dim() == 3
         if self.seperate_tokenizer is not None:
             gen_tokens_song = gen_tokens[:, [0], :]
             gen_tokens_vocal = gen_tokens[:, [1], :]
             gen_tokens_bgm = gen_tokens[:, [2], :]
+            if gen_type == "bgm":
+                gen_tokens_vocal = torch.full_like(gen_tokens_vocal, 3142)
+                vocal_prompt = None
+            elif gen_type == "vocal":
+                gen_tokens_bgm = torch.full_like(gen_tokens_bgm, 9670)
+                bgm_prompt = None
             # gen_audio_song = self.audiotokenizer.decode(gen_tokens_song, prompt)
             gen_audio_seperate = self.seperate_tokenizer.decode([gen_tokens_vocal, gen_tokens_bgm], vocal_prompt, bgm_prompt, chunked=chunked)
             return gen_audio_seperate

generate.py

@@ -70,6 +70,7 @@ if __name__ == "__main__":
     ckpt_path = sys.argv[1]
     input_jsonl = sys.argv[2]
     save_dir = sys.argv[3]
+    gen_type = sys.argv[4] if len(sys.argv) > 4 else "all"
     cfg_path = os.path.join(ckpt_path, 'config.yaml')
     ckpt_path = os.path.join(ckpt_path, 'model.pt')
     cfg = OmegaConf.load(cfg_path)
@@ -146,15 +147,15 @@ if __name__ == "__main__":
         with torch.autocast(device_type="cuda", dtype=torch.float16):
             tokens = model.generate(**generate_inp, return_tokens=True)
         mid_time = time.time()
-            
+
         with torch.no_grad():
             if melody_is_wav:   
-                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav)
+                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type)
             else:
-                wav_seperate = model.generate_audio(tokens)
+                wav_seperate = model.generate_audio(tokens, gen_type=gen_type)
         end_time = time.time()
         torchaudio.save(target_wav_name, wav_seperate[0].cpu().float(), cfg.sample_rate)
-        print(f"process{item['idx']}, lm cost {mid_time - start_time}s, diffusion cost {end_time - mid_time}")
+        print(f"process{item['idx']} {gen_type}, lm cost {mid_time - start_time}s, diffusion cost {end_time - mid_time}")
 
         item["idx"] = f"{item['idx']}"
         item["wav_path"] = target_wav_name

generate.sh

@@ -7,4 +7,5 @@ export PYTHONPATH="$(pwd)/codeclm/tokenizer/":"$(pwd)":"$(pwd)/codeclm/tokenizer
 CKPT_PATH=$1
 JSONL=$2
 SAVE_DIR=$3
-python3 generate.py $CKPT_PATH $JSONL $SAVE_DIR
+GEN_TYEP=$4
+python3 generate.py $CKPT_PATH $JSONL $SAVE_DIR $GEN_TYEP

generate_lowmem.py

@@ -71,6 +71,7 @@ if __name__ == "__main__":
     ckpt_path = sys.argv[1]
     input_jsonl = sys.argv[2]
     save_dir = sys.argv[3]
+    gen_type = sys.argv[4] if len(sys.argv) > 4 else "all"
     cfg_path = os.path.join(ckpt_path, 'config.yaml')
     ckpt_path = os.path.join(ckpt_path, 'model.pt')
     cfg = OmegaConf.load(cfg_path)
@@ -220,12 +221,12 @@ if __name__ == "__main__":
     for item in new_items:
         with torch.no_grad():
             if 'raw_pmt_wav' in item:   
-                wav_seperate = model.generate_audio(item['tokens'], item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True)
+                wav_seperate = model.generate_audio(item['tokens'], item['raw_pmt_wav'], item['raw_vocal_wav'], item['raw_bgm_wav'], chunked=True, gen_type=gen_type)
                 del item['raw_pmt_wav']
                 del item['raw_vocal_wav']
                 del item['raw_bgm_wav']
             else:
-                wav_seperate = model.generate_audio(item['tokens'], chunked=True)
+                wav_seperate = model.generate_audio(item['tokens'], chunked=True, gen_type=gen_type)
         torchaudio.save(item['wav_path'], wav_seperate[0].cpu().float(), cfg.sample_rate)
         del item['tokens']
         del item['pmt_wav']

generate_lowmem.sh

@@ -7,4 +7,5 @@ export PYTHONPATH="$(pwd)/codeclm/tokenizer/":"$(pwd)":"$(pwd)/codeclm/tokenizer
 CKPT_PATH=$1
 JSONL=$2
 SAVE_DIR=$3
-python3 generate_lowmem.py $CKPT_PATH $JSONL $SAVE_DIR
+GEN_TYEP=$4
+python3 generate_lowmem.py $CKPT_PATH $JSONL $SAVE_DIR $GEN_TYEP

levo_inference.py

@@ -67,7 +67,7 @@ class LeVoInference(torch.nn.Module):
 
         self.model.set_generation_params(**self.default_params)
 
-    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, params = dict()):
+    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "all", params = dict()):
         params = {**self.default_params, **params}
         self.model.set_generation_params(**params)
 
@@ -105,8 +105,8 @@ class LeVoInference(torch.nn.Module):
             
         with torch.no_grad():
             if melody_is_wav:
-                wav_seperate = self.model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav)
+                wav_seperate = self.model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type)
             else:
-                wav_seperate = self.model.generate_audio(tokens)
+                wav_seperate = self.model.generate_audio(tokens, gen_type=gen_type)
 
         return wav_seperate[0]

tools/gradio/app.py

@@ -49,7 +49,7 @@ with open(op.join(APP_DIR, 'conf/vocab.yaml'), 'r', encoding='utf-8') as file:
     STRUCTS = yaml.safe_load(file)
 
 
-def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_coef=None, temperature=None, top_k=None, progress=gr.Progress(track_tqdm=True)):
+def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_coef=None, temperature=None, top_k=None, gen_type="all", progress=gr.Progress(track_tqdm=True)):
     global MODEL
     global STRUCTS
     params = {'cfg_coef':cfg_coef, 'temperature':temperature, 'top_k':top_k}
@@ -98,7 +98,7 @@ def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_co
     progress(0.0, "Start Generation")
     start = time.time()
     
-    audio_data = MODEL(lyric_norm, description, prompt_audio, genre, op.join(APP_DIR, "ckpt/prompt.pt"), params).cpu().permute(1, 0).float().numpy()
+    audio_data = MODEL(lyric_norm, description, prompt_audio, genre, op.join(APP_DIR, "ckpt/prompt.pt"), gen_type, params).cpu().permute(1, 0).float().numpy()
 
     end = time.time()
     
@@ -119,7 +119,7 @@ def generate_song(lyric, description=None, prompt_audio=None, genre=None, cfg_co
 # 创建Gradio界面
 with gr.Blocks(title="SongGeneration Demo Space") as demo:
     gr.Markdown("# 🎵 SongGeneration 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.")
+    gr.Markdown("Demo interface for the song generation model. Provide a lyrics, and optionally an audio or text prompt, to generate a custom song. The code is in [GIT](https://github.com/tencent-ailab/SongGeneration)")
     
     with gr.Row():
         with gr.Column():
@@ -197,7 +197,9 @@ lyrics
                     interactive=True,
                     elem_id="top_k",
                 )
-            generate_btn = gr.Button("Generate Song", variant="primary")
+            with gr.Row():
+                generate_btn = gr.Button("Generate Song", variant="primary")
+                generate_bgm_btn = gr.Button("Generate Pure Music", variant="primary")
         
         with gr.Column():
             output_audio = gr.Audio(label="Generated Song", type="numpy")
@@ -228,6 +230,11 @@ lyrics
         inputs=[lyric, description, prompt_audio, genre, cfg_coef, temperature, top_k],
         outputs=[output_audio, output_json]
     )
+    generate_bgm_btn.click(
+        fn=generate_song,
+        inputs=[lyric, description, prompt_audio, genre, cfg_coef, temperature, top_k, gr.State("bgm")],
+        outputs=[output_audio, output_json]
+    )
     
 
 # 启动应用

tools/gradio/levo_inference.py

@@ -62,7 +62,7 @@ class LeVoInference(torch.nn.Module):
 
         self.model.set_generation_params(**self.default_params)
 
-    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, params = dict()):
+    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "all", params = dict()):
         params = {**self.default_params, **params}
         self.model.set_generation_params(**params)
 
@@ -97,14 +97,11 @@ class LeVoInference(torch.nn.Module):
 
         with torch.autocast(device_type="cuda", dtype=torch.float16):
             tokens = self.model.generate(**generate_inp, return_tokens=True)
-
-        if tokens.shape[-1] > 3000:
-            tokens = tokens[..., :3000]
             
         with torch.no_grad():
             if melody_is_wav:
-                wav_seperate = self.model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav)
+                wav_seperate = self.model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type)
             else:
-                wav_seperate = self.model.generate_audio(tokens)
+                wav_seperate = self.model.generate_audio(tokens, gen_type=gen_type)
 
         return wav_seperate[0]

tools/gradio/levo_inference_lowmem.py

@@ -40,7 +40,7 @@ class LeVoInference(torch.nn.Module):
         )
 
 
-    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, params = dict()):
+    def forward(self, lyric: str, description: str = None, prompt_audio_path: os.PathLike = None, genre: str = None, auto_prompt_path: os.PathLike = None, gen_type: str = "all", params = dict()):
         if prompt_audio_path is not None and os.path.exists(prompt_audio_path):
             separator = Separator()
             audio_tokenizer = builders.get_audio_tokenizer_model(self.cfg.audio_tokenizer_checkpoint, self.cfg)
@@ -112,15 +112,12 @@ class LeVoInference(torch.nn.Module):
             max_duration = self.max_duration,
             seperate_tokenizer = seperate_tokenizer,
         )
-
-        if tokens.shape[-1] > 3000:
-            tokens = tokens[..., :3000]
             
         with torch.no_grad():
             if melody_is_wav:
-                wav_seperate = model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav)
+                wav_seperate = self.model.generate_audio(tokens, pmt_wav, vocal_wav, bgm_wav, gen_type=gen_type)
             else:
-                wav_seperate = model.generate_audio(tokens)
+                wav_seperate = self.model.generate_audio(tokens, gen_type=gen_type)
 
         del seperate_tokenizer
         del model