230 lines
9.4 KiB
Python
230 lines
9.4 KiB
Python
# File: ComfyUI/custom_nodes/batch_camera_shake.py
|
|
|
|
import math
|
|
import random
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
class BatchCameraShake:
|
|
"""
|
|
Apply subtle per-frame affine transforms (translate/rotate/scale) to an IMAGE batch.
|
|
Designed for 'handheld' micro-movement.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", {
|
|
"tooltip": "Input image batch to apply camera shake effect to"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0, "min": 0, "max": 2**31-1,
|
|
"tooltip": "Random seed for reproducible shake patterns. Same seed = same shake motion"
|
|
}),
|
|
"translate_px": ("FLOAT", {
|
|
"default": 2.0, "min": 0.0, "max": 50.0, "step": 0.1,
|
|
"tooltip": "Maximum horizontal/vertical movement in pixels. Higher = more pronounced shake"
|
|
}),
|
|
"rotate_deg": ("FLOAT", {
|
|
"default": 0.15, "min": 0.0, "max": 10.0, "step": 0.01,
|
|
"tooltip": "Maximum rotation angle in degrees. Subtle values (0.1-0.3) look natural"
|
|
}),
|
|
"scale_pct": ("FLOAT", {
|
|
"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.1,
|
|
"tooltip": "Maximum zoom variation as a percentage. Simulates camera moving closer/further"
|
|
}),
|
|
"freq": ("FLOAT", {
|
|
"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1,
|
|
"tooltip": "Shake frequency multiplier. Higher = faster, more frequent shake changes"
|
|
}),
|
|
"smoothing": ("FLOAT", {
|
|
"default": 0.85, "min": 0.0, "max": 0.99, "step": 0.01,
|
|
"tooltip": "Motion smoothness. Higher = smoother, floaty movement. Lower = jittery, nervous shake"
|
|
}),
|
|
"distribution": (["even", "center", "ramp_up", "ramp_down", "ease_in", "ease_out", "ease_in_out"], {
|
|
"default": "even",
|
|
"tooltip": "Intensity distribution over time. even=uniform, center=peaks in middle, ease_out=settles to origin"
|
|
}),
|
|
"min_intensity": ("FLOAT", {
|
|
"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
|
|
"tooltip": "Minimum intensity as fraction of maximum (0-1). Only active when distribution is not 'even'. 0=full range, 1=no variation"
|
|
}),
|
|
"border": (["reflect", "replicate", "zeros"], {
|
|
"default": "reflect",
|
|
"tooltip": "Edge handling: reflect mirrors edges, replicate extends edge pixels, zeros fills with black"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "shake"
|
|
CATEGORY = "image/transform"
|
|
|
|
def _smooth_noise(self, n, rng, smoothing):
|
|
# Simple 1D noise smoothed by exponential moving average
|
|
x = np.array([rng.uniform(-1.0, 1.0) for _ in range(n)], dtype=np.float32)
|
|
y = np.zeros_like(x)
|
|
a = smoothing
|
|
y[0] = x[0]
|
|
for i in range(1, n):
|
|
y[i] = a * y[i-1] + (1 - a) * x[i]
|
|
# Normalize using standard deviation for consistent amplitude across seeds
|
|
# This ensures different seeds produce similar shake intensity
|
|
std = float(np.std(y)) if n > 1 else 1.0
|
|
if std < 1e-6:
|
|
std = 1.0
|
|
# Scale so ~95% of values fall within [-1, 1] (2 standard deviations)
|
|
return np.clip(y / (2.0 * std), -1.0, 1.0)
|
|
|
|
def _distribution_envelope(self, n, mode, min_envelope=0.5):
|
|
"""
|
|
Create an intensity envelope based on the distribution mode.
|
|
Returns array of shape (n,) with values in [min_envelope, 1.0].
|
|
The min_envelope parameter controls the lowest point as a fraction of maximum.
|
|
When mode is 'even', min_envelope is ignored and all values are 1.0.
|
|
"""
|
|
if n <= 1:
|
|
return np.ones(n, dtype=np.float32)
|
|
|
|
# Normalized position from 0.0 (first frame) to 1.0 (last frame)
|
|
t = np.linspace(0.0, 1.0, n, dtype=np.float32)
|
|
|
|
if mode == "even":
|
|
# Even distribution ignores min_envelope - always uniform at 1.0
|
|
return np.ones(n, dtype=np.float32)
|
|
elif mode == "center":
|
|
# Bell curve: peaks at middle
|
|
envelope = np.sin(np.pi * t)
|
|
elif mode == "ramp_up":
|
|
# Linear increase
|
|
envelope = t
|
|
elif mode == "ramp_down":
|
|
# Linear decrease
|
|
envelope = 1.0 - t
|
|
elif mode == "ease_in":
|
|
# Quadratic: slow start, accelerates
|
|
envelope = t ** 2
|
|
elif mode == "ease_out":
|
|
# Quadratic: fast start, decelerates (settles)
|
|
envelope = (1.0 - t) ** 2
|
|
elif mode == "ease_in_out":
|
|
# Smooth S-curve
|
|
envelope = 0.5 * (1.0 - np.cos(np.pi * t))
|
|
else:
|
|
# Fallback to even
|
|
return np.ones(n, dtype=np.float32)
|
|
|
|
# Scale from [0, 1] to [min_envelope, 1] so minimum is min_envelope of maximum
|
|
return min_envelope + (1.0 - min_envelope) * envelope
|
|
|
|
def shake(self, images, seed, translate_px, rotate_deg, scale_pct, freq, smoothing, distribution, min_intensity, border):
|
|
"""
|
|
images: torch float32 in [0,1], shape [B,H,W,C] (ComfyUI IMAGE)
|
|
freq: higher values = faster/more frequent shake changes
|
|
distribution: intensity curve over time (even, center, ramp_up, etc.)
|
|
min_intensity: minimum intensity as fraction of max (ignored when distribution='even')
|
|
"""
|
|
device = images.device
|
|
B, H, W, C = images.shape
|
|
|
|
rng = random.Random(int(seed))
|
|
|
|
# Generate smooth per-frame signals for tx, ty, rot, scale
|
|
# freq controls how quickly shake changes: higher freq = more variation over time
|
|
# We generate noise at an effective length, then interpolate to B frames
|
|
effective_len = max(int(B * freq), 2)
|
|
|
|
nx = self._smooth_noise(effective_len, rng, smoothing)
|
|
ny = self._smooth_noise(effective_len, rng, smoothing)
|
|
nr = self._smooth_noise(effective_len, rng, smoothing)
|
|
ns = self._smooth_noise(effective_len, rng, smoothing)
|
|
|
|
# Convert user amplitudes into actual values
|
|
tx_raw = (translate_px * nx).astype(np.float32)
|
|
ty_raw = (translate_px * ny).astype(np.float32)
|
|
rot_raw = (rotate_deg * nr).astype(np.float32)
|
|
sc_raw = ((scale_pct / 100.0) * ns).astype(np.float32) # Store as offset from 1.0
|
|
|
|
# Interpolate from effective_len samples to B output frames
|
|
# This ensures smooth transitions regardless of freq value
|
|
u = np.linspace(0, effective_len - 1, B, dtype=np.float32)
|
|
|
|
def lerp(a):
|
|
i0 = np.floor(u).astype(np.int32)
|
|
i1 = np.clip(i0 + 1, 0, len(a) - 1)
|
|
f = (u - i0).astype(np.float32)
|
|
return a[i0] * (1 - f) + a[i1] * f
|
|
|
|
tx, ty, rot, sc_offset = lerp(tx_raw), lerp(ty_raw), lerp(rot_raw), lerp(sc_raw)
|
|
|
|
# Apply distribution envelope to shape intensity over time
|
|
envelope = self._distribution_envelope(B, distribution, min_intensity)
|
|
tx = tx * envelope
|
|
ty = ty * envelope
|
|
rot = rot * envelope
|
|
sc = 1.0 + (sc_offset * envelope) # Scale returns to 1.0 where envelope is 0
|
|
|
|
# Prepare for grid_sample: NCHW
|
|
x = images.permute(0, 3, 1, 2).contiguous()
|
|
|
|
# Border handling
|
|
if border in ("reflect", "replicate"):
|
|
pad = int(math.ceil(max(translate_px, 2.0) + max(W, H) * (scale_pct / 100.0)))
|
|
# keep pad reasonable
|
|
pad = max(4, min(pad, 256))
|
|
mode = "reflect" if border == "reflect" else "replicate"
|
|
x = F.pad(x, (pad, pad, pad, pad), mode=mode)
|
|
Hp, Wp = x.shape[-2], x.shape[-1]
|
|
else:
|
|
pad = 0
|
|
Hp, Wp = H, W
|
|
|
|
# Build per-frame affine matrices in normalized coords
|
|
# grid_sample expects theta mapping output -> input in normalized [-1,1] coords.
|
|
thetas = []
|
|
for i in range(B):
|
|
a = math.radians(float(rot[i]))
|
|
s = float(sc[i])
|
|
|
|
cos_a = math.cos(a) / s
|
|
sin_a = math.sin(a) / s
|
|
|
|
# translation in pixels -> normalized
|
|
# Note: after padding, use Wp/Hp
|
|
tx_n = 2.0 * float(tx[i]) / max(Wp - 1, 1)
|
|
ty_n = 2.0 * float(ty[i]) / max(Hp - 1, 1)
|
|
|
|
theta = torch.tensor([
|
|
[cos_a, -sin_a, tx_n],
|
|
[sin_a, cos_a, ty_n],
|
|
], dtype=torch.float32, device=device)
|
|
thetas.append(theta)
|
|
|
|
theta = torch.stack(thetas, dim=0) # [B,2,3]
|
|
|
|
# Create sampling grid and apply transform
|
|
grid = F.affine_grid(theta, size=x.shape, align_corners=False)
|
|
y = F.grid_sample(
|
|
x, grid,
|
|
mode="bilinear",
|
|
padding_mode="zeros", # already padded if needed
|
|
align_corners=False
|
|
)
|
|
|
|
# Remove padding and return to ComfyUI IMAGE format
|
|
if pad > 0:
|
|
y = y[:, :, pad:pad+H, pad:pad+W]
|
|
|
|
out = y.permute(0, 2, 3, 1).contiguous().clamp(0.0, 1.0)
|
|
return (out,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"BatchCameraShake": BatchCameraShake
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"BatchCameraShake": "Batch Camera Shake (Subtle)"
|
|
}
|