Function _cache_fn

utils/dataset.py:1047–1127

(datasets, queue, preprocess_media_file_fn, num_text_encoders, regenerate_cache, trust_cache, caching_batch_size)

def _cache_fn(datasets, queue, preprocess_media_file_fn, num_text_encoders, regenerate_cache, trust_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, trust_cache=trust_cache)

    pipes = {}

    def latents_map_fn(example, rank):
        is_edit_dataset = ('control_file' in example)
        first_size_bucket = example['size_bucket'][0]
        tensors_and_masks = []
        image_specs = []
        captions = []
        control_tensors_and_masks = []
        for i, (image_spec, mask_path, size_bucket, caption) in enumerate(
            zip(example['image_spec'], example['mask_file'], example['size_bucket'], example['caption'])
        ):
            assert size_bucket == first_size_bucket
            items = preprocess_media_file_fn(image_spec, mask_path, size_bucket)
            tensors_and_masks.extend(items)
            image_specs.extend([image_spec] * len(items))
            captions.extend([caption] * len(items))
            if is_edit_dataset:
                control_file = example['control_file'][i]
                control_items = preprocess_media_file_fn((None, control_file), None, size_bucket)
                assert len(control_items) == 1
                assert len(items) == 1
                control_tensors_and_masks.append(control_items[0])
            else:
                control_tensors_and_masks.append(None)

        if len(tensors_and_masks) == 0:
            assert not is_edit_dataset
            return {'latents': [], 'mask': [], 'image_spec': [], 'caption': []}

        caching_batch_size = len(example['image_spec'])
        results = defaultdict(list)
        for i in range(0, len(tensors_and_masks), caching_batch_size):
            tensor = torch.stack([t[0] for t in tensors_and_masks[i:i+caching_batch_size]])
            c_tensor = torch.stack([t[0] for t in control_tensors_and_masks[i:i+caching_batch_size]]) if is_edit_dataset else None
            if rank not in pipes:
                pipes[rank] = mp.Pipe(duplex=False)
            parent_conn, child_conn = pipes[rank]
            queue.put((0, tensor, c_tensor, 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)