Merge pull request #3728 from vladmandic/dev

merge dev to master

.github/workflows/build_readme.yaml

@@ -2,8 +2,6 @@ name: update-readme
 
 on:
   workflow_dispatch:
-  schedule:
-  - cron: '0 */4 * * *'
 
 jobs:
   deploy:

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
 

CONTRIBUTING

@@ -4,17 +4,24 @@ Pull requests from everyone are welcome
 
 Procedure for contributing:
 
+- Select SD.Next `dev` branch:  
+  <https://github.com/vladmandic/automatic/tree/dev>  
 - 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/<your-username>/<your-fork>
+  `git clone https://github.com/<your-username>/<your-fork>`  
 - Make your changes  
-- Test your changes
-- Test your changes against code guidelines  
-  - `ruff check`
-  - `pylint <folder>/<filename>.py`
+- Test your changes  
+- Lint your changes against code guidelines  
+  - `ruff check`  
+  - `pylint <folder>/<filename>.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  

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
-
-<div align="center">
-<!-- sponsors --><a href="https://github.com/allangrant"><img src="https://github.com/allangrant.png" width="60px" alt="Allan Grant" /></a><a href="https://github.com/mantzaris"><img src="https://github.com/mantzaris.png" width="60px" alt="a.v.mantzaris" /></a><a href="https://github.com/CurseWave"><img src="https://github.com/CurseWave.png" width="60px" alt="" /></a><a href="https://github.com/smlbiobot"><img src="https://github.com/smlbiobot.png" width="60px" alt="SML (See-ming Lee)" /></a><!-- sponsors -->
-</div>
-
 <br>

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)

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; }

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;

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;

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 {

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;

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;

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 = `
+  <div id="loginDiv" style="margin: 15% auto; max-width: 200px; padding: 2em; background: var(--background-fill-secondary);">
+    <h2>Login</h2> 
+    <label for="username" style="margin-top: 0.5em">Username</label>
+    <input type="text" id="loginUsername" name="username" style="width: 92%; padding: 0.5em; margin-top: 0.5em">
+    <label for="password" style="margin-top: 0.5em">Password</label>
+    <input type="text" id="loginPassword" name="password" style="width: 92%; padding: 0.5em; margin-top: 0.5em">
+    <div id="loginStatus" style="margin-top: 0.5em"></div>
+    <button type="submit" style="width: 100%; padding: 0.5em; margin-top: 0.5em; background: var(--button-primary-background-fill); color: var(--button-primary-text-color); border: var(--button-primary-border-color);">Login</button>
+  </div>
+`;
+
+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;

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); }

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;

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;

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):

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

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 <b>{ver['hash']} {ver['updated']}</b> <span style='color: var(--primary-500)'>{ver['url'].split('/')[-1]}</span><br>"
+    if not shared.native:
+        motd += "<span style='color: orange'>Legacy mode</span><br>"
     if shared.opts.motd:
         try:
             res = requests.get('https://vladmandic.github.io/automatic/motd', timeout=3)

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

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(

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)

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}')

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 + '<br>'
+        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}<br>"
+    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

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("<span>&nbsp Face: Multiple ID Transfers</span><br>")
         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('<a href="https://github.com/somanchiu/ReSwapper" target="_blank">&nbsp ReSwapper</a><br>')
+            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('<a href="https://huggingface.co/h94/IP-Adapter-FaceID" target="_blank">&nbsp Tencent AI Lab IP-Adapter FaceID</a><br>')
@@ -77,6 +88,7 @@ class Script(scripts.Script):
             with gr.Row():
                 gr.HTML('<a href="https://photo-maker.github.io/" target="_blank">&nbsp Tenecent ARC Lab PhotoMaker</a><br>')
             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]

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()
 

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

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

modules/face/photomaker_model.py

@@ -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)

modules/face/photomaker_model_v1.py

@@ -0,0 +1,107 @@
+### original <https://github.com/TencentARC/PhotoMaker/blob/main/photomaker/model.py>
+
+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

modules/face/photomaker_model_v2.py

@@ -0,0 +1,337 @@
+### original <https://github.com/TencentARC/PhotoMaker/blob/main/photomaker/model_v2.py>
+
+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

modules/face/photomaker_pipeline.py

@@ -0,0 +1,884 @@
+### original <https://github.com/TencentARC/PhotoMaker/blob/main/photomaker/pipeline.py>
+
+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)

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

modules/face/reswapper_model.py

@@ -0,0 +1,127 @@
+# original: <https://github.com/somanchiu/ReSwapper/blob/main/StyleTransferModel_128.py>
+
+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

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)

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))

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)

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

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

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

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

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":

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

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

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',

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

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

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 + '<br>'
+    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')

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

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)):

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:

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]

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)

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

modules/perflow/__init__.py

@@ -0,0 +1,4 @@
+### original: <https://github.com/magic-research/piecewise-rectified-flow>
+
+from .scheduler_perflow import PeRFlowScheduler
+from .utils_perflow import merge_delta_weights_into_unet

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

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)<len(self.timesteps) else 0
+        dt = (prev_step - timestep) / self.config.num_train_timesteps # pylint: disable=no-member
+        dt = dt.to(sample.device, sample.dtype)
+
+        prev_sample = sample + dt * pred_velocity
+
+        if not return_dict:
+            return (prev_sample,)
+        return PeRFlowSchedulerOutput(prev_sample=prev_sample, pred_original_sample=None)
+
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> 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

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

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

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.

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

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,

modules/postprocess/yolo.py

@@ -9,13 +9,17 @@ from modules import shared, processing, devices, processing_class, ui_common
 from modules.detailer import Detailer
 
 
-PREDEFINED = [ # <https://huggingface.co/vladmandic/yolo-detailers/tree/main>
+predefined = [ # <https://huggingface.co/vladmandic/yolo-detailers/tree/main>
     '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):

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)

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):

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":

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

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

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')

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
 
 

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

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):

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)
 

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()}')

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

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

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)

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

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

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("<h2>Execution Precision</h2>", "", 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("<h3>Sub-quadratic options</h3>", "", 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("<h2>BitsAndBytes</h2>", "", gr.HTML),
+    "bnb_quantization_sep": OptionInfo("<h2>BitsAndBytes</h2>", "", 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("<h2>Optimum Quanto</h2>", "", 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("<h2>TorchAO</h2>", "", 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("<h2>NNCF</h2>", "", gr.HTML),
+
+    "nncf_compress_sep": OptionInfo("<h2>NNCF</h2>", "", 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("<h2>Layerwise Casting</h2>", "", 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("<h2>Perturbed-Attention Guidance</h2>", "", gr.HTML),
+    "pag_sep": OptionInfo("<h2>PAG: Perturbed attention guidance</h2>", "", gr.HTML),
     "pag_apply_layers": OptionInfo("m0", "PAG layer names"),
 
+    "pab_sep": OptionInfo("<h2>PAB: Pyramid attention broadcast </h2>", "", 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("<h2>Para-attention</h2>", "", 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("<h2>HyperTile</h2>", "", 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("<h2>Models</h2>", "", gr.HTML),
-    "extra_network_reference": OptionInfo(False, "Use reference values when available", gr.Checkbox),
+    "extra_networks_model_sep": OptionInfo("<h2>Rerefence models</h2>", "", 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("<h2>LoRA</h2>", "", 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

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))

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))

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:

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'<script type="text/javascript" src="{webpath(script_js)}"></script>\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'<script type="text/javascript" src="{webpath(script.path)}"></script>\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'<script type="text/javascript">{inline}</script>\n'
+    return js
+
+
 def html_css(css: str):
     def stylesheet(fn):
         return f'<link rel="stylesheet" property="stylesheet" href="{webpath(fn)}">'
@@ -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 = '<title>SD.Next</title>'
     manifest = f'<link rel="manifest" href="{webpath(os.path.join(script_path, "html", "manifest.json"))}">'
+    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'<head>', f'<head>{title}'.encode("utf8"))
-        res.body = res.body.replace(b'</head>', f'{head}</head>'.encode("utf8"))
         res.body = res.body.replace(b'</head>', f'{manifest}</head>'.encode("utf8"))
+        res.body = res.body.replace(b'</head>', f'{login}</head>'.encode("utf8"))
+        res.body = res.body.replace(b'</head>', f'{js}</head>'.encode("utf8"))
         res.body = res.body.replace(b'</body>', f'{css}{body}</body>'.encode("utf8"))
         lines = res.body.decode("utf8").split('\n')
         for line in lines:

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<br><span style="color: var(--body-text-color-subdued)">Specify the components to include<br>Paths can be relative or absolute</span><br>')
+                    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<br>')
+                    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<br>')
+                    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('<br>')
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        gr.HTML('Model metadata<br>')
+                    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('<h2>Search for models</h2>Select a model from the search results to download<br><br>')
                     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)

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')

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:

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

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()

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('<a href="https://huggingface.co/tencent/HunyuanVideo">&nbsp Hunyuan Video</a><br>')
+        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)

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'

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):

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")),

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):

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):

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

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