scripts/data_scaling_mid.py

#!/usr/bin/env python3
"""DiffThinker C6: Data scaling — MidDiT (~2M params) maze path-prediction.
Paper claim: accuracy scales with training data size (20B MMDiT).
We test N=20,50,100,200,400,800 on T4 ($0.59/hr).
A larger DiT should produce non-zero accuracy and show scaling.
"""

import modal, json, os, sys, time, torch, torch.nn as nn, torch.nn.functional as F
import numpy as np
from torch.utils.data import Dataset, DataLoader

app = modal.App("diffthinker-data-scaling")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch>=2.1.0", "torchvision", "Pillow", "numpy")
)


def make_maze_dataset(num_samples, grid_size=8):
    from PIL import Image, ImageDraw
    data = []
    for _ in range(num_samples):
        grid = np.zeros((grid_size, grid_size), dtype=np.uint8)
        start = (0, np.random.randint(0, grid_size))
        goal = (grid_size - 1, np.random.randint(0, grid_size))
        for _ in range(int(grid_size * grid_size * 0.15)):
            wx, wy = np.random.randint(1, grid_size-1, 2)
            grid[wx, wy] = 1
        input_img = Image.new("RGB", (64, 64), (255, 255, 255))
        draw = ImageDraw.Draw(input_img)
        cw = 64 // grid_size
        for r in range(grid_size):
            for c in range(grid_size):
                if grid[r, c] == 1:
                    draw.rectangle([c*cw, r*cw, (c+1)*cw, (r+1)*cw], fill=(100, 100, 100))
        draw.rectangle([start[1]*cw, start[0]*cw, (start[1]+1)*cw, (start[0]+1)*cw], fill=(0, 255, 0))
        draw.rectangle([goal[1]*cw, goal[0]*cw, (goal[1]+1)*cw, (goal[0]+1)*cw], fill=(255, 0, 0))
        output_img = Image.new("RGB", (64, 64), (255, 255, 255))
        draw = ImageDraw.Draw(output_img)
        for r in range(grid_size):
            for c in range(grid_size):
                if grid[r, c] == 1:
                    draw.rectangle([c*cw, r*cw, (c+1)*cw, (r+1)*cw], fill=(100, 100, 100))
        draw.rectangle([start[1]*cw, start[0]*cw, (start[1]+1)*cw, (start[0]+1)*cw], fill=(0, 255, 0))
        draw.rectangle([goal[1]*cw, goal[0]*cw, (goal[1]+1)*cw, (goal[0]+1)*cw], fill=(255, 0, 0))
        path_y = np.linspace(start[0], goal[0], grid_size).astype(int)
        path_x = np.linspace(start[1], goal[1], grid_size).astype(int)
        for px, py in zip(path_x, path_y):
            if 0 <= px < grid_size and 0 <= py < grid_size and grid[py, px] != 1:
                draw.rectangle([px*cw, py*cw, (px+1)*cw, (py+1)*cw], fill=(0, 0, 255))
        data.append({"input": input_img, "output": output_img})
    return data


class MidDiT(nn.Module):
    def __init__(self, ic=3, ims=64):
        super().__init__()
        self.ims = ims; d = 128
        self.tp = nn.Linear(d, ic)
        self.te = nn.Sequential(nn.Linear(1, d), nn.SiLU(), nn.Linear(d, d))
        self.ce = nn.Sequential(nn.Conv2d(ic, 32, 3, padding=1), nn.SiLU(),
                                nn.Conv2d(32, 64, 3, padding=1), nn.SiLU(),
                                nn.Conv2d(64, d, 3, padding=1))
        self.d1 = nn.Conv2d(ic+ic, 64, 3, padding=1)
        self.d2 = nn.Conv2d(64, d, 3, stride=2, padding=1)
        self.d3 = nn.Conv2d(d, d, 3, stride=2, padding=1)
        self.mid = nn.Sequential(nn.Conv2d(d, d, 3, padding=1), nn.SiLU(),
                                 nn.Conv2d(d, d, 3, padding=1), nn.SiLU(),
                                 nn.Conv2d(d, d, 3, padding=1))
        self.u3 = nn.ConvTranspose2d(d, d, 4, stride=2, padding=1)
        self.u2 = nn.ConvTranspose2d(d, 64, 4, stride=2, padding=1)
        self.u1 = nn.Conv2d(64, ic, 3, padding=1)

    def forward(self, x, t, c):
        B = x.shape[0]
        te = self.te(t.view(-1, 1).float() / 100.0)
        te = self.tp(te).view(B, -1, 1, 1).expand(-1, -1, self.ims, self.ims)
        h = torch.cat([x, te], dim=1)
        c = self.ce(c)
        c = F.interpolate(c, size=(self.ims, self.ims), mode='bilinear', align_corners=False)
        h = self.d1(h)
        h = self.d2(h)
        h = self.d3(h)
        h = self.mid(h)
        h = self.u3(h)
        h = self.u2(h)
        h = self.u1(h)
        return h


