"""Load KORMo-Web-10B with transformers 5.

KORMo-Web-10B is a LLaVA-OneVision model (SigLIP vision tower + projector) with the KORMo-10B text model.
transformers does not ship the KORMo text model, so two small steps are needed:

    import kormo_load
    kormo_load.register(MODEL_DIR)                 # 1. register the KORMo text model from this repo
    model = AutoModelForImageTextToText.from_pretrained(MODEL_DIR, dtype="bfloat16")
    kormo_load.patch_vision(model)                 # 2. use the post-layernorm SigLIP features, as in training
"""
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from transformers.dynamic_module_utils import get_class_from_dynamic_module


def register(model_dir):
    """Register `kormo` (from `_configuration_kormo.py` / `_modeling_kormo.py` in `model_dir`) with the Auto classes."""
    cfg = get_class_from_dynamic_module("_configuration_kormo.KORMoConfig", model_dir)
    base = get_class_from_dynamic_module("_modeling_kormo.KORMoModel", model_dir)
    lm = get_class_from_dynamic_module("_modeling_kormo.KORMoForCausalLM", model_dir)
    for fn in (lambda: AutoConfig.register("kormo", cfg), lambda: AutoModel.register(cfg, base), lambda: AutoModelForCausalLM.register(cfg, lm)):
        try:
            fn()
        except ValueError:  # already registered
            pass


def patch_vision(model):
    """The model was trained on SigLIP features *after* post_layernorm (transformers 4.57 behaviour).
    In transformers 5, hidden_states[-1] is the value *before* post_layernorm, so replace it with last_hidden_state.
    Call once after from_pretrained; without it the model still runs but misreads screenshots."""
    vt = model.model.vision_tower if hasattr(model, "model") else model.vision_tower

    def hook(mod, args, out):
        if getattr(out, "hidden_states", None):
            out.hidden_states = tuple(out.hidden_states[:-1]) + (out.last_hidden_state,)
        return out

    vt.register_forward_hook(hook)
    return model