diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py index 50d645eb57..0b03c8e729 100644 --- a/invokeai/app/api/routers/models.py +++ b/invokeai/app/api/routers/models.py @@ -1,13 +1,13 @@ # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) and 2023 Kent Keirsey (https://github.com/hipsterusername) -from typing import Annotated, Literal, Optional, Union, Dict +from typing import Literal, Optional, Union from fastapi import Query from fastapi.routing import APIRouter, HTTPException from pydantic import BaseModel, Field, parse_obj_as from ..dependencies import ApiDependencies from invokeai.backend import BaseModelType, ModelType -from invokeai.backend.model_management.models import OPENAPI_MODEL_CONFIGS +from invokeai.backend.model_management.models import OPENAPI_MODEL_CONFIGS, SchedulerPredictionType MODEL_CONFIGS = Union[tuple(OPENAPI_MODEL_CONFIGS)] models_router = APIRouter(prefix="/v1/models", tags=["models"]) @@ -51,11 +51,14 @@ class CreateModelResponse(BaseModel): info: Union[CkptModelInfo, DiffusersModelInfo] = Field(discriminator="format", description="The model info") status: str = Field(description="The status of the API response") +class ImportModelRequest(BaseModel): + name: str = Field(description="A model path, repo_id or URL to import") + prediction_type: Optional[Literal['epsilon','v_prediction','sample']] = Field(description='Prediction type for SDv2 checkpoint files') + class ConversionRequest(BaseModel): name: str = Field(description="The name of the new model") info: CkptModelInfo = Field(description="The converted model info") save_location: str = Field(description="The path to save the converted model weights") - class ConvertedModelResponse(BaseModel): name: str = Field(description="The name of the new model") @@ -105,6 +108,28 @@ async def update_model( return model_response +@models_router.post( + "/", + operation_id="import_model", + responses={200: {"status": "success"}}, +) +async def import_model( + model_request: ImportModelRequest +) -> None: + """ Add Model """ + items_to_import = set([model_request.name]) + prediction_types = { x.value: x for x in SchedulerPredictionType } + logger = ApiDependencies.invoker.services.logger + + installed_models = ApiDependencies.invoker.services.model_manager.heuristic_import( + items_to_import = items_to_import, + prediction_type_helper = lambda x: prediction_types.get(model_request.prediction_type) + ) + if len(installed_models) > 0: + logger.info(f'Successfully imported {model_request.name}') + else: + logger.error(f'Model {model_request.name} not imported') + raise HTTPException(status_code=500, detail=f'Model {model_request.name} not imported') @models_router.delete( "/{model_name}", diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py index 2955f72252..d37f67ba91 100644 --- a/invokeai/backend/install/model_install_backend.py +++ b/invokeai/backend/install/model_install_backend.py @@ -93,9 +93,10 @@ class ModelInstall(object): def __init__(self, config:InvokeAIAppConfig, prediction_type_helper: Callable[[Path],SchedulerPredictionType]=None, + model_manager: ModelManager = None, access_token:str = None): self.config = config - self.mgr = ModelManager(config.model_conf_path) + self.mgr = model_manager or ModelManager(config.model_conf_path) self.datasets = OmegaConf.load(Dataset_path) self.prediction_helper = prediction_type_helper self.access_token = access_token or HfFolder.get_token() diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index c7fb7569db..eafc283b1d 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -151,13 +151,11 @@ import os import hashlib import textwrap from dataclasses import dataclass -from packaging import version from pathlib import Path -from typing import Dict, Optional, List, Tuple, Union, types +from typing import Optional, List, Tuple, Union, Set, Callable, types from shutil import rmtree import torch -from huggingface_hub import scan_cache_dir from omegaconf import OmegaConf from omegaconf.dictconfig import DictConfig @@ -165,9 +163,13 @@ from pydantic import BaseModel import invokeai.backend.util.logging as logger from invokeai.app.services.config import InvokeAIAppConfig -from invokeai.backend.util import CUDA_DEVICE, download_with_resume +from invokeai.backend.util import CUDA_DEVICE from .model_cache import ModelCache, ModelLocker -from .models import BaseModelType, ModelType, SubModelType, ModelError, MODEL_CLASSES +from .models import ( + BaseModelType, ModelType, SubModelType, + ModelError, SchedulerPredictionType, MODEL_CLASSES, + ModelConfigBase, + ) # We are only starting to number the config file with release 3. # The config file version doesn't have to start at release version, but it will help @@ -686,3 +688,34 @@ class ModelManager(object): if new_models_found and self.config_path: self.commit() + + + def heuristic_import(self, + items_to_import: Set[str], + prediction_type_helper: Callable[[Path],SchedulerPredictionType]=None, + )->Set[str]: + ''' + Import a list of paths, repo_ids or URLs. Returns the + set of successfully imported items. The prediction_type_helper + is a callback that receives the Path of a checkpoint or diffusers + model and returns a SchedulerPredictionType (or None). + ''' + # avoid circular import here + from invokeai.backend.install.model_install_backend import ModelInstall + successfully_installed = set() + + installer = ModelInstall(config = self.globals, + prediction_type_helper = prediction_type_helper, + model_manager = self) + for thing in items_to_import: + try: + installer.heuristic_install(thing) + successfully_installed.add(thing) + except Exception as e: + self.logger.warning(f'{thing} could not be imported: {str(e)}') + + self.commit() + return successfully_installed + + + diff --git a/invokeai/backend/model_management/model_probe.py b/invokeai/backend/model_management/model_probe.py index ec3146a00e..9aeb829a25 100644 --- a/invokeai/backend/model_management/model_probe.py +++ b/invokeai/backend/model_management/model_probe.py @@ -1,17 +1,14 @@ import json -import traceback import torch import safetensors.torch from dataclasses import dataclass -from enum import Enum -from diffusers import ModelMixin, ConfigMixin, StableDiffusionPipeline, AutoencoderKL, ControlNetModel +from diffusers import ModelMixin, ConfigMixin from pathlib import Path from typing import Callable, Literal, Union, Dict from picklescan.scanner import scan_file_path -import invokeai.backend.util.logging as logger from .models import BaseModelType, ModelType, ModelVariantType, SchedulerPredictionType, SilenceWarnings @dataclass @@ -102,7 +99,7 @@ class ModelProbe(object): and prediction_type==SchedulerPredictionType.VPrediction \ ) else 512, ) - except Exception as e: + except Exception: return None return model_info @@ -115,6 +112,9 @@ class ModelProbe(object): return ModelType.TextualInversion checkpoint = checkpoint or cls._scan_and_load_checkpoint(model_path) state_dict = checkpoint.get("state_dict") or checkpoint + + if len(checkpoint) < 10 and all(isinstance(v, torch.Tensor) for v in checkpoint.values()): + return ModelType.TextualInversion if any([x.startswith("model.diffusion_model") for x in state_dict.keys()]): return ModelType.Pipeline if any([x.startswith("encoder.conv_in") for x in state_dict.keys()]): @@ -326,13 +326,9 @@ class PipelineFolderProbe(FolderProbeBase): def get_base_type(self)->BaseModelType: if self.model: unet_conf = self.model.unet.config - scheduler_conf = self.model.scheduler.config else: with open(self.folder_path / 'unet' / 'config.json','r') as file: unet_conf = json.load(file) - with open(self.folder_path / 'scheduler' / 'scheduler_config.json','r') as file: - scheduler_conf = json.load(file) - if unet_conf['cross_attention_dim'] == 768: return BaseModelType.StableDiffusion1 elif unet_conf['cross_attention_dim'] == 1024: diff --git a/invokeai/backend/model_management/models/base.py b/invokeai/backend/model_management/models/base.py index ef354ecc07..0b36cd4235 100644 --- a/invokeai/backend/model_management/models/base.py +++ b/invokeai/backend/model_management/models/base.py @@ -56,7 +56,6 @@ class ModelConfigBase(BaseModel): class Config: use_enum_values = True - class EmptyConfigLoader(ConfigMixin): @classmethod def load_config(cls, *args, **kwargs): diff --git a/invokeai/configs/INITIAL_MODELS.yaml b/invokeai/configs/INITIAL_MODELS.yaml index e2288c43fd..9ff9efc3f5 100644 --- a/invokeai/configs/INITIAL_MODELS.yaml +++ b/invokeai/configs/INITIAL_MODELS.yaml @@ -45,6 +45,7 @@ sd-1/pipeline/portraitplus: repo_id: wavymulder/portraitplus recommended: False sd-1/pipeline/seek.art_MEGA: + repo_id: coreco/seek.art_MEGA description: A general use SD-1.5 "anything" model that supports multiple styles (2.1 GB) recommended: False sd-1/pipeline/trinart_stable_diffusion_v2: