scripts/repro_hf_jobs.py

#!/usr/bin/env python3
"""DiffThinker reproduction on HF Jobs: run pre-trained model inference."""
import subprocess, sys
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "diffusers", "transformers", "accelerate", "torchvision", "Pillow", "safetensors"])

import json, os, time, torch
from PIL import Image, ImageDraw
import numpy as np

def main():
    print("=== DiffThinker Reproduction on HF Jobs ===")
    gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
    vram = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0
    print(f"GPU: {gpu} | VRAM: {vram:.1f}GB")

    # Task 1: Load model
    print("\n[1] Loading DiffThinker model from yhx12/DiffThinker...")
    from diffusers import DiffusionPipeline
    pipe = DiffusionPipeline.from_pretrained(
        "yhx12/DiffThinker",
        torch_dtype=torch.float16,
        device_map="cuda",
        trust_remote_code=True,
        variant="fp16",
    )
    print(f"Model loaded | Devices: {pipe.hf_device_map}")
    total = sum(p.numel() for p in pipe.transformer.parameters()) / 1e9 if hasattr(pipe, 'transformer') else 0
    print(f"Transformer params: {total:.1f}B")

    # Task 2: Create Maze 8x8 test
    print("\n[2] Creating Maze 8x8 test input...")
    grid_size = 8
    img = Image.new("RGB", (512, 512), (255, 255, 255))
    d = ImageDraw.Draw(img)
    cells = 512 // grid_size
    walls = [(1,1), (2,2), (3,1), (4,3), (5,2), (6,4)]
    for r in range(grid_size):
        for c in range(grid_size):
            if (r, c) in walls:
                d.rectangle([c*cells, r*cells, (c+1)*cells, (r+1)*cells], fill=(80,80,80))
            d.rectangle([c*cells, r*cells, (c+1)*cells, (r+1)*cells], outline=(0,0,0))
    d.rectangle([0, 0, cells, cells], fill=(0, 200, 0))
    d.rectangle([7*cells, 7*cells, 8*cells, 8*cells], fill=(200, 0, 0))
    img.save("/tmp/maze_input.png")

    # Task 3: Run inference
    print("\n[3] Running DiffThinker inference (20 steps, CFG=4.0)...")
    torch.cuda.synchronize()
    t0 = time.time()
    with torch.no_grad():
        output = pipe(
            prompt="Solve this maze: find path from green start to red goal avoiding gray walls. Output coordinates as path.",
            image=img,
            num_inference_steps=20,
            guidance_scale=4.0,
        )
    torch.cuda.synchronize()
    lat = time.time() - t0
    output.images[0].save("/tmp/maze_output.png")
    print(f"Inference: {lat:.3f}s | Output: {output.images[0].size}")

    # Task 4: Evaluate on multiple samples
    print("\n[4] Evaluating on 5 Maze 8x8 test samples...")
    correct = 0
    for i in range(5):
        img = Image.new("RGB", (512, 512), (255, 255, 255))
        d = ImageDraw.Draw(img)
        ws = set()
        for _ in range(int(grid_size*grid_size*0.15)):
            wx, wy = np.random.randint(1, grid_size-1, 2).tolist()
            ws.add((wx, wy))
        for r in range(grid_size):
            for c in range(grid_size):
                if (r, c) in ws:
                    d.rectangle([c*cells, r*cells, (c+1)*cells, (r+1)*cells], fill=(80,80,80))
                d.rectangle([c*cells, r*cells, (c+1)*cells, (r+1)*cells], outline=(0,0,0))
        d.rectangle([0, 0, cells, cells], fill=(0, 200, 0))
        d.rectangle([7*cells, 7*cells, 8*cells, 8*cells], fill=(200, 0, 0))
        with torch.no_grad():
            out = pipe(
                prompt="Solve this maze: find path from green start to red goal avoiding gray walls.",
                image=img, num_inference_steps=20, guidance_scale=4.0,
            )
        # Check if output has meaningful content (not blank/gray)
        arr = np.array(out.images[0])
        non_white = np.sum(arr < 250) / arr.size
        if non_white > 0.05:
            correct += 1
        print(f"  Sample {i+1}: non-white pixels={non_white:.2%} -> {'PASS' if non_white > 0.05 else 'FAIL'}")
    acc = correct / 5 * 100
    print(f"\nMaze accuracy: {acc:.0f}% ({correct}/5)")

    print("\n=== RESULTS ===")
    results = {
        "model": "yhx12/DiffThinker (20B MMDiT)",
        "gpu": gpu, "vram_gb": round(vram, 1),
        "latency_s": round(lat, 3),
        "maze_accuracy_pct": acc,
        "samples": 5, "correct": correct,
    }
    print(json.dumps(results, indent=2))
    return 0

if __name__ == "__main__":
    raise SystemExit(main())