feat(app): cleanup of board queries

chore(ui): typegen

.github/workflows/python-checks.yml

@@ -62,7 +62,7 @@ jobs:
 
       - name: install ruff
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
-        run: pip install ruff==0.6.0
+        run: pip install ruff
         shell: bash
 
       - name: ruff check

docker/Dockerfile

@@ -55,7 +55,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 FROM node:20-slim AS web-builder
 ENV PNPM_HOME="/pnpm"
 ENV PATH="$PNPM_HOME:$PATH"
-RUN corepack use pnpm@8.x
 RUN corepack enable
 
 WORKDIR /build

docker/README.md

@@ -1,22 +1,20 @@
 # Invoke in Docker
 
-First things first:
-
-- Ensure that Docker can use your [NVIDIA][nvidia docker docs] or [AMD][amd docker docs] GPU.
-- This document assumes a Linux system, but should work similarly under Windows with WSL2.
+- Ensure that Docker can use the GPU on your system
+- This documentation assumes Linux, but should work similarly under Windows with WSL2
 - We don't recommend running Invoke in Docker on macOS at this time. It works, but very slowly.
 
-## Quickstart
+## Quickstart :lightning:
 
-No `docker compose`, no persistence, single command, using the official images:
+No `docker compose`, no persistence, just a simple one-liner using the official images:
 
-**CUDA (NVIDIA GPU):**
+**CUDA:**
 
 ```bash
 docker run --runtime=nvidia --gpus=all --publish 9090:9090 ghcr.io/invoke-ai/invokeai
 ```
 
-**ROCm (AMD GPU):**
+**ROCm:**
 
 ```bash
 docker run --device /dev/kfd --device /dev/dri --publish 9090:9090 ghcr.io/invoke-ai/invokeai:main-rocm
@@ -24,20 +22,12 @@ docker run --device /dev/kfd --device /dev/dri --publish 9090:9090 ghcr.io/invok
 
 Open `http://localhost:9090` in your browser once the container finishes booting, install some models, and generate away!
 
-### Data persistence
-
-To persist your generated images and downloaded models outside of the container, add a `--volume/-v` flag to the above command, e.g.:
-
-```bash
-docker run --volume /some/local/path:/invokeai {...etc...}
-```
-
-`/some/local/path/invokeai` will contain all your data.
-It can *usually* be reused between different installs of Invoke. Tread with caution and read the release notes!
+> [!TIP]
+> To persist your data (including downloaded models) outside of the container, add a `--volume/-v` flag to the above command, e.g.: `docker run --volume /some/local/path:/invokeai <...the rest of the command>`
 
 ## Customize the container
 
-The included `run.sh` script is a convenience wrapper around `docker compose`. It can be helpful for passing additional build arguments to `docker compose`. Alternatively, the familiar `docker compose` commands work just as well.
+We ship the `run.sh` script, which is a convenient wrapper around `docker compose` for cases where custom image build args are needed. Alternatively, the familiar `docker compose` commands work just as well.
 
 ```bash
 cd docker
@@ -48,14 +38,11 @@ cp .env.sample .env
 
 It will take a few minutes to build the image the first time. Once the application starts up, open `http://localhost:9090` in your browser to invoke!
 
->[!TIP]
->When using the `run.sh` script, the container will continue running after Ctrl+C. To shut it down, use the `docker compose down` command.
-
 ## Docker setup in detail
 
 #### Linux
 
