生成AI完全ガイド: GAN・VAE・拡散モデルをマスターする

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

class Autoencoder(nn.Module):
    """MNIST用の基本的なオートエンコーダ"""
    def __init__(self, input_dim=784, latent_dim=32):
        super().__init__()

        # エンコーダ
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, latent_dim),
            nn.ReLU()
        )

        # デコーダ
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, input_dim),
            nn.Sigmoid()  # ピクセル値を[0, 1]に制限
        )

    def forward(self, x):
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_reconstructed = self.decoder(z)
        return x_reconstructed, z

    def encode(self, x):
        x = x.view(x.size(0), -1)
        return self.encoder(x)

    def decode(self, z):
        return self.decoder(z).view(-1, 1, 28, 28)


def train_autoencoder(epochs=20):
    """オートエンコーダの学習"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    transform = transforms.Compose([transforms.ToTensor()])
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

    model = Autoencoder().to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.BCELoss()

    for epoch in range(epochs):
        total_loss = 0
        for data, _ in train_loader:
            data = data.to(device)
            optimizer.zero_grad()
            reconstructed, z = model(data)
            target = data.view(data.size(0), -1)
            loss = criterion(reconstructed, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(train_loader)
        print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")

    return model

log P(x) >= E_q[log P(x|z)] - KL[q(z|x) || P(z)]
                 ↑                     ↑
        再構成損失              KL発散

z = μ + σ * ε,  ε ~ N(0, I)

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np

class VAE(nn.Module):
    """変分オートエンコーダ"""
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super(VAE, self).__init__()
        self.latent_dim = latent_dim

        # エンコーダ: 入力 -> μ, log(σ²)
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)       # 平均
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)   # 対数分散

        # デコーダ: 潜在ベクトル -> 再構成
        self.fc3 = nn.Linear(latent_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, input_dim)

    def encode(self, x):
        """エンコーダ: x -> (μ, log_var)"""
        h = F.relu(self.fc1(x))
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        """再パラメータ化トリック: z = μ + ε * σ"""
        if self.training:
            std = torch.exp(0.5 * logvar)  # σ = exp(0.5 * log σ²)
            eps = torch.randn_like(std)     # ε ~ N(0, I)
            return mu + eps * std
        else:
            return mu  # 推論時は平均のみ使用

    def decode(self, z):
        """デコーダ: z -> x'"""
        h = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h))

    def forward(self, x):
        x_flat = x.view(-1, 784)
        mu, logvar = self.encode(x_flat)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z)
        return x_recon, mu, logvar

    def generate(self, num_samples, device):
        """標準正規分布からサンプリングして画像を生成"""
        with torch.no_grad():
            z = torch.randn(num_samples, self.latent_dim).to(device)
            samples = self.decode(z)
            return samples.view(num_samples, 1, 28, 28)


def vae_loss(x_recon, x, mu, logvar, beta=1.0):
    """
    VAE損失 = 再構成損失 + β * KL発散
    beta=1: 標準VAE
    beta>1: β-VAE（よりディスエンタングルな表現）
    """
    # 再構成損失（BCE）
    recon_loss = F.binary_cross_entropy(
        x_recon, x.view(-1, 784),
        reduction='sum'
    )

    # KL発散: KL[N(μ, σ²) || N(0, 1)]
    # = -0.5 * Σ(1 + log σ² - μ² - σ²)
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return recon_loss + beta * kl_loss


