return version == VERSION_FLUX ||

sd_version_is_flux2(version) ||

sd_version_is_qwen_image(version) ||

sd_version_is_longcat(version) ||

sd_version_is_z_image(version);

float linear_start = 0.00085f,

float linear_end = 0.0120f,

int timesteps = TIMESTEPS) {

float ls_sqrt = sqrtf(linear_start);

float le_sqrt = sqrtf(linear_end);

float amount = le_sqrt - ls_sqrt;

float product = 1.0f;

for (int i = 0; i < timesteps; i++) {

float beta = ls_sqrt + amount * ((float)i / (timesteps - 1));

product *= 1.0f - powf(beta, 2.0f);

alphas_cumprod[i] = product;

}

float reuse_threshold = params.reuse_threshold;

if (reuse_threshold == INFINITY) {

if (params.mode == SD_CACHE_EASYCACHE) {

reuse_threshold = 0.2f;

} else if (params.mode == SD_CACHE_UCACHE) {

reuse_threshold = 1.0f;

}

}

return std::max(0.0f, reuse_threshold);

std::vector<MmapTensorStore> mmap_tensor_store;

SDBackendManager backend_manager;

SDVersion version;

bool vae_decode_only = false;

bool external_vae_is_invalid = false;

bool free_params_immediately = false;

bool circular_x = false;

bool circular_y = false;

std::shared_ptr<RNG> rng = std::make_shared<PhiloxRNG>();

std::shared_ptr<RNG> sampler_rng = nullptr;

int n_threads = -1;

float default_flow_shift = INFINITY;

std::shared_ptr<Conditioner> cond_stage_model;

std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision; // for svd or wan2.1 i2v

std::shared_ptr<DiffusionModelRunner> diffusion_model;

std::shared_ptr<DiffusionModelRunner> high_noise_diffusion_model;

std::shared_ptr<VAE> first_stage_model;

std::shared_ptr<VAE> preview_vae;

std::shared_ptr<LTXV::LTXAudioVAERunner> audio_vae_model;

std::shared_ptr<ControlNet> control_net;

std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;

std::vector<std::shared_ptr<LoraModel>> cond_stage_lora_models;

std::vector<std::shared_ptr<LoraModel>> diffusion_lora_models;

std::vector<std::shared_ptr<LoraModel>> first_stage_lora_models;

bool apply_lora_immediately = false;

std::string taesd_path;

sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};

bool offload_params_to_cpu = false;

float max_vram = 0.f;

bool stream_layers = false;

std::string backend_spec;

std::string params_backend_spec;

bool is_using_v_parameterization = false;

bool is_using_edm_v_parameterization = false;

std::map<std::string, ggml_tensor*> tensors;

// lora_name => multiplier

std::unordered_map<std::string, float> curr_lora_state;

std::shared_ptr<Denoiser> denoiser = std::make_shared<CompVisDenoiser>();

StableDiffusionGGML() = default;

~StableDiffusionGGML() = default;

ggml_backend_t backend_for(SDBackendModule module) {

ggml_backend_t module_backend = backend_manager.runtime_backend(module);

if (module_backend == nullptr) {

LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));

}

return module_backend;

}

ggml_backend_t params_backend_for(SDBackendModule module) {

ggml_backend_t module_backend = backend_manager.params_backend(module);

if (module_backend == nullptr) {

LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));

}

return module_backend;

}

bool ensure_backend_pair(SDBackendModule module) {

if (backend_for(module) == nullptr) {

return false;

}

return params_backend_for(module) != nullptr;

}

bool init_backend(const sd_ctx_params_t* sd_ctx_params) {

std::string error;

if (!backend_manager.init(sd_ctx_params->backend,

sd_ctx_params->params_backend,

offload_params_to_cpu,

sd_ctx_params->keep_clip_on_cpu,

sd_ctx_params->keep_vae_on_cpu,

sd_ctx_params->keep_control_net_on_cpu,

&error)) {

LOG_ERROR("backend config failed: %s", error.c_str());

return false;

}

return ensure_backend_pair(SDBackendModule::DIFFUSION);

}

