Upload 4 files

comparemodels.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+
+import argparse
+import torch
+from safetensors import safe_open
+
+
+def compare_safetensors(filepath1: str, filepath2: str):
+    """
+    Compares two .safetensors files, ignoring a specific prefix on layer names,
+    and prints a summary of the differences.
+
+    Args:
+        filepath1 (str): Path to the first .safetensors file.
+        filepath2 (str): Path to the second .safetensors file.
+    """
+    # The prefix to ignore on layer names
+    prefix_to_ignore = "model.diffusion_model."
+
+    # Dictionaries to hold results
+    results = {
+        "only_in_file1": [],
+        "only_in_file2": [],
+        "different_content": [],
+    }
+
+    print("\nLoading files and preparing for comparison...")
+    print(f"Ignoring prefix: '{prefix_to_ignore}'")
+
+    try:
+        # Use 'with' to ensure files are closed properly
+        with safe_open(filepath1, framework="pt", device="cpu") as f1, \
+                safe_open(filepath2, framework="pt", device="cpu") as f2:
+
+            # Create maps from the normalized key (suffix) to the original key
+            map1 = {key.removeprefix(prefix_to_ignore): key for key in f1.keys()}
+            map2 = {key.removeprefix(prefix_to_ignore): key for key in f2.keys()}
+
+            # Get the set of normalized tensor keys from each file
+            normalized_keys1 = set(map1.keys())
+            normalized_keys2 = set(map2.keys())
+
+            # 1. Find normalized keys (layers) unique to each file
+            results["only_in_file1"] = sorted(list(normalized_keys1 - normalized_keys2))
+            results["only_in_file2"] = sorted(list(normalized_keys2 - normalized_keys1))
+
+            # 2. Find normalized keys present in both files to compare their content
+            common_normalized_keys = normalized_keys1.intersection(normalized_keys2)
+            print(f"Comparing {len(common_normalized_keys)} common tensors...")
+
+            for norm_key in sorted(list(common_normalized_keys)):
+                # Get the original key for each file using the maps
+                original_key1 = map1[norm_key]
+                original_key2 = map2[norm_key]
+
+                # Get the tensor from each file using its original key
+                tensor1 = f1.get_tensor(original_key1)
+                tensor2 = f2.get_tensor(original_key2)
+
+                # Compare tensors for equality
+                if not torch.equal(tensor1, tensor2):
+                    # Store the normalized key if content differs
+                    results["different_content"].append(norm_key)
+
+        # --- Print the results ---
+        print("\n" + "=" * 60)
+        print("🔍 Safetensor Comparison Results")
+        print("=" * 60)
+        print(f"File 1: {filepath1}")
+        print(f"File 2: {filepath2}")
+        print("-" * 60)
+
+        # Check if any differences were found at all
+        total_diffs = len(results["only_in_file1"]) + len(results["only_in_file2"]) + len(results["different_content"])
+        if total_diffs == 0:
+            print("\n✅ The files are identical after normalization. No differences found.")
+            print("=" * 60 + "\n")
+            return
+
+        # Report tensors with different content
+        if results["different_content"]:
+            print(f"\n↔️ Tensors with Different Content ({len(results['different_content'])}):")
+            for norm_key in results["different_content"]:
+                print(f"  - Normalized Key: {norm_key}")
+                print(f"    (File 1 Original: {map1[norm_key]})")
+                print(f"    (File 2 Original: {map2[norm_key]})")
+
+        # Report tensors only in file 1
+        if results["only_in_file1"]:
+            print(f"\n→ Tensors Only in File 1 ({len(results['only_in_file1'])}):")
+            for norm_key in results["only_in_file1"]:
+                print(f"  - Normalized Key: {norm_key} (Original: {map1[norm_key]})")
+
+        # Report tensors only in file 2
+        if results["only_in_file2"]:
+            print(f"\n← Tensors Only in File 2 ({len(results['only_in_file2'])}):")
+            for norm_key in results["only_in_file2"]:
+                print(f"  - Normalized Key: {norm_key} (Original: {map2[norm_key]})")
+
+        print("\n" + "=" * 60 + "\n")
+
+    except FileNotFoundError as e:
+        print(f"❌ Error: Could not find a file. Details: {e}")
+    except Exception as e:
+        print(f"❌ An error occurred: {e}")
+        print("Please ensure both files are valid .safetensors files.")
+
+
+if __name__ == "__main__":
+    # --- Argument Parser Setup ---
+    parser = argparse.ArgumentParser(
+        description="Compares two .safetensors files and lists the differences in their layers (tensors), ignoring a specific prefix.",
+        formatter_class=argparse.RawTextHelpFormatter
+    )
+
+    parser.add_argument(
+        "file1",
+        type=str,
+        help="Path to the first .safetensors file."
+    )
+    parser.add_argument(
+        "file2",
+        type=str,
+        help="Path to the second .safetensors file."
+    )
+
+    args = parser.parse_args()
+
+    # --- Run the function ---
+    compare_safetensors(args.file1, args.file2)

