From 8e9fb0e11c6cfeb67270f9435ffd6974f6537f37 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 2 Jul 2025 16:47:19 +0300
Subject: [PATCH] Update nodes.py

---
 nodes.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index 2e9c709..09a5c4d 100644
--- a/nodes.py
+++ b/nodes.py
@@ -3837,6 +3837,10 @@ class WanVideoSampler:
                 indices = (torch.arange(4 + 1) - 2) * 1
 
                 # start video generation iteratively
+                estimated_iterations = max(1, min(max_frames_num, len(multitalk_audio_embedding)) // (frame_num - motion_frame) + 1)
+                loop_pbar = tqdm(total=estimated_iterations, desc="Generating video clips")
+                iteration_count = 0
+
                 audio_embedding = [multitalk_audio_embedding]
                 while True:
                     audio_embs = []
@@ -3949,7 +3953,6 @@ class WanVideoSampler:
                     for i in tqdm(range(len(timesteps)-1)):
                         timestep = timesteps[i]
                         latent_model_input = latent.to(device)
-                        print("latent_model_input shape: ", latent_model_input.shape)
 
                         noise_pred, self.cache_state = predict_with_cfg(
                             latent_model_input, 
@@ -4022,7 +4025,10 @@ class WanVideoSampler:
                         gen_video_list.append(videos[:, :, cur_motion_frames_num:])
 
                     # decide whether is done
-                    if arrive_last_frame: break
+                    if arrive_last_frame: 
+                        loop_pbar.update(estimated_iterations - iteration_count)
+                        loop_pbar.close()
+                        break
 
                     # update next condition frames
                     is_first_clip = False
@@ -4032,6 +4038,10 @@ class WanVideoSampler:
                     audio_start_idx += (frame_num - cur_motion_frames_num)
                     audio_end_idx = audio_start_idx + clip_length
 
+                    # Update progress bar
+                    iteration_count += 1
+                    loop_pbar.update(1)
+
                     # Repeat audio emb
                     if audio_end_idx >= min(max_frames_num, len(audio_embedding[0])):
                         arrive_last_frame = True
@@ -4048,8 +4058,10 @@ class WanVideoSampler:
                             else:
                                 miss_lengths.append(0)
                     
-                    if max_frames_num <= frame_num: break
-                    
+                    if max_frames_num <= frame_num: 
+                        loop_pbar.update(estimated_iterations - iteration_count)
+                        loop_pbar.close()
+                        break
                 
                 gen_video_samples = torch.cat(gen_video_list, dim=2)[:, :, :int(max_frames_num)] 
                 gen_video_samples = gen_video_samples.to(torch.float32)