std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {

if (rng_type == STD_DEFAULT_RNG) {

return std::make_shared<STDDefaultRNG>();

} else if (rng_type == CPU_RNG) {

return std::make_shared<MT19937RNG>();

} else { // default: CUDA_RNG

return std::make_shared<PhiloxRNG>();

}

}

bool init(const sd_ctx_params_t* sd_ctx_params) {

n_threads = sd_ctx_params->n_threads;

vae_decode_only = sd_ctx_params->vae_decode_only;

free_params_immediately = sd_ctx_params->free_params_immediately;

offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu;

max_vram = sd_ctx_params->max_vram;

stream_layers = sd_ctx_params->stream_layers;

backend_spec = SAFE_STR(sd_ctx_params->backend);

params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);

if (stream_layers && max_vram == 0.f) {

LOG_WARN("--stream-layers has no effect without --max-vram set; ignoring");

stream_layers = false;

}

if (stream_layers && !offload_params_to_cpu && params_backend_spec.empty()) {

// Streaming needs CPU-resident params.

LOG_WARN("--stream-layers has no effect without --offload-to-cpu (or --params-backend); ignoring");

stream_layers = false;

}

bool use_tae = false;

bool use_audio_vae = false;

rng = get_rng(sd_ctx_params->rng_type);

if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) {

sampler_rng = get_rng(sd_ctx_params->sampler_rng_type);

} else {

sampler_rng = rng;

}

ggml_log_set(ggml_log_callback_default, nullptr);

if (!init_backend(sd_ctx_params)) {

return false;

}

max_vram = sd::ggml_graph_cut::resolve_max_vram_gib(max_vram, backend_for(SDBackendModule::DIFFUSION));

ModelLoader model_loader;

if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) {

LOG_INFO("loading model from '%s'", sd_ctx_params->model_path);

if (!model_loader.init_from_file(sd_ctx_params->model_path)) {

LOG_ERROR("init model loader from file failed: '%s'", sd_ctx_params->model_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->diffusion_model_path)) > 0) {

LOG_INFO("loading diffusion model from '%s'", sd_ctx_params->diffusion_model_path);

if (!model_loader.init_from_file(sd_ctx_params->diffusion_model_path, "model.diffusion_model.")) {

LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->diffusion_model_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) {

LOG_INFO("loading high noise diffusion model from '%s'", sd_ctx_params->high_noise_diffusion_model_path);

if (!model_loader.init_from_file(sd_ctx_params->high_noise_diffusion_model_path, "model.high_noise_diffusion_model.")) {

LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->high_noise_diffusion_model_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->uncond_diffusion_model_path)) > 0) {

LOG_INFO("loading unconditional diffusion model from '%s'", sd_ctx_params->uncond_diffusion_model_path);

if (!model_loader.init_from_file(sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) {

LOG_WARN("loading unconditional diffusion model from '%s' failed", sd_ctx_params->uncond_diffusion_model_path);

}

}

bool is_unet = sd_version_is_unet(model_loader.get_sd_version());

if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) {

LOG_INFO("loading clip_l from '%s'", sd_ctx_params->clip_l_path);

std::string prefix = is_unet ? "cond_stage_model.transformer." : "text_encoders.clip_l.transformer.";

if (!model_loader.init_from_file(sd_ctx_params->clip_l_path, prefix)) {

LOG_WARN("loading clip_l from '%s' failed", sd_ctx_params->clip_l_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->clip_g_path)) > 0) {

LOG_INFO("loading clip_g from '%s'", sd_ctx_params->clip_g_path);

std::string prefix = is_unet ? "cond_stage_model.1.transformer." : "text_encoders.clip_g.transformer.";

if (!model_loader.init_from_file(sd_ctx_params->clip_g_path, prefix)) {

LOG_WARN("loading clip_g from '%s' failed", sd_ctx_params->clip_g_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->clip_vision_path)) > 0) {

LOG_INFO("loading clip_vision from '%s'", sd_ctx_params->clip_vision_path);

std::string prefix = "cond_stage_model.transformer.";

if (!model_loader.init_from_file(sd_ctx_params->clip_vision_path, prefix)) {

LOG_WARN("loading clip_vision from '%s' failed", sd_ctx_params->clip_vision_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->t5xxl_path)) > 0) {

LOG_INFO("loading t5xxl from '%s'", sd_ctx_params->t5xxl_path);

if (!model_loader.init_from_file(sd_ctx_params->t5xxl_path, "text_encoders.t5xxl.transformer.")) {

LOG_WARN("loading t5xxl from '%s' failed", sd_ctx_params->t5xxl_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->llm_path)) > 0) {

LOG_INFO("loading llm from '%s'", sd_ctx_params->llm_path);

if (!model_loader.init_from_file(sd_ctx_params->llm_path, "text_encoders.llm.")) {

LOG_WARN("loading llm from '%s' failed", sd_ctx_params->llm_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->llm_vision_path)) > 0) {

LOG_INFO("loading llm vision from '%s'", sd_ctx_params->llm_vision_path);

if (!model_loader.init_from_file(sd_ctx_params->llm_vision_path, "text_encoders.llm.visual.")) {

LOG_WARN("loading llm vision from '%s' failed", sd_ctx_params->llm_vision_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->vae_path)) > 0) {

LOG_INFO("loading vae from '%s'", sd_ctx_params->vae_path);

if (!model_loader.init_from_file(sd_ctx_params->vae_path, "vae.")) {

LOG_WARN("loading vae from '%s' failed", sd_ctx_params->vae_path);

external_vae_is_invalid = true;

}

}

