Update README.md

README.md

@@ -5,9 +5,67 @@ This is a simple attempt. I trained with CIFAR-10 dataset.
 ## Usage
 
 ```python
-from diffusers import DDPMPipeline
+# 生成图像有误...以下代码需修改！！！
 
-pipeline = DDPMPipeline.from_pretrained('BackTo2014/DDPM-test')
-image = pipeline().images[0]
-image
+import torch
+from diffusers import DDPMPipeline, DDPMScheduler
+from diffusers.models import UNet2DModel
+from PIL import Image
+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:
+    unet = UNet2DModel.from_pretrained(
+        model_id,
+        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")
 ```