force einops 0.6.1 or greater and call allow_ops_in_compiled_graph

dalle2_pytorch/__init__.py

@@ -1,9 +1,5 @@
-import torch
-from packaging import version
-
-if version.parse(torch.__version__) >= version.parse('2.0.0'):
-    from einops._torch_specific import allow_ops_in_compiled_graph
-    allow_ops_in_compiled_graph()
+from einops._torch_specific import allow_ops_in_compiled_graph
+allow_ops_in_compiled_graph()
 
 from dalle2_pytorch.version import __version__
 from dalle2_pytorch.dalle2_pytorch import DALLE2, DiffusionPriorNetwork, DiffusionPrior, Unet, Decoder

dalle2_pytorch/train_configs.py

@@ -1,6 +1,6 @@
 import json
 from torchvision import transforms as T
-from pydantic import BaseModel, validator, model_validator
+from pydantic import BaseModel, validator, root_validator
 from typing import List, Optional, Union, Tuple, Dict, Any, TypeVar
 
 from x_clip import CLIP as XCLIP
@@ -38,9 +38,9 @@ class TrainSplitConfig(BaseModel):
     val: float = 0.15
     test: float = 0.1
 
-    @model_validator(mode = 'after')
-    def validate_all(self):
-        actual_sum = sum([*dict(self).values()])
+    @root_validator
+    def validate_all(cls, fields):
+        actual_sum = sum([*fields.values()])
         if actual_sum != 1.:
             raise ValueError(f'{fields.keys()} must sum to 1.0. Found: {actual_sum}')
         return fields
@@ -59,7 +59,6 @@ class TrackerLogConfig(BaseModel):
         kwargs = self.dict()
         return create_logger(self.log_type, data_path, **kwargs)
 
-
 class TrackerLoadConfig(BaseModel):
     load_from: Optional[str] = None
     only_auto_resume: bool = False  # Only attempt to load if the logger is auto-resuming
@@ -349,7 +348,7 @@ class TrainDecoderConfig(BaseModel):
             config = json.load(f)
         return cls(**config)
     
-    @model_validator(mode = 'after')
+    @root_validator
     def check_has_embeddings(cls, values):
         # Makes sure that enough information is provided to get the embeddings specified for training
         data_config, decoder_config = values.get('data'), values.get('decoder')

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.15.0'
+__version__ = '1.14.1'

setup.py

@@ -36,7 +36,7 @@ setup(
     'numpy',
     'packaging',
     'pillow',
-    'pydantic>=2',
+    'pydantic',
     'pytorch-warmup',
     'resize-right>=0.0.2',
     'rotary-embedding-torch',