Make quantized loading fast for both T5XXL and FLUX transformer.

.github/workflows/build-container.yml

@@ -13,12 +13,6 @@ on:
     tags:
       - 'v*.*.*'
   workflow_dispatch:
-    inputs:
-      push-to-registry:
-        description: Push the built image to the container registry
-        required: false
-        type: boolean
-        default: false
 
 permissions:
   contents: write
@@ -56,15 +50,16 @@ jobs:
           df -h
 
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v3
 
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@v4
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           images: |
             ghcr.io/${{ github.repository }}
+            ${{ env.DOCKERHUB_REPOSITORY }}
           tags: |
             type=ref,event=branch
             type=ref,event=tag
@@ -77,33 +72,49 @@ jobs:
             suffix=-${{ matrix.gpu-driver }},onlatest=false
 
       - name: Set up QEMU
-        uses: docker/setup-qemu-action@v3
+        uses: docker/setup-qemu-action@v2
 
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
+        uses: docker/setup-buildx-action@v2
         with:
           platforms: ${{ env.PLATFORMS }}
 
       - name: Login to GitHub Container Registry
         if: github.event_name != 'pull_request'
-        uses: docker/login-action@v3
+        uses: docker/login-action@v2
         with:
           registry: ghcr.io
           username: ${{ github.repository_owner }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
+      # - name: Login to Docker Hub
+      #   if: github.event_name != 'pull_request' && vars.DOCKERHUB_REPOSITORY != ''
+      #   uses: docker/login-action@v2
+      #   with:
+      #     username: ${{ secrets.DOCKERHUB_USERNAME }}
+      #     password: ${{ secrets.DOCKERHUB_TOKEN }}
+
       - name: Build container
         timeout-minutes: 40
         id: docker_build
-        uses: docker/build-push-action@v6
+        uses: docker/build-push-action@v4
         with:
           context: .
           file: docker/Dockerfile
           platforms: ${{ env.PLATFORMS }}
-          push: ${{ github.ref == 'refs/heads/main' || github.ref_type == 'tag' || github.event.inputs.push-to-registry }}
+          push: ${{ github.ref == 'refs/heads/main' || github.ref_type == 'tag' }}
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
           cache-from: |
             type=gha,scope=${{ github.ref_name }}-${{ matrix.gpu-driver }}
             type=gha,scope=main-${{ matrix.gpu-driver }}
           cache-to: type=gha,mode=max,scope=${{ github.ref_name }}-${{ matrix.gpu-driver }}
+
+      # - name: Docker Hub Description
+      #   if: github.ref == 'refs/heads/main' || github.ref == 'refs/tags/*' && vars.DOCKERHUB_REPOSITORY != ''
+      #   uses: peter-evans/dockerhub-description@v3
+      #   with:
+      #     username: ${{ secrets.DOCKERHUB_USERNAME }}
+      #     password: ${{ secrets.DOCKERHUB_TOKEN }}
+      #     repository: ${{ vars.DOCKERHUB_REPOSITORY }}
+      #     short-description: ${{ github.event.repository.description }}

.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

.github/workflows/python-tests.yml

@@ -60,7 +60,7 @@ jobs:
             extra-index-url: 'https://download.pytorch.org/whl/cpu'
             github-env: $GITHUB_ENV
           - platform: macos-default
-            os: macOS-14
+            os: macOS-12
             github-env: $GITHUB_ENV
           - platform: windows-cpu
             os: windows-2022

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
+```

docs/help/FAQ.md

@@ -196,22 +196,6 @@ tips to reduce the problem:
     === "12GB VRAM GPU"
 
         This should be sufficient to generate larger images up to about 1280x1280.
-		
-## Checkpoint Models Load Slowly or Use Too Much RAM
-
-The difference between diffusers models (a folder containing multiple
-subfolders) and checkpoint models (a file ending with .safetensors or
-.ckpt) is that InvokeAI is able to load diffusers models into memory
-incrementally, while checkpoint models must be loaded all at
-once. With very large models, or systems with limited RAM, you may
-experience slowdowns and other memory-related issues when loading
-checkpoint models.
-
-To solve this, go to the Model Manager tab (the cube), select the
-checkpoint model that's giving you trouble, and press the "Convert"
-button in the upper right of your browser window. This will conver the
-checkpoint into a diffusers model, after which loading should be
-faster and less memory-intensive.
 
 ## Memory Leak (Linux)
 

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/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

@@ -3,10 +3,8 @@
 
 import io
 import pathlib
-import shutil
 import traceback
 from copy import deepcopy
-from enum import Enum
 from tempfile import TemporaryDirectory
 from typing import List, Optional, Type
 
@@ -19,7 +17,6 @@ from starlette.exceptions import HTTPException
 from typing_extensions import Annotated
 
 from invokeai.app.api.dependencies import ApiDependencies
-from invokeai.app.services.config import get_config
 from invokeai.app.services.model_images.model_images_common import ModelImageFileNotFoundException
 from invokeai.app.services.model_install.model_install_common import ModelInstallJob
 from invokeai.app.services.model_records import (
@@ -34,7 +31,6 @@ from invokeai.backend.model_manager.config import (
     ModelFormat,
     ModelType,
 )
-from invokeai.backend.model_manager.load.model_cache.model_cache_base import CacheStats
 from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch
 from invokeai.backend.model_manager.metadata.metadata_base import ModelMetadataWithFiles, UnknownMetadataException
 from invokeai.backend.model_manager.search import ModelSearch
@@ -54,13 +50,6 @@ class ModelsList(BaseModel):
     model_config = ConfigDict(use_enum_values=True)
 
 
-class CacheType(str, Enum):
-    """Cache type - one of vram or ram."""
-
-    RAM = "RAM"
-    VRAM = "VRAM"
-
-
 def add_cover_image_to_model_config(config: AnyModelConfig, dependencies: Type[ApiDependencies]) -> AnyModelConfig:
     """Add a cover image URL to a model configuration."""
     cover_image = dependencies.invoker.services.model_images.get_url(config.key)
@@ -808,83 +797,3 @@ async def get_starter_models() -> list[StarterModel]:
         model.dependencies = missing_deps
 
     return starter_models
-
-
-@model_manager_router.get(
-    "/model_cache",
-    operation_id="get_cache_size",
-    response_model=float,
-    summary="Get maximum size of model manager RAM or VRAM cache.",
-)
-async def get_cache_size(cache_type: CacheType = Query(description="The cache type", default=CacheType.RAM)) -> float:
-    """Return the current RAM or VRAM cache size setting (in GB)."""
-    cache = ApiDependencies.invoker.services.model_manager.load.ram_cache
-    value = 0.0
-    if cache_type == CacheType.RAM:
-        value = cache.max_cache_size
-    elif cache_type == CacheType.VRAM:
-        value = cache.max_vram_cache_size
-    return value
-
-
-@model_manager_router.put(
-    "/model_cache",
-    operation_id="set_cache_size",
-    response_model=float,
-    summary="Set maximum size of model manager RAM or VRAM cache, optionally writing new value out to invokeai.yaml config file.",
-)
-async def set_cache_size(
-    value: float = Query(description="The new value for the maximum cache size"),
-    cache_type: CacheType = Query(description="The cache type", default=CacheType.RAM),
-    persist: bool = Query(description="Write new value out to invokeai.yaml", default=False),
-) -> float:
-    """Set the current RAM or VRAM cache size setting (in GB). ."""
-    cache = ApiDependencies.invoker.services.model_manager.load.ram_cache
-    app_config = get_config()
-    # Record initial state.
-    vram_old = app_config.vram
-    ram_old = app_config.ram
-
-    # Prepare target state.
-    vram_new = vram_old
-    ram_new = ram_old
-    if cache_type == CacheType.RAM:
-        ram_new = value
-    elif cache_type == CacheType.VRAM:
-        vram_new = value
-    else:
-        raise ValueError(f"Unexpected {cache_type=}.")
-
-    config_path = app_config.config_file_path
-    new_config_path = config_path.with_suffix(".yaml.new")
-
-    try:
-        # Try to apply the target state.
-        cache.max_vram_cache_size = vram_new
-        cache.max_cache_size = ram_new
-        app_config.ram = ram_new
-        app_config.vram = vram_new
-        if persist:
-            app_config.write_file(new_config_path)
-            shutil.move(new_config_path, config_path)
-    except Exception as e:
-        # If there was a failure, restore the initial state.
-        cache.max_cache_size = ram_old
-        cache.max_vram_cache_size = vram_old
-        app_config.ram = ram_old
-        app_config.vram = vram_old
-
-        raise RuntimeError("Failed to update cache size") from e
-    return value
-
-
-@model_manager_router.get(
-    "/stats",
-    operation_id="get_stats",
-    response_model=Optional[CacheStats],
-    summary="Get model manager RAM cache performance statistics.",
-)
-async def get_stats() -> Optional[CacheStats]:
-    """Return performance statistics on the model manager's RAM cache. Will return null if no models have been loaded."""
-
-    return ApiDependencies.invoker.services.model_manager.load.ram_cache.stats

invokeai/app/api/routers/session_queue.py

@@ -11,7 +11,6 @@ from invokeai.app.services.session_queue.session_queue_common import (
     Batch,
     BatchStatus,
     CancelByBatchIDsResult,
-    CancelByDestinationResult,
     ClearResult,
     EnqueueBatchResult,
     PruneResult,
@@ -106,21 +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_destination",
-    operation_id="cancel_by_destination",
-    responses={200: {"model": CancelByBatchIDsResult}},
-)
-async def cancel_by_destination(
-    queue_id: str = Path(description="The queue id to perform this operation on"),
-    destination: str = Query(description="The destination to cancel all queue items for"),
-) -> CancelByDestinationResult:
-    """Immediately cancels all queue items with the given origin"""
-    return ApiDependencies.invoker.services.session_queue.cancel_by_destination(
-        queue_id=queue_id, destination=destination
-    )
-
-
 @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/baseinvocation.py

@@ -20,6 +20,7 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    cast,
 )
 
 import semver
@@ -79,7 +80,7 @@ class UIConfigBase(BaseModel):
     version: str = Field(
         description='The node\'s version. Should be a valid semver string e.g. "1.0.0" or "3.8.13".',
     )
-    node_pack: str = Field(description="The node pack that this node belongs to, will be 'invokeai' for built-in nodes")
+    node_pack: Optional[str] = Field(default=None, description="Whether or not this is a custom node")
     classification: Classification = Field(default=Classification.Stable, description="The node's classification")
 
     model_config = ConfigDict(
@@ -229,16 +230,18 @@ class BaseInvocation(ABC, BaseModel):
     @staticmethod
     def json_schema_extra(schema: dict[str, Any], model_class: Type[BaseInvocation]) -> None:
         """Adds various UI-facing attributes to the invocation's OpenAPI schema."""
-        if title := model_class.UIConfig.title:
-            schema["title"] = title
-        if tags := model_class.UIConfig.tags:
-            schema["tags"] = tags
-        if category := model_class.UIConfig.category:
-            schema["category"] = category
-        if node_pack := model_class.UIConfig.node_pack:
-            schema["node_pack"] = node_pack
-        schema["classification"] = model_class.UIConfig.classification
-        schema["version"] = model_class.UIConfig.version
+        uiconfig = cast(UIConfigBase | None, getattr(model_class, "UIConfig", None))
+        if uiconfig is not None:
+            if uiconfig.title is not None:
+                schema["title"] = uiconfig.title
+            if uiconfig.tags is not None:
+                schema["tags"] = uiconfig.tags
+            if uiconfig.category is not None:
+                schema["category"] = uiconfig.category
+            if uiconfig.node_pack is not None:
+                schema["node_pack"] = uiconfig.node_pack
+            schema["classification"] = uiconfig.classification
+            schema["version"] = uiconfig.version
         if "required" not in schema or not isinstance(schema["required"], list):
             schema["required"] = []
         schema["class"] = "invocation"
@@ -309,7 +312,7 @@ class BaseInvocation(ABC, BaseModel):
         json_schema_extra={"field_kind": FieldKind.NodeAttribute},
     )
 
-    UIConfig: ClassVar[UIConfigBase]
+    UIConfig: ClassVar[Type[UIConfigBase]]
 
     model_config = ConfigDict(
         protected_namespaces=(),
@@ -438,25 +441,30 @@ def invocation(
         validate_fields(cls.model_fields, invocation_type)
 
         # Add OpenAPI schema extras
-        uiconfig: dict[str, Any] = {}
-        uiconfig["title"] = title
-        uiconfig["tags"] = tags
-        uiconfig["category"] = category
-        uiconfig["classification"] = classification
-        # The node pack is the module name - will be "invokeai" for built-in nodes
-        uiconfig["node_pack"] = cls.__module__.split(".")[0]
+        uiconfig_name = cls.__qualname__ + ".UIConfig"
+        if not hasattr(cls, "UIConfig") or cls.UIConfig.__qualname__ != uiconfig_name:
+            cls.UIConfig = type(uiconfig_name, (UIConfigBase,), {})
+        cls.UIConfig.title = title
+        cls.UIConfig.tags = tags
+        cls.UIConfig.category = category
+        cls.UIConfig.classification = classification
+
+        # Grab the node pack's name from the module name, if it's a custom node
+        is_custom_node = cls.__module__.rsplit(".", 1)[0] == "invokeai.app.invocations"
+        if is_custom_node:
+            cls.UIConfig.node_pack = cls.__module__.split(".")[0]
+        else:
+            cls.UIConfig.node_pack = None
 
         if version is not None:
             try:
                 semver.Version.parse(version)
             except ValueError as e:
                 raise InvalidVersionError(f'Invalid version string for node "{invocation_type}": "{version}"') from e
-            uiconfig["version"] = version
+            cls.UIConfig.version = version
         else:
             logger.warn(f'No version specified for node "{invocation_type}", using "1.0.0"')
-            uiconfig["version"] = "1.0.0"
-
-        cls.UIConfig = UIConfigBase(**uiconfig)
+            cls.UIConfig.version = "1.0.0"
 
         if use_cache is not None:
             cls.model_fields["use_cache"].default = use_cache

invokeai/app/invocations/compel.py

@@ -19,8 +19,7 @@ from invokeai.app.invocations.model import CLIPField
 from invokeai.app.invocations.primitives import ConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.ti_utils import generate_ti_list
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
-from invokeai.backend.lora.lora_patcher import LoraPatcher
+from invokeai.backend.lora import LoRAModelRaw
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     BasicConditioningInfo,
@@ -83,10 +82,9 @@ class CompelInvocation(BaseInvocation):
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
             tokenizer_info as tokenizer,
-            LoraPatcher.apply_lora_patches(
-                model=text_encoder,
-                patches=_lora_loader(),
-                prefix="lora_te_",
+            ModelPatcher.apply_lora_text_encoder(
+                text_encoder,
+                loras=_lora_loader(),
                 cached_weights=cached_weights,
             ),
             # Apply CLIP Skip after LoRA to prevent LoRA application from failing on skipped layers.
@@ -179,9 +177,9 @@ class SDXLPromptInvocationBase:
             # apply all patches while the model is on the target device
             text_encoder_info.model_on_device() as (cached_weights, text_encoder),
             tokenizer_info as tokenizer,
-            LoraPatcher.apply_lora_patches(
+            ModelPatcher.apply_lora(
                 text_encoder,
-                patches=_lora_loader(),
+                loras=_lora_loader(),
                 prefix=lora_prefix,
                 cached_weights=cached_weights,
             ),

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/denoise_latents.py

@@ -36,8 +36,7 @@ from invokeai.app.invocations.t2i_adapter import T2IAdapterField
 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.lora_model_raw import LoRAModelRaw
-from invokeai.backend.lora.lora_patcher import LoraPatcher
+from invokeai.backend.lora import LoRAModelRaw
 from invokeai.backend.model_manager import BaseModelType, ModelVariantType
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion import PipelineIntermediateState
@@ -186,7 +185,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
     )
     denoise_mask: Optional[DenoiseMaskField] = InputField(
         default=None,
-        description=FieldDescriptions.denoise_mask,
+        description=FieldDescriptions.mask,
         input=Input.Connection,
         ui_order=8,
     )
@@ -980,10 +979,9 @@ class DenoiseLatentsInvocation(BaseInvocation):
             ModelPatcher.apply_freeu(unet, self.unet.freeu_config),
             SeamlessExt.static_patch_model(unet, self.unet.seamless_axes),  # FIXME
             # Apply the LoRA after unet has been moved to its target device for faster patching.
-            LoraPatcher.apply_lora_patches(
-                model=unet,
-                patches=_lora_loader(),
-                prefix="lora_unet_",
+            ModelPatcher.apply_lora_unet(
+                unet,
+                loras=_lora_loader(),
                 cached_weights=cached_weights,
             ),
         ):

invokeai/app/invocations/fields.py

@@ -40,18 +40,14 @@ class UIType(str, Enum, metaclass=MetaEnum):
 
     # region Model Field Types
     MainModel = "MainModelField"
-    FluxMainModel = "FluxMainModelField"
     SDXLMainModel = "SDXLMainModelField"
     SDXLRefinerModel = "SDXLRefinerModelField"
     ONNXModel = "ONNXModelField"
     VAEModel = "VAEModelField"
-    FluxVAEModel = "FluxVAEModelField"
     LoRAModel = "LoRAModelField"
     ControlNetModel = "ControlNetModelField"
     IPAdapterModel = "IPAdapterModelField"
     T2IAdapterModel = "T2IAdapterModelField"
-    T5EncoderModel = "T5EncoderModelField"
-    CLIPEmbedModel = "CLIPEmbedModelField"
     SpandrelImageToImageModel = "SpandrelImageToImageModelField"
     # endregion
 
@@ -129,17 +125,13 @@ class FieldDescriptions:
     negative_cond = "Negative conditioning tensor"
     noise = "Noise tensor"
     clip = "CLIP (tokenizer, text encoder, LoRAs) and skipped layer count"
-    t5_encoder = "T5 tokenizer and text encoder"
-    clip_embed_model = "CLIP Embed loader"
     unet = "UNet (scheduler, LoRAs)"
-    transformer = "Transformer"
     vae = "VAE"
     cond = "Conditioning tensor"
     controlnet_model = "ControlNet model to load"
     vae_model = "VAE model to load"
     lora_model = "LoRA model to load"
     main_model = "Main model (UNet, VAE, CLIP) to load"
-    flux_model = "Flux model (Transformer) to load"
     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"
@@ -181,7 +173,7 @@ class FieldDescriptions:
     )
     num_1 = "The first number"
     num_2 = "The second number"
-    denoise_mask = "A mask of the region to apply the denoising process to."
+    mask = "The mask to use for the operation"
     board = "The board to save the image to"
     image = "The image to process"
     tile_size = "Tile size"
@@ -239,12 +231,6 @@ class ColorField(BaseModel):
         return (self.r, self.g, self.b, self.a)
 
 
-class FluxConditioningField(BaseModel):
-    """A conditioning tensor primitive value"""
-
-    conditioning_name: str = Field(description="The name of conditioning tensor")
-
-
 class ConditioningField(BaseModel):
     """A conditioning tensor primitive value"""
 

invokeai/app/invocations/flux_denoise.py

@@ -1,267 +0,0 @@
-from typing import Callable, Iterator, Optional, Tuple
-
-import torch
-import torchvision.transforms as tv_transforms
-from torchvision.transforms.functional import resize as tv_resize
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
-from invokeai.app.invocations.fields import (
-    DenoiseMaskField,
-    FieldDescriptions,
-    FluxConditioningField,
-    Input,
-    InputField,
-    LatentsField,
-    WithBoard,
-    WithMetadata,
-)
-from invokeai.app.invocations.model import TransformerField
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.flux.denoise import denoise
-from invokeai.backend.flux.inpaint_extension import InpaintExtension
-from invokeai.backend.flux.model import Flux
-from invokeai.backend.flux.sampling_utils import (
-    clip_timestep_schedule,
-    generate_img_ids,
-    get_noise,
-    get_schedule,
-    pack,
-    unpack,
-)
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
-from invokeai.backend.lora.lora_patcher import LoraPatcher
-from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
-from invokeai.backend.util.devices import TorchDevice
-
-
-@invocation(
-    "flux_denoise",
-    title="FLUX Denoise",
-    tags=["image", "flux"],
-    category="image",
-    version="1.0.0",
-    classification=Classification.Prototype,
-)
-class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
-    """Run denoising process with a FLUX transformer model."""
-
-    # If latents is provided, this means we are doing image-to-image.
-    latents: Optional[LatentsField] = InputField(
-        default=None,
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    # denoise_mask is used for image-to-image inpainting. Only the masked region is modified.
-    denoise_mask: Optional[DenoiseMaskField] = InputField(
-        default=None,
-        description=FieldDescriptions.denoise_mask,
-        input=Input.Connection,
-    )
-    denoising_start: float = InputField(
-        default=0.0,
-        ge=0,
-        le=1,
-        description=FieldDescriptions.denoising_start,
-    )
-    denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
-    transformer: TransformerField = InputField(
-        description=FieldDescriptions.flux_model,
-        input=Input.Connection,
-        title="Transformer",
-    )
-    positive_text_conditioning: FluxConditioningField = InputField(
-        description=FieldDescriptions.positive_cond, input=Input.Connection
-    )
-    width: int = InputField(default=1024, multiple_of=16, description="Width of the generated image.")
-    height: int = InputField(default=1024, multiple_of=16, description="Height of the generated image.")
-    num_steps: int = InputField(
-        default=4, description="Number of diffusion steps. Recommended values are schnell: 4, dev: 50."
-    )
-    guidance: float = InputField(
-        default=4.0,
-        description="The guidance strength. Higher values adhere more strictly to the prompt, and will produce less diverse images. FLUX dev only, ignored for schnell.",
-    )
-    seed: int = InputField(default=0, description="Randomness seed for reproducibility.")
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        latents = self._run_diffusion(context)
-        latents = latents.detach().to("cpu")
-
-        name = context.tensors.save(tensor=latents)
-        return LatentsOutput.build(latents_name=name, latents=latents, seed=None)
-
-    def _run_diffusion(
-        self,
-        context: InvocationContext,
-    ):
-        inference_dtype = torch.bfloat16
-
-        # Load the conditioning data.
-        cond_data = context.conditioning.load(self.positive_text_conditioning.conditioning_name)
-        assert len(cond_data.conditionings) == 1
-        flux_conditioning = cond_data.conditionings[0]
-        assert isinstance(flux_conditioning, FLUXConditioningInfo)
-        flux_conditioning = flux_conditioning.to(dtype=inference_dtype)
-        t5_embeddings = flux_conditioning.t5_embeds
-        clip_embeddings = flux_conditioning.clip_embeds
-
-        # Load the input latents, if provided.
-        init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None
-        if init_latents is not None:
-            init_latents = init_latents.to(device=TorchDevice.choose_torch_device(), dtype=inference_dtype)
-
-        # Prepare input noise.
-        noise = get_noise(
-            num_samples=1,
-            height=self.height,
-            width=self.width,
-            device=TorchDevice.choose_torch_device(),
-            dtype=inference_dtype,
-            seed=self.seed,
-        )
-
-        transformer_info = context.models.load(self.transformer.transformer)
-        is_schnell = "schnell" in transformer_info.config.config_path
-
-        # Calculate the timestep schedule.
-        image_seq_len = noise.shape[-1] * noise.shape[-2] // 4
-        timesteps = get_schedule(
-            num_steps=self.num_steps,
-            image_seq_len=image_seq_len,
-            shift=not is_schnell,
-        )
-
-        # Clip the timesteps schedule based on denoising_start and denoising_end.
-        timesteps = clip_timestep_schedule(timesteps, self.denoising_start, self.denoising_end)
-
-        # Prepare input latent image.
-        if init_latents is not None:
-            # If init_latents is provided, we are doing image-to-image.
-
-            if is_schnell:
-                context.logger.warning(
-                    "Running image-to-image with a FLUX schnell model. This is not recommended. The results are likely "
-                    "to be poor. Consider using a FLUX dev model instead."
-                )
-
-            # Noise the orig_latents by the appropriate amount for the first timestep.
-            t_0 = timesteps[0]
-            x = t_0 * noise + (1.0 - t_0) * init_latents
-        else:
-            # init_latents are not provided, so we are not doing image-to-image (i.e. we are starting from pure noise).
-            if self.denoising_start > 1e-5:
-                raise ValueError("denoising_start should be 0 when initial latents are not provided.")
-
-            x = noise
-
-        # If len(timesteps) == 1, then short-circuit. We are just noising the input latents, but not taking any
-        # denoising steps.
-        if len(timesteps) <= 1:
-            return x
-
-        inpaint_mask = self._prep_inpaint_mask(context, x)
-
-        b, _c, h, w = x.shape
-        img_ids = generate_img_ids(h=h, w=w, batch_size=b, device=x.device, dtype=x.dtype)
-
-        bs, t5_seq_len, _ = t5_embeddings.shape
-        txt_ids = torch.zeros(bs, t5_seq_len, 3, dtype=inference_dtype, device=TorchDevice.choose_torch_device())
-
-        # Pack all latent tensors.
-        init_latents = pack(init_latents) if init_latents is not None else None
-        inpaint_mask = pack(inpaint_mask) if inpaint_mask is not None else None
-        noise = pack(noise)
-        x = pack(x)
-
-        # Now that we have 'packed' the latent tensors, verify that we calculated the image_seq_len correctly.
-        assert image_seq_len == x.shape[1]
-
-        # Prepare inpaint extension.
-        inpaint_extension: InpaintExtension | None = None
-        if inpaint_mask is not None:
-            assert init_latents is not None
-            inpaint_extension = InpaintExtension(
-                init_latents=init_latents,
-                inpaint_mask=inpaint_mask,
-                noise=noise,
-            )
-
-        with (
-            transformer_info.model_on_device() as (cached_weights, transformer),
-            # Apply the LoRA after transformer has been moved to its target device for faster patching.
-            LoraPatcher.apply_lora_patches(
-                model=transformer,
-                patches=self._lora_iterator(context),
-                prefix="",
-                cached_weights=cached_weights,
-            ),
-        ):
-            assert isinstance(transformer, Flux)
-
-            x = denoise(
-                model=transformer,
-                img=x,
-                img_ids=img_ids,
-                txt=t5_embeddings,
-                txt_ids=txt_ids,
-                vec=clip_embeddings,
-                timesteps=timesteps,
-                step_callback=self._build_step_callback(context),
-                guidance=self.guidance,
-                inpaint_extension=inpaint_extension,
-            )
-
-        x = unpack(x.float(), self.height, self.width)
-        return x
-
-    def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> torch.Tensor | None:
-        """Prepare the inpaint mask.
-
-        - Loads the mask
-        - Resizes if necessary
-        - Casts to same device/dtype as latents
-        - Expands mask to the same shape as latents so that they line up after 'packing'
-
-        Args:
-            context (InvocationContext): The invocation context, for loading the inpaint mask.
-            latents (torch.Tensor): A latent image tensor. In 'unpacked' format. Used to determine the target shape,
-                device, and dtype for the inpaint mask.
-
-        Returns:
-            torch.Tensor | None: Inpaint mask.
-        """
-        if self.denoise_mask is None:
-            return None
-
-        mask = context.tensors.load(self.denoise_mask.mask_name)
-
-        _, _, latent_height, latent_width = latents.shape
-        mask = tv_resize(
-            img=mask,
-            size=[latent_height, latent_width],
-            interpolation=tv_transforms.InterpolationMode.BILINEAR,
-            antialias=False,
-        )
-
-        mask = mask.to(device=latents.device, dtype=latents.dtype)
-
-        # Expand the inpaint mask to the same shape as `latents` so that when we 'pack' `mask` it lines up with
-        # `latents`.
-        return mask.expand_as(latents)
-
-    def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[LoRAModelRaw, float]]:
-        for lora in self.transformer.loras:
-            lora_info = context.models.load(lora.lora)
-            assert isinstance(lora_info.model, LoRAModelRaw)
-            yield (lora_info.model, lora.weight)
-            del lora_info
-
-    def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
-        def step_callback(state: PipelineIntermediateState) -> None:
-            state.latents = unpack(state.latents.float(), self.height, self.width).squeeze()
-            context.util.flux_step_callback(state)
-
-        return step_callback

invokeai/app/invocations/flux_lora_loader.py

@@ -1,53 +0,0 @@
-from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
-from invokeai.app.invocations.model import LoRAField, ModelIdentifierField, TransformerField
-from invokeai.app.services.shared.invocation_context import InvocationContext
-
-
-@invocation_output("flux_lora_loader_output")
-class FluxLoRALoaderOutput(BaseInvocationOutput):
-    """FLUX LoRA Loader Output"""
-
-    transformer: TransformerField = OutputField(
-        default=None, description=FieldDescriptions.transformer, title="FLUX Transformer"
-    )
-
-
-@invocation(
-    "flux_lora_loader",
-    title="FLUX LoRA",
-    tags=["lora", "model", "flux"],
-    category="model",
-    version="1.0.0",
-)
-class FluxLoRALoaderInvocation(BaseInvocation):
-    """Apply a LoRA model to a FLUX transformer."""
-
-    lora: ModelIdentifierField = InputField(
-        description=FieldDescriptions.lora_model, title="LoRA", ui_type=UIType.LoRAModel
-    )
-    weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight)
-    transformer: TransformerField = InputField(
-        description=FieldDescriptions.transformer,
-        input=Input.Connection,
-        title="FLUX Transformer",
-    )
-
-    def invoke(self, context: InvocationContext) -> FluxLoRALoaderOutput:
-        lora_key = self.lora.key
-
-        if not context.models.exists(lora_key):
-            raise ValueError(f"Unknown lora: {lora_key}!")
-
-        if any(lora.lora.key == lora_key for lora in self.transformer.loras):
-            raise Exception(f'LoRA "{lora_key}" already applied to transformer.')
-
-        transformer = self.transformer.model_copy(deep=True)
-        transformer.loras.append(
-            LoRAField(
-                lora=self.lora,
-                weight=self.weight,
-            )
-        )
-
-        return FluxLoRALoaderOutput(transformer=transformer)

invokeai/app/invocations/flux_text_encoder.py

@@ -1,92 +0,0 @@
-from typing import Literal
-
-import torch
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
-from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField
-from invokeai.app.invocations.model import CLIPField, T5EncoderField
-from invokeai.app.invocations.primitives import FluxConditioningOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.flux.modules.conditioner import HFEncoder
-from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, FLUXConditioningInfo
-
-
-@invocation(
-    "flux_text_encoder",
-    title="FLUX Text Encoding",
-    tags=["prompt", "conditioning", "flux"],
-    category="conditioning",
-    version="1.0.0",
-    classification=Classification.Prototype,
-)
-class FluxTextEncoderInvocation(BaseInvocation):
-    """Encodes and preps a prompt for a flux image."""
-
-    clip: CLIPField = InputField(
-        title="CLIP",
-        description=FieldDescriptions.clip,
-        input=Input.Connection,
-    )
-    t5_encoder: T5EncoderField = InputField(
-        title="T5Encoder",
-        description=FieldDescriptions.t5_encoder,
-        input=Input.Connection,
-    )
-    t5_max_seq_len: Literal[256, 512] = InputField(
-        description="Max sequence length for the T5 encoder. Expected to be 256 for FLUX schnell models and 512 for FLUX dev models."
-    )
-    prompt: str = InputField(description="Text prompt to encode.")
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> FluxConditioningOutput:
-        # Note: The T5 and CLIP encoding are done in separate functions to ensure that all model references are locally
-        # scoped. This ensures that the T5 model can be freed and gc'd before loading the CLIP model (if necessary).
-        t5_embeddings = self._t5_encode(context)
-        clip_embeddings = self._clip_encode(context)
-        conditioning_data = ConditioningFieldData(
-            conditionings=[FLUXConditioningInfo(clip_embeds=clip_embeddings, t5_embeds=t5_embeddings)]
-        )
-
-        conditioning_name = context.conditioning.save(conditioning_data)
-        return FluxConditioningOutput.build(conditioning_name)
-
-    def _t5_encode(self, context: InvocationContext) -> torch.Tensor:
-        t5_tokenizer_info = context.models.load(self.t5_encoder.tokenizer)
-        t5_text_encoder_info = context.models.load(self.t5_encoder.text_encoder)
-
-        prompt = [self.prompt]
-
-        with (
-            t5_text_encoder_info as t5_text_encoder,
-            t5_tokenizer_info as t5_tokenizer,
-        ):
-            assert isinstance(t5_text_encoder, T5EncoderModel)
-            assert isinstance(t5_tokenizer, T5Tokenizer)
-
-            t5_encoder = HFEncoder(t5_text_encoder, t5_tokenizer, False, self.t5_max_seq_len)
-
-            prompt_embeds = t5_encoder(prompt)
-
-        assert isinstance(prompt_embeds, torch.Tensor)
-        return prompt_embeds
-
-    def _clip_encode(self, context: InvocationContext) -> torch.Tensor:
-        clip_tokenizer_info = context.models.load(self.clip.tokenizer)
-        clip_text_encoder_info = context.models.load(self.clip.text_encoder)
-
-        prompt = [self.prompt]
-
-        with (
-            clip_text_encoder_info as clip_text_encoder,
-            clip_tokenizer_info as clip_tokenizer,
-        ):
-            assert isinstance(clip_text_encoder, CLIPTextModel)
-            assert isinstance(clip_tokenizer, CLIPTokenizer)
-
-            clip_encoder = HFEncoder(clip_text_encoder, clip_tokenizer, True, 77)
-
-            pooled_prompt_embeds = clip_encoder(prompt)
-
-        assert isinstance(pooled_prompt_embeds, torch.Tensor)
-        return pooled_prompt_embeds

invokeai/app/invocations/flux_text_to_image.py

@@ -0,0 +1,278 @@
+from pathlib import Path
+from typing import Literal
+
+import torch
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
+from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
+from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
+from optimum.quanto import qfloat8
+from PIL import Image
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from transformers.models.auto import AutoModelForTextEncoding
+
+from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
+from invokeai.app.invocations.fields import InputField, WithBoard, WithMetadata
+from invokeai.app.invocations.primitives import ImageOutput
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.quantization.fast_quantized_diffusion_model import FastQuantizedDiffusersModel
+from invokeai.backend.quantization.fast_quantized_transformers_model import FastQuantizedTransformersModel
+from invokeai.backend.util.devices import TorchDevice
+
+TFluxModelKeys = Literal["flux-schnell"]
+FLUX_MODELS: dict[TFluxModelKeys, str] = {"flux-schnell": "black-forest-labs/FLUX.1-schnell"}
+
+
+class QuantizedFluxTransformer2DModel(FastQuantizedDiffusersModel):
+    base_class = FluxTransformer2DModel
+
+
+class QuantizedModelForTextEncoding(FastQuantizedTransformersModel):
+    auto_class = AutoModelForTextEncoding
+
+
+@invocation(
+    "flux_text_to_image",
+    title="FLUX Text to Image",
+    tags=["image"],
+    category="image",
+    version="1.0.0",
+)
+class FluxTextToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Text-to-image generation using a FLUX model."""
+
+    model: TFluxModelKeys = InputField(description="The FLUX model to use for text-to-image generation.")
+    use_8bit: bool = InputField(
+        default=False, description="Whether to quantize the transformer model to 8-bit precision."
+    )
+    positive_prompt: str = InputField(description="Positive prompt for text-to-image generation.")
+    width: int = InputField(default=1024, multiple_of=16, description="Width of the generated image.")
+    height: int = InputField(default=1024, multiple_of=16, description="Height of the generated image.")
+    num_steps: int = InputField(default=4, description="Number of diffusion steps.")
+    guidance: float = InputField(
+        default=4.0,
+        description="The guidance strength. Higher values adhere more strictly to the prompt, and will produce less diverse images.",
+    )
+    seed: int = InputField(default=0, description="Randomness seed for reproducibility.")
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        model_path = context.models.download_and_cache_model(FLUX_MODELS[self.model])
+
+        t5_embeddings, clip_embeddings = self._encode_prompt(context, model_path)
+        latents = self._run_diffusion(context, model_path, clip_embeddings, t5_embeddings)
+        image = self._run_vae_decoding(context, model_path, latents)
+        image_dto = context.images.save(image=image)
+        return ImageOutput.build(image_dto)
+
+    def _encode_prompt(self, context: InvocationContext, flux_model_dir: Path) -> tuple[torch.Tensor, torch.Tensor]:
+        # Determine the T5 max sequence length based on the model.
+        if self.model == "flux-schnell":
+            max_seq_len = 256
+        # elif self.model == "flux-dev":
+        #     max_seq_len = 512
+        else:
+            raise ValueError(f"Unknown model: {self.model}")
+
+        # Load the CLIP tokenizer.
+        clip_tokenizer_path = flux_model_dir / "tokenizer"
+        clip_tokenizer = CLIPTokenizer.from_pretrained(clip_tokenizer_path, local_files_only=True)
+        assert isinstance(clip_tokenizer, CLIPTokenizer)
+
+        # Load the T5 tokenizer.
+        t5_tokenizer_path = flux_model_dir / "tokenizer_2"
+        t5_tokenizer = T5TokenizerFast.from_pretrained(t5_tokenizer_path, local_files_only=True)
+        assert isinstance(t5_tokenizer, T5TokenizerFast)
+
+        clip_text_encoder_path = flux_model_dir / "text_encoder"
+        t5_text_encoder_path = flux_model_dir / "text_encoder_2"
+        with (
+            context.models.load_local_model(
+                model_path=clip_text_encoder_path, loader=self._load_flux_text_encoder
+            ) as clip_text_encoder,
+            context.models.load_local_model(
+                model_path=t5_text_encoder_path, loader=self._load_flux_text_encoder_2
+            ) as t5_text_encoder,
+        ):
+            assert isinstance(clip_text_encoder, CLIPTextModel)
+            assert isinstance(t5_text_encoder, T5EncoderModel)
+            pipeline = FluxPipeline(
+                scheduler=None,
+                vae=None,
+                text_encoder=clip_text_encoder,
+                tokenizer=clip_tokenizer,
+                text_encoder_2=t5_text_encoder,
+                tokenizer_2=t5_tokenizer,
+                transformer=None,
+            )
+
+            # prompt_embeds: T5 embeddings
+            # pooled_prompt_embeds: CLIP embeddings
+            prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
+                prompt=self.positive_prompt,
+                prompt_2=self.positive_prompt,
+                device=TorchDevice.choose_torch_device(),
+                max_sequence_length=max_seq_len,
+            )
+
+        assert isinstance(prompt_embeds, torch.Tensor)
+        assert isinstance(pooled_prompt_embeds, torch.Tensor)
+        return prompt_embeds, pooled_prompt_embeds
+
+    def _run_diffusion(
+        self,
+        context: InvocationContext,
+        flux_model_dir: Path,
+        clip_embeddings: torch.Tensor,
+        t5_embeddings: torch.Tensor,
+    ):
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(flux_model_dir / "scheduler", local_files_only=True)
+
+        # HACK(ryand): Manually empty the cache. Currently we don't check the size of the model before loading it from
+        # disk. Since the transformer model is large (24GB), there's a good chance that it will OOM on 32GB RAM systems
+        # if the cache is not empty.
+        context.models._services.model_manager.load.ram_cache.make_room(24 * 2**30)
+
+        transformer_path = flux_model_dir / "transformer"
+        with context.models.load_local_model(
+            model_path=transformer_path, loader=self._load_flux_transformer
+        ) as transformer:
+            assert isinstance(transformer, FluxTransformer2DModel)
+
+            flux_pipeline_with_transformer = FluxPipeline(
+                scheduler=scheduler,
+                vae=None,
+                text_encoder=None,
+                tokenizer=None,
+                text_encoder_2=None,
+                tokenizer_2=None,
+                transformer=transformer,
+            )
+
+            t5_embeddings = t5_embeddings.to(dtype=transformer.dtype)
+            clip_embeddings = clip_embeddings.to(dtype=transformer.dtype)
+
+            latents = flux_pipeline_with_transformer(
+                height=self.height,
+                width=self.width,
+                num_inference_steps=self.num_steps,
+                guidance_scale=self.guidance,
+                generator=torch.Generator().manual_seed(self.seed),
+                prompt_embeds=t5_embeddings,
+                pooled_prompt_embeds=clip_embeddings,
+                output_type="latent",
+                return_dict=False,
+            )[0]
+
+        assert isinstance(latents, torch.Tensor)
+        return latents
+
+    def _run_vae_decoding(
+        self,
+        context: InvocationContext,
+        flux_model_dir: Path,
+        latents: torch.Tensor,
+    ) -> Image.Image:
+        vae_path = flux_model_dir / "vae"
+        with context.models.load_local_model(model_path=vae_path, loader=self._load_flux_vae) as vae:
+            assert isinstance(vae, AutoencoderKL)
+
+            flux_pipeline_with_vae = FluxPipeline(
+                scheduler=None,
+                vae=vae,
+                text_encoder=None,
+                tokenizer=None,
+                text_encoder_2=None,
+                tokenizer_2=None,
+                transformer=None,
+            )
+
+            latents = flux_pipeline_with_vae._unpack_latents(
+                latents, self.height, self.width, flux_pipeline_with_vae.vae_scale_factor
+            )
+            latents = (
+                latents / flux_pipeline_with_vae.vae.config.scaling_factor
+            ) + flux_pipeline_with_vae.vae.config.shift_factor
+            latents = latents.to(dtype=vae.dtype)
+            image = flux_pipeline_with_vae.vae.decode(latents, return_dict=False)[0]
+            image = flux_pipeline_with_vae.image_processor.postprocess(image, output_type="pil")[0]
+
+        assert isinstance(image, Image.Image)
+        return image
+
+    @staticmethod
+    def _load_flux_text_encoder(path: Path) -> CLIPTextModel:
+        model = CLIPTextModel.from_pretrained(path, local_files_only=True)
+        assert isinstance(model, CLIPTextModel)
+        return model
+
+    def _load_flux_text_encoder_2(self, path: Path) -> T5EncoderModel:
+        if self.use_8bit:
+            model_8bit_path = path / "quantized"
+            if model_8bit_path.exists():
+                # The quantized model exists, load it.
+                # TODO(ryand): The requantize(...) operation in from_pretrained(...) is very slow. This seems like
+                # something that we should be able to make much faster.
+                q_model = QuantizedModelForTextEncoding.from_pretrained(model_8bit_path)
+
+                # Access the underlying wrapped model.
+                # We access the wrapped model, even though it is private, because it simplifies the type checking by
+                # always returning a T5EncoderModel from this function.
+                model = q_model._wrapped
+            else:
+                # The quantized model does not exist yet, quantize and save it.
+                # TODO(ryand): dtype?
+                model = T5EncoderModel.from_pretrained(path, local_files_only=True)
+                assert isinstance(model, T5EncoderModel)
+
+                q_model = QuantizedModelForTextEncoding.quantize(model, weights=qfloat8)
+
+                model_8bit_path.mkdir(parents=True, exist_ok=True)
+                q_model.save_pretrained(model_8bit_path)
+
+                # (See earlier comment about accessing the wrapped model.)
+                model = q_model._wrapped
+        else:
+            model = T5EncoderModel.from_pretrained(path, local_files_only=True)
+
+        assert isinstance(model, T5EncoderModel)
+        return model
+
+    def _load_flux_transformer(self, path: Path) -> FluxTransformer2DModel:
+        if self.use_8bit:
+            model_8bit_path = path / "quantized"
+            if model_8bit_path.exists():
+                # The quantized model exists, load it.
+                # TODO(ryand): The requantize(...) operation in from_pretrained(...) is very slow. This seems like
+                # something that we should be able to make much faster.
+                q_model = QuantizedFluxTransformer2DModel.from_pretrained(model_8bit_path)
+
+                # Access the underlying wrapped model.
+                # We access the wrapped model, even though it is private, because it simplifies the type checking by
+                # always returning a FluxTransformer2DModel from this function.
+                model = q_model._wrapped
+            else:
+                # The quantized model does not exist yet, quantize and save it.
+                # TODO(ryand): Loading in float16 and then quantizing seems to result in NaNs. In order to run this on
+                # GPUs that don't support bfloat16, we would need to host the quantized model instead of generating it
+                # here.
+                model = FluxTransformer2DModel.from_pretrained(path, local_files_only=True, torch_dtype=torch.bfloat16)
+                assert isinstance(model, FluxTransformer2DModel)
+
+                q_model = QuantizedFluxTransformer2DModel.quantize(model, weights=qfloat8)
+
+                model_8bit_path.mkdir(parents=True, exist_ok=True)
+                q_model.save_pretrained(model_8bit_path)
+
+                # (See earlier comment about accessing the wrapped model.)
+                model = q_model._wrapped
+        else:
+            model = FluxTransformer2DModel.from_pretrained(path, local_files_only=True, torch_dtype=torch.bfloat16)
+
+        assert isinstance(model, FluxTransformer2DModel)
+        return model
+
+    @staticmethod
+    def _load_flux_vae(path: Path) -> AutoencoderKL:
+        model = AutoencoderKL.from_pretrained(path, local_files_only=True)
+        assert isinstance(model, AutoencoderKL)
+        return model

invokeai/app/invocations/flux_vae_decode.py

@@ -1,60 +0,0 @@
-import torch
-from einops import rearrange
-from PIL import Image
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    Input,
-    InputField,
-    LatentsField,
-    WithBoard,
-    WithMetadata,
-)
-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.flux.modules.autoencoder import AutoEncoder
-from invokeai.backend.model_manager.load.load_base import LoadedModel
-from invokeai.backend.util.devices import TorchDevice
-
-
-@invocation(
-    "flux_vae_decode",
-    title="FLUX Latents to Image",
-    tags=["latents", "image", "vae", "l2i", "flux"],
-    category="latents",
-    version="1.0.0",
-)
-class FluxVaeDecodeInvocation(BaseInvocation, WithMetadata, WithBoard):
-    """Generates an image from latents."""
-
-    latents: LatentsField = InputField(
-        description=FieldDescriptions.latents,
-        input=Input.Connection,
-    )
-    vae: VAEField = InputField(
-        description=FieldDescriptions.vae,
-        input=Input.Connection,
-    )
-
-    def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Image:
-        with vae_info as vae:
-            assert isinstance(vae, AutoEncoder)
-            latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=TorchDevice.choose_torch_dtype())
-            img = vae.decode(latents)
-
-        img = img.clamp(-1, 1)
-        img = rearrange(img[0], "c h w -> h w c")  # noqa: F821
-        img_pil = Image.fromarray((127.5 * (img + 1.0)).byte().cpu().numpy())
-        return img_pil
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> ImageOutput:
-        latents = context.tensors.load(self.latents.latents_name)
-        vae_info = context.models.load(self.vae.vae)
-        image = self._vae_decode(vae_info=vae_info, latents=latents)
-
-        TorchDevice.empty_cache()
-        image_dto = context.images.save(image=image)
-        return ImageOutput.build(image_dto)

invokeai/app/invocations/flux_vae_encode.py

@@ -1,67 +0,0 @@
-import einops
-import torch
-
-from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.fields import (
-    FieldDescriptions,
-    ImageField,
-    Input,
-    InputField,
-)
-from invokeai.app.invocations.model import VAEField
-from invokeai.app.invocations.primitives import LatentsOutput
-from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.flux.modules.autoencoder import AutoEncoder
-from invokeai.backend.model_manager import LoadedModel
-from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-from invokeai.backend.util.devices import TorchDevice
-
-
-@invocation(
-    "flux_vae_encode",
-    title="FLUX Image to Latents",
-    tags=["latents", "image", "vae", "i2l", "flux"],
-    category="latents",
-    version="1.0.0",
-)
-class FluxVaeEncodeInvocation(BaseInvocation):
-    """Encodes an image into latents."""
-
-    image: ImageField = InputField(
-        description="The image to encode.",
-    )
-    vae: VAEField = InputField(
-        description=FieldDescriptions.vae,
-        input=Input.Connection,
-    )
-
-    @staticmethod
-    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
-        # TODO(ryand): Expose seed parameter at the invocation level.
-        # TODO(ryand): Write a util function for generating random tensors that is consistent across devices / dtypes.
-        # There's a starting point in get_noise(...), but it needs to be extracted and generalized. This function
-        # should be used for VAE encode sampling.
-        generator = torch.Generator(device=TorchDevice.choose_torch_device()).manual_seed(0)
-        with vae_info as vae:
-            assert isinstance(vae, AutoEncoder)
-            image_tensor = image_tensor.to(
-                device=TorchDevice.choose_torch_device(), dtype=TorchDevice.choose_torch_dtype()
-            )
-            latents = vae.encode(image_tensor, sample=True, generator=generator)
-            return latents
-
-    @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> LatentsOutput:
-        image = context.images.get_pil(self.image.image_name)
-
-        vae_info = context.models.load(self.vae.vae)
-
-        image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
-        if image_tensor.dim() == 3:
-            image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
-
-        latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
-
-        latents = latents.to("cpu")
-        name = context.tensors.save(tensor=latents)
-        return LatentsOutput.build(latents_name=name, latents=latents, seed=None)

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/mask.py

@@ -126,7 +126,7 @@ class ImageMaskToTensorInvocation(BaseInvocation, WithMetadata):
     title="Tensor Mask to Image",
     tags=["mask"],
     category="mask",
-    version="1.1.0",
+    version="1.0.0",
 )
 class MaskTensorToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Convert a mask tensor to an image."""
@@ -135,11 +135,6 @@ class MaskTensorToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
         mask = context.tensors.load(self.mask.tensor_name)
-
-        # Squeeze the channel dimension if it exists.
-        if mask.dim() == 3:
-            mask = mask.squeeze(0)
-
         # Ensure that the mask is binary.
         if mask.dtype != torch.bool:
             mask = mask > 0.5

invokeai/app/invocations/model.py

@@ -1,5 +1,5 @@
 import copy
-from typing import List, Literal, Optional
+from typing import List, Optional
 
 from pydantic import BaseModel, Field
 
@@ -13,14 +13,7 @@ from invokeai.app.invocations.baseinvocation import (
 from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.shared.models import FreeUConfig
-from invokeai.backend.flux.util import max_seq_lengths
-from invokeai.backend.model_manager.config import (
-    AnyModelConfig,
-    BaseModelType,
-    CheckpointConfigBase,
-    ModelType,
-    SubModelType,
-)
+from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelType, SubModelType
 
 
 class ModelIdentifierField(BaseModel):
@@ -67,16 +60,6 @@ class CLIPField(BaseModel):
     loras: List[LoRAField] = Field(description="LoRAs to apply on model loading")
 
 
-class TransformerField(BaseModel):
-    transformer: ModelIdentifierField = Field(description="Info to load Transformer submodel")
-    loras: List[LoRAField] = Field(description="LoRAs to apply on model loading")
-
-
-class T5EncoderField(BaseModel):
-    tokenizer: ModelIdentifierField = Field(description="Info to load tokenizer submodel")
-    text_encoder: ModelIdentifierField = Field(description="Info to load text_encoder submodel")
-
-
 class VAEField(BaseModel):
     vae: ModelIdentifierField = Field(description="Info to load vae submodel")
     seamless_axes: List[str] = Field(default_factory=list, description='Axes("x" and "y") to which apply seamless')
@@ -139,78 +122,6 @@ class ModelIdentifierInvocation(BaseInvocation):
         return ModelIdentifierOutput(model=self.model)
 
 
-@invocation_output("flux_model_loader_output")
-class FluxModelLoaderOutput(BaseInvocationOutput):
-    """Flux base model loader output"""
-
-    transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer")
-    clip: CLIPField = OutputField(description=FieldDescriptions.clip, title="CLIP")
-    t5_encoder: T5EncoderField = OutputField(description=FieldDescriptions.t5_encoder, title="T5 Encoder")
-    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
-    max_seq_len: Literal[256, 512] = OutputField(
-        description="The max sequence length to used for the T5 encoder. (256 for schnell transformer, 512 for dev transformer)",
-        title="Max Seq Length",
-    )
-
-
-@invocation(
-    "flux_model_loader",
-    title="Flux Main Model",
-    tags=["model", "flux"],
-    category="model",
-    version="1.0.4",
-    classification=Classification.Prototype,
-)
-class FluxModelLoaderInvocation(BaseInvocation):
-    """Loads a flux base model, outputting its submodels."""
-
-    model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.flux_model,
-        ui_type=UIType.FluxMainModel,
-        input=Input.Direct,
-    )
-
-    t5_encoder_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.t5_encoder, ui_type=UIType.T5EncoderModel, input=Input.Direct, title="T5 Encoder"
-    )
-
-    clip_embed_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.clip_embed_model,
-        ui_type=UIType.CLIPEmbedModel,
-        input=Input.Direct,
-        title="CLIP Embed",
-    )
-
-    vae_model: ModelIdentifierField = InputField(
-        description=FieldDescriptions.vae_model, ui_type=UIType.FluxVAEModel, title="VAE"
-    )
-
-    def invoke(self, context: InvocationContext) -> FluxModelLoaderOutput:
-        for key in [self.model.key, self.t5_encoder_model.key, self.clip_embed_model.key, self.vae_model.key]:
-            if not context.models.exists(key):
-                raise ValueError(f"Unknown model: {key}")
-
-        transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
-        vae = self.vae_model.model_copy(update={"submodel_type": SubModelType.VAE})
-
-        tokenizer = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
-        clip_encoder = self.clip_embed_model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
-
-        tokenizer2 = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer2})
-        t5_encoder = self.t5_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder2})
-
-        transformer_config = context.models.get_config(transformer)
-        assert isinstance(transformer_config, CheckpointConfigBase)
-
-        return FluxModelLoaderOutput(
-            transformer=TransformerField(transformer=transformer, loras=[]),
-            clip=CLIPField(tokenizer=tokenizer, text_encoder=clip_encoder, loras=[], skipped_layers=0),
-            t5_encoder=T5EncoderField(tokenizer=tokenizer2, text_encoder=t5_encoder),
-            vae=VAEField(vae=vae),
-            max_seq_len=max_seq_lengths[transformer_config.config_path],
-        )
-
-
 @invocation(
     "main_model_loader",
     title="Main Model",

invokeai/app/invocations/primitives.py

@@ -12,7 +12,6 @@ from invokeai.app.invocations.fields import (
     ConditioningField,
     DenoiseMaskField,
     FieldDescriptions,
-    FluxConditioningField,
     ImageField,
     Input,
     InputField,
@@ -415,17 +414,6 @@ class MaskOutput(BaseInvocationOutput):
     height: int = OutputField(description="The height of the mask in pixels.")
 
 
-@invocation_output("flux_conditioning_output")
-class FluxConditioningOutput(BaseInvocationOutput):
-    """Base class for nodes that output a single conditioning tensor"""
-
-    conditioning: FluxConditioningField = OutputField(description=FieldDescriptions.cond)
-
-    @classmethod
-    def build(cls, conditioning_name: str) -> "FluxConditioningOutput":
-        return cls(conditioning=FluxConditioningField(conditioning_name=conditioning_name))
-
-
 @invocation_output("conditioning_output")
 class ConditioningOutput(BaseInvocationOutput):
     """Base class for nodes that output a single conditioning tensor"""

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

@@ -22,8 +22,8 @@ from invokeai.app.invocations.fields import (
 from invokeai.app.invocations.model import UNetField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
-from invokeai.backend.lora.lora_patcher import LoraPatcher
+from invokeai.backend.lora import LoRAModelRaw
+from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.stable_diffusion.diffusers_pipeline import ControlNetData, PipelineIntermediateState
 from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
     MultiDiffusionPipeline,
@@ -204,11 +204,7 @@ class TiledMultiDiffusionDenoiseLatents(BaseInvocation):
         # Load the UNet model.
         unet_info = context.models.load(self.unet.unet)
 
-        with (
-            ExitStack() as exit_stack,
-            unet_info as unet,
-            LoraPatcher.apply_lora_patches(model=unet, patches=_lora_loader(), prefix="lora_unet_"),
-        ):
+        with ExitStack() as exit_stack, unet_info as unet, ModelPatcher.apply_lora_unet(unet, _lora_loader()):
             assert isinstance(unet, UNet2DConditionModel)
             latents = latents.to(device=unet.device, dtype=unet.dtype)
             if noise is not None:

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,8 +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 queue item")
-    destination: str | None = Field(default=None, description="The destination of the queue item")
 
 
 class InvocationEventBase(QueueItemEventBase):
@@ -97,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")
@@ -114,8 +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,
-            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -149,8 +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,
-            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -188,8 +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,
-            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -222,8 +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,
-            destination=queue_item.destination,
             session_id=queue_item.session_id,
             invocation=invocation,
             invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
@@ -261,8 +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,
-            destination=queue_item.destination,
             session_id=queue_item.session_id,
             status=queue_item.status,
             error_type=queue_item.error_type,
@@ -289,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/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_install/model_install_common.py

@@ -103,7 +103,7 @@ class HFModelSource(StringLikeSource):
         if self.variant:
             base += f":{self.variant or ''}"
         if self.subfolder:
-            base += f"::{self.subfolder.as_posix()}"
+            base += f":{self.subfolder}"
         return base
 
 

invokeai/app/services/model_install/model_install_default.py

@@ -783,9 +783,8 @@ class ModelInstallService(ModelInstallServiceBase):
         # So what we do is to synthesize a folder named "sdxl-turbo_vae" here.
         if subfolder:
             top = Path(remote_files[0].path.parts[0])  # e.g. "sdxl-turbo/"
-            path_to_remove = top / subfolder  # sdxl-turbo/vae/
-            subfolder_rename = subfolder.name.replace("/", "_").replace("\\", "_")
-            path_to_add = Path(f"{top}_{subfolder_rename}")
+            path_to_remove = top / subfolder.parts[-1]  # sdxl-turbo/vae/
+            path_to_add = Path(f"{top}_{subfolder}")
         else:
             path_to_remove = Path(".")
             path_to_add = Path(".")

invokeai/app/services/model_records/model_records_base.py

@@ -77,7 +77,6 @@ class ModelRecordChanges(BaseModelExcludeNull):
     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)
-    format: Optional[str] = Field(description="format 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,
-    CancelByDestinationResult,
     CancelByQueueIDResult,
     ClearResult,
     EnqueueBatchResult,
@@ -96,11 +95,6 @@ class SessionQueueBase(ABC):
         """Cancels all queue items with matching batch IDs"""
         pass
 
-    @abstractmethod
-    def cancel_by_destination(self, queue_id: str, destination: str) -> CancelByDestinationResult:
-        """Cancels all queue items with the given batch destination"""
-        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,14 +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 queue item. This data is used by the frontend to determine how to handle results.",
-    )
-    destination: str | None = Field(
-        default=None,
-        description="The origin of this queue item. This data is used by the frontend to determine how to handle results",
-    )
     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(
@@ -203,14 +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. This data is used by the frontend to determine how to handle results.",
-    )
-    destination: str | None = Field(
-        default=None,
-        description="The origin of this queue item. This data is used by the frontend to determine how to handle results",
-    )
     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."
     )
@@ -310,8 +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")
-    destination: str | None = Field(..., description="The destination 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'")
@@ -346,12 +328,6 @@ class CancelByBatchIDsResult(BaseModel):
     canceled: int = Field(..., description="Number of queue items canceled")
 
 
-class CancelByDestinationResult(CancelByBatchIDsResult):
-    """Result of canceling by a destination"""
-
-    pass
-
-
 class CancelByQueueIDResult(CancelByBatchIDsResult):
     """Result of canceling by queue id"""
 
@@ -457,8 +433,6 @@ class SessionQueueValueToInsert(NamedTuple):
     field_values: Optional[str]  # field_values json
     priority: int  # priority
     workflow: Optional[str]  # workflow json
-    origin: str | None
-    destination: str | None
 
 
 ValuesToInsert: TypeAlias = list[SessionQueueValueToInsert]
@@ -479,8 +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
-                batch.destination,  # destination
             )
         )
     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,
-    CancelByDestinationResult,
     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, destination)
-                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_destination(self, queue_id: str, destination: str) -> CancelByDestinationResult:
-        try:
-            current_queue_item = self.get_current(queue_id)
-            self.__lock.acquire()
-            where = """--sql
-                WHERE
-                  queue_id == ?
-                  AND destination == ?
-                  AND status != 'canceled'
-                  AND status != 'completed'
-                  AND status != 'failed'
-                """
-            params = (queue_id, destination)
-            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.destination == destination:
-                self._set_queue_item_status(current_queue_item.item_id, "canceled")
-        except Exception:
-            self.__conn.rollback()
-            raise
-        finally:
-            self.__lock.release()
-        return CancelByDestinationResult(canceled=count)
-
     def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
         try:
             current_queue_item = self.get_current(queue_id)
@@ -578,9 +541,7 @@ class SqliteSessionQueue(SessionQueueBase):
                     started_at,
                     session_id,
                     batch_id,
-                    queue_id,
-                    origin,
-                    destination
+                    queue_id
                 FROM session_queue
                 WHERE queue_id = ?
             """
@@ -660,7 +621,7 @@ class SqliteSessionQueue(SessionQueueBase):
             self.__lock.acquire()
             self.__cursor.execute(
                 """--sql
-                SELECT status, count(*), origin, destination
+                SELECT status, count(*)
                 FROM session_queue
                 WHERE
                   queue_id = ?
@@ -672,8 +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
-            destination = result[0]["destination"] if result else None
         except Exception:
             self.__conn.rollback()
             raise
@@ -682,8 +641,6 @@ class SqliteSessionQueue(SessionQueueBase):
 
         return BatchStatus(
             batch_id=batch_id,
-            origin=origin,
-            destination=destination,
             queue_id=queue_id,
             pending=counts.get("pending", 0),
             in_progress=counts.get("in_progress", 0),

invokeai/app/services/shared/invocation_context.py

@@ -14,7 +14,7 @@ from invokeai.app.services.image_records.image_records_common import ImageCatego
 from invokeai.app.services.images.images_common import ImageDTO
 from invokeai.app.services.invocation_services import InvocationServices
 from invokeai.app.services.model_records.model_records_base import UnknownModelException
-from invokeai.app.util.step_callback import flux_step_callback, stable_diffusion_step_callback
+from invokeai.app.util.step_callback import stable_diffusion_step_callback
 from invokeai.backend.model_manager.config import (
     AnyModel,
     AnyModelConfig,
@@ -557,24 +557,6 @@ class UtilInterface(InvocationContextInterface):
             is_canceled=self.is_canceled,
         )
 
-    def flux_step_callback(self, intermediate_state: PipelineIntermediateState) -> None:
-        """
-        The step callback emits a progress event with the current step, the total number of
-        steps, a preview image, and some other internal metadata.
-
-        This should be called after each denoising step.
-
-        Args:
-            intermediate_state: The intermediate state of the diffusion pipeline.
-        """
-
-        flux_step_callback(
-            context_data=self._data,
-            intermediate_state=intermediate_state,
-            events=self._services.events,
-            is_canceled=self.is_canceled,
-        )
-
 
 class InvocationContext:
     """Provides access to various services and data for the current invocation.

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,34 +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.
-        - Adds `destination` column to the session queue table.
-        """
-
-        cursor.execute("ALTER TABLE session_queue ADD COLUMN origin TEXT;")
-        cursor.execute("ALTER TABLE session_queue ADD COLUMN destination 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.
-        - Adds `destination` 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/services/workflow_records/default_workflows/FLUX Image to Image.json

@@ -1,407 +0,0 @@
-{
-  "name": "FLUX Image to Image",
-  "author": "InvokeAI",
-  "description": "A simple image-to-image workflow using a FLUX dev model. ",
-  "version": "1.0.4",
-  "contact": "",
-  "tags": "image2image, flux, image-to-image",
-  "notes": "Prerequisite model downloads: T5 Encoder, CLIP-L Encoder, and FLUX VAE. Quantized and un-quantized versions can be found in the starter models tab within your Model Manager. We recommend using FLUX dev models for image-to-image workflows. The image-to-image performance with FLUX schnell models is poor.",
-  "exposedFields": [
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "model"
-    },
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "t5_encoder_model"
-    },
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "clip_embed_model"
-    },
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "vae_model"
-    },
-    {
-      "nodeId": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "fieldName": "denoising_start"
-    },
-    {
-      "nodeId": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "fieldName": "prompt"
-    },
-    {
-      "nodeId": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "fieldName": "num_steps"
-    }
-  ],
-  "meta": {
-    "version": "3.0.0",
-    "category": "default"
-  },
-  "nodes": [
-    {
-      "id": "2981a67c-480f-4237-9384-26b68dbf912b",
-      "type": "invocation",
-      "data": {
-        "id": "2981a67c-480f-4237-9384-26b68dbf912b",
-        "type": "flux_vae_encode",
-        "version": "1.0.0",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": true,
-        "inputs": {
-          "image": {
-            "name": "image",
-            "label": "",
-            "value": {
-              "image_name": "8a5c62aa-9335-45d2-9c71-89af9fc1f8d4.png"
-            }
-          },
-          "vae": {
-            "name": "vae",
-            "label": ""
-          }
-        }
-      },
-      "position": {
-        "x": 732.7680166609682,
-        "y": -24.37398171806909
-      }
-    },
-    {
-      "id": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "type": "invocation",
-      "data": {
-        "id": "ace0258f-67d7-4eee-a218-6fff27065214",
-        "type": "flux_denoise",
-        "version": "1.0.0",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": true,
-        "inputs": {
-          "board": {
-            "name": "board",
-            "label": ""
-          },
-          "metadata": {
-            "name": "metadata",
-            "label": ""
-          },
-          "latents": {
-            "name": "latents",
-            "label": ""
-          },
-          "denoise_mask": {
-            "name": "denoise_mask",
-            "label": ""
-          },
-          "denoising_start": {
-            "name": "denoising_start",
-            "label": "",
-            "value": 0.04
-          },
-          "denoising_end": {
-            "name": "denoising_end",
-            "label": "",
-            "value": 1
-          },
-          "transformer": {
-            "name": "transformer",
-            "label": ""
-          },
-          "positive_text_conditioning": {
-            "name": "positive_text_conditioning",
-            "label": ""
-          },
-          "width": {
-            "name": "width",
-            "label": "",
-            "value": 1024
-          },
-          "height": {
-            "name": "height",
-            "label": "",
-            "value": 1024
-          },
-          "num_steps": {
-            "name": "num_steps",
-            "label": "Steps (Recommend 30 for Dev, 4 for Schnell)",
-            "value": 30
-          },
-          "guidance": {
-            "name": "guidance",
-            "label": "",
-            "value": 4
-          },
-          "seed": {
-            "name": "seed",
-            "label": "",
-            "value": 0
-          }
-        }
-      },
-      "position": {
-        "x": 1182.8836633018684,
-        "y": -251.38882958913183
-      }
-    },
-    {
-      "id": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-      "type": "invocation",
-      "data": {
-        "id": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-        "type": "flux_vae_decode",
-        "version": "1.0.0",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": false,
-        "useCache": true,
-        "inputs": {
-          "board": {
-            "name": "board",
-            "label": ""
-          },
-          "metadata": {
-            "name": "metadata",
-            "label": ""
-          },
-          "latents": {
-            "name": "latents",
-            "label": ""
-          },
-          "vae": {
-            "name": "vae",
-            "label": ""
-          }
-        }
-      },
-      "position": {
-        "x": 1575.5797431839133,
-        "y": -209.00150975507415
-      }
-    },
-    {
-      "id": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "type": "invocation",
-      "data": {
-        "id": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-        "type": "flux_model_loader",
-        "version": "1.0.4",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": false,
-        "inputs": {
-          "model": {
-            "name": "model",
-            "label": "Model (dev variant recommended for Image-to-Image)"
-          },
-          "t5_encoder_model": {
-            "name": "t5_encoder_model",
-            "label": ""
-          },
-          "clip_embed_model": {
-            "name": "clip_embed_model",
-            "label": "",
-            "value": {
-              "key": "fa23a584-b623-415d-832a-21b5098ff1a1",
-              "hash": "blake3:17c19f0ef941c3b7609a9c94a659ca5364de0be364a91d4179f0e39ba17c3b70",
-              "name": "clip-vit-large-patch14",
-              "base": "any",
-              "type": "clip_embed"
-            }
-          },
-          "vae_model": {
-            "name": "vae_model",
-            "label": "",
-            "value": {
-              "key": "74fc82ba-c0a8-479d-a890-2126f82da758",
-              "hash": "blake3:ce21cb76364aa6e2421311cf4a4b5eb052a76c4f1cd207b50703d8978198a068",
-              "name": "FLUX.1-schnell_ae",
-              "base": "flux",
-              "type": "vae"
-            }
-          }
-        }
-      },
-      "position": {
-        "x": 328.1809894659957,
-        "y": -90.2241133566946
-      }
-    },
-    {
-      "id": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "type": "invocation",
-      "data": {
-        "id": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-        "type": "flux_text_encoder",
-        "version": "1.0.0",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": true,
-        "inputs": {
-          "clip": {
-            "name": "clip",
-            "label": ""
-          },
-          "t5_encoder": {
-            "name": "t5_encoder",
-            "label": ""
-          },
-          "t5_max_seq_len": {
-            "name": "t5_max_seq_len",
-            "label": "T5 Max Seq Len",
-            "value": 256
-          },
-          "prompt": {
-            "name": "prompt",
-            "label": "",
-            "value": "a cat wearing a birthday hat"
-          }
-        }
-      },
-      "position": {
-        "x": 745.8823365057267,
-        "y": -299.60249175851914
-      }
-    },
-    {
-      "id": "4754c534-a5f3-4ad0-9382-7887985e668c",
-      "type": "invocation",
-      "data": {
-        "id": "4754c534-a5f3-4ad0-9382-7887985e668c",
-        "type": "rand_int",
-        "version": "1.0.1",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": false,
-        "inputs": {
-          "low": {
-            "name": "low",
-            "label": "",
-            "value": 0
-          },
-          "high": {
-            "name": "high",
-            "label": "",
-            "value": 2147483647
-          }
-        }
-      },
-      "position": {
-        "x": 725.834098928012,
-        "y": 496.2710031089931
-      }
-    }
-  ],
-  "edges": [
-    {
-      "id": "reactflow__edge-2981a67c-480f-4237-9384-26b68dbf912bheight-ace0258f-67d7-4eee-a218-6fff27065214height",
-      "type": "default",
-      "source": "2981a67c-480f-4237-9384-26b68dbf912b",
-      "target": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "sourceHandle": "height",
-      "targetHandle": "height"
-    },
-    {
-      "id": "reactflow__edge-2981a67c-480f-4237-9384-26b68dbf912bwidth-ace0258f-67d7-4eee-a218-6fff27065214width",
-      "type": "default",
-      "source": "2981a67c-480f-4237-9384-26b68dbf912b",
-      "target": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "sourceHandle": "width",
-      "targetHandle": "width"
-    },
-    {
-      "id": "reactflow__edge-2981a67c-480f-4237-9384-26b68dbf912blatents-ace0258f-67d7-4eee-a218-6fff27065214latents",
-      "type": "default",
-      "source": "2981a67c-480f-4237-9384-26b68dbf912b",
-      "target": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "sourceHandle": "latents",
-      "targetHandle": "latents"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90vae-2981a67c-480f-4237-9384-26b68dbf912bvae",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "2981a67c-480f-4237-9384-26b68dbf912b",
-      "sourceHandle": "vae",
-      "targetHandle": "vae"
-    },
-    {
-      "id": "reactflow__edge-ace0258f-67d7-4eee-a218-6fff27065214latents-7e5172eb-48c1-44db-a770-8fd83e1435d1latents",
-      "type": "default",
-      "source": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "target": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-      "sourceHandle": "latents",
-      "targetHandle": "latents"
-    },
-    {
-      "id": "reactflow__edge-4754c534-a5f3-4ad0-9382-7887985e668cvalue-ace0258f-67d7-4eee-a218-6fff27065214seed",
-      "type": "default",
-      "source": "4754c534-a5f3-4ad0-9382-7887985e668c",
-      "target": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "sourceHandle": "value",
-      "targetHandle": "seed"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90transformer-ace0258f-67d7-4eee-a218-6fff27065214transformer",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "sourceHandle": "transformer",
-      "targetHandle": "transformer"
-    },
-    {
-      "id": "reactflow__edge-01f674f8-b3d1-4df1-acac-6cb8e0bfb63cconditioning-ace0258f-67d7-4eee-a218-6fff27065214positive_text_conditioning",
-      "type": "default",
-      "source": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "target": "ace0258f-67d7-4eee-a218-6fff27065214",
-      "sourceHandle": "conditioning",
-      "targetHandle": "positive_text_conditioning"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90vae-7e5172eb-48c1-44db-a770-8fd83e1435d1vae",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-      "sourceHandle": "vae",
-      "targetHandle": "vae"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90max_seq_len-01f674f8-b3d1-4df1-acac-6cb8e0bfb63ct5_max_seq_len",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "sourceHandle": "max_seq_len",
-      "targetHandle": "t5_max_seq_len"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90t5_encoder-01f674f8-b3d1-4df1-acac-6cb8e0bfb63ct5_encoder",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "sourceHandle": "t5_encoder",
-      "targetHandle": "t5_encoder"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90clip-01f674f8-b3d1-4df1-acac-6cb8e0bfb63cclip",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "sourceHandle": "clip",
-      "targetHandle": "clip"
-    }
-  ]
-}

invokeai/app/services/workflow_records/default_workflows/Flux Text to Image.json

@@ -1,326 +0,0 @@
-{
-  "name": "FLUX Text to Image",
-  "author": "InvokeAI",
-  "description": "A simple text-to-image workflow using FLUX dev or schnell models.",
-  "version": "1.0.4",
-  "contact": "",
-  "tags": "text2image, flux",
-  "notes": "Prerequisite model downloads: T5 Encoder, CLIP-L Encoder, and FLUX VAE. Quantized and un-quantized versions can be found in the starter models tab within your Model Manager. We recommend 4 steps for FLUX schnell models and 30 steps for FLUX dev models.",
-  "exposedFields": [
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "model"
-    },
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "t5_encoder_model"
-    },
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "clip_embed_model"
-    },
-    {
-      "nodeId": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "fieldName": "vae_model"
-    },
-    {
-      "nodeId": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "fieldName": "prompt"
-    },
-    {
-      "nodeId": "4fe24f07-f906-4f55-ab2c-9beee56ef5bd",
-      "fieldName": "num_steps"
-    }
-  ],
-  "meta": {
-    "version": "3.0.0",
-    "category": "default"
-  },
-  "nodes": [
-    {
-      "id": "4fe24f07-f906-4f55-ab2c-9beee56ef5bd",
-      "type": "invocation",
-      "data": {
-        "id": "4fe24f07-f906-4f55-ab2c-9beee56ef5bd",
-        "type": "flux_denoise",
-        "version": "1.0.0",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": true,
-        "inputs": {
-          "board": {
-            "name": "board",
-            "label": ""
-          },
-          "metadata": {
-            "name": "metadata",
-            "label": ""
-          },
-          "latents": {
-            "name": "latents",
-            "label": ""
-          },
-          "denoise_mask": {
-            "name": "denoise_mask",
-            "label": ""
-          },
-          "denoising_start": {
-            "name": "denoising_start",
-            "label": "",
-            "value": 0
-          },
-          "denoising_end": {
-            "name": "denoising_end",
-            "label": "",
-            "value": 1
-          },
-          "transformer": {
-            "name": "transformer",
-            "label": ""
-          },
-          "positive_text_conditioning": {
-            "name": "positive_text_conditioning",
-            "label": ""
-          },
-          "width": {
-            "name": "width",
-            "label": "",
-            "value": 1024
-          },
-          "height": {
-            "name": "height",
-            "label": "",
-            "value": 1024
-          },
-          "num_steps": {
-            "name": "num_steps",
-            "label": "Steps (Recommend 30 for Dev, 4 for Schnell)",
-            "value": 30
-          },
-          "guidance": {
-            "name": "guidance",
-            "label": "",
-            "value": 4
-          },
-          "seed": {
-            "name": "seed",
-            "label": "",
-            "value": 0
-          }
-        }
-      },
-      "position": {
-        "x": 1186.1868226120378,
-        "y": -214.9459927686657
-      }
-    },
-    {
-      "id": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-      "type": "invocation",
-      "data": {
-        "id": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-        "type": "flux_vae_decode",
-        "version": "1.0.0",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": false,
-        "useCache": true,
-        "inputs": {
-          "board": {
-            "name": "board",
-            "label": ""
-          },
-          "metadata": {
-            "name": "metadata",
-            "label": ""
-          },
-          "latents": {
-            "name": "latents",
-            "label": ""
-          },
-          "vae": {
-            "name": "vae",
-            "label": ""
-          }
-        }
-      },
-      "position": {
-        "x": 1575.5797431839133,
-        "y": -209.00150975507415
-      }
-    },
-    {
-      "id": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "type": "invocation",
-      "data": {
-        "id": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-        "type": "flux_model_loader",
-        "version": "1.0.4",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": false,
-        "inputs": {
-          "model": {
-            "name": "model",
-            "label": ""
-          },
-          "t5_encoder_model": {
-            "name": "t5_encoder_model",
-            "label": ""
-          },
-          "clip_embed_model": {
-            "name": "clip_embed_model",
-            "label": ""
-          },
-          "vae_model": {
-            "name": "vae_model",
-            "label": ""
-          }
-        }
-      },
-      "position": {
-        "x": 381.1882713063478,
-        "y": -95.89663532854017
-      }
-    },
-    {
-      "id": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "type": "invocation",
-      "data": {
-        "id": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-        "type": "flux_text_encoder",
-        "version": "1.0.0",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": true,
-        "inputs": {
-          "clip": {
-            "name": "clip",
-            "label": ""
-          },
-          "t5_encoder": {
-            "name": "t5_encoder",
-            "label": ""
-          },
-          "t5_max_seq_len": {
-            "name": "t5_max_seq_len",
-            "label": "T5 Max Seq Len",
-            "value": 256
-          },
-          "prompt": {
-            "name": "prompt",
-            "label": "",
-            "value": "a cat"
-          }
-        }
-      },
-      "position": {
-        "x": 778.4899149328337,
-        "y": -100.36469216659502
-      }
-    },
-    {
-      "id": "4754c534-a5f3-4ad0-9382-7887985e668c",
-      "type": "invocation",
-      "data": {
-        "id": "4754c534-a5f3-4ad0-9382-7887985e668c",
-        "type": "rand_int",
-        "version": "1.0.1",
-        "label": "",
-        "notes": "",
-        "isOpen": true,
-        "isIntermediate": true,
-        "useCache": false,
-        "inputs": {
-          "low": {
-            "name": "low",
-            "label": "",
-            "value": 0
-          },
-          "high": {
-            "name": "high",
-            "label": "",
-            "value": 2147483647
-          }
-        }
-      },
-      "position": {
-        "x": 800.9667463219505,
-        "y": 285.8297267547506
-      }
-    }
-  ],
-  "edges": [
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90transformer-4fe24f07-f906-4f55-ab2c-9beee56ef5bdtransformer",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "4fe24f07-f906-4f55-ab2c-9beee56ef5bd",
-      "sourceHandle": "transformer",
-      "targetHandle": "transformer"
-    },
-    {
-      "id": "reactflow__edge-01f674f8-b3d1-4df1-acac-6cb8e0bfb63cconditioning-4fe24f07-f906-4f55-ab2c-9beee56ef5bdpositive_text_conditioning",
-      "type": "default",
-      "source": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "target": "4fe24f07-f906-4f55-ab2c-9beee56ef5bd",
-      "sourceHandle": "conditioning",
-      "targetHandle": "positive_text_conditioning"
-    },
-    {
-      "id": "reactflow__edge-4754c534-a5f3-4ad0-9382-7887985e668cvalue-4fe24f07-f906-4f55-ab2c-9beee56ef5bdseed",
-      "type": "default",
-      "source": "4754c534-a5f3-4ad0-9382-7887985e668c",
-      "target": "4fe24f07-f906-4f55-ab2c-9beee56ef5bd",
-      "sourceHandle": "value",
-      "targetHandle": "seed"
-    },
-    {
-      "id": "reactflow__edge-4fe24f07-f906-4f55-ab2c-9beee56ef5bdlatents-7e5172eb-48c1-44db-a770-8fd83e1435d1latents",
-      "type": "default",
-      "source": "4fe24f07-f906-4f55-ab2c-9beee56ef5bd",
-      "target": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-      "sourceHandle": "latents",
-      "targetHandle": "latents"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90vae-7e5172eb-48c1-44db-a770-8fd83e1435d1vae",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "7e5172eb-48c1-44db-a770-8fd83e1435d1",
-      "sourceHandle": "vae",
-      "targetHandle": "vae"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90max_seq_len-01f674f8-b3d1-4df1-acac-6cb8e0bfb63ct5_max_seq_len",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "sourceHandle": "max_seq_len",
-      "targetHandle": "t5_max_seq_len"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90t5_encoder-01f674f8-b3d1-4df1-acac-6cb8e0bfb63ct5_encoder",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "sourceHandle": "t5_encoder",
-      "targetHandle": "t5_encoder"
-    },
-    {
-      "id": "reactflow__edge-f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90clip-01f674f8-b3d1-4df1-acac-6cb8e0bfb63cclip",
-      "type": "default",
-      "source": "f8d9d7c8-9ed7-4bd7-9e42-ab0e89bfac90",
-      "target": "01f674f8-b3d1-4df1-acac-6cb8e0bfb63c",
-      "sourceHandle": "clip",
-      "targetHandle": "clip"
-    }
-  ]
-}

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/app/util/step_callback.py

@@ -38,25 +38,6 @@ SD1_5_LATENT_RGB_FACTORS = [
     [-0.1307, -0.1874, -0.7445],  # L4
 ]
 
-FLUX_LATENT_RGB_FACTORS = [
-    [-0.0412, 0.0149, 0.0521],
-    [0.0056, 0.0291, 0.0768],
-    [0.0342, -0.0681, -0.0427],
-    [-0.0258, 0.0092, 0.0463],
-    [0.0863, 0.0784, 0.0547],
-    [-0.0017, 0.0402, 0.0158],
-    [0.0501, 0.1058, 0.1152],
-    [-0.0209, -0.0218, -0.0329],
-    [-0.0314, 0.0083, 0.0896],
-    [0.0851, 0.0665, -0.0472],
-    [-0.0534, 0.0238, -0.0024],
-    [0.0452, -0.0026, 0.0048],
-    [0.0892, 0.0831, 0.0881],
-    [-0.1117, -0.0304, -0.0789],
-    [0.0027, -0.0479, -0.0043],
-    [-0.1146, -0.0827, -0.0598],
-]
-
 
 def sample_to_lowres_estimated_image(
     samples: torch.Tensor, latent_rgb_factors: torch.Tensor, smooth_matrix: Optional[torch.Tensor] = None
@@ -113,32 +94,3 @@ def stable_diffusion_step_callback(
         intermediate_state,
         ProgressImage(dataURL=dataURL, width=width, height=height),
     )
-
-
-def flux_step_callback(
-    context_data: "InvocationContextData",
-    intermediate_state: PipelineIntermediateState,
-    events: "EventServiceBase",
-    is_canceled: Callable[[], bool],
-) -> None:
-    if is_canceled():
-        raise CanceledException
-    sample = intermediate_state.latents
-    latent_rgb_factors = torch.tensor(FLUX_LATENT_RGB_FACTORS, dtype=sample.dtype, device=sample.device)
-    latent_image_perm = sample.permute(1, 2, 0).to(dtype=sample.dtype, device=sample.device)
-    latent_image = latent_image_perm @ latent_rgb_factors
-    latents_ubyte = (
-        ((latent_image + 1) / 2).clamp(0, 1).mul(0xFF)  # change scale from -1..1 to 0..1  # to 0..255
-    ).to(device="cpu", dtype=torch.uint8)
-    image = Image.fromarray(latents_ubyte.cpu().numpy())
-    (width, height) = image.size
-    width *= 8
-    height *= 8
-    dataURL = image_to_dataURL(image, image_format="JPEG")
-
-    events.emit_invocation_denoise_progress(
-        context_data.queue_item,
-        context_data.invocation,
-        intermediate_state,
-        ProgressImage(dataURL=dataURL, width=width, height=height),
-    )

invokeai/backend/flux/denoise.py

@@ -1,56 +0,0 @@
-from typing import Callable
-
-import torch
-from tqdm import tqdm
-
-from invokeai.backend.flux.inpaint_extension import InpaintExtension
-from invokeai.backend.flux.model import Flux
-from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
-
-
-def denoise(
-    model: Flux,
-    # model input
-    img: torch.Tensor,
-    img_ids: torch.Tensor,
-    txt: torch.Tensor,
-    txt_ids: torch.Tensor,
-    vec: torch.Tensor,
-    # sampling parameters
-    timesteps: list[float],
-    step_callback: Callable[[PipelineIntermediateState], None],
-    guidance: float,
-    inpaint_extension: InpaintExtension | None,
-):
-    step = 0
-    # guidance_vec is ignored for schnell.
-    guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
-    for t_curr, t_prev in tqdm(list(zip(timesteps[:-1], timesteps[1:], strict=True))):
-        t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
-        pred = model(
-            img=img,
-            img_ids=img_ids,
-            txt=txt,
-            txt_ids=txt_ids,
-            y=vec,
-            timesteps=t_vec,
-            guidance=guidance_vec,
-        )
-        preview_img = img - t_curr * pred
-        img = img + (t_prev - t_curr) * pred
-
-        if inpaint_extension is not None:
-            img = inpaint_extension.merge_intermediate_latents_with_init_latents(img, t_prev)
-
-        step_callback(
-            PipelineIntermediateState(
-                step=step,
-                order=1,
-                total_steps=len(timesteps),
-                timestep=int(t_curr),
-                latents=preview_img,
-            ),
-        )
-        step += 1
-
-    return img

invokeai/backend/flux/inpaint_extension.py

@@ -1,35 +0,0 @@
-import torch
-
-
-class InpaintExtension:
-    """A class for managing inpainting with FLUX."""
-
-    def __init__(self, init_latents: torch.Tensor, inpaint_mask: torch.Tensor, noise: torch.Tensor):
-        """Initialize InpaintExtension.
-
-        Args:
-            init_latents (torch.Tensor): The initial latents (i.e. un-noised at timestep 0). In 'packed' format.
-            inpaint_mask (torch.Tensor): A mask specifying which elements to inpaint. Range [0, 1]. Values of 1 will be
-                re-generated. Values of 0 will remain unchanged. Values between 0 and 1 can be used to blend the
-                inpainted region with the background. In 'packed' format.
-            noise (torch.Tensor): The noise tensor used to noise the init_latents. In 'packed' format.
-        """
-        assert init_latents.shape == inpaint_mask.shape == noise.shape
-        self._init_latents = init_latents
-        self._inpaint_mask = inpaint_mask
-        self._noise = noise
-
-    def merge_intermediate_latents_with_init_latents(
-        self, intermediate_latents: torch.Tensor, timestep: float
-    ) -> torch.Tensor:
-        """Merge the intermediate latents with the initial latents for the current timestep using the inpaint mask. I.e.
-        update the intermediate latents to keep the regions that are not being inpainted on the correct noise
-        trajectory.
-
-        This function should be called after each denoising step.
-        """
-        # Noise the init latents for the current timestep.
-        noised_init_latents = self._noise * timestep + (1.0 - timestep) * self._init_latents
-
-        # Merge the intermediate latents with the noised_init_latents using the inpaint_mask.
-        return intermediate_latents * self._inpaint_mask + noised_init_latents * (1.0 - self._inpaint_mask)

invokeai/backend/flux/math.py

@@ -1,32 +0,0 @@
-# Initially pulled from https://github.com/black-forest-labs/flux
-
-import torch
-from einops import rearrange
-from torch import Tensor
-
-
-def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
-    q, k = apply_rope(q, k, pe)
-
-    x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
-    x = rearrange(x, "B H L D -> B L (H D)")
-
-    return x
-
-
-def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
-    assert dim % 2 == 0
-    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
-    omega = 1.0 / (theta**scale)
-    out = torch.einsum("...n,d->...nd", pos, omega)
-    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
-    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
-    return out.float()
-
-
-def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
-    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
-    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
-    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
-    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
-    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

invokeai/backend/flux/model.py

@@ -1,117 +0,0 @@
-# Initially pulled from https://github.com/black-forest-labs/flux
-
-from dataclasses import dataclass
-
-import torch
-from torch import Tensor, nn
-
-from invokeai.backend.flux.modules.layers import (
-    DoubleStreamBlock,
-    EmbedND,
-    LastLayer,
-    MLPEmbedder,
-    SingleStreamBlock,
-    timestep_embedding,
-)
-
-
-@dataclass
-class FluxParams:
-    in_channels: int
-    vec_in_dim: int
-    context_in_dim: int
-    hidden_size: int
-    mlp_ratio: float
-    num_heads: int
-    depth: int
-    depth_single_blocks: int
-    axes_dim: list[int]
-    theta: int
-    qkv_bias: bool
-    guidance_embed: bool
-
-
-class Flux(nn.Module):
-    """
-    Transformer model for flow matching on sequences.
-    """
-
-    def __init__(self, params: FluxParams):
-        super().__init__()
-
-        self.params = params
-        self.in_channels = params.in_channels
-        self.out_channels = self.in_channels
-        if params.hidden_size % params.num_heads != 0:
-            raise ValueError(f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}")
-        pe_dim = params.hidden_size // params.num_heads
-        if sum(params.axes_dim) != pe_dim:
-            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
-        self.hidden_size = params.hidden_size
-        self.num_heads = params.num_heads
-        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
-        self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
-        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
-        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
-        self.guidance_in = (
-            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
-        )
-        self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
-
-        self.double_blocks = nn.ModuleList(
-            [
-                DoubleStreamBlock(
-                    self.hidden_size,
-                    self.num_heads,
-                    mlp_ratio=params.mlp_ratio,
-                    qkv_bias=params.qkv_bias,
-                )
-                for _ in range(params.depth)
-            ]
-        )
-
-        self.single_blocks = nn.ModuleList(
-            [
-                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
-                for _ in range(params.depth_single_blocks)
-            ]
-        )
-
-        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
-
-    def forward(
-        self,
-        img: Tensor,
-        img_ids: Tensor,
-        txt: Tensor,
-        txt_ids: Tensor,
-        timesteps: Tensor,
-        y: Tensor,
-        guidance: Tensor | None = None,
-    ) -> Tensor:
-        if img.ndim != 3 or txt.ndim != 3:
-            raise ValueError("Input img and txt tensors must have 3 dimensions.")
-
-        # running on sequences img
-        img = self.img_in(img)
-        vec = self.time_in(timestep_embedding(timesteps, 256))
-        if self.params.guidance_embed:
-            if guidance is None:
-                raise ValueError("Didn't get guidance strength for guidance distilled model.")
-            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
-        vec = vec + self.vector_in(y)
-        txt = self.txt_in(txt)
-
-        ids = torch.cat((txt_ids, img_ids), dim=1)
-        pe = self.pe_embedder(ids)
-
-        for block in self.double_blocks:
-            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
-
-        img = torch.cat((txt, img), 1)
-        for block in self.single_blocks:
-            img = block(img, vec=vec, pe=pe)
-        img = img[:, txt.shape[1] :, ...]
-
-        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
-        return img

invokeai/backend/flux/modules/autoencoder.py

@@ -1,324 +0,0 @@
-# Initially pulled from https://github.com/black-forest-labs/flux
-
-from dataclasses import dataclass
-
-import torch
-from einops import rearrange
-from torch import Tensor, nn
-
-
-@dataclass
-class AutoEncoderParams:
-    resolution: int
-    in_channels: int
-    ch: int
-    out_ch: int
-    ch_mult: list[int]
-    num_res_blocks: int
-    z_channels: int
-    scale_factor: float
-    shift_factor: float
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
-
-    def attention(self, h_: Tensor) -> Tensor:
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        b, c, h, w = q.shape
-        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
-        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
-        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
-        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
-
-        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return x + self.proj_out(self.attention(x))
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-
-        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
-        if self.in_channels != self.out_channels:
-            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x):
-        h = x
-        h = self.norm1(h)
-        h = torch.nn.functional.silu(h)
-        h = self.conv1(h)
-
-        h = self.norm2(h)
-        h = torch.nn.functional.silu(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            x = self.nin_shortcut(x)
-
-        return x + h
-
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        # no asymmetric padding in torch conv, must do it ourselves
-        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
-
-    def forward(self, x: Tensor):
-        pad = (0, 1, 0, 1)
-        x = nn.functional.pad(x, pad, mode="constant", value=0)
-        x = self.conv(x)
-        return x
-
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels: int):
-        super().__init__()
-        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, x: Tensor):
-        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        x = self.conv(x)
-        return x
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        resolution: int,
-        in_channels: int,
-        ch: int,
-        ch_mult: list[int],
-        num_res_blocks: int,
-        z_channels: int,
-    ):
-        super().__init__()
-        self.ch = ch
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        # downsampling
-        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,) + tuple(ch_mult)
-        self.in_ch_mult = in_ch_mult
-        self.down = nn.ModuleList()
-        block_in = self.ch
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch * in_ch_mult[i_level]
-            block_out = ch * ch_mult[i_level]
-            for _ in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
-                block_in = block_out
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions - 1:
-                down.downsample = Downsample(block_in)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
-        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, x: Tensor) -> Tensor:
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1])
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions - 1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h)
-        # end
-        h = self.norm_out(h)
-        h = torch.nn.functional.silu(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        ch: int,
-        out_ch: int,
-        ch_mult: list[int],
-        num_res_blocks: int,
-        in_channels: int,
-        resolution: int,
-        z_channels: int,
-    ):
-        super().__init__()
-        self.ch = ch
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.ffactor = 2 ** (self.num_resolutions - 1)
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        block_in = ch * ch_mult[self.num_resolutions - 1]
-        curr_res = resolution // 2 ** (self.num_resolutions - 1)
-        self.z_shape = (1, z_channels, curr_res, curr_res)
-
-        # z to block_in
-        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch * ch_mult[i_level]
-            for _ in range(self.num_res_blocks + 1):
-                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
-                block_in = block_out
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in)
-                curr_res = curr_res * 2
-            self.up.insert(0, up)  # prepend to get consistent order
-
-        # end
-        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
-        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, z: Tensor) -> Tensor:
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                h = self.up[i_level].block[i_block](h)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = torch.nn.functional.silu(h)
-        h = self.conv_out(h)
-        return h
-
-
-class DiagonalGaussian(nn.Module):
-    def __init__(self, chunk_dim: int = 1):
-        super().__init__()
-        self.chunk_dim = chunk_dim
-
-    def forward(self, z: Tensor, sample: bool = True, generator: torch.Generator | None = None) -> Tensor:
-        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
-        if sample:
-            std = torch.exp(0.5 * logvar)
-            # Unfortunately, torch.randn_like(...) does not accept a generator argument at the time of writing, so we
-            # have to use torch.randn(...) instead.
-            return mean + std * torch.randn(size=mean.size(), generator=generator, dtype=mean.dtype, device=mean.device)
-        else:
-            return mean
-
-
-class AutoEncoder(nn.Module):
-    def __init__(self, params: AutoEncoderParams):
-        super().__init__()
-        self.encoder = Encoder(
-            resolution=params.resolution,
-            in_channels=params.in_channels,
-            ch=params.ch,
-            ch_mult=params.ch_mult,
-            num_res_blocks=params.num_res_blocks,
-            z_channels=params.z_channels,
-        )
-        self.decoder = Decoder(
-            resolution=params.resolution,
-            in_channels=params.in_channels,
-            ch=params.ch,
-            out_ch=params.out_ch,
-            ch_mult=params.ch_mult,
-            num_res_blocks=params.num_res_blocks,
-            z_channels=params.z_channels,
-        )
-        self.reg = DiagonalGaussian()
-
-        self.scale_factor = params.scale_factor
-        self.shift_factor = params.shift_factor
-
-    def encode(self, x: Tensor, sample: bool = True, generator: torch.Generator | None = None) -> Tensor:
-        """Run VAE encoding on input tensor x.
-
-        Args:
-            x (Tensor): Input image tensor. Shape: (batch_size, in_channels, height, width).
-            sample (bool, optional): If True, sample from the encoded distribution, else, return the distribution mean.
-                Defaults to True.
-            generator (torch.Generator | None, optional): Optional random number generator for reproducibility.
-                Defaults to None.
-
-        Returns:
-            Tensor: Encoded latent tensor. Shape: (batch_size, z_channels, latent_height, latent_width).
-        """
-
-        z = self.reg(self.encoder(x), sample=sample, generator=generator)
-        z = self.scale_factor * (z - self.shift_factor)
-        return z
-
-    def decode(self, z: Tensor) -> Tensor:
-        z = z / self.scale_factor + self.shift_factor
-        return self.decoder(z)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.decode(self.encode(x))

invokeai/backend/flux/modules/conditioner.py

@@ -1,33 +0,0 @@
-# Initially pulled from https://github.com/black-forest-labs/flux
-
-from torch import Tensor, nn
-from transformers import PreTrainedModel, PreTrainedTokenizer
-
-
-class HFEncoder(nn.Module):
-    def __init__(self, encoder: PreTrainedModel, tokenizer: PreTrainedTokenizer, is_clip: bool, max_length: int):
-        super().__init__()
-        self.max_length = max_length
-        self.is_clip = is_clip
-        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
-        self.tokenizer = tokenizer
-        self.hf_module = encoder
-        self.hf_module = self.hf_module.eval().requires_grad_(False)
-
-    def forward(self, text: list[str]) -> Tensor:
-        batch_encoding = self.tokenizer(
-            text,
-            truncation=True,
-            max_length=self.max_length,
-            return_length=False,
-            return_overflowing_tokens=False,
-            padding="max_length",
-            return_tensors="pt",
-        )
-
-        outputs = self.hf_module(
-            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
-            attention_mask=None,
-            output_hidden_states=False,
-        )
-        return outputs[self.output_key]

invokeai/backend/flux/modules/layers.py

@@ -1,253 +0,0 @@
-# Initially pulled from https://github.com/black-forest-labs/flux
-
-import math
-from dataclasses import dataclass
-
-import torch
-from einops import rearrange
-from torch import Tensor, nn
-
-from invokeai.backend.flux.math import attention, rope
-
-
-class EmbedND(nn.Module):
-    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
-        super().__init__()
-        self.dim = dim
-        self.theta = theta
-        self.axes_dim = axes_dim
-
-    def forward(self, ids: Tensor) -> Tensor:
-        n_axes = ids.shape[-1]
-        emb = torch.cat(
-            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
-            dim=-3,
-        )
-
-        return emb.unsqueeze(1)
-
-
-def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
-    """
-    Create sinusoidal timestep embeddings.
-    :param t: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an (N, D) Tensor of positional embeddings.
-    """
-    t = time_factor * t
-    half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(t.device)
-
-    args = t[:, None].float() * freqs[None]
-    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-    if dim % 2:
-        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    if torch.is_floating_point(t):
-        embedding = embedding.to(t)
-    return embedding
-
-
-class MLPEmbedder(nn.Module):
-    def __init__(self, in_dim: int, hidden_dim: int):
-        super().__init__()
-        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
-        self.silu = nn.SiLU()
-        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.out_layer(self.silu(self.in_layer(x)))
-
-
-class RMSNorm(torch.nn.Module):
-    def __init__(self, dim: int):
-        super().__init__()
-        self.scale = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x: Tensor):
-        x_dtype = x.dtype
-        x = x.float()
-        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
-        return (x * rrms).to(dtype=x_dtype) * self.scale
-
-
-class QKNorm(torch.nn.Module):
-    def __init__(self, dim: int):
-        super().__init__()
-        self.query_norm = RMSNorm(dim)
-        self.key_norm = RMSNorm(dim)
-
-    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
-        q = self.query_norm(q)
-        k = self.key_norm(k)
-        return q.to(v), k.to(v)
-
-
-class SelfAttention(nn.Module):
-    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.norm = QKNorm(head_dim)
-        self.proj = nn.Linear(dim, dim)
-
-    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
-        qkv = self.qkv(x)
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        q, k = self.norm(q, k, v)
-        x = attention(q, k, v, pe=pe)
-        x = self.proj(x)
-        return x
-
-
-@dataclass
-class ModulationOut:
-    shift: Tensor
-    scale: Tensor
-    gate: Tensor
-
-
-class Modulation(nn.Module):
-    def __init__(self, dim: int, double: bool):
-        super().__init__()
-        self.is_double = double
-        self.multiplier = 6 if double else 3
-        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
-
-    def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
-        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
-
-        return (
-            ModulationOut(*out[:3]),
-            ModulationOut(*out[3:]) if self.is_double else None,
-        )
-
-
-class DoubleStreamBlock(nn.Module):
-    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
-        super().__init__()
-
-        mlp_hidden_dim = int(hidden_size * mlp_ratio)
-        self.num_heads = num_heads
-        self.hidden_size = hidden_size
-        self.img_mod = Modulation(hidden_size, double=True)
-        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
-
-        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.img_mlp = nn.Sequential(
-            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
-            nn.GELU(approximate="tanh"),
-            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
-        )
-
-        self.txt_mod = Modulation(hidden_size, double=True)
-        self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
-
-        self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.txt_mlp = nn.Sequential(
-            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
-            nn.GELU(approximate="tanh"),
-            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
-        )
-
-    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
-        img_mod1, img_mod2 = self.img_mod(vec)
-        txt_mod1, txt_mod2 = self.txt_mod(vec)
-
-        # prepare image for attention
-        img_modulated = self.img_norm1(img)
-        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
-        img_qkv = self.img_attn.qkv(img_modulated)
-        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
-
-        # prepare txt for attention
-        txt_modulated = self.txt_norm1(txt)
-        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
-        txt_qkv = self.txt_attn.qkv(txt_modulated)
-        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
-
-        # run actual attention
-        q = torch.cat((txt_q, img_q), dim=2)
-        k = torch.cat((txt_k, img_k), dim=2)
-        v = torch.cat((txt_v, img_v), dim=2)
-
-        attn = attention(q, k, v, pe=pe)
-        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
-
-        # calculate the img bloks
-        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
-        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
-
-        # calculate the txt bloks
-        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
-        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
-        return img, txt
-
-
-class SingleStreamBlock(nn.Module):
-    """
-    A DiT block with parallel linear layers as described in
-    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
-    """
-
-    def __init__(
-        self,
-        hidden_size: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        qk_scale: float | None = None,
-    ):
-        super().__init__()
-        self.hidden_dim = hidden_size
-        self.num_heads = num_heads
-        head_dim = hidden_size // num_heads
-        self.scale = qk_scale or head_dim**-0.5
-
-        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
-        # qkv and mlp_in
-        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
-        # proj and mlp_out
-        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
-
-        self.norm = QKNorm(head_dim)
-
-        self.hidden_size = hidden_size
-        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-
-        self.mlp_act = nn.GELU(approximate="tanh")
-        self.modulation = Modulation(hidden_size, double=False)
-
-    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
-        mod, _ = self.modulation(vec)
-        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
-        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
-
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        q, k = self.norm(q, k, v)
-
-        # compute attention
-        attn = attention(q, k, v, pe=pe)
-        # compute activation in mlp stream, cat again and run second linear layer
-        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
-        return x + mod.gate * output
-
-
-class LastLayer(nn.Module):
-    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
-        super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
-
-    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
-        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
-        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
-        x = self.linear(x)
-        return x

invokeai/backend/flux/sampling_utils.py

@@ -1,135 +0,0 @@
-# Initially pulled from https://github.com/black-forest-labs/flux
-
-import math
-from typing import Callable
-
-import torch
-from einops import rearrange, repeat
-
-
-def get_noise(
-    num_samples: int,
-    height: int,
-    width: int,
-    device: torch.device,
-    dtype: torch.dtype,
-    seed: int,
-):
-    # We always generate noise on the same device and dtype then cast to ensure consistency across devices/dtypes.
-    rand_device = "cpu"
-    rand_dtype = torch.float16
-    return torch.randn(
-        num_samples,
-        16,
-        # allow for packing
-        2 * math.ceil(height / 16),
-        2 * math.ceil(width / 16),
-        device=rand_device,
-        dtype=rand_dtype,
-        generator=torch.Generator(device=rand_device).manual_seed(seed),
-    ).to(device=device, dtype=dtype)
-
-
-def time_shift(mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
-    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
-
-
-def get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15) -> Callable[[float], float]:
-    m = (y2 - y1) / (x2 - x1)
-    b = y1 - m * x1
-    return lambda x: m * x + b
-
-
-def get_schedule(
-    num_steps: int,
-    image_seq_len: int,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-    shift: bool = True,
-) -> list[float]:
-    # extra step for zero
-    timesteps = torch.linspace(1, 0, num_steps + 1)
-
-    # shifting the schedule to favor high timesteps for higher signal images
-    if shift:
-        # estimate mu based on linear estimation between two points
-        mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
-        timesteps = time_shift(mu, 1.0, timesteps)
-
-    return timesteps.tolist()
-
-
-def _find_last_index_ge_val(timesteps: list[float], val: float, eps: float = 1e-6) -> int:
-    """Find the last index in timesteps that is >= val.
-
-    We use epsilon-close equality to avoid potential floating point errors.
-    """
-    idx = len(list(filter(lambda t: t >= (val - eps), timesteps))) - 1
-    assert idx >= 0
-    return idx
-
-
-def clip_timestep_schedule(timesteps: list[float], denoising_start: float, denoising_end: float) -> list[float]:
-    """Clip the timestep schedule to the denoising range.
-
-    Args:
-        timesteps (list[float]): The original timestep schedule: [1.0, ..., 0.0].
-        denoising_start (float): A value in [0, 1] specifying the start of the denoising process. E.g. a value of 0.2
-            would mean that the denoising process start at the last timestep in the schedule >= 0.8.
-        denoising_end (float): A value in [0, 1] specifying the end of the denoising process. E.g. a value of 0.8 would
-            mean that the denoising process end at the last timestep in the schedule >= 0.2.
-
-    Returns:
-        list[float]: The clipped timestep schedule.
-    """
-    assert 0.0 <= denoising_start <= 1.0
-    assert 0.0 <= denoising_end <= 1.0
-    assert denoising_start <= denoising_end
-
-    t_start_val = 1.0 - denoising_start
-    t_end_val = 1.0 - denoising_end
-
-    t_start_idx = _find_last_index_ge_val(timesteps, t_start_val)
-    t_end_idx = _find_last_index_ge_val(timesteps, t_end_val)
-
-    clipped_timesteps = timesteps[t_start_idx : t_end_idx + 1]
-
-    return clipped_timesteps
-
-
-def unpack(x: torch.Tensor, height: int, width: int) -> torch.Tensor:
-    """Unpack flat array of patch embeddings to latent image."""
-    return rearrange(
-        x,
-        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
-        h=math.ceil(height / 16),
-        w=math.ceil(width / 16),
-        ph=2,
-        pw=2,
-    )
-
-
-def pack(x: torch.Tensor) -> torch.Tensor:
-    """Pack latent image to flattented array of patch embeddings."""
-    # Pixel unshuffle with a scale of 2, and flatten the height/width dimensions to get an array of patches.
-    return rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-
-
-def generate_img_ids(h: int, w: int, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-    """Generate tensor of image position ids.
-
-    Args:
-        h (int): Height of image in latent space.
-        w (int): Width of image in latent space.
-        batch_size (int): Batch size.
-        device (torch.device): Device.
-        dtype (torch.dtype): dtype.
-
-    Returns:
-        torch.Tensor: Image position ids.
-    """
-    img_ids = torch.zeros(h // 2, w // 2, 3, device=device, dtype=dtype)
-    img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2, device=device, dtype=dtype)[:, None]
-    img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2, device=device, dtype=dtype)[None, :]
-    img_ids = repeat(img_ids, "h w c -> b (h w) c", b=batch_size)
-    return img_ids

invokeai/backend/flux/util.py

@@ -1,71 +0,0 @@
-# Initially pulled from https://github.com/black-forest-labs/flux
-
-from dataclasses import dataclass
-from typing import Dict, Literal
-
-from invokeai.backend.flux.model import FluxParams
-from invokeai.backend.flux.modules.autoencoder import AutoEncoderParams
-
-
-@dataclass
-class ModelSpec:
-    params: FluxParams
-    ae_params: AutoEncoderParams
-    ckpt_path: str | None
-    ae_path: str | None
-    repo_id: str | None
-    repo_flow: str | None
-    repo_ae: str | None
-
-
-max_seq_lengths: Dict[str, Literal[256, 512]] = {
-    "flux-dev": 512,
-    "flux-schnell": 256,
-}
-
-
-ae_params = {
-    "flux": AutoEncoderParams(
-        resolution=256,
-        in_channels=3,
-        ch=128,
-        out_ch=3,
-        ch_mult=[1, 2, 4, 4],
-        num_res_blocks=2,
-        z_channels=16,
-        scale_factor=0.3611,
-        shift_factor=0.1159,
-    )
-}
-
-
-params = {
-    "flux-dev": FluxParams(
-        in_channels=64,
-        vec_in_dim=768,
-        context_in_dim=4096,
-        hidden_size=3072,
-        mlp_ratio=4.0,
-        num_heads=24,
-        depth=19,
-        depth_single_blocks=38,
-        axes_dim=[16, 56, 56],
-        theta=10_000,
-        qkv_bias=True,
-        guidance_embed=True,
-    ),
-    "flux-schnell": FluxParams(
-        in_channels=64,
-        vec_in_dim=768,
-        context_in_dim=4096,
-        hidden_size=3072,
-        mlp_ratio=4.0,
-        num_heads=24,
-        depth=19,
-        depth_single_blocks=38,
-        axes_dim=[16, 56, 56],
-        theta=10_000,
-        qkv_bias=True,
-        guidance_embed=False,
-    ),
-}

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/load_flux_model.py

@@ -0,0 +1,129 @@
+import json
+import os
+import time
+from pathlib import Path
+from typing import Union
+
+import torch
+from diffusers.models.model_loading_utils import load_state_dict
+from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
+from diffusers.utils import (
+    CONFIG_NAME,
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFETENSORS_WEIGHTS_NAME,
+    _get_checkpoint_shard_files,
+    is_accelerate_available,
+)
+from optimum.quanto import qfloat8
+from optimum.quanto.models import QuantizedDiffusersModel
+from optimum.quanto.models.shared_dict import ShardedStateDict
+
+from invokeai.backend.requantize import requantize
+
+
+class QuantizedFluxTransformer2DModel(QuantizedDiffusersModel):
+    base_class = FluxTransformer2DModel
+
+    @classmethod
+    def from_pretrained(cls, model_name_or_path: Union[str, os.PathLike]):
+        if cls.base_class is None:
+            raise ValueError("The `base_class` attribute needs to be configured.")
+
+        if not is_accelerate_available():
+            raise ValueError("Reloading a quantized diffusers model requires the accelerate library.")
+        from accelerate import init_empty_weights
+
+        if os.path.isdir(model_name_or_path):
+            # Look for a quantization map
+            qmap_path = os.path.join(model_name_or_path, cls._qmap_name())
+            if not os.path.exists(qmap_path):
+                raise ValueError(f"No quantization map found in {model_name_or_path}: is this a quantized model ?")
+
+            # Look for original model config file.
+            model_config_path = os.path.join(model_name_or_path, CONFIG_NAME)
+            if not os.path.exists(model_config_path):
+                raise ValueError(f"{CONFIG_NAME} not found in {model_name_or_path}.")
+
+            with open(qmap_path, "r", encoding="utf-8") as f:
+                qmap = json.load(f)
+
+            with open(model_config_path, "r", encoding="utf-8") as f:
+                original_model_cls_name = json.load(f)["_class_name"]
+            configured_cls_name = cls.base_class.__name__
+            if configured_cls_name != original_model_cls_name:
+                raise ValueError(
+                    f"Configured base class ({configured_cls_name}) differs from what was derived from the provided configuration ({original_model_cls_name})."
+                )
+
+            # Create an empty model
+            config = cls.base_class.load_config(model_name_or_path)
+            with init_empty_weights():
+                model = cls.base_class.from_config(config)
+
+            # Look for the index of a sharded checkpoint
+            checkpoint_file = os.path.join(model_name_or_path, SAFE_WEIGHTS_INDEX_NAME)
+            if os.path.exists(checkpoint_file):
+                # Convert the checkpoint path to a list of shards
+                _, sharded_metadata = _get_checkpoint_shard_files(model_name_or_path, checkpoint_file)
+                # Create a mapping for the sharded safetensor files
+                state_dict = ShardedStateDict(model_name_or_path, sharded_metadata["weight_map"])
+            else:
+                # Look for a single checkpoint file
+                checkpoint_file = os.path.join(model_name_or_path, SAFETENSORS_WEIGHTS_NAME)
+                if not os.path.exists(checkpoint_file):
+                    raise ValueError(f"No safetensor weights found in {model_name_or_path}.")
+                # Get state_dict from model checkpoint
+                state_dict = load_state_dict(checkpoint_file)
+
+            # Requantize and load quantized weights from state_dict
+            requantize(model, state_dict=state_dict, quantization_map=qmap)
+            model.eval()
+            return cls(model)
+        else:
+            raise NotImplementedError("Reloading quantized models directly from the hub is not supported yet.")
+
+
+def load_flux_transformer(path: Path) -> FluxTransformer2DModel:
+    # model = FluxTransformer2DModel.from_pretrained(path, local_files_only=True, torch_dtype=torch.bfloat16)
+    model_8bit_path = path / "quantized"
+    if model_8bit_path.exists():
+        # The quantized model exists, load it.
+        # TODO(ryand): The requantize(...) operation in from_pretrained(...) is very slow. This seems like
+        # something that we should be able to make much faster.
+        q_model = QuantizedFluxTransformer2DModel.from_pretrained(model_8bit_path)
+
+        # Access the underlying wrapped model.
+        # We access the wrapped model, even though it is private, because it simplifies the type checking by
+        # always returning a FluxTransformer2DModel from this function.
+        model = q_model._wrapped
+    else:
+        # The quantized model does not exist yet, quantize and save it.
+        # TODO(ryand): Loading in float16 and then quantizing seems to result in NaNs. In order to run this on
+        # GPUs that don't support bfloat16, we would need to host the quantized model instead of generating it
+        # here.
+        model = FluxTransformer2DModel.from_pretrained(path, local_files_only=True, torch_dtype=torch.bfloat16)
+        assert isinstance(model, FluxTransformer2DModel)
+
+        q_model = QuantizedFluxTransformer2DModel.quantize(model, weights=qfloat8)
+
+        model_8bit_path.mkdir(parents=True, exist_ok=True)
+        q_model.save_pretrained(model_8bit_path)
+
+        # (See earlier comment about accessing the wrapped model.)
+        model = q_model._wrapped
+
+    assert isinstance(model, FluxTransformer2DModel)
+    return model
+
+
+def main():
+    start = time.time()
+    model = load_flux_transformer(
+        Path("/data/invokeai/models/.download_cache/black-forest-labs_flux.1-schnell/FLUX.1-schnell/transformer/")
+    )
+    print(f"Time to load: {time.time() - start}s")
+    print("hi")
+
+
+if __name__ == "__main__":
+    main()

invokeai/backend/lora.py

@@ -0,0 +1,630 @@
+# Copyright (c) 2024 The InvokeAI Development team
+"""LoRA model support."""
+
+import bisect
+from pathlib import Path
+from typing import Dict, List, Optional, Set, 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
+
+
+class LoRALayerBase:
+    # rank: Optional[int]
+    # alpha: Optional[float]
+    # bias: Optional[torch.Tensor]
+    # layer_key: str
+
+    # @property
+    # def scale(self):
+    #    return self.alpha / self.rank if (self.alpha and self.rank) else 1.0
+
+    def __init__(
+        self,
+        layer_key: str,
+        values: Dict[str, torch.Tensor],
+    ):
+        if "alpha" in values:
+            self.alpha = values["alpha"].item()
+        else:
+            self.alpha = None
+
+        if "bias_indices" in values and "bias_values" in values and "bias_size" in values:
+            self.bias: Optional[torch.Tensor] = torch.sparse_coo_tensor(
+                values["bias_indices"],
+                values["bias_values"],
+                tuple(values["bias_size"]),
+            )
+
+        else:
+            self.bias = None
+
+        self.rank = None  # set in layer implementation
+        self.layer_key = layer_key
+
+    def get_weight(self, orig_weight: 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]:
+            if val is not None:
+                model_size += val.nelement() * val.element_size()
+        return model_size
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> 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}"
+            )
+
+
+# TODO: find and debug lora/locon with bias
+class LoRALayer(LoRALayerBase):
+    # up: torch.Tensor
+    # mid: Optional[torch.Tensor]
+    # down: torch.Tensor
+
+    def __init__(
+        self,
+        layer_key: str,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(layer_key, values)
+
+        self.up = values["lora_up.weight"]
+        self.down = values["lora_down.weight"]
+        self.mid = values.get("lora_mid.weight", 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:
+        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])
+            weight = torch.einsum("m n w h, i m, n j -> i j w h", self.mid, up, down)
+        else:
+            weight = self.up.reshape(self.up.shape[0], -1) @ self.down.reshape(self.down.shape[0], -1)
+
+        return weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        for val in [self.up, self.mid, self.down]:
+            if val is not None:
+                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)
+
+        self.up = self.up.to(device=device, dtype=dtype)
+        self.down = self.down.to(device=device, dtype=dtype)
+
+        if self.mid is not None:
+            self.mid = self.mid.to(device=device, dtype=dtype)
+
+
+class LoHALayer(LoRALayerBase):
+    # w1_a: torch.Tensor
+    # w1_b: torch.Tensor
+    # w2_a: torch.Tensor
+    # w2_b: torch.Tensor
+    # t1: Optional[torch.Tensor] = None
+    # t2: Optional[torch.Tensor] = None
+
+    def __init__(self, layer_key: str, values: Dict[str, torch.Tensor]):
+        super().__init__(layer_key, values)
+
+        self.w1_a = values["hada_w1_a"]
+        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)
+
+        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:
+        if self.t1 is None:
+            weight: torch.Tensor = (self.w1_a @ self.w1_b) * (self.w2_a @ self.w2_b)
+
+        else:
+            rebuild1 = torch.einsum("i j k l, j r, i p -> p r k l", self.t1, self.w1_b, self.w1_a)
+            rebuild2 = torch.einsum("i j k l, j r, i p -> p r k l", self.t2, self.w2_b, self.w2_a)
+            weight = rebuild1 * rebuild2
+
+        return weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        for val in [self.w1_a, self.w1_b, self.w2_a, self.w2_b, self.t1, self.t2]:
+            if val is not None:
+                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)
+
+        self.w1_a = self.w1_a.to(device=device, dtype=dtype)
+        self.w1_b = self.w1_b.to(device=device, dtype=dtype)
+        if self.t1 is not None:
+            self.t1 = self.t1.to(device=device, dtype=dtype)
+
+        self.w2_a = self.w2_a.to(device=device, dtype=dtype)
+        self.w2_b = self.w2_b.to(device=device, dtype=dtype)
+        if self.t2 is not None:
+            self.t2 = self.t2.to(device=device, dtype=dtype)
+
+
+class LoKRLayer(LoRALayerBase):
+    # w1: Optional[torch.Tensor] = None
+    # w1_a: Optional[torch.Tensor] = None
+    # w1_b: Optional[torch.Tensor] = None
+    # w2: Optional[torch.Tensor] = None
+    # w2_a: Optional[torch.Tensor] = None
+    # w2_b: Optional[torch.Tensor] = None
+    # t2: Optional[torch.Tensor] = None
+
+    def __init__(
+        self,
+        layer_key: str,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(layer_key, values)
+
+        self.w1 = values.get("lokr_w1", None)
+        if self.w1 is None:
+            self.w1_a = values["lokr_w1_a"]
+            self.w1_b = values["lokr_w1_b"]
+
+        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"]
+
+        self.t2 = values.get("lokr_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]
+        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:
+        w1: Optional[torch.Tensor] = self.w1
+        if w1 is None:
+            assert self.w1_a is not None
+            assert self.w1_b is not None
+            w1 = self.w1_a @ self.w1_b
+
+        w2 = self.w2
+        if w2 is None:
+            if self.t2 is None:
+                assert self.w2_a is not None
+                assert self.w2_b is not None
+                w2 = self.w2_a @ self.w2_b
+            else:
+                w2 = torch.einsum("i j k l, i p, j r -> p r k l", self.t2, self.w2_a, self.w2_b)
+
+        if len(w2.shape) == 4:
+            w1 = w1.unsqueeze(2).unsqueeze(2)
+        w2 = w2.contiguous()
+        assert w1 is not None
+        assert w2 is not None
+        weight = torch.kron(w1, w2)
+
+        return weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        for val in [self.w1, self.w1_a, self.w1_b, self.w2, self.w2_a, self.w2_b, self.t2]:
+            if val is not None:
+                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)
+
+        if self.w1 is not None:
+            self.w1 = self.w1.to(device=device, dtype=dtype)
+        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)
+
+        if self.w2 is not None:
+            self.w2 = self.w2.to(device=device, dtype=dtype)
+        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)
+
+        if self.t2 is not None:
+            self.t2 = self.t2.to(device=device, dtype=dtype)
+
+
+class FullLayer(LoRALayerBase):
+    # bias handled in LoRALayerBase(calc_size, to)
+    # weight: torch.Tensor
+    # bias: Optional[torch.Tensor]
+
+    def __init__(
+        self,
+        layer_key: str,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(layer_key, values)
+
+        self.weight = values["diff"]
+        self.bias = values.get("diff_b", None)
+
+        self.rank = None  # unscaled
+        self.check_keys(values, {"diff", "diff_b"})
+
+    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)
+
+
+class IA3Layer(LoRALayerBase):
+    # weight: torch.Tensor
+    # on_input: torch.Tensor
+
+    def __init__(
+        self,
+        layer_key: str,
+        values: Dict[str, torch.Tensor],
+    ):
+        super().__init__(layer_key, values)
+
+        self.weight = values["weight"]
+        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:
+        weight = self.weight
+        if not self.on_input:
+            weight = weight.reshape(-1, 1)
+        assert orig_weight is not None
+        return orig_weight * weight
+
+    def calc_size(self) -> int:
+        model_size = super().calc_size()
+        model_size += self.weight.nelement() * self.weight.element_size()
+        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)
+
+
+AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer]
+
+
+class LoRAModelRaw(RawModel):  # (torch.nn.Module):
+    _name: str
+    layers: Dict[str, AnyLoRALayer]
+
+    def __init__(
+        self,
+        name: str,
+        layers: Dict[str, AnyLoRALayer],
+    ):
+        self._name = name
+        self.layers = layers
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
+        # TODO: try revert if exception?
+        for _key, layer in self.layers.items():
+            layer.to(device=device, dtype=dtype)
+
+    def calc_size(self) -> int:
+        model_size = 0
+        for _, layer in self.layers.items():
+            model_size += layer.calc_size()
+        return model_size
+
+    @classmethod
+    def _convert_sdxl_keys_to_diffusers_format(cls, state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """Convert the keys of an SDXL LoRA state_dict to diffusers format.
+
+        The input state_dict can be in either Stability AI format or diffusers format. If the state_dict is already in
+        diffusers format, then this function will have no effect.
+
+        This function is adapted from:
+        https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L385-L409
+
+        Args:
+            state_dict (Dict[str, Tensor]): The SDXL LoRA state_dict.
+
+        Raises:
+            ValueError: If state_dict contains an unrecognized key, or not all keys could be converted.
+
+        Returns:
+            Dict[str, Tensor]: The diffusers-format state_dict.
+        """
+        converted_count = 0  # The number of Stability AI keys converted to diffusers format.
+        not_converted_count = 0  # The number of keys that were not converted.
+
+        # Get a sorted list of Stability AI UNet keys so that we can efficiently search for keys with matching prefixes.
+        # For example, we want to efficiently find `input_blocks_4_1` in the list when searching for
+        # `input_blocks_4_1_proj_in`.
+        stability_unet_keys = list(SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP)
+        stability_unet_keys.sort()
+
+        new_state_dict = {}
+        for full_key, value in state_dict.items():
+            if full_key.startswith("lora_unet_"):
+                search_key = full_key.replace("lora_unet_", "")
+                # Use bisect to find the key in stability_unet_keys that *may* match the search_key's prefix.
+                position = bisect.bisect_right(stability_unet_keys, search_key)
+                map_key = stability_unet_keys[position - 1]
+                # Now, check if the map_key *actually* matches the search_key.
+                if search_key.startswith(map_key):
+                    new_key = full_key.replace(map_key, SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP[map_key])
+                    new_state_dict[new_key] = value
+                    converted_count += 1
+                else:
+                    new_state_dict[full_key] = value
+                    not_converted_count += 1
+            elif full_key.startswith("lora_te1_") or full_key.startswith("lora_te2_"):
+                # The CLIP text encoders have the same keys in both Stability AI and diffusers formats.
+                new_state_dict[full_key] = value
+                continue
+            else:
+                raise ValueError(f"Unrecognized SDXL LoRA key prefix: '{full_key}'.")
+
+        if converted_count > 0 and not_converted_count > 0:
+            raise ValueError(
+                f"The SDXL LoRA could only be partially converted to diffusers format. converted={converted_count},"
+                f" not_converted={not_converted_count}"
+            )
+
+        return new_state_dict
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        file_path: Union[str, Path],
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+        base_model: Optional[BaseModelType] = None,
+    ) -> Self:
+        device = device or torch.device("cpu")
+        dtype = dtype or torch.float32
+
+        if isinstance(file_path, str):
+            file_path = Path(file_path)
+
+        model = cls(
+            name=file_path.stem,
+            layers={},
+        )
+
+        if file_path.suffix == ".safetensors":
+            sd = load_file(file_path.absolute().as_posix(), device="cpu")
+        else:
+            sd = torch.load(file_path, map_location="cpu")
+
+        state_dict = cls._group_state(sd)
+
+        if base_model == BaseModelType.StableDiffusionXL:
+            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:
+                layer: AnyLoRALayer = LoRALayer(layer_key, values)
+
+            # loha
+            elif "hada_w1_a" in values:
+                layer = LoHALayer(layer_key, values)
+
+            # lokr
+            elif "lokr_w1" in values or "lokr_w1_a" in values:
+                layer = LoKRLayer(layer_key, values)
+
+            # diff
+            elif "diff" in values:
+                layer = FullLayer(layer_key, values)
+
+            # ia3
+            elif "on_input" in values:
+                layer = IA3Layer(layer_key, values)
+
+            else:
+                print(f">> Encountered unknown lora layer module in {model.name}: {layer_key} - {list(values.keys())}")
+                raise Exception("Unknown lora format!")
+
+            # lower memory consumption by removing already parsed layer values
+            state_dict[layer_key].clear()
+
+            layer.to(device=device, dtype=dtype)
+            model.layers[layer_key] = layer
+
+        return model
+
+    @staticmethod
+    def _group_state(state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]:
+        state_dict_groupped: Dict[str, Dict[str, torch.Tensor]] = {}
+
+        for key, value in state_dict.items():
+            stem, leaf = key.split(".", 1)
+            if stem not in state_dict_groupped:
+                state_dict_groupped[stem] = {}
+            state_dict_groupped[stem][leaf] = value
+
+        return state_dict_groupped
+
+
+# code from
+# https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L15C1-L97C32
+def make_sdxl_unet_conversion_map() -> List[Tuple[str, str]]:
+    """Create a dict mapping state_dict keys from Stability AI SDXL format to diffusers SDXL format."""
+    unet_conversion_map_layer = []
+
+    for i in range(3):  # num_blocks is 3 in sdxl
+        # loop over downblocks/upblocks
+        for j in range(2):
+            # loop over resnets/attentions for downblocks
+            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
+            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
+            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
+
+            if i < 3:
+                # no attention layers in down_blocks.3
+                hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
+                sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
+                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
+
+        for j in range(3):
+            # loop over resnets/attentions for upblocks
+            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
+            sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
+            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
+
+            # if i > 0: commentout for sdxl
+            # no attention layers in up_blocks.0
+            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
+            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
+            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
+
+        if i < 3:
+            # no downsample in down_blocks.3
+            hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
+            sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
+            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
+
+            # no upsample in up_blocks.3
+            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
+            sd_upsample_prefix = f"output_blocks.{3*i + 2}.{2}."  # change for sdxl
+            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
+
+    hf_mid_atn_prefix = "mid_block.attentions.0."
+    sd_mid_atn_prefix = "middle_block.1."
+    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
+
+    for j in range(2):
+        hf_mid_res_prefix = f"mid_block.resnets.{j}."
+        sd_mid_res_prefix = f"middle_block.{2*j}."
+        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
+
+    unet_conversion_map_resnet = [
+        # (stable-diffusion, HF Diffusers)
+        ("in_layers.0.", "norm1."),
+        ("in_layers.2.", "conv1."),
+        ("out_layers.0.", "norm2."),
+        ("out_layers.3.", "conv2."),
+        ("emb_layers.1.", "time_emb_proj."),
+        ("skip_connection.", "conv_shortcut."),
+    ]
+
+    unet_conversion_map = []
+    for sd, hf in unet_conversion_map_layer:
+        if "resnets" in hf:
+            for sd_res, hf_res in unet_conversion_map_resnet:
+                unet_conversion_map.append((sd + sd_res, hf + hf_res))
+        else:
+            unet_conversion_map.append((sd, hf))
+
+    for j in range(2):
+        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
+        sd_time_embed_prefix = f"time_embed.{j*2}."
+        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
+
+    for j in range(2):
+        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
+        sd_label_embed_prefix = f"label_emb.0.{j*2}."
+        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
+
+    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
+    unet_conversion_map.append(("out.0.", "conv_norm_out."))
+    unet_conversion_map.append(("out.2.", "conv_out."))
+
+    return unet_conversion_map
+
+
+SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP = {
+    sd.rstrip(".").replace(".", "_"): hf.rstrip(".").replace(".", "_") for sd, hf in make_sdxl_unet_conversion_map()
+}

invokeai/backend/lora/conversions/flux_diffusers_lora_conversion_utils.py

@@ -1,210 +0,0 @@
-from typing import Dict
-
-import torch
-
-from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
-from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
-from invokeai.backend.lora.layers.lora_layer import LoRALayer
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
-
-
-def is_state_dict_likely_in_flux_diffusers_format(state_dict: Dict[str, torch.Tensor]) -> bool:
-    """Checks if the provided state dict is likely in the Diffusers FLUX LoRA format.
-
-    This is intended to be a reasonably high-precision detector, but it is not guaranteed to have perfect precision. (A
-    perfect-precision detector would require checking all keys against a whitelist and verifying tensor shapes.)
-    """
-    # First, check that all keys end in "lora_A.weight" or "lora_B.weight" (i.e. are in PEFT format).
-    all_keys_in_peft_format = all(k.endswith(("lora_A.weight", "lora_B.weight")) for k in state_dict.keys())
-
-    # Next, check that this is likely a FLUX model by spot-checking a few keys.
-    expected_keys = [
-        "transformer.single_transformer_blocks.0.attn.to_q.lora_A.weight",
-        "transformer.single_transformer_blocks.0.attn.to_q.lora_B.weight",
-        "transformer.transformer_blocks.0.attn.add_q_proj.lora_A.weight",
-        "transformer.transformer_blocks.0.attn.add_q_proj.lora_B.weight",
-    ]
-    all_expected_keys_present = all(k in state_dict for k in expected_keys)
-
-    return all_keys_in_peft_format and all_expected_keys_present
-
-
-# TODO(ryand): What alpha should we use? 1.0? Rank of the LoRA?
-def lora_model_from_flux_diffusers_state_dict(state_dict: Dict[str, torch.Tensor], alpha: float = 1.0) -> LoRAModelRaw:  # pyright: ignore[reportRedeclaration] (state_dict is intentionally re-declared)
-    """Loads a state dict in the Diffusers FLUX LoRA format into a LoRAModelRaw object.
-
-    This function is based on:
-    https://github.com/huggingface/diffusers/blob/55ac421f7bb12fd00ccbef727be4dc2f3f920abb/scripts/convert_flux_to_diffusers.py
-    """
-    # Group keys by layer.
-    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = _group_by_layer(state_dict)
-
-    # Remove the "transformer." prefix from all keys.
-    grouped_state_dict = {k.replace("transformer.", ""): v for k, v in grouped_state_dict.items()}
-
-    # Constants for FLUX.1
-    num_double_layers = 19
-    num_single_layers = 38
-    # inner_dim = 3072
-    # mlp_ratio = 4.0
-
-    layers: dict[str, AnyLoRALayer] = {}
-
-    def add_lora_layer_if_present(src_key: str, dst_key: str) -> None:
-        if src_key in grouped_state_dict:
-            src_layer_dict = grouped_state_dict.pop(src_key)
-            layers[dst_key] = LoRALayer(
-                values={
-                    "lora_down.weight": src_layer_dict.pop("lora_A.weight"),
-                    "lora_up.weight": src_layer_dict.pop("lora_B.weight"),
-                    "alpha": torch.tensor(alpha),
-                },
-            )
-            assert len(src_layer_dict) == 0
-
-    def add_qkv_lora_layer_if_present(src_keys: list[str], dst_qkv_key: str) -> None:
-        """Handle the Q, K, V matrices for a transformer block. We need special handling because the diffusers format
-        stores them in separate matrices, whereas the BFL format used internally by InvokeAI concatenates them.
-        """
-        # We expect that either all src keys are present or none of them are. Verify this.
-        keys_present = [key in grouped_state_dict for key in src_keys]
-        assert all(keys_present) or not any(keys_present)
-
-        # If none of the keys are present, return early.
-        if not any(keys_present):
-            return
-
-        src_layer_dicts = [grouped_state_dict.pop(key) for key in src_keys]
-        sub_layers: list[LoRALayerBase] = []
-        for src_layer_dict in src_layer_dicts:
-            sub_layers.append(
-                LoRALayer(
-                    values={
-                        "lora_down.weight": src_layer_dict.pop("lora_A.weight"),
-                        "lora_up.weight": src_layer_dict.pop("lora_B.weight"),
-                        "alpha": torch.tensor(alpha),
-                    },
-                )
-            )
-            assert len(src_layer_dict) == 0
-        layers[dst_qkv_key] = ConcatenatedLoRALayer(lora_layers=sub_layers, concat_axis=0)
-
-    # time_text_embed.timestep_embedder -> time_in.
-    add_lora_layer_if_present("time_text_embed.timestep_embedder.linear_1", "time_in.in_layer")
-    add_lora_layer_if_present("time_text_embed.timestep_embedder.linear_2", "time_in.out_layer")
-
-    # time_text_embed.text_embedder -> vector_in.
-    add_lora_layer_if_present("time_text_embed.text_embedder.linear_1", "vector_in.in_layer")
-    add_lora_layer_if_present("time_text_embed.text_embedder.linear_2", "vector_in.out_layer")
-
-    # time_text_embed.guidance_embedder -> guidance_in.
-    add_lora_layer_if_present("time_text_embed.guidance_embedder.linear_1", "guidance_in")
-    add_lora_layer_if_present("time_text_embed.guidance_embedder.linear_2", "guidance_in")
-
-    # context_embedder -> txt_in.
-    add_lora_layer_if_present("context_embedder", "txt_in")
-
-    # x_embedder -> img_in.
-    add_lora_layer_if_present("x_embedder", "img_in")
-
-    # Double transformer blocks.
-    for i in range(num_double_layers):
-        # norms.
-        add_lora_layer_if_present(f"transformer_blocks.{i}.norm1.linear", f"double_blocks.{i}.img_mod.lin")
-        add_lora_layer_if_present(f"transformer_blocks.{i}.norm1_context.linear", f"double_blocks.{i}.txt_mod.lin")
-
-        # Q, K, V
-        add_qkv_lora_layer_if_present(
-            [
-                f"transformer_blocks.{i}.attn.to_q",
-                f"transformer_blocks.{i}.attn.to_k",
-                f"transformer_blocks.{i}.attn.to_v",
-            ],
-            f"double_blocks.{i}.img_attn.qkv",
-        )
-        add_qkv_lora_layer_if_present(
-            [
-                f"transformer_blocks.{i}.attn.add_q_proj",
-                f"transformer_blocks.{i}.attn.add_k_proj",
-                f"transformer_blocks.{i}.attn.add_v_proj",
-            ],
-            f"double_blocks.{i}.txt_attn.qkv",
-        )
-
-        # ff img_mlp
-        add_lora_layer_if_present(
-            f"transformer_blocks.{i}.ff.net.0.proj",
-            f"double_blocks.{i}.img_mlp.0",
-        )
-        add_lora_layer_if_present(
-            f"transformer_blocks.{i}.ff.net.2",
-            f"double_blocks.{i}.img_mlp.2",
-        )
-
-        # ff txt_mlp
-        add_lora_layer_if_present(
-            f"transformer_blocks.{i}.ff_context.net.0.proj",
-            f"double_blocks.{i}.txt_mlp.0",
-        )
-        add_lora_layer_if_present(
-            f"transformer_blocks.{i}.ff_context.net.2",
-            f"double_blocks.{i}.txt_mlp.2",
-        )
-
-        # output projections.
-        add_lora_layer_if_present(
-            f"transformer_blocks.{i}.attn.to_out.0",
-            f"double_blocks.{i}.img_attn.proj",
-        )
-        add_lora_layer_if_present(
-            f"transformer_blocks.{i}.attn.to_add_out",
-            f"double_blocks.{i}.txt_attn.proj",
-        )
-
-    # Single transformer blocks.
-    for i in range(num_single_layers):
-        # norms
-        add_lora_layer_if_present(
-            f"single_transformer_blocks.{i}.norm.linear",
-            f"single_blocks.{i}.modulation.lin",
-        )
-
-        # Q, K, V, mlp
-        add_qkv_lora_layer_if_present(
-            [
-                f"single_transformer_blocks.{i}.attn.to_q",
-                f"single_transformer_blocks.{i}.attn.to_k",
-                f"single_transformer_blocks.{i}.attn.to_v",
-                f"single_transformer_blocks.{i}.proj_mlp",
-            ],
-            f"single_blocks.{i}.linear1",
-        )
-
-        # Output projections.
-        add_lora_layer_if_present(
-            f"single_transformer_blocks.{i}.proj_out",
-            f"single_blocks.{i}.linear2",
-        )
-
-    # Final layer.
-    add_lora_layer_if_present("proj_out", "final_layer.linear")
-
-    # Assert that all keys were processed.
-    assert len(grouped_state_dict) == 0
-
-    return LoRAModelRaw(layers=layers)
-
-
-def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
-    """Groups the keys in the state dict by layer."""
-    layer_dict: dict[str, dict[str, torch.Tensor]] = {}
-    for key in state_dict:
-        # Split the 'lora_A.weight' or 'lora_B.weight' suffix from the layer name.
-        parts = key.rsplit(".", maxsplit=2)
-        layer_name = parts[0]
-        key_name = ".".join(parts[1:])
-        if layer_name not in layer_dict:
-            layer_dict[layer_name] = {}
-        layer_dict[layer_name][key_name] = state_dict[key]
-    return layer_dict

invokeai/backend/lora/conversions/flux_kohya_lora_conversion_utils.py

@@ -1,80 +0,0 @@
-import re
-from typing import Any, Dict, TypeVar
-
-import torch
-
-from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
-from invokeai.backend.lora.layers.utils import any_lora_layer_from_state_dict
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
-
-# A regex pattern that matches all of the keys in the Kohya FLUX LoRA format.
-# Example keys:
-#   lora_unet_double_blocks_0_img_attn_proj.alpha
-#   lora_unet_double_blocks_0_img_attn_proj.lora_down.weight
-#   lora_unet_double_blocks_0_img_attn_proj.lora_up.weight
-FLUX_KOHYA_KEY_REGEX = (
-    r"lora_unet_(\w+_blocks)_(\d+)_(img_attn|img_mlp|img_mod|txt_attn|txt_mlp|txt_mod|linear1|linear2|modulation)_?(.*)"
-)
-
-
-def is_state_dict_likely_in_flux_kohya_format(state_dict: Dict[str, Any]) -> bool:
-    """Checks if the provided state dict is likely in the Kohya FLUX LoRA format.
-
-    This is intended to be a high-precision detector, but it is not guaranteed to have perfect precision. (A
-    perfect-precision detector would require checking all keys against a whitelist and verifying tensor shapes.)
-    """
-    return all(re.match(FLUX_KOHYA_KEY_REGEX, k) for k in state_dict.keys())
-
-
-def lora_model_from_flux_kohya_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRAModelRaw:
-    # Group keys by layer.
-    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {}
-    for key, value in state_dict.items():
-        layer_name, param_name = key.split(".", 1)
-        if layer_name not in grouped_state_dict:
-            grouped_state_dict[layer_name] = {}
-        grouped_state_dict[layer_name][param_name] = value
-
-    # Convert the state dict to the InvokeAI format.
-    grouped_state_dict = convert_flux_kohya_state_dict_to_invoke_format(grouped_state_dict)
-
-    # Create LoRA layers.
-    layers: dict[str, AnyLoRALayer] = {}
-    for layer_key, layer_state_dict in grouped_state_dict.items():
-        layers[layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
-
-    # Create and return the LoRAModelRaw.
-    return LoRAModelRaw(layers=layers)
-
-
-T = TypeVar("T")
-
-
-def convert_flux_kohya_state_dict_to_invoke_format(state_dict: Dict[str, T]) -> Dict[str, T]:
-    """Converts a state dict from the Kohya FLUX LoRA format to LoRA weight format used internally by InvokeAI.
-
-    Example key conversions:
-    "lora_unet_double_blocks_0_img_attn_proj" -> "double_blocks.0.img_attn.proj"
-    "lora_unet_double_blocks_0_img_attn_proj" -> "double_blocks.0.img_attn.proj"
-    "lora_unet_double_blocks_0_img_attn_proj" -> "double_blocks.0.img_attn.proj"
-    "lora_unet_double_blocks_0_img_attn_qkv" -> "double_blocks.0.img_attn.qkv"
-    "lora_unet_double_blocks_0_img_attn_qkv" -> "double_blocks.0.img.attn.qkv"
-    "lora_unet_double_blocks_0_img_attn_qkv" -> "double_blocks.0.img.attn.qkv"
-    """
-
-    def replace_func(match: re.Match[str]) -> str:
-        s = f"{match.group(1)}.{match.group(2)}.{match.group(3)}"
-        if match.group(4):
-            s += f".{match.group(4)}"
-        return s
-
-    converted_dict: dict[str, T] = {}
-    for k, v in state_dict.items():
-        match = re.match(FLUX_KOHYA_KEY_REGEX, k)
-        if match:
-            new_key = re.sub(FLUX_KOHYA_KEY_REGEX, replace_func, k)
-            converted_dict[new_key] = v
-        else:
-            raise ValueError(f"Key '{k}' does not match the expected pattern for FLUX LoRA weights.")
-
-    return converted_dict

invokeai/backend/lora/conversions/sd_lora_conversion_utils.py

@@ -1,29 +0,0 @@
-from typing import Dict
-
-import torch
-
-from invokeai.backend.lora.layers.any_lora_layer import AnyLoRALayer
-from invokeai.backend.lora.layers.utils import any_lora_layer_from_state_dict
-from invokeai.backend.lora.lora_model_raw import LoRAModelRaw
-
-
-def lora_model_from_sd_state_dict(state_dict: Dict[str, torch.Tensor]) -> LoRAModelRaw:
-    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = _group_state(state_dict)
-
-    layers: dict[str, AnyLoRALayer] = {}
-    for layer_key, values in grouped_state_dict.items():
-        layers[layer_key] = any_lora_layer_from_state_dict(values)
-
-    return LoRAModelRaw(layers=layers)
-
-
-def _group_state(state_dict: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, torch.Tensor]]:
-    state_dict_groupped: Dict[str, Dict[str, torch.Tensor]] = {}
-
-    for key, value in state_dict.items():
-        stem, leaf = key.split(".", 1)
-        if stem not in state_dict_groupped:
-            state_dict_groupped[stem] = {}
-        state_dict_groupped[stem][leaf] = value
-
-    return state_dict_groupped

invokeai/backend/lora/conversions/sdxl_lora_conversion_utils.py

@@ -1,154 +0,0 @@
-import bisect
-from typing import Dict, List, Tuple, TypeVar
-
-T = TypeVar("T")
-
-
-def convert_sdxl_keys_to_diffusers_format(state_dict: Dict[str, T]) -> dict[str, T]:
-    """Convert the keys of an SDXL LoRA state_dict to diffusers format.
-
-    The input state_dict can be in either Stability AI format or diffusers format. If the state_dict is already in
-    diffusers format, then this function will have no effect.
-
-    This function is adapted from:
-    https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L385-L409
-
-    Args:
-        state_dict (Dict[str, Tensor]): The SDXL LoRA state_dict.
-
-    Raises:
-        ValueError: If state_dict contains an unrecognized key, or not all keys could be converted.
-
-    Returns:
-        Dict[str, Tensor]: The diffusers-format state_dict.
-    """
-    converted_count = 0  # The number of Stability AI keys converted to diffusers format.
-    not_converted_count = 0  # The number of keys that were not converted.
-
-    # Get a sorted list of Stability AI UNet keys so that we can efficiently search for keys with matching prefixes.
-    # For example, we want to efficiently find `input_blocks_4_1` in the list when searching for
-    # `input_blocks_4_1_proj_in`.
-    stability_unet_keys = list(SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP)
-    stability_unet_keys.sort()
-
-    new_state_dict: dict[str, T] = {}
-    for full_key, value in state_dict.items():
-        if full_key.startswith("lora_unet_"):
-            search_key = full_key.replace("lora_unet_", "")
-            # Use bisect to find the key in stability_unet_keys that *may* match the search_key's prefix.
-            position = bisect.bisect_right(stability_unet_keys, search_key)
-            map_key = stability_unet_keys[position - 1]
-            # Now, check if the map_key *actually* matches the search_key.
-            if search_key.startswith(map_key):
-                new_key = full_key.replace(map_key, SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP[map_key])
-                new_state_dict[new_key] = value
-                converted_count += 1
-            else:
-                new_state_dict[full_key] = value
-                not_converted_count += 1
-        elif full_key.startswith("lora_te1_") or full_key.startswith("lora_te2_"):
-            # The CLIP text encoders have the same keys in both Stability AI and diffusers formats.
-            new_state_dict[full_key] = value
-            continue
-        else:
-            raise ValueError(f"Unrecognized SDXL LoRA key prefix: '{full_key}'.")
-
-    if converted_count > 0 and not_converted_count > 0:
-        raise ValueError(
-            f"The SDXL LoRA could only be partially converted to diffusers format. converted={converted_count},"
-            f" not_converted={not_converted_count}"
-        )
-
-    return new_state_dict
-
-
-# code from
-# https://github.com/bmaltais/kohya_ss/blob/2accb1305979ba62f5077a23aabac23b4c37e935/networks/lora_diffusers.py#L15C1-L97C32
-def _make_sdxl_unet_conversion_map() -> List[Tuple[str, str]]:
-    """Create a dict mapping state_dict keys from Stability AI SDXL format to diffusers SDXL format."""
-    unet_conversion_map_layer: list[tuple[str, str]] = []
-
-    for i in range(3):  # num_blocks is 3 in sdxl
-        # loop over downblocks/upblocks
-        for j in range(2):
-            # loop over resnets/attentions for downblocks
-            hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
-            sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
-            unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
-
-            if i < 3:
-                # no attention layers in down_blocks.3
-                hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
-                sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
-                unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
-
-        for j in range(3):
-            # loop over resnets/attentions for upblocks
-            hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
-            sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
-            unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
-
-            # if i > 0: commentout for sdxl
-            # no attention layers in up_blocks.0
-            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
-            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
-            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
-
-        if i < 3:
-            # no downsample in down_blocks.3
-            hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
-            sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
-            unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
-
-            # no upsample in up_blocks.3
-            hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
-            sd_upsample_prefix = f"output_blocks.{3*i + 2}.{2}."  # change for sdxl
-            unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
-
-    hf_mid_atn_prefix = "mid_block.attentions.0."
-    sd_mid_atn_prefix = "middle_block.1."
-    unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
-
-    for j in range(2):
-        hf_mid_res_prefix = f"mid_block.resnets.{j}."
-        sd_mid_res_prefix = f"middle_block.{2*j}."
-        unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-    unet_conversion_map_resnet = [
-        # (stable-diffusion, HF Diffusers)
-        ("in_layers.0.", "norm1."),
-        ("in_layers.2.", "conv1."),
-        ("out_layers.0.", "norm2."),
-        ("out_layers.3.", "conv2."),
-        ("emb_layers.1.", "time_emb_proj."),
-        ("skip_connection.", "conv_shortcut."),
-    ]
-
-    unet_conversion_map: list[tuple[str, str]] = []
-    for sd, hf in unet_conversion_map_layer:
-        if "resnets" in hf:
-            for sd_res, hf_res in unet_conversion_map_resnet:
-                unet_conversion_map.append((sd + sd_res, hf + hf_res))
-        else:
-            unet_conversion_map.append((sd, hf))
-
-    for j in range(2):
-        hf_time_embed_prefix = f"time_embedding.linear_{j+1}."
-        sd_time_embed_prefix = f"time_embed.{j*2}."
-        unet_conversion_map.append((sd_time_embed_prefix, hf_time_embed_prefix))
-
-    for j in range(2):
-        hf_label_embed_prefix = f"add_embedding.linear_{j+1}."
-        sd_label_embed_prefix = f"label_emb.0.{j*2}."
-        unet_conversion_map.append((sd_label_embed_prefix, hf_label_embed_prefix))
-
-    unet_conversion_map.append(("input_blocks.0.0.", "conv_in."))
-    unet_conversion_map.append(("out.0.", "conv_norm_out."))
-    unet_conversion_map.append(("out.2.", "conv_out."))
-
-    return unet_conversion_map
-
-
-SDXL_UNET_STABILITY_TO_DIFFUSERS_MAP = {
-    sd.rstrip(".").replace(".", "_"): hf.rstrip(".").replace(".", "_") for sd, hf in _make_sdxl_unet_conversion_map()
-}

invokeai/backend/lora/layers/any_lora_layer.py

@@ -1,11 +0,0 @@
-from typing import Union
-
-from invokeai.backend.lora.layers.concatenated_lora_layer import ConcatenatedLoRALayer
-from invokeai.backend.lora.layers.full_layer import FullLayer
-from invokeai.backend.lora.layers.ia3_layer import IA3Layer
-from invokeai.backend.lora.layers.loha_layer import LoHALayer
-from invokeai.backend.lora.layers.lokr_layer import LoKRLayer
-from invokeai.backend.lora.layers.lora_layer import LoRALayer
-from invokeai.backend.lora.layers.norm_layer import NormLayer
-
-AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer, NormLayer, ConcatenatedLoRALayer]

invokeai/backend/lora/layers/concatenated_lora_layer.py

@@ -1,46 +0,0 @@
-from typing import List, Optional
-
-import torch
-
-from invokeai.backend.lora.layers.lora_layer_base import LoRALayerBase
-
-
-class ConcatenatedLoRALayer(LoRALayerBase):
-    """A LoRA layer that is composed of multiple LoRA layers concatenated along a specified axis.
-
-    This class was created to handle a special case with FLUX LoRA models. In the BFL FLUX model format, the attention
-    Q, K, V matrices are concatenated along the first dimension. In the diffusers LoRA format, the Q, K, V matrices are
-    stored as separate tensors. This class enables diffusers LoRA layers to be used in BFL FLUX models.
-    """
-
-    def __init__(self, lora_layers: List[LoRALayerBase], concat_axis: int = 0):
-        # Note: We pass values={} to the base class, because the values are handled by the individual LoRA layers.
-        super().__init__(values={})
-
-        self._lora_layers = lora_layers
-        self._concat_axis = concat_axis
-
-    def get_weight(self, orig_weight: torch.Tensor) -> torch.Tensor:
-        # TODO(ryand): Currently, we pass orig_weight=None to the sub-layers. If we want to support sub-layers that
-        # require this value, we will need to implement chunking of the original weight tensor here.
-        layer_weights = [lora_layer.get_weight(None) for lora_layer in self._lora_layers]  # pyright: ignore[reportArgumentType]
-        return torch.cat(layer_weights, dim=self._concat_axis)
-
-    def get_bias(self, orig_bias: torch.Tensor) -> Optional[torch.Tensor]:
-        # TODO(ryand): Currently, we pass orig_bias=None to the sub-layers. If we want to support sub-layers that
-        # require this value, we will need to implement chunking of the original bias tensor here.
-        layer_biases = [lora_layer.get_bias(None) for lora_layer in self._lora_layers]  # pyright: ignore[reportArgumentType]
-        layer_bias_is_none = [layer_bias is None for layer_bias in layer_biases]
-        if any(layer_bias_is_none):
-            assert all(layer_bias_is_none)
-            return None
-
-        # Ignore the type error, because we have just verified that all layer biases are non-None.
-        return torch.cat(layer_biases, dim=self._concat_axis)
-
-    def calc_size(self) -> int:
-        return sum(lora_layer.calc_size() for lora_layer in self._lora_layers)
-
-    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
-        for lora_layer in self._lora_layers:
-            lora_layer.to(device=device, dtype=dtype)