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;

}

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;

FrozenCLIPEmbedderWithCustomWords cond_stage_model;

UNetModel diffusion_model;

AutoEncoderKL first_stage_model;

bool use_tiny_autoencoder = false;

bool vae_tiling = 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::map<std::string, LoraModel> loras;

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

ggml_backend_t backend = NULL; // general backend

ggml_type model_data_type = GGML_TYPE_COUNT;

TinyAutoEncoder tae_first_stage;

std::string taesd_path;

ControlNet control_net;

StableDiffusionGGML() = default;

StableDiffusionGGML(int n_threads,

bool vae_decode_only,

bool free_params_immediately,

std::string lora_model_dir,

rng_type_t rng_type)

: n_threads(n_threads),

vae_decode_only(vae_decode_only),

free_params_immediately(free_params_immediately),

lora_model_dir(lora_model_dir) {

first_stage_model.decode_only = vae_decode_only;

tae_first_stage.decode_only = vae_decode_only;

if (rng_type == STD_DEFAULT_RNG) {

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

} else if (rng_type == CUDA_RNG) {

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

}

}

~StableDiffusionGGML() {

ggml_backend_free(backend);

}

bool load_from_file(const std::string& model_path,

const std::string& vae_path,

const std::string control_net_path,

const std::string embeddings_path,

const std::string& taesd_path,

bool vae_tiling_,

ggml_type wtype,

schedule_t schedule,

bool control_net_cpu) {

use_tiny_autoencoder = taesd_path.size() > 0;

#ifdef SD_USE_CUBLAS

LOG_DEBUG("Using CUDA backend");

backend = ggml_backend_cuda_init(0);

#endif

#ifdef SD_USE_METAL

LOG_DEBUG("Using Metal backend");

ggml_metal_log_set_callback(ggml_log_callback_default, nullptr);

backend = ggml_backend_metal_init();

#endif

if (!backend) {

LOG_DEBUG("Using CPU backend");

backend = ggml_backend_cpu_init();

}

#ifdef SD_USE_FLASH_ATTENTION

#if defined(SD_USE_CUBLAS) || defined(SD_USE_METAL)

LOG_WARN("Flash Attention not supported with GPU Backend");

#else

LOG_INFO("Flash Attention enabled");

#endif

#endif

LOG_INFO("loading model from '%s'", model_path.c_str());

ModelLoader model_loader;

vae_tiling = vae_tiling_;

if (!model_loader.init_from_file(model_path)) {

LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());

return false;

}

if (vae_path.size() > 0) {

LOG_INFO("loading vae from '%s'", vae_path.c_str());

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

LOG_WARN("loading vae from '%s' failed", vae_path.c_str());

}

}

version = model_loader.get_sd_version();

if (version == VERSION_COUNT) {

LOG_ERROR("get sd version from file failed: '%s'", model_path.c_str());

return false;

}

if (version == VERSION_XL) {

scale_factor = 0.13025f;

}

cond_stage_model = FrozenCLIPEmbedderWithCustomWords(version);

diffusion_model = UNetModel(version);

LOG_INFO("Stable Diffusion %s ", model_version_to_str[version]);

if (wtype == GGML_TYPE_COUNT) {

model_data_type = model_loader.get_sd_wtype();

} else {

model_data_type = wtype;

}

LOG_INFO("Stable Diffusion weight type: %s", ggml_type_name(model_data_type));

LOG_DEBUG("loading vocab");

std::string merges_utf8_str = model_loader.load_merges();

if (merges_utf8_str.size() == 0) {

LOG_ERROR("get merges failed: '%s'", model_path.c_str());

return false;

}

cond_stage_model.tokenizer.load_from_merges(merges_utf8_str);

// create the ggml context for network params

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

if (

!cond_stage_model.alloc_params_buffer(backend, model_data_type) ||

!diffusion_model.alloc_params_buffer(backend, model_data_type)) {

return false;

}

cond_stage_model.text_model.embd_dir = embeddings_path;

