From 9c0cf36cd7964978850fd40f8cf76ac2cf9f9e15 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Sun, 1 Jun 2025 17:07:41 +0800
Subject: [PATCH] weight adapter fixes for training node

---
 comfy/weight_adapter/base.py |  5 +++--
 comfy/weight_adapter/lora.py | 14 ++++++++------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/comfy/weight_adapter/base.py b/comfy/weight_adapter/base.py
index c4ee67e9a..b5c7db423 100644
--- a/comfy/weight_adapter/base.py
+++ b/comfy/weight_adapter/base.py
@@ -12,13 +12,14 @@ class WeightAdapterBase:
     weights: list[torch.Tensor]
 
     @classmethod
-    def load(cls, x: str, lora: dict[str, torch.Tensor]) -> Optional["WeightAdapterBase"]:
+    def load(cls, x: str, lora: dict[str, torch.Tensor], alpha: float, dora_scale: torch.Tensor) -> Optional["WeightAdapterBase"]:
         raise NotImplementedError
 
     def to_train(self) -> "WeightAdapterTrainBase":
         raise NotImplementedError
 
-    def create_train(self, weight, *args) -> "WeightAdapterTrainBase":
+    @classmethod
+    def create_train(cls, weight, *args) -> "WeightAdapterTrainBase":
         """
         weight: The original weight tensor to be modified.
         *args: Additional arguments for configuration, such as rank, alpha etc.
diff --git a/comfy/weight_adapter/lora.py b/comfy/weight_adapter/lora.py
index 75d4e90cd..60145d5b1 100644
--- a/comfy/weight_adapter/lora.py
+++ b/comfy/weight_adapter/lora.py
@@ -29,17 +29,18 @@ class LoraDiff(WeightAdapterTrainBase):
             layer = torch.nn.Linear
         self.lora_up = layer(rank, out_dim, bias=False)
         self.lora_down = layer(in_dim, rank, bias=False)
-        self.lora_up.weight.copy_(mat1)
-        self.lora_down.weight.copy_(mat2)
+        self.lora_up.weight.data.copy_(mat1)
+        self.lora_down.weight.data.copy_(mat2)
         if mid is not None:
             self.lora_mid = layer(mid, rank, bias=False)
-            self.lora_mid.weight.copy_(mid)
+            self.lora_mid.weight.data.copy_(mid)
         else:
             self.lora_mid = None
         self.rank = rank
         self.alpha = torch.nn.Parameter(torch.tensor(alpha), requires_grad=False)
 
     def __call__(self, w):
+        org_dtype = w.dtype
         if self.lora_mid is None:
             diff = self.lora_up.weight @ self.lora_down.weight
         else:
@@ -48,7 +49,7 @@ class LoraDiff(WeightAdapterTrainBase):
             )
         scale = self.alpha / self.rank
         weight = w + scale * diff.reshape(w.shape)
-        return weight
+        return weight.to(org_dtype)
 
     def passive_memory_usage(self):
         return sum(param.numel() * param.element_size() for param in self.parameters())
@@ -61,13 +62,14 @@ class LoRAAdapter(WeightAdapterBase):
         self.loaded_keys = loaded_keys
         self.weights = weights
 
-    def create_train(self, weight, rank=1, alpha=1.0):
+    @classmethod
+    def create_train(cls, weight, rank=1, alpha=1.0):
         out_dim = weight.shape[0]
         in_dim = weight.shape[1:].numel()
         mat1 = torch.empty(out_dim, rank, device=weight.device, dtype=weight.dtype)
         mat2 = torch.empty(rank, in_dim, device=weight.device, dtype=weight.dtype)
         torch.nn.init.kaiming_uniform_(mat1, a=5**0.5)
-        torch.nn.init.constant__(mat2, 0.0)
+        torch.nn.init.constant_(mat2, 0.0)
         return LoraDiff(
             (mat1, mat2, alpha, None, None, None)
         )