safetensorsファイルをpythonで使用する方法

最近、HuggingFace では、モデルの効率と安全性を向上させる新しいファイル形式である SafeTensors が使用されています。SafeTensors を使用してモデルを保存および読み込むことで、処理速度が向上し、悪意のあるコードを実行する可能性が低くなります。

今回は、anything-v5.0のsafetensorsファイルを取り扱い、pythonでAIイラスト生成を行ってみます。

実行環境はGoogle Colabとなります。

実行時は、下記設定で実行します。なぜかGPUがT4じゃないとうまくいかなかったです。何か理由があるのだろうけれど、理由を追えなかったです。

!pip install torch
!pip install --upgrade diffusers transformers scipy
!pip install datetime
!pip install safetensors
!pip install omegaconf
!pip install hydra-core --upgrade

from diffusers import StableDiffusionPipeline
import torch
import datetime

!wget https://huggingface.co/ckpt/anything-v5.0/resolve/main/AnythingV5V3_v5PrtRE.safetensors
pipe = StableDiffusionPipeline.from_ckpt(
    "./AnythingV5V3_v5PrtRE.safetensors"
)
pipe.to("cuda")

prompt = "(((super realistic))), (((best quality))),((masterpiece)), ((ultra-detailed)), a girl, smile, (((super realistic black hair))), shirt, black eyes, an anime style"

for i in range(10):
  image = pipe(prompt).images[0]
  timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
  image.save(f"/content/image_{timestamp}.png")

!wgetで取得したいsafetensorsファイルを指定します。

StableDiffusionPipeline.from_chkpt()内にダウンロードしたsafetensorsファイル名を指定します。