ggml_type vae_type = model_data_type;

if (version == VERSION_XL) {

vae_type = GGML_TYPE_F32; // avoid nan, not work...

}

if (!use_tiny_autoencoder && !first_stage_model.alloc_params_buffer(backend, vae_type)) {

return false;

}

LOG_DEBUG("preparing memory for the weights");

// prepare memory for the weights

{

// cond_stage_model(FrozenCLIPEmbedder)

cond_stage_model.init_params();

cond_stage_model.map_by_name(tensors, "cond_stage_model.");

// diffusion_model(UNetModel)

diffusion_model.init_params();

diffusion_model.map_by_name(tensors, "model.diffusion_model.");

if (!use_tiny_autoencoder) {

// firest_stage_model(AutoEncoderKL)

first_stage_model.init_params();

}

first_stage_model.map_by_name(tensors, "first_stage_model.");

}

struct ggml_init_params params;

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

params.mem_buffer = NULL;

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

if (!ctx) {

LOG_ERROR("ggml_init() failed");

return false;

}

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

int64_t t0 = ggml_time_ms();

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

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

tensors_need_to_load["alphas_cumprod"] = alphas_cumprod_tensor;

for (auto& pair : tensors) {

const std::string& name = pair.first;

if (use_tiny_autoencoder && starts_with(name, "first_stage_model.")) {

ignore_tensors.insert(name);

continue;

}

if (vae_decode_only && (starts_with(name, "first_stage_model.encoder") || starts_with(name, "first_stage_model.quant"))) {

ignore_tensors.insert(name);

continue;

}

tensors_need_to_load.insert(pair);

}

bool success = model_loader.load_tensors(tensors_need_to_load, backend, ignore_tensors);

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 total_params_size =

cond_stage_model.params_buffer_size +

diffusion_model.params_buffer_size +

first_stage_model.params_buffer_size;

LOG_INFO("total memory buffer size = %.2fMB (clip %.2fMB, unet %.2fMB, vae %.2fMB)",

total_params_size / 1024.0 / 1024.0,

cond_stage_model.params_buffer_size / 1024.0 / 1024.0,

diffusion_model.params_buffer_size / 1024.0 / 1024.0,

first_stage_model.params_buffer_size / 1024.0 / 1024.0);

int64_t t1 = ggml_time_ms();

LOG_INFO("loading model from '%s' completed, taking %.2fs", model_path.c_str(), (t1 - t0) * 1.0f / 1000);

// check is_using_v_parameterization_for_sd2

bool is_using_v_parameterization = false;

if (version == VERSION_2_x) {

if (is_using_v_parameterization_for_sd2(ctx)) {

is_using_v_parameterization = true;

}

}

if (is_using_v_parameterization) {

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

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

} else {

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

}

if (schedule != DEFAULT) {

switch (schedule) {

case DISCRETE:

LOG_INFO("running with discrete schedule");

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

break;

case KARRAS:

LOG_INFO("running with Karras schedule");

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

break;

case DEFAULT:

// Don't touch anything.

break;

default:

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

abort();

}

}

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

denoiser->schedule->alphas_cumprod[i] = ((float*)alphas_cumprod_tensor->data)[i];

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

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

}

LOG_DEBUG("finished loaded file");

ggml_free(ctx);

if(control_net_path.size() > 0) {

ggml_backend_t cn_backend = NULL;

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

LOG_DEBUG("ControlNet: Using CPU backend");

cn_backend = ggml_backend_cpu_init();

} else {

cn_backend = backend;

}

if(!control_net.load_from_file(control_net_path, cn_backend, GGML_TYPE_F16 /* just f16 controlnet models */)) {

return false;

}

}

if (use_tiny_autoencoder) {

return tae_first_stage.load_from_file(taesd_path, backend);

}

return true;

}

bool is_using_v_parameterization_for_sd2(ggml_context* work_ctx) {

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); // [N, ]

struct ggml_tensor* t_emb = new_timestep_embedding(work_ctx, NULL, timesteps, diffusion_model.model_channels); // [N, model_channels]

