601 lines
27 KiB
Python
601 lines
27 KiB
Python
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.dataloader_mixins import (
|
|
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)
|
|
self.dataset_relative_item_id: str = kwargs.get("dataset_relative_item_id", '')
|
|
self.caption_sources_raw: dict = kwargs.get("caption_sources_raw", {})
|
|
self.caption_source_templates: dict = kwargs.get("caption_source_templates", {})
|
|
self.caption_source_paths: dict = kwargs.get("caption_source_paths", {})
|
|
# 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
|
|
)
|
|
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
|
|
)
|
|
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")
|
|
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)
|
|
self.caption_template: str = kwargs.get("caption_template", None)
|
|
self.dataset_relative_item_id = kwargs.get("dataset_relative_item_id", self.dataset_relative_item_id)
|
|
self.caption_sources_raw = kwargs.get("caption_sources_raw", self.caption_sources_raw)
|
|
self.caption_source_templates = kwargs.get("caption_source_templates", self.caption_source_templates)
|
|
self.caption_source_paths = kwargs.get("caption_source_paths", self.caption_source_paths)
|
|
# 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()
|
|
|
|
|
|
class DataLoaderBatchDTO:
|
|
def get_caption_source_templates(self, source_names: List[str]) -> List[str]:
|
|
if len(source_names) != len(self.file_items):
|
|
raise ValueError('caption source selection count must match batch size')
|
|
templates = []
|
|
for item, source_name in zip(self.file_items, source_names):
|
|
if source_name not in item.caption_source_templates:
|
|
raise ValueError(f'caption source {source_name!r} is unavailable for {item.dataset_relative_item_id}')
|
|
templates.append(item.caption_source_templates[source_name])
|
|
return templates
|
|
|
|
def __init__(self, **kwargs):
|
|
try:
|
|
self.file_items: List["FileItemDTO"] = kwargs.get("file_items", None)
|
|
self.dataset_relative_item_ids = [item.dataset_relative_item_id for item in self.file_items]
|
|
is_latents_cached = self.file_items[0].is_latent_cached
|
|
self.tensor: 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
|
|
self.audio_latents: Union[torch.Tensor, None] = None
|
|
|
|
# just for holding noise and preds during training
|
|
self.audio_target: Union[torch.Tensor, None] = None
|
|
self.audio_pred: Union[torch.Tensor, None] = None
|
|
# the noise drawn for the audio stream on the primary (grad enabled)
|
|
# prediction. Secondary passes (cfg / guidance loss / prior preds)
|
|
# reuse it so their noisy audio matches the stored audio_target.
|
|
self.audio_noise: Union[torch.Tensor, None] = None
|
|
# audio predictions from the non primary passes. Kept separate so
|
|
# they cannot stomp the primary pred we backprop through.
|
|
self.audio_pred_uncond: Union[torch.Tensor, None] = None
|
|
self.audio_pred_prior: Union[torch.Tensor, None] = None
|
|
self.audio_pred_preservation: Union[torch.Tensor, None] = None
|
|
# which of the above the current secondary pass writes to. None (the
|
|
# default) means no secondary pass is in flight: any grad-enabled
|
|
# prediction is a primary one and writes audio_pred (and the
|
|
# noisy/sigma bookkeeping) directly. The trainer sets this around
|
|
# its prior / guidance-unconditional / preservation passes.
|
|
self.audio_pred_slot: Union[str, None] = None
|
|
# noisy audio rows and audio sigma of the primary pass, used to
|
|
# rebuild the clean audio estimate for perceptual losses
|
|
self.audio_noisy: Union[torch.Tensor, None] = None
|
|
self.audio_sigma: Union[torch.Tensor, None] = None
|
|
|
|
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
|
|
self.tensor: torch.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
|
|
self.latents = torch.cat(
|
|
[x.get_latent().unsqueeze(0) 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
|
|
]
|
|
)
|
|
if any([x._cached_audio_latent is not None for x in self.file_items]):
|
|
# find one to use as a base; item 0 may not have one
|
|
base_audio_latent = None
|
|
for x in self.file_items:
|
|
if x._cached_audio_latent is not None:
|
|
base_audio_latent = x._cached_audio_latent
|
|
break
|
|
self.audio_latents = torch.cat(
|
|
[
|
|
x._cached_audio_latent.unsqueeze(0)
|
|
if x._cached_audio_latent is not None
|
|
else torch.zeros_like(base_audio_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
|
|
# 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([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 set_secondary_audio_pred(self, pred):
|
|
"""Route an audio prediction from a non primary pass (prior,
|
|
unconditional/guidance, preservation) to its own slot so it cannot
|
|
stomp the primary prediction the loss backprops through. Passes that
|
|
did not declare a slot (e.g. a trainer's extra no_grad prediction)
|
|
are simply not stored."""
|
|
if self.audio_pred_slot is not None:
|
|
setattr(self, self.audio_pred_slot, pred)
|
|
|
|
def cleanup(self):
|
|
del self.latents
|
|
del self.tensor
|
|
del self.control_tensor
|
|
del self.audio_tensor
|
|
del self.audio_data
|
|
del self.audio_target
|
|
del self.audio_pred
|
|
del self.audio_noise
|
|
del self.audio_pred_uncond
|
|
del self.audio_pred_prior
|
|
del self.audio_pred_preservation
|
|
del self.audio_noisy
|
|
del self.audio_sigma
|
|
del self.first_frame_latents
|
|
del self.audio_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
|