Upload lora_dims.py

lora_dims.py

@@ -0,0 +1,91 @@
+import os
+import argparse
+from safetensors import safe_open
+
+def get_lora_dimensions_from_directory(directory_path):
+    """
+    Scans a directory for .safetensors files and extracts the LoRA network dimension
+    from their metadata, falling back to inspecting tensor shapes if metadata is absent.
+
+    Args:
+        directory_path (str): The path to the directory to scan.
+    """
+    print(f"Scanning for LoRA models in: '{directory_path}'...\n")
+    
+    found_models = 0
+    
+    # Walk through the directory and its subdirectories
+    for root, _, files in os.walk(directory_path):
+        for filename in sorted(files):
+            if filename.lower().endswith(".safetensors"):
+                file_path = os.path.join(root, filename)
+                try:
+                    # Use safe_open to read metadata without loading the whole file
+                    with safe_open(file_path, framework="pt", device="cpu") as f:
+                        metadata = f.metadata()
+                        
+                        if not metadata:
+                            print(f"- {filename}: No metadata found. Checking weights...")
+                            # Fallthrough to weight checking
+                        
+                        # LoRA training scripts like Kohya's SS store the dimension here
+                        network_dim = metadata.get("ss_network_dim") if metadata else None
+                        
+                        if network_dim:
+                            print(f"- {filename}: Dimension = {network_dim} (from metadata)")
+                            found_models += 1
+                        else:
+                            # Fallback: try to determine dimension from tensor shapes
+                            dim_from_weights = None
+                            for key in f.keys():
+                                # Typically, the rank is the first dimension of the 'lora_down' tensor
+                                if key.endswith("lora_down.weight"):
+                                    tensor = f.get_tensor(key)
+                                    # The shape of lora_down.weight is (rank, in_features)
+                                    dim_from_weights = tensor.shape[0]
+                                    break  # Found it, no need to check other keys
+                                
+                                # Alternative naming uses lora_B or lora_up for the up-projection
+                                if key.endswith(("lora_B.weight", "lora_up.weight")):
+                                    tensor = f.get_tensor(key)
+                                    # The shape of lora_up/lora_B is (out_features, rank)
+                                    dim_from_weights = tensor.shape[1]
+                                    break  # Found it, no need to check other keys
+                            
+                            if dim_from_weights is not None:
+                                print(f"- {filename}: Dimension = {dim_from_weights} (from weights)")
+                                found_models += 1
+                            else:
+                                print(f"- {filename}: (Could not determine dimension from metadata or weights)")
+
+                except Exception as e:
+                    print(f"Could not process {filename}. Error: {e}")
+
+    if found_models == 0:
+        print("\nNo LoRA models with dimension information were found in the specified directory.")
+    else:
+        print(f"\nScan complete. Found {found_models} models with dimension info.")
+
+if __name__ == "__main__":
+    # Set up command-line argument parsing
+    parser = argparse.ArgumentParser(
+        description="Get the network dimensions of LoRA models in a directory.",
+        formatter_class=argparse.RawTextHelpFormatter
+    )
+    
+    parser.add_argument(
+        "directory",
+        type=str,
+        help="The path to the directory containing your LoRA (.safetensors) files."
+    )
+
+    args = parser.parse_args()
+
+    # Check if the provided path is a valid directory
+    if not os.path.isdir(args.directory):
+        print(f"Error: The path '{args.directory}' is not a valid directory.")
+    else:
+        get_lora_dimensions_from_directory(args.directory)
+
+
+