-1. Ensure buildkit is enabled in the Docker daemon settings (`/etc/docker/daemon.json`)
+1. Ensure builkit is enabled in the Docker daemon settings (`/etc/docker/daemon.json`)
 2. Install the `docker compose` plugin using your package manager, or follow a [tutorial](https://docs.docker.com/compose/install/linux/#install-using-the-repository).
     - The deprecated `docker-compose` (hyphenated) CLI probably won't work. Update to a recent version.
 3. Ensure docker daemon is able to access the GPU.
@@ -111,7 +98,25 @@ GPU_DRIVER=cuda
 
 Any environment variables supported by InvokeAI can be set here. See the [Configuration docs](https://invoke-ai.github.io/InvokeAI/features/CONFIGURATION/) for further detail.
 
----
+## Even More Customizing!
 
-[nvidia docker docs]: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html
-[amd docker docs]: https://rocm.docs.amd.com/projects/install-on-linux/en/latest/how-to/docker.html
+See the `docker-compose.yml` file. The `command` instruction can be uncommented and used to run arbitrary startup commands. Some examples below.
+
+### Reconfigure the runtime directory
+
+Can be used to download additional models from the supported model list
+
+In conjunction with `INVOKEAI_ROOT` can be also used to initialize a runtime directory
+
+```yaml
+command:
+  - invokeai-configure
+  - --yes
+```
+
+Or install models:
+
+```yaml
+command:
+  - invokeai-model-install
+```

installer/templates/invoke.sh.in

@@ -17,7 +17,7 @@
 set -eu
 
 # Ensure we're in the correct folder in case user's CWD is somewhere else
-scriptdir=$(dirname $(readlink -f "$0"))
+scriptdir=$(dirname "$0")
 cd "$scriptdir"
 
 . .venv/bin/activate

invokeai/app/api/dependencies.py

@@ -1,6 +1,5 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
-import asyncio
 from logging import Logger
 
 import torch
@@ -32,8 +31,6 @@ from invokeai.app.services.session_processor.session_processor_default import (
 )
 from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue
 from invokeai.app.services.shared.sqlite.sqlite_util import init_db
-from invokeai.app.services.style_preset_images.style_preset_images_disk import StylePresetImageFileStorageDisk
-from invokeai.app.services.style_preset_records.style_preset_records_sqlite import SqliteStylePresetRecordsStorage
 from invokeai.app.services.urls.urls_default import LocalUrlService
 from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
@@ -66,12 +63,7 @@ class ApiDependencies:
     invoker: Invoker
 
     @staticmethod
-    def initialize(
-        config: InvokeAIAppConfig,
-        event_handler_id: int,
-        loop: asyncio.AbstractEventLoop,
-        logger: Logger = logger,
-    ) -> None:
+    def initialize(config: InvokeAIAppConfig, event_handler_id: int, logger: Logger = logger) -> None:
         logger.info(f"InvokeAI version {__version__}")
         logger.info(f"Root directory = {str(config.root_path)}")
 
@@ -82,7 +74,6 @@ class ApiDependencies:
         image_files = DiskImageFileStorage(f"{output_folder}/images")
 
         model_images_folder = config.models_path
-        style_presets_folder = config.style_presets_path
 
         db = init_db(config=config, logger=logger, image_files=image_files)
 
@@ -93,7 +84,7 @@ class ApiDependencies:
         board_images = BoardImagesService()
         board_records = SqliteBoardRecordStorage(db=db)
         boards = BoardService()
-        events = FastAPIEventService(event_handler_id, loop=loop)
+        events = FastAPIEventService(event_handler_id)
         bulk_download = BulkDownloadService()
         image_records = SqliteImageRecordStorage(db=db)
         images = ImageService()
@@ -118,8 +109,6 @@ class ApiDependencies:
         session_queue = SqliteSessionQueue(db=db)
         urls = LocalUrlService()
         workflow_records = SqliteWorkflowRecordsStorage(db=db)
-        style_preset_records = SqliteStylePresetRecordsStorage(db=db)
-        style_preset_image_files = StylePresetImageFileStorageDisk(style_presets_folder / "images")
 
         services = InvocationServices(
             board_image_records=board_image_records,
@@ -145,8 +134,6 @@ class ApiDependencies:
             workflow_records=workflow_records,
             tensors=tensors,
             conditioning=conditioning,
-            style_preset_records=style_preset_records,
-            style_preset_image_files=style_preset_image_files,
         )
 
         ApiDependencies.invoker = Invoker(services)

invokeai/app/api/routers/boards.py

@@ -5,7 +5,7 @@ from fastapi.routing import APIRouter
 from pydantic import BaseModel, Field
 
 from invokeai.app.api.dependencies import ApiDependencies
-from invokeai.app.services.board_records.board_records_common import BoardChanges
+from invokeai.app.services.board_records.board_records_common import BoardChanges, UncategorizedImageCounts
 from invokeai.app.services.boards.boards_common import BoardDTO
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 
@@ -146,3 +146,14 @@ async def list_all_board_image_names(
         board_id,
     )
     return image_names
+
+
+@boards_router.get(
+    "/uncategorized/counts",
+    operation_id="get_uncategorized_image_counts",
+    response_model=UncategorizedImageCounts,
+)
+async def get_uncategorized_image_counts() -> UncategorizedImageCounts:
+    """Gets count of images and assets for uncategorized images (images with no board assocation)"""
+
+    return ApiDependencies.invoker.services.board_records.get_uncategorized_image_counts()

invokeai/app/api/routers/images.py

@@ -218,8 +218,9 @@ async def get_image_workflow(
         raise HTTPException(status_code=404)
 
 
-@images_router.get(
+@images_router.api_route(
     "/i/{image_name}/full",
+    methods=["GET", "HEAD"],
     operation_id="get_image_full",
     response_class=Response,
     responses={
@@ -230,18 +231,6 @@ async def get_image_workflow(
         404: {"description": "Image not found"},
     },
 )
-@images_router.head(
-    "/i/{image_name}/full",
-    operation_id="get_image_full_head",
-    response_class=Response,
-    responses={
-        200: {
-            "description": "Return the full-resolution image",
-            "content": {"image/png": {}},
-        },
-        404: {"description": "Image not found"},
-    },
-)
 async def get_image_full(
     image_name: str = Path(description="The name of full-resolution image file to get"),
 ) -> Response:
@@ -253,7 +242,6 @@ async def get_image_full(
             content = f.read()
         response = Response(content, media_type="image/png")
         response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
-        response.headers["Content-Disposition"] = f'inline; filename="{image_name}"'
         return response
     except Exception:
         raise HTTPException(status_code=404)

invokeai/app/api/routers/model_manager.py

@@ -6,7 +6,7 @@ import pathlib
 import traceback
 from copy import deepcopy
 from tempfile import TemporaryDirectory
-from typing import List, Optional, Type
+from typing import Any, Dict, List, Optional, Type
 
 from fastapi import Body, Path, Query, Response, UploadFile
 from fastapi.responses import FileResponse, HTMLResponse
@@ -430,11 +430,13 @@ async def delete_model_image(
 async def install_model(
     source: str = Query(description="Model source to install, can be a local path, repo_id, or remote URL"),
     inplace: Optional[bool] = Query(description="Whether or not to install a local model in place", default=False),
-    access_token: Optional[str] = Query(description="access token for the remote resource", default=None),
-    config: ModelRecordChanges = Body(
-        description="Object containing fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
+    # TODO(MM2): Can we type this?
+    config: Optional[Dict[str, Any]] = Body(
+        description="Dict of fields that override auto-probed values in the model config record, such as name, description and prediction_type ",
+        default=None,
         example={"name": "string", "description": "string"},
     ),
+    access_token: Optional[str] = None,
 ) -> ModelInstallJob:
     """Install a model using a string identifier.
 
@@ -449,9 +451,8 @@ async def install_model(
        - model/name:fp16:path/to/model.safetensors
        - model/name::path/to/model.safetensors
 
-    `config` is a ModelRecordChanges object. Fields in this object will override
-    the ones that are probed automatically. Pass an empty object to accept
-    all the defaults.
+    `config` is an optional dict containing model configuration values that will override
+    the ones that are probed automatically.
 
     `access_token` is an optional access token for use with Urls that require
     authentication.
@@ -736,7 +737,7 @@ async def convert_model(
         # write the converted file to the convert path
         raw_model = converted_model.model
         assert hasattr(raw_model, "save_pretrained")
-        raw_model.save_pretrained(convert_path)  # type: ignore
+        raw_model.save_pretrained(convert_path)
         assert convert_path.exists()
 
         # temporarily rename the original safetensors file so that there is no naming conflict
@@ -749,12 +750,12 @@ async def convert_model(
         try:
             new_key = installer.install_path(
                 convert_path,
-                config=ModelRecordChanges(
-                    name=original_name,
-                    description=model_config.description,
-                    hash=model_config.hash,
-                    source=model_config.source,
-                ),
+                config={
+                    "name": original_name,
+                    "description": model_config.description,
+                    "hash": model_config.hash,
+                    "source": model_config.source,
+                },
             )
         except Exception as e:
             logger.error(str(e))

invokeai/app/api/routers/session_queue.py

@@ -11,7 +11,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
-    CancelByOriginResult,
     ClearResult,
     EnqueueBatchResult,
     PruneResult,
@@ -106,19 +105,6 @@ async def cancel_by_batch_ids(
     return ApiDependencies.invoker.services.session_queue.cancel_by_batch_ids(queue_id=queue_id, batch_ids=batch_ids)
 
 
-@session_queue_router.put(
-    "/{queue_id}/cancel_by_origin",
-    operation_id="cancel_by_origin",
-    responses={200: {"model": CancelByBatchIDsResult}},
-)
-async def cancel_by_origin(
-    queue_id: str = Path(description="The queue id to perform this operation on"),
-    origin: str = Query(description="The origin to cancel all queue items for"),
-) -> CancelByOriginResult:
-    """Immediately cancels all queue items with the given origin"""
-    return ApiDependencies.invoker.services.session_queue.cancel_by_origin(queue_id=queue_id, origin=origin)
-
-
 @session_queue_router.put(
     "/{queue_id}/clear",
     operation_id="clear",

invokeai/app/api/routers/style_presets.py

@@ -1,274 +0,0 @@
-import csv
-import io
-import json
-import traceback
-from typing import Optional
-
-import pydantic
-from fastapi import APIRouter, File, Form, HTTPException, Path, Response, UploadFile
-from fastapi.responses import FileResponse
-from PIL import Image
-from pydantic import BaseModel, Field
-
-from invokeai.app.api.dependencies import ApiDependencies
-from invokeai.app.api.routers.model_manager import IMAGE_MAX_AGE
-from invokeai.app.services.style_preset_images.style_preset_images_common import StylePresetImageFileNotFoundException
-from invokeai.app.services.style_preset_records.style_preset_records_common import (
-    InvalidPresetImportDataError,
-    PresetData,
-    PresetType,
-    StylePresetChanges,
-    StylePresetNotFoundError,
-    StylePresetRecordWithImage,
-    StylePresetWithoutId,
-    UnsupportedFileTypeError,
-    parse_presets_from_file,
-)
-
-
-class StylePresetFormData(BaseModel):
-    name: str = Field(description="Preset name")
-    positive_prompt: str = Field(description="Positive prompt")
-    negative_prompt: str = Field(description="Negative prompt")
-    type: PresetType = Field(description="Preset type")
-
-
-style_presets_router = APIRouter(prefix="/v1/style_presets", tags=["style_presets"])
-
-
-@style_presets_router.get(
-    "/i/{style_preset_id}",
-    operation_id="get_style_preset",
-    responses={
-        200: {"model": StylePresetRecordWithImage},
-    },
-)
-async def get_style_preset(
-    style_preset_id: str = Path(description="The style preset to get"),
-) -> StylePresetRecordWithImage:
-    """Gets a style preset"""
-    try:
-        image = ApiDependencies.invoker.services.style_preset_image_files.get_url(style_preset_id)
-        style_preset = ApiDependencies.invoker.services.style_preset_records.get(style_preset_id)
-        return StylePresetRecordWithImage(image=image, **style_preset.model_dump())
-    except StylePresetNotFoundError:
-        raise HTTPException(status_code=404, detail="Style preset not found")
-
-
-@style_presets_router.patch(
-    "/i/{style_preset_id}",
-    operation_id="update_style_preset",
-    responses={
-        200: {"model": StylePresetRecordWithImage},
-    },
-)
-async def update_style_preset(
-    image: Optional[UploadFile] = File(description="The image file to upload", default=None),
-    style_preset_id: str = Path(description="The id of the style preset to update"),
-    data: str = Form(description="The data of the style preset to update"),
-) -> StylePresetRecordWithImage:
-    """Updates a style preset"""
-    if image is not None:
-        if not image.content_type or not image.content_type.startswith("image"):
-            raise HTTPException(status_code=415, detail="Not an image")
-
-        contents = await image.read()
-        try:
-            pil_image = Image.open(io.BytesIO(contents))
-
-        except Exception:
-            ApiDependencies.invoker.services.logger.error(traceback.format_exc())
-            raise HTTPException(status_code=415, detail="Failed to read image")
-
-        try:
-            ApiDependencies.invoker.services.style_preset_image_files.save(style_preset_id, pil_image)
-        except ValueError as e:
-            raise HTTPException(status_code=409, detail=str(e))
-    else:
-        try:
-            ApiDependencies.invoker.services.style_preset_image_files.delete(style_preset_id)
-        except StylePresetImageFileNotFoundException:
-            pass
-
-    try:
-        parsed_data = json.loads(data)
-        validated_data = StylePresetFormData(**parsed_data)
-
-        name = validated_data.name
-        type = validated_data.type
-        positive_prompt = validated_data.positive_prompt
-        negative_prompt = validated_data.negative_prompt
-
-    except pydantic.ValidationError:
-        raise HTTPException(status_code=400, detail="Invalid preset data")
-
-    preset_data = PresetData(positive_prompt=positive_prompt, negative_prompt=negative_prompt)
-    changes = StylePresetChanges(name=name, preset_data=preset_data, type=type)
-
-    style_preset_image = ApiDependencies.invoker.services.style_preset_image_files.get_url(style_preset_id)
-    style_preset = ApiDependencies.invoker.services.style_preset_records.update(
-        style_preset_id=style_preset_id, changes=changes
-    )
-    return StylePresetRecordWithImage(image=style_preset_image, **style_preset.model_dump())
-
-
-@style_presets_router.delete(
-    "/i/{style_preset_id}",
-    operation_id="delete_style_preset",
-)
-async def delete_style_preset(
-    style_preset_id: str = Path(description="The style preset to delete"),
-) -> None:
-    """Deletes a style preset"""
-    try:
-        ApiDependencies.invoker.services.style_preset_image_files.delete(style_preset_id)
-    except StylePresetImageFileNotFoundException:
-        pass
-
-    ApiDependencies.invoker.services.style_preset_records.delete(style_preset_id)
-
-
-@style_presets_router.post(
-    "/",
-    operation_id="create_style_preset",
-    responses={
-        200: {"model": StylePresetRecordWithImage},
-    },
-)
-async def create_style_preset(
-    image: Optional[UploadFile] = File(description="The image file to upload", default=None),
-    data: str = Form(description="The data of the style preset to create"),
-) -> StylePresetRecordWithImage:
-    """Creates a style preset"""
-
-    try:
-        parsed_data = json.loads(data)
-        validated_data = StylePresetFormData(**parsed_data)
-
-        name = validated_data.name
-        type = validated_data.type
-        positive_prompt = validated_data.positive_prompt
-        negative_prompt = validated_data.negative_prompt
-
-    except pydantic.ValidationError:
-        raise HTTPException(status_code=400, detail="Invalid preset data")
-
-    preset_data = PresetData(positive_prompt=positive_prompt, negative_prompt=negative_prompt)
-    style_preset = StylePresetWithoutId(name=name, preset_data=preset_data, type=type)
-    new_style_preset = ApiDependencies.invoker.services.style_preset_records.create(style_preset=style_preset)
-
-    if image is not None:
-        if not image.content_type or not image.content_type.startswith("image"):
-            raise HTTPException(status_code=415, detail="Not an image")
-
-        contents = await image.read()
-        try:
-            pil_image = Image.open(io.BytesIO(contents))
-
-        except Exception:
-            ApiDependencies.invoker.services.logger.error(traceback.format_exc())
-            raise HTTPException(status_code=415, detail="Failed to read image")
-
-        try:
-            ApiDependencies.invoker.services.style_preset_image_files.save(new_style_preset.id, pil_image)
-        except ValueError as e:
-            raise HTTPException(status_code=409, detail=str(e))
-
-    preset_image = ApiDependencies.invoker.services.style_preset_image_files.get_url(new_style_preset.id)
-    return StylePresetRecordWithImage(image=preset_image, **new_style_preset.model_dump())
-
-
-@style_presets_router.get(
-    "/",
-    operation_id="list_style_presets",
-    responses={
-        200: {"model": list[StylePresetRecordWithImage]},
-    },
-)
-async def list_style_presets() -> list[StylePresetRecordWithImage]:
-    """Gets a page of style presets"""
-    style_presets_with_image: list[StylePresetRecordWithImage] = []
-    style_presets = ApiDependencies.invoker.services.style_preset_records.get_many()
-    for preset in style_presets:
-        image = ApiDependencies.invoker.services.style_preset_image_files.get_url(preset.id)
-        style_preset_with_image = StylePresetRecordWithImage(image=image, **preset.model_dump())
-        style_presets_with_image.append(style_preset_with_image)
-
-    return style_presets_with_image
-
-
-@style_presets_router.get(
-    "/i/{style_preset_id}/image",
-    operation_id="get_style_preset_image",
-    responses={
-        200: {
-            "description": "The style preset image was fetched successfully",
-        },
-        400: {"description": "Bad request"},
-        404: {"description": "The style preset image could not be found"},
-    },
-    status_code=200,
-)
-async def get_style_preset_image(
-    style_preset_id: str = Path(description="The id of the style preset image to get"),
-) -> FileResponse:
-    """Gets an image file that previews the model"""
-
-    try:
-        path = ApiDependencies.invoker.services.style_preset_image_files.get_path(style_preset_id)
-
-        response = FileResponse(
-            path,
-            media_type="image/png",
-            filename=style_preset_id + ".png",
-            content_disposition_type="inline",
-        )
-        response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
-        return response
-    except Exception:
-        raise HTTPException(status_code=404)
-
-
-@style_presets_router.get(
-    "/export",
-    operation_id="export_style_presets",
-    responses={200: {"content": {"text/csv": {}}, "description": "A CSV file with the requested data."}},
-    status_code=200,
-)
-async def export_style_presets():
-    # Create an in-memory stream to store the CSV data
-    output = io.StringIO()
-    writer = csv.writer(output)
-
-    # Write the header
-    writer.writerow(["name", "prompt", "negative_prompt"])
-
-    style_presets = ApiDependencies.invoker.services.style_preset_records.get_many(type=PresetType.User)
-
-    for preset in style_presets:
-        writer.writerow([preset.name, preset.preset_data.positive_prompt, preset.preset_data.negative_prompt])
-
-    csv_data = output.getvalue()
-    output.close()
-
-    return Response(
-        content=csv_data,
-        media_type="text/csv",
-        headers={"Content-Disposition": "attachment; filename=prompt_templates.csv"},
-    )
-
-
-@style_presets_router.post(
-    "/import",
-    operation_id="import_style_presets",
-)
-async def import_style_presets(file: UploadFile = File(description="The file to import")):
-    try:
-        style_presets = await parse_presets_from_file(file)
-        ApiDependencies.invoker.services.style_preset_records.create_many(style_presets)
-    except InvalidPresetImportDataError as e:
-        ApiDependencies.invoker.services.logger.error(traceback.format_exc())
-        raise HTTPException(status_code=400, detail=str(e))
-    except UnsupportedFileTypeError as e:
-        ApiDependencies.invoker.services.logger.error(traceback.format_exc())
-        raise HTTPException(status_code=415, detail=str(e))

invokeai/app/api_app.py

@@ -30,7 +30,6 @@ from invokeai.app.api.routers import (
     images,
     model_manager,
     session_queue,
-    style_presets,
     utilities,
     workflows,
 )
@@ -56,13 +55,11 @@ mimetypes.add_type("text/css", ".css")
 torch_device_name = TorchDevice.get_torch_device_name()
 logger.info(f"Using torch device: {torch_device_name}")
 
-loop = asyncio.new_event_loop()
-
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Add startup event to load dependencies
-    ApiDependencies.initialize(config=app_config, event_handler_id=event_handler_id, loop=loop, logger=logger)
+    ApiDependencies.initialize(config=app_config, event_handler_id=event_handler_id, logger=logger)
     yield
     # Shut down threads
     ApiDependencies.shutdown()
@@ -109,7 +106,6 @@ app.include_router(board_images.board_images_router, prefix="/api")
 app.include_router(app_info.app_router, prefix="/api")
 app.include_router(session_queue.session_queue_router, prefix="/api")
 app.include_router(workflows.workflows_router, prefix="/api")
-app.include_router(style_presets.style_presets_router, prefix="/api")
 
 app.openapi = get_openapi_func(app)
 
@@ -188,6 +184,8 @@ def invoke_api() -> None:
 
     check_cudnn(logger)
 
+    # Start our own event loop for eventing usage
+    loop = asyncio.new_event_loop()
     config = uvicorn.Config(
         app=app,
         host=app_config.host,

invokeai/app/invocations/compel.py

@@ -80,12 +80,12 @@ class CompelInvocation(BaseInvocation):
 
         with (
             # apply all patches while the model is on the target device
-            text_encoder_info.model_on_device() as (cached_weights, text_encoder),
+            text_encoder_info.model_on_device() as (model_state_dict, text_encoder),
             tokenizer_info as tokenizer,
             ModelPatcher.apply_lora_text_encoder(
                 text_encoder,
                 loras=_lora_loader(),
-                cached_weights=cached_weights,
+                model_state_dict=model_state_dict,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
             ModelPatcher.apply_clip_skip(text_encoder, self.clip.skipped_layers),
@@ -175,13 +175,13 @@ class SDXLPromptInvocationBase:
 
         with (
             # apply all patches while the model is on the target device
-            text_encoder_info.model_on_device() as (cached_weights, text_encoder),
+            text_encoder_info.model_on_device() as (state_dict, text_encoder),
             tokenizer_info as tokenizer,
             ModelPatcher.apply_lora(
                 text_encoder,
                 loras=_lora_loader(),
                 prefix=lora_prefix,
-                cached_weights=cached_weights,
+                model_state_dict=state_dict,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
             ModelPatcher.apply_clip_skip(text_encoder, clip_field.skipped_layers),

invokeai/app/invocations/controlnet_image_processors.py

@@ -21,8 +21,6 @@ from controlnet_aux import (
 from controlnet_aux.util import HWC3, ade_palette
 from PIL import Image
 from pydantic import BaseModel, Field, field_validator, model_validator
-from transformers import pipeline
-from transformers.pipelines import DepthEstimationPipeline
 
 from invokeai.app.invocations.baseinvocation import (
     BaseInvocation,
@@ -46,12 +44,13 @@ from invokeai.app.invocations.util import validate_begin_end_step, validate_weig
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES, heuristic_resize
 from invokeai.backend.image_util.canny import get_canny_edges
-from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import DepthAnythingPipeline
+from invokeai.backend.image_util.depth_anything import DEPTH_ANYTHING_MODELS, DepthAnythingDetector
 from invokeai.backend.image_util.dw_openpose import DWPOSE_MODELS, DWOpenposeDetector
 from invokeai.backend.image_util.hed import HEDProcessor
 from invokeai.backend.image_util.lineart import LineartProcessor
 from invokeai.backend.image_util.lineart_anime import LineartAnimeProcessor
 from invokeai.backend.image_util.util import np_to_pil, pil_to_np
+from invokeai.backend.util.devices import TorchDevice
 
 
 class ControlField(BaseModel):
@@ -593,14 +592,7 @@ class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
         return color_map
 
 
-DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small", "small_v2"]
-# DepthAnything V2 Small model is licensed under Apache 2.0 but not the base and large models.
-DEPTH_ANYTHING_MODELS = {
-    "large": "LiheYoung/depth-anything-large-hf",
-    "base": "LiheYoung/depth-anything-base-hf",
-    "small": "LiheYoung/depth-anything-small-hf",
-    "small_v2": "depth-anything/Depth-Anything-V2-Small-hf",
-}
+DEPTH_ANYTHING_MODEL_SIZES = Literal["large", "base", "small"]
 
 
 @invocation(
@@ -608,33 +600,28 @@ DEPTH_ANYTHING_MODELS = {
     title="Depth Anything Processor",
     tags=["controlnet", "depth", "depth anything"],
     category="controlnet",
-    version="1.1.3",
+    version="1.1.2",
 )
 class DepthAnythingImageProcessorInvocation(ImageProcessorInvocation):
     """Generates a depth map based on the Depth Anything algorithm"""
 
     model_size: DEPTH_ANYTHING_MODEL_SIZES = InputField(
-        default="small_v2", description="The size of the depth model to use"
+        default="small", description="The size of the depth model to use"
     )
     resolution: int = InputField(default=512, ge=1, description=FieldDescriptions.image_res)
 
     def run_processor(self, image: Image.Image) -> Image.Image:
-        def load_depth_anything(model_path: Path):
-            depth_anything_pipeline = pipeline(model=str(model_path), task="depth-estimation", local_files_only=True)
-            assert isinstance(depth_anything_pipeline, DepthEstimationPipeline)
-            return DepthAnythingPipeline(depth_anything_pipeline)
+        def loader(model_path: Path):
+            return DepthAnythingDetector.load_model(
+                model_path, model_size=self.model_size, device=TorchDevice.choose_torch_device()
+            )
 
         with self._context.models.load_remote_model(
-            source=DEPTH_ANYTHING_MODELS[self.model_size], loader=load_depth_anything
-        ) as depth_anything_detector:
-            assert isinstance(depth_anything_detector, DepthAnythingPipeline)
-            depth_map = depth_anything_detector.generate_depth(image)
-
-            # Resizing to user target specified size
-            new_height = int(image.size[1] * (self.resolution / image.size[0]))
-            depth_map = depth_map.resize((self.resolution, new_height))
-
-            return depth_map
+            source=DEPTH_ANYTHING_MODELS[self.model_size], loader=loader
+        ) as model:
+            depth_anything_detector = DepthAnythingDetector(model, TorchDevice.choose_torch_device())
+            processed_image = depth_anything_detector(image=image, resolution=self.resolution)
+            return processed_image
 
 
 @invocation(

invokeai/app/invocations/create_gradient_mask.py

@@ -39,7 +39,7 @@ class GradientMaskOutput(BaseInvocationOutput):
     title="Create Gradient Mask",
     tags=["mask", "denoise"],
     category="latents",
-    version="1.2.0",
+    version="1.1.0",
 )
 class CreateGradientMaskInvocation(BaseInvocation):
     """Creates mask for denoising model run."""
@@ -93,7 +93,6 @@ class CreateGradientMaskInvocation(BaseInvocation):
 
             # redistribute blur so that the original edges are 0 and blur outwards to 1
             blur_tensor = (blur_tensor - 0.5) * 2
-            blur_tensor[blur_tensor < 0] = 0.0
 
             threshold = 1 - self.minimum_denoise
 

invokeai/app/invocations/denoise_latents.py

@@ -1,6 +1,5 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
 import inspect
-import os
 from contextlib import ExitStack
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
@@ -37,10 +36,9 @@ from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.controlnet_utils import prepare_control_image
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
 from invokeai.backend.lora import LoRAModelRaw
-from invokeai.backend.model_manager import BaseModelType, ModelVariantType
+from invokeai.backend.model_manager import BaseModelType
 from invokeai.backend.model_patcher import ModelPatcher
-from invokeai.backend.stable_diffusion import PipelineIntermediateState
-from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext, DenoiseInputs
+from invokeai.backend.stable_diffusion import PipelineIntermediateState, set_seamless
 from invokeai.backend.stable_diffusion.diffusers_pipeline import (
     ControlNetData,
     StableDiffusionGeneratorPipeline,
@@ -55,19 +53,6 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     TextConditioningData,
     TextConditioningRegions,
 )
-from invokeai.backend.stable_diffusion.diffusion.custom_atttention import CustomAttnProcessor2_0
-from invokeai.backend.stable_diffusion.diffusion_backend import StableDiffusionBackend
-from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
-from invokeai.backend.stable_diffusion.extensions.controlnet import ControlNetExt
-from invokeai.backend.stable_diffusion.extensions.freeu import FreeUExt
-from invokeai.backend.stable_diffusion.extensions.inpaint import InpaintExt
-from invokeai.backend.stable_diffusion.extensions.inpaint_model import InpaintModelExt
-from invokeai.backend.stable_diffusion.extensions.lora import LoRAExt
-from invokeai.backend.stable_diffusion.extensions.preview import PreviewExt
-from invokeai.backend.stable_diffusion.extensions.rescale_cfg import RescaleCFGExt
-from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
-from invokeai.backend.stable_diffusion.extensions.t2i_adapter import T2IAdapterExt
-from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager
 from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
 from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
 from invokeai.backend.util.devices import TorchDevice
@@ -329,10 +314,9 @@ class DenoiseLatentsInvocation(BaseInvocation):
         context: InvocationContext,
         positive_conditioning_field: Union[ConditioningField, list[ConditioningField]],
         negative_conditioning_field: Union[ConditioningField, list[ConditioningField]],
+        unet: UNet2DConditionModel,
         latent_height: int,
         latent_width: int,
-        device: torch.device,
-        dtype: torch.dtype,
         cfg_scale: float | list[float],
         steps: int,
         cfg_rescale_multiplier: float,
@@ -346,10 +330,10 @@ class DenoiseLatentsInvocation(BaseInvocation):
             uncond_list = [uncond_list]
 
         cond_text_embeddings, cond_text_embedding_masks = DenoiseLatentsInvocation._get_text_embeddings_and_masks(
-            cond_list, context, device, dtype
+            cond_list, context, unet.device, unet.dtype
         )
         uncond_text_embeddings, uncond_text_embedding_masks = DenoiseLatentsInvocation._get_text_embeddings_and_masks(
-            uncond_list, context, device, dtype
+            uncond_list, context, unet.device, unet.dtype
         )
 
         cond_text_embedding, cond_regions = DenoiseLatentsInvocation._concat_regional_text_embeddings(
@@ -357,14 +341,14 @@ class DenoiseLatentsInvocation(BaseInvocation):
             masks=cond_text_embedding_masks,
             latent_height=latent_height,
             latent_width=latent_width,
-            dtype=dtype,
+            dtype=unet.dtype,
         )
         uncond_text_embedding, uncond_regions = DenoiseLatentsInvocation._concat_regional_text_embeddings(
             text_conditionings=uncond_text_embeddings,
             masks=uncond_text_embedding_masks,
             latent_height=latent_height,
             latent_width=latent_width,
-            dtype=dtype,
+            dtype=unet.dtype,
         )
 
         if isinstance(cfg_scale, list):
@@ -471,65 +455,6 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         return controlnet_data
 
-    @staticmethod
-    def parse_controlnet_field(
-        exit_stack: ExitStack,
-        context: InvocationContext,
-        control_input: ControlField | list[ControlField] | None,
-        ext_manager: ExtensionsManager,
-    ) -> None:
-        # Normalize control_input to a list.
-        control_list: list[ControlField]
-        if isinstance(control_input, ControlField):
-            control_list = [control_input]
-        elif isinstance(control_input, list):
-            control_list = control_input
-        elif control_input is None:
-            control_list = []
-        else:
-            raise ValueError(f"Unexpected control_input type: {type(control_input)}")
-
-        for control_info in control_list:
-            model = exit_stack.enter_context(context.models.load(control_info.control_model))
-            ext_manager.add_extension(
-                ControlNetExt(
-                    model=model,
-                    image=context.images.get_pil(control_info.image.image_name),
-                    weight=control_info.control_weight,
-                    begin_step_percent=control_info.begin_step_percent,
-                    end_step_percent=control_info.end_step_percent,
-                    control_mode=control_info.control_mode,
-                    resize_mode=control_info.resize_mode,
-                )
-            )
-
-    @staticmethod
-    def parse_t2i_adapter_field(
-        exit_stack: ExitStack,
-        context: InvocationContext,
-        t2i_adapters: Optional[Union[T2IAdapterField, list[T2IAdapterField]]],
-        ext_manager: ExtensionsManager,
-    ) -> None:
-        if t2i_adapters is None:
-            return
-
-        # Handle the possibility that t2i_adapters could be a list or a single T2IAdapterField.
-        if isinstance(t2i_adapters, T2IAdapterField):
-            t2i_adapters = [t2i_adapters]
-
-        for t2i_adapter_field in t2i_adapters:
-            ext_manager.add_extension(
-                T2IAdapterExt(
-                    node_context=context,
-                    model_id=t2i_adapter_field.t2i_adapter_model,
-                    image=context.images.get_pil(t2i_adapter_field.image.image_name),
-                    weight=t2i_adapter_field.weight,
-                    begin_step_percent=t2i_adapter_field.begin_step_percent,
-                    end_step_percent=t2i_adapter_field.end_step_percent,
-                    resize_mode=t2i_adapter_field.resize_mode,
-                )
-            )
-
     def prep_ip_adapter_image_prompts(
         self,
         context: InvocationContext,
@@ -739,7 +664,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
         else:
             masked_latents = torch.where(mask < 0.5, 0.0, latents)
 
-        return mask, masked_latents, self.denoise_mask.gradient
+        return 1 - mask, masked_latents, self.denoise_mask.gradient
 
     @staticmethod
     def prepare_noise_and_latents(
@@ -782,157 +707,12 @@ class DenoiseLatentsInvocation(BaseInvocation):
 
         return seed, noise, latents
 
+    @torch.no_grad()
+    @SilenceWarnings()  # This quenches the NSFW nag from diffusers.
     def invoke(self, context: InvocationContext) -> LatentsOutput:
-        if os.environ.get("USE_MODULAR_DENOISE", False):
-            return self._new_invoke(context)
-        else:
-            return self._old_invoke(context)
-
-    @torch.no_grad()
-    @SilenceWarnings()  # This quenches the NSFW nag from diffusers.
-    def _new_invoke(self, context: InvocationContext) -> LatentsOutput:
-        ext_manager = ExtensionsManager(is_canceled=context.util.is_canceled)
-
-        device = TorchDevice.choose_torch_device()
-        dtype = TorchDevice.choose_torch_dtype()
-
-        seed, noise, latents = self.prepare_noise_and_latents(context, self.noise, self.latents)
-        _, _, latent_height, latent_width = latents.shape
-
-        conditioning_data = self.get_conditioning_data(
-            context=context,
-            positive_conditioning_field=self.positive_conditioning,
-            negative_conditioning_field=self.negative_conditioning,
-            cfg_scale=self.cfg_scale,
-            steps=self.steps,
-            latent_height=latent_height,
-            latent_width=latent_width,
-            device=device,
-            dtype=dtype,
-            # TODO: old backend, remove
-            cfg_rescale_multiplier=self.cfg_rescale_multiplier,
-        )
-
-        scheduler = get_scheduler(
-            context=context,
-            scheduler_info=self.unet.scheduler,
-            scheduler_name=self.scheduler,
-            seed=seed,
-        )
-
-        timesteps, init_timestep, scheduler_step_kwargs = self.init_scheduler(
-            scheduler,
-            seed=seed,
-            device=device,
-            steps=self.steps,
-            denoising_start=self.denoising_start,
-            denoising_end=self.denoising_end,
-        )
-
-        # get the unet's config so that we can pass the base to sd_step_callback()
-        unet_config = context.models.get_config(self.unet.unet.key)
-
-        ### preview
-        def step_callback(state: PipelineIntermediateState) -> None:
-            context.util.sd_step_callback(state, unet_config.base)
-
-        ext_manager.add_extension(PreviewExt(step_callback))
-
-        ### cfg rescale
-        if self.cfg_rescale_multiplier > 0:
-            ext_manager.add_extension(RescaleCFGExt(self.cfg_rescale_multiplier))
-
-        ### freeu
-        if self.unet.freeu_config:
-            ext_manager.add_extension(FreeUExt(self.unet.freeu_config))
-
-        ### lora
-        if self.unet.loras:
-            for lora_field in self.unet.loras:
-                ext_manager.add_extension(
-                    LoRAExt(
-                        node_context=context,
-                        model_id=lora_field.lora,
-                        weight=lora_field.weight,
-                    )
-                )
-        ### seamless
-        if self.unet.seamless_axes:
-            ext_manager.add_extension(SeamlessExt(self.unet.seamless_axes))
-
-        ### inpaint
-        mask, masked_latents, is_gradient_mask = self.prep_inpaint_mask(context, latents)
-        # NOTE: We used to identify inpainting models by inpecting the shape of the loaded UNet model weights. Now we
-        # use the ModelVariantType config. During testing, there was a report of a user with models that had an
-        # incorrect ModelVariantType value. Re-installing the model fixed the issue. If this issue turns out to be
-        # prevalent, we will have to revisit how we initialize the inpainting extensions.
-        if unet_config.variant == ModelVariantType.Inpaint:
-            ext_manager.add_extension(InpaintModelExt(mask, masked_latents, is_gradient_mask))
-        elif mask is not None:
-            ext_manager.add_extension(InpaintExt(mask, is_gradient_mask))
-
-        # Initialize context for modular denoise
-        latents = latents.to(device=device, dtype=dtype)
-        if noise is not None:
-            noise = noise.to(device=device, dtype=dtype)
-        denoise_ctx = DenoiseContext(
-            inputs=DenoiseInputs(
-                orig_latents=latents,
-                timesteps=timesteps,
-                init_timestep=init_timestep,
-                noise=noise,
-                seed=seed,
-                scheduler_step_kwargs=scheduler_step_kwargs,
-                conditioning_data=conditioning_data,
-                attention_processor_cls=CustomAttnProcessor2_0,
-            ),
-            unet=None,
-            scheduler=scheduler,
-        )
-
-        # context for loading additional models
-        with ExitStack() as exit_stack:
-            # later should be smth like:
-            # for extension_field in self.extensions:
-            #    ext = extension_field.to_extension(exit_stack, context, ext_manager)
-            #    ext_manager.add_extension(ext)
-            self.parse_controlnet_field(exit_stack, context, self.control, ext_manager)
-            self.parse_t2i_adapter_field(exit_stack, context, self.t2i_adapter, ext_manager)
-
-            # ext: t2i/ip adapter
-            ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx)
-
-            unet_info = context.models.load(self.unet.unet)
-            assert isinstance(unet_info.model, UNet2DConditionModel)
-            with (
-                unet_info.model_on_device() as (cached_weights, unet),
-                ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls),
-                # ext: controlnet
-                ext_manager.patch_extensions(denoise_ctx),
-                # ext: freeu, seamless, ip adapter, lora
-                ext_manager.patch_unet(unet, cached_weights),
-            ):
-                sd_backend = StableDiffusionBackend(unet, scheduler)
-                denoise_ctx.unet = unet
-                result_latents = sd_backend.latents_from_embeddings(denoise_ctx, ext_manager)
-
-        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
-        result_latents = result_latents.detach().to("cpu")
-        TorchDevice.empty_cache()
-
-        name = context.tensors.save(tensor=result_latents)
-        return LatentsOutput.build(latents_name=name, latents=result_latents, seed=None)
-
-    @torch.no_grad()
-    @SilenceWarnings()  # This quenches the NSFW nag from diffusers.
-    def _old_invoke(self, context: InvocationContext) -> LatentsOutput:
         seed, noise, latents = self.prepare_noise_and_latents(context, self.noise, self.latents)
 
         mask, masked_latents, gradient_mask = self.prep_inpaint_mask(context, latents)
-        # At this point, the mask ranges from 0 (leave unchanged) to 1 (inpaint).
-        # We invert the mask here for compatibility with the old backend implementation.
-        if mask is not None:
-            mask = 1 - mask
 
         # TODO(ryand): I have hard-coded `do_classifier_free_guidance=True` to mirror the behaviour of ControlNets,
         # below. Investigate whether this is appropriate.
@@ -975,14 +755,14 @@ class DenoiseLatentsInvocation(BaseInvocation):
         assert isinstance(unet_info.model, UNet2DConditionModel)
         with (
             ExitStack() as exit_stack,
-            unet_info.model_on_device() as (cached_weights, unet),
+            unet_info.model_on_device() as (model_state_dict, unet),
             ModelPatcher.apply_freeu(unet, self.unet.freeu_config),
-            SeamlessExt.static_patch_model(unet, self.unet.seamless_axes),  # FIXME
+            set_seamless(unet, self.unet.seamless_axes),  # FIXME
             # Apply the LoRA after unet has been moved to its target device for faster patching.
             ModelPatcher.apply_lora_unet(
                 unet,
                 loras=_lora_loader(),
-                cached_weights=cached_weights,
+                model_state_dict=model_state_dict,
             ),
         ):
             assert isinstance(unet, UNet2DConditionModel)
@@ -1008,8 +788,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
                 context=context,
                 positive_conditioning_field=self.positive_conditioning,
                 negative_conditioning_field=self.negative_conditioning,
-                device=unet.device,
-                dtype=unet.dtype,
+                unet=unet,
                 latent_height=latent_height,
                 latent_width=latent_width,
                 cfg_scale=self.cfg_scale,

invokeai/app/invocations/fields.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Any, Callable, Optional, Tuple
 
-from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, model_validator
+from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter
 from pydantic.fields import _Unset
 from pydantic_core import PydanticUndefined
 
@@ -48,7 +48,6 @@ class UIType(str, Enum, metaclass=MetaEnum):
     ControlNetModel = "ControlNetModelField"
     IPAdapterModel = "IPAdapterModelField"
     T2IAdapterModel = "T2IAdapterModelField"
-    SpandrelImageToImageModel = "SpandrelImageToImageModelField"
     # endregion
 
     # region Misc Field Types
@@ -135,7 +134,6 @@ class FieldDescriptions:
     sdxl_main_model = "SDXL Main model (UNet, VAE, CLIP1, CLIP2) to load"
     sdxl_refiner_model = "SDXL Refiner Main Modde (UNet, VAE, CLIP2) to load"
     onnx_main_model = "ONNX Main model (UNet, VAE, CLIP) to load"
-    spandrel_image_to_image_model = "Image-to-Image model"
     lora_weight = "The weight at which the LoRA is applied to each model"
     compel_prompt = "Prompt to be parsed by Compel to create a conditioning tensor"
     raw_prompt = "Raw prompt text (no parsing)"
@@ -242,31 +240,6 @@ class ConditioningField(BaseModel):
     )
 
 
-class BoundingBoxField(BaseModel):
-    """A bounding box primitive value."""
-
-    x_min: int = Field(ge=0, description="The minimum x-coordinate of the bounding box (inclusive).")
-    x_max: int = Field(ge=0, description="The maximum x-coordinate of the bounding box (exclusive).")
-    y_min: int = Field(ge=0, description="The minimum y-coordinate of the bounding box (inclusive).")
-    y_max: int = Field(ge=0, description="The maximum y-coordinate of the bounding box (exclusive).")
-
-    score: Optional[float] = Field(
-        default=None,
-        ge=0.0,
-        le=1.0,
-        description="The score associated with the bounding box. In the range [0, 1]. This value is typically set "
-        "when the bounding box was produced by a detector and has an associated confidence score.",
-    )
-
-    @model_validator(mode="after")
-    def check_coords(self):
-        if self.x_min > self.x_max:
-            raise ValueError(f"x_min ({self.x_min}) is greater than x_max ({self.x_max}).")
-        if self.y_min > self.y_max:
-            raise ValueError(f"y_min ({self.y_min}) is greater than y_max ({self.y_max}).")
-        return self
-
-
 class MetadataField(RootModel[dict[str, Any]]):
     """
     Pydantic model for metadata with custom root of type dict[str, Any].

invokeai/app/invocations/grounding_dino.py

@@ -1,100 +0,0 @@
-from pathlib import Path
-from typing import Literal
-
-import torch
-from PIL import Image
-from transformers import pipeline
-from transformers.pipelines import ZeroShotObjectDetectionPipeline
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.fields import BoundingBoxField, ImageField, InputField
-from invokeai.app.invocations.primitives import BoundingBoxCollectionOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.image_util.grounding_dino.detection_result import DetectionResult
-from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
-
-GroundingDinoModelKey = Literal["grounding-dino-tiny", "grounding-dino-base"]
-GROUNDING_DINO_MODEL_IDS: dict[GroundingDinoModelKey, str] = {
-    "grounding-dino-tiny": "IDEA-Research/grounding-dino-tiny",
-    "grounding-dino-base": "IDEA-Research/grounding-dino-base",
-}
-
-
-@invocation(
-    "grounding_dino",
-    title="Grounding DINO (Text Prompt Object Detection)",
-    tags=["prompt", "object detection"],
-    category="image",
-    version="1.0.0",
-)
-class GroundingDinoInvocation(BaseInvocation):
-    """Runs a Grounding DINO model. Performs zero-shot bounding-box object detection from a text prompt."""
-
-    # Reference:
-    # - https://arxiv.org/pdf/2303.05499
-    # - https://huggingface.co/docs/transformers/v4.43.3/en/model_doc/grounding-dino#grounded-sam
-    # - https://github.com/NielsRogge/Transformers-Tutorials/blob/a39f33ac1557b02ebfb191ea7753e332b5ca933f/Grounding%20DINO/GroundingDINO_with_Segment_Anything.ipynb
-
-    model: GroundingDinoModelKey = InputField(description="The Grounding DINO model to use.")
-    prompt: str = InputField(description="The prompt describing the object to segment.")
-    image: ImageField = InputField(description="The image to segment.")
-    detection_threshold: float = InputField(
-        description="The detection threshold for the Grounding DINO model. All detected bounding boxes with scores above this threshold will be returned.",
-        ge=0.0,
-        le=1.0,
-        default=0.3,
-    )
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> BoundingBoxCollectionOutput:
-        # The model expects a 3-channel RGB image.
-        image_pil = context.images.get_pil(self.image.image_name, mode="RGB")
-
-        detections = self._detect(
-            context=context, image=image_pil, labels=[self.prompt], threshold=self.detection_threshold
-        )
-
-        # Convert detections to BoundingBoxCollectionOutput.
-        bounding_boxes: list[BoundingBoxField] = []
-        for detection in detections:
-            bounding_boxes.append(
-                BoundingBoxField(
-                    x_min=detection.box.xmin,
-                    x_max=detection.box.xmax,
-                    y_min=detection.box.ymin,
-                    y_max=detection.box.ymax,
-                    score=detection.score,
-                )
-            )
-        return BoundingBoxCollectionOutput(collection=bounding_boxes)
-
-    @staticmethod
-    def _load_grounding_dino(model_path: Path):
-        grounding_dino_pipeline = pipeline(
-            model=str(model_path),
-            task="zero-shot-object-detection",
-            local_files_only=True,
-            # TODO(ryand): Setting the torch_dtype here doesn't work. Investigate whether fp16 is supported by the
-            # model, and figure out how to make it work in the pipeline.
-            # torch_dtype=TorchDevice.choose_torch_dtype(),
-        )
-        assert isinstance(grounding_dino_pipeline, ZeroShotObjectDetectionPipeline)
-        return GroundingDinoPipeline(grounding_dino_pipeline)
-
-    def _detect(
-        self,
-        context: InvocationContext,
-        image: Image.Image,
-        labels: list[str],
-        threshold: float = 0.3,
-    ) -> list[DetectionResult]:
-        """Use Grounding DINO to detect bounding boxes for a set of labels in an image."""
-        # TODO(ryand): I copied this "."-handling logic from the transformers example code. Test it and see if it
-        # actually makes a difference.
-        labels = [label if label.endswith(".") else label + "." for label in labels]
-
-        with context.models.load_remote_model(
-            source=GROUNDING_DINO_MODEL_IDS[self.model], loader=GroundingDinoInvocation._load_grounding_dino
-        ) as detector:
-            assert isinstance(detector, GroundingDinoPipeline)
-            return detector.detect(image=image, candidate_labels=labels, threshold=threshold)

invokeai/app/invocations/image.py

@@ -6,19 +6,13 @@ import cv2
 import numpy
 from PIL import Image, ImageChops, ImageFilter, ImageOps
 
-from invokeai.app.invocations.baseinvocation import (
-    BaseInvocation,
-    Classification,
-    invocation,
-    invocation_output,
-)
+from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
 from invokeai.app.invocations.constants import IMAGE_MODES
 from invokeai.app.invocations.fields import (
     ColorField,
     FieldDescriptions,
     ImageField,
     InputField,
-    OutputField,
     WithBoard,
     WithMetadata,
 )
@@ -1013,62 +1007,3 @@ class MaskFromIDInvocation(BaseInvocation, WithMetadata, WithBoard):
         image_dto = context.images.save(image=mask, image_category=ImageCategory.MASK)
 
         return ImageOutput.build(image_dto)
-
-
-@invocation_output("canvas_v2_mask_and_crop_output")
-class CanvasV2MaskAndCropOutput(ImageOutput):
-    offset_x: int = OutputField(description="The x offset of the image, after cropping")
-    offset_y: int = OutputField(description="The y offset of the image, after cropping")
-
-
-@invocation(
-    "canvas_v2_mask_and_crop",
-    title="Canvas V2 Mask and Crop",
-    tags=["image", "mask", "id"],
-    category="image",
-    version="1.0.0",
-    classification=Classification.Prototype,
-)
-class CanvasV2MaskAndCropInvocation(BaseInvocation, WithMetadata, WithBoard):
-    """Handles Canvas V2 image output masking and cropping"""
-
-    source_image: ImageField | None = InputField(
-        default=None,
-        description="The source image onto which the masked generated image is pasted. If omitted, the masked generated image is returned with transparency.",
-    )
-    generated_image: ImageField = InputField(description="The image to apply the mask to")
-    mask: ImageField = InputField(description="The mask to apply")
-    mask_blur: int = InputField(default=0, ge=0, description="The amount to blur the mask by")
-
-    def _prepare_mask(self, mask: Image.Image) -> Image.Image:
-        mask_array = numpy.array(mask)
-        kernel = numpy.ones((self.mask_blur, self.mask_blur), numpy.uint8)
-        dilated_mask_array = cv2.erode(mask_array, kernel, iterations=3)
-        dilated_mask = Image.fromarray(dilated_mask_array)
-        if self.mask_blur > 0:
-            mask = dilated_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))
-        return ImageOps.invert(mask.convert("L"))
-
-    def invoke(self, context: InvocationContext) -> CanvasV2MaskAndCropOutput:
-        mask = self._prepare_mask(context.images.get_pil(self.mask.image_name))
-
-        if self.source_image:
-            generated_image = context.images.get_pil(self.generated_image.image_name)
-            source_image = context.images.get_pil(self.source_image.image_name)
-            source_image.paste(generated_image, (0, 0), mask)
-            image_dto = context.images.save(image=source_image)
-        else:
-            generated_image = context.images.get_pil(self.generated_image.image_name)
-            generated_image.putalpha(mask)
-            image_dto = context.images.save(image=generated_image)
-
-        # bbox = image.getbbox()
-        # image = image.crop(bbox)
-
-        return CanvasV2MaskAndCropOutput(
-            image=ImageField(image_name=image_dto.image_name),
-            offset_x=0,
-            offset_y=0,
-            width=image_dto.width,
-            height=image_dto.height,
-        )

invokeai/app/invocations/latents_to_image.py

@@ -24,7 +24,7 @@ from invokeai.app.invocations.fields import (
 from invokeai.app.invocations.model import VAEField
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
+from invokeai.backend.stable_diffusion import set_seamless
 from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params
 from invokeai.backend.util.devices import TorchDevice
 
@@ -59,7 +59,7 @@ class LatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
 
         vae_info = context.models.load(self.vae.vae)
         assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny))
-        with SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), vae_info as vae:
+        with set_seamless(vae_info.model, self.vae.seamless_axes), vae_info as vae:
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
             latents = latents.to(vae.device)
             if self.fp32:

invokeai/app/invocations/mask.py

@@ -1,10 +1,9 @@
 import numpy as np
 import torch
-from PIL import Image
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, InvocationContext, invocation
-from invokeai.app.invocations.fields import ImageField, InputField, TensorField, WithBoard, WithMetadata
-from invokeai.app.invocations.primitives import ImageOutput, MaskOutput
+from invokeai.app.invocations.fields import ImageField, InputField, TensorField, WithMetadata
+from invokeai.app.invocations.primitives import MaskOutput
 
 
 @invocation(
@@ -119,27 +118,3 @@ class ImageMaskToTensorInvocation(BaseInvocation, WithMetadata):
             height=mask.shape[1],
             width=mask.shape[2],
         )
-
-
-@invocation(
-    "tensor_mask_to_image",
-    title="Tensor Mask to Image",
-    tags=["mask"],
-    category="mask",
-    version="1.0.0",
-)
-class MaskTensorToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
-    """Convert a mask tensor to an image."""
-
-    mask: TensorField = InputField(description="The mask tensor to convert.")
-
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        mask = context.tensors.load(self.mask.tensor_name)
-        # Ensure that the mask is binary.
-        if mask.dtype != torch.bool:
-            mask = mask > 0.5
-        mask_np = (mask.float() * 255).byte().cpu().numpy()
-
-        mask_pil = Image.fromarray(mask_np, mode="L")
-        image_dto = context.images.save(image=mask_pil)
-        return ImageOutput.build(image_dto)

invokeai/app/invocations/primitives.py

@@ -7,7 +7,6 @@ import torch
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
 from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
 from invokeai.app.invocations.fields import (
-    BoundingBoxField,
     ColorField,
     ConditioningField,
     DenoiseMaskField,
@@ -470,42 +469,3 @@ class ConditioningCollectionInvocation(BaseInvocation):
 
 
 # endregion
-
-# region BoundingBox
-
-
-@invocation_output("bounding_box_output")
-class BoundingBoxOutput(BaseInvocationOutput):
-    """Base class for nodes that output a single bounding box"""
-
-    bounding_box: BoundingBoxField = OutputField(description="The output bounding box.")
-
-
-@invocation_output("bounding_box_collection_output")
-class BoundingBoxCollectionOutput(BaseInvocationOutput):
-    """Base class for nodes that output a collection of bounding boxes"""
-
-    collection: list[BoundingBoxField] = OutputField(description="The output bounding boxes.", title="Bounding Boxes")
-
-
-@invocation(
-    "bounding_box",
-    title="Bounding Box",
-    tags=["primitives", "segmentation", "collection", "bounding box"],
-    category="primitives",
-    version="1.0.0",
-)
-class BoundingBoxInvocation(BaseInvocation):
-    """Create a bounding box manually by supplying box coordinates"""
-
-    x_min: int = InputField(default=0, description="x-coordinate of the bounding box's top left vertex")
-    y_min: int = InputField(default=0, description="y-coordinate of the bounding box's top left vertex")
-    x_max: int = InputField(default=0, description="x-coordinate of the bounding box's bottom right vertex")
-    y_max: int = InputField(default=0, description="y-coordinate of the bounding box's bottom right vertex")
-
-    def invoke(self, context: InvocationContext) -> BoundingBoxOutput:
-        bounding_box = BoundingBoxField(x_min=self.x_min, y_min=self.y_min, x_max=self.x_max, y_max=self.y_max)
-        return BoundingBoxOutput(bounding_box=bounding_box)
-
-
-# endregion

invokeai/app/invocations/segment_anything.py

@@ -1,161 +0,0 @@
-from pathlib import Path
-from typing import Literal
-
-import numpy as np
-import torch
-from PIL import Image
-from transformers import AutoModelForMaskGeneration, AutoProcessor
-from transformers.models.sam import SamModel
-from transformers.models.sam.processing_sam import SamProcessor
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.fields import BoundingBoxField, ImageField, InputField, TensorField
-from invokeai.app.invocations.primitives import MaskOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.image_util.segment_anything.mask_refinement import mask_to_polygon, polygon_to_mask
-from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
-
-SegmentAnythingModelKey = Literal["segment-anything-base", "segment-anything-large", "segment-anything-huge"]
-SEGMENT_ANYTHING_MODEL_IDS: dict[SegmentAnythingModelKey, str] = {
-    "segment-anything-base": "facebook/sam-vit-base",
-    "segment-anything-large": "facebook/sam-vit-large",
-    "segment-anything-huge": "facebook/sam-vit-huge",
-}
-
-
-@invocation(
-    "segment_anything",
-    title="Segment Anything",
-    tags=["prompt", "segmentation"],
-    category="segmentation",
-    version="1.0.0",
-)
-class SegmentAnythingInvocation(BaseInvocation):
-    """Runs a Segment Anything Model."""
-
-    # Reference:
-    # - https://arxiv.org/pdf/2304.02643
-    # - https://huggingface.co/docs/transformers/v4.43.3/en/model_doc/grounding-dino#grounded-sam
-    # - https://github.com/NielsRogge/Transformers-Tutorials/blob/a39f33ac1557b02ebfb191ea7753e332b5ca933f/Grounding%20DINO/GroundingDINO_with_Segment_Anything.ipynb
-
-    model: SegmentAnythingModelKey = InputField(description="The Segment Anything model to use.")
-    image: ImageField = InputField(description="The image to segment.")
-    bounding_boxes: list[BoundingBoxField] = InputField(description="The bounding boxes to prompt the SAM model with.")
-    apply_polygon_refinement: bool = InputField(
-        description="Whether to apply polygon refinement to the masks. This will smooth the edges of the masks slightly and ensure that each mask consists of a single closed polygon (before merging).",
-        default=True,
-    )
-    mask_filter: Literal["all", "largest", "highest_box_score"] = InputField(
-        description="The filtering to apply to the detected masks before merging them into a final output.",
-        default="all",
-    )
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> MaskOutput:
-        # The models expect a 3-channel RGB image.
-        image_pil = context.images.get_pil(self.image.image_name, mode="RGB")
-
-        if len(self.bounding_boxes) == 0:
-            combined_mask = torch.zeros(image_pil.size[::-1], dtype=torch.bool)
-        else:
-            masks = self._segment(context=context, image=image_pil)
-            masks = self._filter_masks(masks=masks, bounding_boxes=self.bounding_boxes)
-
-            # masks contains bool values, so we merge them via max-reduce.
-            combined_mask, _ = torch.stack(masks).max(dim=0)
-
-        mask_tensor_name = context.tensors.save(combined_mask)
-        height, width = combined_mask.shape
-        return MaskOutput(mask=TensorField(tensor_name=mask_tensor_name), width=width, height=height)
-
-    @staticmethod
-    def _load_sam_model(model_path: Path):
-        sam_model = AutoModelForMaskGeneration.from_pretrained(
-            model_path,
-            local_files_only=True,
-            # TODO(ryand): Setting the torch_dtype here doesn't work. Investigate whether fp16 is supported by the
-            # model, and figure out how to make it work in the pipeline.
-            # torch_dtype=TorchDevice.choose_torch_dtype(),
-        )
-        assert isinstance(sam_model, SamModel)
-
-        sam_processor = AutoProcessor.from_pretrained(model_path, local_files_only=True)
-        assert isinstance(sam_processor, SamProcessor)
-        return SegmentAnythingPipeline(sam_model=sam_model, sam_processor=sam_processor)
-
-    def _segment(
-        self,
-        context: InvocationContext,
-        image: Image.Image,
-    ) -> list[torch.Tensor]:
-        """Use Segment Anything (SAM) to generate masks given an image + a set of bounding boxes."""
-        # Convert the bounding boxes to the SAM input format.
-        sam_bounding_boxes = [[bb.x_min, bb.y_min, bb.x_max, bb.y_max] for bb in self.bounding_boxes]
-
-        with (
-            context.models.load_remote_model(
-                source=SEGMENT_ANYTHING_MODEL_IDS[self.model], loader=SegmentAnythingInvocation._load_sam_model
-            ) as sam_pipeline,
-        ):
-            assert isinstance(sam_pipeline, SegmentAnythingPipeline)
-            masks = sam_pipeline.segment(image=image, bounding_boxes=sam_bounding_boxes)
-
-        masks = self._process_masks(masks)
-        if self.apply_polygon_refinement:
-            masks = self._apply_polygon_refinement(masks)
-
-        return masks
-
-    def _process_masks(self, masks: torch.Tensor) -> list[torch.Tensor]:
-        """Convert the tensor output from the Segment Anything model from a tensor of shape
-        [num_masks, channels, height, width] to a list of tensors of shape [height, width].
-        """
-        assert masks.dtype == torch.bool
-        # [num_masks, channels, height, width] -> [num_masks, height, width]
-        masks, _ = masks.max(dim=1)
-        # Split the first dimension into a list of masks.
-        return list(masks.cpu().unbind(dim=0))
-
-    def _apply_polygon_refinement(self, masks: list[torch.Tensor]) -> list[torch.Tensor]:
-        """Apply polygon refinement to the masks.
-
-        Convert each mask to a polygon, then back to a mask. This has the following effect:
-        - Smooth the edges of the mask slightly.
-        - Ensure that each mask consists of a single closed polygon
-            - Removes small mask pieces.
-            - Removes holes from the mask.
-        """
-        # Convert tensor masks to np masks.
-        np_masks = [mask.cpu().numpy().astype(np.uint8) for mask in masks]
-
-        # Apply polygon refinement.
-        for idx, mask in enumerate(np_masks):
-            shape = mask.shape
-            assert len(shape) == 2  # Assert length to satisfy type checker.
-            polygon = mask_to_polygon(mask)
-            mask = polygon_to_mask(polygon, shape)
-            np_masks[idx] = mask
-
-        # Convert np masks back to tensor masks.
-        masks = [torch.tensor(mask, dtype=torch.bool) for mask in np_masks]
-
-        return masks
-
-    def _filter_masks(self, masks: list[torch.Tensor], bounding_boxes: list[BoundingBoxField]) -> list[torch.Tensor]:
-        """Filter the detected masks based on the specified mask filter."""
-        assert len(masks) == len(bounding_boxes)
-
-        if self.mask_filter == "all":
-            return masks
-        elif self.mask_filter == "largest":
-            # Find the largest mask.
-            return [max(masks, key=lambda x: float(x.sum()))]
-        elif self.mask_filter == "highest_box_score":
-            # Find the index of the bounding box with the highest score.
-            # Note that we fallback to -1.0 if the score is None. This is mainly to satisfy the type checker. In most
-            # cases the scores should all be non-None when using this filtering mode. That being said, -1.0 is a
-            # reasonable fallback since the expected score range is [0.0, 1.0].
-            max_score_idx = max(range(len(bounding_boxes)), key=lambda i: bounding_boxes[i].score or -1.0)
-            return [masks[max_score_idx]]
-        else:
-            raise ValueError(f"Invalid mask filter: {self.mask_filter}")

invokeai/app/invocations/spandrel_image_to_image.py

@@ -1,253 +0,0 @@
-from typing import Callable
-
-import numpy as np
-import torch
-from PIL import Image
-from tqdm import tqdm
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    ImageField,
-    InputField,
-    UIType,
-    WithBoard,
-    WithMetadata,
-)
-from invokeai.app.invocations.model import ModelIdentifierField
-from invokeai.app.invocations.primitives import ImageOutput
-from invokeai.app.services.session_processor.session_processor_common import CanceledException
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
-from invokeai.backend.tiles.tiles import calc_tiles_min_overlap
-from invokeai.backend.tiles.utils import TBLR, Tile
-
-
-@invocation("spandrel_image_to_image", title="Image-to-Image", tags=["upscale"], category="upscale", version="1.3.0")
-class SpandrelImageToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
-    """Run any spandrel image-to-image model (https://github.com/chaiNNer-org/spandrel)."""
-
-    image: ImageField = InputField(description="The input image")
-    image_to_image_model: ModelIdentifierField = InputField(
-        title="Image-to-Image Model",
-        description=FieldDescriptions.spandrel_image_to_image_model,
-        ui_type=UIType.SpandrelImageToImageModel,
-    )
-    tile_size: int = InputField(
-        default=512, description="The tile size for tiled image-to-image. Set to 0 to disable tiling."
-    )
-
-    @classmethod
-    def scale_tile(cls, tile: Tile, scale: int) -> Tile:
-        return Tile(
-            coords=TBLR(
-                top=tile.coords.top * scale,
-                bottom=tile.coords.bottom * scale,
-                left=tile.coords.left * scale,
-                right=tile.coords.right * scale,
-            ),
-            overlap=TBLR(
-                top=tile.overlap.top * scale,
-                bottom=tile.overlap.bottom * scale,
-                left=tile.overlap.left * scale,
-                right=tile.overlap.right * scale,
-            ),
-        )
-
-    @classmethod
-    def upscale_image(
-        cls,
-        image: Image.Image,
-        tile_size: int,
-        spandrel_model: SpandrelImageToImageModel,
-        is_canceled: Callable[[], bool],
-    ) -> Image.Image:
-        # Compute the image tiles.
-        if tile_size > 0:
-            min_overlap = 20
-            tiles = calc_tiles_min_overlap(
-                image_height=image.height,
-                image_width=image.width,
-                tile_height=tile_size,
-                tile_width=tile_size,
-                min_overlap=min_overlap,
-            )
-        else:
-            # No tiling. Generate a single tile that covers the entire image.
-            min_overlap = 0
-            tiles = [
-                Tile(
-                    coords=TBLR(top=0, bottom=image.height, left=0, right=image.width),
-                    overlap=TBLR(top=0, bottom=0, left=0, right=0),
-                )
-            ]
-
-        # Sort tiles first by left x coordinate, then by top y coordinate. During tile processing, we want to iterate
-        # over tiles left-to-right, top-to-bottom.
-        tiles = sorted(tiles, key=lambda x: x.coords.left)
-        tiles = sorted(tiles, key=lambda x: x.coords.top)
-
-        # Prepare input image for inference.
-        image_tensor = SpandrelImageToImageModel.pil_to_tensor(image)
-
-        # Scale the tiles for re-assembling the final image.
-        scale = spandrel_model.scale
-        scaled_tiles = [cls.scale_tile(tile, scale=scale) for tile in tiles]
-
-        # Prepare the output tensor.
-        _, channels, height, width = image_tensor.shape
-        output_tensor = torch.zeros(
-            (height * scale, width * scale, channels), dtype=torch.uint8, device=torch.device("cpu")
-        )
-
-        image_tensor = image_tensor.to(device=spandrel_model.device, dtype=spandrel_model.dtype)
-
-        # Run the model on each tile.
-        for tile, scaled_tile in tqdm(list(zip(tiles, scaled_tiles, strict=True)), desc="Upscaling Tiles"):
-            # Exit early if the invocation has been canceled.
-            if is_canceled():
-                raise CanceledException
-
-            # Extract the current tile from the input tensor.
-            input_tile = image_tensor[
-                :, :, tile.coords.top : tile.coords.bottom, tile.coords.left : tile.coords.right
-            ].to(device=spandrel_model.device, dtype=spandrel_model.dtype)
-
-            # Run the model on the tile.
-            output_tile = spandrel_model.run(input_tile)
-
-            # Convert the output tile into the output tensor's format.
-            # (N, C, H, W) -> (C, H, W)
-            output_tile = output_tile.squeeze(0)
-            # (C, H, W) -> (H, W, C)
-            output_tile = output_tile.permute(1, 2, 0)
-            output_tile = output_tile.clamp(0, 1)
-            output_tile = (output_tile * 255).to(dtype=torch.uint8, device=torch.device("cpu"))
-
-            # Merge the output tile into the output tensor.
-            # We only keep half of the overlap on the top and left side of the tile. We do this in case there are
-            # edge artifacts. We don't bother with any 'blending' in the current implementation - for most upscalers
-            # it seems unnecessary, but we may find a need in the future.
-            top_overlap = scaled_tile.overlap.top // 2
-            left_overlap = scaled_tile.overlap.left // 2
-            output_tensor[
-                scaled_tile.coords.top + top_overlap : scaled_tile.coords.bottom,
-                scaled_tile.coords.left + left_overlap : scaled_tile.coords.right,
-                :,
-            ] = output_tile[top_overlap:, left_overlap:, :]
-
-        # Convert the output tensor to a PIL image.
-        np_image = output_tensor.detach().numpy().astype(np.uint8)
-        pil_image = Image.fromarray(np_image)
-
-        return pil_image
-
-    @torch.inference_mode()
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        # Images are converted to RGB, because most models don't support an alpha channel. In the future, we may want to
-        # revisit this.
-        image = context.images.get_pil(self.image.image_name, mode="RGB")
-
-        # Load the model.
-        spandrel_model_info = context.models.load(self.image_to_image_model)
-
-        # Do the upscaling.
-        with spandrel_model_info as spandrel_model:
-            assert isinstance(spandrel_model, SpandrelImageToImageModel)
-
-            # Upscale the image
-            pil_image = self.upscale_image(image, self.tile_size, spandrel_model, context.util.is_canceled)
-
-        image_dto = context.images.save(image=pil_image)
-        return ImageOutput.build(image_dto)
-
-
-@invocation(
-    "spandrel_image_to_image_autoscale",
-    title="Image-to-Image (Autoscale)",
-    tags=["upscale"],
-    category="upscale",
-    version="1.0.0",
-)
-class SpandrelImageToImageAutoscaleInvocation(SpandrelImageToImageInvocation):
-    """Run any spandrel image-to-image model (https://github.com/chaiNNer-org/spandrel) until the target scale is reached."""
-
-    scale: float = InputField(
-        default=4.0,
-        gt=0.0,
-        le=16.0,
-        description="The final scale of the output image. If the model does not upscale the image, this will be ignored.",
-    )
-    fit_to_multiple_of_8: bool = InputField(
-        default=False,
-        description="If true, the output image will be resized to the nearest multiple of 8 in both dimensions.",
-    )
-
-    @torch.inference_mode()
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        # Images are converted to RGB, because most models don't support an alpha channel. In the future, we may want to
-        # revisit this.
-        image = context.images.get_pil(self.image.image_name, mode="RGB")
-
-        # Load the model.
-        spandrel_model_info = context.models.load(self.image_to_image_model)
-
-        # The target size of the image, determined by the provided scale. We'll run the upscaler until we hit this size.
-        # Later, we may mutate this value if the model doesn't upscale the image or if the user requested a multiple of 8.
-        target_width = int(image.width * self.scale)
-        target_height = int(image.height * self.scale)
-
-        # Do the upscaling.
-        with spandrel_model_info as spandrel_model:
-            assert isinstance(spandrel_model, SpandrelImageToImageModel)
-
-            # First pass of upscaling. Note: `pil_image` will be mutated.
-            pil_image = self.upscale_image(image, self.tile_size, spandrel_model, context.util.is_canceled)
-
-            # Some models don't upscale the image, but we have no way to know this in advance. We'll check if the model
-            # upscaled the image and run the loop below if it did. We'll require the model to upscale both dimensions
-            # to be considered an upscale model.
-            is_upscale_model = pil_image.width > image.width and pil_image.height > image.height
-
-            if is_upscale_model:
-                # This is an upscale model, so we should keep upscaling until we reach the target size.
-                iterations = 1
-                while pil_image.width < target_width or pil_image.height < target_height:
-                    pil_image = self.upscale_image(pil_image, self.tile_size, spandrel_model, context.util.is_canceled)
-                    iterations += 1
-
-                    # Sanity check to prevent excessive or infinite loops. All known upscaling models are at least 2x.
-                    # Our max scale is 16x, so with a 2x model, we should never exceed 16x == 2^4 -> 4 iterations.
-                    # We'll allow one extra iteration "just in case" and bail at 5 upscaling iterations. In practice,
-                    # we should never reach this limit.
-                    if iterations >= 5:
-                        context.logger.warning(
-                            "Upscale loop reached maximum iteration count of 5, stopping upscaling early."
-                        )
-                        break
-            else:
-                # This model doesn't upscale the image. We should ignore the scale parameter, modifying the output size
-                # to be the same as the processed image size.
-
-                # The output size is now the size of the processed image.
-                target_width = pil_image.width
-                target_height = pil_image.height
-
-                # Warn the user if they requested a scale greater than 1.
-                if self.scale > 1:
-                    context.logger.warning(
-                        "Model does not increase the size of the image, but a greater scale than 1 was requested. Image will not be scaled."
-                    )
-
-        # We may need to resize the image to a multiple of 8. Use floor division to ensure we don't scale the image up
-        # in the final resize
-        if self.fit_to_multiple_of_8:
-            target_width = int(target_width // 8 * 8)
-            target_height = int(target_height // 8 * 8)
-
-        # Final resize. Per PIL documentation, Lanczos provides the best quality for both upscale and downscale.
-        # See: https://pillow.readthedocs.io/en/stable/handbook/concepts.html#filters-comparison-table
-        pil_image = pil_image.resize((target_width, target_height), resample=Image.Resampling.LANCZOS)
-
-        image_dto = context.images.save(image=pil_image)
-        return ImageOutput.build(image_dto)

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

@@ -175,10 +175,6 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
         _, _, latent_height, latent_width = latents.shape
 
         # Calculate the tile locations to cover the latent-space image.
-        # TODO(ryand): In the future, we may want to revisit the tile overlap strategy. Things to consider:
-        # - How much overlap 'context' to provide for each denoising step.
-        # - How much overlap to use during merging/blending.
-        # - Should we 'jitter' the tile locations in each step so that the seams are in different places?
         tiles = calc_tiles_min_overlap(
             image_height=latent_height,
             image_width=latent_width,
@@ -222,8 +218,7 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
                 context=context,
                 positive_conditioning_field=self.positive_conditioning,
                 negative_conditioning_field=self.negative_conditioning,
-                device=unet.device,
-                dtype=unet.dtype,
+                unet=unet,
                 latent_height=latent_tile_height,
                 latent_width=latent_tile_width,
                 cfg_scale=self.cfg_scale,

invokeai/app/services/board_records/board_records_base.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord
+from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord, UncategorizedImageCounts
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 
 
@@ -48,3 +48,8 @@ class BoardRecordStorageBase(ABC):
     def get_all(self, include_archived: bool = False) -> list[BoardRecord]:
         """Gets all board records."""
         pass
+
+    @abstractmethod
+    def get_uncategorized_image_counts(self) -> UncategorizedImageCounts:
+        """Gets count of images and assets for uncategorized images (images with no board assocation)."""
+        pass

invokeai/app/services/board_records/board_records_common.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 from pydantic import BaseModel, Field
 
@@ -26,21 +26,25 @@ class BoardRecord(BaseModelExcludeNull):
     """Whether or not the board is archived."""
     is_private: Optional[bool] = Field(default=None, description="Whether the board is private.")
     """Whether the board is private."""
+    image_count: int = Field(description="The number of images in the board.")
+    asset_count: int = Field(description="The number of assets in the board.")
 
 
-def deserialize_board_record(board_dict: dict) -> BoardRecord:
+def deserialize_board_record(board_dict: dict[str, Any]) -> BoardRecord:
     """Deserializes a board record."""
 
     # Retrieve all the values, setting "reasonable" defaults if they are not present.
 
     board_id = board_dict.get("board_id", "unknown")
     board_name = board_dict.get("board_name", "unknown")
-    cover_image_name = board_dict.get("cover_image_name", "unknown")
+    cover_image_name = board_dict.get("cover_image_name", None)
     created_at = board_dict.get("created_at", get_iso_timestamp())
     updated_at = board_dict.get("updated_at", get_iso_timestamp())
     deleted_at = board_dict.get("deleted_at", get_iso_timestamp())
     archived = board_dict.get("archived", False)
     is_private = board_dict.get("is_private", False)
+    image_count = board_dict.get("image_count", 0)
+    asset_count = board_dict.get("asset_count", 0)
 
     return BoardRecord(
         board_id=board_id,
@@ -51,6 +55,8 @@ def deserialize_board_record(board_dict: dict) -> BoardRecord:
         deleted_at=deleted_at,
         archived=archived,
         is_private=is_private,
+        image_count=image_count,
+        asset_count=asset_count,
     )
 
 
@@ -63,19 +69,24 @@ class BoardChanges(BaseModel, extra="forbid"):
 class BoardRecordNotFoundException(Exception):
     """Raised when an board record is not found."""
 
-    def __init__(self, message="Board record not found"):
+    def __init__(self, message: str = "Board record not found"):
         super().__init__(message)
 
 
 class BoardRecordSaveException(Exception):
     """Raised when an board record cannot be saved."""
 
-    def __init__(self, message="Board record not saved"):
+    def __init__(self, message: str = "Board record not saved"):
         super().__init__(message)
 
 
 class BoardRecordDeleteException(Exception):
     """Raised when an board record cannot be deleted."""
 
-    def __init__(self, message="Board record not deleted"):
+    def __init__(self, message: str = "Board record not deleted"):
         super().__init__(message)
+
+
+class UncategorizedImageCounts(BaseModel):
+    image_count: int = Field(description="The number of uncategorized images.")
+    asset_count: int = Field(description="The number of uncategorized assets.")

invokeai/app/services/board_records/board_records_sqlite.py

@@ -1,5 +1,6 @@
 import sqlite3
 import threading
+from dataclasses import dataclass
 from typing import Union, cast
 
 from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase
@@ -9,12 +10,108 @@ from invokeai.app.services.board_records.board_records_common import (
     BoardRecordDeleteException,
     BoardRecordNotFoundException,
     BoardRecordSaveException,
+    UncategorizedImageCounts,
     deserialize_board_record,
 )
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
 from invokeai.app.util.misc import uuid_string
 
+BASE_BOARD_RECORD_QUERY = """
+    -- This query retrieves board records, joining with the board_images and images tables to get image counts and cover image names.
+    -- It is not a complete query, as it is missing a GROUP BY or WHERE clause (and is unterminated).
+    SELECT b.board_id,
+        b.board_name,
+        b.created_at,
+        b.updated_at,
+        b.archived,
+        -- Count the number of images in the board, alias image_count
+        COUNT(
+            CASE
+                WHEN i.image_category in ('general') -- "Images" are images in the 'general' category
+                AND i.is_intermediate = 0 THEN 1 -- Intermediates are not counted
+            END
+        ) AS image_count,
+        -- Count the number of assets in the board, alias asset_count
+        COUNT(
+            CASE
+                WHEN i.image_category in ('control', 'mask', 'user', 'other') -- "Assets" are images in any of the other categories ('control', 'mask', 'user', 'other')
+                AND i.is_intermediate = 0 THEN 1 -- Intermediates are not counted
+            END
+        ) AS asset_count,
+        -- Get the name of the the most recent image in the board, alias cover_image_name
+        (
+            SELECT bi.image_name
+            FROM board_images bi
+                JOIN images i ON bi.image_name = i.image_name
+            WHERE bi.board_id = b.board_id
+                AND i.is_intermediate = 0 -- Intermediates cannot be cover images
+            ORDER BY i.created_at DESC -- Sort by created_at to get the most recent image
+            LIMIT 1
+        ) AS cover_image_name
+    FROM boards b
+        LEFT JOIN board_images bi ON b.board_id = bi.board_id
+        LEFT JOIN images i ON bi.image_name = i.image_name
+    """
+
+
+@dataclass
+class PaginatedBoardRecordsQueries:
+    main_query: str
+    total_count_query: str
+
+
+def get_paginated_list_board_records_queries(include_archived: bool) -> PaginatedBoardRecordsQueries:
+    """Gets a query to retrieve a paginated list of board records."""
+
+    archived_condition = "WHERE b.archived = 0" if not include_archived else ""
+
+    # The GROUP BY must be added _after_ the WHERE clause!
+    main_query = f"""
+        {BASE_BOARD_RECORD_QUERY}
+        {archived_condition}
+        GROUP BY b.board_id,
+            b.board_name,
+            b.created_at,
+            b.updated_at
+        ORDER BY b.created_at DESC
+        LIMIT ? OFFSET ?;
+        """
+
+    total_count_query = f"""
+        SELECT COUNT(*)
+        FROM boards b
+        {archived_condition};
+        """
+
+    return PaginatedBoardRecordsQueries(main_query=main_query, total_count_query=total_count_query)
+
+
+def get_list_all_board_records_query(include_archived: bool) -> str:
+    """Gets a query to retrieve all board records."""
+
+    archived_condition = "WHERE b.archived = 0" if not include_archived else ""
+
+    # The GROUP BY must be added _after_ the WHERE clause!
+    return f"""
+        {BASE_BOARD_RECORD_QUERY}
+        {archived_condition}
+        GROUP BY b.board_id,
+            b.board_name,
+            b.created_at,
+            b.updated_at
+        ORDER BY b.created_at DESC;
+        """
+
+
+def get_board_record_query() -> str:
+    """Gets a query to retrieve a board record."""
+
+    return f"""
+        {BASE_BOARD_RECORD_QUERY}
+        WHERE b.board_id = ?;
+        """
+
 
 class SqliteBoardRecordStorage(BoardRecordStorageBase):
     _conn: sqlite3.Connection
@@ -76,11 +173,7 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
         try:
             self._lock.acquire()
             self._cursor.execute(
-                """--sql
-                SELECT *
-                FROM boards
-                WHERE board_id = ?;
-                """,
+                get_board_record_query(),
                 (board_id,),
             )
 
@@ -92,7 +185,7 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
             self._lock.release()
         if result is None:
             raise BoardRecordNotFoundException
-        return BoardRecord(**dict(result))
+        return deserialize_board_record(dict(result))
 
     def update(
         self,
@@ -149,45 +242,17 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
         try:
             self._lock.acquire()
 
-            # Build base query
-            base_query = """
-                SELECT *
-                FROM boards
-                {archived_filter}
-                ORDER BY created_at DESC
-                LIMIT ? OFFSET ?;
-            """
+            queries = get_paginated_list_board_records_queries(include_archived=include_archived)
 
-            # Determine archived filter condition
-            if include_archived:
-                archived_filter = ""
-            else:
-                archived_filter = "WHERE archived = 0"
-
-            final_query = base_query.format(archived_filter=archived_filter)
-
-            # Execute query to fetch boards
-            self._cursor.execute(final_query, (limit, offset))
+            self._cursor.execute(
+                queries.main_query,
+                (limit, offset),
+            )
 
             result = cast(list[sqlite3.Row], self._cursor.fetchall())
             boards = [deserialize_board_record(dict(r)) for r in result]
 
-            # Determine count query
-            if include_archived:
-                count_query = """
-                    SELECT COUNT(*)
-                    FROM boards;
-                """
-            else:
-                count_query = """
-                    SELECT COUNT(*)
-                    FROM boards
-                    WHERE archived = 0;
-                """
-
-            # Execute count query
-            self._cursor.execute(count_query)
-
+            self._cursor.execute(queries.total_count_query)
             count = cast(int, self._cursor.fetchone()[0])
 
             return OffsetPaginatedResults[BoardRecord](items=boards, offset=offset, limit=limit, total=count)
@@ -201,26 +266,9 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
     def get_all(self, include_archived: bool = False) -> list[BoardRecord]:
         try:
             self._lock.acquire()
-
-            base_query = """
-                SELECT *
-                FROM boards
-                {archived_filter}
-                ORDER BY created_at DESC
-            """
-
-            if include_archived:
-                archived_filter = ""
-            else:
-                archived_filter = "WHERE archived = 0"
-
-            final_query = base_query.format(archived_filter=archived_filter)
-
-            self._cursor.execute(final_query)
-
+            self._cursor.execute(get_list_all_board_records_query(include_archived=include_archived))
             result = cast(list[sqlite3.Row], self._cursor.fetchall())
             boards = [deserialize_board_record(dict(r)) for r in result]
-
             return boards
 
         except sqlite3.Error as e:
@@ -228,3 +276,28 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
             raise e
         finally:
             self._lock.release()
+
+    def get_uncategorized_image_counts(self) -> UncategorizedImageCounts:
+        try:
+            self._lock.acquire()
+            query = """
+                -- Get the count of uncategorized images and assets.
+                SELECT
+                    CASE
+                        WHEN i.image_category = 'general' THEN 'image_count' -- "Images" are images in the 'general' category
+                        ELSE 'asset_count' -- "Assets" are images in any of the other categories ('control', 'mask', 'user', 'other')
+                    END AS category_type,
+                    COUNT(*) AS unassigned_count
+                FROM images i
+                LEFT JOIN board_images bi ON i.image_name = bi.image_name
+                WHERE bi.board_id IS NULL -- Uncategorized images have no board association
+                AND i.is_intermediate = 0 -- Omit intermediates from the counts
+                GROUP BY category_type; -- Group by category_type alias, as derived from the image_category column earlier
+                """
+            self._cursor.execute(query)
+            results = self._cursor.fetchall()
+            image_count = dict(results)["image_count"]
+            asset_count = dict(results)["asset_count"]
+            return UncategorizedImageCounts(image_count=image_count, asset_count=asset_count)
+        finally:
+            self._lock.release()

invokeai/app/services/boards/boards_common.py

@@ -1,23 +1,8 @@
-from typing import Optional
-
-from pydantic import Field
-
 from invokeai.app.services.board_records.board_records_common import BoardRecord
 
 
+# TODO(psyche): BoardDTO is now identical to BoardRecord. We should consider removing it.
 class BoardDTO(BoardRecord):
-    """Deserialized board record with cover image URL and image count."""
+    """Deserialized board record."""
 
-    cover_image_name: Optional[str] = Field(description="The name of the board's cover image.")
-    """The URL of the thumbnail of the most recent image in the board."""
-    image_count: int = Field(description="The number of images in the board.")
-    """The number of images in the board."""
-
-
-def board_record_to_dto(board_record: BoardRecord, cover_image_name: Optional[str], image_count: int) -> BoardDTO:
-    """Converts a board record to a board DTO."""
-    return BoardDTO(
-        **board_record.model_dump(exclude={"cover_image_name"}),
-        cover_image_name=cover_image_name,
-        image_count=image_count,
-    )
+    pass

invokeai/app/services/boards/boards_default.py

@@ -1,6 +1,6 @@
 from invokeai.app.services.board_records.board_records_common import BoardChanges
 from invokeai.app.services.boards.boards_base import BoardServiceABC
-from invokeai.app.services.boards.boards_common import BoardDTO, board_record_to_dto
+from invokeai.app.services.boards.boards_common import BoardDTO
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.shared.pagination import OffsetPaginatedResults
 
@@ -16,17 +16,11 @@ class BoardService(BoardServiceABC):
         board_name: str,
     ) -> BoardDTO:
         board_record = self.__invoker.services.board_records.save(board_name)
-        return board_record_to_dto(board_record, None, 0)
+        return BoardDTO.model_validate(board_record.model_dump())
 
     def get_dto(self, board_id: str) -> BoardDTO:
         board_record = self.__invoker.services.board_records.get(board_id)
-        cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(board_record.board_id)
-        if cover_image:
-            cover_image_name = cover_image.image_name
-        else:
-            cover_image_name = None
-        image_count = self.__invoker.services.board_image_records.get_image_count_for_board(board_id)
-        return board_record_to_dto(board_record, cover_image_name, image_count)
+        return BoardDTO.model_validate(board_record.model_dump())
 
     def update(
         self,
@@ -34,14 +28,7 @@ class BoardService(BoardServiceABC):
         changes: BoardChanges,
     ) -> BoardDTO:
         board_record = self.__invoker.services.board_records.update(board_id, changes)
-        cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(board_record.board_id)
-        if cover_image:
-            cover_image_name = cover_image.image_name
-        else:
-            cover_image_name = None
-
-        image_count = self.__invoker.services.board_image_records.get_image_count_for_board(board_id)
-        return board_record_to_dto(board_record, cover_image_name, image_count)
+        return BoardDTO.model_validate(board_record.model_dump())
 
     def delete(self, board_id: str) -> None:
         self.__invoker.services.board_records.delete(board_id)
@@ -50,30 +37,10 @@ class BoardService(BoardServiceABC):
         self, offset: int = 0, limit: int = 10, include_archived: bool = False
     ) -> OffsetPaginatedResults[BoardDTO]:
         board_records = self.__invoker.services.board_records.get_many(offset, limit, include_archived)
-        board_dtos = []
-        for r in board_records.items:
-            cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(r.board_id)
-            if cover_image:
-                cover_image_name = cover_image.image_name
-            else:
-                cover_image_name = None
-
-            image_count = self.__invoker.services.board_image_records.get_image_count_for_board(r.board_id)
-            board_dtos.append(board_record_to_dto(r, cover_image_name, image_count))
-
+        board_dtos = [BoardDTO.model_validate(r.model_dump()) for r in board_records.items]
         return OffsetPaginatedResults[BoardDTO](items=board_dtos, offset=offset, limit=limit, total=len(board_dtos))
 
     def get_all(self, include_archived: bool = False) -> list[BoardDTO]:
         board_records = self.__invoker.services.board_records.get_all(include_archived)
-        board_dtos = []
-        for r in board_records:
-            cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(r.board_id)
-            if cover_image:
-                cover_image_name = cover_image.image_name
-            else:
-                cover_image_name = None
-
-            image_count = self.__invoker.services.board_image_records.get_image_count_for_board(r.board_id)
-            board_dtos.append(board_record_to_dto(r, cover_image_name, image_count))
-
+        board_dtos = [BoardDTO.model_validate(r.model_dump()) for r in board_records]
         return board_dtos

invokeai/app/services/config/config_default.py

@@ -91,7 +91,6 @@ class InvokeAIAppConfig(BaseSettings):
         db_dir: Path to InvokeAI databases directory.
         outputs_dir: Path to directory for outputs.
         custom_nodes_dir: Path to directory for custom nodes.
-        style_presets_dir: Path to directory for style presets.
         log_handlers: Log handler. Valid options are "console", "file=<path>", "syslog=path|address:host:port", "http=<url>".
         log_format: Log format. Use "plain" for text-only, "color" for colorized output, "legacy" for 2.3-style logging and "syslog" for syslog-style.<br>Valid values: `plain`, `color`, `syslog`, `legacy`
         log_level: Emit logging messages at this level or higher.<br>Valid values: `debug`, `info`, `warning`, `error`, `critical`
@@ -154,7 +153,6 @@ class InvokeAIAppConfig(BaseSettings):
     db_dir:                        Path = Field(default=Path("databases"),  description="Path to InvokeAI databases directory.")
     outputs_dir:                   Path = Field(default=Path("outputs"),    description="Path to directory for outputs.")
     custom_nodes_dir:              Path = Field(default=Path("nodes"),      description="Path to directory for custom nodes.")
-    style_presets_dir:      Path = Field(default=Path("style_presets"),      description="Path to directory for style presets.")
 
     # LOGGING
     log_handlers:             list[str] = Field(default=["console"],        description='Log handler. Valid options are "console", "file=<path>", "syslog=path|address:host:port", "http=<url>".')
@@ -302,11 +300,6 @@ class InvokeAIAppConfig(BaseSettings):
         """Path to the models directory, resolved to an absolute path.."""
         return self._resolve(self.models_dir)
 
-    @property
-    def style_presets_path(self) -> Path:
-        """Path to the style presets directory, resolved to an absolute path.."""
-        return self._resolve(self.style_presets_dir)
-
     @property
     def convert_cache_path(self) -> Path:
         """Path to the converted cache models directory, resolved to an absolute path.."""

invokeai/app/services/events/events_common.py

@@ -88,7 +88,6 @@ class QueueItemEventBase(QueueEventBase):
 
     item_id: int = Field(description="The ID of the queue item")
     batch_id: str = Field(description="The ID of the queue batch")
-    origin: str | None = Field(default=None, description="The origin of the batch")
 
 
 class InvocationEventBase(QueueItemEventBase):
@@ -96,6 +95,8 @@ class InvocationEventBase(QueueItemEventBase):
 
     session_id: str = Field(description="The ID of the session (aka graph execution state)")
     queue_id: str = Field(description="The ID of the queue")
+    item_id: int = Field(description="The ID of the queue item")
+    batch_id: str = Field(description="The ID of the queue batch")
     session_id: str = Field(description="The ID of the session (aka graph execution state)")
     invocation: AnyInvocation = Field(description="The ID of the invocation")
     invocation_source_id: str = Field(description="The ID of the prepared invocation's source node")
@@ -113,7 +114,6 @@ class InvocationStartedEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
-            origin=queue_item.origin,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -147,7 +147,6 @@ class InvocationDenoiseProgressEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
-            origin=queue_item.origin,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -185,7 +184,6 @@ class InvocationCompleteEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
-            origin=queue_item.origin,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -218,7 +216,6 @@ class InvocationErrorEvent(InvocationEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
-            origin=queue_item.origin,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -256,7 +253,6 @@ class QueueItemStatusChangedEvent(QueueItemEventBase):
             queue_id=queue_item.queue_id,
             item_id=queue_item.item_id,
             batch_id=queue_item.batch_id,
-            origin=queue_item.origin,
             session_id=queue_item.session_id,
             status=queue_item.status,
             error_type=queue_item.error_type,
@@ -283,14 +279,12 @@ class BatchEnqueuedEvent(QueueEventBase):
         description="The number of invocations initially requested to be enqueued (may be less than enqueued if queue was full)"
     )
     priority: int = Field(description="The priority of the batch")
-    origin: str | None = Field(default=None, description="The origin of the batch")
 
     @classmethod
     def build(cls, enqueue_result: EnqueueBatchResult) -> "BatchEnqueuedEvent":
         return cls(
             queue_id=enqueue_result.queue_id,
             batch_id=enqueue_result.batch.batch_id,
-            origin=enqueue_result.batch.origin,
             enqueued=enqueue_result.enqueued,
             requested=enqueue_result.requested,
             priority=enqueue_result.priority,

invokeai/app/services/events/events_fastapievents.py

@@ -1,44 +1,46 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
 import asyncio
 import threading
+from queue import Empty, Queue
 
 from fastapi_events.dispatcher import dispatch
 
 from invokeai.app.services.events.events_base import EventServiceBase
-from invokeai.app.services.events.events_common import EventBase
+from invokeai.app.services.events.events_common import (
+    EventBase,
+)
 
 
 class FastAPIEventService(EventServiceBase):
-    def __init__(self, event_handler_id: int, loop: asyncio.AbstractEventLoop) -> None:
+    def __init__(self, event_handler_id: int) -> None:
         self.event_handler_id = event_handler_id
-        self._queue = asyncio.Queue[EventBase | None]()
+        self._queue = Queue[EventBase | None]()
         self._stop_event = threading.Event()
-        self._loop = loop
-
-        # We need to store a reference to the task so it doesn't get GC'd
-        # See: https://docs.python.org/3/library/asyncio-task.html#creating-tasks
-        self._background_tasks: set[asyncio.Task[None]] = set()
-        task = self._loop.create_task(self._dispatch_from_queue(stop_event=self._stop_event))
-        self._background_tasks.add(task)
-        task.add_done_callback(self._background_tasks.remove)
+        asyncio.create_task(self._dispatch_from_queue(stop_event=self._stop_event))
 
         super().__init__()
 
     def stop(self, *args, **kwargs):
         self._stop_event.set()
-        self._loop.call_soon_threadsafe(self._queue.put_nowait, None)
+        self._queue.put(None)
 
     def dispatch(self, event: EventBase) -> None:
-        self._loop.call_soon_threadsafe(self._queue.put_nowait, event)
+        self._queue.put(event)
 
     async def _dispatch_from_queue(self, stop_event: threading.Event):
         """Get events on from the queue and dispatch them, from the correct thread"""
         while not stop_event.is_set():
             try:
-                event = await self._queue.get()
+                event = self._queue.get(block=False)
                 if not event:  # Probably stopping
                     continue
                 # Leave the payloads as live pydantic models
                 dispatch(event, middleware_id=self.event_handler_id, payload_schema_dump=False)
 
+            except Empty:
+                await asyncio.sleep(0.1)
+                pass
+
             except asyncio.CancelledError as e:
                 raise e  # Raise a proper error

invokeai/app/services/image_files/image_files_disk.py

@@ -1,10 +1,11 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team
 from pathlib import Path
 from queue import Queue
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 
 from PIL import Image, PngImagePlugin
 from PIL.Image import Image as PILImageType
+from send2trash import send2trash
 
 from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase
 from invokeai.app.services.image_files.image_files_common import (
@@ -19,12 +20,18 @@ from invokeai.app.util.thumbnails import get_thumbnail_name, make_thumbnail
 class DiskImageFileStorage(ImageFileStorageBase):
     """Stores images on disk"""
 
+    __output_folder: Path
+    __cache_ids: Queue  # TODO: this is an incredibly naive cache
+    __cache: Dict[Path, PILImageType]
+    __max_cache_size: int
+    __invoker: Invoker
+
     def __init__(self, output_folder: Union[str, Path]):
-        self.__cache: dict[Path, PILImageType] = {}
-        self.__cache_ids = Queue[Path]()
+        self.__cache = {}
+        self.__cache_ids = Queue()
         self.__max_cache_size = 10  # TODO: get this from config
 
-        self.__output_folder = output_folder if isinstance(output_folder, Path) else Path(output_folder)
+        self.__output_folder: Path = output_folder if isinstance(output_folder, Path) else Path(output_folder)
         self.__thumbnails_folder = self.__output_folder / "thumbnails"
         # Validate required output folders at launch
         self.__validate_storage_folders()
@@ -96,7 +103,7 @@ class DiskImageFileStorage(ImageFileStorageBase):
             image_path = self.get_path(image_name)
 
             if image_path.exists():
-                image_path.unlink()
+                send2trash(image_path)
             if image_path in self.__cache:
                 del self.__cache[image_path]
 
@@ -104,7 +111,7 @@ class DiskImageFileStorage(ImageFileStorageBase):
             thumbnail_path = self.get_path(thumbnail_name, True)
 
             if thumbnail_path.exists():
-                thumbnail_path.unlink()
+                send2trash(thumbnail_path)
             if thumbnail_path in self.__cache:
                 del self.__cache[thumbnail_path]
         except Exception as e:

invokeai/app/services/invocation_services.py

@@ -4,8 +4,6 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from invokeai.app.services.object_serializer.object_serializer_base import ObjectSerializerBase
-from invokeai.app.services.style_preset_images.style_preset_images_base import StylePresetImageFileStorageBase
-from invokeai.app.services.style_preset_records.style_preset_records_base import StylePresetRecordsStorageBase
 
 if TYPE_CHECKING:
     from logging import Logger
@@ -63,8 +61,6 @@ class InvocationServices:
         workflow_records: "WorkflowRecordsStorageBase",
         tensors: "ObjectSerializerBase[torch.Tensor]",
         conditioning: "ObjectSerializerBase[ConditioningFieldData]",
-        style_preset_records: "StylePresetRecordsStorageBase",
-        style_preset_image_files: "StylePresetImageFileStorageBase",
     ):
         self.board_images = board_images
         self.board_image_records = board_image_records
@@ -89,5 +85,3 @@ class InvocationServices:
         self.workflow_records = workflow_records
         self.tensors = tensors
         self.conditioning = conditioning
-        self.style_preset_records = style_preset_records
-        self.style_preset_image_files = style_preset_image_files

invokeai/app/services/model_images/model_images_default.py

@@ -2,6 +2,7 @@ from pathlib import Path
 
 from PIL import Image
 from PIL.Image import Image as PILImageType
+from send2trash import send2trash
 
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_images.model_images_base import ModelImageFileStorageBase
@@ -69,7 +70,7 @@ class ModelImageFileStorageDisk(ModelImageFileStorageBase):
             if not self._validate_path(path):
                 raise ModelImageFileNotFoundException
 
-            path.unlink()
+            send2trash(path)
 
         except Exception as e:
             raise ModelImageFileDeleteException from e

invokeai/app/services/model_install/model_install_base.py

@@ -3,7 +3,7 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic.networks import AnyHttpUrl
 
@@ -12,7 +12,7 @@ from invokeai.app.services.download import DownloadQueueServiceBase
 from invokeai.app.services.events.events_base import EventServiceBase
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
-from invokeai.app.services.model_records import ModelRecordChanges, ModelRecordServiceBase
+from invokeai.app.services.model_records import ModelRecordServiceBase
 from invokeai.backend.model_manager import AnyModelConfig
 
 
@@ -64,7 +64,7 @@ class ModelInstallServiceBase(ABC):
     def register_path(
         self,
         model_path: Union[Path, str],
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> str:
         """
         Probe and register the model at model_path.
@@ -72,7 +72,7 @@ class ModelInstallServiceBase(ABC):
         This keeps the model in its current location.
 
         :param model_path: Filesystem Path to the model.
-        :param config: ModelRecordChanges object that will override autoassigned model record values.
+        :param config: Dict of attributes that will override autoassigned values.
         :returns id: The string ID of the registered model.
         """
 
@@ -92,7 +92,7 @@ class ModelInstallServiceBase(ABC):
     def install_path(
         self,
         model_path: Union[Path, str],
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> str:
         """
         Probe, register and install the model in the models directory.
@@ -101,7 +101,7 @@ class ModelInstallServiceBase(ABC):
         the models directory handled by InvokeAI.
 
         :param model_path: Filesystem Path to the model.
-        :param config: ModelRecordChanges object that will override autoassigned model record values.
+        :param config: Dict of attributes that will override autoassigned values.
         :returns id: The string ID of the registered model.
         """
 
@@ -109,14 +109,14 @@ class ModelInstallServiceBase(ABC):
     def heuristic_import(
         self,
         source: str,
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
         access_token: Optional[str] = None,
         inplace: Optional[bool] = False,
     ) -> ModelInstallJob:
         r"""Install the indicated model using heuristics to interpret user intentions.
 
         :param source: String source
-        :param config: Optional ModelRecordChanges object. Any fields in this object
+        :param config: Optional dict. Any fields in this dict
          will override corresponding autoassigned probe fields in the
          model's config record as described in `import_model()`.
         :param access_token: Optional access token for remote sources.
@@ -147,7 +147,7 @@ class ModelInstallServiceBase(ABC):
     def import_model(
         self,
         source: ModelSource,
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> ModelInstallJob:
         """Install the indicated model.
 

invokeai/app/services/model_install/model_install_common.py

@@ -2,14 +2,13 @@ import re
 import traceback
 from enum import Enum
 from pathlib import Path
-from typing import Literal, Optional, Set, Union
+from typing import Any, Dict, Literal, Optional, Set, Union
 
 from pydantic import BaseModel, Field, PrivateAttr, field_validator
 from pydantic.networks import AnyHttpUrl
 from typing_extensions import Annotated
 
 from invokeai.app.services.download import DownloadJob, MultiFileDownloadJob
-from invokeai.app.services.model_records import ModelRecordChanges
 from invokeai.backend.model_manager import AnyModelConfig, ModelRepoVariant
 from invokeai.backend.model_manager.config import ModelSourceType
 from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
@@ -134,9 +133,8 @@ class ModelInstallJob(BaseModel):
     id: int = Field(description="Unique ID for this job")
     status: InstallStatus = Field(default=InstallStatus.WAITING, description="Current status of install process")
     error_reason: Optional[str] = Field(default=None, description="Information about why the job failed")
-    config_in: ModelRecordChanges = Field(
-        default_factory=ModelRecordChanges,
-        description="Configuration information (e.g. 'description') to apply to model.",
+    config_in: Dict[str, Any] = Field(
+        default_factory=dict, description="Configuration information (e.g. 'description') to apply to model."
     )
     config_out: Optional[AnyModelConfig] = Field(
         default=None, description="After successful installation, this will hold the configuration object."

invokeai/app/services/model_install/model_install_default.py

@@ -163,27 +163,26 @@ class ModelInstallService(ModelInstallServiceBase):
     def register_path(
         self,
         model_path: Union[Path, str],
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> str:  # noqa D102
         model_path = Path(model_path)
-        config = config or ModelRecordChanges()
-        if not config.source:
-            config.source = model_path.resolve().as_posix()
-        config.source_type = ModelSourceType.Path
+        config = config or {}
+        if not config.get("source"):
+            config["source"] = model_path.resolve().as_posix()
+        config["source_type"] = ModelSourceType.Path
         return self._register(model_path, config)
 
     def install_path(
         self,
         model_path: Union[Path, str],
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> str:  # noqa D102
         model_path = Path(model_path)
-        config = config or ModelRecordChanges()
-        info: AnyModelConfig = ModelProbe.probe(
-            Path(model_path), config.model_dump(), hash_algo=self._app_config.hashing_algorithm
-        )  # type: ignore
+        config = config or {}
 
-        if preferred_name := config.name:
+        info: AnyModelConfig = ModelProbe.probe(Path(model_path), config, hash_algo=self._app_config.hashing_algorithm)
+
+        if preferred_name := config.get("name"):
             preferred_name = Path(preferred_name).with_suffix(model_path.suffix)
 
         dest_path = (
@@ -205,7 +204,7 @@ class ModelInstallService(ModelInstallServiceBase):
     def heuristic_import(
         self,
         source: str,
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
         access_token: Optional[str] = None,
         inplace: Optional[bool] = False,
     ) -> ModelInstallJob:
@@ -217,7 +216,7 @@ class ModelInstallService(ModelInstallServiceBase):
             source_obj.access_token = access_token
         return self.import_model(source_obj, config)
 
-    def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] = None) -> ModelInstallJob:  # noqa D102
+    def import_model(self, source: ModelSource, config: Optional[Dict[str, Any]] = None) -> ModelInstallJob:  # noqa D102
         similar_jobs = [x for x in self.list_jobs() if x.source == source and not x.in_terminal_state]
         if similar_jobs:
             self._logger.warning(f"There is already an active install job for {source}. Not enqueuing.")
@@ -319,17 +318,16 @@ class ModelInstallService(ModelInstallServiceBase):
                         model_path = self._app_config.models_path / model_path
                     model_path = model_path.resolve()
 
-                    config = ModelRecordChanges(
-                        name=model_name,
-                        description=stanza.get("description"),
-                    )
+                    config: dict[str, Any] = {}
+                    config["name"] = model_name
+                    config["description"] = stanza.get("description")
                     legacy_config_path = stanza.get("config")
                     if legacy_config_path:
                         # In v3, these paths were relative to the root. Migrate them to be relative to the legacy_conf_dir.
                         legacy_config_path = self._app_config.root_path / legacy_config_path
                         if legacy_config_path.is_relative_to(self._app_config.legacy_conf_path):
                             legacy_config_path = legacy_config_path.relative_to(self._app_config.legacy_conf_path)
-                        config.config_path = str(legacy_config_path)
+                        config["config_path"] = str(legacy_config_path)
                     try:
                         id = self.register_path(model_path=model_path, config=config)
                         self._logger.info(f"Migrated {model_name} with id {id}")
@@ -502,11 +500,11 @@ class ModelInstallService(ModelInstallServiceBase):
         job.total_bytes = self._stat_size(job.local_path)
         job.bytes = job.total_bytes
         self._signal_job_running(job)
-        job.config_in.source = str(job.source)
-        job.config_in.source_type = MODEL_SOURCE_TO_TYPE_MAP[job.source.__class__]
+        job.config_in["source"] = str(job.source)
+        job.config_in["source_type"] = MODEL_SOURCE_TO_TYPE_MAP[job.source.__class__]
         # enter the metadata, if there is any
         if isinstance(job.source_metadata, (HuggingFaceMetadata)):
-            job.config_in.source_api_response = job.source_metadata.api_response
+            job.config_in["source_api_response"] = job.source_metadata.api_response
 
         if job.inplace:
             key = self.register_path(job.local_path, job.config_in)
@@ -641,11 +639,11 @@ class ModelInstallService(ModelInstallServiceBase):
         return new_path
 
     def _register(
-        self, model_path: Path, config: Optional[ModelRecordChanges] = None, info: Optional[AnyModelConfig] = None
+        self, model_path: Path, config: Optional[Dict[str, Any]] = None, info: Optional[AnyModelConfig] = None
     ) -> str:
-        config = config or ModelRecordChanges()
+        config = config or {}
 
-        info = info or ModelProbe.probe(model_path, config.model_dump(), hash_algo=self._app_config.hashing_algorithm)  # type: ignore
+        info = info or ModelProbe.probe(model_path, config, hash_algo=self._app_config.hashing_algorithm)
 
         model_path = model_path.resolve()
 
@@ -676,13 +674,11 @@ class ModelInstallService(ModelInstallServiceBase):
         precision = TorchDevice.choose_torch_dtype()
         return ModelRepoVariant.FP16 if precision == torch.float16 else None
 
-    def _import_local_model(
-        self, source: LocalModelSource, config: Optional[ModelRecordChanges] = None
-    ) -> ModelInstallJob:
+    def _import_local_model(self, source: LocalModelSource, config: Optional[Dict[str, Any]]) -> ModelInstallJob:
         return ModelInstallJob(
             id=self._next_id(),
             source=source,
-            config_in=config or ModelRecordChanges(),
+            config_in=config or {},
             local_path=Path(source.path),
             inplace=source.inplace or False,
         )
@@ -690,7 +686,7 @@ class ModelInstallService(ModelInstallServiceBase):
     def _import_from_hf(
         self,
         source: HFModelSource,
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> ModelInstallJob:
         # Add user's cached access token to HuggingFace requests
         if source.access_token is None:
@@ -706,7 +702,7 @@ class ModelInstallService(ModelInstallServiceBase):
     def _import_from_url(
         self,
         source: URLModelSource,
-        config: Optional[ModelRecordChanges] = None,
+        config: Optional[Dict[str, Any]],
     ) -> ModelInstallJob:
         remote_files, metadata = self._remote_files_from_source(source)
         return self._import_remote_model(
@@ -721,7 +717,7 @@ class ModelInstallService(ModelInstallServiceBase):
         source: HFModelSource | URLModelSource,
         remote_files: List[RemoteModelFile],
         metadata: Optional[AnyModelRepoMetadata],
-        config: Optional[ModelRecordChanges],
+        config: Optional[Dict[str, Any]],
     ) -> ModelInstallJob:
         if len(remote_files) == 0:
             raise ValueError(f"{source}: No downloadable files found")
@@ -734,7 +730,7 @@ class ModelInstallService(ModelInstallServiceBase):
         install_job = ModelInstallJob(
             id=self._next_id(),
             source=source,
-            config_in=config or ModelRecordChanges(),
+            config_in=config or {},
             source_metadata=metadata,
             local_path=destdir,  # local path may change once the download has started due to content-disposition handling
             bytes=0,

invokeai/app/services/model_records/model_records_base.py

@@ -18,7 +18,6 @@ from invokeai.backend.model_manager.config import (
     ControlAdapterDefaultSettings,
     MainModelDefaultSettings,
     ModelFormat,
-    ModelSourceType,
     ModelType,
     ModelVariantType,
     SchedulerPredictionType,
@@ -67,16 +66,10 @@ class ModelRecordChanges(BaseModelExcludeNull):
     """A set of changes to apply to a model."""
 
     # Changes applicable to all models
-    source: Optional[str] = Field(description="original source of the model", default=None)
-    source_type: Optional[ModelSourceType] = Field(description="type of model source", default=None)
-    source_api_response: Optional[str] = Field(description="metadata from remote source", default=None)
     name: Optional[str] = Field(description="Name of the model.", default=None)
     path: Optional[str] = Field(description="Path to the model.", default=None)
     description: Optional[str] = Field(description="Model description", default=None)
     base: Optional[BaseModelType] = Field(description="The base model.", default=None)
-    type: Optional[ModelType] = Field(description="Type of model", default=None)
-    key: Optional[str] = Field(description="Database ID for this model", default=None)
-    hash: Optional[str] = Field(description="hash of model file", default=None)
     trigger_phrases: Optional[set[str]] = Field(description="Set of trigger phrases for this model", default=None)
     default_settings: Optional[MainModelDefaultSettings | ControlAdapterDefaultSettings] = Field(
         description="Default settings for this model", default=None

invokeai/app/services/session_queue/session_queue_base.py

@@ -6,7 +6,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
-    CancelByOriginResult,
     CancelByQueueIDResult,
     ClearResult,
     EnqueueBatchResult,
@@ -96,11 +95,6 @@ class SessionQueueBase(ABC):
         """Cancels all queue items with matching batch IDs"""
         pass
 
-    @abstractmethod
-    def cancel_by_origin(self, queue_id: str, origin: str) -> CancelByOriginResult:
-        """Cancels all queue items with the given batch origin"""
-        pass
-
     @abstractmethod
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
         """Cancels all queue items with matching queue ID"""

invokeai/app/services/session_queue/session_queue_common.py

@@ -77,7 +77,6 @@ BatchDataCollection: TypeAlias = list[list[BatchDatum]]
 
 class Batch(BaseModel):
     batch_id: str = Field(default_factory=uuid_string, description="The ID of the batch")
-    origin: str | None = Field(default=None, description="The origin of this batch.")
     data: Optional[BatchDataCollection] = Field(default=None, description="The batch data collection.")
     graph: Graph = Field(description="The graph to initialize the session with")
     workflow: Optional[WorkflowWithoutID] = Field(
@@ -196,7 +195,6 @@ class SessionQueueItemWithoutGraph(BaseModel):
     status: QUEUE_ITEM_STATUS = Field(default="pending", description="The status of this queue item")
     priority: int = Field(default=0, description="The priority of this queue item")
     batch_id: str = Field(description="The ID of the batch associated with this queue item")
-    origin: str | None = Field(default=None, description="The origin of this queue item. ")
     session_id: str = Field(
         description="The ID of the session associated with this queue item. The session doesn't exist in graph_executions until the queue item is executed."
     )
@@ -296,7 +294,6 @@ class SessionQueueStatus(BaseModel):
 class BatchStatus(BaseModel):
     queue_id: str = Field(..., description="The ID of the queue")
     batch_id: str = Field(..., description="The ID of the batch")
-    origin: str | None = Field(..., description="The origin of the batch")
     pending: int = Field(..., description="Number of queue items with status 'pending'")
     in_progress: int = Field(..., description="Number of queue items with status 'in_progress'")
     completed: int = Field(..., description="Number of queue items with status 'complete'")
@@ -331,12 +328,6 @@ class CancelByBatchIDsResult(BaseModel):
     canceled: int = Field(..., description="Number of queue items canceled")
 
 
-class CancelByOriginResult(BaseModel):
-    """Result of canceling by list of batch ids"""
-
-    canceled: int = Field(..., description="Number of queue items canceled")
-
-
 class CancelByQueueIDResult(CancelByBatchIDsResult):
     """Result of canceling by queue id"""
 
@@ -442,7 +433,6 @@ class SessionQueueValueToInsert(NamedTuple):
     field_values: Optional[str]  # field_values json
     priority: int  # priority
     workflow: Optional[str]  # workflow json
-    origin: str | None
 
 
 ValuesToInsert: TypeAlias = list[SessionQueueValueToInsert]
@@ -463,7 +453,6 @@ def prepare_values_to_insert(queue_id: str, batch: Batch, priority: int, max_new
                 json.dumps(field_values, default=to_jsonable_python) if field_values else None,  # field_values (json)
                 priority,  # priority
                 json.dumps(workflow, default=to_jsonable_python) if workflow else None,  # workflow (json)
-                batch.origin,  # origin
             )
         )
     return values_to_insert

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -10,7 +10,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
-    CancelByOriginResult,
     CancelByQueueIDResult,
     ClearResult,
     EnqueueBatchResult,
@@ -128,8 +127,8 @@ class SqliteSessionQueue(SessionQueueBase):
 
             self.__cursor.executemany(
                 """--sql
-                INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
                 """,
                 values_to_insert,
             )
@@ -418,7 +417,11 @@ class SqliteSessionQueue(SessionQueueBase):
             )
             self.__conn.commit()
             if current_queue_item is not None and current_queue_item.batch_id in batch_ids:
-                self._set_queue_item_status(current_queue_item.item_id, "canceled")
+                batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id)
+                queue_status = self.get_queue_status(queue_id=queue_id)
+                self.__invoker.services.events.emit_queue_item_status_changed(
+                    current_queue_item, batch_status, queue_status
+                )
         except Exception:
             self.__conn.rollback()
             raise
@@ -426,46 +429,6 @@ class SqliteSessionQueue(SessionQueueBase):
             self.__lock.release()
         return CancelByBatchIDsResult(canceled=count)
 
-    def cancel_by_origin(self, queue_id: str, origin: str) -> CancelByOriginResult:
-        try:
-            current_queue_item = self.get_current(queue_id)
-            self.__lock.acquire()
-            where = """--sql
-                WHERE
-                  queue_id == ?
-                  AND origin == ?
-                  AND status != 'canceled'
-                  AND status != 'completed'
-                  AND status != 'failed'
-                """
-            params = (queue_id, origin)
-            self.__cursor.execute(
-                f"""--sql
-                SELECT COUNT(*)
-                FROM session_queue
-                {where};
-                """,
-                params,
-            )
-            count = self.__cursor.fetchone()[0]
-            self.__cursor.execute(
-                f"""--sql
-                UPDATE session_queue
-                SET status = 'canceled'
-                {where};
-                """,
-                params,
-            )
-            self.__conn.commit()
-            if current_queue_item is not None and current_queue_item.origin == origin:
-                self._set_queue_item_status(current_queue_item.item_id, "canceled")
-        except Exception:
-            self.__conn.rollback()
-            raise
-        finally:
-            self.__lock.release()
-        return CancelByOriginResult(canceled=count)
-
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
         try:
             current_queue_item = self.get_current(queue_id)
@@ -578,8 +541,7 @@ class SqliteSessionQueue(SessionQueueBase):
                     started_at,
                     session_id,
                     batch_id,
-                    queue_id,
-                    origin
+                    queue_id
                 FROM session_queue
                 WHERE queue_id = ?
             """
@@ -659,7 +621,7 @@ class SqliteSessionQueue(SessionQueueBase):
             self.__lock.acquire()
             self.__cursor.execute(
                 """--sql
-                SELECT status, count(*), origin
+                SELECT status, count(*)
                 FROM session_queue
                 WHERE
                   queue_id = ?
@@ -671,7 +633,6 @@ class SqliteSessionQueue(SessionQueueBase):
             result = cast(list[sqlite3.Row], self.__cursor.fetchall())
             total = sum(row[1] for row in result)
             counts: dict[str, int] = {row[0]: row[1] for row in result}
-            origin = result[0]["origin"] if result else None
         except Exception:
             self.__conn.rollback()
             raise
@@ -680,7 +641,6 @@ class SqliteSessionQueue(SessionQueueBase):
 
         return BatchStatus(
             batch_id=batch_id,
-            origin=origin,
             queue_id=queue_id,
             pending=counts.get("pending", 0),
             in_progress=counts.get("in_progress", 0),

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -16,8 +16,6 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_10 import
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_11 import build_migration_11
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_12 import build_migration_12
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_13 import build_migration_13
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_14 import build_migration_14
-from invokeai.app.services.shared.sqlite_migrator.migrations.migration_15 import build_migration_15
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -51,8 +49,6 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_11(app_config=config, logger=logger))
     migrator.register_migration(build_migration_12(app_config=config))
     migrator.register_migration(build_migration_13())
-    migrator.register_migration(build_migration_14())
-    migrator.register_migration(build_migration_15())
     migrator.run_migrations()
 
     return db

invokeai/app/services/shared/sqlite_migrator/migrations/migration_14.py

@@ -1,61 +0,0 @@
-import sqlite3
-
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-
-class Migration14Callback:
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._create_style_presets(cursor)
-
-    def _create_style_presets(self, cursor: sqlite3.Cursor) -> None:
-        """Create the table used to store style presets."""
-        tables = [
-            """--sql
-            CREATE TABLE IF NOT EXISTS style_presets (
-                id TEXT NOT NULL PRIMARY KEY,
-                name TEXT NOT NULL,
-                preset_data TEXT NOT NULL,
-                type TEXT NOT NULL DEFAULT "user",
-                created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
-                -- Updated via trigger
-                updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW'))
-            );
-            """
-        ]
-
-        # Add trigger for `updated_at`.
-        triggers = [
-            """--sql
-            CREATE TRIGGER IF NOT EXISTS style_presets
-            AFTER UPDATE
-            ON style_presets FOR EACH ROW
-            BEGIN
-                UPDATE style_presets SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
-                    WHERE id = old.id;
-            END;
-            """
-        ]
-
-        # Add indexes for searchable fields
-        indices = [
-            "CREATE INDEX IF NOT EXISTS idx_style_presets_name ON style_presets(name);",
-        ]
-
-        for stmt in tables + indices + triggers:
-            cursor.execute(stmt)
-
-
-def build_migration_14() -> Migration:
-    """
-    Build the migration from database version 13 to 14..
-
-    This migration does the following:
-    - Create the table used to store style presets.
-    """
-    migration_14 = Migration(
-        from_version=13,
-        to_version=14,
-        callback=Migration14Callback(),
-    )
-
-    return migration_14

invokeai/app/services/shared/sqlite_migrator/migrations/migration_15.py

@@ -1,31 +0,0 @@
-import sqlite3
-
-from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
-
-
-class Migration15Callback:
-    def __call__(self, cursor: sqlite3.Cursor) -> None:
-        self._add_origin_col(cursor)
-
-    def _add_origin_col(self, cursor: sqlite3.Cursor) -> None:
-        """
-        - Adds `origin` column to the session queue table.
-        """
-
-        cursor.execute("ALTER TABLE session_queue ADD COLUMN origin TEXT;")
-
-
-def build_migration_15() -> Migration:
-    """
-    Build the migration from database version 14 to 15.
-
-    This migration does the following:
-        - Adds `origin` column to the session queue table.
-    """
-    migration_15 = Migration(
-        from_version=14,
-        to_version=15,
-        callback=Migration15Callback(),
-    )
-
-    return migration_15

invokeai/app/services/style_preset_images/style_preset_images_base.py

@@ -1,33 +0,0 @@
-from abc import ABC, abstractmethod
-from pathlib import Path
-
-from PIL.Image import Image as PILImageType
-
-
-class StylePresetImageFileStorageBase(ABC):
-    """Low-level service responsible for storing and retrieving image files."""
-
-    @abstractmethod
-    def get(self, style_preset_id: str) -> PILImageType:
-        """Retrieves a style preset image as PIL Image."""
-        pass
-
-    @abstractmethod
-    def get_path(self, style_preset_id: str) -> Path:
-        """Gets the internal path to a style preset image."""
-        pass
-
-    @abstractmethod
-    def get_url(self, style_preset_id: str) -> str | None:
-        """Gets the URL to fetch a style preset image."""
-        pass
-
-    @abstractmethod
-    def save(self, style_preset_id: str, image: PILImageType) -> None:
-        """Saves a style preset image."""
-        pass
-
-    @abstractmethod
-    def delete(self, style_preset_id: str) -> None:
-        """Deletes a style preset image."""
-        pass

invokeai/app/services/style_preset_images/style_preset_images_common.py

@@ -1,19 +0,0 @@
-class StylePresetImageFileNotFoundException(Exception):
-    """Raised when an image file is not found in storage."""
-
-    def __init__(self, message: str = "Style preset image file not found"):
-        super().__init__(message)
-
-
-class StylePresetImageFileSaveException(Exception):
-    """Raised when an image cannot be saved."""
-
-    def __init__(self, message: str = "Style preset image file not saved"):
-        super().__init__(message)
-
-
-class StylePresetImageFileDeleteException(Exception):
-    """Raised when an image cannot be deleted."""
-
-    def __init__(self, message: str = "Style preset image file not deleted"):
-        super().__init__(message)

invokeai/app/services/style_preset_images/style_preset_images_disk.py

@@ -1,88 +0,0 @@
-from pathlib import Path
-
-from PIL import Image
-from PIL.Image import Image as PILImageType
-
-from invokeai.app.services.invoker import Invoker
-from invokeai.app.services.style_preset_images.style_preset_images_base import StylePresetImageFileStorageBase
-from invokeai.app.services.style_preset_images.style_preset_images_common import (
-    StylePresetImageFileDeleteException,
-    StylePresetImageFileNotFoundException,
-    StylePresetImageFileSaveException,
-)
-from invokeai.app.services.style_preset_records.style_preset_records_common import PresetType
-from invokeai.app.util.misc import uuid_string
-from invokeai.app.util.thumbnails import make_thumbnail
-
-
-class StylePresetImageFileStorageDisk(StylePresetImageFileStorageBase):
-    """Stores images on disk"""
-
-    def __init__(self, style_preset_images_folder: Path):
-        self._style_preset_images_folder = style_preset_images_folder
-        self._validate_storage_folders()
-
-    def start(self, invoker: Invoker) -> None:
-        self._invoker = invoker
-
-    def get(self, style_preset_id: str) -> PILImageType:
-        try:
-            path = self.get_path(style_preset_id)
-
-            return Image.open(path)
-        except FileNotFoundError as e:
-            raise StylePresetImageFileNotFoundException from e
-
-    def save(self, style_preset_id: str, image: PILImageType) -> None:
-        try:
-            self._validate_storage_folders()
-            image_path = self._style_preset_images_folder / (style_preset_id + ".webp")
-            thumbnail = make_thumbnail(image, 256)
-            thumbnail.save(image_path, format="webp")
-
-        except Exception as e:
-            raise StylePresetImageFileSaveException from e
-
-    def get_path(self, style_preset_id: str) -> Path:
-        style_preset = self._invoker.services.style_preset_records.get(style_preset_id)
-        if style_preset.type is PresetType.Default:
-            default_images_dir = Path(__file__).parent / Path("default_style_preset_images")
-            path = default_images_dir / (style_preset.name + ".png")
-        else:
-            path = self._style_preset_images_folder / (style_preset_id + ".webp")
-
-        return path
-
-    def get_url(self, style_preset_id: str) -> str | None:
-        path = self.get_path(style_preset_id)
-        if not self._validate_path(path):
-            return
-
-        url = self._invoker.services.urls.get_style_preset_image_url(style_preset_id)
-
-        # The image URL never changes, so we must add random query string to it to prevent caching
-        url += f"?{uuid_string()}"
-
-        return url
-
-    def delete(self, style_preset_id: str) -> None:
-        try:
-            path = self.get_path(style_preset_id)
-
-            if not self._validate_path(path):
-                raise StylePresetImageFileNotFoundException
-
-            path.unlink()
-
-        except StylePresetImageFileNotFoundException as e:
-            raise StylePresetImageFileNotFoundException from e
-        except Exception as e:
-            raise StylePresetImageFileDeleteException from e
-
-    def _validate_path(self, path: Path) -> bool:
-        """Validates the path given for an image."""
-        return path.exists()
-
-    def _validate_storage_folders(self) -> None:
-        """Checks if the required folders exist and create them if they don't"""
-        self._style_preset_images_folder.mkdir(parents=True, exist_ok=True)

invokeai/app/services/style_preset_records/default_style_presets.json

@@ -1,146 +0,0 @@
-[
-  {
-    "name": "Photography (General)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt}. photography. f/2.8 macro photo, bokeh, photorealism",
-      "negative_prompt": "painting, digital art. sketch, blurry"
-    }
-  },
-  {
-    "name": "Photography (Studio Lighting)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt}, photography. f/8 photo. centered subject, studio lighting.",
-      "negative_prompt": "painting, digital art. sketch, blurry"
-    }
-  },
-  {
-    "name": "Photography (Landscape)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt}, landscape photograph, f/12, lifelike,  highly detailed.",
-      "negative_prompt": "painting, digital art. sketch, blurry"
-    }
-  },
-  {
-    "name": "Photography (Portrait)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt}. photography. portraiture. catch light in eyes. one flash. rembrandt lighting. Soft box. dark shadows. High contrast. 80mm lens. F2.8.",
-      "negative_prompt": "painting, digital art. sketch, blurry"
-    }
-  },
-  {
-    "name": "Photography (Black and White)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} photography. natural light. 80mm lens. F1.4. strong contrast, hard light. dark contrast. blurred background. black and white",
-      "negative_prompt": "painting, digital art. sketch, colour+"
-    }
-  },
-  {
-    "name": "Architectural Visualization",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt}. architectural photography, f/12, luxury, aesthetically pleasing form and function.",
-      "negative_prompt": "painting, digital art. sketch, blurry"
-    }
-  },
-  {
-    "name": "Concept Art (Fantasy)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "concept artwork of a {prompt}. (digital painterly art style)++, mythological, (textured 2d dry media brushpack)++, glazed brushstrokes, otherworldly. painting+, illustration+",
-      "negative_prompt": "photo. distorted, blurry, out of focus. sketch. (cgi, 3d.)++"
-    }
-  },
-  {
-    "name": "Concept Art (Sci-Fi)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "(concept art)++, {prompt}, (sleek futurism)++, (textured 2d dry media)++, metallic highlights, digital painting style",
-      "negative_prompt": "photo. distorted, blurry, out of focus. sketch. (cgi, 3d.)++"
-    }
-  },
-  {
-    "name": "Concept Art (Character)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "(character concept art)++, stylized painterly digital painting of {prompt}, (painterly, impasto. Dry brush.)++",
-      "negative_prompt": "photo. distorted, blurry, out of focus. sketch. (cgi, 3d.)++"
-    }
-  },
-  {
-    "name": "Concept Art (Painterly)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} oil painting. high contrast. impasto. sfumato. chiaroscuro. Palette knife.",
-      "negative_prompt": "photo. smooth. border. frame"
-    }
-  },
-  {
-    "name": "Environment Art",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} environment artwork, hyper-realistic digital painting style with cinematic composition, atmospheric, depth and detail, voluminous. textured dry brush 2d media",
-      "negative_prompt": "photo, distorted, blurry, out of focus. sketch."
-    }
-  },
-  {
-    "name": "Interior Design (Visualization)",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} interior design photo, gentle shadows, light mid-tones, dimension, mix of smooth and textured surfaces, focus on negative space and clean lines, focus",
-      "negative_prompt": "photo, distorted. sketch."
-    }
-  },
-  {
-    "name": "Product Rendering",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} high quality product photography, 3d rendering with key lighting, shallow depth of field, simple plain background, studio lighting.",
-      "negative_prompt": "blurry, sketch, messy, dirty. unfinished."
-    }
-  },
-  {
-    "name": "Sketch",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} black and white pencil drawing, off-center composition, cross-hatching for shadows, bold strokes, textured paper. sketch+++",
-      "negative_prompt": "blurry, photo, painting, color. messy, dirty. unfinished. frame, borders."
-    }
-  },
-  {
-    "name": "Line Art",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} Line art. bold outline. simplistic. white background. 2d",
-      "negative_prompt": "photo. digital art. greyscale. solid black. painting"
-    }
-  },
-  {
-    "name": "Anime",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} anime++, bold outline, cel-shaded coloring, shounen, seinen",
-      "negative_prompt": "(photo)+++. greyscale. solid black. painting"
-    }
-  },
-  {
-    "name": "Illustration",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "{prompt} illustration, bold linework, illustrative details, vector art style, flat coloring",
-      "negative_prompt": "(photo)+++. greyscale. painting, black and white."
-    }
-  },
-  {
-    "name": "Vehicles",
-    "type": "default",
-    "preset_data": {
-      "positive_prompt": "A weird futuristic normal auto, {prompt} elegant design, nice color, nice wheels",
-      "negative_prompt": "sketch. digital art. greyscale. painting"
-    }
-  }
-]