int64_t t0 = ggml_time_ms();

ggml_set_f32(timesteps, 999);

set_timestep_embedding(timesteps, t_emb, diffusion_model.model_channels);

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

std::vector<struct ggml_tensor*> controls;

diffusion_model.alloc_compute_buffer(x_t, c, controls, t_emb);

diffusion_model.compute(out, n_threads, x_t, NULL, c, controls, 1.0f, t_emb);

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(const std::string& lora_name, float multiplier) {

int64_t t0 = ggml_time_ms();

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

if (!lora.load_from_file(backend)) {

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

return;

}

lora.multiplier = multiplier;

lora.apply(tensors, n_threads);

loras[lora_name] = lora;

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

if (lora_state.size() > 0 && model_data_type != GGML_TYPE_F16 && model_data_type != GGML_TYPE_F32) {

LOG_WARN("In quantized models when applying LoRA, the images have poor quality.");

}

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;

if (curr_lora_state.find(lora_name) != curr_lora_state.end()) {

float curr_multiplier = curr_lora_state[lora_name];

float multiplier_diff = multiplier - curr_multiplier;

if (multiplier_diff != 0.f) {

lora_state_diff[lora_name] = multiplier_diff;

}

} else {

lora_state_diff[lora_name] = multiplier;

}

}

for (auto& kv : lora_state_diff) {

apply_lora(kv.first, kv.second);

}

curr_lora_state = lora_state;

}

std::pair<ggml_tensor*, ggml_tensor*> get_learned_condition(ggml_context* work_ctx,

const std::string& text,

int clip_skip,

int width,

int height,

bool force_zero_embeddings = false) {

cond_stage_model.set_clip_skip(clip_skip);

auto tokens_and_weights = cond_stage_model.tokenize(text, true);

std::vector<int>& tokens = tokens_and_weights.first;

std::vector<float>& weights = tokens_and_weights.second;

int64_t t0 = ggml_time_ms();

struct ggml_tensor* pooled = NULL;

size_t total_hidden_size = cond_stage_model.text_model.hidden_size;

if (version == VERSION_XL) {

total_hidden_size += cond_stage_model.text_model2.hidden_size;

pooled = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, cond_stage_model.text_model2.projection_dim);

}

struct ggml_tensor* hidden_states = ggml_new_tensor_2d(work_ctx,

GGML_TYPE_F32,

total_hidden_size,

cond_stage_model.text_model.max_position_embeddings); // [N, n_token, hidden_size]

cond_stage_model.alloc_compute_buffer(work_ctx, (int)tokens.size());

cond_stage_model.compute(n_threads, tokens, hidden_states, pooled);

cond_stage_model.free_compute_buffer();

// if (pooled != NULL) {

// print_ggml_tensor(hidden_states);

// print_ggml_tensor(pooled);

// }

int64_t t1 = ggml_time_ms();

LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);

ggml_tensor* result = ggml_dup_tensor(work_ctx, hidden_states);

{

float original_mean = ggml_tensor_mean(hidden_states);

for (int i2 = 0; i2 < hidden_states->ne[2]; i2++) {

for (int i1 = 0; i1 < hidden_states->ne[1]; i1++) {

for (int i0 = 0; i0 < hidden_states->ne[0]; i0++) {

float value = ggml_tensor_get_f32(hidden_states, i0, i1, i2);

value *= weights[i1];

ggml_tensor_set_f32(result, value, i0, i1, i2);

}

}

}

float new_mean = ggml_tensor_mean(result);

ggml_tensor_scale(result, (original_mean / new_mean));

}

if (force_zero_embeddings) {

float* vec = (float*)result->data;

for (int i = 0; i < ggml_nelements(result); i++) {

vec[i] = 0;

}

}

ggml_tensor* vec = NULL;

if (version == VERSION_XL) {

int out_dim = 256;

vec = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, diffusion_model.adm_in_channels);

// [0:1280]

size_t offset = 0;

