Script to merge safetensor part files into a single file

import json
import safetensors
import safetensors.torch
import torch
import os

def merge_safetensors_shards(index_file_path, output_file_path):
    """
    Merges safetensors shards into a single safetensors file.

    Args:
        index_file_path (str): Path to the model.safetensors.index.json file.
        output_file_path (str): Path to save the merged safetensors file.
    """

    try:
        with open(index_file_path, 'r') as f:
            index_data = json.load(f)
    except FileNotFoundError:
        print(f"Error: Index file not found at {index_file_path}")
        return

    shard_file_directory = os.path.dirname(index_file_path)
    merged_weights = {}
    total_size = 0

    print("Starting to merge safetensors shards...")

    for weight_name, shard_file in index_data['weight_map'].items():
        shard_path = os.path.join(shard_file_directory, shard_file)
        try:
            with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
                tensor = f.get_tensor(weight_name)
                merged_weights[weight_name] = tensor
                total_size += tensor.numel() * tensor.element_size()  # Calculate size in bytes
            print(f"Loaded weight '{weight_name}' from shard '{shard_file}'")
        except FileNotFoundError:
            print(f"Error: Shard file not found at {shard_path}")
            return

    print(f"All shards loaded. Total size of merged weights: {total_size} bytes")
    print(f"Saving merged safetensors file to: {output_file_path}")

    try:
        safetensors.torch.save_file(merged_weights, output_file_path, metadata={"format": "pt"})
        print(f"Successfully merged safetensors shards to {output_file_path}.")
    except Exception as e:
        print(f"Error saving merged safetensors file: {e}")


if __name__ == "__main__":
    index_file_path = "diffusion_pytorch_model.safetensors.index.json" # index.json file from the repo
    output_file_path = "your_path/merged_diffusion_pytorch_model.safetensors" # You can change the output name if needed

    merge_safetensors_shards(index_file_path, output_file_path)

    print("\nMerge process completed.")