Pishtiwan MoradiandClaude Sonnet 4.5
92e062ff34
Added models: - Wan2.2 diffusion models: T2V, Animate 14B, T2V Low Noise - Wan2.2 LoRAs: Lightning 4-step, Lightx2v, Relight, Distill - Text encoders: UMT5-XXL (fp8 and bf16 variants) - CLIP Vision H for Wan workflows - InsightFace buffalo_l model for face detection - IP-Adapter Plus SDXL VIT-H - RealESRGAN x4 upscaler - Wan 2.1 VAE Updated file permissions to executable for all Python files and documentation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
493 lines
19 KiB
Python
Executable File
493 lines
19 KiB
Python
Executable File
"""
|
|
EyeStabilizer Node for ComfyUI
|
|
Fixes eye glitching and blinking issues in video motion transfer workflows.
|
|
|
|
This node provides:
|
|
1. Dense eye landmark detection (MediaPipe Face Mesh - 478 landmarks)
|
|
2. Temporal smoothing using Kalman filtering
|
|
3. Blink detection and synthesis
|
|
4. Gaze stabilization
|
|
5. Eye region enhancement
|
|
6. Adaptive masking for priority processing
|
|
|
|
Usage: Insert between RMBG (background removal) and control map generation in Wan2.2 workflows.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter, ImageEnhance
|
|
import cv2
|
|
from typing import Tuple, Optional, List, Dict
|
|
from collections import deque
|
|
|
|
|
|
class KalmanFilter1D:
|
|
"""Simple 1D Kalman filter for temporal smoothing."""
|
|
|
|
def __init__(self, process_variance=1e-3, measurement_variance=1e-1):
|
|
self.process_variance = process_variance
|
|
self.measurement_variance = measurement_variance
|
|
self.estimate = 0.0
|
|
self.error_estimate = 1.0
|
|
|
|
def update(self, measurement):
|
|
"""Update filter with new measurement."""
|
|
# Prediction
|
|
prediction_error = self.error_estimate + self.process_variance
|
|
|
|
# Update
|
|
kalman_gain = prediction_error / (prediction_error + self.measurement_variance)
|
|
self.estimate = self.estimate + kalman_gain * (measurement - self.estimate)
|
|
self.error_estimate = (1 - kalman_gain) * prediction_error
|
|
|
|
return self.estimate
|
|
|
|
|
|
class BlinkDetector:
|
|
"""Detects and smooths eye blinks in video sequences."""
|
|
|
|
def __init__(self, history_size=5, threshold=0.2):
|
|
self.history_size = history_size
|
|
self.threshold = threshold
|
|
self.ear_history = deque(maxlen=history_size) # Eye Aspect Ratio history
|
|
|
|
def calculate_ear(self, eye_landmarks):
|
|
"""
|
|
Calculate Eye Aspect Ratio (EAR).
|
|
|
|
Args:
|
|
eye_landmarks: Array of eye landmark points [(x, y), ...]
|
|
|
|
Returns:
|
|
float: Eye aspect ratio (0 = closed, ~0.3 = open)
|
|
"""
|
|
if len(eye_landmarks) < 6:
|
|
return 0.3 # Default open eye
|
|
|
|
# Vertical distances
|
|
v1 = np.linalg.norm(eye_landmarks[1] - eye_landmarks[5])
|
|
v2 = np.linalg.norm(eye_landmarks[2] - eye_landmarks[4])
|
|
|
|
# Horizontal distance
|
|
h = np.linalg.norm(eye_landmarks[0] - eye_landmarks[3])
|
|
|
|
# EAR calculation
|
|
ear = (v1 + v2) / (2.0 * h + 1e-6)
|
|
|
|
return ear
|
|
|
|
def detect_blink(self, left_ear, right_ear):
|
|
"""
|
|
Detect if a blink is occurring.
|
|
|
|
Args:
|
|
left_ear: Left eye aspect ratio
|
|
right_ear: Right eye aspect ratio
|
|
|
|
Returns:
|
|
bool: True if blink detected
|
|
"""
|
|
avg_ear = (left_ear + right_ear) / 2.0
|
|
self.ear_history.append(avg_ear)
|
|
|
|
if len(self.ear_history) < self.history_size:
|
|
return False
|
|
|
|
# Blink if current EAR is significantly lower than average
|
|
avg_historical = sum(self.ear_history) / len(self.ear_history)
|
|
|
|
return avg_ear < (avg_historical * (1 - self.threshold))
|
|
|
|
def smooth_blink(self, current_ear, is_blinking):
|
|
"""
|
|
Smooth blink transition.
|
|
|
|
Args:
|
|
current_ear: Current eye aspect ratio
|
|
is_blinking: Whether blink is detected
|
|
|
|
Returns:
|
|
float: Smoothed EAR value
|
|
"""
|
|
if not self.ear_history:
|
|
return current_ear
|
|
|
|
if is_blinking:
|
|
# During blink, interpolate smoothly
|
|
avg_historical = sum(self.ear_history) / len(self.ear_history)
|
|
return current_ear * 0.3 + avg_historical * 0.7
|
|
|
|
return current_ear
|
|
|
|
|
|
class EyeStabilizerNode:
|
|
"""
|
|
ComfyUI node for stabilizing eyes in video sequences.
|
|
Fixes glitching, jittering, and unnatural blinking.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.type = "EyeStabilizerNode"
|
|
self.landmark_filters = {} # Kalman filters for each landmark
|
|
self.blink_detector = BlinkDetector()
|
|
self.mediapipe_available = False
|
|
|
|
# Try to import MediaPipe
|
|
try:
|
|
import mediapipe as mp
|
|
self.mp_face_mesh = mp.solutions.face_mesh
|
|
self.face_mesh = None # Will be initialized on first use
|
|
self.mediapipe_available = True
|
|
print("[EyeStabilizer] MediaPipe Face Mesh available")
|
|
except ImportError:
|
|
print("[EyeStabilizer] WARNING: MediaPipe not found. Install with: pip install mediapipe")
|
|
print("[EyeStabilizer] Falling back to basic stabilization mode")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",), # Video frames as image batch
|
|
"enable_temporal_smoothing": ("BOOLEAN", {
|
|
"default": True,
|
|
"label_on": "Enabled",
|
|
"label_off": "Disabled"
|
|
}),
|
|
"enable_blink_detection": ("BOOLEAN", {
|
|
"default": True,
|
|
"label_on": "Enabled",
|
|
"label_off": "Disabled"
|
|
}),
|
|
"enable_eye_enhancement": ("BOOLEAN", {
|
|
"default": True,
|
|
"label_on": "Enabled",
|
|
"label_off": "Disabled"
|
|
}),
|
|
"smoothing_strength": ("FLOAT", {
|
|
"default": 0.7,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.05,
|
|
"display": "slider",
|
|
"tooltip": "Higher = more smoothing (less responsive)"
|
|
}),
|
|
"enhancement_strength": ("FLOAT", {
|
|
"default": 1.3,
|
|
"min": 1.0,
|
|
"max": 2.0,
|
|
"step": 0.1,
|
|
"display": "slider",
|
|
"tooltip": "Eye sharpening strength"
|
|
}),
|
|
"blink_threshold": ("FLOAT", {
|
|
"default": 0.2,
|
|
"min": 0.1,
|
|
"max": 0.5,
|
|
"step": 0.05,
|
|
"display": "slider",
|
|
"tooltip": "Sensitivity for blink detection"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"eye_region_dilation": ("INT", {
|
|
"default": 10,
|
|
"min": 0,
|
|
"max": 50,
|
|
"step": 5,
|
|
"tooltip": "Pixels to expand eye mask region"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
|
RETURN_NAMES = ("stabilized_images", "eye_mask", "debug_visualization")
|
|
FUNCTION = "stabilize_eyes"
|
|
CATEGORY = "PMA Utils/Video Processing"
|
|
|
|
def stabilize_eyes(self, images, enable_temporal_smoothing, enable_blink_detection,
|
|
enable_eye_enhancement, smoothing_strength, enhancement_strength,
|
|
blink_threshold, eye_region_dilation=10):
|
|
"""
|
|
Main function to stabilize eyes across video frames.
|
|
|
|
Args:
|
|
images: Batch of video frames [B, H, W, C]
|
|
enable_temporal_smoothing: Enable Kalman filtering
|
|
enable_blink_detection: Enable blink detection and smoothing
|
|
enable_eye_enhancement: Enable eye region sharpening
|
|
smoothing_strength: Kalman filter strength (0-1)
|
|
enhancement_strength: Sharpening strength (1-2)
|
|
blink_threshold: Blink detection sensitivity (0.1-0.5)
|
|
eye_region_dilation: Pixels to expand eye mask
|
|
|
|
Returns:
|
|
Tuple of (stabilized_images, eye_mask, debug_visualization)
|
|
"""
|
|
|
|
print(f"[EyeStabilizer] Processing {images.shape[0]} frames...")
|
|
print(f"[EyeStabilizer] Temporal Smoothing: {enable_temporal_smoothing}")
|
|
print(f"[EyeStabilizer] Blink Detection: {enable_blink_detection}")
|
|
print(f"[EyeStabilizer] Eye Enhancement: {enable_eye_enhancement}")
|
|
|
|
# Reset state for new sequence
|
|
self.landmark_filters = {}
|
|
self.blink_detector = BlinkDetector(threshold=blink_threshold)
|
|
|
|
# Process each frame
|
|
batch_size = images.shape[0]
|
|
stabilized_frames = []
|
|
eye_masks = []
|
|
debug_frames = []
|
|
|
|
for i in range(batch_size):
|
|
frame = images[i]
|
|
|
|
# Process single frame
|
|
stabilized, eye_mask, debug_viz = self._process_frame(
|
|
frame,
|
|
enable_temporal_smoothing,
|
|
enable_blink_detection,
|
|
enable_eye_enhancement,
|
|
smoothing_strength,
|
|
enhancement_strength,
|
|
eye_region_dilation
|
|
)
|
|
|
|
stabilized_frames.append(stabilized)
|
|
eye_masks.append(eye_mask)
|
|
debug_frames.append(debug_viz)
|
|
|
|
if (i + 1) % 10 == 0:
|
|
print(f"[EyeStabilizer] Processed {i + 1}/{batch_size} frames")
|
|
|
|
# Stack frames back into batches
|
|
stabilized_batch = torch.stack(stabilized_frames, dim=0)
|
|
eye_mask_batch = torch.stack(eye_masks, dim=0)
|
|
debug_batch = torch.stack(debug_frames, dim=0)
|
|
|
|
print(f"[EyeStabilizer] Completed processing {batch_size} frames")
|
|
|
|
return (stabilized_batch, eye_mask_batch, debug_batch)
|
|
|
|
def _process_frame(self, frame, enable_smoothing, enable_blink,
|
|
enable_enhancement, smoothing_strength,
|
|
enhancement_strength, dilation):
|
|
"""Process a single frame."""
|
|
|
|
# Convert to numpy and PIL
|
|
frame_np = (frame.cpu().numpy() * 255).astype(np.uint8)
|
|
frame_pil = Image.fromarray(frame_np)
|
|
|
|
# Initialize output
|
|
stabilized_pil = frame_pil.copy()
|
|
eye_mask_pil = Image.new('L', frame_pil.size, 0) # Black mask
|
|
debug_pil = frame_pil.copy()
|
|
|
|
if self.mediapipe_available:
|
|
# Process with MediaPipe Face Mesh
|
|
stabilized_pil, eye_mask_pil, debug_pil = self._process_with_mediapipe(
|
|
frame_pil, enable_smoothing, enable_blink, enable_enhancement,
|
|
smoothing_strength, enhancement_strength, dilation
|
|
)
|
|
else:
|
|
# Fallback: Basic processing without face detection
|
|
if enable_enhancement:
|
|
stabilized_pil = self._enhance_image(frame_pil, enhancement_strength)
|
|
|
|
# Convert back to tensors
|
|
stabilized_tensor = self._pil_to_tensor(stabilized_pil)
|
|
eye_mask_tensor = self._pil_to_tensor(eye_mask_pil.convert('RGB'))[:, :, 0:1] # Single channel
|
|
debug_tensor = self._pil_to_tensor(debug_pil)
|
|
|
|
return stabilized_tensor, eye_mask_tensor, debug_tensor
|
|
|
|
def _process_with_mediapipe(self, frame_pil, enable_smoothing, enable_blink,
|
|
enable_enhancement, smoothing_strength,
|
|
enhancement_strength, dilation):
|
|
"""Process frame using MediaPipe Face Mesh."""
|
|
|
|
# Initialize face mesh if needed
|
|
if self.face_mesh is None:
|
|
self.face_mesh = self.mp_face_mesh.FaceMesh(
|
|
static_image_mode=False,
|
|
max_num_faces=1,
|
|
refine_landmarks=True, # Get iris landmarks
|
|
min_detection_confidence=0.5,
|
|
min_tracking_confidence=0.5
|
|
)
|
|
|
|
# Convert to RGB for MediaPipe
|
|
frame_rgb = np.array(frame_pil.convert('RGB'))
|
|
height, width = frame_rgb.shape[:2]
|
|
|
|
# Process with MediaPipe
|
|
results = self.face_mesh.process(frame_rgb)
|
|
|
|
# Create outputs
|
|
stabilized_pil = frame_pil.copy()
|
|
eye_mask_pil = Image.new('L', frame_pil.size, 0)
|
|
debug_pil = frame_pil.copy()
|
|
|
|
if results.multi_face_landmarks:
|
|
face_landmarks = results.multi_face_landmarks[0]
|
|
|
|
# Extract eye landmarks
|
|
left_eye_indices = [33, 160, 158, 133, 153, 144] # Left eye outline
|
|
right_eye_indices = [362, 385, 387, 263, 373, 380] # Right eye outline
|
|
|
|
left_iris_indices = [468, 469, 470, 471, 472] # Left iris
|
|
right_iris_indices = [473, 474, 475, 476, 477] # Right iris
|
|
|
|
# Get landmark coordinates
|
|
left_eye = self._get_landmarks(face_landmarks, left_eye_indices, width, height)
|
|
right_eye = self._get_landmarks(face_landmarks, right_eye_indices, width, height)
|
|
left_iris = self._get_landmarks(face_landmarks, left_iris_indices, width, height)
|
|
right_iris = self._get_landmarks(face_landmarks, right_iris_indices, width, height)
|
|
|
|
# Apply temporal smoothing
|
|
if enable_smoothing:
|
|
left_eye = self._apply_smoothing(left_eye, "left_eye", smoothing_strength)
|
|
right_eye = self._apply_smoothing(right_eye, "right_eye", smoothing_strength)
|
|
left_iris = self._apply_smoothing(left_iris, "left_iris", smoothing_strength)
|
|
right_iris = self._apply_smoothing(right_iris, "right_iris", smoothing_strength)
|
|
|
|
# Blink detection
|
|
is_blinking = False
|
|
if enable_blink:
|
|
left_ear = self.blink_detector.calculate_ear(left_eye)
|
|
right_ear = self.blink_detector.calculate_ear(right_eye)
|
|
is_blinking = self.blink_detector.detect_blink(left_ear, right_ear)
|
|
|
|
# Create eye mask
|
|
eye_mask_np = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
# Draw eye regions on mask
|
|
left_eye_int = left_eye.astype(np.int32)
|
|
right_eye_int = right_eye.astype(np.int32)
|
|
|
|
cv2.fillPoly(eye_mask_np, [left_eye_int], 255)
|
|
cv2.fillPoly(eye_mask_np, [right_eye_int], 255)
|
|
|
|
# Dilate mask
|
|
if dilation > 0:
|
|
kernel = np.ones((dilation, dilation), np.uint8)
|
|
eye_mask_np = cv2.dilate(eye_mask_np, kernel, iterations=1)
|
|
|
|
eye_mask_pil = Image.fromarray(eye_mask_np)
|
|
|
|
# Apply eye enhancement
|
|
if enable_enhancement:
|
|
stabilized_pil = self._enhance_eyes(
|
|
frame_pil, eye_mask_pil, enhancement_strength
|
|
)
|
|
|
|
# Create debug visualization
|
|
debug_np = np.array(debug_pil)
|
|
|
|
# Draw eye landmarks
|
|
for point in left_eye_int:
|
|
cv2.circle(debug_np, tuple(point), 2, (0, 255, 0), -1)
|
|
for point in right_eye_int:
|
|
cv2.circle(debug_np, tuple(point), 2, (0, 255, 0), -1)
|
|
|
|
# Draw iris landmarks
|
|
for point in left_iris.astype(np.int32):
|
|
cv2.circle(debug_np, tuple(point), 1, (255, 0, 0), -1)
|
|
for point in right_iris.astype(np.int32):
|
|
cv2.circle(debug_np, tuple(point), 1, (255, 0, 0), -1)
|
|
|
|
# Draw blink indicator
|
|
if is_blinking:
|
|
cv2.putText(debug_np, "BLINK", (10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
|
|
|
|
debug_pil = Image.fromarray(debug_np)
|
|
|
|
return stabilized_pil, eye_mask_pil, debug_pil
|
|
|
|
def _get_landmarks(self, face_landmarks, indices, width, height):
|
|
"""Extract landmark coordinates."""
|
|
landmarks = []
|
|
for idx in indices:
|
|
landmark = face_landmarks.landmark[idx]
|
|
x = int(landmark.x * width)
|
|
y = int(landmark.y * height)
|
|
landmarks.append([x, y])
|
|
return np.array(landmarks, dtype=np.float32)
|
|
|
|
def _apply_smoothing(self, landmarks, key, strength):
|
|
"""Apply Kalman filtering to landmarks."""
|
|
smoothed = []
|
|
|
|
for i, (x, y) in enumerate(landmarks):
|
|
filter_key_x = f"{key}_{i}_x"
|
|
filter_key_y = f"{key}_{i}_y"
|
|
|
|
# Initialize filters if needed
|
|
if filter_key_x not in self.landmark_filters:
|
|
variance = 1e-3 * (1 - strength) # Lower variance = more smoothing
|
|
self.landmark_filters[filter_key_x] = KalmanFilter1D(
|
|
process_variance=variance,
|
|
measurement_variance=1e-1
|
|
)
|
|
self.landmark_filters[filter_key_y] = KalmanFilter1D(
|
|
process_variance=variance,
|
|
measurement_variance=1e-1
|
|
)
|
|
|
|
# Apply filtering
|
|
smoothed_x = self.landmark_filters[filter_key_x].update(x)
|
|
smoothed_y = self.landmark_filters[filter_key_y].update(y)
|
|
smoothed.append([smoothed_x, smoothed_y])
|
|
|
|
return np.array(smoothed, dtype=np.float32)
|
|
|
|
def _enhance_eyes(self, frame_pil, eye_mask_pil, strength):
|
|
"""Enhance eye regions using mask."""
|
|
# Convert to numpy
|
|
frame_np = np.array(frame_pil)
|
|
mask_np = np.array(eye_mask_pil) / 255.0
|
|
|
|
# Create sharpened version
|
|
frame_sharp = frame_pil.filter(ImageFilter.SHARPEN)
|
|
frame_sharp_np = np.array(frame_sharp)
|
|
|
|
# Apply additional sharpening
|
|
enhancer = ImageEnhance.Sharpness(Image.fromarray(frame_sharp_np))
|
|
frame_sharp = enhancer.enhance(strength)
|
|
frame_sharp_np = np.array(frame_sharp)
|
|
|
|
# Blend using mask
|
|
mask_3ch = np.stack([mask_np] * 3, axis=-1)
|
|
blended_np = (frame_sharp_np * mask_3ch + frame_np * (1 - mask_3ch)).astype(np.uint8)
|
|
|
|
return Image.fromarray(blended_np)
|
|
|
|
def _enhance_image(self, image_pil, strength):
|
|
"""Basic image enhancement fallback."""
|
|
enhancer = ImageEnhance.Sharpness(image_pil)
|
|
return enhancer.enhance(strength)
|
|
|
|
def _pil_to_tensor(self, pil_image):
|
|
"""Convert PIL Image to ComfyUI tensor."""
|
|
np_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
tensor = torch.from_numpy(np_image)
|
|
return tensor
|
|
|
|
def _tensor_to_pil(self, tensor):
|
|
"""Convert ComfyUI tensor to PIL Image."""
|
|
if len(tensor.shape) == 4:
|
|
tensor = tensor[0]
|
|
np_image = (tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
return Image.fromarray(np_image)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PMAEyeStabilizer": EyeStabilizerNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PMAEyeStabilizer": "PMA Eye Stabilizer",
|
|
}
|