Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 60 additions & 52 deletions src/denoiser.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -549,8 +549,48 @@ struct Denoiser {
};

struct CompVisDenoiser : public Denoiser {
float sigmas[TIMESTEPS];
float log_sigmas[TIMESTEPS];

private:
struct Constants {
const float beta_start = 0.00085f;
const float beta_end = 0.0120f;
float alphas_cumprods[TIMESTEPS];
float sigmas[TIMESTEPS];
float log_sigmas[TIMESTEPS];
Constants() {
double ls_sqrt = std::sqrt(static_cast<double>(beta_start));
double le_sqrt = std::sqrt(static_cast<double>(beta_end));
double step = (le_sqrt - ls_sqrt) / (TIMESTEPS - 1);
double alphas_cumprod = 1.0;

for (int i = 0; i < TIMESTEPS; ++i) {
double sqrt_beta = ls_sqrt + step * i;
alphas_cumprod *= (1.0 - (sqrt_beta * sqrt_beta));
double sigma = std::sqrt((1.0 - alphas_cumprod) / alphas_cumprod);
alphas_cumprods[i] = static_cast<float>(alphas_cumprod);
sigmas[i] = static_cast<float>(sigma);
log_sigmas[i] = static_cast<float>(std::log(sigma));
}
}
static const Constants& get_instance() {
static Constants instance;
return instance;
}
};

const float *sigmas = get_sigmas();
const float *log_sigmas = get_log_sigmas();

public:
static const float* get_sigmas() {
return Constants::get_instance().sigmas;
}
static const float* get_log_sigmas() {
return Constants::get_instance().log_sigmas;
}
static const float* get_alphas_cumprods() {
return Constants::get_instance().alphas_cumprods;
}

float sigma_data = 1.0f;

Expand All @@ -564,20 +604,12 @@ struct CompVisDenoiser : public Denoiser {

float sigma_to_t(float sigma) override {
float log_sigma = std::log(sigma);
std::vector<float> dists;
dists.reserve(TIMESTEPS);
for (float log_sigma_val : log_sigmas) {
dists.push_back(log_sigma - log_sigma_val);
}
const float* high_ptr = std::upper_bound(log_sigmas, log_sigmas + TIMESTEPS, log_sigma);

int low_idx = 0;
for (size_t i = 0; i < TIMESTEPS; i++) {
if (dists[i] >= 0) {
low_idx++;
}
}
low_idx = std::min(std::max(low_idx - 1, 0), TIMESTEPS - 2);
int high_idx = low_idx + 1;
int high_idx = static_cast<int>(high_ptr - log_sigmas);
int low_idx = high_idx - 1;
low_idx = std::clamp(low_idx, 0, TIMESTEPS - 2);
high_idx = low_idx + 1;

float low = log_sigmas[low_idx];
float high = log_sigmas[high_idx];
Expand Down Expand Up @@ -1242,27 +1274,15 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
float beta_start = 0.00085f;
float beta_end = 0.0120f;
std::vector<double> alphas_cumprod(TIMESTEPS);
std::vector<double> compvis_sigmas(TIMESTEPS);
for (int i = 0; i < TIMESTEPS; i++) {
alphas_cumprod[i] =
(i == 0 ? 1.0f : alphas_cumprod[i - 1]) *
(1.0f -
std::pow(sqrtf(beta_start) +
(sqrtf(beta_end) - sqrtf(beta_start)) *
((float)i / (TIMESTEPS - 1)),
2));
compvis_sigmas[i] =
std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
}

const float* alphas_cumprod = CompVisDenoiser::get_alphas_cumprods();
const float* compvis_sigmas = CompVisDenoiser::get_sigmas();

int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
int timestep = static_cast<int>(roundf(TIMESTEPS - i * ((float)TIMESTEPS / steps))) - 1;
int prev_timestep = timestep - TIMESTEPS / steps;
float sigma = static_cast<float>(compvis_sigmas[timestep]);
float sigma = compvis_sigmas[timestep];
if (i == 0) {
x *= std::sqrt(sigma * sigma + 1) / sigma;
} else {
Expand All @@ -1276,8 +1296,8 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
sd::Tensor<float> model_output = std::move(model_output_opt);
model_output = (x - model_output) * (1.0f / sigma);

float alpha_prod_t = static_cast<float>(alphas_cumprod[timestep]);
float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
float alpha_prod_t = alphas_cumprod[timestep];
float alpha_prod_t_prev = prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0];
float beta_prod_t = 1.0f - alpha_prod_t;

sd::Tensor<float> pred_original_sample = ((x / std::sqrt(sigma * sigma + 1)) -
Expand All @@ -1304,29 +1324,17 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
float beta_start = 0.00085f;
float beta_end = 0.0120f;
std::vector<double> alphas_cumprod(TIMESTEPS);
std::vector<double> compvis_sigmas(TIMESTEPS);
for (int i = 0; i < TIMESTEPS; i++) {
alphas_cumprod[i] =
(i == 0 ? 1.0f : alphas_cumprod[i - 1]) *
(1.0f -
std::pow(sqrtf(beta_start) +
(sqrtf(beta_end) - sqrtf(beta_start)) *
((float)i / (TIMESTEPS - 1)),
2));
compvis_sigmas[i] =
std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
}

const float* alphas_cumprod = CompVisDenoiser::get_alphas_cumprods();
const float* compvis_sigmas = CompVisDenoiser::get_sigmas();

int original_steps = 50;
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
int timestep = TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor(i * ((float)original_steps / steps));
int prev_timestep = i >= steps - 1 ? 0 : TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor((i + 1) * ((float)original_steps / steps));
int timestep_s = (int)floor((1 - eta) * prev_timestep);
float sigma = static_cast<float>(compvis_sigmas[timestep]);
float sigma = compvis_sigmas[timestep];

if (i == 0) {
x *= std::sqrt(sigma * sigma + 1) / sigma;
Expand All @@ -1341,10 +1349,10 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
sd::Tensor<float> model_output = std::move(model_output_opt);
model_output = (x - model_output) * (1.0f / sigma);

float alpha_prod_t = static_cast<float>(alphas_cumprod[timestep]);
float alpha_prod_t = alphas_cumprod[timestep];
float beta_prod_t = 1.0f - alpha_prod_t;
float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
float alpha_prod_s = static_cast<float>(alphas_cumprod[timestep_s]);
float alpha_prod_t_prev = prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0];
float alpha_prod_s = alphas_cumprod[timestep_s];
float beta_prod_s = 1.0f - alpha_prod_s;

sd::Tensor<float> pred_original_sample = ((x / std::sqrt(sigma * sigma + 1)) -
Expand Down
23 changes: 3 additions & 20 deletions src/stable-diffusion.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -73,19 +73,9 @@ const char* sampling_methods_str[] = {

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

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;
}
void calculate_alphas_cumprod(float* alphas_cumprod) {
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've preserved this, but I'm not sure the alphas_cumprod tensor is really being used.

const float * src = CompVisDenoiser::get_alphas_cumprods();
std::copy(src, src + TIMESTEPS, alphas_cumprod);
}

static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
Expand Down Expand Up @@ -980,13 +970,6 @@ class StableDiffusionGGML {
}
}

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

ggml_free(ctx);
Expand Down
Loading