Blame for toolkit/data_transfer_object/data_loader.py

import os
from typing import TYPE_CHECKING, List, Union
import cv2
import torch

from PIL import Image
from PIL.ImageOps import exif_transpose
import av
            
from toolkit import image_utils
from toolkit.basic import get_quick_signature_string
from toolkit.dto import DTO
from toolkit.dataloader_mixins import (
    _latent_from_uint8,
    CaptionProcessingDTOMixin,
    ImageProcessingDTOMixin,
    LatentCachingFileItemDTOMixin,
    ControlFileItemDTOMixin,
    ArgBreakMixin,
    MaskFileItemDTOMixin,
    AugmentationFileItemDTOMixin,
    UnconditionalFileItemDTOMixin,
    ClipImageFileItemDTOMixin,
    InpaintControlFileItemDTOMixin,
    TextEmbeddingFileItemDTOMixin,
    AudioProcessingDTOMixin,
)
from toolkit.prompt_utils import PromptEmbeds, concat_prompt_embeds

if TYPE_CHECKING:
    from toolkit.config_modules import DatasetConfig

printed_messages = []

# keep in sync with video_extensions in toolkit/data_loader.py (importing it
# here would be circular)
video_extensions = ['.mp4', '.avi', '.mov', '.webm', '.mkv', '.wmv', '.m4v', '.flv']


def print_once(msg):
    global printed_messages
    if msg not in printed_messages:
        print(msg)
        printed_messages.append(msg)


