Running YuE2 and SheetSage2 on <6GB VRAM

[xorbitr]

Hi everyone,

For those of you who are concerned with high VRAM usage in the these models, I made a mod which allows them to run below 6GB VRAM. It is implemented as a session option, so the old behavior is preserved by default.
I am new to GitHub, so I am just leaving this here. If anybody is interested, I am ready to share the code in a way that you find appropriate.

Below are VRAM usage graphs from tests completed on 1060 6GB.

YuE2 generation config:
Q8_0 main model, FP32 VAE
Peak VRAM usage: 5.6GB, average below 4GB

SheetSage2 ABC decoding, 3.5 minute input:
Original weights
Peak VRAM usage: 3.3GB

A small note on how it is achieved. First of all, by removing some potential duplicate allocations of VRAM for graphs, next by separating the processing stages (AR/NAR/VAE in YuE2 and encode/decode in SheetSage2). When min VRAM usage mode is turned on, the AR model prepares all data for NAR stage without loading the next model (currently they work in turns), also between all stages the weights are unloaded and all allocations related to the previous stage are released.
With SheetSage2, since it spills beyond 6GB anyway at encode stage (unless you use FP16 weight format), I additionally tested flash attention instead of explicit, and on all my test songs the output ABC files were the same.

How to activate the minimum VRAM mode: supply session option "yue2.min_vram_usage=true" or "sheetsage2.min_vram_usage=true" respectively.

The diffs in attachments are made versus v0.9.2-drama3-base-small.

Thanks for reading!

YuE2 VRAM trace:

YuE2 WAV result:

YuE2 generation log:
yue2_generation_log.txt

SheetSage2 VRAM trace:

SheetSage2 compare ABC against reference:

changes_versus_v0.9.2.tar.gz