class MazeDS(Dataset):
    def __init__(self, data):
        self.data = data
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        item = self.data[idx]
        inp = torch.tensor(np.array(item["input"]).transpose(2,0,1), dtype=torch.float32) / 255.0
        out = torch.tensor(np.array(item["output"]).transpose(2,0,1), dtype=torch.float32) / 255.0
        return inp, out


@app.function(image=image, gpu="T4", timeout=7200)
def train_and_eval(num_train, num_test=50, epochs=30):
    device = torch.device("cuda")
    gpu_name = torch.cuda.get_device_name(0)
    print(f"=== N={num_train} samples | T4 | {epochs} epochs ===")

    train_data = make_maze_dataset(num_train)
    test_data = make_maze_dataset(num_test)
    train_loader = DataLoader(MazeDS(train_data), batch_size=8, shuffle=True)
    test_ds = MazeDS(test_data)

    model = MidDiT().to(device)
    total_p = sum(p.numel() for p in model.parameters())
    optim = torch.optim.AdamW(model.parameters(), lr=1e-4)

    t_start = time.time()
    for ep in range(epochs):
        model.train()
        loss_sum = 0.0
        for cond, target in train_loader:
            cond, target = cond.to(device), target.to(device)
            B = cond.shape[0]
            t = torch.randint(0, 100, (B,), device=device)
            noise = torch.randn_like(target)
            alpha = t.view(-1,1,1,1).float() / 100.0
            x_t = (1 - alpha) * target + alpha * noise
            v = model(x_t, t, cond)
            loss = F.mse_loss(v, noise - target)
            optim.zero_grad(); loss.backward(); optim.step()
            loss_sum += loss.item()
        if (ep+1) % 10 == 0 or ep == epochs-1:
            print(f"  Ep {ep+1}/{epochs} loss={loss_sum/len(train_loader):.4f}")
    elapsed = time.time() - t_start

    model.eval()
    correct = 0
    with torch.no_grad():
        for inp, target in test_ds:
            inp = inp.unsqueeze(0).to(device)
            target_np = target.numpy().transpose(1,2,0) * 255
            x = torch.randn(1,3,64,64,device=device)
            for s in range(20):
                tv = torch.full((1,), s*5, device=device, dtype=torch.long)
                vc = model(x, tv, inp)
                vu = model(x, tv, torch.zeros_like(inp))
                x = x + (1.0/20) * (vu + 4.0*(vc-vu))
            out_np = (x.squeeze(0).cpu().numpy().transpose(1,2,0) * 255).clip(0,255).astype(np.uint8)
            mse = np.mean((out_np.astype(float) - target_np.astype(float))**2)
            if mse < 500:
                correct += 1

    acc = correct / num_test * 100
    final_loss = loss_sum / len(train_loader)
    print(f"  Acc: {acc:.1f}% ({correct}/{num_test}) | Loss: {final_loss:.4f} | Time: {elapsed:.1f}s")

    return {"num_train": num_train, "params": total_p, "final_loss": round(final_loss, 4),
            "accuracy": round(acc, 1), "epochs": epochs, "time_sec": round(elapsed, 1),"gpu": gpu_name}


@app.function(image=image, gpu="T4", timeout=7200)
def data_scaling_sweep():
    ns = [20, 50, 100, 200, 400, 800]
    results = []
    for n in ns:
        r = train_and_eval.remote(num_train=n)
        results.append(r)
        print(json.dumps(r))

    print("\n=== DATA SCALING RESULTS ===")
    print(f"{'N':>6} | {'Loss':>8} | {'Acc%':>6} | {'Time':>8}")
    print("-"*35)
    for r in results:
        print(f"{r['num_train']:>6} | {r['final_loss']:>8.4f} | {r['accuracy']:>5.1f}% | {r['time_sec']:>7.1f}s")

    return results


@app.local_entrypoint()
def main():
    print("="*60)
    print("C6: Data Scaling — MidDiT (~2M params) on T4")
    print("Paper: 20B MMDiT accuracy scales with data (20→800+ samples)")
    print("="*60)
    r = data_scaling_sweep.remote()
    print(json.dumps(r, indent=2))