Run a tiny U-Net denoiser on a batch of images and timesteps.
eu=lin2(SiLU(lin1(emb(t)))),h1=Rdown(conv_in(x),e)=up(Rmid(down(h1),e)),out=conv_out(SiLU(GN(Rup([u,h1],e))))Here emb(t) lists the cosines and then the sines of tfi with fi=10000−2i/E for i below E / 2. Each lin is a linear layer, every conv except the 1x1 skip is 3x3 with padding 1, and down has stride 2. The up step repeats each pixel 2x2, then applies its conv. [u,h1] joins channels with u first. Each R is GN, SiLU, conv1, plus the projection of SiLU(e), then GN, SiLU, conv2, plus a skip (a 1x1 conv if its skip_weight exists, else the input). Every GN has num_groups groups and eps 1e-5.
| keys | shapes |
|---|---|
| time_linear1_weight, time_linear1_bias | (T, E), (T,) |
| time_linear2_weight, time_linear2_bias | (T, T), (T,) |
| conv_in_weight, conv_in_bias | (C, C_in, 3, 3), (C,) |
| downsample_*, upsample_* | (C, C, 3, 3), (C,) |
| P_norm1_*, P_conv1_* | (K,), (C, K, 3, 3) and (C,) |
| P_time_weight, P_time_bias | (C, T), (C,) |
| P_norm2_*, P_conv2_* | (C,), (C, C, 3, 3) and (C,) |
| up_res_skip_weight, up_res_skip_bias | (C, 2C, 1, 1), (C,) |
| out_norm_*, conv_out_* | (C,), (C_out, C, 3, 3) and (C_out,) |
Each star stands for weight and bias. P is down_res, mid_res or up_res, with K = C, C and 2C. Return a PyTorch tensor of shape (B, C_out, H, W) with the dtype of x.
Input: x = [[[[-2.472,0.9,-0.625,0.914],[0.653,-1.959,1.286,1.071],[1.099,0.303,0.519,0.639],[1.021,0.924,-0.373,-0.135]]]], timesteps = [10], params = {46 tensors, C_in=1, C=2, E=4, T=4, C_out=1; values in the test case}, num_groups = 1
Output: [[[[0.173095,0.133803,0.209918,0.032938],[0.920273,0.109064,-0.202224,-0.039853],[0.060771,-0.149412,-0.246185,-0.407389],[0.291618,-0.07764,-0.03258,0.106188]]]]
Explanation: The 1x4x4 image is lifted to 2 channels, halved to 2x2 by the strided convolution, upsampled back to 4x4, joined with the 2 encoder channels into 4 channels and projected to 1 output channel. Timestep 10 enters every residual block through the time MLP.
Input: x = [[[[-0.414,0.108],[0.542,0.966]]],[[[0.525,-0.057],[0.714,0.92]]]], timesteps = [0,999], params = {46 tensors, C_in=1, C=2, E=2, T=4, C_out=1; values in the test case}, num_groups = 2
Output: [[[[-0.084252,0.668281],[0.307617,0.6261]]],[[[-0.356636,0.737723],[0.08915,0.923358]]]]
Input: x = [[[[-0.254,-0.791,-1.184,-0.345],[-1.762,1.166,0.514,0.976]],[[0.588,0.5,0.835,-1.168],[0.479,-0.329,-0.433,0.593]]]], timesteps = [500], params = {46 tensors, C_in=2, C=2, E=4, T=4, C_out=2; values in the test case}, num_groups = 1
Output: [[[[-0.151334,0.247716,-0.374217,0.249671],[0.988048,-0.733239,0.497515,0.286534]],[[-0.422448,-1.006463,-0.727194,-0.265752],[0.56916,-0.451766,0.857516,-0.556124]]]]
F.conv2d(h, w, b, stride=2, padding=1) halves H and W.
torch.cat([u, h1], dim=1) joins channels with u first.
params['time_linear1_weight'].shape[1] gives E.
Accepts: array
Accepts: array
Accepts: object
Accepts: integer
Run a tiny U-Net denoiser on a batch of images and timesteps.
eu=lin2(SiLU(lin1(emb(t)))),h1=Rdown(conv_in(x),e)=up(Rmid(down(h1),e)),out=conv_out(SiLU(GN(Rup([u,h1],e))))Here emb(t) lists the cosines and then the sines of tfi with fi=10000−2i/E for i below E / 2. Each lin is a linear layer, every conv except the 1x1 skip is 3x3 with padding 1, and down has stride 2. The up step repeats each pixel 2x2, then applies its conv. [u,h1] joins channels with u first. Each R is GN, SiLU, conv1, plus the projection of SiLU(e), then GN, SiLU, conv2, plus a skip (a 1x1 conv if its skip_weight exists, else the input). Every GN has num_groups groups and eps 1e-5.
| keys | shapes |
|---|---|
| time_linear1_weight, time_linear1_bias | (T, E), (T,) |
| time_linear2_weight, time_linear2_bias | (T, T), (T,) |
| conv_in_weight, conv_in_bias | (C, C_in, 3, 3), (C,) |
| downsample_*, upsample_* | (C, C, 3, 3), (C,) |
| P_norm1_*, P_conv1_* | (K,), (C, K, 3, 3) and (C,) |
| P_time_weight, P_time_bias | (C, T), (C,) |
| P_norm2_*, P_conv2_* | (C,), (C, C, 3, 3) and (C,) |
| up_res_skip_weight, up_res_skip_bias | (C, 2C, 1, 1), (C,) |
| out_norm_*, conv_out_* | (C,), (C_out, C, 3, 3) and (C_out,) |
Each star stands for weight and bias. P is down_res, mid_res or up_res, with K = C, C and 2C. Return a PyTorch tensor of shape (B, C_out, H, W) with the dtype of x.
Input: x = [[[[-2.472,0.9,-0.625,0.914],[0.653,-1.959,1.286,1.071],[1.099,0.303,0.519,0.639],[1.021,0.924,-0.373,-0.135]]]], timesteps = [10], params = {46 tensors, C_in=1, C=2, E=4, T=4, C_out=1; values in the test case}, num_groups = 1
Output: [[[[0.173095,0.133803,0.209918,0.032938],[0.920273,0.109064,-0.202224,-0.039853],[0.060771,-0.149412,-0.246185,-0.407389],[0.291618,-0.07764,-0.03258,0.106188]]]]
Explanation: The 1x4x4 image is lifted to 2 channels, halved to 2x2 by the strided convolution, upsampled back to 4x4, joined with the 2 encoder channels into 4 channels and projected to 1 output channel. Timestep 10 enters every residual block through the time MLP.
Input: x = [[[[-0.414,0.108],[0.542,0.966]]],[[[0.525,-0.057],[0.714,0.92]]]], timesteps = [0,999], params = {46 tensors, C_in=1, C=2, E=2, T=4, C_out=1; values in the test case}, num_groups = 2
Output: [[[[-0.084252,0.668281],[0.307617,0.6261]]],[[[-0.356636,0.737723],[0.08915,0.923358]]]]
Input: x = [[[[-0.254,-0.791,-1.184,-0.345],[-1.762,1.166,0.514,0.976]],[[0.588,0.5,0.835,-1.168],[0.479,-0.329,-0.433,0.593]]]], timesteps = [500], params = {46 tensors, C_in=2, C=2, E=4, T=4, C_out=2; values in the test case}, num_groups = 1
Output: [[[[-0.151334,0.247716,-0.374217,0.249671],[0.988048,-0.733239,0.497515,0.286534]],[[-0.422448,-1.006463,-0.727194,-0.265752],[0.56916,-0.451766,0.857516,-0.556124]]]]
F.conv2d(h, w, b, stride=2, padding=1) halves H and W.
torch.cat([u, h1], dim=1) joins channels with u first.
params['time_linear1_weight'].shape[1] gives E.
Accepts: array
Accepts: array
Accepts: object
Accepts: integer
