Update nodes.py

This commit is contained in:
kijai 2024-09-01 20:40:00 +03:00
parent 181f153120
commit 7aa591b3a6

View File

@ -1865,19 +1865,20 @@ class FluxBlockLoraLoader:
temp = self.loaded_lora temp = self.loaded_lora
self.loaded_lora = None self.loaded_lora = None
del temp del temp
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" # Find the first key that ends with "weight"
weight_key = next((key for key in lora.keys() if key.endswith('weight')), None) rank = "unknown"
# Print the shape of the value corresponding to the key weight_key = next((key for key in lora.keys() if key.endswith('weight')), None)
if weight_key: # Print the shape of the value corresponding to the key
print(f"Shape of the first 'weight' key ({weight_key}): {lora[weight_key].shape}") if weight_key:
rank = str(lora[weight_key].shape[0]) print(f"Shape of the first 'weight' key ({weight_key}): {lora[weight_key].shape}")
else: rank = str(lora[weight_key].shape[0])
print("No key ending with 'weight' found.") else:
rank = "Couldn't find rank" print("No key ending with 'weight' found.")
self.loaded_lora = (lora_path, lora) rank = "Couldn't find rank"
self.loaded_lora = (lora_path, lora)
key_map = {} key_map = {}
if model is not None: if model is not None: