fix: harden for large files (#643)

Thank you for your contribution.

README.md

@@ -334,9 +334,9 @@ arguments:
   --skip-layers LAYERS               Layers to skip for SLG steps: (default: [7,8,9])
   --skip-layer-start START           SLG enabling point: (default: 0.01)
   --skip-layer-end END               SLG disabling point: (default: 0.2)
-  --scheduler {discrete, karras, exponential, ays, gits} Denoiser sigma scheduler (default: discrete)
+  --scheduler {discrete, karras, exponential, ays, gits, smoothstep} Denoiser sigma scheduler (default: discrete)
   --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}
-                                     sampling method (default: "euler_a")
+                                     sampling method (default: "euler" for Flux/SD3/Wan, "euler_a" otherwise)
   --steps  STEPS                     number of sample steps (default: 20)
   --high-noise-cfg-scale SCALE       (high noise) unconditional guidance scale: (default: 7.0)
   --high-noise-img-cfg-scale SCALE   (high noise) image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)
@@ -347,7 +347,7 @@ arguments:
   --high-noise-skip-layers LAYERS    (high noise) Layers to skip for SLG steps: (default: [7,8,9])
   --high-noise-skip-layer-start      (high noise) SLG enabling point: (default: 0.01)
   --high-noise-skip-layer-end END    (high noise) SLG disabling point: (default: 0.2)
-  --high-noise-scheduler {discrete, karras, exponential, ays, gits} Denoiser sigma scheduler (default: discrete)
+  --high-noise-scheduler {discrete, karras, exponential, ays, gits, smoothstep} Denoiser sigma scheduler (default: discrete)
   --high-noise-sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}
                                      (high noise) sampling method (default: "euler_a")
   --high-noise-steps  STEPS          (high noise) number of sample steps (default: -1 = auto)

examples/cli/main.cpp

@@ -240,7 +240,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --skip-layer-end END               SLG disabling point: (default: 0.2)\n");
     printf("  --scheduler {discrete, karras, exponential, ays, gits, smoothstep} Denoiser sigma scheduler (default: discrete)\n");
     printf("  --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd}\n");
-    printf("                                     sampling method (default: \"euler_a\")\n");
+    printf("                                     sampling method (default: \"euler\" for Flux/SD3/Wan, \"euler_a\" otherwise)\n");
     printf("  --steps  STEPS                     number of sample steps (default: 20)\n");
     printf("  --high-noise-cfg-scale SCALE       (high noise) unconditional guidance scale: (default: 7.0)\n");
     printf("  --high-noise-img-cfg-scale SCALE   (high noise) image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)\n");
@@ -1202,6 +1202,10 @@ int main(int argc, const char* argv[]) {
         return 1;
     }
 