def train_vae(epochs=50, latent_dim=20, beta=1.0):
    """VAEの学習"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    transform = transforms.Compose([transforms.ToTensor()])
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

    model = VAE(latent_dim=latent_dim).to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    train_losses = []

    for epoch in range(epochs):
        model.train()
        total_loss = 0

        for data, _ in train_loader:
            data = data.to(device)
            optimizer.zero_grad()
            x_recon, mu, logvar = model(data)
            loss = vae_loss(x_recon, data, mu, logvar, beta)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(train_dataset)
        train_losses.append(avg_loss)

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")

    return model, train_losses


def interpolate_latent_space(model, img1, img2, steps=10, device='cpu'):
    """潜在空間内の2枚の画像を補間"""
    model.eval()
    with torch.no_grad():
        z1_flat = img1.view(-1, 784).to(device)
        z2_flat = img2.view(-1, 784).to(device)

        mu1, _ = model.encode(z1_flat)
        mu2, _ = model.encode(z2_flat)

        # 線形補間
        interpolated_images = []
        for alpha in np.linspace(0, 1, steps):
            z_interp = (1 - alpha) * mu1 + alpha * mu2
            img_recon = model.decode(z_interp)
            interpolated_images.append(img_recon.view(1, 28, 28))

    return interpolated_images

class ConvVAE(nn.Module):
    """カラー画像用の畳み込みVAE"""
    def __init__(self, latent_dim=128):
        super().__init__()
        self.latent_dim = latent_dim

        # 畳み込みエンコーダ
        self.encoder_conv = nn.Sequential(
            nn.Conv2d(3, 32, 4, stride=2, padding=1),   # 16x16
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2, padding=1),  # 8x8
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, stride=2, padding=1), # 4x4
            nn.ReLU(),
        )
        self.fc_mu = nn.Linear(128 * 4 * 4, latent_dim)
        self.fc_logvar = nn.Linear(128 * 4 * 4, latent_dim)

        # 転置畳み込みデコーダ
        self.decoder_fc = nn.Linear(latent_dim, 128 * 4 * 4)
        self.decoder_conv = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),  # 8x8
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),   # 16x16
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 4, stride=2, padding=1),    # 32x32
            nn.Sigmoid()
        )

    def encode(self, x):
        h = self.encoder_conv(x).view(x.size(0), -1)
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        h = F.relu(self.decoder_fc(z)).view(-1, 128, 4, 4)
        return self.decoder_conv(h)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

min_G max_D [E_x[log D(x)] + E_z[log(1 - D(G(z)))]]

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np

class Generator(nn.Module):
    """GANジェネレータ"""
    def __init__(self, noise_dim=100, output_dim=784):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(noise_dim, 256),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(256),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(512),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(1024),
            nn.Linear(1024, output_dim),
            nn.Tanh()  # [-1, 1]の範囲
        )

    def forward(self, z):
        return self.model(z).view(-1, 1, 28, 28)


class Discriminator(nn.Module):
    """GANディスクリミネータ"""
    def __init__(self, input_dim=784):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, 1024),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(1024, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.model(x.view(x.size(0), -1))


def train_gan(epochs=200, noise_dim=100, lr=2e-4):
    """GANの学習"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])  # [-1, 1]に正規化
    ])
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)

    G = Generator(noise_dim).to(device)
    D = Discriminator().to(device)

    optimizer_G = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))
    optimizer_D = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))

    criterion = nn.BCELoss()

    for epoch in range(epochs):
        for i, (real_imgs, _) in enumerate(dataloader):
            batch_size = real_imgs.size(0)
            real_imgs = real_imgs.to(device)

            real_labels = torch.ones(batch_size, 1).to(device)
            fake_labels = torch.zeros(batch_size, 1).to(device)

            # === ディスクリミネータの更新 ===
            optimizer_D.zero_grad()

            d_real = D(real_imgs)
            d_loss_real = criterion(d_real, real_labels)

            z = torch.randn(batch_size, noise_dim).to(device)
            fake_imgs = G(z).detach()  # Gの勾配をデタッチ
            d_fake = D(fake_imgs)
            d_loss_fake = criterion(d_fake, fake_labels)

            d_loss = d_loss_real + d_loss_fake
            d_loss.backward()
            optimizer_D.step()

            # === ジェネレータの更新 ===
            optimizer_G.zero_grad()

            z = torch.randn(batch_size, noise_dim).to(device)
            fake_imgs = G(z)
            g_pred = D(fake_imgs)
            # GはDが偽物を本物として分類するよう望む
            g_loss = criterion(g_pred, real_labels)

            g_loss.backward()
            optimizer_G.step()

        if (epoch + 1) % 20 == 0:
            print(f"Epoch {epoch+1}/{epochs} | D Loss: {d_loss.item():.4f} | G Loss: {g_loss.item():.4f}")

    return G, D

