Stable Diffusionのもっとも単純な学習プログラムを作って動かしてみる

Stable Diffusion の仕組みを理解するには、実際に小さなモデルを自分で動かしてみるのが一番の近道です。この記事では、100×100 の画像を生成する最小構成の拡散モデルを PyTorch で作ります。テキスト条件は「flower」に固定し、学習用と生成用のファイルを分けた構成です。

この記事で作るもの

• 画像サイズは 100×100 固定
• ピクセル空間で動く DDPM(VAE は使わない)
• UNet は ResBlock、Self-Attention、Cross-Attention で構成
• テキスト条件は「flower」をエンコードした固定ベクトル
• 学習(train.py)と生成(generate.py)を別ファイルに分離

Stable Diffusion との違い

準備

必要なライブラリをインストールします。scipy は Flowers102 データセットの読み込みに使います。

pip install torch torchvision pillow scipy

学習データには Oxford 102 Flowers(8,189枚)を使います。train.py が自動でダウンロードします。画像アーカイブは約330MBで、展開後の画像も残るため、保存先の flowers_data フォルダは合計で約700MBが目安です。自分で集めた花の画像フォルダを使うこともできます。

ファイル構成

• model.py:固定パラメータ、UNet、DDPM の共通部分
• train.py:学習
• generate.py:画像生成

3つとも同じフォルダに置いてください。

model.py:モデルと拡散プロセス

学習と生成の両方から読み込まれる共通ファイルです。ポイントは次のとおりです。

• 固定パラメータ:画像サイズ、チャンネル数、拡散ステップ数(1000)などをファイル先頭の定数で固定しています。
• 固定テキスト埋め込み:encode_text() が「flower」から毎回同じベクトル(4トークン×128次元)を作ります。学習はされません。本物の CLIP などに差し替えるときは、この関数を置き換えるだけで済みます。
• UNet:100→50→25 と縮小し、スキップ接続でつなぎながら 100 に戻します。100 は2回のダウンサンプリングで割り切れるので、サイズ調整は不要です。
• Attention:最も低解像度の 25×25 の層と中間層に、Self-Attention と Cross-Attention を入れています。Cross-Attention では、画像特徴が Query、テキスト埋め込みが Key/Value になります。
• 学習の損失:画像にノイズを加え、そのノイズを UNet に当てさせる MSE 損失です。
• サンプリング:純ノイズから 1000 ステップかけてノイズを除去します(DDPM ancestral sampling)。

