# examples/mae_export.py """Export real MAE reconstructions with strict weights and author-forward parity. No network is used. Supply the official checkpoint and the pinned author files models_mae.py and util_pos_embed.py in --official-dir. The official source is executed after substituting its unavailable timm imports with native primitives. Author masking/encoder/decoder/loss/positions remain unchanged; reference attention explicitly computes softmax(QK^T), independently of the SDPA path. """ import argparse import ast import hashlib import json from pathlib import Path import platform import sys import numpy as np from PIL import Image, __version__ as pillow_version import torch from torch.nn import functional as F from mae_training import Block, MAE, PatchEmbed REVISION = 'efb2a8062c206524e35e47d04501ed4f544c0ae8' WEIGHTS_URL = 'https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_large.pth' WEIGHTS_HASH = 'd0072183ec21f3b3e50874a957254ea4d03c2d7255d0ad47aebb85a868a6bdb6' MEAN = [.485, .456, .406] STD = [.229, .224, .225] def digest(path): """Hashes a file without buffering the full checkpoint.""" with Path(path).open('rb') as stream: return hashlib.file_digest(stream, 'sha256').hexdigest() class ReferenceBlock(Block): """Explicit reference attention and MLP under the author's forward loop.""" def forward(self, tokens): normalized = self.norm1(tokens) b, n, c = normalized.shape heads = self.attn.heads qkv = self.attn.qkv(normalized).reshape(b, n, 3, heads, c // heads) q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) attention = ((q * ((c // heads) ** -.5)) @ k.transpose(-2, -1)).softmax(-1) tokens = tokens + self.attn.proj((attention @ v).transpose(1, 2).reshape(b, n, c)) return tokens + self.mlp.fc2(F.gelu(self.mlp.fc1(self.norm2(tokens)))) def author_model(directory): """Loads author orchestration, replacing only timm and helper imports.""" directory = Path(directory) expected_hashes = { 'models_mae.py': '973c01f31fdf15febb5976baecd639b83a5cc011161af852e54ca68c6c2bafb4', 'util_pos_embed.py': 'a9e02a9a5e2587a42476bd827055b8f56ab3eaefef5e096d2fd8cadae5ddd099', } for name, expected in expected_hashes.items(): if digest(directory / name) != expected: raise ValueError(f'Official source checksum mismatch: {name}') positions = {} # NumPy 2 removed np.float; builtin float has the same float64 semantics. position_source = (directory / 'util_pos_embed.py').read_text().replace( 'dtype=np.float)', 'dtype=float)') exec(compile(position_source, str(directory / 'util_pos_embed.py'), 'exec'), positions) source = ast.parse((directory / 'models_mae.py').read_text()) source.body = [node for node in source.body if not ( isinstance(node, ast.ImportFrom) and node.module in ( 'timm.models.vision_transformer', 'util.pos_embed'))] namespace = {'PatchEmbed': PatchEmbed, 'Block': ReferenceBlock, 'get_2d_sincos_pos_embed': positions['get_2d_sincos_pos_embed']} exec(compile(source, str(directory / 'models_mae.py'), 'exec'), namespace) return namespace['mae_vit_large_patch16'](norm_pix_loss=False) def save_image(tensor, destination): """De-normalizes, clamps, and truncates to uint8 as in the author demo.""" array = tensor[0].permute(1, 2, 0).cpu().numpy() pixels = np.clip((array * np.array(STD) + np.array(MEAN)) * 255, 0, 255) Image.fromarray(pixels.astype(np.uint8)).save(destination) return {'file': destination.name, 'sha256': digest(destination)} def export(checkpoint, official_dir, photos_dir, output): """Exports two licensed photos, recording independently checked predictions.""" torch.set_num_threads(2) torch.manual_seed(7) if digest(checkpoint) != WEIGHTS_HASH: raise ValueError('Checkpoint SHA-256 does not match the official visualization release.') output = Path(output) output.mkdir(parents=True, exist_ok=True) state = torch.load(checkpoint, map_location='cpu', weights_only=True, mmap=True)['model'] model = MAE(image_size=224, patch_size=16, width=1024, depth=24, heads=16, decoder_width=512, decoder_depth=8, decoder_heads=16, norm_pix_loss=False) model.load_state_dict(state, strict=True) model.eval() reference = author_model(official_dir) reference.load_state_dict(state, strict=True) reference.eval() manifest = json.loads((Path(photos_dir) / 'predictions.json').read_text()) result = {'provenance': { 'architecture': 'ViT-L/16: encoder 24 x 1024, 16 heads; decoder 8 x 512, 16 heads', 'officialRevision': REVISION, 'weightsUrl': WEIGHTS_URL, 'compatibility': 'np.float dtype alias replaced by builtin float in author position helper', 'weightsSha256': WEIGHTS_HASH, 'strictTensorCount': len(state), 'checkpointLicense': 'https://github.com/facebookresearch/mae/blob/' + REVISION + '/LICENSE', 'normPixLoss': False, 'seed': 7, 'threads': 2, 'python': platform.python_version(), 'torch': torch.__version__, 'numpy': np.__version__, 'pillow': pillow_version, 'programSha256': digest(Path(__file__).with_name('mae_training.py')), 'exportSha256': digest(__file__), 'officialSourceSha256': {name: digest(Path(official_dir) / name) for name in ['models_mae.py', 'util_pos_embed.py']}, 'command': 'python ' + ' '.join(sys.argv), 'parity': 'Pinned author forward/positions, native PatchEmbed and Block storage; ' 'independent explicit softmax attention versus native SDPA; atol=0.0001, rtol=0.0001', }, 'preprocessing': {'size': [224, 224], 'resize': 'Pillow RGB bicubic, direct square resize; no crop', 'scale': 'float64 RGB / 255, subtract mean/divide std, then float32', 'mean': MEAN, 'std': STD}, 'photos': []} for photo in manifest['photos']: source = Path(photos_dir) / photo['originalImage'] if digest(source) != photo['source']['sha256']: raise ValueError('Original photo SHA-256 mismatch.') array = np.asarray(Image.open(source).convert('RGB').resize( (224, 224), Image.Resampling.BICUBIC)) / 255. normalized = (array - np.array(MEAN)) / np.array(STD) images = torch.from_numpy(normalized).permute(2, 0, 1)[None].float() with torch.inference_mode(): torch.manual_seed(7) loss, prediction, mask = model(images) torch.manual_seed(7) expected_loss, expected, expected_mask = reference(images) torch.testing.assert_close(prediction, expected, atol=1e-4, rtol=1e-4) torch.testing.assert_close(loss, expected_loss, atol=1e-5, rtol=1e-5) torch.testing.assert_close(mask, expected_mask, atol=0, rtol=0) pixel_mask = model.unpatchify(mask[..., None].expand(-1, -1, 768)) versions = {'original': images, 'masked': images * (1 - pixel_mask), 'prediction': model.unpatchify(prediction), 'composite': model.reconstruct(images, prediction, mask)} artifacts = {name: save_image(value, output / f'{photo["id"]}-{name}.png') for name, value in versions.items()} result['photos'].append({ 'id': photo['id'], 'title': photo['title'], 'source': photo['source'], 'artifacts': artifacts, 'maskedLoss': loss.item(), 'referenceLoss': expected_loss.item(), 'referenceMaxAbsError': (prediction - expected).abs().max().item(), 'mask': mask[0].int().tolist(), 'maskRatio': .75, 'inputShape': list(images.shape), 'predictionShape': list(prediction.shape), 'encoderTokensIncludingCls': 50, 'decoderTokensIncludingCls': 197}) (output / 'reconstruction.json').write_text(json.dumps(result, indent=2) + '\n') print(json.dumps({'photos': len(result['photos']), 'strictTensorCount': len(state), 'maxAbsErrors': [p['referenceMaxAbsError'] for p in result['photos']]}, indent=2)) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--checkpoint', required=True) parser.add_argument('--official-dir', required=True) parser.add_argument('--photos-dir', default='public/examples/vit-demo') parser.add_argument('--output', default='public/examples/mae-demo') args = parser.parse_args() export(args.checkpoint, args.official_dir, args.photos_dir, args.output) if __name__ == '__main__': main()