Spaces:
Running on Zero
Running on Zero
| 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() | |
| 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 | |
| def tensor(self, value): | |
| self._tensor = value | |
| 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] | |
| def latents(self) -> Union[torch.Tensor, None]: | |
| return self._latents | |
| 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 | |
| def latents(self): | |
| self._latents = None | |
| 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 | |
| def tensor(self, value): | |
| self._tensor = value | |
| self._tensor_uint8 = None | |
| def tensor(self): | |
| self._tensor = None | |
| self._tensor_uint8 = None | |
| 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() | |
| def dataset_config(self) -> "DatasetConfig": | |
| if len(self.file_items) > 0: | |
| return self.file_items[0].dataset_config | |
| else: | |
| return None | |