invokeai/app/services/style_preset_records/style_preset_records_base.py

@@ -1,42 +0,0 @@
-from abc import ABC, abstractmethod
-
-from invokeai.app.services.style_preset_records.style_preset_records_common import (
-    PresetType,
-    StylePresetChanges,
-    StylePresetRecordDTO,
-    StylePresetWithoutId,
-)
-
-
-class StylePresetRecordsStorageBase(ABC):
-    """Base class for style preset storage services."""
-
-    @abstractmethod
-    def get(self, style_preset_id: str) -> StylePresetRecordDTO:
-        """Get style preset by id."""
-        pass
-
-    @abstractmethod
-    def create(self, style_preset: StylePresetWithoutId) -> StylePresetRecordDTO:
-        """Creates a style preset."""
-        pass
-
-    @abstractmethod
-    def create_many(self, style_presets: list[StylePresetWithoutId]) -> None:
-        """Creates many style presets."""
-        pass
-
-    @abstractmethod
-    def update(self, style_preset_id: str, changes: StylePresetChanges) -> StylePresetRecordDTO:
-        """Updates a style preset."""
-        pass
-
-    @abstractmethod
-    def delete(self, style_preset_id: str) -> None:
-        """Deletes a style preset."""
-        pass
-
-    @abstractmethod
-    def get_many(self, type: PresetType | None = None) -> list[StylePresetRecordDTO]:
-        """Gets many workflows."""
-        pass

