| 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 |
|
|
