RFD4ProteinaScheduler

RFD4ProteinaScheduler is the temporal-score-rescaled flow sampler of RFD4-Proteína for one modality. The pipeline uses two instances, one for the backbone (bb_ca, coordinates in nanometers) and one for the per-token local latents (local_latents), each with its own time schedule, noise schedule and rescaling. See RFD4ProteinaPipeline for the sampling loop.

RFD4ProteinaScheduler

class diffusers.RFD4ProteinaScheduler

< source >

( schedule_mode: strnoise_mode: strtsr_k: floattsr_sigma: floattail_steps: int = 10tail_step_multiplier: int = 2t_lim_ode: float = 0.95t_lim_ode_below: float = 0.02 )

Parameters

schedule_mode (str) — Spacing of the times between noise and data: "log" or "power".
noise_mode (str) — Shape of the noise injected in the stochastic window: "1/t" or "tan".
tsr_k (float) — How low the sampling temperature goes; higher is lower.
tsr_sigma (float) — How fast the low temperature is reached; higher is faster.
tail_steps (int, defaults to 10) — Number of final intervals that are finer than the rest by tail_step_multiplier.
tail_step_multiplier (int, defaults to 2) — How much finer the final intervals are.
t_lim_ode (float, defaults to 0.95) — Above this time the step is deterministic.
t_lim_ode_below (float, defaults to 0.02) — Below this time the step is deterministic.

Temporal-score-rescaled flow sampler for one modality of RFD4-Proteína.

The pipeline integrates two modalities, the backbone and the per-token local latents, with one scheduler each. They differ in their time schedule, noise schedule and rescaling.

Call set_timesteps before sampling. Times increase from noise toward clean data. A scheduler tracks the current interval and should not be shared by concurrent generations.

set_timesteps

< source >

( num_inference_steps: intdevice: typing.Union[str, torch.device, NoneType] = Nonestart_time: float | None = None )

Parameters

num_inference_steps (int) — Number of integration intervals.
device (str or torch.device, optional) — Device for timesteps. The interval lengths and noise coefficients stay on the CPU.
start_time (float, optional) — Partial-diffusion start time in [0, 1). Omit for full generation.

Resolve the schedule and reset the integration cursor.

add_noise

< source >

( original_samples: Tensornoise: Tensortimesteps: typing.Union[float, torch.Tensor] )

Parameters

original_samples (torch.Tensor) — Clean state (batch, tokens, D).
noise (torch.Tensor) — Noise of the same shape, drawn by the caller.
timesteps (float or torch.Tensor) — Time in [0, 1), a scalar or (batch,).

Sample the forward marginal (1 - t) * noise + t * clean.

step

< source >

( model_output: Tensortimestep: Tensorsample: Tensoremit_mask: Tensorgenerator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonereturn_dict: bool = True ) → RFD4ProteinaSchedulerOutput or tuple

Parameters

model_output (torch.Tensor) — Predicted velocity, shaped like sample.
timestep (torch.Tensor) — The next element of timesteps. Out-of-order steps are rejected.
sample (torch.Tensor) — Current state (B, N, D).
emit_mask (torch.Tensor) — Boolean (B, N), true for the tokens that are updated (not context or padding).
generator (torch.Generator or list[torch.Generator], optional) — A generator for a single row or distinct generators for each batch row. Omitting them uses the global stream.
return_dict (bool, defaults to True) — Return RFD4ProteinaSchedulerOutput; otherwise a one-item tuple holding the state.

Returns

RFD4ProteinaSchedulerOutput or tuple

The updated state.

Advance one integration interval and update step_index.