ControlNet

ControlNetは、Stable Diffusionなどの事前学習済みテキストから画像への拡散モデルに空間制御を追加するニューラルネットワークアーキテクチャです。テキストプロンプトとともに、エッジマップ、深度マップ、人体姿勢スケルトン、セグメンテーションマップなどの条件画像を入力します。生成画像はその構造に従い、プロンプトが内容とスタイルを決定します。「銀色の鎧を着た騎士」というプロンプトと姿勢スケルトンを組み合わせると、その姿勢を正確に取った騎士が生成されます。Stanford UniversityのLvmin Zhang、Anyi Rao、Maneesh Agrawalaが2023年にAdding Conditional Control to Text-to-Image Diffusion Modelsで発表し、この論文はICCV 2023の最優秀論文に贈られるMarr Prizeを共同受賞しました。ベースモデルの重みは凍結されたままなので、ControlNetはモデルがすでに持つ知識を損なうことなく制御機能を追加します。

ControlNetの仕組み#

Stable Diffusionは、潜在表現を繰り返しノイズ除去することで画像を生成します。この処理にはU-Netを使います。ControlNetはそのU-Netを固定し、エンコーダーの学習可能なコピーを作成します。このコピーには、12個のエンコーディングブロックと中央のブロックが含まれます。まず、小さな畳み込みネットワークで512×512の条件画像を、潜在表現に対応する64×64の特徴マップに変換します。学習可能なコピーは潜在表現とこれらの条件特徴を処理し、その出力を、固定されたU-Netの12個のスキップ接続および中央のブロックに加算することで、固定されたデコーダーに制御信号を渡します。

コピーはゼロ畳み込みのみを介して固定モデルに接続します。ゼロ畳み込みは1×1の畳み込みで、重みとバイアスがゼロから始まります。最初の訓練ステップではゼロを出力するため、結合モデルは元のモデルとまったく同じように動作し、ランダムノイズが事前学習済みの層に届くことはありません。出力側のゼロ畳み込みには、入力がゼロではないため、ゼロ以外の重み勾配が適用されます。そのため、最初の更新で重みがゼロから変化し、そこから制御信号が大きくなります。このPyTorchコードは、1つのブロックを示しています。

import copy

import torch
from torch import nn

# A frozen block from the pretrained model and a trainable copy of it
torch.manual_seed(0)
locked = nn.Conv2d(64, 64, 3, padding=1).requires_grad_(False)
trainable = copy.deepcopy(locked).requires_grad_(True)

# Zero convolutions: 1x1 convolutions whose weights and biases start at zero
zero_in, zero_out = nn.Conv2d(64, 64, 1), nn.Conv2d(64, 64, 1)
for conv in (zero_in, zero_out):
    nn.init.zeros_(conv.weight)
    nn.init.zeros_(conv.bias)

x = torch.randn(1, 64, 32, 32)  # features entering the block
condition = torch.randn(1, 64, 32, 32)  # encoded edge, depth, or pose map
y = locked(x) + zero_out(trainable(x + zero_in(condition)))

# At step 0 the output equals the frozen block's, and only the outer zero conv gets a gradient
print(torch.equal(y, locked(x)))
y.sum().backward()
print(f"zero_in grad: {zero_in.weight.grad.norm():.1f}, zero_out grad: {zero_out.weight.grad.norm():.1f}")

Trueとzero_in grad: 0.0, zero_out grad: 2040.6を出力します。ブロックの出力は凍結モデルの出力と同一です。外側のゼロ畳み込みは、入力がゼロではないため勾配を受け取ります。内側のゼロ畳み込みは、外側のゼロ畳み込みがゼロから変化した後に学習を始めます。

条件の種類#

元の論文では、Cannyエッジ、深度マップ、表面法線マップ、M-LSDの直線、HEDのソフトエッジ、ADE20Kのセグメンテーションマップ、OpenPoseの人体姿勢、ユーザースケッチ用に、それぞれ個別のControlNetを訓練しました。各条件画像は、OpenCVのCannyエッジ検出器や深度推定ネットワークなど、既製のモデルを使って参照写真から抽出します。複数のControlNetを、それぞれ異なる強度で同時に適用することもできます。Ultralytics YOLO26の姿勢推定モデルはこの種の姿勢条件を抽出し、人物ごとに17個のCOCOキーポイントを予測します。これらは空白のキャンバス上にスケルトンとして描画できます。

import cv2
import numpy as np
from ultralytics import YOLO

# 人物ごとに17個のキーポイントを予測するYOLO26姿勢推定モデルを読み込みます
model = YOLO("yolo26n-pose.pt")
results = model("https://ultralytics.com/images/bus.jpg")
result = results[0]

# 写真と同じサイズの黒いキャンバス上に、スケルトンのみを描画します
skeletons = result.plot(img=np.zeros_like(result.orig_img), boxes=False)
cv2.imwrite("pose_condition.png", skeletons)

保存された画像では、各人物の手足が黒い背景上に色付きの線で表示され、写真の色、テクスチャ、背景は取り除かれています。

ControlNetと関連手法#

• LoRA: 少数の追加重みをファインチューニングして、モデルに新しいスタイル、キャラクター、オブジェクトを学習させます。LoRAはモデルが描く内容を変え、ControlNetは各要素の位置を固定します。両者はよく併用されます。
• T2I-Adapter: 約7,700万個のパラメーターを持つ、より軽量な代替手法です。エンコーダーの完全なコピーを実行する代わりに、条件特徴を凍結されたU-Netのエンコーダー特徴に追加します。生成時には、異なる条件に対応するアダプターを組み合わせられます。
• 画像から画像への生成: 入力写真にノイズを加えたものからノイズ除去を開始するため、出力にはその写真の色や大まかなレイアウトが引き継がれます。ControlNetはエッジなど抽出した構造のみを制約し、色やテクスチャはプロンプトに委ねます。
• 画像インペインティング: 画像のマスクされた部分を、残りの部分に合うように再生成します。変更する領域を制御する一方、ControlNetは生成されるものの構造を制御します。

実際のアプリケーション#

• 建築・インテリアデザイン: 部屋のスケッチ、CADの線画、M-LSDのラインマップを使い、形状を保ったまま、さまざまな素材やスタイルのフォトリアリスティックなレンダリングを生成します。
• キャラクターアート・アニメーション: 姿勢スケルトンによってスタイルを変えてもキャラクターの姿勢を固定できます。また、参照動画から抽出した各フレームの姿勢を使って、一貫性のあるアニメーションを生成できます。
• 製品・マーケティング画像: 製品写真の深度マップやエッジマップによって形状と配置を保ちながら、プロンプトで設定、照明、背景を変えられます。
• 合成データ: セグメンテーションマップや姿勢から画像を生成すると、生成前にレイアウトが指定されているため、ラベルがあらかじめ分かっている新しい訓練画像を作成できます。

ControlNetの訓練#

訓練には画像と条件マップのペアが必要です。キャプション付き画像データセットに条件抽出器を適用することで、このペアを自動生成します。固定された重みには勾配更新が不要なため、ControlNetを使用しないStable Diffusionの最適化と比べて、各訓練イテレーションで必要なGPUメモリは約23%増加し、時間は34%増加するにとどまります。論文では、5万枚未満から100万枚超まで、さまざまな規模のデータセットで安定して訓練できると報告されています。訓練では「突然の収束」も見られます。モデルは何千ステップもの間、条件を無視した後、突然それに従い始めます。通常は最適化ステップが10,000回以内のことで、論文の例では6,133ステップ目でした。