chore: fix cuda CI (#286)

.github/workflows/build.yml

@@ -146,7 +146,7 @@ jobs:
             sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip
 
   windows-latest-cmake:
-    runs-on: windows-latest
+    runs-on: windows-2019
 
     strategy:
       matrix:

CMakeLists.txt

@@ -34,7 +34,7 @@ option(SD_BUILD_SHARED_LIBS          "sd: build shared libs" OFF)
 
 if(SD_CUBLAS)
 	message("Use CUBLAS as backend stable-diffusion")
-    set(GGML_CUBLAS ON)
+    set(GGML_CUDA ON)
     add_definitions(-DSD_USE_CUBLAS)
 endif()
 

clip.hpp

@@ -679,8 +679,8 @@ public:
         class_embedding                     = ggml_repeat(ctx, class_embed_weight, class_embedding);      // [N, embed_dim]
         class_embedding                     = ggml_reshape_4d(ctx, class_embedding, 1, embed_dim, 1, N);  // [N, 1, embed_dim, 1]
 
-        struct ggml_tensor* x = ggml_concat(ctx, class_embedding, patch_embedding);    // [N, num_positions, embed_dim, 1]
-        x                     = ggml_reshape_3d(ctx, x, embed_dim, num_positions, N);  // [N, num_positions, embed_dim]
+        struct ggml_tensor* x = ggml_concat(ctx, class_embedding, patch_embedding, 2);  // [N, num_positions, embed_dim, 1]
+        x                     = ggml_reshape_3d(ctx, x, embed_dim, num_positions, N);   // [N, num_positions, embed_dim]
         x                     = ggml_add(ctx, x, position_embed_weight);
         return x;  // [N, num_positions, embed_dim]
     }
@@ -1036,7 +1036,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
                                              hidden_states2->ne[3]);
             hidden_states2 = ggml_cont(ctx, ggml_permute(ctx, hidden_states2, 2, 0, 1, 3));
 
-            hidden_states = ggml_concat(ctx, hidden_states, hidden_states2);  // [N, n_token, hidden_size + hidden_size2]
+            hidden_states = ggml_concat(ctx, hidden_states, hidden_states2, 2);  // [N, n_token, hidden_size + hidden_size2]
 
             hidden_states = ggml_cont(ctx, ggml_permute(ctx, hidden_states, 1, 2, 0, 3));
         }
@@ -1069,7 +1069,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
             auto token_embed_weight = text_model.get_token_embed_weight();
             token_embed_weight      = ggml_reshape_3d(compute_ctx, token_embed_weight, token_embed_weight->ne[0], 1, token_embed_weight->ne[1]);
             // concatenate custom embeddings
-            embeddings = ggml_concat(compute_ctx, token_embed_weight, custom_embeddings);
+            embeddings = ggml_concat(compute_ctx, token_embed_weight, custom_embeddings, 2);
             embeddings = ggml_reshape_2d(compute_ctx, embeddings, embeddings->ne[0], embeddings->ne[2]);
         }
 

denoiser.hpp

@@ -261,4 +261,401 @@ struct CompVisVDenoiser : public Denoiser {
     }
 };
 
