| 51 | self.vace_patch_embedding = torch.nn.Conv3d(vace_in_dim, dim, kernel_size=patch_size, stride=patch_size) |
| 52 | |
| 53 | def forward(self, x, vace_context, context, t_mod, freqs): |
| 54 | c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context] |
| 55 | c = [u.flatten(2).transpose(1, 2) for u in c] |
| 56 | c = torch.cat([ |
| 57 | torch.cat([u, u.new_zeros(1, x.shape[1] - u.size(1), u.size(2))], |
| 58 | dim=1) for u in c |
| 59 | ]) |
| 60 | |
| 61 | for block in self.vace_blocks: |
| 62 | c = block(c, x, context, t_mod, freqs) |
| 63 | hints = torch.unbind(c)[:-1] |
| 64 | return hints |
| 65 | |
| 66 | @staticmethod |
| 67 | def state_dict_converter(): |