extract_layers.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+
+import argparse
+from safetensors import safe_open
+from safetensors.torch import save_file
+from tqdm import tqdm
+
+# This set contains the specific layers we want to find and save.
+# Using a set is efficient for checking if a layer should be kept.
+KEYS_TO_KEEP = {
+    "distilled_guidance_layer.in_proj.bias",
+    "distilled_guidance_layer.in_proj.weight",
+    "distilled_guidance_layer.layers.0.in_layer.bias",
+    "distilled_guidance_layer.layers.0.in_layer.weight",
+    "distilled_guidance_layer.layers.0.out_layer.bias",
+    "distilled_guidance_layer.layers.0.out_layer.weight",
+    "distilled_guidance_layer.layers.1.in_layer.bias",
+    "distilled_guidance_layer.layers.1.in_layer.weight",
+    "distilled_guidance_layer.layers.1.out_layer.bias",
+    "distilled_guidance_layer.layers.1.out_layer.weight",
+    "distilled_guidance_layer.layers.2.in_layer.bias",
+    "distilled_guidance_layer.layers.2.in_layer.weight",
+    "distilled_guidance_layer.layers.2.out_layer.bias",
+    "distilled_guidance_layer.layers.2.out_layer.weight",
+    "distilled_guidance_layer.layers.3.in_layer.bias",
+    "distilled_guidance_layer.layers.3.in_layer.weight",
+    "distilled_guidance_layer.layers.3.out_layer.bias",
+    "distilled_guidance_layer.layers.3.out_layer.weight",
+    "distilled_guidance_layer.layers.4.in_layer.bias",
+    "distilled_guidance_layer.layers.4.in_layer.weight",
+    "distilled_guidance_layer.layers.4.out_layer.bias",
+    "distilled_guidance_layer.layers.4.out_layer.weight",
+    "distilled_guidance_layer.norms.0.scale",
+    "distilled_guidance_layer.norms.1.scale",
+    "distilled_guidance_layer.norms.2.scale",
+    "distilled_guidance_layer.norms.3.scale",
+    "distilled_guidance_layer.norms.4.scale",
+    "distilled_guidance_layer.out_proj.bias",
+    "distilled_guidance_layer.out_proj.weight",
+}
+
+def extract_safetensors_layers(input_path: str, output_path: str, keys_to_keep: set):
+    """
+    Reads a safetensors file, extracts specified layers, and saves them to a new file.
+    """
+    tensors_to_save = {}
+    
+    print(f"▶️  Reading from: {input_path}")
+    try:
+        # Open the source file for reading
+        with safe_open(input_path, framework="pt", device="cpu") as f:
+            
+            # Iterate through the keys we want to keep
+            for key in tqdm(keys_to_keep, desc="Extracting Tensors"):
+                try:
+                    # If the key exists in the file, add its tensor to our dictionary
+                    tensors_to_save[key] = f.get_tensor(key)
+                except KeyError:
+                    print(f"\n⚠️  Warning: Key not found in source file and will be skipped: {key}")
+
+        print("\n--- Summary ---")
+        print(f"Layers specified to keep: {len(keys_to_keep)}")
+        print(f"Layers found and extracted: {len(tensors_to_save)}")
+
+        if not tensors_to_save:
+            print("\n❌ Error: No matching layers were found. The output file will not be created.")
+            return
+
+        # Save the extracted tensors to the new file
+        print(f"\n▶️  Saving {len(tensors_to_save)} tensors to: {output_path}")
+        save_file(tensors_to_save, output_path)
+
+        print("\n✅ Extraction complete!")
+
+    except FileNotFoundError:
+        print(f"\n❌ Error: The file '{input_path}' was not found.")
+    except Exception as e:
+        print(f"\n❌ An error occurred: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Extracts a predefined list of layers from a .safetensors file and saves them to a new file.",
+        formatter_class=argparse.RawTextHelpFormatter
+    )
+
+    parser.add_argument(
+        "input_file",
+        type=str,
+        help="Path to the source .safetensors file."
+    )
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help="Path to save the new, extracted .safetensors file."
+    )
+
+    args = parser.parse_args()
+    
+    extract_safetensors_layers(args.input_file, args.output_file, KEYS_TO_KEEP)

