Skip to content

Commit c133d7b

Browse files
daniandthewebdaniandtheweb
committed
feat: added ipndm and ipndm_v samplers
1 parent 697d000 commit c133d7b

File tree

5 files changed

+159
-2
lines changed

5 files changed

+159
-2
lines changed

‎README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,7 @@ arguments:
217217
1.0 corresponds to full destruction of information in init image
218218
-H, --height H image height, in pixel space (default: 512)
219219
-W, --width W image width, in pixel space (default: 512)
220-
--sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, lcm}
220+
--sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm}
221221
sampling method (default: "euler_a")
222222
--steps STEPS number of sample steps (default: 20)
223223
--rng {std_default, cuda} RNG (default: cuda)

‎denoiser.hpp

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#ifndef __DENOISER_HPP__
22
#define __DENOISER_HPP__
33

4+
#include "ggml.h"
45
#include "ggml_extend.hpp"
56

67
/*================================================= CompVisDenoiser ==================================================*/
@@ -705,6 +706,156 @@ static void sample_k_diffusion(sample_method_t method,
705706
}
706707
}
707708
} break;
709+
case IPNDM: {
710+
int max_order = 4;
711+
ggml_tensor* x_next = x;
712+
std::vector<ggml_tensor*> buffer_model;
713+
714+
for (int i = 0; i < steps; i++) {
715+
float sigma = sigmas[i];
716+
float sigma_next = sigmas[i + 1];
717+
718+
ggml_tensor* x_cur = x_next;
719+
float* vec_x_cur = (float*)x_cur->data;
720+
float* vec_x_next = (float*)x_next->data;
721+
722+
// Denoising step
723+
ggml_tensor* denoised = model(x_cur, sigma, i + 1);
724+
float* vec_denoised = (float*)denoised->data;
725+
// d_cur = (x_cur - denoised) / sigma
726+
struct ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x_cur);
727+
float* vec_d_cur = (float*)d_cur->data;
728+
729+
for (int j = 0; j < ggml_nelements(d_cur); j++) {
730+
vec_d_cur[j] = (vec_x_cur[j] - vec_denoised[j]) / sigma;
731+
}
732+
733+
int order = std::min(max_order, i + 1);
734+
735+
// Calculate vec_x_next based on the order
736+
switch (order) {
737+
case 1: // First Euler step
738+
for (int j = 0; j < ggml_nelements(x_next); j++) {
739+
vec_x_next[j] = vec_x_cur[j] + (sigma_next - sigma) * vec_d_cur[j];
740+
}
741+
break;
742+
743+
case 2: // Use one history point
744+
{
745+
float* vec_d_prev1 = (float*)buffer_model.back()->data;
746+
for (int j = 0; j < ggml_nelements(x_next); j++) {
747+
vec_x_next[j] = vec_x_cur[j] + (sigma_next - sigma) * (3 * vec_d_cur[j] - vec_d_prev1[j]) / 2;
748+
}
749+
}
750+
break;
751+
752+
case 3: // Use two history points
753+
{
754+
float* vec_d_prev1 = (float*)buffer_model.back()->data;
755+
float* vec_d_prev2 = (float*)buffer_model[buffer_model.size() - 2]->data;
756+
for (int j = 0; j < ggml_nelements(x_next); j++) {
757+
vec_x_next[j] = vec_x_cur[j] + (sigma_next - sigma) * (23 * vec_d_cur[j] - 16 * vec_d_prev1[j] + 5 * vec_d_prev2[j]) / 12;
758+
}
759+
}
760+
break;
761+
762+
case 4: // Use three history points
763+
{
764+
float* vec_d_prev1 = (float*)buffer_model.back()->data;
765+
float* vec_d_prev2 = (float*)buffer_model[buffer_model.size() - 2]->data;
766+
float* vec_d_prev3 = (float*)buffer_model[buffer_model.size() - 3]->data;
767+
for (int j = 0; j < ggml_nelements(x_next); j++) {
768+
vec_x_next[j] = vec_x_cur[j] + (sigma_next - sigma) * (55 * vec_d_cur[j] - 59 * vec_d_prev1[j] + 37 * vec_d_prev2[j] - 9 * vec_d_prev3[j]) / 24;
769+
}
770+
}
771+
break;
772+
}
773+
774+
// Manage buffer_model
775+
if (buffer_model.size() == max_order - 1) {
776+
// Shift elements to the left
777+
for (int k = 0; k < max_order - 2; k++) {
778+
buffer_model[k] = buffer_model[k + 1];
779+
}
780+
buffer_model.back() = d_cur; // Replace the last element with d_cur
781+
} else {
782+
buffer_model.push_back(d_cur);
783+
}
784+
}
785+
} break;
786+
case IPNDM_V: {
787+
int max_order = 4;
788+
std::vector<ggml_tensor*> buffer_model;
789+
ggml_tensor* x_next = x;
790+
791+
for (int i = 0; i < steps; i++) {
792+
float sigma = sigmas[i];
793+
float t_next = sigmas[i + 1];
794+
795+
// Denoising step
796+
ggml_tensor* denoised = model(x, sigma, i + 1);
797+
float* vec_denoised = (float*)denoised->data;
798+
struct ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x);
799+
float* vec_d_cur = (float*)d_cur->data;
800+
float* vec_x = (float*)x->data;
801+
802+
// d_cur = (x - denoised) / sigma
803+
for (int j = 0; j < ggml_nelements(d_cur); j++) {
804+
vec_d_cur[j] = (vec_x[j] - vec_denoised[j]) / sigma;
805+
}
806+
807+
int order = std::min(max_order, i + 1);
808+
float h_n = t_next - sigma;
809+
float h_n_1 = (i > 0) ? (sigma - sigmas[i - 1]) : h_n;
810+
811+
switch (order) {
812+
case 1:
813+
for (int j = 0; j < ggml_nelements(x_next); j++) {
814+
vec_x[j] += vec_d_cur[j] * h_n;
815+
}
816+
break;
817+
818+
case 2: {
819+
float* vec_d_prev1 = (float*)buffer_model.back()->data;
820+
for (int j = 0; j < ggml_nelements(x_next); j++) {
821+
vec_x[j] += h_n * ((2 + (h_n / h_n_1)) * vec_d_cur[j] - (h_n / h_n_1) * vec_d_prev1[j]) / 2;
822+
}
823+
break;
824+
}
825+
826+
case 3: {
827+
float h_n_2 = (i > 1) ? (sigmas[i - 1] - sigmas[i - 2]) : h_n_1;
828+
float* vec_d_prev1 = (float*)buffer_model.back()->data;
829+
float* vec_d_prev2 = (buffer_model.size() > 1) ? (float*)buffer_model[buffer_model.size() - 2]->data : vec_d_prev1;
830+
for (int j = 0; j < ggml_nelements(x_next); j++) {
831+
vec_x[j] += h_n * ((23 * vec_d_cur[j] - 16 * vec_d_prev1[j] + 5 * vec_d_prev2[j]) / 12);
832+
}
833+
break;
834+
}
835+
836+
case 4: {
837+
float h_n_2 = (i > 1) ? (sigmas[i - 1] - sigmas[i - 2]) : h_n_1;
838+
float h_n_3 = (i > 2) ? (sigmas[i - 2] - sigmas[i - 3]) : h_n_2;
839+
float* vec_d_prev1 = (float*)buffer_model.back()->data;
840+
float* vec_d_prev2 = (buffer_model.size() > 1) ? (float*)buffer_model[buffer_model.size() - 2]->data : vec_d_prev1;
841+
float* vec_d_prev3 = (buffer_model.size() > 2) ? (float*)buffer_model[buffer_model.size() - 3]->data : vec_d_prev2;
842+
for (int j = 0; j < ggml_nelements(x_next); j++) {
843+
vec_x[j] += h_n * ((55 * vec_d_cur[j] - 59 * vec_d_prev1[j] + 37 * vec_d_prev2[j] - 9 * vec_d_prev3[j]) / 24);
844+
}
845+
break;
846+
}
847+
}
848+
849+
// Manage buffer_model
850+
if (buffer_model.size() == max_order - 1) {
851+
buffer_model.erase(buffer_model.begin());
852+
}
853+
buffer_model.push_back(d_cur);
854+
855+
// Prepare the next d tensor
856+
d_cur = ggml_dup_tensor(work_ctx, x_next);
857+
}
858+
} break;
708859
case LCM: // Latent Consistency Models
709860
{
710861
struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);

