"""Model configuration constants for Z-Image."""

ADALN_EMBED_DIM = 256
SEQ_MULTI_OF = 32

ROPE_THETA = 256.0
ROPE_AXES_DIMS = [32, 48, 48]
ROPE_AXES_LENS = [1536, 512, 512]

FREQUENCY_EMBEDDING_SIZE = 256
MAX_PERIOD = 10000

BASE_IMAGE_SEQ_LEN = 256
MAX_IMAGE_SEQ_LEN = 4096
BASE_SHIFT = 0.5
MAX_SHIFT = 1.15

DEFAULT_VAE_SCALE_FACTOR = 8
DEFAULT_VAE_IN_CHANNELS = 3
DEFAULT_VAE_OUT_CHANNELS = 3
DEFAULT_VAE_LATENT_CHANNELS = 4
DEFAULT_VAE_NORM_NUM_GROUPS = 32
DEFAULT_VAE_SCALING_FACTOR = 0.18215

DEFAULT_TRANSFORMER_PATCH_SIZE = (2,)
DEFAULT_TRANSFORMER_F_PATCH_SIZE = (1,)
DEFAULT_TRANSFORMER_IN_CHANNELS = 16
DEFAULT_TRANSFORMER_DIM = 3840
DEFAULT_TRANSFORMER_N_LAYERS = 30
DEFAULT_TRANSFORMER_N_REFINER_LAYERS = 2
DEFAULT_TRANSFORMER_N_HEADS = 30
DEFAULT_TRANSFORMER_N_KV_HEADS = 30
DEFAULT_TRANSFORMER_NORM_EPS = 1e-5
DEFAULT_TRANSFORMER_QK_NORM = True
DEFAULT_TRANSFORMER_CAP_FEAT_DIM = 2560
DEFAULT_TRANSFORMER_T_SCALE = 1000.0

DEFAULT_SCHEDULER_NUM_TRAIN_TIMESTEPS = 1000
DEFAULT_SCHEDULER_SHIFT = 3.0
DEFAULT_SCHEDULER_USE_DYNAMIC_SHIFTING = False

DEFAULT_LOAD_DEVICE = "cuda"
DEFAULT_LOAD_DTYPE_STR = "bfloat16"

BYTES_PER_GB = 2**30