mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
from pathlib import Path
|
|
from typing import Optional, Union
|
|
|
|
import torch
|
|
|
|
from invokeai.backend.model_manager.config import BaseModelType
|
|
from invokeai.backend.peft.sdxl_format_utils import convert_sdxl_keys_to_diffusers_format
|
|
from invokeai.backend.util.serialization import load_state_dict
|
|
|
|
|
|
class PeftModel:
|
|
"""A class for loading and managing parameter-efficient fine-tuning models."""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
state_dict: dict[str, torch.Tensor],
|
|
network_alphas: dict[str, torch.Tensor],
|
|
):
|
|
self.name = name
|
|
self.state_dict = state_dict
|
|
self.network_alphas = network_alphas
|
|
|
|
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)
|
|
|
|
state_dict = load_state_dict(file_path, device=str(device))
|
|
if base_model == BaseModelType.StableDiffusionXL:
|
|
state_dict = convert_sdxl_keys_to_diffusers_format(state_dict)
|
|
|
|
# TODO(ryand): We shouldn't be using an unexported function from diffusers here. Consider opening an upstream PR
|
|
# to move this function to state_dict_utils.py.
|
|
# state_dict, network_alphas = _convert_kohya_lora_to_diffusers(state_dict)
|
|
return cls(name=file_path.stem, state_dict=state_dict, network_alphas=network_alphas)
|