A sequential module that passes timestep embeddings to the children that support it as an extra input.
| 135 | |
| 136 | |
| 137 | class TimestepEmbedSequential(nn.Sequential, TimestepBlock): |
| 138 | """ |
| 139 | A sequential module that passes timestep embeddings to the children that |
| 140 | support it as an extra input. |
| 141 | """ |
| 142 | |
| 143 | def forward(self, x, emb): |
| 144 | for layer in self: |
| 145 | if isinstance(layer, TimestepBlock): |
| 146 | x = layer(x, emb) |
| 147 | else: |
| 148 | x = layer(x) |
| 149 | return x |
| 150 | |
| 151 | |
| 152 | class Upsample(nn.Module): |