class FileItemDTO(
    LatentCachingFileItemDTOMixin,
    TextEmbeddingFileItemDTOMixin,
    CaptionProcessingDTOMixin,
    ImageProcessingDTOMixin,
    AudioProcessingDTOMixin,
    ControlFileItemDTOMixin,
    InpaintControlFileItemDTOMixin,
    ClipImageFileItemDTOMixin,
    MaskFileItemDTOMixin,
    AugmentationFileItemDTOMixin,
    UnconditionalFileItemDTOMixin,
    ArgBreakMixin,
):
    def __init__(self, *args, **kwargs):
        self.path = kwargs.get("path", "")
        self.dataset_config: "DatasetConfig" = kwargs.get("dataset_config", None)
        # a video dataset can contain both videos and images. Images are
        # treated as single-frame items and bucketed separately from videos
        dataset_is_video = self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count
        self.is_video = dataset_is_video and os.path.splitext(self.path)[1].lower() in video_extensions
        self.is_audio_model = kwargs.get("is_audio_model", False)
        self.sample_rate = kwargs.get("sample_rate", 48000)
        self.num_frames = self.dataset_config.num_frames if self.is_video else 1
        self.temporal_compression = kwargs.get("temporal_compression", 8)
        # module-level function (picklable) for models whose valid frame
        # counts are not temporal_compression * n + 1; None = default math
        _sd = kwargs.get("sd", None)
        self.frame_count_snapper = (
            _sd.get_frame_count_snapper()
            if _sd is not None and hasattr(_sd, "get_frame_count_snapper")
            else None
        )
        # models with an RGBA VAE want the alpha channel loaded and encoded
        self.load_rgba = _sd is not None and _sd.load_rgba
        size_database = kwargs.get("size_database", {})
        dataset_root = kwargs.get("dataset_root", None)
        self.encode_control_in_text_embeddings = kwargs.get(
            "encode_control_in_text_embeddings", False
        )
        self.encode_first_frame_in_text_embeddings = kwargs.get(
            "encode_first_frame_in_text_embeddings", False
        )
        # D-OPSD: also cache teacher embeds with the item's own media as reference 1
        self.dopsd_self_ref = kwargs.get("dopsd_self_ref", False)
        self.te_padding_side = kwargs.get("te_padding_side", "right")
        self.latent_space_version = kwargs.get("latent_space_version", "sd1")
        self.text_embedding_space_version = kwargs.get("text_embedding_space_version", "sd1")
        # model sizes references against the item's bucket, so the cache key must carry it
        self.text_embedding_uses_target_size = kwargs.get("text_embedding_uses_target_size", False)
        if dataset_root is not None:
            # remove dataset root from path
            file_key = self.path.replace(dataset_root, "")
        else:
            file_key = os.path.basename(self.path)

        file_signature = get_quick_signature_string(self.path)
        if file_signature is None:
            raise Exception("Error: Could not get file signature for {self.path}")

        use_db_entry = False
        db_entry = None
        if file_key in size_database:
            db_entry = size_database[file_key]
            if (
                db_entry is not None
                and len(db_entry) >= 3
                and db_entry[2] == file_signature
            ):
                use_db_entry = True
        video_total_frames = None
        video_fps = None
        if self.is_audio_model:
            # get the length of the audio file in ms
            with av.open(self.path) as c:
                if c.duration is not None:
                    w =  int(c.duration / 1_000)
                else:
                    s = c.streams.audio[0]
                    w = int(float(s.duration * s.time_base) * 1_000)
            h = 1
        elif self.is_video:
            # video entries also carry (total_frames, fps); older 3-item entries
            # get re-read and upgraded here
            if use_db_entry and len(db_entry) >= 5:
                w, h, _, video_total_frames, video_fps = db_entry[:5]
            else:
                # Open the video file
                video = cv2.VideoCapture(self.path)

                # Check if video opened successfully
                if not video.isOpened():
                    raise Exception(f"Error: Could not open video file {self.path}")

                # Get width and height
                width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
                height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
                w, h = width, height
                video_total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
                video_fps = video.get(cv2.CAP_PROP_FPS)

                # Release the video capture object immediately
                video.release()
                size_database[file_key] = (width, height, file_signature, video_total_frames, video_fps)
        elif use_db_entry:
            w, h, _ = db_entry[:3]
        else:
            if self.dataset_config.fast_image_size:
                # original method is significantly faster, but some images are read sideways. Not sure why. Do slow method by default.
                try:
                    w, h = image_utils.get_image_size(self.path)
                except image_utils.UnknownImageFormat:
                    print_once(
                        f"Warning: Some images in the dataset cannot be fast read. "
                        + f"This process is faster for png, jpeg"
                    )
                    img = exif_transpose(Image.open(self.path))
                    w, h = img.size
            else:
                img = exif_transpose(Image.open(self.path))
                w, h = img.size
            size_database[file_key] = (w, h, file_signature)
        self.width: int = w
        self.height: int = h
        if self.is_video and self.dataset_config.auto_frame_count:
            # compute the real frame count now (same math as load time) so buckets
            # are keyed on the frame count this video will actually train at
            self.num_frames = self.get_auto_frame_count(video_total_frames, video_fps)
        self.dataloader_transforms = kwargs.get("dataloader_transforms", None)
        super().__init__(*args, **kwargs)

        # self.caption_path: str = kwargs.get('caption_path', None)
        self.raw_caption: str = kwargs.get("raw_caption", None)
        # we scale first, then crop
        self.scale_to_width: int = kwargs.get(
            "scale_to_width", int(self.width * self.dataset_config.scale)
        )
        self.scale_to_height: int = kwargs.get(
            "scale_to_height", int(self.height * self.dataset_config.scale)
        )
        # crop values are from scaled size
        self.crop_x: int = kwargs.get("crop_x", 0)
        self.crop_y: int = kwargs.get("crop_y", 0)
        self.crop_width: int = kwargs.get("crop_width", self.scale_to_width)
        self.crop_height: int = kwargs.get("crop_height", self.scale_to_height)
        self.flip_x: bool = kwargs.get("flip_x", False)
        self.flip_y: bool = kwargs.get("flip_x", False)
        self.augments: List[str] = self.dataset_config.augments
        self.loss_multiplier: float = self.dataset_config.loss_multiplier

        self.network_weight: float = self.dataset_config.network_weight
        self.is_reg = self.dataset_config.is_reg
        self.prior_reg = self.dataset_config.prior_reg
        self.tensor: Union[torch.Tensor, None] = None
        self.audio_data = None
        self.audio_tensor = None

    def cleanup(self):
        self.tensor = None
        self.audio_data = None
        self.audio_tensor = None
        self.cleanup_latent()
        self.cleanup_text_embedding()
        self.cleanup_control()
        self.cleanup_inpaint()
        self.cleanup_clip_image()
        self.cleanup_mask()
        self.cleanup_unconditional()

    @property
    def tensor(self) -> Union[torch.Tensor, None]:
        t = self.__dict__.get('_tensor')
        if t is None:
            cached = getattr(self, '_cached_tensor_uint8', None)
            if cached is not None:
                # cached pixels stay uint8 until first read, so workers ship 1 byte/px over shm;
                # fp16 on read (consumers cast up themselves)
                t = _latent_from_uint8(cached, dtype=torch.float16)
                self._tensor = t
        return t

    @tensor.setter
    def tensor(self, value):
        self._tensor = value

    @tensor.deleter
    def tensor(self):
        self._tensor = None