memcpy(vec->data, pooled->data, ggml_nbytes(pooled));

offset += ggml_nbytes(pooled);

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

// original_size_as_tuple

float orig_width = (float)width;

float orig_height = (float)height;

ggml_tensor_set_f32(timesteps, orig_height, 0);

ggml_tensor_set_f32(timesteps, orig_width, 1);

ggml_tensor* embed_view = ggml_view_2d(work_ctx, vec, out_dim, 2, ggml_type_size(GGML_TYPE_F32) * out_dim, offset);

offset += ggml_nbytes(embed_view);

set_timestep_embedding(timesteps, embed_view, out_dim);

// print_ggml_tensor(ggml_reshape_1d(work_ctx, embed_view, out_dim * 2));

// crop_coords_top_left

float crop_coord_top = 0.f;

float crop_coord_left = 0.f;

ggml_tensor_set_f32(timesteps, crop_coord_top, 0);

ggml_tensor_set_f32(timesteps, crop_coord_left, 1);

embed_view = ggml_view_2d(work_ctx, vec, out_dim, 2, ggml_type_size(GGML_TYPE_F32) * out_dim, offset);

offset += ggml_nbytes(embed_view);

set_timestep_embedding(timesteps, embed_view, out_dim);

// print_ggml_tensor(ggml_reshape_1d(work_ctx, embed_view, out_dim * 2));

// target_size_as_tuple

float target_width = (float)width;

float target_height = (float)height;

ggml_tensor_set_f32(timesteps, target_height, 0);

ggml_tensor_set_f32(timesteps, target_width, 1);

embed_view = ggml_view_2d(work_ctx, vec, out_dim, 2, ggml_type_size(GGML_TYPE_F32) * out_dim, offset);

offset += ggml_nbytes(embed_view);

set_timestep_embedding(timesteps, embed_view, out_dim);

// print_ggml_tensor(ggml_reshape_1d(work_ctx, embed_view, out_dim * 2));

GGML_ASSERT(offset == ggml_nbytes(vec));

}

// print_ggml_tensor(result);

return {result, vec};

}

ggml_tensor* sample(ggml_context* work_ctx,

ggml_tensor* x_t,

ggml_tensor* noise,

ggml_tensor* c,

ggml_tensor* c_vector,

ggml_tensor* uc,

ggml_tensor* uc_vector,

ggml_tensor* control_hint,

float cfg_scale,

sample_method_t method,

const std::vector<float>& sigmas,

float control_strength) {

size_t steps = sigmas.size() - 1;

// x_t = load_tensor_from_file(work_ctx, "./rand0.bin");

// print_ggml_tensor(x_t);

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

copy_ggml_tensor(x, x_t);

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

struct ggml_tensor* timesteps = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 1); // [N, ]

struct ggml_tensor* t_emb = new_timestep_embedding(work_ctx, NULL, timesteps, diffusion_model.model_channels); // [N, model_channels]

struct ggml_tensor* guided_hint = NULL;

if(control_hint != NULL) {

guided_hint = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, noised_input->ne[0], noised_input->ne[1], diffusion_model.model_channels, 1);

control_net.process_hint(guided_hint, n_threads, control_hint);

control_net.alloc_compute_buffer(noised_input, guided_hint, c, t_emb);

}

diffusion_model.alloc_compute_buffer(noised_input, c, control_net.controls, t_emb, c_vector);

bool has_unconditioned = cfg_scale != 1.0 && uc != NULL;

if (noise == NULL) {

// x = x * sigmas[0]

ggml_tensor_scale(x, sigmas[0]);

} else {

// xi = x + noise * sigma_sched[0]

ggml_tensor_scale(noise, sigmas[0]);

ggml_tensor_add(x, noise);

}

// denoise wrapper

struct ggml_tensor* out_cond = ggml_dup_tensor(work_ctx, x);

struct ggml_tensor* out_uncond = NULL;

if (has_unconditioned) {

out_uncond = ggml_dup_tensor(work_ctx, x);

}