if (strlen(SAFE_STR(sd_ctx_params->taesd_path)) > 0) {

LOG_INFO("loading tae from '%s'", sd_ctx_params->taesd_path);

if (!model_loader.init_from_file(sd_ctx_params->taesd_path, "tae.")) {

LOG_WARN("loading tae from '%s' failed", sd_ctx_params->taesd_path);

} else {

use_tae = true;

}

}

if (strlen(SAFE_STR(sd_ctx_params->embeddings_connectors_path)) > 0) {

LOG_INFO("loading embeddings connectors from '%s'", sd_ctx_params->embeddings_connectors_path);

if (!model_loader.init_from_file(sd_ctx_params->embeddings_connectors_path)) {

LOG_WARN("loading embeddings connectors from '%s' failed", sd_ctx_params->embeddings_connectors_path);

}

}

if (strlen(SAFE_STR(sd_ctx_params->audio_vae_path)) > 0) {

LOG_INFO("loading LTX audio VAE from '%s'", sd_ctx_params->audio_vae_path);

if (!model_loader.init_from_file(sd_ctx_params->audio_vae_path)) {

LOG_WARN("loading LTX audio VAE weights from '%s' failed", sd_ctx_params->audio_vae_path);

} else {

use_audio_vae = true;

}

}

model_loader.convert_tensors_name();

version = model_loader.get_sd_version();

if (version == VERSION_COUNT) {

LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path));

return false;

}

auto& tensor_storage_map = model_loader.get_tensor_storage_map();

LOG_INFO("Version: %s ", model_version_to_str[version]);

ggml_type wtype = (int)sd_ctx_params->wtype < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT)

? (ggml_type)sd_ctx_params->wtype

: GGML_TYPE_COUNT;

std::string tensor_type_rules = SAFE_STR(sd_ctx_params->tensor_type_rules);

if (wtype != GGML_TYPE_COUNT || tensor_type_rules.size() > 0) {

model_loader.set_wtype_override(wtype, tensor_type_rules);

}

std::map<ggml_type, uint32_t> wtype_stat = model_loader.get_wtype_stat();

std::map<ggml_type, uint32_t> conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat();

std::map<ggml_type, uint32_t> diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat();

std::map<ggml_type, uint32_t> vae_wtype_stat = model_loader.get_vae_wtype_stat();

