47 lines
1.1 KiB
Python
47 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
from .base import LycorisBaseModule
|
|
from .locon import LoConModule
|
|
from .loha import LohaModule
|
|
from .lokr import LokrModule
|
|
from .full import FullModule
|
|
from .norms import NormModule
|
|
from .diag_oft import DiagOFTModule
|
|
from .boft import ButterflyOFTModule
|
|
from .glora import GLoRAModule
|
|
from .dylora import DyLoraModule
|
|
from .ia3 import IA3Module
|
|
|
|
from ..functional.general import factorization
|
|
|
|
|
|
MODULE_LIST = [
|
|
LoConModule,
|
|
LohaModule,
|
|
IA3Module,
|
|
LokrModule,
|
|
FullModule,
|
|
NormModule,
|
|
DiagOFTModule,
|
|
ButterflyOFTModule,
|
|
GLoRAModule,
|
|
DyLoraModule,
|
|
]
|
|
|
|
|
|
def get_module(lyco_state_dict, lora_name):
|
|
for module in MODULE_LIST:
|
|
if module.algo_check(lyco_state_dict, lora_name):
|
|
return module, tuple(module.extract_state_dict(lyco_state_dict, lora_name))
|
|
return None, None
|
|
|
|
|
|
@torch.no_grad()
|
|
def make_module(lyco_type: LycorisBaseModule, params, lora_name, orig_module):
|
|
try:
|
|
module = lyco_type.make_module_from_state_dict(lora_name, orig_module, *params)
|
|
except NotImplementedError:
|
|
module = None
|
|
return module
|