Optimize RAM to VRAM transfer

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

@@ -46,6 +46,7 @@ class CacheRecord(Generic[T]):
 
     key: str
     model: T
+    state_dict: Optional[Dict[str, torch.Tensor]]  # this is a copy that stays in CPU
     size: int
     loaded: bool = False
     _locks: int = 0

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

@@ -162,7 +162,8 @@ def put(
         if key in self._cached_models:
             return
         self.make_room(size)
-        cache_record = CacheRecord(key, model, size)
+        state_dict = model.state_dict() if hasattr(model, "state_dict") else None
+        cache_record = CacheRecord(key=key, model=model, state_dict=state_dict, size=size)
         self._cached_models[key] = cache_record
         self._cache_stack.append(key)
 
@@ -267,6 +268,13 @@ def move_model_to_device(self, cache_entry: CacheRecord[AnyModel], target_device
         start_model_to_time = time.time()
         snapshot_before = self._capture_memory_snapshot()
         try:
+            if target_device == self.storage_device:
+                cache_entry.model.load_state_dict(cache_entry.state_dict, assign=True)
+            else:
+                new_dict: Dict[str, torch.Tensor] = {}
+                for k, v in cache_entry.state_dict.items():
+                    new_dict[k] = v.to(torch.device(target_device), copy=True)
+                cache_entry.model.load_state_dict(new_dict, assign=True)
             cache_entry.model.to(target_device)
         except Exception as e:  # blow away cache entry
             self._delete_cache_entry(cache_entry)