BISAM20andClaude Opus 4.8
5ac1aed2a8
Drop the "status" text widget (editor JS) and the "refresh" int input (Python node) to declutter the node. generate_masks gains **kwargs to absorb stray values from older saved workflows. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
564 lines
21 KiB
Python
564 lines
21 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
import json
|
|
import os
|
|
import folder_paths
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
import base64
|
|
import io
|
|
|
|
# --- Persistent per-node shape storage -------------------------------------
|
|
# Disk backup lives next to the node (NOT in ComfyUI's temp dir, which is wiped
|
|
# on restart). The widget value stored in the workflow is the primary store;
|
|
# these files are the "both" safety net and a fallback for headless runs.
|
|
BACKUP_DIR = os.path.join(os.path.dirname(__file__), "roto_data")
|
|
DEFAULT_NODE_KEY = "_default"
|
|
|
|
# In-memory mirror of the latest saved data per node, populated by the editor's
|
|
# /save route so /load can serve it without a disk round-trip.
|
|
roto_store = {}
|
|
|
|
|
|
def node_key_for(unique_id):
|
|
"""Normalize a node id (or missing id) into a safe filesystem/dict key."""
|
|
if unique_id is None or unique_id == "":
|
|
return DEFAULT_NODE_KEY
|
|
return "".join(c if (c.isalnum() or c in "-_") else "_" for c in str(unique_id))
|
|
|
|
|
|
def get_backup_path(node_key):
|
|
"""Persistent per-node backup file path."""
|
|
return os.path.join(BACKUP_DIR, f"roto_{node_key}.json")
|
|
|
|
|
|
def parse_mask_data(mask_data):
|
|
"""Parse the serialized widget string. Returns dict or None if empty/invalid."""
|
|
if not mask_data or not str(mask_data).strip():
|
|
return None
|
|
try:
|
|
data = json.loads(mask_data)
|
|
if isinstance(data, dict) and "shapes" in data:
|
|
return data
|
|
except (json.JSONDecodeError, TypeError):
|
|
pass
|
|
return None
|
|
|
|
|
|
def load_backup(node_key):
|
|
"""Load the per-node disk backup. Returns dict or None."""
|
|
path = get_backup_path(node_key)
|
|
if os.path.exists(path):
|
|
try:
|
|
with open(path, 'r') as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, OSError):
|
|
return None
|
|
return None
|
|
|
|
|
|
def save_backup(node_key, data):
|
|
"""Write the per-node disk backup and update the in-memory mirror."""
|
|
os.makedirs(BACKUP_DIR, exist_ok=True)
|
|
roto_store[node_key] = data
|
|
with open(get_backup_path(node_key), 'w') as f:
|
|
json.dump(data, f, indent=2)
|
|
|
|
|
|
def resolve_size(value, dim, default):
|
|
"""Resolve a shape size (radius/width/height) to output pixels.
|
|
|
|
Sizes are stored normalized (0-1, a fraction of the frame) so masks render
|
|
at the correct size for ANY output resolution. Legacy data stored absolute
|
|
pixels (> 1); those are passed through unchanged for backward compatibility.
|
|
"""
|
|
if value is None:
|
|
value = default
|
|
return value * dim if value <= 1.0 else value
|
|
|
|
|
|
class AnimatedMaskDrawer:
|
|
"""
|
|
ComfyUI Custom Node for drawing animated masks with keyframes
|
|
Includes integrated web UI that launches from the node
|
|
"""
|
|
|
|
# Per-node video store (keyed by node unique_id). Replaces the old single
|
|
# global so multiple Roto nodes don't clobber each other's frames.
|
|
current_videos = {}
|
|
last_video = None # fallback for calls without a node id
|
|
last_params_hash = None
|
|
|
|
def __init__(self):
|
|
self.mask_data = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video": ("IMAGE", {"tooltip": "Input video or image batch to create masks for"}),
|
|
},
|
|
"optional": {
|
|
"mask_data": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Serialized roto shapes/keyframes (hidden - managed by the editor). Stored in the workflow so shapes survive restarts."
|
|
}),
|
|
"width": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 4096,
|
|
"tooltip": "Manual output mask width. Only used when auto_resolution is off (0 = use video width)."
|
|
}),
|
|
"height": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 4096,
|
|
"tooltip": "Manual output mask height. Only used when auto_resolution is off (0 = use video height)."
|
|
}),
|
|
"feather": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 100,
|
|
"tooltip": "Blur radius for soft edges (10-20 recommended for smooth masks)"
|
|
}),
|
|
"invert": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Invert mask (swap masked/unmasked areas)"
|
|
}),
|
|
# NOTE: keep new widgets LAST. ComfyUI serializes widget values
|
|
# positionally, so inserting one mid-list shifts every saved
|
|
# value in existing workflows and breaks them.
|
|
"auto_resolution": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Automatically match the mask output to the input video's resolution. Turn off to set width/height manually."
|
|
}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("mask_sequence", "mask_preview", "masked_video")
|
|
OUTPUT_TOOLTIPS = (
|
|
"Animated mask sequence (batch) - Use for inpainting, effects, compositing",
|
|
"Video with red mask overlay - Connect to Preview Image to see masks",
|
|
"Input video with white masks overlaid (100% opacity)"
|
|
)
|
|
FUNCTION = "generate_masks"
|
|
CATEGORY = "mask/animation"
|
|
OUTPUT_NODE = False
|
|
|
|
def generate_masks(self, video, width=0, height=0, feather=0, invert=False,
|
|
auto_resolution=True, mask_data="", unique_id=None, **kwargs):
|
|
"""
|
|
Generate animated masks based on keyframe data
|
|
"""
|
|
# Store video reference per-node for the editor UI to fetch frames.
|
|
node_key = node_key_for(unique_id)
|
|
AnimatedMaskDrawer.current_videos[node_key] = video
|
|
AnimatedMaskDrawer.last_video = video
|
|
|
|
print(f"[AnimatedMaskDrawer] Video stored for node {node_key}: {video.shape}")
|
|
|
|
# Resolve shape data. Priority:
|
|
# 1. mask_data widget (serialized in the workflow -> survives restarts)
|
|
# 2. per-node disk backup file
|
|
# 3. empty
|
|
data = parse_mask_data(mask_data)
|
|
if data is None:
|
|
data = load_backup(node_key)
|
|
if data is None:
|
|
data = {"shapes": [], "keyframes": {}}
|
|
|
|
# Get video dimensions
|
|
batch_size = video.shape[0]
|
|
video_height, video_width = video.shape[1], video.shape[2]
|
|
|
|
# Resolve the output resolution. Auto-detect from the video by default;
|
|
# otherwise honour the manual width/height (0, None, or legacy 512 fall
|
|
# back to the video dimension so old workflows keep auto-detecting).
|
|
if auto_resolution is None:
|
|
auto_resolution = True
|
|
if auto_resolution or not width or width == 512:
|
|
width = video_width
|
|
if auto_resolution or not height or height == 512:
|
|
height = video_height
|
|
|
|
masks = []
|
|
preview_frames = []
|
|
masked_video_frames = []
|
|
|
|
for frame_idx in range(batch_size):
|
|
# Create blank mask - START WITH BLACK (0)
|
|
mask_img = Image.new('L', (width, height), 0)
|
|
draw = ImageDraw.Draw(mask_img)
|
|
|
|
# Get interpolated shapes for this frame
|
|
frame_shapes = self.interpolate_frame(data, frame_idx, batch_size)
|
|
|
|
# Draw each shape - shapes will be WHITE on black background
|
|
for shape in frame_shapes:
|
|
self.draw_shape(draw, shape, width, height)
|
|
|
|
# Convert to tensor
|
|
mask_array = np.array(mask_img).astype(np.float32) / 255.0
|
|
|
|
# Apply feathering if needed
|
|
if feather > 0:
|
|
mask_array = self.apply_feather(mask_array, feather)
|
|
|
|
# Invert if needed (AFTER drawing, so default shows drawn areas as white)
|
|
if invert:
|
|
mask_array = 1.0 - mask_array
|
|
|
|
masks.append(mask_array)
|
|
|
|
# Create preview frame (video with red mask overlay)
|
|
# Resize mask to match video dimensions if needed
|
|
if mask_array.shape != (video_height, video_width):
|
|
from PIL import Image as PILImage
|
|
mask_pil = PILImage.fromarray((mask_array * 255).astype(np.uint8))
|
|
mask_pil = mask_pil.resize((video_width, video_height), PILImage.LANCZOS)
|
|
mask_array_resized = np.array(mask_pil).astype(np.float32) / 255.0
|
|
else:
|
|
mask_array_resized = mask_array
|
|
|
|
frame = video[frame_idx].cpu().numpy()
|
|
|
|
# Red overlay preview
|
|
overlay = frame.copy()
|
|
overlay[:, :, 0] = np.clip(overlay[:, :, 0] + mask_array_resized * 0.4, 0, 1)
|
|
preview_frames.append(overlay)
|
|
|
|
# White mask overlay on video (100% opacity)
|
|
masked_video = frame.copy()
|
|
# Where mask is active (white), make the video white
|
|
mask_3channel = np.stack([mask_array_resized] * 3, axis=-1)
|
|
masked_video = masked_video * (1 - mask_3channel) + mask_3channel
|
|
masked_video_frames.append(masked_video)
|
|
|
|
# Stack masks into batch
|
|
mask_tensor = torch.from_numpy(np.stack(masks))
|
|
|
|
# Stack preview frames
|
|
preview_tensor = torch.from_numpy(np.stack(preview_frames))
|
|
|
|
# Stack masked video frames
|
|
masked_video_tensor = torch.from_numpy(np.stack(masked_video_frames))
|
|
|
|
return (mask_tensor, preview_tensor, masked_video_tensor)
|
|
|
|
def interpolate_frame(self, data, frame_idx, total_frames):
|
|
"""
|
|
Interpolate shape properties based on keyframes
|
|
"""
|
|
shapes = data.get("shapes", [])
|
|
keyframes = data.get("keyframes", {})
|
|
|
|
interpolated = []
|
|
|
|
for shape in shapes:
|
|
shape_id = str(shape.get("id"))
|
|
shape_keyframes = keyframes.get(shape_id, [])
|
|
|
|
if not shape_keyframes:
|
|
# No keyframes, use default values
|
|
interpolated.append(shape)
|
|
continue
|
|
|
|
# Find surrounding keyframes
|
|
kf_before = None
|
|
kf_after = None
|
|
|
|
for kf in sorted(shape_keyframes, key=lambda x: x["frame"]):
|
|
if kf["frame"] <= frame_idx:
|
|
kf_before = kf
|
|
elif kf["frame"] > frame_idx and kf_after is None:
|
|
kf_after = kf
|
|
|
|
# Interpolate properties
|
|
if kf_before and kf_after:
|
|
# Linear interpolation
|
|
progress = (frame_idx - kf_before["frame"]) / (kf_after["frame"] - kf_before["frame"])
|
|
interpolated_shape = self.lerp_shape(shape, kf_before, kf_after, progress)
|
|
elif kf_before:
|
|
interpolated_shape = {**shape, **kf_before.get("properties", {})}
|
|
else:
|
|
interpolated_shape = shape
|
|
|
|
interpolated.append(interpolated_shape)
|
|
|
|
return interpolated
|
|
|
|
def lerp_shape(self, base_shape, kf1, kf2, t):
|
|
"""
|
|
Linear interpolation between two keyframes
|
|
"""
|
|
result = base_shape.copy()
|
|
props1 = kf1.get("properties", {})
|
|
props2 = kf2.get("properties", {})
|
|
|
|
# Interpolate numeric properties
|
|
for key in ["x", "y", "width", "height", "radius", "opacity", "rotation"]:
|
|
if key in props1 and key in props2:
|
|
result[key] = props1[key] + (props2[key] - props1[key]) * t
|
|
|
|
return result
|
|
|
|
def draw_shape(self, draw, shape, img_width, img_height):
|
|
"""
|
|
Draw a shape on the mask with rotation support
|
|
"""
|
|
shape_type = shape.get("type", "circle")
|
|
x = shape.get("x", 0.5) * img_width
|
|
y = shape.get("y", 0.5) * img_height
|
|
opacity = int(shape.get("opacity", 1.0) * 255)
|
|
rotation = shape.get("rotation", 0)
|
|
|
|
if shape_type == "circle":
|
|
# Radius is stored normalized (fraction of width); scale to output px.
|
|
radius = resolve_size(shape.get("radius"), img_width, 50)
|
|
# Create a temporary image for rotation
|
|
if rotation != 0:
|
|
size = int(radius * 2.5)
|
|
temp_img = Image.new('L', (size, size), 0)
|
|
temp_draw = ImageDraw.Draw(temp_img)
|
|
temp_draw.ellipse([size//2 - radius, size//2 - radius,
|
|
size//2 + radius, size//2 + radius], fill=opacity)
|
|
temp_img = temp_img.rotate(rotation, expand=False)
|
|
# Paste the rotated image
|
|
paste_x = int(x - size // 2)
|
|
paste_y = int(y - size // 2)
|
|
draw._image.paste(temp_img, (paste_x, paste_y), temp_img)
|
|
else:
|
|
draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=opacity)
|
|
|
|
elif shape_type == "rectangle":
|
|
# Width/height stored normalized (fraction of each dim); scale to px.
|
|
w = resolve_size(shape.get("width"), img_width, 100)
|
|
h = resolve_size(shape.get("height"), img_height, 100)
|
|
|
|
if rotation != 0:
|
|
# Create temporary image for rotation
|
|
size = int(max(w, h) * 1.5)
|
|
temp_img = Image.new('L', (size, size), 0)
|
|
temp_draw = ImageDraw.Draw(temp_img)
|
|
temp_draw.rectangle([size//2 - w//2, size//2 - h//2,
|
|
size//2 + w//2, size//2 + h//2], fill=opacity)
|
|
temp_img = temp_img.rotate(rotation, expand=False)
|
|
paste_x = int(x - size // 2)
|
|
paste_y = int(y - size // 2)
|
|
draw._image.paste(temp_img, (paste_x, paste_y), temp_img)
|
|
else:
|
|
draw.rectangle([x - w/2, y - h/2, x + w/2, y + h/2], fill=opacity)
|
|
|
|
elif shape_type == "bezier" or shape_type == "polygon":
|
|
points = shape.get("points", [])
|
|
if points:
|
|
# Handle both dict and tuple formats
|
|
if isinstance(points[0], dict):
|
|
scaled_points = [(p['x'] * img_width, p['y'] * img_height) for p in points]
|
|
else:
|
|
scaled_points = [(p[0] * img_width, p[1] * img_height) for p in points]
|
|
|
|
if len(scaled_points) >= 3:
|
|
if rotation != 0:
|
|
# Rotate points around center
|
|
cx = x
|
|
cy = y
|
|
rot_rad = np.radians(rotation)
|
|
cos_r = np.cos(rot_rad)
|
|
sin_r = np.sin(rot_rad)
|
|
|
|
rotated_points = []
|
|
for px, py in scaled_points:
|
|
# Translate to origin
|
|
px -= cx
|
|
py -= cy
|
|
# Rotate
|
|
new_x = px * cos_r - py * sin_r
|
|
new_y = px * sin_r + py * cos_r
|
|
# Translate back
|
|
rotated_points.append((new_x + cx, new_y + cy))
|
|
|
|
draw.polygon(rotated_points, fill=opacity)
|
|
else:
|
|
draw.polygon(scaled_points, fill=opacity)
|
|
|
|
def apply_feather(self, mask, feather_radius):
|
|
"""
|
|
Apply Gaussian blur for feathering effect
|
|
"""
|
|
try:
|
|
from scipy.ndimage import gaussian_filter
|
|
return gaussian_filter(mask, sigma=feather_radius/3)
|
|
except:
|
|
return mask
|
|
|
|
|
|
# Web server routes for the integrated UI
|
|
@PromptServer.instance.routes.get("/animated_mask_editor")
|
|
async def get_editor(request):
|
|
"""Serve the mask editor UI"""
|
|
html_path = os.path.join(os.path.dirname(__file__), "animated_mask_editor.html")
|
|
|
|
if os.path.exists(html_path):
|
|
with open(html_path, 'r', encoding='utf-8') as f:
|
|
html_content = f.read()
|
|
else:
|
|
# Fallback: serve embedded HTML
|
|
html_content = get_embedded_html()
|
|
|
|
return web.Response(text=html_content, content_type='text/html')
|
|
|
|
def _request_node_key(request):
|
|
"""Resolve the per-node storage key from the ?node= query param."""
|
|
return node_key_for(request.query.get('node'))
|
|
|
|
|
|
def _get_node_video(node_key):
|
|
"""Return the stored video tensor for a node, falling back to the last seen."""
|
|
video = AnimatedMaskDrawer.current_videos.get(node_key)
|
|
if video is None:
|
|
video = AnimatedMaskDrawer.last_video
|
|
return video
|
|
|
|
|
|
@PromptServer.instance.routes.post("/animated_mask_editor/save")
|
|
async def save_mask_data(request):
|
|
"""Save mask data from the editor to this node's persistent backup."""
|
|
node_key = _request_node_key(request)
|
|
data = await request.json()
|
|
|
|
save_backup(node_key, data)
|
|
|
|
return web.json_response({"success": True, "node": node_key, "path": get_backup_path(node_key)})
|
|
|
|
@PromptServer.instance.routes.post("/animated_mask_editor/stash")
|
|
async def stash_mask_data(request):
|
|
"""Receive the node's current widget value (from the JS extension) so the
|
|
editor can load exactly what's stored in the workflow, even after a restart."""
|
|
body = await request.json()
|
|
node_key = node_key_for(body.get('node'))
|
|
parsed = parse_mask_data(body.get('data'))
|
|
if parsed is not None:
|
|
roto_store[node_key] = parsed
|
|
return web.json_response({"success": True, "node": node_key})
|
|
|
|
@PromptServer.instance.routes.get("/animated_mask_editor/load")
|
|
async def load_mask_data(request):
|
|
"""Load existing mask data for a node (in-memory mirror, else disk backup)."""
|
|
node_key = _request_node_key(request)
|
|
data = roto_store.get(node_key)
|
|
if data is None:
|
|
data = load_backup(node_key)
|
|
if data is None:
|
|
data = {"shapes": [], "keyframes": {}}
|
|
return web.json_response(data)
|
|
|
|
@PromptServer.instance.routes.get("/animated_mask_editor/video_info")
|
|
async def get_video_info(request):
|
|
"""Get video information (frame count, dimensions) for a node."""
|
|
try:
|
|
node_key = _request_node_key(request)
|
|
video = _get_node_video(node_key)
|
|
|
|
if video is not None:
|
|
info = {
|
|
"total_frames": int(video.shape[0]),
|
|
"width": int(video.shape[2]),
|
|
"height": int(video.shape[1]),
|
|
"has_video": True
|
|
}
|
|
print(f"[AnimatedMaskDrawer] video_info for node {node_key}: {info}")
|
|
return web.json_response(info)
|
|
else:
|
|
print(f"[AnimatedMaskDrawer] No video for node {node_key} - queue the workflow first")
|
|
return web.json_response({"has_video": False, "error": "No video loaded. Please queue the workflow in ComfyUI first."})
|
|
except Exception as e:
|
|
print(f"[AnimatedMaskDrawer] Error getting video info: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.json_response({"has_video": False, "error": str(e)})
|
|
|
|
@PromptServer.instance.routes.get("/animated_mask_editor/video_frame")
|
|
async def get_video_frame(request):
|
|
"""Get specific video frame as base64 JPEG for a node."""
|
|
frame_idx = int(request.query.get('frame', 0))
|
|
node_key = _request_node_key(request)
|
|
|
|
video = _get_node_video(node_key)
|
|
if video is None:
|
|
return web.json_response({"error": "No video loaded"})
|
|
|
|
# Clamp frame index
|
|
frame_idx = max(0, min(frame_idx, video.shape[0] - 1))
|
|
|
|
# Get frame (tensor to numpy)
|
|
frame = video[frame_idx].cpu().numpy()
|
|
|
|
# Convert to PIL Image (0-1 float to 0-255 uint8)
|
|
frame_uint8 = (frame * 255).astype(np.uint8)
|
|
img = Image.fromarray(frame_uint8)
|
|
|
|
# Convert to JPEG base64
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='JPEG', quality=85)
|
|
img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
return web.json_response({
|
|
"frame": frame_idx,
|
|
"image": f"data:image/jpeg;base64,{img_base64}"
|
|
})
|
|
|
|
|
|
def get_embedded_html():
|
|
"""Returns embedded HTML if external file not found"""
|
|
return """
|
|
<!DOCTYPE html>
|
|
<html>
|
|
<head>
|
|
<title>Animated Mask Editor</title>
|
|
<style>
|
|
body {
|
|
margin: 0;
|
|
padding: 20px;
|
|
font-family: Arial, sans-serif;
|
|
background: #1a1a1a;
|
|
color: white;
|
|
}
|
|
h1 { text-align: center; }
|
|
.info {
|
|
max-width: 600px;
|
|
margin: 0 auto;
|
|
padding: 20px;
|
|
background: #2a2a2a;
|
|
border-radius: 8px;
|
|
}
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<h1>Animated Mask Editor</h1>
|
|
<div class="info">
|
|
<p>Please place the animated_mask_editor.html file in the same directory as this node.</p>
|
|
<p>The editor UI will load here once the file is available.</p>
|
|
</div>
|
|
</body>
|
|
</html>
|
|
"""
|
|
|
|
|
|
# Node registration for ComfyUI
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AnimatedMaskDrawer": AnimatedMaskDrawer
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AnimatedMaskDrawer": "Animated Mask Drawer 🎨"
|
|
}
|
|
|
|
# Add web directory for serving static files
|
|
WEB_DIRECTORY = "./web"
|