class DCGANGenerator(nn.Module):
    """DCGANジェネレータ（64x64画像用）"""
    def __init__(self, noise_dim=100, ngf=64, nc=3):
        super().__init__()
        self.main = nn.Sequential(
            # 入力: noise_dim x 1 x 1
            nn.ConvTranspose2d(noise_dim, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # 状態: (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # 状態: (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # 状態: (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # 状態: (ngf) x 32 x 32
            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # 出力: nc x 64 x 64
        )

    def forward(self, z):
        z = z.view(z.size(0), -1, 1, 1)
        return self.main(z)


class DCGANDiscriminator(nn.Module):
    """DCGANディスクリミネータ（64x64画像用）"""
    def __init__(self, ndf=64, nc=3):
        super().__init__()
        self.main = nn.Sequential(
            # 入力: nc x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.main(x).view(-1, 1)

class WGANDiscriminator(nn.Module):
    """WGANクリティック: Sigmoid出力なし"""
    def __init__(self, input_dim=784):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1)
            # Sigmoidなし！Wasserstein距離計算に必要
        )

    def forward(self, x):
        return self.model(x.view(x.size(0), -1))


def train_wgan(G, D, dataloader, device, epochs=100,
               n_critic=5, clip_value=0.01, lr=5e-5):
    """WGANの学習"""
    optimizer_G = optim.RMSprop(G.parameters(), lr=lr)
    optimizer_D = optim.RMSprop(D.parameters(), lr=lr)

    for epoch in range(epochs):
        for i, (real_imgs, _) in enumerate(dataloader):
            real_imgs = real_imgs.to(device)
            batch_size = real_imgs.size(0)

            # クリティックをn_critic回更新
            for _ in range(n_critic):
                optimizer_D.zero_grad()

                z = torch.randn(batch_size, 100).to(device)
                fake_imgs = G(z).detach()

                # Wasserstein損失: E[D(real)] - E[D(fake)]
                d_loss = -torch.mean(D(real_imgs)) + torch.mean(D(fake_imgs))
                d_loss.backward()
                optimizer_D.step()

                # ウェイトクリッピング（Lipschitz制約）
                for p in D.parameters():
                    p.data.clamp_(-clip_value, clip_value)

            # ジェネレータの更新
            optimizer_G.zero_grad()
            z = torch.randn(batch_size, 100).to(device)
            fake_imgs = G(z)
            g_loss = -torch.mean(D(fake_imgs))
            g_loss.backward()
            optimizer_G.step()

def gradient_penalty(D, real_imgs, fake_imgs, device):
    """勾配ペナルティの計算"""
    batch_size = real_imgs.size(0)
    # 本物と偽物の画像間のランダムな補間
    alpha = torch.rand(batch_size, 1, 1, 1).to(device)
    interpolated = alpha * real_imgs + (1 - alpha) * fake_imgs
    interpolated.requires_grad_(True)

    d_interpolated = D(interpolated)

    gradients = torch.autograd.grad(
        outputs=d_interpolated,
        inputs=interpolated,
        grad_outputs=torch.ones_like(d_interpolated),
        create_graph=True,
        retain_graph=True
    )[0]

    gradients = gradients.view(batch_size, -1)
    gradient_norm = gradients.norm(2, dim=1)
    penalty = ((gradient_norm - 1) ** 2).mean()
    return penalty


def wgan_gp_d_loss(D, real_imgs, fake_imgs, device, lambda_gp=10):
    """WGAN-GPディスクリミネータ損失"""
    d_real = D(real_imgs).mean()
    d_fake = D(fake_imgs).mean()
    gp = gradient_penalty(D, real_imgs, fake_imgs, device)
    return -d_real + d_fake + lambda_gp * gp

class ConditionalGenerator(nn.Module):
    """条件付きGANジェネレータ"""
    def __init__(self, noise_dim=100, num_classes=10, embed_dim=50):
        super().__init__()
        self.label_embedding = nn.Embedding(num_classes, embed_dim)
        self.model = nn.Sequential(
            nn.Linear(noise_dim + embed_dim, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 784),
            nn.Tanh()
        )

    def forward(self, z, labels):
        label_embed = self.label_embedding(labels)
        x = torch.cat([z, label_embed], dim=1)
        return self.model(x).view(-1, 1, 28, 28)

