fix: apply LoRAs at runtime where a merge crashes, and size FP8 on Metal

CMakeLists.txt

@@ -76,7 +76,8 @@ function(sd_native_carry_patch patch)
 endfunction()
 
 sd_native_carry_patch(pickle_line_span.patch src/model_io/pickle_io.cpp)
-sd_native_carry_patch(fp8_metal_f16_load.patch src/model_manager.cpp)
+sd_native_carry_patch(fp8_metal_f16_load.patch src/model_manager.cpp src/core/backend_fit.cpp)
+sd_native_carry_patch(lora_immediate_fallback.patch src/pipeline/diffusion_engine.cpp)
 
 # Wrapper exports are compiled into the upstream library target, so they share
 # its SD_API export macro and reach upstream without patching it.

README.md

@@ -610,10 +610,14 @@ the macOS probe and builds the iOS device and simulator slices.
 Upstream 1017 has its
 [ABI audit, wrapper review, FP8 evidence and outstanding gates](docs/upstream_1017_qualification.md);
 the carry inventory is in the [upstream 982 record](docs/upstream_982_qualification.md).
-It carries three fixes the pin lacks: unbatched CFG in Metal builds (below),
+It carries four fixes the pin lacks: unbatched CFG in Metal builds (below),
 the [`pickle_line_span.patch`](patches/README.md) for legacy checkpoints above
-2 GiB, and the [`fp8_metal_f16_load.patch`](patches/README.md), which loads FP8
-weights as F16 on Metal instead of aborting.
+2 GiB, the [`fp8_metal_f16_load.patch`](patches/README.md), which loads FP8
+weights as F16 on Metal instead of aborting, so they take twice their file size
+there, and the [`lora_immediate_fallback.patch`](patches/README.md), which
+applies a LoRA at runtime instead of crashing where the `immediately` mode
+cannot merge it: into FP8 weights, and on Metal, which cannot add into the F16
+convolution weights every model has.
 Upstream's `generate_image` now returns an owned result with per-image
 generation metadata; `sd_dart_generate_image` keeps returning an image array
 that `free_sd_images` frees and drops that metadata. Since upstream 982 a null

patches/README.md

