#include <algorithm>

#include <cmath>

#include <cstdlib>

#include <set>

#include <unordered_set>

#include <vector>

#include "core/ggml_extend.hpp"

#include "core/ggml_graph_cut.h"

#include "core/rng.hpp"

#include "core/rng_mt19937.hpp"

#include "core/rng_philox.hpp"

#include "core/util.h"

#include "model_loader.h"

#include "model_manager.h"

#include "stable-diffusion.h"

#include "conditioning/conditioner.hpp"

#include "extensions/generation_extension.h"

#include "model/adapter/lora.hpp"

#include "model/diffusion/anima.hpp"

#include "model/diffusion/control.hpp"

#include "model/diffusion/ernie_image.hpp"

#include "model/diffusion/flux.hpp"

#include "model/diffusion/hidream_o1.hpp"

#include "model/diffusion/ideogram4.hpp"

#include "model/diffusion/lens.hpp"

#include "model/diffusion/ltxv.hpp"

#include "model/diffusion/mmdit.hpp"

#include "model/diffusion/model.hpp"

#include "model/diffusion/pid.hpp"

#include "model/diffusion/qwen_image.hpp"

#include "model/diffusion/unet.hpp"

#include "model/diffusion/wan.hpp"

#include "model/diffusion/z_image.hpp"

#include "model/upscaler/esrgan.hpp"

#include "model/upscaler/ltx_latent_upscaler.hpp"

#include "model/vae/auto_encoder_kl.hpp"

#include "model/vae/ltx_audio_vae.hpp"

#include "model/vae/ltx_vae.hpp"

#include "model/vae/tae.hpp"

#include "model/vae/vae.hpp"

#include "model/vae/wan_vae.hpp"

#include "runtime/denoiser.hpp"

#include "runtime/guidance.h"

#include "runtime/sample-cache.h"

#include "upscaler.h"

#include "name_conversion.h"

#include "runtime/latent-preview.h"

const char* sd_vae_format_name(enum sd_vae_format_t format);

static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback);

#include <atomic>

const char* model_version_to_str[] = {

"SD 1.x",

"SD 1.x Inpaint",

"Instruct-Pix2Pix",

"SD 1.x Tiny UNet",

"SD 2.x",

"SD 2.x Inpaint",

"SD 2.x Tiny UNet",

"SDXS (512-DS)",

"SDXS (09)",

"SDXL",

"SDXL Inpaint",

"SDXL Instruct-Pix2Pix",

"SDXL (Vega)",

"SDXL (SSD1B)",

"SVD",

"SD3.x",

"Flux",

"Flux Fill",

"Flux Control",

"Flex.2",

"Chroma Radiance",

"Wan 2.x",

"Wan 2.2 I2V",

"Wan 2.2 TI2V",

"Qwen Image",

"Anima",

"Flux.2",

"Flux.2 klein",

"LTXAV",

"HiDream O1",

"Z-Image",

"Ovis Image",

"Ernie Image",

"Lens",

"Longcat-Image",

"PiD",

"Ideogram 4",

"ESRGAN",

};

const char* sampling_methods_str[] = {

"Euler",

"Euler A",

"Heun",

"DPM2",

"DPM++ (2s)",

"DPM++ (2M)",

"modified DPM++ (2M)",

"iPNDM",

"iPNDM_v",

"LCM",

"DDIM \"trailing\"",

"TCD",

"Res Multistep",

"Res 2s",

"ER-SDE",

"Euler CFG++",

"Euler A CFG++",

"Euler GE",

};

/*================================================== Helper Functions ================================================*/

static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {

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);

}

static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) {

return sd_version_is_inpaint_or_unet_edit(version) ||

(has_ref_images && sd_version_supports_ref_latent_img_cfg(version));

}

void calculate_alphas_cumprod(float* alphas_cumprod,

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;

}

}

static float get_cache_reuse_threshold(const sd_cache_params_t& params) {

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);

}

/*=============================================== StableDiffusionGGML ================================================*/

static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,

"sd_cancel_mode_t must be lock-free");

class StableDiffusionGGML {

public:

SDBackendManager backend_manager;

SDVersion version;

bool external_vae_is_invalid = 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>> runtime_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 enable_mmap = false;

sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;

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::shared_ptr<ModelManager> model_manager;

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

std::vector<float> file_alphas_cumprod;

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;

}

std::atomic<sd_cancel_mode_t> cancellation_flag = SD_CANCEL_RESET;

void set_cancel_flag(enum sd_cancel_mode_t flag) {

cancellation_flag.store(flag, std::memory_order_release);

}

void reset_cancel_flag() {

set_cancel_flag(SD_CANCEL_RESET);

}

enum sd_cancel_mode_t get_cancel_flag() {

return cancellation_flag.load(std::memory_order_acquire);

}