struct ggml_tensor* denoised = ggml_dup_tensor(work_ctx, x);

auto denoise = [&](ggml_tensor* input, float sigma, int step) {

if (step == 1) {

pretty_progress(0, (int)steps, 0);

}

int64_t t0 = ggml_time_us();

float c_skip = 1.0f;

float c_out = 1.0f;

float c_in = 1.0f;

std::vector<float> scaling = denoiser->get_scalings(sigma);

if (scaling.size() == 3) { // CompVisVDenoiser

c_skip = scaling[0];

c_out = scaling[1];

c_in = scaling[2];

} else { // CompVisDenoiser

c_out = scaling[0];

c_in = scaling[1];

}

float t = denoiser->schedule->sigma_to_t(sigma);

ggml_set_f32(timesteps, t);

set_timestep_embedding(timesteps, t_emb, diffusion_model.model_channels);

copy_ggml_tensor(noised_input, input);

// noised_input = noised_input * c_in

ggml_tensor_scale(noised_input, c_in);

// cond

if(control_hint != NULL) {

control_net.compute(n_threads, noised_input, guided_hint, c, t_emb);

}

diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c, control_net.controls, control_strength, t_emb, c_vector);

float* negative_data = NULL;

if (has_unconditioned) {

// uncond

if(control_hint != NULL) {

control_net.compute(n_threads, noised_input, guided_hint, uc, t_emb);

}

diffusion_model.compute(out_uncond, n_threads, noised_input, NULL, uc, control_net.controls, control_strength, t_emb, uc_vector);

negative_data = (float*)out_uncond->data;

}

float* vec_denoised = (float*)denoised->data;

float* vec_input = (float*)input->data;

float* positive_data = (float*)out_cond->data;

int ne_elements = (int)ggml_nelements(denoised);

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

float latent_result = positive_data[i];

if (has_unconditioned) {

// out_uncond + cfg_scale * (out_cond - out_uncond)

latent_result = negative_data[i] + cfg_scale * (positive_data[i] - negative_data[i]);

}

// v = latent_result, eps = latent_result

// denoised = (v * c_out + input * c_skip) or (input + eps * c_out)

vec_denoised[i] = latent_result * c_out + vec_input[i] * c_skip;

}

int64_t t1 = ggml_time_us();

if (step > 0) {

pretty_progress(step, (int)steps, (t1 - t0) / 1000000.f);

// LOG_INFO("step %d sampling completed taking %.2fs", step, (t1 - t0) * 1.0f / 1000000);

}

};

// sample_euler_ancestral

switch (method) {

case EULER_A: {

struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);

struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);

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

float sigma = sigmas[i];

// denoise

denoise(x, sigma, i + 1);

// d = (x - denoised) / sigma

{

float* vec_d = (float*)d->data;

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

float* vec_denoised = (float*)denoised->data;

for (int i = 0; i < ggml_nelements(d); i++) {

vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma;

}

}

// get_ancestral_step

float sigma_up = std::min(sigmas[i + 1],

std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i])));

float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up);

// Euler method

float dt = sigma_down - sigmas[i];

// x = x + d * dt

{

float* vec_d = (float*)d->data;

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

for (int i = 0; i < ggml_nelements(x); i++) {

vec_x[i] = vec_x[i] + vec_d[i] * dt;

}

}

if (sigmas[i + 1] > 0) {

// x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up

ggml_tensor_set_f32_randn(noise, rng);

// noise = load_tensor_from_file(work_ctx, "./rand" + std::to_string(i+1) + ".bin");

{

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

float* vec_noise = (float*)noise->data;

for (int i = 0; i < ggml_nelements(x); i++) {

vec_x[i] = vec_x[i] + vec_noise[i] * sigma_up;

}

}

}

}

} break;

case EULER: // Implemented without any sigma churn

{

struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);

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

float sigma = sigmas[i];

// denoise

denoise(x, sigma, i + 1);

// d = (x - denoised) / sigma