class DataLoaderBatchDTO:
    def __init__(self, **kwargs):
        try:
            self.file_items: List["FileItemDTO"] = kwargs.get("file_items", None)
            is_latents_cached = self.file_items[0].is_latent_cached
            self._tensor: Union[torch.Tensor, None] = None
            self._tensor_uint8: Union[torch.Tensor, None] = None
            self.latents: Union[torch.Tensor, None] = None
            self.control_tensor: Union[torch.Tensor, None] = None
            self.control_tensor_list: Union[List[List[torch.Tensor]], None] = None
            self.clip_image_tensor: Union[torch.Tensor, None] = None
            self.mask_tensor: Union[torch.Tensor, None] = None
            self.unaugmented_tensor: Union[torch.Tensor, None] = None
            self.unconditional_tensor: Union[torch.Tensor, None] = None
            self.unconditional_latents: Union[torch.Tensor, None] = None
            self.clip_image_embeds: Union[List[dict], None] = None
            self.clip_image_embeds_unconditional: Union[List[dict], None] = None
            self.sigmas: Union[torch.Tensor, None] = (
                None  # can be added elseware and passed along training code
            )
            self.extra_values: Union[torch.Tensor, None] = (
                torch.tensor([x.extra_values for x in self.file_items])
                if len(self.file_items[0].extra_values) > 0
                else None
            )
            self.audio_data: Union[List, None] = (
                [x.audio_data for x in self.file_items]
                if self.file_items[0].audio_data is not None
                else None
            )
            self.audio_tensor: Union[torch.Tensor, None] = None
            self.first_frame_latents: Union[torch.Tensor, None] = None
            # control-video reference paths (encoded + disk-cached lazily by
            # models with supports_video_control_images)
            self.control_video_paths_list: Union[List, None] = None
            if any(getattr(x, 'control_video_paths', None) for x in self.file_items):
                self.control_video_paths_list = [
                    list(getattr(x, 'control_video_paths', None) or [])
                    for x in self.file_items
                ]

            # set by the trainer around the D-OPSD teacher pass
            self.dopsd_teacher_pass: bool = False

            self.num_frames: int = self.file_items[0].num_frames

            if (
                not is_latents_cached
                or self.file_items[0].dataset_config.load_image_when_caching_latents
                or self.file_items[0].dataset_config.cache_tensors_to_disk
            ):
                # only return a tensor if latents are not cached, or if we are explicitly
                # loading the raw image alongside the cached latents
                if all(x.__dict__.get('_tensor') is None and getattr(x, '_cached_tensor_uint8', None) is not None for x in self.file_items):
                    # cached pixels cross the worker boundary as uint8; `tensor` expands them on first read
                    self._tensor_uint8 = torch.cat([x._cached_tensor_uint8.unsqueeze(0) for x in self.file_items])
                else:
                    self.tensor = torch.cat([x.tensor.unsqueeze(0) for x in self.file_items])
            # if we have encoded latents, we concatenate them
            self.latents: Union[torch.Tensor, None] = None
            if is_latents_cached:
                # this get_latent call with trigger loading all cached items from the disk.
                # DTO.stack keeps any extra streams (audio, ...) riding on the batch latents
                self.latents = DTO.stack([x.get_latent() for x in self.file_items])
                if any(
                    [x._cached_first_frame_latent is not None for x in self.file_items]
                ):
                    # find one to use as a base; item 0 may not have one
                    base_first_frame_latent = None
                    for x in self.file_items:
                        if x._cached_first_frame_latent is not None:
                            base_first_frame_latent = x._cached_first_frame_latent
                            break
                    self.first_frame_latents = torch.cat(
                        [
                            x._cached_first_frame_latent.unsqueeze(0)
                            if x._cached_first_frame_latent is not None
                            else torch.zeros_like(base_first_frame_latent).unsqueeze(0)
                            for x in self.file_items
                        ]
                    )
            self.prompt_embeds: Union[PromptEmbeds, None] = None
            # diff output preservation embeds (trigger word replaced with class)
            self.dop_prompt_embeds: Union[PromptEmbeds, None] = None
            # D-OPSD teacher embeds (trigger word replaced with the self-reference token)
            self.dopsd_prompt_embeds: Union[PromptEmbeds, None] = None
            # if self.file_items[0].control_tensor is not None:
            # if any have a control tensor, we concatenate them
            if any([x.control_tensor is not None for x in self.file_items]):
                # find one to use as a base
                base_control_tensor = None
                for x in self.file_items:
                    if x.control_tensor is not None:
                        base_control_tensor = x.control_tensor
                        break
                control_tensors = []
                for x in self.file_items:
                    if x.control_tensor is None:
                        control_tensors.append(torch.zeros_like(base_control_tensor))
                    else:
                        control_tensors.append(x.control_tensor)
                self.control_tensor = torch.cat(
                    [x.unsqueeze(0) for x in control_tensors]
                )

            # handle control tensor list
            if any([x.control_tensor_list is not None for x in self.file_items]):
                self.control_tensor_list = []
                for x in self.file_items:
                    if x.control_tensor_list is not None:
                        self.control_tensor_list.append(x.control_tensor_list)
                    else:
                        raise Exception(
                            f"Could not find control tensors for all file items, missing for {x.path}"
                        )

            self.inpaint_tensor: Union[torch.Tensor, None] = None
            if any([x.inpaint_tensor is not None for x in self.file_items]):
                # find one to use as a base
                base_inpaint_tensor = None
                for x in self.file_items:
                    if x.inpaint_tensor is not None:
                        base_inpaint_tensor = x.inpaint_tensor
                        break
                inpaint_tensors = []
                for x in self.file_items:
                    if x.inpaint_tensor is None:
                        inpaint_tensors.append(torch.zeros_like(base_inpaint_tensor))
                    else:
                        inpaint_tensors.append(x.inpaint_tensor)
                self.inpaint_tensor = torch.cat(
                    [x.unsqueeze(0) for x in inpaint_tensors]
                )

            self.loss_multiplier_list: List[float] = [
                x.loss_multiplier for x in self.file_items
            ]

            if any([x.clip_image_tensor is not None for x in self.file_items]):
                # find one to use as a base
                base_clip_image_tensor = None
                for x in self.file_items:
                    if x.clip_image_tensor is not None:
                        base_clip_image_tensor = x.clip_image_tensor
                        break
                clip_image_tensors = []
                for x in self.file_items:
                    if x.clip_image_tensor is None:
                        clip_image_tensors.append(
                            torch.zeros_like(base_clip_image_tensor)
                        )
                    else:
                        clip_image_tensors.append(x.clip_image_tensor)
                self.clip_image_tensor = torch.cat(
                    [x.unsqueeze(0) for x in clip_image_tensors]
                )

            if any([x.mask_tensor is not None for x in self.file_items]):
                # find one to use as a base
                base_mask_tensor = None
                for x in self.file_items:
                    if x.mask_tensor is not None:
                        base_mask_tensor = x.mask_tensor
                        break
                mask_tensors = []
                for x in self.file_items:
                    if x.mask_tensor is None:
                        mask_tensors.append(torch.zeros_like(base_mask_tensor))
                    else:
                        mask_tensors.append(x.mask_tensor)
                self.mask_tensor = torch.cat([x.unsqueeze(0) for x in mask_tensors])

            # add unaugmented tensors for ones with augments
            if any([x.unaugmented_tensor is not None for x in self.file_items]):
                # find one to use as a base
                base_unaugmented_tensor = None
                for x in self.file_items:
                    if x.unaugmented_tensor is not None:
                        base_unaugmented_tensor = x.unaugmented_tensor
                        break
                unaugmented_tensor = []
                for x in self.file_items:
                    if x.unaugmented_tensor is None:
                        unaugmented_tensor.append(
                            torch.zeros_like(base_unaugmented_tensor)
                        )
                    else:
                        unaugmented_tensor.append(x.unaugmented_tensor)
                self.unaugmented_tensor = torch.cat(
                    [x.unsqueeze(0) for x in unaugmented_tensor]
                )

            # add unconditional tensors
            if any([x.unconditional_tensor is not None for x in self.file_items]):
                # find one to use as a base
                base_unconditional_tensor = None
                for x in self.file_items:
                    if x.unaugmented_tensor is not None:
                        base_unconditional_tensor = x.unconditional_tensor
                        break
                unconditional_tensor = []
                for x in self.file_items:
                    if x.unconditional_tensor is None:
                        unconditional_tensor.append(
                            torch.zeros_like(base_unconditional_tensor)
                        )
                    else:
                        unconditional_tensor.append(x.unconditional_tensor)
                self.unconditional_tensor = torch.cat(
                    [x.unsqueeze(0) for x in unconditional_tensor]
                )

            if any([x.clip_image_embeds is not None for x in self.file_items]):
                self.clip_image_embeds = []
                for x in self.file_items:
                    if x.clip_image_embeds is not None:
                        self.clip_image_embeds.append(x.clip_image_embeds)
                    else:
                        raise Exception("clip_image_embeds is None for some file items")

            if any(
                [x.clip_image_embeds_unconditional is not None for x in self.file_items]
            ):
                self.clip_image_embeds_unconditional = []
                for x in self.file_items:
                    if x.clip_image_embeds_unconditional is not None:
                        self.clip_image_embeds_unconditional.append(
                            x.clip_image_embeds_unconditional
                        )
                    else:
                        raise Exception(
                            "clip_image_embeds_unconditional is None for some file items"
                        )

            if any([x.prompt_embeds is not None for x in self.file_items]):
                # find one to use as a base
                base_prompt_embeds = None
                for x in self.file_items:
                    if x.prompt_embeds is not None:
                        base_prompt_embeds = x.prompt_embeds
                        break
                prompt_embeds_list = []
                for x in self.file_items:
                    if x.prompt_embeds is None:
                        y = base_prompt_embeds
                    else:
                        y = x.prompt_embeds
                    if x.text_embedding_space_version == "zimage":
                        # z image needs to be a list if it is not already
                        if not isinstance(y.text_embeds, list):
                            y.text_embeds = [y.text_embeds]
                    prompt_embeds_list.append(y)
                padding_side = self.file_items[0].te_padding_side

                self.prompt_embeds = concat_prompt_embeds(prompt_embeds_list, padding_side=padding_side)

            if any([x.dop_prompt_embeds is not None for x in self.file_items]):
                # find one to use as a base
                base_dop_prompt_embeds = None
                for x in self.file_items:
                    if x.dop_prompt_embeds is not None:
                        base_dop_prompt_embeds = x.dop_prompt_embeds
                        break
                dop_prompt_embeds_list = []
                for x in self.file_items:
                    if x.dop_prompt_embeds is None:
                        y = base_dop_prompt_embeds
                    else:
                        y = x.dop_prompt_embeds
                    if x.text_embedding_space_version == "zimage":
                        # z image needs to be a list if it is not already
                        if not isinstance(y.text_embeds, list):
                            y.text_embeds = [y.text_embeds]
                    dop_prompt_embeds_list.append(y)
                padding_side = self.file_items[0].te_padding_side

                self.dop_prompt_embeds = concat_prompt_embeds(dop_prompt_embeds_list, padding_side=padding_side)

            if any([getattr(x, 'dopsd_prompt_embeds', None) is not None for x in self.file_items]):
                # find one to use as a base
                base_dopsd_prompt_embeds = None
                for x in self.file_items:
                    if x.dopsd_prompt_embeds is not None:
                        base_dopsd_prompt_embeds = x.dopsd_prompt_embeds
                        break
                dopsd_prompt_embeds_list = []
                for x in self.file_items:
                    if x.dopsd_prompt_embeds is None:
                        y = base_dopsd_prompt_embeds
                    else:
                        y = x.dopsd_prompt_embeds
                    dopsd_prompt_embeds_list.append(y)
                padding_side = self.file_items[0].te_padding_side

                self.dopsd_prompt_embeds = concat_prompt_embeds(dopsd_prompt_embeds_list, padding_side=padding_side)

            if any([x.audio_tensor is not None for x in self.file_items]):
                # find one to use as a base
                base_audio_tensor = None
                for x in self.file_items:
                    if x.audio_tensor is not None:
                        base_audio_tensor = x.audio_tensor
                        break
                audio_tensors = []
                for x in self.file_items:
                    if x.audio_tensor is None:
                        audio_tensors.append(torch.zeros_like(base_audio_tensor))
                    else:
                        audio_tensors.append(x.audio_tensor)
                self.audio_tensor = torch.cat([x.unsqueeze(0) for x in audio_tensors])

        except Exception as e:
            print(e)
            raise e

    def get_is_reg_list(self):
        return [x.is_reg for x in self.file_items]

    def get_network_weight_list(self):
        return [x.network_weight for x in self.file_items]

    def get_caption_list(
        self, trigger=None, to_replace_list=None, add_if_not_present=True
    ):
        return [x.caption for x in self.file_items]

    def get_caption_short_list(
        self, trigger=None, to_replace_list=None, add_if_not_present=True
    ):
        return [x.caption_short for x in self.file_items]

    @property
    def latents(self) -> Union[torch.Tensor, None]:
        return self._latents

    @latents.setter
    def latents(self, value):
        # math on a DTO returns a plain tensor, so trainer code that rescales
        # or re-noises the latents would silently drop the extra streams
        # (audio, ...) riding on them. Re-attach them here so assignments stay
        # plain tensor code everywhere else.
        prev = getattr(self, '_latents', None)
        if isinstance(prev, DTO) and value is not None and not isinstance(value, DTO):
            value = DTO(value, **prev.extras)
        self._latents = value

    @latents.deleter
    def latents(self):
        self._latents = None

    @property
    def tensor(self) -> Union[torch.Tensor, None]:
        if self._tensor is None and self._tensor_uint8 is not None:
            self._tensor = _latent_from_uint8(self._tensor_uint8, dtype=torch.float16)
            self._tensor_uint8 = None
        return self._tensor

    @tensor.setter
    def tensor(self, value):
        self._tensor = value
        self._tensor_uint8 = None

    @tensor.deleter
    def tensor(self):
        self._tensor = None
        self._tensor_uint8 = None

    @property
    def audio_latents(self) -> Union[torch.Tensor, None]:
        # cached audio rides inside the latents DTO
        if isinstance(self.latents, DTO):
            return self.latents.get('audio')
        return None

    def cleanup(self):
        del self.latents
        del self.tensor
        del self.control_tensor
        del self.audio_tensor
        del self.audio_data
        del self.first_frame_latents
        for file_item in self.file_items:
            file_item.cleanup()

    @property
    def dataset_config(self) -> "DatasetConfig":
        if len(self.file_items) > 0:
            return self.file_items[0].dataset_config
        else:
            return None