Update rp.py

rp.py

@@ -1,26 +1,55 @@
-import torch
 from safetensors.torch import load_file, save_file
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm  # Ensure tqdm is installed
+
+def load_model(file_path):
+    return load_file(file_path)
+
+def save_model(merged_model, output_file):
+    print(f"Saving merged model to {output_file}")
+    save_file(merged_model, output_file)
+
+def resize_tensor_shapes(tensor1, tensor2):
+    if tensor1.size() == tensor2.size():
+        return tensor1, tensor2
+
+    max_shape = [max(s1, s2) for s1, s2 in zip(tensor1.shape, tensor2.shape)]
+    tensor1_resized = F.pad(tensor1, (0, max_shape[-1] - tensor1.size(-1)))
+    tensor2_resized = F.pad(tensor2, (0, max_shape[-1] - tensor2.size(-1)))
+
+    return tensor1_resized, tensor2_resized
+
+def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.5):
+    print(f"Merging checkpoints with blend ratio: {blend_ratio}")
+    merged = {}
+    all_keys = set(ckpt1.keys()).union(set(ckpt2.keys()))
+
+    for key in tqdm(all_keys, desc="Merging Checkpoints", unit="layer"):
+        t1, t2 = ckpt1.get(key), ckpt2.get(key)
+        if t1 is not None and t2 is not None:
+            t1, t2 = resize_tensor_shapes(t1, t2)
+            merged[key] = blend_ratio * t1 + (1 - blend_ratio) * t2
+        elif t1 is not None:
+            merged[key] = t1
+        else:
+            merged[key] = t2
+
+    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
+
+    # Load the models
+    model1 = load_model(model1_path)
+    model2 = load_model(model2_path)
+
+    # Merge the models
+    merged_model = merge_checkpoints(model1, model2, blend_ratio)
 
-def reduce_key_size(input_file, output_file, reduction_factor=0.50):
-    # Load the model
-    model_data = load_file(input_file)
-
-    # Iterate through all the tensors and reduce their size
-    for key in model_data.keys():
-        original_tensor = model_data[key]
-        
-        # Calculate the new size
-        new_size = int(original_tensor.size(0) * (1 - reduction_factor))
-        
-        # Resize the tensor (this could vary depending on your requirements)
-        if new_size > 0:  # Ensure new size is positive
-            reduced_tensor = original_tensor[:new_size]
-            model_data[key] = reduced_tensor
-    
-    # Save the modified model
-    save_file(model_data, output_file)
-
-# Usage example
-input_file = 'model-00002-of-00002.safetensors'  # Replace with your input model file
-output_file = 'model-00002-of-00002.safetensors'  # Desired output file name
-reduce_key_size(input_file, output_file)
+    # Save the merged model
+    save_model(merged_model, output_file)