scripts/repro_flow_matching.py

#!/usr/bin/env python3
"""DiffThinker Reproduction: Simplified Flow Matching on Maze & Sudoku.
Uses lightweight DiT backbone on Modal T4 ($0.59/hr).
Paper 2512.24165 — Claim verification at toy scale with honest reporting.
"""

import modal
import json, os, sys, time

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

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

def make_maze_data(grid_size=8, num_samples=50):
    import numpy as np
    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))
        wall_count = int(grid_size * grid_size * 0.15)
        for _ in range(wall_count):
            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)
        cell_w = 64 // grid_size
        for r in range(grid_size):
            for c in range(grid_size):
                if grid[r, c] == 1:
                    draw.rectangle([c*cell_w, r*cell_w, (c+1)*cell_w, (r+1)*cell_w], fill=(100, 100, 100))
        draw.rectangle([start[1]*cell_w, start[0]*cell_w, (start[1]+1)*cell_w, (start[0]+1)*cell_w], fill=(0, 255, 0))
        draw.rectangle([goal[1]*cell_w, goal[0]*cell_w, (goal[1]+1)*cell_w, (goal[0]+1)*cell_w], 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*cell_w, r*cell_w, (c+1)*cell_w, (r+1)*cell_w], fill=(100, 100, 100))
        draw.rectangle([start[1]*cell_w, start[0]*cell_w, (start[1]+1)*cell_w, (start[0]+1)*cell_w], fill=(0, 255, 0))
        draw.rectangle([goal[1]*cell_w, goal[0]*cell_w, (goal[1]+1)*cell_w, (goal[0]+1)*cell_w], 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*cell_w, py*cell_w, (px+1)*cell_w, (py+1)*cell_w], fill=(0, 0, 255))
        data.append({"input": input_img, "output": output_img, "grid_size": grid_size})
    return data

def make_sudoku_data(num_samples=20):
    import numpy as np
    from PIL import Image, ImageDraw
    data = []
    for _ in range(num_samples):
        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)
        puzzle = np.zeros((4, 4), dtype=int)
        given = np.random.choice(16, 8, replace=False)
        for idx in given:
            r, c = divmod(idx, 4)
            puzzle[r, c] = sol[r, c]
        input_img = Image.new("RGB", (16, 16), (255, 255, 255))
        draw = ImageDraw.Draw(input_img)
        for r in range(4):
            for c in range(4):
                x, y = c*4, r*4
                draw.rectangle([x, y, x+3, y+3], outline=(0, 0, 0))
                if puzzle[r, c] > 0:
                    draw.text((x+1, y), str(puzzle[r, c]), fill=(0, 0, 0))
        output_img = Image.new("RGB", (16, 16), (255, 255, 255))
        draw = ImageDraw.Draw(output_img)
        for r in range(4):
            for c in range(4):
                x, y = c*4, r*4
                draw.rectangle([x, y, x+3, y+3], outline=(0, 0, 0))
                draw.text((x+1, y), str(sol[r, c]), fill=(0, 0, 0))
        data.append({"input": input_img, "output": output_img})
    return data