+typedef std::function<ggml_tensor*(ggml_tensor*, float, int)> denoise_cb_t;
+
+// k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
+static void sample_k_diffusion(sample_method_t method,
+                               denoise_cb_t model,
+                               ggml_context* work_ctx,
+                               ggml_tensor* x,
+                               std::vector<float> sigmas,
+                               std::shared_ptr<RNG> rng) {
+    size_t steps = sigmas.size() - 1;
+    // 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
+                ggml_tensor* denoised = model(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
+                ggml_tensor* denoised = model(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
+                ggml_tensor* denoised = model(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;
+                    }
+
+                    ggml_tensor* denoised = model(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
+                ggml_tensor* denoised = model(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;
+                    }
+
+                    ggml_tensor* denoised = model(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
+                ggml_tensor* denoised = model(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];
+                    }
+
+                    ggml_tensor* denoised = model(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
+                ggml_tensor* denoised = model(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
+                ggml_tensor* denoised = model(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
+                ggml_tensor* denoised = model(x, sigma, i + 1);
+
+                // x = denoised
+                {
+                    float* vec_x        = (float*)x->data;
+                    float* vec_denoised = (float*)denoised->data;
+                    for (int j = 0; j < ggml_nelements(x); j++) {
+                        vec_x[j] = vec_denoised[j];
+                    }
+                }
+
+                if (sigmas[i + 1] > 0) {
+                    // x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
+                    ggml_tensor_set_f32_randn(noise, rng);
+                    // noise = load_tensor_from_file(res_ctx, "./rand" + std::to_string(i+1) + ".bin");
+                    {
+                        float* vec_x     = (float*)x->data;
+                        float* vec_noise = (float*)noise->data;
+
+                        for (int j = 0; j < ggml_nelements(x); j++) {
+                            vec_x[j] = vec_x[j] + sigmas[i + 1] * vec_noise[j];
+                        }
+                    }
+                }
+            }
+        } break;
+
+        default:
+            LOG_ERROR("Attempting to sample with nonexisting sample method %i", method);
+            abort();
+    }
+}
+
 #endif  // __DENOISER_HPP__

esrgan.hpp

@@ -42,13 +42,13 @@ public:
         auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
 
         auto x1    = lrelu(ctx, conv1->forward(ctx, x));
-        auto x_cat = ggml_concat(ctx, x, x1);
+        auto x_cat = ggml_concat(ctx, x, x1, 2);
         auto x2    = lrelu(ctx, conv2->forward(ctx, x_cat));
-        x_cat      = ggml_concat(ctx, x_cat, x2);
+        x_cat      = ggml_concat(ctx, x_cat, x2, 2);
         auto x3    = lrelu(ctx, conv3->forward(ctx, x_cat));
-        x_cat      = ggml_concat(ctx, x_cat, x3);
+        x_cat      = ggml_concat(ctx, x_cat, x3, 2);
         auto x4    = lrelu(ctx, conv4->forward(ctx, x_cat));
-        x_cat      = ggml_concat(ctx, x_cat, x4);
+        x_cat      = ggml_concat(ctx, x_cat, x4, 2);
         auto x5    = conv5->forward(ctx, x_cat);
 
         x5 = ggml_add(ctx, ggml_scale(ctx, x5, 0.2f), x);

ggml_extend.hpp

@@ -917,7 +917,7 @@ public:
             return NULL;
         }
         // it's performing a compute, check if backend isn't cpu
-        if (!ggml_backend_is_cpu(backend) && tensor->backend == GGML_BACKEND_TYPE_CPU) {
+        if (!ggml_backend_is_cpu(backend) && (tensor->buffer == NULL || ggml_backend_buffer_is_host(tensor->buffer))) {
             // pass input tensors to gpu memory
             auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
 

model.cpp

@@ -571,10 +571,9 @@ void convert_tensor(void* src,
         if (dst_type == GGML_TYPE_F16) {
             ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
         } else {
-            int64_t hist[16];
             std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
             const float* im = imatrix.data();
-            ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, hist, im);
+            ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im);
         }
     } else if (dst_type == GGML_TYPE_F32) {
         if (src_type == GGML_TYPE_F16) {
@@ -602,10 +601,9 @@ void convert_tensor(void* src,
         if (dst_type == GGML_TYPE_F16) {
             ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
         } else {
-            int64_t hist[16];
             std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
             const float* im = imatrix.data();
-            ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, hist, im);
+            ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im);
         }
     }
 }

pmid.hpp

@@ -64,7 +64,7 @@ public:
         auto prompt_embeds0 = ggml_cont(ctx, ggml_permute(ctx, prompt_embeds, 2, 0, 1, 3));
         auto id_embeds0     = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 2, 0, 1, 3));
         // concat is along dim 2
-        auto stacked_id_embeds = ggml_concat(ctx, prompt_embeds0, id_embeds0);
+        auto stacked_id_embeds = ggml_concat(ctx, prompt_embeds0, id_embeds0, 2);
         stacked_id_embeds      = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 1, 2, 0, 3));
 
         // stacked_id_embeds = mlp1.forward(ctx, stacked_id_embeds);
@@ -102,12 +102,12 @@ public:
 
         stacked_id_embeds = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 0, 2, 1, 3));
         if (left && right) {
-            stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds);
-            stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right);
+            stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds, 2);
+            stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right, 2);
         } else if (left) {
-            stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds);
+            stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds, 2);
         } else if (right) {
-            stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right);
+            stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right, 2);
         }
         stacked_id_embeds                         = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 0, 2, 1, 3));
         class_tokens_mask                         = ggml_cont(ctx, ggml_transpose(ctx, class_tokens_mask));
@@ -146,7 +146,7 @@ struct PhotoMakerIDEncoderBlock : public CLIPVisionModelProjection {
         id_embeds   = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 2, 0, 1, 3));
         id_embeds_2 = ggml_cont(ctx, ggml_permute(ctx, id_embeds_2, 2, 0, 1, 3));
 
-        id_embeds = ggml_concat(ctx, id_embeds, id_embeds_2);  // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right
+        id_embeds = ggml_concat(ctx, id_embeds, id_embeds_2, 2);  // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right
         id_embeds = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 1, 2, 0, 3));
 
         struct ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,

stable-diffusion.cpp

@@ -877,7 +877,7 @@ public:
         }
         struct ggml_tensor* denoised = ggml_dup_tensor(work_ctx, x);
 
