apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix SDXL node names once

This commit is contained in:
Sean Sube 2023-11-24 16:51:03 -06:00
parent 74832fc61b
commit d7c95a4a4f
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 15 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
api/onnx_web/convert/diffusion/lora.py
View File

@ -73,7 +73,7 @@ def fix_node_name(key: str):
def fix_xl_names(keys: Dict[str, Any], nodes: List[NodeProto]) -> Dict[str, Any]:
fixed = {}
remaining = list(nodes)
names = [fix_node_name(node.name) for node in nodes]
for key, value in keys.items():
root, *rest = key.split(".")
@ -128,28 +128,28 @@ def fix_xl_names(keys: Dict[str, Any], nodes: List[NodeProto]) -> Dict[str, Any]
continue
logger.trace("searching for XL node: %s -> /%s/*/%s", root, block, suffix)
match: Optional[NodeProto] = None
match: Optional[str] = None
if "conv" in suffix:
match = next(
node for node in remaining if fix_node_name(node.name) == f"{root}_Conv"
node for node in names if node == f"{root}_Conv"
)
elif "time_emb_proj" in root:
match = next(
node for node in remaining if fix_node_name(node.name) == f"{root}_Gemm"
node for node in names if node == f"{root}_Gemm"
)
elif block == "text_model" or simple:
match = next(
node
for node in remaining
if fix_node_name(node.name) == f"{root}_MatMul"
for node in names
if node == f"{root}_MatMul"
)
else:
# search in order. one side has sparse indices, so they will not match.
match = next(
node
for node in remaining
if node.name.startswith(f"/{block}")
and fix_node_name(node.name).endswith(
for node in names
if node.startswith(block)
and node.endswith(
f"{suffix}_MatMul"
) # needs to be fixed because some places use to_out.0
)
@ -158,10 +158,9 @@ def fix_xl_names(keys: Dict[str, Any], nodes: List[NodeProto]) -> Dict[str, Any]
logger.warning("no matches for XL key: %s", root)
continue
else:
logger.trace("matched key: %s -> %s", key, match.name)
logger.trace("matched key: %s -> %s", key, match)
name: str = match.name
name = fix_node_name(name)
name = match
if name.endswith("_MatMul"):
name = name[:-7]
elif name.endswith("_Gemm"):
@ -169,15 +168,15 @@ def fix_xl_names(keys: Dict[str, Any], nodes: List[NodeProto]) -> Dict[str, Any]
elif name.endswith("_Conv"):
name = name[:-5]
logger.trace("matching XL key with node: %s -> %s, %s", key, match.name, name)
logger.trace("matching XL key with node: %s -> %s, %s", key, match, name)
fixed[name] = value
remaining.remove(match)
names.remove(match)
logger.debug(
"SDXL LoRA key fixup matched %s keys, %s remaining",
len(fixed.keys()),
len(remaining),
len(names),
)
return fixed
@ -487,7 +486,7 @@ def blend_loras(
blended[base_key] = np_weights
# rewrite node names for XL and flatten layers
weights = Dict[str, np.ndarray] = {}
weights: Dict[str, np.ndarray] = {}
for blended in layers:
if xl: