codecs/base.py

from abc import ABC, abstractmethod

import torch


class VideoTokenizer(ABC):
    @abstractmethod
    def encode(self, video: torch.Tensor) -> list[torch.Tensor]:
        ...

    @abstractmethod
    def decode(self, tokens: list[torch.Tensor]) -> torch.Tensor:
        ...

    @property
    @abstractmethod
    def num_layers(self) -> int:
        ...

    @property
    @abstractmethod
    def layer_token_counts(self) -> list[int]:
        ...

    @property
    @abstractmethod
    def name(self) -> str:
        ...


class AudioTokenizer(ABC):
    @abstractmethod
    def encode(self, audio: torch.Tensor) -> torch.Tensor:
        ...

    @abstractmethod
    def decode(self, tokens: torch.Tensor) -> torch.Tensor:
        ...

    @property
    @abstractmethod
    def sample_rate(self) -> int:
        ...

    @property
    @abstractmethod
    def num_codebooks(self) -> int:
        ...

    @property
    @abstractmethod
    def name(self) -> str:
        ...