Parameters ---------- x: torch.Tensor (batch, seqlen1, hidden_dim) (where hidden_dim = num heads * head dim)

Forward process. Args: x: tensor of shape [..., dim] Returns: embedding: an embedding of `x` of shape [...,

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created f

Args: pc (torch.FloatTensor): [B, N, 3] feats (torch.FloatTensor or None): [B, N, C] Returns:

See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 Args: timesteps (`torch.T

Forward process in flow-matching Args: sample (`torch.FloatTensor`): The input sample. times

Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_inde

Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_inde

Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`)

The index counter for current timestep. It will increase 1 after each scheduler step.

The index counter for current timestep. It will increase 1 after each scheduler step.

Update target parameters to be closer to those of source parameters using an exponential moving average. :param target_params: the targe