| |
| |
| """Optional modular entry points for the separately installed RFD4 CIF preparation package.""" |
|
|
| from pathlib import Path |
| from typing import ClassVar |
|
|
| import torch |
| from diffusers.modular_pipelines import ( |
| AutoPipelineBlocks, |
| InputParam, |
| ModularPipelineBlocks, |
| OutputParam, |
| ) |
|
|
|
|
| class RFD4PrepareCIFStep(ModularPipelineBlocks): |
| """Prepare a single annotated input without loading generator weights.""" |
|
|
| _requirements: ClassVar = {"rfd4_proteina_prep": "==0.1.0", "atomworks": "==3.0.3"} |
|
|
| @property |
| def description(self): |
| return "Prepare tensors and retain the atom mapping for annotated CIF export." |
|
|
| @property |
| def inputs(self): |
| return [ |
| InputParam( |
| "cif_path", |
| required=True, |
| type_hint=str, |
| description="Local annotated CIF filename.", |
| ), |
| InputParam( |
| "seed", |
| default=42, |
| type_hint=int, |
| description="Seed for expandable segments.", |
| ), |
| InputParam( |
| "device", |
| default="cpu", |
| type_hint=str, |
| description="Device for returned tensors.", |
| ), |
| InputParam( |
| "path_length_cutoff", |
| default=64, |
| type_hint=int, |
| description="Topology path-length cutoff.", |
| ), |
| ] |
|
|
| @property |
| def intermediate_outputs(self): |
| return [ |
| OutputParam( |
| "features", |
| type_hint=dict, |
| description="Batched tensors, with coordinates in nanometers.", |
| ), |
| OutputParam( |
| "atom_array", description="Original atom mapping; retain it for export." |
| ), |
| ] |
|
|
| def __call__(self, components, state): |
| block_state = self.get_block_state(state) |
| path = Path(block_state.cif_path) |
| if not path.is_file(): |
| raise FileNotFoundError(path) |
| if type(block_state.seed) is not int: |
| raise TypeError("seed must be an integer.") |
| if ( |
| type(block_state.path_length_cutoff) is not int |
| or block_state.path_length_cutoff < 1 |
| ): |
| raise ValueError("path_length_cutoff must be a positive integer.") |
|
|
| from rfd4_proteina_prep import prepare_cif |
|
|
| features, atoms = prepare_cif( |
| path, |
| seed=block_state.seed, |
| device=torch.device(block_state.device), |
| path_length_cutoff=block_state.path_length_cutoff, |
| ) |
| block_state.features = features |
| block_state.atom_array = atoms |
| self.set_block_state(state, block_state) |
| return components, state |
|
|
|
|
| class RFD4ExportCIFStep(ModularPipelineBlocks): |
| """Export one decoded design through the same package used for preparation.""" |
|
|
| _requirements: ClassVar = {"rfd4_proteina_prep": "==0.1.0", "atomworks": "==3.0.3"} |
|
|
| @property |
| def description(self): |
| return "Write a new annotated CIF, preserving conditions and insertion mapping; never overwrite a file." |
|
|
| @property |
| def inputs(self): |
| return [ |
| InputParam( |
| "pipeline_output", |
| required=True, |
| description="RFD4ProteinaPipelineOutput with torch tensors.", |
| ), |
| InputParam( |
| "features", |
| required=True, |
| type_hint=dict, |
| description="The features used for this design batch.", |
| ), |
| InputParam( |
| "atom_array", |
| required=True, |
| description="Atom mapping returned during preparation.", |
| ), |
| InputParam( |
| "output_path", |
| required=True, |
| type_hint=str, |
| description="New .cif filename; parent must exist.", |
| ), |
| InputParam( |
| "row", |
| default=0, |
| type_hint=int, |
| description="Batch row to export using this atom mapping.", |
| ), |
| ] |
|
|
| @property |
| def intermediate_outputs(self): |
| return [ |
| OutputParam( |
| "cif_path", |
| type_hint=str, |
| description="Written CIF path; coordinates are in angstroms.", |
| ) |
| ] |
|
|
| def __call__(self, components, state): |
| block_state = self.get_block_state(state) |
| path = Path(block_state.output_path) |
| if path.suffix.lower() != ".cif": |
| raise ValueError("output_path must end in .cif.") |
| if path.exists() or path.is_symlink(): |
| raise FileExistsError(path) |
| if not path.parent.is_dir(): |
| raise FileNotFoundError(path.parent) |
| batch_size, tokens = block_state.features["padding_mask"].shape |
| if type(block_state.row) is not int or not 0 <= block_state.row < batch_size: |
| raise ValueError("row must select a batch row in features.") |
| output = block_state.pipeline_output |
| for name, shape in { |
| "coords_nm": (batch_size, tokens, 37, 3), |
| "residue_type": (batch_size, tokens), |
| "atom_mask": (batch_size, tokens, 37), |
| "insertion_logits": (batch_size, tokens, tokens), |
| }.items(): |
| value = getattr(output, name) |
| if not isinstance(value, torch.Tensor) or tuple(value.shape) != shape: |
| raise ValueError( |
| f"pipeline_output.{name} must be a torch tensor with shape {shape}." |
| ) |
|
|
| from rfd4_proteina_prep import write_cif |
|
|
| block_state.cif_path = write_cif( |
| output, |
| block_state.features, |
| block_state.atom_array, |
| path, |
| row=block_state.row, |
| ) |
| self.set_block_state(state, block_state) |
| return components, state |
|
|
|
|
| class RFD4CIFBlocks(AutoPipelineBlocks): |
| """Prepare a CIF, or export when pipeline_output is provided; never perform model inference.""" |
|
|
| block_classes: ClassVar = [RFD4ExportCIFStep, RFD4PrepareCIFStep] |
| block_names: ClassVar = ["export", "prepare"] |
| block_trigger_inputs: ClassVar = ["pipeline_output", None] |
|
|
