mirror of
https://github.com/vladmandic/automatic
synced 2026-08-26 15:16:01 +02:00
Preliminary support for sd1.5/sdxl OMI lora
This commit is contained in:
@@ -26,7 +26,7 @@ re_compiled = {}
|
||||
def make_unet_conversion_map() -> Dict[str, str]:
|
||||
unet_conversion_map_layer = []
|
||||
|
||||
for i in range(3): # num_blocks is 3 in sdxl
|
||||
for i in range(4): # num_blocks is 3 in sdxl
|
||||
# loop over downblocks/upblocks
|
||||
for j in range(2):
|
||||
# loop over resnets/attentions for downblocks
|
||||
@@ -108,7 +108,7 @@ def make_unet_conversion_map() -> Dict[str, str]:
|
||||
class KeyConvert:
|
||||
def __init__(self):
|
||||
self.is_sdxl = True if shared.sd_model_type == "sdxl" else False
|
||||
self.UNET_CONVERSION_MAP = make_unet_conversion_map() if self.is_sdxl else None
|
||||
self.UNET_CONVERSION_MAP = make_unet_conversion_map()
|
||||
self.LORA_PREFIX_UNET = "lora_unet_"
|
||||
self.LORA_PREFIX_TEXT_ENCODER = "lora_te_"
|
||||
self.OFT_PREFIX_UNET = "oft_unet_"
|
||||
@@ -117,16 +117,15 @@ class KeyConvert:
|
||||
self.LORA_PREFIX_TEXT_ENCODER2 = "lora_te2_"
|
||||
|
||||
def __call__(self, key):
|
||||
if self.is_sdxl:
|
||||
if "diffusion_model" in key: # Fix NTC Slider naming error
|
||||
key = key.replace("diffusion_model", "lora_unet")
|
||||
map_keys = list(self.UNET_CONVERSION_MAP.keys()) # prefix of U-Net modules
|
||||
map_keys.sort()
|
||||
search_key = key.replace(self.LORA_PREFIX_UNET, "").replace(self.OFT_PREFIX_UNET, "").replace(self.LORA_PREFIX_TEXT_ENCODER1, "").replace(self.LORA_PREFIX_TEXT_ENCODER2, "")
|
||||
position = bisect.bisect_right(map_keys, search_key)
|
||||
map_key = map_keys[position - 1]
|
||||
if search_key.startswith(map_key):
|
||||
key = key.replace(map_key, self.UNET_CONVERSION_MAP[map_key]).replace("oft", "lora") # pylint: disable=unsubscriptable-object
|
||||
if "diffusion_model" in key: # Fix NTC Slider naming error
|
||||
key = key.replace("diffusion_model", "lora_unet")
|
||||
map_keys = list(self.UNET_CONVERSION_MAP.keys()) # prefix of U-Net modules
|
||||
map_keys.sort()
|
||||
search_key = key.replace(self.LORA_PREFIX_UNET, "").replace(self.OFT_PREFIX_UNET, "").replace(self.LORA_PREFIX_TEXT_ENCODER1, "").replace(self.LORA_PREFIX_TEXT_ENCODER2, "")
|
||||
position = bisect.bisect_right(map_keys, search_key)
|
||||
map_key = map_keys[position - 1]
|
||||
if search_key.startswith(map_key):
|
||||
key = key.replace(map_key, self.UNET_CONVERSION_MAP[map_key]).replace("oft", "lora") # pylint: disable=unsubscriptable-object
|
||||
if "lycoris" in key and "transformer" in key:
|
||||
key = key.replace("lycoris", "lora_transformer")
|
||||
sd_module = shared.sd_model.network_layer_mapping.get(key, None)
|
||||
|
||||
@@ -109,7 +109,19 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
|
||||
emb_dict[vec_name] = weight
|
||||
bundle_embeddings[emb_name] = emb_dict
|
||||
continue
|
||||
if len(parts) > 5: # messy handler for diffusers peft lora
|
||||
if parts[0] in ["clip_l","clip_g","t5","unet","transformer"]:
|
||||
network_part = []
|
||||
while parts[-1] in ["alpha","weight","lora_up","lora_down"]:
|
||||
network_part.insert(0,parts[-1])
|
||||
parts = parts[0:-1]
|
||||
network_part = ".".join(network_part)
|
||||
key_network_without_network_parts = "_".join(parts)
|
||||
if key_network_without_network_parts.startswith("unet") or key_network_without_network_parts.startswith("transformer"):
|
||||
key_network_without_network_parts = "lora_" + key_network_without_network_parts
|
||||
key_network_without_network_parts = key_network_without_network_parts.replace("clip_g","lora_te2").replace("clip_l","lora_te")
|
||||
#TODO Add t5 key support for SD3.5/f1 here?
|
||||
|
||||
elif len(parts) > 5: # messy handler for diffusers peft lora
|
||||
key_network_without_network_parts = '_'.join(parts[:-2])
|
||||
if not key_network_without_network_parts.startswith('lora_'):
|
||||
key_network_without_network_parts = 'lora_' + key_network_without_network_parts
|
||||
|
||||
Reference in New Issue
Block a user