Hugging Face's logo

Diffusers documentation

RFD4ProteinaScheduler

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.41.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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

< >

( 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

< >

( 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

< >

( 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

< >

( 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.

RFD4ProteinaSchedulerOutput

class diffusers.schedulers.scheduling_rfd4_proteina.RFD4ProteinaSchedulerOutput

< >

( prev_sample: Tensor )

Parameters

  • prev_sample (torch.Tensor) — The state after the step, shaped (B, N, D): D is 3 for the backbone (coordinates in nanometers) and the latent width for the local latents.

Updated state of one modality.

Update on GitHub

←RePaintScheduler ScoreSdeVeScheduler→