Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import torch | |
| from .padding import pad_video_frames | |
| class TemporalExpansion: | |
| video: torch.Tensor | |
| source_frames: int | |
| effective_frames: int | |
| expanded_frames: int | |
| repeat_factor: int | |
| def expand_for_temporal_compression( | |
| video: torch.Tensor, | |
| temporal_compression: int, | |
| *, | |
| model_compression: int = 4, | |
| ) -> TemporalExpansion: | |
| """Expand an input timeline onto the model's fixed temporal grid. | |
| The first frame remains unique. Every later frame is repeated | |
| ``model_compression // temporal_compression`` times. | |
| """ | |
| if video.ndim != 5: | |
| raise ValueError("video must be [B,C,F,H,W]") | |
| if temporal_compression not in {1, 2, 4}: | |
| raise ValueError("temporal_compression must be one of 1, 2, or 4") | |
| if model_compression <= 0: | |
| raise ValueError("model_compression must be positive") | |
| if model_compression % temporal_compression: | |
| raise ValueError("model_compression must be divisible by temporal_compression") | |
| source_frames = video.shape[2] | |
| # Keep every requested source frame. The model still requires an aligned | |
| # temporal grid, so repeat the final frame only for model input and crop | |
| # the decoded result back to ``source_frames`` in the pipeline. | |
| video, _ = pad_video_frames(video, temporal_compression, mode="repeat") | |
| effective_frames = video.shape[2] | |
| repeat_factor = model_compression // temporal_compression | |
| if repeat_factor == 1 or effective_frames == 1: | |
| expanded = video | |
| else: | |
| expanded = torch.cat( | |
| [video[:, :, :1], video[:, :, 1:].repeat_interleave(repeat_factor, dim=2)], | |
| dim=2, | |
| ) | |
| return TemporalExpansion( | |
| video=expanded, | |
| source_frames=source_frames, | |
| effective_frames=effective_frames, | |
| expanded_frames=expanded.shape[2], | |
| repeat_factor=repeat_factor, | |
| ) | |
| def _gaussian_weights( | |
| count: int, | |
| sigma: float, | |
| *, | |
| device: torch.device, | |
| ) -> torch.Tensor: | |
| positions = torch.arange(count, device=device, dtype=torch.float32) | |
| positions = positions - (count - 1) / 2 | |
| weights = torch.exp(-0.5 * (positions / sigma).square()) | |
| return weights / weights.sum() | |
| def merge_temporal_output( | |
| video: torch.Tensor, | |
| expansion: TemporalExpansion, | |
| *, | |
| mode: str = "gaussian", | |
| gaussian_sigma: float = 1.0, | |
| ) -> torch.Tensor: | |
| """Collapse decoded duplicate groups back to the effective input length.""" | |
| if video.ndim != 5: | |
| raise ValueError("video must be [B,C,F,H,W]") | |
| mode = mode.lower() | |
| if mode not in {"gaussian", "first"}: | |
| raise ValueError("temporal merge mode must be 'gaussian' or 'first'") | |
| if gaussian_sigma <= 0: | |
| raise ValueError("gaussian_sigma must be positive") | |
| required = expansion.expanded_frames | |
| if video.shape[2] < required: | |
| raise ValueError( | |
| f"decoded video has {video.shape[2]} frames, expected at least {required}" | |
| ) | |
| video = video[:, :, :required] | |
| factor = expansion.repeat_factor | |
| if factor == 1 or expansion.effective_frames == 1: | |
| return video[:, :, : expansion.effective_frames] | |
| first = video[:, :, :1] | |
| tail = video[:, :, 1:].reshape( | |
| video.shape[0], | |
| video.shape[1], | |
| expansion.effective_frames - 1, | |
| factor, | |
| video.shape[3], | |
| video.shape[4], | |
| ) | |
| if mode == "first": | |
| merged = tail[:, :, :, 0] | |
| else: | |
| weights = _gaussian_weights(factor, gaussian_sigma, device=video.device) | |
| merged = (tail.float() * weights.view(1, 1, 1, factor, 1, 1)).sum( | |
| dim=3 | |
| ).to(video.dtype) | |
| return torch.cat([first, merged], dim=2) | |