From 566815931298d1abd76faea3c6429499e78bfaee Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Fri, 10 May 2024 10:42:34 +1000 Subject: [PATCH 1/2] fix(api): add cover image to update model response Fixes a bug where the image _appears_ to be reset when editing a model. See: https://old.reddit.com/r/StableDiffusion/comments/1cnx40d/invoke_42_control_layers_regional_guidance_w_text/l3asdej/ --- invokeai/app/api/routers/model_manager.py | 24 +++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 7bb0f23dc8d..2571c50507b 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -6,7 +6,7 @@ import shutil import traceback from copy import deepcopy -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Type from fastapi import Body, Path, Query, Response, UploadFile from fastapi.responses import FileResponse @@ -52,6 +52,13 @@ class ModelsList(BaseModel): model_config = ConfigDict(use_enum_values=True) +def add_cover_image_to_model_config(config: AnyModelConfig, dependencies: Type[ApiDependencies]) -> AnyModelConfig: + """Add a cover image URL to a model configuration.""" + cover_image = dependencies.invoker.services.model_images.get_url(config.key) + config.cover_image = cover_image + return config + + ############################################################################## # These are example inputs and outputs that are used in places where Swagger # is unable to generate a correct example. @@ -118,8 +125,7 @@ async def list_model_records( record_store.search_by_attr(model_type=model_type, model_name=model_name, model_format=model_format) ) for model in found_models: - cover_image = ApiDependencies.invoker.services.model_images.get_url(model.key) - model.cover_image = cover_image + model = add_cover_image_to_model_config(model, ApiDependencies) return ModelsList(models=found_models) @@ -160,12 +166,9 @@ async def get_model_record( key: str = Path(description="Key of the model record to fetch."), ) -> AnyModelConfig: """Get a model record""" - record_store = ApiDependencies.invoker.services.model_manager.store try: - config: AnyModelConfig = record_store.get_model(key) - cover_image = ApiDependencies.invoker.services.model_images.get_url(key) - config.cover_image = cover_image - return config + config = ApiDependencies.invoker.services.model_manager.store.get_model(key) + return add_cover_image_to_model_config(config, ApiDependencies) except UnknownModelException as e: raise HTTPException(status_code=404, detail=str(e)) @@ -294,14 +297,15 @@ async def update_model_record( installer = ApiDependencies.invoker.services.model_manager.install try: record_store.update_model(key, changes=changes) - model_response: AnyModelConfig = installer.sync_model_path(key) + config = installer.sync_model_path(key) + config = add_cover_image_to_model_config(config, ApiDependencies) logger.info(f"Updated model: {key}") except UnknownModelException as e: raise HTTPException(status_code=404, detail=str(e)) except ValueError as e: logger.error(str(e)) raise HTTPException(status_code=409, detail=str(e)) - return model_response + return config @model_manager_router.get( From a0dfddcefab56588960bb35208720d05a7923048 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Fri, 10 May 2024 11:17:35 +1000 Subject: [PATCH 2/2] fix(api): retain cover image when converting model to diffusers We need to retrieve and re-save the image, because a conversion to diffusers creates a new model record, with a new key. See: https://old.reddit.com/r/StableDiffusion/comments/1cnx40d/invoke_42_control_layers_regional_guidance_w_text/l3bv152/ --- invokeai/app/api/routers/model_manager.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 2571c50507b..1ba3e30e07d 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -16,6 +16,7 @@ from starlette.exceptions import HTTPException from typing_extensions import Annotated +from invokeai.app.services.model_images.model_images_common import ModelImageFileNotFoundException from invokeai.app.services.model_install import ModelInstallJob from invokeai.app.services.model_records import ( DuplicateModelException, @@ -652,6 +653,14 @@ async def convert_model( logger.error(str(e)) raise HTTPException(status_code=409, detail=str(e)) + # Update the model image if the model had one + try: + model_image = ApiDependencies.invoker.services.model_images.get(key) + ApiDependencies.invoker.services.model_images.save(model_image, new_key) + ApiDependencies.invoker.services.model_images.delete(key) + except ModelImageFileNotFoundException: + pass + # delete the original safetensors file installer.delete(key) @@ -659,7 +668,8 @@ async def convert_model( shutil.rmtree(cache_path) # return the config record for the new diffusers directory - new_config: AnyModelConfig = store.get_model(new_key) + new_config = store.get_model(new_key) + new_config = add_cover_image_to_model_config(new_config, ApiDependencies) return new_config