float linear_start = 0.00085f,

float linear_end = 0.0120,

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;

}

ggml_backend_t backend = nullptr; // general backend

ggml_backend_t clip_backend = nullptr;

ggml_backend_t control_net_backend = nullptr;

ggml_backend_t vae_backend = nullptr;

SDVersion version;

bool vae_decode_only = false;

bool free_params_immediately = false;

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

int n_threads = -1;

float scale_factor = 0.18215f;

float shift_factor = 0.f;

std::shared_ptr<Conditioner> cond_stage_model;

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

std::shared_ptr<DiffusionModel> diffusion_model;

std::shared_ptr<DiffusionModel> high_noise_diffusion_model;

std::shared_ptr<VAE> first_stage_model;

std::shared_ptr<TinyAutoEncoder> tae_first_stage;

std::shared_ptr<ControlNet> control_net = nullptr;

std::shared_ptr<PhotoMakerIDEncoder> pmid_model;

std::shared_ptr<LoraModel> pmid_lora;

std::shared_ptr<PhotoMakerIDEmbed> pmid_id_embeds;

std::string taesd_path;

bool use_tiny_autoencoder = false;

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

bool offload_params_to_cpu = false;

bool stacked_id = false;

bool is_using_v_parameterization = false;

bool is_using_edm_v_parameterization = false;

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

std::string lora_model_dir;

// lora_name => multiplier

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

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

StableDiffusionGGML() = default;

~StableDiffusionGGML() {

if (clip_backend != backend) {

ggml_backend_free(clip_backend);

}

if (control_net_backend != backend) {

ggml_backend_free(control_net_backend);

}

if (vae_backend != backend) {

ggml_backend_free(vae_backend);

}

ggml_backend_free(backend);

}

void init_backend() {

#ifdef SD_USE_CUDA

LOG_DEBUG("Using CUDA backend");

backend = ggml_backend_cuda_init(0);

#endif

#ifdef SD_USE_METAL

LOG_DEBUG("Using Metal backend");

backend = ggml_backend_metal_init();

#endif

#ifdef SD_USE_VULKAN

LOG_DEBUG("Using Vulkan backend");

for (int device = 0; device < ggml_backend_vk_get_device_count(); ++device) {

backend = ggml_backend_vk_init(device);

}

if (!backend) {

LOG_WARN("Failed to initialize Vulkan backend");

}

#endif

#ifdef SD_USE_OPENCL

LOG_DEBUG("Using OpenCL backend");

// ggml_log_set(ggml_log_callback_default, nullptr); // Optional ggml logs

backend = ggml_backend_opencl_init();

if (!backend) {

LOG_WARN("Failed to initialize OpenCL backend");

}

#endif

#ifdef SD_USE_SYCL

LOG_DEBUG("Using SYCL backend");

backend = ggml_backend_sycl_init(0);

#endif

if (!backend) {

LOG_DEBUG("Using CPU backend");

backend = ggml_backend_cpu_init();

}

}

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;

lora_model_dir = SAFE_STR(sd_ctx_params->lora_model_dir);

taesd_path = SAFE_STR(sd_ctx_params->taesd_path);

use_tiny_autoencoder = taesd_path.size() > 0;

offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu;

if (sd_ctx_params->rng_type == STD_DEFAULT_RNG) {

rng = std::make_shared<STDDefaultRNG>();

} else if (sd_ctx_params->rng_type == CUDA_RNG) {

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

}

ggml_log_set(ggml_log_callback_default, nullptr);

init_backend();

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

}

}

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->qwen2vl_path)) > 0) {

LOG_INFO("loading qwen2vl from '%s'", sd_ctx_params->qwen2vl_path);

if (!model_loader.init_from_file(sd_ctx_params->qwen2vl_path, "text_encoders.qwen2vl.")) {

LOG_WARN("loading qwen2vl from '%s' failed", sd_ctx_params->qwen2vl_path);

}

}

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

LOG_INFO("loading qwen2vl vision from '%s'", sd_ctx_params->qwen2vl_vision_path);

if (!model_loader.init_from_file(sd_ctx_params->qwen2vl_vision_path, "text_encoders.qwen2vl.visual.")) {

LOG_WARN("loading qwen2vl vision from '%s' failed", sd_ctx_params->qwen2vl_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);

}

}

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

