Merge branch 'lucidrains:main' into main

.github/workflows/test.yml

@@ -26,7 +26,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu
+        python -m pip install torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cpu
         python -m pip install scipy
     - name: Test with pytest
       run: |

meshgpt_pytorch/data.py

@@ -14,8 +14,8 @@
 from einops import rearrange, reduce
 from torch import nn, Tensor
 
-from beartype import beartype
 from beartype.typing import Tuple, List, Callable, Dict
+from meshgpt_pytorch.typing import typecheck
 
 from torchtyping import TensorType
 
@@ -39,7 +39,7 @@ def identity(t):
 # you would decorate your Dataset class with this
 # and then change your `data_kwargs = ["text_embeds", "vertices", "faces"]`
 
-@beartype
+@typecheck
 def cache_text_embeds_for_dataset(
     embed_texts_fn: Callable[[List[str]], Tensor],
     max_text_len: int,
@@ -158,7 +158,7 @@ def __getitem__(self, idx):
 # you would decorate your Dataset class with this function
 # and then change your `data_kwargs = ["vertices", "faces", "face_edges"]`
 
-@beartype
+@typecheck
 def cache_face_edges_for_dataset(
     max_edges_len: int,
     cache_path: str = './face_edges_cache',
@@ -262,7 +262,7 @@ def __getitem__(self, idx):
 # dataset
 
 class DatasetFromTransforms(Dataset):
-    @beartype
+    @typecheck
     def __init__(
         self,
         folder: str,

meshgpt_pytorch/trainer.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from pathlib import Path
 from functools import partial
 from packaging import version

meshgpt_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.4.0'
+__version__ = '1.4.1'

setup.py

@@ -24,10 +24,10 @@
     'huggingface_hub>=0.21.4',
     'beartype',
     "huggingface_hub>=0.21.4",
-    'classifier-free-guidance-pytorch>=0.6.2',
-    'einops>=0.7.0',
-    'einx[torch]>=0.1.3',
-    'ema-pytorch>=0.5.0',
+    'classifier-free-guidance-pytorch>=0.6.7',
+    'einops>=0.8.0',
+    'einx[torch]>=0.3.0',
+    'ema-pytorch>=0.5.1',
     'environs',
     'gateloop-transformer>=0.2.2',
     'jaxtyping',