The model definition file’ mindone/models/dit.py, with flash-attention support, based on @geniuspatrick 's implementation. When enabling FA, the inference results on 910B with ms 2.2.10 were ok.
Usage:
from mindone.models.dit import DiT_models
model = DiT_models["DiT_XL_2"](input_size=32, block_kwargs={"enable_flash_attention": True})
TODO:
[ ] Automatic downloading of pretrained checkpoints
The model definition file’
mindone/models/dit.py, with flash-attention support, based on @geniuspatrick 's implementation. When enabling FA, the inference results on 910B with ms 2.2.10 were ok.Usage:
TODO: