Synchronized timer to count the inference time of `nn.Module.forward`. Supports both context manager and decorator usage. Example a

ODE solver class

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

在训练 batch 结束后，记录所有 HunYuanDiTBlock 内所有 RMSNorm 的 scale 最大值。

We rename some layer names to align with flash attention

The final layer of HunYuanDiT.

NOTE: this is very similar to diffusers.FlowMatchEulerDiscreteScheduler. Except our timesteps are reversed Euler scheduler. This model

Output class for the scheduler's `step` function output. Args: prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels,

The sin/cosine positional embedding. Given an input tensor `x` of shape [n_batch, ..., c_dim], it converts each feature dimension of `x[..., i]` i

用于进行 skip_connection 的ablation测试 先做norm 再 linear

只在训练 batch 结束后，记录最后一个 HunYuanDiTBlock 内所有 RMSNorm 的 scale 最大值。

An iterable dataset yielding a list of urls.

Embeds scalar timesteps into vector representations.

ULIP similarity score calculator This class wraps the ULIP-PointBERT model, providing functionality to compute similarity scores be

SDE solver class

The trick function of adding auxiliary (aux) loss, which includes the gradient of the aux loss during backpropagation.

Linear Coupling Plan

note: use with a base_lr of 1.0

view order: front, front clockwise 90, back, front clockwise 270

Which type of output the model predicts.

Which type of path to use.

Sampler class for the transport model

class for VP path flow matching

Which type of weighting to use.