From c9e3487ec3c197b38f7661ff51895aa75323aacf Mon Sep 17 00:00:00 2001 From: gdamms Date: Wed, 29 May 2024 14:03:31 +0200 Subject: messing with latent diff --- plots.py | 162 ++++++++++++++++++++++++++++++++++----------------------------- 1 file changed, 87 insertions(+), 75 deletions(-) (limited to 'plots.py') diff --git a/plots.py b/plots.py index 9137f24..0513495 100644 --- a/plots.py +++ b/plots.py @@ -7,6 +7,7 @@ import os from rich.progress import track from main import UNet, q_xt_xt_1, p_xt_1_xt, tensor_to_image +from autoencoder import Autoencoder os.makedirs('plots', exist_ok=True) @@ -27,6 +28,16 @@ BIN_MIN = -4 BIN_MAX = 4 + +model = UNet().to(DEVICE) +model.load_state_dict(torch.load('model.pth')) + +autoencoder = Autoencoder(input_dim=(1, 28, 28), latent_dim=(1, 8, 8)).to(DEVICE) +autoencoder.load_state_dict(torch.load('autoencoder.pth')) + + + + # plt.figure() # plt.plot(BETA, label='beta') # plt.plot(ALPHA, label='alpha') @@ -39,98 +50,99 @@ BIN_MAX = 4 mnist = datasets.MNIST('data', train=True, download=True) img, label = mnist[np.random.randint(0, len(mnist))] img = np.array(img) / 255 * 2 - 1 +img = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) -# plt.figure() -# plt.imshow(img, cmap='gray') -# plt.title('Image') -# plt.axis('off') -# plt.savefig('plots/img.tmp.png') +encoded = autoencoder.encode(img) +img = encoded.squeeze().cpu().detach().numpy() -# plt.figure() -# plt.hist(img.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) -# plt.yscale('log') -# plt.title('Image histogram') -# plt.savefig('plots/img_hist.tmp.png') +plt.figure() +plt.imshow(img, cmap='gray') +plt.title('Image') +plt.axis('off') +plt.savefig('plots/img.tmp.png') -# def norm_dist(x, mean, std): -# return np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi)) +plt.figure() +plt.hist(img.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) +plt.yscale('log') +plt.title('Image histogram') +plt.savefig('plots/img_hist.tmp.png') -# x_norm = np.linspace(BIN_MIN, BIN_MAX, 100) -# y_norm = norm_dist(x_norm, 0, 1) * 28**2 / NB_BINS * (BIN_MAX - BIN_MIN) +def norm_dist(x, mean, std): + return np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi)) -# fig = plt.figure(figsize=(10, 5)) -# fig.suptitle('Diffusion naturelle') -# gs = GridSpec(1, 3, figure=fig) -# ax1 = fig.add_subplot(gs[0, 0]) -# ax2 = fig.add_subplot(gs[0, 1:]) +x_norm = np.linspace(BIN_MIN, BIN_MAX, 100) +y_norm = norm_dist(x_norm, 0, 1) * 8**2 / NB_BINS * (BIN_MAX - BIN_MIN) -# plots_to_save = np.linspace(1, DIFFU_STEPS, 100).astype(int) +fig = plt.figure(figsize=(10, 5)) +fig.suptitle('Diffusion naturelle') +gs = GridSpec(1, 3, figure=fig) +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1:]) -# xt = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) -# for t in track(range(1, DIFFU_STEPS+1)): -# xt, eps = q_xt_xt_1(xt, t) +plots_to_save = np.linspace(1, DIFFU_STEPS, 100).astype(int) -# if t not in plots_to_save: -# continue +xt = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) +for t in track(range(1, DIFFU_STEPS+1)): + xt, eps = q_xt_xt_1(xt, t) -# xt_numpy = xt.cpu().detach().numpy()[0, 0] + if t not in plots_to_save: + continue -# ax1.clear() -# ax1.imshow(xt_numpy, cmap='gray') -# ax1.set_title(f'xt at t={t:04d}') -# ax1.axis('off') + xt_numpy = xt.cpu().detach().numpy()[0, 0] -# ax2.clear() -# ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) -# ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') -# ax2.legend() -# ax2.set_yscale('log') -# ax2.set_ylim(y_norm.min(), 1e3) -# ax2.set_title(f'xt histogram') + ax1.clear() + ax1.imshow(xt_numpy, cmap='gray') + ax1.set_title(f'xt at t={t:04d}') + ax1.axis('off') -# fig.savefig(f'plots/diffusion/{t:04d}.tmp.png') -# os.system('convert -delay 20 -loop 0 plots/diffusion/*.png plots/diffusion.tmp.gif') + ax2.clear() + ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) + ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') + ax2.legend() + ax2.set_yscale('log') + ax2.set_ylim(y_norm.min(), 1e3) + ax2.set_title(f'xt histogram') + fig.savefig(f'plots/diffusion/{t:04d}.tmp.png') +os.system('convert -delay 20 -loop 0 plots/diffusion/*.png plots/diffusion.tmp.gif') +fig = plt.figure(figsize=(10, 5)) +fig.suptitle('Diffusion inverse') +gs = GridSpec(1, 3, figure=fig) +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1:]) -model = UNet().to(DEVICE) -model.load_state_dict(torch.load('mnist_model.pth')) - -# fig = plt.figure(figsize=(10, 5)) -# fig.suptitle('Diffusion inverse') -# gs = GridSpec(1, 3, figure=fig) -# ax1 = fig.add_subplot(gs[0, 0]) -# ax2 = fig.add_subplot(gs[0, 1:]) - -# xt = torch.randn(1, 1, 28, 28, device=DEVICE) -# vec = torch.zeros(1, 10).to(DEVICE) -# vec[0, label] = 1 -# for t in track(range(DIFFU_STEPS, 0, -1)): -# t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32) -# xt = p_xt_1_xt(model, xt, t_tensor, vec) - -# if t not in plots_to_save: -# continue - -# xt_numpy = xt.cpu().detach().numpy()[0, 0] - -# ax1.clear() -# ax1.imshow(xt_numpy, cmap='gray') -# ax1.set_title(f'xt at t={t:04d}') -# ax1.axis('off') - -# ax2.clear() -# ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) -# ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') -# ax2.legend() -# ax2.set_yscale('log') -# ax2.set_ylim(y_norm.min(), 1e3) -# ax2.set_title(f'xt histogram') - -# fig.savefig(f'plots/diffusion_inverse/{t:04d}.tmp.png') -# os.system('convert -delay 20 -loop 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif') +xt = torch.randn(1, 1, 28, 28, device=DEVICE) +vec = torch.zeros(1, 10).to(DEVICE) +vec[0, label] = 1 +for t in track(range(DIFFU_STEPS, 0, -1)): + t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32) + xt = p_xt_1_xt(model, xt, t_tensor, vec) + + if t not in plots_to_save: + continue + + xt_numpy = xt.cpu().detach().numpy()[0, 0] + + ax1.clear() + ax1.imshow(xt_numpy, cmap='gray') + ax1.set_title(f'xt at t={t:04d}') + ax1.axis('off') + + ax2.clear() + ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) + ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') + ax2.legend() + ax2.set_yscale('log') + ax2.set_ylim(y_norm.min(), 1e3) + ax2.set_title(f'xt histogram') + + fig.savefig(f'plots/diffusion_inverse/{t:04d}.tmp.png') +os.system('convert -delay 20 -loop 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif') + +exit(0) tpause = {150: 'xt', 20: 'mu', 50: 'xt_1'} -- cgit v1.3.1