progress bar fixed

OmniAvatar/wan_video.py

@@ -223,7 +223,7 @@ class WanVideoPipeline(BasePipeline):
         tile_stride=(15, 26),
         tea_cache_l1_thresh=None,
         tea_cache_model_id="",
-        progress_bar_cmd=tqdm,
+        progress_bar_cmd=None,
         return_latent=False,
     ):
         tiler_kwargs = {"tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride}
@@ -249,7 +249,7 @@ class WanVideoPipeline(BasePipeline):
         
         # Denoise
         self.load_models_to_device(["dit"])
-        for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)):
+        for progress_id, timestep in enumerate(tqdm(self.scheduler.timesteps) if progress_bar_cmd is None else self.scheduler.timesteps ):
             if fixed_frame > 0: # new
                 latents[:, :, :fixed_frame] = lat[:, :, :fixed_frame]
             timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device)
@@ -273,6 +273,10 @@ class WanVideoPipeline(BasePipeline):
                 noise_pred = noise_pred_posi
             # Scheduler
             latents = self.scheduler.step(noise_pred, self.scheduler.timesteps[progress_id], latents)
+
+            if progress_bar_cmd is not None:
+                progress_bar_cmd.update(1)
+            
             
         if fixed_frame > 0: # new
             latents[:, :, :fixed_frame] = lat[:, :, :fixed_frame]

app.py

@@ -11,6 +11,7 @@ import librosa
 import numpy as np
 import uuid
 import shutil
+from tqdm import tqdm
 
 import importlib, site, sys
 from huggingface_hub import hf_hub_download, snapshot_download
@@ -443,51 +444,54 @@ class WanInferencePipeline(nn.Module):
             msk[:, :, 1:] = 1
             image_emb["y"] = torch.cat([image_cat, msk], dim=1)
 
-        for t in range(times):
-            print(f"[{t+1}/{times}]")
-            audio_emb = {}
-            if t == 0:
-                overlap = first_fixed_frame
-            else:
-                overlap = fixed_frame
-                image_emb["y"][:, -1:, :prefix_lat_frame] = 0 # 第一次推理是mask只有1，往后都是mask overlap
-            prefix_overlap = (3 + overlap) // 4
-            if audio_embeddings is not None:
+        total_iterations = times * num_steps
+
+        with tqdm(total=total_iterations) as pbar:
+            for t in range(times):
+                print(f"[{t+1}/{times}]")
+                audio_emb = {}
                 if t == 0:
-                    audio_tensor = audio_embeddings[
-                            :min(L - overlap, audio_embeddings.shape[0])
+                    overlap = first_fixed_frame
+                else:
+                    overlap = fixed_frame
+                    image_emb["y"][:, -1:, :prefix_lat_frame] = 0 # 第一次推理是mask只有1，往后都是mask overlap
+                prefix_overlap = (3 + overlap) // 4
+                if audio_embeddings is not None:
+                    if t == 0:
+                        audio_tensor = audio_embeddings[
+                                :min(L - overlap, audio_embeddings.shape[0])
+                            ]
+                    else:
+                        audio_start = L - first_fixed_frame + (t - 1) * (L - overlap)
+                        audio_tensor = audio_embeddings[
+                            audio_start: min(audio_start + L - overlap, audio_embeddings.shape[0])
                         ]
+                        
+                    audio_tensor = torch.cat([audio_prefix, audio_tensor], dim=0)
+                    audio_prefix = audio_tensor[-fixed_frame:]
+                    audio_tensor = audio_tensor.unsqueeze(0).to(device=self.device, dtype=self.dtype)
+                    audio_emb["audio_emb"] = audio_tensor
                 else:
-                    audio_start = L - first_fixed_frame + (t - 1) * (L - overlap)
-                    audio_tensor = audio_embeddings[
-                        audio_start: min(audio_start + L - overlap, audio_embeddings.shape[0])
-                    ]
-                    
-                audio_tensor = torch.cat([audio_prefix, audio_tensor], dim=0)
-                audio_prefix = audio_tensor[-fixed_frame:]
-                audio_tensor = audio_tensor.unsqueeze(0).to(device=self.device, dtype=self.dtype)
-                audio_emb["audio_emb"] = audio_tensor
-            else:
-                audio_prefix = None
-            if image is not None and img_lat is None:
-                self.pipe.load_models_to_device(['vae'])
-                img_lat = self.pipe.encode_video(image.to(dtype=self.dtype)).to(self.device, dtype=self.dtype)
-                assert img_lat.shape[2] == prefix_overlap
-            img_lat = torch.cat([img_lat, torch.zeros_like(img_lat[:, :, :1].repeat(1, 1, T - prefix_overlap, 1, 1), dtype=self.dtype)], dim=2)
-            frames, _, latents = self.pipe.log_video(img_lat, prompt, prefix_overlap, image_emb, audio_emb,
-                                                 negative_prompt, num_inference_steps=num_steps, 
-                                                 cfg_scale=guidance_scale, audio_cfg_scale=audio_scale if audio_scale is not None else guidance_scale,
-                                                 return_latent=True,
-                                                 tea_cache_l1_thresh=self.args.tea_cache_l1_thresh,tea_cache_model_id="Wan2.1-T2V-14B")
-
-            torch.cuda.empty_cache()
-            img_lat = None
-            image = (frames[:, -fixed_frame:].clip(0, 1) * 2.0 - 1.0).permute(0, 2, 1, 3, 4).contiguous()
-            
-            if t == 0:
-                video.append(frames)
-            else:
-                video.append(frames[:, overlap:])
+                    audio_prefix = None
+                if image is not None and img_lat is None:
+                    self.pipe.load_models_to_device(['vae'])
+                    img_lat = self.pipe.encode_video(image.to(dtype=self.dtype)).to(self.device, dtype=self.dtype)
+                    assert img_lat.shape[2] == prefix_overlap
+                img_lat = torch.cat([img_lat, torch.zeros_like(img_lat[:, :, :1].repeat(1, 1, T - prefix_overlap, 1, 1), dtype=self.dtype)], dim=2)
+                frames, _, latents = self.pipe.log_video(img_lat, prompt, prefix_overlap, image_emb, audio_emb,
+                                                     negative_prompt, num_inference_steps=num_steps, 
+                                                     cfg_scale=guidance_scale, audio_cfg_scale=audio_scale if audio_scale is not None else guidance_scale,
+                                                     return_latent=True,
+                                                     tea_cache_l1_thresh=self.args.tea_cache_l1_thresh,tea_cache_model_id="Wan2.1-T2V-14B", progress_bar_cmd=pbar)
+    
+                torch.cuda.empty_cache()
+                img_lat = None
+                image = (frames[:, -fixed_frame:].clip(0, 1) * 2.0 - 1.0).permute(0, 2, 1, 3, 4).contiguous()
+                
+                if t == 0:
+                    video.append(frames)
+                else:
+                    video.append(frames[:, overlap:])
         video = torch.cat(video, dim=1)
         video = video[:, :ori_audio_len + 1]