YuE2 VAE: a grown decoder (YuE2-VAE-RS x3) loses 210 tensors silently under strict=False; a 35-line patch is ready

Open
  • Is something wrong?
#17,697 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
43/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
backend, machine-learning

Research direction

Read comfy/ldm/audio/autoencoder.py and the YuE2 branch in comfy/sd.py; the proposed patch adds configurable decoder units and selects float32 for the grown decoder. First inspect how Standard decoder shapes and checkpoint keys are loaded, then verify the new checkpoint has no missing or unexpected keys and matches the reference decode. The issue also raises unresolved TF32 and tiled decoding behavior, which need maintainer guidance.

Written by the indexing model from the issue text.

Description

Hello, and thank you for native YuE2 (#16250).

We are Ruach Studio (https://github.com/igrbible/Ruach_Studio, AGPL-3.0), an open local studio built on YuE2-3B.

We published YuE2-VAE-RS x3 v1 and v2: m-a-p's Standard decoder grown to 200.6 M parameters (three residual units at the latent rate, three more in each upsampling block, Standard's own operators), trained on music; it decodes the same latents. Weights CC BY-NC 4.0, float32 only; code Apache-2.0: https://huggingface.co/goldhub/YuE2-RS-Finetunes/tree/main/VAE-RS

What stock ComfyUI does with it today, in the YuE2 branch of comfy/sd.py: AudioOobleckVAE is built with Standard's shape and filled with load_state_dict(..., strict=False), so a checkpoint carrying x3 under vae.decoder.* loses the 210 tensors of the new units without an error. Measured: 23 % of the signal's RMS away from the real x3. And working_dtypes puts float16 first, while x3 is float32 only (guide §5).

What we offer: VAE-RS/kit/comfyui-native.patch (against master of November 2026, 29 lines added, 6 changed). DecoderBlock(extra_units=0) and OobleckDecoder(pre_units=0, extra_units=0) in comfy/ldm/audio/autoencoder.py, zero by default so Standard is untouched; in comfy/sd.py the YuE2 branch counts decoder.pre.<i> and the units of decoder.layers.1 and pins working_dtypes = [torch.float32] for the grown decoder. Tested: no missing or unexpected key, and a decode bit for bit equal to our reference.

Two points we would like your view on before a PR: TF32 (cuDNN convolutions run TF32 by default on Ampere and newer, which misses our verify.py tolerance of 1e-4 tenfold), and tiling (VAEDecodeAudioTiled blends a 64-frame overlap, while x3 needs 60 frames of cropped context each side to equal the whole decode; Standard needs 12).

Integration guide: https://huggingface.co/goldhub/YuE2-RS-Finetunes/blob/main/docs/VAE-RS-x3-integration.md
Study: https://ruachstudio.art/#vae-study · announcement: https://github.com/igrbible/Ruach_Studio/discussions/6

We are glad to open the PR (rebased on #16955 if it lands first, since both touch comfy/sd.py) and to answer questions.

Viktor Zhuromskyy and Claude (Anthropic), Ruach Studio

Dominant language
Python
Stars
146k
Forks
17.4k
Avg merge
19h 44m
Merged PRs (30d)
178

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from Comfy-Org/ComfyUI

All issues in Comfy-Org/ComfyUI

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.