File size: 1,203 Bytes
aa45001 | 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 29 30 31 32 33 34 35 36 | from pathlib import Path
import os
import sys
ROOT = Path(__file__).resolve().parents[1]
BACKBONE = ROOT/'assets/esm2_t33_650M_UR50D'
PREDICTOR = ROOT/'checkpoints/metalatte_predictor.ckpt'
DIFFUSION = ROOT/'checkpoints/metalorian_diffusion.pt'
TRAIN = ROOT/'data/metalatte/train'
VALIDATION = ROOT/'data/metalatte/validation'
os.environ['USE_TF'] = '0'
os.environ.setdefault('HF_HUB_OFFLINE', '1')
os.environ.setdefault('WANDB_MODE', 'disabled')
os.environ.setdefault('MPLCONFIGDIR', str(ROOT/'.cache/matplotlib'))
sys.dont_write_bytecode = True
_configured = False
def configure():
global _configured
if _configured:
return
import torch
original = torch.load
def load(file, *args, **kwargs):
candidate = file if isinstance(file, (str, Path)) else getattr(file, 'name', getattr(file, 'path', None))
if isinstance(candidate, (str, Path)) and Path(candidate).resolve() in {PREDICTOR.resolve(), DIFFUSION.resolve()}:
file = str(Path(candidate).resolve())
kwargs.setdefault('weights_only', False)
kwargs.setdefault('mmap', True)
return original(file, *args, **kwargs)
torch.load = load
_configured = True
|