From 909fe87da2b0a3db0bc60e1a0bf9a56dca2ee96d Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Sat, 26 Jul 2025 15:29:04 +0800
Subject: [PATCH] Update codeclm.py

---
 SongGeneration/codeclm/models/codeclm.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/SongGeneration/codeclm/models/codeclm.py b/SongGeneration/codeclm/models/codeclm.py
index b8a364f..73c075b 100644
--- a/SongGeneration/codeclm/models/codeclm.py
+++ b/SongGeneration/codeclm/models/codeclm.py
@@ -281,10 +281,12 @@ class CodecLM:
             gen_tokens_bgm = gen_tokens[:, [2], :]
             if gen_type == "bgm":
                 gen_tokens_vocal = torch.full_like(gen_tokens_vocal, 3142)
-                vocal_prompt = None
+                if vocal_prompt is not None:
+                    vocal_prompt = torch.zeros_like(vocal_prompt)
             elif gen_type == "vocal":
                 gen_tokens_bgm = torch.full_like(gen_tokens_bgm, 9670)
-                bgm_prompt = None
+                if bgm_prompt is not None:
+                    bgm_prompt = torch.zeros_like(bgm_prompt)
             # 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