Update README.md

README.md

@@ -8,64 +8,63 @@ This is a simple attempt. I trained with CIFAR-10 dataset.
 # 生成图像有误...以下代码需修改！！！
 
 import torch
-from diffusers import DDPMPipeline, DDPMScheduler
-from diffusers.models import UNet2DModel
+from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
 from PIL import Image
+import os
 import matplotlib.pyplot as plt
 
-# 模型ID
-model_id = "BackTo2014/DDPM-test"
-
-# 检查设备
+# 设备选择
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# 加载UNet模型和配置文件
-try:
+model_id = "BackTo2014/DDPM-test"
+
+def load_and_eval(checkpoint_path, output_dir="./generated_images"):
+    # 加载 UNet 模型
     unet = UNet2DModel.from_pretrained(
-        model_id,
+        model_id,  # 替换为你的模型存储库名称
+        filename=checkpoint_path,  # 使用传入的检查点文件名
         ignore_mismatched_sizes=True,
         low_cpu_mem_usage=False,
-    ).to(device)  # 将模型移动到GPU上
-except ValueError as e:
-    print(f"Error loading model: {e}")
-    
-    # 获取模型的state_dict
-    state_dict = unet.state_dict()
-    
-    # 手动初始化缺失的权重
-    for key in e.args[0].split(': ')[1].split(', '):
-        name, size = key.split('.')
-        size = tuple(map(int, size.replace(')', '').replace('(', '').split(',')))
-        
-        # 创建随机权重
-        new_weight = torch.randn(size).to(device)  # 将权重移动到GPU上
-        
-        # 更新state_dict
-        state_dict[name] = new_weight
-    
-    # 加载更新后的state_dict
-    unet.load_state_dict(state_dict).to(device)  # 将模型移动到GPU上
-
-# 如果sample_size未定义，则手动设置
-if unet.config.sample_size is None:
-    # 假设样本尺寸为 32x32
-    unet.config.sample_size = (32, 32)
-
-# 初始化Scheduler
-scheduler = DDPMScheduler.from_config(model_id)
-
-# 创建DDPMPipeline
-pipeline = DDPMPipeline(unet=unet, scheduler=scheduler)
-
-# 生成图像
-generator = torch.manual_seed(0)
-image = pipeline(num_inference_steps=1000, generator=generator).images[0]
-
-# 使用matplotlib显示图像
-plt.imshow(image)
-plt.axis('off')  # 不显示坐标轴
-plt.show()
-
-# 保存图像
-image.save("generated_image.png")
+    ).to(device)
+
+    # 确保 sample_size 是一个有效的尺寸信息
+    if unet.config.sample_size is None:
+        # 假设样本尺寸为 32x32 或者根据你的需求设置
+        unet.config.sample_size = (32, 32)
+
+    # 初始化调度器
+    scheduler = DDPMScheduler.from_config(model_id)  # 替换为你的调度器存储库名称
+
+    # 创建管道
+    pipeline = DDPMPipeline(unet=unet, scheduler=scheduler)
+
+    # 设置生成参数
+    num_images = 4  # 生成4张图像
+    generator = torch.manual_seed(0)  # 固定随机种子
+    num_inference_steps = 999  # 推理步数
+
+    # 生成图像
+    images = []
+    for _ in range(num_images):
+        image = pipeline(generator=generator, num_inference_steps=num_inference_steps).images[0]
+        images.append(image)
+
+    # 创建输出目录
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    # 保存图像
+    for i, img in enumerate(images):
+        img.save(os.path.join(output_dir, f"generated_image_{i}.png"))
+
+    # 使用 Matplotlib 显示图像
+    fig, axs = plt.subplots(1, len(images), figsize=(len(images) * 5, 5))
+    for ax, img in zip(axs.flatten(), images):
+        ax.imshow(img)
+        ax.axis('off')
+    plt.show()
+
+if __name__ == "__main__":
+    checkpoint_path = "ckpt_141_.pt"  # 检查点文件名
+    load_and_eval(checkpoint_path)
 ```