fixup unit tests and remove debugging statements

Author: Lincoln Stein

invokeai/app/api/dependencies.py

@@ -4,7 +4,6 @@
 
 import torch
 
-import invokeai.backend.util.devices  # horrible hack
 from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
 from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
 from invokeai.app.services.shared.sqlite.sqlite_util import init_db

invokeai/app/invocations/compel.py

@@ -99,6 +99,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
                 textual_inversion_manager=ti_manager,
                 dtype_for_device_getter=TorchDevice.choose_torch_dtype,
                 truncate_long_prompts=False,
+                device=TorchDevice.choose_torch_device(),
             )
 
             conjunction = Compel.parse_prompt_string(self.prompt)
@@ -113,6 +114,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
         conditioning_data = ConditioningFieldData(conditionings=[BasicConditioningInfo(embeds=c)])
 
         conditioning_name = context.conditioning.save(conditioning_data)
+
         return ConditioningOutput(
             conditioning=ConditioningField(
                 conditioning_name=conditioning_name,

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -74,9 +74,9 @@ def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: st
             )
             self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
 
-    def reset_stats(self):
-        self._stats = {}
-        self._cache_stats = {}
+    def reset_stats(self, graph_execution_state_id: str):
+        self._stats.pop(graph_execution_state_id)
+        self._cache_stats.pop(graph_execution_state_id)
 
     def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
         graph_stats_summary = self._get_graph_summary(graph_execution_state_id)

invokeai/app/services/model_manager/model_manager_default.py

@@ -76,8 +76,6 @@ def build_model_manager(
 
         ram_cache = ModelCache(
             max_cache_size=app_config.ram,
-            max_vram_cache_size=app_config.vram,
-            lazy_offloading=app_config.lazy_offload,
             logger=logger,
         )
         convert_cache = ModelConvertCache(cache_path=app_config.convert_cache_path, max_size=app_config.convert_cache)

invokeai/app/services/session_processor/session_processor_default.py

@@ -1,7 +1,7 @@
 import traceback
 from contextlib import suppress
 from queue import Queue
-from threading import BoundedSemaphore, Thread, Lock
+from threading import BoundedSemaphore, Lock, Thread
 from threading import Event as ThreadEvent
 from typing import Optional, Set
 
@@ -61,7 +61,9 @@ def __init__(
         self._on_after_run_session_callbacks = on_after_run_session_callbacks or []
         self._process_lock = Lock()
 
-    def start(self, services: InvocationServices, cancel_event: ThreadEvent, profiler: Optional[Profiler] = None) -> None:
+    def start(
+        self, services: InvocationServices, cancel_event: ThreadEvent, profiler: Optional[Profiler] = None
+    ) -> None:
         self._services = services
         self._cancel_event = cancel_event
         self._profiler = profiler
@@ -214,7 +216,7 @@ def _on_after_run_session(self, queue_item: SessionQueueItem) -> None:
             # we don't care about that - suppress the error.
             with suppress(GESStatsNotFoundError):
                 self._services.performance_statistics.log_stats(queue_item.session.id)
-                self._services.performance_statistics.reset_stats()
+                self._services.performance_statistics.reset_stats(queue_item.session.id)
 
             for callback in self._on_after_run_session_callbacks:
                 callback(queue_item=queue_item)
@@ -384,7 +386,6 @@ def start(self, invoker: Invoker) -> None:
             )
             worker.start()
 
-
     def stop(self, *args, **kwargs) -> None:
         self._stop_event.set()
 
@@ -465,7 +466,7 @@ def _process(
                     # Run the graph
                     # self.session_runner.run(queue_item=self._queue_item)
 
-                except Exception as e:
+                except Exception:
                     # Wait for next polling interval or event to try again
                     poll_now_event.wait(self._polling_interval)
                     continue
@@ -494,7 +495,7 @@ def _process_next_session(self) -> None:
                 with self._invoker.services.model_manager.load.ram_cache.reserve_execution_device():
                     # Run the session on the reserved GPU
                     self.session_runner.run(queue_item=queue_item)
-            except Exception as e:
+            except Exception:
                 continue
             finally:
                 self._active_queue_items.remove(queue_item)

invokeai/app/services/session_queue/session_queue_common.py

@@ -239,6 +239,7 @@ def queue_item_dto_from_dict(cls, queue_item_dict: dict) -> "SessionQueueItemDTO
     def __hash__(self) -> int:
         return self.item_id
 
+
 class SessionQueueItemDTO(SessionQueueItemWithoutGraph):
     pass
 

invokeai/app/services/shared/invocation_context.py

@@ -325,7 +325,6 @@ def load(self, name: str) -> ConditioningFieldData:
         Returns:
             The loaded conditioning data.
         """
-
         return self._services.conditioning.load(name)
 
 

invokeai/backend/model_manager/load/model_cache/model_cache_base.py

@@ -43,26 +43,9 @@ def model(self) -> AnyModel:
 
 @dataclass
 class CacheRecord(Generic[T]):
-    """
-    Elements of the cache:
-
-    key: Unique key for each model, same as used in the models database.
-    model: Model in memory.
-    state_dict: A read-only copy of the model's state dict in RAM. It will be
-                used as a template for creating a copy in the VRAM.
-    size: Size of the model
-    loaded: True if the model's state dict is currently in VRAM
-
-    Before a model is executed, the state_dict template is copied into VRAM,
-    and then injected into the model. When the model is finished, the VRAM
-    copy of the state dict is deleted, and the RAM version is reinjected
-    into the model.
-    """
+    """Elements of the cache."""
 
     key: str
-    model: T
-    device: torch.device
-    state_dict: Optional[Dict[str, torch.Tensor]]
     size: int
     model: T
     loaded: bool = False
@@ -130,28 +113,12 @@ def get_execution_device(self) -> torch.device:
         """
         pass
 
-    @property
-    @abstractmethod
-    def lazy_offloading(self) -> bool:
-        """Return true if the cache is configured to lazily offload models in VRAM."""
-        pass
-
     @property
     @abstractmethod
     def max_cache_size(self) -> float:
         """Return true if the cache is configured to lazily offload models in VRAM."""
         pass
 
-    @abstractmethod
-    def offload_unlocked_models(self, size_required: int) -> None:
-        """Offload from VRAM any models not actively in use."""
-        pass
-
-    @abstractmethod
-    def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device: torch.device) -> None:
-        """Move model into the indicated device."""
-        pass
-
     @property
     @abstractmethod
     def stats(self) -> Optional[CacheStats]: