# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""Load the pretrained MoGe-2 checkpoint into the torch reference model.

The official checkpoint ``Ruicheng/moge-2-vitl-normal`` ships a single
``model.pt`` (~1.3 GB) holding both the architecture config (``model_config``)
and the weights (``model``). The repo ships no modeling code, so the
architecture is provided by the vendored upstream ``moge`` package.

``load_moge2`` returns an eval-mode ``MoGeModel``. Checkpoint resolution order
(``resolve_checkpoint``):

1. an explicit ``ckpt_path`` argument;
2. ``<weights_dir>/model.pt`` where ``weights_dir`` is the argument, else
   ``$MOGE_WEIGHTS_DIR`` (the tt-model local override for offline/host use);
3. the legacy dev layout ``$TT_MOGE_WEIGHTS/model.pt`` (default ``./weights``,
   what ``scripts/download_weights.sh`` populates) when that file exists;
4. ``huggingface_hub.hf_hub_download`` from ``repo_id`` (argument, else
   ``$HF_MODEL`` -- exported by tt-model serve -- else the default repo) at
   ``revision`` (argument, else ``$TT_WEIGHTS_REVISION``, else the default
   branch). tt-model pre-downloads the pinned revision into the HF cache, so
   inside the container this is a cache hit and no ``refs/main`` lookup happens.

No environment variable is read at import time.
"""
from __future__ import annotations

import os
from typing import Optional

# Ensure the vendored `moge` package is importable.
from . import _HERE  # noqa: F401  (side effect: sys.path insert)
from moge.model.v2 import MoGeModel

_HF_REPO = "Ruicheng/moge-2-vitl-normal"
_CKPT_FILENAME = "model.pt"


def default_repo_id() -> str:
    """The weights repo the loader falls back to (``$HF_MODEL`` or the upstream id)."""
    return os.environ.get("HF_MODEL") or _HF_REPO


def default_revision() -> Optional[str]:
    """The pinned weights revision, if any (``$TT_WEIGHTS_REVISION``)."""
    return os.environ.get("TT_WEIGHTS_REVISION") or None


def resolve_checkpoint(
    ckpt_path: Optional[str] = None,
    *,
    repo_id: Optional[str] = None,
    revision: Optional[str] = None,
    weights_dir: Optional[str] = None,
) -> str:
    """Return a local path to ``model.pt`` (see the module docstring for the order)."""
    if ckpt_path is not None:
        return ckpt_path

    weights_dir = weights_dir or os.environ.get("MOGE_WEIGHTS_DIR") or None
    if weights_dir:
        local = os.path.join(weights_dir, _CKPT_FILENAME)
        if not os.path.isfile(local):
            raise FileNotFoundError(f"MOGE_WEIGHTS_DIR={weights_dir!r} does not contain {_CKPT_FILENAME}")
        return local

    legacy = os.path.join(os.environ.get("TT_MOGE_WEIGHTS", "./weights"), _CKPT_FILENAME)
    if os.path.isfile(legacy):
        return legacy

    from huggingface_hub import hf_hub_download

    return hf_hub_download(
        repo_id=repo_id or default_repo_id(),
        filename=_CKPT_FILENAME,
        revision=revision or default_revision(),
    )


def load_moge2(ckpt_path: Optional[str] = None, verbose: bool = True, **resolve_kwargs) -> MoGeModel:
    """Load the checkpoint into an eval-mode ``MoGeModel`` (fp32, CPU)."""
    path = resolve_checkpoint(ckpt_path, **resolve_kwargs)
    if verbose:
        print(f"MoGe-2: loading checkpoint {path}", flush=True)
    model = MoGeModel.from_pretrained(path).eval()
    return model