-        auto denoise = [&](ggml_tensor* input, float sigma, int step) {
+        auto denoise = [&](ggml_tensor* input, float sigma, int step) -> ggml_tensor* {
             if (step == 1) {
                 pretty_progress(0, (int)steps, 0);
             }
@@ -983,393 +983,11 @@ public:
                 pretty_progress(step, (int)steps, (t1 - t0) / 1000000.f);
                 // LOG_INFO("step %d sampling completed taking %.2fs", step, (t1 - t0) * 1.0f / 1000000);
             }
+            return denoised;
         };
 
-        // 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);
+        sample_k_diffusion(method, denoise, work_ctx, x, sigmas, rng);
 
-                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);
-
-                    // x = denoised
-                    {
-                        float* vec_x        = (float*)x->data;
-                        float* vec_denoised = (float*)denoised->data;
-                        for (int j = 0; j < ggml_nelements(x); j++) {
-                            vec_x[j] = vec_denoised[j];
-                        }
-                    }
-
-                    if (sigmas[i + 1] > 0) {
-                        // x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
-                        ggml_tensor_set_f32_randn(noise, rng);
-                        // noise = load_tensor_from_file(res_ctx, "./rand" + std::to_string(i+1) + ".bin");
-                        {
-                            float* vec_x     = (float*)x->data;
-                            float* vec_noise = (float*)noise->data;
-
-                            for (int j = 0; j < ggml_nelements(x); j++) {
-                                vec_x[j] = vec_x[j] + sigmas[i + 1] * vec_noise[j];
-                            }
-                        }
-                    }
-                }
-            } break;
-
-            default:
-                LOG_ERROR("Attempting to sample with nonexisting sample method %i", method);
-                abort();
-        }
         if (control_net) {
             control_net->free_control_ctx();
             control_net->free_compute_buffer();
@@ -1717,7 +1335,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx,
     for (int b = 0; b < batch_count; b++) {
         int64_t sampling_start = ggml_time_ms();
         int64_t cur_seed       = seed + b;
-        LOG_INFO("generating image: %i/%i - seed %i", b + 1, batch_count, cur_seed);
+        LOG_INFO("generating image: %i/%i - seed %" PRId64, b + 1, batch_count, cur_seed);
 
         sd_ctx->sd->rng->manual_seed(cur_seed);
         struct ggml_tensor* x_t   = NULL;

stable-diffusion.h

@@ -60,12 +60,11 @@ enum sd_type_t {
     SD_TYPE_Q4_0 = 2,
     SD_TYPE_Q4_1 = 3,
     // SD_TYPE_Q4_2 = 4, support has been removed
-    // SD_TYPE_Q4_3 (5) support has been removed
-    SD_TYPE_Q5_0 = 6,
-    SD_TYPE_Q5_1 = 7,
-    SD_TYPE_Q8_0 = 8,
-    SD_TYPE_Q8_1 = 9,
-    // k-quantizations
+    // SD_TYPE_Q4_3 = 5, support has been removed
+    SD_TYPE_Q5_0    = 6,
+    SD_TYPE_Q5_1    = 7,
+    SD_TYPE_Q8_0    = 8,
+    SD_TYPE_Q8_1    = 9,
     SD_TYPE_Q2_K    = 10,
     SD_TYPE_Q3_K    = 11,
     SD_TYPE_Q4_K    = 12,
@@ -80,9 +79,13 @@ enum sd_type_t {
     SD_TYPE_IQ3_S   = 21,
     SD_TYPE_IQ2_S   = 22,
     SD_TYPE_IQ4_XS  = 23,
-    SD_TYPE_I8,
-    SD_TYPE_I16,
-    SD_TYPE_I32,
+    SD_TYPE_I8      = 24,
+    SD_TYPE_I16     = 25,
+    SD_TYPE_I32     = 26,
+    SD_TYPE_I64     = 27,
+    SD_TYPE_F64     = 28,
+    SD_TYPE_IQ1_M   = 29,
+    SD_TYPE_BF16    = 30,
     SD_TYPE_COUNT,
 };
 

unet.hpp

@@ -396,7 +396,7 @@ public:
             if (c_concat->ne[3] != x->ne[3]) {
                 c_concat = ggml_repeat(ctx, c_concat, x);
             }
-            x = ggml_concat(ctx, x, c_concat);
+            x = ggml_concat(ctx, x, c_concat, 2);
         }
 
         if (y != NULL) {
@@ -491,7 +491,7 @@ public:
                     control_offset--;
                 }
 
-                h = ggml_concat(ctx, h, h_skip);
+                h = ggml_concat(ctx, h, h_skip, 2);
 
                 std::string name = "output_blocks." + std::to_string(output_block_idx) + ".0";