scripts/repro_fix.py

#!/usr/bin/env python3
"""Fixed reproduction: text-based MLLM + Flow Matching.
Uses Modal T4. Loads Qwen2.5-7B-Instruct (text) in 8-bit for MLLM proxy.
Paper: 2512.24165
"""
import modal, json, os, time

app = modal.App("diffthinker-v2")

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

def make_maze_text(grid_size=8, num=50):
    """Generate maze tasks as text descriptions + solution paths."""
    import numpy as np
    data = []
    for _ in range(num):
        g = np.zeros((grid_size, grid_size), dtype=int)
        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)
            g[wx, wy] = 1
        desc = f"Grid ({grid_size}x{grid_size}):\n"
        for r in range(grid_size):
            row = ""
            for c in range(grid_size):
                if (r,c) == start: row += "S "
                elif (r,c) == goal: row += "G "
                elif g[r,c] == 1: row += "# "
                else: row += ". "
            desc += row + "\n"
        desc += f"Start at {start}, goal at {goal}. Find path avoiding #."
        # Simple BFS solution
        from collections import deque
        q = deque([start])
        parent = {start: None}
        solved = False
        while q:
            r, c = q.popleft()
            if (r,c) == goal:
                solved = True
                break
            for dr, dc in [(0,1),(0,-1),(1,0),(-1,0)]:
                nr, nc = r+dr, c+dc
                if 0 <= nr < grid_size and 0 <= nc < grid_size and g[nr,nc] != 1 and (nr,nc) not in parent:
                    parent[(nr,nc)] = (r,c)
                    q.append((nr,nc))
        sol = []
        if solved:
            cur = goal
            while cur:
                sol.append(cur)
                cur = parent[cur]
            sol.reverse()
        data.append({"input": desc, "output": str(sol) if sol else "unsolvable", "solvable": solved})
    return data


