add it completely into the unet package?

or keep it separate?