q(x_t | x_{t-1}) = N(x_t; sqrt(1-β_t) * x_{t-1}, β_t * I)

q(x_t | x_0) = N(x_t; sqrt(ā_t) * x_0, (1-ā_t) * I)

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class NoiseSchedule:
    """ノイズスケジュールを管理"""
    def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.timesteps = timesteps

        # 線形ノイズスケジュール
        self.betas = torch.linspace(beta_start, beta_end, timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)

        # 順方向プロセスの係数
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)

        # 逆方向プロセスの係数
        self.posterior_variance = (
            self.betas * (1.0 - self.alphas_cumprod_prev) /
            (1.0 - self.alphas_cumprod)
        )

    def q_sample(self, x_start, t, noise=None):
        """順方向プロセス: x_0にtステップのノイズを追加"""
        if noise is None:
            noise = torch.randn_like(x_start)

        sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
        sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)

        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise

class SinusoidalPositionEmbeddings(nn.Module):
    """タイムステップの正弦波位置埋め込み"""
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, time):
        device = time.device
        half_dim = self.dim // 2
        embeddings = np.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        return embeddings


class ResidualBlock(nn.Module):
    """タイムステップ条件付きの残差ブロック"""
    def __init__(self, in_channels, out_channels, time_emb_dim):
        super().__init__()
        self.time_mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(time_emb_dim, out_channels)
        )
        self.block1 = nn.Sequential(
            nn.GroupNorm(8, in_channels),
            nn.SiLU(),
            nn.Conv2d(in_channels, out_channels, 3, padding=1)
        )
        self.block2 = nn.Sequential(
            nn.GroupNorm(8, out_channels),
            nn.SiLU(),
            nn.Conv2d(out_channels, out_channels, 3, padding=1)
        )
        self.residual_conv = (
            nn.Conv2d(in_channels, out_channels, 1)
            if in_channels != out_channels else nn.Identity()
        )

    def forward(self, x, time_emb):
        h = self.block1(x)
        time_emb = self.time_mlp(time_emb)[:, :, None, None]
        h = h + time_emb
        h = self.block2(h)
        return h + self.residual_conv(x)


class SimpleUNet(nn.Module):
    """DDPMのための簡略化したU-Net"""
    def __init__(self, in_channels=1, model_channels=64, time_emb_dim=256):
        super().__init__()

        # タイムステップ埋め込み
        self.time_mlp = nn.Sequential(
            SinusoidalPositionEmbeddings(model_channels),
            nn.Linear(model_channels, time_emb_dim),
            nn.SiLU(),
            nn.Linear(time_emb_dim, time_emb_dim)
        )

        # エンコーダ
        self.down1 = ResidualBlock(in_channels, model_channels, time_emb_dim)
        self.down2 = ResidualBlock(model_channels, model_channels * 2, time_emb_dim)
        self.pool = nn.MaxPool2d(2)

        # ボトルネック
        self.bottleneck = ResidualBlock(model_channels * 2, model_channels * 2, time_emb_dim)

        # デコーダ
        self.up1 = ResidualBlock(model_channels * 4, model_channels, time_emb_dim)
        self.up2 = ResidualBlock(model_channels * 2, model_channels, time_emb_dim)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        # 出力
        self.final = nn.Conv2d(model_channels, in_channels, 1)

    def forward(self, x, t):
        time_emb = self.time_mlp(t)

        # エンコーダパス
        x1 = self.down1(x, time_emb)
        x2 = self.down2(self.pool(x1), time_emb)

        # ボトルネック
        x_mid = self.bottleneck(x2, time_emb)

        # デコーダパス（スキップ接続付き）
        x = self.up1(torch.cat([self.upsample(x_mid), x2], dim=1), time_emb)
        x = self.up2(torch.cat([self.upsample(x), x1], dim=1), time_emb)

        return self.final(x)

