mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
119 lines
4.5 KiB
Python
119 lines
4.5 KiB
Python
# Copyright (c) 2023 The InvokeAI Development Team
|
|
"""Utilities used by the Model Manager"""
|
|
|
|
|
|
def lora_token_vector_length(checkpoint: dict) -> int:
|
|
"""
|
|
Given a checkpoint in memory, return the lora token vector length
|
|
|
|
:param checkpoint: The checkpoint
|
|
"""
|
|
|
|
def _handle_unet_key(key, tensor, checkpoint):
|
|
lora_token_vector_length = None
|
|
if "_attn2_to_k." not in key and "_attn2_to_v." not in key:
|
|
return lora_token_vector_length
|
|
|
|
# check lora/locon
|
|
if ".lora_up.weight" in key:
|
|
lora_token_vector_length = tensor.shape[0]
|
|
elif ".lora_down.weight" in key:
|
|
lora_token_vector_length = tensor.shape[1]
|
|
|
|
# check loha (don't worry about hada_t1/hada_t2 as it used only in 4d shapes)
|
|
elif ".hada_w1_b" in key or ".hada_w2_b" in key:
|
|
lora_token_vector_length = tensor.shape[1]
|
|
|
|
# check lokr (don't worry about lokr_t2 as it used only in 4d shapes)
|
|
elif ".lokr_" in key:
|
|
_lokr_key = key.split(".")[0]
|
|
|
|
if _lokr_key + ".lokr_w1" in checkpoint:
|
|
_lokr_w1 = checkpoint[_lokr_key + ".lokr_w1"]
|
|
elif _lokr_key + "lokr_w1_b" in checkpoint:
|
|
_lokr_w1 = checkpoint[_lokr_key + ".lokr_w1_b"]
|
|
else:
|
|
return lora_token_vector_length # unknown format
|
|
|
|
if _lokr_key + ".lokr_w2" in checkpoint:
|
|
_lokr_w2 = checkpoint[_lokr_key + ".lokr_w2"]
|
|
elif _lokr_key + "lokr_w2_b" in checkpoint:
|
|
_lokr_w2 = checkpoint[_lokr_key + ".lokr_w2_b"]
|
|
else:
|
|
return lora_token_vector_length # unknown format
|
|
|
|
lora_token_vector_length = _lokr_w1.shape[1] * _lokr_w2.shape[1]
|
|
|
|
elif ".diff" in key:
|
|
lora_token_vector_length = tensor.shape[1]
|
|
|
|
return lora_token_vector_length
|
|
|
|
def _handle_te_key(key, tensor, checkpoint):
|
|
lora_token_vector_length = None
|
|
if "text_model_encoder_layers_" not in key:
|
|
return lora_token_vector_length
|
|
|
|
# skip detect by mlp
|
|
if "_self_attn_" not in key:
|
|
return lora_token_vector_length
|
|
|
|
# check lora/locon
|
|
if ".lora_up.weight" in key:
|
|
lora_token_vector_length = tensor.shape[0]
|
|
elif ".lora_down.weight" in key:
|
|
lora_token_vector_length = tensor.shape[1]
|
|
|
|
# check loha (don't worry about hada_t1/hada_t2 as it used only in 4d shapes)
|
|
elif ".hada_w1_a" in key or ".hada_w2_a" in key:
|
|
lora_token_vector_length = tensor.shape[0]
|
|
elif ".hada_w1_b" in key or ".hada_w2_b" in key:
|
|
lora_token_vector_length = tensor.shape[1]
|
|
|
|
# check lokr (don't worry about lokr_t2 as it used only in 4d shapes)
|
|
elif ".lokr_" in key:
|
|
_lokr_key = key.split(".")[0]
|
|
|
|
if _lokr_key + ".lokr_w1" in checkpoint:
|
|
_lokr_w1 = checkpoint[_lokr_key + ".lokr_w1"]
|
|
elif _lokr_key + "lokr_w1_b" in checkpoint:
|
|
_lokr_w1 = checkpoint[_lokr_key + ".lokr_w1_b"]
|
|
else:
|
|
return lora_token_vector_length # unknown format
|
|
|
|
if _lokr_key + ".lokr_w2" in checkpoint:
|
|
_lokr_w2 = checkpoint[_lokr_key + ".lokr_w2"]
|
|
elif _lokr_key + "lokr_w2_b" in checkpoint:
|
|
_lokr_w2 = checkpoint[_lokr_key + ".lokr_w2_b"]
|
|
else:
|
|
return lora_token_vector_length # unknown format
|
|
|
|
lora_token_vector_length = _lokr_w1.shape[1] * _lokr_w2.shape[1]
|
|
|
|
elif ".diff" in key:
|
|
lora_token_vector_length = tensor.shape[1]
|
|
|
|
return lora_token_vector_length
|
|
|
|
lora_token_vector_length = None
|
|
lora_te1_length = None
|
|
lora_te2_length = None
|
|
for key, tensor in checkpoint.items():
|
|
if key.startswith("lora_unet_"):
|
|
lora_token_vector_length = _handle_unet_key(key, tensor, checkpoint)
|
|
elif key.startswith("lora_te_"):
|
|
lora_token_vector_length = _handle_te_key(key, tensor, checkpoint)
|
|
|
|
elif key.startswith("lora_te1_"):
|
|
lora_te1_length = _handle_te_key(key, tensor, checkpoint)
|
|
elif key.startswith("lora_te2_"):
|
|
lora_te2_length = _handle_te_key(key, tensor, checkpoint)
|
|
|
|
if lora_te1_length is not None and lora_te2_length is not None:
|
|
lora_token_vector_length = lora_te1_length + lora_te2_length
|
|
|
|
if lora_token_vector_length is not None:
|
|
break
|
|
|
|
return lora_token_vector_length
|