MkItYs: ComfyUI で、画像処理のノードを作る：Python (PIL, Torchvision), ComfyUI

ComfyUI で、画像処理のノードを作ってみます。［※１］［※２］

※: 生成画像：Counterfeit V3.0

※１: ここで実装するのは、画像のメインとマスクを受け取り、マスク部分にべつの画像を合成する、シンプルなカスタムノードです（同様のことは、マスクを反転させた i2i や（より高機能の）IC-Light などでも可能です）。

※２: たんにアルファ値のある透過画像を LoadImage ノードにロードしてもいいのですが（アルファ部分がマスクに変換されます）ーーここでは RemBg を使い、通常のキャラ画像から背景を検出し、マスクに変換したものを使っています。

検証

・: クラウド：GCP
・: コンテナ：Docker
・: ゲストＯＳ：Ubuntu 22.04
・: ウェブブラウザ：Chrome
・: 画像生成アプリ：ComfyUI

概要

◯: 背景

ComfyUI は、画像生成のアプリです。

画像生成では、潜在空間の画像から、通常の画像を生成しますーーそしてこの通常の画像に対しては、ロードやセーブだけでなく、さまざまな画像処理が必要になることがあります。

◯: 対応

ここでは ComfyUI で、かんたんな画像処理のカスタムノードを実装してみます。

◯: 仕様：形式

ComfyUI は、通常の画像を PyTorch のテンソル形式でやり取りしますーー例外もありますが。［後述］

画像の形式は HWC/RGB ですーーこれはたんに、PIL (Pillow) の配列を正規化したものですーーなので単純に、256 の割り算／掛け算で、PIL の画像形式と PyTorch のテンソル形式を変換できます。［※１］［※２］

◯: 仕様：構成

ComfyUI の通常の画像は、メインとマスク（透過処理／アルファ）に分かれますーーこれも例外がありますが。［※３］［後述］

このときメインの画像は (B, H, W, C) 、マスクの画像は (B, H, W) の形式になります。［※４］［※５］

※１: PyTorch の Torchvision の標準形式は CHW ですが、それよりは PIL をかんたんにあつかえるようにするためでしょうね。

※２: この変換には Torchvision の ToPILImage / ToTensor を使うこともできますが、すこし面倒です（CHW の形状にいったん変えないといけないので）ーーまた現在は、この関数そのものが非推奨のようです：
・: UserWarning: The transform `ToTensor()` is deprecated and will be removed in a future release. Instead, please use `v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])`.

※３: Stable Diffusion の仕様では、アルファ値による透過画像をあつかえないためです（それで Lvmin Zhang (lllyasviel) 氏が LayerDiffusion を公開したりしていますが）。

※４: ただしカスタムノードによっては、テンソルの形式ではなく、リスト（list）の形式で返したり、そのまま PIL の形式で返すものもありますーー使う方は混乱しますが、形式の統一を強制できるとも思えませんしーーそれぞれのカスタムノードの記述を読んで、個別に対応していくしかないでしょうね。

※５: バッチ処理そのものを何度も繰り返したいなら、API を使うことになります（ノードによる工夫もありますが、トリッキーです）。

実装

以下、画像のメインとマスクを受け取り、マスク部分にべつの画像を合成する、シンプルなカスタムノードの実装です（キャラと背景を組み合わせる、など）：

◯: 呼出

・: comfyui_xoxxox/__init__.py

from .setimg_cmm import *

NODE_CLASS_MAPPINGS = {
  "ImgCmb_001": ImgCmb_001,
  "ImgCmb_002": ImgCmb_002,
}

◯: 本体

・: comfyui_xoxxox/setimg_cmm.py

import numpy as np
from PIL import Image
import torch

#---------------------------------------------------------------------------
# 画像の変換

