dn6 HF Staff

[RFD3] Feed MPNN real backbone atoms and the schema it validates (#2)

51e5b34 about 1 month ago

6.01 kB

# Copyright 2025 Dhruv Nair. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
RFDiffusion3 Scheduler.
A thin diffusers-compatible wrapper around the foundry EDM noise schedule
and stochastic sampling logic from `rfd3.model.inference_sampler`.
"""

from typing import Optional

import torch

from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin

# Reuse the original noise schedule and sampling config directly
from rfd3.model.inference_sampler import SampleDiffusionWithMotif

# `ComponentSpec.load` only strips `dtype` for a `type_hint` that is not a torch module, and
# modular_model_index.json records this scheduler as `AutoModel`, so the loader forwards its
# weight-placement kwargs here. None of them are constructor arguments, and letting
# `register_to_config` capture `dtype` would write a torch.dtype into config.json.
_LOADER_ONLY_KWARGS = ("dtype", "torch_dtype", "device_map", "variant", "trust_remote_code", "low_cpu_mem_usage")


class RFDiffusionScheduler(SchedulerMixin, ConfigMixin):
    """
    Diffusers-compatible scheduler wrapping the foundry EDM sampler.
    Delegates noise schedule construction and sampling parameters to
    `rfd3.model.inference_sampler.SampleDiffusionWithMotif`.
    """

    config_name = "config.json"

    @register_to_config
    def __init__(
        self,
        num_timesteps: int = 200,
        sigma_data: float = 16.0,
        s_min: float = 4e-4,
        s_max: float = 160.0,
        p: float = 7.0,
        gamma_0: float = 0.6,
        gamma_min: float = 1.0,
        noise_scale: float = 1.003,
        step_scale: float = 1.5,
    ):
        # Instantiate the foundry sampler with matching parameters
        self._sampler = SampleDiffusionWithMotif(
            num_timesteps=num_timesteps,
            sigma_data=sigma_data,
            s_min=s_min,
            s_max=s_max,
            p=p,
            gamma_0=gamma_0,
            gamma_min=gamma_min,
            noise_scale=noise_scale,
            step_scale=step_scale,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path=None, subfolder=None, **kwargs):
        for key in _LOADER_ONLY_KWARGS:
            kwargs.pop(key, None)

        return super().from_pretrained(pretrained_model_name_or_path, subfolder=subfolder, **kwargs)

    @property
    def sampler(self) -> SampleDiffusionWithMotif:
        return self._sampler

    def get_noise_schedule(self, device: torch.device = None) -> torch.Tensor:
        """
        Construct the EDM noise schedule using the foundry implementation.
        Returns:
            torch.Tensor: Noise schedule [num_timesteps] from high to low noise.
        """
        return self._sampler._construct_inference_noise_schedule(
            device=device or torch.device("cpu")
        )

    def get_initial_noise_level(self, device: torch.device = None) -> torch.Tensor:
        """Get the first (largest) noise level from the schedule."""
        return self.get_noise_schedule(device=device)[0]

    def step(
        self,
        xyz_pred: torch.Tensor,
        xyz_noisy: torch.Tensor,
        c_t_minus_1: torch.Tensor,
        c_t: torch.Tensor,
        motif_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Perform one Euler denoising step matching the foundry sampler.
        The foundry ``sample_diffusion_like_af3`` does NOT clamp motif
        coordinates after the Euler update — it relies on noise injection
        having zeroed epsilon for motif atoms so the model's delta is ~0
        there.  We replicate that behaviour here.
        Args:
            xyz_pred: Model's denoised prediction X_denoised_L [B, L, 3]
            xyz_noisy: Noise-injected coordinates X_noisy_L [B, L, 3]
            c_t_minus_1: Previous noise level
            c_t: Next (lower) noise level
            motif_mask: Boolean mask for fixed positions (True = fixed) [L]
                        (unused — kept for API compatibility)
        Returns:
            Updated coordinates X_L [B, L, 3]
        """
        gamma = self._sampler.gamma_0 if c_t > self._sampler.gamma_min else 0.0
        t_hat = c_t_minus_1 * (gamma + 1.0)

        delta_L = (xyz_noisy - xyz_pred) / t_hat
        d_t = c_t - t_hat
        xyz_next = xyz_noisy + self._sampler.step_scale * d_t * delta_L

        return xyz_next

    def add_noise(
        self,
        xyz: torch.Tensor,
        c_t_minus_1: torch.Tensor,
        c_t: torch.Tensor,
        motif_mask: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Inject stochastic noise before the model call, matching the foundry sampler.
        Args:
            xyz: Current coordinates X_L [B, L, 3]
            c_t_minus_1: Previous noise level
            c_t: Current (next lower) noise level
            motif_mask: Boolean mask for fixed positions (True = fixed) [L]
        Returns:
            Tuple of (noisy coordinates X_noisy_L, t_hat scalar)
        """
        gamma = self._sampler.gamma_0 if c_t > self._sampler.gamma_min else 0.0
        t_hat = c_t_minus_1 * (gamma + 1.0)

        noise_std = self._sampler.noise_scale * torch.sqrt(t_hat**2 - c_t_minus_1**2)
        epsilon = noise_std * torch.randn_like(xyz)

        if motif_mask is not None:
            epsilon[:, motif_mask] = 0.0

        xyz_noisy = xyz + epsilon
        return xyz_noisy, t_hat

	# Copyright 2025 Dhruv Nair. All rights reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	"""
	RFDiffusion3 Scheduler.

	A thin diffusers-compatible wrapper around the foundry EDM noise schedule
	and stochastic sampling logic from `rfd3.model.inference_sampler`.
	"""

	from typing import Optional

	import torch

	from diffusers.configuration_utils import ConfigMixin, register_to_config
	from diffusers.schedulers.scheduling_utils import SchedulerMixin

	# Reuse the original noise schedule and sampling config directly
	from rfd3.model.inference_sampler import SampleDiffusionWithMotif

	# `ComponentSpec.load` only strips `dtype` for a `type_hint` that is not a torch module, and
	# modular_model_index.json records this scheduler as `AutoModel`, so the loader forwards its
	# weight-placement kwargs here. None of them are constructor arguments, and letting
	# `register_to_config` capture `dtype` would write a torch.dtype into config.json.
	_LOADER_ONLY_KWARGS = ("dtype", "torch_dtype", "device_map", "variant", "trust_remote_code", "low_cpu_mem_usage")


	class RFDiffusionScheduler(SchedulerMixin, ConfigMixin):
	"""
	Diffusers-compatible scheduler wrapping the foundry EDM sampler.

	Delegates noise schedule construction and sampling parameters to
	`rfd3.model.inference_sampler.SampleDiffusionWithMotif`.
	"""

	config_name = "config.json"

	@register_to_config
	def __init__(
	self,
	num_timesteps: int = 200,
	sigma_data: float = 16.0,
	s_min: float = 4e-4,
	s_max: float = 160.0,
	p: float = 7.0,
	gamma_0: float = 0.6,
	gamma_min: float = 1.0,
	noise_scale: float = 1.003,
	step_scale: float = 1.5,
	):
	# Instantiate the foundry sampler with matching parameters
	self._sampler = SampleDiffusionWithMotif(
	num_timesteps=num_timesteps,
	sigma_data=sigma_data,
	s_min=s_min,
	s_max=s_max,
	p=p,
	gamma_0=gamma_0,
	gamma_min=gamma_min,
	noise_scale=noise_scale,
	step_scale=step_scale,
	)

	@classmethod
	def from_pretrained(cls, pretrained_model_name_or_path=None, subfolder=None, **kwargs):
	for key in _LOADER_ONLY_KWARGS:
	kwargs.pop(key, None)

	return super().from_pretrained(pretrained_model_name_or_path, subfolder=subfolder, **kwargs)

	@property
	def sampler(self) -> SampleDiffusionWithMotif:
	return self._sampler

	def get_noise_schedule(self, device: torch.device = None) -> torch.Tensor:
	"""
	Construct the EDM noise schedule using the foundry implementation.

	Returns:
	torch.Tensor: Noise schedule [num_timesteps] from high to low noise.
	"""
	return self._sampler._construct_inference_noise_schedule(
	device=device or torch.device("cpu")
	)

	def get_initial_noise_level(self, device: torch.device = None) -> torch.Tensor:
	"""Get the first (largest) noise level from the schedule."""
	return self.get_noise_schedule(device=device)[0]

	def step(
	self,
	xyz_pred: torch.Tensor,
	xyz_noisy: torch.Tensor,
	c_t_minus_1: torch.Tensor,
	c_t: torch.Tensor,
	motif_mask: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""
	Perform one Euler denoising step matching the foundry sampler.

	The foundry ``sample_diffusion_like_af3`` does NOT clamp motif
	coordinates after the Euler update — it relies on noise injection
	having zeroed epsilon for motif atoms so the model's delta is ~0
	there. We replicate that behaviour here.

	Args:
	xyz_pred: Model's denoised prediction X_denoised_L [B, L, 3]
	xyz_noisy: Noise-injected coordinates X_noisy_L [B, L, 3]
	c_t_minus_1: Previous noise level
	c_t: Next (lower) noise level
	motif_mask: Boolean mask for fixed positions (True = fixed) [L]
	(unused — kept for API compatibility)

	Returns:
	Updated coordinates X_L [B, L, 3]
	"""
	gamma = self._sampler.gamma_0 if c_t > self._sampler.gamma_min else 0.0
	t_hat = c_t_minus_1 * (gamma + 1.0)

	delta_L = (xyz_noisy - xyz_pred) / t_hat
	d_t = c_t - t_hat
	xyz_next = xyz_noisy + self._sampler.step_scale * d_t * delta_L

	return xyz_next

	def add_noise(
	self,
	xyz: torch.Tensor,
	c_t_minus_1: torch.Tensor,
	c_t: torch.Tensor,
	motif_mask: Optional[torch.Tensor] = None,
	) -> tuple[torch.Tensor, torch.Tensor]:
	"""
	Inject stochastic noise before the model call, matching the foundry sampler.

	Args:
	xyz: Current coordinates X_L [B, L, 3]
	c_t_minus_1: Previous noise level
	c_t: Current (next lower) noise level
	motif_mask: Boolean mask for fixed positions (True = fixed) [L]

	Returns:
	Tuple of (noisy coordinates X_noisy_L, t_hat scalar)
	"""
	gamma = self._sampler.gamma_0 if c_t > self._sampler.gamma_min else 0.0
	t_hat = c_t_minus_1 * (gamma + 1.0)

	noise_std = self._sampler.noise_scale * torch.sqrt(t_hat2 - c_t_minus_12)
	epsilon = noise_std * torch.randn_like(xyz)

	if motif_mask is not None:
	epsilon[:, motif_mask] = 0.0

	xyz_noisy = xyz + epsilon
	return xyz_noisy, t_hat