| 1 | import itertools |
| 2 | from dataclasses import dataclass |
| 3 | from typing import Callable, List, NamedTuple, Tuple |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | |
| 8 | def compute_trapezoidal_mask_1d( |
| 9 | length: int, |
| 10 | ramp_left: int, |
| 11 | ramp_right: int, |
| 12 | left_starts_from_0: bool = False, |
| 13 | ) -> torch.Tensor: |
| 14 | """ |
| 15 | Generate a 1D trapezoidal blending mask with linear ramps. |
| 16 | Args: |
| 17 | length: Output length of the mask. |
| 18 | ramp_left: Fade-in length on the left. |
| 19 | ramp_right: Fade-out length on the right. |
| 20 | left_starts_from_0: Whether the ramp starts from 0 or first non-zero value. |
| 21 | Useful for temporal tiles where the first tile is causal. |
| 22 | Returns: |
| 23 | A 1D tensor of shape `(length,)` with values in [0, 1]. |
| 24 | """ |
| 25 | if length <= 0: |
| 26 | raise ValueError("Mask length must be positive.") |
| 27 | |
| 28 | ramp_left = max(0, min(ramp_left, length)) |
| 29 | ramp_right = max(0, min(ramp_right, length)) |
| 30 | |
| 31 | mask = torch.ones(length) |
| 32 | |
| 33 | if ramp_left > 0: |
| 34 | interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2 |
| 35 | fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1] |
| 36 | if not left_starts_from_0: |
| 37 | fade_in = fade_in[1:] |
| 38 | mask[:ramp_left] *= fade_in |
| 39 | |
| 40 | if ramp_right > 0: |
| 41 | fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1] |
| 42 | mask[-ramp_right:] *= fade_out |
| 43 | |
| 44 | return mask.clamp_(0, 1) |
| 45 | |
| 46 | |
| 47 | def compute_rectangular_mask_1d( |
| 48 | length: int, |
| 49 | left_ramp: int, |
| 50 | right_ramp: int, |
| 51 | ) -> torch.Tensor: |
| 52 | """ |
| 53 | Generate a 1D rectangular (pulse) mask. |
| 54 | Args: |
| 55 | length: Output length of the mask. |
| 56 | left_ramp: Number of elements at the start of the mask to set to 0. |
| 57 | right_ramp: Number of elements at the end of the mask to set to 0. |
| 58 | Returns: |
| 59 | A 1D tensor of shape `(length,)` with values 0 or 1. |
| 60 | """ |
| 61 | if length <= 0: |
| 62 | raise ValueError("Mask length must be positive.") |
| 63 | |
| 64 | mask = torch.ones(length) |
| 65 | if left_ramp > 0: |
| 66 | mask[:left_ramp] = 0 |
| 67 | if right_ramp > 0: |
| 68 | mask[-right_ramp:] = 0 |
| 69 | return mask |
| 70 | |
| 71 | |
| 72 | @dataclass(frozen=True) |
| 73 | class SpatialTilingConfig: |
| 74 | """Configuration for dividing each frame into spatial tiles with optional overlap. |
| 75 | Args: |
| 76 | tile_size_in_pixels (int): Size of each tile in pixels. Must be at least 64 and divisible by 32. |
| 77 | tile_overlap_in_pixels (int, optional): Overlap between tiles in pixels. Must be divisible by 32. Defaults to 0. |
| 78 | """ |
| 79 | |
| 80 | tile_size_in_pixels: int |
| 81 | tile_overlap_in_pixels: int = 0 |
| 82 | |
| 83 | def __post_init__(self) -> None: |
| 84 | if self.tile_size_in_pixels < 64: |
| 85 | raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}") |
| 86 | if self.tile_size_in_pixels % 32 != 0: |
| 87 | raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}") |
| 88 | if self.tile_overlap_in_pixels % 32 != 0: |
| 89 | raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}") |
| 90 | if self.tile_overlap_in_pixels >= self.tile_size_in_pixels: |
| 91 | raise ValueError( |
| 92 | f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}" |
| 93 | ) |
| 94 | |
| 95 | |
| 96 | @dataclass(frozen=True) |
| 97 | class TemporalTilingConfig: |
| 98 | """Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap. |
| 99 | Args: |
| 100 | tile_size_in_frames (int): Number of frames in each tile. Must be at least 16 and divisible by 8. |
| 101 | tile_overlap_in_frames (int, optional): Number of overlapping frames between consecutive tiles. |
| 102 | Must be divisible by 8. Defaults to 0. |
| 103 | """ |
| 104 | |
| 105 | tile_size_in_frames: int |
| 106 | tile_overlap_in_frames: int = 0 |
| 107 | |
| 108 | def __post_init__(self) -> None: |
| 109 | if self.tile_size_in_frames < 16: |
| 110 | raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}") |
| 111 | if self.tile_size_in_frames % 8 != 0: |
| 112 | raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}") |
| 113 | if self.tile_overlap_in_frames % 8 != 0: |
| 114 | raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}") |
| 115 | if self.tile_overlap_in_frames >= self.tile_size_in_frames: |
| 116 | raise ValueError( |
| 117 | f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}" |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | @dataclass(frozen=True) |
| 122 | class TilingConfig: |
| 123 | """Configuration for splitting video into tiles with optional overlap. |
| 124 | Attributes: |
| 125 | spatial_config: Configuration for splitting spatial dimensions into tiles. |
| 126 | temporal_config: Configuration for splitting temporal dimension into tiles. |
| 127 | """ |
| 128 | |
| 129 | spatial_config: SpatialTilingConfig | None = None |
| 130 | temporal_config: TemporalTilingConfig | None = None |
| 131 | |
| 132 | @classmethod |
| 133 | def default(cls) -> "TilingConfig": |
| 134 | return cls( |
| 135 | spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64), |
| 136 | temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24), |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | @dataclass(frozen=True) |
| 141 | class DimensionIntervals: |
| 142 | """Defines how a single dimension is split into overlapping intervals (tiles). |
| 143 | Each list has length N where N is the number of intervals. The i-th element |
| 144 | of each list describes the i-th interval. |
| 145 | Attributes: |
| 146 | starts: Start index of each interval (inclusive). |
| 147 | ends: End index of each interval (exclusive). |
| 148 | left_ramps: Length of the left blend ramp for each interval. |
| 149 | Used to create masks that fade in from 0 to 1. |
| 150 | right_ramps: Length of the right blend ramp for each interval. |
| 151 | Used to create masks that fade out from 1 to 0. |
| 152 | """ |
| 153 | |
| 154 | starts: List[int] |
| 155 | ends: List[int] |
| 156 | left_ramps: List[int] |
| 157 | right_ramps: List[int] |
| 158 | |
| 159 | |
| 160 | @dataclass(frozen=True) |
| 161 | class TensorTilingSpec: |
| 162 | """Specifies how a tensor of a given shape is split into intervals (tiles) along each dimension. |
| 163 | Attributes: |
| 164 | original_shape: Shape of the tensor being tiled. |
| 165 | dimension_intervals: Per-dimension intervals (starts, ends, ramps) for each axis. |
| 166 | """ |
| 167 | |
| 168 | original_shape: torch.Size |
| 169 | dimension_intervals: Tuple[DimensionIntervals, ...] |
| 170 | |
| 171 | |
| 172 | # Operation to split a single dimension of the tensor into intervals based on the length along the dimension. |
| 173 | SplitOperation = Callable[[int], DimensionIntervals] |
| 174 | # Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension. |
| 175 | MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]] |
| 176 | |
| 177 | |
| 178 | def default_split_operation(length: int) -> DimensionIntervals: |
| 179 | return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0]) |
| 180 | |
| 181 | |
| 182 | DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation |
| 183 | |
| 184 | |
| 185 | def default_mapping_operation( |
| 186 | _intervals: DimensionIntervals, |
| 187 | ) -> tuple[list[slice], list[torch.Tensor | None]]: |
| 188 | return [slice(0, None)], [None] |
| 189 | |
| 190 | |
| 191 | DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation |
| 192 | |
| 193 | |
| 194 | class Tile(NamedTuple): |
| 195 | """ |
| 196 | Represents a single tile. |
| 197 | Attributes: |
| 198 | in_coords: |
| 199 | Tuple of slices specifying where to cut the tile from the INPUT tensor. |
| 200 | out_coords: |
| 201 | Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor. |
| 202 | masks_1d: |
| 203 | Per-dimension masks in OUTPUT units. |
| 204 | These are used to create all-dimensional blending mask. |
| 205 | Methods: |
| 206 | blend_mask: |
| 207 | Create a single N-D mask from the per-dimension masks. |
| 208 | """ |
| 209 | |
| 210 | in_coords: Tuple[slice, ...] |
| 211 | out_coords: Tuple[slice, ...] |
| 212 | masks_1d: Tuple[Tuple[torch.Tensor, ...]] |
| 213 | |
| 214 | @property |
| 215 | def blend_mask(self) -> torch.Tensor: |
| 216 | num_dims = len(self.out_coords) |
| 217 | per_dimension_masks: List[torch.Tensor] = [] |
| 218 | |
| 219 | for dim_idx in range(num_dims): |
| 220 | mask_1d = self.masks_1d[dim_idx] |
| 221 | view_shape = [1] * num_dims |
| 222 | if mask_1d is None: |
| 223 | # Broadcast mask along this dimension (length 1). |
| 224 | one = torch.ones(1) |
| 225 | |
| 226 | view_shape[dim_idx] = 1 |
| 227 | per_dimension_masks.append(one.view(*view_shape)) |
| 228 | continue |
| 229 | |
| 230 | # Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply. |
| 231 | view_shape[dim_idx] = mask_1d.shape[0] |
| 232 | per_dimension_masks.append(mask_1d.view(*view_shape)) |
| 233 | |
| 234 | # Multiply per-dimension masks to form the full N-D mask (separable blending window). |
| 235 | combined_mask = per_dimension_masks[0] |
| 236 | for mask in per_dimension_masks[1:]: |
| 237 | combined_mask = combined_mask * mask |
| 238 | |
| 239 | return combined_mask |
| 240 | |
| 241 | |
| 242 | def create_tiles_from_intervals_and_mappers( |
| 243 | intervals: TensorTilingSpec, |
| 244 | mappers: List[MappingOperation], |
| 245 | ) -> List[Tile]: |
| 246 | full_dim_input_slices = [] |
| 247 | full_dim_output_slices = [] |
| 248 | full_dim_masks_1d = [] |
| 249 | for axis_index in range(len(intervals.original_shape)): |
| 250 | dimension_intervals = intervals.dimension_intervals[axis_index] |
| 251 | starts = dimension_intervals.starts |
| 252 | ends = dimension_intervals.ends |
| 253 | input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)] |
| 254 | output_slices, masks_1d = mappers[axis_index](dimension_intervals) |
| 255 | full_dim_input_slices.append(input_slices) |
| 256 | full_dim_output_slices.append(output_slices) |
| 257 | full_dim_masks_1d.append(masks_1d) |
| 258 | |
| 259 | tiles = [] |
| 260 | tile_in_coords = list(itertools.product(*full_dim_input_slices)) |
| 261 | tile_out_coords = list(itertools.product(*full_dim_output_slices)) |
| 262 | tile_mask_1ds = list(itertools.product(*full_dim_masks_1d)) |
| 263 | for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True): |
| 264 | tiles.append( |
| 265 | Tile( |
| 266 | in_coords=in_coord, |
| 267 | out_coords=out_coord, |
| 268 | masks_1d=mask_1d, |
| 269 | ) |
| 270 | ) |
| 271 | return tiles |
| 272 | |
| 273 | |
| 274 | def create_tiles( |
| 275 | tensor_shape: torch.Size, |
| 276 | splitters: List[SplitOperation], |
| 277 | mappers: List[MappingOperation], |
| 278 | ) -> List[Tile]: |
| 279 | if len(splitters) != len(tensor_shape): |
| 280 | raise ValueError( |
| 281 | f"Number of splitters must be equal to number of dimensions in tensor shape, " |
| 282 | f"got {len(splitters)} and {len(tensor_shape)}" |
| 283 | ) |
| 284 | if len(mappers) != len(tensor_shape): |
| 285 | raise ValueError( |
| 286 | f"Number of mappers must be equal to number of dimensions in tensor shape, " |
| 287 | f"got {len(mappers)} and {len(tensor_shape)}" |
| 288 | ) |
| 289 | intervals = [splitter(length) for splitter, length in zip(splitters, tensor_shape, strict=True)] |
| 290 | tiling_spec = TensorTilingSpec(original_shape=tensor_shape, dimension_intervals=tuple(intervals)) |
| 291 | return create_tiles_from_intervals_and_mappers(tiling_spec, mappers) |
| 292 |