for (auto& [name, tensor_storage] : tensor_storage_map) {

if (contains(name, "qwen2vl") &&

ends_with(name, "weight") &&

(tensor_storage.type == GGML_TYPE_F32 || tensor_storage.type == GGML_TYPE_BF16)) {

tensor_storage.expected_type = GGML_TYPE_F16;

}

}

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;

if (wtype != GGML_TYPE_COUNT) {

model_loader.set_wtype_override(wtype);

}

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_version_is_sdxl(version)) {

scale_factor = 0.13025f;

} else if (sd_version_is_sd3(version)) {

scale_factor = 1.5305f;

shift_factor = 0.0609f;

} else if (sd_version_is_flux(version)) {

scale_factor = 0.3611f;

shift_factor = 0.1159f;

} else if (sd_version_is_wan(version) || sd_version_is_qwen_image(version)) {

scale_factor = 1.0f;

}

if (sd_version_is_control(version)) {

// Might need vae encode for control cond

vae_decode_only = false;

}

bool clip_on_cpu = sd_ctx_params->keep_clip_on_cpu;

{

clip_backend = backend;

if (clip_on_cpu && !ggml_backend_is_cpu(backend)) {

LOG_INFO("CLIP: Using CPU backend");

clip_backend = ggml_backend_cpu_init();

}

if (sd_version_is_sd3(version)) {

cond_stage_model = std::make_shared<SD3CLIPEmbedder>(clip_backend,

offload_params_to_cpu,

tensor_storage_map);

diffusion_model = std::make_shared<MMDiTModel>(backend,

offload_params_to_cpu,

tensor_storage_map);

} 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->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>(clip_backend,

offload_params_to_cpu,

tensor_storage_map,

sd_ctx_params->chroma_use_t5_mask,

sd_ctx_params->chroma_t5_mask_pad);

} else {

cond_stage_model = std::make_shared<FluxCLIPEmbedder>(clip_backend,

offload_params_to_cpu,

tensor_storage_map);

}

diffusion_model = std::make_shared<FluxModel>(backend,

offload_params_to_cpu,

tensor_storage_map,

version,

sd_ctx_params->chroma_use_dit_mask);

} else if (sd_version_is_wan(version)) {

cond_stage_model = std::make_shared<T5CLIPEmbedder>(clip_backend,

offload_params_to_cpu,

tensor_storage_map,

true,

1,

true);

diffusion_model = std::make_shared<WanModel>(backend,

offload_params_to_cpu,

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<WanModel>(backend,

offload_params_to_cpu,

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

clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(backend,

offload_params_to_cpu,

tensor_storage_map);

clip_vision->alloc_params_buffer();

clip_vision->get_param_tensors(tensors);

}

} 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<Qwen2_5_VLCLIPEmbedder>(clip_backend,

offload_params_to_cpu,

tensor_storage_map,

"",

enable_vision);

diffusion_model = std::make_shared<QwenImageModel>(backend,

offload_params_to_cpu,

tensor_storage_map,

"model.diffusion_model",

version);

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

if (strstr(SAFE_STR(sd_ctx_params->photo_maker_path), "v2")) {

cond_stage_model = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(clip_backend,

offload_params_to_cpu,

tensor_storage_map,

SAFE_STR(sd_ctx_params->embedding_dir),

version,

PM_VERSION_2);

} else {

cond_stage_model = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(clip_backend,

offload_params_to_cpu,

tensor_storage_map,

SAFE_STR(sd_ctx_params->embedding_dir),

version);

}

diffusion_model = std::make_shared<UNetModel>(backend,

offload_params_to_cpu,

tensor_storage_map,

version);

if (sd_ctx_params->diffusion_conv_direct) {

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

std::dynamic_pointer_cast<UNetModel>(diffusion_model)->unet.set_conv2d_direct_enabled(true);

}

}

if (sd_ctx_params->diffusion_flash_attn) {

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

diffusion_model->set_flash_attn_enabled(true);

}

cond_stage_model->alloc_params_buffer();

cond_stage_model->get_param_tensors(tensors);

diffusion_model->alloc_params_buffer();

diffusion_model->get_param_tensors(tensors);

if (sd_version_is_unet_edit(version)) {

vae_decode_only = false;

}

if (high_noise_diffusion_model) {

high_noise_diffusion_model->alloc_params_buffer();

high_noise_diffusion_model->get_param_tensors(tensors);

}

if (sd_ctx_params->keep_vae_on_cpu && !ggml_backend_is_cpu(backend)) {

LOG_INFO("VAE Autoencoder: Using CPU backend");

vae_backend = ggml_backend_cpu_init();

} else {

vae_backend = backend;

}

