156 lines
5.9 KiB
Python
156 lines
5.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from toolkit.network_mixins import ToolkitNetworkMixin
|
|
from toolkit.models.loha import LoHaModule
|
|
import re
|
|
import weakref
|
|
|
|
class LoHaNetwork(ToolkitNetworkMixin, nn.Module):
|
|
def __init__(
|
|
self,
|
|
text_encoder,
|
|
unet,
|
|
lora_dim=4,
|
|
alpha=1,
|
|
dropout=None,
|
|
rank_dropout=None,
|
|
module_dropout=None,
|
|
multiplier=1.0,
|
|
train_unet=True,
|
|
train_text_encoder=True,
|
|
**kwargs
|
|
):
|
|
nn.Module.__init__(self)
|
|
super().__init__(
|
|
text_encoder=text_encoder,
|
|
unet=unet,
|
|
train_text_encoder=train_text_encoder,
|
|
train_unet=train_unet,
|
|
**kwargs
|
|
)
|
|
self.network_type = "loha"
|
|
self.peft_format = None
|
|
|
|
# Force None to bypass incompatible base mixin save hooks
|
|
self.base_model_ref = None
|
|
|
|
# FIX: Explicitly set is_pixart.
|
|
# The base Mixin uses this in load_weights but doesn't initialize it itself.
|
|
self.is_pixart = kwargs.get('is_pixart', False)
|
|
|
|
self.lora_dim = lora_dim
|
|
self.alpha = alpha
|
|
self.dropout = dropout
|
|
self.rank_dropout = rank_dropout
|
|
self.module_dropout = module_dropout
|
|
self._multiplier = multiplier
|
|
|
|
# Use ModuleList so PyTorch registers the parameters for saving
|
|
self.loha_modules = nn.ModuleList()
|
|
|
|
if train_unet:
|
|
self.create_modules(unet, "lora_unet")
|
|
if train_text_encoder:
|
|
if isinstance(text_encoder, list):
|
|
for i, te in enumerate(text_encoder):
|
|
self.create_modules(te, f"lora_te_{i+1}")
|
|
else:
|
|
self.create_modules(text_encoder, "lora_te")
|
|
|
|
def create_modules(self, root_module, prefix):
|
|
for name, module in root_module.named_modules():
|
|
if module.__class__.__name__ in ["Linear", "Conv2d", "LoRACompatibleLinear", "LoRACompatibleConv"]:
|
|
if not module.weight.requires_grad and "refiner" not in prefix:
|
|
pass
|
|
|
|
lora_name = f"{prefix}_{name}".replace('.', '_')
|
|
|
|
loha_mod = LoHaModule(
|
|
lora_name=lora_name,
|
|
network=self,
|
|
org_module=module,
|
|
multiplier=self._multiplier,
|
|
lora_dim=self.lora_dim,
|
|
alpha=self.alpha,
|
|
dropout=self.dropout,
|
|
rank_dropout=self.rank_dropout,
|
|
module_dropout=self.module_dropout,
|
|
)
|
|
|
|
# Append works the same way with ModuleList
|
|
self.loha_modules.append(loha_mod)
|
|
|
|
# Direct injection
|
|
module.forward = loha_mod.forward
|
|
module.loha_module = loha_mod
|
|
|
|
def apply_to(self, text_encoder, unet, train_text_encoder, train_unet):
|
|
pass
|
|
|
|
def prepare_grad_etc(self, text_encoder, unet):
|
|
for module in self.get_all_modules():
|
|
for param in module.parameters():
|
|
param.requires_grad = True
|
|
|
|
def get_all_modules(self):
|
|
return self.loha_modules
|
|
|
|
def save_weights(self, file, dtype=torch.float16, metadata=None, extra_state_dict=None):
|
|
if metadata is None:
|
|
metadata = {}
|
|
|
|
metadata["ss_network_module"] = "lycoris.kohya"
|
|
metadata["ss_network_dim"] = str(self.lora_dim)
|
|
metadata["ss_network_alpha"] = str(self.alpha)
|
|
metadata["ss_network_args"] = str({'algo': 'loha'})
|
|
|
|
super().save_weights(file, dtype, metadata, extra_state_dict)
|
|
|
|
def prepare_optimizer_params(self, text_encoder_lr=None, unet_lr=None, default_lr=1e-4):
|
|
all_params = []
|
|
|
|
def enumerate_params(modules, lr):
|
|
params = []
|
|
for mod in modules:
|
|
for param in mod.parameters():
|
|
if param.requires_grad:
|
|
params.append(param)
|
|
return {"params": params, "lr": lr}
|
|
|
|
if self.train_unet:
|
|
unet_mods = [m for m in self.loha_modules if "lora_unet" in m.lora_name]
|
|
all_params.append(enumerate_params(unet_mods, unet_lr or default_lr))
|
|
|
|
if self.train_text_encoder:
|
|
te_mods = [m for m in self.loha_modules if "lora_te" in m.lora_name]
|
|
all_params.append(enumerate_params(te_mods, text_encoder_lr or default_lr))
|
|
|
|
return all_params
|
|
|
|
@torch.no_grad()
|
|
def _update_torch_multiplier(self):
|
|
multiplier = self._multiplier
|
|
try:
|
|
first_module = self.get_all_modules()[0]
|
|
except IndexError:
|
|
return
|
|
|
|
if hasattr(first_module, 'hada_w1_a'):
|
|
device = first_module.hada_w1_a.device
|
|
dtype = first_module.hada_w1_a.dtype
|
|
if hasattr(first_module.hada_w1_a, '_memory_management_device'):
|
|
device = first_module.hada_w1_a._memory_management_device
|
|
else:
|
|
raise ValueError(f"Unknown module type: {type(first_module)}")
|
|
|
|
with torch.no_grad():
|
|
tensor_multiplier = None
|
|
if isinstance(multiplier, int) or isinstance(multiplier, float):
|
|
# Create scalar tensor
|
|
tensor_multiplier = torch.tensor(multiplier).to(device, dtype=dtype)
|
|
elif isinstance(multiplier, list):
|
|
tensor_multiplier = torch.tensor(multiplier).to(device, dtype=dtype)
|
|
elif isinstance(multiplier, torch.Tensor):
|
|
tensor_multiplier = multiplier.clone().detach().to(device, dtype=dtype)
|
|
|
|
self.torch_multiplier = tensor_multiplier.clone().detach() |