"""
model.py : 学習(train.py)と生成(generate.py)で共有するモジュール
- 画像サイズ 100x100 固定、ハイパーパラメータはすべて下の定数で固定
- テキスト条件は「flower」をエンコードした固定ベクトル (Cross-Attention で UNet に注入)
"""
import hashlib
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
# ---------------------------------------------------------------------------
# 固定パラメータ
# ---------------------------------------------------------------------------
IMG_SIZE = 100          # 画像サイズ (100x100)
CHANNELS = 3            # RGB
BASE_CH = 64            # 基本チャンネル数
CH_MULT = (1, 2, 2)     # 解像度ごとの倍率: 100 -> 50 -> 25
NUM_RES = 2             # 各解像度の ResBlock 数
TIME_DIM = 256          # 時刻埋め込みの次元
GROUPS = 8              # GroupNorm のグループ数
TEXT_PROMPT = "flower"  # 固定テキスト
TEXT_LEN = 4            # テキスト埋め込みのトークン数 (疑似)
TEXT_DIM = 128          # テキスト埋め込みの次元
T = 1000                # 拡散ステップ数
BETA_START = 1e-4
BETA_END = 0.02
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CKPT_PATH = "ddpm100_flower.pt"
# ---------------------------------------------------------------------------
# 固定テキスト埋め込み ("flower" をエンコードした結果の代わり)
# ---------------------------------------------------------------------------
def encode_text(prompt=TEXT_PROMPT):
"""プロンプト文字列から決定論的に作る固定ベクトル (TEXT_LEN, TEXT_DIM)。
実際の CLIP / T5 などに差し替える場合は、この関数だけ置き換えればよい。"""
seed = int(hashlib.md5(prompt.encode()).hexdigest()[:8], 16)
g = torch.Generator().manual_seed(seed)
return torch.randn(TEXT_LEN, TEXT_DIM, generator=g)
# ---------------------------------------------------------------------------
# UNet
# ---------------------------------------------------------------------------
def timestep_embedding(t, dim):
half = dim // 2
freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
args = t[:, None].float() * freqs[None]
return torch.cat([args.sin(), args.cos()], dim=-1)
class ResBlock(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.norm1 = nn.GroupNorm(GROUPS, in_ch)
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
self.t_proj = nn.Linear(TIME_DIM, out_ch)
self.norm2 = nn.GroupNorm(GROUPS, out_ch)
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
def forward(self, x, temb):
h = self.conv1(F.silu(self.norm1(x)))
h = h + self.t_proj(F.silu(temb))[:, :, None, None]
h = self.conv2(F.silu(self.norm2(h)))
return h + self.skip(x)
class SelfAttention(nn.Module):
def __init__(self, ch):
super().__init__()
self.norm = nn.GroupNorm(GROUPS, ch)
self.qkv = nn.Conv2d(ch, ch * 3, 1)
self.out = nn.Conv2d(ch, ch, 1)
def forward(self, x):
b, c, h, w = x.shape
q, k, v = self.qkv(self.norm(x)).reshape(b, 3, c, h * w).transpose(2, 3).unbind(1)
a = F.scaled_dot_product_attention(q[:, None], k[:, None], v[:, None])[:, 0]
return x + self.out(a.transpose(1, 2).reshape(b, c, h, w))
class CrossAttention(nn.Module):
"""画像特徴 (Query) がテキスト埋め込み (Key/Value) を参照する"""
def __init__(self, ch):
super().__init__()
self.norm = nn.GroupNorm(GROUPS, ch)
self.q = nn.Conv2d(ch, ch, 1)
self.k = nn.Linear(TEXT_DIM, ch)
self.v = nn.Linear(TEXT_DIM, ch)
self.out = nn.Conv2d(ch, ch, 1)
def forward(self, x, ctx):
b, c, h, w = x.shape
q = self.q(self.norm(x)).reshape(b, c, h * w).transpose(1, 2)  # (b, hw, c)
k, v = self.k(ctx), self.v(ctx)                                # (b, L, c)
a = F.scaled_dot_product_attention(q[:, None], k[:, None], v[:, None])[:, 0]
return x + self.out(a.transpose(1, 2).reshape(b, c, h, w))
class AttnBlock(nn.Module):
"""Self-Attention -> Cross-Attention (Stable Diffusion の Transformer ブロック相当を簡略化)"""
def __init__(self, ch):
super().__init__()
self.self_attn = SelfAttention(ch)
self.cross_attn = CrossAttention(ch)
def forward(self, x, ctx):
return self.cross_attn(self.self_attn(x), ctx)
class ResAttn(nn.Module):
def __init__(self, in_ch, out_ch, use_attn):
super().__init__()
self.res = ResBlock(in_ch, out_ch)
self.attn = AttnBlock(out_ch) if use_attn else None
def forward(self, x, temb, ctx):
x = self.res(x, temb)
return self.attn(x, ctx) if self.attn is not None else x
class Downsample(nn.Module):
def __init__(self, ch):
super().__init__()
self.conv = nn.Conv2d(ch, ch, 3, stride=2, padding=1)  # 100->50, 50->25
def forward(self, x):
return self.conv(x)
class Upsample(nn.Module):
def __init__(self, ch):
super().__init__()
self.conv = nn.Conv2d(ch, ch, 3, padding=1)
def forward(self, x):
return self.conv(F.interpolate(x, scale_factor=2, mode="nearest"))  # 25->50, 50->100
class UNet(nn.Module):
def __init__(self):
super().__init__()
# 固定テキスト埋め込み (学習しない・state_dict にも含めない)
self.register_buffer("text_emb", encode_text()[None], persistent=False)
self.time_mlp = nn.Sequential(
nn.Linear(TIME_DIM, TIME_DIM), nn.SiLU(), nn.Linear(TIME_DIM, TIME_DIM)
)
self.conv_in = nn.Conv2d(CHANNELS, BASE_CH, 3, padding=1)
last = len(CH_MULT) - 1  # Attention は最も低解像度 (25x25) の層だけに入れる
# ---- Encoder ----
self.down = nn.ModuleList()
skip_chs = [BASE_CH]
cur = BASE_CH
for i, m in enumerate(CH_MULT):
out = BASE_CH * m
for _ in range(NUM_RES):
self.down.append(ResAttn(cur, out, use_attn=(i == last)))
cur = out
skip_chs.append(cur)
if i != last:
self.down.append(Downsample(cur))
skip_chs.append(cur)
# ---- Middle ----
self.mid1 = ResBlock(cur, cur)
self.mid_attn = AttnBlock(cur)
self.mid2 = ResBlock(cur, cur)
# ---- Decoder ----
self.up = nn.ModuleList()
for i in reversed(range(len(CH_MULT))):
out = BASE_CH * CH_MULT[i]
for _ in range(NUM_RES + 1):
self.up.append(ResAttn(cur + skip_chs.pop(), out, use_attn=(i == last)))
cur = out
if i != 0:
self.up.append(Upsample(cur))
self.norm_out = nn.GroupNorm(GROUPS, cur)
self.conv_out = nn.Conv2d(cur, CHANNELS, 3, padding=1)
def forward(self, x, t):
ctx = self.text_emb.expand(x.shape[0], -1, -1)
temb = self.time_mlp(timestep_embedding(t, TIME_DIM))
h = self.conv_in(x)
skips = [h]
for layer in self.down:
h = layer(h, temb, ctx) if isinstance(layer, ResAttn) else layer(h)
skips.append(h)
h = self.mid1(h, temb)
h = self.mid_attn(h, ctx)
h = self.mid2(h, temb)
for layer in self.up:
if isinstance(layer, ResAttn):
h = layer(torch.cat([h, skips.pop()], dim=1), temb, ctx)
else:
h = layer(h)
return self.conv_out(F.silu(self.norm_out(h)))  # ノイズ ε を予測
# ---------------------------------------------------------------------------
# 拡散プロセス (DDPM)
# ---------------------------------------------------------------------------
betas = torch.linspace(BETA_START, BETA_END, T)
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, dim=0)
def diffusion_loss(model, x0):
"""x0 にノイズを加え、そのノイズを予測させる MSE 損失"""
t = torch.randint(0, T, (x0.shape[0],), device=x0.device)
noise = torch.randn_like(x0)
ab = alphas_cumprod.to(x0.device)[t][:, None, None, None]
xt = ab.sqrt() * x0 + (1 - ab).sqrt() * noise
return F.mse_loss(model(xt, t), noise)
@torch.no_grad()
def sample(model, n):
"""純ノイズから T ステップかけて画像を生成 (DDPM ancestral sampling)。出力は [0,1]"""
model.eval()
x = torch.randn(n, CHANNELS, IMG_SIZE, IMG_SIZE, device=DEVICE)
for i in reversed(range(T)):
t = torch.full((n,), i, device=DEVICE, dtype=torch.long)
eps = model(x, t)
a, ab, beta = alphas[i].item(), alphas_cumprod[i].item(), betas[i].item()
x = (x - beta / math.sqrt(1 - ab) * eps) / math.sqrt(a)
if i > 0:
x = x + math.sqrt(beta) * torch.randn_like(x)
return (x.clamp(-1, 1) + 1) / 2

