From e320c3c577c20b33810400269522569755673ae9 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Fri, 2 May 2025 20:53:57 +0800
Subject: [PATCH] Correct calculate_weight and load for OFT

---
 comfy/weight_adapter/oft.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/comfy/weight_adapter/oft.py b/comfy/weight_adapter/oft.py
index 0ea229b79..25009eca3 100644
--- a/comfy/weight_adapter/oft.py
+++ b/comfy/weight_adapter/oft.py
@@ -32,17 +32,18 @@ class OFTAdapter(WeightAdapterBase):
             blocks = lora[blocks_name]
             if blocks.ndim == 3:
                 loaded_keys.add(blocks_name)
+            else:
+                blocks = None
+        if blocks is None:
+            return None
 
         rescale = None
         if rescale_name in lora.keys():
             rescale = lora[rescale_name]
             loaded_keys.add(rescale_name)
 
-        if blocks is not None:
-            weights = (blocks, rescale, alpha, dora_scale)
-            return cls(loaded_keys, weights)
-        else:
-            return None
+        weights = (blocks, rescale, alpha, dora_scale)
+        return cls(loaded_keys, weights)
 
     def calculate_weight(
         self,
@@ -79,16 +80,17 @@ class OFTAdapter(WeightAdapterBase):
                     normed_q = q * alpha / q_norm
             # use float() to prevent unsupported type in .inverse()
             r = (I + normed_q) @ (I - normed_q).float().inverse()
-            r = r.to(original_weight)
+            r = r.to(weight)
+            _, *shape = weight.shape
             lora_diff = torch.einsum(
                 "k n m, k n ... -> k m ...",
                 (r * strength) - strength * I,
-                original_weight,
-            )
+                weight.view(block_num, block_size, *shape),
+            ).view(-1, *shape)
             if dora_scale is not None:
                 weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength, intermediate_dtype, function)
             else:
-                weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
+                weight += function((strength * lora_diff).type(weight.dtype))
         except Exception as e:
             logging.error("ERROR {} {} {}".format(self.name, key, e))
         return weight