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

+50 -142
View File
@@ -142,14 +142,13 @@ public:
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
virtual void set_flash_attention_enabled(bool enabled) = 0;
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
virtual void runner_done() {}
virtual void runner_end() {}
};
// ldm.modules.encoders.modules.FrozenCLIPEmbedder
@@ -202,13 +201,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
text_model->set_stream_layers_enabled(enabled);
if (sd_version_is_sdxl(version)) {
text_model2->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
text_model->set_runtime_backends(backends);
if (sd_version_is_sdxl(version)) {
@@ -244,10 +236,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void runner_done() override {
text_model->runner_done();
void runner_end() override {
text_model->runner_end();
if (sd_version_is_sdxl(version)) {
text_model2->runner_done();
text_model2->runner_end();
}
}
@@ -461,9 +453,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states.empty());
if (sd_version_is_sdxl(version)) {
auto chunk_hidden_states2 = text_model2->compute(n_threads,
@@ -473,9 +463,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states2.empty());
chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0);
@@ -487,9 +475,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled.empty());
}
}
@@ -608,7 +594,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(pixel_values, return_pooled, clip_skip);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
}
};
@@ -675,18 +661,6 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_stream_layers_enabled(enabled);
}
if (clip_g) {
clip_g->set_stream_layers_enabled(enabled);
}
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
@@ -753,15 +727,15 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (clip_l) {
clip_l->runner_done();
clip_l->runner_end();
}
if (clip_g) {
clip_g->runner_done();
clip_g->runner_end();
}
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -881,9 +855,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states_l.empty());
chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights);
@@ -897,9 +869,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled_l.empty());
}
} else {
@@ -928,9 +898,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states_g.empty());
chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights);
@@ -944,9 +912,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled_g.empty());
}
} else {
@@ -969,9 +935,7 @@ struct SD3CLIPEmbedder : public Conditioner {
chunk_hidden_states_t5 = t5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states_t5.empty());
chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights);
} else {
@@ -1079,15 +1043,6 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_stream_layers_enabled(enabled);
}
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
@@ -1139,12 +1094,12 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (clip_l) {
clip_l->runner_done();
clip_l->runner_end();
}
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -1247,9 +1202,7 @@ struct FluxCLIPEmbedder : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled.empty());
} else {
pooled = sd::Tensor<float>::zeros({768});
@@ -1268,9 +1221,7 @@ struct FluxCLIPEmbedder : public Conditioner {
chunk_hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states.empty());
chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
if (zero_out_masked) {
@@ -1366,12 +1317,6 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
@@ -1408,9 +1353,9 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -1508,9 +1453,7 @@ struct T5CLIPEmbedder : public Conditioner {
auto chunk_hidden_states = t5->compute(n_threads,
input_ids,
t5_attn_mask_chunk,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states.empty());
chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
@@ -1582,12 +1525,6 @@ struct MiniT2IConditioner : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
@@ -1624,9 +1561,9 @@ struct MiniT2IConditioner : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -1657,9 +1594,7 @@ struct MiniT2IConditioner : public Conditioner {
sd::Tensor<float> hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>::from_vector(t5_mask),
false,
true,
true);
false);
GGML_ASSERT(!hidden_states.empty());
result.c_crossattn = std::move(hidden_states);
result.c_vector = sd::Tensor<float>::from_vector(mask);
@@ -1696,10 +1631,6 @@ struct AnimaConditioner : public Conditioner {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
void set_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
@@ -1724,8 +1655,8 @@ struct AnimaConditioner : public Conditioner {
llm->set_weight_adapter(adapter);
}
void runner_done() override {
llm->runner_done();
void runner_end() override {
llm->runner_end();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<int>, std::vector<float>> tokenize(std::string text) {
@@ -1787,9 +1718,7 @@ struct AnimaConditioner : public Conditioner {
{},
{},
false,
false,
true,
true);
false);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), qwen_weights);
auto t5_ids_tensor = sd::Tensor<int32_t>::from_vector(t5_tokens);
@@ -1887,13 +1816,6 @@ struct LLMEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
if (byt5) {
byt5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
if (byt5) {
@@ -1942,12 +1864,12 @@ struct LLMEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (llm) {
llm->runner_done();
llm->runner_end();
}
if (byt5) {
byt5->runner_done();
byt5->runner_end();
}
}
@@ -2051,8 +1973,6 @@ struct LLMEmbedder : public Conditioner {
out_layers,
false,
false,
true,
true,
deepstack_image_embeds,
image_grids);
GGML_ASSERT(!hidden_states.empty());
@@ -2220,9 +2140,7 @@ struct LLMEmbedder : public Conditioner {
prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
continue;
@@ -2250,9 +2168,7 @@ struct LLMEmbedder : public Conditioner {
auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
add_vision_outputs(llm->encode_video_block_outputs(n_threads,
pair,
false,
true,
true),
false),
static_cast<int>(first.shape()[1]) / patch_size,
static_cast<int>(first.shape()[0]) / patch_size);
}
@@ -2263,9 +2179,7 @@ struct LLMEmbedder : public Conditioner {
prompt += "<Picture " + std::to_string(i + 1) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
}
@@ -2352,7 +2266,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
@@ -2411,7 +2325,7 @@ struct LLMEmbedder : public Conditioner {
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
@@ -2494,7 +2408,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>";
@@ -2562,7 +2476,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
@@ -2625,7 +2539,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
@@ -2857,9 +2771,7 @@ struct LLMEmbedder : public Conditioner {
auto byt5_hidden_states = byt5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false,
true,
true);
false);
GGML_ASSERT(!byt5_hidden_states.empty());
extra_hidden_states_vec.push_back(std::move(byt5_hidden_states));
}
@@ -2960,13 +2872,11 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool auto_runner_end = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
}
};
@@ -3060,9 +2970,9 @@ struct LTXAVEmbedder : public Conditioner {
projector->set_weight_adapter(adapter);
}
void runner_done() override {
llm->runner_done();
projector->runner_done();
void runner_end() override {
llm->runner_end();
projector->runner_end();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
@@ -3129,9 +3039,7 @@ struct LTXAVEmbedder : public Conditioner {
{},
{},
true,
false,
true,
true);
false);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), weights);
@@ -3190,7 +3098,7 @@ struct LTXAVEmbedder : public Conditioner {
}
hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens});
return projector->compute(n_threads, hidden_states, false, true, true);
return projector->compute(n_threads, hidden_states, false);
}
SDCondition get_learned_condition(int n_threads,