From f311f16d988d22fecca5532974b652316fa066d6 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 9 Jan 2025 12:43:03 +0530 Subject: [PATCH] [core] Fix loading of precomputed conditions and latents (#199) * support precomputation. * fixes --- finetrainers/dataset.py | 11 ++++++++--- finetrainers/trainer.py | 8 ++++++-- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/finetrainers/dataset.py b/finetrainers/dataset.py index 6054e49..19ebb69 100644 --- a/finetrainers/dataset.py +++ b/finetrainers/dataset.py @@ -353,13 +353,18 @@ class ImageOrVideoDatasetWithResizeAndRectangleCrop(ImageOrVideoDataset): class PrecomputedDataset(Dataset): - def __init__(self, data_root: str) -> None: + def __init__(self, data_root: str, model_name: str = None, cleaned_model_id: str = None) -> None: super().__init__() self.data_root = Path(data_root) - self.latents_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_LATENTS_DIR_NAME - self.conditions_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME + if model_name and cleaned_model_id: + precomputation_dir = self.data_root / f"{model_name}_{cleaned_model_id}_{PRECOMPUTED_DIR_NAME}" + self.latents_path = precomputation_dir / PRECOMPUTED_LATENTS_DIR_NAME + self.conditions_path = precomputation_dir / PRECOMPUTED_CONDITIONS_DIR_NAME + else: + self.latents_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_LATENTS_DIR_NAME + self.conditions_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME self.latent_conditions = sorted(os.listdir(self.latents_path)) self.text_conditions = sorted(os.listdir(self.conditions_path)) diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index a05a8f0..81c8264 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -225,7 +225,9 @@ class Trainer: if not should_precompute: logger.info("Precomputed conditions and latents found. Loading precomputed data.") self.dataloader = torch.utils.data.DataLoader( - PrecomputedDataset(self.args.data_root), + PrecomputedDataset( + data_root=self.args.data_root, model_name=self.args.model_name, cleaned_model_id=cleaned_model_id + ), batch_size=self.args.batch_size, shuffle=True, collate_fn=collate_fn, @@ -353,7 +355,9 @@ class Trainer: # Update dataloader to use precomputed conditions and latents self.dataloader = torch.utils.data.DataLoader( - PrecomputedDataset(self.args.data_root), + PrecomputedDataset( + data_root=self.args.data_root, model_name=self.args.model_name, cleaned_model_id=cleaned_model_id + ), batch_size=self.args.batch_size, shuffle=True, collate_fn=collate_fn,