    def forward(self, x, emb):
        """
        Apply the block to a Tensor, conditioned on a timestep embedding.

            x: an [N x C x ...] Tensor of features.
            emb: an [N x emb_channels x ...] Tensor of timestep embeddings.
        :return: an [N x C x ...] Tensor of outputs.
        """
        if self.updown:
            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
            h = in_rest(x)
            h = self.h_upd(h)
            x = self.x_upd(x)
            h = in_conv(h)
        else:
            h = self.in_layers(x)
        emb_out = self.emb_layers(emb)
        emb_out = append_dims(emb_out, h.dim())
        if self.use_scale_shift_norm:
            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
            scale, shift = torch.chunk(emb_out, 2, dim=1)
            h = out_norm(h) * (1 + scale) + shift
            h = out_rest(h)
        else:
            h = h + emb_out
            h = self.out_layers(h)
        return self.skip_connection(x) + h
Apply the block to a Tensor, conditioned on a timestep embedding.

x: an [N x C x ...] Tensor of features.
emb: an [N x emb_channels x ...] Tensor of timestep embeddings.
:return: an [N x C x ...] Tensor of outputs.
modules/diffusion/unet/resblock.py