finetune/datasets/i2v_camera_dataset.py

import hashlib
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Tuple

import torch
from accelerate.logging import get_logger
from safetensors.torch import load_file, save_file
from torch.utils.data import Dataset
from torchvision import transforms
from typing_extensions import override

from finetune.constants import LOG_LEVEL, LOG_NAME

from .utils import (
    load_images,
    load_images_from_videos,
    load_prompts,
    load_videos,
    preprocess_image_with_resize,
    preprocess_video_with_buckets,
    preprocess_video_with_resize,
)

if TYPE_CHECKING:
    from finetune.trainer import Trainer

# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error
# Very few bug reports but it happens. Look in decord Github issues for more relevant information.
import decord  # isort:skip
import numpy as np
import os
from finetune.utils.camera_utils import get_camera_condition

decord.bridge.set_bridge("torch")

logger = get_logger(LOG_NAME, LOG_LEVEL)


class BaseI2VDataset(Dataset):
    """
    Base dataset class for Image-to-Video (I2V) training.

    This dataset loads prompts, videos and corresponding conditioning images for I2V training.

    Args:
        data_root (str): Root directory containing the dataset files
        caption_column (str): Path to file containing text prompts/captions
        video_column (str): Path to file containing video paths
        image_column (str): Path to file containing image paths
        device (torch.device): Device to load the data on
        encode_video_fn (Callable[[torch.Tensor], torch.Tensor], optional): Function to encode videos
    """

    def __init__(
            self,
            data_root: str,
            cache_root: str,
            metadata_path: str,
            enable_align_factor: bool,
            device: torch.device = torch.device("cpu"),
            trainer: "Trainer" = None,
            *args,
            **kwargs,
    ) -> None:
        super().__init__()

        self.trainer = trainer
        self.data_root = data_root
        self.cache_root = cache_root
        self.enable_align_factor = enable_align_factor

        self.train_resolution_str = "x".join(str(x) for x in self.trainer.args.train_resolution)
        self.video_latent_dir = cache_root / "video_latent" / self.trainer.args.model_name / self.train_resolution_str
        self.prompt_embeddings_dir = cache_root / "prompt_embeddings"
        self.video_latent_dir.mkdir(parents=True, exist_ok=True)
        self.prompt_embeddings_dir.mkdir(parents=True, exist_ok=True)

        self.all_metadata = np.load(os.path.join(data_root, metadata_path), allow_pickle=True)["arr_0"].tolist()
        logger.info(f"Data Count (all): {len(self.all_metadata)}", main_process_only=True)

        self.all_metadata = list(filter(lambda x: x['camera_extrinsics'].shape[0] > self.trainer.args.train_resolution[0], self.all_metadata))
        logger.info(f"Data Count (num_frames > {self.trainer.args.train_resolution[0]}): {len(self.all_metadata)}", main_process_only=True)

        logger.info(f"Data Count (final): {len(self.all_metadata)}", main_process_only=True)

        self.device = device
        self.encode_video = trainer.encode_video
        self.encode_text = trainer.encode_text

    def __len__(self) -> int:
        return len(self.all_metadata)

    def __getitem__(self, index: int) -> Dict[str, Any]:
        if isinstance(index, list):
            # Here, index is actually a list of data objects that we need to return.
            # The BucketSampler should ideally return indices. But, in the sampler, we'd like
            # to have information about num_frames, height and width. Since this is not stored
            # as metadata, we need to read the video to get this information. You could read this
            # information without loading the full video in memory, but we do it anyway. In order
            # to not load the video twice (once to get the metadata, and once to return the loaded video
            # based on sampled indices), we cache it in the BucketSampler. When the sampler is
            # to yield, we yield the cache data instead of indices. So, this special check ensures
            # that data is not loaded a second time. PRs are welcome for improvements.
            return index

        metadata = self.all_metadata[index % len(self.all_metadata)]
        video_path = Path(os.path.join(self.data_root, metadata['video_path']))
        image = load_images_from_videos([video_path])[0]
        prompt: str = metadata['long_caption']
        camera_extrinsics = torch.from_numpy(metadata['camera_extrinsics'])  # [F, 4, 4]
        fx, fy, cx, cy = metadata['camera_intrinsics']
        camera_intrinsics = torch.tensor([
            [fx, 0, cx],
            [0, fy, cy],
            [0, 0, 1]
        ])  # 3x3)  # [3, 3]
        align_factor = metadata['align_factor'] if self.enable_align_factor else 1.0

        prompt_hash = str(hashlib.sha256(prompt.encode()).hexdigest())
        prompt_embedding_path = self.prompt_embeddings_dir / (prompt_hash + ".safetensors")
        encoded_video_path = self.video_latent_dir / (video_path.stem + ".safetensors")

        if prompt_embedding_path.exists():
            prompt_embedding = load_file(prompt_embedding_path)["prompt_embedding"]
            logger.debug(
                f"process {self.trainer.accelerator.process_index}: Loaded prompt embedding from {prompt_embedding_path}",
                main_process_only=False,
            )
        else:
            prompt_embedding = self.encode_text(prompt)
            prompt_embedding = prompt_embedding.to("cpu")
            # [1, seq_len, hidden_size] -> [seq_len, hidden_size]
            prompt_embedding = prompt_embedding[0]
            try:
                save_file({"prompt_embedding": prompt_embedding}, prompt_embedding_path)
            except:
                pass
            logger.info(f"Saved prompt embedding to {prompt_embedding_path}", main_process_only=False)

        (
            frames,
            image,
            camera_extrinsics,
            camera_intrinsics,
        ) = self.preprocess(
            video_path,
            image,
            camera_extrinsics,
            camera_intrinsics,
            self.trainer.args.use_precompute_video_latents,
        )
        H, W = frames.shape[-2:]
        image = self.image_transform(image)
        frames = self.video_transform(frames)
        video = frames  # F, C, H, W

        if encoded_video_path.exists() and self.trainer.args.use_precompute_video_latents:
            encoded_video = load_file(encoded_video_path)["encoded_video"]
            logger.debug(f"Loaded encoded video from {encoded_video_path}", main_process_only=False)
        else:
            frames = frames.unsqueeze(0).permute(0, 2, 1, 3, 4).contiguous()  # [F, C, H, W] -> [B, C, F, H, W], value in [-1,1]
            encoded_video = self.encode_video(frames)
            encoded_video = encoded_video[0].to("cpu")  # [1, C, F, H, W] -> [C, F, H, W]
            if self.trainer.args.precompute:
                try:
                    save_file({"encoded_video": encoded_video}, encoded_video_path)
                except:
                    pass
                logger.info(f"Saved encoded video to {encoded_video_path}", main_process_only=False)

        if not self.trainer.args.precompute:
            # plucker embedding
            cond_frame_index = torch.zeros(1, device=camera_extrinsics.device, dtype=torch.long)
            plucker_embedding, relative_c2w_RT_4x4 = get_camera_condition(  # B 6 F H W
                H, W, camera_intrinsics.unsqueeze(0), camera_extrinsics.unsqueeze(0), mode="w2c",
                cond_frame_index=cond_frame_index, align_factor=align_factor
            )  # [B=1, C=6, F, H, W]
            plucker_embedding = plucker_embedding[0].contiguous()
        else:
            plucker_embedding = None

        ret = {
            "image": image,
            "prompt_embedding": prompt_embedding,  # [C, H, W]
            "prompt": prompt,
            "video": video,  # F, C, H, W
            "encoded_video": encoded_video,  # [C, F//4, H//8, W//8]
            "plucker_embedding": plucker_embedding,  # [B=1, C=6, F, H, W]
            "video_metadata": {
                "num_frames": encoded_video.shape[1],
                "height": encoded_video.shape[2],
                "width": encoded_video.shape[3],
            },
        }

        return ret

    def preprocess(self, video_path: Path | None, image_path: Path | None) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Loads and preprocesses a video and an image.
        If either path is None, no preprocessing will be done for that input.

        Args:
            video_path: Path to the video file to load
            image_path: Path to the image file to load

        Returns:
            A tuple containing:
                - video(torch.Tensor) of shape [F, C, H, W] where F is number of frames,
                  C is number of channels, H is height and W is width
                - image(torch.Tensor) of shape [C, H, W]
        """
        raise NotImplementedError("Subclass must implement this method")

    def video_transform(self, frames: torch.Tensor) -> torch.Tensor:
        """
        Applies transformations to a video.

        Args:
            frames (torch.Tensor): A 4D tensor representing a video
                with shape [F, C, H, W] where:
                - F is number of frames
                - C is number of channels (3 for RGB)
                - H is height
                - W is width

        Returns:
            torch.Tensor: The transformed video tensor
        """
        raise NotImplementedError("Subclass must implement this method")

    def image_transform(self, image: torch.Tensor) -> torch.Tensor:
        """
        Applies transformations to an image.

        Args:
            image (torch.Tensor): A 3D tensor representing an image
                with shape [C, H, W] where:
                - C is number of channels (3 for RGB)
                - H is height
                - W is width

        Returns:
            torch.Tensor: The transformed image tensor
        """
        raise NotImplementedError("Subclass must implement this method")


class I2VDatasetWithResize(BaseI2VDataset):
    """
    A dataset class for image-to-video generation that resizes inputs to fixed dimensions.

    This class preprocesses videos and images by resizing them to specified dimensions:
    - Videos are resized to max_num_frames x height x width
    - Images are resized to height x width

    Args:
        max_num_frames (int): Maximum number of frames to extract from videos
        height (int): Target height for resizing videos and images
        width (int): Target width for resizing videos and images
    """

    def __init__(self, max_num_frames: int, height: int, width: int, keep_aspect_ratio: bool, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self.max_num_frames = max_num_frames
        self.height = height
        self.width = width
        self.keep_aspect_ratio = keep_aspect_ratio

        self.__frame_transforms = transforms.Compose([transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)])
        self.__image_transforms = self.__frame_transforms

    def _resize_for_rectangle_crop(self, frames, H, W):
        '''
        :param frames: C,F,H,W
        :param image_size: H,W
        :return: frames: C,F,crop_H,crop_W;  camera_intrinsics: F,3,3
        '''
        ori_H, ori_W = frames.shape[-2:]

        # if ori_W / ori_H < 1.0:
        #     tmp_H, tmp_W = int(H), int(W)
        #     H, W = tmp_W, tmp_H

        if ori_W / ori_H > W / H:
            frames = transforms.functional.resize(
                frames,
                size=[H, int(ori_W * H / ori_H)],
            )
        else:
            frames = transforms.functional.resize(
                frames,
                size=[int(ori_H * W / ori_W), W],
            )

        resized_H, resized_W = frames.shape[2], frames.shape[3]
        frames = frames.squeeze(0)

        delta_H = resized_H - H
        delta_W = resized_W - W

        top, left = delta_H // 2, delta_W // 2
        frames = transforms.functional.crop(frames, top=top, left=left, height=H, width=W)

        return frames, resized_H, resized_W

    @override
    def preprocess(self, video_path: Path | None, image_path: Path | None, camera_pose_4x4, camera_intrinsics, use_precompute_video_latents=True):
        if video_path is not None:
            video, indices = preprocess_video_with_resize(
                video_path, self.max_num_frames, self.height, self.width,
                keep_aspect_ratio=self.keep_aspect_ratio,
                use_precompute_video_latents=use_precompute_video_latents
            )
            if self.keep_aspect_ratio:
                video, resized_H, resized_W = self._resize_for_rectangle_crop(video, self.height, self.width)
            else:
                resized_H, resized_W = video.shape[-2:]

            camera_pose_4x4 = camera_pose_4x4[indices]
            camera_intrinsics = camera_intrinsics.clone()
            cur_H, cur_W = video.shape[-2:]
            camera_intrinsics[0, 0] *= resized_W
            camera_intrinsics[0, 2] *= cur_W
            camera_intrinsics[1, 1] *= resized_H
            camera_intrinsics[1, 2] *= cur_H
            camera_intrinsics = camera_intrinsics.unsqueeze(0).repeat(camera_pose_4x4.shape[0], 1, 1)  # f,3,3
        else:
            video = None
        if image_path is not None and use_precompute_video_latents:
            image = preprocess_image_with_resize(image_path, self.height, self.width, keep_aspect_ratio=self.keep_aspect_ratio)
            if self.keep_aspect_ratio:
                image, resized_H, resized_W = self._resize_for_rectangle_crop(image.unsqueeze(0), self.height, self.width)
        elif not use_precompute_video_latents:
            image = video[0, :, :, :].clone()
        else:
            image = None
        
        return video, image, camera_pose_4x4, camera_intrinsics

    @override
    def video_transform(self, frames: torch.Tensor) -> torch.Tensor:
        return torch.stack([self.__frame_transforms(f) for f in frames], dim=0)

    @override
    def image_transform(self, image: torch.Tensor) -> torch.Tensor:
        return self.__image_transforms(image)


class I2VDatasetWithBuckets(BaseI2VDataset):
    def __init__(
            self,
            video_resolution_buckets: List[Tuple[int, int, int]],
            vae_temporal_compression_ratio: int,
            vae_height_compression_ratio: int,
            vae_width_compression_ratio: int,
            *args,
            **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)

        self.video_resolution_buckets = [
            (
                int(b[0] / vae_temporal_compression_ratio),
                int(b[1] / vae_height_compression_ratio),
                int(b[2] / vae_width_compression_ratio),
            )
            for b in video_resolution_buckets
        ]
        self.__frame_transforms = transforms.Compose([transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)])
        self.__image_transforms = self.__frame_transforms

    @override
    def preprocess(self, video_path: Path, image_path: Path) -> Tuple[torch.Tensor, torch.Tensor]:
        video = preprocess_video_with_buckets(video_path, self.video_resolution_buckets)
        image = preprocess_image_with_resize(image_path, video.shape[2], video.shape[3])
        return video, image

    @override
    def video_transform(self, frames: torch.Tensor) -> torch.Tensor:
        return torch.stack([self.__frame_transforms(f) for f in frames], dim=0)

    @override
    def image_transform(self, image: torch.Tensor) -> torch.Tensor:
        return self.__image_transforms(image)