SDPose: resize input always

Author: kijai

comfy/ldm/modules/sdpose.py

@@ -90,7 +90,7 @@ def forward(self, x): # Decode heatmaps to keypoints
                 origin_max = np.max(hm[k])
                 dr = np.zeros((H + 2 * border, W + 2 * border), dtype=np.float32)
                 dr[border:-border, border:-border] = hm[k].copy()
-                dr = gaussian_filter(dr, sigma=2.0)
+                dr = gaussian_filter(dr, sigma=2.0, truncate=2.5)
                 hm[k] = dr[border:-border, border:-border].copy()
                 cur_max = np.max(hm[k])
                 if cur_max > 0:

comfy_extras/nodes_rtdetr.py

@@ -32,10 +32,12 @@ def define_schema(cls):
     def execute(cls, model, image, threshold, class_name, max_detections) -> io.NodeOutput:
         B, H, W, C = image.shape
 
-        image_in = comfy.utils.common_upscale(image.movedim(-1, 1), 640, 640, "bilinear", crop="disabled")
-
         comfy.model_management.load_model_gpu(model)
-        results = model.model.diffusion_model(image_in, (W, H))  # list of B dicts
+        results = []
+        for i in range(0, B, 32):
+            batch = image[i:i + 32]
+            image_in = comfy.utils.common_upscale(batch.movedim(-1, 1), 640, 640, "bilinear", crop="disabled")
+            results.extend(model.model.diffusion_model(image_in, (W, H)))
 
         all_bbox_dicts = []
 

comfy_extras/nodes_sdpose.py

@@ -1,5 +1,6 @@
 import torch
 import comfy.utils
+import comfy.model_management
 import numpy as np
 import math
 import colorsys
@@ -410,7 +411,9 @@ def _zeros(n):
             pose_outputs.append(canvas)
 
         pose_outputs_np = np.stack(pose_outputs) if len(pose_outputs) > 1 else np.expand_dims(pose_outputs[0], 0)
-        final_pose_output = torch.from_numpy(pose_outputs_np).float() / 255.0
+        final_pose_output = torch.from_numpy(pose_outputs_np).to(
+            device=comfy.model_management.intermediate_device(),
+            dtype=comfy.model_management.intermediate_dtype()) / 255.0
         return io.NodeOutput(final_pose_output)
 
 class SDPoseKeypointExtractor(io.ComfyNode):
@@ -459,6 +462,27 @@ def output_patch(h, hsp, transformer_options):
         model_h = int(head.heatmap_size[0]) * 4   # e.g. 192 * 4 = 768
         model_w = int(head.heatmap_size[1]) * 4   # e.g. 256 * 4 = 1024
 
+        def _resize_to_model(imgs):
+            """Aspect-preserving resize + zero-pad BHWC images to (model_h, model_w). Returns (resized_bhwc, scale, pad_top, pad_left)."""
+            h, w = imgs.shape[-3], imgs.shape[-2]
+            scale = min(model_h / h, model_w / w)
+            sh, sw = int(round(h * scale)), int(round(w * scale))
+            pt, pl = (model_h - sh) // 2, (model_w - sw) // 2
+            chw = imgs.permute(0, 3, 1, 2).float()
+            scaled = comfy.utils.common_upscale(chw, sw, sh, upscale_method="bilinear", crop="disabled")
+            padded = torch.zeros(scaled.shape[0], scaled.shape[1], model_h, model_w, dtype=scaled.dtype, device=scaled.device)
+            padded[:, :, pt:pt + sh, pl:pl + sw] = scaled
+            return padded.permute(0, 2, 3, 1), scale, pt, pl
+
+        def _remap_keypoints(kp, scale, pad_top, pad_left, offset_x=0, offset_y=0):
+            """Remap keypoints from model space back to original image space."""
+            kp = kp.copy() if isinstance(kp, np.ndarray) else np.array(kp, dtype=np.float32)
+            invalid = kp[..., 0] < 0
+            kp[..., 0] = (kp[..., 0] - pad_left) / scale + offset_x
+            kp[..., 1] = (kp[..., 1] - pad_top)  / scale + offset_y
+            kp[invalid] = -1
+            return kp
+
         def _run_on_latent(latent_batch):
             """Run one forward pass and return (keypoints_list, scores_list) for the batch."""
             nonlocal captured_feat
