From 0940d2f4a3deb502cd82cf5ae49a7481b09a4232 Mon Sep 17 00:00:00 2001
From: Hawk Lee <l.y@live.cn>
Date: Tue, 20 Jan 2026 20:28:00 +0800
Subject: [PATCH] feat: Expose more Ditto parameters (eye, smoothing, head
 rotation)

---
 aiia_ditto_nodes.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/aiia_ditto_nodes.py b/aiia_ditto_nodes.py
index d6a1dd8..b84d32a 100644
--- a/aiia_ditto_nodes.py
+++ b/aiia_ditto_nodes.py
@@ -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
         )