File size: 3,629 Bytes
38b4eff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from diffusers import UNet2DModel, DDPMScheduler
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
import os
from PIL import Image
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1. Config
image_size = 28
train_batch_size = 128
num_epochs = 1
learning_rate = 1e-4
# 2. Model: UNet2DModel from diffusers (small version for toy)
model = UNet2DModel(
sample_size=image_size,
in_channels=1,
out_channels=1,
layers_per_block=2,
block_out_channels=(32, 64, 64),
down_block_types=(
"DownBlock2D",
"AttnDownBlock2D",
"DownBlock2D",
),
up_block_types=(
"UpBlock2D",
"AttnUpBlock2D",
"UpBlock2D",
),
).to(device)
# 3. Scheduler
noise_scheduler = DDPMScheduler(num_train_timesteps=100)
# 4. Data Loading (MNIST)
preprocess = transforms.Compose([
transforms.Resize((image_size, image_size)),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=preprocess)
train_dataloader = DataLoader(dataset, batch_size=train_batch_size, shuffle=True)
# 5. Training
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
print(f"Starting NeuralAI Diffusion Toy V2 Training on {device}...")
for epoch in range(num_epochs):
losses = []
for step, (images, _) in enumerate(tqdm(train_dataloader)):
images = images.to(device)
noise = torch.randn(images.shape).to(device)
bs = images.shape[0]
# Sample a random timestep for each image
timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bs,), device=device).long()
# Add noise to the clean images according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_images = noise_scheduler.add_noise(images, noise, timesteps)
# Predict the noise residual
noise_pred = model(noisy_images, timesteps).sample
loss = F.mse_loss(noise_pred, noise)
loss.backward()
optimizer.step()
optimizer.zero_grad()
losses.append(loss.item())
print(f"Epoch {epoch} | Loss: {np.mean(losses):.6f}")
# 6. Sampling Logic
def generate_samples(model, scheduler, num_samples=1):
model.eval()
# Start from random noise
sample = torch.randn(num_samples, 1, image_size, image_size).to(device)
for t in tqdm(scheduler.timesteps):
with torch.no_grad():
residual = model(sample, t).sample
# Compute previous image: x_t -> x_t-1
sample = scheduler.step(residual, t, sample).prev_sample
return sample
print("Generating NeuralAI Diffusion sample...")
generated = generate_samples(model, noise_scheduler)
# Save result
output_dir = "/home/workspace/Projects/NeuralAI/storage/images"
os.makedirs(output_dir, exist_ok=True)
# Convert to PIL and save
gen_img = (generated[0] / 2 + 0.5).clamp(0, 1).cpu().numpy().squeeze()
gen_img = (gen_img * 255).astype(np.uint8)
img = Image.fromarray(gen_img)
img.save(os.path.join(output_dir, "toy_v2_sample.png"))
# Save model checkpoint
checkpoint_dir = "/home/workspace/Projects/NeuralAI/checkpoints/diffusion_toy"
os.makedirs(checkpoint_dir, exist_ok=True)
torch.save(model.state_dict(), os.path.join(checkpoint_dir, "unet_toy.pt"))
print(f"Sample saved to {output_dir}/toy_v2_sample.png")
print(f"Model saved to {checkpoint_dir}/unet_toy.pt")
|