Stable Diffusion web UIをnixで管理して実行できるようにしました
リポジトリ
ncaq/stable-diffusion-nix: Stable Diffusion WebUI managed by Nix
背景
かなり前にStable Diffusion WebUIを技術的な興味から触ってみたのですが、セットアップ手順でPython関係の依存環境を平然とpipなどでグローバルにインストールするようになっていたり、今後のアップデートに耐えられるとは到底思えませんでした。再インストールする時に手順を確認するのも面倒ですし。
Python界隈にありがちな悪い方法ですね。せめてpoetryなどで管理してくれれば良いのですが。割と活発なリポジトリなので色々と政治的な合意を取る必要がありそうなので、英語が苦手な自分は抜本的な改善に入るのは躊躇してしまいます。
有志の情報によるとDockerを使う方法が書いてありましたが、その方法は私の環境ではうまく動きませんでした。
そこでnixで管理してnixの環境があればnix runを実行するだけで実行できるようにしたら美しいと考えて構築してみることにしました。
そんなに自分で画像生成AIを使うわけではないのですが、インフラが整っていないところを整備したい欲求がありました。
大本のリポジトリ
stable-diffusion-webuiはなぜかmasterブランチで活発に開発が行われていない状態になっています。なのでdevブランチを参照することにしました。
inputs = {
nixpkgs.url = "github:NixOS/nixpkgs/nixos-25.05";
flake-parts.url = "github:hercules-ci/flake-parts";
treefmt-nix = {
url = "github:numtide/treefmt-nix";
inputs.nixpkgs.follows = "nixpkgs";
};
stable-diffusion-webui = {
url = "github:AUTOMATIC1111/stable-diffusion-webui/dev";
flake = false;
};
};
のように定義していればflake.lockでコミットのリビジョンは固定されるので、不安定さは許容できると考えています。
マイナーなPythonパッケージ
NixOS/nixpkgs: Nix Packages collection & NixOS
に登録されていないPythonパッケージがあるようなので、それらはこちらのflake.nixでパッケージを作成します。長期的にメンテナンスをするならばnixpkgsに登録すべきですが、これでちゃんと動いている保証がないので、まずはこのリポジトリで管理することにしました。
例えば以下のように定義しています。
blendmodes = pyFinal.buildPythonPackage rec {
pname = "blendmodes";
version = "2022";
src = pyFinal.fetchPypi {
inherit pname version;
sha256 = "sha256-k2jxwekOhzS0lo8MrANpbg5acU/VfnS8bu/SXAQY/QM=";
};
propagatedBuildInputs = with pyFinal; [
numpy
aenum
];
};
あまりちゃんとバージョン管理されていなくて、最新バージョンだと正常に動かないこともよくあるので、こちらで細かいバージョンを指定したいという思惑もあります。
一部はdoCheck = false;にしないと動かないのでテストをスキップできるようにしたいという意図もあります。
真面目にこれをメンテナンスするならば将来的には問題を解決してnixpkgsに登録するかもしれません。まだ実験段階なので許容していますが。
pydanticのバージョン不整合問題
stable-diffusion-webuiではpydanticのバージョンが1.x系でないと動かないのですが、それが依存しているfastapiなどのパッケージが間接的にpydanticの2.x系を要求しているため、バージョンのコンフリクトが発生してしまいます。
そこで以下のようにfastapiのパッケージ定義を上書きして、 pydantic v2への依存を削除してpydantic v1を使うようにしました。
# Override FastAPI to avoid pydantic v2.
fastapi = pyPrev.fastapi.overridePythonAttrs (_old: {
dependencies = with pyFinal; [
pydantic_1
starlette
typing-extensions
];
nativeCheckInputs =
with pyFinal;
[
anyio
dirty-equals
flask
inline-snapshot
passlib
pyjwt
pytest-asyncio
pytestCheckHook
sqlalchemy
trio
]
++ anyio.optional-dependencies.trio
++ passlib.optional-dependencies.bcrypt;
doCheck = false;
});
pydanticのバージョンが1.x系だとテストは動かないのでdoCheck = false;にしています。
そしてgradioの方も同じくpydantic_1に依存するようにパッケージ定義を上書きしています。
モンキーパッチ
真面目に追っていないので不可解なのですが、一部のコードはバージョンの不整合を解決するためにこちらで修正してやる必要があります。なので以下のようにパッチを適用する設定を書きます。
stable-diffusion-stability-ai = pkgs.stdenv.mkDerivation {
pname = "stable-diffusion-stability-ai";
version = "unstable-2022-11-23";
src = pkgs.fetchFromGitHub {
owner = "Stability-AI";
repo = "stablediffusion";
rev = "cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf";
sha256 = "sha256-yEtrz/JTq53JDI4NZI26KsD8LAgiViwiNaB2i1CBs/I=";
};
patches = [
./patch/01-pytorch-lightning-import-fix-repo.patch
];
patchFlags = [ "-p0" ];
installPhase = ''
mkdir -p $out
cp -r . $out/
'';
};
今現在あるpatchは以下のものです。
--- ldm/models/diffusion/ddpm.py
+++ ldm/models/diffusion/ddpm.py
@@ -17,7 +17,7 @@
import itertools
from tqdm import tqdm
from torchvision.utils import make_grid
-from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
from omegaconf import ListConfig
from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
--- modules/paths_internal.py
+++ modules/paths_internal.py
@@ -32,5 +32,5 @@ data_path = cmd_opts_pre.data_dir
models_path = cmd_opts_pre.models_dir if cmd_opts_pre.models_dir else os.path.join(data_path, "models")
extensions_dir = os.path.join(data_path, "extensions")
extensions_builtin_dir = os.path.join(script_path, "extensions-builtin")
-config_states_dir = os.path.join(script_path, "config_states")
+config_states_dir = os.path.join(data_path, "config_states")
default_output_dir = os.path.join(data_path, "outputs")
roboto_ttf_file = os.path.join(modules_path, 'Roboto-Regular.ttf')
パス周りはnixのパスになることで何か不整合が起きるのはわかるのですが、 PyTorchのimportが異なる問題はなぜ起きているのかよく分かりません。まだあまり真面目に調査はしていません。
インストールと実行コマンドを定義
もともとシェルスクリプトからPythonのエントリーポイントを起動するようになっているので、実行ファイルをこちらで定義する必要があります。
またいくつかのリポジトリをサブモジュール的にcloneしてくる必要があります。それも自動化します。
とりあえずは以下のようになっています。
buildPhase = ''
# Copy source
cp -r . $TMPDIR/webui
cd $TMPDIR/webui
# Copy repositories and make them writable
mkdir -p repositories
cp -r ${generative-models} repositories/generative-models
cp -r ${k-diffusion} repositories/k-diffusion
cp -r ${stable-diffusion-stability-ai} repositories/stable-diffusion-stability-ai
cp -r ${stable-diffusion-webui-assets} repositories/stable-diffusion-webui-assets
chmod -R u+w repositories
cd -
'';
installPhase = ''
mkdir -p $out/share/stable-diffusion-webui
cp -r $TMPDIR/webui/* $out/share/stable-diffusion-webui/
# Create wrapper script
mkdir -p $out/bin
cat > $out/bin/stable-diffusion-webui << EOF
#!/usr/bin/env bash
set -euo pipefail
cd $out/share/stable-diffusion-webui
export PYTHONPATH=$out/share/stable-diffusion-webui:$out/share/stable-diffusion-webui/modules
export GRADIO_ANALYTICS_ENABLED=False
exec ${pythonEnv}/bin/python webui.py \
--data-dir "\$HOME/Desktop/stable-diffusion-nix/data" \
"\$@"
EOF
chmod +x $out/bin/stable-diffusion-webui
'';
データディレクトリの苦渋の選択
modelやoutputsなどのデータを保存するディレクトリを指定する必要があります。
現在リポジトリが$HOME/Desktop/stable-diffusion-nix/にcloneされていることを前提にそれ以下のdataディレクトリを使うことにしています。
本当は$XDG_DATA_HOME以下の、具体的には$XDG_DATA_HOME/stable-diffusion-webuiに置きたいのですが、デフォルトで$XDG_DATA_HOMEは$HOME/.local/shareになり、
gradioの制約でドットで始めるディレクトリ名はsecret fileとして扱われてしまってプレビューが見られなくなってしまう問題があります。これの回避がちょっと思いつかなかったので、仕方なくパスをハードコーディングしています。
もう少し真面目にやるならせめてリポジトリの場所を環境から自動的に特定して、とりあえずclone場所はどこでも良いようにしますが、まずはnixで構築できるか実験したかったので後回しにしました。
結果
このような手順で構築することでnix runするだけでStable Diffusion web UIが立ち上がり、正常に画像が生成されることを確認しました。
警告が残っている
upstreamでPythonの新しいバージョンへの対応ができていないのか、シンタックスなどで警告が出てしまいます。もともとの方法ならどうせすぐ壊れるので放置していましたが、 nixでちゃんと再現可能なら修正しても良いかもしれません。
xformersが使えない
最適化などをしてくれるらしいxformersは有効にしていません。
nixpkgsにもpython312Packages.xformersは存在するようですが、これを依存関係に入れるとビルドが通らないので先送りすることにしました。別に必須ではないですし。
拡張機能
pipによる自動インストールが動かない
拡張機能はgit cloneして内部でセットアップするようになっていますが、これが内部でpipを実行して依存関係を解決しようとするため、 nixと相性が悪くてうまく動かないようです。
拡張機能も依存関係を含めてパッケージ化してnixで管理するようにすることを検討しています。
一定の拡張機能をデフォルトで有効化したい
例えば日本語化の拡張機能などは最初からインストールしておきたいです。これもパッケージ化して管理すれば良さそうです。
ただ手動でweb UI上でインストールする拡張機能と、ちゃんとマネージドにした拡張機能が競合しないか心配です。
マネージドな拡張機能にはextensionsとは違うディレクトリが使えれば良いのですが。
モデル
モデルのダウンロード作業も自動化したほうが当然嬉しいのでメジャーなモデルはコマンド一発でダウンロードできるようにしておきたいです。その他のモデルも選択でダウンロードできるようにしたいですね。
設定
Stable Diffusion web UIの設定もある程度nixで管理できるようにしたいです。現在config.jsonなどで設定されているようなので、
Firefoxのuser.jsみたいに大部分をnixの方でバージョン管理して常に必須の設定を書けるようにしたいです。
やるなら私用の設定はプリセットにしてコマンドで一発で設定できるようにするでしょう。現状どうせ私しか使わないのでデフォルトをそれにしても構わないのですが、一応リポジトリは名目上万人が使えるようにしているので。
感想
nixで管理することでセットアップのし直しが大幅に楽になりました。
Pydanticのバージョン不整合は結構解決が面倒でしたね。 v2に移行してほしいけど作業は面倒なのでしょうか。私がやるべきでしょうか。
Python界隈の依存関係の解決と管理は極めて面倒ですが、 nixを使うことで割と立ち向かえることがわかりました。
hatena-bookmark