{

float* vec_d = (float*)d->data;

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

float* vec_denoised = (float*)denoised->data;

for (int j = 0; j < ggml_nelements(d); j++) {

vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigma;

}

}

float dt = sigmas[i + 1] - sigma;

// x = x + d * dt

{

float* vec_d = (float*)d->data;

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

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x[j] = vec_x[j] + vec_d[j] * dt;

}

}

}

} break;

case HEUN: {

struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);

struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);

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

// denoise

denoise(x, sigmas[i], -(i + 1));

// d = (x - denoised) / sigma

{

float* vec_d = (float*)d->data;

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

float* vec_denoised = (float*)denoised->data;

for (int j = 0; j < ggml_nelements(x); j++) {

vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i];

}

}

float dt = sigmas[i + 1] - sigmas[i];

if (sigmas[i + 1] == 0) {

// Euler step

// x = x + d * dt

float* vec_d = (float*)d->data;

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

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x[j] = vec_x[j] + vec_d[j] * dt;

}

} else {

// Heun step

float* vec_d = (float*)d->data;

float* vec_d2 = (float*)d->data;

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

float* vec_x2 = (float*)x2->data;

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x2[j] = vec_x[j] + vec_d[j] * dt;

}

denoise(x2, sigmas[i + 1], i + 1);

float* vec_denoised = (float*)denoised->data;

for (int j = 0; j < ggml_nelements(x); j++) {

float d2 = (vec_x2[j] - vec_denoised[j]) / sigmas[i + 1];

vec_d[j] = (vec_d[j] + d2) / 2;

vec_x[j] = vec_x[j] + vec_d[j] * dt;

}

}

}

} break;

case DPM2: {

struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);

struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);

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

// denoise

denoise(x, sigmas[i], i + 1);

// d = (x - denoised) / sigma

{

float* vec_d = (float*)d->data;

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

float* vec_denoised = (float*)denoised->data;

for (int j = 0; j < ggml_nelements(x); j++) {

vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i];

}

}

if (sigmas[i + 1] == 0) {

// Euler step

// x = x + d * dt

float dt = sigmas[i + 1] - sigmas[i];

float* vec_d = (float*)d->data;

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

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x[j] = vec_x[j] + vec_d[j] * dt;

}

} else {

// DPM-Solver-2

float sigma_mid = exp(0.5f * (log(sigmas[i]) + log(sigmas[i + 1])));

float dt_1 = sigma_mid - sigmas[i];

float dt_2 = sigmas[i + 1] - sigmas[i];

float* vec_d = (float*)d->data;

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

float* vec_x2 = (float*)x2->data;

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x2[j] = vec_x[j] + vec_d[j] * dt_1;

}

denoise(x2, sigma_mid, i + 1);

float* vec_denoised = (float*)denoised->data;

for (int j = 0; j < ggml_nelements(x); j++) {

float d2 = (vec_x2[j] - vec_denoised[j]) / sigma_mid;

vec_x[j] = vec_x[j] + d2 * dt_2;

}

}

}

} break;

case DPMPP2S_A: {

struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);

struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);

struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);

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

// denoise

denoise(x, sigmas[i], i + 1);

// get_ancestral_step

float sigma_up = std::min(sigmas[i + 1],

std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i])));

float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up);

auto t_fn = [](float sigma) -> float { return -log(sigma); };

auto sigma_fn = [](float t) -> float { return exp(-t); };

if (sigma_down == 0) {

// Euler step

float* vec_d = (float*)d->data;

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

float* vec_denoised = (float*)denoised->data;

for (int j = 0; j < ggml_nelements(d); j++) {

vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i];

}

// TODO: If sigma_down == 0, isn't this wrong?

// But

// https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L525

// has this exactly the same way.

float dt = sigma_down - sigmas[i];

