機械学習モデルのファイル形式と、ComfyUIでの利用可能性

機械学習モデルのファイル形式と、ComfyUIでの利用可能性

近年のAIの発展は目覚ましく、日夜新たなオープンモデルが公開され、その派生モデルが個人勢から公開されている。
そのファイル形式には複数の形式があるため、本記事では、特にマルチモーダル系のモデルに使われることが多い形式についてまとめる。
また、個人的によく使うComfyUIでの対応状況についても現状を把握する。

一般的なファイル形式とその特徴

.pth, .pt, .bin

• 概要
  • PyTorchの標準的な形式
• 中身
  • PyTorchのstate_dict等が含まれるpickleデータ
• 特徴
  • 比較的古くから使われる
  • 任意コードの実行リスク有

.safetensors

• 概要
  • Hugging Face主導の比較的安全な形式
• 中身
  • 重みとメタデータ
• 特徴
  • 近年一般的な形式
  • .pth系と比べ高速化

.ckpt

• 概要
  • 学習時のデータを含む包括的なデータ（であることが多い）
  • ファイルの中身は、.pth, .pt等と同じ(pickle)であるが、オプティマイザやエポック数などの学習時の情報も含んでいることを明示する際に使われることが多い
• 中身
  • state_dictの他、エポック数等の学習時のデータ
• 特徴
  • 学習時のスナップショットとなる
  • 学習時のバックアップや再開時に利用

.onnx

• 概要
  • フレームワーク非依存の中間表現
• 中身
  • 計算グラフや重み、メタデータなどの完全な構造
• 特徴
  • PyTorch / TensorFlow / JAX等複数の環境で読み込み可能

.gguf (旧: .ggml)

• 概要
  • 量子化に注力した形式
• 中身
  • 重みやアーキテクチャ、メタデータ
• 特徴
  • 様々な量子化レベルに対応し、VRAM制限のあるローカルでの実行がしやすい
  • 推論に特化
    • gguf形式のモデルでのファインチューニングなどは他形式と比べ環境が整っていない
  • llama.cpp, Ollamaでの利用を想定して作成された形式
  • 過去にコード実行の脆弱性があった模様（reddit）

そのほか、iOSやmacOS系を想定した.mlmodel, モバイル向け.tfliteが存在する。

ComfyUIでの対応状況

ComfyUIはPyTorchを利用しているため、基本的にPyTorchで読み込めるファイル形式であれば対応している。そのため、.pth, .pt, .bin, .safetensors, .ckpt形式については標準で利用できる。
.ggufについては、カスタムノードを入れると利用可能（ComfyUI-GGUF）。
.onnxに関しては、カスタムノードの独自実装により、一部のモデルが利用可能である（例1、例2、例3）。