File size: 1,179 Bytes
b5731a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """Upstream PyTorch codec (CPU) loader. Needs torch, omegaconf, pyyaml and NCM_UPSTREAM=<HazyResearch/ncm checkout>."""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from paths import UPSTREAM
sys.path.insert(0, UPSTREAM) # checkout of github.com/HazyResearch/ncm (not vendored); provides models.vqvae
import torch, yaml
from omegaconf import OmegaConf
def load_state(path):
sd = torch.load(path, map_location="cpu", weights_only=True)
for k in ("state_dict", "model", "vqvae", "tokenizer"):
if isinstance(sd, dict) and k in sd and isinstance(sd[k], dict) and not torch.is_tensor(sd[k]): sd = sd[k]
return sd
def load_ref(d):
from models.vqvae import VQVAE
cfg = OmegaConf.load(f"{d}/config.yaml")
kw = OmegaConf.to_container(cfg.tokenizer.vqvae, resolve=True)
kw.pop("checkpoint_path", None)
kw["vector_quantizer_config"] = cfg.tokenizer.vqvae.vector_quantizer_config
m = VQVAE(**kw)
sd = load_state(f"{d}/tokenizer.pt")
sd = {k.removeprefix("module.").removeprefix("_orig_mod."): v for k, v in sd.items()}
print("load:", m.load_state_dict(sd, strict=False))
return m.eval(), cfg
|