auto wtype_stat_to_str = [](const std::map<ggml_type, uint32_t>& m, int key_width = 8, int value_width = 5) -> std::string {

std::ostringstream oss;

bool first = true;

for (const auto& [type, count] : m) {

if (!first)

oss << "|";

first = false;

oss << std::right << std::setw(key_width) << ggml_type_name(type)

<< ": "

<< std::left << std::setw(value_width) << count;

}

return oss.str();

};

LOG_INFO("Weight type stat: %s", wtype_stat_to_str(wtype_stat).c_str());

LOG_INFO("Conditioner weight type stat: %s", wtype_stat_to_str(conditioner_wtype_stat).c_str());

LOG_INFO("Diffusion model weight type stat: %s", wtype_stat_to_str(diffusion_model_wtype_stat).c_str());

LOG_INFO("VAE weight type stat: %s", wtype_stat_to_str(vae_wtype_stat).c_str());

LOG_DEBUG("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor));

if (sd_ctx_params->lora_apply_mode == LORA_APPLY_AUTO) {

bool have_quantized_weight = false;

if (wtype != GGML_TYPE_COUNT && ggml_is_quantized(wtype)) {

have_quantized_weight = true;

} else {

for (const auto& [type, _] : wtype_stat) {

if (ggml_is_quantized(type)) {

have_quantized_weight = true;

break;

}

}

}

// Avoid full-model LoRA merge buffers on constrained setups.

const bool streaming_constrained = stream_layers ||

sd_ctx_params->offload_params_to_cpu;

if (have_quantized_weight || streaming_constrained) {

apply_lora_immediately = false;

} else {

apply_lora_immediately = true;

}

} else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) {

apply_lora_immediately = true;

} else {

apply_lora_immediately = false;

}

std::map<std::string, ggml_tensor*> mmap_able_tensors;

bool enable_mmap_tensors = false;

bool needs_writable_mmap = false;

if (sd_ctx_params->enable_mmap) {

if (apply_lora_immediately) {

needs_writable_mmap = true;

LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap");

}

enable_mmap_tensors = true;

}

// split definition to avoid msvc choking on the extra parameter handling

auto module_can_mmap = [&](SDBackendModule module) {

return enable_mmap_tensors &&

(backend_manager.runtime_backend_is_cpu(module) ||

backend_manager.params_backend_is_cpu(module) ||

backend_manager.runtime_backend_supports_host_buffer(module));

};

auto get_param_tensors_p = [&](auto&& model, bool do_mmap, const char* prefix) {

std::map<std::string, ggml_tensor*> temp;

model->get_param_tensors(temp, prefix);

for (const auto& [key, tensor] : temp) {

tensors[key] = tensor;

if (do_mmap) {

mmap_able_tensors[key] = tensor;

}

}

};

auto get_param_tensors = [&](auto&& model, bool do_mmap) {

std::map<std::string, ggml_tensor*> temp;

model->get_param_tensors(temp);

for (const auto& [key, tensor] : temp) {

tensors[key] = tensor;

if (do_mmap) {

mmap_able_tensors[key] = tensor;

}

}

};

if (sd_version_is_control(version)) {

// Might need vae encode for control cond

vae_decode_only = false;

}

bool tae_preview_only = sd_ctx_params->tae_preview_only;

if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {

tae_preview_only = false;

use_tae = true;

}

if (sd_ctx_params->circular_x || sd_ctx_params->circular_y) {

LOG_INFO("Using circular padding for convolutions");

}

const size_t max_graph_vram_bytes = sd::ggml_graph_cut::max_vram_gib_to_bytes(max_vram);

{

if (!ensure_backend_pair(SDBackendModule::TE) ||

!ensure_backend_pair(SDBackendModule::DIFFUSION)) {

return false;

}

if (sd_version_is_sd3(version)) {

cond_stage_model = std::make_shared<SD3CLIPEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map);

diffusion_model = std::make_shared<MMDiTRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model");

} else if (sd_version_is_pid(version)) {

vae_decode_only = false;

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version);

