diff --git a/kalign/README b/kalign/README
new file mode 100644
index 00000000..13b9617b
--- /dev/null
+++ b/kalign/README
@@ -0,0 +1,17 @@
+Binaries for kalign2.  Mac and Linux binaries come from Anaconda, and Windows comes from NIH.
+
+Mac
+
+https://anaconda.org/channels/bioconda/packages/kalign2/files
+bioconda/kalign2/2.04/osx-64/kalign2-2.04-h18d8cf3_8.tar.bz2  Dec 12, 2024
+bioconda/kalign2/2.04/osx-arm64/kalign2-2.04-hbdacb55_8.tar.bz2  Dec 10, 2024
+
+Linux
+
+https://anaconda.org/channels/bioconda/packages/kalign2/files
+bioconda/kalign2/2.04/linux-64/kalign2-2.04-h7b50bb2_8.tar.bz2  Dec 13, 2024
+bioconda/kalign2/2.04/linux-aarch64/kalign2-2.04-h661a7b0_8.tar.bz2 Dec 10, 2024
+
+Windows
+
+https://ftp.ncbi.nlm.nih.gov/toolbox/gbench/kalign
\ No newline at end of file
diff --git a/kalign/kalign_darwin_arm64 b/kalign/kalign_darwin_arm64
new file mode 100755
index 00000000..e78de373
Binary files /dev/null and b/kalign/kalign_darwin_arm64 differ
diff --git a/kalign/kalign_darwin_x86_64 b/kalign/kalign_darwin_x86_64
new file mode 100755
index 00000000..c1c2a6bf
Binary files /dev/null and b/kalign/kalign_darwin_x86_64 differ
diff --git a/kalign/kalign_linux_aarch64 b/kalign/kalign_linux_aarch64
new file mode 100755
index 00000000..103eae40
Binary files /dev/null and b/kalign/kalign_linux_aarch64 differ
diff --git a/kalign/kalign_linux_x86_64 b/kalign/kalign_linux_x86_64
new file mode 100755
index 00000000..8eed8f2b
Binary files /dev/null and b/kalign/kalign_linux_x86_64 differ
diff --git a/kalign/kalign_windows_x86_64.exe b/kalign/kalign_windows_x86_64.exe
new file mode 100755
index 00000000..a0d07029
Binary files /dev/null and b/kalign/kalign_windows_x86_64.exe differ
diff --git a/openfold3/core/config/config_utils.py b/openfold3/core/config/config_utils.py
index 6999678e..547c591a 100644
--- a/openfold3/core/config/config_utils.py
+++ b/openfold3/core/config/config_utils.py
@@ -51,7 +51,7 @@ def load_json(path: Path | str) -> dict[str, Any]:
 
 def _ensure_list(value: Any) -> Any:
     if not isinstance(value, list):
-        logging.info("Single value: {value} will be converted to a list")
+        logging.debug(f"Single value: {value} will be converted to a list")
         return [value]
     else:
         return value
diff --git a/openfold3/core/data/framework/data_module.py b/openfold3/core/data/framework/data_module.py
index e0b85891..051c0430 100644
--- a/openfold3/core/data/framework/data_module.py
+++ b/openfold3/core/data/framework/data_module.py
@@ -540,6 +540,44 @@ class DataModule(pl.LightningDataModule):
             )
         self.next_dataset_indices = state_dict["next_dataset_indices"]
 
