DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 03:24:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: centralize runner weight staging and cleanup (#1644)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-06-13 13:19:13 +08:00
1 parent 3a54597776
commit 563137a592
43 files changed
+2335 -1729
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/conditioning/conditioner.hpp
+175 -209
View File
@@ -113,14 +113,13 @@ struct Conditioner {
public:
virtual SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) = 0;
virtual bool alloc_params_buffer() = 0;
virtual void free_params_buffer() = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual size_t get_params_buffer_size() = 0;
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_flash_attention_enabled(bool enabled) = 0;
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
virtual void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) {}
virtual void runner_done() {}
};
// ldm.modules.encoders.modules.FrozenCLIPEmbedder
@@ -167,33 +166,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
bool alloc_params_buffer() override {
if (!text_model->alloc_params_buffer()) {
return false;
}
if (sd_version_is_sdxl(version)) {
if (!text_model2->alloc_params_buffer()) {
return false;
}
}
return true;
}
void free_params_buffer() override {
text_model->free_params_buffer();
if (sd_version_is_sdxl(version)) {
text_model2->free_params_buffer();
}
}
size_t get_params_buffer_size() override {
size_t buffer_size = text_model->get_params_buffer_size();
if (sd_version_is_sdxl(version)) {
buffer_size += text_model2->get_params_buffer_size();
}
return buffer_size;
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
text_model->set_max_graph_vram_bytes(max_vram_bytes);
if (sd_version_is_sdxl(version)) {
@@ -222,6 +194,20 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
text_model->set_weight_manager(manager);
if (sd_version_is_sdxl(version)) {
text_model2->set_weight_manager(manager);
}
}
void runner_done() override {
text_model->runner_done();
if (sd_version_is_sdxl(version)) {
text_model2->runner_done();
}
}
bool load_embedding(std::string embd_name, std::string embd_path, std::vector<int32_t>& bpe_tokens) {
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(embd_path)) {
@@ -263,7 +249,8 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
return true;
};
model_loader.load_tensors(on_load, 1);
model_loader.set_n_threads(1);
model_loader.load_tensors(on_load);
int pos_start = num_custom_embeddings;
if (embd) {
int64_t hidden_size = text_model->model.hidden_size;
@@ -432,7 +419,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
token_embed_custom.data(),
max_token_idx,
false,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!chunk_hidden_states.empty());
if (sd_version_is_sdxl(version)) {
auto chunk_hidden_states2 = text_model2->compute(n_threads,
@@ -441,7 +431,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
token_embed_custom.data(),
max_token_idx,
false,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!chunk_hidden_states2.empty());
chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0);
@@ -452,7 +445,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
token_embed_custom.data(),
max_token_idx,
true,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!pooled.empty());
}
}
@@ -523,15 +519,15 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
struct FrozenCLIPVisionEmbedder : public GGMLRunner {
CLIPVisionModelProjection vision_model;
std::string weight_prefix = "cond_stage_model.transformer";
FrozenCLIPVisionEmbedder(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {})
: GGMLRunner(backend, params_backend) {
std::string prefix = "cond_stage_model.transformer";
bool proj_in = false;
bool proj_in = false;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
if (!starts_with(name, weight_prefix)) {
continue;
}
if (contains(name, "self_attn.in_proj")) {
@@ -540,7 +536,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
}
}
vision_model = CLIPVisionModelProjection(OPEN_CLIP_VIT_H_14, false, proj_in);
vision_model.init(params_ctx, tensor_storage_map, prefix);
vision_model.init(params_ctx, tensor_storage_map, weight_prefix);
}
std::string get_desc() override {
@@ -548,7 +544,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
vision_model.get_param_tensors(tensors, "cond_stage_model.transformer");
vision_model.get_param_tensors(tensors, weight_prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& pixel_values_tensor, bool return_pooled, int clip_skip) {
@@ -571,7 +567,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));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
}
};
@@ -626,51 +622,6 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
bool alloc_params_buffer() override {
if (clip_l) {
if (!clip_l->alloc_params_buffer()) {
return false;
}
}
if (clip_g) {
if (!clip_g->alloc_params_buffer()) {
return false;
}
}
if (t5) {
if (!t5->alloc_params_buffer()) {
return false;
}
}
return true;
}
void free_params_buffer() override {
if (clip_l) {
clip_l->free_params_buffer();
}
if (clip_g) {
clip_g->free_params_buffer();
}
if (t5) {
t5->free_params_buffer();
}
}
size_t get_params_buffer_size() override {
size_t buffer_size = 0;
if (clip_l) {
buffer_size += clip_l->get_params_buffer_size();
}
if (clip_g) {
buffer_size += clip_g->get_params_buffer_size();
}
if (t5) {
buffer_size += t5->get_params_buffer_size();
}
return buffer_size;
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
if (clip_l) {
clip_l->set_max_graph_vram_bytes(max_vram_bytes);
@@ -719,6 +670,30 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
if (clip_l) {
clip_l->set_weight_manager(manager);
}
if (clip_g) {
clip_g->set_weight_manager(manager);
}
if (t5) {
t5->set_weight_manager(manager);
}
}
void runner_done() override {
if (clip_l) {
clip_l->runner_done();
}
if (clip_g) {
clip_g->runner_done();
}
if (t5) {
t5->runner_done();
}
}
std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
size_t min_length = 0,
size_t max_length = 0,
@@ -834,7 +809,10 @@ struct SD3CLIPEmbedder : public Conditioner {
nullptr,
max_token_idx,
false,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!chunk_hidden_states_l.empty());
chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights);
@@ -847,7 +825,10 @@ struct SD3CLIPEmbedder : public Conditioner {
nullptr,
max_token_idx,
true,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!pooled_l.empty());
}
} else {
@@ -875,7 +856,10 @@ struct SD3CLIPEmbedder : public Conditioner {
nullptr,
max_token_idx,
false,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!chunk_hidden_states_g.empty());
chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights);
@@ -888,7 +872,10 @@ struct SD3CLIPEmbedder : public Conditioner {
nullptr,
max_token_idx,
true,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!pooled_g.empty());
}
} else {
@@ -910,7 +897,10 @@ struct SD3CLIPEmbedder : public Conditioner {
chunk_hidden_states_t5 = t5->compute(n_threads,
input_ids,
sd::Tensor<float>());
sd::Tensor<float>(),
false,
true,
true);
GGML_ASSERT(!chunk_hidden_states_t5.empty());
chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights);
} else {
@@ -1009,40 +999,6 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
bool alloc_params_buffer() override {
if (clip_l) {
if (!clip_l->alloc_params_buffer()) {
return false;
}
}
if (t5) {
if (!t5->alloc_params_buffer()) {
return false;
}
}
return true;
}
void free_params_buffer() override {
if (clip_l) {
clip_l->free_params_buffer();
}
if (t5) {
t5->free_params_buffer();
}
}
size_t get_params_buffer_size() override {
size_t buffer_size = 0;
if (clip_l) {
buffer_size += clip_l->get_params_buffer_size();
}
if (t5) {
buffer_size += t5->get_params_buffer_size();
}
return buffer_size;
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
if (clip_l) {
clip_l->set_max_graph_vram_bytes(max_vram_bytes);
@@ -1070,7 +1026,7 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
if (clip_l) {
clip_l->set_weight_adapter(adapter);
}
@@ -1079,6 +1035,24 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
if (clip_l) {
clip_l->set_weight_manager(manager);
}
if (t5) {
t5->set_weight_manager(manager);
}
}
void runner_done() override {
if (clip_l) {
clip_l->runner_done();
}
if (t5) {
t5->runner_done();
}
}
std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
size_t min_length = 0,
size_t max_length = 0) {
@@ -1177,7 +1151,10 @@ struct FluxCLIPEmbedder : public Conditioner {
nullptr,
max_token_idx,
true,
clip_skip);
clip_skip,
false,
true,
true);
GGML_ASSERT(!pooled.empty());
} else {
pooled = sd::Tensor<float>::zeros({768});
@@ -1195,7 +1172,10 @@ struct FluxCLIPEmbedder : public Conditioner {
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
chunk_hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>());
sd::Tensor<float>(),
false,
true,
true);
GGML_ASSERT(!chunk_hidden_states.empty());
chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
if (zero_out_masked) {
@@ -1266,29 +1246,6 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
bool alloc_params_buffer() override {
if (t5) {
if (!t5->alloc_params_buffer()) {
return false;
}
}
return true;
}
void free_params_buffer() override {
if (t5) {
t5->free_params_buffer();
}
}
size_t get_params_buffer_size() override {
size_t buffer_size = 0;
if (t5) {
buffer_size += t5->get_params_buffer_size();
}
return buffer_size;
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
if (t5) {
t5->set_max_graph_vram_bytes(max_vram_bytes);
@@ -1313,6 +1270,18 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
if (t5) {
t5->set_weight_manager(manager);
}
}
void runner_done() override {
if (t5) {
t5->runner_done();
}
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
size_t min_length = 0,
size_t max_length = 0) {
@@ -1406,7 +1375,10 @@ struct T5CLIPEmbedder : public Conditioner {
auto chunk_hidden_states = t5->compute(n_threads,
input_ids,
t5_attn_mask_chunk);
t5_attn_mask_chunk,
false,
true,
true);
GGML_ASSERT(!chunk_hidden_states.empty());
chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
@@ -1465,21 +1437,6 @@ struct AnimaConditioner : public Conditioner {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
bool alloc_params_buffer() override {
if (!llm->alloc_params_buffer()) {
return false;
}
return true;
}
void free_params_buffer() override {
llm->free_params_buffer();
}
size_t get_params_buffer_size() override {
return llm->get_params_buffer_size();
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
@@ -1496,6 +1453,14 @@ struct AnimaConditioner : public Conditioner {
llm->set_weight_adapter(adapter);
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
llm->set_weight_manager(manager);
}
void runner_done() override {
llm->runner_done();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<int>, std::vector<float>> tokenize(std::string text) {
auto parsed_attention = parse_prompt_attention(text);
@@ -1553,7 +1518,11 @@ struct AnimaConditioner : public Conditioner {
input_ids,
sd::Tensor<float>(),
{},
{});
{},
false,
false,
true,
true);
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);
@@ -1617,23 +1586,6 @@ struct LLMEmbedder : public Conditioner {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
bool alloc_params_buffer() override {
if (!llm->alloc_params_buffer()) {
return false;
}
return true;
}
void free_params_buffer() override {
llm->free_params_buffer();
}
size_t get_params_buffer_size() override {
size_t buffer_size = 0;
buffer_size += llm->get_params_buffer_size();
return buffer_size;
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
@@ -1652,6 +1604,18 @@ struct LLMEmbedder : public Conditioner {
}
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
if (llm) {
llm->set_weight_manager(manager);
}
}
void runner_done() override {
if (llm) {
llm->runner_done();
}
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
const std::pair<int, int>& attn_range,
size_t min_length = 0,
@@ -1747,7 +1711,11 @@ struct LLMEmbedder : public Conditioner {
input_ids,
attention_mask,
image_embeds,
out_layers);
out_layers,
false,
false,
true,
true);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), weights);
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
@@ -1825,7 +1793,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);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
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;
@@ -1895,7 +1863,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);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
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;
@@ -2163,11 +2131,15 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
return gf;
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& x) {
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) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
}
};
@@ -2205,25 +2177,6 @@ struct LTXAVEmbedder : public Conditioner {
projector->get_param_tensors(tensors, "text_embedding_projection");
}
bool alloc_params_buffer() override {
if (!llm->alloc_params_buffer()) {
return false;
}
if (!projector->alloc_params_buffer()) {
return false;
}
return true;
}
void free_params_buffer() override {
llm->free_params_buffer();
projector->free_params_buffer();
}
size_t get_params_buffer_size() override {
return llm->get_params_buffer_size() + projector->get_params_buffer_size();
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
projector->set_flash_attention_enabled(enabled);
@@ -2239,6 +2192,16 @@ struct LTXAVEmbedder : public Conditioner {
projector->set_weight_adapter(adapter);
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
llm->set_weight_manager(manager);
projector->set_weight_manager(manager);
}
void runner_done() override {
llm->runner_done();
projector->runner_done();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
const std::pair<int, int>& attn_range) {
std::vector<std::pair<std::string, float>> parsed_attention;
@@ -2302,6 +2265,9 @@ struct LTXAVEmbedder : public Conditioner {
attention_mask,
{},
{},
true,
false,
true,
true);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), weights);
@@ -2361,7 +2327,7 @@ struct LTXAVEmbedder : public Conditioner {
}
hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens});
return projector->compute(n_threads, hidden_states);
return projector->compute(n_threads, hidden_states, false, true, true);
}
SDCondition get_learned_condition(int n_threads,