Unified wrapper to load N frames. - bidirectional=True : Loads N frames [0..N-1] and 2*(N-2) flows from center frames. -

Resize flow back to original resolution. Handles both 4D (B, 2, H, W) and 5D (B, T, 2, H, W) inputs.

Load a pretrained MegaFlow model from HuggingFace Hub. Args: model_name: One of 'megaflow-flow' (optical flow), 'megaflow-chairs-

Args: images (torch.Tensor): Input images with shape [B, S, 3, H, W], in range [0, 1]. B: batch size, S: sequence

Expects a two dimensional flow image of shape. Args: flow_uv (np.ndarray): Flow UV image of shape [H,W,2] clip_flow (float,

Custom interpolate to avoid INT_MAX issues in nn.functional.interpolate.

Custom interpolate to avoid INT_MAX issues in nn.functional.interpolate.

Parse command line arguments and merge them with JSON configuration. This function first parses command line arguments, then loads configuration

Applies 1D rotary position embeddings along one dimension. Args: tokens: Input token features. positions: Position in

Apply positional embedding to tensor x.

Implementation of the forward pass through the DPT head. This method processes a specific chunk of frames from the sequence.

r"""Sample a tensor using bilinear interpolation `bilinear_sampler(input, coords)` samples a tensor :attr:`input` at coordinates :attr:`coord

Calculates new dimensions by fixing the width and preserving the aspect ratio. Ensures the height is divisible by the patch_size constraint.

Calculates new dimensions preserving the original aspect ratio. Aligns the longest edge to the target_fix_width and ensures divisibility.

1x1 convolution without padding

3x3 convolution with padding

Create the data loader for the corresponding training set

this will perform the index select, cat the tensors, and provide the attn_bias from cache

Load checkpoint into a LightningModule or nn.Module. Supports: - None: load default pretrained weights (e.g., safetensors) -

This function generates a 1D positional embedding from a given grid using sine and cosine functions. Args: - embed_dim: The embedding di

Loss function defined over sequence of flow predictions

Build the patch embed layer. If 'conv', we use a simple PatchEmbed conv layer. Otherwise, we use a vision transformer.

Initializes the position generator with an empty cache.

Computes frequency components for rotary embeddings. Args: dim: Feature dimension (must be even). seq_len: Maximum se

Process frame attention blocks. We keep tokens in shape (B*S, P, C).

Process global attention blocks. We keep tokens in shape (B, S*P, C).

Performs feature rotation by splitting and recombining feature dimensions. Args: x: Input tensor to rotate. Returns:

Photometric augmentation

Convert a self-attention qkv projection layer (dim -> 3*dim) into separate q_proj (dim -> dim) and kv_proj (dim -> 2*dim) layers. Re

Create a normalized UV grid of shape (width, height, 2). The grid spans horizontally and vertically according to an aspect ratio, ensuri

Occlusion augmentation (only for last frame)

Compute hierarchical flow loss with optional initial prediction loss. Parameters ---------- outputs : Dict[str, List[torch.Tenso

Applies the flow color wheel to (possibly clipped) flow components u and v. According to the C++ source code of Daniel Scharstein Accord

x_list contains a list of tensors to nest together and run

Generator that yields (RGB frame, original_shape_tuples, native_fps). If fix_width is provided, resizes the frame preserving aspect ratio before y