Spaces:
Running on Zero
Running on Zero
File size: 2,518 Bytes
a61978a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | from diffsynth.core import UnifiedDataset
import os
class GeoNeXtDataset(UnifiedDataset):
def __init__(self, data_profile,base_path, metadata_path, *args, **kwargs):
self.profile = data_profile
self._bad_sample_count = 0
# Call the parent class, but disable metadata loading there.
super().__init__(metadata_path=None, *args, **kwargs)
self.load_from_cache = False
# Core behavior: let the profile load and transform dataset entries.
full_meta_path = metadata_path
# 1) Load normalized data entries via profile.
self.data = self.profile.load_and_transform(full_meta_path)
# 2) Register profile-defined operators.
self.special_operator_map = self.profile.get_operator_map()
self.data_file_keys = self.profile.get_data_keys()
# 3) Disable cache-read mode for online preprocessing.
print(f"[Dataset] Initialized with profile: {self.profile.__class__.__name__}")
print(f"[Dataset] Total samples: {len(self.data)}")
# Override parent metadata loader (unused in profile-driven mode).
def load_metadata(self, metadata_path):
pass
def __getitem__(self, data_id):
try:
if self.load_from_cache:
data = self.cached_data[data_id % len(self.cached_data)]
data = self.cached_data_operator(data)
else:
data = self.data[data_id % len(self.data)].copy()
for key in self.data_file_keys:
if key in data:
if key in self.special_operator_map:
data[key] = self.special_operator_map[key](data[key])
elif key in self.data_file_keys:
data[key] = self.main_data_operator(data[key])
if hasattr(self.profile, "postprocess_sample"):
data = self.profile.postprocess_sample(data)
return data
except Exception as exc:
# Robust fallback: skip corrupted samples instead of crashing DDP training.
self._bad_sample_count += 1
if self._bad_sample_count <= 20 or self._bad_sample_count % 100 == 0:
print(
f"[Dataset][Warn] Skip bad sample data_id={data_id} "
f"(count={self._bad_sample_count}): {exc}"
)
return None
# Backward-compatible alias.
FlexibleDataset = GeoNeXtDataset
|