Utilize weight adapter scheme in basic training node

comfy/weight_adapter/__init__.py

@@ -1,4 +1,4 @@
-from .base import WeightAdapterBase
+from .base import WeightAdapterBase, WeightAdapterTrainBase
 from .lora import LoRAAdapter
 from .loha import LoHaAdapter
 from .lokr import LoKrAdapter

comfy/weight_adapter/base.py

@@ -18,6 +18,13 @@ class WeightAdapterBase:
     def to_train(self) -> "WeightAdapterTrainBase":
         raise NotImplementedError
 
+    def create_train(self, weight, *args) -> "WeightAdapterTrainBase":
+        """
+        weight: The original weight tensor to be modified.
+        *args: Additional arguments for configuration, such as rank, alpha etc.
+        """
+        raise NotImplementedError
+
     def calculate_weight(
         self,
         weight,

comfy/weight_adapter/lora.py

@@ -29,8 +29,11 @@ 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)
         if mid is not None:
             self.lora_mid = layer(mid, rank, bias=False)
+            self.lora_mid.weight.copy_(mid)
         else:
             self.lora_mid = None
         self.rank = rank
@@ -44,7 +47,7 @@ class LoraDiff(WeightAdapterTrainBase):
                 self.lora_up.weight, self.lora_down.weight, self.lora_mid.weight
             )
         scale = self.alpha / self.rank
-        weight = w + scale * diff
+        weight = w + scale * diff.reshape(w.shape)
         return weight
 
     def passive_memory_usage(self):
@@ -58,6 +61,17 @@ class LoRAAdapter(WeightAdapterBase):
         self.loaded_keys = loaded_keys
         self.weights = weights
 
+    def create_train(self, 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)
+        return LoraDiff(
+            (mat1, mat2, alpha, None, None, None)
+        )
+
     @classmethod
     def load(
         cls,

comfy_extras/nodes_train.py

@@ -17,6 +17,7 @@ import folder_paths
 import node_helpers
 from comfy.cli_args import args
 from comfy.comfy_types.node_typing import IO
+from comfy.weight_adapter import WeightAdapterBase, WeightAdapterTrainBase, adapters
 
 
 class TrainSampler(comfy.samplers.Sampler):
@@ -70,23 +71,6 @@ class BiasDiff(torch.nn.Module):
         return self.passive_memory_usage()
 
 
-class LoraDiff(torch.nn.Module):
-    def __init__(self, lora_down, lora_up):
-        super().__init__()
-        self.lora_down = lora_down
-        self.lora_up = lora_up
-
-    def __call__(self, w):
-        return w + (self.lora_up @ self.lora_down).reshape(w.shape)
-
-    def passive_memory_usage(self):
-        return self.lora_down.nelement() * self.lora_down.element_size() + self.lora_up.nelement() * self.lora_up.element_size()
-
-    def move_to(self, device):
-        self.to(device=device)
-        return self.passive_memory_usage()
-
-
 def load_and_process_images(image_files, input_dir, resize_method="None"):
     """Utility function to load and process a list of images.
 
@@ -384,52 +368,25 @@ class TrainLoraNode:
                         key = "{}.weight".format(n)
                         shape = m.weight.shape
                         if len(shape) >= 2:
-                            in_dim = math.prod(shape[1:])
-                            out_dim = shape[0]
-
-                            # Check if we have existing weights for this layer
-                            lora_up_key = "{}.lora_up.weight".format(n)
-                            lora_down_key = "{}.lora_down.weight".format(n)
-
-                            if existing_lora != "[None]" and (
-                                    lora_up_key in existing_weights
-                                    and lora_down_key in existing_weights
-                            ):
-                                # Initialize with existing weights
-                                lora_up = torch.nn.Parameter(
-                                        existing_weights[lora_up_key].to(dtype=dtype),
-                                        requires_grad=True,
-                                    )
-                                lora_down = torch.nn.Parameter(
-                                        existing_weights[lora_down_key].to(dtype=dtype),
-                                        requires_grad=True,
-                                    )
+                            existing_adapter = None
+                            for adapter_cls in adapters:
+                                existing_adapter = adapter_cls.load(
+                                    n, existing_weights
+                                )
+                                if existing_adapter is not None:
+                                    break
+                            
+                            if existing_adapter is not None:
+                                train_adapter = existing_adapter.to_train()
+                                for name, parameter in train_adapter.named_parameters():
+                                    lora_sd[f"{n}.{name}"] = parameter
                             else:
-                                if existing_lora != "[None]":
-                                    logging.info(f"Warning: No existing weights found for {lora_up_key} or {lora_down_key}")
-                                # Initialize new weights
-                                lora_down = torch.nn.Parameter(
-                                    torch.zeros(
-                                        (
-                                            rank,
-                                            in_dim,
-                                        ),
-                                        dtype=dtype,
-                                    ),
-                                    requires_grad=True,
-                                )
-                                lora_up = torch.nn.Parameter(
-                                    torch.zeros((out_dim, rank), dtype=dtype),
-                                    requires_grad=True,
-                                )
-                                torch.nn.init.zeros_(lora_up)
-                                torch.nn.init.kaiming_uniform_(
-                                    lora_down, a=math.sqrt(5), generator=generator
+                                # Use LoRA with alpha=1.0 by default
+                                train_adapter = adapter_cls[0].create_train(
+                                    m.weight, rank=rank, alpha=1.0
                                 )
 
-                            lora_sd[lora_up_key] = lora_up
-                            lora_sd[lora_down_key] = lora_down
-                            mp.add_weight_wrapper(key, LoraDiff(lora_down, lora_up))
+                            mp.add_weight_wrapper(key, train_adapter)
                         else:
                             diff = torch.nn.Parameter(
                                 torch.zeros(