diffusion_model = std::make_shared<Pid::PiDRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model.net");

} else if (sd_version_is_ideogram4(version)) {

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version,

"",

false);

diffusion_model = std::make_shared<Ideogram4::Ideogram4Runner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model");

} else if (sd_version_is_flux(version)) {

bool is_chroma = false;

for (auto pair : tensor_storage_map) {

if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {

is_chroma = true;

break;

}

}

if (is_chroma) {

if ((sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) && sd_ctx_params->chroma_use_dit_mask) {

LOG_WARN(

"!!!It looks like you are using Chroma with flash attention. "

"This is currently unsupported. "

"If you find that the generated images are broken, "

"try either disabling flash attention or specifying "

"--chroma-disable-dit-mask as a workaround.");

}

cond_stage_model = std::make_shared<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

sd_ctx_params->chroma_use_t5_mask,

sd_ctx_params->chroma_t5_mask_pad);

} else if (version == VERSION_OVIS_IMAGE) {

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version,

"",

false);

} else {

cond_stage_model = std::make_shared<FluxCLIPEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map);

}

diffusion_model = std::make_shared<Flux::FluxRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->chroma_use_dit_mask);

} else if (sd_version_is_flux2(version)) {

bool is_chroma = false;

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version);

diffusion_model = std::make_shared<Flux::FluxRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->chroma_use_dit_mask);

} else if (sd_version_is_ltxav(version)) {

cond_stage_model = std::make_shared<LTXAVEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map);

diffusion_model = std::make_shared<LTXV::LTXAVRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model");

} else if (sd_version_is_wan(version)) {

cond_stage_model = std::make_shared<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

true,

0,

true);

diffusion_model = std::make_shared<WAN::WanRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model",

version);

if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) {

high_noise_diffusion_model = std::make_shared<WAN::WanRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.high_noise_diffusion_model",

version);

}

if (diffusion_model->get_desc() == "Wan2.1-I2V-14B" ||

diffusion_model->get_desc() == "Wan2.1-FLF2V-14B" ||

diffusion_model->get_desc() == "Wan2.1-I2V-1.3B") {

if (!ensure_backend_pair(SDBackendModule::CLIP_VISION)) {

return false;

}

clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(backend_for(SDBackendModule::CLIP_VISION),

params_backend_for(SDBackendModule::CLIP_VISION),

tensor_storage_map);

clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes);

get_param_tensors(clip_vision, module_can_mmap(SDBackendModule::CLIP_VISION));

}

} else if (sd_version_is_qwen_image(version)) {

bool enable_vision = false;

if (!vae_decode_only) {

enable_vision = true;

}

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version,

"",

enable_vision);

diffusion_model = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->qwen_image_zero_cond_t);

} else if (sd_version_is_longcat(version)) {

bool enable_vision = false;

if (!vae_decode_only) {

enable_vision = true;

}

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version,

"",

enable_vision);

diffusion_model = std::make_shared<Flux::FluxRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->chroma_use_dit_mask);

} else if (version == VERSION_HIDREAM_O1) {

cond_stage_model = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map);

diffusion_model = std::make_shared<HiDreamO1::HiDreamO1Runner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model");

} else if (sd_version_is_anima(version)) {

cond_stage_model = std::make_shared<AnimaConditioner>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map);

diffusion_model = std::make_shared<Anima::AnimaRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model");

} else if (sd_version_is_z_image(version)) {

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version);

diffusion_model = std::make_shared<ZImage::ZImageRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model",

version);

} else if (sd_version_is_ernie_image(version)) {

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version);

diffusion_model = std::make_shared<ErnieImage::ErnieImageRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model");

} else if (sd_version_is_lens(version)) {

cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

version);

diffusion_model = std::make_shared<Lens::LensRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model");

} else { // SD1.x SD2.x SDXL

std::map<std::string, std::string> embbeding_map;

for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {

embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path));

}

cond_stage_model = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(backend_for(SDBackendModule::TE),

params_backend_for(SDBackendModule::TE),

tensor_storage_map,

embbeding_map,

version);

