 and upsampling blocks used to create the UNet architecture. block_out_channels: the number of output channels of the downsampling blocks; also used in reverse order for the number of input channels of the upsampling blocks. layers_per_block: the number of ResNet blocks present in each UNet block. To use the model for inference, create the image shape with random Gaussian noise. It should have a batch axis because the model can receive multiple