scripts/hf_jobs_repro.py

#!/usr/bin/env python3
"""DiffThinker reproduction on HF Jobs (T4). 
Reproduces Claims 4-6 with toy-scale Flow Matching + CFG ablation + data scaling.
"""
import json, os, sys, time, io, base64
import numpy as np
from PIL import Image, ImageDraw
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0
print(f"GPU: {gpu_name} | VRAM: {gpu_mem:.1f}GB")
print(f"PyTorch: {torch.__version__}")

def make_maze(grid_size=8, num=50):
    data = []
    for _ in range(num):
        g = np.zeros((grid_size, grid_size), dtype=np.uint8)
        s = (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)
            g[wx, wy] = 1
        inp = Image.new("RGB", (64, 64), (255, 255, 255))
        d = ImageDraw.Draw(inp)
        cw = 64 // grid_size
        for r in range(grid_size):
            for c in range(grid_size):
                if g[r, c] == 1:
                    d.rectangle([c*cw, r*cw, (c+1)*cw, (r+1)*cw], fill=(100,)*3)
        d.rectangle([s[1]*cw, s[0]*cw, (s[1]+1)*cw, (s[0]+1)*cw], fill=(0, 255, 0))
        d.rectangle([goal[1]*cw, goal[0]*cw, (goal[1]+1)*cw, (goal[0]+1)*cw], fill=(255, 0, 0))
        out = Image.new("RGB", (64, 64), (255, 255, 255))
        d = ImageDraw.Draw(out)
        for r in range(grid_size):
            for c in range(grid_size):
                if g[r, c] == 1:
                    d.rectangle([c*cw, r*cw, (c+1)*cw, (r+1)*cw], fill=(100,)*3)
        d.rectangle([s[1]*cw, s[0]*cw, (s[1]+1)*cw, (s[0]+1)*cw], fill=(0, 255, 0))
        d.rectangle([goal[1]*cw, goal[0]*cw, (goal[1]+1)*cw, (goal[0]+1)*cw], fill=(255, 0, 0))
        py = np.linspace(s[0], goal[0], grid_size).astype(int)
        px = np.linspace(s[1], goal[1], grid_size).astype(int)
        for x, y in zip(px, py):
            if 0 <= x < grid_size and 0 <= y < grid_size and g[y, x] != 1:
                d.rectangle([x*cw, y*cw, (x+1)*cw, (y+1)*cw], fill=(0, 0, 255))
        data.append({"input": inp, "output": out})
    return data

class SimpleDiT(nn.Module):
    def __init__(self):
        super().__init__()
        self.time_embed = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, 64))
        self.time_proj = nn.Linear(64, 3)
        self.cond_encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.SiLU(), nn.Conv2d(16, 32, 3, padding=1), nn.SiLU(), nn.Conv2d(32, 64, 3, padding=1))
        self.down1 = nn.Conv2d(6, 32, 3, padding=1)
        self.down2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
        self.down3 = nn.Conv2d(64, 64, 3, stride=2, padding=1)
        self.mid = nn.Sequential(nn.Conv2d(64, 64, 3, padding=1), nn.SiLU(), nn.Conv2d(64, 64, 3, padding=1))
        self.up3 = nn.ConvTranspose2d(64, 64, 4, stride=2, padding=1)
        self.up2 = nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1)
        self.up1 = nn.Conv2d(32, 3, 3, padding=1)
    def forward(self, x_t, t, cond):
        B = x_t.shape[0]
        t_emb = self.time_embed(t.view(-1, 1).float() / 100.0)
        t_emb = self.time_proj(t_emb).view(B, -1, 1, 1).expand(-1, -1, 64, 64)
        h = torch.cat([x_t, t_emb], dim=1)
        c = self.cond_encoder(cond)
        c = F.interpolate(c, size=(64, 64), mode='bilinear', align_corners=False)
        h = self.down1(h)
        h = self.down2(h)
        h = self.down3(h)
        h = self.mid(h)
        h = self.up3(h)
        h = self.up2(h)
        h = self.up1(h)
        return h

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

print("\n=== DiffThinker Reproduction on HF Jobs ===\n")

# 1. Flow Matching training + CFG sweep
print("Claim 4: Flow Matching reformulation (training + inference)")
train_samples = 200
test_samples = 20
train_data = make_maze(8, train_samples)
test_data = make_maze(8, test_samples)
train_loader = DataLoader(TaskDataset(train_data), batch_size=4, shuffle=True)
test_dataset = TaskDataset(test_data)

model = SimpleDiT().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
params = sum(p.numel() for p in model.parameters())
print(f"Model params: {params:,}")
print(f"Training samples: {train_samples}, Test samples: {test_samples}")

