From de433c40d3c01a3af320183341a79681cb63e690 Mon Sep 17 00:00:00 2001
From: Towsty <ianjmcarr@gmail.com>
Date: Thu, 3 Sep 2026 17:46:54 -0500
Subject: [PATCH] Lock YuE to a 16GB Comfy-safe recipe instead of chasing
 single knobs.

Keep mmgp profile 2 (no quanto), node-legal token floor, two Stage-A segments, and document why profiles 3-5 blow up under InferenceMode.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 server/assets/workflow_yue.json |  4 ++--
 server/utils/musicWorkflow.ts   | 17 ++++++++++++++++-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/server/assets/workflow_yue.json b/server/assets/workflow_yue.json
index 5e3b00b..efc3f0f 100644
--- a/server/assets/workflow_yue.json
+++ b/server/assets/workflow_yue.json
@@ -7,7 +7,7 @@
       "use_mmgp": true,
       "stage1_cache_size": 8192,
       "exllamav2_cache_mode": "FP16",
-      "mmgp_profile": 4
+      "mmgp_profile": 2
     },
     "class_type": "YUE_Stage_A_Loader",
     "_meta": {
@@ -22,7 +22,7 @@
       "run_n_segment": 2,
       "repetition_penalty": 1.1,
       "prompt_start_time": 0,
-      "prompt_end_time": 30,
+      "prompt_end_time": 45,
       "max_new_tokens": 2944,
       "use_dual_tracks_prompt": false,
       "use_audio_prompt": false,
diff --git a/server/utils/musicWorkflow.ts b/server/utils/musicWorkflow.ts
index db2d3cc..01d0f80 100644
--- a/server/utils/musicWorkflow.ts
+++ b/server/utils/musicWorkflow.ts
@@ -116,15 +116,26 @@ function buildYueWorkflow(params: MusicWorkflowParams): WorkflowGraph {
   const prefix = params.filenamePrefix || musicFilenamePrefix()
   const tags = params.tags
   const lyrics = params.lyrics || '[instrumental]\n'
-  // 5080 16GB OOMs past ~3 long Stage-A segments on fp16+mmgp. Cap hard.
+  // Beast 5080 16GB recipe (keep in sync with workflow_yue.json):
+  // - mmgp profile 2 only (3–5 enable quanto and crash under Comfy InferenceMode)
+  // - max 2 Stage-A segments (3+ OOMs)
+  // - max_new_tokens floor is 2944 (node min); do not go lower
+  // - Triton/torch.compile off on Windows
   const sections = lyrics.split(/\[[^\]]+\]/g).filter(part => part.trim()).length || 2
   const segments = Math.max(1, Math.min(2, sections))
   for (const node of Object.values(graph)) {
     if (node.class_type === 'YUE_Stage_A_Loader') {
       node.inputs.stage_A_repo = yueStageARepo()
+      node.inputs.quantization_model = 'fp16'
+      node.inputs.use_mmgp = true
+      node.inputs.stage1_cache_size = 8192
+      node.inputs.mmgp_profile = 2
     }
     if (node.class_type === 'YUE_Stage_B_Loader') {
       node.inputs.stage_B_repo = yueStageBRepo()
+      node.inputs.use_mmgp = true
+      node.inputs.stage2_cache_size = 4096
+      node.inputs.stage2_batch_size = 1
     }
     if (node.class_type === 'YUE_Stage_A_Sampler') {
       node.inputs.genres_prompt = tags.replace(/,/g, ' ').replace(/\s+/g, ' ').trim()
@@ -133,6 +144,10 @@ function buildYueWorkflow(params: MusicWorkflowParams): WorkflowGraph {
       node.inputs.prompt_end_time = Math.min(params.duration, 45)
       node.inputs.run_n_segment = segments
       node.inputs.max_new_tokens = 2944
+      node.inputs.use_dual_tracks_prompt = false
+      node.inputs.use_audio_prompt = false
+      node.inputs.offload_model = true
+      node.inputs.stage1_no_guidance = true
     }
   }
   applyPrefix(graph, prefix)