MCPcopy Create free account
hub / github.com/LuChengTHU/dpm-solver / load_datapool

Function load_datapool

examples/stable-diffusion/scripts/train_searcher.py:28–59  ·  view source on GitHub ↗
(dpath)

Source from the content-addressed store, hash-verified

26 reorder_k).build()
27
28def load_datapool(dpath):
29
30
31 def load_single_file(saved_embeddings):
32 compressed = np.load(saved_embeddings)
33 database = {key: compressed[key] for key in compressed.files}
34 return database
35
36 def load_multi_files(data_archive):
37 database = {key: [] for key in data_archive[0].files}
38 for d in tqdm(data_archive, desc=f'Loading datapool from {len(data_archive)} individual files.'):
39 for key in d.files:
40 database[key].append(d[key])
41
42 return database
43
44 print(f'Load saved patch embedding from "{dpath}"')
45 file_content = glob.glob(os.path.join(dpath, '*.npz'))
46
47 if len(file_content) == 1:
48 data_pool = load_single_file(file_content[0])
49 elif len(file_content) > 1:
50 data = [np.load(f) for f in file_content]
51 prefetched_data = parallel_data_prefetch(load_multi_files, data,
52 n_proc=min(len(data), cpu_count()), target_data_type='dict')
53
54 data_pool = {key: np.concatenate([od[key] for od in prefetched_data], axis=1)[0] for key in prefetched_data[0].keys()}
55 else:
56 raise ValueError(f'No npz-files in specified path "{dpath}" is this directory existing?')
57
58 print(f'Finished loading of retrieval database of length {data_pool["embedding"].shape[0]}.')
59 return data_pool
60
61
62def train_searcher(opt,

Callers 1

train_searcherFunction · 0.85

Calls 2

parallel_data_prefetchFunction · 0.90
load_single_fileFunction · 0.85

Tested by

no test coverage detected