fix merge conflicts

invokeai/backend/model_manager/load/load_default.py

@@ -37,7 +37,7 @@ class ModelLoader(ModelLoaderBase):
         self._logger = logger
         self._ram_cache = ram_cache
         self._convert_cache = convert_cache
-        self._torch_dtype = torch_dtype(choose_torch_device(), app_config)
+        self._torch_dtype = torch_dtype(choose_torch_device())
 
     def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
         """

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

@@ -117,7 +117,7 @@ class ModelCacheBase(ABC, Generic[T]):
 
     @property
     @abstractmethod
-    def stats(self) -> CacheStats:
+    def stats(self) -> Optional[CacheStats]:
         """Return collected CacheStats object."""
         pass
 

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

@@ -270,12 +270,14 @@ class ModelCache(ModelCacheBase[AnyModel]):
         if torch.device(source_device).type == torch.device(target_device).type:
             return
 
-        # may raise an exception here if insufficient GPU VRAM
-        self._check_free_vram(target_device, cache_entry.size)
-
         start_model_to_time = time.time()
         snapshot_before = self._capture_memory_snapshot()
-        cache_entry.model.to(target_device)
+        try:
+            cache_entry.model.to(target_device)
+        except Exception as e:  # blow away cache entry
+            self._delete_cache_entry(cache_entry)
+            raise e
+
         snapshot_after = self._capture_memory_snapshot()
         end_model_to_time = time.time()
         self.logger.debug(
@@ -330,11 +332,11 @@ class ModelCache(ModelCacheBase[AnyModel]):
                     f" {in_ram_models}/{in_vram_models}({locked_in_vram_models})"
                 )
 
-    def make_room(self, model_size: int) -> None:
+    def make_room(self, size: int) -> None:
         """Make enough room in the cache to accommodate a new model of indicated size."""
         # calculate how much memory this model will require
         # multiplier = 2 if self.precision==torch.float32 else 1
-        bytes_needed = model_size
+        bytes_needed = size
         maximum_size = self.max_cache_size * GIG  # stored in GB, convert to bytes
         current_size = self.cache_size()
 
@@ -389,12 +391,11 @@ class ModelCache(ModelCacheBase[AnyModel]):
             # 1 from onnx runtime object
             if not cache_entry.locked and refs <= (3 if "onnx" in model_key else 2):
                 self.logger.debug(
-                    f"Removing {model_key} from RAM cache to free at least {(model_size/GIG):.2f} GB (-{(cache_entry.size/GIG):.2f} GB)"
+                    f"Removing {model_key} from RAM cache to free at least {(size/GIG):.2f} GB (-{(cache_entry.size/GIG):.2f} GB)"
                 )
                 current_size -= cache_entry.size
                 models_cleared += 1
-                del self._cache_stack[pos]
-                del self._cached_models[model_key]
+                self._delete_cache_entry(cache_entry)
                 del cache_entry
 
             else:
@@ -422,16 +423,6 @@ class ModelCache(ModelCacheBase[AnyModel]):
 
         self.logger.debug(f"After making room: cached_models={len(self._cached_models)}")
 
-    def _check_free_vram(self, target_device: torch.device, needed_size: int) -> None:
-        if target_device.type != "cuda":
-            return
-        vram_device = (  # mem_get_info() needs an indexed device
-            target_device if target_device.index is not None else torch.device(str(target_device), index=0)
-        )
-        free_mem, _ = torch.cuda.mem_get_info(torch.device(vram_device))
-        if needed_size > free_mem:
-            needed_gb = round(needed_size / GIG, 2)
-            free_gb = round(free_mem / GIG, 2)
-            raise torch.cuda.OutOfMemoryError(
-                f"Insufficient VRAM to load model, requested {needed_gb}GB but only had {free_gb}GB free"
-            )
+    def _delete_cache_entry(self, cache_entry: CacheRecord[AnyModel]) -> None:
+        self._cache_stack.remove(cache_entry.key)
+        del self._cached_models[cache_entry.key]

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

@@ -34,7 +34,6 @@ class ModelLocker(ModelLockerBase):
 
         # NOTE that the model has to have the to() method in order for this code to move it into GPU!
         self._cache_entry.lock()
-
         try:
             if self._cache.lazy_offloading:
                 self._cache.offload_unlocked_models(self._cache_entry.size)
@@ -51,6 +50,7 @@ class ModelLocker(ModelLockerBase):
         except Exception:
             self._cache_entry.unlock()
             raise
+
         return self.model
 
     def unlock(self) -> None: