README.md

Original model github address：DenoisingDiffusionProbabilityModel-ddpm-

This is a simple attempt. I trained with CIFAR-10 dataset.

Usage

# 生成图像有误...以下代码需修改！！！

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")