Update opensora/models/ae/videobase/causal_vae/modeling_causalvae.py

opensora/models/ae/videobase/causal_vae/modeling_causalvae.py

@@ -316,7 +316,8 @@ class CausalVAEModel(VideoBaseAE_PL):
         self.tile_sample_min_size = 256
         self.tile_sample_min_size_t = 65
         self.tile_latent_min_size = int(self.tile_sample_min_size / (2 ** (len(hidden_size_mult) - 1)))
-        # self.tile_latent_min_size_t = int((self.tile_sample_min_size_t-1) / (2 ** self.time_compress)) + 1
+        t_down_ratio = [i for i in encoder_temporal_downsample if len(i) > 0]
+        self.tile_latent_min_size_t = int((self.tile_sample_min_size_t-1) / (2 ** len(t_down_ratio))) + 1
         self.tile_overlap_factor = 0.25
         self.use_tiling = False
 
@@ -374,8 +375,9 @@ class CausalVAEModel(VideoBaseAE_PL):
         if self.use_tiling and (
             x.shape[-1] > self.tile_sample_min_size
             or x.shape[-2] > self.tile_sample_min_size
+            or x.shape[-3] > self.tile_sample_min_size_t
         ):
-            return self.tiled_encode2d(x)
+            return self.tiled_encode(x)
         h = self.encoder(x)
         moments = self.quant_conv(h)
         posterior = DiagonalGaussianDistribution(moments)
@@ -385,8 +387,9 @@ class CausalVAEModel(VideoBaseAE_PL):
         if self.use_tiling and (
             z.shape[-1] > self.tile_latent_min_size
             or z.shape[-2] > self.tile_latent_min_size
+            or z.shape[-3] > self.tile_latent_min_size_t
         ):
-            return self.tiled_decode2d(z)
+            return self.tiled_decode(z)
         z = self.post_quant_conv(z)
         dec = self.decoder(z)
         return dec
@@ -554,7 +557,54 @@ class CausalVAEModel(VideoBaseAE_PL):
             ) + b[:, :, :, :, x] * (x / blend_extent)
         return b
 
-    def tiled_encode2d(self, x):
+    def tiled_encode(self, x):
+        t = x.shape[2]
+        t_chunk_idx = [i for i in range(0, t, self.tile_sample_min_size_t-1)]
+        if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
+            t_chunk_start_end = [[0, t]]
+        else:
+            t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i+1]+1] for i in range(len(t_chunk_idx)-1)]
+            if t_chunk_start_end[-1][-1] > t:
+                t_chunk_start_end[-1][-1] = t
+            elif t_chunk_start_end[-1][-1] < t:
+                last_start_end = [t_chunk_idx[-1], t]
+                t_chunk_start_end.append(last_start_end)
+        moments = []
+        for idx, (start, end) in enumerate(t_chunk_start_end):
+            chunk_x = x[:, :, start: end]
+            if idx != 0:
+                moment = self.tiled_encode2d(chunk_x, return_moments=True)[:, :, 1:]
+            else:
+                moment = self.tiled_encode2d(chunk_x, return_moments=True)
+            moments.append(moment)
+        moments = torch.cat(moments, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior
+    
+    def tiled_decode(self, x):
+        t = x.shape[2]
+        t_chunk_idx = [i for i in range(0, t, self.tile_latent_min_size_t-1)]
+        if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
+            t_chunk_start_end = [[0, t]]
+        else:
+            t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i+1]+1] for i in range(len(t_chunk_idx)-1)]
+            if t_chunk_start_end[-1][-1] > t:
+                t_chunk_start_end[-1][-1] = t
+            elif t_chunk_start_end[-1][-1] < t:
+                last_start_end = [t_chunk_idx[-1], t]
+                t_chunk_start_end.append(last_start_end)
+        dec_ = []
+        for idx, (start, end) in enumerate(t_chunk_start_end):
+            chunk_x = x[:, :, start: end]
+            if idx != 0:
+                dec = self.tiled_decode2d(chunk_x)[:, :, 1:]
+            else:
+                dec = self.tiled_decode2d(chunk_x)
+            dec_.append(dec)
+        dec_ = torch.cat(dec_, dim=2)
+        return dec_
+
+    def tiled_encode2d(self, x, return_moments=False):
         overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
         blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
         row_limit = self.tile_latent_min_size - blend_extent
@@ -590,7 +640,8 @@ class CausalVAEModel(VideoBaseAE_PL):
 
         moments = torch.cat(result_rows, dim=3)
         posterior = DiagonalGaussianDistribution(moments)
-
+        if return_moments:
+            return moments
         return posterior
 
     def tiled_decode2d(self, z):