Add IPAdapter to model_management __init__.py

This commit is contained in:
Ryan Dick 2023-09-12 13:31:52 -04:00
parent 7ee55489bb
commit 0d823901ef

View File

@ -1,29 +1,31 @@
import inspect import inspect
from enum import Enum from enum import Enum
from pydantic import BaseModel
from typing import Literal, get_origin from typing import Literal, get_origin
from pydantic import BaseModel
from .base import ( # noqa: F401 from .base import ( # noqa: F401
BaseModelType, BaseModelType,
ModelType, DuplicateModelException,
SubModelType, InvalidModelException,
ModelBase, ModelBase,
ModelConfigBase, ModelConfigBase,
ModelError,
ModelNotFoundException,
ModelType,
ModelVariantType, ModelVariantType,
SchedulerPredictionType, SchedulerPredictionType,
ModelError,
SilenceWarnings, SilenceWarnings,
ModelNotFoundException, SubModelType,
InvalidModelException,
DuplicateModelException,
) )
from .stable_diffusion import StableDiffusion1Model, StableDiffusion2Model
from .sdxl import StableDiffusionXLModel
from .vae import VaeModel
from .lora import LoRAModel
from .controlnet import ControlNetModel # TODO: from .controlnet import ControlNetModel # TODO:
from .textual_inversion import TextualInversionModel from .ip_adapter import IPAdapterModel
from .lora import LoRAModel
from .sdxl import StableDiffusionXLModel
from .stable_diffusion import StableDiffusion1Model, StableDiffusion2Model
from .stable_diffusion_onnx import ONNXStableDiffusion1Model, ONNXStableDiffusion2Model from .stable_diffusion_onnx import ONNXStableDiffusion1Model, ONNXStableDiffusion2Model
from .textual_inversion import TextualInversionModel
from .vae import VaeModel
MODEL_CLASSES = { MODEL_CLASSES = {
BaseModelType.StableDiffusion1: { BaseModelType.StableDiffusion1: {
@ -33,6 +35,7 @@ MODEL_CLASSES = {
ModelType.Lora: LoRAModel, ModelType.Lora: LoRAModel,
ModelType.ControlNet: ControlNetModel, ModelType.ControlNet: ControlNetModel,
ModelType.TextualInversion: TextualInversionModel, ModelType.TextualInversion: TextualInversionModel,
ModelType.IPAdapter: IPAdapterModel,
}, },
BaseModelType.StableDiffusion2: { BaseModelType.StableDiffusion2: {
ModelType.ONNX: ONNXStableDiffusion2Model, ModelType.ONNX: ONNXStableDiffusion2Model,
@ -41,6 +44,7 @@ MODEL_CLASSES = {
ModelType.Lora: LoRAModel, ModelType.Lora: LoRAModel,
ModelType.ControlNet: ControlNetModel, ModelType.ControlNet: ControlNetModel,
ModelType.TextualInversion: TextualInversionModel, ModelType.TextualInversion: TextualInversionModel,
ModelType.IPAdapter: IPAdapterModel,
}, },
BaseModelType.StableDiffusionXL: { BaseModelType.StableDiffusionXL: {
ModelType.Main: StableDiffusionXLModel, ModelType.Main: StableDiffusionXLModel,
@ -50,6 +54,7 @@ MODEL_CLASSES = {
ModelType.ControlNet: ControlNetModel, ModelType.ControlNet: ControlNetModel,
ModelType.TextualInversion: TextualInversionModel, ModelType.TextualInversion: TextualInversionModel,
ModelType.ONNX: ONNXStableDiffusion2Model, ModelType.ONNX: ONNXStableDiffusion2Model,
ModelType.IPAdapter: IPAdapterModel,
}, },
BaseModelType.StableDiffusionXLRefiner: { BaseModelType.StableDiffusionXLRefiner: {
ModelType.Main: StableDiffusionXLModel, ModelType.Main: StableDiffusionXLModel,
@ -59,6 +64,7 @@ MODEL_CLASSES = {
ModelType.ControlNet: ControlNetModel, ModelType.ControlNet: ControlNetModel,
ModelType.TextualInversion: TextualInversionModel, ModelType.TextualInversion: TextualInversionModel,
ModelType.ONNX: ONNXStableDiffusion2Model, ModelType.ONNX: ONNXStableDiffusion2Model,
ModelType.IPAdapter: IPAdapterModel,
}, },
# BaseModelType.Kandinsky2_1: { # BaseModelType.Kandinsky2_1: {
# ModelType.Main: Kandinsky2_1Model, # ModelType.Main: Kandinsky2_1Model,