class CnvImg:
  # HWC (256) -> HWC (0-1)
  @classmethod
  def cnvtsr(cls, imgpil):
    g = np.array(imgpil).astype(np.float32) / 255.0
    imgtsr = torch.from_numpy(g)
    return imgtsr

  # HWC (0-1) -> HWC (256)
  @classmethod
  def cnvpil(cls, imgtsr):
    g = imgtsr.cpu().numpy() * 255.0
    g = np.clip(g, 0, 255).astype(np.uint8)
    imgpil = Image.fromarray(g)
    return imgpil

#---------------------------------------------------------------------------
# 背景の画像と本体の画像とそのマスクを受け取り、背景の画像と合成する

class ImgCmb_001:
  @classmethod
  def INPUT_TYPES(s):
    return {
      "required": {
        "imgtsr_bak": ("IMAGE",),
        "imgtsr_con": ("IMAGE",),
        "imgtsr_msk": ("MASK",),
      },
    }
  RETURN_TYPES = ("IMAGE",)
  FUNCTION = "anchor"
  CATEGORY = "xoxxox"

  def anchor(self, imgtsr_bak, imgtsr_con, imgtsr_msk):
    lstimg = []
    imgpil_bak = CnvImg.cnvpil(imgtsr_bak[0])

    for (i, m) in enumerate(imgtsr_con):
      imgpil_con = CnvImg.cnvpil(imgtsr_con[i])
      imgpil_msk = CnvImg.cnvpil(imgtsr_msk[i])
      imgpil_cmp = Image.composite(imgpil_con, imgpil_bak, imgpil_msk)
      lstimg.append(CnvImg.cnvtsr(imgpil_cmp)[None,])
    lsttsr = torch.cat(lstimg)
    return (lsttsr,)

#---------------------------------------------------------------------------
# 背景の画像と本体の画像とそのマスクを受け取り、背景の画像と合成する（使用：ToPILImage, ToTensor）

from torchvision.transforms import ToPILImage, ToTensor

class ImgCmb_002:
  @classmethod
  def INPUT_TYPES(s):
    return {
      "required": {
        "imgtsr_bak": ("IMAGE",),
        "imgtsr_con": ("IMAGE",),
        "imgtsr_msk": ("MASK",),
      },
    }
  RETURN_TYPES = ("IMAGE",)
  FUNCTION = "anchor"
  CATEGORY = "xoxxox"

  def anchor(self, imgtsr_bak, imgtsr_con, imgtsr_msk):
    cnvpil = ToPILImage()
    cnvtsr = ToTensor()
    lstimg = []
    print("imgtsr_bak[" + str(imgtsr_bak[0].shape) + "]") # DBG (> imgtsr_bak[torch.Size([1216, 832, 3])]) # HWC (0-1)
    imgpil_bak = cnvpil(imgtsr_bak[0].permute(2, 0, 1)) # HWC (0-1) -> CHW (0-1) -> HWC (256)

    for (i, m) in enumerate(imgtsr_con):
      print("imgtsr_con[" + str(imgtsr_con[i].shape) + "]") # DBG (> imgtsr_con[torch.Size([1216, 832, 3])]) # HWC (0-1)
      print("imgtsr_msk[" + str(imgtsr_msk[i].shape) + "]") # DBG (> imgtsr_msk[torch.Size([1216, 832])]) # HW (0-1)
      imgpil_con = cnvpil(imgtsr_con[i].permute(2, 0, 1)) # HWC (0-1) -> CHW (0-1) -> HWC (256)
      imgpil_msk = cnvpil(imgtsr_msk[i]) # HW (0-1) -> HW (256)
      imgpil_cmp = Image.composite(imgpil_con, imgpil_bak, imgpil_msk)
      lstimg.append(cnvtsr(imgpil_cmp).permute(1, 2, 0)[None,]) # HWC (256)-> CHW (0-1) -> HWC (0-1)
    lsttsr = torch.cat(lstimg)
    return (lsttsr,)

#---------------------------------------------------------------------------

関連

検証

概要

実装