mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-05-08 16:28:53 +00:00
Compare commits
No commits in common. "master" and "master-592-b8079e2" have entirely different histories.
master
...
master-592
@ -54,8 +54,6 @@ Context Options:
|
|||||||
-t, --threads <int> number of threads to use during computation (default: -1). If threads <= 0,
|
-t, --threads <int> number of threads to use during computation (default: -1). If threads <= 0,
|
||||||
then threads will be set to the number of CPU physical cores
|
then threads will be set to the number of CPU physical cores
|
||||||
--chroma-t5-mask-pad <int> t5 mask pad size of chroma
|
--chroma-t5-mask-pad <int> t5 mask pad size of chroma
|
||||||
--max-vram <float> maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
|
|
||||||
graph splitting
|
|
||||||
--force-sdxl-vae-conv-scale force use of conv scale on sdxl vae
|
--force-sdxl-vae-conv-scale force use of conv scale on sdxl vae
|
||||||
--offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM
|
--offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM
|
||||||
when needed
|
when needed
|
||||||
|
|||||||
@ -394,12 +394,7 @@ ArgOptions SDContextParams::get_options() {
|
|||||||
&chroma_t5_mask_pad},
|
&chroma_t5_mask_pad},
|
||||||
};
|
};
|
||||||
|
|
||||||
options.float_options = {
|
options.float_options = {};
|
||||||
{"",
|
|
||||||
"--max-vram",
|
|
||||||
"maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting",
|
|
||||||
&max_vram},
|
|
||||||
};
|
|
||||||
|
|
||||||
options.bool_options = {
|
options.bool_options = {
|
||||||
{"",
|
{"",
|
||||||
@ -675,7 +670,6 @@ std::string SDContextParams::to_string() const {
|
|||||||
<< " rng_type: " << sd_rng_type_name(rng_type) << ",\n"
|
<< " rng_type: " << sd_rng_type_name(rng_type) << ",\n"
|
||||||
<< " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
|
<< " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
|
||||||
<< " offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
|
<< " offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
|
||||||
<< " max_vram: " << max_vram << ",\n"
|
|
||||||
<< " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
|
<< " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
|
||||||
<< " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
|
<< " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
|
||||||
<< " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
|
<< " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
|
||||||
@ -750,7 +744,6 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
|
|||||||
chroma_use_t5_mask,
|
chroma_use_t5_mask,
|
||||||
chroma_t5_mask_pad,
|
chroma_t5_mask_pad,
|
||||||
qwen_image_zero_cond_t,
|
qwen_image_zero_cond_t,
|
||||||
max_vram,
|
|
||||||
};
|
};
|
||||||
return sd_ctx_params;
|
return sd_ctx_params;
|
||||||
}
|
}
|
||||||
|
|||||||
@ -109,7 +109,6 @@ struct SDContextParams {
|
|||||||
rng_type_t rng_type = CUDA_RNG;
|
rng_type_t rng_type = CUDA_RNG;
|
||||||
rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
|
rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
|
||||||
bool offload_params_to_cpu = false;
|
bool offload_params_to_cpu = false;
|
||||||
float max_vram = 0.f;
|
|
||||||
bool enable_mmap = false;
|
bool enable_mmap = false;
|
||||||
bool control_net_cpu = false;
|
bool control_net_cpu = false;
|
||||||
bool clip_on_cpu = false;
|
bool clip_on_cpu = false;
|
||||||
|
|||||||
@ -156,8 +156,6 @@ Context Options:
|
|||||||
-t, --threads <int> number of threads to use during computation (default: -1). If threads <= 0,
|
-t, --threads <int> number of threads to use during computation (default: -1). If threads <= 0,
|
||||||
then threads will be set to the number of CPU physical cores
|
then threads will be set to the number of CPU physical cores
|
||||||
--chroma-t5-mask-pad <int> t5 mask pad size of chroma
|
--chroma-t5-mask-pad <int> t5 mask pad size of chroma
|
||||||
--max-vram <float> maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
|
|
||||||
graph splitting
|
|
||||||
--force-sdxl-vae-conv-scale force use of conv scale on sdxl vae
|
--force-sdxl-vae-conv-scale force use of conv scale on sdxl vae
|
||||||
--offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM
|
--offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM
|
||||||
when needed
|
when needed
|
||||||
|
|||||||
@ -203,7 +203,6 @@ typedef struct {
|
|||||||
bool chroma_use_t5_mask;
|
bool chroma_use_t5_mask;
|
||||||
int chroma_t5_mask_pad;
|
int chroma_t5_mask_pad;
|
||||||
bool qwen_image_zero_cond_t;
|
bool qwen_image_zero_cond_t;
|
||||||
float max_vram;
|
|
||||||
} sd_ctx_params_t;
|
} sd_ctx_params_t;
|
||||||
|
|
||||||
typedef struct {
|
typedef struct {
|
||||||
|
|||||||
@ -499,15 +499,9 @@ namespace Anima {
|
|||||||
encoder_hidden_states = adapted_context;
|
encoder_hidden_states = adapted_context;
|
||||||
}
|
}
|
||||||
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "anima.prelude", "x");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(embedded_timestep, "anima.prelude", "embedded_timestep");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(temb, "anima.prelude", "temb");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(encoder_hidden_states, "anima.prelude", "context");
|
|
||||||
|
|
||||||
for (int i = 0; i < num_layers; i++) {
|
for (int i = 0; i < num_layers; i++) {
|
||||||
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
|
||||||
x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
|
x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C]
|
x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C]
|
||||||
|
|||||||
@ -328,7 +328,6 @@ public:
|
|||||||
auto conv_out = std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"]);
|
auto conv_out = std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"]);
|
||||||
|
|
||||||
auto h = conv_in->forward(ctx, x); // [N, ch, h, w]
|
auto h = conv_in->forward(ctx, x); // [N, ch, h, w]
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.prelude", "h");
|
|
||||||
|
|
||||||
// downsampling
|
// downsampling
|
||||||
size_t num_resolutions = ch_mult.size();
|
size_t num_resolutions = ch_mult.size();
|
||||||
@ -338,14 +337,12 @@ public:
|
|||||||
auto down_block = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
|
auto down_block = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
|
||||||
|
|
||||||
h = down_block->forward(ctx, h);
|
h = down_block->forward(ctx, h);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".block." + std::to_string(j), "h");
|
|
||||||
}
|
}
|
||||||
if (i != num_resolutions - 1) {
|
if (i != num_resolutions - 1) {
|
||||||
std::string name = "down." + std::to_string(i) + ".downsample";
|
std::string name = "down." + std::to_string(i) + ".downsample";
|
||||||
auto down_sample = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
|
auto down_sample = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
|
||||||
|
|
||||||
h = down_sample->forward(ctx, h);
|
h = down_sample->forward(ctx, h);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".downsample", "h");
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -353,7 +350,6 @@ public:
|
|||||||
h = mid_block_1->forward(ctx, h);
|
h = mid_block_1->forward(ctx, h);
|
||||||
h = mid_attn_1->forward(ctx, h);
|
h = mid_attn_1->forward(ctx, h);
|
||||||
h = mid_block_2->forward(ctx, h); // [N, block_in, h, w]
|
h = mid_block_2->forward(ctx, h); // [N, block_in, h, w]
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.mid", "h");
|
|
||||||
|
|
||||||
// end
|
// end
|
||||||
h = norm_out->forward(ctx, h);
|
h = norm_out->forward(ctx, h);
|
||||||
@ -454,7 +450,6 @@ public:
|
|||||||
|
|
||||||
// conv_in
|
// conv_in
|
||||||
auto h = conv_in->forward(ctx, z); // [N, block_in, h, w]
|
auto h = conv_in->forward(ctx, z); // [N, block_in, h, w]
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.prelude", "h");
|
|
||||||
|
|
||||||
// middle
|
// middle
|
||||||
h = mid_block_1->forward(ctx, h);
|
h = mid_block_1->forward(ctx, h);
|
||||||
@ -462,7 +457,6 @@ public:
|
|||||||
|
|
||||||
h = mid_attn_1->forward(ctx, h);
|
h = mid_attn_1->forward(ctx, h);
|
||||||
h = mid_block_2->forward(ctx, h); // [N, block_in, h, w]
|
h = mid_block_2->forward(ctx, h); // [N, block_in, h, w]
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.mid", "h");
|
|
||||||
|
|
||||||
// upsampling
|
// upsampling
|
||||||
int num_resolutions = static_cast<int>(ch_mult.size());
|
int num_resolutions = static_cast<int>(ch_mult.size());
|
||||||
@ -472,14 +466,12 @@ public:
|
|||||||
auto up_block = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
|
auto up_block = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
|
||||||
|
|
||||||
h = up_block->forward(ctx, h);
|
h = up_block->forward(ctx, h);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".block." + std::to_string(j), "h");
|
|
||||||
}
|
}
|
||||||
if (i != 0) {
|
if (i != 0) {
|
||||||
std::string name = "up." + std::to_string(i) + ".upsample";
|
std::string name = "up." + std::to_string(i) + ".upsample";
|
||||||
auto up_sample = std::dynamic_pointer_cast<UpSampleBlock>(blocks[name]);
|
auto up_sample = std::dynamic_pointer_cast<UpSampleBlock>(blocks[name]);
|
||||||
|
|
||||||
h = up_sample->forward(ctx, h);
|
h = up_sample->forward(ctx, h);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".upsample", "h");
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -607,7 +599,6 @@ public:
|
|||||||
if (use_quant) {
|
if (use_quant) {
|
||||||
auto post_quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["post_quant_conv"]);
|
auto post_quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["post_quant_conv"]);
|
||||||
z = post_quant_conv->forward(ctx, z); // [N, z_channels, h, w]
|
z = post_quant_conv->forward(ctx, z); // [N, z_channels, h, w]
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(z, "vae.decode.prelude", "z");
|
|
||||||
}
|
}
|
||||||
auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
|
auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
|
||||||
|
|
||||||
@ -625,7 +616,6 @@ public:
|
|||||||
if (use_quant) {
|
if (use_quant) {
|
||||||
auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["quant_conv"]);
|
auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["quant_conv"]);
|
||||||
z = quant_conv->forward(ctx, z); // [N, 2*embed_dim, h/8, w/8]
|
z = quant_conv->forward(ctx, z); // [N, 2*embed_dim, h/8, w/8]
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(z, "vae.encode.final", "z");
|
|
||||||
}
|
}
|
||||||
if (sd_version_uses_flux2_vae(version)) {
|
if (sd_version_uses_flux2_vae(version)) {
|
||||||
z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
|
z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
|
||||||
|
|||||||
14
src/clip.hpp
14
src/clip.hpp
@ -95,9 +95,8 @@ public:
|
|||||||
|
|
||||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
int clip_skip = -1,
|
int clip_skip = -1) {
|
||||||
const std::string& graph_cut_prefix = "") {
|
|
||||||
// x: [N, n_token, d_model]
|
// x: [N, n_token, d_model]
|
||||||
int layer_idx = n_layer - 1;
|
int layer_idx = n_layer - 1;
|
||||||
// LOG_DEBUG("clip_skip %d", clip_skip);
|
// LOG_DEBUG("clip_skip %d", clip_skip);
|
||||||
@ -113,9 +112,6 @@ public:
|
|||||||
std::string name = "layers." + std::to_string(i);
|
std::string name = "layers." + std::to_string(i);
|
||||||
auto layer = std::dynamic_pointer_cast<CLIPLayer>(blocks[name]);
|
auto layer = std::dynamic_pointer_cast<CLIPLayer>(blocks[name]);
|
||||||
x = layer->forward(ctx, x, mask); // [N, n_token, d_model]
|
x = layer->forward(ctx, x, mask); // [N, n_token, d_model]
|
||||||
if (!graph_cut_prefix.empty()) {
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x");
|
|
||||||
}
|
|
||||||
// LOG_DEBUG("layer %d", i);
|
// LOG_DEBUG("layer %d", i);
|
||||||
}
|
}
|
||||||
return x;
|
return x;
|
||||||
@ -308,8 +304,7 @@ public:
|
|||||||
auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
|
auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
|
||||||
|
|
||||||
auto x = embeddings->forward(ctx, input_ids, tkn_embeddings); // [N, n_token, hidden_size]
|
auto x = embeddings->forward(ctx, input_ids, tkn_embeddings); // [N, n_token, hidden_size]
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "clip_text.prelude", "x");
|
x = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip);
|
||||||
x = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip, "clip_text");
|
|
||||||
if (return_pooled || with_final_ln) {
|
if (return_pooled || with_final_ln) {
|
||||||
x = final_layer_norm->forward(ctx, x);
|
x = final_layer_norm->forward(ctx, x);
|
||||||
}
|
}
|
||||||
@ -373,8 +368,7 @@ public:
|
|||||||
|
|
||||||
auto x = embeddings->forward(ctx, pixel_values); // [N, num_positions, embed_dim]
|
auto x = embeddings->forward(ctx, pixel_values); // [N, num_positions, embed_dim]
|
||||||
x = pre_layernorm->forward(ctx, x);
|
x = pre_layernorm->forward(ctx, x);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "clip_vision.prelude", "x");
|
x = encoder->forward(ctx, x, nullptr, clip_skip);
|
||||||
x = encoder->forward(ctx, x, nullptr, clip_skip, "clip_vision");
|
|
||||||
|
|
||||||
auto last_hidden_state = x;
|
auto last_hidden_state = x;
|
||||||
|
|
||||||
|
|||||||
@ -85,8 +85,7 @@ public:
|
|||||||
virtual void free_params_buffer() = 0;
|
virtual void free_params_buffer() = 0;
|
||||||
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
|
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
|
||||||
virtual size_t get_params_buffer_size() = 0;
|
virtual size_t get_params_buffer_size() = 0;
|
||||||
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
|
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
||||||
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
|
||||||
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
|
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
|
||||||
virtual std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(int n_threads,
|
virtual std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(int n_threads,
|
||||||
const ConditionerParams& conditioner_params) {
|
const ConditionerParams& conditioner_params) {
|
||||||
@ -166,13 +165,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
|||||||
return 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)) {
|
|
||||||
text_model2->set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_flash_attention_enabled(bool enabled) override {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
text_model->set_flash_attention_enabled(enabled);
|
text_model->set_flash_attention_enabled(enabled);
|
||||||
if (sd_version_is_sdxl(version)) {
|
if (sd_version_is_sdxl(version)) {
|
||||||
@ -789,18 +781,6 @@ struct SD3CLIPEmbedder : public Conditioner {
|
|||||||
return 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);
|
|
||||||
}
|
|
||||||
if (clip_g) {
|
|
||||||
clip_g->set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
if (t5) {
|
|
||||||
t5->set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_flash_attention_enabled(bool enabled) override {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
if (clip_l) {
|
if (clip_l) {
|
||||||
clip_l->set_flash_attention_enabled(enabled);
|
clip_l->set_flash_attention_enabled(enabled);
|
||||||
@ -1144,15 +1124,6 @@ struct FluxCLIPEmbedder : public Conditioner {
|
|||||||
return 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);
|
|
||||||
}
|
|
||||||
if (t5) {
|
|
||||||
t5->set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_flash_attention_enabled(bool enabled) override {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
if (clip_l) {
|
if (clip_l) {
|
||||||
clip_l->set_flash_attention_enabled(enabled);
|
clip_l->set_flash_attention_enabled(enabled);
|
||||||
@ -1378,12 +1349,6 @@ struct T5CLIPEmbedder : public Conditioner {
|
|||||||
return 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);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_flash_attention_enabled(bool enabled) override {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
if (t5) {
|
if (t5) {
|
||||||
t5->set_flash_attention_enabled(enabled);
|
t5->set_flash_attention_enabled(enabled);
|
||||||
@ -1560,10 +1525,6 @@ struct AnimaConditioner : public Conditioner {
|
|||||||
return llm->get_params_buffer_size();
|
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);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_flash_attention_enabled(bool enabled) override {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
llm->set_flash_attention_enabled(enabled);
|
llm->set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
@ -1696,10 +1657,6 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
return 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);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_flash_attention_enabled(bool enabled) override {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
llm->set_flash_attention_enabled(enabled);
|
llm->set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|||||||
@ -808,18 +808,6 @@ static std::tuple<float, float, float> get_ancestral_step_flow(float sigma_from,
|
|||||||
return {sigma_down, sigma_up, alpha_scale};
|
return {sigma_down, sigma_up, alpha_scale};
|
||||||
}
|
}
|
||||||
|
|
||||||
static std::tuple<float, float, float> get_ancestral_step(float sigma_from,
|
|
||||||
float sigma_to,
|
|
||||||
float eta,
|
|
||||||
bool is_flow_denoiser) {
|
|
||||||
if (is_flow_denoiser) {
|
|
||||||
return get_ancestral_step_flow(sigma_from, sigma_to, eta);
|
|
||||||
} else {
|
|
||||||
auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
|
|
||||||
return {sigma_down, sigma_up, 1.0f};
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
|
static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
|
||||||
sd::Tensor<float> x,
|
sd::Tensor<float> x,
|
||||||
const std::vector<float>& sigmas,
|
const std::vector<float>& sigmas,
|
||||||
@ -1259,7 +1247,6 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
|
|||||||
sd::Tensor<float> x,
|
sd::Tensor<float> x,
|
||||||
const std::vector<float>& sigmas,
|
const std::vector<float>& sigmas,
|
||||||
std::shared_ptr<RNG> rng,
|
std::shared_ptr<RNG> rng,
|
||||||
bool is_flow_denoiser,
|
|
||||||
float eta) {
|
float eta) {
|
||||||
sd::Tensor<float> old_denoised = x;
|
sd::Tensor<float> old_denoised = x;
|
||||||
bool have_old_sigma = false;
|
bool have_old_sigma = false;
|
||||||
@ -1291,8 +1278,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
|
|||||||
|
|
||||||
float sigma_from = sigmas[i];
|
float sigma_from = sigmas[i];
|
||||||
float sigma_to = sigmas[i + 1];
|
float sigma_to = sigmas[i + 1];
|
||||||
|
auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
|
||||||
auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step(sigma_from, sigma_to, eta, is_flow_denoiser);
|
|
||||||
|
|
||||||
if (sigma_down == 0.0f || !have_old_sigma) {
|
if (sigma_down == 0.0f || !have_old_sigma) {
|
||||||
x += ((x - denoised) / sigma_from) * (sigma_down - sigma_from);
|
x += ((x - denoised) / sigma_from) * (sigma_down - sigma_from);
|
||||||
@ -1319,10 +1305,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
|
|||||||
x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised);
|
x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised);
|
||||||
}
|
}
|
||||||
|
|
||||||
if (sigma_to > 0.0f && sigma_up > 0.0f) {
|
if (sigmas[i + 1] > 0 && sigma_up > 0.0f) {
|
||||||
if (is_flow_denoiser) {
|
|
||||||
x *= alpha_scale;
|
|
||||||
}
|
|
||||||
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
|
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1337,7 +1320,6 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
|
|||||||
sd::Tensor<float> x,
|
sd::Tensor<float> x,
|
||||||
const std::vector<float>& sigmas,
|
const std::vector<float>& sigmas,
|
||||||
std::shared_ptr<RNG> rng,
|
std::shared_ptr<RNG> rng,
|
||||||
bool is_flow_denoiser,
|
|
||||||
float eta) {
|
float eta) {
|
||||||
const float c2 = 0.5f;
|
const float c2 = 0.5f;
|
||||||
auto t_fn = [](float sigma) -> float { return -logf(sigma); };
|
auto t_fn = [](float sigma) -> float { return -logf(sigma); };
|
||||||
@ -1366,7 +1348,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
|
|||||||
}
|
}
|
||||||
sd::Tensor<float> denoised = std::move(denoised_opt);
|
sd::Tensor<float> denoised = std::move(denoised_opt);
|
||||||
|
|
||||||
auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step(sigma_from, sigma_to, eta, is_flow_denoiser);
|
auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
|
||||||
|
|
||||||
sd::Tensor<float> x0 = x;
|
sd::Tensor<float> x0 = x;
|
||||||
if (sigma_down == 0.0f || sigma_from == 0.0f) {
|
if (sigma_down == 0.0f || sigma_from == 0.0f) {
|
||||||
@ -1395,10 +1377,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
|
|||||||
x = x0 + h * (b1 * eps1 + b2 * eps2);
|
x = x0 + h * (b1 * eps1 + b2 * eps2);
|
||||||
}
|
}
|
||||||
|
|
||||||
if (sigma_to > 0.0f && sigma_up > 0.0f) {
|
if (sigmas[i + 1] > 0 && sigma_up > 0.0f) {
|
||||||
if (is_flow_denoiser) {
|
|
||||||
x *= alpha_scale;
|
|
||||||
}
|
|
||||||
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
|
x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -1685,9 +1664,9 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
|
|||||||
case IPNDM_V_SAMPLE_METHOD:
|
case IPNDM_V_SAMPLE_METHOD:
|
||||||
return sample_ipndm_v(model, std::move(x), sigmas);
|
return sample_ipndm_v(model, std::move(x), sigmas);
|
||||||
case RES_MULTISTEP_SAMPLE_METHOD:
|
case RES_MULTISTEP_SAMPLE_METHOD:
|
||||||
return sample_res_multistep(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
|
return sample_res_multistep(model, std::move(x), sigmas, rng, eta);
|
||||||
case RES_2S_SAMPLE_METHOD:
|
case RES_2S_SAMPLE_METHOD:
|
||||||
return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
|
return sample_res_2s(model, std::move(x), sigmas, rng, eta);
|
||||||
case ER_SDE_SAMPLE_METHOD:
|
case ER_SDE_SAMPLE_METHOD:
|
||||||
return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
|
return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
|
||||||
case DDIM_TRAILING_SAMPLE_METHOD:
|
case DDIM_TRAILING_SAMPLE_METHOD:
|
||||||
|
|||||||
@ -49,7 +49,6 @@ struct DiffusionModel {
|
|||||||
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
|
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
|
||||||
virtual int64_t get_adm_in_channels() = 0;
|
virtual int64_t get_adm_in_channels() = 0;
|
||||||
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
||||||
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) = 0;
|
|
||||||
virtual void set_circular_axes(bool circular_x, bool circular_y) = 0;
|
virtual void set_circular_axes(bool circular_x, bool circular_y) = 0;
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -99,10 +98,6 @@ struct UNetModel : public DiffusionModel {
|
|||||||
unet.set_flash_attention_enabled(enabled);
|
unet.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
unet.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
unet.set_circular_axes(circular_x, circular_y);
|
unet.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
@ -169,10 +164,6 @@ struct MMDiTModel : public DiffusionModel {
|
|||||||
mmdit.set_flash_attention_enabled(enabled);
|
mmdit.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
mmdit.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
mmdit.set_circular_axes(circular_x, circular_y);
|
mmdit.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
@ -238,10 +229,6 @@ struct FluxModel : public DiffusionModel {
|
|||||||
flux.set_flash_attention_enabled(enabled);
|
flux.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
flux.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
flux.set_circular_axes(circular_x, circular_y);
|
flux.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
@ -312,10 +299,6 @@ struct AnimaModel : public DiffusionModel {
|
|||||||
anima.set_flash_attention_enabled(enabled);
|
anima.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
anima.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
anima.set_circular_axes(circular_x, circular_y);
|
anima.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
@ -381,10 +364,6 @@ struct WanModel : public DiffusionModel {
|
|||||||
wan.set_flash_attention_enabled(enabled);
|
wan.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
wan.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
wan.set_circular_axes(circular_x, circular_y);
|
wan.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
@ -454,10 +433,6 @@ struct QwenImageModel : public DiffusionModel {
|
|||||||
qwen_image.set_flash_attention_enabled(enabled);
|
qwen_image.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
qwen_image.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
qwen_image.set_circular_axes(circular_x, circular_y);
|
qwen_image.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
@ -524,10 +499,6 @@ struct ZImageModel : public DiffusionModel {
|
|||||||
z_image.set_flash_attention_enabled(enabled);
|
z_image.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
z_image.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
z_image.set_circular_axes(circular_x, circular_y);
|
z_image.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
@ -593,10 +564,6 @@ struct ErnieImageModel : public DiffusionModel {
|
|||||||
ernie_image.set_flash_attention_enabled(enabled);
|
ernie_image.set_flash_attention_enabled(enabled);
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
|
||||||
ernie_image.set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||||
ernie_image.set_circular_axes(circular_x, circular_y);
|
ernie_image.set_circular_axes(circular_x, circular_y);
|
||||||
}
|
}
|
||||||
|
|||||||
@ -295,9 +295,7 @@ namespace ErnieImage {
|
|||||||
auto c = time_embedding->forward(ctx, sample); // [N, hidden_size]
|
auto c = time_embedding->forward(ctx, sample); // [N, hidden_size]
|
||||||
|
|
||||||
auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 6 * hidden_size]
|
auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 6 * hidden_size]
|
||||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.prelude", "hidden_states");
|
auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(mod_params, "ernie_image.prelude", "mod_params");
|
|
||||||
auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
|
|
||||||
std::vector<ggml_tensor*> temb;
|
std::vector<ggml_tensor*> temb;
|
||||||
temb.reserve(6);
|
temb.reserve(6);
|
||||||
for (auto chunk : chunks) {
|
for (auto chunk : chunks) {
|
||||||
@ -307,7 +305,6 @@ namespace ErnieImage {
|
|||||||
for (int i = 0; i < params.num_layers; i++) {
|
for (int i = 0; i < params.num_layers; i++) {
|
||||||
auto layer = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
|
auto layer = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
|
||||||
hidden_states = layer->forward(ctx, hidden_states, pe, temb);
|
hidden_states = layer->forward(ctx, hidden_states, pe, temb);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.layers." + std::to_string(i), "hidden_states");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
hidden_states = final_norm->forward(ctx, hidden_states, c);
|
hidden_states = final_norm->forward(ctx, hidden_states, c);
|
||||||
|
|||||||
@ -124,33 +124,27 @@ public:
|
|||||||
auto conv_hr = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
|
auto conv_hr = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
|
||||||
auto conv_last = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
|
auto conv_last = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
|
||||||
|
|
||||||
auto feat = conv_first->forward(ctx, x);
|
auto feat = conv_first->forward(ctx, x);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat");
|
|
||||||
auto body_feat = feat;
|
auto body_feat = feat;
|
||||||
for (int i = 0; i < num_block; i++) {
|
for (int i = 0; i < num_block; i++) {
|
||||||
std::string name = "body." + std::to_string(i);
|
std::string name = "body." + std::to_string(i);
|
||||||
auto block = std::dynamic_pointer_cast<RRDB>(blocks[name]);
|
auto block = std::dynamic_pointer_cast<RRDB>(blocks[name]);
|
||||||
|
|
||||||
body_feat = block->forward(ctx, body_feat);
|
body_feat = block->forward(ctx, body_feat);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat");
|
|
||||||
}
|
}
|
||||||
body_feat = conv_body->forward(ctx, body_feat);
|
body_feat = conv_body->forward(ctx, body_feat);
|
||||||
feat = ggml_add(ctx->ggml_ctx, feat, body_feat);
|
feat = ggml_add(ctx->ggml_ctx, feat, body_feat);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat");
|
|
||||||
// upsample
|
// upsample
|
||||||
if (scale >= 2) {
|
if (scale >= 2) {
|
||||||
auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
|
auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
|
||||||
feat = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
|
feat = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
|
||||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat");
|
|
||||||
if (scale == 4) {
|
if (scale == 4) {
|
||||||
auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
|
auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
|
||||||
feat = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
|
feat = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
|
||||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat");
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// for all scales
|
// for all scales
|
||||||
auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
|
auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
|
||||||
sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out");
|
|
||||||
return out;
|
return out;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
@ -928,9 +928,6 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
|
|
||||||
txt = txt_in->forward(ctx, txt);
|
txt = txt_in->forward(ctx, txt);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(img, "flux.prelude", "img");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt, "flux.prelude", "txt");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(vec, "flux.prelude", "vec");
|
|
||||||
|
|
||||||
for (int i = 0; i < params.depth; i++) {
|
for (int i = 0; i < params.depth; i++) {
|
||||||
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
|
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
|
||||||
@ -942,8 +939,6 @@ namespace Flux {
|
|||||||
auto img_txt = block->forward(ctx, img, txt, vec, pe, txt_img_mask, ds_img_mods, ds_txt_mods);
|
auto img_txt = block->forward(ctx, img, txt, vec, pe, txt_img_mask, ds_img_mods, ds_txt_mods);
|
||||||
img = img_txt.first; // [N, n_img_token, hidden_size]
|
img = img_txt.first; // [N, n_img_token, hidden_size]
|
||||||
txt = img_txt.second; // [N, n_txt_token, hidden_size]
|
txt = img_txt.second; // [N, n_txt_token, hidden_size]
|
||||||
sd::ggml_graph_cut::mark_graph_cut(img, "flux.double_blocks." + std::to_string(i), "img");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt, "flux.double_blocks." + std::to_string(i), "txt");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
|
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
|
||||||
@ -954,7 +949,6 @@ namespace Flux {
|
|||||||
auto block = std::dynamic_pointer_cast<SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
|
||||||
|
|
||||||
txt_img = block->forward(ctx, txt_img, vec, pe, txt_img_mask, ss_mods);
|
txt_img = block->forward(ctx, txt_img, vec, pe, txt_img_mask, ss_mods);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.single_blocks." + std::to_string(i), "txt_img");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
img = ggml_view_3d(ctx->ggml_ctx,
|
img = ggml_view_3d(ctx->ggml_ctx,
|
||||||
|
|||||||
@ -6,7 +6,6 @@
|
|||||||
#include <stdarg.h>
|
#include <stdarg.h>
|
||||||
#include <algorithm>
|
#include <algorithm>
|
||||||
#include <atomic>
|
#include <atomic>
|
||||||
#include <cstdlib>
|
|
||||||
#include <cstring>
|
#include <cstring>
|
||||||
#include <fstream>
|
#include <fstream>
|
||||||
#include <functional>
|
#include <functional>
|
||||||
@ -27,7 +26,6 @@
|
|||||||
#include "ggml-backend.h"
|
#include "ggml-backend.h"
|
||||||
#include "ggml.h"
|
#include "ggml.h"
|
||||||
#include "ggml_extend_backend.hpp"
|
#include "ggml_extend_backend.hpp"
|
||||||
#include "ggml_graph_cut.h"
|
|
||||||
|
|
||||||
#include "model.h"
|
#include "model.h"
|
||||||
#include "tensor.hpp"
|
#include "tensor.hpp"
|
||||||
@ -1710,8 +1708,6 @@ struct GGMLRunnerContext {
|
|||||||
struct GGMLRunner {
|
struct GGMLRunner {
|
||||||
protected:
|
protected:
|
||||||
typedef std::function<ggml_cgraph*()> get_graph_cb_t;
|
typedef std::function<ggml_cgraph*()> get_graph_cb_t;
|
||||||
using GraphCutSegment = sd::ggml_graph_cut::Segment;
|
|
||||||
using GraphCutPlan = sd::ggml_graph_cut::Plan;
|
|
||||||
|
|
||||||
ggml_backend_t params_backend = nullptr;
|
ggml_backend_t params_backend = nullptr;
|
||||||
ggml_backend_t runtime_backend = nullptr;
|
ggml_backend_t runtime_backend = nullptr;
|
||||||
@ -1728,11 +1724,6 @@ protected:
|
|||||||
ggml_context* compute_ctx = nullptr;
|
ggml_context* compute_ctx = nullptr;
|
||||||
ggml_gallocr* compute_allocr = nullptr;
|
ggml_gallocr* compute_allocr = nullptr;
|
||||||
|
|
||||||
ggml_context* partial_offload_ctx = nullptr;
|
|
||||||
ggml_backend_buffer_t partial_runtime_params_buffer = nullptr;
|
|
||||||
std::vector<std::pair<ggml_tensor*, ggml_tensor*>> partial_offload_pairs;
|
|
||||||
size_t max_graph_vram_bytes = 0;
|
|
||||||
|
|
||||||
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
|
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
|
||||||
|
|
||||||
std::vector<float> one_vec = {1.f};
|
std::vector<float> one_vec = {1.f};
|
||||||
@ -1750,9 +1741,6 @@ protected:
|
|||||||
bool circular_x_enabled = false;
|
bool circular_x_enabled = false;
|
||||||
bool circular_y_enabled = false;
|
bool circular_y_enabled = false;
|
||||||
|
|
||||||
sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
|
|
||||||
std::unordered_set<const ggml_tensor*> params_tensor_set_;
|
|
||||||
|
|
||||||
template <typename T>
|
template <typename T>
|
||||||
static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
|
static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
|
||||||
if (!tensor.has_value()) {
|
if (!tensor.has_value()) {
|
||||||
@ -1787,7 +1775,6 @@ protected:
|
|||||||
|
|
||||||
params_ctx = ggml_init(params);
|
params_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(params_ctx != nullptr);
|
GGML_ASSERT(params_ctx != nullptr);
|
||||||
params_tensor_set_.clear();
|
|
||||||
if (params_backend != runtime_backend) {
|
if (params_backend != runtime_backend) {
|
||||||
offload_ctx = ggml_init(params);
|
offload_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(offload_ctx != nullptr);
|
GGML_ASSERT(offload_ctx != nullptr);
|
||||||
@ -1799,15 +1786,10 @@ protected:
|
|||||||
ggml_free(params_ctx);
|
ggml_free(params_ctx);
|
||||||
params_ctx = nullptr;
|
params_ctx = nullptr;
|
||||||
}
|
}
|
||||||
params_tensor_set_.clear();
|
|
||||||
if (offload_ctx != nullptr) {
|
if (offload_ctx != nullptr) {
|
||||||
ggml_free(offload_ctx);
|
ggml_free(offload_ctx);
|
||||||
offload_ctx = nullptr;
|
offload_ctx = nullptr;
|
||||||
}
|
}
|
||||||
if (partial_offload_ctx != nullptr) {
|
|
||||||
ggml_free(partial_offload_ctx);
|
|
||||||
partial_offload_ctx = nullptr;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
void alloc_cache_ctx() {
|
void alloc_cache_ctx() {
|
||||||
@ -1842,17 +1824,6 @@ protected:
|
|||||||
ggml_free(compute_ctx);
|
ggml_free(compute_ctx);
|
||||||
compute_ctx = nullptr;
|
compute_ctx = nullptr;
|
||||||
}
|
}
|
||||||
backend_tensor_data_map.clear();
|
|
||||||
}
|
|
||||||
|
|
||||||
void rebuild_params_tensor_set() {
|
|
||||||
params_tensor_set_.clear();
|
|
||||||
if (params_ctx == nullptr) {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) {
|
|
||||||
params_tensor_set_.insert(t);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
void prepare_build_in_tensor_before() {
|
void prepare_build_in_tensor_before() {
|
||||||
@ -1888,25 +1859,13 @@ protected:
|
|||||||
return gf;
|
return gf;
|
||||||
}
|
}
|
||||||
|
|
||||||
bool prepare_compute_graph(get_graph_cb_t get_graph,
|
bool alloc_compute_buffer(get_graph_cb_t get_graph) {
|
||||||
ggml_cgraph** gf_out) {
|
|
||||||
GGML_ASSERT(gf_out != nullptr);
|
|
||||||
|
|
||||||
reset_compute_ctx();
|
|
||||||
ggml_cgraph* gf = get_compute_graph(get_graph);
|
|
||||||
if (gf == nullptr) {
|
|
||||||
free_compute_ctx();
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
|
|
||||||
*gf_out = gf;
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
bool alloc_compute_buffer(ggml_cgraph* gf) {
|
|
||||||
if (compute_allocr != nullptr) {
|
if (compute_allocr != nullptr) {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
reset_compute_ctx();
|
||||||
|
ggml_cgraph* gf = get_compute_graph(get_graph);
|
||||||
|
backend_tensor_data_map.clear();
|
||||||
compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_backend));
|
compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_backend));
|
||||||
|
|
||||||
if (!ggml_gallocr_reserve(compute_allocr, gf)) {
|
if (!ggml_gallocr_reserve(compute_allocr, gf)) {
|
||||||
@ -1932,132 +1891,47 @@ protected:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
bool copy_cache_tensors_to_cache_buffer(const std::unordered_set<std::string>* cache_keep_names = nullptr) {
|
void copy_cache_tensors_to_cache_buffer() {
|
||||||
ggml_context* old_cache_ctx = cache_ctx;
|
if (cache_tensor_map.size() == 0) {
|
||||||
ggml_backend_buffer_t old_cache_buffer = cache_buffer;
|
return;
|
||||||
cache_ctx = nullptr;
|
|
||||||
cache_buffer = nullptr;
|
|
||||||
std::map<std::string, ggml_tensor*> merged_cache_sources;
|
|
||||||
if (old_cache_ctx != nullptr) {
|
|
||||||
for (ggml_tensor* tensor = ggml_get_first_tensor(old_cache_ctx); tensor != nullptr; tensor = ggml_get_next_tensor(old_cache_ctx, tensor)) {
|
|
||||||
if (cache_keep_names != nullptr && cache_keep_names->find(tensor->name) == cache_keep_names->end()) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
merged_cache_sources[tensor->name] = tensor;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
for (const auto& kv : cache_tensor_map) {
|
free_cache_ctx_and_buffer();
|
||||||
if (cache_keep_names != nullptr && cache_keep_names->find(kv.first) == cache_keep_names->end()) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
merged_cache_sources[kv.first] = kv.second;
|
|
||||||
}
|
|
||||||
cache_tensor_map.clear();
|
|
||||||
if (merged_cache_sources.empty()) {
|
|
||||||
if (old_cache_buffer != nullptr) {
|
|
||||||
ggml_backend_buffer_free(old_cache_buffer);
|
|
||||||
}
|
|
||||||
if (old_cache_ctx != nullptr) {
|
|
||||||
ggml_free(old_cache_ctx);
|
|
||||||
}
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
alloc_cache_ctx();
|
alloc_cache_ctx();
|
||||||
std::vector<std::pair<ggml_tensor*, ggml_tensor*>> source_to_cache_tensors;
|
GGML_ASSERT(cache_buffer == nullptr);
|
||||||
source_to_cache_tensors.reserve(merged_cache_sources.size());
|
std::map<ggml_tensor*, ggml_tensor*> runtime_tensor_to_cache_tensor;
|
||||||
for (const auto& kv : merged_cache_sources) {
|
for (auto kv : cache_tensor_map) {
|
||||||
ggml_tensor* source_tensor = sd::ggml_graph_cut::cache_source_tensor(kv.second);
|
auto cache_tensor = ggml_dup_tensor(cache_ctx, kv.second);
|
||||||
auto cache_tensor = ggml_dup_tensor(cache_ctx, source_tensor);
|
|
||||||
ggml_set_name(cache_tensor, kv.first.c_str());
|
ggml_set_name(cache_tensor, kv.first.c_str());
|
||||||
source_to_cache_tensors.push_back({source_tensor, cache_tensor});
|
runtime_tensor_to_cache_tensor[kv.second] = cache_tensor;
|
||||||
}
|
}
|
||||||
size_t num_tensors = ggml_tensor_num(cache_ctx);
|
size_t num_tensors = ggml_tensor_num(cache_ctx);
|
||||||
cache_buffer = ggml_backend_alloc_ctx_tensors(cache_ctx, runtime_backend);
|
cache_buffer = ggml_backend_alloc_ctx_tensors(cache_ctx, runtime_backend);
|
||||||
GGML_ASSERT(cache_buffer != nullptr);
|
GGML_ASSERT(cache_buffer != nullptr);
|
||||||
for (const auto& kv : source_to_cache_tensors) {
|
for (auto kv : runtime_tensor_to_cache_tensor) {
|
||||||
ggml_tensor* src = kv.first;
|
ggml_backend_tensor_copy(kv.first, kv.second);
|
||||||
ggml_tensor* dst = kv.second;
|
|
||||||
ggml_backend_buffer_t src_buf = sd::ggml_graph_cut::tensor_buffer(src);
|
|
||||||
ggml_backend_buffer_t dst_buf = sd::ggml_graph_cut::tensor_buffer(dst);
|
|
||||||
if (src_buf == nullptr || dst_buf == nullptr) {
|
|
||||||
LOG_ERROR("%s cache copy tensor buffer missing: name=%s src_buffer=%p src_view_src=%p src_view_src_buffer=%p dst_buffer=%p",
|
|
||||||
get_desc().c_str(),
|
|
||||||
src && src->name[0] != '\0' ? src->name : "<unnamed>",
|
|
||||||
src ? src->buffer : nullptr,
|
|
||||||
src ? src->view_src : nullptr,
|
|
||||||
(src && src->view_src) ? src->view_src->buffer : nullptr,
|
|
||||||
dst ? dst->buffer : nullptr);
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
const bool use_staging_copy = src->view_src != nullptr || !ggml_is_contiguous(src) || src->buffer == nullptr;
|
|
||||||
if (use_staging_copy) {
|
|
||||||
std::vector<uint8_t> host_data(ggml_nbytes(src));
|
|
||||||
ggml_backend_tensor_get(src, host_data.data(), 0, host_data.size());
|
|
||||||
ggml_backend_tensor_set(dst, host_data.data(), 0, host_data.size());
|
|
||||||
} else {
|
|
||||||
ggml_backend_tensor_copy(src, dst);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
ggml_backend_synchronize(runtime_backend);
|
ggml_backend_synchronize(runtime_backend);
|
||||||
|
cache_tensor_map.clear();
|
||||||
size_t cache_buffer_size = ggml_backend_buffer_get_size(cache_buffer);
|
size_t cache_buffer_size = ggml_backend_buffer_get_size(cache_buffer);
|
||||||
LOG_DEBUG("%s cache backend buffer size = % 6.2f MB(%s) (%i tensors)",
|
LOG_DEBUG("%s cache backend buffer size = % 6.2f MB(%s) (%i tensors)",
|
||||||
get_desc().c_str(),
|
get_desc().c_str(),
|
||||||
cache_buffer_size / (1024.f * 1024.f),
|
cache_buffer_size / (1024.f * 1024.f),
|
||||||
ggml_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM",
|
ggml_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM",
|
||||||
num_tensors);
|
num_tensors);
|
||||||
if (old_cache_buffer != nullptr) {
|
|
||||||
ggml_backend_buffer_free(old_cache_buffer);
|
|
||||||
}
|
|
||||||
if (old_cache_ctx != nullptr) {
|
|
||||||
ggml_free(old_cache_ctx);
|
|
||||||
}
|
|
||||||
return true;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true) {
|
void copy_data_to_backend_tensor() {
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
std::unordered_set<const ggml_tensor*> graph_tensor_set;
|
|
||||||
const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
|
|
||||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
|
||||||
graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
|
|
||||||
for (int i = 0; i < n_leafs; ++i) {
|
|
||||||
graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
|
|
||||||
}
|
|
||||||
for (int i = 0; i < n_nodes; ++i) {
|
|
||||||
graph_tensor_set.insert(ggml_graph_node(gf, i));
|
|
||||||
}
|
|
||||||
|
|
||||||
for (auto& kv : backend_tensor_data_map) {
|
for (auto& kv : backend_tensor_data_map) {
|
||||||
auto tensor = kv.first;
|
auto tensor = kv.first;
|
||||||
auto data = kv.second;
|
auto data = kv.second;
|
||||||
|
|
||||||
if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
|
|
||||||
if (buf == nullptr) {
|
|
||||||
LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
|
|
||||||
get_desc().c_str(),
|
|
||||||
tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
|
|
||||||
tensor ? ggml_op_name(tensor->op) : "<null>",
|
|
||||||
data,
|
|
||||||
tensor ? tensor->view_src : nullptr,
|
|
||||||
(tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
|
ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
|
||||||
}
|
}
|
||||||
|
|
||||||
if (clear_after_copy) {
|
backend_tensor_data_map.clear();
|
||||||
backend_tensor_data_map.clear();
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
bool offload_all_params() {
|
bool offload_params_to_runtime_backend() {
|
||||||
restore_partial_params();
|
|
||||||
if (params_backend == runtime_backend) {
|
if (params_backend == runtime_backend) {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
@ -2084,7 +1958,6 @@ protected:
|
|||||||
num_tensors);
|
num_tensors);
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
ggml_backend_buffer_set_usage(runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
|
|
||||||
|
|
||||||
ggml_tensor* t = ggml_get_first_tensor(params_ctx);
|
ggml_tensor* t = ggml_get_first_tensor(params_ctx);
|
||||||
ggml_tensor* offload_t = ggml_get_first_tensor(offload_ctx);
|
ggml_tensor* offload_t = ggml_get_first_tensor(offload_ctx);
|
||||||
@ -2114,85 +1987,7 @@ protected:
|
|||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
|
||||||
bool offload_partial_params(const std::vector<ggml_tensor*>& tensors) {
|
void offload_params_to_params_backend() {
|
||||||
restore_partial_params();
|
|
||||||
if (params_backend == runtime_backend) {
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
if (tensors.empty()) {
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
GGML_ASSERT(!params_on_runtime_backend);
|
|
||||||
GGML_ASSERT(partial_runtime_params_buffer == nullptr);
|
|
||||||
|
|
||||||
std::vector<ggml_tensor*> unique_tensors;
|
|
||||||
std::unordered_set<ggml_tensor*> seen_tensors;
|
|
||||||
unique_tensors.reserve(tensors.size());
|
|
||||||
seen_tensors.reserve(tensors.size());
|
|
||||||
for (ggml_tensor* tensor : tensors) {
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
if (seen_tensors.insert(tensor).second) {
|
|
||||||
unique_tensors.push_back(tensor);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if (unique_tensors.empty()) {
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_init_params params;
|
|
||||||
params.mem_size = std::max<size_t>(1, unique_tensors.size()) * ggml_tensor_overhead();
|
|
||||||
params.mem_buffer = nullptr;
|
|
||||||
params.no_alloc = true;
|
|
||||||
|
|
||||||
partial_offload_ctx = ggml_init(params);
|
|
||||||
GGML_ASSERT(partial_offload_ctx != nullptr);
|
|
||||||
|
|
||||||
partial_offload_pairs.clear();
|
|
||||||
partial_offload_pairs.reserve(unique_tensors.size());
|
|
||||||
|
|
||||||
for (ggml_tensor* tensor : unique_tensors) {
|
|
||||||
GGML_ASSERT(tensor->view_src == nullptr);
|
|
||||||
ggml_tensor* offload_tensor = ggml_dup_tensor(partial_offload_ctx, tensor);
|
|
||||||
ggml_set_name(offload_tensor, tensor->name);
|
|
||||||
partial_offload_pairs.push_back({tensor, offload_tensor});
|
|
||||||
}
|
|
||||||
|
|
||||||
partial_runtime_params_buffer = ggml_backend_alloc_ctx_tensors(partial_offload_ctx, runtime_backend);
|
|
||||||
if (partial_runtime_params_buffer == nullptr) {
|
|
||||||
LOG_ERROR("%s alloc partial runtime params backend buffer failed, num_tensors = %zu",
|
|
||||||
get_desc().c_str(),
|
|
||||||
partial_offload_pairs.size());
|
|
||||||
ggml_free(partial_offload_ctx);
|
|
||||||
partial_offload_ctx = nullptr;
|
|
||||||
partial_offload_pairs.clear();
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
ggml_backend_buffer_set_usage(partial_runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
|
|
||||||
|
|
||||||
for (auto& pair : partial_offload_pairs) {
|
|
||||||
ggml_tensor* tensor = pair.first;
|
|
||||||
ggml_tensor* offload_tensor = pair.second;
|
|
||||||
|
|
||||||
ggml_backend_tensor_copy(tensor, offload_tensor);
|
|
||||||
std::swap(tensor->buffer, offload_tensor->buffer);
|
|
||||||
std::swap(tensor->data, offload_tensor->data);
|
|
||||||
std::swap(tensor->extra, offload_tensor->extra);
|
|
||||||
}
|
|
||||||
|
|
||||||
size_t params_buffer_size = ggml_backend_buffer_get_size(partial_runtime_params_buffer);
|
|
||||||
LOG_DEBUG("%s offload partial params (%6.2f MB, %zu tensors) to runtime backend (%s)",
|
|
||||||
get_desc().c_str(),
|
|
||||||
params_buffer_size / (1024.f * 1024.f),
|
|
||||||
partial_offload_pairs.size(),
|
|
||||||
ggml_backend_name(runtime_backend));
|
|
||||||
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
void restore_all_params() {
|
|
||||||
restore_partial_params();
|
|
||||||
if (!params_on_runtime_backend) {
|
if (!params_on_runtime_backend) {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@ -2218,323 +2013,17 @@ protected:
|
|||||||
params_on_runtime_backend = false;
|
params_on_runtime_backend = false;
|
||||||
}
|
}
|
||||||
|
|
||||||
void restore_partial_params() {
|
|
||||||
if (partial_offload_pairs.empty()) {
|
|
||||||
if (partial_runtime_params_buffer != nullptr) {
|
|
||||||
ggml_backend_buffer_free(partial_runtime_params_buffer);
|
|
||||||
partial_runtime_params_buffer = nullptr;
|
|
||||||
}
|
|
||||||
if (partial_offload_ctx != nullptr) {
|
|
||||||
ggml_free(partial_offload_ctx);
|
|
||||||
partial_offload_ctx = nullptr;
|
|
||||||
}
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
for (auto& pair : partial_offload_pairs) {
|
|
||||||
ggml_tensor* tensor = pair.first;
|
|
||||||
ggml_tensor* offload_tensor = pair.second;
|
|
||||||
|
|
||||||
tensor->buffer = offload_tensor->buffer;
|
|
||||||
tensor->data = offload_tensor->data;
|
|
||||||
tensor->extra = offload_tensor->extra;
|
|
||||||
offload_tensor->buffer = nullptr;
|
|
||||||
offload_tensor->data = nullptr;
|
|
||||||
offload_tensor->extra = nullptr;
|
|
||||||
}
|
|
||||||
|
|
||||||
if (partial_runtime_params_buffer != nullptr) {
|
|
||||||
ggml_backend_buffer_free(partial_runtime_params_buffer);
|
|
||||||
partial_runtime_params_buffer = nullptr;
|
|
||||||
}
|
|
||||||
partial_offload_pairs.clear();
|
|
||||||
|
|
||||||
if (partial_offload_ctx != nullptr) {
|
|
||||||
ggml_free(partial_offload_ctx);
|
|
||||||
partial_offload_ctx = nullptr;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
bool should_use_graph_cut_segmented_compute(const GraphCutPlan& plan) {
|
|
||||||
return plan.has_cuts &&
|
|
||||||
plan.valid &&
|
|
||||||
max_graph_vram_bytes > 0 &&
|
|
||||||
plan.segments.size() > 1 &&
|
|
||||||
params_backend != runtime_backend &&
|
|
||||||
!ggml_backend_is_cpu(runtime_backend);
|
|
||||||
}
|
|
||||||
|
|
||||||
bool can_attempt_graph_cut_segmented_compute() const {
|
|
||||||
return max_graph_vram_bytes > 0 &&
|
|
||||||
params_backend != runtime_backend &&
|
|
||||||
!ggml_backend_is_cpu(runtime_backend);
|
|
||||||
}
|
|
||||||
|
|
||||||
bool resolve_graph_cut_plan(ggml_cgraph* gf,
|
|
||||||
GraphCutPlan* plan_out) {
|
|
||||||
GGML_ASSERT(plan_out != nullptr);
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
*plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
|
|
||||||
gf,
|
|
||||||
&graph_cut_plan_cache_,
|
|
||||||
max_graph_vram_bytes,
|
|
||||||
params_tensor_set_,
|
|
||||||
get_desc().c_str());
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
void reset_segment_runtime_tensors(const GraphCutSegment& segment,
|
|
||||||
ggml_cgraph* gf) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
|
|
||||||
for (const auto& input : segment.input_refs) {
|
|
||||||
ggml_tensor* input_tensor = sd::ggml_graph_cut::input_tensor(gf, input);
|
|
||||||
if (input_tensor == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
switch (input.type) {
|
|
||||||
case GraphCutSegment::INPUT_PREVIOUS_CUT:
|
|
||||||
case GraphCutSegment::INPUT_EXTERNAL:
|
|
||||||
input_tensor->buffer = nullptr;
|
|
||||||
input_tensor->data = nullptr;
|
|
||||||
input_tensor->extra = nullptr;
|
|
||||||
break;
|
|
||||||
case GraphCutSegment::INPUT_PARAM:
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
for (int node_idx : segment.internal_node_indices) {
|
|
||||||
ggml_tensor* node = ggml_graph_node(gf, node_idx);
|
|
||||||
if (node == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
node->buffer = nullptr;
|
|
||||||
node->data = nullptr;
|
|
||||||
node->extra = nullptr;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
bool bind_segment_cached_inputs(ggml_cgraph* gf, const GraphCutSegment& segment) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
for (const auto& input : segment.input_refs) {
|
|
||||||
ggml_tensor* input_tensor = sd::ggml_graph_cut::input_tensor(gf, input);
|
|
||||||
if (input_tensor == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
switch (input.type) {
|
|
||||||
case GraphCutSegment::INPUT_PREVIOUS_CUT: {
|
|
||||||
ggml_tensor* cache_tensor = get_cache_tensor_by_name(input.display_name);
|
|
||||||
if (cache_tensor == nullptr) {
|
|
||||||
LOG_ERROR("%s missing graph cut cache tensor: %s",
|
|
||||||
get_desc().c_str(),
|
|
||||||
input.display_name.c_str());
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
if (input_tensor->view_src != nullptr) {
|
|
||||||
input_tensor->view_src = cache_tensor;
|
|
||||||
input_tensor->buffer = nullptr;
|
|
||||||
input_tensor->data = cache_tensor->data == nullptr
|
|
||||||
? nullptr
|
|
||||||
: static_cast<void*>(static_cast<char*>(cache_tensor->data) + input_tensor->view_offs);
|
|
||||||
input_tensor->extra = cache_tensor->extra;
|
|
||||||
} else {
|
|
||||||
input_tensor->buffer = cache_tensor->buffer;
|
|
||||||
input_tensor->data = cache_tensor->data;
|
|
||||||
input_tensor->extra = cache_tensor->extra;
|
|
||||||
}
|
|
||||||
for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
|
|
||||||
input_tensor->src[src_idx] = nullptr;
|
|
||||||
}
|
|
||||||
input_tensor->op = GGML_OP_NONE;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
case GraphCutSegment::INPUT_EXTERNAL:
|
|
||||||
case GraphCutSegment::INPUT_PARAM:
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
template <typename T>
|
|
||||||
std::optional<sd::Tensor<T>> execute_graph(ggml_cgraph* gf,
|
|
||||||
int n_threads,
|
|
||||||
bool free_compute_buffer_immediately,
|
|
||||||
const std::vector<ggml_tensor*>& runtime_param_tensors,
|
|
||||||
bool preserve_backend_tensor_data_map,
|
|
||||||
bool no_return = false,
|
|
||||||
const std::unordered_set<std::string>* cache_keep_names = nullptr) {
|
|
||||||
int64_t t_execute_begin = ggml_time_ms();
|
|
||||||
const bool use_partial_param_offload = !runtime_param_tensors.empty();
|
|
||||||
int64_t t_offload_begin = ggml_time_ms();
|
|
||||||
if (use_partial_param_offload) {
|
|
||||||
if (!offload_partial_params(runtime_param_tensors)) {
|
|
||||||
LOG_ERROR("%s offload partial params to runtime backend failed", get_desc().c_str());
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
if (!offload_all_params()) {
|
|
||||||
LOG_ERROR("%s offload params to runtime backend failed", get_desc().c_str());
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
int64_t t_offload_end = ggml_time_ms();
|
|
||||||
|
|
||||||
int64_t t_alloc_begin = ggml_time_ms();
|
|
||||||
if (!alloc_compute_buffer(gf)) {
|
|
||||||
LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str());
|
|
||||||
if (use_partial_param_offload) {
|
|
||||||
restore_partial_params();
|
|
||||||
}
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
|
|
||||||
if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
|
|
||||||
LOG_ERROR("%s alloc compute graph failed", get_desc().c_str());
|
|
||||||
if (free_compute_buffer_immediately) {
|
|
||||||
free_compute_buffer();
|
|
||||||
} else if (use_partial_param_offload) {
|
|
||||||
restore_partial_params();
|
|
||||||
}
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
int64_t t_alloc_end = ggml_time_ms();
|
|
||||||
|
|
||||||
int64_t t_copy_begin = ggml_time_ms();
|
|
||||||
copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map);
|
|
||||||
int64_t t_copy_end = ggml_time_ms();
|
|
||||||
if (ggml_backend_is_cpu(runtime_backend)) {
|
|
||||||
ggml_backend_cpu_set_n_threads(runtime_backend, n_threads);
|
|
||||||
}
|
|
||||||
|
|
||||||
int64_t t_compute_begin = ggml_time_ms();
|
|
||||||
ggml_status status = ggml_backend_graph_compute(runtime_backend, gf);
|
|
||||||
int64_t t_compute_end = ggml_time_ms();
|
|
||||||
if (status != GGML_STATUS_SUCCESS) {
|
|
||||||
LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
|
|
||||||
if (free_compute_buffer_immediately) {
|
|
||||||
free_compute_buffer();
|
|
||||||
} else if (use_partial_param_offload) {
|
|
||||||
restore_partial_params();
|
|
||||||
}
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
|
|
||||||
int64_t t_cache_begin = ggml_time_ms();
|
|
||||||
if (!copy_cache_tensors_to_cache_buffer(cache_keep_names)) {
|
|
||||||
if (free_compute_buffer_immediately) {
|
|
||||||
free_compute_buffer();
|
|
||||||
} else if (use_partial_param_offload) {
|
|
||||||
restore_partial_params();
|
|
||||||
}
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
int64_t t_cache_end = ggml_time_ms();
|
|
||||||
auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
|
|
||||||
std::optional<sd::Tensor<T>> output;
|
|
||||||
if (!no_return) {
|
|
||||||
output = sd::make_sd_tensor_from_ggml<T>(result);
|
|
||||||
} else {
|
|
||||||
output = sd::Tensor<T>();
|
|
||||||
}
|
|
||||||
|
|
||||||
if (free_compute_buffer_immediately) {
|
|
||||||
free_compute_buffer();
|
|
||||||
} else if (use_partial_param_offload) {
|
|
||||||
restore_partial_params();
|
|
||||||
}
|
|
||||||
if (use_partial_param_offload) {
|
|
||||||
LOG_DEBUG("%s execute_graph timing: offload=%lld ms alloc=%lld ms copy_in=%lld ms compute=%lld ms cache=%lld ms total=%lld ms",
|
|
||||||
get_desc().c_str(),
|
|
||||||
t_offload_end - t_offload_begin,
|
|
||||||
t_alloc_end - t_alloc_begin,
|
|
||||||
t_copy_end - t_copy_begin,
|
|
||||||
t_compute_end - t_compute_begin,
|
|
||||||
t_cache_end - t_cache_begin,
|
|
||||||
ggml_time_ms() - t_execute_begin);
|
|
||||||
}
|
|
||||||
return output;
|
|
||||||
}
|
|
||||||
|
|
||||||
template <typename T>
|
|
||||||
std::optional<sd::Tensor<T>> compute_with_graph_cuts(ggml_cgraph* gf,
|
|
||||||
const GraphCutPlan& plan,
|
|
||||||
int n_threads,
|
|
||||||
bool free_compute_buffer_immediately,
|
|
||||||
bool no_return = false) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
|
|
||||||
free_compute_buffer();
|
|
||||||
free_cache_ctx_and_buffer();
|
|
||||||
|
|
||||||
std::optional<sd::Tensor<T>> output = sd::Tensor<T>();
|
|
||||||
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) {
|
|
||||||
int64_t t_segment_begin = ggml_time_ms();
|
|
||||||
const auto& segment = plan.segments[seg_idx];
|
|
||||||
auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx);
|
|
||||||
LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s",
|
|
||||||
get_desc().c_str(),
|
|
||||||
seg_idx + 1,
|
|
||||||
plan.segments.size(),
|
|
||||||
segment.group_name.c_str());
|
|
||||||
|
|
||||||
reset_segment_runtime_tensors(segment, gf);
|
|
||||||
if (!bind_segment_cached_inputs(gf, segment)) {
|
|
||||||
free_cache_ctx_and_buffer();
|
|
||||||
free_compute_buffer();
|
|
||||||
free_compute_ctx();
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
|
|
||||||
const bool is_last_segment = seg_idx + 1 == plan.segments.size();
|
|
||||||
if (!is_last_segment) {
|
|
||||||
for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) {
|
|
||||||
ggml_tensor* output_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx);
|
|
||||||
if (output_tensor != nullptr &&
|
|
||||||
sd::ggml_graph_cut::is_graph_cut_tensor(output_tensor) &&
|
|
||||||
future_cut_names.find(output_tensor->name) != future_cut_names.end()) {
|
|
||||||
cache(output_tensor->name, output_tensor);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_context* segment_graph_ctx = nullptr;
|
|
||||||
ggml_cgraph* segment_graph = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx);
|
|
||||||
auto segment_output = execute_graph<T>(segment_graph,
|
|
||||||
n_threads,
|
|
||||||
true,
|
|
||||||
sd::ggml_graph_cut::runtime_param_tensors(gf, segment, get_desc().c_str()),
|
|
||||||
true,
|
|
||||||
!is_last_segment || no_return,
|
|
||||||
&future_cut_names);
|
|
||||||
ggml_free(segment_graph_ctx);
|
|
||||||
if (!segment_output.has_value()) {
|
|
||||||
free_cache_ctx_and_buffer();
|
|
||||||
free_compute_buffer();
|
|
||||||
free_compute_ctx();
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
output = std::move(segment_output);
|
|
||||||
}
|
|
||||||
|
|
||||||
backend_tensor_data_map.clear();
|
|
||||||
free_cache_ctx_and_buffer();
|
|
||||||
free_compute_ctx();
|
|
||||||
return output;
|
|
||||||
}
|
|
||||||
|
|
||||||
public:
|
public:
|
||||||
virtual std::string get_desc() = 0;
|
virtual std::string get_desc() = 0;
|
||||||
|
|
||||||
GGMLRunner(ggml_backend_t backend, bool offload_params_to_cpu = false)
|
GGMLRunner(ggml_backend_t backend, bool offload_params_to_cpu = false)
|
||||||
: runtime_backend(backend) {
|
: runtime_backend(backend) {
|
||||||
|
alloc_params_ctx();
|
||||||
if (!ggml_backend_is_cpu(runtime_backend) && offload_params_to_cpu) {
|
if (!ggml_backend_is_cpu(runtime_backend) && offload_params_to_cpu) {
|
||||||
params_backend = ggml_backend_cpu_init();
|
params_backend = ggml_backend_cpu_init();
|
||||||
} else {
|
} else {
|
||||||
params_backend = runtime_backend;
|
params_backend = runtime_backend;
|
||||||
}
|
}
|
||||||
alloc_params_ctx();
|
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual ~GGMLRunner() {
|
virtual ~GGMLRunner() {
|
||||||
@ -2574,8 +2063,6 @@ public:
|
|||||||
num_tensors);
|
num_tensors);
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
rebuild_params_tensor_set();
|
|
||||||
ggml_backend_buffer_set_usage(params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
|
|
||||||
size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer);
|
size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer);
|
||||||
LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)",
|
LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)",
|
||||||
get_desc().c_str(),
|
get_desc().c_str(),
|
||||||
@ -2609,8 +2096,7 @@ public:
|
|||||||
ggml_gallocr_free(compute_allocr);
|
ggml_gallocr_free(compute_allocr);
|
||||||
compute_allocr = nullptr;
|
compute_allocr = nullptr;
|
||||||
}
|
}
|
||||||
restore_partial_params();
|
offload_params_to_params_backend();
|
||||||
restore_all_params();
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// do copy after alloc graph
|
// do copy after alloc graph
|
||||||
@ -2674,36 +2160,41 @@ public:
|
|||||||
int n_threads,
|
int n_threads,
|
||||||
bool free_compute_buffer_immediately,
|
bool free_compute_buffer_immediately,
|
||||||
bool no_return = false) {
|
bool no_return = false) {
|
||||||
ggml_cgraph* gf = nullptr;
|
if (!offload_params_to_runtime_backend()) {
|
||||||
if (!prepare_compute_graph(get_graph, &gf)) {
|
LOG_ERROR("%s offload params to runtime backend failed", get_desc().c_str());
|
||||||
return std::nullopt;
|
return std::nullopt;
|
||||||
}
|
}
|
||||||
GGML_ASSERT(gf != nullptr);
|
if (!alloc_compute_buffer(get_graph)) {
|
||||||
|
|
||||||
if (can_attempt_graph_cut_segmented_compute()) {
|
|
||||||
GraphCutPlan plan;
|
|
||||||
if (!resolve_graph_cut_plan(gf, &plan)) {
|
|
||||||
free_compute_ctx();
|
|
||||||
return std::nullopt;
|
|
||||||
}
|
|
||||||
if (should_use_graph_cut_segmented_compute(plan)) {
|
|
||||||
return compute_with_graph_cuts<T>(gf,
|
|
||||||
plan,
|
|
||||||
n_threads,
|
|
||||||
free_compute_buffer_immediately,
|
|
||||||
no_return);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if (!alloc_compute_buffer(gf)) {
|
|
||||||
LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str());
|
LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str());
|
||||||
return std::nullopt;
|
return std::nullopt;
|
||||||
}
|
}
|
||||||
return execute_graph<T>(gf,
|
reset_compute_ctx();
|
||||||
n_threads,
|
ggml_cgraph* gf = get_compute_graph(get_graph);
|
||||||
free_compute_buffer_immediately,
|
if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
|
||||||
{},
|
LOG_ERROR("%s alloc compute graph failed", get_desc().c_str());
|
||||||
false,
|
return std::nullopt;
|
||||||
no_return);
|
}
|
||||||
|
copy_data_to_backend_tensor();
|
||||||
|
if (ggml_backend_is_cpu(runtime_backend)) {
|
||||||
|
ggml_backend_cpu_set_n_threads(runtime_backend, n_threads);
|
||||||
|
}
|
||||||
|
|
||||||
|
ggml_status status = ggml_backend_graph_compute(runtime_backend, gf);
|
||||||
|
if (status != GGML_STATUS_SUCCESS) {
|
||||||
|
LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
|
||||||
|
return std::nullopt;
|
||||||
|
}
|
||||||
|
copy_cache_tensors_to_cache_buffer();
|
||||||
|
auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
|
||||||
|
std::optional<sd::Tensor<T>> output;
|
||||||
|
if (!no_return) {
|
||||||
|
output = sd::make_sd_tensor_from_ggml<T>(result);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (free_compute_buffer_immediately) {
|
||||||
|
free_compute_buffer();
|
||||||
|
}
|
||||||
|
return output;
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_flash_attention_enabled(bool enabled) {
|
void set_flash_attention_enabled(bool enabled) {
|
||||||
@ -2723,10 +2214,6 @@ public:
|
|||||||
weight_adapter = adapter;
|
weight_adapter = adapter;
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) {
|
|
||||||
max_graph_vram_bytes = max_vram_bytes;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_backend_t get_runtime_backend() {
|
ggml_backend_t get_runtime_backend() {
|
||||||
return runtime_backend;
|
return runtime_backend;
|
||||||
}
|
}
|
||||||
|
|||||||
@ -1,676 +0,0 @@
|
|||||||
#include "ggml_graph_cut.h"
|
|
||||||
|
|
||||||
#include <algorithm>
|
|
||||||
#include <cstring>
|
|
||||||
#include <map>
|
|
||||||
#include <set>
|
|
||||||
#include <sstream>
|
|
||||||
#include <stack>
|
|
||||||
#include <unordered_map>
|
|
||||||
|
|
||||||
#include "ggml-alloc.h"
|
|
||||||
#include "ggml-backend.h"
|
|
||||||
#include "util.h"
|
|
||||||
|
|
||||||
#include "../ggml/src/ggml-impl.h"
|
|
||||||
|
|
||||||
namespace sd::ggml_graph_cut {
|
|
||||||
|
|
||||||
static std::string graph_cut_tensor_display_name(const ggml_tensor* tensor) {
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
return "<null>";
|
|
||||||
}
|
|
||||||
if (tensor->name[0] != '\0') {
|
|
||||||
return tensor->name;
|
|
||||||
}
|
|
||||||
return sd_format("<tensor@%p>", (const void*)tensor);
|
|
||||||
}
|
|
||||||
|
|
||||||
static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
GGML_ASSERT(tensor != nullptr);
|
|
||||||
for (int i = 0; i < gf->n_leafs; ++i) {
|
|
||||||
if (gf->leafs[i] == tensor) {
|
|
||||||
return i;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return -1;
|
|
||||||
}
|
|
||||||
|
|
||||||
static bool is_params_tensor(const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const ggml_tensor* tensor) {
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
return params_tensor_set.find(tensor) != params_tensor_set.end();
|
|
||||||
}
|
|
||||||
|
|
||||||
static Plan::InputShape input_shape(const ggml_tensor* tensor) {
|
|
||||||
Plan::InputShape shape;
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
return shape;
|
|
||||||
}
|
|
||||||
shape.type = tensor->type;
|
|
||||||
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
|
|
||||||
shape.ne[static_cast<size_t>(i)] = tensor->ne[i];
|
|
||||||
}
|
|
||||||
return shape;
|
|
||||||
}
|
|
||||||
|
|
||||||
static size_t graph_cut_segment_vram_bytes(const Segment& segment) {
|
|
||||||
return segment.compute_buffer_size +
|
|
||||||
segment.input_param_bytes +
|
|
||||||
segment.input_previous_cut_bytes +
|
|
||||||
segment.output_bytes;
|
|
||||||
}
|
|
||||||
|
|
||||||
static Segment make_segment_seed(const Plan& plan,
|
|
||||||
size_t start_segment_index,
|
|
||||||
size_t end_segment_index) {
|
|
||||||
GGML_ASSERT(start_segment_index < plan.segments.size());
|
|
||||||
GGML_ASSERT(end_segment_index < plan.segments.size());
|
|
||||||
GGML_ASSERT(start_segment_index <= end_segment_index);
|
|
||||||
|
|
||||||
Segment seed;
|
|
||||||
const auto& start_segment = plan.segments[start_segment_index];
|
|
||||||
const auto& target_segment = plan.segments[end_segment_index];
|
|
||||||
std::unordered_set<int> seen_output_node_indices;
|
|
||||||
for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
|
|
||||||
for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
|
|
||||||
if (seen_output_node_indices.insert(output_node_index).second) {
|
|
||||||
seed.output_node_indices.push_back(output_node_index);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if (start_segment_index == end_segment_index) {
|
|
||||||
seed.group_name = target_segment.group_name;
|
|
||||||
} else {
|
|
||||||
seed.group_name = sd_format("%s..%s",
|
|
||||||
start_segment.group_name.c_str(),
|
|
||||||
target_segment.group_name.c_str());
|
|
||||||
}
|
|
||||||
return seed;
|
|
||||||
}
|
|
||||||
|
|
||||||
static void build_segment(ggml_cgraph* gf,
|
|
||||||
Plan& plan,
|
|
||||||
Segment& segment,
|
|
||||||
const std::unordered_map<const ggml_tensor*, int>& producer_index,
|
|
||||||
std::unordered_set<int>& available_cut_output_node_indices,
|
|
||||||
ggml_backend_t backend,
|
|
||||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const char* log_desc) {
|
|
||||||
std::set<int> internal_nodes;
|
|
||||||
std::unordered_set<const ggml_tensor*> input_seen;
|
|
||||||
std::vector<Segment::InputRef> input_refs;
|
|
||||||
|
|
||||||
std::stack<ggml_tensor*> work_stack;
|
|
||||||
for (int output_node_index : segment.output_node_indices) {
|
|
||||||
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
|
|
||||||
if (output != nullptr) {
|
|
||||||
work_stack.push(output);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
while (!work_stack.empty()) {
|
|
||||||
ggml_tensor* tensor = work_stack.top();
|
|
||||||
work_stack.pop();
|
|
||||||
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
auto producer_it = producer_index.find(tensor);
|
|
||||||
if (producer_it == producer_index.end()) {
|
|
||||||
if (input_seen.insert(tensor).second) {
|
|
||||||
Segment::InputRef input_ref;
|
|
||||||
input_ref.type = is_params_tensor(params_tensor_set, tensor) ? Segment::INPUT_PARAM : Segment::INPUT_EXTERNAL;
|
|
||||||
input_ref.display_name = graph_cut_tensor_display_name(tensor);
|
|
||||||
input_ref.leaf_index = graph_leaf_index(gf, tensor);
|
|
||||||
input_refs.push_back(std::move(input_ref));
|
|
||||||
}
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
int node_idx = producer_it->second;
|
|
||||||
if (available_cut_output_node_indices.find(node_idx) != available_cut_output_node_indices.end()) {
|
|
||||||
if (input_seen.insert(tensor).second) {
|
|
||||||
Segment::InputRef input_ref;
|
|
||||||
input_ref.type = Segment::INPUT_PREVIOUS_CUT;
|
|
||||||
input_ref.display_name = graph_cut_tensor_display_name(tensor);
|
|
||||||
input_ref.node_index = node_idx;
|
|
||||||
input_refs.push_back(std::move(input_ref));
|
|
||||||
}
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
if (!internal_nodes.insert(node_idx).second) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_tensor* node = ggml_graph_node(gf, node_idx);
|
|
||||||
for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
|
|
||||||
if (node->src[src_idx] != nullptr) {
|
|
||||||
work_stack.push(node->src[src_idx]);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if (!internal_nodes.empty()) {
|
|
||||||
segment.internal_node_indices.assign(internal_nodes.begin(), internal_nodes.end());
|
|
||||||
}
|
|
||||||
|
|
||||||
std::sort(input_refs.begin(),
|
|
||||||
input_refs.end(),
|
|
||||||
[](const Segment::InputRef& a, const Segment::InputRef& b) {
|
|
||||||
if (a.type != b.type) {
|
|
||||||
return a.type < b.type;
|
|
||||||
}
|
|
||||||
return a.display_name < b.display_name;
|
|
||||||
});
|
|
||||||
segment.input_refs = input_refs;
|
|
||||||
for (const auto& input : input_refs) {
|
|
||||||
ggml_tensor* current_input = input_tensor(gf, input);
|
|
||||||
size_t tensor_bytes = current_input == nullptr
|
|
||||||
? 0
|
|
||||||
: (input.type == Segment::INPUT_PREVIOUS_CUT
|
|
||||||
? cache_tensor_bytes(current_input)
|
|
||||||
: ggml_nbytes(current_input));
|
|
||||||
switch (input.type) {
|
|
||||||
case Segment::INPUT_PREVIOUS_CUT:
|
|
||||||
segment.input_previous_cut_bytes += tensor_bytes;
|
|
||||||
break;
|
|
||||||
case Segment::INPUT_PARAM:
|
|
||||||
segment.input_param_bytes += tensor_bytes;
|
|
||||||
break;
|
|
||||||
case Segment::INPUT_EXTERNAL:
|
|
||||||
default:
|
|
||||||
segment.input_external_bytes += tensor_bytes;
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for (int output_node_index : segment.output_node_indices) {
|
|
||||||
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
|
|
||||||
segment.output_bytes += cache_tensor_bytes(output);
|
|
||||||
}
|
|
||||||
segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc);
|
|
||||||
|
|
||||||
for (int output_node_index : segment.output_node_indices) {
|
|
||||||
available_cut_output_node_indices.insert(output_node_index);
|
|
||||||
}
|
|
||||||
plan.segments.push_back(std::move(segment));
|
|
||||||
}
|
|
||||||
|
|
||||||
bool is_graph_cut_tensor(const ggml_tensor* tensor) {
|
|
||||||
if (tensor == nullptr || tensor->name[0] == '\0') {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
return std::strncmp(tensor->name, GGML_RUNNER_CUT_PREFIX, std::strlen(GGML_RUNNER_CUT_PREFIX)) == 0;
|
|
||||||
}
|
|
||||||
|
|
||||||
std::string make_graph_cut_name(const std::string& group, const std::string& output) {
|
|
||||||
return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output;
|
|
||||||
}
|
|
||||||
|
|
||||||
void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output) {
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
auto name = make_graph_cut_name(group, output);
|
|
||||||
ggml_set_name(tensor, name.c_str());
|
|
||||||
}
|
|
||||||
|
|
||||||
int leaf_count(ggml_cgraph* gf) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
return gf->n_leafs;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
if (leaf_index < 0 || leaf_index >= gf->n_leafs) {
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
return gf->leafs[leaf_index];
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) {
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
return tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_tensor* cache_source_tensor(ggml_tensor* tensor) {
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
return tensor->view_src ? tensor->view_src : tensor;
|
|
||||||
}
|
|
||||||
|
|
||||||
size_t cache_tensor_bytes(const ggml_tensor* tensor) {
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
const ggml_tensor* cache_src = tensor->view_src ? tensor->view_src : tensor;
|
|
||||||
return ggml_nbytes(cache_src);
|
|
||||||
}
|
|
||||||
|
|
||||||
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
for (const auto& input_shape_ref : plan.input_shapes) {
|
|
||||||
if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index];
|
|
||||||
if (leaf == nullptr || input_shape_ref.type != leaf->type) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
|
|
||||||
if (input_shape_ref.ne[static_cast<size_t>(d)] != leaf->ne[d]) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return true;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
if (output_index >= segment.output_node_indices.size()) {
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
int node_index = segment.output_node_indices[output_index];
|
|
||||||
if (node_index < 0 || node_index >= ggml_graph_n_nodes(gf)) {
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
return ggml_graph_node(gf, node_index);
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
if (input_ref.type == Segment::INPUT_PREVIOUS_CUT) {
|
|
||||||
if (input_ref.node_index < 0 || input_ref.node_index >= ggml_graph_n_nodes(gf)) {
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
return ggml_graph_node(gf, input_ref.node_index);
|
|
||||||
}
|
|
||||||
if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) {
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
return leaf_tensor(gf, input_ref.leaf_index);
|
|
||||||
}
|
|
||||||
|
|
||||||
std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
std::vector<ggml_tensor*> tensors;
|
|
||||||
std::unordered_set<ggml_tensor*> seen_tensors;
|
|
||||||
tensors.reserve(segment.input_refs.size());
|
|
||||||
seen_tensors.reserve(segment.input_refs.size());
|
|
||||||
for (const auto& input_ref : segment.input_refs) {
|
|
||||||
if (input_ref.type != Segment::INPUT_PARAM) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
ggml_tensor* tensor = input_tensor(gf, input_ref);
|
|
||||||
if (tensor == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
if (seen_tensors.insert(tensor).second) {
|
|
||||||
tensors.push_back(tensor);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return tensors;
|
|
||||||
}
|
|
||||||
|
|
||||||
std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc) {
|
|
||||||
std::vector<ggml_tensor*> tensors = param_tensors(gf, segment);
|
|
||||||
std::vector<ggml_tensor*> filtered_tensors;
|
|
||||||
filtered_tensors.reserve(tensors.size());
|
|
||||||
for (ggml_tensor* tensor : tensors) {
|
|
||||||
if (tensor_buffer(tensor) == nullptr) {
|
|
||||||
LOG_WARN("%s graph cut skipping param input without buffer: segment=%s tensor=%s",
|
|
||||||
log_desc == nullptr ? "unknown" : log_desc,
|
|
||||||
segment.group_name.c_str(),
|
|
||||||
tensor->name);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
filtered_tensors.push_back(tensor);
|
|
||||||
}
|
|
||||||
return filtered_tensors;
|
|
||||||
}
|
|
||||||
|
|
||||||
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
|
|
||||||
const Plan& plan,
|
|
||||||
size_t current_segment_index) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
std::unordered_set<std::string> future_input_names;
|
|
||||||
for (size_t seg_idx = current_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
|
|
||||||
const auto& segment = plan.segments[seg_idx];
|
|
||||||
for (const auto& input_ref : segment.input_refs) {
|
|
||||||
if (input_ref.type != Segment::INPUT_PREVIOUS_CUT) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
ggml_tensor* current_input = input_tensor(gf, input_ref);
|
|
||||||
if (current_input != nullptr && current_input->name[0] != '\0') {
|
|
||||||
future_input_names.insert(current_input->name);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return future_input_names;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
|
|
||||||
const Segment& segment,
|
|
||||||
ggml_context** graph_ctx_out) {
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
GGML_ASSERT(graph_ctx_out != nullptr);
|
|
||||||
|
|
||||||
const size_t graph_size = segment.internal_node_indices.size() + segment.input_refs.size() + 8;
|
|
||||||
ggml_init_params params = {
|
|
||||||
/*.mem_size =*/ggml_graph_overhead_custom(graph_size, false) + 1024,
|
|
||||||
/*.mem_buffer =*/nullptr,
|
|
||||||
/*.no_alloc =*/true,
|
|
||||||
};
|
|
||||||
ggml_context* graph_ctx = ggml_init(params);
|
|
||||||
GGML_ASSERT(graph_ctx != nullptr);
|
|
||||||
ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false);
|
|
||||||
GGML_ASSERT(segment_graph != nullptr);
|
|
||||||
|
|
||||||
for (const auto& input : segment.input_refs) {
|
|
||||||
ggml_tensor* current_input = input_tensor(gf, input);
|
|
||||||
if (current_input == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
|
|
||||||
segment_graph->leafs[segment_graph->n_leafs++] = current_input;
|
|
||||||
}
|
|
||||||
|
|
||||||
for (int output_node_index : segment.output_node_indices) {
|
|
||||||
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
|
|
||||||
if (output == nullptr) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
ggml_set_output(output);
|
|
||||||
}
|
|
||||||
for (int node_idx : segment.internal_node_indices) {
|
|
||||||
ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
|
|
||||||
}
|
|
||||||
*graph_ctx_out = graph_ctx;
|
|
||||||
return segment_graph;
|
|
||||||
}
|
|
||||||
|
|
||||||
size_t measure_segment_compute_buffer(ggml_backend_t backend,
|
|
||||||
ggml_cgraph* gf,
|
|
||||||
const Segment& segment,
|
|
||||||
const char* log_desc) {
|
|
||||||
GGML_ASSERT(backend != nullptr);
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
if (segment.internal_node_indices.empty()) {
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_context* graph_ctx = nullptr;
|
|
||||||
ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx);
|
|
||||||
ggml_gallocr_t allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
|
|
||||||
|
|
||||||
size_t sizes[1] = {0};
|
|
||||||
ggml_gallocr_reserve_n_size(
|
|
||||||
allocr,
|
|
||||||
segment_graph,
|
|
||||||
nullptr,
|
|
||||||
nullptr,
|
|
||||||
sizes);
|
|
||||||
size_t buffer_size = sizes[0];
|
|
||||||
|
|
||||||
ggml_gallocr_free(allocr);
|
|
||||||
ggml_free(graph_ctx);
|
|
||||||
return buffer_size;
|
|
||||||
}
|
|
||||||
|
|
||||||
Plan build_plan(ggml_backend_t backend,
|
|
||||||
ggml_cgraph* gf,
|
|
||||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const char* log_desc) {
|
|
||||||
GGML_ASSERT(backend != nullptr);
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
Plan plan;
|
|
||||||
plan.available = true;
|
|
||||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
|
||||||
if (n_nodes <= 0) {
|
|
||||||
return plan;
|
|
||||||
}
|
|
||||||
plan.n_nodes = n_nodes;
|
|
||||||
plan.n_leafs = gf->n_leafs;
|
|
||||||
for (int i = 0; i < gf->n_leafs; ++i) {
|
|
||||||
ggml_tensor* leaf = gf->leafs[i];
|
|
||||||
if (is_params_tensor(params_tensor_set, leaf)) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
auto shape = input_shape(leaf);
|
|
||||||
shape.leaf_index = i;
|
|
||||||
plan.input_shapes.push_back(shape);
|
|
||||||
}
|
|
||||||
|
|
||||||
std::unordered_map<const ggml_tensor*, int> producer_index;
|
|
||||||
producer_index.reserve(static_cast<size_t>(n_nodes));
|
|
||||||
for (int i = 0; i < n_nodes; ++i) {
|
|
||||||
producer_index[ggml_graph_node(gf, i)] = i;
|
|
||||||
}
|
|
||||||
|
|
||||||
std::vector<Segment> grouped_segments;
|
|
||||||
std::unordered_map<std::string, size_t> group_to_segment;
|
|
||||||
for (int i = 0; i < n_nodes; ++i) {
|
|
||||||
ggml_tensor* node = ggml_graph_node(gf, i);
|
|
||||||
if (!is_graph_cut_tensor(node)) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
plan.has_cuts = true;
|
|
||||||
std::string full_name(node->name);
|
|
||||||
std::string payload = full_name.substr(std::strlen(GGML_RUNNER_CUT_PREFIX));
|
|
||||||
size_t sep = payload.find('|');
|
|
||||||
std::string group = sep == std::string::npos ? payload : payload.substr(0, sep);
|
|
||||||
|
|
||||||
auto it = group_to_segment.find(group);
|
|
||||||
if (it == group_to_segment.end()) {
|
|
||||||
Segment segment;
|
|
||||||
segment.group_name = group;
|
|
||||||
segment.output_node_indices.push_back(i);
|
|
||||||
group_to_segment[group] = grouped_segments.size();
|
|
||||||
grouped_segments.push_back(std::move(segment));
|
|
||||||
} else {
|
|
||||||
auto& segment = grouped_segments[it->second];
|
|
||||||
segment.output_node_indices.push_back(i);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if (!plan.has_cuts) {
|
|
||||||
return plan;
|
|
||||||
}
|
|
||||||
|
|
||||||
std::unordered_set<int> available_cut_output_node_indices;
|
|
||||||
available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
|
|
||||||
for (auto& segment : grouped_segments) {
|
|
||||||
build_segment(gf,
|
|
||||||
plan,
|
|
||||||
segment,
|
|
||||||
producer_index,
|
|
||||||
available_cut_output_node_indices,
|
|
||||||
backend,
|
|
||||||
params_tensor_set,
|
|
||||||
log_desc);
|
|
||||||
}
|
|
||||||
|
|
||||||
ggml_tensor* final_output = ggml_graph_node(gf, -1);
|
|
||||||
if (final_output != nullptr && available_cut_output_node_indices.find(n_nodes - 1) == available_cut_output_node_indices.end()) {
|
|
||||||
Segment final_segment;
|
|
||||||
final_segment.group_name = "ggml_runner.final";
|
|
||||||
final_segment.output_node_indices.push_back(n_nodes - 1);
|
|
||||||
build_segment(gf,
|
|
||||||
plan,
|
|
||||||
final_segment,
|
|
||||||
producer_index,
|
|
||||||
available_cut_output_node_indices,
|
|
||||||
backend,
|
|
||||||
params_tensor_set,
|
|
||||||
log_desc);
|
|
||||||
}
|
|
||||||
|
|
||||||
return plan;
|
|
||||||
}
|
|
||||||
|
|
||||||
Plan apply_max_vram_budget(ggml_cgraph* gf,
|
|
||||||
const Plan& base_plan,
|
|
||||||
size_t max_graph_vram_bytes,
|
|
||||||
ggml_backend_t backend,
|
|
||||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const char* log_desc) {
|
|
||||||
GGML_ASSERT(backend != nullptr);
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
int64_t t_budget_begin = ggml_time_ms();
|
|
||||||
if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) {
|
|
||||||
return base_plan;
|
|
||||||
}
|
|
||||||
|
|
||||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
|
||||||
std::unordered_map<const ggml_tensor*, int> producer_index;
|
|
||||||
producer_index.reserve(static_cast<size_t>(n_nodes));
|
|
||||||
for (int i = 0; i < n_nodes; ++i) {
|
|
||||||
producer_index[ggml_graph_node(gf, i)] = i;
|
|
||||||
}
|
|
||||||
|
|
||||||
Plan merged_plan;
|
|
||||||
merged_plan.available = true;
|
|
||||||
merged_plan.has_cuts = base_plan.has_cuts;
|
|
||||||
merged_plan.valid = base_plan.valid;
|
|
||||||
merged_plan.n_nodes = base_plan.n_nodes;
|
|
||||||
merged_plan.n_leafs = base_plan.n_leafs;
|
|
||||||
|
|
||||||
std::unordered_set<int> available_cut_output_node_indices;
|
|
||||||
available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
|
|
||||||
|
|
||||||
size_t start_segment_index = 0;
|
|
||||||
while (start_segment_index < base_plan.segments.size()) {
|
|
||||||
Plan single_plan;
|
|
||||||
auto single_available_cut_output_node_indices = available_cut_output_node_indices;
|
|
||||||
auto single_seed = make_segment_seed(base_plan,
|
|
||||||
start_segment_index,
|
|
||||||
start_segment_index);
|
|
||||||
build_segment(gf,
|
|
||||||
single_plan,
|
|
||||||
single_seed,
|
|
||||||
producer_index,
|
|
||||||
single_available_cut_output_node_indices,
|
|
||||||
backend,
|
|
||||||
params_tensor_set,
|
|
||||||
log_desc);
|
|
||||||
GGML_ASSERT(!single_plan.segments.empty());
|
|
||||||
|
|
||||||
size_t best_end_segment_index = start_segment_index;
|
|
||||||
bool can_merge_next_segment = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes;
|
|
||||||
|
|
||||||
while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) {
|
|
||||||
const size_t next_end_segment_index = best_end_segment_index + 1;
|
|
||||||
Plan candidate_plan;
|
|
||||||
auto candidate_available_cut_output_node_indices = available_cut_output_node_indices;
|
|
||||||
auto candidate_seed = make_segment_seed(base_plan,
|
|
||||||
start_segment_index,
|
|
||||||
next_end_segment_index);
|
|
||||||
build_segment(gf,
|
|
||||||
candidate_plan,
|
|
||||||
candidate_seed,
|
|
||||||
producer_index,
|
|
||||||
candidate_available_cut_output_node_indices,
|
|
||||||
backend,
|
|
||||||
params_tensor_set,
|
|
||||||
log_desc);
|
|
||||||
GGML_ASSERT(!candidate_plan.segments.empty());
|
|
||||||
|
|
||||||
const auto& candidate_segment = candidate_plan.segments.back();
|
|
||||||
if (graph_cut_segment_vram_bytes(candidate_segment) > max_graph_vram_bytes) {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
best_end_segment_index = next_end_segment_index;
|
|
||||||
}
|
|
||||||
|
|
||||||
auto best_seed = make_segment_seed(base_plan,
|
|
||||||
start_segment_index,
|
|
||||||
best_end_segment_index);
|
|
||||||
build_segment(gf,
|
|
||||||
merged_plan,
|
|
||||||
best_seed,
|
|
||||||
producer_index,
|
|
||||||
available_cut_output_node_indices,
|
|
||||||
backend,
|
|
||||||
params_tensor_set,
|
|
||||||
log_desc);
|
|
||||||
start_segment_index = best_end_segment_index + 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) {
|
|
||||||
LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments",
|
|
||||||
log_desc,
|
|
||||||
max_graph_vram_bytes / 1024.0 / 1024.0,
|
|
||||||
base_plan.segments.size(),
|
|
||||||
merged_plan.segments.size());
|
|
||||||
}
|
|
||||||
|
|
||||||
if (log_desc != nullptr) {
|
|
||||||
LOG_INFO("%s graph cut max_vram budget merge took %lld ms",
|
|
||||||
log_desc,
|
|
||||||
ggml_time_ms() - t_budget_begin);
|
|
||||||
}
|
|
||||||
|
|
||||||
return merged_plan;
|
|
||||||
}
|
|
||||||
|
|
||||||
Plan resolve_plan(ggml_backend_t backend,
|
|
||||||
ggml_cgraph* gf,
|
|
||||||
PlanCache* cache,
|
|
||||||
size_t max_graph_vram_bytes,
|
|
||||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const char* log_desc) {
|
|
||||||
GGML_ASSERT(backend != nullptr);
|
|
||||||
GGML_ASSERT(gf != nullptr);
|
|
||||||
GGML_ASSERT(cache != nullptr);
|
|
||||||
|
|
||||||
int64_t t_prepare_begin = ggml_time_ms();
|
|
||||||
Plan base_plan;
|
|
||||||
int64_t t_plan_begin = ggml_time_ms();
|
|
||||||
if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) {
|
|
||||||
base_plan = cache->graph_cut_plan;
|
|
||||||
} else {
|
|
||||||
base_plan = build_plan(backend, gf, params_tensor_set, log_desc);
|
|
||||||
cache->graph_cut_plan = base_plan;
|
|
||||||
cache->graph_cut_plan.available = true;
|
|
||||||
cache->budgeted_graph_cut_plan.available = false;
|
|
||||||
if (log_desc != nullptr) {
|
|
||||||
LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
Plan resolved_plan = base_plan;
|
|
||||||
if (max_graph_vram_bytes > 0 && base_plan.has_cuts) {
|
|
||||||
if (cache->budgeted_graph_cut_plan.available &&
|
|
||||||
cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes &&
|
|
||||||
plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) {
|
|
||||||
resolved_plan = cache->budgeted_graph_cut_plan;
|
|
||||||
} else {
|
|
||||||
resolved_plan = apply_max_vram_budget(gf,
|
|
||||||
base_plan,
|
|
||||||
max_graph_vram_bytes,
|
|
||||||
backend,
|
|
||||||
params_tensor_set,
|
|
||||||
log_desc);
|
|
||||||
cache->budgeted_graph_cut_plan = resolved_plan;
|
|
||||||
cache->budgeted_graph_cut_plan.available = true;
|
|
||||||
cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return resolved_plan;
|
|
||||||
}
|
|
||||||
|
|
||||||
} // namespace sd::ggml_graph_cut
|
|
||||||
@ -1,104 +0,0 @@
|
|||||||
#ifndef __SD_GGML_GRAPH_CUT_H__
|
|
||||||
#define __SD_GGML_GRAPH_CUT_H__
|
|
||||||
|
|
||||||
#include <array>
|
|
||||||
#include <string>
|
|
||||||
#include <unordered_set>
|
|
||||||
#include <vector>
|
|
||||||
|
|
||||||
#include "ggml-backend.h"
|
|
||||||
#include "ggml.h"
|
|
||||||
|
|
||||||
namespace sd::ggml_graph_cut {
|
|
||||||
|
|
||||||
struct Segment {
|
|
||||||
enum InputType {
|
|
||||||
INPUT_EXTERNAL = 0,
|
|
||||||
INPUT_PREVIOUS_CUT,
|
|
||||||
INPUT_PARAM,
|
|
||||||
};
|
|
||||||
|
|
||||||
struct InputRef {
|
|
||||||
InputType type = INPUT_EXTERNAL;
|
|
||||||
std::string display_name;
|
|
||||||
int leaf_index = -1;
|
|
||||||
int node_index = -1;
|
|
||||||
};
|
|
||||||
|
|
||||||
size_t compute_buffer_size = 0;
|
|
||||||
size_t output_bytes = 0;
|
|
||||||
size_t input_external_bytes = 0;
|
|
||||||
size_t input_previous_cut_bytes = 0;
|
|
||||||
size_t input_param_bytes = 0;
|
|
||||||
std::string group_name;
|
|
||||||
std::vector<int> internal_node_indices;
|
|
||||||
std::vector<int> output_node_indices;
|
|
||||||
std::vector<InputRef> input_refs;
|
|
||||||
};
|
|
||||||
|
|
||||||
struct Plan {
|
|
||||||
struct InputShape {
|
|
||||||
int leaf_index = -1;
|
|
||||||
ggml_type type = GGML_TYPE_COUNT;
|
|
||||||
std::array<int64_t, GGML_MAX_DIMS> ne = {0, 0, 0, 0};
|
|
||||||
};
|
|
||||||
|
|
||||||
bool available = false;
|
|
||||||
bool has_cuts = false;
|
|
||||||
bool valid = true;
|
|
||||||
int n_nodes = 0;
|
|
||||||
int n_leafs = 0;
|
|
||||||
std::vector<InputShape> input_shapes;
|
|
||||||
std::vector<Segment> segments;
|
|
||||||
};
|
|
||||||
|
|
||||||
struct PlanCache {
|
|
||||||
Plan graph_cut_plan;
|
|
||||||
Plan budgeted_graph_cut_plan;
|
|
||||||
size_t budgeted_graph_cut_plan_max_vram_bytes = 0;
|
|
||||||
};
|
|
||||||
|
|
||||||
static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";
|
|
||||||
|
|
||||||
bool is_graph_cut_tensor(const ggml_tensor* tensor);
|
|
||||||
std::string make_graph_cut_name(const std::string& group, const std::string& output);
|
|
||||||
void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output);
|
|
||||||
int leaf_count(ggml_cgraph* gf);
|
|
||||||
ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index);
|
|
||||||
ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor);
|
|
||||||
ggml_tensor* cache_source_tensor(ggml_tensor* tensor);
|
|
||||||
size_t cache_tensor_bytes(const ggml_tensor* tensor);
|
|
||||||
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan);
|
|
||||||
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
|
|
||||||
ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
|
|
||||||
std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
|
|
||||||
std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc);
|
|
||||||
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
|
|
||||||
const Plan& plan,
|
|
||||||
size_t current_segment_index);
|
|
||||||
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
|
|
||||||
const Segment& segment,
|
|
||||||
ggml_context** graph_ctx_out);
|
|
||||||
size_t measure_segment_compute_buffer(ggml_backend_t backend,
|
|
||||||
ggml_cgraph* gf,
|
|
||||||
const Segment& segment,
|
|
||||||
const char* log_desc);
|
|
||||||
Plan build_plan(ggml_backend_t backend,
|
|
||||||
ggml_cgraph* gf,
|
|
||||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const char* log_desc);
|
|
||||||
Plan apply_max_vram_budget(ggml_cgraph* gf,
|
|
||||||
const Plan& base_plan,
|
|
||||||
size_t max_graph_vram_bytes,
|
|
||||||
ggml_backend_t backend,
|
|
||||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const char* log_desc);
|
|
||||||
Plan resolve_plan(ggml_backend_t backend,
|
|
||||||
ggml_cgraph* gf,
|
|
||||||
PlanCache* cache,
|
|
||||||
size_t max_graph_vram_bytes,
|
|
||||||
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
|
|
||||||
const char* log_desc);
|
|
||||||
} // namespace sd::ggml_graph_cut
|
|
||||||
|
|
||||||
#endif
|
|
||||||
@ -346,7 +346,6 @@ namespace LLM {
|
|||||||
auto merger = std::dynamic_pointer_cast<PatchMerger>(blocks["merger"]);
|
auto merger = std::dynamic_pointer_cast<PatchMerger>(blocks["merger"]);
|
||||||
|
|
||||||
auto x = patch_embed->forward(ctx, pixel_values);
|
auto x = patch_embed->forward(ctx, pixel_values);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.prelude", "x");
|
|
||||||
|
|
||||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]);
|
x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]);
|
||||||
x = ggml_get_rows(ctx->ggml_ctx, x, window_index);
|
x = ggml_get_rows(ctx->ggml_ctx, x, window_index);
|
||||||
@ -360,11 +359,9 @@ namespace LLM {
|
|||||||
mask = nullptr;
|
mask = nullptr;
|
||||||
}
|
}
|
||||||
x = block->forward(ctx, x, pe, mask);
|
x = block->forward(ctx, x, pe, mask);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
x = merger->forward(ctx, x);
|
x = merger->forward(ctx, x);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.final", "x");
|
|
||||||
|
|
||||||
x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index);
|
x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index);
|
||||||
|
|
||||||
@ -509,7 +506,6 @@ namespace LLM {
|
|||||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
||||||
|
|
||||||
auto x = embed_tokens->forward(ctx, input_ids);
|
auto x = embed_tokens->forward(ctx, input_ids);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.prelude", "x");
|
|
||||||
|
|
||||||
std::vector<ggml_tensor*> intermediate_outputs;
|
std::vector<ggml_tensor*> intermediate_outputs;
|
||||||
|
|
||||||
@ -556,10 +552,6 @@ namespace LLM {
|
|||||||
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
|
||||||
|
|
||||||
x = block->forward(ctx, x, input_pos, attention_mask);
|
x = block->forward(ctx, x, input_pos, attention_mask);
|
||||||
if (out_layers.size() > 1) {
|
|
||||||
x = ggml_cont(ctx->ggml_ctx, x);
|
|
||||||
}
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.layers." + std::to_string(i), "x");
|
|
||||||
if (out_layers.find(i + 1) != out_layers.end()) {
|
if (out_layers.find(i + 1) != out_layers.end()) {
|
||||||
intermediate_outputs.push_back(x);
|
intermediate_outputs.push_back(x);
|
||||||
}
|
}
|
||||||
|
|||||||
@ -767,8 +767,6 @@ public:
|
|||||||
auto context_x = block->forward(ctx, context, x, c_mod);
|
auto context_x = block->forward(ctx, context, x, c_mod);
|
||||||
context = context_x.first;
|
context = context_x.first;
|
||||||
x = context_x.second;
|
x = context_x.second;
|
||||||
sd::ggml_graph_cut::mark_graph_cut(context, "mmdit.joint_blocks." + std::to_string(i), "context");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "mmdit.joint_blocks." + std::to_string(i), "x");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
x = final_layer->forward(ctx, x, c_mod); // (N, T, patch_size ** 2 * out_channels)
|
x = final_layer->forward(ctx, x, c_mod); // (N, T, patch_size ** 2 * out_channels)
|
||||||
@ -811,11 +809,6 @@ public:
|
|||||||
|
|
||||||
context = context_embedder->forward(ctx, context); // [N, L, D] aka [N, L, 1536]
|
context = context_embedder->forward(ctx, context); // [N, L, D] aka [N, L, 1536]
|
||||||
}
|
}
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "mmdit.prelude", "x");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(c, "mmdit.prelude", "c");
|
|
||||||
if (context != nullptr) {
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(context, "mmdit.prelude", "context");
|
|
||||||
}
|
|
||||||
|
|
||||||
x = forward_core_with_concat(ctx, x, c, context, skip_layers); // (N, H*W, patch_size ** 2 * out_channels)
|
x = forward_core_with_concat(ctx, x, c, context, skip_layers); // (N, H*W, patch_size ** 2 * out_channels)
|
||||||
|
|
||||||
|
|||||||
@ -24,75 +24,6 @@ static inline void preprocessing_set_4d(sd::Tensor<float>& tensor, float value,
|
|||||||
tensor.values()[static_cast<size_t>(preprocessing_offset_4d(tensor, i0, i1, i2, i3))] = value;
|
tensor.values()[static_cast<size_t>(preprocessing_offset_4d(tensor, i0, i1, i2, i3))] = value;
|
||||||
}
|
}
|
||||||
|
|
||||||
static inline uint8_t preprocessing_float_to_u8(float value) {
|
|
||||||
if (value <= 0.0f) {
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
if (value >= 1.0f) {
|
|
||||||
return 255;
|
|
||||||
}
|
|
||||||
return static_cast<uint8_t>(value * 255.0f + 0.5f);
|
|
||||||
}
|
|
||||||
|
|
||||||
static inline void preprocessing_tensor_frame_to_sd_image(const sd::Tensor<float>& tensor, int frame_index, uint8_t* image_data) {
|
|
||||||
const auto& shape = tensor.shape();
|
|
||||||
GGML_ASSERT(shape.size() == 4 || shape.size() == 5);
|
|
||||||
GGML_ASSERT(image_data != nullptr);
|
|
||||||
|
|
||||||
const int width = static_cast<int>(shape[0]);
|
|
||||||
const int height = static_cast<int>(shape[1]);
|
|
||||||
const int channel = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
|
|
||||||
const size_t pixels = static_cast<size_t>(width) * static_cast<size_t>(height);
|
|
||||||
const float* src = tensor.data();
|
|
||||||
|
|
||||||
if (shape.size() == 4) {
|
|
||||||
GGML_ASSERT(frame_index >= 0 && frame_index < shape[3]);
|
|
||||||
const size_t frame_stride = pixels * static_cast<size_t>(channel);
|
|
||||||
const float* frame_ptr = src + static_cast<size_t>(frame_index) * frame_stride;
|
|
||||||
if (channel == 3) {
|
|
||||||
const float* c0 = frame_ptr;
|
|
||||||
const float* c1 = frame_ptr + pixels;
|
|
||||||
const float* c2 = frame_ptr + pixels * 2;
|
|
||||||
for (size_t i = 0; i < pixels; ++i) {
|
|
||||||
image_data[i * 3 + 0] = preprocessing_float_to_u8(c0[i]);
|
|
||||||
image_data[i * 3 + 1] = preprocessing_float_to_u8(c1[i]);
|
|
||||||
image_data[i * 3 + 2] = preprocessing_float_to_u8(c2[i]);
|
|
||||||
}
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
for (size_t i = 0; i < pixels; ++i) {
|
|
||||||
for (int c = 0; c < channel; ++c) {
|
|
||||||
image_data[i * static_cast<size_t>(channel) + static_cast<size_t>(c)] =
|
|
||||||
preprocessing_float_to_u8(frame_ptr[i + pixels * static_cast<size_t>(c)]);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
GGML_ASSERT(frame_index >= 0 && frame_index < shape[2]);
|
|
||||||
const size_t channel_stride = pixels * static_cast<size_t>(shape[2]);
|
|
||||||
const float* frame_ptr = src + static_cast<size_t>(frame_index) * pixels;
|
|
||||||
if (channel == 3) {
|
|
||||||
const float* c0 = frame_ptr;
|
|
||||||
const float* c1 = frame_ptr + channel_stride;
|
|
||||||
const float* c2 = frame_ptr + channel_stride * 2;
|
|
||||||
for (size_t i = 0; i < pixels; ++i) {
|
|
||||||
image_data[i * 3 + 0] = preprocessing_float_to_u8(c0[i]);
|
|
||||||
image_data[i * 3 + 1] = preprocessing_float_to_u8(c1[i]);
|
|
||||||
image_data[i * 3 + 2] = preprocessing_float_to_u8(c2[i]);
|
|
||||||
}
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
for (size_t i = 0; i < pixels; ++i) {
|
|
||||||
for (int c = 0; c < channel; ++c) {
|
|
||||||
image_data[i * static_cast<size_t>(channel) + static_cast<size_t>(c)] =
|
|
||||||
preprocessing_float_to_u8(frame_ptr[i + channel_stride * static_cast<size_t>(c)]);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
static inline sd::Tensor<float> sd_image_to_preprocessing_tensor(sd_image_t image) {
|
static inline sd::Tensor<float> sd_image_to_preprocessing_tensor(sd_image_t image) {
|
||||||
sd::Tensor<float> tensor({static_cast<int64_t>(image.width), static_cast<int64_t>(image.height), static_cast<int64_t>(image.channel), 1});
|
sd::Tensor<float> tensor({static_cast<int64_t>(image.width), static_cast<int64_t>(image.height), static_cast<int64_t>(image.channel), 1});
|
||||||
for (uint32_t y = 0; y < image.height; ++y) {
|
for (uint32_t y = 0; y < image.height; ++y) {
|
||||||
@ -108,7 +39,20 @@ static inline sd::Tensor<float> sd_image_to_preprocessing_tensor(sd_image_t imag
|
|||||||
static inline void preprocessing_tensor_to_sd_image(const sd::Tensor<float>& tensor, uint8_t* image_data) {
|
static inline void preprocessing_tensor_to_sd_image(const sd::Tensor<float>& tensor, uint8_t* image_data) {
|
||||||
GGML_ASSERT(tensor.dim() == 4);
|
GGML_ASSERT(tensor.dim() == 4);
|
||||||
GGML_ASSERT(tensor.shape()[3] == 1);
|
GGML_ASSERT(tensor.shape()[3] == 1);
|
||||||
preprocessing_tensor_frame_to_sd_image(tensor, 0, image_data);
|
GGML_ASSERT(image_data != nullptr);
|
||||||
|
|
||||||
|
int width = static_cast<int>(tensor.shape()[0]);
|
||||||
|
int height = static_cast<int>(tensor.shape()[1]);
|
||||||
|
int channel = static_cast<int>(tensor.shape()[2]);
|
||||||
|
for (int y = 0; y < height; ++y) {
|
||||||
|
for (int x = 0; x < width; ++x) {
|
||||||
|
for (int c = 0; c < channel; ++c) {
|
||||||
|
float value = preprocessing_get_4d(tensor, x, y, c, 0);
|
||||||
|
value = std::min(1.0f, std::max(0.0f, value));
|
||||||
|
image_data[(y * width + x) * channel + c] = static_cast<uint8_t>(std::round(value * 255.0f));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
static inline sd::Tensor<float> gaussian_kernel_tensor(int kernel_size) {
|
static inline sd::Tensor<float> gaussian_kernel_tensor(int kernel_size) {
|
||||||
|
|||||||
@ -412,9 +412,6 @@ namespace Qwen {
|
|||||||
auto img = img_in->forward(ctx, x);
|
auto img = img_in->forward(ctx, x);
|
||||||
auto txt = txt_norm->forward(ctx, context);
|
auto txt = txt_norm->forward(ctx, context);
|
||||||
txt = txt_in->forward(ctx, txt);
|
txt = txt_in->forward(ctx, txt);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.prelude", "img");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.prelude", "txt");
|
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb");
|
|
||||||
|
|
||||||
for (int i = 0; i < params.num_layers; i++) {
|
for (int i = 0; i < params.num_layers; i++) {
|
||||||
auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
|
||||||
@ -422,8 +419,6 @@ namespace Qwen {
|
|||||||
auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index);
|
auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index);
|
||||||
img = result.first;
|
img = result.first;
|
||||||
txt = result.second;
|
txt = result.second;
|
||||||
sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.transformer_blocks." + std::to_string(i), "img");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.transformer_blocks." + std::to_string(i), "txt");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if (params.zero_cond_t) {
|
if (params.zero_cond_t) {
|
||||||
|
|||||||
@ -144,7 +144,6 @@ public:
|
|||||||
std::string taesd_path;
|
std::string taesd_path;
|
||||||
sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0, 0};
|
sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0, 0};
|
||||||
bool offload_params_to_cpu = false;
|
bool offload_params_to_cpu = false;
|
||||||
float max_vram = 0.f;
|
|
||||||
bool use_pmid = false;
|
bool use_pmid = false;
|
||||||
|
|
||||||
bool is_using_v_parameterization = false;
|
bool is_using_v_parameterization = false;
|
||||||
@ -191,7 +190,6 @@ public:
|
|||||||
vae_decode_only = sd_ctx_params->vae_decode_only;
|
vae_decode_only = sd_ctx_params->vae_decode_only;
|
||||||
free_params_immediately = sd_ctx_params->free_params_immediately;
|
free_params_immediately = sd_ctx_params->free_params_immediately;
|
||||||
offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu;
|
offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu;
|
||||||
max_vram = sd_ctx_params->max_vram;
|
|
||||||
|
|
||||||
bool use_tae = false;
|
bool use_tae = false;
|
||||||
|
|
||||||
@ -377,10 +375,6 @@ public:
|
|||||||
|
|
||||||
bool clip_on_cpu = sd_ctx_params->keep_clip_on_cpu;
|
bool clip_on_cpu = sd_ctx_params->keep_clip_on_cpu;
|
||||||
|
|
||||||
const size_t max_graph_vram_bytes = max_vram <= 0.f
|
|
||||||
? 0
|
|
||||||
: static_cast<size_t>(static_cast<double>(max_vram) * 1024.0 * 1024.0 * 1024.0);
|
|
||||||
|
|
||||||
{
|
{
|
||||||
clip_backend = backend;
|
clip_backend = backend;
|
||||||
if (clip_on_cpu && !ggml_backend_is_cpu(backend)) {
|
if (clip_on_cpu && !ggml_backend_is_cpu(backend)) {
|
||||||
@ -470,7 +464,6 @@ public:
|
|||||||
clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(backend,
|
clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(backend,
|
||||||
offload_params_to_cpu,
|
offload_params_to_cpu,
|
||||||
tensor_storage_map);
|
tensor_storage_map);
|
||||||
clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
clip_vision->alloc_params_buffer();
|
clip_vision->alloc_params_buffer();
|
||||||
clip_vision->get_param_tensors(tensors);
|
clip_vision->get_param_tensors(tensors);
|
||||||
}
|
}
|
||||||
@ -547,11 +540,9 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
cond_stage_model->alloc_params_buffer();
|
cond_stage_model->alloc_params_buffer();
|
||||||
cond_stage_model->get_param_tensors(tensors);
|
cond_stage_model->get_param_tensors(tensors);
|
||||||
|
|
||||||
diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
diffusion_model->alloc_params_buffer();
|
diffusion_model->alloc_params_buffer();
|
||||||
diffusion_model->get_param_tensors(tensors);
|
diffusion_model->get_param_tensors(tensors);
|
||||||
|
|
||||||
@ -560,7 +551,6 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (high_noise_diffusion_model) {
|
if (high_noise_diffusion_model) {
|
||||||
high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
high_noise_diffusion_model->alloc_params_buffer();
|
high_noise_diffusion_model->alloc_params_buffer();
|
||||||
high_noise_diffusion_model->get_param_tensors(tensors);
|
high_noise_diffusion_model->get_param_tensors(tensors);
|
||||||
}
|
}
|
||||||
@ -633,19 +623,16 @@ public:
|
|||||||
} else if (use_tae && !tae_preview_only) {
|
} else if (use_tae && !tae_preview_only) {
|
||||||
LOG_INFO("using TAE for encoding / decoding");
|
LOG_INFO("using TAE for encoding / decoding");
|
||||||
first_stage_model = create_tae();
|
first_stage_model = create_tae();
|
||||||
first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
first_stage_model->alloc_params_buffer();
|
first_stage_model->alloc_params_buffer();
|
||||||
first_stage_model->get_param_tensors(tensors, "tae");
|
first_stage_model->get_param_tensors(tensors, "tae");
|
||||||
} else {
|
} else {
|
||||||
LOG_INFO("using VAE for encoding / decoding");
|
LOG_INFO("using VAE for encoding / decoding");
|
||||||
first_stage_model = create_vae();
|
first_stage_model = create_vae();
|
||||||
first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
first_stage_model->alloc_params_buffer();
|
first_stage_model->alloc_params_buffer();
|
||||||
first_stage_model->get_param_tensors(tensors, "first_stage_model");
|
first_stage_model->get_param_tensors(tensors, "first_stage_model");
|
||||||
if (use_tae && tae_preview_only) {
|
if (use_tae && tae_preview_only) {
|
||||||
LOG_INFO("using TAE for preview");
|
LOG_INFO("using TAE for preview");
|
||||||
preview_vae = create_tae();
|
preview_vae = create_tae();
|
||||||
preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
preview_vae->alloc_params_buffer();
|
preview_vae->alloc_params_buffer();
|
||||||
preview_vae->get_param_tensors(tensors, "tae");
|
preview_vae->get_param_tensors(tensors, "tae");
|
||||||
}
|
}
|
||||||
@ -1117,13 +1104,8 @@ public:
|
|||||||
cond_stage_lora_models.push_back(lora);
|
cond_stage_lora_models.push_back(lora);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// Only attach the adapter when there are LoRAs targeting the cond_stage model.
|
auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(cond_stage_lora_models);
|
||||||
// An empty MultiLoraAdapter still routes every linear/conv through
|
cond_stage_model->set_weight_adapter(multi_lora_adapter);
|
||||||
// forward_with_lora() instead of the direct kernel path — slower for no benefit.
|
|
||||||
if (!cond_stage_lora_models.empty()) {
|
|
||||||
auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(cond_stage_lora_models);
|
|
||||||
cond_stage_model->set_weight_adapter(multi_lora_adapter);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
if (diffusion_model) {
|
if (diffusion_model) {
|
||||||
std::vector<std::shared_ptr<LoraModel>> lora_models;
|
std::vector<std::shared_ptr<LoraModel>> lora_models;
|
||||||
@ -1154,12 +1136,10 @@ public:
|
|||||||
diffusion_lora_models.push_back(lora);
|
diffusion_lora_models.push_back(lora);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if (!diffusion_lora_models.empty()) {
|
auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(diffusion_lora_models);
|
||||||
auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(diffusion_lora_models);
|
diffusion_model->set_weight_adapter(multi_lora_adapter);
|
||||||
diffusion_model->set_weight_adapter(multi_lora_adapter);
|
if (high_noise_diffusion_model) {
|
||||||
if (high_noise_diffusion_model) {
|
high_noise_diffusion_model->set_weight_adapter(multi_lora_adapter);
|
||||||
high_noise_diffusion_model->set_weight_adapter(multi_lora_adapter);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1192,10 +1172,8 @@ public:
|
|||||||
first_stage_lora_models.push_back(lora);
|
first_stage_lora_models.push_back(lora);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if (!first_stage_lora_models.empty()) {
|
auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(first_stage_lora_models);
|
||||||
auto multi_lora_adapter = std::make_shared<MultiLoraAdapter>(first_stage_lora_models);
|
first_stage_model->set_weight_adapter(multi_lora_adapter);
|
||||||
first_stage_model->set_weight_adapter(multi_lora_adapter);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -2164,7 +2142,6 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
|
|||||||
sd_ctx_params->prediction = PREDICTION_COUNT;
|
sd_ctx_params->prediction = PREDICTION_COUNT;
|
||||||
sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO;
|
sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO;
|
||||||
sd_ctx_params->offload_params_to_cpu = false;
|
sd_ctx_params->offload_params_to_cpu = false;
|
||||||
sd_ctx_params->max_vram = 0.f;
|
|
||||||
sd_ctx_params->enable_mmap = false;
|
sd_ctx_params->enable_mmap = false;
|
||||||
sd_ctx_params->keep_clip_on_cpu = false;
|
sd_ctx_params->keep_clip_on_cpu = false;
|
||||||
sd_ctx_params->keep_control_net_on_cpu = false;
|
sd_ctx_params->keep_control_net_on_cpu = false;
|
||||||
@ -2206,7 +2183,6 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
|
|||||||
"sampler_rng_type: %s\n"
|
"sampler_rng_type: %s\n"
|
||||||
"prediction: %s\n"
|
"prediction: %s\n"
|
||||||
"offload_params_to_cpu: %s\n"
|
"offload_params_to_cpu: %s\n"
|
||||||
"max_vram: %.3f\n"
|
|
||||||
"keep_clip_on_cpu: %s\n"
|
"keep_clip_on_cpu: %s\n"
|
||||||
"keep_control_net_on_cpu: %s\n"
|
"keep_control_net_on_cpu: %s\n"
|
||||||
"keep_vae_on_cpu: %s\n"
|
"keep_vae_on_cpu: %s\n"
|
||||||
@ -2239,7 +2215,6 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
|
|||||||
sd_rng_type_name(sd_ctx_params->sampler_rng_type),
|
sd_rng_type_name(sd_ctx_params->sampler_rng_type),
|
||||||
sd_prediction_name(sd_ctx_params->prediction),
|
sd_prediction_name(sd_ctx_params->prediction),
|
||||||
BOOL_STR(sd_ctx_params->offload_params_to_cpu),
|
BOOL_STR(sd_ctx_params->offload_params_to_cpu),
|
||||||
sd_ctx_params->max_vram,
|
|
||||||
BOOL_STR(sd_ctx_params->keep_clip_on_cpu),
|
BOOL_STR(sd_ctx_params->keep_clip_on_cpu),
|
||||||
BOOL_STR(sd_ctx_params->keep_control_net_on_cpu),
|
BOOL_STR(sd_ctx_params->keep_control_net_on_cpu),
|
||||||
BOOL_STR(sd_ctx_params->keep_vae_on_cpu),
|
BOOL_STR(sd_ctx_params->keep_vae_on_cpu),
|
||||||
@ -3457,13 +3432,9 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
|
|||||||
std::unique_ptr<UpscalerGGML> hires_upscaler;
|
std::unique_ptr<UpscalerGGML> hires_upscaler;
|
||||||
if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
|
if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
|
||||||
LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
|
LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
|
||||||
hires_upscaler = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
|
hires_upscaler = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
|
||||||
false,
|
false,
|
||||||
request.hires.upscale_tile_size);
|
request.hires.upscale_tile_size);
|
||||||
const size_t max_graph_vram_bytes = sd_ctx->sd->max_vram <= 0.f
|
|
||||||
? 0
|
|
||||||
: static_cast<size_t>(static_cast<double>(sd_ctx->sd->max_vram) * 1024.0 * 1024.0 * 1024.0);
|
|
||||||
hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
if (!hires_upscaler->load_from_file(request.hires.model_path,
|
if (!hires_upscaler->load_from_file(request.hires.model_path,
|
||||||
sd_ctx->sd->offload_params_to_cpu,
|
sd_ctx->sd->offload_params_to_cpu,
|
||||||
sd_ctx->sd->n_threads)) {
|
sd_ctx->sd->n_threads)) {
|
||||||
|
|||||||
@ -251,8 +251,7 @@ public:
|
|||||||
ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
ggml_tensor* past_bias = nullptr,
|
ggml_tensor* past_bias = nullptr,
|
||||||
ggml_tensor* attention_mask = nullptr,
|
ggml_tensor* attention_mask = nullptr,
|
||||||
ggml_tensor* relative_position_bucket = nullptr,
|
ggml_tensor* relative_position_bucket = nullptr) {
|
||||||
const std::string& graph_cut_prefix = "") {
|
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
for (int i = 0; i < num_layers; i++) {
|
for (int i = 0; i < num_layers; i++) {
|
||||||
auto block = std::dynamic_pointer_cast<T5Block>(blocks["block." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<T5Block>(blocks["block." + std::to_string(i)]);
|
||||||
@ -260,9 +259,6 @@ public:
|
|||||||
auto ret = block->forward(ctx, x, past_bias, attention_mask, relative_position_bucket);
|
auto ret = block->forward(ctx, x, past_bias, attention_mask, relative_position_bucket);
|
||||||
x = ret.first;
|
x = ret.first;
|
||||||
past_bias = ret.second;
|
past_bias = ret.second;
|
||||||
if (!graph_cut_prefix.empty()) {
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".block." + std::to_string(i), "x");
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
auto final_layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["final_layer_norm"]);
|
auto final_layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["final_layer_norm"]);
|
||||||
@ -309,8 +305,7 @@ public:
|
|||||||
auto encoder = std::dynamic_pointer_cast<T5Stack>(blocks["encoder"]);
|
auto encoder = std::dynamic_pointer_cast<T5Stack>(blocks["encoder"]);
|
||||||
|
|
||||||
auto x = shared->forward(ctx, input_ids);
|
auto x = shared->forward(ctx, input_ids);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "t5.prelude", "x");
|
x = encoder->forward(ctx, x, past_bias, attention_mask, relative_position_bucket);
|
||||||
x = encoder->forward(ctx, x, past_bias, attention_mask, relative_position_bucket, "t5");
|
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
@ -482,14 +482,12 @@ public:
|
|||||||
|
|
||||||
emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim]
|
emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim]
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(emb, "unet.prelude", "emb");
|
|
||||||
|
|
||||||
// input_blocks
|
// input_blocks
|
||||||
std::vector<ggml_tensor*> hs;
|
std::vector<ggml_tensor*> hs;
|
||||||
|
|
||||||
// input block 0
|
// input block 0
|
||||||
auto h = input_blocks_0_0->forward(ctx, x);
|
auto h = input_blocks_0_0->forward(ctx, x);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks.0", "h");
|
|
||||||
|
|
||||||
ggml_set_name(h, "bench-start");
|
ggml_set_name(h, "bench-start");
|
||||||
hs.push_back(h);
|
hs.push_back(h);
|
||||||
@ -507,7 +505,6 @@ public:
|
|||||||
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
|
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
|
||||||
h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w]
|
h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w]
|
||||||
}
|
}
|
||||||
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
|
|
||||||
hs.push_back(h);
|
hs.push_back(h);
|
||||||
}
|
}
|
||||||
if (tiny_unet) {
|
if (tiny_unet) {
|
||||||
@ -521,7 +518,6 @@ public:
|
|||||||
auto block = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
|
auto block = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
|
||||||
|
|
||||||
h = block->forward(ctx, h); // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
|
h = block->forward(ctx, h); // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
|
|
||||||
hs.push_back(h);
|
hs.push_back(h);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -535,7 +531,6 @@ public:
|
|||||||
h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
|
h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
|
|
||||||
if (controls.size() > 0) {
|
if (controls.size() > 0) {
|
||||||
auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true);
|
auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true);
|
||||||
h = ggml_add(ctx->ggml_ctx, h, cs); // middle control
|
h = ggml_add(ctx->ggml_ctx, h, cs); // middle control
|
||||||
@ -586,7 +581,6 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
output_block_idx += 1;
|
output_block_idx += 1;
|
||||||
sd::ggml_graph_cut::mark_graph_cut(h, "unet.output_blocks." + std::to_string(output_block_idx - 1), "h");
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -12,13 +12,6 @@ UpscalerGGML::UpscalerGGML(int n_threads,
|
|||||||
tile_size(tile_size) {
|
tile_size(tile_size) {
|
||||||
}
|
}
|
||||||
|
|
||||||
void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
|
|
||||||
max_graph_vram_bytes = max_vram_bytes;
|
|
||||||
if (esrgan_upscaler) {
|
|
||||||
esrgan_upscaler->set_max_graph_vram_bytes(max_vram_bytes);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
|
bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
|
||||||
bool offload_params_to_cpu,
|
bool offload_params_to_cpu,
|
||||||
int n_threads) {
|
int n_threads) {
|
||||||
@ -37,7 +30,6 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
|
|||||||
}
|
}
|
||||||
LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
|
LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
|
||||||
esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
|
esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
|
||||||
esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
|
|
||||||
if (direct) {
|
if (direct) {
|
||||||
esrgan_upscaler->set_conv2d_direct_enabled(true);
|
esrgan_upscaler->set_conv2d_direct_enabled(true);
|
||||||
}
|
}
|
||||||
|
|||||||
@ -14,9 +14,8 @@ struct UpscalerGGML {
|
|||||||
std::shared_ptr<ESRGAN> esrgan_upscaler;
|
std::shared_ptr<ESRGAN> esrgan_upscaler;
|
||||||
std::string esrgan_path;
|
std::string esrgan_path;
|
||||||
int n_threads;
|
int n_threads;
|
||||||
bool direct = false;
|
bool direct = false;
|
||||||
int tile_size = 128;
|
int tile_size = 128;
|
||||||
size_t max_graph_vram_bytes = 0;
|
|
||||||
|
|
||||||
UpscalerGGML(int n_threads,
|
UpscalerGGML(int n_threads,
|
||||||
bool direct = false,
|
bool direct = false,
|
||||||
@ -25,7 +24,6 @@ struct UpscalerGGML {
|
|||||||
bool load_from_file(const std::string& esrgan_path,
|
bool load_from_file(const std::string& esrgan_path,
|
||||||
bool offload_params_to_cpu,
|
bool offload_params_to_cpu,
|
||||||
int n_threads);
|
int n_threads);
|
||||||
void set_max_graph_vram_bytes(size_t max_vram_bytes);
|
|
||||||
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
|
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
|
||||||
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
|
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
|
||||||
};
|
};
|
||||||
|
|||||||
12
src/util.cpp
12
src/util.cpp
@ -505,7 +505,17 @@ sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index)
|
|||||||
int channel = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
|
int channel = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
|
||||||
uint8_t* data = (uint8_t*)malloc(static_cast<size_t>(width * height * channel));
|
uint8_t* data = (uint8_t*)malloc(static_cast<size_t>(width * height * channel));
|
||||||
GGML_ASSERT(data != nullptr);
|
GGML_ASSERT(data != nullptr);
|
||||||
preprocessing_tensor_frame_to_sd_image(tensor, frame_index, data);
|
|
||||||
|
for (int iw = 0; iw < width; ++iw) {
|
||||||
|
for (int ih = 0; ih < height; ++ih) {
|
||||||
|
for (int ic = 0; ic < channel; ++ic) {
|
||||||
|
float value = shape.size() == 5 ? tensor.index(iw, ih, frame_index, ic, 0)
|
||||||
|
: tensor.index(iw, ih, ic, frame_index);
|
||||||
|
value = std::clamp(value, 0.0f, 1.0f);
|
||||||
|
data[(ih * width + iw) * channel + ic] = static_cast<uint8_t>(std::round(value * 255.0f));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
return {
|
return {
|
||||||
static_cast<uint32_t>(width),
|
static_cast<uint32_t>(width),
|
||||||
static_cast<uint32_t>(height),
|
static_cast<uint32_t>(height),
|
||||||
|
|||||||
25
src/wan.hpp
25
src/wan.hpp
@ -692,7 +692,6 @@ namespace WAN {
|
|||||||
} else {
|
} else {
|
||||||
x = conv1->forward(ctx, x);
|
x = conv1->forward(ctx, x);
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.prelude", "x");
|
|
||||||
|
|
||||||
// downsamples
|
// downsamples
|
||||||
std::vector<int64_t> dims = {dim};
|
std::vector<int64_t> dims = {dim};
|
||||||
@ -718,14 +717,12 @@ namespace WAN {
|
|||||||
x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx);
|
x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.down." + std::to_string(i), "x");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// middle
|
// middle
|
||||||
x = middle_0->forward(ctx, x, b, feat_cache, feat_idx);
|
x = middle_0->forward(ctx, x, b, feat_cache, feat_idx);
|
||||||
x = middle_1->forward(ctx, x, b);
|
x = middle_1->forward(ctx, x, b);
|
||||||
x = middle_2->forward(ctx, x, b, feat_cache, feat_idx);
|
x = middle_2->forward(ctx, x, b, feat_cache, feat_idx);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.mid", "x");
|
|
||||||
|
|
||||||
// head
|
// head
|
||||||
x = head_0->forward(ctx, x);
|
x = head_0->forward(ctx, x);
|
||||||
@ -866,13 +863,11 @@ namespace WAN {
|
|||||||
} else {
|
} else {
|
||||||
x = conv1->forward(ctx, x);
|
x = conv1->forward(ctx, x);
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.prelude", "x");
|
|
||||||
|
|
||||||
// middle
|
// middle
|
||||||
x = middle_0->forward(ctx, x, b, feat_cache, feat_idx);
|
x = middle_0->forward(ctx, x, b, feat_cache, feat_idx);
|
||||||
x = middle_1->forward(ctx, x, b);
|
x = middle_1->forward(ctx, x, b);
|
||||||
x = middle_2->forward(ctx, x, b, feat_cache, feat_idx);
|
x = middle_2->forward(ctx, x, b, feat_cache, feat_idx);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.mid", "x");
|
|
||||||
|
|
||||||
// upsamples
|
// upsamples
|
||||||
std::vector<int64_t> dims = {dim_mult[dim_mult.size() - 1] * dim};
|
std::vector<int64_t> dims = {dim_mult[dim_mult.size() - 1] * dim};
|
||||||
@ -898,7 +893,6 @@ namespace WAN {
|
|||||||
x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx);
|
x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.up." + std::to_string(i), "x");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// head
|
// head
|
||||||
@ -1037,7 +1031,6 @@ namespace WAN {
|
|||||||
if (wan2_2) {
|
if (wan2_2) {
|
||||||
x = patchify(ctx->ggml_ctx, x, 2, b);
|
x = patchify(ctx->ggml_ctx, x, 2, b);
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
|
|
||||||
|
|
||||||
auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
|
auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
|
||||||
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
|
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
|
||||||
@ -1058,7 +1051,6 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
out = conv1->forward(ctx, out);
|
out = conv1->forward(ctx, out);
|
||||||
auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0];
|
auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0];
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu");
|
|
||||||
clear_cache();
|
clear_cache();
|
||||||
return mu;
|
return mu;
|
||||||
}
|
}
|
||||||
@ -1076,7 +1068,6 @@ namespace WAN {
|
|||||||
|
|
||||||
int64_t iter_ = z->ne[2];
|
int64_t iter_ = z->ne[2];
|
||||||
auto x = conv2->forward(ctx, z);
|
auto x = conv2->forward(ctx, z);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode.prelude", "x");
|
|
||||||
ggml_tensor* out;
|
ggml_tensor* out;
|
||||||
for (int i = 0; i < iter_; i++) {
|
for (int i = 0; i < iter_; i++) {
|
||||||
_conv_idx = 0;
|
_conv_idx = 0;
|
||||||
@ -1092,7 +1083,6 @@ namespace WAN {
|
|||||||
if (wan2_2) {
|
if (wan2_2) {
|
||||||
out = unpatchify(ctx->ggml_ctx, out, 2, b);
|
out = unpatchify(ctx->ggml_ctx, out, 2, b);
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
|
|
||||||
clear_cache();
|
clear_cache();
|
||||||
return out;
|
return out;
|
||||||
}
|
}
|
||||||
@ -1107,15 +1097,13 @@ namespace WAN {
|
|||||||
auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
|
auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
|
||||||
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
|
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
|
||||||
|
|
||||||
auto x = conv2->forward(ctx, z);
|
auto x = conv2->forward(ctx, z);
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode_partial.prelude", "x");
|
|
||||||
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w]
|
auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w]
|
||||||
_conv_idx = 0;
|
_conv_idx = 0;
|
||||||
auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
|
auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
|
||||||
if (wan2_2) {
|
if (wan2_2) {
|
||||||
out = unpatchify(ctx->ggml_ctx, out, 2, b);
|
out = unpatchify(ctx->ggml_ctx, out, 2, b);
|
||||||
}
|
}
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
|
|
||||||
return out;
|
return out;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -1996,13 +1984,6 @@ namespace WAN {
|
|||||||
c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N); // [N, dim, t_len*h_len*w_len]
|
c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N); // [N, dim, t_len*h_len*w_len]
|
||||||
c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim]
|
c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim]
|
||||||
}
|
}
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x");
|
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e");
|
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(e0, "wan.prelude", "e0");
|
|
||||||
// sd::ggml_graph_cut::mark_graph_cut(context, "wan.prelude", "context");
|
|
||||||
if (c != nullptr) {
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(c, "wan.prelude", "c");
|
|
||||||
}
|
|
||||||
|
|
||||||
auto x_orig = x;
|
auto x_orig = x;
|
||||||
|
|
||||||
@ -2023,10 +2004,6 @@ namespace WAN {
|
|||||||
c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
|
c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
|
||||||
x = ggml_add(ctx->ggml_ctx, x, c_skip);
|
x = ggml_add(ctx->ggml_ctx, x, c_skip);
|
||||||
}
|
}
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x");
|
|
||||||
if (c != nullptr) {
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c");
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
x = head->forward(ctx, x, e); // [N, t_len*h_len*w_len, pt*ph*pw*out_dim]
|
x = head->forward(ctx, x, e); // [N, t_len*h_len*w_len, pt*ph*pw*out_dim]
|
||||||
|
|||||||
@ -371,9 +371,6 @@ namespace ZImage {
|
|||||||
|
|
||||||
auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size]
|
auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size]
|
||||||
auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size]
|
auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size]
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.prelude", "txt");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(img, "z_image.prelude", "img");
|
|
||||||
sd::ggml_graph_cut::mark_graph_cut(t_emb, "z_image.prelude", "t_emb");
|
|
||||||
|
|
||||||
int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), SEQ_MULTI_OF);
|
int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), SEQ_MULTI_OF);
|
||||||
if (n_txt_pad_token > 0) {
|
if (n_txt_pad_token > 0) {
|
||||||
@ -396,24 +393,20 @@ namespace ZImage {
|
|||||||
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
|
||||||
|
|
||||||
txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
|
txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.context_refiner." + std::to_string(i), "txt");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
|
for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
|
||||||
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
|
||||||
|
|
||||||
img = block->forward(ctx, img, img_pe, nullptr, t_emb);
|
img = block->forward(ctx, img, img_pe, nullptr, t_emb);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(img, "z_image.noise_refiner." + std::to_string(i), "img");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size]
|
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size]
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img");
|
|
||||||
|
|
||||||
for (int i = 0; i < z_image_params.num_layers; i++) {
|
for (int i = 0; i < z_image_params.num_layers; i++) {
|
||||||
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
|
||||||
|
|
||||||
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
|
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
|
||||||
sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.layers." + std::to_string(i), "txt_img");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
txt_img = final_layer->forward(ctx, txt_img, t_emb); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C]
|
txt_img = final_layer->forward(ctx, txt_img, t_emb); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C]
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user