InvokeAI 4.2.9.dev20240824__tar.gz → 4.2.9rc2__tar.gz

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/InvokeAI.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: InvokeAI
-Version: 4.2.9.dev20240824
+Version: 4.2.9rc2
 Summary: An implementation of Stable Diffusion which provides various new features and options to aid the image generation process
 Author-email: The InvokeAI Project <lincoln.stein@gmail.com>
 License:                                  Apache License

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/InvokeAI.egg-info/SOURCES.txt

@@ -40,6 +40,10 @@ invokeai/app/invocations/cv.py
 invokeai/app/invocations/denoise_latents.py
 invokeai/app/invocations/facetools.py
 invokeai/app/invocations/fields.py
+invokeai/app/invocations/flux_denoise.py
+invokeai/app/invocations/flux_text_encoder.py
+invokeai/app/invocations/flux_vae_decode.py
+invokeai/app/invocations/flux_vae_encode.py
 invokeai/app/invocations/grounding_dino.py
 invokeai/app/invocations/ideal_size.py
 invokeai/app/invocations/image.py
@@ -174,7 +178,6 @@ invokeai/app/services/shared/sqlite_migrator/migrations/migration_11.py
 invokeai/app/services/shared/sqlite_migrator/migrations/migration_12.py
 invokeai/app/services/shared/sqlite_migrator/migrations/migration_13.py
 invokeai/app/services/shared/sqlite_migrator/migrations/migration_14.py
-invokeai/app/services/shared/sqlite_migrator/migrations/migration_15.py
 invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py
 invokeai/app/services/shared/sqlite_migrator/migrations/migration_3.py
 invokeai/app/services/shared/sqlite_migrator/migrations/migration_4.py
@@ -218,7 +221,9 @@ invokeai/app/services/workflow_records/workflow_records_base.py
 invokeai/app/services/workflow_records/workflow_records_common.py
 invokeai/app/services/workflow_records/workflow_records_sqlite.py
 invokeai/app/services/workflow_records/default_workflows/ESRGAN Upscaling with Canny ControlNet.json
+invokeai/app/services/workflow_records/default_workflows/FLUX Image to Image.json
 invokeai/app/services/workflow_records/default_workflows/Face Detailer with IP-Adapter & Canny (See Note in Details).json
+invokeai/app/services/workflow_records/default_workflows/Flux Text to Image.json
 invokeai/app/services/workflow_records/default_workflows/Multi ControlNet (Canny & Depth).json
 invokeai/app/services/workflow_records/default_workflows/MultiDiffusion SD1.5.json
 invokeai/app/services/workflow_records/default_workflows/MultiDiffusion SDXL.json
@@ -247,6 +252,15 @@ invokeai/backend/model_patcher.py
 invokeai/backend/raw_model.py
 invokeai/backend/spandrel_image_to_image_model.py
 invokeai/backend/textual_inversion.py
+invokeai/backend/flux/denoise.py
+invokeai/backend/flux/inpaint_extension.py
+invokeai/backend/flux/math.py
+invokeai/backend/flux/model.py
+invokeai/backend/flux/sampling_utils.py
+invokeai/backend/flux/util.py
+invokeai/backend/flux/modules/autoencoder.py
+invokeai/backend/flux/modules/conditioner.py
+invokeai/backend/flux/modules/layers.py
 invokeai/backend/image_util/__init__.py
 invokeai/backend/image_util/canny.py
 invokeai/backend/image_util/hed.py
@@ -303,6 +317,7 @@ invokeai/backend/model_manager/load/model_cache/model_cache_default.py
 invokeai/backend/model_manager/load/model_cache/model_locker.py
 invokeai/backend/model_manager/load/model_loaders/__init__.py
 invokeai/backend/model_manager/load/model_loaders/controlnet.py
+invokeai/backend/model_manager/load/model_loaders/flux.py
 invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py
 invokeai/backend/model_manager/load/model_loaders/ip_adapter.py
 invokeai/backend/model_manager/load/model_loaders/lora.py
@@ -320,6 +335,12 @@ invokeai/backend/model_manager/util/libc_util.py
 invokeai/backend/model_manager/util/model_util.py
 invokeai/backend/model_manager/util/select_hf_files.py
 invokeai/backend/onnx/onnx_runtime.py
