| 393 | self.abondon_cond = abondon_cond |
| 394 | |
| 395 | def forward(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None, |
| 396 | uncondi_length = None, uncondi_pe = None, mask_uncond = None) -> Tensor: |
| 397 | if self.abondon_cond: |
| 398 | x = [ix[:u_l, :] for ix, u_l in zip(x, uncondi_length)] |
| 399 | x = pad_sequence(x, batch_first=True) |
| 400 | if not x.shape[1] == pe.shape[2]: |
| 401 | pe = uncondi_pe |
| 402 | mask = mask_uncond |
| 403 | # print("single stream block", x.shape, pe.shape) |
| 404 | x = super().forward(x, vec, pe, mask) |
| 405 | return x |
| 406 | |
| 407 | |
| 408 | class DoubleStreamBlockD(DoubleStreamBlock): |