invokeai/app/services/style_preset_records/style_preset_records_common.py

@@ -1,139 +0,0 @@
-import codecs
-import csv
-import json
-from enum import Enum
-from typing import Any, Optional
-
-import pydantic
-from fastapi import UploadFile
-from pydantic import AliasChoices, BaseModel, ConfigDict, Field, TypeAdapter
-
-from invokeai.app.util.metaenum import MetaEnum
-
-
-class StylePresetNotFoundError(Exception):
-    """Raised when a style preset is not found"""
-
-
-class PresetData(BaseModel, extra="forbid"):
-    positive_prompt: str = Field(description="Positive prompt")
-    negative_prompt: str = Field(description="Negative prompt")
-
-
-PresetDataValidator = TypeAdapter(PresetData)
-
-
-class PresetType(str, Enum, metaclass=MetaEnum):
-    User = "user"
-    Default = "default"
-    Project = "project"
-
-
-class StylePresetChanges(BaseModel, extra="forbid"):
-    name: Optional[str] = Field(default=None, description="The style preset's new name.")
-    preset_data: Optional[PresetData] = Field(default=None, description="The updated data for style preset.")
-    type: Optional[PresetType] = Field(description="The updated type of the style preset")
-
-
-class StylePresetWithoutId(BaseModel):
-    name: str = Field(description="The name of the style preset.")
-    preset_data: PresetData = Field(description="The preset data")
-    type: PresetType = Field(description="The type of style preset")
-
-
-class StylePresetRecordDTO(StylePresetWithoutId):
-    id: str = Field(description="The style preset ID.")
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "StylePresetRecordDTO":
-        data["preset_data"] = PresetDataValidator.validate_json(data.get("preset_data", ""))
-        return StylePresetRecordDTOValidator.validate_python(data)
-
-
-StylePresetRecordDTOValidator = TypeAdapter(StylePresetRecordDTO)
-
-
-class StylePresetRecordWithImage(StylePresetRecordDTO):
-    image: Optional[str] = Field(description="The path for image")
-
-
-class StylePresetImportRow(BaseModel):
-    name: str = Field(min_length=1, description="The name of the preset.")
-    positive_prompt: str = Field(
-        default="",
-        description="The positive prompt for the preset.",
-        validation_alias=AliasChoices("positive_prompt", "prompt"),
-    )
-    negative_prompt: str = Field(default="", description="The negative prompt for the preset.")
-
-    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
-
-
-StylePresetImportList = list[StylePresetImportRow]
-StylePresetImportListTypeAdapter = TypeAdapter(StylePresetImportList)
-
-
-class UnsupportedFileTypeError(ValueError):
-    """Raised when an unsupported file type is encountered"""
-
-    pass
-
-
-class InvalidPresetImportDataError(ValueError):
-    """Raised when invalid preset import data is encountered"""
-
-    pass
-
-
-async def parse_presets_from_file(file: UploadFile) -> list[StylePresetWithoutId]:
-    """Parses style presets from a file. The file must be a CSV or JSON file.
-
-    If CSV, the file must have the following columns:
-    - name
-    - prompt (or positive_prompt)
-    - negative_prompt
-
-    If JSON, the file must be a list of objects with the following keys:
-    - name
-    - prompt (or positive_prompt)
-    - negative_prompt
-
-    Args:
-        file (UploadFile): The file to parse.
-
-    Returns:
-        list[StylePresetWithoutId]: The parsed style presets.
-
-    Raises:
-        UnsupportedFileTypeError: If the file type is not supported.
-        InvalidPresetImportDataError: If the data in the file is invalid.
-    """
-    if file.content_type not in ["text/csv", "application/json"]:
-        raise UnsupportedFileTypeError()
-
-    if file.content_type == "text/csv":
-        csv_reader = csv.DictReader(codecs.iterdecode(file.file, "utf-8"))
-        data = list(csv_reader)
-    else:  # file.content_type == "application/json":
-        json_data = await file.read()
-        data = json.loads(json_data)
-
-    try:
-        imported_presets = StylePresetImportListTypeAdapter.validate_python(data)
-
-        style_presets: list[StylePresetWithoutId] = []
-
-        for imported in imported_presets:
-            preset_data = PresetData(positive_prompt=imported.positive_prompt, negative_prompt=imported.negative_prompt)
-            style_preset = StylePresetWithoutId(name=imported.name, preset_data=preset_data, type=PresetType.User)
-            style_presets.append(style_preset)
-    except pydantic.ValidationError as e:
-        if file.content_type == "text/csv":
-            msg = "Invalid CSV format: must include columns 'name', 'prompt', and 'negative_prompt' and name cannot be blank"
-        else:  # file.content_type == "application/json":
-            msg = "Invalid JSON format: must be a list of objects with keys 'name', 'prompt', and 'negative_prompt' and name cannot be blank"
-        raise InvalidPresetImportDataError(msg) from e
-    finally:
-        file.file.close()
-
-    return style_presets

