|
|
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 |
|
|
|
