Dreanhunter30/Light_Multimodal_Learning_The_SpectralMix_Head

1# app.py — SpectralMix Super UI: One Head, Two Modalities + Simulator2# HF Spaces-friendly, CPU-ready (Gradio 4.x)3 4import io5import json6import platform7from pathlib import Path8 9import gradio as gr10import numpy as np11import torch12from PIL import Image13import matplotlib.pyplot as plt14 15# -------- local imports (должны быть в проекте) --------16# infer.py обязан экспонировать: load_classes, load_head, predict_image, routing_for_text17from infer import load_classes, load_head, predict_image, routing_for_text18 19# ---------------- Config ----------------20ROOT = Path(".")21CLASS_CANDIDATES = [ROOT / "classes.json", ROOT / "классы.json"]22WEIGHT_CANDIDATES = [ROOT / "head.pt", ROOT / "weights.pt", ROOT / "голова.pt", ROOT / "веса.pt"]23TAU_DEFAULT = 0.5  # стартовая температура (меняется слайдером)24 25# ---------------- Utilities ----------------26def pick_first_existing(paths):27    for p in paths:28        if p.exists():29            return p30    return None31 32def load_classes_fallback():33    p = pick_first_existing(CLASS_CANDIDATES)34    if p:35        try:36            return load_classes(str(p))37        except Exception:38            pass39    # fallback: try meta.json (root or 'artifacts'), else CIFAR-10 defaults40    for mp in [ROOT / "meta.json", ROOT / "artifacts" / "meta.json"]:41        if mp.exists():42            try:43                with open(mp, "r", encoding="utf-8") as f:44                    meta = json.load(f)45                    if isinstance(meta.get("classes"), list):46                        return meta["classes"]47            except Exception:48                pass49    return ["airplane","automobile","bird","cat","deer","dog","frog","horse","ship","truck"]50 51# Classes52CLASSES = load_classes_fallback()53 54# Weights (root only)55_weight_path = pick_first_existing(WEIGHT_CANDIDATES)56if _weight_path is None:57    print("[error] Model weights not found in repo root. Expected one of:",58          ", ".join([p.name for p in WEIGHT_CANDIDATES]))59 60# Загружаем голову (и переводим в eval)61HEAD = None62if _weight_path is not None:63    HEAD = load_head(str(_weight_path), num_classes=len(CLASSES))64    if hasattr(HEAD, "eval"):65        HEAD.eval()66    print(f"[info] Using local weights: {_weight_path}")67 68# ---------------- Plotting helpers (matplotlib; no custom colors) ----------------69def _fig_to_pil(fig):70    buf = io.BytesIO()71    fig.savefig(buf, format="png", bbox_inches="tight")72    plt.close(fig)73    buf.seek(0)74    return Image.open(buf)75 76def plot_topk(labels, probs, title="Top-5 classes"):77    fig = plt.figure()78    if labels and probs:79        x = list(range(len(labels)))80        plt.bar(x, probs)81        plt.xticks(x, labels, rotation=30, ha="right")82    plt.ylabel("Probability")83    plt.title(title)84    return _fig_to_pil(fig)85 86def plot_expert_bars(vec, title="Expert activations (predicted class)"):87    vec = np.asarray(vec, dtype=np.float64)88    fig = plt.figure()89    x = list(range(len(vec)))90    plt.bar(x, vec)91    plt.xticks(x, [f"E{i}" for i in x], rotation=0)92    plt.ylabel("Routing weight")93    plt.title(title)94    return _fig_to_pil(fig)95 96def plot_routing_heatmap(routing, class_labels, title="Routing heatmap (classes × experts)"):97    arr = np.array(routing, dtype=np.float32)98    fig = plt.figure()99    im = plt.imshow(arr, aspect="auto", vmin=0.0, vmax=1.0)100    plt.colorbar(im)101    plt.yticks(range(len(class_labels)), class_labels)102    plt.xlabel("Experts")103    plt.ylabel("Classes")104    plt.title(title)105    return _fig_to_pil(fig)106 107def _softmax_np(x, tau=1.0):108    x = np.asarray(x, dtype=np.float64) / max(tau, 1e-12)109    x = x - x.max()110    ex = np.exp(x)111    s = ex.sum()112    return (ex / (s if s > 0 else 1.0)).astype(np.float64)113 114def cosine_sim(a, b, eps=1e-12):115    a = np.asarray(a, dtype=np.float64)116    b = np.asarray(b, dtype=np.float64)117    na = np.linalg.norm(a)118    nb = np.linalg.norm(b)119    if na < eps or nb < eps:120        return 0.0121    return float(np.dot(a, b) / (na * nb))122 123# ---------------- Inference wrappers ----------------124def _ensure_ready():125    if HEAD is None:126        raise gr.Error(127            "Model weights not found in repo root. "128            "Place `head.pt` (или `weights.pt` / `голова.pt` / `веса.pt`) рядом с app.py."129        )130 131def _preprocess_image(img: Image.Image) -> Image.Image:132    if img is None:133        raise gr.Error("Upload an image.")134    img = img.convert("RGB")135    # лёгкий ресайз, чтобы не ронять CPU136    max_side = max(img.size)137    if max_side > 1024:138        scale = 1024.0 / max_side139        img = img.resize((int(img.width * scale), int(img.height * scale)))140    return img141 142@torch.no_grad()143def run_image(img: Image.Image, tau: float):144    _ensure_ready()145    img = _preprocess_image(img)146    # Предсказание (детерминированно, без Gumbel-шума — в infer.py выключен шум на eval)147    idxs, vals, routing = predict_image(HEAD, img, CLASSES, tau=float(tau))  # routing [K,R]148    if len(idxs) == 0:149        empty = {"class": [], "prob": []}150        blank = plot_topk([], [], title="Top-5 (image)")151        return empty, blank, blank, blank, "No prediction."152    labels = [CLASSES[i] for i in idxs]153    probs = [float(v) for v in vals]154    pred_idx = int(idxs[0])155    expert_vec = routing[pred_idx].tolist()156 157    top5_plot = plot_topk(labels, probs, title="Top-5 (image)")158    experts_plot = plot_expert_bars(expert_vec, title="Experts (image, predicted class)")159    heatmap_plot = plot_routing_heatmap(routing, CLASSES, title="Routing heatmap (image)")160 161    top_ex = torch.topk(torch.tensor(expert_vec), k=min(5, len(expert_vec)))162    details = (163        f"Predicted: {CLASSES[pred_idx]}  |  τ={tau:.3f}\n"164        f"Top experts (idx:weight): " +165        ", ".join([f"{int(i)}:{float(w):.3f}" for i, w in zip(top_ex.indices.tolist(), top_ex.values.tolist())]) +166        f"\nSum(weights)={float(np.sum(expert_vec)):.3f}"167    )168    return {"class": labels, "prob": probs}, top5_plot, experts_plot, heatmap_plot, details169 170@torch.no_grad()171def run_text(text: str, tau: float):172    _ensure_ready()173    if not text or not text.strip():174        empty = {"class": [], "prob": []}175        blank = plot_topk([], [], title="Top-5 (text)")176        return empty, blank, blank, blank, "Enter a non-empty text."177    probs, routing = routing_for_text(HEAD, text, tau=float(tau))  # probs [K], routing [K,R]178    if probs is None or routing is None:179        blank = plot_topk([], [], title="Top-5 (text)")180        return {"class": [], "prob": []}, blank, blank, blank, "No output from text pipeline."181    top = torch.topk(probs, k=min(5, len(CLASSES)))182    labels = [CLASSES[i] for i in top.indices.tolist()]183    pvals = [float(v) for v in top.values.tolist()]184    pred_idx = int(torch.argmax(probs).item())185    expert_vec = routing[pred_idx].tolist()186 187    top5_plot = plot_topk(labels, pvals, title="Top-5 (text)")188    experts_plot = plot_expert_bars(expert_vec, title="Experts (text, predicted class)")189    heatmap_plot = plot_routing_heatmap(routing, CLASSES, title="Routing heatmap (text)")190 191    top_ex = torch.topk(torch.tensor(expert_vec), k=min(5, len(expert_vec)))192    details = (193        f"Predicted: {CLASSES[pred_idx]}  |  τ={tau:.3f}\n"194        f"Top experts (idx:weight): " +195        ", ".join([f"{int(i)}:{float(w):.3f}" for i, w in zip(top_ex.indices.tolist(), top_ex.values.tolist())]) +196        f"\nSum(weights)={float(np.sum(expert_vec)):.3f}"197    )198    return {"class": labels, "prob": pvals}, top5_plot, experts_plot, heatmap_plot, details199 200@torch.no_grad()201def run_compare(img: Image.Image, text: str, tau: float):202    _ensure_ready()203    if img is None or not text or not text.strip():204        blank = plot_topk([], [], title="Top-5")205        return blank, blank, "Provide both image and text."206    img = _preprocess_image(img)207    # image208    idxs_i, _, routing_i = predict_image(HEAD, img, CLASSES, tau=float(tau))209    if len(idxs_i) == 0:210        blank = plot_topk([], [], title="Top-5")211        return blank, blank, "No image prediction."212    pred_i = int(idxs_i[0])213    vec_i = routing_i[pred_i].tolist()214    # text215    probs_t, routing_t = routing_for_text(HEAD, text, tau=float(tau))216    if probs_t is None:217        blank = plot_topk([], [], title="Top-5")218        return blank, blank, "No text prediction."219    pred_t = int(torch.argmax(probs_t).item())220    vec_t = routing_t[pred_t].tolist()221 222    sim = cosine_sim(vec_i, vec_t)223    bar_i = plot_expert_bars(vec_i, title=f"Experts (image → {CLASSES[pred_i]})")224    bar_t = plot_expert_bars(vec_t, title=f"Experts (text  → {CLASSES[pred_t]})")225    info = (226        f"Predicted (image): {CLASSES[pred_i]} | Predicted (text): {CLASSES[pred_t]}\n"227        f"Cosine similarity of expert vectors: {sim:.3f}  (1.0 = identical, 0 = orthogonal)  |  τ={tau:.3f}"228    )229    return bar_i, bar_t, info230 231# ---------------- Simulator (how the head works) ----------------232_rng_global = np.random.default_rng(12345)233 234def _gumbel_softmax_np(logits, tau=1.0, hard=False, rng=None):235    rng = rng or _rng_global236    U = rng.uniform(low=1e-8, high=1-1e-8, size=len(logits))237    g = -np.log(-np.log(U))238    y = _softmax_np(np.asarray(logits) + g, tau=tau)239    if hard:240        hard_vec = np.zeros_like(y)241        hard_vec[int(np.argmax(y))] = 1.0242        return hard_vec243    return y244 245def simulate_mix(R=4, C=10, tau=0.5, hard=False, couple=0.5, seed=42, use_gumbel=False):246    """247    Возвращает:248      W: [R, C] — компонентные логиты для классов249      alpha_img, alpha_txt: [R] — смеси экспертов250      z_img, z_txt: [C] — итоговые логиты по классам251    """252    rng = np.random.default_rng(int(seed))253    # компонентные "головы": W_r \in R^{C}254    W = rng.normal(loc=0.0, scale=1.0, size=(R, C)).astype(np.float64)255    # гейтовые логиты256    gate_img = rng.normal(size=R)257    noise = rng.normal(size=R)258    gate_txt = couple * gate_img + (1.0 - couple) * noise259 260    if use_gumbel:261        alpha_img = _gumbel_softmax_np(gate_img, tau=tau, hard=hard, rng=rng)262        alpha_txt = _gumbel_softmax_np(gate_txt, tau=tau, hard=hard, rng=rng)263    else:264        alpha_img = _softmax_np(gate_img, tau=tau)265        alpha_txt = _softmax_np(gate_txt, tau=tau)266        if hard:267            h = np.zeros_like(alpha_img); h[int(np.argmax(alpha_img))] = 1.0; alpha_img = h268            h = np.zeros_like(alpha_txt); h[int(np.argmax(alpha_txt))] = 1.0; alpha_txt = h269 270    z_img = alpha_img @ W271    z_txt = alpha_txt @ W272    return W, alpha_img, alpha_txt, z_img, z_txt273 274def plot_class_logits(logits, class_labels, title):275    probs = _softmax_np(np.asarray(logits), tau=1.0)276    fig = plt.figure()277    x = list(range(len(class_labels)))278    plt.bar(x, probs)279    plt.xticks(x, class_labels, rotation=30, ha="right")280    plt.ylabel("Probability")281    plt.title(title)282    return _fig_to_pil(fig)283 284def plot_W_heatmap(W, title="Component heads W (R × C)"):285    arr = np.asarray(W, dtype=np.float32)286    if arr.shape[1] > 0:287        col_std = arr.std(axis=0, keepdims=True) + 1e-9288        arr = (arr - arr.mean(axis=0, keepdims=True)) / col_std289    fig = plt.figure()290    im = plt.imshow(arr, aspect="auto")291    plt.colorbar(im)292    plt.xlabel("Classes")293    plt.ylabel("Experts")294    plt.title(title)295    return _fig_to_pil(fig)296 297def run_simulator(R, C, tau, couple, seed, hard, use_gumbel):298    # валидные диапазоны299    R = int(max(2, min(16, R)))300    C = int(max(2, min(20, C)))301    # классы для подписи302    if len(CLASSES) >= C:303        cls = CLASSES[:C]304    else:305        cls = (CLASSES + [f"class_{i}" for i in range(C - len(CLASSES))])[:C]306 307    W, a_img, a_txt, z_img, z_txt = simulate_mix(308        R=R, C=C, tau=float(tau), hard=bool(hard),309        couple=float(couple), seed=int(seed), use_gumbel=bool(use_gumbel)310    )311    bars_img = plot_expert_bars(a_img, title="α (image)")312    bars_txt = plot_expert_bars(a_txt, title="α (text)")313    heat_W  = plot_W_heatmap(W, title="Components W (experts × classes)")314    cls_img = plot_class_logits(z_img, cls, title="Mixed logits → probs (image)")315    cls_txt = plot_class_logits(z_txt, cls, title="Mixed logits → probs (text)")316 317    sim = cosine_sim(a_img, a_txt)318    info = (319        f"Cosine(α_image, α_text) = {sim:.3f}  |  τ={tau:.3f} "320        f"| {'Gumbel' if use_gumbel else 'Softmax'} | {'hard' if hard else 'soft'} | couple={couple:.2f}"321    )322    top_ex_img = ", ".join([f"E{i}:{w:.3f}" for i, w in enumerate(a_img)])323    top_ex_txt = ", ".join([f"E{i}:{w:.3f}" for i, w in enumerate(a_txt)])324    details = f"α_image: [{top_ex_img}]\nα_text : [{top_ex_txt}]"325    return bars_img, bars_txt, heat_W, cls_img, cls_txt, info, details326 327# ---------------- UI helpers ----------------328def _env_banner(weights_path: str, classes_len: int):329    import importlib330    pkgs = {}331    try:332        import torchvision as _tv333        pkgs["torchvision"] = _tv.__version__334    except Exception:335        pkgs["torchvision"] = "n/a"336    try:337        import transformers as _tf338        pkgs["transformers"] = _tf.__version__339    except Exception:340        pkgs["transformers"] = "n/a"341    try:342        import open_clip_torch as _oc343        pkgs["open_clip_torch"] = getattr(_oc, "__version__", "present")344    except Exception:345        pkgs["open_clip_torch"] = "n/a"346 347    rows = [348        f"- **torch**: {torch.__version__}",349        f"- **torchvision**: {pkgs['torchvision']}",350        f"- **gradio**: {gr.__version__}",351        f"- **numpy**: {np.__version__}",352        f"- **matplotlib**: {plt.matplotlib.__version__}",353        f"- **transformers**: {pkgs['transformers']}",354        f"- **open_clip_torch**: {pkgs['open_clip_torch']}",355    ]356    sysrow = f"- **Python**: {platform.python_version()}  |  **Device**: {'CUDA' if torch.cuda.is_available() else 'CPU'}"357    wrow = f"**Using weights**: `{weights_path}`  |  **#classes**: {classes_len}"358    return wrow + "<br>" + sysrow + "<br>" + "<br>".join(rows)359 360# ---------------- UI ----------------361with gr.Blocks(title="SpectralMix — One Head, Two Modalities (Image/Text)") as demo:362    gr.HTML("""363    <style>.notranslate { translate: no; }</style>364    <div class="notranslate" style="font-size:28px; font-weight:700;">🧠 SpectralMix — One Head, Two Modalities</div>365    <div>One classifier head for <i>image</i> and <i>text</i>. We mix a few <b>orthogonal experts</b> with (Gumbel-)Softmax.<br>366    Same class ⇒ same expert (similarity≈1). Different classes ⇒ different experts (≈0). This avoids catastrophic forgetting across stages.</div>367    """)368    gr.Markdown(_env_banner(str(_weight_path) if _weight_path else "N/A", len(CLASSES)))369 370    # Общая температура для инференса371    tau_slider = gr.Slider(0.1, 2.0, value=TAU_DEFAULT, step=0.05, label="Temperature τ (inference)")372 373    with gr.Tab("Image → classes & experts"):374        with gr.Row():375            img_in = gr.Image(type="pil", label="Image", height=260)376            with gr.Column():377                top5_out = gr.JSON(label="Top-5 classes")378                details = gr.Textbox(label="Details", lines=3)379        with gr.Row():380            probs_plot   = gr.Image(label="Top-5 probabilities", height=260)381            experts_plot = gr.Image(label="Expert activations (predicted class)", height=260)382            routing_plot = gr.Image(label="Routing heatmap (classes × experts)", height=260)383 384        gr.Button("Predict").click(385            fn=run_image,386            inputs=[img_in, tau_slider],387            outputs=[top5_out, probs_plot, experts_plot, routing_plot, details],388            concurrency_limit=2389        )390 391    with gr.Tab("Text → classes & experts"):392        with gr.Row():393            txt_in = gr.Textbox(label="Text prompt (e.g., 'a photo of a dog')", lines=2)394            with gr.Column():395                txt_top = gr.JSON(label="Top-5 classes")396                txt_details = gr.Textbox(label="Details", lines=3)397        with gr.Row():398            txt_probs_plot   = gr.Image(label="Top-5 probabilities", height=260)399            txt_experts_plot = gr.Image(label="Expert activations (predicted class)", height=260)400            txt_routing_plot = gr.Image(label="Routing heatmap (classes × experts)", height=260)401 402        gr.Button("Run").click(403            fn=run_text, inputs=[txt_in, tau_slider],404            outputs=[txt_top, txt_probs_plot, txt_experts_plot, txt_routing_plot, txt_details],405            concurrency_limit=2406        )407 408    with gr.Tab("Compare (Image vs Text)"):409        with gr.Row():410            c_img = gr.Image(type="pil", label="Image", height=240)411            c_txt = gr.Textbox(label="Text prompt", lines=2)412        with gr.Row():413            bar_i = gr.Image(label="Experts (image)", height=260)414            bar_t = gr.Image(label="Experts (text)", height=260)415        info = gr.Textbox(label="Similarity", lines=2)416        gr.Button("Compare").click(417            run_compare, inputs=[c_img, c_txt, tau_slider], outputs=[bar_i, bar_t, info],418            concurrency_limit=2419        )420        gr.Markdown(421            "Tip: same-class pairs should give similarity ≈ **1.0**, different-class pairs → ≈ **0.0**."422        )423 424    with gr.Tab("Simulator (how the head works)"):425        gr.Markdown("Interactive simulator of routing and mixing. "426                    "Adjust #experts, temperature, and modality coupling; see α-vectors and class probabilities.")427        with gr.Row():428            R_in   = gr.Slider(2, 16, value=4, step=1, label="#Experts (R)")429            C_in   = gr.Slider(2, 20, value=min(10, len(CLASSES)), step=1, label="#Classes (C)")430            tau_in = gr.Slider(0.1, 2.0, value=0.5, step=0.05, label="Temperature τ (sim)")431        with gr.Row():432            couple_in  = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Coupling of modalities (0=different, 1=identical)")433            seed_in    = gr.Number(value=42, precision=0, label="Seed")434            hard_in    = gr.Checkbox(False, label="Hard routing (one-hot)")435            gumbel_in  = gr.Checkbox(False, label="Use Gumbel-Softmax (stochastic)")436 437        with gr.Row():438            sim_alpha_img = gr.Image(label="α (image)", height=220)439            sim_alpha_txt = gr.Image(label="α (text)",  height=220)440            sim_W_heat    = gr.Image(label="Components (W)", height=220)441        with gr.Row():442            sim_cls_img = gr.Image(label="Image → probs", height=240)443            sim_cls_txt = gr.Image(label="Text  → probs", height=240)444 445        sim_info    = gr.Textbox(label="Summary", lines=2)446        sim_details = gr.Textbox(label="Details", lines=3)447 448        gr.Button("Simulate").click(449            run_simulator,450            inputs=[R_in, C_in, tau_in, couple_in, seed_in, hard_in, gumbel_in],451            outputs=[sim_alpha_img, sim_alpha_txt, sim_W_heat, sim_cls_img, sim_cls_txt, sim_info, sim_details],452            concurrency_limit=2453        )454 455    # Healthcheck (скрытый блок, удобно при отладке)456    with gr.Row(visible=False):457        btn = gr.Button("Healthcheck")458        txt = gr.Textbox()459        def _healthcheck():460            try:461                _ensure_ready()462                return "ok"463            except Exception as e:464                return f"error: {e}"465        btn.click(lambda: _healthcheck(), outputs=txt, concurrency_limit=1)466 467# очередь и запуск (Gradio 4.x — без concurrency_count)468demo.queue()469 470if __name__ == "__main__":471    demo.launch(server_name="0.0.0.0", server_port=7860)472