if (sd_version_is_wan(version) || sd_version_is_qwen_image(version)) {

first_stage_model = std::make_shared<WAN::WanVAERunner>(vae_backend,

offload_params_to_cpu,

tensor_storage_map,

"first_stage_model",

vae_decode_only,

version);

first_stage_model->alloc_params_buffer();

first_stage_model->get_param_tensors(tensors, "first_stage_model");

} else if (version == VERSION_CHROMA_RADIANCE) {

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

offload_params_to_cpu);

} else if (!use_tiny_autoencoder) {

first_stage_model = std::make_shared<AutoEncoderKL>(vae_backend,

offload_params_to_cpu,

tensor_storage_map,

"first_stage_model",

vae_decode_only,

false,

version);

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 (version == VERSION_SDXL &&

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

float vae_conv_2d_scale = 1.f / 32.f;

LOG_WARN(

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

"using Conv2D scale %.3f",

vae_conv_2d_scale);

first_stage_model->set_conv2d_scale(vae_conv_2d_scale);

}

first_stage_model->alloc_params_buffer();

first_stage_model->get_param_tensors(tensors, "first_stage_model");

} else {

tae_first_stage = std::make_shared<TinyAutoEncoder>(vae_backend,

offload_params_to_cpu,

tensor_storage_map,

"decoder.layers",

vae_decode_only,

version);

if (sd_ctx_params->vae_conv_direct) {

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

tae_first_stage->set_conv2d_direct_enabled(true);

}

}

// first_stage_model->get_param_tensors(tensors, "first_stage_model.");

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

ggml_backend_t controlnet_backend = nullptr;

if (sd_ctx_params->keep_control_net_on_cpu && !ggml_backend_is_cpu(backend)) {

LOG_DEBUG("ControlNet: Using CPU backend");

controlnet_backend = ggml_backend_cpu_init();

} else {

controlnet_backend = backend;

}

control_net = std::make_shared<ControlNet>(controlnet_backend,

offload_params_to_cpu,

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

}

}

if (strstr(SAFE_STR(sd_ctx_params->photo_maker_path), "v2")) {

pmid_model = std::make_shared<PhotoMakerIDEncoder>(backend,

offload_params_to_cpu,

tensor_storage_map,

"pmid",

version,

PM_VERSION_2);

LOG_INFO("using PhotoMaker Version 2");

} else {

pmid_model = std::make_shared<PhotoMakerIDEncoder>(backend,

offload_params_to_cpu,

tensor_storage_map,

"pmid",

version);

}

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

pmid_lora = std::make_shared<LoraModel>(backend, sd_ctx_params->photo_maker_path, "");

if (!pmid_lora->load_from_file(true, n_threads)) {

LOG_WARN("load photomaker lora tensors from %s failed", sd_ctx_params->photo_maker_path);

return false;

}

LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", sd_ctx_params->photo_maker_path);

if (!model_loader.init_from_file(sd_ctx_params->photo_maker_path, "pmid.")) {

LOG_WARN("loading stacked ID embedding from '%s' failed", sd_ctx_params->photo_maker_path);

} else {

stacked_id = true;

}

}

if (stacked_id) {

if (!pmid_model->alloc_params_buffer()) {

LOG_ERROR(" pmid model params buffer allocation failed");

return false;

}

pmid_model->get_param_tensors(tensors, "pmid");

}

}

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

struct 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_tiny_autoencoder) {

ignore_tensors.insert("first_stage_model.");

}

if (stacked_id) {

ignore_tensors.insert("lora.");

}

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("text_encoders.qwen2vl.visual.");

}

if (version == VERSION_SVD) {

ignore_tensors.insert("conditioner.embedders.3");

}

bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads);

if (!success) {

LOG_ERROR("load tensors from model loader failed");

ggml_free(ctx);

return false;

}

// LOG_DEBUG("model size = %.2fMB", total_size / 1024.0 / 1024.0);