listlayers.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+
+import argparse
+from safetensors import safe_open
+
+def list_safetensor_layers(filepath: str):
+    """
+    Opens a .safetensors file and prints the name and shape of each tensor.
+
+    Args:
+        filepath (str): The path to the .safetensors file.
+    """
+    try:
+        print(f"\n📄 Tensors in: {filepath}\n" + "="*50)
+        
+        total_tensors = 0
+        with safe_open(filepath, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                tensor = f.get_tensor(key)
+                print(f"- {key:<50} | Shape: {tensor.shape}")
+                total_tensors += 1
+        
+        print("="*50 + f"\n✅ Found {total_tensors} total tensors.\n")
+
+    except FileNotFoundError:
+        print(f"❌ Error: The file '{filepath}' was not found.")
+    except Exception as e:
+        print(f"❌ An error occurred: {e}")
+        print("Please ensure the file is a valid .safetensors file.")
+
+if __name__ == "__main__":
+    # --- Argument Parser Setup ---
+    parser = argparse.ArgumentParser(
+        description="List all layers (tensors) and their shapes in a .safetensors file.",
+        formatter_class=argparse.RawTextHelpFormatter
+    )
+
+    parser.add_argument(
+        "filepath",
+        type=str,
+        help="Path to the .safetensors file."
+    )
+
+    args = parser.parse_args()
+    
+    # --- Run the function ---
+    list_safetensor_layers(args.filepath)

prune_layers.py

@@ -0,0 +1,238 @@
+#!/usr/bin/env python
+
+import argparse
+from safetensors import safe_open
+from safetensors.torch import save_file
+from tqdm import tqdm
+
+# This is the set of layers that will be removed from the file.
+# Using a set provides fast lookups.
+KEYS_TO_REMOVE = {
+    "model.diffusion_model.double_blocks.0.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.0.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.0.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.0.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.1.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.1.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.1.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.1.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.10.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.10.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.10.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.10.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.11.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.11.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.11.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.11.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.12.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.12.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.12.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.12.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.13.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.13.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.13.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.13.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.14.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.14.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.14.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.14.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.15.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.15.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.15.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.15.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.16.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.16.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.16.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.16.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.17.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.17.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.17.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.17.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.18.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.18.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.18.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.18.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.2.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.2.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.2.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.2.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.3.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.3.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.3.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.3.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.4.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.4.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.4.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.4.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.5.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.5.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.5.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.5.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.6.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.6.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.6.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.6.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.7.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.7.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.7.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.7.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.8.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.8.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.8.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.8.txt_mod.lin.weight",
+    "model.diffusion_model.double_blocks.9.img_mod.lin.bias",
+    "model.diffusion_model.double_blocks.9.img_mod.lin.weight",
+    "model.diffusion_model.double_blocks.9.txt_mod.lin.bias",
+    "model.diffusion_model.double_blocks.9.txt_mod.lin.weight",
+    "model.diffusion_model.final_layer.adaLN_modulation.1.bias",
+    "model.diffusion_model.final_layer.adaLN_modulation.1.weight",
+    # "model.diffusion_model.guidance_in.in_layer.bias",
+    # "model.diffusion_model.guidance_in.in_layer.weight",
+    # "model.diffusion_model.guidance_in.out_layer.bias",
+    # "model.diffusion_model.guidance_in.out_layer.weight",
+    "model.diffusion_model.single_blocks.0.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.0.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.1.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.1.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.10.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.10.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.11.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.11.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.12.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.12.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.13.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.13.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.14.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.14.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.15.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.15.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.16.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.16.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.17.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.17.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.18.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.18.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.19.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.19.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.2.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.2.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.20.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.20.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.21.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.21.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.22.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.22.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.23.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.23.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.24.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.24.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.25.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.25.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.26.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.26.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.27.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.27.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.28.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.28.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.29.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.29.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.3.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.3.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.30.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.30.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.31.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.31.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.32.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.32.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.33.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.33.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.34.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.34.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.35.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.35.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.36.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.36.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.37.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.37.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.4.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.4.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.5.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.5.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.6.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.6.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.7.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.7.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.8.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.8.modulation.lin.weight",
+    "model.diffusion_model.single_blocks.9.modulation.lin.bias",
+    "model.diffusion_model.single_blocks.9.modulation.lin.weight",
+    # "model.diffusion_model.time_in.in_layer.bias",
+    # "model.diffusion_model.time_in.in_layer.weight",
+    # "model.diffusion_model.time_in.out_layer.bias",
+    # "model.diffusion_model.time_in.out_layer.weight",
+    # "model.diffusion_model.vector_in.in_layer.bias",
+    # "model.diffusion_model.vector_in.in_layer.weight",
+    # "model.diffusion_model.vector_in.out_layer.bias",
+    # "model.diffusion_model.vector_in.out_layer.weight",
+}
+
+def prune_safetensors_file(input_path: str, output_path: str, keys_to_remove: set):
+    """
+    Reads a safetensors file, removes specified layers, and saves to a new file.
+    """
+    tensors_to_keep = {}
+    original_count = 0
+
+    print(f"▶️  Reading from: {input_path}")
+    try:
+        # Open the source file for reading
+        with safe_open(input_path, framework="pt", device="cpu") as f:
+            all_keys = f.keys()
+            original_count = len(all_keys)
+
+            # Iterate through all tensors, keeping only the ones not in the removal list
+            for key in tqdm(all_keys, desc="Filtering Tensors"):
+                if 'modulation' not in key :
+                    tensors_to_keep[key] = f.get_tensor(key)
+
+        removed_count = original_count - len(tensors_to_keep)
+        
+        print("\n--- Summary ---")
+        print(f"Original tensor count: {original_count}")
+        print(f"Tensors removed: {removed_count}")
+        print(f"New tensor count: {len(tensors_to_keep)}")
+
+        if removed_count == 0 and len(keys_to_remove) > 0:
+            print("\n⚠️  Warning: None of the specified layers to remove were found in the input file.")
+
+        # Save the filtered tensors to the new file
+        print(f"\n▶️  Saving {len(tensors_to_keep)} tensors to: {output_path}")
+        save_file(tensors_to_keep, output_path)
+
+        print("\n✅ Pruning complete!")
+
+    except FileNotFoundError:
+        print(f"\n❌ Error: The file '{input_path}' was not found.")
+    except Exception as e:
+        print(f"\n❌ An error occurred: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Removes a predefined list of layers from a .safetensors file and saves it as a new file.",
+        formatter_class=argparse.RawTextHelpFormatter
+    )
+
+    parser.add_argument(
+        "input_file",
+        type=str,
+        help="Path to the source .safetensors file."
+    )
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help="Path to save the new, pruned .safetensors file."
+    )
+
+    args = parser.parse_args()
+    
+    prune_safetensors_file(args.input_file, args.output_file, KEYS_TO_REMOVE)