Diffusers-5.py
· 749 B · Python
Исходник
from diffusers import DDPMScheduler, UNet2DModel
from PIL import Image
import torch
scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda")
scheduler.set_timesteps(50)
sample_size = model.config.sample_size
noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")
input = noise
for t in scheduler.timesteps:
with torch.no_grad():
noisy_residual = model(input, t).sample
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
input = prev_noisy_sample
image = (input / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
image = Image.fromarray((image * 255).round().astype("uint8"))
image
| 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 |