chore: prep for v5.9.1

Before FLUX Fill was merged, we didn't do any checks for the model variant. We always returned "normal".

To determine if a model is a FLUX Fill model, we need to check the state dict for a specific key. Initially, this logic was too strict and rejected quantized FLUX models. This issue was resolved, but it turns out there is another failure mode - some fine-tunes use a different key.

This change further reduces the strictness, handling the alternate key and also falling back to "normal" if we don't see either key. This effectively restores the previous probing behaviour for all FLUX models.

Closes #7856
Closes #7859

.github/workflows/python-checks.yml

@@ -34,6 +34,9 @@ on:
 
 jobs:
   python-checks:
+    env:
+      # uv requires a venv by default - but for this, we can simply use the system python
+      UV_SYSTEM_PYTHON: 1
     runs-on: ubuntu-latest
     timeout-minutes: 5 # expected run time: <1 min
     steps:
@@ -57,25 +60,19 @@ jobs:
               - '!invokeai/frontend/web/**'
               - 'tests/**'
 
-      - name: setup python
+      - name: setup uv
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
-        uses: actions/setup-python@v5
+        uses: astral-sh/setup-uv@v5
         with:
-          python-version: '3.10'
-          cache: pip
-          cache-dependency-path: pyproject.toml
-
-      - name: install ruff
-        if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
-        run: pip install ruff==0.9.9
-        shell: bash
+          version: '0.6.10'
+          enable-cache: true
 
       - name: ruff check
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
-        run: ruff check --output-format=github .
+        run: uv tool run ruff@0.11.2 check --output-format=github .
         shell: bash
 
       - name: ruff format
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
-        run: ruff format --check .
+        run: uv tool run ruff@0.11.2 format --check .
         shell: bash

.github/workflows/python-tests.yml

@@ -39,24 +39,15 @@ jobs:
     strategy:
       matrix:
         python-version:
-          - '3.10'
           - '3.11'
+          - '3.12'
         platform:
-          - linux-cuda-11_7
-          - linux-rocm-5_2
           - linux-cpu
           - macos-default
           - windows-cpu
         include:
-          - platform: linux-cuda-11_7
-            os: ubuntu-22.04
-            github-env: $GITHUB_ENV
-          - platform: linux-rocm-5_2
-            os: ubuntu-22.04
-            extra-index-url: 'https://download.pytorch.org/whl/rocm5.2'
-            github-env: $GITHUB_ENV
           - platform: linux-cpu
-            os: ubuntu-22.04
+            os: ubuntu-24.04
             extra-index-url: 'https://download.pytorch.org/whl/cpu'
             github-env: $GITHUB_ENV
           - platform: macos-default
@@ -70,6 +61,8 @@ jobs:
     timeout-minutes: 15 # expected run time: 2-6 min, depending on platform
     env:
       PIP_USE_PEP517: '1'
+      UV_SYSTEM_PYTHON: 1
+
     steps:
       - name: checkout
         # https://github.com/nschloe/action-cached-lfs-checkout
@@ -92,20 +85,25 @@ jobs:
               - '!invokeai/frontend/web/**'
               - 'tests/**'
 
+      - name: setup uv
+        if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: '0.6.10'
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+
       - name: setup python
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
-          cache: pip
-          cache-dependency-path: pyproject.toml
 
       - name: install dependencies
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}
         env:
-          PIP_EXTRA_INDEX_URL: ${{ matrix.extra-index-url }}
-        run: >
-          pip3 install --editable=".[test]"
+          UV_INDEX: ${{ matrix.extra-index-url }}
+        run: uv pip install --editable ".[test]"
 
       - name: run pytest
         if: ${{ steps.changed-files.outputs.python_any_changed == 'true' || inputs.always_run == true }}

.github/workflows/typegen-checks.yml

@@ -54,17 +54,25 @@ jobs:
               - 'pyproject.toml'
               - 'invokeai/**'
 
+      - name: setup uv
+        if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: '0.6.10'
+          enable-cache: true
+          python-version: '3.11'
+
       - name: setup python
         if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
         uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
-          cache: pip
-          cache-dependency-path: pyproject.toml
+          python-version: '3.11'
 
-      - name: install python dependencies
+      - name: install dependencies
         if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        run: pip3 install --use-pep517 --editable="."
+        env:
+          UV_INDEX: ${{ matrix.extra-index-url }}
+        run: uv pip install --editable .
 
       - name: install frontend dependencies
         if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
@@ -77,7 +85,7 @@ jobs:
 
       - name: generate schema
         if: ${{ steps.changed-files.outputs.src_any_changed == 'true' || inputs.always_run == true }}
-        run: make frontend-typegen
+        run: cd invokeai/frontend/web && uv run ../../../scripts/generate_openapi_schema.py | pnpm typegen
         shell: bash
 
       - name: compare files

invokeai/app/api/routers/images.py