invokeai/app/services/style_preset_records/style_preset_records_sqlite.py

@@ -1,215 +0,0 @@
-import json
-from pathlib import Path
-
-from invokeai.app.services.invoker import Invoker
-from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
-from invokeai.app.services.style_preset_records.style_preset_records_base import StylePresetRecordsStorageBase
-from invokeai.app.services.style_preset_records.style_preset_records_common import (
-    PresetType,
-    StylePresetChanges,
-    StylePresetNotFoundError,
-    StylePresetRecordDTO,
-    StylePresetWithoutId,
-)
-from invokeai.app.util.misc import uuid_string
-
-
-class SqliteStylePresetRecordsStorage(StylePresetRecordsStorageBase):
-    def __init__(self, db: SqliteDatabase) -> None:
-        super().__init__()
-        self._lock = db.lock
-        self._conn = db.conn
-        self._cursor = self._conn.cursor()
-
-    def start(self, invoker: Invoker) -> None:
-        self._invoker = invoker
-        self._sync_default_style_presets()
-
-    def get(self, style_preset_id: str) -> StylePresetRecordDTO:
-        """Gets a style preset by ID."""
-        try:
-            self._lock.acquire()
-            self._cursor.execute(
-                """--sql
-                SELECT *
-                FROM style_presets
-                WHERE id = ?;
-                """,
-                (style_preset_id,),
-            )
-            row = self._cursor.fetchone()
-            if row is None:
-                raise StylePresetNotFoundError(f"Style preset with id {style_preset_id} not found")
-            return StylePresetRecordDTO.from_dict(dict(row))
-        except Exception:
-            self._conn.rollback()
-            raise
-        finally:
-            self._lock.release()
-
-    def create(self, style_preset: StylePresetWithoutId) -> StylePresetRecordDTO:
-        style_preset_id = uuid_string()
-        try:
-            self._lock.acquire()
-            self._cursor.execute(
-                """--sql
-                INSERT OR IGNORE INTO style_presets (
-                    id,
-                    name,
-                    preset_data,
-                    type
-                )
-                VALUES (?, ?, ?, ?);
-                """,
-                (
-                    style_preset_id,
-                    style_preset.name,
-                    style_preset.preset_data.model_dump_json(),
-                    style_preset.type,
-                ),
-            )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
-        finally:
-            self._lock.release()
-        return self.get(style_preset_id)
-
-    def create_many(self, style_presets: list[StylePresetWithoutId]) -> None:
-        style_preset_ids = []
-        try:
-            self._lock.acquire()
-            for style_preset in style_presets:
-                style_preset_id = uuid_string()
-                style_preset_ids.append(style_preset_id)
-                self._cursor.execute(
-                    """--sql
-                    INSERT OR IGNORE INTO style_presets (
-                        id,
-                        name,
-                        preset_data,
-                        type
-                    )
-                    VALUES (?, ?, ?, ?);
-                    """,
-                    (
-                        style_preset_id,
-                        style_preset.name,
-                        style_preset.preset_data.model_dump_json(),
-                        style_preset.type,
-                    ),
-                )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
-        finally:
-            self._lock.release()
-
-        return None
-
-    def update(self, style_preset_id: str, changes: StylePresetChanges) -> StylePresetRecordDTO:
-        try:
-            self._lock.acquire()
-            # Change the name of a style preset
-            if changes.name is not None:
-                self._cursor.execute(
-                    """--sql
-                    UPDATE style_presets
-                    SET name = ?
-                    WHERE id = ?;
-                    """,
-                    (changes.name, style_preset_id),
-                )
-
-            # Change the preset data for a style preset
-            if changes.preset_data is not None:
-                self._cursor.execute(
-                    """--sql
-                    UPDATE style_presets
-                    SET preset_data = ?
-                    WHERE id = ?;
-                    """,
-                    (changes.preset_data.model_dump_json(), style_preset_id),
-                )
-
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
-        finally:
-            self._lock.release()
-        return self.get(style_preset_id)
-
-    def delete(self, style_preset_id: str) -> None:
-        try:
-            self._lock.acquire()
-            self._cursor.execute(
-                """--sql
-                DELETE from style_presets
-                WHERE id = ?;
-                """,
-                (style_preset_id,),
-            )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
-        finally:
-            self._lock.release()
-        return None
-
-    def get_many(self, type: PresetType | None = None) -> list[StylePresetRecordDTO]:
-        try:
-            self._lock.acquire()
-            main_query = """
-                SELECT
-                    *
-                FROM style_presets
-                """
-
-            if type is not None:
-                main_query += "WHERE type = ? "
-
-            main_query += "ORDER BY LOWER(name) ASC"
-
-            if type is not None:
-                self._cursor.execute(main_query, (type,))
-            else:
-                self._cursor.execute(main_query)
-
-            rows = self._cursor.fetchall()
-            style_presets = [StylePresetRecordDTO.from_dict(dict(row)) for row in rows]
-
-            return style_presets
-        except Exception:
-            self._conn.rollback()
-            raise
-        finally:
-            self._lock.release()
-
-    def _sync_default_style_presets(self) -> None:
-        """Syncs default style presets to the database. Internal use only."""
-
-        # First delete all existing default style presets
-        try:
-            self._lock.acquire()
-            self._cursor.execute(
-                """--sql
-                DELETE FROM style_presets
-                WHERE type = "default";
-                """
-            )
-            self._conn.commit()
-        except Exception:
-            self._conn.rollback()
-            raise
-        finally:
-            self._lock.release()
-        # Next, parse and create the default style presets
-        with self._lock, open(Path(__file__).parent / Path("default_style_presets.json"), "r") as file:
-            presets = json.load(file)
-            for preset in presets:
-                style_preset = StylePresetWithoutId.model_validate(preset)
-                self.create(style_preset)

