diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index f6c8eb49..ec6f6239 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -633,8 +633,11 @@ namespace MiniMaxH3VAE { auto plan = make_vae_temporal_tile_plan(input.shape()[2], {17, 0}); auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& chunk, const VAETemporalTile& 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()) { return {}; } @@ -695,13 +698,16 @@ namespace MiniMaxH3VAE { {static_cast(tokens_per_chunk + token_overlap), static_cast(token_overlap)}); GGML_ASSERT(plan.tiles.size() == static_cast(num_chunks)); auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& 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, chunk, tiling, true, circular_x, circular_y, - silent); + silent, + false); if (decoded.empty()) { return sd::Tensor(); } @@ -728,6 +734,7 @@ namespace MiniMaxH3VAE { } return first; }); + runner_end(); if (result.empty()) { return {}; } diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index 5d6e4eee..d5ac6f58 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -241,6 +241,15 @@ public: sd_tiling_params_t tiling_params, bool circular_x = false, bool circular_y = false) { + return encode(n_threads, x, tiling_params, circular_x, circular_y, true); + } + + sd::Tensor encode(int n_threads, + const sd::Tensor& x, + sd_tiling_params_t tiling_params, + bool circular_x, + bool circular_y, + bool auto_end) { int64_t t0 = ggml_time_ms(); tiling_params = resolve_tiling_params(tiling_params); sd::Tensor input = x; @@ -280,7 +289,9 @@ public: tiling_params); } - runner_end(); + if (auto_end) { + runner_end(); + } if (output.empty()) { LOG_ERROR("vae encode compute failed"); @@ -298,6 +309,17 @@ public: bool circular_x = false, bool circular_y = false, bool silent = false) { + return decode(n_threads, x, tiling_params, decode_video, circular_x, circular_y, silent, true); + } + + sd::Tensor decode(int n_threads, + const sd::Tensor& 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(); tiling_params = resolve_tiling_params(tiling_params); sd::Tensor input = x; @@ -338,7 +360,9 @@ public: tiling_params); } - runner_end(); + if (auto_end) { + runner_end(); + } if (output.empty()) { LOG_ERROR("vae decode compute failed");