microdiffusion.py

"""
The most atomic way to train and sample from a Diffusion Model in pure, dependency-free Python.
This file is the complete algorithm. Everything else is just efficiency.
Inspired by @karpathy's microgpt.py

The key ideas of diffusion:
1. Forward process: gradually destroy data by adding Gaussian noise over T timesteps
2. Reverse process: train a neural network to predict and remove that noise
3. Sampling: start from pure noise, iteratively denoise to generate new images
4. Conditioning: tell the network WHAT to generate using label embeddings (like text in Stable Diffusion)
"""

import math     # math.log, math.exp, math.cos, math.pi, math.sqrt
import os       # os.path.exists
import random   # random.seed, random.gauss, random.randint, random.choice
random.seed(42) # Let there be order among chaos

# =============================================================================
# Hyperparameters & Constants (edit these to experiment)
# =============================================================================
IMG_SIZE = 8            # image width and height in pixels
T = 20                  # total diffusion timesteps (kept small for scalar autograd)
TIME_EMB_DIM = 16       # dimension of sinusoidal time embeddings
LABEL_EMB_DIM = 16      # dimension of learned label embeddings
HIDDEN_DIM = 128        # width of hidden layers in the denoising MLP
CFG_DROP_PROB = 0.1     # probability of dropping label during training (classifier-free guidance)
GUIDANCE_SCALE = 3.0    # how strongly labels steer sampling (1.0 = no guidance)
NUM_STEPS = 1500        # number of training steps
LEARNING_RATE = 0.001   # initial learning rate
BETA1, BETA2 = 0.9, 0.999  # Adam optimizer momentum parameters
EPS_ADAM = 1e-8         # Adam epsilon for numerical stability
SCHEDULE_OFFSET = 0.008 # small offset in cosine schedule to prevent alpha_bar(0) = 1 exactly

IMG_DIM = IMG_SIZE * IMG_SIZE  # 64 pixels per image
INPUT_DIM = IMG_DIM + TIME_EMB_DIM + LABEL_EMB_DIM  # 64 + 16 + 16 = 96
OUTPUT_DIM = IMG_DIM    # predict noise for each pixel (64)

# =============================================================================
# Dataset: 8x8 grayscale images loaded from data.txt
# =============================================================================
def rows_to_pixels(rows):
    """Convert '#'/'.' strings into a flat list of floats (0.0=black, 1.0=white)"""
    pixels = []
    for row in rows[:IMG_SIZE]:
        for ch in row[:IMG_SIZE]:
            pixels.append(1.0 if ch == '.' else 0.0)
        pixels.extend([1.0] * (IMG_SIZE - min(len(row), IMG_SIZE)))
    pixels.extend([1.0] * (IMG_DIM - len(pixels)))
    return pixels

def load_dataset(path):
    """Load 8x8 images from text file. '#'=black, '.'=white, '# label' headers."""
    dataset = {}
    with open(path) as f:
        lines = f.read().strip().split('\n')
    label, rows = None, []
    for line in lines + ['# __end__']:
        line = line.rstrip()
        if line.startswith('# '):
            if label is not None and rows:
                dataset[label] = rows_to_pixels(rows)
            label = line[2:].strip()
            rows = []
        elif line == '':
            continue
        else:
            rows.append(line)
    return dataset

if not os.path.exists('data.txt'):
    print("data.txt not found! Please create it with 8x8 images.")
    exit(1)

dataset = load_dataset('data.txt')
images = list(dataset.values())
labels = list(dataset.keys())
label_to_idx = {label: i for i, label in enumerate(labels)}
num_labels = len(labels)
print(f"num images: {len(images)} ({', '.join(labels)})")
print(f"image size: {IMG_SIZE}x{IMG_SIZE} = {IMG_DIM} pixels")


# =============================================================================
# ASCII art visualization
# =============================================================================
ASCII_RAMP = " .:-=+*#%@"

def print_image(pixels, label=""):
    """Print an 8x8 grayscale image as ASCII art. pixels: flat list of floats in [0,1]"""
    if label:
        print(label)
    for row in range(IMG_SIZE):
        line = ""
        for col in range(IMG_SIZE):
            val = max(0.0, min(1.0, pixels[row * IMG_SIZE + col]))
            line += ASCII_RAMP[int((1.0 - val) * (len(ASCII_RAMP) - 1))] + " "
        print(line)
    print()

print("--- training data ---")
for name, img in list(dataset.items())[:3]:
    print_image(img, label=f"{name}:")

