def __init__(self, channels, emb_channels, dropout, out_channels=None, use_conv=False, use_scale_shift_norm=False, dims=2, use_checkpoint=False, up=False, down=False):
    super().__init__()
    self.channels = channels
    self.emb_channels = emb_channels
    self.dropout = dropout
    self.out_channels = (out_channels or channels)
    self.use_conv = use_conv
    self.use_checkpoint = use_checkpoint
    self.use_scale_shift_norm = use_scale_shift_norm
    self.in_layers = nn.Sequential(normalization(channels), nn.SiLU(), conv_nd(dims, channels, self.out_channels, 3, padding=1))
    self.updown = (up or down)
    if up:
        self.h_upd = Upsample(channels, False, dims)
        self.x_upd = Upsample(channels, False, dims)
    elif down:
        self.h_upd = Downsample(channels, False, dims)
        self.x_upd = Downsample(channels, False, dims)
    else:
        self.h_upd = self.x_upd = nn.Identity()
    self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, ((2 * self.out_channels) if use_scale_shift_norm else self.out_channels)))
    self.out_layers = nn.Sequential(normalization(self.out_channels), nn.SiLU(), nn.Dropout(p=dropout), zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)))
    if (self.out_channels == channels):
        self.skip_connection = nn.Identity()
    elif use_conv:
        self.skip_connection = conv_nd(dims, channels, self.out_channels, 3, padding=1)
    else:
        self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
def __init__(self, channels, emb_channels, dropout, out_channels=None, use_conv=False, use_scale_shift_norm=False, dims=2, use_checkpoint=False, up=False, down=False):
    super().__init__()
    self.channels = channels
    self.emb_channels = emb_channels
    self.dropout = dropout
    self.out_channels = (out_channels or channels)
    self.use_conv = use_conv
    self.use_checkpoint = use_checkpoint
    self.use_scale_shift_norm = use_scale_shift_norm
    self.in_layers = nn.Sequential(normalization(channels), nn.SiLU(), conv_nd(dims, channels, self.out_channels, 3, padding=1))
    self.updown = (up or down)
    if up:
        self.h_upd = Upsample(channels, False, dims)
        self.x_upd = Upsample(channels, False, dims)
    elif down:
        self.h_upd = Downsample(channels, False, dims)
        self.x_upd = Downsample(channels, False, dims)
    else:
        self.h_upd = self.x_upd = nn.Identity()
    self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, ((2 * self.out_channels) if use_scale_shift_norm else self.out_channels)))
    self.out_layers = nn.Sequential(normalization(self.out_channels), nn.SiLU(), nn.Dropout(p=dropout), zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)))
    if (self.out_channels == channels):
        self.skip_connection = nn.Identity()
    elif use_conv:
        self.skip_connection = conv_nd(dims, channels, self.out_channels, 3, padding=1)
    else:
        self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)