Create runner.py

50k_results/runner.py

@@ -0,0 +1,591 @@
+"""
+Large-Scale Geometric Similarity - Cell 10
+============================================
+50,000 synthetic character images → FLUX VAE → Geometric Features
+Categories from generator_type field (15 types).
+
+Streams from HuggingFace datasets, encodes in batches,
+extracts gate vectors + patch features, computes similarity.
+
+Requires Cell 1 (generator.py) and Cell 2 (model.py) in namespace.
+"""
+
+import os, json, gc, time
+from pathlib import Path
+from collections import Counter, defaultdict
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.patheffects as pe
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+DATASET_ID   = "AbstractPhil/synthetic-characters"
+SUBSET       = "schnell_full_1_512"
+MODEL_REPO   = "AbstractPhil/grid-geometric-multishape"
+MODEL_FILE   = "checkpoint_v10/best_model_epoch200.pt"
+VAE_REPO     = "black-forest-labs/FLUX.1-schnell"
+
+OUTPUT_DIR   = "/content/results_50k"
+DEVICE       = "cuda" if torch.cuda.is_available() else "cpu"
+IMAGE_SIZE   = 128
+FLUX_SCALE   = 0.3611
+
+# Batch sizes — tuned for L4 (24GB VRAM)
+VAE_BATCH    = 128      # images per VAE encode
+FEAT_BATCH   = 256      # adapted latents per model forward
+
+MIN_CATEGORY_SIZE = 50  # drop categories smaller than this
+
+img_transform = transforms.Compose([
+    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+])
+
+# ── Load Models ───────────────────────────────────────────────────────────────
+
+def load_vae():
+    from diffusers import AutoencoderKL
+    print("Loading FLUX VAE...")
+    vae = AutoencoderKL.from_pretrained(
+        VAE_REPO, subfolder="vae", torch_dtype=torch.float16,
+    ).to(DEVICE).eval()
+    print("✓ VAE ready")
+    return vae
+
+
+def load_model():
+    from huggingface_hub import hf_hub_download
+    print("Loading geometric model...")
+    path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+    ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
+    config = ckpt["config"]
+    model = SuperpositionPatchClassifier(
+        embed_dim=config["embed_dim"],
+        patch_dim=config["patch_dim"],
+        n_bootstrap=config["n_bootstrap"],
+        n_geometric=config["n_geometric"],
+        n_heads=config["n_heads"],
+        dropout=0.0,
+    ).to(DEVICE).eval()
+    model.load_state_dict(ckpt["model_state_dict"])
+    print(f"✓ Model ready (epoch {ckpt['epoch']})")
+    return model
+
+
+# ── Streaming Encode + Extract ────────────────────────────────────────────────
+
+def process_image(img_pil):
+    """PIL Image → tensor ready for VAE."""
+    return img_transform(img_pil.convert("RGB"))
+
+
+def adapt_latent(z):
+    """(B, 16, H, W) → (B, 8, 16, 16)"""
+    B, C, H, W = z.shape
+    if H != 16 or W != 16:
+        z = F.interpolate(z, size=(16, 16), mode='bilinear', align_corners=False)
+    if C == 16:
+        z = z.view(B, 8, 2, 16, 16).mean(dim=2)
+    return z
+
+
+@torch.no_grad()
+def extract_gate_vectors(adapted, model):
+    """
+    adapted: (B, 8, 16, 16)
+    Returns: gate_vectors (B, 64, 17), patch_features (B, 64, 256)
+    """
+    out = model(adapted)
+
+    local_gates = torch.cat([
+        F.softmax(out["local_dim_logits"], dim=-1),
+        F.softmax(out["local_curv_logits"], dim=-1),
+        torch.sigmoid(out["local_bound_logits"]),
+        torch.sigmoid(out["local_axis_logits"]),
+    ], dim=-1)
+
+    struct_gates = torch.cat([
+        F.softmax(out["struct_topo_logits"], dim=-1),
+        torch.sigmoid(out["struct_neighbor_logits"]),
+        F.softmax(out["struct_role_logits"], dim=-1),
+    ], dim=-1)
+
+    gates = torch.cat([local_gates, struct_gates], dim=-1)
+    return gates.cpu(), out["patch_features"].cpu()
+
+
+# ── Dataset wrapper for DataLoader ────────────────────────────────────────────
+
+class HFImageDataset(torch.utils.data.Dataset):
+    """Wraps HF dataset for PyTorch DataLoader with parallel workers."""
+    def __init__(self, hf_ds):
+        self.ds = hf_ds
+        self.N = len(hf_ds)
+
+    def __len__(self):
+        return self.N
+
+    def __getitem__(self, idx):
+        row = self.ds[idx]
+        try:
+            tensor = img_transform(row["image"].convert("RGB"))
+        except:
+            tensor = torch.zeros(3, IMAGE_SIZE, IMAGE_SIZE)
+        cat = row.get("generator_type", "unknown")
+        rid = row.get("id", idx)
+        return tensor, cat, rid
+
+
+def collate_fn(batch):
+    tensors, cats, ids = zip(*batch)
+    return torch.stack(tensors), list(cats), list(ids)
+
+
+def _save_checkpoint(all_gates, all_patch, all_cats, all_ids, n):
+    g = torch.cat(all_gates) if isinstance(all_gates[0], torch.Tensor) and all_gates[0].dim() == 3 else torch.cat(all_gates)
+    p = torch.cat(all_patch) if isinstance(all_patch[0], torch.Tensor) and all_patch[0].dim() == 3 else torch.cat(all_patch)
+    path = os.path.join(OUTPUT_DIR, f"checkpoint_{n}.pt")
+    torch.save({"gates": g, "patch_feats": p, "categories": all_cats, "ids": all_ids}, path)
+    print(f"\n  💾 Checkpoint: {path} ({g.shape[0]} samples)")
+
+
+def find_latest_checkpoint(output_dir=OUTPUT_DIR):
+    """Find highest numbered checkpoint file."""
+    import glob
+    pattern = os.path.join(output_dir, "checkpoint_*.pt")
+    files = glob.glob(pattern)
+    if not files:
+        return None, 0
+    # Extract numbers
+    best_n, best_f = 0, None
+    for f in files:
+        try:
+            n = int(os.path.basename(f).replace("checkpoint_", "").replace(".pt", ""))
+            if n > best_n:
+                best_n, best_f = n, f
+        except:
+            pass
+    return best_f, best_n
+
+
+def run_extraction(ds, vae, model):
+    """
+    DataLoader with workers → VAE encode → geometric extract.
+    Resumes from latest checkpoint if available.
+    Returns: gates (N, 64, 17), patch_feats (N, 64, 256), categories list
+    """
+    from tqdm import tqdm
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    # Check for existing checkpoint
+    ckpt_path, resume_from = find_latest_checkpoint()
+    if ckpt_path:
+        print(f"\n🔄 Resuming from checkpoint: {ckpt_path} ({resume_from} samples)")
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        all_gates = [ckpt["gates"]]
+        all_patch = [ckpt["patch_feats"]]
+        all_cats  = list(ckpt["categories"])
+        all_ids   = list(ckpt["ids"])
+        processed = resume_from
+        del ckpt
+        gc.collect()
+        print(f"  ✓ Loaded {processed} cached samples")
+    else:
+        all_gates = []
+        all_patch = []
+        all_cats = []
+        all_ids = []
+        processed = 0
+
+    # Skip already-processed samples
+    N = len(ds)
+    remaining = N - resume_from
+
+    if remaining <= 0:
+        print(f"✓ All {N} samples already extracted")
+        gates = torch.cat(all_gates)
+        patch_feats = torch.cat(all_patch)
+        return gates, patch_feats, all_cats, all_ids
+
+    # Subset dataset to remaining samples
+    if resume_from > 0:
+        ds_remaining = ds.select(range(resume_from, N))
+        print(f"  Extracting remaining {remaining} samples...")
+    else:
+        ds_remaining = ds
+
+    dataset = HFImageDataset(ds_remaining)
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=VAE_BATCH,
+        shuffle=False,
+        num_workers=8,
+        pin_memory=True,
+        prefetch_factor=4,
+        collate_fn=collate_fn,
+        persistent_workers=True,
+    )
+
+    pbar = tqdm(total=remaining, unit="img", desc=f"Extracting (from {resume_from})")
+
+    for batch_pixels, cats, ids in loader:
+        batch_pixels = batch_pixels.to(DEVICE, non_blocking=True)
+
+        # VAE encode (fp16)
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            latents = vae.encode(batch_pixels.half()).latent_dist.sample() * FLUX_SCALE
+        adapted = adapt_latent(latents.float())  # geometric model expects fp32
+
+        # Extract in sub-batches
+        for fstart in range(0, adapted.shape[0], FEAT_BATCH):
+            fend = min(fstart + FEAT_BATCH, adapted.shape[0])
+            gates, patch_feats = extract_gate_vectors(adapted[fstart:fend], model)
+            all_gates.append(gates)
+            all_patch.append(patch_feats)
+
+        all_cats.extend(cats)
+        all_ids.extend(ids)
+        processed += len(cats)
+
+        pbar.update(len(cats))
+
+        # Periodic checkpoint
+        if processed % SAVE_EVERY < VAE_BATCH and processed >= SAVE_EVERY:
+            _save_checkpoint(all_gates, all_patch, all_cats, all_ids, processed)
+
+    pbar.close()
+    print(f"✓ Processed {processed} images total")
+
+    # Final checkpoint
+    _save_checkpoint(all_gates, all_patch, all_cats, all_ids, processed)
+
+    gates = torch.cat(all_gates)
+    patch_feats = torch.cat(all_patch)
+
+    return gates, patch_feats, all_cats, all_ids
+
+
+# ── Build Representations ─────────────────────────────────────────────────────
+
+def build_reps(gates, patch_feats):
+    N = gates.shape[0]
+
+    # Mean pool on GPU (49k × 64 × 256 is 3.2GB — fits L4)
+    global_feats = patch_feats.to(DEVICE).mean(dim=1).cpu()  # (N, 256)
+    torch.cuda.empty_cache()
+
+    # Normalize on GPU per-rep
+    reps = {
+        "gate_vectors": F.normalize(gates.reshape(N, -1).to(DEVICE), dim=-1).cpu(),
+        "patch_feat":   F.normalize(patch_feats.reshape(N, -1).to(DEVICE), dim=-1).cpu(),
+        "global_feat":  F.normalize(global_feats.to(DEVICE), dim=-1).cpu(),
+    }
+    torch.cuda.empty_cache()
+    return reps, global_feats
+
+
+# ── Category Similarity (size-weighted) ───────────────────────────────────────
+
+def compute_similarity(reps, cat_indices, cat_names):
+    """
+    GPU-accelerated chunked similarity.
+    Computes only the category blocks needed.
+    """
+    results = {}
+
+    for rep_name, features in reps.items():
+        print(f"  Computing: {rep_name}...")
+        features_gpu = features.to(DEVICE)
+        n_cats = len(cat_names)
+        cat_matrix = np.zeros((n_cats, n_cats))
+
+        for i, ci in enumerate(cat_names):
+            fi = features_gpu[cat_indices[ci]]  # (ni, D) on GPU
+            for j, cj in enumerate(cat_names):
+                if j < i:
+                    # Symmetric — reuse
+                    cat_matrix[i, j] = cat_matrix[j, i]
+                    continue
+
+                fj = features_gpu[cat_indices[cj]]  # (nj, D) on GPU
+
+                # Chunked matmul on GPU
+                chunk = 4000
+                block_sums = 0.0
+                block_count = 0
+                diag_sum = 0.0
+                diag_count = 0
+
+                for s in range(0, fi.shape[0], chunk):
+                    sim = fi[s:s+chunk] @ fj.T  # (chunk, nj) on GPU
+                    if i == j:
+                        # Exclude self-similarity on diagonal
+                        row_offset = s
+                        for r in range(sim.shape[0]):
+                            global_r = row_offset + r
+                            if global_r < sim.shape[1]:
+                                diag_sum += sim[r, global_r].item()
+                                diag_count += 1
+                        block_sums += sim.sum().item()
+                        block_count += sim.numel()
+                    else:
+                        block_sums += sim.sum().item()
+                        block_count += sim.numel()
+
+                if i == j:
+                    # Within: total minus diagonal, divided by off-diagonal count
+                    val = (block_sums - diag_sum) / max(block_count - diag_count, 1)
+                else:
+                    val = block_sums / max(block_count, 1)
+
+                cat_matrix[i, j] = float(val)
+                if j > i:
+                    cat_matrix[j, i] = float(val)
+
+        del features_gpu
+        torch.cuda.empty_cache()
+
+        # Size-weighted between
+        sizes = {c: len(cat_indices[c]) for c in cat_names}
+        total = sum(sizes.values())
+
+        between_sum, between_pairs = 0.0, 0
+        for i, ci in enumerate(cat_names):
+            for j, cj in enumerate(cat_names):
+                if i != j:
+                    n_pairs = sizes[ci] * sizes[cj]
+                    between_sum += cat_matrix[i, j] * n_pairs
+                    between_pairs += n_pairs
+        between_mean = between_sum / max(between_pairs, 1)
+
+        discriminability = {}
+        for i, ci in enumerate(cat_names):
+            cross_sum, cross_n = 0.0, 0
+            for j, cj in enumerate(cat_names):
+                if i != j:
+                    cross_sum += cat_matrix[i, j] * sizes[cj]
+                    cross_n += sizes[cj]
+            cat_between = cross_sum / max(cross_n, 1)
+            discriminability[ci] = float(cat_matrix[i, i] - cat_between)
+
+        overall = sum(discriminability[c] * sizes[c] / total for c in cat_names)
+
+        results[rep_name] = {
+            "matrix": cat_matrix,
+            "within": {c: float(cat_matrix[i, i]) for i, c in enumerate(cat_names)},
+            "between_mean": float(between_mean),
+            "discriminability": discriminability,
+            "overall_discriminability": float(overall),
+            "sizes": sizes,
+        }
+
+    return results
+
+
+# ── Display ───────────────────────────────────────────────────────────────────
+
+def print_results(results, cat_names):
+    first = next(iter(results.values()))
+    sizes = first["sizes"]
+    total = sum(sizes.values())
+
+    print(f"\nCategories ({len(cat_names)}, {total} total):")
+    for c in cat_names:
+        print(f"  {c:30s} n={sizes[c]:5d}  ({sizes[c]/total*100:5.1f}%)")
+
+    for rep_name, data in results.items():
+        print(f"\n{'='*80}")
+        print(f"  {rep_name}")
+        print(f"{'='*80}")
+
+        # Top/bottom within
+        within_sorted = sorted(data["within"].items(), key=lambda x: -x[1])
+        print(f"\n  Within-category similarity (top 5 / bottom 5):")
+        for c, v in within_sorted[:5]:
+            print(f"    {c:30s} {v:.4f}  (n={sizes[c]})")
+        print(f"    ...")
+        for c, v in within_sorted[-5:]:
+            print(f"    {c:30s} {v:.4f}  (n={sizes[c]})")
+
+        print(f"\n  Between-category mean: {data['between_mean']:.4f}")
+
+        # Discriminability ranked
+        disc_sorted = sorted(data["discriminability"].items(), key=lambda x: -x[1])
+        print(f"\n  Discriminability (within − weighted between):")
+        print(f"  {'Top 5':>36s}")
+        for c, d in disc_sorted[:5]:
+            sign = "+" if d > 0 else ""
+            print(f"    {c:30s} {sign}{d:.4f}")
+        print(f"  {'Bottom 5':>36s}")
+        for c, d in disc_sorted[-5:]:
+            sign = "+" if d > 0 else ""
+            print(f"    {c:30s} {sign}{d:.4f}")
+        print(f"    {'OVERALL':30s} {'+' if data['overall_discriminability'] > 0 else ''}{data['overall_discriminability']:.4f}")
+
+
+def plot_results(results, cat_names, output_dir=OUTPUT_DIR):
+    os.makedirs(output_dir, exist_ok=True)
+
+    for rep_name, data in results.items():
+        mat = data["matrix"]
+        n = len(cat_names)
+
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7), facecolor='#0a0a0a')
+
+        # Similarity matrix
+        im = ax1.imshow(mat, cmap='magma', vmin=mat.min() * 0.95, vmax=mat.max(), aspect='equal')
+        ax1.set_xticks(range(n))
+        ax1.set_yticks(range(n))
+        short = [c.replace("character_", "").replace("_", "\n") for c in cat_names]
+        ax1.set_xticklabels(short, fontsize=6, color='white', rotation=45, ha='right')
+        ax1.set_yticklabels(short, fontsize=6, color='white')
+        for i in range(n):
+            for j in range(n):
+                ax1.text(j, i, f'{mat[i,j]:.3f}', ha='center', va='center',
+                        fontsize=5, color='white' if mat[i,j] < np.median(mat) else 'black')
+        ax1.set_title(f"{rep_name} — Similarity Matrix", color='white', fontsize=10, fontweight='bold')
+        ax1.tick_params(colors='white')
+        plt.colorbar(im, ax=ax1, fraction=0.046, pad=0.04)
+
+        # Discriminability bar chart
+        ax2.set_facecolor('#0a0a0a')
+        disc = data["discriminability"]
+        disc_sorted = sorted(disc.items(), key=lambda x: -x[1])
+        names_d = [x[0].replace("character_", "") for x in disc_sorted]
+        vals_d = [x[1] for x in disc_sorted]
+        colors = ['#00b894' if v > 0 else '#e17055' for v in vals_d]
+
+        ax2.barh(range(len(names_d)), vals_d, color=colors, edgecolor='white', linewidth=0.3)
+        ax2.set_yticks(range(len(names_d)))
+        ax2.set_yticklabels(names_d, fontsize=7, color='white')
+        ax2.axvline(0, color='white', linewidth=0.5, alpha=0.5)
+        ax2.axvline(data["overall_discriminability"], color='#fdcb6e',
+                    linewidth=1, linestyle='--', alpha=0.8, label=f'overall={data["overall_discriminability"]:.4f}')
+        ax2.set_xlabel("Discriminability", color='white', fontsize=9)
+        ax2.set_title(f"{rep_name} — Discriminability", color='white', fontsize=10, fontweight='bold')
+        ax2.tick_params(colors='white', labelsize=7)
+        ax2.spines['bottom'].set_color('white')
+        ax2.spines['left'].set_color('white')
+        ax2.spines['top'].set_visible(False)
+        ax2.spines['right'].set_visible(False)
+        ax2.legend(fontsize=7, framealpha=0.7, facecolor='#1a1a2e', labelcolor='white')
+
+        safe_name = rep_name.replace(" ", "_").replace("(", "").replace(")", "")
+        path = os.path.join(output_dir, f"{safe_name}.png")
+        fig.savefig(path, dpi=150, bbox_inches='tight', facecolor=fig.get_facecolor())
+        plt.close(fig)
+        print(f"✓ Plot: {path}")
+
+
+# ── Save ──────────────────────────────────────────────────────────────────────
+
+def save_final(gates, patch_feats, categories, ids, results, cat_names):
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    # Features
+    tpath = os.path.join(OUTPUT_DIR, "geometric_features_50k.pt")
+    torch.save({
+        "gate_vectors":   gates,
+        "patch_features": patch_feats,
+        "global_features": patch_feats.to(DEVICE).mean(dim=1).cpu(),
+        "categories":     categories,
+        "ids":            ids,
+        "cat_names":      cat_names,
+    }, tpath)
+    print(f"✓ Saved: {tpath}")
+    print(f"  gates: {gates.shape}, patch_feats: {patch_feats.shape}")
+
+    # Similarity JSON
+    out = {}
+    for rep_name, data in results.items():
+        out[rep_name] = {
+            "within": data["within"],
+            "between_mean": data["between_mean"],
+            "discriminability": data["discriminability"],
+            "overall_discriminability": data["overall_discriminability"],
+            "sizes": data["sizes"],
+            "matrix": data["matrix"].tolist(),
+        }
+    jpath = os.path.join(OUTPUT_DIR, "similarity_results_50k.json")
+    with open(jpath, "w") as f:
+        json.dump(out, f, indent=2)
+    print(f"✓ Saved: {jpath}")
+
+
+# ── Main ────────────────────────────────────────��─────────────────────────────
+
+def run_50k():
+    from datasets import load_dataset
+
+    # 1. Load dataset
+    print(f"Loading dataset: {DATASET_ID} / {SUBSET}...")
+    ds = load_dataset(DATASET_ID, SUBSET, split="train")
+    print(f"✓ {len(ds)} samples loaded")
+
+    # Show category distribution
+    cats = ds["generator_type"]
+    cat_counts = Counter(cats)
+    print(f"\nGenerator type distribution:")
+    for c, n in cat_counts.most_common():
+        print(f"  {c:30s} {n:6d}  ({n/len(ds)*100:5.1f}%)")
+
+    # 2. Load models
+    vae = load_vae()
+    model = load_model()
+
+    # 3. Stream encode + extract
+    gates, patch_feats, categories, ids = run_extraction(ds, vae, model)
+
+    # 4. Free VAE
+    del vae
+    gc.collect()
+    torch.cuda.empty_cache()
+    print("✓ Freed VAE memory")
+
+    # 5. Build category indices (with minimum size filter)
+    cat_counts_final = Counter(categories)
+    cat_names = sorted([c for c, n in cat_counts_final.items() if n >= MIN_CATEGORY_SIZE])
+    dropped = [c for c, n in cat_counts_final.items() if n < MIN_CATEGORY_SIZE]
+    if dropped:
+        print(f"\n⚠ Dropping {len(dropped)} categories with < {MIN_CATEGORY_SIZE} samples: {dropped}")
+
+    # Build index mapping (vectorized)
+    cat_indices = {}
+    cat_array = np.array(categories)
+    for c in cat_names:
+        cat_indices[c] = torch.from_numpy(np.where(cat_array == c)[0]).long()
+
+    total_used = sum(len(v) for v in cat_indices.values())
+    print(f"\nUsing {len(cat_names)} categories, {total_used}/{len(categories)} samples")
+
+    # 6. Build representations
+    print("\nBuilding representations...")
+    reps, global_feats = build_reps(gates, patch_feats)
+
+    # 7. Compute similarity
+    print("Computing category similarity (chunked)...")
+    sim_results = compute_similarity(reps, cat_indices, cat_names)
+
+    # 8. Display
+    print_results(sim_results, cat_names)
+
+    # 9. Plot
+    plot_results(sim_results, cat_names)
+
+    # 10. Save
+    save_final(gates, patch_feats, categories, ids, sim_results, cat_names)
+
+    return sim_results, gates, patch_feats, cat_indices, cat_names
+
+
+# ── Run ───────────────────────────────────────────────────────────────────────
+sim_results, gates, patch_feats, cat_indices, cat_names = run_50k()