(self, video: Tensor, quantize=False, cond: Optional[Tensor] = None, video_contains_first_frame=True)
| 1381 | |
| 1382 | @beartype |
| 1383 | def encode(self, video: Tensor, quantize=False, cond: Optional[Tensor] = None, video_contains_first_frame=True): |
| 1384 | encode_first_frame_separately = self.separate_first_frame_encoding and video_contains_first_frame |
| 1385 | |
| 1386 | # whether to pad video or not |
| 1387 | |
| 1388 | if video_contains_first_frame: |
| 1389 | video_len = video.shape[2] |
| 1390 | |
| 1391 | video = pad_at_dim(video, (self.time_padding, 0), value=0.0, dim=2) |
| 1392 | video_packed_shape = [torch.Size([self.time_padding]), torch.Size([]), torch.Size([video_len - 1])] |
| 1393 | |
| 1394 | # conditioning, if needed |
| 1395 | |
| 1396 | assert (not self.has_cond) or exists( |
| 1397 | cond |
| 1398 | ), "`cond` must be passed into tokenizer forward method since conditionable layers were specified" |
| 1399 | |
| 1400 | if exists(cond): |
| 1401 | assert cond.shape == (video.shape[0], self.dim_cond) |
| 1402 | |
| 1403 | cond = self.encoder_cond_in(cond) |
| 1404 | cond_kwargs = dict(cond=cond) |
| 1405 | |
| 1406 | # initial conv |
| 1407 | # taking into account whether to encode first frame separately |
| 1408 | |
| 1409 | if encode_first_frame_separately: |
| 1410 | pad, first_frame, video = unpack(video, video_packed_shape, "b c * h w") |
| 1411 | first_frame = self.conv_in_first_frame(first_frame) |
| 1412 | |
| 1413 | video = self.conv_in(video) |
| 1414 | |
| 1415 | if encode_first_frame_separately: |
| 1416 | video, _ = pack([first_frame, video], "b c * h w") |
| 1417 | video = pad_at_dim(video, (self.time_padding, 0), dim=2) |
| 1418 | |
| 1419 | # encoder layers |
| 1420 | |
| 1421 | for fn, has_cond in zip(self.encoder_layers, self.has_cond_across_layers): |
| 1422 | layer_kwargs = dict() |
| 1423 | |
| 1424 | if has_cond: |
| 1425 | layer_kwargs = cond_kwargs |
| 1426 | |
| 1427 | video = fn(video, **layer_kwargs) |
| 1428 | |
| 1429 | maybe_quantize = identity if not quantize else self.quantizers |
| 1430 | |
| 1431 | return maybe_quantize(video) |
| 1432 | |
| 1433 | @beartype |
| 1434 | def decode_from_code_indices(self, codes: Tensor, cond: Optional[Tensor] = None, video_contains_first_frame=True): |
no test coverage detected