From 0046311633bf838adbd4d30f99f9a4f0b6bd687f Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Sun, 18 Jan 2026 17:38:56 +0800
Subject: [PATCH] add 1.5 test
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

我没时间测试，先加上
---
 generate.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/generate.py b/generate.py
index 68fc020..eb6b47a 100644
--- a/generate.py
+++ b/generate.py
@@ -91,6 +91,17 @@ def build_model(Weigths_Path,infer_model_path):
     audiolm = builders.get_lm_model(cfg)
     checkpoint = torch.load(infer_model_path, map_location='cpu',weights_only=False)
     audiolm_state_dict = {k.replace('audiolm.', ''): v for k, v in checkpoint.items() if k.startswith('audiolm')}
+    
+    #### @tuolaku  https://github.com/smthemex/ComfyUI_SongGeneration/issues/37 #####
+    # add 1.5 support，test。。。。
+    key = "condition_provider.conditioners.type_info.output_proj.weight"
+    expected_vocab_size = 151646
+    if key in audiolm_state_dict:
+        weight = audiolm_state_dict[key]
+        if weight.size(0) > expected_vocab_size:
+            print(f"[SongGeneration] Trimming {key} from {weight.size(0)} to {expected_vocab_size}")
+            audiolm_state_dict[key] = weight[:expected_vocab_size, :]
+    #####
     audiolm.load_state_dict(audiolm_state_dict, strict=False)
     audiolm = audiolm.eval()
     #audiolm = audiolm.cuda().to(torch.float16)