| |
| |
|
|
| """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 |
|
|
| |
| from . import _HERE |
| 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 |
|
|
