chore: fix cuda CI (#286 )

fix: sample_k_diffusion should be static
sync: update ggml
2026-10-03 02:38:25 +08:00 · 2024-06-12 23:13:24 +08:00 · 2024-06-10 23:04:02 +08:00 · 2024-06-10 22:59:36 +08:00 · 2024-06-10 22:42:15 +08:00 · 2024-06-01 14:01:58 +08:00
17 changed files with 1130 additions and 703 deletions
@@ -4,17 +4,36 @@ on:
  workflow_dispatch: # allows manual triggering
    inputs:
      create_release:
-        description: 'Create new release'
+        description: "Create new release"
        required: true
        type: boolean
  push:
    branches:
      - master
      - ci
-    paths: ['.github/workflows/**', '**/CMakeLists.txt', '**/Makefile', '**/*.h', '**/*.hpp', '**/*.c', '**/*.cpp', '**/*.cu']
+    paths:
+      [
+        ".github/workflows/**",
+        "**/CMakeLists.txt",
+        "**/Makefile",
+        "**/*.h",
+        "**/*.hpp",
+        "**/*.c",
+        "**/*.cpp",
+        "**/*.cu",
+      ]
  pull_request:
    types: [opened, synchronize, reopened]
-    paths: ['**/CMakeLists.txt', '**/Makefile', '**/*.h', '**/*.hpp', '**/*.c', '**/*.cpp', '**/*.cu']
+    paths:
+      [
+        "**/CMakeLists.txt",
+        "**/Makefile",
+        "**/*.h",
+        "**/*.hpp",
+        "**/*.c",
+        "**/*.cpp",
+        "**/*.cu",
+      ]

 env:
  BRANCH_NAME: ${{ github.head_ref || github.ref_name }}
@@ -67,12 +86,12 @@ jobs:

      - name: Upload artifacts
        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
        with:
+          name: 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
          path: |
            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

-
  macOS-latest-cmake:
    runs-on: macos-latest

@@ -120,29 +139,29 @@ jobs:

      - name: Upload artifacts
        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
        with:
+          name: 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
          path: |
            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:
        include:
-          - build: 'noavx'
-            defines: '-DGGML_AVX=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF -DSD_BUILD_SHARED_LIBS=ON'
-          - build: 'avx2'
-            defines: '-DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON'
-          - build: 'avx'
-            defines: '-DGGML_AVX2=OFF -DSD_BUILD_SHARED_LIBS=ON'
-          - build: 'avx512'
-            defines: '-DGGML_AVX512=ON -DSD_BUILD_SHARED_LIBS=ON'
-          - build: 'cuda12'
-            defines: '-DSD_CUBLAS=ON -DSD_BUILD_SHARED_LIBS=ON'
-          - build: 'rocm5.5'
+          - build: "noavx"
+            defines: "-DGGML_AVX=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF -DSD_BUILD_SHARED_LIBS=ON"
+          - build: "avx2"
+            defines: "-DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON"
+          - build: "avx"
+            defines: "-DGGML_AVX2=OFF -DSD_BUILD_SHARED_LIBS=ON"
+          - build: "avx512"
+            defines: "-DGGML_AVX512=ON -DSD_BUILD_SHARED_LIBS=ON"
+          - build: "cuda12"
+            defines: "-DSD_CUBLAS=ON -DSD_BUILD_SHARED_LIBS=ON"
+          - build: "rocm5.5"
            defines: '-G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS="gfx1100;gfx1102;gfx1030" -DSD_BUILD_SHARED_LIBS=ON'
    steps:
      - name: Clone
@@ -156,8 +175,8 @@ jobs:
        if: ${{ matrix.build == 'cuda12' }}
        uses: Jimver/cuda-toolkit@v0.2.11
        with:
-          cuda: '12.2.0'
-          method: 'network'
+          cuda: "12.2.0"
+          method: "network"
          sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "thrust", "visual_studio_integration"]'

      - name: Install rocm-toolkit
@@ -165,7 +184,7 @@ jobs:
        if: ${{ matrix.build == 'rocm5.5' }}
        uses: Cyberhan123/rocm-toolkit@v0.1.0
        with:
-          rocm: '5.5.0'
+          rocm: "5.5.0"

      - name: Install Ninja
        id: install-ninja
@@ -231,15 +250,17 @@ jobs:

      - name: Upload Cuda runtime
        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' && matrix.build == 'cuda12' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
        with:
+          name: sd-cudart-sd-bin-win-cu12-x64.zip
          path: |
            cudart-sd-bin-win-cu12-x64.zip

      - name: Upload artifacts
        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
        with:
+          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip
          path: |
            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip

@@ -256,7 +277,11 @@ jobs:
    steps:
      - name: Download artifacts
        id: download-artifact
-        uses: actions/download-artifact@v3
+        uses: actions/download-artifact@v4
+        with:
+          path: ./artifact
+          pattern: sd-*
+          merge-multiple: true

      - name: Get commit hash
        id: commit
@@ -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()

@@ -190,12 +190,13 @@ arguments:
  --rng {std_default, cuda}          RNG (default: cuda)
  -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)
  -b, --batch-count COUNT            number of images to generate.
-  --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)
+  --schedule {discrete, karras, ays} Denoiser sigma schedule (default: discrete)
  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)
                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
  --vae-tiling                       process vae in tiles to reduce memory usage
  --control-net-cpu                  keep controlnet in cpu (for low vram)
  --canny                            apply canny preprocessor (edge detection)
