mlx-community/SEA-RAFT-Tartan480x640-S-mlx

SEA-RAFT-S (Wang, Lipson, Deng — "SEA-RAFT: Simple, Efficient, Accurate RAFT for Optical Flow", ECCV 2024) converted to Apple MLX format (NHWC), consumed by xocialize/sea-raft-mlx (Python) and the upcoming MLX-Swift port.

This is the TartanAir-stage checkpoint — trained on TartanAir (CC-BY) only, giving it the cleanest data-provenance chain of the SEA-RAFT releases. For higher accuracy (at the cost of research-restricted dataset stages in the chain), see SEA-RAFT-Tartan-C-T-TSKH-spring540x960-S-mlx.

Parity

Converted weights verified against the PyTorch reference (princeton-vl/SEA-RAFT@9137517) on an identical input pair: final flow cosine 1.0000, max_abs 1.2e-3 px, EPE 1.1e-4 px (fp32 CPU).

Model

• Variant S: ResNet-18-style FPN backbone, dim 128, 4-level all-pairs correlation (r=4), 2 ConvNeXt update blocks, 4 iterations, convex 8× upsample.
• Input [N, H, W, 3] RGB in 0–255; output [N, H, W, 2] flow in pixels.

Files

model.safetensors (MLX NHWC, fp32) · config.json.

License & provenance

• Code + weights: BSD-3-Clause — the upstream repository LICENSE (Princeton Vision & Learning Lab, 2024). The original checkpoint is the first author's upload MemorySlices/Tartan480x640-S (tagged bsd-3-clause).
• Training data (this stage): TartanAir (CC-BY-4.0).
• Author confirmation (2026-06-10): the authors confirmed in princeton-vl/SEA-RAFT#31 that "we intend the BSD-3-Clause license to cover the checkpoints as well" — the weights are BSD-3-Clause by explicit author statement. Conversion is format-only (NCHW→NHWC), with attribution retained.