feat: Expose more Ditto parameters (eye, smoothing, head rotation)

This commit is contained in:
Hawk Lee
2026-01-20 20:28:00 +08:00
parent 733593d893
commit 0940d2f4a3
+20 -2
View File
@@ -155,7 +155,7 @@ class ComfyStreamSDK(StreamSDK):
# We just leave it as is. setup() will be called by the Sampler node.
pass
def setup(self, source_image_pil, **kwargs):
def setup(self, source_image_pil, total_frames=0, **kwargs):
# Override setup to accept PIL Image instead of path
# ======== Prepare Options ========
@@ -195,6 +195,7 @@ class ComfyStreamSDK(StreamSDK):
self.N_d = kwargs.get("N_d", -1)
self.use_d_keys = kwargs.get("use_d_keys", None)
self.relative_d = kwargs.get("relative_d", True)
# drive_eye might be passed in kwargs
self.drive_eye = kwargs.get("drive_eye", None) # None: true4image, false4video
self.delta_eye_arr = kwargs.get("delta_eye_arr", None)
self.delta_eye_open_n = kwargs.get("delta_eye_open_n", 0)
@@ -361,6 +362,12 @@ class AIIA_DittoSampler:
"fps": ("INT", {"default": 25, "min": 15, "max": 60}),
"crop_scale": ("FLOAT", {"default": 2.3, "min": 1.0, "max": 5.0, "step": 0.1}),
"emo": (["Neutral", "Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Contempt"], {"default": "Neutral"}),
"drive_eye": ("BOOLEAN", {"default": True}),
"chk_eye_blink": ("BOOLEAN", {"default": True}),
"smo_k_d": ("INT", {"default": 3, "min": 1, "max": 9}),
"hd_rot_p": ("FLOAT", {"default": 0.0, "min": -30.0, "max": 30.0, "step": 1.0}),
"hd_rot_y": ("FLOAT", {"default": 0.0, "min": -30.0, "max": 30.0, "step": 1.0}),
"hd_rot_r": ("FLOAT", {"default": 0.0, "min": -30.0, "max": 30.0, "step": 1.0}),
}
}
@@ -369,7 +376,7 @@ class AIIA_DittoSampler:
FUNCTION = "generate"
CATEGORY = "AIIA/Ditto"
def generate(self, pipe, ref_image, audio, sampling_steps, fps, crop_scale, emo):
def generate(self, pipe, ref_image, audio, sampling_steps, fps, crop_scale, emo, drive_eye, chk_eye_blink, smo_k_d, hd_rot_p, hd_rot_y, hd_rot_r):
# pipe is the dict we returned in Loader
master_sdk = pipe["sdk"]
cfg_pkl = pipe["cfg_pkl"]
@@ -405,6 +412,13 @@ class AIIA_DittoSampler:
}
emo_idx = emo_map.get(emo, 4)
# Prepare Controls
overall_ctrl_info = {
"delta_pitch": hd_rot_p,
"delta_yaw": hd_rot_y,
"delta_roll": hd_rot_r,
}
# Calling setup() creates fresh threads and queues
master_sdk.setup(
source_image_pil=ref_image_pil,
@@ -413,6 +427,10 @@ class AIIA_DittoSampler:
sampling_timesteps=sampling_steps,
crop_scale=crop_scale,
emo=emo_idx,
drive_eye=drive_eye,
delta_eye_open_n=0 if chk_eye_blink else -1, # 0=random, -1=none
smo_k_d=smo_k_d,
overall_ctrl_info=overall_ctrl_info,
total_frames=num_frames # For pbar
)