knox 修订了这个 Gist . 跳至此修订
1 file changed, 22 insertions
Diffusers-5.py(file created)
| @@ -0,0 +1,22 @@ | |||
| 1 | + | from diffusers import DDPMScheduler, UNet2DModel | |
| 2 | + | from PIL import Image | |
| 3 | + | import torch | |
| 4 | + | ||
| 5 | + | scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256") | |
| 6 | + | model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda") | |
| 7 | + | scheduler.set_timesteps(50) | |
| 8 | + | ||
| 9 | + | sample_size = model.config.sample_size | |
| 10 | + | noise = torch.randn((1, 3, sample_size, sample_size), device="cuda") | |
| 11 | + | input = noise | |
| 12 | + | ||
| 13 | + | for t in scheduler.timesteps: | |
| 14 | + | with torch.no_grad(): | |
| 15 | + | noisy_residual = model(input, t).sample | |
| 16 | + | prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample | |
| 17 | + | input = prev_noisy_sample | |
| 18 | + | ||
| 19 | + | image = (input / 2 + 0.5).clamp(0, 1) | |
| 20 | + | image = image.cpu().permute(0, 2, 3, 1).numpy()[0] | |
| 21 | + | image = Image.fromarray((image * 255).round().astype("uint8")) | |
| 22 | + | image | |
更新
更早