rc-foundry 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

foundry/trainers/fabric.py

@@ -36,6 +36,11 @@ from foundry.utils.weights import (
     freeze_parameters_with_config,
     load_weights_with_policies,
 )
+from foundry.utils.xpu import (
+    SingleXPUStrategy,
+    XPUAccelerator,
+    XPUMixedPrecision,
+)
 
 ranked_logger = RankedLogger(__name__, rank_zero_only=True)
 
@@ -122,8 +127,19 @@ class FabricTrainer(ABC):
             (3) Fabric Loggers (https://lightning.ai/docs/fabric/2.4.0/api/loggers.html)
             (4) Efficient Gradient Accumulation (https://lightning.ai/docs/fabric/2.4.0/advanced/gradient_accumulation.html)
         """
-        # Use custom DDP strategy only for multi-device, non-interactive environments
-        if (
+        # Handle XPU accelerator
+        is_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
+        if accelerator == "xpu" or (
+            accelerator == "auto" and is_xpu and not torch.cuda.is_available()
+        ):
+            accelerator = XPUAccelerator()
+            precision_plugin = None
+            if precision in ("16-mixed", "bf16-mixed"):
+                precision_plugin = XPUMixedPrecision(precision=precision)
+                precision = None  # Handled by plugin
+            strategy = SingleXPUStrategy(precision_plugin=precision_plugin)
+            ranked_logger.info("Using Intel XPU with SingleXPUStrategy")
+        elif (
             strategy == "ddp"
             and not is_interactive_environment()
             and not (num_nodes == 1 and devices_per_node == 1)

foundry/utils/components.py

@@ -5,11 +5,11 @@ from typing import List
 import numpy as np
 from atomworks.ml.encoding_definitions import AF3SequenceEncoding
 from biotite.structure import AtomArray
-from rfd3.constants import (
-    TIP_BY_RESTYPE,
-)
 
 from foundry.common import exists
+from foundry.constants import (
+    TIP_BY_RESTYPE,
+)
 from foundry.utils.ddp import RankedLogger
 
 global_logger = RankedLogger(__name__, rank_zero_only=False)

foundry/utils/ddp.py

@@ -20,7 +20,7 @@ def is_rank_zero() -> bool:
 
 
 def set_accelerator_based_on_availability(cfg: dict | DictConfig):
-    """Set training accelerator to CPU if no GPUs are available.
+    """Set training accelerator based on available hardware.
 
     Args:
         cfg: Hydra object with trainer settings "accelerator", "devices_per_node", and "num_nodes".
@@ -28,22 +28,25 @@ def set_accelerator_based_on_availability(cfg: dict | DictConfig):
     Returns:
         None; modifies the input `cfg` object in place.
     """
-    if not torch.cuda.is_available():
+    assert "trainer" in cfg, "Configuration object must have a 'trainer' key."
+    for key in ["accelerator", "devices_per_node", "num_nodes"]:
+        assert (
+            key in cfg.trainer
+        ), f"Configuration object must have a 'trainer.{key}' key."
+
+    if torch.cuda.is_available():
+        cfg.trainer.accelerator = "gpu"
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        logger.info("Intel XPU detected - using XPU accelerator")
+        cfg.trainer.accelerator = "xpu"
+    else:
         logger.error(
-            "No GPUs available - Setting accelerator to 'cpu'. Are you sure you are using the correct configs?"
+            "No GPUs/XPUs available - Setting accelerator to 'cpu'. Are you sure you are using the correct configs?"
         )
-        assert "trainer" in cfg, "Configuration object must have a 'trainer' key."
-        for key in ["accelerator", "devices_per_node", "num_nodes"]:
-            assert (
-                key in cfg.trainer
-            ), f"Configuration object must have a 'trainer.{key}' key."
-
-        # Override accelerator settings
         cfg.trainer.accelerator = "cpu"
         cfg.trainer.devices_per_node = 1
         cfg.trainer.num_nodes = 1
-    else:
-        cfg.trainer.accelerator = "gpu"
+
     return cfg
 
 

foundry/utils/xpu/__init__.py

@@ -0,0 +1,27 @@
+"""XPU utilities for Intel GPU support.
+
+XPU support in PyTorch is now native (torch.xpu.is_available()), but Lightning Fabric
+requires custom Accelerator, Strategy, and Precision plugins for proper XPU handling.
+
+These components are used directly (not via registry) when XPU is detected:
+- XPUAccelerator: Custom accelerator for XPU devices
+- SingleXPUStrategy: Strategy for single-device XPU training/inference
+- XPUMixedPrecision: Precision plugin with proper XPU autocast support
+
+Usage:
+    from foundry.utils.xpu import XPUAccelerator, SingleXPUStrategy, XPUMixedPrecision
+
+    # Check availability
+    if XPUAccelerator.is_available():
+        strategy = SingleXPUStrategy(precision_plugin=XPUMixedPrecision("bf16-mixed"))
+
+Note:
+    The FabricTrainer automatically uses these components when XPU is detected.
+    You typically don't need to use them directly unless customizing behavior.
+"""
+
+from .single_xpu_strategy import SingleXPUStrategy
+from .xpu_accelerator import XPUAccelerator
+from .xpu_precision import XPUMixedPrecision
+
+__all__ = ["SingleXPUStrategy", "XPUAccelerator", "XPUMixedPrecision"]

foundry/utils/xpu/single_xpu_strategy.py

@@ -0,0 +1,47 @@
+"""Lightning Fabric strategy for single XPU device.
+
+https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html
+"""
+
+import torch
+from lightning.fabric.plugins import CheckpointIO
+from lightning.fabric.plugins.precision import Precision
+from lightning.fabric.strategies import SingleDeviceStrategy
+from lightning.fabric.utilities.types import _DEVICE
+
+
+class SingleXPUStrategy(SingleDeviceStrategy):
+    """Strategy for training/inference on a single Intel XPU device.
+
+    This strategy extends SingleDeviceStrategy to properly handle XPU devices.
+    """
+
+    strategy_name = "xpu_single"
+
+    def __init__(
+        self,
+        device: _DEVICE = "xpu:0",
+        checkpoint_io: CheckpointIO | None = None,
+        precision_plugin: Precision | None = None,
+    ) -> None:
+        """Initialize the single XPU strategy.
+
+        Args:
+            device: The XPU device to use. Defaults to "xpu:0".
+            checkpoint_io: Plugin for checkpoint I/O.
+            precision_plugin: Plugin for precision handling (set via _precision property).
+
+        Raises:
+            RuntimeError: If XPU devices are not available.
+        """
+        if not (hasattr(torch, "xpu") and torch.xpu.is_available()):
+            msg = "`SingleXPUStrategy` requires XPU devices to run"
+            raise RuntimeError(msg)
+
+        super().__init__(
+            device=device,
+            checkpoint_io=checkpoint_io,
+        )
+        # Precision is handled via the _precision property in newer Lightning versions
+        if precision_plugin is not None:
+            self._precision = precision_plugin

foundry/utils/xpu/xpu_accelerator.py

@@ -0,0 +1,91 @@
+"""XPU Accelerator for Intel XPU devices."""
+
+from typing import Any
+
+import torch
+from lightning.fabric.accelerators import Accelerator
+
+
+class XPUAccelerator(Accelerator):
+    """Accelerator for Intel XPU devices.
+
+    This accelerator enables training and inference on Intel GPUs using
+    PyTorch's native XPU support (torch.xpu).
+    """
+
+    @property
+    def name(self) -> str:
+        """Return the name of this accelerator."""
+        return "xpu"
+
+    @staticmethod
+    def setup_device(device: torch.device) -> None:
+        """Set up the specified XPU device.
+
+        Args:
+            device: The torch device to set up.
+
+        Raises:
+            RuntimeError: If device is not an XPU device.
+        """
+        if device.type != "xpu":
+            msg = f"Device should be xpu, got {device} instead"
+            raise RuntimeError(msg)
+
+        torch.xpu.set_device(device)
+
+    @staticmethod
+    def parse_devices(devices: str | list | torch.device) -> list:
+        """Parse devices specification for XPU training.
+
+        Args:
+            devices: Device specification (int, list of ints, or string).
+
+        Returns:
+            List of device indices.
+        """
+        if isinstance(devices, list):
+            return devices
+        return [devices]
+
+    @staticmethod
+    def get_parallel_devices(devices: list) -> list[torch.device]:
+        """Generate a list of parallel XPU devices.
+
+        Args:
+            devices: List of device indices.
+
+        Returns:
+            List of torch.device objects for XPU.
+        """
+        return [torch.device("xpu", idx) for idx in devices]
+
+    @staticmethod
+    def auto_device_count() -> int:
+        """Return the number of available XPU devices."""
+        return torch.xpu.device_count()
+
+    @staticmethod
+    def is_available() -> bool:
+        """Check if XPU is available."""
+        return hasattr(torch, "xpu") and torch.xpu.is_available()
+
+    @staticmethod
+    def get_device_stats(device: str | torch.device) -> dict[str, Any]:
+        """Return XPU device statistics.
+
+        Currently returns an empty dict as XPU stats API may vary.
+
+        Args:
+            device: The device to get stats for.
+
+        Returns:
+            Dictionary of device statistics.
+        """
+        del device  # Unused
+        return {}
+
+    def teardown(self) -> None:
+        """Clean up XPU accelerator resources."""
+        # Empty implementation required by base class
+        pass

foundry/utils/xpu/xpu_precision.py

@@ -0,0 +1,72 @@
+"""XPU Precision Plugin for Lightning Fabric."""
+
+from contextlib import contextmanager
+from typing import Any, Generator, Literal
+
+import torch
+from lightning.fabric.plugins.precision import MixedPrecision
+from torch import Tensor
+
+
+class XPUMixedPrecision(MixedPrecision):
+    """Mixed precision plugin for Intel XPU devices.
+
+    This overrides the default MixedPrecision plugin to use 'xpu' as the
+    device type for torch.autocast instead of 'cuda'.
+    """
+
+    def __init__(
+        self,
+        precision: Literal["16-mixed", "bf16-mixed"] = "bf16-mixed",
+    ) -> None:
+        """Initialize XPU mixed precision.
+
+        Args:
+            precision: The precision mode. "16-mixed" uses float16,
+                "bf16-mixed" uses bfloat16. Defaults to "bf16-mixed".
+
+        Raises:
+            ValueError: If precision is not "16-mixed" or "bf16-mixed".
+        """
+        # Determine dtype from precision string
+        if precision == "16-mixed":
+            dtype = torch.float16
+        elif precision == "bf16-mixed":
+            dtype = torch.bfloat16
+        else:
+            msg = f"Invalid precision: {precision}. Must be '16-mixed' or 'bf16-mixed'"
+            raise ValueError(msg)
+
+        # Initialize with xpu device type
+        super().__init__(precision=precision, device="xpu")
+        self._desired_input_dtype = dtype
+
+    @contextmanager
+    def forward_context(self) -> Generator[None, None, None]:
+        """Context manager for forward pass with XPU autocast."""
+        with torch.autocast(device_type="xpu", dtype=self._desired_input_dtype):
+            yield
+
+    def convert_input(self, data: Any) -> Any:
+        """Convert input data to the appropriate precision.
+
+        Args:
+            data: Input data to convert.
+
+        Returns:
+            Converted data.
+        """
+        return self._convert_fp_tensor(data)
+
+    def _convert_fp_tensor(self, data: Any) -> Any:
+        """Convert floating point tensors to the desired dtype.
+
+        Args:
+            data: Data to convert.
+
+        Returns:
+            Converted data if it's a floating point tensor, otherwise unchanged.
+        """
+        if isinstance(data, Tensor) and data.is_floating_point():
+            return data.to(self._desired_input_dtype)
+        return data

foundry/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.9'
-__version_tuple__ = version_tuple = (0, 1, 9)
+__version__ = version = '0.1.11'
+__version_tuple__ = version_tuple = (0, 1, 11)
 
 __commit_id__ = commit_id = None

mpnn/inference_engines/mpnn.py

@@ -67,11 +67,15 @@ class MPNNInferenceEngine:
             else checkpoint_path
         )
 
-        # Determine the device.
+        # Determine the device (supports XPU, CUDA, and CPU).
         if device is not None:
             self.device = torch.device(device)
+        elif torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        elif hasattr(torch, "xpu") and torch.xpu.is_available():
+            self.device = torch.device("xpu")
         else:
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.device = torch.device("cpu")
 
         # Set up allowed model types.
         self.allowed_model_types = {"protein_mpnn", "ligand_mpnn"}

{rc_foundry-0.1.9.dist-info → rc_foundry-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rc-foundry
-Version: 0.1.9
+Version: 0.1.11
 Summary: Shared utilities and training infrastructure for biomolecular structure prediction models.
 Author-email: Institute for Protein Design <contact@ipd.uw.edu>
 License: BSD 3-Clause License
@@ -108,6 +108,16 @@ All models within Foundry rely on [AtomWorks](https://github.com/RosettaCommons/
 pip install "rc-foundry[all]"
 ```
 
+**Intel XPU Installation**
+
+For Intel XPU devices, install PyTorch with XPU support first, then install Foundry.
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/xpu
+pip install "rc-foundry[all]"
+```
+> [!NOTE]
+> Use `pip` (not `uv`) for XPU installs since UV re-resolves dependencies and may replace your XPU torch with the standard PyPI version.
+
 **Downloading weights** Models can be downloaded to a target folder with:
 ```
 foundry install base-models --checkpoint-dir <path/to/ckpt/dir>

{rc_foundry-0.1.9.dist-info → rc_foundry-0.1.11.dist-info}/RECORD

@@ -1,7 +1,7 @@
 foundry/__init__.py,sha256=H8S1nl5v6YeW8ggn1jKy4GdtH7c-FGS-j7CqUCAEnAU,1926
 foundry/common.py,sha256=Aur8mH-CNmcUqSsw7VgaCQSW5sH1Bqf8Da91jzxPV1Y,3035
 foundry/constants.py,sha256=0n1wBKCvNuw3QaQehSbmsHYkIdaGn3tLeRFItBrdeHY,913
-foundry/version.py,sha256=ib8ckvf-NNDfacXd8unW0p5cf-gl57XyQvjoEMc_pvc,704
+foundry/version.py,sha256=0-Ruc52ECccw_8Ef0d7jMkzrb8fkobUkZLqGGvcm1ik,706
 foundry/callbacks/__init__.py,sha256=VsRT1e4sqlJHPcTCsfupMEx82Iz-LoOAGPpwvf_OJeE,126
 foundry/callbacks/callback.py,sha256=xZBo_suP4bLrP6gl5uJPbaXm00DXigePa6dMeDxucgg,3890
 foundry/callbacks/health_logging.py,sha256=tEtkByOlaAA7nnelxb7PbM9_dcIgOsdbxCdQY3K5pMc,16664
@@ -18,14 +18,14 @@ foundry/model/layers/blocks.py,sha256=ihbbP_1fOlrkrcrQSk9thCrNWjK8mtxD3WxcBng9Ht
 foundry/testing/__init__.py,sha256=BnrU7fZ4l0Dm1vrGcNPQYTAw83PW4DGYz7TGhGqgrfQ,223
 foundry/testing/fixtures.py,sha256=j27a8CAonygjlWsUjZ-95M5MF4Rjp9nw7JskqiZlweI,486
 foundry/testing/pytest_hooks.py,sha256=5Ebw1GXYO2XqS9Jvpzty7g3gCXIdXu16jqg53XcuUx4,450
-foundry/trainers/fabric.py,sha256=cjaTHbGuJEQwaGBvIAXD_il4bHtY-crsTY14Xn77uXA,40401
+foundry/trainers/fabric.py,sha256=qZVO3C48-Won57TbxMO9qb8VOh8JPNG8kqXSW6nh-YU,41057
 foundry/training/EMA.py,sha256=3OWA9Pz7XuDr-SRxbz24tZf55DmhSa2fKy9r5v2IXqA,2651
 foundry/training/checkpoint.py,sha256=mUiObg-qcF3tvMfVu77sD9m3yVRp71czv07ccliU7qQ,1791
 foundry/training/schedulers.py,sha256=StmXegPfIdLAv31FreCTrDh9dsOvNUfzG4YGa61Y4oE,3647
 foundry/utils/alignment.py,sha256=2anqy0mn9zeFEiVWS_EG7zHiyPk1C_gbUu-SRvQ5mAM,2502
-foundry/utils/components.py,sha256=Piw2TfQF26uuxC3hXG3iv_4rgud1lVO-cv6N-p05EDY,15200
+foundry/utils/components.py,sha256=gCRjmZXqltOqU6tV_jISX6T-BX2z9zzr-H0Xgo7oojw,15203
 foundry/utils/datasets.py,sha256=pLBxVezm-TSrYuC5gFnJZdGnNWV7aPH2QiWIVE2hkdQ,16629
-foundry/utils/ddp.py,sha256=202_7qqm4ihPjpB5Q9NhUjDl4u22pu5JvY0ui0UkRUQ,3970
+foundry/utils/ddp.py,sha256=gJY3w0pwelr6UoxOkHsttgcZsTN97kfHFtfGN2VKLM8,4075
 foundry/utils/instantiators.py,sha256=oGCp6hrmY-QPPPEjxKxe5uVFL125fH1RaLxjMKWCD_8,2169
 foundry/utils/logging.py,sha256=ywV75MBlQsothV0IBvqoAQTNg6pjo2-Cib7Uo080nzQ,9312
 foundry/utils/rigid.py,sha256=_Z1pmitb6xgxyguLj_TukKscUBJjQsU4bsBD24GVS84,44444
@@ -33,13 +33,17 @@ foundry/utils/rotation_augmentation.py,sha256=7q1WEX2iJ0i7-2aV-M97nEaEdpqexDTaZn
 foundry/utils/squashfs.py,sha256=QlcwuJyVe-QVfIOS7o1QfLhaCQPNzzox7ln4n8dcYEg,5234
 foundry/utils/torch.py,sha256=OLsqoxw4CTXbGzWUHernLUT7uQjLu0tVPtD8h8747DI,11211
 foundry/utils/weights.py,sha256=btz4S02xff2vgiq4xMfiXuhK1ERafqQPtmimo1DmoWY,10381
+foundry/utils/xpu/__init__.py,sha256=i8_uiVj5b0-_Z4bX4RZmaE01e0NHbo-eTNXMYF1M0mM,1114
+foundry/utils/xpu/single_xpu_strategy.py,sha256=28iEo-6e0gbpApmFxZdfhLO-RZb2QlOf9XKlRdTeXXc,1592
+foundry/utils/xpu/xpu_accelerator.py,sha256=LPVXvgBKY3OE1ghyhZvotP6aWNVZXMup1Swi1v6U5GY,2467
+foundry/utils/xpu/xpu_precision.py,sha256=zuFfLkBIDLbF-omJ6mTiGoVU0UgXg84A2j3ZjYyoAOc,2285
 foundry_cli/__init__.py,sha256=0BxY2RUKJLaMXUGgypPCwlTskTEFdVnkhTR4C4ft2Kw,52
 foundry_cli/download_checkpoints.py,sha256=CxU9dKBa1vAkVd450tfH5aZAlQIUTrHsDGTbmxzd_JQ,8922
 mpnn/__init__.py,sha256=hgQcXFaCbAxFrhydVAy0xj8yC7UJF-GCCFhqD0sZ7I4,57
 mpnn/inference.py,sha256=wPtGR325eVRVeesXoWtBK6b_-VcU8BZae5IfQN3-mvA,1669
 mpnn/train.py,sha256=9eQGBd3rdNF5Zr2w8oUgETbqxBavNBajtA6Vbc5zESE,10239
 mpnn/collate/feature_collator.py,sha256=LpzAFWo1VMa06dJLmfUWZsKe4xvLZjHbx4RICg2lgbQ,10510
-mpnn/inference_engines/mpnn.py,sha256=PmDEsIFipdk2fY57FA-vCp4evoU83DVVuUVmlViUtWk,21725
+mpnn/inference_engines/mpnn.py,sha256=Wg7FMI5KKdq14qQzKs6TkX-rvUi1fB6lSnJMdwWuKKw,21912
 mpnn/loss/nll_loss.py,sha256=KmdNe-BCzGYtijjappzBArQcT1gHVlJnKdY1PYQ4mhU,5947
 mpnn/metrics/nll.py,sha256=T6oMeUOEeHZzOMTH8NHFtsH9vUwLAsHQDPszzj4YKXI,15299
 mpnn/metrics/sequence_recovery.py,sha256=YDw_LmH-a3ajBYWK0mucJEQvw0_VEyxvrBN7da4vX8Q,19034
@@ -64,7 +68,7 @@ rf3/_version.py,sha256=fCfpbI5aeA6yHqjo3tK78-l2dPGxhp-AyKSoCXp34Nc,739
 rf3/alignment.py,sha256=BvvwMqQGCVxV20xIsTighD1kXMadXXL2SkckLjTerx0,2102
 rf3/chemical.py,sha256=VECnRPgVm-icXbZeUG4svcENzdUiIupP6dhka_8zCrg,26572
 rf3/cli.py,sha256=jxjq8u77J8itIK4aNTfIpnMsNgg8brW1A3NfVjlgE0s,2743
-rf3/inference.py,sha256=yUjSqCjIqsGErUezq42-8K9570Opm_2eYa-bNnwltwA,2494
+rf3/inference.py,sha256=R_bzJe2K6Ek9uhrkZ97IiVaGKhBaObWc5SWh6MVDLLg,2530
 rf3/kinematics.py,sha256=V3yjalPupu1X2FEp7l3XZR-qzLKrhWLZyECk6RgIkcs,10901
 rf3/scoring.py,sha256=dTllswE-6Fgli2eLiNzLFc2Rhz4ouDT4WL-sVbvLTGU,41541
 rf3/train.py,sha256=V4nqCC_1JKLI3WQ-nErNa8sqFpvb1mFhXSe6ZPpEheM,7945
@@ -119,18 +123,18 @@ rfd3/__init__.py,sha256=2Wto2IsUIj2lGag9m_gqgdCwBNl5p21-Xnr7W_RpU3c,348
 rfd3/callbacks.py,sha256=Zjt8RiaYWquoKOwRmC_wCUbRbov-V4zd2_73zjhgDHE,2783
 rfd3/cli.py,sha256=ka3K5H117fzDYIDXFpOpJV21w_XBrHYJZdFE0thsGBI,1644
 rfd3/constants.py,sha256=wLvDzrThpOrK8T3wGFNQeGrhAXOJQze8l3v_7pjIdMM,13141
-rfd3/engine.py,sha256=viXzVGYkHPyzv5d0Ifg8zFZ6Wqg-U4I8Y4ldQLPe9x4,21536
+rfd3/engine.py,sha256=cj7BIP35HMycWV3rBp9n9_ik5t7eNMK1wkVfJ58xFmI,21903
 rfd3/run_inference.py,sha256=HfRMQ30_SAHfc-VFzBV52F-aLaNdG6PW8VkdMyB__wE,1264
 rfd3/train.py,sha256=rHswffIUhOae3_iYyvAiQ3jALoFuzrcRUgMlbJLinlI,7947
 rfd3/inference/datasets.py,sha256=9VLbzl7dpG8mk_pjs0R5C2wFYUoRIgXXoZcS9IohSy0,6510
-rfd3/inference/input_parsing.py,sha256=pocqnnhE3-szeBbCL9gy9E3kZJSP_CGHXd6FFRxfv0c,46563
+rfd3/inference/input_parsing.py,sha256=qUMsMyeh-kW757xX10zi07kKg4RS0wM8jSoIxu7UsBI,46794
 rfd3/inference/legacy_input_parsing.py,sha256=G2XxkrjdIpL6i1YY7xEmkFitVv__Pc45ow6IKKPHw64,28855
 rfd3/inference/parsing.py,sha256=ktAMUuZE3Pe4bKAjjV3zjqcEDmGlMZ-cotIUhJsEQQA,5402
-rfd3/inference/symmetry/atom_array.py,sha256=j8yhtZAjRNm4d06KyS4gk2XCquHPCwR0k9WiNmxz7WA,12941
+rfd3/inference/symmetry/atom_array.py,sha256=HfFagFUB5yB-Y4IfUM5nuVGWHC5AEkyHqt0JcIqTQ_E,10922
 rfd3/inference/symmetry/checks.py,sha256=ZWpC1JjrAjXY__xDt8EFYb5WUdSF0kobZyxxmacFU7U,10076
 rfd3/inference/symmetry/contigs.py,sha256=6OvbZ2dJg-a0mvvKAC0VkzUH5HpUDxOJvkByIst_roU,2127
-rfd3/inference/symmetry/frames.py,sha256=kwog5jU_wgv6ACcUER2iU5qz-mdfdOe0cpi2OTEDKMU,18894
-rfd3/inference/symmetry/symmetry_utils.py,sha256=CGUzMI5CKVIcNi5_l2-YRu-ExroZX54GndxLb5P7RtY,14680
+rfd3/inference/symmetry/frames.py,sha256=aEwkmlUsYexERX9hu09JMhisC8QTpHPVhfITbL80-EE,10819
+rfd3/inference/symmetry/symmetry_utils.py,sha256=ibYeOT-4z_Pxr5i8cKbNAAqjRf4YfKOQCpSkuZQqG0I,14554
 rfd3/metrics/design_metrics.py,sha256=O1RqZdjQPNlAWYRg6UJTERYg_gUI1_hVleKsm9xbWBY,16836
 rfd3/metrics/hbonds_hbplus_metrics.py,sha256=Sewy9KzmrA1OnfkasN-fmWrQ9IRx9G7Yyhe2ua0mk28,11518
 rfd3/metrics/hbonds_metrics.py,sha256=SIR4BnDhYdpVSqwXXRYpQ_tB-M0_fVyugGl08WivCmE,15257
@@ -140,7 +144,7 @@ rfd3/metrics/sidechain_metrics.py,sha256=EGZuFuWQ0cCe83EVPAf4eysN8vP9ifNjfnmE0o5
 rfd3/model/RFD3.py,sha256=95aKzye-XzuDyLGgost-Wsfu8eT635zHIRky-pNoHSA,3569
 rfd3/model/RFD3_diffusion_module.py,sha256=BPjKGyQpbnqdzii3gXMKLhhijNqV8Xh4bSosmfDBt8w,12094
 rfd3/model/cfg_utils.py,sha256=XPBLyoB_bQRLmdrJ1Z0hCjcVvgUMGIPuw4rxTlHjB_s,2575
-rfd3/model/inference_sampler.py,sha256=5k_UIkJCL6QO4BEqAxkR9GKxNAOz-NRq5Jh51Wm5MU8,25152
+rfd3/model/inference_sampler.py,sha256=zyA4pU2MIr2PPTuQm9k70UXUm3PRp644va6XgmB3aJ8,25210
 rfd3/model/layers/attention.py,sha256=XuNA7WyFlRfLnAgky1PtGvXFCnDGv7GeEcXz8hodTBo,19472
 rfd3/model/layers/block_utils.py,sha256=oN0aD-vZiH4JbIFs2CzDmb2B74GNPKzdFurmGd-dirE,21244
 rfd3/model/layers/blocks.py,sha256=MOjJ53THxM2MMM27Ap7xiIXRCdI_SHzqKzLLQVX6FEc,24888
@@ -161,21 +165,21 @@ rfd3/transforms/conditioning_utils.py,sha256=9Pn9AFbih2FCzp5OOM9y7z6KH7HPxVibxTr
 rfd3/transforms/design_transforms.py,sha256=ePvnLsuKUOsE4LLcmF0bbkx1vf2AiD-35rzF4zUEcEE,30944
 rfd3/transforms/dna_crop.py,sha256=JeOsG0tXghJvgzEimfzBvlFN_lVd9TrvjnC929Abz5A,18214
 rfd3/transforms/hbonds.py,sha256=ijfJapFlhsh3JktpDoT3VFqKTTg6ynrqMlD7dU2xFsA,16415
-rfd3/transforms/hbonds_hbplus.py,sha256=xyDP-CyVl2OsUY90HsrPoKw1VycBXUrq00WfrX8HJVM,8364
+rfd3/transforms/hbonds_hbplus.py,sha256=i-57ZGZjofxWvygDidwWeKrkrALz5LcwJCCcw6rCOfQ,8654
 rfd3/transforms/ncaa_transforms.py,sha256=Lz4L8OGuOOG53sKJHcLSdV7WPQ3YzOzwd5tJG4CHqP0,4983
 rfd3/transforms/pipelines.py,sha256=FGH-XH3taTWQ6k1zpDO_d-097EQdXmL6uqXZXw4HIMs,22086
 rfd3/transforms/ppi_transforms.py,sha256=7rXyf-tn2TLz6ybYR_YVDtSDG7hOgqhYY4shNviA_Sw,23493
 rfd3/transforms/rasa.py,sha256=a4IPFvVMMxldoGLyJQiSlGg7IyUkcBASbRZLWmguAKk,4156
-rfd3/transforms/symmetry.py,sha256=9I9gzAZkk5vMUJm7x8XCDSHtNPYYLAHt4meXxOczGT0,2970
+rfd3/transforms/symmetry.py,sha256=GSnMF7oAnUxPozfafsRuHEv0yKXW0BpLTI6wsKGZrbc,2658
 rfd3/transforms/training_conditions.py,sha256=UXiUPjDwrNKM95tRe0eXrMeRN8XlTPc_MXUvo6UpePo,19510
 rfd3/transforms/util_transforms.py,sha256=2AcLkzx-73ZFgcWD1cIHv7NyniRPI4_zThHK8azyQaY,18119
 rfd3/transforms/virtual_atoms.py,sha256=UpmxzPPd5FaJigcRoxgLSHHrLLOqsCvZ5PPZfQSGqII,12547
-rfd3/utils/inference.py,sha256=Yf3aAUk_YZi58uIJr5Y2wfVnQ-2bh3S5GHLBPzCRjUs,26448
+rfd3/utils/inference.py,sha256=A0yGOsO3cXiTqwLYP0uTGTB9MR6vsmnYMEiroYgvnDE,26578
 rfd3/utils/io.py,sha256=wbdjUTQkDc3RCSM7gdogA-XOKR68HeQ-cfvyN4pP90w,9849
 rfd3/utils/vizualize.py,sha256=HPlczrA3zkOuxV5X05eOvy_Oga9e3cPnFUXOEP4RR_g,11046
-rf3/configs/inference.yaml,sha256=JmEZdkAnbnOrX79lGS5xrYYho9aBFfVxfUp-8KjJV5I,309
-rf3/configs/train.yaml,sha256=4KW2fKc9a_gjg8yMoQfOpfkC-nJ5mdQEfoikOKxbnKc,1573
-rf3/configs/validate.yaml,sha256=3LkhXyneEuuH-ueFH9FyYY5cCDi1_0KoHNwEceuQPwI,1581
+rf3/configs/inference.yaml,sha256=uvg-ppF7-ANMeecb60Vodx6GNlEx0W3HPiGgiajUZ5Q,375
+rf3/configs/train.yaml,sha256=aQWrr2RXsFo97CYSFOyTMu7C8F6Nz7ymZU3BuX1gkec,1639
+rf3/configs/validate.yaml,sha256=6FUaD-kLbu0jfUSgFIXfAuWhHtVH-4w_swGJSf2AQKc,1647
 rf3/configs/callbacks/default.yaml,sha256=MkxOj7dMXh4jJRIE62gLjoOYecGuZLWiJrr780_nubA,89
 rf3/configs/callbacks/dump_validation_structures.yaml,sha256=EYEibR25v7KZJtadvCFLFMEPTf0FvKFNW2ocx4wm57A,259
 rf3/configs/callbacks/metrics_logging.yaml,sha256=MNm4OpvOHxvDJofVUA27NVaiDkp1NzqOYCzl6l_7ceo,432
@@ -224,6 +228,7 @@ rf3/configs/trainer/cpu.yaml,sha256=J1WbK2SQ_VMoEOOqv0XTg0FYKPaec1TNySvWklMaE4k,
 rf3/configs/trainer/ddp.yaml,sha256=uClrdTzEMNxgq4IQhMgm8okC16wUS2I5i3rKnm5SktU,65
 rf3/configs/trainer/rf3.yaml,sha256=WBjnaYofmEV7OfJisvFGN6y_UJGNEskAdmbV9oO4ICI,500
 rf3/configs/trainer/rf3_with_confidence.yaml,sha256=meDTw0S2nTcuAj5tefGwJteo2715x6kO143FMd3db14,346
+rf3/configs/trainer/xpu.yaml,sha256=GlUmQcnWtZOcp7wHip5DikKWfOVZTwrTX_HU05A4kc8,73
 rf3/configs/trainer/loss/structure_prediction.yaml,sha256=XPH2RcIo6m1BDrIoBvd2xBgJ0c4KEOnbWc5oEcBKVAM,164
 rf3/configs/trainer/loss/structure_prediction_with_confidence.yaml,sha256=OZxtrQIqYGDFlhp3-1fFVZ2Ek6QRMP74ImsrNOJcelo,52
 rf3/configs/trainer/loss/losses/confidence_loss.yaml,sha256=c9xI7Q32Ddl5R_GpoWq4RS7AxiJf4hsnGeMZxWs9mBU,501
@@ -231,10 +236,10 @@ rf3/configs/trainer/loss/losses/diffusion_loss.yaml,sha256=ktBy3bUjslklPuRhGWuHb
 rf3/configs/trainer/loss/losses/distogram_loss.yaml,sha256=-dZWPeQ1mWwi9pUtB7xJW57qKgUNMQ4vao1PXo07RC0,56
 rf3/configs/trainer/metrics/structure_prediction.yaml,sha256=xWp2DqoqlofbgRLMNi0LKuAaDSCbW1tAEuEk83fmc0w,439
 rfd3/configs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rfd3/configs/dev.yaml,sha256=Y3PiHaCRIdRXf2Ea2aUDJk6hFxyNEljJf-O2VKQRNH0,151
-rfd3/configs/inference.yaml,sha256=tUo5-G1-rsCU8RLclSP0ZKu5-O6JC5ofgI6bHb5kwmw,119
-rfd3/configs/train.yaml,sha256=SbP5136VjJH_i2fc_4b0U0ZOqz3jeby1YORXSQtiFk0,450
-rfd3/configs/validate.yaml,sha256=TSToLqzuP8hecj0K6TAfCywtpTZI1u1-THsX0jxCG1s,537
+rfd3/configs/dev.yaml,sha256=ulil1Am_romDNUOGeQY3HLDRyPMiVKHOZrh59H4uAMA,176
+rfd3/configs/inference.yaml,sha256=zzrhijUPplk--wjKk5d_Ht3Nnv-qLav5DQcfHcmLSNY,144
+rfd3/configs/train.yaml,sha256=BJ3owsuFDVWiVgZdmA0VaQkPtMS9ws8FEyvBZeI4YQ8,476
+rfd3/configs/validate.yaml,sha256=E78w-K6bE0Kl27V-wz9BuR9QyZVXTupxjT0UpYIA7Ao,562
 rfd3/configs/callbacks/design_callbacks.yaml,sha256=JWgE1-v_spzUy7JH3_6dHct_-oX4DevozRH-pM5Ds2k,196
 rfd3/configs/callbacks/metrics_logging.yaml,sha256=pZPePSYGKEV560e3WatuLvJiHlz1CIGFOaOWoRmBh8g,694
 rfd3/configs/callbacks/train_logging.yaml,sha256=Z25GVLTHo1HvQUjnBdayaozmNww8UhAh9DczStYaZig,1050
@@ -301,11 +306,12 @@ rfd3/configs/paths/data/default.yaml,sha256=jfs1dbbcOqHja4_6lXheyRg4t0YExqVn2w0r
 rfd3/configs/trainer/cpu.yaml,sha256=rJf5LHf6x5fN5EKg8mFEn2SwfGW5dV1JdYaHqWMfpXc,74
 rfd3/configs/trainer/ddp.yaml,sha256=uClrdTzEMNxgq4IQhMgm8okC16wUS2I5i3rKnm5SktU,65
 rfd3/configs/trainer/rfd3_base.yaml,sha256=R3lZxdyjUirjlLU31qWlnZgHaz4GcWTGGIz4fUl7AyM,1016
+rfd3/configs/trainer/xpu.yaml,sha256=GlUmQcnWtZOcp7wHip5DikKWfOVZTwrTX_HU05A4kc8,73
 rfd3/configs/trainer/loss/losses/diffusion_loss.yaml,sha256=FE4FCEfurE0ekwZ4YfS6wCvPSNqxClwg_kc73cPql5Y,323
 rfd3/configs/trainer/loss/losses/sequence_loss.yaml,sha256=kezbQcqwAZ0VKQPUBr2MsNr9DcDL3ENIP1i-j7h-6Co,64
 rfd3/configs/trainer/metrics/design_metrics.yaml,sha256=xVDpClhHqSHvsf-8StL26z51Vn-iuWMDG9KMB-kqOI0,719
-rc_foundry-0.1.9.dist-info/METADATA,sha256=SBHGrkr8RsLTCTOhJnbusBs8G7C8HxJUPajCbmU6OyE,11502
-rc_foundry-0.1.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-rc_foundry-0.1.9.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
-rc_foundry-0.1.9.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
-rc_foundry-0.1.9.dist-info/RECORD,,
+rc_foundry-0.1.11.dist-info/METADATA,sha256=lK8MYr2naArWrcTOenHGERoQTliCfzPDgDave3bpWJI,11873
+rc_foundry-0.1.11.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rc_foundry-0.1.11.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
+rc_foundry-0.1.11.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
+rc_foundry-0.1.11.dist-info/RECORD,,

rf3/configs/inference.yaml

@@ -2,6 +2,11 @@
 # ^ The "package" determines where the content of the config is placed in the output config
 # For more information about overriding configs, see: https://hydra.cc/docs/advanced/overriding_packages/#overriding-packages-using-the-defaults-list
 
+hydra:
+  searchpath:
+    - pkg://rf3.configs
+    - pkg://configs
+
 defaults:
   - inference_engine: rf3
   - _self_

rf3/configs/train.yaml

@@ -2,6 +2,11 @@
 # ^ The "package" determines where the content of the config is placed in the output config
 # For more information about overriding configs, see: https://hydra.cc/docs/advanced/overriding_packages/#overriding-packages-using-the-defaults-list
 
+hydra:
+  searchpath:
+    - pkg://rf3.configs
+    - pkg://configs
+
 # NOTE: order of defaults determines the order in which configs override each other (higher up items are overridden by lower items)
 defaults:
   - callbacks: default

rf3/configs/trainer/xpu.yaml

@@ -0,0 +1,6 @@
+strategy: xpu_single
+
+accelerator: xpu
+devices_per_node: 1
+num_nodes: 1
+

rf3/configs/validate.yaml

@@ -2,6 +2,11 @@
 # ^ The "package" determines where the content of the config is placed in the output config
 # For more information about overriding configs, see: https://hydra.cc/docs/advanced/overriding_packages/#overriding-packages-using-the-defaults-list
 
+hydra:
+  searchpath:
+    - pkg://rf3.configs
+    - pkg://configs
+
 # NOTE: order of defaults determines the order in which configs override each other (higher up items are overridden by lower items)
 defaults:
   - callbacks: default

rf3/inference.py

@@ -12,7 +12,10 @@ from foundry.utils.logging import suppress_warnings
 
 # Setup root dir and environment variables (more info: https://github.com/ashleve/rootutils)
 # NOTE: Sets the `PROJECT_ROOT` environment variable to the root directory of the project (where `.project-root` is located)
-rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+try:
+    rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+except Exception:
+    pass
 
 load_dotenv(override=True)
 

rfd3/configs/dev.yaml

@@ -2,6 +2,7 @@
 # Inference engine config for development
 hydra:
   searchpath:
+    - pkg://rfd3.configs
     - pkg://configs
 
 defaults:

rfd3/configs/inference.yaml

@@ -2,6 +2,7 @@
 
 hydra:
   searchpath:
+    - pkg://rfd3.configs
     - pkg://configs
 
 defaults:

rfd3/configs/train.yaml

@@ -1,6 +1,7 @@
 # @package _global_
 hydra:
   searchpath:
+    - pkg://rfd3.configs
     - pkg://configs
 
 defaults:
@@ -25,4 +26,4 @@ ckpt_path: null
 
 # Placeholders
 name: aa_design
-tags: [aa_design]
+tags: [aa_design]

rfd3/configs/trainer/xpu.yaml

@@ -0,0 +1,6 @@
+strategy: xpu_single
+
+accelerator: xpu
+devices_per_node: 1
+num_nodes: 1
+