From affe9fa1da7ffa9ff6a5e49b21b25b80fc71cda0 Mon Sep 17 00:00:00 2001 From: Statistical Impossibility <220770210+Statistical-Impossibility@users.noreply.github.com> Date: Sun, 14 Jun 2026 20:18:40 +0300 Subject: [PATCH] feat: add Cadence + ComputeFlow nodes (3 -> 5), bump to 0.3.0 --- NODES.md | 72 +++++++++++++++++++ README.md | 53 +++++++++++--- __init__.py | 12 ++++ cadence_core.py | 75 ++++++++++++++++++++ cadence_nodes.py | 70 +++++++++++++++++++ flow_core.py | 73 ++++++++++++++++++++ flow_nodes.py | 83 ++++++++++++++++++++++ pyproject.toml | 6 +- raft_core.py | 37 ++++++++++ tests/conftest.py | 4 ++ tests/test_cadence_core.py | 134 ++++++++++++++++++++++++++++++++++++ tests/test_cadence_nodes.py | 118 +++++++++++++++++++++++++++++++ tests/test_flow_core.py | 35 ++++++++++ tests/test_flow_nodes.py | 61 ++++++++++++++++ tests/test_flow_warp.py | 33 +++++++++ tests/test_raft_core.py | 19 +++++ 16 files changed, 873 insertions(+), 12 deletions(-) create mode 100644 cadence_core.py create mode 100644 cadence_nodes.py create mode 100644 flow_core.py create mode 100644 flow_nodes.py create mode 100644 raft_core.py create mode 100644 tests/test_cadence_core.py create mode 100644 tests/test_cadence_nodes.py create mode 100644 tests/test_flow_core.py create mode 100644 tests/test_flow_nodes.py create mode 100644 tests/test_flow_warp.py create mode 100644 tests/test_raft_core.py diff --git a/NODES.md b/NODES.md index 8d16663..8c5831b 100644 --- a/NODES.md +++ b/NODES.md @@ -125,3 +125,75 @@ independent; at `strength 0` the node is a pass-through. richer image as the bootstrap/reference and the whole loop inherits its colors. - `preserve_luminance` keeps brightness, so this node controls **color only**; it will not stop a loop that is darkening on its own (that's a denoise/luminance issue). + +--- + +## `Cadence` + +Menu group: **Stat-Imp / Deforum**. + +Vanilla-Deforum **cadence**: only diffuse every *N*th frame, and fill the in-betweens +("tweens") by warping the two image lineages along the motion schedule and cross-dissolving +them. This is the classic Deforum speed/smoothness trick — at `cadence = N` the expensive +sampler runs *N×* fewer times. `cadence = 1` is a mathematical no-op (every frame is a gen +frame), so it's safe to leave wired. + +It consolidates the scalar math + switch/blend logic of the old multi-node CADENCE group into +one node. The actual warps stay **outside** the node (so any transform can feed it — 3D depth +warp today, optical flow tomorrow), and the sampler gate stays graph-level; this node only does +the dissolve math and emits a boolean telling the graph which kind of frame this is. + +> This is a **loop-internal** node: its scalar inputs come from the Deforum harness +> (`loop_index`, `num_frames`) and a cadence knob, not from hand-typed widgets. It has **zero +> widgets** by design (every input is `forceInput`), which makes the frontend widget-serialization +> trap impossible. Use the example workflow's wiring as the reference. + +### Inputs (required) +- `loop_index` (INT) — current frame index from the `ForLoop`. +- `cadence` (INT) — diffuse every *N*th frame (`1` = off / every frame). +- `num_frames` (INT) — total frames (for the tail divisor, so the last fade completes on the final frame). +- `is_start` (BOOLEAN) — true on the loop's first iteration (bootstrap). +- `warped_old` (IMAGE, *lazy*) — the previous anchor lineage, warped along the schedule. Only pulled on **tween** frames. +- `corrected` (IMAGE, *lazy*) — the current warped+color-corrected frame. +- `fresh` (IMAGE, *lazy*) — the freshly diffused image. Only pulled on **gen** frames (so the sampler is skipped on tweens). + +### Inputs (optional — A2 v2 optical-flow morph) +- `flow` (FLOW) — displacement field from **Compute Optical Flow** below. Leave unconnected for the plain v1 dissolve (bit-identical). +- `flow_factor` (FLOAT, 0–1) — how far to morph the lineages into alignment before dissolving (tween frames only). + +### Outputs +- `frame` (IMAGE) — the emitted frame (a dissolve on tweens, the fresh diffusion on gen frames). +- `next_anchor` (IMAGE) — the anchor lineage to carry to the next iteration (wire to `ForLoopClose`). +- `is_gen` (BOOLEAN) — true on gen frames; drives the graph-level sampler gate. + +> Key (counter-intuitive, but vanilla-identical) semantic: new diffused content arrives **≤N +> frames after** its schedule index — it fades in across the *following* span, not instantly. + +--- + +## `Compute Optical Flow` + +Menu group: **Stat-Imp / Deforum**. + +Estimates an **optical-flow** field (per-pixel displacement) between two frames, to feed the +`Cadence` node's `flow` input for motion-aware in-betweens. Kept as a separate node so the flow +backend is swappable (and a ground-truth Blender motion-vector pass could replace it later). + +### Inputs +- `image_from` (IMAGE) — source frame. +- `image_to` (IMAGE) — target frame. Output flow maps `image_from → image_to`, sized to `image_to`. +- `method` — `None` (zero flow = identity, ≡ cadence v1), `RAFT`, `DIS Medium`, `DIS Fine`, `Farneback`. +- `device` — `AUTO` / `cuda` / `cpu` (only affects RAFT). + +### Outputs +- `flow` (FLOW) — `(1, H, W, 2)` displacement field; wire into `Cadence` → `flow`. + +### Backends & dependencies +- `DIS Medium` / `DIS Fine` / `Farneback` — **OpenCV** (`cv2`), CPU. OpenCV is usually already + bundled with ComfyUI. +- `RAFT` — **torchvision** (weights auto-download on first use), GPU or CPU. +- `None` — no dependency; returns identity flow. + +> Practical note: flow-based morphing shines on **clean, high-quality frames**. On heavily +> distilled / low-step models the per-frame texture noise can make estimated flow noisy; in that +> regime `None` (plain dissolve) or `Farneback`/`DIS` are safer than `RAFT`. diff --git a/README.md b/README.md index 3492d85..cba5d99 100644 --- a/README.md +++ b/README.md @@ -1,26 +1,45 @@ # ComfyUI Stat-Imp Nodes Clean, MIT-licensed nodes for Deforum-style animation loops in ComfyUI. -This release ships three nodes: a **depth-aware 3D camera warp** and two -**color-coherence** nodes (a fixed-anchor matcher and a rolling/EMA matcher) -that stop the feedback loop's color from drifting into neon "acid". +This release ships following nodes: -Pure torch — no extra Python dependencies. See **[NODES.md](NODES.md)** for the -full per-node reference (inputs, outputs, motion directions, tuning guidance). +- a **depth-aware 3D camera warp** (exact forward port of Deforum's `transform_image_3d`), +- two **color-coherence** nodes — a fixed-anchor matcher and a rolling/EMA matcher — + that stop the feedback loop's color from drifting into neon "acid", +- a **cadence** node (diffuse every Nth frame, dissolve the in-betweens — vanilla Deforum semantics), and +- an **optical-flow** estimator that feeds the cadence node for motion-aware in-betweens. + +The core warp + color nodes are **pure torch** — no extra Python dependencies. The optical-flow +node additionally uses **OpenCV** (DIS / Farneback backends, usually already bundled with ComfyUI) +and, only for its optional RAFT backend, **torchvision**. See **[NODES.md](NODES.md)** for the full +per-node reference (inputs, outputs, motion directions, tuning guidance). + +## Requirements + +1. **A working ComfyUI install.** If you don't have one yet, get it from the official site / + repo first: **https://github.com/comfyanonymous/ComfyUI** (follow its install guide). +2. **The Deforum harness — our fork.** The example workflow runs inside a Deforum + `ForLoop` animation loop. It needs the harness nodes (ForLoop, ValueSchedule, Set/Get bus), + and specifically **our fork** of `deforum-comfy-nodes`, which widens the loop's carry state + (10 `value` sockets instead of 1) and fixes the value-schedule formula parser so `sin/cos` + and multi-argument functions work. Install **our fork**, not the upstream: + **https://github.com/Statistical-Impossibility/deforum-comfy-nodes** + (See *"Why a fork"* below.) The five nodes in *this* pack work standalone, but the shipped + example workflow assumes the fork. ## Install -**Manual / git:** clone into your ComfyUI custom-nodes folder, then restart ComfyUI -and hard-refresh the browser: +Clone into your ComfyUI custom-nodes folder, then restart ComfyUI and hard-refresh the browser: ```bash cd ComfyUI/custom_nodes git clone https://github.com/Statistical-Impossibility/comfyui-Stat-Imp-nodes +# and the harness fork the example workflow needs: +git clone https://github.com/Statistical-Impossibility/deforum-comfy-nodes ``` The nodes appear under the **Stat-Imp / Deforum** menu group. - ## How to use These nodes are building blocks for a Deforum-style 3D animation loop. A ready-to-run @@ -32,6 +51,21 @@ a full depth-warp + color-coherence animation loop for **SDXL Turbo / Lightning Open it in ComfyUI, pick a distilled SDXL checkpoint, and queue. The embedded note explains the inputs and settings. +## Why a fork (of deforum-comfy-nodes) + +The upstream [deforum-comfy-nodes](https://github.com/deforum/deforum-comfy-nodes) harness is great +but two things blocked this pack: + +1. **Loop carry width.** Cadence needs to thread several extra images through the animation loop + (the old anchor lineage, the captured frame-0, …). Upstream `ForLoop` exposes a single + carry socket; our fork widens it to **10** (`value1`…`value10`). +2. **Value-schedule formula parser.** Upstream `_sanitize_value` stripped *all* parentheses, so + `0:(cos(t))` collapsed to `cost` → `NameError`, and a comma inside a value (`pow(t,2)`) was + misread as a keyframe separator. Our fork keeps inner parens and splits keyframes only at + paren-depth 0, so `sin/cos/abs/pow/min/max/...` work in value schedules. + +Both are clean, additive changes; the fork stays drop-in compatible with upstream workflows. + ## Credits This pack stands on the original **Deforum** work — full credit to its authors: @@ -41,4 +75,5 @@ This pack stands on the original **Deforum** work — full credit to its authors `transform_image_3d`, and the color matcher is a clean-room reimplementation of its `maintain_colors` (`helpers/colors.py`). - **[Deforum-Comfy-Nodes](https://github.com/deforum/deforum-comfy-nodes)** - — the main ComfyUI Deforum harness and node pack (ForLoop, schedules, Set/Get bus). + — the main ComfyUI Deforum harness and node pack (ForLoop, schedules, Set/Get bus); + our fork above extends it. diff --git a/__init__.py b/__init__.py index 4a1a95e..99e7c91 100644 --- a/__init__.py +++ b/__init__.py @@ -1,9 +1,13 @@ try: from .nodes import DepthWarp3D_StatImp from .color_nodes import NODE_CLASS_MAPPINGS as COLOR_CLASS, NODE_DISPLAY_NAME_MAPPINGS as COLOR_DISPLAY + from .cadence_nodes import NODE_CLASS_MAPPINGS as CADENCE_CLASS, NODE_DISPLAY_NAME_MAPPINGS as CADENCE_DISPLAY + from .flow_nodes import NODE_CLASS_MAPPINGS as FLOW_CLASS, NODE_DISPLAY_NAME_MAPPINGS as FLOW_DISPLAY except ImportError: from nodes import DepthWarp3D_StatImp from color_nodes import NODE_CLASS_MAPPINGS as COLOR_CLASS, NODE_DISPLAY_NAME_MAPPINGS as COLOR_DISPLAY + from cadence_nodes import NODE_CLASS_MAPPINGS as CADENCE_CLASS, NODE_DISPLAY_NAME_MAPPINGS as CADENCE_DISPLAY + from flow_nodes import NODE_CLASS_MAPPINGS as FLOW_CLASS, NODE_DISPLAY_NAME_MAPPINGS as FLOW_DISPLAY # Distinct node keys (the type IDs saved in workflows) so this pack coexists with # the old deforum-comfy-nodes -- install both, swap a node for the other, no collision. @@ -14,4 +18,12 @@ NODE_DISPLAY_NAME_MAPPINGS = {"DepthWarp3D| Deforum_Stat-Imp": "Depth Warp 3D"} NODE_CLASS_MAPPINGS.update(COLOR_CLASS) NODE_DISPLAY_NAME_MAPPINGS.update(COLOR_DISPLAY) +# Cadence (brick A2): vanilla-semantics dissolve, lazy inputs gate sampler/warp pulls. +NODE_CLASS_MAPPINGS.update(CADENCE_CLASS) +NODE_DISPLAY_NAME_MAPPINGS.update(CADENCE_DISPLAY) + +# Optical-flow estimator (A2 v2): feeds the Cadence node's flow socket. +NODE_CLASS_MAPPINGS.update(FLOW_CLASS) +NODE_DISPLAY_NAME_MAPPINGS.update(FLOW_DISPLAY) + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/cadence_core.py b/cadence_core.py new file mode 100644 index 0000000..e955f9c --- /dev/null +++ b/cadence_core.py @@ -0,0 +1,75 @@ +"""Cadence core (brick A2): pure math + emit for vanilla-semantics cadence. + +The dissolve pairs the two most recent PAST diffusion lineages (content arrives +<=N frames after its schedule index, same as vanilla Deforum). Tail divisor +D = min(N, F - p) makes the final fade complete exactly at frame F-1. +""" +import torch + + +def cadence_math(loop_index, cadence, num_frames): + """Per-iteration scalars. Returns (is_gen, w, first_span). + + is_gen -- this iteration diffuses (i % N == 0); drives the sampler gate. + w -- dissolve weight ((i%N)+1)/D with tail divisor D = min(N, F-p). + first_span -- i < N: no old lineage exists yet, emit the corrected image as-is. + """ + n = max(1, int(cadence)) + i = int(loop_index) + m = i % n + p = i - m + d = min(n, int(num_frames) - p) + return m == 0, (m + 1) / d, i < n + + +def cadence_emit(is_start, first_span, is_gen, w, warped_old, corrected, fresh, + flow=None, flow_factor=1.0): + """Emit (frame, next_anchor). Unused image args may be None (lazy inputs). + + frame: fresh at frame 0; corrected during the first span; otherwise the + dissolve OLD*(1-w) + NEW*w where OLD/NEW are the two most recent lineages. + next_anchor: the value5 roll — corrected on gen iterations, warped_old on tweens. + + Optical-flow morph (A2 v2): on TWEEN frames only (is_gen False), when `flow` is given + and flow_factor != 0, the two blended lineages (warped_old, corrected) are flow-aligned + before the dissolve — + OLD_aligned = flow_warp(warped_old, flow * ( w * flow_factor)) + NEW_aligned = flow_warp(corrected, flow * (-(1-w) * flow_factor)) + Gen frames are left as a plain blend (their pair is corrected->fresh, a different field; + vanilla likewise flows only the in-between frames). flow=None / flow_factor=0 / zero-field + / gen-frame => bit-identical to the plain dissolve. + """ + next_anchor = corrected if is_gen else warped_old + if is_start: + return fresh, next_anchor + if first_span: + return corrected, next_anchor + old = corrected if is_gen else warped_old + new = fresh if is_gen else corrected + if not is_gen and flow is not None and float(flow_factor) != 0.0: + old = flow_warp(old, flow * (w * float(flow_factor))) + new = flow_warp(new, flow * (-(1.0 - w) * float(flow_factor))) + # clamp mirrors ComfyUI ImageBlend's post-blend clamp (no-op for in-range inputs) + frame = (old * (1.0 - w) + new * w).clamp(0.0, 1.0) + return frame, next_anchor + + +def flow_warp(image, flow): + """Backward-warp image (B,H,W,C) by flow (B,H,W,2) abs-pixel: out(x) = image(x - flow(x)). + + Mirrors vanilla image_transform_optical_flow (cv2.remap of x-flow) in torch grid_sample, + bilinear + reflection padding (no black edges). Sampling-only; no estimation. + """ + import torch.nn.functional as F + b, h, w, c = image.shape + ys, xs = torch.meshgrid(torch.arange(h, device=image.device, dtype=image.dtype), + torch.arange(w, device=image.device, dtype=image.dtype), + indexing="ij") + src_x = xs[None] - flow[..., 0] # sample location = x - flow + src_y = ys[None] - flow[..., 1] + gx = 2.0 * src_x / max(w - 1, 1) - 1.0 # -> normalized [-1,1] for grid_sample + gy = 2.0 * src_y / max(h - 1, 1) - 1.0 + grid = torch.stack((gx, gy), dim=-1) # (B,H,W,2) + chw = image.permute(0, 3, 1, 2) + out = F.grid_sample(chw, grid, mode="bilinear", padding_mode="reflection", align_corners=True) + return out.permute(0, 2, 3, 1).contiguous() diff --git a/cadence_nodes.py b/cadence_nodes.py new file mode 100644 index 0000000..dcda69a --- /dev/null +++ b/cadence_nodes.py @@ -0,0 +1,70 @@ +"""ComfyUI node: Cadence — vanilla-semantics cadence dissolve (brick A2). + +Consolidates the CADENCE group's scalar math + switch/blend logic. The 2nd +warp (old-anchor) stays OUTSIDE the node so any transform can feed warped_old +(3D depth warp today, optical flow later). The sampler gate also stays +graph-level (LazySwitch on the LATENT path); this node only emits its boolean. + +Lazy contract: tween iterations never request `fresh` (so the sampler branch +is never executed); gen iterations never request `warped_old` (warp skipped). +""" +try: + from .cadence_core import cadence_math, cadence_emit +except ImportError: + from cadence_core import cadence_math, cadence_emit + + +class Cadence: + CATEGORY = "Stat-Imp/Deforum" + RETURN_TYPES = ("IMAGE", "IMAGE", "BOOLEAN") + RETURN_NAMES = ("frame", "next_anchor", "is_gen") + FUNCTION = "run" + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + # all scalars forceInput: the node has ZERO widgets, so saved + # widgets_values=[] can never mis-align in the frontend (FINDINGS trap) + "loop_index": ("INT", {"default": 0, "forceInput": True}), + "cadence": ("INT", {"default": 1, "forceInput": True}), + "num_frames": ("INT", {"default": 1, "forceInput": True}), + "is_start": ("BOOLEAN", {"default": False, "forceInput": True}), + "warped_old": ("IMAGE", {"lazy": True}), + "corrected": ("IMAGE", {"lazy": True}), + "fresh": ("IMAGE", {"lazy": True}), + }, "optional": { + # A2 v2 optical-flow morph (omit both => bit-identical v1 dissolve). + # NON-lazy on purpose: a lazy optional input that is left UNCONNECTED makes + # ComfyUI raise "needs input flow, but there is no input" — which would break + # every v1 workflow that has no ComputeFlow wired. Non-lazy => unconnected is + # simply None. flow is link-only + flow_factor is forceInput => still ZERO widgets. + "flow": ("FLOW",), + "flow_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, + "forceInput": True}), + }} + + def check_lazy_status(self, loop_index, cadence, num_frames, is_start, + warped_old=None, corrected=None, fresh=None, + flow=None, flow_factor=1.0): + is_gen, _, _ = cadence_math(loop_index, cadence, num_frames) + needed = [] + if corrected is None: + needed.append("corrected") + if is_gen and fresh is None: + needed.append("fresh") + if not is_gen and warped_old is None: + needed.append("warped_old") + return needed + + def run(self, loop_index, cadence, num_frames, is_start, + warped_old=None, corrected=None, fresh=None, + flow=None, flow_factor=1.0): + is_gen, w, first_span = cadence_math(loop_index, cadence, num_frames) + frame, next_anchor = cadence_emit(bool(is_start), first_span, is_gen, w, + warped_old, corrected, fresh, + flow=flow, flow_factor=float(flow_factor)) + return (frame, next_anchor, is_gen) + + +NODE_CLASS_MAPPINGS = {"Cadence| Deforum_Stat-Imp": Cadence} +NODE_DISPLAY_NAME_MAPPINGS = {"Cadence| Deforum_Stat-Imp": "Cadence"} diff --git a/flow_core.py b/flow_core.py new file mode 100644 index 0000000..c19d895 --- /dev/null +++ b/flow_core.py @@ -0,0 +1,73 @@ +"""Optical-flow primitives for cadence (Legacy = OpenCV DIS / Farneback). + +Pure functions, no ComfyUI imports, so they run headless under pytest. +Clean-room reimplementation of the flow compute + flow-warp primitives +(mirrors the MIT seanlynch/comfyui-optical-flow approach). All images are +HxWx3 uint8 BGR; flow is (H,W,2) float32 absolute pixel displacement. +""" +import cv2 +import numpy as np + +def compute_flow_dis(i1, i2, method="DIS Medium"): + """Dense optical flow i1->i2 via OpenCV DIS. Returns (H,W,2) float32. + + Vanilla's two cadence presets (hybrid_video.py:362-381): DIS Medium = the stock + MEDIUM preset; DIS Fine = custom (finest scale 0, 192 grad-descent iters, patch 8/stride 4). + """ + g1 = cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY) + g2 = cv2.cvtColor(i2, cv2.COLOR_BGR2GRAY) + if method == "DIS Medium": + dis = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_MEDIUM) + elif method == "DIS Fine": + dis = cv2.DISOpticalFlow_create(None) + dis.setFinestScale(0) + dis.setGradientDescentIterations(192) + dis.setPatchSize(8) + dis.setPatchStride(4) + else: + raise ValueError(f"unknown DIS method: {method}") + return dis.calc(g1, g2, None) + + +def compute_flow_farneback(i1, i2): + """Dense optical flow i1->i2 via Farneback (the other original-Deforum method).""" + g1 = cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY) + g2 = cv2.cvtColor(i2, cv2.COLOR_BGR2GRAY) + return cv2.calcOpticalFlowFarneback( + g1, g2, None, + pyr_scale=0.5, levels=3, winsize=15, iterations=3, + poly_n=5, poly_sigma=1.2, flags=0, + ) + + +def _extend_flow(flow, w, h): + """Pad a flow field up to (h,w), filling the border with identity coords.""" + fh, fw = flow.shape[:2] + ox, oy = (w - fw) // 2, (h - fh) // 2 + out = np.zeros((h, w, 2), np.float32) + out[:, :, 0] = np.arange(w)[None, :] + out[:, :, 1] = np.arange(h)[:, None] + out[oy:oy + fh, ox:ox + fw] = flow + np.array([ox, oy], np.float32) + return out + + +def _remap(img, flow): + """Resample img by an absolute-coordinate map, with reflect padding to avoid black edges.""" + border = cv2.BORDER_REFLECT_101 + h, w = img.shape[:2] + dy, dx = int(h * 0.25), int(w * 0.25) # 25% reflect pad stops black edges + big = cv2.copyMakeBorder(img, dy, dy, dx, dx, border) + lh, lw = big.shape[:2] + big_flow = _extend_flow(flow, lw, lh) + out = cv2.remap(big, big_flow, None, cv2.INTER_LINEAR, borderMode=border) + return out[dy:dy + h, dx:dx + w] # center-crop back + + +def image_transform_optical_flow(img, flow, flow_factor=1.0): + """Warp img along `flow` (scaled by flow_factor). img HxWx3 uint8, flow (H,W,2).""" + flow = flow * flow_factor if flow_factor != 1.0 else flow.copy() + flow = -flow # reversed for remap (sample from source) + h, w = img.shape[:2] + flow[:, :, 0] += np.arange(w) + flow[:, :, 1] += np.arange(h)[:, np.newaxis] + return _remap(img, flow) diff --git a/flow_nodes.py b/flow_nodes.py new file mode 100644 index 0000000..d1fffa7 --- /dev/null +++ b/flow_nodes.py @@ -0,0 +1,83 @@ +"""ComfyUI node: ComputeFlow — optical-flow estimator for cadence (A2 v2). + +Estimator lives OUTSIDE the Cadence node (transform-agnostic boundary). Outputs FLOW = +torch (1,H,W,2) abs-pixel displacement image_from->image_to; "None" => zero flow (identity). +Backends: RAFT (torchvision, GPU/CPU), DIS Medium/Fine + Farneback (OpenCV, CPU). +""" +import numpy as np +import torch + +try: + from .flow_core import compute_flow_dis, compute_flow_farneback +except ImportError: + from flow_core import compute_flow_dis, compute_flow_farneback + + +def _to_bgr_uint8(image): + """ComfyUI IMAGE (1,H,W,3) RGB float 0..1 -> HxWx3 uint8 BGR for OpenCV/RAFT cores.""" + arr = (image[0].clamp(0, 1).cpu().numpy() * 255.0).round().astype(np.uint8) # HWC RGB + return arr[:, :, ::-1].copy() # -> BGR + + +def _pick_device(device): + if device == "cpu": + return "cpu" + if device == "cuda": + return "cuda" + return "cuda" if torch.cuda.is_available() else "cpu" + + +def estimate_flow(image_from, image_to, method, device): + """Return FLOW (1,H,W,2) float32 abs-pixel, image_from->image_to, sized to image_to. + + image_from is resized to image_to's resolution first: on bootstrap/first-span frames the + old-anchor lineage can be a different size than the corrected image, and the flow backends + (RAFT/DIS/Farneback) all assert equal-size inputs. On real tween frames both are already the + render resolution, so the resize is a no-op. (Cadence ignores the result on guarded frames.) + """ + import torch.nn.functional as F + _, h, w, _ = image_to.shape + if image_from.shape[1:3] != image_to.shape[1:3]: + chw = image_from.permute(0, 3, 1, 2) + chw = F.interpolate(chw, size=(h, w), mode="bilinear", align_corners=False) + image_from = chw.permute(0, 2, 3, 1).contiguous() + if method == "None": + return torch.zeros(1, h, w, 2, dtype=torch.float32) + a_bgr = _to_bgr_uint8(image_from) + b_bgr = _to_bgr_uint8(image_to) + if method in ("DIS Medium", "DIS Fine"): + f = compute_flow_dis(a_bgr, b_bgr, method) + elif method == "Farneback": + f = compute_flow_farneback(a_bgr, b_bgr) + elif method == "RAFT": + try: + from .raft_core import raft_flow + except ImportError: + from raft_core import raft_flow + f = raft_flow(a_bgr, b_bgr, device=_pick_device(device)) + else: + raise ValueError(f"unknown flow method: {method}") + return torch.from_numpy(np.ascontiguousarray(f, dtype=np.float32))[None] # (1,H,W,2) + + +class ComputeFlow: + CATEGORY = "Stat-Imp/Deforum" + RETURN_TYPES = ("FLOW",) + RETURN_NAMES = ("flow",) + FUNCTION = "run" + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "image_from": ("IMAGE",), + "image_to": ("IMAGE",), + "method": (["None", "RAFT", "DIS Medium", "DIS Fine", "Farneback"],), + "device": (["AUTO", "cuda", "cpu"], {"default": "AUTO"}), + }} + + def run(self, image_from, image_to, method, device): + return (estimate_flow(image_from, image_to, method, device),) + + +NODE_CLASS_MAPPINGS = {"ComputeFlow| Deforum_Stat-Imp": ComputeFlow} +NODE_DISPLAY_NAME_MAPPINGS = {"ComputeFlow| Deforum_Stat-Imp": "Compute Optical Flow"} diff --git a/pyproject.toml b/pyproject.toml index 6f5e0e7..fdd817b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,12 @@ [project] name = "comfyui-stat-imp-nodes" -description = "Clean pure-torch Deforum-style nodes for ComfyUI animation loops: an exact-Deforum depth-aware 3D camera warp (pixel-identical to transform_image_3d) plus LAB/EMA color-coherence nodes that stop feedback-loop color drift." -version = "0.2.0" +description = "Clean pure-torch Deforum-style nodes for ComfyUI animation loops: an exact-Deforum depth-aware 3D camera warp (pixel-identical to transform_image_3d), LAB/EMA color-coherence nodes that stop feedback-loop color drift, a vanilla-semantics cadence node, and an optical-flow estimator." +version = "0.3.0" license = { text = "MIT" } dependencies = [] [project.urls] -Repository = "https://github.com/GITHUB_HANDLE/comfyui-Stat-Imp-nodes" +Repository = "https://github.com/Statistical-Impossibility/comfyui-Stat-Imp-nodes" [tool.comfy] PublisherId = "PUBLISHER_ID" diff --git a/raft_core.py b/raft_core.py new file mode 100644 index 0000000..755bf82 --- /dev/null +++ b/raft_core.py @@ -0,0 +1,37 @@ +"""RAFT optical flow (Improved variant) — torchvision (BSD-3), opt-in. + +Not used by the Legacy DIS/Farneback path. Requires `torchvision`; the weights +(~20 MB) auto-download on first use, so this stays opt-in. The model is lazily +constructed and cached. Returns (H,W,2) float32 absolute pixel flow. +""" +import numpy as np + +_RAFT = None + + +def _model(device): + global _RAFT + if _RAFT is None: + import torch # noqa: F401 + from torchvision.models.optical_flow import raft_large, Raft_Large_Weights + _RAFT = raft_large(weights=Raft_Large_Weights.DEFAULT, progress=False).eval().to(device) + return _RAFT + + +def raft_flow(i1, i2, device="cuda"): + """i1,i2: HxWx3 uint8 BGR. Returns (H,W,2) float32 absolute pixel flow.""" + import torch + import torch.nn.functional as F + + def prep(img): + t = torch.from_numpy(img[:, :, ::-1].copy()).permute(2, 0, 1)[None].float() / 255.0 # BGR->RGB BCHW + h, w = t.shape[-2:] + t = F.pad(t, (0, (-w) % 8, 0, (-h) % 8), mode="replicate") + return (t * 2 - 1).to(device), (h, w) + + a, (h, w) = prep(i1) + b, _ = prep(i2) + with torch.inference_mode(): + flow = _model(device)(a, b)[-1] # (1,2,H,W), last = finest + flow = flow[0, :, :h, :w].permute(1, 2, 0).cpu().numpy() # (H,W,2) + return np.ascontiguousarray(flow, dtype=np.float32) diff --git a/tests/conftest.py b/tests/conftest.py index 8bd89de..d78b763 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,3 +4,7 @@ import sys # Make the pack modules (color_core, color_nodes, nodes, ...) importable when # running `pytest` from the pack dir, without installing the package. sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def pytest_configure(config): + config.addinivalue_line("markers", "slow: long/needs-download (RAFT weights)") diff --git a/tests/test_cadence_core.py b/tests/test_cadence_core.py new file mode 100644 index 0000000..42b6137 --- /dev/null +++ b/tests/test_cadence_core.py @@ -0,0 +1,134 @@ +"""Tests for cadence_core (brick A2 consolidation).""" +import pytest +import torch + +from cadence_core import cadence_math, cadence_emit + + +# (i, N, F) -> (is_gen, w, first_span) +MATH_TABLE = [ + # cadence=1 degenerates: always gen, w=1, only i=0 in first span + (0, 1, 8, True, 1.0, True), + (3, 1, 8, True, 1.0, False), + (7, 1, 8, True, 1.0, False), + # N=4, F=12 (no tail truncation: D=N=4 mid-video) + (0, 4, 12, True, 1 / 4, True), + (1, 4, 12, False, 2 / 4, True), + (3, 4, 12, False, 4 / 4, True), + (4, 4, 12, True, 1 / 4, False), + (6, 4, 12, False, 3 / 4, False), + (7, 4, 12, False, 4 / 4, False), + (11, 4, 12, False, 4 / 4, False), + # tail divisor: F=9, N=6 -> last span p=6, D=min(6, 9-6)=3, fade completes at F-1 + (6, 6, 9, True, 1 / 3, False), + (7, 6, 9, False, 2 / 3, False), + (8, 6, 9, False, 3 / 3, False), + # cadence=0 guard: treated as 1 (no div-by-zero) + (2, 0, 8, True, 1.0, False), +] + + +@pytest.mark.parametrize("i,n,f,is_gen,w,first_span", MATH_TABLE) +def test_cadence_math_table(i, n, f, is_gen, w, first_span): + g, ww, fs = cadence_math(i, n, f) + assert g is is_gen + assert ww == pytest.approx(w) + assert fs is first_span + + +def _img(val, h=4, w=4): + return torch.full((1, h, w, 3), float(val)) + + +def test_emit_frame0_is_fresh(): + fresh, corrected = _img(0.9), _img(0.5) + frame, anchor = cadence_emit(True, True, True, 1.0, None, corrected, fresh) + assert torch.equal(frame, fresh) + assert torch.equal(anchor, corrected) # is_gen at i=0 -> anchor rolls to corrected + + +def test_emit_first_span_tween_is_corrected(): + corrected, old = _img(0.5), _img(0.2) + # i=1, N=4: first_span tween -> frame = corrected, anchor = warped_old + frame, anchor = cadence_emit(False, True, False, 0.5, old, corrected, None) + assert torch.equal(frame, corrected) + assert torch.equal(anchor, old) + + +def test_emit_gen_blends_corrected_to_fresh(): + # gen iteration mid-video: OLD=corrected, NEW=fresh, w=1/4 + corrected, fresh = _img(0.4), _img(0.8) + frame, anchor = cadence_emit(False, False, True, 0.25, None, corrected, fresh) + expected = corrected * 0.75 + fresh * 0.25 + assert torch.allclose(frame, expected) + assert torch.equal(anchor, corrected) + + +def test_emit_tween_blends_old_to_corrected(): + # tween mid-video: OLD=warped_old, NEW=corrected, w=3/4 + old, corrected = _img(0.2), _img(0.6) + frame, anchor = cadence_emit(False, False, False, 0.75, old, corrected, None) + expected = old * 0.25 + corrected * 0.75 + assert torch.allclose(frame, expected) + assert torch.equal(anchor, old) + + +def test_emit_cadence1_every_frame_fresh(): + # N=1: is_gen always, w=1 -> frame == fresh exactly (no-cadence degeneration) + corrected, fresh = _img(0.4), _img(0.8) + frame, _ = cadence_emit(False, False, True, 1.0, None, corrected, fresh) + assert torch.equal(frame, fresh) + + +def test_emit_output_clamped(): + old, corrected = _img(0.0), _img(1.0) + frame, _ = cadence_emit(False, False, False, 0.5, old, corrected, None) + assert frame.min() >= 0.0 and frame.max() <= 1.0 + + +def test_emit_flow_none_equals_v1(): + # flow=None must reproduce the plain-dissolve path exactly + old, corrected = _img(0.2), _img(0.6) + a = cadence_emit(False, False, False, 0.75, old, corrected, None) + b = cadence_emit(False, False, False, 0.75, old, corrected, None, flow=None, flow_factor=1.0) + assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) + + +def test_emit_flow_factor_zero_equals_v1(): + old, corrected = _img(0.2), _img(0.6) + flow = torch.ones(1, 4, 4, 2) # non-trivial flow, but factor 0 disables morph + a = cadence_emit(False, False, False, 0.5, old, corrected, None) + b = cadence_emit(False, False, False, 0.5, old, corrected, None, flow=flow, flow_factor=0.0) + assert torch.allclose(a[0], b[0], atol=1e-5) + + +def test_emit_flow_zero_field_equals_v1(): + old, corrected = _img(0.2), _img(0.6) + flow = torch.zeros(1, 4, 4, 2) + a = cadence_emit(False, False, False, 0.5, old, corrected, None) + b = cadence_emit(False, False, False, 0.5, old, corrected, None, flow=flow, flow_factor=1.0) + assert torch.allclose(a[0], b[0], atol=1e-5) + + +def _ramp(h=8, w=8): + # horizontal ramp so a horizontal warp visibly changes pixels + x = torch.arange(w, dtype=torch.float32).repeat(h, 1) / (w - 1) + return x[None, :, :, None].repeat(1, 1, 1, 3) + + +def test_emit_flow_morph_differs_from_plain_blend(): + old, corrected = _ramp(), _ramp() + flow = torch.zeros(1, 8, 8, 2) + flow[..., 0] = 2.0 # real horizontal motion -> morph shifts content, plain blend does not + plain = cadence_emit(False, False, False, 0.5, old, corrected, None)[0] + morph = cadence_emit(False, False, False, 0.5, old, corrected, None, flow=flow, flow_factor=1.0)[0] + assert not torch.allclose(plain, morph, atol=1e-3) + + +def test_emit_flow_ignored_on_first_span_and_frame0(): + corrected, fresh, old = _img(0.5), _img(0.9), _img(0.2) + flow = torch.ones(1, 4, 4, 2) + # frame 0 -> fresh regardless of flow + assert torch.equal(cadence_emit(True, True, True, 1.0, old, corrected, fresh, flow=flow)[0], fresh) + # first span -> corrected regardless of flow + assert torch.equal(cadence_emit(False, True, False, 0.5, old, corrected, None, flow=flow)[0], corrected) diff --git a/tests/test_cadence_nodes.py b/tests/test_cadence_nodes.py new file mode 100644 index 0000000..8e14003 --- /dev/null +++ b/tests/test_cadence_nodes.py @@ -0,0 +1,118 @@ +"""Tests for the Cadence node class (lazy-input contract + run).""" +import torch + +from cadence_nodes import Cadence + + +def _img(val): + return torch.full((1, 4, 4, 3), float(val)) + + +def _lazy(i, n, f, start): + """check_lazy_status with no image inputs resolved yet.""" + return set(Cadence().check_lazy_status(i, n, f, start)) + + +def test_lazy_gen_iteration_never_requests_warped_old(): + # gen (i%N==0): needs corrected+fresh; warp branch must NOT be pulled + assert _lazy(4, 4, 12, False) == {"corrected", "fresh"} + + +def test_lazy_tween_iteration_never_requests_fresh(): + # tween: needs corrected+warped_old(+flow for v2); the sampler branch must NOT be pulled + needed = _lazy(5, 4, 12, False) + assert {"corrected", "warped_old"} <= needed + assert "fresh" not in needed + + +def test_lazy_frame0_requests_fresh_not_warp(): + assert _lazy(0, 4, 12, True) == {"corrected", "fresh"} + + +def test_lazy_cadence1_always_gen(): + assert _lazy(3, 1, 8, False) == {"corrected", "fresh"} + + +def test_lazy_requests_only_missing_inputs(): + # corrected already supplied + a zero flow supplied => only warped_old still missing + flow = torch.zeros(1, 4, 4, 2) + needed = Cadence().check_lazy_status(5, 4, 12, False, corrected=_img(0.5), flow=flow) + assert set(needed) == {"warped_old"} + + +def test_run_tween_returns_blend_anchor_and_is_gen(): + old, corrected = _img(0.2), _img(0.6) + frame, anchor, is_gen = Cadence().run(6, 4, 12, False, warped_old=old, corrected=corrected) + assert is_gen is False + expected = (old * 0.25 + corrected * 0.75).clamp(0, 1) # w=3/4 at i=6,N=4 + assert torch.allclose(frame, expected) + assert torch.equal(anchor, old) + + +def test_run_gen_returns_is_gen_true(): + corrected, fresh = _img(0.4), _img(0.8) + frame, anchor, is_gen = Cadence().run(4, 4, 12, False, corrected=corrected, fresh=fresh) + assert is_gen is True + assert torch.equal(anchor, corrected) + + +def test_node_has_no_widgets(): + # serialization-trap guard: every input is forceInput or link-only + # -> the frontend derives ZERO widgets -> widgets_values=[] always aligns + for name, spec in Cadence.INPUT_TYPES()["required"].items(): + opts = spec[1] if len(spec) > 1 else {} + assert spec[0] == "IMAGE" or opts.get("forceInput") is True, name + + +def test_lazy_flags_on_all_image_inputs(): + req = Cadence.INPUT_TYPES()["required"] + for name in ("warped_old", "corrected", "fresh"): + assert req[name][1].get("lazy") is True, name + + +def test_pack_registers_cadence(): + import __init__ as pack + assert "Cadence| Deforum_Stat-Imp" in pack.NODE_CLASS_MAPPINGS + assert pack.NODE_CLASS_MAPPINGS["Cadence| Deforum_Stat-Imp"] is Cadence + + +def test_flow_optional_present_and_non_lazy(): + # flow is NON-lazy on purpose so an UNCONNECTED flow doesn't crash v1 workflows + opt = Cadence.INPUT_TYPES().get("optional", {}) + assert opt["flow"][0] == "FLOW" + assert (len(opt["flow"]) < 2) or (opt["flow"][1].get("lazy") is not True) + assert opt["flow_factor"][0] == "FLOAT" + + +def test_lazy_never_requests_flow(): + # flow is non-lazy -> check_lazy_status must never list it (would crash if unconnected) + for args in [(5, 4, 12, False), (4, 4, 12, False), (0, 4, 12, True)]: + assert "flow" not in Cadence().check_lazy_status(*args) + + +def test_gen_frame_ignores_flow(): + import torch + ramp = (torch.arange(8, dtype=torch.float32).repeat(8, 1) / 7.0)[None, :, :, None].repeat(1, 1, 1, 3) + flow = torch.zeros(1, 8, 8, 2); flow[..., 0] = 2.0 + # gen iteration (i=4,N=4): flow must be ignored -> equals plain blend + plain = Cadence().run(4, 4, 12, False, corrected=ramp, fresh=ramp)[0] + withflow = Cadence().run(4, 4, 12, False, corrected=ramp, fresh=ramp, flow=flow, flow_factor=1.0)[0] + assert torch.equal(plain, withflow) + + +def test_run_with_flow_none_matches_v1(): + old, corrected = _img(0.2), _img(0.6) + a = Cadence().run(6, 4, 12, False, warped_old=old, corrected=corrected)[0] + b = Cadence().run(6, 4, 12, False, warped_old=old, corrected=corrected, + flow=None, flow_factor=1.0)[0] + assert torch.equal(a, b) + + +def test_run_with_flow_morphs(): + # ramp content so a horizontal warp changes pixels + ramp = (torch.arange(8, dtype=torch.float32).repeat(8, 1) / 7.0)[None, :, :, None].repeat(1, 1, 1, 3) + flow = torch.zeros(1, 8, 8, 2); flow[..., 0] = 2.0 + plain = Cadence().run(6, 4, 12, False, warped_old=ramp, corrected=ramp)[0] + morph = Cadence().run(6, 4, 12, False, warped_old=ramp, corrected=ramp, + flow=flow, flow_factor=1.0)[0] + assert not torch.allclose(plain, morph, atol=1e-3) diff --git a/tests/test_flow_core.py b/tests/test_flow_core.py new file mode 100644 index 0000000..71a0ab2 --- /dev/null +++ b/tests/test_flow_core.py @@ -0,0 +1,35 @@ +"""Tests for OpenCV optical-flow estimators (DIS/Farneback) — numpy BGR uint8 in, (H,W,2) out.""" +import cv2 +import numpy as np + +from flow_core import compute_flow_dis, compute_flow_farneback + + +def _shifted_pair(dx=4): + # blurred noise = spatially-coherent texture that DIS/Farneback can actually track + # (per-pixel white noise is pathological for optical flow — aperture problem everywhere) + rng = np.random.default_rng(0) + base = (rng.random((64, 64)) * 255).astype(np.uint8) + base = cv2.GaussianBlur(base, (0, 0), sigmaX=3.0) + base = np.stack([base] * 3, axis=-1) # HxWx3 BGR + shifted = np.roll(base, dx, axis=1) # content moves right by dx + return base, shifted + + +def test_dis_medium_recovers_horizontal_shift(): + a, b = _shifted_pair(4) + flow = compute_flow_dis(a, b, "DIS Medium") # flow a->b + cx = flow[20:44, 20:44, 0].mean() # central region avoids wrap edges + assert cx > 2.0 # detects rightward motion of several px + + +def test_dis_fine_runs_and_shapes(): + a, b = _shifted_pair(4) + flow = compute_flow_dis(a, b, "DIS Fine") + assert flow.shape == (64, 64, 2) and flow.dtype == np.float32 + + +def test_farneback_recovers_shift(): + a, b = _shifted_pair(4) + flow = compute_flow_farneback(a, b) + assert flow[20:44, 20:44, 0].mean() > 2.0 diff --git a/tests/test_flow_nodes.py b/tests/test_flow_nodes.py new file mode 100644 index 0000000..7f5db62 --- /dev/null +++ b/tests/test_flow_nodes.py @@ -0,0 +1,61 @@ +"""Tests for ComputeFlow node + estimate_flow dispatch (IMAGE<->core conversion, FLOW out).""" +import cv2 +import numpy as np +import torch + +from flow_nodes import ComputeFlow, estimate_flow + + +def _img_pair(dx=4): + rng = np.random.default_rng(2) + base = (rng.random((48, 48)).astype(np.float32) * 255) + base = cv2.GaussianBlur(base, (0, 0), sigmaX=3.0) / 255.0 + base = base[None, :, :, None].repeat(3, axis=-1) # ComfyUI IMAGE BHWC RGB 0..1 + a = torch.from_numpy(base.copy()) + b = torch.from_numpy(np.roll(base, dx, axis=2).copy()) # shift right in x + return a, b + + +def test_none_returns_zero_flow_matching_size(): + a, b = _img_pair() + flow = estimate_flow(a, b, "None", "cpu") + assert flow.shape == (1, 48, 48, 2) + assert torch.count_nonzero(flow) == 0 + + +def test_dis_medium_returns_flow_tensor_with_motion(): + a, b = _img_pair(4) + flow = estimate_flow(a, b, "DIS Medium", "cpu") + assert flow.shape == (1, 48, 48, 2) + assert flow[0, 16:32, 16:32, 0].mean().item() > 1.5 + + +def test_node_run_returns_flow_tuple(): + a, b = _img_pair(4) + (flow,) = ComputeFlow().run(a, b, "None", "AUTO") + assert isinstance(flow, torch.Tensor) and flow.shape[-1] == 2 + + +def test_node_return_type_is_flow(): + assert ComputeFlow.RETURN_TYPES == ("FLOW",) + + +def test_method_choices_match_vanilla(): + choices = ComputeFlow.INPUT_TYPES()["required"]["method"][0] + assert choices == ["None", "RAFT", "DIS Medium", "DIS Fine", "Farneback"] + + +def test_mismatched_sizes_do_not_crash(): + # bootstrap/first-span frames can feed a smaller old-anchor than the corrected image; + # estimate_flow must resize image_from -> image_to and return flow sized to image_to. + a = torch.rand(1, 24, 24, 3) + b = torch.rand(1, 48, 48, 3) + for method in ("None", "DIS Medium", "Farneback"): + flow = estimate_flow(a, b, method, "cpu") + assert flow.shape == (1, 48, 48, 2), method + + +def test_pack_registers_computeflow(): + import __init__ as pack + assert "ComputeFlow| Deforum_Stat-Imp" in pack.NODE_CLASS_MAPPINGS + assert pack.NODE_CLASS_MAPPINGS["ComputeFlow| Deforum_Stat-Imp"] is ComputeFlow diff --git a/tests/test_flow_warp.py b/tests/test_flow_warp.py new file mode 100644 index 0000000..e1edb20 --- /dev/null +++ b/tests/test_flow_warp.py @@ -0,0 +1,33 @@ +"""Tests for the torch flow-warp (backward sampling, vanilla remap semantics).""" +import torch + +from cadence_core import flow_warp + + +def _img(): + # 1x5x5x3 ramp so horizontal shifts are visible + x = torch.arange(5, dtype=torch.float32).repeat(5, 1) # cols vary 0..4 + return (x / 4.0)[None, :, :, None].repeat(1, 1, 1, 3) + + +def test_zero_flow_is_identity(): + img = _img() + flow = torch.zeros(1, 5, 5, 2) + out = flow_warp(img, flow) + assert torch.allclose(out, img, atol=1e-5) + + +def test_constant_shift_moves_content_left(): + # flow_warp samples img at (x - f); f=+1 in x => out(x)=img(x-1) => content shifts RIGHT by 1 + img = _img() + flow = torch.zeros(1, 5, 5, 2) + flow[..., 0] = 1.0 + out = flow_warp(img, flow) + # interior column 2 should now equal original column 1 + assert torch.allclose(out[0, 2, 2], img[0, 2, 1], atol=1e-5) + + +def test_preserves_shape_and_dtype(): + img = _img() + out = flow_warp(img, torch.zeros(1, 5, 5, 2)) + assert out.shape == img.shape and out.dtype == img.dtype diff --git a/tests/test_raft_core.py b/tests/test_raft_core.py new file mode 100644 index 0000000..c36abe5 --- /dev/null +++ b/tests/test_raft_core.py @@ -0,0 +1,19 @@ +"""RAFT smoke test (torchvision, CPU). Downloads weights on first run.""" +import cv2 +import numpy as np +import pytest + +raft = pytest.importorskip("raft_core") + + +@pytest.mark.slow +def test_raft_recovers_horizontal_shift(): + # RAFT downsamples by 8 and needs feature maps >= 16 -> images must be >= 128px + rng = np.random.default_rng(1) + base = (rng.random((128, 128)) * 255).astype(np.uint8) + base = cv2.GaussianBlur(base, (0, 0), sigmaX=4.0) + base = np.stack([base] * 3, axis=-1) + shifted = np.roll(base, 4, axis=1) + flow = raft.raft_flow(base, shifted, device="cpu") # a->b + assert flow.shape == (128, 128, 2) + assert flow[40:88, 40:88, 0].mean() > 1.5