invokeai/app/services/urls/urls_base.py

@@ -13,8 +13,3 @@ class UrlServiceBase(ABC):
     def get_model_image_url(self, model_key: str) -> str:
         """Gets the URL for a model image"""
         pass
-
-    @abstractmethod
-    def get_style_preset_image_url(self, style_preset_id: str) -> str:
-        """Gets the URL for a style preset image"""
-        pass

invokeai/app/services/urls/urls_default.py

@@ -19,6 +19,3 @@ class LocalUrlService(UrlServiceBase):
 
     def get_model_image_url(self, model_key: str) -> str:
         return f"{self._base_url_v2}/models/i/{model_key}/image"
-
-    def get_style_preset_image_url(self, style_preset_id: str) -> str:
-        return f"{self._base_url}/style_presets/i/{style_preset_id}/image"

invokeai/app/util/custom_openapi.py

@@ -81,7 +81,7 @@ def get_openapi_func(
         # Add the output map to the schema
         openapi_schema["components"]["schemas"]["InvocationOutputMap"] = {
             "type": "object",
-            "properties": dict(sorted(invocation_output_map_properties.items())),
+            "properties": invocation_output_map_properties,
             "required": invocation_output_map_required,
         }
 

invokeai/backend/image_util/depth_anything/__init__.py

@@ -0,0 +1,90 @@
+from pathlib import Path
+from typing import Literal
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import repeat
+from PIL import Image
+from torchvision.transforms import Compose
+
+from invokeai.app.services.config.config_default import get_config
+from invokeai.backend.image_util.depth_anything.model.dpt import DPT_DINOv2
+from invokeai.backend.image_util.depth_anything.utilities.util import NormalizeImage, PrepareForNet, Resize
+from invokeai.backend.util.logging import InvokeAILogger
+
+config = get_config()
+logger = InvokeAILogger.get_logger(config=config)
+
+DEPTH_ANYTHING_MODELS = {
+    "large": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitl14.pth?download=true",
+    "base": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vitb14.pth?download=true",
+    "small": "https://huggingface.co/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vits14.pth?download=true",
+}
+
+
+transform = Compose(
+    [
+        Resize(
+            width=518,
+            height=518,
+            resize_target=False,
+            keep_aspect_ratio=True,
+            ensure_multiple_of=14,
+            resize_method="lower_bound",
+            image_interpolation_method=cv2.INTER_CUBIC,
+        ),
+        NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        PrepareForNet(),
+    ]
+)
+
+
+class DepthAnythingDetector:
+    def __init__(self, model: DPT_DINOv2, device: torch.device) -> None:
+        self.model = model
+        self.device = device
+
+    @staticmethod
+    def load_model(
+        model_path: Path, device: torch.device, model_size: Literal["large", "base", "small"] = "small"
+    ) -> DPT_DINOv2:
+        match model_size:
+            case "small":
+                model = DPT_DINOv2(encoder="vits", features=64, out_channels=[48, 96, 192, 384])
+            case "base":
+                model = DPT_DINOv2(encoder="vitb", features=128, out_channels=[96, 192, 384, 768])
+            case "large":
+                model = DPT_DINOv2(encoder="vitl", features=256, out_channels=[256, 512, 1024, 1024])
+
+        model.load_state_dict(torch.load(model_path.as_posix(), map_location="cpu"))
+        model.eval()
+
+        model.to(device)
+        return model
+
+    def __call__(self, image: Image.Image, resolution: int = 512) -> Image.Image:
+        if not self.model:
+            logger.warn("DepthAnything model was not loaded. Returning original image")
+            return image
+
+        np_image = np.array(image, dtype=np.uint8)
+        np_image = np_image[:, :, ::-1] / 255.0
+
+        image_height, image_width = np_image.shape[:2]
+        np_image = transform({"image": np_image})["image"]
+        tensor_image = torch.from_numpy(np_image).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            depth = self.model(tensor_image)
+            depth = F.interpolate(depth[None], (image_height, image_width), mode="bilinear", align_corners=False)[0, 0]
+            depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
+
+        depth_map = repeat(depth, "h w -> h w 3").cpu().numpy().astype(np.uint8)
+        depth_map = Image.fromarray(depth_map)
+
+        new_height = int(image_height * (resolution / image_width))
+        depth_map = depth_map.resize((resolution, new_height))
+
+        return depth_map

invokeai/backend/image_util/depth_anything/depth_anything_pipeline.py

@@ -1,31 +0,0 @@
-from typing import Optional
-
-import torch
-from PIL import Image
-from transformers.pipelines import DepthEstimationPipeline
-
-from invokeai.backend.raw_model import RawModel
-
-
-class DepthAnythingPipeline(RawModel):
-    """Custom wrapper for the Depth Estimation pipeline from transformers adding compatibility
-    for Invoke's Model Management System"""
-
-    def __init__(self, pipeline: DepthEstimationPipeline) -> None:
-        self._pipeline = pipeline
-
-    def generate_depth(self, image: Image.Image) -> Image.Image:
-        depth_map = self._pipeline(image)["depth"]
-        assert isinstance(depth_map, Image.Image)
-        return depth_map
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
-        if device is not None and device.type not in {"cpu", "cuda"}:
-            device = None
-        self._pipeline.model.to(device=device, dtype=dtype)
-        self._pipeline.device = self._pipeline.model.device
-
-    def calc_size(self) -> int:
-        from invokeai.backend.model_manager.load.model_util import calc_module_size
-
-        return calc_module_size(self._pipeline.model)

invokeai/backend/image_util/depth_anything/model/blocks.py

@@ -0,0 +1,145 @@
+import torch.nn as nn
+
+
+def _make_scratch(in_shape, out_shape, groups=1, expand=False):
+    scratch = nn.Module()
+
+    out_shape1 = out_shape
+    out_shape2 = out_shape
+    out_shape3 = out_shape
+    if len(in_shape) >= 4:
+        out_shape4 = out_shape
+
+    if expand:
+        out_shape1 = out_shape
+        out_shape2 = out_shape * 2
+        out_shape3 = out_shape * 4
+        if len(in_shape) >= 4:
+            out_shape4 = out_shape * 8
+
+    scratch.layer1_rn = nn.Conv2d(
+        in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    scratch.layer2_rn = nn.Conv2d(
+        in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    scratch.layer3_rn = nn.Conv2d(
+        in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    if len(in_shape) >= 4:
+        scratch.layer4_rn = nn.Conv2d(
+            in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+        )
+
+    return scratch
+
+
+class ResidualConvUnit(nn.Module):
+    """Residual convolution module."""
+
+    def __init__(self, features, activation, bn):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super().__init__()
+
+        self.bn = bn
+
+        self.groups = 1
+
+        self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+
+        self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+
+        if self.bn:
+            self.bn1 = nn.BatchNorm2d(features)
+            self.bn2 = nn.BatchNorm2d(features)
+
+        self.activation = activation
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (tensor): input
+
+        Returns:
+            tensor: output
+        """
+
+        out = self.activation(x)
+        out = self.conv1(out)
+        if self.bn:
+            out = self.bn1(out)
+
+        out = self.activation(out)
+        out = self.conv2(out)
+        if self.bn:
+            out = self.bn2(out)
+
+        if self.groups > 1:
+            out = self.conv_merge(out)
+
+        return self.skip_add.add(out, x)
+
+
+class FeatureFusionBlock(nn.Module):
+    """Feature fusion block."""
+
+    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
+        """Init.
+
+        Args:
+            features (int): number of features
+        """
+        super(FeatureFusionBlock, self).__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+
+        self.groups = 1
+
+        self.expand = expand
+        out_features = features
+        if self.expand:
+            out_features = features // 2
+
+        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
+
+        self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
+        self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+        self.size = size
+
+    def forward(self, *xs, size=None):
+        """Forward pass.
+
+        Returns:
+            tensor: output
+        """
+        output = xs[0]
+
+        if len(xs) == 2:
+            res = self.resConfUnit1(xs[1])
+            output = self.skip_add.add(output, res)
+
+        output = self.resConfUnit2(output)
+
+        if (size is None) and (self.size is None):
+            modifier = {"scale_factor": 2}
+        elif size is None:
+            modifier = {"size": self.size}
+        else:
+            modifier = {"size": size}
+
+        output = nn.functional.interpolate(output, **modifier, mode="bilinear", align_corners=self.align_corners)
+
+        output = self.out_conv(output)
+
+        return output

invokeai/backend/image_util/depth_anything/model/dpt.py

@@ -0,0 +1,183 @@
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from invokeai.backend.image_util.depth_anything.model.blocks import FeatureFusionBlock, _make_scratch
+
+torchhub_path = Path(__file__).parent.parent / "torchhub"
+
+
+def _make_fusion_block(features, use_bn, size=None):
+    return FeatureFusionBlock(
+        features,
+        nn.ReLU(False),
+        deconv=False,
+        bn=use_bn,
+        expand=False,
+        align_corners=True,
+        size=size,
+    )
+
+
+class DPTHead(nn.Module):
+    def __init__(self, nclass, in_channels, features, out_channels, use_bn=False, use_clstoken=False):
+        super(DPTHead, self).__init__()
+
+        self.nclass = nclass
+        self.use_clstoken = use_clstoken
+
+        self.projects = nn.ModuleList(
+            [
+                nn.Conv2d(
+                    in_channels=in_channels,
+                    out_channels=out_channel,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                )
+                for out_channel in out_channels
+            ]
+        )
+
+        self.resize_layers = nn.ModuleList(
+            [
+                nn.ConvTranspose2d(
+                    in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=4, stride=4, padding=0
+                ),
+                nn.ConvTranspose2d(
+                    in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0
+                ),
+                nn.Identity(),
+                nn.Conv2d(
+                    in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=3, stride=2, padding=1
+                ),
+            ]
+        )
+
+        if use_clstoken:
+            self.readout_projects = nn.ModuleList()
+            for _ in range(len(self.projects)):
+                self.readout_projects.append(nn.Sequential(nn.Linear(2 * in_channels, in_channels), nn.GELU()))
+
+        self.scratch = _make_scratch(
+            out_channels,
+            features,
+            groups=1,
+            expand=False,
+        )
+
+        self.scratch.stem_transpose = None
+
+        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
+        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
+
+        head_features_1 = features
+        head_features_2 = 32
+
+        if nclass > 1:
+            self.scratch.output_conv = nn.Sequential(
+                nn.Conv2d(head_features_1, head_features_1, kernel_size=3, stride=1, padding=1),
+                nn.ReLU(True),
+                nn.Conv2d(head_features_1, nclass, kernel_size=1, stride=1, padding=0),
+            )
+        else:
+            self.scratch.output_conv1 = nn.Conv2d(
+                head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1
+            )
+
+            self.scratch.output_conv2 = nn.Sequential(
+                nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
+                nn.ReLU(True),
+                nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
+                nn.ReLU(True),
+                nn.Identity(),
+            )
+
+    def forward(self, out_features, patch_h, patch_w):
+        out = []
+        for i, x in enumerate(out_features):
+            if self.use_clstoken:
+                x, cls_token = x[0], x[1]
+                readout = cls_token.unsqueeze(1).expand_as(x)
+                x = self.readout_projects[i](torch.cat((x, readout), -1))
+            else:
+                x = x[0]
+
+            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+
+            x = self.projects[i](x)
+            x = self.resize_layers[i](x)
+
+            out.append(x)
+
+        layer_1, layer_2, layer_3, layer_4 = out
+
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+        path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
+        path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
+        path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
+        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
+
+        out = self.scratch.output_conv1(path_1)
+        out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
+        out = self.scratch.output_conv2(out)
+
+        return out
+
+
+class DPT_DINOv2(nn.Module):
+    def __init__(
+        self,
+        features,
+        out_channels,
+        encoder="vitl",
+        use_bn=False,
+        use_clstoken=False,
+    ):
+        super(DPT_DINOv2, self).__init__()
+
+        assert encoder in ["vits", "vitb", "vitl"]
+
+        # # in case the Internet connection is not stable, please load the DINOv2 locally
+        # if use_local:
+        #     self.pretrained = torch.hub.load(
+        #         torchhub_path / "facebookresearch_dinov2_main",
+        #         "dinov2_{:}14".format(encoder),
+        #         source="local",
+        #         pretrained=False,
+        #     )
+        # else:
+        #     self.pretrained = torch.hub.load(
+        #         "facebookresearch/dinov2",
+        #         "dinov2_{:}14".format(encoder),
+        #     )
+
+        self.pretrained = torch.hub.load(
+            "facebookresearch/dinov2",
+            "dinov2_{:}14".format(encoder),
+        )
+
+        dim = self.pretrained.blocks[0].attn.qkv.in_features
+
+        self.depth_head = DPTHead(1, dim, features, out_channels=out_channels, use_bn=use_bn, use_clstoken=use_clstoken)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+
+        features = self.pretrained.get_intermediate_layers(x, 4, return_class_token=True)
+
+        patch_h, patch_w = h // 14, w // 14
+
+        depth = self.depth_head(features, patch_h, patch_w)
+        depth = F.interpolate(depth, size=(h, w), mode="bilinear", align_corners=True)
+        depth = F.relu(depth)
+
+        return depth.squeeze(1)

invokeai/backend/image_util/depth_anything/utilities/util.py

@@ -0,0 +1,227 @@
+import math
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
+    """Rezise the sample to ensure the given size. Keeps aspect ratio.
+
+    Args:
+        sample (dict): sample
+        size (tuple): image size
+
+    Returns:
+        tuple: new size
+    """
+    shape = list(sample["disparity"].shape)
+
+    if shape[0] >= size[0] and shape[1] >= size[1]:
+        return sample
+
+    scale = [0, 0]
+    scale[0] = size[0] / shape[0]
+    scale[1] = size[1] / shape[1]
+
+    scale = max(scale)
+
+    shape[0] = math.ceil(scale * shape[0])
+    shape[1] = math.ceil(scale * shape[1])
+
+    # resize
+    sample["image"] = cv2.resize(sample["image"], tuple(shape[::-1]), interpolation=image_interpolation_method)
+
+    sample["disparity"] = cv2.resize(sample["disparity"], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST)
+    sample["mask"] = cv2.resize(
+        sample["mask"].astype(np.float32),
+        tuple(shape[::-1]),
+        interpolation=cv2.INTER_NEAREST,
+    )
+    sample["mask"] = sample["mask"].astype(bool)
+
+    return tuple(shape)
+
+
+class Resize(object):
+    """Resize sample to given size (width, height)."""
+
+    def __init__(
+        self,
+        width,
+        height,
+        resize_target=True,
+        keep_aspect_ratio=False,
+        ensure_multiple_of=1,
+        resize_method="lower_bound",
+        image_interpolation_method=cv2.INTER_AREA,
+    ):
+        """Init.
+
+        Args:
+            width (int): desired output width
+            height (int): desired output height
+            resize_target (bool, optional):
+                True: Resize the full sample (image, mask, target).
+                False: Resize image only.
+                Defaults to True.
+            keep_aspect_ratio (bool, optional):
+                True: Keep the aspect ratio of the input sample.
+                Output sample might not have the given width and height, and
+                resize behaviour depends on the parameter 'resize_method'.
+                Defaults to False.
+            ensure_multiple_of (int, optional):
+                Output width and height is constrained to be multiple of this parameter.
+                Defaults to 1.
+            resize_method (str, optional):
+                "lower_bound": Output will be at least as large as the given size.
+                "upper_bound": Output will be at max as large as the given size. (Output size might be smaller
+                    than given size.)
+                "minimal": Scale as least as possible.  (Output size might be smaller than given size.)
+                Defaults to "lower_bound".
+        """
+        self.__width = width
+        self.__height = height
+
+        self.__resize_target = resize_target
+        self.__keep_aspect_ratio = keep_aspect_ratio
+        self.__multiple_of = ensure_multiple_of
+        self.__resize_method = resize_method
+        self.__image_interpolation_method = image_interpolation_method
+
+    def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
+        y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if max_val is not None and y > max_val:
+            y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        if y < min_val:
+            y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
+
+        return y
+
+    def get_size(self, width, height):
+        # determine new height and width
+        scale_height = self.__height / height
+        scale_width = self.__width / width
+
+        if self.__keep_aspect_ratio:
+            if self.__resize_method == "lower_bound":
+                # scale such that output size is lower bound
+                if scale_width > scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "upper_bound":
+                # scale such that output size is upper bound
+                if scale_width < scale_height:
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            elif self.__resize_method == "minimal":
+                # scale as least as possbile
+                if abs(1 - scale_width) < abs(1 - scale_height):
+                    # fit width
+                    scale_height = scale_width
+                else:
+                    # fit height
+                    scale_width = scale_height
+            else:
+                raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        if self.__resize_method == "lower_bound":
+            new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
+            new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
+        elif self.__resize_method == "upper_bound":
+            new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
+            new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
+        elif self.__resize_method == "minimal":
+            new_height = self.constrain_to_multiple_of(scale_height * height)
+            new_width = self.constrain_to_multiple_of(scale_width * width)
+        else:
+            raise ValueError(f"resize_method {self.__resize_method} not implemented")
+
+        return (new_width, new_height)
+
+    def __call__(self, sample):
+        width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
+
+        # resize sample
+        sample["image"] = cv2.resize(
+            sample["image"],
+            (width, height),
+            interpolation=self.__image_interpolation_method,
+        )
+
+        if self.__resize_target:
+            if "disparity" in sample:
+                sample["disparity"] = cv2.resize(
+                    sample["disparity"],
+                    (width, height),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+
+            if "depth" in sample:
+                sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
+
+            if "semseg_mask" in sample:
+                # sample["semseg_mask"] = cv2.resize(
+                #     sample["semseg_mask"], (width, height), interpolation=cv2.INTER_NEAREST
+                # )
+                sample["semseg_mask"] = F.interpolate(
+                    torch.from_numpy(sample["semseg_mask"]).float()[None, None, ...], (height, width), mode="nearest"
+                ).numpy()[0, 0]
+
+            if "mask" in sample:
+                sample["mask"] = cv2.resize(
+                    sample["mask"].astype(np.float32),
+                    (width, height),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+                # sample["mask"] = sample["mask"].astype(bool)
+
+        # print(sample['image'].shape, sample['depth'].shape)
+        return sample
+
+
+class NormalizeImage(object):
+    """Normlize image by given mean and std."""
+
+    def __init__(self, mean, std):
+        self.__mean = mean
+        self.__std = std
+
+    def __call__(self, sample):
+        sample["image"] = (sample["image"] - self.__mean) / self.__std
+
+        return sample
+
+
+class PrepareForNet(object):
+    """Prepare sample for usage as network input."""
+
+    def __init__(self):
+        pass
+
+    def __call__(self, sample):
+        image = np.transpose(sample["image"], (2, 0, 1))
+        sample["image"] = np.ascontiguousarray(image).astype(np.float32)
+
+        if "mask" in sample:
+            sample["mask"] = sample["mask"].astype(np.float32)
+            sample["mask"] = np.ascontiguousarray(sample["mask"])
+
+        if "depth" in sample:
+            depth = sample["depth"].astype(np.float32)
+            sample["depth"] = np.ascontiguousarray(depth)
+
+        if "semseg_mask" in sample:
+            sample["semseg_mask"] = sample["semseg_mask"].astype(np.float32)
+            sample["semseg_mask"] = np.ascontiguousarray(sample["semseg_mask"])
+
+        return sample

invokeai/backend/image_util/grounding_dino/detection_result.py

@@ -1,22 +0,0 @@
-from pydantic import BaseModel, ConfigDict
-
-
-class BoundingBox(BaseModel):
-    """Bounding box helper class."""
-
-    xmin: int
-    ymin: int
-    xmax: int
-    ymax: int
-
-
-class DetectionResult(BaseModel):
-    """Detection result from Grounding DINO."""
-
-    score: float
-    label: str
-    box: BoundingBox
-    model_config = ConfigDict(
-        # Allow arbitrary types for mask, since it will be a numpy array.
-        arbitrary_types_allowed=True
-    )

invokeai/backend/image_util/grounding_dino/grounding_dino_pipeline.py

@@ -1,37 +0,0 @@
-from typing import Optional
-
-import torch
-from PIL import Image
-from transformers.pipelines import ZeroShotObjectDetectionPipeline
-
-from invokeai.backend.image_util.grounding_dino.detection_result import DetectionResult
-from invokeai.backend.raw_model import RawModel
-
-
-class GroundingDinoPipeline(RawModel):
-    """A wrapper class for a ZeroShotObjectDetectionPipeline that makes it compatible with the model manager's memory
-    management system.
-    """
-
-    def __init__(self, pipeline: ZeroShotObjectDetectionPipeline):
-        self._pipeline = pipeline
-
-    def detect(self, image: Image.Image, candidate_labels: list[str], threshold: float = 0.1) -> list[DetectionResult]:
-        results = self._pipeline(image=image, candidate_labels=candidate_labels, threshold=threshold)
-        assert results is not None
-        results = [DetectionResult.model_validate(result) for result in results]
-        return results
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
-        # HACK(ryand): The GroundingDinoPipeline does not work on MPS devices. We only allow it to be moved to CPU or
-        # CUDA.
-        if device is not None and device.type not in {"cpu", "cuda"}:
-            device = None
-        self._pipeline.model.to(device=device, dtype=dtype)
-        self._pipeline.device = self._pipeline.model.device
-
-    def calc_size(self) -> int:
-        # HACK(ryand): Fix the circular import issue.
-        from invokeai.backend.model_manager.load.model_util import calc_module_size
-
-        return calc_module_size(self._pipeline.model)

invokeai/backend/image_util/segment_anything/mask_refinement.py

@@ -1,50 +0,0 @@
-# This file contains utilities for Grounded-SAM mask refinement based on:
-# https://github.com/NielsRogge/Transformers-Tutorials/blob/a39f33ac1557b02ebfb191ea7753e332b5ca933f/Grounding%20DINO/GroundingDINO_with_Segment_Anything.ipynb
-
-
-import cv2
-import numpy as np
-import numpy.typing as npt
-
-
-def mask_to_polygon(mask: npt.NDArray[np.uint8]) -> list[tuple[int, int]]:
-    """Convert a binary mask to a polygon.
-
-    Returns:
-        list[list[int]]: List of (x, y) coordinates representing the vertices of the polygon.
-    """
-    # Find contours in the binary mask.
-    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-    # Find the contour with the largest area.
-    largest_contour = max(contours, key=cv2.contourArea)
-
-    # Extract the vertices of the contour.
-    polygon = largest_contour.reshape(-1, 2).tolist()
-
-    return polygon
-
-
-def polygon_to_mask(
-    polygon: list[tuple[int, int]], image_shape: tuple[int, int], fill_value: int = 1
-) -> npt.NDArray[np.uint8]:
-    """Convert a polygon to a segmentation mask.
-
-    Args:
-        polygon (list): List of (x, y) coordinates representing the vertices of the polygon.
-        image_shape (tuple): Shape of the image (height, width) for the mask.
-        fill_value (int): Value to fill the polygon with.
-
-    Returns:
-        np.ndarray: Segmentation mask with the polygon filled (with value 255).
-    """
-    # Create an empty mask.
-    mask = np.zeros(image_shape, dtype=np.uint8)
-
-    # Convert polygon to an array of points.
-    pts = np.array(polygon, dtype=np.int32)
-
-    # Fill the polygon with white color (255).
-    cv2.fillPoly(mask, [pts], color=(fill_value,))
-
-    return mask

invokeai/backend/image_util/segment_anything/segment_anything_pipeline.py

@@ -1,53 +0,0 @@
-from typing import Optional
-
-import torch
-from PIL import Image
-from transformers.models.sam import SamModel
-from transformers.models.sam.processing_sam import SamProcessor
-
-from invokeai.backend.raw_model import RawModel
-
-
-class SegmentAnythingPipeline(RawModel):
-    """A wrapper class for the transformers SAM model and processor that makes it compatible with the model manager."""
-
-    def __init__(self, sam_model: SamModel, sam_processor: SamProcessor):
-        self._sam_model = sam_model
-        self._sam_processor = sam_processor
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
-        # HACK(ryand): The SAM pipeline does not work on MPS devices. We only allow it to be moved to CPU or CUDA.
-        if device is not None and device.type not in {"cpu", "cuda"}:
-            device = None
-        self._sam_model.to(device=device, dtype=dtype)
-
-    def calc_size(self) -> int:
-        # HACK(ryand): Fix the circular import issue.
-        from invokeai.backend.model_manager.load.model_util import calc_module_size
-
-        return calc_module_size(self._sam_model)
-
-    def segment(self, image: Image.Image, bounding_boxes: list[list[int]]) -> torch.Tensor:
-        """Run the SAM model.
-
-        Args:
-            image (Image.Image): The image to segment.
-            bounding_boxes (list[list[int]]): The bounding box prompts. Each bounding box is in the format
-                [xmin, ymin, xmax, ymax].
-
-        Returns:
-            torch.Tensor: The segmentation masks. dtype: torch.bool. shape: [num_masks, channels, height, width].
-        """
-        # Add batch dimension of 1 to the bounding boxes.
-        boxes = [bounding_boxes]
-        inputs = self._sam_processor(images=image, input_boxes=boxes, return_tensors="pt").to(self._sam_model.device)
-        outputs = self._sam_model(**inputs)
-        masks = self._sam_processor.post_process_masks(
-            masks=outputs.pred_masks,
-            original_sizes=inputs.original_sizes,
-            reshaped_input_sizes=inputs.reshaped_input_sizes,
-        )
-
-        # There should be only one batch.
-        assert len(masks) == 1
-        return masks[0]

invokeai/backend/ip_adapter/ip_adapter.py

@@ -124,14 +124,16 @@ class IPAdapter(RawModel):
             self.device, dtype=self.dtype
         )
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+    def to(
+        self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, non_blocking: bool = False
+    ):
         if device is not None:
             self.device = device
         if dtype is not None:
             self.dtype = dtype
 
-        self._image_proj_model.to(device=self.device, dtype=self.dtype)
-        self.attn_weights.to(device=self.device, dtype=self.dtype)
+        self._image_proj_model.to(device=self.device, dtype=self.dtype, non_blocking=non_blocking)
+        self.attn_weights.to(device=self.device, dtype=self.dtype, non_blocking=non_blocking)
 
     def calc_size(self) -> int:
         # HACK(ryand): Fix this issue with circular imports.

invokeai/backend/lora.py

@@ -3,15 +3,15 @@
 
 import bisect
 from pathlib import Path
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from safetensors.torch import load_file
 from typing_extensions import Self
 
-import invokeai.backend.util.logging as logger
 from invokeai.backend.model_manager import BaseModelType
 from invokeai.backend.raw_model import RawModel
+from invokeai.backend.util.devices import TorchDevice
 
 
 class LoRALayerBase:
@@ -47,19 +47,9 @@ class LoRALayerBase:
         self.rank = None  # set in layer implementation
         self.layer_key = layer_key
 
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+    def get_weight(self, orig_weight: Optional[torch.Tensor]) -> torch.Tensor:
         raise NotImplementedError()
 
-    def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
-        return self.bias
-
-    def get_parameters(self, orig_module: torch.nn.Module) -> Dict[str, torch.Tensor]:
-        params = {"weight": self.get_weight(orig_module.weight)}
-        bias = self.get_bias(orig_module.bias)
-        if bias is not None:
-            params["bias"] = bias
-        return params
-
     def calc_size(self) -> int:
         model_size = 0
         for val in [self.bias]:
@@ -67,20 +57,14 @@ class LoRALayerBase:
                 model_size += val.nelement() * val.element_size()
         return model_size
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
         if self.bias is not None:
-            self.bias = self.bias.to(device=device, dtype=dtype)
-
-    def check_keys(self, values: Dict[str, torch.Tensor], known_keys: Set[str]):
-        """Log a warning if values contains unhandled keys."""
-        # {"alpha", "bias_indices", "bias_values", "bias_size"} are hard-coded, because they are handled by
-        # `LoRALayerBase`. Sub-classes should provide the known_keys that they handled.
-        all_known_keys = known_keys | {"alpha", "bias_indices", "bias_values", "bias_size"}
-        unknown_keys = set(values.keys()) - all_known_keys
-        if unknown_keys:
-            logger.warning(
-                f"Unexpected keys found in LoRA/LyCORIS layer, model might work incorrectly! Keys: {unknown_keys}"
-            )
+            self.bias = self.bias.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
 
 # TODO: find and debug lora/locon with bias
@@ -98,19 +82,14 @@ class LoRALayer(LoRALayerBase):
 
         self.up = values["lora_up.weight"]
         self.down = values["lora_down.weight"]
-        self.mid = values.get("lora_mid.weight", None)
+        if "lora_mid.weight" in values:
+            self.mid: Optional[torch.Tensor] = values["lora_mid.weight"]
+        else:
+            self.mid = None
 
         self.rank = self.down.shape[0]
-        self.check_keys(
-            values,
-            {
-                "lora_up.weight",
-                "lora_down.weight",
-                "lora_mid.weight",
-            },
-        )
 
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+    def get_weight(self, orig_weight: Optional[torch.Tensor]) -> torch.Tensor:
         if self.mid is not None:
             up = self.up.reshape(self.up.shape[0], self.up.shape[1])
             down = self.down.reshape(self.down.shape[0], self.down.shape[1])
@@ -127,14 +106,19 @@ class LoRALayer(LoRALayerBase):
                 model_size += val.nelement() * val.element_size()
         return model_size
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        super().to(device=device, dtype=dtype)
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
+        super().to(device=device, dtype=dtype, non_blocking=non_blocking)
 
-        self.up = self.up.to(device=device, dtype=dtype)
-        self.down = self.down.to(device=device, dtype=dtype)
+        self.up = self.up.to(device=device, dtype=dtype, non_blocking=non_blocking)
+        self.down = self.down.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
         if self.mid is not None:
-            self.mid = self.mid.to(device=device, dtype=dtype)
+            self.mid = self.mid.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
 
 class LoHALayer(LoRALayerBase):
@@ -152,23 +136,20 @@ class LoHALayer(LoRALayerBase):
         self.w1_b = values["hada_w1_b"]
         self.w2_a = values["hada_w2_a"]
         self.w2_b = values["hada_w2_b"]
-        self.t1 = values.get("hada_t1", None)
-        self.t2 = values.get("hada_t2", None)
+
+        if "hada_t1" in values:
+            self.t1: Optional[torch.Tensor] = values["hada_t1"]
+        else:
+            self.t1 = None
+
+        if "hada_t2" in values:
+            self.t2: Optional[torch.Tensor] = values["hada_t2"]
+        else:
+            self.t2 = None
 
         self.rank = self.w1_b.shape[0]
-        self.check_keys(
-            values,
-            {
-                "hada_w1_a",
-                "hada_w1_b",
-                "hada_w2_a",
-                "hada_w2_b",
-                "hada_t1",
-                "hada_t2",
-            },
-        )
 
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+    def get_weight(self, orig_weight: Optional[torch.Tensor]) -> torch.Tensor:
         if self.t1 is None:
             weight: torch.Tensor = (self.w1_a @ self.w1_b) * (self.w2_a @ self.w2_b)
 
@@ -186,18 +167,23 @@ class LoHALayer(LoRALayerBase):
                 model_size += val.nelement() * val.element_size()
         return model_size
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
         super().to(device=device, dtype=dtype)
 
-        self.w1_a = self.w1_a.to(device=device, dtype=dtype)
-        self.w1_b = self.w1_b.to(device=device, dtype=dtype)
+        self.w1_a = self.w1_a.to(device=device, dtype=dtype, non_blocking=non_blocking)
+        self.w1_b = self.w1_b.to(device=device, dtype=dtype, non_blocking=non_blocking)
         if self.t1 is not None:
-            self.t1 = self.t1.to(device=device, dtype=dtype)
+            self.t1 = self.t1.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
-        self.w2_a = self.w2_a.to(device=device, dtype=dtype)
-        self.w2_b = self.w2_b.to(device=device, dtype=dtype)
+        self.w2_a = self.w2_a.to(device=device, dtype=dtype, non_blocking=non_blocking)
+        self.w2_b = self.w2_b.to(device=device, dtype=dtype, non_blocking=non_blocking)
         if self.t2 is not None:
-            self.t2 = self.t2.to(device=device, dtype=dtype)
+            self.t2 = self.t2.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
 
 class LoKRLayer(LoRALayerBase):
@@ -216,45 +202,37 @@ class LoKRLayer(LoRALayerBase):
     ):
         super().__init__(layer_key, values)
 
-        self.w1 = values.get("lokr_w1", None)
-        if self.w1 is None:
+        if "lokr_w1" in values:
+            self.w1: Optional[torch.Tensor] = values["lokr_w1"]
+            self.w1_a = None
+            self.w1_b = None
+        else:
+            self.w1 = None
             self.w1_a = values["lokr_w1_a"]
             self.w1_b = values["lokr_w1_b"]
-        else:
-            self.w1_b = None
-            self.w1_a = None
 
-        self.w2 = values.get("lokr_w2", None)
-        if self.w2 is None:
-            self.w2_a = values["lokr_w2_a"]
-            self.w2_b = values["lokr_w2_b"]
-        else:
+        if "lokr_w2" in values:
+            self.w2: Optional[torch.Tensor] = values["lokr_w2"]
             self.w2_a = None
             self.w2_b = None
+        else:
+            self.w2 = None
+            self.w2_a = values["lokr_w2_a"]
+            self.w2_b = values["lokr_w2_b"]
 
-        self.t2 = values.get("lokr_t2", None)
+        if "lokr_t2" in values:
+            self.t2: Optional[torch.Tensor] = values["lokr_t2"]
+        else:
+            self.t2 = None
 
-        if self.w1_b is not None:
-            self.rank = self.w1_b.shape[0]
-        elif self.w2_b is not None:
-            self.rank = self.w2_b.shape[0]
+        if "lokr_w1_b" in values:
+            self.rank = values["lokr_w1_b"].shape[0]
+        elif "lokr_w2_b" in values:
+            self.rank = values["lokr_w2_b"].shape[0]
         else:
             self.rank = None  # unscaled
 
-        self.check_keys(
-            values,
-            {
-                "lokr_w1",
-                "lokr_w1_a",
-                "lokr_w1_b",
-                "lokr_w2",
-                "lokr_w2_a",
-                "lokr_w2_b",
-                "lokr_t2",
-            },
-        )
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+    def get_weight(self, orig_weight: Optional[torch.Tensor]) -> torch.Tensor:
         w1: Optional[torch.Tensor] = self.w1
         if w1 is None:
             assert self.w1_a is not None
@@ -286,7 +264,12 @@ class LoKRLayer(LoRALayerBase):
                 model_size += val.nelement() * val.element_size()
         return model_size
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
         super().to(device=device, dtype=dtype)
 
         if self.w1 is not None:
@@ -294,25 +277,23 @@ class LoKRLayer(LoRALayerBase):
         else:
             assert self.w1_a is not None
             assert self.w1_b is not None
-            self.w1_a = self.w1_a.to(device=device, dtype=dtype)
-            self.w1_b = self.w1_b.to(device=device, dtype=dtype)
+            self.w1_a = self.w1_a.to(device=device, dtype=dtype, non_blocking=non_blocking)
+            self.w1_b = self.w1_b.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
         if self.w2 is not None:
-            self.w2 = self.w2.to(device=device, dtype=dtype)
+            self.w2 = self.w2.to(device=device, dtype=dtype, non_blocking=non_blocking)
         else:
             assert self.w2_a is not None
             assert self.w2_b is not None
-            self.w2_a = self.w2_a.to(device=device, dtype=dtype)
-            self.w2_b = self.w2_b.to(device=device, dtype=dtype)
+            self.w2_a = self.w2_a.to(device=device, dtype=dtype, non_blocking=non_blocking)
+            self.w2_b = self.w2_b.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
         if self.t2 is not None:
-            self.t2 = self.t2.to(device=device, dtype=dtype)
+            self.t2 = self.t2.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
 
 class FullLayer(LoRALayerBase):
-    # bias handled in LoRALayerBase(calc_size, to)
     # weight: torch.Tensor
-    # bias: Optional[torch.Tensor]
 
     def __init__(
         self,
@@ -322,12 +303,15 @@ class FullLayer(LoRALayerBase):
         super().__init__(layer_key, values)
 
         self.weight = values["diff"]
-        self.bias = values.get("diff_b", None)
+
+        if len(values.keys()) > 1:
+            _keys = list(values.keys())
+            _keys.remove("diff")
+            raise NotImplementedError(f"Unexpected keys in lora diff layer: {_keys}")
 
         self.rank = None  # unscaled
-        self.check_keys(values, {"diff", "diff_b"})
 
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+    def get_weight(self, orig_weight: Optional[torch.Tensor]) -> torch.Tensor:
         return self.weight
 
     def calc_size(self) -> int:
@@ -335,10 +319,15 @@ class FullLayer(LoRALayerBase):
         model_size += self.weight.nelement() * self.weight.element_size()
         return model_size
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
         super().to(device=device, dtype=dtype)
 
-        self.weight = self.weight.to(device=device, dtype=dtype)
+        self.weight = self.weight.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
 
 class IA3Layer(LoRALayerBase):
@@ -356,9 +345,8 @@ class IA3Layer(LoRALayerBase):
         self.on_input = values["on_input"]
 
         self.rank = None  # unscaled
-        self.check_keys(values, {"weight", "on_input"})
 
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
+    def get_weight(self, orig_weight: Optional[torch.Tensor]) -> torch.Tensor:
         weight = self.weight
         if not self.on_input:
             weight = weight.reshape(-1, 1)
@@ -371,46 +359,19 @@ class IA3Layer(LoRALayerBase):
         model_size += self.on_input.nelement() * self.on_input.element_size()
         return model_size
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
-        super().to(device=device, dtype=dtype)
-
-        self.weight = self.weight.to(device=device, dtype=dtype)
-        self.on_input = self.on_input.to(device=device, dtype=dtype)
-
-
-class NormLayer(LoRALayerBase):
-    # bias handled in LoRALayerBase(calc_size, to)
-    # weight: torch.Tensor
-    # bias: Optional[torch.Tensor]
-
-    def __init__(
+    def to(
         self,
-        layer_key: str,
-        values: Dict[str, torch.Tensor],
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
     ):
-        super().__init__(layer_key, values)
-
-        self.weight = values["w_norm"]
-        self.bias = values.get("b_norm", None)
-
-        self.rank = None  # unscaled
-        self.check_keys(values, {"w_norm", "b_norm"})
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        return self.weight
-
-    def calc_size(self) -> int:
-        model_size = super().calc_size()
-        model_size += self.weight.nelement() * self.weight.element_size()
-        return model_size
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
         super().to(device=device, dtype=dtype)
 
-        self.weight = self.weight.to(device=device, dtype=dtype)
+        self.weight = self.weight.to(device=device, dtype=dtype, non_blocking=non_blocking)
+        self.on_input = self.on_input.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
 
-AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer, NormLayer]
+AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer]
 
 
 class LoRAModelRaw(RawModel):  # (torch.nn.Module):
@@ -429,10 +390,15 @@ class LoRAModelRaw(RawModel):  # (torch.nn.Module):
     def name(self) -> str:
         return self._name
 
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
         # TODO: try revert if exception?
         for _key, layer in self.layers.items():
-            layer.to(device=device, dtype=dtype)
+            layer.to(device=device, dtype=dtype, non_blocking=non_blocking)
 
     def calc_size(self) -> int:
         model_size = 0
@@ -528,19 +494,16 @@ class LoRAModelRaw(RawModel):  # (torch.nn.Module):
             state_dict = cls._convert_sdxl_keys_to_diffusers_format(state_dict)
 
         for layer_key, values in state_dict.items():
-            # Detect layers according to LyCORIS detection logic(`weight_list_det`)
-            # https://github.com/KohakuBlueleaf/LyCORIS/tree/8ad8000efb79e2b879054da8c9356e6143591bad/lycoris/modules
-
             # lora and locon
-            if "lora_up.weight" in values:
+            if "lora_down.weight" in values:
                 layer: AnyLoRALayer = LoRALayer(layer_key, values)
 
             # loha
-            elif "hada_w1_a" in values:
+            elif "hada_w1_b" in values:
                 layer = LoHALayer(layer_key, values)
 
             # lokr
-            elif "lokr_w1" in values or "lokr_w1_a" in values:
+            elif "lokr_w1_b" in values or "lokr_w1" in values:
                 layer = LoKRLayer(layer_key, values)
 
             # diff
@@ -548,13 +511,9 @@ class LoRAModelRaw(RawModel):  # (torch.nn.Module):
                 layer = FullLayer(layer_key, values)
 
             # ia3
-            elif "on_input" in values:
+            elif "weight" in values and "on_input" in values:
                 layer = IA3Layer(layer_key, values)
 
-            # norms
-            elif "w_norm" in values:
-                layer = NormLayer(layer_key, values)
-
             else:
                 print(f">> Encountered unknown lora layer module in {model.name}: {layer_key} - {list(values.keys())}")
                 raise Exception("Unknown lora format!")
@@ -562,7 +521,7 @@ class LoRAModelRaw(RawModel):  # (torch.nn.Module):
             # lower memory consumption by removing already parsed layer values
             state_dict[layer_key].clear()
 
-            layer.to(device=device, dtype=dtype)
+            layer.to(device=device, dtype=dtype, non_blocking=TorchDevice.get_non_blocking(device))
             model.layers[layer_key] = layer
 
         return model

invokeai/backend/model_manager/config.py

@@ -67,7 +67,6 @@ class ModelType(str, Enum):
     IPAdapter = "ip_adapter"
     CLIPVision = "clip_vision"
     T2IAdapter = "t2i_adapter"
-    SpandrelImageToImage = "spandrel_image_to_image"
 
 
 class SubModelType(str, Enum):
@@ -354,7 +353,7 @@ class CLIPVisionDiffusersConfig(DiffusersConfigBase):
     """Model config for CLIPVision."""
 
     type: Literal[ModelType.CLIPVision] = ModelType.CLIPVision
-    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+    format: Literal[ModelFormat.Diffusers]
 
     @staticmethod
     def get_tag() -> Tag:
@@ -365,24 +364,13 @@ class T2IAdapterConfig(DiffusersConfigBase, ControlAdapterConfigBase):
     """Model config for T2I."""
 
     type: Literal[ModelType.T2IAdapter] = ModelType.T2IAdapter
-    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+    format: Literal[ModelFormat.Diffusers]
 
     @staticmethod
     def get_tag() -> Tag:
         return Tag(f"{ModelType.T2IAdapter.value}.{ModelFormat.Diffusers.value}")
 
 
-class SpandrelImageToImageConfig(ModelConfigBase):
-    """Model config for Spandrel Image to Image models."""
-
-    type: Literal[ModelType.SpandrelImageToImage] = ModelType.SpandrelImageToImage
-    format: Literal[ModelFormat.Checkpoint] = ModelFormat.Checkpoint
-
-    @staticmethod
-    def get_tag() -> Tag:
-        return Tag(f"{ModelType.SpandrelImageToImage.value}.{ModelFormat.Checkpoint.value}")
-
-
 def get_model_discriminator_value(v: Any) -> str:
     """
     Computes the discriminator value for a model config.
@@ -419,7 +407,6 @@ AnyModelConfig = Annotated[
         Annotated[IPAdapterInvokeAIConfig, IPAdapterInvokeAIConfig.get_tag()],
         Annotated[IPAdapterCheckpointConfig, IPAdapterCheckpointConfig.get_tag()],
         Annotated[T2IAdapterConfig, T2IAdapterConfig.get_tag()],
-        Annotated[SpandrelImageToImageConfig, SpandrelImageToImageConfig.get_tag()],
         Annotated[CLIPVisionDiffusersConfig, CLIPVisionDiffusersConfig.get_tag()],
     ],
     Discriminator(get_model_discriminator_value),

invokeai/backend/model_manager/load/model_cache/model_cache_default.py

@@ -167,8 +167,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
         size = calc_model_size_by_data(self.logger, model)
         self.make_room(size)
 
-        running_on_cpu = self.execution_device == torch.device("cpu")
-        state_dict = model.state_dict() if isinstance(model, torch.nn.Module) and not running_on_cpu else None
+        state_dict = model.state_dict() if isinstance(model, torch.nn.Module) else None
         cache_record = CacheRecord(key=key, model=model, device=self.storage_device, state_dict=state_dict, size=size)
         self._cached_models[key] = cache_record
         self._cache_stack.append(key)
@@ -290,9 +289,11 @@ class ModelCache(ModelCacheBase[AnyModel]):
                 else:
                     new_dict: Dict[str, torch.Tensor] = {}
                     for k, v in cache_entry.state_dict.items():
-                        new_dict[k] = v.to(target_device, copy=True)
+                        new_dict[k] = v.to(
+                            target_device, copy=True, non_blocking=TorchDevice.get_non_blocking(target_device)
+                        )
                     cache_entry.model.load_state_dict(new_dict, assign=True)
-            cache_entry.model.to(target_device)
+            cache_entry.model.to(target_device, non_blocking=TorchDevice.get_non_blocking(target_device))
             cache_entry.device = target_device
         except Exception as e:  # blow away cache entry
             self._delete_cache_entry(cache_entry)

invokeai/backend/model_manager/load/model_loaders/spandrel_image_to_image.py

@@ -1,45 +0,0 @@
-from pathlib import Path
-from typing import Optional
-
-import torch
-
-from invokeai.backend.model_manager.config import (
-    AnyModel,
-    AnyModelConfig,
-    BaseModelType,
-    ModelFormat,
-    ModelType,
-    SubModelType,
-)
-from invokeai.backend.model_manager.load.load_default import ModelLoader
-from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
-from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
-
-
-@ModelLoaderRegistry.register(
-    base=BaseModelType.Any, type=ModelType.SpandrelImageToImage, format=ModelFormat.Checkpoint
-)
-class SpandrelImageToImageModelLoader(ModelLoader):
-    """Class for loading Spandrel Image-to-Image models (i.e. models wrapped by spandrel.ImageModelDescriptor)."""
-
-    def _load_model(
-        self,
-        config: AnyModelConfig,
-        submodel_type: Optional[SubModelType] = None,
-    ) -> AnyModel:
-        if submodel_type is not None:
-            raise ValueError("Unexpected submodel requested for Spandrel model.")
-
-        model_path = Path(config.path)
-        model = SpandrelImageToImageModel.load_from_file(model_path)
-
-        torch_dtype = self._torch_dtype
-        if not model.supports_dtype(torch_dtype):
-            self._logger.warning(
-                f"The configured dtype ('{self._torch_dtype}') is not supported by the {model.get_model_type_name()} "
-                "model. Falling back to 'float32'."
-            )
-            torch_dtype = torch.float32
-        model.to(dtype=torch_dtype)
-
-        return model

invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py

@@ -98,9 +98,6 @@ class StableDiffusionDiffusersModel(GenericDiffusersLoader):
                 ModelVariantType.Normal: StableDiffusionXLPipeline,
                 ModelVariantType.Inpaint: StableDiffusionXLInpaintPipeline,
             },
-            BaseModelType.StableDiffusionXLRefiner: {
-                ModelVariantType.Normal: StableDiffusionXLPipeline,
-            },
         }
         assert isinstance(config, MainCheckpointConfig)
         try:

invokeai/backend/model_manager/load/model_util.py

@@ -11,14 +11,10 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from transformers import CLIPTokenizer
 
-from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import DepthAnythingPipeline
-from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
-from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
 from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
 from invokeai.backend.lora import LoRAModelRaw
 from invokeai.backend.model_manager.config import AnyModel
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
-from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
 from invokeai.backend.textual_inversion import TextualInversionModelRaw
 
 
@@ -37,18 +33,7 @@ def calc_model_size_by_data(logger: logging.Logger, model: AnyModel) -> int:
     elif isinstance(model, CLIPTokenizer):
         # TODO(ryand): Accurately calculate the tokenizer's size. It's small enough that it shouldn't matter for now.
         return 0
-    elif isinstance(
-        model,
-        (
-            TextualInversionModelRaw,
-            IPAdapter,
-            LoRAModelRaw,
-            SpandrelImageToImageModel,
-            GroundingDinoPipeline,
-            SegmentAnythingPipeline,
-            DepthAnythingPipeline,
-        ),
-    ):
+    elif isinstance(model, (TextualInversionModelRaw, IPAdapter, LoRAModelRaw)):
         return model.calc_size()
     else:
         # TODO(ryand): Promote this from a log to an exception once we are confident that we are handling all of the

invokeai/backend/model_manager/probe.py

@@ -4,7 +4,6 @@ from pathlib import Path
 from typing import Any, Dict, Literal, Optional, Union
 
 import safetensors.torch
-import spandrel
 import torch
 from picklescan.scanner import scan_file_path
 
@@ -26,7 +25,6 @@ from invokeai.backend.model_manager.config import (
     SchedulerPredictionType,
 )
 from invokeai.backend.model_manager.util.model_util import lora_token_vector_length, read_checkpoint_meta
-from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
 from invokeai.backend.util.silence_warnings import SilenceWarnings
 
 CkptType = Dict[str | int, Any]
@@ -222,46 +220,24 @@ class ModelProbe(object):
         ckpt = ckpt.get("state_dict", ckpt)
 
         for key in [str(k) for k in ckpt.keys()]:
-            if key.startswith(("cond_stage_model.", "first_stage_model.", "model.diffusion_model.")):
+            if any(key.startswith(v) for v in {"cond_stage_model.", "first_stage_model.", "model.diffusion_model."}):
                 return ModelType.Main
-            elif key.startswith(("encoder.conv_in", "decoder.conv_in")):
+            elif any(key.startswith(v) for v in {"encoder.conv_in", "decoder.conv_in"}):
                 return ModelType.VAE
-            elif key.startswith(("lora_te_", "lora_unet_")):
+            elif any(key.startswith(v) for v in {"lora_te_", "lora_unet_"}):
                 return ModelType.LoRA
-            elif key.endswith(("to_k_lora.up.weight", "to_q_lora.down.weight")):
+            elif any(key.endswith(v) for v in {"to_k_lora.up.weight", "to_q_lora.down.weight"}):
                 return ModelType.LoRA
-            elif key.startswith(("controlnet", "control_model", "input_blocks")):
+            elif any(key.startswith(v) for v in {"controlnet", "control_model", "input_blocks"}):
                 return ModelType.ControlNet
-            elif key.startswith(("image_proj.", "ip_adapter.")):
+            elif any(key.startswith(v) for v in {"image_proj.", "ip_adapter."}):
                 return ModelType.IPAdapter
             elif key in {"emb_params", "string_to_param"}:
                 return ModelType.TextualInversion
-
-        # diffusers-ti
-        if len(ckpt) < 10 and all(isinstance(v, torch.Tensor) for v in ckpt.values()):
-            return ModelType.TextualInversion
-
-        # Check if the model can be loaded as a SpandrelImageToImageModel.
-        # This check is intentionally performed last, as it can be expensive (it requires loading the model from disk).
-        try:
-            # It would be nice to avoid having to load the Spandrel model from disk here. A couple of options were
-            # explored to avoid this:
-            # 1. Call `SpandrelImageToImageModel.load_from_state_dict(ckpt)`, where `ckpt` is a state_dict on the meta
-            #    device. Unfortunately, some Spandrel models perform operations during initialization that are not
-            #    supported on meta tensors.
-            # 2. Spandrel has internal logic to determine a model's type from its state_dict before loading the model.
-            #    This logic is not exposed in spandrel's public API. We could copy the logic here, but then we have to
-            #    maintain it, and the risk of false positive detections is higher.
-            SpandrelImageToImageModel.load_from_file(model_path)
-            return ModelType.SpandrelImageToImage
-        except spandrel.UnsupportedModelError:
-            pass
-        except RuntimeError as e:
-            if "No such file or directory" in str(e):
-                # This error is expected if the model_path does not exist (which is the case in some unit tests).
-                pass
-            else:
-                raise e
+        else:
+            # diffusers-ti
+            if len(ckpt) < 10 and all(isinstance(v, torch.Tensor) for v in ckpt.values()):
+                return ModelType.TextualInversion
 
         raise InvalidModelConfigException(f"Unable to determine model type for {model_path}")
 
@@ -593,11 +569,6 @@ class T2IAdapterCheckpointProbe(CheckpointProbeBase):
         raise NotImplementedError()
 
 
-class SpandrelImageToImageCheckpointProbe(CheckpointProbeBase):
-    def get_base_type(self) -> BaseModelType:
-        return BaseModelType.Any
-
-
 ########################################################
 # classes for probing folders
 #######################################################
@@ -805,11 +776,6 @@ class CLIPVisionFolderProbe(FolderProbeBase):
         return BaseModelType.Any
 
 
-class SpandrelImageToImageFolderProbe(FolderProbeBase):
-    def get_base_type(self) -> BaseModelType:
-        raise NotImplementedError()
-
-
 class T2IAdapterFolderProbe(FolderProbeBase):
     def get_base_type(self) -> BaseModelType:
         config_file = self.model_path / "config.json"
@@ -839,7 +805,6 @@ ModelProbe.register_probe("diffusers", ModelType.ControlNet, ControlNetFolderPro
 ModelProbe.register_probe("diffusers", ModelType.IPAdapter, IPAdapterFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.CLIPVision, CLIPVisionFolderProbe)
 ModelProbe.register_probe("diffusers", ModelType.T2IAdapter, T2IAdapterFolderProbe)
-ModelProbe.register_probe("diffusers", ModelType.SpandrelImageToImage, SpandrelImageToImageFolderProbe)
 
 ModelProbe.register_probe("checkpoint", ModelType.Main, PipelineCheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.VAE, VaeCheckpointProbe)
@@ -849,6 +814,5 @@ ModelProbe.register_probe("checkpoint", ModelType.ControlNet, ControlNetCheckpoi
 ModelProbe.register_probe("checkpoint", ModelType.IPAdapter, IPAdapterCheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.CLIPVision, CLIPVisionCheckpointProbe)
 ModelProbe.register_probe("checkpoint", ModelType.T2IAdapter, T2IAdapterCheckpointProbe)
-ModelProbe.register_probe("checkpoint", ModelType.SpandrelImageToImage, SpandrelImageToImageCheckpointProbe)
 
 ModelProbe.register_probe("onnx", ModelType.ONNX, ONNXFolderProbe)

invokeai/backend/model_manager/starter_models.py

@@ -187,171 +187,157 @@ STARTER_MODELS: list[StarterModel] = [
     # endregion
     # region ControlNet
     StarterModel(
-        name="QRCode Monster v2 (SD1.5)",
+        name="QRCode Monster",
         base=BaseModelType.StableDiffusion1,
-        source="monster-labs/control_v1p_sd15_qrcode_monster::v2",
-        description="ControlNet model that generates scannable creative QR codes",
-        type=ModelType.ControlNet,
-    ),
-    StarterModel(
-        name="QRCode Monster (SDXL)",
-        base=BaseModelType.StableDiffusionXL,
-        source="monster-labs/control_v1p_sdxl_qrcode_monster",
-        description="ControlNet model that generates scannable creative QR codes",
+        source="monster-labs/control_v1p_sd15_qrcode_monster",
+        description="Controlnet model that generates scannable creative QR codes",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="canny",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_canny",
-        description="ControlNet weights trained on sd-1.5 with canny conditioning.",
+        description="Controlnet weights trained on sd-1.5 with canny conditioning.",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="inpaint",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_inpaint",
-        description="ControlNet weights trained on sd-1.5 with canny conditioning, inpaint version",
+        description="Controlnet weights trained on sd-1.5 with canny conditioning, inpaint version",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="mlsd",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_mlsd",
-        description="ControlNet weights trained on sd-1.5 with canny conditioning, MLSD version",
+        description="Controlnet weights trained on sd-1.5 with canny conditioning, MLSD version",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="depth",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11f1p_sd15_depth",
-        description="ControlNet weights trained on sd-1.5 with depth conditioning",
+        description="Controlnet weights trained on sd-1.5 with depth conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="normal_bae",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_normalbae",
-        description="ControlNet weights trained on sd-1.5 with normalbae image conditioning",
+        description="Controlnet weights trained on sd-1.5 with normalbae image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="seg",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_seg",
-        description="ControlNet weights trained on sd-1.5 with seg image conditioning",
+        description="Controlnet weights trained on sd-1.5 with seg image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="lineart",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_lineart",
-        description="ControlNet weights trained on sd-1.5 with lineart image conditioning",
+        description="Controlnet weights trained on sd-1.5 with lineart image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="lineart_anime",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15s2_lineart_anime",
-        description="ControlNet weights trained on sd-1.5 with anime image conditioning",
+        description="Controlnet weights trained on sd-1.5 with anime image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="openpose",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_openpose",
-        description="ControlNet weights trained on sd-1.5 with openpose image conditioning",
+        description="Controlnet weights trained on sd-1.5 with openpose image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="scribble",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_scribble",
-        description="ControlNet weights trained on sd-1.5 with scribble image conditioning",
+        description="Controlnet weights trained on sd-1.5 with scribble image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="softedge",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11p_sd15_softedge",
-        description="ControlNet weights trained on sd-1.5 with soft edge conditioning",
+        description="Controlnet weights trained on sd-1.5 with soft edge conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="shuffle",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11e_sd15_shuffle",
-        description="ControlNet weights trained on sd-1.5 with shuffle image conditioning",
+        description="Controlnet weights trained on sd-1.5 with shuffle image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="tile",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11f1e_sd15_tile",
-        description="ControlNet weights trained on sd-1.5 with tiled image conditioning",
+        description="Controlnet weights trained on sd-1.5 with tiled image conditioning",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="ip2p",
         base=BaseModelType.StableDiffusion1,
         source="lllyasviel/control_v11e_sd15_ip2p",
-        description="ControlNet weights trained on sd-1.5 with ip2p conditioning.",
+        description="Controlnet weights trained on sd-1.5 with ip2p conditioning.",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="canny-sdxl",
         base=BaseModelType.StableDiffusionXL,
-        source="xinsir/controlNet-canny-sdxl-1.0",
-        description="ControlNet weights trained on sdxl-1.0 with canny conditioning, by Xinsir.",
+        source="xinsir/controlnet-canny-sdxl-1.0",
+        description="Controlnet weights trained on sdxl-1.0 with canny conditioning, by Xinsir.",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="depth-sdxl",
         base=BaseModelType.StableDiffusionXL,
-        source="diffusers/controlNet-depth-sdxl-1.0",
-        description="ControlNet weights trained on sdxl-1.0 with depth conditioning.",
+        source="diffusers/controlnet-depth-sdxl-1.0",
+        description="Controlnet weights trained on sdxl-1.0 with depth conditioning.",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="softedge-dexined-sdxl",
         base=BaseModelType.StableDiffusionXL,
-        source="SargeZT/controlNet-sd-xl-1.0-softedge-dexined",
-        description="ControlNet weights trained on sdxl-1.0 with dexined soft edge preprocessing.",
+        source="SargeZT/controlnet-sd-xl-1.0-softedge-dexined",
+        description="Controlnet weights trained on sdxl-1.0 with dexined soft edge preprocessing.",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="depth-16bit-zoe-sdxl",
         base=BaseModelType.StableDiffusionXL,
-        source="SargeZT/controlNet-sd-xl-1.0-depth-16bit-zoe",
-        description="ControlNet weights trained on sdxl-1.0 with Zoe's preprocessor (16 bits).",
+        source="SargeZT/controlnet-sd-xl-1.0-depth-16bit-zoe",
+        description="Controlnet weights trained on sdxl-1.0 with Zoe's preprocessor (16 bits).",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="depth-zoe-sdxl",
         base=BaseModelType.StableDiffusionXL,
-        source="diffusers/controlNet-zoe-depth-sdxl-1.0",
-        description="ControlNet weights trained on sdxl-1.0 with Zoe's preprocessor (32 bits).",
+        source="diffusers/controlnet-zoe-depth-sdxl-1.0",
+        description="Controlnet weights trained on sdxl-1.0 with Zoe's preprocessor (32 bits).",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="openpose-sdxl",
         base=BaseModelType.StableDiffusionXL,
-        source="xinsir/controlNet-openpose-sdxl-1.0",
-        description="ControlNet weights trained on sdxl-1.0 compatible with the DWPose processor by Xinsir.",
+        source="xinsir/controlnet-openpose-sdxl-1.0",
+        description="Controlnet weights trained on sdxl-1.0 compatible with the DWPose processor by Xinsir.",
         type=ModelType.ControlNet,
     ),
     StarterModel(
         name="scribble-sdxl",
         base=BaseModelType.StableDiffusionXL,
-        source="xinsir/controlNet-scribble-sdxl-1.0",
-        description="ControlNet weights trained on sdxl-1.0 compatible with various lineart processors and black/white sketches by Xinsir.",
-        type=ModelType.ControlNet,
-    ),
-    StarterModel(
-        name="tile-sdxl",
-        base=BaseModelType.StableDiffusionXL,
-        source="xinsir/controlNet-tile-sdxl-1.0",
-        description="ControlNet weights trained on sdxl-1.0 with tiled image conditioning",
+        source="xinsir/controlnet-scribble-sdxl-1.0",
+        description="Controlnet weights trained on sdxl-1.0 compatible with various lineart processors and black/white sketches by Xinsir.",
         type=ModelType.ControlNet,
     ),
     # endregion
@@ -413,43 +399,6 @@ STARTER_MODELS: list[StarterModel] = [
         type=ModelType.T2IAdapter,
     ),
     # endregion
-    # region SpandrelImageToImage
-    StarterModel(
-        name="RealESRGAN_x4plus_anime_6B",
-        base=BaseModelType.Any,
-        source="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
-        description="A Real-ESRGAN 4x upscaling model (optimized for anime images).",
-        type=ModelType.SpandrelImageToImage,
-    ),
-    StarterModel(
-        name="RealESRGAN_x4plus",
-        base=BaseModelType.Any,
-        source="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
-        description="A Real-ESRGAN 4x upscaling model (general-purpose).",
-        type=ModelType.SpandrelImageToImage,
-    ),
-    StarterModel(
-        name="ESRGAN_SRx4_DF2KOST_official",
-        base=BaseModelType.Any,
-        source="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth",
-        description="The official ESRGAN 4x upscaling model.",
-        type=ModelType.SpandrelImageToImage,
-    ),
-    StarterModel(
-        name="RealESRGAN_x2plus",
-        base=BaseModelType.Any,
-        source="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
-        description="A Real-ESRGAN 2x upscaling model (general-purpose).",
-        type=ModelType.SpandrelImageToImage,
-    ),
-    StarterModel(
-        name="SwinIR - realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN",
-        base=BaseModelType.Any,
-        source="https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN-with-dict-keys-params-and-params_ema.pth",
-        description="A SwinIR 4x upscaling model.",
-        type=ModelType.SpandrelImageToImage,
-    ),
-    # endregion
 ]
 
 assert len(STARTER_MODELS) == len({m.source for m in STARTER_MODELS}), "Duplicate starter models"

invokeai/backend/model_patcher.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import pickle
 from contextlib import contextmanager
-from typing import Any, Dict, Generator, Iterator, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, Generator, Iterator, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -17,9 +17,8 @@ from invokeai.backend.lora import LoRAModelRaw
 from invokeai.backend.model_manager import AnyModel
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
-from invokeai.backend.stable_diffusion.extensions.lora import LoRAExt
 from invokeai.backend.textual_inversion import TextualInversionManager, TextualInversionModelRaw
-from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+from invokeai.backend.util.devices import TorchDevice
 
 """
 loras = [
@@ -33,27 +32,8 @@ with LoRAHelper.apply_lora_unet(unet, loras):
 """
 
 
+# TODO: rename smth like ModelPatcher and add TI method?
 class ModelPatcher:
-    @staticmethod
-    @contextmanager
-    def patch_unet_attention_processor(unet: UNet2DConditionModel, processor_cls: Type[Any]):
-        """A context manager that patches `unet` with the provided attention processor.
-
-        Args:
-            unet (UNet2DConditionModel): The UNet model to patch.
-            processor (Type[Any]): Class which will be initialized for each key and passed to set_attn_processor(...).
-        """
-        unet_orig_processors = unet.attn_processors
-
-        # create separate instance for each attention, to be able modify each attention separately
-        unet_new_processors = {key: processor_cls() for key in unet_orig_processors.keys()}
-        try:
-            unet.set_attn_processor(unet_new_processors)
-            yield None
-
-        finally:
-            unet.set_attn_processor(unet_orig_processors)
-
     @staticmethod
     def _resolve_lora_key(model: torch.nn.Module, lora_key: str, prefix: str) -> Tuple[str, torch.nn.Module]:
         assert "." not in lora_key
@@ -86,13 +66,13 @@ class ModelPatcher:
         cls,
         unet: UNet2DConditionModel,
         loras: Iterator[Tuple[LoRAModelRaw, float]],
-        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
+        model_state_dict: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Generator[None, None, None]:
         with cls.apply_lora(
             unet,
             loras=loras,
             prefix="lora_unet_",
-            cached_weights=cached_weights,
+            model_state_dict=model_state_dict,
         ):
             yield
 
@@ -102,9 +82,9 @@ class ModelPatcher:
         cls,
         text_encoder: CLIPTextModel,
         loras: Iterator[Tuple[LoRAModelRaw, float]],
-        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
+        model_state_dict: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Generator[None, None, None]:
-        with cls.apply_lora(text_encoder, loras=loras, prefix="lora_te_", cached_weights=cached_weights):
+        with cls.apply_lora(text_encoder, loras=loras, prefix="lora_te_", model_state_dict=model_state_dict):
             yield
 
     @classmethod
@@ -114,7 +94,7 @@ class ModelPatcher:
         model: AnyModel,
         loras: Iterator[Tuple[LoRAModelRaw, float]],
         prefix: str,
-        cached_weights: Optional[Dict[str, torch.Tensor]] = None,
+        model_state_dict: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Generator[None, None, None]:
         """
         Apply one or more LoRAs to a model.
@@ -122,26 +102,71 @@ class ModelPatcher:
         :param model: The model to patch.
         :param loras: An iterator that returns the LoRA to patch in and its patch weight.
         :param prefix: A string prefix that precedes keys used in the LoRAs weight layers.
-        :cached_weights: Read-only copy of the model's state dict in CPU, for unpatching purposes.
+        :model_state_dict: Read-only copy of the model's state dict in CPU, for unpatching purposes.
         """
-        original_weights = OriginalWeightsStorage(cached_weights)
+        original_weights = {}
         try:
-            for lora_model, lora_weight in loras:
-                LoRAExt.patch_model(
-                    model=model,
-                    prefix=prefix,
-                    lora=lora_model,
-                    lora_weight=lora_weight,
-                    original_weights=original_weights,
-                )
-                del lora_model
+            with torch.no_grad():
+                for lora, lora_weight in loras:
+                    # assert lora.device.type == "cpu"
+                    for layer_key, layer in lora.layers.items():
+                        if not layer_key.startswith(prefix):
+                            continue
 
-            yield
+                        # TODO(ryand): A non-negligible amount of time is currently spent resolving LoRA keys. This
+                        # should be improved in the following ways:
+                        # 1. The key mapping could be more-efficiently pre-computed. This would save time every time a
+                        #    LoRA model is applied.
+                        # 2. From an API perspective, there's no reason that the `ModelPatcher` should be aware of the
+                        #    intricacies of Stable Diffusion key resolution. It should just expect the input LoRA
+                        #    weights to have valid keys.
+                        assert isinstance(model, torch.nn.Module)
+                        module_key, module = cls._resolve_lora_key(model, layer_key, prefix)
+
+                        # All of the LoRA weight calculations will be done on the same device as the module weight.
+                        # (Performance will be best if this is a CUDA device.)
+                        device = module.weight.device
+                        dtype = module.weight.dtype
+
+                        if module_key not in original_weights:
+                            if model_state_dict is not None:  # we were provided with the CPU copy of the state dict
+                                original_weights[module_key] = model_state_dict[module_key + ".weight"]
+                            else:
+                                original_weights[module_key] = module.weight.detach().to(device="cpu", copy=True)
+
+                        layer_scale = layer.alpha / layer.rank if (layer.alpha and layer.rank) else 1.0
+
+                        # We intentionally move to the target device first, then cast. Experimentally, this was found to
+                        # be significantly faster for 16-bit CPU tensors being moved to a CUDA device than doing the
+                        # same thing in a single call to '.to(...)'.
+                        layer.to(device=device, non_blocking=TorchDevice.get_non_blocking(device))
+                        layer.to(dtype=torch.float32, non_blocking=TorchDevice.get_non_blocking(device))
+                        # TODO(ryand): Using torch.autocast(...) over explicit casting may offer a speed benefit on CUDA
+                        # devices here. Experimentally, it was found to be very slow on CPU. More investigation needed.
+                        layer_weight = layer.get_weight(module.weight) * (lora_weight * layer_scale)
+                        layer.to(
+                            device=TorchDevice.CPU_DEVICE,
+                            non_blocking=TorchDevice.get_non_blocking(TorchDevice.CPU_DEVICE),
+                        )
+
+                        assert isinstance(layer_weight, torch.Tensor)  # mypy thinks layer_weight is a float|Any ??!
+                        if module.weight.shape != layer_weight.shape:
+                            # TODO: debug on lycoris
+                            assert hasattr(layer_weight, "reshape")
+                            layer_weight = layer_weight.reshape(module.weight.shape)
+
+                        assert isinstance(layer_weight, torch.Tensor)  # mypy thinks layer_weight is a float|Any ??!
+                        module.weight += layer_weight.to(dtype=dtype, non_blocking=TorchDevice.get_non_blocking(device))
+
+            yield  # wait for context manager exit
 
         finally:
+            assert hasattr(model, "get_submodule")  # mypy not picking up fact that torch.nn.Module has get_submodule()
             with torch.no_grad():
-                for param_key, weight in original_weights.get_changed_weights():
-                    model.get_parameter(param_key).copy_(weight)
+                for module_key, weight in original_weights.items():
+                    model.get_submodule(module_key).weight.copy_(
+                        weight, non_blocking=TorchDevice.get_non_blocking(weight.device)
+                    )
 
     @classmethod
     @contextmanager

invokeai/backend/onnx/onnx_runtime.py

@@ -190,7 +190,12 @@ class IAIOnnxRuntimeModel(RawModel):
         return self.session.run(None, inputs)
 
     # compatability with RawModel ABC
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
         pass
 
     # compatability with diffusers load code

invokeai/backend/raw_model.py

@@ -1,3 +1,15 @@
+"""Base class for 'Raw' models.
+
+The RawModel class is the base class of LoRAModelRaw and TextualInversionModelRaw,
+and is used for type checking of calls to the model patcher. Its main purpose
+is to avoid a circular import issues when lora.py tries to import BaseModelType
+from invokeai.backend.model_manager.config, and the latter tries to import LoRAModelRaw
+from lora.py.
+
+The term 'raw' was introduced to describe a wrapper around a torch.nn.Module
+that adds additional methods and attributes.
+"""
+
 from abc import ABC, abstractmethod
 from typing import Optional
 
@@ -5,18 +17,13 @@ import torch
 
 
 class RawModel(ABC):
-    """Base class for 'Raw' models.
-
-    The RawModel class is the base class of LoRAModelRaw, TextualInversionModelRaw, etc.
-    and is used for type checking of calls to the model patcher. Its main purpose
-    is to avoid a circular import issues when lora.py tries to import BaseModelType
-    from invokeai.backend.model_manager.config, and the latter tries to import LoRAModelRaw
-    from lora.py.
-
-    The term 'raw' was introduced to describe a wrapper around a torch.nn.Module
-    that adds additional methods and attributes.
-    """
+    """Abstract base class for 'Raw' model wrappers."""
 
     @abstractmethod
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+    def to(
+        self,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        non_blocking: bool = False,
+    ) -> None:
         pass