|
Original model github address:[DenoisingDiffusionProbabilityModel-ddpm-](https://github.com/zoubohao/DenoisingDiffusionProbabilityModel-ddpm-) |
|
|
|
This is a simple attempt. I trained with CIFAR-10 dataset. |
|
|
|
## Usage |
|
|
|
```python |
|
# 生成图像有误...以下代码需修改!!! |
|
|
|
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") |
|
``` |