fixes

Author: ablattmann

ldm/util.py

@@ -2,6 +2,13 @@
 
 import torch
 import numpy as np
+from collections import abc
+from einops import rearrange
+from functools import partial
+
+import multiprocessing as mp
+from threading import Thread
+from queue import Queue
 
 from inspect import isfunction
 from PIL import Image, ImageDraw, ImageFont

scripts/knn2img.py

@@ -72,7 +72,7 @@ def __init__(self, database, retriever_version='ViT-L/14'):
                           'img_id': [],
                           'patch_coords': []}
         self.load_database()
-        self.searcher = self.load_searcher()
+        self.load_searcher()
 
     def train_searcher(self, k,
                        metric='dot_product',
@@ -117,7 +117,7 @@ def load_database(self):
             prefetched_data = parallel_data_prefetch(self.load_multi_files, data,
                                                      n_proc=min(len(data), cpu_count()), target_data_type='dict')
 
-            self.database = {key: np.concatenate([od[key] for od in prefetched_data], axis=1)[0] for key in self.data_pool}
+            self.database = {key: np.concatenate([od[key] for od in prefetched_data], axis=1)[0] for key in self.database}
         else:
             raise ValueError(f'No npz-files in specified path "{self.database_path}" is this directory existing?')