From 6970adaff742ac89adb3d85c803689210dc030e2 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sat, 7 Jan 2023 15:05:39 +0100 Subject: Made aspect ratio bucketing configurable --- training/util.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) (limited to 'training/util.py') diff --git a/training/util.py b/training/util.py index 6f42228..2b7f71d 100644 --- a/training/util.py +++ b/training/util.py @@ -1,6 +1,7 @@ from pathlib import Path import json import copy +import itertools from typing import Iterable, Optional from contextlib import contextmanager @@ -71,13 +72,7 @@ class CheckpointerBase: file_path = samples_path.joinpath(pool, f"step_{step}.jpg") file_path.parent.mkdir(parents=True, exist_ok=True) - data_enum = enumerate(data) - - batches = [ - batch - for j, batch in data_enum - if j * data.batch_size < self.sample_batch_size * self.sample_batches - ] + batches = list(itertools.islice(itertools.cycle(data), self.sample_batch_size * self.sample_batches)) prompts = [ prompt for batch in batches -- cgit v1.2.3-54-g00ecf