File size: 2,810 Bytes
742c169 930b76b 742c169 930b76b 742c169 930b76b 742c169 930b76b 742c169 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Hugging Face ``transformers`` wrapper around the TT-VidT encoder.
``scripts/tools/export_hf.py`` copies this file (as ``modeling_ttvidt.py``) and the
encoder sources next to an exported model, so the same directory / Hub repo loads with
from transformers import AutoModel
model = AutoModel.from_pretrained("<repo>", trust_remote_code=True)
motion = model(video).motion_output # video: [B, T, 3, 256, 256] in [-1, 1]
and with ``ttvidt.hub.load_model("<repo>")``: ``config.json`` holds the release format
(``format``, ``model`` = TTVidTrainer arguments) plus the ``transformers`` fields, and
``model.safetensors`` uses the TTVidTrainer key names (``encoder.*``).
"""
from typing import ClassVar
from transformers import PretrainedConfig, PreTrainedModel
from .dinov3_vit import DINOv3VidTConfig, DINOv3VidTModel
class TTVidTConfig(PretrainedConfig):
"""``model``: the TTVidTrainer keyword arguments of the release (``backbone_config``
holds the full DINOv3VidTConfig, so no DINOv3 download is needed)."""
model_type = "ttvidt"
def __init__(
self,
format: str | None = None,
model: dict | None = None,
encoder_only: bool = False,
**kwargs,
):
self.format = format
self.model = model or {}
self.encoder_only = encoder_only
super().__init__(**kwargs)
def encoder_config(self) -> DINOv3VidTConfig:
if self.model.get("backbone_arch", "ttvidt") != "ttvidt":
raise ValueError(
f"only TT-VidT encoders are supported, got {self.model.get('backbone_arch')!r}"
)
cfg = dict(self.model["backbone_config"])
for k, v in (self.model.get("motion_encoder_config") or {}).items():
cfg.setdefault(k, v) # same merge as TTVidTrainer
return DINOv3VidTConfig(**cfg)
class TTVidTModel(PreTrainedModel):
"""The TT-VidT video encoder (``DINOv3VidTModel``); ``forward`` returns its output
(``motion_output``: [B, T, K, D] motion tokens per frame)."""
config_class = TTVidTConfig
base_model_prefix = "ttvidt"
main_input_name = "pixel_values"
# the decoder and the frame-VAE latent statistics of full releases are not used
_keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [
r"^decoder\.",
r"^latent_(mean|std)$",
]
def __init__(self, config: TTVidTConfig):
super().__init__(config)
self.encoder = DINOv3VidTModel(config.encoder_config())
self.post_init()
def _init_weights(self, module):
pass # DINOv3VidTModel initialises itself; pretrained weights come from the file
def forward(self, pixel_values, **kwargs):
"""pixel_values: [B, T, 3, H, W] in [-1, 1]."""
return self.encoder(pixel_values, **kwargs)
|