| 218 | |
| 219 | |
| 220 | class PopBlock(torch.nn.Module): |
| 221 | def __init__(self): |
| 222 | super().__init__() |
| 223 | |
| 224 | def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): |
| 225 | res_hidden_states = res_stack.pop() |
| 226 | hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) |
| 227 | return hidden_states, time_emb, text_emb, res_stack |
| 228 | |
| 229 | |
| 230 | class SDUNet(torch.nn.Module): |