24/11/21 Windows、ComfyUI、FluxでTorch.Compileして画像生成を20%高速化しよう

最初に結論を書いておくと、この方法は実際やってみたら言うほど便利じゃないし活用できるシチュエーションは限定的だと思う。

前置き

今までのSDXLとかに比べると、Fluxの画像生成にはとかく時間がかかる。当初はRTX4090で1024x1024の画像生成に16秒かかっていた。その後、ComfyUIのfastモードによってfp8モデルなら生成時間は10秒になって37%高速化された。

それはいいんだが、できればもっと高速化して欲しい。有力な方法としてモデルをTorch.Compileする事でさらに高速化できるらしい。

しかしTorch.Compileが使えるのはLinuxだけだった。Torch.Compileを使うにはTritonというライブラリが必要なのだが、TritonチームはWindowsサポートを拒否し続けている。

Windowsで使いたければWSL2で動かすしかない。しかし私はEagleという画像管理ソフトとの連携が必須だったので、WSL2には移行できなかった。

そうこうしてる内に、TritonをWindowsサポートするフォークリポジトリを作ってくれた人が現れた。

WindowsでTritonが動くという事は、WindowsでTorch.Compileできるようになったという事であり、つまりWindowsでFlux推論が高速化できるようになったという事だ。

あとTorch.CompileするとLoRAが効かなくなるデメリットがあったらしいが、それも今は解決されているらしい↓

あらかじめ書いておくが、WindowsでTorch.Compileする導入は結構めんどい。カスタムノード入れてオシマイとかではない。

また、私の場合は上手く動かなくてアレコレ試行錯誤するハメになったので、割とハチャメチャな導入をしてしまった。確実なインスコ方法は書けないので、せいぜいヒント程度に捉えて欲しい。