File size: 770 Bytes
1b703d5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | """Position encoding options used by exported dense DiT blocks."""
from __future__ import annotations
from enum import Enum
class DiTPositionEncoding(Enum):
"""Position encoding strategy used inside exported DiT blocks."""
ROPE_2D_AXIAL_DILATED = "rope_2d_axial_dilated"
ROPE_2D_AXIAL_UNNORMALIZED_DILATED = "rope_2d_axial_unnormalized_dilated"
ROPE_2D_AXIAL_NORMALIZED = "rope_2d_axial_normalized"
ROPE_2D_AXIAL_UNNORMALIZED = "rope_2d_axial_unnormalized"
ROPE_2D_AXIAL_FREQ_AWARE = "rope_2d_axial_freq_aware"
ROPE_2D_AXIAL_BETA_WARP = "rope_2d_axial_beta_warp"
ROPE_2D_AXIAL_ALPHA_WARP = "rope_2d_axial_alpha_warp"
ROPE_3D_ZIMAGE = "rope_3d_zimage"
ROPE_1D = "rope_1d"
NONE = "none"
__all__ = ["DiTPositionEncoding"]
|