@app.function(image=image, gpu="T4", timeout=1800)
def mllm_text_eval(task="maze", num_samples=5):
    """Evaluate Qwen2.5-7B-Instruct (text-only, 8-bit) on maze via text description.
    This tests the paper's claim that text-centric MLLMs struggle with spatial reasoning."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

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

    bnb = BitsAndBytesConfig(load_in_8bit=True, bnb_8bit_compute_dtype=torch.bfloat16)

    model = AutoModelForCausalLM.from_pretrained(
        model_id, quantization_config=bnb, device_map="auto",
        trust_remote_code=True, torch_dtype=torch.bfloat16,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    print("Model loaded")

    test_data = make_maze_text(8, num_samples)

    results = []
    correct = 0
    for idx, item in enumerate(test_data):
        messages = [{"role": "user", "content": f"{item['input']}\nOutput the path as a list of (row,col) coordinates."}]
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(text, 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=0.0)
        torch.cuda.synchronize()
        lat = time.time() - t0
        resp = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        # Check if answer contains the solution
        sol_str = str(item["output"])
        has_solution = any(coord in resp for coord in [str(p) for p in eval(sol_str)]) if item["solvable"] else False
        if item["solvable"] and has_solution:
            correct += 1
        results.append({"latency": lat, "response": resp[:200], "correct": has_solution, "solvable": item["solvable"]})
        print(f"  [{idx+1}] lat={lat:.3f}s correct={has_solution}")

    acc = correct / sum(1 for r in results if r["solvable"]) * 100 if sum(1 for r in results if r["solvable"]) else 0
    avg_lat = sum(r["latency"] for r in results) / len(results)
    print(f"\nMLLM text: acc={acc:.1f}% ({correct}/{sum(1 for r in results if r['solvable'])}), lat={avg_lat:.3f}s")
    return {"model": f"{model_id} (8-bit)", "avg_latency_s": avg_lat, "accuracy_pct": acc,
            "correct": correct, "total_solvable": sum(1 for r in results if r["solvable"]), "results": results}


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

    device = torch.device("cuda")

    def make_maze_img(grid_size=8, num=num_train):
        data = []
        for _ in range(num):
            g = np.zeros((grid_size, grid_size), dtype=np.uint8)
            s = (0, np.random.randint(0, grid_size))
            gl = (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([gl[1]*cw, gl[0]*cw, (gl[1]+1)*cw, (gl[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([gl[1]*cw, gl[0]*cw, (gl[1]+1)*cw, (gl[0]+1)*cw], fill=(255, 0, 0))
            py = np.linspace(s[0], gl[0], grid_size).astype(int)
            px = np.linspace(s[1], gl[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_proj = nn.Linear(64, 3)
            self.time_embed = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, 64))
            self.ce = 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.d1 = nn.Conv2d(6, 32, 3, padding=1)
            self.d2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
            self.d3 = nn.Conv2d(64, 64, 3, stride=2, padding=1)
            self.m = nn.Sequential(nn.Conv2d(64, 64, 3, padding=1), nn.SiLU(), nn.Conv2d(64, 64, 3, padding=1))
            self.u3 = nn.ConvTranspose2d(64, 64, 4, stride=2, padding=1)
            self.u2 = nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1)
            self.u1 = nn.Conv2d(32, 3, 3, padding=1)
        def forward(self, x_t, t, cond):
            B = x_t.shape[0]
            te = self.time_embed(t.view(-1, 1).float() / 100.0)
            te = self.time_proj(te).view(B, -1, 1, 1).expand(-1, -1, 64, 64)
            h = torch.cat([x_t, te], dim=1)
            c = self.ce(cond)
            c = F.interpolate(c, size=(64, 64), mode='bilinear', align_corners=False)
            h = self.d1(h)
            h = self.d2(h)
            h = self.d3(h)
            h = self.m(h)
            h = self.u3(h)
            h = self.u2(h)
            h = self.u1(h)
            return h

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

    train = make_maze_img()
    test = make_maze_img(num=20)
    tl = DataLoader(D(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()):,}")

    for ep in range(num_epochs):
        model.train()
        el = 0
        for c, t in tl:
            c, t = c.to(device), t.to(device)
            t_vec = torch.randint(0, 100, (c.shape[0],), device=device)
            n = torch.randn_like(t)
            a = t_vec.view(-1,1,1,1).float() / 100.0
            xt = (1-a) * t + a * n
            v = model(xt, t_vec, c)
            loss = F.mse_loss(v, n - t)
            opt.zero_grad()
            loss.backward()
            opt.step()
            el += loss.item()
        if (ep+1) % 5 == 0:
            print(f"Epoch {ep+1}/{num_epochs} | Loss: {el/len(tl):.6f}")

    model.eval()
    td = D(test)
    cfg_w = 4.0
    correct = 0
    lats = []
    for item in td:
        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)
            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 += (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)
        if np.mean((out.astype(float) - tn.astype(float))**2) < 500:
            correct += 1

    acc = correct / len(test) * 100
    avg_lat = sum(lats) / len(lats)
    print(f"\nFlowMatching: acc={acc:.1f}% ({correct}/{len(test)}), lat={avg_lat:.3f}s")
    return {"num_train": num_train, "accuracy": acc, "avg_latency": avg_lat, "correct": correct, "total": len(test)}


@app.function(image=image, gpu="T4", timeout=3600)
def data_scaling_v2():
    """Claim 6: Train with 20,50,100,200,400 samples, measure accuracy."""
    import torch, torch.nn as nn, torch.nn.functional as F
    import numpy as np
    from torch.utils.data import Dataset, DataLoader
    from PIL import Image, ImageDraw

    device = torch.device("cuda")

    class SimpleDiT(nn.Module):
        def __init__(self):
            super().__init__()
            self.time_proj = nn.Linear(64, 3)
            self.time_embed = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, 64))
            self.ce = 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.d1 = nn.Conv2d(6, 32, 3, padding=1)
            self.d2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
            self.d3 = nn.Conv2d(64, 64, 3, stride=2, padding=1)
            self.m = nn.Sequential(nn.Conv2d(64, 64, 3, padding=1), nn.SiLU(), nn.Conv2d(64, 64, 3, padding=1))
            self.u3 = nn.ConvTranspose2d(64, 64, 4, stride=2, padding=1)
            self.u2 = nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1)
            self.u1 = nn.Conv2d(32, 3, 3, padding=1)
        def forward(self, x_t, t, cond):
            B = x_t.shape[0]
            te = self.time_embed(t.view(-1, 1).float() / 100.0)
            te = self.time_proj(te).view(B, -1, 1, 1).expand(-1, -1, 64, 64)
            h = torch.cat([x_t, te], dim=1)
            c = self.ce(cond)
            c = F.interpolate(c, size=(64, 64), mode='bilinear', align_corners=False)
            h = self.d1(h)
            h = self.d2(h)
            h = self.d3(h)
            h = self.m(h)
            h = self.u3(h)
            h = self.u2(h)
            h = self.u1(h)
            return h

    def gen(n):
        data = []
        for _ in range(n):
            g = np.zeros((8,8), dtype=np.uint8)
            s = (0, np.random.randint(0, 8))
            gl = (7, np.random.randint(0, 8))
            for _ in range(9):
                wx, wy = np.random.randint(1, 7, 2)
                g[wx, wy] = 1
            inp = Image.new("RGB", (64, 64), (255, 255, 255))
            d = ImageDraw.Draw(inp)
            cw = 8
            for r in range(8):
                for c in range(8):
                    if g[r, c]: 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([gl[1]*cw, gl[0]*cw, (gl[1]+1)*cw, (gl[0]+1)*cw], fill=(255, 0, 0))
            out = Image.new("RGB", (64, 64), (255, 255, 255))
            d = ImageDraw.Draw(out)
            for r in range(8):
                for c in range(8):
                    if g[r, c]: 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([gl[1]*cw, gl[0]*cw, (gl[1]+1)*cw, (gl[0]+1)*cw], fill=(255, 0, 0))
            py = np.linspace(s[0], gl[0], 8).astype(int)
            px = np.linspace(s[1], gl[1], 8).astype(int)
            for x, y in zip(px, py):
                if 0 <= x < 8 and 0 <= y < 8 and g[y, x] != 1:
                    d.rectangle([x*cw, y*cw, (x+1)*cw, (y+1)*cw], fill=(0, 0, 255))
            data.append({"i": inp, "o": out})
        return data

    test = gen(20)
    test_t = [(torch.tensor(np.array(it["i"]).transpose(2,0,1), dtype=torch.float32)/255.0,
               torch.tensor(np.array(it["o"]).transpose(2,0,1), dtype=torch.float32)/255.0) for it in test]

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

    sizes = [20, 50, 100, 200, 400]
    results = {}
    for sz in sizes:
        print(f"\n--- N={sz} ---")
        train = gen(sz)
        tl = DataLoader(D(train), batch_size=4, shuffle=True)
        model = SimpleDiT().to(device)
        opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
        for ep in range(25):
            model.train()
            el = 0
            for c, t in tl:
                c, t = c.to(device), t.to(device)
                tv = torch.randint(0, 100, (c.shape[0],), device=device)
                n = torch.randn_like(t)
                a = tv.view(-1,1,1,1).float() / 100.0
                v = model((1-a)*t + a*n, tv, c)
                loss = F.mse_loss(v, n - t)
                opt.zero_grad()
                loss.backward()
                opt.step()
                el += loss.item()
        model.eval()
        correct = 0
        for inp, tgt in test_t:
            inp, tgt = inp.unsqueeze(0).to(device), tgt
            tn = tgt.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 = model(x, tv, inp)
                    vu = model(x, tv, torch.zeros_like(inp))
                    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)
            if np.mean((out.astype(float) - tn.astype(float))**2) < 500:
                correct += 1
        acc = correct / 20 * 100
        results[f"N={sz}"] = {"accuracy": acc, "correct": correct, "total": 20}
        print(f"  N={sz}: acc={acc:.1f}%")

    return results

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

    # Step 1: MLLM text proxy (tests text-centric reasoning limitations)
    print("[1/3] MLLM text-only baseline (Qwen2.5-7B, 8-bit)...")
    mllm_r = mllm_text_eval.remote(task="maze", num_samples=5)
    print(json.dumps(mllm_r, indent=2))

    # Step 2: Flow Matching with 200 samples
    print("\n[2/3] Flow Matching (200 samples, 40 epochs)...")
    fm_r = flow_matching_v2.remote(task="maze", num_train=200, num_epochs=40)
    print(json.dumps(fm_r, indent=2))

    # Step 3: Data scaling
    print("\n[3/3] Data scaling ablation...")
    ds_r = data_scaling_v2.remote()
    print(json.dumps(ds_r, indent=2))

    print("\n=== Done ===")