+invokeai/backend/quantization/__init__.py
+invokeai/backend/quantization/bnb_llm_int8.py
+invokeai/backend/quantization/bnb_nf4.py
+invokeai/backend/quantization/scripts/load_flux_model_bnb_llm_int8.py
+invokeai/backend/quantization/scripts/load_flux_model_bnb_nf4.py
+invokeai/backend/quantization/scripts/quantize_t5_xxl_bnb_llm_int8.py
 invokeai/backend/stable_diffusion/__init__.py
 invokeai/backend/stable_diffusion/denoise_context.py
 invokeai/backend/stable_diffusion/diffusers_pipeline.py
@@ -386,11 +407,11 @@ invokeai/frontend/install/__init__.py
 invokeai/frontend/install/import_images.py
 invokeai/frontend/web/__init__.py
 invokeai/frontend/web/dist/index.html
-invokeai/frontend/web/dist/assets/App-BoqVb0rn.js
+invokeai/frontend/web/dist/assets/App-D1NtDI7i.js
 invokeai/frontend/web/dist/assets/App-DEu4J2pT.css
 invokeai/frontend/web/dist/assets/ThemeLocaleProvider-B2NxUfIp.css
-invokeai/frontend/web/dist/assets/ThemeLocaleProvider-CbAimOPZ.js
-invokeai/frontend/web/dist/assets/index-BNntoaVx.js
+invokeai/frontend/web/dist/assets/ThemeLocaleProvider-CFdiQhFL.js
+invokeai/frontend/web/dist/assets/index-CsEnX8zB.js
 invokeai/frontend/web/dist/assets/inter-cyrillic-ext-wght-normal-B2xhLi22.woff2
 invokeai/frontend/web/dist/assets/inter-cyrillic-wght-normal-CMZtQduZ.woff2
 invokeai/frontend/web/dist/assets/inter-greek-ext-wght-normal-CGAr0uHJ.woff2
@@ -398,7 +419,6 @@ invokeai/frontend/web/dist/assets/inter-greek-wght-normal-CaVNZxsx.woff2
 invokeai/frontend/web/dist/assets/inter-latin-ext-wght-normal-CFHvXkgd.woff2
 invokeai/frontend/web/dist/assets/inter-latin-wght-normal-C2S99t-D.woff2
 invokeai/frontend/web/dist/assets/inter-vietnamese-wght-normal-CBcvBZtf.woff2
-invokeai/frontend/web/dist/assets/worker-Dywf1EVc.js
 invokeai/frontend/web/dist/assets/images/invoke-alert-favicon.svg
 invokeai/frontend/web/dist/assets/images/invoke-avatar-circle.svg
 invokeai/frontend/web/dist/assets/images/invoke-avatar-square.svg

InvokeAI-4.2.9rc2/InvokeAI.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+invokeai-web = invokeai.app.run_app:run_app

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/InvokeAI.egg-info/requires.txt

@@ -11,6 +11,7 @@ onnxruntime==1.16.3
 opencv-python==4.9.0.80
 pytorch-lightning==2.1.3
 safetensors==0.4.3
+sentencepiece==0.2.0
 spandrel==0.3.4
 timm==0.6.13
 torch==2.2.2
@@ -52,6 +53,9 @@ scikit-image~=0.21.0
 semver~=3.0.1
 test-tube~=0.7.5
 
+[:sys_platform != "darwin"]
+bitsandbytes==0.43.3
+
 [:sys_platform == "win32"]
 windows-curses
 

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: InvokeAI
-Version: 4.2.9.dev20240824
+Version: 4.2.9rc2
 Summary: An implementation of Stable Diffusion which provides various new features and options to aid the image generation process
 Author-email: The InvokeAI Project <lincoln.stein@gmail.com>
 License:                                  Apache License

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/invokeai/app/api/routers/model_manager.py

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

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/invokeai/app/api/routers/session_queue.py

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

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/invokeai/app/invocations/denoise_latents.py

@@ -185,7 +185,7 @@ class DenoiseLatentsInvocation(BaseInvocation):
     )
     denoise_mask: Optional[DenoiseMaskField] = InputField(
         default=None,
-        description=FieldDescriptions.mask,
+        description=FieldDescriptions.denoise_mask,
         input=Input.Connection,
         ui_order=8,
     )

{InvokeAI-4.2.9.dev20240824 → InvokeAI-4.2.9rc2}/invokeai/app/invocations/fields.py

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

InvokeAI-4.2.9rc2/invokeai/app/invocations/flux_denoise.py

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