diffusion_model = std::make_shared<UNetModelRunner>(backend_for(SDBackendModule::DIFFUSION),

params_backend_for(SDBackendModule::DIFFUSION),

tensor_storage_map,

"model.diffusion_model",

version);

if (sd_ctx_params->diffusion_conv_direct) {

LOG_INFO("Using Conv2d direct in the diffusion model");

diffusion_model->set_conv2d_direct_enabled(true);

}

}

cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);

get_param_tensors(cond_stage_model, module_can_mmap(SDBackendModule::TE));

diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);

diffusion_model->set_stream_layers_enabled(stream_layers);

get_param_tensors(diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION));

if (sd_version_is_unet_edit(version)) {

vae_decode_only = false;

}

if (high_noise_diffusion_model) {

high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);

high_noise_diffusion_model->set_stream_layers_enabled(stream_layers);

get_param_tensors(high_noise_diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION));

}

if (!ensure_backend_pair(SDBackendModule::VAE)) {

return false;

}

auto create_tae = [&]() -> std::shared_ptr<VAE> {

if (sd_version_is_wan(version) ||

sd_version_is_qwen_image(version) ||

sd_version_is_anima(version) ||

sd_version_is_ltxav(version)) {

return std::make_shared<TinyVideoAutoEncoder>(backend_for(SDBackendModule::VAE),

params_backend_for(SDBackendModule::VAE),

tensor_storage_map,

"decoder",

vae_decode_only,

version);

} else {

auto model = std::make_shared<TinyImageAutoEncoder>(backend_for(SDBackendModule::VAE),

params_backend_for(SDBackendModule::VAE),

tensor_storage_map,

"decoder.layers",

vae_decode_only,

version);

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 (sd_version_is_pid(version) && 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 (sd_version_is_ltxav(version)) {

return std::make_shared<LTXVideoVAE>(backend_for(SDBackendModule::VAE),

params_backend_for(SDBackendModule::VAE),

tensor_storage_map,

"first_stage_model",

vae_decode_only,

version);

} else if (sd_version_is_wan(version) ||

sd_version_is_qwen_image(version) ||

sd_version_is_anima(version)) {

return std::make_shared<WAN::WanVAERunner>(backend_for(SDBackendModule::VAE),

params_backend_for(SDBackendModule::VAE),

tensor_storage_map,

"first_stage_model",

vae_decode_only,

version);

} else {

auto model = std::make_shared<AutoEncoderKL>(backend_for(SDBackendModule::VAE),

params_backend_for(SDBackendModule::VAE),

tensor_storage_map,

"first_stage_model",

vae_decode_only,

false,

vae_version);

if (sd_version_is_sdxl(version) &&

(strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || external_vae_is_invalid)) {

float vae_conv_2d_scale = 1.f / 32.f;

LOG_WARN(

"No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, "

"using Conv2D scale %.3f",

vae_conv_2d_scale);

model->set_conv2d_scale(vae_conv_2d_scale);

}

return model;

}

};

bool vae_mmap = module_can_mmap(SDBackendModule::VAE);

if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {

LOG_INFO("using FakeVAE");

first_stage_model = std::make_shared<FakeVAE>(version,

backend_for(SDBackendModule::VAE),

params_backend_for(SDBackendModule::VAE));

} else if (use_tae && !tae_preview_only) {

LOG_INFO("using TAE for encoding / decoding");

first_stage_model = create_tae();

first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);

get_param_tensors_p(first_stage_model, vae_mmap, "tae");

} else {

LOG_INFO("using VAE for encoding / decoding");

first_stage_model = create_vae();

first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);

get_param_tensors_p(first_stage_model, vae_mmap, "first_stage_model");

if (use_tae && tae_preview_only) {

LOG_INFO("using TAE for preview");

preview_vae = create_tae();

preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes);

get_param_tensors_p(preview_vae, vae_mmap, "tae");

}

}

