From 7aa591b3a64a3f83ec2c3e92758d0bb0926a6fe0 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 1 Sep 2024 20:40:00 +0300
Subject: [PATCH] Update nodes.py

---
 nodes/nodes.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/nodes/nodes.py b/nodes/nodes.py
index 9fff99f..d84947d 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -1865,19 +1865,20 @@ class FluxBlockLoraLoader:
                 temp = self.loaded_lora
                 self.loaded_lora = None
                 del temp
-
+        
         if lora is None:
             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)
+        rank = "unknown"
+        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)
 
         key_map = {}
         if model is not None: