Skip to content

Commit 9ef6e73

Browse files
leejetleejet
authored
feat: drive layer split from graph-cut segments (#1762)
1 parent bb84971 commit 9ef6e73

8 files changed

Lines changed: 610 additions & 205 deletions

File tree

‎src/conditioning/conditioner.hpp‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,8 @@ struct Conditioner {
118118
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
119119
virtual void set_stream_layers_enabled(bool enabled) {}
120120
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
121+
virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
122+
virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
121123
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
122124
virtual void set_flash_attention_enabled(bool enabled) = 0;
123125
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
@@ -181,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
181183
}
182184
}
183185

186+
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
187+
text_model->set_runtime_backends(backends);
188+
if (sd_version_is_sdxl(version)) {
189+
text_model2->set_runtime_backends(backends);
190+
}
191+
}
192+
193+
void set_graph_cut_layer_split_enabled(bool enabled) override {
194+
text_model->set_graph_cut_layer_split_enabled(enabled);
195+
if (sd_version_is_sdxl(version)) {
196+
text_model2->set_graph_cut_layer_split_enabled(enabled);
197+
}
198+
}
199+
200+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
201+
text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
202+
if (sd_version_is_sdxl(version)) {
203+
text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
204+
}
205+
}
206+
184207
void set_flash_attention_enabled(bool enabled) override {
185208
text_model->set_flash_attention_enabled(enabled);
186209
if (sd_version_is_sdxl(version)) {
@@ -639,11 +662,41 @@ struct SD3CLIPEmbedder : public Conditioner {
639662
}
640663

641664
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
665+
if (clip_l) {
666+
clip_l->set_runtime_backends(backends);
667+
}
668+
if (clip_g) {
669+
clip_g->set_runtime_backends(backends);
670+
}
642671
if (t5) {
643672
t5->set_runtime_backends(backends);
644673
}
645674
}
646675

676+
void set_graph_cut_layer_split_enabled(bool enabled) override {
677+
if (clip_l) {
678+
clip_l->set_graph_cut_layer_split_enabled(enabled);
679+
}
680+
if (clip_g) {
681+
clip_g->set_graph_cut_layer_split_enabled(enabled);
682+
}
683+
if (t5) {
684+
t5->set_graph_cut_layer_split_enabled(enabled);
685+
}
686+
}
687+
688+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
689+
if (clip_l) {
690+
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
691+
}
692+
if (clip_g) {
693+
clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
694+
}
695+
if (t5) {
696+
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
697+
}
698+
}
699+
647700
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
648701
if (t5) {
649702
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1010,11 +1063,32 @@ struct FluxCLIPEmbedder : public Conditioner {
10101063
}
10111064

10121065
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
1066+
if (clip_l) {
1067+
clip_l->set_runtime_backends(backends);
1068+
}
10131069
if (t5) {
10141070
t5->set_runtime_backends(backends);
10151071
}
10161072
}
10171073

1074+
void set_graph_cut_layer_split_enabled(bool enabled) override {
1075+
if (clip_l) {
1076+
clip_l->set_graph_cut_layer_split_enabled(enabled);
1077+
}
1078+
if (t5) {
1079+
t5->set_graph_cut_layer_split_enabled(enabled);
1080+
}
1081+
}
1082+
1083+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
1084+
if (clip_l) {
1085+
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
1086+
}
1087+
if (t5) {
1088+
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
1089+
}
1090+
}
1091+
10181092
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
10191093
if (t5) {
10201094
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1278,6 +1352,18 @@ struct T5CLIPEmbedder : public Conditioner {
12781352
}
12791353
}
12801354

1355+
void set_graph_cut_layer_split_enabled(bool enabled) override {
1356+
if (t5) {
1357+
t5->set_graph_cut_layer_split_enabled(enabled);
1358+
}
1359+
}
1360+
1361+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
1362+
if (t5) {
1363+
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
1364+
}
1365+
}
1366+
12811367
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
12821368
if (t5) {
12831369
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1482,6 +1568,18 @@ struct MiniT2IConditioner : public Conditioner {
14821568
}
14831569
}
14841570

1571+
void set_graph_cut_layer_split_enabled(bool enabled) override {
1572+
if (t5) {
1573+
t5->set_graph_cut_layer_split_enabled(enabled);
1574+
}
1575+
}
1576+
1577+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
1578+
if (t5) {
1579+
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
1580+
}
1581+
}
1582+
14851583
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
14861584
if (t5) {
14871585
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1576,6 +1674,14 @@ struct AnimaConditioner : public Conditioner {
15761674
llm->set_runtime_backends(backends);
15771675
}
15781676

1677+
void set_graph_cut_layer_split_enabled(bool enabled) override {
1678+
llm->set_graph_cut_layer_split_enabled(enabled);
1679+
}
1680+
1681+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
1682+
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
1683+
}
1684+
15791685
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
15801686
llm->get_param_tensors(tensors, "text_encoders.llm");
15811687
}
@@ -1729,6 +1835,18 @@ struct LLMEmbedder : public Conditioner {
17291835
llm->set_runtime_backends(backends);
17301836
}
17311837

1838+
void set_graph_cut_layer_split_enabled(bool enabled) override {
1839+
if (llm) {
1840+
llm->set_graph_cut_layer_split_enabled(enabled);
1841+
}
1842+
}
1843+
1844+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
1845+
if (llm) {
1846+
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
1847+
}
1848+
}
1849+
17321850
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
17331851
llm->get_param_tensors(tensors, "text_encoders.llm");
17341852
}
@@ -2406,6 +2524,14 @@ struct LTXAVEmbedder : public Conditioner {
24062524
llm->set_runtime_backends(backends);
24072525
}
24082526

2527+
void set_graph_cut_layer_split_enabled(bool enabled) override {
2528+
llm->set_graph_cut_layer_split_enabled(enabled);
2529+
}
2530+
2531+
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
2532+
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
2533+
}
2534+
24092535
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
24102536
llm->get_param_tensors(tensors, "text_encoders.llm");
24112537
}

0 commit comments

Comments
 (0)