tqdm progress bars

liveportrait/live_portrait_pipeline.py

@@ -6,7 +6,7 @@ Pipeline of LivePortrait
 
 import cv2
 import comfy.utils
-
+from tqdm import tqdm
 import os.path as osp
 import numpy as np
 from .config.inference_config import InferenceConfig
@@ -77,7 +77,7 @@ class LivePortraitPipeline(object):
         if inference_cfg.flag_eye_retargeting or inference_cfg.flag_lip_retargeting:
             driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_images_np)
 
-        for i in range(total_frames):
+        for i in tqdm(range(total_frames), desc='Animating...', total=total_frames):
             source_frame_rgb = self._get_source_frame(
                 source_np, i, total_frames, mismatch_method
             )

liveportrait/template_maker.py

@@ -8,7 +8,7 @@ import os
 import cv2
 import numpy as np
 import pickle
-from rich.progress import track
+from tqdm import tqdm
 from .utils.cropper import Cropper
 
 from .utils.io import load_driving_info
@@ -41,7 +41,7 @@ class TemplateMaker:
         templates = []
 
 
-        for i in track(range(n_frames), description='Making templates...', total=n_frames):
+        for i in tqdm(range(n_frames), desc='Making templates...', total=n_frames):
             I_d_i = I_d_lst[i]
             x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
             R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])

liveportrait/utils/video.py

@@ -10,7 +10,7 @@ import subprocess
 import imageio
 import cv2
 
-from rich.progress import track
+from tqdm import tqdm
 from .helper import prefix
 from .rprint import rprint as print
 
@@ -35,7 +35,7 @@ def images2video(images, wfp, **kwargs):
     )
 
     n = len(images)
-    for i in track(range(n), description='writing', transient=True):
+    for i in tqdm(range(n), description='writing', transient=True):
         if image_mode.lower() == 'bgr':
             writer.append_data(images[i][..., ::-1])
         else:
@@ -83,7 +83,7 @@ def blend(img: np.ndarray, mask: np.ndarray, background_color=(255, 255, 255)):
 def concat_frames(I_p_lst, driving_rgb_lst, img_rgb):
     # TODO: add more concat style, e.g., left-down corner driving
     out_lst = []
-    for idx, _ in track(enumerate(I_p_lst), total=len(I_p_lst), description='Concatenating result...'):
+    for idx, _ in tqdm(enumerate(I_p_lst), total=len(I_p_lst), desc='Concatenating result...'):
         source_image_drived = I_p_lst[idx]
         image_drive = driving_rgb_lst[idx]