What are Diffusion Models?

from dataclasses import dataclass

@dataclass
class TrainingConfig:
    image_size = 128
    train_batch_size = 32
    eval_batch_size = 8  # how many images to sample during evaluation
    num_epochs = 50
    gradient_accumulation_steps = 1
    learning_rate = 1e-4
    lr_warmup_steps = 500
    save_image_epochs = 10
    save_model_epochs = 25
    mixed_precision = "fp16"  # `no` for float32, `fp16` for automatic mixed precision
    output_dir = "output"  # the model name locally and on the HF Hub
    seed = 42
    device = "cuda"

config = TrainingConfig()

from datasets import load_dataset

config.dataset_name = "huggan/smithsonian_butterflies_subset"
dataset = load_dataset(config.dataset_name, split="train")

import matplotlib.pyplot as plt

fig, axs = plt.subplots(1, 4, figsize=(16, 4))
for i, image in enumerate(dataset[:4]["image"]):
    axs[i].imshow(image)
    axs[i].set_axis_off()
fig.show()

from torchvision import transforms

preprocess = transforms.Compose([
    transforms.Resize((config.image_size, config.image_size)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5]),
])

def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"img": images}

dataset.set_transform(transform)
print(dataset[0]['img'].shape)

# torch.Size([3, 128, 128])

import torch
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)

from diffusers import UNet2DModel
model = UNet2DModel(
    sample_size=config.image_size,  # the target image resolution
    in_channels=3,  # the number of input channels, 3 for RGB images
    out_channels=3,  # the number of output channels
    layers_per_block=2,  # how many ResNet layers to use per UNet block
    block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block
    down_block_types=(
        "DownBlock2D",  # a regular ResNet downsampling block
        "DownBlock2D",
        "DownBlock2D",
        "DownBlock2D",
        "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
        "DownBlock2D",
    ),
    up_block_types=(
        "UpBlock2D",  # a regular ResNet upsampling block
        "AttnUpBlock2D",  # a ResNet upsampling block with spatial self-attention
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
    ),
)

# Check input and output shapes
sample_image = dataset[0]['img'].unsqueeze(0)
print("Input Shape", sample_image.shape)
print("Output Shape", model(sample_image, timestep=0).sample.shape)

# Create a DDPM scheduler
import torch
from PIL import Image
from diffusers import DDPMScheduler

noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
noise = torch.randn(sample_image.shape)
timesteps = torch.LongTensor([50])
noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)

Image.fromarray(((noisy_image.permute(0,2,3,1)+1.0)*127.5).type(torch.uint8).numpy()[0])

# Create optim and lr scheduler
from diffusers.optimization import get_cosine_schedule_with_warmup

optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
lr_scheduler = get_cosine_schedule_with_warmup(optimizer, 
                                            num_warmup_steps=config.lr_warmup_steps, 
                                            num_training_steps= (len(train_dataloader)*config.num_epochs)
                                            )

noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
loss = F.mse_loss(noise_pred, noise)

from accelerate import Accelerator
from tqdm.auto import tqdm
from pathlib import Path
import torch.nn.functional as F
import os

def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler):
    # Initialize accelerator
    accelerator = Accelerator(
        mixed_precision=config.mixed_precision,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        project_dir=os.path.join(config.output_dir, "logs")
    )
    if accelerator.is_main_process:
        if config.output_dir is not None:
            os.makedirs(config.output_dir, exist_ok=True)
        
    # Prepare everything
    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, lr_scheduler
    )

    global_step = 0

    # Train!
    for epoch in range(config.num_epochs):
        progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
        progress_bar.set_description(f"Epoch {epoch}")

        for step, batch in enumerate(train_dataloader):
            clean_images = batch["img"]

            # Sample noise to add to the clean image
            noise = torch.randn(clean_images.shape, device=config.device)
            bs = clean_images.shape[0]

            # Sample a random timestep for each image
            timesteps = torch.randint(
                0, noise_scheduler.config.num_train_timesteps, (bs,), device=config.device, dtype=torch.int64
            )

            # forward diffusion process
            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

            with accelerator.accumulate(model):
                # Predict Noise residual
                noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
                loss = F.mse_loss(noise_pred, noise)
                accelerator.backward(loss)

                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), 1.0) # Gradient clipping
                
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
            
            progress_bar.update(1)
            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
            progress_bar.set_postfix(**logs)
            accelerator.log(logs, step=global_step)
            global_step += 1

        # Evaluation
        if accelerator.is_main_process:
            pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)

            if (epoch + 1) % config.save_image_epochs == 0:
                evaluate(config, epoch, pipeline)

            if (epoch + 1) % config.save_model_epochs == 0:
                pipeline.save_pretrained(os.path.join(config.output_dir, f"epoch_{epoch}"))
            elif (epoch + 1) == config.num_epochs:
                pipeline.save_pretrained(os.path.join(config.output_dir, f"final"))