{

size_t clip_params_mem_size = cond_stage_model->get_params_buffer_size();

size_t unet_params_mem_size = diffusion_model->get_params_buffer_size();

if (high_noise_diffusion_model) {

unet_params_mem_size += high_noise_diffusion_model->get_params_buffer_size();

}

size_t vae_params_mem_size = 0;

if (!use_tiny_autoencoder) {

vae_params_mem_size = first_stage_model->get_params_buffer_size();

} else {

if (!tae_first_stage->load_from_file(taesd_path, n_threads)) {

return false;

}

vae_params_mem_size = tae_first_stage->get_params_buffer_size();

}

size_t control_net_params_mem_size = 0;

if (control_net) {

if (!control_net->load_from_file(SAFE_STR(sd_ctx_params->control_net_path), n_threads)) {

return false;

}

control_net_params_mem_size = control_net->get_params_buffer_size();

}

size_t pmid_params_mem_size = 0;

if (stacked_id) {

pmid_params_mem_size = pmid_model->get_params_buffer_size();

}

size_t total_params_ram_size = 0;

size_t total_params_vram_size = 0;

if (ggml_backend_is_cpu(clip_backend)) {

total_params_ram_size += clip_params_mem_size + pmid_params_mem_size;

} else {

total_params_vram_size += clip_params_mem_size + pmid_params_mem_size;

}

if (ggml_backend_is_cpu(backend)) {

total_params_ram_size += unet_params_mem_size;

} else {

total_params_vram_size += unet_params_mem_size;

}

if (ggml_backend_is_cpu(vae_backend)) {

total_params_ram_size += vae_params_mem_size;

} else {

total_params_vram_size += vae_params_mem_size;

}

if (ggml_backend_is_cpu(control_net_backend)) {

total_params_ram_size += control_net_params_mem_size;

} else {

total_params_vram_size += control_net_params_mem_size;

}

size_t total_params_size = total_params_ram_size + total_params_vram_size;

LOG_INFO(

"total params memory size = %.2fMB (VRAM %.2fMB, RAM %.2fMB): "

"text_encoders %.2fMB(%s), diffusion_model %.2fMB(%s), vae %.2fMB(%s), controlnet %.2fMB(%s), pmid %.2fMB(%s)",

total_params_size / 1024.0 / 1024.0,

total_params_vram_size / 1024.0 / 1024.0,

total_params_ram_size / 1024.0 / 1024.0,

clip_params_mem_size / 1024.0 / 1024.0,

ggml_backend_is_cpu(clip_backend) ? "RAM" : "VRAM",

unet_params_mem_size / 1024.0 / 1024.0,

ggml_backend_is_cpu(backend) ? "RAM" : "VRAM",

vae_params_mem_size / 1024.0 / 1024.0,

ggml_backend_is_cpu(vae_backend) ? "RAM" : "VRAM",

control_net_params_mem_size / 1024.0 / 1024.0,

ggml_backend_is_cpu(control_net_backend) ? "RAM" : "VRAM",

pmid_params_mem_size / 1024.0 / 1024.0,

ggml_backend_is_cpu(clip_backend) ? "RAM" : "VRAM");

}

if (sd_ctx_params->prediction != DEFAULT_PRED) {

switch (sd_ctx_params->prediction) {

case EPS_PRED:

LOG_INFO("running in eps-prediction mode");

break;

case V_PRED:

LOG_INFO("running in v-prediction mode");

denoiser = std::make_shared<CompVisVDenoiser>();

break;

case EDM_V_PRED:

LOG_INFO("running in v-prediction EDM mode");

denoiser = std::make_shared<EDMVDenoiser>();

break;

case SD3_FLOW_PRED: {

LOG_INFO("running in FLOW mode");

float shift = sd_ctx_params->flow_shift;

if (shift == INFINITY) {

shift = 3.0;

}

denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);

break;

}

case FLUX_FLOW_PRED: {

LOG_INFO("running in Flux FLOW mode");

float shift = sd_ctx_params->flow_shift;

if (shift == INFINITY) {

shift = 3.0;

}

denoiser = std::make_shared<FluxFlowDenoiser>(shift);

break;

}

default: {

LOG_ERROR("Unknown parametrization %i", sd_ctx_params->prediction);

return false;

}

}

} else {

if (sd_version_is_sd2(version)) {

// check is_using_v_parameterization_for_sd2

if (is_using_v_parameterization_for_sd2(ctx, sd_version_is_inpaint(version))) {

is_using_v_parameterization = true;

}

} else if (sd_version_is_sdxl(version)) {

if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) {

// CosXL models

// TODO: get sigma_min and sigma_max values from file

is_using_edm_v_parameterization = true;

}

if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) {

is_using_v_parameterization = true;

}

} else if (version == VERSION_SVD) {

// TODO: V_PREDICTION_EDM

is_using_v_parameterization = true;

}

