Merge branch 'main' into dev/pytorch2

'hotdot' to 'hotdog'; the world's least important PR :)

.coveragerc

@@ -1,6 +0,0 @@
-[run]
-omit='.env/*'
-source='.'
-
-[report]
-show_missing = true

.dockerignore

@@ -4,22 +4,22 @@
 !ldm
 !pyproject.toml
 
-# Guard against pulling in any models that might exist in the directory tree
-**/*.pt*
-**/*.ckpt
-
-# ignore frontend but whitelist dist
-invokeai/frontend/
-!invokeai/frontend/dist/
+# ignore frontend/web but whitelist dist
+invokeai/frontend/web/
+!invokeai/frontend/web/dist/
 
 # ignore invokeai/assets but whitelist invokeai/assets/web
 invokeai/assets/
 !invokeai/assets/web/
 
+# Guard against pulling in any models that might exist in the directory tree
+**/*.pt*
+**/*.ckpt
+
 # Byte-compiled / optimized / DLL files
 **/__pycache__/
 **/*.py[cod]
 
 # Distribution / packaging
-*.egg-info/
-*.egg
+**/*.egg-info/
+**/*.egg

.git-blame-ignore-revs

@@ -0,0 +1 @@
+b3dccfaeb636599c02effc377cdd8a87d658256c

.github/CODEOWNERS

@@ -3,49 +3,32 @@
 
 # documentation
 /docs/ @lstein @mauwii @tildebyte @blessedcoolant
-mkdocs.yml @lstein @mauwii @blessedcoolant
+/mkdocs.yml @lstein @mauwii @blessedcoolant
+
+# nodes
+/invokeai/app/ @Kyle0654 @blessedcoolant
 
 # installation and configuration
-/pyproject.toml @mauwii @lstein @ebr @blessedcoolant
+/pyproject.toml @mauwii @lstein @blessedcoolant
 /docker/ @mauwii @lstein @blessedcoolant
-/scripts/ @ebr @lstein @blessedcoolant
-/installer/ @ebr @lstein @tildebyte  @blessedcoolant
-ldm/invoke/config @lstein @ebr  @blessedcoolant
-invokeai/assets @lstein @ebr  @blessedcoolant
-invokeai/configs @lstein @ebr  @blessedcoolant
-/ldm/invoke/_version.py @lstein @blessedcoolant
+/scripts/ @ebr @lstein
+/installer/ @lstein @ebr
+/invokeai/assets @lstein @ebr
+/invokeai/configs @lstein
+/invokeai/version @lstein @blessedcoolant
 
 # web ui
 /invokeai/frontend @blessedcoolant @psychedelicious @lstein
 /invokeai/backend @blessedcoolant @psychedelicious @lstein
 
-# generation and model management
-/ldm/*.py @lstein @blessedcoolant
-/ldm/generate.py @lstein @keturn @blessedcoolant
-/ldm/invoke/args.py @lstein @blessedcoolant
-/ldm/invoke/ckpt* @lstein @blessedcoolant
-/ldm/invoke/ckpt_generator @lstein @blessedcoolant
-/ldm/invoke/CLI.py @lstein @blessedcoolant
-/ldm/invoke/config @lstein @ebr @mauwii @blessedcoolant
-/ldm/invoke/generator @keturn @damian0815 @blessedcoolant
-/ldm/invoke/globals.py @lstein @blessedcoolant 
-/ldm/invoke/merge_diffusers.py @lstein @blessedcoolant
-/ldm/invoke/model_manager.py @lstein @blessedcoolant
-/ldm/invoke/txt2mask.py @lstein @blessedcoolant
-/ldm/invoke/patchmatch.py @Kyle0654 @blessedcoolant @lstein
-/ldm/invoke/restoration @lstein @blessedcoolant
+# generation, model management, postprocessing
+/invokeai/backend @keturn @damian0815 @lstein @blessedcoolant @jpphoto
 
-# attention, textual inversion, model configuration
-/ldm/models @damian0815 @keturn @lstein @blessedcoolant
-/ldm/modules @damian0815 @keturn @lstein @blessedcoolant
+# front ends
+/invokeai/frontend/CLI @lstein
+/invokeai/frontend/install @lstein @ebr @mauwii 
+/invokeai/frontend/merge @lstein @blessedcoolant @hipsterusername
+/invokeai/frontend/training @lstein @blessedcoolant @hipsterusername
+/invokeai/frontend/web @psychedelicious @blessedcoolant
 
-# Nodes
-apps/ @Kyle0654 @lstein @blessedcoolant
-
-# legacy REST API
-# is CapableWeb still engaged?
-/ldm/invoke/pngwriter.py @CapableWeb @lstein @blessedcoolant
-/ldm/invoke/server_legacy.py @CapableWeb @lstein @blessedcoolant
-/scripts/legacy_api.py @CapableWeb @lstein @blessedcoolant
-/tests/legacy_tests.sh @CapableWeb @lstein @blessedcoolant
 

.github/ISSUE_TEMPLATE/BUG_REPORT.yml

@@ -65,6 +65,16 @@ body:
       placeholder: 8GB
     validations:
       required: false
+      
+  - type: input
+    id: version-number
+    attributes:
+      label: What version did you experience this issue on?
+      description: |
+        Please share the version of Invoke AI that you experienced the issue on. If this is not the latest version, please update first to confirm the issue still exists. If you are testing main, please include the commit hash instead.
+      placeholder: X.X.X
+    validations:
+      required: true
 
   - type: textarea
     id: what-happened

.github/workflows/build-container.yml

@@ -5,17 +5,20 @@ on:
       - 'main'
       - 'update/ci/docker/*'
       - 'update/docker/*'
+      - 'dev/ci/docker/*'
+      - 'dev/docker/*'
     paths:
       - 'pyproject.toml'
-      - 'ldm/**'
-      - 'invokeai/backend/**'
-      - 'invokeai/configs/**'
-      - 'invokeai/frontend/dist/**'
+      - '.dockerignore'
+      - 'invokeai/**'
       - 'docker/Dockerfile'
     tags:
       - 'v*.*.*'
   workflow_dispatch:
 
+permissions:
+  contents: write
+  packages: write
 
 jobs:
   docker:
@@ -24,11 +27,11 @@ jobs:
       fail-fast: false
       matrix:
         flavor:
-          - amd
+          - rocm
           - cuda
           - cpu
         include:
-          - flavor: amd
+          - flavor: rocm
             pip-extra-index-url: 'https://download.pytorch.org/whl/rocm5.2'
           - flavor: cuda
             pip-extra-index-url: ''
@@ -54,9 +57,9 @@ jobs:
           tags: |
             type=ref,event=branch
             type=ref,event=tag
-            type=semver,pattern={{version}}
-            type=semver,pattern={{major}}.{{minor}}
-            type=semver,pattern={{major}}
+            type=pep440,pattern={{version}}
+            type=pep440,pattern={{major}}.{{minor}}
+            type=pep440,pattern={{major}}
             type=sha,enable=true,prefix=sha-,format=short
           flavor: |
             latest=${{ matrix.flavor == 'cuda' && github.ref == 'refs/heads/main' }}
@@ -92,7 +95,7 @@ jobs:
           context: .
           file: ${{ env.DOCKERFILE }}
           platforms: ${{ env.PLATFORMS }}
-          push: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/tags/*' }}
+          push: ${{ github.ref == 'refs/heads/main' || github.ref_type == 'tag' }}
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
           build-args: PIP_EXTRA_INDEX_URL=${{ matrix.pip-extra-index-url }}

.github/workflows/close-inactive-issues.yml

@@ -0,0 +1,27 @@
+name: Close inactive issues
+on:
+  schedule:
+    - cron: "00 6 * * *"
+
+env:
+  DAYS_BEFORE_ISSUE_STALE: 14
+  DAYS_BEFORE_ISSUE_CLOSE: 28
+
+jobs:
+  close-issues:
+    runs-on: ubuntu-latest
+    permissions:
+      issues: write
+      pull-requests: write
+    steps:
+      - uses: actions/stale@v5
+        with:
+          days-before-issue-stale: ${{ env.DAYS_BEFORE_ISSUE_STALE }}
+          days-before-issue-close: ${{ env.DAYS_BEFORE_ISSUE_CLOSE }}
+          stale-issue-label: "Inactive Issue"
+          stale-issue-message: "There has been no activity in this issue for ${{ env.DAYS_BEFORE_ISSUE_STALE }} days. If this issue is still being experienced, please reply with an updated confirmation that the issue is still being experienced with the latest release."
+          close-issue-message: "Due to inactivity, this issue was automatically closed. If you are still experiencing the issue, please recreate the issue."
+          days-before-pr-stale: -1
+          days-before-pr-close: -1
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          operations-per-run: 500

.github/workflows/lint-frontend.yml

@@ -3,14 +3,22 @@ name: Lint frontend
 on:
   pull_request:
     paths:
-      - 'invokeai/frontend/**'
+      - 'invokeai/frontend/web/**'
+    types:
+      - 'ready_for_review'
+      - 'opened'
+      - 'synchronize'
   push:
+    branches:
+      - 'main'
     paths:
-      - 'invokeai/frontend/**'
+      - 'invokeai/frontend/web/**'
+  merge_group:
+  workflow_dispatch:
 
 defaults:
   run:
-    working-directory: invokeai/frontend
+    working-directory: invokeai/frontend/web
 
 jobs:
   lint-frontend:
@@ -23,7 +31,7 @@ jobs:
           node-version: '18'
       - uses: actions/checkout@v3
       - run: 'yarn install --frozen-lockfile'
-      - run: 'yarn tsc'
-      - run: 'yarn run madge'
-      - run: 'yarn run lint --max-warnings=0'
-      - run: 'yarn run prettier --check'
+      - run: 'yarn run lint:tsc'
+      - run: 'yarn run lint:madge'
+      - run: 'yarn run lint:eslint'
+      - run: 'yarn run lint:prettier'

.github/workflows/mkdocs-material.yml

@@ -5,6 +5,9 @@ on:
       - 'main'
       - 'development'
 
+permissions:
+    contents: write
+
 jobs:
   mkdocs-material:
     if: github.event.pull_request.draft == false

.github/workflows/pypi-release.yml

@@ -3,7 +3,7 @@ name: PyPI Release
 on:
   push:
     paths:
-      - 'ldm/invoke/_version.py'
+      - 'invokeai/version/invokeai_version.py'
   workflow_dispatch:
 
 jobs:

.github/workflows/test-invoke-pip-skip.yml

@@ -1,12 +1,11 @@
 name: Test invoke.py pip
 on:
   pull_request:
-    paths-ignore:
-      - 'pyproject.toml'
-      - 'ldm/**'
-      - 'invokeai/backend/**'
-      - 'invokeai/configs/**'
-      - 'invokeai/frontend/dist/**'
+    paths:
+      - '**'
+      - '!pyproject.toml'
+      - '!invokeai/**'
+      - 'invokeai/frontend/web/**'
   merge_group:
   workflow_dispatch:
 

.github/workflows/test-invoke-pip.yml

@@ -5,17 +5,13 @@ on:
       - 'main'
     paths:
       - 'pyproject.toml'
-      - 'ldm/**'
-      - 'invokeai/backend/**'
-      - 'invokeai/configs/**'
-      - 'invokeai/frontend/dist/**'
+      - 'invokeai/**'
+      - '!invokeai/frontend/web/**'
   pull_request:
     paths:
       - 'pyproject.toml'
-      - 'ldm/**'
-      - 'invokeai/backend/**'
-      - 'invokeai/configs/**'
-      - 'invokeai/frontend/dist/**'
+      - 'invokeai/**'
+      - '!invokeai/frontend/web/**'
     types:
       - 'ready_for_review'
       - 'opened'
@@ -112,7 +108,7 @@ jobs:
       - name: set INVOKEAI_OUTDIR
         run: >
           python -c
-          "import os;from ldm.invoke.globals import Globals;OUTDIR=os.path.join(Globals.root,str('outputs'));print(f'INVOKEAI_OUTDIR={OUTDIR}')"
+          "import os;from invokeai.backend.globals import Globals;OUTDIR=os.path.join(Globals.root,str('outputs'));print(f'INVOKEAI_OUTDIR={OUTDIR}')"
           >> ${{ matrix.github-env }}
 
       - name: run invokeai-configure

.gitignore

@@ -63,6 +63,7 @@ pip-delete-this-directory.txt
 htmlcov/
 .tox/
 .nox/
+.coveragerc
 .coverage
 .coverage.*
 .cache
@@ -73,6 +74,7 @@ cov.xml
 *.py,cover
 .hypothesis/
 .pytest_cache/
+.pytest.ini
 cover/
 junit/
 
@@ -198,7 +200,7 @@ checkpoints
 .DS_Store
 
 # Let the frontend manage its own gitignore
-!invokeai/frontend/*
+!invokeai/frontend/web/*
 
 # Scratch folder
 .scratch/
@@ -213,11 +215,6 @@ gfpgan/
 # config file (will be created by installer)
 configs/models.yaml
 
-# weights (will be created by installer)
-models/ldm/stable-diffusion-v1/*.ckpt
-models/clipseg
-models/gfpgan
-
 # ignore initfile
 .invokeai
 
@@ -232,6 +229,3 @@ installer/install.bat
 installer/install.sh
 installer/update.bat
 installer/update.sh
-
-# no longer stored in source directory
-models

.pytest.ini

@@ -1,5 +0,0 @@
-[pytest]
-DJANGO_SETTINGS_MODULE = webtas.settings
-; python_files = tests.py test_*.py *_tests.py
-
-addopts = --cov=. --cov-config=.coveragerc --cov-report xml:cov.xml

README.md

@@ -139,13 +139,13 @@ not supported.
     _For Windows/Linux with an NVIDIA GPU:_
 
     ```terminal
-    pip install InvokeAI[xformers] --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu117
+    pip install "InvokeAI[xformers]" --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu117
     ```
 
     _For Linux with an AMD GPU:_
 
     ```sh
-    pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.2
+    pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.4.2
     ```
 
     _For Macintoshes, either Intel or M1/M2:_

coverage/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

docker/Dockerfile

@@ -4,15 +4,15 @@ ARG PYTHON_VERSION=3.9
 ##################
 ##  base image  ##
 ##################
-FROM python:${PYTHON_VERSION}-slim AS python-base
+FROM --platform=${TARGETPLATFORM} python:${PYTHON_VERSION}-slim AS python-base
 
 LABEL org.opencontainers.image.authors="mauwii@outlook.de"
 
-# prepare for buildkit cache
+# Prepare apt for buildkit cache
 RUN rm -f /etc/apt/apt.conf.d/docker-clean \
   && echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' >/etc/apt/apt.conf.d/keep-cache
 
-# Install necessary packages
+# Install dependencies
 RUN \
   --mount=type=cache,target=/var/cache/apt,sharing=locked \
   --mount=type=cache,target=/var/lib/apt,sharing=locked \
@@ -23,7 +23,7 @@ RUN \
     libglib2.0-0=2.66.* \
     libopencv-dev=4.5.*
 
-# set working directory and env
+# Set working directory and env
 ARG APPDIR=/usr/src
 ARG APPNAME=InvokeAI
 WORKDIR ${APPDIR}
@@ -32,7 +32,7 @@ ENV PATH ${APPDIR}/${APPNAME}/bin:$PATH
 ENV PYTHONDONTWRITEBYTECODE 1
 # Turns off buffering for easier container logging
 ENV PYTHONUNBUFFERED 1
-# don't fall back to legacy build system
+# Don't fall back to legacy build system
 ENV PIP_USE_PEP517=1
 
 #######################
@@ -40,7 +40,7 @@ ENV PIP_USE_PEP517=1
 #######################
 FROM python-base AS pyproject-builder
 
-# Install dependencies
+# Install build dependencies
 RUN \
   --mount=type=cache,target=/var/cache/apt,sharing=locked \
   --mount=type=cache,target=/var/lib/apt,sharing=locked \
@@ -51,26 +51,30 @@ RUN \
     gcc=4:10.2.* \
     python3-dev=3.9.*
 
-# prepare pip for buildkit cache
+# Prepare pip for buildkit cache
 ARG PIP_CACHE_DIR=/var/cache/buildkit/pip
 ENV PIP_CACHE_DIR ${PIP_CACHE_DIR}
 RUN mkdir -p ${PIP_CACHE_DIR}
 
-# create virtual environment
-RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
+# Create virtual environment
+RUN --mount=type=cache,target=${PIP_CACHE_DIR} \
   python3 -m venv "${APPNAME}" \
   --upgrade-deps
 
-# copy sources
-COPY --link . .
-
-# install pyproject.toml
+# Install requirements
+COPY --link pyproject.toml .
+COPY --link invokeai/version/invokeai_version.py invokeai/version/__init__.py invokeai/version/
 ARG PIP_EXTRA_INDEX_URL
 ENV PIP_EXTRA_INDEX_URL ${PIP_EXTRA_INDEX_URL}
-RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
+RUN --mount=type=cache,target=${PIP_CACHE_DIR} \
+  "${APPNAME}"/bin/pip install .
+
+# Install pyproject.toml
+COPY --link . .
+RUN --mount=type=cache,target=${PIP_CACHE_DIR} \
   "${APPNAME}/bin/pip" install .
 
-# build patchmatch
+# Build patchmatch
 RUN python3 -c "from patchmatch import patch_match"
 
 #####################
@@ -86,14 +90,14 @@ RUN useradd \
   -U \
   "${UNAME}"
 
-# create volume directory
+# Create volume directory
 ARG VOLUME_DIR=/data
 RUN mkdir -p "${VOLUME_DIR}" \
-  && chown -R "${UNAME}" "${VOLUME_DIR}"
+  && chown -hR "${UNAME}:${UNAME}" "${VOLUME_DIR}"
 
-# setup runtime environment
-USER ${UNAME}
-COPY --chown=${UNAME} --from=pyproject-builder ${APPDIR}/${APPNAME} ${APPNAME}
+# Setup runtime environment
+USER ${UNAME}:${UNAME}
+COPY --chown=${UNAME}:${UNAME} --from=pyproject-builder ${APPDIR}/${APPNAME} ${APPNAME}
 ENV INVOKEAI_ROOT ${VOLUME_DIR}
 ENV TRANSFORMERS_CACHE ${VOLUME_DIR}/.cache
 ENV INVOKE_MODEL_RECONFIGURE "--yes --default_only"

docker/build.sh

@@ -41,7 +41,7 @@ else
 fi
 
 # Build Container
-DOCKER_BUILDKIT=1 docker build \
+docker build \
     --platform="${PLATFORM:-linux/amd64}" \
     --tag="${CONTAINER_IMAGE:-invokeai}" \
     ${CONTAINER_FLAVOR:+--build-arg="CONTAINER_FLAVOR=${CONTAINER_FLAVOR}"} \

docker/env.sh

@@ -49,3 +49,6 @@ CONTAINER_FLAVOR="${CONTAINER_FLAVOR-cuda}"
 CONTAINER_TAG="${CONTAINER_TAG-"${INVOKEAI_BRANCH##*/}-${CONTAINER_FLAVOR}"}"
 CONTAINER_IMAGE="${CONTAINER_REGISTRY}/${CONTAINER_REPOSITORY}:${CONTAINER_TAG}"
 CONTAINER_IMAGE="${CONTAINER_IMAGE,,}"
+
+# enable docker buildkit
+export DOCKER_BUILDKIT=1

docker/run.sh

@@ -21,10 +21,10 @@ docker run \
   --tty \
   --rm \
   --platform="${PLATFORM}" \
-  --name="${REPOSITORY_NAME,,}" \
-  --hostname="${REPOSITORY_NAME,,}" \
-  --mount=source="${VOLUMENAME}",target=/data \
-  --mount type=bind,source="$(pwd)"/outputs,target=/data/outputs \
+  --name="${REPOSITORY_NAME}" \
+  --hostname="${REPOSITORY_NAME}" \
+  --mount type=volume,volume-driver=local,source="${VOLUMENAME}",target=/data \
+  --mount type=bind,source="$(pwd)"/outputs/,target=/data/outputs/ \
   ${MODELSPATH:+--mount="type=bind,source=${MODELSPATH},target=/data/models"} \
   ${HUGGING_FACE_HUB_TOKEN:+--env="HUGGING_FACE_HUB_TOKEN=${HUGGING_FACE_HUB_TOKEN}"} \
   --publish=9090:9090 \
@@ -32,7 +32,7 @@ docker run \
   ${GPU_FLAGS:+--gpus="${GPU_FLAGS}"} \
   "${CONTAINER_IMAGE}" ${@:+$@}
 
-# Remove Trash folder
+echo -e "\nCleaning trash folder ..."
 for f in outputs/.Trash*; do
   if [ -e "$f" ]; then
     rm -Rf "$f"

docs/contributing/LOCAL_DEVELOPMENT.md

@@ -0,0 +1,83 @@
+# Local Development
+
+If you are looking to contribute you will need to have a local development
+environment. See the
+[Developer Install](../installation/020_INSTALL_MANUAL.md#developer-install) for
+full details.
+
+Broadly this involves cloning the repository, installing the pre-reqs, and
+InvokeAI (in editable form). Assuming this is working, choose your area of
+focus.
+
+## Documentation
+
+We use [mkdocs](https://www.mkdocs.org) for our documentation with the
+[material theme](https://squidfunk.github.io/mkdocs-material/). Documentation is
+written in markdown files under the `./docs` folder and then built into a static
+website for hosting with GitHub Pages at
+[invoke-ai.github.io/InvokeAI](https://invoke-ai.github.io/InvokeAI).
+
+To contribute to the documentation you'll need to install the dependencies. Note
+the use of `"`.
+
+```zsh
+pip install ".[docs]"
+```
+
+Now, to run the documentation locally with hot-reloading for changes made.
+
+```zsh
+mkdocs serve
+```
+
+You'll then be prompted to connect to `http://127.0.0.1:8080` in order to
+access.
+
+## Backend
+
+The backend is contained within the `./invokeai/backend` folder structure. To
+get started however please install the development dependencies.
+
+From the root of the repository run the following command. Note the use of `"`.
+
+```zsh
+pip install ".[test]"
+```
+
+This in an optional group of packages which is defined within the
+`pyproject.toml` and will be required for testing the changes you make the the
+code.
+
+### Running Tests
+
+We use [pytest](https://docs.pytest.org/en/7.2.x/) for our test suite. Tests can
+be found under the `./tests` folder and can be run with a single `pytest`
+command. Optionally, to review test coverage you can append `--cov`.
+
+```zsh
+pytest --cov
+```
+
+Test outcomes and coverage will be reported in the terminal. In addition a more
+detailed report is created in both XML and HTML format in the `./coverage`
+folder. The HTML one in particular can help identify missing statements
+requiring tests to ensure coverage. This can be run by opening
+`./coverage/html/index.html`.
+
+For example.
+
+```zsh
+pytest --cov; open ./coverage/html/index.html
+```
+
+??? info "HTML coverage report output"
+
+    ![html-overview](../assets/contributing/html-overview.png)
+
+    ![html-detail](../assets/contributing/html-detail.png)
+
+## Front End
+
+<!--#TODO: get input from blessedcoolant here, for the moment inserted the frontend README via snippets extension.-->
+
+--8<-- "invokeai/frontend/web/README.md"

docs/features/INPAINTING.md

@@ -168,11 +168,15 @@ used by Stable Diffusion 1.4 and 1.5.
 After installation, your `models.yaml` should contain an entry that looks like
 this one:
 
-inpainting-1.5: weights: models/ldm/stable-diffusion-v1/sd-v1-5-inpainting.ckpt
-description: SD inpainting v1.5 config:
-configs/stable-diffusion/v1-inpainting-inference.yaml vae:
-models/ldm/stable-diffusion-v1/vae-ft-mse-840000-ema-pruned.ckpt width: 512
-height: 512
+```yml
+inpainting-1.5:
+  weights: models/ldm/stable-diffusion-v1/sd-v1-5-inpainting.ckpt
+  description: SD inpainting v1.5
+  config: configs/stable-diffusion/v1-inpainting-inference.yaml
+  vae: models/ldm/stable-diffusion-v1/vae-ft-mse-840000-ema-pruned.ckpt
+  width: 512
+  height: 512
+```
 
 As shown in the example, you may include a VAE fine-tuning weights file as well.
 This is strongly recommended.

docs/features/PROMPTS.md

@@ -268,7 +268,7 @@ model is so good at inpainting, a good substitute is to use the `clipseg` text
 masking option:
 
 ```bash
-invoke> a fluffy cat eating a hotdot
+invoke> a fluffy cat eating a hotdog
 Outputs:
 [1010] outputs/000025.2182095108.png: a fluffy cat eating a hotdog
 invoke> a smiling dog eating a hotdog -I 000025.2182095108.png -tm cat

docs/features/TEXTUAL_INVERSION.md

@@ -17,7 +17,7 @@ notebooks.
 
 You will need a GPU to perform training in a reasonable length of
 time, and at least 12 GB of VRAM. We recommend using the [`xformers`
-library](../installation/070_INSTALL_XFORMERS) to accelerate the
+library](../installation/070_INSTALL_XFORMERS.md) to accelerate the
 training process further. During training, about ~8 GB is temporarily
 needed in order to store intermediate models, checkpoints and logs.
 

docs/installation/010_INSTALL_AUTOMATED.md

@@ -417,7 +417,7 @@ Then type the following commands:
 
 === "AMD System"
     ```bash
-    pip install torch torchvision --force-reinstall --extra-index-url https://download.pytorch.org/whl/rocm5.2
+    pip install torch torchvision --force-reinstall --extra-index-url https://download.pytorch.org/whl/rocm5.4.2
     ```
 
 ### Corrupted configuration file

docs/installation/020_INSTALL_MANUAL.md

@@ -148,13 +148,13 @@ manager, please follow these steps:
     === "CUDA (NVidia)"
 
         ```bash
-        pip install InvokeAI[xformers] --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu117
+        pip install "InvokeAI[xformers]" --use-pep517 --extra-index-url https://download.pytorch.org/whl/cu117
         ```
 
     === "ROCm (AMD)"
 
         ```bash
-        pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.2
+        pip install InvokeAI --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.4.2
         ```
 
     === "CPU (Intel Macs & non-GPU systems)"
@@ -315,7 +315,7 @@ installation protocol (important!)
 
     === "ROCm (AMD)"
         ```bash
-        pip install -e . --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.2
+        pip install -e . --use-pep517 --extra-index-url https://download.pytorch.org/whl/rocm5.4.2
         ```
 
     === "CPU (Intel Macs & non-GPU systems)"

docs/installation/030_INSTALL_CUDA_AND_ROCM.md

@@ -110,7 +110,7 @@ recipes are available
 
 When installing torch and torchvision manually with `pip`, remember to provide
 the argument `--extra-index-url
-https://download.pytorch.org/whl/rocm5.2` as described in the [Manual
+https://download.pytorch.org/whl/rocm5.4.2` as described in the [Manual
 Installation Guide](020_INSTALL_MANUAL.md).
 
 This will be done automatically for you if you use the installer

docs/installation/060_INSTALL_PATCHMATCH.md

@@ -24,7 +24,7 @@ You need to have opencv installed so that pypatchmatch can be built:
 brew install opencv
 ```
 
-The next time you start `invoke`, after sucesfully installing opencv, pypatchmatch will be built.
+The next time you start `invoke`, after successfully installing opencv, pypatchmatch will be built.
 
 ## Linux
 
@@ -56,7 +56,7 @@ Prior to installing PyPatchMatch, you need to take the following steps:
 
 5. Confirm that pypatchmatch is installed. At the command-line prompt enter
    `python`, and then at the `>>>` line type
-   `from patchmatch import patch_match`: It should look like the follwing:
+   `from patchmatch import patch_match`: It should look like the following:
 
     ```py
     Python 3.9.5 (default, Nov 23 2021, 15:27:38)
@@ -108,4 +108,4 @@ Prior to installing PyPatchMatch, you need to take the following steps:
 
 [**Next, Follow Steps 4-6 from the Debian Section above**](#linux)
 
-If you see no errors, then you're ready to go!
+If you see no errors you're ready to go!

installer/create_installer.sh

@@ -11,10 +11,10 @@ if [[ -v "VIRTUAL_ENV" ]]; then
     exit -1
 fi
 
-VERSION=$(cd ..; python -c "from ldm.invoke import __version__ as version; print(version)")
+VERSION=$(cd ..; python -c "from invokeai.version import __version__ as version; print(version)")
 PATCH=""
 VERSION="v${VERSION}${PATCH}"
-LATEST_TAG="v2.3-latest"
+LATEST_TAG="v3.0-latest"
 
 echo Building installer for version $VERSION
 echo "Be certain that you're in the 'installer' directory before continuing."

installer/lib/installer.py

@@ -291,7 +291,7 @@ class InvokeAiInstance:
                     src = Path(__file__).parents[1].expanduser().resolve()
                     # if the above directory contains one of these files, we'll do a source install
                     next(src.glob("pyproject.toml"))
-                    next(src.glob("ldm"))
+                    next(src.glob("invokeai"))
                 except StopIteration:
                     print("Unable to find a wheel or perform a source install. Giving up.")
 
@@ -342,14 +342,14 @@ class InvokeAiInstance:
 
         introduction()
 
-        from ldm.invoke.config import invokeai_configure
+        from invokeai.frontend.install import invokeai_configure
 
         # NOTE: currently the config script does its own arg parsing! this means the command-line switches
         # from the installer will also automatically propagate down to the config script.
         # this may change in the future with config refactoring!
         succeeded = False
         try:
-            invokeai_configure.main()
+            invokeai_configure()
             succeeded = True
         except requests.exceptions.ConnectionError as e:
             print(f'\nA network error was encountered during configuration and download: {str(e)}')
@@ -456,13 +456,12 @@ def get_torch_source() -> (Union[str, None],str):
     optional_modules = None
     if OS == "Linux":
         if device == "rocm":
-            url = "https://download.pytorch.org/whl/rocm5.2"
+            url = "https://download.pytorch.org/whl/rocm5.4.2"
         elif device == "cpu":
             url = "https://download.pytorch.org/whl/cpu"
 
     if device == 'cuda':
-        url = 'https://download.pytorch.org/whl/cu117'
-        optional_modules = '[xformers]'
+        url = 'https://download.pytorch.org/whl/cu118'
 
     # in all other cases, Torch wheels should be coming from PyPi as of Torch 1.13
 

installer/templates/invoke.sh.in

@@ -24,9 +24,9 @@ if [ "$(uname -s)" == "Darwin" ]; then
     export PYTORCH_ENABLE_MPS_FALLBACK=1
 fi
 
-while true
-do
 if [ "$0" != "bash" ]; then
+  while true
+  do
     echo "Do you want to generate images using the"
     echo "1. command-line interface"
     echo "2. browser-based UI"
@@ -67,29 +67,29 @@ if [ "$0" != "bash" ]; then
             ;;
         7)
             invokeai-configure --root ${INVOKEAI_ROOT} --yes --default_only
-	    ;;
-	8)
-	    echo "Developer Console:"
+            ;;
+        8)
+            echo "Developer Console:"
             file_name=$(basename "${BASH_SOURCE[0]}")
             bash --init-file "$file_name"
             ;;
         9)
-	    echo "Update:"
+            echo "Update:"
             invokeai-update
             ;;
         10)
             invokeai --help
             ;;
-	[qQ])
+        [qQ])
             exit 0
             ;;
         *)
             echo "Invalid selection"
             exit;;
     esac
+ done
 else # in developer console
     python --version
     echo "Press ^D to exit"
     export PS1="(InvokeAI) \u@\h \w> "
 fi
-done

invokeai/README

@@ -1,3 +1,11 @@
-After version 2.3 is released, the ldm/invoke modules will be migrated to this location
-so that we have a proper invokeai distribution. Currently it is only being used for
-data files.
+Organization of the source tree:
+
+app -- Home of nodes invocations and services
+assets -- Images and other data files used by InvokeAI
+backend -- Non-user facing libraries, including the rendering
+	core.
+configs -- Configuration files used at install and run times
+frontend -- User-facing scripts, including the CLI and the WebUI
+version -- Current InvokeAI version string, stored
+	in version/invokeai_version.py
+	

invokeai/app/api/dependencies.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+import os
+from argparse import Namespace
+
+from ..services.latent_storage import DiskLatentsStorage, ForwardCacheLatentsStorage
+
+from ...backend import Globals
+from ..services.model_manager_initializer import get_model_manager
+from ..services.restoration_services import RestorationServices
+from ..services.graph import GraphExecutionState
+from ..services.image_storage import DiskImageStorage
+from ..services.invocation_queue import MemoryInvocationQueue
+from ..services.invocation_services import InvocationServices
+from ..services.invoker import Invoker
+from ..services.processor import DefaultInvocationProcessor
+from ..services.sqlite import SqliteItemStorage
+from .events import FastAPIEventService
+
+
+# TODO: is there a better way to achieve this?
+def check_internet() -> bool:
+    """
+    Return true if the internet is reachable.
+    It does this by pinging huggingface.co.
+    """
+    import urllib.request
+
+    host = "http://huggingface.co"
+    try:
+        urllib.request.urlopen(host, timeout=1)
+        return True
+    except:
+        return False
+
+
+class ApiDependencies:
+    """Contains and initializes all dependencies for the API"""
+
+    invoker: Invoker = None
+
+    @staticmethod
+    def initialize(config, event_handler_id: int):
+        Globals.try_patchmatch = config.patchmatch
+        Globals.always_use_cpu = config.always_use_cpu
+        Globals.internet_available = config.internet_available and check_internet()
+        Globals.disable_xformers = not config.xformers
+        Globals.ckpt_convert = config.ckpt_convert
+
+        # TODO: Use a logger
+        print(f">> Internet connectivity is {Globals.internet_available}")
+
+        events = FastAPIEventService(event_handler_id)
+
+        output_folder = os.path.abspath(
+            os.path.join(os.path.dirname(__file__), "../../../../outputs")
+        )
+
+        latents = ForwardCacheLatentsStorage(DiskLatentsStorage(f'{output_folder}/latents'))
+
+        images = DiskImageStorage(f'{output_folder}/images')
+
+        # TODO: build a file/path manager?
+        db_location = os.path.join(output_folder, "invokeai.db")
+
+        services = InvocationServices(
+            model_manager=get_model_manager(config),
+            events=events,
+            latents=latents,
+            images=images,
+            queue=MemoryInvocationQueue(),
+            graph_execution_manager=SqliteItemStorage[GraphExecutionState](
+                filename=db_location, table_name="graph_executions"
+            ),
+            processor=DefaultInvocationProcessor(),
+            restoration=RestorationServices(config),
+        )
+
+        ApiDependencies.invoker = Invoker(services)
+
+    @staticmethod
+    def shutdown():
+        if ApiDependencies.invoker:
+            ApiDependencies.invoker.stop()

invokeai/app/api/events.py

@@ -1,11 +1,14 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
 import asyncio
+import threading
 from queue import Empty, Queue
 from typing import Any
+
 from fastapi_events.dispatcher import dispatch
+
 from ..services.events import EventServiceBase
-import threading
+
 
 class FastAPIEventService(EventServiceBase):
     event_handler_id: int
@@ -16,39 +19,34 @@ class FastAPIEventService(EventServiceBase):
         self.event_handler_id = event_handler_id
         self.__queue = Queue()
         self.__stop_event = threading.Event()
-        asyncio.create_task(self.__dispatch_from_queue(stop_event = self.__stop_event))
+        asyncio.create_task(self.__dispatch_from_queue(stop_event=self.__stop_event))
 
         super().__init__()
 
-
     def stop(self, *args, **kwargs):
         self.__stop_event.set()
         self.__queue.put(None)
 
-
     def dispatch(self, event_name: str, payload: Any) -> None:
-        self.__queue.put(dict(
-            event_name = event_name,
-            payload = payload
-        ))
-
+        self.__queue.put(dict(event_name=event_name, payload=payload))
 
     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 = self.__queue.get(block = False)
-                if not event: # Probably stopping
+                event = self.__queue.get(block=False)
+                if not event:  # Probably stopping
                     continue
 
                 dispatch(
-                    event.get('event_name'),
-                    payload       = event.get('payload'),
-                    middleware_id = self.event_handler_id)
+                    event.get("event_name"),
+                    payload=event.get("payload"),
+                    middleware_id=self.event_handler_id,
+                )
 
             except Empty:
                 await asyncio.sleep(0.001)
                 pass
 
             except asyncio.CancelledError as e:
-                raise e # Raise a proper error
+                raise e  # Raise a proper error

invokeai/app/api/routers/images.py

@@ -0,0 +1,66 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+from datetime import datetime, timezone
+
+from fastapi import Path, Request, UploadFile
+from fastapi.responses import FileResponse, Response
+from fastapi.routing import APIRouter
+from PIL import Image
+
+from ...services.image_storage import ImageType
+from ..dependencies import ApiDependencies
+
+images_router = APIRouter(prefix="/v1/images", tags=["images"])
+
+
+@images_router.get("/{image_type}/{image_name}", operation_id="get_image")
+async def get_image(
+    image_type: ImageType = Path(description="The type of image to get"),
+    image_name: str = Path(description="The name of the image to get"),
+):
+    """Gets a result"""
+    # TODO: This is not really secure at all. At least make sure only output results are served
+    filename = ApiDependencies.invoker.services.images.get_path(image_type, image_name)
+    return FileResponse(filename)
+
+@images_router.get("/{image_type}/thumbnails/{image_name}", operation_id="get_thumbnail")
+async def get_thumbnail(
+    image_type: ImageType = Path(description="The type of image to get"),
+    image_name: str = Path(description="The name of the image to get"),
+):
+    """Gets a thumbnail"""
+    # TODO: This is not really secure at all. At least make sure only output results are served
+    filename = ApiDependencies.invoker.services.images.get_path(image_type, 'thumbnails/' + image_name)
+    return FileResponse(filename)
+
+
+@images_router.post(
+    "/uploads/",
+    operation_id="upload_image",
+    responses={
+        201: {"description": "The image was uploaded successfully"},
+        404: {"description": "Session not found"},
+    },
+)
+async def upload_image(file: UploadFile, request: Request):
+    if not file.content_type.startswith("image"):
+        return Response(status_code=415)
+
+    contents = await file.read()
+    try:
+        im = Image.open(contents)
+    except:
+        # Error opening the image
+        return Response(status_code=415)
+
+    filename = f"{str(int(datetime.now(timezone.utc).timestamp()))}.png"
+    ApiDependencies.invoker.services.images.save(ImageType.UPLOAD, filename, im)
+
+    return Response(
+        status_code=201,
+        headers={
+            "Location": request.url_for(
+                "get_image", image_type=ImageType.UPLOAD, image_name=filename
+            )
+        },
+    )

invokeai/app/api/routers/models.py

@@ -0,0 +1,279 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from typing import Annotated, Any, List, Literal, Optional, Union
+
+from fastapi.routing import APIRouter
+from pydantic import BaseModel, Field, parse_obj_as
+
+from ..dependencies import ApiDependencies
+
+models_router = APIRouter(prefix="/v1/models", tags=["models"])
+
+
+class VaeRepo(BaseModel):
+    repo_id: str = Field(description="The repo ID to use for this VAE")
+    path: Optional[str] = Field(description="The path to the VAE")
+    subfolder: Optional[str] = Field(description="The subfolder to use for this VAE")
+
+
+class ModelInfo(BaseModel):
+    description: Optional[str] = Field(description="A description of the model")
+    
+
+class CkptModelInfo(ModelInfo):
+    format: Literal['ckpt'] = 'ckpt'
+
+    config: str = Field(description="The path to the model config")
+    weights: str = Field(description="The path to the model weights")
+    vae: str = Field(description="The path to the model VAE")
+    width: Optional[int] = Field(description="The width of the model")
+    height: Optional[int] = Field(description="The height of the model")
+
+
+class DiffusersModelInfo(ModelInfo):
+    format: Literal['diffusers'] = 'diffusers'
+
+    vae: Optional[VaeRepo] = Field(description="The VAE repo to use for this model")
+    repo_id: Optional[str] = Field(description="The repo ID to use for this model")
+    path: Optional[str] = Field(description="The path to the model")
+
+
+class ModelsList(BaseModel):
+    models: dict[str, Annotated[Union[(CkptModelInfo,DiffusersModelInfo)], Field(discriminator="format")]]
+
+
+
+@models_router.get(
+    "/",
+    operation_id="list_models",
+    responses={200: {"model": ModelsList }},
+)
+async def list_models() -> ModelsList:
+    """Gets a list of models"""
+    models_raw = ApiDependencies.invoker.services.model_manager.list_models()
+    models = parse_obj_as(ModelsList, { "models": models_raw })
+    return models
+
+        # @socketio.on("requestSystemConfig")
+        # def handle_request_capabilities():
+        #     print(">> System config requested")
+        #     config = self.get_system_config()
+        #     config["model_list"] = self.generate.model_manager.list_models()
+        #     config["infill_methods"] = infill_methods()
+        #     socketio.emit("systemConfig", config)
+
+        # @socketio.on("searchForModels")
+        # def handle_search_models(search_folder: str):
+        #     try:
+        #         if not search_folder:
+        #             socketio.emit(
+        #                 "foundModels",
+        #                 {"search_folder": None, "found_models": None},
+        #             )
+        #         else:
+        #             (
+        #                 search_folder,
+        #                 found_models,
+        #             ) = self.generate.model_manager.search_models(search_folder)
+        #             socketio.emit(
+        #                 "foundModels",
+        #                 {"search_folder": search_folder, "found_models": found_models},
+        #             )
+        #     except Exception as e:
+        #         self.handle_exceptions(e)
+        #         print("\n")
+
+        # @socketio.on("addNewModel")
+        # def handle_add_model(new_model_config: dict):
+        #     try:
+        #         model_name = new_model_config["name"]
+        #         del new_model_config["name"]
+        #         model_attributes = new_model_config
+        #         if len(model_attributes["vae"]) == 0:
+        #             del model_attributes["vae"]
+        #         update = False
+        #         current_model_list = self.generate.model_manager.list_models()
+        #         if model_name in current_model_list:
+        #             update = True
+
+        #         print(f">> Adding New Model: {model_name}")
+
+        #         self.generate.model_manager.add_model(
+        #             model_name=model_name,
+        #             model_attributes=model_attributes,
+        #             clobber=True,
+        #         )
+        #         self.generate.model_manager.commit(opt.conf)
+
+        #         new_model_list = self.generate.model_manager.list_models()
+        #         socketio.emit(
+        #             "newModelAdded",
+        #             {
+        #                 "new_model_name": model_name,
+        #                 "model_list": new_model_list,
+        #                 "update": update,
+        #             },
+        #         )
+        #         print(f">> New Model Added: {model_name}")
+        #     except Exception as e:
+        #         self.handle_exceptions(e)
+
+        # @socketio.on("deleteModel")
+        # def handle_delete_model(model_name: str):
+        #     try:
+        #         print(f">> Deleting Model: {model_name}")
+        #         self.generate.model_manager.del_model(model_name)
+        #         self.generate.model_manager.commit(opt.conf)
+        #         updated_model_list = self.generate.model_manager.list_models()
+        #         socketio.emit(
+        #             "modelDeleted",
+        #             {
+        #                 "deleted_model_name": model_name,
+        #                 "model_list": updated_model_list,
+        #             },
+        #         )
+        #         print(f">> Model Deleted: {model_name}")
+        #     except Exception as e:
+        #         self.handle_exceptions(e)
+
+        # @socketio.on("requestModelChange")
+        # def handle_set_model(model_name: str):
+        #     try:
+        #         print(f">> Model change requested: {model_name}")
+        #         model = self.generate.set_model(model_name)
+        #         model_list = self.generate.model_manager.list_models()
+        #         if model is None:
+        #             socketio.emit(
+        #                 "modelChangeFailed",
+        #                 {"model_name": model_name, "model_list": model_list},
+        #             )
+        #         else:
+        #             socketio.emit(
+        #                 "modelChanged",
+        #                 {"model_name": model_name, "model_list": model_list},
+        #             )
+        #     except Exception as e:
+        #         self.handle_exceptions(e)
+
+        # @socketio.on("convertToDiffusers")
+        # def convert_to_diffusers(model_to_convert: dict):
+        #     try:
+        #         if model_info := self.generate.model_manager.model_info(
+        #             model_name=model_to_convert["model_name"]
+        #         ):
+        #             if "weights" in model_info:
+        #                 ckpt_path = Path(model_info["weights"])
+        #                 original_config_file = Path(model_info["config"])
+        #                 model_name = model_to_convert["model_name"]
+        #                 model_description = model_info["description"]
+        #             else:
+        #                 self.socketio.emit(
+        #                     "error", {"message": "Model is not a valid checkpoint file"}
+        #                 )
+        #         else:
+        #             self.socketio.emit(
+        #                 "error", {"message": "Could not retrieve model info."}
+        #             )
+
+        #         if not ckpt_path.is_absolute():
+        #             ckpt_path = Path(Globals.root, ckpt_path)
+
+        #         if original_config_file and not original_config_file.is_absolute():
+        #             original_config_file = Path(Globals.root, original_config_file)
+
+        #         diffusers_path = Path(
+        #             ckpt_path.parent.absolute(), f"{model_name}_diffusers"
+        #         )
+
+        #         if model_to_convert["save_location"] == "root":
+        #             diffusers_path = Path(
+        #                 global_converted_ckpts_dir(), f"{model_name}_diffusers"
+        #             )
+
+        #         if (
+        #             model_to_convert["save_location"] == "custom"
+        #             and model_to_convert["custom_location"] is not None
+        #         ):
+        #             diffusers_path = Path(
+        #                 model_to_convert["custom_location"], f"{model_name}_diffusers"
+        #             )
+
+        #         if diffusers_path.exists():
+        #             shutil.rmtree(diffusers_path)
+
+        #         self.generate.model_manager.convert_and_import(
+        #             ckpt_path,
+        #             diffusers_path,
+        #             model_name=model_name,
+        #             model_description=model_description,
+        #             vae=None,
+        #             original_config_file=original_config_file,
+        #             commit_to_conf=opt.conf,
+        #         )
+
+        #         new_model_list = self.generate.model_manager.list_models()
+        #         socketio.emit(
+        #             "modelConverted",
+        #             {
+        #                 "new_model_name": model_name,
+        #                 "model_list": new_model_list,
+        #                 "update": True,
+        #             },
+        #         )
+        #         print(f">> Model Converted: {model_name}")
+        #     except Exception as e:
+        #         self.handle_exceptions(e)
+
+        # @socketio.on("mergeDiffusersModels")
+        # def merge_diffusers_models(model_merge_info: dict):
+        #     try:
+        #         models_to_merge = model_merge_info["models_to_merge"]
+        #         model_ids_or_paths = [
+        #             self.generate.model_manager.model_name_or_path(x)
+        #             for x in models_to_merge
+        #         ]
+        #         merged_pipe = merge_diffusion_models(
+        #             model_ids_or_paths,
+        #             model_merge_info["alpha"],
+        #             model_merge_info["interp"],
+        #             model_merge_info["force"],
+        #         )
+
+        #         dump_path = global_models_dir() / "merged_models"
+        #         if model_merge_info["model_merge_save_path"] is not None:
+        #             dump_path = Path(model_merge_info["model_merge_save_path"])
+
+        #         os.makedirs(dump_path, exist_ok=True)
+        #         dump_path = dump_path / model_merge_info["merged_model_name"]
+        #         merged_pipe.save_pretrained(dump_path, safe_serialization=1)
+
+        #         merged_model_config = dict(
+        #             model_name=model_merge_info["merged_model_name"],
+        #             description=f'Merge of models {", ".join(models_to_merge)}',
+        #             commit_to_conf=opt.conf,
+        #         )
+
+        #         if vae := self.generate.model_manager.config[models_to_merge[0]].get(
+        #             "vae", None
+        #         ):
+        #             print(f">> Using configured VAE assigned to {models_to_merge[0]}")
+        #             merged_model_config.update(vae=vae)
+
+        #         self.generate.model_manager.import_diffuser_model(
+        #             dump_path, **merged_model_config
+        #         )
+        #         new_model_list = self.generate.model_manager.list_models()
+
+        #         socketio.emit(
+        #             "modelsMerged",
+        #             {
+        #                 "merged_models": models_to_merge,
+        #                 "merged_model_name": model_merge_info["merged_model_name"],
+        #                 "model_list": new_model_list,
+        #                 "update": True,
+        #             },
+        #         )
+        #         print(f">> Models Merged: {models_to_merge}")
+        #         print(f">> New Model Added: {model_merge_info['merged_model_name']}")
+        #     except Exception as e:
+        #         self.handle_exceptions(e)

invokeai/app/api/routers/sessions.py

@@ -0,0 +1,287 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+from typing import Annotated, List, Optional, Union
+
+from fastapi import Body, Path, Query
+from fastapi.responses import Response
+from fastapi.routing import APIRouter
+from pydantic.fields import Field
+
+from ...invocations import *
+from ...invocations.baseinvocation import BaseInvocation
+from ...services.graph import (
+    Edge,
+    EdgeConnection,
+    Graph,
+    GraphExecutionState,
+    NodeAlreadyExecutedError,
+)
+from ...services.item_storage import PaginatedResults
+from ..dependencies import ApiDependencies
+
+session_router = APIRouter(prefix="/v1/sessions", tags=["sessions"])
+
+
+@session_router.post(
+    "/",
+    operation_id="create_session",
+    responses={
+        200: {"model": GraphExecutionState},
+        400: {"description": "Invalid json"},
+    },
+)
+async def create_session(
+    graph: Optional[Graph] = Body(
+        default=None, description="The graph to initialize the session with"
+    )
+) -> GraphExecutionState:
+    """Creates a new session, optionally initializing it with an invocation graph"""
+    session = ApiDependencies.invoker.create_execution_state(graph)
+    return session
+
+
+@session_router.get(
+    "/",
+    operation_id="list_sessions",
+    responses={200: {"model": PaginatedResults[GraphExecutionState]}},
+)
+async def list_sessions(
+    page: int = Query(default=0, description="The page of results to get"),
+    per_page: int = Query(default=10, description="The number of results per page"),
+    query: str = Query(default="", description="The query string to search for"),
+) -> PaginatedResults[GraphExecutionState]:
+    """Gets a list of sessions, optionally searching"""
+    if query == "":
+        result = ApiDependencies.invoker.services.graph_execution_manager.list(
+            page, per_page
+        )
+    else:
+        result = ApiDependencies.invoker.services.graph_execution_manager.search(
+            query, page, per_page
+        )
+    return result
+
+
+@session_router.get(
+    "/{session_id}",
+    operation_id="get_session",
+    responses={
+        200: {"model": GraphExecutionState},
+        404: {"description": "Session not found"},
+    },
+)
+async def get_session(
+    session_id: str = Path(description="The id of the session to get"),
+) -> GraphExecutionState:
+    """Gets a session"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        return Response(status_code=404)
+    else:
+        return session
+
+
+@session_router.post(
+    "/{session_id}/nodes",
+    operation_id="add_node",
+    responses={
+        200: {"model": str},
+        400: {"description": "Invalid node or link"},
+        404: {"description": "Session not found"},
+    },
+)
+async def add_node(
+    session_id: str = Path(description="The id of the session"),
+    node: Annotated[
+        Union[BaseInvocation.get_invocations()], Field(discriminator="type") # type: ignore
+    ] = Body(description="The node to add"),
+) -> str:
+    """Adds a node to the graph"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        return Response(status_code=404)
+
+    try:
+        session.add_node(node)
+        ApiDependencies.invoker.services.graph_execution_manager.set(
+            session
+        )  # TODO: can this be done automatically, or add node through an API?
+        return session.id
+    except NodeAlreadyExecutedError:
+        return Response(status_code=400)
+    except IndexError:
+        return Response(status_code=400)
+
+
+@session_router.put(
+    "/{session_id}/nodes/{node_path}",
+    operation_id="update_node",
+    responses={
+        200: {"model": GraphExecutionState},
+        400: {"description": "Invalid node or link"},
+        404: {"description": "Session not found"},
+    },
+)
+async def update_node(
+    session_id: str = Path(description="The id of the session"),
+    node_path: str = Path(description="The path to the node in the graph"),
+    node: Annotated[
+        Union[BaseInvocation.get_invocations()], Field(discriminator="type") # type: ignore
+    ] = Body(description="The new node"),
+) -> GraphExecutionState:
+    """Updates a node in the graph and removes all linked edges"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        return Response(status_code=404)
+
+    try:
+        session.update_node(node_path, node)
+        ApiDependencies.invoker.services.graph_execution_manager.set(
+            session
+        )  # TODO: can this be done automatically, or add node through an API?
+        return session
+    except NodeAlreadyExecutedError:
+        return Response(status_code=400)
+    except IndexError:
+        return Response(status_code=400)
+
+
+@session_router.delete(
+    "/{session_id}/nodes/{node_path}",
+    operation_id="delete_node",
+    responses={
+        200: {"model": GraphExecutionState},
+        400: {"description": "Invalid node or link"},
+        404: {"description": "Session not found"},
+    },
+)
+async def delete_node(
+    session_id: str = Path(description="The id of the session"),
+    node_path: str = Path(description="The path to the node to delete"),
+) -> GraphExecutionState:
+    """Deletes a node in the graph and removes all linked edges"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        return Response(status_code=404)
+
+    try:
+        session.delete_node(node_path)
+        ApiDependencies.invoker.services.graph_execution_manager.set(
+            session
+        )  # TODO: can this be done automatically, or add node through an API?
+        return session
+    except NodeAlreadyExecutedError:
+        return Response(status_code=400)
+    except IndexError:
+        return Response(status_code=400)
+
+
+@session_router.post(
+    "/{session_id}/edges",
+    operation_id="add_edge",
+    responses={
+        200: {"model": GraphExecutionState},
+        400: {"description": "Invalid node or link"},
+        404: {"description": "Session not found"},
+    },
+)
+async def add_edge(
+    session_id: str = Path(description="The id of the session"),
+    edge: Edge = Body(description="The edge to add"),
+) -> GraphExecutionState:
+    """Adds an edge to the graph"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        return Response(status_code=404)
+
+    try:
+        session.add_edge(edge)
+        ApiDependencies.invoker.services.graph_execution_manager.set(
+            session
+        )  # TODO: can this be done automatically, or add node through an API?
+        return session
+    except NodeAlreadyExecutedError:
+        return Response(status_code=400)
+    except IndexError:
+        return Response(status_code=400)
+
+
+# TODO: the edge being in the path here is really ugly, find a better solution
+@session_router.delete(
+    "/{session_id}/edges/{from_node_id}/{from_field}/{to_node_id}/{to_field}",
+    operation_id="delete_edge",
+    responses={
+        200: {"model": GraphExecutionState},
+        400: {"description": "Invalid node or link"},
+        404: {"description": "Session not found"},
+    },
+)
+async def delete_edge(
+    session_id: str = Path(description="The id of the session"),
+    from_node_id: str = Path(description="The id of the node the edge is coming from"),
+    from_field: str = Path(description="The field of the node the edge is coming from"),
+    to_node_id: str = Path(description="The id of the node the edge is going to"),
+    to_field: str = Path(description="The field of the node the edge is going to"),
+) -> GraphExecutionState:
+    """Deletes an edge from the graph"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        return Response(status_code=404)
+
+    try:
+        edge = Edge(
+            source=EdgeConnection(node_id=from_node_id, field=from_field),
+            destination=EdgeConnection(node_id=to_node_id, field=to_field)
+        )
+        session.delete_edge(edge)
+        ApiDependencies.invoker.services.graph_execution_manager.set(
+            session
+        )  # TODO: can this be done automatically, or add node through an API?
+        return session
+    except NodeAlreadyExecutedError:
+        return Response(status_code=400)
+    except IndexError:
+        return Response(status_code=400)
+
+
+@session_router.put(
+    "/{session_id}/invoke",
+    operation_id="invoke_session",
+    responses={
+        200: {"model": None},
+        202: {"description": "The invocation is queued"},
+        400: {"description": "The session has no invocations ready to invoke"},
+        404: {"description": "Session not found"},
+    },
+)
+async def invoke_session(
+    session_id: str = Path(description="The id of the session to invoke"),
+    all: bool = Query(
+        default=False, description="Whether or not to invoke all remaining invocations"
+    ),
+) -> None:
+    """Invokes a session"""
+    session = ApiDependencies.invoker.services.graph_execution_manager.get(session_id)
+    if session is None:
+        return Response(status_code=404)
+
+    if session.is_complete():
+        return Response(status_code=400)
+
+    ApiDependencies.invoker.invoke(session, invoke_all=all)
+    return Response(status_code=202)
+
+
+@session_router.delete(
+    "/{session_id}/invoke",
+    operation_id="cancel_session_invoke",
+    responses={
+        202: {"description": "The invocation is canceled"}
+    },
+)
+async def cancel_session_invoke(
+    session_id: str = Path(description="The id of the session to cancel"),
+) -> None:
+    """Invokes a session"""
+    ApiDependencies.invoker.cancel(session_id)
+    return Response(status_code=202)

invokeai/app/api/sockets.py

@@ -1,36 +1,38 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
 from fastapi import FastAPI
-from fastapi_socketio import SocketManager
 from fastapi_events.handlers.local import local_handler
 from fastapi_events.typing import Event
+from fastapi_socketio import SocketManager
+
 from ..services.events import EventServiceBase
 
+
 class SocketIO:
     __sio: SocketManager
 
     def __init__(self, app: FastAPI):
-        self.__sio = SocketManager(app = app)
-        self.__sio.on('subscribe', handler=self._handle_sub)
-        self.__sio.on('unsubscribe', handler=self._handle_unsub)
-        
+        self.__sio = SocketManager(app=app)
+        self.__sio.on("subscribe", handler=self._handle_sub)
+        self.__sio.on("unsubscribe", handler=self._handle_unsub)
+
         local_handler.register(
-            event_name = EventServiceBase.session_event,
-            _func=self._handle_session_event
+            event_name=EventServiceBase.session_event, _func=self._handle_session_event
         )
 
     async def _handle_session_event(self, event: Event):
         await self.__sio.emit(
-            event = event[1]['event'],
-            data = event[1]['data'],
-            room = event[1]['data']['graph_execution_state_id']
+            event=event[1]["event"],
+            data=event[1]["data"],
+            room=event[1]["data"]["graph_execution_state_id"],
         )
 
     async def _handle_sub(self, sid, data, *args, **kwargs):
-        if 'session' in data:
-            self.__sio.enter_room(sid, data['session'])
-        
+        if "session" in data:
+            self.__sio.enter_room(sid, data["session"])
+
         # @app.sio.on('unsubscribe')
+
     async def _handle_unsub(self, sid, data, *args, **kwargs):
-        if 'session' in data:
-            self.__sio.leave_room(sid, data['session'])
+        if "session" in data:
+            self.__sio.leave_room(sid, data["session"])

invokeai/app/api_app.py

@@ -1,37 +1,37 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
-
 import asyncio
 from inspect import signature
-from fastapi import FastAPI
-from fastapi.openapi.utils import get_openapi
-from fastapi.openapi.docs import get_swagger_ui_html, get_redoc_html
-from fastapi.staticfiles import StaticFiles
-from fastapi_events.middleware import EventHandlerASGIMiddleware
-from fastapi_events.handlers.local import local_handler
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic.schema import schema
+
 import uvicorn
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
+from fastapi.openapi.utils import get_openapi
+from fastapi.staticfiles import StaticFiles
+from fastapi_events.handlers.local import local_handler
+from fastapi_events.middleware import EventHandlerASGIMiddleware
+from pydantic.schema import schema
+
+from ..backend import Args
+from .api.dependencies import ApiDependencies
+from .api.routers import images, sessions, models
 from .api.sockets import SocketIO
 from .invocations import *
 from .invocations.baseinvocation import BaseInvocation
-from .api.routers import images, sessions
-from .api.dependencies import ApiDependencies
-from ..args import Args
 
 # Create the app
 # TODO: create this all in a method so configuration/etc. can be passed in?
-app = FastAPI(
-    title     = "Invoke AI",
-    docs_url  = None,
-    redoc_url = None
-)
+app = FastAPI(title="Invoke AI", docs_url=None, redoc_url=None)
 
 # Add event handler
 event_handler_id: int = id(app)
 app.add_middleware(
     EventHandlerASGIMiddleware,
-    handlers      = [local_handler], # TODO: consider doing this in services to support different configurations
-    middleware_id = event_handler_id)
+    handlers=[
+        local_handler
+    ],  # TODO: consider doing this in services to support different configurations
+    middleware_id=event_handler_id,
+)
 
 # Add CORS
 # TODO: use configuration for this
@@ -48,38 +48,36 @@ socket_io = SocketIO(app)
 
 config = {}
 
+
 # Add startup event to load dependencies
-@app.on_event('startup')
+@app.on_event("startup")
 async def startup_event():
-    args = Args()
-    config = args.parse_args()
+    config = Args()
+    config.parse_args()
 
     ApiDependencies.initialize(
-        args             = args,
-        config           = config,
-        event_handler_id = event_handler_id
+        config=config, event_handler_id=event_handler_id
     )
 
+
 # Shut down threads
-@app.on_event('shutdown')
+@app.on_event("shutdown")
 async def shutdown_event():
     ApiDependencies.shutdown()
 
+
 # Include all routers
 # TODO: REMOVE
 # app.include_router(
 #     invocation.invocation_router,
 #     prefix = '/api')
 
-app.include_router(
-    sessions.session_router,
-    prefix = '/api'
-)
+app.include_router(sessions.session_router, prefix="/api")
+
+app.include_router(images.images_router, prefix="/api")
+
+app.include_router(models.models_router, prefix="/api")
 
-app.include_router(
-    images.images_router,
-    prefix = '/api'
-)
 
 # Build a custom OpenAPI to include all outputs
 # TODO: can outputs be included on metadata of invocation schemas somehow?
@@ -87,10 +85,10 @@ def custom_openapi():
     if app.openapi_schema:
         return app.openapi_schema
     openapi_schema = get_openapi(
-        title       = app.title,
-        description = "An API for invoking AI image operations",
-        version     = "1.0.0",
-        routes      = app.routes
+        title=app.title,
+        description="An API for invoking AI image operations",
+        version="1.0.0",
+        routes=app.routes,
     )
 
     # Add all outputs
@@ -102,12 +100,12 @@ def custom_openapi():
         output_types.add(output_type)
 
     output_schemas = schema(output_types, ref_prefix="#/components/schemas/")
-    for schema_key, output_schema in output_schemas['definitions'].items():
+    for schema_key, output_schema in output_schemas["definitions"].items():
         openapi_schema["components"]["schemas"][schema_key] = output_schema
 
         # TODO: note that we assume the schema_key here is the TYPE.__name__
         # This could break in some cases, figure out a better way to do it
-        output_type_titles[schema_key] = output_schema['title']
+        output_type_titles[schema_key] = output_schema["title"]
 
     # Add a reference to the output type to additionalProperties of the invoker schema
     for invoker in all_invocations:
@@ -115,47 +113,45 @@ def custom_openapi():
         output_type = signature(invoker.invoke).return_annotation
         output_type_title = output_type_titles[output_type.__name__]
         invoker_schema = openapi_schema["components"]["schemas"][invoker_name]
-        outputs_ref = { '$ref': f'#/components/schemas/{output_type_title}' }
-        if 'additionalProperties' not in invoker_schema:
-            invoker_schema['additionalProperties'] = {}
+        outputs_ref = {"$ref": f"#/components/schemas/{output_type_title}"}
+
+        invoker_schema["output"] = outputs_ref
 
-        invoker_schema['additionalProperties']['outputs'] = outputs_ref
-    
     app.openapi_schema = openapi_schema
     return app.openapi_schema
 
+
 app.openapi = custom_openapi
 
 # Override API doc favicons
-app.mount('/static', StaticFiles(directory='static/dream_web'), name='static')
+app.mount("/static", StaticFiles(directory="static/dream_web"), name="static")
+
 
 @app.get("/docs", include_in_schema=False)
 def overridden_swagger():
-	return get_swagger_ui_html(
+    return get_swagger_ui_html(
         openapi_url=app.openapi_url,
         title=app.title,
-        swagger_favicon_url="/static/favicon.ico"
+        swagger_favicon_url="/static/favicon.ico",
     )
 
+
 @app.get("/redoc", include_in_schema=False)
 def overridden_redoc():
-	return get_redoc_html(
+    return get_redoc_html(
         openapi_url=app.openapi_url,
         title=app.title,
-        redoc_favicon_url="/static/favicon.ico"
+        redoc_favicon_url="/static/favicon.ico",
     )
 
+
 def invoke_api():
     # Start our own event loop for eventing usage
     # TODO: determine if there's a better way to do this
     loop = asyncio.new_event_loop()
-    config = uvicorn.Config(
-        app = app,
-        host = "0.0.0.0",
-        port = 9090,
-        loop = loop)
-        # Use access_log to turn off logging
-    
+    config = uvicorn.Config(app=app, host="0.0.0.0", port=9090, loop=loop)
+    # Use access_log to turn off logging
+
     server = uvicorn.Server(config)
     loop.run_until_complete(server.serve())
 

invokeai/app/cli/commands.py

@@ -0,0 +1,239 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from abc import ABC, abstractmethod
+import argparse
+from typing import Any, Callable, Iterable, Literal, get_args, get_origin, get_type_hints
+from pydantic import BaseModel, Field
+import networkx as nx
+import matplotlib.pyplot as plt
+from ..invocations.image import ImageField
+from ..services.graph import GraphExecutionState
+from ..services.invoker import Invoker
+
+
+def add_parsers(
+    subparsers,
+    commands: list[type],
+    command_field: str = "type",
+    exclude_fields: list[str] = ["id", "type"],
+    add_arguments: Callable[[argparse.ArgumentParser], None]|None = None
+    ):
+    """Adds parsers for each command to the subparsers"""
+
+    # Create subparsers for each command
+    for command in commands:
+        hints = get_type_hints(command)
+        cmd_name = get_args(hints[command_field])[0]
+        command_parser = subparsers.add_parser(cmd_name, help=command.__doc__)
+
+        if add_arguments is not None:
+            add_arguments(command_parser)
+
+        # Convert all fields to arguments
+        fields = command.__fields__ # type: ignore
+        for name, field in fields.items():
+            if name in exclude_fields:
+                continue
+
+            if get_origin(field.type_) == Literal:
+                allowed_values = get_args(field.type_)
+                allowed_types = set()
+                for val in allowed_values:
+                    allowed_types.add(type(val))
+                allowed_types_list = list(allowed_types)
+                field_type = allowed_types_list[0] if len(allowed_types) == 1 else Union[allowed_types_list]  # type: ignore
+
+                command_parser.add_argument(
+                    f"--{name}",
+                    dest=name,
+                    type=field_type,
+                    default=field.default if field.default_factory is None else field.default_factory(),
+                    choices=allowed_values,
+                    help=field.field_info.description,
+                )
+            else:
+                command_parser.add_argument(
+                    f"--{name}",
+                    dest=name,
+                    type=field.type_,
+                    default=field.default if field.default_factory is None else field.default_factory(),
+                    help=field.field_info.description,
+                )
+
+
+class CliContext:
+    invoker: Invoker
+    session: GraphExecutionState
+    parser: argparse.ArgumentParser
+    defaults: dict[str, Any]
+
+    def __init__(self, invoker: Invoker, session: GraphExecutionState, parser: argparse.ArgumentParser):
+        self.invoker = invoker
+        self.session = session
+        self.parser = parser
+        self.defaults = dict()
+
+    def get_session(self):
+        self.session = self.invoker.services.graph_execution_manager.get(self.session.id)
+        return self.session
+
+
+class ExitCli(Exception):
+    """Exception to exit the CLI"""
+    pass
+
+
+class BaseCommand(ABC, BaseModel):
+    """A CLI command"""
+
+    # All commands must include a type name like this:
+    # type: Literal['your_command_name'] = 'your_command_name'
+
+    @classmethod
+    def get_all_subclasses(cls):
+        subclasses = []
+        toprocess = [cls]
+        while len(toprocess) > 0:
+            next = toprocess.pop(0)
+            next_subclasses = next.__subclasses__()
+            subclasses.extend(next_subclasses)
+            toprocess.extend(next_subclasses)
+        return subclasses
+
+    @classmethod
+    def get_commands(cls):
+        return tuple(BaseCommand.get_all_subclasses())
+
+    @classmethod
+    def get_commands_map(cls):
+        # Get the type strings out of the literals and into a dictionary
+        return dict(map(lambda t: (get_args(get_type_hints(t)['type'])[0], t),BaseCommand.get_all_subclasses()))
+
+    @abstractmethod
+    def run(self, context: CliContext) -> None:
+        """Run the command. Raise ExitCli to exit."""
+        pass
+
+
+class ExitCommand(BaseCommand):
+    """Exits the CLI"""
+    type: Literal['exit'] = 'exit'
+
+    def run(self, context: CliContext) -> None:
+        raise ExitCli()
+
+
+class HelpCommand(BaseCommand):
+    """Shows help"""
+    type: Literal['help'] = 'help'
+
+    def run(self, context: CliContext) -> None:
+        context.parser.print_help()
+
+
+def get_graph_execution_history(
+    graph_execution_state: GraphExecutionState,
+) -> Iterable[str]:
+    """Gets the history of fully-executed invocations for a graph execution"""
+    return (
+        n
+        for n in reversed(graph_execution_state.executed_history)
+        if n in graph_execution_state.graph.nodes
+    )
+
+
+def get_invocation_command(invocation) -> str:
+    fields = invocation.__fields__.items()
+    type_hints = get_type_hints(type(invocation))
+    command = [invocation.type]
+    for name, field in fields:
+        if name in ["id", "type"]:
+            continue
+
+        # TODO: add links
+
+        # Skip image fields when serializing command
+        type_hint = type_hints.get(name) or None
+        if type_hint is ImageField or ImageField in get_args(type_hint):
+            continue
+
+        field_value = getattr(invocation, name)
+        field_default = field.default
+        if field_value != field_default:
+            if type_hint is str or str in get_args(type_hint):
+                command.append(f'--{name} "{field_value}"')
+            else:
+                command.append(f"--{name} {field_value}")
+
+    return " ".join(command)
+
+
+class HistoryCommand(BaseCommand):
+    """Shows the invocation history"""
+    type: Literal['history'] = 'history'
+
+    # Inputs
+    # fmt: off
+    count: int = Field(default=5, gt=0, description="The number of history entries to show")
+    # fmt: on
+
+    def run(self, context: CliContext) -> None:
+        history = list(get_graph_execution_history(context.get_session()))
+        for i in range(min(self.count, len(history))):
+            entry_id = history[-1 - i]
+            entry = context.get_session().graph.get_node(entry_id)
+            print(f"{entry_id}: {get_invocation_command(entry)}")
+
+
+class SetDefaultCommand(BaseCommand):
+    """Sets a default value for a field"""
+    type: Literal['default'] = 'default'
+
+    # Inputs
+    # fmt: off
+    field: str = Field(description="The field to set the default for")
+    value: str = Field(description="The value to set the default to, or None to clear the default")
+    # fmt: on
+
+    def run(self, context: CliContext) -> None:
+        if self.value is None:
+            if self.field in context.defaults:
+                del context.defaults[self.field]
+        else:
+            context.defaults[self.field] = self.value
+
+
+class DrawGraphCommand(BaseCommand):
+    """Debugs a graph"""
+    type: Literal['draw_graph'] = 'draw_graph'
+
+    def run(self, context: CliContext) -> None:
+        session: GraphExecutionState = context.invoker.services.graph_execution_manager.get(context.session.id)
+        nxgraph = session.graph.nx_graph_flat()
+
+        # Draw the networkx graph
+        plt.figure(figsize=(20, 20))
+        pos = nx.spectral_layout(nxgraph)
+        nx.draw_networkx_nodes(nxgraph, pos, node_size=1000)
+        nx.draw_networkx_edges(nxgraph, pos, width=2)
+        nx.draw_networkx_labels(nxgraph, pos, font_size=20, font_family="sans-serif")
+        plt.axis("off")
+        plt.show()
+
+
+class DrawExecutionGraphCommand(BaseCommand):
+    """Debugs an execution graph"""
+    type: Literal['draw_xgraph'] = 'draw_xgraph'
+
+    def run(self, context: CliContext) -> None:
+        session: GraphExecutionState = context.invoker.services.graph_execution_manager.get(context.session.id)
+        nxgraph = session.execution_graph.nx_graph_flat()
+
+        # Draw the networkx graph
+        plt.figure(figsize=(20, 20))
+        pos = nx.spectral_layout(nxgraph)
+        nx.draw_networkx_nodes(nxgraph, pos, node_size=1000)
+        nx.draw_networkx_edges(nxgraph, pos, width=2)
+        nx.draw_networkx_labels(nxgraph, pos, font_size=20, font_family="sans-serif")
+        plt.axis("off")
+        plt.show()

invokeai/app/cli/completer.py

@@ -0,0 +1,167 @@
+"""
+Readline helper functions for cli_app.py
+You may import the global singleton `completer` to get access to the
+completer object.
+"""
+import atexit
+import readline
+import shlex
+
+from pathlib import Path
+from typing import List, Dict, Literal, get_args, get_type_hints, get_origin
+
+from ...backend import ModelManager, Globals
+from ..invocations.baseinvocation import BaseInvocation
+from .commands import BaseCommand
+
+# singleton object, class variable
+completer = None
+
+class Completer(object):
+    
+    def __init__(self, model_manager: ModelManager):
+        self.commands = self.get_commands()
+        self.matches = None
+        self.linebuffer = None
+        self.manager = model_manager
+        return
+
+    def complete(self, text, state):
+        """
+        Complete commands and switches fromm the node CLI command line.
+        Switches are determined in a context-specific manner.
+        """
+
+        buffer = readline.get_line_buffer()
+        if state == 0:
+            options = None
+            try:
+                current_command, current_switch = self.get_current_command(buffer)
+                options = self.get_command_options(current_command, current_switch)
+            except IndexError:
+                pass
+            options = options or list(self.parse_commands().keys())
+            
+            if not text:  # first time
+                self.matches = options
+            else:
+                self.matches = [s for s in options if s and s.startswith(text)]
+
+        try:
+            match = self.matches[state]
+        except IndexError:
+            match = None
+        return match
+
+    @classmethod
+    def get_commands(self)->List[object]:
+        """
+        Return a list of all the client commands and invocations.
+        """
+        return BaseCommand.get_commands() + BaseInvocation.get_invocations()
+
+    def get_current_command(self, buffer: str)->tuple[str, str]:
+        """
+        Parse the readline buffer to find the most recent command and its switch.
+        """
+        if len(buffer)==0:
+            return None, None
+        tokens = shlex.split(buffer)
+        command = None
+        switch = None
+        for t in tokens:
+            if t[0].isalpha():
+                if switch is None:
+                    command = t
+            else:
+                switch = t
+        # don't try to autocomplete switches that are already complete
+        if switch and buffer.endswith(' '):
+            switch=None
+        return command or '', switch or ''
+
+    def parse_commands(self)->Dict[str, List[str]]:
+        """
+        Return a dict in which the keys are the command name
+        and the values are the parameters the command takes.
+        """
+        result = dict()
+        for command in self.commands:
+            hints = get_type_hints(command)
+            name = get_args(hints['type'])[0]
+            result.update({name:hints})
+        return result
+
+    def get_command_options(self, command: str, switch: str)->List[str]:
+        """
+        Return all the parameters that can be passed to the command as
+        command-line switches. Returns None if the command is unrecognized.
+        """
+        parsed_commands = self.parse_commands()
+        if command not in parsed_commands:
+            return None
+        
+        # handle switches in the format "-foo=bar"
+        argument = None
+        if switch and '=' in switch:
+            switch, argument = switch.split('=')
+            
+        parameter = switch.strip('-')
+        if parameter in parsed_commands[command]:
+            if argument is None:
+                return self.get_parameter_options(parameter, parsed_commands[command][parameter])
+            else:
+                return [f"--{parameter}={x}" for x in self.get_parameter_options(parameter, parsed_commands[command][parameter])]
+        else:
+            return [f"--{x}" for x in parsed_commands[command].keys()]
+
+    def get_parameter_options(self, parameter: str, typehint)->List[str]:
+        """
+        Given a parameter type (such as Literal), offers autocompletions.
+        """
+        if get_origin(typehint) == Literal:
+            return get_args(typehint)
+        if parameter == 'model':
+            return self.manager.model_names()
+        
+    def _pre_input_hook(self):
+        if self.linebuffer:
+            readline.insert_text(self.linebuffer)
+            readline.redisplay()
+            self.linebuffer = None
+    
+def set_autocompleter(model_manager: ModelManager) -> Completer:
+    global completer
+    
+    if completer:
+        return completer
+    
+    completer = Completer(model_manager)
+
+    readline.set_completer(completer.complete)
+    # pyreadline3 does not have a set_auto_history() method
+    try:
+        readline.set_auto_history(True)
+    except:
+        pass
+    readline.set_pre_input_hook(completer._pre_input_hook)
+    readline.set_completer_delims(" ")
+    readline.parse_and_bind("tab: complete")
+    readline.parse_and_bind("set print-completions-horizontally off")
+    readline.parse_and_bind("set page-completions on")
+    readline.parse_and_bind("set skip-completed-text on")
+    readline.parse_and_bind("set show-all-if-ambiguous on")
+
+    histfile = Path(Globals.root, ".invoke_history")
+    try:
+        readline.read_history_file(histfile)
+        readline.set_history_length(1000)
+    except FileNotFoundError:
+        pass
+    except OSError:  # file likely corrupted
+        newname = f"{histfile}.old"
+        print(
+            f"## Your history file {histfile} couldn't be loaded and may be corrupted. Renaming it to {newname}"
+        )
+        histfile.replace(Path(newname))
+    atexit.register(readline.write_history_file, histfile)

invokeai/app/cli_app.py

@@ -0,0 +1,300 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+import argparse
+import os
+import re
+import shlex
+import time
+from typing import (
+    Union,
+    get_type_hints,
+)
+
+from pydantic import BaseModel
+from pydantic.fields import Field
+
+from .services.latent_storage import DiskLatentsStorage, ForwardCacheLatentsStorage
+
+from ..backend import Args
+from .cli.commands import BaseCommand, CliContext, ExitCli, add_parsers, get_graph_execution_history
+from .cli.completer import set_autocompleter
+from .invocations import *
+from .invocations.baseinvocation import BaseInvocation
+from .services.events import EventServiceBase
+from .services.model_manager_initializer import get_model_manager
+from .services.restoration_services import RestorationServices
+from .services.graph import Edge, EdgeConnection, GraphExecutionState, are_connection_types_compatible
+from .services.image_storage import DiskImageStorage
+from .services.invocation_queue import MemoryInvocationQueue
+from .services.invocation_services import InvocationServices
+from .services.invoker import Invoker
+from .services.processor import DefaultInvocationProcessor
+from .services.sqlite import SqliteItemStorage
+
+
+class CliCommand(BaseModel):
+    command: Union[BaseCommand.get_commands() + BaseInvocation.get_invocations()] = Field(discriminator="type")  # type: ignore
+
+
+class InvalidArgs(Exception):
+    pass
+
+
+def add_invocation_args(command_parser):
+    # Add linking capability
+    command_parser.add_argument(
+        "--link",
+        "-l",
+        action="append",
+        nargs=3,
+        help="A link in the format 'source_node source_field dest_field'. source_node can be relative to history (e.g. -1)",
+    )
+
+    command_parser.add_argument(
+        "--link_node",
+        "-ln",
+        action="append",
+        help="A link from all fields in the specified node. Node can be relative to history (e.g. -1)",
+    )
+
+
+def get_command_parser() -> argparse.ArgumentParser:
+    # Create invocation parser
+    parser = argparse.ArgumentParser()
+
+    def exit(*args, **kwargs):
+        raise InvalidArgs
+
+    parser.exit = exit
+    subparsers = parser.add_subparsers(dest="type")
+
+    # Create subparsers for each invocation
+    invocations = BaseInvocation.get_all_subclasses()
+    add_parsers(subparsers, invocations, add_arguments=add_invocation_args)
+
+    # Create subparsers for each command
+    commands = BaseCommand.get_all_subclasses()
+    add_parsers(subparsers, commands, exclude_fields=["type"])
+
+    return parser
+
+
+def generate_matching_edges(
+    a: BaseInvocation, b: BaseInvocation
+) -> list[Edge]:
+    """Generates all possible edges between two invocations"""
+    atype = type(a)
+    btype = type(b)
+
+    aoutputtype = atype.get_output_type()
+
+    afields = get_type_hints(aoutputtype)
+    bfields = get_type_hints(btype)
+
+    matching_fields = set(afields.keys()).intersection(bfields.keys())
+
+    # Remove invalid fields
+    invalid_fields = set(["type", "id"])
+    matching_fields = matching_fields.difference(invalid_fields)
+
+    # Validate types
+    matching_fields = [f for f in matching_fields if are_connection_types_compatible(afields[f], bfields[f])]
+
+    edges = [
+        Edge(
+            source=EdgeConnection(node_id=a.id, field=field),
+            destination=EdgeConnection(node_id=b.id, field=field)
+        )
+        for field in matching_fields
+    ]
+    return edges
+
+
+class SessionError(Exception):
+    """Raised when a session error has occurred"""
+    pass
+
+
+def invoke_all(context: CliContext):
+    """Runs all invocations in the specified session"""
+    context.invoker.invoke(context.session, invoke_all=True)
+    while not context.get_session().is_complete():
+        # Wait some time
+        time.sleep(0.1)
+
+    # Print any errors
+    if context.session.has_error():
+        for n in context.session.errors:
+            print(
+                f"Error in node {n} (source node {context.session.prepared_source_mapping[n]}): {context.session.errors[n]}"
+            )
+        
+        raise SessionError()
+
+
+def invoke_cli():
+    config = Args()
+    config.parse_args()
+    model_manager = get_model_manager(config)
+
+    # This initializes the autocompleter and returns it.
+    # Currently nothing is done with the returned Completer
+    # object, but the object can be used to change autocompletion
+    # behavior on the fly, if desired.
+    completer = set_autocompleter(model_manager)
+
+    events = EventServiceBase()
+
+    output_folder = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "../../../outputs")
+    )
+
+    # TODO: build a file/path manager?
+    db_location = os.path.join(output_folder, "invokeai.db")
+
+    services = InvocationServices(
+        model_manager=model_manager,
+        events=events,
+        latents = ForwardCacheLatentsStorage(DiskLatentsStorage(f'{output_folder}/latents')),
+        images=DiskImageStorage(f'{output_folder}/images'),
+        queue=MemoryInvocationQueue(),
+        graph_execution_manager=SqliteItemStorage[GraphExecutionState](
+            filename=db_location, table_name="graph_executions"
+        ),
+        processor=DefaultInvocationProcessor(),
+        restoration=RestorationServices(config),
+    )
+
+    invoker = Invoker(services)
+    session: GraphExecutionState = invoker.create_execution_state()
+    parser = get_command_parser()
+
+    re_negid = re.compile('^-[0-9]+$')
+
+    # Uncomment to print out previous sessions at startup
+    # print(services.session_manager.list())
+
+    context = CliContext(invoker, session, parser)
+
+    while True:
+        try:
+            cmd_input = input("invoke> ")
+        except (KeyboardInterrupt, EOFError):
+            # Ctrl-c exits
+            break
+
+        try:
+            # Refresh the state of the session
+            history = list(get_graph_execution_history(context.session))
+
+            # Split the command for piping
+            cmds = cmd_input.split("|")
+            start_id = len(history)
+            current_id = start_id
+            new_invocations = list()
+            for cmd in cmds:
+                if cmd is None or cmd.strip() == "":
+                    raise InvalidArgs("Empty command")
+
+                # Parse args to create invocation
+                args = vars(context.parser.parse_args(shlex.split(cmd.strip())))
+
+                # Override defaults
+                for field_name, field_default in context.defaults.items():
+                    if field_name in args:
+                        args[field_name] = field_default
+
+                # Parse invocation
+                args["id"] = current_id
+                command = CliCommand(command=args)
+
+                # Run any CLI commands immediately
+                if isinstance(command.command, BaseCommand):
+                    # Invoke all current nodes to preserve operation order
+                    invoke_all(context)
+
+                    # Run the command
+                    command.command.run(context)
+                    continue
+
+                # Pipe previous command output (if there was a previous command)
+                edges: list[Edge] = list()
+                if len(history) > 0 or current_id != start_id:
+                    from_id = (
+                        history[0] if current_id == start_id else str(current_id - 1)
+                    )
+                    from_node = (
+                        next(filter(lambda n: n[0].id == from_id, new_invocations))[0]
+                        if current_id != start_id
+                        else context.session.graph.get_node(from_id)
+                    )
+                    matching_edges = generate_matching_edges(
+                        from_node, command.command
+                    )
+                    edges.extend(matching_edges)
+
+                # Parse provided links
+                if "link_node" in args and args["link_node"]:
+                    for link in args["link_node"]:
+                        node_id = link
+                        if re_negid.match(node_id):
+                            node_id = str(current_id + int(node_id))
+
+                        link_node = context.session.graph.get_node(node_id)
+                        matching_edges = generate_matching_edges(
+                            link_node, command.command
+                        )
+                        matching_destinations = [e.destination for e in matching_edges]
+                        edges = [e for e in edges if e.destination not in matching_destinations]
+                        edges.extend(matching_edges)
+
+                if "link" in args and args["link"]:
+                    for link in args["link"]:
+                        edges = [e for e in edges if e.destination.node_id != command.command.id or e.destination.field != link[2]]
+
+                        node_id = link[0]
+                        if re_negid.match(node_id):
+                            node_id = str(current_id + int(node_id))
+
+                        edges.append(
+                            Edge(
+                                source=EdgeConnection(node_id=node_id, field=link[1]),
+                                destination=EdgeConnection(
+                                    node_id=command.command.id, field=link[2]
+                                )
+                            )
+                        )
+
+                new_invocations.append((command.command, edges))
+
+                current_id = current_id + 1
+
+                # Add the node to the session
+                context.session.add_node(command.command)
+                for edge in edges:
+                    print(edge)
+                    context.session.add_edge(edge)
+
+            # Execute all remaining nodes
+            invoke_all(context)
+
+        except InvalidArgs:
+            print('Invalid command, use "help" to list commands')
+            continue
+
+        except SessionError:
+            # Start a new session
+            print("Session error: creating a new session")
+            context.session = context.invoker.create_execution_state()
+
+        except ExitCli:
+            break
+
+        except SystemExit:
+            continue
+
+    invoker.stop()
+
+
+if __name__ == "__main__":
+    invoke_cli()

invokeai/app/invocations/__init__.py

@@ -4,5 +4,9 @@ __all__ = []
 
 dirname = os.path.dirname(os.path.abspath(__file__))
 for f in os.listdir(dirname):
-    if f != "__init__.py" and os.path.isfile("%s/%s" % (dirname, f)) and f[-3:] == ".py":
+    if (
+        f != "__init__.py"
+        and os.path.isfile("%s/%s" % (dirname, f))
+        and f[-3:] == ".py"
+    ):
         __all__.append(f[:-3])

invokeai/app/invocations/baseinvocation.py

@@ -3,7 +3,9 @@
 from abc import ABC, abstractmethod
 from inspect import signature
 from typing import get_args, get_type_hints
+
 from pydantic import BaseModel, Field
+
 from ..services.invocation_services import InvocationServices
 
 
@@ -70,5 +72,7 @@ class BaseInvocation(ABC, BaseModel):
     def invoke(self, context: InvocationContext) -> BaseInvocationOutput:
         """Invoke with provided context and return outputs."""
         pass
-
+    
+    #fmt: off
     id: str = Field(description="The id of this node. Must be unique among all nodes.")
+    #fmt: on

invokeai/app/invocations/collections.py

@@ -0,0 +1,50 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from typing import Literal
+
+import cv2 as cv
+import numpy as np
+import numpy.random
+from PIL import Image, ImageOps
+from pydantic import Field
+
+from ..services.image_storage import ImageType
+from .baseinvocation import BaseInvocation, InvocationContext, BaseInvocationOutput
+from .image import ImageField, ImageOutput
+
+
+class IntCollectionOutput(BaseInvocationOutput):
+    """A collection of integers"""
+
+    type: Literal["int_collection"] = "int_collection"
+
+    # Outputs
+    collection: list[int] = Field(default=[], description="The int collection")
+
+
+class RangeInvocation(BaseInvocation):
+    """Creates a range"""
+
+    type: Literal["range"] = "range"
+
+    # Inputs
+    start: int = Field(default=0, description="The start of the range")
+    stop: int = Field(default=10, description="The stop of the range")
+    step: int = Field(default=1, description="The step of the range")
+
+    def invoke(self, context: InvocationContext) -> IntCollectionOutput:
+        return IntCollectionOutput(collection=list(range(self.start, self.stop, self.step)))
+
+
+class RandomRangeInvocation(BaseInvocation):
+    """Creates a collection of random numbers"""
+
+    type: Literal["random_range"] = "random_range"
+
+    # Inputs
+    low: int = Field(default=0, description="The inclusive low value")
+    high: int = Field(default=np.iinfo(np.int32).max, description="The exclusive high value")
+    size: int = Field(default=1, description="The number of values to generate")
+
+    def invoke(self, context: InvocationContext) -> IntCollectionOutput:
+        return IntCollectionOutput(collection=list(numpy.random.randint(self.low, self.high, size=self.size)))

invokeai/app/invocations/cv.py

@@ -1,30 +1,36 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
 from typing import Literal
-import numpy
-from pydantic import Field
-from PIL import Image, ImageOps
+
 import cv2 as cv
-from .image import ImageField, ImageOutput
-from .baseinvocation import BaseInvocation, InvocationContext
+import numpy
+from PIL import Image, ImageOps
+from pydantic import Field
+
 from ..services.image_storage import ImageType
+from .baseinvocation import BaseInvocation, InvocationContext
+from .image import ImageField, ImageOutput
 
 
 class CvInpaintInvocation(BaseInvocation):
     """Simple inpaint using opencv."""
-    type: Literal['cv_inpaint'] = 'cv_inpaint'
+    #fmt: off
+    type: Literal["cv_inpaint"] = "cv_inpaint"
 
     # Inputs
     image: ImageField = Field(default=None, description="The image to inpaint")
     mask: ImageField = Field(default=None, description="The mask to use when inpainting")
+    #fmt: on
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
         mask = context.services.images.get(self.mask.image_type, self.mask.image_name)
 
         # Convert to cv image/mask
         # TODO: consider making these utility functions
-        cv_image = cv.cvtColor(numpy.array(image.convert('RGB')), cv.COLOR_RGB2BGR)
+        cv_image = cv.cvtColor(numpy.array(image.convert("RGB")), cv.COLOR_RGB2BGR)
         cv_mask = numpy.array(ImageOps.invert(mask))
 
         # Inpaint
@@ -35,8 +41,10 @@ class CvInpaintInvocation(BaseInvocation):
         image_inpainted = Image.fromarray(cv.cvtColor(cv_inpainted, cv.COLOR_BGR2RGB))
 
         image_type = ImageType.INTERMEDIATE
-        image_name = context.services.images.create_name(context.graph_execution_state_id, self.id)
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
         context.services.images.save(image_type, image_name, image_inpainted)
         return ImageOutput(
-            image = ImageField(image_type = image_type, image_name = image_name)
+            image=ImageField(image_type=image_type, image_name=image_name)
         )

invokeai/app/invocations/generate.py

@@ -0,0 +1,242 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+from functools import partial
+from typing import Literal, Optional, Union
+
+import numpy as np
+from torch import Tensor
+
+from pydantic import Field
+
+from ..services.image_storage import ImageType
+from .baseinvocation import BaseInvocation, InvocationContext
+from .image import ImageField, ImageOutput
+from ...backend.generator import Txt2Img, Img2Img, Inpaint, InvokeAIGenerator
+from ...backend.stable_diffusion import PipelineIntermediateState
+from ..util.util import diffusers_step_callback_adapter, CanceledException
+
+SAMPLER_NAME_VALUES = Literal[
+    tuple(InvokeAIGenerator.schedulers())
+]
+
+# Text to image
+class TextToImageInvocation(BaseInvocation):
+    """Generates an image using text2img."""
+
+    type: Literal["txt2img"] = "txt2img"
+
+    # Inputs
+    # TODO: consider making prompt optional to enable providing prompt through a link
+    # fmt: off
+    prompt: Optional[str] = Field(description="The prompt to generate an image from")
+    seed:        int = Field(default=-1,ge=-1, le=np.iinfo(np.uint32).max, description="The seed to use (-1 for a random seed)", )
+    steps:       int = Field(default=10, gt=0, description="The number of steps to use to generate the image")
+    width:       int = Field(default=512, multiple_of=64, gt=0, description="The width of the resulting image", )
+    height:      int = Field(default=512, multiple_of=64, gt=0, description="The height of the resulting image", )
+    cfg_scale: float = Field(default=7.5, gt=0, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", )
+    sampler_name: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The sampler to use" )
+    seamless:   bool = Field(default=False, description="Whether or not to generate an image that can tile without seams", )
+    model:       str = Field(default="", description="The model to use (currently ignored)")
+    progress_images: bool = Field(default=False, description="Whether or not to produce progress images during generation",  )
+    # fmt: on
+
+    # TODO: pass this an emitter method or something? or a session for dispatching?
+    def dispatch_progress(
+        self, context: InvocationContext, intermediate_state: PipelineIntermediateState
+    ) -> None:
+        if (context.services.queue.is_canceled(context.graph_execution_state_id)):
+            raise CanceledException
+
+        step = intermediate_state.step
+        if intermediate_state.predicted_original is not None:
+            # Some schedulers report not only the noisy latents at the current timestep,
+            # but also their estimate so far of what the de-noised latents will be.
+            sample = intermediate_state.predicted_original
+        else:
+            sample = intermediate_state.latents
+        
+        diffusers_step_callback_adapter(sample, step, steps=self.steps, id=self.id, context=context)
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        # def step_callback(state: PipelineIntermediateState):
+        #     if (context.services.queue.is_canceled(context.graph_execution_state_id)):
+        #         raise CanceledException
+        #     self.dispatch_progress(context, state.latents, state.step)
+
+        # Handle invalid model parameter
+        # TODO: figure out if this can be done via a validator that uses the model_cache
+        # TODO: How to get the default model name now?
+        #       (right now uses whatever current model is set in model manager)
+        model= context.services.model_manager.get_model()
+        outputs = Txt2Img(model).generate(
+            prompt=self.prompt,
+            step_callback=partial(self.dispatch_progress, context),
+            **self.dict(
+                exclude={"prompt"}
+            ),  # Shorthand for passing all of the parameters above manually
+        )
+        # Outputs is an infinite iterator that will return a new InvokeAIGeneratorOutput object
+        # each time it is called. We only need the first one.
+        generate_output = next(outputs)
+
+        # Results are image and seed, unwrap for now and ignore the seed
+        # TODO: pre-seed?
+        # TODO: can this return multiple results? Should it?
+        image_type = ImageType.RESULT
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
+        context.services.images.save(image_type, image_name, generate_output.image)
+        return ImageOutput(
+            image=ImageField(image_type=image_type, image_name=image_name)
+        )
+
+
+class ImageToImageInvocation(TextToImageInvocation):
+    """Generates an image using img2img."""
+
+    type: Literal["img2img"] = "img2img"
+
+    # Inputs
+    image: Union[ImageField, None] = Field(description="The input image")
+    strength: float = Field(
+        default=0.75, gt=0, le=1, description="The strength of the original image"
+    )
+    fit: bool = Field(
+        default=True,
+        description="Whether or not the result should be fit to the aspect ratio of the input image",
+    )
+
+    def dispatch_progress(
+        self, context: InvocationContext, intermediate_state: PipelineIntermediateState
+    ) -> None:  
+        if (context.services.queue.is_canceled(context.graph_execution_state_id)):
+            raise CanceledException
+
+        step = intermediate_state.step
+        if intermediate_state.predicted_original is not None:
+            # Some schedulers report not only the noisy latents at the current timestep,
+            # but also their estimate so far of what the de-noised latents will be.
+            sample = intermediate_state.predicted_original
+        else:
+            sample = intermediate_state.latents
+
+        diffusers_step_callback_adapter(sample, step, steps=self.steps, id=self.id, context=context)
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = (
+            None
+            if self.image is None
+            else context.services.images.get(
+                self.image.image_type, self.image.image_name
+            )
+        )
+        mask = None
+
+        # Handle invalid model parameter
+        # TODO: figure out if this can be done via a validator that uses the model_cache
+        # TODO: How to get the default model name now?
+        model = context.services.model_manager.get_model()
+        outputs = Img2Img(model).generate(
+                prompt=self.prompt,
+                init_image=image,
+                init_mask=mask,
+                step_callback=partial(self.dispatch_progress, context),
+                **self.dict(
+                    exclude={"prompt", "image", "mask"}
+                ),  # Shorthand for passing all of the parameters above manually
+            )
+
+        # Outputs is an infinite iterator that will return a new InvokeAIGeneratorOutput object
+        # each time it is called. We only need the first one.
+        generator_output = next(outputs)
+
+        result_image = generator_output.image
+
+        # Results are image and seed, unwrap for now and ignore the seed
+        # TODO: pre-seed?
+        # TODO: can this return multiple results? Should it?
+        image_type = ImageType.RESULT
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
+        context.services.images.save(image_type, image_name, result_image)
+        return ImageOutput(
+            image=ImageField(image_type=image_type, image_name=image_name)
+        )
+
+class InpaintInvocation(ImageToImageInvocation):
+    """Generates an image using inpaint."""
+
+    type: Literal["inpaint"] = "inpaint"
+
+    # Inputs
+    mask: Union[ImageField, None] = Field(description="The mask")
+    inpaint_replace: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="The amount by which to replace masked areas with latent noise",
+    )
+
+    def dispatch_progress(
+        self, context: InvocationContext, intermediate_state: PipelineIntermediateState
+    ) -> None:  
+        if (context.services.queue.is_canceled(context.graph_execution_state_id)):
+            raise CanceledException
+
+        step = intermediate_state.step
+        if intermediate_state.predicted_original is not None:
+            # Some schedulers report not only the noisy latents at the current timestep,
+            # but also their estimate so far of what the de-noised latents will be.
+            sample = intermediate_state.predicted_original
+        else:
+            sample = intermediate_state.latents
+
+        diffusers_step_callback_adapter(sample, step, steps=self.steps, id=self.id, context=context)
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = (
+            None
+            if self.image is None
+            else context.services.images.get(
+                self.image.image_type, self.image.image_name
+            )
+        )
+        mask = (
+            None
+            if self.mask is None
+            else context.services.images.get(self.mask.image_type, self.mask.image_name)
+        )
+
+        # Handle invalid model parameter
+        # TODO: figure out if this can be done via a validator that uses the model_cache
+        # TODO: How to get the default model name now?
+        model = context.services.model_manager.get_model()
+        outputs = Inpaint(model).generate(
+                prompt=self.prompt,
+                init_img=image,
+                init_mask=mask,
+                step_callback=partial(self.dispatch_progress, context),
+                **self.dict(
+                    exclude={"prompt", "image", "mask"}
+                ),  # Shorthand for passing all of the parameters above manually
+            )
+
+        # Outputs is an infinite iterator that will return a new InvokeAIGeneratorOutput object
+        # each time it is called. We only need the first one.
+        generator_output = next(outputs)
+
+        result_image = generator_output.image
+
+        # Results are image and seed, unwrap for now and ignore the seed
+        # TODO: pre-seed?
+        # TODO: can this return multiple results? Should it?
+        image_type = ImageType.RESULT
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
+        context.services.images.save(image_type, image_name, result_image)
+        return ImageOutput(
+            image=ImageField(image_type=image_type, image_name=image_name)
+        )

invokeai/app/invocations/image.py

@@ -2,108 +2,156 @@
 
 from datetime import datetime, timezone
 from typing import Literal, Optional
+
 import numpy
-from pydantic import Field, BaseModel
-from PIL import Image, ImageOps, ImageFilter
-from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext
+from PIL import Image, ImageFilter, ImageOps
+from pydantic import BaseModel, Field
+
 from ..services.image_storage import ImageType
 from ..services.invocation_services import InvocationServices
+from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext
 
 
 class ImageField(BaseModel):
     """An image field used for passing image objects between invocations"""
-    image_type: str = Field(default=ImageType.RESULT, description="The type of the image")
+
+    image_type: str = Field(
+        default=ImageType.RESULT, description="The type of the image"
+    )
     image_name: Optional[str] = Field(default=None, description="The name of the image")
 
 
 class ImageOutput(BaseInvocationOutput):
     """Base class for invocations that output an image"""
-    type: Literal['image'] = 'image'
-
-    image: ImageField = Field(default=None, description="The output image")
+    #fmt: off
+    type: Literal["image"] = "image"
+    image:      ImageField = Field(default=None, description="The output image")
+    #fmt: on
 
+    class Config:
+        schema_extra = {
+            'required': [
+                'type',
+                'image',
+            ]
+        }
 
 class MaskOutput(BaseInvocationOutput):
     """Base class for invocations that output a mask"""
-    type: Literal['mask'] = 'mask'
-
-    mask: ImageField = Field(default=None, description="The output mask")
+    #fmt: off
+    type: Literal["mask"] = "mask"
+    mask:      ImageField = Field(default=None, description="The output mask")
+    #fmt: on
 
+    class Config:
+        schema_extra = {
+            'required': [
+                'type',
+                'mask',
+            ]
+        }
 
 # TODO: this isn't really necessary anymore
 class LoadImageInvocation(BaseInvocation):
     """Load an image from a filename and provide it as output."""
-    type: Literal['load_image'] = 'load_image'
+    #fmt: off
+    type: Literal["load_image"] = "load_image"
 
     # Inputs
     image_type: ImageType = Field(description="The type of the image")
-    image_name: str = Field(description="The name of the image")
+    image_name:       str = Field(description="The name of the image")
+    #fmt: on
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
         return ImageOutput(
-            image = ImageField(image_type = self.image_type, image_name = self.image_name)
+            image=ImageField(image_type=self.image_type, image_name=self.image_name)
         )
 
 
 class ShowImageInvocation(BaseInvocation):
     """Displays a provided image, and passes it forward in the pipeline."""
-    type: Literal['show_image'] = 'show_image'
+
+    type: Literal["show_image"] = "show_image"
 
     # Inputs
     image: ImageField = Field(default=None, description="The image to show")
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
         if image:
             image.show()
 
         # TODO: how to handle failure?
 
         return ImageOutput(
-            image = ImageField(image_type = self.image.image_type, image_name = self.image.image_name)
+            image=ImageField(
+                image_type=self.image.image_type, image_name=self.image.image_name
+            )
         )
 
 
 class CropImageInvocation(BaseInvocation):
     """Crops an image to a specified box. The box can be outside of the image."""
-    type: Literal['crop'] = 'crop'
+    #fmt: off
+    type: Literal["crop"] = "crop"
 
     # Inputs
     image: ImageField = Field(default=None, description="The image to crop")
-    x: int      = Field(default=0, description="The left x coordinate of the crop rectangle")
-    y: int      = Field(default=0, description="The top y coordinate of the crop rectangle")
-    width: int  = Field(default=512, gt=0, description="The width of the crop rectangle")
+    x:      int = Field(default=0, description="The left x coordinate of the crop rectangle")
+    y:      int = Field(default=0, description="The top y coordinate of the crop rectangle")
+    width:  int = Field(default=512, gt=0, description="The width of the crop rectangle")
     height: int = Field(default=512, gt=0, description="The height of the crop rectangle")
+    #fmt: on
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
 
-        image_crop = Image.new(mode = 'RGBA', size = (self.width, self.height), color = (0, 0, 0, 0))
+        image_crop = Image.new(
+            mode="RGBA", size=(self.width, self.height), color=(0, 0, 0, 0)
+        )
         image_crop.paste(image, (-self.x, -self.y))
 
         image_type = ImageType.INTERMEDIATE
-        image_name = context.services.images.create_name(context.graph_execution_state_id, self.id)
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
         context.services.images.save(image_type, image_name, image_crop)
         return ImageOutput(
-            image = ImageField(image_type = image_type, image_name = image_name)
+            image=ImageField(image_type=image_type, image_name=image_name)
         )
 
 
 class PasteImageInvocation(BaseInvocation):
     """Pastes an image into another image."""
-    type: Literal['paste'] = 'paste'
+    #fmt: off
+    type: Literal["paste"] = "paste"
 
     # Inputs
-    base_image: ImageField     = Field(default=None, description="The base image")
-    image: ImageField          = Field(default=None, description="The image to paste")
+    base_image:     ImageField = Field(default=None, description="The base image")
+    image:          ImageField = Field(default=None, description="The image to paste")
     mask: Optional[ImageField] = Field(default=None, description="The mask to use when pasting")
-    x: int                     = Field(default=0, description="The left x coordinate at which to paste the image")
-    y: int                     = Field(default=0, description="The top y coordinate at which to paste the image")
+    x:                     int = Field(default=0, description="The left x coordinate at which to paste the image")
+    y:                     int = Field(default=0, description="The top y coordinate at which to paste the image")
+    #fmt: on
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        base_image = context.services.images.get(self.base_image.image_type, self.base_image.image_name)
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
-        mask = None if self.mask is None else ImageOps.invert(services.images.get(self.mask.image_type, self.mask.image_name))
+        base_image = context.services.images.get(
+            self.base_image.image_type, self.base_image.image_name
+        )
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
+        mask = (
+            None
+            if self.mask is None
+            else ImageOps.invert(
+                services.images.get(self.mask.image_type, self.mask.image_name)
+            )
+        )
         # TODO: probably shouldn't invert mask here... should user be required to do it?
 
         min_x = min(0, self.x)
@@ -111,75 +159,98 @@ class PasteImageInvocation(BaseInvocation):
         max_x = max(base_image.width, image.width + self.x)
         max_y = max(base_image.height, image.height + self.y)
 
-        new_image = Image.new(mode = 'RGBA', size = (max_x - min_x, max_y - min_y), color = (0, 0, 0, 0))
+        new_image = Image.new(
+            mode="RGBA", size=(max_x - min_x, max_y - min_y), color=(0, 0, 0, 0)
+        )
         new_image.paste(base_image, (abs(min_x), abs(min_y)))
-        new_image.paste(image, (max(0, self.x), max(0, self.y)), mask = mask)
+        new_image.paste(image, (max(0, self.x), max(0, self.y)), mask=mask)
 
         image_type = ImageType.RESULT
-        image_name = context.services.images.create_name(context.graph_execution_state_id, self.id)
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
         context.services.images.save(image_type, image_name, new_image)
         return ImageOutput(
-            image = ImageField(image_type = image_type, image_name = image_name)
+            image=ImageField(image_type=image_type, image_name=image_name)
         )
 
 
 class MaskFromAlphaInvocation(BaseInvocation):
     """Extracts the alpha channel of an image as a mask."""
-    type: Literal['tomask'] = 'tomask'
+    #fmt: off
+    type: Literal["tomask"] = "tomask"
 
     # Inputs
     image: ImageField = Field(default=None, description="The image to create the mask from")
-    invert: bool = Field(default=False, description="Whether or not to invert the mask")
+    invert:      bool = Field(default=False, description="Whether or not to invert the mask")
+    #fmt: on
 
     def invoke(self, context: InvocationContext) -> MaskOutput:
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
 
         image_mask = image.split()[-1]
         if self.invert:
             image_mask = ImageOps.invert(image_mask)
 
         image_type = ImageType.INTERMEDIATE
-        image_name = context.services.images.create_name(context.graph_execution_state_id, self.id)
-        context.services.images.save(image_type, image_name, image_mask)
-        return MaskOutput(
-            mask = ImageField(image_type = image_type, image_name = image_name)
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
         )
+        context.services.images.save(image_type, image_name, image_mask)
+        return MaskOutput(mask=ImageField(image_type=image_type, image_name=image_name))
 
 
 class BlurInvocation(BaseInvocation):
     """Blurs an image"""
-    type: Literal['blur'] = 'blur'
+
+    #fmt: off
+    type: Literal["blur"] = "blur"
 
     # Inputs
     image: ImageField = Field(default=None, description="The image to blur")
-    radius: float     = Field(default=8.0, ge=0, description="The blur radius")
-    blur_type: Literal['gaussian', 'box'] = Field(default='gaussian', description="The type of blur")
-
+    radius:     float = Field(default=8.0, ge=0, description="The blur radius")
+    blur_type: Literal["gaussian", "box"] = Field(default="gaussian", description="The type of blur")
+    #fmt: on
+    
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
 
-        blur = ImageFilter.GaussianBlur(self.radius) if self.blur_type == 'gaussian' else ImageFilter.BoxBlur(self.radius)
+        blur = (
+            ImageFilter.GaussianBlur(self.radius)
+            if self.blur_type == "gaussian"
+            else ImageFilter.BoxBlur(self.radius)
+        )
         blur_image = image.filter(blur)
 
         image_type = ImageType.INTERMEDIATE
-        image_name = context.services.images.create_name(context.graph_execution_state_id, self.id)
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
         context.services.images.save(image_type, image_name, blur_image)
         return ImageOutput(
-            image = ImageField(image_type = image_type, image_name = image_name)
+            image=ImageField(image_type=image_type, image_name=image_name)
         )
 
 
 class LerpInvocation(BaseInvocation):
     """Linear interpolation of all pixels of an image"""
-    type: Literal['lerp'] = 'lerp'
+    #fmt: off
+    type: Literal["lerp"] = "lerp"
 
     # Inputs
     image: ImageField = Field(default=None, description="The image to lerp")
     min: int = Field(default=0, ge=0, le=255, description="The minimum output value")
     max: int = Field(default=255, ge=0, le=255, description="The maximum output value")
+    #fmt: on
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
 
         image_arr = numpy.asarray(image, dtype=numpy.float32) / 255
         image_arr = image_arr * (self.max - self.min) + self.max
@@ -187,33 +258,46 @@ class LerpInvocation(BaseInvocation):
         lerp_image = Image.fromarray(numpy.uint8(image_arr))
 
         image_type = ImageType.INTERMEDIATE
-        image_name = context.services.images.create_name(context.graph_execution_state_id, self.id)
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
         context.services.images.save(image_type, image_name, lerp_image)
         return ImageOutput(
-            image = ImageField(image_type = image_type, image_name = image_name)
+            image=ImageField(image_type=image_type, image_name=image_name)
         )
 
 
 class InverseLerpInvocation(BaseInvocation):
     """Inverse linear interpolation of all pixels of an image"""
-    type: Literal['ilerp'] = 'ilerp'
+    #fmt: off
+    type: Literal["ilerp"] = "ilerp"
 
     # Inputs
     image: ImageField = Field(default=None, description="The image to lerp")
     min: int = Field(default=0, ge=0, le=255, description="The minimum input value")
     max: int = Field(default=255, ge=0, le=255, description="The maximum input value")
-
+    #fmt: on
+    
     def invoke(self, context: InvocationContext) -> ImageOutput:
-        image = context.services.images.get(self.image.image_type, self.image.image_name)
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
 
         image_arr = numpy.asarray(image, dtype=numpy.float32)
-        image_arr = numpy.minimum(numpy.maximum(image_arr - self.min, 0) / float(self.max - self.min), 1) * 255
+        image_arr = (
+            numpy.minimum(
+                numpy.maximum(image_arr - self.min, 0) / float(self.max - self.min), 1
+            )
+            * 255
+        )
 
         ilerp_image = Image.fromarray(numpy.uint8(image_arr))
 
         image_type = ImageType.INTERMEDIATE
-        image_name = context.services.images.create_name(context.graph_execution_state_id, self.id)
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
         context.services.images.save(image_type, image_name, ilerp_image)
         return ImageOutput(
-            image = ImageField(image_type = image_type, image_name = image_name)
+            image=ImageField(image_type=image_type, image_name=image_name)
         )

invokeai/app/invocations/latent.py

@@ -0,0 +1,321 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from typing import Literal, Optional
+from pydantic import BaseModel, Field
+from torch import Tensor
+import torch
+
+from ...backend.model_management.model_manager import ModelManager
+from ...backend.util.devices import CUDA_DEVICE, torch_dtype
+from ...backend.stable_diffusion.diffusion.shared_invokeai_diffusion import PostprocessingSettings
+from ...backend.image_util.seamless import configure_model_padding
+from ...backend.prompting.conditioning import get_uc_and_c_and_ec
+from ...backend.stable_diffusion.diffusers_pipeline import ConditioningData, StableDiffusionGeneratorPipeline
+from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext
+import numpy as np
+from accelerate.utils import set_seed
+from ..services.image_storage import ImageType
+from .baseinvocation import BaseInvocation, InvocationContext
+from .image import ImageField, ImageOutput
+from ...backend.generator import Generator
+from ...backend.stable_diffusion import PipelineIntermediateState
+from ...backend.util.util import image_to_dataURL
+from diffusers.schedulers import SchedulerMixin as Scheduler
+import diffusers
+from diffusers import DiffusionPipeline
+
+
+class LatentsField(BaseModel):
+    """A latents field used for passing latents between invocations"""
+
+    latents_name: Optional[str] = Field(default=None, description="The name of the latents")
+
+
+class LatentsOutput(BaseInvocationOutput):
+    """Base class for invocations that output latents"""
+    #fmt: off
+    type: Literal["latent_output"] = "latent_output"
+    latents: LatentsField            = Field(default=None, description="The output latents")
+    #fmt: on
+
+class NoiseOutput(BaseInvocationOutput):
+    """Invocation noise output"""
+    #fmt: off
+    type: Literal["noise_output"] = "noise_output"
+    noise: LatentsField            = Field(default=None, description="The output noise")
+    #fmt: on
+
+
+# TODO: this seems like a hack
+scheduler_map = dict(
+    ddim=diffusers.DDIMScheduler,
+    dpmpp_2=diffusers.DPMSolverMultistepScheduler,
+    k_dpm_2=diffusers.KDPM2DiscreteScheduler,
+    k_dpm_2_a=diffusers.KDPM2AncestralDiscreteScheduler,
+    k_dpmpp_2=diffusers.DPMSolverMultistepScheduler,
+    k_euler=diffusers.EulerDiscreteScheduler,
+    k_euler_a=diffusers.EulerAncestralDiscreteScheduler,
+    k_heun=diffusers.HeunDiscreteScheduler,
+    k_lms=diffusers.LMSDiscreteScheduler,
+    plms=diffusers.PNDMScheduler,
+)
+
+
+SAMPLER_NAME_VALUES = Literal[
+    tuple(list(scheduler_map.keys()))
+]
+
+
+def get_scheduler(scheduler_name:str, model: StableDiffusionGeneratorPipeline)->Scheduler:
+    scheduler_class = scheduler_map.get(scheduler_name,'ddim')
+    scheduler = scheduler_class.from_config(model.scheduler.config)
+    # hack copied over from generate.py
+    if not hasattr(scheduler, 'uses_inpainting_model'):
+        scheduler.uses_inpainting_model = lambda: False
+    return scheduler
+
+
+def get_noise(width:int, height:int, device:torch.device, seed:int = 0, latent_channels:int=4, use_mps_noise:bool=False, downsampling_factor:int = 8):
+    # limit noise to only the diffusion image channels, not the mask channels
+    input_channels = min(latent_channels, 4)
+    use_device = "cpu" if (use_mps_noise or device.type == "mps") else device
+    generator = torch.Generator(device=use_device).manual_seed(seed)
+    x = torch.randn(
+        [
+            1,
+            input_channels,
+            height // downsampling_factor,
+            width //  downsampling_factor,
+        ],
+        dtype=torch_dtype(device),
+        device=use_device,
+        generator=generator,
+    ).to(device)
+    # if self.perlin > 0.0:
+    #     perlin_noise = self.get_perlin_noise(
+    #         width // self.downsampling_factor, height // self.downsampling_factor
+    #     )
+    #     x = (1 - self.perlin) * x + self.perlin * perlin_noise
+    return x
+
+
+class NoiseInvocation(BaseInvocation):
+    """Generates latent noise."""
+
+    type: Literal["noise"] = "noise"
+
+    # Inputs
+    seed:        int = Field(default=0, ge=0, le=np.iinfo(np.uint32).max, description="The seed to use", )
+    width:       int = Field(default=512, multiple_of=64, gt=0, description="The width of the resulting noise", )
+    height:      int = Field(default=512, multiple_of=64, gt=0, description="The height of the resulting noise", )
+
+    def invoke(self, context: InvocationContext) -> NoiseOutput:
+        device = torch.device(CUDA_DEVICE)
+        noise = get_noise(self.width, self.height, device, self.seed)
+
+        name = f'{context.graph_execution_state_id}__{self.id}'
+        context.services.latents.set(name, noise)
+        return NoiseOutput(
+            noise=LatentsField(latents_name=name)
+        )
+
+
+# Text to image
+class TextToLatentsInvocation(BaseInvocation):
+    """Generates latents from a prompt."""
+
+    type: Literal["t2l"] = "t2l"
+
+    # Inputs
+    # TODO: consider making prompt optional to enable providing prompt through a link
+    # fmt: off
+    prompt: Optional[str] = Field(description="The prompt to generate an image from")
+    seed:        int = Field(default=-1,ge=-1, le=np.iinfo(np.uint32).max, description="The seed to use (-1 for a random seed)", )
+    noise: Optional[LatentsField] = Field(description="The noise to use")
+    steps:       int = Field(default=10, gt=0, description="The number of steps to use to generate the image")
+    width:       int = Field(default=512, multiple_of=64, gt=0, description="The width of the resulting image", )
+    height:      int = Field(default=512, multiple_of=64, gt=0, description="The height of the resulting image", )
+    cfg_scale: float = Field(default=7.5, gt=0, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", )
+    sampler_name: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The sampler to use" )
+    seamless:   bool = Field(default=False, description="Whether or not to generate an image that can tile without seams", )
+    seamless_axes: str = Field(default="", description="The axes to tile the image on, 'x' and/or 'y'")
+    model:       str = Field(default="", description="The model to use (currently ignored)")
+    progress_images: bool = Field(default=False, description="Whether or not to produce progress images during generation",  )
+    # fmt: on
+
+    # TODO: pass this an emitter method or something? or a session for dispatching?
+    def dispatch_progress(
+        self, context: InvocationContext, sample: Tensor, step: int
+    ) -> None:  
+        # TODO: only output a preview image when requested
+        image = Generator.sample_to_lowres_estimated_image(sample)
+
+        (width, height) = image.size
+        width *= 8
+        height *= 8
+
+        dataURL = image_to_dataURL(image, image_format="JPEG")
+
+        context.services.events.emit_generator_progress(
+            context.graph_execution_state_id,
+            self.id,
+            {
+                "width": width,
+                "height": height,
+                "dataURL": dataURL
+            },
+            step,
+            self.steps,
+        )
+    
+    def get_model(self, model_manager: ModelManager) -> StableDiffusionGeneratorPipeline:
+        model_info = model_manager.get_model(self.model)
+        model_name = model_info['model_name']
+        model_hash = model_info['hash']
+        model: StableDiffusionGeneratorPipeline = model_info['model']
+        model.scheduler = get_scheduler(
+            model=model,
+            scheduler_name=self.sampler_name
+        )
+        
+        if isinstance(model, DiffusionPipeline):
+            for component in [model.unet, model.vae]:
+                configure_model_padding(component,
+                                        self.seamless,
+                                        self.seamless_axes
+                                        )
+        else:
+            configure_model_padding(model,
+                                    self.seamless,
+                                    self.seamless_axes
+                                    )
+
+        return model
+
+
+    def get_conditioning_data(self, model: StableDiffusionGeneratorPipeline) -> ConditioningData:
+        uc, c, extra_conditioning_info = get_uc_and_c_and_ec(self.prompt, model=model)
+        conditioning_data = ConditioningData(
+            uc,
+            c,
+            self.cfg_scale,
+            extra_conditioning_info,
+            postprocessing_settings=PostprocessingSettings(
+                threshold=0.0,#threshold,
+                warmup=0.2,#warmup,
+                h_symmetry_time_pct=None,#h_symmetry_time_pct,
+                v_symmetry_time_pct=None#v_symmetry_time_pct,
+            ),
+        ).add_scheduler_args_if_applicable(model.scheduler, eta=None)#ddim_eta)
+        return conditioning_data
+
+
+    def invoke(self, context: InvocationContext) -> LatentsOutput:
+        noise = context.services.latents.get(self.noise.latents_name)
+
+        def step_callback(state: PipelineIntermediateState):
+            self.dispatch_progress(context, state.latents, state.step)
+
+        model = self.get_model(context.services.model_manager)
+        conditioning_data = self.get_conditioning_data(model)
+
+        # TODO: Verify the noise is the right size
+
+        result_latents, result_attention_map_saver = model.latents_from_embeddings(
+            latents=torch.zeros_like(noise, dtype=torch_dtype(model.device)),
+            noise=noise,
+            num_inference_steps=self.steps,
+            conditioning_data=conditioning_data,
+            callback=step_callback
+        )
+
+        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
+        torch.cuda.empty_cache()
+
+        name = f'{context.graph_execution_state_id}__{self.id}'
+        context.services.latents.set(name, result_latents)
+        return LatentsOutput(
+            latents=LatentsField(latents_name=name)
+        )
+
+
+class LatentsToLatentsInvocation(TextToLatentsInvocation):
+    """Generates latents using latents as base image."""
+
+    type: Literal["l2l"] = "l2l"
+
+    # Inputs
+    latents: Optional[LatentsField] = Field(description="The latents to use as a base image")
+    strength: float = Field(default=0.5, description="The strength of the latents to use")
+
+    def invoke(self, context: InvocationContext) -> LatentsOutput:
+        noise = context.services.latents.get(self.noise.latents_name)
+        latent = context.services.latents.get(self.latents.latents_name)
+
+        def step_callback(state: PipelineIntermediateState):
+            self.dispatch_progress(context, state.latents, state.step)
+
+        model = self.get_model(context.services.model_manager)
+        conditioning_data = self.get_conditioning_data(model)
+
+        # TODO: Verify the noise is the right size
+
+        initial_latents = latent if self.strength < 1.0 else torch.zeros_like(
+            latent, device=model.device, dtype=latent.dtype
+        )
+        
+        timesteps, _ = model.get_img2img_timesteps(
+            self.steps,
+            self.strength,
+            device=model.device,
+        )
+
+        result_latents, result_attention_map_saver = model.latents_from_embeddings(
+            latents=initial_latents,
+            timesteps=timesteps,
+            noise=noise,
+            num_inference_steps=self.steps,
+            conditioning_data=conditioning_data,
+            callback=step_callback
+        )
+
+        # https://discuss.huggingface.co/t/memory-usage-by-later-pipeline-stages/23699
+        torch.cuda.empty_cache()
+
+        name = f'{context.graph_execution_state_id}__{self.id}'
+        context.services.latents.set(name, result_latents)
+        return LatentsOutput(
+            latents=LatentsField(latents_name=name)
+        )
+
+
+# Latent to image
+class LatentsToImageInvocation(BaseInvocation):
+    """Generates an image from latents."""
+
+    type: Literal["l2i"] = "l2i"
+
+    # Inputs
+    latents: Optional[LatentsField] = Field(description="The latents to generate an image from")
+    model: str = Field(default="", description="The model to use")
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        latents = context.services.latents.get(self.latents.latents_name)
+
+        # TODO: this only really needs the vae
+        model_info = context.services.model_manager.get_model(self.model)
+        model: StableDiffusionGeneratorPipeline = model_info['model']
+
+        with torch.inference_mode():
+            np_image = model.decode_latents(latents)
+            image = model.numpy_to_pil(np_image)[0]
+
+            image_type = ImageType.RESULT
+            image_name = context.services.images.create_name(
+                context.graph_execution_state_id, self.id
+            )
+            context.services.images.save(image_type, image_name, image)
+            return ImageOutput(
+                image=ImageField(image_type=image_type, image_name=image_name)
+            )

invokeai/app/invocations/math.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+from datetime import datetime, timezone
+from typing import Literal, Optional
+
+import numpy
+from PIL import Image, ImageFilter, ImageOps
+from pydantic import BaseModel, Field
+
+from ..services.image_storage import ImageType
+from ..services.invocation_services import InvocationServices
+from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext
+
+
+class IntOutput(BaseInvocationOutput):
+    """An integer output"""
+    #fmt: off
+    type: Literal["int_output"] = "int_output"
+    a: int = Field(default=None, description="The output integer")
+    #fmt: on
+
+
+class AddInvocation(BaseInvocation):
+    """Adds two numbers"""
+    #fmt: off
+    type: Literal["add"] = "add"
+    a: int = Field(default=0, description="The first number")
+    b: int = Field(default=0, description="The second number")
+    #fmt: on
+
+    def invoke(self, context: InvocationContext) -> IntOutput:
+        return IntOutput(a=self.a + self.b)
+
+
+class SubtractInvocation(BaseInvocation):
+    """Subtracts two numbers"""
+    #fmt: off
+    type: Literal["sub"] = "sub"
+    a: int = Field(default=0, description="The first number")
+    b: int = Field(default=0, description="The second number")
+    #fmt: on
+
+    def invoke(self, context: InvocationContext) -> IntOutput:
+        return IntOutput(a=self.a - self.b)
+
+
+class MultiplyInvocation(BaseInvocation):
+    """Multiplies two numbers"""
+    #fmt: off
+    type: Literal["mul"] = "mul"
+    a: int = Field(default=0, description="The first number")
+    b: int = Field(default=0, description="The second number")
+    #fmt: on
+
+    def invoke(self, context: InvocationContext) -> IntOutput:
+        return IntOutput(a=self.a * self.b)
+
+
+class DivideInvocation(BaseInvocation):
+    """Divides two numbers"""
+    #fmt: off
+    type: Literal["div"] = "div"
+    a: int = Field(default=0, description="The first number")
+    b: int = Field(default=0, description="The second number")
+    #fmt: on
+
+    def invoke(self, context: InvocationContext) -> IntOutput:
+        return IntOutput(a=int(self.a / self.b))

invokeai/app/invocations/prompt.py

@@ -1,9 +1,22 @@
 from typing import Literal
+
 from pydantic.fields import Field
+
 from .baseinvocation import BaseInvocationOutput
 
+
 class PromptOutput(BaseInvocationOutput):
     """Base class for invocations that output a prompt"""
-    type: Literal['prompt'] = 'prompt'
+    #fmt: off
+    type: Literal["prompt"] = "prompt"
 
     prompt: str = Field(default=None, description="The output prompt")
+    #fmt: on
+
+    class Config:
+        schema_extra = {
+            'required': [
+                'type',
+                'prompt',
+            ]
+        }

invokeai/app/invocations/reconstruct.py

@@ -0,0 +1,42 @@
+from datetime import datetime, timezone
+from typing import Literal, Union
+
+from pydantic import Field
+
+from ..services.image_storage import ImageType
+from ..services.invocation_services import InvocationServices
+from .baseinvocation import BaseInvocation, InvocationContext
+from .image import ImageField, ImageOutput
+
+class RestoreFaceInvocation(BaseInvocation):
+    """Restores faces in an image."""
+    #fmt: off
+    type:  Literal["restore_face"] = "restore_face"
+
+    # Inputs
+    image: Union[ImageField, None] = Field(description="The input image")
+    strength:                float = Field(default=0.75, gt=0, le=1, description="The strength of the restoration"  )
+    #fmt: on
+    
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
+        results = context.services.restoration.upscale_and_reconstruct(
+            image_list=[[image, 0]],
+            upscale=None,
+            strength=self.strength,  # GFPGAN strength
+            save_original=False,
+            image_callback=None,
+        )
+
+        # Results are image and seed, unwrap for now
+        # TODO: can this return multiple results?
+        image_type = ImageType.RESULT
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
+        context.services.images.save(image_type, image_name, results[0][0])
+        return ImageOutput(
+            image=ImageField(image_type=image_type, image_name=image_name)
+        )

invokeai/app/invocations/upscale.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+from datetime import datetime, timezone
+from typing import Literal, Union
+
+from pydantic import Field
+
+from ..services.image_storage import ImageType
+from ..services.invocation_services import InvocationServices
+from .baseinvocation import BaseInvocation, InvocationContext
+from .image import ImageField, ImageOutput
+
+
+class UpscaleInvocation(BaseInvocation):
+    """Upscales an image."""
+    #fmt: off
+    type: Literal["upscale"] = "upscale"
+
+    # Inputs
+    image: Union[ImageField, None] = Field(description="The input image", default=None)
+    strength: float = Field(default=0.75, gt=0, le=1, description="The strength")
+    level: Literal[2, 4] = Field(default=2, description="The upscale level")
+    #fmt: on
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.services.images.get(
+            self.image.image_type, self.image.image_name
+        )
+        results = context.services.restoration.upscale_and_reconstruct(
+            image_list=[[image, 0]],
+            upscale=(self.level, self.strength),
+            strength=0.0,  # GFPGAN strength
+            save_original=False,
+            image_callback=None,
+        )
+
+        # Results are image and seed, unwrap for now
+        # TODO: can this return multiple results?
+        image_type = ImageType.RESULT
+        image_name = context.services.images.create_name(
+            context.graph_execution_state_id, self.id
+        )
+        context.services.images.save(image_type, image_name, results[0][0])
+        return ImageOutput(
+            image=ImageField(image_type=image_type, image_name=image_name)
+        )

invokeai/app/services/events.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+from typing import Any, Dict, TypedDict
+
+ProgressImage = TypedDict(
+    "ProgressImage", {"dataURL": str, "width": int, "height": int}
+)
+
+class EventServiceBase:
+    session_event: str = "session_event"
+
+    """Basic event bus, to have an empty stand-in when not needed"""
+
+    def dispatch(self, event_name: str, payload: Any) -> None:
+        pass
+
+    def __emit_session_event(self, event_name: str, payload: Dict) -> None:
+        self.dispatch(
+            event_name=EventServiceBase.session_event,
+            payload=dict(event=event_name, data=payload),
+        )
+
+    # Define events here for every event in the system.
+    # This will make them easier to integrate until we find a schema generator.
+    def emit_generator_progress(
+        self,
+        graph_execution_state_id: str,
+        invocation_id: str,
+        progress_image: ProgressImage | None,
+        step: int,
+        total_steps: int,
+    ) -> None:
+        """Emitted when there is generation progress"""
+        self.__emit_session_event(
+            event_name="generator_progress",
+            payload=dict(
+                graph_execution_state_id=graph_execution_state_id,
+                invocation_id=invocation_id,
+                progress_image=progress_image,
+                step=step,
+                total_steps=total_steps,
+            ),
+        )
+
+    def emit_invocation_complete(
+        self, graph_execution_state_id: str, invocation_id: str, result: Dict
+    ) -> None:
+        """Emitted when an invocation has completed"""
+        self.__emit_session_event(
+            event_name="invocation_complete",
+            payload=dict(
+                graph_execution_state_id=graph_execution_state_id,
+                invocation_id=invocation_id,
+                result=result,
+            ),
+        )
+
+    def emit_invocation_error(
+        self, graph_execution_state_id: str, invocation_id: str, error: str
+    ) -> None:
+        """Emitted when an invocation has completed"""
+        self.__emit_session_event(
+            event_name="invocation_error",
+            payload=dict(
+                graph_execution_state_id=graph_execution_state_id,
+                invocation_id=invocation_id,
+                error=error,
+            ),
+        )
+
+    def emit_invocation_started(
+        self, graph_execution_state_id: str, invocation_id: str
+    ) -> None:
+        """Emitted when an invocation has started"""
+        self.__emit_session_event(
+            event_name="invocation_started",
+            payload=dict(
+                graph_execution_state_id=graph_execution_state_id,
+                invocation_id=invocation_id,
+            ),
+        )
+
+    def emit_graph_execution_complete(self, graph_execution_state_id: str) -> None:
+        """Emitted when a session has completed all invocations"""
+        self.__emit_session_event(
+            event_name="graph_execution_state_complete",
+            payload=dict(graph_execution_state_id=graph_execution_state_id),
+        )

invokeai/app/services/image_storage.py

@@ -1,20 +1,23 @@
 # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
 
-from abc import ABC, abstractmethod
-from enum import Enum
 import datetime
 import os
+from abc import ABC, abstractmethod
+from enum import Enum
 from pathlib import Path
 from queue import Queue
 from typing import Dict
+
 from PIL.Image import Image
-from ...pngwriter import PngWriter
+from invokeai.app.util.save_thumbnail import save_thumbnail
+
+from invokeai.backend.image_util import PngWriter
 
 
 class ImageType(str, Enum):
-    RESULT = 'results'
-    INTERMEDIATE = 'intermediates'
-    UPLOAD = 'uploads'
+    RESULT = "results"
+    INTERMEDIATE = "intermediates"
+    UPLOAD = "uploads"
 
 
 class ImageStorageBase(ABC):
@@ -38,14 +41,15 @@ class ImageStorageBase(ABC):
         pass
 
     def create_name(self, context_id: str, node_id: str) -> str:
-        return f'{context_id}_{node_id}_{str(int(datetime.datetime.now(datetime.timezone.utc).timestamp()))}.png'
+        return f"{context_id}_{node_id}_{str(int(datetime.datetime.now(datetime.timezone.utc).timestamp()))}.png"
 
 
 class DiskImageStorage(ImageStorageBase):
     """Stores images on disk"""
+
     __output_folder: str
     __pngWriter: PngWriter
-    __cache_ids: Queue # TODO: this is an incredibly naive cache
+    __cache_ids: Queue  # TODO: this is an incredibly naive cache
     __cache: Dict[str, Image]
     __max_cache_size: int
 
@@ -54,13 +58,18 @@ class DiskImageStorage(ImageStorageBase):
         self.__pngWriter = PngWriter(output_folder)
         self.__cache = dict()
         self.__cache_ids = Queue()
-        self.__max_cache_size = 10 # TODO: get this from config
+        self.__max_cache_size = 10  # TODO: get this from config
 
         Path(output_folder).mkdir(parents=True, exist_ok=True)
 
         # TODO: don't hard-code. get/save/delete should maybe take subpath?
         for image_type in ImageType:
-            Path(os.path.join(output_folder, image_type)).mkdir(parents=True, exist_ok=True)
+            Path(os.path.join(output_folder, image_type)).mkdir(
+                parents=True, exist_ok=True
+            )
+            Path(os.path.join(output_folder, image_type, "thumbnails")).mkdir(
+                parents=True, exist_ok=True
+            )
 
     def get(self, image_type: ImageType, image_name: str) -> Image:
         image_path = self.get_path(image_type, image_name)
@@ -79,8 +88,14 @@ class DiskImageStorage(ImageStorageBase):
 
     def save(self, image_type: ImageType, image_name: str, image: Image) -> None:
         image_subpath = os.path.join(image_type, image_name)
-        self.__pngWriter.save_image_and_prompt_to_png(image, "", image_subpath, None) # TODO: just pass full path to png writer
-
+        self.__pngWriter.save_image_and_prompt_to_png(
+            image, "", image_subpath, None
+        )  # TODO: just pass full path to png writer
+        save_thumbnail(
+            image=image,
+            filename=image_name,
+            path=os.path.join(self.__output_folder, image_type, "thumbnails"),
+        )
         image_path = self.get_path(image_type, image_name)
         self.__set_cache(image_path, image)
 
@@ -88,7 +103,7 @@ class DiskImageStorage(ImageStorageBase):
         image_path = self.get_path(image_type, image_name)
         if os.path.exists(image_path):
             os.remove(image_path)
-        
+
         if image_path in self.__cache:
             del self.__cache[image_path]
 
@@ -98,7 +113,9 @@ class DiskImageStorage(ImageStorageBase):
     def __set_cache(self, image_name: str, image: Image):
         if not image_name in self.__cache:
             self.__cache[image_name] = image
-            self.__cache_ids.put(image_name) # TODO: this should refresh position for LRU cache
+            self.__cache_ids.put(
+                image_name
+            )  # TODO: this should refresh position for LRU cache
             if len(self.__cache) > self.__max_cache_size:
                 cache_id = self.__cache_ids.get()
                 del self.__cache[cache_id]

invokeai/app/services/invocation_queue.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+
+from abc import ABC, abstractmethod
+from queue import Queue
+import time
+
+
+# TODO: make this serializable
+class InvocationQueueItem:
+    # session_id: str
+    graph_execution_state_id: str
+    invocation_id: str
+    invoke_all: bool
+    timestamp: float
+
+    def __init__(
+        self,
+        # session_id: str,
+        graph_execution_state_id: str,
+        invocation_id: str,
+        invoke_all: bool = False,
+    ):
+        # self.session_id = session_id
+        self.graph_execution_state_id = graph_execution_state_id
+        self.invocation_id = invocation_id
+        self.invoke_all = invoke_all
+        self.timestamp = time.time()
+
+
+class InvocationQueueABC(ABC):
+    """Abstract base class for all invocation queues"""
+
+    @abstractmethod
+    def get(self) -> InvocationQueueItem:
+        pass
+
+    @abstractmethod
+    def put(self, item: InvocationQueueItem | None) -> None:
+        pass
+
+    @abstractmethod
+    def cancel(self, graph_execution_state_id: str) -> None:
+        pass
+
+    @abstractmethod
+    def is_canceled(self, graph_execution_state_id: str) -> bool:
+        pass
+
+
+class MemoryInvocationQueue(InvocationQueueABC):
+    __queue: Queue
+    __cancellations: dict[str, float]
+
+    def __init__(self):
+        self.__queue = Queue()
+        self.__cancellations = dict()
+
+    def get(self) -> InvocationQueueItem:
+        item = self.__queue.get()
+
+        while isinstance(item, InvocationQueueItem) \
+            and item.graph_execution_state_id in self.__cancellations \
+            and self.__cancellations[item.graph_execution_state_id] > item.timestamp:
+            item = self.__queue.get()
+
+        # Clear old items
+        for graph_execution_state_id in list(self.__cancellations.keys()):
+            if self.__cancellations[graph_execution_state_id] < item.timestamp:
+                del self.__cancellations[graph_execution_state_id]
+
+        return item
+
+    def put(self, item: InvocationQueueItem | None) -> None:
+        self.__queue.put(item)
+
+    def cancel(self, graph_execution_state_id: str) -> None:
+        if graph_execution_state_id not in self.__cancellations:
+            self.__cancellations[graph_execution_state_id] = time.time()
+
+    def is_canceled(self, graph_execution_state_id: str) -> bool:
+        return graph_execution_state_id in self.__cancellations

invokeai/app/services/invocation_services.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
+from invokeai.backend import ModelManager
+
+from .events import EventServiceBase
+from .latent_storage import LatentsStorageBase
+from .image_storage import ImageStorageBase
+from .restoration_services import RestorationServices
+from .invocation_queue import InvocationQueueABC
+from .item_storage import ItemStorageABC
+
+class InvocationServices:
+    """Services that can be used by invocations"""
+
+    events: EventServiceBase
+    latents: LatentsStorageBase
+    images: ImageStorageBase
+    queue: InvocationQueueABC
+    model_manager: ModelManager
+    restoration: RestorationServices
+
+    # NOTE: we must forward-declare any types that include invocations, since invocations can use services
+    graph_execution_manager: ItemStorageABC["GraphExecutionState"]
+    processor: "InvocationProcessorABC"
+
+    def __init__(
+            self,
+            model_manager: ModelManager,
+            events: EventServiceBase,
+            latents: LatentsStorageBase,
+            images: ImageStorageBase,
+            queue: InvocationQueueABC,
+            graph_execution_manager: ItemStorageABC["GraphExecutionState"],
+            processor: "InvocationProcessorABC",
+            restoration: RestorationServices,
+    ):
+        self.model_manager = model_manager
+        self.events = events
+        self.latents = latents
+        self.images = images
+        self.queue = queue
+        self.graph_execution_manager = graph_execution_manager
+        self.processor = processor
+        self.restoration = restoration

invokeai/app/services/invoker.py

@@ -2,11 +2,12 @@
 
 from abc import ABC
 from threading import Event, Thread
-from .graph import Graph, GraphExecutionState
-from .item_storage import ItemStorageABC
+
 from ..invocations.baseinvocation import InvocationContext
-from .invocation_services import InvocationServices
+from .graph import Graph, GraphExecutionState
 from .invocation_queue import InvocationQueueABC, InvocationQueueItem
+from .invocation_services import InvocationServices
+from .item_storage import ItemStorageABC
 
 
 class Invoker:
@@ -14,14 +15,13 @@ class Invoker:
 
     services: InvocationServices
 
-    def __init__(self,
-        services: InvocationServices
-    ):
+    def __init__(self, services: InvocationServices):
         self.services = services
         self._start()
 
-
-    def invoke(self, graph_execution_state: GraphExecutionState, invoke_all: bool = False) -> str|None:
+    def invoke(
+        self, graph_execution_state: GraphExecutionState, invoke_all: bool = False
+    ) -> str | None:
         """Determines the next node to invoke and returns the id of the invoked node, or None if there are no nodes to execute"""
 
         # Get the next invocation
@@ -33,38 +33,39 @@ class Invoker:
         self.services.graph_execution_manager.set(graph_execution_state)
 
         # Queue the invocation
-        print(f'queueing item {invocation.id}')
-        self.services.queue.put(InvocationQueueItem(
-            #session_id    = session.id,
-            graph_execution_state_id = graph_execution_state.id,
-            invocation_id = invocation.id,
-            invoke_all    = invoke_all
-        ))
+        self.services.queue.put(
+            InvocationQueueItem(
+                # session_id    = session.id,
+                graph_execution_state_id=graph_execution_state.id,
+                invocation_id=invocation.id,
+                invoke_all=invoke_all,
+            )
+        )
 
         return invocation.id
 
-
-    def create_execution_state(self, graph: Graph|None = None) -> GraphExecutionState:
+    def create_execution_state(self, graph: Graph | None = None) -> GraphExecutionState:
         """Creates a new execution state for the given graph"""
-        new_state = GraphExecutionState(graph = Graph() if graph is None else graph)
+        new_state = GraphExecutionState(graph=Graph() if graph is None else graph)
         self.services.graph_execution_manager.set(new_state)
         return new_state
-
+    
+    def cancel(self, graph_execution_state_id: str) -> None:
+        """Cancels the given execution state"""
+        self.services.queue.cancel(graph_execution_state_id)
 
     def __start_service(self, service) -> None:
         # Call start() method on any services that have it
-        start_op = getattr(service, 'start', None)
+        start_op = getattr(service, "start", None)
         if callable(start_op):
             start_op(self)
 
-
     def __stop_service(self, service) -> None:
         # Call stop() method on any services that have it
-        stop_op = getattr(service, 'stop', None)
+        stop_op = getattr(service, "stop", None)
         if callable(stop_op):
             stop_op(self)
 
-    
     def _start(self) -> None:
         """Starts the invoker. This is called automatically when the invoker is created."""
         for service in vars(self.services):
@@ -73,7 +74,6 @@ class Invoker:
         for service in vars(self.services):
             self.__start_service(getattr(self.services, service))
 
-
     def stop(self) -> None:
         """Stops the invoker. A new invoker will have to be created to execute further."""
         # First stop all services
@@ -87,4 +87,4 @@ class Invoker:
 
 
 class InvocationProcessorABC(ABC):
-    pass
+    pass

invokeai/app/services/item_storage.py

@@ -1,19 +1,21 @@
+from abc import ABC, abstractmethod
+from typing import Callable, Generic, TypeVar
 
-from typing import Callable, TypeVar, Generic
 from pydantic import BaseModel, Field
 from pydantic.generics import GenericModel
-from abc import ABC, abstractmethod
 
-T = TypeVar('T', bound=BaseModel)
+T = TypeVar("T", bound=BaseModel)
+
 
 class PaginatedResults(GenericModel, Generic[T]):
     """Paginated results"""
-    items: list[T]   = Field(description = "Items")
-    page: int        = Field(description = "Current Page")
-    pages: int       = Field(description = "Total number of pages")
-    per_page: int    = Field(description = "Number of items per page")
-    total: int       = Field(description = "Total number of items in result")
-
+    #fmt: off
+    items: list[T] = Field(description="Items")
+    page: int = Field(description="Current Page")
+    pages: int = Field(description="Total number of pages")
+    per_page: int = Field(description="Number of items per page")
+    total: int = Field(description="Total number of items in result")
+    #fmt: on
 
 class ItemStorageABC(ABC, Generic[T]):
     _on_changed_callbacks: list[Callable[[T], None]]
@@ -24,6 +26,7 @@ class ItemStorageABC(ABC, Generic[T]):
         self._on_deleted_callbacks = list()
 
     """Base item storage class"""
+
     @abstractmethod
     def get(self, item_id: str) -> T:
         pass
@@ -37,7 +40,9 @@ class ItemStorageABC(ABC, Generic[T]):
         pass
 
     @abstractmethod
-    def search(self, query: str, page: int = 0, per_page: int = 10) -> PaginatedResults[T]:
+    def search(
+        self, query: str, page: int = 0, per_page: int = 10
+    ) -> PaginatedResults[T]:
         pass
 
     def on_changed(self, on_changed: Callable[[T], None]) -> None:
@@ -51,7 +56,7 @@ class ItemStorageABC(ABC, Generic[T]):
     def _on_changed(self, item: T) -> None:
         for callback in self._on_changed_callbacks:
             callback(item)
-    
+
     def _on_deleted(self, item_id: str) -> None:
         for callback in self._on_deleted_callbacks:
             callback(item_id)

invokeai/app/services/latent_storage.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
+
+import os
+from abc import ABC, abstractmethod
+from pathlib import Path
+from queue import Queue
+from typing import Dict
+
+import torch
+
+class LatentsStorageBase(ABC):
+    """Responsible for storing and retrieving latents."""
+
+    @abstractmethod
+    def get(self, name: str) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def set(self, name: str, data: torch.Tensor) -> None:
+        pass
+
+    @abstractmethod
+    def delete(self, name: str) -> None:
+        pass
+
+
+class ForwardCacheLatentsStorage(LatentsStorageBase):
+    """Caches the latest N latents in memory, writing-thorugh to and reading from underlying storage"""
+    
+    __cache: Dict[str, torch.Tensor]
+    __cache_ids: Queue
+    __max_cache_size: int
+    __underlying_storage: LatentsStorageBase
+
+    def __init__(self, underlying_storage: LatentsStorageBase, max_cache_size: int = 20):
+        self.__underlying_storage = underlying_storage
+        self.__cache = dict()
+        self.__cache_ids = Queue()
+        self.__max_cache_size = max_cache_size
+
+    def get(self, name: str) -> torch.Tensor:
+        cache_item = self.__get_cache(name)
+        if cache_item is not None:
+            return cache_item
+
+        latent = self.__underlying_storage.get(name)
+        self.__set_cache(name, latent)
+        return latent
+
+    def set(self, name: str, data: torch.Tensor) -> None:
+        self.__underlying_storage.set(name, data)
+        self.__set_cache(name, data)
+
+    def delete(self, name: str) -> None:
+        self.__underlying_storage.delete(name)
+        if name in self.__cache:
+            del self.__cache[name]
+
+    def __get_cache(self, name: str) -> torch.Tensor|None:
+        return None if name not in self.__cache else self.__cache[name]
+
+    def __set_cache(self, name: str, data: torch.Tensor):
+        if not name in self.__cache:
+            self.__cache[name] = data
+            self.__cache_ids.put(name)
+            if self.__cache_ids.qsize() > self.__max_cache_size:
+                self.__cache.pop(self.__cache_ids.get())
+
+
+class DiskLatentsStorage(LatentsStorageBase):
+    """Stores latents in a folder on disk without caching"""
+
+    __output_folder: str
+
+    def __init__(self, output_folder: str):
+        self.__output_folder = output_folder
+        Path(output_folder).mkdir(parents=True, exist_ok=True)
+
+    def get(self, name: str) -> torch.Tensor:
+        latent_path = self.get_path(name)
+        return torch.load(latent_path)
+
+    def set(self, name: str, data: torch.Tensor) -> None:
+        latent_path = self.get_path(name)
+        torch.save(data, latent_path)
+
+    def delete(self, name: str) -> None:
+        latent_path = self.get_path(name)
+        os.remove(latent_path)
+
+    def get_path(self, name: str) -> str:
+        return os.path.join(self.__output_folder, name)
+    

invokeai/app/services/model_manager_initializer.py

@@ -0,0 +1,120 @@
+import os
+import sys
+import torch
+from argparse import Namespace
+from invokeai.backend import Args
+from omegaconf import OmegaConf
+from pathlib import Path
+
+import invokeai.version
+from ...backend import ModelManager
+from ...backend.util import choose_precision, choose_torch_device
+from ...backend import Globals
+
+# TODO: Replace with an abstract class base ModelManagerBase
+def get_model_manager(config: Args) -> ModelManager:
+    if not config.conf:
+        config_file = os.path.join(Globals.root, "configs", "models.yaml")
+        if not os.path.exists(config_file):
+            report_model_error(
+                config, FileNotFoundError(f"The file {config_file} could not be found.")
+            )
+
+    print(f">> {invokeai.version.__app_name__}, version {invokeai.version.__version__}")
+    print(f'>> InvokeAI runtime directory is "{Globals.root}"')
+
+    # these two lines prevent a horrible warning message from appearing
+    # when the frozen CLIP tokenizer is imported
+    import transformers  # type: ignore
+
+    transformers.logging.set_verbosity_error()
+    import diffusers
+
+    diffusers.logging.set_verbosity_error()
+
+    # normalize the config directory relative to root
+    if not os.path.isabs(config.conf):
+        config.conf = os.path.normpath(os.path.join(Globals.root, config.conf))
+
+    if config.embeddings:
+        if not os.path.isabs(config.embedding_path):
+            embedding_path = os.path.normpath(
+                os.path.join(Globals.root, config.embedding_path)
+            )
+        else:
+            embedding_path = config.embedding_path
+    else:
+        embedding_path = None
+
+    # migrate legacy models
+    ModelManager.migrate_models()
+
+    # creating the model manager
+    try:
+        device = torch.device(choose_torch_device())
+        precision = 'float16' if config.precision=='float16' \
+        else 'float32' if config.precision=='float32' \
+        else choose_precision(device)
+        
+        model_manager = ModelManager(
+            OmegaConf.load(config.conf),
+            precision=precision,
+            device_type=device,
+            max_loaded_models=config.max_loaded_models,
+            embedding_path = Path(embedding_path),
+        )
+    except (FileNotFoundError, TypeError, AssertionError) as e:
+        report_model_error(config, e)
+    except (IOError, KeyError) as e:
+        print(f"{e}. Aborting.")
+        sys.exit(-1)
+
+    # try to autoconvert new models
+    # autoimport new .ckpt files
+    if path := config.autoconvert:
+        model_manager.autoconvert_weights(
+            conf_path=config.conf,
+            weights_directory=path,
+        )
+
+    return model_manager
+
+def report_model_error(opt: Namespace, e: Exception):
+    print(f'** An error occurred while attempting to initialize the model: "{str(e)}"')
+    print(
+        "** This can be caused by a missing or corrupted models file, and can sometimes be fixed by (re)installing the models."
+    )
+    yes_to_all = os.environ.get("INVOKE_MODEL_RECONFIGURE")
+    if yes_to_all:
+        print(
+            "** Reconfiguration is being forced by environment variable INVOKE_MODEL_RECONFIGURE"
+        )
+    else:
+        response = input(
+            "Do you want to run invokeai-configure script to select and/or reinstall models? [y] "
+        )
+        if response.startswith(("n", "N")):
+            return
+
+    print("invokeai-configure is launching....\n")
+
+    # Match arguments that were set on the CLI
+    # only the arguments accepted by the configuration script are parsed
+    root_dir = ["--root", opt.root_dir] if opt.root_dir is not None else []
+    config = ["--config", opt.conf] if opt.conf is not None else []
+    previous_config = sys.argv
+    sys.argv = ["invokeai-configure"]
+    sys.argv.extend(root_dir)
+    sys.argv.extend(config.to_dict())
+    if yes_to_all is not None:
+        for arg in yes_to_all.split():
+            sys.argv.append(arg)
+
+    from invokeai.frontend.install import invokeai_configure
+
+    invokeai_configure()
+    # TODO: Figure out how to restart
+    # print('** InvokeAI will now restart')
+    # sys.argv = previous_args
+    # main() # would rather do a os.exec(), but doesn't exist?
+    # sys.exit(0)

invokeai/app/services/processor.py

@@ -0,0 +1,124 @@
+import traceback
+from threading import Event, Thread
+
+from ..invocations.baseinvocation import InvocationContext
+from .invocation_queue import InvocationQueueItem
+from .invoker import InvocationProcessorABC, Invoker
+from ..util.util import CanceledException
+
+class DefaultInvocationProcessor(InvocationProcessorABC):
+    __invoker_thread: Thread
+    __stop_event: Event
+    __invoker: Invoker
+
+    def start(self, invoker) -> None:
+        self.__invoker = invoker
+        self.__stop_event = Event()
+        self.__invoker_thread = Thread(
+            name="invoker_processor",
+            target=self.__process,
+            kwargs=dict(stop_event=self.__stop_event),
+        )
+        self.__invoker_thread.daemon = (
+            True  # TODO: probably better to just not use threads?
+        )
+        self.__invoker_thread.start()
+
+    def stop(self, *args, **kwargs) -> None:
+        self.__stop_event.set()
+
+    def __process(self, stop_event: Event):
+        try:
+            while not stop_event.is_set():
+                queue_item: InvocationQueueItem = self.__invoker.services.queue.get()
+                if not queue_item:  # Probably stopping
+                    continue
+
+                graph_execution_state = (
+                    self.__invoker.services.graph_execution_manager.get(
+                        queue_item.graph_execution_state_id
+                    )
+                )
+                invocation = graph_execution_state.execution_graph.get_node(
+                    queue_item.invocation_id
+                )
+
+                # Send starting event
+                self.__invoker.services.events.emit_invocation_started(
+                    graph_execution_state_id=graph_execution_state.id,
+                    invocation_id=invocation.id,
+                )
+
+                # Invoke
+                try:
+                    outputs = invocation.invoke(
+                        InvocationContext(
+                            services=self.__invoker.services,
+                            graph_execution_state_id=graph_execution_state.id,
+                        )
+                    )
+
+                    # Check queue to see if this is canceled, and skip if so
+                    if self.__invoker.services.queue.is_canceled(
+                        graph_execution_state.id
+                    ):
+                        continue
+
+                    # Save outputs and history
+                    graph_execution_state.complete(invocation.id, outputs)
+
+                    # Save the state changes
+                    self.__invoker.services.graph_execution_manager.set(
+                        graph_execution_state
+                    )
+
+                    # Send complete event
+                    self.__invoker.services.events.emit_invocation_complete(
+                        graph_execution_state_id=graph_execution_state.id,
+                        invocation_id=invocation.id,
+                        result=outputs.dict(),
+                    )
+
+                except KeyboardInterrupt:
+                    pass
+
+                except CanceledException:
+                    pass
+
+                except Exception as e:
+                    error = traceback.format_exc()
+
+                    # Save error
+                    graph_execution_state.set_node_error(invocation.id, error)
+
+                    # Save the state changes
+                    self.__invoker.services.graph_execution_manager.set(
+                        graph_execution_state
+                    )
+
+                    # Send error event
+                    self.__invoker.services.events.emit_invocation_error(
+                        graph_execution_state_id=graph_execution_state.id,
+                        invocation_id=invocation.id,
+                        error=error,
+                    )
+
+                    pass
+                
+                # Check queue to see if this is canceled, and skip if so
+                if self.__invoker.services.queue.is_canceled(
+                    graph_execution_state.id
+                ):
+                    continue
+
+                # Queue any further commands if invoking all
+                is_complete = graph_execution_state.is_complete()
+                if queue_item.invoke_all and not is_complete:
+                    self.__invoker.invoke(graph_execution_state, invoke_all=True)
+                elif is_complete:
+                    self.__invoker.services.events.emit_graph_execution_complete(
+                        graph_execution_state.id
+                    )
+
+        except KeyboardInterrupt:
+            ...  # Log something?

invokeai/app/services/restoration_services.py

@@ -0,0 +1,109 @@
+import sys
+import traceback
+import torch
+from ...backend.restoration import Restoration
+from ...backend.util import choose_torch_device, CPU_DEVICE, MPS_DEVICE
+
+# This should be a real base class for postprocessing functions,
+# but right now we just instantiate the existing gfpgan, esrgan
+# and codeformer functions.
+class RestorationServices:
+    '''Face restoration and upscaling'''
+    
+    def __init__(self,args):
+        try:
+            gfpgan, codeformer, esrgan = None, None, None
+            if args.restore or args.esrgan:
+                restoration = Restoration()
+                if args.restore:
+                    gfpgan, codeformer = restoration.load_face_restore_models(
+                        args.gfpgan_model_path
+                    )
+                else:
+                    print(">> Face restoration disabled")
+                    if args.esrgan:
+                        esrgan = restoration.load_esrgan(args.esrgan_bg_tile)
+                    else:
+                        print(">> Upscaling disabled")
+            else:
+                print(">> Face restoration and upscaling disabled")
+        except (ModuleNotFoundError, ImportError):
+            print(traceback.format_exc(), file=sys.stderr)
+            print(">> You may need to install the ESRGAN and/or GFPGAN modules")
+        self.device = torch.device(choose_torch_device())
+        self.gfpgan = gfpgan
+        self.codeformer = codeformer
+        self.esrgan = esrgan
+
+    # note that this one method does gfpgan and codepath reconstruction, as well as
+    # esrgan upscaling
+    # TO DO: refactor into separate methods
+    def upscale_and_reconstruct(
+        self,
+        image_list,
+        facetool="gfpgan",
+        upscale=None,
+        upscale_denoise_str=0.75,
+        strength=0.0,
+        codeformer_fidelity=0.75,
+        save_original=False,
+        image_callback=None,
+        prefix=None,
+    ):
+        results = []
+        for r in image_list:
+            image, seed = r
+            try:
+                if strength > 0:
+                    if self.gfpgan is not None or self.codeformer is not None:
+                        if facetool == "gfpgan":
+                            if self.gfpgan is None:
+                                print(
+                                    ">> GFPGAN not found. Face restoration is disabled."
+                                )
+                            else:
+                                image = self.gfpgan.process(image, strength, seed)
+                        if facetool == "codeformer":
+                            if self.codeformer is None:
+                                print(
+                                    ">> CodeFormer not found. Face restoration is disabled."
+                                )
+                            else:
+                                cf_device = (
+                                    CPU_DEVICE if self.device == MPS_DEVICE else self.device
+                                )
+                                image = self.codeformer.process(
+                                    image=image,
+                                    strength=strength,
+                                    device=cf_device,
+                                    seed=seed,
+                                    fidelity=codeformer_fidelity,
+                                )
+                    else:
+                        print(">> Face Restoration is disabled.")
+                if upscale is not None:
+                    if self.esrgan is not None:
+                        if len(upscale) < 2:
+                            upscale.append(0.75)
+                        image = self.esrgan.process(
+                            image,
+                            upscale[1],
+                            seed,
+                            int(upscale[0]),
+                            denoise_str=upscale_denoise_str,
+                        )
+                    else:
+                        print(">> ESRGAN is disabled. Image not upscaled.")
+            except Exception as e:
+                print(
+                    f">> Error running RealESRGAN or GFPGAN. Your image was not upscaled.\n{e}"
+                )
+
+            if image_callback is not None:
+                image_callback(image, seed, upscaled=True, use_prefix=prefix)
+            else:
+                r[0] = image
+
+            results.append([image, seed])
+
+        return results

invokeai/app/services/sqlite.py

@@ -1,12 +1,15 @@
 import sqlite3
 from threading import Lock
 from typing import Generic, TypeVar, Union, get_args
+
 from pydantic import BaseModel, parse_raw_as
+
 from .item_storage import ItemStorageABC, PaginatedResults
 
-T = TypeVar('T', bound=BaseModel)
+T = TypeVar("T", bound=BaseModel)
+
+sqlite_memory = ":memory:"
 
-sqlite_memory = ':memory:'
 
 class SqliteItemStorage(ItemStorageABC, Generic[T]):
     _filename: str
@@ -16,15 +19,17 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
     _id_field: str
     _lock: Lock
 
-    def __init__(self, filename: str, table_name: str, id_field: str = 'id'):
+    def __init__(self, filename: str, table_name: str, id_field: str = "id"):
         super().__init__()
 
         self._filename = filename
         self._table_name = table_name
-        self._id_field = id_field # TODO: validate that T has this field
+        self._id_field = id_field  # TODO: validate that T has this field
         self._lock = Lock()
 
-        self._conn = sqlite3.connect(self._filename, check_same_thread=False) # TODO: figure out a better threading solution
+        self._conn = sqlite3.connect(
+            self._filename, check_same_thread=False
+        )  # TODO: figure out a better threading solution
         self._cursor = self._conn.cursor()
 
         self._create_table()
@@ -32,10 +37,14 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
     def _create_table(self):
         try:
             self._lock.acquire()
-            self._cursor.execute(f'''CREATE TABLE IF NOT EXISTS {self._table_name} (
+            self._cursor.execute(
+                f"""CREATE TABLE IF NOT EXISTS {self._table_name} (
                 item TEXT,
-                id TEXT GENERATED ALWAYS AS (json_extract(item, '$.{self._id_field}')) VIRTUAL NOT NULL);''')
-            self._cursor.execute(f'''CREATE UNIQUE INDEX IF NOT EXISTS {self._table_name}_id ON {self._table_name}(id);''')
+                id TEXT GENERATED ALWAYS AS (json_extract(item, '$.{self._id_field}')) VIRTUAL NOT NULL);"""
+            )
+            self._cursor.execute(
+                f"""CREATE UNIQUE INDEX IF NOT EXISTS {self._table_name}_id ON {self._table_name}(id);"""
+            )
         finally:
             self._lock.release()
 
@@ -46,7 +55,11 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
     def set(self, item: T):
         try:
             self._lock.acquire()
-            self._cursor.execute(f'''INSERT OR REPLACE INTO {self._table_name} (item) VALUES (?);''', (item.json(),))
+            self._cursor.execute(
+                f"""INSERT OR REPLACE INTO {self._table_name} (item) VALUES (?);""",
+                (item.json(),),
+            )
+            self._conn.commit()
         finally:
             self._lock.release()
         self._on_changed(item)
@@ -54,7 +67,9 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
     def get(self, id: str) -> Union[T, None]:
         try:
             self._lock.acquire()
-            self._cursor.execute(f'''SELECT item FROM {self._table_name} WHERE id = ?;''', (str(id),))
+            self._cursor.execute(
+                f"""SELECT item FROM {self._table_name} WHERE id = ?;""", (str(id),)
+            )
             result = self._cursor.fetchone()
         finally:
             self._lock.release()
@@ -67,7 +82,10 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
     def delete(self, id: str):
         try:
             self._lock.acquire()
-            self._cursor.execute(f'''DELETE FROM {self._table_name} WHERE id = ?;''', (str(id),))
+            self._cursor.execute(
+                f"""DELETE FROM {self._table_name} WHERE id = ?;""", (str(id),)
+            )
+            self._conn.commit()
         finally:
             self._lock.release()
         self._on_deleted(id)
@@ -75,12 +93,15 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
     def list(self, page: int = 0, per_page: int = 10) -> PaginatedResults[T]:
         try:
             self._lock.acquire()
-            self._cursor.execute(f'''SELECT item FROM {self._table_name} LIMIT ? OFFSET ?;''', (per_page, page * per_page))
+            self._cursor.execute(
+                f"""SELECT item FROM {self._table_name} LIMIT ? OFFSET ?;""",
+                (per_page, page * per_page),
+            )
             result = self._cursor.fetchall()
 
             items = list(map(lambda r: self._parse_item(r[0]), result))
 
-            self._cursor.execute(f'''SELECT count(*) FROM {self._table_name};''')
+            self._cursor.execute(f"""SELECT count(*) FROM {self._table_name};""")
             count = self._cursor.fetchone()[0]
         finally:
             self._lock.release()
@@ -88,22 +109,26 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
         pageCount = int(count / per_page) + 1
 
         return PaginatedResults[T](
-            items = items,
-            page = page,
-            pages = pageCount,
-            per_page = per_page,
-            total = count
+            items=items, page=page, pages=pageCount, per_page=per_page, total=count
         )
-    
-    def search(self, query: str, page: int = 0, per_page: int = 10) -> PaginatedResults[T]:
+
+    def search(
+        self, query: str, page: int = 0, per_page: int = 10
+    ) -> PaginatedResults[T]:
         try:
             self._lock.acquire()
-            self._cursor.execute(f'''SELECT item FROM {self._table_name} WHERE item LIKE ? LIMIT ? OFFSET ?;''', (f'%{query}%', per_page, page * per_page))
+            self._cursor.execute(
+                f"""SELECT item FROM {self._table_name} WHERE item LIKE ? LIMIT ? OFFSET ?;""",
+                (f"%{query}%", per_page, page * per_page),
+            )
             result = self._cursor.fetchall()
 
             items = list(map(lambda r: self._parse_item(r[0]), result))
 
-            self._cursor.execute(f'''SELECT count(*) FROM {self._table_name} WHERE item LIKE ?;''', (f'%{query}%',))
+            self._cursor.execute(
+                f"""SELECT count(*) FROM {self._table_name} WHERE item LIKE ?;""",
+                (f"%{query}%",),
+            )
             count = self._cursor.fetchone()[0]
         finally:
             self._lock.release()
@@ -111,9 +136,5 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
         pageCount = int(count / per_page) + 1
 
         return PaginatedResults[T](
-            items = items,
-            page = page,
-            pages = pageCount,
-            per_page = per_page,
-            total = count
+            items=items, page=page, pages=pageCount, per_page=per_page, total=count
         )

invokeai/app/util/save_thumbnail.py

@@ -0,0 +1,25 @@
+import os
+from PIL import Image
+
+
+def save_thumbnail(
+    image: Image.Image,
+    filename: str,
+    path: str,
+    size: int = 256,
+) -> str:
+    """
+    Saves a thumbnail of an image, returning its path.
+    """
+    base_filename = os.path.splitext(filename)[0]
+    thumbnail_path = os.path.join(path, base_filename + ".webp")
+
+    if os.path.exists(thumbnail_path):
+        return thumbnail_path
+
+    image_copy = image.copy()
+    image_copy.thumbnail(size=(size, size))
+
+    image_copy.save(thumbnail_path, "WEBP")
+
+    return thumbnail_path

invokeai/app/util/util.py

@@ -0,0 +1,42 @@
+import torch
+from PIL import Image
+from ..invocations.baseinvocation import InvocationContext
+from ...backend.util.util import image_to_dataURL
+from ...backend.generator.base import Generator
+from ...backend.stable_diffusion import PipelineIntermediateState
+
+class CanceledException(Exception):
+    pass
+
+def fast_latents_step_callback(sample: torch.Tensor, step: int, steps: int, id: str, context: InvocationContext, ):
+    # TODO: only output a preview image when requested
+    image = Generator.sample_to_lowres_estimated_image(sample)
+
+    (width, height) = image.size
+    width *= 8
+    height *= 8
+
+    dataURL = image_to_dataURL(image, image_format="JPEG")
+
+    context.services.events.emit_generator_progress(
+        context.graph_execution_state_id,
+        id,
+        {
+            "width": width,
+            "height": height,
+            "dataURL": dataURL
+        },
+        step,
+        steps,
+    )
+
+def diffusers_step_callback_adapter(*cb_args, **kwargs):
+    """
+    txt2img gives us a Tensor in the step_callbak, while img2img gives us a PipelineIntermediateState.
+    This adapter grabs the needed data and passes it along to the callback function.
+    """
+    if isinstance(cb_args[0], PipelineIntermediateState):
+        progress_state: PipelineIntermediateState = cb_args[0]
+        return fast_latents_step_callback(progress_state.latents, progress_state.step, **kwargs)
+    else:
+        return fast_latents_step_callback(*cb_args, **kwargs)

invokeai/backend/__init__.py

@@ -1,5 +1,16 @@
-'''
+"""
 Initialization file for invokeai.backend
-'''
-from .invoke_ai_web_server import InvokeAIWebServer
-
+"""
+from .generate import Generate
+from .generator import (
+    InvokeAIGeneratorBasicParams,
+    InvokeAIGenerator,
+    InvokeAIGeneratorOutput,
+    Txt2Img,
+    Img2Img,
+    Inpaint
+)
+from .model_management import ModelManager
+from .safety_checker import SafetyChecker
+from .args import Args
+from .globals import Globals

invokeai/backend/config/invokeai_configure.py

@@ -6,19 +6,20 @@
 #
 # Coauthor: Kevin Turner http://github.com/keturn
 #
-print("Loading Python libraries...\n")
+import sys
+print("Loading Python libraries...\n",file=sys.stderr)
+
 import argparse
 import io
 import os
 import re
 import shutil
-import sys
 import traceback
 import warnings
 from argparse import Namespace
 from pathlib import Path
-from urllib import request
 from shutil import get_terminal_size
+from urllib import request
 
 import npyscreen
 import torch
@@ -37,17 +38,20 @@ from transformers import (
 
 import invokeai.configs as configs
 
+from ...frontend.install.model_install import addModelsForm, process_and_execute
+from ...frontend.install.widgets import (
+    CenteredButtonPress,
+    IntTitleSlider,
+    set_min_terminal_size,
+)
 from ..args import PRECISION_CHOICES, Args
-from ..globals import Globals, global_config_dir, global_config_file, global_cache_dir
-from .model_install import addModelsForm, process_and_execute
+from ..globals import Globals, global_cache_dir, global_config_dir, global_config_file
 from .model_install_backend import (
     default_dataset,
     download_from_hf,
-    recommended_datasets,
     hf_download_with_resume,
+    recommended_datasets,
 )
-from .widgets import IntTitleSlider, CenteredButtonPress, set_min_terminal_size
-
 
 warnings.filterwarnings("ignore")
 
@@ -82,6 +86,7 @@ INIT_FILE_PREAMBLE = """# InvokeAI initialization file
 # -Ak_euler_a -C10.0
 """
 
+
 # --------------------------------------------
 def postscript(errors: None):
     if not any(errors):
@@ -180,13 +185,11 @@ def download_with_progress_bar(model_url: str, model_dest: str, label: str = "th
 # ---------------------------------------------
 # this will preload the Bert tokenizer fles
 def download_bert():
-    print(
-        "Installing bert tokenizer...",
-        file=sys.stderr
-    )
+    print("Installing bert tokenizer...", file=sys.stderr)
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", category=DeprecationWarning)
         from transformers import BertTokenizerFast
+
         download_from_hf(BertTokenizerFast, "bert-base-uncased")
 
 
@@ -197,12 +200,14 @@ def download_sd1_clip():
     download_from_hf(CLIPTokenizer, version)
     download_from_hf(CLIPTextModel, version)
 
+
 # ---------------------------------------------
 def download_sd2_clip():
-    version = 'stabilityai/stable-diffusion-2'
+    version = "stabilityai/stable-diffusion-2"
     print("Installing SD2 clip model...", file=sys.stderr)
-    download_from_hf(CLIPTokenizer, version, subfolder='tokenizer')
-    download_from_hf(CLIPTextModel, version, subfolder='text_encoder')
+    download_from_hf(CLIPTokenizer, version, subfolder="tokenizer")
+    download_from_hf(CLIPTextModel, version, subfolder="text_encoder")
+
 
 # ---------------------------------------------
 def download_realesrgan():
@@ -290,7 +295,7 @@ def download_vaes():
         # first the diffusers version
         repo_id = "stabilityai/sd-vae-ft-mse"
         args = dict(
-            cache_dir=global_cache_dir("diffusers"),
+            cache_dir=global_cache_dir("hub"),
         )
         if not AutoencoderKL.from_pretrained(repo_id, **args):
             raise Exception(f"download of {repo_id} failed")
@@ -323,13 +328,13 @@ def get_root(root: str = None) -> str:
 class editOptsForm(npyscreen.FormMultiPage):
     # for responsive resizing - disabled
     # FIX_MINIMUM_SIZE_WHEN_CREATED = False
-    
+
     def create(self):
         program_opts = self.parentApp.program_opts
         old_opts = self.parentApp.invokeai_opts
         first_time = not (Globals.root / Globals.initfile).exists()
         access_token = HfFolder.get_token()
-        window_width,window_height = get_terminal_size()
+        window_width, window_height = get_terminal_size()
         for i in [
             "Configure startup settings. You can come back and change these later.",
             "Use ctrl-N and ctrl-P to move to the <N>ext and <P>revious fields.",
@@ -681,6 +686,7 @@ def run_console_ui(
     else:
         return (editApp.new_opts, editApp.user_selections)
 
+
 # -------------------------------------
 def write_opts(opts: Namespace, init_file: Path):
     """
@@ -701,8 +707,8 @@ def write_opts(opts: Namespace, init_file: Path):
         "^--?(o|out|no-xformer|xformer|no-ckpt|ckpt|free|no-nsfw|nsfw|prec|max_load|embed|always|ckpt|free_gpu)"
     )
     # fix windows paths
-    opts.outdir = opts.outdir.replace('\\','/')
-    opts.embedding_path = opts.embedding_path.replace('\\','/')
+    opts.outdir = opts.outdir.replace("\\", "/")
+    opts.embedding_path = opts.embedding_path.replace("\\", "/")
     new_file = f"{init_file}.new"
     try:
         lines = [x.strip() for x in open(init_file, "r").readlines()]
@@ -855,6 +861,7 @@ def main():
     except KeyboardInterrupt:
         print("\nGoodbye! Come back soon.")
 
+
 # -------------------------------------
 if __name__ == "__main__":
     main()

invokeai/backend/config/model_install_backend.py

@@ -8,6 +8,7 @@ import sys
 import warnings
 from pathlib import Path
 from tempfile import TemporaryFile
+from typing import List
 
 import requests
 from diffusers import AutoencoderKL
@@ -15,12 +16,12 @@ from huggingface_hub import hf_hub_url
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
 from tqdm import tqdm
-from typing import List
 
 import invokeai.configs as configs
-from ..generator.diffusers_pipeline import StableDiffusionGeneratorPipeline
+
 from ..globals import Globals, global_cache_dir, global_config_dir
-from ..model_manager import ModelManager
+from ..model_management import ModelManager
+from ..stable_diffusion import StableDiffusionGeneratorPipeline
 
 warnings.filterwarnings("ignore")
 
@@ -44,45 +45,49 @@ Config_preamble = """
 # was trained on.
 """
 
+
 def default_config_file():
     return Path(global_config_dir()) / "models.yaml"
 
+
 def sd_configs():
     return Path(global_config_dir()) / "stable-diffusion"
 
+
 def initial_models():
     global Datasets
     if Datasets:
         return Datasets
     return (Datasets := OmegaConf.load(Dataset_path))
 
+
 def install_requested_models(
-        install_initial_models: List[str] = None,
-        remove_models: List[str] = None,
-        scan_directory: Path = None,
-        external_models: List[str] = None,
-        scan_at_startup: bool = False,
-        convert_to_diffusers: bool = False,
-        precision: str = "float16",
-        purge_deleted: bool = False,
-        config_file_path: Path = None,
+    install_initial_models: List[str] = None,
+    remove_models: List[str] = None,
+    scan_directory: Path = None,
+    external_models: List[str] = None,
+    scan_at_startup: bool = False,
+    convert_to_diffusers: bool = False,
+    precision: str = "float16",
+    purge_deleted: bool = False,
+    config_file_path: Path = None,
 ):
-    '''
+    """
     Entry point for installing/deleting starter models, or installing external models.
-    '''
-    config_file_path=config_file_path or default_config_file()
+    """
+    config_file_path = config_file_path or default_config_file()
     if not config_file_path.exists():
-        open(config_file_path,'w')
-            
-    model_manager= ModelManager(OmegaConf.load(config_file_path),precision=precision)
-    
+        open(config_file_path, "w")
+
+    model_manager = ModelManager(OmegaConf.load(config_file_path), precision=precision)
+
     if remove_models and len(remove_models) > 0:
         print("== DELETING UNCHECKED STARTER MODELS ==")
         for model in remove_models:
-            print(f'{model}...')
+            print(f"{model}...")
             model_manager.del_model(model, delete_files=purge_deleted)
         model_manager.commit(config_file_path)
-    
+
     if install_initial_models and len(install_initial_models) > 0:
         print("== INSTALLING SELECTED STARTER MODELS ==")
         successfully_downloaded = download_weight_datasets(
@@ -96,20 +101,20 @@ def install_requested_models(
 
     # due to above, we have to reload the model manager because conf file
     # was changed behind its back
-    model_manager= ModelManager(OmegaConf.load(config_file_path),precision=precision)
+    model_manager = ModelManager(OmegaConf.load(config_file_path), precision=precision)
 
     external_models = external_models or list()
     if scan_directory:
         external_models.append(str(scan_directory))
 
-    if len(external_models)>0:
+    if len(external_models) > 0:
         print("== INSTALLING EXTERNAL MODELS ==")
         for path_url_or_repo in external_models:
             try:
                 model_manager.heuristic_import(
                     path_url_or_repo,
                     convert=convert_to_diffusers,
-                    commit_to_conf=config_file_path
+                    commit_to_conf=config_file_path,
                 )
             except KeyboardInterrupt:
                 sys.exit(-1)
@@ -117,17 +122,18 @@ def install_requested_models(
                 pass
 
     if scan_at_startup and scan_directory.is_dir():
-        argument = '--autoconvert' if convert_to_diffusers else '--autoimport'
+        argument = "--autoconvert" if convert_to_diffusers else "--autoimport"
         initfile = Path(Globals.root, Globals.initfile)
-        replacement = Path(Globals.root, f'{Globals.initfile}.new')
-        directory = str(scan_directory).replace('\\','/')
-        with open(initfile,'r') as input:
-            with open(replacement,'w') as output:
+        replacement = Path(Globals.root, f"{Globals.initfile}.new")
+        directory = str(scan_directory).replace("\\", "/")
+        with open(initfile, "r") as input:
+            with open(replacement, "w") as output:
                 while line := input.readline():
                     if not line.startswith(argument):
                         output.writelines([line])
-                output.writelines([f'{argument} {directory}'])
-        os.replace(replacement,initfile)
+                output.writelines([f"{argument} {directory}"])
+        os.replace(replacement, initfile)
+
 
 # -------------------------------------
 def yes_or_no(prompt: str, default_yes=True):
@@ -183,7 +189,9 @@ def migrate_models_ckpt():
     if not os.path.exists(os.path.join(model_path, "model.ckpt")):
         return
     new_name = initial_models()["stable-diffusion-1.4"]["file"]
-    print('The Stable Diffusion v4.1 "model.ckpt" is already installed. The name will be changed to {new_name} to avoid confusion.')
+    print(
+        'The Stable Diffusion v4.1 "model.ckpt" is already installed. The name will be changed to {new_name} to avoid confusion.'
+    )
     print(f"model.ckpt => {new_name}")
     os.replace(
         os.path.join(model_path, "model.ckpt"), os.path.join(model_path, new_name)
@@ -262,7 +270,6 @@ def _download_diffusion_weights(
             path = download_from_hf(
                 model_class,
                 repo_id,
-                cache_subdir="diffusers",
                 safety_checker=None,
                 **extra_args,
             )
@@ -383,7 +390,8 @@ def update_config_file(successfully_downloaded: dict, config_file: Path):
 
 # ---------------------------------------------
 def new_config_file_contents(
-        successfully_downloaded: dict, config_file: Path,
+    successfully_downloaded: dict,
+    config_file: Path,
 ) -> str:
     if config_file.exists():
         conf = OmegaConf.load(str(config_file.expanduser().resolve()))
@@ -413,7 +421,9 @@ def new_config_file_contents(
             stanza["weights"] = os.path.relpath(
                 successfully_downloaded[model], start=Globals.root
             )
-            stanza["config"] = os.path.normpath(os.path.join(sd_configs(), mod["config"]))
+            stanza["config"] = os.path.normpath(
+                os.path.join(sd_configs(), mod["config"])
+            )
         if "vae" in mod:
             if "file" in mod["vae"]:
                 stanza["vae"] = os.path.normpath(
@@ -445,7 +455,7 @@ def delete_weights(model_name: str, conf_stanza: dict):
     print(
         f"\n** The checkpoint version of {model_name} is superseded by the diffusers version. Deleting the original file {weights}?"
     )
-          
+
     weights = Path(weights)
     if not weights.is_absolute():
         weights = Path(Globals.root) / weights

invokeai/backend/generate.py

@@ -5,6 +5,7 @@
 
 import gc
 import importlib
+import logging
 import os
 import random
 import re
@@ -19,28 +20,23 @@ import numpy as np
 import skimage
 import torch
 import transformers
+from PIL import Image, ImageOps
+from accelerate.utils import set_seed
 from diffusers.pipeline_utils import DiffusionPipeline
 from diffusers.utils.import_utils import is_xformers_available
 from omegaconf import OmegaConf
-from PIL import Image, ImageOps
-from pytorch_lightning import logging, seed_everything
-
-import ldm.invoke.conditioning
-from ldm.invoke.args import metadata_from_png
-from ldm.invoke.concepts_lib import HuggingFaceConceptsLibrary
-from ldm.invoke.conditioning import get_uc_and_c_and_ec
-from ldm.invoke.devices import choose_precision, choose_torch_device
-from ldm.invoke.generator.inpaint import infill_methods
-from ldm.invoke.globals import Globals, global_cache_dir
-from ldm.invoke.image_util import InitImageResizer
-from ldm.invoke.model_manager import ModelManager
-from ldm.invoke.pngwriter import PngWriter
-from ldm.invoke.seamless import configure_model_padding
-from ldm.invoke.txt2mask import Txt2Mask
-from ldm.models.diffusion.ddim import DDIMSampler
-from ldm.models.diffusion.ksampler import KSampler
-from ldm.models.diffusion.plms import PLMSSampler
+from pathlib import Path
 
+from .args import metadata_from_png
+from .generator import infill_methods
+from .globals import Globals, global_cache_dir
+from .image_util import InitImageResizer, PngWriter, Txt2Mask, configure_model_padding
+from .model_management import ModelManager
+from .safety_checker import SafetyChecker
+from .prompting import get_uc_and_c_and_ec
+from .prompting.conditioning import log_tokenization
+from .stable_diffusion import HuggingFaceConceptsLibrary
+from .util import choose_precision, choose_torch_device
 
 def fix_func(orig):
     if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
@@ -227,6 +223,7 @@ class Generate:
             self.precision,
             max_loaded_models=max_loaded_models,
             sequential_offload=self.free_gpu_mem,
+            embedding_path=Path(self.embedding_path),
         )
         # don't accept invalid models
         fallback = self.model_manager.default_model() or FALLBACK_MODEL_NAME
@@ -249,31 +246,8 @@ class Generate:
 
         # load safety checker if requested
         if safety_checker:
-            try:
-                print(">> Initializing NSFW checker")
-                from diffusers.pipelines.stable_diffusion.safety_checker import (
-                    StableDiffusionSafetyChecker,
-                )
-                from transformers import AutoFeatureExtractor
-
-                safety_model_id = "CompVis/stable-diffusion-safety-checker"
-                safety_model_path = global_cache_dir("hub")
-                self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
-                    safety_model_id,
-                    local_files_only=True,
-                    cache_dir=safety_model_path,
-                )
-                self.safety_feature_extractor = AutoFeatureExtractor.from_pretrained(
-                    safety_model_id,
-                    local_files_only=True,
-                    cache_dir=safety_model_path,
-                )
-                self.safety_checker.to(self.device)
-            except Exception:
-                print(
-                    "** An error was encountered while installing the safety checker:"
-                )
-                print(traceback.format_exc())
+            print(">> Initializing NSFW checker")
+            self.safety_checker = SafetyChecker(self.device)
         else:
             print(">> NSFW checker is disabled")
 
@@ -328,8 +302,8 @@ class Generate:
         variation_amount=0.0,
         threshold=0.0,
         perlin=0.0,
-        h_symmetry_time_pct = None,
-        v_symmetry_time_pct = None,
+        h_symmetry_time_pct=None,
+        v_symmetry_time_pct=None,
         karras_max=None,
         outdir=None,
         # these are specific to img2img and inpaint
@@ -486,7 +460,7 @@ class Generate:
 
         if sampler_name and (sampler_name != self.sampler_name):
             self.sampler_name = sampler_name
-            self._set_sampler()
+            self._set_scheduler()
 
         # apply the concepts library to the prompt
         prompt = self.huggingface_concepts_library.replace_concepts_with_triggers(
@@ -495,28 +469,11 @@ class Generate:
             self.model.textual_inversion_manager.get_all_trigger_strings(),
         )
 
-        # bit of a hack to change the cached sampler's karras threshold to
-        # whatever the user asked for
-        if karras_max is not None and isinstance(self.sampler, KSampler):
-            self.sampler.adjust_settings(karras_max=karras_max)
-
         tic = time.time()
         if self._has_cuda():
             torch.cuda.reset_peak_memory_stats()
 
         results = list()
-        init_image = None
-        mask_image = None
-
-        try:
-            if (
-                self.free_gpu_mem
-                and self.model.cond_stage_model.device != self.model.device
-            ):
-                self.model.cond_stage_model.device = self.model.device
-                self.model.cond_stage_model.to(self.model.device)
-        except AttributeError:
-            pass
 
         try:
             uc, c, extra_conditioning_info = get_uc_and_c_and_ec(
@@ -545,15 +502,6 @@ class Generate:
             generator.set_variation(self.seed, variation_amount, with_variations)
             generator.use_mps_noise = use_mps_noise
 
-            checker = (
-                {
-                    "checker": self.safety_checker,
-                    "extractor": self.safety_feature_extractor,
-                }
-                if self.safety_checker
-                else None
-            )
-
             results = generator.generate(
                 prompt,
                 iterations=iterations,
@@ -580,7 +528,7 @@ class Generate:
                 embiggen_strength=embiggen_strength,
                 inpaint_replace=inpaint_replace,
                 mask_blur_radius=mask_blur_radius,
-                safety_checker=checker,
+                safety_checker=self.safety_checker,
                 seam_size=seam_size,
                 seam_blur=seam_blur,
                 seam_strength=seam_strength,
@@ -717,7 +665,7 @@ class Generate:
             prompt,
             model=self.model,
             skip_normalize_legacy_blend=opt.skip_normalize,
-            log_tokens=ldm.invoke.conditioning.log_tokenization,
+            log_tokens=log_tokenization,
         )
 
         if tool in ("gfpgan", "codeformer", "upscale"):
@@ -741,7 +689,7 @@ class Generate:
             )
 
         elif tool == "outcrop":
-            from ldm.invoke.restoration.outcrop import Outcrop
+            from .restoration.outcrop import Outcrop
 
             extend_instructions = {}
             for direction, pixels in _pairwise(opt.outcrop):
@@ -794,7 +742,7 @@ class Generate:
                 clear_cuda_cache=self.clear_cuda_cache,
             )
         elif tool == "outpaint":
-            from ldm.invoke.restoration.outpaint import Outpaint
+            from .restoration.outpaint import Outpaint
 
             restorer = Outpaint(image, self)
             return restorer.process(opt, args, image_callback=callback, prefix=prefix)
@@ -816,17 +764,12 @@ class Generate:
         hires_fix: bool = False,
         force_outpaint: bool = False,
     ):
-        inpainting_model_in_use = self.sampler.uses_inpainting_model()
-
         if hires_fix:
             return self._make_txt2img2img()
 
         if embiggen is not None:
             return self._make_embiggen()
 
-        if inpainting_model_in_use:
-            return self._make_omnibus()
-
         if ((init_image is not None) and (mask_image is not None)) or force_outpaint:
             return self._make_inpaint()
 
@@ -903,16 +846,9 @@ class Generate:
     def _make_inpaint(self):
         return self._load_generator(".inpaint", "Inpaint")
 
-    def _make_omnibus(self):
-        return self._load_generator(".omnibus", "Omnibus")
-
     def _load_generator(self, module, class_name):
-        if self.is_legacy_model(self.model_name):
-            mn = f"ldm.invoke.ckpt_generator{module}"
-            cn = f"Ckpt{class_name}"
-        else:
-            mn = f"ldm.invoke.generator{module}"
-            cn = class_name
+        mn = f"invokeai.backend.generator{module}"
+        cn = class_name
         module = importlib.import_module(mn)
         constructor = getattr(module, cn)
         return constructor(self.model, self.precision)
@@ -973,21 +909,9 @@ class Generate:
         # uncache generators so they pick up new models
         self.generators = {}
 
-        seed_everything(random.randrange(0, np.iinfo(np.uint32).max))
-        if self.embedding_path is not None:
-            print(f'>> Loading embeddings from {self.embedding_path}')
-            for root, _, files in os.walk(self.embedding_path):
-                for name in files:
-                    ti_path = os.path.join(root, name)
-                    self.model.textual_inversion_manager.load_textual_inversion(
-                        ti_path, defer_injecting_tokens=True
-                    )
-            print(
-                f'>> Textual inversion triggers: {", ".join(sorted(self.model.textual_inversion_manager.get_all_trigger_strings()))}'
-            )
-
+        set_seed(random.randrange(0, np.iinfo(np.uint32).max))
         self.model_name = model_name
-        self._set_sampler()  # requires self.model_name to be set first
+        self._set_scheduler()  # requires self.model_name to be set first
         return self.model
 
     def load_huggingface_concepts(self, concepts: list[str]):
@@ -1030,10 +954,9 @@ class Generate:
         image_callback=None,
         prefix=None,
     ):
-
         results = []
         for r in image_list:
-            image, seed = r
+            image, seed, _ = r
             try:
                 if strength > 0:
                     if self.gfpgan is not None or self.codeformer is not None:
@@ -1120,44 +1043,6 @@ class Generate:
     def is_legacy_model(self, model_name) -> bool:
         return self.model_manager.is_legacy(model_name)
 
-    def _set_sampler(self):
-        if isinstance(self.model, DiffusionPipeline):
-            return self._set_scheduler()
-        else:
-            return self._set_sampler_legacy()
-
-    # very repetitive code - can this be simplified? The KSampler names are
-    # consistent, at least
-    def _set_sampler_legacy(self):
-        msg = f">> Setting Sampler to {self.sampler_name}"
-        if self.sampler_name == "plms":
-            self.sampler = PLMSSampler(self.model, device=self.device)
-        elif self.sampler_name == "ddim":
-            self.sampler = DDIMSampler(self.model, device=self.device)
-        elif self.sampler_name == "k_dpm_2_a":
-            self.sampler = KSampler(self.model, "dpm_2_ancestral", device=self.device)
-        elif self.sampler_name == "k_dpm_2":
-            self.sampler = KSampler(self.model, "dpm_2", device=self.device)
-        elif self.sampler_name == "k_dpmpp_2_a":
-            self.sampler = KSampler(
-                self.model, "dpmpp_2s_ancestral", device=self.device
-            )
-        elif self.sampler_name == "k_dpmpp_2":
-            self.sampler = KSampler(self.model, "dpmpp_2m", device=self.device)
-        elif self.sampler_name == "k_euler_a":
-            self.sampler = KSampler(self.model, "euler_ancestral", device=self.device)
-        elif self.sampler_name == "k_euler":
-            self.sampler = KSampler(self.model, "euler", device=self.device)
-        elif self.sampler_name == "k_heun":
-            self.sampler = KSampler(self.model, "heun", device=self.device)
-        elif self.sampler_name == "k_lms":
-            self.sampler = KSampler(self.model, "lms", device=self.device)
-        else:
-            msg = f">> Unsupported Sampler: {self.sampler_name}, Defaulting to plms"
-            self.sampler = PLMSSampler(self.model, device=self.device)
-
-        print(msg)
-
     def _set_scheduler(self):
         default = self.model.scheduler
 

invokeai/backend/generator/__init__.py

@@ -0,0 +1,13 @@
+"""
+Initialization file for the invokeai.generator package
+"""
+from .base import (
+    InvokeAIGenerator,
+    InvokeAIGeneratorBasicParams,
+    InvokeAIGeneratorOutput,
+    Txt2Img,
+    Img2Img,
+    Inpaint,
+    Generator,
+)
+from .inpaint import infill_methods

invokeai/backend/generator/base.py

@@ -0,0 +1,649 @@
+"""
+Base class for invokeai.backend.generator.*
+including img2img, txt2img, and inpaint
+"""
+from __future__ import annotations
+
+import itertools
+import dataclasses
+import diffusers
+import os
+import random
+import traceback
+from abc import ABCMeta
+from argparse import Namespace
+from contextlib import nullcontext
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image, ImageChops, ImageFilter
+from accelerate.utils import set_seed
+from diffusers import DiffusionPipeline
+from tqdm import trange
+from typing import Callable, List, Iterator, Optional, Type
+from dataclasses import dataclass, field
+from diffusers.schedulers import SchedulerMixin as Scheduler
+
+from ..image_util import configure_model_padding
+from ..util.util import rand_perlin_2d
+from ..safety_checker import SafetyChecker
+from ..prompting.conditioning import get_uc_and_c_and_ec
+from ..stable_diffusion.diffusers_pipeline import StableDiffusionGeneratorPipeline
+
+downsampling = 8
+
+@dataclass
+class InvokeAIGeneratorBasicParams:
+    seed: Optional[int]=None
+    width: int=512
+    height: int=512
+    cfg_scale: float=7.5
+    steps: int=20
+    ddim_eta: float=0.0
+    scheduler: str='ddim'
+    precision: str='float16'
+    perlin: float=0.0
+    threshold: float=0.0
+    seamless: bool=False
+    seamless_axes: List[str]=field(default_factory=lambda: ['x', 'y'])
+    h_symmetry_time_pct: Optional[float]=None
+    v_symmetry_time_pct: Optional[float]=None
+    variation_amount: float = 0.0
+    with_variations: list=field(default_factory=list)
+    safety_checker: Optional[SafetyChecker]=None
+
+@dataclass
+class InvokeAIGeneratorOutput:
+    '''
+    InvokeAIGeneratorOutput is a dataclass that contains the outputs of a generation
+    operation, including the image, its seed, the model name used to generate the image
+    and the model hash, as well as all the generate() parameters that went into
+    generating the image (in .params, also available as attributes)
+    '''
+    image: Image.Image
+    seed: int
+    model_hash: str
+    attention_maps_images: List[Image.Image]
+    params: Namespace
+
+# we are interposing a wrapper around the original Generator classes so that
+# old code that calls Generate will continue to work.
+class InvokeAIGenerator(metaclass=ABCMeta):
+    scheduler_map = dict(
+        ddim=diffusers.DDIMScheduler,
+        dpmpp_2=diffusers.DPMSolverMultistepScheduler,
+        k_dpm_2=diffusers.KDPM2DiscreteScheduler,
+        k_dpm_2_a=diffusers.KDPM2AncestralDiscreteScheduler,
+        k_dpmpp_2=diffusers.DPMSolverMultistepScheduler,
+        k_euler=diffusers.EulerDiscreteScheduler,
+        k_euler_a=diffusers.EulerAncestralDiscreteScheduler,
+        k_heun=diffusers.HeunDiscreteScheduler,
+        k_lms=diffusers.LMSDiscreteScheduler,
+        plms=diffusers.PNDMScheduler,
+    )
+
+    def __init__(self,
+                 model_info: dict,
+                 params: InvokeAIGeneratorBasicParams=InvokeAIGeneratorBasicParams(),
+                 ):
+        self.model_info=model_info
+        self.params=params
+
+    def generate(self,
+                 prompt: str='',
+                 callback: Optional[Callable]=None,
+                 step_callback: Optional[Callable]=None,
+                 iterations: int=1,
+                 **keyword_args,
+                 )->Iterator[InvokeAIGeneratorOutput]:
+        '''
+        Return an iterator across the indicated number of generations.
+        Each time the iterator is called it will return an InvokeAIGeneratorOutput
+        object. Use like this:
+
+           outputs = txt2img.generate(prompt='banana sushi', iterations=5)
+           for result in outputs:
+               print(result.image, result.seed)
+
+        In the typical case of wanting to get just a single image, iterations
+        defaults to 1 and do:
+
+           output = next(txt2img.generate(prompt='banana sushi')
+
+        Pass None to get an infinite iterator.
+
+           outputs = txt2img.generate(prompt='banana sushi', iterations=None)
+           for o in outputs:
+               print(o.image, o.seed)
+
+        '''
+        generator_args = dataclasses.asdict(self.params)
+        generator_args.update(keyword_args)
+
+        model_info = self.model_info
+        model_name = model_info['model_name']
+        model:StableDiffusionGeneratorPipeline = model_info['model']
+        model_hash = model_info['hash']
+        scheduler: Scheduler = self.get_scheduler(
+            model=model,
+            scheduler_name=generator_args.get('scheduler')
+        )
+        uc, c, extra_conditioning_info = get_uc_and_c_and_ec(prompt,model=model)
+        gen_class = self._generator_class()
+        generator = gen_class(model, self.params.precision)
+        if self.params.variation_amount > 0:
+            generator.set_variation(generator_args.get('seed'),
+                                    generator_args.get('variation_amount'),
+                                    generator_args.get('with_variations')
+                                    )
+
+        if isinstance(model, DiffusionPipeline):
+            for component in [model.unet, model.vae]:
+                configure_model_padding(component,
+                                        generator_args.get('seamless',False),
+                                        generator_args.get('seamless_axes')
+                                        )
+        else:
+            configure_model_padding(model,
+                                    generator_args.get('seamless',False),
+                                    generator_args.get('seamless_axes')
+                                    )
+
+        iteration_count = range(iterations) if iterations else itertools.count(start=0, step=1)
+        for i in iteration_count:
+            results = generator.generate(prompt,
+                                         conditioning=(uc, c, extra_conditioning_info),
+                                         step_callback=step_callback,
+                                         sampler=scheduler,
+                                         **generator_args,
+                                         )
+            output = InvokeAIGeneratorOutput(
+                image=results[0][0],
+                seed=results[0][1],
+                attention_maps_images=results[0][2],
+                model_hash = model_hash,
+                params=Namespace(model_name=model_name,**generator_args),
+            )
+            if callback:
+                callback(output)
+            yield output
+
+    @classmethod
+    def schedulers(self)->List[str]:
+        '''
+        Return list of all the schedulers that we currently handle.
+        '''
+        return list(self.scheduler_map.keys())
+
+    def load_generator(self, model: StableDiffusionGeneratorPipeline, generator_class: Type[Generator]):
+        return generator_class(model, self.params.precision)
+
+    def get_scheduler(self, scheduler_name:str, model: StableDiffusionGeneratorPipeline)->Scheduler:
+        scheduler_class = self.scheduler_map.get(scheduler_name,'ddim')
+        scheduler = scheduler_class.from_config(model.scheduler.config)
+        # hack copied over from generate.py
+        if not hasattr(scheduler, 'uses_inpainting_model'):
+            scheduler.uses_inpainting_model = lambda: False
+        return scheduler
+
+    @classmethod
+    def _generator_class(cls)->Type[Generator]:
+        '''
+        In derived classes return the name of the generator to apply.
+        If you don't override will return the name of the derived
+        class, which nicely parallels the generator class names.
+        '''
+        return Generator
+
+# ------------------------------------
+class Txt2Img(InvokeAIGenerator):
+    @classmethod
+    def _generator_class(cls):
+        from .txt2img import Txt2Img
+        return Txt2Img
+
+# ------------------------------------
+class Img2Img(InvokeAIGenerator):
+    def generate(self,
+               init_image: Image.Image | torch.FloatTensor,
+               strength: float=0.75,
+               **keyword_args
+               )->Iterator[InvokeAIGeneratorOutput]:
+        return super().generate(init_image=init_image,
+                                strength=strength,
+                                **keyword_args
+                                )
+    @classmethod
+    def _generator_class(cls):
+        from .img2img import Img2Img
+        return Img2Img
+
+# ------------------------------------
+# Takes all the arguments of Img2Img and adds the mask image and the seam/infill stuff
+class Inpaint(Img2Img):
+    def generate(self,
+                 mask_image: Image.Image | torch.FloatTensor,
+                 # Seam settings - when 0, doesn't fill seam
+                 seam_size: int = 0,
+                 seam_blur: int = 0,
+                 seam_strength: float = 0.7,
+                 seam_steps: int = 10,
+                 tile_size: int = 32,
+                 inpaint_replace=False,
+                 infill_method=None,
+                 inpaint_width=None,
+                 inpaint_height=None,
+                 inpaint_fill: tuple(int) = (0x7F, 0x7F, 0x7F, 0xFF),
+                 **keyword_args
+                 )->Iterator[InvokeAIGeneratorOutput]:
+        return super().generate(
+            mask_image=mask_image,
+            seam_size=seam_size,
+            seam_blur=seam_blur,
+            seam_strength=seam_strength,
+            seam_steps=seam_steps,
+            tile_size=tile_size,
+            inpaint_replace=inpaint_replace,
+            infill_method=infill_method,
+            inpaint_width=inpaint_width,
+            inpaint_height=inpaint_height,
+            inpaint_fill=inpaint_fill,
+            **keyword_args
+        )
+    @classmethod
+    def _generator_class(cls):
+        from .inpaint import Inpaint
+        return Inpaint
+
+# ------------------------------------
+class Embiggen(Txt2Img):
+    def generate(
+            self,
+            embiggen: list=None,
+            embiggen_tiles: list = None,
+            strength: float=0.75,
+            **kwargs)->Iterator[InvokeAIGeneratorOutput]:
+        return super().generate(embiggen=embiggen,
+                                embiggen_tiles=embiggen_tiles,
+                                strength=strength,
+                                **kwargs)
+
+    @classmethod
+    def _generator_class(cls):
+        from .embiggen import Embiggen
+        return Embiggen
+
+
+class Generator:
+    downsampling_factor: int
+    latent_channels: int
+    precision: str
+    model: DiffusionPipeline
+
+    def __init__(self, model: DiffusionPipeline, precision: str):
+        self.model = model
+        self.precision = precision
+        self.seed = None
+        self.latent_channels = model.channels
+        self.downsampling_factor = downsampling  # BUG: should come from model or config
+        self.safety_checker = None
+        self.perlin = 0.0
+        self.threshold = 0
+        self.variation_amount = 0
+        self.with_variations = []
+        self.use_mps_noise = False
+        self.free_gpu_mem = None
+
+    # this is going to be overridden in img2img.py, txt2img.py and inpaint.py
+    def get_make_image(self, prompt, **kwargs):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it
+        """
+        raise NotImplementedError(
+            "image_iterator() must be implemented in a descendent class"
+        )
+
+    def set_variation(self, seed, variation_amount, with_variations):
+        self.seed = seed
+        self.variation_amount = variation_amount
+        self.with_variations = with_variations
+
+    def generate(
+        self,
+        prompt,
+        width,
+        height,
+        sampler,
+        init_image=None,
+        iterations=1,
+        seed=None,
+        image_callback=None,
+        step_callback=None,
+        threshold=0.0,
+        perlin=0.0,
+        h_symmetry_time_pct=None,
+        v_symmetry_time_pct=None,
+        safety_checker: SafetyChecker=None,
+        free_gpu_mem: bool = False,
+        **kwargs,
+    ):
+        scope = nullcontext
+        self.safety_checker = safety_checker
+        self.free_gpu_mem = free_gpu_mem
+        attention_maps_images = []
+        attention_maps_callback = lambda saver: attention_maps_images.append(
+            saver.get_stacked_maps_image()
+        )
+        make_image = self.get_make_image(
+            prompt,
+            sampler=sampler,
+            init_image=init_image,
+            width=width,
+            height=height,
+            step_callback=step_callback,
+            threshold=threshold,
+            perlin=perlin,
+            h_symmetry_time_pct=h_symmetry_time_pct,
+            v_symmetry_time_pct=v_symmetry_time_pct,
+            attention_maps_callback=attention_maps_callback,
+            **kwargs,
+        )
+        results = []
+        seed = seed if seed is not None and seed >= 0 else self.new_seed()
+        first_seed = seed
+        seed, initial_noise = self.generate_initial_noise(seed, width, height)
+
+        # There used to be an additional self.model.ema_scope() here, but it breaks
+        # the inpaint-1.5 model. Not sure what it did.... ?
+        with scope(self.model.device.type):
+            for n in trange(iterations, desc="Generating"):
+                x_T = None
+                if self.variation_amount > 0:
+                    set_seed(seed)
+                    target_noise = self.get_noise(width, height)
+                    x_T = self.slerp(self.variation_amount, initial_noise, target_noise)
+                elif initial_noise is not None:
+                    # i.e. we specified particular variations
+                    x_T = initial_noise
+                else:
+                    set_seed(seed)
+                    try:
+                        x_T = self.get_noise(width, height)
+                    except:
+                        print("** An error occurred while getting initial noise **")
+                        print(traceback.format_exc())
+
+                # Pass on the seed in case a layer beneath us needs to generate noise on its own.
+                image = make_image(x_T, seed)
+
+                if self.safety_checker is not None:
+                    image = self.safety_checker.check(image)
+
+                results.append([image, seed, attention_maps_images])
+
+                if image_callback is not None:
+                    attention_maps_image = (
+                        None
+                        if len(attention_maps_images) == 0
+                        else attention_maps_images[-1]
+                    )
+                    image_callback(
+                        image,
+                        seed,
+                        first_seed=first_seed,
+                        attention_maps_image=attention_maps_image,
+                    )
+
+                seed = self.new_seed()
+
+                # Free up memory from the last generation.
+                clear_cuda_cache = (
+                    kwargs["clear_cuda_cache"] if "clear_cuda_cache" in kwargs else None
+                )
+                if clear_cuda_cache is not None:
+                    clear_cuda_cache()
+
+        return results
+
+    def sample_to_image(self, samples) -> Image.Image:
+        """
+        Given samples returned from a sampler, converts
+        it into a PIL Image
+        """
+        with torch.inference_mode():
+            image = self.model.decode_latents(samples)
+        return self.model.numpy_to_pil(image)[0]
+
+    def repaste_and_color_correct(
+        self,
+        result: Image.Image,
+        init_image: Image.Image,
+        init_mask: Image.Image,
+        mask_blur_radius: int = 8,
+    ) -> Image.Image:
+        if init_image is None or init_mask is None:
+            return result
+
+        # Get the original alpha channel of the mask if there is one.
+        # Otherwise it is some other black/white image format ('1', 'L' or 'RGB')
+        pil_init_mask = (
+            init_mask.getchannel("A")
+            if init_mask.mode == "RGBA"
+            else init_mask.convert("L")
+        )
+        pil_init_image = init_image.convert(
+            "RGBA"
+        )  # Add an alpha channel if one doesn't exist
+
+        # Build an image with only visible pixels from source to use as reference for color-matching.
+        init_rgb_pixels = np.asarray(init_image.convert("RGB"), dtype=np.uint8)
+        init_a_pixels = np.asarray(pil_init_image.getchannel("A"), dtype=np.uint8)
+        init_mask_pixels = np.asarray(pil_init_mask, dtype=np.uint8)
+
+        # Get numpy version of result
+        np_image = np.asarray(result, dtype=np.uint8)
+
+        # Mask and calculate mean and standard deviation
+        mask_pixels = init_a_pixels * init_mask_pixels > 0
+        np_init_rgb_pixels_masked = init_rgb_pixels[mask_pixels, :]
+        np_image_masked = np_image[mask_pixels, :]
+
+        if np_init_rgb_pixels_masked.size > 0:
+            init_means = np_init_rgb_pixels_masked.mean(axis=0)
+            init_std = np_init_rgb_pixels_masked.std(axis=0)
+            gen_means = np_image_masked.mean(axis=0)
+            gen_std = np_image_masked.std(axis=0)
+
+            # Color correct
+            np_matched_result = np_image.copy()
+            np_matched_result[:, :, :] = (
+                (
+                    (
+                        (
+                            np_matched_result[:, :, :].astype(np.float32)
+                            - gen_means[None, None, :]
+                        )
+                        / gen_std[None, None, :]
+                    )
+                    * init_std[None, None, :]
+                    + init_means[None, None, :]
+                )
+                .clip(0, 255)
+                .astype(np.uint8)
+            )
+            matched_result = Image.fromarray(np_matched_result, mode="RGB")
+        else:
+            matched_result = Image.fromarray(np_image, mode="RGB")
+
+        # Blur the mask out (into init image) by specified amount
+        if mask_blur_radius > 0:
+            nm = np.asarray(pil_init_mask, dtype=np.uint8)
+            nmd = cv2.erode(
+                nm,
+                kernel=np.ones((3, 3), dtype=np.uint8),
+                iterations=int(mask_blur_radius / 2),
+            )
+            pmd = Image.fromarray(nmd, mode="L")
+            blurred_init_mask = pmd.filter(ImageFilter.BoxBlur(mask_blur_radius))
+        else:
+            blurred_init_mask = pil_init_mask
+
+        multiplied_blurred_init_mask = ImageChops.multiply(
+            blurred_init_mask, self.pil_image.split()[-1]
+        )
+
+        # Paste original on color-corrected generation (using blurred mask)
+        matched_result.paste(init_image, (0, 0), mask=multiplied_blurred_init_mask)
+        return matched_result
+
+    @staticmethod
+    def sample_to_lowres_estimated_image(samples):
+        # origingally adapted from code by @erucipe and @keturn here:
+        # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7
+
+        # these updated numbers for v1.5 are from @torridgristle
+        v1_5_latent_rgb_factors = torch.tensor(
+            [
+                #    R        G        B
+                [0.3444, 0.1385, 0.0670],  # L1
+                [0.1247, 0.4027, 0.1494],  # L2
+                [-0.3192, 0.2513, 0.2103],  # L3
+                [-0.1307, -0.1874, -0.7445],  # L4
+            ],
+            dtype=samples.dtype,
+            device=samples.device,
+        )
+
+        latent_image = samples[0].permute(1, 2, 0) @ v1_5_latent_rgb_factors
+        latents_ubyte = (
+            ((latent_image + 1) / 2)
+            .clamp(0, 1)  # change scale from -1..1 to 0..1
+            .mul(0xFF)  # to 0..255
+            .byte()
+        ).cpu()
+
+        return Image.fromarray(latents_ubyte.numpy())
+
+    def generate_initial_noise(self, seed, width, height):
+        initial_noise = None
+        if self.variation_amount > 0 or len(self.with_variations) > 0:
+            # use fixed initial noise plus random noise per iteration
+            set_seed(seed)
+            initial_noise = self.get_noise(width, height)
+            for v_seed, v_weight in self.with_variations:
+                seed = v_seed
+                set_seed(seed)
+                next_noise = self.get_noise(width, height)
+                initial_noise = self.slerp(v_weight, initial_noise, next_noise)
+            if self.variation_amount > 0:
+                random.seed()  # reset RNG to an actually random state, so we can get a random seed for variations
+                seed = random.randrange(0, np.iinfo(np.uint32).max)
+        return (seed, initial_noise)
+
+    def get_perlin_noise(self, width, height):
+        fixdevice = "cpu" if (self.model.device.type == "mps") else self.model.device
+        # limit noise to only the diffusion image channels, not the mask channels
+        input_channels = min(self.latent_channels, 4)
+        # round up to the nearest block of 8
+        temp_width = int((width + 7) / 8) * 8
+        temp_height = int((height + 7) / 8) * 8
+        noise = torch.stack(
+            [
+                rand_perlin_2d(
+                    (temp_height, temp_width), (8, 8), device=self.model.device
+                ).to(fixdevice)
+                for _ in range(input_channels)
+            ],
+            dim=0,
+        ).to(self.model.device)
+        return noise[0:4, 0:height, 0:width]
+
+    def new_seed(self):
+        self.seed = random.randrange(0, np.iinfo(np.uint32).max)
+        return self.seed
+
+    def slerp(self, t, v0, v1, DOT_THRESHOLD=0.9995):
+        """
+        Spherical linear interpolation
+        Args:
+            t (float/np.ndarray): Float value between 0.0 and 1.0
+            v0 (np.ndarray): Starting vector
+            v1 (np.ndarray): Final vector
+            DOT_THRESHOLD (float): Threshold for considering the two vectors as
+                                colineal. Not recommended to alter this.
+        Returns:
+            v2 (np.ndarray): Interpolation vector between v0 and v1
+        """
+        inputs_are_torch = False
+        if not isinstance(v0, np.ndarray):
+            inputs_are_torch = True
+            v0 = v0.detach().cpu().numpy()
+        if not isinstance(v1, np.ndarray):
+            inputs_are_torch = True
+            v1 = v1.detach().cpu().numpy()
+
+        dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
+        if np.abs(dot) > DOT_THRESHOLD:
+            v2 = (1 - t) * v0 + t * v1
+        else:
+            theta_0 = np.arccos(dot)
+            sin_theta_0 = np.sin(theta_0)
+            theta_t = theta_0 * t
+            sin_theta_t = np.sin(theta_t)
+            s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+            s1 = sin_theta_t / sin_theta_0
+            v2 = s0 * v0 + s1 * v1
+
+        if inputs_are_torch:
+            v2 = torch.from_numpy(v2).to(self.model.device)
+
+        return v2
+
+    # this is a handy routine for debugging use. Given a generated sample,
+    # convert it into a PNG image and store it at the indicated path
+    def save_sample(self, sample, filepath):
+        image = self.sample_to_image(sample)
+        dirname = os.path.dirname(filepath) or "."
+        if not os.path.exists(dirname):
+            print(f"** creating directory {dirname}")
+            os.makedirs(dirname, exist_ok=True)
+        image.save(filepath, "PNG")
+
+    def torch_dtype(self) -> torch.dtype:
+        return torch.float16 if self.precision == "float16" else torch.float32
+
+    # returns a tensor filled with random numbers from a normal distribution
+    def get_noise(self, width, height):
+        device = self.model.device
+        # limit noise to only the diffusion image channels, not the mask channels
+        input_channels = min(self.latent_channels, 4)
+        if self.use_mps_noise or device.type == "mps":
+            x = torch.randn(
+                [
+                    1,
+                    input_channels,
+                    height // self.downsampling_factor,
+                    width // self.downsampling_factor,
+                ],
+                dtype=self.torch_dtype(),
+                device="cpu",
+            ).to(device)
+        else:
+            x = torch.randn(
+                [
+                    1,
+                    input_channels,
+                    height // self.downsampling_factor,
+                    width // self.downsampling_factor,
+                ],
+                dtype=self.torch_dtype(),
+                device=device,
+            )
+        if self.perlin > 0.0:
+            perlin_noise = self.get_perlin_noise(
+                width // self.downsampling_factor, height // self.downsampling_factor
+            )
+            x = (1 - self.perlin) * x + self.perlin * perlin_noise
+        return x

invokeai/backend/generator/embiggen.py

@@ -1,37 +1,38 @@
-'''
-ldm.invoke.generator.embiggen descends from ldm.invoke.generator
-and generates with ldm.invoke.generator.img2img
-'''
+"""
+invokeai.backend.generator.embiggen descends from .generator
+and generates with .generator.img2img
+"""
 
 import numpy as np
 import torch
 from PIL import Image
 from tqdm import trange
 
-from ldm.invoke.generator.base import Generator
-from ldm.invoke.generator.img2img import Img2Img
+from .base import Generator
+from .img2img import Img2Img
 
 
 class Embiggen(Generator):
     def __init__(self, model, precision):
         super().__init__(model, precision)
-        self.init_latent         = None
+        self.init_latent = None
 
     # Replace generate because Embiggen doesn't need/use most of what it does normallly
-    def generate(self,prompt,iterations=1,seed=None,
-                 image_callback=None, step_callback=None,
-                 **kwargs):
-
-        make_image = self.get_make_image(
-            prompt,
-            step_callback = step_callback,
-            **kwargs
-        )
-        results             = []
-        seed                = seed if seed else self.new_seed()
+    def generate(
+        self,
+        prompt,
+        iterations=1,
+        seed=None,
+        image_callback=None,
+        step_callback=None,
+        **kwargs,
+    ):
+        make_image = self.get_make_image(prompt, step_callback=step_callback, **kwargs)
+        results = []
+        seed = seed if seed else self.new_seed()
 
         # Noise will be generated by the Img2Img generator when called
-        for _ in trange(iterations, desc='Generating'):
+        for _ in trange(iterations, desc="Generating"):
             # make_image will call Img2Img which will do the equivalent of get_noise itself
             image = make_image()
             results.append([image, seed])
@@ -56,13 +57,15 @@ class Embiggen(Generator):
         embiggen,
         embiggen_tiles,
         step_callback=None,
-        **kwargs
+        **kwargs,
     ):
         """
         Returns a function returning an image derived from the prompt and multi-stage twice-baked potato layering over the img2img on the initial image
         Return value depends on the seed at the time you call it
         """
-        assert not sampler.uses_inpainting_model(), "--embiggen is not supported by inpainting models"
+        assert (
+            not sampler.uses_inpainting_model()
+        ), "--embiggen is not supported by inpainting models"
 
         # Construct embiggen arg array, and sanity check arguments
         if embiggen == None:  # embiggen can also be called with just embiggen_tiles
@@ -70,48 +73,57 @@ class Embiggen(Generator):
         elif embiggen[0] < 0:
             embiggen[0] = 1.0
             print(
-                '>> Embiggen scaling factor cannot be negative, fell back to the default of 1.0 !')
+                ">> Embiggen scaling factor cannot be negative, fell back to the default of 1.0 !"
+            )
         if len(embiggen) < 2:
             embiggen.append(0.75)
         elif embiggen[1] > 1.0 or embiggen[1] < 0:
             embiggen[1] = 0.75
-            print('>> Embiggen upscaling strength for ESRGAN must be between 0 and 1, fell back to the default of 0.75 !')
+            print(
+                ">> Embiggen upscaling strength for ESRGAN must be between 0 and 1, fell back to the default of 0.75 !"
+            )
         if len(embiggen) < 3:
             embiggen.append(0.25)
         elif embiggen[2] < 0:
             embiggen[2] = 0.25
-            print('>> Overlap size for Embiggen must be a positive ratio between 0 and 1 OR a number of pixels, fell back to the default of 0.25 !')
+            print(
+                ">> Overlap size for Embiggen must be a positive ratio between 0 and 1 OR a number of pixels, fell back to the default of 0.25 !"
+            )
 
         # Convert tiles from their user-freindly count-from-one to count-from-zero, because we need to do modulo math
         # and then sort them, because... people.
         if embiggen_tiles:
-            embiggen_tiles = list(map(lambda n: n-1, embiggen_tiles))
+            embiggen_tiles = list(map(lambda n: n - 1, embiggen_tiles))
             embiggen_tiles.sort()
 
         if strength >= 0.5:
-            print(f'* WARNING: Embiggen may produce mirror motifs if the strength (-f) is too high (currently {strength}). Try values between 0.35-0.45.')
+            print(
+                f"* WARNING: Embiggen may produce mirror motifs if the strength (-f) is too high (currently {strength}). Try values between 0.35-0.45."
+            )
 
         # Prep img2img generator, since we wrap over it
-        gen_img2img = Img2Img(self.model,self.precision)
+        gen_img2img = Img2Img(self.model, self.precision)
 
         # Open original init image (not a tensor) to manipulate
         initsuperimage = Image.open(init_img)
 
         with Image.open(init_img) as img:
-            initsuperimage = img.convert('RGB')
+            initsuperimage = img.convert("RGB")
 
         # Size of the target super init image in pixels
         initsuperwidth, initsuperheight = initsuperimage.size
 
         # Increase by scaling factor if not already resized, using ESRGAN as able
         if embiggen[0] != 1.0:
-            initsuperwidth = round(initsuperwidth*embiggen[0])
-            initsuperheight = round(initsuperheight*embiggen[0])
+            initsuperwidth = round(initsuperwidth * embiggen[0])
+            initsuperheight = round(initsuperheight * embiggen[0])
             if embiggen[1] > 0:  # No point in ESRGAN upscaling if strength is set zero
-                from ldm.invoke.restoration.realesrgan import ESRGAN
+                from ..restoration.realesrgan import ESRGAN
+
                 esrgan = ESRGAN()
                 print(
-                    f'>> ESRGAN upscaling init image prior to cutting with Embiggen with strength {embiggen[1]}')
+                    f">> ESRGAN upscaling init image prior to cutting with Embiggen with strength {embiggen[1]}"
+                )
                 if embiggen[0] > 2:
                     initsuperimage = esrgan.process(
                         initsuperimage,
@@ -130,7 +142,8 @@ class Embiggen(Generator):
             #   but from personal experiance it doesn't greatly improve anything after 4x
             # Resize to target scaling factor resolution
             initsuperimage = initsuperimage.resize(
-                (initsuperwidth, initsuperheight), Image.Resampling.LANCZOS)
+                (initsuperwidth, initsuperheight), Image.Resampling.LANCZOS
+            )
 
         # Use width and height as tile widths and height
         # Determine buffer size in pixels
@@ -153,23 +166,24 @@ class Embiggen(Generator):
         emb_tiles_x = 1
         emb_tiles_y = 1
         if (initsuperwidth - width) > 0:
-            emb_tiles_x = ceildiv(initsuperwidth - width,
-                                  width - overlap_size_x) + 1
+            emb_tiles_x = ceildiv(initsuperwidth - width, width - overlap_size_x) + 1
         if (initsuperheight - height) > 0:
-            emb_tiles_y = ceildiv(initsuperheight - height,
-                                  height - overlap_size_y) + 1
+            emb_tiles_y = ceildiv(initsuperheight - height, height - overlap_size_y) + 1
         # Sanity
-        assert emb_tiles_x > 1 or emb_tiles_y > 1, f'ERROR: Based on the requested dimensions of {initsuperwidth}x{initsuperheight} and tiles of {width}x{height} you don\'t need to Embiggen! Check your arguments.'
+        assert (
+            emb_tiles_x > 1 or emb_tiles_y > 1
+        ), f"ERROR: Based on the requested dimensions of {initsuperwidth}x{initsuperheight} and tiles of {width}x{height} you don't need to Embiggen! Check your arguments."
 
         # Prep alpha layers --------------
         # https://stackoverflow.com/questions/69321734/how-to-create-different-transparency-like-gradient-with-python-pil
         # agradientL is Left-side transparent
-        agradientL = Image.linear_gradient('L').rotate(
-            90).resize((overlap_size_x, height))
+        agradientL = (
+            Image.linear_gradient("L").rotate(90).resize((overlap_size_x, height))
+        )
         # agradientT is Top-side transparent
-        agradientT = Image.linear_gradient('L').resize((width, overlap_size_y))
+        agradientT = Image.linear_gradient("L").resize((width, overlap_size_y))
         # radial corner is the left-top corner, made full circle then cut to just the left-top quadrant
-        agradientC = Image.new('L', (256, 256))
+        agradientC = Image.new("L", (256, 256))
         for y in range(256):
             for x in range(256):
                 # Find distance to lower right corner (numpy takes arrays)
@@ -177,16 +191,16 @@ class Embiggen(Generator):
                 # Clamp values to max 255
                 if distanceToLR > 255:
                     distanceToLR = 255
-                #Place the pixel as invert of distance
+                # Place the pixel as invert of distance
                 agradientC.putpixel((x, y), round(255 - distanceToLR))
 
         # Create alternative asymmetric diagonal corner to use on "tailing" intersections to prevent hard edges
         # Fits for a left-fading gradient on the bottom side and full opacity on the right side.
-        agradientAsymC = Image.new('L', (256, 256))
+        agradientAsymC = Image.new("L", (256, 256))
         for y in range(256):
             for x in range(256):
-                value = round(max(0, x-(255-y)) * (255 / max(1,y)))
-                #Clamp values
+                value = round(max(0, x - (255 - y)) * (255 / max(1, y)))
+                # Clamp values
                 value = max(0, value)
                 value = min(255, value)
                 agradientAsymC.putpixel((x, y), value)
@@ -204,80 +218,91 @@ class Embiggen(Generator):
         # make masks with an asymmetric upper-right corner so when the curved transparent corner of the next tile
         # to its right is placed it doesn't reveal a hard trailing semi-transparent edge in the overlapping space
         alphaLayerTaC = alphaLayerT.copy()
-        alphaLayerTaC.paste(agradientAsymC.rotate(270).resize((overlap_size_x, overlap_size_y)), (width - overlap_size_x, 0))
+        alphaLayerTaC.paste(
+            agradientAsymC.rotate(270).resize((overlap_size_x, overlap_size_y)),
+            (width - overlap_size_x, 0),
+        )
         alphaLayerLTaC = alphaLayerLTC.copy()
-        alphaLayerLTaC.paste(agradientAsymC.rotate(270).resize((overlap_size_x, overlap_size_y)), (width - overlap_size_x, 0))
+        alphaLayerLTaC.paste(
+            agradientAsymC.rotate(270).resize((overlap_size_x, overlap_size_y)),
+            (width - overlap_size_x, 0),
+        )
 
         if embiggen_tiles:
             # Individual unconnected sides
             alphaLayerR = Image.new("L", (width, height), 255)
-            alphaLayerR.paste(agradientL.rotate(
-                180), (width - overlap_size_x, 0))
+            alphaLayerR.paste(agradientL.rotate(180), (width - overlap_size_x, 0))
             alphaLayerB = Image.new("L", (width, height), 255)
-            alphaLayerB.paste(agradientT.rotate(
-                180), (0, height - overlap_size_y))
+            alphaLayerB.paste(agradientT.rotate(180), (0, height - overlap_size_y))
             alphaLayerTB = Image.new("L", (width, height), 255)
             alphaLayerTB.paste(agradientT, (0, 0))
-            alphaLayerTB.paste(agradientT.rotate(
-                180), (0, height - overlap_size_y))
+            alphaLayerTB.paste(agradientT.rotate(180), (0, height - overlap_size_y))
             alphaLayerLR = Image.new("L", (width, height), 255)
             alphaLayerLR.paste(agradientL, (0, 0))
-            alphaLayerLR.paste(agradientL.rotate(
-                180), (width - overlap_size_x, 0))
+            alphaLayerLR.paste(agradientL.rotate(180), (width - overlap_size_x, 0))
 
             # Sides and corner Layers
             alphaLayerRBC = Image.new("L", (width, height), 255)
-            alphaLayerRBC.paste(agradientL.rotate(
-                180), (width - overlap_size_x, 0))
-            alphaLayerRBC.paste(agradientT.rotate(
-                180), (0, height - overlap_size_y))
-            alphaLayerRBC.paste(agradientC.rotate(180).resize(
-                (overlap_size_x, overlap_size_y)), (width - overlap_size_x, height - overlap_size_y))
+            alphaLayerRBC.paste(agradientL.rotate(180), (width - overlap_size_x, 0))
+            alphaLayerRBC.paste(agradientT.rotate(180), (0, height - overlap_size_y))
+            alphaLayerRBC.paste(
+                agradientC.rotate(180).resize((overlap_size_x, overlap_size_y)),
+                (width - overlap_size_x, height - overlap_size_y),
+            )
             alphaLayerLBC = Image.new("L", (width, height), 255)
             alphaLayerLBC.paste(agradientL, (0, 0))
-            alphaLayerLBC.paste(agradientT.rotate(
-                180), (0, height - overlap_size_y))
-            alphaLayerLBC.paste(agradientC.rotate(90).resize(
-                (overlap_size_x, overlap_size_y)), (0, height - overlap_size_y))
+            alphaLayerLBC.paste(agradientT.rotate(180), (0, height - overlap_size_y))
+            alphaLayerLBC.paste(
+                agradientC.rotate(90).resize((overlap_size_x, overlap_size_y)),
+                (0, height - overlap_size_y),
+            )
             alphaLayerRTC = Image.new("L", (width, height), 255)
-            alphaLayerRTC.paste(agradientL.rotate(
-                180), (width - overlap_size_x, 0))
+            alphaLayerRTC.paste(agradientL.rotate(180), (width - overlap_size_x, 0))
             alphaLayerRTC.paste(agradientT, (0, 0))
-            alphaLayerRTC.paste(agradientC.rotate(270).resize(
-                (overlap_size_x, overlap_size_y)), (width - overlap_size_x, 0))
+            alphaLayerRTC.paste(
+                agradientC.rotate(270).resize((overlap_size_x, overlap_size_y)),
+                (width - overlap_size_x, 0),
+            )
 
             # All but X layers
             alphaLayerABT = Image.new("L", (width, height), 255)
             alphaLayerABT.paste(alphaLayerLBC, (0, 0))
-            alphaLayerABT.paste(agradientL.rotate(
-                180), (width - overlap_size_x, 0))
-            alphaLayerABT.paste(agradientC.rotate(180).resize(
-                (overlap_size_x, overlap_size_y)), (width - overlap_size_x, height - overlap_size_y))
+            alphaLayerABT.paste(agradientL.rotate(180), (width - overlap_size_x, 0))
+            alphaLayerABT.paste(
+                agradientC.rotate(180).resize((overlap_size_x, overlap_size_y)),
+                (width - overlap_size_x, height - overlap_size_y),
+            )
             alphaLayerABL = Image.new("L", (width, height), 255)
             alphaLayerABL.paste(alphaLayerRTC, (0, 0))
-            alphaLayerABL.paste(agradientT.rotate(
-                180), (0, height - overlap_size_y))
-            alphaLayerABL.paste(agradientC.rotate(180).resize(
-                (overlap_size_x, overlap_size_y)), (width - overlap_size_x, height - overlap_size_y))
+            alphaLayerABL.paste(agradientT.rotate(180), (0, height - overlap_size_y))
+            alphaLayerABL.paste(
+                agradientC.rotate(180).resize((overlap_size_x, overlap_size_y)),
+                (width - overlap_size_x, height - overlap_size_y),
+            )
             alphaLayerABR = Image.new("L", (width, height), 255)
             alphaLayerABR.paste(alphaLayerLBC, (0, 0))
             alphaLayerABR.paste(agradientT, (0, 0))
-            alphaLayerABR.paste(agradientC.resize(
-                (overlap_size_x, overlap_size_y)), (0, 0))
+            alphaLayerABR.paste(
+                agradientC.resize((overlap_size_x, overlap_size_y)), (0, 0)
+            )
             alphaLayerABB = Image.new("L", (width, height), 255)
             alphaLayerABB.paste(alphaLayerRTC, (0, 0))
             alphaLayerABB.paste(agradientL, (0, 0))
-            alphaLayerABB.paste(agradientC.resize(
-                (overlap_size_x, overlap_size_y)), (0, 0))
+            alphaLayerABB.paste(
+                agradientC.resize((overlap_size_x, overlap_size_y)), (0, 0)
+            )
 
             # All-around layer
             alphaLayerAA = Image.new("L", (width, height), 255)
             alphaLayerAA.paste(alphaLayerABT, (0, 0))
             alphaLayerAA.paste(agradientT, (0, 0))
-            alphaLayerAA.paste(agradientC.resize(
-                (overlap_size_x, overlap_size_y)), (0, 0))
-            alphaLayerAA.paste(agradientC.rotate(270).resize(
-                (overlap_size_x, overlap_size_y)), (width - overlap_size_x, 0))
+            alphaLayerAA.paste(
+                agradientC.resize((overlap_size_x, overlap_size_y)), (0, 0)
+            )
+            alphaLayerAA.paste(
+                agradientC.rotate(270).resize((overlap_size_x, overlap_size_y)),
+                (width - overlap_size_x, 0),
+            )
 
         # Clean up temporary gradients
         del agradientL
@@ -287,17 +312,20 @@ class Embiggen(Generator):
         def make_image():
             # Make main tiles -------------------------------------------------
             if embiggen_tiles:
-                print(f'>> Making {len(embiggen_tiles)} Embiggen tiles...')
+                print(f">> Making {len(embiggen_tiles)} Embiggen tiles...")
             else:
                 print(
-                    f'>> Making {(emb_tiles_x * emb_tiles_y)} Embiggen tiles ({emb_tiles_x}x{emb_tiles_y})...')
+                    f">> Making {(emb_tiles_x * emb_tiles_y)} Embiggen tiles ({emb_tiles_x}x{emb_tiles_y})..."
+                )
 
             emb_tile_store = []
             # Although we could use the same seed for every tile for determinism, at higher strengths this may
             # produce duplicated structures for each tile and make the tiling effect more obvious
             # instead track and iterate a local seed we pass to Img2Img
             seed = self.seed
-            seedintlimit = np.iinfo(np.uint32).max - 1 # only retreive this one from numpy
+            seedintlimit = (
+                np.iinfo(np.uint32).max - 1
+            )  # only retreive this one from numpy
 
             for tile in range(emb_tiles_x * emb_tiles_y):
                 # Don't iterate on first tile
@@ -334,37 +362,38 @@ class Embiggen(Generator):
 
                 if embiggen_tiles:
                     print(
-                        f'Making tile #{tile + 1} ({embiggen_tiles.index(tile) + 1} of {len(embiggen_tiles)} requested)')
+                        f"Making tile #{tile + 1} ({embiggen_tiles.index(tile) + 1} of {len(embiggen_tiles)} requested)"
+                    )
                 else:
-                    print(
-                        f'Starting {tile + 1} of {(emb_tiles_x * emb_tiles_y)} tiles')
+                    print(f"Starting {tile + 1} of {(emb_tiles_x * emb_tiles_y)} tiles")
 
                 # create a torch tensor from an Image
-                newinitimage = np.array(
-                    newinitimage).astype(np.float32) / 255.0
+                newinitimage = np.array(newinitimage).astype(np.float32) / 255.0
                 newinitimage = newinitimage[None].transpose(0, 3, 1, 2)
                 newinitimage = torch.from_numpy(newinitimage)
                 newinitimage = 2.0 * newinitimage - 1.0
                 newinitimage = newinitimage.to(self.model.device)
-                clear_cuda_cache = kwargs['clear_cuda_cache'] if 'clear_cuda_cache' in kwargs else None
+                clear_cuda_cache = (
+                    kwargs["clear_cuda_cache"] if "clear_cuda_cache" in kwargs else None
+                )
 
                 tile_results = gen_img2img.generate(
                     prompt,
-                    iterations     = 1,
-                    seed           = seed,
-                    sampler        = sampler,
-                    steps          = steps,
-                    cfg_scale      = cfg_scale,
-                    conditioning   = conditioning,
-                    ddim_eta       = ddim_eta,
-                    image_callback = None,  # called only after the final image is generated
-                    step_callback  = step_callback,   # called after each intermediate image is generated
-                    width          = width,
-                    height         = height,
-                    init_image     = newinitimage,    # notice that init_image is different from init_img
-                    mask_image     = None,
-                    strength       = strength,
-                    clear_cuda_cache     = clear_cuda_cache
+                    iterations=1,
+                    seed=seed,
+                    sampler=sampler,
+                    steps=steps,
+                    cfg_scale=cfg_scale,
+                    conditioning=conditioning,
+                    ddim_eta=ddim_eta,
+                    image_callback=None,  # called only after the final image is generated
+                    step_callback=step_callback,  # called after each intermediate image is generated
+                    width=width,
+                    height=height,
+                    init_image=newinitimage,  # notice that init_image is different from init_img
+                    mask_image=None,
+                    strength=strength,
+                    clear_cuda_cache=clear_cuda_cache,
                 )
 
                 emb_tile_store.append(tile_results[0][0])
@@ -373,12 +402,14 @@ class Embiggen(Generator):
                 del newinitimage
 
             # Sanity check we have them all
-            if len(emb_tile_store) == (emb_tiles_x * emb_tiles_y) or (embiggen_tiles != [] and len(emb_tile_store) == len(embiggen_tiles)):
-                outputsuperimage = Image.new(
-                    "RGBA", (initsuperwidth, initsuperheight))
+            if len(emb_tile_store) == (emb_tiles_x * emb_tiles_y) or (
+                embiggen_tiles != [] and len(emb_tile_store) == len(embiggen_tiles)
+            ):
+                outputsuperimage = Image.new("RGBA", (initsuperwidth, initsuperheight))
                 if embiggen_tiles:
                     outputsuperimage.alpha_composite(
-                        initsuperimage.convert('RGBA'), (0, 0))
+                        initsuperimage.convert("RGBA"), (0, 0)
+                    )
                 for tile in range(emb_tiles_x * emb_tiles_y):
                     if embiggen_tiles:
                         if tile in embiggen_tiles:
@@ -387,7 +418,7 @@ class Embiggen(Generator):
                             continue
                     else:
                         intileimage = emb_tile_store[tile]
-                    intileimage = intileimage.convert('RGBA')
+                    intileimage = intileimage.convert("RGBA")
                     # Get row and column entries
                     emb_row_i = tile // emb_tiles_x
                     emb_column_i = tile % emb_tiles_x
@@ -399,8 +430,7 @@ class Embiggen(Generator):
                         if emb_column_i + 1 == emb_tiles_x:
                             left = initsuperwidth - width
                         else:
-                            left = round(emb_column_i *
-                                         (width - overlap_size_x))
+                            left = round(emb_column_i * (width - overlap_size_x))
                         if emb_row_i + 1 == emb_tiles_y:
                             top = initsuperheight - height
                         else:
@@ -411,33 +441,43 @@ class Embiggen(Generator):
                             # top of image
                             if emb_row_i == 0:
                                 if emb_column_i == 0:
-                                    if (tile+1) in embiggen_tiles:  # Look-ahead right
-                                        if (tile+emb_tiles_x) not in embiggen_tiles:  # Look-ahead down
+                                    if (tile + 1) in embiggen_tiles:  # Look-ahead right
+                                        if (
+                                            tile + emb_tiles_x
+                                        ) not in embiggen_tiles:  # Look-ahead down
                                             intileimage.putalpha(alphaLayerB)
                                         # Otherwise do nothing on this tile
-                                    elif (tile+emb_tiles_x) in embiggen_tiles:  # Look-ahead down only
+                                    elif (
+                                        tile + emb_tiles_x
+                                    ) in embiggen_tiles:  # Look-ahead down only
                                         intileimage.putalpha(alphaLayerR)
                                     else:
                                         intileimage.putalpha(alphaLayerRBC)
                                 elif emb_column_i == emb_tiles_x - 1:
-                                    if (tile+emb_tiles_x) in embiggen_tiles:  # Look-ahead down
+                                    if (
+                                        tile + emb_tiles_x
+                                    ) in embiggen_tiles:  # Look-ahead down
                                         intileimage.putalpha(alphaLayerL)
                                     else:
                                         intileimage.putalpha(alphaLayerLBC)
                                 else:
-                                    if (tile+1) in embiggen_tiles:  # Look-ahead right
-                                        if (tile+emb_tiles_x) in embiggen_tiles:  # Look-ahead down
+                                    if (tile + 1) in embiggen_tiles:  # Look-ahead right
+                                        if (
+                                            tile + emb_tiles_x
+                                        ) in embiggen_tiles:  # Look-ahead down
                                             intileimage.putalpha(alphaLayerL)
                                         else:
                                             intileimage.putalpha(alphaLayerLBC)
-                                    elif (tile+emb_tiles_x) in embiggen_tiles:  # Look-ahead down only
+                                    elif (
+                                        tile + emb_tiles_x
+                                    ) in embiggen_tiles:  # Look-ahead down only
                                         intileimage.putalpha(alphaLayerLR)
                                     else:
                                         intileimage.putalpha(alphaLayerABT)
                             # bottom of image
                             elif emb_row_i == emb_tiles_y - 1:
                                 if emb_column_i == 0:
-                                    if (tile+1) in embiggen_tiles: # Look-ahead right
+                                    if (tile + 1) in embiggen_tiles:  # Look-ahead right
                                         intileimage.putalpha(alphaLayerTaC)
                                     else:
                                         intileimage.putalpha(alphaLayerRTC)
@@ -445,34 +485,44 @@ class Embiggen(Generator):
                                     # No tiles to look ahead to
                                     intileimage.putalpha(alphaLayerLTC)
                                 else:
-                                    if (tile+1) in embiggen_tiles: # Look-ahead right
+                                    if (tile + 1) in embiggen_tiles:  # Look-ahead right
                                         intileimage.putalpha(alphaLayerLTaC)
                                     else:
                                         intileimage.putalpha(alphaLayerABB)
                             # vertical middle of image
                             else:
                                 if emb_column_i == 0:
-                                    if (tile+1) in embiggen_tiles: # Look-ahead right
-                                        if (tile+emb_tiles_x) in embiggen_tiles: # Look-ahead down
+                                    if (tile + 1) in embiggen_tiles:  # Look-ahead right
+                                        if (
+                                            tile + emb_tiles_x
+                                        ) in embiggen_tiles:  # Look-ahead down
                                             intileimage.putalpha(alphaLayerTaC)
                                         else:
                                             intileimage.putalpha(alphaLayerTB)
-                                    elif (tile+emb_tiles_x) in embiggen_tiles:  # Look-ahead down only
+                                    elif (
+                                        tile + emb_tiles_x
+                                    ) in embiggen_tiles:  # Look-ahead down only
                                         intileimage.putalpha(alphaLayerRTC)
                                     else:
                                         intileimage.putalpha(alphaLayerABL)
                                 elif emb_column_i == emb_tiles_x - 1:
-                                    if (tile+emb_tiles_x) in embiggen_tiles:  # Look-ahead down
+                                    if (
+                                        tile + emb_tiles_x
+                                    ) in embiggen_tiles:  # Look-ahead down
                                         intileimage.putalpha(alphaLayerLTC)
                                     else:
                                         intileimage.putalpha(alphaLayerABR)
                                 else:
-                                    if (tile+1) in embiggen_tiles: # Look-ahead right
-                                        if (tile+emb_tiles_x) in embiggen_tiles: # Look-ahead down
+                                    if (tile + 1) in embiggen_tiles:  # Look-ahead right
+                                        if (
+                                            tile + emb_tiles_x
+                                        ) in embiggen_tiles:  # Look-ahead down
                                             intileimage.putalpha(alphaLayerLTaC)
                                         else:
                                             intileimage.putalpha(alphaLayerABR)
-                                    elif (tile+emb_tiles_x) in embiggen_tiles:  # Look-ahead down only
+                                    elif (
+                                        tile + emb_tiles_x
+                                    ) in embiggen_tiles:  # Look-ahead down only
                                         intileimage.putalpha(alphaLayerABB)
                                     else:
                                         intileimage.putalpha(alphaLayerAA)
@@ -481,21 +531,28 @@ class Embiggen(Generator):
                             if emb_row_i == 0 and emb_column_i >= 1:
                                 intileimage.putalpha(alphaLayerL)
                             elif emb_row_i >= 1 and emb_column_i == 0:
-                                if emb_column_i + 1 == emb_tiles_x: # If we don't have anything that can be placed to the right
+                                if (
+                                    emb_column_i + 1 == emb_tiles_x
+                                ):  # If we don't have anything that can be placed to the right
                                     intileimage.putalpha(alphaLayerT)
                                 else:
                                     intileimage.putalpha(alphaLayerTaC)
                             else:
-                                if emb_column_i + 1 == emb_tiles_x: # If we don't have anything that can be placed to the right
+                                if (
+                                    emb_column_i + 1 == emb_tiles_x
+                                ):  # If we don't have anything that can be placed to the right
                                     intileimage.putalpha(alphaLayerLTC)
                                 else:
                                     intileimage.putalpha(alphaLayerLTaC)
                     # Layer tile onto final image
                     outputsuperimage.alpha_composite(intileimage, (left, top))
             else:
-                print('Error: could not find all Embiggen output tiles in memory? Something must have gone wrong with img2img generation.')
+                print(
+                    "Error: could not find all Embiggen output tiles in memory? Something must have gone wrong with img2img generation."
+                )
 
             # after internal loops and patching up return Embiggen image
             return outputsuperimage
+
         # end of function declaration
         return make_image

invokeai/backend/generator/img2img.py

@@ -0,0 +1,101 @@
+"""
+invokeai.backend.generator.img2img descends from .generator
+"""
+from typing import Optional
+
+import torch
+from accelerate.utils import set_seed
+from diffusers import logging
+
+from ..stable_diffusion import (
+    ConditioningData,
+    PostprocessingSettings,
+    StableDiffusionGeneratorPipeline,
+)
+from .base import Generator
+
+
+class Img2Img(Generator):
+    def __init__(self, model, precision):
+        super().__init__(model, precision)
+        self.init_latent = None  # by get_noise()
+
+    def get_make_image(
+        self,
+        prompt,
+        sampler,
+        steps,
+        cfg_scale,
+        ddim_eta,
+        conditioning,
+        init_image,
+        strength,
+        step_callback=None,
+        threshold=0.0,
+        warmup=0.2,
+        perlin=0.0,
+        h_symmetry_time_pct=None,
+        v_symmetry_time_pct=None,
+        attention_maps_callback=None,
+        **kwargs,
+    ):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it.
+        """
+        self.perlin = perlin
+
+        # noinspection PyTypeChecker
+        pipeline: StableDiffusionGeneratorPipeline = self.model
+        pipeline.scheduler = sampler
+
+        uc, c, extra_conditioning_info = conditioning
+        conditioning_data = ConditioningData(
+            uc,
+            c,
+            cfg_scale,
+            extra_conditioning_info,
+            postprocessing_settings=PostprocessingSettings(
+                threshold=threshold,
+                warmup=warmup,
+                h_symmetry_time_pct=h_symmetry_time_pct,
+                v_symmetry_time_pct=v_symmetry_time_pct,
+            ),
+        ).add_scheduler_args_if_applicable(pipeline.scheduler, eta=ddim_eta)
+
+        def make_image(x_T: torch.Tensor, seed: int):
+            # FIXME: use x_T for initial seeded noise
+            # We're not at the moment because the pipeline automatically resizes init_image if
+            # necessary, which the x_T input might not match.
+            # In the meantime, reset the seed prior to generating pipeline output so we at least get the same result.
+            logging.set_verbosity_error()  # quench safety check warnings
+            pipeline_output = pipeline.img2img_from_embeddings(
+                init_image,
+                strength,
+                steps,
+                conditioning_data,
+                noise_func=self.get_noise_like,
+                callback=step_callback,
+                seed=seed,
+            )
+            if (
+                pipeline_output.attention_map_saver is not None
+                and attention_maps_callback is not None
+            ):
+                attention_maps_callback(pipeline_output.attention_map_saver)
+            return pipeline.numpy_to_pil(pipeline_output.images)[0]
+
+        return make_image
+
+    def get_noise_like(self, like: torch.Tensor):
+        device = like.device
+        if device.type == "mps":
+            x = torch.randn_like(like, device="cpu").to(device)
+        else:
+            x = torch.randn_like(like, device=device)
+        if self.perlin > 0.0:
+            shape = like.shape
+            x = (1 - self.perlin) * x + self.perlin * self.get_perlin_noise(
+                shape[3], shape[2]
+            )
+        return x

invokeai/backend/generator/inpaint.py

@@ -1,32 +1,35 @@
-'''
-ldm.invoke.generator.inpaint descends from ldm.invoke.generator
-'''
+"""
+invokeai.backend.generator.inpaint descends from .generator
+"""
 from __future__ import annotations
 
 import math
 
-import PIL
 import cv2
 import numpy as np
+import PIL
 import torch
-from PIL import Image, ImageFilter, ImageOps, ImageChops
+from PIL import Image, ImageChops, ImageFilter, ImageOps
 
-from ldm.invoke.generator.diffusers_pipeline import image_resized_to_grid_as_tensor, StableDiffusionGeneratorPipeline, \
-    ConditioningData
-from ldm.invoke.generator.img2img import Img2Img
-from ldm.invoke.patchmatch import PatchMatch
-from ldm.util import debug_image
+from ..image_util import PatchMatch, debug_image
+from ..stable_diffusion.diffusers_pipeline import (
+    ConditioningData,
+    StableDiffusionGeneratorPipeline,
+    image_resized_to_grid_as_tensor,
+)
+from .img2img import Img2Img
 
 
-def infill_methods()->list[str]:
+def infill_methods() -> list[str]:
     methods = [
         "tile",
         "solid",
     ]
     if PatchMatch.patchmatch_available():
-        methods.insert(0, 'patchmatch')
+        methods.insert(0, "patchmatch")
     return methods
 
+
 class Inpaint(Img2Img):
     def __init__(self, model, precision):
         self.inpaint_height = 0
@@ -53,11 +56,11 @@ class Inpaint(Img2Img):
             np.ravel(image),
             shape=(nrows, ncols, height, width, depth),
             strides=(height * _strides[0], width * _strides[1], *_strides),
-            writeable=False
+            writeable=False,
         )
 
     def infill_patchmatch(self, im: Image.Image) -> Image:
-        if im.mode != 'RGBA':
+        if im.mode != "RGBA":
             return im
 
         # Skip patchmatch if patchmatch isn't available
@@ -65,13 +68,17 @@ class Inpaint(Img2Img):
             return im
 
         # Patchmatch (note, we may want to expose patch_size? Increasing it significantly impacts performance though)
-        im_patched_np = PatchMatch.inpaint(im.convert('RGB'), ImageOps.invert(im.split()[-1]), patch_size = 3)
-        im_patched = Image.fromarray(im_patched_np, mode = 'RGB')
+        im_patched_np = PatchMatch.inpaint(
+            im.convert("RGB"), ImageOps.invert(im.split()[-1]), patch_size=3
+        )
+        im_patched = Image.fromarray(im_patched_np, mode="RGB")
         return im_patched
 
-    def tile_fill_missing(self, im: Image.Image, tile_size: int = 16, seed: int = None) -> Image:
+    def tile_fill_missing(
+        self, im: Image.Image, tile_size: int = 16, seed: int = None
+    ) -> Image:
         # Only fill if there's an alpha layer
-        if im.mode != 'RGBA':
+        if im.mode != "RGBA":
             return im
 
         a = np.asarray(im, dtype=np.uint8)
@@ -79,21 +86,21 @@ class Inpaint(Img2Img):
         tile_size = (tile_size, tile_size)
 
         # Get the image as tiles of a specified size
-        tiles = self.get_tile_images(a,*tile_size).copy()
+        tiles = self.get_tile_images(a, *tile_size).copy()
 
         # Get the mask as tiles
-        tiles_mask = tiles[:,:,:,:,3]
+        tiles_mask = tiles[:, :, :, :, 3]
 
         # Find any mask tiles with any fully transparent pixels (we will be replacing these later)
         tmask_shape = tiles_mask.shape
         tiles_mask = tiles_mask.reshape(math.prod(tiles_mask.shape))
-        n,ny = (math.prod(tmask_shape[0:2])), math.prod(tmask_shape[2:])
-        tiles_mask = (tiles_mask > 0)
-        tiles_mask = tiles_mask.reshape((n,ny)).all(axis = 1)
+        n, ny = (math.prod(tmask_shape[0:2])), math.prod(tmask_shape[2:])
+        tiles_mask = tiles_mask > 0
+        tiles_mask = tiles_mask.reshape((n, ny)).all(axis=1)
 
         # Get RGB tiles in single array and filter by the mask
         tshape = tiles.shape
-        tiles_all = tiles.reshape((math.prod(tiles.shape[0:2]), * tiles.shape[2:]))
+        tiles_all = tiles.reshape((math.prod(tiles.shape[0:2]), *tiles.shape[2:]))
         filtered_tiles = tiles_all[tiles_mask]
 
         if len(filtered_tiles) == 0:
@@ -101,23 +108,32 @@ class Inpaint(Img2Img):
 
         # Find all invalid tiles and replace with a random valid tile
         replace_count = (tiles_mask == False).sum()
-        rng = np.random.default_rng(seed = seed)
-        tiles_all[np.logical_not(tiles_mask)] = filtered_tiles[rng.choice(filtered_tiles.shape[0], replace_count),:,:,:]
+        rng = np.random.default_rng(seed=seed)
+        tiles_all[np.logical_not(tiles_mask)] = filtered_tiles[
+            rng.choice(filtered_tiles.shape[0], replace_count), :, :, :
+        ]
 
         # Convert back to an image
         tiles_all = tiles_all.reshape(tshape)
-        tiles_all = tiles_all.swapaxes(1,2)
-        st = tiles_all.reshape((math.prod(tiles_all.shape[0:2]), math.prod(tiles_all.shape[2:4]), tiles_all.shape[4]))
-        si = Image.fromarray(st, mode='RGBA')
+        tiles_all = tiles_all.swapaxes(1, 2)
+        st = tiles_all.reshape(
+            (
+                math.prod(tiles_all.shape[0:2]),
+                math.prod(tiles_all.shape[2:4]),
+                tiles_all.shape[4],
+            )
+        )
+        si = Image.fromarray(st, mode="RGBA")
 
         return si
 
-
     def mask_edge(self, mask: Image, edge_size: int, edge_blur: int) -> Image:
         npimg = np.asarray(mask, dtype=np.uint8)
 
         # Detect any partially transparent regions
-        npgradient = np.uint8(255 * (1.0 - np.floor(np.abs(0.5 - np.float32(npimg) / 255.0) * 2.0)))
+        npgradient = np.uint8(
+            255 * (1.0 - np.floor(np.abs(0.5 - np.float32(npimg) / 255.0) * 2.0))
+        )
 
         # Detect hard edges
         npedge = cv2.Canny(npimg, threshold1=100, threshold2=200)
@@ -126,7 +142,9 @@ class Inpaint(Img2Img):
         npmask = npgradient + npedge
 
         # Expand
-        npmask = cv2.dilate(npmask, np.ones((3,3), np.uint8), iterations = int(edge_size / 2))
+        npmask = cv2.dilate(
+            npmask, np.ones((3, 3), np.uint8), iterations=int(edge_size / 2)
+        )
 
         new_mask = Image.fromarray(npmask)
 
@@ -135,9 +153,23 @@ class Inpaint(Img2Img):
 
         return ImageOps.invert(new_mask)
 
-
-    def seam_paint(self, im: Image.Image, seam_size: int, seam_blur: int, prompt, sampler, steps, cfg_scale, ddim_eta,
-                   conditioning, strength, noise, infill_method, step_callback) -> Image.Image:
+    def seam_paint(
+        self,
+        im: Image.Image,
+        seam_size: int,
+        seam_blur: int,
+        prompt,
+        seed,
+        sampler,
+        steps,
+        cfg_scale,
+        ddim_eta,
+        conditioning,
+        strength,
+        noise,
+        infill_method,
+        step_callback,
+    ) -> Image.Image:
         hard_mask = self.pil_image.split()[-1].copy()
         mask = self.mask_edge(hard_mask, seam_size, seam_blur)
 
@@ -148,45 +180,52 @@ class Inpaint(Img2Img):
             cfg_scale,
             ddim_eta,
             conditioning,
-            init_image = im.copy().convert('RGBA'),
-            mask_image = mask,
-            strength = strength,
-            mask_blur_radius = 0,
-            seam_size = 0,
-            step_callback = step_callback,
-            inpaint_width = im.width,
-            inpaint_height = im.height,
-            infill_method = infill_method
+            init_image=im.copy().convert("RGBA"),
+            mask_image=mask,
+            strength=strength,
+            mask_blur_radius=0,
+            seam_size=0,
+            step_callback=step_callback,
+            inpaint_width=im.width,
+            inpaint_height=im.height,
+            infill_method=infill_method,
         )
 
         seam_noise = self.get_noise(im.width, im.height)
 
-        result = make_image(seam_noise)
+        result = make_image(seam_noise, seed)
 
         return result
 
-
     @torch.no_grad()
-    def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta,
-                       conditioning,
-                       init_image: PIL.Image.Image | torch.FloatTensor,
-                       mask_image: PIL.Image.Image | torch.FloatTensor,
-                       strength: float,
-                       mask_blur_radius: int = 8,
-                       # Seam settings - when 0, doesn't fill seam
-                       seam_size: int = 0,
-                       seam_blur: int = 0,
-                       seam_strength: float = 0.7,
-                       seam_steps: int = 10,
-                       tile_size: int = 32,
-                       step_callback=None,
-                       inpaint_replace=False, enable_image_debugging=False,
-                       infill_method = None,
-                       inpaint_width=None,
-                       inpaint_height=None,
-                       inpaint_fill:tuple(int)=(0x7F, 0x7F, 0x7F, 0xFF),
-                       attention_maps_callback=None,
-                       **kwargs):
+    def get_make_image(
+        self,
+        prompt,
+        sampler,
+        steps,
+        cfg_scale,
+        ddim_eta,
+        conditioning,
+        init_image: PIL.Image.Image | torch.FloatTensor,
+        mask_image: PIL.Image.Image | torch.FloatTensor,
+        strength: float,
+        mask_blur_radius: int = 8,
+        # Seam settings - when 0, doesn't fill seam
+        seam_size: int = 0,
+        seam_blur: int = 0,
+        seam_strength: float = 0.7,
+        seam_steps: int = 10,
+        tile_size: int = 32,
+        step_callback=None,
+        inpaint_replace=False,
+        enable_image_debugging=False,
+        infill_method=None,
+        inpaint_width=None,
+        inpaint_height=None,
+        inpaint_fill: tuple(int) = (0x7F, 0x7F, 0x7F, 0xFF),
+        attention_maps_callback=None,
+        **kwargs,
+    ):
         """
         Returns a function returning an image derived from the prompt and
         the initial image + mask.  Return value depends on the seed at
@@ -204,33 +243,39 @@ class Inpaint(Img2Img):
             self.pil_image = init_image.copy()
 
             # Do infill
-            if infill_method == 'patchmatch' and PatchMatch.patchmatch_available():
+            if infill_method == "patchmatch" and PatchMatch.patchmatch_available():
                 init_filled = self.infill_patchmatch(self.pil_image.copy())
-            elif infill_method == 'tile':
+            elif infill_method == "tile":
                 init_filled = self.tile_fill_missing(
-                    self.pil_image.copy(),
-                    seed = self.seed,
-                    tile_size = tile_size
+                    self.pil_image.copy(), seed=self.seed, tile_size=tile_size
                 )
-            elif infill_method == 'solid':
+            elif infill_method == "solid":
                 solid_bg = PIL.Image.new("RGBA", init_image.size, inpaint_fill)
                 init_filled = PIL.Image.alpha_composite(solid_bg, init_image)
             else:
-                raise ValueError(f"Non-supported infill type {infill_method}", infill_method)
-            init_filled.paste(init_image, (0,0), init_image.split()[-1])
+                raise ValueError(
+                    f"Non-supported infill type {infill_method}", infill_method
+                )
+            init_filled.paste(init_image, (0, 0), init_image.split()[-1])
 
             # Resize if requested for inpainting
             if inpaint_width and inpaint_height:
                 init_filled = init_filled.resize((inpaint_width, inpaint_height))
 
-            debug_image(init_filled, "init_filled", debug_status=self.enable_image_debugging)
+            debug_image(
+                init_filled, "init_filled", debug_status=self.enable_image_debugging
+            )
 
             # Create init tensor
-            init_image = image_resized_to_grid_as_tensor(init_filled.convert('RGB'))
+            init_image = image_resized_to_grid_as_tensor(init_filled.convert("RGB"))
 
         if isinstance(mask_image, PIL.Image.Image):
             self.pil_mask = mask_image.copy()
-            debug_image(mask_image, "mask_image BEFORE multiply with pil_image", debug_status=self.enable_image_debugging)
+            debug_image(
+                mask_image,
+                "mask_image BEFORE multiply with pil_image",
+                debug_status=self.enable_image_debugging,
+            )
 
             init_alpha = self.pil_image.getchannel("A")
             if mask_image.mode != "L":
@@ -243,8 +288,14 @@ class Inpaint(Img2Img):
             if inpaint_width and inpaint_height:
                 mask_image = mask_image.resize((inpaint_width, inpaint_height))
 
-            debug_image(mask_image, "mask_image AFTER multiply with pil_image", debug_status=self.enable_image_debugging)
-            mask: torch.FloatTensor = image_resized_to_grid_as_tensor(mask_image, normalize=False)
+            debug_image(
+                mask_image,
+                "mask_image AFTER multiply with pil_image",
+                debug_status=self.enable_image_debugging,
+            )
+            mask: torch.FloatTensor = image_resized_to_grid_as_tensor(
+                mask_image, normalize=False
+            )
         else:
             mask: torch.FloatTensor = mask_image
 
@@ -256,11 +307,11 @@ class Inpaint(Img2Img):
 
         # todo: support cross-attention control
         uc, c, _ = conditioning
-        conditioning_data = (ConditioningData(uc, c, cfg_scale)
-                             .add_scheduler_args_if_applicable(pipeline.scheduler, eta=ddim_eta))
+        conditioning_data = ConditioningData(
+            uc, c, cfg_scale
+        ).add_scheduler_args_if_applicable(pipeline.scheduler, eta=ddim_eta)
 
-
-        def make_image(x_T):
+        def make_image(x_T: torch.Tensor, seed: int):
             pipeline_output = pipeline.inpaint_from_embeddings(
                 init_image=init_image,
                 mask=1 - mask,  # expects white means "paint here."
@@ -269,45 +320,75 @@ class Inpaint(Img2Img):
                 conditioning_data=conditioning_data,
                 noise_func=self.get_noise_like,
                 callback=step_callback,
+                seed=seed,
             )
 
-            if pipeline_output.attention_map_saver is not None and attention_maps_callback is not None:
+            if (
+                pipeline_output.attention_map_saver is not None
+                and attention_maps_callback is not None
+            ):
                 attention_maps_callback(pipeline_output.attention_map_saver)
 
-            result = self.postprocess_size_and_mask(pipeline.numpy_to_pil(pipeline_output.images)[0])
+            result = self.postprocess_size_and_mask(
+                pipeline.numpy_to_pil(pipeline_output.images)[0]
+            )
 
             # Seam paint if this is our first pass (seam_size set to 0 during seam painting)
             if seam_size > 0:
                 old_image = self.pil_image or init_image
                 old_mask = self.pil_mask or mask_image
 
-                result = self.seam_paint(result, seam_size, seam_blur, prompt, sampler, seam_steps, cfg_scale, ddim_eta,
-                                         conditioning, seam_strength, x_T, infill_method, step_callback)
+                result = self.seam_paint(
+                    result,
+                    seam_size,
+                    seam_blur,
+                    prompt,
+                    seed,
+                    sampler,
+                    seam_steps,
+                    cfg_scale,
+                    ddim_eta,
+                    conditioning,
+                    seam_strength,
+                    x_T,
+                    infill_method,
+                    step_callback,
+                )
 
                 # Restore original settings
-                self.get_make_image(prompt,sampler,steps,cfg_scale,ddim_eta,
-                       conditioning,
-                       old_image,
-                       old_mask,
-                       strength,
-                       mask_blur_radius, seam_size, seam_blur, seam_strength,
-                       seam_steps, tile_size, step_callback,
-                       inpaint_replace, enable_image_debugging,
-                       inpaint_width = inpaint_width,
-                       inpaint_height = inpaint_height,
-                       infill_method = infill_method,
-                       **kwargs)
+                self.get_make_image(
+                    prompt,
+                    sampler,
+                    steps,
+                    cfg_scale,
+                    ddim_eta,
+                    conditioning,
+                    old_image,
+                    old_mask,
+                    strength,
+                    mask_blur_radius,
+                    seam_size,
+                    seam_blur,
+                    seam_strength,
+                    seam_steps,
+                    tile_size,
+                    step_callback,
+                    inpaint_replace,
+                    enable_image_debugging,
+                    inpaint_width=inpaint_width,
+                    inpaint_height=inpaint_height,
+                    infill_method=infill_method,
+                    **kwargs,
+                )
 
             return result
 
         return make_image
 
-
-    def sample_to_image(self, samples)->Image.Image:
-        gen_result = super().sample_to_image(samples).convert('RGB')
+    def sample_to_image(self, samples) -> Image.Image:
+        gen_result = super().sample_to_image(samples).convert("RGB")
         return self.postprocess_size_and_mask(gen_result)
 
-
     def postprocess_size_and_mask(self, gen_result: Image.Image) -> Image.Image:
         debug_image(gen_result, "gen_result", debug_status=self.enable_image_debugging)
 
@@ -318,7 +399,13 @@ class Inpaint(Img2Img):
         if self.pil_image is None or self.pil_mask is None:
             return gen_result
 
-        corrected_result = self.repaste_and_color_correct(gen_result, self.pil_image, self.pil_mask, self.mask_blur_radius)
-        debug_image(corrected_result, "corrected_result", debug_status=self.enable_image_debugging)
+        corrected_result = self.repaste_and_color_correct(
+            gen_result, self.pil_image, self.pil_mask, self.mask_blur_radius
+        )
+        debug_image(
+            corrected_result,
+            "corrected_result",
+            debug_status=self.enable_image_debugging,
+        )
 
         return corrected_result

invokeai/backend/generator/txt2img.py

@@ -0,0 +1,81 @@
+"""
+invokeai.backend.generator.txt2img inherits from invokeai.backend.generator
+"""
+import PIL.Image
+import torch
+
+from ..stable_diffusion import (
+    ConditioningData,
+    PostprocessingSettings,
+    StableDiffusionGeneratorPipeline,
+)
+from .base import Generator
+
+
+class Txt2Img(Generator):
+    def __init__(self, model, precision):
+        super().__init__(model, precision)
+
+    @torch.no_grad()
+    def get_make_image(
+        self,
+        prompt,
+        sampler,
+        steps,
+        cfg_scale,
+        ddim_eta,
+        conditioning,
+        width,
+        height,
+        step_callback=None,
+        threshold=0.0,
+        warmup=0.2,
+        perlin=0.0,
+        h_symmetry_time_pct=None,
+        v_symmetry_time_pct=None,
+        attention_maps_callback=None,
+        **kwargs,
+    ):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it
+        kwargs are 'width' and 'height'
+        """
+        self.perlin = perlin
+
+        # noinspection PyTypeChecker
+        pipeline: StableDiffusionGeneratorPipeline = self.model
+        pipeline.scheduler = sampler
+
+        uc, c, extra_conditioning_info = conditioning
+        conditioning_data = ConditioningData(
+            uc,
+            c,
+            cfg_scale,
+            extra_conditioning_info,
+            postprocessing_settings=PostprocessingSettings(
+                threshold=threshold,
+                warmup=warmup,
+                h_symmetry_time_pct=h_symmetry_time_pct,
+                v_symmetry_time_pct=v_symmetry_time_pct,
+            ),
+        ).add_scheduler_args_if_applicable(pipeline.scheduler, eta=ddim_eta)
+
+        def make_image(x_T: torch.Tensor, _: int) -> PIL.Image.Image:
+            pipeline_output = pipeline.image_from_embeddings(
+                latents=torch.zeros_like(x_T, dtype=self.torch_dtype()),
+                noise=x_T,
+                num_inference_steps=steps,
+                conditioning_data=conditioning_data,
+                callback=step_callback,
+            )
+
+            if (
+                pipeline_output.attention_map_saver is not None
+                and attention_maps_callback is not None
+            ):
+                attention_maps_callback(pipeline_output.attention_map_saver)
+
+            return pipeline.numpy_to_pil(pipeline_output.images)[0]
+
+        return make_image

invokeai/backend/generator/txt2img2img.py

@@ -1,6 +1,6 @@
-'''
-ldm.invoke.generator.txt2img inherits from ldm.invoke.generator
-'''
+"""
+invokeai.backend.generator.txt2img inherits from invokeai.backend.generator
+"""
 
 import math
 from typing import Callable, Optional
@@ -8,21 +8,37 @@ from typing import Callable, Optional
 import torch
 from diffusers.utils.logging import get_verbosity, set_verbosity, set_verbosity_error
 
-from ldm.invoke.generator.base import Generator
-from ldm.invoke.generator.diffusers_pipeline import trim_to_multiple_of, StableDiffusionGeneratorPipeline, \
-    ConditioningData
-from ldm.models.diffusion.shared_invokeai_diffusion import PostprocessingSettings
-
+from ..stable_diffusion import PostprocessingSettings
+from .base import Generator
+from ..stable_diffusion.diffusers_pipeline import StableDiffusionGeneratorPipeline
+from ..stable_diffusion.diffusers_pipeline import ConditioningData
+from ..stable_diffusion.diffusers_pipeline import trim_to_multiple_of
 
 class Txt2Img2Img(Generator):
     def __init__(self, model, precision):
         super().__init__(model, precision)
-        self.init_latent = None    # for get_noise()
+        self.init_latent = None  # for get_noise()
 
-    def get_make_image(self, prompt:str, sampler, steps:int, cfg_scale:float, ddim_eta,
-                       conditioning, width:int, height:int, strength:float,
-                       step_callback:Optional[Callable]=None, threshold=0.0, warmup=0.2, perlin=0.0,
-                       h_symmetry_time_pct=None, v_symmetry_time_pct=None, attention_maps_callback=None, **kwargs):
+    def get_make_image(
+        self,
+        prompt: str,
+        sampler,
+        steps: int,
+        cfg_scale: float,
+        ddim_eta,
+        conditioning,
+        width: int,
+        height: int,
+        strength: float,
+        step_callback: Optional[Callable] = None,
+        threshold=0.0,
+        warmup=0.2,
+        perlin=0.0,
+        h_symmetry_time_pct=None,
+        v_symmetry_time_pct=None,
+        attention_maps_callback=None,
+        **kwargs,
+    ):
         """
         Returns a function returning an image derived from the prompt and the initial image
         Return value depends on the seed at the time you call it
@@ -35,19 +51,20 @@ class Txt2Img2Img(Generator):
         pipeline.scheduler = sampler
 
         uc, c, extra_conditioning_info = conditioning
-        conditioning_data = (
-            ConditioningData(
-                uc, c, cfg_scale, extra_conditioning_info,
-                postprocessing_settings = PostprocessingSettings(
-                    threshold=threshold,
-                    warmup=0.2,
-                    h_symmetry_time_pct=h_symmetry_time_pct,
-                    v_symmetry_time_pct=v_symmetry_time_pct
-                )
-            ).add_scheduler_args_if_applicable(pipeline.scheduler, eta=ddim_eta))
-
-        def make_image(x_T):
+        conditioning_data = ConditioningData(
+            uc,
+            c,
+            cfg_scale,
+            extra_conditioning_info,
+            postprocessing_settings=PostprocessingSettings(
+                threshold=threshold,
+                warmup=0.2,
+                h_symmetry_time_pct=h_symmetry_time_pct,
+                v_symmetry_time_pct=v_symmetry_time_pct,
+            ),
+        ).add_scheduler_args_if_applicable(pipeline.scheduler, eta=ddim_eta)
 
+        def make_image(x_T: torch.Tensor, _: int):
             first_pass_latent_output, _ = pipeline.latents_from_embeddings(
                 latents=torch.zeros_like(x_T),
                 num_inference_steps=steps,
@@ -61,28 +78,40 @@ class Txt2Img2Img(Generator):
             init_width = first_pass_latent_output.size()[3] * self.downsampling_factor
             init_height = first_pass_latent_output.size()[2] * self.downsampling_factor
             print(
-                  f"\n>> Interpolating from {init_width}x{init_height} to {width}x{height} using DDIM sampling"
-                 )
+                f"\n>> Interpolating from {init_width}x{init_height} to {width}x{height} using DDIM sampling"
+            )
 
             # resizing
             resized_latents = torch.nn.functional.interpolate(
                 first_pass_latent_output,
-                size=(height // self.downsampling_factor, width // self.downsampling_factor),
-                mode="bilinear"
+                size=(
+                    height // self.downsampling_factor,
+                    width // self.downsampling_factor,
+                ),
+                mode="bilinear",
             )
 
             # Free up memory from the last generation.
-            clear_cuda_cache = kwargs['clear_cuda_cache'] or None
+            clear_cuda_cache = kwargs["clear_cuda_cache"] or None
             if clear_cuda_cache is not None:
                 clear_cuda_cache()
 
-            second_pass_noise = self.get_noise_like(resized_latents, override_perlin=True)
+            second_pass_noise = self.get_noise_like(
+                resized_latents, override_perlin=True
+            )
 
             # Clear symmetry for the second pass
             from dataclasses import replace
-            new_postprocessing_settings = replace(conditioning_data.postprocessing_settings, h_symmetry_time_pct=None)
-            new_postprocessing_settings = replace(new_postprocessing_settings, v_symmetry_time_pct=None)
-            new_conditioning_data = replace(conditioning_data, postprocessing_settings=new_postprocessing_settings)
+
+            new_postprocessing_settings = replace(
+                conditioning_data.postprocessing_settings, h_symmetry_time_pct=None
+            )
+            new_postprocessing_settings = replace(
+                new_postprocessing_settings, v_symmetry_time_pct=None
+            )
+            new_conditioning_data = replace(
+                conditioning_data, postprocessing_settings=new_postprocessing_settings
+            )
 
             verbosity = get_verbosity()
             set_verbosity_error()
@@ -92,15 +121,18 @@ class Txt2Img2Img(Generator):
                 conditioning_data=new_conditioning_data,
                 strength=strength,
                 noise=second_pass_noise,
-                callback=step_callback)
+                callback=step_callback,
+            )
             set_verbosity(verbosity)
 
-            if pipeline_output.attention_map_saver is not None and attention_maps_callback is not None:
+            if (
+                pipeline_output.attention_map_saver is not None
+                and attention_maps_callback is not None
+            ):
                 attention_maps_callback(pipeline_output.attention_map_saver)
 
             return pipeline.numpy_to_pil(pipeline_output.images)[0]
 
-
         # FIXME: do we really need something entirely different for the inpainting model?
 
         # in the case of the inpainting model being loaded, the trick of
@@ -111,19 +143,23 @@ class Txt2Img2Img(Generator):
 
         return make_image
 
-    def get_noise_like(self, like: torch.Tensor, override_perlin: bool=False):
+    def get_noise_like(self, like: torch.Tensor, override_perlin: bool = False):
         device = like.device
-        if device.type == 'mps':
-            x = torch.randn_like(like, device='cpu', dtype=self.torch_dtype()).to(device)
+        if device.type == "mps":
+            x = torch.randn_like(like, device="cpu", dtype=self.torch_dtype()).to(
+                device
+            )
         else:
             x = torch.randn_like(like, device=device, dtype=self.torch_dtype())
         if self.perlin > 0.0 and override_perlin == False:
             shape = like.shape
-            x = (1-self.perlin)*x + self.perlin*self.get_perlin_noise(shape[3], shape[2])
+            x = (1 - self.perlin) * x + self.perlin * self.get_perlin_noise(
+                shape[3], shape[2]
+            )
         return x
 
     # returns a tensor filled with random numbers from a normal distribution
-    def get_noise(self,width,height,scale = True):
+    def get_noise(self, width, height, scale=True):
         # print(f"Get noise: {width}x{height}")
         if scale:
             # Scale the input width and height for the initial generation
@@ -133,7 +169,9 @@ class Txt2Img2Img(Generator):
             aspect = width / height
             dimension = self.model.unet.config.sample_size * self.model.vae_scale_factor
             min_dimension = math.floor(dimension * 0.5)
-            model_area = dimension * dimension # hardcoded for now since all models are trained on square images
+            model_area = (
+                dimension * dimension
+            )  # hardcoded for now since all models are trained on square images
 
             if aspect > 1.0:
                 init_height = max(min_dimension, math.sqrt(model_area / aspect))
@@ -142,7 +180,9 @@ class Txt2Img2Img(Generator):
                 init_width = max(min_dimension, math.sqrt(model_area * aspect))
                 init_height = init_width / aspect
 
-            scaled_width, scaled_height = trim_to_multiple_of(math.floor(init_width), math.floor(init_height))
+            scaled_width, scaled_height = trim_to_multiple_of(
+                math.floor(init_width), math.floor(init_height)
+            )
 
         else:
             scaled_width = width
@@ -152,10 +192,14 @@ class Txt2Img2Img(Generator):
         channels = self.latent_channels
         if channels == 9:
             channels = 4  # we don't really want noise for all the mask channels
-        shape = (1, channels,
-                 scaled_height // self.downsampling_factor, scaled_width // self.downsampling_factor)
-        if self.use_mps_noise or device.type == 'mps':
-            tensor = torch.empty(size=shape, device='cpu')
+        shape = (
+            1,
+            channels,
+            scaled_height // self.downsampling_factor,
+            scaled_width // self.downsampling_factor,
+        )
+        if self.use_mps_noise or device.type == "mps":
+            tensor = torch.empty(size=shape, device="cpu")
             tensor = self.get_noise_like(like=tensor).to(device)
         else:
             tensor = torch.empty(size=shape, device=device)

invokeai/backend/globals.py

@@ -1,5 +1,5 @@
-'''
-ldm.invoke.globals defines a small number of global variables that would
+"""
+invokeai.backend.globals defines a small number of global variables that would
 otherwise have to be passed through long and complex call chains.
 
 It defines a Namespace object named "Globals" that contains
@@ -9,7 +9,7 @@ the attributes:
   - initfile       - path to the initialization file
   - try_patchmatch - option to globally disable loading of 'patchmatch' module
   - always_use_cpu - force use of CPU even if GPU is available
-'''
+"""
 
 import os
 import os.path as osp
@@ -20,12 +20,12 @@ from typing import Union
 Globals = Namespace()
 
 # Where to look for the initialization file and other key components
-Globals.initfile = 'invokeai.init'
-Globals.models_file = 'models.yaml'
-Globals.models_dir = 'models'
-Globals.config_dir = 'configs'
-Globals.autoscan_dir = 'weights'
-Globals.converted_ckpts_dir = 'converted_ckpts'
+Globals.initfile = "invokeai.init"
+Globals.models_file = "models.yaml"
+Globals.models_dir = "models"
+Globals.config_dir = "configs"
+Globals.autoscan_dir = "weights"
+Globals.converted_ckpts_dir = "converted_ckpts"
 
 # Set the default root directory. This can be overwritten by explicitly
 # passing the `--root <directory>` argument on the command line.
@@ -34,12 +34,15 @@ Globals.converted_ckpts_dir = 'converted_ckpts'
 # 2) use VIRTUAL_ENV environment variable, with a check for initfile being there
 # 3) use ~/invokeai
 
-if os.environ.get('INVOKEAI_ROOT'):
-    Globals.root = osp.abspath(os.environ.get('INVOKEAI_ROOT'))
-elif os.environ.get('VIRTUAL_ENV') and Path(os.environ.get('VIRTUAL_ENV'),'..',Globals.initfile).exists():
-    Globals.root = osp.abspath(osp.join(os.environ.get('VIRTUAL_ENV'), '..'))
+if os.environ.get("INVOKEAI_ROOT"):
+    Globals.root = osp.abspath(os.environ.get("INVOKEAI_ROOT"))
+elif (
+    os.environ.get("VIRTUAL_ENV")
+    and Path(os.environ.get("VIRTUAL_ENV"), "..", Globals.initfile).exists()
+):
+    Globals.root = osp.abspath(osp.join(os.environ.get("VIRTUAL_ENV"), ".."))
 else:
-    Globals.root = osp.abspath(osp.expanduser('~/invokeai'))
+    Globals.root = osp.abspath(osp.expanduser("~/invokeai"))
 
 # Try loading patchmatch
 Globals.try_patchmatch = True
@@ -61,55 +64,59 @@ Globals.sequential_guidance = False
 Globals.full_precision = False
 
 # whether we should convert ckpt files into diffusers models on the fly
-Globals.ckpt_convert = False
+Globals.ckpt_convert = True
 
 # logging tokenization everywhere
 Globals.log_tokenization = False
 
-def global_config_file()->Path:
+
+def global_config_file() -> Path:
     return Path(Globals.root, Globals.config_dir, Globals.models_file)
 
-def global_config_dir()->Path:
+
+def global_config_dir() -> Path:
     return Path(Globals.root, Globals.config_dir)
 
-def global_models_dir()->Path:
+
+def global_models_dir() -> Path:
     return Path(Globals.root, Globals.models_dir)
 
-def global_autoscan_dir()->Path:
+
+def global_autoscan_dir() -> Path:
     return Path(Globals.root, Globals.autoscan_dir)
 
-def global_converted_ckpts_dir()->Path:
+
+def global_converted_ckpts_dir() -> Path:
     return Path(global_models_dir(), Globals.converted_ckpts_dir)
 
-def global_set_root(root_dir:Union[str,Path]):
+
+def global_set_root(root_dir: Union[str, Path]):
     Globals.root = root_dir
 
-def global_cache_dir(subdir:Union[str,Path]='')->Path:
-    '''
+
+def global_cache_dir(subdir: Union[str, Path] = "") -> Path:
+    """
     Returns Path to the model cache directory. If a subdirectory
     is provided, it will be appended to the end of the path, allowing
-    for huggingface-style conventions:
-         global_cache_dir('diffusers')
+    for Hugging Face-style conventions. Currently, Hugging Face has
+    moved all models into the "hub" subfolder, so for any pretrained
+    HF model, use:
          global_cache_dir('hub')
-    Current HuggingFace documentation (mid-Jan 2023) indicates that
-    transformers models will be cached into a "transformers" subdirectory,
-    but in practice they seem to go into "hub". But if needed:
-         global_cache_dir('transformers')
-    One other caveat is that HuggingFace is moving some diffusers models
-    into the "hub" subdirectory as well, so this will need to be revisited
-    from time to time.
-    '''
-    home: str = os.getenv('HF_HOME')
+
+    The legacy location for transformers used to be global_cache_dir('transformers')
+    and global_cache_dir('diffusers') for diffusers.
+    """
+    home: str = os.getenv("HF_HOME")
 
     if home is None:
-        home = os.getenv('XDG_CACHE_HOME')
+        home = os.getenv("XDG_CACHE_HOME")
 
         if home is not None:
-            # Set `home` to $XDG_CACHE_HOME/huggingface, which is the default location mentioned in HuggingFace Hub Client Library.
+            # Set `home` to $XDG_CACHE_HOME/huggingface, which is the default location mentioned in Hugging Face Hub Client Library.
             # See: https://huggingface.co/docs/huggingface_hub/main/en/package_reference/environment_variables#xdgcachehome
-            home += os.sep + 'huggingface'
+            home += os.sep + "huggingface"
 
     if home is not None:
-        return Path(home,subdir)
+        return Path(home, subdir)
     else:
-        return Path(Globals.root,'models',subdir)
+        return Path(Globals.root, "models", subdir)

invokeai/backend/image_util/__init__.py

@@ -0,0 +1,24 @@
+"""
+Initialization file for invokeai.backend.image_util methods.
+"""
+from .patchmatch import PatchMatch
+from .pngwriter import PngWriter, PromptFormatter, retrieve_metadata, write_metadata
+from .seamless import configure_model_padding
+from .txt2mask import Txt2Mask
+from .util import InitImageResizer, make_grid
+
+
+def debug_image(
+    debug_image, debug_text, debug_show=True, debug_result=False, debug_status=False
+):
+    if not debug_status:
+        return
+
+    image_copy = debug_image.copy().convert("RGBA")
+    ImageDraw.Draw(image_copy).text((5, 5), debug_text, (255, 0, 0))
+
+    if debug_show:
+        image_copy.show()
+
+    if debug_result:
+        return image_copy

invokeai/backend/image_util/patchmatch.py

@@ -1,20 +1,22 @@
-'''
+"""
 This module defines a singleton object, "patchmatch" that
 wraps the actual patchmatch object. It respects the global
 "try_patchmatch" attribute, so that patchmatch loading can
 be suppressed or deferred
-'''
-from ldm.invoke.globals import Globals
-import numpy as  np
+"""
+import numpy as np
+
+from invokeai.backend.globals import Globals
+
 
 class PatchMatch:
-    '''
+    """
     Thin class wrapper around the patchmatch function.
-    '''
+    """
 
     patch_match = None
-    tried_load:bool = False
-    
+    tried_load: bool = False
+
     def __init__(self):
         super().__init__()
 
@@ -24,21 +26,22 @@ class PatchMatch:
             return
         if Globals.try_patchmatch:
             from patchmatch import patch_match as pm
+
             if pm.patchmatch_available:
-                print('>> Patchmatch initialized')
+                print(">> Patchmatch initialized")
             else:
-                print('>> Patchmatch not loaded (nonfatal)')
+                print(">> Patchmatch not loaded (nonfatal)")
             self.patch_match = pm
         else:
-            print('>> Patchmatch loading disabled')
+            print(">> Patchmatch loading disabled")
         self.tried_load = True
 
     @classmethod
-    def patchmatch_available(self)->bool:
+    def patchmatch_available(self) -> bool:
         self._load_patch_match()
         return self.patch_match and self.patch_match.patchmatch_available
 
     @classmethod
-    def inpaint(self,*args,**kwargs)->np.ndarray:
+    def inpaint(self, *args, **kwargs) -> np.ndarray:
         if self.patchmatch_available():
-            return self.patch_match.inpaint(*args,**kwargs)
+            return self.patch_match.inpaint(*args, **kwargs)

invokeai/backend/image_util/pngwriter.py

@@ -6,10 +6,11 @@ PngWriter -- Converts Images generated by T2I into PNGs, finds
 
 Exports function retrieve_metadata(path)
 """
+import json
 import os
 import re
-import json
-from PIL import PngImagePlugin, Image
+
+from PIL import Image, PngImagePlugin
 
 # -------------------image generation utils-----
 
@@ -25,52 +26,57 @@ class PngWriter:
         dirlist = sorted(os.listdir(self.outdir), reverse=True)
         # find the first filename that matches our pattern or return 000000.0.png
         existing_name = next(
-            (f for f in dirlist if re.match('^(\d+)\..*\.png', f)),
-            '0000000.0.png',
+            (f for f in dirlist if re.match("^(\d+)\..*\.png", f)),
+            "0000000.0.png",
         )
-        basecount = int(existing_name.split('.', 1)[0]) + 1
-        return f'{basecount:06}'
+        basecount = int(existing_name.split(".", 1)[0]) + 1
+        return f"{basecount:06}"
 
     # saves image named _image_ to outdir/name, writing metadata from prompt
     # returns full path of output
-    def save_image_and_prompt_to_png(self, image, dream_prompt, name, metadata=None, compress_level=6):
+    def save_image_and_prompt_to_png(
+        self, image, dream_prompt, name, metadata=None, compress_level=6
+    ):
         path = os.path.join(self.outdir, name)
         info = PngImagePlugin.PngInfo()
-        info.add_text('Dream', dream_prompt)
+        info.add_text("Dream", dream_prompt)
         if metadata:
-            info.add_text('sd-metadata', json.dumps(metadata))
-        image.save(path, 'PNG', pnginfo=info, compress_level=compress_level)
+            info.add_text("sd-metadata", json.dumps(metadata))
+        image.save(path, "PNG", pnginfo=info, compress_level=compress_level)
         return path
 
-    def retrieve_metadata(self,img_basename):
-        '''
+    def retrieve_metadata(self, img_basename):
+        """
         Given a PNG filename stored in outdir, returns the "sd-metadata"
         metadata stored there, as a dict
-        '''
-        path = os.path.join(self.outdir,img_basename)
+        """
+        path = os.path.join(self.outdir, img_basename)
         all_metadata = retrieve_metadata(path)
-        return all_metadata['sd-metadata']
+        return all_metadata["sd-metadata"]
+
 
 def retrieve_metadata(img_path):
-    '''
+    """
     Given a path to a PNG image, returns the "sd-metadata"
     metadata stored there, as a dict
-    '''
+    """
     im = Image.open(img_path)
-    if hasattr(im, 'text'):
-        md = im.text.get('sd-metadata', '{}')
-        dream_prompt = im.text.get('Dream', '')
+    if hasattr(im, "text"):
+        md = im.text.get("sd-metadata", "{}")
+        dream_prompt = im.text.get("Dream", "")
     else:
         # When trying to retrieve metadata from images without a 'text' payload, such as JPG images.
-        md = '{}'
-        dream_prompt = ''
-    return {'sd-metadata': json.loads(md), 'Dream': dream_prompt}
+        md = "{}"
+        dream_prompt = ""
+    return {"sd-metadata": json.loads(md), "Dream": dream_prompt}
 
-def write_metadata(img_path:str, meta:dict):
+
+def write_metadata(img_path: str, meta: dict):
     im = Image.open(img_path)
     info = PngImagePlugin.PngInfo()
-    info.add_text('sd-metadata', json.dumps(meta))
-    im.save(img_path,'PNG',pnginfo=info)
+    info.add_text("sd-metadata", json.dumps(meta))
+    im.save(img_path, "PNG", pnginfo=info)
+
 
 class PromptFormatter:
     def __init__(self, t2i, opt):
@@ -86,28 +92,30 @@ class PromptFormatter:
 
         switches = list()
         switches.append(f'"{opt.prompt}"')
-        switches.append(f'-s{opt.steps        or t2i.steps}')
-        switches.append(f'-W{opt.width        or t2i.width}')
-        switches.append(f'-H{opt.height       or t2i.height}')
-        switches.append(f'-C{opt.cfg_scale    or t2i.cfg_scale}')
-        switches.append(f'-A{opt.sampler_name or t2i.sampler_name}')
-# to do: put model name into the t2i object
-#        switches.append(f'--model{t2i.model_name}')
+        switches.append(f"-s{opt.steps        or t2i.steps}")
+        switches.append(f"-W{opt.width        or t2i.width}")
+        switches.append(f"-H{opt.height       or t2i.height}")
+        switches.append(f"-C{opt.cfg_scale    or t2i.cfg_scale}")
+        switches.append(f"-A{opt.sampler_name or t2i.sampler_name}")
+        # to do: put model name into the t2i object
+        #        switches.append(f'--model{t2i.model_name}')
         if opt.seamless or t2i.seamless:
-            switches.append(f'--seamless')
+            switches.append(f"--seamless")
         if opt.init_img:
-            switches.append(f'-I{opt.init_img}')
+            switches.append(f"-I{opt.init_img}")
         if opt.fit:
-            switches.append(f'--fit')
+            switches.append(f"--fit")
         if opt.strength and opt.init_img is not None:
-            switches.append(f'-f{opt.strength or t2i.strength}')
+            switches.append(f"-f{opt.strength or t2i.strength}")
         if opt.gfpgan_strength:
-            switches.append(f'-G{opt.gfpgan_strength}')
+            switches.append(f"-G{opt.gfpgan_strength}")
         if opt.upscale:
             switches.append(f'-U {" ".join([str(u) for u in opt.upscale])}')
         if opt.variation_amount > 0:
-            switches.append(f'-v{opt.variation_amount}')
+            switches.append(f"-v{opt.variation_amount}")
         if opt.with_variations:
-            formatted_variations = ','.join(f'{seed}:{weight}' for seed, weight in opt.with_variations)
-            switches.append(f'-V{formatted_variations}')
-        return ' '.join(switches)
+            formatted_variations = ",".join(
+                f"{seed}:{weight}" for seed, weight in opt.with_variations
+            )
+            switches.append(f"-V{formatted_variations}")
+        return " ".join(switches)

invokeai/backend/image_util/seamless.py

@@ -0,0 +1,59 @@
+import torch.nn as nn
+
+
+def _conv_forward_asymmetric(self, input, weight, bias):
+    """
+    Patch for Conv2d._conv_forward that supports asymmetric padding
+    """
+    working = nn.functional.pad(
+        input, self.asymmetric_padding["x"], mode=self.asymmetric_padding_mode["x"]
+    )
+    working = nn.functional.pad(
+        working, self.asymmetric_padding["y"], mode=self.asymmetric_padding_mode["y"]
+    )
+    return nn.functional.conv2d(
+        working,
+        weight,
+        bias,
+        self.stride,
+        nn.modules.utils._pair(0),
+        self.dilation,
+        self.groups,
+    )
+
+
+def configure_model_padding(model, seamless, seamless_axes):
+    """
+    Modifies the 2D convolution layers to use a circular padding mode based on the `seamless` and `seamless_axes` options.
+    """
+    # TODO: get an explicit interface for this in diffusers: https://github.com/huggingface/diffusers/issues/556
+    for m in model.modules():
+        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+            if seamless:
+                m.asymmetric_padding_mode = {}
+                m.asymmetric_padding = {}
+                m.asymmetric_padding_mode["x"] = (
+                    "circular" if ("x" in seamless_axes) else "constant"
+                )
+                m.asymmetric_padding["x"] = (
+                    m._reversed_padding_repeated_twice[0],
+                    m._reversed_padding_repeated_twice[1],
+                    0,
+                    0,
+                )
+                m.asymmetric_padding_mode["y"] = (
+                    "circular" if ("y" in seamless_axes) else "constant"
+                )
+                m.asymmetric_padding["y"] = (
+                    0,
+                    0,
+                    m._reversed_padding_repeated_twice[2],
+                    m._reversed_padding_repeated_twice[3],
+                )
+                m._conv_forward = _conv_forward_asymmetric.__get__(m, nn.Conv2d)
+            else:
+                m._conv_forward = nn.Conv2d._conv_forward.__get__(m, nn.Conv2d)
+                if hasattr(m, "asymmetric_padding_mode"):
+                    del m.asymmetric_padding_mode
+                if hasattr(m, "asymmetric_padding"):
+                    del m.asymmetric_padding

invokeai/backend/image_util/txt2mask.py

@@ -1,9 +1,9 @@
-'''Makes available the Txt2Mask class, which assists in the automatic
+"""Makes available the Txt2Mask class, which assists in the automatic
 assignment of masks via text prompt using clipseg.
 
 Here is typical usage:
 
-    from ldm.invoke.txt2mask import Txt2Mask, SegmentedGrayscale
+    from invokeai.backend.image_util.txt2mask import Txt2Mask, SegmentedGrayscale
     from PIL import Image
 
     txt2mask = Txt2Mask(self.device)
@@ -25,31 +25,39 @@ the mask that exceed the indicated confidence threshold. Values range
 from 0.0 to 1.0. The higher the threshold, the more confident the
 algorithm is. In limited testing, I have found that values around 0.5
 work fine.
-'''
+"""
 
+import numpy as np
 import torch
-import numpy as  np
-from transformers import AutoProcessor, CLIPSegForImageSegmentation
 from PIL import Image, ImageOps
 from torchvision import transforms
-from ldm.invoke.globals import global_cache_dir
+from transformers import AutoProcessor, CLIPSegForImageSegmentation
 
-CLIPSEG_MODEL = 'CIDAS/clipseg-rd64-refined'
+from invokeai.backend.globals import global_cache_dir
+
+CLIPSEG_MODEL = "CIDAS/clipseg-rd64-refined"
 CLIPSEG_SIZE = 352
 
+
 class SegmentedGrayscale(object):
-    def __init__(self, image:Image, heatmap:torch.Tensor):
+    def __init__(self, image: Image, heatmap: torch.Tensor):
         self.heatmap = heatmap
         self.image = image
 
-    def to_grayscale(self,invert:bool=False)->Image:
-        return self._rescale(Image.fromarray(np.uint8(255 - self.heatmap * 255 if invert else self.heatmap * 255)))
+    def to_grayscale(self, invert: bool = False) -> Image:
+        return self._rescale(
+            Image.fromarray(
+                np.uint8(255 - self.heatmap * 255 if invert else self.heatmap * 255)
+            )
+        )
 
-    def to_mask(self,threshold:float=0.5)->Image:
+    def to_mask(self, threshold: float = 0.5) -> Image:
         discrete_heatmap = self.heatmap.lt(threshold).int()
-        return self._rescale(Image.fromarray(np.uint8(discrete_heatmap*255),mode='L'))
+        return self._rescale(
+            Image.fromarray(np.uint8(discrete_heatmap * 255), mode="L")
+        )
 
-    def to_transparent(self,invert:bool=False)->Image:
+    def to_transparent(self, invert: bool = False) -> Image:
         transparent_image = self.image.copy()
         # For img2img, we want the selected regions to be transparent,
         # but to_grayscale() returns the opposite. Thus invert.
@@ -58,70 +66,77 @@ class SegmentedGrayscale(object):
         return transparent_image
 
     # unscales and uncrops the 352x352 heatmap so that it matches the image again
-    def _rescale(self, heatmap:Image)->Image:
-        size = self.image.width if (self.image.width > self.image.height) else self.image.height
-        resized_image = heatmap.resize(
-            (size,size),
-            resample=Image.Resampling.LANCZOS
+    def _rescale(self, heatmap: Image) -> Image:
+        size = (
+            self.image.width
+            if (self.image.width > self.image.height)
+            else self.image.height
         )
-        return resized_image.crop((0,0,self.image.width,self.image.height))
+        resized_image = heatmap.resize((size, size), resample=Image.Resampling.LANCZOS)
+        return resized_image.crop((0, 0, self.image.width, self.image.height))
+
 
 class Txt2Mask(object):
-    '''
+    """
     Create new Txt2Mask object. The optional device argument can be one of
     'cuda', 'mps' or 'cpu'.
-    '''
-    def __init__(self,device='cpu',refined=False):
-        print('>> Initializing clipseg model for text to mask inference')
+    """
+
+    def __init__(self, device="cpu", refined=False):
+        print(">> Initializing clipseg model for text to mask inference")
 
         # BUG: we are not doing anything with the device option at this time
         self.device = device
-        self.processor = AutoProcessor.from_pretrained(CLIPSEG_MODEL,
-                                                       cache_dir=global_cache_dir('hub')
-                                                       )
-        self.model = CLIPSegForImageSegmentation.from_pretrained(CLIPSEG_MODEL,
-                                                                 cache_dir=global_cache_dir('hub')
-                                                                 )
+        self.processor = AutoProcessor.from_pretrained(
+            CLIPSEG_MODEL, cache_dir=global_cache_dir("hub")
+        )
+        self.model = CLIPSegForImageSegmentation.from_pretrained(
+            CLIPSEG_MODEL, cache_dir=global_cache_dir("hub")
+        )
 
     @torch.no_grad()
-    def segment(self, image, prompt:str) -> SegmentedGrayscale:
-        '''
+    def segment(self, image, prompt: str) -> SegmentedGrayscale:
+        """
         Given a prompt string such as "a bagel", tries to identify the object in the
         provided image and returns a SegmentedGrayscale object in which the brighter
         pixels indicate where the object is inferred to be.
-        '''
-        transform = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-            transforms.Resize((CLIPSEG_SIZE, CLIPSEG_SIZE)), # must be multiple of 64...
-        ])
+        """
+        transform = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                ),
+                transforms.Resize(
+                    (CLIPSEG_SIZE, CLIPSEG_SIZE)
+                ),  # must be multiple of 64...
+            ]
+        )
 
         if type(image) is str:
-            image = Image.open(image).convert('RGB')
+            image = Image.open(image).convert("RGB")
 
         image = ImageOps.exif_transpose(image)
         img = self._scale_and_crop(image)
 
-        inputs = self.processor(text=[prompt],
-                                images=[img],
-                                padding=True,
-                                return_tensors='pt')
+        inputs = self.processor(
+            text=[prompt], images=[img], padding=True, return_tensors="pt"
+        )
         outputs = self.model(**inputs)
         heatmap = torch.sigmoid(outputs.logits)
         return SegmentedGrayscale(image, heatmap)
 
-    def _scale_and_crop(self, image:Image)->Image:
-        scaled_image = Image.new('RGB',(CLIPSEG_SIZE,CLIPSEG_SIZE))
-        if image.width > image.height: # width is constraint
+    def _scale_and_crop(self, image: Image) -> Image:
+        scaled_image = Image.new("RGB", (CLIPSEG_SIZE, CLIPSEG_SIZE))
+        if image.width > image.height:  # width is constraint
             scale = CLIPSEG_SIZE / image.width
         else:
             scale = CLIPSEG_SIZE / image.height
         scaled_image.paste(
             image.resize(
-                (int(scale * image.width),
-                 int(scale * image.height)
-                ),
-                resample=Image.Resampling.LANCZOS
-            ),box=(0,0)
+                (int(scale * image.width), int(scale * image.height)),
+                resample=Image.Resampling.LANCZOS,
+            ),
+            box=(0, 0),
         )
         return scaled_image

invokeai/backend/image_util/util.py

@@ -1,12 +1,15 @@
-from math import sqrt, floor, ceil
+from math import ceil, floor, sqrt
+
 from PIL import Image
 
-class InitImageResizer():
+
+class InitImageResizer:
     """Simple class to create resized copies of an Image while preserving the aspect ratio."""
-    def __init__(self,Image):
+
+    def __init__(self, Image):
         self.image = Image
 
-    def resize(self,width=None,height=None) -> Image:
+    def resize(self, width=None, height=None) -> Image:
         """
         Return a copy of the image resized to fit within
         a box width x height. The aspect ratio is
@@ -18,37 +21,36 @@ class InitImageResizer():
         Everything is floored to the nearest multiple of 64 so
         that it can be passed to img2img()
         """
-        im    = self.image
+        im = self.image
 
-        ar = im.width/float(im.height)
+        ar = im.width / float(im.height)
 
         # Infer missing values from aspect ratio
-        if not(width or height): # both missing
-            width  = im.width
+        if not (width or height):  # both missing
+            width = im.width
             height = im.height
-        elif not height:           # height missing
-            height = int(width/ar)
-        elif not width:            # width missing
-            width  = int(height*ar)
+        elif not height:  # height missing
+            height = int(width / ar)
+        elif not width:  # width missing
+            width = int(height * ar)
 
-        w_scale = width/im.width
-        h_scale = height/im.height
-        scale = min(w_scale,h_scale)
-        (rw,rh) = (int(scale*im.width),int(scale*im.height))
+        w_scale = width / im.width
+        h_scale = height / im.height
+        scale = min(w_scale, h_scale)
+        (rw, rh) = (int(scale * im.width), int(scale * im.height))
 
-        #round everything to multiples of 64
-        width,height,rw,rh = map(
-            lambda x: x-x%64, (width,height,rw,rh)
-        )
+        # round everything to multiples of 64
+        width, height, rw, rh = map(lambda x: x - x % 64, (width, height, rw, rh))
 
         # no resize necessary, but return a copy
         if im.width == width and im.height == height:
             return im.copy()
 
         # otherwise resize the original image so that it fits inside the bounding box
-        resized_image = self.image.resize((rw,rh),resample=Image.Resampling.LANCZOS)
+        resized_image = self.image.resize((rw, rh), resample=Image.Resampling.LANCZOS)
         return resized_image
 
+
 def make_grid(image_list, rows=None, cols=None):
     image_cnt = len(image_list)
     if None in (rows, cols):
@@ -57,7 +59,7 @@ def make_grid(image_list, rows=None, cols=None):
     width = image_list[0].width
     height = image_list[0].height
 
-    grid_img = Image.new('RGB', (width * cols, height * rows))
+    grid_img = Image.new("RGB", (width * cols, height * rows))
     i = 0
     for r in range(0, rows):
         for c in range(0, cols):
@@ -67,4 +69,3 @@ def make_grid(image_list, rows=None, cols=None):
             i = i + 1
 
     return grid_img
-

invokeai/backend/model_management/__init__.py

@@ -0,0 +1,9 @@
+"""
+Initialization file for invokeai.backend.model_management
+"""
+from .convert_ckpt_to_diffusers import (
+    convert_ckpt_to_diffusers,
+    load_pipeline_from_original_stable_diffusion_ckpt,
+)
+from .model_manager import ModelManager
+

invokeai/backend/prompting/__init__.py

@@ -0,0 +1,9 @@
+"""
+Initialization file for invokeai.backend.prompting
+"""
+from .conditioning import (
+    get_prompt_structure,
+    get_tokens_for_prompt_object,
+    get_uc_and_c_and_ec,
+    split_weighted_subprompts,
+)

invokeai/backend/prompting/conditioning.py

@@ -1,110 +1,138 @@
-'''
+"""
 This module handles the generation of the conditioning tensors.
 
 Useful function exports:
 
 get_uc_and_c_and_ec()           get the conditioned and unconditioned latent, and edited conditioning if we're doing cross-attention control
 
-'''
+"""
 import re
-from typing import Union, Optional, Any
-
-from transformers import CLIPTokenizer, CLIPTextModel
+from typing import Optional, Union
 
 from compel import Compel
-from compel.prompt_parser import FlattenedPrompt, Blend, Fragment, CrossAttentionControlSubstitute, PromptParser
-from .devices import torch_dtype
-from ..models.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent
-from ldm.invoke.globals import Globals
+from compel.prompt_parser import (
+    Blend,
+    CrossAttentionControlSubstitute,
+    FlattenedPrompt,
+    Fragment,
+    PromptParser,
+)
 
-def get_tokenizer(model) -> CLIPTokenizer:
-    # TODO remove legacy ckpt fallback handling
-    return (getattr(model, 'tokenizer', None) # diffusers
-            or model.cond_stage_model.tokenizer) # ldm
+from invokeai.backend.globals import Globals
 
-def get_text_encoder(model) -> Any:
-    # TODO remove legacy ckpt fallback handling
-    return (getattr(model, 'text_encoder', None)  # diffusers
-            or UnsqueezingLDMTransformer(model.cond_stage_model.transformer)) # ldm
-
-class UnsqueezingLDMTransformer:
-    def __init__(self, ldm_transformer):
-        self.ldm_transformer = ldm_transformer
-
-    @property
-    def device(self):
-        return self.ldm_transformer.device
-
-    def __call__(self, *args, **kwargs):
-        insufficiently_unsqueezed_tensor = self.ldm_transformer(*args, **kwargs)
-        return insufficiently_unsqueezed_tensor.unsqueeze(0)
+from ..stable_diffusion import InvokeAIDiffuserComponent
+from ..util import torch_dtype
 
 
-def get_uc_and_c_and_ec(prompt_string, model, log_tokens=False, skip_normalize_legacy_blend=False):
+def get_uc_and_c_and_ec(
+    prompt_string, model, log_tokens=False, skip_normalize_legacy_blend=False
+):
     # lazy-load any deferred textual inversions.
     # this might take a couple of seconds the first time a textual inversion is used.
-    model.textual_inversion_manager.create_deferred_token_ids_for_any_trigger_terms(prompt_string)
+    model.textual_inversion_manager.create_deferred_token_ids_for_any_trigger_terms(
+        prompt_string
+    )
 
-    tokenizer = get_tokenizer(model)
-    text_encoder = get_text_encoder(model)
-    compel = Compel(tokenizer=tokenizer,
-                    text_encoder=text_encoder,
-                    textual_inversion_manager=model.textual_inversion_manager,
-                    dtype_for_device_getter=torch_dtype)
+    tokenizer = model.tokenizer
+    compel = Compel(
+        tokenizer=tokenizer,
+        text_encoder=model.text_encoder,
+        textual_inversion_manager=model.textual_inversion_manager,
+        dtype_for_device_getter=torch_dtype,
+        truncate_long_prompts=False
+    )
 
-    positive_prompt_string, negative_prompt_string = split_prompt_to_positive_and_negative(prompt_string)
-    legacy_blend = try_parse_legacy_blend(positive_prompt_string, skip_normalize_legacy_blend)
-    positive_prompt: FlattenedPrompt|Blend
+    # get rid of any newline characters
+    prompt_string = prompt_string.replace("\n", " ")
+    (
+        positive_prompt_string,
+        negative_prompt_string,
+    ) = split_prompt_to_positive_and_negative(prompt_string)
+    legacy_blend = try_parse_legacy_blend(
+        positive_prompt_string, skip_normalize_legacy_blend
+    )
+    positive_prompt: Union[FlattenedPrompt, Blend]
     if legacy_blend is not None:
         positive_prompt = legacy_blend
     else:
         positive_prompt = Compel.parse_prompt_string(positive_prompt_string)
-    negative_prompt: FlattenedPrompt|Blend = Compel.parse_prompt_string(negative_prompt_string)
+    negative_prompt: Union[FlattenedPrompt, Blend] = Compel.parse_prompt_string(
+        negative_prompt_string
+    )
 
     if log_tokens or getattr(Globals, "log_tokenization", False):
         log_tokenization(positive_prompt, negative_prompt, tokenizer=tokenizer)
 
     c, options = compel.build_conditioning_tensor_for_prompt_object(positive_prompt)
     uc, _ = compel.build_conditioning_tensor_for_prompt_object(negative_prompt)
+    [c, uc] = compel.pad_conditioning_tensors_to_same_length([c, uc])
 
     tokens_count = get_max_token_count(tokenizer, positive_prompt)
 
-    ec = InvokeAIDiffuserComponent.ExtraConditioningInfo(tokens_count_including_eos_bos=tokens_count,
-                                                         cross_attention_control_args=options.get(
-                                                             'cross_attention_control', None))
+    ec = InvokeAIDiffuserComponent.ExtraConditioningInfo(
+        tokens_count_including_eos_bos=tokens_count,
+        cross_attention_control_args=options.get("cross_attention_control", None),
+    )
     return uc, c, ec
 
 
-def get_prompt_structure(prompt_string, skip_normalize_legacy_blend: bool = False) -> (
-    Union[FlattenedPrompt, Blend], FlattenedPrompt):
-    positive_prompt_string, negative_prompt_string = split_prompt_to_positive_and_negative(prompt_string)
-    legacy_blend = try_parse_legacy_blend(positive_prompt_string, skip_normalize_legacy_blend)
-    positive_prompt: FlattenedPrompt|Blend
+def get_prompt_structure(
+    prompt_string, skip_normalize_legacy_blend: bool = False
+) -> (Union[FlattenedPrompt, Blend], FlattenedPrompt):
+    (
+        positive_prompt_string,
+        negative_prompt_string,
+    ) = split_prompt_to_positive_and_negative(prompt_string)
+    legacy_blend = try_parse_legacy_blend(
+        positive_prompt_string, skip_normalize_legacy_blend
+    )
+    positive_prompt: Union[FlattenedPrompt, Blend]
     if legacy_blend is not None:
         positive_prompt = legacy_blend
     else:
         positive_prompt = Compel.parse_prompt_string(positive_prompt_string)
-    negative_prompt: FlattenedPrompt|Blend = Compel.parse_prompt_string(negative_prompt_string)
+    negative_prompt: Union[FlattenedPrompt, Blend] = Compel.parse_prompt_string(
+        negative_prompt_string
+    )
 
     return positive_prompt, negative_prompt
 
-def get_max_token_count(tokenizer, prompt: Union[FlattenedPrompt, Blend], truncate_if_too_long=True) -> int:
+
+def get_max_token_count(
+    tokenizer, prompt: Union[FlattenedPrompt, Blend], truncate_if_too_long=False
+) -> int:
     if type(prompt) is Blend:
         blend: Blend = prompt
-        return max([get_max_token_count(tokenizer, c, truncate_if_too_long) for c in blend.prompts])
+        return max(
+            [
+                get_max_token_count(tokenizer, c, truncate_if_too_long)
+                for c in blend.prompts
+            ]
+        )
     else:
-        return len(get_tokens_for_prompt_object(tokenizer, prompt, truncate_if_too_long))
+        return len(
+            get_tokens_for_prompt_object(tokenizer, prompt, truncate_if_too_long)
+        )
 
 
-def get_tokens_for_prompt_object(tokenizer, parsed_prompt: FlattenedPrompt, truncate_if_too_long=True) -> [str]:
-
+def get_tokens_for_prompt_object(
+    tokenizer, parsed_prompt: FlattenedPrompt, truncate_if_too_long=True
+) -> [str]:
     if type(parsed_prompt) is Blend:
-        raise ValueError("Blend is not supported here - you need to get tokens for each of its .children")
+        raise ValueError(
+            "Blend is not supported here - you need to get tokens for each of its .children"
+        )
 
-    text_fragments = [x.text if type(x) is Fragment else
-                      (" ".join([f.text for f in x.original]) if type(x) is CrossAttentionControlSubstitute else
-                       str(x))
-                      for x in parsed_prompt.children]
+    text_fragments = [
+        x.text
+        if type(x) is Fragment
+        else (
+            " ".join([f.text for f in x.original])
+            if type(x) is CrossAttentionControlSubstitute
+            else str(x)
+        )
+        for x in parsed_prompt.children
+    ]
     text = " ".join(text_fragments)
     tokens = tokenizer.tokenize(text)
     if truncate_if_too_long:
@@ -113,40 +141,48 @@ def get_tokens_for_prompt_object(tokenizer, parsed_prompt: FlattenedPrompt, trun
     return tokens
 
 
-def split_prompt_to_positive_and_negative(prompt_string_uncleaned):
-    unconditioned_words = ''
-    unconditional_regex = r'\[(.*?)\]'
+def split_prompt_to_positive_and_negative(prompt_string_uncleaned: str):
+    unconditioned_words = ""
+    unconditional_regex = r"\[(.*?)\]"
     unconditionals = re.findall(unconditional_regex, prompt_string_uncleaned)
     if len(unconditionals) > 0:
-        unconditioned_words = ' '.join(unconditionals)
+        unconditioned_words = " ".join(unconditionals)
 
         # Remove Unconditioned Words From Prompt
         unconditional_regex_compile = re.compile(unconditional_regex)
-        clean_prompt = unconditional_regex_compile.sub(' ', prompt_string_uncleaned)
-        prompt_string_cleaned = re.sub(' +', ' ', clean_prompt)
+        clean_prompt = unconditional_regex_compile.sub(" ", prompt_string_uncleaned)
+        prompt_string_cleaned = re.sub(" +", " ", clean_prompt)
     else:
         prompt_string_cleaned = prompt_string_uncleaned
     return prompt_string_cleaned, unconditioned_words
 
 
-def log_tokenization(positive_prompt: Union[Blend, FlattenedPrompt],
-                     negative_prompt: Union[Blend, FlattenedPrompt],
-                     tokenizer):
+def log_tokenization(
+    positive_prompt: Union[Blend, FlattenedPrompt],
+    negative_prompt: Union[Blend, FlattenedPrompt],
+    tokenizer,
+):
     print(f"\n>> [TOKENLOG] Parsed Prompt: {positive_prompt}")
     print(f"\n>> [TOKENLOG] Parsed Negative Prompt: {negative_prompt}")
 
     log_tokenization_for_prompt_object(positive_prompt, tokenizer)
-    log_tokenization_for_prompt_object(negative_prompt, tokenizer, display_label_prefix="(negative prompt)")
+    log_tokenization_for_prompt_object(
+        negative_prompt, tokenizer, display_label_prefix="(negative prompt)"
+    )
 
 
-def log_tokenization_for_prompt_object(p: Union[Blend, FlattenedPrompt], tokenizer, display_label_prefix=None):
+def log_tokenization_for_prompt_object(
+    p: Union[Blend, FlattenedPrompt], tokenizer, display_label_prefix=None
+):
     display_label_prefix = display_label_prefix or ""
     if type(p) is Blend:
         blend: Blend = p
         for i, c in enumerate(blend.prompts):
             log_tokenization_for_prompt_object(
-                c, tokenizer,
-                display_label_prefix=f"{display_label_prefix}(blend part {i + 1}, weight={blend.weights[i]})")
+                c,
+                tokenizer,
+                display_label_prefix=f"{display_label_prefix}(blend part {i + 1}, weight={blend.weights[i]})",
+            )
     elif type(p) is FlattenedPrompt:
         flattened_prompt: FlattenedPrompt = p
         if flattened_prompt.wants_cross_attention_control:
@@ -161,18 +197,26 @@ def log_tokenization_for_prompt_object(p: Union[Blend, FlattenedPrompt], tokeniz
                     edited_fragments.append(f)
 
             original_text = " ".join([x.text for x in original_fragments])
-            log_tokenization_for_text(original_text, tokenizer,
-                                      display_label=f"{display_label_prefix}(.swap originals)")
+            log_tokenization_for_text(
+                original_text,
+                tokenizer,
+                display_label=f"{display_label_prefix}(.swap originals)",
+            )
             edited_text = " ".join([x.text for x in edited_fragments])
-            log_tokenization_for_text(edited_text, tokenizer,
-                                      display_label=f"{display_label_prefix}(.swap replacements)")
+            log_tokenization_for_text(
+                edited_text,
+                tokenizer,
+                display_label=f"{display_label_prefix}(.swap replacements)",
+            )
         else:
             text = " ".join([x.text for x in flattened_prompt.children])
-            log_tokenization_for_text(text, tokenizer, display_label=display_label_prefix)
+            log_tokenization_for_text(
+                text, tokenizer, display_label=display_label_prefix
+            )
 
 
-def log_tokenization_for_text(text, tokenizer, display_label=None):
-    """ shows how the prompt is tokenized
+def log_tokenization_for_text(text, tokenizer, display_label=None, truncate_if_too_long=False):
+    """shows how the prompt is tokenized
     # usually tokens have '</w>' to indicate end-of-word,
     # but for readability it has been replaced with ' '
     """
@@ -183,25 +227,25 @@ def log_tokenization_for_text(text, tokenizer, display_label=None):
     totalTokens = len(tokens)
 
     for i in range(0, totalTokens):
-        token = tokens[i].replace('</w>', ' ')
+        token = tokens[i].replace("</w>", " ")
         # alternate color
         s = (usedTokens % 6) + 1
-        if i < tokenizer.model_max_length:
+        if truncate_if_too_long and i >= tokenizer.model_max_length:
+            discarded = discarded + f"\x1b[0;3{s};40m{token}"
+        else:
             tokenized = tokenized + f"\x1b[0;3{s};40m{token}"
             usedTokens += 1
-        else:  # over max token length
-            discarded = discarded + f"\x1b[0;3{s};40m{token}"
 
     if usedTokens > 0:
         print(f'\n>> [TOKENLOG] Tokens {display_label or ""} ({usedTokens}):')
-        print(f'{tokenized}\x1b[0m')
+        print(f"{tokenized}\x1b[0m")
 
     if discarded != "":
-        print(f'\n>> [TOKENLOG] Tokens Discarded ({totalTokens - usedTokens}):')
-        print(f'{discarded}\x1b[0m')
+        print(f"\n>> [TOKENLOG] Tokens Discarded ({totalTokens - usedTokens}):")
+        print(f"{discarded}\x1b[0m")
 
 
-def try_parse_legacy_blend(text: str, skip_normalize: bool=False) -> Optional[Blend]:
+def try_parse_legacy_blend(text: str, skip_normalize: bool = False) -> Optional[Blend]:
     weighted_subprompts = split_weighted_subprompts(text, skip_normalize=skip_normalize)
     if len(weighted_subprompts) <= 1:
         return None
@@ -212,10 +256,12 @@ def try_parse_legacy_blend(text: str, skip_normalize: bool=False) -> Optional[Bl
     parsed_conjunctions = [pp.parse_conjunction(x) for x in strings]
     flattened_prompts = [x.prompts[0] for x in parsed_conjunctions]
 
-    return Blend(prompts=flattened_prompts, weights=weights, normalize_weights=not skip_normalize)
+    return Blend(
+        prompts=flattened_prompts, weights=weights, normalize_weights=not skip_normalize
+    )
 
 
-def split_weighted_subprompts(text, skip_normalize=False)->list:
+def split_weighted_subprompts(text, skip_normalize=False) -> list:
     """
     Legacy blend parsing.
 
@@ -224,7 +270,8 @@ def split_weighted_subprompts(text, skip_normalize=False)->list:
     if ':' has no value defined, defaults to 1.0
     repeats until no text remaining
     """
-    prompt_parser = re.compile("""
+    prompt_parser = re.compile(
+        """
             (?P<prompt>     # capture group for 'prompt'
             (?:\\\:|[^:])+  # match one or more non ':' characters or escaped colons '\:'
             )               # end 'prompt'
@@ -237,16 +284,20 @@ def split_weighted_subprompts(text, skip_normalize=False)->list:
             |               # OR
             $               # else, if no ':' then match end of line
             )               # end non-capture group
-            """, re.VERBOSE)
-    parsed_prompts = [(match.group("prompt").replace("\\:", ":"), float(
-        match.group("weight") or 1)) for match in re.finditer(prompt_parser, text)]
+            """,
+        re.VERBOSE,
+    )
+    parsed_prompts = [
+        (match.group("prompt").replace("\\:", ":"), float(match.group("weight") or 1))
+        for match in re.finditer(prompt_parser, text)
+    ]
     if skip_normalize:
         return parsed_prompts
     weight_sum = sum(map(lambda x: x[1], parsed_prompts))
     if weight_sum == 0:
         print(
-            "* Warning: Subprompt weights add up to zero. Discarding and using even weights instead.")
+            "* Warning: Subprompt weights add up to zero. Discarding and using even weights instead."
+        )
         equal_weight = 1 / max(len(parsed_prompts), 1)
         return [(x[0], equal_weight) for x in parsed_prompts]
     return [(x[0], x[1] / weight_sum) for x in parsed_prompts]
-

invokeai/backend/restoration/__init__.py

@@ -0,0 +1,4 @@
+"""
+Initialization file for the invokeai.backend.restoration package
+"""
+from .base import Restoration

invokeai/backend/restoration/base.py

@@ -1,38 +1,43 @@
-class Restoration():
+class Restoration:
     def __init__(self) -> None:
         pass
 
-    def load_face_restore_models(self, gfpgan_model_path='./models/gfpgan/GFPGANv1.4.pth'):
+    def load_face_restore_models(
+        self, gfpgan_model_path="./models/gfpgan/GFPGANv1.4.pth"
+    ):
         # Load GFPGAN
         gfpgan = self.load_gfpgan(gfpgan_model_path)
         if gfpgan.gfpgan_model_exists:
-            print('>> GFPGAN Initialized')
+            print(">> GFPGAN Initialized")
         else:
-            print('>> GFPGAN Disabled')
+            print(">> GFPGAN Disabled")
             gfpgan = None
 
         # Load CodeFormer
         codeformer = self.load_codeformer()
         if codeformer.codeformer_model_exists:
-            print('>> CodeFormer Initialized')
+            print(">> CodeFormer Initialized")
         else:
-            print('>> CodeFormer Disabled')
+            print(">> CodeFormer Disabled")
             codeformer = None
 
         return gfpgan, codeformer
 
     # Face Restore Models
     def load_gfpgan(self, gfpgan_model_path):
-        from ldm.invoke.restoration.gfpgan import GFPGAN
+        from .gfpgan import GFPGAN
+
         return GFPGAN(gfpgan_model_path)
 
     def load_codeformer(self):
-        from ldm.invoke.restoration.codeformer import CodeFormerRestoration
+        from .codeformer import CodeFormerRestoration
+
         return CodeFormerRestoration()
 
     # Upscale Models
     def load_esrgan(self, esrgan_bg_tile=400):
-        from ldm.invoke.restoration.realesrgan import ESRGAN
+        from .realesrgan import ESRGAN
+
         esrgan = ESRGAN(esrgan_bg_tile)
-        print('>> ESRGAN Initialized')
-        return esrgan;
+        print(">> ESRGAN Initialized")
+        return esrgan

invokeai/backend/restoration/codeformer.py

@@ -1,17 +1,21 @@
 import os
-import torch
-import numpy as np
-import warnings
 import sys
-from ldm.invoke.globals import Globals
+import warnings
 
-pretrained_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
+import numpy as np
+import torch
 
-class CodeFormerRestoration():
-    def __init__(self,
-            codeformer_dir='models/codeformer',
-            codeformer_model_path='codeformer.pth') -> None:
+from ..globals import Globals
 
+pretrained_model_url = (
+    "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth"
+)
+
+
+class CodeFormerRestoration:
+    def __init__(
+        self, codeformer_dir="models/codeformer", codeformer_model_path="codeformer.pth"
+    ) -> None:
         if not os.path.isabs(codeformer_dir):
             codeformer_dir = os.path.join(Globals.root, codeformer_dir)
 
@@ -19,22 +23,23 @@ class CodeFormerRestoration():
         self.codeformer_model_exists = os.path.isfile(self.model_path)
 
         if not self.codeformer_model_exists:
-            print('## NOT FOUND: CodeFormer model not found at ' + self.model_path)
+            print("## NOT FOUND: CodeFormer model not found at " + self.model_path)
         sys.path.append(os.path.abspath(codeformer_dir))
 
     def process(self, image, strength, device, seed=None, fidelity=0.75):
         if seed is not None:
-            print(f'>> CodeFormer - Restoring Faces for image seed:{seed}')
+            print(f">> CodeFormer - Restoring Faces for image seed:{seed}")
         with warnings.catch_warnings():
-            warnings.filterwarnings('ignore', category=DeprecationWarning)
-            warnings.filterwarnings('ignore', category=UserWarning)
+            warnings.filterwarnings("ignore", category=DeprecationWarning)
+            warnings.filterwarnings("ignore", category=UserWarning)
 
-            from basicsr.utils.download_util import load_file_from_url
             from basicsr.utils import img2tensor, tensor2img
+            from basicsr.utils.download_util import load_file_from_url
             from facexlib.utils.face_restoration_helper import FaceRestoreHelper
-            from ldm.invoke.restoration.codeformer_arch import CodeFormer
-            from torchvision.transforms.functional import normalize
             from PIL import Image
+            from torchvision.transforms.functional import normalize
+
+            from .codeformer_arch import CodeFormer
 
             cf_class = CodeFormer
 
@@ -43,28 +48,31 @@ class CodeFormerRestoration():
                 codebook_size=1024,
                 n_head=8,
                 n_layers=9,
-                connect_list=['32', '64', '128', '256']
+                connect_list=["32", "64", "128", "256"],
             ).to(device)
 
             # note that this file should already be downloaded and cached at
             # this point
-            checkpoint_path = load_file_from_url(url=pretrained_model_url,
-                                                 model_dir=os.path.abspath(os.path.dirname(self.model_path)),
-                                                 progress=True
+            checkpoint_path = load_file_from_url(
+                url=pretrained_model_url,
+                model_dir=os.path.abspath(os.path.dirname(self.model_path)),
+                progress=True,
             )
-            checkpoint = torch.load(checkpoint_path)['params_ema']
+            checkpoint = torch.load(checkpoint_path)["params_ema"]
             cf.load_state_dict(checkpoint)
             cf.eval()
 
-            image = image.convert('RGB')
+            image = image.convert("RGB")
             # Codeformer expects a BGR np array; make array and flip channels
-            bgr_image_array = np.array(image, dtype=np.uint8)[...,::-1]
+            bgr_image_array = np.array(image, dtype=np.uint8)[..., ::-1]
 
             face_helper = FaceRestoreHelper(
                 upscale_factor=1,
                 use_parse=True,
                 device=device,
-                model_rootpath=os.path.join(Globals.root,'models','gfpgan','weights'),
+                model_rootpath=os.path.join(
+                    Globals.root, "models", "gfpgan", "weights"
+                ),
             )
             face_helper.clean_all()
             face_helper.read_image(bgr_image_array)
@@ -72,30 +80,35 @@ class CodeFormerRestoration():
             face_helper.align_warp_face()
 
             for idx, cropped_face in enumerate(face_helper.cropped_faces):
-                cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
-                normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+                cropped_face_t = img2tensor(
+                    cropped_face / 255.0, bgr2rgb=True, float32=True
+                )
+                normalize(
+                    cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True
+                )
                 cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
 
                 try:
                     with torch.no_grad():
                         output = cf(cropped_face_t, w=fidelity, adain=True)[0]
-                        restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
+                        restored_face = tensor2img(
+                            output.squeeze(0), rgb2bgr=True, min_max=(-1, 1)
+                        )
                     del output
                     torch.cuda.empty_cache()
                 except RuntimeError as error:
-                    print(f'\tFailed inference for CodeFormer: {error}.')
+                    print(f"\tFailed inference for CodeFormer: {error}.")
                     restored_face = cropped_face
 
-                restored_face = restored_face.astype('uint8')
+                restored_face = restored_face.astype("uint8")
                 face_helper.add_restored_face(restored_face)
 
-
             face_helper.get_inverse_affine(None)
 
             restored_img = face_helper.paste_faces_to_input_image()
 
             # Flip the channels back to RGB
-            res = Image.fromarray(restored_img[...,::-1])
+            res = Image.fromarray(restored_img[..., ::-1])
 
             if strength < 1.0:
                 # Resize the image to the new image if the sizes have changed

invokeai/backend/restoration/codeformer_arch.py

@@ -1,13 +1,15 @@
 import math
+from typing import List, Optional
+
 import numpy as np
 import torch
-from torch import nn, Tensor
 import torch.nn.functional as F
-from typing import Optional, List
-
-from ldm.invoke.restoration.vqgan_arch import *
 from basicsr.utils import get_root_logger
 from basicsr.utils.registry import ARCH_REGISTRY
+from torch import Tensor, nn
+
+from .vqgan_arch import *
+
 
 def calc_mean_std(feat, eps=1e-5):
     """Calculate mean and std for adaptive_instance_normalization.
@@ -18,7 +20,7 @@ def calc_mean_std(feat, eps=1e-5):
             divide-by-zero. Default: 1e-5.
     """
     size = feat.size()
-    assert len(size) == 4, 'The input feature should be 4D tensor.'
+    assert len(size) == 4, "The input feature should be 4D tensor."
     b, c = size[:2]
     feat_var = feat.view(b, c, -1).var(dim=2) + eps
     feat_std = feat_var.sqrt().view(b, c, 1, 1)
@@ -39,7 +41,9 @@ def adaptive_instance_normalization(content_feat, style_feat):
     size = content_feat.size()
     style_mean, style_std = calc_mean_std(style_feat)
     content_mean, content_std = calc_mean_std(content_feat)
-    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
+    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(
+        size
+    )
     return normalized_feat * style_std.expand(size) + style_mean.expand(size)
 
 
@@ -49,7 +53,9 @@ class PositionEmbeddingSine(nn.Module):
     used by the Attention is all you need paper, generalized to work on images.
     """
 
-    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+    def __init__(
+        self, num_pos_feats=64, temperature=10000, normalize=False, scale=None
+    ):
         super().__init__()
         self.num_pos_feats = num_pos_feats
         self.temperature = temperature
@@ -62,7 +68,9 @@ class PositionEmbeddingSine(nn.Module):
 
     def forward(self, x, mask=None):
         if mask is None:
-            mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool)
+            mask = torch.zeros(
+                (x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool
+            )
         not_mask = ~mask
         y_embed = not_mask.cumsum(1, dtype=torch.float32)
         x_embed = not_mask.cumsum(2, dtype=torch.float32)
@@ -85,6 +93,7 @@ class PositionEmbeddingSine(nn.Module):
         pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
         return pos
 
+
 def _get_activation_fn(activation):
     """Return an activation function given a string"""
     if activation == "relu":
@@ -93,11 +102,13 @@ def _get_activation_fn(activation):
         return F.gelu
     if activation == "glu":
         return F.glu
-    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
 
 
 class TransformerSALayer(nn.Module):
-    def __init__(self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu"):
+    def __init__(
+        self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu"
+    ):
         super().__init__()
         self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)
         # Implementation of Feedforward model - MLP
@@ -115,16 +126,19 @@ class TransformerSALayer(nn.Module):
     def with_pos_embed(self, tensor, pos: Optional[Tensor]):
         return tensor if pos is None else tensor + pos
 
-    def forward(self, tgt,
-                tgt_mask: Optional[Tensor] = None,
-                tgt_key_padding_mask: Optional[Tensor] = None,
-                query_pos: Optional[Tensor] = None):
-
+    def forward(
+        self,
+        tgt,
+        tgt_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+    ):
         # self attention
         tgt2 = self.norm1(tgt)
         q = k = self.with_pos_embed(tgt2, query_pos)
-        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
-                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt2 = self.self_attn(
+            q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )[0]
         tgt = tgt + self.dropout1(tgt2)
 
         # ffn
@@ -133,20 +147,23 @@ class TransformerSALayer(nn.Module):
         tgt = tgt + self.dropout2(tgt2)
         return tgt
 
+
 class Fuse_sft_block(nn.Module):
     def __init__(self, in_ch, out_ch):
         super().__init__()
-        self.encode_enc = ResBlock(2*in_ch, out_ch)
+        self.encode_enc = ResBlock(2 * in_ch, out_ch)
 
         self.scale = nn.Sequential(
-                    nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
-                    nn.LeakyReLU(0.2, True),
-                    nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
+            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
+        )
 
         self.shift = nn.Sequential(
-                    nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
-                    nn.LeakyReLU(0.2, True),
-                    nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1))
+            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
+        )
 
     def forward(self, enc_feat, dec_feat, w=1):
         enc_feat = self.encode_enc(torch.cat([enc_feat, dec_feat], dim=1))
@@ -159,11 +176,19 @@ class Fuse_sft_block(nn.Module):
 
 @ARCH_REGISTRY.register()
 class CodeFormer(VQAutoEncoder):
-    def __init__(self, dim_embd=512, n_head=8, n_layers=9,
-                codebook_size=1024, latent_size=256,
-                connect_list=['32', '64', '128', '256'],
-                fix_modules=['quantize','generator']):
-        super(CodeFormer, self).__init__(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',2, [16], codebook_size)
+    def __init__(
+        self,
+        dim_embd=512,
+        n_head=8,
+        n_layers=9,
+        codebook_size=1024,
+        latent_size=256,
+        connect_list=["32", "64", "128", "256"],
+        fix_modules=["quantize", "generator"],
+    ):
+        super(CodeFormer, self).__init__(
+            512, 64, [1, 2, 2, 4, 4, 8], "nearest", 2, [16], codebook_size
+        )
 
         if fix_modules is not None:
             for module in fix_modules:
@@ -173,33 +198,53 @@ class CodeFormer(VQAutoEncoder):
         self.connect_list = connect_list
         self.n_layers = n_layers
         self.dim_embd = dim_embd
-        self.dim_mlp = dim_embd*2
+        self.dim_mlp = dim_embd * 2
 
         self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd))
         self.feat_emb = nn.Linear(256, self.dim_embd)
 
         # transformer
-        self.ft_layers = nn.Sequential(*[TransformerSALayer(embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0)
-                                    for _ in range(self.n_layers)])
+        self.ft_layers = nn.Sequential(
+            *[
+                TransformerSALayer(
+                    embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0
+                )
+                for _ in range(self.n_layers)
+            ]
+        )
 
         # logits_predict head
         self.idx_pred_layer = nn.Sequential(
-            nn.LayerNorm(dim_embd),
-            nn.Linear(dim_embd, codebook_size, bias=False))
+            nn.LayerNorm(dim_embd), nn.Linear(dim_embd, codebook_size, bias=False)
+        )
 
         self.channels = {
-            '16': 512,
-            '32': 256,
-            '64': 256,
-            '128': 128,
-            '256': 128,
-            '512': 64,
+            "16": 512,
+            "32": 256,
+            "64": 256,
+            "128": 128,
+            "256": 128,
+            "512": 64,
         }
 
         # after second residual block for > 16, before attn layer for ==16
-        self.fuse_encoder_block = {'512':2, '256':5, '128':8, '64':11, '32':14, '16':18}
+        self.fuse_encoder_block = {
+            "512": 2,
+            "256": 5,
+            "128": 8,
+            "64": 11,
+            "32": 14,
+            "16": 18,
+        }
         # after first residual block for > 16, before attn layer for ==16
-        self.fuse_generator_block = {'16':6, '32': 9, '64':12, '128':15, '256':18, '512':21}
+        self.fuse_generator_block = {
+            "16": 6,
+            "32": 9,
+            "64": 12,
+            "128": 15,
+            "256": 18,
+            "512": 21,
+        }
 
         # fuse_convs_dict
         self.fuse_convs_dict = nn.ModuleDict()
@@ -228,20 +273,20 @@ class CodeFormer(VQAutoEncoder):
         lq_feat = x
         # ################# Transformer ###################
         # quant_feat, codebook_loss, quant_stats = self.quantize(lq_feat)
-        pos_emb = self.position_emb.unsqueeze(1).repeat(1,x.shape[0],1)
+        pos_emb = self.position_emb.unsqueeze(1).repeat(1, x.shape[0], 1)
         # BCHW -> BC(HW) -> (HW)BC
-        feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2,0,1))
+        feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2, 0, 1))
         query_emb = feat_emb
         # Transformer encoder
         for layer in self.ft_layers:
             query_emb = layer(query_emb, query_pos=pos_emb)
 
         # output logits
-        logits = self.idx_pred_layer(query_emb) # (hw)bn
-        logits = logits.permute(1,0,2) # (hw)bn -> b(hw)n
+        logits = self.idx_pred_layer(query_emb)  # (hw)bn
+        logits = logits.permute(1, 0, 2)  # (hw)bn -> b(hw)n
 
-        if code_only: # for training stage II
-          # logits doesn't need softmax before cross_entropy loss
+        if code_only:  # for training stage II
+            # logits doesn't need softmax before cross_entropy loss
             return logits, lq_feat
 
         # ################# Quantization ###################
@@ -252,12 +297,14 @@ class CodeFormer(VQAutoEncoder):
         # ------------
         soft_one_hot = F.softmax(logits, dim=2)
         _, top_idx = torch.topk(soft_one_hot, 1, dim=2)
-        quant_feat = self.quantize.get_codebook_feat(top_idx, shape=[x.shape[0],16,16,256])
+        quant_feat = self.quantize.get_codebook_feat(
+            top_idx, shape=[x.shape[0], 16, 16, 256]
+        )
         # preserve gradients
         # quant_feat = lq_feat + (quant_feat - lq_feat).detach()
 
         if detach_16:
-            quant_feat = quant_feat.detach() # for training stage III
+            quant_feat = quant_feat.detach()  # for training stage III
         if adain:
             quant_feat = adaptive_instance_normalization(quant_feat, lq_feat)
 
@@ -267,10 +314,12 @@ class CodeFormer(VQAutoEncoder):
 
         for i, block in enumerate(self.generator.blocks):
             x = block(x)
-            if i in fuse_list: # fuse after i-th block
+            if i in fuse_list:  # fuse after i-th block
                 f_size = str(x.shape[-1])
-                if w>0:
-                    x = self.fuse_convs_dict[f_size](enc_feat_dict[f_size].detach(), x, w)
+                if w > 0:
+                    x = self.fuse_convs_dict[f_size](
+                        enc_feat_dict[f_size].detach(), x, w
+                    )
         out = x
         # logits doesn't need softmax before cross_entropy loss
         return out, logits, lq_feat