@@ -504,56 +528,36 @@ def _run_on_latent(latent_batch):
                         if x2 <= x1 or y2 <= y1:
                             continue
 
-                        crop_h_px, crop_w_px = y2 - y1, x2 - x1
                         crop = img[:, y1:y2, x1:x2, :]  # (1, crop_h, crop_w, C)
-
-                        # scale to fit inside (model_h, model_w) while preserving aspect ratio, then pad to exact model size.
-                        scale = min(model_h / crop_h_px, model_w / crop_w_px)
-                        scaled_h, scaled_w = int(round(crop_h_px * scale)), int(round(crop_w_px * scale))
-                        pad_top, pad_left  = (model_h - scaled_h) // 2, (model_w - scaled_w) // 2
-
-                        crop_chw = crop.permute(0, 3, 1, 2).float()  # BHWC → BCHW
-                        scaled = comfy.utils.common_upscale(crop_chw, scaled_w, scaled_h, upscale_method="bilinear", crop="disabled")
-                        padded = torch.zeros(1, scaled.shape[1], model_h, model_w, dtype=scaled.dtype, device=scaled.device)
-                        padded[:, :, pad_top:pad_top + scaled_h, pad_left:pad_left + scaled_w] = scaled
-                        crop_resized = padded.permute(0, 2, 3, 1)  # BCHW → BHWC
+                        crop_resized, scale, pad_top, pad_left = _resize_to_model(crop)
 
                         latent_crop = vae.encode(crop_resized)
                         kp_batch, sc_batch = _run_on_latent(latent_crop)
-                        kp, sc = kp_batch[0], sc_batch[0]  # (K, 2), coords in model pixel space
-
-                        # remove padding offset, undo scale, offset to full-image coordinates.
-                        kp = kp.copy() if isinstance(kp, np.ndarray) else np.array(kp, dtype=np.float32)
-                        kp[..., 0] = (kp[..., 0] - pad_left) / scale + x1
-                        kp[..., 1] = (kp[..., 1] - pad_top)  / scale + y1
-
+                        kp = _remap_keypoints(kp_batch[0], scale, pad_top, pad_left, x1, y1)
                         img_keypoints.append(kp)
-                        img_scores.append(sc)
+                        img_scores.append(sc_batch[0])
                 else:
-                    # No bboxes for this image – run on the full image
-                    latent_img = vae.encode(img)
+                    img_resized, scale, pad_top, pad_left = _resize_to_model(img)
+                    latent_img = vae.encode(img_resized)
                     kp_batch, sc_batch = _run_on_latent(latent_img)
-                    img_keypoints.append(kp_batch[0])
+                    img_keypoints.append(_remap_keypoints(kp_batch[0], scale, pad_top, pad_left))
                     img_scores.append(sc_batch[0])
 
                 all_keypoints.append(img_keypoints)
                 all_scores.append(img_scores)
                 pbar.update(1)
 
         else: # full-image mode, batched
-            tqdm_pbar = tqdm(total=total_images, desc="Extracting keypoints")
-            for batch_start in range(0, total_images, batch_size):
-                batch_end = min(batch_start + batch_size, total_images)
-                latent_batch = vae.encode(image[batch_start:batch_end])
-
+            for batch_start in tqdm(range(0, total_images, batch_size), desc="Extracting keypoints"):
+                batch_resized, scale, pad_top, pad_left = _resize_to_model(image[batch_start:batch_start + batch_size])
+                latent_batch = vae.encode(batch_resized)
                 kp_batch, sc_batch = _run_on_latent(latent_batch)
 
                 for kp, sc in zip(kp_batch, sc_batch):
-                    all_keypoints.append([kp])
+                    all_keypoints.append([_remap_keypoints(kp, scale, pad_top, pad_left)])
                     all_scores.append([sc])
-                    tqdm_pbar.update(1)
 
-                pbar.update(batch_end - batch_start)
+                pbar.update(len(kp_batch))
 
         openpose_frames = _to_openpose_frames(all_keypoints, all_scores, height, width)
         return io.NodeOutput(openpose_frames)