Massive speed increases and ram optimizations

toolkit/data_loader.py

@@ -448,7 +448,7 @@ class AiToolkitDataset(LatentCachingMixin, BucketsMixin, CaptionMixin, Dataset):
         return len(self.file_list)
 
     def _get_single_item(self, index) -> 'FileItemDTO':
-        file_item = self.file_list[index]
+        file_item = copy.deepcopy(self.file_list[index])
         file_item.load_and_process_image(self.transform)
         file_item.load_caption(self.caption_dict)
         return file_item
@@ -529,14 +529,14 @@ def get_dataloader_from_datasets(
             drop_last=False,
             shuffle=True,
             collate_fn=dto_collation,  # Use the custom collate function
-            num_workers=0
+            num_workers=4
         )
     else:
         data_loader = DataLoader(
             concatenated_dataset,
             batch_size=batch_size,
             shuffle=True,
-            num_workers=0,
+            num_workers=4,
             collate_fn=dto_collation
         )
     return data_loader

toolkit/dataloader_mixins.py

@@ -670,7 +670,8 @@ class LatentCachingFileItemDTOMixin:
             # load it from disk
             state_dict = load_file(
                 self.get_latent_path(),
-                device=device if device is not None else self.latent_load_device
+                # device=device if device is not None else self.latent_load_device
+                device='cpu'
             )
             self._encoded_latent = state_dict['latent']
         return self._encoded_latent

toolkit/network_mixins.py

@@ -53,7 +53,7 @@ class ToolkitModuleMixin:
             super().__init__(*args, **kwargs)
         self.network_ref: weakref.ref = weakref.ref(network)
         self.is_checkpointing = False
-        self.is_normalizing = False
+        # self.is_normalizing = False
         self.normalize_scaler = 1.0
         self._multiplier: Union[float, list, torch.Tensor] = None
 
@@ -118,7 +118,7 @@ class ToolkitModuleMixin:
             multiplier = multiplier.repeat_interleave(num_interleaves)
         # multiplier = 1.0
 
-        if self.is_normalizing:
+        if self.network_ref().is_normalizing:
             with torch.no_grad():
 
                 # do this calculation without set multiplier and instead use same polarity, but with 1.0 multiplier
@@ -390,8 +390,8 @@ class ToolkitNetworkMixin:
     @is_normalizing.setter
     def is_normalizing(self: Network, value: bool):
         self._is_normalizing = value
-        for module in self.get_all_modules():
-            module.is_normalizing = self._is_normalizing
+        # for module in self.get_all_modules():
+        #     module.is_normalizing = self._is_normalizing
 
     def apply_stored_normalizer(self: Network, target_normalize_scaler: float = 1.0):
         for module in self.get_all_modules():

toolkit/timer.py

@@ -0,0 +1,65 @@
+import time
+from collections import OrderedDict, deque
+
+
+class Timer:
+    def __init__(self, name='Timer', max_buffer=10):
+        self.name = name
+        self.max_buffer = max_buffer
+        self.timers = OrderedDict()
+        self.active_timers = {}
+        self.current_timer = None  # Used for the context manager functionality
+
+    def start(self, timer_name):
+        if timer_name not in self.timers:
+            self.timers[timer_name] = deque(maxlen=self.max_buffer)
+        self.active_timers[timer_name] = time.time()
+
+    def cancel(self, timer_name):
+        """Cancel an active timer."""
+        if timer_name in self.active_timers:
+            del self.active_timers[timer_name]
+
+    def stop(self, timer_name):
+        if timer_name not in self.active_timers:
+            raise ValueError(f"Timer '{timer_name}' was not started!")
+
+        elapsed_time = time.time() - self.active_timers[timer_name]
+        self.timers[timer_name].append(elapsed_time)
+
+        # Clean up active timers
+        del self.active_timers[timer_name]
+
+        # Check if this timer's buffer exceeds max_buffer and remove the oldest if it does
+        if len(self.timers[timer_name]) > self.max_buffer:
+            self.timers[timer_name].popleft()
+
+    def print(self):
+        print(f"\nTimer '{self.name}':")
+        # sort by longest at top
+        for timer_name, timings in sorted(self.timers.items(), key=lambda x: sum(x[1]), reverse=True):
+            avg_time = sum(timings) / len(timings)
+            print(f" - {avg_time:.4f}s avg - {timer_name}, num = {len(timings)}")
+
+        print('')
+
+    def reset(self):
+        self.timers.clear()
+        self.active_timers.clear()
+
+    def __call__(self, timer_name):
+        """Enable the use of the Timer class as a context manager."""
+        self.current_timer = timer_name
+        self.start(timer_name)
+        return self
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            # No exceptions, stop the timer normally
+            self.stop(self.current_timer)
+        else:
+            # There was an exception, cancel the timer
+            self.cancel(self.current_timer)