Update A.py

A.py

@@ -34,7 +34,7 @@ def resize_tensor_shapes(tensor1, tensor2):
 
     return tensor1_resized, tensor2_resized
 
-def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.75):
+def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.6):
     print(f"Merging checkpoints with blend ratio: {blend_ratio}")
     merged = {}
     all_keys = set(ckpt1.keys()).union(set(ckpt2.keys()))
@@ -49,7 +49,7 @@ def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.75):
         else:
             merged[key] = t2
 
-    # Control the final size to be approximately 26 GB
+    # Control the final size to be strictly 26 GB
     control_output_size(merged, target_size_gb=26)
 
     return merged
@@ -64,14 +64,20 @@ def control_output_size(merged, target_size_gb):
         excess_size = current_size_bytes - target_size_bytes
         print(f"Current size exceeds target by {excess_size / (1024**2):.2f} MB. Adjusting...")
         
-        # Adjusting the tensors to meet the target size
+        # Calculate the total number of elements to reduce
+        elements_to_reduce = excess_size // 4  # Assuming 4 bytes per float32 tensor
+        total_elements = sum(tensor.numel() for tensor in merged.values())
+
+        # Distribute the reduction uniformly across all tensors
         for key in merged.keys():
             tensor = merged[key]
-            # Calculate how much we can reduce
-            reduce_size = excess_size // tensor.element_size()  # Number of elements to reduce
-            if tensor.numel() > reduce_size:
-                # Truncate the tensor
-                merged[key] = tensor.flatten()[:tensor.numel() - reduce_size].view(tensor.shape)
+            num_elements = tensor.numel()
+            # Calculate how much to reduce from this tensor
+            reduction = min(elements_to_reduce, num_elements)
+            merged[key] = tensor.flatten()[:num_elements - reduction].view(tensor.shape)
+            elements_to_reduce -= reduction
+            if elements_to_reduce <= 0:
+                break
 
 def cleanup_files(*file_paths):
     for file_path in file_paths:
@@ -83,8 +89,8 @@ if __name__ == "__main__":
     try:
         model1_path = "mangledMergeFlux_v0Bfloat16Dev.safetensors"
         model2_path = "output_checkpoint.safetensors"
-        blend_ratio = 0.75  # Adjust ratio based on requirement
-        output_file = "output_checkpoints.safetensors"
+        blend_ratio = 0.6  # Set to 60%
+        output_file = "output_checkpoint.safetensors"
 
         # Loading models
         model1 = load_model(model1_path)