From 9a1420280e3e5dc3ba636c13c6c6ef5094a56cfa Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Sun, 21 Jul 2024 17:33:43 +0300 Subject: [PATCH 01/61] Add rescale cfg support to denoise --- invokeai/app/invocations/denoise_latents.py | 5 +++ .../stable_diffusion/diffusion_backend.py | 8 ++--- .../extension_callback_type.py | 2 +- .../extensions/rescale_cfg.py | 36 +++++++++++++++++++ 4 files changed, 46 insertions(+), 5 deletions(-) create mode 100644 invokeai/backend/stable_diffusion/extensions/rescale_cfg.py diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index ccacc3303c..68deb68445 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -59,6 +59,7 @@ from invokeai.backend.stable_diffusion.diffusion.custom_atttention import Custom from invokeai.backend.stable_diffusion.diffusion_backend import StableDiffusionBackend from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType from invokeai.backend.stable_diffusion.extensions.preview import PreviewExt +from invokeai.backend.stable_diffusion.extensions.rescale_cfg import RescaleCFGExt from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES @@ -790,6 +791,10 @@ class DenoiseLatentsInvocation(BaseInvocation): ext_manager.add_extension(PreviewExt(step_callback)) + ### cfg rescale + if self.cfg_rescale_multiplier > 0: + ext_manager.add_extension(RescaleCFGExt(self.cfg_rescale_multiplier)) + # ext: t2i/ip adapter ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx) diff --git a/invokeai/backend/stable_diffusion/diffusion_backend.py b/invokeai/backend/stable_diffusion/diffusion_backend.py index 806deb5e03..d41aa63c60 100644 --- a/invokeai/backend/stable_diffusion/diffusion_backend.py +++ b/invokeai/backend/stable_diffusion/diffusion_backend.py @@ -76,12 +76,12 @@ class StableDiffusionBackend: both_noise_pred = self.run_unet(ctx, ext_manager, ConditioningMode.Both) ctx.negative_noise_pred, ctx.positive_noise_pred = both_noise_pred.chunk(2) - # ext: override apply_cfg - ctx.noise_pred = self.apply_cfg(ctx) + # ext: override combine_noise_preds + ctx.noise_pred = self.combine_noise_preds(ctx) # ext: cfg_rescale [modify_noise_prediction] # TODO: rename - ext_manager.run_callback(ExtensionCallbackType.POST_APPLY_CFG, ctx) + ext_manager.run_callback(ExtensionCallbackType.POST_COMBINE_NOISE_PREDS, ctx) # compute the previous noisy sample x_t -> x_t-1 step_output = ctx.scheduler.step(ctx.noise_pred, ctx.timestep, ctx.latents, **ctx.inputs.scheduler_step_kwargs) @@ -95,7 +95,7 @@ class StableDiffusionBackend: return step_output @staticmethod - def apply_cfg(ctx: DenoiseContext) -> torch.Tensor: + def combine_noise_preds(ctx: DenoiseContext) -> torch.Tensor: guidance_scale = ctx.inputs.conditioning_data.guidance_scale if isinstance(guidance_scale, list): guidance_scale = guidance_scale[ctx.step_index] diff --git a/invokeai/backend/stable_diffusion/extension_callback_type.py b/invokeai/backend/stable_diffusion/extension_callback_type.py index aaefbd7ed0..e4c365007b 100644 --- a/invokeai/backend/stable_diffusion/extension_callback_type.py +++ b/invokeai/backend/stable_diffusion/extension_callback_type.py @@ -9,4 +9,4 @@ class ExtensionCallbackType(Enum): POST_STEP = "post_step" PRE_UNET = "pre_unet" POST_UNET = "post_unet" - POST_APPLY_CFG = "post_apply_cfg" + POST_COMBINE_NOISE_PREDS = "post_combine_noise_preds" diff --git a/invokeai/backend/stable_diffusion/extensions/rescale_cfg.py b/invokeai/backend/stable_diffusion/extensions/rescale_cfg.py new file mode 100644 index 0000000000..51fad975e7 --- /dev/null +++ b/invokeai/backend/stable_diffusion/extensions/rescale_cfg.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType +from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback + +if TYPE_CHECKING: + from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext + + +class RescaleCFGExt(ExtensionBase): + def __init__(self, rescale_multiplier: float): + super().__init__() + self.rescale_multiplier = rescale_multiplier + + @staticmethod + def _rescale_cfg(total_noise_pred: torch.Tensor, pos_noise_pred: torch.Tensor, multiplier: float = 0.7): + """Implementation of Algorithm 2 from https://arxiv.org/pdf/2305.08891.pdf.""" + ro_pos = torch.std(pos_noise_pred, dim=(1, 2, 3), keepdim=True) + ro_cfg = torch.std(total_noise_pred, dim=(1, 2, 3), keepdim=True) + + x_rescaled = total_noise_pred * (ro_pos / ro_cfg) + x_final = multiplier * x_rescaled + (1.0 - multiplier) * total_noise_pred + return x_final + + @callback(ExtensionCallbackType.POST_COMBINE_NOISE_PREDS) + def rescale_noise_pred(self, ctx: DenoiseContext): + if self.rescale_multiplier > 0: + ctx.noise_pred = self._rescale_cfg( + ctx.noise_pred, + ctx.positive_noise_pred, + self.rescale_multiplier, + ) From e046e60e1cfb3552e27d97bed0acfd6345adc0fe Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Sun, 21 Jul 2024 18:31:10 +0300 Subject: [PATCH 02/61] Add FreeU support to denoise --- invokeai/app/invocations/denoise_latents.py | 9 +++- .../stable_diffusion/extensions/freeu.py | 42 +++++++++++++++++++ .../stable_diffusion/extensions_manager.py | 10 +++-- 3 files changed, 56 insertions(+), 5 deletions(-) create mode 100644 invokeai/backend/stable_diffusion/extensions/freeu.py diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index ccacc3303c..e043e884f9 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -58,6 +58,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( from invokeai.backend.stable_diffusion.diffusion.custom_atttention import CustomAttnProcessor2_0 from invokeai.backend.stable_diffusion.diffusion_backend import StableDiffusionBackend from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType +from invokeai.backend.stable_diffusion.extensions.freeu import FreeUExt from invokeai.backend.stable_diffusion.extensions.preview import PreviewExt from invokeai.backend.stable_diffusion.extensions_manager import ExtensionsManager from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP @@ -790,18 +791,22 @@ class DenoiseLatentsInvocation(BaseInvocation): ext_manager.add_extension(PreviewExt(step_callback)) + ### freeu + if self.unet.freeu_config: + ext_manager.add_extension(FreeUExt(self.unet.freeu_config)) + # ext: t2i/ip adapter ext_manager.run_callback(ExtensionCallbackType.SETUP, denoise_ctx) unet_info = context.models.load(self.unet.unet) assert isinstance(unet_info.model, UNet2DConditionModel) with ( - unet_info.model_on_device() as (model_state_dict, unet), + unet_info.model_on_device() as (cached_weights, unet), ModelPatcher.patch_unet_attention_processor(unet, denoise_ctx.inputs.attention_processor_cls), # ext: controlnet ext_manager.patch_extensions(unet), # ext: freeu, seamless, ip adapter, lora - ext_manager.patch_unet(model_state_dict, unet), + ext_manager.patch_unet(unet, cached_weights), ): sd_backend = StableDiffusionBackend(unet, scheduler) denoise_ctx.unet = unet diff --git a/invokeai/backend/stable_diffusion/extensions/freeu.py b/invokeai/backend/stable_diffusion/extensions/freeu.py new file mode 100644 index 0000000000..c723aaee0b --- /dev/null +++ b/invokeai/backend/stable_diffusion/extensions/freeu.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Dict, Optional + +from diffusers import UNet2DConditionModel + +from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType +from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase + +if TYPE_CHECKING: + from invokeai.app.shared.models import FreeUConfig + + +class FreeUExt(ExtensionBase): + def __init__( + self, + freeu_config: Optional[FreeUConfig], + ): + super().__init__() + self.freeu_config = freeu_config + + @contextmanager + def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None): + did_apply_freeu = False + try: + assert hasattr(unet, "enable_freeu") # mypy doesn't pick up this attribute? + if self.freeu_config is not None: + unet.enable_freeu( + b1=self.freeu_config.b1, + b2=self.freeu_config.b2, + s1=self.freeu_config.s1, + s2=self.freeu_config.s2, + ) + did_apply_freeu = True + + yield + + finally: + assert hasattr(unet, "disable_freeu") # mypy doesn't pick up this attribute? + if did_apply_freeu: + unet.disable_freeu() diff --git a/invokeai/backend/stable_diffusion/extensions_manager.py b/invokeai/backend/stable_diffusion/extensions_manager.py index 1cae2e4219..9c4347a56c 100644 --- a/invokeai/backend/stable_diffusion/extensions_manager.py +++ b/invokeai/backend/stable_diffusion/extensions_manager.py @@ -63,9 +63,13 @@ class ExtensionsManager: yield None @contextmanager - def patch_unet(self, state_dict: Dict[str, torch.Tensor], unet: UNet2DConditionModel): + def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None): if self._is_canceled and self._is_canceled(): raise CanceledException - # TODO: create logic in PR with extension which uses it - yield None + # TODO: create weight patch logic in PR with extension which uses it + with ExitStack() as exit_stack: + for ext in self._extensions: + exit_stack.enter_context(ext.patch_unet(unet, cached_weights)) + + yield None From 5772965f092645f89768fa9be0a7d174168bfc24 Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Sun, 21 Jul 2024 18:31:30 +0300 Subject: [PATCH 03/61] Fix slightly different output with old backend --- invokeai/backend/stable_diffusion/diffusion_backend.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/invokeai/backend/stable_diffusion/diffusion_backend.py b/invokeai/backend/stable_diffusion/diffusion_backend.py index 806deb5e03..60a21bdc02 100644 --- a/invokeai/backend/stable_diffusion/diffusion_backend.py +++ b/invokeai/backend/stable_diffusion/diffusion_backend.py @@ -100,8 +100,10 @@ class StableDiffusionBackend: if isinstance(guidance_scale, list): guidance_scale = guidance_scale[ctx.step_index] - return torch.lerp(ctx.negative_noise_pred, ctx.positive_noise_pred, guidance_scale) - # return ctx.negative_noise_pred + guidance_scale * (ctx.positive_noise_pred - ctx.negative_noise_pred) + # Note: Although logically it same, it seams that precision errors differs. + # This sometimes results in slightly different output. + # return torch.lerp(ctx.negative_noise_pred, ctx.positive_noise_pred, guidance_scale) + return ctx.negative_noise_pred + guidance_scale * (ctx.positive_noise_pred - ctx.negative_noise_pred) def run_unet(self, ctx: DenoiseContext, ext_manager: ExtensionsManager, conditioning_mode: ConditioningMode): sample = ctx.latent_model_input From 1748848b7b13f93f58400af36a9ed4e280ce639c Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Sun, 21 Jul 2024 18:37:20 +0300 Subject: [PATCH 04/61] Ruff fixes --- invokeai/backend/stable_diffusion/extensions/freeu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/backend/stable_diffusion/extensions/freeu.py b/invokeai/backend/stable_diffusion/extensions/freeu.py index c723aaee0b..0f6c47a773 100644 --- a/invokeai/backend/stable_diffusion/extensions/freeu.py +++ b/invokeai/backend/stable_diffusion/extensions/freeu.py @@ -3,9 +3,9 @@ from __future__ import annotations from contextlib import contextmanager from typing import TYPE_CHECKING, Dict, Optional +import torch from diffusers import UNet2DConditionModel -from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase if TYPE_CHECKING: From d0435575c1024f753fe79817c07fc8a83fb189f6 Mon Sep 17 00:00:00 2001 From: Eugene Brodsky Date: Sun, 21 Jul 2024 23:18:36 -0400 Subject: [PATCH 05/61] chore(deps): bump fastapi-events to the next minor version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9953c1c1a0..9acaa17e44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,7 @@ dependencies = [ "transformers==4.41.1", # Core application dependencies, pinned for reproducible builds. - "fastapi-events==0.11.0", + "fastapi-events==0.11.1", "fastapi==0.111.0", "huggingface-hub==0.23.1", "pydantic-settings==2.2.1", From 1b359b55cb57c7ad9de3a5a026f393f3f82fc7e2 Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Mon, 22 Jul 2024 22:17:29 +0300 Subject: [PATCH 06/61] Suggested changes Co-Authored-By: Ryan Dick <14897797+RyanJDick@users.noreply.github.com> --- .../stable_diffusion/denoise_context.py | 20 +++++++++---------- .../extensions/rescale_cfg.py | 6 +++--- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/invokeai/backend/stable_diffusion/denoise_context.py b/invokeai/backend/stable_diffusion/denoise_context.py index 2b43d3fb0f..9060d54977 100644 --- a/invokeai/backend/stable_diffusion/denoise_context.py +++ b/invokeai/backend/stable_diffusion/denoise_context.py @@ -83,47 +83,47 @@ class DenoiseContext: unet: Optional[UNet2DConditionModel] = None # Current state of latent-space image in denoising process. - # None until `pre_denoise_loop` callback. + # None until `PRE_DENOISE_LOOP` callback. # Shape: [batch, channels, latent_height, latent_width] latents: Optional[torch.Tensor] = None # Current denoising step index. - # None until `pre_step` callback. + # None until `PRE_STEP` callback. step_index: Optional[int] = None # Current denoising step timestep. - # None until `pre_step` callback. + # None until `PRE_STEP` callback. timestep: Optional[torch.Tensor] = None # Arguments which will be passed to UNet model. - # Available in `pre_unet`/`post_unet` callbacks, otherwise will be None. + # Available in `PRE_UNET`/`POST_UNET` callbacks, otherwise will be None. unet_kwargs: Optional[UNetKwargs] = None # SchedulerOutput class returned from step function(normally, generated by scheduler). - # Supposed to be used only in `post_step` callback, otherwise can be None. + # Supposed to be used only in `POST_STEP` callback, otherwise can be None. step_output: Optional[SchedulerOutput] = None # Scaled version of `latents`, which will be passed to unet_kwargs initialization. - # Available in events inside step(between `pre_step` and `post_stop`). + # Available in events inside step(between `PRE_STEP` and `POST_STEP`). # Shape: [batch, channels, latent_height, latent_width] latent_model_input: Optional[torch.Tensor] = None # [TMP] Defines on which conditionings current unet call will be runned. - # Available in `pre_unet`/`post_unet` callbacks, otherwise will be None. + # Available in `PRE_UNET`/`POST_UNET` callbacks, otherwise will be None. conditioning_mode: Optional[ConditioningMode] = None # [TMP] Noise predictions from negative conditioning. - # Available in `apply_cfg` and `post_apply_cfg` callbacks, otherwise will be None. + # Available in `POST_COMBINE_NOISE_PREDS` callback, otherwise will be None. # Shape: [batch, channels, latent_height, latent_width] negative_noise_pred: Optional[torch.Tensor] = None # [TMP] Noise predictions from positive conditioning. - # Available in `apply_cfg` and `post_apply_cfg` callbacks, otherwise will be None. + # Available in `POST_COMBINE_NOISE_PREDS` callback, otherwise will be None. # Shape: [batch, channels, latent_height, latent_width] positive_noise_pred: Optional[torch.Tensor] = None # Combined noise prediction from passed conditionings. - # Available in `apply_cfg` and `post_apply_cfg` callbacks, otherwise will be None. + # Available in `POST_COMBINE_NOISE_PREDS` callback, otherwise will be None. # Shape: [batch, channels, latent_height, latent_width] noise_pred: Optional[torch.Tensor] = None diff --git a/invokeai/backend/stable_diffusion/extensions/rescale_cfg.py b/invokeai/backend/stable_diffusion/extensions/rescale_cfg.py index 51fad975e7..7cccbb8a2b 100644 --- a/invokeai/backend/stable_diffusion/extensions/rescale_cfg.py +++ b/invokeai/backend/stable_diffusion/extensions/rescale_cfg.py @@ -14,7 +14,7 @@ if TYPE_CHECKING: class RescaleCFGExt(ExtensionBase): def __init__(self, rescale_multiplier: float): super().__init__() - self.rescale_multiplier = rescale_multiplier + self._rescale_multiplier = rescale_multiplier @staticmethod def _rescale_cfg(total_noise_pred: torch.Tensor, pos_noise_pred: torch.Tensor, multiplier: float = 0.7): @@ -28,9 +28,9 @@ class RescaleCFGExt(ExtensionBase): @callback(ExtensionCallbackType.POST_COMBINE_NOISE_PREDS) def rescale_noise_pred(self, ctx: DenoiseContext): - if self.rescale_multiplier > 0: + if self._rescale_multiplier > 0: ctx.noise_pred = self._rescale_cfg( ctx.noise_pred, ctx.positive_noise_pred, - self.rescale_multiplier, + self._rescale_multiplier, ) From 339dddd018442c4efe49686be871f09724ca476f Mon Sep 17 00:00:00 2001 From: chainchompa Date: Mon, 22 Jul 2024 16:03:01 -0400 Subject: [PATCH 07/61] update uncategorized board totals when deleting and moving images --- .../web/src/services/api/endpoints/images.ts | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/invokeai/frontend/web/src/services/api/endpoints/images.ts b/invokeai/frontend/web/src/services/api/endpoints/images.ts index 2040021d6d..06ec9fc162 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/images.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/images.ts @@ -104,6 +104,10 @@ export const imagesApi = api.injectEndpoints({ type: 'Board', id: boardId, }, + { + type: 'BoardImagesTotal', + id: boardId, + }, ]; }, }), @@ -136,6 +140,10 @@ export const imagesApi = api.injectEndpoints({ type: 'Board', id: boardId, }, + { + type: 'BoardImagesTotal', + id: boardId, + }, ]; return tags; @@ -169,6 +177,10 @@ export const imagesApi = api.injectEndpoints({ type: 'Board', id: boardId, }, + { + type: 'BoardImagesTotal', + id: boardId, + }, ]; }, }), @@ -300,6 +312,10 @@ export const imagesApi = api.injectEndpoints({ type: 'Board', id: boardId, }, + { + type: 'BoardImagesTotal', + id: boardId, + }, ]; }, }), @@ -362,6 +378,10 @@ export const imagesApi = api.injectEndpoints({ }, { type: 'Board', id: board_id }, { type: 'Board', id: imageDTO.board_id ?? 'none' }, + { + type: 'BoardImagesTotal', + id: imageDTO.board_id ?? 'none', + }, ]; }, }), @@ -393,6 +413,10 @@ export const imagesApi = api.injectEndpoints({ }, { type: 'Board', id: imageDTO.board_id ?? 'none' }, { type: 'Board', id: 'none' }, + { + type: 'BoardImagesTotal', + id: imageDTO.board_id ?? 'none', + }, ]; }, }), @@ -434,6 +458,10 @@ export const imagesApi = api.injectEndpoints({ tags.push({ type: 'Image', id: imageDTO.image_name }); } tags.push({ type: 'Board', id: board_id }); + tags.push({ + type: 'BoardImagesTotal', + id: board_id ?? 'none', + }); return tags; }, }), @@ -480,6 +508,10 @@ export const imagesApi = api.injectEndpoints({ } tags.push({ type: 'Image', id: image_name }); tags.push({ type: 'Board', id: board_id }); + tags.push({ + type: 'BoardImagesTotal', + id: board_id ?? 'none', + }); }); return tags; From 5f0fe3c8a986cc32ec20a40c3df56145bb0222ab Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Mon, 22 Jul 2024 23:09:11 +0300 Subject: [PATCH 08/61] Suggested changes Co-Authored-By: Ryan Dick <14897797+RyanJDick@users.noreply.github.com> --- .../stable_diffusion/diffusion_backend.py | 4 +-- .../stable_diffusion/extensions/base.py | 4 +-- .../stable_diffusion/extensions/freeu.py | 27 +++++++------------ 3 files changed, 14 insertions(+), 21 deletions(-) diff --git a/invokeai/backend/stable_diffusion/diffusion_backend.py b/invokeai/backend/stable_diffusion/diffusion_backend.py index 60a21bdc02..5d0a68513f 100644 --- a/invokeai/backend/stable_diffusion/diffusion_backend.py +++ b/invokeai/backend/stable_diffusion/diffusion_backend.py @@ -100,8 +100,8 @@ class StableDiffusionBackend: if isinstance(guidance_scale, list): guidance_scale = guidance_scale[ctx.step_index] - # Note: Although logically it same, it seams that precision errors differs. - # This sometimes results in slightly different output. + # Note: Although this `torch.lerp(...)` line is logically equivalent to the current CFG line, it seems to result + # in slightly different outputs. It is suspected that this is caused by small precision differences. # return torch.lerp(ctx.negative_noise_pred, ctx.positive_noise_pred, guidance_scale) return ctx.negative_noise_pred + guidance_scale * (ctx.positive_noise_pred - ctx.negative_noise_pred) diff --git a/invokeai/backend/stable_diffusion/extensions/base.py b/invokeai/backend/stable_diffusion/extensions/base.py index 802af86e6d..6a85a2e441 100644 --- a/invokeai/backend/stable_diffusion/extensions/base.py +++ b/invokeai/backend/stable_diffusion/extensions/base.py @@ -2,7 +2,7 @@ from __future__ import annotations from contextlib import contextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Dict, List +from typing import TYPE_CHECKING, Callable, Dict, List, Optional import torch from diffusers import UNet2DConditionModel @@ -56,5 +56,5 @@ class ExtensionBase: yield None @contextmanager - def patch_unet(self, state_dict: Dict[str, torch.Tensor], unet: UNet2DConditionModel): + def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None): yield None diff --git a/invokeai/backend/stable_diffusion/extensions/freeu.py b/invokeai/backend/stable_diffusion/extensions/freeu.py index 0f6c47a773..6ec4fea3fa 100644 --- a/invokeai/backend/stable_diffusion/extensions/freeu.py +++ b/invokeai/backend/stable_diffusion/extensions/freeu.py @@ -15,28 +15,21 @@ if TYPE_CHECKING: class FreeUExt(ExtensionBase): def __init__( self, - freeu_config: Optional[FreeUConfig], + freeu_config: FreeUConfig, ): super().__init__() - self.freeu_config = freeu_config + self._freeu_config = freeu_config @contextmanager def patch_unet(self, unet: UNet2DConditionModel, cached_weights: Optional[Dict[str, torch.Tensor]] = None): - did_apply_freeu = False + unet.enable_freeu( + b1=self._freeu_config.b1, + b2=self._freeu_config.b2, + s1=self._freeu_config.s1, + s2=self._freeu_config.s2, + ) + try: - assert hasattr(unet, "enable_freeu") # mypy doesn't pick up this attribute? - if self.freeu_config is not None: - unet.enable_freeu( - b1=self.freeu_config.b1, - b2=self.freeu_config.b2, - s1=self.freeu_config.s1, - s2=self.freeu_config.s2, - ) - did_apply_freeu = True - yield - finally: - assert hasattr(unet, "disable_freeu") # mypy doesn't pick up this attribute? - if did_apply_freeu: - unet.disable_freeu() + unet.disable_freeu() From e92af52fb8bf227c8fca6a94265aba39035613ad Mon Sep 17 00:00:00 2001 From: chainchompa Date: Mon, 22 Jul 2024 16:11:36 -0400 Subject: [PATCH 09/61] fix moving items to uncategorized updating --- invokeai/frontend/web/src/services/api/endpoints/images.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/invokeai/frontend/web/src/services/api/endpoints/images.ts b/invokeai/frontend/web/src/services/api/endpoints/images.ts index 06ec9fc162..6f36866dce 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/images.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/images.ts @@ -417,6 +417,7 @@ export const imagesApi = api.injectEndpoints({ type: 'BoardImagesTotal', id: imageDTO.board_id ?? 'none', }, + { type: 'BoardImagesTotal', id: 'none' }, ]; }, }), From 43b3e242b0b61ba74c2733bd19ba156b47c59590 Mon Sep 17 00:00:00 2001 From: Mary Hipp Date: Tue, 16 Jul 2024 15:58:40 -0400 Subject: [PATCH 10/61] tidy(ui): refactor parameters panel components to be 1:1 with tabs --- .../src/features/ui/components/InvokeTabs.tsx | 37 +++++++++---------- .../ParametersPanelCanvas.tsx} | 10 ++--- .../ParametersPanelTextToImage.tsx | 0 .../ui/components/tabs/UpscalingTab.tsx | 15 ++++++++ .../web/src/features/ui/store/tabMap.tsx | 2 +- 5 files changed, 37 insertions(+), 27 deletions(-) rename invokeai/frontend/web/src/features/ui/components/{ParametersPanel.tsx => ParametersPanels/ParametersPanelCanvas.tsx} (86%) rename invokeai/frontend/web/src/features/ui/components/{ => ParametersPanels}/ParametersPanelTextToImage.tsx (100%) create mode 100644 invokeai/frontend/web/src/features/ui/components/tabs/UpscalingTab.tsx diff --git a/invokeai/frontend/web/src/features/ui/components/InvokeTabs.tsx b/invokeai/frontend/web/src/features/ui/components/InvokeTabs.tsx index 1968c64161..2774079ed4 100644 --- a/invokeai/frontend/web/src/features/ui/components/InvokeTabs.tsx +++ b/invokeai/frontend/web/src/features/ui/components/InvokeTabs.tsx @@ -11,7 +11,7 @@ import StatusIndicator from 'features/system/components/StatusIndicator'; import { selectConfigSlice } from 'features/system/store/configSlice'; import FloatingGalleryButton from 'features/ui/components/FloatingGalleryButton'; import FloatingParametersPanelButtons from 'features/ui/components/FloatingParametersPanelButtons'; -import ParametersPanelTextToImage from 'features/ui/components/ParametersPanelTextToImage'; +import ParametersPanelTextToImage from 'features/ui/components/ParametersPanels/ParametersPanelTextToImage'; import ModelManagerTab from 'features/ui/components/tabs/ModelManagerTab'; import NodesTab from 'features/ui/components/tabs/NodesTab'; import QueueTab from 'features/ui/components/tabs/QueueTab'; @@ -28,19 +28,22 @@ import type { CSSProperties, MouseEvent, ReactElement, ReactNode } from 'react'; import { memo, useCallback, useMemo, useRef } from 'react'; import { useHotkeys } from 'react-hotkeys-hook'; import { useTranslation } from 'react-i18next'; +import { MdZoomOutMap } from 'react-icons/md'; import { PiFlowArrowBold } from 'react-icons/pi'; import { RiBox2Line, RiBrushLine, RiInputMethodLine, RiPlayList2Fill } from 'react-icons/ri'; import type { ImperativePanelGroupHandle } from 'react-resizable-panels'; import { Panel, PanelGroup } from 'react-resizable-panels'; -import ParametersPanel from './ParametersPanel'; +import ParametersPanelCanvas from './ParametersPanels/ParametersPanelCanvas'; import ResizeHandle from './tabs/ResizeHandle'; +import UpscalingTab from './tabs/UpscalingTab'; type TabData = { id: InvokeTabName; translationKey: string; icon: ReactElement; content: ReactNode; + parametersPanel?: ReactNode; }; const TAB_DATA: Record = { @@ -49,18 +52,27 @@ const TAB_DATA: Record = { translationKey: 'ui.tabs.generation', icon: , content: , + parametersPanel: , }, canvas: { id: 'canvas', translationKey: 'ui.tabs.canvas', icon: , content: , + parametersPanel: , + }, + upscaling: { + id: 'upscaling', + translationKey: 'ui.tabs.upscaling', + icon: , + content: , }, workflows: { id: 'workflows', translationKey: 'ui.tabs.workflows', icon: , content: , + parametersPanel: , }, models: { id: 'models', @@ -81,7 +93,6 @@ const enabledTabsSelector = createMemoizedSelector(selectConfigSlice, (config) = ); const NO_GALLERY_PANEL_TABS: InvokeTabName[] = ['models', 'queue']; -const NO_OPTIONS_PANEL_TABS: InvokeTabName[] = ['models', 'queue']; const panelStyles: CSSProperties = { height: '100%', width: '100%' }; const GALLERY_MIN_SIZE_PX = 310; const GALLERY_MIN_SIZE_PCT = 20; @@ -103,7 +114,6 @@ const InvokeTabs = () => { e.target.blur(); } }, []); - const shouldShowOptionsPanel = useMemo(() => !NO_OPTIONS_PANEL_TABS.includes(activeTabName), [activeTabName]); const shouldShowGalleryPanel = useMemo(() => !NO_GALLERY_PANEL_TABS.includes(activeTabName), [activeTabName]); const tabs = useMemo( @@ -232,7 +242,7 @@ const InvokeTabs = () => { style={panelStyles} storage={panelStorage} > - {shouldShowOptionsPanel && ( + {!!TAB_DATA[activeTabName].parametersPanel && ( <> { onExpand={optionsPanel.onExpand} collapsible > - + {TAB_DATA[activeTabName].parametersPanel} { )} - {shouldShowOptionsPanel && } + {!!TAB_DATA[activeTabName].parametersPanel && } {shouldShowGalleryPanel && } ); }; export default memo(InvokeTabs); - -const ParametersPanelComponent = memo(() => { - const activeTabName = useAppSelector(activeTabNameSelector); - - if (activeTabName === 'workflows') { - return ; - } - if (activeTabName === 'generation') { - return ; - } - return ; -}); -ParametersPanelComponent.displayName = 'ParametersPanelComponent'; diff --git a/invokeai/frontend/web/src/features/ui/components/ParametersPanel.tsx b/invokeai/frontend/web/src/features/ui/components/ParametersPanels/ParametersPanelCanvas.tsx similarity index 86% rename from invokeai/frontend/web/src/features/ui/components/ParametersPanel.tsx rename to invokeai/frontend/web/src/features/ui/components/ParametersPanels/ParametersPanelCanvas.tsx index e8f73fd786..622ed96696 100644 --- a/invokeai/frontend/web/src/features/ui/components/ParametersPanel.tsx +++ b/invokeai/frontend/web/src/features/ui/components/ParametersPanels/ParametersPanelCanvas.tsx @@ -10,7 +10,6 @@ import { ControlSettingsAccordion } from 'features/settingsAccordions/components import { GenerationSettingsAccordion } from 'features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion'; import { ImageSettingsAccordion } from 'features/settingsAccordions/components/ImageSettingsAccordion/ImageSettingsAccordion'; import { RefinerSettingsAccordion } from 'features/settingsAccordions/components/RefinerSettingsAccordion/RefinerSettingsAccordion'; -import { activeTabNameSelector } from 'features/ui/store/uiSelectors'; import { OverlayScrollbarsComponent } from 'overlayscrollbars-react'; import type { CSSProperties } from 'react'; import { memo } from 'react'; @@ -20,8 +19,7 @@ const overlayScrollbarsStyles: CSSProperties = { width: '100%', }; -const ParametersPanel = () => { - const activeTabName = useAppSelector(activeTabNameSelector); +const ParametersPanelCanvas = () => { const isSDXL = useAppSelector((s) => s.generation.model?.base === 'sdxl'); return ( @@ -34,8 +32,8 @@ const ParametersPanel = () => { {isSDXL ? : } - {activeTabName !== 'generation' && } - {activeTabName === 'canvas' && } + + {isSDXL && } @@ -46,4 +44,4 @@ const ParametersPanel = () => { ); }; -export default memo(ParametersPanel); +export default memo(ParametersPanelCanvas); diff --git a/invokeai/frontend/web/src/features/ui/components/ParametersPanelTextToImage.tsx b/invokeai/frontend/web/src/features/ui/components/ParametersPanels/ParametersPanelTextToImage.tsx similarity index 100% rename from invokeai/frontend/web/src/features/ui/components/ParametersPanelTextToImage.tsx rename to invokeai/frontend/web/src/features/ui/components/ParametersPanels/ParametersPanelTextToImage.tsx diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/UpscalingTab.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/UpscalingTab.tsx new file mode 100644 index 0000000000..8d4c916c2a --- /dev/null +++ b/invokeai/frontend/web/src/features/ui/components/tabs/UpscalingTab.tsx @@ -0,0 +1,15 @@ +import { Box } from '@invoke-ai/ui-library'; +import { ImageViewer } from 'features/gallery/components/ImageViewer/ImageViewer'; +import { useImageViewer } from 'features/gallery/components/ImageViewer/useImageViewer'; +import { memo } from 'react'; + +const UpscalingTab = () => { + const imageViewer = useImageViewer(); + return ( + + {imageViewer.isOpen && } + + ); +}; + +export default memo(UpscalingTab); diff --git a/invokeai/frontend/web/src/features/ui/store/tabMap.tsx b/invokeai/frontend/web/src/features/ui/store/tabMap.tsx index 526a55b069..5cf97b2d3e 100644 --- a/invokeai/frontend/web/src/features/ui/store/tabMap.tsx +++ b/invokeai/frontend/web/src/features/ui/store/tabMap.tsx @@ -1,3 +1,3 @@ -export const TAB_NUMBER_MAP = ['generation', 'canvas', 'workflows', 'models', 'queue'] as const; +export const TAB_NUMBER_MAP = ['generation', 'canvas', 'upscaling', 'workflows', 'models', 'queue'] as const; export type InvokeTabName = (typeof TAB_NUMBER_MAP)[number]; From a0a54348e8e704a45fd5179c4f4d4d48199b7be3 Mon Sep 17 00:00:00 2001 From: Mary Hipp Date: Wed, 17 Jul 2024 10:50:33 -0400 Subject: [PATCH 11/61] removed upscale button, created spandrel model dropdown, created upscale initial image that works with dnd --- .../listeners/imageDropped.ts | 15 +++++ .../listeners/imageUploaded.ts | 10 ++++ .../listeners/modelsLoaded.ts | 25 ++++++++- invokeai/frontend/web/src/app/store/store.ts | 3 + .../src/common/hooks/useFullscreenDropzone.ts | 4 ++ .../web/src/features/dnd/types/index.ts | 19 +++++-- .../web/src/features/dnd/util/isValidDrop.ts | 2 + .../ImageViewer/CurrentImageButtons.tsx | 27 +-------- .../Upscale/ParamRealESRGANModel.tsx | 2 +- .../components/Upscale/ParamSpandrelModel.tsx | 47 ++++++++++++++++ .../Upscale/ParamUpscaleSettings.tsx | 9 ++- .../parameters/hooks/useIsAllowedToUpscale.ts | 1 + .../src/features/parameters/store/types.ts | 2 +- .../features/parameters/store/upscaleSlice.ts | 50 +++++++++++++++++ .../parameters/types/parameterSchemas.ts | 5 ++ .../UpscaleInitialImage.tsx | 55 +++++++++++++++++++ .../UpscaleSettingsAccordion.tsx | 27 +++++++++ .../src/features/ui/components/InvokeTabs.tsx | 2 + .../ParametersPanelUpscale.tsx | 44 +++++++++++++++ .../frontend/web/src/services/api/types.ts | 7 ++- 20 files changed, 317 insertions(+), 39 deletions(-) create mode 100644 invokeai/frontend/web/src/features/parameters/components/Upscale/ParamSpandrelModel.tsx create mode 100644 invokeai/frontend/web/src/features/parameters/store/upscaleSlice.ts create mode 100644 invokeai/frontend/web/src/features/settingsAccordions/components/UpscaleSettingsAccordion/UpscaleInitialImage.tsx create mode 100644 invokeai/frontend/web/src/features/settingsAccordions/components/UpscaleSettingsAccordion/UpscaleSettingsAccordion.tsx create mode 100644 invokeai/frontend/web/src/features/ui/components/ParametersPanels/ParametersPanelUpscale.tsx diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts index 0cd77dc2e7..eac8922a79 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageDropped.ts @@ -24,6 +24,7 @@ import { import { fieldImageValueChanged } from 'features/nodes/store/nodesSlice'; import { selectOptimalDimension } from 'features/parameters/store/generationSlice'; import { imagesApi } from 'services/api/endpoints/images'; +import { upscaleInitialImageChanged } from '../../../../../features/parameters/store/upscaleSlice'; export const dndDropped = createAction<{ overData: TypesafeDroppableData; @@ -243,6 +244,20 @@ export const addImageDroppedListener = (startAppListening: AppStartListening) => return; } + /** + * Image dropped on upscale initial image + */ + if ( + overData.actionType === 'SET_UPSCALE_INITIAL_IMAGE' && + activeData.payloadType === 'IMAGE_DTO' && + activeData.payload.imageDTO + ) { + const { imageDTO } = activeData.payload; + + dispatch(upscaleInitialImageChanged(imageDTO)); + return; + } + /** * Multiple images dropped on user board */ diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageUploaded.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageUploaded.ts index bd1ee47825..5a4d5dc078 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageUploaded.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/imageUploaded.ts @@ -19,6 +19,7 @@ import { t } from 'i18next'; import { omit } from 'lodash-es'; import { boardsApi } from 'services/api/endpoints/boards'; import { imagesApi } from 'services/api/endpoints/images'; +import { upscaleInitialImageChanged } from '../../../../../features/parameters/store/upscaleSlice'; export const addImageUploadedFulfilledListener = (startAppListening: AppStartListening) => { startAppListening({ @@ -89,6 +90,15 @@ export const addImageUploadedFulfilledListener = (startAppListening: AppStartLis return; } + if (postUploadAction?.type === 'SET_UPSCALE_INITIAL_IMAGE') { + dispatch(upscaleInitialImageChanged(imageDTO)); + toast({ + ...DEFAULT_UPLOADED_TOAST, + description: "set as upscale initial image", + }); + return; + } + if (postUploadAction?.type === 'SET_CONTROL_ADAPTER_IMAGE') { const { id } = postUploadAction; dispatch( diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts index eb86f54c84..eef1f8f49f 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts @@ -17,7 +17,8 @@ import { forEach } from 'lodash-es'; import type { Logger } from 'roarr'; import { modelConfigsAdapterSelectors, modelsApi } from 'services/api/endpoints/models'; import type { AnyModelConfig } from 'services/api/types'; -import { isNonRefinerMainModelConfig, isRefinerMainModelModelConfig, isVAEModelConfig } from 'services/api/types'; +import { isNonRefinerMainModelConfig, isRefinerMainModelModelConfig, isSpandrelImageToImageModelConfig, isVAEModelConfig } from 'services/api/types'; +import { upscaleModelChanged } from '../../../../../features/parameters/store/upscaleSlice'; export const addModelsLoadedListener = (startAppListening: AppStartListening) => { startAppListening({ @@ -36,6 +37,7 @@ export const addModelsLoadedListener = (startAppListening: AppStartListening) => handleVAEModels(models, state, dispatch, log); handleLoRAModels(models, state, dispatch, log); handleControlAdapterModels(models, state, dispatch, log); + handleSpandrelImageToImageModels(models, state, dispatch, log); }, }); }; @@ -177,3 +179,24 @@ const handleControlAdapterModels: ModelHandler = (models, state, dispatch, _log) dispatch(controlAdapterModelCleared({ id: ca.id })); }); }; + +const handleSpandrelImageToImageModels: ModelHandler = (models, state, dispatch, _log) => { + const currentUpscaleModel = state.upscale.upscaleModel; + const upscaleModels = models.filter(isSpandrelImageToImageModelConfig); + + if (currentUpscaleModel) { + const isCurrentUpscaleModelAvailable = upscaleModels.some((m) => m.key === currentUpscaleModel.key); + if (isCurrentUpscaleModelAvailable) { + return; + } + } + + const firstModel = upscaleModels[0]; + if (firstModel) { + dispatch(upscaleModelChanged(firstModel)) + return + } + + dispatch(upscaleModelChanged(null)) + +}; diff --git a/invokeai/frontend/web/src/app/store/store.ts b/invokeai/frontend/web/src/app/store/store.ts index 062cdc1cbf..804d3629ab 100644 --- a/invokeai/frontend/web/src/app/store/store.ts +++ b/invokeai/frontend/web/src/app/store/store.ts @@ -46,6 +46,7 @@ import { actionSanitizer } from './middleware/devtools/actionSanitizer'; import { actionsDenylist } from './middleware/devtools/actionsDenylist'; import { stateSanitizer } from './middleware/devtools/stateSanitizer'; import { listenerMiddleware } from './middleware/listenerMiddleware'; +import { upscalePersistConfig, upscaleSlice } from '../../features/parameters/store/upscaleSlice'; const allReducers = { [canvasSlice.name]: canvasSlice.reducer, @@ -69,6 +70,7 @@ const allReducers = { [controlLayersSlice.name]: undoable(controlLayersSlice.reducer, controlLayersUndoableConfig), [workflowSettingsSlice.name]: workflowSettingsSlice.reducer, [api.reducerPath]: api.reducer, + [upscaleSlice.name]: upscaleSlice.reducer }; const rootReducer = combineReducers(allReducers); @@ -114,6 +116,7 @@ const persistConfigs: { [key in keyof typeof allReducers]?: PersistConfig } = { [hrfPersistConfig.name]: hrfPersistConfig, [controlLayersPersistConfig.name]: controlLayersPersistConfig, [workflowSettingsPersistConfig.name]: workflowSettingsPersistConfig, + [upscalePersistConfig.name]: upscalePersistConfig }; const unserialize: UnserializeFunction = (data, key) => { diff --git a/invokeai/frontend/web/src/common/hooks/useFullscreenDropzone.ts b/invokeai/frontend/web/src/common/hooks/useFullscreenDropzone.ts index 5b1bf1f5b3..d8e7d70a8c 100644 --- a/invokeai/frontend/web/src/common/hooks/useFullscreenDropzone.ts +++ b/invokeai/frontend/web/src/common/hooks/useFullscreenDropzone.ts @@ -21,6 +21,10 @@ const selectPostUploadAction = createMemoizedSelector(activeTabNameSelector, (ac postUploadAction = { type: 'SET_CANVAS_INITIAL_IMAGE' }; } + if (activeTabName === 'upscaling') { + postUploadAction = { type: 'SET_UPSCALE_INITIAL_IMAGE' }; + } + return postUploadAction; }); diff --git a/invokeai/frontend/web/src/features/dnd/types/index.ts b/invokeai/frontend/web/src/features/dnd/types/index.ts index 6fcf18421e..1e72123b75 100644 --- a/invokeai/frontend/web/src/features/dnd/types/index.ts +++ b/invokeai/frontend/web/src/features/dnd/types/index.ts @@ -62,6 +62,10 @@ export type CanvasInitialImageDropData = BaseDropData & { actionType: 'SET_CANVAS_INITIAL_IMAGE'; }; +export type UpscaleInitialImageDropData = BaseDropData & { + actionType: 'SET_UPSCALE_INITIAL_IMAGE'; +}; + type NodesImageDropData = BaseDropData & { actionType: 'SET_NODES_IMAGE'; context: { @@ -87,6 +91,8 @@ export type SelectForCompareDropData = BaseDropData & { }; }; + + export type TypesafeDroppableData = | CurrentImageDropData | ControlAdapterDropData @@ -98,7 +104,8 @@ export type TypesafeDroppableData = | IPALayerImageDropData | RGLayerIPAdapterImageDropData | IILayerImageDropData - | SelectForCompareDropData; + | SelectForCompareDropData + | UpscaleInitialImageDropData; type BaseDragData = { id: string; @@ -159,11 +166,11 @@ interface DragEvent { over: TypesafeOver | null; } -export interface DragStartEvent extends Pick {} -interface DragMoveEvent extends DragEvent {} -interface DragOverEvent extends DragMoveEvent {} -export interface DragEndEvent extends DragEvent {} -interface DragCancelEvent extends DragEndEvent {} +export interface DragStartEvent extends Pick { } +interface DragMoveEvent extends DragEvent { } +interface DragOverEvent extends DragMoveEvent { } +export interface DragEndEvent extends DragEvent { } +interface DragCancelEvent extends DragEndEvent { } export interface DndContextTypesafeProps extends Omit { diff --git a/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts b/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts index 6dec862345..3f8fe5ab73 100644 --- a/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts +++ b/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts @@ -27,6 +27,8 @@ export const isValidDrop = (overData?: TypesafeDroppableData | null, activeData? return payloadType === 'IMAGE_DTO'; case 'SET_CANVAS_INITIAL_IMAGE': return payloadType === 'IMAGE_DTO'; + case 'SET_UPSCALE_INITIAL_IMAGE': + return payloadType === 'IMAGE_DTO'; case 'SET_NODES_IMAGE': return payloadType === 'IMAGE_DTO'; case 'SELECT_FOR_COMPARE': diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImageButtons.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImageButtons.tsx index d500d692fe..da0253cc36 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImageButtons.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImageButtons.tsx @@ -11,12 +11,8 @@ import SingleSelectionMenuItems from 'features/gallery/components/ImageContextMe import { useImageActions } from 'features/gallery/hooks/useImageActions'; import { sentImageToImg2Img } from 'features/gallery/store/actions'; import { selectLastSelectedImage } from 'features/gallery/store/gallerySelectors'; -import { selectGallerySlice } from 'features/gallery/store/gallerySlice'; import { parseAndRecallImageDimensions } from 'features/metadata/util/handlers'; import { $templates } from 'features/nodes/store/nodesSlice'; -import ParamUpscalePopover from 'features/parameters/components/Upscale/ParamUpscaleSettings'; -import { useIsQueueMutationInProgress } from 'features/queue/hooks/useIsQueueMutationInProgress'; -import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus'; import { selectSystemSlice } from 'features/system/store/systemSlice'; import { setActiveTab } from 'features/ui/store/uiSlice'; import { useGetAndLoadEmbeddedWorkflow } from 'features/workflowLibrary/hooks/useGetAndLoadEmbeddedWorkflow'; @@ -37,9 +33,8 @@ import { useGetImageDTOQuery } from 'services/api/endpoints/images'; const selectShouldDisableToolbarButtons = createSelector( selectSystemSlice, - selectGallerySlice, selectLastSelectedImage, - (system, gallery, lastSelectedImage) => { + (system, lastSelectedImage) => { const hasProgressImage = Boolean(system.denoiseProgress?.progress_image); return hasProgressImage || !lastSelectedImage; } @@ -47,13 +42,10 @@ const selectShouldDisableToolbarButtons = createSelector( const CurrentImageButtons = () => { const dispatch = useAppDispatch(); - const isConnected = useAppSelector((s) => s.system.isConnected); const lastSelectedImage = useAppSelector(selectLastSelectedImage); const selection = useAppSelector((s) => s.gallery.selection); const shouldDisableToolbarButtons = useAppSelector(selectShouldDisableToolbarButtons); const templates = useStore($templates); - const isUpscalingEnabled = useFeatureStatus('upscaling'); - const isQueueMutationInProgress = useIsQueueMutationInProgress(); const { t } = useTranslation(); const { currentData: imageDTO } = useGetImageDTOQuery(lastSelectedImage?.image_name ?? skipToken); @@ -107,17 +99,6 @@ const CurrentImageButtons = () => { dispatch(imagesToDeleteSelected(selection)); }, [dispatch, imageDTO, selection]); - useHotkeys( - 'Shift+U', - () => { - handleClickUpscale(); - }, - { - enabled: () => Boolean(isUpscalingEnabled && !shouldDisableToolbarButtons && isConnected), - }, - [isUpscalingEnabled, imageDTO, shouldDisableToolbarButtons, isConnected] - ); - useHotkeys( 'delete', () => { @@ -191,12 +172,6 @@ const CurrentImageButtons = () => { /> - {isUpscalingEnabled && ( - - {isUpscalingEnabled && } - - )} - diff --git a/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamRealESRGANModel.tsx b/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamRealESRGANModel.tsx index 6c1a3ab3a7..d02bfd2b03 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamRealESRGANModel.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamRealESRGANModel.tsx @@ -63,7 +63,7 @@ const ParamESRGANModel = () => { return ( - {t('models.esrganModel')} + {t('models.esrganModel')} ); diff --git a/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamSpandrelModel.tsx b/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamSpandrelModel.tsx new file mode 100644 index 0000000000..b16a77feae --- /dev/null +++ b/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamSpandrelModel.tsx @@ -0,0 +1,47 @@ +import { Combobox, FormControl, FormLabel } from '@invoke-ai/ui-library'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { memo, useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useSpandrelImageToImageModels } from '../../../../services/api/hooks/modelsByType'; +import { useModelCombobox } from '../../../../common/hooks/useModelCombobox'; +import { SpandrelImageToImageModelConfig } from '../../../../services/api/types'; +import { upscaleModelChanged } from '../../store/upscaleSlice'; + +const ParamSpandrelModel = () => { + const { t } = useTranslation(); + const [modelConfigs, { isLoading }] = useSpandrelImageToImageModels(); + + const model = useAppSelector((s) => s.upscale.upscaleModel); + + const dispatch = useAppDispatch(); + + const _onChange = useCallback( + (v: SpandrelImageToImageModelConfig | null) => { + dispatch(upscaleModelChanged(v)); + }, + [dispatch] + ); + + const { options, value, onChange, placeholder, noOptionsMessage } = useModelCombobox({ + modelConfigs, + onChange: _onChange, + selectedModel: model, + isLoading, + }); + + return ( + + Upscale Model + + + ); +}; + +export default memo(ParamSpandrelModel); diff --git a/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamUpscaleSettings.tsx b/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamUpscaleSettings.tsx index c0309bebe4..15fec1d214 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamUpscaleSettings.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Upscale/ParamUpscaleSettings.tsx @@ -17,7 +17,8 @@ import { useTranslation } from 'react-i18next'; import { PiFrameCornersBold } from 'react-icons/pi'; import type { ImageDTO } from 'services/api/types'; -import ParamESRGANModel from './ParamRealESRGANModel'; +import ParamSpandrelModel from './ParamSpandrelModel'; +import { useSpandrelImageToImageModels } from '../../../../services/api/hooks/modelsByType'; type Props = { imageDTO?: ImageDTO }; @@ -28,6 +29,7 @@ const ParamUpscalePopover = (props: Props) => { const { t } = useTranslation(); const { isOpen, onOpen, onClose } = useDisclosure(); const { isAllowedToUpscale, detail } = useIsAllowedToUpscale(imageDTO); + const [modelConfigs] = useSpandrelImageToImageModels(); const handleClickUpscale = useCallback(() => { onClose(); @@ -45,16 +47,17 @@ const ParamUpscalePopover = (props: Props) => { onClick={onOpen} icon={} aria-label={t('parameters.upscale')} + isDisabled={!modelConfigs.length} /> - +