chore: import upstream snapshot with attribution

invokeai/backend/util/__init__.py

@@ -0,0 +1,12 @@
+"""
+Initialization file for invokeai.backend.util
+"""
+
+from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.util.util import Chdir, directory_size
+
+__all__ = [
+    "directory_size",
+    "Chdir",
+    "InvokeAILogger",
+]

invokeai/backend/util/attention.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2023 Lincoln Stein and the InvokeAI Team
+"""
+Utility routine used for autodetection of optimal slice size
+for attention mechanism.
+"""
+
+import psutil
+import torch
+
+
+def auto_detect_slice_size(latents: torch.Tensor) -> str:
+    bytes_per_element_needed_for_baddbmm_duplication = latents.element_size() + 4
+    max_size_required_for_baddbmm = (
+        16
+        * latents.size(dim=2)
+        * latents.size(dim=3)
+        * latents.size(dim=2)
+        * latents.size(dim=3)
+        * bytes_per_element_needed_for_baddbmm_duplication
+    )
+    if latents.device.type in {"cpu", "mps"}:
+        mem_free = psutil.virtual_memory().free
+    elif latents.device.type == "cuda":
+        mem_free, _ = torch.cuda.mem_get_info(latents.device)
+    else:
+        raise ValueError(f"unrecognized device {latents.device}")
+
+    if max_size_required_for_baddbmm > (mem_free * 3.0 / 4.0):
+        return "max"
+    elif torch.backends.mps.is_available():
+        return "max"
+    else:
+        return "balanced"

invokeai/backend/util/build_line.py

@@ -0,0 +1,6 @@
+from typing import Callable
+
+
+def build_line(x1: float, y1: float, x2: float, y2: float) -> Callable[[float], float]:
+    """Build a linear function given two points on the line (x1, y1) and (x2, y2)."""
+    return lambda x: (y2 - y1) / (x2 - x1) * (x - x1) + y1

invokeai/backend/util/calc_tensor_size.py

@@ -0,0 +1,11 @@
+import torch
+
+
+def calc_tensor_size(t: torch.Tensor) -> int:
+    """Calculate the size of a tensor in bytes."""
+    return t.nelement() * t.element_size()
+
+
+def calc_tensors_size(tensors: list[torch.Tensor | None]) -> int:
+    """Calculate the size of a list of tensors in bytes."""
+    return sum(calc_tensor_size(t) for t in tensors if t is not None)

invokeai/backend/util/catch_sigint.py

@@ -0,0 +1,29 @@
+"""
+This module defines a context manager `catch_sigint()` which temporarily replaces
+the sigINT handler defined by the ASGI in order to allow the user to ^C the application
+and shut it down immediately. This was implemented in order to allow the user to interrupt
+slow model hashing during startup.
+
+Use like this:
+
+  from invokeai.backend.util.catch_sigint import catch_sigint
+  with catch_sigint():
+      run_some_hard_to_interrupt_process()
+"""
+
+import signal
+from contextlib import contextmanager
+from typing import Generator
+
+
+def sigint_handler(signum, frame):  # type: ignore
+    signal.signal(signal.SIGINT, signal.SIG_DFL)
+    signal.raise_signal(signal.SIGINT)
+
+
+@contextmanager
+def catch_sigint() -> Generator[None, None, None]:
+    original_handler = signal.getsignal(signal.SIGINT)
+    signal.signal(signal.SIGINT, sigint_handler)
+    yield
+    signal.signal(signal.SIGINT, original_handler)

invokeai/backend/util/devices.py

@@ -0,0 +1,153 @@
+from typing import Dict, Literal, Optional, Union
+
+import torch
+from deprecated import deprecated
+
+from invokeai.app.services.config.config_default import get_config
+
+# legacy APIs
+TorchPrecisionNames = Literal["float32", "float16", "bfloat16"]
+CPU_DEVICE = torch.device("cpu")
+CUDA_DEVICE = torch.device("cuda")
+MPS_DEVICE = torch.device("mps")
+
+
+@deprecated("Use TorchDevice.choose_torch_dtype() instead.")  # type: ignore
+def choose_precision(device: torch.device) -> TorchPrecisionNames:
+    """Return the string representation of the recommended torch device."""
+    torch_dtype = TorchDevice.choose_torch_dtype(device)
+    return PRECISION_TO_NAME[torch_dtype]
+
+
+@deprecated("Use TorchDevice.choose_torch_device() instead.")  # type: ignore
+def choose_torch_device() -> torch.device:
+    """Return the torch.device to use for accelerated inference."""
+    return TorchDevice.choose_torch_device()
+
+
+@deprecated("Use TorchDevice.choose_torch_dtype() instead.")  # type: ignore
+def torch_dtype(device: torch.device) -> torch.dtype:
+    """Return the torch precision for the recommended torch device."""
+    return TorchDevice.choose_torch_dtype(device)
+
+
+NAME_TO_PRECISION: Dict[TorchPrecisionNames, torch.dtype] = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+}
+PRECISION_TO_NAME: Dict[torch.dtype, TorchPrecisionNames] = {v: k for k, v in NAME_TO_PRECISION.items()}
+
+
+class TorchDevice:
+    """Abstraction layer for torch devices."""
+
+    CPU_DEVICE = torch.device("cpu")
+    CUDA_DEVICE = torch.device("cuda")
+    MPS_DEVICE = torch.device("mps")
+
+    @classmethod
+    def choose_torch_device(cls) -> torch.device:
+        """Return the torch.device to use for accelerated inference."""
+        app_config = get_config()
+        if app_config.device != "auto":
+            device = torch.device(app_config.device)
+        elif torch.cuda.is_available():
+            device = CUDA_DEVICE
+        elif torch.backends.mps.is_available():
+            device = MPS_DEVICE
+        else:
+            device = CPU_DEVICE
+        return cls.normalize(device)
+
+    @classmethod
+    def choose_torch_dtype(cls, device: Optional[torch.device] = None) -> torch.dtype:
+        """Return the precision to use for accelerated inference."""
+        device = device or cls.choose_torch_device()
+        config = get_config()
+        if device.type == "cuda" and torch.cuda.is_available():
+            device_name = torch.cuda.get_device_name(device)
+            if "GeForce GTX 1660" in device_name or "GeForce GTX 1650" in device_name:
+                # These GPUs have limited support for float16
+                return cls._to_dtype("float32")
+            elif config.precision == "auto":
+                # Default to float16 for CUDA devices
+                return cls._to_dtype("float16")
+            else:
+                # Use the user-defined precision
+                return cls._to_dtype(config.precision)
+
+        elif device.type == "mps" and torch.backends.mps.is_available():
+            if config.precision == "auto":
+                # Default to float16 for MPS devices
+                return cls._to_dtype("float16")
+            else:
+                # Use the user-defined precision
+                return cls._to_dtype(config.precision)
+        # CPU / safe fallback
+        return cls._to_dtype("float32")
+
+    @classmethod
+    def get_torch_device_name(cls) -> str:
+        """Return the device name for the current torch device."""
+        device = cls.choose_torch_device()
+        return torch.cuda.get_device_name(device) if device.type == "cuda" else device.type.upper()
+
+    @classmethod
+    def normalize(cls, device: Union[str, torch.device]) -> torch.device:
+        """Add the device index to CUDA devices."""
+        device = torch.device(device)
+        if device.index is None and device.type == "cuda" and torch.cuda.is_available():
+            device = torch.device(device.type, torch.cuda.current_device())
+        return device
+
+    @classmethod
+    def empty_cache(cls) -> None:
+        """Clear the GPU device cache."""
+        if torch.backends.mps.is_available():
+            torch.mps.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    @classmethod
+    def _to_dtype(cls, precision_name: TorchPrecisionNames) -> torch.dtype:
+        return NAME_TO_PRECISION[precision_name]
+
+    @classmethod
+    def choose_bfloat16_safe_dtype(cls, device: Optional[torch.device] = None) -> torch.dtype:
+        """Return bfloat16 if supported on the device, else fallback to float16/float32.
+
+        This is useful for models that require bfloat16 precision (e.g., Z-Image, Flux)
+        but need to run on hardware that may not support bfloat16.
+
+        Args:
+            device: The target device. If None, uses choose_torch_device().
+
+        Returns:
+            torch.bfloat16 if supported, torch.float16 for CUDA without bfloat16 support,
+            or torch.float32 for CPU/MPS.
+        """
+        device = device or cls.choose_torch_device()
+        try:
+            # Test if bfloat16 is supported on this device
+            torch.tensor([1.0], dtype=torch.bfloat16, device=device)
+            return torch.bfloat16
+        except TypeError:
+            # bfloat16 not supported - fallback based on device type
+            if device.type == "cuda":
+                return torch.float16
+            return torch.float32
+
+    @classmethod
+    def choose_anima_inference_dtype(cls, device: Optional[torch.device] = None) -> torch.dtype:
+        """Choose the inference dtype for Anima models, honoring config.precision.
+
+        When precision is 'auto', delegates to choose_bfloat16_safe_dtype (current
+        behavior). When precision is set to a specific value (float16, bfloat16,
+        float32), returns that dtype directly without hardware probing.
+        """
+        device = device or cls.choose_torch_device()
+        config = get_config()
+        if config.precision == "auto":
+            return cls.choose_bfloat16_safe_dtype(device)
+        return NAME_TO_PRECISION[config.precision]

invokeai/backend/util/gallery_maintenance.py

