scripts/repro_full.py

#!/usr/bin/env python3
"""Full DiffThinker reproduction: MLLM baseline + Flow Matching + ablations.
Uses Modal T4 ($0.59/hr) with 4-bit quantized models.
Paper: 2512.24165
"""
import modal, json, os, time, io, base64

app = modal.App("diffthinker-full-repro")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch>=2.1.0", "torchvision", "transformers", "accelerate",
        "bitsandbytes", "Pillow", "numpy", "matplotlib", "safetensors",
        "einops", "scipy",
    )
)

def make_maze(grid_size=8, num=50):
    import numpy as np
    from PIL import Image, ImageDraw
    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

def make_sudoku(num=20):
    import numpy as np
    from PIL import Image, ImageDraw
    data = []
    for _ in range(num):
        nums = list(range(1, 5))
        np.random.shuffle(nums)
        sol = np.zeros((4, 4), dtype=int)
        for i in range(4):
            sol[i] = np.roll(nums, i)
        pz = np.zeros((4, 4), dtype=int)
        for idx in np.random.choice(16, 8, replace=False):
            r, c = divmod(idx, 4)
            pz[r, c] = sol[r, c]
        inp = Image.new("RGB", (64, 64), (255, 255, 255))
        d = ImageDraw.Draw(inp)
        for r in range(4):
            for c in range(4):
                x, y = c*16, r*16
                d.rectangle([x, y, x+15, y+15], outline=(0,)*3)
                if pz[r, c] > 0:
                    d.text((x+4, y), str(pz[r, c]), fill=(0,)*3)
        out = Image.new("RGB", (64, 64), (255, 255, 255))
        d = ImageDraw.Draw(out)
        for r in range(4):
            for c in range(4):
                x, y = c*16, r*16
                d.rectangle([x, y, x+15, y+15], outline=(0,)*3)
                d.text((x+4, y), str(sol[r, c]), fill=(0,)*3)
        data.append({"input": inp, "output": out})
    return data


@app.function(image=image, gpu="T4", timeout=1800)
def mllm_eval(task="maze", num_samples=5):
    """Evaluate Qwen2.5-VL-7B (4-bit) as MLLM baseline on T4."""
    import torch
    import numpy as np
    from transformers import AutoProcessor, BitsAndBytesConfig, Qwen2VLForConditionalGeneration

    device = torch.device("cuda")
    gpu_name = torch.cuda.get_device_name(0)
    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
    print(f"GPU: {gpu_name} | VRAM: {gpu_mem:.1f}GB")

    model_id = "Qwen/Qwen2.5-VL-7B-Instruct"
    print(f"Loading {model_id} with 4-bit quantization...")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
    )

    model = Qwen2VLForConditionalGeneration.from_pretrained(
        model_id,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )
    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    print("Model loaded successfully")

    if task == "maze":
        test_data = make_maze(8, num_samples)
        prompt = "Solve this maze. The green cell is start, red is goal, gray are walls. Find a path from start to goal. Output coordinates as a list of (row, col)."
    elif task == "sudoku":
        test_data = make_sudoku(num_samples)
        prompt = "Solve this 4x4 Sudoku puzzle. Fill in the missing numbers. Output the complete grid row by row."
    elif task == "tsp":
        test_data = make_maze(8, num_samples)
        prompt = "Find the shortest path from the green start to the red goal avoiding walls."

    results = []
    for idx, item in enumerate(test_data):
        msg = [{"role": "user", "content": [{"type": "image", "image": item["input"]}, {"type": "text", "text": prompt}]}]
        text = processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=[item["input"]], padding=True, return_tensors="pt").to(device)
        torch.cuda.synchronize()
        t0 = time.time()
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=256, do_sample=False, temperature=1.0)
        torch.cuda.synchronize()
        lat = time.time() - t0
        resp = processor.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        print(f"  [{idx+1}/{num_samples}] latency={lat:.3f}s, resp={resp[:100]}...")
        results.append({"latency": lat, "response": resp, "sample_idx": idx})

    avg_lat = sum(r["latency"] for r in results) / len(results)
    print(f"\nAvg latency: {avg_lat:.3f}s | GPU: T4 (4-bit quantized)")
    return {"task": task, "model": f"{model_id} (4-bit)", "avg_latency_s": avg_lat,
            "num_samples": num_samples, "gpu": "T4", "results": results}


