Upload combine.py

combine.py

@@ -0,0 +1,51 @@
+import os
+import json
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+def merge_safetensors(input_dir, output_file, config_file):
+    # Dictionary to store all tensors
+    merged_tensors = {}
+
+    # Load config
+    with open(config_file, 'r') as f:
+        config = json.load(f)
+
+    # Prepare metadata
+    metadata = {
+        "format": "pt",
+        "total_size": "", #str(total_size),  # Notice we stringify this!
+        "_diffusers_version": config.get("_diffusers_version", ""),
+        "_class_name": config.get("_class_name", ""),
+        # Add other fields at this level
+    }
+
+    total_size = 0
+
+    # Iterate through all files in the input directory
+    for filename in os.listdir(input_dir):
+        if filename.endswith('.safetensors'):
+            file_path = os.path.join(input_dir, filename)
+
+            # Load tensors and metadata from each file
+            with safe_open(file_path, framework="pt", device="cpu") as f:
+                file_metadata = f.metadata()
+                if file_metadata and "__metadata__" in file_metadata:
+                    total_size += int(file_metadata["__metadata__"].get("total_size", 0))
+
+                for key in f.keys():
+                    tensor = f.get_tensor(key)
+                    merged_tensors[key] = tensor
+
+    # Add total size to metadata
+    metadata["total_size"] = str(total_size)
+
+    # Save the merged tensors to a single file with metadata
+    save_file(merged_tensors, output_file, metadata)
+
+
+input_directory = './10_1'
+output_file = './10_1/flux1-merge-S10_D1.safetensors'
+config_file = './10_1/config.json'
+merge_safetensors(input_directory, output_file, config_file)