fix

final tweaks to flow version, still not good enough, will try another…

… day

bring in axial positional embeddings

allow for training under a bit of noise

seeing it with autoregressive flow. putting more work into it

just improvise a version with flow matching

Merge pull request #4 from lucidrains/image-trainer

Image trainer

0.2.1

deviate from paper and use EDM from Karras et al.

should be using mlp width, and default to transformer model dimension…

…s if not given