@app.function(image=image, gpu="T4", timeout=3600)
def flow_matching_full(task="maze", num_train=200, num_epochs=30):
    """Train Flow Matching with more data, eval + CFG sweep."""
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import numpy as np
    from torch.utils.data import Dataset, DataLoader

    device = torch.device("cuda")
    print(f"Device: {device} | Training samples: {num_train} | Epochs: {num_epochs}")

    class SimpleDiT(nn.Module):
        def __init__(self):
            super().__init__()
            self.img_size = 64
            self.time_proj = nn.Linear(64, 3)
            self.time_embed = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, 64))
            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

    train = make_maze(8, num_train) if task == "maze" else make_sudoku(num_train)
    test = make_maze(8, 20) if task == "maze" else make_sudoku(10)
    train_loader = DataLoader(TaskDataset(train), batch_size=4, shuffle=True)

    model = SimpleDiT().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")

    # Train
    losses = []
    for ep in range(num_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}/{num_epochs} | Loss: {avg:.6f}")

    # Eval
    model.eval()
    results = {}
    test_dataset = TaskDataset(test)

    # CFG sweep
    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) * 100
        avg_lat = sum(lats) / len(lats)
        results[f"w={w}"] = {"accuracy": acc, "avg_latency": avg_lat, "correct": correct, "total": len(test)}
        print(f"CFG w={w}: acc={acc:.1f}% ({correct}/{len(test)}), lat={avg_lat:.3f}s")

    return {
        "task": task,
        "num_train": num_train,
        "num_epochs": num_epochs,
        "final_loss": losses[-1],
        "losses": losses[:5] + ["..."] + losses[-3:],
        "cfg_results": results,
    }


@app.function(image=image, gpu="T4", timeout=3600)
def data_scaling_ablation(task="maze"):
    """Train Flow Matching with 20, 50, 100, 200 samples. Verify Claim 6."""
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import numpy as np
    from torch.utils.data import Dataset, DataLoader

    device = torch.device("cuda")

    class SimpleDiT(nn.Module):
        def __init__(self):
            super().__init__()
            self.img_size = 64
            self.time_proj = nn.Linear(64, 3)
            self.time_embed = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, 64))
            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

    test = make_maze(8, 20)
    test_dataset = TaskDataset(test)
    sizes = [20, 50, 100, 200]
    scaling_results = {}

    for sz in sizes:
        print(f"\n--- Training with {sz} samples ---")
        train = make_maze(8, sz)
        train_loader = DataLoader(TaskDataset(train), batch_size=4, shuffle=True)

        model = SimpleDiT().to(device)
        opt = torch.optim.AdamW(model.parameters(), lr=1e-4)

        for ep in range(20):
            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()

        model.eval()
        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()
                cfg_w = 4.0
                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 + cfg_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) * 100
        scaling_results[f"N={sz}"] = {"accuracy": acc, "correct": correct, "total": len(test), "avg_latency": sum(lats)/len(lats)}
        print(f"  N={sz}: acc={acc:.1f}% ({correct}/{len(test)})")

    return scaling_results


@app.local_entrypoint()
def main():
    print("=== DiffThinker: Full Reproduction ===\n")

    # Step 1: MLLM baseline on T4 (4-bit)
    print("[1/4] MLLM baseline evaluation (Qwen2.5-VL-7B, 4-bit on T4)...")
    mllm_r = mllm_eval.remote(task="maze", num_samples=5)
    print(f"MLLM: latency={mllm_r['avg_latency_s']:.3f}s")

    # Step 2: Flow Matching with 200 samples
    print("\n[2/4] Flow Matching training (200 samples, 30 epochs)...")
    fm_r = flow_matching_full.remote(task="maze", num_train=200, num_epochs=30)
    print(f"FlowMatching: final_loss={fm_r['final_loss']:.4f}, cfg_results={json.dumps(fm_r['cfg_results'], indent=2)}")

    # Step 3: Data scaling ablation
    print("\n[3/4] Data scaling ablation...")
    scaling_r = data_scaling_ablation.remote(task="maze")
    print(f"Scaling: {json.dumps(scaling_r, indent=2)}")

    # Step 4: Summary
    print("\n[4/4] Summary")
    print(json.dumps({"mllm": mllm_r, "flow_matching": fm_r, "scaling": scaling_r}, indent=2))
    print("\n=== Done ===")