train.py:学習

画像を短辺 100 にリサイズし、中央を 100×100 に切り出して、[-1, 1] に正規化します。ランダムな左右反転も加えています。学習中は重みの指数移動平均(EMA)も保持し、10エポックごとに EMA 重みとサンプル画像を保存します。生成には EMA 重みを使うほうが安定します。

"""
train.py : 「flower」固定テキスト条件の拡散モデルを学習する
使い方:
python train.py                    # Oxford 102 Flowers を自動ダウンロードして学習
python train.py --data ./my_flowers   # 自前の花画像フォルダで学習
必要: pip install torch torchvision pillow scipy   (scipy は Flowers102 の読み込みに必要)
"""
import argparse
import copy
import glob
import os
import torch
import torch.nn as nn
from PIL import Image
from torch.utils.data import ConcatDataset, DataLoader, Dataset
from torchvision import transforms
from torchvision.datasets import Flowers102
from torchvision.utils import save_image
from model import CKPT_PATH, DEVICE, IMG_SIZE, UNet, diffusion_loss, sample
# 学習用の固定パラメータ
BATCH_SIZE = 32
LR = 2e-4
EPOCHS = 200
EMA_DECAY = 0.999
SAMPLE_EVERY = 10  # 何エポックごとに保存・サンプル生成するか
TF = transforms.Compose([
transforms.Resize(IMG_SIZE),
transforms.CenterCrop(IMG_SIZE),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5] * 3, [0.5] * 3),  # [-1, 1]
])
class FolderDataset(Dataset):
def __init__(self, root):
exts = ("png", "jpg", "jpeg", "webp", "bmp")
self.files = sorted(
f for e in exts for f in glob.glob(os.path.join(root, "**", f"*.{e}"), recursive=True)
)
if not self.files:
raise RuntimeError(f"{root} に画像が見つかりません")
def __len__(self):
return len(self.files)
def __getitem__(self, i):
return TF(Image.open(self.files[i]).convert("RGB"))
class FlowersDataset(Dataset):
"""Oxford 102 Flowers (train+val+test 合計 8189 枚) をラベルなしで使う"""
def __init__(self, root="./flowers_data"):
self.base = ConcatDataset(
[Flowers102(root=root, split=s, download=True) for s in ("train", "val", "test")]
)
def __len__(self):
return len(self.base)
def __getitem__(self, i):
img, _ = self.base[i]
return TF(img.convert("RGB"))
def main():
p = argparse.ArgumentParser()
p.add_argument("--data", default=None, help="花画像のフォルダ (省略時は Flowers102 を自動取得)")
args = p.parse_args()
ds = FolderDataset(args.data) if args.data else FlowersDataset()
dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2,
drop_last=len(ds) >= BATCH_SIZE, pin_memory=True)
model = UNet().to(DEVICE)
ema = copy.deepcopy(model).eval()
opt = torch.optim.AdamW(model.parameters(), lr=LR)
print(f"画像数: {len(ds)} / パラメータ数: {sum(p.numel() for p in model.parameters()):,} / device: {DEVICE}")
for epoch in range(1, EPOCHS + 1):
model.train()
total = 0.0
for x0 in dl:
x0 = x0.to(DEVICE)
loss = diffusion_loss(model, x0)
opt.zero_grad(set_to_none=True)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
with torch.no_grad():  # EMA 更新
for pe, pm in zip(ema.parameters(), model.parameters()):
pe.mul_(EMA_DECAY).add_(pm, alpha=1 - EMA_DECAY)
total += loss.item()
print(f"epoch {epoch}/{EPOCHS}  loss {total / len(dl):.4f}")
if epoch % SAMPLE_EVERY == 0 or epoch == EPOCHS:
torch.save(ema.state_dict(), CKPT_PATH)
save_image(sample(ema, 16), f"sample_epoch{epoch:04d}.png", nrow=4)
print(f"  -> {CKPT_PATH} とサンプル画像を保存")
if __name__ == "__main__":
main()

