config.json

{
  "architectures": [
    "STDiT2"
  ],
  "auto_map": {
    "AutoConfig": "configuration_stdit2.STDiT2Config",
    "AutoModel": "modeling_stdit2.STDiT2"
  },
  "caption_channels": 4096,
  "class_dropout_prob": 0.1,
  "depth": 28,
  "drop_path": 0.0,
  "enable_flash_attn": false,
  "enable_layernorm_kernel": false,
  "enable_sequence_parallelism": false,
  "freeze": null,
  "hidden_size": 1152,
  "in_channels": 4,
  "input_size": [
    32,
    60,
    106
  ],
  "input_sq_size": 512,
  "mlp_ratio": 4.0,
  "model_max_length": 200,
  "model_type": "stdit2",
  "no_temporal_pos_emb": false,
  "num_heads": 16,
  "patch_size": [
    1,
    2,
    2
  ],
  "pred_sigma": true,
  "qk_norm": true,
  "torch_dtype": "float32",
  "transformers_version": "4.40.1"
}