Merge pull request #859 from krahets/main

Fix batch decoding for Wan-Video-VAE

Author: Artiprocher

diffsynth/models/wan_video_vae.py

@@ -1216,7 +1216,6 @@ def single_decode(self, hidden_state, device):
 
 
     def encode(self, videos, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)):
-
         videos = [video.to("cpu") for video in videos]
         hidden_states = []
         for video in videos:
@@ -1234,11 +1233,18 @@ def encode(self, videos, device, tiled=False, tile_size=(34, 34), tile_stride=(1
 
 
     def decode(self, hidden_states, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)):
-        if tiled:
-            video = self.tiled_decode(hidden_states, device, tile_size, tile_stride)
-        else:
-            video = self.single_decode(hidden_states, device)
-        return video
+        hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states]
+        videos = []
+        for hidden_state in hidden_states:
+            hidden_state = hidden_state.unsqueeze(0)
+            if tiled:
+                video = self.tiled_decode(hidden_state, device, tile_size, tile_stride)
+            else:
+                video = self.single_decode(hidden_state, device)
+            video = video.squeeze(0)
+            videos.append(video)
+        videos = torch.stack(videos)
+        return videos
 
 
     @staticmethod