diff --git a/README.md b/README.md index edddf68..9fb4af7 100755 --- a/README.md +++ b/README.md @@ -1175,6 +1175,11 @@ B: 太神奇了！那我们快去生成试试吧！ ## Changelog +### [1.10.13] - 2026-02-02 + +- **Ditto Sampler**: 修复了采样过程中无法正常响应 ComfyUI 中断/取消信号的问题。 + - 为所有工作线程增加了超时检测和状态轮询，支持在长任务执行期间即时退出。 + ### [1.10.12] - 2026-02-02 - **Debug & Utilities**: 新增 **Text Debug Splicer** 节点。 diff --git a/aiia_ditto_nodes.py b/aiia_ditto_nodes.py index 256a136..fa68a37 100644 --- a/aiia_ditto_nodes.py +++ b/aiia_ditto_nodes.py @@ -5,6 +5,7 @@ import numpy as np from PIL import Image import folder_paths import logging +import comfy.model_management # Added for interruption check # Add Ditto library to path @@ -972,8 +973,18 @@ class AIIA_DittoSampler: aud_feat = master_sdk.wav2feat.wav2feat(audio_np) master_sdk.audio2motion_queue.put(aud_feat) - # 5. Wait for completion - # master_sdk.close() joins threads. + # 5. Wait for completion with Interruption Check + # Replace blocking master_sdk.close() with polling + import time + while any(t.is_alive() for t in master_sdk.thread_list): + if comfy.model_management.should_stop(): + logger.info("[Ditto] Interruption detected. Stopping workers...") + master_sdk.stop_event.set() + # Flush queues to unblock workers waiting on put + # (Though we already added timeout in SDK, this is faster) + break + time.sleep(0.1) + master_sdk.close() except Exception as e: logging.error(f"Error during Ditto inference: {e}") diff --git a/libs/Ditto/stream_pipeline_offline.py b/libs/Ditto/stream_pipeline_offline.py index 3c724ee..f25b999 100644 --- a/libs/Ditto/stream_pipeline_offline.py +++ b/libs/Ditto/stream_pipeline_offline.py @@ -295,13 +295,24 @@ class StreamSDK: except queue.Empty: continue if item is None: - self.writer_queue.put(None) + while not self.stop_event.is_set(): + try: + self.writer_queue.put(None, timeout=1) + break + except queue.Full: + continue break frame_idx, render_img = item frame_rgb = self.source_info["img_rgb_lst"][frame_idx] M_c2o = self.source_info["M_c2o_lst"][frame_idx] res_frame_rgb = self.putback(frame_rgb, render_img, M_c2o) - self.writer_queue.put(res_frame_rgb) + + while not self.stop_event.is_set(): + try: + self.writer_queue.put(res_frame_rgb, timeout=1) + break + except queue.Full: + continue def decode_f3d_worker(self): try: @@ -317,11 +328,22 @@ class StreamSDK: except queue.Empty: continue if item is None: - self.putback_queue.put(None) + while not self.stop_event.is_set(): + try: + self.putback_queue.put(None, timeout=1) + break + except queue.Full: + continue break frame_idx, f_3d = item render_img = self.decode_f3d(f_3d) - self.putback_queue.put([frame_idx, render_img]) + + while not self.stop_event.is_set(): + try: + self.putback_queue.put([frame_idx, render_img], timeout=1) + break + except queue.Full: + continue def warp_f3d_worker(self): try: @@ -337,12 +359,23 @@ class StreamSDK: except queue.Empty: continue if item is None: - self.decode_f3d_queue.put(None) + while not self.stop_event.is_set(): + try: + self.decode_f3d_queue.put(None, timeout=1) + break + except queue.Full: + continue break frame_idx, x_s, x_d = item f_s = self.source_info["f_s_lst"][frame_idx] f_3d = self.warp_f3d(f_s, x_s, x_d) - self.decode_f3d_queue.put([frame_idx, f_3d]) + + while not self.stop_event.is_set(): + try: + self.decode_f3d_queue.put([frame_idx, f_3d], timeout=1) + break + except queue.Full: + continue def motion_stitch_worker(self): try: @@ -358,13 +391,24 @@ class StreamSDK: except queue.Empty: continue if item is None: - self.warp_f3d_queue.put(None) + while not self.stop_event.is_set(): + try: + self.warp_f3d_queue.put(None, timeout=1) + break + except queue.Full: + continue break frame_idx, x_d_info, ctrl_kwargs = item x_s_info = self.source_info["x_s_info_lst"][frame_idx] x_s, x_d = self.motion_stitch(x_s_info, x_d_info, **ctrl_kwargs) - self.warp_f3d_queue.put([frame_idx, x_s, x_d]) + + while not self.stop_event.is_set(): + try: + self.warp_f3d_queue.put([frame_idx, x_s, x_d], timeout=1) + break + except queue.Full: + continue def audio2motion_worker(self): try: @@ -401,6 +445,8 @@ class StreamSDK: min_silence_for_reset = 30 # ~1.2s @ 25fps. Avoid resetting on standard sentence pauses. while idx < num_frames: + if self.stop_event.is_set(): + break pbar.update() vad_timeline = getattr(self, "vad_timeline", None) @@ -515,7 +561,12 @@ class StreamSDK: break - self.motion_stitch_queue.put(None) + while not self.stop_event.is_set(): + try: + self.motion_stitch_queue.put(None, timeout=1) + break + except queue.Full: + continue def _audio2motion_worker(self): @@ -615,14 +666,30 @@ class StreamSDK: if is_end: break - self.motion_stitch_queue.put(None) + while not self.stop_event.is_set(): + try: + self.motion_stitch_queue.put(None, timeout=1) + break + except queue.Full: + continue def close(self): - # flush frames - self.audio2motion_queue.put(None) + # flush frames - using non-blocking put as it might be called during interruption + while not self.stop_event.is_set(): + try: + self.audio2motion_queue.put(None, timeout=1) + break + except queue.Full: + continue + # Wait for worker threads to finish for thread in self.thread_list: - thread.join() + while thread.is_alive(): + thread.join(timeout=0.1) + # Check ComfyUI interruption if needed (caller should handle this usually) + # but we just want to ensure we don't block forever if stop_event is set elsewhere + if self.stop_event.is_set() and not thread.is_alive(): + break try: self.writer.close() diff --git a/pyproject.toml b/pyproject.toml index d9b75d7..8661dc1 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "aiia" description = "The Ultimate AI Audio/Video toolkit for ComfyUI. Features an enhanced Ditto (with optimizations that outperform official demos and other SOTA talking head models in lip-sync accuracy and natural motion), EchoMimic V3 & FLOAT, VibeVoice & CosyVoice 3.0 (Zero-Shot Voice Cloning), Multi-Role Podcast Generation, and a powerful Media Browser." -version = "1.10.12" +version = "1.10.13" license = {file = "LICENSE"} readme = "README.md" authors = [