From 9b76a7825ea977dd6905cc271e9f8303a0ef49ac Mon Sep 17 00:00:00 2001 From: Aryan Date: Tue, 15 Oct 2024 15:36:19 +0200 Subject: [PATCH 1/2] update --- .../models/autoencoders/autoencoder_kl_cogvideox.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py index 7834206ddb4a..8fcfb01e5b9c 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py @@ -1182,7 +1182,8 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: frame_batch_size = self.num_sample_frames_batch_size # Note: We expect the number of frames to be either `1` or `frame_batch_size * k` or `frame_batch_size * k + 1` for some k. - num_batches = num_frames // frame_batch_size if num_frames > 1 else 1 + # So, it is okay to not round up as the extra remaining frame is handled in the loop + num_batches = max(num_frames // frame_batch_size, 1) conv_cache = None enc = [] @@ -1330,7 +1331,8 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: row = [] for j in range(0, width, overlap_width): # Note: We expect the number of frames to be either `1` or `frame_batch_size * k` or `frame_batch_size * k + 1` for some k. - num_batches = num_frames // frame_batch_size if num_frames > 1 else 1 + # So, it is okay to not round up as the extra remaining frame is handled in the loop + num_batches = max(num_frames // frame_batch_size, 1) conv_cache = None time = [] @@ -1409,7 +1411,7 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[Decod for i in range(0, height, overlap_height): row = [] for j in range(0, width, overlap_width): - num_batches = num_frames // frame_batch_size + num_batches = max(num_frames // frame_batch_size, 1) conv_cache = None time = [] From 8d9f7da80d8c93ec870cfa968bf2038de2d4356f Mon Sep 17 00:00:00 2001 From: Aryan Date: Wed, 16 Oct 2024 12:41:06 +0200 Subject: [PATCH 2/2] apply suggestions from review --- src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py index 8fcfb01e5b9c..68b49d72acc5 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py @@ -1182,7 +1182,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor: frame_batch_size = self.num_sample_frames_batch_size # Note: We expect the number of frames to be either `1` or `frame_batch_size * k` or `frame_batch_size * k + 1` for some k. - # So, it is okay to not round up as the extra remaining frame is handled in the loop + # As the extra single frame is handled inside the loop, it is not required to round up here. num_batches = max(num_frames // frame_batch_size, 1) conv_cache = None enc = [] @@ -1331,7 +1331,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: row = [] for j in range(0, width, overlap_width): # Note: We expect the number of frames to be either `1` or `frame_batch_size * k` or `frame_batch_size * k + 1` for some k. - # So, it is okay to not round up as the extra remaining frame is handled in the loop + # As the extra single frame is handled inside the loop, it is not required to round up here. num_batches = max(num_frames // frame_batch_size, 1) conv_cache = None time = []