File size: 1,149 Bytes
f545704 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """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
|