+  --color                            colors the logging tags according to level
  -v, --verbose                      print extra info
 ```

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

@@ -13,6 +13,7 @@ struct SigmaSchedule {
    float alphas_cumprod[TIMESTEPS];
    float sigmas[TIMESTEPS];
    float log_sigmas[TIMESTEPS];
+    int version = 0;

    virtual std::vector<float> get_sigmas(uint32_t n) = 0;

@@ -75,6 +76,144 @@ struct DiscreteSchedule : SigmaSchedule {
    }
 };

+/*
+https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html
+*/
+struct AYSSchedule : SigmaSchedule {
+    /* interp and linear_interp adapted from dpilger26's NumCpp library:
+     * https://github.com/dpilger26/NumCpp/tree/5e40aab74d14e257d65d3dc385c9ff9e2120c60e */
+    constexpr double interp(double left, double right, double perc) noexcept {
+        return (left * (1. - perc)) + (right * perc);
+    }
+
+    /* This will make the assumption that the reference x and y values are
+     * already sorted in ascending order because they are being generated as
+     * such in the calling function */
+    std::vector<double> linear_interp(std::vector<float> new_x,
+                                      const std::vector<float> ref_x,
+                                      const std::vector<float> ref_y) {
+        const size_t len_x = new_x.size();
+        size_t i           = 0;
+        size_t j           = 0;
+        std::vector<double> new_y(len_x);
+
+        if (ref_x.size() != ref_y.size()) {
+            LOG_ERROR("Linear Interoplation Failed: length mismatch");
+            return new_y;
+        }
+
+        /* serves as the bounds checking for the below while loop */
+        if ((new_x[0] < ref_x[0]) || (new_x[new_x.size() - 1] > ref_x[ref_x.size() - 1])) {
+            LOG_ERROR("Linear Interpolation Failed: bad bounds");
+            return new_y;
+        }
+
+        while (i < len_x) {
+            if ((ref_x[j] > new_x[i]) || (new_x[i] > ref_x[j + 1])) {
+                j++;
+                continue;
+            }
+
+            const double perc = static_cast<double>(new_x[i] - ref_x[j]) / static_cast<double>(ref_x[j + 1] - ref_x[j]);
+
+            new_y[i] = interp(ref_y[j], ref_y[j + 1], perc);
+            i++;
+        }
+
+        return new_y;
+    }
+
+    std::vector<float> linear_space(const float start, const float end, const size_t num_points) {
+        std::vector<float> result(num_points);
+        const float inc = (end - start) / (static_cast<float>(num_points - 1));
+
+        if (num_points > 0) {
+            result[0] = start;
+
+            for (size_t i = 1; i < num_points; i++) {
+                result[i] = result[i - 1] + inc;
+            }
+        }
+
+        return result;
+    }
+
+    std::vector<float> log_linear_interpolation(std::vector<float> sigma_in,
+                                                const size_t new_len) {
+        const size_t s_len        = sigma_in.size();
+        std::vector<float> x_vals = linear_space(0.f, 1.f, s_len);
+        std::vector<float> y_vals(s_len);
+
+        /* Reverses the input array to be ascending instead of descending,
+         * also hits it with a log, it is log-linear interpolation after all */
+        for (size_t i = 0; i < s_len; i++) {
+            y_vals[i] = std::log(sigma_in[s_len - i - 1]);
+        }
+
+        std::vector<float> new_x_vals  = linear_space(0.f, 1.f, new_len);
+        std::vector<double> new_y_vals = linear_interp(new_x_vals, x_vals, y_vals);
+        std::vector<float> results(new_len);
+
+        for (size_t i = 0; i < new_len; i++) {
+            results[i] = static_cast<float>(std::exp(new_y_vals[new_len - i - 1]));
+        }
+
+        return results;
+    }
+
+    std::vector<float> get_sigmas(uint32_t len) {
+        const std::vector<float> noise_levels[] = {
+            /* SD1.5 */
+            {14.6146412293f, 6.4745760956f, 3.8636745985f, 2.6946151520f,
+             1.8841921177f, 1.3943805092f, 0.9642583904f, 0.6523686016f,
+             0.3977456272f, 0.1515232662f, 0.0291671582f},
+            /* SDXL */
+            {14.6146412293f, 6.3184485287f, 3.7681790315f, 2.1811480769f,
+             1.3405244945f, 0.8620721141f, 0.5550693289f, 0.3798540708f,
+             0.2332364134f, 0.1114188177f, 0.0291671582f},
+            /* SVD */
+            {700.00f, 54.5f, 15.886f, 7.977f, 4.248f, 1.789f, 0.981f, 0.403f,
+             0.173f, 0.034f, 0.002f},
+        };
+
+        std::vector<float> inputs;
+        std::vector<float> results(len + 1);
+
+        switch (version) {
+            case VERSION_2_x: /* fallthrough */
+                LOG_WARN("AYS not designed for SD2.X models");
+            case VERSION_1_x:
+                LOG_INFO("AYS using SD1.5 noise levels");
+                inputs = noise_levels[0];
+                break;
+            case VERSION_XL:
+                LOG_INFO("AYS using SDXL noise levels");
+                inputs = noise_levels[1];
+                break;
+            case VERSION_SVD:
+                LOG_INFO("AYS using SVD noise levels");
+                inputs = noise_levels[2];
+                break;
+            default:
+                LOG_ERROR("Version not compatable with AYS scheduler");
+                return results;
+        }
+
+        /* Stretches those pre-calculated reference levels out to the desired
+         * size using log-linear interpolation */
+        if ((len + 1) != inputs.size()) {
+            results = log_linear_interpolation(inputs, len + 1);
+        } else {
+            results = inputs;
+        }
+
+        /* Not sure if this is strictly neccessary */
+        results[len] = 0.0f;
+
+        return results;
+    }
+};
+
 struct KarrasSchedule : SigmaSchedule {
    std::vector<float> get_sigmas(uint32_t n) {
        // These *COULD* be function arguments here,
@@ -122,4 +261,401 @@ struct CompVisVDenoiser : public Denoiser {
    }
 };

-#endif  // __DENOISER_HPP__
+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__
@@ -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);
@@ -43,6 +43,7 @@ const char* schedule_str[] = {
    "default",
    "discrete",
    "karras",
+    "ays",
 };

 const char* modes_str[] = {
@@ -190,12 +191,13 @@ void print_usage(int argc, const char* argv[]) {
    printf("  --rng {std_default, cuda}          RNG (default: cuda)\n");
    printf("  -s SEED, --seed SEED               RNG seed (default: 42, use random seed for < 0)\n");
    printf("  -b, --batch-count COUNT            number of images to generate.\n");
-    printf("  --schedule {discrete, karras}      Denoiser sigma schedule (default: discrete)\n");
+    printf("  --schedule {discrete, karras, ays} Denoiser sigma schedule (default: discrete)\n");
    printf("  --clip-skip N                      ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n");
    printf("                                     <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x\n");
    printf("  --vae-tiling                       process vae in tiles to reduce memory usage\n");
    printf("  --control-net-cpu                  keep controlnet in cpu (for low vram)\n");
    printf("  --canny                            apply canny preprocessor (edge detection)\n");
+    printf("  --color                            Colors the logging tags according to level\n");
    printf("  -v, --verbose                      print extra info\n");
 }

@@ -656,13 +658,16 @@ int main(int argc, const char* argv[]) {
        return 1;
    }

-    bool vae_decode_only        = true;
-    uint8_t* input_image_buffer = NULL;
+    bool vae_decode_only          = true;
+    uint8_t* input_image_buffer   = NULL;
+    uint8_t* control_image_buffer = NULL;
    if (params.mode == IMG2IMG || params.mode == IMG2VID) {
        vae_decode_only = false;

        int c              = 0;
-        input_image_buffer = stbi_load(params.input_path.c_str(), &params.width, &params.height, &c, 3);
+        int width          = 0;
+        int height         = 0;
+        input_image_buffer = stbi_load(params.input_path.c_str(), &width, &height, &c, 3);
        if (input_image_buffer == NULL) {
            fprintf(stderr, "load image from '%s' failed\n", params.input_path.c_str());
            return 1;
@@ -672,21 +677,22 @@ int main(int argc, const char* argv[]) {
            free(input_image_buffer);
            return 1;
        }
-        if (params.width <= 0) {
+        if (width <= 0) {
            fprintf(stderr, "error: the width of image must be greater than 0\n");
            free(input_image_buffer);
            return 1;
        }
-        if (params.height <= 0) {
+        if (height <= 0) {
            fprintf(stderr, "error: the height of image must be greater than 0\n");
            free(input_image_buffer);
            return 1;
        }

        // Resize input image ...
-        if (params.height % 64 != 0 || params.width % 64 != 0) {
-            int resized_height = params.height + (64 - params.height % 64);
-            int resized_width  = params.width + (64 - params.width % 64);
+        if (params.height != height || params.width != width) {
+            printf("resize input image from %dx%d to %dx%d\n", width, height, params.width, params.height);
+            int resized_height = params.height;
+            int resized_width  = params.width;

            uint8_t* resized_image_buffer = (uint8_t*)malloc(resized_height * resized_width * 3);
            if (resized_image_buffer == NULL) {
@@ -694,7 +700,7 @@ int main(int argc, const char* argv[]) {
                free(input_image_buffer);
                return 1;
            }
-            stbir_resize(input_image_buffer, params.width, params.height, 0,
+            stbir_resize(input_image_buffer, width, height, 0,
                         resized_image_buffer, resized_width, resized_height, 0, STBIR_TYPE_UINT8,
                         3 /*RGB channel*/, STBIR_ALPHA_CHANNEL_NONE, 0,
                         STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
@@ -704,8 +710,6 @@ int main(int argc, const char* argv[]) {
            // Save resized result
            free(input_image_buffer);
            input_image_buffer = resized_image_buffer;
-            params.height      = resized_height;
-            params.width       = resized_width;
        }
    }

@@ -732,31 +736,32 @@ int main(int argc, const char* argv[]) {
        return 1;
    }

+    sd_image_t* control_image = NULL;
+    if (params.controlnet_path.size() > 0 && params.control_image_path.size() > 0) {
+        int c                = 0;
+        control_image_buffer = stbi_load(params.control_image_path.c_str(), &params.width, &params.height, &c, 3);
+        if (control_image_buffer == NULL) {
+            fprintf(stderr, "load image from '%s' failed\n", params.control_image_path.c_str());
+            return 1;
+        }
+        control_image = new sd_image_t{(uint32_t)params.width,
+                                       (uint32_t)params.height,
+                                       3,
+                                       control_image_buffer};
+        if (params.canny_preprocess) {  // apply preprocessor
+            control_image->data = preprocess_canny(control_image->data,
+                                                   control_image->width,
+                                                   control_image->height,
+                                                   0.08f,
+                                                   0.08f,
+                                                   0.8f,
+                                                   1.0f,
+                                                   false);
+        }
+    }
+
    sd_image_t* results;
    if (params.mode == TXT2IMG) {
-        sd_image_t* control_image = NULL;
-        if (params.controlnet_path.size() > 0 && params.control_image_path.size() > 0) {
-            int c              = 0;
-            input_image_buffer = stbi_load(params.control_image_path.c_str(), &params.width, &params.height, &c, 3);
-            if (input_image_buffer == NULL) {
-                fprintf(stderr, "load image from '%s' failed\n", params.control_image_path.c_str());
-                return 1;
-            }
-            control_image = new sd_image_t{(uint32_t)params.width,
-                                           (uint32_t)params.height,
-                                           3,
-                                           input_image_buffer};
-            if (params.canny_preprocess) {  // apply preprocessor
-                control_image->data = preprocess_canny(control_image->data,
-                                                       control_image->width,
-                                                       control_image->height,
-                                                       0.08f,
-                                                       0.08f,
-                                                       0.8f,
-                                                       1.0f,
-                                                       false);
-            }
-        }
        results = txt2img(sd_ctx,
                          params.prompt.c_str(),
                          params.negative_prompt.c_str(),
@@ -828,7 +833,12 @@ int main(int argc, const char* argv[]) {
                              params.sample_steps,
                              params.strength,
                              params.seed,
-                              params.batch_count);
+                              params.batch_count,
+                              control_image,
+                              params.control_strength,
+                              params.style_ratio,
+                              params.normalize_input,
+                              params.input_id_images_path.c_str());
        }
    }

@@ -881,6 +891,8 @@ int main(int argc, const char* argv[]) {
    }
    free(results);
    free_sd_ctx(sd_ctx);
+    free(control_image_buffer);
+    free(input_image_buffer);

    return 0;
 }
@@ -454,7 +454,7 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const
    ggml_tensor* input_tile  = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size, tile_size, input->ne[2], 1);
    ggml_tensor* output_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size * scale, tile_size * scale, output->ne[2], 1);
    on_processing(input_tile, NULL, true);
-    int num_tiles = (input_width * input_height) / (non_tile_overlap * non_tile_overlap);
+    int num_tiles = ceil((float)input_width / non_tile_overlap) * ceil((float)input_height / non_tile_overlap);
    LOG_INFO("processing %i tiles", num_tiles);
    pretty_progress(1, num_tiles, 0.0f);
    int tile_count = 1;
@@ -752,10 +752,9 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_timestep_embedding(
    return ggml_timestep_embedding(ctx, timesteps, dim, max_period);
 }

-
-__STATIC_INLINE__ size_t ggml_tensor_num(ggml_context * ctx) {
+__STATIC_INLINE__ size_t ggml_tensor_num(ggml_context* ctx) {
    size_t num = 0;
-    for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+    for (ggml_tensor* t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
        num++;
    }
    return num;
@@ -851,7 +850,7 @@ protected:
    }

 public:
-    virtual std::string get_desc()       = 0;
+    virtual std::string get_desc() = 0;

    GGMLModule(ggml_backend_t backend, ggml_type wtype = GGML_TYPE_F32)
        : backend(backend), wtype(wtype) {
@@ -918,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);

@@ -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);
        }
    }
 }
@@ -888,6 +886,11 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
            }
        }

+        // ggml_n_dims returns 1 for scalars
+        if (n_dims == 0) {
+            n_dims = 1;
+        }
+
        TensorStorage tensor_storage(prefix + name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin);
        tensor_storage.reverse_ne();

@@ -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,
@@ -450,6 +450,11 @@ public:
                    LOG_INFO("running with Karras schedule");
                    denoiser->schedule = std::make_shared<KarrasSchedule>();
                    break;
+                case AYS:
+                    LOG_INFO("Running with Align-Your-Steps schedule");
+                    denoiser->schedule          = std::make_shared<AYSSchedule>();
+                    denoiser->schedule->version = version;
+                    break;
                case DEFAULT:
                    // Don't touch anything.
                    break;
@@ -852,7 +857,6 @@ public:
        copy_ggml_tensor(x, x_t);

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

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

@@ -873,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);
            }
@@ -979,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();
@@ -1536,60 +1158,35 @@ void free_sd_ctx(sd_ctx_t* sd_ctx) {
    free(sd_ctx);
 }

-sd_image_t* txt2img(sd_ctx_t* sd_ctx,
-                    const char* prompt_c_str,
-                    const char* negative_prompt_c_str,
-                    int clip_skip,
-                    float cfg_scale,
-                    int width,
-                    int height,
-                    enum sample_method_t sample_method,
-                    int sample_steps,
-                    int64_t seed,
-                    int batch_count,
-                    const sd_image_t* control_cond,
-                    float control_strength,
-                    float style_ratio,
-                    bool normalize_input,
-                    const char* input_id_images_path_c_str) {
-    LOG_DEBUG("txt2img %dx%d", width, height);
-    if (sd_ctx == NULL) {
-        return NULL;
-    }
-    // LOG_DEBUG("%s %s %f %d %d %d", prompt_c_str, negative_prompt_c_str, cfg_scale, sample_steps, seed, batch_count);
-    std::string prompt(prompt_c_str);
-    std::string negative_prompt(negative_prompt_c_str);
-    std::string input_id_images_path(input_id_images_path_c_str);
-
-    // preprocess input id images
-    std::vector<sd_image_t*> input_id_images;
-    if (sd_ctx->sd->pmid_model && input_id_images_path.size() > 0) {
-        std::vector<std::string> img_files = get_files_from_dir(input_id_images_path);
-        for (std::string img_file : img_files) {
-            int c = 0;
-            int width, height;
-            uint8_t* input_image_buffer = stbi_load(img_file.c_str(), &width, &height, &c, 3);
-            if (input_image_buffer == NULL) {
-                LOG_ERROR("PhotoMaker load image from '%s' failed", img_file.c_str());
-                continue;
-            } else {
-                LOG_INFO("PhotoMaker loaded image from '%s'", img_file.c_str());
-            }
-            sd_image_t* input_image = NULL;
-            input_image             = new sd_image_t{(uint32_t)width,
-                                         (uint32_t)height,
-                                         3,
-                                         input_image_buffer};
-            input_image             = preprocess_id_image(input_image);
-            if (input_image == NULL) {
-                LOG_ERROR("preprocess input id image from '%s' failed", img_file.c_str());
-                continue;
-            }
-            input_id_images.push_back(input_image);
-        }
+sd_image_t* generate_image(sd_ctx_t* sd_ctx,
+                           struct ggml_context* work_ctx,
+                           ggml_tensor* init_latent,
+                           std::string prompt,
+                           std::string negative_prompt,
+                           int clip_skip,
+                           float cfg_scale,
+                           int width,
+                           int height,
+                           enum sample_method_t sample_method,
+                           const std::vector<float>& sigmas,
+                           int64_t seed,
+                           int batch_count,
+                           const sd_image_t* control_cond,
+                           float control_strength,
+                           float style_ratio,
+                           bool normalize_input,
+                           std::string input_id_images_path) {
+    if (seed < 0) {
+        // Generally, when using the provided command line, the seed is always >0.
+        // However, to prevent potential issues if 'stable-diffusion.cpp' is invoked as a library
+        // by a third party with a seed <0, let's incorporate randomization here.
+        srand((int)time(NULL));
+        seed = rand();
    }

-    // extract and remove lora
+    int sample_steps = sigmas.size() - 1;
+
+    // Apply lora
    auto result_pair                                = extract_and_remove_lora(prompt);
    std::unordered_map<std::string, float> lora_f2m = result_pair.first;  // lora_name -> multiplier

@@ -1605,49 +1202,50 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
    int64_t t1 = ggml_time_ms();
    LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);

-    if (sd_ctx->sd->stacked_id && !sd_ctx->sd->pmid_lora->applied) {
-        t0 = ggml_time_ms();
-        sd_ctx->sd->pmid_lora->apply(sd_ctx->sd->tensors, sd_ctx->sd->n_threads);
-        t1                             = ggml_time_ms();
-        sd_ctx->sd->pmid_lora->applied = true;
-        LOG_INFO("pmid_lora apply completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
-        if (sd_ctx->sd->free_params_immediately) {
-            sd_ctx->sd->pmid_lora->free_params_buffer();
-        }
-    }
-
-    struct ggml_init_params params;
-    params.mem_size = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
-    if (sd_ctx->sd->stacked_id) {
-        params.mem_size += static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
-    }
-    params.mem_size += width * height * 3 * sizeof(float);
-    params.mem_size *= batch_count;
-    params.mem_buffer = NULL;
-    params.no_alloc   = false;
-    // LOG_DEBUG("mem_size %u ", params.mem_size);
-
-    struct ggml_context* work_ctx = ggml_init(params);
-    if (!work_ctx) {
-        LOG_ERROR("ggml_init() failed");
-        return NULL;
-    }
-
-    if (seed < 0) {
-        // Generally, when using the provided command line, the seed is always >0.
-        // However, to prevent potential issues if 'stable-diffusion.cpp' is invoked as a library
-        // by a third party with a seed <0, let's incorporate randomization here.
-        srand((int)time(NULL));
-        seed = rand();
-    }
-
+    // Photo Maker
    std::string prompt_text_only;
    ggml_tensor* init_img              = NULL;
    ggml_tensor* prompts_embeds        = NULL;
    ggml_tensor* pooled_prompts_embeds = NULL;
-    // ggml_tensor* class_tokens_mask = NULL;
    std::vector<bool> class_tokens_mask;
    if (sd_ctx->sd->stacked_id) {
+        if (!sd_ctx->sd->pmid_lora->applied) {
+            t0 = ggml_time_ms();
+            sd_ctx->sd->pmid_lora->apply(sd_ctx->sd->tensors, sd_ctx->sd->n_threads);
+            t1                             = ggml_time_ms();
+            sd_ctx->sd->pmid_lora->applied = true;
+            LOG_INFO("pmid_lora apply completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
+            if (sd_ctx->sd->free_params_immediately) {
+                sd_ctx->sd->pmid_lora->free_params_buffer();
+            }
+        }
+        // preprocess input id images
+        std::vector<sd_image_t*> input_id_images;
+        if (sd_ctx->sd->pmid_model && input_id_images_path.size() > 0) {
+            std::vector<std::string> img_files = get_files_from_dir(input_id_images_path);
+            for (std::string img_file : img_files) {
+                int c = 0;
+                int width, height;
+                uint8_t* input_image_buffer = stbi_load(img_file.c_str(), &width, &height, &c, 3);
+                if (input_image_buffer == NULL) {
+                    LOG_ERROR("PhotoMaker load image from '%s' failed", img_file.c_str());
+                    continue;
+                } else {
+                    LOG_INFO("PhotoMaker loaded image from '%s'", img_file.c_str());
+                }
+                sd_image_t* input_image = NULL;
+                input_image             = new sd_image_t{(uint32_t)width,
+                                             (uint32_t)height,
+                                             3,
+                                             input_image_buffer};
+                input_image             = preprocess_id_image(input_image);
+                if (input_image == NULL) {
+                    LOG_ERROR("preprocess input id image from '%s' failed", img_file.c_str());
+                    continue;
+                }
+                input_id_images.push_back(input_image);
+            }
+        }
        if (input_id_images.size() > 0) {
            sd_ctx->sd->pmid_model->style_strength = style_ratio;
            int32_t w                              = input_id_images[0]->width;
@@ -1682,21 +1280,22 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
            prompt_text_only = sd_ctx->sd->remove_trigger_from_prompt(work_ctx, prompt);
            // printf("%s || %s \n", prompt.c_str(), prompt_text_only.c_str());
            prompt = prompt_text_only;  //
-            if (sample_steps < 50) {
-                LOG_INFO("sampling steps increases from %d to 50 for PHOTOMAKER", sample_steps);
-                sample_steps = 50;
-            }
+            // if (sample_steps < 50) {
+            //     LOG_INFO("sampling steps increases from %d to 50 for PHOTOMAKER", sample_steps);
+            //     sample_steps = 50;
+            // }
        } else {
            LOG_WARN("Provided PhotoMaker model file, but NO input ID images");
            LOG_WARN("Turn off PhotoMaker");
            sd_ctx->sd->stacked_id = false;
        }
+        for (sd_image_t* img : input_id_images) {
+            free(img->data);
+        }
+        input_id_images.clear();
    }
-    for (sd_image_t* img : input_id_images) {
-        free(img->data);
-    }
-    input_id_images.clear();

+    // Get learned condition
    t0                    = ggml_time_ms();
    auto cond_pair        = sd_ctx->sd->get_learned_condition(work_ctx, prompt, clip_skip, width, height);
    ggml_tensor* c        = cond_pair.first;
@@ -1720,12 +1319,14 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
        sd_ctx->sd->cond_stage_model->free_params_buffer();
    }

+    // Control net hint
    struct ggml_tensor* image_hint = NULL;
    if (control_cond != NULL) {
        image_hint = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
        sd_image_to_tensor(control_cond->data, image_hint);
    }

+    // Sample
    std::vector<struct ggml_tensor*> final_latents;  // collect latents to decode
    int C = 4;
    int W = width / 8;
@@ -1734,25 +1335,31 @@ sd_image_t* txt2img(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 = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, W, H, C, 1);
-        ggml_tensor_set_f32_randn(x_t, sd_ctx->sd->rng);
-
-        std::vector<float> sigmas = sd_ctx->sd->denoiser->schedule->get_sigmas(sample_steps);
+        struct ggml_tensor* x_t   = NULL;
+        struct ggml_tensor* noise = NULL;
+        if (init_latent == NULL) {
+            x_t = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, W, H, C, 1);
+            ggml_tensor_set_f32_randn(x_t, sd_ctx->sd->rng);
+        } else {
+            x_t   = init_latent;
+            noise = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, W, H, C, 1);
+            ggml_tensor_set_f32_randn(noise, sd_ctx->sd->rng);
+        }

        int start_merge_step = -1;
        if (sd_ctx->sd->stacked_id) {
            start_merge_step = int(sd_ctx->sd->pmid_model->style_strength / 100.f * sample_steps);
-            if (start_merge_step > 30)
-                start_merge_step = 30;
+            // if (start_merge_step > 30)
+            //     start_merge_step = 30;
            LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step);
        }

        struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx,
                                                     x_t,
-                                                     NULL,
+                                                     noise,
                                                     c,
                                                     NULL,
                                                     c_vector,
@@ -1781,6 +1388,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
    int64_t t3 = ggml_time_ms();
    LOG_INFO("generating %" PRId64 " latent images completed, taking %.2fs", final_latents.size(), (t3 - t1) * 1.0f / 1000);

+    // Decode to image
    LOG_INFO("decoding %zu latents", final_latents.size());
    std::vector<struct ggml_tensor*> decoded_images;  // collect decoded images
    for (size_t i = 0; i < final_latents.size(); i++) {
@@ -1812,9 +1420,74 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
        result_images[i].data    = sd_tensor_to_image(decoded_images[i]);
    }
    ggml_free(work_ctx);
-    LOG_INFO(
-        "txt2img completed in %.2fs",
-        (t4 - t0) * 1.0f / 1000);
+
+    return result_images;
+}
+
+sd_image_t* txt2img(sd_ctx_t* sd_ctx,
+                    const char* prompt_c_str,
+                    const char* negative_prompt_c_str,
+                    int clip_skip,
+                    float cfg_scale,
+                    int width,
+                    int height,
+                    enum sample_method_t sample_method,
+                    int sample_steps,
+                    int64_t seed,
+                    int batch_count,
+                    const sd_image_t* control_cond,
+                    float control_strength,
+                    float style_ratio,
+                    bool normalize_input,
+                    const char* input_id_images_path_c_str) {
+    LOG_DEBUG("txt2img %dx%d", width, height);
+    if (sd_ctx == NULL) {
+        return NULL;
+    }
+
+    struct ggml_init_params params;
+    params.mem_size = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
+    if (sd_ctx->sd->stacked_id) {
+        params.mem_size += static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
+    }
+    params.mem_size += width * height * 3 * sizeof(float);
+    params.mem_size *= batch_count;
+    params.mem_buffer = NULL;
+    params.no_alloc   = false;
+    // LOG_DEBUG("mem_size %u ", params.mem_size);
+
+    struct ggml_context* work_ctx = ggml_init(params);
+    if (!work_ctx) {
+        LOG_ERROR("ggml_init() failed");
+        return NULL;
+    }
+
+    size_t t0 = ggml_time_ms();
+
+    std::vector<float> sigmas = sd_ctx->sd->denoiser->schedule->get_sigmas(sample_steps);
+
+    sd_image_t* result_images = generate_image(sd_ctx,
+                                               work_ctx,
+                                               NULL,
+                                               prompt_c_str,
+                                               negative_prompt_c_str,
+                                               clip_skip,
+                                               cfg_scale,
+                                               width,
+                                               height,
+                                               sample_method,
+                                               sigmas,
+                                               seed,
+                                               batch_count,
+                                               control_cond,
+                                               control_strength,
+                                               style_ratio,
+                                               normalize_input,
+                                               input_id_images_path_c_str);
+
+    size_t t1 = ggml_time_ms();
+
+    LOG_INFO("txt2img completed in %.2fs", (t1 - t0) * 1.0f / 1000);

    return result_images;
 }
@@ -1831,59 +1504,44 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
                    int sample_steps,
                    float strength,
                    int64_t seed,
-                    int batch_count) {
+                    int batch_count,
+                    const sd_image_t* control_cond,
+                    float control_strength,
+                    float style_ratio,
+                    bool normalize_input,
+                    const char* input_id_images_path_c_str) {
+    LOG_DEBUG("img2img %dx%d", width, height);
    if (sd_ctx == NULL) {
        return NULL;
    }
-    std::string prompt(prompt_c_str);
-    std::string negative_prompt(negative_prompt_c_str);
-
-    LOG_INFO("img2img %dx%d", width, height);
-
-    std::vector<float> sigmas = sd_ctx->sd->denoiser->schedule->get_sigmas(sample_steps);
-    size_t t_enc              = static_cast<size_t>(sample_steps * strength);
-    LOG_INFO("target t_enc is %zu steps", t_enc);
-    std::vector<float> sigma_sched;
-    sigma_sched.assign(sigmas.begin() + sample_steps - t_enc - 1, sigmas.end());

    struct ggml_init_params params;
-    params.mem_size = static_cast<size_t>(10 * 1024) * 1024;  // 10 MB
+    params.mem_size = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
+    if (sd_ctx->sd->stacked_id) {
+        params.mem_size += static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
+    }
    params.mem_size += width * height * 3 * sizeof(float) * 2;
+    params.mem_size *= batch_count;
    params.mem_buffer = NULL;
    params.no_alloc   = false;
    // LOG_DEBUG("mem_size %u ", params.mem_size);

-    // draft context
    struct ggml_context* work_ctx = ggml_init(params);
    if (!work_ctx) {
        LOG_ERROR("ggml_init() failed");
        return NULL;
    }

+    size_t t0 = ggml_time_ms();
+
    if (seed < 0) {
-        seed = (int)time(NULL);
+        srand((int)time(NULL));
+        seed = rand();
    }
-
    sd_ctx->sd->rng->manual_seed(seed);

-    // extract and remove lora
-    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);
-    }
-    prompt = result_pair.second;
-    LOG_DEBUG("prompt after extract and remove lora: \"%s\"", prompt.c_str());
-
-    // load lora from file
-    int64_t t0 = ggml_time_ms();
-    sd_ctx->sd->apply_loras(lora_f2m);
-    int64_t t1 = ggml_time_ms();
-    LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
-
    ggml_tensor* init_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
    sd_image_to_tensor(init_image.data, init_img);
-    t0                       = ggml_time_ms();
    ggml_tensor* init_latent = NULL;
    if (!sd_ctx->sd->use_tiny_autoencoder) {
        ggml_tensor* moments = sd_ctx->sd->encode_first_stage(work_ctx, init_img);
@@ -1892,87 +1550,37 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
        init_latent = sd_ctx->sd->encode_first_stage(work_ctx, init_img);
    }
    // print_ggml_tensor(init_latent);
-    t1 = ggml_time_ms();
+    size_t t1 = ggml_time_ms();
    LOG_INFO("encode_first_stage completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);

-    auto cond_pair                = sd_ctx->sd->get_learned_condition(work_ctx, prompt, clip_skip, width, height);
-    ggml_tensor* c                = cond_pair.first;
-    ggml_tensor* c_vector         = cond_pair.second;  // [adm_in_channels, ]
-    struct ggml_tensor* uc        = NULL;
-    struct ggml_tensor* uc_vector = NULL;
-    if (cfg_scale != 1.0) {
-        bool force_zero_embeddings = false;
-        if (sd_ctx->sd->version == VERSION_XL && negative_prompt.size() == 0) {
-            force_zero_embeddings = true;
-        }
-        auto uncond_pair = sd_ctx->sd->get_learned_condition(work_ctx, negative_prompt, clip_skip, width, height, force_zero_embeddings);
-        uc               = uncond_pair.first;
-        uc_vector        = uncond_pair.second;  // [adm_in_channels, ]
-    }
-    int64_t t2 = ggml_time_ms();
-    LOG_INFO("get_learned_condition completed, taking %" PRId64 " ms", t2 - t1);
-    if (sd_ctx->sd->free_params_immediately) {
-        sd_ctx->sd->cond_stage_model->free_params_buffer();
-    }
+    std::vector<float> sigmas = sd_ctx->sd->denoiser->schedule->get_sigmas(sample_steps);
+    size_t t_enc              = static_cast<size_t>(sample_steps * strength);
+    LOG_INFO("target t_enc is %zu steps", t_enc);
+    std::vector<float> sigma_sched;
+    sigma_sched.assign(sigmas.begin() + sample_steps - t_enc - 1, sigmas.end());

-    sd_ctx->sd->rng->manual_seed(seed);
-    struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, init_latent);
-    ggml_tensor_set_f32_randn(noise, sd_ctx->sd->rng);
+    sd_image_t* result_images = generate_image(sd_ctx,
+                                               work_ctx,
+                                               init_latent,
+                                               prompt_c_str,
+                                               negative_prompt_c_str,
+                                               clip_skip,
+                                               cfg_scale,
+                                               width,
+                                               height,
+                                               sample_method,
+                                               sigma_sched,
+                                               seed,
+                                               batch_count,
+                                               control_cond,
+                                               control_strength,
+                                               style_ratio,
+                                               normalize_input,
+                                               input_id_images_path_c_str);

-    LOG_INFO("sampling using %s method", sampling_methods_str[sample_method]);
-    struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx,
-                                                 init_latent,
-                                                 noise,
-                                                 c,
-                                                 NULL,
-                                                 c_vector,
-                                                 uc,
-                                                 NULL,
-                                                 uc_vector,
-                                                 {},
-                                                 0.f,
-                                                 cfg_scale,
-                                                 cfg_scale,
-                                                 sample_method,
-                                                 sigma_sched,
-                                                 -1,
-                                                 NULL,
-                                                 NULL);
-    // struct ggml_tensor *x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
-    // print_ggml_tensor(x_0);
-    int64_t t3 = ggml_time_ms();
-    LOG_INFO("sampling completed, taking %.2fs", (t3 - t2) * 1.0f / 1000);
-    if (sd_ctx->sd->free_params_immediately) {
-        sd_ctx->sd->diffusion_model->free_params_buffer();
-    }
+    size_t t2 = ggml_time_ms();

-    struct ggml_tensor* img = sd_ctx->sd->decode_first_stage(work_ctx, x_0);
-    if (sd_ctx->sd->free_params_immediately && !sd_ctx->sd->use_tiny_autoencoder) {
-        sd_ctx->sd->first_stage_model->free_params_buffer();
-    }
-    if (img == NULL) {
-        ggml_free(work_ctx);
-        return NULL;
-    }
-
-    sd_image_t* result_images = (sd_image_t*)calloc(1, sizeof(sd_image_t));
-    if (result_images == NULL) {
-        ggml_free(work_ctx);
-        return NULL;
-    }
-
-    for (size_t i = 0; i < 1; i++) {
-        result_images[i].width   = width;
-        result_images[i].height  = height;
-        result_images[i].channel = 3;
-        result_images[i].data    = sd_tensor_to_image(img);
-    }
-    ggml_free(work_ctx);
-
-    int64_t t4 = ggml_time_ms();
-    LOG_INFO("decode_first_stage completed, taking %.2fs", (t4 - t3) * 1.0f / 1000);
-
-    LOG_INFO("img2img completed in %.2fs", (t4 - t0) * 1.0f / 1000);
+    LOG_INFO("img2img completed in %.2fs", (t1 - t0) * 1.0f / 1000);

    return result_images;
 }
@@ -49,6 +49,7 @@ enum schedule_t {
    DEFAULT,
    DISCRETE,
    KARRAS,
+    AYS,
    N_SCHEDULES
 };

@@ -59,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,
@@ -79,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,
 };

@@ -160,7 +164,12 @@ SD_API sd_image_t* img2img(sd_ctx_t* sd_ctx,
                           int sample_steps,
                           float strength,
                           int64_t seed,
-                           int batch_count);
+                           int batch_count,
+                           const sd_image_t* control_cond,
+                           float control_strength,
+                           float style_strength,
+                           bool normalize_input,
+                           const char* input_id_images_path);

 SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
                           sd_image_t init_image,
@@ -201,7 +201,7 @@ struct TinyAutoEncoder : public GGMLModule {
    }

    bool load_from_file(const std::string& file_path) {
-        LOG_INFO("loading taesd from '%s'", file_path.c_str());
+        LOG_INFO("loading taesd from '%s', decode_only = %s", file_path.c_str(), decode_only ? "true" : "false");
        alloc_params_buffer();
        std::map<std::string, ggml_tensor*> taesd_tensors;
        taesd.get_param_tensors(taesd_tensors);
@@ -0,0 +1,2 @@
+DisableFormat: true
+SortIncludes: Never
@@ -36,6 +36,7 @@
 #include <unistd.h>
 #endif

+#define USE_EXTERNAL_MZCRC
 #include "miniz.h"
 #include "zip.h"

@@ -1834,3 +1835,234 @@ int zip_extract(const char *zipname, const char *dir,

  return zip_archive_extract(&zip_archive, dir, on_extract, arg);
 }
+
+#if defined(__SSE4_2__) || defined(__AVX512F__)
+#include <immintrin.h>
+#endif
+
+// Phil Katz 32-Bit Cyclic Redundancy Check Uber Alles
+// Goes 73 GiB/s on an AMD Ryzen Threadripper PRO 7995WX
+// "Fast CRC Computation for Generic Polynomials Using PCLMULQDQ Instruction"
+//  V. Gopal, E. Ozturk, et al., 2009, http://intel.ly/2ySEwL0
+mz_ulong mz_crc32(mz_ulong init, const uint8_t *buf, size_t len) {
+  uint32_t crc = ~init;
+#if defined(__AVX512F__) && defined(__VPCLMULQDQ__) && defined(__PCLMUL__)
+  if (len >= 256) {
+    _Alignas(__m512) static const uint64_t k1k2[] = {
+        0x011542778a, 0x01322d1430, 0x011542778a, 0x01322d1430,
+        0x011542778a, 0x01322d1430, 0x011542778a, 0x01322d1430,
+    };
+    _Alignas(__m512) static const uint64_t k3k4[] = {
+        0x0154442bd4, 0x01c6e41596, 0x0154442bd4, 0x01c6e41596,
+        0x0154442bd4, 0x01c6e41596, 0x0154442bd4, 0x01c6e41596,
+    };
+    _Alignas(__m512) static const uint64_t k5k6[] = {
+        0x01751997d0,
+        0x00ccaa009e,
+    };
+    _Alignas(__m512) static const uint64_t k7k8[] = {
+        0x0163cd6124,
+        0x0000000000,
+    };
+    _Alignas(__m512) static const uint64_t poly[] = {
+        0x01db710641,
+        0x01f7011641,
+    };
+    __m512i x0, x1, x2, x3, x4, x5, x6, x7, x8, y5, y6, y7, y8;
+    __m128i a0, a1, a2, a3;
+    x1 = _mm512_loadu_si512((__m512i *)(buf + 0x00));
+    x2 = _mm512_loadu_si512((__m512i *)(buf + 0x40));
+    x3 = _mm512_loadu_si512((__m512i *)(buf + 0x80));
+    x4 = _mm512_loadu_si512((__m512i *)(buf + 0xC0));
+    x1 = _mm512_xor_si512(x1, _mm512_castsi128_si512(_mm_cvtsi32_si128(crc)));
+    x0 = _mm512_load_si512((__m512i *)k1k2);
+    buf += 256;
+    len -= 256;
+    while (len >= 256) {
+      x5 = _mm512_clmulepi64_epi128(x1, x0, 0x00);
+      x6 = _mm512_clmulepi64_epi128(x2, x0, 0x00);
+      x7 = _mm512_clmulepi64_epi128(x3, x0, 0x00);
+      x8 = _mm512_clmulepi64_epi128(x4, x0, 0x00);
+      x1 = _mm512_clmulepi64_epi128(x1, x0, 0x11);
+      x2 = _mm512_clmulepi64_epi128(x2, x0, 0x11);
+      x3 = _mm512_clmulepi64_epi128(x3, x0, 0x11);
+      x4 = _mm512_clmulepi64_epi128(x4, x0, 0x11);
+      y5 = _mm512_loadu_si512((__m512i *)(buf + 0x00));
+      y6 = _mm512_loadu_si512((__m512i *)(buf + 0x40));
+      y7 = _mm512_loadu_si512((__m512i *)(buf + 0x80));
+      y8 = _mm512_loadu_si512((__m512i *)(buf + 0xC0));
+      x1 = _mm512_xor_si512(x1, x5);
+      x2 = _mm512_xor_si512(x2, x6);
+      x3 = _mm512_xor_si512(x3, x7);
+      x4 = _mm512_xor_si512(x4, x8);
+      x1 = _mm512_xor_si512(x1, y5);
+      x2 = _mm512_xor_si512(x2, y6);
+      x3 = _mm512_xor_si512(x3, y7);
+      x4 = _mm512_xor_si512(x4, y8);
+      buf += 256;
+      len -= 256;
+    }
+    x0 = _mm512_load_si512((__m512i *)k3k4);
+    x5 = _mm512_clmulepi64_epi128(x1, x0, 0x00);
+    x1 = _mm512_clmulepi64_epi128(x1, x0, 0x11);
+    x1 = _mm512_xor_si512(x1, x2);
+    x1 = _mm512_xor_si512(x1, x5);
+    x5 = _mm512_clmulepi64_epi128(x1, x0, 0x00);
+    x1 = _mm512_clmulepi64_epi128(x1, x0, 0x11);
+    x1 = _mm512_xor_si512(x1, x3);
+    x1 = _mm512_xor_si512(x1, x5);
+    x5 = _mm512_clmulepi64_epi128(x1, x0, 0x00);
+    x1 = _mm512_clmulepi64_epi128(x1, x0, 0x11);
+    x1 = _mm512_xor_si512(x1, x4);
+    x1 = _mm512_xor_si512(x1, x5);
+    while (len >= 64) {
+      x2 = _mm512_loadu_si512((__m512i *)buf);
+      x5 = _mm512_clmulepi64_epi128(x1, x0, 0x00);
+      x1 = _mm512_clmulepi64_epi128(x1, x0, 0x11);
+      x1 = _mm512_xor_si512(x1, x2);
+      x1 = _mm512_xor_si512(x1, x5);
+      buf += 64;
+      len -= 64;
+    }
+    a0 = _mm_load_si128((__m128i *)k5k6);
+    a1 = _mm512_extracti32x4_epi32(x1, 0);
+    a2 = _mm512_extracti32x4_epi32(x1, 1);
+    a3 = _mm_clmulepi64_si128(a1, a0, 0x00);
+    a1 = _mm_clmulepi64_si128(a1, a0, 0x11);
+    a1 = _mm_xor_si128(a1, a3);
+    a1 = _mm_xor_si128(a1, a2);
+    a2 = _mm512_extracti32x4_epi32(x1, 2);
+    a3 = _mm_clmulepi64_si128(a1, a0, 0x00);
+    a1 = _mm_clmulepi64_si128(a1, a0, 0x11);
+    a1 = _mm_xor_si128(a1, a3);
+    a1 = _mm_xor_si128(a1, a2);
+    a2 = _mm512_extracti32x4_epi32(x1, 3);
+    a3 = _mm_clmulepi64_si128(a1, a0, 0x00);
+    a1 = _mm_clmulepi64_si128(a1, a0, 0x11);
+    a1 = _mm_xor_si128(a1, a3);
+    a1 = _mm_xor_si128(a1, a2);
+    a2 = _mm_clmulepi64_si128(a1, a0, 0x10);
+    a3 = _mm_setr_epi32(~0, 0, ~0, 0);
+    a1 = _mm_srli_si128(a1, 8);
+    a1 = _mm_xor_si128(a1, a2);
+    a0 = _mm_loadl_epi64((__m128i *)k7k8);
+    a2 = _mm_srli_si128(a1, 4);
+    a1 = _mm_and_si128(a1, a3);
+    a1 = _mm_clmulepi64_si128(a1, a0, 0x00);
+    a1 = _mm_xor_si128(a1, a2);
+    a0 = _mm_load_si128((__m128i *)poly);
+    a2 = _mm_and_si128(a1, a3);
+    a2 = _mm_clmulepi64_si128(a2, a0, 0x10);
+    a2 = _mm_and_si128(a2, a3);
+    a2 = _mm_clmulepi64_si128(a2, a0, 0x00);
+    a1 = _mm_xor_si128(a1, a2);
+    crc = _mm_extract_epi32(a1, 1);
+  }
+#endif
+#if defined(__SSE4_2__) && defined(__PCLMUL__)
+  if (len >= 64) {
+    _Alignas(__m128) static const uint64_t k1k2[] = {
+        0x0154442bd4,
+        0x01c6e41596,
+    };
+    _Alignas(__m128) static const uint64_t k3k4[] = {
+        0x01751997d0,
+        0x00ccaa009e,
+    };
+    _Alignas(__m128) static const uint64_t k5k0[] = {
+        0x0163cd6124,
+        0x0000000000,
+    };
+    _Alignas(__m128) static const uint64_t poly[] = {
+        0x01db710641,
+        0x01f7011641,
+    };
+    __m128i x0, x1, x2, x3, x4, x5, x6, x7, x8, y5, y6, y7, y8;
+    x1 = _mm_loadu_si128((__m128i *)(buf + 0x00));
+    x2 = _mm_loadu_si128((__m128i *)(buf + 0x10));
+    x3 = _mm_loadu_si128((__m128i *)(buf + 0x20));
+    x4 = _mm_loadu_si128((__m128i *)(buf + 0x30));
+    x1 = _mm_xor_si128(x1, _mm_cvtsi32_si128(crc));
+    x0 = _mm_load_si128((__m128i *)k1k2);
+    buf += 64;
+    len -= 64;
+    while (len >= 64) {
+      x5 = _mm_clmulepi64_si128(x1, x0, 0x00);
+      x6 = _mm_clmulepi64_si128(x2, x0, 0x00);
+      x7 = _mm_clmulepi64_si128(x3, x0, 0x00);
+      x8 = _mm_clmulepi64_si128(x4, x0, 0x00);
+      x1 = _mm_clmulepi64_si128(x1, x0, 0x11);
+      x2 = _mm_clmulepi64_si128(x2, x0, 0x11);
+      x3 = _mm_clmulepi64_si128(x3, x0, 0x11);
+      x4 = _mm_clmulepi64_si128(x4, x0, 0x11);
+      y5 = _mm_loadu_si128((__m128i *)(buf + 0x00));
+      y6 = _mm_loadu_si128((__m128i *)(buf + 0x10));
+      y7 = _mm_loadu_si128((__m128i *)(buf + 0x20));
+      y8 = _mm_loadu_si128((__m128i *)(buf + 0x30));
+      x1 = _mm_xor_si128(x1, x5);
+      x2 = _mm_xor_si128(x2, x6);
+      x3 = _mm_xor_si128(x3, x7);
+      x4 = _mm_xor_si128(x4, x8);
+      x1 = _mm_xor_si128(x1, y5);
+      x2 = _mm_xor_si128(x2, y6);
+      x3 = _mm_xor_si128(x3, y7);
+      x4 = _mm_xor_si128(x4, y8);
+      buf += 64;
+      len -= 64;
+    }
+    x0 = _mm_load_si128((__m128i *)k3k4);
+    x5 = _mm_clmulepi64_si128(x1, x0, 0x00);
+    x1 = _mm_clmulepi64_si128(x1, x0, 0x11);
+    x1 = _mm_xor_si128(x1, x2);
+    x1 = _mm_xor_si128(x1, x5);
+    x5 = _mm_clmulepi64_si128(x1, x0, 0x00);
+    x1 = _mm_clmulepi64_si128(x1, x0, 0x11);
+    x1 = _mm_xor_si128(x1, x3);
+    x1 = _mm_xor_si128(x1, x5);
+    x5 = _mm_clmulepi64_si128(x1, x0, 0x00);
+    x1 = _mm_clmulepi64_si128(x1, x0, 0x11);
+    x1 = _mm_xor_si128(x1, x4);
+    x1 = _mm_xor_si128(x1, x5);
+    while (len >= 16) {
+      x2 = _mm_loadu_si128((__m128i *)buf);
+      x5 = _mm_clmulepi64_si128(x1, x0, 0x00);
+      x1 = _mm_clmulepi64_si128(x1, x0, 0x11);
+      x1 = _mm_xor_si128(x1, x2);
+      x1 = _mm_xor_si128(x1, x5);
+      buf += 16;
+      len -= 16;
+    }
+    x2 = _mm_clmulepi64_si128(x1, x0, 0x10);
+    x3 = _mm_setr_epi32(~0, 0, ~0, 0);
+    x1 = _mm_srli_si128(x1, 8);
+    x1 = _mm_xor_si128(x1, x2);
+    x0 = _mm_loadl_epi64((__m128i *)k5k0);
+    x2 = _mm_srli_si128(x1, 4);
+    x1 = _mm_and_si128(x1, x3);
+    x1 = _mm_clmulepi64_si128(x1, x0, 0x00);
+    x1 = _mm_xor_si128(x1, x2);
+    x0 = _mm_load_si128((__m128i *)poly);
+    x2 = _mm_and_si128(x1, x3);
+    x2 = _mm_clmulepi64_si128(x2, x0, 0x10);
+    x2 = _mm_and_si128(x2, x3);
+    x2 = _mm_clmulepi64_si128(x2, x0, 0x00);
+    x1 = _mm_xor_si128(x1, x2);
+    crc = _mm_extract_epi32(x1, 1);
+  }
+#endif
+  static uint32_t tab[256];
+  if (!tab[255]) {
+    // generates table for byte-wise crc calculation on the polynomial
+    // x^32+x^26+x^23+x^22+x^16+x^12+x^11+x^10+x^8+x^7+x^5+x^4+x^2+x+1
+    uint32_t polynomial = 0xedb88320;  // bits are reversed
+    for (int d = 0; d < 256; ++d) {
+      uint32_t r = d;
+      for (int i = 0; i < 8; ++i)
+        r = r >> 1 ^ (r & 1 ? polynomial : 0);
+      tab[d] = r;
+    }
+  }
+  for (size_t i = 0; i < len; ++i)
+    crc = crc >> 8 ^ tab[(crc & 255) ^ buf[i]];
+  return ~crc & 0xffffffff;
+}
@@ -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";
Author	SHA1	Message	Date
Phu Tran	9c51d8787f	chore: fix cuda CI (#286 )	2024-06-12 23:13:24 +08:00
leejet	f9f0d4685b	fix: sample_k_diffusion should be static	2024-06-10 23:04:02 +08:00
leejet	8d2050a5cf	sync: update ggml	2024-06-10 22:59:36 +08:00
leejet	08f5b41956	refector: make the sampling module more independent	2024-06-10 22:42:15 +08:00
Eugene	b6daf5c55b	fix: use PRI64 instead of %i for some log (#269 )	2024-06-01 14:01:58 +08:00
leejet	be6cd1a4bf	sync: update ggml	2024-06-01 13:44:09 +08:00
Justine Tunney	e1384defca	perf: make crc32 100x faster on x86-64 (#278 ) This change makes checkpoints load significantly faster by optimizing pkzip's cyclic redundancy check. This code was developed by Intel and Google and Mozilla. See Chromium's zlib codebase for further details.	2024-06-01 12:58:30 +08:00
Phu Tran	814280343c	chore: update artifact actions (#267 )	2024-06-01 12:33:13 +08:00
leejet	1d2af5ca3f	fix: set n_dims of tensor storage to 1 when it's 0	2024-05-14 23:06:52 +08:00
Grauhoandleejet	ce1bcc74a6	feat: add AYS(Align Your Steps) scheduler (#241 ) Added NVIDEA's new "Align Your Steps" style scheduler in accordance with their quick start guide. Currently has handling for SD1.5, SDXL, and SVD, using the noise levels from their paper to generate the sigma values. Can be selected using the --schedule ays command line switch. Updates the main.cpp help message and README to reflect this option, also they now inform the user of the --color switch as well. --------- Co-authored-by: leejet <leejet714@gmail.com>	2024-04-29 23:21:32 +08:00
Eugene	760cfaa618	fix: ignore tensors with the particular dim while loading (#233 )	2024-04-29 23:04:27 +08:00
Eugene	6d16f6853e	fix: correct upscale progressbar (#232 )	2024-04-29 22:59:46 +08:00
leejet	036ba9e6d8	feat: enable controlnet and photo maker for img2img mode	2024-04-14 16:36:08 +08:00