Update A.py

A.py

@@ -27,13 +27,14 @@ def resize_tensor_shapes(tensor1, tensor2):
     if tensor1.size() == tensor2.size():
         return tensor1, tensor2
 
+    # Resize tensor2 to match tensor1's size (Base size)
     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):
+def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.75):
     print(f"Merging checkpoints with blend ratio: {blend_ratio}")
     merged = {}
     all_keys = set(ckpt1.keys()).union(set(ckpt2.keys()))
@@ -48,8 +49,30 @@ def merge_checkpoints(ckpt1, ckpt2, blend_ratio=0.5):
         else:
             merged[key] = t2
 
+    # Control the final size to be approximately 26 GB
+    control_output_size(merged, target_size_gb=26)
+
     return merged
 
+def control_output_size(merged, target_size_gb):
+    # Estimate the size in bytes
+    target_size_bytes = target_size_gb * 1024**3  # Convert GB to bytes
+    current_size_bytes = sum(tensor.numel() * tensor.element_size() for tensor in merged.values())
+
+    # If the current size exceeds the target, truncate the tensors
+    if current_size_bytes > target_size_bytes:
+        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
+        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)
+
 def cleanup_files(*file_paths):
     for file_path in file_paths:
         if os.path.exists(file_path):
@@ -58,14 +81,10 @@ def cleanup_files(*file_paths):
 
 if __name__ == "__main__":
     try:
-        
-        model1_path = "flux1-dev.safetensors"
-        model2_path = "brainflux_v10.safetensors"
-        blend_ratio = 0.4
-        output_file = "output_checkpoint.safetensors"
-
-        # Downloading files
-        
+        model1_path = "mangledMergeFlux_v0Bfloat16Dev.safetensors"
+        model2_path = "output_checkpoint.safetensors"
+        blend_ratio = 0.75  # Adjust ratio based on requirement
+        output_file = "output_checkpoints.safetensors"
 
         # Loading models
         model1 = load_model(model1_path)