#!/usr/bin/env python3 """Génération par COMPRÉHENSION — le modèle CRÉE des images depuis les règles grokkées. PRINCIPE : comprendre → générer. Le modèle grok les primitives visuelles et leurs règles de composition → il peut GÉNÉRER de nouvelles images en composant les primitives selon les règles comprises. Pas une copie — une CRÉATION depuis la compréhension. Pipeline : 1. Cross-modal grok : image + texte → concept ID (compréhension, déjà fait) 2. Flow-matching decoder : concept AMV → intégrer du bruit vers l'image Le decoder apprend : concept → comment composer les primitives visuelles → signal 3. Génération : concept ID → AMV → flow-matching → image générée 4. Vérification : l'image générée est classifiée par le cross-modal → bon concept ? C'est le crown-jewel à l'envers : grok (a,b)→op(a,b) permet de CALCULER tout résultat. Grok les règles visuelles → GÉNÉRER toute image du concept. """ import torch, torch.nn as nn, torch.nn.functional as F, glob, os, numpy as np, time from PIL import Image from sklearn.cluster import MiniBatchKMeans from ocm26400.spectral_core import SpectralCoreBlock from ocm26400.amv import D_MODEL, PART from ocm26400.learned_vocab import LearnedVocab device = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(0) IMG_DIR = "/media/akone/SAVENVME2/Datasets/vision_tinyimagenet" N_CLUSTERS = 10 PATCH_DIM = 48 # 4x4x3 def extract_patches(path): """Image → 4 patches (4x4x3 = 48-dim chacun), puis moyenne.""" im = Image.open(path).convert("RGB").resize((8, 8)) arr = np.array(im, dtype=np.float32) / 255.0 return np.mean([arr[:4, :4].flatten(), arr[:4, 4:].flatten(), arr[4:, :4].flatten(), arr[4:, 4:].flatten()], axis=0) # ============================================================ # Flow-matching decoder : concept AMV → génère l'image # ============================================================ class FlowDecoder(nn.Module): """Apprend : concept AMV + temps t + bruit x_t → vélocité v. À l'inférence : intégrer du bruit vers l'image, conditionné par le concept.""" def __init__(self, cond_dim=D_MODEL, x_dim=PATCH_DIM, hidden=256): super().__init__() self.net = nn.Sequential( nn.Linear(cond_dim + x_dim + 1, hidden), nn.GELU(), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, x_dim) ) def forward(self, x, cond, t): """x: (B, x_dim), cond: (B, cond_dim), t: (B, 1) → vélocité (B, x_dim).""" return self.net(torch.cat([x, cond, t], dim=-1)) @torch.no_grad() def sample(self, cond, n_steps=20): """Génère depuis le bruit, conditionné par le concept AMV.""" B = cond.shape[0] x = torch.randn(B, PATCH_DIM, device=cond.device) # start: bruit for i in range(n_steps): t = torch.full((B, 1), i / n_steps, device=cond.device) v = self.forward(x, cond, t) x = x + v / n_steps # Euler integration return x # image générée def train_generation(): # 1. charger images + clustering (concept IDs) paths = sorted(glob.glob(os.path.join(IMG_DIR, "*.png")))[:1500] all_patches = np.array([extract_patches(p) for p in paths], dtype=np.float64) km = MiniBatchKMeans(n_clusters=N_CLUSTERS, batch_size=256, random_state=0, n_init=3) km.fit(all_patches) labels = km.predict(all_patches) all_patches = all_patches.astype(np.float32) # back to float32 for torch # concept canonical embeddings cv = LearnedVocab(n=N_CLUSTERS, dim=PART, init="ortho", seed=0) cv.freeze() canon = cv._matrix().to(device) # (N_CLUSTERS, PART) # 2. cross-modal classifier (pour vérifier la génération) class Classifier(nn.Module): def __init__(self): super().__init__() self.proj = nn.Linear(PATCH_DIM, D_MODEL) self.core = SpectralCoreBlock(d_model=D_MODEL, seq_len=1) self.head = nn.Linear(D_MODEL, PART) def forward(self, x): return self.head(self.core(self.proj(x).unsqueeze(1)).squeeze(1)) clf = Classifier().to(device) opt_c = torch.optim.Adam(clf.parameters(), lr=3e-3) # 3. flow decoder (générateur) decoder = FlowDecoder(cond_dim=PART, x_dim=PATCH_DIM).to(device) opt_d = torch.optim.Adam(decoder.parameters(), lr=3e-3) # split perm = np.random.permutation(len(paths)) ntr = int(len(paths) * 0.8) tr_idx, te_idx = perm[:ntr], perm[ntr:] patches_t = torch.tensor(all_patches, dtype=torch.float32).to(device) labels_t = torch.tensor(labels, dtype=torch.long).to(device) # PHASE 1 : entraîner classifier + decoder SIMULTANÉMENT print(f"[génération par compréhension] {len(paths)} images, {N_CLUSTERS} concepts", flush=True) print(f" Phase 1: classifier (comprendre) + flow decoder (générer) simultanés", flush=True) t0 = time.time() for step in range(4000): bi = torch.tensor(np.random.choice(tr_idx, 48, replace=False)) x_real = patches_t[bi] # (48, 48) patches réels y = labels_t[bi] # concept IDs cond = canon[y] # (48, PART) concept AMV # --- classifier loss (comprendre : image → concept) --- out_c = clf(x_real) loss_cls = (1 - F.cosine_similarity(out_c, cond).clamp(-1, 1)).mean() # --- flow-matching loss (générer : concept → image) --- x_0 = torch.randn_like(x_real) # bruit (48, 48) t = torch.rand(48, 1, device=device) # (48, 1) x_t = (1 - t) * x_0 + t * x_real # interpolation (48,48) v_target = x_real - x_0 # vélocité cible (48,48) v_pred = decoder(x_t, cond, t) loss_gen = F.mse_loss(v_pred, v_target) loss = loss_cls + loss_gen opt_c.zero_grad(); opt_d.zero_grad() loss.backward() opt_c.step(); opt_d.step() if step % 1000 == 0: clf.eval() with torch.no_grad(): ok = sum(1 for i in te_idx[:100] if (clf(patches_t[i:i+1]) @ canon.t()).argmax(1).item() == labels[i]) print(f" step {step} cls={loss_cls.item():.3f} gen={loss_gen.item():.3f} " f"clf_test={ok}% t={time.time()-t0:.0f}s", flush=True) clf.train() # PHASE 2 : GÉNÉRER des images depuis les concepts print(f"\n[Phase 2: GÉNÉRATION — concept AMV → flow-matching → image créée]") decoder.eval(); clf.eval() n_gen = 5 # générer 5 images par concept generated = {} verify_ok = 0; verify_tot = 0 with torch.no_grad(): for ci in range(N_CLUSTERS): cond = canon[ci:ci+1].expand(n_gen, -1) # (5, PART) gen_patches = decoder.sample(cond, n_steps=25) # (5, 48) patches générés generated[ci] = gen_patches.cpu().numpy() # VÉRIFICATION : le classifier reconnaît-il l'image générée comme le bon concept ? pred = (clf(gen_patches) @ canon.t()).argmax(1) hits = (pred == ci).sum().item() verify_ok += hits; verify_tot += n_gen gen_acc = verify_ok / max(verify_tot, 1) print(f"\n=== GÉNÉRATION PAR COMPRÉHENSION ===") print(f" images générées: {verify_tot} ({n_gen} par concept × {N_CLUSTERS})") print(f" vérification (généré reconnu comme bon concept): {verify_ok}/{verify_tot} = {gen_acc*100:.0f}%") print(f" hasard: {100/N_CLUSTERS:.0f}%") print(f" temps: {time.time()-t0:.0f}s") print(f" méthode: concept ID → flow-matching → image CRÉÉE (pas copiée)") print(f" preuve: l'image générée est reconnue par le classifier = génération cohérente") # save ckpt = "/media/akone/SAVENVME2/Datasets/ocm26400/generation_trained.pt" torch.save({"decoder": decoder.state_dict(), "classifier": clf.state_dict(), "canon": canon, "gen_acc": gen_acc, "method": "flow-matching generation from comprehension"}, ckpt) print(f" [SAUVÉ] {ckpt}") return gen_acc if __name__ == "__main__": print("="*60) print("GÉNÉRATION PAR COMPRÉHENSION — créer depuis les règles grokkées") print("="*60) acc = train_generation() print(f"\nGénération vérifiée: {acc*100:.0f}%")