class DDPM:
    """DDPMの学習とサンプリング"""
    def __init__(self, model, noise_schedule, device):
        self.model = model
        self.schedule = noise_schedule
        self.device = device

    def get_loss(self, x_start, t):
        """学習損失: ノイズ予測誤差"""
        noise = torch.randn_like(x_start)
        x_noisy = self.schedule.q_sample(x_start, t, noise)

        # ニューラルネットワークが追加されたノイズを予測
        predicted_noise = self.model(x_noisy, t)

        return F.mse_loss(noise, predicted_noise)

    @torch.no_grad()
    def p_sample(self, x_t, t):
        """逆方向プロセス: x_tからx_{t-1}へ1ステップデノイズ"""
        betas_t = self.schedule.betas[t].view(-1, 1, 1, 1).to(self.device)
        sqrt_one_minus_alphas_cumprod_t = (
            self.schedule.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1).to(self.device)
        )
        sqrt_recip_alphas_t = torch.sqrt(
            1.0 / self.schedule.alphas[t]
        ).view(-1, 1, 1, 1).to(self.device)

        predicted_noise = self.model(x_t, t)

        model_mean = sqrt_recip_alphas_t * (
            x_t - betas_t * predicted_noise / sqrt_one_minus_alphas_cumprod_t
        )

        if t[0] == 0:
            return model_mean
        else:
            posterior_variance_t = (
                self.schedule.posterior_variance[t].view(-1, 1, 1, 1).to(self.device)
            )
            noise = torch.randn_like(x_t)
            return model_mean + torch.sqrt(posterior_variance_t) * noise

    @torch.no_grad()
    def sample(self, batch_size, image_shape):
        """純粋なノイズから始めて画像を生成"""
        img = torch.randn(batch_size, *image_shape).to(self.device)

        for i in reversed(range(self.schedule.timesteps)):
            t = torch.full((batch_size,), i, dtype=torch.long, device=self.device)
            img = self.p_sample(img, t)

        return img


