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

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

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

No files matched your search

src/model/adapter/lora.hpp
+41 -10
View File
@@ -71,7 +71,8 @@ struct LoraModel : public GGMLRunner {
return true;
};
model_loader.load_tensors(on_new_tensor_cb, n_threads);
model_loader.set_n_threads(n_threads);
model_loader.load_tensors(on_new_tensor_cb);
if (tensors_to_create.empty()) {
return true;
@@ -93,19 +94,39 @@ struct LoraModel : public GGMLRunner {
}
dry_run = false;
model_loader.load_tensors(on_new_tensor_cb, n_threads);
model_loader.load_tensors(on_new_tensor_cb);
LOG_DEBUG("finished loaded lora");
return true;
}
void preprocess_lora_tensors(const std::map<std::string, ggml_tensor*>& model_tensors) {
void release_loaded_tensors() {
free_compute_buffer();
free_params_buffer();
free_params_ctx();
alloc_params_ctx();
lora_tensors.clear();
original_tensor_to_final_tensor.clear();
applied_lora_tensors.clear();
applied = false;
tensor_preprocessed = false;
}
static std::set<std::string> tensor_names(const std::map<std::string, ggml_tensor*>& model_tensors) {
std::set<std::string> names;
for (const auto& item : model_tensors) {
names.insert(item.first);
}
return names;
}
void preprocess_lora_tensors(const std::set<std::string>& model_tensor_names) {
if (tensor_preprocessed) {
return;
}
tensor_preprocessed = true;
// I really hate these hardcoded processes.
if (model_tensors.find("cond_stage_model.1.transformer.text_model.encoder.layers.0.self_attn.in_proj.weight") != model_tensors.end()) {
if (model_tensor_names.find("cond_stage_model.1.transformer.text_model.encoder.layers.0.self_attn.in_proj.weight") != model_tensor_names.end()) {
std::unordered_map<std::string, ggml_tensor*> new_lora_tensors;
for (auto& [old_name, tensor] : lora_tensors) {
std::string new_name = old_name;
@@ -753,11 +774,13 @@ struct LoraModel : public GGMLRunner {
return out_diff;
}
ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors,
const std::set<std::string>& model_tensor_names,
SDVersion version) {
size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
preprocess_lora_tensors(model_tensors);
preprocess_lora_tensors(model_tensor_names);
original_tensor_to_final_tensor.clear();
applied_lora_tensors.clear();
@@ -794,12 +817,16 @@ struct LoraModel : public GGMLRunner {
return gf;
}
void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
void apply(std::map<std::string, ggml_tensor*> model_tensors,
const std::set<std::string>& model_tensor_names,
SDVersion version,
int n_threads,
bool warn_unused = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_lora_graph(model_tensors, version);
return build_lora_graph(model_tensors, model_tensor_names, version);
};
GGMLRunner::compute<float>(get_graph, n_threads, false, true);
stat();
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
stat(!warn_unused);
for (auto item : original_tensor_to_final_tensor) {
ggml_tensor* original_tensor = item.first;
ggml_tensor* final_tensor = item.second;
@@ -810,6 +837,10 @@ struct LoraModel : public GGMLRunner {
GGMLRunner::free_compute_buffer();
}
void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
}
void stat(bool at_runntime = false) {
size_t total_lora_tensors_count = 0;
size_t applied_lora_tensors_count = 0;
src/model/adapter/pmid.hpp
+4 -3
View File
@@ -558,7 +558,7 @@ public:
return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
}
};
@@ -616,14 +616,15 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
return true;
};
model_loader->load_tensors(on_new_tensor_cb, n_threads);
model_loader->set_n_threads(n_threads);
model_loader->load_tensors(on_new_tensor_cb);
if (!alloc_params_buffer()) {
LOG_ERROR("PhotoMaker ID embeds buffer allocation failed");
return false;
}
dry_run = false;
model_loader->load_tensors(on_new_tensor_cb, n_threads);
model_loader->load_tensors(on_new_tensor_cb);
LOG_DEBUG("finished loading PhotoMaker ID Embeds ");
return true;
src/model/diffusion/anima.hpp
+1 -1
View File
@@ -697,7 +697,7 @@ namespace Anima {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, t5_ids, t5_weights);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/control.hpp
+10 -7
View File
@@ -309,6 +309,7 @@ public:
struct ControlNet : public GGMLRunner {
SDVersion version = VERSION_SD1;
ControlNetBlock control_net;
std::string weight_prefix;
ggml_backend_buffer_t control_buffer = nullptr;
ggml_context* control_ctx = nullptr;
@@ -321,9 +322,10 @@ struct ControlNet : public GGMLRunner {
ControlNet(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
SDVersion version = VERSION_SD1)
: GGMLRunner(backend, params_backend), control_net(version) {
control_net.init(params_ctx, tensor_storage_map, "");
SDVersion version = VERSION_SD1,
const std::string& prefix = "")
: GGMLRunner(backend, params_backend), version(version), control_net(version), weight_prefix(prefix) {
control_net.init(params_ctx, tensor_storage_map, prefix);
}
~ControlNet() override {
@@ -374,8 +376,8 @@ struct ControlNet : public GGMLRunner {
return "control_net";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
control_net.get_param_tensors(tensors, prefix);
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
control_net.get_param_tensors(tensors, weight_prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@@ -435,7 +437,7 @@ struct ControlNet : public GGMLRunner {
return build_graph(x, hint, timesteps, context, y);
};
auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false);
auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, false, false);
if (!compute_result.has_value()) {
return std::nullopt;
}
@@ -472,7 +474,8 @@ struct ControlNet : public GGMLRunner {
return false;
}
bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads);
model_loader.set_n_threads(n_threads);
bool success = model_loader.load_tensors(tensors, ignore_tensors);
if (!success) {
LOG_ERROR("load control net tensors from model loader failed");
src/model/diffusion/ernie_image.hpp
+1 -1
View File
@@ -440,7 +440,7 @@ namespace ErnieImage {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/flux.hpp
+1 -1
View File
@@ -1500,7 +1500,7 @@ namespace Flux {
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return result;
}
src/model/diffusion/hidream_o1.hpp
+16 -19
View File
@@ -323,11 +323,15 @@ namespace HiDreamO1 {
return gf;
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image) {
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& image,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
auto get_graph = [&]() {
return build_graph(image);
};
auto output = GGMLRunner::compute<float>(get_graph, n_threads, false);
auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
}
};
@@ -455,7 +459,7 @@ namespace HiDreamO1 {
auto get_graph = [&]() {
return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
@@ -494,21 +498,6 @@ namespace HiDreamO1 {
vision_runner->get_param_tensors(tensors);
}
bool alloc_params_buffer() override {
if (!vision_runner->alloc_params_buffer()) {
return false;
}
return true;
}
void free_params_buffer() override {
vision_runner->free_params_buffer();
}
size_t get_params_buffer_size() override {
return vision_runner->get_params_buffer_size();
}
void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override {
vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes);
}
@@ -521,6 +510,14 @@ namespace HiDreamO1 {
vision_runner->set_weight_adapter(adapter);
}
void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
vision_runner->set_weight_manager(manager);
}
void runner_done() override {
vision_runner->runner_done();
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
SDCondition result;
@@ -666,7 +663,7 @@ namespace HiDreamO1 {
result.c_vinput_mask = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
result.c_image_embeds.reserve(vlm_images.size());
for (const auto& vlm_image : vlm_images) {
auto image_embed = vision_runner->compute(n_threads, vlm_image.second);
auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true);
if (image_embed.empty()) {
LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
return SDCondition();
src/model/diffusion/ideogram4.hpp
+1 -1
View File
@@ -537,7 +537,7 @@ namespace Ideogram4 {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, use_uncond_model);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/lens.hpp
+1 -1
View File
@@ -408,7 +408,7 @@ namespace Lens {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/ltxv.hpp
+1 -1
View File
@@ -1939,7 +1939,7 @@ namespace LTXV {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions);
};
auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return out;
}
src/model/diffusion/mmdit.hpp
+1 -1
View File
@@ -935,7 +935,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
return build_graph(x, timesteps, context, y, skip_layers);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/pid.hpp
+1 -1
View File
@@ -823,7 +823,7 @@ namespace Pid {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/qwen_image.hpp
+1 -1
View File
@@ -627,7 +627,7 @@ namespace Qwen {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/unet.hpp
+1 -1
View File
@@ -772,7 +772,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/wan.hpp
+1 -1
View File
@@ -950,7 +950,7 @@ namespace WAN {
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/z_image.hpp
+1 -1
View File
@@ -634,7 +634,7 @@ namespace ZImage {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/te/clip.hpp
+5 -2
View File
@@ -567,11 +567,14 @@ struct CLIPTextModelRunner : public GGMLRunner {
void* custom_embeddings_data,
size_t max_token_idx,
bool return_pooled,
int clip_skip) {
int clip_skip,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
};
auto result = GGMLRunner::compute<float>(get_graph, n_threads, true);
auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
if (return_pooled) {
return take_or_empty(std::move(result));
}
src/model/te/llm.hpp
+10 -4
View File
@@ -1733,7 +1733,10 @@ namespace LLM {
const sd::Tensor<float>& attention_mask,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids,
attention_mask,
@@ -1741,7 +1744,7 @@ namespace LLM {
out_layers,
return_all_hidden_states);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params),
input_ids.dim() + 1);
}
@@ -1802,11 +1805,14 @@ namespace LLM {
}
sd::Tensor<float> encode_image(const int n_threads,
const sd::Tensor<float>& image) {
const sd::Tensor<float>& image,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_image_graph(image);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
}
};
src/model/te/t5.hpp
+5 -2
View File
@@ -394,11 +394,14 @@ struct T5Runner : public GGMLRunner {
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<float>& attention_mask) {
const sd::Tensor<float>& attention_mask,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids, attention_mask);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true), 3);
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3);
}
static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,
src/model/upscaler/esrgan.hpp
+5 -3
View File
@@ -336,9 +336,11 @@ struct ESRGAN : public GGMLRunner {
}
}
success = model_loader.load_tensors(model_tensors, {}, n_threads);
model_loader.set_n_threads(n_threads);
success = model_loader.load_tensors(model_tensors);
} else {
success = model_loader.load_tensors(esrgan_tensors, {}, n_threads);
model_loader.set_n_threads(n_threads);
success = model_loader.load_tensors(esrgan_tensors);
}
if (!success) {
@@ -367,7 +369,7 @@ struct ESRGAN : public GGMLRunner {
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<float>& x) {
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return result;
}
};
src/model/upscaler/ltx_latent_upscaler.hpp
+17 -36
View File
@@ -240,20 +240,25 @@ namespace LTXVUpsampler {
protected:
int64_t channels;
int stride;
ggml_tensor* kernel = nullptr;
std::vector<float> kernel_data;
std::string kernel_name;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(ctx);
SD_UNUSED(tensor_storage_map);
if (stride == 1) {
return;
}
kernel = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 5, 1, channels);
std::string name = prefix + "kernel";
ggml_set_name(kernel, name.c_str());
kernel_name = prefix + "kernel";
}
public:
BlurDownsample(int64_t channels, int stride)
: channels(channels),
stride(stride) {
GGML_ASSERT(stride >= 1);
static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f};
kernel_data.resize(static_cast<size_t>(5 * 5 * channels));
for (int64_t c = 0; c < channels; ++c) {
@@ -266,26 +271,16 @@ namespace LTXVUpsampler {
}
}
public:
BlurDownsample(int64_t channels, int stride)
: channels(channels),
stride(stride) {
GGML_ASSERT(stride >= 1);
}
void load_fixed_tensors() {
if (kernel == nullptr || kernel_data.empty()) {
return;
}
ggml_backend_tensor_set(kernel, kernel_data.data(), 0, kernel_data.size() * sizeof(float));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
if (stride == 1) {
return x;
}
GGML_ASSERT(kernel != nullptr);
GGML_ASSERT(ctx != nullptr);
GGML_ASSERT(!kernel_data.empty());
GGML_ASSERT(x->ne[2] == channels);
ggml_tensor* kernel = ggml_new_tensor_4d(ctx->ggml_ctx, GGML_TYPE_F32, 5, 5, 1, channels);
ggml_set_name(kernel, kernel_name.empty() ? "blur_down.kernel" : kernel_name.c_str());
ctx->bind_backend_tensor_data(kernel, kernel_data.data());
if (ctx->conv2d_direct_enabled) {
return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
}
@@ -311,11 +306,6 @@ namespace LTXVUpsampler {
blocks["blur_down"] = std::shared_ptr<GGMLBlock>(new BlurDownsample(mid_channels, den));
}
void load_fixed_tensors() {
auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
blur_down->load_fixed_tensors();
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["pixel_shuffle"]);
@@ -426,14 +416,6 @@ namespace LTXVUpsampler {
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.final", "x");
return x;
}
void load_fixed_tensors() {
if (!config.rational_resampler) {
return;
}
auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
upsampler->load_fixed_tensors();
}
};
struct LatentUpsamplerRunner : public GGMLRunner {
@@ -490,12 +472,11 @@ namespace LTXVUpsampler {
if (config.rational_resampler) {
ignore_tensors.insert("upsampler.blur_down.kernel");
}
if (!model_loader.load_tensors(tensors, ignore_tensors, n_threads)) {
model_loader.set_n_threads(n_threads);
if (!model_loader.load_tensors(tensors, ignore_tensors)) {
LOG_ERROR("load LTX latent upsampler tensors failed");
return false;
}
model->load_fixed_tensors();
LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d, scale=%.3f, temporal_factor=%d, rational=%d",
config.in_channels,
config.mid_channels,
@@ -542,7 +523,7 @@ namespace LTXVUpsampler {
}
size_t expected_dim = static_cast<size_t>(x.dim());
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
}
};
src/model/vae/auto_encoder_kl.hpp
+4 -4
View File
@@ -670,7 +670,7 @@ struct AutoEncoderKL : public VAE {
bool decode_only = false,
bool use_video_decoder = false,
SDVersion version = VERSION_SD1)
: decode_only(decode_only), VAE(version, backend, params_backend) {
: VAE(version, backend, params_backend, prefix), decode_only(decode_only) {
if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
scale_factor = 0.18215f;
shift_factor = 0.f;
@@ -718,8 +718,8 @@ struct AutoEncoderKL : public VAE {
return "vae";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
ae.get_param_tensors(tensors, prefix);
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
ae.get_param_tensors(tensors, weight_prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
@@ -742,7 +742,7 @@ struct AutoEncoderKL : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(z, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z.dim());
}
sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {
src/model/vae/ltx_audio_vae.hpp
+6 -4
View File
@@ -997,6 +997,7 @@ namespace LTXV {
struct LTXAudioVAERunner : public GGMLRunner {
LTXAudioVAEConfig config;
LTXAudioVAE model;
std::string weight_prefix;
sd::Tensor<float> bwe_skip_filter_tensor;
LTXAudioVAERunner(ggml_backend_t backend,
@@ -1004,6 +1005,7 @@ namespace LTXV {
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "")
: GGMLRunner(backend, params_backend),
weight_prefix(prefix),
config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
model(config) {
model.init(params_ctx, tensor_storage_map, prefix);
@@ -1013,8 +1015,8 @@ namespace LTXV {
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
model.get_param_tensors(tensors, prefix);
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
model.get_param_tensors(tensors, weight_prefix);
}
size_t get_params_buffer_size() {
@@ -1037,7 +1039,7 @@ namespace LTXV {
ggml_build_forward_expand(gf, waveform);
return gf;
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), 4);
int64_t t1 = ggml_time_ms();
LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
return result;
@@ -1082,7 +1084,7 @@ namespace LTXV {
}
std::map<std::string, ggml_tensor*> tensors;
ltx_audio_vae->get_param_tensors(tensors, "");
ltx_audio_vae->get_param_tensors(tensors);
if (!model_loader.load_tensors(tensors)) {
LOG_ERROR("load tensors from model loader failed");
src/model/vae/ltx_vae.hpp
+7 -7
View File
@@ -1239,7 +1239,7 @@ struct LTXVideoVAE : public VAE {
patch_size,
tensor_storage_map,
prefix),
VAE(version, backend, params_backend) {
VAE(version, backend, params_backend, prefix) {
vae.init(params_ctx, tensor_storage_map, prefix);
decode_timestep_tensor.values()[0] = vae.decode_timestep;
}
@@ -1271,8 +1271,8 @@ struct LTXVideoVAE : public VAE {
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
vae.get_param_tensors(tensors, prefix);
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
vae.get_param_tensors(tensors, weight_prefix);
}
struct TemporalTilePlan {
@@ -1396,7 +1396,7 @@ struct LTXVideoVAE : public VAE {
static_cast<int>(start),
chunk_overlap);
};
auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
expected_dim);
if (chunk.empty()) {
free_cache_ctx_and_buffer();
@@ -1452,7 +1452,7 @@ struct LTXVideoVAE : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
if (result.empty()) {
return {};
}
@@ -1465,7 +1465,7 @@ struct LTXVideoVAE : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_latent_statistics_graph(z, normalize);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
static_cast<size_t>(z.dim()));
}
@@ -1541,7 +1541,7 @@ struct LTXVideoVAE : public VAE {
}
std::map<std::string, ggml_tensor*> tensors;
vae->get_param_tensors(tensors, "first_stage_model");
vae->get_param_tensors(tensors);
if (!model_loader.load_tensors(tensors)) {
LOG_ERROR("load tensors from model loader failed");
src/model/vae/tae.hpp
+11 -11
View File
@@ -628,9 +628,9 @@ struct TinyImageAutoEncoder : public VAE {
const std::string prefix,
bool decoder_only = true,
SDVersion version = VERSION_SD1)
: decode_only(decoder_only),
taesd(decoder_only, version),
VAE(version, backend, params_backend) {
: VAE(version, backend, params_backend, "tae"),
decode_only(decoder_only),
taesd(decoder_only, version) {
scale_input = false;
taesd.init(params_ctx, tensor_storage_map, prefix);
}
@@ -639,8 +639,8 @@ struct TinyImageAutoEncoder : public VAE {
return "taesd";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
taesd.get_param_tensors(tensors, prefix);
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
taesd.get_param_tensors(tensors, weight_prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
@@ -676,7 +676,7 @@ struct TinyImageAutoEncoder : public VAE {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
}
};
@@ -691,8 +691,8 @@ struct TinyVideoAutoEncoder : public VAE {
const std::string prefix,
bool decoder_only = true,
SDVersion version = VERSION_WAN2)
: decode_only(decoder_only),
VAE(version, backend, params_backend) {
: VAE(version, backend, params_backend, "tae"),
decode_only(decoder_only) {
for (auto tensor_storage : tensor_storage_map) {
if (tensor_storage.first.find(prefix + ".3.conv.6.weight") != std::string::npos) {
is_wide = true;
@@ -708,8 +708,8 @@ struct TinyVideoAutoEncoder : public VAE {
return "taehv";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
taehv.get_param_tensors(tensors, prefix);
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
taehv.get_param_tensors(tensors, weight_prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
@@ -746,7 +746,7 @@ struct TinyVideoAutoEncoder : public VAE {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
}
};
src/model/vae/vae.hpp
+5 -4
View File
@@ -7,6 +7,7 @@
struct VAE : public GGMLRunner {
protected:
SDVersion version;
std::string weight_prefix;
bool scale_input = true;
virtual sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z,
@@ -62,8 +63,8 @@ protected:
}
public:
VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
: version(version), GGMLRunner(backend, params_backend) {}
VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend, const std::string& weight_prefix = "")
: version(version), weight_prefix(weight_prefix), GGMLRunner(backend, params_backend) {}
int get_scale_factor() {
int scale_factor = 8;
@@ -214,7 +215,7 @@ public:
virtual sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) = 0;
virtual sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) = 0;
virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
virtual void set_tiling_params(const sd_tiling_params_t& params) {
@@ -251,7 +252,7 @@ struct FakeVAE : public VAE {
return latents;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {}
std::string get_desc() override {
return "fake_vae";
src/model/vae/wan_vae.hpp
+8 -8
View File
@@ -1129,7 +1129,7 @@ namespace WAN {
const std::string prefix = "",
bool decode_only = false,
SDVersion version = VERSION_WAN2)
: decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V), VAE(version, backend, params_backend) {
: VAE(version, backend, params_backend, prefix), decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V) {
ae.init(params_ctx, tensor_storage_map, prefix);
}
@@ -1137,8 +1137,8 @@ namespace WAN {
return "wan_vae";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
ae.get_param_tensors(tensors, prefix);
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
ae.get_param_tensors(tensors, weight_prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
@@ -1255,7 +1255,7 @@ namespace WAN {
return build_graph(input, decode_graph);
}
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
input.empty() ? z.dim() : input.dim());
if (!result.empty() && z.dim() == 4) {
result.squeeze_(2);
@@ -1268,7 +1268,7 @@ namespace WAN {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph_partial(z, decode_graph, i);
};
auto out_opt = GGMLRunner::compute<float>(get_graph, n_threads, true);
auto out_opt = GGMLRunner::compute<float>(get_graph, n_threads, true, true, true);
if (!out_opt.has_value()) {
return {};
}
@@ -1281,7 +1281,7 @@ namespace WAN {
sd::Tensor<float> output = std::move(out);
for (i = 1; i < t; i++) {
auto chunk_opt = GGMLRunner::compute<float>(get_graph, n_threads, true);
auto chunk_opt = GGMLRunner::compute<float>(get_graph, n_threads, true, true, true);
if (!chunk_opt.has_value()) {
return {};
}
@@ -1327,7 +1327,7 @@ namespace WAN {
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_F16;
std::shared_ptr<WanVAERunner> vae = std::make_shared<WanVAERunner>(backend, backend, String2TensorStorage{}, "", false, VERSION_WAN2_2_TI2V);
std::shared_ptr<WanVAERunner> vae = std::make_shared<WanVAERunner>(backend, backend, String2TensorStorage{}, "first_stage_model", false, VERSION_WAN2_2_TI2V);
{
LOG_INFO("loading from '%s'", file_path.c_str());
@@ -1336,7 +1336,7 @@ namespace WAN {
return;
}
std::map<std::string, ggml_tensor*> tensors;
vae->get_param_tensors(tensors, "first_stage_model");
vae->get_param_tensors(tensors);
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "vae.")) {