if (sd_version_is_sd3(version)) {

LOG_INFO("running in FLOW mode");

float shift = sd_ctx_params->flow_shift;

if (shift == INFINITY) {

shift = 3.0;

}

denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);

} else if (sd_version_is_flux(version)) {

LOG_INFO("running in Flux FLOW mode");

float shift = sd_ctx_params->flow_shift;

if (shift == INFINITY) {

shift = 1.0f; // TODO: validate

for (const auto& [name, tensor_storage] : tensor_storage_map) {

if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) {

shift = 1.15f;

}

}

}

denoiser = std::make_shared<FluxFlowDenoiser>(shift);

} else if (sd_version_is_wan(version)) {

LOG_INFO("running in FLOW mode");

float shift = sd_ctx_params->flow_shift;

if (shift == INFINITY) {

shift = 5.0;

}

denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);

} else if (sd_version_is_qwen_image(version)) {

LOG_INFO("running in FLOW mode");

float shift = sd_ctx_params->flow_shift;

if (shift == INFINITY) {

shift = 3.0;

}

denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);

} else if (is_using_v_parameterization) {

LOG_INFO("running in v-prediction mode");

denoiser = std::make_shared<CompVisVDenoiser>();

} else if (is_using_edm_v_parameterization) {

LOG_INFO("running in v-prediction EDM mode");

denoiser = std::make_shared<EDMVDenoiser>();

} else {

LOG_INFO("running in eps-prediction mode");

}

}

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

if (comp_vis_denoiser) {

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

comp_vis_denoiser->sigmas[i] = std::sqrt((1 - ((float*)alphas_cumprod_tensor->data)[i]) / ((float*)alphas_cumprod_tensor->data)[i]);

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

}

}

LOG_DEBUG("finished loaded file");

ggml_free(ctx);

return true;

}

void init_scheduler(scheduler_t scheduler) {

switch (scheduler) {

case DISCRETE:

LOG_INFO("running with discrete scheduler");

denoiser->scheduler = std::make_shared<DiscreteSchedule>();

break;

case KARRAS:

LOG_INFO("running with Karras scheduler");

denoiser->scheduler = std::make_shared<KarrasSchedule>();

break;

case EXPONENTIAL:

LOG_INFO("running exponential scheduler");

denoiser->scheduler = std::make_shared<ExponentialSchedule>();

break;

case AYS:

LOG_INFO("Running with Align-Your-Steps scheduler");

denoiser->scheduler = std::make_shared<AYSSchedule>();

denoiser->scheduler->version = version;

break;

case GITS:

LOG_INFO("Running with GITS scheduler");

denoiser->scheduler = std::make_shared<GITSSchedule>();

denoiser->scheduler->version = version;

break;

case SGM_UNIFORM:

LOG_INFO("Running with SGM Uniform schedule");

denoiser->scheduler = std::make_shared<SGMUniformSchedule>();

denoiser->scheduler->version = version;

break;

case SIMPLE:

LOG_INFO("Running with Simple schedule");

denoiser->scheduler = std::make_shared<SimpleSchedule>();

denoiser->scheduler->version = version;

break;

case SMOOTHSTEP:

LOG_INFO("Running with SmoothStep scheduler");

denoiser->scheduler = std::make_shared<SmoothStepSchedule>();

break;

case DEFAULT:

// Don't touch anything.

break;

default:

LOG_ERROR("Unknown scheduler %i", scheduler);

abort();

}

}

bool is_using_v_parameterization_for_sd2(ggml_context* work_ctx, bool is_inpaint = false) {

struct ggml_tensor* x_t = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 4, 1);

ggml_set_f32(x_t, 0.5);

struct ggml_tensor* c = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 1024, 2, 1, 1);

ggml_set_f32(c, 0.5);

struct ggml_tensor* timesteps = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 1);

ggml_set_f32(timesteps, 999);

struct ggml_tensor* concat = is_inpaint ? ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 5, 1) : nullptr;

if (concat != nullptr) {

ggml_set_f32(concat, 0);

}

int64_t t0 = ggml_time_ms();

struct ggml_tensor* out = ggml_dup_tensor(work_ctx, x_t);

DiffusionParams diffusion_params;

diffusion_params.x = x_t;

diffusion_params.timesteps = timesteps;

diffusion_params.context = c;

diffusion_params.c_concat = concat;

diffusion_model->compute(n_threads, diffusion_params, &out);

diffusion_model->free_compute_buffer();

double result = 0.f;

