xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2025-12-10 05:15:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update nodes.py

Browse Source
This commit is contained in:
kijai 2024-09-01 20:37:18 +03:00
parent a219601b2a
commit 181f153120
1 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
nodes/nodes.py
View File

@ -1842,8 +1842,9 @@ class FluxBlockLoraLoader:
} }
} }
RETURN_TYPES = ("MODEL", ) RETURN_TYPES = ("MODEL", "STRING", )
OUTPUT_TOOLTIPS = ("The modified diffusion model.",) RETURN_NAMES = ("model", "rank", )
OUTPUT_TOOLTIPS = ("The modified diffusion model.", "possible rank of the LoRA.")
FUNCTION = "load_lora" FUNCTION = "load_lora"
CATEGORY = "KJNodes/experimental" CATEGORY = "KJNodes/experimental"
@ -1867,6 +1868,15 @@ class FluxBlockLoraLoader:
if lora is None: if lora is None:
lora = load_torch_file(lora_path, safe_load=True) lora = load_torch_file(lora_path, safe_load=True)
# Find the first key that ends with "weight"
weight_key = next((key for key in lora.keys() if key.endswith('weight')), None)
# Print the shape of the value corresponding to the key
if weight_key:
print(f"Shape of the first 'weight' key ({weight_key}): {lora[weight_key].shape}")
rank = str(lora[weight_key].shape[0])
else:
print("No key ending with 'weight' found.")
rank = "Couldn't find rank"
self.loaded_lora = (lora_path, lora) self.loaded_lora = (lora_path, lora)
key_map = {} key_map = {}
@ -1932,4 +1942,4 @@ class FluxBlockLoraLoader:
if (x not in k): if (x not in k):
print("NOT LOADED {}".format(x)) print("NOT LOADED {}".format(x))
return (new_modelpatcher,) return (new_modelpatcher, rank)