keep the base alphafold3 un-mixined

Author: lucidrains

alphafold3_pytorch/__init__.py

@@ -29,7 +29,8 @@
     InputFeatureEmbedder,
     ConfidenceHead,
     DistogramHead,
-    Alphafold3
+    Alphafold3,
+    Alphafold3WithHubMixin
 )
 
 from alphafold3_pytorch.inputs import (
@@ -79,6 +80,7 @@
     ConfidenceHead,
     DistogramHead,
     Alphafold3,
+    Alphafold3WithHubMixin,
     Alphafold3Config,
     AtomInput,
     Trainer,

alphafold3_pytorch/alphafold3.py

@@ -2830,7 +2830,7 @@ class LossBreakdown(NamedTuple):
     diffusion_bond: Float['']
     diffusion_smooth_lddt: Float['']
 
-class Alphafold3(Module, PyTorchModelHubMixin):
+class Alphafold3(Module):
     """ Algorithm 1 """
 
     @save_args_and_kwargs
@@ -3089,46 +3089,6 @@ def __init__(
 
         self.register_buffer('zero', torch.tensor(0.), persistent = False)
 
-    @classmethod
-    def _from_pretrained(
-        cls,
-        *,
-        model_id: str,
-        revision: str | None,
-        cache_dir: str | Path | None,
-        force_download: bool,
-        proxies: Dict | None,
-        resume_download: bool,
-        local_files_only: bool,
-        token: str | bool | None,
-        map_location: str = 'cpu',
-        strict: bool = False,
-        **model_kwargs,
-    ):
-        model_filename = "alphafold3.bin"
-        model_file = Path(model_id) / model_filename
-
-        if not model_file.exists():
-            model_file = hf_hub_download(
-                repo_id = model_id,
-                filename = model_filename,
-                revision = revision,
-                cache_dir = cache_dir,
-                force_download = force_download,
-                proxies = proxies,
-                resume_download = resume_download,
-                token = token,
-                local_files_only = local_files_only,
-            )
-
-        model = cls.init_and_load(
-            model_file,
-            strict = strict,
-            map_location = map_location
-        )
-
-        return model
-
     @property
     def device(self):
         return self.zero.device
@@ -3623,3 +3583,46 @@ def forward(
         )
 
         return loss, loss_breakdown
+
+# an alphafold3 that can download pretrained weights from huggingface
+
+class Alphafold3WithHubMixin(Alphafold3, PyTorchModelHubMixin):
+    @classmethod
+    def _from_pretrained(
+        cls,
+        *,
+        model_id: str,
+        revision: str | None,
+        cache_dir: str | Path | None,
+        force_download: bool,
+        proxies: Dict | None,
+        resume_download: bool,
+        local_files_only: bool,
+        token: str | bool | None,
+        map_location: str = 'cpu',
+        strict: bool = False,
+        **model_kwargs,
+    ):
+        model_filename = "alphafold3.bin"
+        model_file = Path(model_id) / model_filename
+
+        if not model_file.exists():
+            model_file = hf_hub_download(
+                repo_id = model_id,
+                filename = model_filename,
+                revision = revision,
+                cache_dir = cache_dir,
+                force_download = force_download,
+                proxies = proxies,
+                resume_download = resume_download,
+                token = token,
+                local_files_only = local_files_only,
+            )
+
+        model = cls.init_and_load(
+            model_file,
+            strict = strict,
+            map_location = map_location
+        )
+
+        return model

alphafold3_pytorch/trainer.py

@@ -3,7 +3,7 @@
 from functools import wraps, partial
 from pathlib import Path
 
-from alphafold3_pytorch.alphafold3 import Alphafold3
+from alphafold3_pytorch.alphafold3 import Alphafold3, Alphafold3WithHubMixin
 from alphafold3_pytorch.attention import pad_at_dim
 
 from typing import TypedDict, List, Callable
@@ -195,7 +195,7 @@ class Trainer:
     @typecheck
     def __init__(
         self,
-        model: Alphafold3,
+        model: Alphafold3 | Alphafold3WithHubMixin,
         *,
         dataset: Dataset,
         num_train_steps: int,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.50"
+version = "0.1.51"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }