import logging from typing import List from typing_extensions import override import comfy.model_management import comfy.utils from comfy_api.latest import ComfyExtension, io from comfy_extras.nodes_sam3d_body import SAM3DBody_Predict from .camera_drive import drive_camera from .shape_fit import ( body_type_profile, build_foreground_masks, exaggerate_profile, fit_profile_to_masks, profile_from_pose_data, regenerate_pose_data, ) MHRPoseData = io.Custom("MHR_POSE_DATA") SAM3DBodyModel = io.Custom("SAM3D_BODY_MODEL") SAM3DBodyShapeProfile = io.Custom("SAM3D_BODY_SHAPE_PROFILE") class SAM3DBodyMod_DetectShape(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="SAM3DBodyMod_DetectShape", display_name="Detect SAM3D Body Shape", description="Estimate a reusable MHR body shape from reference images, with optional silhouette fitting.", category="image/detection/SAM3D BodyMod", search_aliases=["SAM3D character shape", "MHR shape from image", "cartoon body fit"], inputs=[ SAM3DBodyModel.Input("sam3d_body_model"), io.Image.Input("image"), io.Combo.Input( "preset", options=["real_human", "cartoon"], default="real_human", tooltip=( "real_human stays close to SAM3D's learned human body space. " "cartoon permits larger silhouette-driven coefficient changes." ), ), io.Int.Input( "person_index", default=0, min=0, max=15, tooltip="Person index to use when the reference contains multiple detections.", ), io.Float.Input( "fit_strength", default=1.0, min=0.0, max=2.0, step=0.05, tooltip="Strength of the silhouette correction. 0 uses native SAM3D; values above 1 exaggerate the correction.", ), io.Int.Input( "fit_iterations", default=24, min=0, max=200, step=1, tooltip="Silhouette fitting steps. Set to 0 for the native SAM3D estimate only.", ), io.Float.Input( "fov", default=0.0, min=0.0, max=180.0, step=0.1, advanced=True, tooltip="Vertical field of view passed to native SAM3D prediction. 0 uses its default.", ), io.Int.Input( "batch_size", default=16, min=1, max=128, step=1, advanced=True, tooltip="Maximum reference crops or regenerated frames processed together.", ), io.Combo.Input( "mask_polarity", options=["auto", "white_is_subject", "comfy_alpha"], default="auto", advanced=True, tooltip="ComfyUI Load Image alpha masks are inverted; auto normally detects this correctly.", ), io.Float.Input( "auto_mask_threshold", default=0.08, min=0.0, max=1.0, step=0.01, advanced=True, tooltip="Color distance from the image-border background used when no foreground mask is connected.", ), io.Float.Input( "recognition_exaggeration", default=1.0, min=0.0, max=3.0, step=0.05, tooltip="Scales all detected shape and proportion deviations from MHR's mean body. 1 keeps them, above 1 exaggerates them.", ), io.Mask.Input( "foreground_mask", optional=True, tooltip="Optional subject silhouette. Auto polarity accepts normal masks and ComfyUI's inverted alpha mask.", ), io.BoundingBox.Input( "bboxes", optional=True, force_input=True, tooltip="Optional reference-person bounding boxes passed to native SAM3D prediction.", ), ], outputs=[ SAM3DBodyShapeProfile.Output("shape_profile"), MHRPoseData.Output("fitted_reference_pose"), ], ) @classmethod def execute( cls, sam3d_body_model, image, preset, person_index, fit_strength, fit_iterations, recognition_exaggeration=1.0, fov=0.0, batch_size=16, mask_polarity="auto", auto_mask_threshold=0.08, foreground_mask=None, bboxes=None, ) -> io.NodeOutput: native = SAM3DBody_Predict.execute( sam3d_body_model, image, bboxes=bboxes, run_hand_refinement=False, fov=float(fov), batch_size=int(batch_size), ) if not native.result: raise RuntimeError("SAM3D BodyMod: native SAM3D prediction returned no pose data") reference_pose = native.result[0] profile = profile_from_pose_data(reference_pose, int(person_index)) profile["preset"] = str(preset) if int(fit_iterations) > 0 and float(fit_strength) > 0.0: masks = build_foreground_masks( image, reference_pose, int(person_index), foreground_mask=foreground_mask, mask_polarity=str(mask_polarity), auto_mask_threshold=float(auto_mask_threshold), ) if not any(mask.any() for mask in masks): logging.warning("SAM3D BodyMod: no usable foreground silhouette; keeping the native shape estimate") else: fit_progress = comfy.utils.ProgressBar(int(fit_iterations)) profile = fit_profile_to_masks( sam3d_body_model.model, reference_pose, profile, masks, int(person_index), str(preset), float(fit_strength), int(fit_iterations), progress=lambda: fit_progress.update(1), ) profile = exaggerate_profile(profile, float(recognition_exaggeration)) selected_count = sum( int(person_index) < len(frame) for frame in reference_pose.get("frames", []) ) regen_progress = comfy.utils.ProgressBar(max(1, selected_count)) fitted_reference = regenerate_pose_data( sam3d_body_model.model, reference_pose, profile, shape_strength=1.0, proportion_strength=1.0, person_index=int(person_index), anchor="camera", batch_size=int(batch_size), progress=lambda: regen_progress.update(1), ) return io.NodeOutput(profile, fitted_reference) class SAM3DBodyMod_BodyTypePreset(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="SAM3DBodyMod_BodyTypePreset", display_name="SAM3D Body Type Preset", description="Apply an ectomorph, mesomorph, endomorph, or detected-feature amplification to a shape profile.", category="image/detection/SAM3D BodyMod", search_aliases=["SAM3D somatotype", "endomorph", "ectomorph", "mesomorph", "maximize body shape"], inputs=[ SAM3DBodyModel.Input("sam3d_body_model"), SAM3DBodyShapeProfile.Input("shape_profile"), io.Combo.Input( "body_type", options=["detected", "ectomorph", "mesomorph", "endomorph", "maximize_detected"], default="detected", tooltip=( "The three named types project semantic width targets into MHR's identity basis. " "maximize_detected amplifies every recognized shape and proportion coefficient." ), ), io.Float.Input( "strength", default=1.0, min=0.0, max=3.0, step=0.05, tooltip="0 leaves the detected profile unchanged. Higher values increase the selected body-type effect.", ), ], outputs=[SAM3DBodyShapeProfile.Output("shape_profile")], ) @classmethod def execute(cls, sam3d_body_model, shape_profile, body_type, strength) -> io.NodeOutput: if body_type not in ("detected", "maximize_detected") and float(strength) > 0.0: comfy.model_management.load_model_gpu(sam3d_body_model) profile = body_type_profile(sam3d_body_model.model, shape_profile, str(body_type), float(strength)) return io.NodeOutput(profile) class SAM3DBodyMod_ApplyShape(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="SAM3DBodyMod_ApplyShape", display_name="Apply SAM3D Body Shape", description="Apply a detected shape profile to animated MHR pose data and rebuild all dependent geometry.", category="image/detection/SAM3D BodyMod", search_aliases=["SAM3D body retarget", "MHR identity transfer", "change animation body"], inputs=[ SAM3DBodyModel.Input("sam3d_body_model"), MHRPoseData.Input("mhr_pose_data"), SAM3DBodyShapeProfile.Input("shape_profile"), io.Float.Input( "shape_strength", default=1.0, min=0.0, max=2.0, step=0.05, tooltip="Surface identity blend. 0 keeps the driver, 1 uses the reference, above 1 exaggerates it.", ), io.Float.Input( "proportion_strength", default=1.0, min=0.0, max=2.0, step=0.05, tooltip="Skeleton/body proportion blend using MHR scale coefficients.", ), io.Int.Input( "person_index", default=-1, min=-1, max=15, tooltip="-1 applies the profile to every person; otherwise only the selected person index.", ), io.Combo.Input( "anchor", options=["camera", "pelvis", "feet"], default="pelvis", tooltip="Keep the original camera, pelvis position, or average ankle position fixed while proportions change.", ), io.Int.Input( "batch_size", default=16, min=1, max=128, step=1, advanced=True, tooltip="Maximum animation frames regenerated together. Lower this if MHR regeneration runs out of VRAM.", ), ], outputs=[MHRPoseData.Output("mhr_pose_data")], ) @classmethod def execute( cls, sam3d_body_model, mhr_pose_data, shape_profile, shape_strength, proportion_strength, person_index, anchor, batch_size=16, ) -> io.NodeOutput: if float(shape_strength) == 0.0 and float(proportion_strength) == 0.0: return io.NodeOutput(mhr_pose_data) comfy.model_management.load_model_gpu(sam3d_body_model) selected_count = sum( len(frame) if int(person_index) < 0 else int(person_index) < len(frame) for frame in mhr_pose_data.get("frames", []) ) progress = comfy.utils.ProgressBar(max(1, selected_count)) out = regenerate_pose_data( sam3d_body_model.model, mhr_pose_data, shape_profile, float(shape_strength), float(proportion_strength), int(person_index), str(anchor), int(batch_size), progress=lambda: progress.update(1), ) return io.NodeOutput(out) class SAM3DBodyMod_CameraDrive(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="SAM3DBodyMod_CameraDrive", display_name="Drive SAM3D Body Camera", description="Orbit a static or keyframed camera around SAM3D Body pose data for native pose rendering.", category="image/detection/SAM3D BodyMod", search_aliases=["SAM3D camera orbit", "MHR camera animation", "azimuth elevation body pose"], inputs=[ MHRPoseData.Input("mhr_pose_data"), io.Float.Input("azimuth", default=0.0, min=-180.0, max=180.0, step=1.0, tooltip="Static camera orbit angle: positive moves the camera to the subject's right."), io.Float.Input("elevation", default=0.0, min=-90.0, max=90.0, step=1.0, tooltip="Static camera elevation in degrees."), io.Float.Input("distance", default=1.0, min=0.2, max=3.0, step=0.05, tooltip="Dolly multiplier relative to SAM3D's estimated subject distance."), io.Float.Input("fov", default=0.0, min=0.0, max=120.0, step=0.1, tooltip="Vertical field of view. 0 preserves each frame's predicted focal length."), io.Boolean.Input("use_keyframes", default=False, tooltip="Use the right-click camera path from the orbit handle instead of the static controls."), io.Int.Input("frame_count", default=0, min=0, max=100000, step=1, tooltip="UI timing aid. Set this to the pose sequence length before placing keyframes; 0 uses 24-frame spacing."), io.String.Input("keyframes", default="", multiline=False, tooltip="Camera path JSON written by right-clicking the orbit handle. Frame numbers are 1-based."), io.Combo.Input( "interp_motion", options=["linear", "ease_in_out", "ease_in", "ease_out"], default="linear", tooltip="Timing curve between camera keyframes.", ), io.Combo.Input( "interpolation", options=["linear", "smooth"], default="linear", tooltip="linear uses straight path segments; smooth uses a Catmull-Rom path through three or more keyframes.", ), ], outputs=[MHRPoseData.Output("mhr_pose_data")], ) @classmethod def execute( cls, mhr_pose_data, azimuth, elevation, distance, fov, use_keyframes=False, frame_count=0, keyframes="", interp_motion="linear", interpolation="linear", ) -> io.NodeOutput: out = drive_camera( mhr_pose_data, float(azimuth), float(elevation), float(distance), float(fov), bool(use_keyframes), int(frame_count), str(keyframes), str(interp_motion), str(interpolation), ) return io.NodeOutput(out) class SAM3DBodyModExtension(ComfyExtension): @override async def get_node_list(self) -> List[type[io.ComfyNode]]: return [ SAM3DBodyMod_DetectShape, SAM3DBodyMod_BodyTypePreset, SAM3DBodyMod_ApplyShape, SAM3DBodyMod_CameraDrive, ]