modular_blocks.py

# 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."
        )