@@ -13,4 +13,5 @@ an upstream issue or pull request to retire it.
 | Patch | Sources | Fixes | Test |
 | --- | --- | --- | --- |
 | `pickle_line_span.patch` | `src/model_io/pickle_io.cpp` | A legacy PyTorch checkpoint above 2 GiB is rejected: the pickle scanner narrows the rest of the file to `int` before looking for a newline ([#23](https://github.com/leehack/stable-diffusion-native/issues/23)). | `tests/test_pickle_span.py` |
-| `fp8_metal_f16_load.patch` | `src/model_manager.cpp` | An FP8 (E4M3 or E5M2) checkpoint aborts the process on Metal: ggml-metal claims FP8 `MUL_MAT` but has no FP8 kernel. Weights a Metal backend computes with now load as F16, which holds every FP8 value; other backends keep FP8 ([#29](https://github.com/leehack/stable-diffusion-native/issues/29)). | `tests/test_fp8_runtime.py` |
+| `fp8_metal_f16_load.patch` | `src/model_manager.cpp`, `src/core/backend_fit.cpp` | An FP8 (E4M3 or E5M2) checkpoint aborts the process on Metal: ggml-metal claims FP8 `MUL_MAT` but has no FP8 kernel. Weights a Metal backend computes with now load as F16, which holds every FP8 value, so they take twice their file size; other backends keep FP8 ([#29](https://github.com/leehack/stable-diffusion-native/issues/29)). Auto-fit plans them at that size, and FP8 LoRA tensors, which crashed ggml-cpu, load as F16 on every backend ([#31](https://github.com/leehack/stable-diffusion-native/issues/31)). | `tests/test_fp8_runtime.py` |
+| `lora_immediate_fallback.patch` | `src/pipeline/diffusion_engine.cpp` | The `immediately` LoRA apply mode, which `auto` selects for unquantized weights, crashes the CPU merging into FP8 weights (ggml-cpu can neither multiply nor write FP8) and aborts Metal merging into F16 weights (ggml-metal adds only equal types), which every model has: convolution weights load as F16. FP8 weights, and every model on a backend that cannot run that add outside host memory, now take the LoRA `at_runtime`, with a warning when `immediately` was requested ([#31](https://github.com/leehack/stable-diffusion-native/issues/31)). | `tests/test_fp8_runtime.py` |

patches/fp8_metal_f16_load.patch

@@ -1,3 +1,38 @@
+diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp
+--- a/src/core/backend_fit.cpp
++++ b/src/core/backend_fit.cpp
+@@ -96,7 +96,19 @@ namespace sd::backend_fit {
+         return false;
+     }
+
++    // model_manager.cpp loads FP8 weights that Metal computes with as F16, at
++    // twice their size. A component the plan keeps off Metal is overestimated.
++    static bool metal_loads_fp8_as_f16() {
++#ifndef SD_USE_UPSTREAM_GGML
++        ggml_backend_reg_t reg = ggml_backend_reg_by_name("MTL");
++        return reg != nullptr && ggml_backend_reg_dev_count(reg) > 0;
++#else
++        return false;
++#endif
++    }
++
+     static std::vector<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) {
++        const bool fp8_as_f16     = metal_loads_fp8_as_f16();
+         int64_t bytes[3]          = {0, 0, 0};
+         int64_t largest_tensor[3] = {0, 0, 0};
+         for (const auto& [name, stored_tensor] : loader.get_tensor_storage_map()) {
+@@ -110,6 +122,11 @@ namespace sd::backend_fit {
+             } else if (override_wtype != GGML_TYPE_COUNT && loader.tensor_should_be_converted(ts, override_wtype)) {
+                 ts.type = override_wtype;
+             }
++#ifndef SD_USE_UPSTREAM_GGML
++            if (fp8_as_f16 && (ts.type == GGML_TYPE_F8_E4M3 || ts.type == GGML_TYPE_F8_E5M2)) {
++                ts.type = GGML_TYPE_F16;
++            }
++#endif
+             const int64_t tensor_bytes = (int64_t)ts.nbytes() + 64;
+             bytes[int(kind)] += tensor_bytes;
+             largest_tensor[int(kind)] = std::max(largest_tensor[int(kind)], tensor_bytes);
 diff --git a/src/model_manager.cpp b/src/model_manager.cpp
 --- a/src/model_manager.cpp
 +++ b/src/model_manager.cpp
@@ -9,20 +44,23 @@ diff --git a/src/model_manager.cpp b/src/model_manager.cpp
  #include <iterator>
  #include <mutex>
  #include <tuple>
-@@ -185,6 +186,36 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors)
+@@ -185,6 +186,40 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors)
      return size;
  }
 
 +// ggml-metal reports FP8 MUL_MAT as supported, so Linear keeps FP8 weights, but
 +// it has no FP8 matrix kernel and aborts the process encoding one. Weights a
 +// Metal backend computes with are loaded as F16 instead, which holds every FP8
-+// value exactly.
-+static void load_fp8_params_as_f16_on_metal(const std::map<std::string, ggml_tensor*>& tensors,
-+                                            ggml_backend_t compute_backend) {
++// value exactly. So are FP8 LoRA tensors on every backend: LoRA multiplies
++// them directly, and ggml-cpu has no FP8 matrix kernel either.
++static void load_fp8_params_as_f16(ModelComponent component,
++                                   const std::map<std::string, ggml_tensor*>& tensors,
++                                   ggml_backend_t compute_backend) {
 +#ifndef SD_USE_UPSTREAM_GGML
 +    ggml_backend_dev_t device = compute_backend != nullptr ? ggml_backend_get_device(compute_backend) : nullptr;
 +    ggml_backend_reg_t reg    = device != nullptr ? ggml_backend_dev_backend_reg(device) : nullptr;
-+    if (reg == nullptr || std::strcmp(ggml_backend_reg_name(reg), "MTL") != 0) {
++    if (component != ModelComponent::LoRA &&
++        (reg == nullptr || std::strcmp(ggml_backend_reg_name(reg), "MTL") != 0)) {
 +        return;
 +    }
 +    for (const auto& pair : tensors) {
@@ -38,6 +76,7 @@ diff --git a/src/model_manager.cpp b/src/model_manager.cpp
 +        }
 +    }
 +#else
++    (void)component;
 +    (void)tensors;
 +    (void)compute_backend;
 +#endif
@@ -46,11 +85,11 @@ diff --git a/src/model_manager.cpp b/src/model_manager.cpp
  void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector<std::pair<ggml_backend_t, size_t>>& device_limits) {
      if (compute_backend == nullptr) {
          return;
-@@ -228,6 +259,7 @@ bool ModelManager::register_param_tensors(ModelComponent component,
+@@ -228,6 +263,7 @@ bool ModelManager::register_param_tensors(ModelComponent component,
          LOG_ERROR("model manager tensor component is invalid");
          return false;
      }
-+    load_fp8_params_as_f16_on_metal(tensors, compute_backend);
++    load_fp8_params_as_f16(component, tensors, compute_backend);
      if (registered_tensor_size != nullptr) {
          *registered_tensor_size += estimate_tensors_size(tensors);
      }

patches/lora_immediate_fallback.patch

@@ -0,0 +1,82 @@
+diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp
+--- a/src/pipeline/diffusion_engine.cpp
++++ b/src/pipeline/diffusion_engine.cpp
+@@ -1012,26 +1012,64 @@ bool StableDiffusionGGML::tae_preview_only() const {
+     return config_->params.tae_preview_only && version != VERSION_SDXS_512_DS && version != VERSION_SDXS_09;
+ }
+
++// Applying a LoRA immediately adds an F32 diff into each F16 weight in place, on
++// the backend holding the weight. ggml-metal adds only equal types, and the
++// scheduler aborts on an op a pre-allocated tensor's backend cannot run, unless
++// the tensor is in host memory, where the CPU runs it. Every model has F16
++// weights: convolutions load as F16 whatever the file holds.
++static bool backend_adds_f32_into_f16(ggml_backend_t backend) {
++    if (ggml_backend_buft_is_host(ggml_backend_get_default_buffer_type(backend))) {
++        return true;
++    }
++    ggml_init_params params = {8 * ggml_tensor_overhead(), nullptr, true};
++    ggml_context* ctx       = ggml_init(params);
++    if (ctx == nullptr) {
++        return false;
++    }
++    ggml_tensor* weight  = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 16);
++    ggml_tensor* diff    = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 16);
++    const bool supported = ggml_backend_supports_op(backend, ggml_add_inplace(ctx, weight, diff));
++    ggml_free(ctx);
++    return supported;
++}
++
+ void StableDiffusionGGML::configure_weight_loading() {
+     const auto* sd_ctx_params = &config_->params;
+     const auto& model_loader  = model_manager->loader();
+     const auto wtype_stat     = model_loader.get_wtype_stat();
+     bool have_int8_tensorwise = false;
+     bool have_quant_type_rule = false;
++    bool have_fp8_weight      = false;
+     for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) {
+         if (tensor_storage.is_int8_tensorwise) {
+             have_int8_tensorwise = true;
+         }
++        const ggml_type load_type = tensor_storage.expected_type < GGML_TYPE_COUNT ? tensor_storage.expected_type
++                                                                                   : tensor_storage.type;
++        if (load_type == GGML_TYPE_F8_E4M3 || load_type == GGML_TYPE_F8_E5M2) {
++            have_fp8_weight = true;
++        }
+         if (tensor_storage.expected_type < GGML_TYPE_COUNT && ggml_is_quantized(tensor_storage.expected_type)) {
+             have_quant_type_rule = true;
+         }
+-        if (have_int8_tensorwise && have_quant_type_rule) {
+-            break;
++    }
++    // Merging into an FP8 weight crashes ggml-cpu, which has neither the FP8
++    // matrix kernel the F32 cast uses nor an F32 to FP8 conversion.
++    std::string immediate_lora_blocker;
++    if (have_fp8_weight) {
++        immediate_lora_blocker = "FP8 weights";
++    } else {
++        for (SDBackendModule module : {SDBackendModule::TE, SDBackendModule::DIFFUSION, SDBackendModule::VAE}) {
++            ggml_backend_t backend = backend_manager.runtime_backend(module);
++            if (backend != nullptr && !backend_adds_f32_into_f16(backend)) {
++                immediate_lora_blocker = std::string("F16 weights on ") + ggml_backend_name(backend);
++                break;
++            }
+         }
+     }
+
+     if (sd_ctx_params->lora_apply_mode == LORA_APPLY_AUTO) {
+-        bool have_quantized_weight = have_int8_tensorwise;
++        bool have_quantized_weight = have_int8_tensorwise || !immediate_lora_blocker.empty();
+         for (const auto& [type, _] : wtype_stat) {
+             if (ggml_is_quantized(type)) {
+                 have_quantized_weight = true;
+@@ -1054,6 +1092,10 @@ void StableDiffusionGGML::configure_weight_loading() {
+                 "INT8 tensorwise weights do not support the immediately LoRA apply mode; "
+                 "using at_runtime instead");
+             apply_lora_immediately = false;
++        } else if (!immediate_lora_blocker.empty()) {
++            LOG_WARN("%s do not support the immediately LoRA apply mode; using at_runtime instead",
++                     immediate_lora_blocker.c_str());
++            apply_lora_immediately = false;
+         } else if (row_split_active()) {
+             LOG_WARN(
+                 "row-split tensors do not support the immediately LoRA apply mode; "

tests/native/fp8_runtime_test.cpp

@@ -2,13 +2,18 @@
 // as llamadart does, and writes its RGB bytes.
 //
 // Usage: fp8_runtime_test devices
-//        fp8_runtime_test generate <model> <backend> <out.rgb>
+//        fp8_runtime_test generate <model> <backend> <out.rgb> [<lora> <mode> [<rules>]]
+//        fp8_runtime_test plan <model> <backend>
 //
 // The model is one that `exit_teardown_runtime_test make-model` wrote, or an
-// FP8 copy of it: it holds its TAESD decoder, so it needs no VAE.
+// FP8 copy of it: it holds its TAESD decoder, so it needs no VAE. <mode> is a
+// LoRA apply mode name: auto, immediately or at_runtime, and <rules> are
+// tensor type rules. Both commands print the log, which holds upstream's
+// auto-fit plan and LoRA apply mode warnings; `plan` only loads the model.
 
 #include "sd_dart_wrapper.h"
 
+#include <cstdint>
 #include <cstdio>
 #include <cstring>
 #include <string>
@@ -25,14 +30,45 @@ int list_devices() {
     return 0;
 }
 
-int generate(const char* model, const char* backend, const char* output) {
+void print_log() {
+    std::vector<char> text(4096);
+    for (uint64_t sequence = 0;
+         (sequence = sd_dart_log_read(sequence, text.data(), text.size(), nullptr, nullptr)) != 0;) {
+        std::printf("%s", text.data());
+    }
+}
+
+int plan(const char* model, const char* backend) {
     sd_dart_log_enable();
     sd_ctx_params_t context_params;
     sd_ctx_params_init(&context_params);
     context_params.model_path = model;
     context_params.taesd_path = model;
     context_params.backend    = backend;
     sd_ctx_t* context         = sd_dart_new_sd_ctx(&context_params);
+    if (context != nullptr) {
+        sd_dart_exit_free(context);
+    }
+    return context != nullptr ? 0 : 3;
+}
+
+int generate(const char* model, const char* backend, const char* output, const char* lora,
+             const char* mode, const char* rules) {
+    sd_dart_log_enable();
+    sd_ctx_params_t context_params;
+    sd_ctx_params_init(&context_params);
+    context_params.model_path        = model;
+    context_params.taesd_path        = model;
+    context_params.backend           = backend;
+    context_params.tensor_type_rules = rules;
+    if (lora != nullptr) {
+        context_params.lora_apply_mode = str_to_lora_apply_mode(mode);
+        if (context_params.lora_apply_mode == LORA_APPLY_MODE_COUNT) {
+            std::fprintf(stderr, "unknown LoRA apply mode %s\n", mode);
+            return 2;
+        }
+    }
+    sd_ctx_t* context         = sd_dart_new_sd_ctx(&context_params);
     if (context == nullptr) {
         std::vector<char> error(4096);
         sd_dart_last_error(error.data(), error.size());
@@ -49,6 +85,11 @@ int generate(const char* model, const char* backend, const char* output) {
     params.seed                           = 7;
     params.sample_params.sample_steps     = 2;
     params.sample_params.guidance.txt_cfg = 1.0f;
+    sd_lora_t adapter                     = {false, 1.0f, lora};
+    if (lora != nullptr) {
+        params.loras      = &adapter;
+        params.lora_count = 1;
+    }
     sd_image_t* images                    = nullptr;
     int count                             = 0;
     if (!sd_dart_generate_image(context, &params, &images, &count)) {
@@ -83,9 +124,19 @@ int main(int argc, char** argv) {
     if (argc == 2 && std::strcmp(argv[1], "devices") == 0) {
         return list_devices();
     }
-    if (argc == 5 && std::strcmp(argv[1], "generate") == 0) {
-        return generate(argv[2], argv[3], argv[4]);
+    int status = 2;
+    if (argc == 4 && std::strcmp(argv[1], "plan") == 0) {
+        status = plan(argv[2], argv[3]);
+    } else if ((argc == 5 || argc == 7 || argc == 8) && std::strcmp(argv[1], "generate") == 0) {
+        status = generate(argv[2], argv[3], argv[4], argc >= 7 ? argv[5] : nullptr,
+                          argc >= 7 ? argv[6] : nullptr, argc == 8 ? argv[7] : nullptr);
+    } else {
+        std::fprintf(stderr,
+                     "usage: %s devices | generate <model> <backend> <out.rgb> [<lora> <mode> [<rules>]] | "
+                     "plan <model> <backend>\n",
+                     argv[0]);
+        return status;
     }
-    std::fprintf(stderr, "usage: %s devices | generate <model> <backend> <out.rgb>\n", argv[0]);
-    return 2;
+    print_log();
+    return status;
 }