use relative imports and update changes involving local disk storage

README.md

@@ -67,15 +67,14 @@ import torch
 
 from huggingface_mae import MAEModel
 
-huggingface_openphenom_model_dir = "."
-# huggingface_modelpath = "recursionpharma/OpenPhenom"
+# huggingface_openphenom_model_dir = "."
+huggingface_modelpath = "recursionpharma/OpenPhenom"
 
 
 @pytest.fixture
 def huggingface_model():
-    # Make sure you have the model/config downloaded from https://huggingface.co/recursionpharma/OpenPhenom to this directory
-    # huggingface-cli download recursionpharma/OpenPhenom --local-dir=.
-    huggingface_model = MAEModel.from_pretrained(huggingface_openphenom_model_dir)
+    # This step downloads the model to a local cache, takes a bit to run
+    huggingface_model = MAEModel.from_pretrained(huggingface_modelpath)
     huggingface_model.eval()
     return huggingface_model
 

huggingface_mae.py

@@ -4,12 +4,13 @@ import torch
 import torch.nn as nn
 
 from transformers import PretrainedConfig, PreTrainedModel
+from transformers.utils import cached_file
 
-from loss import FourierLoss
-from normalizer import Normalizer
-from mae_modules import CAMAEDecoder, MAEDecoder, MAEEncoder
-from mae_utils import flatten_images
-from vit import (
+from .loss import FourierLoss
+from .normalizer import Normalizer
+from .mae_modules import CAMAEDecoder, MAEDecoder, MAEEncoder
+from .mae_utils import flatten_images
+from .vit import (
     generate_2d_sincos_pos_embeddings,
     sincos_positional_encoding_vit,
     vit_small_patch16_256,
@@ -285,8 +286,8 @@ class MAEModel(PreTrainedModel):
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
         filename = kwargs.pop("filename", "model.safetensors")
 
-        modelpath = f"{pretrained_model_name_or_path}/{filename}"
         config = MAEConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        modelpath = cached_file(pretrained_model_name_or_path, filename=filename)
         state_dict = torch.load(modelpath, map_location="cpu")
         model = cls(config)
         model.load_state_dict(state_dict["state_dict"])

mae_modules.py

@@ -7,8 +7,8 @@ import torch.nn as nn
 from timm.models.helpers import checkpoint_seq
 from timm.models.vision_transformer import Block, Mlp, VisionTransformer
 
-from masking import transformer_random_masking
-from vit import channel_agnostic_vit
+from .masking import transformer_random_masking
+from .vit import channel_agnostic_vit
 
 # If interested in training new MAEs, combine an encoder and decoder into a new module, and you should
 # leverage the flattening and unflattening utilities as needed from mae_utils.py.

test_huggingface_mae.py

@@ -1,17 +1,16 @@
 import pytest
 import torch
 
-from huggingface_mae import MAEModel
+# huggingface_openphenom_model_dir = "."
+huggingface_modelpath = "recursionpharma/OpenPhenom"
 
-huggingface_openphenom_model_dir = "."
-# huggingface_modelpath = "recursionpharma/test-pb-model"
+from .huggingface_mae import MAEModel
 
 
 @pytest.fixture
 def huggingface_model():
-    # Make sure you have the model/config downloaded from https://huggingface.co/recursionpharma/test-pb-model to this directory
-    # huggingface-cli download recursionpharma/test-pb-model --local-dir=.
-    huggingface_model = MAEModel.from_pretrained(huggingface_openphenom_model_dir)
+    # This step downloads the model to a local cache, takes a bit to run
+    huggingface_model = MAEModel.from_pretrained(huggingface_modelpath)
     huggingface_model.eval()
     return huggingface_model