+    if (params.sample_params.sample_method == SAMPLE_METHOD_DEFAULT) {
+        params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx);
+    }
+
     sd_image_t* results;
     int num_results = 1;
     if (params.mode == IMG_GEN) {

model.cpp

@@ -107,7 +107,7 @@ const char* unused_tensors[] = {
 };
 
 bool is_unused_tensor(std::string name) {
-    for (int i = 0; i < sizeof(unused_tensors) / sizeof(const char*); i++) {
+    for (size_t i = 0; i < sizeof(unused_tensors) / sizeof(const char*); i++) {
         if (starts_with(name, unused_tensors[i])) {
             return true;
         }

model.h

@@ -119,7 +119,7 @@ struct TensorStorage {
 
     size_t file_index = 0;
     int index_in_zip  = -1;  // >= means stored in a zip file
-    size_t offset     = 0;   // offset in file
+    uint64_t offset   = 0;   // offset in file
 
     TensorStorage() = default;
 
@@ -164,10 +164,10 @@ struct TensorStorage {
 
     std::vector<TensorStorage> chunk(size_t n) {
         std::vector<TensorStorage> chunks;
-        size_t chunk_size = nbytes_to_read() / n;
+        uint64_t chunk_size = nbytes_to_read() / n;
         // printf("%d/%d\n", chunk_size, nbytes_to_read());
         reverse_ne();
-        for (int i = 0; i < n; i++) {
+        for (size_t i = 0; i < n; i++) {
             TensorStorage chunk_i = *this;
             chunk_i.ne[0]         = ne[0] / n;
             chunk_i.offset        = offset + i * chunk_size;

stable-diffusion.cpp

@@ -43,7 +43,7 @@ const char* model_version_to_str[] = {
 };
 
 const char* sampling_methods_str[] = {
-    "Euler A",
+    "default",
     "Euler",
     "Heun",
     "DPM2",
@@ -55,6 +55,7 @@ const char* sampling_methods_str[] = {
     "LCM",
     "DDIM \"trailing\"",
     "TCD",
+    "Euler A",
 };
 
 /*================================================== Helper Functions ================================================*/
@@ -295,11 +296,6 @@ public:
             model_loader.set_wtype_override(wtype);
         }
 
-        if (sd_version_is_sdxl(version)) {
-            vae_wtype = GGML_TYPE_F32;
-            model_loader.set_wtype_override(GGML_TYPE_F32, "vae.");
-        }
-
         LOG_INFO("Weight type:                 %s", ggml_type_name(model_wtype));
         LOG_INFO("Conditioner weight type:     %s", ggml_type_name(conditioner_wtype));
         LOG_INFO("Diffusion model weight type: %s", ggml_type_name(diffusion_model_wtype));
@@ -1505,7 +1501,7 @@ enum rng_type_t str_to_rng_type(const char* str) {
 }
 
 const char* sample_method_to_str[] = {
-    "euler_a",
+    "default",
     "euler",
     "heun",
     "dpm2",
@@ -1517,6 +1513,7 @@ const char* sample_method_to_str[] = {
     "lcm",
     "ddim_trailing",
     "tcd",
+    "euler_a",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {
@@ -1655,7 +1652,7 @@ void sd_sample_params_init(sd_sample_params_t* sample_params) {
     sample_params->guidance.slg.layer_end      = 0.2f;
     sample_params->guidance.slg.scale          = 0.f;
     sample_params->scheduler                   = DEFAULT;
-    sample_params->sample_method               = EULER_A;
+    sample_params->sample_method               = SAMPLE_METHOD_DEFAULT;
     sample_params->sample_steps                = 20;
 }
 
@@ -1797,6 +1794,17 @@ void free_sd_ctx(sd_ctx_t* sd_ctx) {
     free(sd_ctx);
 }
 
+enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
+    if (sd_ctx != NULL && sd_ctx->sd != NULL) {
+        SDVersion version = sd_ctx->sd->version;
+        if (sd_version_is_dit(version))
+            return EULER;
+        else
+            return EULER_A;
+    }
+    return SAMPLE_METHOD_COUNT;
+}
+
 sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx,
                                     struct ggml_context* work_ctx,
                                     ggml_tensor* init_latent,
@@ -2361,6 +2369,11 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
         LOG_INFO("encode_first_stage completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
     }
 
+    enum sample_method_t sample_method = sd_img_gen_params->sample_params.sample_method;
+    if (sample_method == SAMPLE_METHOD_DEFAULT) {
+        sample_method = sd_get_default_sample_method(sd_ctx);
+    }
+
     sd_image_t* result_images = generate_image_internal(sd_ctx,
                                                         work_ctx,
                                                         init_latent,
@@ -2371,7 +2384,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
                                                         sd_img_gen_params->sample_params.eta,
                                                         width,
                                                         height,
-                                                        sd_img_gen_params->sample_params.sample_method,
+                                                        sample_method,
                                                         sigmas,
                                                         seed,
                                                         sd_img_gen_params->batch_count,

stable-diffusion.h

@@ -35,7 +35,7 @@ enum rng_type_t {
 };
 
 enum sample_method_t {
-    EULER_A,
+    SAMPLE_METHOD_DEFAULT,
     EULER,
     HEUN,
     DPM2,
@@ -47,6 +47,7 @@ enum sample_method_t {
     LCM,
     DDIM_TRAILING,
     TCD,
+    EULER_A,
     SAMPLE_METHOD_COUNT
 };
 
@@ -238,6 +239,7 @@ SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);
 
 SD_API sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params);
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
+SD_API enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx);
 
 SD_API void sd_sample_params_init(sd_sample_params_t* sample_params);
 SD_API char* sd_sample_params_to_str(const sd_sample_params_t* sample_params);