support training for eligen and nexusgen

This commit is contained in:
mi804
2025-07-29 13:28:42 +08:00
parent 2861ec4d9f
commit 8ef91b3672
14 changed files with 218 additions and 17 deletions

View File

@@ -120,8 +120,13 @@ class ImageDataset(torch.utils.data.Dataset):
data = self.data[data_id % len(self.data)].copy()
for key in self.data_file_keys:
if key in data:
path = os.path.join(self.base_path, data[key])
data[key] = self.load_data(path)
if isinstance(data[key], list):
print(f"Loading multiple files for key '{key}'.")
path = [os.path.join(self.base_path, p) for p in data[key]]
data[key] = [self.load_data(p) for p in path]
else:
path = os.path.join(self.base_path, data[key])
data[key] = self.load_data(path)
if data[key] is None:
warnings.warn(f"cannot load file {data[key]}.")
return None