@@ -96,6 +96,22 @@ async def upload_image(
         raise HTTPException(status_code=500, detail="Failed to create image")
 
 
+class ImageUploadEntry(BaseModel):
+    image_dto: ImageDTO = Body(description="The image DTO")
+    presigned_url: str = Body(description="The URL to get the presigned URL for the image upload")
+
+
+@images_router.post("/", operation_id="create_image_upload_entry")
+async def create_image_upload_entry(
+    width: int = Body(description="The width of the image"),
+    height: int = Body(description="The height of the image"),
+    board_id: Optional[str] = Body(default=None, description="The board to add this image to, if any"),
+) -> ImageUploadEntry:
+    """Uploads an image from a URL, not implemented"""
+
+    raise HTTPException(status_code=501, detail="Not implemented")
+
+
 @images_router.delete("/i/{image_name}", operation_id="delete_image")
 async def delete_image(
     image_name: str = Path(description="The name of the image to delete"),

invokeai/app/invocations/image.py

@@ -1089,12 +1089,13 @@ class CanvasV2MaskAndCropInvocation(BaseInvocation, WithMetadata, WithBoard):
 
 
 @invocation(
-    "expand_mask_with_fade", title="Expand Mask with Fade", tags=["image", "mask"], category="image", version="1.0.0"
+    "expand_mask_with_fade", title="Expand Mask with Fade", tags=["image", "mask"], category="image", version="1.0.1"
 )
 class ExpandMaskWithFadeInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Expands a mask with a fade effect. The mask uses black to indicate areas to keep from the generated image and white for areas to discard.
     The mask is thresholded to create a binary mask, and then a distance transform is applied to create a fade effect.
     The fade size is specified in pixels, and the mask is expanded by that amount. The result is a mask with a smooth transition from black to white.
+    If the fade size is 0, the mask is returned as-is.
     """
 
     mask: ImageField = InputField(description="The mask to expand")
@@ -1104,6 +1105,11 @@ class ExpandMaskWithFadeInvocation(BaseInvocation, WithMetadata, WithBoard):
     def invoke(self, context: InvocationContext) -> ImageOutput:
         pil_mask = context.images.get_pil(self.mask.image_name, mode="L")
 
+        if self.fade_size_px == 0:
+            # If the fade size is 0, just return the mask as-is.
+            image_dto = context.images.save(image=pil_mask, image_category=ImageCategory.MASK)
+            return ImageOutput.build(image_dto)
+
         np_mask = numpy.array(pil_mask)
 
         # Threshold the mask to create a binary mask - 0 for black, 255 for white
@@ -1141,8 +1147,21 @@ class ExpandMaskWithFadeInvocation(BaseInvocation, WithMetadata, WithBoard):
         coeffs = numpy.polyfit(x_control, y_control, 3)
         poly = numpy.poly1d(coeffs)
 
-        # Evaluate and clip the smooth mapping
-        feather = numpy.clip(poly(d_norm), 0, 1)
+        # Evaluate the polynomial
+        feather = poly(d_norm)
+
+        # The polynomial fit isn't perfect. Points beyond the fade distance are likely to be slightly less than 1.0,
+        # even though the control points indicate that they should be exactly 1.0. This is due to the nature of the
+        # polynomial fit, which is a best approximation of the control points but not an exact match.
+
+        # When this occurs, the area outside the mask and fade-out will not be 100% transparent. For example, it may
+        # have an alpha value of 1 instead of 0. So we must force pixels at or beyond the fade distance to exactly 1.0.
+
+        # Force pixels at or beyond the fade distance to exactly 1.0
+        feather = numpy.where(d_norm >= 1.0, 1.0, feather)
+
+        # Clip any other values to ensure they're in the valid range [0,1]
+        feather = numpy.clip(feather, 0, 1)
 
         # Build final image.
         np_result = numpy.where(black_mask == 1, 0, (feather * 255).astype(numpy.uint8))

invokeai/backend/flux/flux_state_dict_utils.py

@@ -0,0 +1,23 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from invokeai.backend.model_manager.legacy_probe import CkptType
+
+
+def get_flux_in_channels_from_state_dict(state_dict: "CkptType") -> int | None:
+    """Gets the in channels from the state dict."""
+
+    # "Standard" FLUX models use "img_in.weight", but some community fine tunes use
+    # "model.diffusion_model.img_in.weight". Known models that use the latter key:
+    # - https://civitai.com/models/885098?modelVersionId=990775
+    # - https://civitai.com/models/1018060?modelVersionId=1596255
+    # - https://civitai.com/models/978314/ultrareal-fine-tune?modelVersionId=1413133
+
+    keys = {"img_in.weight", "model.diffusion_model.img_in.weight"}
+
+    for key in keys:
+        val = state_dict.get(key)
+        if val is not None:
+            return val.shape[1]
+
+    return None

invokeai/backend/model_manager/config.py

@@ -67,6 +67,11 @@ class InvalidModelConfigException(Exception):
 DEFAULTS_PRECISION = Literal["fp16", "fp32"]
 
 
+class FSLayout(Enum):
+    FILE = "file"
+    DIRECTORY = "directory"
+
+
 class SubmodelDefinition(BaseModel):
     path_or_prefix: str
     model_type: ModelType
@@ -102,29 +107,31 @@ class ModelOnDisk:
 
     def __init__(self, path: Path, hash_algo: HASHING_ALGORITHMS = "blake3_single"):
         self.path = path
-        self.format_type = ModelFormat.Diffusers if path.is_dir() else ModelFormat.Checkpoint
+        # TODO: Revisit checkpoint vs diffusers terminology
+        self.layout = FSLayout.DIRECTORY if path.is_dir() else FSLayout.FILE
         if self.path.suffix in {".safetensors", ".bin", ".pt", ".ckpt"}:
             self.name = path.stem
         else:
             self.name = path.name
         self.hash_algo = hash_algo
+        self._state_dict_cache = {}
 
-    def hash(self):
+    def hash(self) -> str:
         return ModelHash(algorithm=self.hash_algo).hash(self.path)
 
-    def size(self):
-        if self.format_type == ModelFormat.Checkpoint:
+    def size(self) -> int:
+        if self.layout == FSLayout.FILE:
             return self.path.stat().st_size
         return sum(file.stat().st_size for file in self.path.rglob("*"))
 
-    def component_paths(self):
-        if self.format_type == ModelFormat.Checkpoint:
+    def component_paths(self) -> set[Path]:
+        if self.layout == FSLayout.FILE:
             return {self.path}
         extensions = {".safetensors", ".pt", ".pth", ".ckpt", ".bin", ".gguf"}
         return {f for f in self.path.rglob("*") if f.suffix in extensions}
 
-    def repo_variant(self):
-        if self.format_type == ModelFormat.Checkpoint:
+    def repo_variant(self) -> Optional[ModelRepoVariant]:
+        if self.layout == FSLayout.FILE:
             return None
 
         weight_files = list(self.path.glob("**/*.safetensors"))
@@ -140,14 +147,30 @@ class ModelOnDisk:
                 return ModelRepoVariant.ONNX
         return ModelRepoVariant.Default
 
-    @staticmethod
-    def load_state_dict(path: Path):
+    def load_state_dict(self, path: Optional[Path] = None) -> Dict[str | int, Any]:
+        if path in self._state_dict_cache:
+            return self._state_dict_cache[path]
+
+        if not path:
+            components = list(self.component_paths())
+            match components:
+                case []:
+                    raise ValueError("No weight files found for this model")
+                case [p]:
+                    path = p
+                case ps if len(ps) >= 2:
+                    raise ValueError(
+                        f"Multiple weight files found for this model: {ps}. "
+                        f"Please specify the intended file using the 'path' argument"
+                    )
+
         with SilenceWarnings():
             if path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin")):
                 scan_result = scan_file_path(path)
                 if scan_result.infected_files != 0 or scan_result.scan_err:
                     raise RuntimeError(f"The model {path.stem} is potentially infected by malware. Aborting import.")
                 checkpoint = torch.load(path, map_location="cpu")
+                assert isinstance(checkpoint, dict)
             elif path.suffix.endswith(".gguf"):
                 checkpoint = gguf_sd_loader(path, compute_dtype=torch.float32)
             elif path.suffix.endswith(".safetensors"):
@@ -156,6 +179,7 @@ class ModelOnDisk:
                 raise ValueError(f"Unrecognized model extension: {path.suffix}")
 
         state_dict = checkpoint.get("state_dict", checkpoint)
+        self._state_dict_cache[path] = state_dict
         return state_dict
 
 
@@ -238,11 +262,13 @@ class ModelConfigBase(ABC, BaseModel):
 
         for config_cls in sorted_by_match_speed:
             try:
-                return config_cls.from_model_on_disk(mod, **overrides)
-            except InvalidModelConfigException:
-                logger.debug(f"ModelConfig '{config_cls.__name__}' failed to parse '{mod.path}', trying next config")
+                if not config_cls.matches(mod):
+                    continue
             except Exception as e:
-                logger.error(f"Unexpected exception while parsing '{config_cls.__name__}': {e}, trying next config")
+                logger.warning(f"Unexpected exception while matching {mod.name} to '{config_cls.__name__}': {e}")
+                continue
+            else:
+                return config_cls.from_model_on_disk(mod, **overrides)
 
         raise InvalidModelConfigException("No valid config found")
 
@@ -285,9 +311,6 @@ class ModelConfigBase(ABC, BaseModel):
     @classmethod
     def from_model_on_disk(cls, mod: ModelOnDisk, **overrides):
         """Creates an instance of this config or raises InvalidModelConfigException."""
-        if not cls.matches(mod):
-            raise InvalidModelConfigException(f"Path {mod.path} does not match {cls.__name__} format")
-
         fields = cls.parse(mod)
         cls.cast_overrides(overrides)
         fields.update(overrides)
@@ -563,7 +586,7 @@ class LlavaOnevisionConfig(DiffusersConfigBase, ModelConfigBase):
 
     @classmethod
     def matches(cls, mod: ModelOnDisk) -> bool:
-        if mod.format_type == ModelFormat.Checkpoint:
+        if mod.layout == FSLayout.FILE:
             return False
 
         config_path = mod.path / "config.json"

invokeai/backend/model_manager/legacy_probe.py

@@ -14,6 +14,7 @@ from invokeai.backend.flux.controlnet.state_dict_utils import (
     is_state_dict_instantx_controlnet,
     is_state_dict_xlabs_controlnet,
 )
+from invokeai.backend.flux.flux_state_dict_utils import get_flux_in_channels_from_state_dict
 from invokeai.backend.flux.ip_adapter.state_dict_utils import is_state_dict_xlabs_ip_adapter
 from invokeai.backend.flux.redux.flux_redux_state_dict_utils import is_state_dict_likely_flux_redux
 from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
@@ -564,7 +565,14 @@ class CheckpointProbeBase(ProbeBase):
         state_dict = self.checkpoint.get("state_dict") or self.checkpoint
 
         if base_type == BaseModelType.Flux:
-            in_channels = state_dict["img_in.weight"].shape[1]
+            in_channels = get_flux_in_channels_from_state_dict(state_dict)
+
+            if in_channels is None:
+                # If we cannot find the in_channels, we assume that this is a normal variant. Log a warning.
+                logger.warning(
+                    f"{self.model_path} does not have img_in.weight or model.diffusion_model.img_in.weight key. Assuming normal variant."
+                )
+                return ModelVariantType.Normal
 
             # FLUX Model variant types are distinguished by input channels:
             # - Unquantized Dev and Schnell have in_channels=64

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageUploaded.ts

@@ -1,6 +1,8 @@
+import { isAnyOf } from '@reduxjs/toolkit';
 import { logger } from 'app/logging/logger';
 import type { AppStartListening } from 'app/store/middleware/listenerMiddleware';
 import type { RootState } from 'app/store/store';
+import { imageUploadedClientSide } from 'features/gallery/store/actions';
 import { selectListBoardsQueryArgs } from 'features/gallery/store/gallerySelectors';
 import { boardIdSelected, galleryViewChanged } from 'features/gallery/store/gallerySlice';
 import { toast } from 'features/toast/toast';
@@ -8,7 +10,8 @@ import { t } from 'i18next';
 import { omit } from 'lodash-es';
 import { boardsApi } from 'services/api/endpoints/boards';
 import { imagesApi } from 'services/api/endpoints/images';
-
+import type { ImageDTO } from 'services/api/types';
+import { getCategories, getListImagesUrl } from 'services/api/util';
 const log = logger('gallery');
 
 /**
@@ -34,19 +37,56 @@ let lastUploadedToastTimeout: number | null = null;
 
 export const addImageUploadedFulfilledListener = (startAppListening: AppStartListening) => {
   startAppListening({
-    matcher: imagesApi.endpoints.uploadImage.matchFulfilled,
+    matcher: isAnyOf(imagesApi.endpoints.uploadImage.matchFulfilled, imageUploadedClientSide),
     effect: (action, { dispatch, getState }) => {
-      const imageDTO = action.payload;
+      let imageDTO: ImageDTO;
+      let silent;
+      let isFirstUploadOfBatch = true;
+
+      if (imageUploadedClientSide.match(action)) {
+        imageDTO = action.payload.imageDTO;
+        silent = action.payload.silent;
+        isFirstUploadOfBatch = action.payload.isFirstUploadOfBatch;
+      } else if (imagesApi.endpoints.uploadImage.matchFulfilled(action)) {
+        imageDTO = action.payload;
+        silent = action.meta.arg.originalArgs.silent;
+        isFirstUploadOfBatch = action.meta.arg.originalArgs.isFirstUploadOfBatch ?? true;
+      } else {
+        return;
+      }
+
+      if (silent || imageDTO.is_intermediate) {
+        // If the image is silent or intermediate, we don't want to show a toast
+        return;
+      }
+
+      if (imageUploadedClientSide.match(action)) {
+        const categories = getCategories(imageDTO);
+        const boardId = imageDTO.board_id ?? 'none';
+        dispatch(
+          imagesApi.util.invalidateTags([
+            {
+              type: 'ImageList',
+              id: getListImagesUrl({
+                board_id: boardId,
+                categories,
+              }),
+            },
+            {
+              type: 'Board',
+              id: boardId,
+            },
+            {
+              type: 'BoardImagesTotal',
+              id: boardId,
+            },
+          ])
+        );
+      }
       const state = getState();
 
       log.debug({ imageDTO }, 'Image uploaded');
 
-      if (action.meta.arg.originalArgs.silent || imageDTO.is_intermediate) {
-        // When a "silent" upload is requested, or the image is intermediate, we can skip all post-upload actions,
-        // like toasts and switching the gallery view
-        return;
-      }
-
       const boardId = imageDTO.board_id ?? 'none';
 
       const DEFAULT_UPLOADED_TOAST = {
@@ -80,7 +120,7 @@ export const addImageUploadedFulfilledListener = (startAppListening: AppStartLis
        *
        * Default to true to not require _all_ image upload handlers to set this value
        */
-      const isFirstUploadOfBatch = action.meta.arg.originalArgs.isFirstUploadOfBatch ?? true;
+
       if (isFirstUploadOfBatch) {
         dispatch(boardIdSelected({ boardId }));
         dispatch(galleryViewChanged('assets'));

invokeai/frontend/web/src/app/types/invokeai.ts

@@ -73,6 +73,7 @@ export type AppConfig = {
   maxUpscaleDimension?: number;
   allowPrivateBoards: boolean;
   allowPrivateStylePresets: boolean;
+  allowClientSideUpload: boolean;
   disabledTabs: TabName[];
   disabledFeatures: AppFeature[];
   disabledSDFeatures: SDFeature[];
@@ -81,7 +82,6 @@ export type AppConfig = {
   metadataFetchDebounce?: number;
   workflowFetchDebounce?: number;
   isLocal?: boolean;
-  maxImageUploadCount?: number;
   sd: {
     defaultModel?: string;
     disabledControlNetModels: string[];

invokeai/frontend/web/src/common/hooks/useClientSideUpload.ts

@@ -0,0 +1,105 @@
+import { useStore } from '@nanostores/react';
+import { $authToken } from 'app/store/nanostores/authToken';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { imageUploadedClientSide } from 'features/gallery/store/actions';
+import { selectAutoAddBoardId } from 'features/gallery/store/gallerySelectors';
+import { useCallback } from 'react';
+import { useCreateImageUploadEntryMutation } from 'services/api/endpoints/images';
+import type { ImageDTO } from 'services/api/types';
+export const useClientSideUpload = () => {
+  const dispatch = useAppDispatch();
+  const autoAddBoardId = useAppSelector(selectAutoAddBoardId);
+  const authToken = useStore($authToken);
+  const [createImageUploadEntry] = useCreateImageUploadEntryMutation();
+
+  const clientSideUpload = useCallback(
+    async (file: File, i: number): Promise<ImageDTO> => {
+      const image = new Image();
+      const objectURL = URL.createObjectURL(file);
+      image.src = objectURL;
+      let width = 0;
+      let height = 0;
+      let thumbnail: Blob | undefined;
+
+      await new Promise<void>((resolve) => {
+        image.onload = () => {
+          width = image.naturalWidth;
+          height = image.naturalHeight;
+
+          // Calculate thumbnail dimensions maintaining aspect ratio
+          let thumbWidth = width;
+          let thumbHeight = height;
+          if (width > height && width > 256) {
+            thumbWidth = 256;
+            thumbHeight = Math.round((height * 256) / width);
+          } else if (height > 256) {
+            thumbHeight = 256;
+            thumbWidth = Math.round((width * 256) / height);
+          }
+
+          const canvas = document.createElement('canvas');
+          canvas.width = thumbWidth;
+          canvas.height = thumbHeight;
+          const ctx = canvas.getContext('2d');
+          ctx?.drawImage(image, 0, 0, thumbWidth, thumbHeight);
+
+          canvas.toBlob(
+            (blob) => {
+              if (blob) {
+                thumbnail = blob;
+                // Clean up resources
+                URL.revokeObjectURL(objectURL);
+                image.src = ''; // Clear image source
+                image.remove(); // Remove the image element
+                canvas.width = 0; // Clear canvas
+                canvas.height = 0;
+                resolve();
+              }
+            },
+            'image/webp',
+            0.8
+          );
+        };
+
+        // Handle load errors
+        image.onerror = () => {
+          URL.revokeObjectURL(objectURL);
+          image.remove();
+          resolve();
+        };
+      });
+      const { presigned_url, image_dto } = await createImageUploadEntry({
+        width,
+        height,
+        board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
+      }).unwrap();
+
+      await fetch(`${presigned_url}/?type=full`, {
+        method: 'PUT',
+        body: file,
+        ...(authToken && {
+          headers: {
+            Authorization: `Bearer ${authToken}`,
+          },
+        }),
+      });
+
+      await fetch(`${presigned_url}/?type=thumbnail`, {
+        method: 'PUT',
+        body: thumbnail,
+        ...(authToken && {
+          headers: {
+            Authorization: `Bearer ${authToken}`,
+          },
+        }),
+      });
+
+      dispatch(imageUploadedClientSide({ imageDTO: image_dto, silent: false, isFirstUploadOfBatch: i === 0 }));
+
+      return image_dto;
+    },
+    [autoAddBoardId, authToken, createImageUploadEntry, dispatch]
+  );
+
+  return clientSideUpload;
+};

invokeai/frontend/web/src/common/hooks/useImageUploadButton.tsx

@@ -3,7 +3,7 @@ import { IconButton } from '@invoke-ai/ui-library';
 import { logger } from 'app/logging/logger';
 import { useAppSelector } from 'app/store/storeHooks';
 import { selectAutoAddBoardId } from 'features/gallery/store/gallerySelectors';
-import { selectMaxImageUploadCount } from 'features/system/store/configSlice';
+import { selectIsClientSideUploadEnabled } from 'features/system/store/configSlice';
 import { toast } from 'features/toast/toast';
 import { useCallback } from 'react';
 import type { FileRejection } from 'react-dropzone';
@@ -15,6 +15,7 @@ import type { ImageDTO } from 'services/api/types';
 import { assert } from 'tsafe';
 import type { SetOptional } from 'type-fest';
 
+import { useClientSideUpload } from './useClientSideUpload';
 type UseImageUploadButtonArgs =
   | {
       isDisabled?: boolean;
@@ -50,8 +51,9 @@ const log = logger('gallery');
  */
 export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: UseImageUploadButtonArgs) => {
   const autoAddBoardId = useAppSelector(selectAutoAddBoardId);
+  const isClientSideUploadEnabled = useAppSelector(selectIsClientSideUploadEnabled);
   const [uploadImage, request] = useUploadImageMutation();
-  const maxImageUploadCount = useAppSelector(selectMaxImageUploadCount);
+  const clientSideUpload = useClientSideUpload();
   const { t } = useTranslation();
 
   const onDropAccepted = useCallback(
@@ -79,22 +81,27 @@ export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: Us
           onUpload(imageDTO);
         }
       } else {
-        const imageDTOs = await uploadImages(
-          files.map((file, i) => ({
-            file,
-            image_category: 'user',
-            is_intermediate: false,
-            board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
-            silent: false,
-            isFirstUploadOfBatch: i === 0,
-          }))
-        );
+        let imageDTOs: ImageDTO[] = [];
+        if (isClientSideUploadEnabled) {
+          imageDTOs = await Promise.all(files.map((file, i) => clientSideUpload(file, i)));
+        } else {
+          imageDTOs = await uploadImages(
+            files.map((file, i) => ({
+              file,
+              image_category: 'user',
+              is_intermediate: false,
+              board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
+              silent: false,
+              isFirstUploadOfBatch: i === 0,
+            }))
+          );
+        }
         if (onUpload) {
           onUpload(imageDTOs);
         }
       }
     },
-    [allowMultiple, autoAddBoardId, onUpload, uploadImage]
+    [allowMultiple, autoAddBoardId, onUpload, uploadImage, isClientSideUploadEnabled, clientSideUpload]
   );
 
   const onDropRejected = useCallback(
@@ -105,10 +112,7 @@ export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: Us
           file: rejection.file.path,
         }));
         log.error({ errors }, 'Invalid upload');
-        const description =
-          maxImageUploadCount === undefined
-            ? t('toast.uploadFailedInvalidUploadDesc')
-            : t('toast.uploadFailedInvalidUploadDesc_withCount', { count: maxImageUploadCount });
+        const description = t('toast.uploadFailedInvalidUploadDesc');
 
         toast({
           id: 'UPLOAD_FAILED',
@@ -120,7 +124,7 @@ export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: Us
         return;
       }
     },
-    [maxImageUploadCount, t]
+    [t]
   );
 
   const {
@@ -137,8 +141,7 @@ export const useImageUploadButton = ({ onUpload, isDisabled, allowMultiple }: Us
     onDropRejected,
     disabled: isDisabled,
     noDrag: true,
-    multiple: allowMultiple && (maxImageUploadCount === undefined || maxImageUploadCount > 1),
-    maxFiles: maxImageUploadCount,
+    multiple: allowMultiple,
   });
 
   return { getUploadButtonProps, getUploadInputProps, openUploader, request };

invokeai/frontend/web/src/features/controlLayers/components/ParamDenoisingStrength.tsx

@@ -14,8 +14,9 @@ import WavyLine from 'common/components/WavyLine';
 import { selectImg2imgStrength, setImg2imgStrength } from 'features/controlLayers/store/paramsSlice';
 import { selectActiveRasterLayerEntities } from 'features/controlLayers/store/selectors';
 import { selectImg2imgStrengthConfig } from 'features/system/store/configSlice';
-import { memo, useCallback } from 'react';
+import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
+import { useSelectedModelConfig } from 'services/api/hooks/useSelectedModelConfig';
 
 const selectHasRasterLayersWithContent = createSelector(
   selectActiveRasterLayerEntities,
@@ -26,6 +27,7 @@ export const ParamDenoisingStrength = memo(() => {
   const img2imgStrength = useAppSelector(selectImg2imgStrength);
   const dispatch = useAppDispatch();
   const hasRasterLayersWithContent = useAppSelector(selectHasRasterLayersWithContent);
+  const selectedModelConfig = useSelectedModelConfig();
 
   const onChange = useCallback(
     (v: number) => {
@@ -39,8 +41,24 @@ export const ParamDenoisingStrength = memo(() => {
 
   const [invokeBlue300] = useToken('colors', ['invokeBlue.300']);
 
+  const isDisabled = useMemo(() => {
+    if (!hasRasterLayersWithContent) {
+      // Denoising strength does nothing if there are no raster layers w/ content
+      return true;
+    }
+    if (
+      selectedModelConfig?.type === 'main' &&
+      selectedModelConfig?.base === 'flux' &&
+      selectedModelConfig.variant === 'inpaint'
+    ) {
+      // Denoising strength is ignored by FLUX Fill, which is indicated by the variant being 'inpaint'
+      return true;
+    }
+    return false;
+  }, [hasRasterLayersWithContent, selectedModelConfig]);
+
   return (
-    <FormControl isDisabled={!hasRasterLayersWithContent} p={1} justifyContent="space-between" h={8}>
+    <FormControl isDisabled={isDisabled} p={1} justifyContent="space-between" h={8}>
       <Flex gap={3} alignItems="center">
         <InformationalPopover feature="paramDenoisingStrength">
           <FormLabel mr={0}>{`${t('parameters.denoisingStrength')}`}</FormLabel>
@@ -49,7 +67,7 @@ export const ParamDenoisingStrength = memo(() => {
           <WavyLine amplitude={img2imgStrength * 10} stroke={invokeBlue300} strokeWidth={1} width={40} height={14} />
         )}
       </Flex>
-      {hasRasterLayersWithContent ? (
+      {!isDisabled ? (
         <>
           <CompositeSlider
             step={config.coarseStep}

invokeai/frontend/web/src/features/dnd/FullscreenDropzone.tsx

@@ -8,12 +8,13 @@ import { useStore } from '@nanostores/react';
 import { getStore } from 'app/store/nanostores/store';
 import { useAppSelector } from 'app/store/storeHooks';
 import { $focusedRegion } from 'common/hooks/focus';
+import { useClientSideUpload } from 'common/hooks/useClientSideUpload';
 import { setFileToPaste } from 'features/controlLayers/components/CanvasPasteModal';
 import { DndDropOverlay } from 'features/dnd/DndDropOverlay';
 import type { DndTargetState } from 'features/dnd/types';
 import { $imageViewer } from 'features/gallery/components/ImageViewer/useImageViewer';
 import { selectAutoAddBoardId } from 'features/gallery/store/gallerySelectors';
-import { selectMaxImageUploadCount } from 'features/system/store/configSlice';
+import { selectIsClientSideUploadEnabled } from 'features/system/store/configSlice';
 import { toast } from 'features/toast/toast';
 import { selectActiveTab } from 'features/ui/store/uiSelectors';
 import { memo, useCallback, useEffect, useRef, useState } from 'react';
@@ -53,13 +54,6 @@ const zUploadFile = z
     (file) => ({ message: `File extension .${file.name.split('.').at(-1)} is not supported` })
   );
 
-const getFilesSchema = (max?: number) => {
-  if (max === undefined) {
-    return z.array(zUploadFile);
-  }
-  return z.array(zUploadFile).max(max);
-};
-
 const sx = {
   position: 'absolute',
   top: 2,
@@ -74,22 +68,19 @@ const sx = {
 export const FullscreenDropzone = memo(() => {
   const { t } = useTranslation();
   const ref = useRef<HTMLDivElement>(null);
-  const maxImageUploadCount = useAppSelector(selectMaxImageUploadCount);
   const [dndState, setDndState] = useState<DndTargetState>('idle');
   const activeTab = useAppSelector(selectActiveTab);
   const isImageViewerOpen = useStore($imageViewer);
+  const isClientSideUploadEnabled = useAppSelector(selectIsClientSideUploadEnabled);
+  const clientSideUpload = useClientSideUpload();
 
   const validateAndUploadFiles = useCallback(
-    (files: File[]) => {
+    async (files: File[]) => {
       const { getState } = getStore();
-      const uploadFilesSchema = getFilesSchema(maxImageUploadCount);
-      const parseResult = uploadFilesSchema.safeParse(files);
+      const parseResult = z.array(zUploadFile).safeParse(files);
 
       if (!parseResult.success) {
-        const description =
-          maxImageUploadCount === undefined
-            ? t('toast.uploadFailedInvalidUploadDesc')
-            : t('toast.uploadFailedInvalidUploadDesc_withCount', { count: maxImageUploadCount });
+        const description = t('toast.uploadFailedInvalidUploadDesc');
 
         toast({
           id: 'UPLOAD_FAILED',
@@ -118,17 +109,23 @@ export const FullscreenDropzone = memo(() => {
 
       const autoAddBoardId = selectAutoAddBoardId(getState());
 
-      const uploadArgs: UploadImageArg[] = files.map((file, i) => ({
-        file,
-        image_category: 'user',
-        is_intermediate: false,
-        board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
-        isFirstUploadOfBatch: i === 0,
-      }));
+      if (isClientSideUploadEnabled) {
+        for (const [i, file] of files.entries()) {
+          await clientSideUpload(file, i);
+        }
+      } else {
+        const uploadArgs: UploadImageArg[] = files.map((file, i) => ({
+          file,
+          image_category: 'user',
+          is_intermediate: false,
+          board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
+          isFirstUploadOfBatch: i === 0,
+        }));
 
-      uploadImages(uploadArgs);
+        uploadImages(uploadArgs);
+      }
     },
-    [activeTab, isImageViewerOpen, maxImageUploadCount, t]
+    [activeTab, isImageViewerOpen, t, isClientSideUploadEnabled, clientSideUpload]
   );
 
   const onPaste = useCallback(

invokeai/frontend/web/src/features/gallery/components/GalleryUploadButton.tsx

@@ -1,31 +1,18 @@
 import { IconButton } from '@invoke-ai/ui-library';
-import { useAppSelector } from 'app/store/storeHooks';
 import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
-import { selectMaxImageUploadCount } from 'features/system/store/configSlice';
 import { t } from 'i18next';
-import { useMemo } from 'react';
 import { PiUploadBold } from 'react-icons/pi';
 
 export const GalleryUploadButton = () => {
-  const maxImageUploadCount = useAppSelector(selectMaxImageUploadCount);
-  const uploadOptions = useMemo(() => ({ allowMultiple: maxImageUploadCount !== 1 }), [maxImageUploadCount]);
-  const uploadApi = useImageUploadButton(uploadOptions);
+  const uploadApi = useImageUploadButton({ allowMultiple: true });
   return (
     <>
       <IconButton
         size="sm"
         alignSelf="stretch"
         variant="link"
-        aria-label={
-          maxImageUploadCount === undefined || maxImageUploadCount > 1
-            ? t('accessibility.uploadImages')
-            : t('accessibility.uploadImage')
-        }
-        tooltip={
-          maxImageUploadCount === undefined || maxImageUploadCount > 1
-            ? t('accessibility.uploadImages')
-            : t('accessibility.uploadImage')
-        }
+        aria-label={t('accessibility.uploadImages')}
+        tooltip={t('accessibility.uploadImages')}
         icon={<PiUploadBold />}
         {...uploadApi.getUploadButtonProps()}
       />

invokeai/frontend/web/src/features/gallery/store/actions.ts

@@ -1,4 +1,5 @@
 import { createAction } from '@reduxjs/toolkit';
+import type { ImageDTO } from 'services/api/types';
 
 export const sentImageToCanvas = createAction('gallery/sentImageToCanvas');
 
@@ -7,3 +8,9 @@ export const imageDownloaded = createAction('gallery/imageDownloaded');
 export const imageCopiedToClipboard = createAction('gallery/imageCopiedToClipboard');
 
 export const imageOpenedInNewTab = createAction('gallery/imageOpenedInNewTab');
+
+export const imageUploadedClientSide = createAction<{
+  imageDTO: ImageDTO;
+  silent: boolean;
+  isFirstUploadOfBatch: boolean;
+}>('gallery/imageUploadedClientSide');

invokeai/frontend/web/src/features/system/store/configSlice.ts

@@ -20,6 +20,7 @@ const initialConfigState: AppConfig = {
   shouldFetchMetadataFromApi: false,
   allowPrivateBoards: false,
   allowPrivateStylePresets: false,
+  allowClientSideUpload: false,
   disabledTabs: [],
   disabledFeatures: ['lightbox', 'faceRestore', 'batches'],
   disabledSDFeatures: ['variation', 'symmetry', 'hires', 'perlinNoise', 'noiseThreshold'],
@@ -218,6 +219,5 @@ export const selectWorkflowFetchDebounce = createConfigSelector((config) => conf
 export const selectMetadataFetchDebounce = createConfigSelector((config) => config.metadataFetchDebounce ?? 300);
 
 export const selectIsModelsTabDisabled = createConfigSelector((config) => config.disabledTabs.includes('models'));
-export const selectMaxImageUploadCount = createConfigSelector((config) => config.maxImageUploadCount);
-
+export const selectIsClientSideUploadEnabled = createConfigSelector((config) => config.allowClientSideUpload);
 export const selectIsLocal = createSelector(selectConfigSlice, (config) => config.isLocal);

invokeai/frontend/web/src/services/api/endpoints/images.ts

@@ -7,6 +7,8 @@ import type {
   DeleteBoardResult,
   GraphAndWorkflowResponse,
   ImageDTO,
+  ImageUploadEntryRequest,
+  ImageUploadEntryResponse,
   ListImagesArgs,
   ListImagesResponse,
   UploadImageArg,
@@ -287,6 +289,7 @@ export const imagesApi = api.injectEndpoints({
           },
         };
       },
+
       invalidatesTags: (result) => {
         if (!result || result.is_intermediate) {
           // Don't add it to anything
@@ -314,7 +317,13 @@ export const imagesApi = api.injectEndpoints({
         ];
       },
     }),
-
+    createImageUploadEntry: build.mutation<ImageUploadEntryResponse, ImageUploadEntryRequest>({
+      query: ({ width, height, board_id }) => ({
+        url: buildImagesUrl(),
+        method: 'POST',
+        body: { width, height, board_id },
+      }),
+    }),
     deleteBoard: build.mutation<DeleteBoardResult, string>({
       query: (board_id) => ({ url: buildBoardsUrl(board_id), method: 'DELETE' }),
       invalidatesTags: () => [
@@ -549,6 +558,7 @@ export const {
   useGetImageWorkflowQuery,
   useLazyGetImageWorkflowQuery,
   useUploadImageMutation,
+  useCreateImageUploadEntryMutation,
   useClearIntermediatesMutation,
   useAddImagesToBoardMutation,
   useRemoveImagesFromBoardMutation,

invokeai/frontend/web/src/services/api/schema.ts

@@ -466,6 +466,30 @@ export type paths = {
         patch?: never;
         trace?: never;
     };
+    "/api/v1/images/": {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        /**
+         * List Image Dtos
+         * @description Gets a list of image DTOs
+         */
+        get: operations["list_image_dtos"];
+        put?: never;
+        /**
+         * Create Image Upload Entry
+         * @description Uploads an image from a URL, not implemented
+         */
+        post: operations["create_image_upload_entry"];
+        delete?: never;
+        options?: never;
+        head?: never;
+        patch?: never;
+        trace?: never;
+    };
     "/api/v1/images/i/{image_name}": {
         parameters: {
             query?: never;
@@ -619,26 +643,6 @@ export type paths = {
         patch?: never;
         trace?: never;
     };
-    "/api/v1/images/": {
-        parameters: {
-            query?: never;
-            header?: never;
-            path?: never;
-            cookie?: never;
-        };
-        /**
-         * List Image Dtos
-         * @description Gets a list of image DTOs
-         */
-        get: operations["list_image_dtos"];
-        put?: never;
-        post?: never;
-        delete?: never;
-        options?: never;
-        head?: never;
-        patch?: never;
-        trace?: never;
-    };
     "/api/v1/images/delete": {
         parameters: {
             query?: never;
@@ -2358,6 +2362,24 @@ export type components = {
              */
             batch_ids: string[];
         };
+        /** Body_create_image_upload_entry */
+        Body_create_image_upload_entry: {
+            /**
+             * Width
+             * @description The width of the image
+             */
+            width: number;
+            /**
+             * Height
+             * @description The height of the image
+             */
+            height: number;
+            /**
+             * Board Id
+             * @description The board to add this image to, if any
+             */
+            board_id?: string | null;
+        };
         /** Body_create_style_preset */
         Body_create_style_preset: {
             /**
@@ -6451,6 +6473,7 @@ export type components = {
          * @description Expands a mask with a fade effect. The mask uses black to indicate areas to keep from the generated image and white for areas to discard.
          *     The mask is thresholded to create a binary mask, and then a distance transform is applied to create a fade effect.
          *     The fade size is specified in pixels, and the mask is expanded by that amount. The result is a mask with a smooth transition from black to white.
+         *     If the fade size is 0, the mask is returned as-is.
          */
         ExpandMaskWithFadeInvocation: {
             /**
@@ -10753,6 +10776,16 @@ export type components = {
              */
             type: "i2l";
         };
+        /** ImageUploadEntry */
+        ImageUploadEntry: {
+            /** @description The image DTO */
+            image_dto: components["schemas"]["ImageDTO"];
+            /**
+             * Presigned Url
+             * @description The URL to get the presigned URL for the image upload
+             */
+            presigned_url: string;
+        };
         /**
          * ImageUrlsDTO
          * @description The URLs for an image and its thumbnail.
@@ -23218,6 +23251,87 @@ export interface operations {
             };
         };
     };
+    list_image_dtos: {
+        parameters: {
+            query?: {
+                /** @description The origin of images to list. */
+                image_origin?: components["schemas"]["ResourceOrigin"] | null;
+                /** @description The categories of image to include. */
+                categories?: components["schemas"]["ImageCategory"][] | null;
+                /** @description Whether to list intermediate images. */
+                is_intermediate?: boolean | null;
+                /** @description The board id to filter by. Use 'none' to find images without a board. */
+                board_id?: string | null;
+                /** @description The page offset */
+                offset?: number;
+                /** @description The number of images per page */
+                limit?: number;
+                /** @description The order of sort */
+                order_dir?: components["schemas"]["SQLiteDirection"];
+                /** @description Whether to sort by starred images first */
+                starred_first?: boolean;
+                /** @description The term to search for */
+                search_term?: string | null;
+            };
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        requestBody?: never;
+        responses: {
+            /** @description Successful Response */
+            200: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["OffsetPaginatedResults_ImageDTO_"];
+                };
+            };
+            /** @description Validation Error */
+            422: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["HTTPValidationError"];
+                };
+            };
+        };
+    };
+    create_image_upload_entry: {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        requestBody: {
+            content: {
+                "application/json": components["schemas"]["Body_create_image_upload_entry"];
+            };
+        };
+        responses: {
+            /** @description Successful Response */
+            200: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["ImageUploadEntry"];
+                };
+            };
+            /** @description Validation Error */
+            422: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["HTTPValidationError"];
+                };
+            };
+        };
+    };
     get_image_dto: {
         parameters: {
             query?: never;
@@ -23571,54 +23685,6 @@ export interface operations {
             };
         };
     };
-    list_image_dtos: {
-        parameters: {
-            query?: {
-                /** @description The origin of images to list. */
-                image_origin?: components["schemas"]["ResourceOrigin"] | null;
-                /** @description The categories of image to include. */
-                categories?: components["schemas"]["ImageCategory"][] | null;
-                /** @description Whether to list intermediate images. */
-                is_intermediate?: boolean | null;
-                /** @description The board id to filter by. Use 'none' to find images without a board. */
-                board_id?: string | null;
-                /** @description The page offset */
-                offset?: number;
-                /** @description The number of images per page */
-                limit?: number;
-                /** @description The order of sort */
-                order_dir?: components["schemas"]["SQLiteDirection"];
-                /** @description Whether to sort by starred images first */
-                starred_first?: boolean;
-                /** @description The term to search for */
-                search_term?: string | null;
-            };
-            header?: never;
-            path?: never;
-            cookie?: never;
-        };
-        requestBody?: never;
-        responses: {
-            /** @description Successful Response */
-            200: {
-                headers: {
-                    [name: string]: unknown;
-                };
-                content: {
-                    "application/json": components["schemas"]["OffsetPaginatedResults_ImageDTO_"];
-                };
-            };
-            /** @description Validation Error */
-            422: {
-                headers: {
-                    [name: string]: unknown;
-                };
-                content: {
-                    "application/json": components["schemas"]["HTTPValidationError"];
-                };
-            };
-        };
-    };
     delete_images_from_list: {
         parameters: {
             query?: never;

invokeai/frontend/web/src/services/api/types.ts

@@ -354,3 +354,6 @@ export type UploadImageArg = {
    */
   isFirstUploadOfBatch?: boolean;
 };
+
+export type ImageUploadEntryResponse = S['ImageUploadEntry'];
+export type ImageUploadEntryRequest = paths['/api/v1/images/']['post']['requestBody']['content']['application/json'];

invokeai/version/invokeai_version.py

@@ -1 +1 @@
-__version__ = "5.9.0"
+__version__ = "5.9.1"

pyproject.toml

@@ -117,7 +117,7 @@ dependencies = [
 ]
 "dev" = ["jurigged", "pudb", "snakeviz", "gprof2dot"]
 "test" = [
-  "ruff~=0.9.9",
+  "ruff~=0.11.2",
   "ruff-lsp~=0.0.62",
   "mypy",
   "pre-commit",

scripts/strip_models.py

@@ -71,7 +71,7 @@ def create_stripped_model(original_model_path: Path, stripped_model_path: Path)
     print(f"Created clone of {original.name} at {stripped.path}")
 
     for component_path in stripped.component_paths():
-        original_state_dict = ModelOnDisk.load_state_dict(component_path)
+        original_state_dict = stripped.load_state_dict(component_path)
         stripped_state_dict = strip(original_state_dict)  # type: ignore
         with open(component_path, "w") as f:
             json.dump(stripped_state_dict, f, indent=4)

tests/backend/flux/controlnet/test_state_dict_utils.py

@@ -24,7 +24,7 @@ from tests.backend.flux.controlnet.xlabs_flux_controlnet_state_dict import xlabs
     ],
 )
 def test_is_state_dict_xlabs_controlnet(sd_shapes: dict[str, list[int]], expected: bool):
-    sd = {k: None for k in sd_shapes}
+    sd = dict.fromkeys(sd_shapes)
     assert is_state_dict_xlabs_controlnet(sd) == expected
 
 
@@ -37,7 +37,7 @@ def test_is_state_dict_xlabs_controlnet(sd_shapes: dict[str, list[int]], expecte
     ],
 )
 def test_is_state_dict_instantx_controlnet(sd_keys: list[str], expected: bool):
-    sd = {k: None for k in sd_keys}
+    sd = dict.fromkeys(sd_keys)
     assert is_state_dict_instantx_controlnet(sd) == expected
 
 

tests/backend/flux/ip_adapter/test_xlabs_ip_adapter_flux.py

@@ -19,7 +19,7 @@ from tests.backend.flux.ip_adapter.xlabs_flux_ip_adapter_v2_state_dict import xl
 @pytest.mark.parametrize("sd_shapes", [xlabs_flux_ip_adapter_sd_shapes, xlabs_flux_ip_adapter_v2_sd_shapes])
 def test_is_state_dict_xlabs_ip_adapter(sd_shapes: dict[str, list[int]]):
     # Construct a dummy state_dict.
-    sd = {k: None for k in sd_shapes}
+    sd = dict.fromkeys(sd_shapes)
 
     assert is_state_dict_xlabs_ip_adapter(sd)