348 lines
15 KiB
Python
348 lines
15 KiB
Python
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,
|
|
]
|