File size: 2,547 Bytes
a376829 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | # Copyright 2025 Dhruv Nair. All rights reserved.
# Licensed under the Apache License, Version 2.0
from diffusers.utils import logging
from diffusers.modular_pipelines import AutoPipelineBlocks, SequentialPipelineBlocks
from .before_denoise import (
RF3InputStep,
RF3PrepareLatentsStep,
RF3RecyclingStep,
RF3SetTimestepsStep,
)
from .decoders import RF3DecodeStep
from .denoise import RF3DenoiseStep
logger = logging.get_logger(__name__)
class RF3BeforeDenoiseStep(SequentialPipelineBlocks):
"""Sequential block for pre-denoising: input → timesteps → recycling → latents."""
block_classes = [
RF3InputStep,
RF3SetTimestepsStep,
RF3RecyclingStep,
RF3PrepareLatentsStep,
]
block_names = ["input", "set_timesteps", "recycling", "prepare_latents"]
@property
def description(self):
return (
"Before denoise step:\n"
" - `RF3InputStep` parses sequence and builds feature dict\n"
" - `RF3SetTimestepsStep` constructs EDM noise schedule\n"
" - `RF3RecyclingStep` runs trunk recycler (pairformer + MSA + templates)\n"
" - `RF3PrepareLatentsStep` samples initial noised coordinates\n"
)
class RF3AutoBeforeDenoiseStep(AutoPipelineBlocks):
block_classes = [RF3BeforeDenoiseStep]
block_names = ["fold"]
block_trigger_inputs = [None]
@property
def description(self):
return "Before denoise step for RF3 structure prediction."
class RF3AutoDenoiseStep(AutoPipelineBlocks):
block_classes = [RF3DenoiseStep]
block_names = ["denoise"]
block_trigger_inputs = [None]
@property
def description(self) -> str:
return "Denoise step for RF3 structure prediction."
class RF3AutoDecodeStep(AutoPipelineBlocks):
block_classes = [RF3DecodeStep]
block_names = ["decode"]
block_trigger_inputs = [None]
@property
def description(self):
return "Decode step for RF3 — coordinates to tensor/PDB/CIF."
class RF3AutoBlocks(SequentialPipelineBlocks):
"""Full RF3 structure prediction pipeline."""
block_classes = [
RF3AutoBeforeDenoiseStep,
RF3AutoDenoiseStep,
RF3AutoDecodeStep,
]
block_names = [
"before_denoise",
"denoise",
"decoder",
]
@property
def description(self):
return (
"Modular pipeline for protein structure prediction using RF3.\n"
"Provide `sequence` to predict a protein's 3D structure."
)
|