def train_ddpm(epochs=100):
    """DDPMの学習"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    ])
    dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

    model = SimpleUNet(in_channels=1).to(device)
    schedule = NoiseSchedule(timesteps=1000)
    ddpm = DDPM(model, schedule, device)

    optimizer = optim.Adam(model.parameters(), lr=2e-4)

    for epoch in range(epochs):
        total_loss = 0
        for batch, (x, _) in enumerate(dataloader):
            x = x.to(device)
            # ランダムなタイムステップサンプリング
            t = torch.randint(0, schedule.timesteps, (x.size(0),), device=device)

            loss = ddpm.get_loss(x, t)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        if (epoch + 1) % 10 == 0:
            avg_loss = total_loss / len(dataloader)
            print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")

    return model, ddpm

@torch.no_grad()
def ddim_sample(model, schedule, batch_size, image_shape,
                ddim_timesteps=50, eta=0.0, device='cpu'):
    """
    DDIMサンプリング
    eta=0.0: 完全に決定論的
    eta=1.0: DDPMと同等
    """
    # 均等に間隔を置いたタイムステップを選択
    c = schedule.timesteps // ddim_timesteps
    timestep_seq = list(range(0, schedule.timesteps, c))[::-1]

    img = torch.randn(batch_size, *image_shape).to(device)

    for i, t in enumerate(timestep_seq):
        t_tensor = torch.full((batch_size,), t, dtype=torch.long, device=device)
        t_prev = timestep_seq[i + 1] if i + 1 < len(timestep_seq) else -1

        alpha_bar = schedule.alphas_cumprod[t].to(device)
        alpha_bar_prev = (
            schedule.alphas_cumprod[t_prev].to(device) if t_prev >= 0
            else torch.tensor(1.0, device=device)
        )

        # ノイズを予測
        pred_noise = model(img, t_tensor)

        # x_0を予測
        pred_x0 = (img - torch.sqrt(1 - alpha_bar) * pred_noise) / torch.sqrt(alpha_bar)
        pred_x0 = torch.clamp(pred_x0, -1, 1)

        # シグマを計算
        sigma = eta * torch.sqrt(
            (1 - alpha_bar_prev) / (1 - alpha_bar) * (1 - alpha_bar / alpha_bar_prev)
        )

        # x_tへの方向
        pred_dir = torch.sqrt(1 - alpha_bar_prev - sigma**2) * pred_noise

        # 次のステップの画像
        noise = torch.randn_like(img) if t_prev >= 0 else 0
        img = torch.sqrt(alpha_bar_prev) * pred_x0 + pred_dir + sigma * noise

    return img

テキストプロンプト → CLIPテキストエンコーダ → テキスト埋め込み
                                                    ↓
純粋ノイズ z_T → [U-Net + クロスアテンション] → 潜在 z_0
                                                    ↓
                                     VAEデコーダ → 最終画像

guided_noise = uncond_noise + guidance_scale * (cond_noise - uncond_noise)

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch

model_id = "stabilityai/stable-diffusion-2-1"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    use_safetensors=True,
)

# 高速スケジューラを使用
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

# 画像を生成
prompt = "a photorealistic landscape of mountains at sunset, 8k, highly detailed"
negative_prompt = "blurry, low quality, distorted"

image = pipe(
    prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=25,    # 1000 DDPMステップ → 25ステップ
    guidance_scale=7.5,        # CFGスケール
    height=512,
    width=512,
    generator=torch.Generator("cuda").manual_seed(42)
).images[0]

image.save("generated_image.png")


# 画像から画像への変換
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image

img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16
).to("cuda")

init_image = Image.open("input.jpg").resize((512, 512))
result = img2img_pipe(
    prompt="a painting in Van Gogh style",
    image=init_image,
    strength=0.75,  # 変換の度合い（0-1）
    guidance_scale=7.5,
    num_inference_steps=50
).images[0]

from diffusers import StableDiffusionInpaintPipeline

inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    torch_dtype=torch.float16
).to("cuda")

image = Image.open("photo.jpg").resize((512, 512))
mask = Image.open("mask.jpg").resize((512, 512))  # 白=インペイント領域

result = inpaint_pipe(
    prompt="a beautiful garden with flowers",
    image=image,
    mask_image=mask,
    num_inference_steps=50
).images[0]

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import cv2
import numpy as np

# Canny ControlNetを読み込み
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-canny",
    torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16
).to("cuda")

# Cannyエッジを抽出
image = load_image("input.jpg")
image_array = np.array(image)
canny = cv2.Canny(image_array, 100, 200)
canny_image = Image.fromarray(canny)

# ControlNetで生成
result = pipe(
    prompt="a beautiful oil painting",
    image=canny_image,   # 制御信号としてのCannyエッジ
    controlnet_conditioning_scale=1.0,
    num_inference_steps=50,
    guidance_scale=7.5
).images[0]

class DiTBlock(nn.Module):
    """拡散トランスフォーマーブロック"""
    def __init__(self, hidden_dim, num_heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(hidden_dim)

        mlp_dim = int(hidden_dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim, mlp_dim),
            nn.GELU(),
            nn.Linear(mlp_dim, hidden_dim)
        )

        # AdaLN（アダプティブ層正規化）: タイムステップ条件付け
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_dim, 6 * hidden_dim)
        )

    def forward(self, x, c):
        # タイムステップ/条件埋め込みから変調パラメータを計算
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.adaLN_modulation(c).chunk(6, dim=-1)
        )

        # 変調された正規化
        x_norm = (1 + scale_msa.unsqueeze(1)) * self.norm1(x) + shift_msa.unsqueeze(1)
        attn_out, _ = self.attn(x_norm, x_norm, x_norm)
        x = x + gate_msa.unsqueeze(1) * attn_out

        x_norm = (1 + scale_mlp.unsqueeze(1)) * self.norm2(x) + shift_mlp.unsqueeze(1)
        x = x + gate_mlp.unsqueeze(1) * self.mlp(x_norm)

        return x

from diffusers import StableDiffusionXLPipeline

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16"
).to("cuda")

image = pipe(
    prompt="a majestic lion in photorealistic style, 4k",
    negative_prompt="cartoon, blurry",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0
).images[0]