(self, quantized: Tensor, cond: Optional[Tensor] = None, video_contains_first_frame=True)
| 1448 | |
| 1449 | @beartype |
| 1450 | def decode(self, quantized: Tensor, cond: Optional[Tensor] = None, video_contains_first_frame=True): |
| 1451 | decode_first_frame_separately = self.separate_first_frame_encoding and video_contains_first_frame |
| 1452 | |
| 1453 | batch = quantized.shape[0] |
| 1454 | |
| 1455 | # conditioning, if needed |
| 1456 | |
| 1457 | assert (not self.has_cond) or exists( |
| 1458 | cond |
| 1459 | ), "`cond` must be passed into tokenizer forward method since conditionable layers were specified" |
| 1460 | |
| 1461 | if exists(cond): |
| 1462 | assert cond.shape == (batch, self.dim_cond) |
| 1463 | |
| 1464 | cond = self.decoder_cond_in(cond) |
| 1465 | cond_kwargs = dict(cond=cond) |
| 1466 | |
| 1467 | # decoder layers |
| 1468 | |
| 1469 | x = quantized |
| 1470 | |
| 1471 | for fn, has_cond in zip(self.decoder_layers, reversed(self.has_cond_across_layers)): |
| 1472 | layer_kwargs = dict() |
| 1473 | |
| 1474 | if has_cond: |
| 1475 | layer_kwargs = cond_kwargs |
| 1476 | |
| 1477 | x = fn(x, **layer_kwargs) |
| 1478 | |
| 1479 | # to pixels |
| 1480 | |
| 1481 | if decode_first_frame_separately: |
| 1482 | left_pad, xff, x = ( |
| 1483 | x[:, :, : self.time_padding], |
| 1484 | x[:, :, self.time_padding], |
| 1485 | x[:, :, (self.time_padding + 1) :], |
| 1486 | ) |
| 1487 | |
| 1488 | out = self.conv_out(x) |
| 1489 | outff = self.conv_out_first_frame(xff) |
| 1490 | |
| 1491 | video, _ = pack([outff, out], "b c * h w") |
| 1492 | |
| 1493 | else: |
| 1494 | video = self.conv_out(x) |
| 1495 | |
| 1496 | # if video were padded, remove padding |
| 1497 | |
| 1498 | if video_contains_first_frame: |
| 1499 | video = video[:, :, self.time_padding :] |
| 1500 | |
| 1501 | return video |
| 1502 | |
| 1503 | @torch.no_grad() |
| 1504 | def tokenize(self, video): |
no test coverage detected