recursionpharma · LeonardoTredese · Feb 3, 2025
diff --git a/generate_reconstructions.ipynb b/generate_reconstructions.ipynb
@@ -26,9 +26,9 @@
     "import numpy as np\n",
     "from PIL import Image\n",
     "from matplotlib import pyplot as plt\n",
-    "from huggingface_mae import MAEModel\n",
+    "from maes_microscopy.huggingface_mae import MAEModel\n",
     "from skimage.transform import resize\n",
-    "from mae_utils import unflatten_tokens"
+    "from maes_microscopy.mae_utils import unflatten_tokens"
    ]
   },
   {

diff --git a/maes_microscopy/__init__.py b/maes_microscopy/__init__.py
diff --git a/huggingface_mae.py → maes_microscopy/huggingface_mae.py b/huggingface_mae.py → maes_microscopy/huggingface_mae.py
@@ -5,11 +5,11 @@
 
 from transformers import PretrainedConfig, PreTrainedModel
 
-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 maes_microscopy.loss import FourierLoss
+from maes_microscopy.normalizer import Normalizer
+from maes_microscopy.mae_modules import CAMAEDecoder, MAEDecoder, MAEEncoder
+from maes_microscopy.mae_utils import flatten_images
+from maes_microscopy.vit import (
     generate_2d_sincos_pos_embeddings,
     sincos_positional_encoding_vit,
     vit_small_patch16_256,

diff --git a/loss.py → maes_microscopy/loss.py b/loss.py → maes_microscopy/loss.py
diff --git a/mae_modules.py → maes_microscopy/mae_modules.py b/mae_modules.py → maes_microscopy/mae_modules.py
@@ -7,8 +7,8 @@
 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 maes_microscopy.masking import transformer_random_masking
+from maes_microscopy.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.

diff --git a/mae_utils.py → maes_microscopy/mae_utils.py b/mae_utils.py → maes_microscopy/mae_utils.py
diff --git a/masking.py → maes_microscopy/masking.py b/masking.py → maes_microscopy/masking.py
diff --git a/normalizer.py → maes_microscopy/normalizer.py b/normalizer.py → maes_microscopy/normalizer.py
diff --git a/test_huggingface_mae.py → maes_microscopy/test_huggingface_mae.py b/test_huggingface_mae.py → maes_microscopy/test_huggingface_mae.py
@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from huggingface_mae import MAEModel
+from maes_microscopy.huggingface_mae import MAEModel
 
 huggingface_openphenom_model_dir = "."
 # huggingface_modelpath = "recursionpharma/OpenPhenom"

diff --git a/vit.py → maes_microscopy/vit.py b/vit.py → maes_microscopy/vit.py
diff --git a/vit_encoder.py → maes_microscopy/vit_encoder.py b/vit_encoder.py → maes_microscopy/vit_encoder.py
diff --git a/pyproject.toml b/pyproject.toml
@@ -30,4 +30,4 @@ dependencies = [
 ]
 
 [tool.setuptools]
-py-modules = []
+packages=["maes_microscopy"]