‎examples/cli/main.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ const char* sample_method_str[] = {
3737
"dpm++2s_a",
3838
"dpm++2m",
3939
"dpm++2mv2",
40+
"ipndm",
41+
"ipndm_v",
4042
"lcm",
4143
};
4244

@@ -187,7 +189,7 @@ void print_usage(int argc, const char* argv[]) {
187189
printf(" 1.0 corresponds to full destruction of information in init image\n");
188190
printf(" -H, --height H image height, in pixel space (default: 512)\n");
189191
printf(" -W, --width W image width, in pixel space (default: 512)\n");
190-
printf(" --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, lcm}\n");
192+
printf(" --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm}\n");
191193
printf(" sampling method (default: \"euler_a\")\n");
192194
printf(" --steps STEPS number of sample steps (default: 20)\n");
193195
printf(" --rng {std_default, cuda} RNG (default: cuda)\n");

‎stable-diffusion.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ const char* sampling_methods_str[] = {
3939
"DPM++ (2s)",
4040
"DPM++ (2M)",
4141
"modified DPM++ (2M)",
42+
"iPNDM",
43+
"iPNDM_v",
4244
"LCM",
4345
};
4446

‎stable-diffusion.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ enum sample_method_t {
4141
DPMPP2S_A,
4242
DPMPP2M,
4343
DPMPP2Mv2,
44+
IPNDM,
45+
IPNDM_V,
4446
LCM,
4547
N_SAMPLE_METHODS
4648
};

0 commit comments

Comments
 (0)