Update rp.py

rp.py

@@ -1,5 +1,6 @@
 from safetensors.torch import load_file, save_file
 import torch
+torch.cuda.empty_cache()
 import torch.nn.functional as F
 from tqdm import tqdm  # Ensure tqdm is installed
 
@@ -38,18 +39,18 @@ def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.5):
     return merged
 
 if __name__ == "__main__":
-    # Set your file paths and blend ratio here
-    model1_path = "flux1-dev.safetensors.1"  # Model 1 path
-    model2_path = "brainflux_v10.safetensors"  # Model 2 path
-    blend_ratio = 0.4  # Blend ratio
-    output_file = "output_checkpoint.safetensors"  # Output file name
+    try:
+        model1_path = "flux1-dev.safetensors.1"
+        model2_path = "brainflux_v10.safetensors"
+        blend_ratio = 0.4
+        output_file = "output_checkpoint.safetensors"
 
-    # Load the models
-    model1 = load_model(model1_path)
-    model2 = load_model(model2_path)
+        model1 = load_model(model1_path)
+        model2 = load_model(model2_path)
 
-    # Merge the models
-    merged_model = merge_checkpoints(model1, model2, blend_ratio)
+        merged_model = merge_checkpoints(model1, model2, blend_ratio)
 
-    # Save the merged model
-    save_model(merged_model, output_file)
+        save_model(merged_model, output_file)
+        
+    except Exception as e:
+        print(f"An error occurred: {e}")