105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
|
import os
|
||
|
|
import argparse
|
||
|
|
import math
|
||
|
|
import torch
|
||
|
|
from torch.nn import functional as F
|
||
|
|
from torchvision import transforms
|
||
|
|
import numpy as np
|
||
|
|
import matplotlib
|
||
|
|
matplotlib.use("Agg")
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from main import UNet
|
||
|
|
|
||
|
|
|
||
|
|
def build_schedule(timesteps, device):
|
||
|
|
betas = torch.linspace(1e-4, 0.02, timesteps).to(device)
|
||
|
|
alphas = 1.0 - betas
|
||
|
|
alpha_bar = torch.cumprod(alphas, dim=0)
|
||
|
|
return betas, alphas, alpha_bar
|
||
|
|
|
||
|
|
|
||
|
|
def denoise_step(model, x, t, alphas, alpha_bar, betas):
|
||
|
|
n = x.size(0)
|
||
|
|
t_batch = torch.full((n,), t, device=x.device, dtype=torch.long)
|
||
|
|
eps_pred = model(x, t_batch)
|
||
|
|
alpha_t = alphas[t]
|
||
|
|
alpha_bar_t = alpha_bar[t]
|
||
|
|
x = (x - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_t)
|
||
|
|
if t > 0:
|
||
|
|
x += torch.sqrt(betas[t]) * torch.randn_like(x)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
def to_grid_image(x, ncols, cmap="gray"):
|
||
|
|
n = x.size(0)
|
||
|
|
nrows = math.ceil(n / ncols)
|
||
|
|
fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 2, nrows * 2))
|
||
|
|
axes = np.atleast_1d(axes).flatten()
|
||
|
|
for i, ax in enumerate(axes):
|
||
|
|
if i < n:
|
||
|
|
ax.imshow(x[i, 0], cmap=cmap)
|
||
|
|
ax.axis("off")
|
||
|
|
plt.subplots_adjust(wspace=0.05, hspace=0.05)
|
||
|
|
fig.canvas.draw()
|
||
|
|
buf = np.asarray(fig.canvas.buffer_rgba())[:, :, :3]
|
||
|
|
plt.close(fig)
|
||
|
|
return Image.fromarray(buf)
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description="DDPM 采样过程逐帧记录")
|
||
|
|
parser.add_argument("--ckpt", type=str, default="model_state_dict.pth")
|
||
|
|
parser.add_argument("--timesteps", type=int, default=600)
|
||
|
|
parser.add_argument("--samples", type=int, default=4)
|
||
|
|
parser.add_argument("--outdir", type=str, default="sampling_frames")
|
||
|
|
parser.add_argument("--gif", action="store_true", help="额外输出 GIF")
|
||
|
|
parser.add_argument("--gif-fps", type=int, default=20)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
|
|
||
|
|
model = UNet().to(device)
|
||
|
|
state = torch.load(args.ckpt, map_location=device)
|
||
|
|
if isinstance(state, dict) and "state_dict" in state:
|
||
|
|
state = state["state_dict"]
|
||
|
|
model.load_state_dict(state)
|
||
|
|
model.eval()
|
||
|
|
|
||
|
|
betas, alphas, alpha_bar = build_schedule(args.timesteps, device)
|
||
|
|
|
||
|
|
os.makedirs(args.outdir, exist_ok=True)
|
||
|
|
|
||
|
|
x = torch.randn(args.samples, 1, 28, 28, device=device)
|
||
|
|
|
||
|
|
frames = []
|
||
|
|
with torch.no_grad():
|
||
|
|
for t in reversed(range(args.timesteps)):
|
||
|
|
x = denoise_step(model, x, t, alphas, alpha_bar, betas)
|
||
|
|
grid = x.cpu().clamp(-1, 1)
|
||
|
|
grid = (grid + 1) / 2
|
||
|
|
img = to_grid_image(grid, ncols=2)
|
||
|
|
path = os.path.join(args.outdir, f"step_{t:04d}.png")
|
||
|
|
img.save(path)
|
||
|
|
frames.append(path)
|
||
|
|
if t % 50 == 0 or t == 0:
|
||
|
|
print(f"step {t} 已保存 -> {path}")
|
||
|
|
|
||
|
|
print(f"全部 {args.timesteps} 帧已保存到 {args.outdir}/")
|
||
|
|
|
||
|
|
if args.gif:
|
||
|
|
images = [Image.open(p) for p in frames]
|
||
|
|
gif_path = os.path.join(args.outdir, "sampling.gif")
|
||
|
|
images[0].save(
|
||
|
|
gif_path,
|
||
|
|
save_all=True,
|
||
|
|
append_images=images[1:],
|
||
|
|
duration=1000 // args.gif_fps,
|
||
|
|
loop=0,
|
||
|
|
)
|
||
|
|
print(f"GIF 已保存 -> {gif_path}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|