(self, x)
| 139 | return x # [n, l, d] |
| 140 | |
| 141 | def unpatchify(self, x): |
| 142 | bsz = x.shape[0] |
| 143 | p = self.patch_size |
| 144 | c = self.vae_embed_dim |
| 145 | h_, w_ = self.seq_h, self.seq_w |
| 146 | |
| 147 | x = x.reshape(bsz, h_, w_, c, p, p) |
| 148 | x = torch.einsum('nhwcpq->nchpwq', x) |
| 149 | x = x.reshape(bsz, c, h_ * p, w_ * p) |
| 150 | return x # [n, c, h, w] |
| 151 | |
| 152 | def sample_orders(self, bsz): |
| 153 | # generate a batch of random generation orders |