409 lines
16 KiB
Python
409 lines
16 KiB
Python
import math
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
|
||
class EditRouterAdvanced:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"images": ("IMAGE",),
|
||
|
||
"fps": ("INT", {"default": 16, "min": 1, "max": 240, "step": 1}),
|
||
|
||
"metric": (["mean_abs", "mean_squared", "max_abs"], {"default": "mean_abs"}),
|
||
"channel_mode": (["luma", "rgb"], {"default": "luma"}),
|
||
"downsample_long_edge": ("INT", {"default": 256, "min": 0, "max": 4096, "step": 16}),
|
||
"blur_kernel": ("INT", {"default": 3, "min": 0, "max": 31, "step": 1}),
|
||
"smooth_window_frames": ("INT", {"default": 3, "min": 1, "max": 31, "step": 1}),
|
||
|
||
"threshold_mode": (["manual", "auto_mad"], {"default": "auto_mad"}),
|
||
|
||
"high_threshold": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||
"low_threshold": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||
|
||
"mad_high_k": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 50.0, "step": 0.1}),
|
||
"mad_low_k": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 50.0, "step": 0.1}),
|
||
|
||
"gate_mode": (["score_only", "score_and_ratio"], {"default": "score_and_ratio"}),
|
||
"pixel_delta_threshold": ("FLOAT", {"default": 0.20, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"ratio_high": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"ratio_low": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
|
||
"stable_time_s": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 5.0, "step": 0.01}),
|
||
"min_transition_time_s": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 5.0, "step": 0.01}),
|
||
|
||
"index_base": (["0", "1"], {"default": "0"}),
|
||
}
|
||
}
|
||
|
||
# Only normal (non-list) outputs to keep IMAGE routing predictable
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "INT", "INT", "FLOAT", "FLOAT", "STRING")
|
||
RETURN_NAMES = ("before_edit", "during_edit", "after_edit",
|
||
"edit_start_frame", "edit_end_frame",
|
||
"peak_score", "peak_ratio", "report")
|
||
|
||
FUNCTION = "split"
|
||
CATEGORY = "video/analysis"
|
||
|
||
def _ensure_batch_bhwc(self, images):
|
||
if isinstance(images, torch.Tensor):
|
||
if images.dim() == 3:
|
||
return images.unsqueeze(0)
|
||
return images
|
||
if isinstance(images, (list, tuple)) and len(images) > 0 and isinstance(images[0], torch.Tensor):
|
||
if images[0].dim() == 3:
|
||
return torch.stack(images, dim=0)
|
||
raise TypeError("Expected 'images' as torch.Tensor [B,H,W,C] or list of [H,W,C] tensors.")
|
||
|
||
def _to_luma(self, bhwc):
|
||
rgb = bhwc[..., :3]
|
||
r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
|
||
y = 0.2126 * r + 0.7152 * g + 0.0722 * b
|
||
return y.unsqueeze(-1)
|
||
|
||
def _downsample(self, bhwc, long_edge):
|
||
if long_edge <= 0:
|
||
return bhwc
|
||
b, h, w, c = bhwc.shape
|
||
cur_long = max(h, w)
|
||
if cur_long <= long_edge:
|
||
return bhwc
|
||
scale = long_edge / float(cur_long)
|
||
nh = max(1, int(round(h * scale)))
|
||
nw = max(1, int(round(w * scale)))
|
||
x = bhwc.permute(0, 3, 1, 2)
|
||
x = F.interpolate(x, size=(nh, nw), mode="area")
|
||
return x.permute(0, 2, 3, 1)
|
||
|
||
def _blur(self, bhwc, k):
|
||
if k <= 1:
|
||
return bhwc
|
||
if (k % 2) == 0:
|
||
k += 1
|
||
x = bhwc.permute(0, 3, 1, 2)
|
||
x = F.avg_pool2d(x, kernel_size=k, stride=1, padding=k // 2)
|
||
return x.permute(0, 2, 3, 1)
|
||
|
||
def _score(self, a, b, metric):
|
||
diff = a - b
|
||
if metric == "mean_abs":
|
||
return diff.abs().mean()
|
||
if metric == "mean_squared":
|
||
return (diff * diff).mean()
|
||
if metric == "max_abs":
|
||
return diff.abs().max()
|
||
raise ValueError(f"Unknown metric: {metric}")
|
||
|
||
def _smooth_1d(self, x, w):
|
||
if w <= 1 or x.numel() == 0:
|
||
return x
|
||
w = int(w)
|
||
if w % 2 == 0:
|
||
w += 1
|
||
k = torch.ones((1, 1, w), device=x.device, dtype=torch.float32) / float(w)
|
||
xx = x.float().view(1, 1, -1)
|
||
left = w // 2
|
||
right = w - 1 - left
|
||
xx = F.pad(xx, (left, right), mode="replicate")
|
||
yy = F.conv1d(xx, k)
|
||
return yy.view(-1)
|
||
|
||
def _median_mad_thresholds(self, scores_1_to_end, mad_high_k, mad_low_k):
|
||
if scores_1_to_end.numel() == 0:
|
||
return 0.0, 0.0, 0.0, 0.0
|
||
med = torch.median(scores_1_to_end)
|
||
mad = torch.median((scores_1_to_end - med).abs())
|
||
sigma = 1.4826 * mad
|
||
high = med + mad_high_k * sigma
|
||
low = med + mad_low_k * sigma
|
||
return float(med.item()), float(mad.item()), float(high.item()), float(low.item())
|
||
|
||
def split(
|
||
self,
|
||
images,
|
||
fps,
|
||
metric,
|
||
channel_mode,
|
||
downsample_long_edge,
|
||
blur_kernel,
|
||
smooth_window_frames,
|
||
threshold_mode,
|
||
high_threshold,
|
||
low_threshold,
|
||
mad_high_k,
|
||
mad_low_k,
|
||
gate_mode,
|
||
pixel_delta_threshold,
|
||
ratio_high,
|
||
ratio_low,
|
||
stable_time_s,
|
||
min_transition_time_s,
|
||
index_base,
|
||
):
|
||
orig = self._ensure_batch_bhwc(images).float().clamp(0.0, 1.0).contiguous()
|
||
b, h, w, c = orig.shape
|
||
|
||
# safer empty batch: preserves dtype/device/layout expectations
|
||
empty = orig[:0]
|
||
|
||
base = 1 if index_base == "1" else 0
|
||
fps = max(1, int(fps))
|
||
|
||
if b < 2:
|
||
report = f"Edit Router (Advanced): batch={b} (not enough frames). Passing through to before_edit."
|
||
return (orig, empty, empty, -1, -1, 0.0, 0.0, report)
|
||
|
||
proc = orig
|
||
if channel_mode == "luma" and proc.shape[-1] >= 3:
|
||
proc = self._to_luma(proc)
|
||
proc = self._downsample(proc, downsample_long_edge)
|
||
proc = self._blur(proc, blur_kernel)
|
||
|
||
scores = torch.zeros((b,), device=proc.device, dtype=torch.float32)
|
||
ratios = torch.zeros((b,), device=proc.device, dtype=torch.float32)
|
||
|
||
for i in range(1, b):
|
||
scores[i] = self._score(proc[i], proc[i - 1], metric)
|
||
abs_diff = (proc[i] - proc[i - 1]).abs()
|
||
ratios[i] = (abs_diff > pixel_delta_threshold).float().mean()
|
||
|
||
scores = self._smooth_1d(scores, smooth_window_frames)
|
||
ratios = self._smooth_1d(ratios, smooth_window_frames)
|
||
|
||
auto_med = auto_mad = None
|
||
if threshold_mode == "auto_mad":
|
||
med, mad, auto_high, auto_low = self._median_mad_thresholds(scores[1:], mad_high_k, mad_low_k)
|
||
auto_med, auto_mad = med, mad
|
||
high = auto_high
|
||
low = min(auto_low, auto_high)
|
||
else:
|
||
high = float(high_threshold)
|
||
low = min(float(low_threshold), float(high_threshold))
|
||
|
||
stable_frames = max(1, int(math.ceil(max(0.0, stable_time_s) * fps)))
|
||
min_transition_frames = max(1, int(math.ceil(max(0.0, min_transition_time_s) * fps)))
|
||
|
||
def is_changed(i):
|
||
if float(scores[i].item()) <= high:
|
||
return False
|
||
if gate_mode == "score_only":
|
||
return True
|
||
return float(ratios[i].item()) >= ratio_high
|
||
|
||
def is_stable(i):
|
||
if float(scores[i].item()) > low:
|
||
return False
|
||
if gate_mode == "score_only":
|
||
return True
|
||
return float(ratios[i].item()) <= ratio_low
|
||
|
||
# Diagnostic peaks (helps tuning even when no edit)
|
||
max_i = int(torch.argmax(scores).item())
|
||
max_score = float(scores[max_i].item())
|
||
max_ratio = float(ratios[max_i].item())
|
||
|
||
# Find start
|
||
start = -1
|
||
for i in range(1, b):
|
||
if is_changed(i):
|
||
start = i
|
||
break
|
||
|
||
# No edit -> pass all through before_edit (your requested behavior)
|
||
if start == -1:
|
||
report = (
|
||
"Edit Router (Advanced):\n"
|
||
f"- frames: {b} | fps: {fps}\n"
|
||
f"- metric: {metric} | channel_mode: {channel_mode}\n"
|
||
f"- downsample_long_edge: {downsample_long_edge} | blur_kernel: {blur_kernel} | smooth_window_frames: {smooth_window_frames}\n"
|
||
f"- threshold_mode: {threshold_mode}\n"
|
||
+ (f"- auto baseline (median/MAD): {auto_med:.6f}/{auto_mad:.6f}\n" if threshold_mode == "auto_mad" else "")
|
||
+ f"- high/low: {high:.6f}/{low:.6f}\n"
|
||
f"- gate_mode: {gate_mode} | pixel_delta_threshold: {pixel_delta_threshold} | ratio_high/low: {ratio_high}/{ratio_low}\n"
|
||
f"- max score observed: {max_score:.6f} at frame {max_i + base} | ratio={max_ratio:.3f}\n"
|
||
"- No edit detected. Passing all frames to before_edit.\n"
|
||
f"- output lens: before={b} during=0 after=0\n"
|
||
)
|
||
return (orig, empty, empty, -1, -1, max_score, max_ratio, report)
|
||
|
||
# Find end
|
||
end = b
|
||
stable_count = 0
|
||
for i in range(start, b):
|
||
if is_stable(i):
|
||
stable_count += 1
|
||
else:
|
||
stable_count = 0
|
||
if stable_count >= stable_frames:
|
||
end = i - stable_frames + 1
|
||
break
|
||
|
||
end = max(end, start)
|
||
transition_len = max(0, end - start)
|
||
is_fade = transition_len >= min_transition_frames
|
||
|
||
if not is_fade:
|
||
before = orig[:start].contiguous()
|
||
during = empty
|
||
after = orig[start:].contiguous()
|
||
out_start = start + base
|
||
out_end = start + base
|
||
kind = "hard cut (during_edit empty)"
|
||
else:
|
||
before = orig[:start].contiguous()
|
||
during = orig[start:end].contiguous()
|
||
after = orig[end:].contiguous()
|
||
out_start = start + base
|
||
out_end = end + base
|
||
kind = f"fade/transition (during_edit len={during.shape[0]})"
|
||
|
||
peak_score = float(scores[start].item())
|
||
peak_ratio = float(ratios[start].item())
|
||
|
||
report = (
|
||
"Edit Router (Advanced):\n"
|
||
f"- frames: {b} | fps: {fps}\n"
|
||
f"- metric: {metric} | channel_mode: {channel_mode}\n"
|
||
f"- downsample_long_edge: {downsample_long_edge} | blur_kernel: {blur_kernel} | smooth_window_frames: {smooth_window_frames}\n"
|
||
f"- threshold_mode: {threshold_mode}\n"
|
||
+ (f"- auto baseline (median/MAD): {auto_med:.6f}/{auto_mad:.6f}\n" if threshold_mode == "auto_mad" else "")
|
||
+ f"- high/low: {high:.6f}/{low:.6f}\n"
|
||
f"- gate_mode: {gate_mode} | pixel_delta_threshold: {pixel_delta_threshold} | ratio_high/low: {ratio_high}/{ratio_low}\n"
|
||
f"- stable_time_s: {stable_time_s} -> stable_frames: {stable_frames}\n"
|
||
f"- min_transition_time_s: {min_transition_time_s} -> min_transition_frames: {min_transition_frames}\n"
|
||
f"- detected: {kind}\n"
|
||
f"- edit_start_frame (index_base={base}): {out_start} (score={peak_score:.6f}, ratio={peak_ratio:.3f})\n"
|
||
f"- edit_end_frame (index_base={base}): {out_end}\n"
|
||
f"- max score observed: {max_score:.6f} at frame {max_i + base} | ratio={max_ratio:.3f}\n"
|
||
f"- output lens: before={before.shape[0]} during={during.shape[0]} after={after.shape[0]}\n"
|
||
)
|
||
|
||
return (before, during, after, int(out_start), int(out_end), peak_score, peak_ratio, report)
|
||
|
||
class EditRouter(EditRouterAdvanced):
|
||
"""
|
||
Minimal splitter:
|
||
- Detect FIRST edit in an IMAGE batch.
|
||
- Output before_edit and after_edit only.
|
||
- Fade frames are REMOVED (not output).
|
||
- If no edit: before_edit = all frames, after_edit = empty.
|
||
|
||
Detection:
|
||
- luma
|
||
- downsample long edge to 256
|
||
- score[i] = mean_abs(frame[i] - frame[i-1])
|
||
- trigger at threshold (high)
|
||
- built-in hysteresis: low = 0.5 * high
|
||
- trims a short fade by rewinding start while score>low
|
||
- ends transition when score<=low for stable_time_s
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"images": ("IMAGE",),
|
||
"threshold": ("FLOAT", {"default": 0.09, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"stable_frames": ("INT", {"default": 3, "min": 0, "max": 81, "step": 1}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE")
|
||
RETURN_NAMES = ("before_edit", "after_edit")
|
||
FUNCTION = "split"
|
||
CATEGORY = "video/analysis"
|
||
|
||
def split(self, images, threshold, stable_frames):
|
||
orig = self._ensure_batch_bhwc(images).float().clamp(0.0, 1.0).contiguous()
|
||
b, h, w, c = orig.shape
|
||
empty = orig[:0]
|
||
|
||
if b < 2:
|
||
return (orig, empty)
|
||
|
||
fps = 16 # Wan default
|
||
high = float(threshold)
|
||
low = 0.5 * high # slightly more aggressive fade trimming than 0.6
|
||
|
||
# detection preprocess only
|
||
proc = orig
|
||
if proc.shape[-1] >= 3:
|
||
proc = self._to_luma(proc)
|
||
proc = self._downsample(proc, 256)
|
||
|
||
scores = torch.zeros((b,), device=proc.device, dtype=torch.float32)
|
||
for i in range(1, b):
|
||
scores[i] = (proc[i] - proc[i - 1]).abs().mean()
|
||
|
||
# find first trigger
|
||
start = -1
|
||
for i in range(1, b):
|
||
if float(scores[i].item()) > high:
|
||
start = i
|
||
break
|
||
|
||
# no edit -> pass-through to before
|
||
if start == -1:
|
||
return (orig, empty)
|
||
|
||
# rewind to include fade ramp (cap rewind so we don't eat into normal motion)
|
||
max_rewind = max(2, int(round(0.5 * fps)))
|
||
rewind = 0
|
||
while start > 1 and rewind < max_rewind and float(scores[start - 1].item()) >= low:
|
||
start -= 1
|
||
rewind += 1
|
||
|
||
# --- robust "normal motion" threshold (independent of user threshold) ---
|
||
s = scores[1:] # ignore dummy 0
|
||
med = torch.median(s)
|
||
mad = torch.median((s - med).abs())
|
||
sigma = 1.4826 * mad
|
||
normal_high = float((med + 4.0 * sigma).item())
|
||
|
||
# don't let normal_high be ridiculously tiny
|
||
normal_high = max(normal_high, low)
|
||
|
||
stable_frames = max(1, int(stable_frames))
|
||
|
||
# If we’re already back to normal motion right after the detected boundary,
|
||
# treat it as a hard cut (keep frame 'start' as the first after frame).
|
||
if (start + 1) >= b or float(scores[start + 1].item()) <= normal_high:
|
||
end = start
|
||
else:
|
||
# Otherwise, treat as a transition and remove frames until motion returns to normal.
|
||
end = b
|
||
stable_count = 0
|
||
for i in range(start + 1, b):
|
||
if float(scores[i].item()) <= normal_high:
|
||
stable_count += 1
|
||
else:
|
||
stable_count = 0
|
||
if stable_count >= stable_frames:
|
||
end = i - stable_frames + 1
|
||
break
|
||
|
||
# If we never found stability, fallback so after isn't empty.
|
||
if end == b:
|
||
end = start
|
||
|
||
end = max(end, start)
|
||
|
||
before = orig[:start].contiguous()
|
||
after = orig[end:].contiguous()
|
||
return (before, after)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
# "CG_EditRouterAdvanced": EditRouterAdvanced,
|
||
"CG_EditRouter": EditRouter,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
# "CG_EditRouterAdvanced": "Edit Router (Advanced)",
|
||
"CG_EditRouter": "Edit Router",
|
||
}
|
||
|
||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|