From b22e3d66fcd6c8ddb541506369923ccb73ed0617 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 17 Jan 2025 09:18:58 -0500 Subject: [PATCH 01/66] update modernui Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 ++- extensions-builtin/sdnext-modernui | 2 +- modules/shared.py | 2 +- wiki | 2 +- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b333b2e24..2be3a84f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-16 +## Update for 2025-01-17 - **Gallery**: - add http fallback for slow/unreliable links @@ -10,6 +10,7 @@ - sana vae tiling - increase gallery timeouts - update ui element ids + - modernui use local font ## Update for 2025-01-15 diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index c17c53928..d97615471 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit c17c5392816e3e6e212c7d10166260728e10d249 +Subproject commit d97615471c0121ff78c092bb85b620980686af38 diff --git a/modules/shared.py b/modules/shared.py index b092c27b5..016ed98df 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -357,7 +357,7 @@ def list_samplers(): def temp_disable_extensions(): - disable_safe = ['sd-webui-controlnet', 'multidiffusion-upscaler-for-automatic1111', 'a1111-sd-webui-lycoris', 'sd-webui-agent-scheduler', 'clip-interrogator-ext', 'stable-diffusion-webui-rembg', 'sd-extension-chainner', 'stable-diffusion-webui-images-browser'] + disable_safe = ['sd-webui-controlnet', 'multidiffusion-upscaler-for-automatic1111', 'a1111-sd-webui-lycoris', 'sd-webui-agent-scheduler', 'clip-interrogator-ext', 'stable-diffusion-webui-images-browser'] disable_diffusers = ['sd-webui-controlnet', 'multidiffusion-upscaler-for-automatic1111', 'a1111-sd-webui-lycoris', 'sd-webui-animatediff'] disable_themes = ['sd-webui-lobe-theme', 'cozy-nest', 'sdnext-modernui'] disable_original = [] diff --git a/wiki b/wiki index 7f072b554..6caea0521 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 7f072b554c6ee2edadc33879e5b4bdbfa48e6282 +Subproject commit 6caea0521a48bd5b99a18566f21797bff1f2c244 From bfe8ece74978714fc4b36e4f847eb9721ec99f05 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 17 Jan 2025 16:20:46 -0500 Subject: [PATCH 02/66] unique font family registration Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + extensions-builtin/sd-extension-system-info | 2 +- extensions-builtin/sdnext-modernui | 2 +- javascript/base.css | 2 +- javascript/black-gray.css | 2 +- javascript/black-orange.css | 2 +- javascript/black-teal-reimagined.css | 2 +- javascript/black-teal.css | 2 +- javascript/light-teal.css | 2 +- javascript/sdnext.css | 2 +- javascript/simple-dark.css | 2 +- javascript/simple-light.css | 2 +- modules/processing_vae.py | 2 +- modules/sd_checkpoint.py | 5 ++++- modules/shared.py | 4 ++-- 15 files changed, 19 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2be3a84f5..06854d872 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ - increase gallery timeouts - update ui element ids - modernui use local font + - unique font family registration ## Update for 2025-01-15 diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index aafc660ba..f55066aa4 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit aafc660ba3cf78724059917f7c3a05f3ac1ed46a +Subproject commit f55066aa419509fa0eafe92e4e98a4637b19b9c5 diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index d97615471..e0b3e5918 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit d97615471c0121ff78c092bb85b620980686af38 +Subproject commit e0b3e5918f49d090c8c340ca44e2c9e6d5bf5127 diff --git a/javascript/base.css b/javascript/base.css index cadfd1868..0a0b621c2 100644 --- a/javascript/base.css +++ b/javascript/base.css @@ -1,4 +1,4 @@ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } /* toolbutton */ .gradio-button.tool { max-width: min-content; min-width: min-content !important; align-self: end; font-size: 1.4em; color: var(--body-text-color) !important; } diff --git a/javascript/black-gray.css b/javascript/black-gray.css index a78305032..c262a3bf4 100644 --- a/javascript/black-gray.css +++ b/javascript/black-gray.css @@ -1,5 +1,5 @@ /* generic html tags */ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } :root, .light, .dark { --font: 'NotoSans'; --font-mono: 'ui-monospace', 'Consolas', monospace; diff --git a/javascript/black-orange.css b/javascript/black-orange.css index 816394d0d..54b98b1df 100644 --- a/javascript/black-orange.css +++ b/javascript/black-orange.css @@ -1,5 +1,5 @@ /* generic html tags */ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } :root, .light, .dark { --font: 'NotoSans'; --font-mono: 'ui-monospace', 'Consolas', monospace; diff --git a/javascript/black-teal-reimagined.css b/javascript/black-teal-reimagined.css index c221d7e4e..be6176ac4 100644 --- a/javascript/black-teal-reimagined.css +++ b/javascript/black-teal-reimagined.css @@ -4,7 +4,7 @@ font-display: swap; font-style: normal; font-weight: 100; - src: local('NotoSans'), url('notosans-nerdfont-regular.ttf'); + src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf'); } html { diff --git a/javascript/black-teal.css b/javascript/black-teal.css index 2ebf32e96..0a6db4fa2 100644 --- a/javascript/black-teal.css +++ b/javascript/black-teal.css @@ -1,5 +1,5 @@ /* generic html tags */ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } :root, .light, .dark { --font: 'NotoSans'; --font-mono: 'ui-monospace', 'Consolas', monospace; diff --git a/javascript/light-teal.css b/javascript/light-teal.css index 174622e52..7dc9e5950 100644 --- a/javascript/light-teal.css +++ b/javascript/light-teal.css @@ -1,5 +1,5 @@ /* generic html tags */ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } :root, .light, .dark { --font: 'NotoSans'; --font-mono: 'ui-monospace', 'Consolas', monospace; diff --git a/javascript/sdnext.css b/javascript/sdnext.css index 139b7d1bb..1db6288a5 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -1,4 +1,4 @@ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } :root { --left-column: 530px; --color-trace: #666666; diff --git a/javascript/simple-dark.css b/javascript/simple-dark.css index c1f99628d..e2d88f098 100644 --- a/javascript/simple-dark.css +++ b/javascript/simple-dark.css @@ -1,5 +1,5 @@ /* generic html tags */ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } :root, .light, .dark { --font: 'NotoSans'; --font-mono: 'ui-monospace', 'Consolas', monospace; diff --git a/javascript/simple-light.css b/javascript/simple-light.css index 769d12b2b..9fc84c4a6 100644 --- a/javascript/simple-light.css +++ b/javascript/simple-light.css @@ -1,5 +1,5 @@ /* generic html tags */ -@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSans'), url('notosans-nerdfont-regular.ttf') } +@font-face { font-family: 'NotoSans'; font-display: swap; font-style: normal; font-weight: 100; src: local('NotoSansNerd'), url('notosans-nerdfont-regular.ttf') } :root, .light, .dark { --font: 'NotoSans'; --font-mono: 'ui-monospace', 'Consolas', monospace; diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 4d988e73b..33319e174 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -135,7 +135,7 @@ def full_vae_decode(latents, model): decoded = model.vae.decode(latents, return_dict=False)[0] except Exception as e: shared.log.error(f'VAE decode: {e}') - if 'out of memory' not in str(e): + if 'out of memory' not in str(e) and 'no data' not in str(e): errors.display(e, 'VAE decode') decoded = [] diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py index ec41a130f..b0cd4d9b8 100644 --- a/modules/sd_checkpoint.py +++ b/modules/sd_checkpoint.py @@ -309,7 +309,8 @@ def read_metadata_from_safetensors(filename): metadata_len = int.from_bytes(metadata_len, "little") json_start = file.read(2) if metadata_len <= 2 or json_start not in (b'{"', b"{'"): - shared.log.error(f'Model metadata invalid: file="{filename}"') + shared.log.error(f'Model metadata invalid: file="{filename}" len={metadata_len} start={json_start}') + return res json_data = json_start + file.read(metadata_len-2) json_obj = json.loads(json_data) for k, v in json_obj.get("__metadata__", {}).items(): @@ -332,6 +333,8 @@ def read_metadata_from_safetensors(filename): res[k] = v except Exception as e: shared.log.error(f'Model metadata: file="{filename}" {e}') + from modules import errors + errors.display(e, 'Model metadata') sd_metadata[filename] = res global sd_metadata_pending # pylint: disable=global-statement sd_metadata_pending += 1 diff --git a/modules/shared.py b/modules/shared.py index 016ed98df..8c625bc24 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -518,9 +518,9 @@ options_templates.update(options_section(('text_encoder', "Text Encoder"), { options_templates.update(options_section(('cuda', "Compute Settings"), { "math_sep": OptionInfo("

Execution Precision

", "", gr.HTML), - "precision": OptionInfo("Autocast", "Precision type", gr.Radio, {"choices": ["Autocast", "Full"]}), + "precision": OptionInfo("Autocast", "Precision type", gr.Radio, {"choices": ["Autocast", "Full"], "visible": not native}), "cuda_dtype": OptionInfo("Auto", "Device precision type", gr.Radio, {"choices": ["Auto", "FP32", "FP16", "BF16"]}), - "no_half": OptionInfo(False if not cmd_opts.use_openvino else True, "Full precision (--no-half)", None, None, None), + "no_half": OptionInfo(False if not cmd_opts.use_openvino else True, "Force full precision (--no-half)", None, None, None), "upcast_sampling": OptionInfo(False if sys.platform != "darwin" else True, "Upcast sampling", gr.Checkbox, {"visible": not native}), "upcast_attn": OptionInfo(False, "Upcast attention layer", gr.Checkbox, {"visible": not native}), "cuda_cast_unet": OptionInfo(False, "Fixed UNet precision", gr.Checkbox, {"visible": not native}), From fc619ea99df19afa8b718ecdf22f2619da892015 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 18 Jan 2025 21:03:01 +0300 Subject: [PATCH 03/66] [Intel] Switch to PyTorch test branch instead of nightlies on Windows --- installer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/installer.py b/installer.py index 1e62ddd2a..40579066b 100644 --- a/installer.py +++ b/installer.py @@ -673,7 +673,7 @@ def install_ipex(torch_command): # torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/test/xpu') # test wheels are stable previews, significantly slower than IPEX # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.1') else: - torch_command = os.environ.get('TORCH_COMMAND', '--pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/xpu') # torchvision doesn't exist on test/stable branch for windows + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+xpu torchvision==0.21.0+xpu --index-url https://download.pytorch.org/whl/test/xpu') install(os.environ.get('OPENVINO_PACKAGE', 'openvino==2024.5.0'), 'openvino', ignore=True) install('nncf==2.7.0', ignore=True, no_deps=True) # requires older pandas install(os.environ.get('ONNXRUNTIME_PACKAGE', 'onnxruntime-openvino'), 'onnxruntime-openvino', ignore=True) From bd9c0d0fbcff801ac0c0dbb7988d135a3ebdd475 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 18 Jan 2025 21:11:10 +0300 Subject: [PATCH 04/66] IPEX update supported python versions --- installer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/installer.py b/installer.py index 40579066b..88df3521a 100644 --- a/installer.py +++ b/installer.py @@ -659,7 +659,7 @@ def install_rocm_zluda(): def install_ipex(torch_command): t_start = time.time() - check_python(supported_minors=[10,11], reason='IPEX backend requires Python 3.10 or 3.11') + check_python(supported_minors=[9, 10, 11, 12], reason='IPEX backend requires Python 3.9, 3.10 or 3.11') args.use_ipex = True # pylint: disable=attribute-defined-outside-init log.info('IPEX: Intel OneAPI toolkit detected') if os.environ.get("NEOReadDebugKeys", None) is None: From 7006e4705b887ea346dbe886232d9bc3643a2b8a Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 18 Jan 2025 21:56:14 +0300 Subject: [PATCH 05/66] IPEX enable Dynamic Atten for Battlemage --- installer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/installer.py b/installer.py index 88df3521a..88b3b30eb 100644 --- a/installer.py +++ b/installer.py @@ -668,13 +668,16 @@ def install_ipex(torch_command): os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100') if os.environ.get("PYTORCH_ENABLE_XPU_FALLBACK", None) is None: os.environ.setdefault('PYTORCH_ENABLE_XPU_FALLBACK', '1') + if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is None: + # Battlemage doesn't support Flash Atten or Memory Atten yet so it goes OOM without this + os.environ.setdefault('IPEX_FORCE_ATTENTION_SLICE', '1') if "linux" in sys.platform: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cxx11.abi torchvision==0.20.1+cxx11.abi intel-extension-for-pytorch==2.5.10+xpu oneccl_bind_pt==2.5.0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/cn/') # torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/test/xpu') # test wheels are stable previews, significantly slower than IPEX # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.1') else: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+xpu torchvision==0.21.0+xpu --index-url https://download.pytorch.org/whl/test/xpu') - install(os.environ.get('OPENVINO_PACKAGE', 'openvino==2024.5.0'), 'openvino', ignore=True) + install(os.environ.get('OPENVINO_PACKAGE', 'openvino==2024.6.0'), 'openvino', ignore=True) install('nncf==2.7.0', ignore=True, no_deps=True) # requires older pandas install(os.environ.get('ONNXRUNTIME_PACKAGE', 'onnxruntime-openvino'), 'onnxruntime-openvino', ignore=True) ts('ipex', t_start) From facfe2be6b7ec1335dcfabbced2059a28181a5ff Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 18 Jan 2025 22:10:34 +0300 Subject: [PATCH 06/66] IPEX make IPEX_FORCE_ATTENTION_SLICE=-1 disable Dynamic Atten --- modules/intel/ipex/hijacks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py index b1c9a1182..5440b7b68 100644 --- a/modules/intel/ipex/hijacks.py +++ b/modules/intel/ipex/hijacks.py @@ -75,7 +75,7 @@ def as_tensor(data, dtype=None, device=None): return original_as_tensor(data, dtype=dtype, device=device) -if device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is None: +if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '-1' or (device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0'): original_torch_bmm = torch.bmm original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention else: From ae9b40c68884dd99041cd049c11469f5636b3f55 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 18 Jan 2025 11:49:43 -0500 Subject: [PATCH 07/66] refactor to unify latent, resize and model based upscalers Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 + modules/hypernetworks/hypernetwork.py | 14 ++-- modules/images_resize.py | 65 ++++++++++-------- modules/modelloader.py | 14 ++-- modules/postprocess/sdupscaler_model.py | 7 +- modules/processing_diffusers.py | 3 +- modules/processing_helpers.py | 19 +++--- modules/processing_original.py | 30 ++++---- modules/processing_vae.py | 17 +++++ modules/sd_hijack_optimizations.py | 12 ++-- modules/shared.py | 22 +----- modules/ui_sections.py | 13 ++-- modules/upscaler.py | 76 ++++++--------------- modules/upscaler_simple.py | 91 +++++++++++++++++++++++++ scripts/xyz_grid_classes.py | 4 +- webui.py | 1 + wiki | 2 +- 17 files changed, 223 insertions(+), 169 deletions(-) create mode 100644 modules/upscaler_simple.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 06854d872..281fc7a01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Update for 2025-01-17 +- **Upscale** + - code refactor to unify latent, resize and model based upscalers - **Gallery**: - add http fallback for slow/unreliable links - **Fixes**: diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py index e15cf7724..fef95ce5d 100644 --- a/modules/hypernetworks/hypernetwork.py +++ b/modules/hypernetworks/hypernetwork.py @@ -9,6 +9,8 @@ from einops import rearrange, repeat from modules import devices, shared, hashes, errors, files_cache +loaded_hypernetworks = [] + class HypernetworkModule(torch.nn.Module): activation_dict = { "linear": torch.nn.Identity, @@ -280,10 +282,10 @@ def load_hypernetwork(name): def load_hypernetworks(names, multipliers=None): already_loaded = {} - for hypernetwork in shared.loaded_hypernetworks: - if hypernetwork.name in names: - already_loaded[hypernetwork.name] = hypernetwork - shared.loaded_hypernetworks.clear() + for hn in loaded_hypernetworks: + if hn.name in names: + already_loaded[hn.name] = hn + loaded_hypernetworks.clear() for i, name in enumerate(names): hypernetwork = already_loaded.get(name, None) if hypernetwork is None: @@ -291,7 +293,7 @@ def load_hypernetworks(names, multipliers=None): if hypernetwork is None: continue hypernetwork.set_multiplier(multipliers[i] if multipliers else 1.0) - shared.loaded_hypernetworks.append(hypernetwork) + loaded_hypernetworks.append(hypernetwork) def find_closest_hypernetwork_name(search: str): @@ -330,7 +332,7 @@ def attention_CrossAttention_forward(self, x, context=None, mask=None): h = self.heads q = self.to_q(x) context = default(context, x) - context_k, context_v = apply_hypernetworks(shared.loaded_hypernetworks, context, self) + context_k, context_v = apply_hypernetworks(loaded_hypernetworks, context, self) k = self.to_k(context_k) v = self.to_v(context_v) q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q, k, v)) diff --git a/modules/images_resize.py b/modules/images_resize.py index 183e1d7f1..a549b5bf9 100644 --- a/modules/images_resize.py +++ b/modules/images_resize.py @@ -1,45 +1,50 @@ +from typing import Union import sys import time import numpy as np +import torch from PIL import Image -from modules import shared +from modules import shared, upscaler -def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, upscaler_name: str=None, output_type: str='image', context: str=None): +def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width: int, height: int, upscaler_name: str=None, output_type: str='image', context: str=None): upscaler_name = upscaler_name or shared.opts.upscaler_for_img2img - def latent(im, w, h, upscaler): - from modules.processing_vae import vae_encode, vae_decode - import torch - latents = vae_encode(im, shared.sd_model, full_quality=False) # TODO resize image: enable full VAE mode for resize-latent - latents = torch.nn.functional.interpolate(latents, size=(int(h // 8), int(w // 8)), mode=upscaler["mode"], antialias=upscaler["antialias"]) - im = vae_decode(latents, shared.sd_model, output_type='pil', full_quality=False)[0] - return im + def latent(im, scale: float, selected_upscaler: upscaler.UpscalerData): + if isinstance(im, torch.Tensor): + im = selected_upscaler.scaler.upscale(im, scale, selected_upscaler.name) + return im + else: + from modules.processing_vae import vae_encode, vae_decode + latents = vae_encode(im, shared.sd_model, full_quality=False) # TODO resize image: enable full VAE mode for resize-latent + latents = selected_upscaler.scaler.upscale(latents, scale, selected_upscaler.name) + im = vae_decode(latents, shared.sd_model, output_type='pil', full_quality=False)[0] + return im - def resize(im, w, h): - w = int(w) - h = int(h) - if upscaler_name is None or upscaler_name == "None" or im.mode == 'L': + def resize(im: Union[Image.Image, torch.Tensor], w, h): + w, h = int(w), int(h) + if upscaler_name is None or upscaler_name == "None" or (hasattr(im, 'mode') and im.mode == 'L'): return im.resize((w, h), resample=Image.Resampling.LANCZOS) # force for mask - scale = max(w / im.width, h / im.height) + if isinstance(im, torch.Tensor): + scale = max(w // 8 / im.shape[-1] , h // 8 / im.shape[-2]) + else: + scale = max(w / im.width, h / im.height) if scale > 1.0: upscalers = [x for x in shared.sd_upscalers if x.name.lower().replace('-', ' ') == upscaler_name.lower().replace('-', ' ')] if len(upscalers) > 0: - upscaler = upscalers[0] - im = upscaler.scaler.upscale(im, scale, upscaler.data_path) - else: - upscaler = shared.latent_upscale_modes.get(upscaler_name, None) - if upscaler is not None: - im = latent(im, w, h, upscaler) + selected_upscaler: upscaler.UpscalerData = upscalers[0] + if selected_upscaler.name.lower().startswith('latent'): + im = latent(im, scale, selected_upscaler) else: - upscaler = shared.sd_upscalers[0] - shared.log.warning(f"Resize upscaler: invalid={upscaler_name} fallback={upscaler.name}") - shared.log.debug(f"Resize upscaler: available={[u.name for u in shared.sd_upscalers]}") - if im.width != w or im.height != h: # probably downsample after upscaler created larger image + im = selected_upscaler.scaler.upscale(im, scale, selected_upscaler.name) + else: + shared.log.warning(f"Resize upscaler: invalid={upscaler_name} fallback={selected_upscaler.name}") + shared.log.debug(f"Resize upscaler: available={[u.name for u in shared.sd_upscalers]}") + if isinstance(im, Image.Image) and (im.width != w or im.height != h): # probably downsample after upscaler created larger image im = im.resize((w, h), resample=Image.Resampling.LANCZOS) return im - def crop(im): + def crop(im: Image.Image): ratio = width / height src_ratio = im.width / im.height src_w = width if ratio > src_ratio else im.width * height // im.height @@ -49,7 +54,7 @@ def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, ups res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) return res - def fill(im, color=None): + def fill(im: Image.Image, color=None): color = color or shared.opts.image_background """ ratio = round(width / height, 1) @@ -77,7 +82,8 @@ def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, ups res.paste(im, box=((width - im.width)//2, (height - im.height)//2)) return res - def context_aware(im, width, height, context): + def context_aware(im: Image.Image, width, height, context): + width, height = int(width), int(height) import seam_carving # https://github.com/li-plus/seam-carving if 'forward' in context.lower(): energy_mode = "forward" @@ -110,7 +116,10 @@ def resize_image(resize_mode: int, im: Image.Image, width: int, height: int, ups t0 = time.time() if resize_mode is None: resize_mode = 0 - if resize_mode == 0 or (im.width == width and im.height == height) or (width == 0 and height == 0): # none + if isinstance(im, torch.Tensor): # latent resize only supports fixed mode + res = resize(im, width, height) + return res + elif (resize_mode == 0) or (im.width == width and im.height == height) or (width == 0 and height == 0): # none res = im.copy() elif resize_mode == 1: # fixed res = resize(im, width, height) diff --git a/modules/modelloader.py b/modules/modelloader.py index b022b4fc6..0c5fa78d8 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -11,7 +11,7 @@ from PIL import Image import rich.progress as p import huggingface_hub as hf from modules import shared, errors, files_cache -from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone +from modules.upscaler import Upscaler from modules.paths import script_path, models_path @@ -575,7 +575,7 @@ def load_upscalers(): importlib.import_module(full_model) except Exception as e: shared.log.error(f'Error loading upscaler: {model_name} {e}') - datas = [] + upscalers = [] commandline_options = vars(shared.cmd_opts) # some of upscaler classes will not go away after reloading their modules, and we'll end up with two copies of those classes. The newest copy will always be the last in the list, so we go from end to beginning and ignore duplicates used_classes = {} @@ -583,7 +583,7 @@ def load_upscalers(): classname = str(cls) if classname not in used_classes: used_classes[classname] = cls - names = [] + upscaler_types = [] for cls in reversed(used_classes.values()): name = cls.__name__ cmd_name = f"{name.lower().replace('upscaler', '')}_models_path" @@ -591,9 +591,9 @@ def load_upscalers(): scaler = cls(commandline_model_path) scaler.user_path = commandline_model_path scaler.model_download_path = commandline_model_path or scaler.model_path - datas += scaler.scalers - names.append(name[8:]) - shared.sd_upscalers = sorted(datas, key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else "") # Special case for UpscalerNone keeps it at the beginning of the list. + upscalers += scaler.scalers + upscaler_types.append(name[8:]) + shared.sd_upscalers = upscalers t1 = time.time() - shared.log.info(f"Available Upscalers: items={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])} time={t1-t0:.2f} types={names}") + shared.log.info(f"Available Upscalers: items={len(shared.sd_upscalers)} downloaded={len([x for x in shared.sd_upscalers if x.data_path is not None and os.path.isfile(x.data_path)])} user={len([x for x in shared.sd_upscalers if x.custom])} time={t1-t0:.2f} types={upscaler_types}") return [x.name for x in shared.sd_upscalers] diff --git a/modules/postprocess/sdupscaler_model.py b/modules/postprocess/sdupscaler_model.py index e5b1c8b45..5ec7168d3 100644 --- a/modules/postprocess/sdupscaler_model.py +++ b/modules/postprocess/sdupscaler_model.py @@ -4,7 +4,8 @@ from PIL import Image from modules import shared, devices from modules.upscaler import Upscaler, UpscalerData -class UpscalerSD(Upscaler): + +class UpscalerDiffusion(Upscaler): def __init__(self, dirname): # pylint: disable=super-init-not-called self.name = "SDUpscale" self.user_path = dirname @@ -12,8 +13,8 @@ class UpscalerSD(Upscaler): super().__init__() return self.scalers = [ - UpscalerData(name="SD Latent 2x", path="stabilityai/sd-x2-latent-upscaler", upscaler=self, model=None, scale=4), - UpscalerData(name="SD Latent 4x", path="stabilityai/stable-diffusion-x4-upscaler", upscaler=self, model=None, scale=4), + UpscalerData(name="Diffusion Latent Upscaler 2x", path="stabilityai/sd-x2-latent-upscaler", upscaler=self, model=None, scale=4), + UpscalerData(name="Diffusion Latent Upscaler 4x", path="stabilityai/stable-diffusion-x4-upscaler", upscaler=self, model=None, scale=4), ] self.pipelines = [ None, diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index c20ba85a8..240ac8670 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -178,9 +178,8 @@ def process_hires(p: processing.StableDiffusionProcessing, output): output.images = resize_hires(p, latents=output.images) if output is not None else [] sd_hijack_hypertile.hypertile_set(p, hr=True) - latent_upscale = shared.latent_upscale_modes.get(p.hr_upscaler, None) strength = p.hr_denoising_strength if p.hr_denoising_strength > 0 else p.denoising_strength - if (latent_upscale is not None or p.hr_force) and strength > 0: + if (p.hr_upscaler.lower().startswith('latent') or p.hr_force) and strength > 0: p.ops.append('hires') sd_models_compile.openvino_recompile_model(p, hires=True, refiner=False) if shared.sd_model.__class__.__name__ == "OnnxRawPipeline": diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 93896a02f..083d3489a 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -409,20 +409,19 @@ def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler shared.log.warning('Hires: input is not tensor') first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height) return first_pass_images - latent_upscaler = shared.latent_upscale_modes.get(p.hr_upscaler, None) - # shared.log.info(f'Hires: upscaler={p.hr_upscaler} width={p.hr_upscale_to_x} height={p.hr_upscale_to_y} images={latents.shape[0]}') - if latent_upscaler is not None: - return torch.nn.functional.interpolate(latents, size=(p.hr_upscale_to_y // 8, p.hr_upscale_to_x // 8), mode=latent_upscaler["mode"], antialias=latent_upscaler["antialias"]) - first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height) - if p.hr_upscale_to_x == 0 or (p.hr_upscale_to_y == 0 and hasattr(p, 'init_hr')): + + if (p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0) and hasattr(p, 'init_hr'): shared.log.error('Hires: missing upscaling dimensions') return first_pass_images + + if p.hr_upscaler.lower().startswith('latent'): + resized_image = images.resize_image(p.hr_resize_mode, latents, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) + return resized_image + + first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height) resized_images = [] for img in first_pass_images: - if latent_upscaler is None: - resized_image = images.resize_image(p.hr_resize_mode, img, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) - else: - resized_image = img + resized_image = images.resize_image(p.hr_resize_mode, img, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) resized_images.append(resized_image) devices.torch_gc() return resized_images diff --git a/modules/processing_original.py b/modules/processing_original.py index 7a0af1b04..261fc8c13 100644 --- a/modules/processing_original.py +++ b/modules/processing_original.py @@ -72,14 +72,6 @@ def process_original(p: processing.StableDiffusionProcessing): def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts): - latent_scale_mode = shared.latent_upscale_modes.get(p.hr_upscaler, None) if p.hr_upscaler is not None else shared.latent_upscale_modes.get(shared.latent_upscale_default_mode, "None") - if latent_scale_mode is not None: - p.hr_force = False # no need to force anything - if p.enable_hr and (latent_scale_mode is None or p.hr_force): - if len([x for x in shared.sd_upscalers if x.name == p.hr_upscaler]) == 0: - shared.log.warning(f"HiRes: upscaler={p.hr_upscaler} unknown") - p.enable_hr = False - p.ops.append('txt2img') hypertile_set(p) p.sampler = sd_samplers.create_sampler(p.sampler_name, p.sd_model) @@ -109,7 +101,16 @@ def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning, info = processing.create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, [], iteration=p.iteration, position_in_batch=i) p.extra_generation_params, p.detailer_enabled = orig_extra_generation_params, orig_detailer images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], shared.opts.samples_format, info=info, suffix="-before-hires") - if latent_scale_mode is None or p.hr_force: # non-latent upscaling + + if p.hr_upscaler.lower().startswith('latent'): # non-latent upscaling + p.hr_force = True + shared.state.job = 'Upscale' + samples = images.resize_image(1, samples, target_width, target_height, upscaler_name=p.hr_upscaler) + if getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0: + image_conditioning = img2img_image_conditioning(p, decode_first_stage(p.sd_model, samples.to(dtype=devices.dtype_vae), p.full_quality), samples) + else: + image_conditioning = txt2img_image_conditioning(p, samples.to(dtype=devices.dtype_vae)) + else: shared.state.job = 'Upscale' if decoded_samples is None: decoded_samples = decode_first_stage(p.sd_model, samples.to(dtype=devices.dtype_vae), p.full_quality) @@ -130,15 +131,8 @@ def sample_txt2img(p: processing.StableDiffusionProcessingTxt2Img, conditioning, else: samples = p.sd_model.get_first_stage_encoding(p.sd_model.encode_first_stage(resized_samples)) image_conditioning = img2img_image_conditioning(p, resized_samples, samples) - else: - samples = torch.nn.functional.interpolate(samples, size=(target_height // 8, target_width // 8), mode=latent_scale_mode["mode"], antialias=latent_scale_mode["antialias"]) - if getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) < 1.0: - image_conditioning = img2img_image_conditioning(p, decode_first_stage(p.sd_model, samples.to(dtype=devices.dtype_vae), p.full_quality), samples) - else: - image_conditioning = txt2img_image_conditioning(p, samples.to(dtype=devices.dtype_vae)) - if p.hr_sampler_name == "PLMS": - p.hr_sampler_name = 'UniPC' - if p.hr_force or latent_scale_mode is not None: + + if p.hr_force: shared.state.job = 'HiRes' if p.denoising_strength > 0: p.ops.append('hires') diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 33319e174..46bf0ad26 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -162,6 +162,7 @@ def full_vae_decode(latents, model): def full_vae_encode(image, model): + t0 = time.time() if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): log_debug('Moving to CPU: model=UNet') unet_device = model.unet.device @@ -170,9 +171,25 @@ def full_vae_encode(image, model): sd_models.move_model(model.vae, devices.device) vae_name = sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "default" log_debug(f'Encode vae="{vae_name}" dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}') + + upcast = (model.vae.dtype == torch.float16) and (getattr(model.vae.config, 'force_upcast', False) or shared.opts.no_half_vae) + if upcast: + if hasattr(model, 'upcast_vae'): # this is done by diffusers automatically if output_type != 'latent' + model.upcast_vae() + else: # manual upcast and we restore it later + model.vae.orig_dtype = model.vae.dtype + model.vae = model.vae.to(dtype=torch.float32) + encoded = model.vae.encode(image.to(model.vae.device, model.vae.dtype)).latent_dist.sample() + + if hasattr(model.vae, "orig_dtype"): + model.vae = model.vae.to(dtype=model.vae.orig_dtype) + del model.vae.orig_dtype + if shared.opts.diffusers_move_unet and not getattr(model, 'has_accelerate', False) and hasattr(model, 'unet'): sd_models.move_model(model.unet, unet_device) + t1 = time.time() + shared.log.debug(f'Encode: vae="{vae_name}" upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} latents={encoded.shape}:{encoded.device}:{encoded.dtype} time={t1-t0:.3f}') return encoded diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 31089551e..33ff274ce 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -52,7 +52,7 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None): # pylint q_in = self.to_q(x) context = default(context, x) # pylint: disable=possibly-used-before-assignment - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) del context, context_k, context_v, x @@ -90,7 +90,7 @@ def split_cross_attention_forward(self, x, context=None, mask=None): # pylint: d q_in = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) @@ -219,7 +219,7 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None): # q = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k = self.to_k(context_k) v = self.to_v(context_v) del context, context_k, context_v, x @@ -248,7 +248,7 @@ def sub_quad_attention_forward(self, x, context=None, mask=None): q = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k = self.to_k(context_k) v = self.to_v(context_v) del context, context_k, context_v, x @@ -329,7 +329,7 @@ def xformers_attention_forward(self, x, context=None, mask=None): # pylint: disa q_in = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) @@ -360,7 +360,7 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None): q_in = self.to_q(x) context = default(context, x) - context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + context_k, context_v = hypernetwork.apply_hypernetworks(hypernetwork.loaded_hypernetworks, context) k_in = self.to_k(context_k) v_in = self.to_v(context_v) diff --git a/modules/shared.py b/modules/shared.py index 8c625bc24..720ca7c55 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -15,7 +15,6 @@ import gradio as gr import fasteners import orjson import diffusers -from rich.console import Console from modules import errors, devices, shared_items, shared_state, cmd_args, theme, history, files_cache from modules.paths import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack @@ -26,22 +25,19 @@ import modules.interrogate import modules.memmon import modules.styles import modules.paths as paths -from installer import print_dict -from installer import log as central_logger # pylint: disable=E0611 +from installer import log, print_dict, console # pylint: disable=unused-import errors.install([gr]) demo: gr.Blocks = None api = None -log = central_logger progress_print_out = sys.stdout parser = cmd_args.parser url = 'https://github.com/vladmandic/automatic' cmd_opts, _ = parser.parse_known_args() hide_dirs = {"visible": not cmd_opts.hide_ui_dir_config} xformers_available = False -locking_available = True -clip_model = None +locking_available = True # used by file read/write locking interrogator = modules.interrogate.InterrogateModels(os.path.join("models", "interrogate")) sd_upscalers = [] detailers = [] @@ -51,20 +47,7 @@ tab_names = [] extra_networks = [] options_templates = {} hypernetworks = {} -loaded_hypernetworks = [] settings_components = None -latent_upscale_default_mode = "None" -latent_upscale_modes = { - "Latent Nearest": {"mode": "nearest", "antialias": False}, - "Latent Nearest-exact": {"mode": "nearest-exact", "antialias": False}, - "Latent Area": {"mode": "area", "antialias": False}, - "Latent Bilinear": {"mode": "bilinear", "antialias": False}, - "Latent Bicubic": {"mode": "bicubic", "antialias": False}, - "Latent Bilinear antialias": {"mode": "bilinear", "antialias": True}, - "Latent Bicubic antialias": {"mode": "bicubic", "antialias": True}, - # "Latent Linear": {"mode": "linear", "antialias": False}, # not supported for latents with channels=4 - # "Latent Trilinear": {"mode": "trilinear", "antialias": False}, # not supported for latents with channels=4 -} restricted_opts = { "samples_filename_pattern", "directories_filename_pattern", @@ -80,7 +63,6 @@ restricted_opts = { } resize_modes = ["None", "Fixed", "Crop", "Fill", "Outpaint", "Context aware"] compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order'] -console = Console(log_time=True, log_time_format='%H:%M:%S-%f') dir_timestamps = {} dir_cache = {} max_workers = 8 diff --git a/modules/ui_sections.py b/modules/ui_sections.py index def7e39b5..17436dd31 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -323,14 +323,6 @@ def create_hires_inputs(tab): with gr.Group(): with gr.Row(elem_id=f"{tab}_hires_row1"): enable_hr = gr.Checkbox(label='Enable refine pass', value=False, elem_id=f"{tab}_enable_hr") - """ - with gr.Row(elem_id=f"{tab}_hires_fix_row1", variant="compact"): - hr_upscaler = gr.Dropdown(label="Upscaler", elem_id=f"{tab}_hr_upscaler", choices=[*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]], value=shared.latent_upscale_default_mode) - hr_scale = gr.Slider(minimum=0.1, maximum=8.0, step=0.05, label="Rescale by", value=2.0, elem_id=f"{tab}_hr_scale") - with gr.Row(elem_id=f"{tab}_hires_fix_row3", variant="compact"): - hr_resize_x = gr.Slider(minimum=0, maximum=4096, step=8, label="Width resize", value=0, elem_id=f"{tab}_hr_resize_x") - hr_resize_y = gr.Slider(minimum=0, maximum=4096, step=8, label="Height resize", value=0, elem_id=f"{tab}_hr_resize_y") - """ hr_resize_mode, hr_upscaler, hr_resize_context, hr_resize_x, hr_resize_y, hr_scale, _selected_scale_tab = create_resize_inputs(tab, None, accordion=False, latent=True, non_zero=False) with gr.Row(elem_id=f"{tab}_hires_fix_row2", variant="compact"): hr_force = gr.Checkbox(label='Force HiRes', value=False, elem_id=f"{tab}_hr_force") @@ -355,8 +347,11 @@ def create_resize_inputs(tab, images, accordion=True, latent=False, non_zero=Tru prefix = f' {prefix}' with gr.Accordion(open=False, label="Resize", elem_classes=["small-accordion"], elem_id=f"{tab}_resize_group") if accordion else gr.Group(): with gr.Row(): + available_upscalers = [x.name for x in shared.sd_upscalers] + if not latent: + available_upscalers = [x for x in available_upscalers if not x.lower().startswith('latent')] resize_mode = gr.Dropdown(label=f"Mode{prefix}" if non_zero else "Resize mode", elem_id=f"{tab}_resize_mode", choices=shared.resize_modes, type="index", value='Fixed') - resize_name = gr.Dropdown(label=f"Method{prefix}", elem_id=f"{tab}_resize_name", choices=([] if not latent else list(shared.latent_upscale_modes)) + [x.name for x in shared.sd_upscalers], value=shared.latent_upscale_default_mode, visible=True) + resize_name = gr.Dropdown(label=f"Method{prefix}", elem_id=f"{tab}_resize_name", choices=available_upscalers, value=available_upscalers[0], visible=True) resize_context_choices = ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"] resize_context = gr.Dropdown(label=f"Context{prefix}", elem_id=f"{tab}_resize_context", choices=resize_context_choices, value=resize_context_choices[0], visible=False) ui_common.create_refresh_button(resize_name, modelloader.load_upscalers, lambda: {"choices": modelloader.load_upscalers()}, 'refresh_upscalers') diff --git a/modules/upscaler.py b/modules/upscaler.py index bda39f858..80c0ddaf9 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -8,10 +8,9 @@ from modules import devices, modelloader, shared from installer import setup_logging -LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.Resampling.LANCZOS) -NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.Resampling.NEAREST) models = None + class Upscaler: name = None folder = None @@ -97,17 +96,24 @@ class Upscaler: orig_state = copy.deepcopy(shared.state) shared.state.begin('Upscale') self.scale = scale - dest_w = int(img.width * scale) - dest_h = int(img.height * scale) - for _ in range(3): - shape = (img.width, img.height) + if isinstance(img, Image.Image): + dest_w = int(img.width * scale) + dest_h = int(img.height * scale) + else: + dest_w = int(img.shape[-1] * scale) + dest_h = int(img.shape[-2] * scale) + if self.name.lower().startswith('latent'): img = self.do_upscale(img, selected_model) - if shape == (img.width, img.height): - break - if img.width >= dest_w and img.height >= dest_h: - break - if img.width != dest_w or img.height != dest_h: - img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS) + else: + for _ in range(3): + shape = (img.width, img.height) + img = self.do_upscale(img, selected_model) + if shape == (img.width, img.height): + break + if img.width >= dest_w and img.height >= dest_h: + break + if img.width != dest_w or img.height != dest_h: + img = img.resize((int(dest_w), int(dest_h)), resample=Image.Resampling.BICUBIC) shared.state.end() shared.state = orig_state return img @@ -125,7 +131,7 @@ class Upscaler: def find_model(self, path): info = None for scaler in self.scalers: - if scaler.data_path == path: + if (scaler.data_path == path) or (scaler.name == path): info = scaler break if info is None: @@ -157,50 +163,6 @@ class UpscalerData: self.model = model -class UpscalerNone(Upscaler): - name = "None" - scalers = [] - - def load_model(self, path): - pass - - def do_upscale(self, img, selected_model=None): - return img - - def __init__(self, dirname=None): # pylint: disable=unused-argument - super().__init__(False) - self.scalers = [UpscalerData("None", None, self)] - - -class UpscalerLanczos(Upscaler): - scalers = [] - - def do_upscale(self, img, selected_model=None): - return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=LANCZOS) - - def load_model(self, _): - pass - - def __init__(self, dirname=None): # pylint: disable=unused-argument - super().__init__(False) - self.name = "Lanczos" - self.scalers = [UpscalerData("Lanczos", None, self)] - - -class UpscalerNearest(Upscaler): - scalers = [] - - def do_upscale(self, img, selected_model=None): - return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=NEAREST) - - def load_model(self, _): - pass - - def __init__(self, dirname=None): # pylint: disable=unused-argument - super().__init__(False) - self.name = "Nearest" - self.scalers = [UpscalerData("Nearest", None, self)] - def compile_upscaler(model): try: if shared.opts.ipex_optimize and "Upscaler" in shared.opts.ipex_optimize: diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py new file mode 100644 index 000000000..fcf4f2410 --- /dev/null +++ b/modules/upscaler_simple.py @@ -0,0 +1,91 @@ +from PIL import Image +from modules.upscaler import Upscaler, UpscalerData + + +class UpscalerNone(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "None" + self.scalers = [UpscalerData("None", None, self)] + + def load_model(self, path): + pass + + def do_upscale(self, img, selected_model=None): + return img + + +class UpscalerResize(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "Resize" + self.scalers = [ + UpscalerData("Resize Nearest", None, self), + UpscalerData("Resize Lanczos", None, self), + UpscalerData("Resize Bicubic", None, self), + UpscalerData("Resize Bilinear", None, self), + UpscalerData("Resize Hamming", None, self), + UpscalerData("Resize Box", None, self), + ] + + def do_upscale(self, img: Image, selected_model=None): + if selected_model is None: + return img + if selected_model == "Resize Nearest": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.NEAREST) + if selected_model == "Resize Lanczos": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.LANCZOS) + if selected_model == "Resize Bicubic": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BICUBIC) + if selected_model == "Resize Bilinear": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BILINEAR) + if selected_model == "Resize Hamming": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.HAMMING) + if selected_model == "Resize Box": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BOX) + + + def load_model(self, _): + pass + + +class UpscalerLatent(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "Latent" + self.scalers = [ + UpscalerData("Latent Nearest", None, self), + UpscalerData("Latent Nearest exact", None, self), + UpscalerData("Latent Area", None, self), + UpscalerData("Latent Bilinear", None, self), + UpscalerData("Latent Bicubic", None, self), + UpscalerData("Latent Bilinear antialias", None, self), + UpscalerData("Latent Bicubic antialias", None, self), + ] + + def do_upscale(self, img: Image, selected_model=None): + import torch + import torch.nn.functional as F + if isinstance(img, torch.Tensor) and (len(img.shape) == 4): + _batch, _channel, h, w = img.shape + else: + raise ValueError(f"Latent upscale: image={img.shape if isinstance(img, torch.Tensor) else img} type={type(img)} if not supported") + h, w = int((8 * h * self.scale) // 8), int((8 * w * self.scale) // 8) + mode, antialias = '', '' + if selected_model == "Latent Nearest": + mode, antialias = 'nearest', False + elif selected_model == "Latent Nearest exact": + mode, antialias = 'nearest-exact', False + elif selected_model == "Latent Area": + mode, antialias = 'area', False + elif selected_model == "Latent Bilinear": + mode, antialias = 'bilinear', False + elif selected_model == "Latent Bicubic": + mode, antialias = 'bicubic', False + elif selected_model == "Latent Bilinear antialias": + mode, antialias = 'bilinear', True + elif selected_model == "Latent Bicubic antialias": + mode, antialias = 'bicubic', True + else: + raise ValueError(f"Latent upscale: model={selected_model} unknown") + return F.interpolate(img, size=(h, w), mode=mode, antialias=antialias) diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py index 2b1fa2d59..cd4df56e8 100644 --- a/scripts/xyz_grid_classes.py +++ b/scripts/xyz_grid_classes.py @@ -128,7 +128,7 @@ axis_options = [ AxisOption("[Sampler] Shift", float, apply_setting("schedulers_shift")), AxisOption("[Sampler] eta delta", float, apply_setting("eta_noise_seed_delta")), AxisOption("[Sampler] eta multiplier", float, apply_setting("scheduler_eta")), - AxisOption("[Refine] Upscaler", str, apply_field("hr_upscaler"), cost=0.3, choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]), + AxisOption("[Refine] Upscaler", str, apply_field("hr_upscaler"), cost=0.3, choices=lambda: [x.name for x in shared.sd_upscalers]), AxisOption("[Refine] Sampler", str, apply_hr_sampler_name, fmt=format_value_add_label, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), AxisOption("[Refine] Denoising strength", float, apply_field("denoising_strength")), AxisOption("[Refine] Hires steps", int, apply_field("hr_second_pass_steps")), @@ -136,7 +136,7 @@ axis_options = [ AxisOption("[Refine] Guidance rescale", float, apply_field("diffusers_guidance_rescale")), AxisOption("[Refine] Refiner start", float, apply_field("refiner_start")), AxisOption("[Refine] Refiner steps", float, apply_field("refiner_steps")), - AxisOption("[Postprocess] Upscaler", str, apply_upscaler, cost=0.4, choices=lambda: [x.name for x in shared.sd_upscalers][1:]), + AxisOption("[Postprocess] Upscaler", str, apply_upscaler, cost=0.4, choices=lambda: [x.name for x in shared.sd_upscalers]), AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]), AxisOption("[Postprocess] Detailer", str, apply_detailer, fmt=format_value_add_label), AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")), diff --git a/webui.py b/webui.py index f104bbee8..23a9eba95 100644 --- a/webui.py +++ b/webui.py @@ -29,6 +29,7 @@ import modules.ui import modules.txt2img import modules.img2img import modules.upscaler +import modules.upscaler_simple import modules.extra_networks import modules.ui_extra_networks import modules.textual_inversion.textual_inversion diff --git a/wiki b/wiki index 6caea0521..3dcc0808d 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 6caea0521a48bd5b99a18566f21797bff1f2c244 +Subproject commit 3dcc0808db1e9c351a845184a3bf0fe7ca783190 From bdf8ca7e278544163bf34e2170a022293b606b53 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 18 Jan 2025 11:54:03 -0500 Subject: [PATCH 08/66] add legacy mode indicator on splash screen Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 8 ++++---- modules/api/server.py | 2 ++ 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 281fc7a01..de7142e81 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,10 @@ ## Update for 2025-01-17 -- **Upscale** - - code refactor to unify latent, resize and model based upscalers -- **Gallery**: - - add http fallback for slow/unreliable links +- **Other**: + - **gallery**: add http fallback for slow/unreliable links + - **upscale**: code refactor to unify latent, resize and model based upscalers + - **splash**: add legacy mode indicator on splash screen - **Fixes**: - non-full vae decode - send-to image transfer diff --git a/modules/api/server.py b/modules/api/server.py index dabbe634c..87de2d897 100644 --- a/modules/api/server.py +++ b/modules/api/server.py @@ -16,6 +16,8 @@ def get_motd(): ver = shared.get_version() if ver.get('updated', None) is not None: motd = f"version {ver['hash']} {ver['updated']} {ver['url'].split('/')[-1]}
" + if not shared.native: + motd += "Legacy mode
" if shared.opts.motd: try: res = requests.get('https://vladmandic.github.io/automatic/motd', timeout=3) From e26de8cdba321c840dc8a56292f5a05f9f48b970 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 18 Jan 2025 14:29:38 -0500 Subject: [PATCH 09/66] detailer support for face restorer models Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 12 +++-- installer.py | 12 +++++ launch.py | 7 +-- modules/detailer.py | 4 +- modules/postprocess/aurasr_model.py | 2 +- modules/postprocess/codeformer_model.py | 2 +- modules/postprocess/restorer.py | 61 +++++++++++++++++++++++++ modules/postprocess/yolo.py | 38 +++++++++++---- modules/shared.py | 2 +- webui.sh | 7 ++- wiki | 2 +- 11 files changed, 126 insertions(+), 23 deletions(-) create mode 100644 modules/postprocess/restorer.py diff --git a/CHANGELOG.md b/CHANGELOG.md index de7142e81..defa805f1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,19 +1,23 @@ # Change Log for SD.Next -## Update for 2025-01-17 +## Update for 2025-01-18 -- **Other**: +- **Detailer**: + - in addition as standard behavior of detect & run-generate, it can now also run face-restore models + - included models are: *CodeFormer, RestoreFormer, GFPGan, GPEN-BFR* +- **Other**: + - **ipex**: update supported torch versions - **gallery**: add http fallback for slow/unreliable links - **upscale**: code refactor to unify latent, resize and model based upscalers - **splash**: add legacy mode indicator on splash screen -- **Fixes**: +- **Fixes**: - non-full vae decode - send-to image transfer - sana vae tiling - increase gallery timeouts - update ui element ids - modernui use local font - - unique font family registration + - unique font family registration ## Update for 2025-01-15 diff --git a/installer.py b/installer.py index 88b3b30eb..ea37f55f0 100644 --- a/installer.py +++ b/installer.py @@ -734,6 +734,17 @@ def install_torch_addons(): ts('addons', t_start) +# check cudnn +def check_cudnn(): + import site + site_packages = site.getsitepackages() + cuda_path = os.environ.get('CUDA_PATH', '') + for site_package in site_packages: + folder = os.path.join(site_package, 'nvidia', 'cudnn', 'lib') + if os.path.exists(folder) and folder not in cuda_path: + os.environ['CUDA_PATH'] = f"{cuda_path}:{folder}" + + # check torch version def check_torch(): t_start = time.time() @@ -845,6 +856,7 @@ def check_torch(): return if not args.skip_all: install_torch_addons() + check_cudnn() if args.profile: pr.disable() print_profile(pr, 'Torch') diff --git a/launch.py b/launch.py index 07fa93697..f92e9bfda 100755 --- a/launch.py +++ b/launch.py @@ -66,9 +66,10 @@ def get_custom_args(): installer.log.trace(f'Environment: {installer.print_dict(env)}') env = [f'{k}={v}' for k, v in os.environ.items() if k.startswith('SD_')] installer.log.debug(f'Env flags: {env}') - ldd = os.environ.get('LD_PRELOAD', None) - if ldd is not None: - installer.log.debug(f'Linker flags: "{ldd}"') + ldpreload = os.environ.get('LD_PRELOAD', None) + ldpath = os.environ.get('LD_LIBRARY_PATH', None) + if ldpreload is not None or ldpath is not None: + installer.log.debug(f'Linker flags: preload="{ldpreload}" path="{ldpath}"') rec('args') diff --git a/modules/detailer.py b/modules/detailer.py index 31908c295..6c31371aa 100644 --- a/modules/detailer.py +++ b/modules/detailer.py @@ -1,3 +1,4 @@ +from abc import abstractmethod from modules import shared @@ -5,6 +6,7 @@ class Detailer: # abstract class used for postprocessing def name(self): return "None" + @abstractmethod def restore(self, np_image): return np_image @@ -13,5 +15,5 @@ def detail(np_image, p=None): # postprocesses the image detailers = [x for x in shared.detailers if x.name() == shared.opts.detailer_model or shared.opts.detailer_model is None] if len(detailers) == 0: return np_image - detailer = detailers[0] + detailer: Detailer = detailers[0] return detailer.restore(np_image, p) diff --git a/modules/postprocess/aurasr_model.py b/modules/postprocess/aurasr_model.py index ab5844d1a..546adf35b 100644 --- a/modules/postprocess/aurasr_model.py +++ b/modules/postprocess/aurasr_model.py @@ -3,7 +3,7 @@ import diffusers from PIL import Image from modules import shared, devices from modules.upscaler import Upscaler, UpscalerData -from installer import install + class UpscalerAuraSR(Upscaler): def __init__(self, dirname): # pylint: disable=super-init-not-called diff --git a/modules/postprocess/codeformer_model.py b/modules/postprocess/codeformer_model.py index c601f2b40..d0509a120 100644 --- a/modules/postprocess/codeformer_model.py +++ b/modules/postprocess/codeformer_model.py @@ -4,7 +4,7 @@ import torch import modules.detailer from modules import shared, devices, modelloader, errors from modules.paths import models_path -from installer import install + # codeformer people made a choice to include modified basicsr library to their project which makes # it utterly impossible to use it alongside with other libraries that also use basicsr, like GFPGAN. diff --git a/modules/postprocess/restorer.py b/modules/postprocess/restorer.py new file mode 100644 index 000000000..827a08336 --- /dev/null +++ b/modules/postprocess/restorer.py @@ -0,0 +1,61 @@ +import time +import cv2 +import numpy as np +from modules import shared, devices + + +face_helper = None + + +def restore(np_image, name, session, strength): # pylint: disable=unused-argument + t0 = time.time() + global face_helper # pylint: disable=global-statement + try: + from facelib.utils.face_restoration_helper import FaceRestoreHelper + from facelib.detection.retinaface import retinaface + except Exception as e: + shared.log.error(f"FaceRestorer error: {e}") + return np_image + if hasattr(retinaface, 'device'): + retinaface.device = devices.device + if face_helper is None: + face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device) + + np_image = np_image[:, :, ::-1] + original_resolution = np_image.shape[0:2] + resolution = session.get_inputs()[0].shape[-2:] + + if face_helper is None or session is None: + return np_image + face_helper.clean_all() + face_helper.read_image(np_image) + face_helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=5) + face_helper.align_warp_face() + + detected_faces = len(face_helper.cropped_faces) + for cropped_face in face_helper.cropped_faces: + cropped_face = cv2.resize(cropped_face, resolution, interpolation=cv2.INTER_LINEAR) + cropped_face = cropped_face.astype(np.float16)[:,:,::-1] / 255.0 + cropped_face = cropped_face.transpose((2, 0, 1)) + cropped_face = (cropped_face - 0.5) / 0.5 + cropped_face = np.expand_dims(cropped_face, axis=0).astype(np.float16) + w = np.array([strength], dtype=np.double) + if 'codeformer' in name: + restored_face = session.run(None, {'x':cropped_face, 'w':w})[0][0] + else: + restored_face = session.run(None, {'input':cropped_face})[0][0] + restored_face = (restored_face.transpose(1,2,0).clip(-1,1) + 1) * 0.5 + restored_face = (restored_face * 255)[:,:,::-1] + restored_face = restored_face.clip(0, 255).astype('uint8') + face_helper.add_restored_face(restored_face) + face_helper.get_inverse_affine(None) + restored_img = face_helper.paste_faces_to_input_image() + restored_img = restored_img[:, :, ::-1] + if original_resolution != restored_img.shape[0:2]: + restored_img = cv2.resize(restored_img, (0, 0), fx=original_resolution[1]/restored_img.shape[1], fy=original_resolution[0]/restored_img.shape[0], interpolation=cv2.INTER_LINEAR) + + face_helper.clean_all() + t1 = time.time() + shared.log.info(f'Detailer: model="{name}" faces={detected_faces} strength={strength} time={t1-t0:.3f}') + + return restored_img diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py index 85a50cc2a..5e6d5ed05 100644 --- a/modules/postprocess/yolo.py +++ b/modules/postprocess/yolo.py @@ -9,13 +9,17 @@ from modules import shared, processing, devices, processing_class, ui_common from modules.detailer import Detailer -PREDEFINED = [ # +predefined = [ # 'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11m.pt', 'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/face-yolo8n.pt', 'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/hand_yolov8n.pt', 'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/person_yolov8n-seg.pt', 'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/eyes-v1.pt', 'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/eyes-full-v1.pt', + 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/codeformer.fp16.onnx', + 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/restoreformer.fp16.onnx', + 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/GFPGANv1.4.fp16.onnx', + 'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/GPEN-BFR-512.fp16.onnx', ] load_lock = threading.Lock() @@ -50,7 +54,7 @@ class YoloRestorer(Detailer): self.list.clear() files = [] downloaded = 0 - for m in PREDEFINED: + for m in predefined: name = os.path.splitext(os.path.basename(m))[0] self.list[name] = m files.append(name) @@ -61,7 +65,7 @@ class YoloRestorer(Detailer): name = os.path.splitext(os.path.basename(f))[0] if name not in files: self.list[name] = os.path.join(shared.opts.yolo_dir, f) - shared.log.info(f'Available Yolo: path="{shared.opts.yolo_dir}" items={len(list(self.list))} downloaded={downloaded}') + shared.log.info(f'Available Detailer: path="{shared.opts.yolo_dir}" items={len(list(self.list))} downloaded={downloaded}') return self.list def dependencies(self): @@ -156,18 +160,30 @@ class YoloRestorer(Detailer): with load_lock: from modules import modelloader model = None - self.dependencies() if model_name is None: model_name = list(self.list)[0] if model_name in self.models: return model_name, self.models[model_name] else: - model_url = self.list.get(model_name) + model_url = self.list.get(model_name, None) + if model_url is None: + shared.log.error(f'Load: type=Detailer name="{model_name}" error="model not found"') + return None, None file_name = os.path.basename(model_url) model_file = None try: model_file = modelloader.load_file_from_url(url=model_url, model_dir=shared.opts.yolo_dir, file_name=file_name) - if model_file is not None: + if model_file is None: + shared.log.error(f'Load: type=Detailer name="{model_name}" url="{model_url}" error="failed to fetch model"') + elif model_file.endswith('.onnx'): + import onnxruntime as ort + options = ort.SessionOptions() + # options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + session = ort.InferenceSession(model_file, sess_options=options, providers=devices.onnx) + self.models[model_name] = session + return model_name, session + else: + self.dependencies() import ultralytics model = ultralytics.YOLO(model_file) classes = list(model.names.values()) @@ -200,6 +216,11 @@ class YoloRestorer(Detailer): shared.log.warning(f'Detailer: model="{name}" not loaded') continue + if name.endswith('.fp16'): + from modules.postprocess import restorer + np_image = restorer.restore(np_image, name, model, p.detailer_strength) + continue + image = Image.fromarray(np_image) items = self.predict(model, image) if len(items) == 0: @@ -262,8 +283,7 @@ class YoloRestorer(Detailer): p.steps = orig_p.get('steps', 0) report = [{'label': i.label, 'score': i.score, 'size': f'{i.width}x{i.height}' } for i in items] - shared.log.info(f'Detailer: model="{name}" items={report} args={items[0].args} denoise={p.denoising_strength} blur={p.mask_blur} width={p.width} height={p.height} padding={p.inpaint_full_res_padding}') - # shared.log.debug(f'Detailer: prompt="{prompt}" negative="{negative}"') + shared.log.info(f'Detailer: model="{name}" items={report} args={items[0].args} strength={p.detailer_strength} blur={p.mask_blur} width={p.width} height={p.height} padding={p.inpaint_full_res_padding}') models_used.append(name) mask_all = [] @@ -304,8 +324,6 @@ class YoloRestorer(Detailer): p.image_mask = blend([np.array(m) for m in mask_all]) p.image_mask = Image.fromarray(p.image_mask) - # if len(models_used) > 0: - # shared.log.debug(f'Detailer processed: models={models_used}') return np_image def ui(self, tab: str): diff --git a/modules/shared.py b/modules/shared.py index 720ca7c55..54204c201 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -1172,7 +1172,7 @@ opts.data['uni_pc_lower_order_final'] = opts.schedulers_use_loworder # compatibi opts.data['uni_pc_order'] = max(2, opts.schedulers_solver_order) # compatibility log.info(f'Engine: backend={backend} compute={devices.backend} device={devices.get_optimal_device_name()} attention="{opts.cross_attention_optimization}" mode={devices.inference_context.__name__}') if not native: - log.warning('Backend=original is in maintainance-only mode') + log.warning('Backend=original: legacy mode / maintainance-only') opts.data['diffusers_offload_mode'] = 'none' prompt_styles = modules.styles.StyleDatabase(opts) diff --git a/webui.sh b/webui.sh index b6ae67847..dc814713b 100755 --- a/webui.sh +++ b/webui.sh @@ -84,7 +84,12 @@ fi # Add venv lib folder to PATH if [ -d "$(realpath "$venv_dir")/lib/" ] && [[ -z "${DISABLE_VENV_LIBS}" ]] then - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(realpath "$venv_dir")/lib/ + if [[ -v LD_LIBRARY_PATH ]] + then + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(realpath "$venv_dir")/lib/ + else + export LD_LIBRARY_PATH=$(realpath "$venv_dir")/lib/ + fi fi # Add ROCm to PATH if it's not already diff --git a/wiki b/wiki index 3dcc0808d..1e3ebab7d 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 3dcc0808db1e9c351a845184a3bf0fe7ca783190 +Subproject commit 1e3ebab7dcb772a6ad4e4c03be97973c72c36fa4 From de1063275c1403572c12441610b22ae205e64112 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 19 Jan 2025 10:04:48 -0500 Subject: [PATCH 10/66] extract thumbnail from model metadata Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 ++- modules/sd_checkpoint.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index defa805f1..56a2732b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-18 +## Update for 2025-01-19 - **Detailer**: - in addition as standard behavior of detect & run-generate, it can now also run face-restore models @@ -10,6 +10,7 @@ - **gallery**: add http fallback for slow/unreliable links - **upscale**: code refactor to unify latent, resize and model based upscalers - **splash**: add legacy mode indicator on splash screen + - **network**: extract thumbnail from model metadata if present - **Fixes**: - non-full vae decode - send-to image transfer diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py index b0cd4d9b8..e5ddd2e85 100644 --- a/modules/sd_checkpoint.py +++ b/modules/sd_checkpoint.py @@ -1,8 +1,11 @@ +import io +import base64 import os import re import time import json import collections +from PIL import Image from modules import shared, paths, modelloader, hashes, sd_hijack_accelerate @@ -289,6 +292,21 @@ def init_metadata(): sd_metadata = shared.readfile(sd_metadata_file, lock=True) if os.path.isfile(sd_metadata_file) else {} +def extract_thumbnail(filename, data): + try: + thumbnail = data.split(",")[1] + thumbnail = base64.b64decode(thumbnail) + thumbnail = io.BytesIO(thumbnail) + thumbnail = Image.open(thumbnail) + thumbnail = thumbnail.convert("RGB") + thumbnail = thumbnail.resize((512, 512), Image.Resampling.HAMMING) + fn = os.path.splitext(filename)[0] + print('HERE', thumbnail, fn) + thumbnail = thumbnail.save(f"{fn}.thumb.jpg", quality=50) + except Exception as e: + shared.log.error(f"Error extracting thumbnail: {filename} {e}") + + def read_metadata_from_safetensors(filename): global sd_metadata # pylint: disable=global-statement if sd_metadata is None: @@ -314,6 +332,8 @@ def read_metadata_from_safetensors(filename): json_data = json_start + file.read(metadata_len-2) json_obj = json.loads(json_data) for k, v in json_obj.get("__metadata__", {}).items(): + if k == 'modelspec.thumbnail' and v.startswith("data:"): + extract_thumbnail(filename, v) if v.startswith("data:"): v = 'data' if k == 'format' and v == 'pt': From 311d402b0ca8874d3e6645de85c8a30f313461f6 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 19 Jan 2025 14:36:09 -0500 Subject: [PATCH 11/66] model component merge Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 13 +- javascript/base.css | 2 +- javascript/sdnext.css | 7 +- modules/extras.py | 216 ++++++++------------- modules/merging/convert_sdxl.py | 297 +++++++++++++++++++++++++++++ modules/merging/modules_sdxl.py | 310 +++++++++++++++++++++++++++++++ modules/sd_checkpoint.py | 9 +- modules/sd_models.py | 32 ++-- modules/sd_samplers.py | 40 ++-- modules/sd_samplers_diffusers.py | 17 +- modules/ui_models.py | 135 ++++++++------ 11 files changed, 833 insertions(+), 245 deletions(-) create mode 100644 modules/merging/convert_sdxl.py create mode 100644 modules/merging/modules_sdxl.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 56a2732b5..eabd560f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,15 +2,26 @@ ## Update for 2025-01-19 +- **Model Merge** + - replace model components and merge LoRAs + in addition to existing model weights merge support + now also having ability to replace model components and merge LoRAs + you can also test merges in-memory without needing to save to disk at all + and you can also use it to convert diffusers to safetensors if you want + *example*: replace vae in your favorite model with a fixed one? replace text encoder? etc. + *note*: limited to sdxl for now, additional models can be added depending on popularity - **Detailer**: - in addition as standard behavior of detect & run-generate, it can now also run face-restore models - included models are: *CodeFormer, RestoreFormer, GFPGan, GPEN-BFR* - **Other**: - **ipex**: update supported torch versions - **gallery**: add http fallback for slow/unreliable links - - **upscale**: code refactor to unify latent, resize and model based upscalers - **splash**: add legacy mode indicator on splash screen - **network**: extract thumbnail from model metadata if present +- **Refactor**: + - **upscale**: code refactor to unify latent, resize and model based upscalers + - **loader**: ability to run in-memory models + - **schedulers**: ability to create model-less schedulers - **Fixes**: - non-full vae decode - send-to image transfer diff --git a/javascript/base.css b/javascript/base.css index 0a0b621c2..8f89685c2 100644 --- a/javascript/base.css +++ b/javascript/base.css @@ -77,7 +77,7 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt #extensions .info { margin: 0; } #extensions .date { opacity: 0.85; font-size: 90%; } -/* extra networks */ +/* networks */ .extra-networks > div { margin: 0; border-bottom: none !important; } .extra-networks .second-line { display: flex; width: -moz-available; width: -webkit-fill-available; gap: 0.3em; box-shadow: var(--input-shadow); margin-bottom: 2px; } .extra-networks .search { flex: 1; } diff --git a/javascript/sdnext.css b/javascript/sdnext.css index 1db6288a5..c3deebad7 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -207,7 +207,7 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt #extensions .info { margin: 0; } #extensions .date { opacity: 0.85; font-size: var(--text-sm); } -/* extra networks */ +/* networks */ .extra_networks_root { width: 0; position: absolute; height: auto; right: 0; top: 13em; z-index: 100; } /* default is sidebar view */ .extra-networks { background: var(--background-color); padding: var(--block-label-padding); } .extra-networks > div { margin: 0; border-bottom: none !important; gap: 0.3em 0; } @@ -269,11 +269,14 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt .ar-dropdown div { margin: 0; background: var(--background-color)} #txt2img_sampler_timesteps, #img2img_sampler_timesteps { max-width: calc(var(--left-column) - 50px); } -/* extras */ +/* models */ .extras { gap: 0.2em 1em !important } #extras_generate, #extras_interrupt, #extras_skip { display: block !important; position: relative; height: 36px; } #extras_upscale { margin-top: 10px } #pnginfo_html_info .gradio-html > div { margin: 0.5em; } +#models_image, #models_image > div { min-height: 0; } +#models_error { font-family: monospace; color: var(--body-text-color-subdued) } + /* log monitor */ .log-monitor { display: none; justify-content: unset !important; overflow: hidden; padding: 0; margin-top: auto; font-family: monospace; font-size: var(--text-xxs); } diff --git a/modules/extras.py b/modules/extras.py index 162491580..e4eb36639 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -4,14 +4,12 @@ import json import time import shutil +from PIL import Image import torch -import tqdm import gradio as gr import safetensors.torch -from modules.merging.merge import merge_models -from modules.merging.merge_utils import TRIPLE_METHODS - -from modules import shared, images, sd_models, sd_vae, sd_models_config, devices +from modules.merging import merge, merge_utils, modules_sdxl +from modules import shared, images, sd_models, sd_vae, sd_samplers, sd_models_config, devices def run_pnginfo(image): @@ -73,9 +71,9 @@ def run_modelmerger(id_task, **kwargs): # pylint: disable=unused-argument if kwargs.get("secondary_model_name", None) in [None, 'None']: return fail("Failed: Merging requires a secondary model.") secondary_model_info = sd_models.get_closet_checkpoint_match(kwargs.get("secondary_model_name", None)) - if kwargs.get("tertiary_model_name", None) in [None, 'None'] and kwargs.get("merge_mode", None) in TRIPLE_METHODS: + if kwargs.get("tertiary_model_name", None) in [None, 'None'] and kwargs.get("merge_mode", None) in merge_utils.TRIPLE_METHODS: return fail(f"Failed: Interpolation method ({kwargs.get('merge_mode', None)}) requires a tertiary model.") - tertiary_model_info = sd_models.get_closet_checkpoint_match(kwargs.get("tertiary_model_name", None)) if kwargs.get("merge_mode", None) in TRIPLE_METHODS else None + tertiary_model_info = sd_models.get_closet_checkpoint_match(kwargs.get("tertiary_model_name", None)) if kwargs.get("merge_mode", None) in merge_utils.TRIPLE_METHODS else None del kwargs["primary_model_name"] del kwargs["secondary_model_name"] @@ -128,7 +126,7 @@ def run_modelmerger(id_task, **kwargs): # pylint: disable=unused-argument sd_models.unload_model_weights() try: - theta_0 = merge_models(**kwargs) + theta_0 = merge.merge_models(**kwargs) except Exception as e: return fail(f"{e}") @@ -205,144 +203,80 @@ def run_modelmerger(id_task, **kwargs): # pylint: disable=unused-argument return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], f"Model saved to {output_modelname}"] -def run_modelconvert(model, checkpoint_formats, precision, conv_type, custom_name, unet_conv, text_encoder_conv, - vae_conv, others_conv, fix_clip): - # position_ids in clip is int64. model_ema.num_updates is int32 - dtypes_to_fp16 = {torch.float32, torch.float64, torch.bfloat16} - dtypes_to_bf16 = {torch.float32, torch.float64, torch.float16} +def run_model_modules(model_type:str, model_name:str, custom_name:str, + comp_unet:str, comp_vae:str, comp_te1:str, comp_te2:str, + precision:str, comp_scheduler:str, comp_prediction:str, + comp_lora:str, comp_fuse:float, + meta_author:str, meta_version:str, meta_license:str, meta_desc:str, meta_hint:str, meta_thumbnail:Image.Image, + create_diffusers:bool, create_safetensors:bool, debug:bool): - def conv_fp16(t: torch.Tensor): - return t.half() if t.dtype in dtypes_to_fp16 else t - - def conv_bf16(t: torch.Tensor): - return t.bfloat16() if t.dtype in dtypes_to_bf16 else t - - def conv_full(t): - return t - - _g_precision_func = { - "full": conv_full, - "fp32": conv_full, - "fp16": conv_fp16, - "bf16": conv_bf16, - } - - def check_weight_type(k: str) -> str: - if k.startswith("model.diffusion_model"): - return "unet" - elif k.startswith("first_stage_model"): - return "vae" - elif k.startswith("cond_stage_model"): - return "clip" - return "other" - - def load_model(path): - if path.endswith(".safetensors"): - m = safetensors.torch.load_file(path, device="cpu") + status = '' + def msg(text, err:bool=False): + nonlocal status + if err: + shared.log.error(f'Modules merge: {text}') else: - m = torch.load(path, map_location="cpu") - state_dict = m["state_dict"] if "state_dict" in m else m - return state_dict + shared.log.info(f'Modules merge: {text}') + status += text + '
' + return status - def fix_model(model, fix_clip=False): - # code from model-toolkit - nai_keys = { - 'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.', - 'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.', - 'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.' - } - for k in list(model.keys()): - for r in nai_keys: - if type(k) == str and k.startswith(r): - new_key = k.replace(r, nai_keys[r]) - model[new_key] = model[k] - del model[k] - shared.log.warning(f"Model convert: fixed NovelAI error key: {k}") - break - if fix_clip: - i = "cond_stage_model.transformer.text_model.embeddings.position_ids" - if i in model: - correct = torch.Tensor([list(range(77))]).to(torch.int64) - now = model[i].to(torch.int64) + if model_type != 'sdxl': + yield msg("only SDXL models are supported", err=True) + return + if len(custom_name) == 0: + yield msg("output name is required", err=True) + return + checkpoint_info = sd_models.get_closet_checkpoint_match(model_name) + if checkpoint_info is None: + yield msg("input model not found", err=True) + return + fn = checkpoint_info.filename + shared.state.begin('Merge') + yield msg("modules merge starting") + yield msg("unload current model") + sd_models.unload_model_weights(op='model') - broken = correct.ne(now) - broken = [i for i in range(77) if broken[0][i]] - model[i] = correct - if len(broken) != 0: - shared.log.warning(f"Model convert: fixed broken CLiP: {broken}") + modules_sdxl.recipe.name = custom_name + modules_sdxl.recipe.author = meta_author + modules_sdxl.recipe.version = meta_version + modules_sdxl.recipe.desc = meta_desc + modules_sdxl.recipe.hint = meta_hint + modules_sdxl.recipe.license = meta_license + modules_sdxl.recipe.thumbnail = meta_thumbnail + modules_sdxl.recipe.base = fn + modules_sdxl.recipe.unet = comp_unet + modules_sdxl.recipe.vae = comp_vae + modules_sdxl.recipe.te1 = comp_te1 + modules_sdxl.recipe.te2 = comp_te2 + modules_sdxl.recipe.prediction = comp_prediction + modules_sdxl.recipe.diffusers = create_diffusers + modules_sdxl.recipe.safetensors = create_safetensors + modules_sdxl.recipe.fuse = float(comp_fuse) + modules_sdxl.recipe.debug = debug - return model - - if model == "": - return "Error: you must choose a model" - if len(checkpoint_formats) == 0: - return "Error: at least choose one model save format" - - extra_opt = { - "unet": unet_conv, - "clip": text_encoder_conv, - "vae": vae_conv, - "other": others_conv - } - shared.state.begin('Convert') - model_info = sd_models.checkpoints_list[model] - shared.state.textinfo = f"Load {model_info.filename}..." - shared.log.info(f"Model convert loading: {model_info.filename}") - state_dict = load_model(model_info.filename) - - ok = {} # {"state_dict": {}} - - conv_func = _g_precision_func[precision] - - def _hf(wk: str, t: torch.Tensor): - if not isinstance(t, torch.Tensor): - return - w_t = check_weight_type(wk) - conv_t = extra_opt[w_t] - if conv_t == "convert": - ok[wk] = conv_func(t) - elif conv_t == "copy": - ok[wk] = t - elif conv_t == "delete": - return - - shared.log.info("Model convert: running") - if conv_type == "ema-only": - for k in tqdm.tqdm(state_dict): - ema_k = "___" - try: - ema_k = "model_ema." + k[6:].replace(".", "") - except Exception: - pass - if ema_k in state_dict: - _hf(k, state_dict[ema_k]) - elif not k.startswith("model_ema.") or k in ["model_ema.num_updates", "model_ema.decay"]: - _hf(k, state_dict[k]) - elif conv_type == "no-ema": - for k, v in tqdm.tqdm(state_dict.items()): - if "model_ema." not in k: - _hf(k, v) + loras = [l.strip() if ':' in l else f'{l.strip()}:1.0' for l in comp_lora.split(',') if len(l.strip()) > 0] + for lora, strength in [l.split(':') for l in loras]: + modules_sdxl.recipe.lora[lora] = float(strength) + scheduler = sd_samplers.create_sampler(comp_scheduler, None) + modules_sdxl.recipe.scheduler = scheduler.__class__.__name__ if scheduler is not None else None + if precision == 'fp32': + modules_sdxl.recipe.precision = torch.float32 + elif precision == 'bf16': + modules_sdxl.recipe.precision = torch.bfloat16 else: - for k, v in tqdm.tqdm(state_dict.items()): - _hf(k, v) + modules_sdxl.recipe.precision = torch.float16 - ok = fix_model(ok, fix_clip=fix_clip) - output = "" - ckpt_dir = shared.cmd_opts.ckpt_dir or sd_models.model_path - save_name = f"{model_info.model_name}-{precision}" - if conv_type != "disabled": - save_name += f"-{conv_type}" - if custom_name != "": - save_name = custom_name - for fmt in checkpoint_formats: - ext = ".safetensors" if fmt == "safetensors" else ".ckpt" - _save_name = save_name + ext - save_path = os.path.join(ckpt_dir, _save_name) - shared.log.info(f"Model convert saving: {save_path}") - if fmt == "safetensors": - safetensors.torch.save_file(ok, save_path) - else: - torch.save({"state_dict": ok}, save_path) - output += f"Checkpoint saved to {save_path}
" + modules_sdxl.status = status + yield from modules_sdxl.merge() + status = modules_sdxl.status + + devices.torch_gc(force=True) + yield msg("modules merge complete") + if modules_sdxl.pipeline is not None: + checkpoint_info = sd_models.CheckpointInfo(filename='None') + shared.sd_model = modules_sdxl.pipeline + sd_models.set_defaults(shared.sd_model, checkpoint_info) + sd_models.set_diffuser_options(shared.sd_model, offload=False) + sd_models.set_diffuser_offload(shared.sd_model) + yield msg("pipeline loaded") shared.state.end() - return output diff --git a/modules/merging/convert_sdxl.py b/modules/merging/convert_sdxl.py new file mode 100644 index 000000000..93fc71f5d --- /dev/null +++ b/modules/merging/convert_sdxl.py @@ -0,0 +1,297 @@ +import io +import os +import re +import hashlib +import torch +from safetensors.torch import load_file, save_file + + +unet_conversion_map = [ + # (stable-diffusion, HF Diffusers) + ("time_embed.0.weight", "time_embedding.linear_1.weight"), + ("time_embed.0.bias", "time_embedding.linear_1.bias"), + ("time_embed.2.weight", "time_embedding.linear_2.weight"), + ("time_embed.2.bias", "time_embedding.linear_2.bias"), + ("input_blocks.0.0.weight", "conv_in.weight"), + ("input_blocks.0.0.bias", "conv_in.bias"), + ("out.0.weight", "conv_norm_out.weight"), + ("out.0.bias", "conv_norm_out.bias"), + ("out.2.weight", "conv_out.weight"), + ("out.2.bias", "conv_out.bias"), + # the following are for sdxl + ("label_emb.0.0.weight", "add_embedding.linear_1.weight"), + ("label_emb.0.0.bias", "add_embedding.linear_1.bias"), + ("label_emb.0.2.weight", "add_embedding.linear_2.weight"), + ("label_emb.0.2.bias", "add_embedding.linear_2.bias"), +] + +unet_conversion_map_resnet = [ + # (stable-diffusion, HF Diffusers) + ("in_layers.0", "norm1"), + ("in_layers.2", "conv1"), + ("out_layers.0", "norm2"), + ("out_layers.3", "conv2"), + ("emb_layers.1", "time_emb_proj"), + ("skip_connection", "conv_shortcut"), +] + +unet_conversion_map_layer = [] +# hardcoded number of downblocks and resnets/attentions... +# would need smarter logic for other networks. +for i in range(3): + # loop over downblocks/upblocks + + for j in range(2): + # loop over resnets/attentions for downblocks + hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}." + sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0." + unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix)) + + if i > 0: + hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}." + sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1." + unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix)) + + for j in range(4): + # loop over resnets/attentions for upblocks + hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}." + sd_up_res_prefix = f"output_blocks.{3*i + j}.0." + unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix)) + + if i < 2: + # no attention layers in up_blocks.0 + hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}." + sd_up_atn_prefix = f"output_blocks.{3 * i + j}.1." + unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix)) + + if i < 3: + # no downsample in down_blocks.3 + hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv." + sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op." + unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix)) + + # no upsample in up_blocks.3 + hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0." + sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}." + unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix)) +unet_conversion_map_layer.append(("output_blocks.2.2.conv.", "output_blocks.2.1.conv.")) + +hf_mid_atn_prefix = "mid_block.attentions.0." +sd_mid_atn_prefix = "middle_block.1." +unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix)) +for j in range(2): + hf_mid_res_prefix = f"mid_block.resnets.{j}." + sd_mid_res_prefix = f"middle_block.{2*j}." + unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix)) + + +def convert_unet_state_dict(unet_state_dict): + # buyer beware: this is a *brittle* function, + # and correct output requires that all of these pieces interact in + # the exact order in which I have arranged them. + mapping = {k: k for k in unet_state_dict.keys()} + for sd_name, hf_name in unet_conversion_map: + mapping[hf_name] = sd_name + for k, v in mapping.items(): + if "resnets" in k: + for sd_part, hf_part in unet_conversion_map_resnet: + v = v.replace(hf_part, sd_part) + mapping[k] = v + for k, v in mapping.items(): + for sd_part, hf_part in unet_conversion_map_layer: + v = v.replace(hf_part, sd_part) + mapping[k] = v + new_state_dict = {sd_name: unet_state_dict[hf_name] for hf_name, sd_name in mapping.items()} + return new_state_dict + + +vae_conversion_map = [ + # (stable-diffusion, HF Diffusers) + ("nin_shortcut", "conv_shortcut"), + ("norm_out", "conv_norm_out"), + ("mid.attn_1.", "mid_block.attentions.0."), +] + +for i in range(4): + # down_blocks have two resnets + for j in range(2): + hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}." + sd_down_prefix = f"encoder.down.{i}.block.{j}." + vae_conversion_map.append((sd_down_prefix, hf_down_prefix)) + + if i < 3: + hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0." + sd_downsample_prefix = f"down.{i}.downsample." + vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix)) + + hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0." + sd_upsample_prefix = f"up.{3-i}.upsample." + vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix)) + + # up_blocks have three resnets + # also, up blocks in hf are numbered in reverse from sd + for j in range(3): + hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}." + sd_up_prefix = f"decoder.up.{3-i}.block.{j}." + vae_conversion_map.append((sd_up_prefix, hf_up_prefix)) + +# this part accounts for mid blocks in both the encoder and the decoder +for i in range(2): + hf_mid_res_prefix = f"mid_block.resnets.{i}." + sd_mid_res_prefix = f"mid.block_{i+1}." + vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix)) + + +vae_conversion_map_attn = [ + # (stable-diffusion, HF Diffusers) + ("norm.", "group_norm."), + # the following are for SDXL + ("q.", "to_q."), + ("k.", "to_k."), + ("v.", "to_v."), + ("proj_out.", "to_out.0."), +] + + +def reshape_weight_for_sd(w): + # convert HF linear weights to SD conv2d weights + if not w.ndim == 1: + return w.reshape(*w.shape, 1, 1) + else: + return w + + +def convert_vae_state_dict(vae_state_dict): + mapping = {k: k for k in vae_state_dict.keys()} + for k, v in mapping.items(): + for sd_part, hf_part in vae_conversion_map: + v = v.replace(hf_part, sd_part) + mapping[k] = v + for k, v in mapping.items(): + if "attentions" in k: + for sd_part, hf_part in vae_conversion_map_attn: + v = v.replace(hf_part, sd_part) + mapping[k] = v + new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()} + weights_to_convert = ["q", "k", "v", "proj_out"] + for k, v in new_state_dict.items(): + for weight_name in weights_to_convert: + if f"mid.attn_1.{weight_name}.weight" in k: + new_state_dict[k] = reshape_weight_for_sd(v) + return new_state_dict + + +textenc_conversion_lst = [ + # (stable-diffusion, HF Diffusers) + ("transformer.resblocks.", "text_model.encoder.layers."), + ("ln_1", "layer_norm1"), + ("ln_2", "layer_norm2"), + (".c_fc.", ".fc1."), + (".c_proj.", ".fc2."), + (".attn", ".self_attn"), + ("ln_final.", "text_model.final_layer_norm."), + ("token_embedding.weight", "text_model.embeddings.token_embedding.weight"), + ("positional_embedding", "text_model.embeddings.position_embedding.weight"), +] +protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst} +textenc_pattern = re.compile("|".join(protected.keys())) + +# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp +code2idx = {"q": 0, "k": 1, "v": 2} + + +def convert_openclip_text_enc_state_dict(text_enc_dict): + new_state_dict = {} + capture_qkv_weight = {} + capture_qkv_bias = {} + for k, v in text_enc_dict.items(): + if ( + k.endswith(".self_attn.q_proj.weight") + or k.endswith(".self_attn.k_proj.weight") + or k.endswith(".self_attn.v_proj.weight") + ): + k_pre = k[: -len(".q_proj.weight")] + k_code = k[-len("q_proj.weight")] + if k_pre not in capture_qkv_weight: + capture_qkv_weight[k_pre] = [None, None, None] + capture_qkv_weight[k_pre][code2idx[k_code]] = v + continue + + if ( + k.endswith(".self_attn.q_proj.bias") + or k.endswith(".self_attn.k_proj.bias") + or k.endswith(".self_attn.v_proj.bias") + ): + k_pre = k[: -len(".q_proj.bias")] + k_code = k[-len("q_proj.bias")] + if k_pre not in capture_qkv_bias: + capture_qkv_bias[k_pre] = [None, None, None] + capture_qkv_bias[k_pre][code2idx[k_code]] = v + continue + + relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k) + new_state_dict[relabelled_key] = v + + for k_pre, tensors in capture_qkv_weight.items(): + if None in tensors: + raise RuntimeError("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") + relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) + new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors) + + for k_pre, tensors in capture_qkv_bias.items(): + if None in tensors: + raise RuntimeError("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") + relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) + new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors) + + return new_state_dict + + +def convert_openai_text_enc_state_dict(text_enc_dict): + return text_enc_dict + + +def calculate_model_hash(state_dict): + func = hashlib.sha256() + for module in state_dict.values(): + buffer = io.BytesIO() + torch.save(module, buffer) + func.update(buffer.getvalue()) + return func.hexdigest() + + +def convert(model_path:str, checkpoint_path:str, metadata:dict={}): + unet_path = os.path.join(model_path, "unet", "diffusion_pytorch_model.safetensors") + vae_path = os.path.join(model_path, "vae", "diffusion_pytorch_model.safetensors") + text_enc_path = os.path.join(model_path, "text_encoder", "model.safetensors") + text_enc_2_path = os.path.join(model_path, "text_encoder_2", "model.safetensors") + + unet_state_dict = load_file(unet_path, device="cpu") + vae_state_dict = load_file(vae_path, device="cpu") + text_enc_dict = load_file(text_enc_path, device="cpu") + text_enc_2_dict = load_file(text_enc_2_path, device="cpu") + + unet_state_dict = convert_unet_state_dict(unet_state_dict) + unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()} + + vae_state_dict = convert_vae_state_dict(vae_state_dict) + vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()} + + text_enc_dict = convert_openai_text_enc_state_dict(text_enc_dict) + text_enc_dict = {"conditioner.embedders.0.transformer." + k: v for k, v in text_enc_dict.items()} + + text_enc_2_dict = convert_openclip_text_enc_state_dict(text_enc_2_dict) + text_enc_2_dict = {"conditioner.embedders.1.model." + k: v for k, v in text_enc_2_dict.items()} + text_enc_2_dict["conditioner.embedders.1.model.text_projection"] = text_enc_2_dict.pop("conditioner.embedders.1.model.text_projection.weight").T.contiguous() + + state_dict = { + **unet_state_dict, + **vae_state_dict, + **text_enc_dict, + **text_enc_2_dict + } + if metadata.get('modelspec.hash_sha256', None) is not None: + metadata['modelspec.hash_sha256'] = calculate_model_hash(state_dict) + + save_file(state_dict, checkpoint_path, metadata=metadata) + return metadata diff --git a/modules/merging/modules_sdxl.py b/modules/merging/modules_sdxl.py new file mode 100644 index 000000000..61f2cf948 --- /dev/null +++ b/modules/merging/modules_sdxl.py @@ -0,0 +1,310 @@ +import io +import os +import json +import base64 +from datetime import datetime +from PIL import Image +import torch +from safetensors.torch import load_file +import diffusers +import transformers +from modules import shared, devices + + +class Recipe: + author = '' + name = '' + version = '' + desc = '' + hint = '' + license = '' + prediction = '' + thumbnail = None + base = None + unet = None + vae = None + te1 = None + te2 = None + scheduler = 'UniPCMultistepScheduler' + dtype = torch.float16 + diffusers = True + safetensors = True + debug = False + lora = { + } + fuse = 1.0 +class Test: + generate = True + prompt = 'astronaut in a diner drinking coffee with burger and french fries on the table' + negative = 'ugly, blurry' + width = 1024 + height = 1024 + guidance = 4 + steps = 20 +recipe = Recipe() +test = Test() +pipeline: diffusers.StableDiffusionXLPipeline = None +status = '' + + +def msg(text, err:bool=False): + global status # pylint: disable=global-statement + if err: + shared.log.error(f'Modules merge: {text}') + else: + shared.log.info(f'Modules merge: {text}') + status += text + '
' + return status + + +def load_base(override:str=None): + global pipeline # pylint: disable=global-statement + fn = override or recipe.base + yield msg(f'base={fn}') + if os.path.isfile(fn): + pipeline = diffusers.StableDiffusionXLPipeline.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, torch_dtype=recipe.dtype, add_watermarker=False) + elif os.path.isdir(fn): + pipeline = diffusers.StableDiffusionXLPipeline.from_pretrained(fn, cache_dir=shared.opts.hfcache_dir, torch_dtype=recipe.dtype, add_watermarker=False) + else: + yield msg('base: not found') + return None + pipeline.vae.register_to_config(force_upcast = False) + + +def load_unet(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): + if (recipe.unet is None or len(recipe.unet) == 0) and override is None: + return None + fn = override or recipe.unet + if not os.path.isabs(fn): + fn = os.path.join(shared.opts.unet_dir, fn) + if not fn.endswith('.safetensors'): + fn += '.safetensors' + yield msg(f'unet={fn}') + if recipe.debug: + yield msg(f'config={pipe.unet.config}') + try: + unet = diffusers.UNet2DConditionModel.from_config(pipe.unet.config).to(recipe.dtype) + state_dict = load_file(fn) + unet.load_state_dict(state_dict) + pipe.unet = unet.to(device=devices.device, dtype=recipe.dtype) + except Exception as e: + yield msg(f'unet: {e}') + +def load_scheduler(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): + if recipe.scheduler is None and override is None: + return None + config = pipe.scheduler.config.__dict__ + scheduler = override or recipe.scheduler + yield msg(f'scheduler={scheduler}') + if recipe.debug: + yield msg(f'config={config}') + try: + pipe.scheduler = getattr(diffusers, scheduler).from_config(config) + except Exception as e: + yield msg(f'scheduler: {e}') + + + +def load_vae(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): + if (recipe.vae is None or len(recipe.vae) == 0)and override is None: + return None + fn = override or recipe.vae + if not os.path.isabs(fn): + fn = os.path.join(shared.opts.vae_dir, fn) + if not fn.endswith('.safetensors'): + fn += '.safetensors' + try: + vae = diffusers.AutoencoderKL.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, torch_dtype=recipe.dtype) + vae.config.force_upcast = False + vae.config.scaling_factor = 0.13025 + vae.config.sample_size = 1024 + yield msg(f'vae={fn}') + if recipe.debug: + yield msg(f'config={pipe.vae.config}') + pipe.vae = vae.to(device=devices.device, dtype=recipe.dtype) + except Exception as e: + yield msg(f'vae: {e}') + + +def load_te1(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): + if (recipe.te1 is None or len(recipe.te1) == 0) and override is None: + return None + config = pipe.text_encoder.config.__dict__ + pretrained_config = transformers.PretrainedConfig.from_dict(config) + fn = override or recipe.te1 + if not os.path.isabs(fn): + fn = os.path.join(shared.opts.te_dir, fn) + if not fn.endswith('.safetensors'): + fn += '.safetensors' + yield msg(f'te1={fn}') + if recipe.debug: + yield msg(f'config={config}') + try: + state_dict = load_file(fn) + te1 = transformers.CLIPTextModel.from_pretrained(pretrained_model_name_or_path=None, state_dict=state_dict, config=pretrained_config, cache_dir=shared.opts.hfcache_dir) + pipe.text_encoder = te1.to(device=devices.device, dtype=recipe.dtype) + except Exception as e: + yield msg(f'te1: {e}') + + +def load_te2(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): + if (recipe.te2 is None or len(recipe.te2) == 0) and override is None: + return None + config = pipe.text_encoder_2.config.__dict__ + pretrained_config = transformers.PretrainedConfig.from_dict(config) + fn = override or recipe.te2 + if not os.path.isabs(fn): + fn = os.path.join(shared.opts.te_dir, fn) + if not fn.endswith('.safetensors'): + fn += '.safetensors' + yield msg(f'te2={recipe.te2}') + if recipe.debug: + yield msg(f'config={config}') + try: + state_dict = load_file(fn) + te2 = transformers.CLIPTextModelWithProjection.from_pretrained(pretrained_model_name_or_path=None, state_dict=state_dict, config=pretrained_config, cache_dir=shared.opts.hfcache_dir) + pipe.text_encoder_2 = te2.to(device=devices.device, dtype=recipe.dtype) + except Exception as e: + yield msg(f'te2: {e}') + + +def load_lora(pipe: diffusers.StableDiffusionXLPipeline, override: dict=None, fuse: float=None): + if recipe.lora is None and override is None: + return None + names = [] + pipe.unfuse_lora() + pipe.unload_lora_weights() + loras = override or recipe.lora + for lora, weight in loras.items(): + try: + fn = lora + if not os.path.isabs(fn): + fn = os.path.join(shared.opts.lora_dir, fn) + if not fn.endswith('.safetensors'): + fn += '.safetensors' + yield msg(f'lora={fn} weight={weight} fuse={fuse or recipe.fuse}') + name = os.path.splitext(os.path.basename(lora))[0].replace('.', '').replace(' ', '').replace('-', '').replace('_', '') + names.append(name) + pipe.load_lora_weights(fn, name) + except Exception as e: + yield msg(f'lora: {e}') + if len(names) > 0: + pipe.set_adapters(adapter_names=names, adapter_weights=list(loras.values())) + pipe.fuse_lora(adapter_names=names, lora_scale=fuse or recipe.fuse, components=["unet", "text_encoder", "text_encoder_2"]) + pipe.unload_lora_weights() + + +def test_model(pipe: diffusers.StableDiffusionXLPipeline, fn: str, **kwargs): + if not test.generate: + return + try: + generator = torch.Generator(devices.device).manual_seed(int(4242)) + args = { + 'prompt': test.prompt, + 'negative_prompt': test.negative, + 'num_inference_steps': test.steps, + 'width': test.width, + 'height': test.height, + 'guidance_scale': test.guidance, + 'generator': generator, + } + args.update(kwargs) + yield msg(f'test={args}') + image = pipe(**args).images[0] + yield msg(f'image={fn} {image}') + image.save(fn) + except Exception as e: + yield msg(f'test: {e}') + + +def get_thumbnail(): + if recipe.thumbnail is None: + return '' + image = Image.open(recipe.thumbnail) + image = image.convert('RGB') + image.thumbnail((512, 512), resample=Image.Resampling.LANCZOS) + buffer = io.BytesIO() + image.save(buffer, format="JPEG", quality=50) + b64encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") + return f'data:image/jpeg;base64,{b64encoded}' + + +def get_metadata(): + return { + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "stable-diffusion-xl-v1-base", + "modelspec.implementation": "diffusers", + "modelspec.title": recipe.name, + "modelspec.version": recipe.version, + "modelspec.description": recipe.desc, + "modelspec.author": recipe.author, + "modelspec.date": datetime.now().isoformat(timespec='minutes'), + "modelspec.license": recipe.license, + "modelspec.usage_hint": recipe.hint, + "modelspec.prediction_type": recipe.prediction, + "modelspec.dtype": str(recipe.dtype).split('.')[1], + "modelspec.hash_sha256": "", + "modelspec.thumbnail": get_thumbnail(), + "recipe": json.dumps({ + "base": os.path.basename(recipe.base) if recipe.base else "default", + "unet": os.path.basename(recipe.unet) if recipe.unet else "default", + "vae": os.path.basename(recipe.vae) if recipe.vae else "default", + "te1": os.path.basename(recipe.te1) if recipe.te1 else "default", + "te2": os.path.basename(recipe.te2) if recipe.te2 else "default", + "scheduler": recipe.scheduler or "default", + "lora": [f'{os.path.basename(k)}:{v}' for k, v in recipe.lora.items()], + }), + } + + +def save_model(pipe: diffusers.StableDiffusionXLPipeline): + author = recipe.author if len(recipe.author) > 0 else 'anonymous' + folder = os.path.join(shared.opts.diffusers_dir, f'models--{author}--{recipe.name}') + if len(recipe.version) > 0: + folder += f'-{recipe.version}' + if not recipe.diffusers or recipe.safetensors: + return + try: + yield msg('save') + yield msg(f'pretrained={folder}') + pipe.save_pretrained(folder, safe_serialization=True, push_to_hub=False) + with open(os.path.join(folder, 'vae', 'config.json'), 'r', encoding='utf8') as f: + vae_config = json.load(f) + vae_config['force_upcast'] = False + vae_config['scaling_factor'] = 0.13025 + vae_config['sample_size'] = 1024 + with open(os.path.join(folder, 'vae', 'config.json'), 'w', encoding='utf8') as f: + json.dump(vae_config, f, indent=2) + if recipe.safetensors: + fn = recipe.name + if len(recipe.version) > 0: + fn += f'-{recipe.version}' + if not os.path.isabs(fn): + fn = os.path.join(shared.opts.ckpt_dir, fn) + if not fn.endswith('.safetensors'): + fn += '.safetensors' + yield msg(f'safetensors={fn}') + from modules.merging import convert_sdxl + metadata = convert_sdxl(model_path=folder, checkpoint_path=fn, metadata=get_metadata()) + if 'modelspec.thumbnail' in metadata: + metadata['modelspec.thumbnail'] = f"{metadata['modelspec.thumbnail'].split(',')[0]}:{len(metadata['modelspec.thumbnail'])}" + yield msg(f'metadata={metadata}') + except Exception as e: + yield msg(f'save: {e}') + + +def merge(): + global pipeline # pylint: disable=global-statement + yield from load_base() + if pipeline is None: + return + pipeline = pipeline.to(device=devices.device, dtype=recipe.dtype) + yield from load_scheduler(pipeline) + yield from load_unet(pipeline) + yield from load_vae(pipeline) + yield from load_te1(pipeline) + yield from load_te2(pipeline) + yield from load_lora(pipeline) + yield from save_model(pipeline) + # pipeline = pipeline.to(device=devices.device, dtype=recipe.dtype) + # test_model(pipeline, '/tmp/merge.png') diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py index e5ddd2e85..c20fa00f4 100644 --- a/modules/sd_checkpoint.py +++ b/modules/sd_checkpoint.py @@ -52,7 +52,12 @@ class CheckpointInfo: relname, ext = os.path.splitext(relname) ext = ext.lower()[1:] - if os.path.isfile(filename): # ckpt or safetensor + if filename.lower() == 'none': + self.name = 'none' + self.relname = 'none' + self.sha256 = None + self.type = 'unknown' + elif os.path.isfile(filename): # ckpt or safetensor self.name = relname self.filename = filename self.sha256 = hashes.sha256_from_cache(self.filename, f"checkpoint/{relname}") @@ -173,7 +178,7 @@ def update_model_hashes(): return txt -def get_closet_checkpoint_match(s: str): +def get_closet_checkpoint_match(s: str) -> CheckpointInfo: if s.startswith('https://huggingface.co/'): model_name = s.replace('https://huggingface.co/', '') checkpoint_info = CheckpointInfo(model_name) # create a virutal model info diff --git a/modules/sd_models.py b/modules/sd_models.py index 9d0349a3f..6234c2b0b 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -437,6 +437,23 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con return sd_model +def set_defaults(sd_model, checkpoint_info): + sd_model.sd_model_hash = checkpoint_info.calculate_shorthash() # pylint: disable=attribute-defined-outside-init + sd_model.sd_checkpoint_info = checkpoint_info # pylint: disable=attribute-defined-outside-init + sd_model.sd_model_checkpoint = checkpoint_info.filename # pylint: disable=attribute-defined-outside-init + if hasattr(sd_model, "prior_pipe"): + sd_model.default_scheduler = copy.deepcopy(sd_model.prior_pipe.scheduler) if hasattr(sd_model.prior_pipe, "scheduler") else None + else: + sd_model.default_scheduler = copy.deepcopy(sd_model.scheduler) if hasattr(sd_model, "scheduler") else None + sd_model.is_sdxl = False # a1111 compatibility item + sd_model.is_sd2 = hasattr(sd_model, 'cond_stage_model') and hasattr(sd_model.cond_stage_model, 'model') # a1111 compatibility item + sd_model.is_sd1 = not sd_model.is_sd2 # a1111 compatibility item + sd_model.logvar = sd_model.logvar.to(devices.device) if hasattr(sd_model, 'logvar') else None # fix for training + shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256 + if hasattr(sd_model, "set_progress_bar_config"): + sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining}', ncols=80, colour='#327fba') + + def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=None, op='model', revision=None): # pylint: disable=unused-argument if timer is None: timer = Timer() @@ -515,20 +532,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No shared.log.error(f'Load {op}: name="{checkpoint_info.name if checkpoint_info is not None else None}" not loaded') return - sd_model.sd_model_hash = checkpoint_info.calculate_shorthash() # pylint: disable=attribute-defined-outside-init - sd_model.sd_checkpoint_info = checkpoint_info # pylint: disable=attribute-defined-outside-init - sd_model.sd_model_checkpoint = checkpoint_info.filename # pylint: disable=attribute-defined-outside-init - if hasattr(sd_model, "prior_pipe"): - sd_model.default_scheduler = copy.deepcopy(sd_model.prior_pipe.scheduler) if hasattr(sd_model.prior_pipe, "scheduler") else None - else: - sd_model.default_scheduler = copy.deepcopy(sd_model.scheduler) if hasattr(sd_model, "scheduler") else None - sd_model.is_sdxl = False # a1111 compatibility item - sd_model.is_sd2 = hasattr(sd_model, 'cond_stage_model') and hasattr(sd_model.cond_stage_model, 'model') # a1111 compatibility item - sd_model.is_sd1 = not sd_model.is_sd2 # a1111 compatibility item - sd_model.logvar = sd_model.logvar.to(devices.device) if hasattr(sd_model, 'logvar') else None # fix for training - shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256 - if hasattr(sd_model, "set_progress_bar_config"): - sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining}', ncols=80, colour='#327fba') + set_defaults(sd_model, checkpoint_info) if "Kandinsky" in sd_model.__class__.__name__: # need a special case sd_model.scheduler.name = 'DDIM' diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 398f608d1..accd6b0ed 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -7,7 +7,6 @@ from modules.sd_samplers_common import samples_to_image_grid, sample_to_image # debug = shared.log.trace if os.environ.get('SD_SAMPLER_DEBUG', None) is not None else lambda *args, **kwargs: None debug('Trace: SAMPLER') all_samplers = [] -all_samplers = [] all_samplers_map = {} samplers = all_samplers samplers_for_img2img = all_samplers @@ -49,7 +48,7 @@ def visible_sampler_names(): def create_sampler(name, model): if name is None or name == 'None': - return model.scheduler + return model.scheduler if model is not None else None try: current = model.scheduler.__class__.__name__ except Exception: @@ -86,28 +85,31 @@ def create_sampler(name, model): if not any(x in model.__class__.__name__ for x in FlowModels) and 'FlowMatch' in name: shared.log.warning(f'Sampler: default={current} target="{name}" class={model.__class__.__name__} flow-match scheduler unsupported') return None - # if any(x in model.__class__.__name__ for x in FlowModels) and 'FlowMatch' not in name: - # shared.log.warning(f'Sampler: default={current} target="{name}" class={model.__class__.__name__} linear scheduler unsupported') - # return None sampler = config.constructor(model) if sampler is None: sampler = config.constructor(model) - if sampler is None or sampler.sampler is None: - model.scheduler = copy.deepcopy(model.default_scheduler) + if model is not None: + if sampler is None or sampler.sampler is None: + model.scheduler = copy.deepcopy(model.default_scheduler) + else: + model.scheduler = sampler.sampler + if not hasattr(model, 'scheduler_config'): + model.scheduler_config = sampler.sampler.config.copy() if hasattr(sampler, 'sampler') and hasattr(sampler.sampler, 'config') else {} + if hasattr(model, "prior_pipe") and hasattr(model.prior_pipe, "scheduler"): + model.prior_pipe.scheduler = sampler.sampler + model.prior_pipe.scheduler.config.clip_sample = False + if "flow" in model.scheduler.__class__.__name__.lower(): + shared.state.prediction_type = "flow_prediction" + elif hasattr(model.scheduler, "config") and hasattr(model.scheduler.config, "prediction_type"): + shared.state.prediction_type = model.scheduler.config.prediction_type + if model is not None: + clean_config = {k: v for k, v in model.scheduler.config.items() if not k.startswith('_') and v is not None and v is not False} + cls = model.scheduler.__class__.__name__ else: - model.scheduler = sampler.sampler - if not hasattr(model, 'scheduler_config'): - model.scheduler_config = sampler.sampler.config.copy() if hasattr(sampler, 'sampler') and hasattr(sampler.sampler, 'config') else {} - if hasattr(model, "prior_pipe") and hasattr(model.prior_pipe, "scheduler"): - model.prior_pipe.scheduler = sampler.sampler - model.prior_pipe.scheduler.config.clip_sample = False - if "flow" in model.scheduler.__class__.__name__.lower(): - shared.state.prediction_type = "flow_prediction" - elif hasattr(model.scheduler, "config") and hasattr(model.scheduler.config, "prediction_type"): - shared.state.prediction_type = model.scheduler.config.prediction_type - clean_config = {k: v for k, v in model.scheduler.config.items() if not k.startswith('_') and v is not None and v is not False} + clean_config = {k: v for k, v in sampler.sampler.config.items() if not k.startswith('_') and v is not None and v is not False} + cls = sampler.sampler.__class__.__name__ name = sampler.name if sampler is not None and sampler.sampler is not None else 'Default' - shared.log.debug(f'Sampler: "{name}" class={model.scheduler.__class__.__name__} config={clean_config}') + shared.log.debug(f'Sampler: "{name}" class={cls} config={clean_config}') return sampler.sampler else: return None diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 6b2de72aa..a8b77bc90 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -15,26 +15,20 @@ try: CMStochasticIterativeScheduler, UniPCMultistepScheduler, DDIMScheduler, - EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, EDMEulerScheduler, FlowMatchEulerDiscreteScheduler, - DEISMultistepScheduler, SASolverScheduler, - DPMSolverSinglestepScheduler, DPMSolverMultistepScheduler, EDMDPMSolverMultistepScheduler, CosineDPMSolverMultistepScheduler, DPMSolverSDEScheduler, - HeunDiscreteScheduler, FlowMatchHeunDiscreteScheduler, - LCMScheduler, - PNDMScheduler, IPNDMScheduler, DDPMScheduler, @@ -172,14 +166,17 @@ class DiffusionSampler: return self.name = name self.config = {} - if not hasattr(model, 'scheduler'): - return - if getattr(model, "default_scheduler", None) is None: # sanity check + self.sampler = None + # if not hasattr(model, 'scheduler'): + # return + if getattr(model, "default_scheduler", None) is None and (model is not None): # sanity check model.default_scheduler = copy.deepcopy(model.scheduler) for key, value in config.get('All', {}).items(): # apply global defaults self.config[key] = value debug_log(f'Sampler: all="{self.config}"') - if hasattr(model.default_scheduler, 'scheduler_config'): # find model defaults + if model is None: + orig_config = {} + elif hasattr(model.default_scheduler, 'scheduler_config'): # find model defaults orig_config = model.default_scheduler.scheduler_config else: orig_config = model.default_scheduler.config diff --git a/modules/ui_models.py b/modules/ui_models.py index 5d5b452e2..23a39f317 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -4,18 +4,16 @@ import json import inspect from datetime import datetime import gradio as gr -from modules import sd_models, sd_vae, extras +from modules import errors, sd_models, sd_vae, extras, sd_samplers, ui_symbols, hashes from modules.ui_components import ToolButton from modules.ui_common import create_refresh_button from modules.call_queue import wrap_gradio_gpu_call from modules.shared import opts, log, req, readfile, max_workers, native -import modules.ui_symbols -import modules.errors -import modules.hashes from modules.merging import merge_methods from modules.merging.merge_utils import BETA_METHODS, TRIPLE_METHODS, interpolate from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS + search_metadata_civit = None extra_ui = [] @@ -32,9 +30,6 @@ def create_ui(): with gr.Column(elem_id='models_input_container', scale=3): - def gr_show(visible=True): - return {"visible": visible, "__type__": "update"} - with gr.Tab(label="Current"): def analyze(): from modules import modelstats @@ -57,45 +52,6 @@ def create_ui(): model_analyze.click(fn=analyze, inputs=[], outputs=[model_desc, model_modules, model_meta]) - with gr.Tab(label="Convert"): - with gr.Row(): - model_name = gr.Dropdown(sd_models.checkpoint_titles(), label="Original model") - create_refresh_button(model_name, sd_models.list_models, lambda: {"choices": sd_models.checkpoint_titles()}, "refresh_checkpoint_Z") - with gr.Row(): - custom_name = gr.Textbox(label="Output model name") - with gr.Row(): - precision = gr.Radio(choices=["fp32", "fp16", "bf16"], value="fp16", label="Model precision") - m_type = gr.Radio(choices=["disabled", "no-ema", "ema-only"], value="disabled", label="Model pruning methods") - with gr.Row(): - checkpoint_formats = gr.CheckboxGroup(choices=["ckpt", "safetensors"], value=["safetensors"], label="Model Format") - with gr.Row(): - show_extra_options = gr.Checkbox(label="Show extra options", value=False) - fix_clip = gr.Checkbox(label="Fix clip", value=False) - with gr.Row(visible=False) as extra_options: - specific_part_conv = ["copy", "convert", "delete"] - unet_conv = gr.Dropdown(specific_part_conv, value="convert", label="unet") - text_encoder_conv = gr.Dropdown(specific_part_conv, value="convert", label="text encoder") - vae_conv = gr.Dropdown(specific_part_conv, value="convert", label="vae") - others_conv = gr.Dropdown(specific_part_conv, value="convert", label="others") - - show_extra_options.change(fn=lambda x: gr_show(x), inputs=[show_extra_options], outputs=[extra_options]) - - model_converter_convert = gr.Button(label="Convert", variant='primary') - model_converter_convert.click( - fn=extras.run_modelconvert, - inputs=[ - model_name, - checkpoint_formats, - precision, m_type, custom_name, - unet_conv, - text_encoder_conv, - vae_conv, - others_conv, - fix_clip - ], - outputs=[models_outcome] - ) - with gr.Tab(label="Merge"): def sd_model_choices(): return ['None'] + sd_models.checkpoint_titles() @@ -222,7 +178,7 @@ def create_ui(): try: results = extras.run_modelmerger(dummy_component, **kwargs) except Exception as e: - modules.errors.display(e, 'Merge') + errors.display(e, 'Merge') sd_models.list_models() # to remove the potentially missing models from the list return [*[gr.Dropdown.update(choices=sd_models.checkpoint_titles()) for _ in range(4)], f"Error merging checkpoints: {e}"] return results @@ -334,6 +290,76 @@ def create_ui(): ] ) + with gr.Tab(label="Modules"): + with gr.Row(): + with gr.Column(scale=3): + model_type = gr.Dropdown(label="Model type", choices=['sd15', 'sdxl', 'sd21', 'sd35', 'flux.1'], value='sdxl', interactive=False) + with gr.Column(scale=5): + with gr.Row(): + model_name = gr.Dropdown(sd_models.checkpoint_titles(), label="Input model") + create_refresh_button(model_name, sd_models.list_models, lambda: {"choices": sd_models.checkpoint_titles()}, "refresh_checkpoint_Z") + with gr.Column(scale=5): + custom_name = gr.Textbox(label="Output model", placeholder="Output model path") + with gr.Row(): + with gr.Column(scale=3): + gr.HTML('Model components
Specify the components to include
Paths can be relative or absolute
') + with gr.Column(scale=5): + comp_unet = gr.Textbox(placeholder="UNet model", show_label=False) + comp_vae = gr.Textbox(placeholder="VAE model", show_label=False) + with gr.Column(scale=5): + comp_te1 = gr.Textbox(placeholder="Text encoder 1", show_label=False) + comp_te2 = gr.Textbox(placeholder="Text encoder 2", show_label=False) + with gr.Row(): + with gr.Column(scale=3): + gr.HTML('Model settings
') + with gr.Column(scale=10): + with gr.Row(): + precision = gr.Dropdown(label="Model precision", choices=["fp32", "fp16", "bf16"], value="fp16") + comp_scheduler = gr.Dropdown(label="Sampler", choices=[s.name for s in sd_samplers.samplers if s.constructor is not None]) + comp_prediction = gr.Dropdown(Label="Prediction type", choices=["epsilon", "v"], value="epsilon") + with gr.Row(): + with gr.Column(scale=3): + gr.HTML('Merge LoRA
') + with gr.Column(scale=9): + comp_lora = gr.Textbox(label="Comma separated list with optional strength per LoRA", placeholder="LoRA models") + with gr.Column(scale=1): + comp_fuse = gr.Number(label="Fuse strength", value=1.0) + + with gr.Row(): + gr.HTML('
') + with gr.Row(): + with gr.Column(scale=2): + gr.HTML('Model metadata
') + with gr.Column(scale=5): + meta_author = gr.Textbox(placeholder="Author name", show_label=False) + meta_version = gr.Textbox(placeholder="Model version", show_label=False) + meta_license = gr.Textbox(placeholder="Model license", show_label=False) + with gr.Column(scale=5): + meta_desc = gr.Textbox(placeholder="Model description", lines=3, show_label=False) + meta_hint = gr.Textbox(placeholder="Model hint", lines=3, show_label=False) + with gr.Column(scale=3): + meta_thumbnail = gr.Image(label="Thumbnail", type='pil', source='upload') + with gr.Row(): + gr.HTML('Note: Save is optional as you can merge in-memory and use newly created model immediately') + with gr.Row(): + create_diffusers = gr.Checkbox(label="Save diffusers", value=True) + create_safetensors = gr.Checkbox(label="Save safetensors", value=True) + debug = gr.Checkbox(label="Debug info", value=False) + + model_modules_btn = gr.Button(label="Modules", variant='primary') + model_modules_btn.click( + fn=extras.run_model_modules, + inputs=[ + model_type, model_name, custom_name, + comp_unet, comp_vae, comp_te1, comp_te2, + precision, comp_scheduler, comp_prediction, + comp_lora, comp_fuse, + meta_author, meta_version, meta_license, meta_desc, meta_hint, meta_thumbnail, + create_diffusers, create_safetensors, debug, + ], + outputs=[models_outcome] + ) + with gr.Tab(label="Validate"): model_headers = ['name', 'type', 'filename', 'hash', 'added', 'size', 'metadata'] model_data = [] @@ -407,7 +433,7 @@ def create_ui(): gr.HTML('

Search for models

Select a model from the search results to download

') with gr.Row(): hf_search_text = gr.Textbox('', label='Search models', placeholder='search huggingface models') - hf_search_btn = ToolButton(value=modules.ui_symbols.search, label="Search") + hf_search_btn = ToolButton(value=ui_symbols.search, label="Search") with gr.Row(): with gr.Column(scale=2): with gr.Row(): @@ -562,7 +588,7 @@ def create_ui(): found = True break if not found and rehash and os.stat(item['filename']).st_size < (1024 * 1024 * 1024): - sha = modules.hashes.calculate_sha256(item['filename'], quiet=True)[:10] + sha = hashes.calculate_sha256(item['filename'], quiet=True)[:10] r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{sha}') log.debug(f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}') if r.status_code == 200: @@ -622,7 +648,7 @@ def create_ui(): with gr.Row(): civit_search_text = gr.Textbox('', label='Search models', placeholder='keyword') civit_search_tag = gr.Textbox('', label='', placeholder='tags') - civit_search_btn = ToolButton(value=modules.ui_symbols.search, label="Search", interactive=True) + civit_search_btn = ToolButton(value=ui_symbols.search, label="Search", interactive=True) with gr.Row(): civit_search_res = gr.HTML('') with gr.Row(): @@ -718,13 +744,12 @@ def create_ui(): def civit_update_metadata(): nonlocal update_data log.debug('CivitAI update metadata: models') - from modules.ui_extra_networks import get_pages - from modules.modelloader import download_civit_meta + from modules import ui_extra_networks, modelloader res = [] - pages = get_pages('Model') + pages = ui_extra_networks.get_pages('Model') if len(pages) == 0: return 'CivitAI update metadata: no models found' - page: modules.ui_extra_networks.ExtraNetworksPage = pages[0] + page: ui_extra_networks.ExtraNetworksPage = pages[0] table_data = [] update_data.clear() all_hashes = [(item.get('hash', None) or 'XXXXXXXX').upper()[:8] for item in page.list_items()] @@ -738,7 +763,7 @@ def create_ui(): if r.status_code == 200: d = r.json() model.id = d['modelId'] - download_civit_meta(model.fn, model.id) + modelloader.download_civit_meta(model.fn, model.id) fn = os.path.splitext(item['filename'])[0] + '.json' model.meta = readfile(fn, silent=True) model.name = model.meta.get('name', model.name) From 0d67a6d979bd8620ae389b2320f242ef6b34e827 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 19 Jan 2025 15:00:26 -0500 Subject: [PATCH 12/66] add asymmetric vae upscaler Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + modules/upscaler_simple.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index eabd560f5..3c7e63102 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ - in addition as standard behavior of detect & run-generate, it can now also run face-restore models - included models are: *CodeFormer, RestoreFormer, GFPGan, GPEN-BFR* - **Other**: + - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method - **ipex**: update supported torch versions - **gallery**: add http fallback for slow/unreliable links - **splash**: add legacy mode indicator on splash screen diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py index fcf4f2410..bd0c6050f 100644 --- a/modules/upscaler_simple.py +++ b/modules/upscaler_simple.py @@ -89,3 +89,31 @@ class UpscalerLatent(Upscaler): else: raise ValueError(f"Latent upscale: model={selected_model} unknown") return F.interpolate(img, size=(h, w), mode=mode, antialias=antialias) + + +class UpscalerAsymmetricVAE(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "Asymmetric VAE" + self.vae = None + self.scalers = [ + UpscalerData("Asymmetric VAE", None, self), + ] + + def do_upscale(self, img: Image, selected_model=None): + import torchvision.transforms.functional as F + import diffusers + from modules import shared, devices + + if self.vae is None: + self.vae = diffusers.AsymmetricAutoencoderKL.from_pretrained("Heasterian/AsymmetricAutoencoderKLUpscaler", cache_dir=shared.opts.hfcache_dir) + self.vae.requires_grad_(False) + self.vae = self.vae.to(device=devices.device, dtype=devices.dtype) + self.vae.eval() + img = img.resize((8 * (img.width // 8), 8 * (img.height // 8)), resample=Image.Resampling.BILINEAR).convert('RGB') + tensor = (F.pil_to_tensor(img).unsqueeze(0) / 255.0).to(device=devices.device, dtype=devices.dtype) + self.vae = self.vae.to(device=devices.device) + tensor = self.vae(tensor).sample + upscaled = F.to_pil_image(tensor.squeeze().clamp(0.0, 1.0).float().cpu()) + self.vae = self.vae.to(device=devices.cpu) + return upscaled From f5a910f7194d3a0f01b0e97907e3d86dc8314030 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 20 Jan 2025 16:57:55 -0500 Subject: [PATCH 13/66] experimental reswapper Signed-off-by: Vladimir Mandic --- modules/face/__init__.py | 26 +++-- modules/face/faceswap.py | 3 +- modules/face/insightface.py | 11 +- modules/face/photomaker.py | 2 +- modules/face/reswapper.py | 111 +++++++++++++++++++++ modules/face/reswapper_model.py | 127 ++++++++++++++++++++++++ modules/face/reswapper_utils.py | 171 ++++++++++++++++++++++++++++++++ modules/loader.py | 4 +- 8 files changed, 441 insertions(+), 14 deletions(-) create mode 100644 modules/face/reswapper.py create mode 100644 modules/face/reswapper_model.py create mode 100644 modules/face/reswapper_utils.py diff --git a/modules/face/__init__.py b/modules/face/__init__.py index c18da6e2e..c851b4ee4 100644 --- a/modules/face/__init__.py +++ b/modules/face/__init__.py @@ -36,6 +36,7 @@ class Script(scripts.Script): def mode_change(self, mode): return [ + gr.update(visible=mode=='ReSwapper'), gr.update(visible=mode=='FaceID'), gr.update(visible=mode=='FaceSwap'), gr.update(visible=mode=='InstantID'), @@ -47,7 +48,17 @@ class Script(scripts.Script): with gr.Row(): gr.HTML("  Face: Multiple ID Transfers
") with gr.Row(): + models = ['None', 'FaceID', 'FaceSwap', 'InstantID', 'PhotoMaker'] + if shared.cmd_opts.experimental: + models.append('ReSwapper') mode = gr.Dropdown(label='Mode', choices=['None', 'FaceID', 'FaceSwap', 'InstantID', 'PhotoMaker'], value='None') + with gr.Group(visible=False) as cfg_reswapper: + with gr.Row(): + gr.HTML('  ReSwapper
') + with gr.Row(): + from modules.face.reswapper import RESWAPPER_MODELS + reswapper_model = gr.Dropdown(choices=list(RESWAPPER_MODELS), label='ReSwapper Model', value='ReSwapper 256 0.2') + reswapper_original = gr.Checkbox(label='Return original images', value=False) with gr.Group(visible=False) as cfg_faceid: with gr.Row(): gr.HTML('  Tencent AI Lab IP-Adapter FaceID
') @@ -85,11 +96,11 @@ class Script(scripts.Script): with gr.Row(): gallery = gr.Gallery(show_label=False, value=[]) files.change(fn=self.load_images, inputs=[files], outputs=[gallery]) - mode.change(fn=self.mode_change, inputs=[mode], outputs=[cfg_faceid, cfg_faceswap, cfg_instantid, cfg_photomaker]) + mode.change(fn=self.mode_change, inputs=[mode], outputs=[cfg_reswapper, cfg_faceid, cfg_faceswap, cfg_instantid, cfg_photomaker]) - return [mode, gallery, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_trigger, pm_strength, pm_start, fs_cache] + return [mode, gallery, reswapper_model, reswapper_original, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_trigger, pm_strength, pm_start, fs_cache] - def run(self, p: processing.StableDiffusionProcessing, mode, input_images, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_trigger, pm_strength, pm_start, fs_cache): # pylint: disable=arguments-differ, unused-argument + def run(self, p: processing.StableDiffusionProcessing, mode, input_images, reswapper_model, reswapper_original, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_trigger, pm_strength, pm_start, fs_cache): # pylint: disable=arguments-differ, unused-argument if not shared.native: return None if mode == 'None': @@ -134,11 +145,12 @@ class Script(scripts.Script): from modules.face.insightface import get_app app=get_app('buffalo_l') from modules.face.faceswap import face_swap - if shared.opts.save_images_before_detailer and not p.do_not_save_samples: - for i, image in enumerate(processed.images): - info = processing.create_infotext(p, index=i) - images.save_image(image, path=p.outpath_samples, seed=p.all_seeds[i], prompt=p.all_prompts[i], info=info, p=p, suffix="-before-faceswap") processed.images = face_swap(p, app=app, input_images=processed.images, source_image=input_images[0], cache=fs_cache) + elif mode == 'ReSwapper': + from modules.face.insightface import get_app + app = get_app('buffalo_l', resolution=512) + from modules.face.reswapper import reswapper + processed.images = reswapper(p, app=app, source_images=processed.images, target_images=input_images, model_name=reswapper_model, original=reswapper_original) processed.info = processed.infotext(p, 0) processed.infotexts = [processed.info] diff --git a/modules/face/faceswap.py b/modules/face/faceswap.py index 9fb2bfb37..d7d0a32a5 100644 --- a/modules/face/faceswap.py +++ b/modules/face/faceswap.py @@ -16,7 +16,8 @@ def face_swap(p: processing.StableDiffusionProcessing, app, input_images: List[I import insightface.model_zoo global swapper # pylint: disable=global-statement if swapper is None: - model_path = hf.hf_hub_download(repo_id='ezioruan/inswapper_128.onnx', filename='inswapper_128.onnx', cache_dir=shared.opts.diffusers_dir) + model_path = hf.hf_hub_download(repo_id='ezioruan/inswapper_128.onnx', filename='inswapper_128.onnx', cache_dir=shared.opts.hfcache_dir) + # model_path = hf.hf_hub_download(repo_id='somanchiu/reswapper', filename='reswapper_256-1567500_originalInswapperClassCompatible.onnx', cache_dir=shared.opts.hfcache_dir) router: insightface.model_zoo.model_zoo.INSwapper = insightface.model_zoo.model_zoo.ModelRouter(model_path) swapper = router.get_model() diff --git a/modules/face/insightface.py b/modules/face/insightface.py index a59e48b41..4f817c8c7 100644 --- a/modules/face/insightface.py +++ b/modules/face/insightface.py @@ -7,7 +7,7 @@ insightface_app = None instightface_mp = None -def get_app(mp_name): +def get_app(mp_name, threshold=0.5, resolution=640): global insightface_app, instightface_mp # pylint: disable=global-statement from installer import install, installed @@ -19,7 +19,10 @@ def get_app(mp_name): install('git+https://github.com/tencent-ailab/IP-Adapter.git', 'ip_adapter', ignore=False) if insightface_app is None or mp_name != instightface_mp: - from insightface.app import FaceAnalysis + from insightface.model_zoo import model_zoo + from insightface.app import face_analysis + model_zoo.print = lambda *args, **kwargs: None + face_analysis.print = lambda *args, **kwargs: None import huggingface_hub as hf import zipfile log.debug(f"InsightFace: mp={mp_name} provider={devices.onnx}") @@ -45,7 +48,7 @@ def get_app(mp_name): 'download': False, 'download_zip': False, } - insightface_app = FaceAnalysis(name=mp_name, providers=devices.onnx, **kwargs) + insightface_app = face_analysis.FaceAnalysis(name=mp_name, providers=devices.onnx, **kwargs) instightface_mp = mp_name - insightface_app.prepare(ctx_id=0, det_thresh=0.5, det_size=(640, 640)) + insightface_app.prepare(ctx_id=0, det_thresh=threshold, det_size=(resolution, resolution)) return insightface_app diff --git a/modules/face/photomaker.py b/modules/face/photomaker.py index b89f28a10..2b950d572 100644 --- a/modules/face/photomaker.py +++ b/modules/face/photomaker.py @@ -54,7 +54,7 @@ def photo_maker(p: processing.StableDiffusionProcessing, input_images, trigger, p.task_args['start_merge_step'] = int(start * p.steps) p.task_args['prompt'] = p.all_prompts[0] if p.all_prompts is not None else p.prompt - photomaker_path = hf.hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model", cache_dir=shared.opts.diffusers_dir) + photomaker_path = hf.hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model", cache_dir=shared.opts.hfcache_dir) shared.log.debug(f'PhotoMaker: model={photomaker_path} images={len(input_images)} trigger={trigger} args={p.task_args}') # load photomaker adapter diff --git a/modules/face/reswapper.py b/modules/face/reswapper.py new file mode 100644 index 000000000..77328a1ff --- /dev/null +++ b/modules/face/reswapper.py @@ -0,0 +1,111 @@ +from typing import List +import os +import cv2 +import torch +import numpy as np +import huggingface_hub as hf +from PIL import Image +from modules import processing, shared, devices + +RESWAPPER_REPO = 'somanchiu/reswapper' +RESWAPPER_MODELS = { + "ReSwapper 256 0.2": "reswapper_256-1567500.pth", + "ReSwapper 256 0.1": "reswapper_256-1399500.pth", + "ReSwapper 128 0.2": "reswapper-429500.pth", + "ReSwapper 128 0.1": "reswapper-1019500.pth", +} +reswapper_model = None +reswapper_name = None +debug = shared.log.trace if os.environ.get("SD_FACE_DEBUG", None) is not None else lambda *args, **kwargs: None +dtype = devices.dtype + +def get_model(model_name: str): + global reswapper_model, reswapper_name # pylint: disable=global-statement + if reswapper_model is None or reswapper_name != model_name: + try: + fn = RESWAPPER_MODELS.get(model_name) + url = hf.hf_hub_download(repo_id=RESWAPPER_REPO, filename=fn, repo_type="model", cache_dir=shared.opts.hfcache_dir) + from modules.face.reswapper_model import ReSwapperModel + reswapper_model = ReSwapperModel() + reswapper_model.load_state_dict(torch.load(url, map_location='cpu'), strict=False) + reswapper_model = reswapper_model.to(device=devices.device, dtype=dtype) + reswapper_model.eval() + reswapper_name = model_name + shared.log.info(f'ReSwapper: model="{model_name}" url="{url}" cls={reswapper_model.__class__.__name__}') + if reswapper_model is None: + shared.log.error(f'ReSwapper: model="{model_name}" fn="{fn}" url="{url}" failed to load model') + return reswapper_model + except Exception as e: + shared.log.error(f'ReSwapper: model="{model_name}" fn="{fn}" url="{url}" {e}') + return reswapper_model + + +def reswapper( + p: processing.StableDiffusionProcessing, + app, + source_images: List[Image.Image], + target_images: List[Image.Image], + model_name: str, + original: bool, +): + from modules.face import reswapper_utils as utils + if source_images is None or len(source_images) == 0: + shared.log.warning('ReSwapper: no input images') + return None + + processed_images = [] + if original: + processed_images += source_images + + model = get_model(model_name) + if model is None: + return source_images + model = model.to(device=devices.device) + + i = 0 + for x, image in enumerate(source_images): + image = image.convert('RGB') + source_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + source_faces = app.get(source_np) + if len(source_faces) == 0: + shared.log.error(f"ReSwapper: image={x+1} no source faces found") + return source_images + if len(source_faces) != len(target_images): + shared.log.warning(f"ReSwapper: image={x+1} source-faces={len(source_faces)} target-images={len(target_images)}") + for y, source_face in enumerate(source_faces): + target_image = target_images[y] if y < len(target_images) else target_images[-1] + target_image = target_image.convert('RGB') + target_np = cv2.cvtColor(np.array(target_image), cv2.COLOR_RGB2BGR) + target_faces = app.get(target_np) + if len(target_faces) != 1: + shared.log.error(f"ReSwapper: image={x+1} source-faces={y+1} target-faces={len(target_faces)} must be exactly one") + return source_images + target_face = target_faces[0] + source_str = f'score:{source_face.det_score:.2f} gender:{"female" if source_face.gender==0 else "male"} age:{source_face.age}' + target_str = f'score:{target_face.det_score:.2f} gender:{"female" if target_face.gender==0 else "male"} age:{target_face.age}' + shared.log.debug(f'ReSwapper image={x+1} face={y+1} source="{source_str}" target="{target_str}"') + + source_latent = utils.getLatent(source_face) + source_tensor = torch.from_numpy(source_latent).to(device=devices.device, dtype=dtype) + + resolution = 256 if '256' in model_name else 128 + target_np = cv2.cvtColor(np.array(target_image), cv2.COLOR_RGB2BGR) + target_aligned, M = utils.norm_crop2(target_np, target_face.kps, resolution) + target_blob = utils.getBlob(target_aligned, (resolution, resolution)) + target_tensor = torch.from_numpy(target_blob).to(device=devices.device, dtype=dtype) + + with devices.inference_context(): + swapped_tensor = model(target_tensor, source_tensor) + swapped_tensor = swapped_tensor.float() + + swapped_face = (swapped_tensor.squeeze().permute(1, 2, 0).cpu().detach().numpy() * 255).astype(np.uint8) + swapped_face = cv2.cvtColor(swapped_face, cv2.COLOR_RGB2BGR) + swapped_np = utils.blend_swapped_image(swapped_face, source_np, M) + swapped_image = Image.fromarray(cv2.cvtColor(swapped_np, cv2.COLOR_BGR2RGB)) + processed_images.append(swapped_image) + i += 1 + + p.extra_generation_params['ReSwapper'] = f'faces={i}' + devices.torch_gc() + + return processed_images diff --git a/modules/face/reswapper_model.py b/modules/face/reswapper_model.py new file mode 100644 index 000000000..de68d8566 --- /dev/null +++ b/modules/face/reswapper_model.py @@ -0,0 +1,127 @@ +# original: + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class ReSwapperModel(nn.Module): + def __init__(self): + super(ReSwapperModel, self).__init__() + + # self.pad = nn.ReflectionPad2d(3) + # Encoder for target face + self.target_encoder = nn.Sequential( + # self.pad, + nn.Conv2d(3, 128, kernel_size=7, stride=1, padding=0), + nn.LeakyReLU(0.2), + nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2), + nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1), + nn.LeakyReLU(0.2), + nn.Conv2d(512, 1024, kernel_size=3, stride=2, padding=1), + nn.LeakyReLU(0.2), + ) + + # for style_block in self.target_encoder: + # for param in style_block.parameters(): + # param.requires_grad = False + + # Style blocks + self.style_blocks = nn.ModuleList([ + StyleBlock(1024, 1024, blockIndex) for blockIndex in range(6) + ]) + + # Decoder (upsampling) + self.decoder = nn.Sequential( + nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2) + ) + + self.decoderPart1 = nn.Sequential( + nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2), + nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2) + ) + + self.decoderPart2 = nn.Sequential( + # self.pad, + nn.Conv2d(128, 3, kernel_size=7, stride=1, padding=0), + nn.Tanh() + ) + + def forward(self, target, source): + # Encode target face + target = F.pad(target, pad=(3, 3, 3, 3), mode='reflect') + + target_features = self.target_encoder(target) + + # Apply style blocks + x = target_features + for style_block in self.style_blocks: + x = style_block(x, source) + + # Decode + # x = F.interpolate(x, scale_factor=2, mode='linear') + x = F.upsample( + x, + scale_factor=2, # specify the desired height and width + mode='bilinear', # 'linear' in 2D is called 'bilinear' + align_corners=False # this is typically False for ONNX compatibility + ) + output = self.decoder(x) + + output = F.upsample( + output, + scale_factor=2, # specify the desired height and width + mode='bilinear', # 'linear' in 2D is called 'bilinear' + align_corners=False # this is typically False for ONNX compatibility + ) + output = self.decoderPart1(output) + + output = F.pad(output, pad=(3, 3, 3, 3), mode='reflect') + + output = self.decoderPart2(output) + + return (output + 1) / 2 + +class StyleBlock(nn.Module): + def __init__(self, in_channels, out_channels, blockIndex): + super(StyleBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0) + self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0) + self.style1 = nn.Linear(512, 2048) + self.style2 = nn.Linear(512, 2048) + self.style = [self.style1, self.style2] + + self.blockIndex = blockIndex + + def normalizeConvRMS(self, conv): + x = conv - torch.mean(conv, dim=[2, 3], keepdim=True) # centeredConv + squareX = x * x + meanSquaredX = torch.mean(squareX, dim=[2, 3], keepdim=True) + rms = torch.sqrt(meanSquaredX + 0.00000001) + return (1 / rms) * x + + def forward(self, residual, style): + # print(f'Forward: {self.blockIndex}') + style1024 = [] + for index in range(2): + style1 = self.style[index](style) + style1 = torch.unsqueeze(style1, 2) + style1 = torch.unsqueeze(style1, 3) + first_half = style1[:, :1024, :, :] + second_half = style1[:, 1024:, :, :] + + style1024.append([first_half, second_half]) + + conv1 = self.normalizeConvRMS(self.conv1(F.pad(residual, pad=(1, 1, 1, 1), mode='reflect'))) + + out = F.relu(conv1 * style1024[0][0] + style1024[0][1]) + + out = F.pad(out, pad=(1, 1, 1, 1), mode='reflect') + + conv2 = self.normalizeConvRMS(self.conv2(out)) + out = conv2 * style1024[1][0] + style1024[1][1] + + return residual + out diff --git a/modules/face/reswapper_utils.py b/modules/face/reswapper_utils.py new file mode 100644 index 000000000..f5dbf0c93 --- /dev/null +++ b/modules/face/reswapper_utils.py @@ -0,0 +1,171 @@ +import cv2 +import numpy as np +from skimage import transform as trans + + +### https://github.com/somanchiu/ReSwapper/blob/GAN/Image.py + +input_std = 255.0 +input_mean = 0.0 + + +def get_emap(): + emap = np.load("modules/face/reswapper_emap.npy") # https://github.com/somanchiu/ReSwapper/blob/GAN/emap.npy + return emap + + +def postprocess_face(face_tensor): + face_tensor = face_tensor.squeeze().cpu().detach() + face_np = (face_tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8) + face_np = cv2.cvtColor(face_np, cv2.COLOR_RGB2BGR) + return face_np + +def getBlob(aimg, input_size = (128, 128)): + blob = cv2.dnn.blobFromImage(aimg, 1.0 / input_std, input_size, (input_mean, input_mean, input_mean), swapRB=True) + return blob + + +def getLatent(source_face): + latent = source_face.normed_embedding.reshape((1,-1)) + emap = get_emap() + latent = np.dot(latent, emap) + latent /= np.linalg.norm(latent) + return latent + + +def blend_swapped_image(swapped_face, target_image, M): + h, w = target_image.shape[:2] + M_inv = cv2.invertAffineTransform(M) + warped_face = cv2.warpAffine(swapped_face, M_inv, (w, h),borderValue=0.0) + img_white = np.full((swapped_face.shape[0], swapped_face.shape[1]), 255, dtype=np.float32) + img_mask = cv2.warpAffine(img_white, M_inv, (w, h), borderValue=0.0) + img_mask[img_mask > 20] = 255 + mask_h_inds, mask_w_inds = np.where(img_mask == 255) + if len(mask_h_inds) > 0 and len(mask_w_inds) > 0: # safety check + mask_h = np.max(mask_h_inds) - np.min(mask_h_inds) + mask_w = np.max(mask_w_inds) - np.min(mask_w_inds) + mask_size = int(np.sqrt(mask_h * mask_w)) + k = max(mask_size // 10, 10) + kernel = np.ones((k, k), np.uint8) + img_mask = cv2.erode(img_mask, kernel, iterations=1) + k = max(mask_size // 20, 5) + kernel_size = (k, k) + blur_size = tuple(2 * i + 1 for i in kernel_size) + img_mask = cv2.GaussianBlur(img_mask, blur_size, 0) + img_mask = img_mask / 255.0 + img_mask = np.reshape(img_mask, [img_mask.shape[0], img_mask.shape[1], 1]) + result = img_mask * warped_face + (1 - img_mask) * target_image.astype(np.float32) + result = result.astype(np.uint8) + return result + + +def drawKeypoints(image, keypoints, colorBGR, keypointsRadius=2): + for kp in keypoints: + x, y = int(kp[0]), int(kp[1]) + cv2.circle(image, (x, y), radius=keypointsRadius, color=colorBGR, thickness=-1) # BGR format, -1 means filled circle + + +### https://github.com/somanchiu/ReSwapper/blob/GAN/face_align.py + +arcface_dst = np.array( + [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], + [41.5493, 92.3655], [70.7299, 92.2041]], + dtype=np.float32) + + +def estimate_norm(lmk, image_size=112,mode='arcface'): # pylint: disable=unused-argument + if image_size%112==0: + ratio = float(image_size)/112.0 + diff_x = 0 + else: + ratio = float(image_size)/128.0 + diff_x = 8.0*ratio + ratio = float(image_size)/112.0 + diff_x = 0 + dst = arcface_dst * ratio + dst[:,0] += diff_x + if image_size%112==0: + ratio = float(image_size)/112.0 + diff_x = 0 + else: + ratio = float(image_size)/128.0 + diff_x = 8.0*ratio + dst = arcface_dst * ratio + dst[:,0] += diff_x + tform = trans.SimilarityTransform() + tform.estimate(lmk, dst) + M = tform.params[0:2, :] + return M + + +def norm_crop(img, landmark, image_size=112, mode='arcface'): + M = estimate_norm(landmark, image_size, mode) + warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0) + return warped + + +def norm_crop2(img, landmark, image_size=112, mode='arcface'): + M = estimate_norm(landmark, image_size, mode) + warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0) + return warped, M + + +def square_crop(im, S): + if im.shape[0] > im.shape[1]: + height = S + width = int(float(im.shape[1]) / im.shape[0] * S) + scale = float(S) / im.shape[0] + else: + width = S + height = int(float(im.shape[0]) / im.shape[1] * S) + scale = float(S) / im.shape[1] + resized_im = cv2.resize(im, (width, height)) + det_im = np.zeros((S, S, 3), dtype=np.uint8) + det_im[:resized_im.shape[0], :resized_im.shape[1], :] = resized_im + return det_im, scale + + +def transform(data, center, output_size, scale, rotation): + scale_ratio = scale + rot = float(rotation) * np.pi / 180.0 + t1 = trans.SimilarityTransform(scale=scale_ratio) + cx = center[0] * scale_ratio + cy = center[1] * scale_ratio + t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy)) + t3 = trans.SimilarityTransform(rotation=rot) + t4 = trans.SimilarityTransform(translation=(output_size / 2, output_size / 2)) + t = t1 + t2 + t3 + t4 + M = t.params[0:2] + cropped = cv2.warpAffine(data, M, (output_size, output_size), borderValue=0.0) + return cropped, M + + +def trans_points2d(pts, M): + new_pts = np.zeros(shape=pts.shape, dtype=np.float32) + for i in range(pts.shape[0]): + pt = pts[i] + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32) + new_pt = np.dot(M, new_pt) + new_pts[i] = new_pt[0:2] + return new_pts + + +def trans_points3d(pts, M): + scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1]) + #print(scale) + new_pts = np.zeros(shape=pts.shape, dtype=np.float32) + for i in range(pts.shape[0]): + pt = pts[i] + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32) + new_pt = np.dot(M, new_pt) + #print('new_pt', new_pt.shape, new_pt) + new_pts[i][0:2] = new_pt[0:2] + new_pts[i][2] = pts[i][2] * scale + return new_pts + + +def trans_points(pts, M): + if pts.shape[1] == 2: + return trans_points2d(pts, M) + else: + return trans_points3d(pts, M) diff --git a/modules/loader.py b/modules/loader.py index c48afa7a9..0ee139c21 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -52,7 +52,9 @@ import accelerate # pylint: disable=W0611,C0411 timer.startup.record("accelerate") import onnxruntime # pylint: disable=W0611,C0411 -onnxruntime.set_default_logger_severity(3) +onnxruntime.set_default_logger_severity(4) +onnxruntime.set_default_logger_verbosity(1) +onnxruntime.disable_telemetry_events() timer.startup.record("onnx") from fastapi import FastAPI # pylint: disable=W0611,C0411 From fba01890b7a21a5ca947c6ee5e9fe369a665ffce Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 10:05:55 -0500 Subject: [PATCH 14/66] update contributing guidelines Signed-off-by: Vladimir Mandic --- CONTRIBUTING | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/CONTRIBUTING b/CONTRIBUTING index 9dc17f199..a322cb97e 100644 --- a/CONTRIBUTING +++ b/CONTRIBUTING @@ -4,17 +4,23 @@ Pull requests from everyone are welcome Procedure for contributing: +- Select SD.Next `dev` branch: + - Create a fork of the repository on github - In a top right corner of a GitHub, select "Fork" - Its recommended to fork latest version from main branch to avoid any possible conflicting code updates + In a top right corner of a GitHub, select "Fork" + Its recommended to fork latest version from main branch to avoid any possible conflicting code updates - Clone your forked repository to your local system - `git clone https://github.com// + `git clone https://github.com//` - Make your changes -- Test your changes -- Test your changes against code guidelines - - `ruff check` - - `pylint /.py` +- Test your changes +- Lint your changes against code guidelines + - `ruff check` + - `pylint /.py` - Push changes to your fork -- Submit a PR (pull request) +- Submit a PR (pull request) + - Make sure that PR is against `dev` branch + - Update your fork before createing PR so that it is based on latest code + - Make sure that PR does NOT include any unrelated edits + - Make sure that PR does not include changes to submodules Your pull request will be reviewed and pending review results, merged into main branch From 73aaaca85888d139f9906eda6b519e4acd0b4013 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 12:41:32 -0500 Subject: [PATCH 15/66] add photomaker v1 and v2 Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 8 +- modules/face/__init__.py | 9 +- modules/face/photomaker.py | 40 +- modules/face/photomaker_model.py | 555 ----------------- modules/face/photomaker_model_v1.py | 107 ++++ modules/face/photomaker_model_v2.py | 337 +++++++++++ modules/face/photomaker_pipeline.py | 885 ++++++++++++++++++++++++++++ modules/sd_checkpoint.py | 1 - 8 files changed, 1373 insertions(+), 569 deletions(-) delete mode 100644 modules/face/photomaker_model.py create mode 100644 modules/face/photomaker_model_v1.py create mode 100644 modules/face/photomaker_model_v2.py create mode 100644 modules/face/photomaker_pipeline.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c7e63102..ac2f26a50 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-19 +## Update for 2025-01-21 - **Model Merge** - replace model components and merge LoRAs @@ -13,6 +13,12 @@ - **Detailer**: - in addition as standard behavior of detect & run-generate, it can now also run face-restore models - included models are: *CodeFormer, RestoreFormer, GFPGan, GPEN-BFR* +- **Face**: + - new [photomaker v2](https://huggingface.co/TencentARC/PhotoMaker-V2) and reimplemented [photomaker v1](https://huggingface.co/TencentARC/PhotoMaker) + compatible with sdxl models, generates pretty good results and its faster than most other methods + select under *scripts -> face -> photomaker* + - new [reswapper](https://github.com/somanchiu/ReSwapper) + todo: experimental-only and unfinished, only noting in changelog for future reference - **Other**: - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method - **ipex**: update supported torch versions diff --git a/modules/face/__init__.py b/modules/face/__init__.py index c851b4ee4..6835bda17 100644 --- a/modules/face/__init__.py +++ b/modules/face/__init__.py @@ -88,6 +88,7 @@ class Script(scripts.Script): with gr.Row(): gr.HTML('  Tenecent ARC Lab PhotoMaker
') with gr.Row(): + pm_model = gr.Dropdown(label='PhotoMaker Model', choices=['PhotoMaker v1', 'PhotoMaker v2'], value='PhotoMaker v2') pm_trigger = gr.Text(label='Trigger word', value="person") pm_strength = gr.Slider(label='Strength', minimum=0.0, maximum=2.0, step=0.01, value=1.0) pm_start = gr.Slider(label='Start', minimum=0.0, maximum=1.0, step=0.01, value=0.5) @@ -98,9 +99,9 @@ class Script(scripts.Script): files.change(fn=self.load_images, inputs=[files], outputs=[gallery]) mode.change(fn=self.mode_change, inputs=[mode], outputs=[cfg_reswapper, cfg_faceid, cfg_faceswap, cfg_instantid, cfg_photomaker]) - return [mode, gallery, reswapper_model, reswapper_original, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_trigger, pm_strength, pm_start, fs_cache] + return [mode, gallery, reswapper_model, reswapper_original, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_model, pm_trigger, pm_strength, pm_start, fs_cache] - def run(self, p: processing.StableDiffusionProcessing, mode, input_images, reswapper_model, reswapper_original, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_trigger, pm_strength, pm_start, fs_cache): # pylint: disable=arguments-differ, unused-argument + def run(self, p: processing.StableDiffusionProcessing, mode, input_images, reswapper_model, reswapper_original, ip_model, ip_override, ip_cache, ip_strength, ip_structure, id_strength, id_conditioning, id_cache, pm_model, pm_trigger, pm_strength, pm_start, fs_cache): # pylint: disable=arguments-differ, unused-argument if not shared.native: return None if mode == 'None': @@ -130,8 +131,10 @@ class Script(scripts.Script): processed_images = face_id(p, app=app, source_images=input_images, model=ip_model, override=ip_override, cache=ip_cache, scale=ip_strength, structure=ip_structure) # run faceid pipeline processed = processing.Processed(p, images_list=processed_images, seed=p.seed, subseed=p.subseed, index_of_first_image=0) # manually created processed object elif mode == 'PhotoMaker': # photomaker creates pipeline and triggers original process_images + from modules.face.insightface import get_app + app = get_app('buffalo_l') from modules.face.photomaker import photo_maker - processed = photo_maker(p, input_images=input_images, trigger=pm_trigger, strength=pm_strength, start=pm_start) + processed = photo_maker(p, app=app, input_images=input_images, model=pm_model, trigger=pm_trigger, strength=pm_strength, start=pm_start) elif mode == 'InstantID': from modules.face.insightface import get_app app=get_app('antelopev2') diff --git a/modules/face/photomaker.py b/modules/face/photomaker.py index 2b950d572..4ad31660d 100644 --- a/modules/face/photomaker.py +++ b/modules/face/photomaker.py @@ -1,10 +1,12 @@ -import os +import cv2 +import numpy as np +import torch import huggingface_hub as hf from modules import shared, processing, sd_models, devices -def photo_maker(p: processing.StableDiffusionProcessing, input_images, trigger, strength, start): # pylint: disable=arguments-differ - from modules.face.photomaker_model import PhotoMakerStableDiffusionXLPipeline +def photo_maker(p: processing.StableDiffusionProcessing, app, model: str, input_images, trigger, strength, start): # pylint: disable=arguments-differ + from modules.face.photomaker_pipeline import PhotoMakerStableDiffusionXLPipeline # prepare pipeline if len(input_images) == 0: @@ -54,22 +56,42 @@ def photo_maker(p: processing.StableDiffusionProcessing, input_images, trigger, p.task_args['start_merge_step'] = int(start * p.steps) p.task_args['prompt'] = p.all_prompts[0] if p.all_prompts is not None else p.prompt - photomaker_path = hf.hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model", cache_dir=shared.opts.hfcache_dir) - shared.log.debug(f'PhotoMaker: model={photomaker_path} images={len(input_images)} trigger={trigger} args={p.task_args}') + is_v2 = 'v2' in model + if is_v2: + repo_id, fn = 'TencentARC/PhotoMaker-V2', 'photomaker-v2.bin' + else: + repo_id, fn = 'TencentARC/PhotoMaker', 'photomaker-v1.bin' + + photomaker_path = hf.hf_hub_download(repo_id=repo_id, filename=fn, repo_type="model", cache_dir=shared.opts.hfcache_dir) + shared.log.debug(f'PhotoMaker: model="{model}" uri="{repo_id}/{fn}" images={len(input_images)} trigger={trigger} args={p.task_args}') # load photomaker adapter shared.sd_model.load_photomaker_adapter( - os.path.dirname(photomaker_path), - subfolder="", - weight_name=os.path.basename(photomaker_path), - trigger_word=trigger + photomaker_path, + trigger_word=trigger, + weight_name='photomaker-v2.bin' if is_v2 else 'photomaker-v1.bin', + pm_version='v2' if is_v2 else 'v1', + cache_dir=shared.opts.hfcache_dir, ) shared.sd_model.set_adapters(["photomaker"], adapter_weights=[strength]) + # analyze faces + if is_v2: + id_embed_list = [] + for i, source_image in enumerate(input_images): + faces = app.get(cv2.cvtColor(np.array(source_image), cv2.COLOR_RGB2BGR)) + face = sorted(faces, key=lambda x:(x['bbox'][2]-x['bbox'][0])*x['bbox'][3]-x['bbox'][1])[-1] # only use the maximum face + id_embed_list.append(torch.from_numpy(face['embedding'])) + shared.log.debug(f'PhotoMaker: face={i+1} score={face.det_score:.2f} gender={"female" if face.gender==0 else "male"} age={face.age} bbox={face.bbox}') + p.task_args['id_embeds'] = torch.stack(id_embed_list) + # run processing processed: processing.Processed = processing.process_images(p) p.extra_generation_params['PhotoMaker'] = f'{strength}' + # unload photomaker adapter + shared.sd_model.unload_lora_weights() + # restore original pipeline shared.opts.data['prompt_attention'] = orig_prompt_attention shared.sd_model = orig_pipeline diff --git a/modules/face/photomaker_model.py b/modules/face/photomaker_model.py deleted file mode 100644 index 3595c6a36..000000000 --- a/modules/face/photomaker_model.py +++ /dev/null @@ -1,555 +0,0 @@ -from typing import Any, Callable, Dict, List, Optional, Union, Tuple -import PIL -import torch -import torch.nn as nn -from safetensors import safe_open -from huggingface_hub.utils import validate_hf_hub_args -from diffusers import StableDiffusionXLPipeline -from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput -from diffusers.utils import _get_model_file -from transformers import CLIPImageProcessor -from transformers.models.clip.modeling_clip import CLIPVisionModelWithProjection -from transformers.models.clip.configuration_clip import CLIPVisionConfig - - -PipelineImageInput = Union[ - PIL.Image.Image, - torch.FloatTensor, - List[PIL.Image.Image], - List[torch.FloatTensor], -] - - -VISION_CONFIG_DICT = { - "hidden_size": 1024, - "intermediate_size": 4096, - "num_attention_heads": 16, - "num_hidden_layers": 24, - "patch_size": 14, - "projection_dim": 768 -} - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg -def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): - """ - Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 - """ - std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) - std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) - # rescale the results from guidance (fixes overexposure) - noise_pred_rescaled = noise_cfg * (std_text / std_cfg) - # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images - noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg - return noise_cfg - - -class MLP(nn.Module): - def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True): - super().__init__() - if use_residual: - assert in_dim == out_dim - self.layernorm = nn.LayerNorm(in_dim) - self.fc1 = nn.Linear(in_dim, hidden_dim) - self.fc2 = nn.Linear(hidden_dim, out_dim) - self.use_residual = use_residual - self.act_fn = nn.GELU() - - def forward(self, x): - residual = x - x = self.layernorm(x) - x = self.fc1(x) - x = self.act_fn(x) - x = self.fc2(x) - if self.use_residual: - x = x + residual - return x - - -class FuseModule(nn.Module): - def __init__(self, embed_dim): - super().__init__() - self.mlp1 = MLP(embed_dim * 2, embed_dim, embed_dim, use_residual=False) - self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True) - self.layer_norm = nn.LayerNorm(embed_dim) - - def fuse_fn(self, prompt_embeds, id_embeds): - stacked_id_embeds = torch.cat([prompt_embeds, id_embeds], dim=-1) - stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds - stacked_id_embeds = self.mlp2(stacked_id_embeds) - stacked_id_embeds = self.layer_norm(stacked_id_embeds) - return stacked_id_embeds - - def forward( - self, - prompt_embeds, - id_embeds, - class_tokens_mask, - ) -> torch.Tensor: - # id_embeds shape: [b, max_num_inputs, 1, 2048] - id_embeds = id_embeds.to(prompt_embeds.dtype) - num_inputs = class_tokens_mask.sum().unsqueeze(0) - batch_size, max_num_inputs = id_embeds.shape[:2] - # seq_length: 77 - seq_length = prompt_embeds.shape[1] - # flat_id_embeds shape: [b*max_num_inputs, 1, 2048] - flat_id_embeds = id_embeds.view( - -1, id_embeds.shape[-2], id_embeds.shape[-1] - ) - # valid_id_mask [b*max_num_inputs] - valid_id_mask = ( - torch.arange(max_num_inputs, device=flat_id_embeds.device)[None, :] - < num_inputs[:, None] - ) - valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()] - - prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1]) - class_tokens_mask = class_tokens_mask.view(-1) - valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1]) - # slice out the image token embeddings - image_token_embeds = prompt_embeds[class_tokens_mask] - stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds) - assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}" - prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype)) - updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1) - return updated_prompt_embeds - -class PhotoMakerIDEncoder(CLIPVisionModelWithProjection): - def __init__(self): - super().__init__(CLIPVisionConfig(**VISION_CONFIG_DICT)) - self.visual_projection_2 = nn.Linear(1024, 1280, bias=False) - self.fuse_module = FuseModule(2048) - - def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask): - b, num_inputs, c, h, w = id_pixel_values.shape - id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w) - - shared_id_embeds = self.vision_model(id_pixel_values)[1] - id_embeds = self.visual_projection(shared_id_embeds) - id_embeds_2 = self.visual_projection_2(shared_id_embeds) - - id_embeds = id_embeds.view(b, num_inputs, 1, -1) - id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1) - - id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1) - updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask) - - return updated_prompt_embeds - - -class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline): - @validate_hf_hub_args - def load_photomaker_adapter( - self, - pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], - weight_name: str, - subfolder: str = '', - trigger_word: str = 'img', - **kwargs, - ): - # Load the main state dict first. - cache_dir = kwargs.pop("cache_dir", None) - force_download = kwargs.pop("force_download", False) - resume_download = kwargs.pop("resume_download", False) - proxies = kwargs.pop("proxies", None) - local_files_only = kwargs.pop("local_files_only", None) - token = kwargs.pop("token", None) - revision = kwargs.pop("revision", None) - - user_agent = { - "file_type": "attn_procs_weights", - "framework": "pytorch", - } - - if not isinstance(pretrained_model_name_or_path_or_dict, dict): - model_file = _get_model_file( - pretrained_model_name_or_path_or_dict, - weights_name=weight_name, - cache_dir=cache_dir, - force_download=force_download, - resume_download=resume_download, - proxies=proxies, - local_files_only=local_files_only, - token=token, - revision=revision, - subfolder=subfolder, - user_agent=user_agent, - ) - if weight_name.endswith(".safetensors"): - state_dict = {"id_encoder": {}, "lora_weights": {}} - with safe_open(model_file, framework="pt", device="cpu") as f: - for key in f.keys(): - if key.startswith("id_encoder."): - state_dict["id_encoder"][key.replace("id_encoder.", "")] = f.get_tensor(key) - elif key.startswith("lora_weights."): - state_dict["lora_weights"][key.replace("lora_weights.", "")] = f.get_tensor(key) - else: - state_dict = torch.load(model_file, map_location="cpu") - else: - state_dict = pretrained_model_name_or_path_or_dict - - keys = list(state_dict.keys()) - if keys != ["id_encoder", "lora_weights"]: - raise ValueError("Required keys are (`id_encoder` and `lora_weights`) missing from the state dict.") - - self.trigger_word = trigger_word - # load finetuned CLIP image encoder and fuse module here if it has not been registered to the pipeline yet - id_encoder = PhotoMakerIDEncoder() - id_encoder.load_state_dict(state_dict["id_encoder"], strict=True) - id_encoder = id_encoder.to(self.device, dtype=self.unet.dtype) - self.id_encoder = id_encoder - self.id_image_processor = CLIPImageProcessor() - - # load lora into models - self.load_lora_weights(state_dict["lora_weights"], adapter_name="photomaker") - - # Add trigger word token - if self.tokenizer is not None: - self.tokenizer.add_tokens([self.trigger_word], special_tokens=True) - self.tokenizer_2.add_tokens([self.trigger_word], special_tokens=True) - - def encode_prompt_with_trigger_word( - self, - prompt: str, - prompt_2: Optional[str] = None, - num_id_images: int = 1, - device: Optional[torch.device] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - class_tokens_mask: Optional[torch.LongTensor] = None, - ): - device = device or self._execution_device - - """ - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - """ - - # Find the token id of the trigger word - image_token_id = self.tokenizer_2.convert_tokens_to_ids(self.trigger_word) - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2] - text_encoders = ( - [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2] - ) - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_embeds_list = [] - prompts = [prompt, prompt_2] - for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): - input_ids = tokenizer.encode(prompt) - clean_index = 0 - clean_input_ids = [] - class_token_index = [] - # Find out the corrresponding class word token based on the newly added trigger word token - for _i, token_id in enumerate(input_ids): - if token_id == image_token_id: - class_token_index.append(clean_index - 1) - else: - clean_input_ids.append(token_id) - clean_index += 1 - - if len(class_token_index) != 1: - raise ValueError( - f"PhotoMaker currently does not support multiple trigger words in a single prompt.\ - Trigger word: {self.trigger_word}, Prompt: {prompt}." - ) - class_token_index = class_token_index[0] - - # Expand the class word token and corresponding mask - class_token = clean_input_ids[class_token_index] - clean_input_ids = clean_input_ids[:class_token_index] + [class_token] * num_id_images + \ - clean_input_ids[class_token_index+1:] - - # Truncation or padding - max_len = tokenizer.model_max_length - if len(clean_input_ids) > max_len: - clean_input_ids = clean_input_ids[:max_len] - else: - clean_input_ids = clean_input_ids + [tokenizer.pad_token_id] * ( - max_len - len(clean_input_ids) - ) - - class_tokens_mask = [True if class_token_index <= i < class_token_index+num_id_images else False \ - for i in range(len(clean_input_ids))] - - clean_input_ids = torch.tensor(clean_input_ids, dtype=torch.long).unsqueeze(0) - class_tokens_mask = torch.tensor(class_tokens_mask, dtype=torch.bool).unsqueeze(0) - - prompt_embeds = text_encoder( - clean_input_ids.to(device), - output_hidden_states=True, - ) - - # We are only ALWAYS interested in the pooled output of the final text encoder - pooled_prompt_embeds = prompt_embeds[0] - prompt_embeds = prompt_embeds.hidden_states[-2] - prompt_embeds_list.append(prompt_embeds) - - prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) - - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - class_tokens_mask = class_tokens_mask.to(device=device) - - return prompt_embeds, pooled_prompt_embeds, class_tokens_mask - - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - denoising_end: Optional[float] = None, - guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, - callback_steps: int = 1, - # Added parameters (for PhotoMaker) - input_id_images: PipelineImageInput = None, - start_merge_step: int = 0, - class_tokens_mask: Optional[torch.LongTensor] = None, - prompt_embeds_text_only: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds_text_only: Optional[torch.FloatTensor] = None, - ): - # 0. Default height and width to unet - height = height or self.unet.config.sample_size * self.vae_scale_factor - width = width or self.unet.config.sample_size * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - # - if prompt_embeds is not None and class_tokens_mask is None: - raise ValueError( - "If `prompt_embeds` are provided, `class_tokens_mask` also have to be passed. Make sure to generate `class_tokens_mask` from the same tokenizer that was used to generate `prompt_embeds`." - ) - # check the input id images - if input_id_images is None: - raise ValueError( - "Provide `input_id_images`. Cannot leave `input_id_images` undefined for PhotoMaker pipeline." - ) - if not isinstance(input_id_images, list): - input_id_images = [input_id_images] - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self._execution_device - - # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) - # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` - # corresponds to doing no classifier free guidance. - do_classifier_free_guidance = guidance_scale > 1.0 - - assert do_classifier_free_guidance - - # 3. Encode input prompt - num_id_images = len(input_id_images) - - ( - prompt_embeds, - pooled_prompt_embeds, - class_tokens_mask, - ) = self.encode_prompt_with_trigger_word( - prompt=prompt, - prompt_2=prompt_2, - device=device, - num_id_images=num_id_images, - prompt_embeds=prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - class_tokens_mask=class_tokens_mask, - ) - - # 4. Encode input prompt without the trigger word for delayed conditioning - prompt_text_only = prompt.replace(" "+self.trigger_word, "") # sensitive to white space - ( - prompt_embeds_text_only, - negative_prompt_embeds, - pooled_prompt_embeds_text_only, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt_text_only, - prompt_2=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=do_classifier_free_guidance, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - prompt_embeds=prompt_embeds_text_only, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds_text_only, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - ) - - # 5. Prepare the input ID images - dtype = next(self.id_encoder.parameters()).dtype - if not isinstance(input_id_images[0], torch.Tensor): - id_pixel_values = self.id_image_processor(input_id_images, return_tensors="pt").pixel_values - - id_pixel_values = id_pixel_values.unsqueeze(0).to(device=device, dtype=dtype) - - # 6. Get the update text embedding with the stacked ID embedding - prompt_embeds = self.id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask) - - bs_embed, seq_len, _ = prompt_embeds.shape - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( - bs_embed * num_images_per_prompt, -1 - ) - - # 7. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 8. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 9. Prepare extra step kwargs. - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 10. Prepare added time ids & embeddings - if self.text_encoder_2 is None: - text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) - else: - text_encoder_projection_dim = self.text_encoder_2.config.projection_dim - - add_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0) - add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) - - # 11. Denoising loop - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - latent_model_input = ( - torch.cat([latents] * 2) if do_classifier_free_guidance else latents - ) - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - - if i <= start_merge_step: - current_prompt_embeds = torch.cat( - [negative_prompt_embeds, prompt_embeds_text_only], dim=0 - ) - add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds_text_only], dim=0) - else: - current_prompt_embeds = torch.cat( - [negative_prompt_embeds, prompt_embeds], dim=0 - ) - add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - # predict the noise residual - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - noise_pred = self.unet( - latent_model_input, - t, - encoder_hidden_states=current_prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False, - )[0] - - # perform guidance - if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - - if do_classifier_free_guidance and guidance_rescale > 0.0: - # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - - # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - callback(i, t, latents) - - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.dtype == torch.float16 and self.vae.config.force_upcast: - self.upcast_vae() - latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) - - if output_type != "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - return StableDiffusionXLPipelineOutput(images=image) - - # apply watermark if available - # if self.watermark is not None: - # image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, output_type=output_type) - - # Offload last model to CPU - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.final_offload_hook.offload() - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/modules/face/photomaker_model_v1.py b/modules/face/photomaker_model_v1.py new file mode 100644 index 000000000..ea7f87ca2 --- /dev/null +++ b/modules/face/photomaker_model_v1.py @@ -0,0 +1,107 @@ +### original + +import torch +import torch.nn as nn +from transformers.models.clip.modeling_clip import CLIPVisionModelWithProjection +from transformers.models.clip.configuration_clip import CLIPVisionConfig + +VISION_CONFIG_DICT = { + "hidden_size": 1024, + "intermediate_size": 4096, + "num_attention_heads": 16, + "num_hidden_layers": 24, + "patch_size": 14, + "projection_dim": 768 +} + +class MLP(nn.Module): + def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True): + super().__init__() + if use_residual: + assert in_dim == out_dim + self.layernorm = nn.LayerNorm(in_dim) + self.fc1 = nn.Linear(in_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, out_dim) + self.use_residual = use_residual + self.act_fn = nn.GELU() + + def forward(self, x): + residual = x + x = self.layernorm(x) + x = self.fc1(x) + x = self.act_fn(x) + x = self.fc2(x) + if self.use_residual: + x = x + residual + return x + + +class FuseModule(nn.Module): + def __init__(self, embed_dim): + super().__init__() + self.mlp1 = MLP(embed_dim * 2, embed_dim, embed_dim, use_residual=False) + self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True) + self.layer_norm = nn.LayerNorm(embed_dim) + + def fuse_fn(self, prompt_embeds, id_embeds): + unstacked_prompt_embeds = prompt_embeds.unbind(0) + stacked_id_embeds = torch.cat([unstacked_prompt_embeds[0].unsqueeze(0), id_embeds], dim=-1) # monkey patch + stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds + stacked_id_embeds = self.mlp2(stacked_id_embeds) + stacked_id_embeds = self.layer_norm(stacked_id_embeds) + return stacked_id_embeds + + def forward( + self, + prompt_embeds, + id_embeds, + class_tokens_mask, + ) -> torch.Tensor: + # id_embeds shape: [b, max_num_inputs, 1, 2048] + id_embeds = id_embeds.to(prompt_embeds.dtype) + num_inputs = class_tokens_mask.sum().unsqueeze(0) + batch_size, max_num_inputs = id_embeds.shape[:2] + # seq_length: 77 + seq_length = prompt_embeds.shape[1] + # flat_id_embeds shape: [b*max_num_inputs, 1, 2048] + flat_id_embeds = id_embeds.view( + -1, id_embeds.shape[-2], id_embeds.shape[-1] + ) + # valid_id_mask [b*max_num_inputs] + valid_id_mask = ( + torch.arange(max_num_inputs, device=flat_id_embeds.device)[None, :] + < num_inputs[:, None] + ) + valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()] + prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1]) + class_tokens_mask = class_tokens_mask.view(-1) + valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1]) + # slice out the image token embeddings + image_token_embeds = prompt_embeds[class_tokens_mask] + stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds) + assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}" + prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype)) + updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1) + return updated_prompt_embeds + +class PhotoMakerIDEncoder(CLIPVisionModelWithProjection): + def __init__(self): + super().__init__(CLIPVisionConfig(**VISION_CONFIG_DICT)) + self.visual_projection_2 = nn.Linear(1024, 1280, bias=False) + self.fuse_module = FuseModule(2048) + + def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask): # pylint: disable=arguments-differ + b, num_inputs, c, h, w = id_pixel_values.shape + id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w) + + shared_id_embeds = self.vision_model(id_pixel_values)[1] + id_embeds = self.visual_projection(shared_id_embeds) + id_embeds_2 = self.visual_projection_2(shared_id_embeds) + + id_embeds = id_embeds.view(b, num_inputs, 1, -1) + id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1) + + id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1) + updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask) + + return updated_prompt_embeds diff --git a/modules/face/photomaker_model_v2.py b/modules/face/photomaker_model_v2.py new file mode 100644 index 000000000..34704376f --- /dev/null +++ b/modules/face/photomaker_model_v2.py @@ -0,0 +1,337 @@ +### original + +import math +import torch +import torch.nn as nn +from transformers.models.clip.modeling_clip import CLIPVisionModelWithProjection +from transformers.models.clip.configuration_clip import CLIPVisionConfig +from einops import rearrange +from einops.layers.torch import Rearrange + + +class FacePerceiverResampler(torch.nn.Module): + def __init__( + self, + *, + dim=768, + depth=4, + dim_head=64, + heads=16, + embedding_dim=1280, + output_dim=768, + ff_mult=4, + ): + super().__init__() + self.proj_in = torch.nn.Linear(embedding_dim, dim) + self.proj_out = torch.nn.Linear(dim, output_dim) + self.norm_out = torch.nn.LayerNorm(output_dim) + self.layers = torch.nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + torch.nn.ModuleList( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, latents, x): + x = self.proj_in(x) + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + latents = self.proj_out(latents) + return self.norm_out(latents) + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, _width = x.shape + # (bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + + +class Resampler(nn.Module): + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output_dim=1024, + ff_mult=4, + max_seq_len: int = 257, # CLIP tokens + CLS token + apply_pos_emb: bool = False, + num_latents_mean_pooled: int = 0, # number of latents derived from mean pooled representation of the sequence + ): + super().__init__() + self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.proj_in = nn.Linear(embedding_dim, dim) + + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.to_latents_from_mean_pooled_seq = ( + nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, dim * num_latents_mean_pooled), + Rearrange("b (n d) -> b n d", n=num_latents_mean_pooled), + ) + if num_latents_mean_pooled > 0 + else None + ) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + if self.pos_emb is not None: + n, device = x.shape[1], x.device + pos_emb = self.pos_emb(torch.arange(n, device=device)) + x = x + pos_emb + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + + if self.to_latents_from_mean_pooled_seq: + meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool)) + meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq) + latents = torch.cat((meanpooled_latents, latents), dim=-2) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + latents = self.proj_out(latents) + return self.norm_out(latents) + + +def masked_mean(t, *, dim, mask=None): + if mask is None: + return t.mean(dim=dim) + + denom = mask.sum(dim=dim, keepdim=True) + mask = rearrange(mask, "b n -> b n 1") + masked_t = t.masked_fill(~mask, 0.0) + + return masked_t.sum(dim=dim) / denom.clamp(min=1e-5) + + +VISION_CONFIG_DICT = { + "hidden_size": 1024, + "intermediate_size": 4096, + "num_attention_heads": 16, + "num_hidden_layers": 24, + "patch_size": 14, + "projection_dim": 768 +} + + +class MLP(nn.Module): + def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True): + super().__init__() + if use_residual: + assert in_dim == out_dim + self.layernorm = nn.LayerNorm(in_dim) + self.fc1 = nn.Linear(in_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, out_dim) + self.use_residual = use_residual + self.act_fn = nn.GELU() + + def forward(self, x): + residual = x + x = self.layernorm(x) + x = self.fc1(x) + x = self.act_fn(x) + x = self.fc2(x) + if self.use_residual: + x = x + residual + return x + + +class QFormerPerceiver(nn.Module): + def __init__(self, id_embeddings_dim, cross_attention_dim, num_tokens, embedding_dim=1024, use_residual=True, ratio=4): + super().__init__() + + self.num_tokens = num_tokens + self.cross_attention_dim = cross_attention_dim + self.use_residual = use_residual + self.token_proj = nn.Sequential( + nn.Linear(id_embeddings_dim, id_embeddings_dim*ratio), + nn.GELU(), + nn.Linear(id_embeddings_dim*ratio, cross_attention_dim*num_tokens), + ) + self.token_norm = nn.LayerNorm(cross_attention_dim) + self.perceiver_resampler = FacePerceiverResampler( + dim=cross_attention_dim, + depth=4, + dim_head=128, + heads=cross_attention_dim // 128, + embedding_dim=embedding_dim, + output_dim=cross_attention_dim, + ff_mult=4, + ) + + def forward(self, x, last_hidden_state): + x = self.token_proj(x) + x = x.reshape(-1, self.num_tokens, self.cross_attention_dim) + x = self.token_norm(x) # cls token + out = self.perceiver_resampler(x, last_hidden_state) # retrieve from patch tokens + if self.use_residual: + out = x + 1.0 * out + return out + + +class FuseModule(nn.Module): + def __init__(self, embed_dim): + super().__init__() + self.mlp1 = MLP(embed_dim * 2, embed_dim, embed_dim, use_residual=False) + self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True) + self.layer_norm = nn.LayerNorm(embed_dim) + + def fuse_fn(self, prompt_embeds, id_embeds): + stacked_id_embeds = torch.cat([prompt_embeds, id_embeds], dim=-1) + stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds + stacked_id_embeds = self.mlp2(stacked_id_embeds) + stacked_id_embeds = self.layer_norm(stacked_id_embeds) + return stacked_id_embeds + + def forward( + self, + prompt_embeds, + id_embeds, + class_tokens_mask, + ) -> torch.Tensor: + # id_embeds shape: [b, max_num_inputs, 1, 2048] + id_embeds = id_embeds.to(prompt_embeds.dtype) + num_inputs = class_tokens_mask.sum().unsqueeze(0) + batch_size, max_num_inputs = id_embeds.shape[:2] + # seq_length: 77 + seq_length = prompt_embeds.shape[1] + # flat_id_embeds shape: [b*max_num_inputs, 1, 2048] + flat_id_embeds = id_embeds.view( + -1, id_embeds.shape[-2], id_embeds.shape[-1] + ) + # valid_id_mask [b*max_num_inputs] + valid_id_mask = ( + torch.arange(max_num_inputs, device=flat_id_embeds.device)[None, :] + < num_inputs[:, None] + ) + valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()] + + prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1]) + class_tokens_mask = class_tokens_mask.view(-1) + valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1]) + # slice out the image token embeddings + image_token_embeds = prompt_embeds[class_tokens_mask] + stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds) + assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}" + prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype)) + updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1) + return updated_prompt_embeds + + +class PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken(CLIPVisionModelWithProjection): + def __init__(self, id_embeddings_dim=512): + super().__init__(CLIPVisionConfig(**VISION_CONFIG_DICT)) + self.fuse_module = FuseModule(2048) + self.visual_projection_2 = nn.Linear(1024, 1280, bias=False) + + cross_attention_dim = 2048 + # projection + self.num_tokens = 2 + self.cross_attention_dim = cross_attention_dim + self.qformer_perceiver = QFormerPerceiver( + id_embeddings_dim, + cross_attention_dim, + self.num_tokens, + ) + + def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds): # pylint: disable=arguments-differ + b, num_inputs, c, h, w = id_pixel_values.shape + id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w) + + last_hidden_state = self.vision_model(id_pixel_values)[0] + id_embeds = id_embeds.view(b * num_inputs, -1) + + id_embeds = self.qformer_perceiver(id_embeds, last_hidden_state) + id_embeds = id_embeds.view(b, num_inputs, self.num_tokens, -1) + updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask) + + return updated_prompt_embeds diff --git a/modules/face/photomaker_pipeline.py b/modules/face/photomaker_pipeline.py new file mode 100644 index 000000000..ec5728961 --- /dev/null +++ b/modules/face/photomaker_pipeline.py @@ -0,0 +1,885 @@ +### original + +import inspect +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +import PIL +import torch +from transformers import CLIPImageProcessor +from safetensors import safe_open +from huggingface_hub.utils import validate_hf_hub_args +from diffusers import StableDiffusionXLPipeline +from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput +from diffusers.loaders import StableDiffusionXLLoraLoaderMixin, TextualInversionLoaderMixin +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.utils import _get_model_file, USE_PEFT_BACKEND, deprecate, is_torch_xla_available, scale_lora_layers, unscale_lora_layers + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +from modules.face.photomaker_model_v1 import PhotoMakerIDEncoder +from modules.face.photomaker_model_v2 import PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken + +PipelineImageInput = Union[ + PIL.Image.Image, + torch.FloatTensor, + List[PIL.Image.Image], + List[torch.FloatTensor], +] + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline): + @validate_hf_hub_args + def load_photomaker_adapter( + self, + pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + weight_name: str, + subfolder: str = '', + trigger_word: str = 'img', + pm_version: str = 'v2', + **kwargs, + ): + """ + Parameters: + pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): + Can be either: + + - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on + the Hub. + - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved + with [`ModelMixin.save_pretrained`]. + - A [torch state + dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). + + weight_name (`str`): + The weight name NOT the path to the weight. + + subfolder (`str`, defaults to `""`): + The subfolder location of a model file within a larger model repository on the Hub or locally. + + trigger_word (`str`, *optional*, defaults to `"img"`): + The trigger word is used to identify the position of class word in the text prompt, + and it is recommended not to set it as a common word. + This trigger word must be placed after the class word when used, otherwise, it will affect the performance of the personalized generation. + """ + + # Load the main state dict first. + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", None) + token = kwargs.pop("token", None) + revision = kwargs.pop("revision", None) + + user_agent = { + "file_type": "attn_procs_weights", + "framework": "pytorch", + } + + if not isinstance(pretrained_model_name_or_path_or_dict, dict): + model_file = _get_model_file( + pretrained_model_name_or_path_or_dict, + weights_name=weight_name, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + ) + if weight_name.endswith(".safetensors"): + state_dict = {"id_encoder": {}, "lora_weights": {}} + with safe_open(model_file, framework="pt", device="cpu") as f: + for key in f.keys(): + if key.startswith("id_encoder."): + state_dict["id_encoder"][key.replace("id_encoder.", "")] = f.get_tensor(key) + elif key.startswith("lora_weights."): + state_dict["lora_weights"][key.replace("lora_weights.", "")] = f.get_tensor(key) + else: + state_dict = torch.load(model_file, map_location="cpu") + else: + state_dict = pretrained_model_name_or_path_or_dict + + keys = list(state_dict.keys()) + if keys != ["id_encoder", "lora_weights"]: + raise ValueError("Required keys are (`id_encoder` and `lora_weights`) missing from the state dict.") + + self.num_tokens =2 # pylint: disable=attribute-defined-outside-init + self.pm_version = pm_version # pylint: disable=attribute-defined-outside-init + self.trigger_word = trigger_word # pylint: disable=attribute-defined-outside-init + # load finetuned CLIP image encoder and fuse module here if it has not been registered to the pipeline yet + self.id_image_processor = CLIPImageProcessor() # pylint: disable=attribute-defined-outside-init + if pm_version == "v1": # PhotoMaker v1 + id_encoder = PhotoMakerIDEncoder() + elif pm_version == "v2": # PhotoMaker v2 + id_encoder = PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken() + else: + raise NotImplementedError(f"The PhotoMaker version [{pm_version}] does not support") + + id_encoder.load_state_dict(state_dict["id_encoder"], strict=True) + id_encoder = id_encoder.to(self.device, dtype=self.unet.dtype) + self.id_encoder = id_encoder # pylint: disable=attribute-defined-outside-init + + # load lora into models + self.load_lora_weights(state_dict["lora_weights"], adapter_name="photomaker") + + # Add trigger word token + if self.tokenizer is not None: + self.tokenizer.add_tokens([self.trigger_word], special_tokens=True) + + self.tokenizer_2.add_tokens([self.trigger_word], special_tokens=True) + + + def encode_prompt_with_trigger_word( + self, + prompt: str, + prompt_2: Optional[str] = None, + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[str] = None, + negative_prompt_2: Optional[str] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + pooled_prompt_embeds: Optional[torch.Tensor] = None, + negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + lora_scale: Optional[float] = None, + clip_skip: Optional[int] = None, + ### Added args + num_id_images: int = 1, + class_tokens_mask: Optional[torch.LongTensor] = None, + ): + device = device or self._execution_device + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin): + self._lora_scale = lora_scale # pylint: disable=attribute-defined-outside-init + + # dynamically adjust the LoRA scale + if self.text_encoder is not None: + if not USE_PEFT_BACKEND: + adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) + else: + scale_lora_layers(self.text_encoder, lora_scale) + + if self.text_encoder_2 is not None: + if not USE_PEFT_BACKEND: + adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale) + else: + scale_lora_layers(self.text_encoder_2, lora_scale) + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # Find the token id of the trigger word + image_token_id = self.tokenizer_2.convert_tokens_to_ids(self.trigger_word) + + # Define tokenizers and text encoders + tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2] + text_encoders = ( + [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2] + ) + + if prompt_embeds is None: + prompt_2 = prompt_2 or prompt + prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 + + # textual inversion: process multi-vector tokens if necessary + prompt_embeds_list = [] + prompts = [prompt, prompt_2] + for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, tokenizer) + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( + text_input_ids, untruncated_ids + ): + _removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1]) + + clean_index = 0 + clean_input_ids = [] + class_token_index = [] + # Find out the corresponding class word token based on the newly added trigger word token + for _i, token_id in enumerate(text_input_ids.tolist()[0]): + if token_id == image_token_id: + class_token_index.append(clean_index - 1) + else: + clean_input_ids.append(token_id) + clean_index += 1 + + if len(class_token_index) != 1: + raise ValueError( + f"PhotoMaker currently does not support multiple trigger words in a single prompt.\ + Trigger word: {self.trigger_word}, Prompt: {prompt}." + ) + class_token_index = class_token_index[0] + + # Expand the class word token and corresponding mask + class_token = clean_input_ids[class_token_index] + clean_input_ids = clean_input_ids[:class_token_index] + [class_token] * num_id_images * self.num_tokens + \ + clean_input_ids[class_token_index+1:] + + # Truncation or padding + max_len = tokenizer.model_max_length + if len(clean_input_ids) > max_len: + clean_input_ids = clean_input_ids[:max_len] + else: + clean_input_ids = clean_input_ids + [tokenizer.pad_token_id] * ( + max_len - len(clean_input_ids) + ) + + class_tokens_mask = [True if class_token_index <= i < class_token_index+(num_id_images * self.num_tokens) else False \ + for i in range(len(clean_input_ids))] + + clean_input_ids = torch.tensor(clean_input_ids, dtype=torch.long).unsqueeze(0) + class_tokens_mask = torch.tensor(class_tokens_mask, dtype=torch.bool).unsqueeze(0) + + prompt_embeds = text_encoder(clean_input_ids.to(device), output_hidden_states=True) + + # We are only ALWAYS interested in the pooled output of the final text encoder + pooled_prompt_embeds = prompt_embeds[0] + if clip_skip is None: + prompt_embeds = prompt_embeds.hidden_states[-2] + else: + # "2" because SDXL always indexes from the penultimate layer. + prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] + + prompt_embeds_list.append(prompt_embeds) + + prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) + + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) + class_tokens_mask = class_tokens_mask.to(device=device) + # get unconditional embeddings for classifier free guidance + zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt # pylint: disable=no-member + if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: + negative_prompt_embeds = torch.zeros_like(prompt_embeds) + negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds) + elif do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt_2 = negative_prompt_2 or negative_prompt + + # normalize str to list + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_2 = ( + batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 + ) + + uncond_tokens: List[str] + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = [negative_prompt, negative_prompt_2] + + negative_prompt_embeds_list = [] + for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): + if isinstance(self, TextualInversionLoaderMixin): + negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer) + + max_length = prompt_embeds.shape[1] + uncond_input = tokenizer( + negative_prompt, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + negative_prompt_embeds = text_encoder( + uncond_input.input_ids.to(device), + output_hidden_states=True, + ) + # We are only ALWAYS interested in the pooled output of the final text encoder + negative_pooled_prompt_embeds = negative_prompt_embeds[0] + negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2] + + negative_prompt_embeds_list.append(negative_prompt_embeds) + + negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1) + + if self.text_encoder_2 is not None: + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) + else: + prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device) + + bs_embed, seq_len, _ = prompt_embeds.shape + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + if self.text_encoder_2 is not None: + negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) + else: + negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( + bs_embed * num_images_per_prompt, -1 + ) + if do_classifier_free_guidance: + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( + bs_embed * num_images_per_prompt, -1 + ) + + if self.text_encoder is not None: + if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder, lora_scale) + + if self.text_encoder_2 is not None: + if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: + # Retrieve the original scale by scaling back the LoRA layers + unscale_lora_layers(self.text_encoder_2, lora_scale) + + return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, class_tokens_mask + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + denoising_end: Optional[float] = None, + guidance_scale: float = 5.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + pooled_prompt_embeds: Optional[torch.Tensor] = None, + negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + original_size: Optional[Tuple[int, int]] = None, + crops_coords_top_left: Tuple[int, int] = (0, 0), + target_size: Optional[Tuple[int, int]] = None, + negative_original_size: Optional[Tuple[int, int]] = None, + negative_crops_coords_top_left: Tuple[int, int] = (0, 0), + negative_target_size: Optional[Tuple[int, int]] = None, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + # Added parameters (for PhotoMaker) + input_id_images: PipelineImageInput = None, + start_merge_step: int = 10, + class_tokens_mask: Optional[torch.LongTensor] = None, + id_embeds: Optional[torch.FloatTensor] = None, + prompt_embeds_text_only: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds_text_only: Optional[torch.FloatTensor] = None, + **kwargs, + ): + r""" + Function invoked when calling the pipeline for generation. + Only the parameters introduced by PhotoMaker are discussed here. + For explanations of the previous parameters in StableDiffusionXLPipeline, please refer to https://github.com/huggingface/diffusers/blob/v0.25.0/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py + + Args: + input_id_images (`PipelineImageInput`, *optional*): + Input ID Image to work with PhotoMaker. + class_tokens_mask (`torch.LongTensor`, *optional*): + Pre-generated class token. When the `prompt_embeds` parameter is provided in advance, it is necessary to prepare the `class_tokens_mask` beforehand for marking out the position of class word. + prompt_embeds_text_only (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + pooled_prompt_embeds_text_only (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + + Returns: + [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`: + [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a + `tuple`. When returning a tuple, the first element is a list with the generated images. + """ + + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 0. Default height and width to unet + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + original_size = original_size or (height, width) + target_size = target_size or (height, width) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + height, + width, + callback_steps, + negative_prompt, + negative_prompt_2, + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ip_adapter_image, + ip_adapter_image_embeds, + callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale # pylint: disable=attribute-defined-outside-init + self._guidance_rescale = guidance_rescale # pylint: disable=attribute-defined-outside-init + self._clip_skip = clip_skip # pylint: disable=attribute-defined-outside-init + self._cross_attention_kwargs = cross_attention_kwargs # pylint: disable=attribute-defined-outside-init + self._denoising_end = denoising_end # pylint: disable=attribute-defined-outside-init + self._interrupt = False # pylint: disable=attribute-defined-outside-init + + if prompt_embeds is not None and class_tokens_mask is None: + raise ValueError( + "If `prompt_embeds` are provided, `class_tokens_mask` also have to be passed. Make sure to generate `class_tokens_mask` from the same tokenizer that was used to generate `prompt_embeds`." + ) + # check the input id images + if input_id_images is None: + raise ValueError( + "Provide `input_id_images`. Cannot leave `input_id_images` undefined for PhotoMaker pipeline." + ) + if not isinstance(input_id_images, list): + input_id_images = [input_id_images] + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. Encode input prompt + lora_scale = ( + self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + ) + + num_id_images = len(input_id_images) + ( + prompt_embeds, + _, + pooled_prompt_embeds, + _, + class_tokens_mask, + ) = self.encode_prompt_with_trigger_word( + prompt=prompt, + prompt_2=prompt_2, + device=device, + num_id_images=num_id_images, + class_tokens_mask=class_tokens_mask, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + ) + + # 4. Encode input prompt without the trigger word for delayed conditioning + # encode, remove trigger word token, then decode + tokens_text_only = self.tokenizer.encode(prompt, add_special_tokens=False) + trigger_word_token = self.tokenizer.convert_tokens_to_ids(self.trigger_word) + tokens_text_only.remove(trigger_word_token) + prompt_text_only = self.tokenizer.decode(tokens_text_only, add_special_tokens=False) + ( + prompt_embeds_text_only, + negative_prompt_embeds, + pooled_prompt_embeds_text_only, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt_text_only, + prompt_2=prompt_2, + device=device, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + prompt_embeds=prompt_embeds_text_only, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds_text_only, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + ) + + # 5. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas + ) + + # 6. Prepare the input ID images + dtype = next(self.id_encoder.parameters()).dtype + if not isinstance(input_id_images[0], torch.Tensor): + id_pixel_values = self.id_image_processor(input_id_images, return_tensors="pt").pixel_values # pylint: disable=used-before-assignment + + id_pixel_values = id_pixel_values.unsqueeze(0).to(device=device, dtype=dtype) # pylint: disable=used-before-assignment + + # 7. Get the update text embedding with the stacked ID embedding + if id_embeds is not None: + id_embeds = id_embeds.unsqueeze(0).to(device=device, dtype=dtype) + prompt_embeds = self.id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds) + else: + prompt_embeds = self.id_encoder(id_pixel_values, prompt_embeds, class_tokens_mask) + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) + + # 8. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 10. Prepare added time ids & embeddings + add_text_embeds = pooled_prompt_embeds + if self.text_encoder_2 is None: + text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) + else: + text_encoder_projection_dim = self.text_encoder_2.config.projection_dim + + add_time_ids = self._get_add_time_ids( + original_size, + crops_coords_top_left, + target_size, + dtype=prompt_embeds.dtype, + text_encoder_projection_dim=text_encoder_projection_dim, + ) + if negative_original_size is not None and negative_target_size is not None: + negative_add_time_ids = self._get_add_time_ids( + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype=prompt_embeds.dtype, + text_encoder_projection_dim=text_encoder_projection_dim, + ) + else: + negative_add_time_ids = add_time_ids + + if self.do_classifier_free_guidance: + add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) + + prompt_embeds = prompt_embeds.to(device) + add_text_embeds = add_text_embeds.to(device) + add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) + + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image, + ip_adapter_image_embeds, + device, + batch_size * num_images_per_prompt, + self.do_classifier_free_guidance, + ) + + # 11. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + # 11.1 Apply denoising_end + if ( + self.denoising_end is not None + and isinstance(self.denoising_end, float) + and self.denoising_end > 0 + and self.denoising_end < 1 + ): + discrete_timestep_cutoff = int( + round( + self.scheduler.config.num_train_timesteps + - (self.denoising_end * self.scheduler.config.num_train_timesteps) + ) + ) + num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) + timesteps = timesteps[:num_inference_steps] + + # 12. Optionally get Guidance Scale Embedding + timestep_cond = None + if self.unet.config.time_cond_proj_dim is not None: + guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) + timestep_cond = self.get_guidance_scale_embedding( + guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim + ).to(device=device, dtype=latents.dtype) + + self._num_timesteps = len(timesteps) # pylint: disable=attribute-defined-outside-init + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + if i <= start_merge_step: + current_prompt_embeds = torch.cat( + [negative_prompt_embeds, prompt_embeds_text_only], dim=0 + ) if self.do_classifier_free_guidance else prompt_embeds_text_only + add_text_embeds = torch.cat( + [negative_pooled_prompt_embeds, pooled_prompt_embeds_text_only], dim=0 + ) if self.do_classifier_free_guidance else pooled_prompt_embeds_text_only + else: + current_prompt_embeds = torch.cat( + [negative_prompt_embeds, prompt_embeds], dim=0 + ) if self.do_classifier_free_guidance else prompt_embeds + add_text_embeds = torch.cat( + [negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0 + ) if self.do_classifier_free_guidance else pooled_prompt_embeds + + added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + added_cond_kwargs["image_embeds"] = image_embeds + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=current_prompt_embeds, + timestep_cond=timestep_cond, + cross_attention_kwargs=self.cross_attention_kwargs, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0] + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) + negative_pooled_prompt_embeds = callback_outputs.pop( + "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds + ) + add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) + negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + if XLA_AVAILABLE: + xm.mark_step() # pylint: disable=possibly-used-before-assignment + + if output_type != "latent": + # make sure the VAE is in float32 mode, as it overflows in float16 + needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast + + if needs_upcasting: + self.upcast_vae() + latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) + elif latents.dtype != self.vae.dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + self.vae = self.vae.to(latents.dtype) # pylint: disable=attribute-defined-outside-init + + # unscale/denormalize the latents + # denormalize with the mean and std if available and not None + has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None + has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None + if has_latents_mean and has_latents_std: + latents_mean = ( + torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) + ) + latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean + else: + latents = latents / self.vae.config.scaling_factor + + image = self.vae.decode(latents, return_dict=False)[0] + + # cast back to fp16 if needed + if needs_upcasting: + self.vae.to(dtype=torch.float16) + else: + image = latents + return StableDiffusionXLPipelineOutput(images=image) + + # apply watermark if available + # if self.watermark is not None: + # image = self.watermark.apply_watermark(image) + + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return StableDiffusionXLPipelineOutput(images=image) diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py index c20fa00f4..92c83e953 100644 --- a/modules/sd_checkpoint.py +++ b/modules/sd_checkpoint.py @@ -306,7 +306,6 @@ def extract_thumbnail(filename, data): thumbnail = thumbnail.convert("RGB") thumbnail = thumbnail.resize((512, 512), Image.Resampling.HAMMING) fn = os.path.splitext(filename)[0] - print('HERE', thumbnail, fn) thumbnail = thumbnail.save(f"{fn}.thumb.jpg", quality=50) except Exception as e: shared.log.error(f"Error extracting thumbnail: {filename} {e}") From cd44fb602da81ff4a0cd9823727f906d2f656a0e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 13:13:30 -0500 Subject: [PATCH 16/66] add fasthunyuan Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 8 ++++++-- scripts/hunyuanvideo.py | 41 ++++++++++++++++++++++++++--------------- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ac2f26a50..222b599b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,11 +14,15 @@ - in addition as standard behavior of detect & run-generate, it can now also run face-restore models - included models are: *CodeFormer, RestoreFormer, GFPGan, GPEN-BFR* - **Face**: - - new [photomaker v2](https://huggingface.co/TencentARC/PhotoMaker-V2) and reimplemented [photomaker v1](https://huggingface.co/TencentARC/PhotoMaker) + - new [PhotoMaker v2](https://huggingface.co/TencentARC/PhotoMaker-V2) and reimplemented [PhotoMaker v1](https://huggingface.co/TencentARC/PhotoMaker) compatible with sdxl models, generates pretty good results and its faster than most other methods select under *scripts -> face -> photomaker* - - new [reswapper](https://github.com/somanchiu/ReSwapper) + - new [ReSwapper](https://github.com/somanchiu/ReSwapper) todo: experimental-only and unfinished, only noting in changelog for future reference +- **Video** + - **hunyuan video** support for [FastHunyuan](https://huggingface.co/FastVideo/FastHunyuan) + simply select model variant and set appropriate parameters + recommended: sampler-shift=17, steps=6, resolution=720x1280, frames=125, guidance>6.0 - **Other**: - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method - **ipex**: update supported torch versions diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py index aac3eb8f0..36edfea25 100644 --- a/scripts/hunyuanvideo.py +++ b/scripts/hunyuanvideo.py @@ -3,10 +3,9 @@ import torch import gradio as gr import transformers import diffusers -from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer +from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, sd_samplers, model_quant, timer -repo_id = 'tencent/HunyuanVideo' default_template = """Describe the video by detailing the following aspects: 1. The main content and theme of the video. 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. @@ -63,27 +62,24 @@ class Script(scripts.Script): def ui(self, is_img2img): with gr.Row(): gr.HTML('  Hunyuan Video
') + with gr.Row(): + repo_id = gr.Dropdown(label='Model', choices=['tencent/HunyuanVideo', 'FastVideo/FastHunyuan'], value='tencent/HunyuanVideo') with gr.Row(): num_frames = gr.Slider(label='Frames', minimum=9, maximum=257, step=1, value=45) tile_frames = gr.Slider(label='Tile frames', minimum=1, maximum=64, step=1, value=16) with gr.Row(): - override_scheduler = gr.Checkbox(label='Override scheduler', value=True) + with gr.Column(): + override_scheduler = gr.Checkbox(label='Override sampler', value=True) + with gr.Column(): + scheduler_shift = gr.Slider(label='Sampler shift', minimum=0.0, maximum=20.0, step=0.1, value=7.0) with gr.Row(): - template = gr.TextArea(label='Prompt processor', lines=3, value=default_template) + template = gr.TextArea(label='Prompt processor', lines=3, value=default_template, visible=False) with gr.Row(): from modules.ui_sections import create_video_inputs video_type, duration, gif_loop, mp4_pad, mp4_interpolate = create_video_inputs(tab='img2img' if is_img2img else 'txt2img') - return [num_frames, tile_frames, override_scheduler, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate] + return [repo_id, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate] - def run(self, p: processing.StableDiffusionProcessing, num_frames, tile_frames, override_scheduler, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument - # set params - num_frames = int(num_frames) - p.width = 16 * int(p.width // 16) - p.height = 16 * int(p.height // 16) - p.do_not_save_grid = True - p.ops.append('video') - - # load model + def load(self, repo_id:str): if shared.sd_model.__class__ != diffusers.HunyuanVideoPipeline: sd_models.unload_model_weights() t0 = time.time() @@ -134,12 +130,27 @@ class Script(scripts.Script): shared.sd_model.vae.enable_slicing() shared.sd_model.vae.enable_tiling() + def run(self, p: processing.StableDiffusionProcessing, repo_id, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument + # set params + num_frames = int(num_frames) + p.width = 16 * int(p.width // 16) + p.height = 16 * int(p.height // 16) + p.do_not_save_grid = True + p.ops.append('video') + + # load model + self.load(repo_id) + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) devices.torch_gc(force=True) if override_scheduler: p.sampler_name = 'Default' - shared.sd_model.scheduler._shift = 7.0 # pylint: disable=protected-access + else: + shared.sd_model.scheduler = sd_samplers.create_sampler(p.sampler_name, shared.sd_model) + p.sampler_name = 'Default' # avoid double creation + if hasattr(shared.sd_model.scheduler, '_shift'): + shared.sd_model.scheduler._shift = scheduler_shift # pylint: disable=protected-access # encode prompt processing.fix_seed(p) From 9fba54e56bfc1b4f7b48de52afc9783ae3ad66c5 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 13:14:07 -0500 Subject: [PATCH 17/66] update diffusers Signed-off-by: Vladimir Mandic --- installer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/installer.py b/installer.py index ea37f55f0..d9b0042e4 100644 --- a/installer.py +++ b/installer.py @@ -492,7 +492,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git: return - sha = 'b785ddb654e4be3ae0066e231734754bdb2a191c' # diffusers commit hash + sha = 'a647682224fed7d65ac4d2a75ed9f2db8e5253e7' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' From fe111d9f1f5ba6482453c982eeac387aeefa2333 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 13:17:06 -0500 Subject: [PATCH 18/66] update diffusers and contributing guidelines Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 ++ CONTRIBUTING | 3 ++- wiki | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 222b599b4..e8b970318 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Update for 2025-01-21 +- **Contributing**: + - if you'd like to contribute, please see updated [contributing](https://github.com/vladmandic/automatic/blob/dev/CONTRIBUTING) guidelines - **Model Merge** - replace model components and merge LoRAs in addition to existing model weights merge support diff --git a/CONTRIBUTING b/CONTRIBUTING index a322cb97e..f6187fa17 100644 --- a/CONTRIBUTING +++ b/CONTRIBUTING @@ -23,4 +23,5 @@ Procedure for contributing: - Make sure that PR does NOT include any unrelated edits - Make sure that PR does not include changes to submodules -Your pull request will be reviewed and pending review results, merged into main branch +Your pull request will be reviewed and pending review results, merged into `dev` branch +Dev merges to main are performed regularly and any PRs that are merged to `dev` will be included in the next main release diff --git a/wiki b/wiki index 1e3ebab7d..faaee0ec3 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 1e3ebab7dcb772a6ad4e4c03be97973c72c36fa4 +Subproject commit faaee0ec3686c8e37fbd55b78385ef343c560847 From 5e0dd3785aa6353c1f84381e89f1a27649448039 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 13:45:16 -0500 Subject: [PATCH 19/66] update wiki Signed-off-by: Vladimir Mandic --- wiki | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wiki b/wiki index faaee0ec3..f8cb064b3 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit faaee0ec3686c8e37fbd55b78385ef343c560847 +Subproject commit f8cb064b38e46fbbea6ae81b7617fb5af0efddf7 From c7a5a24d2d37b02697fd3afae315b472707a979e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 16:17:28 -0500 Subject: [PATCH 20/66] fix mochi num frames Signed-off-by: Vladimir Mandic --- scripts/mochivideo.py | 2 +- wiki | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/mochivideo.py b/scripts/mochivideo.py index a3ed431bb..cbc9dad20 100644 --- a/scripts/mochivideo.py +++ b/scripts/mochivideo.py @@ -28,7 +28,7 @@ class Script(scripts.Script): def run(self, p: processing.StableDiffusionProcessing, num_frames, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument # set params - num_frames = int(num_frames // 8) + num_frames = int(num_frames) p.width = 32 * int(p.width // 32) p.height = 32 * int(p.height // 32) p.task_args['output_type'] = 'pil' diff --git a/wiki b/wiki index f8cb064b3..b53588154 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit f8cb064b38e46fbbea6ae81b7617fb5af0efddf7 +Subproject commit b5358815413223447dfb6fd27e3011072d3853ea From e9853ec0ca8210c4225ace5a1d53310095d99ac1 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 21 Jan 2025 17:10:12 -0500 Subject: [PATCH 21/66] add para-attention first-block-cache Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 8 ++++++++ modules/modeldata.py | 2 ++ modules/para_attention.py | 20 ++++++++++++++++++++ modules/processing.py | 4 +++- modules/sd_offload.py | 2 +- modules/shared.py | 4 ++++ 6 files changed, 38 insertions(+), 2 deletions(-) create mode 100644 modules/para_attention.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e8b970318..275240251 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,12 @@ - **hunyuan video** support for [FastHunyuan](https://huggingface.co/FastVideo/FastHunyuan) simply select model variant and set appropriate parameters recommended: sampler-shift=17, steps=6, resolution=720x1280, frames=125, guidance>6.0 +- [ParaAttention](https://github.com/chengzeyi/ParaAttention) + - first-block caching that can significantly speed up generation by dynamically reusing partial outputs between steps + - available for: flux, hunyuan-video, ltx-video, mochi + - enable in *settings -> pipeline modifiers -> para-attention* + - adjust residual diff threshold to balance the speedup and the accuracy: + higher values leads to more cache hits and speedups, but might also lead to a higher accuracy drop - **Other**: - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method - **ipex**: update supported torch versions @@ -43,6 +49,8 @@ - update ui element ids - modernui use local font - unique font family registration + - mochi video number of frames + - mark large models that should offload ## Update for 2025-01-15 diff --git a/modules/modeldata.py b/modules/modeldata.py index deb4ac49a..63130d041 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -29,6 +29,8 @@ def get_model_type(pipe): model_type = 'auraflow' elif "Flux" in name: model_type = 'f1' + elif "Mochi" in name: + model_type = 'mochi' elif "Lumina" in name: model_type = 'lumina' elif "OmniGen" in name: diff --git a/modules/para_attention.py b/modules/para_attention.py new file mode 100644 index 000000000..5ca53962c --- /dev/null +++ b/modules/para_attention.py @@ -0,0 +1,20 @@ +from modules import shared + + +supported_models = ['Flux', 'HunyuanVideo', 'CogVideoX', 'Mochi'] + + +def apply_first_block_cache(p): + if not shared.opts.para_cache_enabled or not shared.native: + return + if not any(p.sd_model.__class__.__name__.startswith(x) for x in supported_models): + return + from installer import install + install('para_attn') + try: + from para_attn.first_block_cache import diffusers_adapters + diffusers_adapters.apply_cache_on_pipe(p.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold) + shared.log.info(f'Applying para-attn first-block-cache: diff-threshold={shared.opts.para_diff_threshold} cls={p.sd_model.__class__.__name__}') + except Exception as e: + shared.log.error(f'Applying para-attn first-block-cache: {e}') + return diff --git a/modules/processing.py b/modules/processing.py index 99fb9f7f3..549fead3e 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -4,7 +4,7 @@ import time from contextlib import nullcontext import numpy as np from PIL import Image, ImageOps -from modules import shared, devices, errors, images, scripts, memstats, lowvram, script_callbacks, extra_networks, detailer, sd_hijack_freeu, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer, face_restoration, token_merge +from modules import shared, devices, errors, images, scripts, memstats, lowvram, script_callbacks, extra_networks, detailer, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer, face_restoration, token_merge from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet from modules.processing_class import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, StableDiffusionProcessingControl # pylint: disable=unused-import from modules.processing_info import create_infotext @@ -168,7 +168,9 @@ def process_images(p: StableDiffusionProcessing) -> Processed: shared.prompt_styles.extract_comments(p) if shared.opts.cuda_compile_backend == 'none': token_merge.apply_token_merging(p.sd_model) + from modules import sd_hijack_freeu, para_attention sd_hijack_freeu.apply_freeu(p, not shared.native) + para_attention.apply_first_block_cache(p) if p.width is not None: p.width = 8 * int(p.width / 8) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index f9d01528c..8a490e110 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -10,7 +10,7 @@ from modules.timer import process as process_timer debug_move = shared.log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None -should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen'] +should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen', 'hunyuanvideo', 'cogvideox', 'mochi'] offload_hook_instance = None diff --git a/modules/shared.py b/modules/shared.py index 54204c201..d087cdc9e 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -597,6 +597,10 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { "pag_sep": OptionInfo("

Perturbed-Attention Guidance

", "", gr.HTML), "pag_apply_layers": OptionInfo("m0", "PAG layer names"), + "para_sep": OptionInfo("

Para-Attention

", "", gr.HTML), + "para_cache_enabled": OptionInfo(False, "First-block cache enabled"), + "para_diff_threshold": OptionInfo(0.1, "Residual diff threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), + "hypertile_sep": OptionInfo("

HyperTile

", "", gr.HTML), "hypertile_unet_enabled": OptionInfo(False, "UNet Enabled"), "hypertile_hires_only": OptionInfo(False, "HiRes pass only"), From 9f3a34d0a353f758a13e1e174278fba015a35687 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 22 Jan 2025 17:02:22 +0300 Subject: [PATCH 22/66] IPEX enable FP64 emulation --- installer.py | 15 +++++++++++++-- modules/intel/ipex/diffusers.py | 2 +- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/installer.py b/installer.py index d9b0042e4..aa4f1e20b 100644 --- a/installer.py +++ b/installer.py @@ -662,21 +662,32 @@ def install_ipex(torch_command): check_python(supported_minors=[9, 10, 11, 12], reason='IPEX backend requires Python 3.9, 3.10 or 3.11') args.use_ipex = True # pylint: disable=attribute-defined-outside-init log.info('IPEX: Intel OneAPI toolkit detected') + if os.environ.get("NEOReadDebugKeys", None) is None: os.environ.setdefault('NEOReadDebugKeys', '1') if os.environ.get("ClDeviceGlobalMemSizeAvailablePercent", None) is None: os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100') + if os.environ.get("SYCL_CACHE_PERSISTENT", None) is None: + os.environ.setdefault('SYCL_CACHE_PERSISTENT', '1') # Jit cache + if os.environ.get("PYTORCH_ENABLE_XPU_FALLBACK", None) is None: - os.environ.setdefault('PYTORCH_ENABLE_XPU_FALLBACK', '1') + os.environ.setdefault('PYTORCH_ENABLE_XPU_FALLBACK', '1') # CPU fallback for unsupported ops + if os.environ.get("OverrideDefaultFP64Settings", None) is None: + os.environ.setdefault('OverrideDefaultFP64Settings', '1') + if os.environ.get("IGC_EnableDPEmulation", None) is None: + os.environ.setdefault('IGC_EnableDPEmulation', '1') # FP64 Emulation if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is None: - # Battlemage doesn't support Flash Atten or Memory Atten yet so it goes OOM without this + # XPU PyTorch doesn't support Flash Atten or Memory Atten yet so Battlemage goes OOM without this + # Also force enabled because of FP64 emulation making auto-detect based on the GPU invalid os.environ.setdefault('IPEX_FORCE_ATTENTION_SLICE', '1') + if "linux" in sys.platform: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cxx11.abi torchvision==0.20.1+cxx11.abi intel-extension-for-pytorch==2.5.10+xpu oneccl_bind_pt==2.5.0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/cn/') # torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/test/xpu') # test wheels are stable previews, significantly slower than IPEX # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.1') else: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+xpu torchvision==0.21.0+xpu --index-url https://download.pytorch.org/whl/test/xpu') + install(os.environ.get('OPENVINO_PACKAGE', 'openvino==2024.6.0'), 'openvino', ignore=True) install('nncf==2.7.0', ignore=True, no_deps=True) # requires older pandas install(os.environ.get('ONNXRUNTIME_PACKAGE', 'onnxruntime-openvino'), 'onnxruntime-openvino', ignore=True) diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py index 5bf5bbe39..4fbe47c3b 100644 --- a/modules/intel/ipex/diffusers.py +++ b/modules/intel/ipex/diffusers.py @@ -345,7 +345,7 @@ class AttnProcessor: def ipex_diffusers(): diffusers.utils.torch_utils.fourier_filter = fourier_filter #ARC GPUs can't allocate more than 4GB to a single block: - if not device_supports_fp64 or os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is not None: + if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '1' or (not device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0'): diffusers.models.attention_processor.SlicedAttnProcessor = SlicedAttnProcessor diffusers.models.attention_processor.AttnProcessor = AttnProcessor if not device_supports_fp64: From af35296a68398ca5d52c2084c5e7026b5d1baabd Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 22 Jan 2025 18:15:25 +0300 Subject: [PATCH 23/66] IPEX 4GB alloc detection and log driver version --- installer.py | 4 +--- modules/devices.py | 8 +++++++- modules/intel/ipex/__init__.py | 4 ++-- modules/intel/ipex/diffusers.py | 5 ++--- modules/intel/ipex/hijacks.py | 15 ++++++++++++++- 5 files changed, 26 insertions(+), 10 deletions(-) diff --git a/installer.py b/installer.py index aa4f1e20b..60e68e9de 100644 --- a/installer.py +++ b/installer.py @@ -678,13 +678,11 @@ def install_ipex(torch_command): os.environ.setdefault('IGC_EnableDPEmulation', '1') # FP64 Emulation if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is None: # XPU PyTorch doesn't support Flash Atten or Memory Atten yet so Battlemage goes OOM without this - # Also force enabled because of FP64 emulation making auto-detect based on the GPU invalid os.environ.setdefault('IPEX_FORCE_ATTENTION_SLICE', '1') if "linux" in sys.platform: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cxx11.abi torchvision==0.20.1+cxx11.abi intel-extension-for-pytorch==2.5.10+xpu oneccl_bind_pt==2.5.0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/cn/') - # torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/test/xpu') # test wheels are stable previews, significantly slower than IPEX - # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.1') + # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.2') else: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+xpu torchvision==0.21.0+xpu --index-url https://download.pytorch.org/whl/test/xpu') diff --git a/modules/devices.py b/modules/devices.py index c98aea696..2123c05ed 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -82,7 +82,12 @@ def get_backend(shared_cmd_opts): def get_gpu_info(): def get_driver(): import subprocess - if torch.cuda.is_available() and torch.version.cuda: + if torch.xpu.is_available(): + try: + return torch.xpu.get_device_properties(torch.xpu.current_device()).driver_version + except Exception: + return '' + elif torch.cuda.is_available() and torch.version.cuda: try: result = subprocess.run('nvidia-smi --query-gpu=driver_version --format=csv,noheader', shell=True, check=False, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE) version = result.stdout.decode(encoding="utf8", errors="ignore").strip() @@ -121,6 +126,7 @@ def get_gpu_info(): return { 'device': f'{torch.xpu.get_device_name(torch.xpu.current_device())} n={torch.xpu.device_count()}', 'ipex': get_package_version('intel-extension-for-pytorch'), + 'driver': get_driver(), } elif backend == 'cuda' or backend == 'zluda': return { diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py index 94b0cd0d4..a765144b3 100644 --- a/modules/intel/ipex/__init__.py +++ b/modules/intel/ipex/__init__.py @@ -206,10 +206,10 @@ def ipex_init(): # pylint: disable=too-many-statements torch.cuda.ipc_collect = lambda *args, **kwargs: None torch.cuda.utilization = lambda *args, **kwargs: 0 - ipex_hijacks(legacy=legacy) + device_supports_fp64, can_allocate_plus_4gb = ipex_hijacks(legacy=legacy) try: from .diffusers import ipex_diffusers - ipex_diffusers() + ipex_diffusers(device_supports_fp64=device_supports_fp64, can_allocate_plus_4gb=can_allocate_plus_4gb) except Exception: # pylint: disable=broad-exception-caught pass torch.cuda.is_xpu_hijacked = True diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py index 4fbe47c3b..734212237 100644 --- a/modules/intel/ipex/diffusers.py +++ b/modules/intel/ipex/diffusers.py @@ -6,7 +6,6 @@ from diffusers.models.attention_processor import Attention # pylint: disable=protected-access, missing-function-docstring, line-too-long -device_supports_fp64 = torch.xpu.has_fp64_dtype() if hasattr(torch.xpu, "has_fp64_dtype") else torch.xpu.get_device_properties("xpu").has_fp64 attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4)) @@ -342,10 +341,10 @@ class AttnProcessor: return hidden_states -def ipex_diffusers(): +def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False): diffusers.utils.torch_utils.fourier_filter = fourier_filter #ARC GPUs can't allocate more than 4GB to a single block: - if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '1' or (not device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0'): + if not can_allocate_plus_4gb: diffusers.models.attention_processor.SlicedAttnProcessor = SlicedAttnProcessor diffusers.models.attention_processor.AttnProcessor = AttnProcessor if not device_supports_fp64: diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py index 5440b7b68..f1beb671d 100644 --- a/modules/intel/ipex/hijacks.py +++ b/modules/intel/ipex/hijacks.py @@ -6,6 +6,16 @@ import numpy as np from modules import devices, errors device_supports_fp64 = torch.xpu.has_fp64_dtype() if hasattr(torch.xpu, "has_fp64_dtype") else torch.xpu.get_device_properties("xpu").has_fp64 +if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0' and (torch.xpu.get_device_properties("xpu").total_memory / 1024 / 1024 / 1024) > 4.1: + try: + x = torch.ones((33000,33000), dtype=torch.float32, device="xpu") + del x + torch.xpu.empty_cache() + can_allocate_plus_4gb = True + except Exception: + can_allocate_plus_4gb = False +else: + can_allocate_plus_4gb = bool(os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '-1') # pylint: disable=protected-access, missing-function-docstring, line-too-long, unnecessary-lambda, no-else-return @@ -75,7 +85,7 @@ def as_tensor(data, dtype=None, device=None): return original_as_tensor(data, dtype=dtype, device=device) -if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '-1' or (device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0'): +if can_allocate_plus_4gb: original_torch_bmm = torch.bmm original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention else: @@ -190,6 +200,7 @@ def functional_pad(input, pad, mode='constant', value=None): original_torch_tensor = torch.tensor @wraps(torch.tensor) def torch_tensor(data, *args, dtype=None, device=None, **kwargs): + global device_supports_fp64 if check_device(device): device = return_xpu(device) if not device_supports_fp64: @@ -313,6 +324,7 @@ def torch_load(f, map_location=None, *args, **kwargs): # Hijack Functions: def ipex_hijacks(legacy=True): + global device_supports_fp64, can_allocate_plus_4gb if legacy and float(torch.__version__[:3]) < 2.5: torch.nn.functional.interpolate = interpolate torch.tensor = torch_tensor @@ -350,3 +362,4 @@ def ipex_hijacks(legacy=True): if not device_supports_fp64: torch.from_numpy = from_numpy torch.as_tensor = as_tensor + return device_supports_fp64, can_allocate_plus_4gb From 0926ac1c05f31c687600652e9ca88dc885c69116 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 22 Jan 2025 18:33:25 +0300 Subject: [PATCH 24/66] Update changelog --- CHANGELOG.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 275240251..421f6207e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,9 +31,13 @@ - enable in *settings -> pipeline modifiers -> para-attention* - adjust residual diff threshold to balance the speedup and the accuracy: higher values leads to more cache hits and speedups, but might also lead to a higher accuracy drop +- **IPEX** + - enable force attention slicing for all GPUs + - enable fp64 emulation and sycl jit cache + - switch to pytorch test branch on windows + - update the supported python versions - **Other**: - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method - - **ipex**: update supported torch versions - **gallery**: add http fallback for slow/unreliable links - **splash**: add legacy mode indicator on splash screen - **network**: extract thumbnail from model metadata if present From bda45d865e154610c6c77de48521d27eb48e270c Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 22 Jan 2025 11:01:58 -0500 Subject: [PATCH 25/66] perflow Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 + modules/perflow/__init__.py | 4 + modules/perflow/pfode_solver.py | 257 +++++++++++++++++++ modules/perflow/scheduler_perflow.py | 368 +++++++++++++++++++++++++++ modules/perflow/utils_perflow.py | 75 ++++++ modules/sd_samplers_diffusers.py | 3 + wiki | 2 +- 7 files changed, 711 insertions(+), 1 deletion(-) create mode 100644 modules/perflow/__init__.py create mode 100644 modules/perflow/pfode_solver.py create mode 100644 modules/perflow/scheduler_perflow.py create mode 100644 modules/perflow/utils_perflow.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 421f6207e..5cd1efa24 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,9 @@ - enable fp64 emulation and sycl jit cache - switch to pytorch test branch on windows - update the supported python versions +- [PerFlow](https://github.com/magic-research/piecewise-rectified-flow) + - piecewise rectified flow as model acceleration + - use `perflow` scheduler combined with one of the available pre-trained [models](https://huggingface.co/hansyan) - **Other**: - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method - **gallery**: add http fallback for slow/unreliable links diff --git a/modules/perflow/__init__.py b/modules/perflow/__init__.py new file mode 100644 index 000000000..32213e36e --- /dev/null +++ b/modules/perflow/__init__.py @@ -0,0 +1,4 @@ +### original: + +from .scheduler_perflow import PeRFlowScheduler +from .utils_perflow import merge_delta_weights_into_unet diff --git a/modules/perflow/pfode_solver.py b/modules/perflow/pfode_solver.py new file mode 100644 index 000000000..db1b8e903 --- /dev/null +++ b/modules/perflow/pfode_solver.py @@ -0,0 +1,257 @@ +import torch +import torch.utils.checkpoint + + +class PFODESolver(): + def __init__(self, scheduler, t_initial=1, t_terminal=0,) -> None: + self.t_initial = t_initial + self.t_terminal = t_terminal + self.scheduler = scheduler + + train_step_terminal = 0 + train_step_initial = train_step_terminal + self.scheduler.config.num_train_timesteps # 0+1000 + self.stepsize = (t_terminal-t_initial) / (train_step_terminal - train_step_initial) #1/1000 + + def get_timesteps(self, t_start, t_end, num_steps): + # (b,) -> (b,1) + t_start = t_start[:, None] + t_end = t_end[:, None] + assert t_start.dim() == 2 + + timepoints = torch.arange(0, num_steps, 1).expand(t_start.shape[0], num_steps).to(device=t_start.device) + interval = (t_end - t_start) / (torch.ones([1], device=t_start.device) * num_steps) + timepoints = t_start + interval * timepoints + + timesteps = (self.scheduler.num_train_timesteps - 1) + (timepoints - self.t_initial) / self.stepsize # correspondint to StableDiffusion indexing system, from 999 (t_init) -> 0 (dt) + return timesteps.round().long() + # return timesteps.floor().long() + + def solve(self, + latents, + unet, + t_start, + t_end, + prompt_embeds, + negative_prompt_embeds, + guidance_scale=1.0, + num_steps = 2, + num_windows = 1, + ): + assert t_start.dim() == 1 + assert guidance_scale >= 1 and torch.all(torch.gt(t_start, t_end)) + + do_classifier_free_guidance = True if guidance_scale > 1 else False + bsz = latents.shape[0] + + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + + timestep_cond = None + if unet.config.time_cond_proj_dim is not None: + guidance_scale_tensor = torch.tensor(guidance_scale - 1).repeat(bsz) + timestep_cond = self.get_guidance_scale_embedding( + guidance_scale_tensor, embedding_dim=unet.config.time_cond_proj_dim + ).to(device=latents.device, dtype=latents.dtype) + + timesteps = self.get_timesteps(t_start, t_end, num_steps).to(device=latents.device) + timestep_interval = self.scheduler.config.num_train_timesteps // (num_windows * num_steps) + + # 7. Denoising loop + with torch.no_grad(): + # for i in tqdm(range(num_steps)): + for i in range(num_steps): + + t = torch.cat([timesteps[:, i]]*2) if do_classifier_free_guidance else timesteps[:, i] + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + noise_pred = unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + timestep_cond=timestep_cond, + return_dict=False, + )[0] + + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # STEP: compute the previous noisy sample x_t -> x_t-1 + # latents = self.scheduler.step(noise_pred, timesteps[:, i].cpu(), latents, return_dict=False)[0] + + batch_timesteps = timesteps[:, i].cpu() + prev_timestep = batch_timesteps - timestep_interval + # prev_timestep = batch_timesteps - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps + + alpha_prod_t = self.scheduler.alphas_cumprod[batch_timesteps] + alpha_prod_t_prev = torch.zeros_like(alpha_prod_t) + for ib in range(prev_timestep.shape[0]): + alpha_prod_t_prev[ib] = self.scheduler.alphas_cumprod[prev_timestep[ib]] if prev_timestep[ib] >= 0 else self.scheduler.final_alpha_cumprod + beta_prod_t = 1 - alpha_prod_t + + alpha_prod_t = alpha_prod_t.to(device=latents.device, dtype=latents.dtype) + alpha_prod_t_prev = alpha_prod_t_prev.to(device=latents.device, dtype=latents.dtype) + beta_prod_t = beta_prod_t.to(device=latents.device, dtype=latents.dtype) + + # 3. compute predicted original sample from predicted noise also called + # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf + if self.scheduler.config.prediction_type == "epsilon": + pred_original_sample = (latents - beta_prod_t[:,None,None,None] ** (0.5) * noise_pred) / alpha_prod_t[:, None,None,None] ** (0.5) + pred_epsilon = noise_pred + # elif self.scheduler.config.prediction_type == "sample": + # pred_original_sample = noise_pred + # pred_epsilon = (latents - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + elif self.scheduler.config.prediction_type == "v_prediction": + pred_original_sample = (alpha_prod_t[:,None,None,None]**0.5) * latents - (beta_prod_t[:,None,None,None]**0.5) * noise_pred + pred_epsilon = (alpha_prod_t[:,None,None,None]**0.5) * noise_pred + (beta_prod_t[:,None,None,None]**0.5) * latents + else: + raise ValueError( + f"prediction_type given as {self.scheduler.config.prediction_type} must be one of `epsilon`, `sample`, or" + " `v_prediction`" + ) + pred_sample_direction = (1 - alpha_prod_t_prev[:,None,None,None]) ** (0.5) * pred_epsilon + latents = alpha_prod_t_prev[:,None,None,None] ** (0.5) * pred_original_sample + pred_sample_direction + + return latents + + +class PFODESolverSDXL(): + def __init__(self, scheduler, t_initial=1, t_terminal=0,) -> None: + self.t_initial = t_initial + self.t_terminal = t_terminal + self.scheduler = scheduler + + train_step_terminal = 0 + train_step_initial = train_step_terminal + self.scheduler.config.num_train_timesteps # 0+1000 + + self.stepsize = (t_terminal-t_initial) / (train_step_terminal - train_step_initial) #1/1000 + + def get_timesteps(self, t_start, t_end, num_steps): + # (b,) -> (b,1) + t_start = t_start[:, None] + t_end = t_end[:, None] + assert t_start.dim() == 2 + + timepoints = torch.arange(0, num_steps, 1).expand(t_start.shape[0], num_steps).to(device=t_start.device) + interval = (t_end - t_start) / (torch.ones([1], device=t_start.device) * num_steps) + timepoints = t_start + interval * timepoints + + timesteps = (self.scheduler.num_train_timesteps - 1) + (timepoints - self.t_initial) / self.stepsize # correspondint to StableDiffusion indexing system, from 999 (t_init) -> 0 (dt) + return timesteps.round().long() + # return timesteps.floor().long() + + def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype): + # Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids + add_time_ids = list(original_size + crops_coords_top_left + target_size) + add_time_ids = torch.tensor([add_time_ids], dtype=dtype) + return add_time_ids + + def solve(self, + latents, + unet, + t_start, + t_end, + prompt_embeds, + pooled_prompt_embeds, + negative_prompt_embeds, + negative_pooled_prompt_embeds, + guidance_scale=1.0, + num_steps = 10, + num_windows = 4, + resolution = 1024, + ): + assert t_start.dim() == 1 + assert guidance_scale >= 1 and torch.all(torch.gt(t_start, t_end)) + dtype = latents.dtype + device = latents.device + bsz = latents.shape[0] + do_classifier_free_guidance = True if guidance_scale > 1 else False + + add_text_embeds = pooled_prompt_embeds + add_time_ids = torch.cat( + # [self._get_add_time_ids((1024, 1024), (0, 0), (1024, 1024), dtype) for _ in range(bsz)] + [self._get_add_time_ids((resolution, resolution), (0, 0), (resolution, resolution), dtype) for _ in range(bsz)] + ).to(device) + negative_add_time_ids = add_time_ids + + if do_classifier_free_guidance: + # prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) + add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) + + timestep_cond = None + if unet.config.time_cond_proj_dim is not None: + guidance_scale_tensor = torch.tensor(guidance_scale - 1).repeat(bsz) + timestep_cond = self.get_guidance_scale_embedding( + guidance_scale_tensor, embedding_dim=unet.config.time_cond_proj_dim + ).to(device=latents.device, dtype=latents.dtype) + + timesteps = self.get_timesteps(t_start, t_end, num_steps).to(device=latents.device) + timestep_interval = self.scheduler.config.num_train_timesteps // (num_windows * num_steps) + + # 7. Denoising loop + with torch.no_grad(): + # for i in tqdm(range(num_steps)): + for i in range(num_steps): + # expand the latents if we are doing classifier free guidance + t = torch.cat([timesteps[:, i]]*2) if do_classifier_free_guidance else timesteps[:, i] + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} + noise_pred = unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + timestep_cond=timestep_cond, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0] + + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + + # STEP: compute the previous noisy sample x_t -> x_t-1 + # latents = self.scheduler.step(noise_pred, timesteps[:, i].cpu(), latents, return_dict=False)[0] + + batch_timesteps = timesteps[:, i].cpu() + prev_timestep = batch_timesteps - timestep_interval + # prev_timestep = batch_timesteps - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps + + alpha_prod_t = self.scheduler.alphas_cumprod[batch_timesteps] + alpha_prod_t_prev = torch.zeros_like(alpha_prod_t) + for ib in range(prev_timestep.shape[0]): + alpha_prod_t_prev[ib] = self.scheduler.alphas_cumprod[prev_timestep[ib]] if prev_timestep[ib] >= 0 else self.scheduler.final_alpha_cumprod + beta_prod_t = 1 - alpha_prod_t + + alpha_prod_t = alpha_prod_t.to(device=latents.device, dtype=latents.dtype) + alpha_prod_t_prev = alpha_prod_t_prev.to(device=latents.device, dtype=latents.dtype) + beta_prod_t = beta_prod_t.to(device=latents.device, dtype=latents.dtype) + + # 3. compute predicted original sample from predicted noise also called + # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf + if self.scheduler.config.prediction_type == "epsilon": + pred_original_sample = (latents - beta_prod_t[:,None,None,None] ** (0.5) * noise_pred) / alpha_prod_t[:, None,None,None] ** (0.5) + pred_epsilon = noise_pred + # elif self.scheduler.config.prediction_type == "sample": + # pred_original_sample = noise_pred + # pred_epsilon = (latents - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + # elif self.scheduler.config.prediction_type == "v_prediction": + # pred_original_sample = (alpha_prod_t**0.5) * latents - (beta_prod_t**0.5) * noise_pred + # pred_epsilon = (alpha_prod_t**0.5) * noise_pred + (beta_prod_t**0.5) * latents + else: + raise ValueError( + f"prediction_type given as {self.scheduler.config.prediction_type} must be one of `epsilon`, `sample`, or" + " `v_prediction`" + ) + pred_sample_direction = (1 - alpha_prod_t_prev[:,None,None,None]) ** (0.5) * pred_epsilon + latents = alpha_prod_t_prev[:,None,None,None] ** (0.5) * pred_original_sample + pred_sample_direction + + return latents diff --git a/modules/perflow/scheduler_perflow.py b/modules/perflow/scheduler_perflow.py new file mode 100644 index 000000000..e3a50feaf --- /dev/null +++ b/modules/perflow/scheduler_perflow.py @@ -0,0 +1,368 @@ +# Copyright 2023 Stanford University Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion +# and https://github.com/hojonathanho/diffusion + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import BaseOutput +from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin + + +class Time_Windows(): + def __init__(self, t_initial=1, t_terminal=0, num_windows=4, precision=1./1000) -> None: + assert t_terminal < t_initial + time_windows = [ 1.*i/num_windows for i in range(1, num_windows+1)][::-1] + + self.window_starts = time_windows # [1.0, 0.75, 0.5, 0.25] + self.window_ends = time_windows[1:] + [t_terminal] # [0.75, 0.5, 0.25, 0] + self.precision = precision + + def get_window(self, tp): + idx = 0 + # robust to numerical error; e.g, (0.6+1/10000) belongs to [0.6, 0.3) + while (tp-0.1*self.precision) <= self.window_ends[idx]: + idx += 1 + return self.window_starts[idx], self.window_ends[idx] + + def lookup_window(self, timepoint): + if timepoint.dim() == 0: + t_start, t_end = self.get_window(timepoint) + t_start = torch.ones_like(timepoint) * t_start + t_end = torch.ones_like(timepoint) * t_end + else: + t_start = torch.zeros_like(timepoint) + t_end = torch.zeros_like(timepoint) + bsz = timepoint.shape[0] + for i in range(bsz): + tp = timepoint[i] + ts, te = self.get_window(tp) + t_start[i] = ts + t_end[i] = te + return t_start, t_end + + +@dataclass +class PeRFlowSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + The predicted denoised sample `(x_{0})` based on the model output from the current timestep. + `pred_original_sample` can be used to preview progress or for guidance. + """ + + prev_sample: torch.FloatTensor + pred_original_sample: Optional[torch.FloatTensor] = None + + +# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar +def betas_for_alpha_bar( + num_diffusion_timesteps, + max_beta=0.999, + alpha_transform_type="cosine", +): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of + (1-beta) over time from t = [0,1]. + + Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up + to that part of the diffusion process. + + + Args: + num_diffusion_timesteps (`int`): the number of betas to produce. + max_beta (`float`): the maximum beta to use; use values lower than 1 to + prevent singularities. + alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. + Choose from `cosine` or `exp` + + Returns: + betas (`np.ndarray`): the betas used by the scheduler to step the model outputs + """ + if alpha_transform_type == "cosine": + + def alpha_bar_fn(t): + return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + + elif alpha_transform_type == "exp": + + def alpha_bar_fn(t): + return math.exp(t * -12.0) + + else: + raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}") + + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) + return torch.tensor(betas, dtype=torch.float32) + + + +class PeRFlowScheduler(SchedulerMixin, ConfigMixin): + """ + `ReFlowScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with + non-Markovian guidance. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + beta_start (`float`, defaults to 0.0001): + The starting `beta` value of inference. + beta_end (`float`, defaults to 0.02): + The final `beta` value. + beta_schedule (`str`, defaults to `"linear"`): + The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from + `linear`, `scaled_linear`, or `squaredcos_cap_v2`. + trained_betas (`np.ndarray`, *optional*): + Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. + set_alpha_to_one (`bool`, defaults to `True`): + Each diffusion step uses the alphas product value at that step and at the previous one. For the final step + there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, + otherwise it uses the alpha value at step 0. + prediction_type (`str`, defaults to `epsilon`, *optional*) + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + beta_start: float = 0.00085, + beta_end: float = 0.012, + beta_schedule: str = "scaled_linear", + trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + set_alpha_to_one: bool = False, + prediction_type: str = "ddim_eps", + t_noise: float = 1, + t_clean: float = 0, + num_time_windows = 4, + ): + if trained_betas is not None: + self.betas = torch.tensor(trained_betas, dtype=torch.float32) + elif beta_schedule == "linear": + self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) + elif beta_schedule == "scaled_linear": + # this schedule is very specific to the latent diffusion model. + self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + elif beta_schedule == "squaredcos_cap_v2": + # Glide cosine schedule + self.betas = betas_for_alpha_bar(num_train_timesteps) + else: + raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}") + + self.alphas = 1.0 - self.betas + self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) + + # At every step in ddim, we are looking into the previous alphas_cumprod + # For the final step, there is no previous alphas_cumprod because we are already at 0 + # `set_alpha_to_one` decides whether we set this parameter simply to one or + # whether we use the final alpha of the "non-previous" one. + self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] + + # # standard deviation of the initial noise distribution + self.init_noise_sigma = 1.0 + + self.time_windows = Time_Windows(t_initial=t_noise, t_terminal=t_clean, + num_windows=num_time_windows, + precision=1./num_train_timesteps) + + assert prediction_type in ["ddim_eps", "diff_eps", "velocity"] + + + def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor: # pylint: disable=unused-argument + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`int`, *optional*): + The current timestep in the diffusion chain. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + return sample + + + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + """ + if num_inference_steps < self.config.num_time_windows: # pylint: disable=no-member + num_inference_steps = self.config.num_time_windows # pylint: disable=no-member + print(f"### We recommend a num_inference_steps not less than num_time_windows. It's set as {self.config.num_time_windows}.") # pylint: disable=no-member + + timesteps = [] + for i in range(self.config.num_time_windows): # pylint: disable=no-member + if i < num_inference_steps%self.config.num_time_windows: # pylint: disable=no-member + num_steps_cur_win = num_inference_steps//self.config.num_time_windows+1 # pylint: disable=no-member + else: + num_steps_cur_win = num_inference_steps//self.config.num_time_windows # pylint: disable=no-member + + t_s = self.time_windows.window_starts[i] + t_e = self.time_windows.window_ends[i] + timesteps_cur_win = np.linspace(t_s, t_e, num=num_steps_cur_win, endpoint=False) + timesteps.append(timesteps_cur_win) + + timesteps = np.concatenate(timesteps) + + self.timesteps = torch.from_numpy( # pylint: disable=attribute-defined-outside-init + (timesteps*self.config.num_train_timesteps).astype(np.int64) # pylint: disable=no-member, + ).to(device) + + def get_window_alpha(self, timepoints): + time_windows = self.time_windows + num_train_timesteps = self.config.num_train_timesteps # pylint: disable=no-member + + t_win_start, t_win_end = time_windows.lookup_window(timepoints) + t_win_len = t_win_end - t_win_start + t_interval = timepoints - t_win_start # NOTE: negative value + + idx_start = (t_win_start*num_train_timesteps - 1 ).long() + alphas_cumprod_start = self.alphas_cumprod[idx_start] + + idx_end = torch.clamp( (t_win_end*num_train_timesteps - 1 ).long(), min=0) + alphas_cumprod_end = self.alphas_cumprod[idx_end] + + alpha_cumprod_s_e = alphas_cumprod_start / alphas_cumprod_end + gamma_s_e = alpha_cumprod_s_e ** 0.5 + + return t_win_start, t_win_end, t_win_len, t_interval, gamma_s_e, alphas_cumprod_start, alphas_cumprod_end + + def step( + self, + model_output: torch.FloatTensor, + timestep: int, + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> Union[PeRFlowSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~schedulers.scheduling_ddim.PeRFlowSchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.PeRFlowSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_ddim.PeRFlowSchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + """ + + if self.config.prediction_type == "ddim_eps": # pylint: disable=no-member + pred_epsilon = model_output + t_c = timestep / self.config.num_train_timesteps # pylint: disable=no-member + t_s, t_e, _, c_to_s, _, alphas_cumprod_start, alphas_cumprod_end = self.get_window_alpha(t_c) + + lambda_s = (alphas_cumprod_end / alphas_cumprod_start)**0.5 + eta_s = (1-alphas_cumprod_end)**0.5 - ( alphas_cumprod_end / alphas_cumprod_start * (1-alphas_cumprod_start) )**0.5 + + lambda_t = ( lambda_s * (t_e - t_s) ) / ( lambda_s *(t_c - t_s) + (t_e - t_c) ) + eta_t = ( eta_s * (t_e - t_c) ) / ( lambda_s *(t_c - t_s) + (t_e - t_c) ) + + pred_win_end = lambda_t * sample + eta_t * pred_epsilon + pred_velocity = (pred_win_end - sample) / (t_e - (t_s + c_to_s)) + + elif self.config.prediction_type == "diff_eps": # pylint: disable=no-member + pred_epsilon = model_output + t_c = timestep / self.config.num_train_timesteps # pylint: disable=no-member + t_s, t_e, _, c_to_s, gamma_s_e, _, _ = self.get_window_alpha(t_c) + + lambda_s = 1 / gamma_s_e + eta_s = -1 * ( 1- gamma_s_e**2)**0.5 / gamma_s_e + + lambda_t = ( lambda_s * (t_e - t_s) ) / ( lambda_s *(t_c - t_s) + (t_e - t_c) ) + eta_t = ( eta_s * (t_e - t_c) ) / ( lambda_s *(t_c - t_s) + (t_e - t_c) ) + + pred_win_end = lambda_t * sample + eta_t * pred_epsilon + pred_velocity = (pred_win_end - sample) / (t_e - (t_s + c_to_s)) + + elif self.config.prediction_type == "velocity": # pylint: disable=no-member + pred_velocity = model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `velocity`." # pylint: disable=no-member + ) + + # get dt + idx = torch.argwhere(torch.where(self.timesteps==timestep, 1,0)) + prev_step = self.timesteps[idx+1] if (idx+1) torch.FloatTensor: + # Make sure alphas_cumprod and timestep have same device and dtype as original_samples + alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype) + timesteps = timesteps.to(original_samples.device) - 1 # indexing from 0 + + sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 + sqrt_alpha_prod = sqrt_alpha_prod.flatten() + while len(sqrt_alpha_prod.shape) < len(original_samples.shape): + sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) + + sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() + while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) + + noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps # pylint: disable=no-member diff --git a/modules/perflow/utils_perflow.py b/modules/perflow/utils_perflow.py new file mode 100644 index 000000000..35af13cf7 --- /dev/null +++ b/modules/perflow/utils_perflow.py @@ -0,0 +1,75 @@ +import os +from collections import OrderedDict +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline +from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_unet_checkpoint, convert_ldm_vae_checkpoint, convert_ldm_clip_checkpoint + + +def merge_delta_weights_into_unet(pipe, delta_weights): + unet_weights = pipe.unet.state_dict() + assert unet_weights.keys() == delta_weights.keys() + for key in delta_weights.keys(): + dtype = unet_weights[key].dtype + unet_weights[key] = unet_weights[key].to(dtype=delta_weights[key].dtype) + delta_weights[key].to(device=unet_weights[key].device) + unet_weights[key] = unet_weights[key].to(dtype) + pipe.unet.load_state_dict(unet_weights, strict=True) + return pipe + + +def load_delta_weights_into_unet( + pipe, + model_path = "hsyan/piecewise-rectified-flow-v0-1", + base_path = "runwayml/stable-diffusion-v1-5", +): + ## load delta_weights + if os.path.exists(os.path.join(model_path, "delta_weights.safetensors")): + print("### delta_weights exists, loading...") + delta_weights = OrderedDict() + with safe_open(os.path.join(model_path, "delta_weights.safetensors"), framework="pt", device="cpu") as f: + for key in f.keys(): + delta_weights[key] = f.get_tensor(key) + + elif os.path.exists(os.path.join(model_path, "diffusion_pytorch_model.safetensors")): + print("### merged_weights exists, loading...") + merged_weights = OrderedDict() + with safe_open(os.path.join(model_path, "diffusion_pytorch_model.safetensors"), framework="pt", device="cpu") as f: + for key in f.keys(): + merged_weights[key] = f.get_tensor(key) + + base_weights = StableDiffusionPipeline.from_pretrained( + base_path, torch_dtype=torch.float16, safety_checker=None).unet.state_dict() + assert base_weights.keys() == merged_weights.keys() + + delta_weights = OrderedDict() + for key in merged_weights.keys(): + delta_weights[key] = merged_weights[key] - base_weights[key].to(device=merged_weights[key].device, dtype=merged_weights[key].dtype) + + print("### saving delta_weights...") + save_file(delta_weights, os.path.join(model_path, "delta_weights.safetensors")) + + else: + raise ValueError(f"{model_path} does not contain delta weights or merged weights") + + ## merge delta_weights to the target pipeline + pipe = merge_delta_weights_into_unet(pipe, delta_weights) + return pipe + + +def load_dreambooth_into_pipeline(pipe, sd_dreambooth): + assert sd_dreambooth.endswith(".safetensors") + state_dict = {} + with safe_open(sd_dreambooth, framework="pt", device="cpu") as f: + for key in f.keys(): + state_dict[key] = f.get_tensor(key) + + unet_config = {} # unet, line 449 in convert_ldm_unet_checkpoint + for key in pipe.unet.config.keys(): + if key != 'num_class_embeds': + unet_config[key] = pipe.unet.config[key] + + pipe.unet.load_state_dict(convert_ldm_unet_checkpoint(state_dict, unet_config), strict=False) + pipe.vae.load_state_dict(convert_ldm_vae_checkpoint(state_dict, pipe.vae.config)) + pipe.text_encoder = convert_ldm_clip_checkpoint(state_dict, text_encoder=pipe.text_encoder) + return pipe diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index a8b77bc90..f24641d27 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -48,6 +48,7 @@ try: from modules.schedulers.scheduler_dpm_flowmatch import FlowMatchDPMSolverMultistepScheduler # pylint: disable=ungrouped-imports from modules.schedulers.scheduler_bdia import BDIA_DDIMScheduler # pylint: disable=ungrouped-imports from modules.schedulers.scheduler_ufogen import UFOGenScheduler # pylint: disable=ungrouped-imports + from modules.perflow import PeRFlowScheduler # pylint: disable=ungrouped-imports except Exception as e: shared.log.error(f'Diffusers import error: version={diffusers.__version__} error: {e}') if os.environ.get('SD_SAMPLER_DEBUG', None) is not None: @@ -94,6 +95,7 @@ config = { 'LCM': { 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'thresholding': False, 'timestep_spacing': 'linspace' }, 'TCD': { 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'beta_schedule': 'scaled_linear' }, 'TDD': { }, + 'PeRFlow': { 'prediction_type': 'ddim_eps' }, 'UFOGen': { }, 'BDIA DDIM': { 'clip_sample': False, 'set_alpha_to_one': True, 'steps_offset': 0, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'leading', 'rescale_betas_zero_snr': False, 'thresholding': False, 'gamma': 1.0 }, @@ -154,6 +156,7 @@ samplers_data_diffusers = [ SamplerData('LCM', lambda model: DiffusionSampler('LCM', LCMScheduler, model), [], {}), SamplerData('TCD', lambda model: DiffusionSampler('TCD', TCDScheduler, model), [], {}), SamplerData('TDD', lambda model: DiffusionSampler('TDD', TDDScheduler, model), [], {}), + SamplerData('PeRFlow', lambda model: DiffusionSampler('PeRFlow', PeRFlowScheduler, model), [], {}), SamplerData('UFOGen', lambda model: DiffusionSampler('UFOGen', UFOGenScheduler, model), [], {}), SamplerData('Same as primary', None, [], {}), diff --git a/wiki b/wiki index b53588154..eb2e03e71 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit b5358815413223447dfb6fd27e3011072d3853ea +Subproject commit eb2e03e719e200ee8c2523a08e75caf7e43df1cd From 40e7dcd5a48ff2294203c71f15eb2101e9fa17de Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 23 Jan 2025 00:15:27 +0300 Subject: [PATCH 26/66] OpenVINO add group_size option to compression --- modules/intel/openvino/__init__.py | 8 ++++++-- modules/shared.py | 3 ++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index 1c89ceb48..a4000c43a 100644 --- a/modules/intel/openvino/__init__.py +++ b/modules/intel/openvino/__init__.py @@ -232,7 +232,9 @@ def openvino_compile(gm: GraphModule, *example_inputs, model_hash_str: str = Non if dont_use_4bit_nncf or shared.opts.nncf_compress_weights_mode == "INT8": om = nncf.compress_weights(om) else: - om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito) + compress_group_size = shared.opts.nncf_compress_weights_group_size if shared.opts.nncf_compress_weights_group_size != 0 else None + compress_ratio = shared.opts.nncf_compress_weights_raito if shared.opts.nncf_compress_weights_raito != 0 else None + om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=compress_group_size, ratio=compress_ratio) hints = {} if shared.opts.openvino_accuracy == "performance": @@ -290,7 +292,9 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs): if dont_use_4bit_nncf or shared.opts.nncf_compress_weights_mode == "INT8": om = nncf.compress_weights(om) else: - om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito) + compress_group_size = shared.opts.nncf_compress_weights_group_size if shared.opts.nncf_compress_weights_group_size != 0 else None + compress_ratio = shared.opts.nncf_compress_weights_raito if shared.opts.nncf_compress_weights_raito != 0 else None + om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=compress_group_size, ratio=compress_ratio) hints = {'CACHE_DIR': shared.opts.openvino_cache_path + '/blob'} if shared.opts.openvino_accuracy == "performance": diff --git a/modules/shared.py b/modules/shared.py index d087cdc9e..2200813ea 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -575,7 +575,8 @@ options_templates.update(options_section(('quantization', "Quantization Settings "nncf_sep": OptionInfo("

NNCF

", "", gr.HTML), "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), "nncf_compress_weights_mode": OptionInfo("INT8", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}), - "nncf_compress_weights_raito": OptionInfo(1.0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), + "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), + "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 512, "step": 1, "visible": cmd_opts.use_openvino}), "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": cmd_opts.use_openvino}), "nncf_quant_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}), "quant_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), From e58b9b50e9cbc923f19a4cfe6d3c97997d8db8af Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 23 Jan 2025 01:10:53 +0300 Subject: [PATCH 27/66] Don't force bnb version outside of cuda --- modules/model_quant.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/modules/model_quant.py b/modules/model_quant.py index d570c6ced..0665fc920 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -70,10 +70,13 @@ def load_torchao(msg='', silent=False): def load_bnb(msg='', silent=False): + from modules import devices global bnb # pylint: disable=global-statement if bnb is not None: return bnb - install('bitsandbytes==0.45.0', quiet=True) + if devices.backend == 'cuda': + # forcing a version will uninstall the multi-backend-refactor branch of bnb + install('bitsandbytes==0.45.0', quiet=True) try: import bitsandbytes bnb = bitsandbytes From 1f377a49957a7f51b32687348e1153aef4add3f8 Mon Sep 17 00:00:00 2001 From: Cade Schlaefli Date: Tue, 21 Jan 2025 19:21:37 -0600 Subject: [PATCH 28/66] update = to == in optional dependencies --- installer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/installer.py b/installer.py index 1e62ddd2a..ea62f115d 100644 --- a/installer.py +++ b/installer.py @@ -1056,7 +1056,7 @@ def install_optional(): install('gfpgan') install('clean-fid') install('pillow-jxl-plugin==1.3.1', ignore=True) - install('optimum-quanto=0.2.6', ignore=True) + install('optimum-quanto==0.2.6', ignore=True) install('bitsandbytes==0.45.0', ignore=True) install('pynvml', ignore=True) install('ultralytics==8.3.40', ignore=True) From 9a67e5aeb298bfda4c19bc7cc7e53eb933b3930c Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 23 Jan 2025 09:33:55 -0500 Subject: [PATCH 29/66] update diffusers Signed-off-by: Vladimir Mandic --- extensions-builtin/sd-extension-system-info | 2 +- installer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index f55066aa4..2e50e314a 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit f55066aa419509fa0eafe92e4e98a4637b19b9c5 +Subproject commit 2e50e314abfab364f41c29b079e37f9658bb1627 diff --git a/installer.py b/installer.py index 7f81e62c6..f10738f13 100644 --- a/installer.py +++ b/installer.py @@ -492,7 +492,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git: return - sha = 'a647682224fed7d65ac4d2a75ed9f2db8e5253e7' # diffusers commit hash + sha = '78bc824729f76a14ff2f211fc7f9a31e5500a41e' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' From 5a7c1f50c100bc1fb2ebf382ffa288687645fc7e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 23 Jan 2025 12:07:08 -0500 Subject: [PATCH 30/66] add native torch fp8 storage dtype Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 16 +++++++--- modules/model_quant.py | 72 ++++++++++++++++++++++++++++++++++-------- modules/sd_models.py | 5 ++- modules/sd_offload.py | 17 ++++++---- modules/shared.py | 8 +++++ 5 files changed, 91 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5cd1efa24..6b1a291a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-21 +## Update for 2025-01-23 - **Contributing**: - if you'd like to contribute, please see updated [contributing](https://github.com/vladmandic/automatic/blob/dev/CONTRIBUTING) guidelines @@ -32,10 +32,14 @@ - adjust residual diff threshold to balance the speedup and the accuracy: higher values leads to more cache hits and speedups, but might also lead to a higher accuracy drop - **IPEX** - - enable force attention slicing for all GPUs - - enable fp64 emulation and sycl jit cache - - switch to pytorch test branch on windows - - update the supported python versions + - enable force attention slicing, fp64 emulation, jit cache + - use pytorch test branch on windows + - extend the supported python versions +- **Torch FP8** + - uses torch `float8_e4m3fn` or `float8_e5m2` as data storage and performs dynamic upcasting to compute `dtype` as needed + - compatible with most `unet` and `transformer` models: e.g. *sd15, sdxl, sd35, flux.1, hunyuan-video, ltx-video, etc.* + this is alternative to `bnb`/`quanto`/`torchao` quantization on models/platforms/gpus where those libraries are not available + - enable in *settings -> quantization -> layerwise casting* - [PerFlow](https://github.com/magic-research/piecewise-rectified-flow) - piecewise rectified flow as model acceleration - use `perflow` scheduler combined with one of the available pre-trained [models](https://huggingface.co/hansyan) @@ -58,6 +62,8 @@ - unique font family registration - mochi video number of frames - mark large models that should offload + - avoid repeated optimum-quanto installation + - avoid reinstalling bnb if not cuda ## Update for 2025-01-15 diff --git a/modules/model_quant.py b/modules/model_quant.py index 0665fc920..0a56805df 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -8,6 +8,24 @@ quanto = None ao = None +def get_quant(name): + if "qint8" in name.lower(): + return 'qint8' + if "qint4" in name.lower(): + return 'qint4' + if "fp8" in name.lower(): + return 'fp8' + if "fp4" in name.lower(): + return 'fp4' + if "nf4" in name.lower(): + return 'nf4' + if name.endswith('.gguf'): + return 'gguf' + return 'none' + + + + def create_bnb_config(kwargs = None, allow_bnb: bool = True): from modules import shared, devices if len(shared.opts.bnb_quantization) > 0 and allow_bnb: @@ -117,17 +135,43 @@ def load_quanto(msg='', silent=False): return None -def get_quant(name): - if "qint8" in name.lower(): - return 'qint8' - if "qint4" in name.lower(): - return 'qint4' - if "fp8" in name.lower(): - return 'fp8' - if "fp4" in name.lower(): - return 'fp4' - if "nf4" in name.lower(): - return 'nf4' - if name.endswith('.gguf'): - return 'gguf' - return 'none' +def apply_layerwise(sd_model, quiet:bool=False): + import torch + from diffusers.quantizers import quantization_config + from modules import shared, devices, sd_models + if shared.opts.layerwise_quantization_storage == 'float8_e4m3fn' and hasattr(torch, 'float8_e4m3fn'): + storage_dtype = torch.float8_e4m3fn + elif shared.opts.layerwise_quantization_storage == 'float8_e5m2' and hasattr(torch, 'float8_e5m2'): + storage_dtype = torch.float8_e5m2 + else: + storage_dtype = None + shared.log.warning(f'Quantization: type=layerwise storage={shared.opts.layerwise_quantization_storage} not supported') + return + non_blocking = False + if not hasattr(quantization_config.QuantizationMethod, 'LAYERWISE'): + setattr(quantization_config.QuantizationMethod, 'LAYERWISE', 'layerwise') # noqa: B010 + for module in sd_models.get_signature(sd_model).keys(): + if not hasattr(sd_model, module): + continue + try: + cls = getattr(sd_model, module).__class__.__name__ + if module.startswith('unet') and ('Model' in shared.opts.layerwise_quantization): + m = getattr(sd_model, module) + if hasattr(m, 'enable_layerwise_casting'): + m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking) + m.quantization_method = 'LayerWise' + log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}') + if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization or 'Transformer' in shared.opts.layerwise_quantization): + m = getattr(sd_model, module) + if hasattr(m, 'enable_layerwise_casting'): + m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking) + m.quantization_method = 'LayerWise' + log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}') + if module.startswith('text_encoder') and ('Model' in shared.opts.layerwise_quantization or 'Text Encoder' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()): + m = getattr(sd_model, module) + if hasattr(m, 'enable_layerwise_casting'): + m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking) + m.quantization_method = quantization_config.QuantizationMethod.LAYERWISE # pylint: disable=no-member + log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}') + except Exception as e: + shared.log.error(f'Quantization: type=layerwise {e}') diff --git a/modules/sd_models.py b/modules/sd_models.py index 6234c2b0b..0b98bb504 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -9,7 +9,7 @@ import diffusers import diffusers.loaders.single_file_utils import torch -from modules import paths, shared, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect +from modules import paths, shared, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect, model_quant from modules.timer import Timer, process as process_timer from modules.memstats import memory_stats from modules.modeldata import model_data @@ -570,9 +570,12 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No sd_model = sd_models_compile.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU if shared.opts.optimum_quanto_weights: sd_model = sd_models_compile.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU + if shared.opts.layerwise_quantization: + model_quant.apply_layerwise(sd_model) timer.record("options") set_diffuser_offload(sd_model, op) + if op == 'model' and not (os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface'): if getattr(shared.sd_model, 'sd_checkpoint_info', None) is not None and vae_file is not None: sd_vae.apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 8a490e110..aecb572ba 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -3,9 +3,9 @@ import sys import time import inspect import torch -import accelerate +import accelerate.hooks -from modules import shared, devices, errors +from modules import shared, devices, errors, model_quant from modules.timer import process as process_timer @@ -20,7 +20,6 @@ def get_signature(cls): def disable_offload(sd_model): - from accelerate.hooks import remove_hook_from_module if not getattr(sd_model, 'has_accelerate', False): return if hasattr(sd_model, "_internal_dict"): @@ -31,7 +30,7 @@ def disable_offload(sd_model): module = getattr(sd_model, module_name, None) if isinstance(module, torch.nn.Module): network_layer_name = getattr(module, "network_layer_name", None) - module = remove_hook_from_module(module, recurse=True) + module = accelerate.hooks.remove_hook_from_module(module, recurse=True) if network_layer_name: module.network_layer_name = network_layer_name sd_model.has_accelerate = False @@ -188,7 +187,7 @@ def apply_balanced_offload(sd_model, exclude=[]): checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else None if checkpoint_name is None: checkpoint_name = sd_model.__class__.__name__ - if offload_hook_instance is None or offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory or offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory or checkpoint_name != offload_hook_instance.checkpoint_name: + if (offload_hook_instance is None) or (offload_hook_instance.min_watermark != shared.opts.diffusers_offload_min_gpu_memory) or (offload_hook_instance.max_watermark != shared.opts.diffusers_offload_max_gpu_memory) or (checkpoint_name != offload_hook_instance.checkpoint_name): cached = False offload_hook_instance = OffloadHook(checkpoint_name) @@ -241,9 +240,11 @@ def apply_balanced_offload(sd_model, exclude=[]): if do_offload: module = module.to(devices.cpu, non_blocking=True) used_gpu -= module_size + cls = module.__class__.__name__ + quant = getattr(module, "quantization_method", None) if not cached: - shared.log.debug(f'Model module={module_name} type={module.__class__.__name__} dtype={module.dtype} quant={getattr(module, "quantization_method", None)} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}') - debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={getattr(module, "quantization_method", None)} module={module.__class__.__name__} size={module_size:.3f}') + shared.log.debug(f'Model module={module_name} type={cls} dtype={module.dtype} quant={quant} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}') + debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}') except Exception as e: if 'out of memory' in str(e): devices.torch_gc(fast=True, force=True, reason='oom') @@ -270,6 +271,8 @@ def apply_balanced_offload(sd_model, exclude=[]): apply_balanced_offload_to_module(sd_model.prior_pipe) if hasattr(sd_model, "decoder_pipe"): apply_balanced_offload_to_module(sd_model.decoder_pipe) + if shared.opts.layerwise_quantization: + model_quant.apply_layerwise(sd_model, quiet=True) # need to reapply since hooks were removed/readded set_accelerate(sd_model) t = time.time() - t0 process_timer.add('offload', t) diff --git a/modules/shared.py b/modules/shared.py index 2200813ea..7f406d950 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -564,14 +564,17 @@ options_templates.update(options_section(('quantization', "Quantization Settings "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ['nf4', 'fp8', 'fp4'], "visible": native}), "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}), + "optimum_quanto_sep": OptionInfo("

Optimum Quanto

", "", gr.HTML), "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}), "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}), + "torchao_sep": OptionInfo("

TorchAO

", "", gr.HTML), "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), "torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}), "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ['int4_weight_only', 'int8_dynamic_activation_int4_weight', 'int8_weight_only', 'int8_dynamic_activation_int8_weight', 'float8_weight_only', 'float8_dynamic_activation_float8_weight', 'float8_static_activation_float8_weight'], "visible": native}), + "nncf_sep": OptionInfo("

NNCF

", "", gr.HTML), "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), "nncf_compress_weights_mode": OptionInfo("INT8", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}), @@ -580,6 +583,11 @@ options_templates.update(options_section(('quantization', "Quantization Settings "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": cmd_opts.use_openvino}), "nncf_quant_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}), "quant_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), + + "layerwise_sep": OptionInfo("

Layerwise Casting

", "", gr.HTML), + "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "Text Encoder"], "visible": native}), + "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"], "visible": native}), + "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox, {"visible": native}), })) options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { From 9b579bfd96d82e5e9cfc9090ff14a67ee36ca4ef Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 23 Jan 2025 21:50:26 +0300 Subject: [PATCH 31/66] Move quant functions to model_quant.py --- modules/model_quant.py | 237 ++++++++++++++++++++++++++++-- modules/model_te.py | 4 +- modules/sd_models.py | 10 +- modules/sd_models_compile.py | 276 +---------------------------------- modules/sd_models_utils.py | 83 ++++++++++- 5 files changed, 318 insertions(+), 292 deletions(-) diff --git a/modules/model_quant.py b/modules/model_quant.py index 0a56805df..24c517dc3 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -1,11 +1,17 @@ import sys +import copy +import time import diffusers -from installer import install, log +from installer import install, log, setup_logging -bnb = None -quanto = None ao = None +bnb = None +intel_nncf = None +optimum_quanto = None + +quant_last_model_name = None +quant_last_model_device = None def get_quant(name): @@ -114,22 +120,42 @@ def load_bnb(msg='', silent=False): def load_quanto(msg='', silent=False): from modules import shared - global quanto # pylint: disable=global-statement - if quanto is not None: - return quanto + global optimum_quanto # pylint: disable=global-statement + if optimum_quanto is not None: + return optimum_quanto install('optimum-quanto==0.2.6', quiet=True) try: - from optimum import quanto as optimum_quanto # pylint: disable=no-name-in-module - quanto = optimum_quanto + from optimum import quanto # pylint: disable=no-name-in-module + optimum_quanto = quanto fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access log.debug(f'Quantization: type=quanto version={quanto.__version__} fn={fn}') # pylint: disable=protected-access if shared.opts.diffusers_offload_mode in {'balanced', 'sequential'}: shared.log.error(f'Quantization: type=quanto offload={shared.opts.diffusers_offload_mode} not supported') - return quanto + return optimum_quanto except Exception as e: if len(msg) > 0: log.error(f"{msg} failed to import optimum.quanto: {e}") - quanto = None + optimum_quanto = None + if not silent: + raise + return None + + +def load_nncf(msg='', silent=False): + global intel_nncf # pylint: disable=global-statement + if intel_nncf is not None: + return intel_nncf + install('nncf==2.7.0', quiet=True) + try: + import nncf + intel_nncf = nncf + fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access + log.debug(f'Quantization: type=nncf version={nncf.__version__} fn={fn}') # pylint: disable=protected-access + return intel_nncf + except Exception as e: + if len(msg) > 0: + log.error(f"{msg} failed to import nncf: {e}") + intel_nncf = None if not silent: raise return None @@ -175,3 +201,194 @@ def apply_layerwise(sd_model, quiet:bool=False): log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}') except Exception as e: shared.log.error(f'Quantization: type=layerwise {e}') + + +def nncf_send_to_device(model, device): + for child in model.children(): + if child.__class__.__name__ == "WeightsDecompressor": + child.scale = child.scale.to(device) + child.zero_point = child.zero_point.to(device) + nncf_send_to_device(child, device) + + +def nncf_compress_model(model, op=None, sd_model=None): + from modules import devices, shared + global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement + nncf = load_nncf('Quantize model: type=NNCF') + model.eval() + backup_embeddings = None + if hasattr(model, "get_input_embeddings"): + backup_embeddings = copy.deepcopy(model.get_input_embeddings()) + model = nncf.compress_weights(model) + nncf_send_to_device(model, devices.device) + if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: + model.set_input_embeddings(backup_embeddings) + if op is not None and shared.opts.quant_shuffle_weights: + if quant_last_model_name is not None: + if "." in quant_last_model_name: + last_model_names = quant_last_model_name.split(".") + getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) + else: + getattr(sd_model, quant_last_model_name).to(quant_last_model_device) + devices.torch_gc(force=True) + if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none": + quant_last_model_name = op + quant_last_model_device = model.device + else: + quant_last_model_name = None + quant_last_model_device = None + model.to(devices.device) + devices.torch_gc(force=True) + return model + + +def nncf_compress_weights(sd_model): + try: + t0 = time.time() + from modules import shared, devices, sd_models + shared.log.info(f"Quantization: type=NNCF modules={shared.opts.nncf_compress_weights}") + global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement + + sd_model = sd_models.apply_function_to_model(sd_model, nncf_compress_model, shared.opts.nncf_compress_weights, op="nncf") + if quant_last_model_name is not None: + if "." in quant_last_model_name: + last_model_names = quant_last_model_name.split(".") + getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) + else: + getattr(sd_model, quant_last_model_name).to(quant_last_model_device) + devices.torch_gc(force=True) + quant_last_model_name = None + quant_last_model_device = None + + t1 = time.time() + shared.log.info(f"Quantization: type=NNCF time={t1-t0:.2f}") + except Exception as e: + shared.log.warning(f"Quantization: type=NNCF {e}") + return sd_model + + +def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activations=None): + from modules import devices, shared + quanto = load_quanto('Quantize model: type=Optimum Quanto') + global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement + if sd_model is not None and "Flux" in sd_model.__class__.__name__: # LayerNorm is not supported + exclude_list = ["transformer_blocks.*.norm1.norm", "transformer_blocks.*.norm2", "transformer_blocks.*.norm1_context.norm", "transformer_blocks.*.norm2_context", "single_transformer_blocks.*.norm.norm", "norm_out.norm"] + else: + exclude_list = None + weights = getattr(quanto, weights) if weights is not None else getattr(quanto, shared.opts.optimum_quanto_weights_type) + if activations is not None: + activations = getattr(quanto, activations) if activations != 'none' else None + elif shared.opts.optimum_quanto_activations_type != 'none': + activations = getattr(quanto, shared.opts.optimum_quanto_activations_type) + else: + activations = None + model.eval() + backup_embeddings = None + if hasattr(model, "get_input_embeddings"): + backup_embeddings = copy.deepcopy(model.get_input_embeddings()) + quanto.quantize(model, weights=weights, activations=activations, exclude=exclude_list) + quanto.freeze(model) + if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: + model.set_input_embeddings(backup_embeddings) + if op is not None and shared.opts.quant_shuffle_weights: + if quant_last_model_name is not None: + if "." in quant_last_model_name: + last_model_names = quant_last_model_name.split(".") + getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) + else: + getattr(sd_model, quant_last_model_name).to(quant_last_model_device) + devices.torch_gc(force=True) + if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none": + quant_last_model_name = op + quant_last_model_device = model.device + else: + quant_last_model_name = None + quant_last_model_device = None + model.to(devices.device) + devices.torch_gc(force=True) + return model + + +def optimum_quanto_weights(sd_model): + try: + t0 = time.time() + from modules import shared, devices, sd_models + if shared.opts.diffusers_offload_mode in {"balanced", "sequential"}: + shared.log.warning(f"Quantization: type=Optimum.quanto offload={shared.opts.diffusers_offload_mode} not compatible") + return sd_model + shared.log.info(f"Quantization: type=Optimum.quanto: modules={shared.opts.optimum_quanto_weights}") + global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement + quanto = load_quanto() + quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs) + + sd_model = sd_models.apply_function_to_model(sd_model, optimum_quanto_model, shared.opts.optimum_quanto_weights, op="optimum-quanto") + if quant_last_model_name is not None: + if "." in quant_last_model_name: + last_model_names = quant_last_model_name.split(".") + getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) + else: + getattr(sd_model, quant_last_model_name).to(quant_last_model_device) + devices.torch_gc(force=True) + quant_last_model_name = None + quant_last_model_device = None + + if shared.opts.optimum_quanto_activations_type != 'none': + activations = getattr(quanto, shared.opts.optimum_quanto_activations_type) + else: + activations = None + + if activations is not None: + def optimum_quanto_freeze(model, op=None, sd_model=None): # pylint: disable=unused-argument + quanto.freeze(model) + return model + if shared.opts.diffusers_offload_mode == "model": + sd_model.enable_model_cpu_offload(device=devices.device) + if hasattr(sd_model, "encode_prompt"): + original_encode_prompt = sd_model.encode_prompt + def encode_prompt(*args, **kwargs): + embeds = original_encode_prompt(*args, **kwargs) + sd_model.maybe_free_model_hooks() # Diffusers keeps the TE on VRAM + return embeds + sd_model.encode_prompt = encode_prompt + else: + sd_models.move_model(sd_model, devices.device) + with quanto.Calibration(momentum=0.9): + sd_model(prompt="dummy prompt", num_inference_steps=10) + sd_model = sd_models.apply_function_to_model(sd_model, optimum_quanto_freeze, shared.opts.optimum_quanto_weights, op="optimum-quanto-freeze") + if shared.opts.diffusers_offload_mode == "model": + sd_models.disable_offload(sd_model) + sd_models.move_model(sd_model, devices.cpu) + if hasattr(sd_model, "encode_prompt"): + sd_model.encode_prompt = original_encode_prompt + devices.torch_gc(force=True) + + t1 = time.time() + shared.log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}") + except Exception as e: + shared.log.warning(f"Quantization: type=Optimum.quanto {e}") + return sd_model + + +def torchao_quantization(sd_model): + from modules import shared, devices, sd_models + torchao = load_torchao() + q = torchao.quantization + + fn = getattr(q, shared.opts.torchao_quantization_type, None) + if fn is None: + shared.log.error(f"Quantization: type=TorchAO type={shared.opts.torchao_quantization_type} not supported") + return sd_model + def torchao_model(model, op=None, sd_model=None): # pylint: disable=unused-argument + q.quantize_(model, fn(), device=devices.device) + return model + + shared.log.info(f"Quantization: type=TorchAO pipe={sd_model.__class__.__name__} quant={shared.opts.torchao_quantization_type} fn={fn} targets={shared.opts.torchao_quantization}") + try: + t0 = time.time() + sd_models.apply_function_to_model(sd_model, torchao_model, shared.opts.torchao_quantization, op="torchao") + t1 = time.time() + shared.log.info(f"Quantization: type=TorchAO time={t1-t0:.2f}") + except Exception as e: + shared.log.error(f"Quantization: type=TorchAO {e}") + setup_logging() # torchao uses dynamo which messes with logging so reset is needed + return sd_model diff --git a/modules/model_te.py b/modules/model_te.py index 16bb6d222..024dda47c 100644 --- a/modules/model_te.py +++ b/modules/model_te.py @@ -64,12 +64,12 @@ def load_t5(name=None, cache_dir=None): t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) elif 'qint8' in name.lower(): model_quant.load_quanto('Load model: type=T5') - from modules.sd_models_compile import optimum_quanto_model + from modules.model_quant import optimum_quanto_model t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype) t5 = optimum_quanto_model(t5, weights="qint8", activations="none") elif 'int8' in name.lower(): install('nncf==2.7.0', quiet=True) - from modules.sd_models_compile import nncf_compress_model + from modules.model_quant import nncf_compress_model from modules.sd_hijack import NNCF_T5DenseGatedActDense t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype) for i in range(len(t5.encoder.block)): diff --git a/modules/sd_models.py b/modules/sd_models.py index 0b98bb504..df49b4918 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -16,7 +16,7 @@ from modules.modeldata import model_data from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import from modules.sd_offload import disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import from modules.sd_models_legacy import get_checkpoint_state_dict, load_model_weights, load_model, repair_config # pylint: disable=unused-import -from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, patch_diffuser_config, convert_to_faketensors, read_state_dict, get_state_dict_from_checkpoint # pylint: disable=unused-import +from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, patch_diffuser_config, convert_to_faketensors, read_state_dict, get_state_dict_from_checkpoint, apply_function_to_model # pylint: disable=unused-import model_dir = "Stable-diffusion" @@ -130,9 +130,9 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True, model.requires_grad_(False) model.eval() return model - sd_model = sd_models_compile.apply_compile_to_model(sd_model, eval_model, ["Model", "VAE", "Text Encoder"], op="eval") + sd_model = apply_function_to_model(sd_model, eval_model, ["Model", "VAE", "Text Encoder"], op="eval") if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode == 'post': - sd_model = sd_models_compile.torchao_quantization(sd_model) + sd_model = model_quant.torchao_quantization(sd_model) if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'): shared.log.quiet(quiet, f'Setting {op}: channels-last=True') @@ -567,9 +567,9 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No set_diffuser_options(sd_model, vae, op, offload=False) if shared.opts.nncf_compress_weights and not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): - sd_model = sd_models_compile.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU + sd_model = model_quant.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU if shared.opts.optimum_quanto_weights: - sd_model = sd_models_compile.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU + sd_model = model_quant.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU if shared.opts.layerwise_quantization: model_quant.apply_layerwise(sd_model) timer.record("options") diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 0972c3350..6754984e5 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -25,90 +25,9 @@ class CompiledModelState: self.partitioned_modules = {} -quant_last_model_name = None -quant_last_model_device = None deepcache_worker = None -def apply_compile_to_model(sd_model, function, options, op=None): - if "Model" in options: - if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'): - sd_model.unet = function(sd_model.unet, op="unet", sd_model=sd_model) - if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'): - sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model) - if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model, 'decoder'): - sd_model.decoder = None - sd_model.decoder = sd_model.decoder_pipe.decoder = function(sd_model.decoder_pipe.decoder, op="decoder_pipe.decoder", sd_model=sd_model) - if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'prior'): - if op == "nncf" and "StableCascade" in sd_model.__class__.__name__: # fixes dtype errors - backup_clip_txt_pooled_mapper = copy.deepcopy(sd_model.prior_pipe.prior.clip_txt_pooled_mapper) - sd_model.prior_pipe.prior = function(sd_model.prior_pipe.prior, op="prior_pipe.prior", sd_model=sd_model) - if op == "nncf" and "StableCascade" in sd_model.__class__.__name__: - sd_model.prior_pipe.prior.clip_txt_pooled_mapper = backup_clip_txt_pooled_mapper - if "Text Encoder" in options: - if hasattr(sd_model, 'text_encoder') and hasattr(sd_model.text_encoder, 'config'): - if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'text_encoder') and hasattr(sd_model.decoder_pipe.text_encoder, 'config'): - sd_model.decoder_pipe.text_encoder = function(sd_model.decoder_pipe.text_encoder, op="decoder_pipe.text_encoder", sd_model=sd_model) - else: - if op == "nncf" and sd_model.text_encoder.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: - from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 - for i in range(len(sd_model.text_encoder.encoder.block)): - sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( - sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense, - dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 - ) - sd_model.text_encoder = function(sd_model.text_encoder, op="text_encoder", sd_model=sd_model) - if hasattr(sd_model, 'text_encoder_2') and hasattr(sd_model.text_encoder_2, 'config'): - if op == "nncf" and sd_model.text_encoder_2.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: - from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 - for i in range(len(sd_model.text_encoder_2.encoder.block)): - sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( - sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense, - dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 - ) - sd_model.text_encoder_2 = function(sd_model.text_encoder_2, op="text_encoder_2", sd_model=sd_model) - if hasattr(sd_model, 'text_encoder_3') and hasattr(sd_model.text_encoder_3, 'config'): - if op == "nncf" and sd_model.text_encoder_3.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: - from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 - for i in range(len(sd_model.text_encoder_3.encoder.block)): - sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( - sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense, - dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 - ) - sd_model.text_encoder_3 = function(sd_model.text_encoder_3, op="text_encoder_3", sd_model=sd_model) - if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder') and hasattr(sd_model.prior_pipe.text_encoder, 'config'): - sd_model.prior_pipe.text_encoder = function(sd_model.prior_pipe.text_encoder, op="prior_pipe.text_encoder", sd_model=sd_model) - if "VAE" in options: - if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'): - if op == "compile": - sd_model.vae.decode = function(sd_model.vae.decode, op="vae_decode", sd_model=sd_model) - sd_model.vae.encode = function(sd_model.vae.encode, op="vae_encode", sd_model=sd_model) - else: - sd_model.vae = function(sd_model.vae, op="vae", sd_model=sd_model) - if hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'): - if op == "compile": - sd_model.movq.decode = function(sd_model.movq.decode, op="movq_decode", sd_model=sd_model) - sd_model.movq.encode = function(sd_model.movq.encode, op="movq_encode", sd_model=sd_model) - else: - sd_model.movq = function(sd_model.movq, op="movq", sd_model=sd_model) - if hasattr(sd_model, 'vqgan') and hasattr(sd_model.vqgan, 'decode'): - if op == "compile": - sd_model.vqgan.decode = function(sd_model.vqgan.decode, op="vqgan_decode", sd_model=sd_model) - sd_model.vqgan.encode = function(sd_model.vqgan.encode, op="vqgan_encode", sd_model=sd_model) - else: - sd_model.vqgan = function(sd_model.vqgan, op="vqgan", sd_model=sd_model) - if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'vqgan'): - if op == "compile": - sd_model.decoder_pipe.vqgan.decode = function(sd_model.decoder_pipe.vqgan.decode, op="vqgan_decode", sd_model=sd_model) - sd_model.decoder_pipe.vqgan.encode = function(sd_model.decoder_pipe.vqgan.encode, op="vqgan_encode", sd_model=sd_model) - else: - sd_model.decoder_pipe.vqgan = sd_model.vqgan - if hasattr(sd_model, 'image_encoder') and hasattr(sd_model.image_encoder, 'config'): - sd_model.image_encoder = function(sd_model.image_encoder, op="image_encoder", sd_model=sd_model) - - return sd_model - - def ipex_optimize(sd_model): try: t0 = time.time() @@ -133,7 +52,7 @@ def ipex_optimize(sd_model): devices.torch_gc() return model - sd_model = apply_compile_to_model(sd_model, ipex_optimize_model, shared.opts.ipex_optimize, op="ipex") + sd_model = sd_models.apply_function_to_model(sd_model, ipex_optimize_model, shared.opts.ipex_optimize, op="ipex") t1 = time.time() shared.log.info(f"IPEX Optimize: time={t1-t0:.2f}") @@ -142,169 +61,6 @@ def ipex_optimize(sd_model): return sd_model -def nncf_send_to_device(model): - for child in model.children(): - if child.__class__.__name__ == "WeightsDecompressor": - child.scale = child.scale.to(devices.device) - child.zero_point = child.zero_point.to(devices.device) - nncf_send_to_device(child) - - -def nncf_compress_model(model, op=None, sd_model=None): - import nncf - global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - model.eval() - backup_embeddings = None - if hasattr(model, "get_input_embeddings"): - backup_embeddings = copy.deepcopy(model.get_input_embeddings()) - model = nncf.compress_weights(model) - nncf_send_to_device(model) - if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: - model.set_input_embeddings(backup_embeddings) - if op is not None and shared.opts.quant_shuffle_weights: - if quant_last_model_name is not None: - if "." in quant_last_model_name: - last_model_names = quant_last_model_name.split(".") - getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) - else: - getattr(sd_model, quant_last_model_name).to(quant_last_model_device) - devices.torch_gc(force=True) - if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none": - quant_last_model_name = op - quant_last_model_device = model.device - else: - quant_last_model_name = None - quant_last_model_device = None - model.to(devices.device) - devices.torch_gc(force=True) - return model - - -def nncf_compress_weights(sd_model): - try: - t0 = time.time() - shared.log.info(f"Quantization: type=NNCF modules={shared.opts.nncf_compress_weights}") - global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - install('nncf==2.7.0', quiet=True) - - sd_model = apply_compile_to_model(sd_model, nncf_compress_model, shared.opts.nncf_compress_weights, op="nncf") - if quant_last_model_name is not None: - if "." in quant_last_model_name: - last_model_names = quant_last_model_name.split(".") - getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) - else: - getattr(sd_model, quant_last_model_name).to(quant_last_model_device) - devices.torch_gc(force=True) - quant_last_model_name = None - quant_last_model_device = None - - t1 = time.time() - shared.log.info(f"Quantization: type=NNCF time={t1-t0:.2f}") - except Exception as e: - shared.log.warning(f"Quantization: type=NNCF {e}") - return sd_model - - -def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activations=None): - quanto = model_quant.load_quanto('Compile model: type=Optimum Quanto') - global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - if sd_model is not None and "Flux" in sd_model.__class__.__name__: # LayerNorm is not supported - exclude_list = ["transformer_blocks.*.norm1.norm", "transformer_blocks.*.norm2", "transformer_blocks.*.norm1_context.norm", "transformer_blocks.*.norm2_context", "single_transformer_blocks.*.norm.norm", "norm_out.norm"] - else: - exclude_list = None - weights = getattr(quanto, weights) if weights is not None else getattr(quanto, shared.opts.optimum_quanto_weights_type) - if activations is not None: - activations = getattr(quanto, activations) if activations != 'none' else None - elif shared.opts.optimum_quanto_activations_type != 'none': - activations = getattr(quanto, shared.opts.optimum_quanto_activations_type) - else: - activations = None - model.eval() - backup_embeddings = None - if hasattr(model, "get_input_embeddings"): - backup_embeddings = copy.deepcopy(model.get_input_embeddings()) - quanto.quantize(model, weights=weights, activations=activations, exclude=exclude_list) - quanto.freeze(model) - if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: - model.set_input_embeddings(backup_embeddings) - if op is not None and shared.opts.quant_shuffle_weights: - if quant_last_model_name is not None: - if "." in quant_last_model_name: - last_model_names = quant_last_model_name.split(".") - getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) - else: - getattr(sd_model, quant_last_model_name).to(quant_last_model_device) - devices.torch_gc(force=True) - if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none": - quant_last_model_name = op - quant_last_model_device = model.device - else: - quant_last_model_name = None - quant_last_model_device = None - model.to(devices.device) - devices.torch_gc(force=True) - return model - - -def optimum_quanto_weights(sd_model): - try: - if shared.opts.diffusers_offload_mode in {"balanced", "sequential"}: - shared.log.warning(f"Quantization: type=Optimum.quanto offload={shared.opts.diffusers_offload_mode} not compatible") - return sd_model - t0 = time.time() - shared.log.info(f"Quantization: type=Optimum.quanto: modules={shared.opts.optimum_quanto_weights}") - global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - quanto = model_quant.load_quanto() - quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs) - - sd_model = apply_compile_to_model(sd_model, optimum_quanto_model, shared.opts.optimum_quanto_weights, op="optimum-quanto") - if quant_last_model_name is not None: - if "." in quant_last_model_name: - last_model_names = quant_last_model_name.split(".") - getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device) - else: - getattr(sd_model, quant_last_model_name).to(quant_last_model_device) - devices.torch_gc(force=True) - quant_last_model_name = None - quant_last_model_device = None - - if shared.opts.optimum_quanto_activations_type != 'none': - activations = getattr(quanto, shared.opts.optimum_quanto_activations_type) - else: - activations = None - - if activations is not None: - def optimum_quanto_freeze(model, op=None, sd_model=None): # pylint: disable=unused-argument - quanto.freeze(model) - return model - if shared.opts.diffusers_offload_mode == "model": - sd_model.enable_model_cpu_offload(device=devices.device) - if hasattr(sd_model, "encode_prompt"): - original_encode_prompt = sd_model.encode_prompt - def encode_prompt(*args, **kwargs): - embeds = original_encode_prompt(*args, **kwargs) - sd_model.maybe_free_model_hooks() # Diffusers keeps the TE on VRAM - return embeds - sd_model.encode_prompt = encode_prompt - else: - sd_models.move_model(sd_model, devices.device) - with quanto.Calibration(momentum=0.9): - sd_model(prompt="dummy prompt", num_inference_steps=10) - sd_model = apply_compile_to_model(sd_model, optimum_quanto_freeze, shared.opts.optimum_quanto_weights, op="optimum-quanto-freeze") - if shared.opts.diffusers_offload_mode == "model": - sd_models.disable_offload(sd_model) - sd_models.move_model(sd_model, devices.cpu) - if hasattr(sd_model, "encode_prompt"): - sd_model.encode_prompt = original_encode_prompt - devices.torch_gc(force=True) - - t1 = time.time() - shared.log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}") - except Exception as e: - shared.log.warning(f"Quantization: type=Optimum.quanto {e}") - return sd_model - - def optimize_openvino(sd_model): try: from modules.intel.openvino import openvino_fx # pylint: disable=unused-import @@ -444,7 +200,7 @@ def compile_torch(sd_model): except Exception as e: shared.log.error(f"Model compile: torch inductor config error: {e}") - sd_model = apply_compile_to_model(sd_model, function=torch_compile_model, options=shared.opts.cuda_compile, op="compile") + sd_model = sd_models.apply_function_to_model(sd_model, function=torch_compile_model, options=shared.opts.cuda_compile, op="compile") setup_logging() # compile messes with logging so reset is needed if shared.opts.cuda_compile_precompile: @@ -503,34 +259,6 @@ def compile_diffusers(sd_model): return sd_model -def torchao_quantization(sd_model): - try: - install('torchao==0.7.0', quiet=True) - from torchao import quantization as q - except Exception as e: - shared.log.error(f"Quantization: type=TorchAO quantization not supported: {e}") - return sd_model - - fn = getattr(q, shared.opts.torchao_quantization_type, None) - if fn is None: - shared.log.error(f"Quantization: type=TorchAO type={shared.opts.torchao_quantization_type} not supported") - return sd_model - def torchao_model(model, op=None, sd_model=None): # pylint: disable=unused-argument - q.quantize_(model, fn(), device=devices.device) - return model - - shared.log.info(f"Quantization: type=TorchAO pipe={sd_model.__class__.__name__} quant={shared.opts.torchao_quantization_type} fn={fn} targets={shared.opts.torchao_quantization}") - try: - t0 = time.time() - apply_compile_to_model(sd_model, torchao_model, shared.opts.torchao_quantization, op="torchao") - t1 = time.time() - shared.log.info(f"Quantization: type=TorchAO time={t1-t0:.2f}") - except Exception as e: - shared.log.error(f"Quantization: type=TorchAO {e}") - setup_logging() # torchao uses dynamo which messes with logging so reset is needed - return sd_model - - def openvino_recompile_model(p, hires=False, refiner=False): # recompile if a parameter changes # pylint: disable=unused-argument if shared.opts.cuda_compile_backend == "openvino_fx" and 'Model' in shared.opts.cuda_compile: compile_height = p.height if not hires and hasattr(p, 'height') else p.hr_upscale_to_y diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py index 0ff903483..a09ca73a7 100644 --- a/modules/sd_models_utils.py +++ b/modules/sd_models_utils.py @@ -1,4 +1,5 @@ import io +import copy import json import inspect import os.path @@ -6,7 +7,7 @@ from rich import progress # pylint: disable=redefined-builtin import torch import safetensors.torch -from modules import paths, shared, errors +from modules import paths, shared, devices, errors from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import from modules.sd_offload import disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import from modules.sd_models_legacy import get_checkpoint_state_dict, load_model_weights, load_model, repair_config # pylint: disable=unused-import @@ -149,3 +150,83 @@ def patch_diffuser_config(sd_model, model_file): component.config[k] = v updated[k] = v return sd_model + + +def apply_function_to_model(sd_model, function, options, op=None): + if "Model" in options or "Transformer" in options: + if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'): + sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model) + if "Model" in options: + if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'): + sd_model.unet = function(sd_model.unet, op="unet", sd_model=sd_model) + if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model, 'decoder'): + sd_model.decoder = None + sd_model.decoder = sd_model.decoder_pipe.decoder = function(sd_model.decoder_pipe.decoder, op="decoder_pipe.decoder", sd_model=sd_model) + if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'prior'): + if op == "nncf" and "StableCascade" in sd_model.__class__.__name__: # fixes dtype errors + backup_clip_txt_pooled_mapper = copy.deepcopy(sd_model.prior_pipe.prior.clip_txt_pooled_mapper) + sd_model.prior_pipe.prior = function(sd_model.prior_pipe.prior, op="prior_pipe.prior", sd_model=sd_model) + if op == "nncf" and "StableCascade" in sd_model.__class__.__name__: + sd_model.prior_pipe.prior.clip_txt_pooled_mapper = backup_clip_txt_pooled_mapper + if "Text Encoder" in options: + if hasattr(sd_model, 'text_encoder') and hasattr(sd_model.text_encoder, 'config'): + if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'text_encoder') and hasattr(sd_model.decoder_pipe.text_encoder, 'config'): + sd_model.decoder_pipe.text_encoder = function(sd_model.decoder_pipe.text_encoder, op="decoder_pipe.text_encoder", sd_model=sd_model) + else: + if op == "nncf" and sd_model.text_encoder.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: + from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 + for i in range(len(sd_model.text_encoder.encoder.block)): + sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( + sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense, + dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 + ) + sd_model.text_encoder = function(sd_model.text_encoder, op="text_encoder", sd_model=sd_model) + if hasattr(sd_model, 'text_encoder_2') and hasattr(sd_model.text_encoder_2, 'config'): + if op == "nncf" and sd_model.text_encoder_2.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: + from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 + for i in range(len(sd_model.text_encoder_2.encoder.block)): + sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( + sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense, + dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 + ) + sd_model.text_encoder_2 = function(sd_model.text_encoder_2, op="text_encoder_2", sd_model=sd_model) + if hasattr(sd_model, 'text_encoder_3') and hasattr(sd_model.text_encoder_3, 'config'): + if op == "nncf" and sd_model.text_encoder_3.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: + from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 + for i in range(len(sd_model.text_encoder_3.encoder.block)): + sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( + sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense, + dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 + ) + sd_model.text_encoder_3 = function(sd_model.text_encoder_3, op="text_encoder_3", sd_model=sd_model) + if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder') and hasattr(sd_model.prior_pipe.text_encoder, 'config'): + sd_model.prior_pipe.text_encoder = function(sd_model.prior_pipe.text_encoder, op="prior_pipe.text_encoder", sd_model=sd_model) + if "VAE" in options: + if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'): + if op == "compile": + sd_model.vae.decode = function(sd_model.vae.decode, op="vae_decode", sd_model=sd_model) + sd_model.vae.encode = function(sd_model.vae.encode, op="vae_encode", sd_model=sd_model) + else: + sd_model.vae = function(sd_model.vae, op="vae", sd_model=sd_model) + if hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'): + if op == "compile": + sd_model.movq.decode = function(sd_model.movq.decode, op="movq_decode", sd_model=sd_model) + sd_model.movq.encode = function(sd_model.movq.encode, op="movq_encode", sd_model=sd_model) + else: + sd_model.movq = function(sd_model.movq, op="movq", sd_model=sd_model) + if hasattr(sd_model, 'vqgan') and hasattr(sd_model.vqgan, 'decode'): + if op == "compile": + sd_model.vqgan.decode = function(sd_model.vqgan.decode, op="vqgan_decode", sd_model=sd_model) + sd_model.vqgan.encode = function(sd_model.vqgan.encode, op="vqgan_encode", sd_model=sd_model) + else: + sd_model.vqgan = function(sd_model.vqgan, op="vqgan", sd_model=sd_model) + if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'vqgan'): + if op == "compile": + sd_model.decoder_pipe.vqgan.decode = function(sd_model.decoder_pipe.vqgan.decode, op="vqgan_decode", sd_model=sd_model) + sd_model.decoder_pipe.vqgan.encode = function(sd_model.decoder_pipe.vqgan.encode, op="vqgan_encode", sd_model=sd_model) + else: + sd_model.decoder_pipe.vqgan = sd_model.vqgan + if hasattr(sd_model, 'image_encoder') and hasattr(sd_model.image_encoder, 'config'): + sd_model.image_encoder = function(sd_model.image_encoder, op="image_encoder", sd_model=sd_model) + + return sd_model From 06ba03cf809956ccae8917c3a4fba30b5998a25b Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 23 Jan 2025 15:19:43 -0500 Subject: [PATCH 32/66] settings option to disable reference models Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 ++ modules/face/photomaker_pipeline.py | 10 +++++----- modules/intel/openvino/__init__.py | 8 ++++---- modules/merging/modules_sdxl.py | 14 +++++++------- modules/model_quant.py | 4 ++-- modules/perflow/pfode_solver.py | 6 +++--- modules/sd_models_compile.py | 5 ++--- modules/shared.py | 18 ++++++++++-------- modules/styles.py | 2 +- modules/ui_extra_networks.py | 4 ++-- modules/ui_extra_networks_checkpoints.py | 2 +- modules/upscaler_simple.py | 14 ++++++++------ 12 files changed, 47 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b1a291a2..f35f7b675 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,10 +48,12 @@ - **gallery**: add http fallback for slow/unreliable links - **splash**: add legacy mode indicator on splash screen - **network**: extract thumbnail from model metadata if present + - **network**: setting value to disable use of reference models - **Refactor**: - **upscale**: code refactor to unify latent, resize and model based upscalers - **loader**: ability to run in-memory models - **schedulers**: ability to create model-less schedulers + - **quantiation**: code refactor into dedicated module - **Fixes**: - non-full vae decode - send-to image transfer diff --git a/modules/face/photomaker_pipeline.py b/modules/face/photomaker_pipeline.py index ec5728961..657e96fc7 100644 --- a/modules/face/photomaker_pipeline.py +++ b/modules/face/photomaker_pipeline.py @@ -272,7 +272,7 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline): # textual inversion: process multi-vector tokens if necessary prompt_embeds_list = [] prompts = [prompt, prompt_2] - for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): + for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): # pylint: disable=redefined-argument-from-local if isinstance(self, TextualInversionLoaderMixin): prompt = self.maybe_convert_prompt(prompt, tokenizer) @@ -367,7 +367,7 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline): f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" f" {type(prompt)}." ) - elif batch_size != len(negative_prompt): + if batch_size != len(negative_prompt): raise ValueError( f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" @@ -377,7 +377,7 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline): uncond_tokens = [negative_prompt, negative_prompt_2] negative_prompt_embeds_list = [] - for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): + for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): # pylint: disable=redefined-argument-from-local if isinstance(self, TextualInversionLoaderMixin): negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer) @@ -735,8 +735,8 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline): ): discrete_timestep_cutoff = int( round( - self.scheduler.config.num_train_timesteps - - (self.denoising_end * self.scheduler.config.num_train_timesteps) + self.scheduler.config.num_train_timesteps # pylint: disable=no-member + - (self.denoising_end * self.scheduler.config.num_train_timesteps) # pylint: disable=no-member ) ) num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index a4000c43a..11fa0426f 100644 --- a/modules/intel/openvino/__init__.py +++ b/modules/intel/openvino/__init__.py @@ -221,10 +221,10 @@ def openvino_compile(gm: GraphModule, *example_inputs, model_hash_str: str = Non for idx, _ in enumerate(example_inputs): new_inputs.append(example_inputs[idx].detach().cpu().numpy()) new_inputs = [new_inputs] - if shared.opts.nncf_quant_mode == "INT8": + if shared.opts.nncf_quantize_mode == "INT8": om = nncf.quantize(om, nncf.Dataset(new_inputs)) else: - om = nncf.quantize(om, nncf.Dataset(new_inputs), mode=getattr(nncf.QuantizationMode, shared.opts.nncf_quant_mode), + om = nncf.quantize(om, nncf.Dataset(new_inputs), mode=getattr(nncf.QuantizationMode, shared.opts.nncf_quantize_mode), advanced_parameters=nncf.quantization.advanced_parameters.AdvancedQuantizationParameters( overflow_fix=nncf.quantization.advanced_parameters.OverflowFix.DISABLE, backend_params=None)) @@ -281,10 +281,10 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs): for idx, _ in enumerate(example_inputs): new_inputs.append(example_inputs[idx].detach().cpu().numpy()) new_inputs = [new_inputs] - if shared.opts.nncf_quant_mode == "INT8": + if shared.opts.nncf_quantize_mode == "INT8": om = nncf.quantize(om, nncf.Dataset(new_inputs)) else: - om = nncf.quantize(om, nncf.Dataset(new_inputs), mode=getattr(nncf.QuantizationMode, shared.opts.nncf_quant_mode), + om = nncf.quantize(om, nncf.Dataset(new_inputs), mode=getattr(nncf.QuantizationMode, shared.opts.nncf_quantize_mode), advanced_parameters=nncf.quantization.advanced_parameters.AdvancedQuantizationParameters( overflow_fix=nncf.quantization.advanced_parameters.OverflowFix.DISABLE, backend_params=None)) diff --git a/modules/merging/modules_sdxl.py b/modules/merging/modules_sdxl.py index 61f2cf948..ec994d567 100644 --- a/modules/merging/modules_sdxl.py +++ b/modules/merging/modules_sdxl.py @@ -67,13 +67,13 @@ def load_base(override:str=None): pipeline = diffusers.StableDiffusionXLPipeline.from_pretrained(fn, cache_dir=shared.opts.hfcache_dir, torch_dtype=recipe.dtype, add_watermarker=False) else: yield msg('base: not found') - return None + return pipeline.vae.register_to_config(force_upcast = False) def load_unet(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): if (recipe.unet is None or len(recipe.unet) == 0) and override is None: - return None + return fn = override or recipe.unet if not os.path.isabs(fn): fn = os.path.join(shared.opts.unet_dir, fn) @@ -92,7 +92,7 @@ def load_unet(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): def load_scheduler(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): if recipe.scheduler is None and override is None: - return None + return config = pipe.scheduler.config.__dict__ scheduler = override or recipe.scheduler yield msg(f'scheduler={scheduler}') @@ -107,7 +107,7 @@ def load_scheduler(pipe: diffusers.StableDiffusionXLPipeline, override:str=None) def load_vae(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): if (recipe.vae is None or len(recipe.vae) == 0)and override is None: - return None + return fn = override or recipe.vae if not os.path.isabs(fn): fn = os.path.join(shared.opts.vae_dir, fn) @@ -128,7 +128,7 @@ def load_vae(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): def load_te1(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): if (recipe.te1 is None or len(recipe.te1) == 0) and override is None: - return None + return config = pipe.text_encoder.config.__dict__ pretrained_config = transformers.PretrainedConfig.from_dict(config) fn = override or recipe.te1 @@ -149,7 +149,7 @@ def load_te1(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): def load_te2(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): if (recipe.te2 is None or len(recipe.te2) == 0) and override is None: - return None + return config = pipe.text_encoder_2.config.__dict__ pretrained_config = transformers.PretrainedConfig.from_dict(config) fn = override or recipe.te2 @@ -170,7 +170,7 @@ def load_te2(pipe: diffusers.StableDiffusionXLPipeline, override:str=None): def load_lora(pipe: diffusers.StableDiffusionXLPipeline, override: dict=None, fuse: float=None): if recipe.lora is None and override is None: - return None + return names = [] pipe.unfuse_lora() pipe.unload_lora_weights() diff --git a/modules/model_quant.py b/modules/model_quant.py index 24c517dc3..6798160a8 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -223,7 +223,7 @@ def nncf_compress_model(model, op=None, sd_model=None): nncf_send_to_device(model, devices.device) if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: model.set_input_embeddings(backup_embeddings) - if op is not None and shared.opts.quant_shuffle_weights: + if op is not None and shared.opts.nncf_quantize_shuffle_weights: if quant_last_model_name is not None: if "." in quant_last_model_name: last_model_names = quant_last_model_name.split(".") @@ -290,7 +290,7 @@ def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activation quanto.freeze(model) if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: model.set_input_embeddings(backup_embeddings) - if op is not None and shared.opts.quant_shuffle_weights: + if op is not None and shared.opts.optimum_quanto_shuffle_weights: if quant_last_model_name is not None: if "." in quant_last_model_name: last_model_names = quant_last_model_name.split(".") diff --git a/modules/perflow/pfode_solver.py b/modules/perflow/pfode_solver.py index db1b8e903..bffdd1494 100644 --- a/modules/perflow/pfode_solver.py +++ b/modules/perflow/pfode_solver.py @@ -49,7 +49,7 @@ class PFODESolver(): timestep_cond = None if unet.config.time_cond_proj_dim is not None: guidance_scale_tensor = torch.tensor(guidance_scale - 1).repeat(bsz) - timestep_cond = self.get_guidance_scale_embedding( + timestep_cond = self.get_guidance_scale_embedding( # pylint: disable=no-member guidance_scale_tensor, embedding_dim=unet.config.time_cond_proj_dim ).to(device=latents.device, dtype=latents.dtype) @@ -124,7 +124,7 @@ class PFODESolverSDXL(): self.t_terminal = t_terminal self.scheduler = scheduler - train_step_terminal = 0 + train_step_terminal = 0 train_step_initial = train_step_terminal + self.scheduler.config.num_train_timesteps # 0+1000 self.stepsize = (t_terminal-t_initial) / (train_step_terminal - train_step_initial) #1/1000 @@ -186,7 +186,7 @@ class PFODESolverSDXL(): timestep_cond = None if unet.config.time_cond_proj_dim is not None: guidance_scale_tensor = torch.tensor(guidance_scale - 1).repeat(bsz) - timestep_cond = self.get_guidance_scale_embedding( + timestep_cond = self.get_guidance_scale_embedding( # pylint: disable=no-member guidance_scale_tensor, embedding_dim=unet.config.time_cond_proj_dim ).to(device=latents.device, dtype=latents.dtype) diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 6754984e5..407eeebd8 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -1,9 +1,8 @@ -import copy import time import logging import torch -from modules import shared, devices, sd_models, model_quant -from installer import install, setup_logging +from modules import shared, devices, sd_models +from installer import setup_logging #Used by OpenVINO, can be used with TensorRT or Olive diff --git a/modules/shared.py b/modules/shared.py index 7f406d950..da9fa3f6d 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -560,7 +560,7 @@ options_templates.update(options_section(('backends', "Backend Settings"), { })) options_templates.update(options_section(('quantization', "Quantization Settings"), { - "bnb_sep": OptionInfo("

BitsAndBytes

", "", gr.HTML), + "bnb_quantization_sep": OptionInfo("

BitsAndBytes

", "", gr.HTML), "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ['nf4', 'fp8', 'fp4'], "visible": native}), "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}), @@ -569,22 +569,23 @@ options_templates.update(options_section(('quantization', "Quantization Settings "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}), "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}), + "optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), "torchao_sep": OptionInfo("

TorchAO

", "", gr.HTML), "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), "torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}), "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ['int4_weight_only', 'int8_dynamic_activation_int4_weight', 'int8_weight_only', 'int8_dynamic_activation_int8_weight', 'float8_weight_only', 'float8_dynamic_activation_float8_weight', 'float8_static_activation_float8_weight'], "visible": native}), - "nncf_sep": OptionInfo("

NNCF

", "", gr.HTML), + "nncf_compress_sep": OptionInfo("

NNCF

", "", gr.HTML), "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), "nncf_compress_weights_mode": OptionInfo("INT8", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}), "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 512, "step": 1, "visible": cmd_opts.use_openvino}), "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": cmd_opts.use_openvino}), - "nncf_quant_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}), - "quant_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), + "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}), + "nncf_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), - "layerwise_sep": OptionInfo("

Layerwise Casting

", "", gr.HTML), + "layerwise_quantization_sep": OptionInfo("

Layerwise Casting

", "", gr.HTML), "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "Text Encoder"], "visible": native}), "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"], "visible": native}), "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox, {"visible": native}), @@ -910,8 +911,9 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "extra_networks_fetch": OptionInfo(True, "UI fetch network info on mouse-over"), "extra_network_skip_indexing": OptionInfo(False, "Build info on first access", gr.Checkbox), - "extra_networks_model_sep": OptionInfo("

Models

", "", gr.HTML), - "extra_network_reference": OptionInfo(False, "Use reference values when available", gr.Checkbox), + "extra_networks_model_sep": OptionInfo("

Rerefence models

", "", gr.HTML), + "extra_network_reference_enable": OptionInfo(True, "Enable use of reference models", gr.Checkbox), + "extra_network_reference_values": OptionInfo(False, "Use reference values when available", gr.Checkbox), "extra_networks_lora_sep": OptionInfo("

LoRA

", "", gr.HTML), "extra_networks_default_multiplier": OptionInfo(1.0, "Default strength", gr.Slider, {"minimum": 0.0, "maximum": 2.0, "step": 0.01}), @@ -1189,7 +1191,7 @@ if not native: opts.data['diffusers_offload_mode'] = 'none' prompt_styles = modules.styles.StyleDatabase(opts) -reference_models = readfile(os.path.join('html', 'reference.json')) +reference_models = readfile(os.path.join('html', 'reference.json')) if opts.extra_network_reference_enable else {} cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or (cmd_opts.server_name or False)) and not cmd_opts.insecure devices.args = cmd_opts devices.opts = opts diff --git a/modules/styles.py b/modules/styles.py index 85ae190e2..1517236be 100644 --- a/modules/styles.py +++ b/modules/styles.py @@ -135,7 +135,7 @@ def apply_styles_to_extra(p, style: Style): 'size', ] reference_style = get_reference_style() - extra = infotext.parse(reference_style) if shared.opts.extra_network_reference else {} + extra = infotext.parse(reference_style) if shared.opts.extra_network_reference_values else {} style_extra = apply_wildcards_to_prompt(style.extra, [style.wildcards], silent=True) extra.update(infotext.parse(style_extra)) diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 769f47581..c9d860ce9 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -227,7 +227,7 @@ class ExtraNetworksPage: for parentdir, dirs in {d: files_cache.walk(d, cached=True, recurse=files_cache.not_hidden) for d in allowed_folders}.items(): for tgt in dirs: tgt = tgt.path - if os.path.join(paths.models_path, 'Reference') in tgt: + if os.path.join(paths.models_path, 'Reference') in tgt and shared.opts.extra_network_reference_enable: subdirs['Reference'] = 1 if shared.native and shared.opts.diffusers_dir in tgt: subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1 @@ -242,7 +242,7 @@ class ExtraNetworksPage: subdirs[subdir] = 1 debug(f"Networks: page='{self.name}' subfolders={list(subdirs)}") subdirs = OrderedDict(sorted(subdirs.items())) - if self.name == 'model': + if self.name == 'model' and shared.opts.extra_network_reference_enable: subdirs['Reference'] = 1 subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1 subdirs.move_to_end(os.path.basename(shared.opts.diffusers_dir)) diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py index 66dd21cb0..c08ef357d 100644 --- a/modules/ui_extra_networks_checkpoints.py +++ b/modules/ui_extra_networks_checkpoints.py @@ -15,7 +15,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): shared.refresh_checkpoints() def list_reference(self): # pylint: disable=inconsistent-return-statements - if not shared.opts.sd_checkpoint_autodownload: + if not shared.opts.sd_checkpoint_autodownload or not shared.opts.extra_network_reference_enable: return [] for k, v in shared.reference_models.items(): if not shared.native: diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py index bd0c6050f..95f2acac2 100644 --- a/modules/upscaler_simple.py +++ b/modules/upscaler_simple.py @@ -31,18 +31,20 @@ class UpscalerResize(Upscaler): def do_upscale(self, img: Image, selected_model=None): if selected_model is None: return img - if selected_model == "Resize Nearest": + elif selected_model == "Resize Nearest": return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.NEAREST) - if selected_model == "Resize Lanczos": + elif selected_model == "Resize Lanczos": return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.LANCZOS) - if selected_model == "Resize Bicubic": + elif selected_model == "Resize Bicubic": return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BICUBIC) - if selected_model == "Resize Bilinear": + elif selected_model == "Resize Bilinear": return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BILINEAR) - if selected_model == "Resize Hamming": + elif selected_model == "Resize Hamming": return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.HAMMING) - if selected_model == "Resize Box": + elif selected_model == "Resize Box": return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BOX) + else: + return img def load_model(self, _): From 8b787ec7213af34115102ac3e51e2104b1342502 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 24 Jan 2025 09:20:10 -0500 Subject: [PATCH 33/66] image metadata civitai compatibility Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 ++- modules/processing_info.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f35f7b675..924e55d14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-23 +## Update for 2025-01-24 - **Contributing**: - if you'd like to contribute, please see updated [contributing](https://github.com/vladmandic/automatic/blob/dev/CONTRIBUTING) guidelines @@ -66,6 +66,7 @@ - mark large models that should offload - avoid repeated optimum-quanto installation - avoid reinstalling bnb if not cuda + - image metadata civitai compatibility ## Update for 2025-01-15 diff --git a/modules/processing_info.py b/modules/processing_info.py index f3c8eea81..b5037fbf4 100644 --- a/modules/processing_info.py +++ b/modules/processing_info.py @@ -39,9 +39,9 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No ops = list(set(p.ops)) args = { # basic + "Steps": p.steps, "Size": f"{p.width}x{p.height}" if hasattr(p, 'width') and hasattr(p, 'height') else None, "Sampler": p.sampler_name if p.sampler_name != 'Default' else None, - "Steps": p.steps, "Seed": all_seeds[index], "Seed resize from": None if p.seed_resize_from_w == 0 or p.seed_resize_from_h == 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}", "CFG scale": p.cfg_scale if p.cfg_scale > 1.0 else None, @@ -179,7 +179,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No del args[k] debug(f'Infotext: args={args}') params_text = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in args.items()]) - negative_prompt_text = f"\nNegative prompt: {all_negative_prompts[index]}" if all_negative_prompts[index] else "" + negative_prompt_text = f"\nNegative prompt: {all_negative_prompts[index] if all_negative_prompts[index] else ''}" infotext = f"{all_prompts[index]}{negative_prompt_text}\n{params_text}".strip() debug(f'Infotext: "{infotext}"') return infotext From 68115410ff460a867bb43281f4022158f66349fa Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 24 Jan 2025 11:09:24 -0500 Subject: [PATCH 34/66] unified standar and modernui auth Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 4 +++ javascript/login.js | 72 +++++++++++++++++++++++++++++++++++++++ modules/api/middleware.py | 8 +++-- modules/ui_javascript.py | 21 +++++++++--- 4 files changed, 98 insertions(+), 7 deletions(-) create mode 100644 javascript/login.js diff --git a/CHANGELOG.md b/CHANGELOG.md index 924e55d14..a18801c09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,6 +54,10 @@ - **loader**: ability to run in-memory models - **schedulers**: ability to create model-less schedulers - **quantiation**: code refactor into dedicated module +- **Authentication**: + - perform auth check on ui startup + - unified standard and modern-ui authentication method + - cleanup auth logging - **Fixes**: - non-full vae decode - send-to image transfer diff --git a/javascript/login.js b/javascript/login.js new file mode 100644 index 000000000..41bb5e941 --- /dev/null +++ b/javascript/login.js @@ -0,0 +1,72 @@ +const loginCSS = ` + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + background: var(--background-fill-primary); + color: var(--body-text-color-subdued); + font-family: monospace; + z-index: 100; +`; + +const loginHTML = ` +
+

Login

+ + + + +
+ +
+`; + +function forceLogin() { + const form = document.createElement('form'); + form.method = 'POST'; + form.action = '/login'; + form.id = 'loginForm'; + form.style.cssText = loginCSS; + form.innerHTML = loginHTML; + document.body.appendChild(form); + const username = form.querySelector('#loginUsername'); + const password = form.querySelector('#loginPassword'); + const status = form.querySelector('#loginStatus'); + + form.addEventListener('submit', (event) => { + event.preventDefault(); + const formData = new FormData(form); + formData.append('username', username.value); + formData.append('password', password.value); + console.warn('login', formData); + fetch('/login', { + method: 'POST', + body: formData, + }) + .then(async (res) => { + const json = await res.json(); + const txt = `${res.status}: ${res.statusText} - ${json.detail}`; + status.textContent = txt; + console.log('login', txt); + if (res.status === 200) location.reload(); + }) + .catch((err) => { + status.textContent = err; + console.error('login', err); + }); + }); +} + +function loginCheck() { + fetch('/login_check', {}) + .then((res) => { + if (res.status === 200) console.log('login ok'); + else forceLogin(); + }) + .catch((err) => { + console.error('login', err); + }); +} + +window.onload = loginCheck; diff --git a/modules/api/middleware.py b/modules/api/middleware.py index 7eb2c40e8..b5f02bd60 100644 --- a/modules/api/middleware.py +++ b/modules/api/middleware.py @@ -45,7 +45,7 @@ def setup_middleware(app: FastAPI, cmd_opts): if (cmd_opts.api_log or cmd_opts.api_only) and endpoint.startswith('/sdapi'): if '/sdapi/v1/log' in endpoint or '/sdapi/v1/browser' in endpoint: return res - log.info('API {user} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string, logging-format-interpolation + log.info('API user={user} code={code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string, logging-format-interpolation user = app.tokens.get(token) if hasattr(app, 'tokens') else None, code = res.status_code, ver = req.scope.get('http_version', '0.0'), @@ -69,10 +69,14 @@ def setup_middleware(app: FastAPI, cmd_opts): "body": vars(e).get('body', ''), "errors": str(e), } + if err['code'] == 401 and 'file=' in req.url.path: # dont spam with unauth + return JSONResponse(status_code=err['code'], content=jsonable_encoder(err)) + log.error(f"API error: {req.method}: {req.url} {err}") + if not isinstance(e, HTTPException) and err['error'] != 'TypeError': # do not print backtrace on known httpexceptions errors.display(e, 'HTTP API', [anyio, fastapi, uvicorn, starlette]) - elif err['code'] == 404 or err['code'] == 401: + elif err['code'] in [404, 401, 400]: pass else: log.debug(e, exc_info=True) # print stack trace diff --git a/modules/ui_javascript.py b/modules/ui_javascript.py index ccf5f8c0d..c8847c7d8 100644 --- a/modules/ui_javascript.py +++ b/modules/ui_javascript.py @@ -17,12 +17,13 @@ def webpath(fn): def html_head(): head = '' main = ['script.js'] + skip = ['login.js'] for js in main: script_js = os.path.join(script_path, "javascript", js) head += f'\n' added = [] for script in modules.scripts.list_scripts("javascript", ".js"): - if script.filename in main: + if script.filename in main or script.filename in skip: continue head += f'\n' added.append(script.path) @@ -43,6 +44,14 @@ def html_body(): return body +def html_login(): + fn = os.path.join(script_path, "javascript", "login.js") + with open(fn, 'r', encoding='utf8') as f: + inline = f.read() + js = f'\n' + return js + + def html_css(css: str): def stylesheet(fn): return f'' @@ -78,17 +87,19 @@ def html_css(css: str): def reload_javascript(): base_css = theme.reload_gradio_theme() - head = html_head() - css = html_css(base_css) - body = html_body() title = 'SD.Next' manifest = f'' + login = html_login() + js = html_head() + css = html_css(base_css) + body = html_body() def template_response(*args, **kwargs): res = shared.GradioTemplateResponseOriginal(*args, **kwargs) res.body = res.body.replace(b'', f'{title}'.encode("utf8")) - res.body = res.body.replace(b'', f'{head}'.encode("utf8")) res.body = res.body.replace(b'', f'{manifest}'.encode("utf8")) + res.body = res.body.replace(b'', f'{login}'.encode("utf8")) + res.body = res.body.replace(b'', f'{js}'.encode("utf8")) res.body = res.body.replace(b'', f'{css}{body}'.encode("utf8")) lines = res.body.decode("utf8").split('\n') for line in lines: From 6d08bf25583e2a29d141e3ce8480afb28c72022e Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 25 Jan 2025 08:37:49 -0500 Subject: [PATCH 35/66] xyz grid handle invalid values Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + installer.py | 2 +- scripts/xyz_grid.py | 71 +++++++++++++++++++++++------------------ scripts/xyz_grid_on.py | 72 ++++++++++++++++++++++++------------------ 4 files changed, 85 insertions(+), 61 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a18801c09..30fe75551 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,6 +71,7 @@ - avoid repeated optimum-quanto installation - avoid reinstalling bnb if not cuda - image metadata civitai compatibility + - xyz grid handle invalid values ## Update for 2025-01-15 diff --git a/installer.py b/installer.py index f10738f13..72a6dd7f5 100644 --- a/installer.py +++ b/installer.py @@ -492,7 +492,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git: return - sha = '78bc824729f76a14ff2f211fc7f9a31e5500a41e' # diffusers commit hash + sha = '07860f991639f35f4b5a152676bd4d590c3e589e' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 109d808b8..240376b3c 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -181,29 +181,35 @@ class Script(scripts.Script): if opt.type == int: valslist_ext = [] for val in valslist: - m = re_range.fullmatch(val) - if m is not None: - start_val = int(m.group(1)) if m.group(1) is not None else val - end_val = int(m.group(2)) if m.group(2) is not None else val - num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) - valslist_ext += [int(x) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] - shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') - else: - valslist_ext.append(int(val)) + try: + m = re_range.fullmatch(val) + if m is not None: + start_val = int(m.group(1)) if m.group(1) is not None else val + end_val = int(m.group(2)) if m.group(2) is not None else val + num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) + valslist_ext += [int(x) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] + shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') + else: + valslist_ext.append(int(val)) + except Exception as e: + shared.log.error(f"XYZ grid: value={val} {e}") valslist.clear() valslist = [x for x in valslist_ext if x not in valslist] elif opt.type == float: valslist_ext = [] for val in valslist: - m = re_range.fullmatch(val) - if m is not None: - start_val = float(m.group(1)) if m.group(1) is not None else val - end_val = float(m.group(2)) if m.group(2) is not None else val - num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) - valslist_ext += [round(float(x), 2) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] - shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') - else: - valslist_ext.append(float(val)) + try: + m = re_range.fullmatch(val) + if m is not None: + start_val = float(m.group(1)) if m.group(1) is not None else val + end_val = float(m.group(2)) if m.group(2) is not None else val + num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) + valslist_ext += [round(float(x), 2) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] + shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') + else: + valslist_ext.append(float(val)) + except Exception as e: + shared.log.error(f"XYZ grid: value={val} {e}") valslist.clear() valslist = [x for x in valslist_ext if x not in valslist] elif opt.type == str_permutations: # pylint: disable=comparison-with-callable @@ -214,18 +220,23 @@ class Script(scripts.Script): opt.confirm(p, valslist) return valslist - x_opt = self.current_axis_options[x_type] - if x_opt.choices is not None and not csv_mode: - x_values = list_to_csv_string(x_values_dropdown) - xs = process_axis(x_opt, x_values, x_values_dropdown) - y_opt = self.current_axis_options[y_type] - if y_opt.choices is not None and not csv_mode: - y_values = list_to_csv_string(y_values_dropdown) - ys = process_axis(y_opt, y_values, y_values_dropdown) - z_opt = self.current_axis_options[z_type] - if z_opt.choices is not None and not csv_mode: - z_values = list_to_csv_string(z_values_dropdown) - zs = process_axis(z_opt, z_values, z_values_dropdown) + try: + x_opt = self.current_axis_options[x_type] + if x_opt.choices is not None and not csv_mode: + x_values = list_to_csv_string(x_values_dropdown) + xs = process_axis(x_opt, x_values, x_values_dropdown) + y_opt = self.current_axis_options[y_type] + if y_opt.choices is not None and not csv_mode: + y_values = list_to_csv_string(y_values_dropdown) + ys = process_axis(y_opt, y_values, y_values_dropdown) + z_opt = self.current_axis_options[z_type] + if z_opt.choices is not None and not csv_mode: + z_values = list_to_csv_string(z_values_dropdown) + zs = process_axis(z_opt, z_values, z_values_dropdown) + except Exception as e: + shared.log.error(f"XYZ grid: invalid axis values {e}") + return None + Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes def fix_axis_seeds(axis_opt, axis_list): diff --git a/scripts/xyz_grid_on.py b/scripts/xyz_grid_on.py index 1b0f79f43..f40a1731b 100644 --- a/scripts/xyz_grid_on.py +++ b/scripts/xyz_grid_on.py @@ -194,29 +194,35 @@ class Script(scripts.Script): if opt.type == int: valslist_ext = [] for val in valslist: - m = re_range.fullmatch(val) - if m is not None: - start_val = int(m.group(1)) if m.group(1) is not None else val - end_val = int(m.group(2)) if m.group(2) is not None else val - num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) - valslist_ext += [int(x) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] - shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') - else: - valslist_ext.append(int(val)) + try: + m = re_range.fullmatch(val) + if m is not None: + start_val = int(m.group(1)) if m.group(1) is not None else val + end_val = int(m.group(2)) if m.group(2) is not None else val + num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) + valslist_ext += [int(x) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] + shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') + else: + valslist_ext.append(int(val)) + except Exception as e: + shared.log.error(f"XYZ grid: value={val} {e}") valslist.clear() valslist = [x for x in valslist_ext if x not in valslist] elif opt.type == float: valslist_ext = [] for val in valslist: - m = re_range.fullmatch(val) - if m is not None: - start_val = float(m.group(1)) if m.group(1) is not None else val - end_val = float(m.group(2)) if m.group(2) is not None else val - num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) - valslist_ext += [round(float(x), 2) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] - shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') - else: - valslist_ext.append(float(val)) + try: + m = re_range.fullmatch(val) + if m is not None: + start_val = float(m.group(1)) if m.group(1) is not None else val + end_val = float(m.group(2)) if m.group(2) is not None else val + num = int(m.group(3)) if m.group(3) is not None else int(end_val-start_val) + valslist_ext += [round(float(x), 2) for x in np.linspace(start=start_val, stop=end_val, num=max(2, num)).tolist()] + shared.log.debug(f'XYZ grid range: start={start_val} end={end_val} num={max(2, num)} list={valslist}') + else: + valslist_ext.append(float(val)) + except Exception as e: + shared.log.error(f"XYZ grid: value={val} {e}") valslist.clear() valslist = [x for x in valslist_ext if x not in valslist] elif opt.type == str_permutations: # pylint: disable=comparison-with-callable @@ -227,18 +233,24 @@ class Script(scripts.Script): opt.confirm(p, valslist) return valslist - x_opt = self.current_axis_options[x_type] - if x_opt.choices is not None and not csv_mode: - x_values = list_to_csv_string(x_values_dropdown) - xs = process_axis(x_opt, x_values, x_values_dropdown) - y_opt = self.current_axis_options[y_type] - if y_opt.choices is not None and not csv_mode: - y_values = list_to_csv_string(y_values_dropdown) - ys = process_axis(y_opt, y_values, y_values_dropdown) - z_opt = self.current_axis_options[z_type] - if z_opt.choices is not None and not csv_mode: - z_values = list_to_csv_string(z_values_dropdown) - zs = process_axis(z_opt, z_values, z_values_dropdown) + try: + x_opt = self.current_axis_options[x_type] + if x_opt.choices is not None and not csv_mode: + x_values = list_to_csv_string(x_values_dropdown) + xs = process_axis(x_opt, x_values, x_values_dropdown) + y_opt = self.current_axis_options[y_type] + if y_opt.choices is not None and not csv_mode: + y_values = list_to_csv_string(y_values_dropdown) + ys = process_axis(y_opt, y_values, y_values_dropdown) + z_opt = self.current_axis_options[z_type] + if z_opt.choices is not None and not csv_mode: + z_values = list_to_csv_string(z_values_dropdown) + zs = process_axis(z_opt, z_values, z_values_dropdown) + except Exception as e: + shared.log.error(f"XYZ grid: invalid axis values {e}") + active = False + return None + Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes def fix_axis_seeds(axis_opt, axis_list): From 843011406ff01f9a3174a3750f310252b2e2b830 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 25 Jan 2025 11:41:34 -0500 Subject: [PATCH 36/66] obj trace Signed-off-by: Vladimir Mandic --- modules/memstats.py | 47 +++++++++++++++++++++++++++++++++++++++++++ modules/sd_models.py | 48 ++++++++++++++++++++++---------------------- 2 files changed, 71 insertions(+), 24 deletions(-) diff --git a/modules/memstats.py b/modules/memstats.py index d43e5bbfa..cb43f246c 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -1,3 +1,5 @@ +import re +import sys import os import psutil import torch @@ -56,3 +58,48 @@ def ram_stats(): return ram except Exception: return { 'used': 0, 'total': 0 } + + +class Object: + pattern = r"'(.*?)'" + + def __init__(self, name, obj): + self.id = id(obj) + self.name = name + self.fn = sys._getframe(2).f_code.co_name + self.size = sys.getsizeof(obj) + self.refcount = sys.getrefcount(obj) + if torch.is_tensor(obj): + self.type = obj.dtype + self.size = obj.element_size() * obj.nelement() + else: + self.type = re.findall(self.pattern, str(type(obj)))[0] + self.size = sys.getsizeof(obj) + def __str__(self): + return f'{self.fn}.{self.name} type={self.type} size={self.size} ref={self.refcount}' + + +def get_objects(gcl={}, threshold:int=0): + objects = [] + seen = [] + + for name, obj in gcl.items(): + if id(obj) in seen: + continue + seen.append(id(obj)) + if name == '__name__': + name = obj + elif name.startswith('__'): + continue + try: + o = Object(name, obj) + if o.size >= threshold: + objects.append(o) + except Exception: + pass + + objects = sorted(objects, key=lambda x: x.size, reverse=True) + for obj in objects: + shared.log.trace(obj) + + return objects diff --git a/modules/sd_models.py b/modules/sd_models.py index df49b4918..2de127347 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1050,27 +1050,27 @@ def unload_model_weights(op='model'): shared.compiled_model_state.compiled_cache.clear() shared.compiled_model_state.req_cache.clear() shared.compiled_model_state.partitioned_modules.clear() - if op == 'model' or op == 'dict': - if model_data.sd_model: - if not shared.native: - from modules import sd_hijack - move_model(model_data.sd_model, devices.cpu) - sd_hijack.model_hijack.undo_hijack(model_data.sd_model) - elif not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): - disable_offload(model_data.sd_model) - move_model(model_data.sd_model, 'meta') - model_data.sd_model = None - devices.torch_gc(force=True) - shared.log.debug(f'Unload weights {op}: {memory_stats()}') - elif op == 'refiner': - if model_data.sd_refiner: - if not shared.native: - from modules import sd_hijack - move_model(model_data.sd_refiner, devices.cpu) - sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner) - else: - disable_offload(model_data.sd_refiner) - move_model(model_data.sd_refiner, 'meta') - model_data.sd_refiner = None - devices.torch_gc(force=True) - shared.log.debug(f'Unload weights {op}: {memory_stats()}') + if (op == 'model' or op == 'dict') and model_data.sd_model: + shared.log.debug(f'Current {op}: {memory_stats()}') + if not shared.native: + from modules import sd_hijack + move_model(model_data.sd_model, devices.cpu) + sd_hijack.model_hijack.undo_hijack(model_data.sd_model) + elif not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): + disable_offload(model_data.sd_model) + move_model(model_data.sd_model, 'meta') + model_data.sd_model = None + devices.torch_gc(force=True) + shared.log.debug(f'Unload {op}: {memory_stats()} after') + elif (op == 'refiner') and model_data.sd_refiner: + shared.log.debug(f'Current {op}: {memory_stats()}') + if not shared.native: + from modules import sd_hijack + move_model(model_data.sd_refiner, devices.cpu) + sd_hijack.model_hijack.undo_hijack(model_data.sd_refiner) + else: + disable_offload(model_data.sd_refiner) + move_model(model_data.sd_refiner, 'meta') + model_data.sd_refiner = None + devices.torch_gc(force=True) + shared.log.debug(f'Unload {op}: {memory_stats()}') From 326e87a7e7b6c3171133ff565216ece17e48bc3d Mon Sep 17 00:00:00 2001 From: Sean Kaim Date: Sat, 25 Jan 2025 12:40:13 -0500 Subject: [PATCH 37/66] Fix Docker status and switch the default IPEX server to US * print out docker status properly * switch the default IPEX server to US --- installer.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/installer.py b/installer.py index 72a6dd7f5..576949b47 100644 --- a/installer.py +++ b/installer.py @@ -447,7 +447,7 @@ def get_platform(): 'system': platform.system(), 'release': release, 'python': platform.python_version(), - 'docker': os.environ.get('SD_INSTALL_DEBUG', None) is not None, + 'docker': os.environ.get('SD_DOCKER', None) is not None, # 'host': platform.node(), # 'version': platform.version(), } @@ -455,6 +455,7 @@ def get_platform(): return { 'error': e } +# TODO Requires pydantic 2.x before fully supporting python 3.12 # check python version def check_python(supported_minors=[9, 10, 11, 12], reason=None): t_start = time.time() @@ -537,6 +538,7 @@ def install_rocm_zluda(): log.info('Using CPU-only torch') return os.environ.get('TORCH_COMMAND', 'torch torchvision') + # TODO Requires pydantic 2.x before fully supporting python 3.12 check_python(supported_minors=[10, 11], reason='ROCm or ZLUDA backends require Python 3.10 or 3.11') log.info('ROCm: AMD toolkit detected') os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512') @@ -659,6 +661,9 @@ def install_rocm_zluda(): def install_ipex(torch_command): t_start = time.time() + # https://pytorch-extension.intel.com/installation?platform=gpu&version=v2.5.10%2Bxpu&os=linux%2Fwsl2&package=pip + # while IPEX supports 3.12, other components (such as pydantic 1.x) do not. Capping at 3.11 still. + # TODO Requires pydantic 2.x before fully supporting python 3.12 check_python(supported_minors=[9, 10, 11, 12], reason='IPEX backend requires Python 3.9, 3.10 or 3.11') args.use_ipex = True # pylint: disable=attribute-defined-outside-init log.info('IPEX: Intel OneAPI toolkit detected') @@ -681,7 +686,8 @@ def install_ipex(torch_command): os.environ.setdefault('IPEX_FORCE_ATTENTION_SLICE', '1') if "linux" in sys.platform: - torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cxx11.abi torchvision==0.20.1+cxx11.abi intel-extension-for-pytorch==2.5.10+xpu oneccl_bind_pt==2.5.0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/cn/') + # default to US server. If The China server is needed, change .../release-whl/stable/xpu/us/ to .../release-whl/stable/xpu/cn/ + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cxx11.abi torchvision==0.20.1+cxx11.abi intel-extension-for-pytorch==2.5.10+xpu oneccl_bind_pt==2.5.0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/') # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.2') else: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+xpu torchvision==0.21.0+xpu --index-url https://download.pytorch.org/whl/test/xpu') @@ -695,6 +701,7 @@ def install_ipex(torch_command): def install_openvino(torch_command): t_start = time.time() + # TODO Requires pydantic 2.x before fully supporting python 3.12 check_python(supported_minors=[9, 10, 11, 12], reason='OpenVINO backend requires Python 3.9, 3.10 or 3.11') log.info('OpenVINO: selected') if sys.platform == 'darwin': From a770b1c888261e6f0cd7a731a3deca3397caf47c Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 25 Jan 2025 21:33:42 +0300 Subject: [PATCH 38/66] More correct Dynamic Atten SDPA implementation and deprecate IPEX Diffusers attention --- installer.py | 10 +- modules/devices.py | 4 +- modules/intel/ipex/attention.py | 226 ++++++++------------- modules/intel/ipex/diffusers.py | 305 +---------------------------- modules/intel/ipex/hijacks.py | 21 +- modules/sd_hijack_dynamic_atten.py | 187 +++++++++++------- 6 files changed, 224 insertions(+), 529 deletions(-) diff --git a/installer.py b/installer.py index 576949b47..af8c00f89 100644 --- a/installer.py +++ b/installer.py @@ -455,7 +455,6 @@ def get_platform(): return { 'error': e } -# TODO Requires pydantic 2.x before fully supporting python 3.12 # check python version def check_python(supported_minors=[9, 10, 11, 12], reason=None): t_start = time.time() @@ -538,7 +537,6 @@ def install_rocm_zluda(): log.info('Using CPU-only torch') return os.environ.get('TORCH_COMMAND', 'torch torchvision') - # TODO Requires pydantic 2.x before fully supporting python 3.12 check_python(supported_minors=[10, 11], reason='ROCm or ZLUDA backends require Python 3.10 or 3.11') log.info('ROCm: AMD toolkit detected') os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512') @@ -661,9 +659,8 @@ def install_rocm_zluda(): def install_ipex(torch_command): t_start = time.time() - # https://pytorch-extension.intel.com/installation?platform=gpu&version=v2.5.10%2Bxpu&os=linux%2Fwsl2&package=pip - # while IPEX supports 3.12, other components (such as pydantic 1.x) do not. Capping at 3.11 still. - # TODO Requires pydantic 2.x before fully supporting python 3.12 + # Python 3.12 will cause compatibility issues with other dependencies + # IPEX supports Python 3.12 so don't block it but don't advertise it in the error message check_python(supported_minors=[9, 10, 11, 12], reason='IPEX backend requires Python 3.9, 3.10 or 3.11') args.use_ipex = True # pylint: disable=attribute-defined-outside-init log.info('IPEX: Intel OneAPI toolkit detected') @@ -701,7 +698,8 @@ def install_ipex(torch_command): def install_openvino(torch_command): t_start = time.time() - # TODO Requires pydantic 2.x before fully supporting python 3.12 + # Python 3.12 will cause compatibility issues with other dependencies. + # OpenVINO supports Python 3.12 so don't block it but don't advertise it in the error message check_python(supported_minors=[9, 10, 11, 12], reason='OpenVINO backend requires Python 3.9, 3.10 or 3.11') log.info('OpenVINO: selected') if sys.platform == 'darwin': diff --git a/modules/devices.py b/modules/devices.py index 2123c05ed..874b24de8 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -421,8 +421,8 @@ def set_sdpa_params(): try: global sdpa_pre_dyanmic_atten # pylint: disable=global-statement sdpa_pre_dyanmic_atten = torch.nn.functional.scaled_dot_product_attention - from modules.sd_hijack_dynamic_atten import sliced_scaled_dot_product_attention - torch.nn.functional.scaled_dot_product_attention = sliced_scaled_dot_product_attention + from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention + torch.nn.functional.scaled_dot_product_attention = dynamic_scaled_dot_product_attention log.debug('SDPA Dynamic Attention Hijacked') except Exception as err: log.error(f'SDPA Dynamic Attention failed: {err}') diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py index 1618045b6..eacb6a7d0 100644 --- a/modules/intel/ipex/attention.py +++ b/modules/intel/ipex/attention.py @@ -1,6 +1,7 @@ import os +import math import torch -from functools import cache +from functools import cache, wraps # pylint: disable=protected-access, missing-function-docstring, line-too-long @@ -11,170 +12,117 @@ attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4)) # Find something divisible with the input_tokens @cache -def find_slice_size(slice_size, slice_block_size): - while (slice_size * slice_block_size) > attention_slice_rate: - slice_size = slice_size // 2 - if slice_size <= 1: - slice_size = 1 +def find_split_size(split_size, slice_block_size): + while (split_size * slice_block_size) > attention_slice_rate: + split_size = split_size // 2 + if split_size <= 1: + split_size = 1 break - return slice_size + return split_size + + +@cache +def find_query_size(query_size, slice_query_size): + while (math.sqrt(query_size) * slice_query_size) > attention_slice_rate: + query_size = query_size // 2 + if query_size <= 1: + query_size = 1 + break + return query_size + # Find slice sizes for SDPA @cache -def find_sdpa_slice_sizes(query_shape, query_element_size): - if len(query_shape) == 3: - batch_size_attention, query_tokens, shape_three = query_shape - shape_four = 1 - else: - batch_size_attention, query_tokens, shape_three, shape_four = query_shape +def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size): + batch_size, attn_heads, query_len, _ = query_shape + _, _, key_len, _ = key_shape + _, _, _, head_dim = value_shape - slice_block_size = query_tokens * shape_three * shape_four / 1024 / 1024 * query_element_size - block_size = batch_size_attention * slice_block_size + slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 - split_slice_size = batch_size_attention - split_2_slice_size = query_tokens - split_3_slice_size = shape_three + split_batch_size = batch_size + split_head_size = attn_heads + split_query_size = query_len - do_split = False - do_split_2 = False - do_split_3 = False + do_batch_split = False + do_head_split = False + do_query_split = False - if block_size > sdpa_slice_trigger_rate: - do_split = True - split_slice_size = find_slice_size(split_slice_size, slice_block_size) - if split_slice_size * slice_block_size > attention_slice_rate: - slice_2_block_size = split_slice_size * shape_three * shape_four / 1024 / 1024 * query_element_size - do_split_2 = True - split_2_slice_size = find_slice_size(split_2_slice_size, slice_2_block_size) - if split_2_slice_size * slice_2_block_size > attention_slice_rate: - slice_3_block_size = split_slice_size * split_2_slice_size * shape_four / 1024 / 1024 * query_element_size - do_split_3 = True - split_3_slice_size = find_slice_size(split_3_slice_size, slice_3_block_size) + if batch_size * slice_batch_size > sdpa_slice_trigger_rate: + do_batch_split = True + split_batch_size = find_split_size(split_batch_size, slice_batch_size) - return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size + if split_batch_size * slice_batch_size > attention_slice_rate: + slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + do_head_split = True + split_head_size = find_split_size(split_head_size, slice_head_size) -# Find slice sizes for BMM -@cache -def find_bmm_slice_sizes(input_shape, input_element_size, mat2_shape): - batch_size_attention, input_tokens, mat2_atten_shape = input_shape[0], input_shape[1], mat2_shape[2] - slice_block_size = input_tokens * mat2_atten_shape / 1024 / 1024 * input_element_size - block_size = batch_size_attention * slice_block_size + if split_batch_size * slice_batch_size > attention_slice_rate: + slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + do_query_split = True + split_query_size = find_query_size(split_query_size, slice_query_size) - split_slice_size = batch_size_attention - split_2_slice_size = input_tokens - split_3_slice_size = mat2_atten_shape + return do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size - do_split = False - do_split_2 = False - do_split_3 = False - - if block_size > attention_slice_rate: - do_split = True - split_slice_size = find_slice_size(split_slice_size, slice_block_size) - if split_slice_size * slice_block_size > attention_slice_rate: - slice_2_block_size = split_slice_size * mat2_atten_shape / 1024 / 1024 * input_element_size - do_split_2 = True - split_2_slice_size = find_slice_size(split_2_slice_size, slice_2_block_size) - if split_2_slice_size * slice_2_block_size > attention_slice_rate: - slice_3_block_size = split_slice_size * split_2_slice_size / 1024 / 1024 * input_element_size - do_split_3 = True - split_3_slice_size = find_slice_size(split_3_slice_size, slice_3_block_size) - - return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size - - -original_torch_bmm = torch.bmm -def torch_bmm_32_bit(input, mat2, *, out=None): - if input.device.type != "xpu": - return original_torch_bmm(input, mat2, out=out) - do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_bmm_slice_sizes(input.shape, input.element_size(), mat2.shape) - - # Slice BMM - if do_split: - batch_size_attention, input_tokens, mat2_atten_shape = input.shape[0], input.shape[1], mat2.shape[2] - hidden_states = torch.zeros(input.shape[0], input.shape[1], mat2.shape[2], device=input.device, dtype=input.dtype) - for i in range(batch_size_attention // split_slice_size): - start_idx = i * split_slice_size - end_idx = (i + 1) * split_slice_size - if do_split_2: - for i2 in range(input_tokens // split_2_slice_size): # pylint: disable=invalid-name - start_idx_2 = i2 * split_2_slice_size - end_idx_2 = (i2 + 1) * split_2_slice_size - if do_split_3: - for i3 in range(mat2_atten_shape // split_3_slice_size): # pylint: disable=invalid-name - start_idx_3 = i3 * split_3_slice_size - end_idx_3 = (i3 + 1) * split_3_slice_size - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = original_torch_bmm( - input[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - mat2[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - out=out - ) - else: - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = original_torch_bmm( - input[start_idx:end_idx, start_idx_2:end_idx_2], - mat2[start_idx:end_idx, start_idx_2:end_idx_2], - out=out - ) - else: - hidden_states[start_idx:end_idx] = original_torch_bmm( - input[start_idx:end_idx], - mat2[start_idx:end_idx], - out=out - ) - torch.xpu.synchronize(input.device) - else: - return original_torch_bmm(input, mat2, out=out) - return hidden_states original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention -def scaled_dot_product_attention_32_bit(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs): +@wraps(torch.nn.functional.scaled_dot_product_attention) +def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs): if query.device.type != "xpu": return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) - do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_sdpa_slice_sizes(query.shape, query.element_size()) + is_unsqueezed = False + if len(query.shape) == 3: + query = query.unsqueeze(0) + is_unsqueezed = True + if len(key.shape) == 3: + key = key.unsqueeze(0) + if len(value.shape) == 3: + value = value.unsqueeze(0) + do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size()) # Slice SDPA - if do_split: - batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2] - hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype) - if attn_mask is not None and attn_mask.shape[:-1] != query.shape[:-1]: - if len(query.shape) == 4: - attn_mask = attn_mask.expand((query.shape[0], query.shape[1], query.shape[2], key.shape[-2])) - else: - attn_mask = attn_mask.expand((query.shape[0], query.shape[1], key.shape[-2])) - for i in range(batch_size_attention // split_slice_size): - start_idx = i * split_slice_size - end_idx = (i + 1) * split_slice_size - if do_split_2: - for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name - start_idx_2 = i2 * split_2_slice_size - end_idx_2 = (i2 + 1) * split_2_slice_size - if do_split_3: - for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name - start_idx_3 = i3 * split_3_slice_size - end_idx_3 = (i3 + 1) * split_3_slice_size - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = original_scaled_dot_product_attention( - query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - attn_mask=attn_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attn_mask is not None else attn_mask, + if do_batch_split: + batch_size, attn_heads, query_len, _ = query.shape + _, _, _, head_dim = value.shape + hidden_states = torch.zeros((batch_size, attn_heads, query_len, head_dim), device=query.device, dtype=query.dtype) + if attn_mask is not None: + attn_mask = attn_mask.expand((query.shape[0], query.shape[1], query.shape[2], key.shape[-2])) + for ib in range(batch_size // split_batch_size): + start_idx = ib * split_batch_size + end_idx = (ib + 1) * split_batch_size + if do_head_split: + for ih in range(attn_heads // split_head_size): # pylint: disable=invalid-name + start_idx_h = ih * split_head_size + end_idx_h = (ih + 1) * split_head_size + if do_query_split: + for iq in range(query_len // split_query_size): # pylint: disable=invalid-name + start_idx_q = iq * split_query_size + end_idx_q = (iq + 1) * split_query_size + hidden_states[start_idx:end_idx, start_idx_h:end_idx_h, start_idx_q:end_idx_q, :] = original_scaled_dot_product_attention( + query[start_idx:end_idx, start_idx_h:end_idx_h, start_idx_q:end_idx_q, :], + key[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + value[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + attn_mask=attn_mask[start_idx:end_idx, start_idx_h:end_idx_h, start_idx_q:end_idx_q, :] if attn_mask is not None else attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs ) else: - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = original_scaled_dot_product_attention( - query[start_idx:end_idx, start_idx_2:end_idx_2], - key[start_idx:end_idx, start_idx_2:end_idx_2], - value[start_idx:end_idx, start_idx_2:end_idx_2], - attn_mask=attn_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attn_mask is not None else attn_mask, + hidden_states[start_idx:end_idx, start_idx_h:end_idx_h, :, :] = original_scaled_dot_product_attention( + query[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + key[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + value[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + attn_mask=attn_mask[start_idx:end_idx, start_idx_h:end_idx_h, :, :] if attn_mask is not None else attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs ) else: - hidden_states[start_idx:end_idx] = original_scaled_dot_product_attention( - query[start_idx:end_idx], - key[start_idx:end_idx], - value[start_idx:end_idx], - attn_mask=attn_mask[start_idx:end_idx] if attn_mask is not None else attn_mask, + hidden_states[start_idx:end_idx, :, :, :] = original_scaled_dot_product_attention( + query[start_idx:end_idx, :, :, :], + key[start_idx:end_idx, :, :, :], + value[start_idx:end_idx, :, :, :], + attn_mask=attn_mask[start_idx:end_idx, :, :, :] if attn_mask is not None else attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs ) + if is_unsqueezed: + hidden_states.squeeze(0) torch.xpu.synchronize(query.device) else: return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py index 734212237..413f7d5eb 100644 --- a/modules/intel/ipex/diffusers.py +++ b/modules/intel/ipex/diffusers.py @@ -1,13 +1,9 @@ -import os -from functools import wraps, cache +from functools import wraps import torch import diffusers # pylint: disable=import-error -from diffusers.models.attention_processor import Attention # pylint: disable=protected-access, missing-function-docstring, line-too-long -attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4)) - # Diffusers FreeU # Diffusers is imported before ipex hijacks so fourier_filter needs hijacking too @@ -46,306 +42,7 @@ class FluxPosEmbed(torch.nn.Module): return freqs_cos, freqs_sin -@cache -def find_slice_size(slice_size, slice_block_size): - while (slice_size * slice_block_size) > attention_slice_rate: - slice_size = slice_size // 2 - if slice_size <= 1: - slice_size = 1 - break - return slice_size - -@cache -def find_attention_slice_sizes(query_shape, query_element_size, query_device_type, slice_size=None): - if len(query_shape) == 3: - batch_size_attention, query_tokens, shape_three = query_shape - shape_four = 1 - else: - batch_size_attention, query_tokens, shape_three, shape_four = query_shape - if slice_size is not None: - batch_size_attention = slice_size - - slice_block_size = query_tokens * shape_three * shape_four / 1024 / 1024 * query_element_size - block_size = batch_size_attention * slice_block_size - - split_slice_size = batch_size_attention - split_2_slice_size = query_tokens - split_3_slice_size = shape_three - - do_split = False - do_split_2 = False - do_split_3 = False - - if query_device_type != "xpu": - return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size - - if block_size > attention_slice_rate: - do_split = True - split_slice_size = find_slice_size(split_slice_size, slice_block_size) - if split_slice_size * slice_block_size > attention_slice_rate: - slice_2_block_size = split_slice_size * shape_three * shape_four / 1024 / 1024 * query_element_size - do_split_2 = True - split_2_slice_size = find_slice_size(split_2_slice_size, slice_2_block_size) - if split_2_slice_size * slice_2_block_size > attention_slice_rate: - slice_3_block_size = split_slice_size * split_2_slice_size * shape_four / 1024 / 1024 * query_element_size - do_split_3 = True - split_3_slice_size = find_slice_size(split_3_slice_size, slice_3_block_size) - - return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size - - -class SlicedAttnProcessor: # pylint: disable=too-few-public-methods - r""" - Processor for implementing sliced attention. - - Args: - slice_size (`int`, *optional*): - The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and - `attention_head_dim` must be a multiple of the `slice_size`. - """ - - def __init__(self, slice_size): - self.slice_size = slice_size - - def __call__(self, attn: Attention, hidden_states: torch.Tensor, - encoder_hidden_states=None, attention_mask=None) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches - - residual = hidden_states - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - dim = query.shape[-1] - query = attn.head_to_batch_dim(query) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - batch_size_attention, query_tokens, shape_three = query.shape - hidden_states = torch.zeros( - (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype - ) - - #################################################################### - # ARC GPUs can't allocate more than 4GB to a single block, Slice it: - _, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_attention_slice_sizes(query.shape, query.element_size(), query.device.type, slice_size=self.slice_size) - - for i in range(batch_size_attention // split_slice_size): - start_idx = i * split_slice_size - end_idx = (i + 1) * split_slice_size - if do_split_2: - for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name - start_idx_2 = i2 * split_2_slice_size - end_idx_2 = (i2 + 1) * split_2_slice_size - if do_split_3: - for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name - start_idx_3 = i3 * split_3_slice_size - end_idx_3 = (i3 + 1) * split_3_slice_size - - query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] - key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] - attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - del query_slice - del key_slice - del attn_mask_slice - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]) - - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = attn_slice - del attn_slice - else: - query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2] - key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2] - attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - del query_slice - del key_slice - del attn_mask_slice - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2]) - - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice - del attn_slice - torch.xpu.synchronize(query.device) - else: - query_slice = query[start_idx:end_idx] - key_slice = key[start_idx:end_idx] - attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - del query_slice - del key_slice - del attn_mask_slice - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx]) - - hidden_states[start_idx:end_idx] = attn_slice - del attn_slice - #################################################################### - - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - -class AttnProcessor: - r""" - Default processor for performing attention-related computations. - """ - - def __call__(self, attn, hidden_states: torch.Tensor, encoder_hidden_states=None, attention_mask=None, - temb=None, *args, **kwargs) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches - - residual = hidden_states - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.head_to_batch_dim(query) - key = attn.head_to_batch_dim(key) - value = attn.head_to_batch_dim(value) - - #################################################################### - # ARC GPUs can't allocate more than 4GB to a single block, Slice it: - batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2] - hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype) - do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_attention_slice_sizes(query.shape, query.element_size(), query.device.type) - - if do_split: - for i in range(batch_size_attention // split_slice_size): - start_idx = i * split_slice_size - end_idx = (i + 1) * split_slice_size - if do_split_2: - for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name - start_idx_2 = i2 * split_2_slice_size - end_idx_2 = (i2 + 1) * split_2_slice_size - if do_split_3: - for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name - start_idx_3 = i3 * split_3_slice_size - end_idx_3 = (i3 + 1) * split_3_slice_size - - query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] - key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] - attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - del query_slice - del key_slice - del attn_mask_slice - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]) - - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = attn_slice - del attn_slice - else: - query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2] - key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2] - attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - del query_slice - del key_slice - del attn_mask_slice - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2]) - - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice - del attn_slice - else: - query_slice = query[start_idx:end_idx] - key_slice = key[start_idx:end_idx] - attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None - - attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) - del query_slice - del key_slice - del attn_mask_slice - attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx]) - - hidden_states[start_idx:end_idx] = attn_slice - del attn_slice - torch.xpu.synchronize(query.device) - else: - attention_probs = attn.get_attention_scores(query, key, attention_mask) - hidden_states = torch.bmm(attention_probs, value) - #################################################################### - hidden_states = attn.batch_to_head_dim(hidden_states) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False): diffusers.utils.torch_utils.fourier_filter = fourier_filter - #ARC GPUs can't allocate more than 4GB to a single block: - if not can_allocate_plus_4gb: - diffusers.models.attention_processor.SlicedAttnProcessor = SlicedAttnProcessor - diffusers.models.attention_processor.AttnProcessor = AttnProcessor if not device_supports_fp64: diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py index f1beb671d..4f3a03b4e 100644 --- a/modules/intel/ipex/hijacks.py +++ b/modules/intel/ipex/hijacks.py @@ -86,25 +86,14 @@ def as_tensor(data, dtype=None, device=None): if can_allocate_plus_4gb: - original_torch_bmm = torch.bmm original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention else: # 32 bit attention workarounds for Alchemist: try: - from .attention import torch_bmm_32_bit as original_torch_bmm - from .attention import scaled_dot_product_attention_32_bit as original_scaled_dot_product_attention + from .attention import dynamic_scaled_dot_product_attention as original_scaled_dot_product_attention except Exception: # pylint: disable=broad-exception-caught - original_torch_bmm = torch.bmm original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention - -# Data Type Errors: -@wraps(torch.bmm) -def torch_bmm(input, mat2, *, out=None): - if input.dtype != mat2.dtype: - mat2 = mat2.to(input.dtype) - return original_torch_bmm(input, mat2, out=out) - @wraps(torch.nn.functional.scaled_dot_product_attention) def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs): if query.dtype != key.dtype: @@ -115,6 +104,14 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0. attn_mask = attn_mask.to(dtype=query.dtype) return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) +# Data Type Errors: +original_torch_bmm = torch.bmm +@wraps(torch.bmm) +def torch_bmm(input, mat2, *, out=None): + if input.dtype != mat2.dtype: + mat2 = mat2.to(input.dtype) + return original_torch_bmm(input, mat2, out=out) + # Diffusers FreeU original_fft_fftn = torch.fft.fftn @wraps(torch.fft.fftn) diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index de39a966f..304f3f355 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -1,21 +1,131 @@ from functools import cache, wraps +import math import torch from diffusers.utils import USE_PEFT_BACKEND # pylint: disable=unused-import from modules import shared, devices @cache -def find_slice_size(slice_size, slice_block_size, slice_rate=4): - while (slice_size * slice_block_size) > slice_rate: - slice_size = slice_size // 2 - if slice_size <= 1: - slice_size = 1 +def find_split_size(split_size, slice_block_size, slice_rate=4): + while (split_size * slice_block_size) > slice_rate: + split_size = split_size // 2 + if split_size <= 1: + split_size = 1 break - return slice_size + return split_size + @cache -def find_slice_sizes(query_shape, query_element_size, slice_rate=4): +def find_query_size(query_size, slice_query_size, slice_rate=4): + while (math.sqrt(query_size) * slice_query_size) > slice_rate: + query_size = query_size // 2 + if query_size <= 1: + query_size = 1 + break + return query_size + + +# Find slice sizes for SDPA +@cache +def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size): + batch_size, attn_heads, query_len, _ = query_shape + _, _, key_len, _ = key_shape + _, _, _, head_dim = value_shape + + slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + + split_batch_size = batch_size + split_head_size = attn_heads + split_query_size = query_len + + do_batch_split = False + do_head_split = False + do_query_split = False + + if batch_size * slice_batch_size > shared.opts.dynamic_attention_slice_rate: + do_batch_split = True + split_batch_size = find_split_size(split_batch_size, slice_batch_size, slice_rate=shared.opts.dynamic_attention_slice_rate) + + if split_batch_size * slice_batch_size > shared.opts.dynamic_attention_slice_rate: + slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + do_head_split = True + split_head_size = find_split_size(split_head_size, slice_head_size, slice_rate=shared.opts.dynamic_attention_slice_rate) + + if split_batch_size * slice_batch_size > shared.opts.dynamic_attention_slice_rate: + slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + do_query_split = True + split_query_size = find_query_size(split_query_size, slice_query_size, slice_rate=shared.opts.dynamic_attention_slice_rate) + + return do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size + + +if devices.sdpa_pre_dyanmic_atten is None: + devices.sdpa_pre_dyanmic_atten = torch.nn.functional.scaled_dot_product_attention +@wraps(devices.sdpa_pre_dyanmic_atten) +def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs): + is_unsqueezed = False + if len(query.shape) == 3: + query = query.unsqueeze(0) + is_unsqueezed = True + if len(key.shape) == 3: + key = key.unsqueeze(0) + if len(value.shape) == 3: + value = value.unsqueeze(0) + do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size()) + + # Slice SDPA + if do_batch_split: + batch_size, attn_heads, query_len, _ = query.shape + _, _, _, head_dim = value.shape + hidden_states = torch.zeros((batch_size, attn_heads, query_len, head_dim), device=query.device, dtype=query.dtype) + if attn_mask is not None: + attn_mask = attn_mask.expand((query.shape[0], query.shape[1], query.shape[2], key.shape[-2])) + for ib in range(batch_size // split_batch_size): + start_idx = ib * split_batch_size + end_idx = (ib + 1) * split_batch_size + if do_head_split: + for ih in range(attn_heads // split_head_size): # pylint: disable=invalid-name + start_idx_h = ih * split_head_size + end_idx_h = (ih + 1) * split_head_size + if do_query_split: + for iq in range(query_len // split_query_size): # pylint: disable=invalid-name + start_idx_q = iq * split_query_size + end_idx_q = (iq + 1) * split_query_size + hidden_states[start_idx:end_idx, start_idx_h:end_idx_h, start_idx_q:end_idx_q, :] = devices.sdpa_pre_dyanmic_atten( + query[start_idx:end_idx, start_idx_h:end_idx_h, start_idx_q:end_idx_q, :], + key[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + value[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + attn_mask=attn_mask[start_idx:end_idx, start_idx_h:end_idx_h, start_idx_q:end_idx_q, :] if attn_mask is not None else attn_mask, + dropout_p=dropout_p, is_causal=is_causal, **kwargs + ) + else: + hidden_states[start_idx:end_idx, start_idx_h:end_idx_h, :, :] = devices.sdpa_pre_dyanmic_atten( + query[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + key[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + value[start_idx:end_idx, start_idx_h:end_idx_h, :, :], + attn_mask=attn_mask[start_idx:end_idx, start_idx_h:end_idx_h, :, :] if attn_mask is not None else attn_mask, + dropout_p=dropout_p, is_causal=is_causal, **kwargs + ) + else: + hidden_states[start_idx:end_idx, :, :, :] = devices.sdpa_pre_dyanmic_atten( + query[start_idx:end_idx, :, :, :], + key[start_idx:end_idx, :, :, :], + value[start_idx:end_idx, :, :, :], + attn_mask=attn_mask[start_idx:end_idx, :, :, :] if attn_mask is not None else attn_mask, + dropout_p=dropout_p, is_causal=is_causal, **kwargs + ) + if is_unsqueezed: + hidden_states.squeeze(0) + if devices.backend != "directml": + getattr(torch, query.device.type).synchronize() + else: + return devices.sdpa_pre_dyanmic_atten(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) + return hidden_states + + +@cache +def find_bmm_slice_sizes(query_shape, query_element_size, slice_rate=4): if len(query_shape) == 3: batch_size_attention, query_tokens, shape_three = query_shape shape_four = 1 @@ -35,73 +145,18 @@ def find_slice_sizes(query_shape, query_element_size, slice_rate=4): if block_size > slice_rate: do_split = True - split_slice_size = find_slice_size(split_slice_size, slice_block_size, slice_rate=slice_rate) + split_slice_size = find_split_size(split_slice_size, slice_block_size, slice_rate=slice_rate) if split_slice_size * slice_block_size > slice_rate: slice_2_block_size = split_slice_size * shape_three * shape_four / 1024 / 1024 * query_element_size do_split_2 = True - split_2_slice_size = find_slice_size(split_2_slice_size, slice_2_block_size, slice_rate=slice_rate) + split_2_slice_size = find_split_size(split_2_slice_size, slice_2_block_size, slice_rate=slice_rate) if split_2_slice_size * slice_2_block_size > slice_rate: slice_3_block_size = split_slice_size * split_2_slice_size * shape_four / 1024 / 1024 * query_element_size do_split_3 = True - split_3_slice_size = find_slice_size(split_3_slice_size, slice_3_block_size, slice_rate=slice_rate) + split_3_slice_size = find_split_size(split_3_slice_size, slice_3_block_size, slice_rate=slice_rate) return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size -if devices.sdpa_pre_dyanmic_atten is None: - devices.sdpa_pre_dyanmic_atten = torch.nn.functional.scaled_dot_product_attention -@wraps(devices.sdpa_pre_dyanmic_atten) -def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs): - do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate) - - # Slice SDPA - if do_split: - batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2] - hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype) - if attn_mask is not None and attn_mask.shape[:-1] != query.shape[:-1]: - if len(query.shape) == 4: - attn_mask = attn_mask.expand((query.shape[0], query.shape[1], query.shape[2], key.shape[-2])) - else: - attn_mask = attn_mask.expand((query.shape[0], query.shape[1], key.shape[-2])) - for i in range(batch_size_attention // split_slice_size): - start_idx = i * split_slice_size - end_idx = (i + 1) * split_slice_size - if do_split_2: - for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name - start_idx_2 = i2 * split_2_slice_size - end_idx_2 = (i2 + 1) * split_2_slice_size - if do_split_3: - for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name - start_idx_3 = i3 * split_3_slice_size - end_idx_3 = (i3 + 1) * split_3_slice_size - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = devices.sdpa_pre_dyanmic_atten( - query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3], - attn_mask=attn_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attn_mask is not None else attn_mask, - dropout_p=dropout_p, is_causal=is_causal, **kwargs - ) - else: - hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = devices.sdpa_pre_dyanmic_atten( - query[start_idx:end_idx, start_idx_2:end_idx_2], - key[start_idx:end_idx, start_idx_2:end_idx_2], - value[start_idx:end_idx, start_idx_2:end_idx_2], - attn_mask=attn_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attn_mask is not None else attn_mask, - dropout_p=dropout_p, is_causal=is_causal, **kwargs - ) - else: - hidden_states[start_idx:end_idx] = devices.sdpa_pre_dyanmic_atten( - query[start_idx:end_idx], - key[start_idx:end_idx], - value[start_idx:end_idx], - attn_mask=attn_mask[start_idx:end_idx] if attn_mask is not None else attn_mask, - dropout_p=dropout_p, is_causal=is_causal, **kwargs - ) - if devices.backend != "directml": - getattr(torch, query.device.type).synchronize() - else: - return devices.sdpa_pre_dyanmic_atten(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) - return hidden_states - class DynamicAttnProcessorBMM: r""" @@ -151,7 +206,7 @@ class DynamicAttnProcessorBMM: # Slicing parts: batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2] hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype) - do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate) + do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_bmm_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate) if do_split: for i in range(batch_size_attention // split_slice_size): From 58a3f7c39364e278181025a71f7009dc199a7c5d Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 25 Jan 2025 21:39:08 +0300 Subject: [PATCH 39/66] Fix OmniGen and update Changelog --- CHANGELOG.md | 5 +++++ modules/omnigen/pipeline.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 30fe75551..18e0de6ee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,8 +33,10 @@ higher values leads to more cache hits and speedups, but might also lead to a higher accuracy drop - **IPEX** - enable force attention slicing, fp64 emulation, jit cache + - use the us server by default on linux - use pytorch test branch on windows - extend the supported python versions + - removed diffusers attention hijack as it is a duplicate of dynamic attention bmm - **Torch FP8** - uses torch `float8_e4m3fn` or `float8_e5m2` as data storage and performs dynamic upcasting to compute `dtype` as needed - compatible with most `unet` and `transformer` models: e.g. *sd15, sdxl, sd35, flux.1, hunyuan-video, ltx-video, etc.* @@ -54,6 +56,7 @@ - **loader**: ability to run in-memory models - **schedulers**: ability to create model-less schedulers - **quantiation**: code refactor into dedicated module + - **dynamic attention sdpa**: more correct implementation - **Authentication**: - perform auth check on ui startup - unified standard and modern-ui authentication method @@ -72,6 +75,8 @@ - avoid reinstalling bnb if not cuda - image metadata civitai compatibility - xyz grid handle invalid values + - omnigen pipeline handle float seeds + - correct logging of docker status on logs, thanks @kmscode ## Update for 2025-01-15 diff --git a/modules/omnigen/pipeline.py b/modules/omnigen/pipeline.py index e17afea5c..a07467543 100644 --- a/modules/omnigen/pipeline.py +++ b/modules/omnigen/pipeline.py @@ -160,7 +160,7 @@ class OmniGenPipeline(): latent_size_h, latent_size_w = height//8, width//8 if seed is not None: - generator = torch.Generator(device=self.device).manual_seed(seed) + generator = torch.Generator(device=self.device).manual_seed(int(seed)) else: generator = None latents = torch.randn(num_prompt, 4, latent_size_h, latent_size_w, device=self.device, generator=generator) From 02ecc6710d387de432f2ce339c933585bbb76ca2 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 25 Jan 2025 21:43:38 +0300 Subject: [PATCH 40/66] Fix missing unsqueeze --- modules/intel/ipex/attention.py | 8 ++++---- modules/sd_hijack_dynamic_atten.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py index eacb6a7d0..4f74a2c2e 100644 --- a/modules/intel/ipex/attention.py +++ b/modules/intel/ipex/attention.py @@ -120,10 +120,10 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop value[start_idx:end_idx, :, :, :], attn_mask=attn_mask[start_idx:end_idx, :, :, :] if attn_mask is not None else attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs - ) - if is_unsqueezed: - hidden_states.squeeze(0) + ) torch.xpu.synchronize(query.device) else: - return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) + hidden_states = original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) + if is_unsqueezed: + hidden_states.squeeze(0) return hidden_states diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index 304f3f355..49f89b414 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -115,12 +115,12 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop attn_mask=attn_mask[start_idx:end_idx, :, :, :] if attn_mask is not None else attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs ) - if is_unsqueezed: - hidden_states.squeeze(0) if devices.backend != "directml": getattr(torch, query.device.type).synchronize() else: return devices.sdpa_pre_dyanmic_atten(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) + if is_unsqueezed: + hidden_states.squeeze(0) return hidden_states From d266143f204b5a76a84290da9e17fef7352751b6 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 25 Jan 2025 21:46:23 +0300 Subject: [PATCH 41/66] Cleanup --- modules/intel/ipex/attention.py | 2 +- modules/sd_hijack_dynamic_atten.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py index 4f74a2c2e..51a9961a7 100644 --- a/modules/intel/ipex/attention.py +++ b/modules/intel/ipex/attention.py @@ -120,7 +120,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop value[start_idx:end_idx, :, :, :], attn_mask=attn_mask[start_idx:end_idx, :, :, :] if attn_mask is not None else attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs - ) + ) torch.xpu.synchronize(query.device) else: hidden_states = original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index 49f89b414..3bdaa0205 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -118,7 +118,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop if devices.backend != "directml": getattr(torch, query.device.type).synchronize() else: - return devices.sdpa_pre_dyanmic_atten(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) + hidden_states = devices.sdpa_pre_dyanmic_atten(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) if is_unsqueezed: hidden_states.squeeze(0) return hidden_states From 3f2c9c4e224c75169e96e5ca58e088a1bc4603cf Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 25 Jan 2025 22:44:29 +0300 Subject: [PATCH 42/66] Fix dyn atten not registering slice rate changes --- modules/sd_hijack_dynamic_atten.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index 3bdaa0205..cce8c13f0 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -28,7 +28,7 @@ def find_query_size(query_size, slice_query_size, slice_rate=4): # Find slice sizes for SDPA @cache -def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size): +def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size, slice_rate=4): batch_size, attn_heads, query_len, _ = query_shape _, _, key_len, _ = key_shape _, _, _, head_dim = value_shape @@ -43,19 +43,19 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz do_head_split = False do_query_split = False - if batch_size * slice_batch_size > shared.opts.dynamic_attention_slice_rate: + if batch_size * slice_batch_size > slice_rate: do_batch_split = True - split_batch_size = find_split_size(split_batch_size, slice_batch_size, slice_rate=shared.opts.dynamic_attention_slice_rate) + split_batch_size = find_split_size(split_batch_size, slice_batch_size, slice_rate=slice_rate) - if split_batch_size * slice_batch_size > shared.opts.dynamic_attention_slice_rate: + if split_batch_size * slice_batch_size > slice_rate: slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 do_head_split = True - split_head_size = find_split_size(split_head_size, slice_head_size, slice_rate=shared.opts.dynamic_attention_slice_rate) + split_head_size = find_split_size(split_head_size, slice_head_size, slice_rate=slice_rate) - if split_batch_size * slice_batch_size > shared.opts.dynamic_attention_slice_rate: + if split_batch_size * slice_batch_size > slice_rate: slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * head_dim * query_element_size / 1024 / 1024 / 2 do_query_split = True - split_query_size = find_query_size(split_query_size, slice_query_size, slice_rate=shared.opts.dynamic_attention_slice_rate) + split_query_size = find_query_size(split_query_size, slice_query_size, slice_rate=slice_rate) return do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size @@ -72,7 +72,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop key = key.unsqueeze(0) if len(value.shape) == 3: value = value.unsqueeze(0) - do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size()) + do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate) # Slice SDPA if do_batch_split: From fad1edd8fa06b7a63b8dbddb478933bb4a4f056c Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 26 Jan 2025 00:18:54 +0300 Subject: [PATCH 43/66] Fix ipex triton --- installer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/installer.py b/installer.py index af8c00f89..61d9c02d2 100644 --- a/installer.py +++ b/installer.py @@ -685,6 +685,8 @@ def install_ipex(torch_command): if "linux" in sys.platform: # default to US server. If The China server is needed, change .../release-whl/stable/xpu/us/ to .../release-whl/stable/xpu/cn/ torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cxx11.abi torchvision==0.20.1+cxx11.abi intel-extension-for-pytorch==2.5.10+xpu oneccl_bind_pt==2.5.0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/') + if os.environ.get('TRITON_COMMAND', None) is None: + os.environ.setdefault('TRITON_COMMAND', '--pre pytorch-triton-xpu==3.1.0+91b14bf559 --index-url https://download.pytorch.org/whl/nightly/xpu') # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.2') else: torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+xpu torchvision==0.21.0+xpu --index-url https://download.pytorch.org/whl/test/xpu') From bb07dd7a8f0611224e7a106ffaadec84e126fb97 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 26 Jan 2025 03:48:01 +0300 Subject: [PATCH 44/66] Add trigger rate control to dyn atten --- CHANGELOG.md | 2 +- modules/sd_hijack_dynamic_atten.py | 12 ++++++------ modules/shared.py | 1 + 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 18e0de6ee..77ac8ec0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -56,7 +56,7 @@ - **loader**: ability to run in-memory models - **schedulers**: ability to create model-less schedulers - **quantiation**: code refactor into dedicated module - - **dynamic attention sdpa**: more correct implementation + - **dynamic attention sdpa**: more correct implementation and new trigger rate control - **Authentication**: - perform auth check on ui startup - unified standard and modern-ui authentication method diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index cce8c13f0..a3e63dd6f 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -28,7 +28,7 @@ def find_query_size(query_size, slice_query_size, slice_rate=4): # Find slice sizes for SDPA @cache -def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size, slice_rate=4): +def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size, slice_rate=4, trigger_rate=6): batch_size, attn_heads, query_len, _ = query_shape _, _, key_len, _ = key_shape _, _, _, head_dim = value_shape @@ -43,7 +43,7 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz do_head_split = False do_query_split = False - if batch_size * slice_batch_size > slice_rate: + if batch_size * slice_batch_size > trigger_rate: do_batch_split = True split_batch_size = find_split_size(split_batch_size, slice_batch_size, slice_rate=slice_rate) @@ -72,7 +72,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop key = key.unsqueeze(0) if len(value.shape) == 3: value = value.unsqueeze(0) - do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate) + do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate, trigger_rate=shared.opts.dynamic_attention_trigger_rate) # Slice SDPA if do_batch_split: @@ -125,7 +125,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop @cache -def find_bmm_slice_sizes(query_shape, query_element_size, slice_rate=4): +def find_bmm_slice_sizes(query_shape, query_element_size, slice_rate=4, trigger_rate=6): if len(query_shape) == 3: batch_size_attention, query_tokens, shape_three = query_shape shape_four = 1 @@ -143,7 +143,7 @@ def find_bmm_slice_sizes(query_shape, query_element_size, slice_rate=4): do_split_2 = False do_split_3 = False - if block_size > slice_rate: + if block_size > trigger_rate: do_split = True split_slice_size = find_split_size(split_slice_size, slice_block_size, slice_rate=slice_rate) if split_slice_size * slice_block_size > slice_rate: @@ -206,7 +206,7 @@ class DynamicAttnProcessorBMM: # Slicing parts: batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2] hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype) - do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_bmm_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate) + do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_bmm_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate, trigger_rate=shared.opts.dynamic_attention_trigger_rate) if do_split: for i in range(batch_size_attention // split_slice_size): diff --git a/modules/shared.py b/modules/shared.py index da9fa3f6d..25bab6fd4 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -515,6 +515,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "sdp_options": OptionInfo(startup_sdp_options, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'Sage attention'], "visible": native}), "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }), "dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory, "step": 0.1, "visible": native}), + "dynamic_attention_trigger_rate": OptionInfo(6, "Dynamic Attention trigger rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory*2, "step": 0.1, "visible": native}), "sub_quad_sep": OptionInfo("

Sub-quadratic options

", "", gr.HTML, {"visible": not native}), "sub_quad_q_chunk_size": OptionInfo(512, "Attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8, "visible": not native}), "sub_quad_kv_chunk_size": OptionInfo(512, "Attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8, "visible": not native}), From a21bc9d9cdcf624f2f52a56158608a6d26c81b83 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 26 Jan 2025 04:45:05 +0300 Subject: [PATCH 45/66] Make IPEX dyn atten care only about the sdpa memory usage --- modules/intel/ipex/attention.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py index 51a9961a7..3e15627e7 100644 --- a/modules/intel/ipex/attention.py +++ b/modules/intel/ipex/attention.py @@ -7,7 +7,7 @@ from functools import cache, wraps # ARC GPUs can't allocate more than 4GB to a single block so we slice the attetion layers -sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 6)) +sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 4)) attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4)) # Find something divisible with the input_tokens @@ -38,7 +38,7 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz _, _, key_len, _ = key_shape _, _, _, head_dim = value_shape - slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024 split_batch_size = batch_size split_head_size = attn_heads @@ -53,12 +53,12 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz split_batch_size = find_split_size(split_batch_size, slice_batch_size) if split_batch_size * slice_batch_size > attention_slice_rate: - slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024 do_head_split = True split_head_size = find_split_size(split_head_size, slice_head_size) if split_batch_size * slice_batch_size > attention_slice_rate: - slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024 do_query_split = True split_query_size = find_query_size(split_query_size, slice_query_size) From b0ecdf2f1ca86dbec44a61fc86416a1dc7da15e5 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 26 Jan 2025 16:57:26 +0300 Subject: [PATCH 46/66] Better dyn atten sdpa usage estimation logic --- modules/intel/ipex/attention.py | 49 ++++++++++++------------------ modules/sd_hijack_dynamic_atten.py | 44 +++++++++++---------------- modules/shared.py | 4 +-- 3 files changed, 40 insertions(+), 57 deletions(-) diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py index 3e15627e7..e8e5b815b 100644 --- a/modules/intel/ipex/attention.py +++ b/modules/intel/ipex/attention.py @@ -7,38 +7,29 @@ from functools import cache, wraps # ARC GPUs can't allocate more than 4GB to a single block so we slice the attetion layers -sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 4)) -attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4)) +sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 3)) +attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 2)) # Find something divisible with the input_tokens @cache -def find_split_size(split_size, slice_block_size): - while (split_size * slice_block_size) > attention_slice_rate: - split_size = split_size // 2 +def find_split_size(original_size, slice_block_size, slice_rate=2): + split_size = original_size + while True: + if (split_size * slice_block_size) <= slice_rate and original_size % split_size == 0: + return split_size + split_size = split_size - 1 if split_size <= 1: - split_size = 1 - break + return 1 return split_size -@cache -def find_query_size(query_size, slice_query_size): - while (math.sqrt(query_size) * slice_query_size) > attention_slice_rate: - query_size = query_size // 2 - if query_size <= 1: - query_size = 1 - break - return query_size - - # Find slice sizes for SDPA @cache -def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size): +def find_sdpa_slice_sizes(query_shape, key_shape, query_element_size, slice_rate=2, trigger_rate=3): batch_size, attn_heads, query_len, _ = query_shape _, _, key_len, _ = key_shape - _, _, _, head_dim = value_shape - slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024 + slice_batch_size = attn_heads * (query_len * key_len) * query_element_size / 1024 / 1024 / 1024 split_batch_size = batch_size split_head_size = attn_heads @@ -48,19 +39,19 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz do_head_split = False do_query_split = False - if batch_size * slice_batch_size > sdpa_slice_trigger_rate: + if batch_size * slice_batch_size >= trigger_rate: do_batch_split = True - split_batch_size = find_split_size(split_batch_size, slice_batch_size) + split_batch_size = find_split_size(batch_size, slice_batch_size, slice_rate=slice_rate) - if split_batch_size * slice_batch_size > attention_slice_rate: - slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024 + if split_batch_size * slice_batch_size > slice_rate: + slice_head_size = split_batch_size * (query_len * key_len) * query_element_size / 1024 / 1024 / 1024 do_head_split = True - split_head_size = find_split_size(split_head_size, slice_head_size) + split_head_size = find_split_size(attn_heads, slice_head_size, slice_rate=slice_rate) - if split_batch_size * slice_batch_size > attention_slice_rate: - slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024 + if split_head_size * slice_head_size > slice_rate: + slice_query_size = split_batch_size * split_head_size * (key_len) * query_element_size / 1024 / 1024 / 1024 do_query_split = True - split_query_size = find_query_size(split_query_size, slice_query_size) + split_query_size = find_split_size(query_len, slice_query_size, slice_rate=slice_rate) return do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size @@ -78,7 +69,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop key = key.unsqueeze(0) if len(value.shape) == 3: value = value.unsqueeze(0) - do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size()) + do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, query.element_size(), slice_rate=attention_slice_rate, trigger_rate=sdpa_slice_trigger_rate) # Slice SDPA if do_batch_split: diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index a3e63dd6f..9ea103036 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -6,34 +6,26 @@ from diffusers.utils import USE_PEFT_BACKEND # pylint: disable=unused-import from modules import shared, devices +# Find something divisible with the input_tokens @cache -def find_split_size(split_size, slice_block_size, slice_rate=4): - while (split_size * slice_block_size) > slice_rate: - split_size = split_size // 2 +def find_split_size(original_size, slice_block_size, slice_rate=2): + split_size = original_size + while True: + if (split_size * slice_block_size) <= slice_rate and original_size % split_size == 0: + return split_size + split_size = split_size - 1 if split_size <= 1: - split_size = 1 - break + return 1 return split_size -@cache -def find_query_size(query_size, slice_query_size, slice_rate=4): - while (math.sqrt(query_size) * slice_query_size) > slice_rate: - query_size = query_size // 2 - if query_size <= 1: - query_size = 1 - break - return query_size - - # Find slice sizes for SDPA @cache -def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_size, slice_rate=4, trigger_rate=6): +def find_sdpa_slice_sizes(query_shape, key_shape, query_element_size, slice_rate=2, trigger_rate=3): batch_size, attn_heads, query_len, _ = query_shape _, _, key_len, _ = key_shape - _, _, _, head_dim = value_shape - slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + slice_batch_size = attn_heads * (query_len * key_len) * query_element_size / 1024 / 1024 / 1024 split_batch_size = batch_size split_head_size = attn_heads @@ -43,19 +35,19 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz do_head_split = False do_query_split = False - if batch_size * slice_batch_size > trigger_rate: + if batch_size * slice_batch_size >= trigger_rate: do_batch_split = True - split_batch_size = find_split_size(split_batch_size, slice_batch_size, slice_rate=slice_rate) + split_batch_size = find_split_size(batch_size, slice_batch_size, slice_rate=slice_rate) if split_batch_size * slice_batch_size > slice_rate: - slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + slice_head_size = split_batch_size * (query_len * key_len) * query_element_size / 1024 / 1024 / 1024 do_head_split = True - split_head_size = find_split_size(split_head_size, slice_head_size, slice_rate=slice_rate) + split_head_size = find_split_size(attn_heads, slice_head_size, slice_rate=slice_rate) - if split_batch_size * slice_batch_size > slice_rate: - slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * head_dim * query_element_size / 1024 / 1024 / 2 + if split_head_size * slice_head_size > slice_rate: + slice_query_size = split_batch_size * split_head_size * (key_len) * query_element_size / 1024 / 1024 / 1024 do_query_split = True - split_query_size = find_query_size(split_query_size, slice_query_size, slice_rate=slice_rate) + split_query_size = find_split_size(query_len, slice_query_size, slice_rate=slice_rate) return do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size @@ -72,7 +64,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop key = key.unsqueeze(0) if len(value.shape) == 3: value = value.unsqueeze(0) - do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, value.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate, trigger_rate=shared.opts.dynamic_attention_trigger_rate) + do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate, trigger_rate=shared.opts.dynamic_attention_trigger_rate) # Slice SDPA if do_batch_split: diff --git a/modules/shared.py b/modules/shared.py index 25bab6fd4..6edb9daf4 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -514,8 +514,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cross_attention_optimization": OptionInfo(startup_cross_attention, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention(native)}), "sdp_options": OptionInfo(startup_sdp_options, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'Sage attention'], "visible": native}), "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }), - "dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory, "step": 0.1, "visible": native}), - "dynamic_attention_trigger_rate": OptionInfo(6, "Dynamic Attention trigger rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory*2, "step": 0.1, "visible": native}), + "dynamic_attention_slice_rate": OptionInfo(2, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory, "step": 0.1, "visible": native}), + "dynamic_attention_trigger_rate": OptionInfo(3, "Dynamic Attention trigger rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory*2, "step": 0.1, "visible": native}), "sub_quad_sep": OptionInfo("

Sub-quadratic options

", "", gr.HTML, {"visible": not native}), "sub_quad_q_chunk_size": OptionInfo(512, "Attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8, "visible": not native}), "sub_quad_kv_chunk_size": OptionInfo(512, "Attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8, "visible": not native}), From f16599324a541230056b91f8b27b543408c8d3bd Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 26 Jan 2025 09:00:57 -0500 Subject: [PATCH 47/66] fix fasthunyuan Signed-off-by: Vladimir Mandic --- scripts/hunyuanvideo.py | 33 +++++++++++++++++++-------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py index 36edfea25..e132a7837 100644 --- a/scripts/hunyuanvideo.py +++ b/scripts/hunyuanvideo.py @@ -15,6 +15,11 @@ default_template = """Describe the video by detailing the following aspects: 6. Thematic and aesthetic concepts associated with the scene, i.e. realistic, futuristic, fairy tale, etc. """ +models = { + 'HunyuanVideo': { 'repo': 'tencent/HunyuanVideo', 'revision': 'refs/pr/18' }, + 'FastHunyuan': { 'repo': 'FastVideo/FastHunyuan', 'revision': None }, +} + def get_template(template: str = None): # diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.DEFAULT_PROMPT_TEMPLATE base_template_pre = "<|start_header_id|>system<|end_header_id|>\n\n" @@ -63,7 +68,7 @@ class Script(scripts.Script): with gr.Row(): gr.HTML('  Hunyuan Video
') with gr.Row(): - repo_id = gr.Dropdown(label='Model', choices=['tencent/HunyuanVideo', 'FastVideo/FastHunyuan'], value='tencent/HunyuanVideo') + model = gr.Dropdown(label='Model', choices=list(models.keys()), value=list(models.keys())[0]) with gr.Row(): num_frames = gr.Slider(label='Frames', minimum=9, maximum=257, step=1, value=45) tile_frames = gr.Slider(label='Tile frames', minimum=1, maximum=64, step=1, value=16) @@ -77,9 +82,9 @@ class Script(scripts.Script): with gr.Row(): from modules.ui_sections import create_video_inputs video_type, duration, gif_loop, mp4_pad, mp4_interpolate = create_video_inputs(tab='img2img' if is_img2img else 'txt2img') - return [repo_id, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate] + return [model, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate] - def load(self, repo_id:str): + def load(self, model:str): if shared.sd_model.__class__ != diffusers.HunyuanVideoPipeline: sd_models.unload_model_weights() t0 = time.time() @@ -92,36 +97,36 @@ class Script(scripts.Script): if quant_args: model_quant.load_torchao(f'Load model: type=HunyuanVideo quant={quant_args}') transformer = diffusers.HunyuanVideoTransformer3DModel.from_pretrained( - repo_id, + pretrained_model_name_or_path='tencent/HunyuanVideo', subfolder="transformer", torch_dtype=devices.dtype, - revision="refs/pr/18", - cache_dir = shared.opts.hfcache_dir, + revision='refs/pr/18', + cache_dir=shared.opts.hfcache_dir, **quant_args ) shared.log.debug(f'Video: module={transformer.__class__.__name__}') text_encoder = transformers.LlamaModel.from_pretrained( - repo_id, + pretrained_model_name_or_path=models.get(model)['repo'], subfolder="text_encoder", - revision="refs/pr/18", + revision=models.get(model)['revision'], cache_dir = shared.opts.hfcache_dir, torch_dtype=devices.dtype, **quant_args ) shared.log.debug(f'Video: module={text_encoder.__class__.__name__}') shared.sd_model = diffusers.HunyuanVideoPipeline.from_pretrained( - repo_id, + pretrained_model_name_or_path=models.get(model)['repo'], transformer=transformer, text_encoder=text_encoder, - revision="refs/pr/18", + revision=models.get(model)['revision'], cache_dir = shared.opts.hfcache_dir, torch_dtype=devices.dtype, **quant_args ) t1 = time.time() - shared.log.debug(f'Video: load cls={shared.sd_model.__class__.__name__} repo="{repo_id}" dtype={devices.dtype} time={t1-t0:.2f}') + shared.log.debug(f'Video: load cls={shared.sd_model.__class__.__name__} model="{model}" repo={models.get(model)['repo']} dtype={devices.dtype} time={t1-t0:.2f}') sd_models.set_diffuser_options(shared.sd_model) - shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(repo_id) + shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(models.get(model)['repo']) shared.sd_model.sd_model_hash = None shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode shared.sd_model.vae.orig_encode_prompt = shared.sd_model.encode_prompt @@ -130,7 +135,7 @@ class Script(scripts.Script): shared.sd_model.vae.enable_slicing() shared.sd_model.vae.enable_tiling() - def run(self, p: processing.StableDiffusionProcessing, repo_id, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument + def run(self, p: processing.StableDiffusionProcessing, model, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument # set params num_frames = int(num_frames) p.width = 16 * int(p.width // 16) @@ -139,7 +144,7 @@ class Script(scripts.Script): p.ops.append('video') # load model - self.load(repo_id) + self.load(model) shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) devices.torch_gc(force=True) From de1344196fc5de35cda413e40d19f0c2114eefd7 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 26 Jan 2025 21:46:48 +0300 Subject: [PATCH 48/66] Skip triton install with rocm --- installer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/installer.py b/installer.py index 61d9c02d2..8d6aeb471 100644 --- a/installer.py +++ b/installer.py @@ -625,6 +625,9 @@ def install_rocm_zluda(): else: torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm{rocm.version}') + if os.environ.get('TRITON_COMMAND', None) is None: + os.environ.setdefault('TRITON_COMMAND', 'skip') # pytorch auto installs pytorch-triton-rocm as a dependency instead + if sys.version_info < (3, 11): ort_version = os.environ.get('ONNXRUNTIME_VERSION', None) if rocm.version is None or float(rocm.version) > 6.0: @@ -745,7 +748,7 @@ def install_torch_addons(): install('optimum-quanto==0.2.6', 'optimum-quanto') if not args.experimental: uninstall('wandb', quiet=True) - if triton_command is not None: + if triton_command is not None and triton_command != 'skip': install(triton_command, 'triton', quiet=True) ts('addons', t_start) From 47224f731125d25ae7f77c116063bcfe3ad51525 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 26 Jan 2025 22:21:59 +0300 Subject: [PATCH 49/66] Fix hunyuanvideo import error --- scripts/hunyuanvideo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py index e132a7837..ec061836f 100644 --- a/scripts/hunyuanvideo.py +++ b/scripts/hunyuanvideo.py @@ -124,7 +124,7 @@ class Script(scripts.Script): **quant_args ) t1 = time.time() - shared.log.debug(f'Video: load cls={shared.sd_model.__class__.__name__} model="{model}" repo={models.get(model)['repo']} dtype={devices.dtype} time={t1-t0:.2f}') + shared.log.debug(f'Video: load cls={shared.sd_model.__class__.__name__} model="{model}" repo={models.get(model)["repo"]} dtype={devices.dtype} time={t1-t0:.2f}') sd_models.set_diffuser_options(shared.sd_model) shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(models.get(model)['repo']) shared.sd_model.sd_model_hash = None From a83497cd823ade89134073b348ec8f29ac44c572 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 26 Jan 2025 14:58:54 -0500 Subject: [PATCH 50/66] moondream2, controlnetunion, etc. Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 8 ++++++-- modules/control/run.py | 7 ++++--- modules/control/units/controlnet.py | 14 ++++++++++---- modules/processing_class.py | 3 ++- modules/vqa.py | 7 ++++++- 5 files changed, 28 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 77ac8ec0a..d027b979a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-24 +## Update for 2025-01-26 - **Contributing**: - if you'd like to contribute, please see updated [contributing](https://github.com/vladmandic/automatic/blob/dev/CONTRIBUTING) guidelines @@ -36,7 +36,7 @@ - use the us server by default on linux - use pytorch test branch on windows - extend the supported python versions - - removed diffusers attention hijack as it is a duplicate of dynamic attention bmm + - improve sdpa dynamic attention - **Torch FP8** - uses torch `float8_e4m3fn` or `float8_e5m2` as data storage and performs dynamic upcasting to compute `dtype` as needed - compatible with most `unet` and `transformer` models: e.g. *sd15, sdxl, sd35, flux.1, hunyuan-video, ltx-video, etc.* @@ -77,6 +77,10 @@ - xyz grid handle invalid values - omnigen pipeline handle float seeds - correct logging of docker status on logs, thanks @kmscode + - fix omnigen + - fix docker status reporting + - vlm/vqa with moondream2 + - rocm do not override triton installation ## Update for 2025-01-15 diff --git a/modules/control/run.py b/modules/control/run.py index 21fc76a4b..bfb3a0248 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -148,7 +148,9 @@ def check_active(p, unit_type, units): active_end.append(float(u.end)) p.guess_mode = u.guess if isinstance(u.mode, str): - p.control_mode = u.choices.index(u.mode) if u.mode in u.choices else 0 + if not hasattr(p, 'control_mode'): + p.control_mode = [] + p.control_mode.append(u.choices.index(u.mode) if u.mode in u.choices else 0) p.is_tile = p.is_tile or 'tile' in u.mode.lower() p.control_tile = u.tile p.extra_generation_params["Control mode"] = u.mode @@ -427,8 +429,6 @@ def control_run(state: str = '', else: original_pipeline = None - possible = sd_models.get_call(pipe).keys() - try: with devices.inference_context(): if isinstance(inputs, str): # only video, the rest is a list @@ -460,6 +460,7 @@ def control_run(state: str = '', if pipe is None: # pipe may have been reset externally pipe = set_pipe(p, has_models, unit_type, selected_models, active_model, active_strength, control_conditioning, control_guidance_start, control_guidance_end, inits) debug_log(f'Control pipeline reinit: class={pipe.__class__.__name__}') + possible = sd_models.get_call(pipe).keys() processed_image = None if frame is not None: inputs = [Image.fromarray(frame)] # cv2 to pil diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index 4837577fe..9c638f648 100644 --- a/modules/control/units/controlnet.py +++ b/modules/control/units/controlnet.py @@ -291,7 +291,7 @@ class ControlNet(): log.debug(f'Control {what} model NNCF Compress: id="{model_id}"') from installer import install install('nncf==2.7.0', quiet=True) - from modules.sd_models_compile import nncf_compress_model + from modules.model_quant import nncf_compress_model self.model = nncf_compress_model(self.model) except Exception as e: log.error(f'Control {what} model NNCF Compression failed: id="{model_id}" {e}') @@ -299,7 +299,7 @@ class ControlNet(): try: log.debug(f'Control {what} model Optimum Quanto: id="{model_id}"') model_quant.load_quanto('Load model: type=ControlNet') - from modules.sd_models_compile import optimum_quanto_model + from modules.model_quant import optimum_quanto_model self.model = optimum_quanto_model(self.model) except Exception as e: log.error(f'Control {what} model Optimum Quanto: id="{model_id}" {e}') @@ -335,9 +335,15 @@ class ControlNetPipeline(): return elif detect.is_sdxl(pipeline) and len(controlnets) > 0: from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetUnionPipeline - if controlnet.__class__.__name__ == 'ControlNetUnionModel': + classes = [c.__class__.__name__ for c in controlnets] + if any(c == 'ControlNetUnionModel' for c in classes): + if not all(c == 'ControlNetUnionModel' for c in classes): + log.warning(f'Control {what}: units={classes} mixed type') cls = StableDiffusionXLControlNetUnionPipeline - controlnets = controlnets[0] # using only first one + if len(controlnets) > 1: + # TODO controlnet-union multi-unit + log.warning(f'Control {what}: units={classes} supports single unit only') + controlnets = controlnets[0] else: cls = StableDiffusionXLControlNetPipeline self.pipeline = cls( diff --git a/modules/processing_class.py b/modules/processing_class.py index 3eeb73622..a01108c18 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -354,7 +354,8 @@ class StableDiffusionProcessing: raise NotImplementedError def close(self): - self.sampler = None # pylint: disable=attribute-defined-outside-init + self.sampler = None + self.scripts = None class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): diff --git a/modules/vqa.py b/modules/vqa.py index ee4197a5e..d0172159b 100644 --- a/modules/vqa.py +++ b/modules/vqa.py @@ -117,7 +117,12 @@ def pix(question: str, image: Image.Image, repo: str = None): def moondream(question: str, image: Image.Image, repo: str = None): global processor, model, loaded # pylint: disable=global-statement if model is None or loaded != repo: - model = transformers.AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, cache_dir=shared.opts.hfcache_dir) # revision = "2024-03-05" + model = transformers.AutoModelForCausalLM.from_pretrained( + repo, + revision="2024-08-26", + trust_remote_code=True, + cache_dir=shared.opts.hfcache_dir + ) processor = transformers.AutoTokenizer.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) loaded = repo model.eval() From 338e3a59dae722ef4d105a2850186abe820c8737 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 27 Jan 2025 09:01:24 -0500 Subject: [PATCH 51/66] port streaming load to diffusers Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + modules/sd_models.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d027b979a..f2f31e4b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -81,6 +81,7 @@ - fix docker status reporting - vlm/vqa with moondream2 - rocm do not override triton installation + - port streaming model load to diffusers ## Update for 2025-01-15 diff --git a/modules/sd_models.py b/modules/sd_models.py index 2de127347..b87efed6f 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -400,6 +400,8 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con diffusers.loaders.single_file_utils.CHECKPOINT_KEY_NAMES["clip"] = "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight" # patch for diffusers==0.28.0 diffusers_load_config['use_safetensors'] = True diffusers_load_config['cache_dir'] = shared.opts.hfcache_dir # use hfcache instead of diffusers dir as this is for config only in case of single-file + if shared.opts.stream_load: + diffusers_load_config['disable_mmap '] = True if shared.opts.disable_accelerate: from diffusers.utils import import_utils import_utils._accelerate_available = False # pylint: disable=protected-access From 79416af9944234cd7868aaacb88344e9e9e0a6fc Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 27 Jan 2025 10:00:06 -0500 Subject: [PATCH 52/66] parse cgroups Signed-off-by: Vladimir Mandic --- extensions-builtin/sd-extension-system-info | 2 +- launch.py | 10 +++------ modules/memstats.py | 24 +++++++++++++++++++++ modules/shared.py | 5 ++--- wiki | 2 +- 5 files changed, 31 insertions(+), 12 deletions(-) diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index 2e50e314a..8c7edb3be 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit 2e50e314abfab364f41c29b079e37f9658bb1627 +Subproject commit 8c7edb3be11b8b8c2d2dcd0421e93345bd20fcae diff --git a/launch.py b/launch.py index f92e9bfda..c80840036 100755 --- a/launch.py +++ b/launch.py @@ -151,13 +151,9 @@ def run_extension_installer(ext_dir): # compatbility function def get_memory_stats(): - import psutil - def gb(val: float): - return round(val / 1024 / 1024 / 1024, 2) - process = psutil.Process(os.getpid()) - res = process.memory_info() - ram_total = 100 * res.rss / process.memory_percent() - return f'{gb(res.rss)}/{gb(ram_total)}' + from modules.memstats import ram_stats + res = ram_stats() + return f'{res["used"]}/{res["total"]}' def start_server(immediate=True, server=None): diff --git a/modules/memstats.py b/modules/memstats.py index cb43f246c..1ff7f721f 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -7,12 +7,34 @@ from modules import shared, errors fail_once = False mem = {} +docker_limit = None +runpod_limit = None def gb(val: float): return round(val / 1024 / 1024 / 1024, 2) +def get_docker_limit(): + global docker_limit # pylint: disable=global-statement + if docker_limit is not None: + return docker_limit + try: + with open('/sys/fs/cgroup/memory/memory.limit_in_bytes', 'r', encoding='utf8') as f: + docker_limit = float(f.read()) + except Exception: + docker_limit = sys.float_info.max + return docker_limit + + +def get_runpod_limit(): + global runpod_limit # pylint: disable=global-statement + if runpod_limit is not None: + return runpod_limit + runpod_limit = float(os.environ.get('RUNPOD_MEM_GB', sys.float_info.max)) + return runpod_limit + + def memory_stats(): global fail_once # pylint: disable=global-statement mem.clear() @@ -20,6 +42,7 @@ def memory_stats(): process = psutil.Process(os.getpid()) res = process.memory_info() ram_total = 100 * res.rss / process.memory_percent() + ram_total = min(ram_total, get_docker_limit(), get_runpod_limit()) ram = { 'used': gb(res.rss), 'total': gb(ram_total) } mem.update({ 'ram': ram }) except Exception as e: @@ -54,6 +77,7 @@ def ram_stats(): process = psutil.Process(os.getpid()) res = process.memory_info() ram_total = 100 * res.rss / process.memory_percent() + ram_total = min(ram_total, docker_limit(), runpod_limit()) ram = { 'used': gb(res.rss), 'total': gb(ram_total) } return ram except Exception: diff --git a/modules/shared.py b/modules/shared.py index 6edb9daf4..38e59ef48 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -7,9 +7,7 @@ import json import threading import contextlib from types import SimpleNamespace -from urllib.parse import urlparse from enum import Enum -import psutil import requests import gradio as gr import fasteners @@ -205,7 +203,7 @@ elif cmd_opts.use_directml: devices.backend = devices.get_backend(cmd_opts) devices.device = devices.get_optimal_device() mem_stat = memory_stats() -cpu_memory = round(psutil.virtual_memory().total / 1024 / 1024 / 1024, 2) +cpu_memory = mem_stat['ram']['total'] if "ram" in mem_stat else 0 gpu_memory = mem_stat['gpu']['total'] if "gpu" in mem_stat else 0 native = backend == Backend.DIFFUSERS if not files_cache.do_cache_folders: @@ -307,6 +305,7 @@ default_checkpoint = list_checkpoint_titles()[0] if len(list_checkpoint_titles() def is_url(string): + from urllib.parse import urlparse parsed_url = urlparse(string) return all([parsed_url.scheme, parsed_url.netloc]) diff --git a/wiki b/wiki index eb2e03e71..4b2c2d0d8 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit eb2e03e719e200ee8c2523a08e75caf7e43df1cd +Subproject commit 4b2c2d0d8e621dfd8c216124a020501b8f4b2528 From 1afa4be02c6007e0065708fe70292fb22ef40784 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 27 Jan 2025 10:25:01 -0500 Subject: [PATCH 53/66] detect external and public ip Signed-off-by: Vladimir Mandic --- webui.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/webui.py b/webui.py index 23a9eba95..bd99bfc1d 100644 --- a/webui.py +++ b/webui.py @@ -209,6 +209,27 @@ def async_policy(): asyncio.set_event_loop_policy(AnyThreadEventLoopPolicy()) +def get_external_ip(): + import socket + try: + ip_address = socket.gethostbyname(socket.gethostname()) + if ip_address.startswith('127.'): + return None + return ip_address + except Exception: + return None + + +def get_remote_ip(): + import requests + try: + response = requests.get('https://api.ipify.org?format=json', timeout=2) + ip_address = response.json()['ip'] + return ip_address + except Exception: + return None + + def start_common(): log.debug('Entering start sequence') if shared.cmd_opts.data_dir is not None and len(shared.cmd_opts.data_dir) > 0: @@ -284,6 +305,14 @@ def start_ui(): if shared.cmd_opts.data_dir is not None: gr_tempdir.register_tmp_file(shared.demo, os.path.join(shared.cmd_opts.data_dir, 'x')) shared.log.info(f'Local URL: {local_url}') + if shared.cmd_opts.listen: + proto = 'https' if shared.cmd_opts.tls_keyfile is not None else 'http' + external_ip = get_external_ip() + if external_ip is not None: + shared.log.info(f'External URL: {proto}://{external_ip}:{shared.cmd_opts.port}') + public_ip = get_remote_ip() + if public_ip is not None: + shared.log.info(f'Public URL: {proto}://{public_ip}:{shared.cmd_opts.port}') if shared.cmd_opts.docs: shared.log.info(f'API Docs: {local_url[:-1]}/docs') # pylint: disable=unsubscriptable-object shared.log.info(f'API ReDocs: {local_url[:-1]}/redocs') # pylint: disable=unsubscriptable-object From 8a28e8b7166f0f081128be5a087a245a1257ca30 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 27 Jan 2025 10:44:24 -0500 Subject: [PATCH 54/66] update changelog Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f2f31e4b9..854622ca7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,10 +57,11 @@ - **schedulers**: ability to create model-less schedulers - **quantiation**: code refactor into dedicated module - **dynamic attention sdpa**: more correct implementation and new trigger rate control -- **Authentication**: +- **Remote access**: - perform auth check on ui startup - - unified standard and modern-ui authentication method - - cleanup auth logging + - unified standard and modern-ui authentication method & cleanup auth logging + - detect & report local/external/public ip addresses if using `listen` mode + - detect *docker* enforced limits instead of system limits if running in a container - **Fixes**: - non-full vae decode - send-to image transfer From 6358276460067b9dbe912e235144b84ee2998e6b Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 27 Jan 2025 11:02:28 -0500 Subject: [PATCH 55/66] warn if using public interface without auth Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + webui.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 854622ca7..defe0e1d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -62,6 +62,7 @@ - unified standard and modern-ui authentication method & cleanup auth logging - detect & report local/external/public ip addresses if using `listen` mode - detect *docker* enforced limits instead of system limits if running in a container + - warn if using public interface without authentication - **Fixes**: - non-full vae decode - send-to image transfer diff --git a/webui.py b/webui.py index bd99bfc1d..d036e6ad3 100644 --- a/webui.py +++ b/webui.py @@ -306,6 +306,8 @@ def start_ui(): gr_tempdir.register_tmp_file(shared.demo, os.path.join(shared.cmd_opts.data_dir, 'x')) shared.log.info(f'Local URL: {local_url}') if shared.cmd_opts.listen: + if not gradio_auth_creds: + shared.log.warning('Public interface enabled without authentication') proto = 'https' if shared.cmd_opts.tls_keyfile is not None else 'http' external_ip = get_external_ip() if external_ip is not None: From e90dfe7e0a126f563c5833955e00f779b5f8d868 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 27 Jan 2025 15:39:59 -0500 Subject: [PATCH 56/66] fix fasthunyuan Signed-off-by: Vladimir Mandic --- installer.py | 2 +- scripts/hunyuanvideo.py | 18 +++++++++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/installer.py b/installer.py index 8d6aeb471..33347a0be 100644 --- a/installer.py +++ b/installer.py @@ -492,7 +492,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git: return - sha = '07860f991639f35f4b5a152676bd4d590c3e589e' # diffusers commit hash + sha = '8ceec90d76767035a63b879d659a2ed8e12c5bba' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py index ec061836f..a4bb148b0 100644 --- a/scripts/hunyuanvideo.py +++ b/scripts/hunyuanvideo.py @@ -19,6 +19,8 @@ models = { 'HunyuanVideo': { 'repo': 'tencent/HunyuanVideo', 'revision': 'refs/pr/18' }, 'FastHunyuan': { 'repo': 'FastVideo/FastHunyuan', 'revision': None }, } +loaded_model = None + def get_template(template: str = None): # diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video.DEFAULT_PROMPT_TEMPLATE @@ -85,7 +87,8 @@ class Script(scripts.Script): return [model, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate] def load(self, model:str): - if shared.sd_model.__class__ != diffusers.HunyuanVideoPipeline: + global loaded_model # pylint: disable=global-statement + if shared.sd_model.__class__ != diffusers.HunyuanVideoPipeline or model != loaded_model: sd_models.unload_model_weights() t0 = time.time() quant_args = {} @@ -113,12 +116,20 @@ class Script(scripts.Script): torch_dtype=devices.dtype, **quant_args ) + text_encoder_2 = transformers.CLIPTextModel.from_pretrained( + pretrained_model_name_or_path=models.get(model)['repo'], + subfolder="text_encoder_2", + revision=models.get(model)['revision'], + cache_dir = shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + ) shared.log.debug(f'Video: module={text_encoder.__class__.__name__}') shared.sd_model = diffusers.HunyuanVideoPipeline.from_pretrained( - pretrained_model_name_or_path=models.get(model)['repo'], + pretrained_model_name_or_path='tencent/HunyuanVideo', transformer=transformer, text_encoder=text_encoder, - revision=models.get(model)['revision'], + text_encoder_2=text_encoder_2, + revision='refs/pr/18', cache_dir = shared.opts.hfcache_dir, torch_dtype=devices.dtype, **quant_args @@ -134,6 +145,7 @@ class Script(scripts.Script): shared.sd_model.encode_prompt = hijack_encode_prompt shared.sd_model.vae.enable_slicing() shared.sd_model.vae.enable_tiling() + loaded_model = model def run(self, p: processing.StableDiffusionProcessing, model, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument # set params From f0a212036825abe0cae1ee6b4cb475e0f8b60527 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 27 Jan 2025 16:13:32 -0500 Subject: [PATCH 57/66] fix typo Signed-off-by: Vladimir Mandic --- modules/sd_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index b87efed6f..a3d749be8 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -401,7 +401,7 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con diffusers_load_config['use_safetensors'] = True diffusers_load_config['cache_dir'] = shared.opts.hfcache_dir # use hfcache instead of diffusers dir as this is for config only in case of single-file if shared.opts.stream_load: - diffusers_load_config['disable_mmap '] = True + diffusers_load_config['disable_mmap'] = True if shared.opts.disable_accelerate: from diffusers.utils import import_utils import_utils._accelerate_available = False # pylint: disable=protected-access From a7c0219577d7eb1f4616650b25f7df5980438ec0 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 28 Jan 2025 08:25:55 -0500 Subject: [PATCH 58/66] pab: pyramid attention broadcast Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 9 +++++++-- installer.py | 2 +- modules/sd_offload.py | 16 +++++++++++++++- modules/shared.py | 10 ++++++++-- scripts/hunyuanvideo.py | 1 + 5 files changed, 32 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index defe0e1d2..acf49e095 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2025-01-26 +## Update for 2025-01-28 - **Contributing**: - if you'd like to contribute, please see updated [contributing](https://github.com/vladmandic/automatic/blob/dev/CONTRIBUTING) guidelines @@ -25,6 +25,11 @@ - **hunyuan video** support for [FastHunyuan](https://huggingface.co/FastVideo/FastHunyuan) simply select model variant and set appropriate parameters recommended: sampler-shift=17, steps=6, resolution=720x1280, frames=125, guidance>6.0 +- [PAB: Pyramid Attention Broadcast](https://oahzxl.github.io/PAB/) + - speed up generation by caching attention results between steps + - enable in *settings -> pipeline modifiers -> pab* + - adjust settings as needed: wider timestep range means more acceleration, but higher accuracy drop + - compatible with most `transformer` based models: e.g. flux.1, hunyuan-video, lyx-video, mochi, etc. - [ParaAttention](https://github.com/chengzeyi/ParaAttention) - first-block caching that can significantly speed up generation by dynamically reusing partial outputs between steps - available for: flux, hunyuan-video, ltx-video, mochi @@ -39,7 +44,7 @@ - improve sdpa dynamic attention - **Torch FP8** - uses torch `float8_e4m3fn` or `float8_e5m2` as data storage and performs dynamic upcasting to compute `dtype` as needed - - compatible with most `unet` and `transformer` models: e.g. *sd15, sdxl, sd35, flux.1, hunyuan-video, ltx-video, etc.* + - compatible with most `unet` and `transformer` based models: e.g. *sd15, sdxl, sd35, flux.1, hunyuan-video, ltx-video, etc.* this is alternative to `bnb`/`quanto`/`torchao` quantization on models/platforms/gpus where those libraries are not available - enable in *settings -> quantization -> layerwise casting* - [PerFlow](https://github.com/magic-research/piecewise-rectified-flow) diff --git a/installer.py b/installer.py index 33347a0be..d189a835a 100644 --- a/installer.py +++ b/installer.py @@ -492,7 +492,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git: return - sha = '8ceec90d76767035a63b879d659a2ed8e12c5bba' # diffusers commit hash + sha = '7b100ce589b917d4c116c9e61a6ec46d4f2ab062' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' diff --git a/modules/sd_offload.py b/modules/sd_offload.py index aecb572ba..94a75496e 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -3,8 +3,8 @@ import sys import time import inspect import torch +import diffusers import accelerate.hooks - from modules import shared, devices, errors, model_quant from modules.timer import process as process_timer @@ -271,8 +271,22 @@ def apply_balanced_offload(sd_model, exclude=[]): apply_balanced_offload_to_module(sd_model.prior_pipe) if hasattr(sd_model, "decoder_pipe"): apply_balanced_offload_to_module(sd_model.decoder_pipe) + if shared.opts.layerwise_quantization: model_quant.apply_layerwise(sd_model, quiet=True) # need to reapply since hooks were removed/readded + if shared.opts.pab_enabled and hasattr(sd_model, 'transformer'): + pab_config = diffusers.PyramidAttentionBroadcastConfig( + spatial_attention_block_skip_range=shared.opts.pab_block_skip_range, + spatial_attention_timestep_skip_range=(int(100 * shared.opts.pab_timestep_skip_start), int(100 * shared.opts.pab_timestep_skip_end)), + current_timestep_callback=lambda: sd_model.current_timestep, # pylint: disable=protected-access + ) + try: + diffusers.apply_pyramid_attention_broadcast(sd_model.transformer, pab_config) + except Exception: # hook may already exist + pass + if not cached: + shared.log.info(f'Applying PAB: cls={sd_model.transformer.__class__.__name__} block={shared.opts.pab_block_skip_range} start={shared.opts.pab_timestep_skip_start} end={shared.opts.pab_timestep_skip_end}') + set_accelerate(sd_model) t = time.time() - t0 process_timer.add('offload', t) diff --git a/modules/shared.py b/modules/shared.py index 38e59ef48..fe77c5a45 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -604,10 +604,16 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { "freeu_s1": OptionInfo(0.9, "1st stage skip", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), "freeu_s2": OptionInfo(0.2, "2nd stage skip", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), - "pag_sep": OptionInfo("

Perturbed-Attention Guidance

", "", gr.HTML), + "pag_sep": OptionInfo("

PAG: Perturbed attention guidance

", "", gr.HTML), "pag_apply_layers": OptionInfo("m0", "PAG layer names"), - "para_sep": OptionInfo("

Para-Attention

", "", gr.HTML), + "pab_sep": OptionInfo("

PAB: Pyramid attention broadcast

", "", gr.HTML), + "pab_enabled": OptionInfo(False, "Attention cache enabled"), + "pab_block_skip_range": OptionInfo(2, "Block skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}), + "pab_timestep_skip_start": OptionInfo(0.1, "Timestep skip start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}), + "pab_timestep_skip_end": OptionInfo(0.8, "Timestep skip end", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}), + + "para_sep": OptionInfo("

Para-attention

", "", gr.HTML), "para_cache_enabled": OptionInfo(False, "First-block cache enabled"), "para_diff_threshold": OptionInfo(0.1, "Residual diff threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py index a4bb148b0..d67fd7d31 100644 --- a/scripts/hunyuanvideo.py +++ b/scripts/hunyuanvideo.py @@ -145,6 +145,7 @@ class Script(scripts.Script): shared.sd_model.encode_prompt = hijack_encode_prompt shared.sd_model.vae.enable_slicing() shared.sd_model.vae.enable_tiling() + shared.sd_model.vae.use_framewise_decoding = True loaded_model = model def run(self, p: processing.StableDiffusionProcessing, model, num_frames, tile_frames, override_scheduler, scheduler_shift, template, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument From 6960b051299fb13c35930a3cae3d19fc5528ae20 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 28 Jan 2025 09:02:44 -0500 Subject: [PATCH 59/66] lint and update changelog Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 19 +++++++++++++++++-- modules/face/photomaker_pipeline.py | 3 +-- modules/sd_hijack_dynamic_atten.py | 1 - 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index acf49e095..c239895e3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,21 @@ # Change Log for SD.Next -## Update for 2025-01-28 +## Highlights for 2025-01-28 + +Two weeks since last release, time for update! + +*What's New?* +- New **Detailer** functionality including ability to use several new face-restore models +- Support for new models/pipelines: + face-swapper with **Photomaker-v2** and video with **Fast-Hunyuan** +- Support for several new optimizations and accelerations: + Many **IPEX** improvements, native *torch fp8* support, + support for **PAB:Pyramid-attention-broadcast**, **ParaAttention** and **PerFlow** +- Fully built-in both model **weights merge** as well as model **component merge** +- Improved remote access control and reliability as well as running inside containers +- And of course, hotfixes for all reported issues... + +## Details for 2025-01-28 - **Contributing**: - if you'd like to contribute, please see updated [contributing](https://github.com/vladmandic/automatic/blob/dev/CONTRIBUTING) guidelines @@ -60,7 +75,7 @@ - **upscale**: code refactor to unify latent, resize and model based upscalers - **loader**: ability to run in-memory models - **schedulers**: ability to create model-less schedulers - - **quantiation**: code refactor into dedicated module + - **quantization**: code refactor into dedicated module - **dynamic attention sdpa**: more correct implementation and new trigger rate control - **Remote access**: - perform auth check on ui startup diff --git a/modules/face/photomaker_pipeline.py b/modules/face/photomaker_pipeline.py index 657e96fc7..9aa5b0a69 100644 --- a/modules/face/photomaker_pipeline.py +++ b/modules/face/photomaker_pipeline.py @@ -373,8 +373,7 @@ class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline): f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" " the batch size of `prompt`." ) - else: - uncond_tokens = [negative_prompt, negative_prompt_2] + uncond_tokens = [negative_prompt, negative_prompt_2] negative_prompt_embeds_list = [] for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): # pylint: disable=redefined-argument-from-local diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index 9ea103036..bba5cbc2a 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -1,6 +1,5 @@ from functools import cache, wraps -import math import torch from diffusers.utils import USE_PEFT_BACKEND # pylint: disable=unused-import from modules import shared, devices From b141d266f61b715631b02250e040b08921c8d967 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 29 Jan 2025 00:31:00 +0300 Subject: [PATCH 60/66] Update dyn atten defaults --- modules/shared.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/shared.py b/modules/shared.py index fe77c5a45..6d4b8be78 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -513,8 +513,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "cross_attention_optimization": OptionInfo(startup_cross_attention, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention(native)}), "sdp_options": OptionInfo(startup_sdp_options, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'Sage attention'], "visible": native}), "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }), - "dynamic_attention_slice_rate": OptionInfo(2, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory, "step": 0.1, "visible": native}), - "dynamic_attention_trigger_rate": OptionInfo(3, "Dynamic Attention trigger rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory*2, "step": 0.1, "visible": native}), + "dynamic_attention_slice_rate": OptionInfo(0.5, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.01, "maximum": gpu_memory, "step": 0.01, "visible": native}), + "dynamic_attention_trigger_rate": OptionInfo(1, "Dynamic Attention trigger rate in GB", gr.Slider, {"minimum": 0.01, "maximum": gpu_memory*2, "step": 0.01, "visible": native}), "sub_quad_sep": OptionInfo("

Sub-quadratic options

", "", gr.HTML, {"visible": not native}), "sub_quad_q_chunk_size": OptionInfo(512, "Attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8, "visible": not native}), "sub_quad_kv_chunk_size": OptionInfo(512, "Attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8, "visible": not native}), From 50f73c9553aa1b6c8612511603338f4ee642b476 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 29 Jan 2025 00:43:22 +0300 Subject: [PATCH 61/66] Update ipex dyn atten defaults --- modules/intel/ipex/attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py index e8e5b815b..42961377a 100644 --- a/modules/intel/ipex/attention.py +++ b/modules/intel/ipex/attention.py @@ -7,8 +7,8 @@ from functools import cache, wraps # ARC GPUs can't allocate more than 4GB to a single block so we slice the attetion layers -sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 3)) -attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 2)) +sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 1)) +attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 0.5)) # Find something divisible with the input_tokens @cache From da6a1979ff16cae8bebc3475849ffcc4114aa2c9 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 28 Jan 2025 18:47:45 -0500 Subject: [PATCH 62/66] cleanup lora overrides Signed-off-by: Vladimir Mandic --- modules/lora/extra_networks_lora.py | 11 ++++++---- modules/lora/network_overrides.py | 31 +++++++++++++++-------------- scripts/xyz_grid_shared.py | 2 +- 3 files changed, 24 insertions(+), 20 deletions(-) diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py index 448b214ac..357c5291f 100644 --- a/modules/lora/extra_networks_lora.py +++ b/modules/lora/extra_networks_lora.py @@ -179,11 +179,14 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): networks.previously_loaded_networks = networks.loaded_networks.copy() debug_log(f'Load network: type=LoRA active={[n.name for n in networks.previously_loaded_networks]} deactivate') if shared.native and len(networks.diffuser_loaded) > 0: - if hasattr(shared.sd_model, "unload_lora_weights") and hasattr(shared.sd_model, "text_encoder"): - if not (shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled is True): + if not (shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled is True): + if hasattr(shared.sd_model, "unfuse_lora"): + try: + shared.sd_model.unfuse_lora() + except Exception: + pass + if hasattr(shared.sd_model, "unload_lora_weights"): try: - if shared.opts.lora_fuse_diffusers: - shared.sd_model.unfuse_lora() shared.sd_model.unload_lora_weights() # fails for non-CLIP models except Exception: pass diff --git a/modules/lora/network_overrides.py b/modules/lora/network_overrides.py index 81bd07ce1..65448ef2d 100644 --- a/modules/lora/network_overrides.py +++ b/modules/lora/network_overrides.py @@ -2,31 +2,32 @@ from modules import shared maybe_diffusers = [ # forced if lora_maybe_diffusers is enabled - 'aaebf6360f7d', # sd15-lcm - '3d18b05e4f56', # sdxl-lcm - 'b71dcb732467', # sdxl-tcd - '813ea5fb1c67', # sdxl-turbo + # 'aaebf6360f7d', # sd15-lcm + # '3d18b05e4f56', # sdxl-lcm + # 'b71dcb732467', # sdxl-tcd + # '813ea5fb1c67', # sdxl-turbo # not really needed, but just in case - '5a48ac366664', # hyper-sd15-1step - 'ee0ff23dcc42', # hyper-sd15-2step - 'e476eb1da5df', # hyper-sd15-4step - 'ecb844c3f3b0', # hyper-sd15-8step - '1ab289133ebb', # hyper-sd15-8step-cfg - '4f494295edb1', # hyper-sdxl-8step - 'ca14a8c621f8', # hyper-sdxl-8step-cfg - '1c88f7295856', # hyper-sdxl-4step - 'fdd5dcd1d88a', # hyper-sdxl-2step - '8cca3706050b', # hyper-sdxl-1step + # '5a48ac366664', # hyper-sd15-1step + # 'ee0ff23dcc42', # hyper-sd15-2step + # 'e476eb1da5df', # hyper-sd15-4step + # 'ecb844c3f3b0', # hyper-sd15-8step + # '1ab289133ebb', # hyper-sd15-8step-cfg + # '4f494295edb1', # hyper-sdxl-8step + # 'ca14a8c621f8', # hyper-sdxl-8step-cfg + # '1c88f7295856', # hyper-sdxl-4step + # 'fdd5dcd1d88a', # hyper-sdxl-2step + # '8cca3706050b', # hyper-sdxl-1step ] force_diffusers = [ # forced always '816d0eed49fd', # flash-sdxl 'c2ec22757b46', # flash-sd15 + '22c8339e7666', # spo-sdxl-10ep ] force_models = [ # forced always - 'sc', # 'sd3', + 'sc', 'kandinsky', 'hunyuandit', 'hunyuanvideo', diff --git a/scripts/xyz_grid_shared.py b/scripts/xyz_grid_shared.py index efd70724b..a3f7a3da9 100644 --- a/scripts/xyz_grid_shared.py +++ b/scripts/xyz_grid_shared.py @@ -213,7 +213,7 @@ def list_lora(): import sys lora = [v for k, v in sys.modules.items() if k == 'networks' or k == 'modules.lora.networks'][0] loras = [v.fullname for v in lora.available_networks.values()] - return ['None'] + loras + return ['None'] + sorted(loras) def apply_lora(p, x, xs): From 7e3a24b640b1c8c2d38ceca97ab37955c99c2959 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 28 Jan 2025 18:48:02 -0500 Subject: [PATCH 63/66] update wiki Signed-off-by: Vladimir Mandic --- wiki | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wiki b/wiki index 4b2c2d0d8..ba2f43a51 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 4b2c2d0d8e621dfd8c216124a020501b8f4b2528 +Subproject commit ba2f43a51370e98be99dadb9c2f19e3551f33cca From fc18cd1bf46ee8cfb43afd8cdf9726cdb135e972 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 29 Jan 2025 08:45:08 -0500 Subject: [PATCH 64/66] update changelog Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c239895e3..1474dce65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,19 +1,21 @@ # Change Log for SD.Next -## Highlights for 2025-01-28 +## Highlights for 2025-01-29 Two weeks since last release, time for update! -*What's New?* -- New **Detailer** functionality including ability to use several new face-restore models +*What's New?* +- New **Detailer** functionality including ability to use several new + face-restore models: *RestoreFormer, CodeFormer, GFPGan, GPEN-BFR* - Support for new models/pipelines: face-swapper with **Photomaker-v2** and video with **Fast-Hunyuan** - Support for several new optimizations and accelerations: Many **IPEX** improvements, native *torch fp8* support, support for **PAB:Pyramid-attention-broadcast**, **ParaAttention** and **PerFlow** -- Fully built-in both model **weights merge** as well as model **component merge** +- Fully built-in both model **merge weights** as well as model **merge component** + Finally replace that pesky VAE in your favorite model with a fixed one! - Improved remote access control and reliability as well as running inside containers -- And of course, hotfixes for all reported issues... +- And of course, hotfixes for all reported issues... ## Details for 2025-01-28 From bf310a3f427ce11b15aea8d7f2dde865dab69af0 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 29 Jan 2025 08:50:06 -0500 Subject: [PATCH 65/66] Update build_readme.yaml --- .github/workflows/build_readme.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/build_readme.yaml b/.github/workflows/build_readme.yaml index 7a685fd4e..d10a563ef 100644 --- a/.github/workflows/build_readme.yaml +++ b/.github/workflows/build_readme.yaml @@ -2,8 +2,6 @@ name: update-readme on: workflow_dispatch: - schedule: - - cron: '0 */4 * * *' jobs: deploy: From b2df5e45ac3e6f752a2b79d9713712dfd03c4db2 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 29 Jan 2025 08:55:26 -0500 Subject: [PATCH 66/66] update readme Signed-off-by: Vladimir Mandic --- README.md | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 4c2ef0ca2..c8f584161 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ - [Documentation](https://vladmandic.github.io/sdnext-docs/) - [SD.Next Features](#sdnext-features) -- [Model support](#model-support) and [Specifications]() +- [Model support](#model-support) - [Platform support](#platform-support) - [Getting started](#getting-started) @@ -32,7 +32,7 @@ All individual features are not listed here, instead check [ChangeLog](CHANGELOG ▹ **Windows | Linux | MacOS | nVidia | AMD | IntelArc/IPEX | DirectML | OpenVINO | ONNX+Olive | ZLUDA** - Platform specific autodetection and tuning performed on install - Optimized processing with latest `torch` developments with built-in support for model compile, quantize and compress - Compile backends: *Triton | StableFast | DeepCache | OneDiff* + Compile backends: *Triton | StableFast | DeepCache | OneDiff | TeaCache | etc.* Quantization and compression methods: *BitsAndBytes | TorchAO | Optimum-Quanto | NNCF* - Built-in queue management - Built in installer with automatic updates and dependency management @@ -82,6 +82,11 @@ SD.Next supports broad range of models: [supported models](https://vladmandic.gi > [!WARNING] > If you run into issues, check out [troubleshooting](https://vladmandic.github.io/sdnext-docs/Troubleshooting/) and [debugging](https://vladmandic.github.io/sdnext-docs/Debug/) guides +### Contributing + +Please see [Contributing](CONTRIBUTING) for details on how to contribute to this project +And for any question, reach out on [Discord](https://discord.gg/VjvR2tabEX) or open an [issue](https://github.com/vladmandic/automatic/issues) or [discussion](https://github.com/vladmandic/automatic/discussions) + ### Credits - Main credit goes to [Automatic1111 WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) for the original codebase @@ -104,10 +109,4 @@ SD.Next supports broad range of models: [supported models](https://vladmandic.gi If you're unsure how to use a feature, best place to start is [Docs](https://vladmandic.github.io/sdnext-docs/) and if its not there, check [ChangeLog](https://vladmandic.github.io/sdnext-docs/CHANGELOG/) for when feature was first introduced as it will always have a short note on how to use it -### Sponsors - -
-Allan Granta.v.mantzarisSML (See-ming Lee) -
-