nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

restructure, add visualization node

Browse Source
This commit is contained in:
kijai
2025-04-11 18:40:50 +03:00
parent 2dc25c150e
commit ea71dfda1a
4 changed files with 285 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+4
View File
@@ -1,3 +1,7 @@
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .recammaster.nodes import NODE_CLASS_MAPPINGS as RECAM_MASTER_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as RECAM_MASTER_NODE_DISPLAY_NAME_MAPPINGS
NODE_CLASS_MAPPINGS.update(RECAM_MASTER_NODE_CLASS_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(RECAM_MASTER_NODE_DISPLAY_NAME_MAPPINGS)
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
nodes.py
-118
View File
@@ -2033,122 +2033,6 @@ class WanVideoVACEStartToEndFrame:
return (out_batch.cpu().float(), masks.cpu().float())
#region ReCamMaster
class Camera(object):
def __init__(self, c2w):
c2w_mat = np.array(c2w).reshape(4, 4)
self.c2w_mat = c2w_mat
self.w2c_mat = np.linalg.inv(c2w_mat)
class WanVideoReCamMasterCameraEmbed:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"camera_type": ([
"pan_right",
"pan_left",
"tilt_up",
"tilt_down",
"zoom_in",
"zoom_out",
"translate_up",
"translate_down",
"arc_left",
"arc_right",
], {"default": "pan_right", "tooltip": "Camera type to use"}),
"latents": ("LATENT", {"tooltip": "source video"}),
},
}
RETURN_TYPES = ("WANVIDIMAGE_EMBEDS", )
RETURN_NAMES = ("camera_embeds",)
FUNCTION = "process"
CATEGORY = "WanVideoWrapper"
DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster"
def process(self, camera_type, latents):
# load camera
import json
from einops import rearrange
camera_data_path = os.path.join(script_directory, "camera_extrinsics.json")
with open(camera_data_path, 'r') as file:
cam_data = json.load(file)
samples = latents["samples"].squeeze(0)
C, T, H, W = samples.shape
num_frames = (T - 1) * 4 + 1
camera_type_map = {
"pan_right": 1,
"pan_left": 2,
"tilt_up": 3,
"tilt_down": 4,
"zoom_in": 5,
"zoom_out": 6,
"translate_up": 7,
"translate_down": 8,
"arc_left": 9,
"arc_right": 10,
}
cam_idx = list(range(num_frames))[::4]
print("cam_idx", cam_idx)
traj = [self.parse_matrix(cam_data[f"frame{idx}"][f"cam{int(camera_type_map[camera_type]):02d}"]) for idx in cam_idx]
traj = np.stack(traj).transpose(0, 2, 1)
c2ws = []
for c2w in traj:
c2w = c2w[:, [1, 2, 0, 3]]
c2w[:3, 1] *= -1.
c2w[:3, 3] /= 100
c2ws.append(c2w)
tgt_cam_params = [Camera(cam_param) for cam_param in c2ws]
relative_poses = []
for i in range(len(tgt_cam_params)):
relative_pose = self.get_relative_pose([tgt_cam_params[0], tgt_cam_params[i]])
relative_poses.append(torch.as_tensor(relative_pose)[:,:3,:][1])
pose_embedding = torch.stack(relative_poses, dim=0) # 21x3x4
pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')
seq_len = math.ceil((H * W) / 4 * ((num_frames - 1) // 4 + 1))
embeds = {
"max_seq_len": seq_len,
"target_shape": samples.shape,
"num_frames": num_frames,
"recammaster": {
"camera_embed": pose_embedding,
"source_latents": samples
}
}
return (embeds,)
def parse_matrix(self, matrix_str):
rows = matrix_str.strip().split('] [')
matrix = []
for row in rows:
row = row.replace('[', '').replace(']', '')
matrix.append(list(map(float, row.split())))
return np.array(matrix)
def get_relative_pose(self, cam_params):
abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]
abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]
cam_to_origin = 0
target_cam_c2w = np.array([
[1, 0, 0, 0],
[0, 1, 0, -cam_to_origin],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
abs2rel = target_cam_c2w @ abs_w2cs[0]
ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]
ret_poses = np.array(ret_poses, dtype=np.float32)
return ret_poses
#region context options
class WanVideoContextOptions:
@@ -3357,7 +3241,6 @@ NODE_CLASS_MAPPINGS = {
"WanVideoVACEEncode": WanVideoVACEEncode,
"WanVideoVACEStartToEndFrame": WanVideoVACEStartToEndFrame,
"WanVideoVACEModelSelect": WanVideoVACEModelSelect,
"WanVideoReCamMasterCameraEmbed": WanVideoReCamMasterCameraEmbed,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"WanVideoSampler": "WanVideo Sampler",
@@ -3393,5 +3276,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"WanVideoVACEEncode": "WanVideo VACE Encode",
"WanVideoVACEStartToEndFrame": "WanVideo VACE Start To End Frame",
"WanVideoVACEModelSelect": "WanVideo VACE Model Select",
"WanVideoReCamMasterCameraEmbed": "WanVideo ReCam Master Camera Embed",
}
camera_extrinsics.json → recammaster/camera_extrinsics.json
View File
recammaster/nodes.py
+281
View File
@@ -0,0 +1,281 @@
import numpy as np
import os, io
import torch
from PIL import Image
import math, time
script_directory = os.path.dirname(os.path.abspath(__file__))
class Camera(object):
def __init__(self, c2w):
c2w_mat = np.array(c2w).reshape(4, 4)
self.c2w_mat = c2w_mat
self.w2c_mat = np.linalg.inv(c2w_mat)
class WanVideoReCamMasterCameraEmbed:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"camera_type": ([
"pan_right",
"pan_left",
"tilt_up",
"tilt_down",
"zoom_in",
"zoom_out",
"translate_up",
"translate_down",
"arc_left",
"arc_right",
], {"default": "pan_right", "tooltip": "Camera type to use"}),
"latents": ("LATENT", {"tooltip": "source video"}),
},
}
RETURN_TYPES = ("WANVIDIMAGE_EMBEDS", "CAMERAPOSES",)
RETURN_NAMES = ("camera_embeds", "camera_poses",)
FUNCTION = "process"
CATEGORY = "WanVideoWrapper"
DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster"
def process(self, camera_type, latents):
# load camera
import json
from einops import rearrange
camera_data_path = os.path.join(script_directory, "camera_extrinsics.json")
with open(camera_data_path, 'r') as file:
cam_data = json.load(file)
samples = latents["samples"].squeeze(0)
C, T, H, W = samples.shape
num_frames = (T - 1) * 4 + 1
camera_type_map = {
"pan_right": 1,
"pan_left": 2,
"tilt_up": 3,
"tilt_down": 4,
"zoom_in": 5,
"zoom_out": 6,
"translate_up": 7,
"translate_down": 8,
"arc_left": 9,
"arc_right": 10,
}
cam_idx = list(range(num_frames))[::4]
traj = [self.parse_matrix(cam_data[f"frame{idx}"][f"cam{int(camera_type_map[camera_type]):02d}"]) for idx in cam_idx]
traj = np.stack(traj).transpose(0, 2, 1)
c2ws = []
for c2w in traj:
c2w = c2w[:, [1, 2, 0, 3]]
c2w[:3, 1] *= -1.
c2w[:3, 3] /= 100
c2ws.append(c2w)
tgt_cam_params = [Camera(cam_param) for cam_param in c2ws]
relative_poses = []
for i in range(len(tgt_cam_params)):
relative_pose = self.get_relative_pose([tgt_cam_params[0], tgt_cam_params[i]])
relative_poses.append(torch.as_tensor(relative_pose)[:,:3,:][1])
pose_embedding = torch.stack(relative_poses, dim=0) # 21x3x4
pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')
seq_len = math.ceil((H * W) / 4 * ((num_frames - 1) // 4 + 1))
embeds = {
"max_seq_len": seq_len,
"target_shape": samples.shape,
"num_frames": num_frames,
"recammaster": {
"camera_embed": pose_embedding,
"source_latents": samples
}
}
return (embeds, traj,)
def parse_matrix(self, matrix_str):
rows = matrix_str.strip().split('] [')
matrix = []
for row in rows:
row = row.replace('[', '').replace(']', '')
matrix.append(list(map(float, row.split())))
return np.array(matrix)
def get_relative_pose(self, cam_params):
abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]
abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]
cam_to_origin = 0
target_cam_c2w = np.array([
[1, 0, 0, 0],
[0, 1, 0, -cam_to_origin],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
abs2rel = target_cam_c2w @ abs_w2cs[0]
ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]
ret_poses = np.array(ret_poses, dtype=np.float32)
return ret_poses
def get_c2w(w2cs, transform_matrix, relative_c2w=True):
if relative_c2w:
target_cam_c2w = np.array([
[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
abs2rel = target_cam_c2w @ w2cs[0]
ret_poses = [target_cam_c2w, ] + [abs2rel @ np.linalg.inv(w2c) for w2c in w2cs[1:]]
else:
ret_poses = [np.linalg.inv(w2c) for w2c in w2cs]
ret_poses = [transform_matrix @ x for x in ret_poses]
return np.array(ret_poses, dtype=np.float32)
class ReCamMasterPoseVisualizer:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"camera_poses": ("CAMERAPOSES",),
"base_xval": ("FLOAT", {"default": 0.2,"min": 0, "max": 100, "step": 0.01}),
"zval": ("FLOAT", {"default": 0.3,"min": 0, "max": 100, "step": 0.01}),
"scale": ("FLOAT", {"default": 1.0,"min": 0.01, "max": 10.0, "step": 0.01}),
"arrow_length": ("FLOAT", {"default": 1,"min": 0, "max": 100, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "plot"
CATEGORY = "WanVideoWrapper"
DESCRIPTION = """
Visualizes the camera poses, from Animatediff-Evolved CameraCtrl Pose
or a .txt file with RealEstate camera intrinsics and coordinates, in a 3D plot.
"""
def plot(self, camera_poses, scale, base_xval, zval, arrow_length):
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from torchvision.transforms import ToTensor
x_min = -2.0 * scale
x_max = 2.0 * scale
y_min = -2.0 * scale
y_max = 2.0 * scale
z_min = -2.0 * scale
z_max = 2.0 * scale
plt.rcParams['text.color'] = '#999999'
self.fig = plt.figure(figsize=(18, 7))
self.fig.patch.set_facecolor('#353535')
self.ax = self.fig.add_subplot(projection='3d')
self.ax.set_facecolor('#353535') # Set the background color here
self.ax.grid(color='#999999', linestyle='-', linewidth=0.5)
self.plotly_data = None # plotly data traces
self.ax.set_aspect("auto")
self.ax.set_xlim(x_min, x_max)
self.ax.set_ylim(y_min, y_max)
self.ax.set_zlim(z_min, z_max)
self.ax.set_xlabel('x', color='#999999')
self.ax.set_ylabel('y', color='#999999')
self.ax.set_zlabel('z', color='#999999')
for text in self.ax.get_xticklabels() + self.ax.get_yticklabels() + self.ax.get_zticklabels():
text.set_color('#999999')
print('initialize camera pose visualizer')
total_frames = len(camera_poses)
w2cs = []
for cam in camera_poses:
if cam.shape[0] == 3:
cam = np.vstack((cam, np.array([[0, 0, 0, 1]])))
cam = cam[:, [1, 2, 0, 3]]
cam[:3, 1] *= -1.
w2cs.append(np.linalg.inv(cam))
transform_matrix = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]])
c2ws = get_c2w(w2cs, transform_matrix, True)
scale = max(max(abs(c2w[:3, 3])) for c2w in c2ws)
if scale > 1e-3: # otherwise, pan or tilt
for c2w in c2ws:
c2w[:3, 3] /= scale
for frame_idx, c2w in enumerate(c2ws):
self.extrinsic2pyramid(c2w, frame_idx / total_frames, hw_ratio=1, base_xval=base_xval, zval=(zval))
if arrow_length > 0:
pos = c2w[:3, 3]
forward = c2w[:3, 2]
arrow_start = pos + forward * base_xval
arrow_length = arrow_length
self.ax.quiver(arrow_start[0], arrow_start[1], arrow_start[2],
forward[0], forward[1], forward[2],
color='black', length=arrow_length, arrow_length_ratio=0.1)
# Create the colorbar
cmap = mpl.cm.rainbow
norm = mpl.colors.Normalize(vmin=0, vmax=total_frames)
colorbar = self.fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), ax=self.ax, orientation='vertical')
# Change the colorbar label
colorbar.set_label('Frame', color='#999999') # Change the label and its color
# Change the tick colors
colorbar.ax.yaxis.set_tick_params(colors='#999999') # Change the tick color
# Change the tick frequency
# Assuming you want to set the ticks at every 10th frame
ticks = np.arange(0, total_frames, 10)
colorbar.ax.yaxis.set_ticks(ticks)
plt.title('')
plt.draw()
buf = io.BytesIO()
plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
buf.seek(0)
img = Image.open(buf)
tensor_img = ToTensor()(img)
buf.close()
tensor_img = tensor_img.permute(1, 2, 0).unsqueeze(0)
return (tensor_img,)
def extrinsic2pyramid(self, extrinsic, color_map='red', hw_ratio=9/16, base_xval=1, zval=3):
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.pyplot as plt
vertex_std = np.array([[0, 0, 0, 1],
[base_xval, -base_xval * hw_ratio, zval, 1],
[base_xval, base_xval * hw_ratio, zval, 1],
[-base_xval, base_xval * hw_ratio, zval, 1],
[-base_xval, -base_xval * hw_ratio, zval, 1]])
vertex_transformed = vertex_std @ extrinsic.T
meshes = [[vertex_transformed[0, :-1], vertex_transformed[1][:-1], vertex_transformed[2, :-1]],
[vertex_transformed[0, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1]],
[vertex_transformed[0, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]],
[vertex_transformed[0, :-1], vertex_transformed[4, :-1], vertex_transformed[1, :-1]],
[vertex_transformed[1, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]]]
color = color_map if isinstance(color_map, str) else plt.cm.rainbow(color_map)
self.ax.add_collection3d(
Poly3DCollection(meshes, facecolors=color, linewidths=0.3, edgecolors=color, alpha=0.35))
def customize_legend(self, list_label):
from matplotlib.patches import Patch
import matplotlib.pyplot as plt
list_handle = []
for idx, label in enumerate(list_label):
color = plt.cm.rainbow(idx / len(list_label))
patch = Patch(color=color, label=label)
list_handle.append(patch)
plt.legend(loc='right', bbox_to_anchor=(1.8, 0.5), handles=list_handle)
NODE_CLASS_MAPPINGS = {
"WanVideoReCamMasterCameraEmbed": WanVideoReCamMasterCameraEmbed,
"ReCamMasterPoseVisualizer": ReCamMasterPoseVisualizer,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"WanVideoReCamMasterCameraEmbed": "WanVideo ReCamMaster Camera Embed",
"ReCamMasterPoseVisualizer": "ReCamMaster Pose Visualizer",
}