feat: set default rng to cuda

examples/main.cpp

@@ -86,8 +86,8 @@ struct Option {
     SampleMethod sample_method = EULAR_A;
     int sample_steps = 20;
     float strength = 0.75f;
-    RNGType rng_type = STD_DEFAULT_RNG;
-    int seed = 42;
+    RNGType rng_type = CUDA_RNG;
+    int64_t seed = 42;
     bool verbose = false;
 
     void print() {
@@ -106,7 +106,7 @@ struct Option {
         printf("    sample_steps:    %d\n", sample_steps);
         printf("    strength:        %.2f\n", strength);
         printf("    rng:             %s\n", rng_type_to_str[rng_type]);
-        printf("    seed:            %d\n", seed);
+        printf("    seed:            %ld\n", seed);
     }
 };
 
@@ -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");
 }
@@ -233,7 +233,7 @@ void parse_args(int argc, const char* argv[], Option* opt) {
                 invalid_arg = true;
                 break;
             }
-            opt->seed = std::stoi(argv[i]);
+            opt->seed = std::stoll(argv[i]);
         } else if (arg == "-h" || arg == "--help") {
             print_usage(argc, argv);
             exit(0);
@@ -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;
         }

rng.h

@@ -6,7 +6,7 @@
 
 class RNG {
    public:
-    virtual void manual_seed(uint32_t seed) = 0;
+    virtual void manual_seed(uint64_t seed) = 0;
     virtual std::vector<float> randn(uint32_t n) = 0;
 };
 
@@ -15,7 +15,7 @@ class STDDefaultRNG : public RNG {
     std::default_random_engine generator;
 
    public:
-    void manual_seed(uint32_t seed) {
+    void manual_seed(uint64_t seed) {
         generator.seed(seed);
     }
 

rng_philox.h

@@ -93,7 +93,7 @@ class PhiloxRNG : public RNG {
         this->offset = 0;
     }
 
-    void manual_seed(uint32_t seed) {
+    void manual_seed(uint64_t seed) {
         this->seed = seed;
         this->offset = 0;
     }

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;
@@ -3029,6 +3038,9 @@ class StableDiffusionGGML {
             }
             bool some_tensor_not_init = false;
             for (auto pair : tensors) {
+                if (pair.first.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos) {
+                    continue;
+                }
                 if (tensor_names_in_file.find(pair.first) == tensor_names_in_file.end()) {
                     LOG_ERROR("tensor '%s' not in model file", pair.first.c_str());
                     some_tensor_not_init = true;
@@ -3096,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;
@@ -3119,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",
@@ -3152,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);
@@ -3162,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;
 
@@ -3198,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;
@@ -3248,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);
 
@@ -3332,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;
@@ -3357,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",
@@ -3390,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);
@@ -3449,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;
 
@@ -3471,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
@@ -3587,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;
@@ -3604,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",
@@ -3629,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
@@ -3716,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;
@@ -3733,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",
@@ -3758,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
@@ -3823,10 +3835,11 @@ std::vector<uint8_t> StableDiffusion::txt2img(const std::string& prompt,
                                               int height,
                                               SampleMethod sample_method,
                                               int sample_steps,
-                                              int seed) {
+                                              int64_t seed) {
     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;
@@ -3911,7 +3924,7 @@ std::vector<uint8_t> StableDiffusion::img2img(const std::vector<uint8_t>& init_i
                                               SampleMethod sample_method,
                                               int sample_steps,
                                               float strength,
-                                              int seed) {
+                                              int64_t seed) {
     std::vector<uint8_t> result;
     if (init_img_vec.size() != width * height * 3) {
         return result;

stable-diffusion.h

@@ -40,7 +40,7 @@ class StableDiffusion {
         int height,
         SampleMethod sample_method,
         int sample_steps,
-        int seed);
+        int64_t seed);
     std::vector<uint8_t> img2img(
         const std::vector<uint8_t>& init_img,
         const std::string& prompt,
@@ -51,7 +51,7 @@ class StableDiffusion {
         SampleMethod sample_method,
         int sample_steps,
         float strength,
-        int seed);
+        int64_t seed);
 };
 
 void set_sd_log_level(SDLogLevel level);