epochs = 40
losses = []
for ep in range(epochs):
    model.train()
    el = 0.0
    for cond, target in train_loader:
        cond, target = cond.to(device), target.to(device)
        t = torch.randint(0, 100, (cond.shape[0],), device=device)
        noise = torch.randn_like(target)
        a = t.view(-1,1,1,1).float() / 100.0
        xt = (1 - a) * target + a * noise
        v_pred = model(xt, t, cond)
        loss = F.mse_loss(v_pred, noise - target)
        opt.zero_grad()
        loss.backward()
        opt.step()
        el += loss.item()
    avg = el / len(train_loader)
    losses.append(avg)
    if (ep+1) % 5 == 0:
        print(f"  Epoch {ep+1}/{epochs} | Loss: {avg:.6f}")

print(f"Final loss: {losses[-1]:.6f}")
print(f"Loss trajectory (first): {losses[0]:.6f} -> {losses[-1]:.6f}")

# 2. CFG sweep (Claim 6)
print("\nClaim 6: CFG ablation")
model.eval()
cfg_results = {}
for w in [1.0, 2.0, 4.0, 7.0]:
    correct = 0
    lats = []
    for item in test_dataset:
        inp = item[0].unsqueeze(0).to(device)
        t_np = item[1].numpy().transpose(1,2,0) * 255
        with torch.no_grad():
            x = torch.randn(1, 3, 64, 64, device=device)
            t0 = time.time()
            for step in range(20):
                tv = torch.full((1,), step * 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 + w * (vc - vu))
            torch.cuda.synchronize()
            lats.append(time.time() - t0)
        out = (x.squeeze(0).cpu().numpy().transpose(1,2,0) * 255).clip(0, 255).astype(np.uint8)
        mse = np.mean((out.astype(float) - t_np.astype(float))**2)
        if mse < 500:
            correct += 1
    acc = correct / len(test_dataset) * 100
    avg_lat = sum(lats) / len(lats)
    cfg_results[f"w={w}"] = {"accuracy": acc, "avg_latency_s": avg_lat, "correct": correct, "total": len(test_dataset)}
    print(f"  CFG w={w}: acc={acc:.1f}% ({correct}/{len(test_dataset)}), lat={avg_lat:.3f}s")

# 3. Data scaling (Claim 6)
print("\nClaim 6: Data scaling ablation")
test_data2 = make_maze(8, 10)
test_dataset2 = TaskDataset(test_data2)
for sz in [20, 50, 100]:
    train = make_maze(8, sz)
    train_loader2 = DataLoader(TaskDataset(train), batch_size=4, shuffle=True)
    m = SimpleDiT().to(device)
    o = torch.optim.AdamW(m.parameters(), lr=1e-4)
    for ep in range(15):
        m.train()
        el = 0.0
        for cond, target in train_loader2:
            cond, target = cond.to(device), target.to(device)
            t = torch.randint(0, 100, (cond.shape[0],), device=device)
            noise = torch.randn_like(target)
            a = t.view(-1,1,1,1).float() / 100.0
            xt = (1 - a) * target + a * noise
            v_pred = m(xt, t, cond)
            loss = F.mse_loss(v_pred, noise - target)
            o.zero_grad()
            loss.backward()
            o.step()
            el += loss.item()
    m.eval()
    c = 0
    for item in test_dataset2:
        inp = item[0].unsqueeze(0).to(device)
        tn = item[1].numpy().transpose(1,2,0) * 255
        with torch.no_grad():
            x = torch.randn(1, 3, 64, 64, device=device)
            for step in range(20):
                tv = torch.full((1,), step * 5, device=device, dtype=torch.long)
                vc = m(x, tv, inp)
                vu = m(x, tv, torch.zeros_like(inp))
                x = x + (1.0/20) * (vu + 4.0 * (vc - vu))
        out = (x.squeeze(0).cpu().numpy().transpose(1,2,0) * 255).clip(0, 255).astype(np.uint8)
        mse = np.mean((out.astype(float) - tn.astype(float))**2)
        if mse < 500:
            c += 1
    print(f"  N={sz}: acc={c/len(test_dataset2)*100:.1f}% ({c}/{len(test_dataset2)})")

print("\n=== Summary ===")
print(json.dumps({
    "model_params": params,
    "final_loss": losses[-1],
    "loss_trajectory": [round(l, 4) for l in [losses[0], losses[-1]]],
    "cfg_results": {k: {sk: sv for sk, sv in v.items()} for k, v in cfg_results.items()},
    "gpu": gpu_name,
    "vram_gb": round(gpu_mem, 1),
}, indent=2))
print("\n=== Done ===")