From 76ea2aa6a664edbd553799f78d6bafdd23e34dc6 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 7 Oct 2025 11:41:54 +0300
Subject: [PATCH] Update wan_video_vae.py

---
 wanvideo/wan_video_vae.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/wanvideo/wan_video_vae.py b/wanvideo/wan_video_vae.py
index e9e8771..713014b 100644
--- a/wanvideo/wan_video_vae.py
+++ b/wanvideo/wan_video_vae.py
@@ -969,10 +969,13 @@ class VideoVAE_(nn.Module):
 
 
     #modification originally by @raindrop313 https://github.com/raindrop313/ComfyUI-WanVideoStartEndFrames
-    def encode_2(self, x):
+    def encode_2(self, x, pbar=True):
         t = x.shape[2]
         iter_ = 2 + (t - 2) // 4
 
+        if pbar:
+            pbar = ProgressBar(iter_)
+
         for i in range(iter_):
             self._enc_conv_idx = [0]
             if i == 0:
@@ -994,6 +997,8 @@ class VideoVAE_(nn.Module):
         out_tail = out[:, :, -1, :, :].unsqueeze(2)
         mu = torch.cat([self.conv1(out_head), self.conv1(out_tail)], dim=2).chunk(2, dim=1)[0]
         mu = (mu - self.mean.to(mu)) * self.inv_std.to(mu)
+        if pbar:
+            pbar.update_absolute(0)
 
         return mu