load_magvit.py

"""Lightweight loader for the MagViT-v2 image (de)tokenizer.

The `showlab/magvitv2` checkpoint is a plain safetensors state dict with an
empty config (the model takes no constructor args). We instantiate the vendored
`MAGVITv2` module and load the weights directly, avoiding the heavyweight
diffusers `ModelMixin.from_pretrained` path (which relies on older diffusers
internal APIs).
"""
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from magvit.modeling_magvitv2 import MAGVITv2


def load_magvit(repo_id: str = "showlab/magvitv2") -> MAGVITv2:
    model = MAGVITv2()
    weights_path = hf_hub_download(repo_id, "pytorch_model.safetensors")
    state_dict = load_file(weights_path)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if missing:
        print(f"[magvit] missing keys: {len(missing)} (e.g. {missing[:5]})", flush=True)
    if unexpected:
        print(f"[magvit] unexpected keys: {len(unexpected)} (e.g. {unexpected[:5]})", flush=True)
    model.eval()
    model.requires_grad_(False)
    return model