From 1fcd91bcc5c206c145d0e9a17eef675c9bdb2df1 Mon Sep 17 00:00:00 2001 From: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com> Date: Thu, 6 Apr 2023 15:17:48 -0400 Subject: [PATCH 01/14] Add/Update and Delete Models --- invokeai/app/api/routers/models.py | 302 ++++++++--------------------- 1 file changed, 79 insertions(+), 223 deletions(-) diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py index 5b3fbebddd..07ecb45003 100644 --- a/invokeai/app/api/routers/models.py +++ b/invokeai/app/api/routers/models.py @@ -2,9 +2,9 @@ from typing import Annotated, Any, List, Literal, Optional, Union -from fastapi.routing import APIRouter +from fastapi.routing import APIRouter, HTTPException from pydantic import BaseModel, Field, parse_obj_as - +from pathlib import Path from ..dependencies import ApiDependencies models_router = APIRouter(prefix="/v1/models", tags=["models"]) @@ -15,11 +15,9 @@ class VaeRepo(BaseModel): path: Optional[str] = Field(description="The path to the VAE") subfolder: Optional[str] = Field(description="The subfolder to use for this VAE") - class ModelInfo(BaseModel): description: Optional[str] = Field(description="A description of the model") - class CkptModelInfo(ModelInfo): format: Literal['ckpt'] = 'ckpt' @@ -29,7 +27,6 @@ class CkptModelInfo(ModelInfo): width: Optional[int] = Field(description="The width of the model") height: Optional[int] = Field(description="The height of the model") - class DiffusersModelInfo(ModelInfo): format: Literal['diffusers'] = 'diffusers' @@ -37,12 +34,27 @@ class DiffusersModelInfo(ModelInfo): repo_id: Optional[str] = Field(description="The repo ID to use for this model") path: Optional[str] = Field(description="The path to the model") +class CreateModelRequest (BaseModel): + name: str = Field(description="The name of the model") + info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model info") + +class CreateModelResponse (BaseModel): + name: str = Field(description="The name of the new model") + info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model info") + status: str = Field(description="The status of the API response") + +class ConvertedModelRequest (BaseModel): + name: str = Field(description="The name of the new model") + info: CkptModelInfo = Field(description="The converted model info") + +class ConvertedModelResponse (BaseModel): + name: str = Field(description="The name of the new model") + info: DiffusersModelInfo = Field(description="The converted model info") class ModelsList(BaseModel): models: dict[str, Annotated[Union[(CkptModelInfo,DiffusersModelInfo)], Field(discriminator="format")]] - @models_router.get( "/", operation_id="list_models", @@ -54,226 +66,70 @@ async def list_models() -> ModelsList: models = parse_obj_as(ModelsList, { "models": models_raw }) return models - # @socketio.on("requestSystemConfig") - # def handle_request_capabilities(): - # print(">> System config requested") - # config = self.get_system_config() - # config["model_list"] = self.generate.model_manager.list_models() - # config["infill_methods"] = infill_methods() - # socketio.emit("systemConfig", config) - # @socketio.on("searchForModels") - # def handle_search_models(search_folder: str): - # try: - # if not search_folder: - # socketio.emit( - # "foundModels", - # {"search_folder": None, "found_models": None}, - # ) - # else: - # ( - # search_folder, - # found_models, - # ) = self.generate.model_manager.search_models(search_folder) - # socketio.emit( - # "foundModels", - # {"search_folder": search_folder, "found_models": found_models}, - # ) - # except Exception as e: - # self.handle_exceptions(e) - # print("\n") +@models_router.post( + "/", + operation_id="update_model", + responses={ + 201: { + "model_response": "Model added", + }, + 202: { + "description": "Model submission is processing. Check back later." + }, + }, +) +async def update_model( + model_request: CreateModelRequest +) -> CreateModelResponse: + """ Add Model """ + try: + model_request_info = model_request.info + print(f">> Checking for {model_request_info}...") + info_dict = model_request_info.dict() - # @socketio.on("addNewModel") - # def handle_add_model(new_model_config: dict): - # try: - # model_name = new_model_config["name"] - # del new_model_config["name"] - # model_attributes = new_model_config - # if len(model_attributes["vae"]) == 0: - # del model_attributes["vae"] - # update = False - # current_model_list = self.generate.model_manager.list_models() - # if model_name in current_model_list: - # update = True + ApiDependencies.invoker.services.model_manager.add_model( + model_name=model_request.name, + model_attributes=info_dict, + clobber=True, + ) + model_response = CreateModelResponse(name=model_request.name, info=model_request.info, status="success") - # print(f">> Adding New Model: {model_name}") + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + return model_response - # self.generate.model_manager.add_model( - # model_name=model_name, - # model_attributes=model_attributes, - # clobber=True, - # ) - # self.generate.model_manager.commit(opt.conf) - # new_model_list = self.generate.model_manager.list_models() - # socketio.emit( - # "newModelAdded", - # { - # "new_model_name": model_name, - # "model_list": new_model_list, - # "update": update, - # }, - # ) - # print(f">> New Model Added: {model_name}") - # except Exception as e: - # self.handle_exceptions(e) +@models_router.delete( + "/{model_name}", + operation_id="del_model", + responses={ + 204: { + "description": "Model deleted" + }, + 404: { + "description": "Model not found" + } + }, +) +async def delete_model(model_name: str) -> None: + """Delete Model""" + model_names = ApiDependencies.invoker.services.model_manager.model_names() + model_exists = model_name in model_names + + try: + # check if model exists + print(f">> Checking for model {model_name}...") - # @socketio.on("deleteModel") - # def handle_delete_model(model_name: str): - # try: - # print(f">> Deleting Model: {model_name}") - # self.generate.model_manager.del_model(model_name) - # self.generate.model_manager.commit(opt.conf) - # updated_model_list = self.generate.model_manager.list_models() - # socketio.emit( - # "modelDeleted", - # { - # "deleted_model_name": model_name, - # "model_list": updated_model_list, - # }, - # ) - # print(f">> Model Deleted: {model_name}") - # except Exception as e: - # self.handle_exceptions(e) - - # @socketio.on("requestModelChange") - # def handle_set_model(model_name: str): - # try: - # print(f">> Model change requested: {model_name}") - # model = self.generate.set_model(model_name) - # model_list = self.generate.model_manager.list_models() - # if model is None: - # socketio.emit( - # "modelChangeFailed", - # {"model_name": model_name, "model_list": model_list}, - # ) - # else: - # socketio.emit( - # "modelChanged", - # {"model_name": model_name, "model_list": model_list}, - # ) - # except Exception as e: - # self.handle_exceptions(e) - - # @socketio.on("convertToDiffusers") - # def convert_to_diffusers(model_to_convert: dict): - # try: - # if model_info := self.generate.model_manager.model_info( - # model_name=model_to_convert["model_name"] - # ): - # if "weights" in model_info: - # ckpt_path = Path(model_info["weights"]) - # original_config_file = Path(model_info["config"]) - # model_name = model_to_convert["model_name"] - # model_description = model_info["description"] - # else: - # self.socketio.emit( - # "error", {"message": "Model is not a valid checkpoint file"} - # ) - # else: - # self.socketio.emit( - # "error", {"message": "Could not retrieve model info."} - # ) - - # if not ckpt_path.is_absolute(): - # ckpt_path = Path(Globals.root, ckpt_path) - - # if original_config_file and not original_config_file.is_absolute(): - # original_config_file = Path(Globals.root, original_config_file) - - # diffusers_path = Path( - # ckpt_path.parent.absolute(), f"{model_name}_diffusers" - # ) - - # if model_to_convert["save_location"] == "root": - # diffusers_path = Path( - # global_converted_ckpts_dir(), f"{model_name}_diffusers" - # ) - - # if ( - # model_to_convert["save_location"] == "custom" - # and model_to_convert["custom_location"] is not None - # ): - # diffusers_path = Path( - # model_to_convert["custom_location"], f"{model_name}_diffusers" - # ) - - # if diffusers_path.exists(): - # shutil.rmtree(diffusers_path) - - # self.generate.model_manager.convert_and_import( - # ckpt_path, - # diffusers_path, - # model_name=model_name, - # model_description=model_description, - # vae=None, - # original_config_file=original_config_file, - # commit_to_conf=opt.conf, - # ) - - # new_model_list = self.generate.model_manager.list_models() - # socketio.emit( - # "modelConverted", - # { - # "new_model_name": model_name, - # "model_list": new_model_list, - # "update": True, - # }, - # ) - # print(f">> Model Converted: {model_name}") - # except Exception as e: - # self.handle_exceptions(e) - - # @socketio.on("mergeDiffusersModels") - # def merge_diffusers_models(model_merge_info: dict): - # try: - # models_to_merge = model_merge_info["models_to_merge"] - # model_ids_or_paths = [ - # self.generate.model_manager.model_name_or_path(x) - # for x in models_to_merge - # ] - # merged_pipe = merge_diffusion_models( - # model_ids_or_paths, - # model_merge_info["alpha"], - # model_merge_info["interp"], - # model_merge_info["force"], - # ) - - # dump_path = global_models_dir() / "merged_models" - # if model_merge_info["model_merge_save_path"] is not None: - # dump_path = Path(model_merge_info["model_merge_save_path"]) - - # os.makedirs(dump_path, exist_ok=True) - # dump_path = dump_path / model_merge_info["merged_model_name"] - # merged_pipe.save_pretrained(dump_path, safe_serialization=1) - - # merged_model_config = dict( - # model_name=model_merge_info["merged_model_name"], - # description=f'Merge of models {", ".join(models_to_merge)}', - # commit_to_conf=opt.conf, - # ) - - # if vae := self.generate.model_manager.config[models_to_merge[0]].get( - # "vae", None - # ): - # print(f">> Using configured VAE assigned to {models_to_merge[0]}") - # merged_model_config.update(vae=vae) - - # self.generate.model_manager.import_diffuser_model( - # dump_path, **merged_model_config - # ) - # new_model_list = self.generate.model_manager.list_models() - - # socketio.emit( - # "modelsMerged", - # { - # "merged_models": models_to_merge, - # "merged_model_name": model_merge_info["merged_model_name"], - # "model_list": new_model_list, - # "update": True, - # }, - # ) - # print(f">> Models Merged: {models_to_merge}") - # print(f">> New Model Added: {model_merge_info['merged_model_name']}") - # except Exception as e: - # self.handle_exceptions(e) \ No newline at end of file + if not model_exists: + print(f">> Model not found") + raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found") + + # delete model + print(f">> Deleting Model: {model_name}") + ApiDependencies.invoker.services.model_manager.del_model(model_name, delete_files=True) + print(f">> Model Deleted: {model_name}") + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + From 9d80b28a4f88b740a1c1be38c6519f225624c112 Mon Sep 17 00:00:00 2001 From: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com> Date: Thu, 6 Apr 2023 16:23:09 -0400 Subject: [PATCH 02/14] Begin Convert Work --- invokeai/app/api/routers/models.py | 204 +++++++++++++++++++++++++++++ 1 file changed, 204 insertions(+) diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py index 07ecb45003..95a0f2817c 100644 --- a/invokeai/app/api/routers/models.py +++ b/invokeai/app/api/routers/models.py @@ -133,3 +133,207 @@ async def delete_model(model_name: str) -> None: except Exception as e: raise HTTPException(status_code=500, detail=str(e)) +@models_router.post( + "/{model_to_convert}", + operation_id="convert_model", + responses={ + 201: { + "model_response": "Model converted successfully.", + }, + 202: { + "description": "Model conversion is processing. Check back later." + }, + }, +) +async def convert_model(convert_request = ConvertedModelRequest) -> ConvertedModelResponse: + """ Convert Model """ + try: + convert_request_info = convert_request.info + info_dict = convert_request_info.dict() + convert_request = ConvertedModelRequest(name=convert_request.name, config=info_dict.config, weights=info_dict.weights, description=info_dict.description) + + if model_info := ApiDependencies.invoker.services.model_manager.model_info( + model_name=convert_request.name + ): + if "weights" in model_info: + ckpt_path = Path(convert_request.weights) + original_config_file = Path(convert_request.config) + model_name = convert_request.weights + model_description = convert_request.description + else: + raise HTTPException(status_code=404, detail=f"Model '{convert_request.name}' is not a valid checkpoint model") + else: + raise HTTPException(status_code=404, detail=f"Unable to retrieve model info") + + if not ckpt_path.is_absolute(): + ckpt_path = Path(Globals.root, ckpt_path) + + if original_config_file and not original_config_file.is_absolute(): + original_config_file = Path(Globals.root, original_config_file) + + diffusers_path = Path( + ckpt_path.parent.absolute(), f"{model_name}_diffusers" + ) + + if model_to_convert["save_location"] == "root": + diffusers_path = Path( + global_converted_ckpts_dir(), f"{model_name}_diffusers" + ) + + if ( + model_to_convert["save_location"] == "custom" + and model_to_convert["custom_location"] is not None + ): + diffusers_path = Path( + model_to_convert["custom_location"], f"{model_name}_diffusers" + ) + + if diffusers_path.exists(): + shutil.rmtree(diffusers_path) + + self.generate.model_manager.convert_and_import( + ckpt_path, + diffusers_path, + model_name=model_name, + model_description=model_description, + vae=None, + original_config_file=original_config_file, + commit_to_conf=opt.conf, + ) + + new_model_list = self.generate.model_manager.list_models() + socketio.emit( + "modelConverted", + { + "new_model_name": model_name, + "model_list": new_model_list, + "update": True, + }, + ) + print(f">> Model Converted: {model_name}") + + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + + # @socketio.on("convertToDiffusers") + # def convert_to_diffusers(model_to_convert: dict): + # try: + # if model_info := self.generate.model_manager.model_info( + # model_name=model_to_convert["model_name"] + # ): + # if "weights" in model_info: + # ckpt_path = Path(model_info["weights"]) + # original_config_file = Path(model_info["config"]) + # model_name = model_to_convert["model_name"] + # model_description = model_info["description"] + # else: + # self.socketio.emit( + # "error", {"message": "Model is not a valid checkpoint file"} + # ) + # else: + # self.socketio.emit( + # "error", {"message": "Could not retrieve model info."} + # ) + + # if not ckpt_path.is_absolute(): + # ckpt_path = Path(Globals.root, ckpt_path) + + # if original_config_file and not original_config_file.is_absolute(): + # original_config_file = Path(Globals.root, original_config_file) + + # diffusers_path = Path( + # ckpt_path.parent.absolute(), f"{model_name}_diffusers" + # ) + + # if model_to_convert["save_location"] == "root": + # diffusers_path = Path( + # global_converted_ckpts_dir(), f"{model_name}_diffusers" + # ) + + # if ( + # model_to_convert["save_location"] == "custom" + # and model_to_convert["custom_location"] is not None + # ): + # diffusers_path = Path( + # model_to_convert["custom_location"], f"{model_name}_diffusers" + # ) + + # if diffusers_path.exists(): + # shutil.rmtree(diffusers_path) + + # self.generate.model_manager.convert_and_import( + # ckpt_path, + # diffusers_path, + # model_name=model_name, + # model_description=model_description, + # vae=None, + # original_config_file=original_config_file, + # commit_to_conf=opt.conf, + # ) + + # new_model_list = self.generate.model_manager.list_models() + # socketio.emit( + # "modelConverted", + # { + # "new_model_name": model_name, + # "model_list": new_model_list, + # "update": True, + # }, + # ) + # print(f">> Model Converted: {model_name}") + # except Exception as e: + # self.handle_exceptions(e) + + # @socketio.on("mergeDiffusersModels") + # def merge_diffusers_models(model_merge_info: dict): + # try: + # models_to_merge = model_merge_info["models_to_merge"] + # model_ids_or_paths = [ + # self.generate.model_manager.model_name_or_path(x) + # for x in models_to_merge + # ] + # merged_pipe = merge_diffusion_models( + # model_ids_or_paths, + # model_merge_info["alpha"], + # model_merge_info["interp"], + # model_merge_info["force"], + # ) + + # dump_path = global_models_dir() / "merged_models" + # if model_merge_info["model_merge_save_path"] is not None: + # dump_path = Path(model_merge_info["model_merge_save_path"]) + + # os.makedirs(dump_path, exist_ok=True) + # dump_path = dump_path / model_merge_info["merged_model_name"] + # merged_pipe.save_pretrained(dump_path, safe_serialization=1) + + # merged_model_config = dict( + # model_name=model_merge_info["merged_model_name"], + # description=f'Merge of models {", ".join(models_to_merge)}', + # commit_to_conf=opt.conf, + # ) + + # if vae := self.generate.model_manager.config[models_to_merge[0]].get( + # "vae", None + # ): + # print(f">> Using configured VAE assigned to {models_to_merge[0]}") + # merged_model_config.update(vae=vae) + + # self.generate.model_manager.import_diffuser_model( + # dump_path, **merged_model_config + # ) + # new_model_list = self.generate.model_manager.list_models() + + # socketio.emit( + # "modelsMerged", + # { + # "merged_models": models_to_merge, + # "merged_model_name": model_merge_info["merged_model_name"], + # "model_list": new_model_list, + # "update": True, + # }, + # ) + # print(f">> Models Merged: {models_to_merge}") + # print(f">> New Model Added: {model_merge_info['merged_model_name']}") + # except Exception as e: \ No newline at end of file From 7919d81fb10d7b0201881775669e438eb727aa45 Mon Sep 17 00:00:00 2001 From: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com> Date: Thu, 6 Apr 2023 22:25:18 -0400 Subject: [PATCH 03/14] Update to address feedback --- invokeai/app/api/routers/models.py | 154 +++++++---------------------- 1 file changed, 33 insertions(+), 121 deletions(-) diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py index 95a0f2817c..aaeb9517d4 100644 --- a/invokeai/app/api/routers/models.py +++ b/invokeai/app/api/routers/models.py @@ -1,11 +1,17 @@ -# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) +# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) and 2023 Kent Keirsey (https://github.com/hipsterusername) +import shutil +import asyncio from typing import Annotated, Any, List, Literal, Optional, Union from fastapi.routing import APIRouter, HTTPException from pydantic import BaseModel, Field, parse_obj_as from pathlib import Path from ..dependencies import ApiDependencies +from invokeai.backend.globals import Globals, global_converted_ckpts_dir +from invokeai.backend.args import Args + + models_router = APIRouter(prefix="/v1/models", tags=["models"]) @@ -34,20 +40,22 @@ class DiffusersModelInfo(ModelInfo): repo_id: Optional[str] = Field(description="The repo ID to use for this model") path: Optional[str] = Field(description="The path to the model") -class CreateModelRequest (BaseModel): +class CreateModelRequest(BaseModel): name: str = Field(description="The name of the model") - info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model info") + info: Union[CkptModelInfo, DiffusersModelInfo] = Field(discriminator="format", description="The model info") -class CreateModelResponse (BaseModel): +class CreateModelResponse(BaseModel): name: str = Field(description="The name of the new model") - info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model info") + info: Union[CkptModelInfo, DiffusersModelInfo] = Field(discriminator="format", description="The model info") status: str = Field(description="The status of the API response") -class ConvertedModelRequest (BaseModel): +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): +class ConvertedModelResponse(BaseModel): name: str = Field(description="The name of the new model") info: DiffusersModelInfo = Field(description="The converted model info") @@ -70,34 +78,22 @@ async def list_models() -> ModelsList: @models_router.post( "/", operation_id="update_model", - responses={ - 201: { - "model_response": "Model added", - }, - 202: { - "description": "Model submission is processing. Check back later." - }, - }, + responses={200: {"status": "success"}}, ) async def update_model( model_request: CreateModelRequest ) -> CreateModelResponse: """ Add Model """ - try: - model_request_info = model_request.info - print(f">> Checking for {model_request_info}...") - info_dict = model_request_info.dict() + model_request_info = model_request.info + info_dict = model_request_info.dict() + model_response = CreateModelResponse(name=model_request.name, info=model_request.info, status="success") - ApiDependencies.invoker.services.model_manager.add_model( - model_name=model_request.name, - model_attributes=info_dict, - clobber=True, - ) - model_response = CreateModelResponse(name=model_request.name, info=model_request.info, status="success") + ApiDependencies.invoker.services.model_manager.add_model( + model_name=model_request.name, + model_attributes=info_dict, + clobber=True, + ) - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) - return model_response @@ -106,7 +102,7 @@ async def update_model( operation_id="del_model", responses={ 204: { - "description": "Model deleted" + "description": "Model deleted successfully" }, 404: { "description": "Model not found" @@ -117,103 +113,19 @@ async def delete_model(model_name: str) -> None: """Delete Model""" model_names = ApiDependencies.invoker.services.model_manager.model_names() model_exists = model_name in model_names - - try: - # check if model exists - print(f">> Checking for model {model_name}...") - if not model_exists: - print(f">> Model not found") - raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found") - - # delete model + # check if model exists + print(f">> Checking for model {model_name}...") + + if model_exists: print(f">> Deleting Model: {model_name}") ApiDependencies.invoker.services.model_manager.del_model(model_name, delete_files=True) print(f">> Model Deleted: {model_name}") - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + raise HTTPException(status_code=200, detail=f"Model '{model_name}' deleted successfully") -@models_router.post( - "/{model_to_convert}", - operation_id="convert_model", - responses={ - 201: { - "model_response": "Model converted successfully.", - }, - 202: { - "description": "Model conversion is processing. Check back later." - }, - }, -) -async def convert_model(convert_request = ConvertedModelRequest) -> ConvertedModelResponse: - """ Convert Model """ - try: - convert_request_info = convert_request.info - info_dict = convert_request_info.dict() - convert_request = ConvertedModelRequest(name=convert_request.name, config=info_dict.config, weights=info_dict.weights, description=info_dict.description) - - if model_info := ApiDependencies.invoker.services.model_manager.model_info( - model_name=convert_request.name - ): - if "weights" in model_info: - ckpt_path = Path(convert_request.weights) - original_config_file = Path(convert_request.config) - model_name = convert_request.weights - model_description = convert_request.description - else: - raise HTTPException(status_code=404, detail=f"Model '{convert_request.name}' is not a valid checkpoint model") - else: - raise HTTPException(status_code=404, detail=f"Unable to retrieve model info") - - if not ckpt_path.is_absolute(): - ckpt_path = Path(Globals.root, ckpt_path) - - if original_config_file and not original_config_file.is_absolute(): - original_config_file = Path(Globals.root, original_config_file) - - diffusers_path = Path( - ckpt_path.parent.absolute(), f"{model_name}_diffusers" - ) - - if model_to_convert["save_location"] == "root": - diffusers_path = Path( - global_converted_ckpts_dir(), f"{model_name}_diffusers" - ) - - if ( - model_to_convert["save_location"] == "custom" - and model_to_convert["custom_location"] is not None - ): - diffusers_path = Path( - model_to_convert["custom_location"], f"{model_name}_diffusers" - ) - - if diffusers_path.exists(): - shutil.rmtree(diffusers_path) - - self.generate.model_manager.convert_and_import( - ckpt_path, - diffusers_path, - model_name=model_name, - model_description=model_description, - vae=None, - original_config_file=original_config_file, - commit_to_conf=opt.conf, - ) - - new_model_list = self.generate.model_manager.list_models() - socketio.emit( - "modelConverted", - { - "new_model_name": model_name, - "model_list": new_model_list, - "update": True, - }, - ) - print(f">> Model Converted: {model_name}") - - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + if not model_exists: + print(f">> Model not found") + raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found") # @socketio.on("convertToDiffusers") From 5fe8cb56fc067cd843bf5bebfb759e554ceace4b Mon Sep 17 00:00:00 2001 From: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com> Date: Thu, 6 Apr 2023 22:26:28 -0400 Subject: [PATCH 04/14] Correct response note --- invokeai/app/api/routers/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py index aaeb9517d4..8332b44a35 100644 --- a/invokeai/app/api/routers/models.py +++ b/invokeai/app/api/routers/models.py @@ -121,7 +121,7 @@ async def delete_model(model_name: str) -> None: print(f">> Deleting Model: {model_name}") ApiDependencies.invoker.services.model_manager.del_model(model_name, delete_files=True) print(f">> Model Deleted: {model_name}") - raise HTTPException(status_code=200, detail=f"Model '{model_name}' deleted successfully") + raise HTTPException(status_code=204, detail=f"Model '{model_name}' deleted successfully") if not model_exists: print(f">> Model not found") From 54d9833db0c92d345c849c684e7f8404b2dc603d Mon Sep 17 00:00:00 2001 From: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com> Date: Fri, 7 Apr 2023 22:25:30 -0400 Subject: [PATCH 05/14] Else. --- invokeai/app/api/routers/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py index 8332b44a35..2de079cd6d 100644 --- a/invokeai/app/api/routers/models.py +++ b/invokeai/app/api/routers/models.py @@ -123,7 +123,7 @@ async def delete_model(model_name: str) -> None: print(f">> Model Deleted: {model_name}") raise HTTPException(status_code=204, detail=f"Model '{model_name}' deleted successfully") - if not model_exists: + else: print(f">> Model not found") raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found") From 34402cc46a50ebc525862de25b2f3b763261d25b Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 4 Apr 2023 11:05:15 +1000 Subject: [PATCH 06/14] feat(nodes): add list_images endpoint - add `list_images` endpoint at `GET api/v1/images` - extend `ImageStorageBase` with `list()` method, implemented it for `DiskImageStorage` - add `ImageReponse` class to for image responses, which includes urls, metadata - add `ImageMetadata` class (basically a stub at the moment) - uploaded images now named `"{uuid}_{timestamp}.png"` - add `models` modules. besides separating concerns more clearly, this helps to mitigate circular dependencies - improve thumbnail handling --- invokeai/app/api/models/images.py | 14 +++ invokeai/app/api/routers/images.py | 26 +++++- invokeai/app/cli/commands.py | 3 +- invokeai/app/invocations/cv.py | 4 +- invokeai/app/invocations/generate.py | 7 +- invokeai/app/invocations/image.py | 12 +-- invokeai/app/invocations/reconstruct.py | 4 +- invokeai/app/invocations/upscale.py | 4 +- invokeai/app/models/__init__.py | 0 invokeai/app/models/exceptions.py | 3 + invokeai/app/models/image.py | 26 ++++++ invokeai/app/models/metadata.py | 11 +++ invokeai/app/services/image_storage.py | 93 ++++++++++++++++--- invokeai/app/services/processor.py | 2 +- invokeai/app/util/__init__.py | 0 .../app/util/{util.py => step_callback.py} | 23 ++--- 16 files changed, 184 insertions(+), 48 deletions(-) create mode 100644 invokeai/app/api/models/images.py create mode 100644 invokeai/app/models/__init__.py create mode 100644 invokeai/app/models/exceptions.py create mode 100644 invokeai/app/models/image.py create mode 100644 invokeai/app/models/metadata.py create mode 100644 invokeai/app/util/__init__.py rename invokeai/app/util/{util.py => step_callback.py} (73%) diff --git a/invokeai/app/api/models/images.py b/invokeai/app/api/models/images.py new file mode 100644 index 0000000000..5ff0a48a44 --- /dev/null +++ b/invokeai/app/api/models/images.py @@ -0,0 +1,14 @@ +from pydantic import BaseModel, Field + +from invokeai.app.models.image import ImageType +from invokeai.app.models.metadata import ImageMetadata + + +class ImageResponse(BaseModel): + """The response type for images""" + + image_type: ImageType = Field(description="The type of the image") + image_name: str = Field(description="The name of the image") + image_url: str = Field(description="The url of the image") + thumbnail_url: str = Field(description="The url of the image's thumbnail") + metadata: ImageMetadata = Field(description="The image's metadata") diff --git a/invokeai/app/api/routers/images.py b/invokeai/app/api/routers/images.py index 453c114a28..bb3aabae6d 100644 --- a/invokeai/app/api/routers/images.py +++ b/invokeai/app/api/routers/images.py @@ -1,18 +1,20 @@ # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) from datetime import datetime, timezone +import uuid -from fastapi import Path, Request, UploadFile +from fastapi import Path, Query, Request, UploadFile from fastapi.responses import FileResponse, Response from fastapi.routing import APIRouter from PIL import Image +from invokeai.app.api.models.images import ImageResponse +from invokeai.app.services.item_storage import PaginatedResults from ...services.image_storage import ImageType from ..dependencies import ApiDependencies images_router = APIRouter(prefix="/v1/images", tags=["images"]) - @images_router.get("/{image_type}/{image_name}", operation_id="get_image") async def get_image( image_type: ImageType = Path(description="The type of image to get"), @@ -53,14 +55,30 @@ async def upload_image(file: UploadFile, request: Request): # Error opening the image return Response(status_code=415) - filename = f"{str(int(datetime.now(timezone.utc).timestamp()))}.png" + filename = f"{uuid.uuid4()}_{str(int(datetime.now(timezone.utc).timestamp()))}.png" ApiDependencies.invoker.services.images.save(ImageType.UPLOAD, filename, im) return Response( status_code=201, headers={ "Location": request.url_for( - "get_image", image_type=ImageType.UPLOAD, image_name=filename + "get_image", image_type=ImageType.UPLOAD.value, image_name=filename ) }, ) + +@images_router.get( + "/", + operation_id="list_images", + responses={200: {"model": PaginatedResults[ImageResponse]}}, +) +async def list_images( + image_type: ImageType = Query(default=ImageType.RESULT, description="The type of images to get"), + page: int = Query(default=0, description="The page of images to get"), + per_page: int = Query(default=10, description="The number of images per page"), +) -> PaginatedResults[ImageResponse]: + """Gets a list of images""" + result = ApiDependencies.invoker.services.images.list( + image_type, page, per_page + ) + return result diff --git a/invokeai/app/cli/commands.py b/invokeai/app/cli/commands.py index 5f4da73303..4e9c9aa581 100644 --- a/invokeai/app/cli/commands.py +++ b/invokeai/app/cli/commands.py @@ -6,7 +6,8 @@ from typing import Any, Callable, Iterable, Literal, get_args, get_origin, get_t from pydantic import BaseModel, Field import networkx as nx import matplotlib.pyplot as plt -from ..invocations.image import ImageField + +from ..models.image import ImageField from ..services.graph import GraphExecutionState from ..services.invoker import Invoker diff --git a/invokeai/app/invocations/cv.py b/invokeai/app/invocations/cv.py index 6a2788131b..ce784313cf 100644 --- a/invokeai/app/invocations/cv.py +++ b/invokeai/app/invocations/cv.py @@ -7,9 +7,9 @@ import numpy from PIL import Image, ImageOps from pydantic import Field -from ..services.image_storage import ImageType +from invokeai.app.models.image import ImageField, ImageType from .baseinvocation import BaseInvocation, InvocationContext -from .image import ImageField, ImageOutput +from .image import ImageOutput class CvInpaintInvocation(BaseInvocation): diff --git a/invokeai/app/invocations/generate.py b/invokeai/app/invocations/generate.py index d6e624b325..153d11189e 100644 --- a/invokeai/app/invocations/generate.py +++ b/invokeai/app/invocations/generate.py @@ -8,12 +8,13 @@ from torch import Tensor from pydantic import Field -from ..services.image_storage import ImageType +from invokeai.app.models.image import ImageField, ImageType from .baseinvocation import BaseInvocation, InvocationContext -from .image import ImageField, ImageOutput +from .image import ImageOutput from ...backend.generator import Txt2Img, Img2Img, Inpaint, InvokeAIGenerator from ...backend.stable_diffusion import PipelineIntermediateState -from ..util.util import diffusers_step_callback_adapter, CanceledException +from ..models.exceptions import CanceledException +from ..util.step_callback import diffusers_step_callback_adapter SAMPLER_NAME_VALUES = Literal[ tuple(InvokeAIGenerator.schedulers()) diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index 65ea4c3edb..491a4895a6 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -7,20 +7,10 @@ import numpy from PIL import Image, ImageFilter, ImageOps from pydantic import BaseModel, Field -from ..services.image_storage import ImageType +from ..models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext - -class ImageField(BaseModel): - """An image field used for passing image objects between invocations""" - - image_type: str = Field( - default=ImageType.RESULT, description="The type of the image" - ) - image_name: Optional[str] = Field(default=None, description="The name of the image") - - class ImageOutput(BaseInvocationOutput): """Base class for invocations that output an image""" #fmt: off diff --git a/invokeai/app/invocations/reconstruct.py b/invokeai/app/invocations/reconstruct.py index c4d8f3ac7c..68449729d6 100644 --- a/invokeai/app/invocations/reconstruct.py +++ b/invokeai/app/invocations/reconstruct.py @@ -3,10 +3,10 @@ from typing import Literal, Union from pydantic import Field -from ..services.image_storage import ImageType +from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices from .baseinvocation import BaseInvocation, InvocationContext -from .image import ImageField, ImageOutput +from .image import ImageOutput class RestoreFaceInvocation(BaseInvocation): """Restores faces in an image.""" diff --git a/invokeai/app/invocations/upscale.py b/invokeai/app/invocations/upscale.py index 4079877fdb..ea3221572e 100644 --- a/invokeai/app/invocations/upscale.py +++ b/invokeai/app/invocations/upscale.py @@ -5,10 +5,10 @@ from typing import Literal, Union from pydantic import Field -from ..services.image_storage import ImageType +from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices from .baseinvocation import BaseInvocation, InvocationContext -from .image import ImageField, ImageOutput +from .image import ImageOutput class UpscaleInvocation(BaseInvocation): diff --git a/invokeai/app/models/__init__.py b/invokeai/app/models/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/invokeai/app/models/exceptions.py b/invokeai/app/models/exceptions.py new file mode 100644 index 0000000000..32ad3b8f03 --- /dev/null +++ b/invokeai/app/models/exceptions.py @@ -0,0 +1,3 @@ +class CanceledException(Exception): + """Execution canceled by user.""" + pass diff --git a/invokeai/app/models/image.py b/invokeai/app/models/image.py new file mode 100644 index 0000000000..1561e6bcc5 --- /dev/null +++ b/invokeai/app/models/image.py @@ -0,0 +1,26 @@ +from enum import Enum +from typing import Optional +from pydantic import BaseModel, Field + + +class ImageType(str, Enum): + RESULT = "results" + INTERMEDIATE = "intermediates" + UPLOAD = "uploads" + + +class ImageField(BaseModel): + """An image field used for passing image objects between invocations""" + + image_type: str = Field( + default=ImageType.RESULT, description="The type of the image" + ) + image_name: Optional[str] = Field(default=None, description="The name of the image") + + class Config: + schema_extra = { + "required": [ + "image_type", + "image_name", + ] + } diff --git a/invokeai/app/models/metadata.py b/invokeai/app/models/metadata.py new file mode 100644 index 0000000000..2531168272 --- /dev/null +++ b/invokeai/app/models/metadata.py @@ -0,0 +1,11 @@ +from typing import Optional +from pydantic import BaseModel, Field + +class ImageMetadata(BaseModel): + """An image's metadata""" + + timestamp: float = Field(description="The creation timestamp of the image") + width: int = Field(description="The width of the image in pixels") + height: int = Field(description="The height of the image in pixels") + # TODO: figure out metadata + sd_metadata: Optional[dict] = Field(default={}, description="The image's SD-specific metadata") diff --git a/invokeai/app/services/image_storage.py b/invokeai/app/services/image_storage.py index c80a4bfb31..80d72efca8 100644 --- a/invokeai/app/services/image_storage.py +++ b/invokeai/app/services/image_storage.py @@ -2,24 +2,25 @@ import datetime import os +from glob import glob from abc import ABC, abstractmethod from enum import Enum from pathlib import Path from queue import Queue -from typing import Dict +from typing import Callable, Dict, List from PIL.Image import Image +import PIL.Image as PILImage +from pydantic import BaseModel +from invokeai.app.api.models.images import ImageResponse +from invokeai.app.models.image import ImageField, ImageType +from invokeai.app.models.metadata import ImageMetadata +from invokeai.app.services.item_storage import PaginatedResults from invokeai.app.util.save_thumbnail import save_thumbnail from invokeai.backend.image_util import PngWriter -class ImageType(str, Enum): - RESULT = "results" - INTERMEDIATE = "intermediates" - UPLOAD = "uploads" - - class ImageStorageBase(ABC): """Responsible for storing and retrieving images.""" @@ -27,9 +28,17 @@ class ImageStorageBase(ABC): def get(self, image_type: ImageType, image_name: str) -> Image: pass + @abstractmethod + def list( + self, image_type: ImageType, page: int = 0, per_page: int = 10 + ) -> PaginatedResults[ImageResponse]: + pass + # TODO: make this a bit more flexible for e.g. cloud storage @abstractmethod - def get_path(self, image_type: ImageType, image_name: str) -> str: + def get_path( + self, image_type: ImageType, image_name: str, is_thumbnail: bool = False + ) -> str: pass @abstractmethod @@ -71,19 +80,74 @@ class DiskImageStorage(ImageStorageBase): parents=True, exist_ok=True ) + def list( + self, image_type: ImageType, page: int = 0, per_page: int = 10 + ) -> PaginatedResults[ImageResponse]: + dir_path = os.path.join(self.__output_folder, image_type) + image_paths = glob(f"{dir_path}/*.png") + count = len(image_paths) + + sorted_image_paths = sorted( + glob(f"{dir_path}/*.png"), key=os.path.getctime, reverse=True + ) + + page_of_image_paths = sorted_image_paths[ + page * per_page : (page + 1) * per_page + ] + + page_of_images: List[ImageResponse] = [] + + for path in page_of_image_paths: + filename = os.path.basename(path) + img = PILImage.open(path) + page_of_images.append( + ImageResponse( + image_type=image_type.value, + image_name=filename, + # TODO: DiskImageStorage should not be building URLs...? + image_url=f"api/v1/images/{image_type.value}/{filename}", + thumbnail_url=f"api/v1/images/{image_type.value}/thumbnails/{os.path.splitext(filename)[0]}.webp", + # TODO: Creation of this object should happen elsewhere, just making it fit here so it works + metadata=ImageMetadata( + timestamp=os.path.getctime(path), + width=img.width, + height=img.height, + ), + ) + ) + + page_count_trunc = int(count / per_page) + page_count_mod = count % per_page + page_count = page_count_trunc if page_count_mod == 0 else page_count_trunc + 1 + + return PaginatedResults[ImageResponse]( + items=page_of_images, + page=page, + pages=page_count, + per_page=per_page, + total=count, + ) + def get(self, image_type: ImageType, image_name: str) -> Image: image_path = self.get_path(image_type, image_name) cache_item = self.__get_cache(image_path) if cache_item: return cache_item - image = Image.open(image_path) + image = PILImage.open(image_path) self.__set_cache(image_path, image) return image # TODO: make this a bit more flexible for e.g. cloud storage - def get_path(self, image_type: ImageType, image_name: str) -> str: - path = os.path.join(self.__output_folder, image_type, image_name) + def get_path( + self, image_type: ImageType, image_name: str, is_thumbnail: bool = False + ) -> str: + if is_thumbnail: + path = os.path.join( + self.__output_folder, image_type, "thumbnails", image_name + ) + else: + path = os.path.join(self.__output_folder, image_type, image_name) return path def save(self, image_type: ImageType, image_name: str, image: Image) -> None: @@ -101,12 +165,19 @@ class DiskImageStorage(ImageStorageBase): def delete(self, image_type: ImageType, image_name: str) -> None: image_path = self.get_path(image_type, image_name) + thumbnail_path = self.get_path(image_type, image_name, True) if os.path.exists(image_path): os.remove(image_path) if image_path in self.__cache: del self.__cache[image_path] + if os.path.exists(thumbnail_path): + os.remove(thumbnail_path) + + if thumbnail_path in self.__cache: + del self.__cache[thumbnail_path] + def __get_cache(self, image_name: str) -> Image: return None if image_name not in self.__cache else self.__cache[image_name] diff --git a/invokeai/app/services/processor.py b/invokeai/app/services/processor.py index b460563278..0125d7eb62 100644 --- a/invokeai/app/services/processor.py +++ b/invokeai/app/services/processor.py @@ -4,7 +4,7 @@ from threading import Event, Thread from ..invocations.baseinvocation import InvocationContext from .invocation_queue import InvocationQueueItem from .invoker import InvocationProcessorABC, Invoker -from ..util.util import CanceledException +from ..models.exceptions import CanceledException class DefaultInvocationProcessor(InvocationProcessorABC): __invoker_thread: Thread diff --git a/invokeai/app/util/__init__.py b/invokeai/app/util/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/invokeai/app/util/util.py b/invokeai/app/util/step_callback.py similarity index 73% rename from invokeai/app/util/util.py rename to invokeai/app/util/step_callback.py index 60a5072cb0..466f78ddb0 100644 --- a/invokeai/app/util/util.py +++ b/invokeai/app/util/step_callback.py @@ -1,14 +1,16 @@ import torch -from PIL import Image from ..invocations.baseinvocation import InvocationContext from ...backend.util.util import image_to_dataURL from ...backend.generator.base import Generator from ...backend.stable_diffusion import PipelineIntermediateState -class CanceledException(Exception): - pass - -def fast_latents_step_callback(sample: torch.Tensor, step: int, steps: int, id: str, context: InvocationContext, ): +def fast_latents_step_callback( + sample: torch.Tensor, + step: int, + steps: int, + id: str, + context: InvocationContext, +): # TODO: only output a preview image when requested image = Generator.sample_to_lowres_estimated_image(sample) @@ -21,15 +23,12 @@ def fast_latents_step_callback(sample: torch.Tensor, step: int, steps: int, id: context.services.events.emit_generator_progress( context.graph_execution_state_id, id, - { - "width": width, - "height": height, - "dataURL": dataURL - }, + {"width": width, "height": height, "dataURL": dataURL}, step, steps, ) + def diffusers_step_callback_adapter(*cb_args, **kwargs): """ txt2img gives us a Tensor in the step_callbak, while img2img gives us a PipelineIntermediateState. @@ -37,6 +36,8 @@ def diffusers_step_callback_adapter(*cb_args, **kwargs): """ if isinstance(cb_args[0], PipelineIntermediateState): progress_state: PipelineIntermediateState = cb_args[0] - return fast_latents_step_callback(progress_state.latents, progress_state.step, **kwargs) + return fast_latents_step_callback( + progress_state.latents, progress_state.step, **kwargs + ) else: return fast_latents_step_callback(*cb_args, **kwargs) From 4463124bddd221c333d4c70e73aa2949ad35453d Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 3 Apr 2023 18:56:48 +1000 Subject: [PATCH 07/14] feat(nodes): mark ImageField properties required, add docs --- docs/contributing/INVOCATIONS.md | 129 +++++++++++++++++++++++++++---- invokeai/app/services/graph.py | 3 - 2 files changed, 113 insertions(+), 19 deletions(-) diff --git a/docs/contributing/INVOCATIONS.md b/docs/contributing/INVOCATIONS.md index c8a97c19e4..1f12cfc8f5 100644 --- a/docs/contributing/INVOCATIONS.md +++ b/docs/contributing/INVOCATIONS.md @@ -1,10 +1,18 @@ # Invocations -Invocations represent a single operation, its inputs, and its outputs. These operations and their outputs can be chained together to generate and modify images. +Invocations represent a single operation, its inputs, and its outputs. These +operations and their outputs can be chained together to generate and modify +images. ## Creating a new invocation -To create a new invocation, either find the appropriate module file in `/ldm/invoke/app/invocations` to add your invocation to, or create a new one in that folder. All invocations in that folder will be discovered and made available to the CLI and API automatically. Invocations make use of [typing](https://docs.python.org/3/library/typing.html) and [pydantic](https://pydantic-docs.helpmanual.io/) for validation and integration into the CLI and API. +To create a new invocation, either find the appropriate module file in +`/ldm/invoke/app/invocations` to add your invocation to, or create a new one in +that folder. All invocations in that folder will be discovered and made +available to the CLI and API automatically. Invocations make use of +[typing](https://docs.python.org/3/library/typing.html) and +[pydantic](https://pydantic-docs.helpmanual.io/) for validation and integration +into the CLI and API. An invocation looks like this: @@ -41,34 +49,54 @@ class UpscaleInvocation(BaseInvocation): Each portion is important to implement correctly. ### Class definition and type + ```py class UpscaleInvocation(BaseInvocation): """Upscales an image.""" type: Literal['upscale'] = 'upscale' ``` -All invocations must derive from `BaseInvocation`. They should have a docstring that declares what they do in a single, short line. They should also have a `type` with a type hint that's `Literal["command_name"]`, where `command_name` is what the user will type on the CLI or use in the API to create this invocation. The `command_name` must be unique. The `type` must be assigned to the value of the literal in the type hint. + +All invocations must derive from `BaseInvocation`. They should have a docstring +that declares what they do in a single, short line. They should also have a +`type` with a type hint that's `Literal["command_name"]`, where `command_name` +is what the user will type on the CLI or use in the API to create this +invocation. The `command_name` must be unique. The `type` must be assigned to +the value of the literal in the type hint. ### Inputs + ```py # Inputs image: Union[ImageField,None] = Field(description="The input image") strength: float = Field(default=0.75, gt=0, le=1, description="The strength") level: Literal[2,4] = Field(default=2, description="The upscale level") ``` -Inputs consist of three parts: a name, a type hint, and a `Field` with default, description, and validation information. For example: -| Part | Value | Description | -| ---- | ----- | ----------- | -| Name | `strength` | This field is referred to as `strength` | -| Type Hint | `float` | This field must be of type `float` | -| Field | `Field(default=0.75, gt=0, le=1, description="The strength")` | The default value is `0.75`, the value must be in the range (0,1], and help text will show "The strength" for this field. | -Notice that `image` has type `Union[ImageField,None]`. The `Union` allows this field to be parsed with `None` as a value, which enables linking to previous invocations. All fields should either provide a default value or allow `None` as a value, so that they can be overwritten with a linked output from another invocation. +Inputs consist of three parts: a name, a type hint, and a `Field` with default, +description, and validation information. For example: -The special type `ImageField` is also used here. All images are passed as `ImageField`, which protects them from pydantic validation errors (since images only ever come from links). +| Part | Value | Description | +| --------- | ------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------- | +| Name | `strength` | This field is referred to as `strength` | +| Type Hint | `float` | This field must be of type `float` | +| Field | `Field(default=0.75, gt=0, le=1, description="The strength")` | The default value is `0.75`, the value must be in the range (0,1], and help text will show "The strength" for this field. | -Finally, note that for all linking, the `type` of the linked fields must match. If the `name` also matches, then the field can be **automatically linked** to a previous invocation by name and matching. +Notice that `image` has type `Union[ImageField,None]`. The `Union` allows this +field to be parsed with `None` as a value, which enables linking to previous +invocations. All fields should either provide a default value or allow `None` as +a value, so that they can be overwritten with a linked output from another +invocation. + +The special type `ImageField` is also used here. All images are passed as +`ImageField`, which protects them from pydantic validation errors (since images +only ever come from links). + +Finally, note that for all linking, the `type` of the linked fields must match. +If the `name` also matches, then the field can be **automatically linked** to a +previous invocation by name and matching. ### Invoke Function + ```py def invoke(self, context: InvocationContext) -> ImageOutput: image = context.services.images.get(self.image.image_type, self.image.image_name) @@ -88,13 +116,22 @@ Finally, note that for all linking, the `type` of the linked fields must match. image = ImageField(image_type = image_type, image_name = image_name) ) ``` -The `invoke` function is the last portion of an invocation. It is provided an `InvocationContext` which contains services to perform work as well as a `session_id` for use as needed. It should return a class with output values that derives from `BaseInvocationOutput`. -Before being called, the invocation will have all of its fields set from defaults, inputs, and finally links (overriding in that order). +The `invoke` function is the last portion of an invocation. It is provided an +`InvocationContext` which contains services to perform work as well as a +`session_id` for use as needed. It should return a class with output values that +derives from `BaseInvocationOutput`. -Assume that this invocation may be running simultaneously with other invocations, may be running on another machine, or in other interesting scenarios. If you need functionality, please provide it as a service in the `InvocationServices` class, and make sure it can be overridden. +Before being called, the invocation will have all of its fields set from +defaults, inputs, and finally links (overriding in that order). + +Assume that this invocation may be running simultaneously with other +invocations, may be running on another machine, or in other interesting +scenarios. If you need functionality, please provide it as a service in the +`InvocationServices` class, and make sure it can be overridden. ### Outputs + ```py class ImageOutput(BaseInvocationOutput): """Base class for invocations that output an image""" @@ -102,4 +139,64 @@ class ImageOutput(BaseInvocationOutput): image: ImageField = Field(default=None, description="The output image") ``` -Output classes look like an invocation class without the invoke method. Prefer to use an existing output class if available, and prefer to name inputs the same as outputs when possible, to promote automatic invocation linking. + +Output classes look like an invocation class without the invoke method. Prefer +to use an existing output class if available, and prefer to name inputs the same +as outputs when possible, to promote automatic invocation linking. + +## Schema Generation + +Invocation, output and related classes are used to generate an OpenAPI schema. + +### Required Properties + +The schema generation treat all properties with default values as optional. This +makes sense internally, but when when using these classes via the generated +schema, we end up with e.g. the `ImageOutput` class having its `image` property +marked as optional. + +We know that this property will always be present, so the additional logic +needed to always check if the property exists adds a lot of extraneous cruft. + +To fix this, we can leverage `pydantic`'s +[schema customisation](https://docs.pydantic.dev/usage/schema/#schema-customization) +to mark properties that we know will always be present as required. + +Here's that `ImageOutput` class, without the needed schema customisation: + +```python +class ImageOutput(BaseInvocationOutput): + """Base class for invocations that output an image""" + + type: Literal["image"] = "image" + image: ImageField = Field(default=None, description="The output image") +``` + +The generated OpenAPI schema, and all clients/types generated from it, will have +the `type` and `image` properties marked as optional, even though we know they +will always have a value by the time we can interact with them via the API. + +Here's the same class, but with the schema customisation added: + +```python +class ImageOutput(BaseInvocationOutput): + """Base class for invocations that output an image""" + + type: Literal["image"] = "image" + image: ImageField = Field(default=None, description="The output image") + + class Config: + schema_extra = { + 'required': [ + 'type', + 'image', + ] + } +``` + +The resultant schema (and any API client or types generated from it) will now +have see `type` as string literal `"image"` and `image` as an `ImageField` +object. + +See this `pydantic` issue for discussion on this solution: + diff --git a/invokeai/app/services/graph.py b/invokeai/app/services/graph.py index 98c2f29308..e286569bcc 100644 --- a/invokeai/app/services/graph.py +++ b/invokeai/app/services/graph.py @@ -794,9 +794,6 @@ class GraphExecutionState(BaseModel): default_factory=dict, ) - # Declare all fields as required; necessary for OpenAPI schema generation build. - # Technically only fields without a `default_factory` need to be listed here. - # See: https://github.com/pydantic/pydantic/discussions/4577 class Config: schema_extra = { 'required': [ From de189f2db6c4d5dbf893b93300b10fe1e3a77d5a Mon Sep 17 00:00:00 2001 From: AbdBarho Date: Sun, 9 Apr 2023 21:53:59 +0200 Subject: [PATCH 08/14] Increase chunk size when computing SHAs --- invokeai/backend/model_management/model_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index a51a2fec22..534b526081 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -1204,7 +1204,7 @@ class ModelManager(object): return self.device.type == "cuda" def _diffuser_sha256( - self, name_or_path: Union[str, Path], chunksize=4096 + self, name_or_path: Union[str, Path], chunksize=16777216 ) -> Union[str, bytes]: path = None if isinstance(name_or_path, Path): From 5bd0bb637f11fb4acc8ac504300a85f347a1df8a Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 10 Apr 2023 18:14:06 +1000 Subject: [PATCH 09/14] fix(nodes): add missing type to `ImageField` --- invokeai/app/models/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/models/image.py b/invokeai/app/models/image.py index 1561e6bcc5..9edb16800d 100644 --- a/invokeai/app/models/image.py +++ b/invokeai/app/models/image.py @@ -12,7 +12,7 @@ class ImageType(str, Enum): class ImageField(BaseModel): """An image field used for passing image objects between invocations""" - image_type: str = Field( + image_type: ImageType = Field( default=ImageType.RESULT, description="The type of the image" ) image_name: Optional[str] = Field(default=None, description="The name of the image") From dad3a7f263ce1774a3454d305f7e62e22a909fb3 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 10 Apr 2023 18:13:23 +1000 Subject: [PATCH 10/14] fix(nodes): `sampler_name` --> `scheduler` the name of this was changed at some point. nodes still used the old name, so scheduler selection did nothing. simple fix. --- invokeai/app/invocations/generate.py | 2 +- invokeai/app/invocations/latent.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/invokeai/app/invocations/generate.py b/invokeai/app/invocations/generate.py index 153d11189e..70c695fd2e 100644 --- a/invokeai/app/invocations/generate.py +++ b/invokeai/app/invocations/generate.py @@ -35,7 +35,7 @@ class TextToImageInvocation(BaseInvocation): width: int = Field(default=512, multiple_of=64, gt=0, description="The width of the resulting image", ) height: int = Field(default=512, multiple_of=64, gt=0, description="The height of the resulting image", ) cfg_scale: float = Field(default=7.5, gt=0, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", ) - sampler_name: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The sampler to use" ) + scheduler: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The scheduler to use" ) seamless: bool = Field(default=False, description="Whether or not to generate an image that can tile without seams", ) model: str = Field(default="", description="The model to use (currently ignored)") progress_images: bool = Field(default=False, description="Whether or not to produce progress images during generation", ) diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index 49c3c4f11e..ca3c7246c7 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -136,7 +136,7 @@ class TextToLatentsInvocation(BaseInvocation): width: int = Field(default=512, multiple_of=64, gt=0, description="The width of the resulting image", ) height: int = Field(default=512, multiple_of=64, gt=0, description="The height of the resulting image", ) cfg_scale: float = Field(default=7.5, gt=0, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", ) - sampler_name: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The sampler to use" ) + scheduler: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The scheduler to use" ) seamless: bool = Field(default=False, description="Whether or not to generate an image that can tile without seams", ) seamless_axes: str = Field(default="", description="The axes to tile the image on, 'x' and/or 'y'") model: str = Field(default="", description="The model to use (currently ignored)") @@ -175,7 +175,7 @@ class TextToLatentsInvocation(BaseInvocation): model: StableDiffusionGeneratorPipeline = model_info['model'] model.scheduler = get_scheduler( model=model, - scheduler_name=self.sampler_name + scheduler_name=self.scheduler ) if isinstance(model, DiffusionPipeline): From 427db7c7e2dd4b54cf10b4f4328a16c99da9151e Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Sun, 9 Apr 2023 22:33:16 +1000 Subject: [PATCH 11/14] feat(nodes): fix typo in PasteImageInvocation --- invokeai/app/invocations/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index 491a4895a6..9f072036a4 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -139,7 +139,7 @@ class PasteImageInvocation(BaseInvocation): None if self.mask is None else ImageOps.invert( - services.images.get(self.mask.image_type, self.mask.image_name) + context.services.images.get(self.mask.image_type, self.mask.image_name) ) ) # TODO: probably shouldn't invert mask here... should user be required to do it? From 07e3a0ec1545ee64ecf76c1d118a8e995b5a4a2c Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 10 Apr 2023 19:07:48 +1000 Subject: [PATCH 12/14] feat(nodes): add invocation schema customisation, add model selection - add invocation schema customisation done via fastapi's `Config` class and `schema_extra`. when using `Config`, inherit from `InvocationConfig` to get type hints. where it makes sense - like for all math invocations - define a `MathInvocationConfig` class and have all invocations inherit from it. this customisation can provide any arbitrary additional data to the UI. currently it provides tags and field type hints. this is necessary for `model` type fields, which are actually string fields. without something like this, we can't reliably differentiate `model` fields from normal `string` fields. can also be used for future field types. all invocations now have tags, and all `model` fields have ui type hints. - fix model handling for invocations added a helper to fall back to the default model if an invalid model name is chosen. model names in graphs now work. - fix latents progress callback noticed this wasn't correct while working on everything else. --- invokeai/app/invocations/cv.py | 17 ++++- invokeai/app/invocations/generate.py | 49 +++++++----- invokeai/app/invocations/image.py | 25 ++++-- invokeai/app/invocations/latent.py | 88 +++++++++++++++------- invokeai/app/invocations/math.py | 29 ++++--- invokeai/app/invocations/models/config.py | 54 +++++++++++++ invokeai/app/invocations/reconstruct.py | 9 +++ invokeai/app/invocations/upscale.py | 10 +++ invokeai/app/invocations/util/get_model.py | 11 +++ 9 files changed, 228 insertions(+), 64 deletions(-) create mode 100644 invokeai/app/invocations/models/config.py create mode 100644 invokeai/app/invocations/util/get_model.py diff --git a/invokeai/app/invocations/cv.py b/invokeai/app/invocations/cv.py index ce784313cf..adcb2405c4 100644 --- a/invokeai/app/invocations/cv.py +++ b/invokeai/app/invocations/cv.py @@ -5,14 +5,27 @@ from typing import Literal import cv2 as cv import numpy from PIL import Image, ImageOps -from pydantic import Field +from pydantic import BaseModel, Field +from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType from .baseinvocation import BaseInvocation, InvocationContext from .image import ImageOutput -class CvInpaintInvocation(BaseInvocation): +class CVInvocation(BaseModel): + """Helper class to provide all OpenCV invocations with additional config""" + + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["cv", "image"], + }, + } + + +class CvInpaintInvocation(BaseInvocation, CVInvocation): """Simple inpaint using opencv.""" #fmt: off type: Literal["cv_inpaint"] = "cv_inpaint" diff --git a/invokeai/app/invocations/generate.py b/invokeai/app/invocations/generate.py index 70c695fd2e..a07b1ac379 100644 --- a/invokeai/app/invocations/generate.py +++ b/invokeai/app/invocations/generate.py @@ -6,9 +6,13 @@ from typing import Literal, Optional, Union import numpy as np from torch import Tensor -from pydantic import Field +from pydantic import BaseModel, Field +from invokeai.app.invocations.models.config import ( + InvocationConfig, +) from invokeai.app.models.image import ImageField, ImageType +from invokeai.app.invocations.util.get_model import choose_model from .baseinvocation import BaseInvocation, InvocationContext from .image import ImageOutput from ...backend.generator import Txt2Img, Img2Img, Inpaint, InvokeAIGenerator @@ -16,12 +20,26 @@ from ...backend.stable_diffusion import PipelineIntermediateState from ..models.exceptions import CanceledException from ..util.step_callback import diffusers_step_callback_adapter -SAMPLER_NAME_VALUES = Literal[ - tuple(InvokeAIGenerator.schedulers()) -] +SAMPLER_NAME_VALUES = Literal[tuple(InvokeAIGenerator.schedulers())] + + +class SDImageInvocation(BaseModel): + """Helper class to provide all Stable Diffusion raster image invocations with additional config""" + + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["stable-diffusion", "image"], + "type_hints": { + "model": "model", + }, + }, + } + # Text to image -class TextToImageInvocation(BaseInvocation): +class TextToImageInvocation(BaseInvocation, SDImageInvocation): """Generates an image using text2img.""" type: Literal["txt2img"] = "txt2img" @@ -59,16 +77,9 @@ class TextToImageInvocation(BaseInvocation): diffusers_step_callback_adapter(sample, step, steps=self.steps, id=self.id, context=context) def invoke(self, context: InvocationContext) -> ImageOutput: - # def step_callback(state: PipelineIntermediateState): - # if (context.services.queue.is_canceled(context.graph_execution_state_id)): - # raise CanceledException - # self.dispatch_progress(context, state.latents, state.step) - # Handle invalid model parameter - # TODO: figure out if this can be done via a validator that uses the model_cache - # TODO: How to get the default model name now? - # (right now uses whatever current model is set in model manager) - model= context.services.model_manager.get_model() + model = choose_model(context.services.model_manager, self.model) + outputs = Txt2Img(model).generate( prompt=self.prompt, step_callback=partial(self.dispatch_progress, context), @@ -135,9 +146,8 @@ class ImageToImageInvocation(TextToImageInvocation): mask = None # Handle invalid model parameter - # TODO: figure out if this can be done via a validator that uses the model_cache - # TODO: How to get the default model name now? - model = context.services.model_manager.get_model() + model = choose_model(context.services.model_manager, self.model) + outputs = Img2Img(model).generate( prompt=self.prompt, init_image=image, @@ -211,9 +221,8 @@ class InpaintInvocation(ImageToImageInvocation): ) # Handle invalid model parameter - # TODO: figure out if this can be done via a validator that uses the model_cache - # TODO: How to get the default model name now? - model = context.services.model_manager.get_model() + model = choose_model(context.services.model_manager, self.model) + outputs = Inpaint(model).generate( prompt=self.prompt, init_img=image, diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index 9f072036a4..0f783b2541 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -7,10 +7,23 @@ import numpy from PIL import Image, ImageFilter, ImageOps from pydantic import BaseModel, Field +from invokeai.app.invocations.models.config import InvocationConfig + from ..models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext + +class PILInvocationConfig(BaseModel): + """Helper class to provide all PIL invocations with additional config""" + + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["PIL", "image"], + }, + } + class ImageOutput(BaseInvocationOutput): """Base class for invocations that output an image""" #fmt: off @@ -82,7 +95,7 @@ class ShowImageInvocation(BaseInvocation): ) -class CropImageInvocation(BaseInvocation): +class CropImageInvocation(BaseInvocation, PILInvocationConfig): """Crops an image to a specified box. The box can be outside of the image.""" #fmt: off type: Literal["crop"] = "crop" @@ -115,7 +128,7 @@ class CropImageInvocation(BaseInvocation): ) -class PasteImageInvocation(BaseInvocation): +class PasteImageInvocation(BaseInvocation, PILInvocationConfig): """Pastes an image into another image.""" #fmt: off type: Literal["paste"] = "paste" @@ -165,7 +178,7 @@ class PasteImageInvocation(BaseInvocation): ) -class MaskFromAlphaInvocation(BaseInvocation): +class MaskFromAlphaInvocation(BaseInvocation, PILInvocationConfig): """Extracts the alpha channel of an image as a mask.""" #fmt: off type: Literal["tomask"] = "tomask" @@ -192,7 +205,7 @@ class MaskFromAlphaInvocation(BaseInvocation): return MaskOutput(mask=ImageField(image_type=image_type, image_name=image_name)) -class BlurInvocation(BaseInvocation): +class BlurInvocation(BaseInvocation, PILInvocationConfig): """Blurs an image""" #fmt: off @@ -226,7 +239,7 @@ class BlurInvocation(BaseInvocation): ) -class LerpInvocation(BaseInvocation): +class LerpInvocation(BaseInvocation, PILInvocationConfig): """Linear interpolation of all pixels of an image""" #fmt: off type: Literal["lerp"] = "lerp" @@ -257,7 +270,7 @@ class LerpInvocation(BaseInvocation): ) -class InverseLerpInvocation(BaseInvocation): +class InverseLerpInvocation(BaseInvocation, PILInvocationConfig): """Inverse linear interpolation of all pixels of an image""" #fmt: off type: Literal["ilerp"] = "ilerp" diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index ca3c7246c7..d59d681f40 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -2,9 +2,13 @@ from typing import Literal, Optional from pydantic import BaseModel, Field -from torch import Tensor import torch +from invokeai.app.invocations.models.config import InvocationConfig +from invokeai.app.models.exceptions import CanceledException +from invokeai.app.invocations.util.get_model import choose_model +from invokeai.app.util.step_callback import diffusers_step_callback_adapter + from ...backend.model_management.model_manager import ModelManager from ...backend.util.devices import choose_torch_device, torch_dtype from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import PostprocessingSettings @@ -13,13 +17,10 @@ from ...backend.prompting.conditioning import get_uc_and_c_and_ec from ...backend.stable_diffusion.diffusers_pipeline import ConditioningData, StableDiffusionGeneratorPipeline from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext import numpy as np -from accelerate.utils import set_seed from ..services.image_storage import ImageType from .baseinvocation import BaseInvocation, InvocationContext from .image import ImageField, ImageOutput -from ...backend.generator import Generator from ...backend.stable_diffusion import PipelineIntermediateState -from ...backend.util.util import image_to_dataURL from diffusers.schedulers import SchedulerMixin as Scheduler import diffusers from diffusers import DiffusionPipeline @@ -109,6 +110,15 @@ class NoiseInvocation(BaseInvocation): width: int = Field(default=512, multiple_of=64, gt=0, description="The width of the resulting noise", ) height: int = Field(default=512, multiple_of=64, gt=0, description="The height of the resulting noise", ) + + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["latents", "noise"], + }, + } + def invoke(self, context: InvocationContext) -> NoiseOutput: device = torch.device(choose_torch_device()) noise = get_noise(self.width, self.height, device, self.seed) @@ -143,33 +153,37 @@ class TextToLatentsInvocation(BaseInvocation): progress_images: bool = Field(default=False, description="Whether or not to produce progress images during generation", ) # fmt: on + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["latents", "image"], + "type_hints": { + "model": "model" + } + }, + } + # TODO: pass this an emitter method or something? or a session for dispatching? def dispatch_progress( - self, context: InvocationContext, sample: Tensor, step: int + self, context: InvocationContext, intermediate_state: PipelineIntermediateState ) -> None: - # TODO: only output a preview image when requested - image = Generator.sample_to_lowres_estimated_image(sample) + if (context.services.queue.is_canceled(context.graph_execution_state_id)): + raise CanceledException - (width, height) = image.size - width *= 8 - height *= 8 + step = intermediate_state.step + if intermediate_state.predicted_original is not None: + # Some schedulers report not only the noisy latents at the current timestep, + # but also their estimate so far of what the de-noised latents will be. + sample = intermediate_state.predicted_original + else: + sample = intermediate_state.latents - dataURL = image_to_dataURL(image, image_format="JPEG") + diffusers_step_callback_adapter(sample, step, steps=self.steps, id=self.id, context=context) - context.services.events.emit_generator_progress( - context.graph_execution_state_id, - self.id, - { - "width": width, - "height": height, - "dataURL": dataURL - }, - step, - self.steps, - ) def get_model(self, model_manager: ModelManager) -> StableDiffusionGeneratorPipeline: - model_info = model_manager.get_model(self.model) + model_info = choose_model(model_manager, self.model) model_name = model_info['model_name'] model_hash = model_info['hash'] model: StableDiffusionGeneratorPipeline = model_info['model'] @@ -214,7 +228,7 @@ class TextToLatentsInvocation(BaseInvocation): noise = context.services.latents.get(self.noise.latents_name) def step_callback(state: PipelineIntermediateState): - self.dispatch_progress(context, state.latents, state.step) + self.dispatch_progress(context, state) model = self.get_model(context.services.model_manager) conditioning_data = self.get_conditioning_data(model) @@ -244,6 +258,17 @@ class LatentsToLatentsInvocation(TextToLatentsInvocation): type: Literal["l2l"] = "l2l" + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["latents"], + "type_hints": { + "model": "model" + } + }, + } + # Inputs latents: Optional[LatentsField] = Field(description="The latents to use as a base image") strength: float = Field(default=0.5, description="The strength of the latents to use") @@ -253,7 +278,7 @@ class LatentsToLatentsInvocation(TextToLatentsInvocation): latent = context.services.latents.get(self.latents.latents_name) def step_callback(state: PipelineIntermediateState): - self.dispatch_progress(context, state.latents, state.step) + self.dispatch_progress(context, state) model = self.get_model(context.services.model_manager) conditioning_data = self.get_conditioning_data(model) @@ -299,12 +324,23 @@ class LatentsToImageInvocation(BaseInvocation): latents: Optional[LatentsField] = Field(description="The latents to generate an image from") model: str = Field(default="", description="The model to use") + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["latents", "image"], + "type_hints": { + "model": "model" + } + }, + } + @torch.no_grad() def invoke(self, context: InvocationContext) -> ImageOutput: latents = context.services.latents.get(self.latents.latents_name) # TODO: this only really needs the vae - model_info = context.services.model_manager.get_model(self.model) + model_info = choose_model(context.services.model_manager, self.model) model: StableDiffusionGeneratorPipeline = model_info['model'] with torch.inference_mode(): diff --git a/invokeai/app/invocations/math.py b/invokeai/app/invocations/math.py index ecdcc834c7..e9d90b1d61 100644 --- a/invokeai/app/invocations/math.py +++ b/invokeai/app/invocations/math.py @@ -1,17 +1,26 @@ # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) -from datetime import datetime, timezone -from typing import Literal, Optional +from typing import Literal -import numpy -from PIL import Image, ImageFilter, ImageOps from pydantic import BaseModel, Field -from ..services.image_storage import ImageType -from ..services.invocation_services import InvocationServices +from invokeai.app.invocations.models.config import InvocationConfig + from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext +class MathInvocationConfig(BaseModel): + """Helper class to provide all math invocations with additional config""" + + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["math"], + } + } + + class IntOutput(BaseInvocationOutput): """An integer output""" #fmt: off @@ -20,7 +29,7 @@ class IntOutput(BaseInvocationOutput): #fmt: on -class AddInvocation(BaseInvocation): +class AddInvocation(BaseInvocation, MathInvocationConfig): """Adds two numbers""" #fmt: off type: Literal["add"] = "add" @@ -32,7 +41,7 @@ class AddInvocation(BaseInvocation): return IntOutput(a=self.a + self.b) -class SubtractInvocation(BaseInvocation): +class SubtractInvocation(BaseInvocation, MathInvocationConfig): """Subtracts two numbers""" #fmt: off type: Literal["sub"] = "sub" @@ -44,7 +53,7 @@ class SubtractInvocation(BaseInvocation): return IntOutput(a=self.a - self.b) -class MultiplyInvocation(BaseInvocation): +class MultiplyInvocation(BaseInvocation, MathInvocationConfig): """Multiplies two numbers""" #fmt: off type: Literal["mul"] = "mul" @@ -56,7 +65,7 @@ class MultiplyInvocation(BaseInvocation): return IntOutput(a=self.a * self.b) -class DivideInvocation(BaseInvocation): +class DivideInvocation(BaseInvocation, MathInvocationConfig): """Divides two numbers""" #fmt: off type: Literal["div"] = "div" diff --git a/invokeai/app/invocations/models/config.py b/invokeai/app/invocations/models/config.py new file mode 100644 index 0000000000..f53bbdda00 --- /dev/null +++ b/invokeai/app/invocations/models/config.py @@ -0,0 +1,54 @@ +from typing import Dict, List, Literal, TypedDict +from pydantic import BaseModel + + +# TODO: when we can upgrade to python 3.11, we can use the`NotRequired` type instead of `total=False` +class UIConfig(TypedDict, total=False): + type_hints: Dict[ + str, + Literal[ + "integer", + "float", + "boolean", + "string", + "enum", + "image", + "latents", + "model", + ], + ] + tags: List[str] + + +class CustomisedSchemaExtra(TypedDict): + ui: UIConfig + + +class InvocationConfig(BaseModel.Config): + """Customizes pydantic's BaseModel.Config class for use by Invocations. + + Provide `schema_extra` a `ui` dict to add hints for generated UIs. + + `tags` + - A list of strings, used to categorise invocations. + + `type_hints` + - A dict of field types which override the types in the invocation definition. + - Each key should be the name of one of the invocation's fields. + - Each value should be one of the valid types: + - `integer`, `float`, `boolean`, `string`, `enum`, `image`, `latents`, `model` + + ```python + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["stable-diffusion", "image"], + "type_hints": { + "initial_image": "image", + }, + }, + } + ``` + """ + + schema_extra: CustomisedSchemaExtra diff --git a/invokeai/app/invocations/reconstruct.py b/invokeai/app/invocations/reconstruct.py index 68449729d6..db0a28c075 100644 --- a/invokeai/app/invocations/reconstruct.py +++ b/invokeai/app/invocations/reconstruct.py @@ -2,6 +2,7 @@ from datetime import datetime, timezone from typing import Literal, Union from pydantic import Field +from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices @@ -18,6 +19,14 @@ class RestoreFaceInvocation(BaseInvocation): strength: float = Field(default=0.75, gt=0, le=1, description="The strength of the restoration" ) #fmt: on + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["restoration", "image"], + }, + } + def invoke(self, context: InvocationContext) -> ImageOutput: image = context.services.images.get( self.image.image_type, self.image.image_name diff --git a/invokeai/app/invocations/upscale.py b/invokeai/app/invocations/upscale.py index ea3221572e..c9aa86cc6d 100644 --- a/invokeai/app/invocations/upscale.py +++ b/invokeai/app/invocations/upscale.py @@ -4,6 +4,7 @@ from datetime import datetime, timezone from typing import Literal, Union from pydantic import Field +from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices @@ -22,6 +23,15 @@ class UpscaleInvocation(BaseInvocation): level: Literal[2, 4] = Field(default=2, description="The upscale level") #fmt: on + + # Schema customisation + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["upscaling", "image"], + }, + } + def invoke(self, context: InvocationContext) -> ImageOutput: image = context.services.images.get( self.image.image_type, self.image.image_name diff --git a/invokeai/app/invocations/util/get_model.py b/invokeai/app/invocations/util/get_model.py new file mode 100644 index 0000000000..d3484a0b9d --- /dev/null +++ b/invokeai/app/invocations/util/get_model.py @@ -0,0 +1,11 @@ +from invokeai.app.invocations.baseinvocation import InvocationContext +from invokeai.backend.model_management.model_manager import ModelManager + + +def choose_model(model_manager: ModelManager, model_name: str): + """Returns the default model if the `model_name` not a valid model, else returns the selected model.""" + if model_manager.valid_model(model_name): + return model_manager.get_model(model_name) + else: + print(f"* Warning: '{model_name}' is not a valid model name. Using default model instead.") + return model_manager.get_model() \ No newline at end of file From 1f2c1e14dbf45220bf46dfd5d05312ebab408a88 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 11 Apr 2023 08:24:48 +1000 Subject: [PATCH 13/14] fix(nodes): move InvocationConfig to baseinvocation.py --- invokeai/app/invocations/baseinvocation.py | 55 +++++++++++++++++++++- invokeai/app/invocations/cv.py | 3 +- invokeai/app/invocations/generate.py | 5 +- invokeai/app/invocations/image.py | 4 +- invokeai/app/invocations/latent.py | 3 +- invokeai/app/invocations/math.py | 4 +- invokeai/app/invocations/models/config.py | 54 --------------------- invokeai/app/invocations/reconstruct.py | 3 +- invokeai/app/invocations/upscale.py | 3 +- 9 files changed, 61 insertions(+), 73 deletions(-) delete mode 100644 invokeai/app/invocations/models/config.py diff --git a/invokeai/app/invocations/baseinvocation.py b/invokeai/app/invocations/baseinvocation.py index 72fe39ed0b..3590129b96 100644 --- a/invokeai/app/invocations/baseinvocation.py +++ b/invokeai/app/invocations/baseinvocation.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from inspect import signature -from typing import get_args, get_type_hints +from typing import get_args, get_type_hints, Dict, List, Literal, TypedDict from pydantic import BaseModel, Field @@ -76,3 +76,56 @@ class BaseInvocation(ABC, BaseModel): #fmt: off id: str = Field(description="The id of this node. Must be unique among all nodes.") #fmt: on + + +# TODO: figure out a better way to provide these hints +# TODO: when we can upgrade to python 3.11, we can use the`NotRequired` type instead of `total=False` +class UIConfig(TypedDict, total=False): + type_hints: Dict[ + str, + Literal[ + "integer", + "float", + "boolean", + "string", + "enum", + "image", + "latents", + "model", + ], + ] + tags: List[str] + + +class CustomisedSchemaExtra(TypedDict): + ui: UIConfig + + +class InvocationConfig(BaseModel.Config): + """Customizes pydantic's BaseModel.Config class for use by Invocations. + + Provide `schema_extra` a `ui` dict to add hints for generated UIs. + + `tags` + - A list of strings, used to categorise invocations. + + `type_hints` + - A dict of field types which override the types in the invocation definition. + - Each key should be the name of one of the invocation's fields. + - Each value should be one of the valid types: + - `integer`, `float`, `boolean`, `string`, `enum`, `image`, `latents`, `model` + + ```python + class Config(InvocationConfig): + schema_extra = { + "ui": { + "tags": ["stable-diffusion", "image"], + "type_hints": { + "initial_image": "image", + }, + }, + } + ``` + """ + + schema_extra: CustomisedSchemaExtra diff --git a/invokeai/app/invocations/cv.py b/invokeai/app/invocations/cv.py index adcb2405c4..9afbbbbcc9 100644 --- a/invokeai/app/invocations/cv.py +++ b/invokeai/app/invocations/cv.py @@ -6,10 +6,9 @@ import cv2 as cv import numpy from PIL import Image, ImageOps from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput diff --git a/invokeai/app/invocations/generate.py b/invokeai/app/invocations/generate.py index a07b1ac379..d0eeeae698 100644 --- a/invokeai/app/invocations/generate.py +++ b/invokeai/app/invocations/generate.py @@ -7,13 +7,10 @@ import numpy as np from torch import Tensor from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import ( - InvocationConfig, -) from invokeai.app.models.image import ImageField, ImageType from invokeai.app.invocations.util.get_model import choose_model -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput from ...backend.generator import Txt2Img, Img2Img, Inpaint, InvokeAIGenerator from ...backend.stable_diffusion import PipelineIntermediateState diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index 0f783b2541..cc5f6b53c7 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -7,11 +7,9 @@ import numpy from PIL import Image, ImageFilter, ImageOps from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import InvocationConfig - from ..models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices -from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext +from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig class PILInvocationConfig(BaseModel): diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index d59d681f40..2da6e451a9 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -4,7 +4,6 @@ from typing import Literal, Optional from pydantic import BaseModel, Field import torch -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.exceptions import CanceledException from invokeai.app.invocations.util.get_model import choose_model from invokeai.app.util.step_callback import diffusers_step_callback_adapter @@ -15,7 +14,7 @@ from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import Post from ...backend.image_util.seamless import configure_model_padding from ...backend.prompting.conditioning import get_uc_and_c_and_ec from ...backend.stable_diffusion.diffusers_pipeline import ConditioningData, StableDiffusionGeneratorPipeline -from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext +from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig import numpy as np from ..services.image_storage import ImageType from .baseinvocation import BaseInvocation, InvocationContext diff --git a/invokeai/app/invocations/math.py b/invokeai/app/invocations/math.py index e9d90b1d61..afb0e75377 100644 --- a/invokeai/app/invocations/math.py +++ b/invokeai/app/invocations/math.py @@ -4,9 +4,7 @@ from typing import Literal from pydantic import BaseModel, Field -from invokeai.app.invocations.models.config import InvocationConfig - -from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext +from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig class MathInvocationConfig(BaseModel): diff --git a/invokeai/app/invocations/models/config.py b/invokeai/app/invocations/models/config.py deleted file mode 100644 index f53bbdda00..0000000000 --- a/invokeai/app/invocations/models/config.py +++ /dev/null @@ -1,54 +0,0 @@ -from typing import Dict, List, Literal, TypedDict -from pydantic import BaseModel - - -# TODO: when we can upgrade to python 3.11, we can use the`NotRequired` type instead of `total=False` -class UIConfig(TypedDict, total=False): - type_hints: Dict[ - str, - Literal[ - "integer", - "float", - "boolean", - "string", - "enum", - "image", - "latents", - "model", - ], - ] - tags: List[str] - - -class CustomisedSchemaExtra(TypedDict): - ui: UIConfig - - -class InvocationConfig(BaseModel.Config): - """Customizes pydantic's BaseModel.Config class for use by Invocations. - - Provide `schema_extra` a `ui` dict to add hints for generated UIs. - - `tags` - - A list of strings, used to categorise invocations. - - `type_hints` - - A dict of field types which override the types in the invocation definition. - - Each key should be the name of one of the invocation's fields. - - Each value should be one of the valid types: - - `integer`, `float`, `boolean`, `string`, `enum`, `image`, `latents`, `model` - - ```python - class Config(InvocationConfig): - schema_extra = { - "ui": { - "tags": ["stable-diffusion", "image"], - "type_hints": { - "initial_image": "image", - }, - }, - } - ``` - """ - - schema_extra: CustomisedSchemaExtra diff --git a/invokeai/app/invocations/reconstruct.py b/invokeai/app/invocations/reconstruct.py index db0a28c075..f6df5a2254 100644 --- a/invokeai/app/invocations/reconstruct.py +++ b/invokeai/app/invocations/reconstruct.py @@ -2,11 +2,10 @@ from datetime import datetime, timezone from typing import Literal, Union from pydantic import Field -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput class RestoreFaceInvocation(BaseInvocation): diff --git a/invokeai/app/invocations/upscale.py b/invokeai/app/invocations/upscale.py index c9aa86cc6d..021f3569e8 100644 --- a/invokeai/app/invocations/upscale.py +++ b/invokeai/app/invocations/upscale.py @@ -4,11 +4,10 @@ from datetime import datetime, timezone from typing import Literal, Union from pydantic import Field -from invokeai.app.invocations.models.config import InvocationConfig from invokeai.app.models.image import ImageField, ImageType from ..services.invocation_services import InvocationServices -from .baseinvocation import BaseInvocation, InvocationContext +from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput From d923d1d66bc4ac0a0d7cf931442ed85b7c1f00ab Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 11 Apr 2023 11:50:28 +1000 Subject: [PATCH 14/14] fix(nodes): fix naming of CvInvocationConfig --- invokeai/app/invocations/cv.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/invokeai/app/invocations/cv.py b/invokeai/app/invocations/cv.py index 9afbbbbcc9..52e59b16ac 100644 --- a/invokeai/app/invocations/cv.py +++ b/invokeai/app/invocations/cv.py @@ -12,7 +12,7 @@ from .baseinvocation import BaseInvocation, InvocationContext, InvocationConfig from .image import ImageOutput -class CVInvocation(BaseModel): +class CvInvocationConfig(BaseModel): """Helper class to provide all OpenCV invocations with additional config""" # Schema customisation @@ -24,7 +24,7 @@ class CVInvocation(BaseModel): } -class CvInpaintInvocation(BaseInvocation, CVInvocation): +class CvInpaintInvocation(BaseInvocation, CvInvocationConfig): """Simple inpaint using opencv.""" #fmt: off type: Literal["cv_inpaint"] = "cv_inpaint"