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

other condtions -> text conditions

Browse Source
This commit is contained in:
Aryan
2024-12-23 13:48:58 +01:00
parent 363841ae3f
commit 86e1d1888a
2 changed files with 24 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/dataset.py
+4 -4
View File
@@ -271,9 +271,9 @@ class PrecomputedDataset(Dataset):
self.conditions_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME
self.latent_conditions = sorted(os.listdir(self.latents_path))
self.other_conditions = sorted(os.listdir(self.conditions_path))
self.text_conditions = sorted(os.listdir(self.conditions_path))
assert len(self.latent_conditions) == len(self.other_conditions), "Number of captions and videos do not match"
assert len(self.latent_conditions) == len(self.text_conditions), "Number of captions and videos do not match"
def __len__(self) -> int:
return len(self.latent_conditions)
@@ -281,9 +281,9 @@ class PrecomputedDataset(Dataset):
def __getitem__(self, index: int) -> Dict[str, Any]:
conditions = {}
latent_path = self.latents_path / self.latent_conditions[index]
condition_path = self.conditions_path / self.other_conditions[index]
condition_path = self.conditions_path / self.text_conditions[index]
conditions["latent_conditions"] = torch.load(latent_path, map_location="cpu", weights_only=True)
conditions["other_conditions"] = torch.load(condition_path, map_location="cpu", weights_only=True)
conditions["text_conditions"] = torch.load(condition_path, map_location="cpu", weights_only=True)
return conditions