Apply the PixArt fix here as well. Move models to checkpoints folder as that makes more sense in this case.