+# This function needs to be at global scope to work with multiprocessing pickler on Mac.
+# Somehow this was not needed on Linux.
+class worker_init_function_with_data_seed:
+    def __init__(self, seed):
+        self.data_seed = seed
+    def __call__(self,
+        worker_id: int, rank: int | None = None
+    ) -> None:
+        """Modified default Lightning worker_init_fn with manual data seed.
+
+        This worker_init_fn enables decoupling stochastic processes in the data
+        pipeline from those in the model. Taken from Pytorch Lightning 2.4.1 source
+        code: https://github.com/Lightning-AI/pytorch-lightning/blob/f3f10d460338ca8b2901d5cd43456992131767ec/src/lightning/fabric/utilities/seed.py#L85
+
+        Args:
+            worker_id (int):
+                Worker id.
+            rank (Optional[int], optional):
+                Worker process rank. Defaults to None.
+        """
+        # implementation notes: https://github.com/pytorch/pytorch/issues/5059#issuecomment-817392562
+        global_rank = rank if rank is not None else rank_zero_only.rank
+        process_seed = self.data_seed
+        # back out the base seed so we can use all the bits
+        base_seed = process_seed - worker_id
+        seed_sequence = _generate_seed_sequence(
+            base_seed, worker_id, global_rank, count=4
+        )
+        torch.manual_seed(seed_sequence[0])  # torch takes a 64-bit seed
+        random.seed(
+            (seed_sequence[1] << 32) | seed_sequence[2]
+        )  # combine two 64-bit seeds
+        if _NUMPY_AVAILABLE:
+            import numpy as np
+
+            np.random.seed(
+                seed_sequence[3] & 0xFFFFFFFF
+            )  # numpy takes 32-bit seed only
 
 class InferenceDataModule(DataModule):
     """LightningDataModule that contains a prepare_data hook for inference."""
diff --git a/openfold3/core/data/framework/single_datasets/inference.py b/openfold3/core/data/framework/single_datasets/inference.py
index d440df19..1727ab2b 100644
--- a/openfold3/core/data/framework/single_datasets/inference.py
+++ b/openfold3/core/data/framework/single_datasets/inference.py
@@ -331,7 +331,9 @@ class InferenceDataset(Dataset):
         is_repeated_sample = bool(datapoint["repeated_sample"])
 
         try:
+            logger.info(f'Creating features for {query_id}')
             features = self.create_all_features(query)
