From 84576241013026d7d0723186c8d6ab65484a6f28 Mon Sep 17 00:00:00 2001
From: stduhpf <stephduh@live.fr>
Date: Sun, 2 Aug 2026 10:08:47 +0200
Subject: [PATCH] feat: support more LoRA models (Kroma-v0.1 support) (#1842)

---
 src/model/diffusion/krea2.hpp | 27 ++++++++++++++++++++-------
 src/name_conversion.cpp       |  2 ++
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp
index 9121069a..b3947b71 100644
--- a/src/model/diffusion/krea2.hpp
+++ b/src/model/diffusion/krea2.hpp
@@ -180,9 +180,12 @@ namespace Krea2 {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             ggml_tensor* scale = params["scale"];
-            scale              = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
-            x                  = ggml_rms_norm(ctx->ggml_ctx, x, eps);
-            x                  = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
+            if (ctx->weight_adapter) {
+                scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight");
+            }
+            scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
+            x     = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+            x     = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
             return x;
         }
     };
@@ -295,10 +298,11 @@ namespace Krea2 {
     class KreaDoubleSharedModulation : public GGMLBlock {
     protected:
         int64_t dim;
+        std::string prefix;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             GGML_UNUSED(tensor_storage_map);
-            GGML_UNUSED(prefix);
+            this->prefix  = prefix;
             params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6);
         }
 
@@ -307,7 +311,11 @@ namespace Krea2 {
             : dim(dim) {}
 
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
-            auto lin = ggml_repeat(ctx->ggml_ctx, params["lin"], vec);
+            auto lin = params["lin"];
+            if (ctx->weight_adapter) {
+                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+            }
+            lin      = ggml_repeat(ctx->ggml_ctx, lin, vec);
             auto out = ggml_add(ctx->ggml_ctx, vec, lin);
             return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
         }
@@ -316,10 +324,11 @@ namespace Krea2 {
     class KreaFinalModulation : public GGMLBlock {
     protected:
         int64_t dim;
+        std::string prefix;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             GGML_UNUSED(tensor_storage_map);
-            GGML_UNUSED(prefix);
+            this->prefix  = prefix;
             params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2);
         }
 
@@ -328,7 +337,11 @@ namespace Krea2 {
             : dim(dim) {}
 
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
-            auto out = ggml_add(ctx->ggml_ctx, params["lin"], vec);
+            auto lin = params["lin"];
+            if (ctx->weight_adapter) {
+                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+            }
+            auto out = ggml_add(ctx->ggml_ctx, lin, vec);
             return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
         }
     };
diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp
index af8a7f57..d6e296e1 100644
--- a/src/name_conversion.cpp
+++ b/src/name_conversion.cpp
@@ -1384,6 +1384,8 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
             {".lora_B.weight", ".weight.lora_up"},
             {".lora_A.default.weight", ".weight.lora_down"},
             {".lora_B.default.weight", ".weight.lora_up"},
+            {".lora_A", ".weight.lora_down"},
+            {".lora_B", ".weight.lora_up"},
             {".lora_linear", ".weight.alpha"},
             {".alpha", ".weight.alpha"},
             {".scale", ".weight.scale"},