技術 約2分で読めます

ComfyUIのUpscaleがMac MPSで壊れる問題をcontiguousで直した

ComfyUI Apple Silicon PyTorch MPS

AntiGravityで出した640px画像をComfyUIで一括拡大していたら、出力が一部ガチで崩壊した。
最初はアップスケーラモデルの相性だと思っていたけど、実際はもっと手前のメモリレイアウトが原因だった。

何が起きていたか

再現条件は以下だった。

  • Mac（Apple Silicon）
  • ComfyUIの Upscale Image (using Model)
  • Load Image から入った画像をそのままUpscale

同じUpscaleでも、生成フロー末尾の画像は壊れにくいのに、Load Image 起点だけ壊れる。
この差を追ったら、原因箇所がはっきりした。

原因

ComfyUIの画像テンソルは基本 BHWC。Upscaleノード内部ではESRGAN系モデルに渡すため BCHW に並べ替える。

この並べ替え後テンソルが non-contiguous のまま conv2d 側に渡るケースがあり、MPS環境で破綻する。

ポイントは tiled_scale_multidim 側でタイルを narrow() した s_in を、そのまま function(s_in) に渡していたこと。

実際に当てた修正

comfy/utils.py の該当行を次のように変更した。

# before
ps = function(s_in).to(output_device)

# after
ps = function(s_in.contiguous()).to(output_device)

この1行で、壊れていた出力が正常化した。

関連Issueはこれ。この記事を書いている時点でまだオープン。
https://github.com/comfyanonymous/ComfyUI/issues/11851

パッチ手順

cd ~/ComfyUI
source venv/bin/activate

# 念のためバックアップ
cp comfy/utils.py comfy/utils.py.bak

# 手修正（エディタで以下を置換）
# ps = function(s_in).to(output_device)
# ->
# ps = function(s_in.contiguous()).to(output_device)

# ComfyUI再起動
python main.py

補足

--force-fp32 や PYTORCH_ENABLE_MPS_FALLBACK=0 は切り分けとしては有効だった。
ただ、今回の本命は精度設定よりメモリレイアウトだった。

なので「MPSだから壊れる」というより、「non-contiguous入力を踏む経路だとMPSで壊れやすい」のほうが実態に近い。

運用メモ

ComfyUIを更新すると comfy/utils.py の手修正は上書きされる。
更新後に再発したら、同じ1行を再適用する。

同様の壊れ方が Sharpen や他の BHWC -> BCHW 変換ノードで出るなら、permute/movedim 直後に contiguous() を挟むと再発防止しやすい。