size_t max_graph_vram_bytes_for_module(SDBackendModule module) {

return max_vram_assignment.bytes_for_backend(backend_for(module));

}

bool ensure_backend_pair(SDBackendModule module) {

if (backend_for(module) == nullptr) {

return false;

}

return params_backend_for(module) != nullptr;

}

template <typename T>

bool register_runner_params(const std::string& desc,

const std::shared_ptr<T>& model,

SDBackendModule module,

size_t* params_mem_size = nullptr) {

if (model == nullptr) {

return true;

}

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

model->get_param_tensors(group_tensors);

if (model_manager == nullptr) {

return true;

}

return model_manager->register_param_tensors(desc,

std::move(group_tensors),

backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,

backend_for(module),

params_backend_for(module),

params_mem_size);

}

bool init_backend() {

std::string error;

if (!backend_manager.init(backend_spec.c_str(),

params_backend_spec.c_str(),

&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>();

}

}

void refresh_compvis_denoiser_sigmas() {

auto comp_vis_denoiser = std::dynamic_pointer_cast<CompVisDenoiser>(denoiser);

if (!comp_vis_denoiser) {

return;

}

std::vector<float> alphas_cumprod(TIMESTEPS);

if (file_alphas_cumprod.size() == TIMESTEPS) {

alphas_cumprod = file_alphas_cumprod;

} else {

calculate_alphas_cumprod(alphas_cumprod.data());

}

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

comp_vis_denoiser->sigmas[i] = std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);

comp_vis_denoiser->log_sigmas[i] = std::log(comp_vis_denoiser->sigmas[i]);

}

}

void load_alphas_cumprod(ModelLoader& model_loader) {

file_alphas_cumprod.clear();

std::vector<float> loaded_alphas;

if (!model_loader.load_float_tensor("alphas_cumprod", loaded_alphas, n_threads, enable_mmap)) {

return;

}

if (loaded_alphas.size() != TIMESTEPS) {

LOG_WARN("ignore alphas_cumprod from model file: expected %d values, got %zu",

TIMESTEPS,

loaded_alphas.size());

return;

}

for (float alpha : loaded_alphas) {

if (!std::isfinite(alpha) || alpha <= 0.0f || alpha > 1.0f) {

LOG_WARN("ignore invalid alphas_cumprod from model file");

return;

}

}

file_alphas_cumprod = std::move(loaded_alphas);

LOG_DEBUG("loaded alphas_cumprod from model file");

}

bool init(const sd_ctx_params_t* sd_ctx_params) {

n_threads = sd_ctx_params->n_threads;

enable_mmap = sd_ctx_params->enable_mmap;

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);

max_vram_assignment.reset(0.f);

{

std::string error;

if (!max_vram_assignment.parse(SAFE_STR(sd_ctx_params->max_vram), &error)) {

LOG_ERROR("%s", error.c_str());

return false;

}

}

std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers);

add_rpc_devices(rpc_servers_spec);

bool use_tae = false;

bool use_audio_vae = false;

bool use_control_net = 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()) {

return false;

}

{

std::string error;

if (!max_vram_assignment.canonicalize_backend_keys(&error)) {

LOG_ERROR("%s", error.c_str());

return false;

}

}

if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) {

LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");

stream_layers = false;

}

model_manager = std::make_shared<ModelManager>();

model_manager->set_n_threads(n_threads);

model_manager->set_enable_mmap(enable_mmap);

ModelLoader& model_loader = model_manager->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->pulid_weights_path)) > 0) {

LOG_INFO("loading PuLID weights from '%s'", sd_ctx_params->pulid_weights_path);

if (!model_loader.init_from_file(sd_ctx_params->pulid_weights_path,

"model.diffusion_model.")) {

LOG_WARN("loading PuLID weights from '%s' failed", sd_ctx_params->pulid_weights_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;

}

}

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

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

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

return false;

} else {

use_control_net = 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);

}

model_loader.process_model_files(enable_mmap, true);

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 params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION);

const bool streaming_constrained = stream_layers || params_offloaded;

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;

}

if (enable_mmap && apply_lora_immediately) {

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

}

load_alphas_cumprod(model_loader);

size_t text_encoder_params_mem_size = 0;

size_t unet_params_mem_size = 0;

size_t vae_params_mem_size = 0;

size_t control_net_params_mem_size = 0;

size_t extension_params_mem_size = 0;

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");

}

{

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),

tensor_storage_map,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

model_manager);

} else if (sd_version_is_pid(version)) {

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

tensor_storage_map,

version,

"",

false,

model_manager);

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

tensor_storage_map,

"model.diffusion_model.net",

