fix: keep MiniMax-H3 VAE weights resident across temporal chunks (#2103)

This commit is contained in:
losewayy authored and GitHub committed 2026-10-06 19:40:48 +08:00
1 parent 1da55ebde8
commit 0963357214
2 files changed
+35 -4

No files matched your search

+9 -2
View File
@@ -633,8 +633,11 @@ namespace MiniMaxH3VAE {
auto plan = make_vae_temporal_tile_plan(input.shape()[2], {17, 0}); auto plan = make_vae_temporal_tile_plan(input.shape()[2], {17, 0});
auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) { auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) {
SD_UNUSED(tile); SD_UNUSED(tile);
return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y); // keep the runner alive across chunks; ending it here would
// evict and reload the encoder weights every chunk
return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y, false);
}); });
runner_end();
if (result.empty()) { if (result.empty()) {
return {}; return {};
} }
@@ -695,13 +698,16 @@ namespace MiniMaxH3VAE {
{static_cast<int>(tokens_per_chunk + token_overlap), static_cast<int>(token_overlap)}); {static_cast<int>(tokens_per_chunk + token_overlap), static_cast<int>(token_overlap)});
GGML_ASSERT(plan.tiles.size() == static_cast<size_t>(num_chunks)); GGML_ASSERT(plan.tiles.size() == static_cast<size_t>(num_chunks));
auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) { auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) {
// keep the runner alive across chunks; ending it here would
// evict and reload the decoder weights every chunk
auto decoded = VAE::decode(n_threads, auto decoded = VAE::decode(n_threads,
chunk, chunk,
tiling, tiling,
true, true,
circular_x, circular_x,
circular_y, circular_y,
silent); silent,
false);
if (decoded.empty()) { if (decoded.empty()) {
return sd::Tensor<float>(); return sd::Tensor<float>();
} }
@@ -728,6 +734,7 @@ namespace MiniMaxH3VAE {
} }
return first; return first;
}); });
runner_end();
if (result.empty()) { if (result.empty()) {
return {}; return {};
} }
+26 -2
View File
@@ -241,6 +241,15 @@ public:
sd_tiling_params_t tiling_params, sd_tiling_params_t tiling_params,
bool circular_x = false, bool circular_x = false,
bool circular_y = false) { bool circular_y = false) {
return encode(n_threads, x, tiling_params, circular_x, circular_y, true);
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x,
bool circular_y,
bool auto_end) {
int64_t t0 = ggml_time_ms(); int64_t t0 = ggml_time_ms();
tiling_params = resolve_tiling_params(tiling_params); tiling_params = resolve_tiling_params(tiling_params);
sd::Tensor<float> input = x; sd::Tensor<float> input = x;
@@ -280,7 +289,9 @@ public:
tiling_params); tiling_params);
} }
runner_end(); if (auto_end) {
runner_end();
}
if (output.empty()) { if (output.empty()) {
LOG_ERROR("vae encode compute failed"); LOG_ERROR("vae encode compute failed");
@@ -298,6 +309,17 @@ public:
bool circular_x = false, bool circular_x = false,
bool circular_y = false, bool circular_y = false,
bool silent = false) { bool silent = false) {
return decode(n_threads, x, tiling_params, decode_video, circular_x, circular_y, silent, true);
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video,
bool circular_x,
bool circular_y,
bool silent,
bool auto_end) {
int64_t t0 = ggml_time_ms(); int64_t t0 = ggml_time_ms();
tiling_params = resolve_tiling_params(tiling_params); tiling_params = resolve_tiling_params(tiling_params);
sd::Tensor<float> input = x; sd::Tensor<float> input = x;
@@ -338,7 +360,9 @@ public:
tiling_params); tiling_params);
} }
runner_end(); if (auto_end) {
runner_end();
}
if (output.empty()) { if (output.empty()) {
LOG_ERROR("vae decode compute failed"); LOG_ERROR("vae decode compute failed");