storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-16 00:13:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[core] Fix loading of precomputed conditions and latents (#199)

Browse Source 
* support precomputation.

* fixes
This commit is contained in:
Sayak Paul
2025-01-09 12:43:03 +05:30
committed by GitHub
parent 6716ccf8c2
commit f311f16d98
2 changed files with 14 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/dataset.py
+8 -3
View File
@@ -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))
finetrainers/trainer.py
+6 -2
View File
@@ -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,