github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-04-23 03:00:31 -04:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into refactor/rename-get-logger

Browse Source
This commit is contained in:
Lincoln Stein
2023-08-20 16:08:32 -04:00
committed by GitHub
parent 3cb6d333f6 0f1b975d0e
commit 45d172d5a8
106 changed files with 473 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
invokeai/backend/model_management/__init__.py
View File

@@ -1,10 +1,10 @@
"""
Initialization file for invokeai.backend.model_management
"""
from .model_manager import ModelManager, ModelInfo, AddModelResult, SchedulerPredictionType
from .model_cache import ModelCache
from .lora import ModelPatcher, ONNXModelPatcher
from .models import (
from .model_manager import ModelManager, ModelInfo, AddModelResult, SchedulerPredictionType # noqa: F401
from .model_cache import ModelCache # noqa: F401
from .lora import ModelPatcher, ONNXModelPatcher # noqa: F401
from .models import ( # noqa: F401
BaseModelType,
ModelType,
SubModelType,
@@ -12,5 +12,4 @@ from .models import (
ModelNotFoundException,
DuplicateModelException,
)
from .model_merge import ModelMerger, MergeInterpolationMethod
from .lora import ModelPatcher
from .model_merge import ModelMerger, MergeInterpolationMethod # noqa: F401

69
invokeai/backend/model_management/convert_ckpt_to_diffusers.py
View File

@@ -20,11 +20,36 @@
import re
from contextlib import nullcontext
from io import BytesIO
from typing import Optional, Union
from pathlib import Path
from typing import Optional, Union
import requests
import torch
from diffusers.models import (
AutoencoderKL,
ControlNetModel,
PriorTransformer,
UNet2DConditionModel,
)
from diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel
from diffusers.pipelines.paint_by_example import PaintByExampleImageEncoder
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.schedulers import (
DDIMScheduler,
DDPMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
UnCLIPScheduler,
)
from diffusers.utils import is_accelerate_available, is_omegaconf_available
from diffusers.utils.import_utils import BACKENDS_MAPPING
from picklescan.scanner import scan_file_path
from transformers import (
AutoFeatureExtractor,
BertTokenizerFast,
@@ -37,35 +62,8 @@ from transformers import (
CLIPVisionModelWithProjection,
)
from diffusers.models import (
AutoencoderKL,
ControlNetModel,
PriorTransformer,
UNet2DConditionModel,
)
from diffusers.schedulers import (
DDIMScheduler,
DDPMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
HeunDiscreteScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
UnCLIPScheduler,
)
from diffusers.utils import is_accelerate_available, is_omegaconf_available, is_safetensors_available
from diffusers.utils.import_utils import BACKENDS_MAPPING
from diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel
from diffusers.pipelines.paint_by_example import PaintByExampleImageEncoder
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from invokeai.backend.util.logging import InvokeAILogger
from invokeai.app.services.config import InvokeAIAppConfig
from picklescan.scanner import scan_file_path
from invokeai.backend.util.logging import InvokeAILogger
from .models import BaseModelType, ModelVariantType
try:
@@ -1221,9 +1219,6 @@ def download_from_original_stable_diffusion_ckpt(
raise ValueError(BACKENDS_MAPPING["omegaconf"][1])
if from_safetensors:
if not is_safetensors_available():
raise ValueError(BACKENDS_MAPPING["safetensors"][1])
from safetensors.torch import load_file as safe_load
checkpoint = safe_load(checkpoint_path, device="cpu")
@@ -1662,9 +1657,6 @@ def download_controlnet_from_original_ckpt(
from omegaconf import OmegaConf
if from_safetensors:
if not is_safetensors_available():
raise ValueError(BACKENDS_MAPPING["safetensors"][1])
from safetensors import safe_open
checkpoint = {}
@@ -1741,7 +1733,7 @@ def convert_ckpt_to_diffusers(
pipe.save_pretrained(
dump_path,
safe_serialization=use_safetensors and is_safetensors_available(),
safe_serialization=use_safetensors,
)
@@ -1757,7 +1749,4 @@ def convert_controlnet_to_diffusers(
"""
pipe = download_controlnet_from_original_ckpt(checkpoint_path, **kwargs)
pipe.save_pretrained(
dump_path,
safe_serialization=is_safetensors_available(),
)
pipe.save_pretrained(dump_path, safe_serialization=True)

20
invokeai/backend/model_management/lora.py
View File

@@ -5,21 +5,16 @@ from contextlib import contextmanager
from typing import Optional, Dict, Tuple, Any, Union, List
from pathlib import Path
import torch
from safetensors.torch import load_file
from torch.utils.hooks import RemovableHandle
from diffusers.models import UNet2DConditionModel
from transformers import CLIPTextModel
from onnx import numpy_helper
from onnxruntime import OrtValue
import numpy as np
import torch
from compel.embeddings_provider import BaseTextualInversionManager
from diffusers.models import UNet2DConditionModel
from safetensors.torch import load_file
from transformers import CLIPTextModel, CLIPTokenizer
from .models.lora import LoRAModel
"""
loras = [
(lora_model1, 0.7),
@@ -52,7 +47,7 @@ class ModelPatcher:
module = module.get_submodule(submodule_name)
module_key += "." + submodule_name
submodule_name = key_parts.pop(0)
except:
except Exception:
submodule_name += "_" + key_parts.pop(0)
module = module.get_submodule(submodule_name)
@@ -312,7 +307,8 @@ class TextualInversionManager(BaseTextualInversionManager):
class ONNXModelPatcher:
from .models.base import IAIOnnxRuntimeModel, OnnxRuntimeModel
from .models.base import IAIOnnxRuntimeModel
from diffusers import OnnxRuntimeModel
@classmethod
@contextmanager
@@ -341,7 +337,7 @@ class ONNXModelPatcher:
def apply_lora(
cls,
model: IAIOnnxRuntimeModel,
loras: List[Tuple[LoraModel, float]],
loras: List[Tuple[LoRAModel, float]],
prefix: str,
):
from .models.base import IAIOnnxRuntimeModel

2
invokeai/backend/model_management/model_cache.py
View File

@@ -273,7 +273,7 @@ class ModelCache(object):
self.cache.logger.debug(f"Locking {self.key} in {self.cache.execution_device}")
self.cache._print_cuda_stats()
except:
except Exception:
self.cache_entry.unlock()
raise

8
invokeai/backend/model_management/model_manager.py
View File

@@ -419,12 +419,12 @@ class ModelManager(object):
base_model_str, model_type_str, model_name = model_key.split("/", 2)
try:
model_type = ModelType(model_type_str)
except:
except Exception:
raise Exception(f"Unknown model type: {model_type_str}")
try:
base_model = BaseModelType(base_model_str)
except:
except Exception:
raise Exception(f"Unknown base model: {base_model_str}")
return (model_name, base_model, model_type)
@@ -855,7 +855,7 @@ class ModelManager(object):
info.pop("config")
result = self.add_model(model_name, base_model, model_type, model_attributes=info, clobber=True)
except:
except Exception:
# something went wrong, so don't leave dangling diffusers model in directory or it will cause a duplicate model error!
rmtree(new_diffusers_path)
raise
@@ -1042,7 +1042,7 @@ class ModelManager(object):
# Patch in the SD VAE from core so that it is available for use by the UI
try:
self.heuristic_import({str(self.resolve_model_path("core/convert/sd-vae-ft-mse"))})
except:
except Exception:
pass
installer = ModelInstall(

6
invokeai/backend/model_management/model_probe.py
View File

@@ -217,9 +217,9 @@ class ModelProbe(object):
raise "The model {model_name} is potentially infected by malware. Aborting import."
###################################################3
# ##################################################3
# Checkpoint probing
###################################################3
# ##################################################3
class ProbeBase(object):
def get_base_type(self) -> BaseModelType:
pass
@@ -431,7 +431,7 @@ class PipelineFolderProbe(FolderProbeBase):
return ModelVariantType.Depth
elif in_channels == 4:
return ModelVariantType.Normal
except:
except Exception:
pass
return ModelVariantType.Normal

4
invokeai/backend/model_management/models/__init__.py
View File

@@ -2,7 +2,7 @@ import inspect
from enum import Enum
from pydantic import BaseModel
from typing import Literal, get_origin
from .base import (
from .base import ( # noqa: F401
BaseModelType,
ModelType,
SubModelType,
@@ -118,7 +118,7 @@ def get_model_config_enums():
fields = model_config.__annotations__
try:
field = fields["model_format"]
except:
except Exception:
raise Exception("format field not found")
# model_format: None

34
invokeai/backend/model_management/models/base.py
View File

@@ -3,27 +3,28 @@ import os
import sys
import typing
import inspect
from enum import Enum
import warnings
from abc import ABCMeta, abstractmethod
from contextlib import suppress
from enum import Enum
from pathlib import Path
from picklescan.scanner import scan_file_path
import torch
import numpy as np
import safetensors.torch
from pathlib import Path
from diffusers import DiffusionPipeline, ConfigMixin, OnnxRuntimeModel
from contextlib import suppress
from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Type, Literal, TypeVar, Generic, Callable, Any, Union
import onnx
import safetensors.torch
from diffusers import DiffusionPipeline, ConfigMixin
from onnx import numpy_helper
from onnxruntime import (
InferenceSession,
SessionOptions,
get_available_providers,
)
from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Type, Literal, TypeVar, Generic, Callable, Any, Union
from diffusers import logging as diffusers_logging
from transformers import logging as transformers_logging
class DuplicateModelException(Exception):
@@ -171,7 +172,7 @@ class ModelBase(metaclass=ABCMeta):
fields = value.__annotations__
try:
field = fields["model_format"]
except:
except Exception:
raise Exception(f"Invalid config definition - format field not found({cls.__qualname__})")
if isinstance(field, type) and issubclass(field, str) and issubclass(field, Enum):
@@ -244,7 +245,7 @@ class DiffusersModel(ModelBase):
try:
config_data = DiffusionPipeline.load_config(self.model_path)
# config_data = json.loads(os.path.join(self.model_path, "model_index.json"))
except:
except Exception:
raise Exception("Invalid diffusers model! (model_index.json not found or invalid)")
config_data.pop("_ignore_files", None)
@@ -343,7 +344,7 @@ def calc_model_size_by_fs(model_path: str, subfolder: Optional[str] = None, vari
with open(os.path.join(model_path, file), "r") as f:
index_data = json.loads(f.read())
return int(index_data["metadata"]["total_size"])
except:
except Exception:
pass
# calculate files size if there is no index file
@@ -440,7 +441,7 @@ def read_checkpoint_meta(path: Union[str, Path], scan: bool = False):
if str(path).endswith(".safetensors"):
try:
checkpoint = _fast_safetensors_reader(path)
except:
except Exception:
# TODO: create issue for support "meta"?
checkpoint = safetensors.torch.load_file(path, device="cpu")
else:
@@ -452,11 +453,6 @@ def read_checkpoint_meta(path: Union[str, Path], scan: bool = False):
return checkpoint
import warnings
from diffusers import logging as diffusers_logging
from transformers import logging as transformers_logging
class SilenceWarnings(object):
def __init__(self):
self.transformers_verbosity = transformers_logging.get_verbosity()
@@ -639,7 +635,7 @@ class IAIOnnxRuntimeModel:
raise Exception("You should call create_session before running model")
inputs = {k: np.array(v) for k, v in kwargs.items()}
output_names = self.session.get_outputs()
# output_names = self.session.get_outputs()
# for k in inputs:
# self.io_binding.bind_cpu_input(k, inputs[k])
# for name in output_names:

6
invokeai/backend/model_management/models/controlnet.py
View File

@@ -43,7 +43,7 @@ class ControlNetModel(ModelBase):
try:
config = EmptyConfigLoader.load_config(self.model_path, config_name="config.json")
# config = json.loads(os.path.join(self.model_path, "config.json"))
except:
except Exception:
raise Exception("Invalid controlnet model! (config.json not found or invalid)")
model_class_name = config.get("_class_name", None)
@@ -53,7 +53,7 @@ class ControlNetModel(ModelBase):
try:
self.model_class = self._hf_definition_to_type(["diffusers", model_class_name])
self.model_size = calc_model_size_by_fs(self.model_path)
except:
except Exception:
raise Exception("Invalid ControlNet model!")
def get_size(self, child_type: Optional[SubModelType] = None):
@@ -78,7 +78,7 @@ class ControlNetModel(ModelBase):
variant=variant,
)
break
except:
except Exception:
pass
if not model:
raise ModelNotFoundException()

2
invokeai/backend/model_management/models/stable_diffusion.py
View File

@@ -330,5 +330,5 @@ def _select_ckpt_config(version: BaseModelType, variant: ModelVariantType):
config_path = config_path.relative_to(app_config.root_path)
return str(config_path)
except:
except Exception:
return None

14
invokeai/backend/model_management/models/stable_diffusion_onnx.py
View File

@@ -1,25 +1,17 @@
import os
import json
from enum import Enum
from pydantic import Field
from pathlib import Path
from typing import Literal, Optional, Union
from typing import Literal
from diffusers import OnnxRuntimeModel
from .base import (
ModelBase,
ModelConfigBase,
BaseModelType,
ModelType,
SubModelType,
ModelVariantType,
DiffusersModel,
SchedulerPredictionType,
SilenceWarnings,
read_checkpoint_meta,
classproperty,
OnnxRuntimeModel,
IAIOnnxRuntimeModel,
)
from invokeai.app.services.config import InvokeAIAppConfig
class StableDiffusionOnnxModelFormat(str, Enum):

7
invokeai/backend/model_management/models/vae.py
View File

@@ -5,7 +5,6 @@ from typing import Optional
import safetensors
import torch
from diffusers.utils import is_safetensors_available
from omegaconf import OmegaConf
from invokeai.app.services.config import InvokeAIAppConfig
@@ -44,14 +43,14 @@ class VaeModel(ModelBase):
try:
config = EmptyConfigLoader.load_config(self.model_path, config_name="config.json")
# config = json.loads(os.path.join(self.model_path, "config.json"))
except:
except Exception:
raise Exception("Invalid vae model! (config.json not found or invalid)")
try:
vae_class_name = config.get("_class_name", "AutoencoderKL")
self.vae_class = self._hf_definition_to_type(["diffusers", vae_class_name])
self.model_size = calc_model_size_by_fs(self.model_path)
except:
except Exception:
raise Exception("Invalid vae model! (Unkown vae type)")
def get_size(self, child_type: Optional[SubModelType] = None):
@@ -175,5 +174,5 @@ def _convert_vae_ckpt_and_cache(
vae_config=config,
image_size=image_size,
)
vae_model.save_pretrained(output_path, safe_serialization=is_safetensors_available())
vae_model.save_pretrained(output_path, safe_serialization=True)
return output_path