generate.py:画像生成

学習済みの重みを読み込み、純ノイズから画像を生成して1枚のグリッド画像として保存します。

"""
generate.py : 学習済みモデルで「flower」の画像を生成する
使い方:
python generate.py                          # 16枚を samples.png に保存
python generate.py --n 36 --out grid.png --seed 0
"""
import argparse
import math
import torch
from torchvision.utils import save_image
from model import CKPT_PATH, DEVICE, UNet, sample
def main():
p = argparse.ArgumentParser()
p.add_argument("--n", type=int, default=16, help="生成枚数")
p.add_argument("--out", default="samples.png", help="出力ファイル")
p.add_argument("--ckpt", default=CKPT_PATH, help="学習済みチェックポイント")
p.add_argument("--seed", type=int, default=None, help="乱数シード")
args = p.parse_args()
if args.seed is not None:
torch.manual_seed(args.seed)
model = UNet().to(DEVICE)
model.load_state_dict(torch.load(args.ckpt, map_location=DEVICE))
imgs = sample(model, args.n)
save_image(imgs, args.out, nrow=max(1, int(math.sqrt(args.n))))
print(f"保存しました: {args.out}")
if __name__ == "__main__":
main()

実行手順

# 学習(Flowers102 を自動ダウンロード)
python train.py
# 自前の画像フォルダを使う場合
python train.py --data ./my_flowers
# 生成
python generate.py --n 16 --out samples.png --seed 0

