refactor: unify runner lifecycles and weight residency (#1940)

This commit is contained in:
leejet authored and GitHub committed 2026-09-06 22:30:45 +08:00
1 parent 6c57cc3b38
commit 462d675018
78 files changed
+2963 -2051

No files matched your search

+1 -1
View File
@@ -744,7 +744,7 @@ struct AutoEncoderKL : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(z, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
}
sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {
+1 -3
View File
@@ -827,9 +827,7 @@ namespace Hunyuan {
};
auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
n_threads,
true,
true,
true),
false),
graph_input.dim());
if (!output.empty() && input.dim() == 4) {
output.squeeze_(2);
+1 -1
View File
@@ -1042,7 +1042,7 @@ namespace LTXV {
ggml_build_forward_expand(gf, waveform);
return gf;
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), 4);
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
int64_t t1 = ggml_time_ms();
LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
return result;
+3 -5
View File
@@ -1334,7 +1334,6 @@ struct LTXVideoVAE : public VAE {
(int)plan.tiles.size());
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& z_chunk, const VAETemporalTile& tile) {
LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d",
@@ -1349,12 +1348,11 @@ struct LTXVideoVAE : public VAE {
static_cast<int>(tile.start),
tile.overlap);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
expected_dim);
});
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
return output;
}
@@ -1407,7 +1405,7 @@ struct LTXVideoVAE : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
if (result.empty()) {
return {};
}
@@ -1420,7 +1418,7 @@ struct LTXVideoVAE : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_latent_statistics_graph(z, normalize);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
static_cast<size_t>(z.dim()));
}
+1 -1
View File
@@ -490,7 +490,7 @@ namespace MageVAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), input.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), input.dim());
}
int get_encoder_output_channels(int input_channels) override {
+2 -2
View File
@@ -480,7 +480,7 @@ namespace MiniMaxH3 {
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
GGMLRunner::compute<float>(get_graph, n_threads, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
@@ -500,7 +500,7 @@ namespace MiniMaxH3 {
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
GGMLRunner::compute<float>(get_graph, n_threads, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",
-2
View File
@@ -793,8 +793,6 @@ namespace MiniMaxH3VAE {
return restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph,
n_threads,
false,
false,
false),
5);
}
+2 -2
View File
@@ -787,7 +787,7 @@ struct TinyImageAutoEncoder : public VAE {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
}
};
@@ -872,7 +872,7 @@ struct TinyVideoAutoEncoder : public VAE {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
}
};
+2 -2
View File
@@ -251,7 +251,7 @@ public:
tiling_params);
}
runner_done();
runner_end();
if (output.empty()) {
LOG_ERROR("vae encode compute failed");
@@ -305,7 +305,7 @@ public:
tiling_params);
}
runner_done();
runner_end();
if (output.empty()) {
LOG_ERROR("vae decode compute failed");
+2 -4
View File
@@ -1415,7 +1415,6 @@ namespace WAN {
(int)plan.tiles.size());
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
ae.clear_cache();
auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& input_tile, const VAETemporalTile& tile) {
@@ -1428,12 +1427,11 @@ namespace WAN {
return build_temporal_tile_graph(input_tile, static_cast<int>(tile.start));
};
return restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
GGMLRunner::compute<float>(get_graph, n_threads, false),
static_cast<size_t>(input.dim()));
});
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
ae.clear_cache();
return output;
}
@@ -1448,7 +1446,7 @@ namespace WAN {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input.empty() ? z : input, decode_graph);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
input.empty() ? z.dim() : input.dim());
if (!result.empty() && z.dim() == 4) {
result.squeeze_(2);