# =============================================================================
# Autograd engine (same as microgpt, the foundation of everything)
# =============================================================================
class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads')

    def __init__(self, data, children=(), local_grads=()):
        self.data = data
        self.grad = 0
        self._children = children
        self._local_grads = local_grads

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data + other.data, (self, other), (1, 1))

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data * other.data, (self, other), (other.data, self.data))

    def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),))
    def log(self): return Value(math.log(self.data + 1e-10), (self,), (1/(self.data + 1e-10),))
    def exp(self): return Value(math.exp(max(-20, min(20, self.data))), (self,), (math.exp(max(-20, min(20, self.data))),))
    def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))
    def __neg__(self): return self * -1
    def __radd__(self, other): return self + other
    def __sub__(self, other): return self + (-other)
    def __rsub__(self, other): return other + (-self)
    def __rmul__(self, other): return self * other
    def __truediv__(self, other): return self * other**-1
    def __rtruediv__(self, other): return other * self**-1

    def backward(self):
        topo, visited = [], set()
        stack = [(self, False)]
        while stack:
            v, processed = stack.pop()
            if processed:
                topo.append(v)
                continue
            if v in visited:
                continue
            visited.add(v)
            stack.append((v, True))
            for child in v._children:
                if child not in visited:
                    stack.append((child, False))
        self.grad = 1
        for v in reversed(topo):
            for child, local_grad in zip(v._children, v._local_grads):
                child.grad += local_grad * v.grad

# =============================================================================
# Noise schedule (cosine, Nichol & Dhariwal 2021)
# =============================================================================
def alpha_bar_fn(t):
    """Cosine noise schedule: cumulative signal retention at timestep t"""
    return math.cos(((t / T) + SCHEDULE_OFFSET) / (1 + SCHEDULE_OFFSET) * math.pi / 2) ** 2

alpha_bars = [alpha_bar_fn(t) for t in range(T + 1)]
print(f"noise schedule: alpha_bar[0]={alpha_bars[0]:.4f} (clean) -> alpha_bar[{T}]={alpha_bars[T]:.4f} (noisy)")

def forward_diffusion(x0, t):
    """Add noise to clean image x0 at timestep t. Returns (noisy_image, noise)."""
    ab = alpha_bars[t]
    sqrt_ab, sqrt_1_ab = math.sqrt(ab), math.sqrt(1 - ab)
    noise = [random.gauss(0, 1) for _ in range(IMG_DIM)]
    x_t = [sqrt_ab * x0[i] + sqrt_1_ab * noise[i] for i in range(IMG_DIM)]
    return x_t, noise

