808 lines
34 KiB
Python
808 lines
34 KiB
Python
from pathlib import Path
|
|
import os.path
|
|
import random
|
|
from collections import defaultdict
|
|
import math
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from deepspeed.utils.logging import logger
|
|
from deepspeed import comm as dist
|
|
import datasets
|
|
from datasets.fingerprint import Hasher
|
|
from PIL import Image
|
|
import imageio
|
|
import multiprocess as mp
|
|
|
|
from utils.common import is_main_process, VIDEO_EXTENSIONS, log_duration
|
|
|
|
|
|
DEBUG = False
|
|
IMAGE_SIZE_ROUND_TO_MULTIPLE = 32
|
|
NUM_PROC = min(8, os.cpu_count())
|
|
|
|
|
|
def shuffle_with_seed(l, seed=None):
|
|
rng_state = random.getstate()
|
|
random.seed(seed)
|
|
random.shuffle(l)
|
|
random.setstate(rng_state)
|
|
|
|
|
|
def process_caption_fn(shuffle_tags=False, caption_prefix=''):
|
|
def fn(example):
|
|
with open(example['caption_file']) as f:
|
|
caption = f.read().strip()
|
|
if shuffle_tags:
|
|
tags = [tag.strip() for tag in caption.split(',')]
|
|
random.shuffle(tags)
|
|
caption = ', '.join(tags)
|
|
caption = caption_prefix + caption
|
|
|
|
example['caption'] = caption
|
|
return example
|
|
return fn
|
|
|
|
|
|
def round_to_multiple(x, multiple):
|
|
return int(round(x / multiple) * multiple)
|
|
|
|
|
|
def _map_and_cache(dataset, map_fn, cache_dir, cache_file_prefix='', new_fingerprint_args=None, regenerate_cache=False, caching_batch_size=1, with_indices=False):
|
|
# Do the fingerprinting ourselves, because otherwise map() does it by serializing the map function.
|
|
# That goes poorly when the function is capturing huge models (slow, OOMs, etc).
|
|
new_fingerprint_args = [] if new_fingerprint_args is None else new_fingerprint_args
|
|
new_fingerprint_args.append(dataset._fingerprint)
|
|
new_fingerprint = Hasher.hash(new_fingerprint_args)
|
|
cache_file = cache_dir / f'{cache_file_prefix}{new_fingerprint}.arrow'
|
|
cache_file = str(cache_file)
|
|
dataset = dataset.map(
|
|
map_fn,
|
|
cache_file_name=cache_file,
|
|
load_from_cache_file=(not regenerate_cache),
|
|
writer_batch_size=100,
|
|
new_fingerprint=new_fingerprint,
|
|
remove_columns=dataset.column_names,
|
|
batched=True,
|
|
batch_size=caching_batch_size,
|
|
with_indices=with_indices,
|
|
num_proc=NUM_PROC,
|
|
)
|
|
dataset.set_format('torch')
|
|
return dataset
|
|
|
|
|
|
# The smallest unit of a dataset. Represents a single size bucket from a single folder of images
|
|
# and captions on disk. Not batched; returns individual items.
|
|
class SizeBucketDataset:
|
|
def __init__(self, metadata_dataset, directory_config, size_bucket, model_name):
|
|
self.metadata_dataset = metadata_dataset
|
|
self.directory_config = directory_config
|
|
self.size_bucket = size_bucket
|
|
self.model_name = model_name
|
|
self.path = Path(self.directory_config['path'])
|
|
self.cache_dir = self.path / 'cache' / self.model_name / f'cache_{size_bucket[0]}x{size_bucket[1]}x{size_bucket[2]}'
|
|
os.makedirs(self.cache_dir, exist_ok=True)
|
|
self.text_embedding_datasets = []
|
|
self.num_repeats = self.directory_config.get('num_repeats', 1)
|
|
|
|
def cache_latents(self, map_fn, regenerate_cache=False, caching_batch_size=1):
|
|
print(f'caching latents: {self.size_bucket}')
|
|
self.latent_dataset = _map_and_cache(
|
|
self.metadata_dataset,
|
|
map_fn,
|
|
self.cache_dir,
|
|
cache_file_prefix='latents_',
|
|
regenerate_cache=regenerate_cache,
|
|
caching_batch_size=caching_batch_size,
|
|
with_indices=True,
|
|
)
|
|
# Shuffle again, since one media file can produce multiple training examples. E.g. video, or maybe
|
|
# in the future data augmentation. Don't need to shuffle text embeddings since those are looked
|
|
# up by index.
|
|
self.latent_dataset = self.latent_dataset.shuffle(seed=123)
|
|
# TODO: should we do dataset.flatten_indices() to make it contiguous on disk again?
|
|
# self.latent_dataset = self.latent_dataset.flatten_indices(
|
|
# cache_file_name=str(self.cache_dir / 'latents_flattened.arrow')
|
|
# )
|
|
|
|
def add_text_embedding_dataset(self, te_dataset):
|
|
self.text_embedding_datasets.append(te_dataset)
|
|
|
|
def __getitem__(self, idx):
|
|
idx = idx % len(self.latent_dataset)
|
|
ret = self.latent_dataset[idx]
|
|
te_idx = ret['te_idx'].item()
|
|
if DEBUG:
|
|
print(Path(self.metadata_dataset[te_idx]['image_file']).stem)
|
|
for ds in self.text_embedding_datasets:
|
|
ret.update(ds[te_idx])
|
|
return ret
|
|
|
|
def __len__(self):
|
|
return len(self.latent_dataset) * self.num_repeats
|
|
|
|
|
|
# Logical concatenation of multiple SizeBucketDataset, for the same size bucket. It returns items
|
|
# as batches.
|
|
class ConcatenatedBatchedDataset:
|
|
def __init__(self, datasets):
|
|
self.datasets = datasets
|
|
self.post_init_called = False
|
|
|
|
def post_init(self, batch_size):
|
|
iteration_order = []
|
|
for i, ds in enumerate(self.datasets):
|
|
print(i, len(ds))
|
|
iteration_order.extend([i]*len(ds))
|
|
shuffle_with_seed(iteration_order, 0)
|
|
cumulative_sums = [0] * len(self.datasets)
|
|
for k, dataset_idx in enumerate(iteration_order):
|
|
iteration_order[k] = (dataset_idx, cumulative_sums[dataset_idx])
|
|
cumulative_sums[dataset_idx] += 1
|
|
self.iteration_order = iteration_order
|
|
assert len(self.iteration_order) > 0, 'ConcatenatedBatchedDataset is empty. Are your file paths correct?'
|
|
self.batch_size = batch_size
|
|
self._make_divisible_by(self.batch_size)
|
|
self.post_init_called = True
|
|
|
|
def __len__(self):
|
|
assert self.post_init_called
|
|
return len(self.iteration_order) // self.batch_size
|
|
|
|
def __getitem__(self, idx):
|
|
assert self.post_init_called
|
|
start = idx * self.batch_size
|
|
end = start + self.batch_size
|
|
return [self.datasets[i][j] for i, j in self.iteration_order[start:end]]
|
|
|
|
def _make_divisible_by(self, n):
|
|
new_length = (len(self.iteration_order) // n) * n
|
|
self.iteration_order = self.iteration_order[:new_length]
|
|
if new_length == 0 and is_main_process():
|
|
logger.warning(f"size bucket {self.datasets[0].size_bucket} is being completely dropped because it doesn't have enough images")
|
|
|
|
|
|
class ARBucketDataset:
|
|
def __init__(self, ar_frames, resolutions, metadata_dataset, directory_config, model_name):
|
|
self.ar_frames = ar_frames
|
|
self.resolutions = resolutions
|
|
self.metadata_dataset = metadata_dataset
|
|
self.directory_config = directory_config
|
|
self.model_name = model_name
|
|
self.size_buckets = []
|
|
self.path = Path(directory_config['path'])
|
|
self.cache_dir = self.path / 'cache' / self.model_name / f'ar_frames_{self.ar_frames[0]:.3f}_{self.ar_frames[1]}'
|
|
os.makedirs(self.cache_dir, exist_ok=True)
|
|
|
|
for res in resolutions:
|
|
area = res**2
|
|
w = math.sqrt(area * self.ar_frames[0])
|
|
h = area / w
|
|
w = round_to_multiple(w, IMAGE_SIZE_ROUND_TO_MULTIPLE)
|
|
h = round_to_multiple(h, IMAGE_SIZE_ROUND_TO_MULTIPLE)
|
|
size_bucket = (w, h, self.ar_frames[1])
|
|
metadata_with_size_bucket = self.metadata_dataset.map(lambda example: {'size_bucket': size_bucket}, keep_in_memory=True)
|
|
sub_dataset = SizeBucketDataset(metadata_with_size_bucket, directory_config, size_bucket, model_name)
|
|
print(size_bucket)
|
|
self.size_buckets.append(
|
|
sub_dataset
|
|
)
|
|
|
|
def get_size_bucket_datasets(self):
|
|
return self.size_buckets
|
|
|
|
def cache_latents(self, map_fn, regenerate_cache=False, caching_batch_size=1):
|
|
print(f'caching latents: {self.ar_frames}')
|
|
for ds in self.size_buckets:
|
|
ds.cache_latents(map_fn, regenerate_cache=regenerate_cache, caching_batch_size=caching_batch_size)
|
|
|
|
def cache_text_embeddings(self, map_fn, i, regenerate_cache=False, caching_batch_size=1):
|
|
print(f'caching text embeddings: {self.ar_frames}')
|
|
te_dataset = _map_and_cache(
|
|
self.metadata_dataset,
|
|
map_fn,
|
|
self.cache_dir,
|
|
cache_file_prefix=f'text_embeddings_{i}_',
|
|
new_fingerprint_args=[i],
|
|
regenerate_cache=regenerate_cache,
|
|
caching_batch_size=caching_batch_size,
|
|
)
|
|
for size_bucket_dataset in self.size_buckets:
|
|
size_bucket_dataset.add_text_embedding_dataset(te_dataset)
|
|
|
|
|
|
class DirectoryDataset:
|
|
def __init__(self, directory_config, dataset_config, model_name):
|
|
self._set_defaults(directory_config, dataset_config)
|
|
self.directory_config = directory_config
|
|
self.dataset_config = dataset_config
|
|
self.model_name = model_name
|
|
self.enable_ar_bucket = directory_config.get('enable_ar_bucket', dataset_config.get('enable_ar_bucket', False))
|
|
self.resolutions = self._process_user_provided_resolutions(
|
|
directory_config.get('resolutions', dataset_config['resolutions'])
|
|
)
|
|
self.path = Path(self.directory_config['path'])
|
|
self.cache_dir = self.path / 'cache' / self.model_name
|
|
|
|
if not self.path.exists() or not self.path.is_dir():
|
|
raise RuntimeError(f'Invalid path: {self.path}')
|
|
|
|
if not self.enable_ar_bucket:
|
|
self.ars = np.array([1.0])
|
|
elif ars := self.directory_config.get('ar_buckets', self.dataset_config.get('ar_buckets', None)):
|
|
self.ars = self._process_user_provided_ars(ars)
|
|
else:
|
|
min_ar = self.directory_config.get('min_ar', self.dataset_config['min_ar'])
|
|
max_ar = self.directory_config.get('max_ar', self.dataset_config['max_ar'])
|
|
num_ar_buckets = self.directory_config.get('num_ar_buckets', self.dataset_config['num_ar_buckets'])
|
|
self.ars = np.geomspace(min_ar, max_ar, num=num_ar_buckets)
|
|
frame_buckets = self.directory_config.get('frame_buckets', self.dataset_config.get('frame_buckets', [1]))
|
|
if 1 not in frame_buckets:
|
|
# always have an image bucket for convenience
|
|
frame_buckets.append(1)
|
|
frame_buckets.sort()
|
|
self.frame_buckets = np.array(frame_buckets)
|
|
|
|
def cache_metadata(self, regenerate_cache=False):
|
|
files = list(self.path.glob('*'))
|
|
# deterministic order
|
|
files.sort()
|
|
|
|
image_files = []
|
|
caption_files = []
|
|
for file in files:
|
|
if not file.is_file() or file.suffix == '.txt' or file.suffix == '.npz':
|
|
continue
|
|
image_file = file
|
|
caption_file = image_file.with_suffix('.txt')
|
|
if not os.path.exists(caption_file):
|
|
logger.warning(f'Image file {image_file} does not have corresponding caption file.')
|
|
caption_file = ''
|
|
image_files.append(str(image_file))
|
|
caption_files.append(str(caption_file))
|
|
assert len(image_files) > 0, f'Directory {self.path} had no images/videos!'
|
|
|
|
metadata_dataset = datasets.Dataset.from_dict({'image_file': image_files, 'caption_file': caption_files})
|
|
# Shuffle the data. Use a fixed seed, so the dataset is identical on all processes.
|
|
# Processes other than rank 0 will then load it from cache.
|
|
metadata_dataset = metadata_dataset.shuffle(seed=0)
|
|
metadata_map_fn = self._metadata_map_fn(self.ars, self.frame_buckets)
|
|
fingerprint = Hasher.hash([metadata_dataset._fingerprint, metadata_map_fn])
|
|
print('caching metadata')
|
|
metadata_dataset = metadata_dataset.map(
|
|
metadata_map_fn,
|
|
cache_file_name=str(self.cache_dir / f'metadata/metadata_{fingerprint}.arrow'),
|
|
load_from_cache_file=(not regenerate_cache),
|
|
batched=True,
|
|
batch_size=1,
|
|
num_proc=NUM_PROC,
|
|
remove_columns=metadata_dataset.column_names,
|
|
)
|
|
grouped_metadata = defaultdict(lambda: defaultdict(list))
|
|
for example in metadata_dataset:
|
|
ar_bucket = example['ar_bucket']
|
|
ar_bucket = (ar_bucket[0], int(ar_bucket[1]))
|
|
d = grouped_metadata[ar_bucket]
|
|
for k, v in example.items():
|
|
d[k].append(v)
|
|
self.ar_buckets = []
|
|
for ar_bucket, metadata in grouped_metadata.items():
|
|
metadata = datasets.Dataset.from_dict(metadata)
|
|
self.ar_buckets.append(
|
|
ARBucketDataset(
|
|
ar_bucket,
|
|
self.resolutions,
|
|
metadata,
|
|
self.directory_config,
|
|
self.model_name,
|
|
)
|
|
)
|
|
|
|
def _set_defaults(self, directory_config, dataset_config):
|
|
directory_config.setdefault('enable_ar_bucket', dataset_config.get('enable_ar_bucket', False))
|
|
directory_config.setdefault('resolutions', dataset_config['resolutions'])
|
|
directory_config.setdefault('shuffle_tags', dataset_config.get('shuffle_tags', False))
|
|
directory_config.setdefault('caption_prefix', dataset_config.get('caption_prefix', ''))
|
|
|
|
def _metadata_map_fn(self, ars, frame_buckets):
|
|
log_ars = np.log(ars)
|
|
def fn(example):
|
|
# batch size always 1
|
|
caption_file = example['caption_file'][0]
|
|
image_file = example['image_file'][0]
|
|
if not caption_file:
|
|
caption = ''
|
|
else:
|
|
with open(caption_file) as f:
|
|
caption = f.read().strip()
|
|
if self.directory_config['shuffle_tags']:
|
|
tags = [tag.strip() for tag in caption.split(',')]
|
|
random.shuffle(tags)
|
|
caption = ', '.join(tags)
|
|
caption = self.directory_config['caption_prefix'] + caption
|
|
empty_return = {'image_file': [], 'caption': [], 'ar_bucket': [], 'is_video': []}
|
|
|
|
image_file = Path(image_file)
|
|
try:
|
|
if image_file.suffix in VIDEO_EXTENSIONS:
|
|
# 100% accurate frame count, but much slower.
|
|
# frames = 0
|
|
# for frame in imageio.v3.imiter(image_file):
|
|
# frames += 1
|
|
# height, width = frame.shape[:2]
|
|
# TODO: this is an estimate of frame count. What happens if variable frame rate? Is
|
|
# it still close enough?
|
|
meta = imageio.v3.immeta(image_file)
|
|
height, width = meta['size']
|
|
frames = int(meta['fps'] * meta['duration'])
|
|
else:
|
|
pil_img = Image.open(image_file)
|
|
width, height = pil_img.size
|
|
frames = 1
|
|
except Exception:
|
|
logger.warning(f'Image file {image_file} could not be opened. Skipping.')
|
|
return empty_return
|
|
is_video = (frames > 1)
|
|
log_ar = np.log(width / height)
|
|
# Best AR bucket is the one with the smallest AR difference in log space.
|
|
i = np.argmin(np.abs(log_ar - log_ars))
|
|
# find closest frame bucket where the number of frames is greater than or equal to the bucket
|
|
diffs = frames - frame_buckets
|
|
positive_diffs = diffs[diffs >= 0]
|
|
if len(positive_diffs) == 0:
|
|
# video not long enough to find any valid frame bucket
|
|
print(f'video with frames={frames} is being skipped because it is too short')
|
|
return empty_return
|
|
j = np.argmin(positive_diffs)
|
|
if is_video and frame_buckets[j] == 1:
|
|
# don't let video be mapped to the image frame bucket
|
|
print(f'video with frames={frames} is being skipped because it is too short')
|
|
return empty_return
|
|
ar_bucket = (ars[i], frame_buckets[j])
|
|
|
|
return {'image_file': [str(image_file)], 'caption': [caption], 'ar_bucket': [ar_bucket], 'is_video': [is_video]}
|
|
return fn
|
|
|
|
def _process_user_provided_ars(self, ars):
|
|
ar_buckets = set()
|
|
for ar in ars:
|
|
if isinstance(ar, (tuple, list)):
|
|
assert len(ar) == 2
|
|
ar = round(ar[0] / ar[1], 6)
|
|
ar_buckets.add(ar)
|
|
ar_buckets = list(ar_buckets)
|
|
ar_buckets.sort()
|
|
return np.array(ar_buckets)
|
|
|
|
def _process_user_provided_resolutions(self, resolutions):
|
|
result = set()
|
|
for res in resolutions:
|
|
if isinstance(res, (tuple, list)):
|
|
assert len(res) == 2
|
|
res = round(math.sqrt(res[0] * res[1]), 6)
|
|
result.add(res)
|
|
result = list(result)
|
|
result.sort()
|
|
return result
|
|
|
|
def get_size_bucket_datasets(self):
|
|
result = []
|
|
for ar_bucket_dataset in self.ar_buckets:
|
|
result.extend(ar_bucket_dataset.get_size_bucket_datasets())
|
|
return result
|
|
|
|
def cache_latents(self, map_fn, regenerate_cache=False, caching_batch_size=1):
|
|
print(f'caching latents: {self.path}')
|
|
for ds in self.ar_buckets:
|
|
ds.cache_latents(map_fn, regenerate_cache=regenerate_cache, caching_batch_size=caching_batch_size)
|
|
|
|
def cache_text_embeddings(self, map_fn, i, regenerate_cache=False, caching_batch_size=1):
|
|
for ds in self.ar_buckets:
|
|
ds.cache_text_embeddings(map_fn, i, regenerate_cache=regenerate_cache, caching_batch_size=caching_batch_size)
|
|
|
|
|
|
# Outermost dataset object that the caller uses. Contains multiple ConcatenatedBatchedDataset. Responsible
|
|
# for returning the correct batch for the process's data parallel rank. Calls model.prepare_inputs so the
|
|
# returned tuple of tensors is whatever the model needs.
|
|
class Dataset:
|
|
def __init__(self, dataset_config, model_name):
|
|
super().__init__()
|
|
self.dataset_config = dataset_config
|
|
self.model_name = model_name
|
|
self.post_init_called = False
|
|
self.eval_quantile = None
|
|
|
|
self.directory_datasets = []
|
|
for directory_config in dataset_config['directory']:
|
|
directory_dataset = DirectoryDataset(directory_config, dataset_config, model_name)
|
|
self.directory_datasets.append(directory_dataset)
|
|
|
|
def post_init(self, data_parallel_rank, data_parallel_world_size, per_device_batch_size, gradient_accumulation_steps):
|
|
self.data_parallel_rank = data_parallel_rank
|
|
self.data_parallel_world_size = data_parallel_world_size
|
|
self.batch_size = per_device_batch_size * gradient_accumulation_steps
|
|
self.global_batch_size = self.data_parallel_world_size * self.batch_size
|
|
|
|
# group same size_bucket together
|
|
datasets_by_size_bucket = defaultdict(list)
|
|
for directory_dataset in self.directory_datasets:
|
|
for size_bucket_dataset in directory_dataset.get_size_bucket_datasets():
|
|
datasets_by_size_bucket[size_bucket_dataset.size_bucket].append(size_bucket_dataset)
|
|
self.buckets = []
|
|
for datasets in datasets_by_size_bucket.values():
|
|
print(len(datasets))
|
|
self.buckets.append(ConcatenatedBatchedDataset(datasets))
|
|
|
|
for bucket in self.buckets:
|
|
bucket.post_init(self.global_batch_size)
|
|
|
|
iteration_order = []
|
|
for i, bucket in enumerate(self.buckets):
|
|
iteration_order.extend([i]*(len(bucket)))
|
|
shuffle_with_seed(iteration_order, 0)
|
|
cumulative_sums = [0] * len(self.buckets)
|
|
for k, dataset_idx in enumerate(iteration_order):
|
|
iteration_order[k] = (dataset_idx, cumulative_sums[dataset_idx])
|
|
cumulative_sums[dataset_idx] += 1
|
|
self.iteration_order = iteration_order
|
|
if DEBUG:
|
|
print(f'Dataset iteration_order: {self.iteration_order}')
|
|
|
|
self.post_init_called = True
|
|
|
|
if subsample_ratio := self.dataset_config.get('subsample_ratio', None):
|
|
new_len = int(len(self) * subsample_ratio)
|
|
self.iteration_order = self.iteration_order[:new_len]
|
|
|
|
def set_eval_quantile(self, quantile):
|
|
self.eval_quantile = quantile
|
|
|
|
def __len__(self):
|
|
assert self.post_init_called
|
|
return len(self.iteration_order)
|
|
|
|
def __getitem__(self, idx):
|
|
assert self.post_init_called
|
|
i, j = self.iteration_order[idx]
|
|
examples = self.buckets[i][j]
|
|
start_idx = self.data_parallel_rank*self.batch_size
|
|
examples_for_this_dp_rank = examples[start_idx:start_idx+self.batch_size]
|
|
if DEBUG:
|
|
print((start_idx, start_idx+self.batch_size))
|
|
batch = self._collate(examples_for_this_dp_rank)
|
|
return batch
|
|
|
|
# collates a list of dictionaries of tensors into a single dictionary of batched tensors
|
|
def _collate(self, examples):
|
|
ret = {}
|
|
for key in examples[0].keys():
|
|
ret[key] = torch.stack([example[key] for example in examples])
|
|
return ret
|
|
|
|
def cache_metadata(self, regenerate_cache=False):
|
|
for ds in self.directory_datasets:
|
|
ds.cache_metadata(regenerate_cache=regenerate_cache)
|
|
|
|
def cache_latents(self, map_fn, regenerate_cache=False, caching_batch_size=1):
|
|
for ds in self.directory_datasets:
|
|
ds.cache_latents(map_fn, regenerate_cache=regenerate_cache, caching_batch_size=caching_batch_size)
|
|
|
|
def cache_text_embeddings(self, map_fn, i, regenerate_cache=False, caching_batch_size=1):
|
|
for ds in self.directory_datasets:
|
|
ds.cache_text_embeddings(map_fn, i, regenerate_cache=regenerate_cache, caching_batch_size=caching_batch_size)
|
|
|
|
|
|
def _cache_fn(datasets, queue, preprocess_media_file_fn, num_text_encoders, regenerate_cache, caching_batch_size):
|
|
# Dataset map() starts a bunch of processes. Make sure torch uses a limited number of threads
|
|
# to avoid CPU contention.
|
|
# TODO: if we ever change Datasets map to use spawn instead of fork, this might not work.
|
|
#torch.set_num_threads(os.cpu_count() // NUM_PROC)
|
|
# HF Datasets map can randomly hang if this is greater than one (???)
|
|
# See https://github.com/pytorch/pytorch/issues/10996
|
|
# Alternatively, we could try fixing this by using spawn instead of fork.
|
|
torch.set_num_threads(1)
|
|
|
|
for ds in datasets:
|
|
ds.cache_metadata(regenerate_cache=regenerate_cache)
|
|
|
|
def latents_map_fn(example, indices):
|
|
first_size_bucket = example['size_bucket'][0]
|
|
tensors = []
|
|
te_idx = []
|
|
for idx, path, size_bucket in zip(indices, example['image_file'], example['size_bucket']):
|
|
assert size_bucket == first_size_bucket
|
|
items = preprocess_media_file_fn(path, size_bucket)
|
|
tensors.extend(items)
|
|
te_idx.extend([idx] * len(items))
|
|
|
|
if len(tensors) == 0:
|
|
return {'latents': [], 'te_idx': []}
|
|
|
|
caching_batch_size = len(example['image_file'])
|
|
results = defaultdict(list)
|
|
for i in range(0, len(tensors), caching_batch_size):
|
|
batched = torch.stack(tensors[i:i+caching_batch_size])
|
|
parent_conn, child_conn = mp.Pipe(duplex=False)
|
|
queue.put((0, batched, child_conn))
|
|
result = parent_conn.recv() # dict
|
|
for k, v in result.items():
|
|
results[k].append(v)
|
|
# concatenate the list of tensors at each key into one batched tensor
|
|
for k, v in results.items():
|
|
results[k] = torch.cat(v)
|
|
results['te_idx'] = te_idx
|
|
return results
|
|
|
|
for ds in datasets:
|
|
ds.cache_latents(latents_map_fn, regenerate_cache=regenerate_cache, caching_batch_size=caching_batch_size)
|
|
|
|
for text_encoder_idx in range(num_text_encoders):
|
|
def text_embedding_map_fn(example):
|
|
parent_conn, child_conn = mp.Pipe(duplex=False)
|
|
queue.put((text_encoder_idx+1, example['caption'], example['is_video'], child_conn))
|
|
result = parent_conn.recv() # dict
|
|
return result
|
|
for ds in datasets:
|
|
ds.cache_text_embeddings(text_embedding_map_fn, text_encoder_idx+1, regenerate_cache=regenerate_cache, caching_batch_size=caching_batch_size)
|
|
|
|
# signal that we're done
|
|
queue.put(None)
|
|
|
|
|
|
# Helper class to make caching multiple datasets more efficient by moving
|
|
# models to GPU as few times as needed.
|
|
class DatasetManager:
|
|
def __init__(self, model, regenerate_cache=False, caching_batch_size=1):
|
|
self.model = model
|
|
self.vae = self.model.get_vae()
|
|
self.text_encoders = self.model.get_text_encoders()
|
|
self.submodels = [self.vae] + list(self.text_encoders)
|
|
self.call_vae_fn = self.model.get_call_vae_fn(self.vae)
|
|
self.call_text_encoder_fns = [self.model.get_call_text_encoder_fn(text_encoder) for text_encoder in self.text_encoders]
|
|
self.regenerate_cache = regenerate_cache
|
|
self.caching_batch_size = caching_batch_size
|
|
self.datasets = []
|
|
|
|
def register(self, dataset):
|
|
self.datasets.append(dataset)
|
|
|
|
# Some notes for myself:
|
|
# Use a manager queue, since that can be pickled and unpickled, and sent to other processes.
|
|
# IMPORTANT: we use multiprocess library (not Python multiprocessing!) just like HF Datasets does.
|
|
# After hours of debugging and looking up related issues, I have concluded multiprocessing is outright bugged
|
|
# for this use case. Something about making a manager queue and sending it to the caching process, and then
|
|
# further sending it to map() workers via the pickled map function, is broken. It gets through a lot of the caching,
|
|
# but eventually, inevitably, queue.put() will fail with BrokenPipeError. Switching from multiprocessing to multiprocess,
|
|
# which has basically the same API, and everything works perfectly. ¯\_(ツ)_/¯
|
|
def cache(self):
|
|
if is_main_process():
|
|
manager = mp.Manager()
|
|
queue = [manager.Queue()]
|
|
else:
|
|
queue = [None]
|
|
torch.distributed.broadcast_object_list(queue, src=0, group=dist.get_world_group())
|
|
queue = queue[0]
|
|
|
|
# start up a process to run through the dataset caching flow
|
|
if is_main_process():
|
|
process = mp.Process(
|
|
target=_cache_fn,
|
|
args=(
|
|
self.datasets,
|
|
queue,
|
|
self.model.get_preprocess_media_file_fn(),
|
|
len(self.text_encoders),
|
|
self.regenerate_cache,
|
|
self.caching_batch_size,
|
|
)
|
|
)
|
|
process.start()
|
|
|
|
# loop on the original processes (one per GPU) to handle tasks requiring GPU models (VAE, text encoders)
|
|
while True:
|
|
task = queue.get()
|
|
if task is None:
|
|
# Propagate None so all worker processes break out of this loop.
|
|
# This is safe because it's a FIFO queue. The first None always comes after all work items.
|
|
queue.put(None)
|
|
break
|
|
self._handle_task(task)
|
|
|
|
# Free memory in all unneeded submodels. This is easier than trying to delete every reference.
|
|
# TODO: check if this is actually freeing memory.
|
|
for model in self.submodels:
|
|
model.to('meta')
|
|
|
|
dist.barrier()
|
|
if is_main_process():
|
|
process.join()
|
|
|
|
# Now load all datasets from cache.
|
|
for ds in self.datasets:
|
|
ds.cache_metadata()
|
|
ds.cache_latents(None)
|
|
for i in range(1, len(self.text_encoders)+1):
|
|
ds.cache_text_embeddings(None, i)
|
|
|
|
@torch.no_grad()
|
|
def _handle_task(self, task):
|
|
id = task[0]
|
|
# moved needed submodel to cuda, and everything else to cpu
|
|
if next(self.submodels[id].parameters()).device.type != 'cuda':
|
|
for i, submodel in enumerate(self.submodels):
|
|
if i != id:
|
|
submodel.to('cpu')
|
|
self.submodels[id].to('cuda')
|
|
if id == 0:
|
|
tensor, pipe = task[1:]
|
|
results = self.call_vae_fn(tensor)
|
|
elif id > 0:
|
|
caption, is_video, pipe = task[1:]
|
|
results = self.call_text_encoder_fns[id-1](caption, is_video=is_video)
|
|
else:
|
|
raise RuntimeError()
|
|
# Need to move to CPU here. If we don't, we get this error:
|
|
# RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
|
|
# I think this is because HF Datasets uses the multiprocess library (different from Python multiprocessing!) so it will always use fork.
|
|
results = {k: v.to('cpu') for k, v in results.items()}
|
|
pipe.send(results)
|
|
|
|
|
|
def split_batch(batch, pieces):
|
|
example_tuple = batch
|
|
split_size = example_tuple[0].size(0) // pieces
|
|
split_examples = zip(*(torch.split(tensor, split_size) for tensor in example_tuple))
|
|
# Deepspeed works with a tuple of (features, labels), even if we don't provide a loss_fn to PipelineEngine,
|
|
# and instead compute the loss ourselves in the model. It's okay to just return None for the labels here.
|
|
return [(ex, None) for ex in split_examples]
|
|
|
|
|
|
# DataLoader that divides batches into microbatches for gradient accumulation steps when doing
|
|
# pipeline parallel training. Iterates indefinitely (deepspeed requirement). Keeps track of epoch.
|
|
# Updates epoch as soon as the final batch is returned (notably different from qlora-pipe).
|
|
class PipelineDataLoader:
|
|
def __init__(self, dataset, gradient_accumulation_steps, model, num_dataloader_workers=2):
|
|
self.model = model
|
|
self.dataset = dataset
|
|
self.gradient_accumulation_steps = gradient_accumulation_steps
|
|
self.num_dataloader_workers = num_dataloader_workers
|
|
self.iter_called = False
|
|
self.eval_quantile = None
|
|
self.epoch = 1
|
|
self.num_batches_pulled = 0
|
|
self.next_micro_batch = None
|
|
self.recreate_dataloader = False
|
|
# Be careful to only create the DataLoader some bounded number of times: https://github.com/pytorch/pytorch/issues/91252
|
|
self._create_dataloader()
|
|
self.data = self._pull_batches_from_dataloader()
|
|
|
|
def reset(self):
|
|
self.epoch = 1
|
|
self.num_batches_pulled = 0
|
|
self.next_micro_batch = None
|
|
self.data = self._pull_batches_from_dataloader()
|
|
|
|
def set_eval_quantile(self, quantile):
|
|
self.eval_quantile = quantile
|
|
|
|
def __iter__(self):
|
|
self.iter_called = True
|
|
return self
|
|
|
|
def __len__(self):
|
|
return len(self.dataset) * self.gradient_accumulation_steps
|
|
|
|
def __next__(self):
|
|
if self.next_micro_batch == None:
|
|
self.next_micro_batch = next(self.data)
|
|
ret = self.next_micro_batch
|
|
try:
|
|
self.next_micro_batch = next(self.data)
|
|
except StopIteration:
|
|
if self.recreate_dataloader:
|
|
self._create_dataloader()
|
|
self.recreate_dataloader = False
|
|
self.data = self._pull_batches_from_dataloader()
|
|
self.num_batches_pulled = 0
|
|
self.next_micro_batch = next(self.data)
|
|
self.epoch += 1
|
|
return ret
|
|
|
|
def _create_dataloader(self, skip_first_n_batches=None):
|
|
if skip_first_n_batches is not None:
|
|
sampler = SkipFirstNSampler(skip_first_n_batches, len(self.dataset))
|
|
else:
|
|
sampler = None
|
|
self.dataloader = torch.utils.data.DataLoader(
|
|
self.dataset,
|
|
pin_memory=True,
|
|
batch_size=None,
|
|
sampler=sampler,
|
|
num_workers=self.num_dataloader_workers,
|
|
persistent_workers=(self.num_dataloader_workers > 0),
|
|
)
|
|
|
|
def _pull_batches_from_dataloader(self):
|
|
for batch in self.dataloader:
|
|
batch = self.model.prepare_inputs(batch, timestep_quantile=self.eval_quantile)
|
|
self.num_batches_pulled += 1
|
|
for micro_batch in split_batch(batch, self.gradient_accumulation_steps):
|
|
yield micro_batch
|
|
|
|
# Only the first and last stages in the pipeline pull from the dataloader. Parts of the code need
|
|
# to know the epoch, so we synchronize the epoch so the processes that don't use the dataloader
|
|
# know the current epoch.
|
|
def sync_epoch(self):
|
|
process_group = dist.get_world_group()
|
|
result = [None] * dist.get_world_size(process_group)
|
|
torch.distributed.all_gather_object(result, self.epoch, group=process_group)
|
|
max_epoch = -1
|
|
for epoch in result:
|
|
max_epoch = max(epoch, max_epoch)
|
|
self.epoch = max_epoch
|
|
|
|
def state_dict(self):
|
|
return {
|
|
'epoch': self.epoch,
|
|
'num_batches_pulled': self.num_batches_pulled,
|
|
}
|
|
|
|
def load_state_dict(self, state_dict):
|
|
assert not self.iter_called
|
|
self.epoch = state_dict['epoch']
|
|
# -1 because by preloading the next micro_batch, it's always going to have one more batch
|
|
# pulled than the actual number of batches iterated by the caller.
|
|
self.num_batches_pulled = state_dict['num_batches_pulled'] - 1
|
|
self._create_dataloader(skip_first_n_batches=self.num_batches_pulled)
|
|
self.data = self._pull_batches_from_dataloader()
|
|
# Recreate the dataloader after the first pass so that it won't skip
|
|
# batches again (we only want it to skip batches the first time).
|
|
self.recreate_dataloader = True
|
|
|
|
|
|
class SkipFirstNSampler(torch.utils.data.Sampler):
|
|
def __init__(self, n, dataset_length):
|
|
super().__init__()
|
|
self.n = n
|
|
self.dataset_length = dataset_length
|
|
|
|
def __len__(self):
|
|
return self.dataset_length
|
|
|
|
def __iter__(self):
|
|
for i in range(self.n, self.dataset_length):
|
|
yield i
|
|
|
|
|
|
if __name__ == '__main__':
|
|
from utils import common
|
|
common.is_main_process = lambda: True
|
|
from contextlib import contextmanager
|
|
@contextmanager
|
|
def _zero_first():
|
|
yield
|
|
common.zero_first = _zero_first
|
|
|
|
from utils import dataset as dataset_util
|
|
dataset_util.DEBUG = True
|
|
|
|
from models import flux
|
|
model = flux.CustomFluxPipeline.from_pretrained('/data2/imagegen_models/FLUX.1-dev', torch_dtype=torch.bfloat16)
|
|
model.model_config = {'guidance': 1.0, 'dtype': torch.bfloat16}
|
|
|
|
import toml
|
|
dataset_manager = dataset_util.DatasetManager(model)
|
|
with open('/home/anon/code/diffusion-pipe-configs/datasets/tiny1.toml') as f:
|
|
dataset_config = toml.load(f)
|
|
train_data = dataset_util.Dataset(dataset_config, model)
|
|
dataset_manager.register(train_data)
|
|
dataset_manager.cache()
|
|
|
|
train_data.post_init(data_parallel_rank=0, data_parallel_world_size=1, per_device_batch_size=1, gradient_accumulation_steps=2)
|
|
print(f'Dataset length: {len(train_data)}')
|
|
|
|
for item in train_data:
|
|
pass
|