@app.function(image=image, gpu="T4", timeout=3600)
def full_repro(task="maze", grid_size=8, num_epochs=15):
    """Train + Evaluate Flow Matching inside a single Modal function."""
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import numpy as np
    from torch.utils.data import Dataset, DataLoader
    from PIL import Image

    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 "none"
    print(f"Device: {device} | GPU: {gpu_name}")
    print(f"Task: {task} | Epochs: {num_epochs} | Grid: {grid_size}")

    class SimpleDiT(nn.Module):
        def __init__(self, img_channels=3, img_size=64, latent_dim=64):
            super().__init__()
            self.img_size = img_size
            self.time_proj = nn.Linear(latent_dim, img_channels)
            self.time_embed = nn.Sequential(nn.Linear(1, latent_dim), nn.SiLU(), nn.Linear(latent_dim, latent_dim))
            self.cond_encoder = nn.Sequential(
                nn.Conv2d(img_channels, 16, 3, padding=1), nn.SiLU(),
                nn.Conv2d(16, 32, 3, padding=1), nn.SiLU(),
                nn.Conv2d(32, latent_dim, 3, padding=1),
            )
            self.down1 = nn.Conv2d(img_channels + img_channels, 32, 3, padding=1)
            self.down2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
            self.down3 = nn.Conv2d(64, latent_dim, 3, stride=2, padding=1)
            self.mid = nn.Sequential(nn.Conv2d(latent_dim, latent_dim, 3, padding=1), nn.SiLU(), nn.Conv2d(latent_dim, latent_dim, 3, padding=1))
            self.up3 = nn.ConvTranspose2d(latent_dim, 64, 4, stride=2, padding=1)
            self.up2 = nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1)
            self.up1 = nn.Conv2d(32, img_channels, 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, self.img_size, self.img_size)
            h = torch.cat([x_t, t_emb], dim=1)
            c = self.cond_encoder(cond)
            c = F.interpolate(c, size=(self.img_size, self.img_size), 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, 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

    print("Generating data...")
    if task == "maze":
        train_data = make_maze_data(grid_size, 50)
        test_data = make_maze_data(grid_size, 10)
    elif task == "sudoku":
        train_data = make_sudoku_data(20)
        test_data = make_sudoku_data(5)
    else:
        raise ValueError(f"Unknown task: {task}")

    train_dataset = TaskDataset(train_data)
    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
    test_dataset = TaskDataset(test_data)

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

    # Training
    for epoch in range(num_epochs):
        model.train()
        epoch_loss = 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_target = noise - target
            v_pred = model(x_t, t, cond)
            loss = F.mse_loss(v_pred, v_target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        print(f"Epoch {epoch+1}/{num_epochs} | Loss: {epoch_loss/len(train_loader):.6f}")

    # Evaluation
    model.eval()
    latencies = []
    correct = 0
    print(f"\nEvaluating {len(test_data)} samples...")
    for idx, item in enumerate(test_data):
        inp = torch.tensor(np.array(item["input"]).transpose(2, 0, 1), dtype=torch.float32).unsqueeze(0).to(device) / 255.0
        target = torch.tensor(np.array(item["output"]).transpose(2, 0, 1), dtype=torch.float32).unsqueeze(0).to(device) / 255.0
        with torch.no_grad():
            x = torch.randn(1, 3, 64, 64, device=device)
            t_start = time.time()
            num_steps = 20
            cfg_w = 4.0
            for step in range(num_steps):
                t_tensor = torch.full((1,), step * (100 // num_steps), device=device, dtype=torch.long)
                v_cond = model(x, t_tensor, inp)
                v_uncond = model(x, t_tensor, torch.zeros_like(inp))
                v = v_uncond + cfg_w * (v_cond - v_uncond)
                x = x + (1.0 / num_steps) * v
            torch.cuda.synchronize()
            infer_time = time.time() - t_start
        latencies.append(infer_time)
        out_np = (x.squeeze(0).cpu().numpy().transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)
        target_np = (target.squeeze(0).cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)
        mse = np.mean((out_np.astype(float) - target_np.astype(float))**2)
        is_correct = mse < 500.0
        if is_correct:
            correct += 1
        print(f"  [{idx+1}/{len(test_data)}] MSE={mse:.1f}, latency={infer_time:.3f}s, correct={is_correct}")

    avg_lat = sum(latencies) / len(latencies)
    acc = correct / len(test_data) * 100
    print(f"\nAvg latency: {avg_lat:.3f}s | Accuracy: {acc:.1f}%")
    print(f"CFG scale used: w={cfg_w} (paper default: w=4)")
    print(f"Inference steps: {num_steps} (paper default: 20)")

    return {
        "task": task,
        "model_params": total_params,
        "num_train": len(train_data),
        "num_test": len(test_data),
        "num_epochs": num_epochs,
        "final_loss": epoch_loss / len(train_loader),
        "accuracy_pct": acc,
        "avg_latency_s": avg_lat,
        "correct": correct,
        "total": len(test_data),
        "latencies": latencies,
        "cfg_scale": cfg_w,
        "inference_steps": num_steps,
        "gpu": gpu_name,
    }


@app.function(image=image, gpu="T4", timeout=1800)
def mllm_baseline(task="maze", grid_size=8):
    """Evaluate Qwen2.5-VL-7B as MLLM baseline."""
    import torch
    from transformers import AutoModelForVision2Seq, AutoProcessor
    import numpy as np
    import time

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_id = "Qwen/Qwen2.5-VL-7B-Instruct"
    print(f"Loading {model_id} onto {device}...")
    print(f"GPU mem: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB" if torch.cuda.is_available() else "CPU")

    try:
        model = AutoModelForVision2Seq.from_pretrained(
            model_id, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True
        )
        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
        print("Model loaded successfully")
    except Exception as e:
        print(f"Failed to load: {e}")
        return {"error": str(e), "model": model_id}

    if task == "maze":
        test_data = make_maze_data(grid_size, 3)

    prompt = "Solve this maze: find a path from green start to red goal. Output coordinates as (row,col) pairs."

    latencies = []
    for item in test_data:
        messages = [{"role": "user", "content": [{"type": "image", "image": item["input"]}, {"type": "text", "text": prompt}]}]
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=[item["input"]], padding=True, return_tensors="pt").to(device)
        with torch.no_grad():
            t_start = time.time()
            outputs = model.generate(**inputs, max_new_tokens=256, do_sample=False)
            torch.cuda.synchronize()
            lat = time.time() - t_start
        latencies.append(lat)
        resp = processor.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        print(f"  Sample {len(latencies)}: latency={lat:.3f}s, response={resp[:80]}...")

    avg_lat = sum(latencies) / len(latencies) if latencies else 0
    results = {
        "task": task, "model": model_id, "avg_latency_s": avg_lat,
        "num_samples": len(test_data), "latencies": latencies
    }
    print(f"MLLM baseline: avg latency = {avg_lat:.3f}s")
    return results


@app.function(image=image, gpu="T4", timeout=3600)
def full_repro_with_cfg_ablation(task="maze", grid_size=8):
    """Run ablation on CFG scales [1, 2, 4, 7] to 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" if torch.cuda.is_available() else "cpu")

    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]
            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

    train_data = make_maze_data(grid_size, 60)
    test_data = make_maze_data(grid_size, 10)
    train_loader = DataLoader(TaskDataset(train_data), batch_size=4, shuffle=True)
    test_dataset = TaskDataset(test_data)

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

    for epoch in range(10):
        model.train()
        loss_total = 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)
            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)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss_total += loss.item()
        print(f"Train epoch {epoch+1}: loss={loss_total/len(train_loader):.6f}")

    model.eval()
    cfg_scales = [1.0, 2.0, 4.0, 7.0]
    ablation = {}

    for w in cfg_scales:
        correct = 0
        latencies = []
        for item in test_dataset:
            inp = item[0].unsqueeze(0).to(device)
            target_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):
                    t_vec = torch.full((1,), step * 5, device=device, dtype=torch.long)
                    v_c = model(x, t_vec, inp)
                    v_u = model(x, t_vec, torch.zeros_like(inp))
                    x = x + (1.0/20) * (v_u + w * (v_c - v_u))
                torch.cuda.synchronize()
                lat = time.time() - t0
            latencies.append(lat)
            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 / len(test_dataset) * 100
        ablation[f"w={w}"] = {"accuracy": acc, "avg_latency": sum(latencies)/len(latencies)}
        print(f"CFG w={w}: acc={acc:.1f}%, lat={sum(latencies)/len(latencies):.3f}s")

    return ablation


@app.local_entrypoint()
def main(task="maze"):
    print(f"=== DiffThinker Reproduction === Task: {task} | Paper: 2512.24165 | GPU: T4\n")

    # Step 1: Full training + evaluation
    print("[1/2] Training + Evaluating Flow Matching...")
    r = full_repro.remote(task=task)
    print(json.dumps(r, indent=2))

    # Step 2: CFG ablation
    print("\n[2/2] CFG Ablation...")
    abl = full_repro_with_cfg_ablation.remote(task=task)
    print(f"Ablation results: {json.dumps(abl, indent=2)}")

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