# =============================================================================
# Denoising network: unified for training (autograd) and inference (float-only)
# =============================================================================
def time_embedding(t):
    """Encode timestep t into a vector using sinusoidal functions"""
    emb = []
    for i in range(TIME_EMB_DIM // 2):
        freq = 1.0 / (10.0 ** (i / (TIME_EMB_DIM // 2)))
        emb.append(math.sin(t * freq))
        emb.append(math.cos(t * freq))
    return emb

# Initialize parameters
matrix = lambda nout, nin, std=None: [[Value(random.gauss(0, std or (2.0/(nin+nout))**0.5)) for _ in range(nin)] for _ in range(nout)]

state_dict = {
    'label_emb': matrix(num_labels, LABEL_EMB_DIM, std=0.1),
    'fc1': matrix(HIDDEN_DIM, INPUT_DIM),
    'fc2': matrix(HIDDEN_DIM, HIDDEN_DIM),
    'fc3': matrix(OUTPUT_DIM, HIDDEN_DIM),
    'skip': matrix(HIDDEN_DIM, INPUT_DIM),
}
params = [p for mat in state_dict.values() for row in mat for p in row]
print(f"num params: {len(params)}")

def linear_val(x, w):
    """Matrix-vector multiply with Value objects (autograd-tracked)."""
    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

def linear_flt(x, w):
    """Matrix-vector multiply with raw floats (no autograd, for inference)."""
    return [sum(wo[j].data * x[j] for j in range(len(x))) for wo in w]

def denoise_network(x_t_flat, t, label_idx=None, use_autograd=True):
    """
    Unified denoising function: epsilon_theta(x_t, t, label).
    use_autograd=True for training (Value objects), False for inference (raw floats).
    """
    linear = linear_val if use_autograd else linear_flt
    relu = (lambda v: v.relu()) if use_autograd else (lambda v: max(0, v))
    wrap = Value if use_autograd else float
    t_emb = [wrap(te) for te in time_embedding(t)]

    if label_idx is not None:
        l_emb = ([v for v in state_dict['label_emb'][label_idx]] if use_autograd
                 else [v.data for v in state_dict['label_emb'][label_idx]])
    else:
        l_emb = [wrap(0.0) for _ in range(LABEL_EMB_DIM)]

    x = x_t_flat + t_emb + l_emb                          # concatenate: 64 + 16 + 16 = 96
    h1 = [relu(hi) for hi in linear(x, state_dict['fc1'])] # layer 1 + relu
    skip = linear(x, state_dict['skip'])                    # skip connection
    h2 = [relu(a + b) for a, b in zip(linear(h1, state_dict['fc2']), skip)]  # layer 2 + skip + relu
    return linear(h2, state_dict['fc3'])                    # output (no activation)

# =============================================================================
# Training loop
# =============================================================================
m_buf = [0.0] * len(params)
v_buf = [0.0] * len(params)
print(f"\n--- training ({NUM_STEPS} steps) ---")

for step in range(NUM_STEPS):
    img_idx = random.randint(0, len(images) - 1)
    x0, label_idx = images[img_idx], img_idx
    t = random.randint(1, T)
    x_t, noise_true = forward_diffusion(x0, t)
    x_t_val = [Value(xi) for xi in x_t]
    use_label = label_idx if random.random() > CFG_DROP_PROB else None

    noise_pred = denoise_network(x_t_val, t, use_label, use_autograd=True)

    loss = sum((pred - Value(true)) ** 2 for pred, true in zip(noise_pred, noise_true))
    loss = loss * (1.0 / IMG_DIM)
    loss.backward()

    lr_t = LEARNING_RATE * (0.1 + 0.9 * (1 - step / NUM_STEPS))
    for i, p in enumerate(params):
        m_buf[i] = BETA1 * m_buf[i] + (1 - BETA1) * p.grad
        v_buf[i] = BETA2 * v_buf[i] + (1 - BETA2) * p.grad ** 2
        m_hat = m_buf[i] / (1 - BETA1 ** (step + 1))
        v_hat = v_buf[i] / (1 - BETA2 ** (step + 1))
        p.data -= lr_t * m_hat / (v_hat ** 0.5 + EPS_ADAM)
        p.grad = 0

    if (step + 1) % 50 == 0 or step == 0:
        print(f"step {step+1:4d}/{NUM_STEPS} | loss {loss.data:.6f}")

# =============================================================================
# Sampling: reverse diffusion with classifier-free guidance
# =============================================================================
def sample(label_idx=None, guidance_scale=GUIDANCE_SCALE):
    """Generate a new image by reversing the diffusion process."""
    x = [random.gauss(0, 1) for _ in range(IMG_DIM)]

    for t in range(T, 0, -1):
        if label_idx is not None and guidance_scale != 1.0:
            noise_cond = denoise_network(x, t, label_idx, use_autograd=False)
            noise_uncond = denoise_network(x, t, None, use_autograd=False)
            noise_pred = [noise_uncond[i] + guidance_scale * (noise_cond[i] - noise_uncond[i])
                          for i in range(IMG_DIM)]
        else:
            noise_pred = denoise_network(x, t, label_idx, use_autograd=False)

        ab_t = alpha_bars[t]
        ab_t_prev = alpha_bars[t - 1]
        beta_t = max(1e-4, min(0.999, 1 - ab_t / max(ab_t_prev, 1e-8)))
        alpha_t = 1 - beta_t
        coeff1 = 1.0 / math.sqrt(alpha_t)
        coeff2 = beta_t / math.sqrt(max(1 - ab_t, 1e-8))

        x = [coeff1 * (x[i] - coeff2 * noise_pred[i]) for i in range(IMG_DIM)]

        if t > 1:
            sigma = math.sqrt(beta_t)
            z = [random.gauss(0, 1) for _ in range(IMG_DIM)]
            x = [x[i] + sigma * z[i] for i in range(IMG_DIM)]
    return x

# --- Generate samples ---
print("\n--- unconditional sampling ---")
for i in range(3):
    print_image(sample(), label=f"sample {i+1}:")

print("--- conditional sampling (label-guided) ---")
for label_name in ['3', 'A', 'heart', 'checkerboard', 'X']:
    if label_name in label_to_idx:
        print_image(sample(label_idx=label_to_idx[label_name]), label=f"generate '{label_name}':")

print("--- forward diffusion visualization (digit '3' being destroyed by noise) ---")
x0 = dataset['3']
for t in [0, 5, 10, 15, 20]:
    if t == 0:
        print_image(x0, label=f"t={t} (clean):")
    else:
        x_t, _ = forward_diffusion(x0, t)
        print_image(x_t, label=f"t={t}:")