mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 11:26:34 +00:00
292 lines
15 KiB
C++
292 lines
15 KiB
C++
#include "model_builders.h"
|
|
|
|
#include <cmath>
|
|
#include <cstring>
|
|
#include <utility>
|
|
|
|
#include "conditioning/conditioner.hpp"
|
|
#include "core/ggml_extend_backend.h"
|
|
#include "core/util.h"
|
|
#include "extensions/generation_extension.h"
|
|
#include "model/adapter/ip_adapter.hpp"
|
|
#include "model/audio/wav2vec2.hpp"
|
|
#include "model/diffusion/control.hpp"
|
|
#include "model/diffusion/model.hpp"
|
|
#include "model/vae/auto_encoder_kl.hpp"
|
|
#include "model/vae/hunyuan_vae.hpp"
|
|
#include "model/vae/ltx_audio_vae.hpp"
|
|
#include "model/vae/ltx_vae.hpp"
|
|
#include "model/vae/mage_vae.hpp"
|
|
#include "model/vae/minimax_h3_audio_vae.hpp"
|
|
#include "model/vae/minimax_h3_vae.hpp"
|
|
#include "model/vae/tae.hpp"
|
|
#include "model/vae/vae.hpp"
|
|
#include "model/vae/wan_vae.hpp"
|
|
#include "model_definition.h"
|
|
|
|
namespace sd::model_builders {
|
|
|
|
bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) {
|
|
if (backends.runtime_backend(module) == nullptr) {
|
|
LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));
|
|
return false;
|
|
}
|
|
if (backends.params_backend(module) == nullptr) {
|
|
LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback) {
|
|
switch (format) {
|
|
case SD_VAE_FORMAT_FLUX:
|
|
return VERSION_FLUX;
|
|
case SD_VAE_FORMAT_SD3:
|
|
return VERSION_SD3;
|
|
case SD_VAE_FORMAT_FLUX2:
|
|
return VERSION_FLUX2;
|
|
case SD_VAE_FORMAT_WAN:
|
|
return VERSION_WAN2;
|
|
case SD_VAE_FORMAT_AUTO:
|
|
default:
|
|
return fallback;
|
|
}
|
|
}
|
|
|
|
bool build_core_runners(const Context& ctx, CoreRunners& runners) try {
|
|
const auto* sd_ctx_params = &ctx.params;
|
|
const auto& tensor_storage_map = ctx.tensor_storage_map;
|
|
const auto version = ctx.version;
|
|
const auto& weight_manager = ctx.weight_manager;
|
|
CoreRunners result;
|
|
TokenizerConfig tokenizers(sd_ctx_params->tokenizer);
|
|
if (!ensure_backend_pair(ctx.backends, SDBackendModule::TE) ||
|
|
!ensure_backend_pair(ctx.backends, SDBackendModule::DIFFUSION)) {
|
|
return false;
|
|
}
|
|
|
|
const auto* definition = find_model_definition(version);
|
|
if (!definition || !definition->build || !definition->build(ctx, tokenizers, result)) {
|
|
LOG_ERROR("failed to build model runners for %s", sd_version_name(version));
|
|
return false;
|
|
}
|
|
|
|
if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && result.clip_vision == nullptr) {
|
|
if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
|
|
return false;
|
|
}
|
|
result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
|
|
tensor_storage_map,
|
|
weight_manager);
|
|
}
|
|
|
|
if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) {
|
|
result.ip_adapter = std::make_shared<IPAdapter::IPAdapterRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
|
|
tensor_storage_map,
|
|
"ip_adapter",
|
|
weight_manager);
|
|
}
|
|
if (result.conditioner) {
|
|
result.conditioner->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
if (result.diffusion) {
|
|
result.diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
if (result.high_noise_diffusion) {
|
|
result.high_noise_diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
if (result.clip_vision) {
|
|
result.clip_vision->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
if (result.ip_adapter) {
|
|
result.ip_adapter->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
tokenizers.validate_usage();
|
|
runners = std::move(result);
|
|
return true;
|
|
} catch (const std::exception& error) {
|
|
LOG_ERROR("failed to build model runners: %s", error.what());
|
|
return false;
|
|
}
|
|
|
|
bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) try {
|
|
const auto* sd_ctx_params = &ctx.params;
|
|
const auto& tensor_storage_map = ctx.tensor_storage_map;
|
|
const auto version = ctx.version;
|
|
const auto& weight_manager = ctx.weight_manager;
|
|
VAERunners result;
|
|
if (!ensure_backend_pair(ctx.backends, SDBackendModule::VAE)) {
|
|
return false;
|
|
}
|
|
|
|
auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
|
|
if (find_model_traits(version)->video_tae) {
|
|
return std::make_shared<TinyVideoAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"decoder",
|
|
decode_only,
|
|
version,
|
|
weight_manager);
|
|
|
|
} else {
|
|
auto model = std::make_shared<TinyImageAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"decoder.layers",
|
|
decode_only,
|
|
version,
|
|
weight_manager);
|
|
return model;
|
|
}
|
|
};
|
|
|
|
sd_vae_format_t vae_format = sd_ctx_params->vae_format;
|
|
if (vae_format < SD_VAE_FORMAT_AUTO || vae_format >= SD_VAE_FORMAT_COUNT) {
|
|
LOG_WARN("invalid VAE format override, using auto");
|
|
vae_format = SD_VAE_FORMAT_AUTO;
|
|
}
|
|
SDVersion vae_version = version;
|
|
if (find_model_traits(version)->vae_format_override && vae_format != SD_VAE_FORMAT_AUTO) {
|
|
vae_version = sd_vae_format_to_version(vae_format, vae_version);
|
|
}
|
|
|
|
auto create_vae = [&]() -> std::shared_ptr<VAE> {
|
|
if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::LTXAV) {
|
|
return std::make_shared<LTXVideoVAE>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"first_stage_model",
|
|
false,
|
|
version,
|
|
weight_manager);
|
|
} else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MINIMAX_H3) {
|
|
if (version == VERSION_P_H3_IMAGE) {
|
|
// The video VAE also loads successfully, but uses different temporal decoding and posterior sampling.
|
|
auto proj = tensor_storage_map.find("first_stage_model.decoder.proj_out.weight");
|
|
if (proj == tensor_storage_map.end() || proj->second.ne[1] != 768) {
|
|
throw std::runtime_error("P-H3-Image requires its single-frame VAE; pass --vae p-h3-image_vae_bf16.safetensors");
|
|
}
|
|
}
|
|
return std::make_shared<MiniMaxH3VAE::MiniMaxH3VideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"first_stage_model",
|
|
weight_manager);
|
|
} else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MAGE) {
|
|
return std::make_shared<MageVAE::MageVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"first_stage_model",
|
|
weight_manager);
|
|
} else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::HUNYUAN) {
|
|
return std::make_shared<Hunyuan::HunyuanVideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"first_stage_model",
|
|
false,
|
|
vae_version,
|
|
weight_manager);
|
|
} else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::WAN) {
|
|
return std::make_shared<WAN::WanVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"first_stage_model",
|
|
false,
|
|
vae_version,
|
|
weight_manager);
|
|
} else {
|
|
auto model = std::make_shared<AutoEncoderKL>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"first_stage_model",
|
|
false,
|
|
false,
|
|
vae_version,
|
|
weight_manager);
|
|
const auto* definition = find_model_definition(version);
|
|
if (definition && definition->configure_vae) {
|
|
definition->configure_vae(ctx, options, *model);
|
|
}
|
|
return model;
|
|
}
|
|
};
|
|
|
|
if (find_model_traits(version)->fake_vae) {
|
|
LOG_INFO("using FakeVAE");
|
|
result.vae = std::make_shared<FakeVAE>(version,
|
|
ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
weight_manager);
|
|
} else if (options.use_tae && !options.tae_preview_only) {
|
|
LOG_INFO("using TAE for encoding / decoding");
|
|
result.vae = create_tae(false);
|
|
} else {
|
|
LOG_INFO("using VAE for encoding / decoding");
|
|
result.vae = create_vae();
|
|
if (options.use_tae && options.tae_preview_only) {
|
|
LOG_INFO("using TAE for preview");
|
|
result.preview = create_tae(true);
|
|
}
|
|
}
|
|
|
|
if (options.use_audio_vae) {
|
|
if (find_model_traits(version)->vae_kind == ModelVAEKind::MINIMAX_H3) {
|
|
result.audio = std::make_shared<MiniMaxH3::AudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"",
|
|
weight_manager);
|
|
} else {
|
|
result.audio = std::make_shared<LTXV::LTXAudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
|
|
tensor_storage_map,
|
|
"",
|
|
weight_manager);
|
|
}
|
|
}
|
|
|
|
if (result.vae) {
|
|
result.vae->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
if (result.preview) {
|
|
result.preview->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
if (result.audio) {
|
|
result.audio->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
}
|
|
runners = std::move(result);
|
|
return true;
|
|
} catch (const std::exception& error) {
|
|
LOG_ERROR("failed to build VAE runners: %s", error.what());
|
|
return false;
|
|
}
|
|
|
|
bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner) {
|
|
const auto* sd_ctx_params = &ctx.params;
|
|
const auto& tensor_storage_map = ctx.tensor_storage_map;
|
|
const auto version = ctx.version;
|
|
const auto& weight_manager = ctx.weight_manager;
|
|
if (sd_version_is_z_image(version)) {
|
|
return true;
|
|
}
|
|
if (!ensure_backend_pair(ctx.backends, SDBackendModule::CONTROL_NET)) {
|
|
return false;
|
|
}
|
|
auto control_net = std::make_shared<ControlNet>(ctx.backends.runtime_backend(SDBackendModule::CONTROL_NET),
|
|
tensor_storage_map,
|
|
version,
|
|
"",
|
|
weight_manager);
|
|
control_net->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
|
|
runner = std::move(control_net);
|
|
return true;
|
|
}
|
|
|
|
bool build_extension_runners(const GenerationExtensionInitContext& ctx,
|
|
std::vector<std::shared_ptr<GenerationExtension>>& extensions) {
|
|
std::vector<std::shared_ptr<GenerationExtension>> result;
|
|
for (auto extension : {create_photomaker_extension(), create_pulid_extension()}) {
|
|
if (!extension->init(ctx)) {
|
|
return false;
|
|
}
|
|
if (extension->is_enabled()) {
|
|
result.push_back(std::move(extension));
|
|
}
|
|
}
|
|
extensions = std::move(result);
|
|
return true;
|
|
}
|
|
|
|
} // namespace sd::model_builders
|