From 5098e9408d1cdd0e085d74f3ea9a402aeb592149 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 22 Apr 2025 14:54:20 +0800 Subject: [PATCH] Utilize weight adapter scheme in basic training node --- comfy/weight_adapter/__init__.py | 2 +- comfy/weight_adapter/base.py | 7 +++ comfy/weight_adapter/lora.py | 16 ++++++- comfy_extras/nodes_train.py | 77 +++++++------------------------- 4 files changed, 40 insertions(+), 62 deletions(-) diff --git a/comfy/weight_adapter/__init__.py b/comfy/weight_adapter/__init__.py index e6cd805b6..d278b032c 100644 --- a/comfy/weight_adapter/__init__.py +++ b/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 diff --git a/comfy/weight_adapter/base.py b/comfy/weight_adapter/base.py index 105cc3007..bea67f9de 100644 --- a/comfy/weight_adapter/base.py +++ b/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, diff --git a/comfy/weight_adapter/lora.py b/comfy/weight_adapter/lora.py index 062511b93..75d4e90cd 100644 --- a/comfy/weight_adapter/lora.py +++ b/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, diff --git a/comfy_extras/nodes_train.py b/comfy_extras/nodes_train.py index 3d6916523..4ffc64345 100644 --- a/comfy_extras/nodes_train.py +++ b/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(