From aca9d7448940e5eb25b09c8d2df2d5520efe88b0 Mon Sep 17 00:00:00 2001 From: Kevin Turner <83819+keturn@users.noreply.github.com> Date: Sun, 19 Feb 2023 15:33:16 -0800 Subject: [PATCH 1/4] refactor(InvokeAIDiffuserComponent): rename internal methods Prefix with `_` as is tradition. --- .../diffusion/shared_invokeai_diffusion.py | 31 ++++++++++++------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index ca3e608fc0..d9163468d2 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -1,4 +1,3 @@ -import math from contextlib import contextmanager from dataclasses import dataclass from math import ceil @@ -6,8 +5,8 @@ from typing import Callable, Optional, Union, Any, Dict import numpy as np import torch - from diffusers.models.cross_attention import AttnProcessor + from ldm.models.diffusion.cross_attention_control import Arguments, \ restore_default_cross_attention, override_cross_attention, Context, get_cross_attention_modules, \ CrossAttentionType, SwapCrossAttnContext @@ -143,11 +142,16 @@ class InvokeAIDiffuserComponent: wants_hybrid_conditioning = isinstance(conditioning, dict) if wants_hybrid_conditioning: - unconditioned_next_x, conditioned_next_x = self.apply_hybrid_conditioning(x, sigma, unconditioning, conditioning) + unconditioned_next_x, conditioned_next_x = self._apply_hybrid_conditioning(x, sigma, unconditioning, + conditioning) elif wants_cross_attention_control: - unconditioned_next_x, conditioned_next_x = self.apply_cross_attention_controlled_conditioning(x, sigma, unconditioning, conditioning, cross_attention_control_types_to_do) + unconditioned_next_x, conditioned_next_x = self._apply_cross_attention_controlled_conditioning(x, sigma, + unconditioning, + conditioning, + cross_attention_control_types_to_do) else: - unconditioned_next_x, conditioned_next_x = self.apply_standard_conditioning(x, sigma, unconditioning, conditioning) + unconditioned_next_x, conditioned_next_x = self._apply_standard_conditioning(x, sigma, unconditioning, + conditioning) combined_next_x = self._combine(unconditioned_next_x, conditioned_next_x, unconditional_guidance_scale) @@ -181,7 +185,7 @@ class InvokeAIDiffuserComponent: # methods below are called from do_diffusion_step and should be considered private to this class. - def apply_standard_conditioning(self, x, sigma, unconditioning, conditioning): + def _apply_standard_conditioning(self, x, sigma, unconditioning, conditioning): # fast batched path x_twice = torch.cat([x] * 2) sigma_twice = torch.cat([sigma] * 2) @@ -194,7 +198,7 @@ class InvokeAIDiffuserComponent: return unconditioned_next_x, conditioned_next_x - def apply_hybrid_conditioning(self, x, sigma, unconditioning, conditioning): + def _apply_hybrid_conditioning(self, x, sigma, unconditioning, conditioning): assert isinstance(conditioning, dict) assert isinstance(unconditioning, dict) x_twice = torch.cat([x] * 2) @@ -212,18 +216,21 @@ class InvokeAIDiffuserComponent: return unconditioned_next_x, conditioned_next_x - def apply_cross_attention_controlled_conditioning(self, + def _apply_cross_attention_controlled_conditioning(self, x: torch.Tensor, sigma, unconditioning, conditioning, cross_attention_control_types_to_do): if self.is_running_diffusers: - return self.apply_cross_attention_controlled_conditioning__diffusers(x, sigma, unconditioning, conditioning, cross_attention_control_types_to_do) + return self._apply_cross_attention_controlled_conditioning__diffusers(x, sigma, unconditioning, + conditioning, + cross_attention_control_types_to_do) else: - return self.apply_cross_attention_controlled_conditioning__compvis(x, sigma, unconditioning, conditioning, cross_attention_control_types_to_do) + return self._apply_cross_attention_controlled_conditioning__compvis(x, sigma, unconditioning, conditioning, + cross_attention_control_types_to_do) - def apply_cross_attention_controlled_conditioning__diffusers(self, + def _apply_cross_attention_controlled_conditioning__diffusers(self, x: torch.Tensor, sigma, unconditioning, @@ -246,7 +253,7 @@ class InvokeAIDiffuserComponent: return unconditioned_next_x, conditioned_next_x - def apply_cross_attention_controlled_conditioning__compvis(self, x:torch.Tensor, sigma, unconditioning, conditioning, cross_attention_control_types_to_do): + def _apply_cross_attention_controlled_conditioning__compvis(self, x:torch.Tensor, sigma, unconditioning, conditioning, cross_attention_control_types_to_do): # print('pct', percent_through, ': doing cross attention control on', cross_attention_control_types_to_do) # slower non-batched path (20% slower on mac MPS) # We are only interested in using attention maps for conditioned_next_x, but batching them with generation of From d0abe13b6060d07e7946df16a64686f9d359ac94 Mon Sep 17 00:00:00 2001 From: Kevin Turner <83819+keturn@users.noreply.github.com> Date: Sun, 19 Feb 2023 16:04:54 -0800 Subject: [PATCH 2/4] performance(InvokeAIDiffuserComponent): add low-memory path for calculating conditioned and unconditioned predictions sequentially Proof of concept. Still needs to be wired up to options or heuristics. --- .../diffusion/shared_invokeai_diffusion.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index d9163468d2..2e513d3f5a 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -29,7 +29,7 @@ class InvokeAIDiffuserComponent: * Hybrid conditioning (used for inpainting) ''' debug_thresholding = False - + sequential_conditioning = False @dataclass class ExtraConditioningInfo: @@ -149,9 +149,13 @@ class InvokeAIDiffuserComponent: unconditioning, conditioning, cross_attention_control_types_to_do) + elif self.sequential_conditioning: + unconditioned_next_x, conditioned_next_x = self._apply_standard_conditioning_sequentially( + x, sigma, unconditioning, conditioning) + else: - unconditioned_next_x, conditioned_next_x = self._apply_standard_conditioning(x, sigma, unconditioning, - conditioning) + unconditioned_next_x, conditioned_next_x = self._apply_standard_conditioning( + x, sigma, unconditioning, conditioning) combined_next_x = self._combine(unconditioned_next_x, conditioned_next_x, unconditional_guidance_scale) @@ -198,6 +202,16 @@ class InvokeAIDiffuserComponent: return unconditioned_next_x, conditioned_next_x + def _apply_standard_conditioning_sequentially(self, x: torch.Tensor, sigma, unconditioning: torch.Tensor, conditioning: torch.Tensor): + # low-memory sequential path + unconditioned_next_x = self.model_forward_callback(x, sigma, unconditioning) + conditioned_next_x = self.model_forward_callback(x, sigma, conditioning) + if conditioned_next_x.device.type == 'mps': + # prevent a result filled with zeros. seems to be a torch bug. + conditioned_next_x = conditioned_next_x.clone() + return unconditioned_next_x, conditioned_next_x + + def _apply_hybrid_conditioning(self, x, sigma, unconditioning, conditioning): assert isinstance(conditioning, dict) assert isinstance(unconditioning, dict) From 6c8d4b091eb0a8752b252905dd471a298c4dcb9a Mon Sep 17 00:00:00 2001 From: Kevin Turner <83819+keturn@users.noreply.github.com> Date: Sun, 19 Feb 2023 16:58:54 -0800 Subject: [PATCH 3/4] dev(InvokeAIDiffuserComponent): mollify type checker's concern about the optional argument --- ldm/models/diffusion/shared_invokeai_diffusion.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index 2e513d3f5a..d962932484 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -6,12 +6,18 @@ from typing import Callable, Optional, Union, Any, Dict import numpy as np import torch from diffusers.models.cross_attention import AttnProcessor +from typing_extensions import TypeAlias from ldm.models.diffusion.cross_attention_control import Arguments, \ restore_default_cross_attention, override_cross_attention, Context, get_cross_attention_modules, \ CrossAttentionType, SwapCrossAttnContext from ldm.models.diffusion.cross_attention_map_saving import AttentionMapSaver +ModelForwardCallback: TypeAlias = Union[ + # x, t, conditioning, Optional[cross-attention kwargs] + Callable[[torch.Tensor, torch.Tensor, torch.Tensor, Optional[dict[str, Any]]], torch.Tensor], + Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] +] @dataclass(frozen=True) class PostprocessingSettings: @@ -42,8 +48,7 @@ class InvokeAIDiffuserComponent: return self.cross_attention_control_args is not None - def __init__(self, model, model_forward_callback: - Callable[[torch.Tensor, torch.Tensor, torch.Tensor, Optional[dict[str,Any]]], torch.Tensor], + def __init__(self, model, model_forward_callback: ModelForwardCallback, is_running_diffusers: bool=False, ): """ From 2dded6826742b7b3be3d09ec1d64231e632e24c3 Mon Sep 17 00:00:00 2001 From: Kevin Turner <83819+keturn@users.noreply.github.com> Date: Sun, 19 Feb 2023 21:21:14 -0800 Subject: [PATCH 4/4] add `--sequential_guidance` option for low-RAM tradeoff --- ldm/invoke/CLI.py | 7 ++++--- ldm/invoke/args.py | 17 ++++++++++++----- ldm/invoke/globals.py | 5 ++++- .../diffusion/shared_invokeai_diffusion.py | 6 ++++-- 4 files changed, 24 insertions(+), 11 deletions(-) diff --git a/ldm/invoke/CLI.py b/ldm/invoke/CLI.py index d639c16640..52d3b33304 100644 --- a/ldm/invoke/CLI.py +++ b/ldm/invoke/CLI.py @@ -60,6 +60,7 @@ def main(): Globals.always_use_cpu = args.always_use_cpu Globals.internet_available = args.internet_available and check_internet() Globals.disable_xformers = not args.xformers + Globals.sequential_guidance = args.sequential_guidance Globals.ckpt_convert = args.ckpt_convert print(f">> Internet connectivity is {Globals.internet_available}") @@ -749,7 +750,7 @@ def import_ckpt_model( base_name = Path(url_attachment_name(path_or_url)).name if is_a_url else Path(path_or_url).name default_name = Path(base_name).stem default_description = f"Imported model {default_name}" - + model_name, model_description = _get_model_name_and_desc( manager, completer, @@ -834,7 +835,7 @@ def _ask_for_config_file(model_path: Union[str,Path], completer, plural: bool=Fa '2': 'v2-inference-v.yaml', '3': 'v1-inpainting-inference.yaml', } - + prompt = '''What type of models are these?: [1] Models based on Stable Diffusion 1.X [2] Models based on Stable Diffusion 2.X @@ -843,7 +844,7 @@ def _ask_for_config_file(model_path: Union[str,Path], completer, plural: bool=Fa [1] A model based on Stable Diffusion 1.X [2] A model based on Stable Diffusion 2.X [3] An inpainting models based on Stable Diffusion 1.X -[4] Something else''' +[4] Something else''' print(prompt) choice = input(f'Your choice: [{default}] ') choice = choice.strip() or default diff --git a/ldm/invoke/args.py b/ldm/invoke/args.py index 1bd1aa46ab..7fa31a0df1 100644 --- a/ldm/invoke/args.py +++ b/ldm/invoke/args.py @@ -91,14 +91,14 @@ import pydoc import re import shlex import sys -import ldm.invoke -import ldm.invoke.pngwriter - -from ldm.invoke.globals import Globals -from ldm.invoke.prompt_parser import split_weighted_subprompts from argparse import Namespace from pathlib import Path +import ldm.invoke +import ldm.invoke.pngwriter +from ldm.invoke.globals import Globals +from ldm.invoke.prompt_parser import split_weighted_subprompts + APP_ID = ldm.invoke.__app_id__ APP_NAME = ldm.invoke.__app_name__ APP_VERSION = ldm.invoke.__version__ @@ -484,6 +484,13 @@ class Args(object): action='store_true', help='Force free gpu memory before final decoding', ) + model_group.add_argument( + '--sequential_guidance', + dest='sequential_guidance', + action='store_true', + help="Calculate guidance in serial instead of in parallel, lowering memory requirement " + "at the expense of speed", + ) model_group.add_argument( '--xformers', action=argparse.BooleanOptionalAction, diff --git a/ldm/invoke/globals.py b/ldm/invoke/globals.py index d00f254da5..eeb69d266b 100644 --- a/ldm/invoke/globals.py +++ b/ldm/invoke/globals.py @@ -13,8 +13,8 @@ the attributes: import os import os.path as osp -from pathlib import Path from argparse import Namespace +from pathlib import Path from typing import Union Globals = Namespace() @@ -48,6 +48,9 @@ Globals.internet_available = True # Whether to disable xformers Globals.disable_xformers = False +# Low-memory tradeoff for guidance calculations. +Globals.sequential_guidance = False + # whether we are forcing full precision Globals.full_precision = False diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index d962932484..d80a3d8bc7 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -8,6 +8,7 @@ import torch from diffusers.models.cross_attention import AttnProcessor from typing_extensions import TypeAlias +from ldm.invoke.globals import Globals from ldm.models.diffusion.cross_attention_control import Arguments, \ restore_default_cross_attention, override_cross_attention, Context, get_cross_attention_modules, \ CrossAttentionType, SwapCrossAttnContext @@ -35,7 +36,7 @@ class InvokeAIDiffuserComponent: * Hybrid conditioning (used for inpainting) ''' debug_thresholding = False - sequential_conditioning = False + sequential_guidance = False @dataclass class ExtraConditioningInfo: @@ -60,6 +61,7 @@ class InvokeAIDiffuserComponent: self.is_running_diffusers = is_running_diffusers self.model_forward_callback = model_forward_callback self.cross_attention_control_context = None + self.sequential_guidance = Globals.sequential_guidance @contextmanager def custom_attention_context(self, @@ -154,7 +156,7 @@ class InvokeAIDiffuserComponent: unconditioning, conditioning, cross_attention_control_types_to_do) - elif self.sequential_conditioning: + elif self.sequential_guidance: unconditioned_next_x, conditioned_next_x = self._apply_standard_conditioning_sequentially( x, sigma, unconditioning, conditioning)