File size: 15,414 Bytes
e0cba20 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 | """Caption processing when caching Text Encoder outputs (issue #2453).
The deterministic part of ``process_caption`` (prefix/suffix, secondary separator, replacements, first-line
truncation) must be applied to the captions passed to ``cache_batch_outputs``.
"""
import logging
import random
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
import pytest
import torch
from library.dataset import BaseDataset, BucketBatchIndex, ImageInfo
from library.strategy_base import TextEncoderOutputsCachingStrategy, TextEncodingStrategy, TokenizeStrategy
from library.subset import DEFAULT_MULTILINE_CAPTION_SEPARATOR
SEP = DEFAULT_MULTILINE_CAPTION_SEPARATOR
def make_subset(**overrides):
params = dict(
enable_multiline_caption=False,
multiline_caption_separator=SEP,
enable_wildcard=False,
caption_prefix=None,
caption_suffix=None,
caption_dropout_rate=0.0,
caption_dropout_every_n_epochs=0,
caption_tag_dropout_rate=0.0,
shuffle_caption=False,
token_warmup_min=1,
token_warmup_step=0,
keep_tokens=0,
keep_tokens_separator="",
caption_separator=",",
secondary_separator=None,
)
params.update(overrides)
return SimpleNamespace(**params)
def make_dataset():
ds = SimpleNamespace(
current_epoch=0, current_step=0, max_train_steps=100, replacements={}, has_multiple_caption_candidates=False
)
# process_caption_for_cache delegates to this method
ds.process_caption_candidates_for_cache = lambda subset, caption: BaseDataset.process_caption_candidates_for_cache(
ds, subset, caption
)
return ds
def for_cache(subset, caption, replacements=None):
ds = make_dataset()
if replacements:
ds.replacements = replacements
return BaseDataset.process_caption_for_cache(ds, subset, caption)
# ---------------------------------------------------------------------------
# process_caption_for_cache
# ---------------------------------------------------------------------------
def test_for_cache_applies_prefix_suffix_and_first_line():
subset = make_subset(caption_prefix="pre", caption_suffix="suf")
assert for_cache(subset, "line1\nline2") == "pre line1 suf"
def test_for_cache_applies_secondary_separator_and_replacements():
subset = make_subset(secondary_separator=";;;")
assert for_cache(subset, "a;;;b, girl", replacements={"girl": "woman"}) == "a,b, woman"
def test_for_cache_is_deterministic_and_ignores_random_options():
# dropout / shuffle / tag dropout are not applied (they are rejected by is_text_encoder_output_cacheable anyway)
subset = make_subset(caption_dropout_rate=1.0, shuffle_caption=True, caption_tag_dropout_rate=1.0, caption_prefix="pre")
assert for_cache(subset, "a, b, c") == "pre a, b, c"
def test_for_cache_multiline_single_candidate_is_normalized():
subset = make_subset(enable_multiline_caption=True, caption_suffix="suf")
assert for_cache(subset, " a \nb\n\n") == " a\nb suf"
def test_for_cache_wildcard_caption_is_kept_as_is():
subset = make_subset(enable_wildcard=True, caption_prefix="pre")
assert for_cache(subset, "a {x|y}\nb") == "pre a {x|y}\nb"
def test_for_cache_matches_process_caption_for_plain_captions():
subset = make_subset(caption_prefix="pre", caption_suffix="suf", secondary_separator=";;;")
ds = make_dataset()
caption = "a;;;b, c"
assert BaseDataset.process_caption_for_cache(ds, subset, caption) == BaseDataset.process_caption(ds, subset, caption)
# ---------------------------------------------------------------------------
# is_text_encoder_output_cacheable: wildcard warning
# ---------------------------------------------------------------------------
def test_wildcard_warns_but_is_cacheable(caplog):
ds = make_dataset()
ds.subsets = [make_subset(enable_wildcard=True)]
with caplog.at_level(logging.WARNING, logger="library.dataset"):
assert BaseDataset.is_text_encoder_output_cacheable(ds) is True
assert any("enable_wildcard" in r.getMessage() for r in caplog.records)
def test_no_wildcard_no_warning(caplog):
ds = make_dataset()
ds.subsets = [make_subset()]
with caplog.at_level(logging.WARNING, logger="library.dataset"):
assert BaseDataset.is_text_encoder_output_cacheable(ds) is True
assert not any("enable_wildcard" in r.getMessage() for r in caplog.records)
# ---------------------------------------------------------------------------
# new_cache_text_encoder_outputs passes processed captions to the caching strategy
# ---------------------------------------------------------------------------
class RecordingCachingStrategy(TextEncoderOutputsCachingStrategy):
def __init__(self, batch_size):
super().__init__(cache_to_disk=False, batch_size=batch_size, skip_disk_cache_validity_check=False)
self.calls = []
def cache_batch_outputs(self, tokenize_strategy, models, text_encoding_strategy, batch, captions):
self.calls.append(([info.image_key for info in batch], list(captions)))
def test_new_cache_text_encoder_outputs_uses_processed_captions():
ds = BaseDataset((64, 64), 1.0, False, False)
ds.batch_size = 2
subset_a = make_subset(caption_prefix="pre", secondary_separator=";;;")
subset_b = make_subset(enable_multiline_caption=True, caption_suffix="suf")
ds.subsets = [subset_a, subset_b]
infos = [
(ImageInfo("img0", 1, "a;;;b\nignored", False, "img0.png"), subset_a),
(ImageInfo("img1", 1, "c", False, "img1.png"), subset_a),
(ImageInfo("img2", 1, "x\ny\n", False, "img2.png"), subset_b),
]
for info, subset in infos:
ds.register_image(info, subset)
caching = RecordingCachingStrategy(batch_size=2)
accelerator = SimpleNamespace(num_processes=1, process_index=0)
with patch.object(TokenizeStrategy, "get_strategy", return_value=object()), patch.object(
TextEncodingStrategy, "get_strategy", return_value=object()
), patch.object(TextEncoderOutputsCachingStrategy, "get_strategy", return_value=caching):
BaseDataset.new_cache_text_encoder_outputs(ds, models=[], accelerator=accelerator)
assert caching.calls == [
(["img0", "img1"], ["pre a,b", "pre c"]),
(["img2"], ["x\ny suf"]),
]
# the raw captions are untouched
assert [info.caption for info, _ in infos] == ["a;;;b\nignored", "c", "x\ny\n"]
# ---------------------------------------------------------------------------
# multiple caption candidates: process_caption_candidates_for_cache
# ---------------------------------------------------------------------------
def candidates_for_cache(subset, caption):
return BaseDataset.process_caption_candidates_for_cache(make_dataset(), subset, caption)
def test_candidates_for_cache_multiline_returns_all_candidates_with_prefix_suffix():
subset = make_subset(enable_multiline_caption=True, caption_prefix="pre", caption_suffix="suf")
caption = f"a\nb\n{SEP} # comment\nc\n{SEP}\n\n"
assert candidates_for_cache(subset, caption) == ["pre a\nb suf", "pre c suf"]
def test_candidates_for_cache_single_candidate_modes():
assert candidates_for_cache(make_subset(), "a\nb") == ["a"]
assert candidates_for_cache(make_subset(enable_wildcard=True), "a {x|y}\nb") == ["a {x|y}\nb"]
assert candidates_for_cache(make_subset(enable_multiline_caption=True), "a\nb") == ["a\nb"]
# process_caption_for_cache is the first candidate
subset = make_subset(enable_multiline_caption=True)
assert for_cache(subset, f"a\n{SEP}\nb") == "a"
def test_cacheable_with_multiple_candidates_only_if_supported(caplog):
ds = make_dataset()
ds.subsets = [make_subset(enable_multiline_caption=True)]
ds.has_multiple_caption_candidates = True
with caplog.at_level(logging.INFO, logger="library.dataset"):
assert BaseDataset.is_text_encoder_output_cacheable(ds) is False
assert BaseDataset.is_text_encoder_output_cacheable(ds, supports_multiple_caption_candidates=True) is True
assert any("multiple candidates" in r.getMessage() and r.levelno == logging.ERROR for r in caplog.records)
assert any("all candidates" in r.getMessage() and r.levelno == logging.INFO for r in caplog.records)
# the other restrictions still apply
ds.subsets = [make_subset(enable_multiline_caption=True, shuffle_caption=True)]
assert BaseDataset.is_text_encoder_output_cacheable(ds, supports_multiple_caption_candidates=True) is False
# ---------------------------------------------------------------------------
# new_cache_text_encoder_outputs with multiple candidates
# ---------------------------------------------------------------------------
class RecordingMultiCachingStrategy(TextEncoderOutputsCachingStrategy):
def __init__(self, batch_size, cache_to_disk=False, valid_npz=()):
super().__init__(cache_to_disk=cache_to_disk, batch_size=batch_size, skip_disk_cache_validity_check=False)
self.calls = []
self.checked = []
self.valid_npz = set(valid_npz)
@property
def supports_multiple_caption_candidates(self):
return True
def get_outputs_npz_path(self, image_abs_path):
return image_abs_path + "_te.npz"
def is_disk_cached_outputs_expected(self, npz_path, captions=None):
self.checked.append((npz_path, captions))
return npz_path in self.valid_npz
def cache_batch_outputs(self, tokenize_strategy, models, text_encoding_strategy, batch, captions):
raise AssertionError("cache_batch_outputs must not be called for a strategy with multiple candidate support")
def cache_batch_outputs_multi(self, tokenize_strategy, models, text_encoding_strategy, batch, captions_list):
self.calls.append(([info.image_key for info in batch], [list(c) for c in captions_list]))
def make_multi_candidate_dataset():
ds = BaseDataset((64, 64), 1.0, False, False)
ds.batch_size = 3
subset = make_subset(enable_multiline_caption=True, caption_prefix="pre")
ds.subsets = [subset]
infos = [
ImageInfo("img0", 1, f"a\n{SEP}\nb\n{SEP}\nc", False, "img0.png"), # 3 candidates
ImageInfo("img1", 1, "d", False, "img1.png"), # 1 candidate
ImageInfo("img2", 1, f"e\n{SEP}\nf", False, "img2.png"), # 2 candidates
ImageInfo("img3", 1, "g", False, "img3.png"), # 1 candidate
]
for info in infos:
ds.register_image(info, subset)
ds.has_multiple_caption_candidates = True
return ds, infos
def run_cache(ds, caching):
accelerator = SimpleNamespace(num_processes=1, process_index=0)
with patch.object(TokenizeStrategy, "get_strategy", return_value=object()), patch.object(
TextEncodingStrategy, "get_strategy", return_value=object()
), patch.object(TextEncoderOutputsCachingStrategy, "get_strategy", return_value=caching):
BaseDataset.new_cache_text_encoder_outputs(ds, models=[], accelerator=accelerator)
def test_new_cache_text_encoder_outputs_batches_by_number_of_candidates():
ds, _ = make_multi_candidate_dataset()
caching = RecordingMultiCachingStrategy(batch_size=3)
run_cache(ds, caching)
# the batch is flushed when the number of candidates reaches batch_size (3): img0 alone, img1+img2, img3
assert caching.calls == [
(["img0"], [["pre a", "pre b", "pre c"]]),
(["img1", "img2"], [["pre d"], ["pre e", "pre f"]]),
(["img3"], [["pre g"]]),
]
def test_new_cache_text_encoder_outputs_checks_disk_cache_with_candidates():
ds, infos = make_multi_candidate_dataset()
caching = RecordingMultiCachingStrategy(batch_size=10, cache_to_disk=True, valid_npz={"img1.png_te.npz"})
run_cache(ds, caching)
assert caching.checked == [
("img0.png_te.npz", ["pre a", "pre b", "pre c"]),
("img1.png_te.npz", ["pre d"]),
("img2.png_te.npz", ["pre e", "pre f"]),
("img3.png_te.npz", ["pre g"]),
]
assert caching.calls == [(["img0", "img2", "img3"], [["pre a", "pre b", "pre c"], ["pre e", "pre f"], ["pre g"]])]
assert [info.text_encoder_outputs_npz for info in infos] == [f"img{i}.png_te.npz" for i in range(4)]
def test_new_cache_text_encoder_outputs_rejects_multiple_candidates_without_support():
ds, _ = make_multi_candidate_dataset()
caching = RecordingCachingStrategy(batch_size=3) # no multiple candidate support
with pytest.raises(ValueError, match="multiple candidates"):
run_cache(ds, caching)
assert caching.calls == []
# ---------------------------------------------------------------------------
# __getitem__ chooses one candidate and returns the matching caption
# ---------------------------------------------------------------------------
class IndexRecordingStrategy(RecordingMultiCachingStrategy):
def __init__(self):
super().__init__(batch_size=1, cache_to_disk=True)
self.loaded = []
def load_outputs_npz(self, npz_path, candidate_index=0):
self.loaded.append((npz_path, candidate_index))
# outputs of candidate k: embeds filled with k, mask of ones, rate
return [
np.full((2, 4), float(candidate_index), dtype=np.float32),
np.ones(2, dtype=np.int32),
np.array(0.0, dtype=np.float32),
]
def make_getitem_dataset(in_memory: bool):
ds = BaseDataset((64, 64), 1.0, False, False)
ds.batch_size = 2
ds.prior_loss_weight = 1.0
subset = make_subset(enable_multiline_caption=True, caption_suffix="suf")
subset.flip_aug = False
subset.custom_attributes = {}
ds.subsets = [subset]
infos = [
ImageInfo("img0", 1, f"a\n{SEP}\nb\n{SEP}\nc", False, "img0.png"),
ImageInfo("img1", 1, "d", False, "img1.png"),
]
strategy = IndexRecordingStrategy()
for info in infos:
ds.register_image(info, subset)
info.latents = torch.zeros(4, 8, 8)
info.latents_original_size = (64, 64)
info.latents_crop_ltrb = (0, 0, 0, 0)
info.alpha_mask = None
info.bucket_reso = (64, 64)
if in_memory:
num = len(BaseDataset.process_caption_candidates_for_cache(ds, subset, info.caption))
info.text_encoder_outputs = [strategy.load_outputs_npz("mem", k) for k in range(num)]
else:
info.text_encoder_outputs_npz = strategy.get_outputs_npz_path(info.absolute_path)
strategy.loaded.clear()
ds.text_encoder_output_caching_strategy = strategy
ds.bucket_manager = SimpleNamespace(buckets=[["img0", "img1"]])
ds.buckets_indices = [BucketBatchIndex(0, 2, 0)]
return ds, strategy
@pytest.mark.parametrize("in_memory", [False, True])
def test_getitem_chooses_candidate_and_matching_caption(in_memory):
ds, strategy = make_getitem_dataset(in_memory)
chosen = set()
for seed in range(20):
random.seed(seed)
example = ds[0]
embeds = example["text_encoder_outputs_list"][0] # [B, L, C]
k0 = int(embeds[0, 0, 0].item())
chosen.add(k0)
assert example["captions"] == [["a suf", "b suf", "c suf"][k0], "d suf"]
assert int(embeds[1, 0, 0].item()) == 0
assert example["input_ids_list"] is None
if not in_memory:
assert strategy.loaded[-2:] == [("img0.png_te.npz", k0), ("img1.png_te.npz", 0)]
assert chosen == {0, 1, 2}
|