Add skeleton of a PeftModel class.

invokeai/backend/peft/peft_model.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+from safetensors.torch import load_file
+
+from invokeai.backend.model_manager.config import BaseModelType
+
+
+class PeftModel:
+    """A class for loading and managing parameter-efficient fine-tuning models."""
+
+    def __init__(
+        self,
+        name: str,
+        state_dict: dict[str, torch.Tensor],
+    ):
+        self._name = name
+        self._state_dict = state_dict
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def calc_size(self) -> int:
+        model_size = 0
+        for tensor in self._state_dict.values():
+            model_size += tensor.nelement() * tensor.element_size()
+        return model_size
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        file_path: Union[str, Path],
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        base_model: Optional[BaseModelType] = None,
+    ):
+        device = device or torch.device("cpu")
+        dtype = dtype or torch.float32
+
+        file_path = Path(file_path)
+
+        # TODO(ryand): Implement a helper function for this. This logic is duplicated repeatedly.
+        if file_path.suffix == ".safetensors":
+            state_dict = load_file(file_path, device="cpu")
+        else:
+            state_dict = torch.load(file_path, map_location="cpu")
+
+        # TODO(ryand):
+        # - Detect state_dict format
+        # - Convert state_dict to diffusers format if necessary
+
+        # if base_model == BaseModelType.StableDiffusionXL:
+        #     state_dict = cls._convert_sdxl_keys_to_diffusers_format(state_dict)
+        return cls(name=file_path.stem, state_dict=state_dict)