restrict data range

Author: KimbingNg

train.py

@@ -564,16 +564,15 @@ def encode_images(self, images):
         """Encode images to vision states (for i2v)"""
         if self.vision_encoder is None:
             return None
-        if isinstance(images, torch.Tensor):
-            images_np = (images.cpu().permute(0, 2, 3, 1).numpy() * 255).astype("uint8")
-        else:
-            images_np = images
+        assert images.max() <= 1.0 and images.min() >= -1.0, f"Images must be in the range [-1, 1], but got {images.min()} {images.max()}"
+        images = (images + 1) / 2 # [-1, 1] -> [0, 1]
+        images_np = (images.cpu().permute(0, 2, 3, 1).numpy() * 255).clip(0, 255).astype("uint8")
         vision_states = self.vision_encoder.encode_images(images_np)
         return vision_states.last_hidden_state.to(device=self.device, dtype=self.transformer.dtype)
 
     def encode_vae(self, images: torch.Tensor) -> torch.Tensor:
-        if images.max() > 1.0:
-            images = images / 255.0
+        if images.max() > 1.0 or images.min() < -1.0:
+            raise ValueError(f"Images must be in the range [-1, 1], but got {images.min()} {images.max()}")
 
         if images.ndim == 4:
             images = images.unsqueeze(2)
@@ -674,7 +673,8 @@ def prepare_batch(self, batch: Dict[str, Any]) -> Dict[str, Any]:
                 byt5_text_mask = torch.cat(byt5_mask_list, dim=0)
 
         vision_states = None
-        if task_type == "i2v" and images is not None:
+        if task_type == "i2v":
+            assert images is not None, '`pixel_values` must be provided for i2v task'
             if images.ndim == 5:
                 first_frame = images[:, :, 0, :, :]
             else: