ed here.   split_head_dim (bool, optional, defaults to False) —
Whether to split the head dimension into a new axis for the self-attention computation. In most cases,
enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL.    A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
shaped output. This model inherits from FlaxModelMixin. Check the s