| """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 |
|
|
