（9-4）基于Diffusion Transformer的文生图系统：生成图像

9.6  生成图像

在本项目中，使用分布式数据并行（DDP）在多个GPU上进行训练，以生成高质量的图像。通过对输入数据进行处理和增强，将图像输入到深度学习模型中，使用自适应动量估计（EMA）来优化模型参数，并最终将生成的图像保存到指定路径。这一流程支持大规模数据集，旨在提升训练效率和图像生成的效果。

9.6.1  预训练生成

文件sample.py的功能是利用训练好的 DiT 模型生成图像样本，该文件允许用户设置模型、图像尺寸、类别数量和采样参数，加载相应的模型和变分自编码器（VAE），并通过生成噪声和条件标签进行采样。最终，生成的图像会被保存为 "sample.png" 文件。通过设置参数，用户可以灵活控制生成的样本质量和样式。

import torch
# 允许使用 TensorFloat-32（TF32）加速计算
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
from torchvision.utils import save_image
from diffusion import create_diffusion
from diffusers.models import AutoencoderKL
from download import find_model
from models import DiT_models
import argparse


def main(args):
    # 设置 PyTorch 随机种子和禁用梯度计算
    torch.manual_seed(args.seed)
    torch.set_grad_enabled(False)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # 检查是否提供了检查点路径
    if args.ckpt is None:
        assert args.model == "DiT-XL/2", "只有 DiT-XL/2 模型可自动下载。"
        assert args.image_size in [256, 512]
        assert args.num_classes == 1000

    # 加载模型：
    latent_size = args.image_size // 8
    model = DiT_models[args.model](
        input_size=latent_size,
        num_classes=args.num_classes
    ).to(device)
    
    # 自动下载预训练模型或加载自定义的 DiT 检查点
    ckpt_path = args.ckpt or f"DiT-XL-2-{args.image_size}x{args.image_size}.pt"
    state_dict = find_model(ckpt_path)
    model.load_state_dict(state_dict)
    model.eval()  # 将模型设置为评估模式
    diffusion = create_diffusion(str(args.num_sampling_steps))
    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)

    # 设置用于条件生成模型的标签（可以根据需要更改）：
    class_labels = [207, 360, 387, 974, 88, 979, 417, 279]

    # 创建采样噪声：
    n = len(class_labels)
    z = torch.randn(n, 4, latent_size, latent_size, device=device)
    y = torch.tensor(class_labels, device=device)

    # 设置无分类器引导：
    z = torch.cat([z, z], 0)
    y_null = torch.tensor([1000] * n, device=device)
    y = torch.cat([y, y_null], 0)
    model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)

    # 进行图像采样：
    samples = diffusion.p_sample_loop(
        model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
    )
    samples, _ = samples.chunk(2, dim=0)  # 移除空类样本
    samples = vae.decode(samples / 0.18215).sample

    # 保存和显示图像：
    save_image(samples, "sample.png", nrow=4, normalize=True, value_range=(-1, 1))


if __name__ == "__main__":
    # 定义命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-XL/2")
    parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="mse")
    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
    parser.add_argument("--num-classes", type=int, default=1000)
    parser.add_argument("--cfg-scale", type=float, default=4.0)
    parser.add_argument("--num-sampling-steps", type=int, default=250)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--ckpt", type=str, default=None,
                        help="可选的 DiT 检查点路径（默认：自动下载预训练的 DiT-XL/2 模型）。")
    args = parser.parse_args()
    main(args)

9.6.2  基于DDP的图像生成

文件sample_ddp.py实现了基于分布式数据并行（DDP）的图像生成，通过加载训练好的 DiT 模型并利用扩散模型进行采样。首先，它设置了分布式环境并初始化随机种子，确保每个 GPU 在采样时能够生成不同的随机数。接着，它加载预训练模型和相应的变换器，创建保存采样结果的文件夹，并计算每个 GPU 需要生成的样本数量。代码通过循环进行图像采样，使用分类无关引导（CFG）可选地增强生成效果，最后将生成的图像保存为 PNG 文件，并在所有进程完成后将图像汇总为一个 NPZ 文件，方便后续使用和分析。