GPU があると学習が速くなります。生成も 1000 ステップを順番に回すため、CPU ではかなり時間がかかります。

なぜ 100×100 なら VAE の decoder が不要なのか

Stable Diffusion が VAE を使う目的は、拡散を行う空間を小さくして計算を軽くすることです。

• 512×512×3(約79万要素)を直接扱うのは重いので、SD は VAE で約48分の1に圧縮してから拡散します。
• 100×100×3 は、もともと SD の潜在空間と同じ桁の大きさです。圧縮しても得られる軽量化は小さくなります。
• Attention の計算量はトークン数の2乗で増えます。25×25(625トークン)は、SD の 64×64(4096トークン)よりずっと軽いです。
• VAE を入れると、先に VAE を学習する手間が増えます。再構成で細部がぼやけると、それが品質の上限にもなります。

なお、ここで不要なのは VAE の decoder です。UNet 自体の後半(拡大してスキップ接続をつなぐ部分)は、入力と同じ解像度のノイズ予測を出すために必要です。

注意点

• テキストが常に「flower」の1種類なので、Cross-Attention は実質「定数を足す」働きしかしません。構造がテキスト条件付きになっているだけで、花を描く能力は学習データで決まります。複数のプロンプトに対応させるには、プロンプトごとに埋め込みを変え、ラベル付きデータで学習する必要があります。
• 100×100 の画像なので、細部は粗くなります。
• 学習がこのコードで正しくできること(lossの低下)を確認しています

発展:ここから先にできること

• encode_text() を CLIP のテキストエンコーダに置き換えて、複数のプロンプトに対応させる
• 画像サイズを 256×256 以上に上げ、VAE で潜在空間化する
• DDIM サンプラーを実装して、生成のステップ数を減らす
• Classifier-free guidance を加えて、テキスト条件の効きを調整する

まとめ

ピクセル空間の DDPM と UNet に、固定テキストの Cross-Attention を加えるだけで、Stable Diffusion の骨格を小さく再現できます。VAE を外し、サイズと条件を固定したことで、コードは3ファイルにまとまりました。まずはこの最小構成を動かし、そこから CLIP の導入や潜在空間化へ進むと、Stable Diffusion との差が理解しやすくなります。