|
import os |
|
import math |
|
import json |
|
|
|
OUTPUT_FILE_NAME = "consolidated.safetensors" |
|
CHUNK_PATHS_FILE = "chunk_paths.json" |
|
|
|
def merge(chunk_paths): |
|
output_path = os.path.join(os.path.dirname(chunk_paths[0]), OUTPUT_FILE_NAME) |
|
|
|
with open(output_path, "wb") as f_out: |
|
for filepath in chunk_paths: |
|
with open(filepath, "rb") as f_in: |
|
f_out.write(f_in.read()) |
|
|
|
print(f"Merged file saved to {output_path}") |
|
|
|
if __name__ == "__main__": |
|
|
|
if os.path.exists(CHUNK_PATHS_FILE): |
|
with open(CHUNK_PATHS_FILE) as f: |
|
chunk_paths = json.load(f) |
|
else: |
|
chunk_paths = split(main_filepath) |
|
|
|
merge(chunk_paths) |