model_manager);

} else if (sd_version_is_ideogram4(version)) {

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

tensor_storage_map,

version,

"",

false,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

model_manager);

} 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) {

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

tensor_storage_map,

sd_ctx_params->chroma_use_t5_mask,

sd_ctx_params->chroma_t5_mask_pad,

false,

model_manager);

} else if (version == VERSION_OVIS_IMAGE) {

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

tensor_storage_map,

version,

"",

false,

model_manager);

} else {

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

tensor_storage_map,

model_manager);

}

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

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->chroma_use_dit_mask,

model_manager);

} else if (sd_version_is_flux2(version)) {

bool is_chroma = false;

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

tensor_storage_map,

version,

"",

false,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->chroma_use_dit_mask,

model_manager);

} else if (sd_version_is_ltxav(version)) {

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

tensor_storage_map,

"text_encoders.llm",

"text_embedding_projection",

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

model_manager);

} else if (sd_version_is_wan(version)) {

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

tensor_storage_map,

true,

0,

true,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

version,

model_manager);

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),

tensor_storage_map,

"model.high_noise_diffusion_model",

version,

model_manager);

}

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),

tensor_storage_map,

model_manager);

clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION));

if (!register_runner_params("CLIP vision",

clip_vision,

SDBackendModule::CLIP_VISION)) {

return false;

}

}

} else if (sd_version_is_qwen_image(version)) {

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

tensor_storage_map,

version,

"",

true,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->qwen_image_zero_cond_t,

model_manager);

} else if (sd_version_is_longcat(version)) {

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

tensor_storage_map,

version,

"",

true,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

version,

sd_ctx_params->chroma_use_dit_mask,

model_manager);

} else if (version == VERSION_HIDREAM_O1) {

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

tensor_storage_map,

model_manager);

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

tensor_storage_map,

"model",

model_manager);

} else if (sd_version_is_anima(version)) {

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

tensor_storage_map,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

model_manager);

} else if (sd_version_is_z_image(version)) {

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

tensor_storage_map,

version,

"",

false,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

version,

model_manager);

} else if (sd_version_is_ernie_image(version)) {

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

tensor_storage_map,

version,

"",

false,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

model_manager);

} else if (sd_version_is_lens(version)) {

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

tensor_storage_map,

version,

"",

false,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

model_manager);

} 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),

tensor_storage_map,

embbeding_map,

version,

model_manager);

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

tensor_storage_map,

"model.diffusion_model",

version,

model_manager);

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_for_module(SDBackendModule::TE));

if (!register_runner_params("Conditioner model",

cond_stage_model,

SDBackendModule::TE,

&text_encoder_params_mem_size)) {

return false;

}

diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));

diffusion_model->set_stream_layers_enabled(stream_layers);

if (!register_runner_params("Diffusion model",

diffusion_model,

SDBackendModule::DIFFUSION,

&unet_params_mem_size)) {

return false;

}

if (high_noise_diffusion_model) {

high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));

high_noise_diffusion_model->set_stream_layers_enabled(stream_layers);

if (!register_runner_params("High noise diffusion model",

high_noise_diffusion_model,

SDBackendModule::DIFFUSION,

&unet_params_mem_size)) {

return false;

}

}

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

return false;

}

auto create_tae = [&](bool decode_only) -> 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),

tensor_storage_map,

"decoder",

decode_only,

version,

model_manager);

} else {

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

tensor_storage_map,

"decoder.layers",

decode_only,

version,

model_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 (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),

tensor_storage_map,

"first_stage_model",

false,

version,

model_manager);

} 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),

tensor_storage_map,

"first_stage_model",

false,

version,

model_manager);

} else {

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

tensor_storage_map,

"first_stage_model",

false,

false,

vae_version,

model_manager);

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;

}

};

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),

model_manager);

if (!register_runner_params("VAE",

first_stage_model,

SDBackendModule::VAE,

&vae_params_mem_size)) {

return false;

}

} else if (use_tae && !tae_preview_only) {

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

first_stage_model = create_tae(false);

first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE));

if (!register_runner_params("VAE",

first_stage_model,

SDBackendModule::VAE,

&vae_params_mem_size)) {

return false;

}

} 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_for_module(SDBackendModule::VAE));

if (!register_runner_params("VAE",

first_stage_model,

SDBackendModule::VAE,

&vae_params_mem_size)) {

return false;

}

if (use_tae && tae_preview_only) {

LOG_INFO("using TAE for preview");

preview_vae = create_tae(true);

preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE));

if (!register_runner_params("preview VAE",

preview_vae,

SDBackendModule::VAE,

&vae_params_mem_size)) {

return false;

}

}

}

if (use_audio_vae) {

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

tensor_storage_map,

"",

model_manager);

if (!register_runner_params("LTX audio VAE",

audio_vae_model,

SDBackendModule::VAE,

&vae_params_mem_size)) {

return false;

}

}

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 (use_control_net) {