DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 19:36:33 +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/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.")) {