I wanna change some of the layers, what approach you recommend? I wanna change the unet locally, any advice?