if (use_audio_vae) {

audio_vae_model = std::make_shared<LTXV::LTXAudioVAERunner>(backend_for(SDBackendModule::VAE),

params_backend_for(SDBackendModule::VAE),

tensor_storage_map);

get_param_tensors_p(audio_vae_model, vae_mmap, "");

}

if (sd_ctx_params->vae_conv_direct) {

LOG_INFO("Using Conv2d direct in the vae model");

first_stage_model->set_conv2d_direct_enabled(true);

if (preview_vae) {

preview_vae->set_conv2d_direct_enabled(true);

}

}

if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) {

if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) {

return false;

}

control_net = std::make_shared<ControlNet>(backend_for(SDBackendModule::CONTROL_NET),

params_backend_for(SDBackendModule::CONTROL_NET),

tensor_storage_map,

version);

if (sd_ctx_params->diffusion_conv_direct) {

LOG_INFO("Using Conv2d direct in the control net");

control_net->set_conv2d_direct_enabled(true);

}

}

{

generation_extensions.clear();

auto photomaker_extension = create_photomaker_extension();

GenerationExtensionInitContext extension_ctx{

sd_ctx_params,

version,

tensor_storage_map,

model_loader,

n_threads,

[this](SDBackendModule module) { return ensure_backend_pair(module); },

[this](SDBackendModule module) { return backend_for(module); },

[this](SDBackendModule module) { return params_backend_for(module); },

};

if (!photomaker_extension->init(extension_ctx)) {

return false;

}

if (photomaker_extension->is_enabled()) {

generation_extensions.push_back(photomaker_extension);

}

}

{

GenerationExtensionTensorContext extension_tensor_ctx{

tensors,

mmap_able_tensors,

module_can_mmap,

};

for (auto& extension : generation_extensions) {

extension->collect_param_tensors(extension_tensor_ctx);

}

}

if (sd_ctx_params->flash_attn) {

LOG_INFO("Using flash attention");

cond_stage_model->set_flash_attention_enabled(true);

if (clip_vision) {

clip_vision->set_flash_attention_enabled(true);

}

if (first_stage_model) {

first_stage_model->set_flash_attention_enabled(true);

}

if (preview_vae) {

preview_vae->set_flash_attention_enabled(true);

}

}

if (sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) {

LOG_INFO("Using flash attention in the diffusion model");

diffusion_model->set_flash_attention_enabled(true);

if (high_noise_diffusion_model) {

high_noise_diffusion_model->set_flash_attention_enabled(true);

}

}

diffusion_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y);

if (high_noise_diffusion_model) {

high_noise_diffusion_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y);

}

if (control_net) {

control_net->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y);

}

circular_x = sd_ctx_params->circular_x;

circular_y = sd_ctx_params->circular_y;

}

ggml_init_params params;

params.mem_size = static_cast<size_t>(10 * 1024) * 1024; // 10M

params.mem_buffer = nullptr;

params.no_alloc = false;

// LOG_DEBUG("mem_size %u ", params.mem_size);

ggml_context* ctx = ggml_init(params); // for alphas_cumprod and is_using_v_parameterization check

GGML_ASSERT(ctx != nullptr);

ggml_tensor* alphas_cumprod_tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, TIMESTEPS);

calculate_alphas_cumprod((float*)alphas_cumprod_tensor->data);

// load weights

LOG_DEBUG("loading weights");

std::set<std::string> ignore_tensors;

tensors["alphas_cumprod"] = alphas_cumprod_tensor;

if (use_tae && !tae_preview_only) {

ignore_tensors.insert("first_stage_model.");

}

for (auto& extension : generation_extensions) {

extension->add_ignore_tensors(ignore_tensors);

}

ignore_tensors.insert("model.diffusion_model.__x0__");

ignore_tensors.insert("model.diffusion_model.__32x32__");

ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__");

if (vae_decode_only) {

ignore_tensors.insert("first_stage_model.encoder");

ignore_tensors.insert("first_stage_model.conv1");

ignore_tensors.insert("first_stage_model.quant");

ignore_tensors.insert("tae.encoder");

ignore_tensors.insert("text_encoders.llm.visual.");