def create_npz_from_sample_folder(sample_dir, num=50_000):
    """
    从文件夹中的 .png 样本构建一个单一的 .npz 文件。
    """
    samples = []
    for i in tqdm(range(num), desc="Building .npz file from samples"):
        sample_pil = Image.open(f"{sample_dir}/{i:06d}.png")  # 打开样本图像
        sample_np = np.asarray(sample_pil).astype(np.uint8)  # 将图像转换为 numpy 数组
        samples.append(sample_np)
    samples = np.stack(samples)  # 将所有样本堆叠成一个数组
    assert samples.shape == (num, samples.shape[1], samples.shape[2], 3)  # 确保样本形状正确
    npz_path = f"{sample_dir}.npz"  # .npz 文件路径
    np.savez(npz_path, arr_0=samples)  # 保存为 .npz 文件
    print(f"Saved .npz file to {npz_path} [shape={samples.shape}].")
    return npz_path


def main(args):
    """
    运行采样。
    """
    torch.backends.cuda.matmul.allow_tf32 = args.tf32  # 允许使用 TF32，提升性能，但可能导致小的数值差异
    assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage"
    torch.set_grad_enabled(False)  # 禁用梯度计算

    # 设置分布式数据并行（DDP）:
    dist.init_process_group("nccl")  # 初始化进程组
    rank = dist.get_rank()  # 获取当前进程的排名
    device = rank % torch.cuda.device_count()  # 根据排名选择设备
    seed = args.global_seed * dist.get_world_size() + rank  # 计算随机种子
    torch.manual_seed(seed)  # 设置随机种子
    torch.cuda.set_device(device)  # 设置当前设备
    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")

    # 检查模型和参数
    if args.ckpt is None:
        assert args.model == "DiT-XL/2", "Only DiT-XL/2 models are available for auto-download."
        assert args.image_size in [256, 512]
        assert args.num_classes == 1000

    # 加载模型:
    latent_size = args.image_size // 8  # 计算潜在空间大小
    model = DiT_models[args.model](
        input_size=latent_size,
        num_classes=args.num_classes
    ).to(device)  # 将模型移动到设备上

    # 自动下载预训练模型或加载自定义检查点:
    ckpt_path = args.ckpt or f"DiT-XL-2-{args.image_size}x{args.image_size}.pt"
    state_dict = find_model(ckpt_path)  # 查找模型权重
    model.load_state_dict(state_dict)  # 加载权重
    model.eval()  # 切换到评估模式
    diffusion = create_diffusion(str(args.num_sampling_steps))  # 创建扩散模型
    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)  # 加载 VAE
    assert args.cfg_scale >= 1.0, "In almost all cases, cfg_scale must be >= 1.0"
    using_cfg = args.cfg_scale > 1.0  # 判断是否使用分类器自由引导

    # 创建保存样本的文件夹:
    model_string_name = args.model.replace("/", "-")  # 替换模型字符串
    ckpt_string_name = os.path.basename(args.ckpt).replace(".pt", "") if args.ckpt else "pretrained"  # 检查点名称
    folder_name = f"{model_string_name}-{ckpt_string_name}-size-{args.image_size}-vae-{args.vae}-" \
                  f"cfg-{args.cfg_scale}-seed-{args.global_seed}"  # 生成文件夹名称
    sample_folder_dir = f"{args.sample_dir}/{folder_name}"  # 保存样本的文件夹路径
    if rank == 0:
        os.makedirs(sample_folder_dir, exist_ok=True)  # 创建文件夹
        print(f"Saving .png samples at {sample_folder_dir}")
    dist.barrier()  # 所有进程同步

    # 计算每个 GPU 需要生成的样本数量以及需要运行的迭代次数:
    n = args.per_proc_batch_size  # 每个进程的批处理大小
    global_batch_size = n * dist.get_world_size()  # 全局批处理大小
    # 为了使数量可被整除，我们将多采样一点，然后丢弃多余的样本:
    total_samples = int(math.ceil(args.num_fid_samples / global_batch_size) * global_batch_size)
    if rank == 0:
        print(f"Total number of images that will be sampled: {total_samples}")
    assert total_samples % dist.get_world_size() == 0, "total_samples must be divisible by world_size"
    samples_needed_this_gpu = int(total_samples // dist.get_world_size())  # 每个 GPU 需要的样本数量
    assert samples_needed_this_gpu % n == 0, "samples_needed_this_gpu must be divisible by the per-GPU batch size"
    iterations = int(samples_needed_this_gpu // n)  # 迭代次数
    pbar = range(iterations)
    pbar = tqdm(pbar) if rank == 0 else pbar  # 进度条
    total = 0
    for _ in pbar:
        # 生成样本输入:
        z = torch.randn(n, model.in_channels, latent_size, latent_size, device=device)  # 随机噪声
        y = torch.randint(0, args.num_classes, (n,), device=device)  # 随机类别标签

        # 设置分类器自由引导:
        if using_cfg:
            z = torch.cat([z, z], 0)  # 复制噪声
            y_null = torch.tensor([1000] * n, device=device)  # 无效类别标签
            y = torch.cat([y, y_null], 0)  # 合并标签
            model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)  # 模型参数
            sample_fn = model.forward_with_cfg  # 使用带有配置的前向方法
        else:
            model_kwargs = dict(y=y)  # 模型参数
            sample_fn = model.forward  # 使用普通的前向方法

        # 采样图像:
        samples = diffusion.p_sample_loop(
            sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=False, device=device
        )
        if using_cfg:
            samples, _ = samples.chunk(2, dim=0)  # 移除无效类别样本

        samples = vae.decode(samples / 0.18215).sample  # 解码样本
        samples = torch.clamp(127.5 * samples + 128.0, 0, 255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()  # 转换为图像格式

        # 将样本保存为单个 .png 文件
        for i, sample in enumerate(samples):
            index = i * dist.get_world_size() + rank + total  # 计算保存的索引
            Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png")  # 保存图像
        total += global_batch_size  # 更新总样本数量

    # 确保所有进程在尝试转换为 .npz 之前已完成保存其样本
    dist.barrier()
    if rank == 0:
        create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)  # 转换为 .npz 文件
        print("Done.")
    dist.barrier()
    dist.destroy_process_group()  # 销毁进程组


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-XL/2")
    parser.add_argument("--vae",  type=str, choices=["ema", "mse"], default="ema")
    parser.add_argument("--sample-dir", type=str, default="samples")
    parser.add_argument("--per-proc-batch-size", type=int, default=32)
    parser.add_argument("--num-fid-samples", type=int, default=50_000)
    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
    parser.add_argument("--num-classes", type=int, default=1000)
    parser.add_argument("--cfg-scale",  type=float, default=1.5)
    parser.add_argument("--num-sampling-steps", type=int, default=250)
    parser.add_argument("--global-seed", type=int, default=0)
    parser.add_argument("--tf32", action=argparse.BooleanOptionalAction, default=True,
                        help="默认使用 TF32 矩阵乘法。这大幅加速了在 Ampere GPU 上的采样。")
    parser.add_argument("--ckpt", type=str, default=None,
                        help="可选的 DiT 检查点路径（默认：自动下载预训练的 DiT-XL/2 模型）。")
    args = parser.parse_args()
    main(args)

9.7  调试运行

1. 设置环境

（1）创建 Conda 环境：在项目中提供了一个 environment.yml 文件，用于创建 Conda 环境。如果您只想在 CPU 上运行预训练模型，可以从该文件中移除 cudatoolkit 和 pytorch-cuda 的依赖项。使用以下命令创建环境：

conda env create -f environment.yml

（2）激活环境：创建完成后，通过如下命令激活新的 Conda 环境：

conda activate DiT

通过上述步骤，将成功设置好运行该项目所需的环境。接下来，可以开始使用预训练模型或进行模型训练。

2. 训练前的准备

要在一个节点的单个 GPU 上提取 ImageNet 特征，请运行以下命令：

torchrun --nnodes=1 --nproc_per_node=1 extract_features.py --model DiT-XL/2 --data-path /path/to/imagenet/train --features-path /path/to/store/features

1. torchrun：用于运行分布式训练的命令。
2. --nnodes=1：指定使用一个节点。
3. --nproc_per_node=1：指定每个节点使用一个进程（GPU）。
4. extract_features.py：提取特征的脚本。
5. --model DiT-XL/2：指定要使用的模型。
6. --data-path /path/to/imagenet/train：指定包含 ImageNet 训练数据的路径。
7. --features-path /path/to/store/features：指定存储提取特征的路径。

通过上述命令，将能够成功提取 ImageNet 数据集的特征，为后续的模型训练做好准备。

3. 训练 DiT

本项目提供了多个训练DiT模型的脚步文件，其中默认的脚本文件是train.py。该文件train.py用于训练基于类别条件的 DiT 模型，但也可以轻松修改以支持其他类型的条件。

（1）使用单个 GPU 在一个节点上启动 DiT-XL/2 (256x256) 训练：

accelerate launch --mixed_precision fp16 train.py --model DiT-XL/2 --features-path /path/to/store/features

（2）使用多个 GPU 在一个节点上启动 DiT-XL/2 (256x256) 训练：

accelerate launch --multi_gpu --num_processes N --mixed_precision fp16 train.py --model DiT-XL/2 --features-path /path/to/store/features

1. accelerate launch：用于启动训练的命令。
2. --mixed_precision fp16：启用混合精度训练以提高性能。
3. --multi_gpu：指示使用多个 GPU 进行训练。
4. --num_processes N：指定使用的 GPU 数量。
5. --features-path /path/to/store/features：指定存储提取特征的路径。

此外，还可以使用“train_options”文件夹中的训练脚本，通过里面的脚本文件，可以根据需要灵活地进行训练。

4. 生成图像

（1）可以使用文件 sample.py 从提供的预训练 DiT 模型进行采样，会根据使用的模型自动下载预训练 DiT 模型的权重。该脚本具有多种参数，可以在 256x256 和 512x512 模型之间切换、调整采样步骤、改变无分类器引导尺度等。例如，要从512x512 DiT-XL/2 模型进行采样，可以使用以下命令实现：

python sample.py --image-size 512 --seed 1

（2）自定义 DiT 检查点

如果使用文件 train.py 训练了新的 DiT 模型，可以添加 --ckpt 参数来使用自己的检查点。例如，要从自定义的 256x256 DiT-L/4 模型的 EMA 权重中进行采样，可以运行以下命令实现：

python sample.py --model DiT-L/4 --image-size 256 --ckpt /path/to/model.pt

通过这些选项，可以灵活地使用预训练模型或自定义模型进行图像采样。

5. 性能评估

通过使用文件sample_ddp.py 脚本，可以并行从 DiT 模型中采样大量图像，并生成一个样本文件夹以及一个 .npz 文件。这个 .npz 文件可以直接用于 ADM 的 TensorFlow 评估工具，计算 FID（Fréchet Inception Distance）、Inception Score 和其他评估指标。具体地，可以通过指定 GPU 数量和要生成的样本数量来执行此操作，以便获取模型的性能评估数据。例如运行如下命令在分布式环境中并行生成图像并进行性能评估。

torchrun --nnodes=1 --nproc_per_node=N sample_ddp.py --model DiT-XL/2 --num-fid-samples 50000

运行上述命令后，将在一个节点上并行启动多个进程，从 DiT-XL/2 模型中生成 50000 张图像，并为后续的 FID 计算做好准备。

最终，本项目通过使用预训练的 DiT 模型，用户可以根据文本输入生成高质量的图像，如图9-8所示。项目中的 sample.py 脚本支持使用训练好的模型进行图像生成，用户可以调整参数以满足不同的需求。整体上，项目提供了从特征提取到训练和图像生成的完整流程。

图9-8  生成的图像