Diffusers for AITemplate
Work in progress
Depends on hlky/AITemplate
Currently base modules (activations, embeddings, resnet, etc.) are covered by unittests. Top level modules (FluxTransformer2DModel, UNet2DConditionModel, etc.) can be considered tested if a corresponding builder exists.
Modules that cannot currently be supported due to missing kernels in AIT, or unsupported parameters of existing kernels, are marked by raising NotImplementedError.
- SD1
- SD1 Inpaint
- SD2
- SDXL
- SDXL Refiner
- SD Cascade (prior)
- SD3
- Kadinsky3
- PixArt
- Flux-dev
- Flux-schnell
- FluxTransformerBlock, FluxSingleTransformerBlock
Modules work with any supported model weights.
Model specific inference functions are not yet implemented. For a general function see aitemplate.testing.benchmark_model - this works with any module.
diffusers_ait IR is based on Diffusers, weight mapping should not be required for model weights in Diffusers format.
Mapping for other weight formats (LDM etc.) is required and will be implemented. See diffusers/scripts for examples of how weights are mapped.