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: diff --git a/CHANGELOG.md b/CHANGELOG.md index b333b2e24..1474dce65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,15 +1,111 @@ # Change Log for SD.Next -## Update for 2025-01-16 +## Highlights for 2025-01-29 -- **Gallery**: - - add http fallback for slow/unreliable links -- **Fixes**: +Two weeks since last release, time for update! + +*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 **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... + +## 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 +- **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* +- **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 +- **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 +- [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 + - 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, fp64 emulation, jit cache + - use the us server by default on linux + - use pytorch test branch on windows + - extend the supported python versions + - 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` 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) + - 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 + - **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 + - **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 + - 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 - sana vae tiling - increase gallery timeouts - update ui element ids + - modernui use local font + - 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 + - image metadata civitai compatibility + - 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 + - port streaming model load to diffusers ## Update for 2025-01-15 diff --git a/CONTRIBUTING b/CONTRIBUTING index 9dc17f199..f6187fa17 100644 --- a/CONTRIBUTING +++ b/CONTRIBUTING @@ -4,17 +4,24 @@ 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 +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/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) -
-
diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index aafc660ba..8c7edb3be 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit aafc660ba3cf78724059917f7c3a05f3ac1ed46a +Subproject commit 8c7edb3be11b8b8c2d2dcd0421e93345bd20fcae diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index c17c53928..e0b3e5918 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit c17c5392816e3e6e212c7d10166260728e10d249 +Subproject commit e0b3e5918f49d090c8c340ca44e2c9e6d5bf5127 diff --git a/installer.py b/installer.py index 1e62ddd2a..d189a835a 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(), } @@ -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 = '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 '' @@ -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: @@ -659,22 +662,39 @@ 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') + # 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') + 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: + # XPU PyTorch doesn't support Flash Atten or Memory Atten yet so Battlemage 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') + # 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', '--pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/xpu') # torchvision doesn't exist on test/stable branch for windows - install(os.environ.get('OPENVINO_PACKAGE', 'openvino==2024.5.0'), 'openvino', ignore=True) + 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) ts('ipex', t_start) @@ -683,6 +703,8 @@ def install_ipex(torch_command): def install_openvino(torch_command): t_start = time.time() + # 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': @@ -726,11 +748,22 @@ 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) +# 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() @@ -842,6 +875,7 @@ def check_torch(): return if not args.skip_all: install_torch_addons() + check_cudnn() if args.profile: pr.disable() print_profile(pr, 'Torch') @@ -1056,7 +1090,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) diff --git a/javascript/base.css b/javascript/base.css index cadfd1868..8f89685c2 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; } @@ -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/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/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/javascript/sdnext.css b/javascript/sdnext.css index 139b7d1bb..c3deebad7 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; @@ -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/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/launch.py b/launch.py index 07fa93697..c80840036 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') @@ -150,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/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/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) 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/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/devices.py b/modules/devices.py index c98aea696..874b24de8 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 { @@ -415,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/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/face/__init__.py b/modules/face/__init__.py index c18da6e2e..6835bda17 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
') @@ -77,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) @@ -85,11 +97,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_model, 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_model, pm_trigger, pm_strength, pm_start, fs_cache): # pylint: disable=arguments-differ, unused-argument if not shared.native: return None if mode == 'None': @@ -119,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') @@ -134,11 +148,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..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.diffusers_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..9aa5b0a69 --- /dev/null +++ b/modules/face/photomaker_pipeline.py @@ -0,0 +1,884 @@ +### 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): # pylint: disable=redefined-argument-from-local + 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)}." + ) + 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" + " the batch size of `prompt`." + ) + 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 + 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 # 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))) + 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/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/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/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/attention.py b/modules/intel/ipex/attention.py index 1618045b6..42961377a 100644 --- a/modules/intel/ipex/attention.py +++ b/modules/intel/ipex/attention.py @@ -1,181 +1,120 @@ 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 # 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)) -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', 1)) +attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 0.5)) # 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 - break - return slice_size +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: + return 1 + return split_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, query_element_size, slice_rate=2, trigger_rate=3): + batch_size, attn_heads, query_len, _ = query_shape + _, _, key_len, _ = key_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 * (query_len * key_len) * query_element_size / 1024 / 1024 / 1024 - 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 >= trigger_rate: + do_batch_split = True + split_batch_size = find_split_size(batch_size, slice_batch_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 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(attn_heads, slice_head_size, slice_rate=slice_rate) -# 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_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_split_size(query_len, slice_query_size, slice_rate=slice_rate) - 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, query.element_size(), slice_rate=attention_slice_rate, trigger_rate=sdpa_slice_trigger_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] = 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 ) 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/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py index 5bf5bbe39..413f7d5eb 100644 --- a/modules/intel/ipex/diffusers.py +++ b/modules/intel/ipex/diffusers.py @@ -1,14 +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 -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)) - # Diffusers FreeU # Diffusers is imported before ipex hijacks so fourier_filter needs hijacking too @@ -47,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(): +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 device_supports_fp64 or os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is not None: - 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 b1c9a1182..4f3a03b4e 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,26 +85,15 @@ 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: - original_torch_bmm = torch.bmm +if can_allocate_plus_4gb: 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: @@ -105,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) @@ -190,6 +197,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 +321,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 +359,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 diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index 1c89ceb48..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)) @@ -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": @@ -279,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)) @@ -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/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 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/modules/memstats.py b/modules/memstats.py index d43e5bbfa..1ff7f721f 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -1,3 +1,5 @@ +import re +import sys import os import psutil import torch @@ -5,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() @@ -18,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: @@ -52,7 +77,53 @@ 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: 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/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..ec994d567 --- /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 + 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 + 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 + 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 + 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 + 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 + 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 + 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/model_quant.py b/modules/model_quant.py index d570c6ced..6798160a8 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -1,11 +1,35 @@ 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): + 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): @@ -70,10 +94,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 @@ -93,38 +120,275 @@ 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 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 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 + + +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}') + + +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.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(".") + 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.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(".") + 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/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/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/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) 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/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..bffdd1494 --- /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( # pylint: disable=no-member + 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( # pylint: disable=no-member + 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/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/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/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/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/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/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_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 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 4d988e73b..46bf0ad26 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 = [] @@ -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_checkpoint.py b/modules/sd_checkpoint.py index ec41a130f..92c83e953 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 @@ -49,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}") @@ -170,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 @@ -289,6 +297,20 @@ 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] + 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: @@ -309,10 +331,13 @@ 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(): + if k == 'modelspec.thumbnail' and v.startswith("data:"): + extract_thumbnail(filename, v) if v.startswith("data:"): v = 'data' if k == 'format' and v == 'pt': @@ -332,6 +357,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/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index de39a966f..bba5cbc2a 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -5,17 +5,118 @@ 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_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 - break - return slice_size +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: + return 1 + return split_size + + +# Find slice sizes for SDPA +@cache +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 + + 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 + split_query_size = query_len + + do_batch_split = False + do_head_split = False + do_query_split = False + + if batch_size * slice_batch_size >= trigger_rate: + do_batch_split = True + 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 * (query_len * key_len) * query_element_size / 1024 / 1024 / 1024 + do_head_split = True + split_head_size = find_split_size(attn_heads, slice_head_size, slice_rate=slice_rate) + + 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_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 + + +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, 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: + 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 devices.backend != "directml": + getattr(torch, query.device.type).synchronize() + else: + 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 + @cache -def find_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 @@ -33,75 +134,20 @@ def find_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_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 +197,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, 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/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/sd_models.py b/modules/sd_models.py index 9d0349a3f..a3d749be8 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -9,14 +9,14 @@ 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 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') @@ -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 @@ -437,6 +439,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 +534,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' @@ -563,12 +569,15 @@ 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") 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) @@ -1043,27 +1052,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()}') diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 0972c3350..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 @@ -25,90 +24,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 +51,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 +60,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 +199,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 +258,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 diff --git a/modules/sd_offload.py b/modules/sd_offload.py index f9d01528c..94a75496e 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -3,14 +3,14 @@ import sys import time import inspect import torch -import accelerate - -from modules import shared, devices, errors +import diffusers +import accelerate.hooks +from modules import shared, devices, errors, model_quant 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 @@ -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,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/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..f24641d27 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, @@ -54,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: @@ -100,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 }, @@ -160,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, [], {}), @@ -172,14 +169,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/shared.py b/modules/shared.py index b092c27b5..6d4b8be78 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -7,15 +7,12 @@ 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 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 +23,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 +45,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 +61,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 @@ -223,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: @@ -325,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]) @@ -357,7 +338,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 = [] @@ -518,9 +499,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}), @@ -532,7 +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(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory, "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}), @@ -578,25 +560,35 @@ 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}), + "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}), + "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(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}), + "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_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}), })) options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { @@ -612,9 +604,19 @@ 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"), + "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}), + "hypertile_sep": OptionInfo("

HyperTile

", "", gr.HTML), "hypertile_unet_enabled": OptionInfo(False, "UNet Enabled"), "hypertile_hires_only": OptionInfo(False, "HiRes pass only"), @@ -915,8 +917,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}), @@ -1190,11 +1193,11 @@ 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) -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/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: 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) 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..95f2acac2 --- /dev/null +++ b/modules/upscaler_simple.py @@ -0,0 +1,121 @@ +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 + elif selected_model == "Resize Nearest": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.NEAREST) + elif selected_model == "Resize Lanczos": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.LANCZOS) + elif selected_model == "Resize Bicubic": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BICUBIC) + elif selected_model == "Resize Bilinear": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BILINEAR) + elif selected_model == "Resize Hamming": + return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.HAMMING) + 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, _): + 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) + + +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 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() diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py index aac3eb8f0..d67fd7d31 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. @@ -16,6 +15,13 @@ 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 }, +} +loaded_model = 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,28 +69,26 @@ class Script(scripts.Script): def ui(self, is_img2img): with gr.Row(): gr.HTML('  Hunyuan Video
') + with gr.Row(): + 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) 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 [model, 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 - if shared.sd_model.__class__ != diffusers.HunyuanVideoPipeline: + def load(self, model:str): + 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 = {} @@ -96,36 +100,44 @@ 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 ) + 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( - repo_id, + pretrained_model_name_or_path='tencent/HunyuanVideo', transformer=transformer, text_encoder=text_encoder, - revision="refs/pr/18", + text_encoder_2=text_encoder_2, + revision='refs/pr/18', 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 @@ -133,13 +145,30 @@ 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 + # 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(model) 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) 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/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_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/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): 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): diff --git a/webui.py b/webui.py index f104bbee8..d036e6ad3 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 @@ -208,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: @@ -283,6 +305,16 @@ 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: + 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: + 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 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 7f072b554..ba2f43a51 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 7f072b554c6ee2edadc33879e5b4bdbfa48e6282 +Subproject commit ba2f43a51370e98be99dadb9c2f19e3551f33cca