From 5924d3e6752aba1554c739377256d64eea9800fa Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Tue, 23 Dec 2025 17:59:18 +0200 Subject: [PATCH] Add alternative to taichi rendering Co-Authored-By: weni2012 <4287334+weni2012@users.noreply.github.com> --- NLFPoseExtract/nlf_render.py | 26 ++- nodes.py | 34 +-- render_3d/render_torch.py | 386 +++++++++++++++++++++++++++++++++++ 3 files changed, 425 insertions(+), 21 deletions(-) create mode 100644 render_3d/render_torch.py diff --git a/NLFPoseExtract/nlf_render.py b/NLFPoseExtract/nlf_render.py index f4dc2ed..1e57e26 100644 --- a/NLFPoseExtract/nlf_render.py +++ b/NLFPoseExtract/nlf_render.py @@ -3,7 +3,8 @@ import torch import copy import logging -from ..render_3d.taichi_cylinder import render_whole +from ..render_3d.taichi_cylinder import render_whole as render_whole_taichi +from ..render_3d.render_torch import render_whole as render_whole_torch from ..pose_draw.draw_pose_utils import draw_pose_to_canvas_np def p3d_single_p2d(points, intrinsic_matrix): @@ -187,7 +188,7 @@ def collect_smpl_poses_samurai(data): -def render_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intrinsic_matrix=None, draw_2d=True, draw_face=True, draw_hands=True): +def render_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intrinsic_matrix=None, draw_2d=True, draw_face=True, draw_hands=True, render_backend="taichi"): """ return a list of images """ base_colors_255_dict = { @@ -282,8 +283,14 @@ def render_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intr cylinder_specs = get_single_pose_cylinder_specs((i, smpl_poses[i], None, None, None, None, colors, limb_seq, draw_seq)) cylinder_specs_list.append(cylinder_specs) - - frames_np_rgba = render_whole(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) + if render_backend == "taichi": + try: + frames_np_rgba = render_whole_taichi(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) + except: + logging.warning("Taichi rendering failed. Falling back to torch rendering.") + frames_np_rgba = render_whole_torch(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) + else: + frames_np_rgba = render_whole_torch(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) if dw_poses is not None and draw_2d: canvas_2d = draw_pose_to_canvas_np(aligned_poses, pool=None, H=height, W=width, reshape_scale=0, show_feet_flag=False, show_body_flag=False, show_cheek_flag=True, dw_hand=True, show_face_flag=draw_face, show_hand_flag=draw_hands) @@ -476,7 +483,7 @@ def get_cylinder_specs_list_from_poses(smpl_poses, include_missing=False): return cylinder_specs_list -def render_multi_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intrinsic_matrix=None, draw_2d=True, draw_face=True, draw_hands=True): +def render_multi_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intrinsic_matrix=None, draw_2d=True, draw_face=True, draw_hands=True, render_backend="taichi"): cylinder_specs_list = get_cylinder_specs_list_from_poses(smpl_poses) @@ -486,7 +493,14 @@ def render_multi_nlf_as_images(smpl_poses, dw_poses, height, width, video_length focal_y = intrinsic_matrix[1,1] princpt = (intrinsic_matrix[0,2], intrinsic_matrix[1,2]) # (cx, cy) - frames_np_rgba = render_whole(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) + if render_backend == "taichi": + try: + frames_np_rgba = render_whole_taichi(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) + except: + logging.warning("Taichi rendering failed. Falling back to torch rendering.") + frames_np_rgba = render_whole_torch(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) + else: + frames_np_rgba = render_whole_torch(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) if dw_poses is not None and draw_2d: aligned_poses = copy.deepcopy(dw_poses) canvas_2d = draw_pose_to_canvas_np(aligned_poses, pool=None, H=height, W=width, reshape_scale=0, show_feet_flag=False, show_body_flag=False, show_cheek_flag=True, dw_hand=True, show_face_flag=draw_face, show_hand_flag=draw_hands) diff --git a/nodes.py b/nodes.py index dddbdbb..422e9ea 100644 --- a/nodes.py +++ b/nodes.py @@ -254,6 +254,7 @@ class RenderNLFPoses: "draw_hands": ("BOOLEAN", {"default": True, "tooltip": "Whether to draw hand keypoints"}), "render_device": (["gpu", "cpu", "opengl", "cuda", "vulkan", "metal"], {"default": "gpu", "tooltip": "Taichi device to use for rendering"}), "scale_hands": ("BOOLEAN", {"default": True, "tooltip": "Whether to scale hand keypoints when aligning DW poses"}), + "render_backend": (["taichi", "torch"], {"default": "taichi", "tooltip": "Rendering backend to use"}), } } @@ -262,22 +263,25 @@ class RenderNLFPoses: FUNCTION = "predict" CATEGORY = "WanVideoWrapper" - def predict(self, nlf_poses, width, height, dw_poses=None, ref_dw_pose=None, draw_face=True, draw_hands=True, render_device="gpu", scale_hands=True): + def predict(self, nlf_poses, width, height, dw_poses=None, ref_dw_pose=None, draw_face=True, draw_hands=True, render_device="gpu", scale_hands=True, render_backend="taichi"): from .NLFPoseExtract.nlf_render import render_nlf_as_images, render_multi_nlf_as_images, shift_dwpose_according_to_nlf, process_data_to_COCO_format, intrinsic_matrix_from_field_of_view from .NLFPoseExtract.align3d import solve_new_camera_params_central, solve_new_camera_params_down - import taichi as ti - - device_map = { - "cpu": ti.cpu, - "gpu": ti.gpu, - "opengl": ti.opengl, - "cuda": ti.cuda, - "vulkan": ti.vulkan, - "metal": ti.metal, - } - - ti.init(arch=device_map.get(render_device.lower())) + if render_backend == "taichi": + try: + import taichi as ti + device_map = { + "cpu": ti.cpu, + "gpu": ti.gpu, + "opengl": ti.opengl, + "cuda": ti.cuda, + "vulkan": ti.vulkan, + "metal": ti.metal, + } + ti.init(arch=device_map.get(render_device.lower())) + except: + logging.warning("Taichi selected but not installed. Falling back to torch rendering.") + render_backend = "torch" if isinstance(nlf_poses, dict): pose_input = nlf_poses['joints3d_nonparam'][0] if 'joints3d_nonparam' in nlf_poses else nlf_poses @@ -346,9 +350,9 @@ class RenderNLFPoses: intrinsic_matrix = ori_camera_pose if pose_input[0].shape[0] > 1: - frames_np = render_multi_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands) + frames_np = render_multi_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands, render_backend = render_backend) else: - frames_np = render_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands) + frames_np = render_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands, render_backend = render_backend) frames_tensor = torch.from_numpy(np.stack(frames_np, axis=0)).contiguous() / 255.0 frames_tensor, mask = frames_tensor[..., :3], frames_tensor[..., -1] > 0.5 diff --git a/render_3d/render_torch.py b/render_3d/render_torch.py new file mode 100644 index 0000000..941f668 --- /dev/null +++ b/render_3d/render_torch.py @@ -0,0 +1,386 @@ +import torch +import numpy as np +import random +import math + + +def flatten_specs(specs_list): + """ + Flatten specs_list into numpy arrays + index tables. + Returns: + starts: (N, 3) float32 + ends: (N, 3) float32 + colors: (N, 4) float32 + frame_offset: (num_frames,) int32 + frame_count: (num_frames,) int32 + """ + starts, ends, colors = [], [], [] + frame_offset, frame_count = [], [] + offset = 0 + for specs in specs_list: + frame_offset.append(offset) + frame_count.append(len(specs)) + for s, e, c in specs: + starts.append(s) + ends.append(e) + colors.append(c) + offset += len(specs) + + # Handle empty case + if len(starts) == 0: + return ( + np.zeros((0, 3), dtype=np.float32), + np.zeros((0, 3), dtype=np.float32), + np.zeros((0, 4), dtype=np.float32), + np.array(frame_offset, dtype=np.int32), + np.array(frame_count, dtype=np.int32), + ) + + return ( + np.array(starts, dtype=np.float32), + np.array(ends, dtype=np.float32), + np.array(colors, dtype=np.float32), + np.array(frame_offset, dtype=np.int32), + np.array(frame_count, dtype=np.int32), + ) + + +def render_whole( + specs_list, H=480, W=640, fx=500, fy=500, cx=240, cy=320, radius=21.5, device=None +): + """ + Render cylinders using PyTorch ray marching. + """ + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + starts_np, ends_np, colors_np, frame_offset_np, frame_count_np = flatten_specs( + specs_list + ) + + # Check if there is anything to render + if len(starts_np) == 0: + return [np.zeros((H, W, 4), dtype=np.uint8) for _ in range(len(specs_list))] + + # Move geometry data to device + all_starts = torch.from_numpy(starts_np).to(device).float() + all_ends = torch.from_numpy(ends_np).to(device).float() + all_colors = torch.from_numpy(colors_np).to(device).float() + + # Calculate global z bounds for simple culling/near-far plane setting + z_min_val = min(starts_np[:, 2].min(), ends_np[:, 2].min()) + z_max_val = max(starts_np[:, 2].max(), ends_np[:, 2].max()) + + znear = 0.1 + zfar = max(min(z_max_val, 25000), 10000) + + # Prepare rays for the whole image + # Grid of coordinates + y_coords, x_coords = torch.meshgrid( + torch.arange(H, device=device).float(), + torch.arange(W, device=device).float(), + indexing="ij", + ) + + # Camera intrinsics to ray directions + u = (x_coords - cx) / fx + v = (y_coords - cy) / fy + z = torch.ones_like(u) + + # Ray directions in camera/world space (assuming identity rotation for camera) + ray_dirs = torch.stack([u, v, z], dim=-1) + ray_dirs = ray_dirs / torch.norm(ray_dirs, dim=-1, keepdim=True) # (H, W, 3) + + ray_origins = torch.zeros( + (H, W, 3), device=device + ) # Camera at (0,0,0) [C variable in taichi] + + light_dir = torch.tensor([0.0, 0.0, 1.0], device=device) + + # Rendering parameters + MAX_STEPS = 100 + EPSILON = 1e-3 + + rendered_frames = [] + + # We render frame by frame to avoid OOM with large cylinder counts per frame + # But batching pixels is implicitly done by operating on full (H, W) tensors. + + for i in range(len(specs_list)): + start_idx = frame_offset_np[i] + count = frame_count_np[i] + + if count == 0: + rendered_frames.append(np.zeros((H, W, 4), dtype=np.uint8)) + continue + + # Get cylinders for this frame + curr_starts = all_starts[start_idx : start_idx + count] # (M, 3) + curr_ends = all_ends[start_idx : start_idx + count] # (M, 3) + curr_colors = all_colors[start_idx : start_idx + count] # (M, 4) + + # --- Ray Marching --- + + # Optimization: Precompute cylinder vectors + ba = curr_ends - curr_starts # (M, 3) + ba_len = torch.sqrt((ba * ba).sum(dim=1)) + ba_norm = ba / ba_len.unsqueeze(1) # Normalized axis + + # We need to find closest cylinder for each pixel. + # Since M (num cylinders) is small (~20-100), we can broadcast. + # But (H*W) is large (480*640 = 307200). + # (H, W, 1, 3) - (1, 1, M, 3) -> Memory heavey. + # So we flatten pixels. + + pixels_shape = (H * W,) + flat_ray_dirs = ray_dirs.view(-1, 3) + flat_ray_origins = ray_origins.view(-1, 3) + + flat_t = torch.ones(pixels_shape[0], device=device) * znear + flat_active = torch.ones(pixels_shape[0], dtype=torch.bool, device=device) + flat_hit = torch.zeros(pixels_shape[0], dtype=torch.bool, device=device) + flat_hit_color = torch.zeros((pixels_shape[0], 4), device=device) + flat_hit_pos = torch.zeros( + (pixels_shape[0], 3), device=device + ) # Store hit pos for normal calc + + # To avoid OOM, checking 300k pixels vs 100 cylinders is fine (30MB matrices). + # Let's verify: + # Points P: (N_pix, 3) + # Cyl Start A: (N_cyl, 3) + # P - A: (N_pix, N_cyl, 3). 300k * 100 * 3 * 4bytes ~= 360MB. + # This fits in standard GPU memory easily. + + depth_near = max(z_min_val, 0.1) + depth_far = min(z_max_val + 6000, 20000) + + for step in range(MAX_STEPS): + if not flat_active.any(): + break + + # Current points for active rays + # Only compute for active rays to save time? + # Indexing might be slower than just masking. Let's try masking. + + p = flat_ray_origins + flat_ray_dirs * flat_t.unsqueeze(1) # (N_pix, 3) + + # --- SDF Calculation --- + # Broadcast p against cylinders + # We only need to compute SDF for active pixels, but let's do all for simplicity first, + # or better: filter indices. + + active_indices = torch.nonzero(flat_active).squeeze() + if active_indices.numel() == 0: + break + + p_active = p[active_indices] # (K, 3) + + pa = p_active.unsqueeze(1) - curr_starts.unsqueeze(0) # (K, M, 3) + + # proj + # ba_norm: (M, 3) -> (1, M, 3) + proj = (pa * ba_norm.unsqueeze(0)).sum(dim=-1) # (K, M) + + # clamp + proj_clamped = proj.clamp(min=0.0).min(ba_len.unsqueeze(0)) # (K, M) + + # vec to closest point on axis + closest_on_axis = curr_starts.unsqueeze(0) + proj_clamped.unsqueeze( + -1 + ) * ba_norm.unsqueeze(0) # (K, M, 3) + + # dist + dist_vec = p_active.unsqueeze(1) - closest_on_axis + dist_euc = torch.norm(dist_vec, dim=-1) # (K, M) + sdf = dist_euc - radius # (K, M) + + # Combine all cylinders (Union = min) + min_sdf, min_idx = sdf.min(dim=1) # (K,) + + # Update t + # If min_sdf < EPSILON, we hit + # If flat_t > zfar, we miss + + # Map back to full arrays + current_t_vals = flat_t[active_indices] + + hit_cond = min_sdf < EPSILON + miss_cond = current_t_vals > zfar + + # For hits + new_hits = hit_cond & (~miss_cond) + # Only update hit info for newly hit rays + + # We need to write back results + # Global indices of new hits + hit_global_idx = active_indices[new_hits] + + if hit_global_idx.numel() > 0: + flat_hit[hit_global_idx] = True + flat_active[hit_global_idx] = False + flat_hit_pos[hit_global_idx] = p_active[new_hits] # Store position + + # Get color of closest cylinder + closest_cyl_idx = min_idx[new_hits] + flat_hit_color[hit_global_idx] = curr_colors[closest_cyl_idx] + + # For misses + miss_global_idx = active_indices[miss_cond] + if miss_global_idx.numel() > 0: + flat_active[miss_global_idx] = False + + # Step t + # Only step remaining active + still_active_local = ~(hit_cond | miss_cond) + if still_active_local.any(): + step_dist = min_sdf[still_active_local] + # Avoid stepping too small to prevent stuck + step_dist = torch.max(step_dist, torch.tensor(1e-4, device=device)) + + active_global_idx = active_indices[still_active_local] + flat_t[active_global_idx] += step_dist + + # --- Shading --- + # Compute normals for all hit pixels + hit_indices = torch.nonzero(flat_hit).squeeze() + + if hit_indices.numel() > 0: + p_hit = flat_hit_pos[hit_indices] # (NumHits, 3) + hit_cols = flat_hit_color[hit_indices] # (NumHits, 4) + + # Finite difference normal + e = 1e-3 + + # We need a function to compute scene SDF at arbitrary points quickly + def get_sdf_batch(points): + # points: (N, 3) + # returns: (N,) min sdf + # Re-use curr_starts, curr_ends logic + + # Chunking if too large? + # Assuming it fits since points are subset of image + + pa = points.unsqueeze(1) - curr_starts.unsqueeze(0) # (N, M, 3) + proj = (pa * ba_norm.unsqueeze(0)).sum(dim=-1) + proj_clamped = proj.clamp(min=0.0).min(ba_len.unsqueeze(0)) + + closest = curr_starts.unsqueeze(0) + proj_clamped.unsqueeze( + -1 + ) * ba_norm.unsqueeze(0) + dist = torch.norm(points.unsqueeze(1) - closest, dim=-1) + sdf = dist - radius + return sdf.min(dim=1)[0] + + def get_normal_batch(points): + # Central difference + dx = get_sdf_batch( + points + torch.tensor([e, 0, 0], device=device) + ) - get_sdf_batch(points - torch.tensor([e, 0, 0], device=device)) + dy = get_sdf_batch( + points + torch.tensor([0, e, 0], device=device) + ) - get_sdf_batch(points - torch.tensor([0, e, 0], device=device)) + dz = get_sdf_batch( + points + torch.tensor([0, 0, e], device=device) + ) - get_sdf_batch(points - torch.tensor([0, 0, e], device=device)) + n = torch.stack([dx, dy, dz], dim=-1) + return n / (torch.norm(n, dim=-1, keepdim=True) + 1e-8) + + normals = get_normal_batch(p_hit) + + # Blinn-Phong + # View dir is -ray_dir + view_dir = -flat_ray_dirs[hit_indices] + view_dir = view_dir / torch.norm(view_dir, dim=-1, keepdim=True) + + # Light dir (0,0,1) + # Diffuse + # max(n.dot(-light_dir), 0) -> note taichi code used -light_dir for diffuse? + # Taichi: diff = max(n.dot(-light_dir), 0.0) where light_dir = [0,0,1] + # So light comes from +Z (camera). + + diff = torch.clamp( + (normals * (-light_dir)).sum(dim=-1), min=0.0 + ) # (NumHits,) + + # Specular + half_dir = (view_dir + (-light_dir)).float() + half_dir = half_dir / (torch.norm(half_dir, dim=-1, keepdim=True) + 1e-8) + + spec = torch.clamp((normals * half_dir).sum(dim=-1), min=0.0) + spec = spec**32 + + # Depth factor + z_vals = p_hit[:, 2] + depth_factor = 1.0 - (z_vals - depth_near) / (depth_far - znear) + depth_factor = depth_factor.clamp(0.0, 1.0) + + # Combine + diffuse_term = 0.3 + 0.7 * diff + base_rgb = ( + hit_cols[:, :3] + * diffuse_term.unsqueeze(-1) + * depth_factor.unsqueeze(-1) + ) + + highlight = ( + torch.tensor([1.0, 1.0, 1.0], device=device) + * (0.5 * spec.unsqueeze(-1)) + * depth_factor.unsqueeze(-1) + ) + + final_rgb = base_rgb + highlight + + # Assign back + flat_hit_color[hit_indices, :3] = final_rgb + flat_hit_color[hit_indices, 3] = hit_cols[:, 3] # Alpha + + # Reshape to image + frame_img = flat_hit_color.view(H, W, 4) + + # Convert to numpy uint8 + frame_np = (frame_img.clamp(0, 1) * 255).byte().cpu().numpy() + rendered_frames.append(frame_np) + + return rendered_frames + + +def random_cylinder(): + """Generate a random cylinder (start, end, color).""" + # Start point [-200,200]^2, z in [300,400] + ax = random.uniform(-200, 200) + ay = random.uniform(-200, 200) + az = random.uniform(300, 400) + start = [ax, ay, az] + + # Random direction and length + theta = random.uniform(0, 2 * math.pi) + phi = random.uniform(-math.pi / 4, math.pi / 4) # Tilt angle + L = 100 + dx = math.cos(phi) * math.cos(theta) + dy = math.cos(phi) * math.sin(theta) + dz = math.sin(phi) + end = [ax + dx * L, ay + dy * L, az + dz * L] + + # Random color (RGB + alpha=1) + color = [random.random(), random.random(), random.random(), 1.0] + + return (start, end, color) + + +def generate_specs_list(num_frames=120, min_cyl=10, max_cyl=120): + """Generate specs_list, each frame has several random cylinders.""" + specs_list = [] + for _ in range(num_frames): + n_cyl = random.randint(min_cyl, max_cyl) + specs = [random_cylinder() for _ in range(n_cyl)] + specs_x_shift = [ + ( + [spec[0][0] + 50, spec[0][1], spec[0][2]], + [spec[1][0] + 50, spec[1][1], spec[1][2]], + spec[2], + ) + for spec in specs + ] + specs_list.append(specs) + specs_list.append(specs_x_shift) + return specs_list