+            logger.info(f'Finished creating features for {query_id}')
             features["query_id"] = query_id
             features["seed"] = torch.tensor([seed])
             features["repeated_sample"] = torch.tensor(
diff --git a/openfold3/core/data/pipelines/preprocessing/template.py b/openfold3/core/data/pipelines/preprocessing/template.py
index 865939f5..f965501d 100644
--- a/openfold3/core/data/pipelines/preprocessing/template.py
+++ b/openfold3/core/data/pipelines/preprocessing/template.py
@@ -15,6 +15,8 @@
 """Preprocessing pipelines for template data ran before training/evaluation."""
 
 import logging
+logger = logging.getLogger(__name__)
+
 import multiprocessing as mp
 import os
 import random
@@ -1756,6 +1758,7 @@ class TemplatePreprocessor:
 
     def __call__(self) -> None:
         # Preprocess template alignments into template cache entries
+        logger.info(f"Preprocessing templates for {len(self.inputs)} chains")
         if len(self.inputs) >= 1:
             manager = mp.Manager()
             self.seq_hash_map = manager.dict()
@@ -1783,6 +1786,8 @@ class TemplatePreprocessor:
         elif isinstance(self.input_set, InferenceQuerySet):
             self._update_inference_query_set()
 
+        logger.info("Finished preprocessing templates")
+
     def preprocess_templates(
         self,
         input_data: TemplatePreprocessorInputTrain | TemplatePreprocessorInputInference,
diff --git a/openfold3/core/data/primitives/structure/cleanup.py b/openfold3/core/data/primitives/structure/cleanup.py
index 84a69ef8..bc0bce17 100644
--- a/openfold3/core/data/primitives/structure/cleanup.py
+++ b/openfold3/core/data/primitives/structure/cleanup.py
@@ -245,7 +245,7 @@ def remove_hydrogens(atom_array: AtomArray) -> AtomArray:
     n_hydrogens_removed = atom_array_len_prev - atom_array_len_new
 
     if n_hydrogens_removed > 0:
-        logger.info(f"Removed {n_hydrogens_removed} hydrogen atoms.")
+        logger.debug(f"Removed {n_hydrogens_removed} hydrogen atoms.")
 
     return atom_array
 
diff --git a/openfold3/core/data/primitives/structure/unresolved.py b/openfold3/core/data/primitives/structure/unresolved.py
index ad183381..5147f7e4 100644
--- a/openfold3/core/data/primitives/structure/unresolved.py
+++ b/openfold3/core/data/primitives/structure/unresolved.py
@@ -450,7 +450,7 @@ def append_unresolved_segment(
     # Log what segment was added
     segment_start = segment.res_id[0]
     segment_end = segment.res_id[-1]
-    logger.info(
+    logger.debug(
         f"Added unresolved segment: chain_id={reference_atom.chain_id}, "
         f"span={segment_start}-{segment_end}, type={segment_type}"
     )
@@ -885,7 +885,7 @@ def add_unresolved_atoms_within_residue(
     # Add bonds within all the added atoms
     update_bond_list(extended_atom_array)
 
-    logger.info(
+    logger.debug(
         f"Added unresolved atoms for {len(modified_res_tuples)} residues: "
         f"{modified_res_tuples}"
     )
diff --git a/openfold3/core/data/tools/colabfold_msa_server.py b/openfold3/core/data/tools/colabfold_msa_server.py
index 907bdfd8..99c83856 100644
--- a/openfold3/core/data/tools/colabfold_msa_server.py
+++ b/openfold3/core/data/tools/colabfold_msa_server.py
@@ -328,7 +328,7 @@ def query_colabfold_msa_server(
     # Extract a3m files
     if any(not os.path.isfile(a3m_file) for a3m_file in a3m_files):
         with tarfile.open(tar_gz_file) as tar_gz:
-            tar_gz.extractall(path)
+            tar_gz.extractall(path, filter='data')
 
     # Process templates
     if use_templates:
@@ -687,9 +687,9 @@ class ColabFoldQueryRunner:
         # TODO: add template alignments fetching code here by setting use_templates=True
         # TODO: replace prints with proper logging
         if len(self.colabfold_mapper.seqs) == 0:
-            print("No protein sequences found for main MSA generation. Skipping...")
+            logger.debug("No protein sequences found for main MSA generation. Skipping...")
             return
-        print(
+        logger.info(
             f"Submitting {len(self.colabfold_mapper.seqs)} sequences to the Colabfold"
             " MSA server for main MSAs..."
         )
@@ -774,9 +774,9 @@ class ColabFoldQueryRunner:
         # Submit queries for paired MSAss
         num_complexes = len(self.colabfold_mapper.complex_id_to_complex_group)
         if num_complexes == 0:
-            print("No complexes found for paired MSA generation. Skipping...")
+            logger.debug("No complexes found for paired MSA generation. Skipping...")
             return
-        print(
+        logger.info(
             f"Submitting {num_complexes} paired MSA queries"
             " to the Colabfold MSA server..."
         )
@@ -1036,7 +1036,7 @@ def preprocess_colabfold_msas(
     # Gather MSA data
     colabfold_mapper = collect_colabfold_msa_data(inference_query_set)
     output_directory = compute_settings.msa_output_directory
-    logger.warning(f"Using output directory: {output_directory} for ColabFold MSAs.")
+    logger.info(f"Using output directory: {output_directory} for ColabFold MSAs.")
 
     # Save mappings to file
     if compute_settings.save_mappings:
diff --git a/openfold3/core/data/tools/kalign.py b/openfold3/core/data/tools/kalign.py
index 34668187..a0482257 100644
--- a/openfold3/core/data/tools/kalign.py
+++ b/openfold3/core/data/tools/kalign.py
@@ -14,13 +14,19 @@
 
 from functools import lru_cache
 
-import kalign
 
 
 @lru_cache(maxsize=512)
 def _run_kalign_cached(sequences: tuple[str, ...]) -> str:
     """Wrapper around kalign.align with caching."""
-    return kalign.align(list(sequences))
+    from sys import platform
+    if platform == 'win32' or platform == 'darwin':
+        # The Python kalign module is not available on Windows.
+        # The Python kalign module aborts on Mac due to duplicate libomp.dylib.
+        return _run_kalign_exe(sequences)
+    else:
+        import kalign
+        return kalign.align(list(sequences))
 
 
 def run_kalign(
@@ -38,3 +44,62 @@ def run_kalign(
             The aligned sequences as a list of strings.
     """
     return _run_kalign_cached(tuple(sequences))
+
+def _run_kalign_exe(
+    sequences: list[str],
+) -> list[str]:
+    """
+    Runs the kalign executable in a subprocess and returns the aligned sequences.
+    This is used on Windows because there is currently (March 2026) no PyPi kalign
+    package for Windows.
+
+    Args:
+        sequences (list[str]):
+            Sequences to be aligned. In the template pipeline,
+            the first sequence is the query, and the rest are templates
+            sequences to be realigned to it from hmmsearch.
+    Returns:
+        list:
+            The aligned sequences as a list of strings.
+    """
+    import sys
+    from os.path import dirname, join, exists
+    kalign_exe = join(dirname(sys.executable), 'kalign')
+    if not exists(kalign_exe):
+        kalign_exe = shutil.which("kalign")
+
+    if kalign_exe is None:
+        raise RuntimeError(
+            "Kalign is not available. Please install it and ensure it is in your PATH."
+        )
+
+    a3m_string = '\n'.join(f'>sequence_{i}\n{seq}' for i,seq in enumerate(sequences))
+    import subprocess
+    try:
+        result = subprocess.run(
+            [kalign_exe], input=a3m_string, capture_output=True, text=True, check=True
+        )
+
+        # The resulting MSA is stored in the variable
+        alignment_result = result.stdout
+
+        # Strip header from kalign version 3.
+        if not alignment_result.startswith('>') and '\n>' in alignment_result:
+            alignment_result = alignment_result[alignment_result.find('\n>'):]
+
+    except subprocess.CalledProcessError as e:
+        print(f"Kalign command failed:\n{e.stderr}")
+
+    # Convert fasta/a3m to list of sequences.
+    aligned_seqs = []
+    seq = ''
+    for line in alignment_result.split('\n'):
+        if line.startswith('>'):
+            if seq:
+                aligned_seqs.append(seq)
+                seq = ''
+        else:
+            seq += line
+    aligned_seqs.append(seq)
+
+    return aligned_seqs
diff --git a/openfold3/core/model/feature_embedders/input_embedders.py b/openfold3/core/model/feature_embedders/input_embedders.py
index 314213ef..b9b25a3b 100644
--- a/openfold3/core/model/feature_embedders/input_embedders.py
+++ b/openfold3/core/model/feature_embedders/input_embedders.py
@@ -124,6 +124,8 @@ class InputEmbedderAllAtom(nn.Module):
             z:
                 [*, N_token, N_token, C_z] Pair representation
         """
+        import warnings
+        warnings.filterwarnings('ignore', message='.*CUDA is not available.*')
         with torch.amp.autocast(device_type="cuda", dtype=torch.float32):
             a, _, _, _ = self.atom_attn_enc(
                 batch=batch,
diff --git a/openfold3/core/runners/writer.py b/openfold3/core/runners/writer.py
index c1628360..2a2e8790 100644
--- a/openfold3/core/runners/writer.py
+++ b/openfold3/core/runners/writer.py
@@ -177,6 +177,7 @@ class OF3OutputWriter(BasePredictionWriter):
     ):
         """Writes confidence scores to disk"""
         plddt = confidence_scores["plddt"]
+        pae = confidence_scores["pae"]
         pde = confidence_scores["pde"]
         gpde = confidence_scores["gpde"]
         aggregated_confidence_scores = {"avg_plddt": np.mean(plddt), "gpde": gpde}
@@ -193,7 +194,7 @@ class OF3OutputWriter(BasePredictionWriter):
         )
 
         # Full confidence scores
-        full_confidence_scores = {"plddt": plddt, "pde": pde}
+        full_confidence_scores = {"plddt": plddt, "pde": pde, "pae": pae}
         out_fmt = self.full_confidence_format
         out_file_full = Path(f"{output_prefix}_confidences.{out_fmt}")
 
diff --git a/openfold3/core/utils/chunk_utils.py b/openfold3/core/utils/chunk_utils.py
index c89e1d1c..8e43ccda 100644
--- a/openfold3/core/utils/chunk_utils.py
+++ b/openfold3/core/utils/chunk_utils.py
@@ -355,7 +355,7 @@ class ChunkSizeTuner:
 
     @staticmethod
     def _determine_favorable_chunk_size(fn, args, min_chunk_size, max_chunk_size):
-        logging.info("Tuning chunk size...")
+        logging.debug("Tuning chunk size...")
 
         if min_chunk_size >= max_chunk_size:
             return min_chunk_size
diff --git a/openfold3/core/utils/script_utils.py b/openfold3/core/utils/script_utils.py
index b4ae73be..f9b60b83 100644
--- a/openfold3/core/utils/script_utils.py
+++ b/openfold3/core/utils/script_utils.py
@@ -14,7 +14,6 @@
 
 import logging
 import platform
-import resource
 
 logging.basicConfig()
 logger = logging.getLogger(__file__)
@@ -36,6 +35,7 @@ def set_ulimits():
     # This is more than is needed,
     no_file_limit = 32768
 
+    import resource
     soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
 
     if no_file_limit > hard:
diff --git a/openfold3/entry_points/experiment_runner.py b/openfold3/entry_points/experiment_runner.py
index 4090b61b..9bd3679c 100644
--- a/openfold3/entry_points/experiment_runner.py
+++ b/openfold3/entry_points/experiment_runner.py
@@ -269,6 +269,7 @@ class ExperimentRunner(ABC):
                 "strategy": self.strategy,
                 "callbacks": self.callbacks,
                 "logger": self.loggers,
+                "enable_progress_bar": False,
             }
         )
 
@@ -681,8 +682,11 @@ class InferenceExperimentRunner(ExperimentRunner):
         self._log_model_config()
         logger.info(f"Loading weights from {self.ckpt_path}")
         ckpt = load_checkpoint(self.ckpt_path)
+        logger.info("Finished loading weights")
+        logger.info('Loading model state dictionary')
         state_dict, _ = get_state_dict_from_checkpoint(ckpt, init_from_ema_weights=True)
         self._warn_on_missing_version_tensor_in_load_statedict(state_dict)
+        logger.info('Finished loading model state dictionary')
 
     def run(self, inference_query_set) -> None:
         """Set up the experiment environment."""
@@ -696,12 +700,13 @@ class InferenceExperimentRunner(ExperimentRunner):
                 return
 
         self.inference_query_set = inference_query_set
-        logger.info("Beginning inference prediction")
+        logger.info("Beginning prediction inference")
         self.trainer.predict(
             model=self.lightning_module,
             datamodule=self.lightning_data_module,
             return_predictions=False,
         )
+        logger.info("Finished prediction inference")
 
     @cached_property
     def callbacks(self):
@@ -766,11 +771,11 @@ class InferenceExperimentRunner(ExperimentRunner):
     def cleanup(self):
         """Cleanup directories from colabfold MSA"""
         if self.is_rank_zero and self.log_dir.is_dir() and not os.listdir(self.log_dir):
-            print("Removing empty log directory...")
+            logger.debug("Removing empty log directory...")
             self.log_dir.rmdir()
 
         if self.use_msa_server and self.is_rank_zero:
-            print("Cleaning up MSA directories...")
+            logger.debug("Cleaning up MSA directories...")
 
             # Always remove raw directory
             # TODO: Change to use ColabFoldQueryRunner.cleanup() when
@@ -779,7 +784,7 @@ class InferenceExperimentRunner(ExperimentRunner):
                 self.experiment_config.msa_computation_settings.msa_output_directory
                 / "raw"
             )
-            self._maybe_remove_dir(raw_colabfold_msa_path)
+            #self._maybe_remove_dir(raw_colabfold_msa_path)
             if self.experiment_config.msa_computation_settings.cleanup_msa_dir:
                 msa_output_dir = (
                     self.experiment_config.msa_computation_settings.msa_output_directory
diff --git a/openfold3/entry_points/validator.py b/openfold3/entry_points/validator.py
index 767cd700..5ce8a2d1 100644
--- a/openfold3/entry_points/validator.py
+++ b/openfold3/entry_points/validator.py
@@ -111,8 +111,8 @@ class DataModuleArgs(BaseModel):
     model_config = PydanticConfigDict(extra="forbid")
     batch_size: int = 1
     data_seed: int | None = None
-    num_workers: int = 10
-    num_workers_validation: int = 4
+    num_workers: int = 0
+    num_workers_validation: int = 0
     epoch_len: int = 4
 
 
@@ -159,7 +159,7 @@ class OutputWritingSettings(BaseModel):
     """
 
     structure_format: Literal["pdb", "cif"] = "cif"
-    full_confidence_output_format: Literal["json", "npz"] = "json"
+    full_confidence_output_format: Literal["json", "npz"] = "npz"
     write_features: bool = False
     write_latent_outputs: bool = False
 
diff --git a/openfold3/projects/of3_all_atom/config/model_config.py b/openfold3/projects/of3_all_atom/config/model_config.py
index 11ba1dba..547d500a 100644
--- a/openfold3/projects/of3_all_atom/config/model_config.py
+++ b/openfold3/projects/of3_all_atom/config/model_config.py
@@ -84,6 +84,9 @@ model_selection_metric_weights_config = mlc.FrozenConfigDict(
     }
 )
 
+from sys import platform
+deepspeed_available = (platform == 'linux')  # Deepspeed is linux only
+
 model_config = mlc.ConfigDict(
     {
         "settings": {
@@ -92,7 +95,7 @@ model_config = mlc.ConfigDict(
                     "chunk_size": None,
                     # Use DeepSpeed memory-efficient attention kernel. Mutually
                     # exclusive with use_lma.
-                    "use_deepspeed_evo_attention": False,
+                    "use_deepspeed_evo_attention": deepspeed_available,
                     "use_cueq_triangle_kernels": False,
                     # Use Staats & Rabe's low-memory attention algorithm. Mutually
                     # exclusive with use_deepspeed_evo_attention.
@@ -104,7 +107,7 @@ model_config = mlc.ConfigDict(
                 },
                 "eval": {
                     "chunk_size": None,
-                    "use_deepspeed_evo_attention": True,
+                    "use_deepspeed_evo_attention":  deepspeed_available,
                     "use_cueq_triangle_kernels": False,
                     "use_lma": False,
                     "msa_module": {
diff --git a/openfold3/projects/of3_all_atom/runner.py b/openfold3/projects/of3_all_atom/runner.py
index 3e72c6f1..7e58d979 100644
--- a/openfold3/projects/of3_all_atom/runner.py
+++ b/openfold3/projects/of3_all_atom/runner.py
@@ -901,13 +901,15 @@ class OpenFold3AllAtom(ModelRunner):
             and num_atoms > self.config.settings.memory.eval.per_sample_atom_cutoff
         )
 
+        logger.info(f'Computing confidence scores for {", ".join(batch["query_id"])}')
         confidence_scores = get_confidence_scores(
             batch=batch,
             outputs=outputs,
             config=self.config,
             compute_per_sample=compute_per_sample,
         )
-
+        logger.info('Finished computing confidence scores')
+        
         return confidence_scores
 
     def predict_step(self, batch, batch_idx):
@@ -926,8 +928,8 @@ class OpenFold3AllAtom(ModelRunner):
         self.reseed(seed[0])  # TODO: assuming we have bs = 1 for now
 
         # Probably need to change the logic
-        logger.debug(
-            f"Started inference for {', '.join(query_id)} on rank {self.global_rank} "
+        logger.info(
+            f"Started inference for {', '.join(query_id)} ({batch_idx+1}) on rank {self.global_rank} "
             f"step {self.global_step}"
         )
         try:
diff --git a/openfold3/run_openfold.py b/openfold3/run_openfold.py
index 732e0078..b363beb3 100644
--- a/openfold3/run_openfold.py
+++ b/openfold3/run_openfold.py
@@ -20,6 +20,13 @@ Main run script for OpenFold3. Please see the README for usage details.
 # ruff: noqa: F821
 
 import logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+import sys, os
+if sys.platform != 'linux' and 'TORCH_LOGS' not in os.environ:
+    # Suppress torch logging of "NOTE: Redirects are currently not supported in Windows or MacOs."
+    os.environ['TORCH_LOGS'] = '-distributed'
+
 from pathlib import Path
 
 import click
@@ -146,6 +153,24 @@ def train(runner_yaml: Path, seed: int | None = None, data_seed: int | None = No
     required=False,
     help="Output directory for writing results",
 )
+@click.option(
+    "--msa_and_templates_only",
+    type=bool,
+    default=False,
+    help="Use ColabFold MSA server to get MSAs and templates but don't predict structures.",
+)
+@click.option(
+    "--device",
+    type=click.Choice(["gpu", "cpu", "tpu"]),
+    help="The device to use for prediction: gpu, cpu, tpu. Default is gpu.",
+    default="gpu",
+)
+@click.option(
+    "--precision",
+    type=click.Choice(["bf16-mixed", "32-true"]),
+    help="The floating point precision used by pytorch_lightning.Trainer. Default is bf16-mixed when using CUDA on Linux or Windows, otherwise 32-true.",
+    default=None,
+)
 def predict(
     query_json: Path,
     inference_ckpt_path: Path | None = None,
@@ -156,6 +181,9 @@ def predict(
     use_msa_server: bool = True,
     use_templates: bool = True,
     output_dir: Path | None = None,
+    msa_and_templates_only: bool = False,
+    device: str = "gpu",
+    precision: str | None = None,
 ):
     """Perform inference on a set of queries defined in the query_json."""
     _torch_gpu_setup()
@@ -178,6 +206,16 @@ def predict(
         inference_ckpt_name=inference_ckpt_name,
         **runner_args,
     )
+    expt_config.pl_trainer_args.accelerator = device
+
+    if precision is None:
+        precision = 32    # bfloat16 is many times slower than float32 except with Nvidia CUDA.
+        if sys.platform in ('linux', 'win32') and device == 'gpu':
+            import torch
+            if torch.cuda.is_available():
+                precision = "bf16-mixed"
+    expt_config.pl_trainer_args.precision = precision
+    
     expt_runner = InferenceExperimentRunner(
         expt_config,
         num_diffusion_samples,
@@ -186,10 +224,17 @@ def predict(
         use_templates,
         output_dir,
     )
-
+    if device == 'cpu':
+        expt_runner.model_config['settings']['memory']['eval']['use_deepspeed_evo_attention'] = False
+    
     # Load inference query set
     query_set = InferenceQuerySet.from_json(query_json)
 
+    if msa_and_templates_only:
+        expt_runner.inference_query_set = query_set
+        expt_runner.lightning_data_module.prepare_data()
+        return
+        
     # Run the forward pass
     expt_runner.setup()
     expt_runner.run(query_set)
diff --git a/openfold3/setup_openfold.py b/openfold3/setup_openfold.py
index 5daa1e47..132daac5 100644
--- a/openfold3/setup_openfold.py
+++ b/openfold3/setup_openfold.py
@@ -42,14 +42,17 @@ logging.basicConfig(level=logging.INFO, format="%(message)s")
 logger = logging.getLogger(__name__)
 
 
-def setup_openfold_cache() -> tuple[Path, Path]:
+def setup_openfold_cache(prompt = True) -> tuple[Path, Path]:
     """Set up the OpenFold cache directory."""
     logger.info("Setting up OpenFold cache directory...")
 
     default_cache = Path.home() / ".openfold3"
-    user_input = input(
-        f"Please specify the OpenFold cache directory (default: {default_cache}): "
-    ).strip()
+    if prompt:
+        user_input = input(
+            f"Please specify the OpenFold cache directory (default: {default_cache}): "
+        ).strip()
+    else:
+        user_input = default_cache
 
     # Use user input if provided, otherwise use default
     if user_input:
@@ -66,7 +69,7 @@ def setup_openfold_cache() -> tuple[Path, Path]:
 
 
 def setup_param_directory(
-    openfold_cache: Path, ckpt_root_file: Path
+        openfold_cache: Path, ckpt_root_file: Path, prompt: bool = True
 ) -> tuple[Path, bool]:
     """Check and set up the parameter directory."""
 
@@ -80,12 +83,15 @@ def setup_param_directory(
         logger.info(
             f"OpenFold3 parameters may already be installed at: {existing_path}"
         )
-        logger.info("Do you want to:")
-        logger.info("1) Use existing parameters (skip download)")
-        logger.info("2) Download to a new location")
-        logger.info("3) Re-download to existing location")
+        if prompt:
+            logger.info("Do you want to:")
+            logger.info("1) Use existing parameters (skip download)")
+            logger.info("2) Download to a new location")
+            logger.info("3) Re-download to existing location")
 
-        choice = input("Enter your choice (1/2/3): ").strip()
+            choice = input("Enter your choice (1/2/3): ").strip()
+        else:
+            choice = "1"
 
         if choice == "1":
             logger.info(f"Using existing parameters at: {existing_path}")
@@ -107,11 +113,13 @@ def setup_param_directory(
     else:
         # First time setup
         logger.info("Downloading OpenFold3 parameters...")
-        user_input = input(
-            "Please specify the directory for parameter download "
-            f"(default: {openfold_cache}): "
-        ).strip()
-
+        if prompt:
+            user_input = input(
+                "Please specify the directory for parameter download "
+                f"(default: {openfold_cache}): "
+            ).strip()
+        else:
+            user_input = openfold_cache
         # Use user input if provided, otherwise use default (the cache directory)
         if user_input:
             param_dir = Path(user_input).expanduser()
@@ -130,7 +138,7 @@ def setup_param_directory(
     return param_dir, True  # Proceed with download
 
 
-def download_parameters(param_dir) -> None:
+def download_parameters(param_dir, prompt = True) -> None:
     """Perform the parameter download."""
     # Exclude incompatible checkpoints:
     all_checkpoints = [
@@ -139,13 +147,16 @@ def download_parameters(param_dir) -> None:
         if name not in LEGACY_CHECKPOINTS
     ]
 
-    logger.info("Select parameters to download:")
-    logger.info(f"1) Download only the default checkpoint ({DEFAULT_CHECKPOINT_NAME})")
-    logger.info(f"2) Download all parameters ({', '.join(all_checkpoints)}) (default)")
-    logger.info("3) Download a specific parameter by name")
-
-    choice = input("Enter your choice (1/2/3, default: 1): ").strip() or "1"
+    if prompt:
+        logger.info("Select parameters to download:")
+        logger.info(f"1) Download only the default checkpoint ({DEFAULT_CHECKPOINT_NAME})")
+        logger.info(f"2) Download all parameters ({', '.join(all_checkpoints)}) (default)")
+        logger.info("3) Download a specific parameter by name")
 
+        choice = input("Enter your choice (1/2/3, default: 1): ").strip() or "1"
+    else:
+        choice = "1"
+        
     logger.info("Starting parameter download...")
 
     if choice == "1":
@@ -239,23 +250,24 @@ def run_integration_tests() -> None:
     logger.info("Integration tests passed!")
 
 
-def main():
+def main(prompt = True, test = True):
     """Main execution."""
     # Step 1: Set up OpenFold cache directory
-    openfold_cache, ckpt_root_file = setup_openfold_cache()
+    openfold_cache, ckpt_root_file = setup_openfold_cache(prompt)
 
     # Step 2: Set up checkpoint directory
-    param_dir, should_download = setup_param_directory(openfold_cache, ckpt_root_file)
+    param_dir, should_download = setup_param_directory(openfold_cache, ckpt_root_file, prompt)
 
     # Step 3: Perform download if needed
     if should_download:
-        download_parameters(param_dir)
+        download_parameters(param_dir, prompt)
 
     # Step 4: Setup CCD with biotite
     setup_biotite_ccd(ccd_path=biotite.setup_ccd.OUTPUT_CCD, force_download=False)
 
     # Step 5: Run tests (always run regardless of download status)
-    run_integration_tests()
+    if test:
+        run_integration_tests()
 
 
 if __name__ == "__main__":
diff --git a/pyproject.toml b/pyproject.toml
index 6837874b..4b3c3401 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,8 +20,8 @@ dependencies = [
   "typing-extensions",
   "wandb",
   "ml-collections",
-  "mkl",
-  "aria2",
+  "mkl ; sys_platform != 'darwin'",
+  "aria2 ; sys_platform != 'darwin'",
   "biotite",
   "rdkit",
   "boto3",
@@ -29,18 +29,20 @@ dependencies = [
   "func_timeout",
   "lmdb", # == 1.6  # This verison might be required to avoid Error during training
   "ijson",
-  "deepspeed",
+  # Deepspeed version 0.18.5 produces the following error during predictions on Linux
+  # "RuntimeError: Unable to JIT load the evoformer_attn op due to it not being compatible due to hardware/software issue."
+  "deepspeed==0.18.4 ; sys_platform != 'win32'",
   "pdbeccdutils",
   # For cutlass >=4, import library is cutlass_cppgen, but deepspeed is
   # currently using cutlass.
-  "nvidia-cutlass <4",
-  "kalign-python",
+  "nvidia-cutlass <4 ; sys_platform != 'darwin'",
+  "kalign-python ; sys_platform != 'win32'",
 
   # additional pins
 
   # Pinned cutlass <4 calls from cuda.__version__, which was removed when
   # cuda-python went fully metapackage
-  "cuda-python <12.9.1",
+  "cuda-python <12.9.1 ; sys_platform != 'darwin'",
 ]
 authors = [
   {name = "OpenFold Team", email = "support@openfold.io" }
