fix minor bugs

Author: a-r-r-o-w

src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py

@@ -1090,7 +1090,7 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
 
         frame_batch_size = self.num_sample_frames_batch_size
         enc = []
-        for i in range(num_frames // frame_batch_size):
+        for i in range((num_frames + frame_batch_size - 1) // frame_batch_size):
             remaining_frames = num_frames % frame_batch_size
             start_frame = frame_batch_size * i + (0 if i == 0 else remaining_frames)
             end_frame = frame_batch_size * (i + 1) + remaining_frames
@@ -1141,7 +1141,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
 
         frame_batch_size = self.num_latent_frames_batch_size
         dec = []
-        for i in range(num_frames // frame_batch_size):
+        for i in range((num_frames + frame_batch_size - 1) // frame_batch_size):
             remaining_frames = num_frames % frame_batch_size
             start_frame = frame_batch_size * i + (0 if i == 0 else remaining_frames)
             end_frame = frame_batch_size * (i + 1) + remaining_frames

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py

@@ -367,6 +367,7 @@ def prepare_latents(
         if latents is None:
             assert image.ndim == 4
             image = image.unsqueeze(2)  # [B, C, F, H, W]
+            print(image.shape)
 
             if isinstance(generator, list):
                 if len(generator) != batch_size:
@@ -392,6 +393,7 @@ def prepare_latents(
                 width // self.vae_scale_factor_spatial,
             )
             latent_padding = torch.zeros(padding_shape, device=device, dtype=dtype)
+            print(init_latents.shape, latent_padding.shape)
             init_latents = torch.cat([init_latents, latent_padding], dim=1)
 
             noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
@@ -723,10 +725,11 @@ def __call__(
         self._num_timesteps = len(timesteps)
 
         # 5. Prepare latents
-        image = self.video_processor.preprocess(image, height=height, width=width).to(device)
-        image = image.unsqueeze(2)  # [B, C, F, H, W]
+        image = self.video_processor.preprocess(image, height=height, width=width).to(
+            device, dtype=prompt_embeds.dtype
+        )
 
-        latent_channels = self.transformer.config.in_channels
+        latent_channels = self.transformer.config.in_channels // 2
         latents = self.prepare_latents(
             image,
             batch_size * num_videos_per_prompt,