| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| --- /usr/local/lib/python3.12/dist-packages/cosmos_predict2/tokenizers/tokenizer.py | |
| +++ /usr/local/lib/python3.12/dist-packages/cosmos_predict2/tokenizers/tokenizer.py | |
| def count_param(self): | |
| return sum(p.numel() for p in self.model.parameters()) | |
| - @torch.no_grad() | |
| def encode(self, videos): | |
| """ | |
| videos: A list of videos each with shape [C, T, H, W]. | |
| latent = latent.to(in_dtype) | |
| return latent | |
| - @torch.no_grad() | |
| def decode(self, zs): | |
| in_dtype = zs.dtype | |
| with self.context: | |
| ) | |
| return num_chunks | |
| - @torch.no_grad() | |
| def encode(self, state: torch.Tensor) -> torch.Tensor: | |
| if not self.squeeze_for_image: | |
| assert state.shape[2] == 1, f"Expect the number of frames {state.shape[2]} to be 1" | |
| latent = latent.to(in_dtype) | |
| return latent | |
| - @torch.no_grad() | |
| def decode(self, state: torch.Tensor) -> torch.Tensor: | |
| assert state.shape[2] == 1, f"Expect the number of frames {state.shape[2]} to be 1" | |
| in_dtype = state.dtype | |