# Evaluation
from diffusers import DDPMPipeline
from diffusers.utils import make_image_grid
import os

def evaluate(config, epoch, pipeline):
    images = pipeline(
        batch_size = config.eval_batch_size,
        generator = torch.Generator(device=config.device).manual_seed(config.seed),
    ).images

    image_grid = make_image_grid(images, rows=2, cols=4)

    test_dir = os.path.join(config.output_dir, "test")
    os.makedirs(test_dir, exist_ok=True)
    image_grid.save(os.path.join(test_dir, f"epoch_{epoch}.png"))

rom accelerate import notebook_launcher

args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)

notebook_launcher(train_loop, args, num_processes=1)

# Generate samples
from diffusers import DDPMPipeline
import random
pipeline = DDPMPipeline.from_pretrained("/openbayes/home/miyan/works/Diffusion-Model-0-1/output/epoch_49").to(config.device)

samples_num = 500
batch_size = 20
for epoch in range(samples_num // batch_size):
    images = pipeline(
        batch_size = batch_size,
        generator = [torch.Generator(device=config.device).manual_seed(random.randint(0, 100000)) for _ in range(batch_size)],
    ).images

    test_dir = os.path.join(config.output_dir, "samples")
    os.makedirs(test_dir, exist_ok=True)
    for i, image in enumerate(images):
        image.save(os.path.join(test_dir, f"{i+epoch*batch_size}.png"))

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import numpy as np
from scipy.linalg import sqrtm

def calculate_inception_score(images, device, batch_size=32, splits=10):
    """
    计算 Inception Score (IS)。
    参数:
        images: torch.Tensor，形状为(N, C, H, W) 的生成图像
        device: torch.device，计算设备（CPU或GPU）
        batch_size: 批大小
        splits: 将生成的图片集分为几份计算IS
    返回:
        (is_mean, is_std): IS的均值和标准差
    """
    # 加载预训练的Inception v3模型，用于分类
    inception = models.inception_v3(pretrained=True, transform_input=True).to(device)
    inception.eval()

    preds = []
    # 分批次计算预测概率分布
    with torch.no_grad():
        for i in range(0, len(images), batch_size):
            batch = images[i:i+batch_size].to(device)
            # Inception v3要求输入为299x299，如果 images 已经是此大小且已标准化则无需再次处理
            logits = inception(batch)
            probs = F.softmax(logits, dim=1)
            preds.append(probs.cpu().numpy())
    preds = np.concatenate(preds, axis=0)  # (N, 1000)

    # 计算IS
    N = preds.shape[0]
    split_scores = []
    for k in range(splits):
        part = preds[k * (N // splits) : (k+1) * (N // splits), :]
        p_y = np.mean(part, axis=0)
        scores = []
        for i in range(part.shape[0]):
            p_yx = part[i]
            scores.append(np.sum(p_yx * (np.log(p_yx + 1e-10) - np.log(p_y + 1e-10))))
        split_scores.append(np.exp(np.mean(scores)))

    is_mean = np.mean(split_scores)
    is_std = np.std(split_scores)
    return is_mean, is_std
  
  import os
from PIL import Image
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#######################
# 数据预处理Transform
#######################
# Inception v3预期输入尺寸为299x299，且通常使用标准化到[-1,1]
transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
    transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
])

#######################
# 加载生成的图像 (fake_images)
#######################
def load_fake_images_from_folder(folder, transform):
    images = []
    for filename in os.listdir(folder):
        if filename.lower().endswith(('png','jpg','jpeg')):
            img_path = os.path.join(folder, filename)
            img = Image.open(img_path).convert('RGB')
            img = transform(img)
            images.append(img)
    # 将所有图像合并为一个Tensor: (N, C, H, W)
    if len(images) > 0:
        images = torch.stack(images, dim=0)
    else:
        images = torch.empty(0)  # 如果没有图像则返回空tensor
    return images

fake_folder = "/openbayes/home/miyan/works/Diffusion-Model-0-1/output/samples"
fake_images = load_fake_images_from_folder(fake_folder, transform)  # (N,3,299,299)

is_mean, is_std = calculate_inception_score(fake_images, device)
print("IS:", is_mean, is_std)

def calculate_frechet_distance(mu1, sigma1, mu2, sigma2):
    """
    计算Fréchet Inception Distance所需的Fréchet距离。
    参数:
        mu1, sigma1: 实际数据特征均值和协方差矩阵
        mu2, sigma2: 生成数据特征均值和协方差矩阵
    返回:
        fid: FID分数
    """
    diff = mu1 - mu2
    diff_sq = diff.dot(diff)

    # 计算矩阵的对称矩阵平方根
    covmean, _ = sqrtm(sigma1.dot(sigma2), disp=False)
    if np.iscomplexobj(covmean):
        covmean = covmean.real

    fid = diff_sq + np.trace(sigma1 + sigma2 - 2 * covmean)
    return fid

def calculate_fid(real_features, fake_features):
    """
    计算 Frechet Inception Distance (FID)。
    参数:
        real_features: np.ndarray, shape (N, 2048)，真实图像特征
        fake_features: np.ndarray, shape (M, 2048)，生成图像特征
    返回:
        fid: FID分数（越低越好）
    """
    mu_real = np.mean(real_features, axis=0)
    sigma_real = np.cov(real_features, rowvar=False)
    mu_fake = np.mean(fake_features, axis=0)
    sigma_fake = np.cov(fake_features, rowvar=False)

    fid = calculate_frechet_distance(mu_real, sigma_real, mu_fake, sigma_fake)
    return fid

transform_img = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
    transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
])

dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")
def transform(examples):
    images = [transform_img(image.convert("RGB")) for image in examples["image"]]
    return {"img": images}
dataset.set_transform(transform)

real_loader = DataLoader(dataset, batch_size=32, shuffle=False)
inception = models.inception_v3(pretrained=True, transform_input=True).to(device)
inception.eval()
# 计算FID需要real_features和fake_features（需要先提取特征）
# 可通过迭代real_loader对真实数据提取特征：
real_features_list = []
with torch.no_grad():
    for batch in real_loader:
        imgs = batch["img"]
        imgs = imgs.to(device)
        feats = inception(imgs)
        real_features_list.append(feats.cpu().numpy())
real_features = np.concatenate(real_features_list, axis=0)

# 对 fake_images 同样提取特征
fake_features = []
with torch.no_grad():
    for i in range(0, len(fake_images), 32):
        batch = fake_images[i:i+32].to(device)
        feats = inception(batch)
        fake_features.append(feats.cpu().numpy())
fake_features = np.concatenate(fake_features, axis=0)

fid_score = calculate_fid(real_features, fake_features)
print("FID:", fid_score)

class TrainingConfig:
    image_size = 128
    train_batch_size = 32
    eval_batch_size = 8  # how many images to sample during evaluation
    num_epochs = 100
    gradient_accumulation_steps = 1
    learning_rate = 1e-4
    lr_warmup_steps = 1000
    save_image_epochs = 10
    save_model_epochs = 500
    mixed_precision = "fp16"  # `no` for float32, `fp16` for automatic mixed precision
    output_dir = "output"  # the model name locally and on the HF Hub
    seed = 42
    device = "cuda"

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --checkpointing_steps=5000 \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model"

# Load scheduler, tokenizer and models.
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
    tokenizer = CLIPTokenizer.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
    )

# Load text encoder and vae
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
        text_encoder = CLIPTextModel.from_pretrained(
            args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, variant=args.variant
        )
        vae = AutoencoderKL.from_pretrained(
            args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, variant=args.variant
        )

# Load U-Net
unet = UNet2DConditionModel.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="unet", 			revision=args.non_ema_revision
    )

vae.requires_grad_(False)
text_encoder.requires_grad_(False)
unet.train()

optimizer = optimizer_cls(
        unet.parameters(),
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

# Convert images to latent space
latents =vae.encode(batch["pixel_values"].to(weight_dtype)).latent_dist.sample()
latents = latents * vae.config.scaling_factor

# Get the text embedding for conditioning
encoder_hidden_states = text_encoder(batch["input_ids"], return_dict=False)[0]

# Predict the noise residual and compute loss
model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]

# Loss
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")