feat: set default rng to cuda

examples/main.cpp

@@ -86,7 +86,7 @@ struct Option {
     SampleMethod sample_method = EULAR_A;
     int sample_steps = 20;
     float strength = 0.75f;
-    RNGType rng_type = STD_DEFAULT_RNG;
+    RNGType rng_type = CUDA_RNG;
     int64_t seed = 42;
     bool verbose = false;
 
@@ -130,7 +130,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  -W, --width W                      image width, in pixel space (default: 512)\n");
     printf("  --sample-method SAMPLE_METHOD      sample method (default: \"eular a\")\n");
     printf("  --steps  STEPS                     number of sample steps (default: 20)\n");
-    printf("  --rng {std_default, cuda}          RNG (default: std_default)\n");
+    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("  -v, --verbose                      print extra info\n");
 }
@@ -285,13 +285,13 @@ void parse_args(int argc, const char* argv[], Option* opt) {
         exit(1);
     }
 
-    if (opt->w <= 0 || opt->w % 32 != 0) {
-        fprintf(stderr, "error: the width must be a multiple of 32\n");
+    if (opt->w <= 0 || opt->w % 64 != 0) {
+        fprintf(stderr, "error: the width must be a multiple of 64\n");
         exit(1);
     }
 
-    if (opt->h <= 0 || opt->h % 32 != 0) {
-        fprintf(stderr, "error: the height must be a multiple of 32\n");
+    if (opt->h <= 0 || opt->h % 64 != 0) {
+        fprintf(stderr, "error: the height must be a multiple of 64\n");
         exit(1);
     }
 
@@ -337,13 +337,13 @@ int main(int argc, const char* argv[]) {
             free(img_data);
             return 1;
         }
-        if (opt.w <= 0 || opt.w % 32 != 0) {
-            fprintf(stderr, "error: the width of image must be a multiple of 32\n");
+        if (opt.w <= 0 || opt.w % 64 != 0) {
+            fprintf(stderr, "error: the width of image must be a multiple of 64\n");
             free(img_data);
             return 1;
         }
-        if (opt.h <= 0 || opt.h % 32 != 0) {
-            fprintf(stderr, "error: the height of image must be a multiple of 32\n");
+        if (opt.h <= 0 || opt.h % 64 != 0) {
+            fprintf(stderr, "error: the height of image must be a multiple of 64\n");
             free(img_data);
             return 1;
         }

stable-diffusion.cpp

@@ -14,9 +14,9 @@
 #include <vector>
 
 #include "ggml/ggml.h"
-#include "stable-diffusion.h"
 #include "rng.h"
 #include "rng_philox.h"
+#include "stable-diffusion.h"
 
 static SDLogLevel log_level = SDLogLevel::INFO;
 
@@ -2659,10 +2659,19 @@ struct DiscreteSchedule {
     float sigmas[TIMESTEPS];
     float log_sigmas[TIMESTEPS];
 
-    std::vector<float> get_sigmas(int n) {
+    std::vector<float> get_sigmas(uint32_t n) {
         std::vector<float> result;
 
         int t_max = TIMESTEPS - 1;
+
+        if (n == 0) {
+            return result;
+        } else if (n == 1) {
+            result.push_back(t_to_sigma(t_max));
+            result.push_back(0);
+            return result;
+        }
+
         float step = static_cast<float>(t_max) / static_cast<float>(n - 1);
         for (int i = 0; i < n; ++i) {
             float t = t_max - step * i;
@@ -3099,7 +3108,7 @@ class StableDiffusionGGML {
         struct ggml_tensor* c = ggml_new_tensor_4d(res_ctx, GGML_TYPE_F32, 1024, 2, 1, 1);
         ggml_set_f32(c, 0.5);
 
-        size_t ctx_size = 1 * 1024 * 1024;  // 1MB
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
             struct ggml_init_params params;
@@ -3122,8 +3131,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* out = diffusion_model.forward(ctx, x_t, NULL, c, t_emb);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-            struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+            struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+            struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("diffusion context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3155,8 +3164,8 @@ class StableDiffusionGGML {
         struct ggml_tensor* out = diffusion_model.forward(ctx, x_t, NULL, c, t_emb);
         ggml_hold_dynamic_tensor(out);
 
-        struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-        struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+        struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+        struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
         ggml_set_dynamic(ctx, false);
         struct ggml_tensor* buf = ggml_new_tensor_1d(ctx, GGML_TYPE_I8, cplan.work_size);
@@ -3165,7 +3174,7 @@ class StableDiffusionGGML {
         cplan.work_data = (uint8_t*)buf->data;
 
         int64_t t0 = ggml_time_ms();
-        ggml_graph_compute(&diffusion_graph, &cplan);
+        ggml_graph_compute(diffusion_graph, &cplan);
 
         double result = 0.f;
 
@@ -3201,7 +3210,7 @@ class StableDiffusionGGML {
                                                             true);
         std::vector<int>& tokens = tokens_and_weights.first;
         std::vector<float>& weights = tokens_and_weights.second;
-        size_t ctx_size = 1 * 1024 * 1024;  // 1MB
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
             struct ggml_init_params params;
@@ -3251,14 +3260,14 @@ class StableDiffusionGGML {
         ggml_set_dynamic(ctx, params.dynamic);
 
         struct ggml_tensor* hidden_states = cond_stage_model.text_model.forward(ctx, input_ids);
-        struct ggml_cgraph cond_graph = ggml_build_forward(hidden_states);
+        struct ggml_cgraph* cond_graph = ggml_build_forward_ctx(ctx, hidden_states);
         LOG_DEBUG("building condition graph completed: %d nodes, %d leafs",
-                  cond_graph.n_nodes, cond_graph.n_leafs);
+                  cond_graph->n_nodes, cond_graph->n_leafs);
 
         memcpy(input_ids->data, tokens.data(), tokens.size() * ggml_element_size(input_ids));
 
         int64_t t0 = ggml_time_ms();
-        ggml_graph_compute_with_ctx(ctx, &cond_graph, n_threads);
+        ggml_graph_compute_with_ctx(ctx, cond_graph, n_threads);
         int64_t t1 = ggml_time_ms();
         LOG_DEBUG("computing condition graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
 
@@ -3335,7 +3344,7 @@ class StableDiffusionGGML {
         struct ggml_tensor* x = ggml_dup_tensor(res_ctx, x_t);
         copy_ggml_tensor(x, x_t);
 
-        size_t ctx_size = 1 * 1024 * 1024;  // 1MB
+        size_t ctx_size = 10 * 1024 * 1024;  // 10MB
         // calculate the amount of memory required
         {
             struct ggml_init_params params;
@@ -3360,8 +3369,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* out = diffusion_model.forward(ctx, noised_input, NULL, context, t_emb);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-            struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+            struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+            struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("diffusion context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3393,8 +3402,8 @@ class StableDiffusionGGML {
         struct ggml_tensor* out = diffusion_model.forward(ctx, noised_input, NULL, context, t_emb);
         ggml_hold_dynamic_tensor(out);
 
-        struct ggml_cgraph diffusion_graph = ggml_build_forward(out);
-        struct ggml_cplan cplan = ggml_graph_plan(&diffusion_graph, n_threads);
+        struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out);
+        struct ggml_cplan cplan = ggml_graph_plan(diffusion_graph, n_threads);
 
         ggml_set_dynamic(ctx, false);
         struct ggml_tensor* buf = ggml_new_tensor_1d(ctx, GGML_TYPE_I8, cplan.work_size);
@@ -3452,12 +3461,12 @@ class StableDiffusionGGML {
             if (cfg_scale != 1.0 && uc != NULL) {
                 // uncond
                 copy_ggml_tensor(context, uc);
-                ggml_graph_compute(&diffusion_graph, &cplan);
+                ggml_graph_compute(diffusion_graph, &cplan);
                 copy_ggml_tensor(out_uncond, out);
 
                 // cond
                 copy_ggml_tensor(context, c);
-                ggml_graph_compute(&diffusion_graph, &cplan);
+                ggml_graph_compute(diffusion_graph, &cplan);
 
                 out_cond = out;
 
@@ -3474,7 +3483,7 @@ class StableDiffusionGGML {
             } else {
                 // cond
                 copy_ggml_tensor(context, c);
-                ggml_graph_compute(&diffusion_graph, &cplan);
+                ggml_graph_compute(diffusion_graph, &cplan);
             }
 
             // v = out, eps = out
@@ -3590,7 +3599,7 @@ class StableDiffusionGGML {
         struct ggml_tensor* result = NULL;
 
         // calculate the amount of memory required
-        size_t ctx_size = 1 * 1024 * 1024;
+        size_t ctx_size = 10 * 1024 * 1024; // 10MB
         {
             struct ggml_init_params params;
             params.mem_size = ctx_size;
@@ -3607,8 +3616,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* moments = first_stage_model.encode(ctx, x);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph vae_graph = ggml_build_forward(moments);
-            struct ggml_cplan cplan = ggml_graph_plan(&vae_graph, n_threads);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments);
+            struct ggml_cplan cplan = ggml_graph_plan(vae_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("vae context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3632,10 +3641,10 @@ class StableDiffusionGGML {
             }
 
             struct ggml_tensor* moments = first_stage_model.encode(ctx, x);
-            struct ggml_cgraph vae_graph = ggml_build_forward(moments);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments);
 
             int64_t t0 = ggml_time_ms();
-            ggml_graph_compute_with_ctx(ctx, &vae_graph, n_threads);
+            ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads);
             int64_t t1 = ggml_time_ms();
 
 #ifdef GGML_PERF
@@ -3719,7 +3728,7 @@ class StableDiffusionGGML {
         }
 
         // calculate the amount of memory required
-        size_t ctx_size = 1 * 1024 * 1024;
+        size_t ctx_size = 10 * 1024 * 1024; // 10MB
         {
             struct ggml_init_params params;
             params.mem_size = ctx_size;
@@ -3736,8 +3745,8 @@ class StableDiffusionGGML {
             struct ggml_tensor* img = first_stage_model.decoder.forward(ctx, z);
             ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx);
 
-            struct ggml_cgraph vae_graph = ggml_build_forward(img);
-            struct ggml_cplan cplan = ggml_graph_plan(&vae_graph, n_threads);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img);
+            struct ggml_cplan cplan = ggml_graph_plan(vae_graph, n_threads);
 
             ctx_size += cplan.work_size;
             LOG_DEBUG("vae context need %.2fMB static memory, with work_size needing %.2fMB",
@@ -3761,10 +3770,10 @@ class StableDiffusionGGML {
             }
 
             struct ggml_tensor* img = first_stage_model.decode(ctx, z);
-            struct ggml_cgraph vae_graph = ggml_build_forward(img);
+            struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img);
 
             int64_t t0 = ggml_time_ms();
-            ggml_graph_compute_with_ctx(ctx, &vae_graph, n_threads);
+            ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads);
             int64_t t1 = ggml_time_ms();
 
 #ifdef GGML_PERF
@@ -3830,6 +3839,7 @@ std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
     std::vector<uint8_t> result;
     struct ggml_init_params params;
     params.mem_size = static_cast<size_t>(10 * 1024) * 1024;  // 10M
+    params.mem_size += width * height * 3 * sizeof(float) * 2;
     params.mem_buffer = NULL;
     params.no_alloc = false;
     params.dynamic = false;