@@ -0,0 +1,594 @@
+# pylint: disable=line-too-long
+# pylint: disable=broad-exception-caught
+# pylint: disable=missing-function-docstring
+"""Script to peform db maintenance and outputs directory management."""
+
+import argparse
+import datetime
+import enum
+import glob
+import locale
+import os
+import shutil
+import sqlite3
+from pathlib import Path
+
+import PIL
+import PIL.ImageOps
+import PIL.PngImagePlugin
+import yaml
+
+
+class ConfigMapper:
+    """Configuration loader."""
+
+    def __init__(self):  # noqa D107
+        pass
+
+    TIMESTAMP_STRING = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
+
+    INVOKE_DIRNAME = "invokeai"
+    YAML_FILENAME = "invokeai.yaml"
+    DATABASE_FILENAME = "invokeai.db"
+
+    DEFAULT_OUTDIR = "outputs"
+    DEFAULT_DB_DIR = "databases"
+
+    database_path = None
+    database_backup_dir = None
+    outputs_path = None
+    archive_path = None
+    thumbnails_path = None
+    thumbnails_archive_path = None
+
+    def load(self):
+        """Read paths from yaml config and validate."""
+        root = "."
+
+        if not self.__load_from_root_config(os.path.abspath(root)):
+            return False
+
+        return True
+
+    def __load_from_root_config(self, invoke_root):
+        """Validate a yaml path exists, confirm the user wants to use it and load config."""
+        yaml_path = os.path.join(invoke_root, self.YAML_FILENAME)
+        if not os.path.exists(yaml_path):
+            print(f"Unable to find invokeai.yaml at {yaml_path}!")
+            return False
+        if os.path.exists(yaml_path):
+            db_dir, outdir = self.__load_paths_from_yaml_file(yaml_path)
+
+            if db_dir is None:
+                db_dir = self.DEFAULT_DB_DIR
+                print(f"The invokeai.yaml file was found but is missing the db_dir setting! Defaulting to {db_dir}")
+            if outdir is None:
+                outdir = self.DEFAULT_OUTDIR
+                print(f"The invokeai.yaml file was found but is missing the outdir setting! Defaulting to {outdir}")
+
+            if os.path.isabs(db_dir):
+                self.database_path = os.path.join(db_dir, self.DATABASE_FILENAME)
+            else:
+                self.database_path = os.path.join(invoke_root, db_dir, self.DATABASE_FILENAME)
+
+            self.database_backup_dir = os.path.join(os.path.dirname(self.database_path), "backup")
+
+            if os.path.isabs(outdir):
+                self.outputs_path = os.path.join(outdir, "images")
+                self.archive_path = os.path.join(outdir, "images-archive")
+            else:
+                self.outputs_path = os.path.join(invoke_root, outdir, "images")
+                self.archive_path = os.path.join(invoke_root, outdir, "images-archive")
+
+            self.thumbnails_path = os.path.join(self.outputs_path, "thumbnails")
+            self.thumbnails_archive_path = os.path.join(self.archive_path, "thumbnails")
+
+            db_exists = os.path.exists(self.database_path)
+            outdir_exists = os.path.exists(self.outputs_path)
+
+            text = f"Found {self.YAML_FILENAME} file at {yaml_path}:"
+            text += f"\n  Database : {self.database_path} - {'Exists!' if db_exists else 'Not Found!'}"
+            text += f"\n  Outputs  : {self.outputs_path}- {'Exists!' if outdir_exists else 'Not Found!'}"
+            print(text)
+
+            if db_exists and outdir_exists:
+                return True
+            else:
+                print(
+                    "\nOne or more paths specified in invoke.yaml do not exist. Please inspect/correct the configuration and ensure the script is run in the developer console mode (option 8) from an Invoke AI root directory."
+                )
+                return False
+        else:
+            print(
+                f"Auto-discovery of configuration failed! Could not find ({yaml_path})!\n\nPlease ensure the script is run in the developer console mode (option 8) from an Invoke AI root directory."
+            )
+            return False
+
+    def __load_paths_from_yaml_file(self, yaml_path):
+        """Load an Invoke AI yaml file and get the database and outputs paths."""
+        try:
+            with open(yaml_path, "rt", encoding=locale.getpreferredencoding()) as file:
+                yamlinfo = yaml.safe_load(file)
+                db_dir = yamlinfo.get("InvokeAI", {}).get("Paths", {}).get("db_dir", None)
+                outdir = yamlinfo.get("InvokeAI", {}).get("Paths", {}).get("outdir", None)
+                return db_dir, outdir
+        except Exception:
+            print(f"Failed to load paths from yaml file! {yaml_path}!")
+            return None, None
+
+
+class MaintenanceStats:
+    """DTO for tracking work progress."""
+
+    def __init__(self):  # noqa D107
+        pass
+
+    time_start = datetime.datetime.utcnow()
+    count_orphaned_db_entries_cleaned = 0
+    count_orphaned_disk_files_cleaned = 0
+    count_orphaned_thumbnails_cleaned = 0
+    count_thumbnails_regenerated = 0
+    count_errors = 0
+
+    @staticmethod
+    def get_elapsed_time_string():
+        """Get a friendly time string for the  time elapsed since processing start."""
+        time_now = datetime.datetime.utcnow()
+        total_seconds = (time_now - MaintenanceStats.time_start).total_seconds()
+        hours = int((total_seconds) / 3600)
+        minutes = int(((total_seconds) % 3600) / 60)
+        seconds = total_seconds % 60
+        out_str = f"{hours} hour(s) -" if hours > 0 else ""
+        out_str += f"{minutes} minute(s) -" if minutes > 0 else ""
+        out_str += f"{seconds:.2f} second(s)"
+        return out_str
+
+
+class DatabaseMapper:
+    """Class to abstract database functionality."""
+
+    def __init__(self, database_path, database_backup_dir):  # noqa D107
+        self.database_path = database_path
+        self.database_backup_dir = database_backup_dir
+        self.connection = None
+        self.cursor = None
+
+    def backup(self, timestamp_string):
+        """Take a backup of the database."""
+        if not os.path.exists(self.database_backup_dir):
+            print(f"Database backup directory {self.database_backup_dir} does not exist -> creating...", end="")
+            os.makedirs(self.database_backup_dir)
+            print("Done!")
+        database_backup_path = os.path.join(self.database_backup_dir, f"backup-{timestamp_string}-invokeai.db")
+        print(f"Making DB Backup at {database_backup_path}...", end="")
+        shutil.copy2(self.database_path, database_backup_path)
+        print("Done!")
+
+    def connect(self):
+        """Open connection to the database."""
+        self.connection = sqlite3.connect(self.database_path)
+        self.cursor = self.connection.cursor()
+
+    def get_all_image_files(self):
+        """Get the full list of image file names from the database."""
+        sql_get_image_by_name = "SELECT image_name FROM images"
+        self.cursor.execute(sql_get_image_by_name)
+        rows = self.cursor.fetchall()
+        db_files = []
+        for row in rows:
+            db_files.append(row[0])
+        return db_files
+
+    def remove_image_file_record(self, filename: str):
+        """Remove an image file reference from the database by filename."""
+        sanitized_filename = str.replace(filename, "'", "''")  # prevent injection
+        sql_command = f"DELETE FROM images WHERE image_name='{sanitized_filename}'"
+        self.cursor.execute(sql_command)
+        self.connection.commit()
+
+    def does_image_exist(self, image_filename):
+        """Check database if a image name already exists and return a boolean."""
+        sanitized_filename = str.replace(image_filename, "'", "''")  # prevent injection
+        sql_get_image_by_name = f"SELECT image_name FROM images WHERE image_name='{sanitized_filename}'"
+        self.cursor.execute(sql_get_image_by_name)
+        rows = self.cursor.fetchall()
+        return True if len(rows) > 0 else False
+
+    def disconnect(self):
+        """Disconnect from the db, cleaning up connections and cursors."""
+        if self.cursor is not None:
+            self.cursor.close()
+        if self.connection is not None:
+            self.connection.close()
+
+
+class PhysicalFileMapper:
+    """Containing class for script functionality."""
+
+    def __init__(self, outputs_path, thumbnails_path, archive_path, thumbnails_archive_path):  # noqa D107
+        self.outputs_path = outputs_path
+        self.archive_path = archive_path
+        self.thumbnails_path = thumbnails_path
+        self.thumbnails_archive_path = thumbnails_archive_path
+
+    def create_archive_directories(self):
+        """Create the directory for archiving orphaned image files."""
+        if not os.path.exists(self.archive_path):
+            print(f"Image archive directory ({self.archive_path}) does not exist -> creating...", end="")
+            os.makedirs(self.archive_path)
+            print("Created!")
+        if not os.path.exists(self.thumbnails_archive_path):
+            print(
+                f"Image thumbnails archive directory ({self.thumbnails_archive_path}) does not exist -> creating...",
+                end="",
+            )
+            os.makedirs(self.thumbnails_archive_path)
+            print("Created!")
+
+    def get_image_path_for_image_name(self, image_filename):  # noqa D102
+        return os.path.join(self.outputs_path, image_filename)
+
+    def image_file_exists(self, image_filename):  # noqa D102
+        return os.path.exists(self.get_image_path_for_image_name(image_filename))
+
+    def get_thumbnail_path_for_image(self, image_filename):  # noqa D102
+        return os.path.join(self.thumbnails_path, os.path.splitext(image_filename)[0]) + ".webp"
+
+    def get_image_name_from_thumbnail_path(self, thumbnail_path):  # noqa D102
+        return os.path.splitext(os.path.basename(thumbnail_path))[0] + ".png"
+
+    def thumbnail_exists_for_filename(self, image_filename):  # noqa D102
+        return os.path.exists(self.get_thumbnail_path_for_image(image_filename))
+
+    def archive_image(self, image_filename):  # noqa D102
+        if self.image_file_exists(image_filename):
+            image_path = self.get_image_path_for_image_name(image_filename)
+            shutil.move(image_path, self.archive_path)
+
+    def archive_thumbnail_by_image_filename(self, image_filename):  # noqa D102
+        if self.thumbnail_exists_for_filename(image_filename):
+            thumbnail_path = self.get_thumbnail_path_for_image(image_filename)
+            shutil.move(thumbnail_path, self.thumbnails_archive_path)
+
+    def get_all_png_filenames_in_directory(self, directory_path):  # noqa D102
+        filepaths = glob.glob(directory_path + "/*.png", recursive=False)
+        filenames = []
+        for filepath in filepaths:
+            filenames.append(os.path.basename(filepath))
+        return filenames
+
+    def get_all_image_filenames_recursive(self):  # noqa D102
+        """Return the set of all image filenames found anywhere under outputs_path.
+
+        When an image_subfolder_strategy is configured, images are written to
+        subdirectories of outputs/images (e.g. by date, type or hash). This walks
+        the entire tree so that images stored in subfolders are not mistaken for
+        missing files. Thumbnails are stored as .webp and so are excluded by the
+        .png glob; the images-archive directory is a sibling of outputs/images and
+        is therefore not traversed.
+        """
+        filepaths = glob.glob(os.path.join(self.outputs_path, "**", "*.png"), recursive=True)
+        return {os.path.basename(filepath) for filepath in filepaths}
+
+    def get_all_thumbnails_with_full_path(self, thumbnails_directory):  # noqa D102
+        return glob.glob(thumbnails_directory + "/*.webp", recursive=False)
+
+    def generate_thumbnail_for_image_name(self, image_filename):  # noqa D102
+        # create thumbnail
+        file_path = self.get_image_path_for_image_name(image_filename)
+        thumb_path = self.get_thumbnail_path_for_image(image_filename)
+        thumb_size = 256, 256
+        with PIL.Image.open(file_path) as source_image:
+            source_image.thumbnail(thumb_size)
+            source_image.save(thumb_path, "webp")
+
+
+class MaintenanceOperation(str, enum.Enum):
+    """Enum class for operations."""
+
+    Ask = "ask"
+    CleanOrphanedDbEntries = "clean"
+    CleanOrphanedDiskFiles = "archive"
+    ReGenerateThumbnails = "thumbnails"
+    All = "all"
+
+
+class InvokeAIDatabaseMaintenanceApp:
+    """Main processor class for the application."""
+
+    _operation: MaintenanceOperation
+    _headless: bool = False
+    __stats: MaintenanceStats = MaintenanceStats()
+
+    def __init__(self, operation: MaintenanceOperation = MaintenanceOperation.Ask):
+        """Initialize maintenance app."""
+        self._operation = MaintenanceOperation(operation)
+        self._headless = operation != MaintenanceOperation.Ask
+
+    def ask_for_operation(self) -> MaintenanceOperation:
+        """Ask user to choose the operation to perform."""
+        while True:
+            print()
+            print("It is recommennded to run these operations as ordered below to avoid additional")
+            print("work being performed that will be discarded in a subsequent step.")
+            print()
+            print("Select maintenance operation:")
+            print()
+            print("1) Clean Orphaned Database Image Entries")
+            print("     Cleans entries in the database where the matching file was removed from")
+            print("     the outputs directory.")
+            print("2) Archive Orphaned Image Files")
+            print("     Files found in the outputs directory without an entry in the database are")
+            print("     moved to an archive directory.")
+            print("3) Re-Generate Missing Thumbnail Files")
+            print("     For files found in the outputs directory, re-generate a thumbnail if it")
+            print("     not found in the thumbnails directory.")
+            print()
+            print("(CTRL-C to quit)")
+
+            try:
+                input_option = int(input("Specify desired operation number (1-3): "))
+
+                operations = [
+                    MaintenanceOperation.CleanOrphanedDbEntries,
+                    MaintenanceOperation.CleanOrphanedDiskFiles,
+                    MaintenanceOperation.ReGenerateThumbnails,
+                ]
+                return operations[input_option - 1]
+            except (IndexError, ValueError):
+                print("\nInvalid selection!")
+
+    def ask_to_continue(self) -> bool:
+        """Ask user whether they want to continue with the operation."""
+        while True:
+            input_choice = input("Do you wish to continue? (Y or N)? ")
+            if str.lower(input_choice) == "y":
+                return True
+            if str.lower(input_choice) == "n":
+                return False
+
+    def clean_orphaned_db_entries(
+        self, config: ConfigMapper, file_mapper: PhysicalFileMapper, db_mapper: DatabaseMapper
+    ):
+        """Clean dangling database entries that no longer point to a file in outputs."""
+        if self._headless:
+            print(f"Removing database references to images that no longer exist in {config.outputs_path}...")
+        else:
+            print()
+            print("===============================================================================")
+            print("= Clean Orphaned Database Entries")
+            print()
+            print("Perform this operation if you have removed files from the outputs/images")
+            print("directory but the database was never updated. You may see this as empty imaages")
+            print("in the app gallery, or images that only show an enlarged version of the")
+            print("thumbnail.")
+            print()
+            print(f"Database File Path               : {config.database_path}")
+            print(f"Database backup will be taken at : {config.database_backup_dir}")
+            print(f"Outputs/Images Directory         : {config.outputs_path}")
+            print(f"Outputs/Images Archive Directory : {config.archive_path}")
+
+            print("\nNotes about this operation:")
+            print("- This operation will find database image file entries that do not exist in the")
+            print("  outputs/images dir and remove those entries from the database.")
+            print("- This operation will target all image types including intermediate files.")
+            print("- If a thumbnail still exists in outputs/images/thumbnails matching the")
+            print("  orphaned entry, it will be moved to the archive directory.")
+            print()
+
+            if not self.ask_to_continue():
+                raise KeyboardInterrupt
+
+        file_mapper.create_archive_directories()
+        db_mapper.backup(config.TIMESTAMP_STRING)
+        db_mapper.connect()
+        db_files = db_mapper.get_all_image_files()
+        # Build an index of every image present anywhere under outputs/images, including any
+        # subfolders created by an image_subfolder_strategy. A db entry is only orphaned if its
+        # file is absent from this entire tree, not just the top-level outputs/images directory.
+        disk_image_filenames = file_mapper.get_all_image_filenames_recursive()
+        for db_file in db_files:
+            try:
+                if db_file not in disk_image_filenames:
+                    print(f"Found orphaned image db entry {db_file}. Cleaning ...", end="")
+                    db_mapper.remove_image_file_record(db_file)
+                    print("Cleaned!")
+                    if file_mapper.thumbnail_exists_for_filename(db_file):
+                        print("A thumbnail was found, archiving ...", end="")
+                        file_mapper.archive_thumbnail_by_image_filename(db_file)
+                        print("Archived!")
+                    self.__stats.count_orphaned_db_entries_cleaned += 1
+            except Exception as ex:
+                print("An error occurred cleaning db entry, error was:")
+                print(ex)
+                self.__stats.count_errors += 1
+
+    def clean_orphaned_disk_files(
+        self, config: ConfigMapper, file_mapper: PhysicalFileMapper, db_mapper: DatabaseMapper
+    ):
+        """Archive image files that no longer have entries in the database."""
+        if self._headless:
+            print(f"Archiving orphaned image files to {config.archive_path}...")
+        else:
+            print()
+            print("===============================================================================")
+            print("= Clean Orphaned Disk Files")
+            print()
+            print("Perform this operation if you have files that were copied into the outputs")
+            print("directory which are not referenced by the database. This can happen if you")
+            print("upgraded to a version with a fresh database, but re-used the outputs directory")
+            print("and now new images are mixed with the files not in the db. The script will")
+            print("archive these files so you can choose to delete them or re-import using the")
+            print("official import script.")
+            print()
+            print(f"Database File Path               : {config.database_path}")
+            print(f"Database backup will be taken at : {config.database_backup_dir}")
+            print(f"Outputs/Images Directory         : {config.outputs_path}")
+            print(f"Outputs/Images Archive Directory : {config.archive_path}")
+
+            print("\nNotes about this operation:")
+            print("- This operation will find image files not referenced by the database and move to an")
+            print("  archive directory.")
+            print("- This operation will target all image types including intermediate references.")
+            print("- The matching thumbnail will also be archived.")
+            print("- Any remaining orphaned thumbnails will also be archived.")
+
+            if not self.ask_to_continue():
+                raise KeyboardInterrupt
+
+            print()
+
+        file_mapper.create_archive_directories()
+        db_mapper.backup(config.TIMESTAMP_STRING)
+        db_mapper.connect()
+        phys_files = file_mapper.get_all_png_filenames_in_directory(config.outputs_path)
+        for phys_file in phys_files:
+            try:
+                if not db_mapper.does_image_exist(phys_file):
+                    print(f"Found orphaned file {phys_file}, archiving...", end="")
+                    file_mapper.archive_image(phys_file)
+                    print("Archived!")
+                    if file_mapper.thumbnail_exists_for_filename(phys_file):
+                        print("Related thumbnail exists, archiving...", end="")
+                        file_mapper.archive_thumbnail_by_image_filename(phys_file)
+                        print("Archived!")
+                    else:
+                        print("No matching thumbnail existed to be cleaned.")
+                    self.__stats.count_orphaned_disk_files_cleaned += 1
+            except Exception as ex:
+                print("Error found trying to archive file or thumbnail, error was:")
+                print(ex)
+                self.__stats.count_errors += 1
+
+        thumb_filepaths = file_mapper.get_all_thumbnails_with_full_path(config.thumbnails_path)
+        # archive any remaining orphaned thumbnails
+        for thumb_filepath in thumb_filepaths:
+            try:
+                thumb_src_image_name = file_mapper.get_image_name_from_thumbnail_path(thumb_filepath)
+                if not file_mapper.image_file_exists(thumb_src_image_name):
+                    print(f"Found orphaned thumbnail {thumb_filepath}, archiving...", end="")
+                    file_mapper.archive_thumbnail_by_image_filename(thumb_src_image_name)
+                    print("Archived!")
+                    self.__stats.count_orphaned_thumbnails_cleaned += 1
+            except Exception as ex:
+                print("Error found trying to archive thumbnail, error was:")
+                print(ex)
+                self.__stats.count_errors += 1
+
+    def regenerate_thumbnails(self, config: ConfigMapper, file_mapper: PhysicalFileMapper, *args):
+        """Create missing thumbnails for any valid general images both in the db and on disk."""
+        if self._headless:
+            print("Regenerating missing image thumbnails...")
+        else:
+            print()
+            print("===============================================================================")
+            print("= Regenerate Thumbnails")
+            print()
+            print("This operation will find files that have no matching thumbnail on disk")
+            print("and regenerate those thumbnail files.")
+            print("NOTE: It is STRONGLY recommended that the user first clean/archive orphaned")
+            print("  disk files from the previous menu to avoid wasting time regenerating")
+            print("  thumbnails for orphaned files.")
+
+            print()
+            print(f"Outputs/Images Directory         : {config.outputs_path}")
+            print(f"Outputs/Images Directory         : {config.thumbnails_path}")
+
+            print("\nNotes about this operation:")
+            print("- This operation will find image files both referenced in the db and on disk")
+            print("  that do not have a matching thumbnail on disk and re-generate the thumbnail")
+            print("  file.")
+
+            if not self.ask_to_continue():
+                raise KeyboardInterrupt
+
+            print()
+
+        phys_files = file_mapper.get_all_png_filenames_in_directory(config.outputs_path)
+        for phys_file in phys_files:
+            try:
+                if not file_mapper.thumbnail_exists_for_filename(phys_file):
+                    print(f"Found file without thumbnail {phys_file}...Regenerating Thumbnail...", end="")
+                    file_mapper.generate_thumbnail_for_image_name(phys_file)
+                    print("Done!")
+                    self.__stats.count_thumbnails_regenerated += 1
+            except Exception as ex:
+                print("Error found trying to regenerate thumbnail, error was:")
+                print(ex)
+                self.__stats.count_errors += 1
+
+    def main(self):  # noqa D107
+        print("\n===============================================================================")
+        print("Database and outputs Maintenance for Invoke AI 3.0.0 +")
+        print("===============================================================================\n")
+
+        config_mapper = ConfigMapper()
+        if not config_mapper.load():
+            print("\nInvalid configuration...exiting.\n")
+            return
+
+        file_mapper = PhysicalFileMapper(
+            config_mapper.outputs_path,
+            config_mapper.thumbnails_path,
+            config_mapper.archive_path,
+            config_mapper.thumbnails_archive_path,
+        )
+        db_mapper = DatabaseMapper(config_mapper.database_path, config_mapper.database_backup_dir)
+
+        op = self._operation
+        operations_to_perform = []
+
+        if op == MaintenanceOperation.Ask:
+            op = self.ask_for_operation()
+
+        if op in [MaintenanceOperation.CleanOrphanedDbEntries, MaintenanceOperation.All]:
+            operations_to_perform.append(self.clean_orphaned_db_entries)
+        if op in [MaintenanceOperation.CleanOrphanedDiskFiles, MaintenanceOperation.All]:
+            operations_to_perform.append(self.clean_orphaned_disk_files)
+        if op in [MaintenanceOperation.ReGenerateThumbnails, MaintenanceOperation.All]:
+            operations_to_perform.append(self.regenerate_thumbnails)
+
+        for operation in operations_to_perform:
+            operation(config_mapper, file_mapper, db_mapper)
+
+        print("\n===============================================================================")
+        print(f"= Maintenance Complete - Elapsed Time: {MaintenanceStats.get_elapsed_time_string()}")
+        print()
+        print(f"Orphaned db entries cleaned             : {self.__stats.count_orphaned_db_entries_cleaned}")
+        print(f"Orphaned disk files archived            : {self.__stats.count_orphaned_disk_files_cleaned}")
+        print(f"Orphaned thumbnail files archived       : {self.__stats.count_orphaned_thumbnails_cleaned}")
+        print(f"Thumbnails regenerated                  : {self.__stats.count_thumbnails_regenerated}")
+        print(f"Errors during operation                 : {self.__stats.count_errors}")
+
+        print()
+
+
+def main():  # noqa D107
+    parser = argparse.ArgumentParser(
+        description="InvokeAI image database maintenance utility",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""Operations:
+  ask         Choose operation from a menu [default]
+  all         Run all maintenance operations
+  clean       Clean database of dangling entries
+  archive     Archive orphaned image files
+  thumbnails  Regenerate missing image thumbnails
+""",
+    )
+    parser.add_argument("--root", default=".", type=Path, help="InvokeAI root directory")
+    parser.add_argument(
+        "--operation", default="ask", choices=[x.value for x in MaintenanceOperation], help="Operation to perform."
+    )
+    args = parser.parse_args()
+    try:
+        os.chdir(args.root)
+        app = InvokeAIDatabaseMaintenanceApp(args.operation)
+        app.main()
+    except KeyboardInterrupt:
+        print("\n\nUser cancelled execution.")
+    except FileNotFoundError:
+        print(f"Invalid root directory '{args.root}'.")
+
+
+if __name__ == "__main__":
+    main()

invokeai/backend/util/hotfixes.py

@@ -0,0 +1,843 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import diffusers
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor
+from diffusers.models.controlnets.controlnet import ControlNetConditioningEmbedding, ControlNetOutput, zero_module
+from diffusers.models.embeddings import (
+    TextImageProjection,
+    TextImageTimeEmbedding,
+    TextTimeEmbedding,
+    TimestepEmbedding,
+    Timesteps,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.unets.unet_2d_blocks import (
+    CrossAttnDownBlock2D,
+    DownBlock2D,
+    UNetMidBlock2DCrossAttn,
+    get_down_block,
+)
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+from torch import nn
+
+from invokeai.backend.model_manager.taxonomy import BaseModelType, SchedulerPredictionType
+from invokeai.backend.util.logging import InvokeAILogger
+
+# TODO: create PR to diffusers
+# Modified ControlNetModel with encoder_attention_mask argument added
+
+
+logger = InvokeAILogger.get_logger(__name__)
+
+
+# NOTE(ryand): I'm not the origina author of this code, but for future reference, it appears that this class was copied
+# from diffusers in order to add support for the encoder_attention_mask argument.
+class ControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    """
+    A ControlNet model.
+
+    Args:
+        in_channels (`int`, defaults to 4):
+            The number of channels in the input sample.
+        flip_sin_to_cos (`bool`, defaults to `True`):
+            Whether to flip the sin to cos in the time embedding.
+        freq_shift (`int`, defaults to 0):
+            The frequency shift to apply to the time embedding.
+        down_block_types (`tuple[str]`, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", \
+            "CrossAttnDownBlock2D", "DownBlock2D")`):
+            The tuple of downsample blocks to use.
+        only_cross_attention (`Union[bool, Tuple[bool]]`, defaults to `False`):
+        block_out_channels (`tuple[int]`, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        layers_per_block (`int`, defaults to 2):
+            The number of layers per block.
+        downsample_padding (`int`, defaults to 1):
+            The padding to use for the downsampling convolution.
+        mid_block_scale_factor (`float`, defaults to 1):
+            The scale factor to use for the mid block.
+        act_fn (`str`, defaults to "silu"):
+            The activation function to use.
+        norm_num_groups (`int`, *optional*, defaults to 32):
+            The number of groups to use for the normalization. If None, normalization and activation layers is skipped
+            in post-processing.
+        norm_eps (`float`, defaults to 1e-5):
+            The epsilon to use for the normalization.
+        cross_attention_dim (`int`, defaults to 1280):
+            The dimension of the cross attention features.
+        transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
+            [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
+            [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
+        encoder_hid_dim (`int`, *optional*, defaults to None):
+            If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
+            dimension to `cross_attention_dim`.
+        encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
+            If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
+            embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
+        attention_head_dim (`Union[int, Tuple[int]]`, defaults to 8):
+            The dimension of the attention heads.
+        use_linear_projection (`bool`, defaults to `False`):
+        class_embed_type (`str`, *optional*, defaults to `None`):
+            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None,
+            `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
+        addition_embed_type (`str`, *optional*, defaults to `None`):
+            Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
+            "text". "text" will use the `TextTimeEmbedding` layer.
+        num_class_embeds (`int`, *optional*, defaults to 0):
+            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
+            class conditioning with `class_embed_type` equal to `None`.
+        upcast_attention (`bool`, defaults to `False`):
+        resnet_time_scale_shift (`str`, defaults to `"default"`):
+            Time scale shift config for ResNet blocks (see `ResnetBlock2D`). Choose from `default` or `scale_shift`.
+        projection_class_embeddings_input_dim (`int`, *optional*, defaults to `None`):
+            The dimension of the `class_labels` input when `class_embed_type="projection"`. Required when
+            `class_embed_type="projection"`.
+        controlnet_conditioning_channel_order (`str`, defaults to `"rgb"`):
+            The channel order of conditional image. Will convert to `rgb` if it's `bgr`.
+        conditioning_embedding_out_channels (`tuple[int]`, *optional*, defaults to `(16, 32, 96, 256)`):
+            The tuple of output channel for each block in the `conditioning_embedding` layer.
+        global_pool_conditions (`bool`, defaults to `False`):
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 4,
+        conditioning_channels: int = 3,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "DownBlock2D",
+        ),
+        only_cross_attention: Union[bool, Tuple[bool]] = False,
+        block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+        layers_per_block: int = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        act_fn: str = "silu",
+        norm_num_groups: Optional[int] = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: int = 1280,
+        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+        encoder_hid_dim: Optional[int] = None,
+        encoder_hid_dim_type: Optional[str] = None,
+        attention_head_dim: Union[int, Tuple[int]] = 8,
+        num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+        use_linear_projection: bool = False,
+        class_embed_type: Optional[str] = None,
+        addition_embed_type: Optional[str] = None,
+        addition_time_embed_dim: Optional[int] = None,
+        num_class_embeds: Optional[int] = None,
+        upcast_attention: bool = False,
+        resnet_time_scale_shift: str = "default",
+        projection_class_embeddings_input_dim: Optional[int] = None,
+        controlnet_conditioning_channel_order: str = "rgb",
+        conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),
+        global_pool_conditions: bool = False,
+        addition_embed_type_num_heads=64,
+    ):
+        super().__init__()
+
+        # If `num_attention_heads` is not defined (which is the case for most models)
+        # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
+        # The reason for this behavior is to correct for incorrectly named variables that were introduced
+        # when this library was created...
+        # The incorrect naming was only discovered much ...
+        # later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
+        # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
+        # which is why we correct for the naming here.
+        num_attention_heads = num_attention_heads or attention_head_dim
+
+        # Check inputs
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. \
+                    `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `only_cross_attention` as `down_block_types`. \
+                    `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `num_attention_heads` as `down_block_types`. \
+                    `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+            )
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        # input
+        conv_in_kernel = 3
+        conv_in_padding = (conv_in_kernel - 1) // 2
+        self.conv_in = nn.Conv2d(
+            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
+        )
+
+        # time
+        time_embed_dim = block_out_channels[0] * 4
+        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+        timestep_input_dim = block_out_channels[0]
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+            act_fn=act_fn,
+        )
+
+        if encoder_hid_dim_type is None and encoder_hid_dim is not None:
+            encoder_hid_dim_type = "text_proj"
+            self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
+            logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
+
+        if encoder_hid_dim is None and encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
+            )
+
+        if encoder_hid_dim_type == "text_proj":
+            self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
+        elif encoder_hid_dim_type == "text_image_proj":
+            # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension ...
+            # for the currently only use
+            # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
+            self.encoder_hid_proj = TextImageProjection(
+                text_embed_dim=encoder_hid_dim,
+                image_embed_dim=cross_attention_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+
+        elif encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
+            )
+        else:
+            self.encoder_hid_proj = None
+
+        # class embedding
+        if class_embed_type is None and num_class_embeds is not None:
+            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+        elif class_embed_type == "timestep":
+            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+        elif class_embed_type == "identity":
+            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        elif class_embed_type == "projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
+            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
+            # 2. it projects from an arbitrary input dimension.
+            #
+            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
+            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
+            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
+            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        else:
+            self.class_embedding = None
+
+        if addition_embed_type == "text":
+            if encoder_hid_dim is not None:
+                text_time_embedding_from_dim = encoder_hid_dim
+            else:
+                text_time_embedding_from_dim = cross_attention_dim
+
+            self.add_embedding = TextTimeEmbedding(
+                text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
+            )
+        elif addition_embed_type == "text_image":
+            # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`.
+            # To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension...
+            # for the currently only use
+            # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
+            self.add_embedding = TextImageTimeEmbedding(
+                text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
+            )
+        elif addition_embed_type == "text_time":
+            self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
+            self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+
+        elif addition_embed_type is not None:
+            raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
+
+        # control net conditioning embedding
+        self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
+            conditioning_embedding_channels=block_out_channels[0],
+            block_out_channels=conditioning_embedding_out_channels,
+            conditioning_channels=conditioning_channels,
+        )
+
+        self.down_blocks = nn.ModuleList([])
+        self.controlnet_down_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        # down
+        output_channel = block_out_channels[0]
+
+        controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+        controlnet_block = zero_module(controlnet_block)
+        self.controlnet_down_blocks.append(controlnet_block)
+
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block,
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                num_attention_heads=num_attention_heads[i],
+                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                downsample_padding=downsample_padding,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+            )
+            self.down_blocks.append(down_block)
+
+            for _ in range(layers_per_block):
+                controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                controlnet_block = zero_module(controlnet_block)
+                self.controlnet_down_blocks.append(controlnet_block)
+
+            if not is_final_block:
+                controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                controlnet_block = zero_module(controlnet_block)
+                self.controlnet_down_blocks.append(controlnet_block)
+
+        # mid
+        mid_block_channel = block_out_channels[-1]
+
+        controlnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)
+        controlnet_block = zero_module(controlnet_block)
+        self.controlnet_mid_block = controlnet_block
+
+        self.mid_block = UNetMidBlock2DCrossAttn(
+            transformer_layers_per_block=transformer_layers_per_block[-1],
+            in_channels=mid_block_channel,
+            temb_channels=time_embed_dim,
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            output_scale_factor=mid_block_scale_factor,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            cross_attention_dim=cross_attention_dim,
+            num_attention_heads=num_attention_heads[-1],
+            resnet_groups=norm_num_groups,
+            use_linear_projection=use_linear_projection,
+            upcast_attention=upcast_attention,
+        )
+
+    @classmethod
+    def from_unet(
+        cls,
+        unet: UNet2DConditionModel,
+        controlnet_conditioning_channel_order: str = "rgb",
+        conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),
+        load_weights_from_unet: bool = True,
+    ):
+        r"""
+        Instantiate a [`ControlNetModel`] from [`UNet2DConditionModel`].
+
+        Parameters:
+            unet (`UNet2DConditionModel`):
+                The UNet model weights to copy to the [`ControlNetModel`]. All configuration options are also copied
+                where applicable.
+        """
+        transformer_layers_per_block = (
+            unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1
+        )
+        encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None
+        encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None
+        addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None
+        addition_time_embed_dim = (
+            unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None
+        )
+
+        controlnet = cls(
+            encoder_hid_dim=encoder_hid_dim,
+            encoder_hid_dim_type=encoder_hid_dim_type,
+            addition_embed_type=addition_embed_type,
+            addition_time_embed_dim=addition_time_embed_dim,
+            transformer_layers_per_block=transformer_layers_per_block,
+            in_channels=unet.config.in_channels,
+            flip_sin_to_cos=unet.config.flip_sin_to_cos,
+            freq_shift=unet.config.freq_shift,
+            down_block_types=unet.config.down_block_types,
+            only_cross_attention=unet.config.only_cross_attention,
+            block_out_channels=unet.config.block_out_channels,
+            layers_per_block=unet.config.layers_per_block,
+            downsample_padding=unet.config.downsample_padding,
+            mid_block_scale_factor=unet.config.mid_block_scale_factor,
+            act_fn=unet.config.act_fn,
+            norm_num_groups=unet.config.norm_num_groups,
+            norm_eps=unet.config.norm_eps,
+            cross_attention_dim=unet.config.cross_attention_dim,
+            attention_head_dim=unet.config.attention_head_dim,
+            num_attention_heads=unet.config.num_attention_heads,
+            use_linear_projection=unet.config.use_linear_projection,
+            class_embed_type=unet.config.class_embed_type,
+            num_class_embeds=unet.config.num_class_embeds,
+            upcast_attention=unet.config.base is BaseModelType.StableDiffusion2
+            and unet.config.prediction_type is SchedulerPredictionType.VPrediction,
+            resnet_time_scale_shift=unet.config.resnet_time_scale_shift,
+            projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,
+            controlnet_conditioning_channel_order=controlnet_conditioning_channel_order,
+            conditioning_embedding_out_channels=conditioning_embedding_out_channels,
+        )
+
+        if load_weights_from_unet:
+            controlnet.conv_in.load_state_dict(unet.conv_in.state_dict())
+            controlnet.time_proj.load_state_dict(unet.time_proj.state_dict())
+            controlnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
+
+            if controlnet.class_embedding:
+                controlnet.class_embedding.load_state_dict(unet.class_embedding.state_dict())
+
+            controlnet.down_blocks.load_state_dict(unet.down_blocks.state_dict())
+            controlnet.mid_block.load_state_dict(unet.mid_block.state_dict())
+
+        return controlnet
+
+    @property
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "set_processor"):
+                processors[f"{name}.processor"] = module.processor
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        self.set_attn_processor(AttnProcessor())
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice
+    def set_attention_slice(self, slice_size):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module splits the input tensor in slices to compute attention in
+        several steps. This is useful for saving some memory in exchange for a small decrease in speed.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
+                `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_sliceable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_sliceable_dims(module)
+
+        num_sliceable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_sliceable_layers * [1]
+
+        slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        controlnet_cond: torch.FloatTensor,
+        conditioning_scale: float = 1.0,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        guess_mode: bool = False,
+        return_dict: bool = True,
+    ) -> Union[ControlNetOutput, Tuple]:
+        """
+        The [`ControlNetModel`] forward method.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The noisy input tensor.
+            timestep (`Union[torch.Tensor, float, int]`):
+                The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.Tensor`):
+                The encoder hidden states.
+            controlnet_cond (`torch.FloatTensor`):
+                The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`.
+            conditioning_scale (`float`, defaults to `1.0`):
+                The scale factor for ControlNet outputs.
+            class_labels (`torch.Tensor`, *optional*, defaults to `None`):
+                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
+            timestep_cond (`torch.Tensor`, *optional*, defaults to `None`):
+            attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
+            added_cond_kwargs (`dict`):
+                Additional conditions for the Stable Diffusion XL UNet.
+            cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`):
+                A kwargs dictionary that if specified is passed along to the `AttnProcessor`.
+            encoder_attention_mask (`torch.Tensor`):
+                A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
+                `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
+                which adds large negative values to the attention scores corresponding to "discard" tokens.
+            guess_mode (`bool`, defaults to `False`):
+                In this mode, the ControlNet encoder tries its best to recognize the input content of the input even if
+                you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a [`~models.controlnet.ControlNetOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.controlnet.ControlNetOutput`] **or** `tuple`:
+                If `return_dict` is `True`, a [`~models.controlnet.ControlNetOutput`] is returned, otherwise a tuple is
+                returned where the first element is the sample tensor.
+        """
+        # check channel order
+        channel_order = self.config.controlnet_conditioning_channel_order
+
+        if channel_order == "rgb":
+            # in rgb order by default
+            ...
+        elif channel_order == "bgr":
+            controlnet_cond = torch.flip(controlnet_cond, dims=[1])
+        else:
+            raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}")
+
+        # prepare attention_mask
+        if attention_mask is not None:
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+            emb = emb + class_emb
+
+        if "addition_embed_type" in self.config:
+            if self.config.addition_embed_type == "text":
+                aug_emb = self.add_embedding(encoder_hidden_states)
+
+            elif self.config.addition_embed_type == "text_time":
+                if "text_embeds" not in added_cond_kwargs:
+                    raise ValueError(
+                        f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which \
+                            requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                    )
+                text_embeds = added_cond_kwargs.get("text_embeds")
+                if "time_ids" not in added_cond_kwargs:
+                    raise ValueError(
+                        f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which \
+                            requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                    )
+                time_ids = added_cond_kwargs.get("time_ids")
+                time_embeds = self.add_time_proj(time_ids.flatten())
+                time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+                add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+                add_embeds = add_embeds.to(emb.dtype)
+                aug_emb = self.add_embedding(add_embeds)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
+        sample = sample + controlnet_cond
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+            down_block_res_samples += res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+                cross_attention_kwargs=cross_attention_kwargs,
+                encoder_attention_mask=encoder_attention_mask,
+            )
+
+        # 5. Control net blocks
+
+        controlnet_down_block_res_samples = ()
+
+        for down_block_res_sample, controlnet_block in zip(
+            down_block_res_samples, self.controlnet_down_blocks, strict=True
+        ):
+            down_block_res_sample = controlnet_block(down_block_res_sample)
+            controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,)
+
+        down_block_res_samples = controlnet_down_block_res_samples
+
+        mid_block_res_sample = self.controlnet_mid_block(sample)
+
+        # 6. scaling
+        if guess_mode and not self.config.global_pool_conditions:
+            scales = torch.logspace(-1, 0, len(down_block_res_samples) + 1, device=sample.device)  # 0.1 to 1.0
+
+            scales = scales * conditioning_scale
+            down_block_res_samples = [
+                sample * scale for sample, scale in zip(down_block_res_samples, scales, strict=False)
+            ]
+            mid_block_res_sample = mid_block_res_sample * scales[-1]  # last one
+        else:
+            down_block_res_samples = [sample * conditioning_scale for sample in down_block_res_samples]
+            mid_block_res_sample = mid_block_res_sample * conditioning_scale
+
+        if self.config.global_pool_conditions:
+            down_block_res_samples = [torch.mean(sample, dim=(2, 3), keepdim=True) for sample in down_block_res_samples]
+            mid_block_res_sample = torch.mean(mid_block_res_sample, dim=(2, 3), keepdim=True)
+
+        if not return_dict:
+            return (down_block_res_samples, mid_block_res_sample)
+
+        return ControlNetOutput(
+            down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample
+        )
+
+
+diffusers.ControlNetModel = ControlNetModel
+diffusers.models.controlnets.controlnet.ControlNetModel = ControlNetModel
+
+
+# patch LoRACompatibleConv to use original Conv2D forward function
+# this needed to make work seamless patch
+# NOTE: with this patch, torch.compile crashes on 2.0 torch(already fixed in nightly)
+# https://github.com/huggingface/diffusers/pull/4315
+# https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/lora.py#L96C18-L96C18
+def new_LoRACompatibleConv_forward(self, hidden_states, scale: float = 1.0):
+    if self.lora_layer is None:
+        return super(diffusers.models.lora.LoRACompatibleConv, self).forward(hidden_states)
+    else:
+        return super(diffusers.models.lora.LoRACompatibleConv, self).forward(hidden_states) + (
+            scale * self.lora_layer(hidden_states)
+        )
+
+
+diffusers.models.lora.LoRACompatibleConv.forward = new_LoRACompatibleConv_forward
+
+try:
+    import xformers
+
+    xformers_available = True
+except Exception:
+    xformers_available = False
+
+
+if xformers_available:
+    # TODO: remove when fixed in diffusers
+    _xformers_memory_efficient_attention = xformers.ops.memory_efficient_attention
+
+    def new_memory_efficient_attention(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_bias=None,
+        p: float = 0.0,
+        scale: Optional[float] = None,
+        *,
+        op=None,
+    ):
+        # diffusers not align shape to 8, which is required by xformers
+        if attn_bias is not None and type(attn_bias) is torch.Tensor:
+            orig_size = attn_bias.shape[-1]
+            new_size = ((orig_size + 7) // 8) * 8
+            aligned_attn_bias = torch.zeros(
+                (attn_bias.shape[0], attn_bias.shape[1], new_size),
+                device=attn_bias.device,
+                dtype=attn_bias.dtype,
+            )
+            aligned_attn_bias[:, :, :orig_size] = attn_bias
+            attn_bias = aligned_attn_bias[:, :, :orig_size]
+
+        return _xformers_memory_efficient_attention(
+            query=query,
+            key=key,
+            value=value,
+            attn_bias=attn_bias,
+            p=p,
+            scale=scale,
+            op=op,
+        )
+
+    xformers.ops.memory_efficient_attention = new_memory_efficient_attention

invokeai/backend/util/logging.py

@@ -0,0 +1,435 @@
+# Copyright (c) 2023 Lincoln D. Stein and The InvokeAI Development Team
+
+"""
+Logging class for InvokeAI that produces console messages.
+
+Usage:
+
+from invokeai.backend.util.logging import InvokeAILogger
+
+logger = InvokeAILogger.get_logger(name='InvokeAI') // Initialization
+(or)
+logger = InvokeAILogger.get_logger(__name__) // To use the filename
+logger.configure()
+
+logger.critical('this is critical') // Critical Message
+logger.error('this is an error') // Error Message
+logger.warning('this is a warning') // Warning Message
+logger.info('this is info') // Info Message
+logger.debug('this is debugging') // Debug Message
+
+Console messages:
+    [12-05-2023 20]::[InvokeAI]::CRITICAL --> This is an info message [In Bold Red]
+    [12-05-2023 20]::[InvokeAI]::ERROR --> This is an info message [In Red]
+    [12-05-2023 20]::[InvokeAI]::WARNING --> This is an info message [In Yellow]
+    [12-05-2023 20]::[InvokeAI]::INFO --> This is an info message [In Grey]
+    [12-05-2023 20]::[InvokeAI]::DEBUG --> This is an info message [In Grey]
+
+Alternate Method (in this case the logger name will be set to InvokeAI):
+import invokeai.backend.util.logging as IAILogger
+IAILogger.debug('this is a debugging message')
+
+## Configuration
+
+The default configuration will print to stderr on the console. To add
+additional logging handlers, call get_logger with an initialized InvokeAIAppConfig
+object:
+
+
+ config = InvokeAIAppConfig.get_config()
+ config.parse_args()
+ logger = InvokeAILogger.get_logger(config=config)
+
+### Three command-line options control logging:
+
+`--log_handlers <handler1> <handler2> ...`
+
+This option activates one or more log handlers. Options are "console", "file", "syslog" and "http". To specify more than one, separate them by spaces:
+
+```
+invokeai-web --log_handlers console syslog=/dev/log file=C:\\Users\\fred\\invokeai.log
+```
+
+The format of these options is described below.
+
+### `--log_format {plain|color|legacy|syslog}`
+
+This controls the format of log messages written to the console. Only the "console" log handler is currently affected by this setting.
+
+* "plain" provides formatted messages like this:
+
+```bash
+
+[2023-05-24 23:18:2[2023-05-24 23:18:50,352]::[InvokeAI]::DEBUG --> this is a debug message
+[2023-05-24 23:18:50,352]::[InvokeAI]::INFO --> this is an informational messages
+[2023-05-24 23:18:50,352]::[InvokeAI]::WARNING --> this is a warning
+[2023-05-24 23:18:50,352]::[InvokeAI]::ERROR --> this is an error
+[2023-05-24 23:18:50,352]::[InvokeAI]::CRITICAL --> this is a critical error
+```
+
+* "color" produces similar output, but the text will be color coded to indicate the severity of the message.
+
+* "legacy" produces output similar to InvokeAI versions 2.3 and earlier:
+
+```
+### this is a critical error
+*** this is an error
+** this is a warning
+>> this is an informational messages
+   | this is a debug message
+```
+
+* "syslog" produces messages suitable for syslog entries:
+
+```bash
+InvokeAI [2691178] <CRITICAL> this is a critical error
+InvokeAI [2691178] <ERROR> this is an error
+InvokeAI [2691178] <WARNING> this is a warning
+InvokeAI [2691178] <INFO> this is an informational messages
+InvokeAI [2691178] <DEBUG> this is a debug message
+```
+
+(note that the date, time and hostname will be added by the syslog system)
+
+### `--log_level {debug|info|warning|error|critical}`
+
+Providing this command-line option will cause only messages at the specified level or above to be emitted.
+
+## Console logging
+
+When "console" is provided to `--log_handlers`, messages will be written to the command line window in which InvokeAI was launched. By default, the color formatter will be used unless overridden by `--log_format`.
+
+## File logging
+
+When "file" is provided to `--log_handlers`, entries will be written to the file indicated in the path argument. By default, the "plain" format will be used:
+
+```bash
+invokeai-web --log_handlers file=/var/log/invokeai.log
+```
+
+## Syslog logging
+
+When "syslog" is requested, entries will be sent to the syslog system. There are a variety of ways to control where the log message is sent:
+
+* Send to the local machine using the `/dev/log` socket:
+
+```
+invokeai-web --log_handlers syslog=/dev/log
+```
+
+* Send to the local machine using a UDP message:
+
+```
+invokeai-web --log_handlers syslog=localhost
+```
+
+* Send to the local machine using a UDP message on a nonstandard port:
+
+```
+invokeai-web --log_handlers syslog=localhost:512
+```
+
+* Send to a remote machine named "loghost" on the local LAN using  facility LOG_USER and UDP packets:
+
+```
+invokeai-web --log_handlers syslog=loghost,facility=LOG_USER,socktype=SOCK_DGRAM
+```
+
+This can be abbreviated `syslog=loghost`, as LOG_USER and SOCK_DGRAM are defaults.
+
+* Send to a remote machine named "loghost" using the facility LOCAL0  and using a TCP socket:
+
+```
+invokeai-web --log_handlers syslog=loghost,facility=LOG_LOCAL0,socktype=SOCK_STREAM
+```
+
+If no arguments are specified (just a bare "syslog"), then the logging system will look for a UNIX socket named `/dev/log`, and if not found try to send a UDP message to `localhost`. The Macintosh OS used to support logging to a socket named `/var/run/syslog`, but this feature has since been disabled.
+
+## Web logging
+
+If you have access to a web server that is configured to log messages when a particular URL is requested, you can log using the "http" method:
+
+```
+invokeai-web --log_handlers http=http://my.server/path/to/logger,method=POST
+```
+
+The optional [,method=] part can be used to specify whether the URL accepts GET (default) or POST messages.
+
+Currently password authentication and SSL are not supported.
+
+## Using the configuration file
+
+You can set and forget logging options by adding a "Logging" section to `invokeai.yaml`:
+
+```
+InvokeAI:
+  [... other settings...]
+  Logging:
+    log_handlers:
+       - console
+       - syslog=/dev/log
+    log_level: info
+    log_format: color
+```
+
+"""
+
+import logging.handlers
+import socket
+import urllib.parse
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.config.config_default import get_config
+
+try:
+    import syslog
+
+    SYSLOG_AVAILABLE = True
+except ImportError:
+    SYSLOG_AVAILABLE = False
+
+
+# module level functions
+def debug(msg: str, *args: str, **kwargs: Any) -> None:  # noqa D103
+    InvokeAILogger.get_logger().debug(msg, *args, **kwargs)
+
+
+def info(msg: str, *args: str, **kwargs: Any) -> None:  # noqa D103
+    InvokeAILogger.get_logger().info(msg, *args, **kwargs)
+
+
+def warning(msg: str, *args: str, **kwargs: Any) -> None:  # noqa D103
+    InvokeAILogger.get_logger().warning(msg, *args, **kwargs)
+
+
+def error(msg: str, *args: str, **kwargs: Any) -> None:  # noqa D103
+    InvokeAILogger.get_logger().error(msg, *args, **kwargs)
+
+
+def critical(msg: str, *args: str, **kwargs: Any) -> None:  # noqa D103
+    InvokeAILogger.get_logger().critical(msg, *args, **kwargs)
+
+
+def log(level: int, msg: str, *args: str, **kwargs: Any) -> None:  # noqa D103
+    InvokeAILogger.get_logger().log(level, msg, *args, **kwargs)
+
+
+def disable(level: int = logging.CRITICAL) -> None:  # noqa D103
+    logging.disable(level)
+
+
+def basicConfig(**kwargs: Any) -> None:  # noqa D103
+    logging.basicConfig(**kwargs)
+
+
+_FACILITY_MAP = (
+    {
+        "LOG_KERN": syslog.LOG_KERN,
+        "LOG_USER": syslog.LOG_USER,
+        "LOG_MAIL": syslog.LOG_MAIL,
+        "LOG_DAEMON": syslog.LOG_DAEMON,
+        "LOG_AUTH": syslog.LOG_AUTH,
+        "LOG_LPR": syslog.LOG_LPR,
+        "LOG_NEWS": syslog.LOG_NEWS,
+        "LOG_UUCP": syslog.LOG_UUCP,
+        "LOG_CRON": syslog.LOG_CRON,
+        "LOG_SYSLOG": syslog.LOG_SYSLOG,
+        "LOG_LOCAL0": syslog.LOG_LOCAL0,
+        "LOG_LOCAL1": syslog.LOG_LOCAL1,
+        "LOG_LOCAL2": syslog.LOG_LOCAL2,
+        "LOG_LOCAL3": syslog.LOG_LOCAL3,
+        "LOG_LOCAL4": syslog.LOG_LOCAL4,
+        "LOG_LOCAL5": syslog.LOG_LOCAL5,
+        "LOG_LOCAL6": syslog.LOG_LOCAL6,
+        "LOG_LOCAL7": syslog.LOG_LOCAL7,
+    }
+    if SYSLOG_AVAILABLE
+    else {}
+)
+
+_SOCK_MAP = {
+    "SOCK_STREAM": socket.SOCK_STREAM,
+    "SOCK_DGRAM": socket.SOCK_DGRAM,
+}
+
+
+class InvokeAIFormatter(logging.Formatter):
+    """Base class for logging formatter."""
+
+    def format(self, record: logging.LogRecord) -> str:  # noqa D102
+        formatter = logging.Formatter(self.log_fmt(record.levelno))
+        return formatter.format(record)
+
+    def log_fmt(self, levelno: int) -> str:  # noqa D102
+        return "[%(asctime)s]::[%(name)s]::%(levelname)s --> %(message)s"
+
+
+class InvokeAISyslogFormatter(InvokeAIFormatter):
+    """Formatting for syslog."""
+
+    def log_fmt(self, levelno: int) -> str:  # noqa D102
+        return "%(name)s [%(process)d] <%(levelname)s> %(message)s"
+
+
+class InvokeAILegacyLogFormatter(InvokeAIFormatter):  # noqa D102
+    """Formatting for the InvokeAI Logger (legacy version)."""
+
+    FORMATS = {
+        logging.DEBUG: "   | %(message)s",
+        logging.INFO: ">> %(message)s",
+        logging.WARNING: "** %(message)s",
+        logging.ERROR: "*** %(message)s",
+        logging.CRITICAL: "### %(message)s",
+    }
+
+    def log_fmt(self, levelno: int) -> str:  # noqa D102
+        format = self.FORMATS.get(levelno)
+        assert format is not None
+        return format
+
+
+class InvokeAIPlainLogFormatter(InvokeAIFormatter):
+    """Custom Formatting for the InvokeAI Logger (plain version)."""
+
+    def log_fmt(self, levelno: int) -> str:  # noqa D102
+        return "[%(asctime)s]::[%(name)s]::%(levelname)s --> %(message)s"
+
+
+class InvokeAIColorLogFormatter(InvokeAIFormatter):
+    """Custom Formatting for the InvokeAI Logger."""
+
+    # Color Codes
+    grey = "\x1b[38;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    cyan = "\x1b[36;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+
+    # Log Format
+    log_format = "[%(asctime)s]::[%(name)s]::%(levelname)s --> %(message)s"
+    ## More Formatting Options: %(pathname)s, %(filename)s, %(module)s, %(lineno)d
+
+    # Format Map
+    FORMATS = {
+        logging.DEBUG: cyan + log_format + reset,
+        logging.INFO: grey + log_format + reset,
+        logging.WARNING: yellow + log_format + reset,
+        logging.ERROR: red + log_format + reset,
+        logging.CRITICAL: bold_red + log_format + reset,
+    }
+
+    def log_fmt(self, levelno: int) -> str:  # noqa D102
+        format = self.FORMATS.get(levelno)
+        assert format is not None
+        return format
+
+
+LOG_FORMATTERS = {
+    "plain": InvokeAIPlainLogFormatter,
+    "color": InvokeAIColorLogFormatter,
+    "syslog": InvokeAISyslogFormatter,
+    "legacy": InvokeAILegacyLogFormatter,
+}
+
+
+class InvokeAILogger(object):  # noqa D102
+    loggers: Dict[str, logging.Logger] = {}
+
+    @classmethod
+    def get_logger(cls, name: str = "InvokeAI", config: Optional[InvokeAIAppConfig] = None) -> logging.Logger:  # noqa D102
+        config = config or get_config()
+        if name in cls.loggers:
+            return cls.loggers[name]
+
+        logger = logging.getLogger(name)
+        logger.setLevel(config.log_level.upper())  # yes, strings work here
+        for ch in cls.get_loggers(config):
+            logger.addHandler(ch)
+        cls.loggers[name] = logger
+        return cls.loggers[name]
+
+    @classmethod
+    def get_loggers(cls, config: InvokeAIAppConfig) -> list[logging.Handler]:  # noqa D102
+        handler_strs = config.log_handlers
+        handlers = []
+        for handler in handler_strs:
+            handler_name, *args = handler.split("=", 2)
+            arg = args[0] if len(args) > 0 else None
+
+            # console and file get the fancy formatter.
+            # syslog gets a simple one
+            # http gets no custom formatter
+            formatter = LOG_FORMATTERS[config.log_format]
+            if handler_name == "console":
+                ch: logging.Handler = logging.StreamHandler()
+                ch.setFormatter(formatter())
+                handlers.append(ch)
+
+            elif handler_name == "syslog":
+                ch = cls._parse_syslog_args(arg)
+                handlers.append(ch)
+
+            elif handler_name == "file":
+                ch = cls._parse_file_args(arg)
+                ch.setFormatter(formatter())
+                handlers.append(ch)
+
+            elif handler_name == "http":
+                ch = cls._parse_http_args(arg)
+                handlers.append(ch)
+        return handlers
+
+    @staticmethod
+    def _parse_syslog_args(args: Optional[str] = None) -> logging.Handler:
+        if not SYSLOG_AVAILABLE:
+            raise ValueError("syslog is not available on this system")
+        if not args:
+            args = "/dev/log" if Path("/dev/log").exists() else "address:localhost:514"
+        syslog_args: Dict[str, Any] = {}
+        try:
+            for a in args.split(","):
+                arg_name, *arg_value = a.split(":", 2)
+                if arg_name == "address":
+                    host, *port_list = arg_value
+                    port = 514 if not port_list else int(port_list[0])
+                    syslog_args["address"] = (host, port)
+                elif arg_name == "facility":
+                    syslog_args["facility"] = _FACILITY_MAP[arg_value[0]]
+                elif arg_name == "socktype":
+                    syslog_args["socktype"] = _SOCK_MAP[arg_value[0]]
+                else:
+                    syslog_args["address"] = arg_name
+        except Exception:
+            raise ValueError(f"{args} is not a value argument list for syslog logging")
+        return logging.handlers.SysLogHandler(**syslog_args)
+
+    @staticmethod
+    def _parse_file_args(args: Optional[str] = None) -> logging.Handler:  # noqa D102
+        if not args:
+            raise ValueError("please provide filename for file logging using format 'file=/path/to/logfile.txt'")
+        return logging.FileHandler(args)
+
+    @staticmethod
+    def _parse_http_args(args: Optional[str] = None) -> logging.Handler:  # noqa D102
+        if not args:
+            raise ValueError("please provide destination for http logging using format 'http=url'")
+        arg_list = args.split(",")
+        url = urllib.parse.urlparse(arg_list.pop(0))
+        if url.scheme != "http":
+            raise ValueError(f"the http logging module can only log to HTTP URLs, but {url.scheme} was specified")
+        host = url.hostname
+        path = url.path
+        port = url.port or 80
+
+        syslog_args: Dict[str, Any] = {}
+        for a in arg_list:
+            arg_name, *arg_value = a.split(":", 2)
+            if arg_name == "method":
+                method = arg_value[0] if len(arg_value) > 0 else "GET"
+                syslog_args[arg_name] = method
+            else:  # TODO: Provide support for SSL context and credentials
+                pass
+        return logging.handlers.HTTPHandler(f"{host}:{port}", path, **syslog_args)

invokeai/backend/util/mask.py

@@ -0,0 +1,53 @@
+import torch
+
+
+def to_standard_mask_dim(mask: torch.Tensor) -> torch.Tensor:
+    """Standardize the dimensions of a mask tensor.
+
+    Args:
+        mask (torch.Tensor): A mask tensor. The shape can be (1, h, w) or (h, w).
+
+    Returns:
+        torch.Tensor: The output mask tensor. The shape is (1, h, w).
+    """
+    # Get the mask height and width.
+    if mask.ndim == 2:
+        mask = mask.unsqueeze(0)
+    elif mask.ndim == 3 and mask.shape[0] == 1:
+        pass
+    else:
+        raise ValueError(f"Unsupported mask shape: {mask.shape}. Expected (1, h, w) or (h, w).")
+
+    return mask
+
+
+def to_standard_float_mask(mask: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
+    """Standardize the format of a mask tensor.
+
+    Args:
+        mask (torch.Tensor): A mask tensor. The dtype can be any bool, float, or int type. The shape must be (1, h, w)
+            or (h, w).
+
+        out_dtype (torch.dtype): The dtype of the output mask tensor. Must be a float type.
+
+    Returns:
+        torch.Tensor: The output mask tensor. The dtype is out_dtype. The shape is (1, h, w). All values are either 0.0
+            or 1.0.
+    """
+
+    if not out_dtype.is_floating_point:
+        raise ValueError(f"out_dtype must be a float type, but got {out_dtype}")
+
+    mask = to_standard_mask_dim(mask)
+    mask = mask.to(out_dtype)
+
+    # Set masked regions to 1.0.
+    if mask.dtype == torch.bool:
+        mask = mask.to(out_dtype)
+    else:
+        mask = mask.to(out_dtype)
+        mask_region = mask > 0.5
+        mask[mask_region] = 1.0
+        mask[~mask_region] = 0.0
+
+    return mask

invokeai/backend/util/original_weights_storage.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import Dict, Iterator, Optional, Tuple
+
+import torch
+
+from invokeai.backend.util.devices import TorchDevice
+
+
+class OriginalWeightsStorage:
+    """A class for tracking the original weights of a model for patch/unpatch operations."""
+
+    def __init__(self, cached_weights: Optional[Dict[str, torch.Tensor]] = None):
+        # The original weights of the model.
+        self._weights: dict[str, torch.Tensor] = {}
+        # The keys of the weights that have been changed (via `save()`) during the lifetime of this instance.
+        self._changed_weights: set[str] = set()
+        if cached_weights:
+            self._weights.update(cached_weights)
+
+    def save(self, key: str, weight: torch.Tensor, copy: bool = True):
+        self._changed_weights.add(key)
+        if key in self._weights:
+            return
+
+        self._weights[key] = weight.detach().to(device=TorchDevice.CPU_DEVICE, copy=copy)
+
+    def get(self, key: str, copy: bool = False) -> Optional[torch.Tensor]:
+        weight = self._weights.get(key, None)
+        if weight is not None and copy:
+            weight = weight.clone()
+        return weight
+
+    def contains(self, key: str) -> bool:
+        return key in self._weights
+
+    def get_changed_weights(self) -> Iterator[Tuple[str, torch.Tensor]]:
+        for key in self._changed_weights:
+            yield key, self._weights[key]

invokeai/backend/util/prefix_logger_adapter.py

@@ -0,0 +1,12 @@
+import logging
+from typing import Any, MutableMapping
+
+
+# Issue with type hints related to LoggerAdapter: https://github.com/python/typeshed/issues/7855
+class PrefixedLoggerAdapter(logging.LoggerAdapter):  # type: ignore
+    def __init__(self, logger: logging.Logger, prefix: str):
+        super().__init__(logger, {})
+        self.prefix = prefix
+
+    def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> tuple[str, MutableMapping[str, Any]]:
+        return f"[{self.prefix}] {msg}", kwargs

invokeai/backend/util/silence_warnings.py

@@ -0,0 +1,36 @@
+import warnings
+from contextlib import ContextDecorator
+
+from diffusers.utils import logging as diffusers_logging
+from transformers import logging as transformers_logging
+
+
+# Inherit from ContextDecorator to allow using SilenceWarnings as both a context manager and a decorator.
+class SilenceWarnings(ContextDecorator):
+    """A context manager that disables warnings from transformers & diffusers modules while active.
+
+    As context manager:
+    ```
+    with SilenceWarnings():
+        # do something
+    ```
+
+    As decorator:
+    ```
+    @SilenceWarnings()
+    def some_function():
+        # do something
+    ```
+    """
+
+    def __enter__(self) -> None:
+        self._transformers_verbosity = transformers_logging.get_verbosity()
+        self._diffusers_verbosity = diffusers_logging.get_verbosity()
+        transformers_logging.set_verbosity_error()
+        diffusers_logging.set_verbosity_error()
+        warnings.simplefilter("ignore")
+
+    def __exit__(self, *args) -> None:
+        transformers_logging.set_verbosity(self._transformers_verbosity)
+        diffusers_logging.set_verbosity(self._diffusers_verbosity)
+        warnings.simplefilter("default")

invokeai/backend/util/test_utils.py

@@ -0,0 +1,64 @@
+import contextlib
+from pathlib import Path
+from typing import Optional, Union
+
+import pytest
+import torch
+
+from invokeai.app.services.model_manager import ModelManagerServiceBase
+from invokeai.app.services.model_records import UnknownModelException
+from invokeai.backend.model_manager.load.load_base import LoadedModel
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, SubModelType
+
+
+@pytest.fixture(scope="session")
+def torch_device():
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def install_and_load_model(
+    model_manager: ModelManagerServiceBase,
+    model_path_id_or_url: Union[str, Path],
+    model_name: str,
+    base_model: BaseModelType,
+    model_type: ModelType,
+    submodel_type: Optional[SubModelType] = None,
+) -> LoadedModel:
+    """Install a model if it is not already installed, then get the LoadedModel for that model.
+
+    This is intended as a utility function for tests.
+
+    Args:
+        mm2_model_manager (ModelManagerServiceBase): The model manager
+        model_path_id_or_url (Union[str, Path]): The path, HF ID, URL, etc. where the model can be installed from if it
+            is not already installed.
+        model_name (str): The model name, forwarded to ModelManager.get_model(...).
+        base_model (BaseModelType): The base model, forwarded to ModelManager.get_model(...).
+        model_type (ModelType): The model type, forwarded to ModelManager.get_model(...).
+        submodel_type (Optional[SubModelType]): The submodel type, forwarded to ModelManager.get_model(...).
+
+    Returns:
+        LoadedModelInfo
+    """
+    # If the requested model is already installed, return its LoadedModel
+    with contextlib.suppress(UnknownModelException):
+        # TODO: Replace with wrapper call
+        configs = model_manager.store.search_by_attr(
+            model_name=model_name, base_model=base_model, model_type=model_type
+        )
+        loaded_model: LoadedModel = model_manager.load.load_model(configs[0])
+        return loaded_model
+
+    # Install the requested model.
+    job = model_manager.install.heuristic_import(model_path_id_or_url)
+    model_manager.install.wait_for_job(job, timeout=10)
+    assert job.complete
+
+    try:
+        loaded_model = model_manager.load.load_model(job.config_out)
+        return loaded_model
+    except UnknownModelException as e:
+        raise Exception(
+            "Failed to get model info after installing it. There could be a mismatch between the requested model and"
+            f" the installation id ('{model_path_id_or_url}'). Error: {e}"
+        )

invokeai/backend/util/util.py

@@ -0,0 +1,74 @@
+import base64
+import io
+import os
+import re
+import unicodedata
+from pathlib import Path
+
+from PIL import Image
+
+
+def slugify(value: str, allow_unicode: bool = False) -> str:
+    """
+    Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated
+    dashes to single dashes. Remove characters that aren't alphanumerics,
+    underscores, or hyphens. Replace slashes with underscores.
+    Convert to lowercase. Also strip leading and
+    trailing whitespace, dashes, and underscores.
+
+    Adapted from Django: https://github.com/django/django/blob/main/django/utils/text.py
+    """
+    value = str(value)
+    if allow_unicode:
+        value = unicodedata.normalize("NFKC", value)
+    else:
+        value = unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii")
+    value = re.sub(r"[/]", "_", value.lower())
+    value = re.sub(r"[^.\w\s-]", "", value.lower())
+    return re.sub(r"[-\s]+", "-", value).strip("-_")
+
+
+def safe_filename(directory: Path, value: str) -> str:
+    """Make a string safe to use as a filename."""
+    escaped_string = slugify(value)
+    max_name_length = os.pathconf(directory, "PC_NAME_MAX") if hasattr(os, "pathconf") else 256
+    return escaped_string[len(escaped_string) - max_name_length :]
+
+
+def directory_size(directory: Path) -> int:
+    """
+    Return the aggregate size of all files in a directory (bytes).
+    """
+    sum = 0
+    for root, _, files in os.walk(directory):
+        for f in files:
+            sum += Path(root, f).stat().st_size
+    return sum
+
+
+def image_to_dataURL(image: Image.Image, image_format: str = "PNG") -> str:
+    """
+    Converts an image into a base64 image dataURL.
+    """
+    buffered = io.BytesIO()
+    image.save(buffered, format=image_format)
+    mime_type = Image.MIME.get(image_format.upper(), "image/" + image_format.lower())
+    image_base64 = f"data:{mime_type};base64," + base64.b64encode(buffered.getvalue()).decode("UTF-8")
+    return image_base64
+
+
+class Chdir(object):
+    """Context manager to chdir to desired directory and change back after context exits:
+    Args:
+        path (Path): The path to the cwd
+    """
+
+    def __init__(self, path: Path):
+        self.path = path
+        self.original = Path().absolute()
+
+    def __enter__(self):
+        os.chdir(self.path)
+
+    def __exit__(self, *args):
+        os.chdir(self.original)

invokeai/backend/util/vae_working_memory.py

@@ -0,0 +1,199 @@
+from typing import Literal
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage
+from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
+from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny
+
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
+from invokeai.backend.flux.modules.autoencoder import AutoEncoder
+
+
+def estimate_vae_working_memory_sd15_sdxl(
+    operation: Literal["encode", "decode"],
+    image_tensor: torch.Tensor,
+    vae: AutoencoderKL | AutoencoderTiny,
+    tile_size: int | None,
+    fp32: bool,
+) -> int:
+    """Estimate the working memory required to encode or decode the given tensor."""
+    # It was found experimentally that the peak working memory scales linearly with the number of pixels and the
+    # element size (precision). This estimate is accurate for both SD1 and SDXL.
+    element_size = 4 if fp32 else 2
+
+    # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414
+    # Encoding uses ~45% the working memory as decoding.
+    scaling_constant = 2200 if operation == "decode" else 1100
+
+    latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
+
+    if tile_size is not None:
+        if tile_size == 0:
+            tile_size = vae.tile_sample_min_size
+            assert isinstance(tile_size, int)
+        h = tile_size
+        w = tile_size
+        working_memory = h * w * element_size * scaling_constant
+
+        # We add 25% to the working memory estimate when tiling is enabled to account for factors like tile overlap
+        # and number of tiles. We could make this more precise in the future, but this should be good enough for
+        # most use cases.
+        working_memory = working_memory * 1.25
+    else:
+        h = latent_scale_factor_for_operation * image_tensor.shape[-2]
+        w = latent_scale_factor_for_operation * image_tensor.shape[-1]
+        working_memory = h * w * element_size * scaling_constant
+
+    if fp32:
+        # If we are running in FP32, then we should account for the likely increase in model size (~250MB).
+        working_memory += 250 * 2**20
+
+    return int(working_memory)
+
+
+def estimate_vae_working_memory_cogview4(
+    operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL
+) -> int:
+    """Estimate the working memory required by the invocation in bytes."""
+    latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
+
+    h = latent_scale_factor_for_operation * image_tensor.shape[-2]
+    w = latent_scale_factor_for_operation * image_tensor.shape[-1]
+    element_size = next(vae.parameters()).element_size()
+
+    # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414
+    # Encoding uses ~45% the working memory as decoding.
+    scaling_constant = 2200 if operation == "decode" else 1100
+    working_memory = h * w * element_size * scaling_constant
+
+    print(f"estimate_vae_working_memory_cogview4: {int(working_memory)}")
+
+    return int(working_memory)
+
+
+def estimate_vae_working_memory_flux(
+    operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoEncoder
+) -> int:
+    """Estimate the working memory required by the invocation in bytes."""
+
+    latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
+
+    out_h = latent_scale_factor_for_operation * image_tensor.shape[-2]
+    out_w = latent_scale_factor_for_operation * image_tensor.shape[-1]
+    element_size = next(vae.parameters()).element_size()
+
+    # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414
+    # Encoding uses ~45% the working memory as decoding.
+    scaling_constant = 2200 if operation == "decode" else 1100
+
+    working_memory = out_h * out_w * element_size * scaling_constant
+
+    print(f"estimate_vae_working_memory_flux: {int(working_memory)}")
+
+    return int(working_memory)
+
+
+def estimate_vae_working_memory_anima(
+    operation: Literal["encode", "decode"],
+    image_tensor: torch.Tensor,
+    vae: AutoencoderKLWan,
+    tile_size: int | None,
+) -> int:
+    """Estimate the working memory required to encode or decode with the Wan 2.1 VAE (Anima).
+
+    The Wan VAE uses 3D convolutions and needs noticeably more working memory per output
+    pixel than the 2D VAEs estimated above. Calibrated empirically on a 1024x1024 fp16
+    decode: peak reserved memory was ~5.95GB for a full decode and ~1.73GB with 512px
+    tiles (384px stride), i.e. ~2900 bytes per output pixel per element byte. Encoding
+    follows the house ratio of ~50% of decode.
+    """
+    element_size = next(vae.parameters()).element_size()
+    scaling_constant = 2900 if operation == "decode" else 1450
+
+    if tile_size is not None:
+        h = tile_size
+        w = tile_size
+        # Add 25% to account for tile overlap.
+        working_memory = h * w * element_size * scaling_constant * 1.25
+    else:
+        latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
+        h = latent_scale_factor_for_operation * image_tensor.shape[-2]
+        w = latent_scale_factor_for_operation * image_tensor.shape[-1]
+        working_memory = h * w * element_size * scaling_constant
+
+    return int(working_memory)
+
+
+def estimate_vae_working_memory_qwen_image(
+    operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKLQwenImage
+) -> int:
+    """Estimate the working memory required by the invocation in bytes.
+
+    The Qwen Image VAE is a video-style autoencoder that operates on 5D tensors of shape
+    (B, C, num_frames, H, W). Tiling is not used, so peak working memory scales with the full
+    spatial output. The two trailing dimensions are the spatial H/W in latent space (decode) or
+    pixel space (encode), matching the convention used by the other estimators here.
+    """
+    latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
+
+    h = latent_scale_factor_for_operation * image_tensor.shape[-2]
+    w = latent_scale_factor_for_operation * image_tensor.shape[-1]
+    element_size = next(vae.parameters()).element_size()
+
+    # The Qwen Image VAE is much heavier than the SD/SDXL VAE and needs correspondingly larger
+    # constants. These were calibrated by measuring peak *reserved* memory growth (not just allocated
+    # -- reserved is what the cache's `free >= estimate` check compares against) across a resolution
+    # grid in fp16, on both an AMD W7900 (ROCm) and an NVIDIA card (CUDA). See
+    # scripts/calibrate_qwen_vae_working_memory.py.
+    #
+    # Implied constant = reserved_bytes / (h * w * element_size). Per-point maxima (fp16):
+    #              512^2  768^2  1024^2  1536^2  1792^2  2048^2    -> ship (max observed + ~8% headroom)
+    #   ROCm decode  5132   4596   4570    3273    3735    4813    -> 5500
+    #   ROCm encode  5864   5858   5858    3532    4364   (OOM)    -> 6300
+    #   CUDA decode  2660   2519   2690    2671    2281   (OOM)    -> 2900
+    #   CUDA encode  1456   1451   1458    1456    1455    1455    -> 1600
+    #
+    # Why this branches on backend (the only estimator here that does):
+    #  - The Qwen VAE is attention-heavy. With Flash/efficient attention (CUDA) the attention memory
+    #    is O(area) and the curve is flat/linear; the ROCm build falls back to math attention, which
+    #    is O(area^2), so ROCm reserves ~2x (decode) to ~4x (encode) more and goes super-linear above
+    #    ~1792^2. The two backends differ far more than any headroom, so a single constant would
+    #    either under-estimate on ROCm (OOM) or massively over-budget on CUDA (needless eviction).
+    #  - "Encoding is half of decoding" (as the sibling estimators assume) is only true on CUDA. On
+    #    ROCm encode reserves >= decode, so the ROCm encode constant is sized accordingly -- this is
+    #    the path Qwen Image Edit exercises.
+    #  - On ROCm the linear model under-estimates for decodes well above 2048^2, but those OOM on a
+    #    48GB card regardless; on CUDA the curve stays linear so no extra term is needed.
+    is_rocm = torch.version.hip is not None
+    if operation == "decode":
+        scaling_constant = 5500 if is_rocm else 2900
+    else:  # encode
+        scaling_constant = 6300 if is_rocm else 1600
+
+    working_memory = h * w * element_size * scaling_constant
+
+    return int(working_memory)
+
+
+def estimate_vae_working_memory_sd3(
+    operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL
+) -> int:
+    """Estimate the working memory required by the invocation in bytes."""
+    # Encode operations use approximately 50% of the memory required for decode operations
+
+    latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
+
+    h = latent_scale_factor_for_operation * image_tensor.shape[-2]
+    w = latent_scale_factor_for_operation * image_tensor.shape[-1]
+    element_size = next(vae.parameters()).element_size()
+
+    # This constant is determined experimentally and takes into consideration both allocated and reserved memory. See #8414
+    # Encoding uses ~45% the working memory as decoding.
+    scaling_constant = 2200 if operation == "decode" else 1100
+
+    working_memory = h * w * element_size * scaling_constant
+
+    print(f"estimate_vae_working_memory_sd3: {int(working_memory)}")
+
+    return int(working_memory)