{

float* vec_x = (float*)x_t->data;

float* vec_out = (float*)out->data;

int64_t n = ggml_nelements(out);

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

result += ((double)vec_out[i] - (double)vec_x[i]);

}

result /= n;

}

int64_t t1 = ggml_time_ms();

LOG_DEBUG("check is_using_v_parameterization_for_sd2, taking %.2fs", (t1 - t0) * 1.0f / 1000);

return result < -1;

}

void apply_lora(std::string lora_name, float multiplier) {

int64_t t0 = ggml_time_ms();

std::string high_noise_tag = "|high_noise|";

bool is_high_noise = false;

if (starts_with(lora_name, high_noise_tag)) {

lora_name = lora_name.substr(high_noise_tag.size());

is_high_noise = true;

LOG_DEBUG("high noise lora: %s", lora_name.c_str());

}

std::string st_file_path = path_join(lora_model_dir, lora_name + ".safetensors");

std::string ckpt_file_path = path_join(lora_model_dir, lora_name + ".ckpt");

std::string file_path;

if (file_exists(st_file_path)) {

file_path = st_file_path;

} else if (file_exists(ckpt_file_path)) {

file_path = ckpt_file_path;

} else {

LOG_WARN("can not find %s or %s for lora %s", st_file_path.c_str(), ckpt_file_path.c_str(), lora_name.c_str());

return;

}

LoraModel lora(backend, file_path, is_high_noise ? "model.high_noise_" : "");

if (!lora.load_from_file(false, n_threads)) {

LOG_WARN("load lora tensors from %s failed", file_path.c_str());

return;

}

lora.multiplier = multiplier;

// TODO: send version?

lora.apply(tensors, version, n_threads);

lora.free_params_buffer();

int64_t t1 = ggml_time_ms();

LOG_INFO("lora '%s' applied, taking %.2fs", lora_name.c_str(), (t1 - t0) * 1.0f / 1000);

}

void apply_loras(const std::unordered_map<std::string, float>& lora_state) {

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

for (auto& kv : lora_state) {

const std::string& lora_name = kv.first;

float multiplier = kv.second;

lora_state_diff[lora_name] += multiplier;

}

for (auto& kv : curr_lora_state) {

const std::string& lora_name = kv.first;

float curr_multiplier = kv.second;

lora_state_diff[lora_name] -= curr_multiplier;

}

size_t rm = lora_state_diff.size() - lora_state.size();

if (rm != 0) {

LOG_INFO("attempting to apply %lu LoRAs (removing %lu applied LoRAs)", lora_state.size(), rm);

} else {

LOG_INFO("attempting to apply %lu LoRAs", lora_state.size());

}

for (auto& kv : lora_state_diff) {

apply_lora(kv.first, kv.second);

}

curr_lora_state = lora_state;

}

std::string apply_loras_from_prompt(const std::string& prompt) {

auto result_pair = extract_and_remove_lora(prompt);

std::unordered_map<std::string, float> lora_f2m = result_pair.first; // lora_name -> multiplier

for (auto& kv : lora_f2m) {

LOG_DEBUG("lora %s:%.2f", kv.first.c_str(), kv.second);

}

int64_t t0 = ggml_time_ms();

apply_loras(lora_f2m);

int64_t t1 = ggml_time_ms();

LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);

LOG_DEBUG("prompt after extract and remove lora: \"%s\"", result_pair.second.c_str());

return result_pair.second;

}

ggml_tensor* id_encoder(ggml_context* work_ctx,

ggml_tensor* init_img,

ggml_tensor* prompts_embeds,

ggml_tensor* id_embeds,

std::vector<bool>& class_tokens_mask) {

ggml_tensor* res = nullptr;

pmid_model->compute(n_threads, init_img, prompts_embeds, id_embeds, class_tokens_mask, &res, work_ctx);

return res;

}

ggml_tensor* get_clip_vision_output(ggml_context* work_ctx,

sd_image_t init_image,

bool return_pooled = true,

int clip_skip = -1,

bool zero_out_masked = false) {

ggml_tensor* output = nullptr;

if (zero_out_masked) {

if (return_pooled) {

output = ggml_new_tensor_1d(work_ctx,

GGML_TYPE_F32,

clip_vision->vision_model.projection_dim);

} else {

output = ggml_new_tensor_2d(work_ctx,

GGML_TYPE_F32,

clip_vision->vision_model.hidden_size,

257);

}

ggml_set_f32(output, 0.f);