Nodes/ComfyUI_YuE/YUE_Stage_B_Loader
ComfyUI Node

YUE_Stage_B_Loader

The quiet middle of the YuE chain

By smthemex·Created 2 years ago·Updated 2 years ago· 190
YUE_Stage_B_Loader
  • info
  • model
stage_B_repom-a-p/YuE-s2-1B-general
stage2_cache_size8192
stage2_batch_size2
exllamav2_cache_mode
use_mmgptrue

YUE_Stage_B_Loader is the third node in the YuE pipeline and the one you'll spend the least time on - which is exactly how a loader should behave. Stage A wrote your song as codec tokens; this node loads the Stage B upsampler that turns those tokens into actual audio. It has one genuinely unusual feature that trips people up: it takes a data wire from the sampler, not from the other loader.

The info input is the whole game

The info input (type quantization_model) has to come from YUE_Stage_A_Sampler's info output. That's not a typo in the example workflow. The sampler is the node that knows how Stage A was actually run, because it reads the quantization choice off the loaded model, and Stage B needs to match: if you generated Stage A in exllamav2 Q8, Stage B should load the same way or you're mixing mismatched formats. So this loader consumes a small dict carrying quantization_model and mmgp_profile and uses it to decide how to load YuE-s2-1B.

That's an odd design for ComfyUI - most loaders are self-contained - and it's easy to miss. Both loaders output a model wire, so wiring them together directly looks natural, but then info sits empty and the node errors or loads blind. Follow the example graph: A Loader → A Sampler → B Loader → B Sampler.

Inputs you'll actually set

  • stage_B_repo - default m-a-p/YuE-s2-1B-general, the 3.65G upsampler. Leave it alone unless you're using one of the community quantized Stage B repos the README links for sub-16G cards.
  • stage2_batch_size (default 2) and stage2_cache_size (default 8192) - the knobs you touch. Bigger batch means faster but thirstier; drop the batch if Stage B OOMs.
  • exllamav2_cache_mode (FP16/Q8/Q6/Q4) - only relevant if your Stage A run used exllamav2.
  • use_mmgp - on by default. Same optional-dependency situation as everywhere in this pack: it only does anything if mmgp is actually installed and you want the CPU/VRAM offload trick.

Output

Just model (MODEL_YUE_B), which feeds YUE_Stage_B_Sampler. Nothing to preview here - the whole node is a handoff.

Install and models

The shared pack install applies - clone into custom_nodes, install requirements, restart:

cd ComfyUI/custom_nodes
git clone https://github.com/smthemex/ComfyUI_YuE
cd ComfyUI_YuE
pip install -r requirements.txt

(ComfyUI Manager → "ComfyUI_YuE" works too.) The Stage B model auto-downloads from Hugging Face on first load. You also need the two vocoder decoder files - decoder_131000.pth and decoder_151000.pth - plus ckpt_00360000.pth in ComfyUI/models/yue; those get consumed one node downstream, not here. If your VRAM is under 16G, the README points at exllamav2-quantized Stage B repos, and exllamav2 mode here is how you'd use them.

Gotchas

  • The info wire is required. Empty info means the loader doesn't know how to load. Wire YUE_Stage_A_Sampler → this node.
  • mmgp only if you use it. It's not in requirements.txt by design; run pip install mmgp only if you're flipping use_mmgp on, otherwise skip it.
  • Don't rename the install folder. Parts of this pack assume the literal custom_nodes/ComfyUI_YuE path (the semantic model is loaded from a hardcoded location). Clone it as-is.
CategoryYUE

Inputs (6)

NameTypeDefaultDescription
infoquantization_model
stage_B_repoSTRINGm-a-p/YuE-s2-1B-general
stage2_cache_sizeINT81924096–2147483647
stage2_batch_sizeINT21–64
exllamav2_cache_modeCOMBO4 options: FP16, Q8, Q6, Q4
use_mmgpBOOLEANtrue

Outputs (1)

NameTypeDescription
modelMODEL_YUE_B