分割されたsafetensorsファイルを一つのファイルに結合

やりたいこと

HuggingFaceなどでModelをダウンロードした際に、model-00001-of-00003.safetensors、model-00002-of-00003.safetensors、model-00003-of-00003.safetensorsのように複数のファイルに分割されていることがあります。

このファイル形式だと使用できないアプリがある（StableDiffusion WebUI Forgeはできないようでした。）ので一つのファイルに結合する方法を紹介します。

注意

結合しても使用できないモデルもありました。やり方の問題なのか、もともとのソフトとモデルの問題なのかはわかりませんでした。

またモデルの編集にあたる可能性があるので自己責任で使用をお願いします。

準備

Pythonを以下のページからダウンロードしインストールします。

Download Python

The official home of the Python Programming Language

www.python.org

環境の作成

作業用フォルダを任意の場所に作成します。

コマンドプロンプトを開き以下のコマンドを実行します。（以下のコマンドは任意なのですが意味が分からない場合は実行を推奨です。）

python -mvenv venv
venv\scripts\activate.bat

コマンドラインのカレントフォルダのまえに(venv)とついていることを確認してください。

(venv)とついていることを確認したら以下のコマンドを実行します。

pip install safetensors
pip3 install torch torchvision torchaudio

環境構築はこれで完了ですが実行でも同じコマンドプロンプトを使用するので終了しないでください。

コードの作成

作業フォルダにcombine.pyを以下の内容で作成します。

import os
import torch
from safetensors.torch import save_file, load_file

def combine_safetensors(input_dir, output_file_name="combined_model.safetensors"):
    print(f"Searching for safetensors files in: {input_dir}")

    files_to_load = sorted([f for f in os.listdir(input_dir) if f.endswith('.safetensors')])

    if not files_to_load:
        print("No .safetensors files found in the specified directory.")
        return

    print(f"Found files: {files_to_load}")

    combined_state_dict = {}
    for i, filename in enumerate(files_to_load):
        file_path = os.path.join(input_dir, filename)
        print(f"Loading part {i+1}/{len(files_to_load)}: {filename}")
        try:
            state_dict_part = load_file(file_path)
            combined_state_dict.update(state_dict_part)
            del state_dict_part
        except Exception as e:
            print(f"Error loading {filename}: {e}")
            return

    output_path = os.path.join(input_dir, output_file_name)
    print(f"Saving combined model to: {output_path}")

    try:
        save_file(combined_state_dict, output_path)
        print("Combination successful!")
    except Exception as e:
        print(f"Error saving combined file: {e}")

if __name__ == "__main__":
    input_folder = input("INPUT FOLDER: ")

    output_name = input("OUTPUT FILE: ")
    if not output_name:
        output_name = "combined_model.safetensors"
    elif not output_name.endswith('.safetensors'):
        output_name += '.safetensors'

    combine_safetensors(input_folder, output_name)

実行

任意のフォルダに分割されたsafetensorsを置きます。（わかりやすいようにフォルダには対象のファイルのみ置くことがおすすめです。）

環境構築に使用したコマンドプロントで以下のコマンドを実行します。

Python combine.py

”INPUT FOLDER:”と表示されるので分割されたファイルを置いたフォルダを入力します。

続いて”OUTPUT FILE: “と表示されるので出力ファイル名を入力します。

しばらく待つとINPUT FOLDER内に結合されたファイルがOUTPUT FILEで指定した名前で保存されます。

結果

分割されたsafetensorsファイルを一つのファイルに結合することができました。