Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP - Begin to integrate SpandreImageToImageModel type into the model manager.

Browse Source
This commit is contained in:
Ryan Dick
2024-06-28 15:01:42 -04:00
parent e6abea7bc5
commit 59ce9cf41c
5 changed files with 123 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
invokeai/backend/model_manager/probe.py
View File

@ -10,6 +10,7 @@ from picklescan.scanner import scan_file_path
import invokeai.backend.util.logging as logger
from invokeai.app.util.misc import uuid_string
from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
from invokeai.backend.util.silence_warnings import SilenceWarnings
from .config import (
@ -240,6 +241,14 @@ class ModelProbe(object):
if len(ckpt) < 10 and all(isinstance(v, torch.Tensor) for v in ckpt.values()):
return ModelType.TextualInversion
# Check if the model can be loaded as a SpandrelImageToImageModel.
try:
_ = SpandrelImageToImageModel.load_from_state_dict(ckpt)
return ModelType.SpandrelImageToImage
except Exception:
# TODO(ryand): Catch a more specific exception type here if we can.
pass
raise InvalidModelConfigException(f"Unable to determine model type for {model_path}")
@classmethod
@ -570,6 +579,11 @@ class T2IAdapterCheckpointProbe(CheckpointProbeBase):
raise NotImplementedError()
class SpandrelImageToImageModelProbe(CheckpointProbeBase):
def get_base_type(self) -> BaseModelType:
raise NotImplementedError()
########################################################
# classes for probing folders
#######################################################