import torch, json, time, re, glob, traceback
import numpy as np
from transformers import AutoModelForImageTextToText, AutoProcessor, AutoConfig
from PIL import Image
MODEL="/model"
cfg=AutoConfig.from_pretrained(MODEL); tc=cfg.text_config
L,E,K=tc.num_hidden_layers,tc.num_experts,tc.num_experts_per_tok
IMG_TOK=getattr(cfg,"image_token_id",None)
print(f"L={L} E={E} K={K} img_token={IMG_TOK}",flush=True)
proc=AutoProcessor.from_pretrained(MODEL,trust_remote_code=True)
print("cargando modelo...",flush=True); t0=time.time()
model=AutoModelForImageTextToText.from_pretrained(MODEL,device_map="auto",dtype=torch.bfloat16); model.eval()
print(f"cargado {time.time()-t0:.0f}s",flush=True)
# contadores: visual (solo tokens de imagen) y total (todos)
counts_vis=np.zeros((L,E),dtype=np.int64)
counts_all=np.zeros((L,E),dtype=np.int64)
state={"imgmask":None}
def mk(li):
    def h(m,i,o):
        lg=o[0] if isinstance(o,(tuple,list)) else o
        if not torch.is_tensor(lg): return
        lg=lg.reshape(-1,lg.shape[-1])
        if lg.shape[-1]!=E: return
        idx=lg.float().topk(K,-1).indices  # [tokens,K]
        counts_all[li]+=torch.bincount(idx.flatten(),minlength=E).cpu().numpy()
        mm=state["imgmask"]
        if mm is not None and mm.shape[0]==idx.shape[0]:
            vis=idx[mm.to(idx.device)]
            if vis.numel(): counts_vis[li]+=torch.bincount(vis.flatten(),minlength=E).cpu().numpy()
    return h
for n,mod in model.named_modules():
    if re.search(r"layers\.(\d+)\.mlp\.gate$",n):
        mod.register_forward_hook(mk(int(re.search(r"layers\.(\d+)\.",n).group(1))))
imgs=sorted(glob.glob("/pimg/*.jpg"))[:55]
print(f"{len(imgs)} imagenes",flush=True)
done=0;errs=0;t0=time.time()
for ii,ip in enumerate(imgs):
    try:
        image=Image.open(ip).convert("RGB")
        msgs=[{"role":"user","content":[{"type":"image"},{"type":"text","text":"Describe esta imagen para un negocio de lenceria."}]}]
        text=proc.apply_chat_template(msgs,add_generation_prompt=True,tokenize=False)
        inputs=proc(text=[text],images=[image],return_tensors="pt").to(model.device)
        ids=inputs["input_ids"][0]
        state["imgmask"]=(ids==IMG_TOK) if IMG_TOK is not None else None
        with torch.no_grad(): model(**inputs,use_cache=False)
        nvis=int(state["imgmask"].sum()) if state["imgmask"] is not None else 0
        state["imgmask"]=None
        done+=1
        if ii<3: print(f"  img{ii}: {ids.shape[0]} tok, {nvis} visuales",flush=True)
    except Exception as e:
        errs+=1
        if errs<=3: traceback.print_exc()
    if (ii+1)%10==0:
        np.save("/counts_vision.npy",counts_vis); np.save("/counts_visall.npy",counts_all)
        print(f"img {ii+1}/{len(imgs)} {time.time()-t0:.0f}s",flush=True)
np.save("/counts_vision.npy",counts_vis); np.save("/counts_visall.npy",counts_all)
nv=(counts_vis>0).sum(1)
print(f"VISION DONE {done} ok {errs} err. expertos por tokens VISUALES: min={nv.min()} max={nv.max()} mean={nv.mean():.0f} de {E}",flush=True)