Merge pull request #3957 from vladmandic/dev

merge dev
2026-09-18 00:34:33 +02:00 · 2025-06-02 15:52:22 +02:00
parent 12ebadccd4 7d05bed459
commit 3ca9f29c8f
35 changed files with 1512 additions and 1159 deletions
@@ -1,5 +1,56 @@
 # Change Log for SD.Next

+## Update for 2025-06-02
+
+### Highlights for 2025-06-02
+
+This release is all about quantization: with new SD.Next own quantization method: **SDNQ**  
+**SDNQ** is based on **NNCF**, but has been re-implemented, optimized and evolved enough to become its own quantization method!  
+It's fully cross-platform, supports all GPUs and includes tons of quantization methods:
+- *8-bit, 6-bit, 4-bit, 2-bit and 1-bit int and uint*
+- *8-bit e5, e4 and fnuz float*
+
+Also unlike most traditional methods, its also applicable to nearly all model types  
+
+*Hint*: Even if you may not need quantization for your current model, it may be worth trying it out as it can significantly improve performance or capabilities of your existing workflow! For example, you may not have issues with SD15 or SDXL, but you may have been limited running at high resolutions or with multiple ControlNet due to VRAM requirements - this will significantly reduce memory requirements. And on-the-fly quantization takes just few seconds during model load, there is no need to have multiple quant models permanently saved.  
+
+On a different topic, **SD.Next Wiki & Docs** and its **UI Hints** and **UI Localization** system are community efforts and any contributions are welcome!  
+You dont need any coding experience, but if you learned something and you find documentation either wrong or insufficient, please do suggest edits!  
+Take a look at [Docs](https://github.com/vladmandic/sdnext/wiki/Docs), [Hints](https://github.com/vladmandic/sdnext/wiki/Hints) and [Localization](https://github.com/vladmandic/sdnext/wiki/Locale) contribution guides
+
+[ReadMe](https://github.com/vladmandic/automatic/blob/master/README.md) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867)
+
+### Details for 2025-06-02
+
+- **SDNQ Quantization**  
+  - Renamed `NNCF` to `SDNQ`  
+  - Renamed quantization scheme names to the underlying dtype names instead of NNCF names  
+    - `INT8_SYM` -> `int8`  
+    - `INT8` -> `uint8`  
+    - `INT4_SYM` -> `int4`  
+    - `INT4` -> `uint4`  
+  - Add `float8_e4m3fn`, `float8_e5m2`, `float8_e4m3fnuz`, `float8_e5m2fnuz`, `int6`, `uint6`, `int2`, `uint2` and `uint1` support  
+  - Add quantized matmul support for `float8_e4m3fn` and `float8_e5m2`  
+  - Set the default quant mode to `pre`  
+  - Use per token input quant with int8 and fp8 quantized matmul  
+  - Implement better layer hijacks  
+  - Add an option to toggle quantize with GPU  
+  - Fix conv quant and add support for conv quant with asym modes  
+  - Fix lora weight change  
+  - Fix high RAM usage with pre mode  
+  - Fix scale and zero_point not being offloaded  
+- **IPEX**  
+  - Disabe Dynamic Attention by default on PyTorch 2.7  
+  - Remove GradScaler hijack and use `torch.amp.GradScaler` instead  
+- **Feature**  
+  - TeaCache support for HiDream I1  
+- **Changes**  
+  - Set the default attention optimizer to Scaled-Dot-Product on all backends  
+  - Enable Dynamic attention for Scaled-Dot-Product with ROCm, DirectML, MPS and CPU backends  
+- **Fixes**
+  - Gallery duplicate entries  
+  - Prompt enhancement args mismatch  
+
 ## Update for 2025-05-17

 *Curious how your system is performing?*  
@@ -163,7 +214,7 @@ What else?
    implemented as an extension for **SD.Next** (for the moment while dev is ongoing)  
    generate high-quality videos with pretty much unlimited duration and with limited VRAM!  
    install as any other extension and for details see extension [README](https://github.com/vladmandic/sd-extension-framepack/blob/main/README.md)  
-    - I2V & FLF2V support with explicit strength controls   
+    - I2V & FLF2V support with explicit strength controls  
    - complex actions: modify prompts for each section of the video  
    - LoRA support: use normal **HunyuanVideo** LoRAs  
    - decode: use local, tiny or remote VAE  
@@ -275,7 +326,7 @@ There are quite a few other performance and quality-of-life improvements in this

 - **Models**  
  - [HiDream-I1](https://huggingface.co/HiDream-ai/HiDream-I1-Full) in fast, dev and full variants!  
-    new absolutely massive image generative foundation model with **17B** parameters and 4 text-encoders with additional **8.3B** parameters   
+    new absolutely massive image generative foundation model with **17B** parameters and 4 text-encoders with additional **8.3B** parameters  
    simply select from *networks -> models -> reference*  
    due to size (over 25B params in 58GB), offloading and on-the-fly quantization are pretty much a necessity  
    see [HiDream Wiki page](https://github.com/vladmandic/sdnext/wiki/HiDream) for details  
@@ -334,7 +385,7 @@ Time for another major release with ~120 commits and [ChangeLog](https://github.
 *Highlights?*  
 Video...Brand new Video processing module with support for all latest models: **WAN21, Hunyuan, LTX, Cog, Allegro, Mochi1, Latte1** in both *T2V* and *I2V* workflows  
 And combined with *on-the-fly quantization*, support for *Local/Tiny/Remote* VAE, acceleration modules such as *FasterCache or PAB*, and more!  
-Models...And support for new models: **CogView-4**, **SANA 1.5**, 
+Models...And support for new models: **CogView-4**, **SANA 1.5**,  

 *Plus...*  
 - New **Prompt Enhance** using LLM,
@@ -633,7 +684,7 @@ We're back with another update with nearly 100 commits!
  - updated **CUDA** receipe to `torch==2.6.0` with `cuda==12.6` and add prebuilt image  
  - added **ROCm** receipe and prebuilt image  
  - added **IPEX** receipe and add prebuilt image  
-  - added **OpenVINO** receipe and prebuilt image   
+  - added **OpenVINO** receipe and prebuilt image  
 - **System**  
  - improve **python==3.12** compatibility  
  - **Torch**  
@@ -706,7 +757,7 @@ Just one week after latest release and what a week it was with over 50 commits!

 - **GitHub**
  - rename core repo from <https://github.com/vladmandic/automatic> to <https://github.com/vladmandic/sdnext>  
-    old repo url should automatically redirect to new one for seamless transition and in-place upgrades   
+    old repo url should automatically redirect to new one for seamless transition and in-place upgrades  
    all internal links have been updated  
    wiki content and docs site have been updated  
 - **Docs**:
@@ -1055,7 +1106,7 @@ We've also added support for several new models such as highly anticipated [NVLa
 And several new SOTA video models: [Lightricks LTX-Video](https://huggingface.co/Lightricks/LTX-Video), [Hunyuan Video](https://huggingface.co/tencent/HunyuanVideo) and [Genmo Mochi.1 Preview](https://huggingface.co/genmo/mochi-1-preview)  

 And a lot of **Control** and **IPAdapter** goodies  
- for **SDXL** there is new [ProMax](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0), improved *Union* and *Tiling* models 
+- for **SDXL** there is new [ProMax](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0), improved *Union* and *Tiling* models  
 - for **FLUX.1** there are [Flux Tools](https://blackforestlabs.ai/flux-1-tools/) as well as official *Canny* and *Depth* models,  
  a cool [Redux](https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev) model as well as [XLabs](https://huggingface.co/XLabs-AI/flux-ip-adapter-v2) IP-adapter
 - for **SD3.5** there are official *Canny*, *Blur* and *Depth* models in addition to existing 3rd party models  
@@ -31,11 +31,11 @@ All individual features are not listed here, instead check [ChangeLog](CHANGELOG
 - Multiple [diffusion models](https://vladmandic.github.io/sdnext-docs/Model-Support/)!  
 - Built-in Control for Text, Image, Batch and Video processing!  
 - Multiplatform!  
- ▹ **Windows | Linux | MacOS | nVidia CUDA | AMD ROCm | IntelArc/IPEX | DirectML | OpenVINO | ONNX+Olive | ZLUDA**
+ ▹ **Windows | Linux | MacOS | nVidia CUDA | AMD ROCm | Intel Arc / IPEX XPU | 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  
+- Optimized processing with latest `torch` developments with built-in support for model compile and quantize  
  Compile backends: *Triton | StableFast | DeepCache | OneDiff | TeaCache | etc.*  
-  Quantization and compression methods: *BitsAndBytes | TorchAO | Optimum-Quanto | NNCF*  
+  Quantization methods: *SDNQ | BitsAndBytes | Optimum-Quanto | TorchAO*  
 - **Interrogate/Captioning** with 150+ **OpenCLiP** models and 20+ built-in **VLMs**  
 - Built-in queue management  
 - Built in installer with automatic updates and dependency management  
@@ -546,7 +546,7 @@ def check_diffusers():
    t_start = time.time()
    if args.skip_all or args.skip_git or args.experimental:
        return
-    sha = '20379d9d1395b8e95977faf80facff43065ba75f' # diffusers commit hash
+    sha = '6508da6f06a0da1054ae6a808d0025c04b70f0e8' # 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 ''
@@ -727,10 +727,6 @@ def install_ipex(torch_command):
    if os.environ.get("PYTORCH_ENABLE_XPU_FALLBACK", None) is None:
        os.environ.setdefault('PYTORCH_ENABLE_XPU_FALLBACK', '1') # CPU fallback for unsupported ops

-    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')
-
    # FP64 emulation causes random UR Errors
    #if os.environ.get("OverrideDefaultFP64Settings", None) is None:
    #    os.environ.setdefault('OverrideDefaultFP64Settings', '1')
@@ -787,8 +783,6 @@ def install_torch_addons():
        install('DeepCache')
    if opts.get('cuda_compile_backend', '') == 'olive-ai':
        install('olive-ai')
-    if opts.get('nncf_compress_weights', False) and not args.use_openvino:
-        install('nncf==2.16.0', 'nncf')
    if opts.get('optimum_quanto_weights', False):
        install('optimum-quanto==0.2.7', 'optimum-quanto')
    if opts.get('torchao_quantization', False):
@@ -1175,7 +1169,6 @@ def install_optional():
    install('albumentations==1.4.3', ignore=True)
    install('pydantic==1.10.21', ignore=True)
    reload('pydantic', '1.10.21')
-    install('nncf==2.16.0', ignore=True)
    install('gguf', ignore=True)
    install('av', ignore=True)
    try:
@@ -1392,6 +1385,7 @@ def check_version(offline=False, reset=True): # pylint: disable=unused-argument
        args.skip_git = True # pylint: disable=attribute-defined-outside-init
    ver = get_version()
    log.info(f'Version: {print_dict(ver)}')
+    branch_name = ver['branch'] if ver is not None and 'branch' in ver else 'master'
    if args.version or args.skip_git:
        return
    check_ui(ver)
@@ -1406,30 +1400,29 @@ def check_version(offline=False, reset=True): # pylint: disable=unused-argument
        return
    commits = None
    try:
-        commits = requests.get('https://api.github.com/repos/vladmandic/sdnext/branches/master', timeout=10).json()
-        if commits['commit']['sha'] != commit:
-            if args.upgrade:
-                global quick_allowed # pylint: disable=global-statement
-                quick_allowed = False
-                log.info('Updating main repository')
-                try:
-                    git('add .')
-                    git('stash')
-                    update('.', keep_branch=True)
-                    # git('git stash pop')
-                    ver = git('log -1 --pretty=format:"%h %ad"')
-                    log.info(f'Repository upgraded: {ver}')
-                except Exception:
-                    if not reset:
-                        log.error('Repository error upgrading')
-                    else:
-                        log.warning('Repository: retrying upgrade...')
-                        git_reset()
-                        check_version(offline=offline, reset=False)
-            else:
-                dt = commits["commit"]["commit"]["author"]["date"]
-                commit = commits["commit"]["sha"][:8]
-                log.info(f'Version: check latest available hash={commit} updated={dt}')
+        commits = requests.get(f'https://api.github.com/repos/vladmandic/sdnext/branches/{branch_name}', timeout=10).json()
+        if commits['commit']['sha'] != commit and args.upgrade:
+            global quick_allowed # pylint: disable=global-statement
+            quick_allowed = False
+            log.info('Updating main repository')
+            try:
+                git('add .')
+                git('stash')
+                update('.', keep_branch=True)
+                # git('git stash pop')
+                ver = git('log -1 --pretty=format:"%h %ad"')
+                log.info(f'Repository upgraded: {ver}')
+            except Exception:
+                if not reset:
+                    log.error('Repository error upgrading')
+                else:
+                    log.warning('Repository: retrying upgrade...')
+                    git_reset()
+                    check_version(offline=offline, reset=False)
+        else:
+            dt = commits["commit"]["commit"]["author"]["date"]
+            commit = commits["commit"]["sha"][:8]
+            log.info(f'Version: app=sd.next latest={dt} hash={commit} branch={branch_name}')
    except Exception as e:
        log.error(f'Repository failed to check version: {e} {commits}')
    ts('latest', t_start)
@@ -12,11 +12,11 @@ const getENActiveTab = () => {
  if (gradioApp().getElementById('video_prompt')?.checkVisibility()) return 'video';
  if (gradioApp().getElementById('framepack_prompt_row')?.checkVisibility()) return 'framepack';
  // legacy method
-  if (gradioApp().getElementById('tab_txt2img').style.display === 'block') tabName = 'txt2img';
-  else if (gradioApp().getElementById('tab_img2img').style.display === 'block') tabName = 'img2img';
-  else if (gradioApp().getElementById('tab_control').style.display === 'block') tabName = 'control';
-  else if (gradioApp().getElementById('tab_video').style.display === 'block') tabName = 'video';
-  else if (gradioApp().getElementById('tab_framepack_tab').style.display === 'block') tabName = 'framepack';
+  if (gradioApp().getElementById('tab_txt2img')?.style.display === 'block') tabName = 'txt2img';
+  else if (gradioApp().getElementById('tab_img2img')?.style.display === 'block') tabName = 'img2img';
+  else if (gradioApp().getElementById('tab_control')?.style.display === 'block') tabName = 'control';
+  else if (gradioApp().getElementById('tab_video')?.style.display === 'block') tabName = 'video';
+  else if (gradioApp().getElementById('tab_framepack_tab')?.style.display === 'block') tabName = 'framepack';
  // log('getENActiveTab', tabName);
  return tabName;
 };
@@ -5,6 +5,8 @@ let url;
 let currentImage;
 let pruneImagesTimer;
 let outstanding = 0;
+let lastSort = 0;
+let lastSortName = 'none';
 const el = {
  folders: undefined,
  files: undefined,
@@ -124,9 +126,14 @@ class GalleryFile extends HTMLElement {
  }

  async connectedCallback() {
-    if (this.shadow.children.length > 0) return;
+    if (this.shadow.children.length > 0) {
+      return;
+    }
    const ext = this.name.split('.').pop().toLowerCase();
-    if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'jxl', 'svg', 'mp4'].includes(ext)) return;
+    if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'jxl', 'svg', 'mp4'].includes(ext)) {
+      console.error(`gallery: type=${ext} file=${this.name} unsupported`);
+      return;
+    }
    this.hash = await getHash(`${this.folder}/${this.name}/${this.size}/${this.mtime}`); // eslint-disable-line no-use-before-define
    const style = document.createElement('style');
    const width = opts.browser_fixed_width ? `${opts.extra_networks_card_size}px` : 'unset';
@@ -143,7 +150,6 @@ class GalleryFile extends HTMLElement {
    `;

    const cache = (this.hash && opts.browser_cache) ? await idbGet(this.hash) : undefined;
-    this.shadow.appendChild(style);
    const img = document.createElement('img');
    img.className = 'gallery-file';
    img.loading = 'lazy';
@@ -196,14 +202,20 @@ class GalleryFile extends HTMLElement {
        img.src = `file=${this.src}`;
      }
    }
-    if (!ok) return;
+    if (!ok) {
+      return;
+    }
    img.onclick = () => {
      currentImage = this.src;
      el.btnSend.click();
    };
    img.title = `Folder: ${this.folder}\nFile: ${this.name}\nSize: ${this.size.toLocaleString()} bytes\nModified: ${this.mtime.toLocaleString()}`;
+    if (this.shadow.children.length > 0) {
+      return; // avoid double-adding
+    }
    this.title = img.title;
    this.style.display = this.title.toLowerCase().includes(el.search.value.toLowerCase()) ? 'unset' : 'none';
+    this.shadow.appendChild(style);
    this.shadow.appendChild(img);
  }
 }
@@ -272,49 +284,68 @@ async function gallerySearch(evt) {
  }, 250);
 }

+const findDuplicates = (arr, key) => {
+  const map = new Map();
+  return arr.filter(item => {
+    const value = item[key];
+    if (map.has(value)) return true;
+    map.set(value, true);
+    return false;
+  });
+};
+
 async function gallerySort(btn) {
  const t0 = performance.now();
  const arr = Array.from(el.files.children).filter((node) => node.name); // filter out separators
+  if (arr.length === 0) return; // no files to sort
+  if (btn) lastSort = btn.charCodeAt(0);
+  lastSortName = 'none';
  const fragment = document.createDocumentFragment();
-  el.files.innerHTML = '';
-  log('gallerySort', btn.charCodeAt(0));
-  switch (btn.charCodeAt(0)) {
+  switch (lastSort) {
    case 61789: // name asc
+      lastSortName = 'name asc';
      arr
        .sort((a, b) => a.name.localeCompare(b.name))
        .forEach((node) => fragment.appendChild(node));
      break;
    case 61790: // name dsc
+      lastSortName = 'name dsc';
      arr
        .sort((b, a) => a.name.localeCompare(b.name))
        .forEach((node) => fragment.appendChild(node));
      break;
    case 61792: // size asc
+      lastSortName = 'size asc';
      arr
        .sort((a, b) => a.size - b.size)
        .forEach((node) => fragment.appendChild(node));
      break;
    case 61793: // size dsc
+      lastSortName = 'size dsc';
      arr
        .sort((b, a) => a.size - b.size)
        .forEach((node) => fragment.appendChild(node));
      break;
    case 61794: // resolution asc
+      lastSortName = 'resolution asc';
      arr
        .sort((a, b) => a.width * a.height - b.width * b.height)
        .forEach((node) => fragment.appendChild(node));
      break;
    case 61795: // resolution dsc
+      lastSortName = 'resolution dsc';
      arr
        .sort((b, a) => a.width * a.height - b.width * b.height)
        .forEach((node) => fragment.appendChild(node));
      break;
    case 61662:
+      lastSortName = 'modified asc';
      arr
        .sort((a, b) => a.mtime - b.mtime)
        .forEach((node) => fragment.appendChild(node));
      break;
    case 61661:
+      lastSortName = 'modified dsc';
      arr
        .sort((b, a) => a.mtime - b.mtime)
        .forEach((node) => fragment.appendChild(node));
@@ -322,14 +353,16 @@ async function gallerySort(btn) {
    default:
      break;
  }
+  if (fragment.children.length === 0) return;
+  el.files.innerHTML = '';
  el.files.appendChild(fragment);
  addSeparators();
  const t1 = performance.now();
-  el.status.innerText = `Sort | ${arr.length.toLocaleString()} images | ${Math.floor(t1 - t0).toLocaleString()}ms`;
+  log(`gallerySort: char=${lastSort} len=${arr.length} time=${Math.floor(t1 - t0)} sort=${lastSortName}`);
+  el.status.innerText = `Sort | ${lastSortName} | ${arr.length.toLocaleString()} images | ${Math.floor(t1 - t0).toLocaleString()}ms`;
 }

 async function fetchFilesHT(evt) {
-  el.status.innerText = `Folder | ${evt.target.name}`;
  const t0 = performance.now();
  const fragment = document.createDocumentFragment();
  el.status.innerText = `Folder | ${evt.target.name} | in-progress`;
@@ -389,7 +422,7 @@ async function fetchFilesWS(evt) { // fetch file-by-file list over websockets
      const file = new GalleryFile(data[0], data[1]);
      fragment.appendChild(file);
      if (numFiles % 100 === 0) {
-        el.status.innerText = `Folder | ${evt.target.name} | ${numFiles.toLocaleString()} images | ${Math.floor(t1 - t0).toLocaleString()}ms`;
+        el.status.innerText = `Folder | ${evt.target.name} | ${numFiles.toLocaleString()} images | in-progress | ${Math.floor(t1 - t0).toLocaleString()}ms`;
        el.files.appendChild(fragment);
        fragment = document.createDocumentFragment();
      }
@@ -397,6 +430,7 @@ async function fetchFilesWS(evt) { // fetch file-by-file list over websockets
  };
  ws.onclose = (event) => {
    el.files.appendChild(fragment);
+    // gallerySort();
    log(`gallery: folder=${evt.target.name} num=${numFiles} time=${Math.floor(t1 - t0)}ms`);
    el.status.innerText = `Folder | ${evt.target.name} | ${numFiles.toLocaleString()} images | ${Math.floor(t1 - t0).toLocaleString()}ms`;
    addSeparators();
@@ -264,8 +264,11 @@ def main():
    installer.check_onnx()
    installer.check_diffusers()
    installer.check_modified_files()
+    if args.test:
+        installer.log.info('Startup: test mode')
+        installer.quick_allowed = False
    if args.reinstall:
-        installer.log.info('Forcing reinstall of all packages')
+        installer.log.info('Startup: force reinstall of all packages')
        installer.quick_allowed = False
    if args.skip_all:
        installer.log.info('Startup: skip all')
@@ -6,15 +6,16 @@
 import torch
 import torch.nn.functional as F
 import numpy as np
-from skimage import filters


 ## Attention Utils
 def get_dynamic_threshold(tensor):
+    from skimage import filters
    return filters.threshold_otsu(tensor.float().cpu().numpy())


 def attn_map_to_binary(attention_map, scaler=1.):
+    from skimage import filters
    attention_map_np = attention_map.float().cpu().numpy()
    threshold_value = filters.threshold_otsu(attention_map_np) * scaler
    binary_mask = (attention_map_np > threshold_value).astype(np.uint8)
@@ -286,13 +286,13 @@ class ControlNet():
                    return
                if self.dtype is not None:
                    self.model.to(self.dtype)
-                if "ControlNet" in opts.nncf_compress_weights:
+                if "ControlNet" in opts.sdnq_quantize_weights:
                    try:
-                        log.debug(f'Control {what} model NNCF Compress: id="{model_id}"')
-                        from modules.model_quant import nncf_compress_model
-                        self.model = nncf_compress_model(self.model)
+                        log.debug(f'Control {what} model SDNQ Compress: id="{model_id}"')
+                        from modules.model_quant import sdnq_quantize_model
+                        self.model = sdnq_quantize_model(self.model)
                    except Exception as e:
-                        log.error(f'Control {what} model NNCF Compression failed: id="{model_id}" {e}')
+                        log.error(f'Control {what} model SDNQ Compression failed: id="{model_id}" {e}')
                elif "ControlNet" in opts.optimum_quanto_weights:
                    try:
                        log.debug(f'Control {what} model Optimum Quanto: id="{model_id}"')
@@ -1,6 +1,5 @@
 import cv2
 import numpy as np
-from skimage import transform as trans


 ### https://github.com/somanchiu/ReSwapper/blob/GAN/Image.py
@@ -74,6 +73,7 @@ arcface_dst = np.array(


 def estimate_norm(lmk, image_size=112,mode='arcface'): # pylint: disable=unused-argument
+    from skimage import transform as trans
    if image_size%112==0:
        ratio = float(image_size)/112.0
        diff_x = 0
@@ -126,6 +126,7 @@ def square_crop(im, S):


 def transform(data, center, output_size, scale, rotation):
+    from skimage import transform as trans
    scale_ratio = scale
    rot = float(rotation) * np.pi / 180.0
    t1 = trans.SimilarityTransform(scale=scale_ratio)
@@ -1,6 +1,5 @@
 import os
 import sys
-import contextlib
 import torch
 try:
    import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
@@ -21,7 +20,7 @@ def ipex_init(): # pylint: disable=too-many-statements
            try:
                # force xpu device on torch compile and triton
                # import inductor utils to get around lazy import
-                from torch._inductor import utils as torch_inductor_utils # pylint: disable=import-error, unused-import
+                from torch._inductor import utils as torch_inductor_utils # pylint: disable=import-error, unused-import # noqa: F401
                torch._inductor.utils.GPU_TYPES = ["xpu"]
                torch._inductor.utils.get_gpu_type = lambda *args, **kwargs: "xpu"
                from triton import backends as triton_backends # pylint: disable=import-error
@@ -161,29 +160,6 @@ def ipex_init(): # pylint: disable=too-many-statements
            torch.cuda.seed_all = torch.xpu.seed_all
            torch.cuda.initial_seed = torch.xpu.initial_seed

-            # AMP:
-            if has_ipex:
-                torch.xpu.amp.custom_fwd = torch.cuda.amp.custom_fwd
-                torch.xpu.amp.custom_bwd = torch.cuda.amp.custom_bwd
-                torch.cuda.amp = torch.xpu.amp
-                if torch_version < 2.3:
-                    torch.is_autocast_enabled = torch.xpu.is_autocast_xpu_enabled
-                    torch.get_autocast_gpu_dtype = torch.xpu.get_autocast_xpu_dtype
-
-                if not hasattr(torch.cuda.amp, "common"):
-                    torch.cuda.amp.common = contextlib.nullcontext()
-                torch.cuda.amp.common.amp_definitely_not_available = lambda: False
-
-                try:
-                    torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
-                except Exception: # pylint: disable=broad-exception-caught
-                    try:
-                        from .gradscaler import gradscaler_init # pylint: disable=import-outside-toplevel, import-error
-                        gradscaler_init()
-                        torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
-                    except Exception: # pylint: disable=broad-exception-caught
-                        torch.cuda.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler
-
            # C
            if torch_version < 2.3:
                torch._C._cuda_getCurrentRawStream = ipex._C._getCurrentRawStream
@@ -217,10 +193,10 @@ def ipex_init(): # pylint: disable=too-many-statements
            torch.cuda.ipc_collect = lambda *args, **kwargs: None
            torch.cuda.utilization = lambda *args, **kwargs: 0

-            device_supports_fp64, can_allocate_plus_4gb = ipex_hijacks()
+            device_supports_fp64 = ipex_hijacks()
            try:
                from .diffusers import ipex_diffusers
-                ipex_diffusers(device_supports_fp64=device_supports_fp64, can_allocate_plus_4gb=can_allocate_plus_4gb)
+                ipex_diffusers(device_supports_fp64=device_supports_fp64)
            except Exception: # pylint: disable=broad-exception-caught
                pass
            torch.cuda.is_xpu_hijacked = True
@@ -81,7 +81,7 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np"):
    return emb


-def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False):
+def ipex_diffusers(device_supports_fp64=False):
    diffusers.utils.torch_utils.fourier_filter = fourier_filter
    if not device_supports_fp64:
        # get around lazy imports
@@ -1,183 +0,0 @@
-from collections import defaultdict
-import torch
-import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
-import intel_extension_for_pytorch._C as core # pylint: disable=import-error, unused-import
-
-# 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
-OptState = ipex.cpu.autocast._grad_scaler.OptState
-_MultiDeviceReplicator = ipex.cpu.autocast._grad_scaler._MultiDeviceReplicator
-_refresh_per_optimizer_state = ipex.cpu.autocast._grad_scaler._refresh_per_optimizer_state
-
-def _unscale_grads_(self, optimizer, inv_scale, found_inf, allow_fp16): # pylint: disable=unused-argument
-    per_device_inv_scale = _MultiDeviceReplicator(inv_scale)
-    per_device_found_inf = _MultiDeviceReplicator(found_inf)
-
-    # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype.
-    # There could be hundreds of grads, so we'd like to iterate through them just once.
-    # However, we don't know their devices or dtypes in advance.
-
-    # https://stackoverflow.com/questions/5029934/defaultdict-of-defaultdict
-    # Google says mypy struggles with defaultdicts type annotations.
-    per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))  # type: ignore[var-annotated]
-    # sync grad to master weight
-    if hasattr(optimizer, "sync_grad"):
-        optimizer.sync_grad()
-    with torch.no_grad():
-        for group in optimizer.param_groups:
-            for param in group["params"]:
-                if param.grad is None:
-                    continue
-                if (not allow_fp16) and param.grad.dtype == torch.float16:
-                    raise ValueError("Attempting to unscale FP16 gradients.")
-                if param.grad.is_sparse:
-                    # is_coalesced() == False means the sparse grad has values with duplicate indices.
-                    # coalesce() deduplicates indices and adds all values that have the same index.
-                    # For scaled fp16 values, there's a good chance coalescing will cause overflow,
-                    # so we should check the coalesced _values().
-                    if param.grad.dtype is torch.float16:
-                        param.grad = param.grad.coalesce()
-                    to_unscale = param.grad._values()
-                else:
-                    to_unscale = param.grad
-
-                # -: is there a way to split by device and dtype without appending in the inner loop?
-                to_unscale = to_unscale.to("cpu")
-                per_device_and_dtype_grads[to_unscale.device][
-                    to_unscale.dtype
-                ].append(to_unscale)
-
-        for _, per_dtype_grads in per_device_and_dtype_grads.items():
-            for grads in per_dtype_grads.values():
-                core._amp_foreach_non_finite_check_and_unscale_(
-                    grads,
-                    per_device_found_inf.get("cpu"),
-                    per_device_inv_scale.get("cpu"),
-                )
-
-    return per_device_found_inf._per_device_tensors
-
-def unscale_(self, optimizer):
-    """
-    Divides ("unscales") the optimizer's gradient tensors by the scale factor.
-    :meth:`unscale_` is optional, serving cases where you need to
-    :ref:`modify or inspect gradients<working-with-unscaled-gradients>`
-    between the backward pass(es) and :meth:`step`.
-    If :meth:`unscale_` is not called explicitly,  gradients will be unscaled  automatically during :meth:`step`.
-    Simple example, using :meth:`unscale_` to enable clipping of unscaled gradients::
-        ...
-        scaler.scale(loss).backward()
-        scaler.unscale_(optimizer)
-        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
-        scaler.step(optimizer)
-        scaler.update()
-    Args:
-        optimizer (torch.optim.Optimizer):  Optimizer that owns the gradients to be unscaled.
-    .. warning::
-        :meth:`unscale_` should only be called once per optimizer per :meth:`step` call,
-        and only after all gradients for that optimizer's assigned parameters have been accumulated.
-        Calling :meth:`unscale_` twice for a given optimizer between each :meth:`step` triggers a RuntimeError.
-    .. warning::
-        :meth:`unscale_` may unscale sparse gradients out of place, replacing the ``.grad`` attribute.
-    """
-    if not self._enabled:
-        return
-
-    self._check_scale_growth_tracker("unscale_")
-
-    optimizer_state = self._per_optimizer_states[id(optimizer)]
-
-    if optimizer_state["stage"] is OptState.UNSCALED: # pylint: disable=no-else-raise
-        raise RuntimeError(
-            "unscale_() has already been called on this optimizer since the last update()."
-        )
-    elif optimizer_state["stage"] is OptState.STEPPED:
-        raise RuntimeError("unscale_() is being called after step().")
-
-    # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64.
-    assert self._scale is not None
-    if device_supports_fp64:
-        inv_scale = self._scale.double().reciprocal().float()
-    else:
-        inv_scale = self._scale.to("cpu").double().reciprocal().float().to(self._scale.device)
-    found_inf = torch.full(
-        (1,), 0.0, dtype=torch.float32, device=self._scale.device
-    )
-
-    optimizer_state["found_inf_per_device"] = self._unscale_grads_(
-        optimizer, inv_scale, found_inf, False
-    )
-    optimizer_state["stage"] = OptState.UNSCALED
-
-def update(self, new_scale=None):
-    """
-    Updates the scale factor.
-    If any optimizer steps were skipped the scale is multiplied by ``backoff_factor``
-    to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively,
-    the scale is multiplied by ``growth_factor`` to increase it.
-    Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not
-    used directly, it's used to fill GradScaler's internal scale tensor. So if
-    ``new_scale`` was a tensor, later in-place changes to that tensor will not further
-    affect the scale GradScaler uses internally.)
-    Args:
-        new_scale (float or :class:`torch.FloatTensor`, optional, default=None):  New scale factor.
-    .. warning::
-        :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has
-        been invoked for all optimizers used this iteration.
-    """
-    if not self._enabled:
-        return
-
-    _scale, _growth_tracker = self._check_scale_growth_tracker("update")
-
-    if new_scale is not None:
-        # Accept a new user-defined scale.
-        if isinstance(new_scale, float):
-            self._scale.fill_(new_scale)  # type: ignore[union-attr]
-        else:
-            reason = "new_scale should be a float or a 1-element torch.FloatTensor with requires_grad=False."
-            assert isinstance(new_scale, torch.FloatTensor), reason  # type: ignore[attr-defined]
-            assert new_scale.numel() == 1, reason
-            assert new_scale.requires_grad is False, reason
-            self._scale.copy_(new_scale)  # type: ignore[union-attr]
-    else:
-        # Consume shared inf/nan data collected from optimizers to update the scale.
-        # If all found_inf tensors are on the same device as self._scale, this operation is asynchronous.
-        found_infs = [
-            found_inf.to(device="cpu", non_blocking=True)
-            for state in self._per_optimizer_states.values()
-            for found_inf in state["found_inf_per_device"].values()
-        ]
-
-        assert len(found_infs) > 0, "No inf checks were recorded prior to update."
-
-        found_inf_combined = found_infs[0]
-        if len(found_infs) > 1:
-            for i in range(1, len(found_infs)):
-                found_inf_combined += found_infs[i]
-
-        to_device = _scale.device
-        _scale = _scale.to("cpu")
-        _growth_tracker = _growth_tracker.to("cpu")
-
-        core._amp_update_scale_(
-            _scale,
-            _growth_tracker,
-            found_inf_combined,
-            self._growth_factor,
-            self._backoff_factor,
-            self._growth_interval,
-        )
-
-        _scale = _scale.to(to_device)
-        _growth_tracker = _growth_tracker.to(to_device)
-    # To prepare for next iteration, clear the data collected from optimizers this iteration.
-    self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
-
-def gradscaler_init():
-    torch.xpu.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler
-    torch.xpu.amp.GradScaler._unscale_grads_ = _unscale_grads_
-    torch.xpu.amp.GradScaler.unscale_ = unscale_
-    torch.xpu.amp.GradScaler.update = update
-    return torch.xpu.amp.GradScaler
@@ -9,16 +9,21 @@ from modules import devices, errors
 torch_version = float(torch.__version__[:3])
 device_supports_fp64 = torch.xpu.has_fp64_dtype() if hasattr(torch.xpu, "has_fp64_dtype") else torch.xpu.get_device_properties(devices.device).has_fp64

-if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0' and (torch.xpu.get_device_properties(devices.device).total_memory / 1024 / 1024 / 1024) > 4.1:
-    try:
-        x = torch.ones((33000,33000), dtype=torch.float32, device=devices.device)
-        del x
-        torch.xpu.empty_cache()
-        can_allocate_plus_4gb = True
-    except Exception:
-        can_allocate_plus_4gb = False
+if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0':
+    if torch_version >= 2.7:
+        use_dynamic_attention = False # torch 2.7 has flash atten support
+    elif (torch.xpu.get_device_properties(devices.device).total_memory / 1024 / 1024 / 1024) > 4.1:
+        try:
+            x = torch.ones((33000,33000), dtype=torch.float32, device=devices.device)
+            del x
+            torch.xpu.empty_cache()
+            use_dynamic_attention = False
+        except Exception:
+            use_dynamic_attention = True
+    else:
+        use_dynamic_attention = True
 else:
-    can_allocate_plus_4gb = bool(os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '-1')
+    use_dynamic_attention = 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

@@ -51,14 +56,42 @@ def return_xpu(device): # keep the device instance type, aka return string if th
 # Autocast
 original_autocast_init = torch.amp.autocast_mode.autocast.__init__
@wraps(torch.amp.autocast_mode.autocast.__init__)
-def autocast_init(self, device_type, dtype=None, enabled=True, cache_enabled=None):
-    if device_type == "cuda" or device_type == "xpu":
+def autocast_init(self, device_type=None, dtype=None, enabled=True, cache_enabled=None):
+    if device_type is None or check_cuda(device_type) or check_device_type(device_type, "xpu"):
        if dtype is None:
            dtype = devices.dtype
        return original_autocast_init(self, device_type="xpu", dtype=dtype, enabled=enabled, cache_enabled=cache_enabled)
    else:
        return original_autocast_init(self, device_type=device_type, dtype=dtype, enabled=enabled, cache_enabled=cache_enabled)

+
+original_grad_scaler_init = torch.amp.grad_scaler.GradScaler.__init__
+@wraps(torch.amp.grad_scaler.GradScaler.__init__)
+def GradScaler_init(self, device: str = None, init_scale: float = 2.0**16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True):
+    if device is None or check_cuda(device):
+        return original_grad_scaler_init(self, device=return_xpu(device), init_scale=init_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, enabled=enabled)
+    else:
+        return original_grad_scaler_init(self, device=device, init_scale=init_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, enabled=enabled)
+
+
+original_is_autocast_enabled = torch.is_autocast_enabled
+@wraps(torch.is_autocast_enabled)
+def torch_is_autocast_enabled(device_type=None):
+    if device_type is None or check_cuda(device_type):
+        return original_is_autocast_enabled(return_xpu(device_type))
+    else:
+        return original_is_autocast_enabled(device_type)
+
+
+original_get_autocast_dtype = torch.get_autocast_dtype
+@wraps(torch.get_autocast_dtype)
+def torch_get_autocast_dtype(device_type=None):
+    if device_type is None or check_cuda(device_type) or check_device_type(device_type, "xpu"):
+        return devices.dtype
+    else:
+        return original_get_autocast_dtype(device_type)
+
+
 # Latent Antialias CPU Offload:
 # IPEX 2.5 and above has partial support but doesn't really work most of the time.
 original_interpolate = torch.nn.functional.interpolate
@@ -94,7 +127,7 @@ def as_tensor(data, dtype=None, device=None):
        return original_as_tensor(data, dtype=dtype, device=device)


-if can_allocate_plus_4gb:
+if not use_dynamic_attention:
    original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
 else:
    # 32 bit attention workarounds for Alchemist:
@@ -360,7 +393,7 @@ class torch_Generator(original_torch_Generator):

 # Hijack Functions:
 def ipex_hijacks():
-    global device_supports_fp64, can_allocate_plus_4gb
+    global device_supports_fp64
    if torch_version >= 2.4:
        torch.UntypedStorage.cuda = UntypedStorage_cuda
        torch.UntypedStorage.to = UntypedStorage_to
@@ -404,4 +437,28 @@ def ipex_hijacks():
    if not device_supports_fp64:
        torch.from_numpy = from_numpy
        torch.as_tensor = as_tensor
-    return device_supports_fp64, can_allocate_plus_4gb
+
+    # AMP:
+    torch.amp.grad_scaler.GradScaler.__init__ = GradScaler_init
+    torch.is_autocast_enabled = torch_is_autocast_enabled
+    torch.get_autocast_gpu_dtype = torch_get_autocast_dtype
+    torch.get_autocast_dtype = torch_get_autocast_dtype
+
+    if hasattr(torch.xpu, "amp"):
+        if not hasattr(torch.xpu.amp, "custom_fwd"):
+            torch.xpu.amp.custom_fwd = torch.cuda.amp.custom_fwd
+            torch.xpu.amp.custom_bwd = torch.cuda.amp.custom_bwd
+        if not hasattr(torch.xpu.amp, "GradScaler"):
+            torch.xpu.amp.GradScaler = torch.amp.grad_scaler.GradScaler
+        torch.cuda.amp = torch.xpu.amp
+    else:
+        if not hasattr(torch.amp, "custom_fwd"):
+            torch.amp.custom_fwd = torch.cuda.amp.custom_fwd
+            torch.amp.custom_bwd = torch.cuda.amp.custom_bwd
+        torch.cuda.amp = torch.amp
+
+    if not hasattr(torch.cuda.amp, "common"):
+        torch.cuda.amp.common = nullcontext()
+    torch.cuda.amp.common.amp_definitely_not_available = lambda: False
+
+    return device_supports_fp64
@@ -52,6 +52,24 @@ DEFAULT_OPENVINO_PYTHON_CONFIG = MappingProxyType(
    },
 )

+dtype_mapping = {
+        torch.float32: Type.f32,
+        torch.float64: Type.f64,
+        torch.float16: Type.f16,
+        torch.bfloat16: Type.bf16,
+        torch.float8_e4m3fn: Type.f8e4m3,
+        torch.float8_e5m2: Type.f8e5m2,
+        torch.int64: Type.i64,
+        torch.uint64: Type.u64,
+        torch.int32: Type.i32,
+        torch.uint32: Type.u32,
+        torch.int8: Type.i8,
+        torch.uint8: Type.u8,
+        torch.bool: Type.boolean
+    }
+if hasattr(torch, "float8_e8m0fnu"):
+    dtype_mapping[torch.float8_e8m0fnu] = Type.f8e8m0
+

 class OpenVINOGraphModule(torch.nn.Module):
    def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None, file_name="", int_inputs=[]):
@@ -219,17 +237,6 @@ def openvino_compile(gm: GraphModule, *example_inputs, model_hash_str: str = Non
                    f.write("\n")
                f.close()

-    dtype_mapping = {
-        torch.float32: Type.f32,
-        torch.float64: Type.f64,
-        torch.float16: Type.f16,
-        torch.int64: Type.i64,
-        torch.int32: Type.i32,
-        torch.uint8: Type.u8,
-        torch.int8: Type.i8,
-        torch.bool: Type.boolean
-    }
-
    idx_minus = 0
    for idx, input_data in enumerate(example_inputs):
        if isinstance(input_data, int):
@@ -283,17 +290,6 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs):
    global dont_use_nncf
    global dont_use_quant

-    dtype_mapping = {
-        torch.float32: Type.f32,
-        torch.float64: Type.f64,
-        torch.float16: Type.f16,
-        torch.int64: Type.i64,
-        torch.int32: Type.i32,
-        torch.uint8: Type.u8,
-        torch.int8: Type.i8,
-        torch.bool: Type.boolean
-    }
-
    for idx, input_data in enumerate(example_inputs):
        om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
        om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
@@ -45,8 +45,8 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n
                    self.network_weights_backup = True
                else:
                    self.network_weights_backup = weight.clone().to(devices.cpu)
-                    if self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1:
-                        self.nncf_decompressor_backup = self.pre_ops["0"].to(devices.cpu)
+                    if hasattr(self, "sdnq_decompressor"):
+                        self.sdnq_decompressor_backup = self.sdnq_decompressor.to(devices.cpu)

        if bias_backup is None:
            if getattr(self, 'bias', None) is not None:
@@ -79,11 +79,8 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
            continue
        try:
            t0 = time.time()
-            if self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1:
-                return_device = self.weight.data.device
-                self.weight.data = self.weight.data.to(devices.device)
-                weight = self.pre_ops["0"].to(devices.device)(self, return_decompressed_only=True)
-                self.weight.data = self.weight.data.to(return_device)
+            if hasattr(self, "sdnq_decompressor"):
+                weight = self.sdnq_decompressor.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_decompressor.use_quantized_matmul)
            else:
                weight = self.weight.to(devices.device) # must perform calc on gpu due to performance
            updown, ex_bias = module.calc_updown(weight)
@@ -139,24 +136,23 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
            # weight._quantize(devices.device) / weight.to(device=device)
        except Exception as e:
            shared.log.error(f'Network load: type=LoRA quant=bnb cls={self.__class__.__name__} type={self.quant_type} blocksize={self.blocksize} state={vars(self.quant_state)} weight={self.weight} bias={lora_weights} {e}')
-    elif not bias and self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1:
-        num_bits = None
-        is_asym_mode = None
+    elif not bias and hasattr(self, "sdnq_decompressor"):
        try:
-            from modules.model_quant_nncf import nncf_compress_layer
-            num_bits = self.pre_ops["0"].num_bits
-            is_asym_mode = self.pre_ops["0"].quantization_mode == "asymmetric"
-            self.weight = torch.nn.Parameter(model_weights.to(devices.device), requires_grad=False)
-            dequant_weight = self.pre_ops["0"](self, return_decompressed_only=True)
+            from modules.model_quant_sdnq import sdnq_quantize_layer
+            if hasattr(self, "sdnq_decompressor_backup"):
+                sdnq_decompressor = self.sdnq_decompressor_backup.to(devices.device)
+            else:
+                sdnq_decompressor = self.sdnq_decompressor.to(devices.device)
+            dequant_weight = sdnq_decompressor(model_weights.to(devices.device), skip_quantized_matmul=sdnq_decompressor.use_quantized_matmul)
            new_weight = dequant_weight.to(devices.device, dtype=torch.float32) + lora_weights.to(devices.device, dtype=torch.float32)
            self.weight = torch.nn.Parameter(new_weight, requires_grad=False)
-            self.pre_ops.pop("0")
-            self._custom_forward_fn = None # pylint: disable=protected-access
-            self = nncf_compress_layer(self, num_bits, is_asym_mode, torch_dtype=devices.dtype, quant_conv=shared.opts.nncf_quantize_conv_layers, group_size=shared.opts.nncf_compress_weights_group_size, use_int8_matmul=shared.opts.nncf_decompress_int8_matmul)
+            self.sdnq_decompressor = None
+            self = sdnq_quantize_layer(self, sdnq_decompressor.weights_dtype, torch_dtype=devices.dtype, group_size=shared.opts.sdnq_quantize_weights_group_size, quant_conv=shared.opts.sdnq_quantize_conv_layers, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, param_name=getattr(self, 'network_layer_name', None))
            self = self.to(device)
+            weight = None
            del dequant_weight
        except Exception as e:
-            shared.log.error(f'Network load: type=LoRA quant=nncf cls={self.__class__.__name__} bits={num_bits} is_asym_mode={is_asym_mode} weight={self.weight} lora_weights={lora_weights} {e}')
+            shared.log.error(f'Network load: type=LoRA quant=sdnq cls={self.__class__.__name__} weight={self.weight} lora_weights={lora_weights} {e}')
    else:
        try:
            new_weight = model_weights.to(devices.device) + lora_weights.to(devices.device)
@@ -218,8 +214,8 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
            network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, device=device, bias=False)
        else:
            self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False)
-            if hasattr(self, "nncf_decompressor_backup"):
-                self.pre_ops["0"] = self.nncf_decompressor_backup.to(device)
+            if hasattr(self, "sdnq_decompressor_backup"):
+                self.sdnq_decompressor = self.sdnq_decompressor_backup.to(device)

    if bias_backup is not None:
        self.bias = None
@@ -94,6 +94,11 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
    load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
    shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')

+    if shared.opts.teacache_enabled:
+        from modules import teacache
+        shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.HiDreamImageTransformer2DModel.__name__}')
+        diffusers.HiDreamImageTransformer2DModel.forward = teacache.teacache_hidream_forward # patch must be done before transformer is loaded
+
    if 'I1' in repo_id:
        cls = diffusers.HiDreamImagePipeline
    elif 'E1' in repo_id:
@@ -9,7 +9,6 @@ from installer import installed, install, log, setup_logging

 ao = None
 bnb = None
-intel_nncf = None
 optimum_quanto = None
 quant_last_model_name = None
 quant_last_model_device = None
@@ -104,38 +103,36 @@ def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str =
    return kwargs


-def create_nncf_config(kwargs = None, allow_nncf: bool = True, module: str = 'Model'):
+def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Model', weights_dtype: str = None):
    from modules import shared
-    if len(shared.opts.nncf_compress_weights) > 0 and (shared.opts.nncf_compress_mode == 'pre') and allow_nncf:
-        if 'Model' in shared.opts.nncf_compress_weights or (module is not None and module in shared.opts.nncf_compress_weights) or module == 'any':
-            load_nncf(silent=True)
-            if intel_nncf is None:
-                return kwargs
-            from modules.model_quant_nncf import NNCFQuantizer, NNCFConfig
-            diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["nncf"] = NNCFQuantizer
-            transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["nncf"] = NNCFQuantizer
-            diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["nncf"] = NNCFConfig
-            transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["nncf"] = NNCFConfig
+    if len(shared.opts.sdnq_quantize_weights) > 0 and (shared.opts.sdnq_quantize_mode == 'pre') and allow_sdnq:
+        if 'Model' in shared.opts.sdnq_quantize_weights or (module is not None and module in shared.opts.sdnq_quantize_weights) or module == 'any':
+            from modules.model_quant_sdnq import SDNQQuantizer, SDNQConfig
+            diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
+            transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
+            diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
+            transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig

-            nncf_config = NNCFConfig(
-                weights_dtype=shared.opts.nncf_compress_weights_mode.lower(),
-                group_size=shared.opts.nncf_compress_weights_group_size,
-                use_int8_matmul=shared.opts.nncf_decompress_int8_matmul,
+            sdnq_config = SDNQConfig(
+                weights_dtype=weights_dtype if weights_dtype is not None else shared.opts.sdnq_quantize_weights_mode,
+                group_size=shared.opts.sdnq_quantize_weights_group_size,
+                quant_conv=shared.opts.sdnq_quantize_conv_layers,
+                use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
            )
-            log.debug(f'Quantization: module="{module}" type=nncf dtype={shared.opts.nncf_compress_weights_mode}')
+            log.debug(f'Quantization: module="{module}" type=sdnq dtype={shared.opts.sdnq_quantize_weights_mode}')
            if kwargs is None:
-                return nncf_config
+                return sdnq_config
            else:
-                kwargs['quantization_config'] = nncf_config
+                kwargs['quantization_config'] = sdnq_config
                return kwargs
    return kwargs


 def check_quant(module: str = ''):
    from modules import shared
-    if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization or 'Model' in shared.opts.nncf_compress_weights:
+    if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization or 'Model' in shared.opts.sdnq_quantize_weights:
        return True
-    if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization or module in shared.opts.nncf_compress_weights:
+    if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization or module in shared.opts.sdnq_quantize_weights:
        return True
    return False

@@ -154,25 +151,25 @@ def check_nunchaku(module: str = ''):
 def create_config(kwargs = None, allow: bool = True, module: str = 'Model'):
    if kwargs is None:
        kwargs = {}
+    kwargs = create_sdnq_config(kwargs, allow_sdnq=allow, module=module)
+    if kwargs is not None and 'quantization_config' in kwargs:
+        if debug:
+            log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}')
+        return kwargs
    kwargs = create_bnb_config(kwargs, allow_bnb=allow, module=module)
    if kwargs is not None and 'quantization_config' in kwargs:
        if debug:
            log.trace(f'Quantization: type=bnb config={kwargs.get("quantization_config", None)}')
        return kwargs
-    kwargs = create_ao_config(kwargs, allow_ao=allow, module=module)
-    if kwargs is not None and 'quantization_config' in kwargs:
-        if debug:
-            log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}')
-        return kwargs
    kwargs = create_quanto_config(kwargs, allow_quanto=allow, module=module)
    if kwargs is not None and 'quantization_config' in kwargs:
        if debug:
            log.trace(f'Quantization: type=quanto config={kwargs.get("quantization_config", None)}')
        return kwargs
-    kwargs = create_nncf_config(kwargs, allow_nncf=allow, module=module)
+    kwargs = create_ao_config(kwargs, allow_ao=allow, module=module)
    if kwargs is not None and 'quantization_config' in kwargs:
        if debug:
-            log.trace(f'Quantization: type=nncf config={kwargs.get("quantization_config", None)}')
+            log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}')
        return kwargs
    return kwargs

@@ -259,35 +256,6 @@ def load_quanto(msg='', silent=False):
    return None


-def load_nncf(msg='', silent=False):
-    global intel_nncf # pylint: disable=global-statement
-    if intel_nncf is not None:
-        return intel_nncf
-    if not installed('nncf'):
-        install('nncf==2.16.0', quiet=True)
-        log.warning('Quantization: nncf installed please restart')
-    install('jstyleson', quiet=True)
-    install('texttable', quiet=True)
-    install('tabulate', quiet=True)
-    try:
-        import nncf
-        intel_nncf = nncf
-        try:
-            nncf.common.logging.logger.warn_bkc_version_mismatch = lambda *args, **kwargs: None # silence the pytorch version warning
-        except Exception:
-            pass
-        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
@@ -331,22 +299,18 @@ def apply_layerwise(sd_model, quiet:bool=False):
                log.error(f'Quantization: type=layerwise {e}')


-def nncf_compress_model(model, op=None, sd_model=None, do_gc=True):
+def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
    global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
    from modules import devices, shared
-    from accelerate import init_empty_weights
-
-    load_nncf('Quantize model: type=NNCF')
-    from modules.model_quant_nncf import apply_nncf_to_module
-    from nncf.torch.nncf_module_replacement import replace_modules_by_nncf_modules # get around lazy import
+    from modules.model_quant_sdnq import apply_sdnq_to_module

    model.eval()

    if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
        import torch
-        from modules.model_quant_nncf import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
+        from modules.model_quant_sdnq import SDNQ_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
        for i in range(len(model.encoder.block)):
-            model.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
+            model.encoder.block[i].layer[1].DenseReluDense = SDNQ_T5DenseGatedActDense(
                model.encoder.block[i].layer[1].DenseReluDense,
                dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
            )
@@ -355,18 +319,21 @@ def nncf_compress_model(model, op=None, sd_model=None, do_gc=True):
    if hasattr(model, "get_input_embeddings"):
        backup_embeddings = copy.deepcopy(model.get_input_embeddings())

-    with init_empty_weights():
-        model, _ = replace_modules_by_nncf_modules(model)
-
-    num_bits = 8 if shared.opts.nncf_compress_weights_mode in {"INT8", "INT8_SYM", "INT8_ASYM"} else 4
-    is_asym_mode = shared.opts.nncf_compress_weights_mode in {"INT8", "INT4", "INT8_ASYM", "INT4_ASYM"}
-    model = apply_nncf_to_module(model, num_bits, is_asym_mode, quant_conv=shared.opts.nncf_quantize_conv_layers)
-    model.quantization_method = 'NNCF'
+    model = apply_sdnq_to_module(
+        model,
+        weights_dtype=shared.opts.sdnq_quantize_weights_mode,
+        torch_dtype=devices.dtype,
+        group_size=shared.opts.sdnq_quantize_weights_group_size,
+        quant_conv=shared.opts.sdnq_quantize_conv_layers,
+        use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+        param_name=op,
+    )
+    model.quantization_method = 'SDNQ'

    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 op is not None and shared.opts.sdnq_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(".")
@@ -382,19 +349,21 @@ def nncf_compress_model(model, op=None, sd_model=None, do_gc=True):
            quant_last_model_name = None
            quant_last_model_device = None
        model.to(devices.device)
+    elif shared.opts.diffusers_offload_mode != "none":
+        model = model.to(devices.cpu)
    if do_gc:
        devices.torch_gc(force=True)
    return model


-def nncf_compress_weights(sd_model):
+def sdnq_quantize_weights(sd_model):
    try:
        t0 = time.time()
        from modules import shared, devices, sd_models
-        log.info(f"Quantization: type=NNCF modules={shared.opts.nncf_compress_weights}")
+        log.info(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_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")
+        sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq")
        if quant_last_model_name is not None:
            if "." in quant_last_model_name:
                last_model_names = quant_last_model_name.split(".")
@@ -406,9 +375,9 @@ def nncf_compress_weights(sd_model):
        quant_last_model_device = None

        t1 = time.time()
-        log.info(f"Quantization: type=NNCF time={t1-t0:.2f}")
+        log.info(f"Quantization: type=SDNQ time={t1-t0:.2f}")
    except Exception as e:
-        log.warning(f"Quantization: type=NNCF {e}")
+        log.warning(f"Quantization: type=SDNQ {e}")
    return sd_model


@@ -554,12 +523,12 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al
    # if 'variant' in config:
    #     del config['variant']
    if device_map:
-        if shared.opts.device_map == 'cpu':
-            config['device_map'] = 'cpu'
-        if shared.opts.device_map == 'gpu':
-            config['device_map'] = devices.device
        if devices.backend == "ipex" and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}:
            config['device_map'] = 'cpu' # alchemist gpus hits the 4GB allocation limit with transformers, UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS emulates above 4GB allocations
+        elif shared.opts.device_map == 'cpu':
+            config['device_map'] = 'cpu'
+        elif shared.opts.device_map == 'gpu':
+            config['device_map'] = devices.device
    if allow_quant:
        quant_args = create_config(module=module)
    else:
@@ -569,8 +538,8 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al

 def do_post_load_quant(sd_model):
    from modules import shared
-    if shared.opts.nncf_compress_weights and shared.opts.nncf_compress_mode == 'post' and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
-        sd_model = nncf_compress_weights(sd_model)
+    if shared.opts.sdnq_quantize_weights and shared.opts.sdnq_quantize_mode == 'post':
+        sd_model = sdnq_quantize_weights(sd_model)
    if shared.opts.optimum_quanto_weights:
        sd_model = optimum_quanto_weights(sd_model)
    if shared.opts.torchao_quantization and shared.opts.torchao_quantization_mode == 'post':
@@ -1,669 +0,0 @@
-# pylint: disable=redefined-builtin,no-member
-
-from typing import Any, Dict, List, Tuple, Optional, Union
-from dataclasses import dataclass
-from enum import Enum
-import os
-import torch
-from diffusers.quantizers.base import DiffusersQuantizer
-from diffusers.quantizers.quantization_config import QuantizationConfigMixin
-from diffusers.utils import get_module_from_name
-from accelerate import init_empty_weights
-from accelerate.utils import CustomDtype
-from modules import devices, shared
-
-
-debug = os.environ.get('SD_QUANT_DEBUG', None) is not None
-torch_dtype_dict = {
-    "int8": torch.int8,
-    "uint8": torch.uint8,
-    "int4": CustomDtype.INT4,
-    "uint4": CustomDtype.INT4,
-}
-weights_dtype_dict = {
-    "int8_asym": "uint8",
-    "int8_sym": "int8",
-    "int4_asym": "uint4",
-    "int4_sym": "int4",
-    "int8": "uint8",
-    "int4": "uint4",
-}
-linear_types = ["NNCFLinear", "Linear"]
-conv_types = ["NNCFConv1d", "NNCFConv2d", "NNCFConv3d", "Conv1d", "Conv2d", "Conv3d"]
-conv_transpose_types = ["NNCFConvTranspose1d", "NNCFConvTranspose2d", "NNCFConvTranspose3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d"]
-allowed_types = []
-allowed_types.extend(linear_types)
-allowed_types.extend(conv_types)
-allowed_types.extend(conv_transpose_types)
-
-
-class QuantizationMethod(str, Enum):
-    NNCF = "nncf"
-
-
-def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_conv=False, group_size=0, use_int8_matmul=False, param_name=None): # pylint: disable=unused-argument
-    if layer.__class__.__name__ in allowed_types:
-        if torch_dtype is None:
-            torch_dtype = devices.dtype
-        result_shape = None
-
-        if layer.__class__.__name__ in conv_types:
-            if is_asym_mode or not quant_conv: # don't quant convs with asym mode
-                return layer
-            reduction_axes = [i for i in range(layer.weight.ndim) if i != 0]
-            use_int8_matmul = False
-        if layer.__class__.__name__ in conv_transpose_types:
-            if is_asym_mode or not quant_conv: # don't quant convs with asym mode
-                return layer
-            reduction_axes = [i for i in range(layer.weight.ndim) if i != 1]
-            use_int8_matmul = False
-        else:
-            reduction_axes = -1
-            channel_size = layer.weight.shape[-1]
-            use_int8_matmul = use_int8_matmul and not is_asym_mode and channel_size >= 32 and layer.weight.shape[0] >= 32
-
-            if not use_int8_matmul and (group_size > 0 or (num_bits == 4 and group_size != -1)):
-                if group_size == 0:
-                    group_size = 64
-                    num_of_groups = channel_size // group_size
-
-                if group_size >= channel_size:
-                    group_size = channel_size
-                    num_of_groups = 1
-                else:
-                    num_of_groups = channel_size // group_size
-                    while channel_size % group_size != 0: # find something divisible
-                        num_of_groups -= 1
-                        if num_of_groups <= 1:
-                            group_size = channel_size
-                            num_of_groups = 1
-                            break
-                        group_size = channel_size / num_of_groups
-
-                if num_of_groups > 1:
-                    result_shape = layer.weight.shape
-                    new_shape = list(result_shape)
-                    last_dim_index = layer.weight.ndim
-                    new_shape[last_dim_index - 1 : last_dim_index] = (int(num_of_groups), int(group_size))
-                    layer.weight.data = layer.weight.reshape(new_shape)
-
-        if shared.opts.diffusers_offload_mode != "none":
-            return_device = layer.weight.data.device
-        else:
-            return_device = devices.device
-        layer.weight.data = layer.weight.data.to(devices.device, dtype=torch.float32)
-
-        if is_asym_mode:
-            scale, zero_point = get_int_scale_asymmetric(layer.weight, reduction_axes, num_bits)
-        else:
-            scale = get_int_scale_symmetric(layer.weight, reduction_axes, num_bits)
-            zero_point = None
-        compressed_weight = quantize_int(layer.weight, scale, zero_point, is_asym_mode, num_bits)
-
-        if not shared.opts.nncf_decompress_fp32:
-            scale = scale.to(torch_dtype)
-            if zero_point is not None:
-                zero_point = zero_point.to(torch_dtype)
-
-        if use_int8_matmul:
-            layer._custom_forward_fn = linear_forward_int8_matmul # pylint: disable=protected-access
-            scale = scale.squeeze(-1)
-            if num_bits == 8:
-                compressed_weight = compressed_weight.transpose(0,1)
-        else:
-            layer._custom_forward_fn = None # pylint: disable=protected-access
-
-        if num_bits == 4:
-            if is_asym_mode:
-                decompressor = INT4AsymmetricWeightsDecompressor(
-                    scale=scale.data,
-                    zero_point=zero_point.data,
-                    compressed_weight_shape=compressed_weight.shape,
-                    result_dtype=torch_dtype,
-                    result_shape=result_shape,
-                )
-            else:
-                decompressor = INT4SymmetricWeightsDecompressor(
-                    scale=scale.data,
-                    compressed_weight_shape=compressed_weight.shape,
-                    result_dtype=torch_dtype,
-                    result_shape=result_shape,
-                    use_int8_matmul=use_int8_matmul,
-                )
-        else:
-            if is_asym_mode:
-                decompressor = INT8AsymmetricWeightsDecompressor(
-                    scale=scale.data,
-                    zero_point=zero_point.data,
-                    result_dtype=torch_dtype,
-                    result_shape=result_shape,
-                )
-            else:
-                decompressor = INT8SymmetricWeightsDecompressor(
-                    scale=scale.data,
-                    result_dtype=torch_dtype,
-                    result_shape=result_shape,
-                    use_int8_matmul=use_int8_matmul,
-                )
-
-        compressed_weight = decompressor.pack_weight(compressed_weight).to(return_device)
-        decompressor = decompressor.to(return_device)
-        layer.register_pre_forward_operation(decompressor)
-        layer.weight.requires_grad = False
-        layer.weight.data = compressed_weight
-    return layer
-
-
-def apply_nncf_to_module(model, num_bits, is_asym_mode, quant_conv=False):
-    has_children = list(model.children())
-    if not has_children:
-        return model
-    for param_name, module in model.named_children():
-        if module.__class__.__name__.startswith("NNCF") and hasattr(module, "weight") and module.weight is not None:
-            module = nncf_compress_layer(
-                module,
-                num_bits,
-                is_asym_mode,
-                torch_dtype=devices.dtype,
-                quant_conv=quant_conv,
-                group_size=shared.opts.nncf_compress_weights_group_size,
-                use_int8_matmul=shared.opts.nncf_decompress_int8_matmul,
-                param_name=param_name,
-            )
-        module = apply_nncf_to_module(module, num_bits, is_asym_mode, quant_conv=quant_conv)
-    return model
-
-
-def nncf_send_to_device(model, device):
-    for child in model.children():
-        if "WeightsDecompressor" in child.__class__.__name__:
-            child.scale = child.scale.to(device)
-            if hasattr(child, "zero_point"):
-                child.zero_point = child.zero_point.to(device)
-        nncf_send_to_device(child, device)
-
-
-class NNCFQuantizer(DiffusersQuantizer):
-    r"""
-    Diffusers Quantizer for NNCF
-    """
-
-    requires_parameters_quantization = True
-    use_keep_in_fp32_modules = True
-    requires_calibration = False
-    required_packages = ["nncf"]
-    torch_dtype = None
-
-    def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation
-        super().__init__(quantization_config, **kwargs)
-
-    def check_if_quantized_param(
-        self,
-        model,
-        param_value: "torch.Tensor",
-        param_name: str,
-        state_dict: Dict[str, Any],
-        **kwargs,
-    ):
-        module, _ = get_module_from_name(model, param_name)
-        return module.__class__.__name__.startswith("NNCF") and param_name.endswith(".weight")
-
-    def check_quantized_param(self, *args, **kwargs) -> bool:
-        """
-        needed for transformers compatibilty, returns self.check_if_quantized_param
-        """
-        return self.check_if_quantized_param(*args, **kwargs)
-
-    def create_quantized_param( # pylint: disable=arguments-differ
-        self,
-        model,
-        param_value: "torch.Tensor",
-        param_name: str,
-        target_device: "torch.device",
-        state_dict: Dict[str, Any], # pylint: disable=unused-argument
-        unexpected_keys: List[str], # pylint: disable=unused-argument
-        **kwargs,
-    ):
-        # load the model params to target_device first
-        layer, tensor_name = get_module_from_name(model, param_name)
-        layer._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device) # pylint: disable=protected-access
-
-        # nncf_padding_value somehow ends up in the meta device with cogvideo even if we don't use init_empty_weights
-        # set it to the default value if it is in the meta device:
-        if layer.__class__.__name__ == "NNCFConv2d" and hasattr(layer, "get_padding_value_ref") and hasattr(layer, "_set_padding_value"):
-            if layer.get_padding_value_ref().device == torch.device("meta"):
-                layer._set_padding_value(torch.zeros([1]))
-
-        split_param_name = param_name.split(".")
-        if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert):
-            layer = nncf_compress_layer(
-                layer,
-                self.quantization_config.num_bits,
-                self.quantization_config.is_asym_mode,
-                torch_dtype=self.torch_dtype,
-                group_size=self.quantization_config.group_size,
-                use_int8_matmul=self.quantization_config.use_int8_matmul,
-                param_name=param_name,
-            )
-
-    def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
-        max_memory = {key: val * 0.70 for key, val in max_memory.items()}
-        return max_memory
-
-    def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": # pylint: disable=unused-argument,arguments-renamed
-        return torch_dtype_dict[self.quantization_config.weights_dtype]
-
-    def update_torch_dtype(self, torch_dtype: "torch.dtype" = None) -> "torch.dtype":
-        if torch_dtype is None:
-            torch_dtype = devices.dtype
-        self.torch_dtype = torch_dtype
-        return torch_dtype
-
-    def _process_model_before_weight_loading( # pylint: disable=arguments-differ
-        self,
-        model,
-        device_map, # pylint: disable=unused-argument
-        keep_in_fp32_modules: List[str] = [],
-        **kwargs,
-    ):
-        from nncf.torch.nncf_module_replacement import replace_modules_by_nncf_modules
-
-        self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
-        if not isinstance(self.modules_to_not_convert, list):
-            self.modules_to_not_convert = [self.modules_to_not_convert]
-        if keep_in_fp32_modules is not None:
-            self.modules_to_not_convert.extend(keep_in_fp32_modules)
-
-        model.config.quantization_config = self.quantization_config
-        with init_empty_weights():
-            model, _ = replace_modules_by_nncf_modules(model)
-
-    def _process_model_after_weight_loading(self, model, **kwargs):
-        return model
-
-    def update_tp_plan(self, config):
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return config
-
-    def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return unexpected_keys
-
-    def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return missing_keys
-
-    def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: # pylint: disable=unused-argument
-        """
-        needed for transformers compatibilty, no-op function
-        """
-        return expected_keys
-
-    @property
-    def is_trainable(self):
-        return False
-
-    @property
-    def is_serializable(self):
-        return False
-
-
-@dataclass
-class NNCFConfig(QuantizationConfigMixin):
-    """
-    This is a wrapper class about all possible attributes and features that you can play with a model that has been
-    loaded using `nncf`.
-
-    Args:
-        weights_dtype (`str`, *optional*, defaults to `"int8"`):
-            The target dtype for the weights after quantization. Supported values are ("int8", "int8_sym", "int4", "int4_sym")
-       modules_to_not_convert (`list`, *optional*, default to `None`):
-            The list of modules to not quantize, useful for quantizing models that explicitly require to have some
-            modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
-    """
-
-    def __init__( # pylint: disable=super-init-not-called
-        self,
-        weights_dtype: str = "int8_sym",
-        group_size: int = 0,
-        use_int8_matmul: bool = False,
-        modules_to_not_convert: Optional[List[str]] = None,
-        **kwargs, # pylint: disable=unused-argument
-    ):
-        self.quant_method = QuantizationMethod.NNCF
-        self.weights_dtype = weights_dtype_dict[weights_dtype.lower()]
-        self.group_size = group_size
-        self.use_int8_matmul = use_int8_matmul
-        self.modules_to_not_convert = modules_to_not_convert
-
-        self.post_init()
-
-        self.num_bits = 8 if self.weights_dtype in {"int8", "uint8"} else 4
-        self.is_asym_mode = self.weights_dtype in {"uint8", "uint4"}
-        self.is_integer = True
-
-    def post_init(self):
-        r"""
-        Safety checker that arguments are correct
-        """
-        accepted_weights = ["int8", "uint8", "int4", "uint4"]
-        if self.weights_dtype not in accepted_weights:
-            raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")
-
-
-class NNCF_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class
-    def __init__(self, T5DenseGatedActDense, dtype):
-        super().__init__()
-        self.wi_0 = T5DenseGatedActDense.wi_0
-        self.wi_1 = T5DenseGatedActDense.wi_1
-        self.wo = T5DenseGatedActDense.wo
-        self.dropout = T5DenseGatedActDense.dropout
-        self.act = T5DenseGatedActDense.act
-        self.torch_dtype = dtype
-
-    def forward(self, hidden_states):
-        hidden_gelu = self.act(self.wi_0(hidden_states))
-        hidden_linear = self.wi_1(hidden_states)
-        hidden_states = hidden_gelu * hidden_linear
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced
-        hidden_states = self.wo(hidden_states)
-        return hidden_states
-
-
-def get_int_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: List[int], num_bits: int) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
-    zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
-    max_values = torch.amax(weight, dim=reduction_axes, keepdims=True)
-    scale = (max_values - zero_point) / (2**num_bits - 1)
-    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
-    scale = torch.where(torch.abs(scale) < eps, eps, scale)
-    return scale, zero_point
-
-
-def get_int_scale_symmetric(weight: torch.FloatTensor, reduction_axes: List[int], num_bits: int) -> torch.FloatTensor:
-    w_abs_min = torch.abs(torch.amin(weight, dim=reduction_axes, keepdims=True))
-    w_max = torch.amax(weight, dim=reduction_axes, keepdims=True)
-    scale = torch.where(w_abs_min >= w_max, w_abs_min, -w_max) / (2 ** (num_bits - 1))
-    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
-    scale = torch.where(torch.abs(scale) < eps, eps, scale)
-    return scale
-
-
-def quantize_int(weight: torch.FloatTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, is_asym_mode: bool, num_bits: int, flatten: Optional[bool] = False) -> torch.ByteTensor:
-    dtype = torch.uint8 if is_asym_mode else torch.int8
-    level_low = 0 if is_asym_mode else -(2 ** (num_bits - 1))
-    level_high = 2**num_bits - 1 if is_asym_mode else 2 ** (num_bits - 1) - 1
-    if zero_point is not None:
-        compressed_weight = torch.sub(weight, zero_point).div_(scale)
-    else:
-        compressed_weight = torch.div(weight, scale)
-    compressed_weight = compressed_weight.round_().clamp_(level_low, level_high).to(dtype)
-    if flatten:
-        compressed_weight = compressed_weight.flatten(0,-2)
-    return compressed_weight
-
-
-def decompress_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
-    result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype)
-    if result_shape is not None:
-        result = result.reshape(result_shape)
-    return result
-
-
-def decompress_symmetric(input: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
-    result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
-    if result_shape is not None:
-        result = result.reshape(result_shape)
-    return result
-
-
-def decompress_int4_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
-    return decompress_asymmetric(unpack_uint4(input, shape), scale, zero_point, dtype, result_shape)
-
-
-def decompress_int4_symmetric(input: torch.Tensor, scale: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
-    return decompress_symmetric(unpack_int4(input, shape, dtype=scale.dtype), scale, dtype, result_shape)
-
-
-def pack_uint4(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
-    packed_tensor = tensor.contiguous().reshape(-1, 2)
-    packed_tensor = torch.bitwise_and(packed_tensor[..., ::2], 15) | packed_tensor[..., 1::2] << 4
-    return packed_tensor
-
-
-def pack_int4(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.int8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.int8 type is supported.")
-    tensor = tensor + 8
-    return pack_uint4(tensor.to(dtype=torch.uint8))
-
-
-def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size, transpose: Optional[bool] = False) -> torch.Tensor:
-    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
-    if transpose:
-        result = result.transpose(0,1)
-    return result
-
-
-def unpack_int4(packed_tensor: torch.Tensor, shape: torch.Size, dtype: Optional[torch.dtype] = torch.int8, transpose: Optional[bool] = False) -> torch.Tensor:
-    result = unpack_uint4(packed_tensor, shape).to(dtype=dtype).sub_(8)
-    if transpose:
-        result = result.transpose(0,1)
-    return result
-
-
-def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]:
-    input_scale = torch.div(input.abs().max(), 127)
-    input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8).flatten(0,-2)
-    scale = torch.mul(input_scale, scale)
-    return input, scale
-
-
-def int8_matmul(
-    input: torch.Tensor,
-    weight: torch.Tensor,
-    scale: torch.Tensor,
-    compressed_weight_shape: torch.Size,
-):
-    if compressed_weight_shape is not None:
-        weight = unpack_int4_compiled(weight, compressed_weight_shape, transpose=True)
-    return_dtype = input.dtype
-    output_shape = list(input.shape)
-    output_shape[-1] = weight.shape[-1]
-    input, scale = quantize_int8_matmul_input_compiled(input, scale)
-    return decompress_symmetric_compiled(torch._int_mm(input, weight), scale, return_dtype, output_shape) # pylint: disable=protected-access
-
-
-class linear_forward_int8_matmul():
-    def __func__(self, input) -> torch.FloatTensor:
-        if self.pre_ops["0"].skip_int8_matmul:
-            return torch.nn.Linear.forward(self, input)
-        result = int8_matmul(input, self.weight, self.pre_ops["0"].scale, getattr(self.pre_ops["0"], "compressed_weight_shape", None))
-        if self.bias is not None:
-            result.add_(self.bias)
-        return result
-
-
-class INT8AsymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        zero_point: torch.Tensor,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-    ):
-        super().__init__()
-        self.num_bits = 8
-        self.quantization_mode = "asymmetric"
-        self.scale = scale
-        self.zero_point = zero_point
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        if debug:
-            if torch.any((weight < 0) | (weight > 255)):
-                raise ValueError("Weight values are not in [0, 255].")
-        return weight.to(dtype=torch.uint8)
-
-    def forward(self, x, input=None, *args, return_decompressed_only=False): # pylint: disable=keyword-arg-before-vararg,unused-argument
-        result = decompress_asymmetric_compiled(x.weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)
-        if return_decompressed_only:
-            return result
-        else:
-            x.weight = result
-
-
-class INT8SymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-        use_int8_matmul: bool,
-    ):
-        super().__init__()
-        self.num_bits = 8
-        self.quantization_mode = "symmetric"
-        self.scale = scale
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-        self.use_int8_matmul = use_int8_matmul
-        self.skip_int8_matmul = False
-        self.input_scale = None
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        if debug:
-            if torch.any((weight < -128) | (weight > 127)):
-                raise ValueError("Weight values are not in [-128, 127].")
-        return weight.to(dtype=torch.int8)
-
-    def forward(self, x, input=None, *args, return_decompressed_only=False): # pylint: disable=unused-argument,keyword-arg-before-vararg
-        if self.use_int8_matmul:
-            if input is not None:
-                if torch.numel(input[0]) / input[0].shape[-1] < 32:
-                    self.skip_int8_matmul = True
-                else:
-                    self.skip_int8_matmul = False
-                    return
-            result = decompress_symmetric_compiled(x.weight.transpose(0,1), self.scale.unsqueeze(-1), self.result_dtype, self.result_shape)
-        else:
-            result = decompress_symmetric_compiled(x.weight, self.scale, self.result_dtype, self.result_shape)
-        if return_decompressed_only:
-            return result
-        else:
-            x.weight = result
-
-
-class INT4AsymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        zero_point: torch.Tensor,
-        compressed_weight_shape: torch.Size,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-    ):
-        super().__init__()
-        self.num_bits = 4
-        self.quantization_mode = "asymmetric"
-        self.scale = scale
-        self.zero_point = zero_point
-        self.compressed_weight_shape = compressed_weight_shape
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        if debug:
-            if torch.any((weight < 0) | (weight > 15)):
-                raise ValueError("Weight values are not in [0, 15].")
-        return pack_uint4(weight.to(dtype=torch.uint8))
-
-    def forward(self, x, input=None, *args, return_decompressed_only=False): # pylint: disable=unused-argument,keyword-arg-before-vararg
-        result = decompress_int4_asymmetric_compiled(x.weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape)
-        if return_decompressed_only:
-            return result
-        else:
-            x.weight = result
-
-
-class INT4SymmetricWeightsDecompressor(torch.nn.Module):
-    def __init__(
-        self,
-        scale: torch.Tensor,
-        compressed_weight_shape: torch.Size,
-        result_dtype: torch.dtype,
-        result_shape: torch.Size,
-        use_int8_matmul: bool,
-    ):
-        super().__init__()
-        self.num_bits = 4
-        self.quantization_mode = "symmetric"
-        self.scale = scale
-        self.compressed_weight_shape = compressed_weight_shape
-        self.result_dtype = result_dtype
-        self.result_shape = result_shape
-        self.use_int8_matmul = use_int8_matmul
-        self.skip_int8_matmul = False
-        self.input_scale = None
-
-    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        if debug:
-            if torch.any((weight < -8) | (weight > 7)):
-                raise ValueError("Tensor values are not in [-8, 7].")
-        return pack_int4(weight.to(dtype=torch.int8))
-
-    def forward(self, x, input=None, *arg, return_decompressed_only=False): # pylint: disable=keyword-arg-before-vararg,unused-argument
-        if self.use_int8_matmul:
-            if input is not None:
-                if torch.numel(input[0]) / input[0].shape[-1] < 32:
-                    self.skip_int8_matmul = True
-                else:
-                    self.skip_int8_matmul = False
-                    return
-            result = decompress_int4_symmetric_compiled(x.weight, self.scale.unsqueeze(-1), self.compressed_weight_shape, self.result_dtype, self.result_shape)
-        else:
-            result = decompress_int4_symmetric_compiled(x.weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape)
-        if return_decompressed_only:
-            return result
-        else:
-            x.weight = result
-
-
-if shared.opts.nncf_decompress_compile:
-    try:
-        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) # pylint: disable=protected-access
-        decompress_asymmetric_compiled = torch.compile(decompress_asymmetric, fullgraph=True)
-        decompress_symmetric_compiled = torch.compile(decompress_symmetric, fullgraph=True)
-        decompress_int4_asymmetric_compiled = torch.compile(decompress_int4_asymmetric, fullgraph=True)
-        decompress_int4_symmetric_compiled = torch.compile(decompress_int4_symmetric, fullgraph=True)
-        if devices.backend != "ipex": # pytorch uses the cpu device in torch._int_mm op with ipex + torch.compile
-            quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
-            unpack_int4_compiled = unpack_int4
-            int8_matmul = torch.compile(int8_matmul, fullgraph=True)
-        else:
-            quantize_int8_matmul_input_compiled = torch.compile(quantize_int8_matmul_input, fullgraph=True)
-            unpack_int4_compiled = torch.compile(unpack_int4, fullgraph=True)
-    except Exception as e:
-        shared.log.warning(f"Quantization: type=nncf Decompress using torch.compile is not available: {e}")
-        decompress_asymmetric_compiled = decompress_asymmetric
-        decompress_symmetric_compiled = decompress_symmetric
-        decompress_int4_asymmetric_compiled = decompress_int4_asymmetric
-        decompress_int4_symmetric_compiled = decompress_int4_symmetric
-        quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
-        unpack_int4_compiled = unpack_int4
-else:
-    decompress_asymmetric_compiled = decompress_asymmetric
-    decompress_symmetric_compiled = decompress_symmetric
-    decompress_int4_asymmetric_compiled = decompress_int4_asymmetric
-    decompress_int4_symmetric_compiled = decompress_int4_symmetric
-    quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
-    unpack_int4_compiled = unpack_int4
@@ -0,0 +1,828 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+from typing import Any, Dict, List, Tuple, Optional, Union
+from dataclasses import dataclass
+from enum import Enum
+import sys
+import torch
+from diffusers.quantizers.base import DiffusersQuantizer
+from diffusers.quantizers.quantization_config import QuantizationConfigMixin
+from diffusers.utils import get_module_from_name
+from accelerate.utils import CustomDtype
+from modules import devices, shared
+
+
+dtype_dict = {
+    "int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True},
+    "uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True},
+    "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False},
+    "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False},
+    "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False},
+    "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False},
+}
+
+quantized_matmul_dtypes = ("int8", "int6", "int4", "int2", "float8_e4m3fn", "float8_e5m2")
+if devices.backend in {"cpu", "openvino"}:
+    quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz")
+
+linear_types = ("Linear",)
+conv_types = ("Conv1d", "Conv2d", "Conv3d")
+conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d")
+allowed_types = linear_types + conv_types + conv_transpose_types
+
+
+class QuantizationMethod(str, Enum):
+    SDNQ = "sdnq"
+
+
+def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None, pre_mode=False): # pylint: disable=unused-argument
+    layer_class_name = layer.__class__.__name__
+    if layer_class_name in allowed_types:
+        is_conv_type = False
+        is_conv_transpose_type = False
+        is_linear_type = False
+        use_tensorwise_fp8_matmul = False
+        result_shape = None
+        if torch_dtype is None:
+            torch_dtype = devices.dtype
+
+        if layer_class_name in conv_types:
+            if not quant_conv:
+                return layer
+            reduction_axes = [i for i in range(layer.weight.ndim) if i != 0]
+            use_quantized_matmul = False
+            is_conv_type = True
+        elif layer_class_name in conv_transpose_types:
+            if not quant_conv:
+                return layer
+            reduction_axes = [i for i in range(layer.weight.ndim) if i != 1]
+            use_quantized_matmul = False
+            is_conv_transpose_type = True
+        else:
+            is_linear_type = True
+            reduction_axes = -1
+            output_channel_size, channel_size = layer.weight.shape
+            if use_quantized_matmul:
+                use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and channel_size >= 32 and output_channel_size >= 32
+                if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
+                    use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
+                    use_tensorwise_fp8_matmul = devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and float(torch.__version__[:3]) <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
+
+            if not use_quantized_matmul and (group_size > 0 or (dtype_dict[weights_dtype]["num_bits"] < 6 and group_size != -1)):
+                if group_size == 0:
+                    if dtype_dict[weights_dtype]["num_bits"] < 4:
+                        group_size = 32
+                    else:
+                        group_size = 64
+                    num_of_groups = channel_size // group_size
+
+                if group_size >= channel_size:
+                    group_size = channel_size
+                    num_of_groups = 1
+                else:
+                    num_of_groups = channel_size // group_size
+                    while channel_size % group_size != 0: # find something divisible
+                        num_of_groups -= 1
+                        if num_of_groups <= 1:
+                            group_size = channel_size
+                            num_of_groups = 1
+                            break
+                        group_size = channel_size / num_of_groups
+
+                if num_of_groups > 1:
+                    result_shape = layer.weight.shape
+                    new_shape = list(result_shape)
+                    last_dim_index = layer.weight.ndim
+                    new_shape[last_dim_index - 1 : last_dim_index] = (int(num_of_groups), int(group_size))
+                    layer.weight.data = layer.weight.reshape(new_shape)
+
+        layer.weight.requires_grad = False
+        if shared.opts.diffusers_offload_mode in {"none", "model"}:
+            return_device = devices.device
+        elif pre_mode:
+            if shared.opts.device_map == "gpu":
+                return_device = devices.device
+            elif shared.opts.sdnq_quantize_with_gpu:
+                return_device = devices.cpu
+            else:
+                return_device = layer.weight.device
+        else:
+            return_device = layer.weight.device
+        if not pre_mode:
+            if shared.opts.sdnq_quantize_with_gpu:
+                layer.weight.data = layer.weight.to(devices.device).to(dtype=torch.float32)
+            else:
+                layer.weight.data = layer.weight.to(dtype=torch.float32)
+
+        if dtype_dict[weights_dtype]["is_unsigned"]:
+            scale, zero_point = get_scale_asymmetric(layer.weight, reduction_axes, weights_dtype)
+        else:
+            scale = get_scale_symmetric(layer.weight, reduction_axes, weights_dtype)
+            zero_point = None
+        layer.weight.data = quantize_weight(layer.weight, scale, zero_point, weights_dtype)
+
+        if not shared.opts.sdnq_decompress_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]):
+            scale = scale.to(torch_dtype)
+            if zero_point is not None:
+                zero_point = zero_point.to(torch_dtype)
+
+        if use_quantized_matmul:
+            scale = scale.transpose(0,1)
+            if dtype_dict[weights_dtype]["num_bits"] == 8:
+                layer.weight.data = layer.weight.transpose(0,1)
+            if not dtype_dict[weights_dtype]["is_integer"]:
+                stride = layer.weight.stride()
+                if stride[0] > stride[1] and stride[1] == 1:
+                    layer.weight.data = layer.weight.t().contiguous().t()
+                if not use_tensorwise_fp8_matmul:
+                    scale = scale.to(torch.float32)
+
+        layer.sdnq_decompressor = decompressor_dict[weights_dtype](
+            scale=scale,
+            zero_point=zero_point,
+            compressed_weight_shape=layer.weight.shape,
+            result_dtype=torch_dtype,
+            result_shape=result_shape,
+            weights_dtype=weights_dtype,
+            use_quantized_matmul=use_quantized_matmul,
+        )
+        layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight).to(return_device)
+        layer.sdnq_decompressor = layer.sdnq_decompressor.to(return_device)
+
+        if is_linear_type:
+            if use_quantized_matmul:
+                if dtype_dict[weights_dtype]["is_integer"]:
+                    layer.forward = quantized_linear_forward_int8_matmul
+                else:
+                    if use_tensorwise_fp8_matmul:
+                        layer.forward = quantized_linear_forward_fp8_matmul_tensorwise
+                    else:
+                        layer.forward = quantized_linear_forward_fp8_matmul
+            else:
+                layer.forward = quantized_linear_forward
+        elif is_conv_type:
+            layer.forward = quantized_conv_forward
+        elif is_conv_transpose_type:
+            if layer_class_name.endswith("1d"):
+                layer.forward = quantized_conv_transpose_1d_forward
+            elif layer_class_name.endswith("2d"):
+                layer.forward = quantized_conv_transpose_2d_forward
+            elif layer_class_name.endswith("3d"):
+                layer.forward = quantized_conv_transpose_3d_forward
+        layer.forward = layer.forward.__get__(layer, layer.__class__)
+        devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}")
+    return layer
+
+
+def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None):
+    has_children = list(model.children())
+    if not has_children:
+        return model
+    for module_param_name, module in model.named_children():
+        if hasattr(module, "weight") and module.weight is not None:
+            module = sdnq_quantize_layer(
+                module,
+                weights_dtype=weights_dtype,
+                torch_dtype=torch_dtype,
+                group_size=group_size,
+                quant_conv=quant_conv,
+                use_quantized_matmul=use_quantized_matmul,
+                param_name=module_param_name,
+            )
+        module = apply_sdnq_to_module(
+                module,
+                weights_dtype=weights_dtype,
+                torch_dtype=torch_dtype,
+                group_size=group_size,
+                quant_conv=quant_conv,
+                use_quantized_matmul=use_quantized_matmul,
+                param_name=module_param_name,
+            )
+    return model
+
+
+def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+    zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
+    scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"])
+    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
+    scale = torch.where(torch.abs(scale) < eps, eps, scale)
+    if dtype_dict[weights_dtype]["min"] != 0:
+        zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"]))
+    return scale, zero_point
+
+
+def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> torch.FloatTensor:
+    abs_min_values = torch.amin(weight, dim=reduction_axes, keepdims=True).abs_()
+    max_values = torch.amax(weight, dim=reduction_axes, keepdims=True)
+    scale = torch.where(abs_min_values >= max_values, abs_min_values, -max_values).div_(dtype_dict[weights_dtype]["max"])
+    eps = torch.finfo(scale.dtype).eps # prevent divison by 0
+    scale = torch.where(torch.abs(scale) < eps, eps, scale)
+    return scale
+
+
+def quantize_weight(weight: torch.FloatTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, weights_dtype: str) -> torch.ByteTensor:
+    if zero_point is not None:
+        compressed_weight = torch.sub(weight, zero_point).div_(scale)
+    else:
+        compressed_weight = torch.div(weight, scale)
+    if dtype_dict[weights_dtype]["is_integer"]:
+        compressed_weight.round_()
+    compressed_weight = compressed_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
+    return compressed_weight
+
+
+def decompress_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor:
+    result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype)
+    if result_shape is not None:
+        result = result.reshape(result_shape)
+    return result
+
+
+def decompress_symmetric(input: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.Tensor:
+    if skip_quantized_matmul:
+        result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
+    else:
+        result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
+    if result_shape is not None:
+        result = result.reshape(result_shape)
+    return result
+
+
+def decompress_packed_int_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.Tensor:
+    return decompress_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape)
+
+
+def decompress_packed_int_symmetric(input: torch.Tensor, scale: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.Tensor:
+    if skip_quantized_matmul:
+        return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
+    else:
+        return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
+
+
+def pack_int_symetric(tensor: torch.Tensor, weights_dtype: str) -> torch.Tensor:
+    return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"]).sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
+
+
+def unpack_int_symetric(packed_tensor: torch.Tensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.Tensor:
+    if dtype is None:
+        dtype = dtype_dict[weights_dtype]["torch_dtype"]
+    result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"])
+    if transpose:
+        result = result.transpose(0,1)
+    return result
+
+
+def pack_uint6(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 4)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)),
+            torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)),
+            torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint4(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 2)
+    packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4))
+    return packed_tensor
+
+
+def pack_uint2(tensor: torch.Tensor) -> torch.Tensor:
+    if tensor.dtype != torch.uint8:
+        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+    packed_tensor = tensor.contiguous().reshape(-1, 4)
+    packed_tensor = torch.bitwise_or(
+        torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)),
+        torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
+    )
+    return packed_tensor
+
+
+def unpack_uint6(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 63),
+            torch.bitwise_and(packed_tensor[:, 1], 63),
+            torch.bitwise_and(packed_tensor[:, 2], 63),
+            torch.bitwise_or(
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48),
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 12),
+                ),
+                torch.bitwise_right_shift(packed_tensor[:, 2], 6)
+            )
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
+    return result
+
+
+def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor, 3),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3),
+            torch.bitwise_right_shift(packed_tensor, 6),
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448)
+    input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
+    input_scale = input_scale.to(torch.float32)
+    return input, input_scale
+
+
+def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448)
+    input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
+    scale = torch.mul(input_scale, scale)
+    if scale.dtype == torch.float16: # fp16 will overflow
+        scale = scale.to(dtype=torch.float32)
+    return input, scale
+
+
+def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 127)
+    input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8)
+    scale = torch.mul(input_scale, scale)
+    if scale.dtype == torch.float16: # fp16 will overflow
+        scale = scale.to(dtype=torch.float32)
+    return input, scale
+
+
+def fp8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    output_shape = list(input.shape)
+    output_shape[-1] = weight.shape[-1]
+    input, input_scale = quantize_fp8_matmul_input(input)
+    return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape)
+
+
+# sm89 doesn't support row wise scale in Windows
+def fp8_matmul_tensorwise(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    output_shape = list(input.shape)
+    output_shape[-1] = weight.shape[-1]
+    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
+    input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
+    result = decompress_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape)
+    if bias is not None:
+        result.add_(bias)
+    return result
+
+
+def int8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    compressed_weight_shape: torch.Size,
+    weights_dtype: str,
+) -> torch.FloatTensor:
+    if compressed_weight_shape is not None:
+        weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
+    return_dtype = input.dtype
+    output_shape = list(input.shape)
+    output_shape[-1] = weight.shape[-1]
+    input, scale = quantize_int8_matmul_input(input, scale)
+    result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
+    if bias is not None:
+        result.add_(bias)
+    return result
+
+
+def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[-1] < 32:
+        return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
+    return fp8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale)
+
+
+def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[-1] < 32:
+        return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
+    return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_decompressor.scale)
+
+
+def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[-1] < 32:
+        return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
+    return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None), self.sdnq_decompressor.weights_dtype)
+
+
+def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias)
+
+
+def quantized_conv_forward(self, input) -> torch.FloatTensor:
+    return self._conv_forward(input, self.sdnq_decompressor(self.weight), self.bias)
+
+
+def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
+    return torch.nn.functional.conv_transpose1d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
+
+
+def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
+    return torch.nn.functional.conv_transpose2d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
+
+
+def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+    output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
+    return torch.nn.functional.conv_transpose3d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
+
+
+class AsymmetricWeightsDecompressor(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        **kwargs,
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = False
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+        self.register_buffer("zero_point", zero_point)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
+
+    def forward(self, weight, **kwargs):
+        return decompress_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)
+
+
+class SymmetricWeightsDecompressor(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        use_quantized_matmul: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = use_quantized_matmul
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
+
+    def forward(self, weight, skip_quantized_matmul=False, **kwargs):
+        return decompress_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
+
+
+class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        compressed_weight_shape: torch.Size,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        **kwargs,
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = False
+        self.compressed_weight_shape = compressed_weight_shape
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+        self.register_buffer("zero_point", zero_point)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]))
+
+    def forward(self, weight, **kwargs):
+        return decompress_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
+
+
+class PackedINTSymmetricWeightsDecompressor(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.Tensor,
+        compressed_weight_shape: torch.Size,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        weights_dtype: str,
+        use_quantized_matmul: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = use_quantized_matmul
+        self.compressed_weight_shape = compressed_weight_shape
+        self.result_dtype = result_dtype
+        self.result_shape = result_shape
+        self.register_buffer("scale", scale)
+
+    def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return pack_int_symetric(weight, self.weights_dtype)
+
+    def forward(self, weight, skip_quantized_matmul=False, **kwargs):
+        return decompress_packed_int_symmetric_compiled(weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
+
+
+decompressor_dict = {
+    "int8": SymmetricWeightsDecompressor,
+    "uint8": AsymmetricWeightsDecompressor,
+    "int6": PackedINTSymmetricWeightsDecompressor,
+    "uint6": PackedINTAsymmetricWeightsDecompressor,
+    "int4": PackedINTSymmetricWeightsDecompressor,
+    "uint4": PackedINTAsymmetricWeightsDecompressor,
+    "int2": PackedINTSymmetricWeightsDecompressor,
+    "uint2": PackedINTAsymmetricWeightsDecompressor,
+    "uint1": AsymmetricWeightsDecompressor,
+    "float8_e4m3fn": SymmetricWeightsDecompressor,
+    "float8_e4m3fnuz": SymmetricWeightsDecompressor,
+    "float8_e5m2": SymmetricWeightsDecompressor,
+    "float8_e5m2fnuz": SymmetricWeightsDecompressor,
+}
+
+
+packed_int_function_dict = {
+    "int6": {"pack": pack_uint6, "unpack": unpack_uint6},
+    "uint6": {"pack": pack_uint6, "unpack": unpack_uint6},
+    "int4": {"pack": pack_uint4, "unpack": unpack_uint4},
+    "uint4": {"pack": pack_uint4, "unpack": unpack_uint4},
+    "int2": {"pack": pack_uint2, "unpack": unpack_uint2},
+    "uint2": {"pack": pack_uint2, "unpack": unpack_uint2},
+}
+
+
+class SDNQQuantizer(DiffusersQuantizer):
+    r"""
+    Diffusers Quantizer for SDNQ
+    """
+
+    requires_parameters_quantization = True
+    use_keep_in_fp32_modules = True
+    requires_calibration = False
+    required_packages = None
+    torch_dtype = None
+
+    def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation
+        super().__init__(quantization_config, **kwargs)
+
+    def check_if_quantized_param(
+        self,
+        model,
+        param_value: "torch.Tensor",
+        param_name: str,
+        state_dict: Dict[str, Any],
+        **kwargs,
+    ):
+        if param_name.endswith(".weight"):
+            split_param_name = param_name.split(".")
+            if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert):
+                layer_class_name = get_module_from_name(model, param_name)[0].__class__.__name__
+                if layer_class_name in allowed_types:
+                    if layer_class_name in conv_types or layer_class_name in conv_transpose_types:
+                        if self.quantization_config.quant_conv:
+                            return True
+                    else:
+                        return True
+        param_value.data = param_value.clone() # safetensors is unable to release the cpu memory without this
+        return False
+
+    def check_quantized_param(self, *args, **kwargs) -> bool:
+        """
+        needed for transformers compatibilty, returns self.check_if_quantized_param
+        """
+        return self.check_if_quantized_param(*args, **kwargs)
+
+    def create_quantized_param( # pylint: disable=arguments-differ
+        self,
+        model,
+        param_value: torch.FloatTensor,
+        param_name: str,
+        target_device: torch.device,
+        state_dict: Dict[str, Any], # pylint: disable=unused-argument
+        unexpected_keys: List[str], # pylint: disable=unused-argument
+        **kwargs,
+    ):
+        # load the model params to target_device first
+        layer, _ = get_module_from_name(model, param_name)
+        if shared.opts.sdnq_quantize_with_gpu:
+            if param_value.dtype == torch.float32 and devices.same_device(param_value.device, devices.device):
+                param_value = param_value.clone()
+            else:
+                param_value = param_value.to(devices.device).to(dtype=torch.float32)
+        else:
+            if param_value.dtype == torch.float32 and devices.same_device(param_value.device, target_device):
+                param_value = param_value.clone()
+            else:
+                param_value = param_value.to(target_device).to(dtype=torch.float32)
+        layer.weight = torch.nn.Parameter(param_value, requires_grad=False)
+        layer = sdnq_quantize_layer(
+            layer,
+            weights_dtype=self.quantization_config.weights_dtype,
+            torch_dtype=self.torch_dtype,
+            group_size=self.quantization_config.group_size,
+            quant_conv=self.quantization_config.quant_conv,
+            use_quantized_matmul=self.quantization_config.use_quantized_matmul,
+            param_name=param_name,
+            pre_mode=True,
+        )
+
+    def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+        max_memory = {key: val * 0.80 for key, val in max_memory.items()}
+        return max_memory
+
+    def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed
+        return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"]
+
+    def update_torch_dtype(self, torch_dtype: torch.dtype = None) -> torch.dtype:
+        if torch_dtype is None:
+            torch_dtype = devices.dtype
+        self.torch_dtype = torch_dtype
+        return torch_dtype
+
+    def _process_model_before_weight_loading( # pylint: disable=arguments-differ
+        self,
+        model,
+        device_map, # pylint: disable=unused-argument
+        keep_in_fp32_modules: List[str] = [],
+        **kwargs,
+    ):
+        model.config.quantization_config = self.quantization_config
+        self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
+        if not isinstance(self.modules_to_not_convert, list):
+            self.modules_to_not_convert = [self.modules_to_not_convert]
+        if keep_in_fp32_modules is not None:
+            self.modules_to_not_convert.extend(keep_in_fp32_modules)
+
+    def _process_model_after_weight_loading(self, model, **kwargs):
+        if shared.opts.diffusers_offload_mode != "none":
+            model = model.to(devices.cpu)
+        devices.torch_gc(force=True)
+        return model
+
+    def get_cuda_warm_up_factor(self):
+        return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"]
+
+    def update_tp_plan(self, config):
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return config
+
+    def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return unexpected_keys
+
+    def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return missing_keys
+
+    def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: # pylint: disable=unused-argument
+        """
+        needed for transformers compatibilty, no-op function
+        """
+        return expected_keys
+
+    @property
+    def is_trainable(self):
+        return False
+
+    @property
+    def is_serializable(self):
+        return False
+
+
+@dataclass
+class SDNQConfig(QuantizationConfigMixin):
+    """
+    This is a wrapper class about all possible attributes and features that you can play with a model that has been
+    loaded using `sdnq`.
+
+    Args:
+        weights_dtype (`str`, *optional*, defaults to `"int8"`):
+            The target dtype for the weights after quantization. Supported values are:
+            ("int8", "uint8", "int6", "uint6", "int4", "uint4", "uint2", "uint1", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
+       modules_to_not_convert (`list`, *optional*, default to `None`):
+            The list of modules to not quantize, useful for quantizing models that explicitly require to have some
+            modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+    """
+
+    def __init__( # pylint: disable=super-init-not-called
+        self,
+        weights_dtype: str = "int8",
+        group_size: int = 0,
+        quant_conv: bool = False,
+        use_quantized_matmul: bool = False,
+        modules_to_not_convert: Optional[List[str]] = None,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        self.weights_dtype = weights_dtype
+        self.quant_method = QuantizationMethod.SDNQ
+        self.group_size = group_size
+        self.quant_conv = quant_conv
+        self.use_quantized_matmul = use_quantized_matmul
+        self.modules_to_not_convert = modules_to_not_convert
+        self.post_init()
+        self.is_integer = dtype_dict[self.weights_dtype]["is_integer"]
+
+    def post_init(self):
+        r"""
+        Safety checker that arguments are correct
+        """
+        accepted_weights = ["int8", "uint8", "int6", "uint6", "int4", "uint4", "uint2", "uint1", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"]
+        if self.weights_dtype not in accepted_weights:
+            raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")
+
+
+class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class
+    def __init__(self, T5DenseGatedActDense, dtype):
+        super().__init__()
+        self.wi_0 = T5DenseGatedActDense.wi_0
+        self.wi_1 = T5DenseGatedActDense.wi_1
+        self.wo = T5DenseGatedActDense.wo
+        self.dropout = T5DenseGatedActDense.dropout
+        self.act = T5DenseGatedActDense.act
+        self.torch_dtype = dtype
+
+    def forward(self, hidden_states):
+        hidden_gelu = self.act(self.wi_0(hidden_states))
+        hidden_linear = self.wi_1(hidden_states)
+        hidden_states = hidden_gelu * hidden_linear
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced
+        hidden_states = self.wo(hidden_states)
+        return hidden_states
+
+
+if shared.opts.sdnq_decompress_compile:
+    try:
+        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
+        decompress_asymmetric_compiled = torch.compile(decompress_asymmetric, fullgraph=True)
+        decompress_symmetric_compiled = torch.compile(decompress_symmetric, fullgraph=True)
+        decompress_packed_int_asymmetric_compiled = torch.compile(decompress_packed_int_asymmetric, fullgraph=True)
+        decompress_packed_int_symmetric_compiled = torch.compile(decompress_packed_int_symmetric, fullgraph=True)
+        fp8_matmul = torch.compile(fp8_matmul, fullgraph=True)
+        fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True)
+        int8_matmul = torch.compile(int8_matmul, fullgraph=True)
+    except Exception as e:
+        shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}")
+        decompress_asymmetric_compiled = decompress_asymmetric
+        decompress_symmetric_compiled = decompress_symmetric
+        decompress_packed_int_asymmetric_compiled = decompress_packed_int_asymmetric
+        decompress_packed_int_symmetric_compiled = decompress_packed_int_symmetric
+else:
+    decompress_asymmetric_compiled = decompress_asymmetric
+    decompress_symmetric_compiled = decompress_symmetric
+    decompress_packed_int_asymmetric_compiled = decompress_packed_int_asymmetric
+    decompress_packed_int_symmetric_compiled = decompress_packed_int_symmetric
@@ -70,8 +70,13 @@ 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 'int8' in name.lower():
-        from modules.model_quant import create_nncf_config
-        quantization_config = create_nncf_config(kwargs=None, allow_nncf=True, module="any")
+        from modules.model_quant import create_sdnq_config
+        quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='int8')
+        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+
+    elif 'uint4' in name.lower():
+        from modules.model_quant import create_sdnq_config
+        quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='uint4')
        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)

    elif 'qint4' in name.lower():
@@ -167,7 +167,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:

        shared.prompt_styles.apply_styles_to_extra(p)
        shared.prompt_styles.extract_comments(p)
-        if shared.opts.cuda_compile_backend == 'none':
+        if 'Model' not in shared.opts.cuda_compile:
            token_merge.apply_token_merging(p.sd_model)
            from modules import sd_hijack_freeu, para_attention, teacache
            sd_hijack_freeu.apply_freeu(p, not shared.native)
@@ -8,7 +8,6 @@ import torch
 import numpy as np
 import cv2
 from PIL import Image
-from skimage import exposure
 from blendmodes.blend import blendLayers, BlendType
 from modules import shared, devices, images, sd_models, sd_samplers, sd_hijack_hypertile, processing_vae, timer

@@ -33,6 +32,7 @@ def setup_color_correction(image):


 def apply_color_correction(correction, original_image):
+    from skimage import exposure
    shared.log.debug(f"Applying color correction: correction={correction.shape} image={original_image}")
    np_image = np.asarray(original_image)
    np_recolor = cv2.cvtColor(np_image, cv2.COLOR_RGB2LAB)
@@ -545,11 +545,15 @@ def apply_circular(enable: bool, model):
    if getattr(model, 'texture_tiling', False) == enable:
        return
    try:
+        i = 0
        for layer in [layer for layer in model.unet.modules() if type(layer) is torch.nn.Conv2d]:
+            i += 1
            layer.padding_mode = 'circular' if enable else 'zeros'
        for layer in [layer for layer in model.vae.modules() if type(layer) is torch.nn.Conv2d]:
+            i += 1
            layer.padding_mode = 'circular' if enable else 'zeros'
        model.texture_tiling = enable
+        shared.log.debug(f'Apply texture tiling: enabled={enable} layers={i} cls={model.__class__.__name__} ')
    except Exception as e:
        debug(f"Diffusers tiling failed: {e}")

@@ -161,10 +161,10 @@ def apply_function_to_model(sd_model, function, options, op=None):
            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
+            if op == "sdnq" 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__:
+            if op == "sdnq" and "StableCascade" in sd_model.__class__.__name__:
                sd_model.prior_pipe.prior.clip_txt_pooled_mapper = backup_clip_txt_pooled_mapper
    if "TE" in options:
        if hasattr(sd_model, 'text_encoder') and hasattr(sd_model.text_encoder, 'config'):
@@ -9,7 +9,7 @@ from modules import shared, devices, processing, images, sd_vae_approx, sd_vae_t

 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
 approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE": 3 }
-flow_models = ['f1', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2', 'cogview4']
+flow_models = ['f1', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2', 'cogview4', 'h1']
 warned = False
 queue_lock = threading.Lock()

@@ -66,7 +66,7 @@ config = {
    # prediction_type is ideally set in model as well, but it maybe needed that we do auto-detect of model type in the future
    'All': { 'num_train_timesteps': 1000, 'beta_start': 0.0001, 'beta_end': 0.02, 'beta_schedule': 'linear', 'prediction_type': 'epsilon' },

-    'UniPC': { 'predict_x0': True, 'sample_max_value': 1.0, 'solver_order': 2, 'solver_type': 'bh2', 'thresholding': False, 'use_beta_sigmas': False, 'use_exponential_sigmas': False, 'use_flow_sigmas': False, 'use_karras_sigmas': False, 'lower_order_final': True, 'timestep_spacing': 'linspace', 'final_sigmas_type': 'zero', 'rescale_betas_zero_snr': False },
+    'UniPC': { 'flow_shift': 1, 'predict_x0': True, 'sample_max_value': 1.0, 'solver_order': 2, 'solver_type': 'bh2', 'thresholding': False, 'use_beta_sigmas': False, 'use_exponential_sigmas': False, 'use_flow_sigmas': False, 'use_karras_sigmas': False, 'lower_order_final': True, 'timestep_spacing': 'linspace', 'final_sigmas_type': 'zero', 'rescale_betas_zero_snr': False },
    '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 },

    'Euler': { 'steps_offset': 0, 'interpolation_type': "linear", 'rescale_betas_zero_snr': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'use_beta_sigmas': False, 'use_exponential_sigmas': False, 'use_karras_sigmas': False },
@@ -269,6 +269,8 @@ class DiffusionSampler:
            self.config['beta_end'] = shared.opts.schedulers_beta_end
        if 'shift' in self.config:
            self.config['shift'] = shared.opts.schedulers_shift if shared.opts.schedulers_shift > 0 else 3
+        if 'flow_shift' in self.config:
+            self.config['flow_shift'] = shared.opts.schedulers_shift if shared.opts.schedulers_shift > 0 else 3
        if 'use_dynamic_shifting' in self.config:
            self.config['use_dynamic_shifting'] = True if shared.opts.schedulers_shift == 0 else shared.opts.schedulers_dynamic_shift
        if 'use_beta_sigmas' in self.config and 'sigma_schedule' in self.config:
@@ -369,17 +369,12 @@ def get_default_modes():
        default_offload_mode = "sequential"
        default_diffusers_offload_min_gpu_memory = 0

-    if devices.backend == "directml": # Force BMM for DirectML instead of SDP
-        default_cross_attention = "Dynamic Attention BMM" if native else "Sub-quadratic"
-    elif devices.backend == "cpu":
-        default_cross_attention = "Scaled-Dot-Product" if native else "Doggettx's"
-    elif devices.backend == "mps":
-        default_cross_attention = "Scaled-Dot-Product" if native else "Doggettx's"
-    else: # cuda, rocm, zluda, ipex, openvino
-        default_cross_attention = "Scaled-Dot-Product"
+    default_cross_attention = "Scaled-Dot-Product"

    if devices.backend == "zluda":
-        default_sdp_options =  ['Math attention', 'Flash attention', 'Dynamic attention']
+        default_sdp_options = ['Flash attention', 'Math attention', 'Dynamic attention']
+    elif devices.backend in {"rocm", "directml", "cpu", "mps"}:
+        default_sdp_options = ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention']
    else:
        default_sdp_options = ['Flash attention', 'Memory attention', 'Math attention']

@@ -510,49 +505,47 @@ options_templates.update(options_section(('backends', "Backend Settings"), {

    "openvino_sep": OptionInfo("<h2>OpenVINO</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
    "openvino_devices": OptionInfo([], "OpenVINO devices to use", gr.CheckboxGroup, {"choices": get_openvino_device_list() if cmd_opts.use_openvino else [], "visible": cmd_opts.use_openvino}), # pylint: disable=E0606
-    "openvino_accuracy": OptionInfo("performance", "OpenVINO accuracy mode", gr.Radio, {"choices": ['performance', 'accuracy'], "visible": cmd_opts.use_openvino}),
+    "openvino_accuracy": OptionInfo("performance", "OpenVINO accuracy mode", gr.Radio, {"choices": ["performance", "accuracy"], "visible": cmd_opts.use_openvino}),
    "openvino_disable_model_caching": OptionInfo(True, "OpenVINO disable model caching", gr.Checkbox, {"visible": cmd_opts.use_openvino}),
    "openvino_disable_memory_cleanup": OptionInfo(True, "OpenVINO disable memory cleanup after compile", gr.Checkbox, {"visible": cmd_opts.use_openvino}),

    "directml_sep": OptionInfo("<h2>DirectML</h2>", "", gr.HTML, {"visible": devices.backend == "directml"}),
-    "directml_memory_provider": OptionInfo(default_memory_provider, 'DirectML memory stats provider', gr.Radio, {"choices": memory_providers, "visible": devices.backend == "directml"}),
+    "directml_memory_provider": OptionInfo(default_memory_provider, "DirectML memory stats provider", gr.Radio, {"choices": memory_providers, "visible": devices.backend == "directml"}),
    "directml_catch_nan": OptionInfo(False, "DirectML retry ops for NaN", gr.Checkbox, {"visible": devices.backend == "directml"}),
 }))

-options_templates.update(options_section(('quantization', "Quantization Settings"), {
+options_templates.update(options_section(("quantization", "Quantization Settings"), {
+    "sdnq_quantize_sep": OptionInfo("<h2>SDNQ: SD.Next Quantization</h2>", "", gr.HTML),
+    "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
+    "sdnq_quantize_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ["pre", "post"], "visible": native}),
+    "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "int6", "uint4", "float8_e4m3fn", "uint8", "uint6", "int4", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "int2", "uint2", "uint1"], "visible": native}),
+    "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}),
+    "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}),
+    "sdnq_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}),
+    "sdnq_decompress_compile": OptionInfo(devices.has_triton(), "Decompress using torch.compile", gr.Checkbox, {"visible": native}),
+    "sdnq_use_quantized_matmul": OptionInfo(False, "Use Quantized MatMul", gr.Checkbox, {"visible": native}),
+    "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize with the GPU", gr.Checkbox, {"visible": native}),
+    "sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
+
    "bnb_quantization_sep": OptionInfo("<h2>BitsAndBytes</h2>", "", gr.HTML),
    "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
-    "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ['nf4', 'fp8', 'fp4'], "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}),

-    "nncf_compress_sep": OptionInfo("<h2>NNCF: Neural Network Compression Framework</h2>", "", gr.HTML),
-    "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
-    "nncf_compress_mode": OptionInfo("post", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native and not cmd_opts.use_openvino}),
-    "nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8', 'INT8_SYM', 'INT4', 'INT4_SYM']}),
-    "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": 4096, "step": 1, "visible": native}),
-    "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE"], "visible": cmd_opts.use_openvino}),
-    "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}),
-    "nncf_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native and not cmd_opts.use_openvino}),
-    "nncf_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native and not cmd_opts.use_openvino}),
-    "nncf_decompress_compile": OptionInfo(devices.has_triton(), "Decompress using torch.compile", gr.Checkbox, {"visible": native and not cmd_opts.use_openvino}),
-    "nncf_decompress_int8_matmul": OptionInfo(False, "Use direct INT8 MatMul", gr.Checkbox, {"visible": native and not cmd_opts.use_openvino}),
-    "nncf_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native and not cmd_opts.use_openvino}),
-
    "quanto_quantization_sep": OptionInfo("<h2>Optimum Quanto</h2>", "", gr.HTML),
    "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
    "quanto_quantization_type": OptionInfo("int8", "Quantization weights type", gr.Dropdown, {"choices": ["float8", "int8", "int4", "int2"], "visible": native}),

    "optimum_quanto_sep": OptionInfo("<h2>Optimum Quanto: post-load</h2>", "", gr.HTML),
    "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "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_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 in post mode", gr.Checkbox, {"visible": native}),

    "torchao_sep": OptionInfo("<h2>TorchAO</h2>", "", gr.HTML),
    "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "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}),
+    "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}),

    "layerwise_quantization_sep": OptionInfo("<h2>Layerwise Casting</h2>", "", gr.HTML),
    "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "TE"], "visible": native}),
@@ -563,6 +556,14 @@ options_templates.update(options_section(('quantization', "Quantization Settings
    "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
    "nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox, {"visible": native}),
    "nunchaku_offload": OptionInfo(False, "Nunchaku offloading", gr.Checkbox, {"visible": native}),
+
+    "nncf_compress_sep": OptionInfo("<h2>NNCF: Neural Network Compression Framework</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
+    "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": cmd_opts.use_openvino}),
+    "nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ["INT8", "INT4_ASYM", "INT8_SYM", "INT4_SYM", "NF4"], "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": 4096, "step": 1, "visible": cmd_opts.use_openvino}),
+    "nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE"], "visible": cmd_opts.use_openvino}),
+    "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ["INT8", "FP8_E4M3", "FP8_E5M2"], "visible": cmd_opts.use_openvino}),
 }))

 options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
@@ -1,10 +1,11 @@
 from .teacache_flux import teacache_flux_forward
+from .teacache_hidream import teacache_hidream_forward
 from .teacache_ltx import teacache_ltx_forward
 from .teacache_mochi import teacache_mochi_forward
 from .teacache_cogvideox import teacache_cog_forward


-supported_models = ['Flux', 'CogVideoX', 'Mochi', 'LTX']
+supported_models = ['Flux', 'CogVideoX', 'Mochi', 'LTX', 'HiDream']


 def apply_teacache(p):
@@ -22,4 +23,7 @@ def apply_teacache(p):
    shared.sd_model.transformer.__class__.accumulated_rel_l1_distance = 0
    shared.sd_model.transformer.__class__.previous_modulated_input = None
    shared.sd_model.transformer.__class__.previous_residual = None
+    if shared.sd_model.__class__.__name__.startswith('HiDream'):
+        shared.sd_model.transformer.__class__.ret_steps = p.steps * 0.1
+        shared.sd_model.transformer.__class__.coefficients = [-3.13605009e+04, -7.12425503e+02, 4.91363285e+01, 8.26515490e+00, 1.08053901e-01]
    shared.log.info(f'Transformers cache: type=teacache cls={shared.sd_model.__class__.__name__} thresh={shared.opts.teacache_thresh}')
@@ -0,0 +1,264 @@
+from typing import Any, Dict, List, Optional, Tuple
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import logging, deprecate, USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+
+import torch
+import numpy as np
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def teacache_hidream_forward(
+    self,
+    hidden_states: torch.Tensor,
+    timesteps: torch.LongTensor = None,
+    encoder_hidden_states_t5: torch.Tensor = None,
+    encoder_hidden_states_llama3: torch.Tensor = None,
+    pooled_embeds: torch.Tensor = None,
+    img_ids: Optional[torch.Tensor] = None,
+    img_sizes: Optional[List[Tuple[int, int]]] = None,
+    hidden_states_masks: Optional[torch.Tensor] = None,
+    attention_kwargs: Optional[Dict[str, Any]] = None,
+    return_dict: bool = True,
+    **kwargs,
+):
+    encoder_hidden_states = kwargs.get("encoder_hidden_states", None)
+
+    if encoder_hidden_states is not None:
+        deprecation_message = "The `encoder_hidden_states` argument is deprecated. Please use `encoder_hidden_states_t5` and `encoder_hidden_states_llama3` instead."
+        deprecate("encoder_hidden_states", "0.35.0", deprecation_message)
+        encoder_hidden_states_t5 = encoder_hidden_states[0]
+        encoder_hidden_states_llama3 = encoder_hidden_states[1]
+
+    if img_ids is not None and img_sizes is not None and hidden_states_masks is None:
+        deprecation_message = (
+            "Passing `img_ids` and `img_sizes` with unpachified `hidden_states` is deprecated and will be ignored."
+        )
+        deprecate("img_ids", "0.35.0", deprecation_message)
+
+    if hidden_states_masks is not None and (img_ids is None or img_sizes is None):
+        raise ValueError("if `hidden_states_masks` is passed, `img_ids` and `img_sizes` must also be passed.")
+    elif hidden_states_masks is not None and hidden_states.ndim != 3:
+        raise ValueError(
+            "if `hidden_states_masks` is passed, `hidden_states` must be a 3D tensors with shape (batch_size, patch_height * patch_width, patch_size * patch_size * channels)"
+        )
+
+    if attention_kwargs is not None:
+        attention_kwargs = attention_kwargs.copy()
+        lora_scale = attention_kwargs.pop("scale", 1.0)
+    else:
+        lora_scale = 1.0
+
+    if USE_PEFT_BACKEND:
+        # weight the lora layers by setting `lora_scale` for each PEFT layer
+        scale_lora_layers(self, lora_scale)
+    else:
+        if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+            logger.warning(
+                "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+            )
+
+    # spatial forward
+    batch_size = hidden_states.shape[0]
+    hidden_states_type = hidden_states.dtype
+
+    # Patchify the input
+    if hidden_states_masks is None:
+        hidden_states, hidden_states_masks, img_sizes, img_ids = self.patchify(hidden_states)
+
+    # Embed the hidden states
+    hidden_states = self.x_embedder(hidden_states)
+
+    # 0. time
+    timesteps = self.t_embedder(timesteps, hidden_states_type)
+    p_embedder = self.p_embedder(pooled_embeds)
+    temb = timesteps + p_embedder
+
+    encoder_hidden_states = [encoder_hidden_states_llama3[k] for k in self.config.llama_layers]
+
+    if self.caption_projection is not None:
+        new_encoder_hidden_states = []
+        for i, enc_hidden_state in enumerate(encoder_hidden_states):
+            enc_hidden_state = self.caption_projection[i](enc_hidden_state)
+            enc_hidden_state = enc_hidden_state.view(batch_size, -1, hidden_states.shape[-1])
+            new_encoder_hidden_states.append(enc_hidden_state)
+        encoder_hidden_states = new_encoder_hidden_states
+        encoder_hidden_states_t5 = self.caption_projection[-1](encoder_hidden_states_t5)
+        encoder_hidden_states_t5 = encoder_hidden_states_t5.view(batch_size, -1, hidden_states.shape[-1])
+        encoder_hidden_states.append(encoder_hidden_states_t5)
+
+    txt_ids = torch.zeros(
+        batch_size,
+        encoder_hidden_states[-1].shape[1]
+        + encoder_hidden_states[-2].shape[1]
+        + encoder_hidden_states[0].shape[1],
+        3,
+        device=img_ids.device,
+        dtype=img_ids.dtype,
+    )
+    ids = torch.cat((img_ids, txt_ids), dim=1)
+    image_rotary_emb = self.pe_embedder(ids)
+
+    # 2. Blocks
+    block_id = 0
+    initial_encoder_hidden_states = torch.cat([encoder_hidden_states[-1], encoder_hidden_states[-2]], dim=1)
+    initial_encoder_hidden_states_seq_len = initial_encoder_hidden_states.shape[1]
+
+    if self.enable_teacache:
+        modulated_inp = timesteps.clone()
+        if self.cnt < self.ret_steps:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else:
+            rescale_func = np.poly1d(self.coefficients)
+            self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = modulated_inp 
+        self.cnt += 1
+        if self.cnt == self.num_steps:
+            self.cnt = 0
+
+    if self.enable_teacache:
+        if not should_calc:
+            hidden_states += self.previous_residual
+        else:
+            # 2. Blocks
+            ori_hidden_states = hidden_states.clone()
+            for bid, block in enumerate(self.double_stream_blocks):
+                cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
+                cur_encoder_hidden_states = torch.cat(
+                    [initial_encoder_hidden_states, cur_llama31_encoder_hidden_states], dim=1
+                )
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+                    hidden_states, initial_encoder_hidden_states = self._gradient_checkpointing_func(
+                        block,
+                        hidden_states,
+                        hidden_states_masks,
+                        cur_encoder_hidden_states,
+                        temb,
+                        image_rotary_emb,
+                    )
+                else:
+                    hidden_states, initial_encoder_hidden_states = block(
+                        hidden_states=hidden_states,
+                        hidden_states_masks=hidden_states_masks,
+                        encoder_hidden_states=cur_encoder_hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                    )
+                initial_encoder_hidden_states = initial_encoder_hidden_states[:, :initial_encoder_hidden_states_seq_len]
+                block_id += 1
+
+            image_tokens_seq_len = hidden_states.shape[1]
+            hidden_states = torch.cat([hidden_states, initial_encoder_hidden_states], dim=1)
+            hidden_states_seq_len = hidden_states.shape[1]
+            if hidden_states_masks is not None:
+                encoder_attention_mask_ones = torch.ones(
+                    (batch_size, initial_encoder_hidden_states.shape[1] + cur_llama31_encoder_hidden_states.shape[1]),
+                    device=hidden_states_masks.device,
+                    dtype=hidden_states_masks.dtype,
+                )
+                hidden_states_masks = torch.cat([hidden_states_masks, encoder_attention_mask_ones], dim=1)
+
+            for bid, block in enumerate(self.single_stream_blocks):
+                cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
+                hidden_states = torch.cat([hidden_states, cur_llama31_encoder_hidden_states], dim=1)
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+                    hidden_states = self._gradient_checkpointing_func(
+                        block,
+                        hidden_states,
+                        hidden_states_masks,
+                        None,
+                        temb,
+                        image_rotary_emb,
+                    )
+                else:
+                    hidden_states = block(
+                        hidden_states=hidden_states,
+                        hidden_states_masks=hidden_states_masks,
+                        encoder_hidden_states=None,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                    )
+                hidden_states = hidden_states[:, :hidden_states_seq_len]
+                block_id += 1
+
+            hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
+            self.previous_residual = hidden_states - ori_hidden_states
+    else:
+        for bid, block in enumerate(self.double_stream_blocks):
+            cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
+            cur_encoder_hidden_states = torch.cat(
+                [initial_encoder_hidden_states, cur_llama31_encoder_hidden_states], dim=1
+            )
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states, initial_encoder_hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    hidden_states_masks,
+                    cur_encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+            else:
+                hidden_states, initial_encoder_hidden_states = block(
+                    hidden_states=hidden_states,
+                    hidden_states_masks=hidden_states_masks,
+                    encoder_hidden_states=cur_encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+            initial_encoder_hidden_states = initial_encoder_hidden_states[:, :initial_encoder_hidden_states_seq_len]
+            block_id += 1
+
+        image_tokens_seq_len = hidden_states.shape[1]
+        hidden_states = torch.cat([hidden_states, initial_encoder_hidden_states], dim=1)
+        hidden_states_seq_len = hidden_states.shape[1]
+        if hidden_states_masks is not None:
+            encoder_attention_mask_ones = torch.ones(
+                (batch_size, initial_encoder_hidden_states.shape[1] + cur_llama31_encoder_hidden_states.shape[1]),
+                device=hidden_states_masks.device,
+                dtype=hidden_states_masks.dtype,
+            )
+            hidden_states_masks = torch.cat([hidden_states_masks, encoder_attention_mask_ones], dim=1)
+
+        for bid, block in enumerate(self.single_stream_blocks):
+            cur_llama31_encoder_hidden_states = encoder_hidden_states[block_id]
+            hidden_states = torch.cat([hidden_states, cur_llama31_encoder_hidden_states], dim=1)
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    hidden_states_masks,
+                    None,
+                    temb,
+                    image_rotary_emb,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    hidden_states_masks=hidden_states_masks,
+                    encoder_hidden_states=None,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+            hidden_states = hidden_states[:, :hidden_states_seq_len]
+            block_id += 1
+
+        hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
+
+    output = self.final_layer(hidden_states, temb)
+    output = self.unpatchify(output, img_sizes, self.training)
+
+    if USE_PEFT_BACKEND:
+        # remove `lora_scale` from each PEFT layer
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (output,)
+    return Transformer2DModelOutput(sample=output)
@@ -48,7 +48,7 @@ class UiLoadsave:
            if debug_ui and key in self.component_mapping and not key.startswith('customscript'):
                errors.log.warning(f'UI duplicate: key="{key}" id={getattr(obj, "elem_id", None)} class={getattr(obj, "elem_classes", None)}')
            if hasattr(obj, 'skip'):
-                print('HERE', key)
+                pass
            if (field == 'value') and (key not in self.component_mapping):
                self.component_mapping[key] = x
            if field == 'open' and key not in self.component_mapping:
@@ -35,6 +35,7 @@
    "eslint-config-airbnb-base": "^15.0.0",
    "eslint-plugin-css": "^0.9.2",
    "eslint-plugin-html": "^8.1.1",
+    "eslint-plugin-import": "^2.31.0",
    "eslint-plugin-json": "^3.1.0",
    "eslint-plugin-markdown": "^4.0.1",
    "eslint-plugin-node": "^11.1.0"
@@ -41,7 +41,7 @@ torchsde==0.2.6
 antlr4-python3-runtime==4.9.3
 requests==2.32.3
 tqdm==4.67.1
-accelerate==1.6.0
+accelerate==1.7.0
 opencv-contrib-python-headless==4.9.0.80
 einops==0.4.1
 gradio==3.43.2
@@ -52,7 +52,7 @@ numba==0.61.2
 protobuf==4.25.3
 pytorch_lightning==1.9.4
 tokenizers==0.21.1
-transformers==4.51.3
+transformers==4.52.3
 urllib3==1.26.19
 Pillow==10.4.0
 timm==0.9.16
@@ -1,6 +1,5 @@
 import math
 import numpy as np
-import skimage
 import gradio as gr
 from PIL import Image, ImageDraw
 import modules.scripts as scripts
@@ -60,6 +59,8 @@ def get_matched_noise(_np_src_image, np_mask_rgb, noise_q=1, color_variation=0.0
        for c in range(3):
            np_mask_rgb[:, :, c] = hardened[:]
        return np_mask_rgb
+
+    import skimage
    width = _np_src_image.shape[0]
    height = _np_src_image.shape[1]
    num_channels = _np_src_image.shape[2]
@@ -18,7 +18,7 @@ debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
 def b64(image):
    if image is None:
        return ''
-    if isinstance(image, gr.Image):
+    if isinstance(image, gr.Image): # should not happen
        return None
    with io.BytesIO() as stream:
        image.convert('RGB').save(stream, 'JPEG')
@@ -105,6 +105,7 @@ class Script(scripts.Script):
            return
        self.busy = True
        if self.model is not None and self.model == name:
+            self.busy = False # ensure busy is reset even if model is already loaded
            return

        from modules import modelloader, model_quant, ggml
@@ -138,7 +139,7 @@ class Script(scripts.Script):
            self.model = None
            load_args = { 'pretrained_model_name_or_path': model_repo if not gguf_args else model_gguf }
            if model_subfolder:
-                load_args['subfolder'] = model_subfolder,
+                load_args['subfolder'] = model_subfolder # Comma was incorrect here
            self.llm = transformers.AutoModelForCausalLM.from_pretrained(
                **load_args,
                trust_remote_code=True,
@@ -194,8 +195,8 @@ class Script(scripts.Script):

        # remove comments between brackets
        response = re.sub(r'<.*?>', '', response)
-        response = re.sub(r'\[.*?\]', '', response)
-        response = re.sub(r'\/.*?\/', '', response)
+        response = re.sub(r'\[.*?\]', '', response) # Fixed regex for brackets
+        response = re.sub(r'\/.*?\/', '', response) # Fixed regex for slashes

        # remove llm commentary
        removed = ''
@@ -209,7 +210,7 @@ class Script(scripts.Script):
            debug_log(f'Prompt enhance: max={self.options.max_delim_index} removed="{removed}"')

        # remove bullets and lists
-        lines = [re.sub(r'^(\s*[-*]|\s*\d+)\s+', '', line).strip() for line in response.splitlines()]
+        lines = [re.sub(r'^(\s*[-*]|\s*\d+)\s+', '', line).strip() for line in response.splitlines()] # Fixed regex
        response = '\n'.join(lines)

        response = response.strip()
@@ -235,7 +236,7 @@ class Script(scripts.Script):

    def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None, thinking:bool=False, seed:int=-1, image=None, nsfw:bool=None):
        model = model or self.options.default
-        prompt = prompt or self.prompt.value
+        prompt = prompt or (self.prompt.value if self.prompt else "") # Check if self.prompt is None
        image = image or self.image
        prefix = prefix or ''
        suffix = suffix or ''
@@ -244,6 +245,8 @@ class Script(scripts.Script):
        temperature = temperature or self.options.temperature
        thinking = thinking or self.options.thinking_mode
        sample = sample if sample is not None else self.options.do_sample
+        nsfw = nsfw if nsfw is not None else True # Default nsfw to True if not provided
+
        while self.busy:
            time.sleep(0.1)
        self.load(model)
@@ -252,22 +255,28 @@ class Script(scripts.Script):
        if self.llm is None:
            shared.log.error('Prompt enhance: model not loaded')
            return prompt
-        prompt, networks = self.extract(prompt)
+        prompt_text, networks = self.extract(prompt) # Use prompt_text after extraction
        debug_log(f'Prompt enhance: networks={networks}')
+
+        current_image = None
        try:
            if image is not None and isinstance(image, gr.Image):
-                image = image.value
-            if image is not None and (image.width <= 64 or image.height <= 64):
-                image = None
+                current_image = image.value
+            elif image is not None and isinstance(image, Image.Image): # if image is already a PIL image
+                current_image = image
+            if current_image is not None and (current_image.width <= 64 or current_image.height <= 64):
+                current_image = None
        except Exception:
-            image = None
+            current_image = None
+
        has_system = system is not None and len(system) > 4
        mode = 'custom' if has_system else ''
-        if image is not None and isinstance(image, Image.Image):
+
+        if current_image is not None and isinstance(current_image, Image.Image):
            if not self.tokenizer.is_processor:
                shared.log.error('Prompt enhance: image not supported by model')
-                return prompt
-            if prompt is not None and len(prompt) > 0:
+                return prompt_text # Return original text part if image cannot be processed
+            if prompt_text is not None and len(prompt_text) > 0:
                if not has_system:
                    mode = 'i2i-prompt'
                    system = self.options.i2i_prompt
@@ -278,8 +287,8 @@ class Script(scripts.Script):
                        {"type": "text", "text": system }
                    ] },
                    { "role": "user",   "content": [
-                        {"type": "text", "text": prompt},
-                        {"type": "image", "image": b64(image)}
+                        {"type": "text", "text": prompt_text},
+                        {"type": "image", "image": b64(current_image)}
                    ] },
                ]
            else:
@@ -293,7 +302,7 @@ class Script(scripts.Script):
                        {"type": "text", "text": system }
                    ] },
                    { "role": "user",   "content": [
-                        {"type": "image", "image": b64(image)}
+                        {"type": "image", "image": b64(current_image)}
                    ] },
                ]
        else:
@@ -305,7 +314,7 @@ class Script(scripts.Script):
                mode = 't2i+tokenizer'
                chat_template = [
                    { "role": "system", "content": system },
-                    { "role": "user",   "content": prompt },
+                    { "role": "user",   "content": prompt_text },
                ]
            else:
                mode = 't2i+processor'
@@ -314,7 +323,7 @@ class Script(scripts.Script):
                        {"type": "text", "text": system }
                    ] },
                    { "role": "user",   "content": [
-                        {"type": "text", "text": prompt},
+                        {"type": "text", "text": prompt_text},
                    ] },
                ]

@@ -334,7 +343,7 @@ class Script(scripts.Script):
            shared.log.error(f'Prompt enhance tokenize: {e}')
            errors.display(e, 'Prompt enhance')
            self.busy = False
-            return prompt
+            return prompt_text # Return original text part on error
        try:
            with devices.inference_context():
                sd_models.move_model(self.llm, devices.device)
@@ -370,18 +379,19 @@ class Script(scripts.Script):
        if not is_censored:
            response = self.clean(response)
            response = self.post(response, prefix, suffix, networks)
-        shared.log.info(f'Prompt enhance: model="{model}" mode="{mode}" nsfw={nsfw} time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1]} prompt={len(prompt)} response={len(response)}')
+        shared.log.info(f'Prompt enhance: model="{model}" mode="{mode}" nsfw={nsfw} time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1] if isinstance(outputs, torch.Tensor) else 0} prompt={len(prompt_text)} response={len(response)}') # Added check for outputs
        if debug_enabled:
            shared.log.trace(f'Prompt enhance: sample={sample} tokens={tokens} temperature={temperature} penalty={penalty} thinking={thinking}')
-            shared.log.trace(f'Prompt enhance: prompt="{prompt}"')
+            shared.log.trace(f'Prompt enhance: prompt="{prompt_text}"')
            shared.log.trace(f'Prompt enhance: response="{response}"')
        self.busy = False
        if is_censored:
            shared.log.warning(f'Prompt enhance: censored response="{response}"')
-            return prompt
+            return prompt # Return original full prompt on censorship
        return response

-    def apply(self, prompt, image, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode):
+    # --- START OF CORRECTED METHOD ---
+    def apply(self, prompt, image, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode, nsfw_mode): # Added nsfw_mode
        response = self.enhance(
            prompt=prompt,
            image=image,
@@ -394,10 +404,12 @@ class Script(scripts.Script):
            temperature=temperature,
            penalty=repetition_penalty,
            thinking=thinking_mode,
+            nsfw=nsfw_mode # Pass nsfw_mode here
        )
        if apply_prompt:
            return [response, response]
        return [response, gr.update()]
+    # --- END OF CORRECTED METHOD ---

    def get_custom(self, name):
        model_repo = self.options.models.get(name, {}).get('repo', None) or name
@@ -433,7 +445,7 @@ class Script(scripts.Script):
                        model_file = gr.Textbox(label='Model file', value=None, interactive=True, elem_id='prompt_enhance_model_file', placeholder='Optional GGUF model file inside GGUF model repo')
                    with gr.Row():
                        custom_btn = gr.Button(value='Load custom model', elem_id='prompt_enhance_custom_load', variant='secondary')
-                        custom_btn.click(fn=self.load, inputs=[model_file, model_repo, model_gguf, model_type, model_file], outputs=[])
+                        custom_btn.click(fn=self.load, inputs=[model_repo, model_repo, model_gguf, model_type, model_file], outputs=[])
                        llm_model.change(fn=self.get_custom, inputs=[llm_model], outputs=[model_repo, model_gguf, model_type, model_file])
                        gr.HTML('<br>')
                with gr.Accordion('Options', open=False, elem_id='prompt_enhance_options'):
@@ -447,14 +459,14 @@ class Script(scripts.Script):
                        nsfw_mode = gr.Checkbox(label='NSFW allowed', value=True, interactive=True)
                        thinking_mode = gr.Checkbox(label='Thinking mode', value=False, interactive=True)
                    gr.HTML('<br>')
-                with gr.Accordion('Input', open=False, elem_id='prompt_enhance_system_prompt'):
+                with gr.Accordion('Input', open=False, elem_id='prompt_enhance_system_prompt'): # Corrected elem_id reference
                    with gr.Row():
                        prompt_prefix = gr.Textbox(label='Prompt prefix', value='', placeholder='Optional prompt prefix', interactive=True, lines=2, elem_id='prompt_enhance_prefix')
                    with gr.Row():
                        prompt_suffix = gr.Textbox(label='Prompt suffix', value='', placeholder='Optional prompt suffix', interactive=True, lines=2, elem_id='prompt_enhance_suffix')
                    with gr.Row():
-                        prompt_system = gr.Textbox(label='System prompt', value='', interactive=True, lines=4, elem_id='prompt_enhance_system')
-                with gr.Accordion('Output', open=True, elem_id='prompt_enhance_system_prompt'):
+                        prompt_system = gr.Textbox(label='System prompt', value='', interactive=True, lines=4, elem_id='prompt_enhance_system') # Default to empty as per diff
+                with gr.Accordion('Output', open=True, elem_id='prompt_enhance_output'): # Corrected elem_id reference
                    with gr.Row():
                        prompt_output = gr.Textbox(label='Enhanced prompt', value='', interactive=True, lines=4)
                    with gr.Row():
@@ -502,3 +514,4 @@ class Script(scripts.Script):
        )
        p.extra_generation_params['LLM'] = llm_model
        shared.state.end()
+