This is a simple attempt. I trained with CIFAR-10 dataset.
1# 生成图像有误...以下代码需修改！！！
2
3import torch
4from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
5from PIL import Image
6import os
7import matplotlib.pyplot as plt
8
9# 设备选择
10device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
11
12model_id = "BackTo2014/DDPM-test"
13
14def load_and_eval(checkpoint_path, output_dir="./generated_images"):
15 # 加载 UNet 模型
16 unet = UNet2DModel.from_pretrained(
17 model_id, # 替换为你的模型存储库名称
18 filename=checkpoint_path, # 使用传入的检查点文件名
19 ignore_mismatched_sizes=True,
20 low_cpu_mem_usage=False,
21 ).to(device)
22
23 # 确保 sample_size 是一个有效的尺寸信息
24 if unet.config.sample_size is None:
25 # 假设样本尺寸为 32x32 或者根据你的需求设置
26 unet.config.sample_size = (32, 32)
27
28 # 初始化调度器
29 scheduler = DDPMScheduler.from_config(model_id) # 替换为你的调度器存储库名称
30
31 # 创建管道
32 pipeline = DDPMPipeline(unet=unet, scheduler=scheduler)
33
34 # 设置生成参数
35 num_images = 4 # 生成4张图像
36 generator = torch.manual_seed(0) # 固定随机种子
37 num_inference_steps = 999 # 推理步数
38
39 # 生成图像
40 images = []
41 for _ in range(num_images):
42 image = pipeline(generator=generator, num_inference_steps=num_inference_steps).images[0]
43 images.append(image)
44
45 # 创建输出目录
46 if not os.path.exists(output_dir):
47 os.makedirs(output_dir)
48
49 # 保存图像
50 for i, img in enumerate(images):
51 img.save(os.path.join(output_dir, f"generated_image_{i}.png"))
52
53 # 使用 Matplotlib 显示图像
54 fig, axs = plt.subplots(1, len(images), figsize=(len(images) * 5, 5))
55 for ax, img in zip(axs.flatten(), images):
56 ax.imshow(img)
57 ax.axis('off')
58 plt.show()
59
60if __name__ == "__main__":
61 checkpoint_path = "ckpt_141_.pt" # 检查点文件名
62 load_and_eval(checkpoint_path)