for (int j = 0; j < ggml_nelements(d); j++) {

vec_x[j] = vec_x[j] + vec_d[j] * dt;

}

} else {

// DPM-Solver++(2S)

float t = t_fn(sigmas[i]);

float t_next = t_fn(sigma_down);

float h = t_next - t;

float s = t + 0.5f * h;

float* vec_d = (float*)d->data;

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

float* vec_x2 = (float*)x2->data;

float* vec_denoised = (float*)denoised->data;

// First half-step

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x2[j] = (sigma_fn(s) / sigma_fn(t)) * vec_x[j] - (exp(-h * 0.5f) - 1) * vec_denoised[j];

}

denoise(x2, sigmas[i + 1], i + 1);

// Second half-step

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x[j] = (sigma_fn(t_next) / sigma_fn(t)) * vec_x[j] - (exp(-h) - 1) * vec_denoised[j];

}

}

// Noise addition

if (sigmas[i + 1] > 0) {

ggml_tensor_set_f32_randn(noise, rng);

{

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

float* vec_noise = (float*)noise->data;

for (int i = 0; i < ggml_nelements(x); i++) {

vec_x[i] = vec_x[i] + vec_noise[i] * sigma_up;

}

}

}

}

} break;

case DPMPP2M: // DPM++ (2M) from Karras et al (2022)

{

struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);

auto t_fn = [](float sigma) -> float { return -log(sigma); };

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

// denoise

denoise(x, sigmas[i], i + 1);

float t = t_fn(sigmas[i]);

float t_next = t_fn(sigmas[i + 1]);

float h = t_next - t;

float a = sigmas[i + 1] / sigmas[i];

float b = exp(-h) - 1.f;

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

float* vec_denoised = (float*)denoised->data;

float* vec_old_denoised = (float*)old_denoised->data;

if (i == 0 || sigmas[i + 1] == 0) {

// Simpler step for the edge cases

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x[j] = a * vec_x[j] - b * vec_denoised[j];

}

} else {

float h_last = t - t_fn(sigmas[i - 1]);

float r = h_last / h;

for (int j = 0; j < ggml_nelements(x); j++) {

float denoised_d = (1.f + 1.f / (2.f * r)) * vec_denoised[j] - (1.f / (2.f * r)) * vec_old_denoised[j];

vec_x[j] = a * vec_x[j] - b * denoised_d;

}

}

// old_denoised = denoised

for (int j = 0; j < ggml_nelements(x); j++) {

vec_old_denoised[j] = vec_denoised[j];

}

}

} break;

case DPMPP2Mv2: // Modified DPM++ (2M) from https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457

{

struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);

auto t_fn = [](float sigma) -> float { return -log(sigma); };

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

// denoise

denoise(x, sigmas[i], i + 1);

float t = t_fn(sigmas[i]);

float t_next = t_fn(sigmas[i + 1]);

float h = t_next - t;

float a = sigmas[i + 1] / sigmas[i];

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

float* vec_denoised = (float*)denoised->data;

float* vec_old_denoised = (float*)old_denoised->data;

if (i == 0 || sigmas[i + 1] == 0) {

// Simpler step for the edge cases

float b = exp(-h) - 1.f;

for (int j = 0; j < ggml_nelements(x); j++) {

vec_x[j] = a * vec_x[j] - b * vec_denoised[j];

}

} else {

float h_last = t - t_fn(sigmas[i - 1]);

float h_min = std::min(h_last, h);

float h_max = std::max(h_last, h);

float r = h_max / h_min;

float h_d = (h_max + h_min) / 2.f;

float b = exp(-h_d) - 1.f;

for (int j = 0; j < ggml_nelements(x); j++) {

float denoised_d = (1.f + 1.f / (2.f * r)) * vec_denoised[j] - (1.f / (2.f * r)) * vec_old_denoised[j];

vec_x[j] = a * vec_x[j] - b * denoised_d;

}

}

// old_denoised = denoised

for (int j = 0; j < ggml_nelements(x); j++) {

vec_old_denoised[j] = vec_denoised[j];

}

}

} break;

case LCM: // Latent Consistency Models

{

struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);

struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);

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

float sigma = sigmas[i];

// denoise

denoise(x, sigma, i + 1);